diff --git a/src/agent/loop_.rs b/src/agent/loop_.rs index 99ca53424..bf52a3592 100644 --- a/src/agent/loop_.rs +++ b/src/agent/loop_.rs @@ -2971,6 +2971,11 @@ pub async fn run( // MCP servers are user-declared external integrations; the built-in allow/deny // filter is not appropriate for them and would silently drop all MCP tools when // a restrictive allowlist is configured. Keep this block after any such filter call. + // + // When `deferred_loading` is enabled, MCP tools are NOT added to the registry + // eagerly. Instead, a `tool_search` built-in is registered so the LLM can + // fetch schemas on demand. This reduces context window waste. + let mut deferred_section = String::new(); if config.mcp.enabled && !config.mcp.servers.is_empty() { tracing::info!( "Initializing MCP client — {} server(s) configured", @@ -2979,28 +2984,51 @@ pub async fn run( match crate::tools::McpRegistry::connect_all(&config.mcp.servers).await { Ok(registry) => { let registry = std::sync::Arc::new(registry); - let names = registry.tool_names(); - let mut registered = 0usize; - for name in names { - if let Some(def) = registry.get_tool_def(&name).await { - let wrapper: std::sync::Arc = - std::sync::Arc::new(crate::tools::McpToolWrapper::new( - name, - def, - std::sync::Arc::clone(®istry), - )); - if let Some(ref handle) = delegate_handle { - handle.write().push(std::sync::Arc::clone(&wrapper)); + if config.mcp.deferred_loading { + // Deferred path: build stubs and register tool_search + let deferred_set = crate::tools::DeferredMcpToolSet::from_registry( + std::sync::Arc::clone(®istry), + ) + .await; + tracing::info!( + "MCP deferred: {} tool stub(s) from {} server(s)", + deferred_set.len(), + registry.server_count() + ); + deferred_section = + crate::tools::mcp_deferred::build_deferred_tools_section(&deferred_set); + let activated = std::sync::Arc::new(std::sync::Mutex::new( + crate::tools::ActivatedToolSet::new(), + )); + tools_registry.push(Box::new(crate::tools::ToolSearchTool::new( + deferred_set, + activated, + ))); + } else { + // Eager path: register all MCP tools directly + let names = registry.tool_names(); + let mut registered = 0usize; + for name in names { + if let Some(def) = registry.get_tool_def(&name).await { + let wrapper: std::sync::Arc = + std::sync::Arc::new(crate::tools::McpToolWrapper::new( + name, + def, + std::sync::Arc::clone(®istry), + )); + if let Some(ref handle) = delegate_handle { + handle.write().push(std::sync::Arc::clone(&wrapper)); + } + tools_registry.push(Box::new(crate::tools::ArcToolRef(wrapper))); + registered += 1; } - tools_registry.push(Box::new(crate::tools::ArcToolRef(wrapper))); - registered += 1; } + tracing::info!( + "MCP: {} tool(s) registered from {} server(s)", + registered, + registry.server_count() + ); } - tracing::info!( - "MCP: {} tool(s) registered from {} server(s)", - registered, - registry.server_count() - ); } Err(e) => { tracing::error!("MCP registry failed to initialize: {e:#}"); @@ -3196,6 +3224,12 @@ pub async fn run( system_prompt.push_str(&build_tool_instructions(&tools_registry)); } + // Append deferred MCP tool names so the LLM knows what is available + if !deferred_section.is_empty() { + system_prompt.push('\n'); + system_prompt.push_str(&deferred_section); + } + // ── Approval manager (supervised mode) ─────────────────────── let approval_manager = if interactive { Some(ApprovalManager::from_config(&config.autonomy)) diff --git a/src/channels/mod.rs b/src/channels/mod.rs index 63e7f7ffb..6c2e3344e 100644 --- a/src/channels/mod.rs +++ b/src/channels/mod.rs @@ -3441,6 +3441,9 @@ pub async fn start_channels(config: Config) -> Result<()> { ); // Wire MCP tools into the registry before freezing — non-fatal. + // When `deferred_loading` is enabled, MCP tools are NOT added eagerly. + // Instead, a `tool_search` built-in is registered for on-demand loading. + let mut deferred_section = String::new(); if config.mcp.enabled && !config.mcp.servers.is_empty() { tracing::info!( "Initializing MCP client — {} server(s) configured", @@ -3449,28 +3452,49 @@ pub async fn start_channels(config: Config) -> Result<()> { match crate::tools::McpRegistry::connect_all(&config.mcp.servers).await { Ok(registry) => { let registry = std::sync::Arc::new(registry); - let names = registry.tool_names(); - let mut registered = 0usize; - for name in names { - if let Some(def) = registry.get_tool_def(&name).await { - let wrapper: std::sync::Arc = - std::sync::Arc::new(crate::tools::McpToolWrapper::new( - name, - def, - std::sync::Arc::clone(®istry), - )); - if let Some(ref handle) = delegate_handle_ch { - handle.write().push(std::sync::Arc::clone(&wrapper)); + if config.mcp.deferred_loading { + let deferred_set = crate::tools::DeferredMcpToolSet::from_registry( + std::sync::Arc::clone(®istry), + ) + .await; + tracing::info!( + "MCP deferred: {} tool stub(s) from {} server(s)", + deferred_set.len(), + registry.server_count() + ); + deferred_section = + crate::tools::mcp_deferred::build_deferred_tools_section(&deferred_set); + let activated = std::sync::Arc::new(std::sync::Mutex::new( + crate::tools::ActivatedToolSet::new(), + )); + built_tools.push(Box::new(crate::tools::ToolSearchTool::new( + deferred_set, + activated, + ))); + } else { + let names = registry.tool_names(); + let mut registered = 0usize; + for name in names { + if let Some(def) = registry.get_tool_def(&name).await { + let wrapper: std::sync::Arc = + std::sync::Arc::new(crate::tools::McpToolWrapper::new( + name, + def, + std::sync::Arc::clone(®istry), + )); + if let Some(ref handle) = delegate_handle_ch { + handle.write().push(std::sync::Arc::clone(&wrapper)); + } + built_tools.push(Box::new(crate::tools::ArcToolRef(wrapper))); + registered += 1; } - built_tools.push(Box::new(crate::tools::ArcToolRef(wrapper))); - registered += 1; } + tracing::info!( + "MCP: {} tool(s) registered from {} server(s)", + registered, + registry.server_count() + ); } - tracing::info!( - "MCP: {} tool(s) registered from {} server(s)", - registered, - registry.server_count() - ); } Err(e) => { // Non-fatal — daemon continues with the tools registered above. @@ -3565,6 +3589,12 @@ pub async fn start_channels(config: Config) -> Result<()> { system_prompt.push_str(&build_tool_instructions(tools_registry.as_ref())); } + // Append deferred MCP tool names so the LLM knows what is available + if !deferred_section.is_empty() { + system_prompt.push('\n'); + system_prompt.push_str(&deferred_section); + } + if !skills.is_empty() { println!( " 🧩 Skills: {}", diff --git a/src/config/schema.rs b/src/config/schema.rs index 54057c1aa..b3107c35d 100644 --- a/src/config/schema.rs +++ b/src/config/schema.rs @@ -514,16 +514,36 @@ pub struct McpServerConfig { } /// External MCP client configuration (`[mcp]` section). -#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, Default)] +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] pub struct McpConfig { /// Enable MCP tool loading. #[serde(default)] pub enabled: bool, + /// Load MCP tool schemas on-demand via `tool_search` instead of eagerly + /// including them in the LLM context window. When `true` (the default), + /// only tool names are listed in the system prompt; the LLM must call + /// `tool_search` to fetch full schemas before invoking a deferred tool. + #[serde(default = "default_deferred_loading")] + pub deferred_loading: bool, /// Configured MCP servers. #[serde(default, alias = "mcpServers")] pub servers: Vec, } +fn default_deferred_loading() -> bool { + true +} + +impl Default for McpConfig { + fn default() -> Self { + Self { + enabled: false, + deferred_loading: default_deferred_loading(), + servers: Vec::new(), + } + } +} + // ── TTS (Text-to-Speech) ───────────────────────────────────────── fn default_tts_provider() -> String { @@ -8843,6 +8863,7 @@ require_otp_to_resume = true let cfg = McpConfig { enabled: true, servers: vec![stdio_server("fs", "/usr/bin/mcp-fs")], + ..Default::default() }; assert!(validate_mcp_config(&cfg).is_ok()); } @@ -8852,6 +8873,7 @@ require_otp_to_resume = true let cfg = McpConfig { enabled: true, servers: vec![http_server("svc", "http://localhost:8080/mcp")], + ..Default::default() }; assert!(validate_mcp_config(&cfg).is_ok()); } @@ -8861,6 +8883,7 @@ require_otp_to_resume = true let cfg = McpConfig { enabled: true, servers: vec![sse_server("svc", "https://example.com/events")], + ..Default::default() }; assert!(validate_mcp_config(&cfg).is_ok()); } @@ -8870,6 +8893,7 @@ require_otp_to_resume = true let cfg = McpConfig { enabled: true, servers: vec![stdio_server("", "/usr/bin/tool")], + ..Default::default() }; let err = validate_mcp_config(&cfg).expect_err("empty name should fail"); assert!( @@ -8883,6 +8907,7 @@ require_otp_to_resume = true let cfg = McpConfig { enabled: true, servers: vec![stdio_server(" ", "/usr/bin/tool")], + ..Default::default() }; let err = validate_mcp_config(&cfg).expect_err("whitespace name should fail"); assert!( @@ -8899,6 +8924,7 @@ require_otp_to_resume = true stdio_server("fs", "/usr/bin/mcp-a"), stdio_server("fs", "/usr/bin/mcp-b"), ], + ..Default::default() }; let err = validate_mcp_config(&cfg).expect_err("duplicate name should fail"); assert!(err.to_string().contains("duplicate name"), "got: {err}"); @@ -8911,6 +8937,7 @@ require_otp_to_resume = true let cfg = McpConfig { enabled: true, servers: vec![server], + ..Default::default() }; let err = validate_mcp_config(&cfg).expect_err("zero timeout should fail"); assert!(err.to_string().contains("greater than 0"), "got: {err}"); @@ -8923,6 +8950,7 @@ require_otp_to_resume = true let cfg = McpConfig { enabled: true, servers: vec![server], + ..Default::default() }; let err = validate_mcp_config(&cfg).expect_err("oversized timeout should fail"); assert!(err.to_string().contains("exceeds max"), "got: {err}"); @@ -8935,6 +8963,7 @@ require_otp_to_resume = true let cfg = McpConfig { enabled: true, servers: vec![server], + ..Default::default() }; assert!(validate_mcp_config(&cfg).is_ok()); } @@ -8944,6 +8973,7 @@ require_otp_to_resume = true let cfg = McpConfig { enabled: true, servers: vec![stdio_server("fs", "")], + ..Default::default() }; let err = validate_mcp_config(&cfg).expect_err("empty command should fail"); assert!( @@ -8962,6 +8992,7 @@ require_otp_to_resume = true url: None, ..Default::default() }], + ..Default::default() }; let err = validate_mcp_config(&cfg).expect_err("http without url should fail"); assert!(err.to_string().contains("requires url"), "got: {err}"); @@ -8977,6 +9008,7 @@ require_otp_to_resume = true url: None, ..Default::default() }], + ..Default::default() }; let err = validate_mcp_config(&cfg).expect_err("sse without url should fail"); assert!(err.to_string().contains("requires url"), "got: {err}"); @@ -8987,6 +9019,7 @@ require_otp_to_resume = true let cfg = McpConfig { enabled: true, servers: vec![http_server("svc", "ftp://example.com/mcp")], + ..Default::default() }; let err = validate_mcp_config(&cfg).expect_err("non-http scheme should fail"); assert!(err.to_string().contains("http/https"), "got: {err}"); @@ -8997,6 +9030,7 @@ require_otp_to_resume = true let cfg = McpConfig { enabled: true, servers: vec![http_server("svc", "not a url at all !!!")], + ..Default::default() }; let err = validate_mcp_config(&cfg).expect_err("invalid url should fail"); assert!(err.to_string().contains("valid URL"), "got: {err}"); diff --git a/src/tools/mcp_deferred.rs b/src/tools/mcp_deferred.rs new file mode 100644 index 000000000..0e97b0cd6 --- /dev/null +++ b/src/tools/mcp_deferred.rs @@ -0,0 +1,366 @@ +//! Deferred MCP tool loading — stubs and activated-tool tracking. +//! +//! When `mcp.deferred_loading` is enabled, MCP tool schemas are NOT eagerly +//! included in the LLM context window. Instead, only lightweight stubs (name + +//! description) are exposed in the system prompt. The LLM must call the built-in +//! `tool_search` tool to fetch full schemas, which moves them into the +//! [`ActivatedToolSet`] for the current conversation. + +use std::collections::HashMap; +use std::sync::Arc; + +use crate::tools::mcp_client::McpRegistry; +use crate::tools::mcp_protocol::McpToolDef; +use crate::tools::mcp_tool::McpToolWrapper; +use crate::tools::traits::{Tool, ToolSpec}; + +// ── DeferredMcpToolStub ────────────────────────────────────────────────── + +/// A lightweight stub representing a known-but-not-yet-loaded MCP tool. +/// Contains only the prefixed name, a human-readable description, and enough +/// information to construct the full [`McpToolWrapper`] on activation. +#[derive(Debug, Clone)] +pub struct DeferredMcpToolStub { + /// Prefixed name: `__`. + pub prefixed_name: String, + /// Human-readable description (extracted from the MCP tool definition). + pub description: String, + /// The full tool definition — stored so we can construct a wrapper later. + def: McpToolDef, +} + +impl DeferredMcpToolStub { + pub fn new(prefixed_name: String, def: McpToolDef) -> Self { + let description = def + .description + .clone() + .unwrap_or_else(|| "MCP tool".to_string()); + Self { + prefixed_name, + description, + def, + } + } + + /// Materialize this stub into a live [`McpToolWrapper`]. + pub fn activate(&self, registry: Arc) -> McpToolWrapper { + McpToolWrapper::new(self.prefixed_name.clone(), self.def.clone(), registry) + } +} + +// ── DeferredMcpToolSet ─────────────────────────────────────────────────── + +/// Collection of all deferred MCP tool stubs discovered at startup. +/// Provides keyword search for `tool_search`. +#[derive(Clone)] +pub struct DeferredMcpToolSet { + /// All stubs — exposed for test construction. + pub stubs: Vec, + /// Shared registry — exposed for test construction. + pub registry: Arc, +} + +impl DeferredMcpToolSet { + /// Build the set from a connected [`McpRegistry`]. + pub async fn from_registry(registry: Arc) -> Self { + let names = registry.tool_names(); + let mut stubs = Vec::with_capacity(names.len()); + for name in names { + if let Some(def) = registry.get_tool_def(&name).await { + stubs.push(DeferredMcpToolStub::new(name, def)); + } + } + Self { stubs, registry } + } + + /// All stub names (for rendering in the system prompt). + pub fn stub_names(&self) -> Vec<&str> { + self.stubs + .iter() + .map(|s| s.prefixed_name.as_str()) + .collect() + } + + /// Number of deferred stubs. + pub fn len(&self) -> usize { + self.stubs.len() + } + + /// Whether the set is empty. + pub fn is_empty(&self) -> bool { + self.stubs.is_empty() + } + + /// Look up stubs by exact name. Used for `select:name1,name2` queries. + pub fn get_by_name(&self, name: &str) -> Option<&DeferredMcpToolStub> { + self.stubs.iter().find(|s| s.prefixed_name == name) + } + + /// Keyword search — returns stubs whose name or description contains any + /// of the query terms (case-insensitive). Results are ranked by number of + /// matching terms (descending). + pub fn search(&self, query: &str, max_results: usize) -> Vec<&DeferredMcpToolStub> { + let terms: Vec = query + .split_whitespace() + .map(|t| t.to_ascii_lowercase()) + .collect(); + if terms.is_empty() { + return self.stubs.iter().take(max_results).collect(); + } + + let mut scored: Vec<(&DeferredMcpToolStub, usize)> = self + .stubs + .iter() + .filter_map(|stub| { + let haystack = format!( + "{} {}", + stub.prefixed_name.to_ascii_lowercase(), + stub.description.to_ascii_lowercase() + ); + let hits = terms + .iter() + .filter(|t| haystack.contains(t.as_str())) + .count(); + if hits > 0 { + Some((stub, hits)) + } else { + None + } + }) + .collect(); + + scored.sort_by(|a, b| b.1.cmp(&a.1)); + scored + .into_iter() + .take(max_results) + .map(|(s, _)| s) + .collect() + } + + /// Activate a stub by name, returning a boxed [`Tool`]. + pub fn activate(&self, name: &str) -> Option> { + self.get_by_name(name).map(|stub| { + let wrapper = stub.activate(Arc::clone(&self.registry)); + Box::new(wrapper) as Box + }) + } + + /// Return the full [`ToolSpec`] for a stub (for inclusion in `tool_search` results). + pub fn tool_spec(&self, name: &str) -> Option { + self.get_by_name(name).map(|stub| { + let wrapper = stub.activate(Arc::clone(&self.registry)); + wrapper.spec() + }) + } +} + +// ── ActivatedToolSet ───────────────────────────────────────────────────── + +/// Per-conversation mutable state tracking which deferred tools have been +/// activated (i.e. their full schemas have been fetched via `tool_search`). +/// The agent loop consults this each iteration to decide which tool_specs +/// to include in the LLM request. +pub struct ActivatedToolSet { + /// name -> activated Tool + tools: HashMap>, +} + +impl ActivatedToolSet { + pub fn new() -> Self { + Self { + tools: HashMap::new(), + } + } + + /// Mark a tool as activated, storing its live wrapper. + pub fn activate(&mut self, name: String, tool: Box) { + self.tools.insert(name, tool); + } + + /// Whether a tool has been activated. + pub fn is_activated(&self, name: &str) -> bool { + self.tools.contains_key(name) + } + + /// Get an activated tool for execution. + pub fn get(&self, name: &str) -> Option<&dyn Tool> { + self.tools.get(name).map(|t| t.as_ref()) + } + + /// All currently activated tool specs (to include in LLM requests). + pub fn tool_specs(&self) -> Vec { + self.tools.values().map(|t| t.spec()).collect() + } + + /// All activated tools for execution dispatch. + pub fn tool_names(&self) -> Vec<&str> { + self.tools.keys().map(|s| s.as_str()).collect() + } +} + +impl Default for ActivatedToolSet { + fn default() -> Self { + Self::new() + } +} + +// ── System prompt helper ───────────────────────────────────────────────── + +/// Build the `` section for the system prompt. +/// Lists only tool names so the LLM knows what is available without +/// consuming context window on full schemas. +pub fn build_deferred_tools_section(deferred: &DeferredMcpToolSet) -> String { + if deferred.is_empty() { + return String::new(); + } + let mut out = String::from("\n"); + for stub in &deferred.stubs { + out.push_str(&stub.prefixed_name); + out.push('\n'); + } + out.push_str("\n"); + out +} + +#[cfg(test)] +mod tests { + use super::*; + + fn make_stub(name: &str, desc: &str) -> DeferredMcpToolStub { + let def = McpToolDef { + name: name.to_string(), + description: Some(desc.to_string()), + input_schema: serde_json::json!({"type": "object", "properties": {}}), + }; + DeferredMcpToolStub::new(name.to_string(), def) + } + + #[test] + fn stub_uses_description_from_def() { + let stub = make_stub("fs__read", "Read a file"); + assert_eq!(stub.description, "Read a file"); + } + + #[test] + fn stub_defaults_description_when_none() { + let def = McpToolDef { + name: "mystery".into(), + description: None, + input_schema: serde_json::json!({}), + }; + let stub = DeferredMcpToolStub::new("srv__mystery".into(), def); + assert_eq!(stub.description, "MCP tool"); + } + + #[test] + fn activated_set_tracks_activation() { + use crate::tools::traits::ToolResult; + use async_trait::async_trait; + + struct FakeTool; + #[async_trait] + impl Tool for FakeTool { + fn name(&self) -> &str { + "fake" + } + fn description(&self) -> &str { + "fake tool" + } + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({}) + } + async fn execute(&self, _: serde_json::Value) -> anyhow::Result { + Ok(ToolResult { + success: true, + output: String::new(), + error: None, + }) + } + } + + let mut set = ActivatedToolSet::new(); + assert!(!set.is_activated("fake")); + set.activate("fake".into(), Box::new(FakeTool)); + assert!(set.is_activated("fake")); + assert!(set.get("fake").is_some()); + assert_eq!(set.tool_specs().len(), 1); + } + + #[test] + fn build_deferred_section_empty_when_no_stubs() { + let set = DeferredMcpToolSet { + stubs: vec![], + registry: std::sync::Arc::new( + tokio::runtime::Runtime::new() + .unwrap() + .block_on(McpRegistry::connect_all(&[])) + .unwrap(), + ), + }; + assert!(build_deferred_tools_section(&set).is_empty()); + } + + #[test] + fn build_deferred_section_lists_names() { + let stubs = vec![ + make_stub("fs__read_file", "Read a file"), + make_stub("git__status", "Git status"), + ]; + let set = DeferredMcpToolSet { + stubs, + registry: std::sync::Arc::new( + tokio::runtime::Runtime::new() + .unwrap() + .block_on(McpRegistry::connect_all(&[])) + .unwrap(), + ), + }; + let section = build_deferred_tools_section(&set); + assert!(section.contains("")); + assert!(section.contains("fs__read_file")); + assert!(section.contains("git__status")); + assert!(section.contains("")); + } + + #[test] + fn keyword_search_ranks_by_hits() { + let stubs = vec![ + make_stub("fs__read_file", "Read a file from disk"), + make_stub("fs__write_file", "Write a file to disk"), + make_stub("git__log", "Show git log"), + ]; + let set = DeferredMcpToolSet { + stubs, + registry: std::sync::Arc::new( + tokio::runtime::Runtime::new() + .unwrap() + .block_on(McpRegistry::connect_all(&[])) + .unwrap(), + ), + }; + + // "file read" should rank fs__read_file highest (2 hits vs 1) + let results = set.search("file read", 5); + assert!(!results.is_empty()); + assert_eq!(results[0].prefixed_name, "fs__read_file"); + } + + #[test] + fn get_by_name_returns_correct_stub() { + let stubs = vec![ + make_stub("a__one", "Tool one"), + make_stub("b__two", "Tool two"), + ]; + let set = DeferredMcpToolSet { + stubs, + registry: std::sync::Arc::new( + tokio::runtime::Runtime::new() + .unwrap() + .block_on(McpRegistry::connect_all(&[])) + .unwrap(), + ), + }; + assert!(set.get_by_name("a__one").is_some()); + assert!(set.get_by_name("nonexistent").is_none()); + } +} diff --git a/src/tools/mod.rs b/src/tools/mod.rs index dc7b66b09..a8bc68587 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -41,6 +41,7 @@ pub mod hardware_memory_read; pub mod http_request; pub mod image_info; pub mod mcp_client; +pub mod mcp_deferred; pub mod mcp_protocol; pub mod mcp_tool; pub mod mcp_transport; @@ -55,6 +56,7 @@ pub mod schedule; pub mod schema; pub mod screenshot; pub mod shell; +pub mod tool_search; pub mod traits; pub mod web_fetch; pub mod web_search_tool; @@ -84,6 +86,7 @@ pub use hardware_memory_read::HardwareMemoryReadTool; pub use http_request::HttpRequestTool; pub use image_info::ImageInfoTool; pub use mcp_client::McpRegistry; +pub use mcp_deferred::{ActivatedToolSet, DeferredMcpToolSet}; pub use mcp_tool::McpToolWrapper; pub use memory_forget::MemoryForgetTool; pub use memory_recall::MemoryRecallTool; @@ -97,6 +100,7 @@ pub use schedule::ScheduleTool; pub use schema::{CleaningStrategy, SchemaCleanr}; pub use screenshot::ScreenshotTool; pub use shell::ShellTool; +pub use tool_search::ToolSearchTool; pub use traits::Tool; #[allow(unused_imports)] pub use traits::{ToolResult, ToolSpec}; diff --git a/src/tools/tool_search.rs b/src/tools/tool_search.rs new file mode 100644 index 000000000..4bf6163dc --- /dev/null +++ b/src/tools/tool_search.rs @@ -0,0 +1,284 @@ +//! Built-in `tool_search` tool for on-demand MCP tool schema loading. +//! +//! When `mcp.deferred_loading` is enabled, this tool lets the LLM discover and +//! activate deferred MCP tools. Supports two query modes: +//! - `select:name1,name2` — fetch exact tools by prefixed name. +//! - Free-text keyword search — returns the best-matching stubs. + +use std::fmt::Write; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; + +use crate::tools::mcp_deferred::{ActivatedToolSet, DeferredMcpToolSet}; +use crate::tools::traits::{Tool, ToolResult}; + +/// Default maximum number of search results. +const DEFAULT_MAX_RESULTS: usize = 5; + +/// Built-in tool that fetches full schemas for deferred MCP tools. +pub struct ToolSearchTool { + deferred: DeferredMcpToolSet, + activated: Arc>, +} + +impl ToolSearchTool { + pub fn new(deferred: DeferredMcpToolSet, activated: Arc>) -> Self { + Self { + deferred, + activated, + } + } +} + +#[async_trait] +impl Tool for ToolSearchTool { + fn name(&self) -> &str { + "tool_search" + } + + fn description(&self) -> &str { + "Fetch full schema definitions for deferred MCP tools so they can be called. \ + Use \"select:name1,name2\" for exact match or keywords to search." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "query": { + "description": "Query to find deferred tools. Use \"select:\" for direct selection, or keywords to search.", + "type": "string" + }, + "max_results": { + "description": "Maximum number of results to return (default: 5)", + "type": "number", + "default": DEFAULT_MAX_RESULTS + } + }, + "required": ["query"] + }) + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + let query = args + .get("query") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .trim(); + + let max_results = args + .get("max_results") + .and_then(|v| v.as_u64()) + .map(|v| usize::try_from(v).unwrap_or(DEFAULT_MAX_RESULTS)) + .unwrap_or(DEFAULT_MAX_RESULTS); + + if query.is_empty() { + return Ok(ToolResult { + success: false, + output: String::new(), + error: Some("query parameter is required".into()), + }); + } + + // Parse query mode + if let Some(names_str) = query.strip_prefix("select:") { + // Exact selection mode + let names: Vec<&str> = names_str.split(',').map(str::trim).collect(); + return self.select_tools(&names); + } + + // Keyword search mode + let results = self.deferred.search(query, max_results); + if results.is_empty() { + return Ok(ToolResult { + success: true, + output: "No matching deferred tools found.".into(), + error: None, + }); + } + + // Activate and return full specs + let mut output = String::from("\n"); + let mut activated_count = 0; + let mut guard = self.activated.lock().unwrap(); + + for stub in &results { + if let Some(spec) = self.deferred.tool_spec(&stub.prefixed_name) { + if !guard.is_activated(&stub.prefixed_name) { + if let Some(tool) = self.deferred.activate(&stub.prefixed_name) { + guard.activate(stub.prefixed_name.clone(), tool); + activated_count += 1; + } + } + let _ = writeln!( + output, + "{{\"name\": \"{}\", \"description\": \"{}\", \"parameters\": {}}}", + spec.name, + spec.description.replace('"', "\\\""), + spec.parameters + ); + } + } + + output.push_str("\n"); + drop(guard); + + tracing::debug!( + "tool_search: query={query:?}, matched={}, activated={activated_count}", + results.len() + ); + + Ok(ToolResult { + success: true, + output, + error: None, + }) + } +} + +impl ToolSearchTool { + fn select_tools(&self, names: &[&str]) -> anyhow::Result { + let mut output = String::from("\n"); + let mut not_found = Vec::new(); + let mut activated_count = 0; + let mut guard = self.activated.lock().unwrap(); + + for name in names { + if name.is_empty() { + continue; + } + match self.deferred.tool_spec(name) { + Some(spec) => { + if !guard.is_activated(name) { + if let Some(tool) = self.deferred.activate(name) { + guard.activate(name.to_string(), tool); + activated_count += 1; + } + } + let _ = writeln!( + output, + "{{\"name\": \"{}\", \"description\": \"{}\", \"parameters\": {}}}", + spec.name, + spec.description.replace('"', "\\\""), + spec.parameters + ); + } + None => { + not_found.push(*name); + } + } + } + + output.push_str("\n"); + drop(guard); + + if !not_found.is_empty() { + let _ = write!(output, "\nNot found: {}", not_found.join(", ")); + } + + tracing::debug!( + "tool_search select: requested={}, activated={activated_count}, not_found={}", + names.len(), + not_found.len() + ); + + Ok(ToolResult { + success: true, + output, + error: None, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::tools::mcp_client::McpRegistry; + use crate::tools::mcp_deferred::DeferredMcpToolStub; + use crate::tools::mcp_protocol::McpToolDef; + + async fn make_deferred_set(stubs: Vec) -> DeferredMcpToolSet { + let registry = Arc::new(McpRegistry::connect_all(&[]).await.unwrap()); + DeferredMcpToolSet { stubs, registry } + } + + fn make_stub(name: &str, desc: &str) -> DeferredMcpToolStub { + let def = McpToolDef { + name: name.to_string(), + description: Some(desc.to_string()), + input_schema: serde_json::json!({"type": "object", "properties": {}}), + }; + DeferredMcpToolStub::new(name.to_string(), def) + } + + #[tokio::test] + async fn tool_metadata() { + let tool = ToolSearchTool::new( + make_deferred_set(vec![]).await, + Arc::new(Mutex::new(ActivatedToolSet::new())), + ); + assert_eq!(tool.name(), "tool_search"); + assert!(!tool.description().is_empty()); + assert!(tool.parameters_schema()["properties"]["query"].is_object()); + } + + #[tokio::test] + async fn empty_query_returns_error() { + let tool = ToolSearchTool::new( + make_deferred_set(vec![]).await, + Arc::new(Mutex::new(ActivatedToolSet::new())), + ); + let result = tool + .execute(serde_json::json!({"query": ""})) + .await + .unwrap(); + assert!(!result.success); + } + + #[tokio::test] + async fn select_nonexistent_tool_reports_not_found() { + let tool = ToolSearchTool::new( + make_deferred_set(vec![]).await, + Arc::new(Mutex::new(ActivatedToolSet::new())), + ); + let result = tool + .execute(serde_json::json!({"query": "select:nonexistent"})) + .await + .unwrap(); + assert!(result.success); + assert!(result.output.contains("Not found")); + } + + #[tokio::test] + async fn keyword_search_no_matches() { + let tool = ToolSearchTool::new( + make_deferred_set(vec![make_stub("fs__read", "Read file")]).await, + Arc::new(Mutex::new(ActivatedToolSet::new())), + ); + let result = tool + .execute(serde_json::json!({"query": "zzzzz_nonexistent"})) + .await + .unwrap(); + assert!(result.success); + assert!(result.output.contains("No matching")); + } + + #[tokio::test] + async fn keyword_search_finds_match() { + let activated = Arc::new(Mutex::new(ActivatedToolSet::new())); + let tool = ToolSearchTool::new( + make_deferred_set(vec![make_stub("fs__read", "Read a file from disk")]).await, + Arc::clone(&activated), + ); + let result = tool + .execute(serde_json::json!({"query": "read file"})) + .await + .unwrap(); + assert!(result.success); + assert!(result.output.contains("")); + assert!(result.output.contains("fs__read")); + // Tool should now be activated + assert!(activated.lock().unwrap().is_activated("fs__read")); + } +}