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
This commit is contained in:
Adam Perry @ OpenAI
2026-07-14 22:16:52 +00:00
committed by copyberry
parent 68c6e5df2b
commit 44954d1b4b
2 changed files with 143 additions and 0 deletions
@@ -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<dyn ExecProcess>,
/// 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<Semaphore>,
/// 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<RoleClient> for ExecutorProcessTransport {
item: TxJsonRpcMessage<RoleClient>,
) -> impl Future<Output = std::result::Result<(), Self::Error>> + 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)?;
@@ -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<Vec<Vec<u8>>>,
release_first_write: AtomicBool,
}
impl BlockingFirstWriteProcess {
fn writes(&self) -> Vec<Vec<u8>> {
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<u64> {
watch::channel(0).1
}
fn subscribe_events(&self) -> ExecProcessEventReceiver {
ExecProcessEventReceiver::empty()
}
fn read(
&self,
_after_seq: Option<u64>,
_max_bytes: Option<usize>,
_wait_ms: Option<u64>,
) -> ExecProcessFuture<'_, ReadResponse> {
Box::pin(async { unreachable!("send test should not read process output") })
}
fn write(&self, chunk: Vec<u8>) -> 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<RoleClient> =
serde_json::from_value(json!({ "jsonrpc": "2.0", "id": 1, "method": "ping" }))
.expect("first MCP message should deserialize");
let second_message: TxJsonRpcMessage<RoleClient> =
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();