diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 4e431d3d..70584074 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -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(), ) }); diff --git a/crates/worker/src/feature/builtin/memory.rs b/crates/worker/src/feature/builtin/memory.rs index 33c5eff3..ea4c68b1 100644 --- a/crates/worker/src/feature/builtin/memory.rs +++ b/crates/worker/src/feature/builtin/memory.rs @@ -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, - pub system_prompt_override: Option, +struct WorkspaceResidentSummarySource { + client: Arc, +} + +#[async_trait] +impl SystemPromptContributionSource for WorkspaceResidentSummarySource { + async fn load(&self) -> Option { + 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>, + pub(crate) system_prompt_override: Option, pub(crate) resolved_config: manifest::ResolvedMemoryFeatureConfig, } impl MemoryFeatureInstallPlan { - pub async fn prepare( + pub fn prepare( manifest: &manifest::WorkerManifest, client: Arc, prompts: Arc, @@ -361,10 +391,9 @@ impl MemoryFeatureInstallPlan { prompts, manifest.profile.clone(), ) - .await } - async fn prepare_resolved( + fn prepare_resolved( config: manifest::ResolvedMemoryFeatureConfig, client: Arc, prompts: Arc, @@ -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 + }); 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") ); } diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index d4a3b85c..55d041ff 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -1087,6 +1087,11 @@ impl WorkerSession { } } +#[async_trait::async_trait] +pub(crate) trait SystemPromptContributionSource: Send + Sync { + async fn load(&self) -> Option; +} + /// An independent agent execution unit. /// /// Holds a [`Engine`] directly and persists session state via @@ -1230,8 +1235,8 @@ pub struct Worker { prompts: Arc>, /// 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, + /// Deferred resident prompt source installed by an enabled Feature. + feature_resident_summary_source: Option>, /// Complete system prompt replacement installed by an enabled Feature. feature_system_prompt_override: Option, /// Typed user submissions in submit order. K-th entry corresponds to @@ -1452,7 +1457,7 @@ impl Worker { 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 Worker { pub(crate) fn install_system_prompt_contribution( &mut self, - resident_summary: Option, + resident_summary_source: Option>, system_prompt_override: Option, ) { - 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 Worker { } } } - 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, + load_count: Arc, + } + + #[async_trait::async_trait] + impl SystemPromptContributionSource for TestSystemPromptContributionSource { + async fn load(&self) -> Option { + self.load_count.fetch_add(1, Ordering::SeqCst); + self.value.clone() + } + } + + fn test_system_prompt_contribution_source( + value: Option, + ) -> (Arc, Arc) { + 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(),