fix: defer resident Memory loading until first run

This commit is contained in:
2026-09-05 07:52:47 +09:00
parent 15bc299987
commit 69824ea45d
3 changed files with 145 additions and 64 deletions
+2 -3
View File
@@ -997,11 +997,10 @@ where
worker.manifest(), worker.manifest(),
worker.workspace_client_handle(), worker.workspace_client_handle(),
worker.prompts().load_full(), worker.prompts().load_full(),
) )?;
.await?;
let memory_prompt_contribution = memory_install_plan.as_ref().map(|plan| { 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(), plan.system_prompt_override.clone(),
) )
}); });
+65 -45
View File
@@ -23,7 +23,8 @@ use crate::feature::{
ToolDeclaration, ToolDeclaration,
}; };
use crate::worker::{ use crate::worker::{
WorkspaceClient, WorkspaceClientError, WorkspaceRequest, WorkspaceRequestMethod, SystemPromptContributionSource, WorkspaceClient, WorkspaceClientError, WorkspaceRequest,
WorkspaceRequestMethod,
}; };
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
@@ -342,15 +343,44 @@ fn query_schema() -> serde_json::Value {
}) })
} }
pub struct MemoryFeatureInstallPlan { struct WorkspaceResidentSummarySource {
pub module: MemoryToolsFeature, client: Arc<dyn WorkspaceClient>,
pub resident_summary: Option<String>, }
pub system_prompt_override: Option<String>,
#[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, pub(crate) resolved_config: manifest::ResolvedMemoryFeatureConfig,
} }
impl MemoryFeatureInstallPlan { impl MemoryFeatureInstallPlan {
pub async fn prepare( pub fn prepare(
manifest: &manifest::WorkerManifest, manifest: &manifest::WorkerManifest,
client: Arc<dyn WorkspaceClient>, client: Arc<dyn WorkspaceClient>,
prompts: Arc<crate::prompt::catalog::PromptCatalog>, prompts: Arc<crate::prompt::catalog::PromptCatalog>,
@@ -361,10 +391,9 @@ impl MemoryFeatureInstallPlan {
prompts, prompts,
manifest.profile.clone(), manifest.profile.clone(),
) )
.await
} }
async fn prepare_resolved( fn prepare_resolved(
config: manifest::ResolvedMemoryFeatureConfig, config: manifest::ResolvedMemoryFeatureConfig,
client: Arc<dyn WorkspaceClient>, client: Arc<dyn WorkspaceClient>,
prompts: Arc<crate::prompt::catalog::PromptCatalog>, prompts: Arc<crate::prompt::catalog::PromptCatalog>,
@@ -405,30 +434,11 @@ impl MemoryFeatureInstallPlan {
)); ));
} }
let resident_summary = if config.profile.resident.inject_summary { let resident_summary_source = config.profile.resident.inject_summary.then(|| {
match client Arc::new(WorkspaceResidentSummarySource {
.execute_memory_backend_operation( client: Arc::clone(&client),
memory::backend::MemoryBackendOperation::ResidentSummary( }) as Arc<dyn SystemPromptContributionSource>
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 system_prompt_override = if memory_consolidation_worker { let system_prompt_override = if memory_consolidation_worker {
let language = settings.language; let language = settings.language;
Some( Some(
@@ -442,7 +452,7 @@ impl MemoryFeatureInstallPlan {
Ok(Some(Self { Ok(Some(Self {
module: MemoryToolsFeature::new(client, config.profile.staging_tools), module: MemoryToolsFeature::new(client, config.profile.staging_tools),
resident_summary, resident_summary_source,
system_prompt_override, system_prompt_override,
resolved_config: config, resolved_config: config,
})) }))
@@ -559,7 +569,6 @@ mod tests {
prompts.clone(), prompts.clone(),
None, None,
) )
.await
.unwrap(); .unwrap();
assert!(disabled.is_none()); assert!(disabled.is_none());
@@ -573,7 +582,6 @@ mod tests {
prompts.clone(), prompts.clone(),
None, None,
) )
.await
.is_err() .is_err()
); );
enabled enabled
@@ -592,7 +600,6 @@ mod tests {
prompts.clone(), prompts.clone(),
None, None,
) )
.await
.is_err() .is_err()
); );
let plan = MemoryFeatureInstallPlan::prepare_resolved( let plan = MemoryFeatureInstallPlan::prepare_resolved(
@@ -601,10 +608,9 @@ mod tests {
prompts.clone(), prompts.clone(),
None, None,
) )
.await
.unwrap() .unwrap()
.unwrap(); .unwrap();
assert!(plan.resident_summary.is_none()); assert!(plan.resident_summary_source.is_none());
assert!(plan.system_prompt_override.is_none()); assert!(plan.system_prompt_override.is_none());
enabled.profile.resident.inject_summary = true; enabled.profile.resident.inject_summary = true;
@@ -614,14 +620,20 @@ mod tests {
prompts, prompts,
None, None,
) )
.await
.unwrap() .unwrap()
.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] #[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 prompts = crate::prompt::catalog::PromptCatalog::builtins_only().unwrap();
let mut config = manifest::ResolvedMemoryFeatureConfig::default(); let mut config = manifest::ResolvedMemoryFeatureConfig::default();
config.profile.enabled = true; config.profile.enabled = true;
@@ -639,7 +651,6 @@ mod tests {
prompts.clone(), prompts.clone(),
None, None,
) )
.await
.unwrap() .unwrap()
.unwrap(); .unwrap();
let restored = MemoryFeatureInstallPlan::prepare_resolved( let restored = MemoryFeatureInstallPlan::prepare_resolved(
@@ -648,16 +659,25 @@ mod tests {
prompts, prompts,
None, None,
) )
.await
.unwrap() .unwrap()
.unwrap(); .unwrap();
assert_eq!( assert_eq!(
first.resident_summary.as_deref(), first
.resident_summary_source
.unwrap()
.load()
.await
.as_deref(),
Some("first resident summary") Some("first resident summary")
); );
assert_eq!( assert_eq!(
restored.resident_summary.as_deref(), restored
.resident_summary_source
.unwrap()
.load()
.await
.as_deref(),
Some("updated resident summary") Some("updated resident summary")
); );
} }
+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. /// An independent agent execution unit.
/// ///
/// Holds a [`Engine`] directly and persists session state via /// Holds a [`Engine`] directly and persists session state via
@@ -1230,8 +1235,8 @@ pub struct Worker<C: LlmClient, St: Store> {
prompts: Arc<ArcSwap<PromptCatalog>>, prompts: Arc<ArcSwap<PromptCatalog>>,
/// Test/internal policy gate for installed resident prompt contributions. /// Test/internal policy gate for installed resident prompt contributions.
inject_resident_summary: bool, inject_resident_summary: bool,
/// Materialized resident prompt context installed by an enabled Feature. /// Deferred resident prompt source installed by an enabled Feature.
feature_resident_summary: Option<String>, feature_resident_summary_source: Option<Arc<dyn SystemPromptContributionSource>>,
/// Complete system prompt replacement installed by an enabled Feature. /// Complete system prompt replacement installed by an enabled Feature.
feature_system_prompt_override: Option<String>, feature_system_prompt_override: Option<String>,
/// Typed user submissions in submit order. K-th entry corresponds to /// 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, runtime_ticket_role: None,
prompts, prompts,
inject_resident_summary: true, inject_resident_summary: true,
feature_resident_summary: None, feature_resident_summary_source: None,
feature_system_prompt_override: None, feature_system_prompt_override: None,
user_segments: Vec::new(), user_segments: Vec::new(),
sink: SegmentLogSink::new(), sink: SegmentLogSink::new(),
@@ -1498,10 +1503,10 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
pub(crate) fn install_system_prompt_contribution( pub(crate) fn install_system_prompt_contribution(
&mut self, &mut self,
resident_summary: Option<String>, resident_summary_source: Option<Arc<dyn SystemPromptContributionSource>>,
system_prompt_override: Option<String>, 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; 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 let resident_summary = if self.inject_resident_summary {
.inject_resident_summary match &self.feature_resident_summary_source {
.then(|| self.feature_resident_summary.clone()) Some(source) => source.load().await,
.flatten(); None => None,
}
} else {
None
};
let worker_language = worker_language(&self.manifest.engine); let worker_language = worker_language(&self.manifest.engine);
let scope_snapshot = self.scope.snapshot(); let scope_snapshot = self.scope.snapshot();
let cwd_for_prompt = self let cwd_for_prompt = self
@@ -4661,7 +4670,7 @@ where
runtime_ticket_role: None, runtime_ticket_role: None,
prompts: common.prompts, prompts: common.prompts,
inject_resident_summary: true, inject_resident_summary: true,
feature_resident_summary: None, feature_resident_summary_source: None,
feature_system_prompt_override: None, feature_system_prompt_override: None,
user_segments: Vec::new(), user_segments: Vec::new(),
sink: SegmentLogSink::new(), sink: SegmentLogSink::new(),
@@ -4745,7 +4754,7 @@ where
runtime_ticket_role: None, runtime_ticket_role: None,
prompts: common.prompts, prompts: common.prompts,
inject_resident_summary: true, inject_resident_summary: true,
feature_resident_summary: None, feature_resident_summary_source: None,
feature_system_prompt_override: None, feature_system_prompt_override: None,
user_segments: Vec::new(), user_segments: Vec::new(),
sink: SegmentLogSink::new(), sink: SegmentLogSink::new(),
@@ -4863,7 +4872,7 @@ where
runtime_ticket_role: None, runtime_ticket_role: None,
prompts: common.prompts, prompts: common.prompts,
inject_resident_summary: true, inject_resident_summary: true,
feature_resident_summary: None, feature_resident_summary_source: None,
feature_system_prompt_override: None, feature_system_prompt_override: None,
user_segments: Vec::new(), user_segments: Vec::new(),
sink: SegmentLogSink::new(), sink: SegmentLogSink::new(),
@@ -5238,7 +5247,7 @@ where
runtime_ticket_role: None, runtime_ticket_role: None,
prompts: common.prompts, prompts: common.prompts,
inject_resident_summary: true, inject_resident_summary: true,
feature_resident_summary: None, feature_resident_summary_source: None,
feature_system_prompt_override: None, feature_system_prompt_override: None,
user_segments: state.user_segments, user_segments: state.user_segments,
// Seed the mirror with the entries we just replayed so a // Seed the mirror with the entries we just replayed so a
@@ -6609,6 +6618,32 @@ permission = "read"
mod build_summary_prompt_tests { mod build_summary_prompt_tests {
use super::*; 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 { fn test_summary_input(items: &[Item]) -> String {
build_summary_input( build_summary_input(
items, 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] #[tokio::test]
async fn memory_consolidation_prompt_uses_bound_workspace_language() { async fn memory_consolidation_prompt_uses_bound_workspace_language() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
@@ -8163,16 +8224,17 @@ mod build_summary_prompt_tests {
.await .await
.unwrap(); .unwrap();
worker.set_resident_memory_injection(gates.summary); worker.set_resident_memory_injection(gates.summary);
let resident_summary = if memory_config let resident_summary_source = if memory_config
.as_ref() .as_ref()
.is_some_and(|cfg| cfg.profile.resident.inject_summary) .is_some_and(|cfg| cfg.profile.resident.inject_summary)
&& gates.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 { } else {
None None
}; };
worker.install_system_prompt_contribution(resident_summary, None); worker.install_system_prompt_contribution(resident_summary_source, None);
let template = SystemPromptTemplate::parse( let template = SystemPromptTemplate::parse(
"default", "default",
crate::prompt::source::PromptCatalogSource::builtins_only(), crate::prompt::source::PromptCatalogSource::builtins_only(),