gemini_adk_rs/tools/mcp/
session_manager.rs

1//! MCP session management — connection params, tool discovery, and tool invocation.
2//!
3//! Implements a real MCP (Model Context Protocol) client speaking JSON-RPC 2.0.
4//! The primary transport is **stdio** (newline-delimited JSON over a subprocess's
5//! stdin/stdout), which works on default features. An optional **HTTP** transport
6//! (single-shot JSON-RPC POST) is available behind the `mcp-http` feature.
7
8use std::collections::HashMap;
9use std::sync::atomic::{AtomicU64, Ordering};
10use std::time::Duration;
11
12use serde_json::{Value, json};
13use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
14use tokio::process::{Child, ChildStdin, ChildStdout, Command};
15use tokio::sync::Mutex;
16
17/// MCP protocol version advertised during the handshake.
18const MCP_PROTOCOL_VERSION: &str = "2024-11-05";
19
20/// Connection parameters for an MCP server.
21#[derive(Debug, Clone)]
22pub enum McpConnectionParams {
23    /// Connect via stdio (subprocess).
24    Stdio {
25        /// The command to execute.
26        command: String,
27        /// Arguments passed to the command.
28        args: Vec<String>,
29        /// Connection timeout.
30        timeout: Option<Duration>,
31    },
32    /// Connect via SSE/StreamableHTTP.
33    Sse {
34        /// The URL of the MCP server.
35        url: String,
36        /// Optional HTTP headers for authentication.
37        headers: Option<HashMap<String, String>>,
38    },
39}
40
41/// Live stdio connection state: the child process plus framed I/O handles.
42struct StdioConnection {
43    /// The child process. Kept alive so the pipes stay open; killed on drop.
44    #[allow(dead_code)]
45    child: Child,
46    /// Subprocess stdin (we write requests here).
47    stdin: ChildStdin,
48    /// Buffered subprocess stdout (we read newline-delimited responses here).
49    stdout: BufReader<ChildStdout>,
50}
51
52/// Manages the MCP client session lifecycle.
53pub struct McpSessionManager {
54    params: McpConnectionParams,
55    /// Lazily-established stdio connection (None until first use, then reused).
56    stdio: Mutex<Option<StdioConnection>>,
57    /// Monotonic JSON-RPC request id counter.
58    next_id: AtomicU64,
59}
60
61impl McpSessionManager {
62    /// Create a new MCP session manager with the given connection params.
63    pub fn new(params: McpConnectionParams) -> Self {
64        Self {
65            params,
66            stdio: Mutex::new(None),
67            next_id: AtomicU64::new(1),
68        }
69    }
70
71    /// Get the connection parameters.
72    pub fn params(&self) -> &McpConnectionParams {
73        &self.params
74    }
75
76    fn next_id(&self) -> u64 {
77        self.next_id.fetch_add(1, Ordering::Relaxed)
78    }
79
80    /// List available tools from the MCP server.
81    ///
82    /// Connects (and performs the MCP handshake) lazily on first use, then issues
83    /// a `tools/list` JSON-RPC request and maps the result into [`McpToolInfo`]s.
84    pub async fn list_tools(&self) -> Result<Vec<McpToolInfo>, McpError> {
85        match &self.params {
86            McpConnectionParams::Stdio { .. } => self.stdio_list_tools().await,
87            #[cfg(feature = "mcp-http")]
88            McpConnectionParams::Sse { .. } => self.http_list_tools().await,
89            #[cfg(not(feature = "mcp-http"))]
90            McpConnectionParams::Sse { .. } => Err(McpError::ConnectionFailed(
91                "mcp-http feature not enabled".to_string(),
92            )),
93        }
94    }
95
96    /// Call a tool on the MCP server via `tools/call`.
97    ///
98    /// Returns the JSON-RPC `result` object on success. Returns
99    /// [`McpError::ToolCallFailed`] on a JSON-RPC error or when the result has
100    /// `isError: true`.
101    pub async fn call_tool(&self, name: &str, args: Value) -> Result<Value, McpError> {
102        match &self.params {
103            McpConnectionParams::Stdio { .. } => self.stdio_call_tool(name, args).await,
104            #[cfg(feature = "mcp-http")]
105            McpConnectionParams::Sse { .. } => self.http_call_tool(name, args).await,
106            #[cfg(not(feature = "mcp-http"))]
107            McpConnectionParams::Sse { .. } => Err(McpError::ConnectionFailed(
108                "mcp-http feature not enabled".to_string(),
109            )),
110        }
111    }
112
113    // ------------------------------------------------------------------
114    // stdio transport
115    // ------------------------------------------------------------------
116
117    async fn stdio_list_tools(&self) -> Result<Vec<McpToolInfo>, McpError> {
118        let timeout = self.stdio_timeout();
119        let mut guard = self.stdio.lock().await;
120        self.ensure_connected(&mut guard, timeout).await?;
121        let conn = guard.as_mut().expect("connection established above");
122
123        let id = self.next_id();
124        let req = json!({
125            "jsonrpc": "2.0",
126            "id": id,
127            "method": "tools/list",
128            "params": {},
129        });
130        let result = stdio_request(conn, id, &req, timeout).await?;
131        parse_tools_list(&result)
132    }
133
134    async fn stdio_call_tool(&self, name: &str, args: Value) -> Result<Value, McpError> {
135        let timeout = self.stdio_timeout();
136        let mut guard = self.stdio.lock().await;
137        self.ensure_connected(&mut guard, timeout).await?;
138        let conn = guard.as_mut().expect("connection established above");
139
140        let id = self.next_id();
141        let arguments = if args.is_null() { json!({}) } else { args };
142        let req = json!({
143            "jsonrpc": "2.0",
144            "id": id,
145            "method": "tools/call",
146            "params": { "name": name, "arguments": arguments },
147        });
148        let result = stdio_request(conn, id, &req, timeout).await?;
149        check_tool_result(&result, name)
150    }
151
152    fn stdio_timeout(&self) -> Option<Duration> {
153        match &self.params {
154            McpConnectionParams::Stdio { timeout, .. } => *timeout,
155            _ => None,
156        }
157    }
158
159    /// Ensure a stdio connection exists and has completed the MCP handshake.
160    async fn ensure_connected(
161        &self,
162        guard: &mut Option<StdioConnection>,
163        timeout: Option<Duration>,
164    ) -> Result<(), McpError> {
165        if guard.is_some() {
166            return Ok(());
167        }
168        let (command, args) = match &self.params {
169            McpConnectionParams::Stdio { command, args, .. } => (command.clone(), args.clone()),
170            _ => {
171                return Err(McpError::ConnectionFailed(
172                    "stdio transport requested for non-stdio params".to_string(),
173                ));
174            }
175        };
176
177        let mut child = Command::new(&command)
178            .args(&args)
179            .stdin(std::process::Stdio::piped())
180            .stdout(std::process::Stdio::piped())
181            .stderr(std::process::Stdio::null())
182            .spawn()
183            .map_err(|e| {
184                McpError::ConnectionFailed(format!("failed to spawn MCP server '{command}': {e}"))
185            })?;
186
187        let stdin = child.stdin.take().ok_or_else(|| {
188            McpError::ConnectionFailed("MCP server stdin not available".to_string())
189        })?;
190        let stdout = child.stdout.take().ok_or_else(|| {
191            McpError::ConnectionFailed("MCP server stdout not available".to_string())
192        })?;
193
194        let mut conn = StdioConnection {
195            child,
196            stdin,
197            stdout: BufReader::new(stdout),
198        };
199
200        // --- Handshake: initialize -> read result -> notifications/initialized ---
201        let id = self.next_id();
202        let init = json!({
203            "jsonrpc": "2.0",
204            "id": id,
205            "method": "initialize",
206            "params": {
207                "protocolVersion": MCP_PROTOCOL_VERSION,
208                "capabilities": {},
209                "clientInfo": { "name": "gemini-adk-rs", "version": "0.6.0" },
210            },
211        });
212        stdio_request(&mut conn, id, &init, timeout)
213            .await
214            .map_err(|e| McpError::ConnectionFailed(format!("MCP initialize failed: {e}")))?;
215
216        let initialized = json!({
217            "jsonrpc": "2.0",
218            "method": "notifications/initialized",
219        });
220        stdio_write(&mut conn, &initialized).await.map_err(|e| {
221            McpError::ConnectionFailed(format!("MCP initialized notify failed: {e}"))
222        })?;
223
224        *guard = Some(conn);
225        Ok(())
226    }
227
228    // ------------------------------------------------------------------
229    // HTTP transport (feature-gated)
230    // ------------------------------------------------------------------
231
232    #[cfg(feature = "mcp-http")]
233    async fn http_list_tools(&self) -> Result<Vec<McpToolInfo>, McpError> {
234        let id = self.next_id();
235        let req = json!({
236            "jsonrpc": "2.0",
237            "id": id,
238            "method": "tools/list",
239            "params": {},
240        });
241        let result = self.http_request(id, &req).await?;
242        parse_tools_list(&result)
243    }
244
245    #[cfg(feature = "mcp-http")]
246    async fn http_call_tool(&self, name: &str, args: Value) -> Result<Value, McpError> {
247        let id = self.next_id();
248        let arguments = if args.is_null() { json!({}) } else { args };
249        let req = json!({
250            "jsonrpc": "2.0",
251            "id": id,
252            "method": "tools/call",
253            "params": { "name": name, "arguments": arguments },
254        });
255        let result = self.http_request(id, &req).await?;
256        check_tool_result(&result, name)
257    }
258
259    #[cfg(feature = "mcp-http")]
260    #[cfg(not(feature = "mcp-http"))]
261    async fn http_request(&self, _id: u64, _req: &Value) -> Result<Value, McpError> {
262        Err(McpError::ConnectionFailed(
263            "SSE/HTTP MCP transport requires the `mcp-http` feature".to_string(),
264        ))
265    }
266
267    #[cfg(feature = "mcp-http")]
268    async fn http_request(&self, id: u64, req: &Value) -> Result<Value, McpError> {
269        let (url, headers) = match &self.params {
270            McpConnectionParams::Sse { url, headers } => (url.clone(), headers.clone()),
271            _ => {
272                return Err(McpError::ConnectionFailed(
273                    "HTTP transport requested for non-SSE params".to_string(),
274                ));
275            }
276        };
277
278        let client = reqwest::Client::new();
279        let mut builder = client
280            .post(&url)
281            .header("content-type", "application/json")
282            .header("accept", "application/json")
283            .json(req);
284        if let Some(hdrs) = headers {
285            for (k, v) in hdrs {
286                builder = builder.header(k, v);
287            }
288        }
289
290        let resp = builder
291            .send()
292            .await
293            .map_err(|e| McpError::ConnectionFailed(format!("MCP HTTP request failed: {e}")))?;
294        if !resp.status().is_success() {
295            return Err(McpError::ConnectionFailed(format!(
296                "MCP HTTP request returned status {}",
297                resp.status()
298            )));
299        }
300        let body: Value = resp
301            .json()
302            .await
303            .map_err(|e| McpError::Other(format!("invalid MCP HTTP response body: {e}")))?;
304        extract_result(&body, id)
305    }
306}
307
308// ----------------------------------------------------------------------
309// stdio framing helpers
310// ----------------------------------------------------------------------
311
312/// Write a single JSON-RPC message as one compact, newline-terminated line.
313async fn stdio_write(conn: &mut StdioConnection, msg: &Value) -> Result<(), McpError> {
314    let mut line = serde_json::to_string(msg)
315        .map_err(|e| McpError::Other(format!("failed to serialize JSON-RPC: {e}")))?;
316    line.push('\n');
317    conn.stdin
318        .write_all(line.as_bytes())
319        .await
320        .map_err(|e| McpError::ConnectionFailed(format!("failed to write to MCP server: {e}")))?;
321    conn.stdin.flush().await.map_err(|e| {
322        McpError::ConnectionFailed(format!("failed to flush MCP server stdin: {e}"))
323    })?;
324    Ok(())
325}
326
327/// Send a JSON-RPC request and read the matching response by `id`, skipping
328/// notifications and unrelated messages. Honors the optional timeout.
329async fn stdio_request(
330    conn: &mut StdioConnection,
331    id: u64,
332    req: &Value,
333    timeout: Option<Duration>,
334) -> Result<Value, McpError> {
335    let fut = async {
336        stdio_write(conn, req).await?;
337        loop {
338            let mut line = String::new();
339            let n = conn.stdout.read_line(&mut line).await.map_err(|e| {
340                McpError::ConnectionFailed(format!("failed to read from MCP server: {e}"))
341            })?;
342            if n == 0 {
343                // EOF: child closed stdout / exited before responding.
344                return Err(McpError::ConnectionFailed(
345                    "MCP server closed connection before responding".to_string(),
346                ));
347            }
348            let trimmed = line.trim();
349            if trimmed.is_empty() {
350                continue;
351            }
352            let msg: Value = match serde_json::from_str(trimmed) {
353                Ok(v) => v,
354                // Non-JSON line (stray log output) — skip it.
355                Err(_) => continue,
356            };
357            // Skip anything that isn't the response to our request id.
358            match msg.get("id").and_then(value_id_as_u64) {
359                Some(resp_id) if resp_id == id => return extract_result(&msg, id),
360                _ => continue,
361            }
362        }
363    };
364
365    match timeout {
366        Some(dur) => match tokio::time::timeout(dur, fut).await {
367            Ok(res) => res,
368            Err(_) => Err(McpError::ConnectionFailed(format!(
369                "MCP request timed out after {dur:?}"
370            ))),
371        },
372        None => fut.await,
373    }
374}
375
376// ----------------------------------------------------------------------
377// JSON-RPC / MCP result parsing (shared by both transports)
378// ----------------------------------------------------------------------
379
380/// Interpret a JSON-RPC `id` value (number or numeric string) as u64.
381fn value_id_as_u64(v: &Value) -> Option<u64> {
382    if let Some(n) = v.as_u64() {
383        return Some(n);
384    }
385    v.as_str().and_then(|s| s.parse::<u64>().ok())
386}
387
388/// Extract the `result` from a JSON-RPC response, mapping `error` to [`McpError`].
389fn extract_result(msg: &Value, id: u64) -> Result<Value, McpError> {
390    if let Some(err) = msg.get("error") {
391        let code = err
392            .get("code")
393            .and_then(serde_json::Value::as_i64)
394            .unwrap_or(0);
395        let message = err
396            .get("message")
397            .and_then(|m| m.as_str())
398            .unwrap_or("unknown error");
399        return Err(McpError::ToolCallFailed(format!(
400            "JSON-RPC error {code}: {message}"
401        )));
402    }
403    match msg.get("result") {
404        Some(result) => Ok(result.clone()),
405        None => Err(McpError::Other(format!(
406            "JSON-RPC response (id {id}) has neither result nor error"
407        ))),
408    }
409}
410
411/// Map a `tools/list` result into [`McpToolInfo`]s.
412fn parse_tools_list(result: &Value) -> Result<Vec<McpToolInfo>, McpError> {
413    let tools = result
414        .get("tools")
415        .and_then(|t| t.as_array())
416        .ok_or_else(|| McpError::Other("tools/list result missing 'tools' array".to_string()))?;
417
418    let mut out = Vec::with_capacity(tools.len());
419    for t in tools {
420        let name = t
421            .get("name")
422            .and_then(|n| n.as_str())
423            .ok_or_else(|| McpError::Other("tool entry missing 'name'".to_string()))?
424            .to_string();
425        let description = t
426            .get("description")
427            .and_then(|d| d.as_str())
428            .unwrap_or("")
429            .to_string();
430        let input_schema = t
431            .get("inputSchema")
432            .cloned()
433            .unwrap_or_else(|| json!({"type": "object"}));
434        out.push(McpToolInfo {
435            name,
436            description,
437            input_schema,
438        });
439    }
440    Ok(out)
441}
442
443/// Validate a `tools/call` result, surfacing `isError: true` as a failure.
444fn check_tool_result(result: &Value, name: &str) -> Result<Value, McpError> {
445    if result
446        .get("isError")
447        .and_then(serde_json::Value::as_bool)
448        .unwrap_or(false)
449    {
450        return Err(McpError::ToolCallFailed(format!(
451            "tool '{name}' reported isError: {result}"
452        )));
453    }
454    Ok(result.clone())
455}
456
457/// Information about an MCP tool.
458#[derive(Debug, Clone)]
459pub struct McpToolInfo {
460    /// Tool name.
461    pub name: String,
462    /// Human-readable tool description.
463    pub description: String,
464    /// JSON Schema for the tool's input parameters.
465    pub input_schema: serde_json::Value,
466}
467
468/// MCP-related errors.
469#[derive(Debug, thiserror::Error)]
470pub enum McpError {
471    /// Failed to connect to the MCP server.
472    #[error("Connection failed: {0}")]
473    ConnectionFailed(String),
474    /// The MCP session is not connected.
475    #[error("Not connected: {0}")]
476    NotConnected(String),
477    /// A tool call to the MCP server failed.
478    #[error("Tool call failed: {0}")]
479    ToolCallFailed(String),
480    /// A catch-all for other MCP errors.
481    #[error("{0}")]
482    Other(String),
483}
gemini_adk_rs/tools/mcp/session_manager.rs

gemini_adk_rs/tools/mcp/
session_manager.rs
