diff --git a/crates/manifest/src/lib.rs b/crates/manifest/src/lib.rs index 5ebdc72f..a7908702 100644 --- a/crates/manifest/src/lib.rs +++ b/crates/manifest/src/lib.rs @@ -956,6 +956,34 @@ impl WorkerManifest { } } +#[derive(Debug, Default, Deserialize)] +#[serde(default, deny_unknown_fields)] +struct LegacyMemoryFeatureConfig { + enabled: bool, + staging: bool, +} + +#[derive(Debug, Default, Deserialize)] +#[serde(default, deny_unknown_fields)] +struct LegacyMemoryConfig { + #[serde(rename = "workspace_root")] + _workspace_root: Option, + #[serde(rename = "query_result_limit")] + _query_result_limit: Option, + #[serde(rename = "query_excerpt_lines")] + _query_excerpt_lines: Option, + inject_summary: Option, + workspace_id: Option, + settings_revision: Option, + language: Option, + extract_model: Option, + extract_threshold: Option, + extract_worker_max_turns: Option, + consolidation_model: Option, + consolidation_threshold_files: Option, + consolidation_threshold_bytes: Option, +} + const RESOLVED_MANIFEST_SNAPSHOT_SCHEMA_VERSION: u64 = 2; /// Serialize a resolved Worker Manifest for durable Worker-specific storage. @@ -1004,12 +1032,37 @@ pub fn read_persisted_worker_manifest_snapshot( "resolved Worker manifest snapshot is missing manifest", )) })?; - return serde_json::from_value(manifest); + if manifest + .as_object() + .is_some_and(|manifest| manifest.contains_key("memory")) + { + return Err(serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "current resolved Worker manifest contains removed top-level memory authority", + ))); + } + return validate_persisted_worker_manifest(serde_json::from_value(manifest)?); } migrate_legacy_resolved_manifest_snapshot(snapshot) } +fn validate_persisted_worker_manifest( + manifest: WorkerManifest, +) -> Result { + manifest + .feature + .memory + .validate_execution() + .map_err(|message| { + serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + message, + )) + })?; + Ok(manifest) +} + fn migrate_legacy_resolved_manifest_snapshot( mut snapshot: serde_json::Value, ) -> Result { @@ -1030,44 +1083,21 @@ fn migrate_legacy_resolved_manifest_snapshot( "legacy resolved Worker manifest feature must be an object", )) })?; - let legacy_feature_memory = feature - .remove("memory") - .unwrap_or_else(|| serde_json::json!({})); - let legacy_feature_memory = legacy_feature_memory.as_object().ok_or_else(|| { - serde_json::Error::io(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "legacy resolved Worker manifest feature.memory must be an object", - )) - })?; - if legacy_feature_memory - .keys() - .any(|key| !matches!(key.as_str(), "enabled" | "staging")) - { - return Err(serde_json::Error::io(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "legacy resolved Worker manifest mixes old and new Memory configuration", - ))); - } - let enabled = legacy_feature_memory - .get("enabled") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false); - let staging_tools = legacy_feature_memory - .get("staging") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false); + let legacy_feature_memory: LegacyMemoryFeatureConfig = serde_json::from_value( + feature + .remove("memory") + .unwrap_or_else(|| serde_json::json!({})), + )?; + let enabled = legacy_feature_memory.enabled; + let staging_tools = legacy_feature_memory.staging; - let legacy_memory = legacy_memory.unwrap_or_else(|| serde_json::json!({})); - let legacy_memory = legacy_memory.as_object().ok_or_else(|| { - serde_json::Error::io(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "legacy resolved Worker manifest memory must be an object", - )) - })?; - let workspace_id = legacy_memory.get("workspace_id").cloned(); - let settings_revision = legacy_memory.get("settings_revision").cloned(); - let language = legacy_memory.get("language").cloned(); - let workspace_settings = match (workspace_id, settings_revision, language) { + let legacy_memory: LegacyMemoryConfig = + serde_json::from_value(legacy_memory.unwrap_or_else(|| serde_json::json!({})))?; + let mut workspace_settings = match ( + legacy_memory.workspace_id, + legacy_memory.settings_revision, + legacy_memory.language, + ) { (Some(workspace_id), Some(settings_revision), Some(language)) => Some(serde_json::json!({ "workspace_id": workspace_id, "settings_revision": settings_revision, @@ -1081,27 +1111,20 @@ fn migrate_legacy_resolved_manifest_snapshot( ))); } }; - let extraction_threshold = legacy_memory - .get("extract_threshold") - .cloned() - .unwrap_or(serde_json::Value::Null); - let extraction_enabled = !extraction_threshold.is_null(); - if legacy_memory - .get("consolidation_model") - .is_some_and(|model| !model.is_null()) - { + if !enabled { + workspace_settings = None; + } + let extraction_enabled = legacy_memory.extract_threshold.is_some(); + if legacy_memory.consolidation_model.is_some() { return Err(serde_json::Error::io(std::io::Error::new( std::io::ErrorKind::InvalidData, "legacy resolved Worker manifest uses a Worker-owned consolidation model that cannot be migrated to Backend authority", ))); } - let threshold_files = legacy_memory - .get("consolidation_threshold_files") - .and_then(serde_json::Value::as_u64); - let threshold_bytes = legacy_memory - .get("consolidation_threshold_bytes") - .and_then(serde_json::Value::as_u64); - let consolidation_enabled = match (threshold_files, threshold_bytes) { + let consolidation_enabled = match ( + legacy_memory.consolidation_threshold_files, + legacy_memory.consolidation_threshold_bytes, + ) { (None, None) => false, (Some(5), Some(50_000)) => true, _ => { @@ -1116,19 +1139,13 @@ fn migrate_legacy_resolved_manifest_snapshot( "enabled": enabled, "staging_tools": staging_tools, "resident": { - "inject_summary": legacy_memory - .get("inject_summary") - .and_then(serde_json::Value::as_bool) - .unwrap_or(true), + "inject_summary": legacy_memory.inject_summary.unwrap_or(true), }, "extraction": { "enabled": extraction_enabled, - "model": legacy_memory.get("extract_model").cloned().unwrap_or(serde_json::Value::Null), - "threshold": extraction_threshold, - "worker_max_turns": legacy_memory - .get("extract_worker_max_turns") - .cloned() - .unwrap_or(serde_json::Value::Null), + "model": serde_json::to_value(legacy_memory.extract_model)?, + "threshold": legacy_memory.extract_threshold, + "worker_max_turns": legacy_memory.extract_worker_max_turns, }, "consolidation": { "request_enabled": consolidation_enabled, @@ -1142,7 +1159,7 @@ fn migrate_legacy_resolved_manifest_snapshot( .insert("workspace_settings".to_string(), workspace_settings); } feature.insert("memory".to_string(), resolved); - serde_json::from_value(snapshot) + validate_persisted_worker_manifest(serde_json::from_value(snapshot)?) } #[cfg(test)] @@ -1566,6 +1583,18 @@ model_id = "claude-sonnet-4-20250514" let current = write_persisted_worker_manifest_snapshot(&migrated).unwrap(); assert_eq!(current["schema_version"], 2); assert!(current["manifest"].get("memory").is_none()); + + let mut disabled = + serde_json::to_value(WorkerManifest::from_toml(MINIMAL_REQUIRED).unwrap()).unwrap(); + disabled["feature"]["memory"] = serde_json::json!({ "enabled": false }); + disabled["memory"] = serde_json::json!({ + "workspace_id": "workspace-1", + "settings_revision": 9, + "language": "Français", + }); + let disabled = read_persisted_worker_manifest_snapshot(disabled).unwrap(); + assert!(!disabled.feature.memory.profile.enabled); + assert!(disabled.feature.memory.workspace_settings.is_none()); } #[test] @@ -1589,6 +1618,25 @@ model_id = "claude-sonnet-4-20250514" }); assert!(read_persisted_worker_manifest_snapshot(custom_policy).is_err()); + let current = WorkerManifest::from_toml(MINIMAL_REQUIRED).unwrap(); + let mut current = write_persisted_worker_manifest_snapshot(¤t).unwrap(); + current["manifest"]["memory"] = serde_json::json!({ + "workspace_id": "workspace-1", + "settings_revision": 1, + "language": "English", + }); + assert!(read_persisted_worker_manifest_snapshot(current).is_err()); + + let mut missing_settings = WorkerManifest::from_toml(MINIMAL_REQUIRED).unwrap(); + missing_settings.feature.memory.profile.enabled = true; + let missing_settings = write_persisted_worker_manifest_snapshot(&missing_settings).unwrap(); + assert!(read_persisted_worker_manifest_snapshot(missing_settings).is_err()); + + let mut malformed_legacy = manifest.clone(); + malformed_legacy["feature"]["memory"] = serde_json::json!({ "enabled": "yes" }); + malformed_legacy["memory"] = serde_json::json!({ "unknown": true }); + assert!(read_persisted_worker_manifest_snapshot(malformed_legacy).is_err()); + assert!( read_persisted_worker_manifest_snapshot(serde_json::json!({ "schema_version": 3, diff --git a/crates/worker/src/feature/builtin/memory.rs b/crates/worker/src/feature/builtin/memory.rs index e04054b0..33c5eff3 100644 --- a/crates/worker/src/feature/builtin/memory.rs +++ b/crates/worker/src/feature/builtin/memory.rs @@ -620,6 +620,48 @@ mod tests { assert_eq!(plan.resident_summary.as_deref(), Some("# Durable Memory")); } + #[tokio::test] + async fn memory_prompt_contribution_rereads_resident_summary_for_each_install() { + let prompts = crate::prompt::catalog::PromptCatalog::builtins_only().unwrap(); + let mut config = manifest::ResolvedMemoryFeatureConfig::default(); + config.profile.enabled = true; + config + .bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot { + workspace_id: "workspace".to_string(), + settings_revision: 1, + language: "English".to_string(), + }) + .unwrap(); + + let first = MemoryFeatureInstallPlan::prepare_resolved( + config.clone(), + resident_client("first resident summary"), + prompts.clone(), + None, + ) + .await + .unwrap() + .unwrap(); + let restored = MemoryFeatureInstallPlan::prepare_resolved( + config, + resident_client("updated resident summary"), + prompts, + None, + ) + .await + .unwrap() + .unwrap(); + + assert_eq!( + first.resident_summary.as_deref(), + Some("first resident summary") + ); + assert_eq!( + restored.resident_summary.as_deref(), + Some("updated resident summary") + ); + } + #[test] fn memory_feature_owns_normal_and_staging_tool_surfaces() { let normal = MemoryToolsFeature::new(test_client(), false); diff --git a/crates/worker/src/feature/builtin/memory_lifecycle.rs b/crates/worker/src/feature/builtin/memory_lifecycle.rs index 7349f842..13d501ea 100644 --- a/crates/worker/src/feature/builtin/memory_lifecycle.rs +++ b/crates/worker/src/feature/builtin/memory_lifecycle.rs @@ -77,7 +77,10 @@ impl MemoryLifecycleFeature { workspace_context: WorkerWorkspaceContext, event_tx: Option>, ) -> std::io::Result> { - if !lifecycle_enabled || !config.profile.enabled || !config.profile.extraction.enabled { + if !lifecycle_enabled + || !config.profile.enabled + || (!config.profile.extraction.enabled && !config.profile.consolidation.request_enabled) + { return Ok(None); } config @@ -521,9 +524,12 @@ impl FeatureBackgroundTask for MemoryLifecycleTask { context: BackgroundTaskContext, cancellation: BackgroundTaskCancellation, ) -> Result<(), HookError> { - let extraction = self - .run_extraction(context.clone(), cancellation.clone()) - .await; + let extraction = if self.config.profile.extraction.enabled { + self.run_extraction(context.clone(), cancellation.clone()) + .await + } else { + Ok(()) + }; if !cancellation.is_cancelled() { context.generation_fence.ensure_current()?; if self.config.profile.consolidation.request_enabled { @@ -1355,6 +1361,38 @@ permission = "write" ); } + #[tokio::test] + async fn lifecycle_task_requests_consolidation_when_extraction_is_disabled() { + let client = ScriptClient::new(Vec::new()); + let extension_writes = Arc::new(Mutex::new(Vec::new())); + let (event_tx, _) = broadcast::channel(16); + let workspace_client = Arc::new(RecordingWorkspaceClient::default()); + let mut interrupted = capture(2, 250); + interrupted.run_exit = CommittedRunExit::Interrupted; + let mut task = test_task( + interrupted, + Box::new(client), + extension_writes.clone(), + event_tx, + workspace_client.clone(), + ); + task.config.profile.extraction.enabled = false; + task.config.profile.consolidation.request_enabled = true; + run_background_task(task).await; + + assert!(extension_writes.lock().unwrap().is_empty()); + let requests = workspace_client.requests.lock().unwrap(); + assert!( + requests.iter().any(|request| { + request + .body + .as_deref() + .is_some_and(|body| body == "{\"force\":false}") + }), + "recorded requests: {requests:?}" + ); + } + fn internal_result( lifecycle: WorkerRunResult, ) -> Result {