From 44954d1b4bd6030790b7f545e636ceeb994ba27a Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Tue, 14 Jul 2026 22:11:51 +0000 Subject: [PATCH] Serialize concurrent MCP stdin writes (#33180) ## What changed - Guard executor-backed MCP stdio sends with a single-permit semaphore so a second JSON-RPC message cannot start writing while the first write is pending. - Add a regression test that drives two concurrent sends and verifies the process receives each newline-delimited message in order without overlap. GitOrigin-RevId: 30c2e51827ecf57b5eb23d9ad1e876734d6e784e --- .../src/executor_process_transport.rs | 11 ++ .../src/executor_process_transport_tests.rs | 132 ++++++++++++++++++ 2 files changed, 143 insertions(+) diff --git a/codex-rs/rmcp-client/src/executor_process_transport.rs b/codex-rs/rmcp-client/src/executor_process_transport.rs index 23bd8c5e65e6..3ac78a4c450b 100644 --- a/codex-rs/rmcp-client/src/executor_process_transport.rs +++ b/codex-rs/rmcp-client/src/executor_process_transport.rs @@ -40,6 +40,7 @@ use rmcp::transport::Transport; use serde_json::from_slice; use serde_json::to_vec; use tokio::runtime::Handle; +use tokio::sync::Semaphore; use tokio::sync::broadcast; use tracing::debug; use tracing::info; @@ -162,6 +163,10 @@ pub(super) struct ExecutorProcessTransport { /// closes the transport. process: Arc, + /// Prevents concurrent rmcp send futures from issuing overlapping stdin writes. + /// The single-slot semaphore gives mutex semantics while its permit can safely cross `.await`. + stdin_write_semaphore: Arc, + /// Pushed output/lifecycle stream for the process. /// /// The executor process API still supports retained-output reads, but MCP @@ -204,6 +209,7 @@ impl ExecutorProcessTransport { let events = process.subscribe_events(); Self { process, + stdin_write_semaphore: Arc::new(Semaphore::new(1)), events, program_name, stdout: LineBuffer::default(), @@ -231,7 +237,12 @@ impl Transport for ExecutorProcessTransport { item: TxJsonRpcMessage, ) -> impl Future> + Send + 'static { let process = Arc::clone(&self.process); + let stdin_write_semaphore = Arc::clone(&self.stdin_write_semaphore); async move { + let _stdin_write_permit = stdin_write_semaphore + .acquire() + .await + .map_err(io::Error::other)?; // rmcp hands us a structured JSON-RPC message. Stdio transport on // the wire is JSON plus one newline delimiter. let mut bytes = to_vec(&item).map_err(io::Error::other)?; diff --git a/codex-rs/rmcp-client/src/executor_process_transport_tests.rs b/codex-rs/rmcp-client/src/executor_process_transport_tests.rs index 5a034f02728b..5b5b7e5e821d 100644 --- a/codex-rs/rmcp-client/src/executor_process_transport_tests.rs +++ b/codex-rs/rmcp-client/src/executor_process_transport_tests.rs @@ -1,10 +1,142 @@ use bytes::BytesMut; +use codex_exec_server::ExecProcess; +use codex_exec_server::ExecProcessEventReceiver; +use codex_exec_server::ExecProcessFuture; +use codex_exec_server::ProcessId; +use codex_exec_server::ProcessSignal; +use codex_exec_server::ReadResponse; +use codex_exec_server::WriteResponse; +use codex_exec_server::WriteStatus; use pretty_assertions::assert_eq; +use rmcp::service::RoleClient; +use rmcp::service::TxJsonRpcMessage; +use rmcp::transport::Transport; +use serde_json::json; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::task::Poll; +use tokio::sync::watch; +use super::ExecutorProcessTransport; use super::LineBuffer; use super::LineTooLong; use super::MAX_MCP_STDOUT_LINE_BYTES; +struct BlockingFirstWriteProcess { + process_id: ProcessId, + writes: StdMutex>>, + release_first_write: AtomicBool, +} + +impl BlockingFirstWriteProcess { + fn writes(&self) -> Vec> { + self.writes.lock().expect("writes lock").clone() + } +} + +impl ExecProcess for BlockingFirstWriteProcess { + fn process_id(&self) -> &ProcessId { + &self.process_id + } + + fn subscribe_wake(&self) -> watch::Receiver { + watch::channel(0).1 + } + + fn subscribe_events(&self) -> ExecProcessEventReceiver { + ExecProcessEventReceiver::empty() + } + + fn read( + &self, + _after_seq: Option, + _max_bytes: Option, + _wait_ms: Option, + ) -> ExecProcessFuture<'_, ReadResponse> { + Box::pin(async { unreachable!("send test should not read process output") }) + } + + fn write(&self, chunk: Vec) -> ExecProcessFuture<'_, WriteResponse> { + let first_write = { + let mut writes = self.writes.lock().expect("writes lock"); + writes.push(chunk); + writes.len() == 1 + }; + Box::pin(std::future::poll_fn(move |_| { + if first_write && !self.release_first_write.load(Ordering::Acquire) { + return Poll::Pending; + } + Poll::Ready(Ok(WriteResponse { + status: WriteStatus::Accepted, + })) + })) + } + + fn signal(&self, _signal: ProcessSignal) -> ExecProcessFuture<'_, ()> { + Box::pin(async { Ok(()) }) + } + + fn terminate(&self) -> ExecProcessFuture<'_, ()> { + Box::pin(async { Ok(()) }) + } +} + +#[tokio::test] +async fn serializes_concurrent_stdin_writes() { + let process = Arc::new(BlockingFirstWriteProcess { + process_id: ProcessId::from("mcp-stdio-test"), + writes: StdMutex::new(Vec::new()), + release_first_write: AtomicBool::new(false), + }); + let mut transport = + ExecutorProcessTransport::new(process.clone(), "mcp-stdio-test".to_string()); + let first_message: TxJsonRpcMessage = + serde_json::from_value(json!({ "jsonrpc": "2.0", "id": 1, "method": "ping" })) + .expect("first MCP message should deserialize"); + let second_message: TxJsonRpcMessage = + serde_json::from_value(json!({ "jsonrpc": "2.0", "id": 2, "method": "ping" })) + .expect("second MCP message should deserialize"); + + // Drive both sends explicitly so task scheduling cannot hide an overlapping write. + let first_send = transport.send(first_message); + tokio::pin!(first_send); + assert!(matches!(futures::poll!(first_send.as_mut()), Poll::Pending)); + assert_eq!( + process.writes(), + vec![b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"ping\"}\n".to_vec()] + ); + + let second_send = transport.send(second_message); + tokio::pin!(second_send); + assert!(matches!( + futures::poll!(second_send.as_mut()), + Poll::Pending + )); + assert_eq!( + process.writes(), + vec![b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"ping\"}\n".to_vec()] + ); + + process.release_first_write.store(true, Ordering::Release); + assert!(matches!( + futures::poll!(first_send.as_mut()), + Poll::Ready(Ok(())) + )); + assert!(matches!( + futures::poll!(second_send.as_mut()), + Poll::Ready(Ok(())) + )); + assert_eq!( + process.writes(), + vec![ + b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"ping\"}\n".to_vec(), + b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"ping\"}\n".to_vec(), + ] + ); +} + #[test] fn searches_only_new_bytes_after_partial_line() { let mut buffer = LineBuffer::default();