Skip to main content

mermaid_cli/providers/tool/
mod.rs

1//! Tool executors — one type per tool the model can call.
2//!
3//! The trait is small: `execute(args, ctx) -> ToolOutcome` for
4//! dispatch, plus `schema() -> ToolDefinition` for advertising the
5//! tool to the model. Everything else (cancellation, progress,
6//! identity, workdir) rides inside `ExecContext`.
7//!
8//! Adding a tool:
9//!   1. New file under `src/providers/tool/`.
10//!   2. Impl `ToolExecutor` for a unit struct — both `execute` and
11//!      `schema`.
12//!   3. Register it in `ToolRegistry::default()`.
13//!
14//! Because `schema()` lives on the same trait as `execute()`, the
15//! name + JSON schema the model sees cannot drift from the handler
16//! that runs when the model calls it. Single source of truth.
17
18pub mod computer_use;
19pub mod exec;
20pub mod filesystem;
21pub mod mcp;
22pub mod memory;
23pub mod path_safety;
24pub mod policy_gate;
25pub mod subagent;
26pub mod web;
27pub mod web_client;
28
29use async_trait::async_trait;
30use std::collections::HashMap;
31use std::sync::Arc;
32
33use crate::domain::{ToolDefinition, ToolOutcome};
34
35use super::ctx::ExecContext;
36
37/// Implemented by every tool that the model can call. All tools are
38/// `Send + Sync` — they run across tokio `select!` branches inside
39/// the effect runner.
40#[async_trait]
41pub trait ToolExecutor: Send + Sync {
42    /// Canonical name the model uses to call this tool. Matches
43    /// `schema().name` exactly.
44    fn name(&self) -> &'static str;
45
46    /// JSON-schema description the model sees in the outgoing
47    /// request. Adapters translate this into provider-native shape
48    /// (Anthropic's `type: "custom"`, Gemini's `function_declarations`,
49    /// OpenAI's flat `tools`, Ollama's function calling). The same
50    /// `ToolDefinition` feeds all four.
51    fn schema(&self) -> ToolDefinition;
52
53    /// True for tools that exist for internal dispatch only and
54    /// should NOT be advertised to the model (e.g. the MCP proxy
55    /// router, which fronts every `mcp__server__tool` call — the
56    /// individual MCP tools are advertised separately from
57    /// `state.mcp.servers`). Default `false`.
58    fn is_internal(&self) -> bool {
59        false
60    }
61
62    /// Run the tool. The returned `ToolOutcome` is passed verbatim
63    /// into `Msg::ToolFinished` — there's no error-to-outcome
64    /// conversion happening outside this function.
65    async fn execute(&self, args: serde_json::Value, ctx: ExecContext) -> ToolOutcome;
66}
67
68/// Registry of dispatchable tools. Single source of truth for what
69/// the model sees AND what handles a call when the model issues it.
70/// Built once at startup; read-only after that.
71pub struct ToolRegistry {
72    entries: HashMap<&'static str, Arc<dyn ToolExecutor>>,
73}
74
75impl ToolRegistry {
76    pub fn new() -> Self {
77        Self {
78            entries: HashMap::new(),
79        }
80    }
81
82    pub fn register(&mut self, tool: Arc<dyn ToolExecutor>) {
83        self.entries.insert(tool.name(), tool);
84    }
85
86    pub fn get(&self, name: &str) -> Option<Arc<dyn ToolExecutor>> {
87        self.entries.get(name).cloned()
88    }
89
90    pub fn len(&self) -> usize {
91        self.entries.len()
92    }
93
94    pub fn is_empty(&self) -> bool {
95        self.entries.is_empty()
96    }
97
98    pub fn names(&self) -> impl Iterator<Item = &'static str> + '_ {
99        self.entries.keys().copied()
100    }
101
102    /// Emit every user-facing tool's schema, for inclusion in an
103    /// outgoing `ChatRequest.tools`. Effect runner calls this before
104    /// dispatching `Cmd::CallModel` so the model always sees the
105    /// same list the runner can dispatch. Internal routers (the MCP
106    /// proxy) are filtered out.
107    pub fn describe_all(&self) -> Vec<ToolDefinition> {
108        self.entries
109            .values()
110            .filter(|t| !t.is_internal())
111            .map(|t| t.schema())
112            .collect()
113    }
114}
115
116impl Default for ToolRegistry {
117    fn default() -> Self {
118        let mut r = Self::new();
119        r.register(Arc::new(filesystem::ReadFileTool));
120        r.register(Arc::new(filesystem::WriteFileTool));
121        r.register(Arc::new(filesystem::EditFileTool));
122        r.register(Arc::new(filesystem::DeleteFileTool));
123        r.register(Arc::new(filesystem::CreateDirectoryTool));
124        r.register(Arc::new(exec::ExecuteCommandTool));
125        r.register(Arc::new(memory::MemoryTool));
126        // MCP proxy is the dispatcher for every mcp__server__tool
127        // call; it's internal (not advertised) but MUST be registered
128        // so runtime lookups succeed.
129        r.register(Arc::new(mcp::McpToolProxy));
130        r
131    }
132}
133
134/// Whether the host mermaid process is running interactively (TUI)
135/// or headlessly (one-shot `mermaid run <prompt>` / CI). Controls
136/// which tools get registered: headless mode never advertises
137/// GUI / computer-use tools even when a display probes alive, because
138/// a CI job has no user to watch the screenshot.
139#[derive(Debug, Clone, Copy, PartialEq, Eq)]
140pub enum TuiMode {
141    Interactive,
142    Headless,
143}
144
145impl ToolRegistry {
146    /// Register the computer-use tools the given backend can actually drive.
147    /// Screenshot works on every usable backend; the five input tools need
148    /// pointer/keyboard injection (X11/Wayland); `list_windows` is X11-only.
149    /// Advertising a tool the backend can't drive means the model is told it
150    /// has a capability that `bail!`s at call time (#35).
151    fn register_computer_use_tools(&mut self, backend: computer_use::Backend) {
152        let driver = Arc::new(computer_use::ComputerUseDriver::new(backend));
153        self.register(Arc::new(computer_use::ScreenshotTool::new(driver.clone())));
154        if backend.supports_input_injection() {
155            self.register(Arc::new(computer_use::ClickTool::new(driver.clone())));
156            self.register(Arc::new(computer_use::TypeTextTool::new(driver.clone())));
157            self.register(Arc::new(computer_use::PressKeyTool::new(driver.clone())));
158            self.register(Arc::new(computer_use::ScrollTool::new(driver.clone())));
159            self.register(Arc::new(computer_use::MouseMoveTool::new(driver.clone())));
160        }
161        if backend.supports_window_listing() {
162            self.register(Arc::new(computer_use::ListWindowsTool::new(driver.clone())));
163        }
164    }
165
166    /// Config-aware factory. Always registers filesystem + exec +
167    /// the MCP proxy + the subagent tool. Conditionally registers:
168    ///
169    ///   - `web_fetch` and `web_search` per `config.web`. `web_fetch`
170    ///     defaults to the native backend (always registered, no key);
171    ///     `web_search` defaults to Ollama Cloud (registered iff
172    ///     `OLLAMA_API_KEY` resolves) or SearXNG (always). See
173    ///     `web::web_fetch_tool` / `web::web_search_tool`.
174    ///   - The computer-use tools the detected backend can drive (see
175    ///     `register_computer_use_tools`) iff `mode == Interactive` AND
176    ///     `computer_use::probe()` returns a usable backend.
177    ///
178    /// `providers` is the shared `ProviderFactory` that the effect
179    /// runner also holds; the `SubagentSpawner` needs it so child
180    /// reducer loops hit the same provider cache.
181    ///
182    /// Returns `Arc<Self>` so the effect runner can share a handle
183    /// across turns without cloning the underlying HashMap.
184    pub fn build(
185        config: &crate::app::Config,
186        mode: TuiMode,
187        providers: Arc<crate::providers::ProviderFactory>,
188    ) -> Arc<Self> {
189        let mut r = Self::new();
190        r.register(Arc::new(filesystem::ReadFileTool));
191        r.register(Arc::new(filesystem::WriteFileTool));
192        r.register(Arc::new(filesystem::EditFileTool));
193        r.register(Arc::new(filesystem::DeleteFileTool));
194        r.register(Arc::new(filesystem::CreateDirectoryTool));
195        r.register(Arc::new(exec::ExecuteCommandTool));
196        r.register(Arc::new(memory::MemoryTool));
197        r.register(Arc::new(mcp::McpToolProxy));
198
199        if let Some(tool) = web::web_fetch_tool(&config.web) {
200            r.register(Arc::new(tool));
201        }
202        if let Some(tool) = web::web_search_tool(&config.web) {
203            r.register(Arc::new(tool));
204        }
205
206        // Computer-use tools only register when (a) the process runs
207        // interactively (Headless CI has no user to watch a screenshot)
208        // AND (b) a display backend passes the startup probe. Failed
209        // probe → tools aren't advertised → model can't call them.
210        if mode == TuiMode::Interactive {
211            let backend = computer_use::probe();
212            if backend.is_usable() {
213                r.register_computer_use_tools(backend);
214            }
215        }
216
217        // Subagents: always register. Depth + breadth caps live on
218        // `SubagentSpawner`; the tool itself is harmless when nobody
219        // calls it. Headless runs do register the agent — a CI prompt
220        // may still delegate to subagents for batched work.
221        let spawner = Arc::new(subagent::SubagentSpawner::new(providers));
222        r.register(Arc::new(subagent::SubagentTool::new(spawner)));
223
224        Arc::new(r)
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231
232    #[test]
233    fn default_registry_has_builtin_tools() {
234        let r = ToolRegistry::default();
235        for name in &[
236            "read_file",
237            "write_file",
238            "edit_file",
239            "delete_file",
240            "create_directory",
241            "execute_command",
242            "memory",
243        ] {
244            assert!(r.get(name).is_some(), "missing: {}", name);
245        }
246        assert!(r.get("not_a_tool").is_none());
247        assert!(r.len() >= 6);
248    }
249
250    #[test]
251    fn computer_use_registration_is_selective_per_backend() {
252        use computer_use::Backend;
253        let reg = |b: Backend| {
254            let mut r = ToolRegistry::new();
255            r.register_computer_use_tools(b);
256            r
257        };
258
259        // macOS: only screenshot — the input verbs + list_windows bail at
260        // runtime, so advertising them just wastes the model's turns (#35).
261        let mac = reg(Backend::MacOS);
262        assert!(mac.get("screenshot").is_some());
263        for t in [
264            "click",
265            "type_text",
266            "press_key",
267            "scroll",
268            "mouse_move",
269            "list_windows",
270        ] {
271            assert!(mac.get(t).is_none(), "macOS must not advertise {t}");
272        }
273
274        // Wayland: input tools, but no list_windows (no portable enumeration).
275        let way = reg(Backend::Wayland);
276        assert!(way.get("click").is_some());
277        assert!(way.get("list_windows").is_none());
278
279        // X11: all seven.
280        let x11 = reg(Backend::X11);
281        for t in [
282            "screenshot",
283            "click",
284            "type_text",
285            "press_key",
286            "scroll",
287            "mouse_move",
288            "list_windows",
289        ] {
290            assert!(x11.get(t).is_some(), "X11 missing {t}");
291        }
292    }
293
294    #[test]
295    fn describe_all_returns_one_per_user_facing_tool() {
296        let r = ToolRegistry::default();
297        let schemas = r.describe_all();
298        // mcp_proxy is registered but internal — filtered out of
299        // describe_all. So len() includes it but schemas don't.
300        let visible = r
301            .names()
302            .filter(|n| r.get(n).map(|t| !t.is_internal()).unwrap_or(false))
303            .count();
304        assert_eq!(schemas.len(), visible);
305        for schema in &schemas {
306            assert!(
307                r.get(&schema.name).is_some(),
308                "schema for unknown tool: {}",
309                schema.name
310            );
311        }
312    }
313
314    #[test]
315    fn mcp_proxy_is_registered_but_internal() {
316        let r = ToolRegistry::default();
317        let proxy = r.get("mcp_proxy").expect("mcp_proxy registered");
318        assert!(proxy.is_internal());
319        assert!(!r.describe_all().iter().any(|s| s.name == "mcp_proxy"));
320    }
321
322    #[test]
323    fn schema_name_matches_executor_name() {
324        let r = ToolRegistry::default();
325        for name in r.names() {
326            let tool = r.get(name).unwrap();
327            assert_eq!(tool.name(), tool.schema().name.as_str());
328        }
329    }
330
331    /// Serialization guard for tests that mutate the `OLLAMA_API_KEY`
332    /// env var. Cargo's default test harness runs tests in parallel
333    /// threads inside one process; without this mutex two env-touching
334    /// tests would race and occasionally flip each other's expectations.
335    static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
336
337    #[test]
338    fn build_registers_zero_config_web_tools_without_key() {
339        // Both web tools register with no OLLAMA_API_KEY: web_fetch is native,
340        // and web_search defaults to `auto`, which falls back to a managed local
341        // SearXNG (the container starts lazily at call time, not here).
342        let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
343        let prior = std::env::var("OLLAMA_API_KEY").ok();
344        unsafe {
345            std::env::remove_var("OLLAMA_API_KEY");
346        }
347        let cfg = crate::app::Config::default();
348        let providers = Arc::new(crate::providers::ProviderFactory::new(cfg.clone()));
349        let r = ToolRegistry::build(&cfg, TuiMode::Headless, providers);
350        assert!(
351            r.get("web_fetch").is_some(),
352            "native web_fetch registers without a key"
353        );
354        assert!(
355            r.get("web_search").is_some(),
356            "auto web_search (managed SearXNG) registers without a key"
357        );
358        assert!(r.get("read_file").is_some());
359        assert!(r.get("execute_command").is_some());
360        unsafe {
361            if let Some(v) = prior {
362                std::env::set_var("OLLAMA_API_KEY", v);
363            }
364        }
365    }
366
367    #[test]
368    fn build_registers_ollama_web_search_with_key() {
369        // With a key, the default Ollama search backend registers too.
370        let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
371        let prior = std::env::var("OLLAMA_API_KEY").ok();
372        unsafe {
373            std::env::set_var("OLLAMA_API_KEY", "test-key-build");
374        }
375        let cfg = crate::app::Config::default();
376        let providers = Arc::new(crate::providers::ProviderFactory::new(cfg.clone()));
377        let r = ToolRegistry::build(&cfg, TuiMode::Interactive, providers);
378        assert!(r.get("web_search").is_some(), "web_search registered");
379        assert!(r.get("web_fetch").is_some(), "web_fetch registered");
380        unsafe {
381            match prior {
382                Some(v) => std::env::set_var("OLLAMA_API_KEY", v),
383                None => std::env::remove_var("OLLAMA_API_KEY"),
384            }
385        }
386    }
387
388    #[test]
389    fn build_registers_searxng_web_search_without_key() {
390        // The SearXNG search backend registers regardless of OLLAMA_API_KEY —
391        // reachability is a call-time concern, not a registration one.
392        let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
393        let prior = std::env::var("OLLAMA_API_KEY").ok();
394        unsafe {
395            std::env::remove_var("OLLAMA_API_KEY");
396        }
397        let mut cfg = crate::app::Config::default();
398        cfg.web.search_backend = crate::app::SearchBackend::Searxng;
399        let providers = Arc::new(crate::providers::ProviderFactory::new(cfg.clone()));
400        let r = ToolRegistry::build(&cfg, TuiMode::Headless, providers);
401        assert!(
402            r.get("web_search").is_some(),
403            "searxng web_search registers without a key"
404        );
405        assert!(
406            r.get("web_fetch").is_some(),
407            "native web_fetch still present"
408        );
409        unsafe {
410            if let Some(v) = prior {
411                std::env::set_var("OLLAMA_API_KEY", v);
412            }
413        }
414    }
415}