From 520a898bd80e734f8e183a57d4e2b60b4f0789ad Mon Sep 17 00:00:00 2001 From: Hare Date: Mon, 12 Oct 2026 02:35:42 +0900 Subject: [PATCH] fix: bound MCP stdio requests and shutdown across stdin waits --- crates/mcp/src/stdio.rs | 130 ++++++-- crates/mcp/src/stdio/deadline_tests.rs | 362 +++++++++++++++++++++++ crates/mcp/tests/fixtures/mock_server.rs | 19 +- crates/mcp/tests/stdio_lifecycle.rs | 45 +++ crates/worker/src/feature/mcp.rs | 76 ++++- 5 files changed, 602 insertions(+), 30 deletions(-) create mode 100644 crates/mcp/src/stdio/deadline_tests.rs diff --git a/crates/mcp/src/stdio.rs b/crates/mcp/src/stdio.rs index 7379d65e8..7465a39ef 100644 --- a/crates/mcp/src/stdio.rs +++ b/crates/mcp/src/stdio.rs @@ -11,11 +11,11 @@ use secrets::SecretStore; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; use thiserror::Error; -use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt, BufReader}; -use tokio::process::{Child, ChildStderr, ChildStdin, ChildStdout, Command}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader}; +use tokio::process::{Child, ChildStderr, ChildStdout, Command}; use tokio::sync::{Mutex, mpsc}; use tokio::task::JoinHandle; -use tokio::time::{Instant, timeout}; +use tokio::time::{Instant, timeout, timeout_at}; const MCP_PROTOCOL_VERSION: &str = "2025-11-25"; const CLIENT_NAME: &str = "yoi"; @@ -31,8 +31,11 @@ pub struct McpStdioLimits { pub max_diagnostic_lines: usize, pub max_protocol_bytes: usize, pub startup_timeout: Duration, + /// One budget for stdin ownership, frame write/flush, and response wait. pub request_timeout: Duration, + /// Shared budget for stdin ownership/close and graceful process exit. pub shutdown_timeout: Duration, + /// Budget for each fallback wait: terminate, then forced kill/reap. pub kill_timeout: Duration, } @@ -570,6 +573,30 @@ pub fn resolve_stdio_server( Ok(spec) } +// The mutex orders entire frames, including replies emitted by the stdout reader. +type SharedStdin = Arc>>>; + +/// Last boundary reached by a failed request. A flushed frame does not prove +/// that the server executed the method; callers must not blindly retry it. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum McpRequestStage { + WaitingForStdin, + WritingFrame, + FlushingFrame, + AwaitingResponse, +} + +impl fmt::Display for McpRequestStage { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::WaitingForStdin => "stdin ownership (method not sent)", + Self::WritingFrame => "frame write (delivery unknown; stdin closed)", + Self::FlushingFrame => "frame flush (delivery unknown; stdin closed)", + Self::AwaitingResponse => "response (frame flushed; execution unknown; do not retry)", + }) + } +} + /// A running initialized stdio MCP client. pub struct McpStdioClient { server_name: String, @@ -577,7 +604,7 @@ pub struct McpStdioClient { redactor: Redactor, diagnostics: Arc>, list_changes: Arc>, - stdin: Arc>>, + stdin: SharedStdin, child: Option, responses: mpsc::Receiver, reader_task: JoinHandle<()>, @@ -663,7 +690,7 @@ impl McpStdioClient { ) })?; - let stdin = Arc::new(Mutex::new(Some(stdin))); + let stdin: SharedStdin = Arc::new(Mutex::new(Some(Box::new(stdin)))); let diagnostics = Arc::new(Mutex::new(BoundedDiagnostics::new( spec.name.clone(), limits.max_diagnostic_lines, @@ -747,6 +774,12 @@ impl McpStdioClient { /// The caller is responsible for applying Yoi tool permissions before this /// method is reached and for bounding/serializing the untrusted result before /// it is exposed to model-visible tool history. + /// + /// Timeout or cancellation after writing begins can leave execution unknown; + /// this method never retries. Interrupted write/flush closes stdin, and the + /// caller must shut down the process rather than reuse that connection. A + /// fully flushed frame permits later *different* requests with distinct IDs, + /// but is not permission to replay this call when its response is missing. pub async fn call_tool( &mut self, request: CallToolRequest, @@ -902,23 +935,35 @@ impl McpStdioClient { method, params, }; - self.write_protocol(phase, &request).await?; - let response = match timeout( - self.limits.request_timeout, - self.wait_for_response(id, phase), - ) + if self.shutdown_started { + return Err(McpClientError::new( + &self.server_name, + phase, + McpErrorKind::Io("client shutdown has started".into()), + )); + } + let deadline = Instant::now() + self.limits.request_timeout; + let mut stage = McpRequestStage::WaitingForStdin; + // One budget covers ownership, write, flush, and the matching response. + let response = match timeout_at(deadline, async { + self.write_protocol(phase, &request, &mut stage).await?; + stage = McpRequestStage::AwaitingResponse; + self.wait_for_response(id, phase).await + }) .await { - Ok(result) => result?, + Ok(Ok(response)) => response, + Ok(Err(err)) => return Err(err.with_request_stage(stage)), Err(_) => { return Err(McpClientError::new( &self.server_name, phase, McpErrorKind::Timeout { - operation: method.to_owned(), + operation: self.redactor.redact(method), elapsed: self.limits.request_timeout, }, ) + .with_request_stage(stage) .with_diagnostics(self.snapshot_diagnostics().await)); } }; @@ -983,6 +1028,7 @@ impl McpStdioClient { method: &str, params: T, ) -> Result<(), McpClientError> { + let mut stage = McpRequestStage::WaitingForStdin; self.write_protocol( phase, &ClientNotification { @@ -990,6 +1036,7 @@ impl McpStdioClient { method, params, }, + &mut stage, ) .await } @@ -998,6 +1045,7 @@ impl McpStdioClient { &mut self, phase: McpPhase, value: &T, + stage: &mut McpRequestStage, ) -> Result<(), McpClientError> { write_json_line( &self.server_name, @@ -1006,6 +1054,7 @@ impl McpStdioClient { value, self.limits.max_protocol_bytes, &self.redactor, + stage, ) .await } @@ -1013,15 +1062,23 @@ impl McpStdioClient { /// Close stdin and wait for process exit, falling back to terminate and kill. pub async fn shutdown(&mut self) -> Result { self.shutdown_started = true; - { - let mut stdin = self.stdin.lock().await; - stdin.take(); - } + let deadline = Instant::now() + self.limits.shutdown_timeout; + // Ownership and EOF share one graceful budget. A reader reply may own + // stdin indefinitely; its lock is never required to terminate/kill. + let closed = timeout_at(deadline, async { + self.stdin.lock().await.take(); + }) + .await + .is_ok(); let mut child = match self.child.take() { Some(child) => child, None => return Ok(ShutdownReport::already_finished()), }; + if !closed { + self.reader_task.abort(); + return self.terminate_then_kill(child).await; + } if let Ok(Some(status)) = child.try_wait() { self.reader_task.abort(); @@ -1033,7 +1090,7 @@ impl McpStdioClient { }); } - match timeout(self.limits.shutdown_timeout, child.wait()).await { + match timeout_at(deadline, child.wait()).await { Ok(Ok(status)) => { self.reader_task.abort(); self.stderr_task.abort(); @@ -1195,12 +1252,13 @@ pub struct McpDiagnostics { } #[derive(Debug, Error, Clone)] -#[error("MCP stdio server `{server_name}` failed during phase `{phase}`: {kind}")] +#[error("MCP stdio server `{server_name}` failed during phase `{phase}`: {kind}{stage_detail}", stage_detail = request_stage.map(|stage| format!("; {stage}")).unwrap_or_default())] pub struct McpClientError { pub server_name: String, pub phase: McpPhase, pub kind: McpErrorKind, diagnostics: Option, + request_stage: Option, } impl McpClientError { @@ -1210,6 +1268,7 @@ impl McpClientError { phase, kind, diagnostics: None, + request_stage: None, } } @@ -1223,6 +1282,15 @@ impl McpClientError { self } + fn with_request_stage(mut self, stage: McpRequestStage) -> Self { + self.request_stage = Some(stage); + self + } + + pub fn request_stage(&self) -> Option { + self.request_stage + } + pub fn diagnostics(&self) -> Option<&McpDiagnostics> { self.diagnostics.as_ref() } @@ -1311,7 +1379,7 @@ struct ErrorObject<'a> { fn spawn_stdout_reader( server_name: String, stdout: ChildStdout, - stdin: Arc>>, + stdin: SharedStdin, tx: mpsc::Sender, limits: McpStdioLimits, redactor: Redactor, @@ -1368,7 +1436,7 @@ fn spawn_stdout_reader( async fn handle_incoming_message( server_name: &str, - stdin: &Arc>>, + stdin: &SharedStdin, tx: &mpsc::Sender, limits: &McpStdioLimits, redactor: &Redactor, @@ -1392,6 +1460,7 @@ async fn handle_incoming_message( &response, limits.max_protocol_bytes, redactor, + &mut McpRequestStage::WaitingForStdin, ) .await; } @@ -1484,10 +1553,11 @@ fn spawn_stderr_reader( async fn write_json_line( server_name: &str, phase: McpPhase, - stdin: &Arc>>, + stdin: &SharedStdin, value: &T, max_protocol_bytes: usize, redactor: &Redactor, + stage: &mut McpRequestStage, ) -> Result<(), McpClientError> { let mut bytes = serde_json::to_vec(value).map_err(|err| { McpClientError::new( @@ -1507,27 +1577,34 @@ async fn write_json_line( } bytes.push(b'\n'); let mut guard = stdin.lock().await; - let Some(stdin) = guard.as_mut() else { + let Some(mut writer) = guard.take() else { return Err(McpClientError::new( server_name, phase, McpErrorKind::Io("child stdin is closed".into()), )); }; - stdin.write_all(&bytes).await.map_err(|err| { + // Keep ownership out of the slot across both awaits. Timeout/cancel or + // I/O error drops the writer, closing the connection instead of allowing + // the next frame to append to an unverified partial JSON line. + *stage = McpRequestStage::WritingFrame; + writer.write_all(&bytes).await.map_err(|err| { McpClientError::new( server_name, phase, McpErrorKind::Io(redactor.redact(&err.to_string())), ) })?; - stdin.flush().await.map_err(|err| { + *stage = McpRequestStage::FlushingFrame; + writer.flush().await.map_err(|err| { McpClientError::new( server_name, phase, McpErrorKind::Io(redactor.redact(&err.to_string())), ) - }) + })?; + *guard = Some(writer); + Ok(()) } async fn read_protocol_line( @@ -1684,3 +1761,6 @@ fn send_terminate(child: &mut Child) -> Result<(), ()> { fn send_terminate(child: &mut Child) -> Result<(), ()> { child.start_kill().map_err(|_| ()) } + +#[cfg(test)] +mod deadline_tests; diff --git a/crates/mcp/src/stdio/deadline_tests.rs b/crates/mcp/src/stdio/deadline_tests.rs new file mode 100644 index 000000000..b0f3ed310 --- /dev/null +++ b/crates/mcp/src/stdio/deadline_tests.rs @@ -0,0 +1,362 @@ +//! Module-boundary fixtures exercise the same writer/request path as ChildStdin. +use super::*; +use std::future::{Future, poll_fn}; +use std::pin::Pin; +use std::sync::Mutex as StdMutex; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::task::{Context, Poll}; + +#[derive(Clone, Copy)] +enum Stall { + PartialWrite, + Flush, + Response, +} + +struct ControlledWriter { + stall: Stall, + bytes: Arc>>, + reached: Option>, + dropped: Arc, +} + +impl ControlledWriter { + fn signal(&mut self) { + if let Some(tx) = self.reached.take() { + tx.send(()).unwrap(); + } + } +} + +impl AsyncWrite for ControlledWriter { + fn poll_write( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + if matches!(self.stall, Stall::PartialWrite) { + if self.bytes.lock().unwrap().is_empty() { + self.bytes.lock().unwrap().extend_from_slice(&buf[..8]); + return Poll::Ready(Ok(8)); + } + self.signal(); + return Poll::Pending; + } + self.bytes.lock().unwrap().extend_from_slice(buf); + Poll::Ready(Ok(buf.len())) + } + + fn poll_flush(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + self.signal(); + if matches!(self.stall, Stall::Flush) { + Poll::Pending + } else { + Poll::Ready(Ok(())) + } + } + + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } +} + +impl Drop for ControlledWriter { + fn drop(&mut self) { + self.dropped.store(true, Ordering::SeqCst); + } +} + +fn fixture_client(stdin: SharedStdin) -> (McpStdioClient, mpsc::Sender) { + let (tx, responses) = mpsc::channel(16); + let redactor = Redactor::new(vec!["fixture-secret".into()]); + let client = McpStdioClient { + server_name: "deadline-fixture".into(), + limits: McpStdioLimits { + request_timeout: Duration::from_millis(100), + ..Default::default() + }, + diagnostics: Arc::new(Mutex::new(BoundedDiagnostics::new( + "deadline-fixture".into(), + 2, + redactor.clone(), + ))), + redactor, + list_changes: Arc::new(Mutex::new(BoundedListChanged::new( + "deadline-fixture".into(), + ))), + stdin, + child: None, + responses, + reader_task: tokio::spawn(std::future::pending()), + stderr_task: tokio::spawn(std::future::pending()), + next_id: 1, + initialized: None, + shutdown_started: false, + }; + (client, tx) +} + +async fn assert_pending(mut future: Pin<&mut F>) { + poll_fn(|cx| match future.as_mut().poll(cx) { + Poll::Pending => Poll::Ready(()), + Poll::Ready(_) => panic!("fixture boundary unexpectedly completed"), + }) + .await; +} + +#[tokio::test] +async fn stdin_ownership_timeout_and_cancel_do_not_send_a_method() { + for cancel in [false, true] { + let bytes = Arc::new(StdMutex::new(Vec::new())); + let dropped = Arc::new(AtomicBool::new(false)); + let stdin: SharedStdin = Arc::new(Mutex::new(Some(Box::new(ControlledWriter { + stall: Stall::Response, + bytes: bytes.clone(), + reached: None, + dropped: dropped.clone(), + })))); + let guard = stdin.lock().await; + let (mut client, _tx) = fixture_client(stdin.clone()); + let mut request = + Box::pin(client.request::(McpPhase::Running, "tools/call", json!({}))); + assert_pending(request.as_mut()).await; + if !cancel { + let err = timeout(Duration::from_secs(1), &mut request) + .await + .unwrap() + .unwrap_err(); + assert!(matches!(err.kind, McpErrorKind::Timeout { .. })); + assert_eq!(err.request_stage(), Some(McpRequestStage::WaitingForStdin)); + assert!(err.to_string().contains("method not sent")); + } + drop(request); + assert!(bytes.lock().unwrap().is_empty()); + assert!(!dropped.load(Ordering::SeqCst)); + assert!( + guard.is_some(), + "cancel before ownership must preserve the stream" + ); + } +} + +#[tokio::test] +async fn interrupted_write_or_flush_closes_stdin_before_the_next_frame() { + for (stall, stage) in [ + (Stall::PartialWrite, McpRequestStage::WritingFrame), + (Stall::Flush, McpRequestStage::FlushingFrame), + ] { + for cancel in [false, true] { + let bytes = Arc::new(StdMutex::new(Vec::new())); + let dropped = Arc::new(AtomicBool::new(false)); + let (reached, mut rx) = mpsc::unbounded_channel(); + let stdin: SharedStdin = Arc::new(Mutex::new(Some(Box::new(ControlledWriter { + stall, + bytes: bytes.clone(), + reached: Some(reached), + dropped: dropped.clone(), + })))); + let (mut client, _tx) = fixture_client(stdin.clone()); + let mut request = Box::pin(client.request::( + McpPhase::Running, + "tools/call", + json!({"secret": "fixture-secret"}), + )); + // Poll all the way to the controlled Pending, then observe its channel. + assert_pending(request.as_mut()).await; + timeout(Duration::from_secs(1), rx.recv()) + .await + .unwrap() + .unwrap(); + if !cancel { + let err = timeout(Duration::from_secs(1), &mut request) + .await + .unwrap() + .unwrap_err(); + assert!(matches!(err.kind, McpErrorKind::Timeout { .. })); + assert_eq!(err.request_stage(), Some(stage)); + let message = err.to_string(); + assert!(message.contains("delivery unknown; stdin closed")); + assert!(!message.contains("fixture-secret")); + } + drop(request); + assert!(dropped.load(Ordering::SeqCst)); + assert!(stdin.lock().await.is_none()); + let before = bytes.lock().unwrap().clone(); + if matches!(stall, Stall::PartialWrite) { + assert!(!before.contains(&b'\n')); + } else { + assert_eq!(before.last(), Some(&b'\n')); + } + let err = client + .request::(McpPhase::Running, "next-method", json!({})) + .await + .unwrap_err(); + assert!(err.to_string().contains("stdin is closed")); + assert_eq!( + *bytes.lock().unwrap(), + before, + "no next frame can append after interruption" + ); + } + } +} + +#[tokio::test] +async fn flushed_call_timeout_or_cancel_is_not_replayed_and_late_response_is_not_mixed() { + for cancel in [false, true] { + let bytes = Arc::new(StdMutex::new(Vec::new())); + let (reached, mut rx) = mpsc::unbounded_channel(); + let stdin: SharedStdin = Arc::new(Mutex::new(Some(Box::new(ControlledWriter { + stall: Stall::Response, + bytes: bytes.clone(), + reached: Some(reached), + dropped: Arc::new(AtomicBool::new(false)), + })))); + let (mut client, tx) = fixture_client(stdin); + let mut request = + Box::pin(client.request::(McpPhase::Running, "tools/call", json!({}))); + assert_pending(request.as_mut()).await; + timeout(Duration::from_secs(1), rx.recv()) + .await + .unwrap() + .unwrap(); + if !cancel { + let err = timeout(Duration::from_secs(1), &mut request) + .await + .unwrap() + .unwrap_err(); + assert_eq!(err.request_stage(), Some(McpRequestStage::AwaitingResponse)); + assert!(err.to_string().contains("execution unknown; do not retry")); + } + drop(request); + // Reuse is safe only because the first frame completed and IDs are + // matched. This is a different method, not a replay of the unknown call. + tx.send(ReaderEvent::Response(ServerResponse { + id: 1, + result: Some(json!({"late": true})), + error: None, + })) + .await + .unwrap(); + tx.send(ReaderEvent::Response(ServerResponse { + id: 2, + result: Some(json!({"next": true})), + error: None, + })) + .await + .unwrap(); + let result = client + .request::(McpPhase::Running, "next-method", json!({})) + .await + .unwrap(); + assert_eq!(result, json!({"next": true})); + let captured = bytes.lock().unwrap().clone(); + let frames: Vec = captured + .split(|byte| *byte == b'\n') + .filter(|line| !line.is_empty()) + .map(|line| serde_json::from_slice(line).unwrap()) + .collect(); + assert_eq!(frames.len(), 2); + assert_eq!(frames[0]["method"], "tools/call"); + assert_eq!(frames[1]["method"], "next-method"); + } +} + +// Observes actual OS-pipe backpressure, not a guessed delay or just a mock +// Pending. It forwards every operation to ChildStdin without changing framing. +struct PendingObserver { + writer: W, + accepted: usize, + reached: Option>, +} +impl AsyncWrite for PendingObserver { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let result = Pin::new(&mut self.writer).poll_write(cx, buf); + if let Poll::Ready(Ok(count)) = &result { + self.accepted += count; + } + if result.is_pending() && self.accepted > 1024 { + if let Some(tx) = self.reached.take() { + tx.send(()).unwrap(); + } + } + result + } + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.writer).poll_flush(cx) + } + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.writer).poll_shutdown(cx) + } +} + +#[cfg(unix)] +#[tokio::test] +async fn blocked_reader_reply_does_not_prevent_shutdown_process_reaping() { + // A local stdio fixture: initialize, consume initialized, emit an unsupported + // server request whose echoed ID exceeds pipe capacity, then never read. + // The shell has no descendants and ignores TERM to also prove kill fallback. + let init = json!({"jsonrpc":"2.0", "id":1, "result": {"protocolVersion": MCP_PROTOCOL_VERSION, "capabilities":{}, "serverInfo":{"name":"fixture", "version":"1"}}}); + let server_request = + json!({"jsonrpc":"2.0", "id":"x".repeat(256 * 1024), "method":"sampling/createMessage"}); + // Avoid argv limits: generate the large protocol ID in the shell itself. + let script = format!( + "trap '' TERM; read -r init; printf '%s\\n' '{init}'; read -r initialized; printf '%s' '{{\"jsonrpc\":\"2.0\",\"method\":\"sampling/createMessage\",\"id\":\"'; i=0; while [ $i -lt 4096 ]; do printf '%s' '{}'; i=$((i+1)); done; printf '%s\\n' '\"}}'; while :; do :; done", + "x".repeat(64) + ); + assert!(server_request.to_string().len() < McpStdioLimits::default().max_stdout_line_bytes); + let limits = McpStdioLimits { + shutdown_timeout: Duration::from_millis(50), + kill_timeout: Duration::from_millis(100), + ..Default::default() + }; + let mut client = McpStdioClient::spawn( + McpStdioServerSpec::new("reader-blocked", "/bin/sh") + .arg("-c") + .arg(script), + limits, + ) + .await + .unwrap(); + let pid = client.child.as_ref().unwrap().id().unwrap(); + let (reached, mut rx) = mpsc::unbounded_channel(); + { + let mut guard = client.stdin.lock().await; + let writer = guard.take().unwrap(); + *guard = Some(Box::new(PendingObserver { + writer, + accepted: 0, + reached: Some(reached), + })); + } + client.initialize().await.unwrap(); + timeout(Duration::from_secs(3), rx.recv()) + .await + .expect("reader reply reached OS pipe Pending") + .unwrap(); + assert!( + client.stdin.try_lock().is_err(), + "reader owns stdin while its reply is blocked" + ); + let report = timeout(Duration::from_secs(1), client.shutdown()) + .await + .unwrap() + .unwrap(); + assert!(report.terminated && report.killed); + assert!(report.exit_status.is_some()); + assert!(client.child.is_none()); + // Tokio wait already reaped the child; an additional wait cannot find it. + let mut status = 0; + assert_eq!( + unsafe { libc::waitpid(pid as libc::pid_t, &mut status, libc::WNOHANG) }, + -1 + ); + assert_eq!( + std::io::Error::last_os_error().raw_os_error(), + Some(libc::ECHILD) + ); +} diff --git a/crates/mcp/tests/fixtures/mock_server.rs b/crates/mcp/tests/fixtures/mock_server.rs index 9f2893e0e..48cee6fbd 100644 --- a/crates/mcp/tests/fixtures/mock_server.rs +++ b/crates/mcp/tests/fixtures/mock_server.rs @@ -6,8 +6,11 @@ use std::time::Duration; use serde_json::{Value, json}; fn main() { - let mode = env::var("YOI_MCP_MOCK_MODE").unwrap_or_else(|_| "success".to_string()); + let mode = env::args() + .nth(1) + .unwrap_or_else(|| env::var("YOI_MCP_MOCK_MODE").unwrap_or_else(|_| "success".to_string())); match mode.as_str() { + "stdin-unread" => stdin_unread(), "success" => success(), "tools" => tools_list(), "tools-call-normal" => tools_call_normal(), @@ -22,6 +25,20 @@ fn main() { } } +fn stdin_unread() { + let init = read_json(); + assert_eq!(init["method"], "initialize"); + write_json(json!({"jsonrpc": "2.0", "id": init["id"], "result": initialize_result()})); + assert_eq!(read_json()["method"], "notifications/initialized"); + // Protocol barrier: the host only writes its oversized call after this ack. + let ready = read_json(); + assert_eq!(ready["method"], "fixture/ready"); + write_json(json!({"jsonrpc": "2.0", "id": ready["id"], "result": {"ready": true}})); + loop { + thread::park(); + } +} + fn success() { let init = read_json(); assert_eq!(init["method"], "initialize"); diff --git a/crates/mcp/tests/stdio_lifecycle.rs b/crates/mcp/tests/stdio_lifecycle.rs index 4cb0e8aa2..77b96ebdc 100644 --- a/crates/mcp/tests/stdio_lifecycle.rs +++ b/crates/mcp/tests/stdio_lifecycle.rs @@ -278,3 +278,48 @@ async fn sampling_requests_fail_closed_and_are_not_advertised() { let shutdown = client.shutdown().await.expect("shutdown succeeds"); assert!(shutdown.exit_status.is_some_and(|status| status.success())); } + +#[tokio::test] +async fn unread_stdin_request_times_out_during_write_and_closes_before_reuse() { + let limits = McpStdioLimits { + request_timeout: Duration::from_millis(100), + ..tight_limits() + }; + let spec = McpStdioServerSpec::new("unread-stdin", env!("CARGO_BIN_EXE_mcp-stdio-mock-server")) + .arg("stdin-unread"); + let mut client = McpStdioClient::connect(spec, limits).await.unwrap(); + let ready: serde_json::Value = client + .request(McpPhase::Running, "fixture/ready", serde_json::json!({})) + .await + .unwrap(); + assert_eq!(ready["ready"], true); + // Far beyond OS pipe capacity, but inside the existing protocol bound. + let call = CallToolRequest::new( + "write-blocked", + serde_json::json!({"payload": "x".repeat(768 * 1024)}), + ); + let err = tokio::time::timeout(Duration::from_secs(1), client.call_tool(call)) + .await + .expect("the whole request, not only response wait, is bounded") + .unwrap_err(); + assert!( + matches!(err.kind, McpErrorKind::Timeout { ref operation, .. } if operation == "tools/call") + ); + assert_eq!( + err.request_stage(), + Some(mcp::stdio::McpRequestStage::WritingFrame) + ); + // The fixture barrier and reported stage establish actual write backpressure; + // the separate controlled fixture also verifies the exact partial bytes. + let next = client + .request::(McpPhase::Running, "fixture/next", serde_json::json!({})) + .await + .unwrap_err(); + assert!(next.to_string().contains("stdin is closed")); + let report = tokio::time::timeout(Duration::from_secs(1), client.shutdown()) + .await + .unwrap() + .unwrap(); + assert!(report.terminated || report.killed); + assert!(report.exit_status.is_some()); +} diff --git a/crates/worker/src/feature/mcp.rs b/crates/worker/src/feature/mcp.rs index ca13ae137..3af0c783b 100644 --- a/crates/worker/src/feature/mcp.rs +++ b/crates/worker/src/feature/mcp.rs @@ -1081,12 +1081,22 @@ impl Tool for McpStdioTool { &self, input_json: &str, _ctx: ToolExecutionContext, + ) -> Result { + self.execute_with_limits(input_json, McpStdioLimits::default()) + .await + } +} + +impl McpStdioTool { + async fn execute_with_limits( + &self, + input_json: &str, + limits: McpStdioLimits, ) -> Result { let arguments = parse_tool_arguments(input_json)?; - let mut client = - McpStdioClient::connect(self.server_spec.clone(), McpStdioLimits::default()) - .await - .map_err(|err| ToolError::ExecutionFailed(mcp_call_error_message(&err)))?; + let mut client = McpStdioClient::connect(self.server_spec.clone(), limits) + .await + .map_err(|err| ToolError::ExecutionFailed(mcp_call_error_message(&err)))?; let call_result = client .call_tool(CallToolRequest::new(self.mcp_tool_name.clone(), arguments)) @@ -2208,6 +2218,64 @@ read exit_notification || true ); } + #[cfg(unix)] + #[tokio::test] + async fn stdio_tool_request_timeout_reports_unknown_execution_without_replay() { + let dir = tempfile::tempdir().unwrap(); + let log = dir.path().join("frames.jsonl"); + // Frame capture is observed only after shutdown has waited for process + // exit. No guessed sleep or process-global test-control variable. + let script = r#" +printf '%s\n' 'connection' >> "$1" +read -r init || exit 1 +printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2025-11-25","capabilities":{"tools":{}},"serverInfo":{"name":"timeout-fixture","version":"1"}}}' +read -r initialized || exit 1 +while IFS= read -r frame; do + printf '%s\n' "$frame" >> "$1" +done +"#; + let tool = McpStdioTool { + server_spec: McpStdioServerSpec::new("timeout-fixture", "/bin/sh").args([ + "-c".to_string(), + script.to_string(), + "fixture".to_string(), + log.to_str().unwrap().to_string(), + ]), + mcp_tool_name: "demo-tool".to_string(), + }; + let limits = McpStdioLimits { + request_timeout: std::time::Duration::from_millis(100), + shutdown_timeout: std::time::Duration::from_millis(100), + kill_timeout: std::time::Duration::from_millis(100), + ..Default::default() + }; + let err = tokio::time::timeout( + std::time::Duration::from_secs(3), + tool.execute_with_limits("{}", limits), + ) + .await + .unwrap() + .expect_err("response intentionally withheld"); + let message = err.to_string(); + assert!( + message.contains("tools/call transport/protocol failure"), + "{message}" + ); + assert!(message.contains("timed out during tools/call"), "{message}"); + assert!( + message.contains("frame flushed; execution unknown; do not retry"), + "{message}" + ); + let capture = std::fs::read_to_string(log).unwrap(); + let mut lines = capture.lines(); + assert_eq!(lines.next(), Some("connection")); + let frames: Vec = lines + .map(|line| serde_json::from_str(line).unwrap()) + .collect(); + assert_eq!(frames.len(), 1, "must not reconnect or replay: {capture}"); + assert_eq!(frames[0]["method"], "tools/call"); + } + #[test] fn call_tool_result_renderer_marks_untrusted_mcp_error() { let result = CallToolResult {