From 0874988e48cbb53509c7fccb19635c6091d42837 Mon Sep 17 00:00:00 2001 From: Lauritz-Timm Date: Thu, 13 Aug 2026 09:33:05 +0200 Subject: [PATCH] feat(mcp): add typed output contracts --- Cargo.lock | 3 + Cargo.toml | 7 +- crates/icm-mcp/src/catalog.rs | 338 +++++++++- crates/icm-mcp/src/lib.rs | 1 + crates/icm-mcp/src/outputs.rs | 621 +++++++++++++++++ crates/icm-mcp/src/protocol.rs | 53 ++ crates/icm-mcp/src/service.rs | 633 ++++++++++++++++-- crates/icm-mcp/src/tools.rs | 5 +- crates/icm-mcp/src/tools/handlers/common.rs | 52 +- crates/icm-mcp/src/tools/handlers/feedback.rs | 45 +- crates/icm-mcp/src/tools/handlers/memory.rs | 139 ++-- .../icm-mcp/src/tools/handlers/transcript.rs | 42 +- crates/icm-mcp/src/tools/registry.rs | 34 +- crates/icm-mcp/src/tools/tests.rs | 1 + 14 files changed, 1797 insertions(+), 177 deletions(-) create mode 100644 crates/icm-mcp/src/outputs.rs diff --git a/Cargo.lock b/Cargo.lock index 8be5e412..fedd3ca8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1470,6 +1470,7 @@ version = "0.10.61" dependencies = [ "anyhow", "axum", + "base64 0.22.1", "chrono", "clap", "crossterm", @@ -3091,6 +3092,7 @@ version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "687274d293b6cdc6e73e0fee520bf2049650090d7164f87672d212a3c530cf4a" dependencies = [ + "chrono", "dyn-clone", "ref-cast", "schemars_derive", @@ -3643,6 +3645,7 @@ dependencies = [ "libc", "mio", "pin-project-lite", + "signal-hook-registry", "socket2", "tokio-macros", "windows-sys 0.61.2", diff --git a/Cargo.toml b/Cargo.toml index 441be270..4412713c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,7 +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" +schemars = { version = "1", features = ["chrono04"] } # Error handling thiserror = "2" @@ -75,9 +75,10 @@ crossterm = "0.28" axum = "0.8" # tokio is only used by the optional `web` feature (axum dashboard). # `full` was overkill — we need the multi-thread runtime, the -# `#[tokio::main]` macro, and `tokio::net` for the TCP listener. +# `#[tokio::main]` macro, `tokio::net` for the TCP listener, and +# `tokio::signal` for graceful daemon shutdown. # Dropping `full` shaves ~3MB off the optimized binary on Linux. -tokio = { version = "1", features = ["rt-multi-thread", "macros", "net"] } +tokio = { version = "1", features = ["rt-multi-thread", "macros", "net", "signal"] } tower-http = { version = "0.6", features = ["trace", "cors"] } rust-embed = "8" mime_guess = "2" diff --git a/crates/icm-mcp/src/catalog.rs b/crates/icm-mcp/src/catalog.rs index 0a960033..ab6d42a1 100644 --- a/crates/icm-mcp/src/catalog.rs +++ b/crates/icm-mcp/src/catalog.rs @@ -4,11 +4,13 @@ //! The ordered registration vector is the only order source, and the lookup //! map points back into that same vector. +use std::any::TypeId; use std::collections::{HashMap, HashSet}; use std::path::Path; use icm_core::Embedder; use icm_store::Store; +use schemars::JsonSchema; use serde::de::DeserializeOwned; use serde_json::{json, Value}; @@ -19,6 +21,7 @@ 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; +type OutputValidator = fn(&Value, &Value) -> Result<(), String>; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum EmbedderRequirement { @@ -144,6 +147,9 @@ pub struct ToolSpec { description: &'static str, legacy_input_schema: Value, modern_input_schema: Value, + modern_output_schema: Option, + modern_output_type: Option, + validate_output: Option, legacy_input_normalizer: Option, annotations: ToolAnnotations, requirements: ToolRequirements, @@ -171,6 +177,9 @@ impl ToolSpec { description, legacy_input_schema, modern_input_schema, + modern_output_schema: None, + modern_output_type: None, + validate_output: None, legacy_input_normalizer, annotations, requirements, @@ -179,6 +188,17 @@ impl ToolSpec { } } + pub(crate) fn with_output(mut self) -> Self + where + O: DeserializeOwned + JsonSchema + 'static, + { + self.modern_output_schema = Some(generated_output_schema::()); + self.requirements.structured_output_from_revision = Some(ProtocolRevision::V2025_06_18); + self.modern_output_type = Some(TypeId::of::()); + self.validate_output = Some(validate_output::); + self + } + fn legacy_definition(&self) -> Value { json!({ "name": self.name, @@ -188,7 +208,7 @@ impl ToolSpec { } fn modern_definition(&self) -> Value { - json!({ + let mut definition = json!({ "name": self.name, "description": self.description, "inputSchema": self.modern_input_schema, @@ -196,7 +216,14 @@ impl ToolSpec { "_meta": { "com.github.rtk-ai.icm/requirements": self.requirements.as_value(), }, - }) + }); + if let Some(output_schema) = &self.modern_output_schema { + definition + .as_object_mut() + .expect("tool definitions have object roots") + .insert("outputSchema".into(), output_schema.clone()); + } + definition } fn validate_modern_input(&self, arguments: &Value) -> Result<(), String> { @@ -229,6 +256,15 @@ where .map_err(|error| bounded_error(error.to_string())) } +fn validate_output(output: &Value, schema: &Value) -> Result<(), String> +where + O: DeserializeOwned, +{ + serde_json::from_value::(output.clone()) + .map_err(|error| bounded_error(error.to_string()))?; + validate_schema_constraints(output, schema, "$", 0) +} + fn generated_input_schema(legacy: &Value) -> Value where I: ModernToolInput, @@ -271,15 +307,100 @@ where generated } +fn generated_output_schema() -> Value +where + O: JsonSchema, +{ + let mut settings = schemars::generate::SchemaSettings::draft2020_12().for_serialize(); + settings.meta_schema = None; + let schema = settings.into_generator().into_root_schema_for::(); + let mut value = serde_json::to_value(schema).expect("generated output schema must serialize"); + normalize_output_schema(&mut value); + value +} + +fn normalize_output_schema(value: &mut Value) { + let Value::Object(object) = value else { + if let Value::Array(values) = value { + values.iter_mut().for_each(normalize_output_schema); + } + return; + }; + + object.remove("title"); + if object.get("format").and_then(Value::as_str) != Some("date-time") { + object.remove("format"); + } + object.values_mut().for_each(normalize_output_schema); + + if object.contains_key("const") { + object.remove("type"); + } + + if let Some(Value::Array(types)) = object.get("type") { + let non_null: Vec<_> = types + .iter() + .filter(|value| value.as_str() != Some("null")) + .cloned() + .collect(); + if non_null.len() == 1 && non_null.len() + 1 == types.len() { + let Some(non_null_type) = non_null.into_iter().next() else { + return; + }; + let mut non_null_schema = std::mem::take(object); + non_null_schema.insert("type".into(), non_null_type); + if let Some(Value::Array(values)) = non_null_schema.get_mut("enum") { + values.retain(|value| !value.is_null()); + } + object.insert("oneOf".into(), json!([non_null_schema, { "type": "null" }])); + return; + } + } + + let nullable_any_of = object + .get("anyOf") + .and_then(Value::as_array) + .is_some_and(|variants| { + variants.len() == 2 + && variants + .iter() + .any(|variant| variant.get("type").and_then(Value::as_str) == Some("null")) + }); + if nullable_any_of { + if let Some(variants) = object.remove("anyOf") { + object.insert("oneOf".into(), variants); + } + } +} + fn validate_schema_constraints( value: &Value, schema: &Value, path: &str, depth: usize, +) -> Result<(), String> { + validate_schema_constraints_at(value, schema, schema, path, depth) +} + +fn validate_schema_constraints_at( + value: &Value, + schema: &Value, + root_schema: &Value, + path: &str, + depth: usize, ) -> Result<(), String> { if depth > 32 { return Err("input nesting exceeds maximum depth".into()); } + if let Some(reference) = schema.get("$ref").and_then(Value::as_str) { + let Some(referenced) = reference + .strip_prefix('#') + .and_then(|pointer| root_schema.pointer(pointer)) + else { + return Err(format!("{path} contains an unresolved schema reference")); + }; + return validate_schema_constraints_at(value, referenced, root_schema, path, depth + 1); + } 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}")); @@ -328,9 +449,10 @@ fn validate_schema_constraints( ) { for (name, child) in object { if let Some(child_schema) = properties.get(name) { - validate_schema_constraints( + validate_schema_constraints_at( child, child_schema, + root_schema, &format!("{path}.{name}"), depth + 1, )?; @@ -339,7 +461,13 @@ fn validate_schema_constraints( } 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)?; + validate_schema_constraints_at( + child, + items, + root_schema, + &format!("{path}[{index}]"), + depth + 1, + )?; } } Ok(()) @@ -456,7 +584,24 @@ impl ToolCatalog { return DispatchResult::InvalidInput(message); } } - DispatchResult::ToolResult((registration.handler)(context, dispatch_arguments)) + let result = (registration.handler)(context, dispatch_arguments); + if validation == InputValidation::Modern && !result.is_error { + let type_matches = result.structured_content_type() == registration.modern_output_type; + let value_matches = registration.validate_output.is_none_or(|validate_output| { + result + .structured_content + .as_deref() + .zip(registration.modern_output_schema.as_ref()) + .is_some_and(|(output, schema)| validate_output(output, schema).is_ok()) + }); + if !type_matches || !value_matches { + return DispatchResult::ToolResult(ToolResult::error(format!( + "tool {} emitted output that does not match its advertised schema", + registration.name + ))); + } + } + DispatchResult::ToolResult(result) } } @@ -588,6 +733,20 @@ mod tests { ), ]; + const STRUCTURED_TOOLS: [&str; 11] = [ + "icm_memory_recall", + "icm_memory_list_topics", + "icm_memory_stats", + "icm_feedback_record", + "icm_feedback_search", + "icm_feedback_stats", + "icm_transcript_start_session", + "icm_transcript_record", + "icm_transcript_search", + "icm_transcript_show", + "icm_transcript_stats", + ]; + #[test] fn annotation_projection_has_all_four_explicit_fields() { let value = ToolAnnotations::new(true, false, true, false).as_value(); @@ -662,7 +821,7 @@ mod tests { } #[test] - fn modern_projection_has_exact_annotations_and_phase_two_schema_shape() { + fn modern_projection_has_exact_annotations_and_typed_output_schemas() { let catalog = crate::tools::build_catalog(true); let projection = catalog.modern_list(); let tools = projection["tools"].as_array().unwrap(); @@ -678,13 +837,18 @@ mod tests { 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["structuredOutputFromRevision"], + if STRUCTURED_TOOLS.contains(&expected_name) { + json!("2025-06-18") + } else { + Value::Null + } + ); assert_eq!(requirements["legacyVisible"], true); let expected_embedder = if expected_name == "icm_memory_embed_all" { "required" @@ -713,6 +877,162 @@ mod tests { "unused" } ); + assert_eq!( + tool.get("outputSchema").is_some(), + STRUCTURED_TOOLS.contains(&expected_name) + ); + if let Some(schema) = tool.get("outputSchema") { + assert_eq!( + schema.get("additionalProperties"), + Some(&Value::Bool(false)) + ); + assert!(!contains_key(schema, "embedding")); + } + } + } + + #[derive(schemars::JsonSchema, serde::Deserialize, serde::Serialize)] + #[serde(deny_unknown_fields)] + struct OutputProbe { + nested: OutputProbeNested, + } + + #[derive(schemars::JsonSchema, serde::Deserialize, serde::Serialize)] + #[serde(deny_unknown_fields)] + struct OutputProbeNested { + #[schemars(length(min = 1))] + value: String, + } + + fn output_probe(name: &'static str, handler: ToolHandler) -> ToolSpec { + ToolSpec::typed::( + name, + "test-only emitted output probe", + json!({"type": "object", "properties": {}}), + None, + ToolAnnotations::new(true, false, true, false), + ToolRequirements::STORE, + handler, + ) + .with_output::() + } + + fn valid_output_probe(_: &ToolContext<'_>, _: &Value) -> ToolResult { + ToolResult::structured( + "legacy".into(), + "modern".into(), + &OutputProbe { + nested: OutputProbeNested { + value: "valid".into(), + }, + }, + ) + } + + fn malformed_output_probe(context: &ToolContext<'_>, arguments: &Value) -> ToolResult { + let mut result = valid_output_probe(context, arguments); + result.structured_content.as_mut().unwrap()["nested"]["value"] = json!(""); + result + } + + fn missing_output_probe(context: &ToolContext<'_>, arguments: &Value) -> ToolResult { + let mut result = valid_output_probe(context, arguments); + result.structured_content = None; + result + } + + #[test] + fn dispatch_validates_the_actual_emitted_output_value() { + let catalog = ToolCatalog::new( + vec![ + output_probe("valid_output", valid_output_probe), + output_probe("malformed_output", malformed_output_probe), + output_probe("missing_output", missing_output_probe), + ], + false, + ) + .unwrap(); + 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, is_error) in [ + ("valid_output", false), + ("malformed_output", true), + ("missing_output", true), + ] { + let DispatchResult::ToolResult(result) = + catalog.dispatch(&context, name, &json!({}), InputValidation::Modern) + else { + panic!("probe should reach its handler"); + }; + assert_eq!(result.is_error, is_error); + } + } + + #[test] + fn dispatch_rejects_structured_output_type_drift() { + let catalog = ToolCatalog::new( + vec![ToolSpec::typed::( + "typed_output_probe", + "test-only output type probe", + json!({"type": "object", "properties": {}}), + None, + ToolAnnotations::new(true, false, true, false), + ToolRequirements::STORE, + |_, _| { + ToolResult::structured( + "legacy".into(), + "modern".into(), + &json!({"unexpected": true}), + ) + }, + ) + .with_output::()], + false, + ) + .unwrap(); + 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 DispatchResult::ToolResult(result) = catalog.dispatch( + &context, + "typed_output_probe", + &json!({}), + InputValidation::Modern, + ) else { + panic!("probe should reach its handler"); + }; + assert!(result.is_error); + assert_eq!( + result.content[0].text, + "tool typed_output_probe emitted output that does not match its advertised schema" + ); + } + + fn contains_key(value: &Value, needle: &str) -> bool { + match value { + Value::Object(object) => { + object.contains_key(needle) + || object.values().any(|value| contains_key(value, needle)) + } + Value::Array(array) => array.iter().any(|value| contains_key(value, needle)), + _ => false, } } diff --git a/crates/icm-mcp/src/lib.rs b/crates/icm-mcp/src/lib.rs index e0fb297f..9da410f9 100644 --- a/crates/icm-mcp/src/lib.rs +++ b/crates/icm-mcp/src/lib.rs @@ -1,5 +1,6 @@ pub mod catalog; mod inputs; +mod outputs; pub mod protocol; pub mod server; pub mod service; diff --git a/crates/icm-mcp/src/outputs.rs b/crates/icm-mcp/src/outputs.rs new file mode 100644 index 00000000..7f9b006b --- /dev/null +++ b/crates/icm-mcp/src/outputs.rs @@ -0,0 +1,621 @@ +//! Typed MCP tool outputs. + +use std::collections::HashSet; + +use chrono::{DateTime, Utc}; +use icm_core::{Importance, Memory, MemorySource, Scope, StoreStats}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +const MAX_RECALL_RAW_EXCERPT_BYTES: usize = 2_048; + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(inline)] +#[serde(transparent)] +struct Nullable(Option); + +#[derive(Clone, Copy, Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(inline)] +#[serde(rename_all = "camelCase")] +pub(crate) enum SearchMode { + Hybrid, + FullText, + Keyword, +} + +#[derive(Clone, Copy, Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(inline)] +#[serde(rename_all = "lowercase")] +enum ImportanceOutput { + Critical, + High, + Medium, + Low, +} + +impl From for ImportanceOutput { + fn from(value: Importance) -> Self { + match value { + Importance::Critical => Self::Critical, + Importance::High => Self::High, + Importance::Medium => Self::Medium, + Importance::Low => Self::Low, + } + } +} + +#[derive(Clone, Copy, Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(inline)] +#[serde(rename_all = "lowercase")] +enum ScopeOutput { + User, + Project, + Org, +} + +impl From for ScopeOutput { + fn from(value: Scope) -> Self { + match value { + Scope::User => Self::User, + Scope::Project => Self::Project, + Scope::Org => Self::Org, + } + } +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(inline)] +#[serde( + tag = "type", + rename_all = "camelCase", + rename_all_fields = "camelCase", + deny_unknown_fields +)] +enum MemorySourceOutput { + Manual, + Conversation { + thread_id: String, + }, + ClaudeCode { + session_id: String, + file_path: Nullable, + }, +} + +impl From<&MemorySource> for MemorySourceOutput { + fn from(value: &MemorySource) -> Self { + match value { + MemorySource::Manual => Self::Manual, + MemorySource::Conversation { thread_id } => Self::Conversation { + thread_id: thread_id.clone(), + }, + MemorySource::ClaudeCode { + session_id, + file_path, + } => Self::ClaudeCode { + session_id: session_id.clone(), + file_path: Nullable(file_path.clone()), + }, + } + } +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(rename = "memory")] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct MemoryOutput { + #[schemars(length(min = 1))] + id: String, + created_at: DateTime, + updated_at: DateTime, + last_accessed: DateTime, + access_count: u32, + weight: f32, + topic: String, + summary: String, + raw_excerpt: Nullable, + raw_excerpt_truncated: bool, + raw_excerpt_bytes: Nullable, + keywords: Vec, + importance: ImportanceOutput, + source: MemorySourceOutput, + related_ids: Vec, + scope: ScopeOutput, + score: Nullable, +} + +impl MemoryOutput { + fn from_memory(memory: &Memory, score: Option, visible_ids: &HashSet<&str>) -> Self { + let (raw_excerpt, raw_excerpt_truncated, raw_excerpt_bytes) = + bounded_raw_excerpt(memory.raw_excerpt.as_deref()); + Self { + id: memory.id.clone(), + created_at: memory.created_at, + updated_at: memory.updated_at, + last_accessed: memory.last_accessed, + access_count: memory.access_count, + weight: memory.weight, + topic: memory.topic.clone(), + summary: memory.summary.clone(), + raw_excerpt: Nullable(raw_excerpt), + raw_excerpt_truncated, + raw_excerpt_bytes: Nullable(raw_excerpt_bytes), + keywords: memory.keywords.clone(), + importance: memory.importance.into(), + source: (&memory.source).into(), + related_ids: memory + .related_ids + .iter() + .filter(|id| visible_ids.contains(id.as_str())) + .cloned() + .collect(), + scope: memory.scope.into(), + score: Nullable(score), + } + } +} + +fn bounded_raw_excerpt(raw: Option<&str>) -> (Option, bool, Option) { + let Some(raw) = raw else { + return (None, false, None); + }; + let (excerpt, truncated) = truncate_recall_raw(raw); + (Some(excerpt.to_owned()), truncated, Some(raw.len())) +} + +pub(crate) fn truncate_recall_raw(raw: &str) -> (&str, bool) { + if raw.len() <= MAX_RECALL_RAW_EXCERPT_BYTES { + return (raw, false); + } + let mut end = MAX_RECALL_RAW_EXCERPT_BYTES; + while !raw.is_char_boundary(end) { + end -= 1; + } + (&raw[..end], true) +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub(crate) struct MemoryRecallOutput { + query: String, + effective_project: Nullable, + search_mode: SearchMode, + memories: Vec, +} + +impl MemoryRecallOutput { + pub(crate) fn new( + query: &str, + effective_project: Option<&str>, + search_mode: SearchMode, + memories: &[(Memory, f32)], + include_scores: bool, + ) -> Self { + let visible_ids: HashSet<&str> = memories + .iter() + .map(|(memory, _)| memory.id.as_str()) + .collect(); + Self { + query: query.to_owned(), + effective_project: Nullable(effective_project.map(str::to_owned)), + search_mode, + memories: memories + .iter() + .map(|(memory, score)| { + MemoryOutput::from_memory( + memory, + include_scores.then_some(*score), + &visible_ids, + ) + }) + .collect(), + } + } + + pub(crate) fn len(&self) -> usize { + self.memories.len() + } +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(inline)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct TopicCountOutput { + topic: String, + count: usize, +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub(crate) struct MemoryTopicsOutput { + topics: Vec, + total_topics: usize, + total_memories: usize, +} + +impl MemoryTopicsOutput { + pub(crate) fn new(topics: &[(String, usize)]) -> Self { + Self { + topics: topics + .iter() + .map(|(topic, count)| TopicCountOutput { + topic: topic.clone(), + count: *count, + }) + .collect(), + total_topics: topics.len(), + total_memories: topics.iter().map(|(_, count)| count).sum(), + } + } + + pub(crate) fn len(&self) -> usize { + self.topics.len() + } +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub(crate) struct MemoryStatsOutput { + total_memories: usize, + total_topics: usize, + average_weight: f32, + oldest_memory: Nullable>, + newest_memory: Nullable>, +} + +impl From for MemoryStatsOutput { + fn from(value: StoreStats) -> Self { + Self { + total_memories: value.total_memories, + total_topics: value.total_topics, + average_weight: value.avg_weight, + oldest_memory: Nullable(value.oldest_memory), + newest_memory: Nullable(value.newest_memory), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn memory_output_bounds_raw_excerpt_and_never_contains_embedding() { + let mut memory = Memory::new("topic".into(), "summary".into(), Importance::Medium); + memory.raw_excerpt = Some(format!("{}é", "x".repeat(MAX_RECALL_RAW_EXCERPT_BYTES - 1))); + memory.embedding = Some(vec![0.1, 0.2]); + let visible = Memory::new("topic".into(), "visible".into(), Importance::Medium); + memory.related_ids = vec![visible.id.clone(), "hidden-id".into()]; + + let value = serde_json::to_value(MemoryRecallOutput::new( + "query", + None, + SearchMode::FullText, + &[(memory, -1.0), (visible, -1.0)], + false, + )) + .unwrap(); + let first = &value["memories"][0]; + assert_eq!(first["rawExcerptTruncated"], true); + assert_eq!(first["rawExcerptBytes"], MAX_RECALL_RAW_EXCERPT_BYTES + 1); + assert!(first["rawExcerpt"] + .as_str() + .unwrap() + .is_char_boundary(2_047)); + assert!(first.get("embedding").is_none()); + assert_eq!( + first["relatedIds"], + serde_json::json!([value["memories"][1]["id"]]) + ); + } +} + +use icm_core::{Feedback, FeedbackStats, Message, Role, Session, TranscriptHit, TranscriptStats}; + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub(crate) struct TranscriptStartOutput { + #[schemars(length(min = 1))] + session_id: String, +} + +impl TranscriptStartOutput { + pub(crate) fn new(session_id: String) -> Self { + Self { session_id } + } +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub(crate) struct TranscriptRecordOutput { + #[schemars(length(min = 1))] + message_id: String, +} + +impl TranscriptRecordOutput { + pub(crate) fn new(message_id: String) -> Self { + Self { message_id } + } +} + +#[derive(Clone, Copy, Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(inline)] +#[serde(rename_all = "lowercase")] +enum TranscriptRoleOutput { + User, + Assistant, + System, + Tool, +} + +impl From for TranscriptRoleOutput { + fn from(value: Role) -> Self { + match value { + Role::User => Self::User, + Role::Assistant => Self::Assistant, + Role::System => Self::System, + Role::Tool => Self::Tool, + } + } +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(rename = "message")] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct TranscriptMessageOutput { + id: String, + session_id: String, + role: TranscriptRoleOutput, + content: String, + tool_name: Nullable, + tokens: Nullable, + timestamp: DateTime, + metadata: String, +} + +impl From<&Message> for TranscriptMessageOutput { + fn from(value: &Message) -> Self { + Self { + id: value.id.clone(), + session_id: value.session_id.clone(), + role: value.role.into(), + content: value.content.clone(), + tool_name: Nullable(value.tool_name.clone()), + tokens: Nullable(value.tokens), + timestamp: value.ts, + metadata: value.metadata.clone(), + } + } +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(rename = "session")] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct TranscriptSessionOutput { + id: String, + agent: String, + project: Nullable, + started_at: DateTime, + updated_at: DateTime, + metadata: String, +} + +impl From<&Session> for TranscriptSessionOutput { + fn from(value: &Session) -> Self { + Self { + id: value.id.clone(), + agent: value.agent.clone(), + project: Nullable(value.project.clone()), + started_at: value.started_at, + updated_at: value.updated_at, + metadata: value.metadata.clone(), + } + } +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(inline)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct TranscriptHitOutput { + message: TranscriptMessageOutput, + session: TranscriptSessionOutput, + score: f64, +} + +impl From<&TranscriptHit> for TranscriptHitOutput { + fn from(value: &TranscriptHit) -> Self { + Self { + message: (&value.message).into(), + session: (&value.session).into(), + score: value.score, + } + } +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub(crate) struct TranscriptSearchOutput { + hits: Vec, +} + +impl TranscriptSearchOutput { + pub(crate) fn new(hits: &[TranscriptHit]) -> Self { + Self { + hits: hits.iter().map(Into::into).collect(), + } + } + + pub(crate) fn len(&self) -> usize { + self.hits.len() + } +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub(crate) struct TranscriptShowOutput { + session: TranscriptSessionOutput, + messages: Vec, +} + +impl TranscriptShowOutput { + pub(crate) fn new(session: &Session, messages: &[Message]) -> Self { + Self { + session: session.into(), + messages: messages.iter().map(Into::into).collect(), + } + } + + pub(crate) fn len(&self) -> usize { + self.messages.len() + } +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(inline)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct RoleCountOutput { + role: String, + count: usize, +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(inline)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct AgentCountOutput { + agent: String, + count: usize, +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(inline)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct SessionCountOutput { + session_id: String, + message_count: usize, +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub(crate) struct TranscriptStatsOutput { + total_sessions: usize, + total_messages: usize, + total_bytes: u64, + by_role: Vec, + by_agent: Vec, + top_sessions: Vec, + oldest: Nullable>, + newest: Nullable>, +} + +impl From for TranscriptStatsOutput { + fn from(value: TranscriptStats) -> Self { + Self { + total_sessions: value.total_sessions, + total_messages: value.total_messages, + total_bytes: value.total_bytes, + by_role: value + .by_role + .into_iter() + .map(|(role, count)| RoleCountOutput { role, count }) + .collect(), + by_agent: value + .by_agent + .into_iter() + .map(|(agent, count)| AgentCountOutput { agent, count }) + .collect(), + top_sessions: value + .top_sessions + .into_iter() + .map(|(session_id, message_count)| SessionCountOutput { + session_id, + message_count, + }) + .collect(), + oldest: Nullable(value.oldest), + newest: Nullable(value.newest), + } + } +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(rename = "feedback")] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub(crate) struct FeedbackOutput { + id: String, + topic: String, + context: String, + predicted: String, + corrected: String, + reason: Nullable, + source: String, + created_at: DateTime, + applied_count: u32, +} + +impl From<&Feedback> for FeedbackOutput { + fn from(value: &Feedback) -> Self { + Self { + id: value.id.clone(), + topic: value.topic.clone(), + context: value.context.clone(), + predicted: value.predicted.clone(), + corrected: value.corrected.clone(), + reason: Nullable(value.reason.clone()), + source: value.source.clone(), + created_at: value.created_at, + applied_count: value.applied_count, + } + } +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub(crate) struct FeedbackSearchOutput { + feedback: Vec, +} + +impl FeedbackSearchOutput { + pub(crate) fn new(feedback: &[Feedback]) -> Self { + Self { + feedback: feedback.iter().map(Into::into).collect(), + } + } + + pub(crate) fn len(&self) -> usize { + self.feedback.len() + } +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[schemars(inline)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct AppliedCountOutput { + feedback_id: String, + count: u32, +} + +#[derive(Debug, Deserialize, JsonSchema, Serialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub(crate) struct FeedbackStatsOutput { + total: usize, + by_topic: Vec, + most_applied: Vec, +} + +impl From for FeedbackStatsOutput { + fn from(value: FeedbackStats) -> Self { + Self { + total: value.total, + by_topic: value + .by_topic + .into_iter() + .map(|(topic, count)| TopicCountOutput { topic, count }) + .collect(), + most_applied: value + .most_applied + .into_iter() + .map(|(feedback_id, count)| AppliedCountOutput { feedback_id, count }) + .collect(), + } + } +} diff --git a/crates/icm-mcp/src/protocol.rs b/crates/icm-mcp/src/protocol.rs index 03d32e0a..770c8245 100644 --- a/crates/icm-mcp/src/protocol.rs +++ b/crates/icm-mcp/src/protocol.rs @@ -1,3 +1,5 @@ +use std::any::TypeId; + use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; @@ -142,8 +144,14 @@ impl JsonRpcResponse { #[derive(Debug, Serialize)] pub struct ToolResult { pub content: Vec, + #[serde(rename = "structuredContent", skip_serializing_if = "Option::is_none")] + pub structured_content: Option>, #[serde(rename = "isError", skip_serializing_if = "std::ops::Not::not")] pub is_error: bool, + #[serde(skip)] + modern_text: Option, + #[serde(skip)] + structured_content_type: Option, } #[derive(Debug, Serialize)] @@ -160,7 +168,29 @@ impl ToolResult { content_type: "text".into(), text, }], + structured_content: None, is_error: false, + modern_text: None, + structured_content_type: None, + } + } + + pub fn structured(legacy_text: String, modern_text: String, output: &T) -> Self + where + T: Serialize + 'static, + { + match serde_json::to_value(output) { + Ok(structured_content) => Self { + content: vec![TextContent { + content_type: "text".into(), + text: legacy_text, + }], + structured_content: Some(Box::new(structured_content)), + is_error: false, + modern_text: Some(modern_text), + structured_content_type: Some(TypeId::of::()), + }, + Err(error) => Self::error(format!("structured output serialization failed: {error}")), } } @@ -170,8 +200,28 @@ impl ToolResult { content_type: "text".into(), text, }], + structured_content: None, is_error: true, + modern_text: None, + structured_content_type: None, + } + } + + pub(crate) fn structured_content_type(&self) -> Option { + self.structured_content_type + } + + pub fn select_projection(&mut self, modern: bool) { + if modern { + if let Some(text) = self.modern_text.take() { + if let Some(content) = self.content.last_mut() { + content.text = text; + } + } + } else { + self.structured_content = None; } + self.modern_text = None; } /// Append a hint to the last text content block. @@ -242,7 +292,10 @@ mod tests { fn test_append_hint_empty_content() { let mut result = ToolResult { content: vec![], + structured_content: None, is_error: false, + modern_text: None, + structured_content_type: None, }; result.append_hint("[hint]"); assert!(result.content.is_empty()); diff --git a/crates/icm-mcp/src/service.rs b/crates/icm-mcp/src/service.rs index 60923c77..68774cf8 100644 --- a/crates/icm-mcp/src/service.rs +++ b/crates/icm-mcp/src/service.rs @@ -3,8 +3,9 @@ use std::collections::HashSet; use std::path::PathBuf; -use icm_core::Embedder; +use icm_core::{project::project_from_path, Embedder, IcmError, IcmResult, Memory}; use icm_store::Store; +use serde::Serialize; use serde_json::{json, Map, Value}; use crate::catalog::{DispatchResult, InputValidation, ToolCatalog, ToolContext}; @@ -23,6 +24,11 @@ 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; +const ACTIVE_PROJECT_CONTEXT_URI: &str = "icm://active-project/context"; +const ACTIVE_PROJECT_CONTEXT_MIME_TYPE: &str = "application/json"; +const RESOURCE_ROW_LIMIT: usize = 64; +const RESOURCE_FIELD_BYTES: usize = 512; +const RESOURCE_MAX_BYTES: usize = 2048; pub const ICM_INSTRUCTIONS: &str = "\ Use ICM (Infinite Context Memory) proactively to maintain long-term memory across sessions.\n\ @@ -80,12 +86,28 @@ impl Default for ConnectionState { } } +impl ConnectionState { + /// Start a stateless HTTP request in the frozen 2024 compatibility era. + /// Stdio still begins uninitialized; only transports that already carry + /// the protocol revision out-of-band should use this constructor. + pub fn legacy_2024_ready() -> Self { + Self { + phase: ConnectionPhase::LegacyReady { + revision: ProtocolRevision::V2024_11_05, + initialized_seen: false, + }, + calls_since_store: 0, + } + } +} + pub struct McpService<'a> { store: &'a Store, embedder: Option<&'a dyn Embedder>, compact: bool, auto_consolidate: AutoConsolidate, working_directory: PathBuf, + active_project: Option, catalog: ToolCatalog, } @@ -113,12 +135,17 @@ impl<'a> McpService<'a> { auto_consolidate: AutoConsolidate, working_directory: PathBuf, ) -> Self { + let active_project = working_directory + .to_str() + .and_then(project_from_path) + .filter(|project| valid_resource_project(project)); Self { store, embedder, compact, auto_consolidate, working_directory, + active_project, catalog: tools::build_catalog(embedder.is_some()), } } @@ -448,50 +475,132 @@ impl<'a> McpService<'a> { "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) + self.list_resources(id, revision, message) } "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 })), - ) + self.read_resource(id, revision, message) } other => JsonRpcResponse::method_not_found(id, other), } } + fn list_resources( + &self, + id: Value, + revision: ProtocolRevision, + message: &JsonRpcMessage, + ) -> JsonRpcResponse { + 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, + "resources/list cursor is not supported".into(), + ); + } + + let resources = self + .active_project + .as_ref() + .map(|_| { + json!({ + "uri": ACTIVE_PROJECT_CONTEXT_URI, + "name": "active-project-context", + "title": "Active Project Context", + "description": "Stored context for the project inferred from the MCP server working directory.", + "mimeType": ACTIVE_PROJECT_CONTEXT_MIME_TYPE, + "annotations": { "audience": ["assistant"], "priority": 1.0 } + }) + }) + .into_iter() + .collect::>(); + let result = if revision == ProtocolRevision::V2026_07_28 { + project_result( + revision, + json!({ "resources": resources }), + Some((3_600_000, "private")), + ) + } else { + json!({ + "resources": resources, + "_meta": { "ttlMs": 0, "cacheScope": "private" } + }) + }; + JsonRpcResponse::ok(id, result) + } + + fn read_resource( + &self, + 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, + "resources/read params must be an object".into(), + ); + }; + if params + .keys() + .any(|key| !matches!(key.as_str(), "uri" | "_meta")) + { + return JsonRpcResponse::err( + id, + -32602, + "resources/read accepts only uri and _meta".into(), + ); + } + let Some(uri) = params.get("uri").and_then(Value::as_str) else { + return JsonRpcResponse::err(id, -32602, "resources/read.uri must be a string".into()); + }; + let Some(project) = self + .active_project + .as_deref() + .filter(|_| uri == ACTIVE_PROJECT_CONTEXT_URI) + else { + let code = if revision == ProtocolRevision::V2026_07_28 { + -32602 + } else { + -32002 + }; + return JsonRpcResponse::err_with_data( + id, + code, + "resource not found".into(), + Some(json!({ "uri": uri })), + ); + }; + + let text = match active_project_context(self.store, project) { + Ok(text) => text, + Err(error) => { + tracing::warn!(%error, "failed to read active-project MCP resource"); + return JsonRpcResponse::err(id, -32603, "failed to read resource".into()); + } + }; + let value = json!({ + "contents": [{ + "uri": ACTIVE_PROJECT_CONTEXT_URI, + "mimeType": ACTIVE_PROJECT_CONTEXT_MIME_TYPE, + "text": text, + }] + }); + let result = if revision == ProtocolRevision::V2026_07_28 { + project_result(revision, value, Some((0, "private"))) + } else { + let mut value = value; + value["_meta"] = json!({ "ttlMs": 0, "cacheScope": "private" }); + value + }; + JsonRpcResponse::ok(id, result) + } + fn list_tools( &self, id: Value, @@ -573,7 +682,10 @@ impl<'a> McpService<'a> { .catalog .dispatch(&context, name, &arguments, validation) { - DispatchResult::ToolResult(result) => result, + DispatchResult::ToolResult(mut result) => { + result.select_projection(revision != ProtocolRevision::V2024_11_05); + result + } DispatchResult::UnknownTool if revision == ProtocolRevision::V2024_11_05 => { crate::protocol::ToolResult::error(format!("unknown tool: {name}")) } @@ -744,15 +856,7 @@ fn validate_modern_request( )); }; 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, - })), - ))); + return Err(Box::new(unsupported_protocol_version_error(id, requested))); } let Some(capabilities) = metadata @@ -778,6 +882,19 @@ fn validate_modern_request( validate_optional_metadata_values(&id, metadata) } +/// Build the protocol-defined error shared by MCP services and transports. +pub fn unsupported_protocol_version_error(id: Value, requested: &str) -> JsonRpcResponse { + JsonRpcResponse::err_with_data( + id, + -32022, + format!("unsupported protocol version: {requested}"), + Some(json!({ + "supported": SUPPORTED_PROTOCOL_VERSIONS, + "requested": requested, + })), + ) +} + 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()); @@ -1363,6 +1480,152 @@ fn trim_ows(value: &str) -> &str { value.trim_matches(|character| matches!(character, ' ' | '\t')) } +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct ActiveProjectContext { + project: String, + topics: Vec, + memories: Vec, + truncated: bool, + truncation_reasons: Vec<&'static str>, + omitted_at_least: usize, + budget: ResourceBudget, +} + +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct ResourceMemory { + id: String, + topic: String, + summary: String, + importance: String, + weight: f32, + updated_at: String, + field_truncated: bool, +} + +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct ResourceBudget { + max_portable_tokens: usize, + used_portable_tokens: usize, + algorithm: &'static str, +} + +fn active_project_context(store: &Store, project: &str) -> IcmResult { + let mut context = empty_resource_context(project); + let topic_refs = context + .topics + .iter() + .map(String::as_str) + .collect::>(); + let fetched = store.get_by_topics_limited(&topic_refs, RESOURCE_ROW_LIMIT + 1)?; + let fetched_count = fetched.len(); + let memories = fetched + .into_iter() + .take(RESOURCE_ROW_LIMIT) + .map(resource_memory) + .collect::>(); + let field_limited = memories.iter().any(|memory| memory.field_truncated); + let row_limited = fetched_count > RESOURCE_ROW_LIMIT; + let mut truncation_reasons = Vec::new(); + if row_limited { + truncation_reasons.push("rowLimit"); + } + if field_limited { + truncation_reasons.push("fieldLimit"); + } + context.omitted_at_least = fetched_count.saturating_sub(memories.len()); + context.truncated = !truncation_reasons.is_empty(); + context.truncation_reasons = truncation_reasons; + context.memories = memories; + + loop { + let text = serialize_resource_context(&mut context)?; + if text.len() <= RESOURCE_MAX_BYTES { + return Ok(text); + } + if !context.truncation_reasons.contains(&"tokenBudget") { + context.truncation_reasons.push("tokenBudget"); + } + context.truncated = true; + if context.memories.pop().is_none() { + return Err(IcmError::InvalidInput( + "active-project resource metadata exceeds its fixed budget".into(), + )); + } + context.omitted_at_least = fetched_count.saturating_sub(context.memories.len()); + } +} + +fn empty_resource_context(project: &str) -> ActiveProjectContext { + ActiveProjectContext { + project: project.into(), + topics: vec![ + format!("context-{project}"), + format!("contexte-{project}"), + format!("decisions-{project}"), + ], + memories: Vec::new(), + truncated: false, + truncation_reasons: Vec::new(), + omitted_at_least: 0, + budget: ResourceBudget { + max_portable_tokens: RESOURCE_MAX_BYTES, + used_portable_tokens: 0, + algorithm: "utf8-bytes-v1", + }, + } +} + +fn resource_memory(memory: Memory) -> ResourceMemory { + let (id, id_truncated) = truncate_resource_field(&memory.id); + let (summary, summary_truncated) = truncate_resource_field(&memory.summary); + ResourceMemory { + id, + topic: memory.topic, + summary, + importance: memory.importance.to_string(), + weight: memory.weight, + updated_at: memory.updated_at.to_rfc3339(), + field_truncated: id_truncated || summary_truncated, + } +} + +fn truncate_resource_field(value: &str) -> (String, bool) { + if value.len() <= RESOURCE_FIELD_BYTES { + return (value.into(), false); + } + let mut end = RESOURCE_FIELD_BYTES; + while !value.is_char_boundary(end) { + end -= 1; + } + (value[..end].into(), true) +} + +fn serialize_resource_context(context: &mut ActiveProjectContext) -> IcmResult { + loop { + let text = serde_json::to_string_pretty(context)?; + let used = text.len(); + if context.budget.used_portable_tokens == used { + return Ok(text); + } + context.budget.used_portable_tokens = used; + } +} + +fn valid_resource_project(project: &str) -> bool { + if project.is_empty() + || project.len() > 246 + || project.trim() != project + || project.chars().any(char::is_control) + { + return false; + } + serialize_resource_context(&mut empty_resource_context(project)) + .is_ok_and(|text| text.len() <= RESOURCE_MAX_BYTES) +} + fn server_info() -> Value { json!({ "name": SERVER_NAME, "version": SERVER_VERSION }) } @@ -2549,11 +2812,27 @@ mod tests { } #[test] - fn resource_core_is_honest_and_empty() { + fn active_project_resource_is_fixed_scoped_and_bounded() { let store = Store::in_memory().unwrap(); - let service = service(&store); + for index in 0..66 { + let mut memory = Memory::new( + "context-test-project".into(), + format!( + "row {index:02} prompt boundary\n--- RESOURCE-FORGE --- {}", + "bounded ".repeat(90) + ), + Importance::High, + ); + memory.id = format!("01R{index:023}"); + memory.weight = 1.0 - index as f32 / 1_000.0; + memory.access_count = 2; + store.store(memory).unwrap(); + } + + let mut service = service(&store); + service.active_project = Some("test-project".into()); let mut state = ConnectionState::default(); - let response = service + let listed = service .handle( &mut state, request(json!({ @@ -2565,7 +2844,179 @@ mod tests { })), ) .unwrap(); - assert_eq!(response.result.unwrap()["resources"], json!([])); + let listed = listed.result.unwrap(); + assert_eq!(listed["ttlMs"], 3_600_000); + assert_eq!(listed["cacheScope"], "private"); + assert_eq!(listed["resources"][0]["uri"], ACTIVE_PROJECT_CONTEXT_URI); + assert_eq!(listed["resources"][0]["mimeType"], "application/json"); + + let mut state_2025 = + initialized_state_for_revision(&service, ProtocolRevision::V2025_11_25); + let listed_2025 = service + .handle( + &mut state_2025, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"resources/list","params":{} + })), + ) + .unwrap() + .result + .unwrap(); + assert_eq!(listed_2025["_meta"]["ttlMs"], 0); + assert_eq!(listed_2025["_meta"]["cacheScope"], "private"); + + let read = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"resources/read", + "params":{"uri":ACTIVE_PROJECT_CONTEXT_URI,"_meta":modern_metadata()} + })), + ) + .unwrap() + .result + .unwrap(); + assert_eq!(read["ttlMs"], 0); + assert_eq!(read["cacheScope"], "private"); + let text = read["contents"][0]["text"].as_str().unwrap(); + assert!(text.len() <= RESOURCE_MAX_BYTES); + assert!(text.contains("\\n--- RESOURCE-FORGE")); + let context: Value = serde_json::from_str(text).unwrap(); + assert_eq!(context["project"], "test-project"); + assert_eq!( + context["topics"], + json!([ + "context-test-project", + "contexte-test-project", + "decisions-test-project" + ]) + ); + assert_eq!(context["budget"]["usedPortableTokens"], text.len()); + assert_eq!(context["budget"]["algorithm"], "utf8-bytes-v1"); + assert_eq!(context["truncated"], true); + for reason in ["rowLimit", "fieldLimit", "tokenBudget"] { + assert!(context["truncationReasons"] + .as_array() + .unwrap() + .contains(&json!(reason))); + } + assert!(context["memories"] + .as_array() + .unwrap() + .iter() + .any(|memory| memory["fieldTruncated"] == true)); + assert_eq!( + store + .get("01R00000000000000000000000") + .unwrap() + .unwrap() + .access_count, + 2 + ); + } + + #[test] + fn active_project_resource_is_empty_and_cross_project_isolated_in_2025() { + let store = Store::in_memory().unwrap(); + let mut service = service(&store); + service.active_project = Some("test-project".into()); + let mut state = initialized_state_for_revision(&service, ProtocolRevision::V2025_11_25); + + let read = |service: &McpService<'_>, state: &mut ConnectionState, id| { + service + .handle( + state, + request(json!({ + "jsonrpc":"2.0","id":id,"method":"resources/read", + "params":{"uri":ACTIVE_PROJECT_CONTEXT_URI} + })), + ) + .unwrap() + .result + .unwrap() + }; + let empty = read(&service, &mut state, 2); + assert_eq!(empty["_meta"], json!({"ttlMs":0,"cacheScope":"private"})); + let context: Value = + serde_json::from_str(empty["contents"][0]["text"].as_str().unwrap()).unwrap(); + assert_eq!(context["memories"], json!([])); + assert_eq!(context["truncated"], false); + + for (topic, summary) in [ + ("context-test-project", "included"), + ("context-other-project", "excluded project"), + ("preferences", "excluded global"), + ] { + store + .store(Memory::new(topic.into(), summary.into(), Importance::High)) + .unwrap(); + } + let populated = read(&service, &mut state, 3); + let context: Value = + serde_json::from_str(populated["contents"][0]["text"].as_str().unwrap()).unwrap(); + let memories = context["memories"].as_array().unwrap(); + assert_eq!(memories.len(), 1); + assert_eq!(memories[0]["summary"], "included"); + } + + #[test] + fn active_project_resource_rejects_bad_uris_and_hides_internal_failures() { + let store = Store::in_memory().unwrap(); + let mut service = service(&store); + service.active_project = Some("test-project".into()); + + let mut legacy = initialized_state_for_revision(&service, ProtocolRevision::V2024_11_05); + let unavailable = service + .handle( + &mut legacy, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"resources/read", + "params":{"uri":ACTIVE_PROJECT_CONTEXT_URI} + })), + ) + .unwrap(); + assert_eq!(unavailable.error.unwrap().code, -32601); + + for (id, mut params) in [ + (3, json!({})), + (4, json!({"uri":42})), + ( + 5, + json!({"uri":"icm://active-project/context?unexpected=1"}), + ), + ] { + params + .as_object_mut() + .unwrap() + .insert("_meta".into(), modern_metadata()); + let mut modern = ConnectionState::default(); + let rejected = service + .handle( + &mut modern, + request(json!({ + "jsonrpc":"2.0","id":id,"method":"resources/read", + "params":params + })), + ) + .unwrap(); + assert_eq!(rejected.error.unwrap().code, -32602); + } + + service.active_project = Some("x".repeat(RESOURCE_MAX_BYTES)); + let mut modern = ConnectionState::default(); + let failed = service + .handle( + &mut modern, + request(json!({ + "jsonrpc":"2.0","id":6,"method":"resources/read", + "params":{"uri":ACTIVE_PROJECT_CONTEXT_URI,"_meta":modern_metadata()} + })), + ) + .unwrap(); + let error = failed.error.unwrap(); + assert_eq!(error.code, -32603); + assert_eq!(error.message, "failed to read resource"); + assert!(error.data.is_none()); } #[test] @@ -2612,9 +3063,9 @@ mod tests { .unwrap() .result .unwrap(); - let text = result["content"][0]["text"].as_str().unwrap(); - assert!(text.contains("from client")); - assert!(!text.contains("from other")); + let memories = result["structuredContent"]["memories"].as_array().unwrap(); + assert_eq!(memories.len(), 1); + assert_eq!(memories[0]["summary"], "shared marker from client"); } #[test] @@ -2722,6 +3173,60 @@ mod tests { ); } + #[test] + fn typed_outputs_follow_the_negotiated_projection() { + 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 result = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":2,"method":"tools/call", + "params":{"name":"icm_memory_stats","arguments":{}} + })), + ) + .unwrap() + .result + .unwrap(); + if revision == ProtocolRevision::V2024_11_05 { + assert!(result["content"][0]["text"] + .as_str() + .unwrap() + .starts_with("Memories: 0\nTopics: 0\n")); + assert!(result.get("structuredContent").is_none()); + } else { + assert_eq!(result["content"][0]["text"], "Returned memory statistics."); + assert_eq!(result["structuredContent"]["totalMemories"], 0); + } + } + + let mut state = ConnectionState::default(); + let result = service + .handle( + &mut state, + request(json!({ + "jsonrpc":"2.0","id":3,"method":"tools/call", + "params":{ + "name":"icm_memory_stats","arguments":{}, + "_meta":modern_metadata() + } + })), + ) + .unwrap() + .result + .unwrap(); + assert_eq!(result["content"][0]["text"], "Returned memory statistics."); + assert_eq!(result["structuredContent"]["totalMemories"], 0); + assert_eq!(result["resultType"], "complete"); + } + #[test] fn topic_schema_and_runtime_enforce_the_multibyte_boundary() { let store = Store::in_memory().unwrap(); @@ -2901,13 +3406,19 @@ mod tests { .unwrap(); let accepted_result = accepted.result.unwrap(); assert_ne!(accepted_result["isError"], true); + assert_eq!(accepted_result["content"][0]["text"], "Found 30 memories."); assert_eq!( - accepted_result["content"][0]["text"] - .as_str() - .unwrap() - .matches("revision limit probe") - .count(), - 30 + accepted_result["structuredContent"]["memories"] + .as_array() + .map(Vec::len), + Some(30) + ); + let first = &accepted_result["structuredContent"]["memories"][0]; + let stored = store.get(first["id"].as_str().unwrap()).unwrap().unwrap(); + assert_eq!(first["accessCount"], stored.access_count); + assert_eq!( + first["lastAccessed"], + serde_json::to_value(stored.last_accessed).unwrap() ); for (id, limit) in [(8, 0), (9, 101)] { diff --git a/crates/icm-mcp/src/tools.rs b/crates/icm-mcp/src/tools.rs index b7948d9a..89794a5a 100644 --- a/crates/icm-mcp/src/tools.rs +++ b/crates/icm-mcp/src/tools.rs @@ -67,7 +67,10 @@ pub fn call_tool_with_config( args, crate::catalog::InputValidation::Legacy2024Unchecked, ) { - DispatchResult::ToolResult(result) => result, + DispatchResult::ToolResult(mut result) => { + result.select_projection(false); + result + } DispatchResult::UnknownTool => ToolResult::error(format!("unknown tool: {name}")), DispatchResult::InvalidInput(_) => { unreachable!("unchecked legacy compatibility dispatch cannot reject typed inputs") diff --git a/crates/icm-mcp/src/tools/handlers/common.rs b/crates/icm-mcp/src/tools/handlers/common.rs index b29fdd2b..de34254d 100644 --- a/crates/icm-mcp/src/tools/handlers/common.rs +++ b/crates/icm-mcp/src/tools/handlers/common.rs @@ -2,7 +2,10 @@ use serde_json::Value; -use icm_core::{Embedder, Memoir, MemoirStore, Memory}; +use icm_core::{ + is_preference_topic, keyword_matches, project_matches, topic_matches, Embedder, Memoir, + MemoirStore, Memory, +}; use icm_store::Store; use crate::protocol::ToolResult; @@ -105,6 +108,27 @@ pub fn get_i64(args: &Value, key: &str, default: i64) -> i64 { args.get(key).and_then(|v| v.as_i64()).unwrap_or(default) } +/// Apply the shared project/topic/keyword scope used by every recall path. +/// Flatten untrusted text before embedding it into line-oriented output. +pub fn flatten_untrusted_text(value: &str) -> String { + value.replace(['\n', '\r'], " ") +} + +pub fn matches_memory_filters( + memory: &Memory, + project: Option<&str>, + topic: Option<&str>, + keyword: Option<&str>, +) -> bool { + if let Some(project) = project { + if !is_preference_topic(&memory.topic) && !project_matches(&memory.topic, Some(project)) { + return false; + } + } + topic.is_none_or(|value| topic_matches(&memory.topic, value)) + && keyword.is_none_or(|value| keyword_matches(&memory.keywords, value)) +} + pub fn resolve_memoir(store: &Store, name: &str) -> Result { store .get_memoir_by_name(name) @@ -120,15 +144,18 @@ pub fn format_memory_output(memories: &[(Memory, f32)], compact: bool) -> String // 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))); + output.push_str(&format!( + "[{}] {}\n", + mem.topic, + flatten_untrusted_text(&mem.summary) + )); } } else { for (mem, score) in memories { - let summary = flatten(&mem.summary); + let summary = flatten_untrusted_text(&mem.summary); if *score >= 0.0 { output.push_str(&format!( "--- {} [score: {:.3}] ---\n topic: {}\n importance: {}\n weight: {:.3}\n summary: {}\n", @@ -141,8 +168,11 @@ pub fn format_memory_output(memories: &[(Memory, f32)], compact: bool) -> String )); } if !mem.keywords.is_empty() { - let flattened_keywords: Vec = - mem.keywords.iter().map(|k| flatten(k)).collect(); + let flattened_keywords: Vec = mem + .keywords + .iter() + .map(|k| flatten_untrusted_text(k)) + .collect(); output.push_str(&format!(" keywords: {}\n", flattened_keywords.join(", "))); } if let Some(ref raw) = mem.raw_excerpt { @@ -150,15 +180,11 @@ pub fn format_memory_output(memories: &[(Memory, f32)], compact: bool) -> String // 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; - } + let (excerpt, truncated) = crate::outputs::truncate_recall_raw(raw); + if truncated { output.push_str(&format!( " raw: {}… [truncated, {} bytes total]\n", - &raw[..cut], + excerpt, raw.len() )); } else { diff --git a/crates/icm-mcp/src/tools/handlers/feedback.rs b/crates/icm-mcp/src/tools/handlers/feedback.rs index 653faa96..fafb6dc1 100644 --- a/crates/icm-mcp/src/tools/handlers/feedback.rs +++ b/crates/icm-mcp/src/tools/handlers/feedback.rs @@ -5,9 +5,10 @@ use serde_json::Value; use icm_core::{Embedder, Feedback, FeedbackStore}; use icm_store::Store; +use crate::outputs::{FeedbackOutput, FeedbackSearchOutput, FeedbackStatsOutput}; use crate::protocol::ToolResult; -use super::common::{get_i64, get_str, MAX_FEEDBACK_FIELD_LEN}; +use super::common::{flatten_untrusted_text, get_i64, get_str, MAX_FEEDBACK_FIELD_LEN}; pub(in crate::tools) fn tool_feedback_record( store: &Store, @@ -67,13 +68,15 @@ pub(in crate::tools) fn tool_feedback_record( } let id = feedback.id.clone(); + let structured = FeedbackOutput::from(&feedback); match store.store_feedback(feedback) { Ok(_) => { - if compact { - ToolResult::text(format!("ok {id}")) + let legacy = if compact { + format!("ok {id}") } else { - ToolResult::text(format!("Feedback recorded: {id}\n topic: {topic}\n predicted: {predicted}\n corrected: {corrected}")) - } + format!("Feedback recorded: {id}\n topic: {topic}\n predicted: {predicted}\n corrected: {corrected}") + }; + ToolResult::structured(legacy, "Feedback recorded.".into(), &structured) } Err(e) => ToolResult::error(format!("failed to store feedback: {e}")), } @@ -94,8 +97,13 @@ pub(in crate::tools) fn tool_feedback_search( match store.search_feedback(query, query_embedding.as_deref(), topic, limit) { Ok(results) => { + let structured = FeedbackSearchOutput::new(&results); if results.is_empty() { - return ToolResult::text("No feedback found.".into()); + return ToolResult::structured( + "No feedback found.".into(), + "Found 0 feedback entries.".into(), + &structured, + ); } // context/predicted/corrected/reason/source can originate from // untrusted content (a feedback entry recorded from tool output @@ -103,28 +111,34 @@ pub(in crate::tools) fn tool_feedback_search( // value can't forge a fake "--- id [topic] ---" delimiter and // inject a spoofed entry into this output (same injection class // already fixed in recall_context/build_consolidate_prompt). - let flatten = |s: &str| s.replace(['\n', '\r'], " "); let mut output = String::new(); for fb in &results { output.push_str(&format!( "--- {} [{}] ---\n context: {}\n predicted: {}\n corrected: {}\n", fb.id, - flatten(&fb.topic), - flatten(&fb.context), - flatten(&fb.predicted), - flatten(&fb.corrected) + flatten_untrusted_text(&fb.topic), + flatten_untrusted_text(&fb.context), + flatten_untrusted_text(&fb.predicted), + flatten_untrusted_text(&fb.corrected) )); if let Some(ref reason) = fb.reason { - output.push_str(&format!(" reason: {}\n", flatten(reason))); + output.push_str(&format!(" reason: {}\n", flatten_untrusted_text(reason))); } if !fb.source.is_empty() { - output.push_str(&format!(" source: {}\n", flatten(&fb.source))); + output.push_str(&format!( + " source: {}\n", + flatten_untrusted_text(&fb.source) + )); } if fb.applied_count > 0 { output.push_str(&format!(" applied: {} times\n", fb.applied_count)); } } - ToolResult::text(output) + ToolResult::structured( + output, + format!("Found {} feedback entries.", structured.len()), + &structured, + ) } Err(e) => ToolResult::error(format!("failed to search feedback: {e}")), } @@ -133,6 +147,7 @@ pub(in crate::tools) fn tool_feedback_search( pub(in crate::tools) fn tool_feedback_stats(store: &Store) -> ToolResult { match store.feedback_stats() { Ok(stats) => { + let structured = FeedbackStatsOutput::from(stats.clone()); let mut output = format!("Feedback total: {}\n", stats.total); if !stats.by_topic.is_empty() { output.push_str("\nBy topic:\n"); @@ -146,7 +161,7 @@ pub(in crate::tools) fn tool_feedback_stats(store: &Store) -> ToolResult { output.push_str(&format!(" {id}: {count} times\n")); } } - ToolResult::text(output) + ToolResult::structured(output, "Returned feedback statistics.".into(), &structured) } Err(e) => ToolResult::error(format!("failed to get feedback stats: {e}")), } diff --git a/crates/icm-mcp/src/tools/handlers/memory.rs b/crates/icm-mcp/src/tools/handlers/memory.rs index 7f8abfaf..d8f3fda9 100644 --- a/crates/icm-mcp/src/tools/handlers/memory.rs +++ b/crates/icm-mcp/src/tools/handlers/memory.rs @@ -5,18 +5,18 @@ 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, - Embedder, Memory, MemoryStore, WakeUpFormat, WakeUpOptions, DEDUP_SIMILARITY_THRESHOLD, - MSG_NO_MEMORIES, + AutoLinkOptions, Embedder, Memory, MemoryStore, WakeUpFormat, WakeUpOptions, + DEDUP_SIMILARITY_THRESHOLD, MSG_NO_MEMORIES, }; use icm_store::Store; use crate::catalog::ToolContext; +use crate::outputs::{MemoryRecallOutput, MemoryStatsOutput, MemoryTopicsOutput, SearchMode}; use crate::protocol::ToolResult; use super::common::{ - format_memory_output, get_i64, get_str, parse_keywords, resolve_memoir, try_auto_consolidate, - AutoConsolidate, MAX_CONTENT_LEN, MAX_TOPIC_LEN, + format_memory_output, get_i64, get_str, matches_memory_filters, parse_keywords, resolve_memoir, + try_auto_consolidate, AutoConsolidate, MAX_CONTENT_LEN, MAX_TOPIC_LEN, }; pub(in crate::tools) fn tool_wake_up(store: &Store, args: &Value) -> ToolResult { @@ -250,6 +250,43 @@ pub(in crate::tools) fn tool_store( } } +fn recall_result( + query: &str, + project: Option<&str>, + search_mode: SearchMode, + memories: &[(Memory, f32)], + compact: bool, + include_scores: bool, +) -> ToolResult { + let legacy = if memories.is_empty() { + MSG_NO_MEMORIES.into() + } else { + format_memory_output(memories, compact) + }; + let output = MemoryRecallOutput::new(query, project, search_mode, memories, include_scores); + ToolResult::structured(legacy, format!("Found {} memories.", output.len()), &output) +} + +fn update_recall_access(store: &Store, memories: &mut [(Memory, f32)]) { + let refreshed = { + let ids: Vec<&str> = memories + .iter() + .map(|(memory, _)| memory.id.as_str()) + .collect(); + match store.batch_update_access(&ids) { + Ok(0) | Err(_) => return, + Ok(_) => store.get_many(&ids), + } + }; + if let Ok(mut refreshed) = refreshed { + for (memory, _) in memories { + if let Some(current) = refreshed.remove(&memory.id) { + *memory = current; + } + } + } +} + pub(in crate::tools) fn tool_recall(context: &ToolContext<'_>, args: &Value) -> ToolResult { let store = context.store; let embedder = context.embedder; @@ -285,12 +322,8 @@ pub(in crate::tools) fn tool_recall(context: &ToolContext<'_>, args: &Value) -> 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)), - } - }; + let memory_filter = + |memory: &Memory| matches_memory_filters(memory, project.as_deref(), topic, keyword); // Audit finding: filters were applied AFTER the store already truncated // to `limit` — if the top-`limit` global hits all belonged to other @@ -312,13 +345,7 @@ pub(in crate::tools) fn tool_recall(context: &ToolContext<'_>, args: &Value) -> 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)); - } + scored_results.retain(|(memory, _)| memory_filter(memory)); // Graph-aware expansion: follow `related_ids` one hop from // each primary hit and fold neighbors into the result set. @@ -334,35 +361,33 @@ pub(in crate::tools) fn tool_recall(context: &ToolContext<'_>, args: &Value) -> 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.retain(|(memory, _)| memory_filter(memory)); 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)); + update_recall_access(store, &mut expanded); + + return recall_result( + query, + project.as_deref(), + SearchMode::Hybrid, + &expanded, + compact, + true, + ); } } } // Fallback: FTS then keywords + let mut search_mode = SearchMode::FullText; 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() { + search_mode = SearchMode::Keyword; let keywords: Vec<&str> = query.split_whitespace().collect(); results = match store.search_by_keywords(&keywords, query_limit) { Ok(r) => r, @@ -370,13 +395,7 @@ pub(in crate::tools) fn tool_recall(context: &ToolContext<'_>, args: &Value) -> }; } - 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.retain(|memory| memory_filter(memory)); results.truncate(limit); // Convert to scored format with a sentinel score of 1.0 (FTS fallback @@ -391,26 +410,22 @@ pub(in crate::tools) fn tool_recall(context: &ToolContext<'_>, args: &Value) -> 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)); - } + expanded.retain(|(memory, _)| memory_filter(memory)); // 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()); - } + update_recall_access(store, &mut expanded); // 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)) + recall_result( + query, + project.as_deref(), + search_mode, + &for_display, + compact, + false, + ) } pub(in crate::tools) fn tool_forget(store: &Store, args: &Value) -> ToolResult { @@ -529,8 +544,13 @@ pub(in crate::tools) fn tool_consolidate( pub(in crate::tools) fn tool_list_topics(store: &Store) -> ToolResult { match store.list_topics() { Ok(topics) => { + let structured = MemoryTopicsOutput::new(&topics); if topics.is_empty() { - return ToolResult::text("No topics yet.".into()); + return ToolResult::structured( + "No topics yet.".into(), + "Found 0 topics.".into(), + &structured, + ); } // Group topics by scope prefix (before ':') @@ -565,7 +585,11 @@ pub(in crate::tools) fn tool_list_topics(store: &Store) -> ToolResult { } } - ToolResult::text(output) + ToolResult::structured( + output, + format!("Found {} topics.", structured.len()), + &structured, + ) } Err(e) => ToolResult::error(format!("failed to list topics: {e}")), } @@ -574,6 +598,7 @@ pub(in crate::tools) fn tool_list_topics(store: &Store) -> ToolResult { pub(in crate::tools) fn tool_stats(store: &Store) -> ToolResult { match store.stats() { Ok(stats) => { + let structured = MemoryStatsOutput::from(stats.clone()); let mut output = format!( "Memories: {}\nTopics: {}\nAvg weight: {:.3}\n", stats.total_memories, stats.total_topics, stats.avg_weight @@ -590,7 +615,7 @@ pub(in crate::tools) fn tool_stats(store: &Store) -> ToolResult { format_local(&newest, "%Y-%m-%d %H:%M") )); } - ToolResult::text(output) + ToolResult::structured(output, "Returned memory statistics.".into(), &structured) } Err(e) => ToolResult::error(format!("failed to get stats: {e}")), } diff --git a/crates/icm-mcp/src/tools/handlers/transcript.rs b/crates/icm-mcp/src/tools/handlers/transcript.rs index 903e0174..927a58e1 100644 --- a/crates/icm-mcp/src/tools/handlers/transcript.rs +++ b/crates/icm-mcp/src/tools/handlers/transcript.rs @@ -2,18 +2,25 @@ use serde_json::{json, Value}; -use icm_core::TranscriptStore; - use icm_store::Store; +use crate::outputs::{ + TranscriptRecordOutput, TranscriptSearchOutput, TranscriptShowOutput, TranscriptStartOutput, + TranscriptStatsOutput, +}; use crate::protocol::ToolResult; pub(in crate::tools) 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}\"}}")), + Ok(id) => { + let legacy = format!("{{\"session_id\":\"{id}\"}}"); + let output = TranscriptStartOutput::new(id); + ToolResult::structured(legacy, "Transcript session started.".into(), &output) + } Err(e) => ToolResult::error(format!("start_session failed: {e}")), } } @@ -44,12 +51,17 @@ pub(in crate::tools) fn tool_transcript_record(store: &Store, args: &Value) -> T 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}\"}}")), + Ok(id) => { + let legacy = format!("{{\"message_id\":\"{id}\"}}"); + let output = TranscriptRecordOutput::new(id); + ToolResult::structured(legacy, "Transcript message recorded.".into(), &output) + } Err(e) => ToolResult::error(format!("record failed: {e}")), } } pub(in crate::tools) 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()), @@ -64,13 +76,19 @@ pub(in crate::tools) fn tool_transcript_search(store: &Store, args: &Value) -> T 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) + let output = TranscriptSearchOutput::new(&hits); + ToolResult::structured( + json, + format!("Found {} transcript messages.", output.len()), + &output, + ) } Err(e) => ToolResult::error(format!("search failed: {e}")), } } pub(in crate::tools) 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()), @@ -89,13 +107,23 @@ pub(in crate::tools) fn tool_transcript_show(store: &Store, args: &Value) -> Too Ok(m) => m, Err(e) => return ToolResult::error(format!("list_messages failed: {e}")), }; + let output = TranscriptShowOutput::new(&sess, &msgs); let body = json!({ "session": sess, "messages": msgs }); - ToolResult::text(body.to_string()) + ToolResult::structured( + body.to_string(), + format!("Returned a transcript with {} messages.", output.len()), + &output, + ) } pub(in crate::tools) 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())), + Ok(stats) => { + let legacy = serde_json::to_string(&stats).unwrap_or_else(|_| "{}".into()); + let output = TranscriptStatsOutput::from(stats); + ToolResult::structured(legacy, "Returned transcript statistics.".into(), &output) + } Err(e) => ToolResult::error(format!("stats failed: {e}")), } } diff --git a/crates/icm-mcp/src/tools/registry.rs b/crates/icm-mcp/src/tools/registry.rs index 93c39c2d..68dce859 100644 --- a/crates/icm-mcp/src/tools/registry.rs +++ b/crates/icm-mcp/src/tools/registry.rs @@ -12,6 +12,7 @@ use crate::inputs::{ MemoryUpdateInput, NameInput, TopicInput, TranscriptRecordInput, TranscriptSearchInput, TranscriptShowInput, TranscriptStartInput, TranscriptStatsInput, WakeUpInput, }; +use crate::outputs::*; use super::handlers::{ tool_consolidate, tool_embed_all, tool_extract_patterns, tool_feedback_record, @@ -217,7 +218,8 @@ pub(crate) fn build_catalog(has_embedder: bool) -> ToolCatalog { requirements: ToolRequirements::STORE.with_optional_embedder(), ToolAnnotations::new(false, true, false, false), tool_recall - ), + ) + .with_output::(), tool_spec!( MemoryForgetInput, json!({ @@ -315,7 +317,8 @@ pub(crate) fn build_catalog(has_embedder: bool) -> ToolCatalog { }), ToolAnnotations::new(true, false, true, false), |context, _| tool_list_topics(context.store) - ), + ) + .with_output::(), tool_spec!( MemoryStatsInput, json!({ @@ -328,7 +331,8 @@ pub(crate) fn build_catalog(has_embedder: bool) -> ToolCatalog { }), ToolAnnotations::new(true, false, true, false), |context, _| tool_stats(context.store) - ), + ) + .with_output::(), tool_spec!( MemoryUpdateInput, json!({ @@ -710,7 +714,8 @@ pub(crate) fn build_catalog(has_embedder: bool) -> ToolCatalog { args, context.compact ) - ), + ) + .with_output::(), tool_spec!( FeedbackSearchInput, json!({ @@ -741,7 +746,8 @@ pub(crate) fn build_catalog(has_embedder: bool) -> ToolCatalog { requirements: ToolRequirements::STORE.with_optional_embedder(), ToolAnnotations::new(true, false, true, false), |context, args| tool_feedback_search(context.store, context.embedder, args) - ), + ) + .with_output::(), tool_spec!( FeedbackStatsInput, json!({ @@ -754,7 +760,8 @@ pub(crate) fn build_catalog(has_embedder: bool) -> ToolCatalog { }), ToolAnnotations::new(true, false, true, false), |context, _| tool_feedback_stats(context.store) - ), + ) + .with_output::(), // --- Transcript tools (verbatim session replay) --- tool_spec!( TranscriptStartInput, @@ -781,7 +788,8 @@ pub(crate) fn build_catalog(has_embedder: bool) -> ToolCatalog { }), ToolAnnotations::new(false, false, false, false), |context, args| tool_transcript_start_session(context.store, args) - ), + ) + .with_output::(), tool_spec!( TranscriptRecordInput, json!({ @@ -821,7 +829,8 @@ pub(crate) fn build_catalog(has_embedder: bool) -> ToolCatalog { }), ToolAnnotations::new(false, false, false, false), |context, args| tool_transcript_record(context.store, args) - ), + ) + .with_output::(), tool_spec!( TranscriptSearchInput, json!({ @@ -854,7 +863,8 @@ pub(crate) fn build_catalog(has_embedder: bool) -> ToolCatalog { }), ToolAnnotations::new(true, false, true, false), |context, args| tool_transcript_search(context.store, args) - ), + ) + .with_output::(), tool_spec!( TranscriptShowInput, json!({ @@ -871,7 +881,8 @@ pub(crate) fn build_catalog(has_embedder: bool) -> ToolCatalog { }), ToolAnnotations::new(true, false, true, false), |context, args| tool_transcript_show(context.store, args) - ), + ) + .with_output::(), tool_spec!( TranscriptStatsInput, json!({ @@ -884,7 +895,8 @@ pub(crate) fn build_catalog(has_embedder: bool) -> ToolCatalog { }), ToolAnnotations::new(true, false, true, false), |context, _| tool_transcript_stats(context.store) - ), + ) + .with_output::(), tool_spec!( WakeUpInput, json!({ diff --git a/crates/icm-mcp/src/tools/tests.rs b/crates/icm-mcp/src/tools/tests.rs index b20245b0..b41493c2 100644 --- a/crates/icm-mcp/src/tools/tests.rs +++ b/crates/icm-mcp/src/tools/tests.rs @@ -513,6 +513,7 @@ fn test_stats_empty() { let result = call_tool(&store, None, "icm_memory_stats", &json!({}), false); assert!(!result.is_error); assert!(result.content[0].text.contains("Memories: 0")); + assert!(result.structured_content.is_none()); } #[test]