Skip to main content

oxicode_agent/mcp/transport/
stdio.rs

1//! Stdio transport for MCP.
2//!
3//! Spawns a child process and communicates over its stdin/stdout using
4//! **newline-delimited JSON** (JSONL) per the MCP stdio transport spec
5//! (modelcontextprotocol.io/specification/2025-03-26/basic/transports).
6//! Each JSON-RPC message is serialized as one line of JSON terminated by
7//! a single `\n`; the transport rejects any line that exceeds `MAX_LINE_SIZE`.
8//!
9//! The transport is owned by a single [`crate::mcp::client::McpClient`] and
10//! is `&mut`-accessed exclusively by that client. The read loop in
11//! [`McpTransport::request`] dispatches inbound notifications and
12//! server→client requests to the installed [`InboundHandler`] inline.
13
14use super::{InboundHandler, McpTransport};
15use crate::mcp::spawn::SpawnValidator;
16use crate::mcp::types::RawJsonRpcMessage;
17use anyhow::{Context, Result};
18use std::process::Stdio;
19use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWriteExt};
20use tokio::process::{Child, ChildStdin, ChildStdout};
21
22/// Default timeout for individual MCP requests (seconds).
23const REQUEST_TIMEOUT_SECS: u64 = 30;
24
25/// Maximum allowed line size from an MCP server (10 MB).
26/// Guards against a buggy or hostile local server sending a runaway line
27/// that would otherwise cause unbounded allocation in the read buffer.
28const MAX_LINE_SIZE: usize = 10 * 1024 * 1024;
29
30/// Environment variables that servers must not override (security).
31const BLOCKED_ENV_VARS: &[&str] = &[
32    "LD_PRELOAD",
33    "LD_LIBRARY_PATH",
34    "DYLD_INSERT_LIBRARIES",
35    "DYLD_LIBRARY_PATH",
36];
37
38/// Stdio transport for a spawned MCP server process.
39pub struct StdioTransport {
40    /// Child process handle (kept alive to prevent process death).
41    /// `None` after `take_child` has been called.
42    child: Option<Child>,
43    /// Writer to the server's stdin.
44    stdin: ChildStdin,
45    /// Buffered reader from the server's stdout.
46    stdout: tokio::io::BufReader<ChildStdout>,
47    /// Inbound handler for notifications and server→client requests.
48    inbound_handler: Option<InboundHandler>,
49}
50
51impl std::fmt::Debug for StdioTransport {
52    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        f.debug_struct("StdioTransport")
54            .field("connected", &self.is_connected())
55            .finish()
56    }
57}
58
59impl StdioTransport {
60    /// Spawn a child process and return a connected transport.
61    ///
62    /// `validator`, when `Some`, is invoked BEFORE the spawn with
63    /// [`SpawnValidator::validate_command`]; an `Err` is converted to
64    /// `anyhow::Error` and short-circuits the spawn. After the SDK's
65    /// hardcoded `BLOCKED_ENV_VARS` filter is applied, the validator's
66    /// [`SpawnValidator::sanitize_env`] runs over the remaining env so
67    /// consumers can apply additional policy (path normalization, custom
68    /// blocked vars, etc.).
69    ///
70    /// `None` preserves the pre-validator behavior: only the SDK's
71    /// `BLOCKED_ENV_VARS` floor applies. Existing call sites pass `None`.
72    pub fn spawn(
73        command: &str,
74        args: &[String],
75        env: &std::collections::HashMap<String, String>,
76        cwd: Option<&str>,
77        debug: bool,
78        validator: Option<&dyn SpawnValidator>,
79    ) -> Result<Self> {
80        // Consumer-supplied validation runs first.
81        if let Some(v) = validator {
82            v.validate_command(command, args)
83                .map_err(|reason| anyhow::anyhow!("MCP spawn validation failed: {reason}"))?;
84        }
85
86        let mut cmd = tokio::process::Command::new(command);
87        cmd.args(args)
88            .stdin(Stdio::piped())
89            .stdout(Stdio::piped())
90            .kill_on_drop(true);
91
92        if debug {
93            cmd.stderr(Stdio::inherit());
94        } else {
95            cmd.stderr(Stdio::null());
96        }
97
98        // SDK hardcoded floor: loader-injection vectors.
99        let mut filtered_env: std::collections::HashMap<String, String> = env
100            .iter()
101            .filter(|(key, _)| {
102                let upper = key.to_uppercase();
103                let blocked = BLOCKED_ENV_VARS.iter().any(|b| upper == *b);
104                if blocked {
105                    tracing::warn!("MCP: blocked dangerous env override: {}", key);
106                }
107                !blocked
108            })
109            .map(|(k, v)| (k.clone(), v.clone()))
110            .collect();
111
112        // Consumer-supplied env scrub on top of the SDK floor.
113        if let Some(v) = validator {
114            v.sanitize_env(&mut filtered_env);
115        }
116
117        for (key, value) in &filtered_env {
118            cmd.env(key, value);
119        }
120
121        if let Some(dir) = cwd {
122            cmd.current_dir(dir);
123        }
124
125        let mut child = cmd
126            .spawn()
127            .with_context(|| format!("Failed to spawn MCP server: {}", command))?;
128
129        let stdin = child
130            .stdin
131            .take()
132            .context("Failed to acquire stdin from MCP server")?;
133        let stdout = child
134            .stdout
135            .take()
136            .context("Failed to acquire stdout from MCP server")?;
137
138        Ok(Self {
139            child: Some(child),
140            stdin,
141            stdout: tokio::io::BufReader::new(stdout),
142            inbound_handler: None,
143        })
144    }
145
146    /// Wrap an already-spawned child process.
147    /// Used internally by the client to allow later `take_child`.
148    pub fn from_parts(child: Child, stdin: ChildStdin, stdout: ChildStdout) -> Self {
149        Self {
150            child: Some(child),
151            stdin,
152            stdout: tokio::io::BufReader::new(stdout),
153            inbound_handler: None,
154        }
155    }
156
157    /// Take the child process out (for graceful shutdown via signal).
158    pub fn take_child(&mut self) -> Option<Child> {
159        self.child.take()
160    }
161
162    /// Write a single JSON-RPC message framed as one line of JSON + '\n'.
163    /// The MCP spec requires messages to be single-line; `serde_json`
164    /// never embeds raw newlines, so a trailing `\n` is sufficient.
165    async fn write_frame(&mut self, json: &str) -> Result<()> {
166        debug_assert!(
167            !json.contains('\n'),
168            "MCP 메시지에 내장 개행 금지 (스펙 위반)"
169        );
170        self.stdin
171            .write_all(json.as_bytes())
172            .await
173            .context("Failed to write MCP body")?;
174        self.stdin
175            .write_all(b"\n")
176            .await
177            .context("Failed to write MCP newline")?;
178        self.stdin
179            .flush()
180            .await
181            .context("Failed to flush MCP stdin")?;
182        Ok(())
183    }
184
185    /// Read one message (one line) from stdout, with a hard size cap.
186    /// Returns `Ok(None)` on clean EOF, `Err` on overflow or I/O failure.
187    async fn read_frame(&mut self) -> Result<Option<RawJsonRpcMessage>> {
188        match read_line_bounded(&mut self.stdout, MAX_LINE_SIZE).await? {
189            None => Ok(None),
190            Some(bytes) => {
191                let msg: RawJsonRpcMessage =
192                    serde_json::from_slice(&bytes).context("Failed to parse JSON-RPC message")?;
193                Ok(Some(msg))
194            }
195        }
196    }
197}
198
199#[async_trait::async_trait]
200impl McpTransport for StdioTransport {
201    async fn request(&mut self, id: u64, json: &str) -> Result<RawJsonRpcMessage> {
202        self.write_frame(json).await?;
203
204        // Read messages until the response with a matching id arrives.
205        // Anything else (notifications, server→client requests) is
206        // dispatched to the inbound handler inline.
207        let timeout = std::time::Duration::from_secs(REQUEST_TIMEOUT_SECS);
208        tokio::time::timeout(timeout, async {
209            loop {
210                let msg = self
211                    .read_frame()
212                    .await
213                    .context("Failed to read MCP response")?
214                    .ok_or_else(|| anyhow::anyhow!("MCP server closed connection"))?;
215
216                let msg_id = msg.id;
217                if let Some(mid) = msg_id {
218                    if mid == id {
219                        return Ok(msg);
220                    }
221                    // Non-matching id — if it has a `method` it's a
222                    // server→client request; dispatch and maybe reply.
223                    if msg.method.is_some() {
224                        let response = match self.inbound_handler.as_mut() {
225                            Some(h) => h(msg),
226                            None => None,
227                        };
228                        if let Some(value) = response {
229                            let reply = serde_json::to_string(&value)
230                                .context("Failed to serialize inbound response")?;
231                            self.write_frame(&reply)
232                                .await
233                                .context("Failed to write response to server→client request")?;
234                        }
235                        continue;
236                    }
237                    // Orphan response (different id, no method). Should not
238                    // happen in normal operation — log and skip.
239                    tracing::warn!(
240                        "MCP: discarding response with non-matching id {} (expected {})",
241                        mid,
242                        id
243                    );
244                    continue;
245                }
246                // No id → notification → dispatch (return value ignored).
247                if let Some(h) = self.inbound_handler.as_mut() {
248                    h(msg);
249                }
250            }
251        })
252        .await
253        .map_err(|_| anyhow::anyhow!("MCP request timed out after {}s", REQUEST_TIMEOUT_SECS))?
254    }
255
256    async fn notify(&mut self, json: &str) -> Result<()> {
257        self.write_frame(json).await
258    }
259
260    fn set_inbound_handler(&mut self, handler: InboundHandler) {
261        self.inbound_handler = Some(handler);
262    }
263
264    async fn close(&mut self) -> Result<()> {
265        let _ = self.stdin.shutdown().await;
266        // Graceful shutdown: SIGTERM then 5s then SIGKILL.
267        #[cfg(unix)]
268        {
269            if let Some(mut child) = self.take_child()
270                && let Some(pid) = child.id()
271            {
272                // SAFETY: libc::kill sends a signal to a process. The PID comes
273                // from child.id() which is a valid running process. SIGTERM
274                // requests graceful termination. On race (process already
275                // exited), kill returns ESRCH harmlessly.
276                unsafe {
277                    libc::kill(pid as libc::pid_t, libc::SIGTERM);
278                }
279                match tokio::time::timeout(std::time::Duration::from_secs(5), child.wait()).await {
280                    Ok(Ok(_)) => return Ok(()),
281                    _ => {
282                        let _ = child.kill().await;
283                    }
284                }
285            }
286        }
287        #[cfg(not(unix))]
288        {
289            if let Some(mut child) = self.take_child() {
290                let _ = child.kill().await;
291            }
292        }
293        Ok(())
294    }
295
296    fn is_connected(&self) -> bool {
297        self.child.is_some()
298    }
299}
300
301/// Bounded line reader over an [`AsyncBufRead`]. Reads until '\n' (included)
302/// or EOF. Caps the returned buffer at `max` bytes; returns an error if the
303/// line would exceed the cap, preventing unbounded allocation.
304async fn read_line_bounded<R: AsyncBufRead + Unpin>(
305    reader: &mut R,
306    max: usize,
307) -> Result<Option<Vec<u8>>> {
308    let mut buf: Vec<u8> = Vec::new();
309    loop {
310        let chunk = reader
311            .fill_buf()
312            .await
313            .context("Failed to read from MCP stdout")?;
314        if chunk.is_empty() {
315            return if buf.is_empty() {
316                Ok(None)
317            } else {
318                Err(anyhow::anyhow!("MCP server closed connection mid-line"))
319            };
320        }
321        if let Some(pos) = chunk.iter().position(|&b| b == b'\n') {
322            let take = pos + 1;
323            if buf.len() + take > max {
324                return Err(anyhow::anyhow!(
325                    "MCP line exceeds {} bytes (mid-line cap hit)",
326                    max
327                ));
328            }
329            buf.extend_from_slice(&chunk[..take]);
330            reader.consume(take);
331            return Ok(Some(buf));
332        }
333        if buf.len() + chunk.len() > max {
334            return Err(anyhow::anyhow!(
335                "MCP line exceeds {} bytes (chunk would overflow)",
336                max
337            ));
338        }
339        let n = chunk.len();
340        buf.extend_from_slice(chunk);
341        reader.consume(n);
342    }
343}