oxicode_agent/mcp/transport/
stdio.rs1use 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
22const REQUEST_TIMEOUT_SECS: u64 = 30;
24
25const MAX_LINE_SIZE: usize = 10 * 1024 * 1024;
29
30const BLOCKED_ENV_VARS: &[&str] = &[
32 "LD_PRELOAD",
33 "LD_LIBRARY_PATH",
34 "DYLD_INSERT_LIBRARIES",
35 "DYLD_LIBRARY_PATH",
36];
37
38pub struct StdioTransport {
40 child: Option<Child>,
43 stdin: ChildStdin,
45 stdout: tokio::io::BufReader<ChildStdout>,
47 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 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 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 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 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 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 pub fn take_child(&mut self) -> Option<Child> {
159 self.child.take()
160 }
161
162 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 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 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 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 tracing::warn!(
240 "MCP: discarding response with non-matching id {} (expected {})",
241 mid,
242 id
243 );
244 continue;
245 }
246 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 #[cfg(unix)]
268 {
269 if let Some(mut child) = self.take_child()
270 && let Some(pid) = child.id()
271 {
272 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
301async 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}