From e997aa2fc6a28c437a88658adeb33dfdd4e0978b Mon Sep 17 00:00:00 2001 From: Hare Date: Sat, 25 Jul 2026 10:56:27 +0900 Subject: [PATCH] memory: gate staging tools by profile --- crates/manifest/src/config.rs | 52 +++++++++++++++++--- crates/manifest/src/lib.rs | 35 ++++++++++++- crates/manifest/src/profile.rs | 34 +++++++++++++ crates/worker/src/controller.rs | 2 +- crates/worker/src/feature/builtin/memory.rs | 41 +++++++++++++++ crates/workspace-server/src/hosts.rs | 28 ++++++++++- resources/profiles/memory-consolidation.dcdl | 22 +++++---- 7 files changed, 194 insertions(+), 20 deletions(-) diff --git a/crates/manifest/src/config.rs b/crates/manifest/src/config.rs index 73534508..54cf647d 100644 --- a/crates/manifest/src/config.rs +++ b/crates/manifest/src/config.rs @@ -18,9 +18,9 @@ use crate::model::{AuthRef, ModelManifest, ReasoningControl}; use crate::plugin::PluginConfig; use crate::{ CompactionConfig, EngineManifest, FeatureConfig, FeatureFlagConfig, FileUploadLimits, - McpConfig, McpEnvValue, McpStdioCwdPolicy, MemoryConfig, ScopeConfig, SessionConfig, - SkillsConfig, TicketFeatureConfig, ToolOutputLimits, ToolPermissionConfig, ToolPermissionRule, - WebConfig, WorkerManifest, WorkerMeta, + McpConfig, McpEnvValue, McpStdioCwdPolicy, MemoryConfig, MemoryFeatureConfig, ScopeConfig, + SessionConfig, SkillsConfig, TicketFeatureConfig, ToolOutputLimits, ToolPermissionConfig, + ToolPermissionRule, WebConfig, WorkerManifest, WorkerMeta, }; /// Partial-form Worker manifest. Every field is optional; one or more @@ -79,7 +79,7 @@ pub struct FeatureConfigPartial { #[serde(default)] pub task: Option, #[serde(default)] - pub memory: Option, + pub memory: Option, #[serde(default)] pub web: Option, #[serde(default)] @@ -94,7 +94,7 @@ impl FeatureConfigPartial { fn merge(self, other: Self) -> Self { Self { task: merge_option(self.task, other.task, FeatureFlagConfigPartial::merge), - memory: merge_option(self.memory, other.memory, FeatureFlagConfigPartial::merge), + memory: merge_option(self.memory, other.memory, MemoryFeatureConfigPartial::merge), web: merge_option(self.web, other.web, FeatureFlagConfigPartial::merge), workers: merge_option(self.workers, other.workers, FeatureFlagConfigPartial::merge), ticket: merge_option(self.ticket, other.ticket, TicketFeatureConfigPartial::merge), @@ -117,6 +117,23 @@ impl FeatureFlagConfigPartial { } } +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct MemoryFeatureConfigPartial { + #[serde(default)] + pub enabled: Option, + #[serde(default)] + pub staging: Option, +} + +impl MemoryFeatureConfigPartial { + fn merge(self, other: Self) -> Self { + Self { + enabled: other.enabled.or(self.enabled), + staging: other.staging.or(self.staging), + } + } +} + #[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Eq)] #[serde(default, deny_unknown_fields)] pub struct TicketFeatureConfigPartial { @@ -145,7 +162,7 @@ impl From for FeatureConfig { task: value.task.map(FeatureFlagConfig::from).unwrap_or_default(), memory: value .memory - .map(FeatureFlagConfig::from) + .map(MemoryFeatureConfig::from) .unwrap_or_default(), web: value.web.map(FeatureFlagConfig::from).unwrap_or_default(), workers: value @@ -180,6 +197,24 @@ impl From for FeatureFlagConfigPartial { } } +impl From for MemoryFeatureConfig { + fn from(value: MemoryFeatureConfigPartial) -> Self { + Self { + enabled: value.enabled.unwrap_or_default(), + staging: value.staging.unwrap_or_default(), + } + } +} + +impl From for MemoryFeatureConfigPartial { + fn from(value: MemoryFeatureConfig) -> Self { + Self { + enabled: Some(value.enabled), + staging: Some(value.staging), + } + } +} + impl From for TicketFeatureConfig { fn from(value: TicketFeatureConfigPartial) -> Self { Self { @@ -1806,6 +1841,7 @@ orchestration_control = false assert!(!manifest.feature.ticket.intake); assert!(!manifest.feature.ticket.orchestration_control); assert!(!manifest.feature.memory.enabled); + assert!(!manifest.feature.memory.staging); } #[test] @@ -1830,6 +1866,9 @@ orchestration_control = false thread = true orchestration_control = true +[feature.memory] +staging = true + [feature.web] enabled = true "#, @@ -1861,6 +1900,7 @@ enabled = true .try_into() .unwrap(); assert!(manifest.feature.memory.enabled); + assert!(manifest.feature.memory.staging); assert!(manifest.feature.ticket.enabled); assert!(!manifest.feature.ticket.authoring); assert!(manifest.feature.ticket.thread); diff --git a/crates/manifest/src/lib.rs b/crates/manifest/src/lib.rs index d4855059..538db092 100644 --- a/crates/manifest/src/lib.rs +++ b/crates/manifest/src/lib.rs @@ -107,7 +107,7 @@ pub struct FeatureConfig { #[serde(default)] pub task: FeatureFlagConfig, #[serde(default)] - pub memory: FeatureFlagConfig, + pub memory: MemoryFeatureConfig, #[serde(default)] pub web: FeatureFlagConfig, #[serde(default)] @@ -122,7 +122,7 @@ impl Default for FeatureConfig { fn default() -> Self { Self { task: FeatureFlagConfig::disabled(), - memory: FeatureFlagConfig::disabled(), + memory: MemoryFeatureConfig::disabled(), web: FeatureFlagConfig::disabled(), workers: FeatureFlagConfig::disabled(), ticket: TicketFeatureConfig::default(), @@ -153,6 +153,37 @@ impl Default for FeatureFlagConfig { } } +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +pub struct MemoryFeatureConfig { + #[serde(default)] + pub enabled: bool, + /// Exposes Memory staging queue tools in addition to normal Memory CRUD/query tools. + #[serde(default)] + pub staging: bool, +} + +impl MemoryFeatureConfig { + pub const fn disabled() -> Self { + Self { + enabled: false, + staging: false, + } + } + + pub const fn enabled() -> Self { + Self { + enabled: true, + staging: false, + } + } +} + +impl Default for MemoryFeatureConfig { + fn default() -> Self { + Self::disabled() + } +} + #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)] pub struct TicketFeatureConfig { #[serde(default)] diff --git a/crates/manifest/src/profile.rs b/crates/manifest/src/profile.rs index 1f259ac4..62eb337f 100644 --- a/crates/manifest/src/profile.rs +++ b/crates/manifest/src/profile.rs @@ -954,6 +954,17 @@ fn builtin_profile_artifact(label: &str) -> Option { ); Some(value) } + "builtin:memory-consolidation" | "memory-consolidation" => { + value["slug"] = serde_json::Value::String("memory-consolidation".to_string()); + value["description"] = + serde_json::Value::String("Memory staging consolidation profile.".to_string()); + value["feature"]["task"] = serde_json::json!({ "enabled": false }); + value["feature"]["memory"] = serde_json::json!({ "enabled": true, "staging": true }); + value["feature"]["web"] = serde_json::json!({ "enabled": false }); + value["feature"]["workers"] = serde_json::json!({ "enabled": false }); + value["feature"]["ticket"] = serde_json::json!({ "enabled": false, "thread": false }); + Some(value) + } _ => None, } } @@ -1425,6 +1436,29 @@ mod tests { } } + #[test] + fn builtin_memory_consolidation_profile_enables_staging_by_feature() { + let tmp = TempDir::new().unwrap(); + let resolved = ProfileResolver::new() + .with_workspace_base(tmp.path()) + .resolve( + &ProfileSelector::source_named( + ProfileRegistrySource::Builtin, + "memory-consolidation", + ), + ProfileResolveOptions::with_worker_name("arbitrary-worker-name"), + ) + .unwrap(); + + assert_eq!( + resolved.profile.as_ref().unwrap().name.as_deref(), + Some("memory-consolidation") + ); + assert_eq!(resolved.manifest.worker.name, "arbitrary-worker-name"); + assert!(resolved.manifest.feature.memory.enabled); + assert!(resolved.manifest.feature.memory.staging); + } + #[test] fn builtin_role_profiles_preserve_role_tool_policy() { let tmp = TempDir::new().unwrap(); diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 5f79c8b9..780f48ec 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -646,7 +646,7 @@ where base_url, } = workspace_client { - let definitions = if spawner_name == "memory-consolidation" { + let definitions = if feature_config.memory.staging { crate::feature::builtin::memory::workspace_http_memory_consolidation_tools( workspace_id, base_url, diff --git a/crates/worker/src/feature/builtin/memory.rs b/crates/worker/src/feature/builtin/memory.rs index d32051ae..b1247f47 100644 --- a/crates/worker/src/feature/builtin/memory.rs +++ b/crates/worker/src/feature/builtin/memory.rs @@ -418,3 +418,44 @@ fn query_schema() -> serde_json::Value { } }) } + +#[cfg(test)] +mod tests { + use super::*; + use llm_engine::tool::ToolDefinition; + + fn tool_names(definitions: Vec) -> Vec { + let mut names = definitions + .into_iter() + .map(|tool| tool().0.name) + .collect::>(); + names.sort(); + names + } + + #[test] + fn normal_workspace_memory_tools_do_not_include_staging_tools() { + let names = tool_names(workspace_http_memory_tools( + "workspace".to_string(), + "http://backend".to_string(), + )); + + assert!(names.contains(&"MemoryQuery".to_string())); + assert!(!names.contains(&"MemoryStagingList".to_string())); + assert!(!names.contains(&"MemoryStagingRead".to_string())); + assert!(!names.contains(&"MemoryStagingClose".to_string())); + } + + #[test] + fn consolidation_workspace_memory_tools_include_staging_tools() { + let names = tool_names(workspace_http_memory_consolidation_tools( + "workspace".to_string(), + "http://backend".to_string(), + )); + + assert!(names.contains(&"MemoryQuery".to_string())); + assert!(names.contains(&"MemoryStagingList".to_string())); + assert!(names.contains(&"MemoryStagingRead".to_string())); + assert!(names.contains(&"MemoryStagingClose".to_string())); + } +} diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index 4fe4fe3d..cf4d3d1d 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -2892,7 +2892,14 @@ fn default_profile_source_archive( "builtin:default".to_string(), "profiles/default.dcdl".to_string(), ); - for slug in ["companion", "intake", "orchestrator", "coder", "reviewer"] { + for slug in [ + "companion", + "intake", + "orchestrator", + "coder", + "reviewer", + "memory-consolidation", + ] { entrypoints.insert(format!("builtin:{slug}"), format!("profiles/{slug}.dcdl")); } entrypoints.insert(selected, selected_path); @@ -2922,9 +2929,20 @@ fn default_profile_source_archive( "profiles/reviewer.dcdl".to_string(), include_str!("../../../resources/profiles/reviewer.dcdl").to_string(), ); + sources.insert( + "profiles/memory-consolidation.dcdl".to_string(), + include_str!("../../../resources/profiles/memory-consolidation.dcdl").to_string(), + ); let mut imports = BTreeMap::new(); - for slug in ["companion", "intake", "orchestrator", "coder", "reviewer"] { + for slug in [ + "companion", + "intake", + "orchestrator", + "coder", + "reviewer", + "memory-consolidation", + ] { imports.insert( format!("profiles/{slug}.dcdl\0./default.dcdl"), "profiles/default.dcdl".to_string(), @@ -2950,6 +2968,7 @@ fn embedded_profile_path(profile: &ProfileSelector) -> Result { "orchestrator" => Ok("profiles/orchestrator.dcdl".to_string()), "coder" => Ok("profiles/coder.dcdl".to_string()), "reviewer" => Ok("profiles/reviewer.dcdl".to_string()), + "memory-consolidation" => Ok("profiles/memory-consolidation.dcdl".to_string()), other => Err(format!("unknown builtin profile selector: builtin:{other}")), }, ProfileSelector::Named(name) => Err(format!("unknown named profile selector: {name}")), @@ -3433,6 +3452,7 @@ mod tests { ProfileSelector::Builtin("builtin:orchestrator".to_string()), ProfileSelector::Builtin("builtin:coder".to_string()), ProfileSelector::Builtin("builtin:reviewer".to_string()), + ProfileSelector::Builtin("builtin:memory-consolidation".to_string()), ] { let bundle = default_embedded_config_bundle( &selector, @@ -3469,6 +3489,10 @@ mod tests { .unwrap(); assert_eq!(manifest.worker.name, "embedded-test-worker"); assert_eq!(manifest.model.ref_.as_deref(), Some("codex-oauth/gpt-5.5")); + if selector_key == "builtin:memory-consolidation" { + assert!(manifest.feature.memory.enabled); + assert!(manifest.feature.memory.staging); + } } } diff --git a/resources/profiles/memory-consolidation.dcdl b/resources/profiles/memory-consolidation.dcdl index a0a33f19..092cc24f 100644 --- a/resources/profiles/memory-consolidation.dcdl +++ b/resources/profiles/memory-consolidation.dcdl @@ -1,9 +1,13 @@ -worker = { name = "memory-consolidation"; } -feature = { - task = { enabled = false; }; - memory = { enabled = true; }; - web = { enabled = false; }; - workers = { enabled = false; }; - ticket = { enabled = false; thread = false; }; -}; -system_prompt = { file = "internal/memory_consolidation_system.md"; }; +import "./default.dcdl" // { + slug = "memory-consolidation"; + description = "Memory staging consolidation profile."; + scope = "workspace_read"; + + feature = { + task = { enabled = false; }; + memory = { enabled = true; staging = true; }; + web = { enabled = false; }; + workers = { enabled = false; }; + ticket = { enabled = false; thread = false; }; + }; +}