telos_agent/integrations/mcp/
client.rs

1use std::io::{BufRead, BufReader, Write};
2use std::process::{Child, Command, Stdio};
3use std::sync::Mutex;
4use std::sync::atomic::{AtomicU64, Ordering};
5
6use serde_json::{Value, json};
7
8use crate::error::AgentError;
9use crate::integrations::mcp::config::McpServerConfig;
10
11/// Monotonically increasing JSON-RPC request ID.
12static NEXT_ID: AtomicU64 = AtomicU64::new(1);
13
14/// A tool definition returned by an MCP server during the `tools/list` handshake.
15#[derive(Debug, Clone)]
16pub struct McpTool {
17    /// The tool's unique name (used as `name` in `tools/call`).
18    pub name: String,
19    /// A human-readable description.
20    pub description: String,
21    /// The JSON Schema for the tool's arguments.
22    pub input_schema: Value,
23}
24
25/// A connected MCP server via stdio transport.
26///
27/// Spawns the configured command as a child process, performs the MCP
28/// `initialize` handshake, fetches the tool list, and provides a
29/// `call_tool` method to invoke server-side tools.
30///
31/// # Thread safety
32///
33/// All mutable state is behind `Mutex`, making `&self` the only
34/// shared reference needed for every operation. The struct is `Sync`.
35pub struct McpClient {
36    config: McpServerConfig,
37    process: Mutex<Option<Child>>,
38    stdin: Mutex<Option<std::process::ChildStdin>>,
39    reader: Mutex<Option<BufReader<std::process::ChildStdout>>>,
40    server_info: Mutex<Option<Value>>,
41    tools: Mutex<Vec<McpTool>>,
42}
43
44impl McpClient {
45    /// Create a new client for the given server configuration.
46    ///
47    /// The server is not spawned until [`connect`](Self::connect) is called.
48    pub fn new(config: McpServerConfig) -> Self {
49        Self {
50            config,
51            process: Mutex::new(None),
52            stdin: Mutex::new(None),
53            reader: Mutex::new(None),
54            server_info: Mutex::new(None),
55            tools: Mutex::new(Vec::new()),
56        }
57    }
58
59    /// Spawn the server process and perform the MCP initialize handshake.
60    ///
61    /// After the handshake succeeds the client sends the `notifications/initialized`
62    /// notification and fetches the tool list via `tools/list`.
63    pub async fn connect(&self) -> Result<(), AgentError> {
64        let mut cmd = Command::new(&self.config.command);
65        cmd.args(&self.config.args);
66        cmd.stdin(Stdio::piped());
67        cmd.stdout(Stdio::piped());
68        cmd.stderr(Stdio::piped());
69        for (k, v) in &self.config.env {
70            cmd.env(k, v);
71        }
72        if let Some(cwd) = &self.config.cwd {
73            cmd.current_dir(cwd);
74        }
75        hide_console_window(&mut cmd);
76
77        let mut child = cmd.spawn().map_err(|e| AgentError::ToolExecution {
78            tool: "McpClient".into(),
79            message: format!("failed to spawn MCP server '{}': {e}", self.config.command),
80        })?;
81
82        let stdin = child.stdin.take().ok_or_else(|| AgentError::ToolExecution {
83            tool: "McpClient".into(),
84            message: "failed to open child stdin".into(),
85        })?;
86        let stdout = child.stdout.take().ok_or_else(|| AgentError::ToolExecution {
87            tool: "McpClient".into(),
88            message: "failed to open child stdout".into(),
89        })?;
90
91        *self.stdin.lock().unwrap() = Some(stdin);
92        *self.reader.lock().unwrap() = Some(BufReader::new(stdout));
93        *self.process.lock().unwrap() = Some(child);
94
95        // MCP initialize handshake
96        let init_response = self
97            .send_request(
98                "initialize",
99                json!({
100                    "protocolVersion": "2024-11-05",
101                    "capabilities": {},
102                    "clientInfo": { "name": "telos-agent", "version": "0.1.0" },
103                }),
104            )
105            .await?;
106
107        self.server_info.lock().unwrap().replace(init_response);
108
109        // Send initialized notification
110        self.send_notification("notifications/initialized", json!({})).await?;
111
112        // Fetch tools
113        let tools_response = self.send_request("tools/list", json!({})).await?;
114        *self.tools.lock().unwrap() = Self::parse_tools(&tools_response);
115
116        Ok(())
117    }
118
119    /// Return the cached tool list.
120    pub fn tools(&self) -> Vec<McpTool> {
121        self.tools.lock().unwrap().clone()
122    }
123
124    /// Call an MCP tool by name with the given arguments.
125    ///
126    /// Returns the `result` portion of the JSON-RPC response on success.
127    pub async fn call_tool(&self, name: &str, arguments: Value) -> Result<Value, AgentError> {
128        self.send_request("tools/call", json!({ "name": name, "arguments": arguments })).await
129    }
130
131    /// Kill the server process and release I/O handles.
132    pub fn disconnect(&self) {
133        if let Some(mut child) = self.process.lock().unwrap().take() {
134            let _ = child.kill();
135            let _ = child.wait();
136        }
137        *self.stdin.lock().unwrap() = None;
138        *self.reader.lock().unwrap() = None;
139    }
140
141    // ── internal helpers ────────────────────────────────────────────
142
143    /// Send a JSON-RPC request and read the matching response.
144    async fn send_request(&self, method: &str, params: Value) -> Result<Value, AgentError> {
145        let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
146        let request = json!({
147            "jsonrpc": "2.0",
148            "id": id,
149            "method": method,
150            "params": params,
151        });
152        let request_str = serde_json::to_string(&request).unwrap();
153
154        // Write request to child stdin
155        if let Some(ref mut stdin) = *self.stdin.lock().unwrap() {
156            writeln!(stdin, "{request_str}").map_err(|e| AgentError::ToolExecution {
157                tool: "McpClient".into(),
158                message: format!("write to MCP server stdin failed: {e}"),
159            })?;
160            stdin.flush().ok();
161        } else {
162            return Err(AgentError::ToolExecution {
163                tool: "McpClient".into(),
164                message: "MCP client not connected (no stdin)".into(),
165            });
166        }
167
168        // Read one response line from child stdout
169        if let Some(ref mut reader) = *self.reader.lock().unwrap() {
170            let mut line = String::new();
171            reader.read_line(&mut line).map_err(|e| AgentError::ToolExecution {
172                tool: "McpClient".into(),
173                message: format!("read from MCP server stdout failed: {e}"),
174            })?;
175
176            if line.trim().is_empty() {
177                return Err(AgentError::ToolExecution {
178                    tool: "McpClient".into(),
179                    message: "MCP server returned empty response".into(),
180                });
181            }
182
183            let response: Value =
184                serde_json::from_str(&line).map_err(|e| AgentError::ToolExecution {
185                    tool: "McpClient".into(),
186                    message: format!("parse MCP server response failed: {e}"),
187                })?;
188
189            if let Some(err) = response.get("error") {
190                return Err(AgentError::ToolExecution {
191                    tool: "McpClient".into(),
192                    message: format!(
193                        "MCP error (code {}): {}",
194                        err.get("code").and_then(|v| v.as_i64()).unwrap_or(0),
195                        err.get("message").and_then(|v| v.as_str()).unwrap_or("unknown"),
196                    ),
197                });
198            }
199
200            Ok(response.get("result").cloned().unwrap_or(Value::Null))
201        } else {
202            Err(AgentError::ToolExecution {
203                tool: "McpClient".into(),
204                message: "MCP client not connected (no reader)".into(),
205            })
206        }
207    }
208
209    /// Send a JSON-RPC notification (no `id` field, no response expected).
210    async fn send_notification(&self, method: &str, params: Value) -> Result<(), AgentError> {
211        let notification = json!({
212            "jsonrpc": "2.0",
213            "method": method,
214            "params": params,
215        });
216        let notif_str = serde_json::to_string(&notification).unwrap();
217
218        if let Some(ref mut stdin) = *self.stdin.lock().unwrap() {
219            writeln!(stdin, "{notif_str}").ok();
220            stdin.flush().ok();
221        }
222        Ok(())
223    }
224
225    /// Parse a `tools/list` response into a `Vec<McpTool>`.
226    fn parse_tools(response: &Value) -> Vec<McpTool> {
227        response
228            .get("tools")
229            .and_then(|v| v.as_array())
230            .map(|arr| {
231                arr.iter()
232                    .map(|t| McpTool {
233                        name: t.get("name").and_then(|v| v.as_str()).unwrap_or("").to_string(),
234                        description: t
235                            .get("description")
236                            .and_then(|v| v.as_str())
237                            .unwrap_or("")
238                            .to_string(),
239                        input_schema: t
240                            .get("inputSchema")
241                            .cloned()
242                            .unwrap_or(json!({"type": "object"})),
243                    })
244                    .collect()
245            })
246            .unwrap_or_default()
247    }
248}
249
250#[cfg_attr(not(windows), allow(unused_variables))]
251fn hide_console_window(command: &mut Command) {
252    #[cfg(windows)]
253    {
254        use std::os::windows::process::CommandExt;
255
256        const CREATE_NO_WINDOW: u32 = 0x0800_0000;
257        command.creation_flags(CREATE_NO_WINDOW);
258    }
259}
260
261#[cfg(all(test, unix))]
262mod tests {
263    use super::*;
264
265    /// Write a minimal MCP echo-script to a temp file and verify the client
266    /// can connect, list tools, and call a tool.
267    #[test]
268    fn mcp_client_connect_and_list_tools() {
269        // Create a temporary shell script that acts as a MCP server.
270        let dir = std::env::temp_dir().join(format!("mcp-test-{}", std::process::id()));
271        let _ = std::fs::remove_dir_all(&dir);
272        std::fs::create_dir_all(&dir).unwrap();
273
274        let script_path = dir.join("echo-mcp.sh");
275        let script_content = r#"#!/bin/bash
276# Minimal MCP echo server — reads one request line, emits two responses.
277read REQUEST_LINE
278echo '{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2024-11-05","serverInfo":{"name":"echo","version":"1.0"},"capabilities":{"tools":{}}}}'
279read REQUEST_LINE
280echo '{"jsonrpc":"2.0","id":2,"result":{"tools":[{"name":"echo","description":"Echo tool","inputSchema":{"type":"object","properties":{"text":{"type":"string"}},"required":["text"]}}]}}'
281# Keep reading and echoing tool calls
282while read REQUEST_LINE; do
283  echo '{"jsonrpc":"2.0","id":3,"result":{"content":[{"type":"text","text":"pong"}]}}'
284done
285"#;
286        std::fs::write(&script_path, script_content).unwrap();
287
288        use std::os::unix::fs::PermissionsExt;
289        std::fs::set_permissions(&script_path, std::fs::Permissions::from_mode(0o755)).unwrap();
290
291        let rt = tokio::runtime::Runtime::new().unwrap();
292        rt.block_on(async {
293            let config = McpServerConfig::new(script_path.to_str().unwrap(), vec![]);
294            let client = McpClient::new(config);
295
296            // connect
297            client.connect().await.expect("MCP connect should succeed");
298
299            // list tools
300            let tools = client.tools();
301            assert_eq!(tools.len(), 1, "expected 1 tool");
302            assert_eq!(tools[0].name, "echo");
303            assert_eq!(tools[0].description, "Echo tool");
304
305            // call tool
306            let result = client
307                .call_tool("echo", json!({"text": "hello"}))
308                .await
309                .expect("call_tool should succeed");
310            assert_eq!(result["content"][0]["text"], "pong");
311
312            client.disconnect();
313        });
314
315        let _ = std::fs::remove_dir_all(&dir);
316    }
317}
telos_agent/integrations/mcp/client.rs

telos_agent/integrations/mcp/
client.rs
