Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: defer resident Memory loading until first run

Browse Source
This commit is contained in:
Hare
2026-09-05 07:52:47 +09:00
parent 15bc299987
commit 69824ea45d
3 changed files with 145 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+2 -3
View File
@@ -997,11 +997,10 @@ where
worker.manifest(),
worker.workspace_client_handle(),
worker.prompts().load_full(),
)
.await?;
)?;
let memory_prompt_contribution = memory_install_plan.as_ref().map(|plan| {
(
plan.resident_summary.clone(),
plan.resident_summary_source.clone(),
plan.system_prompt_override.clone(),
)
});
crates/worker/src/feature/builtin/memory.rs
+65 -45
View File
@@ -23,7 +23,8 @@ use crate::feature::{
ToolDeclaration,
};
use crate::worker::{
WorkspaceClient, WorkspaceClientError, WorkspaceRequest, WorkspaceRequestMethod,
SystemPromptContributionSource, WorkspaceClient, WorkspaceClientError, WorkspaceRequest,
WorkspaceRequestMethod,
};
#[derive(Clone, Debug)]
@@ -342,15 +343,44 @@ fn query_schema() -> serde_json::Value {
})
}
pub struct MemoryFeatureInstallPlan {
pub module: MemoryToolsFeature,
pub resident_summary: Option<String>,
pub system_prompt_override: Option<String>,
struct WorkspaceResidentSummarySource {
client: Arc<dyn WorkspaceClient>,
}
#[async_trait]
impl SystemPromptContributionSource for WorkspaceResidentSummarySource {
async fn load(&self) -> Option<String> {
match self
.client
.execute_memory_backend_operation(
memory::backend::MemoryBackendOperation::ResidentSummary(
memory::backend::MemoryResidentSummaryOperation::default(),
),
)
.await
{
Ok(memory::backend::MemoryBackendOperationResult::ToolOutput(output)) => output.content,
Ok(other) => {
tracing::debug!(?other, "unexpected resident Memory Backend result");
None
}
Err(error) => {
tracing::debug!(%error, "resident Memory summary unavailable");
None
}
}
}
}
pub(crate) struct MemoryFeatureInstallPlan {
pub(crate) module: MemoryToolsFeature,
pub(crate) resident_summary_source: Option<Arc<dyn SystemPromptContributionSource>>,
pub(crate) system_prompt_override: Option<String>,
pub(crate) resolved_config: manifest::ResolvedMemoryFeatureConfig,
}
impl MemoryFeatureInstallPlan {
pub async fn prepare(
pub fn prepare(
manifest: &manifest::WorkerManifest,
client: Arc<dyn WorkspaceClient>,
prompts: Arc<crate::prompt::catalog::PromptCatalog>,
@@ -361,10 +391,9 @@ impl MemoryFeatureInstallPlan {
prompts,
manifest.profile.clone(),
)
.await
}
async fn prepare_resolved(
fn prepare_resolved(
config: manifest::ResolvedMemoryFeatureConfig,
client: Arc<dyn WorkspaceClient>,
prompts: Arc<crate::prompt::catalog::PromptCatalog>,
@@ -405,30 +434,11 @@ impl MemoryFeatureInstallPlan {
));
}
let resident_summary = if config.profile.resident.inject_summary {
match client
.execute_memory_backend_operation(
memory::backend::MemoryBackendOperation::ResidentSummary(
memory::backend::MemoryResidentSummaryOperation::default(),
),
)
.await
{
Ok(memory::backend::MemoryBackendOperationResult::ToolOutput(output)) => {
output.content
}
Ok(other) => {
tracing::debug!(?other, "unexpected resident Memory Backend result");
None
}
Err(error) => {
tracing::debug!(%error, "resident Memory summary unavailable");
None
}
}
} else {
None
};
let resident_summary_source = config.profile.resident.inject_summary.then(|| {
Arc::new(WorkspaceResidentSummarySource {
client: Arc::clone(&client),
}) as Arc<dyn SystemPromptContributionSource>
});
let system_prompt_override = if memory_consolidation_worker {
let language = settings.language;
Some(
@@ -442,7 +452,7 @@ impl MemoryFeatureInstallPlan {
Ok(Some(Self {
module: MemoryToolsFeature::new(client, config.profile.staging_tools),
resident_summary,
resident_summary_source,
system_prompt_override,
resolved_config: config,
}))
@@ -559,7 +569,6 @@ mod tests {
prompts.clone(),
None,
)
.await
.unwrap();
assert!(disabled.is_none());
@@ -573,7 +582,6 @@ mod tests {
prompts.clone(),
None,
)
.await
.is_err()
);
enabled
@@ -592,7 +600,6 @@ mod tests {
prompts.clone(),
None,
)
.await
.is_err()
);
let plan = MemoryFeatureInstallPlan::prepare_resolved(
@@ -601,10 +608,9 @@ mod tests {
prompts.clone(),
None,
)
.await
.unwrap()
.unwrap();
assert!(plan.resident_summary.is_none());
assert!(plan.resident_summary_source.is_none());
assert!(plan.system_prompt_override.is_none());
enabled.profile.resident.inject_summary = true;
@@ -614,14 +620,20 @@ mod tests {
prompts,
None,
)
.await
.unwrap()
.unwrap();
assert_eq!(plan.resident_summary.as_deref(), Some("# Durable Memory"));
assert_eq!(
plan.resident_summary_source
.unwrap()
.load()
.await
.as_deref(),
Some("# Durable Memory")
);
}
#[tokio::test]
async fn memory_prompt_contribution_rereads_resident_summary_for_each_install() {
async fn memory_prompt_contribution_defers_resident_summary_until_loaded() {
let prompts = crate::prompt::catalog::PromptCatalog::builtins_only().unwrap();
let mut config = manifest::ResolvedMemoryFeatureConfig::default();
config.profile.enabled = true;
@@ -639,7 +651,6 @@ mod tests {
prompts.clone(),
None,
)
.await
.unwrap()
.unwrap();
let restored = MemoryFeatureInstallPlan::prepare_resolved(
@@ -648,16 +659,25 @@ mod tests {
prompts,
None,
)
.await
.unwrap()
.unwrap();
assert_eq!(
first.resident_summary.as_deref(),
first
.resident_summary_source
.unwrap()
.load()
.await
.as_deref(),
Some("first resident summary")
);
assert_eq!(
restored.resident_summary.as_deref(),
restored
.resident_summary_source
.unwrap()
.load()
.await
.as_deref(),
Some("updated resident summary")
);
}
crates/worker/src/worker.rs
+78 -16
View File
@@ -1087,6 +1087,11 @@ impl WorkerSession {
}
}
#[async_trait::async_trait]
pub(crate) trait SystemPromptContributionSource: Send + Sync {
async fn load(&self) -> Option<String>;
}
/// An independent agent execution unit.
///
/// Holds a [`Engine`] directly and persists session state via
@@ -1230,8 +1235,8 @@ pub struct Worker<C: LlmClient, St: Store> {
prompts: Arc<ArcSwap<PromptCatalog>>,
/// Test/internal policy gate for installed resident prompt contributions.
inject_resident_summary: bool,
/// Materialized resident prompt context installed by an enabled Feature.
feature_resident_summary: Option<String>,
/// Deferred resident prompt source installed by an enabled Feature.
feature_resident_summary_source: Option<Arc<dyn SystemPromptContributionSource>>,
/// Complete system prompt replacement installed by an enabled Feature.
feature_system_prompt_override: Option<String>,
/// Typed user submissions in submit order. K-th entry corresponds to
@@ -1452,7 +1457,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
runtime_ticket_role: None,
prompts,
inject_resident_summary: true,
feature_resident_summary: None,
feature_resident_summary_source: None,
feature_system_prompt_override: None,
user_segments: Vec::new(),
sink: SegmentLogSink::new(),
@@ -1498,10 +1503,10 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
pub(crate) fn install_system_prompt_contribution(
&mut self,
resident_summary: Option<String>,
resident_summary_source: Option<Arc<dyn SystemPromptContributionSource>>,
system_prompt_override: Option<String>,
) {
self.feature_resident_summary = resident_summary;
self.feature_resident_summary_source = resident_summary_source;
self.feature_system_prompt_override = system_prompt_override;
}
@@ -2490,10 +2495,14 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
}
}
}
let resident_summary = self
.inject_resident_summary
.then(|| self.feature_resident_summary.clone())
.flatten();
let resident_summary = if self.inject_resident_summary {
match &self.feature_resident_summary_source {
Some(source) => source.load().await,
None => None,
}
} else {
None
};
let worker_language = worker_language(&self.manifest.engine);
let scope_snapshot = self.scope.snapshot();
let cwd_for_prompt = self
@@ -4661,7 +4670,7 @@ where
runtime_ticket_role: None,
prompts: common.prompts,
inject_resident_summary: true,
feature_resident_summary: None,
feature_resident_summary_source: None,
feature_system_prompt_override: None,
user_segments: Vec::new(),
sink: SegmentLogSink::new(),
@@ -4745,7 +4754,7 @@ where
runtime_ticket_role: None,
prompts: common.prompts,
inject_resident_summary: true,
feature_resident_summary: None,
feature_resident_summary_source: None,
feature_system_prompt_override: None,
user_segments: Vec::new(),
sink: SegmentLogSink::new(),
@@ -4863,7 +4872,7 @@ where
runtime_ticket_role: None,
prompts: common.prompts,
inject_resident_summary: true,
feature_resident_summary: None,
feature_resident_summary_source: None,
feature_system_prompt_override: None,
user_segments: Vec::new(),
sink: SegmentLogSink::new(),
@@ -5238,7 +5247,7 @@ where
runtime_ticket_role: None,
prompts: common.prompts,
inject_resident_summary: true,
feature_resident_summary: None,
feature_resident_summary_source: None,
feature_system_prompt_override: None,
user_segments: state.user_segments,
// Seed the mirror with the entries we just replayed so a
@@ -6609,6 +6618,32 @@ permission = "read"
mod build_summary_prompt_tests {
use super::*;
struct TestSystemPromptContributionSource {
value: Option<String>,
load_count: Arc<AtomicUsize>,
}
#[async_trait::async_trait]
impl SystemPromptContributionSource for TestSystemPromptContributionSource {
async fn load(&self) -> Option<String> {
self.load_count.fetch_add(1, Ordering::SeqCst);
self.value.clone()
}
}
fn test_system_prompt_contribution_source(
value: Option<String>,
) -> (Arc<dyn SystemPromptContributionSource>, Arc<AtomicUsize>) {
let load_count = Arc::new(AtomicUsize::new(0));
(
Arc::new(TestSystemPromptContributionSource {
value,
load_count: Arc::clone(&load_count),
}),
load_count,
)
}
fn test_summary_input(items: &[Item]) -> String {
build_summary_input(
items,
@@ -8065,6 +8100,32 @@ mod build_summary_prompt_tests {
}
}
#[tokio::test]
async fn worker_without_initial_system_prompt_does_not_load_feature_contribution() {
let dir = tempfile::tempdir().unwrap();
let cwd = dir.path().join("workspace");
std::fs::create_dir_all(&cwd).unwrap();
let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap();
let mut worker = Worker::new(
minimal_manifest(),
Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient),
store,
WorkerWorkspaceContext::no_workspace(),
WorkerFilesystemAuthority::local(cwd.clone(), cwd.clone()),
Scope::writable(&cwd).unwrap(),
)
.await
.unwrap();
let (source, load_count) =
test_system_prompt_contribution_source(Some("# Durable Memory".to_string()));
worker.install_system_prompt_contribution(Some(source), None);
worker.ensure_system_prompt_materialized().await.unwrap();
assert_eq!(load_count.load(Ordering::SeqCst), 0);
assert!(worker.history().is_empty());
}
#[tokio::test]
async fn memory_consolidation_prompt_uses_bound_workspace_language() {
let dir = tempfile::tempdir().unwrap();
@@ -8163,16 +8224,17 @@ mod build_summary_prompt_tests {
.await
.unwrap();
worker.set_resident_memory_injection(gates.summary);
let resident_summary = if memory_config
let resident_summary_source = if memory_config
.as_ref()
.is_some_and(|cfg| cfg.profile.resident.inject_summary)
&& gates.summary
{
summary_doc.and_then(summary_content_for_backend)
let summary = summary_doc.and_then(summary_content_for_backend);
Some(test_system_prompt_contribution_source(summary).0)
} else {
None
};
worker.install_system_prompt_contribution(resident_summary, None);
worker.install_system_prompt_contribution(resident_summary_source, None);
let template = SystemPromptTemplate::parse(
"default",
crate::prompt::source::PromptCatalogSource::builtins_only(),