memory: gate staging tools by profile

This commit is contained in:
Keisuke Hirata 2026-07-25 10:56:27 +09:00
parent a44673fa6b
commit e997aa2fc6
No known key found for this signature in database
7 changed files with 194 additions and 20 deletions

View File

@ -18,9 +18,9 @@ use crate::model::{AuthRef, ModelManifest, ReasoningControl};
use crate::plugin::PluginConfig; use crate::plugin::PluginConfig;
use crate::{ use crate::{
CompactionConfig, EngineManifest, FeatureConfig, FeatureFlagConfig, FileUploadLimits, CompactionConfig, EngineManifest, FeatureConfig, FeatureFlagConfig, FileUploadLimits,
McpConfig, McpEnvValue, McpStdioCwdPolicy, MemoryConfig, ScopeConfig, SessionConfig, McpConfig, McpEnvValue, McpStdioCwdPolicy, MemoryConfig, MemoryFeatureConfig, ScopeConfig,
SkillsConfig, TicketFeatureConfig, ToolOutputLimits, ToolPermissionConfig, ToolPermissionRule, SessionConfig, SkillsConfig, TicketFeatureConfig, ToolOutputLimits, ToolPermissionConfig,
WebConfig, WorkerManifest, WorkerMeta, ToolPermissionRule, WebConfig, WorkerManifest, WorkerMeta,
}; };
/// Partial-form Worker manifest. Every field is optional; one or more /// Partial-form Worker manifest. Every field is optional; one or more
@ -79,7 +79,7 @@ pub struct FeatureConfigPartial {
#[serde(default)] #[serde(default)]
pub task: Option<FeatureFlagConfigPartial>, pub task: Option<FeatureFlagConfigPartial>,
#[serde(default)] #[serde(default)]
pub memory: Option<FeatureFlagConfigPartial>, pub memory: Option<MemoryFeatureConfigPartial>,
#[serde(default)] #[serde(default)]
pub web: Option<FeatureFlagConfigPartial>, pub web: Option<FeatureFlagConfigPartial>,
#[serde(default)] #[serde(default)]
@ -94,7 +94,7 @@ impl FeatureConfigPartial {
fn merge(self, other: Self) -> Self { fn merge(self, other: Self) -> Self {
Self { Self {
task: merge_option(self.task, other.task, FeatureFlagConfigPartial::merge), 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), web: merge_option(self.web, other.web, FeatureFlagConfigPartial::merge),
workers: merge_option(self.workers, other.workers, FeatureFlagConfigPartial::merge), workers: merge_option(self.workers, other.workers, FeatureFlagConfigPartial::merge),
ticket: merge_option(self.ticket, other.ticket, TicketFeatureConfigPartial::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<bool>,
#[serde(default)]
pub staging: Option<bool>,
}
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)] #[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Eq)]
#[serde(default, deny_unknown_fields)] #[serde(default, deny_unknown_fields)]
pub struct TicketFeatureConfigPartial { pub struct TicketFeatureConfigPartial {
@ -145,7 +162,7 @@ impl From<FeatureConfigPartial> for FeatureConfig {
task: value.task.map(FeatureFlagConfig::from).unwrap_or_default(), task: value.task.map(FeatureFlagConfig::from).unwrap_or_default(),
memory: value memory: value
.memory .memory
.map(FeatureFlagConfig::from) .map(MemoryFeatureConfig::from)
.unwrap_or_default(), .unwrap_or_default(),
web: value.web.map(FeatureFlagConfig::from).unwrap_or_default(), web: value.web.map(FeatureFlagConfig::from).unwrap_or_default(),
workers: value workers: value
@ -180,6 +197,24 @@ impl From<FeatureFlagConfig> for FeatureFlagConfigPartial {
} }
} }
impl From<MemoryFeatureConfigPartial> for MemoryFeatureConfig {
fn from(value: MemoryFeatureConfigPartial) -> Self {
Self {
enabled: value.enabled.unwrap_or_default(),
staging: value.staging.unwrap_or_default(),
}
}
}
impl From<MemoryFeatureConfig> for MemoryFeatureConfigPartial {
fn from(value: MemoryFeatureConfig) -> Self {
Self {
enabled: Some(value.enabled),
staging: Some(value.staging),
}
}
}
impl From<TicketFeatureConfigPartial> for TicketFeatureConfig { impl From<TicketFeatureConfigPartial> for TicketFeatureConfig {
fn from(value: TicketFeatureConfigPartial) -> Self { fn from(value: TicketFeatureConfigPartial) -> Self {
Self { Self {
@ -1806,6 +1841,7 @@ orchestration_control = false
assert!(!manifest.feature.ticket.intake); assert!(!manifest.feature.ticket.intake);
assert!(!manifest.feature.ticket.orchestration_control); assert!(!manifest.feature.ticket.orchestration_control);
assert!(!manifest.feature.memory.enabled); assert!(!manifest.feature.memory.enabled);
assert!(!manifest.feature.memory.staging);
} }
#[test] #[test]
@ -1830,6 +1866,9 @@ orchestration_control = false
thread = true thread = true
orchestration_control = true orchestration_control = true
[feature.memory]
staging = true
[feature.web] [feature.web]
enabled = true enabled = true
"#, "#,
@ -1861,6 +1900,7 @@ enabled = true
.try_into() .try_into()
.unwrap(); .unwrap();
assert!(manifest.feature.memory.enabled); assert!(manifest.feature.memory.enabled);
assert!(manifest.feature.memory.staging);
assert!(manifest.feature.ticket.enabled); assert!(manifest.feature.ticket.enabled);
assert!(!manifest.feature.ticket.authoring); assert!(!manifest.feature.ticket.authoring);
assert!(manifest.feature.ticket.thread); assert!(manifest.feature.ticket.thread);

View File

@ -107,7 +107,7 @@ pub struct FeatureConfig {
#[serde(default)] #[serde(default)]
pub task: FeatureFlagConfig, pub task: FeatureFlagConfig,
#[serde(default)] #[serde(default)]
pub memory: FeatureFlagConfig, pub memory: MemoryFeatureConfig,
#[serde(default)] #[serde(default)]
pub web: FeatureFlagConfig, pub web: FeatureFlagConfig,
#[serde(default)] #[serde(default)]
@ -122,7 +122,7 @@ impl Default for FeatureConfig {
fn default() -> Self { fn default() -> Self {
Self { Self {
task: FeatureFlagConfig::disabled(), task: FeatureFlagConfig::disabled(),
memory: FeatureFlagConfig::disabled(), memory: MemoryFeatureConfig::disabled(),
web: FeatureFlagConfig::disabled(), web: FeatureFlagConfig::disabled(),
workers: FeatureFlagConfig::disabled(), workers: FeatureFlagConfig::disabled(),
ticket: TicketFeatureConfig::default(), 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)] #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
pub struct TicketFeatureConfig { pub struct TicketFeatureConfig {
#[serde(default)] #[serde(default)]

View File

@ -954,6 +954,17 @@ fn builtin_profile_artifact(label: &str) -> Option<serde_json::Value> {
); );
Some(value) 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, _ => 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] #[test]
fn builtin_role_profiles_preserve_role_tool_policy() { fn builtin_role_profiles_preserve_role_tool_policy() {
let tmp = TempDir::new().unwrap(); let tmp = TempDir::new().unwrap();

View File

@ -646,7 +646,7 @@ where
base_url, base_url,
} = workspace_client } = 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( crate::feature::builtin::memory::workspace_http_memory_consolidation_tools(
workspace_id, workspace_id,
base_url, base_url,

View File

@ -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<ToolDefinition>) -> Vec<String> {
let mut names = definitions
.into_iter()
.map(|tool| tool().0.name)
.collect::<Vec<_>>();
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()));
}
}

View File

@ -2892,7 +2892,14 @@ fn default_profile_source_archive(
"builtin:default".to_string(), "builtin:default".to_string(),
"profiles/default.dcdl".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(format!("builtin:{slug}"), format!("profiles/{slug}.dcdl"));
} }
entrypoints.insert(selected, selected_path); entrypoints.insert(selected, selected_path);
@ -2922,9 +2929,20 @@ fn default_profile_source_archive(
"profiles/reviewer.dcdl".to_string(), "profiles/reviewer.dcdl".to_string(),
include_str!("../../../resources/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(); 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( imports.insert(
format!("profiles/{slug}.dcdl\0./default.dcdl"), format!("profiles/{slug}.dcdl\0./default.dcdl"),
"profiles/default.dcdl".to_string(), "profiles/default.dcdl".to_string(),
@ -2950,6 +2968,7 @@ fn embedded_profile_path(profile: &ProfileSelector) -> Result<String, String> {
"orchestrator" => Ok("profiles/orchestrator.dcdl".to_string()), "orchestrator" => Ok("profiles/orchestrator.dcdl".to_string()),
"coder" => Ok("profiles/coder.dcdl".to_string()), "coder" => Ok("profiles/coder.dcdl".to_string()),
"reviewer" => Ok("profiles/reviewer.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}")), other => Err(format!("unknown builtin profile selector: builtin:{other}")),
}, },
ProfileSelector::Named(name) => Err(format!("unknown named profile selector: {name}")), 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:orchestrator".to_string()),
ProfileSelector::Builtin("builtin:coder".to_string()), ProfileSelector::Builtin("builtin:coder".to_string()),
ProfileSelector::Builtin("builtin:reviewer".to_string()), ProfileSelector::Builtin("builtin:reviewer".to_string()),
ProfileSelector::Builtin("builtin:memory-consolidation".to_string()),
] { ] {
let bundle = default_embedded_config_bundle( let bundle = default_embedded_config_bundle(
&selector, &selector,
@ -3469,6 +3489,10 @@ mod tests {
.unwrap(); .unwrap();
assert_eq!(manifest.worker.name, "embedded-test-worker"); assert_eq!(manifest.worker.name, "embedded-test-worker");
assert_eq!(manifest.model.ref_.as_deref(), Some("codex-oauth/gpt-5.5")); 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);
}
} }
} }

View File

@ -1,9 +1,13 @@
worker = { name = "memory-consolidation"; } import "./default.dcdl" // {
slug = "memory-consolidation";
description = "Memory staging consolidation profile.";
scope = "workspace_read";
feature = { feature = {
task = { enabled = false; }; task = { enabled = false; };
memory = { enabled = true; }; memory = { enabled = true; staging = true; };
web = { enabled = false; }; web = { enabled = false; };
workers = { enabled = false; }; workers = { enabled = false; };
ticket = { enabled = false; thread = false; }; ticket = { enabled = false; thread = false; };
}; };
system_prompt = { file = "internal/memory_consolidation_system.md"; }; }