From ef770f15b91136664222c77bd5611e3706befae8 Mon Sep 17 00:00:00 2001 From: Argenis Date: Fri, 13 Mar 2026 17:25:19 -0400 Subject: [PATCH] feat(tool): on-demand MCP tool loading via tool_search (#3446) Add deferred MCP tool activation to reduce context window waste. When mcp.deferred_loading is true (the default), MCP tool schemas are not eagerly included in the LLM context. Instead, only tool names appear in an system prompt section, and the LLM calls the built-in tool_search tool to fetch full schemas on demand. Setting deferred_loading to false preserves the existing eager behavior. Closes #3095 --- src/agent/loop_.rs | 72 ++++++-- src/channels/mod.rs | 68 +++++-- src/config/schema.rs | 36 +++- src/tools/mcp_deferred.rs | 366 ++++++++++++++++++++++++++++++++++++++ src/tools/mod.rs | 4 + src/tools/tool_search.rs | 284 +++++++++++++++++++++++++++++ 6 files changed, 791 insertions(+), 39 deletions(-) create mode 100644 src/tools/mcp_deferred.rs create mode 100644 src/tools/tool_search.rs 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")); + } +}