mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
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:
committed by
copyberry
parent
68c6e5df2b
commit
44954d1b4b
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user