| 1 | use std::collections::VecDeque; |
| 2 | use std::sync::Arc; |
| 3 | use std::time::Duration; |
| 4 | |
| 5 | use anyhow::{Context, Result}; |
| 6 | use tokio::io::{AsyncBufReadExt, AsyncWriteExt}; |
| 7 | use tokio::process::{Child, ChildStdin, ChildStdout}; |
| 8 | use tokio::sync::Mutex as TokioMutex; |
| 9 | |
| 10 | use super::{McpServerConfig, McpTransport}; |
| 11 | use crate::child_env; |
| 12 | |
| 13 | pub(super) struct StdioTransport { |
| 14 | pub(super) child: Arc<TokioMutex<Child>>, |
| 15 | pub(super) stdin: ChildStdin, |
| 16 | pub(super) reader: tokio::io::BufReader<ChildStdout>, |
| 17 | /// Tail of stderr lines from the spawned MCP server. A background task |
| 18 | /// drains the child's stderr into this buffer so a mid-run crash leaves |
| 19 | /// some context behind instead of `Stdio::null` swallowing it. |
| 20 | pub(super) stderr_tail: Arc<StderrTail>, |
| 21 | /// Plugin authority can change in another process while this child is |
| 22 | /// idle. The connection-level cancellation token therefore also owns a |
| 23 | /// process watcher instead of waiting for a later tool call to drop the |
| 24 | /// transport. |
| 25 | pub(super) authority_cancel_watch: Option<tokio::task::JoinHandle<()>>, |
| 26 | /// Holds reviewed executable/script handles for the process lifetime. |
| 27 | pub(super) _reviewed_launch: Option<super::ReviewedStdioLaunch>, |
| 28 | } |
| 29 | |
| 30 | /// How long `StdioTransport::shutdown` waits for the child to exit on SIGTERM |
| 31 | /// before `kill_on_drop` fires SIGKILL. Tuned short so a hung MCP server |
| 32 | /// can't stall TUI exit; well-behaved servers almost always exit within |
| 33 | /// a few hundred ms. |
| 34 | pub(super) const STDIO_SHUTDOWN_GRACE: Duration = Duration::from_millis(2_000); |
| 35 | |
| 36 | /// How many lines of MCP-server stderr to keep around for crash diagnostics. |
| 37 | /// Bounded so a chatty server can't grow this without limit; large enough to |
| 38 | /// catch typical Node/Python startup or panic output. |
| 39 | const STDERR_TAIL_CAPACITY: usize = 64; |
| 40 | |
| 41 | /// Bounded ring buffer for the most recent stderr lines from a spawned MCP |
| 42 | /// server. Used by `StdioTransport` to surface server-side context when the |
| 43 | /// transport read side fails (server crashed, exited early, etc). |
| 44 | #[derive(Default)] |
| 45 | pub(super) struct StderrTail { |
| 46 | lines: TokioMutex<VecDeque<String>>, |
| 47 | } |
| 48 | |
| 49 | impl StderrTail { |
| 50 | pub(super) fn new() -> Arc<Self> { |
| 51 | Arc::new(Self { |
| 52 | lines: TokioMutex::new(VecDeque::with_capacity(STDERR_TAIL_CAPACITY)), |
| 53 | }) |
| 54 | } |
| 55 | |
| 56 | pub(super) async fn push(&self, line: String) { |
| 57 | let mut buf = self.lines.lock().await; |
| 58 | if buf.len() >= STDERR_TAIL_CAPACITY { |
| 59 | buf.pop_front(); |
| 60 | } |
| 61 | buf.push_back(line); |
| 62 | } |
| 63 | |
| 64 | async fn snapshot(&self) -> Vec<String> { |
| 65 | self.lines.lock().await.iter().cloned().collect() |
| 66 | } |
| 67 | } |
| 68 | |
| 69 | impl StdioTransport { |
| 70 | pub(super) fn spawn( |
| 71 | server_name: &str, |
| 72 | command: &str, |
| 73 | config: &McpServerConfig, |
| 74 | cancel_token: tokio_util::sync::CancellationToken, |
| 75 | ) -> Result<Self> { |
| 76 | let reviewed_launch = if let Some(reviewed_plugin) = config.reviewed_plugin.as_ref() { |
| 77 | // This is deliberately the last trust check before constructing |
| 78 | // and spawning the lazy stdio child. It re-reads only the |
| 79 | // Codewhale-owned plugin bundle, never user MCP/provider config or |
| 80 | // credential files, and fails closed on any content/capability |
| 81 | // drift after pool construction. |
| 82 | Some(reviewed_plugin.prepare_stdio_launch( |
| 83 | server_name, |
| 84 | command, |
| 85 | &config.args, |
| 86 | config.cwd.as_deref(), |
| 87 | )?) |
| 88 | } else { |
| 89 | None |
| 90 | }; |
| 91 | let mut cmd = reviewed_launch.as_ref().map_or_else( |
| 92 | || { |
| 93 | let mut command_process = tokio::process::Command::new(command); |
| 94 | command_process.args(&config.args); |
| 95 | command_process |
| 96 | }, |
| 97 | |launch| { |
| 98 | let mut command_process = tokio::process::Command::new(&launch.command); |
| 99 | command_process.args(&launch.args); |
| 100 | command_process |
| 101 | }, |
| 102 | ); |
| 103 | crate::utils::suppress_tokio_console_window(&mut cmd); |
| 104 | cmd.stdin(std::process::Stdio::piped()) |
| 105 | .stdout(std::process::Stdio::piped()) |
| 106 | .stderr(std::process::Stdio::piped()) |
| 107 | .kill_on_drop(true); |
| 108 | let launch_cwd = reviewed_launch |
| 109 | .as_ref() |
| 110 | .and_then(|launch| launch.cwd.as_ref()) |
| 111 | .or(config.cwd.as_ref().filter(|_| reviewed_launch.is_none())); |
| 112 | if let Some(cwd) = launch_cwd { |
| 113 | cmd.current_dir(cwd); |
| 114 | } |
| 115 | #[cfg(unix)] |
| 116 | if let Some(cwd_fd) = reviewed_launch |
| 117 | .as_ref() |
| 118 | .and_then(|launch| launch.cwd_fd.as_ref()) |
| 119 | { |
| 120 | use std::os::fd::AsRawFd as _; |
| 121 | let fd = cwd_fd.as_raw_fd(); |
| 122 | // SAFETY: the closure calls only async-signal-safe `fchdir` on an |
| 123 | // inherited directory descriptor before exec. |
| 124 | unsafe { |
| 125 | cmd.pre_exec(move || { |
| 126 | if libc::fchdir(fd) == 0 { |
| 127 | Ok(()) |
| 128 | } else { |
| 129 | Err(std::io::Error::last_os_error()) |
| 130 | } |
| 131 | }); |
| 132 | } |
| 133 | } |
| 134 | |
| 135 | // Expand `${NAME}` placeholders so secret env values can be sourced |
| 136 | // from the process environment instead of being stored in cleartext |
| 137 | // in the MCP config. The child env is allowlist-sanitized below, so |
| 138 | // these vars would not otherwise be inherited by the child. |
| 139 | let expanded_env = super::expanded_mcp_stdio_env(config) |
| 140 | .with_context(|| format!("MCP server '{server_name}' env expansion failed"))?; |
| 141 | |
| 142 | // User-configured MCP keeps the compatibility-oriented Node/Python |
| 143 | // bootstrap allowlist (#1244). Reviewed plugins receive only the base |
| 144 | // secret-scrubbed child environment plus their explicitly reviewed |
| 145 | // mappings, so namespaces such as NPM_CONFIG_* are never inherited |
| 146 | // ambiently across the consent boundary. |
| 147 | if let Some(reviewed_plugin) = config.reviewed_plugin.as_ref() { |
| 148 | cmd.env_clear(); |
| 149 | for (key, value) in child_env::sanitized_plugin_mcp_env_from( |
| 150 | reviewed_plugin.host_environment.entries().iter().cloned(), |
| 151 | child_env::string_map_env(&expanded_env), |
| 152 | ) { |
| 153 | cmd.env(key, value); |
| 154 | } |
| 155 | } else { |
| 156 | child_env::apply_to_tokio_command_mcp( |
| 157 | &mut cmd, |
| 158 | child_env::string_map_env(&expanded_env), |
| 159 | ); |
| 160 | } |
| 161 | |
| 162 | let mut child = cmd.spawn().with_context(|| { |
| 163 | if config.reviewed_plugin.is_some() { |
| 164 | format!( |
| 165 | "MCP stdio spawn failed (transport=stdio server={server_name} reviewed-plugin argv_count={} env_count={})", |
| 166 | config.args.len(), |
| 167 | expanded_env.len(), |
| 168 | ) |
| 169 | } else { |
| 170 | let env_keys: Vec<&str> = expanded_env.keys().map(String::as_str).collect(); |
| 171 | format!( |
| 172 | "MCP stdio spawn failed (transport=stdio server={server_name} cmd={command:?} args={:?} env_keys={env_keys:?})", |
| 173 | config.args, |
| 174 | ) |
| 175 | } |
| 176 | })?; |
| 177 | |
| 178 | let stdin = child.stdin.take().context("Failed to get MCP stdin")?; |
| 179 | let stdout = child.stdout.take().context("Failed to get MCP stdout")?; |
| 180 | let stderr = child.stderr.take().context("Failed to get MCP stderr")?; |
| 181 | |
| 182 | // Drain stderr into a bounded ring buffer so a crash mid-run leaves |
| 183 | // diagnostic breadcrumbs instead of disappearing into `Stdio::null`. |
| 184 | // The task exits naturally when the child closes its stderr |
| 185 | // (kill_on_drop / exit / explicit shutdown). |
| 186 | let stderr_tail = StderrTail::new(); |
| 187 | { |
| 188 | let tail = Arc::clone(&stderr_tail); |
| 189 | // A reviewed plugin child receives environment-backed values that |
| 190 | // are intentionally absent from its manifest. Still drain its |
| 191 | // stderr to avoid blocking, but do not retain or surface arbitrary |
| 192 | // child output that could echo those credentials into a chat or |
| 193 | // persisted transcript. |
| 194 | let capture_lines = config.reviewed_plugin.is_none(); |
| 195 | tokio::spawn(async move { |
| 196 | let mut lines = tokio::io::BufReader::new(stderr).lines(); |
| 197 | while let Ok(Some(line)) = lines.next_line().await { |
| 198 | if capture_lines { |
| 199 | tail.push(line).await; |
| 200 | } |
| 201 | } |
| 202 | }); |
| 203 | } |
| 204 | |
| 205 | let child = Arc::new(TokioMutex::new(child)); |
| 206 | let authority_cancel_watch = config.reviewed_plugin.as_ref().map(|_| { |
| 207 | let watched_child = Arc::clone(&child); |
| 208 | tokio::spawn(async move { |
| 209 | cancel_token.cancelled().await; |
| 210 | terminate_child_for_authority_change(&watched_child).await; |
| 211 | }) |
| 212 | }); |
| 213 | |
| 214 | Ok(Self { |
| 215 | child, |
| 216 | stdin, |
| 217 | reader: tokio::io::BufReader::new(stdout), |
| 218 | stderr_tail, |
| 219 | authority_cancel_watch, |
| 220 | _reviewed_launch: reviewed_launch, |
| 221 | }) |
| 222 | } |
| 223 | } |
| 224 | |
| 225 | /// Format the captured stderr tail for inclusion in an error message. Empty |
| 226 | /// tails return `None` so the caller can fall back to its original message. |
| 227 | async fn format_stderr_context(tail: &StderrTail) -> Option<String> { |
| 228 | let lines = tail.snapshot().await; |
| 229 | if lines.is_empty() { |
| 230 | return None; |
| 231 | } |
| 232 | Some(format!( |
| 233 | "MCP server stderr (last {} line{}):\n{}", |
| 234 | lines.len(), |
| 235 | if lines.len() == 1 { "" } else { "s" }, |
| 236 | lines.join("\n"), |
| 237 | )) |
| 238 | } |
| 239 | |
| 240 | /// Best-effort SIGTERM. On Unix uses `libc::kill`; on Windows there's no |
| 241 | /// equivalent so we let `kill_on_drop` (TerminateProcess) handle it via the |
| 242 | /// subsequent Drop. Returns whether a signal was actually sent. |
| 243 | fn send_sigterm(child: &Child) -> bool { |
| 244 | #[cfg(unix)] |
| 245 | { |
| 246 | if let Some(pid) = child.id() { |
| 247 | // SAFETY: pid was just obtained from `child.id()`. `libc::kill` |
| 248 | // with `SIGTERM` is async-signal-safe and never observes invalid |
| 249 | // memory. Worst case (pid wrap / process already gone) returns |
| 250 | // ESRCH, which we deliberately ignore. |
| 251 | unsafe { |
| 252 | let _ = libc::kill(pid as i32, libc::SIGTERM); |
| 253 | } |
| 254 | return true; |
| 255 | } |
| 256 | false |
| 257 | } |
| 258 | #[cfg(not(unix))] |
| 259 | { |
| 260 | let _ = child; |
| 261 | false |
| 262 | } |
| 263 | } |
| 264 | |
| 265 | async fn terminate_child_for_authority_change(child: &Arc<TokioMutex<Child>>) { |
| 266 | let mut child = child.lock().await; |
| 267 | terminate_child(&mut child).await; |
| 268 | } |
| 269 | |
| 270 | async fn terminate_child(child: &mut Child) { |
| 271 | // Reap an already-exited child before resolving its PID. Until it is |
| 272 | // reaped, the OS cannot recycle that identity; after it is reaped there is |
| 273 | // nothing left to signal. This avoids a PID-only watcher ever targeting an |
| 274 | // unrelated process after rapid PID reuse. |
| 275 | if child.try_wait().is_ok_and(|status| status.is_some()) { |
| 276 | return; |
| 277 | } |
| 278 | |
| 279 | #[cfg(unix)] |
| 280 | send_sigterm(child); |
| 281 | |
| 282 | #[cfg(not(unix))] |
| 283 | let _ = child.start_kill(); |
| 284 | |
| 285 | match tokio::time::timeout(STDIO_SHUTDOWN_GRACE, child.wait()).await { |
| 286 | Ok(Ok(_)) => {} |
| 287 | Ok(Err(_)) | Err(_) => { |
| 288 | // SIGTERM is advisory. Revocation and explicit shutdown must not |
| 289 | // leave the reviewed child alive indefinitely. |
| 290 | let _ = child.start_kill(); |
| 291 | let _ = child.wait().await; |
| 292 | } |
| 293 | } |
| 294 | } |
| 295 | |
| 296 | #[async_trait::async_trait] |
| 297 | impl McpTransport for StdioTransport { |
| 298 | async fn send(&mut self, mut msg: Vec<u8>) -> Result<()> { |
| 299 | msg.push(b'\n'); |
| 300 | self.stdin.write_all(&msg).await?; |
| 301 | self.stdin.flush().await?; |
| 302 | Ok(()) |
| 303 | } |
| 304 | |
| 305 | async fn recv(&mut self) -> Result<Vec<u8>> { |
| 306 | let mut line_bytes: Vec<u8> = Vec::new(); |
| 307 | loop { |
| 308 | // Bounded read: a server emitting a newline-free multi-GB "line" |
| 309 | // must not OOM us (read_line is unbounded). |
| 310 | let bytes = match read_line_capped( |
| 311 | &mut self.reader, |
| 312 | &mut line_bytes, |
| 313 | super::MAX_MCP_RESPONSE_BYTES, |
| 314 | ) |
| 315 | .await |
| 316 | { |
| 317 | Ok(b) => b, |
| 318 | Err(err) => { |
| 319 | if let Some(stderr) = format_stderr_context(&self.stderr_tail).await { |
| 320 | anyhow::bail!("Stdio transport read error: {err}\n{stderr}"); |
| 321 | } |
| 322 | return Err(err.into()); |
| 323 | } |
| 324 | }; |
| 325 | if bytes == 0 { |
| 326 | if let Some(stderr) = format_stderr_context(&self.stderr_tail).await { |
| 327 | anyhow::bail!("Stdio transport closed\n{stderr}"); |
| 328 | } |
| 329 | anyhow::bail!("Stdio transport closed"); |
| 330 | } |
| 331 | |
| 332 | let line = String::from_utf8_lossy(&line_bytes); |
| 333 | let trimmed = line.trim(); |
| 334 | if trimmed.is_empty() { |
| 335 | continue; |
| 336 | } |
| 337 | |
| 338 | return Ok(trimmed.as_bytes().to_vec()); |
| 339 | } |
| 340 | } |
| 341 | |
| 342 | /// Send SIGTERM and wait up to `STDIO_SHUTDOWN_GRACE` for graceful exit, |
| 343 | /// then force termination and reap the child as the backstop. |
| 344 | async fn shutdown(&mut self) { |
| 345 | let mut child = self.child.lock().await; |
| 346 | terminate_child(&mut child).await; |
| 347 | } |
| 348 | } |
| 349 | |
| 350 | /// Drop fallback (#420): if `shutdown` was never called explicitly, still |
| 351 | /// fire SIGTERM before tokio's `kill_on_drop` sends SIGKILL. The two |
| 352 | /// signals arrive back-to-back so well-behaved servers at least see the |
| 353 | /// SIGTERM first; misbehaving ones get SIGKILL'd anyway. |
| 354 | impl Drop for StdioTransport { |
| 355 | fn drop(&mut self) { |
| 356 | if let Some(watch) = self.authority_cancel_watch.take() { |
| 357 | watch.abort(); |
| 358 | } |
| 359 | if let Ok(mut child) = self.child.try_lock() |
| 360 | && !child.try_wait().is_ok_and(|status| status.is_some()) |
| 361 | { |
| 362 | send_sigterm(&child); |
| 363 | } |
| 364 | } |
| 365 | } |
| 366 | |
| 367 | /// Read one newline-terminated line into `out` (cleared first), aborting if it |
| 368 | /// exceeds `max` bytes without a newline. Bounds an otherwise-unbounded |
| 369 | /// `read_line` so a misbehaving MCP server cannot OOM the client. Returns the |
| 370 | /// number of bytes accumulated; 0 means EOF. |
| 371 | async fn read_line_capped<R>( |
| 372 | reader: &mut R, |
| 373 | out: &mut Vec<u8>, |
| 374 | max: usize, |
| 375 | ) -> std::io::Result<usize> |
| 376 | where |
| 377 | R: tokio::io::AsyncBufRead + Unpin, |
| 378 | { |
| 379 | use tokio::io::AsyncBufReadExt; |
| 380 | out.clear(); |
| 381 | loop { |
| 382 | let (chunk, consumed, done) = { |
| 383 | let available = reader.fill_buf().await?; |
| 384 | if available.is_empty() { |
| 385 | (Vec::new(), 0usize, true) |
| 386 | } else if let Some(pos) = available.iter().position(|&b| b == b'\n') { |
| 387 | (available[..=pos].to_vec(), pos + 1, true) |
| 388 | } else { |
| 389 | (available.to_vec(), available.len(), false) |
| 390 | } |
| 391 | }; |
| 392 | if consumed > 0 { |
| 393 | reader.consume(consumed); |
| 394 | } |
| 395 | out.extend_from_slice(&chunk); |
| 396 | if done { |
| 397 | break; |
| 398 | } |
| 399 | if out.len() > max { |
| 400 | return Err(std::io::Error::new( |
| 401 | std::io::ErrorKind::InvalidData, |
| 402 | format!("MCP stdio line exceeded {max} bytes without a newline"), |
| 403 | )); |
| 404 | } |
| 405 | } |
| 406 | Ok(out.len()) |
| 407 | } |
| 408 | |
| 409 | #[cfg(test)] |
| 410 | mod read_cap_tests { |
| 411 | use super::read_line_capped; |
| 412 | |
| 413 | #[tokio::test] |
| 414 | async fn reads_a_line_and_reports_eof() { |
| 415 | let data = b"hello\nworld\n".to_vec(); |
| 416 | let mut reader = tokio::io::BufReader::new(std::io::Cursor::new(data)); |
| 417 | let mut out = Vec::new(); |
| 418 | assert_eq!( |
| 419 | read_line_capped(&mut reader, &mut out, 1024).await.unwrap(), |
| 420 | 6 |
| 421 | ); |
| 422 | assert_eq!(out, b"hello\n"); |
| 423 | assert_eq!( |
| 424 | read_line_capped(&mut reader, &mut out, 1024).await.unwrap(), |
| 425 | 6 |
| 426 | ); |
| 427 | assert_eq!(out, b"world\n"); |
| 428 | // EOF. |
| 429 | assert_eq!( |
| 430 | read_line_capped(&mut reader, &mut out, 1024).await.unwrap(), |
| 431 | 0 |
| 432 | ); |
| 433 | } |
| 434 | |
| 435 | #[tokio::test] |
| 436 | async fn aborts_on_newline_free_line_over_cap() { |
| 437 | let data = vec![b'x'; 4096]; // no newline |
| 438 | let mut reader = tokio::io::BufReader::new(std::io::Cursor::new(data)); |
| 439 | let mut out = Vec::new(); |
| 440 | let err = read_line_capped(&mut reader, &mut out, 1024) |
| 441 | .await |
| 442 | .unwrap_err(); |
| 443 | assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); |
| 444 | } |
| 445 | } |
| 446 |