zeroclaw_tools/
mcp_deferred.rs

1//! Deferred MCP tool loading — stubs and activated-tool tracking.
2//!
3//! When `mcp.deferred_loading` is enabled, MCP tool schemas are NOT eagerly
4//! included in the LLM context window. Instead, only lightweight stubs (name +
5//! description) are exposed in the system prompt. The LLM must call the built-in
6//! `tool_search` tool to fetch full schemas, which moves them into the
7//! [`ActivatedToolSet`] for the current conversation.
8
9use std::collections::HashMap;
10use std::sync::Arc;
11
12use crate::mcp_client::McpRegistry;
13use crate::mcp_protocol::McpToolDef;
14use crate::mcp_tool::McpToolWrapper;
15use zeroclaw_api::tool::{Tool, ToolSpec};
16
17// ── DeferredMcpToolStub ──────────────────────────────────────────────────
18
19/// A lightweight stub representing a known-but-not-yet-loaded MCP tool.
20/// Contains only the prefixed name, a human-readable description, and enough
21/// information to construct the full [`McpToolWrapper`] on activation.
22#[derive(Debug, Clone)]
23pub struct DeferredMcpToolStub {
24    /// Prefixed name: `<server_name>__<tool_name>`.
25    pub prefixed_name: String,
26    /// Human-readable description (extracted from the MCP tool definition).
27    pub description: String,
28    /// The full tool definition — stored so we can construct a wrapper later.
29    def: McpToolDef,
30}
31
32impl DeferredMcpToolStub {
33    pub fn new(prefixed_name: String, def: McpToolDef) -> Self {
34        let description = def
35            .description
36            .clone()
37            .unwrap_or_else(|| "MCP tool".to_string());
38        Self {
39            prefixed_name,
40            description,
41            def,
42        }
43    }
44
45    /// Materialize this stub into a live [`McpToolWrapper`].
46    pub fn activate(&self, registry: Arc<McpRegistry>) -> McpToolWrapper {
47        McpToolWrapper::new(self.prefixed_name.clone(), self.def.clone(), registry)
48    }
49}
50
51// ── DeferredMcpToolSet ───────────────────────────────────────────────────
52
53/// Collection of all deferred MCP tool stubs discovered at startup.
54/// Provides keyword search for `tool_search`.
55#[derive(Clone)]
56pub struct DeferredMcpToolSet {
57    /// All stubs — exposed for test construction.
58    pub stubs: Vec<DeferredMcpToolStub>,
59    /// Shared registry — exposed for test construction.
60    pub registry: Arc<McpRegistry>,
61}
62
63impl DeferredMcpToolSet {
64    /// Build the set from a connected [`McpRegistry`].
65    pub async fn from_registry(registry: Arc<McpRegistry>) -> Self {
66        let names = registry.tool_names();
67        let mut stubs = Vec::with_capacity(names.len());
68        for name in names {
69            if let Some(def) = registry.get_tool_def(&name).await {
70                stubs.push(DeferredMcpToolStub::new(name, def));
71            }
72        }
73        Self { stubs, registry }
74    }
75
76    /// All stub names (for rendering in the system prompt).
77    pub fn stub_names(&self) -> Vec<&str> {
78        self.stubs
79            .iter()
80            .map(|s| s.prefixed_name.as_str())
81            .collect()
82    }
83
84    /// Number of deferred stubs.
85    pub fn len(&self) -> usize {
86        self.stubs.len()
87    }
88
89    /// Whether the set is empty.
90    pub fn is_empty(&self) -> bool {
91        self.stubs.is_empty()
92    }
93
94    /// Look up stubs by exact name. Used for `select:name1,name2` queries.
95    pub fn get_by_name(&self, name: &str) -> Option<&DeferredMcpToolStub> {
96        self.stubs.iter().find(|s| s.prefixed_name == name)
97    }
98
99    /// Keyword search — returns stubs whose name or description contains any
100    /// of the query terms (case-insensitive). Results are ranked by number of
101    /// matching terms (descending).
102    pub fn search(&self, query: &str, max_results: usize) -> Vec<&DeferredMcpToolStub> {
103        let terms: Vec<String> = query
104            .split_whitespace()
105            .map(|t| t.to_ascii_lowercase())
106            .collect();
107        if terms.is_empty() {
108            return self.stubs.iter().take(max_results).collect();
109        }
110
111        let mut scored: Vec<(&DeferredMcpToolStub, usize)> = self
112            .stubs
113            .iter()
114            .filter_map(|stub| {
115                let haystack = format!(
116                    "{} {}",
117                    stub.prefixed_name.to_ascii_lowercase(),
118                    stub.description.to_ascii_lowercase()
119                );
120                let hits = terms
121                    .iter()
122                    .filter(|t| haystack.contains(t.as_str()))
123                    .count();
124                if hits > 0 { Some((stub, hits)) } else { None }
125            })
126            .collect();
127
128        scored.sort_by_key(|entry| std::cmp::Reverse(entry.1));
129        scored
130            .into_iter()
131            .take(max_results)
132            .map(|(s, _)| s)
133            .collect()
134    }
135
136    /// Activate a stub by name, returning a boxed [`Tool`].
137    pub fn activate(&self, name: &str) -> Option<Box<dyn Tool>> {
138        self.get_by_name(name).map(|stub| {
139            let wrapper = stub.activate(Arc::clone(&self.registry));
140            Box::new(wrapper) as Box<dyn Tool>
141        })
142    }
143
144    /// Return the full [`ToolSpec`] for a stub (for inclusion in `tool_search` results).
145    pub fn tool_spec(&self, name: &str) -> Option<ToolSpec> {
146        self.get_by_name(name).map(|stub| {
147            let wrapper = stub.activate(Arc::clone(&self.registry));
148            wrapper.spec()
149        })
150    }
151}
152
153// ── ActivatedToolSet ─────────────────────────────────────────────────────
154
155/// Per-conversation mutable state tracking which deferred tools have been
156/// activated (i.e. their full schemas have been fetched via `tool_search`).
157/// The agent loop consults this each iteration to decide which tool_specs
158/// to include in the LLM request.
159pub struct ActivatedToolSet {
160    tools: HashMap<String, Arc<dyn Tool>>,
161}
162
163impl ActivatedToolSet {
164    pub fn new() -> Self {
165        Self {
166            tools: HashMap::new(),
167        }
168    }
169
170    pub fn activate(&mut self, name: String, tool: Arc<dyn Tool>) {
171        self.tools.insert(name, tool);
172    }
173
174    pub fn is_activated(&self, name: &str) -> bool {
175        self.tools.contains_key(name)
176    }
177
178    /// Clone the Arc so the caller can drop the mutex guard before awaiting.
179    pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
180        self.tools.get(name).cloned()
181    }
182
183    /// Resolve an activated tool by exact name first, then by unique MCP suffix.
184    ///
185    /// Some model_providers occasionally strip the `<server>__` prefix when calling a
186    /// deferred MCP tool after `tool_search` activation. When the suffix maps to
187    /// exactly one activated tool, allow that call to proceed.
188    pub fn get_resolved(&self, name: &str) -> Option<Arc<dyn Tool>> {
189        if let Some(tool) = self.get(name) {
190            return Some(tool);
191        }
192        if name.contains("__") {
193            return None;
194        }
195
196        let mut resolved = None;
197        for (tool_name, tool) in &self.tools {
198            let Some((_, suffix)) = tool_name.split_once("__") else {
199                continue;
200            };
201            if suffix != name {
202                continue;
203            }
204            if resolved.is_some() {
205                return None;
206            }
207            resolved = Some(Arc::clone(tool));
208        }
209
210        resolved
211    }
212
213    pub fn tool_specs(&self) -> Vec<ToolSpec> {
214        self.tools.values().map(|t| t.spec()).collect()
215    }
216
217    pub fn tool_names(&self) -> Vec<&str> {
218        self.tools.keys().map(|s| s.as_str()).collect()
219    }
220}
221
222impl Default for ActivatedToolSet {
223    fn default() -> Self {
224        Self::new()
225    }
226}
227
228// ── System prompt helper ─────────────────────────────────────────────────
229
230/// Build the `<available-deferred-tools>` section for the system prompt.
231/// Lists only tool names so the LLM knows what is available without
232/// consuming context window on full schemas. Includes an instruction
233/// block that tells the LLM to call `tool_search` to activate them.
234pub fn build_deferred_tools_section(deferred: &DeferredMcpToolSet) -> String {
235    build_deferred_tools_section_filtered(deferred, None)
236}
237
238pub fn build_deferred_tools_section_filtered(
239    deferred: &DeferredMcpToolSet,
240    policy: Option<&crate::tool_search::ToolAccessPolicy>,
241) -> String {
242    if deferred.is_empty() {
243        return String::new();
244    }
245    let mut out = String::new();
246    out.push_str("## Deferred Tools\n\n");
247    out.push_str(
248        "The tools listed below are available but NOT yet loaded. \
249         To use any of them you MUST first call the `tool_search` tool \
250         to fetch their full schemas. Use `\"select:name1,name2\"` for \
251         exact tools or keywords to search. Once activated, the tools \
252         become callable for the rest of the conversation.\n\n",
253    );
254    out.push_str("<available-deferred-tools>\n");
255    let mut count = 0;
256    for stub in &deferred.stubs {
257        if let Some(p) = policy
258            && !p.is_tool_allowed(&stub.prefixed_name)
259        {
260            continue;
261        }
262        out.push_str(&stub.prefixed_name);
263        out.push_str(" - ");
264        out.push_str(&stub.description);
265        out.push('\n');
266        count += 1;
267    }
268    out.push_str("</available-deferred-tools>\n");
269    if count == 0 {
270        return String::new();
271    }
272    out
273}
274
275#[cfg(test)]
276mod tests {
277    use super::*;
278
279    fn make_stub(name: &str, desc: &str) -> DeferredMcpToolStub {
280        let def = McpToolDef {
281            name: name.to_string(),
282            description: Some(desc.to_string()),
283            input_schema: serde_json::json!({"type": "object", "properties": {}}),
284        };
285        DeferredMcpToolStub::new(name.to_string(), def)
286    }
287
288    #[test]
289    fn stub_uses_description_from_def() {
290        let stub = make_stub("fs__read", "Read a file");
291        assert_eq!(stub.description, "Read a file");
292    }
293
294    #[test]
295    fn stub_defaults_description_when_none() {
296        let def = McpToolDef {
297            name: "mystery".into(),
298            description: None,
299            input_schema: serde_json::json!({}),
300        };
301        let stub = DeferredMcpToolStub::new("srv__mystery".into(), def);
302        assert_eq!(stub.description, "MCP tool");
303    }
304
305    #[test]
306    fn activated_set_tracks_activation() {
307        use async_trait::async_trait;
308        use zeroclaw_api::tool::ToolResult;
309
310        struct FakeTool;
311        impl ::zeroclaw_api::attribution::Attributable for FakeTool {
312            fn role(&self) -> ::zeroclaw_api::attribution::Role {
313                ::zeroclaw_api::attribution::Role::Tool(
314                    ::zeroclaw_api::attribution::ToolKind::Plugin,
315                )
316            }
317            fn alias(&self) -> &str {
318                <Self as Tool>::name(self)
319            }
320        }
321        #[async_trait]
322        impl Tool for FakeTool {
323            fn name(&self) -> &str {
324                "fake"
325            }
326            fn description(&self) -> &str {
327                "fake tool"
328            }
329            fn parameters_schema(&self) -> serde_json::Value {
330                serde_json::json!({})
331            }
332            async fn execute(&self, _: serde_json::Value) -> anyhow::Result<ToolResult> {
333                Ok(ToolResult {
334                    success: true,
335                    output: String::new(),
336                    error: None,
337                })
338            }
339        }
340
341        let mut set = ActivatedToolSet::new();
342        assert!(!set.is_activated("fake"));
343        set.activate("fake".into(), Arc::new(FakeTool));
344        assert!(set.is_activated("fake"));
345        assert!(set.get("fake").is_some());
346        assert_eq!(set.tool_specs().len(), 1);
347    }
348
349    #[test]
350    fn activated_set_resolves_unique_suffix() {
351        use async_trait::async_trait;
352        use zeroclaw_api::tool::ToolResult;
353
354        struct FakeTool;
355        impl ::zeroclaw_api::attribution::Attributable for FakeTool {
356            fn role(&self) -> ::zeroclaw_api::attribution::Role {
357                ::zeroclaw_api::attribution::Role::Tool(
358                    ::zeroclaw_api::attribution::ToolKind::Plugin,
359                )
360            }
361            fn alias(&self) -> &str {
362                <Self as Tool>::name(self)
363            }
364        }
365        #[async_trait]
366        impl Tool for FakeTool {
367            fn name(&self) -> &str {
368                "docker-mcp__extract_text"
369            }
370            fn description(&self) -> &str {
371                "fake tool"
372            }
373            fn parameters_schema(&self) -> serde_json::Value {
374                serde_json::json!({})
375            }
376            async fn execute(&self, _: serde_json::Value) -> anyhow::Result<ToolResult> {
377                Ok(ToolResult {
378                    success: true,
379                    output: String::new(),
380                    error: None,
381                })
382            }
383        }
384
385        let mut set = ActivatedToolSet::new();
386        set.activate("docker-mcp__extract_text".into(), Arc::new(FakeTool));
387        assert!(set.get_resolved("extract_text").is_some());
388    }
389
390    #[test]
391    fn activated_set_rejects_ambiguous_suffix() {
392        use async_trait::async_trait;
393        use zeroclaw_api::tool::ToolResult;
394
395        struct FakeTool(&'static str);
396        impl ::zeroclaw_api::attribution::Attributable for FakeTool {
397            fn role(&self) -> ::zeroclaw_api::attribution::Role {
398                ::zeroclaw_api::attribution::Role::Tool(
399                    ::zeroclaw_api::attribution::ToolKind::Plugin,
400                )
401            }
402            fn alias(&self) -> &str {
403                <Self as Tool>::name(self)
404            }
405        }
406        #[async_trait]
407        impl Tool for FakeTool {
408            fn name(&self) -> &str {
409                self.0
410            }
411            fn description(&self) -> &str {
412                "fake tool"
413            }
414            fn parameters_schema(&self) -> serde_json::Value {
415                serde_json::json!({})
416            }
417            async fn execute(&self, _: serde_json::Value) -> anyhow::Result<ToolResult> {
418                Ok(ToolResult {
419                    success: true,
420                    output: String::new(),
421                    error: None,
422                })
423            }
424        }
425
426        let mut set = ActivatedToolSet::new();
427        set.activate(
428            "docker-mcp__extract_text".into(),
429            Arc::new(FakeTool("docker-mcp__extract_text")),
430        );
431        set.activate(
432            "ocr-mcp__extract_text".into(),
433            Arc::new(FakeTool("ocr-mcp__extract_text")),
434        );
435        assert!(set.get_resolved("extract_text").is_none());
436    }
437
438    #[test]
439    fn build_deferred_section_empty_when_no_stubs() {
440        let set = DeferredMcpToolSet {
441            stubs: vec![],
442            registry: std::sync::Arc::new(
443                tokio::runtime::Runtime::new()
444                    .unwrap()
445                    .block_on(McpRegistry::connect_all(&[]))
446                    .unwrap(),
447            ),
448        };
449        assert!(build_deferred_tools_section(&set).is_empty());
450    }
451
452    #[test]
453    fn build_deferred_section_lists_names() {
454        let stubs = vec![
455            make_stub("fs__read_file", "Read a file"),
456            make_stub("git__status", "Git status"),
457        ];
458        let set = DeferredMcpToolSet {
459            stubs,
460            registry: std::sync::Arc::new(
461                tokio::runtime::Runtime::new()
462                    .unwrap()
463                    .block_on(McpRegistry::connect_all(&[]))
464                    .unwrap(),
465            ),
466        };
467        let section = build_deferred_tools_section(&set);
468        assert!(section.contains("<available-deferred-tools>"));
469        assert!(section.contains("fs__read_file - Read a file"));
470        assert!(section.contains("git__status - Git status"));
471        assert!(section.contains("</available-deferred-tools>"));
472    }
473
474    #[test]
475    fn build_deferred_section_includes_tool_search_instruction() {
476        let stubs = vec![make_stub("fs__read_file", "Read a file")];
477        let set = DeferredMcpToolSet {
478            stubs,
479            registry: std::sync::Arc::new(
480                tokio::runtime::Runtime::new()
481                    .unwrap()
482                    .block_on(McpRegistry::connect_all(&[]))
483                    .unwrap(),
484            ),
485        };
486        let section = build_deferred_tools_section(&set);
487        assert!(
488            section.contains("tool_search"),
489            "deferred section must instruct the LLM to use tool_search"
490        );
491        assert!(
492            section.contains("## Deferred Tools"),
493            "deferred section must include a heading"
494        );
495    }
496
497    #[test]
498    fn build_deferred_section_multiple_servers() {
499        let stubs = vec![
500            make_stub("server_a__list", "List items"),
501            make_stub("server_a__create", "Create item"),
502            make_stub("server_b__query", "Query records"),
503        ];
504        let set = DeferredMcpToolSet {
505            stubs,
506            registry: std::sync::Arc::new(
507                tokio::runtime::Runtime::new()
508                    .unwrap()
509                    .block_on(McpRegistry::connect_all(&[]))
510                    .unwrap(),
511            ),
512        };
513        let section = build_deferred_tools_section(&set);
514        assert!(section.contains("server_a__list"));
515        assert!(section.contains("server_a__create"));
516        assert!(section.contains("server_b__query"));
517        assert!(
518            section.contains("tool_search"),
519            "section must mention tool_search for multi-server setups"
520        );
521    }
522
523    #[test]
524    fn keyword_search_ranks_by_hits() {
525        let stubs = vec![
526            make_stub("fs__read_file", "Read a file from disk"),
527            make_stub("fs__write_file", "Write a file to disk"),
528            make_stub("git__log", "Show git log"),
529        ];
530        let set = DeferredMcpToolSet {
531            stubs,
532            registry: std::sync::Arc::new(
533                tokio::runtime::Runtime::new()
534                    .unwrap()
535                    .block_on(McpRegistry::connect_all(&[]))
536                    .unwrap(),
537            ),
538        };
539
540        // "file read" should rank fs__read_file highest (2 hits vs 1)
541        let results = set.search("file read", 5);
542        assert!(!results.is_empty());
543        assert_eq!(results[0].prefixed_name, "fs__read_file");
544    }
545
546    #[test]
547    fn get_by_name_returns_correct_stub() {
548        let stubs = vec![
549            make_stub("a__one", "Tool one"),
550            make_stub("b__two", "Tool two"),
551        ];
552        let set = DeferredMcpToolSet {
553            stubs,
554            registry: std::sync::Arc::new(
555                tokio::runtime::Runtime::new()
556                    .unwrap()
557                    .block_on(McpRegistry::connect_all(&[]))
558                    .unwrap(),
559            ),
560        };
561        assert!(set.get_by_name("a__one").is_some());
562        assert!(set.get_by_name("nonexistent").is_none());
563    }
564
565    #[test]
566    fn search_across_multiple_servers() {
567        let stubs = vec![
568            make_stub("server_a__read_file", "Read a file from disk"),
569            make_stub("server_b__read_config", "Read configuration from database"),
570        ];
571        let set = DeferredMcpToolSet {
572            stubs,
573            registry: std::sync::Arc::new(
574                tokio::runtime::Runtime::new()
575                    .unwrap()
576                    .block_on(McpRegistry::connect_all(&[]))
577                    .unwrap(),
578            ),
579        };
580
581        // "read" should match stubs from both servers
582        let results = set.search("read", 10);
583        assert_eq!(results.len(), 2);
584
585        // "file" should match only server_a
586        let results = set.search("file", 10);
587        assert_eq!(results.len(), 1);
588        assert_eq!(results[0].prefixed_name, "server_a__read_file");
589
590        // "config database" should rank server_b highest (2 hits)
591        let results = set.search("config database", 10);
592        assert!(!results.is_empty());
593        assert_eq!(results[0].prefixed_name, "server_b__read_config");
594    }
595}
zeroclaw_tools/mcp_deferred.rs

zeroclaw_tools/
mcp_deferred.rs
