From 16c4b41e1b56e76eabba52406463b12169f996b2 Mon Sep 17 00:00:00 2001 From: Lauritz-Timm Date: Fri, 7 Aug 2026 22:10:29 +0200 Subject: [PATCH 01/20] feat(mcp): centralize typed tool catalog --- Cargo.lock | 150 ++- Cargo.toml | 1 + crates/icm-mcp/Cargo.toml | 1 + crates/icm-mcp/src/catalog.rs | 860 +++++++++++++++++ crates/icm-mcp/src/inputs.rs | 483 ++++++++++ crates/icm-mcp/src/lib.rs | 2 + crates/icm-mcp/src/tools.rs | 1680 ++++++++++++++++++++------------- 7 files changed, 2468 insertions(+), 709 deletions(-) create mode 100644 crates/icm-mcp/src/catalog.rs create mode 100644 crates/icm-mcp/src/inputs.rs diff --git a/Cargo.lock b/Cargo.lock index 4a0037bf..8be5e412 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -137,7 +137,7 @@ checksum = "0ae92a5119aa49cdbcf6b9f893fe4e1d98b04ccbf82ee0584ad948a44a734dea" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -163,7 +163,7 @@ checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -458,7 +458,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -710,7 +710,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn", + "syn 2.0.117", ] [[package]] @@ -723,7 +723,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn", + "syn 2.0.117", ] [[package]] @@ -734,7 +734,7 @@ checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" dependencies = [ "darling_core 0.20.11", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -745,7 +745,7 @@ checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" dependencies = [ "darling_core 0.23.0", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -765,7 +765,7 @@ checksum = "1e567bd82dcff979e4b03460c307b3cdc9e96fde3d73bed1496d2bc75d9dd62a" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -786,7 +786,7 @@ dependencies = [ "darling 0.20.11", "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -796,7 +796,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" dependencies = [ "derive_builder_core", - "syn", + "syn 2.0.117", ] [[package]] @@ -859,9 +859,15 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + [[package]] name = "either" version = "1.15.0" @@ -900,7 +906,7 @@ checksum = "44f23cf4b44bfce11a86ace86f8a73ffdec849c9fd00a386a53d278bd9e81fb3" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -998,7 +1004,7 @@ checksum = "a0aca10fb742cb43f9e7bb8467c91aa9bcb8e3ffbc6a6f7389bb93ffc920577d" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1109,7 +1115,7 @@ checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1460,7 +1466,7 @@ dependencies = [ [[package]] name = "icm-cli" -version = "0.10.54" +version = "0.10.61" dependencies = [ "anyhow", "axum", @@ -1521,6 +1527,7 @@ dependencies = [ "chrono", "icm-core", "icm-store", + "schemars", "serde", "serde_json", "tempfile", @@ -1746,7 +1753,7 @@ dependencies = [ "indoc", "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1757,7 +1764,7 @@ checksum = "c34819042dc3d3971c46c2190835914dfbe0c3c13f61449b2997f4e9722dfa60" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2083,7 +2090,7 @@ checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2195,7 +2202,7 @@ checksum = "ed3955f1a9c7c0c15e092f9c887db08b1fc683305fdf6eb6684f22555355e202" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2308,7 +2315,7 @@ checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2561,7 +2568,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ "proc-macro2", - "syn", + "syn 2.0.117", ] [[package]] @@ -2589,7 +2596,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52717f9a02b6965224f95ca2a81e2e0c5c43baacd28ca057577988930b6c3d5b" dependencies = [ "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2817,6 +2824,26 @@ dependencies = [ "thiserror", ] +[[package]] +name = "ref-cast" +version = "1.0.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "216e8f773d7923bcba9ceb86a86c93cabb3903a11872fc3f138c49630e50b96d" +dependencies = [ + "ref-cast-impl", +] + +[[package]] +name = "ref-cast-impl" +version = "1.0.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9283685feec7d69af75fb0e858d5e7378f33fe4fc699383b2916ab9273e03c" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "regex" version = "1.12.3" @@ -2953,7 +2980,7 @@ dependencies = [ "proc-macro2", "quote", "rust-embed-utils", - "syn", + "syn 2.0.117", "walkdir", ] @@ -3058,6 +3085,31 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "schemars" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "687274d293b6cdc6e73e0fee520bf2049650090d7164f87672d212a3c530cf4a" +dependencies = [ + "dyn-clone", + "ref-cast", + "schemars_derive", + "serde", + "serde_json", +] + +[[package]] +name = "schemars_derive" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98c67716b46af2f0b8cf752abc930f6f9aecfbf671ecfb531db8a31dbe4e2ba" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn 3.0.3", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -3120,7 +3172,18 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", +] + +[[package]] +name = "serde_derive_internals" +version = "0.30.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f852137cce035d6a4df67ccce505ff6b3e9fd3a10e3e52b24dc71e650bb1a9bd" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", ] [[package]] @@ -3373,7 +3436,7 @@ dependencies = [ "proc-macro2", "quote", "rustversion", - "syn", + "syn 2.0.117", ] [[package]] @@ -3393,6 +3456,17 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "sync_wrapper" version = "1.0.2" @@ -3410,7 +3484,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -3475,7 +3549,7 @@ checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -3582,7 +3656,7 @@ checksum = "5c55a2eff8b69ce66c84f85e1da1c233edc36ceb85a2058d11b0d6a3c7e7569c" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -3752,7 +3826,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -4085,7 +4159,7 @@ dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn", + "syn 2.0.117", "wasm-bindgen-shared", ] @@ -4254,7 +4328,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -4265,7 +4339,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -4507,7 +4581,7 @@ dependencies = [ "heck", "indexmap", "prettyplease", - "syn", + "syn 2.0.117", "wasm-metadata", "wit-bindgen-core", "wit-component", @@ -4523,7 +4597,7 @@ dependencies = [ "prettyplease", "proc-macro2", "quote", - "syn", + "syn 2.0.117", "wit-bindgen-core", "wit-bindgen-rust", ] @@ -4606,7 +4680,7 @@ checksum = "b659052874eb698efe5b9e8cf382204678a0086ebf46982b79d6ca3182927e5d" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", "synstructure", ] @@ -4627,7 +4701,7 @@ checksum = "f65c489a7071a749c849713807783f70672b28094011623e200cb86dcb835953" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -4647,7 +4721,7 @@ checksum = "d71e5d6e06ab090c67b5e44993ec16b72dcbaabc526db883a360057678b48502" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", "synstructure", ] @@ -4687,7 +4761,7 @@ checksum = "eadce39539ca5cb3985590102671f2567e659fca9666581ad3411d59207951f3" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 58876618..441be270 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,6 +31,7 @@ serde = { version = "1", features = ["derive"] } serde_json = { version = "1", features = ["preserve_order"] } serde_json_lenient = { version = "0.2", features = ["preserve_order"] } toml = "0.8" +schemars = "1" # Error handling thiserror = "2" diff --git a/crates/icm-mcp/Cargo.toml b/crates/icm-mcp/Cargo.toml index 392e17bf..d1f6c47f 100644 --- a/crates/icm-mcp/Cargo.toml +++ b/crates/icm-mcp/Cargo.toml @@ -23,6 +23,7 @@ icm-store = { path = "../icm-store", default-features = false } chrono = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } +schemars = { workspace = true } anyhow = { workspace = true } tracing = { workspace = true } diff --git a/crates/icm-mcp/src/catalog.rs b/crates/icm-mcp/src/catalog.rs new file mode 100644 index 00000000..9bbd995f --- /dev/null +++ b/crates/icm-mcp/src/catalog.rs @@ -0,0 +1,860 @@ +//! The immutable MCP tool catalog. +//! +//! A registration contains every fact needed to list and dispatch a tool. +//! The ordered registration vector is the only order source, and the lookup +//! map points back into that same vector. + +use std::collections::{HashMap, HashSet}; +use std::path::Path; + +use icm_core::Embedder; +use icm_store::Store; +use serde::de::DeserializeOwned; +use serde_json::{json, Value}; + +use crate::inputs::ModernToolInput; +use crate::protocol::{ProtocolRevision, ToolResult}; +use crate::tools::AutoConsolidate; + +pub type ToolHandler = for<'a> fn(&ToolContext<'a>, &Value) -> ToolResult; +type InputValidator = fn(&Value) -> Result<(), String>; +type InputNormalizer = fn(&Value) -> Value; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum EmbedderRequirement { + Unused, + Optional, + Required, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum InputValidation { + Legacy2024Unchecked, + Legacy2024, + Modern, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct ToolRequirements { + minimum_protocol_revision: ProtocolRevision, + store: bool, + embedder: EmbedderRequirement, + filesystem_read: bool, + required_client_capabilities: &'static [&'static str], + structured_output_from_revision: Option, + legacy_visible: bool, +} + +impl ToolRequirements { + pub(crate) const STORE: Self = Self { + minimum_protocol_revision: ProtocolRevision::V2024_11_05, + store: true, + embedder: EmbedderRequirement::Unused, + filesystem_read: false, + required_client_capabilities: &[], + structured_output_from_revision: None, + legacy_visible: true, + }; + + pub(crate) const fn with_optional_embedder(mut self) -> Self { + self.embedder = EmbedderRequirement::Optional; + self + } + + pub(crate) const fn with_required_embedder(mut self) -> Self { + self.embedder = EmbedderRequirement::Required; + self + } + + pub(crate) const fn with_filesystem_read(mut self) -> Self { + self.filesystem_read = true; + self + } + + fn is_available(self, has_embedder: bool) -> bool { + self.embedder != EmbedderRequirement::Required || has_embedder + } + + fn as_value(self) -> Value { + let embedder = match self.embedder { + EmbedderRequirement::Unused => "unused", + EmbedderRequirement::Optional => "optional", + EmbedderRequirement::Required => "required", + }; + json!({ + "minimumProtocolRevision": self.minimum_protocol_revision.as_str(), + "serverFacilities": { + "store": if self.store { "required" } else { "unused" }, + "embedder": embedder, + "filesystemRead": if self.filesystem_read { "required" } else { "unused" }, + }, + "requiredClientCapabilities": self.required_client_capabilities, + "structuredOutputFromRevision": self + .structured_output_from_revision + .map(ProtocolRevision::as_str), + "legacyVisible": self.legacy_visible, + }) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct ToolAnnotations { + pub read_only: bool, + pub destructive: bool, + pub idempotent: bool, + pub open_world: bool, +} + +impl ToolAnnotations { + pub const fn new( + read_only: bool, + destructive: bool, + idempotent: bool, + open_world: bool, + ) -> Self { + Self { + read_only, + destructive, + idempotent, + open_world, + } + } + + fn as_value(self) -> Value { + json!({ + "readOnlyHint": self.read_only, + "destructiveHint": self.destructive, + "idempotentHint": self.idempotent, + "openWorldHint": self.open_world, + }) + } +} + +pub struct ToolContext<'a> { + pub store: &'a Store, + pub embedder: Option<&'a dyn Embedder>, + pub compact: bool, + pub auto_consolidate: AutoConsolidate, + pub working_directory: &'a Path, + pub enforce_directory_boundary: bool, +} + +pub struct ToolSpec { + name: &'static str, + description: &'static str, + legacy_input_schema: Value, + modern_input_schema: Value, + legacy_input_normalizer: Option, + annotations: ToolAnnotations, + requirements: ToolRequirements, + validate_input: InputValidator, + handler: ToolHandler, +} + +impl ToolSpec { + pub fn typed( + name: &'static str, + description: &'static str, + legacy_input_schema: Value, + legacy_input_normalizer: Option, + annotations: ToolAnnotations, + requirements: ToolRequirements, + handler: ToolHandler, + ) -> Self + where + I: ModernToolInput, + { + let modern_input_schema = generated_input_schema::(&legacy_input_schema); + + Self { + name, + description, + legacy_input_schema, + modern_input_schema, + legacy_input_normalizer, + annotations, + requirements, + validate_input: deserialize_input::, + handler, + } + } + + fn legacy_definition(&self) -> Value { + json!({ + "name": self.name, + "description": self.description, + "inputSchema": self.legacy_input_schema, + }) + } + + fn modern_definition(&self) -> Value { + json!({ + "name": self.name, + "description": self.description, + "inputSchema": self.modern_input_schema, + "annotations": self.annotations.as_value(), + "_meta": { + "com.github.rtk-ai.icm/requirements": self.requirements.as_value(), + }, + }) + } + + fn validate_modern_input(&self, arguments: &Value) -> Result<(), String> { + (self.validate_input)(arguments)?; + validate_schema_constraints(arguments, &self.modern_input_schema, "$", 0) + } + + fn legacy_validation_input(&self, arguments: &Value) -> Value { + let mut filtered = arguments.clone(); + let Some(object) = filtered.as_object_mut() else { + return filtered; + }; + let declared_properties = self + .legacy_input_schema + .get("properties") + .and_then(Value::as_object); + object.retain(|name, _| { + declared_properties.is_some_and(|properties| properties.contains_key(name)) + }); + filtered + } +} + +fn deserialize_input(arguments: &Value) -> Result<(), String> +where + I: DeserializeOwned, +{ + serde_json::from_value::(arguments.clone()) + .map(|_| ()) + .map_err(|error| bounded_error(error.to_string())) +} + +fn generated_input_schema(legacy: &Value) -> Value +where + I: ModernToolInput, +{ + let mut generated = serde_json::to_value(schemars::schema_for!(I)) + .expect("generated tool input schema must serialize"); + let object = generated + .as_object_mut() + .expect("tool input schema root must be an object"); + object.insert("additionalProperties".into(), Value::Bool(false)); + object + .entry("required") + .or_insert_with(|| Value::Array(Vec::new())); + + // The frozen 2024 projection carries carefully worded descriptions, + // defaults, and numeric bounds. Copy those annotations onto the schema + // generated from the Rust DTO; field shape and requiredness still come + // solely from the type. + if let (Some(modern_properties), Some(legacy_properties)) = ( + object.get_mut("properties").and_then(Value::as_object_mut), + legacy.get("properties").and_then(Value::as_object), + ) { + for (name, legacy_property) in legacy_properties { + let Some(modern_property) = modern_properties + .get_mut(name) + .and_then(Value::as_object_mut) + else { + continue; + }; + for metadata_key in ["description", "default", "minimum", "maximum"] { + if let Some(value) = legacy_property.get(metadata_key) { + modern_property.insert(metadata_key.into(), value.clone()); + } + } + } + } + + I::refine_schema(&mut generated); + + generated +} + +fn validate_schema_constraints( + value: &Value, + schema: &Value, + path: &str, + depth: usize, +) -> Result<(), String> { + if depth > 32 { + return Err("input nesting exceeds maximum depth".into()); + } + if let Some(minimum) = schema.get("minimum").and_then(Value::as_i64) { + if value.as_i64().is_some_and(|actual| actual < minimum) { + return Err(format!("{path} must be at least {minimum}")); + } + } + if let Some(maximum) = schema.get("maximum").and_then(Value::as_i64) { + if value.as_i64().is_some_and(|actual| actual > maximum) { + return Err(format!("{path} must be at most {maximum}")); + } + } + if let Some(minimum) = schema.get("minLength").and_then(Value::as_u64) { + if value + .as_str() + .is_some_and(|actual| actual.chars().count() < minimum as usize) + { + return Err(format!("{path} is shorter than {minimum} characters")); + } + } + if let Some(maximum) = schema.get("maxLength").and_then(Value::as_u64) { + if value + .as_str() + .is_some_and(|actual| actual.chars().count() > maximum as usize) + { + return Err(format!("{path} is longer than {maximum} characters")); + } + } + if let Some(maximum) = schema.get("x-icm-maxUtf8Bytes").and_then(Value::as_u64) { + if value + .as_str() + .is_some_and(|actual| actual.len() > maximum as usize) + { + return Err(format!("{path} exceeds {maximum} UTF-8 bytes")); + } + } + if schema.get("x-icm-trimmedNonEmpty") == Some(&Value::Bool(true)) + && value + .as_str() + .is_some_and(|actual| actual.trim().is_empty()) + { + return Err(format!("{path} must not be empty or whitespace")); + } + + if let (Some(properties), Some(object)) = ( + schema.get("properties").and_then(Value::as_object), + value.as_object(), + ) { + for (name, child) in object { + if let Some(child_schema) = properties.get(name) { + validate_schema_constraints( + child, + child_schema, + &format!("{path}.{name}"), + depth + 1, + )?; + } + } + } + if let (Some(items), Some(array)) = (schema.get("items"), value.as_array()) { + for (index, child) in array.iter().enumerate() { + validate_schema_constraints(child, items, &format!("{path}[{index}]"), depth + 1)?; + } + } + Ok(()) +} + +fn bounded_error(mut message: String) -> String { + const MAX_ERROR_BYTES: usize = 512; + if message.len() > MAX_ERROR_BYTES { + let mut end = MAX_ERROR_BYTES; + while !message.is_char_boundary(end) { + end -= 1; + } + message.truncate(end); + message.push('…'); + } + message +} + +pub enum DispatchResult { + UnknownTool, + InvalidInput(String), + ToolResult(ToolResult), +} + +pub struct ToolCatalog { + registrations: Vec, + by_name: HashMap<&'static str, usize>, + legacy_list: Value, + modern_list: Value, +} + +impl ToolCatalog { + pub fn new(mut registrations: Vec, has_embedder: bool) -> Result { + let mut names = HashSet::with_capacity(registrations.len()); + for registration in ®istrations { + if !names.insert(registration.name) { + return Err(format!( + "duplicate MCP tool registration: {}", + registration.name + )); + } + } + + registrations.retain(|registration| registration.requirements.is_available(has_embedder)); + let by_name = registrations + .iter() + .enumerate() + .map(|(index, registration)| (registration.name, index)) + .collect(); + + let legacy_tools: Vec = registrations + .iter() + .filter(|registration| registration.requirements.legacy_visible) + .map(ToolSpec::legacy_definition) + .collect(); + let modern_tools: Vec = registrations + .iter() + .map(ToolSpec::modern_definition) + .collect(); + + Ok(Self { + registrations, + by_name, + legacy_list: json!({ "tools": legacy_tools }), + modern_list: json!({ "tools": modern_tools }), + }) + } + + pub fn legacy_list(&self) -> Value { + self.legacy_list.clone() + } + + pub fn modern_list(&self) -> Value { + self.modern_list.clone() + } + + pub fn dispatch( + &self, + context: &ToolContext<'_>, + name: &str, + arguments: &Value, + validation: InputValidation, + ) -> DispatchResult { + let Some(index) = self.by_name.get(name) else { + return DispatchResult::UnknownTool; + }; + let registration = &self.registrations[*index]; + let normalized_arguments = matches!( + validation, + InputValidation::Legacy2024Unchecked | InputValidation::Legacy2024 + ) + .then_some(registration.legacy_input_normalizer) + .flatten() + .map(|normalize| normalize(arguments)); + let dispatch_arguments = normalized_arguments.as_ref().unwrap_or(arguments); + if matches!( + validation, + InputValidation::Legacy2024 | InputValidation::Modern + ) { + let legacy_arguments = (validation == InputValidation::Legacy2024) + .then(|| registration.legacy_validation_input(dispatch_arguments)); + let validation_arguments = legacy_arguments.as_ref().unwrap_or(dispatch_arguments); + if let Err(message) = registration.validate_modern_input(validation_arguments) { + return DispatchResult::InvalidInput(message); + } + } + DispatchResult::ToolResult((registration.handler)(context, dispatch_arguments)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const EXPECTED_TOOLS: [(&str, ToolAnnotations); 31] = [ + ( + "icm_memory_store", + ToolAnnotations::new(false, true, false, false), + ), + ( + "icm_memory_recall", + ToolAnnotations::new(false, true, false, false), + ), + ( + "icm_memory_forget", + ToolAnnotations::new(false, true, true, false), + ), + ( + "icm_memory_forget_topic", + ToolAnnotations::new(false, true, true, false), + ), + ("icm_learn", ToolAnnotations::new(false, true, false, true)), + ( + "icm_memory_consolidate", + ToolAnnotations::new(false, true, false, false), + ), + ( + "icm_memory_list_topics", + ToolAnnotations::new(true, false, true, false), + ), + ( + "icm_memory_stats", + ToolAnnotations::new(true, false, true, false), + ), + ( + "icm_memory_update", + ToolAnnotations::new(false, true, false, false), + ), + ( + "icm_memory_health", + ToolAnnotations::new(true, false, true, false), + ), + ( + "icm_memoir_create", + ToolAnnotations::new(false, false, false, false), + ), + ( + "icm_memoir_list", + ToolAnnotations::new(true, false, true, false), + ), + ( + "icm_memoir_show", + ToolAnnotations::new(true, false, true, false), + ), + ( + "icm_memoir_add_concept", + ToolAnnotations::new(false, false, false, false), + ), + ( + "icm_memoir_refine", + ToolAnnotations::new(false, true, false, false), + ), + ( + "icm_memoir_search", + ToolAnnotations::new(true, false, true, false), + ), + ( + "icm_memoir_link", + ToolAnnotations::new(false, false, false, false), + ), + ( + "icm_memoir_inspect", + ToolAnnotations::new(true, false, true, false), + ), + ( + "icm_memoir_export", + ToolAnnotations::new(true, false, true, false), + ), + ( + "icm_memory_extract_patterns", + ToolAnnotations::new(false, false, false, false), + ), + ( + "icm_memoir_search_all", + ToolAnnotations::new(true, false, true, false), + ), + ( + "icm_feedback_record", + ToolAnnotations::new(false, false, false, false), + ), + ( + "icm_feedback_search", + ToolAnnotations::new(true, false, true, false), + ), + ( + "icm_feedback_stats", + ToolAnnotations::new(true, false, true, false), + ), + ( + "icm_transcript_start_session", + ToolAnnotations::new(false, false, false, false), + ), + ( + "icm_transcript_record", + ToolAnnotations::new(false, false, false, false), + ), + ( + "icm_transcript_search", + ToolAnnotations::new(true, false, true, false), + ), + ( + "icm_transcript_show", + ToolAnnotations::new(true, false, true, false), + ), + ( + "icm_transcript_stats", + ToolAnnotations::new(true, false, true, false), + ), + ( + "icm_wake_up", + ToolAnnotations::new(true, false, true, false), + ), + ( + "icm_memory_embed_all", + ToolAnnotations::new(false, false, true, false), + ), + ]; + + #[test] + fn annotation_projection_has_all_four_explicit_fields() { + let value = ToolAnnotations::new(true, false, true, false).as_value(); + assert_eq!(value.as_object().map(serde_json::Map::len), Some(4)); + assert_eq!(value["readOnlyHint"], true); + assert_eq!(value["destructiveHint"], false); + assert_eq!(value["idempotentHint"], true); + assert_eq!(value["openWorldHint"], false); + } + + #[test] + fn catalog_is_the_single_order_list_and_dispatch_source() { + let without_embedder = crate::tools::build_catalog(false); + let expected_without_embedder: Vec<&str> = + EXPECTED_TOOLS[..30].iter().map(|(name, _)| *name).collect(); + let without_projection = without_embedder.legacy_list(); + assert_eq!( + without_projection["tools"] + .as_array() + .unwrap() + .iter() + .map(|tool| tool["name"].as_str().unwrap()) + .collect::>(), + expected_without_embedder + ); + + let with_embedder = crate::tools::build_catalog(true); + let with_projection = with_embedder.legacy_list(); + assert_eq!( + with_projection["tools"] + .as_array() + .unwrap() + .iter() + .map(|tool| tool["name"].as_str().unwrap()) + .collect::>(), + EXPECTED_TOOLS + .iter() + .map(|(name, _)| *name) + .collect::>() + ); + + let store = Store::in_memory().unwrap(); + let working_directory = std::env::current_dir().unwrap(); + let context = ToolContext { + store: &store, + embedder: None, + compact: false, + auto_consolidate: AutoConsolidate::default(), + working_directory: &working_directory, + enforce_directory_boundary: true, + }; + for (name, _) in EXPECTED_TOOLS { + assert!(matches!( + with_embedder.dispatch( + &context, + name, + &json!({"__catalog_probe": true}), + InputValidation::Modern + ), + DispatchResult::InvalidInput(_) + )); + } + assert!(matches!( + without_embedder.dispatch( + &context, + "icm_memory_embed_all", + &json!({"__catalog_probe": true}), + InputValidation::Modern + ), + DispatchResult::UnknownTool + )); + } + + #[test] + fn modern_projection_has_exact_annotations_and_phase_two_schema_shape() { + let catalog = crate::tools::build_catalog(true); + let projection = catalog.modern_list(); + let tools = projection["tools"].as_array().unwrap(); + assert_eq!(tools.len(), EXPECTED_TOOLS.len()); + + for (tool, (expected_name, expected_annotations)) in tools.iter().zip(EXPECTED_TOOLS) { + assert_eq!(tool["name"], expected_name); + assert_eq!(tool["annotations"], expected_annotations.as_value()); + assert_eq!( + tool.pointer("/inputSchema/additionalProperties"), + Some(&Value::Bool(false)) + ); + assert!(tool + .pointer("/inputSchema/required") + .is_some_and(Value::is_array)); + assert!(tool.get("outputSchema").is_none()); + + let requirements = &tool["_meta"]["com.github.rtk-ai.icm/requirements"]; + assert_eq!(requirements["minimumProtocolRevision"], "2024-11-05"); + assert_eq!(requirements["serverFacilities"]["store"], "required"); + assert_eq!(requirements["requiredClientCapabilities"], json!([])); + assert!(requirements["structuredOutputFromRevision"].is_null()); + assert_eq!(requirements["legacyVisible"], true); + let expected_embedder = if expected_name == "icm_memory_embed_all" { + "required" + } else if matches!( + expected_name, + "icm_memory_store" + | "icm_memory_recall" + | "icm_memory_consolidate" + | "icm_memory_update" + | "icm_feedback_record" + | "icm_feedback_search" + ) { + "optional" + } else { + "unused" + }; + assert_eq!( + requirements["serverFacilities"]["embedder"], + expected_embedder + ); + assert_eq!( + requirements["serverFacilities"]["filesystemRead"], + if expected_name == "icm_learn" { + "required" + } else { + "unused" + } + ); + } + } + + #[test] + fn modern_schema_and_validator_share_the_frozen_core_bounds() { + let catalog = crate::tools::build_catalog(false); + let projection = catalog.modern_list(); + let tools = projection["tools"].as_array().unwrap(); + let store = tools + .iter() + .find(|tool| tool["name"] == "icm_memory_store") + .unwrap(); + assert_eq!( + store.pointer("/inputSchema/properties/topic/maxLength"), + Some(&json!(255)) + ); + assert_eq!( + store.pointer("/inputSchema/properties/content/x-icm-maxUtf8Bytes"), + Some(&json!(65_536)) + ); + assert_eq!( + store.pointer("/inputSchema/properties/topic/x-icm-maxUtf8Bytes"), + Some(&json!(255)) + ); + let recall = tools + .iter() + .find(|tool| tool["name"] == "icm_memory_recall") + .unwrap(); + assert_eq!( + recall.pointer("/inputSchema/properties/limit/minimum"), + Some(&json!(1)) + ); + assert_eq!( + recall.pointer("/inputSchema/properties/limit/maximum"), + Some(&json!(100)) + ); + + for (tool_name, field, maximum) in [ + ("icm_memoir_create", "name", 255), + ("icm_memoir_create", "description", 10_000), + ("icm_memoir_add_concept", "name", 255), + ("icm_memoir_add_concept", "definition", 10_000), + ("icm_memoir_refine", "name", 255), + ("icm_memoir_refine", "definition", 10_000), + ] { + let tool = tools.iter().find(|tool| tool["name"] == tool_name).unwrap(); + assert_eq!( + tool.pointer(&format!("/inputSchema/properties/{field}/maxLength")), + Some(&json!(maximum)) + ); + assert_eq!( + tool.pointer(&format!( + "/inputSchema/properties/{field}/x-icm-maxUtf8Bytes" + )), + Some(&json!(maximum)) + ); + } + } + + #[test] + fn memoir_schema_byte_limits_match_catalog_runtime_validation() { + let catalog = crate::tools::build_catalog(false); + let store = Store::in_memory().unwrap(); + let working_directory = std::env::current_dir().unwrap(); + let context = ToolContext { + store: &store, + embedder: None, + compact: false, + auto_consolidate: AutoConsolidate::default(), + working_directory: &working_directory, + enforce_directory_boundary: true, + }; + + let exact_name = "n".repeat(255); + let exact_description = "é".repeat(5_000); + assert!(matches!( + catalog.dispatch( + &context, + "icm_memoir_create", + &json!({"name":exact_name,"description":exact_description}), + InputValidation::Modern + ), + DispatchResult::ToolResult(_) + )); + + for arguments in [ + json!({"name":"n".repeat(256)}), + json!({"name":"é".repeat(128)}), + json!({"name":"short","description":"d".repeat(10_001)}), + json!({"name":"short","description":"é".repeat(5_001)}), + ] { + assert!(matches!( + catalog.dispatch( + &context, + "icm_memoir_create", + &arguments, + InputValidation::Modern + ), + DispatchResult::InvalidInput(_) + )); + } + + let exact_concept_name = "c".repeat(255); + assert!(matches!( + catalog.dispatch( + &context, + "icm_memoir_add_concept", + &json!({ + "memoir":exact_name, + "name":exact_concept_name, + "definition":"d".repeat(10_000) + }), + InputValidation::Modern + ), + DispatchResult::ToolResult(_) + )); + assert!(matches!( + catalog.dispatch( + &context, + "icm_memoir_refine", + &json!({ + "memoir":exact_name, + "name":exact_concept_name, + "definition":"é".repeat(5_000) + }), + InputValidation::Modern + ), + DispatchResult::ToolResult(_) + )); + + for (tool, arguments) in [ + ( + "icm_memoir_add_concept", + json!({"memoir":"m","name":"é".repeat(128),"definition":"valid"}), + ), + ( + "icm_memoir_add_concept", + json!({"memoir":"m","name":"valid","definition":"é".repeat(5_001)}), + ), + ( + "icm_memoir_refine", + json!({"memoir":"m","name":"n".repeat(256),"definition":"valid"}), + ), + ( + "icm_memoir_refine", + json!({"memoir":"m","name":"valid","definition":"d".repeat(10_001)}), + ), + ] { + assert!(matches!( + catalog.dispatch(&context, tool, &arguments, InputValidation::Modern), + DispatchResult::InvalidInput(_) + )); + } + } +} diff --git a/crates/icm-mcp/src/inputs.rs b/crates/icm-mcp/src/inputs.rs new file mode 100644 index 00000000..d1df079d --- /dev/null +++ b/crates/icm-mcp/src/inputs.rs @@ -0,0 +1,483 @@ +//! Typed MCP tool inputs. +//! +//! These types are the modern input contract. The catalog derives closed JSON +//! Schemas from them and deserializes every modern call before dispatch. The +//! separate legacy schema projection remains a frozen compatibility artifact. + +#![allow(dead_code)] + +use schemars::JsonSchema; +use serde::de::DeserializeOwned; +use serde::Deserialize; +use serde_json::{json, Value}; + +pub trait ModernToolInput: DeserializeOwned + JsonSchema { + fn refine_schema(_schema: &mut Value) {} +} + +fn property<'a>(schema: &'a mut Value, name: &str) -> &'a mut serde_json::Map { + schema + .pointer_mut(&format!("/properties/{name}")) + .and_then(Value::as_object_mut) + .unwrap_or_else(|| panic!("generated input schema is missing property {name}")) +} + +fn string_bounds( + schema: &mut Value, + name: &str, + minimum_code_points: Option, + maximum_code_points: Option, + maximum_utf8_bytes: Option, +) { + let property = property(schema, name); + if let Some(minimum) = minimum_code_points { + property.insert("minLength".into(), json!(minimum)); + } + if let Some(maximum) = maximum_code_points { + property.insert("maxLength".into(), json!(maximum)); + } + if let Some(maximum) = maximum_utf8_bytes { + // JSON Schema has no UTF-8 byte-length keyword. Keep the portable + // code-point ceiling and publish the exact byte contract as a local + // annotation that the catalog validator enforces before dispatch. + property.insert("x-icm-maxUtf8Bytes".into(), json!(maximum)); + } +} + +fn integer_bounds(schema: &mut Value, name: &str, minimum: i64, maximum: i64) { + let property = property(schema, name); + property.insert("minimum".into(), json!(minimum)); + property.insert("maximum".into(), json!(maximum)); +} + +fn topic_bounds(schema: &mut Value, name: &str) { + string_bounds(schema, name, Some(1), Some(255), Some(255)); + property(schema, name).insert("x-icm-trimmedNonEmpty".into(), Value::Bool(true)); +} + +fn content_bounds(schema: &mut Value, name: &str) { + string_bounds(schema, name, Some(1), Some(65_536), Some(65_536)); + property(schema, name).insert("x-icm-trimmedNonEmpty".into(), Value::Bool(true)); +} + +macro_rules! default_contract { + ($($name:ty),+ $(,)?) => { + $(impl ModernToolInput for $name {})+ + }; +} + +macro_rules! empty_input { + ($name:ident) => { + #[derive(Debug, Deserialize, JsonSchema)] + #[serde(deny_unknown_fields)] + pub struct $name {} + }; +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(rename_all = "lowercase")] +pub enum ImportanceInput { + Critical, + High, + Medium, + Low, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum RelationInput { + PartOf, + DependsOn, + RelatedTo, + Contradicts, + Refines, + AlternativeTo, + CausedBy, + InstanceOf, + SupersededBy, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(rename_all = "lowercase")] +pub enum ExportFormatInput { + Json, + Dot, + Ascii, + Ai, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(rename_all = "lowercase")] +pub enum TranscriptRoleInput { + User, + Assistant, + System, + Tool, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(rename_all = "lowercase")] +pub enum WakeUpFormatInput { + Markdown, + Plain, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct MemoryStoreInput { + pub topic: String, + pub content: String, + pub importance: Option, + pub keywords: Option>, + pub raw_excerpt: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct MemoryRecallInput { + pub query: String, + pub topic: Option, + pub limit: Option, + pub keyword: Option, + pub project: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct MemoryForgetInput { + pub id: String, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct TopicInput { + pub topic: String, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct LearnInput { + pub directory: Option, + pub name: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct MemoryConsolidateInput { + pub topic: String, + pub summary: String, +} + +empty_input!(MemoryListTopicsInput); +empty_input!(MemoryStatsInput); + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct MemoryUpdateInput { + pub id: String, + pub content: String, + pub importance: Option, + pub keywords: Option>, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct MemoryHealthInput { + pub topic: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct MemoirCreateInput { + pub name: String, + pub description: Option, +} + +empty_input!(MemoirListInput); + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct NameInput { + pub name: String, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct MemoirAddConceptInput { + pub memoir: String, + pub name: String, + pub definition: String, + pub labels: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct MemoirRefineInput { + pub memoir: String, + pub name: String, + pub definition: String, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct MemoirSearchInput { + pub memoir: String, + pub query: String, + pub label: Option, + pub limit: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct MemoirLinkInput { + pub memoir: String, + pub r#from: String, + pub to: String, + pub relation: RelationInput, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct MemoirInspectInput { + pub memoir: String, + pub name: String, + pub depth: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct MemoirExportInput { + pub name: String, + pub format: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct ExtractPatternsInput { + pub topic: String, + pub memoir: Option, + pub min_cluster_size: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct MemoirSearchAllInput { + pub query: String, + pub limit: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct FeedbackRecordInput { + pub topic: String, + pub context: String, + pub predicted: String, + pub corrected: String, + pub reason: Option, + pub source: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct FeedbackSearchInput { + pub query: String, + pub topic: Option, + pub limit: Option, +} + +empty_input!(FeedbackStatsInput); + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct TranscriptStartInput { + pub agent: Option, + pub project: Option, + pub metadata: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct TranscriptRecordInput { + pub session_id: String, + pub role: TranscriptRoleInput, + pub content: String, + pub tool_name: Option, + pub tokens: Option, + pub metadata: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct TranscriptSearchInput { + pub query: String, + pub session_id: Option, + pub project: Option, + pub limit: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct TranscriptShowInput { + pub session_id: String, + pub limit: Option, +} + +empty_input!(TranscriptStatsInput); + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct WakeUpInput { + pub project: Option, + pub max_tokens: Option, + pub format: Option, + pub include_preferences: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct EmbedAllInput { + pub topic: Option, +} + +impl ModernToolInput for MemoryStoreInput { + fn refine_schema(schema: &mut Value) { + topic_bounds(schema, "topic"); + content_bounds(schema, "content"); + string_bounds(schema, "raw_excerpt", None, Some(65_536), Some(65_536)); + } +} + +impl ModernToolInput for MemoryRecallInput { + fn refine_schema(schema: &mut Value) { + string_bounds(schema, "query", Some(1), Some(65_536), Some(65_536)); + property(schema, "query").insert("x-icm-trimmedNonEmpty".into(), Value::Bool(true)); + topic_bounds(schema, "topic"); + integer_bounds(schema, "limit", 1, 100); + } +} + +impl ModernToolInput for TopicInput { + fn refine_schema(schema: &mut Value) { + topic_bounds(schema, "topic"); + } +} + +impl ModernToolInput for MemoryConsolidateInput { + fn refine_schema(schema: &mut Value) { + topic_bounds(schema, "topic"); + content_bounds(schema, "summary"); + } +} + +impl ModernToolInput for MemoryUpdateInput { + fn refine_schema(schema: &mut Value) { + content_bounds(schema, "content"); + } +} + +impl ModernToolInput for MemoryHealthInput { + fn refine_schema(schema: &mut Value) { + topic_bounds(schema, "topic"); + } +} + +impl ModernToolInput for MemoirCreateInput { + fn refine_schema(schema: &mut Value) { + string_bounds(schema, "name", None, Some(255), Some(255)); + string_bounds(schema, "description", None, Some(10_000), Some(10_000)); + } +} + +impl ModernToolInput for MemoirAddConceptInput { + fn refine_schema(schema: &mut Value) { + string_bounds(schema, "name", None, Some(255), Some(255)); + string_bounds(schema, "definition", None, Some(10_000), Some(10_000)); + } +} + +impl ModernToolInput for MemoirRefineInput { + fn refine_schema(schema: &mut Value) { + string_bounds(schema, "name", None, Some(255), Some(255)); + string_bounds(schema, "definition", None, Some(10_000), Some(10_000)); + } +} + +impl ModernToolInput for MemoirSearchInput { + fn refine_schema(schema: &mut Value) { + integer_bounds(schema, "limit", 1, 100); + } +} + +impl ModernToolInput for MemoirInspectInput { + fn refine_schema(schema: &mut Value) { + integer_bounds(schema, "depth", 1, 3); + } +} + +impl ModernToolInput for ExtractPatternsInput { + fn refine_schema(schema: &mut Value) { + topic_bounds(schema, "topic"); + integer_bounds(schema, "min_cluster_size", 2, 50); + } +} + +impl ModernToolInput for MemoirSearchAllInput { + fn refine_schema(schema: &mut Value) { + integer_bounds(schema, "limit", 1, 100); + } +} + +impl ModernToolInput for FeedbackRecordInput { + fn refine_schema(schema: &mut Value) { + topic_bounds(schema, "topic"); + for name in ["context", "predicted", "corrected", "reason"] { + string_bounds(schema, name, None, Some(20_000), Some(20_000)); + } + } +} + +impl ModernToolInput for FeedbackSearchInput { + fn refine_schema(schema: &mut Value) { + topic_bounds(schema, "topic"); + integer_bounds(schema, "limit", 1, 100); + } +} + +impl ModernToolInput for TranscriptSearchInput { + fn refine_schema(schema: &mut Value) { + integer_bounds(schema, "limit", 1, 50); + } +} + +impl ModernToolInput for TranscriptShowInput { + fn refine_schema(schema: &mut Value) { + integer_bounds(schema, "limit", 1, 2_000); + } +} + +impl ModernToolInput for WakeUpInput { + fn refine_schema(schema: &mut Value) { + integer_bounds(schema, "max_tokens", 20, 4_000); + } +} + +impl ModernToolInput for EmbedAllInput { + fn refine_schema(schema: &mut Value) { + topic_bounds(schema, "topic"); + } +} + +default_contract!( + MemoryForgetInput, + LearnInput, + MemoryListTopicsInput, + MemoryStatsInput, + MemoirListInput, + NameInput, + MemoirLinkInput, + MemoirExportInput, + FeedbackStatsInput, + TranscriptStartInput, + TranscriptRecordInput, + TranscriptStatsInput, +); diff --git a/crates/icm-mcp/src/lib.rs b/crates/icm-mcp/src/lib.rs index 7b40295e..0b01c441 100644 --- a/crates/icm-mcp/src/lib.rs +++ b/crates/icm-mcp/src/lib.rs @@ -1,3 +1,5 @@ +pub mod catalog; +mod inputs; pub mod protocol; pub mod server; pub mod tools; diff --git a/crates/icm-mcp/src/tools.rs b/crates/icm-mcp/src/tools.rs index 33e7bcf6..2612ffd2 100644 --- a/crates/icm-mcp/src/tools.rs +++ b/crates/icm-mcp/src/tools.rs @@ -10,6 +10,10 @@ use icm_core::{ }; use icm_store::Store; +use crate::catalog::{ + DispatchResult, ToolAnnotations, ToolCatalog, ToolContext, ToolRequirements, ToolSpec, +}; +use crate::inputs::*; use crate::protocol::ToolResult; /// Historical default threshold for auto-consolidation. The live value comes @@ -44,8 +48,8 @@ impl Default for AutoConsolidate { } } -/// Maximum allowed length for topic names. Must stay <= the store -/// layer's `MAX_TOPIC_BYTES` so the MCP-level rejection happens +/// Maximum allowed UTF-8 byte length for topic names. Must stay <= the +/// store layer's `MAX_TOPIC_BYTES` so the MCP-level rejection happens /// *before* the store's lower-level validation does. const MAX_TOPIC_LEN: usize = 255; @@ -106,653 +110,936 @@ fn try_auto_consolidate( // Tool schemas for tools/list // --------------------------------------------------------------------------- -pub fn tool_definitions(has_embedder: bool) -> Value { - let mut tools = vec![ - // --- Memory tools --- +fn normalize_legacy_recall_input(arguments: &Value) -> Value { + let mut normalized = arguments.clone(); + let Some(object) = normalized.as_object_mut() else { + return normalized; + }; + let Some(limit) = object.get("limit").filter(|limit| limit.is_number()) else { + return normalized; + }; + + // The frozen 2024 handler read limits as i64, defaulted unrepresentable + // numeric values to five, and clamped the result to its advertised 1..20 + // range. Normalize only for that catalog projection; the modern DTO keeps + // its strict 1..100 contract and reaches the handler unchanged. + let normalized_limit = limit.as_i64().unwrap_or(5).clamp(1, 20); + object.insert("limit".into(), json!(normalized_limit)); + normalized +} + +macro_rules! tool_spec { + ( + $input:ty, json!({ - "name": "icm_memory_store", - "description": "Store important information in ICM long-term memory. Use to save decisions, preferences, project context, resolved errors — anything that should persist between sessions.", - "inputSchema": { - "type": "object", - "properties": { - "topic": { - "type": "string", - "description": "Category/namespace. Use the canonical topics from the server instructions: 'decisions-{project}', 'preferences', 'errors-resolved', 'context-{project}' — mixed-language topic names fragment the memory." - }, - "content": { - "type": "string", - "description": "Information to memorize — be concise but complete" - }, - "importance": { - "type": "string", - "enum": ["critical", "high", "medium", "low"], - "default": "medium", - "description": "critical=never forgotten, high=slow decay, medium=normal, low=fast decay" - }, - "keywords": { - "type": "array", - "items": { "type": "string" }, - "description": "Keywords to improve search" - }, - "raw_excerpt": { - "type": "string", - "description": "Optional verbatim (code, exact error message, etc.)" - } - }, - "required": ["topic", "content"] - } + "name": $name:literal, + "description": $description:literal, + "inputSchema": $input_schema:tt }), + $annotations:expr, + $handler:expr + ) => { + ToolSpec::typed::<$input>( + $name, + $description, + json!($input_schema), + None, + $annotations, + ToolRequirements::STORE, + $handler, + ) + }; + ( + $input:ty, json!({ - "name": "icm_memory_recall", - "description": "Search ICM long-term memory. Use to find past decisions, project context, preferences, or solutions to previously encountered problems.", - "inputSchema": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "Natural language search query" - }, - "topic": { - "type": "string", - "description": "Filter by specific topic (optional)" - }, - "limit": { - "type": "integer", - "default": 5, - "minimum": 1, - "maximum": 20, - "description": "Max number of results" - }, - "keyword": { - "type": "string", - "description": "Filter results by keyword (exact match on memory keywords)" - }, - "project": { - "type": "string", - "description": "Project filter (segment-aware). Defaults to the server's cwd directory name. Pass an empty string to disable the filter and search across all projects." - } - }, - "required": ["query"] - } + "name": $name:literal, + "description": $description:literal, + "inputSchema": $input_schema:tt }), + legacy_normalizer: $legacy_normalizer:expr, + $annotations:expr, + $handler:expr + ) => { + ToolSpec::typed::<$input>( + $name, + $description, + json!($input_schema), + Some($legacy_normalizer), + $annotations, + ToolRequirements::STORE, + $handler, + ) + }; + ( + $input:ty, json!({ - "name": "icm_memory_forget", - "description": "Delete a specific memory by its ID. Use when information is obsolete or incorrect.", - "inputSchema": { - "type": "object", - "properties": { - "id": { - "type": "string", - "description": "Memory ID to delete" - } - }, - "required": ["id"] - } + "name": $name:literal, + "description": $description:literal, + "inputSchema": $input_schema:tt }), + requirements: $requirements:expr, + $annotations:expr, + $handler:expr + ) => { + ToolSpec::typed::<$input>( + $name, + $description, + json!($input_schema), + None, + $annotations, + $requirements, + $handler, + ) + }; + ( + $input:ty, json!({ - "name": "icm_memory_forget_topic", - "description": "Delete ALL memories in a topic. Use to clear an entire topic at once.", - "inputSchema": { - "type": "object", - "properties": { - "topic": { - "type": "string", - "description": "Topic whose memories should all be deleted" - } - }, - "required": ["topic"] - } + "name": $name:literal, + "description": $description:literal, + "inputSchema": $input_schema:tt }), - json!({ - "name": "icm_learn", - "description": "Scan a project directory and create a Memoir knowledge graph with its structure, dependencies, modules, and config files.", - "inputSchema": { - "type": "object", - "properties": { - "directory": { - "type": "string", - "description": "Project directory to scan (default: current working directory)" + legacy_normalizer: $legacy_normalizer:expr, + requirements: $requirements:expr, + $annotations:expr, + $handler:expr + ) => { + ToolSpec::typed::<$input>( + $name, + $description, + json!($input_schema), + Some($legacy_normalizer), + $annotations, + $requirements, + $handler, + ) + }; +} + +pub(crate) fn build_catalog(has_embedder: bool) -> ToolCatalog { + let tools = vec![ + // --- Memory tools --- + tool_spec!( + MemoryStoreInput, + json!({ + "name": "icm_memory_store", + "description": "Store important information in ICM long-term memory. Use to save decisions, preferences, project context, resolved errors — anything that should persist between sessions.", + "inputSchema": { + "type": "object", + "properties": { + "topic": { + "type": "string", + "description": "Category/namespace. Use the canonical topics from the server instructions: 'decisions-{project}', 'preferences', 'errors-resolved', 'context-{project}' — mixed-language topic names fragment the memory." + }, + "content": { + "type": "string", + "description": "Information to memorize — be concise but complete" + }, + "importance": { + "type": "string", + "enum": ["critical", "high", "medium", "low"], + "default": "medium", + "description": "critical=never forgotten, high=slow decay, medium=normal, low=fast decay" + }, + "keywords": { + "type": "array", + "items": { "type": "string" }, + "description": "Keywords to improve search" + }, + "raw_excerpt": { + "type": "string", + "description": "Optional verbatim (code, exact error message, etc.)" + } }, - "name": { - "type": "string", - "description": "Memoir name (default: directory name)" - } + "required": ["topic", "content"] } - } - }), - json!({ - "name": "icm_memory_consolidate", - "description": "Consolidate all memories of a topic into a single summary. Useful when a topic accumulates too many entries.", - "inputSchema": { - "type": "object", - "properties": { - "topic": { - "type": "string", - "description": "Topic to consolidate" - }, - "summary": { - "type": "string", - "description": "Consolidated summary to replace all memories in the topic" - } - }, - "required": ["topic", "summary"] - } - }), - json!({ - "name": "icm_memory_list_topics", - "description": "List all available topics in memory with their counts.", - "inputSchema": { - "type": "object", - "properties": {} - } - }), - json!({ - "name": "icm_memory_stats", - "description": "Get global ICM memory statistics.", - "inputSchema": { - "type": "object", - "properties": {} - } - }), - json!({ - "name": "icm_memory_update", - "description": "Update an existing memory in-place. Use to correct, refresh, or extend a memory without creating a duplicate.", - "inputSchema": { - "type": "object", - "properties": { - "id": { - "type": "string", - "description": "Memory ID to update" - }, - "content": { - "type": "string", - "description": "New content (replaces existing summary)" - }, - "importance": { - "type": "string", - "enum": ["critical", "high", "medium", "low"], - "description": "New importance level (optional, keeps existing if not set)" + }), + requirements: ToolRequirements::STORE.with_optional_embedder(), + ToolAnnotations::new(false, true, false, false), + |context, args| tool_store( + context.store, + context.embedder, + args, + context.compact, + context.auto_consolidate + ) + ), + tool_spec!( + MemoryRecallInput, + json!({ + "name": "icm_memory_recall", + "description": "Search ICM long-term memory. Use to find past decisions, project context, preferences, or solutions to previously encountered problems.", + "inputSchema": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Natural language search query" + }, + "topic": { + "type": "string", + "description": "Filter by specific topic (optional)" + }, + "limit": { + "type": "integer", + "default": 5, + "minimum": 1, + "maximum": 20, + "description": "Max number of results" + }, + "keyword": { + "type": "string", + "description": "Filter results by keyword (exact match on memory keywords)" + }, + "project": { + "type": "string", + "description": "Project filter (segment-aware). Defaults to the server's cwd directory name. Pass an empty string to disable the filter and search across all projects." + } }, - "keywords": { - "type": "array", - "items": { "type": "string" }, - "description": "New keywords (optional, keeps existing if not set)" - } - }, - "required": ["id", "content"] - } - }), - json!({ - "name": "icm_memory_health", - "description": "Get health stats for all topics: entry count, staleness, consolidation needs. Use to audit memory hygiene.", - "inputSchema": { - "type": "object", - "properties": { - "topic": { - "type": "string", - "description": "Check a specific topic (optional — checks all if omitted)" - } + "required": ["query"] } - } - }), - // --- Memoir tools --- - json!({ - "name": "icm_memoir_create", - "description": "Create a new memoir — a permanent knowledge container. Memoirs hold concepts that never decay.", - "inputSchema": { - "type": "object", - "properties": { - "name": { - "type": "string", - "description": "Unique human-readable name for the memoir" - }, - "description": { - "type": "string", - "description": "Description of what this memoir is for" - } - }, - "required": ["name"] - } - }), - json!({ - "name": "icm_memoir_list", - "description": "List all memoirs with their concept counts.", - "inputSchema": { - "type": "object", - "properties": {} - } - }), - json!({ - "name": "icm_memoir_show", - "description": "Show a memoir's stats, labels, and all its concepts.", - "inputSchema": { - "type": "object", - "properties": { - "name": { - "type": "string", - "description": "Memoir name" - } - }, - "required": ["name"] - } - }), - json!({ - "name": "icm_memoir_add_concept", - "description": "Add a permanent concept to a memoir. Concepts are knowledge nodes that get refined, never decayed.", - "inputSchema": { - "type": "object", - "properties": { - "memoir": { - "type": "string", - "description": "Memoir name" - }, - "name": { - "type": "string", - "description": "Concept name (unique within memoir)" + }), + legacy_normalizer: normalize_legacy_recall_input, + requirements: ToolRequirements::STORE.with_optional_embedder(), + ToolAnnotations::new(false, true, false, false), + |context, args| tool_recall(context.store, context.embedder, args, context.compact) + ), + tool_spec!( + MemoryForgetInput, + json!({ + "name": "icm_memory_forget", + "description": "Delete a specific memory by its ID. Use when information is obsolete or incorrect.", + "inputSchema": { + "type": "object", + "properties": { + "id": { + "type": "string", + "description": "Memory ID to delete" + } }, - "definition": { - "type": "string", - "description": "Dense description of the concept" + "required": ["id"] + } + }), + ToolAnnotations::new(false, true, true, false), + |context, args| tool_forget(context.store, args) + ), + tool_spec!( + TopicInput, + json!({ + "name": "icm_memory_forget_topic", + "description": "Delete ALL memories in a topic. Use to clear an entire topic at once.", + "inputSchema": { + "type": "object", + "properties": { + "topic": { + "type": "string", + "description": "Topic whose memories should all be deleted" + } }, - "labels": { - "type": "string", - "description": "Comma-separated labels (namespace:value or plain tag). E.g. 'domain:arch,type:decision'" + "required": ["topic"] + } + }), + ToolAnnotations::new(false, true, true, false), + |context, args| tool_forget_topic(context.store, args) + ), + tool_spec!( + LearnInput, + json!({ + "name": "icm_learn", + "description": "Scan a project directory and create a Memoir knowledge graph with its structure, dependencies, modules, and config files.", + "inputSchema": { + "type": "object", + "properties": { + "directory": { + "type": "string", + "description": "Project directory to scan (default: current working directory)" + }, + "name": { + "type": "string", + "description": "Memoir name (default: directory name)" + } } - }, - "required": ["memoir", "name", "definition"] - } - }), - json!({ - "name": "icm_memoir_refine", - "description": "Refine an existing concept with a new, improved definition. Bumps revision and boosts confidence.", - "inputSchema": { - "type": "object", - "properties": { - "memoir": { - "type": "string", - "description": "Memoir name" + } + }), + requirements: ToolRequirements::STORE.with_filesystem_read(), + ToolAnnotations::new(false, true, false, true), + tool_learn_bounded + ), + tool_spec!( + MemoryConsolidateInput, + json!({ + "name": "icm_memory_consolidate", + "description": "Consolidate all memories of a topic into a single summary. Useful when a topic accumulates too many entries.", + "inputSchema": { + "type": "object", + "properties": { + "topic": { + "type": "string", + "description": "Topic to consolidate" + }, + "summary": { + "type": "string", + "description": "Consolidated summary to replace all memories in the topic" + } }, - "name": { - "type": "string", - "description": "Concept name" + "required": ["topic", "summary"] + } + }), + requirements: ToolRequirements::STORE.with_optional_embedder(), + ToolAnnotations::new(false, true, false, false), + |context, args| tool_consolidate(context.store, context.embedder, args) + ), + tool_spec!( + MemoryListTopicsInput, + json!({ + "name": "icm_memory_list_topics", + "description": "List all available topics in memory with their counts.", + "inputSchema": { + "type": "object", + "properties": {} + } + }), + ToolAnnotations::new(true, false, true, false), + |context, _| tool_list_topics(context.store) + ), + tool_spec!( + MemoryStatsInput, + json!({ + "name": "icm_memory_stats", + "description": "Get global ICM memory statistics.", + "inputSchema": { + "type": "object", + "properties": {} + } + }), + ToolAnnotations::new(true, false, true, false), + |context, _| tool_stats(context.store) + ), + tool_spec!( + MemoryUpdateInput, + json!({ + "name": "icm_memory_update", + "description": "Update an existing memory in-place. Use to correct, refresh, or extend a memory without creating a duplicate.", + "inputSchema": { + "type": "object", + "properties": { + "id": { + "type": "string", + "description": "Memory ID to update" + }, + "content": { + "type": "string", + "description": "New content (replaces existing summary)" + }, + "importance": { + "type": "string", + "enum": ["critical", "high", "medium", "low"], + "description": "New importance level (optional, keeps existing if not set)" + }, + "keywords": { + "type": "array", + "items": { "type": "string" }, + "description": "New keywords (optional, keeps existing if not set)" + } }, - "definition": { - "type": "string", - "description": "New, refined definition" + "required": ["id", "content"] + } + }), + requirements: ToolRequirements::STORE.with_optional_embedder(), + ToolAnnotations::new(false, true, false, false), + |context, args| tool_update(context.store, context.embedder, args) + ), + tool_spec!( + MemoryHealthInput, + json!({ + "name": "icm_memory_health", + "description": "Get health stats for all topics: entry count, staleness, consolidation needs. Use to audit memory hygiene.", + "inputSchema": { + "type": "object", + "properties": { + "topic": { + "type": "string", + "description": "Check a specific topic (optional — checks all if omitted)" + } } - }, - "required": ["memoir", "name", "definition"] - } - }), - json!({ - "name": "icm_memoir_search", - "description": "Full-text search concepts within a memoir.", - "inputSchema": { - "type": "object", - "properties": { - "memoir": { - "type": "string", - "description": "Memoir name" - }, - "query": { - "type": "string", - "description": "Search query" - }, - "label": { - "type": "string", - "description": "Filter by label (e.g. 'domain:tech')" + } + }), + ToolAnnotations::new(true, false, true, false), + |context, args| tool_health(context.store, args) + ), + // --- Memoir tools --- + tool_spec!( + MemoirCreateInput, + json!({ + "name": "icm_memoir_create", + "description": "Create a new memoir — a permanent knowledge container. Memoirs hold concepts that never decay.", + "inputSchema": { + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Unique human-readable name for the memoir" + }, + "description": { + "type": "string", + "description": "Description of what this memoir is for" + } }, - "limit": { - "type": "integer", - "default": 10, - "description": "Max results" - } - }, - "required": ["memoir", "query"] - } - }), - json!({ - "name": "icm_memoir_link", - "description": "Create a directed, typed edge between two concepts in the same memoir.", - "inputSchema": { - "type": "object", - "properties": { - "memoir": { - "type": "string", - "description": "Memoir name" + "required": ["name"] + } + }), + ToolAnnotations::new(false, false, false, false), + |context, args| tool_memoir_create(context.store, args) + ), + tool_spec!( + MemoirListInput, + json!({ + "name": "icm_memoir_list", + "description": "List all memoirs with their concept counts.", + "inputSchema": { + "type": "object", + "properties": {} + } + }), + ToolAnnotations::new(true, false, true, false), + |context, _| tool_memoir_list(context.store) + ), + tool_spec!( + NameInput, + json!({ + "name": "icm_memoir_show", + "description": "Show a memoir's stats, labels, and all its concepts.", + "inputSchema": { + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Memoir name" + } }, - "from": { - "type": "string", - "description": "Source concept name" + "required": ["name"] + } + }), + ToolAnnotations::new(true, false, true, false), + |context, args| tool_memoir_show(context.store, args) + ), + tool_spec!( + MemoirAddConceptInput, + json!({ + "name": "icm_memoir_add_concept", + "description": "Add a permanent concept to a memoir. Concepts are knowledge nodes that get refined, never decayed.", + "inputSchema": { + "type": "object", + "properties": { + "memoir": { + "type": "string", + "description": "Memoir name" + }, + "name": { + "type": "string", + "description": "Concept name (unique within memoir)" + }, + "definition": { + "type": "string", + "description": "Dense description of the concept" + }, + "labels": { + "type": "string", + "description": "Comma-separated labels (namespace:value or plain tag). E.g. 'domain:arch,type:decision'" + } }, - "to": { - "type": "string", - "description": "Target concept name" + "required": ["memoir", "name", "definition"] + } + }), + ToolAnnotations::new(false, false, false, false), + |context, args| tool_memoir_add_concept(context.store, args) + ), + tool_spec!( + MemoirRefineInput, + json!({ + "name": "icm_memoir_refine", + "description": "Refine an existing concept with a new, improved definition. Bumps revision and boosts confidence.", + "inputSchema": { + "type": "object", + "properties": { + "memoir": { + "type": "string", + "description": "Memoir name" + }, + "name": { + "type": "string", + "description": "Concept name" + }, + "definition": { + "type": "string", + "description": "New, refined definition" + } }, - "relation": { - "type": "string", - "enum": ["part_of", "depends_on", "related_to", "contradicts", "refines", "alternative_to", "caused_by", "instance_of", "superseded_by"], - "description": "Relation type" - } - }, - "required": ["memoir", "from", "to", "relation"] - } - }), - json!({ - "name": "icm_memoir_inspect", - "description": "Inspect a concept and its graph neighborhood (BFS).", - "inputSchema": { - "type": "object", - "properties": { - "memoir": { - "type": "string", - "description": "Memoir name" + "required": ["memoir", "name", "definition"] + } + }), + ToolAnnotations::new(false, true, false, false), + |context, args| tool_memoir_refine(context.store, args) + ), + tool_spec!( + MemoirSearchInput, + json!({ + "name": "icm_memoir_search", + "description": "Full-text search concepts within a memoir.", + "inputSchema": { + "type": "object", + "properties": { + "memoir": { + "type": "string", + "description": "Memoir name" + }, + "query": { + "type": "string", + "description": "Search query" + }, + "label": { + "type": "string", + "description": "Filter by label (e.g. 'domain:tech')" + }, + "limit": { + "type": "integer", + "default": 10, + "description": "Max results" + } }, - "name": { - "type": "string", - "description": "Concept name" + "required": ["memoir", "query"] + } + }), + ToolAnnotations::new(true, false, true, false), + |context, args| tool_memoir_search(context.store, args) + ), + tool_spec!( + MemoirLinkInput, + json!({ + "name": "icm_memoir_link", + "description": "Create a directed, typed edge between two concepts in the same memoir.", + "inputSchema": { + "type": "object", + "properties": { + "memoir": { + "type": "string", + "description": "Memoir name" + }, + "from": { + "type": "string", + "description": "Source concept name" + }, + "to": { + "type": "string", + "description": "Target concept name" + }, + "relation": { + "type": "string", + "enum": ["part_of", "depends_on", "related_to", "contradicts", "refines", "alternative_to", "caused_by", "instance_of", "superseded_by"], + "description": "Relation type" + } }, - "depth": { - "type": "integer", - "default": 1, - "description": "BFS depth" - } - }, - "required": ["memoir", "name"] - } - }), - json!({ - "name": "icm_memoir_export", - "description": "Export a memoir's full concept graph. Formats: json (structured), dot (Graphviz), ascii (visual), ai (compact markdown for LLM context).", - "inputSchema": { - "type": "object", - "properties": { - "name": { - "type": "string", - "description": "Memoir name" + "required": ["memoir", "from", "to", "relation"] + } + }), + ToolAnnotations::new(false, false, false, false), + |context, args| tool_memoir_link(context.store, args) + ), + tool_spec!( + MemoirInspectInput, + json!({ + "name": "icm_memoir_inspect", + "description": "Inspect a concept and its graph neighborhood (BFS).", + "inputSchema": { + "type": "object", + "properties": { + "memoir": { + "type": "string", + "description": "Memoir name" + }, + "name": { + "type": "string", + "description": "Concept name" + }, + "depth": { + "type": "integer", + "default": 1, + "description": "BFS depth" + } }, - "format": { - "type": "string", - "enum": ["json", "dot", "ascii", "ai"], - "default": "json", - "description": "Output format: json (structured), dot (Graphviz), ascii (visual graph), ai (compact markdown for LLM)" - } - }, - "required": ["name"] - } - }), - json!({ - "name": "icm_memory_extract_patterns", - "description": "Detect recurring patterns in a topic by keyword similarity. Optionally create concepts in a memoir from detected patterns.", - "inputSchema": { - "type": "object", - "properties": { - "topic": { - "type": "string", - "description": "Topic to analyze for patterns" + "required": ["memoir", "name"] + } + }), + ToolAnnotations::new(true, false, true, false), + |context, args| tool_memoir_inspect(context.store, args) + ), + tool_spec!( + MemoirExportInput, + json!({ + "name": "icm_memoir_export", + "description": "Export a memoir's full concept graph. Formats: json (structured), dot (Graphviz), ascii (visual), ai (compact markdown for LLM context).", + "inputSchema": { + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Memoir name" + }, + "format": { + "type": "string", + "enum": ["json", "dot", "ascii", "ai"], + "default": "json", + "description": "Output format: json (structured), dot (Graphviz), ascii (visual graph), ai (compact markdown for LLM)" + } }, - "memoir": { - "type": "string", - "description": "Memoir name — if provided, creates concepts from detected patterns" + "required": ["name"] + } + }), + ToolAnnotations::new(true, false, true, false), + |context, args| tool_memoir_export(context.store, args) + ), + tool_spec!( + ExtractPatternsInput, + json!({ + "name": "icm_memory_extract_patterns", + "description": "Detect recurring patterns in a topic by keyword similarity. Optionally create concepts in a memoir from detected patterns.", + "inputSchema": { + "type": "object", + "properties": { + "topic": { + "type": "string", + "description": "Topic to analyze for patterns" + }, + "memoir": { + "type": "string", + "description": "Memoir name — if provided, creates concepts from detected patterns" + }, + "min_cluster_size": { + "type": "integer", + "default": 3, + "minimum": 2, + "description": "Minimum number of similar memories to form a pattern (default: 3)" + } }, - "min_cluster_size": { - "type": "integer", - "default": 3, - "minimum": 2, - "description": "Minimum number of similar memories to form a pattern (default: 3)" - } - }, - "required": ["topic"] - } - }), - json!({ - "name": "icm_memoir_search_all", - "description": "Full-text search concepts across all memoirs.", - "inputSchema": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "Search query" + "required": ["topic"] + } + }), + ToolAnnotations::new(false, false, false, false), + |context, args| tool_extract_patterns(context.store, args) + ), + tool_spec!( + MemoirSearchAllInput, + json!({ + "name": "icm_memoir_search_all", + "description": "Full-text search concepts across all memoirs.", + "inputSchema": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Search query" + }, + "limit": { + "type": "integer", + "default": 10, + "description": "Max results" + } }, - "limit": { - "type": "integer", - "default": 10, - "description": "Max results" - } - }, - "required": ["query"] - } - }), + "required": ["query"] + } + }), + ToolAnnotations::new(true, false, true, false), + |context, args| tool_memoir_search_all(context.store, args) + ), // --- Feedback tools --- - json!({ - "name": "icm_feedback_record", - "description": "Record a correction/feedback when an AI prediction was wrong. Helps improve future predictions by learning from mistakes.", - "inputSchema": { - "type": "object", - "properties": { - "topic": { - "type": "string", - "description": "Category/namespace for this feedback (e.g. 'triage-owner/repo', 'pr-analysis')" - }, - "context": { - "type": "string", - "description": "What was the situation / input that led to the prediction" - }, - "predicted": { - "type": "string", - "description": "What the AI predicted or did" - }, - "corrected": { - "type": "string", - "description": "What the correct answer/action should have been" - }, - "reason": { - "type": "string", - "description": "Why the correction was made (optional)" - }, - "source": { - "type": "string", - "description": "Which tool/pipeline generated the prediction (optional)" - } - }, - "required": ["topic", "context", "predicted", "corrected"] - } - }), - json!({ - "name": "icm_feedback_search", - "description": "Search past feedback/corrections to inform current predictions. Use before making predictions to learn from past mistakes.", - "inputSchema": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "Search query to find relevant past corrections" + tool_spec!( + FeedbackRecordInput, + json!({ + "name": "icm_feedback_record", + "description": "Record a correction/feedback when an AI prediction was wrong. Helps improve future predictions by learning from mistakes.", + "inputSchema": { + "type": "object", + "properties": { + "topic": { + "type": "string", + "description": "Category/namespace for this feedback (e.g. 'triage-owner/repo', 'pr-analysis')" + }, + "context": { + "type": "string", + "description": "What was the situation / input that led to the prediction" + }, + "predicted": { + "type": "string", + "description": "What the AI predicted or did" + }, + "corrected": { + "type": "string", + "description": "What the correct answer/action should have been" + }, + "reason": { + "type": "string", + "description": "Why the correction was made (optional)" + }, + "source": { + "type": "string", + "description": "Which tool/pipeline generated the prediction (optional)" + } }, - "topic": { - "type": "string", - "description": "Filter by topic (optional)" + "required": ["topic", "context", "predicted", "corrected"] + } + }), + requirements: ToolRequirements::STORE.with_optional_embedder(), + ToolAnnotations::new(false, false, false, false), + |context, args| tool_feedback_record( + context.store, + context.embedder, + args, + context.compact + ) + ), + tool_spec!( + FeedbackSearchInput, + json!({ + "name": "icm_feedback_search", + "description": "Search past feedback/corrections to inform current predictions. Use before making predictions to learn from past mistakes.", + "inputSchema": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Search query to find relevant past corrections" + }, + "topic": { + "type": "string", + "description": "Filter by topic (optional)" + }, + "limit": { + "type": "integer", + "default": 5, + "minimum": 1, + "maximum": 20, + "description": "Max number of results" + } }, - "limit": { - "type": "integer", - "default": 5, - "minimum": 1, - "maximum": 20, - "description": "Max number of results" - } - }, - "required": ["query"] - } - }), - json!({ - "name": "icm_feedback_stats", - "description": "Get feedback statistics: total count, breakdown by topic, most applied corrections.", - "inputSchema": { - "type": "object", - "properties": {} - } - }), + "required": ["query"] + } + }), + requirements: ToolRequirements::STORE.with_optional_embedder(), + ToolAnnotations::new(true, false, true, false), + |context, args| tool_feedback_search(context.store, context.embedder, args) + ), + tool_spec!( + FeedbackStatsInput, + json!({ + "name": "icm_feedback_stats", + "description": "Get feedback statistics: total count, breakdown by topic, most applied corrections.", + "inputSchema": { + "type": "object", + "properties": {} + } + }), + ToolAnnotations::new(true, false, true, false), + |context, _| tool_feedback_stats(context.store) + ), // --- Transcript tools (verbatim session replay) --- - json!({ - "name": "icm_transcript_start_session", - "description": "Create a new transcript session for verbatim message capture. Returns the session_id used by subsequent icm_transcript_record calls. Use once per conversation or debugging session.", - "inputSchema": { - "type": "object", - "properties": { - "agent": { - "type": "string", - "description": "Agent identifier (e.g. 'claude-code', 'cursor', 'gemini-cli'). Default: 'mcp'." - }, - "project": { - "type": "string", - "description": "Project name (optional; usually cwd basename or repo slug)" - }, - "metadata": { - "type": "string", - "description": "Arbitrary JSON metadata (optional)" + tool_spec!( + TranscriptStartInput, + json!({ + "name": "icm_transcript_start_session", + "description": "Create a new transcript session for verbatim message capture. Returns the session_id used by subsequent icm_transcript_record calls. Use once per conversation or debugging session.", + "inputSchema": { + "type": "object", + "properties": { + "agent": { + "type": "string", + "description": "Agent identifier (e.g. 'claude-code', 'cursor', 'gemini-cli'). Default: 'mcp'." + }, + "project": { + "type": "string", + "description": "Project name (optional; usually cwd basename or repo slug)" + }, + "metadata": { + "type": "string", + "description": "Arbitrary JSON metadata (optional)" + } } } - } - }), - json!({ - "name": "icm_transcript_record", - "description": "Append a verbatim message to a transcript session. Stores the raw content with no summarization. Use once per user turn, assistant reply, or tool call for full replay fidelity.", - "inputSchema": { - "type": "object", - "properties": { - "session_id": { - "type": "string", - "description": "Session id from icm_transcript_start_session" - }, - "role": { - "type": "string", - "enum": ["user", "assistant", "system", "tool"], - "description": "Message role" - }, - "content": { - "type": "string", - "description": "Raw message content (stored verbatim)" - }, - "tool_name": { - "type": "string", - "description": "Tool name if role=tool (optional)" - }, - "tokens": { - "type": "integer", - "description": "Token count for billing / stats (optional)" - }, - "metadata": { - "type": "string", - "description": "Arbitrary JSON metadata (optional)" - } - }, - "required": ["session_id", "role", "content"] - } - }), - json!({ - "name": "icm_transcript_search", - "description": "Full-text search across recorded transcript messages (FTS5 BM25). Supports boolean operators, phrase matches, and prefix queries. Use to recall exact quotes or debug past decisions.", - "inputSchema": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "FTS5 query: 'postgres OR mysql', '\"exact phrase\"', 'auth*'" - }, - "session_id": { - "type": "string", - "description": "Restrict to one session (optional)" - }, - "project": { - "type": "string", - "description": "Restrict to one project (optional)" - }, - "limit": { - "type": "integer", - "default": 10, - "minimum": 1, - "maximum": 50 - } - }, - "required": ["query"] - } - }), - json!({ - "name": "icm_transcript_show", - "description": "Replay the full message thread of a transcript session, chronologically. Returns up to `limit` messages with role, content, tool name, timestamp.", - "inputSchema": { - "type": "object", - "properties": { - "session_id": { "type": "string" }, - "limit": { "type": "integer", "default": 200, "minimum": 1, "maximum": 2000 } - }, - "required": ["session_id"] - } - }), - json!({ - "name": "icm_transcript_stats", - "description": "Global transcript statistics: session count, message count, total bytes, breakdown by role and agent, top sessions by message count.", - "inputSchema": { - "type": "object", - "properties": {} - } - }), - json!({ - "name": "icm_wake_up", - "description": "Build a compact critical-facts pack for LLM system-prompt injection. Selects critical/high memories (and preferences) optionally scoped by project, ranks by importance × recency × weight, and truncates to a token budget. Use at session start to hydrate an agent with the most load-bearing context.", - "inputSchema": { - "type": "object", - "properties": { - "project": { - "type": "string", - "description": "Project name filter (substring match against topic). Preferences/identity memories are always included." + }), + ToolAnnotations::new(false, false, false, false), + |context, args| tool_transcript_start_session(context.store, args) + ), + tool_spec!( + TranscriptRecordInput, + json!({ + "name": "icm_transcript_record", + "description": "Append a verbatim message to a transcript session. Stores the raw content with no summarization. Use once per user turn, assistant reply, or tool call for full replay fidelity.", + "inputSchema": { + "type": "object", + "properties": { + "session_id": { + "type": "string", + "description": "Session id from icm_transcript_start_session" + }, + "role": { + "type": "string", + "enum": ["user", "assistant", "system", "tool"], + "description": "Message role" + }, + "content": { + "type": "string", + "description": "Raw message content (stored verbatim)" + }, + "tool_name": { + "type": "string", + "description": "Tool name if role=tool (optional)" + }, + "tokens": { + "type": "integer", + "description": "Token count for billing / stats (optional)" + }, + "metadata": { + "type": "string", + "description": "Arbitrary JSON metadata (optional)" + } }, - "max_tokens": { - "type": "integer", - "default": 200, - "minimum": 20, - "maximum": 4000, - "description": "Approximate token budget (1 token ≈ 4 characters)" + "required": ["session_id", "role", "content"] + } + }), + ToolAnnotations::new(false, false, false, false), + |context, args| tool_transcript_record(context.store, args) + ), + tool_spec!( + TranscriptSearchInput, + json!({ + "name": "icm_transcript_search", + "description": "Full-text search across recorded transcript messages (FTS5 BM25). Supports boolean operators, phrase matches, and prefix queries. Use to recall exact quotes or debug past decisions.", + "inputSchema": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "FTS5 query: 'postgres OR mysql', '\"exact phrase\"', 'auth*'" + }, + "session_id": { + "type": "string", + "description": "Restrict to one session (optional)" + }, + "project": { + "type": "string", + "description": "Restrict to one project (optional)" + }, + "limit": { + "type": "integer", + "default": 10, + "minimum": 1, + "maximum": 50 + } }, - "format": { - "type": "string", - "enum": ["markdown", "plain"], - "default": "markdown", - "description": "Output format" + "required": ["query"] + } + }), + ToolAnnotations::new(true, false, true, false), + |context, args| tool_transcript_search(context.store, args) + ), + tool_spec!( + TranscriptShowInput, + json!({ + "name": "icm_transcript_show", + "description": "Replay the full message thread of a transcript session, chronologically. Returns up to `limit` messages with role, content, tool name, timestamp.", + "inputSchema": { + "type": "object", + "properties": { + "session_id": { "type": "string" }, + "limit": { "type": "integer", "default": 200, "minimum": 1, "maximum": 2000 } }, - "include_preferences": { - "type": "boolean", - "default": true, - "description": "Include global preferences/identity memories regardless of the project filter" + "required": ["session_id"] + } + }), + ToolAnnotations::new(true, false, true, false), + |context, args| tool_transcript_show(context.store, args) + ), + tool_spec!( + TranscriptStatsInput, + json!({ + "name": "icm_transcript_stats", + "description": "Global transcript statistics: session count, message count, total bytes, breakdown by role and agent, top sessions by message count.", + "inputSchema": { + "type": "object", + "properties": {} + } + }), + ToolAnnotations::new(true, false, true, false), + |context, _| tool_transcript_stats(context.store) + ), + tool_spec!( + WakeUpInput, + json!({ + "name": "icm_wake_up", + "description": "Build a compact critical-facts pack for LLM system-prompt injection. Selects critical/high memories (and preferences) optionally scoped by project, ranks by importance × recency × weight, and truncates to a token budget. Use at session start to hydrate an agent with the most load-bearing context.", + "inputSchema": { + "type": "object", + "properties": { + "project": { + "type": "string", + "description": "Project name filter (substring match against topic). Preferences/identity memories are always included." + }, + "max_tokens": { + "type": "integer", + "default": 200, + "minimum": 20, + "maximum": 4000, + "description": "Approximate token budget (1 token ≈ 4 characters)" + }, + "format": { + "type": "string", + "enum": ["markdown", "plain"], + "default": "markdown", + "description": "Output format" + }, + "include_preferences": { + "type": "boolean", + "default": true, + "description": "Include global preferences/identity memories regardless of the project filter" + } } } - } - }), - ]; - - if has_embedder { - tools.push(json!({ - "name": "icm_memory_embed_all", - "description": "Generate embeddings for all memories that don't have one yet. Use this to backfill vector search capability.", - "inputSchema": { - "type": "object", - "properties": { - "topic": { - "type": "string", - "description": "Only embed memories in this topic (optional)" + }), + ToolAnnotations::new(true, false, true, false), + |context, args| tool_wake_up(context.store, args) + ), + tool_spec!( + EmbedAllInput, + json!({ + "name": "icm_memory_embed_all", + "description": "Generate embeddings for all memories that don't have one yet. Use this to backfill vector search capability.", + "inputSchema": { + "type": "object", + "properties": { + "topic": { + "type": "string", + "description": "Only embed memories in this topic (optional)" + } } } - } - })); - } + }), + requirements: ToolRequirements::STORE.with_required_embedder(), + ToolAnnotations::new(false, false, true, false), + |context, args| tool_embed_all(context.store, context.embedder, args) + ), + ]; - json!({ "tools": tools }) + ToolCatalog::new(tools, has_embedder).expect("static MCP tool registrations must be valid") } // --------------------------------------------------------------------------- // Tool dispatch // --------------------------------------------------------------------------- +/// Frozen 2024 tool-list projection retained for callers and compatibility +/// tests. Production service instances cache this projection in their catalog. +pub fn tool_definitions(has_embedder: bool) -> Value { + build_catalog(has_embedder).legacy_list() +} + pub fn call_tool( store: &Store, embedder: Option<&dyn Embedder>, @@ -781,45 +1068,32 @@ pub fn call_tool_with_config( compact: bool, auto_consolidate: AutoConsolidate, ) -> ToolResult { - match name { - // Memory tools - "icm_memory_store" => tool_store(store, embedder, args, compact, auto_consolidate), - "icm_memory_recall" => tool_recall(store, embedder, args, compact), - "icm_memory_forget" => tool_forget(store, args), - "icm_memory_forget_topic" => tool_forget_topic(store, args), - "icm_memory_update" => tool_update(store, embedder, args), - "icm_memory_consolidate" => tool_consolidate(store, embedder, args), - "icm_memory_list_topics" => tool_list_topics(store), - "icm_memory_stats" => tool_stats(store), - "icm_memory_health" => tool_health(store, args), - "icm_memory_extract_patterns" => tool_extract_patterns(store, args), - "icm_memory_embed_all" => tool_embed_all(store, embedder, args), - // Memoir tools - "icm_memoir_create" => tool_memoir_create(store, args), - "icm_memoir_list" => tool_memoir_list(store), - "icm_memoir_show" => tool_memoir_show(store, args), - "icm_memoir_add_concept" => tool_memoir_add_concept(store, args), - "icm_memoir_refine" => tool_memoir_refine(store, args), - "icm_memoir_search" => tool_memoir_search(store, args), - "icm_memoir_search_all" => tool_memoir_search_all(store, args), - "icm_memoir_link" => tool_memoir_link(store, args), - "icm_memoir_inspect" => tool_memoir_inspect(store, args), - "icm_memoir_export" => tool_memoir_export(store, args), - // Learn tool - "icm_learn" => tool_learn(store, args), - // Feedback tools - "icm_feedback_record" => tool_feedback_record(store, embedder, args, compact), - "icm_feedback_search" => tool_feedback_search(store, embedder, args), - "icm_feedback_stats" => tool_feedback_stats(store), - // Transcript tools - "icm_transcript_start_session" => tool_transcript_start_session(store, args), - "icm_transcript_record" => tool_transcript_record(store, args), - "icm_transcript_search" => tool_transcript_search(store, args), - "icm_transcript_show" => tool_transcript_show(store, args), - "icm_transcript_stats" => tool_transcript_stats(store), - // Wake-up tool - "icm_wake_up" => tool_wake_up(store, args), - _ => ToolResult::error(format!("unknown tool: {name}")), + // This public helper is the frozen pre-catalog compatibility dispatcher. + // Keep unavailable tools dispatchable here so their established handler + // errors remain stable; production service discovery and dispatch use the + // capability-filtered catalog stored by `McpService`. + let catalog = build_catalog(true); + let working_directory = + std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from(".")); + let context = ToolContext { + store, + embedder, + compact, + auto_consolidate, + working_directory: &working_directory, + enforce_directory_boundary: false, + }; + match catalog.dispatch( + &context, + name, + args, + crate::catalog::InputValidation::Legacy2024Unchecked, + ) { + DispatchResult::ToolResult(result) => result, + DispatchResult::UnknownTool => ToolResult::error(format!("unknown tool: {name}")), + DispatchResult::InvalidInput(_) => { + unreachable!("unchecked legacy compatibility dispatch cannot reject typed inputs") + } } } @@ -1012,13 +1286,13 @@ fn tool_store( // Input length validation if topic.len() > MAX_TOPIC_LEN { return ToolResult::error(format!( - "topic exceeds maximum length ({} > {MAX_TOPIC_LEN} chars)", + "topic exceeds maximum length ({} > {MAX_TOPIC_LEN} UTF-8 bytes)", topic.len() )); } if content.len() > MAX_CONTENT_LEN { return ToolResult::error(format!( - "content exceeds maximum length ({} > {MAX_CONTENT_LEN} chars)", + "content exceeds maximum length ({} > {MAX_CONTENT_LEN} UTF-8 bytes)", content.len() )); } @@ -1255,9 +1529,10 @@ fn tool_recall( Some(q) => q, None => return ToolResult::error("missing required field: query".into()), }; - // Clamp to the schema's advertised maximum (20) — the code previously - // accepted up to 100, silently diverging from the published contract. - let limit = get_i64(args, "limit", 5).clamp(1, 20) as usize; + // The modern input contract extends the historical advertised maximum + // from 20 to the frozen Phase 2 boundary of 100. Keep the handler cap in + // lockstep so valid modern calls are not silently truncated. + let limit = get_i64(args, "limit", 5).clamp(1, 100) as usize; let topic = get_str(args, "topic"); let keyword = get_str(args, "keyword"); @@ -1454,6 +1729,41 @@ fn tool_learn(store: &Store, args: &Value) -> ToolResult { } } +fn tool_learn_bounded(context: &ToolContext<'_>, args: &Value) -> ToolResult { + if !context.enforce_directory_boundary { + return tool_learn(context.store, args); + } + let requested = get_str(args, "directory").unwrap_or("."); + let requested = std::path::Path::new(requested); + let candidate = if requested.is_absolute() { + requested.to_path_buf() + } else { + context.working_directory.join(requested) + }; + let root = match context.working_directory.canonicalize() { + Ok(root) => root, + Err(_) => { + return ToolResult::error( + "server working directory could not be resolved safely".into(), + ) + } + }; + let candidate = match candidate.canonicalize() { + Ok(candidate) => candidate, + Err(_) => { + return ToolResult::error(format!("directory not found: {}", candidate.display())) + } + }; + if !candidate.starts_with(&root) { + return ToolResult::error( + "directory must remain within the server working directory".into(), + ); + } + let mut bounded_args = args.clone(); + bounded_args["directory"] = Value::String(candidate.to_string_lossy().into_owned()); + tool_learn(context.store, &bounded_args) +} + fn tool_consolidate(store: &Store, embedder: Option<&dyn Embedder>, args: &Value) -> ToolResult { let topic = match get_str(args, "topic") { Some(t) => t, @@ -1779,12 +2089,15 @@ fn tool_memoir_create(store: &Store, args: &Value) -> ToolResult { None => return ToolResult::error("missing required field: name".into()), }; if name.len() > 255 { - return ToolResult::error(format!("name too long: {} chars (max 255)", name.len())); + return ToolResult::error(format!( + "name too long: {} UTF-8 bytes (max 255)", + name.len() + )); } let description = get_str(args, "description").unwrap_or(""); if description.len() > 10_000 { return ToolResult::error(format!( - "description too long: {} chars (max 10000)", + "description too long: {} UTF-8 bytes (max 10000)", description.len() )); } @@ -1886,7 +2199,7 @@ fn tool_memoir_add_concept(store: &Store, args: &Value) -> ToolResult { }; if name.len() > 255 { return ToolResult::error(format!( - "concept name too long: {} chars (max 255)", + "concept name too long: {} UTF-8 bytes (max 255)", name.len() )); } @@ -1896,7 +2209,7 @@ fn tool_memoir_add_concept(store: &Store, args: &Value) -> ToolResult { }; if definition.len() > 10_000 { return ToolResult::error(format!( - "definition too long: {} chars (max 10000)", + "definition too long: {} UTF-8 bytes (max 10000)", definition.len() )); } @@ -1932,13 +2245,19 @@ fn tool_memoir_refine(store: &Store, args: &Value) -> ToolResult { Some(n) => n, None => return ToolResult::error("missing required field: name".into()), }; + if name.len() > 255 { + return ToolResult::error(format!( + "concept name too long: {} UTF-8 bytes (max 255)", + name.len() + )); + } let definition = match get_str(args, "definition") { Some(d) => d, None => return ToolResult::error("missing required field: definition".into()), }; if definition.len() > 10_000 { return ToolResult::error(format!( - "definition too long: {} chars (max 10000)", + "definition too long: {} UTF-8 bytes (max 10000)", definition.len() )); } @@ -2424,7 +2743,7 @@ fn tool_feedback_record( ] { if field_value.len() > MAX_FEEDBACK_FIELD_LEN { return ToolResult::error(format!( - "{field_name} exceeds maximum length ({} > {MAX_FEEDBACK_FIELD_LEN} chars)", + "{field_name} exceeds maximum length ({} > {MAX_FEEDBACK_FIELD_LEN} UTF-8 bytes)", field_value.len() )); } @@ -2665,6 +2984,21 @@ mod tests { assert!(result.content[0].text.contains("unknown tool")); } + #[test] + fn embed_all_without_embedder_preserves_legacy_handler_error() { + let store = test_store(); + let listed = tool_definitions(false); + assert!(!listed["tools"] + .as_array() + .unwrap() + .iter() + .any(|tool| tool["name"] == "icm_memory_embed_all")); + + let result = call_tool(&store, None, "icm_memory_embed_all", &json!({}), false); + assert!(result.is_error); + assert_eq!(result.content[0].text, "embeddings not available"); + } + #[test] fn test_store_missing_topic() { let store = test_store(); @@ -2795,37 +3129,41 @@ mod tests { ); } - /// Audit regression: the schema advertises limit <= 20 but the code - /// accepted 100 — the clamp must match the published contract. + /// The public helper remains the frozen unchecked 2024 compatibility + /// dispatch, including its historical lower and upper recall clamps. #[test] - fn test_recall_limit_clamped_to_schema_max() { + fn test_legacy_call_tool_recall_limits_remain_unchecked_and_clamped() { let store = test_store(); + let consolidation_off = AutoConsolidate { + enabled: false, + threshold: 10, + }; for i in 0..30 { - let r = call_tool( + let r = call_tool_with_config( &store, None, "icm_memory_store", &json!({"topic": "t", "content": format!("clamp probe entry number {i}")}), false, + consolidation_off, ); assert!(!r.is_error); } - let recall_result = call_tool( - &store, - None, - "icm_memory_recall", - &json!({"query": "clamp probe entry", "project": "", "limit": 100}), - false, - ); - assert!(!recall_result.is_error); - let hits = recall_result.content[0] - .text - .matches("clamp probe entry") - .count(); - assert!( - hits <= 20, - "limit must clamp to the schema max of 20, got {hits} hits" - ); + for (limit, expected_hits) in [(0, 1), (100, 20), (101, 20)] { + let recall_result = call_tool( + &store, + None, + "icm_memory_recall", + &json!({"query": "clamp probe entry", "project": "", "limit": limit}), + false, + ); + assert!(!recall_result.is_error); + let hits = recall_result.content[0] + .text + .matches("clamp probe entry") + .count(); + assert_eq!(hits, expected_hits); + } } /// Audit regression: filtering was previously applied AFTER the store From 2e2493dcc1a14ede69c9019099a550cbaf7082a1 Mon Sep 17 00:00:00 2001 From: Lauritz-Timm Date: Thu, 6 Aug 2026 21:28:05 +0200 Subject: [PATCH 02/20] feat(mcp): model supported protocol revisions --- crates/icm-mcp/src/protocol.rs | 67 +++++++++++++++++++++++++++++++++- 1 file changed, 65 insertions(+), 2 deletions(-) diff --git a/crates/icm-mcp/src/protocol.rs b/crates/icm-mcp/src/protocol.rs index 0235334d..03d32e0a 100644 --- a/crates/icm-mcp/src/protocol.rs +++ b/crates/icm-mcp/src/protocol.rs @@ -1,5 +1,56 @@ use serde::{Deserialize, Serialize}; -use serde_json::Value; +use serde_json::{Map, Value}; + +pub const SUPPORTED_PROTOCOL_VERSIONS: [&str; 4] = [ + ProtocolRevision::V2026_07_28.as_str(), + ProtocolRevision::V2025_11_25.as_str(), + ProtocolRevision::V2025_06_18.as_str(), + ProtocolRevision::V2024_11_05.as_str(), +]; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ProtocolEra { + InitializationBased, + PerRequest, +} + +impl ProtocolEra { + pub const fn as_str(self) -> &'static str { + match self { + Self::InitializationBased => "initialization-based", + Self::PerRequest => "per-request", + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ProtocolRevision { + V2024_11_05, + V2025_06_18, + V2025_11_25, + V2026_07_28, +} + +impl ProtocolRevision { + pub const fn parse_exact(value: &str) -> Option { + match value.as_bytes() { + b"2024-11-05" => Some(Self::V2024_11_05), + b"2025-06-18" => Some(Self::V2025_06_18), + b"2025-11-25" => Some(Self::V2025_11_25), + b"2026-07-28" => Some(Self::V2026_07_28), + _ => None, + } + } + + pub const fn as_str(self) -> &'static str { + match self { + Self::V2024_11_05 => "2024-11-05", + Self::V2025_06_18 => "2025-06-18", + Self::V2025_11_25 => "2025-11-25", + Self::V2026_07_28 => "2026-07-28", + } + } +} // --------------------------------------------------------------------------- // JSON-RPC 2.0 message types @@ -23,6 +74,8 @@ pub struct JsonRpcMessage { pub method: Option, #[serde(default)] pub params: Option, + #[serde(flatten)] + pub extra: Map, } fn deserialize_some<'de, D>(deserializer: D) -> Result, D::Error> @@ -46,6 +99,8 @@ pub struct JsonRpcResponse { pub struct JsonRpcError { pub code: i64, pub message: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub data: Option, } impl JsonRpcResponse { @@ -59,11 +114,19 @@ impl JsonRpcResponse { } pub fn err(id: Value, code: i64, message: String) -> Self { + Self::err_with_data(id, code, message, None) + } + + pub fn err_with_data(id: Value, code: i64, message: String, data: Option) -> Self { Self { jsonrpc: "2.0".into(), id, result: None, - error: Some(JsonRpcError { code, message }), + error: Some(JsonRpcError { + code, + message, + data, + }), } } From a0ddfd196db84013b60b14b4ada5efa00802ed23 Mon Sep 17 00:00:00 2001 From: Lauritz-Timm Date: Fri, 7 Aug 2026 22:11:26 +0200 Subject: [PATCH 03/20] feat(mcp): add revision-aware core service --- crates/icm-mcp/src/lib.rs | 1 + crates/icm-mcp/src/server.rs | 343 ++-- crates/icm-mcp/src/service.rs | 3111 +++++++++++++++++++++++++++++++++ crates/icm-mcp/src/tools.rs | 17 +- 4 files changed, 3285 insertions(+), 187 deletions(-) create mode 100644 crates/icm-mcp/src/service.rs diff --git a/crates/icm-mcp/src/lib.rs b/crates/icm-mcp/src/lib.rs index 0b01c441..e0fb297f 100644 --- a/crates/icm-mcp/src/lib.rs +++ b/crates/icm-mcp/src/lib.rs @@ -2,6 +2,7 @@ pub mod catalog; mod inputs; pub mod protocol; pub mod server; +pub mod service; pub mod tools; pub use server::run_server; diff --git a/crates/icm-mcp/src/server.rs b/crates/icm-mcp/src/server.rs index 34823402..016388f9 100644 --- a/crates/icm-mcp/src/server.rs +++ b/crates/icm-mcp/src/server.rs @@ -1,49 +1,34 @@ -use std::io::{self, BufRead, Read, Write}; +//! Bounded stdio framing for the transport-neutral MCP service. -use serde_json::{json, Value}; -use tracing::{debug, error}; +use std::io::{self, BufRead, Read, Write}; use icm_core::Embedder; use icm_store::Store; +use serde_json::Value; +use tracing::error; use crate::protocol::{JsonRpcMessage, JsonRpcResponse}; -use crate::tools::{self, AutoConsolidate}; - -const SERVER_NAME: &str = "icm"; -const SERVER_VERSION: &str = env!("CARGO_PKG_VERSION"); -const PROTOCOL_VERSION: &str = "2024-11-05"; +use crate::service::{ConnectionState, McpService}; +use crate::tools::AutoConsolidate; -/// Number of non-store tool calls before we nudge the agent to store. -const STORE_NUDGE_THRESHOLD: u32 = 10; +/// Maximum allowed line length (10 MiB). The cap is applied while reading, +/// before the complete caller-controlled frame can be allocated. +pub const MAX_LINE_LEN: usize = 10 * 1024 * 1024; -/// Maximum allowed line length (10 MB). The cap is enforced *while reading* -/// (bounded `take` + `read_until`), so an oversized line is never fully -/// buffered — previously the whole line was allocated by `lines()` before -/// the length check ran, defeating the cap (audit finding; same class of -/// bug as the CLI hook-stdin fix in e551c27). -const MAX_LINE_LEN: usize = 10 * 1024 * 1024; - -/// Read one `\n`-terminated line into `buf` without ever buffering more than -/// `MAX_LINE_LEN + 1` bytes of it. Returns `Ok(None)` on EOF, `Ok(Some(true))` -/// for a within-limit line, `Ok(Some(false))` for an oversized line (whose -/// remainder has been drained and discarded in bounded chunks). -fn read_capped_line(reader: &mut impl BufRead, buf: &mut Vec) -> io::Result> { - buf.clear(); - let n = reader +fn read_capped_line(reader: &mut impl BufRead, buffer: &mut Vec) -> io::Result> { + buffer.clear(); + let bytes_read = reader .take(MAX_LINE_LEN as u64 + 1) - .read_until(b'\n', buf)?; - if n == 0 { - return Ok(None); // EOF + .read_until(b'\n', buffer)?; + if bytes_read == 0 { + return Ok(None); } - // Oversized iff we exhausted the read budget without hitting the newline. - if buf.last() != Some(&b'\n') && n == MAX_LINE_LEN + 1 { - // Drain the rest of the line in bounded chunks so the next read - // starts on a fresh line. + if buffer.last() != Some(&b'\n') && bytes_read == MAX_LINE_LEN + 1 { let mut scratch = Vec::with_capacity(64 * 1024); loop { scratch.clear(); - let m = reader.take(1024 * 1024).read_until(b'\n', &mut scratch)?; - if m == 0 || scratch.last() == Some(&b'\n') { + let drained = reader.take(1024 * 1024).read_until(b'\n', &mut scratch)?; + if drained == 0 || scratch.last() == Some(&b'\n') { break; } } @@ -52,7 +37,7 @@ fn read_capped_line(reader: &mut impl BufRead, buf: &mut Vec) -> io::Result< Ok(Some(true)) } -/// Run the MCP server on stdio. Blocks until stdin is closed. +/// Run the MCP server on stdio until stdin closes. pub fn run_server( store: &Store, embedder: Option<&dyn Embedder>, @@ -61,175 +46,179 @@ pub fn run_server( ) -> anyhow::Result<()> { let stdin = io::stdin(); let mut reader = stdin.lock(); - let mut stdout = io::stdout(); - let mut calls_since_store: u32 = 0; - let mut buf: Vec = Vec::new(); + let stdout = io::stdout(); + let mut writer = stdout.lock(); + run_server_with_io( + store, + embedder, + compact, + auto_consolidate, + &mut reader, + &mut writer, + ) +} + +/// Generic framing adapter used by stdio and hermetic transport tests. +pub fn run_server_with_io( + store: &Store, + embedder: Option<&dyn Embedder>, + compact: bool, + auto_consolidate: AutoConsolidate, + reader: &mut impl BufRead, + writer: &mut impl Write, +) -> anyhow::Result<()> { + let service = McpService::new(store, embedder, compact, auto_consolidate); + let mut state = ConnectionState::default(); + let mut buffer = Vec::new(); loop { - let within_limit = match read_capped_line(&mut reader, &mut buf) { - Ok(Some(ok)) => ok, - Ok(None) => break, // EOF - Err(e) => { - error!("stdin read error: {e}"); + let within_limit = match read_capped_line(reader, &mut buffer) { + Ok(Some(within_limit)) => within_limit, + Ok(None) => break, + Err(read_error) => { + error!("stdin read error: {read_error}"); break; } }; - if !within_limit { error!("line too long (max {MAX_LINE_LEN} bytes)"); - let resp = JsonRpcResponse::err( - Value::Null, - -32600, - format!("line too long (max {MAX_LINE_LEN} bytes)"), - ); - write_response(&mut stdout, &resp)?; + write_response( + writer, + &JsonRpcResponse::err( + Value::Null, + -32600, + format!("line too long (max {MAX_LINE_LEN} bytes)"), + ), + )?; continue; } - let line_owned = String::from_utf8_lossy(&buf); - let line = line_owned.trim(); - if line.is_empty() { + if buffer.last() == Some(&b'\n') { + buffer.pop(); + if buffer.last() == Some(&b'\r') { + buffer.pop(); + } + } + if buffer.iter().all(u8::is_ascii_whitespace) { continue; } - - let msg: JsonRpcMessage = match serde_json::from_str(line) { - Ok(m) => m, - Err(e) => { - error!("invalid JSON-RPC: {e}"); - // Send parse error if we can - let resp = JsonRpcResponse::err(Value::Null, -32700, format!("parse error: {e}")); - write_response(&mut stdout, &resp)?; + let line = match std::str::from_utf8(&buffer) { + Ok(line) => line, + Err(parse_error) => { + write_response( + writer, + &JsonRpcResponse::err( + Value::Null, + -32700, + format!("parse error: {parse_error}"), + ), + )?; continue; } }; - - let method = msg.method.as_deref().unwrap_or(""); - debug!("MCP request: {method}"); - - // Notifications have no id — don't respond - let id = match msg.id { - Some(id) => id, - None => continue, + let wire_value: Value = match serde_json::from_str(line) { + Ok(value) => value, + Err(parse_error) => { + error!("invalid JSON-RPC: {parse_error}"); + write_response( + writer, + &JsonRpcResponse::err( + Value::Null, + -32700, + format!("parse error: {parse_error}"), + ), + )?; + continue; + } }; - - let response = match method { - "initialize" => handle_initialize(id), - "ping" => JsonRpcResponse::ok(id, json!({})), - "tools/list" => handle_tools_list(id, embedder.is_some()), - "tools/call" => handle_tools_call( - id, - &msg.params, - store, - embedder, - compact, - auto_consolidate, - &mut calls_since_store, - ), - other => JsonRpcResponse::method_not_found(id, other), + let message: JsonRpcMessage = match serde_json::from_value(wire_value) { + Ok(message) => message, + Err(request_error) => { + write_response( + writer, + &JsonRpcResponse::err( + Value::Null, + -32600, + format!("invalid request: {request_error}"), + ), + )?; + continue; + } }; - write_response(&mut stdout, &response)?; + if let Some(response) = service.handle(&mut state, message) { + write_response(writer, &response)?; + } } - Ok(()) } -fn write_response(stdout: &mut io::Stdout, resp: &JsonRpcResponse) -> anyhow::Result<()> { - let json = serde_json::to_string(resp)?; - writeln!(stdout, "{json}")?; - stdout.flush()?; +fn write_response(writer: &mut impl Write, response: &JsonRpcResponse) -> anyhow::Result<()> { + serde_json::to_writer(&mut *writer, response)?; + writer.write_all(b"\n")?; + writer.flush()?; Ok(()) } -fn handle_initialize(id: Value) -> JsonRpcResponse { - JsonRpcResponse::ok( - id, - json!({ - "protocolVersion": PROTOCOL_VERSION, - "capabilities": { - "tools": {} - }, - "serverInfo": { - "name": SERVER_NAME, - "version": SERVER_VERSION - }, - "instructions": ICM_INSTRUCTIONS - }), - ) -} - -const ICM_INSTRUCTIONS: &str = "\ -Use ICM (Infinite Context Memory) proactively to maintain long-term memory across sessions.\n\ -\n\ -RECALL (icm_memory_recall): At the start of a task, search for relevant past context — decisions, \ -resolved errors, user preferences. Search only what is relevant, do not dump everything.\n\ -\n\ -STORE (icm_memory_store): You MUST store when ANY of these triggers occur:\n\ -1. Error resolved → topic: \"errors-resolved\", importance: high\n\ -2. Architecture/design decision made → topic: \"decisions-{project}\", importance: high\n\ -3. User preference discovered (correction, feedback) → topic: \"preferences\", importance: critical\n\ -4. Significant task completed (feature, fix, config, review) → topic: \"context-{project}\", importance: high\n\ -5. Conversation exceeds ~20 tool calls without a store → store a progress summary\n\ -\n\ -Do this BEFORE responding to the user. Not after. Not later. Immediately.\n\ -\n\ -Do NOT store: trivial details, information already in CLAUDE.md, ephemeral state.\n\ -\n\ -Importance levels: critical (never forgotten), high (slow decay), medium (normal), low (fast decay)."; - -fn handle_tools_list(id: Value, has_embedder: bool) -> JsonRpcResponse { - JsonRpcResponse::ok(id, tools::tool_definitions(has_embedder)) -} - -fn handle_tools_call( - id: Value, - params: &Option, - store: &Store, - embedder: Option<&dyn Embedder>, - compact: bool, - auto_consolidate: AutoConsolidate, - calls_since_store: &mut u32, -) -> JsonRpcResponse { - let params = match params { - Some(p) => p, - None => { - return JsonRpcResponse::err(id, -32602, "missing params".into()); - } - }; - - let tool_name = match params.get("name").and_then(|v| v.as_str()) { - Some(n) => n, - None => { - return JsonRpcResponse::err(id, -32602, "missing tool name".into()); - } - }; - - let args = params.get("arguments").cloned().unwrap_or(json!({})); - - // Track store calls to nudge the agent - if tool_name == "icm_memory_store" { - *calls_since_store = 0; - } else { - *calls_since_store += 1; +#[cfg(test)] +mod tests { + use std::io::Cursor; + + use super::*; + use serde_json::json; + + #[test] + fn in_memory_transport_runs_complete_2024_sequence() { + let store = Store::in_memory().unwrap(); + let input = [ + json!({ + "jsonrpc":"2.0","id":1,"method":"initialize", + "params":{ + "protocolVersion":"2024-11-05","capabilities":{}, + "clientInfo":{"name":"test","version":"1"} + } + }), + json!({"jsonrpc":"2.0","method":"notifications/initialized","params":{}}), + json!({"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}), + ] + .into_iter() + .map(|request| format!("{}\n", serde_json::to_string(&request).unwrap())) + .collect::(); + let mut reader = Cursor::new(input.into_bytes()); + let mut output = Vec::new(); + run_server_with_io( + &store, + None, + false, + AutoConsolidate::default(), + &mut reader, + &mut output, + ) + .unwrap(); + let responses: Vec = String::from_utf8(output) + .unwrap() + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect(); + assert_eq!(responses.len(), 2); + assert_eq!(responses[0]["result"]["protocolVersion"], "2024-11-05"); + assert!(responses[1]["result"]["tools"].is_array()); } - let mut result = - tools::call_tool_with_config(store, embedder, tool_name, &args, compact, auto_consolidate); - - // Nudge: remind the agent to store on every THRESHOLD-th call without a - // store (10, 20, 30, …) — previously the hint was appended to *every* - // response past the threshold, a recurring token tax on the client LLM - // (audit finding). - if tool_name != "icm_memory_store" - && *calls_since_store >= STORE_NUDGE_THRESHOLD - && calls_since_store.is_multiple_of(STORE_NUDGE_THRESHOLD) - { - result.append_hint(&format!( - "\n[ICM: {} tool calls since last store. \ - Consider saving important context with icm_memory_store before it is lost.]", - calls_since_store - )); + #[test] + fn oversized_frame_is_drained_before_the_next_request() { + let mut input = vec![b'x'; MAX_LINE_LEN + 1]; + input.extend_from_slice(b"\n{}"); + let mut reader = Cursor::new(input); + let mut buffer = Vec::new(); + assert_eq!( + read_capped_line(&mut reader, &mut buffer).unwrap(), + Some(false) + ); + assert_eq!( + read_capped_line(&mut reader, &mut buffer).unwrap(), + Some(true) + ); + assert_eq!(buffer, b"{}"); } - - JsonRpcResponse::ok(id, serde_json::to_value(result).unwrap_or(json!(null))) } diff --git a/crates/icm-mcp/src/service.rs b/crates/icm-mcp/src/service.rs new file mode 100644 index 00000000..d5a9c376 --- /dev/null +++ b/crates/icm-mcp/src/service.rs @@ -0,0 +1,3111 @@ +//! Transport-neutral MCP request service. + +use std::collections::HashSet; +use std::path::PathBuf; + +use icm_core::Embedder; +use icm_store::Store; +use serde_json::{json, Map, Value}; + +use crate::catalog::{DispatchResult, InputValidation, ToolCatalog, ToolContext}; +use crate::protocol::{ + JsonRpcMessage, JsonRpcResponse, ProtocolEra, ProtocolRevision, SUPPORTED_PROTOCOL_VERSIONS, +}; +use crate::tools::{self, AutoConsolidate}; + +const SERVER_NAME: &str = "icm"; +const SERVER_VERSION: &str = env!("CARGO_PKG_VERSION"); +const STORE_NUDGE_THRESHOLD: u32 = 10; +const MODERN_PROTOCOL_VERSION_KEY: &str = "io.modelcontextprotocol/protocolVersion"; +const MODERN_CLIENT_CAPABILITIES_KEY: &str = "io.modelcontextprotocol/clientCapabilities"; +const MODERN_CLIENT_INFO_KEY: &str = "io.modelcontextprotocol/clientInfo"; +const MODERN_SERVER_INFO_KEY: &str = "io.modelcontextprotocol/serverInfo"; +const MODERN_LOG_LEVEL_KEY: &str = "io.modelcontextprotocol/logLevel"; +const MODERN_SUBSCRIPTION_ID_KEY: &str = "io.modelcontextprotocol/subscriptionId"; +const MAX_STORED_LIFECYCLE_METHOD_BYTES: usize = 256; + +pub const ICM_INSTRUCTIONS: &str = "\ +Use ICM (Infinite Context Memory) proactively to maintain long-term memory across sessions.\n\ +\n\ +RECALL (icm_memory_recall): At the start of a task, search for relevant past context — decisions, \ +resolved errors, user preferences. Search only what is relevant, do not dump everything.\n\ +\n\ +STORE (icm_memory_store): You MUST store when ANY of these triggers occur:\n\ +1. Error resolved → topic: \"errors-resolved\", importance: high\n\ +2. Architecture/design decision made → topic: \"decisions-{project}\", importance: high\n\ +3. User preference discovered (correction, feedback) → topic: \"preferences\", importance: critical\n\ +4. Significant task completed (feature, fix, config, review) → topic: \"context-{project}\", importance: high\n\ +5. Conversation exceeds ~20 tool calls without a store → store a progress summary\n\ +\n\ +Do this BEFORE responding to the user. Not after. Not later. Immediately.\n\ +\n\ +Do NOT store: trivial details, information already in CLAUDE.md, ephemeral state.\n\ +\n\ +Importance levels: critical (never forgotten), high (slow decay), medium (normal), low (fast decay)."; + +#[derive(Clone, Debug)] +struct LifecycleViolation { + kind: &'static str, + state: &'static str, + method: String, +} + +#[derive(Clone, Debug)] +enum ConnectionPhase { + Uninitialized, + /// The 2024 compatibility projection is ready immediately after + /// initialize because the frozen legacy clients do not send the + /// initialized notification. One optional notification is still accepted. + LegacyReady { + revision: ProtocolRevision, + initialized_seen: bool, + }, + LegacyAwaitingInitialized(ProtocolRevision), + Modern, + Poisoned(LifecycleViolation), +} + +#[derive(Clone, Debug)] +pub struct ConnectionState { + phase: ConnectionPhase, + calls_since_store: u32, +} + +impl Default for ConnectionState { + fn default() -> Self { + Self { + phase: ConnectionPhase::Uninitialized, + calls_since_store: 0, + } + } +} + +pub struct McpService<'a> { + store: &'a Store, + embedder: Option<&'a dyn Embedder>, + compact: bool, + auto_consolidate: AutoConsolidate, + working_directory: PathBuf, + catalog: ToolCatalog, +} + +impl<'a> McpService<'a> { + pub fn new( + store: &'a Store, + embedder: Option<&'a dyn Embedder>, + compact: bool, + auto_consolidate: AutoConsolidate, + ) -> Self { + let working_directory = std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")); + Self::with_working_directory( + store, + embedder, + compact, + auto_consolidate, + working_directory, + ) + } + + pub fn with_working_directory( + store: &'a Store, + embedder: Option<&'a dyn Embedder>, + compact: bool, + auto_consolidate: AutoConsolidate, + working_directory: PathBuf, + ) -> Self { + Self { + store, + embedder, + compact, + auto_consolidate, + working_directory, + catalog: tools::build_catalog(embedder.is_some()), + } + } + + pub fn handle( + &self, + state: &mut ConnectionState, + message: JsonRpcMessage, + ) -> Option { + let response_id = message.id.clone().unwrap_or(Value::Null); + if message.jsonrpc != "2.0" { + return Some(JsonRpcResponse::err( + response_id, + -32600, + "invalid JSON-RPC version; expected 2.0".into(), + )); + } + let Some(method) = message + .method + .as_deref() + .filter(|method| !method.is_empty()) + else { + return Some(JsonRpcResponse::err( + response_id, + -32600, + "invalid request: method must be a non-empty string".into(), + )); + }; + + if message.id.is_none() { + self.handle_notification(state, method, &message); + return None; + } + + if !valid_request_id(message.id.as_ref().expect("checked above")) + && !matches!( + state.phase, + ConnectionPhase::LegacyReady { + revision: ProtocolRevision::V2024_11_05, + .. + } + ) + { + return Some(JsonRpcResponse::err( + Value::Null, + -32600, + "invalid request id; expected a string or integer".into(), + )); + } + + Some(self.handle_request(state, response_id, method, &message)) + } + + fn handle_request( + &self, + state: &mut ConnectionState, + id: Value, + method: &str, + message: &JsonRpcMessage, + ) -> JsonRpcResponse { + if let ConnectionPhase::Poisoned(violation) = &state.phase { + return lifecycle_error(id, violation); + } + + if method == "initialize" { + return match state.phase.clone() { + ConnectionPhase::Uninitialized => self.initialize(state, id, message), + ConnectionPhase::Modern => era_locked_error( + id, + ProtocolEra::PerRequest, + ProtocolEra::InitializationBased, + ), + ConnectionPhase::LegacyReady { .. } + | ConnectionPhase::LegacyAwaitingInitialized(_) => self.lifecycle_violation( + state, + id, + "initialize-already-completed", + phase_name(&state.phase), + method, + ), + ConnectionPhase::Poisoned(_) => unreachable!("handled above"), + }; + } + + match state.phase.clone() { + ConnectionPhase::Uninitialized => { + if method == "server/discover" || requests_modern_era(message) { + if let Err(response) = validate_modern_request(id.clone(), message) { + return *response; + } + state.phase = ConnectionPhase::Modern; + self.dispatch(state, id, ProtocolRevision::V2026_07_28, method, message) + } else { + self.lifecycle_violation( + state, + id, + "initialize-required", + "uninitialized", + method, + ) + } + } + ConnectionPhase::LegacyAwaitingInitialized(revision) => { + if requests_modern_era(message) { + return era_locked_error( + id, + ProtocolEra::InitializationBased, + ProtocolEra::PerRequest, + ); + } + if let Err(response) = validate_legacy_request(id.clone(), message) { + return *response; + } + if method == "ping" { + self.dispatch(state, id, revision, method, message) + } else { + self.lifecycle_violation( + state, + id, + "initialized-notification-required", + "initialize-responded", + method, + ) + } + } + ConnectionPhase::LegacyReady { revision, .. } => { + if requests_modern_era(message) { + era_locked_error( + id, + ProtocolEra::InitializationBased, + ProtocolEra::PerRequest, + ) + } else { + if let Err(response) = validate_legacy_request(id.clone(), message) { + return *response; + } + self.dispatch(state, id, revision, method, message) + } + } + ConnectionPhase::Modern => { + if let Err(response) = validate_modern_request(id.clone(), message) { + return *response; + } + self.dispatch(state, id, ProtocolRevision::V2026_07_28, method, message) + } + ConnectionPhase::Poisoned(_) => unreachable!("handled above"), + } + } + + fn handle_notification( + &self, + state: &mut ConnectionState, + method: &str, + message: &JsonRpcMessage, + ) { + if method == "notifications/initialized" { + if let Err(error) = validate_initialized_notification(message) { + tracing::warn!(error, "ignored malformed MCP initialized notification"); + return; + } + } + if method != "notifications/initialized" { + if matches!(state.phase, ConnectionPhase::Modern) { + if let Err(error) = validate_modern_notification(message) { + tracing::warn!(method, error, "ignored malformed modern MCP notification"); + } + } + return; + } + + match state.phase.clone() { + ConnectionPhase::Uninitialized => { + state.phase = ConnectionPhase::Poisoned(LifecycleViolation { + kind: "initialized-before-initialize", + state: "protocol-error", + method: method.into(), + }); + } + ConnectionPhase::LegacyAwaitingInitialized(revision) => { + state.phase = ConnectionPhase::LegacyReady { + revision, + initialized_seen: true, + }; + } + ConnectionPhase::LegacyReady { + revision, + initialized_seen: false, + } if revision == ProtocolRevision::V2024_11_05 => { + state.phase = ConnectionPhase::LegacyReady { + revision, + initialized_seen: true, + }; + } + ConnectionPhase::LegacyReady { .. } => { + state.phase = ConnectionPhase::Poisoned(LifecycleViolation { + kind: "initialized-already-received", + state: "protocol-error", + method: method.into(), + }); + } + ConnectionPhase::Modern => { + if let Err(error) = validate_modern_notification(message) { + tracing::warn!(method, error, "ignored malformed modern MCP notification"); + } + } + ConnectionPhase::Poisoned(_) => {} + } + } + + fn initialize( + &self, + state: &mut ConnectionState, + id: Value, + message: &JsonRpcMessage, + ) -> JsonRpcResponse { + if let Err(response) = validate_legacy_request(id.clone(), message) { + return *response; + } + let Some(params) = message.params.as_ref().and_then(Value::as_object) else { + return JsonRpcResponse::err(id, -32602, "initialize params must be an object".into()); + }; + let Some(requested) = params.get("protocolVersion").and_then(Value::as_str) else { + return JsonRpcResponse::err( + id, + -32602, + "initialize.protocolVersion must be a string".into(), + ); + }; + + let revision = match ProtocolRevision::parse_exact(requested) { + Some(ProtocolRevision::V2026_07_28) => { + return JsonRpcResponse::err( + id, + -32602, + "2026-07-28 does not use initialize; send per-request metadata or call server/discover" + .into(), + ) + } + Some(revision) => revision, + None => ProtocolRevision::V2025_11_25, + }; + + if !params.get("capabilities").is_some_and(|capabilities| { + valid_initialize_client_capabilities(revision, capabilities) + }) { + return JsonRpcResponse::err( + id, + -32602, + "initialize.capabilities must be a valid client capabilities object".into(), + ); + } + if !params + .get("clientInfo") + .is_some_and(|identity| valid_initialize_implementation_identity(revision, identity)) + { + return JsonRpcResponse::err( + id, + -32602, + "initialize.clientInfo must be valid for the negotiated protocol revision".into(), + ); + } + + state.phase = if revision == ProtocolRevision::V2024_11_05 { + ConnectionPhase::LegacyReady { + revision, + initialized_seen: false, + } + } else { + ConnectionPhase::LegacyAwaitingInitialized(revision) + }; + + let capabilities = if revision == ProtocolRevision::V2024_11_05 { + json!({ "tools": {} }) + } else { + json!({ "tools": {}, "resources": {} }) + }; + JsonRpcResponse::ok( + id, + json!({ + "protocolVersion": revision.as_str(), + "capabilities": capabilities, + "serverInfo": server_info(), + "instructions": ICM_INSTRUCTIONS, + }), + ) + } + + fn lifecycle_violation( + &self, + state: &mut ConnectionState, + id: Value, + kind: &'static str, + state_name: &'static str, + method: &str, + ) -> JsonRpcResponse { + let violation = LifecycleViolation { + kind, + state: state_name, + method: bounded_lifecycle_method(method), + }; + let response = lifecycle_error(id, &violation); + state.phase = ConnectionPhase::Poisoned(violation); + response + } + + fn dispatch( + &self, + state: &mut ConnectionState, + id: Value, + revision: ProtocolRevision, + method: &str, + message: &JsonRpcMessage, + ) -> JsonRpcResponse { + match method { + "ping" => JsonRpcResponse::ok(id, project_result(revision, json!({}), None)), + "server/discover" if revision == ProtocolRevision::V2026_07_28 => { + JsonRpcResponse::ok(id, discovery_result()) + } + "tools/list" => self.list_tools(id, revision, message), + "tools/call" => self.call_tool(state, id, revision, message), + "resources/list" if revision != ProtocolRevision::V2024_11_05 => { + let result = if revision == ProtocolRevision::V2026_07_28 { + project_result( + revision, + json!({ "resources": [] }), + Some((3_600_000, "private")), + ) + } else { + json!({ + "resources": [], + "_meta": { "ttlMs": 0, "cacheScope": "private" } + }) + }; + JsonRpcResponse::ok(id, result) + } + "resources/read" if revision != ProtocolRevision::V2024_11_05 => { + let uri = message + .params + .as_ref() + .and_then(Value::as_object) + .and_then(|params| params.get("uri")) + .and_then(Value::as_str); + let Some(uri) = uri else { + return JsonRpcResponse::err( + id, + -32602, + "resources/read.uri must be a string".into(), + ); + }; + let code = if revision == ProtocolRevision::V2026_07_28 { + -32602 + } else { + -32002 + }; + JsonRpcResponse::err_with_data( + id, + code, + "resource not found".into(), + Some(json!({ "uri": uri })), + ) + } + other => JsonRpcResponse::method_not_found(id, other), + } + } + + fn list_tools( + &self, + id: Value, + revision: ProtocolRevision, + message: &JsonRpcMessage, + ) -> JsonRpcResponse { + if revision != ProtocolRevision::V2024_11_05 { + let cursor = message + .params + .as_ref() + .and_then(Value::as_object) + .and_then(|params| params.get("cursor")); + if cursor.is_some_and(|cursor| !cursor.is_null() && cursor.as_str() != Some("")) { + return JsonRpcResponse::err( + id, + -32602, + "tools/list cursor is not supported for the immutable catalog".into(), + ); + } + } + let result = match revision { + ProtocolRevision::V2024_11_05 => self.catalog.legacy_list(), + ProtocolRevision::V2025_06_18 | ProtocolRevision::V2025_11_25 => { + self.catalog.modern_list() + } + ProtocolRevision::V2026_07_28 => project_result( + revision, + self.catalog.modern_list(), + Some((3_600_000, "private")), + ), + }; + JsonRpcResponse::ok(id, result) + } + + fn call_tool( + &self, + state: &mut ConnectionState, + id: Value, + revision: ProtocolRevision, + message: &JsonRpcMessage, + ) -> JsonRpcResponse { + let Some(params) = message.params.as_ref().and_then(Value::as_object) else { + return JsonRpcResponse::err(id, -32602, "missing params".into()); + }; + let Some(name) = params.get("name").and_then(Value::as_str) else { + return JsonRpcResponse::err(id, -32602, "missing tool name".into()); + }; + let arguments = params + .get("arguments") + .cloned() + .unwrap_or_else(|| json!({})); + if !arguments.is_object() { + return JsonRpcResponse::err(id, -32602, "tool arguments must be an object".into()); + } + + if name == "icm_memory_store" { + state.calls_since_store = 0; + } else { + state.calls_since_store = state.calls_since_store.saturating_add(1); + } + + let context = ToolContext { + store: self.store, + embedder: self.embedder, + compact: self.compact, + auto_consolidate: self.auto_consolidate, + working_directory: &self.working_directory, + enforce_directory_boundary: true, + }; + let validation = if revision == ProtocolRevision::V2024_11_05 { + InputValidation::Legacy2024 + } else { + InputValidation::Modern + }; + let mut result = match self + .catalog + .dispatch(&context, name, &arguments, validation) + { + DispatchResult::ToolResult(result) => result, + DispatchResult::UnknownTool if revision == ProtocolRevision::V2024_11_05 => { + crate::protocol::ToolResult::error(format!("unknown tool: {name}")) + } + DispatchResult::UnknownTool => { + return JsonRpcResponse::err(id, -32602, format!("unknown tool: {name}")) + } + DispatchResult::InvalidInput(message) => { + return self.invalid_tool_arguments_result(id, revision, message) + } + }; + + if revision != ProtocolRevision::V2026_07_28 + && name != "icm_memory_store" + && state.calls_since_store >= STORE_NUDGE_THRESHOLD + && state + .calls_since_store + .is_multiple_of(STORE_NUDGE_THRESHOLD) + { + result.append_hint(&format!( + "\n[ICM: {} tool calls since last store. Consider saving important context with \ + icm_memory_store before it is lost.]", + state.calls_since_store + )); + } + + let value = serde_json::to_value(result).unwrap_or(Value::Null); + JsonRpcResponse::ok(id, project_result(revision, value, None)) + } + + fn invalid_tool_arguments_result( + &self, + id: Value, + revision: ProtocolRevision, + message: String, + ) -> JsonRpcResponse { + if revision != ProtocolRevision::V2024_11_05 { + return JsonRpcResponse::err(id, -32602, format!("invalid arguments: {message}")); + } + let result = crate::protocol::ToolResult::error(format!("invalid arguments: {message}")); + let value = serde_json::to_value(result).unwrap_or(Value::Null); + JsonRpcResponse::ok(id, project_result(revision, value, None)) + } +} + +fn valid_request_id(id: &Value) -> bool { + id.is_string() || id.as_i64().is_some() || id.as_u64().is_some() +} + +fn phase_name(phase: &ConnectionPhase) -> &'static str { + match phase { + ConnectionPhase::Uninitialized => "uninitialized", + ConnectionPhase::LegacyAwaitingInitialized(_) => "initialize-responded", + ConnectionPhase::LegacyReady { .. } => "initialized", + ConnectionPhase::Modern => "modern", + ConnectionPhase::Poisoned(_) => "protocol-error", + } +} + +fn bounded_lifecycle_method(method: &str) -> String { + if method.len() <= MAX_STORED_LIFECYCLE_METHOD_BYTES { + method.into() + } else { + format!("", method.len()) + } +} + +fn lifecycle_error(id: Value, violation: &LifecycleViolation) -> JsonRpcResponse { + JsonRpcResponse::err_with_data( + id, + -31011, + "protocol lifecycle violation; open a new connection".into(), + Some(json!({ + "kind": violation.kind, + "state": violation.state, + "method": violation.method, + })), + ) +} + +fn era_locked_error(id: Value, selected: ProtocolEra, requested: ProtocolEra) -> JsonRpcResponse { + JsonRpcResponse::err_with_data( + id, + -31010, + "protocol era is locked for this connection; open a new connection".into(), + Some(json!({ + "kind": "protocolEraLocked", + "selectedEra": selected.as_str(), + "requestedEra": requested.as_str(), + })), + ) +} + +fn requests_modern_era(message: &JsonRpcMessage) -> bool { + message + .params + .as_ref() + .and_then(Value::as_object) + .and_then(|params| params.get("_meta")) + .and_then(Value::as_object) + .is_some_and(|metadata| { + [ + MODERN_PROTOCOL_VERSION_KEY, + MODERN_CLIENT_CAPABILITIES_KEY, + MODERN_CLIENT_INFO_KEY, + ] + .into_iter() + .any(|key| metadata.contains_key(key)) + }) +} + +fn validate_legacy_request( + id: Value, + message: &JsonRpcMessage, +) -> Result<(), Box> { + if message.extra.contains_key("_meta") { + return Err(invalid_params( + id, + "request metadata must be nested at params._meta", + )); + } + let Some(params) = message.params.as_ref() else { + return Ok(()); + }; + let Some(params) = params.as_object() else { + return Err(invalid_params(id, "request params must be an object")); + }; + let Some(raw_metadata) = params.get("_meta") else { + return Ok(()); + }; + let Some(metadata) = raw_metadata.as_object() else { + return Err(invalid_params(id, "params._meta must be an object")); + }; + validate_legacy_metadata(&id, metadata) +} + +fn validate_modern_request( + id: Value, + message: &JsonRpcMessage, +) -> Result<(), Box> { + if message.extra.contains_key("_meta") { + return Err(invalid_params( + id, + "request metadata must be nested at params._meta", + )); + } + let Some(params) = message.params.as_ref().and_then(Value::as_object) else { + return Err(invalid_params( + id, + "modern request params must be an object", + )); + }; + let Some(metadata) = params.get("_meta").and_then(Value::as_object) else { + return Err(invalid_params(id, "params._meta must be an object")); + }; + validate_metadata_shape(&id, metadata)?; + + let Some(requested) = metadata + .get(MODERN_PROTOCOL_VERSION_KEY) + .and_then(Value::as_str) + else { + return Err(invalid_params( + id, + format!("missing or invalid {MODERN_PROTOCOL_VERSION_KEY}"), + )); + }; + if requested != ProtocolRevision::V2026_07_28.as_str() { + return Err(Box::new(JsonRpcResponse::err_with_data( + id, + -32022, + format!("unsupported protocol version: {requested}"), + Some(json!({ + "supported": SUPPORTED_PROTOCOL_VERSIONS, + "requested": requested, + })), + ))); + } + + let Some(capabilities) = metadata + .get(MODERN_CLIENT_CAPABILITIES_KEY) + .filter(|capabilities| valid_client_capabilities(capabilities)) + else { + return Err(invalid_params( + id, + format!("missing or invalid {MODERN_CLIENT_CAPABILITIES_KEY}"), + )); + }; + debug_assert!(capabilities.is_object()); + + if metadata + .get(MODERN_CLIENT_INFO_KEY) + .is_some_and(|identity| !valid_implementation_identity(identity)) + { + return Err(invalid_params( + id, + format!("invalid {MODERN_CLIENT_INFO_KEY}"), + )); + } + validate_optional_metadata_values(&id, metadata) +} + +fn validate_modern_notification(message: &JsonRpcMessage) -> Result<(), String> { + if message.extra.contains_key("_meta") { + return Err("notification metadata must be nested at params._meta".into()); + } + let Some(params) = message.params.as_ref() else { + return Ok(()); + }; + let Some(params) = params.as_object() else { + return Err("notification params must be an object".into()); + }; + let Some(raw_metadata) = params.get("_meta") else { + return Ok(()); + }; + let Some(metadata) = raw_metadata.as_object() else { + return Err("notification params._meta must be an object".into()); + }; + let null_id = Value::Null; + validate_metadata_shape(&null_id, metadata) + .map_err(|_| "notification metadata has an invalid key or size".to_owned())?; + if metadata + .get(MODERN_PROTOCOL_VERSION_KEY) + .is_some_and(|version| version.as_str() != Some(ProtocolRevision::V2026_07_28.as_str())) + { + return Err(format!("invalid {MODERN_PROTOCOL_VERSION_KEY}")); + } + if metadata + .get(MODERN_CLIENT_CAPABILITIES_KEY) + .is_some_and(|capabilities| !valid_client_capabilities(capabilities)) + { + return Err(format!("invalid {MODERN_CLIENT_CAPABILITIES_KEY}")); + } + if metadata + .get(MODERN_CLIENT_INFO_KEY) + .is_some_and(|identity| !valid_implementation_identity(identity)) + { + return Err(format!("invalid {MODERN_CLIENT_INFO_KEY}")); + } + validate_optional_metadata_values(&null_id, metadata) + .map_err(|_| "notification metadata has an invalid value".to_owned()) +} + +fn validate_initialized_notification(message: &JsonRpcMessage) -> Result<(), String> { + if message.extra.contains_key("_meta") { + return Err("initialized metadata must be nested at params._meta".into()); + } + let Some(params) = message.params.as_ref() else { + return Ok(()); + }; + let Some(params) = params.as_object() else { + return Err("initialized params must be an object".into()); + }; + let Some(raw_metadata) = params.get("_meta") else { + return Ok(()); + }; + let Some(metadata) = raw_metadata.as_object() else { + return Err("initialized params._meta must be an object".into()); + }; + let null_id = Value::Null; + validate_legacy_metadata(&null_id, metadata) + .map_err(|_| "initialized metadata has an invalid shape or value".to_owned()) +} + +fn invalid_params(id: Value, message: impl Into) -> Box { + Box::new(JsonRpcResponse::err(id, -32602, message.into())) +} + +fn validate_metadata_shape( + id: &Value, + metadata: &Map, +) -> Result<(), Box> { + if metadata.len() > 64 + || serde_json::to_vec(metadata).is_ok_and(|encoded| encoded.len() > 65_536) + || metadata + .values() + .any(|value| !metadata_value_within_depth(value, 32)) + { + return Err(invalid_params( + id.clone(), + "params._meta exceeds the supported size or nesting depth", + )); + } + for key in metadata.keys() { + if !valid_metadata_key(key) { + return Err(invalid_params( + id.clone(), + format!("invalid metadata key: {key}"), + )); + } + } + Ok(()) +} + +fn validate_legacy_metadata( + id: &Value, + metadata: &Map, +) -> Result<(), Box> { + if metadata.len() > 64 + || serde_json::to_vec(metadata).is_ok_and(|encoded| encoded.len() > 65_536) + || metadata + .values() + .any(|value| !metadata_value_within_depth(value, 32)) + { + return Err(invalid_params( + id.clone(), + "params._meta exceeds the supported size or nesting depth", + )); + } + if metadata + .get("progressToken") + .is_some_and(|token| !(token.is_string() || token.is_number())) + { + return Err(invalid_params( + id.clone(), + "progressToken must be a string or number", + )); + } + if metadata + .get(MODERN_SUBSCRIPTION_ID_KEY) + .is_some_and(|subscription_id| !valid_request_id(subscription_id)) + { + return Err(invalid_params( + id.clone(), + format!("{MODERN_SUBSCRIPTION_ID_KEY} must be a string or integer"), + )); + } + Ok(()) +} + +fn metadata_value_within_depth(value: &Value, remaining: usize) -> bool { + match value { + Value::Array(values) => { + remaining > 0 + && values + .iter() + .all(|value| metadata_value_within_depth(value, remaining - 1)) + } + Value::Object(values) => { + remaining > 0 + && values + .values() + .all(|value| metadata_value_within_depth(value, remaining - 1)) + } + _ => true, + } +} + +fn validate_optional_metadata_values( + id: &Value, + metadata: &Map, +) -> Result<(), Box> { + if metadata + .get("progressToken") + .is_some_and(|token| !(token.is_string() || token.is_number())) + { + return Err(invalid_params( + id.clone(), + "progressToken must be a string or number", + )); + } + if metadata + .get(MODERN_SUBSCRIPTION_ID_KEY) + .is_some_and(|subscription_id| !valid_request_id(subscription_id)) + { + return Err(invalid_params( + id.clone(), + format!("{MODERN_SUBSCRIPTION_ID_KEY} must be a string or integer"), + )); + } + for (key, validator) in [ + ("traceparent", valid_traceparent as fn(&str) -> bool), + ("tracestate", valid_tracestate), + ("baggage", valid_baggage), + ] { + if metadata + .get(key) + .is_some_and(|value| !value.as_str().is_some_and(validator)) + { + return Err(invalid_params(id.clone(), format!("invalid {key}"))); + } + } + if metadata.get(MODERN_LOG_LEVEL_KEY).is_some_and(|level| { + !matches!( + level.as_str(), + Some( + "debug" + | "info" + | "notice" + | "warning" + | "error" + | "critical" + | "alert" + | "emergency" + ) + ) + }) { + return Err(invalid_params( + id.clone(), + format!("invalid {MODERN_LOG_LEVEL_KEY}"), + )); + } + Ok(()) +} + +fn valid_metadata_key(key: &str) -> bool { + match key.split_once('/') { + Some((prefix, name)) => { + !name.contains('/') && valid_metadata_prefix(prefix) && valid_metadata_name(name) + } + None => valid_metadata_name(key), + } +} + +fn valid_metadata_prefix(prefix: &str) -> bool { + !prefix.is_empty() && prefix.split('.').all(valid_metadata_prefix_label) +} + +fn valid_metadata_prefix_label(label: &str) -> bool { + let mut characters = label.chars(); + let Some(first) = characters.next() else { + return false; + }; + if !first.is_ascii_alphabetic() { + return false; + } + let Some(last) = label.chars().next_back() else { + return false; + }; + last.is_ascii_alphanumeric() + && label + .chars() + .all(|character| character.is_ascii_alphanumeric() || character == '-') +} + +fn valid_metadata_name(name: &str) -> bool { + if name.is_empty() { + return true; + } + let Some(first) = name.chars().next() else { + unreachable!("empty metadata names are handled above"); + }; + let Some(last) = name.chars().next_back() else { + return false; + }; + first.is_ascii_alphanumeric() + && last.is_ascii_alphanumeric() + && name.chars().all(|character| { + character.is_ascii_alphanumeric() || matches!(character, '-' | '_' | '.') + }) +} + +fn valid_initialize_client_capabilities(revision: ProtocolRevision, value: &Value) -> bool { + if revision == ProtocolRevision::V2026_07_28 { + return valid_client_capabilities(value); + } + let Some(capabilities) = bounded_capabilities(value) else { + return false; + }; + capabilities + .iter() + .all(|(name, capability)| match name.as_str() { + "experimental" => valid_experimental_capability(capability), + "roots" => valid_roots_capability(capability), + "sampling" + if matches!( + revision, + ProtocolRevision::V2024_11_05 | ProtocolRevision::V2025_06_18 + ) => + { + capability.is_object() + } + "sampling" => valid_capability_fields(capability, &["context", "tools"]), + "elicitation" if revision == ProtocolRevision::V2025_06_18 => capability.is_object(), + "elicitation" if revision == ProtocolRevision::V2025_11_25 => { + valid_capability_fields(capability, &["form", "url"]) + } + "tasks" if revision == ProtocolRevision::V2025_11_25 => { + valid_tasks_capability(capability) + } + _ => true, + }) +} + +fn valid_client_capabilities(value: &Value) -> bool { + let Some(capabilities) = bounded_capabilities(value) else { + return false; + }; + capabilities + .iter() + .all(|(name, capability)| match name.as_str() { + "experimental" => valid_experimental_capability(capability), + "roots" => valid_roots_capability(capability), + "sampling" => valid_capability_fields(capability, &["context", "tools"]), + "elicitation" => valid_capability_fields(capability, &["form", "url"]), + "extensions" => capability.as_object().is_some_and(|extensions| { + extensions.iter().all(|(identifier, settings)| { + valid_prefixed_metadata_key(identifier) && settings.is_object() + }) + }), + _ => true, + }) +} + +fn bounded_capabilities(value: &Value) -> Option<&Map> { + let capabilities = value.as_object()?; + (capabilities.len() <= 64 + && serde_json::to_vec(capabilities).is_ok_and(|encoded| encoded.len() <= 65_536)) + .then_some(capabilities) +} + +fn valid_experimental_capability(value: &Value) -> bool { + value + .as_object() + .is_some_and(|entries| entries.values().all(Value::is_object)) +} + +fn valid_roots_capability(value: &Value) -> bool { + value + .as_object() + .is_some_and(|fields| fields.get("listChanged").is_none_or(Value::is_boolean)) +} + +fn valid_capability_fields(value: &Value, allowed: &[&str]) -> bool { + value.as_object().is_some_and(|fields| { + fields + .iter() + .all(|(name, value)| !allowed.contains(&name.as_str()) || value.is_object()) + }) +} + +fn valid_tasks_capability(value: &Value) -> bool { + let Some(tasks) = value.as_object() else { + return false; + }; + if ["cancel", "list"] + .into_iter() + .any(|field| tasks.get(field).is_some_and(|value| !value.is_object())) + { + return false; + } + tasks.get("requests").is_none_or(|requests| { + requests.as_object().is_some_and(|requests| { + requests + .get("elicitation") + .is_none_or(|value| valid_capability_fields(value, &["create"])) + && requests + .get("sampling") + .is_none_or(|value| valid_capability_fields(value, &["createMessage"])) + }) + }) +} + +fn valid_prefixed_metadata_key(key: &str) -> bool { + key.split_once('/').is_some_and(|(prefix, name)| { + !name.contains('/') && valid_metadata_prefix(prefix) && valid_metadata_name(name) + }) +} + +fn valid_implementation_identity(value: &Value) -> bool { + let Some(identity) = value.as_object() else { + return false; + }; + if !valid_base_implementation_identity(identity) { + return false; + } + if ["title", "description", "websiteUrl"] + .into_iter() + .any(|field| identity.get(field).is_some_and(|value| !value.is_string())) + { + return false; + } + identity.get("icons").is_none_or(|icons| { + icons + .as_array() + .is_some_and(|icons| icons.iter().all(valid_icon)) + }) +} + +fn valid_initialize_implementation_identity(revision: ProtocolRevision, value: &Value) -> bool { + let Some(identity) = value.as_object() else { + return false; + }; + if !valid_base_implementation_identity(identity) { + return false; + } + match revision { + ProtocolRevision::V2024_11_05 => true, + ProtocolRevision::V2025_06_18 => identity.get("title").is_none_or(Value::is_string), + ProtocolRevision::V2025_11_25 | ProtocolRevision::V2026_07_28 => { + valid_implementation_identity(value) + } + } +} + +fn valid_base_implementation_identity(identity: &Map) -> bool { + ["name", "version"] + .into_iter() + .all(|field| identity.get(field).is_some_and(Value::is_string)) +} + +fn valid_icon(value: &Value) -> bool { + let Some(icon) = value.as_object() else { + return false; + }; + if !icon.get("src").is_some_and(Value::is_string) { + return false; + } + if icon.get("mimeType").is_some_and(|value| !value.is_string()) { + return false; + } + if icon.get("sizes").is_some_and(|sizes| { + !sizes + .as_array() + .is_some_and(|sizes| sizes.iter().all(Value::is_string)) + }) { + return false; + } + !icon + .get("theme") + .is_some_and(|theme| !matches!(theme.as_str(), Some("light" | "dark"))) +} + +fn valid_traceparent(value: &str) -> bool { + let bytes = value.as_bytes(); + if bytes.len() < 55 + || bytes.get(2) != Some(&b'-') + || bytes.get(35) != Some(&b'-') + || bytes.get(52) != Some(&b'-') + || !bytes[0..2].iter().copied().all(is_lower_hex) + || !bytes[3..35].iter().copied().all(is_lower_hex) + || !bytes[36..52].iter().copied().all(is_lower_hex) + || !bytes[53..55].iter().copied().all(is_lower_hex) + || &bytes[0..2] == b"ff" + || bytes[3..35].iter().all(|byte| *byte == b'0') + || bytes[36..52].iter().all(|byte| *byte == b'0') + { + return false; + } + + if &bytes[0..2] == b"00" { + bytes.len() == 55 + } else { + bytes.len() == 55 || bytes.get(55) == Some(&b'-') + } +} + +fn is_lower_hex(byte: u8) -> bool { + byte.is_ascii_digit() || matches!(byte, b'a'..=b'f') +} + +fn valid_tracestate(value: &str) -> bool { + let members: Vec<&str> = value.split(',').collect(); + if !(1..=32).contains(&members.len()) { + return false; + } + let mut keys = HashSet::new(); + for raw_member in members { + let member = trim_ows(raw_member); + if member.is_empty() { + return false; + } + let Some((key, value)) = member.split_once('=') else { + return false; + }; + if !valid_tracestate_key(key) || !valid_tracestate_value(value) || !keys.insert(key) { + return false; + } + } + true +} + +fn valid_tracestate_key(key: &str) -> bool { + if let Some((tenant_id, system_id)) = key.split_once('@') { + !tenant_id.contains('@') + && !system_id.contains('@') + && valid_tracestate_identifier(tenant_id, 241, true) + && valid_tracestate_identifier(system_id, 14, false) + } else { + valid_tracestate_identifier(key, 256, false) + } +} + +fn valid_tracestate_identifier(value: &str, maximum: usize, digit_start: bool) -> bool { + let bytes = value.as_bytes(); + (1..=maximum).contains(&bytes.len()) + && (bytes[0].is_ascii_lowercase() || (digit_start && bytes[0].is_ascii_digit())) + && bytes.iter().copied().all(|byte| { + byte.is_ascii_lowercase() + || byte.is_ascii_digit() + || matches!(byte, b'_' | b'-' | b'*' | b'/') + }) +} + +fn valid_tracestate_value(value: &str) -> bool { + let bytes = value.as_bytes(); + (1..=256).contains(&bytes.len()) + && bytes.last().is_some_and(|byte| *byte != b' ') + && bytes + .iter() + .copied() + .all(|byte| (0x20..=0x7e).contains(&byte) && !matches!(byte, b',' | b'=')) +} + +fn valid_baggage(value: &str) -> bool { + if value.len() > 8_192 { + return false; + } + let members: Vec<&str> = value.split(',').collect(); + (1..=64).contains(&members.len()) + && members + .into_iter() + .all(|member| valid_baggage_member(trim_ows(member))) +} + +fn valid_baggage_member(member: &str) -> bool { + let mut parts = member.split(';'); + parts + .next() + .is_some_and(|pair| valid_baggage_pair(pair, false)) + && parts.all(|property| valid_baggage_pair(property, true)) +} + +fn valid_baggage_pair(part: &str, key_only_allowed: bool) -> bool { + let part = trim_ows(part); + match part.split_once('=') { + Some((key, value)) => { + valid_http_token(trim_ows(key)) && valid_baggage_value(trim_ows(value)) + } + None => key_only_allowed && valid_http_token(part), + } +} + +fn valid_http_token(value: &str) -> bool { + !value.is_empty() + && value.bytes().all(|byte| { + byte.is_ascii_alphanumeric() + || matches!( + byte, + b'!' | b'#' + | b'$' + | b'%' + | b'&' + | b'\'' + | b'*' + | b'+' + | b'-' + | b'.' + | b'^' + | b'_' + | b'`' + | b'|' + | b'~' + ) + }) +} + +fn valid_baggage_value(value: &str) -> bool { + let bytes = value.as_bytes(); + let mut index = 0; + while index < bytes.len() { + let byte = bytes[index]; + if !matches!( + byte, + 0x21 | 0x23..=0x2b | 0x2d..=0x3a | 0x3c..=0x5b | 0x5d..=0x7e + ) { + return false; + } + if byte == b'%' { + if index + 2 >= bytes.len() + || !bytes[index + 1].is_ascii_hexdigit() + || !bytes[index + 2].is_ascii_hexdigit() + { + return false; + } + index += 3; + } else { + index += 1; + } + } + true +} + +fn trim_ows(value: &str) -> &str { + value.trim_matches(|character| matches!(character, ' ' | '\t')) +} + +fn server_info() -> Value { + json!({ "name": SERVER_NAME, "version": SERVER_VERSION }) +} + +fn project_result( + revision: ProtocolRevision, + mut value: Value, + cache: Option<(u64, &'static str)>, +) -> Value { + if revision != ProtocolRevision::V2026_07_28 { + return value; + } + let object = value + .as_object_mut() + .expect("MCP result projections must have object roots"); + object.insert("resultType".into(), Value::String("complete".into())); + let metadata = object + .entry("_meta") + .or_insert_with(|| Value::Object(Map::new())); + let metadata = metadata + .as_object_mut() + .expect("MCP result metadata must have an object root"); + metadata.insert(MODERN_SERVER_INFO_KEY.into(), server_info()); + if let Some((ttl_ms, scope)) = cache { + object.insert("ttlMs".into(), json!(ttl_ms)); + object.insert("cacheScope".into(), Value::String(scope.into())); + } + value +} + +fn discovery_result() -> Value { + project_result( + ProtocolRevision::V2026_07_28, + json!({ + "supportedVersions": SUPPORTED_PROTOCOL_VERSIONS, + "capabilities": { "tools": {}, "resources": {} }, + "instructions": ICM_INSTRUCTIONS, + }), + Some((3_600_000, "private")), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use icm_core::{Importance, Memory, MemoryStore}; + + fn service(store: &Store) -> McpService<'_> { + McpService::new(store, None, false, AutoConsolidate::default()) + } + + fn request(value: Value) -> JsonRpcMessage { + serde_json::from_value(value).unwrap() + } + + fn initialize_response( + service: &McpService<'_>, + state: &mut ConnectionState, + protocol_version: &str, + capabilities: Value, + client_info: Value, + ) -> JsonRpcResponse { + service + .handle( + state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"initialize", + "params":{ + "protocolVersion":protocol_version, + "capabilities":capabilities, + "clientInfo":client_info + } + })), + ) + .unwrap() + } + + fn initialize_2025(service: &McpService<'_>, state: &mut ConnectionState) { + service.handle( + state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"initialize", + "params":{ + "protocolVersion":"2025-11-25","capabilities":{}, + "clientInfo":{"name":"test","version":"1"} + } + })), + ); + assert!(service + .handle( + state, + request(json!({ + "jsonrpc":"2.0","method":"notifications/initialized","params":{} + })), + ) + .is_none()); + } + + fn initialized_state_for_revision( + service: &McpService<'_>, + revision: ProtocolRevision, + ) -> ConnectionState { + assert_ne!(revision, ProtocolRevision::V2026_07_28); + let mut state = ConnectionState::default(); + let response = initialize_response( + service, + &mut state, + revision.as_str(), + json!({}), + json!({"name":"test","version":"1"}), + ); + assert!(response.error.is_none()); + if revision != ProtocolRevision::V2024_11_05 { + assert!(service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","method":"notifications/initialized","params":{} + })), + ) + .is_none()); + } + state + } + + fn modern_metadata() -> Value { + json!({ + MODERN_PROTOCOL_VERSION_KEY: "2026-07-28", + MODERN_CLIENT_CAPABILITIES_KEY: {} + }) + } + + #[test] + fn frozen_2024_is_ready_without_initialized_notification() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + let mut state = ConnectionState::default(); + let initialized = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"initialize", + "params":{ + "protocolVersion":"2024-11-05","capabilities":{}, + "clientInfo":{"name":"test","version":"1"} + } + })), + ) + .unwrap(); + assert_eq!(initialized.result.unwrap()["protocolVersion"], "2024-11-05"); + let listed = service + .handle( + &mut state, + request(json!({"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}})), + ) + .unwrap(); + assert!(listed.result.unwrap()["tools"].is_array()); + } + + #[test] + fn poisoned_lifecycle_state_does_not_retain_or_reemit_oversized_method_names() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + let mut state = ConnectionState::default(); + let oversized_method = "m".repeat(MAX_STORED_LIFECYCLE_METHOD_BYTES + 1); + + let first = service + .handle( + &mut state, + JsonRpcMessage { + jsonrpc: "2.0".into(), + id: Some(json!(1)), + method: Some(oversized_method), + params: Some(json!({})), + extra: Map::new(), + }, + ) + .unwrap(); + let first_method = first + .error + .as_ref() + .and_then(|error| error.data.as_ref()) + .and_then(|data| data.get("method")) + .and_then(Value::as_str) + .unwrap(); + assert_eq!( + first_method, + format!( + "", + MAX_STORED_LIFECYCLE_METHOD_BYTES + 1 + ) + ); + + let repeated = service + .handle( + &mut state, + request(json!({"jsonrpc":"2.0","id":2,"method":"ping","params":{}})), + ) + .unwrap(); + let encoded = serde_json::to_vec(&repeated).unwrap(); + assert!( + encoded.len() < 1_024, + "poisoned response was {} bytes", + encoded.len() + ); + } + + #[test] + fn initialized_2025_requires_notification() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + let mut state = ConnectionState::default(); + service.handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"initialize", + "params":{ + "protocolVersion":"2025-11-25","capabilities":{}, + "clientInfo":{"name":"test","version":"1"} + } + })), + ); + let response = service + .handle( + &mut state, + request(json!({"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}})), + ) + .unwrap(); + assert_eq!(response.error.unwrap().code, -31011); + } + + #[test] + fn initialize_2025_validates_request_metadata_before_state_transition() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + let mut state = ConnectionState::default(); + + for (id, metadata) in [ + (1, json!([])), + (2, json!({"progressToken":{"wrong":"type"}})), + ] { + let response = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":id,"method":"initialize", + "params":{ + "protocolVersion":"2025-11-25", + "capabilities":{}, + "clientInfo":{"name":"test","version":"1"}, + "_meta":metadata + } + })), + ) + .unwrap(); + assert_eq!(response.error.unwrap().code, -32602); + assert!(matches!(state.phase, ConnectionPhase::Uninitialized)); + } + + let valid = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":3,"method":"initialize", + "params":{ + "protocolVersion":"2025-11-25", + "capabilities":{}, + "clientInfo":{"name":"test","version":"1"}, + "_meta":{"progressToken":"initializing"} + } + })), + ) + .unwrap(); + assert_eq!(valid.result.unwrap()["protocolVersion"], "2025-11-25"); + assert!(matches!( + state.phase, + ConnectionPhase::LegacyAwaitingInitialized(ProtocolRevision::V2025_11_25) + )); + } + + #[test] + fn initialize_capabilities_use_the_negotiated_revision_schema() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + + let mut june_state = ConnectionState::default(); + let june = initialize_response( + &service, + &mut june_state, + "2025-06-18", + json!({ + "experimental":{"future":{}}, + "roots":{"listChanged":true,"future":5}, + "sampling":{"tools":5}, + "elicitation":{"form":5}, + "futureCapability":5 + }), + json!({"name":"test","version":"1"}), + ); + assert_eq!(june.result.unwrap()["protocolVersion"], "2025-06-18"); + assert!(matches!( + june_state.phase, + ConnectionPhase::LegacyAwaitingInitialized(ProtocolRevision::V2025_06_18) + )); + + let mut november_state = ConnectionState::default(); + let november = initialize_response( + &service, + &mut november_state, + "2025-11-25", + json!({"sampling":{"tools":5}}), + json!({"name":"test","version":"1"}), + ); + assert_eq!(november.error.unwrap().code, -32602); + assert!(matches!( + november_state.phase, + ConnectionPhase::Uninitialized + )); + + let mut legacy_state = ConnectionState::default(); + let legacy = initialize_response( + &service, + &mut legacy_state, + "2024-11-05", + json!({ + "sampling":{"tools":5}, + "elicitation":5, + "futureCapability":[1,2,3] + }), + json!({"name":"test","version":"1"}), + ); + assert_eq!(legacy.result.unwrap()["protocolVersion"], "2024-11-05"); + assert!(matches!( + legacy_state.phase, + ConnectionPhase::LegacyReady { + revision: ProtocolRevision::V2024_11_05, + initialized_seen: false + } + )); + + for (revision, capabilities) in [ + ("2024-11-05", json!({"roots":{"listChanged":"yes"}})), + ("2025-06-18", json!({"elicitation":5})), + ("2025-11-25", json!({"experimental":{"future":5}})), + ] { + let mut state = ConnectionState::default(); + let response = initialize_response( + &service, + &mut state, + revision, + capabilities, + json!({"name":"test","version":"1"}), + ); + assert_eq!(response.error.unwrap().code, -32602); + assert!(matches!(state.phase, ConnectionPhase::Uninitialized)); + } + } + + #[test] + fn initialize_2025_11_validates_tasks_capability_shapes() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + + let mut valid_state = ConnectionState::default(); + let valid = initialize_response( + &service, + &mut valid_state, + "2025-11-25", + json!({ + "tasks":{ + "cancel":{"future":5}, + "list":{}, + "requests":{ + "elicitation":{"create":{},"future":5}, + "sampling":{"createMessage":{},"future":false}, + "future":5 + }, + "future":true + } + }), + json!({"name":"test","version":"1"}), + ); + assert!(valid.error.is_none()); + assert!(matches!( + valid_state.phase, + ConnectionPhase::LegacyAwaitingInitialized(ProtocolRevision::V2025_11_25) + )); + + for malformed_tasks in [ + json!(5), + json!({"cancel":5}), + json!({"list":"yes"}), + json!({"requests":5}), + json!({"requests":{"elicitation":5}}), + json!({"requests":{"elicitation":{"create":5}}}), + json!({"requests":{"sampling":5}}), + json!({"requests":{"sampling":{"createMessage":5}}}), + ] { + let mut state = ConnectionState::default(); + let response = initialize_response( + &service, + &mut state, + "2025-11-25", + json!({"tasks":malformed_tasks}), + json!({"name":"test","version":"1"}), + ); + assert_eq!(response.error.unwrap().code, -32602); + assert!(matches!(state.phase, ConnectionPhase::Uninitialized)); + } + + let mut june_state = ConnectionState::default(); + let june = initialize_response( + &service, + &mut june_state, + "2025-06-18", + json!({"tasks":5}), + json!({"name":"test","version":"1"}), + ); + assert!(june.error.is_none()); + assert!(matches!( + june_state.phase, + ConnectionPhase::LegacyAwaitingInitialized(ProtocolRevision::V2025_06_18) + )); + } + + #[test] + fn initialize_identity_fields_are_revision_specific() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + + let mut legacy_state = ConnectionState::default(); + let legacy = initialize_response( + &service, + &mut legacy_state, + "2024-11-05", + json!({}), + json!({ + "name":"test","version":"1","title":5,"description":5,"icons":5 + }), + ); + assert!(legacy.error.is_none()); + assert!(matches!( + legacy_state.phase, + ConnectionPhase::LegacyReady { + revision: ProtocolRevision::V2024_11_05, + initialized_seen: false + } + )); + + let mut june_state = ConnectionState::default(); + let june = initialize_response( + &service, + &mut june_state, + "2025-06-18", + json!({}), + json!({"name":"test","version":"1","title":"Test","description":5}), + ); + assert!(june.error.is_none()); + + let mut malformed_june_state = ConnectionState::default(); + let malformed_june = initialize_response( + &service, + &mut malformed_june_state, + "2025-06-18", + json!({}), + json!({"name":"test","version":"1","title":5}), + ); + assert_eq!(malformed_june.error.unwrap().code, -32602); + assert!(matches!( + malformed_june_state.phase, + ConnectionPhase::Uninitialized + )); + + let mut november_state = ConnectionState::default(); + let november = initialize_response( + &service, + &mut november_state, + "2025-11-25", + json!({}), + json!({"name":"test","version":"1","description":5}), + ); + assert_eq!(november.error.unwrap().code, -32602); + assert!(matches!( + november_state.phase, + ConnectionPhase::Uninitialized + )); + + let mut unknown_state = ConnectionState::default(); + let malformed_unknown = initialize_response( + &service, + &mut unknown_state, + "2099-01-01", + json!({}), + json!({"name":"test","version":"1","description":5}), + ); + assert_eq!(malformed_unknown.error.unwrap().code, -32602); + assert!(matches!( + unknown_state.phase, + ConnectionPhase::Uninitialized + )); + + let negotiated_unknown = initialize_response( + &service, + &mut unknown_state, + "2099-01-01", + json!({"sampling":{"tools":{}}}), + json!({"name":"test","version":"1","description":"valid"}), + ); + assert_eq!( + negotiated_unknown.result.unwrap()["protocolVersion"], + "2025-11-25" + ); + assert!(matches!( + unknown_state.phase, + ConnectionPhase::LegacyAwaitingInitialized(ProtocolRevision::V2025_11_25) + )); + + let mut modern_state = ConnectionState::default(); + let modern = initialize_response( + &service, + &mut modern_state, + "2026-07-28", + json!(5), + json!(5), + ); + let modern_error = modern.error.unwrap(); + assert_eq!(modern_error.code, -32602); + assert!(modern_error.message.contains("does not use initialize")); + assert!(matches!(modern_state.phase, ConnectionPhase::Uninitialized)); + } + + #[test] + fn malformed_initialized_notifications_do_not_advance_or_poison_state() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + + let mut still_awaiting = ConnectionState::default(); + service.handle( + &mut still_awaiting, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"initialize", + "params":{ + "protocolVersion":"2025-11-25","capabilities":{}, + "clientInfo":{"name":"test","version":"1"} + } + })), + ); + assert!(service + .handle( + &mut still_awaiting, + request(json!({ + "jsonrpc":"2.0","method":"notifications/initialized","params":[] + })), + ) + .is_none()); + let required = service + .handle( + &mut still_awaiting, + request(json!({"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}})), + ) + .unwrap(); + assert_eq!(required.error.as_ref().unwrap().code, -31011); + assert_eq!( + required.error.unwrap().data.unwrap()["kind"], + "initialized-notification-required" + ); + + let mut recoverable = ConnectionState::default(); + service.handle( + &mut recoverable, + request(json!({ + "jsonrpc":"2.0","id":3,"method":"initialize", + "params":{ + "protocolVersion":"2025-11-25","capabilities":{}, + "clientInfo":{"name":"test","version":"1"} + } + })), + ); + assert!(service + .handle( + &mut recoverable, + request(json!({ + "jsonrpc":"2.0","method":"notifications/initialized", + "params":{"_meta":[]} + })), + ) + .is_none()); + assert!(service + .handle( + &mut recoverable, + request(json!({ + "jsonrpc":"2.0","method":"notifications/initialized","params":{} + })), + ) + .is_none()); + let listed = service + .handle( + &mut recoverable, + request(json!({"jsonrpc":"2.0","id":4,"method":"tools/list","params":{}})), + ) + .unwrap(); + assert!(listed.error.is_none()); + } + + #[test] + fn initialized_2025_allows_progress_only_request_metadata() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + let mut state = ConnectionState::default(); + service.handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"initialize", + "params":{ + "protocolVersion":"2025-11-25","capabilities":{}, + "clientInfo":{"name":"test","version":"1"} + } + })), + ); + assert!(service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","method":"notifications/initialized","params":{} + })), + ) + .is_none()); + let response = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"tools/list", + "params":{"_meta":{"progressToken":1}} + })), + ) + .unwrap(); + assert!(response.error.is_none()); + assert!(response.result.unwrap()["tools"].is_array()); + } + + #[test] + fn initialized_2025_validates_optional_metadata_without_switching_era() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + let mut state = ConnectionState::default(); + initialize_2025(&service, &mut state); + + let malformed = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"tools/list", + "params":{"_meta":[]} + })), + ) + .unwrap(); + assert_eq!(malformed.error.unwrap().code, -32602); + + let valid = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":3,"method":"tools/list", + "params":{"_meta":{ + "progressToken":1.5, + "com.example/evaluation":{"opaque":true}, + "foo":1, + "arbitrary":[{"nested":true}] + }} + })), + ) + .unwrap(); + assert!(valid.error.is_none()); + } + + #[test] + fn discovery_accepts_optional_client_info() { + let store = Store::in_memory().unwrap(); + let initialized_service = service(&store); + let mut state = ConnectionState::default(); + let response = initialized_service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"server/discover", + "params":{"_meta":{ + MODERN_PROTOCOL_VERSION_KEY:"2026-07-28", + MODERN_CLIENT_CAPABILITIES_KEY:{} + }} + })), + ) + .unwrap(); + let result = response.result.unwrap(); + assert_eq!(result["resultType"], "complete"); + assert_eq!(result["_meta"][MODERN_SERVER_INFO_KEY]["name"], SERVER_NAME); + } + + #[test] + fn malformed_2026_identity_does_not_lock_the_connection() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + let mut state = ConnectionState::default(); + let malformed = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"server/discover", + "params":{"_meta":{ + MODERN_PROTOCOL_VERSION_KEY:"2026-07-28", + MODERN_CLIENT_CAPABILITIES_KEY:{}, + MODERN_CLIENT_INFO_KEY:{ + "name":"test","version":"1","icons":"not-an-array" + } + }} + })), + ) + .unwrap(); + assert_eq!(malformed.error.unwrap().code, -32602); + + let valid = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"server/discover", + "params":{"_meta":modern_metadata()} + })), + ) + .unwrap(); + assert!(valid.error.is_none()); + } + + #[test] + fn client_capabilities_are_open_but_validate_known_final_shapes() { + let store = Store::in_memory().unwrap(); + let initialized_service = service(&store); + let mut state = ConnectionState::default(); + let response = initialized_service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"initialize", + "params":{ + "protocolVersion":"2025-11-25", + "capabilities":{"tools":{},"resources":{}}, + "clientInfo":{"name":"test","version":"1"} + } + })), + ) + .unwrap(); + assert_eq!(response.result.unwrap()["protocolVersion"], "2025-11-25"); + + let modern_store = Store::in_memory().unwrap(); + let modern_service = service(&modern_store); + let mut modern_state = ConnectionState::default(); + let modern = modern_service + .handle( + &mut modern_state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"server/discover", + "params":{"_meta":{ + MODERN_PROTOCOL_VERSION_KEY:"2026-07-28", + MODERN_CLIENT_CAPABILITIES_KEY:{ + "sampling":{"context":{},"tools":{},"future":5}, + "elicitation":{"form":{},"url":{},"future":false}, + "experimental":{"x":{}}, + "unknownCapability":5, + "extensions":{"com.example/feature":{}} + } + }} + })), + ) + .unwrap(); + assert!(modern.error.is_none()); + + for capabilities in [ + json!({"sampling":{"tools":5}}), + json!({"elicitation":{"form":"yes"}}), + json!({"experimental":{"x":5}}), + json!({"extensions":{"unprefixed":{}}}), + json!({"extensions":{"com.example/feature":5}}), + json!({"roots":5}), + json!({"roots":{"listChanged":"yes"}}), + ] { + assert!(!valid_client_capabilities(&capabilities)); + } + assert!(valid_client_capabilities(&json!({ + "sampling":{"unknown":5}, + "elicitation":{"unknown":"opaque"}, + "unknownCapability":[1,2,3] + }))); + } + + #[test] + fn modern_metadata_reuses_the_legacy_nesting_bound() { + let mut nested = Value::Null; + for _ in 0..32 { + nested = json!({"next": nested}); + } + let mut metadata = Map::from_iter([("future".to_owned(), nested)]); + assert!(validate_metadata_shape(&Value::Null, &metadata).is_ok()); + + let nested = metadata.remove("future").unwrap(); + metadata.insert("future".into(), json!({"next": nested})); + assert!(validate_metadata_shape(&Value::Null, &metadata).is_err()); + } + + #[test] + fn implementation_identity_validates_final_field_types_without_invented_bounds() { + assert!(valid_implementation_identity(&json!({ + "name":"", + "version":"", + "title":"", + "description":"", + "websiteUrl":"not interpreted as a URI by structural validation", + "icons":[{ + "src":"data:,", + "mimeType":"", + "sizes":["", "not-a-size"], + "theme":"dark", + "future":{"opaque":true} + }], + "futureField":5 + }))); + for identity in [ + json!({"name":"test"}), + json!({"name":5,"version":"1"}), + json!({"name":"test","version":"1","description":5}), + json!({"name":"test","version":"1","icons":[{}]}), + json!({"name":"test","version":"1","icons":[{"src":"x","sizes":[5]}]}), + json!({"name":"test","version":"1","icons":[{"src":"x","theme":"auto"}]}), + ] { + assert!(!valid_implementation_identity(&identity)); + } + } + + #[test] + fn metadata_keys_follow_the_final_optional_prefix_grammar() { + for key in [ + "", + "progressToken", + "traceparent", + "tracestate", + "baggage", + "invalid", + "vendor_hint", + "io.modelcontextprotocol/futureField", + "dev.mcp/future", + "com.example/key_name.v2", + "com.example/", + ] { + assert!( + valid_metadata_key(key), + "expected valid metadata key: {key}" + ); + } + + for key in [ + "1bad/foo", + "bad_/foo", + "bad-/foo", + ".bad/foo", + "bad..name/foo", + "/foo", + "com.example/_bad", + "com.example/-bad", + "com.example/.bad", + "com.example/bad_", + "com.example/bad-", + "com.example/bad.", + "com.example/bad/extra", + "com.example/💥", + ] { + assert!( + !valid_metadata_key(key), + "expected invalid metadata key: {key}" + ); + } + } + + #[test] + fn opaque_metadata_and_subscription_ids_do_not_change_dispatch() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + let mut state = ConnectionState::default(); + assert!(service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"server/discover", + "params":{"_meta":modern_metadata()} + })), + ) + .unwrap() + .error + .is_none()); + + let baseline = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"ping", + "params":{"_meta":modern_metadata()} + })), + ) + .unwrap() + .result + .unwrap(); + + let mut enriched_metadata = modern_metadata(); + let enriched = enriched_metadata.as_object_mut().unwrap(); + enriched.insert("invalid".into(), json!({"opaque":true})); + enriched.insert("vendor_hint".into(), json!([1, 2, 3])); + enriched.insert( + "io.modelcontextprotocol/futureField".into(), + json!({"mustNotAuthorize":true}), + ); + enriched.insert("dev.mcp/future".into(), Value::Bool(true)); + enriched.insert("com.example/key_name.v2".into(), Value::Null); + enriched.insert("com.example/".into(), json!("empty-name")); + enriched.insert(MODERN_SUBSCRIPTION_ID_KEY.into(), json!("subscription")); + let enriched_result = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":3,"method":"ping", + "params":{"_meta":enriched_metadata} + })), + ) + .unwrap() + .result + .unwrap(); + assert_eq!(enriched_result, baseline); + + let mut integer_subscription = modern_metadata(); + integer_subscription + .as_object_mut() + .unwrap() + .insert(MODERN_SUBSCRIPTION_ID_KEY.into(), json!(42)); + assert!(service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":4,"method":"ping", + "params":{"_meta":integer_subscription} + })), + ) + .unwrap() + .error + .is_none()); + + for (id, invalid_subscription) in [(5, json!(1.5)), (6, json!(true)), (7, json!({}))] { + let mut metadata = modern_metadata(); + metadata + .as_object_mut() + .unwrap() + .insert(MODERN_SUBSCRIPTION_ID_KEY.into(), invalid_subscription); + let rejected = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":id,"method":"ping", + "params":{"_meta":metadata} + })), + ) + .unwrap(); + assert_eq!(rejected.error.unwrap().code, -32602); + } + + let mut invalid_key = modern_metadata(); + invalid_key + .as_object_mut() + .unwrap() + .insert("1bad/foo".into(), Value::Null); + let rejected = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":8,"method":"ping", + "params":{"_meta":invalid_key} + })), + ) + .unwrap(); + assert_eq!(rejected.error.unwrap().code, -32602); + } + + #[test] + fn modern_tracing_metadata_is_w3c_validated_and_typed() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + let mut state = ConnectionState::default(); + let accepted = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"server/discover", + "params":{"_meta":{ + MODERN_PROTOCOL_VERSION_KEY:"2026-07-28", + MODERN_CLIENT_CAPABILITIES_KEY:{}, + "traceparent":"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "tracestate":"vendor=value", + "baggage":"project=icm" + }} + })), + ) + .unwrap(); + assert!(accepted.error.is_none()); + + let rejected = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"ping", + "params":{"_meta":{ + MODERN_PROTOCOL_VERSION_KEY:"2026-07-28", + MODERN_CLIENT_CAPABILITIES_KEY:{}, + "traceparent":{"wrong":"type"} + }} + })), + ) + .unwrap(); + assert_eq!(rejected.error.unwrap().code, -32602); + } + + #[test] + fn w3c_trace_context_and_baggage_boundaries_are_exact() { + let trace_id = "4bf92f3577b34da6a3ce929d0e0e4736"; + let parent_id = "00f067aa0ba902b7"; + let current = format!("00-{trace_id}-{parent_id}-01"); + assert!(valid_traceparent(¤t)); + assert!(valid_traceparent(&format!("01-{trace_id}-{parent_id}-01"))); + assert!(valid_traceparent(&format!( + "01-{trace_id}-{parent_id}-01-future-fields-are-opaque" + ))); + for invalid in [ + format!("ff-{trace_id}-{parent_id}-01"), + format!("00-{}-{parent_id}-01", "0".repeat(32)), + format!("00-{trace_id}-{}-01", "0".repeat(16)), + format!("00-{trace_id}-{parent_id}-0A"), + format!("00-{}-{parent_id}-01", trace_id.to_ascii_uppercase()), + format!("00-{trace_id}-{parent_id}-01-extra"), + format!("01-{trace_id}-{parent_id}-01extra"), + ] { + assert!( + !valid_traceparent(&invalid), + "accepted invalid traceparent: {invalid}" + ); + } + + assert!(valid_tracestate("vendor=value")); + assert!(valid_tracestate("1tenant@system=value")); + assert!(valid_tracestate(&format!( + "one={},two={},three={}", + "a".repeat(200), + "b".repeat(200), + "c".repeat(200) + ))); + assert!(valid_tracestate(&format!("{}@s=value", "1".repeat(241)))); + assert!(valid_tracestate(&format!( + "tenant@{}=value", + "s".repeat(14) + ))); + assert!(valid_tracestate(&format!("vendor={}", "v".repeat(256)))); + let thirty_two_members = (0..32) + .map(|index| format!("k{index}=v")) + .collect::>() + .join(","); + assert!(valid_tracestate(&thirty_two_members)); + + for invalid in [ + "".to_owned(), + " \t ".to_owned(), + ",vendor=value".to_owned(), + "vendor=value,".to_owned(), + "one=value,,two=value".to_owned(), + "1simple=value".to_owned(), + "a@@b=value".to_owned(), + "a@1bad=value".to_owned(), + format!("{}@s=value", "1".repeat(242)), + format!("tenant@{}=value", "s".repeat(15)), + "duplicate=one,duplicate=two".to_owned(), + format!("vendor={}", "v".repeat(257)), + "vendor=bad=value".to_owned(), + "vendor=bad\nvalue".to_owned(), + (0..33) + .map(|index| format!("k{index}=v")) + .collect::>() + .join(","), + ] { + assert!( + !valid_tracestate(&invalid), + "accepted invalid tracestate: {invalid}" + ); + } + + assert!(valid_baggage("key=")); + assert!(valid_baggage( + "key=value;property;second=x%20y, other = value" + )); + assert!(valid_baggage(&format!("k={}", "a".repeat(8_190)))); + let sixty_four_members = (0..64) + .map(|index| format!("k{index}=v")) + .collect::>() + .join(","); + assert!(valid_baggage(&sixty_four_members)); + for invalid in [ + "".to_owned(), + "key=unencoded space".to_owned(), + "key=bad%ZZ".to_owned(), + "key=value;".to_owned(), + "key=\"quoted\"".to_owned(), + format!("k={}", "a".repeat(8_191)), + (0..65) + .map(|index| format!("k{index}=v")) + .collect::>() + .join(","), + ] { + assert!( + !valid_baggage(&invalid), + "accepted invalid baggage: {invalid}" + ); + } + + for key in ["traceparent", "tracestate", "baggage"] { + let metadata = Map::from_iter([(key.to_owned(), json!({"wrong":"type"}))]); + assert!(validate_optional_metadata_values(&Value::Null, &metadata).is_err()); + } + } + + #[test] + fn modern_notifications_have_optional_metadata_and_never_poison_requests() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + let mut state = ConnectionState::default(); + assert!(service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"server/discover", + "params":{"_meta":modern_metadata()} + })), + ) + .unwrap() + .error + .is_none()); + assert!(service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","method":"notifications/progress", + "params":{"progressToken":"work","progress":0.5} + })), + ) + .is_none()); + assert!(service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","method":"notifications/cancelled", + "params":{"requestId":1,"reason":"test"} + })), + ) + .is_none()); + assert!(service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","method":"notifications/progress", + "params":{"_meta":[]} + })), + ) + .is_none()); + + let response = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"ping", + "params":{"_meta":modern_metadata()} + })), + ) + .unwrap(); + assert!(response.error.is_none()); + assert_eq!(response.result.unwrap()["resultType"], "complete"); + } + + #[test] + fn resource_core_is_honest_and_empty() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + let mut state = ConnectionState::default(); + let response = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"resources/list", + "params":{"_meta":{ + MODERN_PROTOCOL_VERSION_KEY:"2026-07-28", + MODERN_CLIENT_CAPABILITIES_KEY:{} + }} + })), + ) + .unwrap(); + assert_eq!(response.result.unwrap()["resources"], json!([])); + } + + #[test] + fn explicit_working_directory_scopes_default_recall() { + let tmp = tempfile::tempdir().unwrap(); + let client_directory = tmp.path().join("client-project"); + std::fs::create_dir(&client_directory).unwrap(); + let store = Store::in_memory().unwrap(); + store + .store(Memory::new( + "context-client-project".into(), + "shared marker from client".into(), + Importance::High, + )) + .unwrap(); + store + .store(Memory::new( + "context-other-project".into(), + "shared marker from other".into(), + Importance::High, + )) + .unwrap(); + let service = McpService::with_working_directory( + &store, + None, + false, + AutoConsolidate::default(), + client_directory, + ); + + let mut state = ConnectionState::default(); + initialize_2025(&service, &mut state); + let result = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"tools/call", + "params":{ + "name":"icm_memory_recall", + "arguments":{"query":"shared marker"} + } + })), + ) + .unwrap() + .result + .unwrap(); + let text = result["content"][0]["text"].as_str().unwrap(); + assert!(text.contains("from client")); + assert!(!text.contains("from other")); + } + + #[test] + fn typed_tool_input_errors_use_revision_appropriate_envelopes() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + + for revision in [ + ProtocolRevision::V2024_11_05, + ProtocolRevision::V2025_06_18, + ProtocolRevision::V2025_11_25, + ] { + let mut state = initialized_state_for_revision(&service, revision); + let response = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"tools/call", + "params":{"name":"icm_memory_recall","arguments":{"query":" "}} + })), + ) + .unwrap(); + if revision == ProtocolRevision::V2024_11_05 { + assert!(response.error.is_none()); + let result = response.result.unwrap(); + assert_eq!(result["isError"], true); + assert!(result["content"][0]["text"] + .as_str() + .unwrap() + .starts_with("invalid arguments: ")); + } else { + assert!(response.result.is_none()); + let error = response.error.unwrap(); + assert_eq!(error.code, -32602); + assert!(error.message.starts_with("invalid arguments: ")); + } + } + + let mut modern_state = ConnectionState::default(); + let modern = service + .handle( + &mut modern_state, + request(json!({ + "jsonrpc":"2.0","id":3,"method":"tools/call", + "params":{ + "name":"icm_memory_recall", + "arguments":{"query":" "}, + "_meta":modern_metadata() + } + })), + ) + .unwrap(); + assert!(modern.result.is_none()); + let modern_error = modern.error.unwrap(); + assert_eq!(modern_error.code, -32602); + assert!(modern_error.message.starts_with("invalid arguments: ")); + } + + #[test] + fn valid_tool_business_errors_remain_tool_results_in_every_revision() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + + for revision in [ + ProtocolRevision::V2024_11_05, + ProtocolRevision::V2025_06_18, + ProtocolRevision::V2025_11_25, + ] { + let mut state = initialized_state_for_revision(&service, revision); + let response = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"tools/call", + "params":{ + "name":"icm_memory_forget", + "arguments":{"id":"does-not-exist"} + } + })), + ) + .unwrap(); + assert!(response.error.is_none()); + let result = response.result.unwrap(); + assert_eq!(result["isError"], true); + assert!(result.get("resultType").is_none()); + } + + let mut modern_state = ConnectionState::default(); + let modern = service + .handle( + &mut modern_state, + request(json!({ + "jsonrpc":"2.0","id":3,"method":"tools/call", + "params":{ + "name":"icm_memory_forget", + "arguments":{"id":"does-not-exist"}, + "_meta":modern_metadata() + } + })), + ) + .unwrap(); + assert!(modern.error.is_none()); + let modern_result = modern.result.unwrap(); + assert_eq!(modern_result["isError"], true); + assert_eq!(modern_result["resultType"], "complete"); + assert_eq!( + modern_result["_meta"][MODERN_SERVER_INFO_KEY]["name"], + SERVER_NAME + ); + } + + #[test] + fn topic_schema_and_runtime_enforce_the_multibyte_boundary() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + let mut state = ConnectionState::default(); + initialize_2025(&service, &mut state); + + let accepted_topic = format!("{}a", "é".repeat(127)); + assert_eq!(accepted_topic.chars().count(), 128); + assert_eq!(accepted_topic.len(), 255); + let accepted = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"tools/call", + "params":{ + "name":"icm_memory_store", + "arguments":{"topic":accepted_topic,"content":"valid"} + } + })), + ) + .unwrap(); + assert!(accepted.error.is_none()); + assert_ne!(accepted.result.unwrap()["isError"], true); + + let rejected_topic = "é".repeat(128); + assert_eq!(rejected_topic.chars().count(), 128); + assert_eq!(rejected_topic.len(), 256); + let rejected = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":3,"method":"tools/call", + "params":{ + "name":"icm_memory_store", + "arguments":{"topic":rejected_topic,"content":"invalid"} + } + })), + ) + .unwrap(); + assert!(rejected.result.is_none()); + let rejected_error = rejected.error.unwrap(); + assert_eq!(rejected_error.code, -32602); + assert!(rejected_error.message.contains("exceeds 255 UTF-8 bytes")); + } + + #[test] + fn recall_limit_preserves_legacy_normalization_and_modern_strictness() { + let store = Store::in_memory().unwrap(); + let consolidation_off = AutoConsolidate { + enabled: false, + threshold: 10, + }; + for index in 0..30 { + let stored = crate::tools::call_tool_with_config( + &store, + None, + "icm_memory_store", + &json!({ + "topic":"limit-probe", + "content":format!("revision limit probe entry {index}") + }), + false, + consolidation_off, + ); + assert!(!stored.is_error); + } + let service = service(&store); + + let mut legacy_state = ConnectionState::default(); + service.handle( + &mut legacy_state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"initialize", + "params":{ + "protocolVersion":"2024-11-05","capabilities":{}, + "clientInfo":{"name":"test","version":"1"} + } + })), + ); + for (id, limit, expected_hits) in [(2, 0, 1), (3, 21, 20), (4, 101, 20)] { + let response = service + .handle( + &mut legacy_state, + request(json!({ + "jsonrpc":"2.0","id":id,"method":"tools/call", + "params":{ + "name":"icm_memory_recall", + "arguments":{ + "query":"revision limit probe", + "project":"", + "limit":limit + } + } + })), + ) + .unwrap(); + assert!(response.error.is_none()); + let result = response.result.unwrap(); + assert_ne!(result["isError"], true); + assert_eq!( + result["content"][0]["text"] + .as_str() + .unwrap() + .matches("revision limit probe") + .count(), + expected_hits + ); + } + + let legacy_unknown = service + .handle( + &mut legacy_state, + request(json!({ + "jsonrpc":"2.0","id":5,"method":"tools/call", + "params":{ + "name":"icm_memory_recall", + "arguments":{ + "query":"revision limit probe","project":"","limit":2, + "futureField":{"opaque":[1,2,3]} + } + } + })), + ) + .unwrap(); + let legacy_unknown_result = legacy_unknown.result.unwrap(); + assert_ne!(legacy_unknown_result["isError"], true); + assert_eq!( + legacy_unknown_result["content"][0]["text"] + .as_str() + .unwrap() + .matches("revision limit probe") + .count(), + 2 + ); + + let legacy_bad_type = service + .handle( + &mut legacy_state, + request(json!({ + "jsonrpc":"2.0","id":6,"method":"tools/call", + "params":{ + "name":"icm_memory_recall", + "arguments":{ + "query":"revision limit probe","project":"","limit":"2" + } + } + })), + ) + .unwrap(); + assert_eq!(legacy_bad_type.result.unwrap()["isError"], true); + + let mut modern_state = ConnectionState::default(); + let accepted = service + .handle( + &mut modern_state, + request(json!({ + "jsonrpc":"2.0","id":7,"method":"tools/call", + "params":{ + "name":"icm_memory_recall", + "arguments":{ + "query":"revision limit probe","project":"","limit":100 + }, + "_meta":modern_metadata() + } + })), + ) + .unwrap(); + let accepted_result = accepted.result.unwrap(); + assert_ne!(accepted_result["isError"], true); + assert_eq!( + accepted_result["content"][0]["text"] + .as_str() + .unwrap() + .matches("revision limit probe") + .count(), + 30 + ); + + for (id, limit) in [(8, 0), (9, 101)] { + let rejected = service + .handle( + &mut modern_state, + request(json!({ + "jsonrpc":"2.0","id":id,"method":"tools/call", + "params":{ + "name":"icm_memory_recall", + "arguments":{ + "query":"revision limit probe","project":"","limit":limit + }, + "_meta":modern_metadata() + } + })), + ) + .unwrap(); + assert!(rejected.result.is_none()); + assert_eq!(rejected.error.unwrap().code, -32602); + } + + let modern_unknown = service + .handle( + &mut modern_state, + request(json!({ + "jsonrpc":"2.0","id":10,"method":"tools/call", + "params":{ + "name":"icm_memory_recall", + "arguments":{ + "query":"revision limit probe","project":"","limit":2, + "futureField":{"opaque":[1,2,3]} + }, + "_meta":modern_metadata() + } + })), + ) + .unwrap(); + assert!(modern_unknown.result.is_none()); + assert_eq!(modern_unknown.error.unwrap().code, -32602); + } + + #[test] + fn malformed_tool_call_and_unknown_name_remain_protocol_errors() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + let mut state = ConnectionState::default(); + initialize_2025(&service, &mut state); + + let malformed = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"tools/call", + "params":{"name":"icm_memory_recall","arguments":[]} + })), + ) + .unwrap(); + assert_eq!(malformed.error.unwrap().code, -32602); + + let unknown = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":3,"method":"tools/call", + "params":{"name":"icm_missing","arguments":{}} + })), + ) + .unwrap(); + assert_eq!(unknown.error.unwrap().code, -32602); + } + + #[test] + fn frozen_2024_unknown_tool_remains_a_legacy_tool_error() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + let mut state = ConnectionState::default(); + service.handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"initialize", + "params":{ + "protocolVersion":"2024-11-05","capabilities":{}, + "clientInfo":{"name":"test","version":"1"} + } + })), + ); + + let response = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"tools/call", + "params":{"name":"icm_missing","arguments":{}} + })), + ) + .unwrap(); + assert!(response.error.is_none()); + assert_eq!( + response.result.unwrap(), + json!({ + "content":[{"type":"text","text":"unknown tool: icm_missing"}], + "isError":true + }) + ); + } + + #[test] + fn unavailable_embedder_tool_stays_hidden_from_service_listing_and_dispatch() { + let store = Store::in_memory().unwrap(); + let service = service(&store); + let mut state = initialized_state_for_revision(&service, ProtocolRevision::V2024_11_05); + let listed = service + .handle( + &mut state, + request(json!({"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}})), + ) + .unwrap() + .result + .unwrap(); + assert!(!listed["tools"] + .as_array() + .unwrap() + .iter() + .any(|tool| tool["name"] == "icm_memory_embed_all")); + + let called = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":3,"method":"tools/call", + "params":{"name":"icm_memory_embed_all","arguments":{}} + })), + ) + .unwrap(); + assert!(called.error.is_none()); + assert_eq!( + called.result.unwrap(), + json!({ + "content":[{ + "type":"text", + "text":"unknown tool: icm_memory_embed_all" + }], + "isError":true + }) + ); + } + + #[test] + fn transcript_show_offset_is_tolerated_only_by_the_2024_projection() { + use icm_core::TranscriptStore; + + let store = Store::in_memory().unwrap(); + let session_id = store.create_session("test", None, None).unwrap(); + let service = service(&store); + let mut legacy_state = ConnectionState::default(); + service.handle( + &mut legacy_state, + request(json!({ + "jsonrpc":"2.0","id":1,"method":"initialize", + "params":{ + "protocolVersion":"2024-11-05","capabilities":{}, + "clientInfo":{"name":"test","version":"1"} + } + })), + ); + let legacy = service + .handle( + &mut legacy_state, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"tools/call", + "params":{ + "name":"icm_transcript_show", + "arguments":{"session_id":session_id,"offset":0} + } + })), + ) + .unwrap(); + assert!(legacy.error.is_none()); + let legacy_result = legacy.result.unwrap(); + assert_ne!(legacy_result["isError"], true); + assert!(legacy_result["content"][0]["text"] + .as_str() + .unwrap() + .contains(&session_id)); + + let mut initialized_state = ConnectionState::default(); + initialize_2025(&service, &mut initialized_state); + let initialized = service + .handle( + &mut initialized_state, + request(json!({ + "jsonrpc":"2.0","id":3,"method":"tools/call", + "params":{ + "name":"icm_transcript_show", + "arguments":{"session_id":session_id,"offset":0} + } + })), + ) + .unwrap(); + assert!(initialized.result.is_none()); + assert_eq!(initialized.error.unwrap().code, -32602); + + let mut modern_state = ConnectionState::default(); + let modern = service + .handle( + &mut modern_state, + request(json!({ + "jsonrpc":"2.0","id":4,"method":"tools/call", + "params":{ + "name":"icm_transcript_show", + "arguments":{"session_id":session_id,"offset":0}, + "_meta":modern_metadata() + } + })), + ) + .unwrap(); + assert!(modern.result.is_none()); + assert_eq!(modern.error.unwrap().code, -32602); + } +} diff --git a/crates/icm-mcp/src/tools.rs b/crates/icm-mcp/src/tools.rs index 2612ffd2..7b18ac2d 100644 --- a/crates/icm-mcp/src/tools.rs +++ b/crates/icm-mcp/src/tools.rs @@ -301,7 +301,7 @@ pub(crate) fn build_catalog(has_embedder: bool) -> ToolCatalog { legacy_normalizer: normalize_legacy_recall_input, requirements: ToolRequirements::STORE.with_optional_embedder(), ToolAnnotations::new(false, true, false, false), - |context, args| tool_recall(context.store, context.embedder, args, context.compact) + tool_recall ), tool_spec!( MemoryForgetInput, @@ -1514,12 +1514,10 @@ fn format_memory_output(memories: &[(Memory, f32)], compact: bool) -> String { output } -fn tool_recall( - store: &Store, - embedder: Option<&dyn Embedder>, - args: &Value, - compact: bool, -) -> ToolResult { +fn tool_recall(context: &ToolContext<'_>, args: &Value) -> ToolResult { + let store = context.store; + let embedder = context.embedder; + let compact = context.compact; // Auto-decay if >24h since last decay if let Err(e) = store.maybe_auto_decay() { tracing::warn!(error = %e, "auto-decay failed during recall"); @@ -1544,9 +1542,8 @@ fn tool_recall( // (git remote first) — the CLI hooks store under that name, so a raw // cwd basename would silently miss on renamed checkouts (audit finding). let project_arg = get_str(args, "project"); - let cwd_project = std::env::current_dir() - .ok() - .and_then(|p| icm_core::project::project_from_path(&p.to_string_lossy())); + let cwd_project = + icm_core::project::project_from_path(&context.working_directory.to_string_lossy()); let project: Option = match project_arg { Some("") => None, Some(p) => Some(p.to_string()), From 6f3ee55e0f0035c919084489e583e73225094681 Mon Sep 17 00:00:00 2001 From: Lauritz-Timm Date: Sun, 9 Aug 2026 23:22:28 +0200 Subject: [PATCH 04/20] fix(mcp): preserve legacy tool dispatch --- crates/icm-mcp/src/catalog.rs | 16 +++- crates/icm-mcp/src/service.rs | 160 ++++++++++++++++++++++++++++++---- 2 files changed, 158 insertions(+), 18 deletions(-) diff --git a/crates/icm-mcp/src/catalog.rs b/crates/icm-mcp/src/catalog.rs index 9bbd995f..0a960033 100644 --- a/crates/icm-mcp/src/catalog.rs +++ b/crates/icm-mcp/src/catalog.rs @@ -367,12 +367,13 @@ pub enum DispatchResult { pub struct ToolCatalog { registrations: Vec, by_name: HashMap<&'static str, usize>, + has_embedder: bool, legacy_list: Value, modern_list: Value, } impl ToolCatalog { - pub fn new(mut registrations: Vec, has_embedder: bool) -> Result { + pub fn new(registrations: Vec, has_embedder: bool) -> Result { let mut names = HashSet::with_capacity(registrations.len()); for registration in ®istrations { if !names.insert(registration.name) { @@ -383,7 +384,6 @@ impl ToolCatalog { } } - registrations.retain(|registration| registration.requirements.is_available(has_embedder)); let by_name = registrations .iter() .enumerate() @@ -392,17 +392,22 @@ impl ToolCatalog { let legacy_tools: Vec = registrations .iter() - .filter(|registration| registration.requirements.legacy_visible) + .filter(|registration| { + registration.requirements.legacy_visible + && registration.requirements.is_available(has_embedder) + }) .map(ToolSpec::legacy_definition) .collect(); let modern_tools: Vec = registrations .iter() + .filter(|registration| registration.requirements.is_available(has_embedder)) .map(ToolSpec::modern_definition) .collect(); Ok(Self { registrations, by_name, + has_embedder, legacy_list: json!({ "tools": legacy_tools }), modern_list: json!({ "tools": modern_tools }), }) @@ -427,6 +432,11 @@ impl ToolCatalog { return DispatchResult::UnknownTool; }; let registration = &self.registrations[*index]; + if validation == InputValidation::Modern + && !registration.requirements.is_available(self.has_embedder) + { + return DispatchResult::UnknownTool; + } let normalized_arguments = matches!( validation, InputValidation::Legacy2024Unchecked | InputValidation::Legacy2024 diff --git a/crates/icm-mcp/src/service.rs b/crates/icm-mcp/src/service.rs index d5a9c376..60923c77 100644 --- a/crates/icm-mcp/src/service.rs +++ b/crates/icm-mcp/src/service.rs @@ -229,7 +229,11 @@ impl<'a> McpService<'a> { ProtocolEra::PerRequest, ); } - if let Err(response) = validate_legacy_request(id.clone(), message) { + if let Err(response) = validate_legacy_request( + id.clone(), + message, + revision == ProtocolRevision::V2024_11_05, + ) { return *response; } if method == "ping" { @@ -252,7 +256,11 @@ impl<'a> McpService<'a> { ProtocolEra::PerRequest, ) } else { - if let Err(response) = validate_legacy_request(id.clone(), message) { + if let Err(response) = validate_legacy_request( + id.clone(), + message, + revision == ProtocolRevision::V2024_11_05, + ) { return *response; } self.dispatch(state, id, revision, method, message) @@ -334,7 +342,7 @@ impl<'a> McpService<'a> { id: Value, message: &JsonRpcMessage, ) -> JsonRpcResponse { - if let Err(response) = validate_legacy_request(id.clone(), message) { + if let Err(response) = validate_legacy_request(id.clone(), message, false) { return *response; } let Some(params) = message.params.as_ref().and_then(Value::as_object) else { @@ -525,9 +533,12 @@ impl<'a> McpService<'a> { revision: ProtocolRevision, message: &JsonRpcMessage, ) -> JsonRpcResponse { - let Some(params) = message.params.as_ref().and_then(Value::as_object) else { + let Some(params) = message.params.as_ref() else { return JsonRpcResponse::err(id, -32602, "missing params".into()); }; + if revision != ProtocolRevision::V2024_11_05 && !params.is_object() { + return JsonRpcResponse::err(id, -32602, "missing params".into()); + } let Some(name) = params.get("name").and_then(Value::as_str) else { return JsonRpcResponse::err(id, -32602, "missing tool name".into()); }; @@ -535,7 +546,7 @@ impl<'a> McpService<'a> { .get("arguments") .cloned() .unwrap_or_else(|| json!({})); - if !arguments.is_object() { + if revision != ProtocolRevision::V2024_11_05 && !arguments.is_object() { return JsonRpcResponse::err(id, -32602, "tool arguments must be an object".into()); } @@ -551,10 +562,10 @@ impl<'a> McpService<'a> { compact: self.compact, auto_consolidate: self.auto_consolidate, working_directory: &self.working_directory, - enforce_directory_boundary: true, + enforce_directory_boundary: revision != ProtocolRevision::V2024_11_05, }; let validation = if revision == ProtocolRevision::V2024_11_05 { - InputValidation::Legacy2024 + InputValidation::Legacy2024Unchecked } else { InputValidation::Modern }; @@ -676,6 +687,7 @@ fn requests_modern_era(message: &JsonRpcMessage) -> bool { fn validate_legacy_request( id: Value, message: &JsonRpcMessage, + allow_non_object_params: bool, ) -> Result<(), Box> { if message.extra.contains_key("_meta") { return Err(invalid_params( @@ -687,6 +699,9 @@ fn validate_legacy_request( return Ok(()); }; let Some(params) = params.as_object() else { + if allow_non_object_params { + return Ok(()); + } return Err(invalid_params(id, "request params must be an object")); }; let Some(raw_metadata) = params.get("_meta") else { @@ -2625,11 +2640,7 @@ mod tests { if revision == ProtocolRevision::V2024_11_05 { assert!(response.error.is_none()); let result = response.result.unwrap(); - assert_eq!(result["isError"], true); - assert!(result["content"][0]["text"] - .as_str() - .unwrap() - .starts_with("invalid arguments: ")); + assert_ne!(result["isError"], true); } else { assert!(response.result.is_none()); let error = response.error.unwrap(); @@ -2861,7 +2872,16 @@ mod tests { })), ) .unwrap(); - assert_eq!(legacy_bad_type.result.unwrap()["isError"], true); + let legacy_bad_type = legacy_bad_type.result.unwrap(); + assert_ne!(legacy_bad_type["isError"], true); + assert_eq!( + legacy_bad_type["content"][0]["text"] + .as_str() + .unwrap() + .matches("revision limit probe") + .count(), + 5 + ); let mut modern_state = ConnectionState::default(); let accepted = service @@ -2928,6 +2948,24 @@ mod tests { .unwrap(); assert!(modern_unknown.result.is_none()); assert_eq!(modern_unknown.error.unwrap().code, -32602); + + let modern_bad_type = service + .handle( + &mut modern_state, + request(json!({ + "jsonrpc":"2.0","id":11,"method":"tools/call", + "params":{ + "name":"icm_memory_recall", + "arguments":{ + "query":"revision limit probe","project":"","limit":"2" + }, + "_meta":modern_metadata() + } + })), + ) + .unwrap(); + assert!(modern_bad_type.result.is_none()); + assert_eq!(modern_bad_type.error.unwrap().code, -32602); } #[test] @@ -2993,10 +3031,34 @@ mod tests { "isError":true }) ); + + let non_object_arguments = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":3,"method":"tools/call", + "params":{"name":"icm_memory_stats","arguments":[]} + })), + ) + .unwrap(); + assert!(non_object_arguments.error.is_none()); + assert_ne!(non_object_arguments.result.unwrap()["isError"], true); + + let non_object_params = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":4,"method":"tools/call","params":[] + })), + ) + .unwrap(); + let error = non_object_params.error.unwrap(); + assert_eq!(error.code, -32602); + assert_eq!(error.message, "missing tool name"); } #[test] - fn unavailable_embedder_tool_stays_hidden_from_service_listing_and_dispatch() { + fn unavailable_embedder_tool_stays_hidden_but_preserves_legacy_dispatch() { let store = Store::in_memory().unwrap(); let service = service(&store); let mut state = initialized_state_for_revision(&service, ProtocolRevision::V2024_11_05); @@ -3029,11 +3091,79 @@ mod tests { json!({ "content":[{ "type":"text", - "text":"unknown tool: icm_memory_embed_all" + "text":"embeddings not available" }], "isError":true }) ); + + let mut modern = initialized_state_for_revision(&service, ProtocolRevision::V2025_11_25); + let modern_called = service + .handle( + &mut modern, + request(json!({ + "jsonrpc":"2.0","id":4,"method":"tools/call", + "params":{"name":"icm_memory_embed_all","arguments":{}} + })), + ) + .unwrap(); + assert_eq!(modern_called.error.unwrap().code, -32602); + } + + #[test] + fn legacy_learn_keeps_caller_selected_paths_while_modern_stays_bounded() { + let root = tempfile::tempdir().unwrap(); + let working_directory = root.path().join("server-project"); + let external_directory = root.path().join("external-project"); + std::fs::create_dir(&working_directory).unwrap(); + std::fs::create_dir(&external_directory).unwrap(); + std::fs::write( + external_directory.join("Cargo.toml"), + "[package]\nname='external-project'\nversion='0.1.0'\n", + ) + .unwrap(); + + let store = Store::in_memory().unwrap(); + let service = McpService::with_working_directory( + &store, + None, + false, + AutoConsolidate::default(), + working_directory, + ); + let arguments = json!({"directory":external_directory}); + + let mut legacy = initialized_state_for_revision(&service, ProtocolRevision::V2024_11_05); + let accepted = service + .handle( + &mut legacy, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"tools/call", + "params":{"name":"icm_learn","arguments":arguments} + })), + ) + .unwrap() + .result + .unwrap(); + assert_ne!(accepted["isError"], true); + + let mut modern = initialized_state_for_revision(&service, ProtocolRevision::V2025_11_25); + let rejected = service + .handle( + &mut modern, + request(json!({ + "jsonrpc":"2.0","id":3,"method":"tools/call", + "params":{"name":"icm_learn","arguments":arguments} + })), + ) + .unwrap() + .result + .unwrap(); + assert_eq!(rejected["isError"], true); + assert!(rejected["content"][0]["text"] + .as_str() + .unwrap() + .contains("within the server working directory")); } #[test] From b4a578957b9d945f0530c49f59ef739aaf1a8a2b Mon Sep 17 00:00:00 2001 From: Lauritz-Timm Date: Wed, 12 Aug 2026 12:32:08 +0200 Subject: [PATCH 05/20] refactor(mcp): split tool registry handlers and tests --- crates/icm-mcp/src/tools.rs | 4622 +---------------- crates/icm-mcp/src/tools/handlers.rs | 26 + crates/icm-mcp/src/tools/handlers/common.rs | 172 + crates/icm-mcp/src/tools/handlers/feedback.rs | 153 + crates/icm-mcp/src/tools/handlers/memoir.rs | 635 +++ crates/icm-mcp/src/tools/handlers/memory.rs | 825 +++ .../icm-mcp/src/tools/handlers/transcript.rs | 101 + crates/icm-mcp/src/tools/registry.rs | 944 ++++ crates/icm-mcp/src/tools/tests.rs | 1839 +++++++ 9 files changed, 4703 insertions(+), 4614 deletions(-) create mode 100644 crates/icm-mcp/src/tools/handlers.rs create mode 100644 crates/icm-mcp/src/tools/handlers/common.rs create mode 100644 crates/icm-mcp/src/tools/handlers/feedback.rs create mode 100644 crates/icm-mcp/src/tools/handlers/memoir.rs create mode 100644 crates/icm-mcp/src/tools/handlers/memory.rs create mode 100644 crates/icm-mcp/src/tools/handlers/transcript.rs create mode 100644 crates/icm-mcp/src/tools/registry.rs create mode 100644 crates/icm-mcp/src/tools/tests.rs diff --git a/crates/icm-mcp/src/tools.rs b/crates/icm-mcp/src/tools.rs index 7b18ac2d..b7948d9a 100644 --- a/crates/icm-mcp/src/tools.rs +++ b/crates/icm-mcp/src/tools.rs @@ -1,1038 +1,16 @@ -use chrono::Utc; -use serde_json::{json, Value}; +use serde_json::Value; -use icm_core::{ - add_backrefs, auto_link_memory, build_wake_up, find_similar_memory, format_local, - is_preference_topic, keyword_matches, project_matches, topic_matches, AutoLinkOptions, Concept, - ConceptLink, Embedder, Feedback, FeedbackStore, Label, Memoir, MemoirStore, Memory, - MemoryStore, Relation, WakeUpFormat, WakeUpOptions, DEDUP_SIMILARITY_THRESHOLD, - MSG_NO_MEMORIES, -}; +use icm_core::Embedder; use icm_store::Store; -use crate::catalog::{ - DispatchResult, ToolAnnotations, ToolCatalog, ToolContext, ToolRequirements, ToolSpec, -}; -use crate::inputs::*; +use crate::catalog::{DispatchResult, ToolContext}; use crate::protocol::ToolResult; -/// Historical default threshold for auto-consolidation. The live value comes -/// from [`AutoConsolidate`] (issue #318); this constant is only the fallback -/// for callers that don't pass a policy. -const AUTO_CONSOLIDATE_THRESHOLD: usize = 10; +mod handlers; +mod registry; -/// Auto-consolidation policy for the MCP store path (issue #318). -/// -/// Previously the MCP `icm_memory_store` handler consolidated a topic past a -/// hardcoded 10 entries **unconditionally**, ignoring `[memory] -/// auto_consolidate_enabled` / `auto_consolidate_threshold` — so an explicit -/// `enabled = false` still destructively rolled up (and deleted) a topic's -/// memories. `icm serve` now threads the loaded config through as one of -/// these, and the handler honors it. -#[derive(Clone, Copy, Debug)] -pub struct AutoConsolidate { - pub enabled: bool, - pub threshold: usize, -} - -impl Default for AutoConsolidate { - /// The historical always-on behavior (threshold 10). Used only by callers - /// that don't supply a policy — e.g. tests via [`call_tool`]. The - /// `icm serve` path passes the user's real config through - /// [`call_tool_with_config`] instead. - fn default() -> Self { - Self { - enabled: true, - threshold: AUTO_CONSOLIDATE_THRESHOLD, - } - } -} - -/// Maximum allowed UTF-8 byte length for topic names. Must stay <= the -/// store layer's `MAX_TOPIC_BYTES` so the MCP-level rejection happens -/// *before* the store's lower-level validation does. -const MAX_TOPIC_LEN: usize = 255; - -/// Maximum allowed length for content/summary text. Aligned with the -/// store layer's `MAX_SUMMARY_BYTES` (64 KB). Letting MCP accept -/// larger inputs only to have the store reject them would be -/// confusing — fail fast at the API surface. -const MAX_CONTENT_LEN: usize = 64 * 1024; - -/// `icm_feedback_record`'s context/predicted/corrected/reason had no length -/// cap at all, unlike icm_memory_store's MAX_CONTENT_LEN (audit finding). -const MAX_FEEDBACK_FIELD_LEN: usize = 20_000; - -/// Parse a JSON keywords array from tool arguments. -fn parse_keywords(args: &Value) -> Vec { - args.get("keywords") - .and_then(|v| v.as_array()) - .map(|arr| { - arr.iter() - .filter_map(|v| v.as_str().map(String::from)) - .collect() - }) - .unwrap_or_default() -} - -/// Try to auto-consolidate a topic if the policy is enabled and the topic -/// exceeds the configured threshold (issue #318). Returns a human-readable -/// message if consolidation happened, or an empty string (including when the -/// policy is disabled — a no-op). -/// -/// Routes through `auto_consolidate_with_embedder` so the consolidated -/// memory is embedded inline (closes audit M2/AC2: previously the -/// rolled-up memory had `embedding = None` and was invisible to hybrid -/// recall until a manual `icm embed` rebuilt it). -fn try_auto_consolidate( - store: &Store, - embedder: Option<&dyn Embedder>, - topic: &str, - auto: AutoConsolidate, -) -> String { - if !auto.enabled { - return String::new(); - } - match store.auto_consolidate_with_embedder(topic, auto.threshold, embedder) { - Ok(true) => format!( - "Auto-consolidated topic '{topic}' (exceeded {} entries).", - auto.threshold - ), - Ok(false) => String::new(), - Err(e) => { - tracing::warn!("auto-consolidation failed for topic '{topic}': {e}"); - String::new() - } - } -} - -// --------------------------------------------------------------------------- -// Tool schemas for tools/list -// --------------------------------------------------------------------------- - -fn normalize_legacy_recall_input(arguments: &Value) -> Value { - let mut normalized = arguments.clone(); - let Some(object) = normalized.as_object_mut() else { - return normalized; - }; - let Some(limit) = object.get("limit").filter(|limit| limit.is_number()) else { - return normalized; - }; - - // The frozen 2024 handler read limits as i64, defaulted unrepresentable - // numeric values to five, and clamped the result to its advertised 1..20 - // range. Normalize only for that catalog projection; the modern DTO keeps - // its strict 1..100 contract and reaches the handler unchanged. - let normalized_limit = limit.as_i64().unwrap_or(5).clamp(1, 20); - object.insert("limit".into(), json!(normalized_limit)); - normalized -} - -macro_rules! tool_spec { - ( - $input:ty, - json!({ - "name": $name:literal, - "description": $description:literal, - "inputSchema": $input_schema:tt - }), - $annotations:expr, - $handler:expr - ) => { - ToolSpec::typed::<$input>( - $name, - $description, - json!($input_schema), - None, - $annotations, - ToolRequirements::STORE, - $handler, - ) - }; - ( - $input:ty, - json!({ - "name": $name:literal, - "description": $description:literal, - "inputSchema": $input_schema:tt - }), - legacy_normalizer: $legacy_normalizer:expr, - $annotations:expr, - $handler:expr - ) => { - ToolSpec::typed::<$input>( - $name, - $description, - json!($input_schema), - Some($legacy_normalizer), - $annotations, - ToolRequirements::STORE, - $handler, - ) - }; - ( - $input:ty, - json!({ - "name": $name:literal, - "description": $description:literal, - "inputSchema": $input_schema:tt - }), - requirements: $requirements:expr, - $annotations:expr, - $handler:expr - ) => { - ToolSpec::typed::<$input>( - $name, - $description, - json!($input_schema), - None, - $annotations, - $requirements, - $handler, - ) - }; - ( - $input:ty, - json!({ - "name": $name:literal, - "description": $description:literal, - "inputSchema": $input_schema:tt - }), - legacy_normalizer: $legacy_normalizer:expr, - requirements: $requirements:expr, - $annotations:expr, - $handler:expr - ) => { - ToolSpec::typed::<$input>( - $name, - $description, - json!($input_schema), - Some($legacy_normalizer), - $annotations, - $requirements, - $handler, - ) - }; -} - -pub(crate) fn build_catalog(has_embedder: bool) -> ToolCatalog { - let tools = vec![ - // --- Memory tools --- - tool_spec!( - MemoryStoreInput, - json!({ - "name": "icm_memory_store", - "description": "Store important information in ICM long-term memory. Use to save decisions, preferences, project context, resolved errors — anything that should persist between sessions.", - "inputSchema": { - "type": "object", - "properties": { - "topic": { - "type": "string", - "description": "Category/namespace. Use the canonical topics from the server instructions: 'decisions-{project}', 'preferences', 'errors-resolved', 'context-{project}' — mixed-language topic names fragment the memory." - }, - "content": { - "type": "string", - "description": "Information to memorize — be concise but complete" - }, - "importance": { - "type": "string", - "enum": ["critical", "high", "medium", "low"], - "default": "medium", - "description": "critical=never forgotten, high=slow decay, medium=normal, low=fast decay" - }, - "keywords": { - "type": "array", - "items": { "type": "string" }, - "description": "Keywords to improve search" - }, - "raw_excerpt": { - "type": "string", - "description": "Optional verbatim (code, exact error message, etc.)" - } - }, - "required": ["topic", "content"] - } - }), - requirements: ToolRequirements::STORE.with_optional_embedder(), - ToolAnnotations::new(false, true, false, false), - |context, args| tool_store( - context.store, - context.embedder, - args, - context.compact, - context.auto_consolidate - ) - ), - tool_spec!( - MemoryRecallInput, - json!({ - "name": "icm_memory_recall", - "description": "Search ICM long-term memory. Use to find past decisions, project context, preferences, or solutions to previously encountered problems.", - "inputSchema": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "Natural language search query" - }, - "topic": { - "type": "string", - "description": "Filter by specific topic (optional)" - }, - "limit": { - "type": "integer", - "default": 5, - "minimum": 1, - "maximum": 20, - "description": "Max number of results" - }, - "keyword": { - "type": "string", - "description": "Filter results by keyword (exact match on memory keywords)" - }, - "project": { - "type": "string", - "description": "Project filter (segment-aware). Defaults to the server's cwd directory name. Pass an empty string to disable the filter and search across all projects." - } - }, - "required": ["query"] - } - }), - legacy_normalizer: normalize_legacy_recall_input, - requirements: ToolRequirements::STORE.with_optional_embedder(), - ToolAnnotations::new(false, true, false, false), - tool_recall - ), - tool_spec!( - MemoryForgetInput, - json!({ - "name": "icm_memory_forget", - "description": "Delete a specific memory by its ID. Use when information is obsolete or incorrect.", - "inputSchema": { - "type": "object", - "properties": { - "id": { - "type": "string", - "description": "Memory ID to delete" - } - }, - "required": ["id"] - } - }), - ToolAnnotations::new(false, true, true, false), - |context, args| tool_forget(context.store, args) - ), - tool_spec!( - TopicInput, - json!({ - "name": "icm_memory_forget_topic", - "description": "Delete ALL memories in a topic. Use to clear an entire topic at once.", - "inputSchema": { - "type": "object", - "properties": { - "topic": { - "type": "string", - "description": "Topic whose memories should all be deleted" - } - }, - "required": ["topic"] - } - }), - ToolAnnotations::new(false, true, true, false), - |context, args| tool_forget_topic(context.store, args) - ), - tool_spec!( - LearnInput, - json!({ - "name": "icm_learn", - "description": "Scan a project directory and create a Memoir knowledge graph with its structure, dependencies, modules, and config files.", - "inputSchema": { - "type": "object", - "properties": { - "directory": { - "type": "string", - "description": "Project directory to scan (default: current working directory)" - }, - "name": { - "type": "string", - "description": "Memoir name (default: directory name)" - } - } - } - }), - requirements: ToolRequirements::STORE.with_filesystem_read(), - ToolAnnotations::new(false, true, false, true), - tool_learn_bounded - ), - tool_spec!( - MemoryConsolidateInput, - json!({ - "name": "icm_memory_consolidate", - "description": "Consolidate all memories of a topic into a single summary. Useful when a topic accumulates too many entries.", - "inputSchema": { - "type": "object", - "properties": { - "topic": { - "type": "string", - "description": "Topic to consolidate" - }, - "summary": { - "type": "string", - "description": "Consolidated summary to replace all memories in the topic" - } - }, - "required": ["topic", "summary"] - } - }), - requirements: ToolRequirements::STORE.with_optional_embedder(), - ToolAnnotations::new(false, true, false, false), - |context, args| tool_consolidate(context.store, context.embedder, args) - ), - tool_spec!( - MemoryListTopicsInput, - json!({ - "name": "icm_memory_list_topics", - "description": "List all available topics in memory with their counts.", - "inputSchema": { - "type": "object", - "properties": {} - } - }), - ToolAnnotations::new(true, false, true, false), - |context, _| tool_list_topics(context.store) - ), - tool_spec!( - MemoryStatsInput, - json!({ - "name": "icm_memory_stats", - "description": "Get global ICM memory statistics.", - "inputSchema": { - "type": "object", - "properties": {} - } - }), - ToolAnnotations::new(true, false, true, false), - |context, _| tool_stats(context.store) - ), - tool_spec!( - MemoryUpdateInput, - json!({ - "name": "icm_memory_update", - "description": "Update an existing memory in-place. Use to correct, refresh, or extend a memory without creating a duplicate.", - "inputSchema": { - "type": "object", - "properties": { - "id": { - "type": "string", - "description": "Memory ID to update" - }, - "content": { - "type": "string", - "description": "New content (replaces existing summary)" - }, - "importance": { - "type": "string", - "enum": ["critical", "high", "medium", "low"], - "description": "New importance level (optional, keeps existing if not set)" - }, - "keywords": { - "type": "array", - "items": { "type": "string" }, - "description": "New keywords (optional, keeps existing if not set)" - } - }, - "required": ["id", "content"] - } - }), - requirements: ToolRequirements::STORE.with_optional_embedder(), - ToolAnnotations::new(false, true, false, false), - |context, args| tool_update(context.store, context.embedder, args) - ), - tool_spec!( - MemoryHealthInput, - json!({ - "name": "icm_memory_health", - "description": "Get health stats for all topics: entry count, staleness, consolidation needs. Use to audit memory hygiene.", - "inputSchema": { - "type": "object", - "properties": { - "topic": { - "type": "string", - "description": "Check a specific topic (optional — checks all if omitted)" - } - } - } - }), - ToolAnnotations::new(true, false, true, false), - |context, args| tool_health(context.store, args) - ), - // --- Memoir tools --- - tool_spec!( - MemoirCreateInput, - json!({ - "name": "icm_memoir_create", - "description": "Create a new memoir — a permanent knowledge container. Memoirs hold concepts that never decay.", - "inputSchema": { - "type": "object", - "properties": { - "name": { - "type": "string", - "description": "Unique human-readable name for the memoir" - }, - "description": { - "type": "string", - "description": "Description of what this memoir is for" - } - }, - "required": ["name"] - } - }), - ToolAnnotations::new(false, false, false, false), - |context, args| tool_memoir_create(context.store, args) - ), - tool_spec!( - MemoirListInput, - json!({ - "name": "icm_memoir_list", - "description": "List all memoirs with their concept counts.", - "inputSchema": { - "type": "object", - "properties": {} - } - }), - ToolAnnotations::new(true, false, true, false), - |context, _| tool_memoir_list(context.store) - ), - tool_spec!( - NameInput, - json!({ - "name": "icm_memoir_show", - "description": "Show a memoir's stats, labels, and all its concepts.", - "inputSchema": { - "type": "object", - "properties": { - "name": { - "type": "string", - "description": "Memoir name" - } - }, - "required": ["name"] - } - }), - ToolAnnotations::new(true, false, true, false), - |context, args| tool_memoir_show(context.store, args) - ), - tool_spec!( - MemoirAddConceptInput, - json!({ - "name": "icm_memoir_add_concept", - "description": "Add a permanent concept to a memoir. Concepts are knowledge nodes that get refined, never decayed.", - "inputSchema": { - "type": "object", - "properties": { - "memoir": { - "type": "string", - "description": "Memoir name" - }, - "name": { - "type": "string", - "description": "Concept name (unique within memoir)" - }, - "definition": { - "type": "string", - "description": "Dense description of the concept" - }, - "labels": { - "type": "string", - "description": "Comma-separated labels (namespace:value or plain tag). E.g. 'domain:arch,type:decision'" - } - }, - "required": ["memoir", "name", "definition"] - } - }), - ToolAnnotations::new(false, false, false, false), - |context, args| tool_memoir_add_concept(context.store, args) - ), - tool_spec!( - MemoirRefineInput, - json!({ - "name": "icm_memoir_refine", - "description": "Refine an existing concept with a new, improved definition. Bumps revision and boosts confidence.", - "inputSchema": { - "type": "object", - "properties": { - "memoir": { - "type": "string", - "description": "Memoir name" - }, - "name": { - "type": "string", - "description": "Concept name" - }, - "definition": { - "type": "string", - "description": "New, refined definition" - } - }, - "required": ["memoir", "name", "definition"] - } - }), - ToolAnnotations::new(false, true, false, false), - |context, args| tool_memoir_refine(context.store, args) - ), - tool_spec!( - MemoirSearchInput, - json!({ - "name": "icm_memoir_search", - "description": "Full-text search concepts within a memoir.", - "inputSchema": { - "type": "object", - "properties": { - "memoir": { - "type": "string", - "description": "Memoir name" - }, - "query": { - "type": "string", - "description": "Search query" - }, - "label": { - "type": "string", - "description": "Filter by label (e.g. 'domain:tech')" - }, - "limit": { - "type": "integer", - "default": 10, - "description": "Max results" - } - }, - "required": ["memoir", "query"] - } - }), - ToolAnnotations::new(true, false, true, false), - |context, args| tool_memoir_search(context.store, args) - ), - tool_spec!( - MemoirLinkInput, - json!({ - "name": "icm_memoir_link", - "description": "Create a directed, typed edge between two concepts in the same memoir.", - "inputSchema": { - "type": "object", - "properties": { - "memoir": { - "type": "string", - "description": "Memoir name" - }, - "from": { - "type": "string", - "description": "Source concept name" - }, - "to": { - "type": "string", - "description": "Target concept name" - }, - "relation": { - "type": "string", - "enum": ["part_of", "depends_on", "related_to", "contradicts", "refines", "alternative_to", "caused_by", "instance_of", "superseded_by"], - "description": "Relation type" - } - }, - "required": ["memoir", "from", "to", "relation"] - } - }), - ToolAnnotations::new(false, false, false, false), - |context, args| tool_memoir_link(context.store, args) - ), - tool_spec!( - MemoirInspectInput, - json!({ - "name": "icm_memoir_inspect", - "description": "Inspect a concept and its graph neighborhood (BFS).", - "inputSchema": { - "type": "object", - "properties": { - "memoir": { - "type": "string", - "description": "Memoir name" - }, - "name": { - "type": "string", - "description": "Concept name" - }, - "depth": { - "type": "integer", - "default": 1, - "description": "BFS depth" - } - }, - "required": ["memoir", "name"] - } - }), - ToolAnnotations::new(true, false, true, false), - |context, args| tool_memoir_inspect(context.store, args) - ), - tool_spec!( - MemoirExportInput, - json!({ - "name": "icm_memoir_export", - "description": "Export a memoir's full concept graph. Formats: json (structured), dot (Graphviz), ascii (visual), ai (compact markdown for LLM context).", - "inputSchema": { - "type": "object", - "properties": { - "name": { - "type": "string", - "description": "Memoir name" - }, - "format": { - "type": "string", - "enum": ["json", "dot", "ascii", "ai"], - "default": "json", - "description": "Output format: json (structured), dot (Graphviz), ascii (visual graph), ai (compact markdown for LLM)" - } - }, - "required": ["name"] - } - }), - ToolAnnotations::new(true, false, true, false), - |context, args| tool_memoir_export(context.store, args) - ), - tool_spec!( - ExtractPatternsInput, - json!({ - "name": "icm_memory_extract_patterns", - "description": "Detect recurring patterns in a topic by keyword similarity. Optionally create concepts in a memoir from detected patterns.", - "inputSchema": { - "type": "object", - "properties": { - "topic": { - "type": "string", - "description": "Topic to analyze for patterns" - }, - "memoir": { - "type": "string", - "description": "Memoir name — if provided, creates concepts from detected patterns" - }, - "min_cluster_size": { - "type": "integer", - "default": 3, - "minimum": 2, - "description": "Minimum number of similar memories to form a pattern (default: 3)" - } - }, - "required": ["topic"] - } - }), - ToolAnnotations::new(false, false, false, false), - |context, args| tool_extract_patterns(context.store, args) - ), - tool_spec!( - MemoirSearchAllInput, - json!({ - "name": "icm_memoir_search_all", - "description": "Full-text search concepts across all memoirs.", - "inputSchema": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "Search query" - }, - "limit": { - "type": "integer", - "default": 10, - "description": "Max results" - } - }, - "required": ["query"] - } - }), - ToolAnnotations::new(true, false, true, false), - |context, args| tool_memoir_search_all(context.store, args) - ), - // --- Feedback tools --- - tool_spec!( - FeedbackRecordInput, - json!({ - "name": "icm_feedback_record", - "description": "Record a correction/feedback when an AI prediction was wrong. Helps improve future predictions by learning from mistakes.", - "inputSchema": { - "type": "object", - "properties": { - "topic": { - "type": "string", - "description": "Category/namespace for this feedback (e.g. 'triage-owner/repo', 'pr-analysis')" - }, - "context": { - "type": "string", - "description": "What was the situation / input that led to the prediction" - }, - "predicted": { - "type": "string", - "description": "What the AI predicted or did" - }, - "corrected": { - "type": "string", - "description": "What the correct answer/action should have been" - }, - "reason": { - "type": "string", - "description": "Why the correction was made (optional)" - }, - "source": { - "type": "string", - "description": "Which tool/pipeline generated the prediction (optional)" - } - }, - "required": ["topic", "context", "predicted", "corrected"] - } - }), - requirements: ToolRequirements::STORE.with_optional_embedder(), - ToolAnnotations::new(false, false, false, false), - |context, args| tool_feedback_record( - context.store, - context.embedder, - args, - context.compact - ) - ), - tool_spec!( - FeedbackSearchInput, - json!({ - "name": "icm_feedback_search", - "description": "Search past feedback/corrections to inform current predictions. Use before making predictions to learn from past mistakes.", - "inputSchema": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "Search query to find relevant past corrections" - }, - "topic": { - "type": "string", - "description": "Filter by topic (optional)" - }, - "limit": { - "type": "integer", - "default": 5, - "minimum": 1, - "maximum": 20, - "description": "Max number of results" - } - }, - "required": ["query"] - } - }), - requirements: ToolRequirements::STORE.with_optional_embedder(), - ToolAnnotations::new(true, false, true, false), - |context, args| tool_feedback_search(context.store, context.embedder, args) - ), - tool_spec!( - FeedbackStatsInput, - json!({ - "name": "icm_feedback_stats", - "description": "Get feedback statistics: total count, breakdown by topic, most applied corrections.", - "inputSchema": { - "type": "object", - "properties": {} - } - }), - ToolAnnotations::new(true, false, true, false), - |context, _| tool_feedback_stats(context.store) - ), - // --- Transcript tools (verbatim session replay) --- - tool_spec!( - TranscriptStartInput, - json!({ - "name": "icm_transcript_start_session", - "description": "Create a new transcript session for verbatim message capture. Returns the session_id used by subsequent icm_transcript_record calls. Use once per conversation or debugging session.", - "inputSchema": { - "type": "object", - "properties": { - "agent": { - "type": "string", - "description": "Agent identifier (e.g. 'claude-code', 'cursor', 'gemini-cli'). Default: 'mcp'." - }, - "project": { - "type": "string", - "description": "Project name (optional; usually cwd basename or repo slug)" - }, - "metadata": { - "type": "string", - "description": "Arbitrary JSON metadata (optional)" - } - } - } - }), - ToolAnnotations::new(false, false, false, false), - |context, args| tool_transcript_start_session(context.store, args) - ), - tool_spec!( - TranscriptRecordInput, - json!({ - "name": "icm_transcript_record", - "description": "Append a verbatim message to a transcript session. Stores the raw content with no summarization. Use once per user turn, assistant reply, or tool call for full replay fidelity.", - "inputSchema": { - "type": "object", - "properties": { - "session_id": { - "type": "string", - "description": "Session id from icm_transcript_start_session" - }, - "role": { - "type": "string", - "enum": ["user", "assistant", "system", "tool"], - "description": "Message role" - }, - "content": { - "type": "string", - "description": "Raw message content (stored verbatim)" - }, - "tool_name": { - "type": "string", - "description": "Tool name if role=tool (optional)" - }, - "tokens": { - "type": "integer", - "description": "Token count for billing / stats (optional)" - }, - "metadata": { - "type": "string", - "description": "Arbitrary JSON metadata (optional)" - } - }, - "required": ["session_id", "role", "content"] - } - }), - ToolAnnotations::new(false, false, false, false), - |context, args| tool_transcript_record(context.store, args) - ), - tool_spec!( - TranscriptSearchInput, - json!({ - "name": "icm_transcript_search", - "description": "Full-text search across recorded transcript messages (FTS5 BM25). Supports boolean operators, phrase matches, and prefix queries. Use to recall exact quotes or debug past decisions.", - "inputSchema": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "FTS5 query: 'postgres OR mysql', '\"exact phrase\"', 'auth*'" - }, - "session_id": { - "type": "string", - "description": "Restrict to one session (optional)" - }, - "project": { - "type": "string", - "description": "Restrict to one project (optional)" - }, - "limit": { - "type": "integer", - "default": 10, - "minimum": 1, - "maximum": 50 - } - }, - "required": ["query"] - } - }), - ToolAnnotations::new(true, false, true, false), - |context, args| tool_transcript_search(context.store, args) - ), - tool_spec!( - TranscriptShowInput, - json!({ - "name": "icm_transcript_show", - "description": "Replay the full message thread of a transcript session, chronologically. Returns up to `limit` messages with role, content, tool name, timestamp.", - "inputSchema": { - "type": "object", - "properties": { - "session_id": { "type": "string" }, - "limit": { "type": "integer", "default": 200, "minimum": 1, "maximum": 2000 } - }, - "required": ["session_id"] - } - }), - ToolAnnotations::new(true, false, true, false), - |context, args| tool_transcript_show(context.store, args) - ), - tool_spec!( - TranscriptStatsInput, - json!({ - "name": "icm_transcript_stats", - "description": "Global transcript statistics: session count, message count, total bytes, breakdown by role and agent, top sessions by message count.", - "inputSchema": { - "type": "object", - "properties": {} - } - }), - ToolAnnotations::new(true, false, true, false), - |context, _| tool_transcript_stats(context.store) - ), - tool_spec!( - WakeUpInput, - json!({ - "name": "icm_wake_up", - "description": "Build a compact critical-facts pack for LLM system-prompt injection. Selects critical/high memories (and preferences) optionally scoped by project, ranks by importance × recency × weight, and truncates to a token budget. Use at session start to hydrate an agent with the most load-bearing context.", - "inputSchema": { - "type": "object", - "properties": { - "project": { - "type": "string", - "description": "Project name filter (substring match against topic). Preferences/identity memories are always included." - }, - "max_tokens": { - "type": "integer", - "default": 200, - "minimum": 20, - "maximum": 4000, - "description": "Approximate token budget (1 token ≈ 4 characters)" - }, - "format": { - "type": "string", - "enum": ["markdown", "plain"], - "default": "markdown", - "description": "Output format" - }, - "include_preferences": { - "type": "boolean", - "default": true, - "description": "Include global preferences/identity memories regardless of the project filter" - } - } - } - }), - ToolAnnotations::new(true, false, true, false), - |context, args| tool_wake_up(context.store, args) - ), - tool_spec!( - EmbedAllInput, - json!({ - "name": "icm_memory_embed_all", - "description": "Generate embeddings for all memories that don't have one yet. Use this to backfill vector search capability.", - "inputSchema": { - "type": "object", - "properties": { - "topic": { - "type": "string", - "description": "Only embed memories in this topic (optional)" - } - } - } - }), - requirements: ToolRequirements::STORE.with_required_embedder(), - ToolAnnotations::new(false, false, true, false), - |context, args| tool_embed_all(context.store, context.embedder, args) - ), - ]; - - ToolCatalog::new(tools, has_embedder).expect("static MCP tool registrations must be valid") -} - -// --------------------------------------------------------------------------- -// Tool dispatch -// --------------------------------------------------------------------------- +pub use handlers::AutoConsolidate; +pub(crate) use registry::build_catalog; /// Frozen 2024 tool-list projection retained for callers and compatibility /// tests. Production service instances cache this projection in their catalog. @@ -1097,3589 +75,5 @@ pub fn call_tool_with_config( } } -// --------------------------------------------------------------------------- -// Transcript tool handlers -// --------------------------------------------------------------------------- - -fn tool_transcript_start_session(store: &Store, args: &Value) -> ToolResult { - use icm_core::TranscriptStore; - let agent = args.get("agent").and_then(|v| v.as_str()).unwrap_or("mcp"); - let project = args.get("project").and_then(|v| v.as_str()); - let metadata = args.get("metadata").and_then(|v| v.as_str()); - match store.create_session(agent, project, metadata) { - Ok(id) => ToolResult::text(format!("{{\"session_id\":\"{id}\"}}")), - Err(e) => ToolResult::error(format!("start_session failed: {e}")), - } -} - -fn tool_transcript_record(store: &Store, args: &Value) -> ToolResult { - use icm_core::{Role, TranscriptStore}; - let session_id = match args.get("session_id").and_then(|v| v.as_str()) { - Some(s) => s, - None => return ToolResult::error("session_id is required".into()), - }; - let role_str = match args.get("role").and_then(|v| v.as_str()) { - Some(s) => s, - None => return ToolResult::error("role is required".into()), - }; - let role = match Role::parse(role_str) { - Some(r) => r, - None => { - return ToolResult::error(format!( - "invalid role '{role_str}'; must be user|assistant|system|tool" - )) - } - }; - let content = match args.get("content").and_then(|v| v.as_str()) { - Some(s) => s, - None => return ToolResult::error("content is required".into()), - }; - let tool_name = args.get("tool_name").and_then(|v| v.as_str()); - let tokens = args.get("tokens").and_then(|v| v.as_i64()); - let metadata = args.get("metadata").and_then(|v| v.as_str()); - match store.record_message(session_id, role, content, tool_name, tokens, metadata) { - Ok(id) => ToolResult::text(format!("{{\"message_id\":\"{id}\"}}")), - Err(e) => ToolResult::error(format!("record failed: {e}")), - } -} - -fn tool_transcript_search(store: &Store, args: &Value) -> ToolResult { - use icm_core::TranscriptStore; - let query = match args.get("query").and_then(|v| v.as_str()) { - Some(s) => s, - None => return ToolResult::error("query is required".into()), - }; - let session_id = args.get("session_id").and_then(|v| v.as_str()); - let project = args.get("project").and_then(|v| v.as_str()); - let limit = args - .get("limit") - .and_then(|v| v.as_u64()) - .unwrap_or(10) - .min(50) as usize; - match store.search_transcripts(query, session_id, project, limit) { - Ok(hits) => { - let json = serde_json::to_string(&hits).unwrap_or_else(|_| "[]".into()); - ToolResult::text(json) - } - Err(e) => ToolResult::error(format!("search failed: {e}")), - } -} - -fn tool_transcript_show(store: &Store, args: &Value) -> ToolResult { - use icm_core::TranscriptStore; - let session_id = match args.get("session_id").and_then(|v| v.as_str()) { - Some(s) => s, - None => return ToolResult::error("session_id is required".into()), - }; - let limit = args - .get("limit") - .and_then(|v| v.as_u64()) - .unwrap_or(200) - .min(2000) as usize; - let sess = match store.get_session(session_id) { - Ok(Some(s)) => s, - Ok(None) => return ToolResult::error(format!("session {session_id} not found")), - Err(e) => return ToolResult::error(format!("get_session failed: {e}")), - }; - let msgs = match store.list_session_messages(session_id, limit, 0) { - Ok(m) => m, - Err(e) => return ToolResult::error(format!("list_messages failed: {e}")), - }; - let body = json!({ "session": sess, "messages": msgs }); - ToolResult::text(body.to_string()) -} - -fn tool_transcript_stats(store: &Store) -> ToolResult { - use icm_core::TranscriptStore; - match store.transcript_stats() { - Ok(s) => ToolResult::text(serde_json::to_string(&s).unwrap_or_else(|_| "{}".into())), - Err(e) => ToolResult::error(format!("stats failed: {e}")), - } -} - -// --------------------------------------------------------------------------- -// Wake-up tool handler -// --------------------------------------------------------------------------- - -fn tool_wake_up(store: &Store, args: &Value) -> ToolResult { - // Normalize the project filter: empty string or "-" both mean "disabled", - // mirroring the CLI convention. - let project = match get_str(args, "project") { - Some("") | Some("-") => None, - other => other, - }; - // Clamp token budget to [20, 4000] to guard against accidental blowups. - let max_tokens = get_i64(args, "max_tokens", 200).clamp(20, 4000) as usize; - let format = match get_str(args, "format").unwrap_or("markdown") { - "plain" => WakeUpFormat::Plain, - _ => WakeUpFormat::Markdown, - }; - let include_preferences = args - .get("include_preferences") - .and_then(|v| v.as_bool()) - .unwrap_or(true); - - let opts = WakeUpOptions { - project, - max_tokens, - format, - include_preferences, - }; - - match build_wake_up(store, &opts) { - Ok(pack) => ToolResult::text(pack), - Err(e) => ToolResult::error(format!("wake_up failed: {e}")), - } -} - -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -fn get_str<'a>(args: &'a Value, key: &str) -> Option<&'a str> { - args.get(key).and_then(|v| v.as_str()) -} - -fn get_i64(args: &Value, key: &str, default: i64) -> i64 { - args.get(key).and_then(|v| v.as_i64()).unwrap_or(default) -} - -fn resolve_memoir(store: &Store, name: &str) -> Result { - store - .get_memoir_by_name(name) - .map_err(|e| ToolResult::error(format!("db error: {e}")))? - .ok_or_else(|| ToolResult::error(format!("memoir not found: {name}"))) -} - -// --------------------------------------------------------------------------- -// Memory tool handlers -// --------------------------------------------------------------------------- - -fn tool_store( - store: &Store, - embedder: Option<&dyn Embedder>, - args: &Value, - compact: bool, - auto_consolidate: AutoConsolidate, -) -> ToolResult { - let topic = match get_str(args, "topic") { - Some(t) => t, - None => return ToolResult::error("missing required field: topic".into()), - }; - let content = match get_str(args, "content") { - Some(c) => c, - None => return ToolResult::error("missing required field: content".into()), - }; - - // Empty-string validation: the inputSchema marks `topic` and - // `content` as required, but JSON allows passing `""` which slips - // past the structural check. Reject explicitly so callers don't - // silently end up with a memory under a blank topic that they - // can't meaningfully recall. - if topic.trim().is_empty() { - return ToolResult::error("topic must not be empty".into()); - } - if content.trim().is_empty() { - return ToolResult::error("content must not be empty".into()); - } - - // Input length validation - if topic.len() > MAX_TOPIC_LEN { - return ToolResult::error(format!( - "topic exceeds maximum length ({} > {MAX_TOPIC_LEN} UTF-8 bytes)", - topic.len() - )); - } - if content.len() > MAX_CONTENT_LEN { - return ToolResult::error(format!( - "content exceeds maximum length ({} > {MAX_CONTENT_LEN} UTF-8 bytes)", - content.len() - )); - } - - let importance_str = get_str(args, "importance").unwrap_or("medium"); - let importance = importance_str - .parse() - .unwrap_or(icm_core::Importance::Medium); - - let mut memory = Memory::new(topic.into(), content.into(), importance); - - let kw = parse_keywords(args); - if !kw.is_empty() { - memory.keywords = kw; - } - - if let Some(raw) = get_str(args, "raw_excerpt") { - memory.raw_excerpt = Some(raw.into()); - } - - // Auto-embed if embedder is available - let embed_text = memory.embed_text(); - let embed_vec = if let Some(emb) = embedder { - match emb.embed(&embed_text) { - Ok(vec) => Some(vec), - Err(e) => { - tracing::warn!("embedding failed: {e}"); - None - } - } - } else { - None - }; - - if let Some(ref vec) = embed_vec { - memory.embedding = Some(vec.clone()); - } - - // Dedup check: if a very similar memory exists in the same topic, update it instead - if let Some(ref query_emb) = embed_vec { - if let Ok(Some((existing, score))) = find_similar_memory( - store, - &embed_text, - query_emb, - topic, - DEDUP_SIMILARITY_THRESHOLD, - ) { - let updated = Memory { - id: existing.id.clone(), - created_at: existing.created_at, - last_accessed: existing.last_accessed, - access_count: existing.access_count, - weight: 1.0, - topic: existing.topic.clone(), - summary: content.to_string(), - raw_excerpt: get_str(args, "raw_excerpt") - .map(|r| r.into()) - .or_else(|| existing.raw_excerpt.clone()), - keywords: { - let kw = parse_keywords(args); - if kw.is_empty() { - existing.keywords.clone() - } else { - kw - } - }, - embedding: Some(query_emb.clone()), - // Never let a near-dup merge downgrade importance: an MCP - // caller that omits `importance` defaults to Medium, which - // would otherwise silently demote an existing Critical - // memory into decay/prune eligibility (audit finding). - importance: icm_core::max_importance(existing.importance, importance), - source: existing.source.clone(), - related_ids: existing.related_ids.clone(), - updated_at: Utc::now(), - scope: existing.scope, - }; - if let Err(e) = store.update(&updated) { - return ToolResult::error(format!("failed to update: {e}")); - } - return if compact { - ToolResult::text(format!("ok:{}", updated.id)) - } else { - ToolResult::text(format!( - "Updated existing memory (similarity {score:.2}): {}", - updated.id - )) - }; - } - } - - // Auto-link: populate `related_ids` with similar existing memories BEFORE - // storing, so the new memory lands in the DB with its forward edges - // already set. Back-refs are added AFTER storing so the linked memories - // point to an id that exists in the DB. - let auto_link_opts = AutoLinkOptions::default(); - let linked_ids = if memory.embedding.is_some() { - auto_link_memory(store, &mut memory, &auto_link_opts).unwrap_or_else(|e| { - tracing::warn!("auto-link failed: {e}"); - Vec::new() - }) - } else { - Vec::new() - }; - - match store.store(memory) { - Ok(id) => { - // Best-effort back-ref update. Failure here leaves an asymmetric - // edge (forward-only) but does not fail the store call. - if !linked_ids.is_empty() { - if let Err(e) = add_backrefs(store, &id, &linked_ids) { - tracing::warn!("auto-link back-ref update failed: {e}"); - } - } - - let link_suffix = if linked_ids.is_empty() { - String::new() - } else { - format!( - " (+{} link{})", - linked_ids.len(), - if linked_ids.len() == 1 { "" } else { "s" } - ) - }; - - if compact { - // Try auto-consolidation even in compact mode - let consolidation_msg = - try_auto_consolidate(store, embedder, topic, auto_consolidate); - if consolidation_msg.is_empty() { - ToolResult::text(format!("ok:{id}{link_suffix}")) - } else { - ToolResult::text(format!("ok:{id}{link_suffix}\n{consolidation_msg}")) - } - } else { - let consolidation_msg = - try_auto_consolidate(store, embedder, topic, auto_consolidate); - if consolidation_msg.is_empty() { - // Still show a nudge if approaching threshold - let hint = if let Ok(count) = store.count_by_topic(topic) { - if count > 7 { - format!( - "\nNote: Topic '{topic}' has {count} entries — consider consolidating with icm_memory_consolidate." - ) - } else { - String::new() - } - } else { - String::new() - }; - ToolResult::text(format!("Stored memory: {id}{link_suffix}{hint}")) - } else { - ToolResult::text(format!( - "Stored memory: {id}{link_suffix}\n{consolidation_msg}" - )) - } - } - } - Err(e) => ToolResult::error(format!("failed to store: {e}")), - } -} - -fn format_memory_output(memories: &[(Memory, f32)], compact: bool) -> String { - // Audit finding: `summary` has no newline/CR validation at the store - // layer (only `topic` is checked — see `validate_fields`), and it can - // be LLM/tool-extracted from untrusted content. Written verbatim, a - // stored summary could forge a fake `--- [score: ...] ---` - // delimiter indistinguishable from a real entry, or (compact mode) a - // fake `[topic] ...` line. `keywords` has no validation at all. Flatten - // both, same fix already applied to recall_context/render_detail. - let flatten = |s: &str| s.replace(['\n', '\r'], " "); - let mut output = String::new(); - if compact { - for (mem, _) in memories { - output.push_str(&format!("[{}] {}\n", mem.topic, flatten(&mem.summary))); - } - } else { - for (mem, score) in memories { - let summary = flatten(&mem.summary); - if *score >= 0.0 { - output.push_str(&format!( - "--- {} [score: {:.3}] ---\n topic: {}\n importance: {}\n weight: {:.3}\n summary: {}\n", - mem.id, score, mem.topic, mem.importance, mem.weight, summary - )); - } else { - output.push_str(&format!( - "--- {} ---\n topic: {}\n importance: {}\n weight: {:.3}\n summary: {}\n", - mem.id, mem.topic, mem.importance, mem.weight, summary - )); - } - if !mem.keywords.is_empty() { - let flattened_keywords: Vec = - mem.keywords.iter().map(|k| flatten(k)).collect(); - output.push_str(&format!(" keywords: {}\n", flattened_keywords.join(", "))); - } - if let Some(ref raw) = mem.raw_excerpt { - // raw_excerpt can hold up to 64 KB per memory; dumping it in - // full for every hit floods the client LLM's context (audit - // finding). Cap the recall view — the full excerpt stays in - // the store. - const MAX_RAW_IN_RECALL: usize = 2048; - if raw.len() > MAX_RAW_IN_RECALL { - let mut cut = MAX_RAW_IN_RECALL; - while !raw.is_char_boundary(cut) { - cut -= 1; - } - output.push_str(&format!( - " raw: {}… [truncated, {} bytes total]\n", - &raw[..cut], - raw.len() - )); - } else { - output.push_str(&format!(" raw: {raw}\n")); - } - } - output.push('\n'); - } - } - output -} - -fn tool_recall(context: &ToolContext<'_>, args: &Value) -> ToolResult { - let store = context.store; - let embedder = context.embedder; - let compact = context.compact; - // Auto-decay if >24h since last decay - if let Err(e) = store.maybe_auto_decay() { - tracing::warn!(error = %e, "auto-decay failed during recall"); - } - - let query = match get_str(args, "query") { - Some(q) => q, - None => return ToolResult::error("missing required field: query".into()), - }; - // The modern input contract extends the historical advertised maximum - // from 20 to the frozen Phase 2 boundary of 100. Keep the handler cap in - // lockstep so valid modern calls are not silently truncated. - let limit = get_i64(args, "limit", 5).clamp(1, 100) as usize; - let topic = get_str(args, "topic"); - let keyword = get_str(args, "keyword"); - - // Project filter: same hard segment-aware filter applied to the CLI - // `recall_context` path (extract.rs) so MCP-side recall can't leak - // memories from other projects. Caller can override via the explicit - // `project` arg (empty string disables the filter); otherwise we - // derive it from the server's cwd via the shared icm-core detection - // (git remote first) — the CLI hooks store under that name, so a raw - // cwd basename would silently miss on renamed checkouts (audit finding). - let project_arg = get_str(args, "project"); - let cwd_project = - icm_core::project::project_from_path(&context.working_directory.to_string_lossy()); - let project: Option = match project_arg { - Some("") => None, - Some(p) => Some(p.to_string()), - None => cwd_project, - }; - let project_filter = |m: &Memory| -> bool { - match project.as_deref() { - None => true, - Some(p) => is_preference_topic(&m.topic) || project_matches(&m.topic, Some(p)), - } - }; - - // Audit finding: filters were applied AFTER the store already truncated - // to `limit` — if the top-`limit` global hits all belonged to other - // projects/topics, filtering left nothing and recall reported "no - // memories" even though relevant matches existed further down the - // ranked list. When any filter is active, request a much larger - // candidate pool so filtering has enough to work with, then truncate to - // the caller's requested `limit` at the very end (capped — this is a - // memory-scoped search, not a paginated export). - let filters_active = project.is_some() || topic.is_some() || keyword.is_some(); - let query_limit = if filters_active { - (limit * 10).min(200) - } else { - limit - }; - - // Try hybrid search if embedder is available - if let Some(emb) = embedder { - if let Ok(query_emb) = emb.embed_query(query) { - if let Ok(results) = store.search_hybrid(query, &query_emb, query_limit) { - let mut scored_results = results; - scored_results.retain(|(m, _)| project_filter(m)); - if let Some(t) = topic { - scored_results.retain(|(m, _)| topic_matches(&m.topic, t)); - } - if let Some(kw) = keyword { - scored_results.retain(|(m, _)| keyword_matches(&m.keywords, kw)); - } - - // Graph-aware expansion: follow `related_ids` one hop from - // each primary hit and fold neighbors into the result set. - // Neighbors carry a discounted score so they rank below - // direct matches but can displace weak primary results. - // - // Audit R13b: neighbors are fetched by id without going - // through the project / topic / keyword filters above, - // so a project-A primary hit can pull in a project-B - // neighbor via auto-linked `related_ids`. Re-apply the - // filters to `expanded` so the caller's scope is honored. - let max_neighbors = (query_limit / 3).max(1); - let mut expanded = store - .expand_with_neighbors(&scored_results, max_neighbors, 0.5, query_limit) - .unwrap_or(scored_results); - expanded.retain(|(m, _)| project_filter(m)); - if let Some(t) = topic { - expanded.retain(|(m, _)| topic_matches(&m.topic, t)); - } - if let Some(kw) = keyword { - expanded.retain(|(m, _)| keyword_matches(&m.keywords, kw)); - } - expanded.truncate(limit); - - // Batch update access counts (includes expanded neighbors) - let ids: Vec<&str> = expanded.iter().map(|(m, _)| m.id.as_str()).collect(); - let _ = store.batch_update_access(&ids); - - if expanded.is_empty() { - return ToolResult::text(MSG_NO_MEMORIES.into()); - } - - return ToolResult::text(format_memory_output(&expanded, compact)); - } - } - } - - // Fallback: FTS then keywords - let mut results = match store.search_fts(query, query_limit) { - Ok(r) => r, - Err(e) => return ToolResult::error(format!("search error: {e}")), - }; - - if results.is_empty() { - let keywords: Vec<&str> = query.split_whitespace().collect(); - results = match store.search_by_keywords(&keywords, query_limit) { - Ok(r) => r, - Err(e) => return ToolResult::error(format!("search error: {e}")), - }; - } - - results.retain(|m| project_filter(m)); - if let Some(t) = topic { - results.retain(|m| topic_matches(&m.topic, t)); - } - if let Some(kw) = keyword { - results.retain(|m| keyword_matches(&m.keywords, kw)); - } - results.truncate(limit); - - // Convert to scored format with a sentinel score of 1.0 (FTS fallback - // doesn't expose a real similarity score, but we still want the graph - // expansion to score neighbors relative to their primary parent). - let scored: Vec<(Memory, f32)> = results.into_iter().map(|m| (m, 1.0)).collect(); - - // Graph-aware expansion also applies in the fallback path so that - // keyword-only deployments benefit from auto-linked memories. - // Same R13b re-filter as the hybrid path. - let max_neighbors = (limit / 3).max(1); - let mut expanded = store - .expand_with_neighbors(&scored, max_neighbors, 0.5, limit) - .unwrap_or(scored); - expanded.retain(|(m, _)| project_filter(m)); - if let Some(t) = topic { - expanded.retain(|(m, _)| topic_matches(&m.topic, t)); - } - if let Some(kw) = keyword { - expanded.retain(|(m, _)| keyword_matches(&m.keywords, kw)); - } - - // Batch update access counts (includes expanded neighbors) - let ids: Vec<&str> = expanded.iter().map(|(m, _)| m.id.as_str()).collect(); - let _ = store.batch_update_access(&ids); - - if expanded.is_empty() { - return ToolResult::text(MSG_NO_MEMORIES.into()); - } - - // FTS-path results have synthetic scores — reset to -1.0 for display - // so we don't claim a hybrid-search confidence we didn't compute. - let for_display: Vec<(Memory, f32)> = expanded.into_iter().map(|(m, _)| (m, -1.0)).collect(); - ToolResult::text(format_memory_output(&for_display, compact)) -} - -fn tool_forget(store: &Store, args: &Value) -> ToolResult { - let id = match get_str(args, "id") { - Some(id) => id, - None => return ToolResult::error("missing required field: id".into()), - }; - - match store.delete(id) { - Ok(()) => ToolResult::text(format!("Deleted memory: {id}")), - Err(e) => ToolResult::error(format!("failed to delete: {e}")), - } -} - -fn tool_forget_topic(store: &Store, args: &Value) -> ToolResult { - let topic = match get_str(args, "topic") { - Some(t) => t, - None => return ToolResult::error("missing required field: topic".into()), - }; - - let memories = match store.get_by_topic(topic) { - Ok(m) => m, - Err(e) => return ToolResult::error(format!("failed to get memories: {e}")), - }; - - let count = memories.len(); - for m in &memories { - if let Err(e) = store.delete(&m.id) { - return ToolResult::error(format!("failed to delete memory {}: {e}", m.id)); - } - } - - ToolResult::text(format!("Deleted {count} memories from topic: {topic}")) -} - -fn tool_learn(store: &Store, args: &Value) -> ToolResult { - let dir_str = get_str(args, "directory").unwrap_or("."); - let dir = std::path::PathBuf::from(dir_str); - - if !dir.exists() || !dir.is_dir() { - return ToolResult::error(format!("directory not found: {}", dir.display())); - } - - let name = get_str(args, "name"); - - match icm_core::learn_project(store, &dir, name) { - Ok(result) => ToolResult::text(result.to_string()), - Err(e) => ToolResult::error(format!("learn failed: {e}")), - } -} - -fn tool_learn_bounded(context: &ToolContext<'_>, args: &Value) -> ToolResult { - if !context.enforce_directory_boundary { - return tool_learn(context.store, args); - } - let requested = get_str(args, "directory").unwrap_or("."); - let requested = std::path::Path::new(requested); - let candidate = if requested.is_absolute() { - requested.to_path_buf() - } else { - context.working_directory.join(requested) - }; - let root = match context.working_directory.canonicalize() { - Ok(root) => root, - Err(_) => { - return ToolResult::error( - "server working directory could not be resolved safely".into(), - ) - } - }; - let candidate = match candidate.canonicalize() { - Ok(candidate) => candidate, - Err(_) => { - return ToolResult::error(format!("directory not found: {}", candidate.display())) - } - }; - if !candidate.starts_with(&root) { - return ToolResult::error( - "directory must remain within the server working directory".into(), - ); - } - let mut bounded_args = args.clone(); - bounded_args["directory"] = Value::String(candidate.to_string_lossy().into_owned()); - tool_learn(context.store, &bounded_args) -} - -fn tool_consolidate(store: &Store, embedder: Option<&dyn Embedder>, args: &Value) -> ToolResult { - let topic = match get_str(args, "topic") { - Some(t) => t, - None => return ToolResult::error("missing required field: topic".into()), - }; - let summary = match get_str(args, "summary") { - Some(s) => s, - None => return ToolResult::error("missing required field: summary".into()), - }; - - let mut consolidated = Memory::new(topic.into(), summary.into(), icm_core::Importance::High); - // Same bug class as #394/#395/cmd_consolidate: this tool never attached - // an embedding to the merged memory it creates. - if let Some(emb) = embedder { - if let Ok(vec) = emb.embed(&consolidated.embed_text()) { - consolidated.embedding = Some(vec); - } - } - - match store.consolidate_topic(topic, consolidated) { - Ok(()) => ToolResult::text(format!("Consolidated topic: {topic}")), - Err(e) => ToolResult::error(format!("failed to consolidate: {e}")), - } -} - -fn tool_list_topics(store: &Store) -> ToolResult { - match store.list_topics() { - Ok(topics) => { - if topics.is_empty() { - return ToolResult::text("No topics yet.".into()); - } - - // Group topics by scope prefix (before ':') - let mut scoped: std::collections::BTreeMap> = - std::collections::BTreeMap::new(); - let mut unscoped: Vec<(String, usize)> = Vec::new(); - - for (topic, count) in &topics { - if let Some((prefix, _rest)) = topic.split_once(':') { - scoped - .entry(prefix.to_string()) - .or_default() - .push((topic.clone(), *count)); - } else { - unscoped.push((topic.clone(), *count)); - } - } - - let mut output = String::from("Topics:\n"); - - // Show unscoped topics first - for (topic, count) in &unscoped { - output.push_str(&format!(" {topic}: {count} memories\n")); - } - - // Show scoped topics grouped by prefix - for (prefix, sub_topics) in &scoped { - let total: usize = sub_topics.iter().map(|(_, c)| c).sum(); - output.push_str(&format!(" [{prefix}] ({total} total):\n")); - for (topic, count) in sub_topics { - output.push_str(&format!(" {topic}: {count} memories\n")); - } - } - - ToolResult::text(output) - } - Err(e) => ToolResult::error(format!("failed to list topics: {e}")), - } -} - -fn tool_stats(store: &Store) -> ToolResult { - match store.stats() { - Ok(stats) => { - let mut output = format!( - "Memories: {}\nTopics: {}\nAvg weight: {:.3}\n", - stats.total_memories, stats.total_topics, stats.avg_weight - ); - if let Some(oldest) = stats.oldest_memory { - output.push_str(&format!( - "Oldest: {}\n", - format_local(&oldest, "%Y-%m-%d %H:%M") - )); - } - if let Some(newest) = stats.newest_memory { - output.push_str(&format!( - "Newest: {}\n", - format_local(&newest, "%Y-%m-%d %H:%M") - )); - } - ToolResult::text(output) - } - Err(e) => ToolResult::error(format!("failed to get stats: {e}")), - } -} - -fn tool_update(store: &Store, embedder: Option<&dyn Embedder>, args: &Value) -> ToolResult { - let id = match get_str(args, "id") { - Some(id) => id, - None => return ToolResult::error("missing required field: id".into()), - }; - let content = match get_str(args, "content") { - Some(c) => c, - None => return ToolResult::error("missing required field: content".into()), - }; - - let mut memory = match store.get(id) { - Ok(Some(m)) => m, - Ok(None) => return ToolResult::error(format!("memory not found: {id}")), - Err(e) => return ToolResult::error(format!("db error: {e}")), - }; - - memory.summary = content.to_string(); - memory.updated_at = Utc::now(); - memory.weight = 1.0; // Reset weight on update (refreshed content) - - if let Some(imp_str) = get_str(args, "importance") { - if let Ok(imp) = imp_str.parse() { - memory.importance = imp; - } - } - - let kw = parse_keywords(args); - if !kw.is_empty() { - memory.keywords = kw; - } - - // Re-embed if embedder available - if let Some(emb) = embedder { - if let Ok(vec) = emb.embed(&memory.embed_text()) { - memory.embedding = Some(vec); - } - } - - match store.update(&memory) { - Ok(()) => ToolResult::text(format!("Updated memory: {id}")), - Err(e) => ToolResult::error(format!("failed to update: {e}")), - } -} - -fn tool_health(store: &Store, args: &Value) -> ToolResult { - let specific_topic = get_str(args, "topic"); - - let topics = if let Some(t) = specific_topic { - vec![(t.to_string(), 0usize)] - } else { - match store.list_topics() { - Ok(t) => t, - Err(e) => return ToolResult::error(format!("failed to list topics: {e}")), - } - }; - - if topics.is_empty() { - return ToolResult::text("No topics yet.".into()); - } - - let mut output = String::from("Memory Health Report:\n\n"); - let mut total_stale = 0usize; - let mut topics_needing_consolidation = 0usize; - - for (topic, _) in &topics { - match store.topic_health(topic) { - Ok(health) => { - let status = health.status(); - - output.push_str(&format!( - " {topic}: {status}\n entries: {} avg_weight: {:.2} stale: {} avg_access: {:.1}\n", - health.entry_count, health.avg_weight, health.stale_count, health.avg_access_count - )); - - if health.needs_consolidation { - topics_needing_consolidation += 1; - } - total_stale += health.stale_count; - } - Err(_) => { - output.push_str(&format!(" {topic}: (error reading)\n")); - } - } - } - - output.push_str(&format!( - "\nSummary: {} topics, {} need consolidation, {} stale entries total\n", - topics.len(), - topics_needing_consolidation, - total_stale - )); - - ToolResult::text(output) -} - -fn tool_extract_patterns(store: &Store, args: &Value) -> ToolResult { - let topic = match get_str(args, "topic") { - Some(t) => t, - None => return ToolResult::error("missing required field: topic".into()), - }; - let min_cluster_size = get_i64(args, "min_cluster_size", 3).clamp(2, 50) as usize; - let memoir_name = get_str(args, "memoir"); - - let patterns = match store.detect_patterns(topic, min_cluster_size) { - Ok(p) => p, - Err(e) => return ToolResult::error(format!("pattern detection failed: {e}")), - }; - - if patterns.is_empty() { - return ToolResult::text(format!( - "No patterns detected in topic '{topic}' (min cluster size: {min_cluster_size})." - )); - } - - let mut output = format!( - "Detected {} pattern(s) in topic '{topic}':\n\n", - patterns.len() - ); - - // If memoir is provided, resolve it and create concepts - let memoir_id = if let Some(mname) = memoir_name { - match resolve_memoir(store, mname) { - Ok(m) => Some(m.id), - Err(e) => return e, - } - } else { - None - }; - - for (i, cluster) in patterns.iter().enumerate() { - output.push_str(&format!( - "Pattern {}: {} memories\n Keywords: {}\n Representative: {}\n", - i + 1, - cluster.count, - cluster.keywords.join(", "), - cluster.representative_summary, - )); - - if let Some(ref mid) = memoir_id { - match store.extract_pattern_as_concept(cluster, mid) { - Ok(concept_id) => { - output.push_str(&format!(" -> Created concept: {concept_id}\n")); - } - Err(e) => { - output.push_str(&format!(" -> Failed to create concept: {e}\n")); - } - } - } - - output.push('\n'); - } - - if memoir_id.is_some() { - output.push_str(&format!( - "Created {} concept(s) in memoir '{}'.\n", - patterns.len(), - memoir_name.unwrap_or("?") - )); - } - - ToolResult::text(output) -} - -fn tool_embed_all(store: &Store, embedder: Option<&dyn Embedder>, args: &Value) -> ToolResult { - let embedder = match embedder { - Some(e) => e, - None => return ToolResult::error("embeddings not available".into()), - }; - - let topic_filter = get_str(args, "topic"); - - // Get all memories in a single query - let memories = if let Some(t) = topic_filter { - match store.get_by_topic(t) { - Ok(m) => m, - Err(e) => return ToolResult::error(format!("failed to list memories: {e}")), - } - } else { - match store.list_all() { - Ok(m) => m, - Err(e) => return ToolResult::error(format!("failed to list memories: {e}")), - } - }; - - // Filter to only those without embeddings - let to_embed: Vec<&Memory> = memories.iter().filter(|m| m.embedding.is_none()).collect(); - - if to_embed.is_empty() { - return ToolResult::text("All memories already have embeddings.".into()); - } - - let total = to_embed.len(); - - // Batch embed all texts at once - let texts: Vec = to_embed.iter().map(|m| m.embed_text()).collect(); - let text_refs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect(); - - let embeddings = match embedder.embed_batch(&text_refs) { - Ok(vecs) => vecs, - Err(e) => return ToolResult::error(format!("batch embedding failed: {e}")), - }; - - let mut embedded = 0; - let mut errors = 0; - - for (mem, vec) in to_embed.iter().zip(embeddings) { - let mut updated = (*mem).clone(); - updated.embedding = Some(vec); - if store.update(&updated).is_ok() { - embedded += 1; - } else { - errors += 1; - } - } - - ToolResult::text(format!( - "Embedded {embedded}/{total} memories ({errors} errors)" - )) -} - -// --------------------------------------------------------------------------- -// Memoir tool handlers -// --------------------------------------------------------------------------- - -fn tool_memoir_create(store: &Store, args: &Value) -> ToolResult { - let name = match get_str(args, "name") { - Some(n) => n, - None => return ToolResult::error("missing required field: name".into()), - }; - if name.len() > 255 { - return ToolResult::error(format!( - "name too long: {} UTF-8 bytes (max 255)", - name.len() - )); - } - let description = get_str(args, "description").unwrap_or(""); - if description.len() > 10_000 { - return ToolResult::error(format!( - "description too long: {} UTF-8 bytes (max 10000)", - description.len() - )); - } - - let memoir = Memoir::new(name.into(), description.into()); - match store.create_memoir(memoir) { - Ok(id) => ToolResult::text(format!("Created memoir '{name}': {id}")), - Err(e) => ToolResult::error(format!("failed to create memoir: {e}")), - } -} - -fn tool_memoir_list(store: &Store) -> ToolResult { - let memoirs = match store.list_memoirs() { - Ok(m) => m, - Err(e) => return ToolResult::error(format!("failed to list memoirs: {e}")), - }; - - if memoirs.is_empty() { - return ToolResult::text("No memoirs yet.".into()); - } - - let counts = store.batch_memoir_concept_counts().unwrap_or_default(); - let mut output = String::from("Memoirs:\n"); - for m in &memoirs { - let concept_count = counts.get(&m.id).copied().unwrap_or(0); - output.push_str(&format!( - " {} ({} concepts) — {}\n", - m.name, concept_count, m.description - )); - } - ToolResult::text(output) -} - -fn tool_memoir_show(store: &Store, args: &Value) -> ToolResult { - let name = match get_str(args, "name") { - Some(n) => n, - None => return ToolResult::error("missing required field: name".into()), - }; - - let memoir = match resolve_memoir(store, name) { - Ok(m) => m, - Err(e) => return e, - }; - let stats = match store.memoir_stats(&memoir.id) { - Ok(s) => s, - Err(e) => return ToolResult::error(format!("failed to get stats: {e}")), - }; - let concepts = match store.list_concepts(&memoir.id) { - Ok(c) => c, - Err(e) => return ToolResult::error(format!("failed to list concepts: {e}")), - }; - - let mut output = format!( - "Memoir: {}\nDescription: {}\nConcepts: {}\nLinks: {}\nAvg confidence: {:.2}\n", - memoir.name, - memoir.description, - stats.total_concepts, - stats.total_links, - stats.avg_confidence - ); - - if !stats.label_counts.is_empty() { - output.push_str("Labels:\n"); - for (label, count) in &stats.label_counts { - output.push_str(&format!(" {label} ({count})\n")); - } - } - - if !concepts.is_empty() { - output.push_str("\nConcepts:\n"); - for c in &concepts { - let labels_str = c.format_labels(); - output.push_str(&format!( - " {} [r{} c{:.2}]{}\n {}\n", - c.name, - c.revision, - c.confidence, - if labels_str.is_empty() { - String::new() - } else { - format!(" ({labels_str})") - }, - c.definition - )); - } - } - - ToolResult::text(output) -} - -fn tool_memoir_add_concept(store: &Store, args: &Value) -> ToolResult { - let memoir_name = match get_str(args, "memoir") { - Some(n) => n, - None => return ToolResult::error("missing required field: memoir".into()), - }; - let name = match get_str(args, "name") { - Some(n) => n, - None => return ToolResult::error("missing required field: name".into()), - }; - if name.len() > 255 { - return ToolResult::error(format!( - "concept name too long: {} UTF-8 bytes (max 255)", - name.len() - )); - } - let definition = match get_str(args, "definition") { - Some(d) => d, - None => return ToolResult::error("missing required field: definition".into()), - }; - if definition.len() > 10_000 { - return ToolResult::error(format!( - "definition too long: {} UTF-8 bytes (max 10000)", - definition.len() - )); - } - - let memoir = match resolve_memoir(store, memoir_name) { - Ok(m) => m, - Err(e) => return e, - }; - - let mut concept = Concept::new(memoir.id, name.into(), definition.into()); - - if let Some(labels_str) = get_str(args, "labels") { - concept.labels = labels_str - .split(',') - .filter_map(|s| s.trim().parse::