thndrs_lib/core/mcp/
manager.rs

1//! MCP client manager and namespaced tool routing.
2//!
3//! The manager keeps MCP protocol and process handling behind the adapter while
4//! exposing the same local tool definition and output shapes used by built-in
5//! tools.
6
7use std::collections::BTreeMap;
8
9use super::adapter::{McpSdkClient, McpSdkError, McpToolDefinition, sdk_error_to_output};
10use super::config::{McpConfig, McpServerConfig};
11use crate::tools::{MAX_LINE_LEN, MAX_OUTPUT_BYTES, ToolDefinition, ToolOutput, ToolUseRequest, shell};
12
13const MAX_MCP_OUTPUT_LINES: usize = 100;
14
15/// Errors returned while building an MCP client or manager.
16#[derive(Debug, thiserror::Error)]
17pub enum McpManagerError {
18    /// The server initialized, but listing tools failed.
19    #[error("mcp server `{server}` failed to list tools: {source}")]
20    ListTools { server: String, source: McpSdkError },
21    /// The server could not be initialized.
22    #[error("mcp server `{server}` failed to initialize: {source}")]
23    Initialize { server: String, source: McpSdkError },
24}
25
26/// Connected MCP client for one configured server.
27pub struct McpClient {
28    name: String,
29    sdk: McpSdkClient,
30    tools: Vec<McpToolDefinition>,
31    tool_routes: BTreeMap<String, String>,
32}
33
34impl McpClient {
35    /// Initialize a server and cache its tool definitions.
36    pub fn connect(name: impl Into<String>, config: &McpServerConfig) -> Result<Self, McpManagerError> {
37        let name = name.into();
38        let sdk = McpSdkClient::connect(config)
39            .map_err(|source| McpManagerError::Initialize { server: name.clone(), source })?;
40        let tools = sdk
41            .list_tool_definitions(&name)
42            .map_err(|source| McpManagerError::ListTools { server: name.clone(), source })?;
43        let tool_routes = tools
44            .iter()
45            .map(|tool| (tool.definition.name.to_string(), tool.original_tool_name.clone()))
46            .collect();
47
48        Ok(Self { name, sdk, tools, tool_routes })
49    }
50
51    /// Cached tool definitions for this server.
52    pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
53        self.tools.iter().map(|tool| tool.definition.clone()).collect()
54    }
55
56    /// Current startup and stderr diagnostics for this server.
57    pub fn diagnostics(&self) -> Vec<String> {
58        let mut diagnostics = self.sdk.server_info().diagnostics.clone();
59        if let Some(stderr) = self.sdk.stderr_diagnostics() {
60            diagnostics.push(format!("mcp server `{}` stderr: {stderr}", self.name));
61        }
62        diagnostics
63    }
64
65    /// Return the original MCP tool name for a namespaced provider-visible name.
66    pub fn original_tool_name(&self, namespaced_tool_name: &str) -> Option<&str> {
67        self.tool_routes.get(namespaced_tool_name).map(String::as_str)
68    }
69
70    /// Call a namespaced MCP tool through this client.
71    pub fn call_tool(&self, request: &ToolUseRequest) -> ToolOutput {
72        let Some(original_name) = self.original_tool_name(&request.name) else {
73            return ToolOutput::failed(&request.name, format!("unknown MCP tool: {}", request.name));
74        };
75        let arguments = match serde_json::from_str::<serde_json::Value>(&request.arguments) {
76            Ok(arguments) => arguments,
77            Err(err) => return ToolOutput::failed(&request.name, format!("invalid tool arguments: {err}")),
78        };
79
80        self.sdk
81            .call_tool(&request.name, original_name, arguments)
82            .map(sanitize_mcp_output)
83            .unwrap_or_else(|err| sdk_error_to_output(&request.name, "call", &err))
84    }
85}
86
87/// Manager for configured MCP servers.
88pub struct McpManager {
89    clients: BTreeMap<String, McpClient>,
90    tool_routes: BTreeMap<String, String>,
91    diagnostics: Vec<String>,
92}
93
94impl std::fmt::Debug for McpManager {
95    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
96        f.debug_struct("McpManager")
97            .field("clients", &self.clients.keys().collect::<Vec<_>>())
98            .field("tool_routes", &self.tool_routes)
99            .field("diagnostics", &self.diagnostics)
100            .finish()
101    }
102}
103
104impl McpManager {
105    /// Initialize all enabled servers with bounded startup and cache their tools.
106    pub fn from_config(config: &McpConfig) -> Self {
107        let mut clients = BTreeMap::new();
108        let mut tool_routes = BTreeMap::new();
109        let mut diagnostics = Vec::new();
110
111        for (name, server) in &config.servers {
112            if !server.enabled {
113                diagnostics.push(format!("mcp server `{name}` disabled"));
114                continue;
115            }
116
117            match McpClient::connect(name.clone(), server) {
118                Ok(client) => {
119                    for tool in &client.tools {
120                        tool_routes.insert(tool.definition.name.to_string(), name.clone());
121                    }
122                    diagnostics.extend(client.diagnostics());
123                    clients.insert(name.clone(), client);
124                }
125                Err(err) => diagnostics.push(err.to_string()),
126            }
127        }
128
129        Self { clients, tool_routes, diagnostics }
130    }
131
132    /// Cached MCP tool definitions for provider prompt assembly.
133    pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
134        self.clients.values().flat_map(McpClient::tool_definitions).collect()
135    }
136
137    /// Startup and non-fatal diagnostics recorded while initializing servers.
138    pub fn diagnostics(&self) -> &[String] {
139        &self.diagnostics
140    }
141
142    /// Append loader diagnostics to this manager.
143    pub fn extend_diagnostics(&mut self, diagnostics: impl IntoIterator<Item = String>) {
144        self.diagnostics.extend(diagnostics);
145    }
146
147    /// Route a namespaced tool request to the correct MCP client.
148    pub fn call_tool(&self, request: &ToolUseRequest) -> ToolOutput {
149        let Some(server_name) = self.tool_routes.get(&request.name) else {
150            return ToolOutput::failed(&request.name, format!("unknown MCP tool: {}", request.name));
151        };
152        let Some(client) = self.clients.get(server_name) else {
153            return ToolOutput::failed(&request.name, format!("MCP server `{server_name}` is not available"));
154        };
155
156        client.call_tool(request)
157    }
158}
159
160fn sanitize_mcp_output(mut output: ToolOutput) -> ToolOutput {
161    let mut remaining_bytes = MAX_OUTPUT_BYTES;
162    let mut sanitized = Vec::new();
163    let mut truncated = false;
164
165    for line in output.display.lines {
166        if sanitized.len() >= MAX_MCP_OUTPUT_LINES {
167            truncated = true;
168            break;
169        }
170
171        let redacted = shell::redact_secrets(&line);
172        let mut capped: String = redacted.chars().take(MAX_LINE_LEN).collect();
173        if capped.len() < redacted.len() {
174            capped.push_str("...");
175            truncated = true;
176        }
177
178        let bytes = capped.len();
179        if bytes > remaining_bytes {
180            let allowed = remaining_bytes.min(capped.len());
181            capped = capped.chars().take(allowed).collect();
182            capped.push_str("...");
183            sanitized.push(capped);
184            truncated = true;
185            break;
186        }
187
188        remaining_bytes = remaining_bytes.saturating_sub(bytes);
189        sanitized.push(capped);
190    }
191
192    if truncated {
193        sanitized.push("[mcp output truncated]".to_string());
194    }
195
196    if let Some(error) = output.error {
197        output.error = Some(shell::redact_secrets(&error));
198    }
199    output.display.lines = sanitized.clone();
200    output.model.lines = sanitized;
201    output.evidence.byte_count = output.display.lines.join("\n").len();
202    output
203}
204
205#[cfg(test)]
206mod tests {
207    use std::fs;
208    use std::path::Path;
209
210    use serde_json::json;
211
212    use super::*;
213    use crate::app::ToolStatus;
214    use crate::mcp::config::McpServerConfig;
215
216    #[test]
217    fn manager_routes_duplicate_tool_names_by_namespace() {
218        let temp = tempfile::tempdir().expect("tempdir");
219        let script = write_fake_server(temp.path(), "echo", 0, false);
220        let mut config = McpConfig::default();
221        config.servers.insert("alpha".to_string(), server_config(&script));
222        config.servers.insert("beta".to_string(), server_config(&script));
223
224        let manager = McpManager::from_config(&config);
225        let names: Vec<String> = manager
226            .tool_definitions()
227            .into_iter()
228            .map(|tool| tool.name.to_string())
229            .collect();
230
231        assert_eq!(names, vec!["mcp__alpha__echo", "mcp__beta__echo"]);
232
233        let output = manager.call_tool(&ToolUseRequest::new(
234            "mcp__beta__echo".to_string(),
235            json!({ "text": "from beta" }).to_string(),
236            "toolu_1".to_string(),
237        ));
238
239        assert_eq!(output, ToolOutput::ok("mcp__beta__echo", vec!["from beta".to_string()]));
240    }
241
242    #[test]
243    fn manager_records_failed_server_diagnostics() {
244        let temp = tempfile::tempdir().expect("tempdir");
245        let script = write_fake_server(temp.path(), "echo", 2, false);
246        let mut config = McpConfig::default();
247        config.servers.insert(
248            "slow".to_string(),
249            McpServerConfig { timeout_secs: 1, ..server_config(&script) },
250        );
251
252        let manager = McpManager::from_config(&config);
253
254        assert!(manager.tool_definitions().is_empty());
255        assert!(
256            manager
257                .diagnostics()
258                .iter()
259                .any(|diagnostic| diagnostic.contains("timed out after 1s"))
260        );
261    }
262
263    #[test]
264    fn manager_returns_stable_failure_for_unknown_tool() {
265        let manager = McpManager { clients: BTreeMap::new(), tool_routes: BTreeMap::new(), diagnostics: Vec::new() };
266
267        let output = manager.call_tool(&ToolUseRequest::new(
268            "mcp__missing__echo".to_string(),
269            "{}".to_string(),
270            "toolu_1".to_string(),
271        ));
272
273        assert_eq!(output.status, ToolStatus::Failed);
274        assert_eq!(output.error.as_deref(), Some("unknown MCP tool: mcp__missing__echo"));
275    }
276
277    #[test]
278    fn mcp_output_is_redacted_and_capped() {
279        let output = sanitize_mcp_output(ToolOutput::ok(
280            "mcp__docs__secret",
281            vec![format!("api_key=sk-{}", "a".repeat(80)), "x".repeat(MAX_LINE_LEN + 20)],
282        ));
283
284        assert!(output.display.lines[0].contains("[REDACTED]"));
285        assert!(output.display.lines[1].ends_with("..."));
286        assert!(output.display.lines.iter().any(|line| line == "[mcp output truncated]"));
287    }
288
289    #[test]
290    fn provider_catalog_includes_cached_mcp_stdio_tools() {
291        let temp = tempfile::tempdir().expect("tempdir");
292        let script = write_fake_server(temp.path(), "echo", 0, false);
293        let mut config = McpConfig::default();
294        config.servers.insert("docs".to_string(), server_config(&script));
295
296        let manager = McpManager::from_config(&config);
297        let definitions = crate::tools::runtime_tool_definitions(Some(&manager));
298        let schemas = crate::tools::tool_catalog_schemas(&definitions);
299        assert!(
300            schemas
301                .as_array()
302                .expect("schemas")
303                .iter()
304                .filter_map(|schema| schema["name"].as_str())
305                .any(|name| name == "mcp__docs__echo")
306        );
307
308        assert_eq!(
309            schemas
310                .as_array()
311                .expect("schemas")
312                .iter()
313                .find(|schema| schema["name"] == "mcp__docs__echo")
314                .expect("mcp schema")["input_schema"]["properties"]["text"]["type"],
315            "string"
316        );
317    }
318
319    fn server_config(script: &Path) -> McpServerConfig {
320        McpServerConfig {
321            command: "python3".to_string(),
322            args: vec![script.display().to_string()],
323            timeout_secs: 5,
324            ..McpServerConfig::default()
325        }
326    }
327
328    fn write_fake_server(
329        dir: &Path, tool_name: &str, initialize_sleep_secs: u64, exit_after_initialize: bool,
330    ) -> std::path::PathBuf {
331        let path = dir.join(format!("fake_{tool_name}.py"));
332        let exit_after_initialize = python_bool(exit_after_initialize);
333        fs::write(
334            &path,
335            format!(
336                r#"#!/usr/bin/env python3
337import json
338import sys
339import time
340
341tool_name = {tool_name:?}
342initialize_sleep_secs = {initialize_sleep_secs}
343exit_after_initialize = {exit_after_initialize}
344
345for line in sys.stdin:
346    msg = json.loads(line)
347    method = msg.get("method")
348    if method == "initialize":
349        if initialize_sleep_secs:
350            time.sleep(initialize_sleep_secs)
351        print(json.dumps({{
352            "jsonrpc": "2.0",
353            "id": msg["id"],
354            "result": {{
355                "protocolVersion": "2025-06-18",
356                "capabilities": {{"tools": {{}}}},
357                "serverInfo": {{"name": "fake", "version": "0.1.0"}}
358            }}
359        }}), flush=True)
360        if exit_after_initialize:
361            sys.exit(0)
362    elif method == "notifications/initialized":
363        continue
364    elif method == "tools/list":
365        print(json.dumps({{
366            "jsonrpc": "2.0",
367            "id": msg["id"],
368            "result": {{
369                "tools": [{{
370                    "name": tool_name,
371                    "description": "Echo text",
372                    "inputSchema": {{"type": "object", "properties": {{"text": {{"type": "string"}}}}}}
373                }}]
374            }}
375        }}), flush=True)
376    elif method == "tools/call":
377        args = msg.get("params", {{}}).get("arguments", {{}})
378        print(json.dumps({{
379            "jsonrpc": "2.0",
380            "id": msg["id"],
381            "result": {{"content": [{{"type": "text", "text": args.get("text", "")}}], "isError": False}}
382        }}), flush=True)
383"#
384            ),
385        )
386        .expect("write fake server");
387        path
388    }
389
390    fn python_bool(value: bool) -> &'static str {
391        if value { "True" } else { "False" }
392    }
393}
thndrs_lib/core/mcp/manager.rs

thndrs_lib/core/mcp/
manager.rs
