mirror of
https://github.com/openai/codex.git
synced 2026-11-13 14:46:14 +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:
1 parent
68c6e5df2b
commit
44954d1b4b
2 files changed
+143
No files matched your search
@@ -40,6 +40,7 @@ use rmcp::transport::Transport;
|
|||||||
use serde_json::from_slice;
|
use serde_json::from_slice;
|
||||||
use serde_json::to_vec;
|
use serde_json::to_vec;
|
||||||
use tokio::runtime::Handle;
|
use tokio::runtime::Handle;
|
||||||
|
use tokio::sync::Semaphore;
|
||||||
use tokio::sync::broadcast;
|
use tokio::sync::broadcast;
|
||||||
use tracing::debug;
|
use tracing::debug;
|
||||||
use tracing::info;
|
use tracing::info;
|
||||||
@@ -162,6 +163,10 @@ pub(super) struct ExecutorProcessTransport {
|
|||||||
/// closes the transport.
|
/// closes the transport.
|
||||||
process: Arc<dyn ExecProcess>,
|
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.
|
/// Pushed output/lifecycle stream for the process.
|
||||||
///
|
///
|
||||||
/// The executor process API still supports retained-output reads, but MCP
|
/// The executor process API still supports retained-output reads, but MCP
|
||||||
@@ -204,6 +209,7 @@ impl ExecutorProcessTransport {
|
|||||||
let events = process.subscribe_events();
|
let events = process.subscribe_events();
|
||||||
Self {
|
Self {
|
||||||
process,
|
process,
|
||||||
|
stdin_write_semaphore: Arc::new(Semaphore::new(1)),
|
||||||
events,
|
events,
|
||||||
program_name,
|
program_name,
|
||||||
stdout: LineBuffer::default(),
|
stdout: LineBuffer::default(),
|
||||||
@@ -231,7 +237,12 @@ impl Transport<RoleClient> for ExecutorProcessTransport {
|
|||||||
item: TxJsonRpcMessage<RoleClient>,
|
item: TxJsonRpcMessage<RoleClient>,
|
||||||
) -> impl Future<Output = std::result::Result<(), Self::Error>> + Send + 'static {
|
) -> impl Future<Output = std::result::Result<(), Self::Error>> + Send + 'static {
|
||||||
let process = Arc::clone(&self.process);
|
let process = Arc::clone(&self.process);
|
||||||
|
let stdin_write_semaphore = Arc::clone(&self.stdin_write_semaphore);
|
||||||
async move {
|
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
|
// rmcp hands us a structured JSON-RPC message. Stdio transport on
|
||||||
// the wire is JSON plus one newline delimiter.
|
// the wire is JSON plus one newline delimiter.
|
||||||
let mut bytes = to_vec(&item).map_err(io::Error::other)?;
|
let mut bytes = to_vec(&item).map_err(io::Error::other)?;
|
||||||
|
|||||||
@@ -1,10 +1,142 @@
|
|||||||
use bytes::BytesMut;
|
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 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::LineBuffer;
|
||||||
use super::LineTooLong;
|
use super::LineTooLong;
|
||||||
use super::MAX_MCP_STDOUT_LINE_BYTES;
|
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]
|
#[test]
|
||||||
fn searches_only_new_bytes_after_partial_line() {
|
fn searches_only_new_bytes_after_partial_line() {
|
||||||
let mut buffer = LineBuffer::default();
|
let mut buffer = LineBuffer::default();
|
||||||
|
|||||||
Reference in new issue
Block a user