codei_mcp/
client.rs

1use std::sync::atomic::{AtomicU64, Ordering};
2
3use codei_config::McpServer;
4use serde::{Deserialize, Serialize};
5use serde_json::{json, Value};
6use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
7use tokio::process::{Child, ChildStdin, Command};
8use tokio::time::{timeout, Duration};
9use tracing::{debug, warn};
10
11use crate::error::McpError;
12
13const PROTOCOL_VERSION: &str = "2024-11-05";
14const REQUEST_TIMEOUT_SECS: u64 = 30;
15
16/// Metadata for a tool exposed by an MCP server.
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct McpToolInfo {
19    pub name: String,
20    pub description: String,
21    #[serde(rename = "inputSchema", default)]
22    pub input_schema: Value,
23}
24
25/// Result of `tools/call`.
26#[derive(Debug, Clone, Deserialize)]
27pub struct McpToolCallResult {
28    #[serde(default)]
29    pub content: Vec<McpContentBlock>,
30    #[serde(default)]
31    pub is_error: bool,
32}
33
34#[derive(Debug, Clone, Deserialize)]
35pub struct McpContentBlock {
36    #[serde(rename = "type", default)]
37    pub kind: String,
38    #[serde(default)]
39    pub text: String,
40}
41
42impl McpToolCallResult {
43    pub fn text(&self) -> String {
44        self.content
45            .iter()
46            .filter(|b| b.kind == "text" || b.kind.is_empty())
47            .map(|b| b.text.as_str())
48            .collect::<Vec<_>>()
49            .join("\n")
50    }
51}
52
53/// JSON-RPC client over stdio transport.
54pub struct McpClient {
55    server_name: String,
56    child: Child,
57    stdin: ChildStdin,
58    reader: BufReader<tokio::process::ChildStdout>,
59    next_id: AtomicU64,
60}
61
62impl McpClient {
63    pub async fn connect(server: &McpServer) -> Result<Self, McpError> {
64        let mut cmd = Command::new(&server.command);
65        cmd.args(&server.args);
66        for (key, value) in &server.env {
67            cmd.env(key, value);
68        }
69        cmd.stdin(std::process::Stdio::piped())
70            .stdout(std::process::Stdio::piped())
71            .stderr(std::process::Stdio::piped())
72            .kill_on_drop(true);
73
74        let mut child = cmd.spawn().map_err(|source| McpError::Spawn {
75            name: server.name.clone(),
76            source,
77        })?;
78
79        let stdin = child.stdin.take().ok_or_else(|| McpError::Protocol {
80            server: server.name.clone(),
81            message: "missing stdin".into(),
82        })?;
83        let stdout = child.stdout.take().ok_or_else(|| McpError::Protocol {
84            server: server.name.clone(),
85            message: "missing stdout".into(),
86        })?;
87
88        let mut client = Self {
89            server_name: server.name.clone(),
90            child,
91            stdin,
92            reader: BufReader::new(stdout),
93            next_id: AtomicU64::new(1),
94        };
95
96        client.initialize().await?;
97        Ok(client)
98    }
99
100    pub fn server_name(&self) -> &str {
101        &self.server_name
102    }
103
104    async fn initialize(&mut self) -> Result<(), McpError> {
105        let result = self
106            .request(
107                "initialize",
108                json!({
109                    "protocolVersion": PROTOCOL_VERSION,
110                    "capabilities": {},
111                    "clientInfo": {
112                        "name": "codei",
113                        "version": env!("CARGO_PKG_VERSION"),
114                    }
115                }),
116            )
117            .await?;
118
119        debug!(
120            server = %self.server_name,
121            result = %result,
122            "MCP initialize complete"
123        );
124
125        self.notify("notifications/initialized", json!({})).await?;
126        Ok(())
127    }
128
129    pub async fn list_tools(&mut self) -> Result<Vec<McpToolInfo>, McpError> {
130        let result = self.request("tools/list", json!({})).await?;
131        let tools = result
132            .get("tools")
133            .and_then(|v| v.as_array())
134            .cloned()
135            .unwrap_or_default();
136        let mut parsed = Vec::new();
137        for tool in tools {
138            let info: McpToolInfo =
139                serde_json::from_value(tool).map_err(|err| McpError::Protocol {
140                    server: self.server_name.clone(),
141                    message: format!("invalid tool definition: {err}"),
142                })?;
143            parsed.push(info);
144        }
145        Ok(parsed)
146    }
147
148    pub async fn call_tool(
149        &mut self,
150        name: &str,
151        arguments: Value,
152    ) -> Result<McpToolCallResult, McpError> {
153        let result = self
154            .request(
155                "tools/call",
156                json!({
157                    "name": name,
158                    "arguments": arguments,
159                }),
160            )
161            .await?;
162        serde_json::from_value(result).map_err(|err| McpError::Protocol {
163            server: self.server_name.clone(),
164            message: format!("invalid tools/call response: {err}"),
165        })
166    }
167
168    async fn notify(&mut self, method: &str, params: Value) -> Result<(), McpError> {
169        let message = json!({
170            "jsonrpc": "2.0",
171            "method": method,
172            "params": params,
173        });
174        self.write_message(&message).await
175    }
176
177    async fn request(&mut self, method: &str, params: Value) -> Result<Value, McpError> {
178        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
179        let message = json!({
180            "jsonrpc": "2.0",
181            "id": id,
182            "method": method,
183            "params": params,
184        });
185        self.write_message(&message).await?;
186
187        let response = timeout(
188            Duration::from_secs(REQUEST_TIMEOUT_SECS),
189            self.read_response(id),
190        )
191        .await
192        .map_err(|_| McpError::Timeout {
193            secs: REQUEST_TIMEOUT_SECS,
194        })??;
195
196        Ok(response)
197    }
198
199    async fn read_response(&mut self, id: u64) -> Result<Value, McpError> {
200        loop {
201            if self.child.try_wait()?.is_some() {
202                return Err(McpError::Exited {
203                    name: self.server_name.clone(),
204                });
205            }
206
207            let line = self.read_line().await?;
208            if line.trim().is_empty() {
209                continue;
210            }
211
212            let value: Value = serde_json::from_str(&line)?;
213            if value.get("method").is_some() && value.get("id").is_none() {
214                debug!(server = %self.server_name, %line, "MCP notification");
215                continue;
216            }
217
218            if value.get("id").and_then(|v| v.as_u64()) != Some(id) {
219                warn!(server = %self.server_name, %line, "unexpected MCP response id");
220                continue;
221            }
222
223            if let Some(error) = value.get("error") {
224                let message = error
225                    .get("message")
226                    .and_then(|v| v.as_str())
227                    .unwrap_or("unknown error");
228                return Err(McpError::Protocol {
229                    server: self.server_name.clone(),
230                    message: message.to_string(),
231                });
232            }
233
234            return Ok(value.get("result").cloned().unwrap_or(Value::Null));
235        }
236    }
237
238    async fn write_message(&mut self, message: &Value) -> Result<(), McpError> {
239        let line = serde_json::to_string(message)?;
240        debug!(server = %self.server_name, %line, "MCP send");
241        self.stdin.write_all(line.as_bytes()).await?;
242        self.stdin.write_all(b"\n").await?;
243        self.stdin.flush().await?;
244        Ok(())
245    }
246
247    async fn read_line(&mut self) -> Result<String, McpError> {
248        let mut line = String::new();
249        self.reader.read_line(&mut line).await?;
250        if line.is_empty() {
251            return Err(McpError::Exited {
252                name: self.server_name.clone(),
253            });
254        }
255        Ok(line)
256    }
257}
codei_mcp/client.rs

codei_mcp/
client.rs
