From d761119c2495e36257a55b1ca887b17638826ed6 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 19 Jun 2026 18:02:37 +0200 Subject: [PATCH] Use scalar concurrency for filesystem batches --- codex-rs/exec-server/src/client.rs | 41 +-- codex-rs/exec-server/src/connection.rs | 48 +-- .../src/noise_relay/executor_stream.rs | 2 +- .../src/noise_relay/executor_stream_tests.rs | 2 +- .../exec-server/src/noise_relay/harness.rs | 2 +- .../src/noise_relay/message_framing.rs | 9 +- .../src/noise_relay/message_framing_tests.rs | 5 +- codex-rs/exec-server/src/protocol.rs | 15 - codex-rs/exec-server/src/relay.rs | 15 +- .../exec-server/src/remote_file_system.rs | 185 ---------- codex-rs/exec-server/src/rpc.rs | 144 ++------ codex-rs/exec-server/src/server/handler.rs | 10 +- codex-rs/exec-server/src/server/processor.rs | 344 ++++++++---------- codex-rs/exec-server/src/server/registry.rs | 8 +- 14 files changed, 209 insertions(+), 621 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 31f2249d85..88c0c80342 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -87,7 +87,6 @@ use crate::protocol::INITIALIZE_METHOD; use crate::protocol::INITIALIZED_METHOD; use crate::protocol::InitializeParams; use crate::protocol::InitializeResponse; -use crate::protocol::InitializeWireResponse; use crate::protocol::ProcessOutputChunk; use crate::protocol::ProcessSignal; use crate::protocol::ReadParams; @@ -98,7 +97,6 @@ use crate::protocol::TerminateParams; use crate::protocol::TerminateResponse; use crate::protocol::WriteParams; use crate::protocol::WriteResponse; -use crate::rpc::RpcBatchCall; use crate::rpc::RpcCallError; use crate::rpc::RpcClient; @@ -201,7 +199,6 @@ struct Inner { http_body_streams_write_lock: Mutex<()>, http_body_stream_next_id: AtomicU64, session_id: OnceLock, - rpc_batch_supported: AtomicBool, reconnect_strategy: Option, } @@ -474,7 +471,7 @@ impl ExecServerClient { } = options; timeout(initialize_timeout, async { - let response: InitializeWireResponse = rpc_client + let response: InitializeResponse = rpc_client .call( INITIALIZE_METHOD, &InitializeParams { @@ -487,9 +484,6 @@ impl ExecServerClient { .inner .session_id .get_or_init(|| response.session_id.clone()); - self.inner - .rpc_batch_supported - .store(response.capabilities.rpc_batch, Ordering::Release); if session_id != &response.session_id { return Err(ExecServerError::Protocol(format!( "exec-server initialized an unexpected session {}", @@ -499,9 +493,7 @@ impl ExecServerClient { rpc_client .notify(INITIALIZED_METHOD, &serde_json::json!({})) .await?; - Ok(InitializeResponse { - session_id: response.session_id, - }) + Ok(response) }) .await .map_err(|_| ExecServerError::InitializeTimedOut { @@ -716,10 +708,6 @@ impl ExecServerClient { self.inner.session_id.get().cloned() } - pub(crate) fn supports_rpc_batch(&self) -> bool { - self.inner.rpc_batch_supported.load(Ordering::Acquire) - } - fn is_disconnected(&self) -> bool { self.inner.is_failed() } @@ -753,7 +741,6 @@ impl ExecServerClient { http_body_streams_write_lock: Mutex::new(()), http_body_stream_next_id: AtomicU64::new(1), session_id, - rpc_batch_supported: AtomicBool::new(false), reconnect_strategy, }); let client = Self { inner }; @@ -774,29 +761,6 @@ impl ExecServerClient { self.call_rpc(&rpc_client, method, params).await } - pub(crate) async fn call_batch( - &self, - calls: Vec, - ) -> Result>, ExecServerError> { - let rpc_client = self.inner.rpc_client().await?; - match rpc_client.call_batch(calls).await { - Ok(results) => Ok(results - .into_iter() - .map(|result| result.map_err(ExecServerError::from)) - .collect()), - Err(error) => { - let error = ExecServerError::from(error); - if is_transport_closed_error(&error) { - Err(ExecServerError::Disconnected(disconnected_message( - /*reason*/ None, - ))) - } else { - Err(error) - } - } - } - } - async fn call_rpc( &self, rpc_client: &Arc, @@ -849,7 +813,6 @@ impl From for ExecServerError { code: error.code, message: error.message, }, - RpcCallError::InvalidBatch(message) => Self::Protocol(message), } } } diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 255eb26ec8..2d2e5afe06 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -12,8 +12,6 @@ use futures::Sink; use futures::SinkExt; use futures::Stream; use futures::StreamExt; -use serde::Deserialize; -use serde::Serialize; use tokio::io::AsyncRead; use tokio::io::AsyncWrite; use tokio::process::Child; @@ -31,7 +29,6 @@ use tokio::io::BufReader; use tokio::io::BufWriter; pub(crate) const CHANNEL_CAPACITY: usize = 128; -pub(crate) const MAX_RPC_BATCH_REQUESTS: usize = 512; const STDIO_TERMINATION_GRACE_PERIOD: Duration = Duration::from_secs(2); #[cfg(test)] pub(crate) const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_millis(25); @@ -41,37 +38,10 @@ pub(crate) const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30 #[derive(Debug)] pub(crate) enum JsonRpcConnectionEvent { Message(JSONRPCMessage), - Batch(Vec), MalformedMessage { reason: String }, Disconnected { reason: Option }, } -/// Exec-server-local wire envelope for one JSON-RPC message or a batch. -/// -/// Keeping this type local avoids expanding the app-server protocol surface while still using -/// standard JSON-RPC batch framing on exec-server transports. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(untagged)] -pub(crate) enum JsonRpcWireMessage { - Single(JSONRPCMessage), - Batch(Vec), -} - -impl JsonRpcWireMessage { - pub(crate) fn into_connection_event(self) -> JsonRpcConnectionEvent { - match self { - Self::Single(message) => JsonRpcConnectionEvent::Message(message), - Self::Batch(messages) => JsonRpcConnectionEvent::Batch(messages), - } - } -} - -impl From for JsonRpcWireMessage { - fn from(message: JSONRPCMessage) -> Self { - Self::Single(message) - } -} - #[derive(Clone)] pub(crate) enum JsonRpcTransport { // Plain means no child process; transport bytes may still be encrypted. @@ -251,7 +221,7 @@ fn log_stdio_child_wait_result(result: std::io::Result } pub(crate) struct JsonRpcConnection { - pub(crate) outgoing_tx: mpsc::Sender, + pub(crate) outgoing_tx: mpsc::Sender, pub(crate) incoming_rx: mpsc::Receiver, pub(crate) disconnected_rx: watch::Receiver, pub(crate) task_handles: Vec>, @@ -279,10 +249,10 @@ impl JsonRpcConnection { if line.trim().is_empty() { continue; } - match serde_json::from_str::(&line) { + match serde_json::from_str::(&line) { Ok(message) => { if incoming_tx_for_reader - .send(message.into_connection_event()) + .send(JsonRpcConnectionEvent::Message(message)) .await .is_err() { @@ -425,7 +395,7 @@ impl JsonRpcConnection { Some(Ok(message)) => match message.parse_jsonrpc_frame() { Ok(JsonRpcWebSocketFrame::Message(message)) => { if incoming_tx - .send(message.into_connection_event()) + .send(JsonRpcConnectionEvent::Message(message)) .await .is_err() { @@ -494,7 +464,7 @@ impl JsonRpcConnection { } enum JsonRpcWebSocketFrame { - Message(JsonRpcWireMessage), + Message(JSONRPCMessage), Close, Ignore, } @@ -579,7 +549,7 @@ async fn send_malformed_message( async fn write_jsonrpc_line_message( writer: &mut BufWriter, - message: &JsonRpcWireMessage, + message: &JSONRPCMessage, ) -> std::io::Result<()> where W: AsyncWrite + Unpin, @@ -594,7 +564,7 @@ where async fn send_websocket_jsonrpc_message( websocket_writer: &mut W, connection_label: &str, - message: &JsonRpcWireMessage, + message: &JSONRPCMessage, ) -> Result<(), String> where W: Sink + Unpin, @@ -614,7 +584,7 @@ where } } -fn serialize_jsonrpc_message(message: &JsonRpcWireMessage) -> Result { +fn serialize_jsonrpc_message(message: &JSONRPCMessage) -> Result { serde_json::to_string(message) } @@ -724,7 +694,7 @@ mod tests { ); let message = test_jsonrpc_message(); - connection.outgoing_tx.send(message.clone().into()).await?; + connection.outgoing_tx.send(message.clone()).await?; control.wait_for_blocked_write().await?; control.send_inbound(Message::Pong(b"check".to_vec().into()))?; assert!( diff --git a/codex-rs/exec-server/src/noise_relay/executor_stream.rs b/codex-rs/exec-server/src/noise_relay/executor_stream.rs index 04f25dc446..57a05579dc 100644 --- a/codex-rs/exec-server/src/noise_relay/executor_stream.rs +++ b/codex-rs/exec-server/src/noise_relay/executor_stream.rs @@ -75,7 +75,7 @@ impl NoiseVirtualStream { }; for message in self.inbound_decoder.push(&plaintext)? { self.incoming_tx - .try_send(message.into_connection_event()) + .try_send(JsonRpcConnectionEvent::Message(message)) .map_err(|_| { ExecServerError::Protocol( "Noise virtual stream inbound queue is full or closed".to_string(), diff --git a/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs b/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs index d757a12393..9119236cfd 100644 --- a/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs +++ b/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs @@ -51,7 +51,7 @@ async fn processor_exit_reports_closed_virtual_stream() -> Result<()> { id: RequestId::Integer(1), result: serde_json::Value::Null, }); - let ciphertext = harness_transport.encrypt(&frame_jsonrpc_message(&message.clone().into())?)?; + let ciphertext = harness_transport.encrypt(&frame_jsonrpc_message(&message)?)?; stream.receive_data(RelayData { seq: 0, segment_index: 0, diff --git a/codex-rs/exec-server/src/noise_relay/harness.rs b/codex-rs/exec-server/src/noise_relay/harness.rs index 62e82e608c..1989525224 100644 --- a/codex-rs/exec-server/src/noise_relay/harness.rs +++ b/codex-rs/exec-server/src/noise_relay/harness.rs @@ -420,7 +420,7 @@ async fn receive_data( // messages; emit only complete, successfully parsed messages. for message in decoder.push(&plaintext)? { incoming_tx - .send(message.into_connection_event()) + .send(JsonRpcConnectionEvent::Message(message)) .await .map_err(|_| ExecServerError::Closed)?; } diff --git a/codex-rs/exec-server/src/noise_relay/message_framing.rs b/codex-rs/exec-server/src/noise_relay/message_framing.rs index b84253695c..24ee48173e 100644 --- a/codex-rs/exec-server/src/noise_relay/message_framing.rs +++ b/codex-rs/exec-server/src/noise_relay/message_framing.rs @@ -1,5 +1,6 @@ +use codex_app_server_protocol::JSONRPCMessage; + use crate::ExecServerError; -use crate::connection::JsonRpcWireMessage; const LENGTH_PREFIX_BYTES: usize = size_of::(); const MAX_NOISE_JSONRPC_MESSAGE_LEN: usize = 64 * 1024 * 1024; @@ -11,9 +12,7 @@ pub(crate) const NOISE_RECORD_PLAINTEXT_LEN: usize = 60 * 1024; /// exec-server responses can be much larger. A four-byte authenticated length /// prefix lets the caller split this byte stream into bounded Noise records and /// lets the receiver reconstruct exact JSON-RPC message boundaries. -pub(crate) fn frame_jsonrpc_message( - message: &JsonRpcWireMessage, -) -> Result, ExecServerError> { +pub(crate) fn frame_jsonrpc_message(message: &JSONRPCMessage) -> Result, ExecServerError> { let mut framed = vec![0; LENGTH_PREFIX_BYTES]; serde_json::to_writer(&mut framed, message)?; let message_len = framed.len() - LENGTH_PREFIX_BYTES; @@ -40,7 +39,7 @@ impl JsonRpcMessageDecoder { pub(crate) fn push( &mut self, plaintext_record: &[u8], - ) -> Result, ExecServerError> { + ) -> Result, ExecServerError> { if plaintext_record.len() > NOISE_RECORD_PLAINTEXT_LEN { return Err(ExecServerError::Protocol( "Noise relay plaintext record exceeds maximum length".to_string(), diff --git a/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs b/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs index 5cc3ed10f1..c3df14bf56 100644 --- a/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs +++ b/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs @@ -7,7 +7,6 @@ use super::MAX_NOISE_JSONRPC_MESSAGE_LEN; use super::NOISE_RECORD_PLAINTEXT_LEN; use super::frame_jsonrpc_message; use crate::ExecServerError; -use crate::connection::JsonRpcWireMessage; #[test] fn fragments_and_reassembles_large_jsonrpc_message() { @@ -17,7 +16,7 @@ fn fragments_and_reassembles_large_jsonrpc_message() { "data": "x".repeat(128 * 1024), })), }); - let framed = frame_jsonrpc_message(&message.clone().into()).unwrap(); + let framed = frame_jsonrpc_message(&message).unwrap(); assert!(framed.len() > 128 * 1024); let mut decoder = JsonRpcMessageDecoder::default(); @@ -26,7 +25,7 @@ fn fragments_and_reassembles_large_jsonrpc_message() { decoded.extend(decoder.push(record).unwrap()); } - assert_eq!(decoded, vec![JsonRpcWireMessage::Single(message)]); + assert_eq!(decoded, vec![message]); } #[test] diff --git a/codex-rs/exec-server/src/protocol.rs b/codex-rs/exec-server/src/protocol.rs index cde670a2ed..97ef86aad2 100644 --- a/codex-rs/exec-server/src/protocol.rs +++ b/codex-rs/exec-server/src/protocol.rs @@ -66,21 +66,6 @@ pub struct InitializeResponse { pub session_id: String, } -#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct ExecServerCapabilities { - #[serde(default)] - pub rpc_batch: bool, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct InitializeWireResponse { - pub session_id: String, - #[serde(default)] - pub capabilities: ExecServerCapabilities, -} - /// Information about an execution/filesystem environment. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index ca467d0d0f..033da4b2a7 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use std::time::Duration; +use codex_app_server_protocol::JSONRPCMessage; use futures::Sink; use futures::SinkExt; use futures::Stream; @@ -24,7 +25,6 @@ use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; use crate::connection::JsonRpcTransport; -use crate::connection::JsonRpcWireMessage; use crate::connection::WEBSOCKET_KEEPALIVE_INTERVAL; use crate::noise_channel::NoiseChannelIdentity; use crate::noise_channel::NoiseChannelPublicKey; @@ -170,7 +170,7 @@ impl RelayMessageFrame { } } - fn into_jsonrpc_message(self) -> Result { + fn into_jsonrpc_message(self) -> Result { let payload = self.into_data()?.payload; serde_json::from_slice(&payload).map_err(ExecServerError::Json) } @@ -211,7 +211,7 @@ pub(crate) fn decode_relay_message_frame( .map_err(|err| ExecServerError::Protocol(format!("invalid relay message frame: {err}"))) } -pub(crate) fn jsonrpc_payload(message: &JsonRpcWireMessage) -> Result, ExecServerError> { +pub(crate) fn jsonrpc_payload(message: &JSONRPCMessage) -> Result, ExecServerError> { serde_json::to_vec(message).map_err(ExecServerError::Json) } @@ -346,7 +346,7 @@ where &mut websocket, &mut keepalive, &incoming_tx, - message.into_connection_event(), + JsonRpcConnectionEvent::Message(message), ) .await { @@ -847,7 +847,6 @@ mod tests { use std::task::Poll; use std::time::Duration; - use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::RequestId; use futures::Sink; @@ -881,7 +880,7 @@ mod tests { encode_relay_message_frame(&RelayMessageFrame::data( stream_id, /*seq*/ 0, - jsonrpc_payload(&message.clone().into())?, + jsonrpc_payload(&message)?, )) .into(), )) @@ -1030,7 +1029,7 @@ mod tests { let message = test_jsonrpc_message(); control.set_write_blocked(); - connection.outgoing_tx.send(message.clone().into()).await?; + connection.outgoing_tx.send(message.clone()).await?; control.wait_for_blocked_write().await?; control.send_inbound(Message::Pong(b"check".to_vec().into()))?; assert!( @@ -1048,7 +1047,7 @@ mod tests { }; let frame = decode_relay_message_frame(data_payload.as_ref())?; assert_eq!(frame.stream_id, stream_id); - assert_eq!(frame.into_jsonrpc_message()?, message.into()); + assert_eq!(frame.into_jsonrpc_message()?, message); drop(connection); Ok(()) } diff --git a/codex-rs/exec-server/src/remote_file_system.rs b/codex-rs/exec-server/src/remote_file_system.rs index eadabba5d9..f9136e1461 100644 --- a/codex-rs/exec-server/src/remote_file_system.rs +++ b/codex-rs/exec-server/src/remote_file_system.rs @@ -10,31 +10,20 @@ use crate::ExecServerError; use crate::ExecutorFileSystem; use crate::ExecutorFileSystemFuture; use crate::FileMetadata; -use crate::FileSystemOperation; -use crate::FileSystemOperationOutput; -use crate::FileSystemOperationResult; use crate::FileSystemReadStream; use crate::FileSystemResult; use crate::FileSystemSandboxContext; use crate::ReadDirectoryEntry; use crate::RemoveOptions; use crate::client::LazyRemoteExecServerClient; -use crate::connection::MAX_RPC_BATCH_REQUESTS; use crate::protocol::FsCanonicalizeParams; -use crate::protocol::FsCanonicalizeResponse; use crate::protocol::FsCopyParams; use crate::protocol::FsCreateDirectoryParams; use crate::protocol::FsGetMetadataParams; -use crate::protocol::FsGetMetadataResponse; use crate::protocol::FsReadDirectoryParams; use crate::protocol::FsReadFileParams; -use crate::protocol::FsReadFileResponse; use crate::protocol::FsRemoveParams; use crate::protocol::FsWriteFileParams; -use crate::rpc::RpcBatchCall; -use codex_file_system::execute_batch_with_scalar_operations; -use serde::Serialize; -use serde::de::DeserializeOwned; const INVALID_REQUEST_ERROR_CODE: i64 = -32600; const NOT_FOUND_ERROR_CODE: i64 = -32004; @@ -234,68 +223,6 @@ impl RemoteFileSystem { .map_err(map_remote_error)?; Ok(()) } - - async fn execute_batch( - &self, - operations: Vec, - sandbox: Option<&FileSystemSandboxContext>, - ) -> FileSystemResult> { - if operations.is_empty() { - return Ok(Vec::new()); - } - - let client = self.client.get().await.map_err(map_remote_error)?; - if !client.supports_rpc_batch() { - return execute_batch_with_scalar_operations(self, operations, sandbox).await; - } - - let sandbox = remote_sandbox_context(sandbox); - let mut output_kinds = Vec::with_capacity(operations.len()); - let mut calls = Vec::with_capacity(operations.len()); - for operation in operations { - let (output_kind, call) = remote_batch_call(operation, sandbox.clone())?; - output_kinds.push(output_kind); - calls.push(call); - } - - let mut decoded = Vec::with_capacity(output_kinds.len()); - let mut output_kinds = output_kinds.into_iter(); - let mut calls = calls.into_iter(); - loop { - let output_kind_chunk = output_kinds - .by_ref() - .take(MAX_RPC_BATCH_REQUESTS) - .collect::>(); - if output_kind_chunk.is_empty() { - break; - } - let call_chunk = calls - .by_ref() - .take(MAX_RPC_BATCH_REQUESTS) - .collect::>(); - let results = client - .call_batch(call_chunk) - .await - .map_err(map_remote_error)?; - if results.len() != output_kind_chunk.len() { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!( - "exec-server returned {} batch results for {} operations", - results.len(), - output_kind_chunk.len() - ), - )); - } - decoded.extend( - output_kind_chunk - .into_iter() - .zip(results) - .map(|(output_kind, result)| decode_remote_batch_result(output_kind, result)), - ); - } - Ok(decoded) - } } impl ExecutorFileSystem for RemoteFileSystem { @@ -383,118 +310,6 @@ impl ExecutorFileSystem for RemoteFileSystem { sandbox, )) } - - fn execute_batch<'a>( - &'a self, - operations: Vec, - sandbox: Option<&'a FileSystemSandboxContext>, - ) -> ExecutorFileSystemFuture<'a, Vec> { - Box::pin(RemoteFileSystem::execute_batch(self, operations, sandbox)) - } -} - -#[derive(Clone, Copy)] -enum RemoteBatchOutputKind { - Canonicalize, - ReadFile, - GetMetadata, - ReadDirectory, -} - -fn remote_batch_call( - operation: FileSystemOperation, - sandbox: Option, -) -> io::Result<(RemoteBatchOutputKind, RpcBatchCall)> { - match operation { - FileSystemOperation::Canonicalize { path } => Ok(( - RemoteBatchOutputKind::Canonicalize, - rpc_batch_call( - crate::protocol::FS_CANONICALIZE_METHOD, - FsCanonicalizeParams { path, sandbox }, - )?, - )), - FileSystemOperation::ReadFile { path } => Ok(( - RemoteBatchOutputKind::ReadFile, - rpc_batch_call( - crate::protocol::FS_READ_FILE_METHOD, - FsReadFileParams { path, sandbox }, - )?, - )), - FileSystemOperation::GetMetadata { path } => Ok(( - RemoteBatchOutputKind::GetMetadata, - rpc_batch_call( - crate::protocol::FS_GET_METADATA_METHOD, - FsGetMetadataParams { path, sandbox }, - )?, - )), - FileSystemOperation::ReadDirectory { path } => Ok(( - RemoteBatchOutputKind::ReadDirectory, - rpc_batch_call( - crate::protocol::FS_READ_DIRECTORY_METHOD, - FsReadDirectoryParams { path, sandbox }, - )?, - )), - } -} - -fn rpc_batch_call(method: &str, params: impl Serialize) -> io::Result { - Ok(RpcBatchCall { - method: method.to_string(), - params: serde_json::to_value(params) - .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?, - }) -} - -fn decode_remote_batch_result( - output_kind: RemoteBatchOutputKind, - result: Result, -) -> FileSystemOperationResult { - let value = result.map_err(map_remote_error)?; - match output_kind { - RemoteBatchOutputKind::Canonicalize => { - let response: FsCanonicalizeResponse = decode_batch_value(value)?; - Ok(FileSystemOperationOutput::Canonicalize(response.path)) - } - RemoteBatchOutputKind::ReadFile => { - let response: FsReadFileResponse = decode_batch_value(value)?; - let contents = STANDARD.decode(response.data_base64).map_err(|error| { - io::Error::new( - io::ErrorKind::InvalidData, - format!("remote fs/readFile returned invalid base64 dataBase64: {error}"), - ) - })?; - Ok(FileSystemOperationOutput::ReadFile(contents)) - } - RemoteBatchOutputKind::GetMetadata => { - let response: FsGetMetadataResponse = decode_batch_value(value)?; - Ok(FileSystemOperationOutput::GetMetadata(FileMetadata { - is_directory: response.is_directory, - is_file: response.is_file, - is_symlink: response.is_symlink, - size: response.size, - created_at_ms: response.created_at_ms, - modified_at_ms: response.modified_at_ms, - })) - } - RemoteBatchOutputKind::ReadDirectory => { - let response: crate::protocol::FsReadDirectoryResponse = decode_batch_value(value)?; - Ok(FileSystemOperationOutput::ReadDirectory( - response - .entries - .into_iter() - .map(|entry| ReadDirectoryEntry { - file_name: entry.file_name, - is_directory: entry.is_directory, - is_file: entry.is_file, - }) - .collect(), - )) - } - } -} - -fn decode_batch_value(value: serde_json::Value) -> io::Result { - serde_json::from_value(value).map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)) } fn remote_sandbox_context( diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index b90d23f7a4..c2a5124541 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -26,8 +26,6 @@ use tokio::task::JoinHandle; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; use crate::connection::JsonRpcTransport; -use crate::connection::JsonRpcWireMessage; -use crate::connection::MAX_RPC_BATCH_REQUESTS; pub(crate) const SESSION_ALREADY_ATTACHED_ERROR_CODE: i64 = -32010; @@ -39,13 +37,6 @@ pub(crate) enum RpcCallError { Json(serde_json::Error), /// The executor returned a JSON-RPC error response for this call. Server(JSONRPCErrorError), - /// The caller attempted to construct a batch that cannot be represented safely. - InvalidBatch(String), -} - -pub(crate) struct RpcBatchCall { - pub(crate) method: String, - pub(crate) params: Value, } type PendingRequest = oneshot::Sender>; @@ -73,7 +64,6 @@ pub(crate) enum RpcServerOutboundMessage { error: JSONRPCErrorError, }, Notification(JSONRPCNotification), - Batch(Vec), } #[derive(Clone)] @@ -117,7 +107,7 @@ impl RpcNotificationSender { pub(crate) struct RpcRouter { request_routes: HashMap<&'static str, RequestRoute>, - batchable_request_methods: HashSet<&'static str>, + concurrent_request_methods: HashSet<&'static str>, notification_routes: HashMap<&'static str, NotificationRoute>, } @@ -125,7 +115,7 @@ impl Default for RpcRouter { fn default() -> Self { Self { request_routes: HashMap::new(), - batchable_request_methods: HashSet::new(), + concurrent_request_methods: HashSet::new(), notification_routes: HashMap::new(), } } @@ -146,7 +136,7 @@ where F: Fn(Arc, P) -> Fut + Send + Sync + 'static, Fut: Future> + Send + 'static, { - self.batchable_request_methods.remove(method); + self.concurrent_request_methods.remove(method); self.request_routes.insert( method, Box::new(move |state, request| { @@ -176,7 +166,7 @@ where ); } - pub(crate) fn batchable_request(&mut self, method: &'static str, handler: F) + pub(crate) fn concurrent_request(&mut self, method: &'static str, handler: F) where P: DeserializeOwned + Send + 'static, R: Serialize + Send + 'static, @@ -184,7 +174,7 @@ where Fut: Future> + Send + 'static, { self.request(method, handler); - self.batchable_request_methods.insert(method); + self.concurrent_request_methods.insert(method); } pub(crate) fn request_with_id(&mut self, method: &'static str, handler: F) @@ -193,7 +183,7 @@ where F: Fn(Arc, RequestId, P) -> Fut + Send + Sync + 'static, Fut: Future> + Send + 'static, { - self.batchable_request_methods.remove(method); + self.concurrent_request_methods.remove(method); self.request_routes.insert( method, Box::new(move |state, request| { @@ -242,8 +232,8 @@ where self.request_routes.get(method) } - pub(crate) fn is_request_batchable(&self, method: &str) -> bool { - self.batchable_request_methods.contains(method) + pub(crate) fn is_request_concurrent(&self, method: &str) -> bool { + self.concurrent_request_methods.contains(method) } pub(crate) fn notification_route(&self, method: &str) -> Option<&NotificationRoute> { @@ -252,7 +242,7 @@ where } pub(crate) struct RpcClient { - write_tx: mpsc::Sender, + write_tx: mpsc::Sender, pending: Arc>>, // Shared transport state from `JsonRpcConnection`. Calls use this to fail // immediately when the socket closes, even if no JSON-RPC error response @@ -282,7 +272,7 @@ impl RpcClient { let closed_for_reader = Arc::clone(&closed); let transport_for_reader = transport.clone(); let reader_task = tokio::spawn(async move { - let disconnect_reason = 'reader: loop { + let disconnect_reason = loop { let Some(event) = incoming_rx.recv().await else { break None; }; @@ -295,16 +285,6 @@ impl RpcClient { break None; } } - JsonRpcConnectionEvent::Batch(messages) => { - for message in messages { - if let Err(err) = - handle_server_message(&pending_for_reader, &event_tx, message).await - { - let _ = err; - break 'reader None; - } - } - } JsonRpcConnectionEvent::MalformedMessage { reason } => { let _ = reason; break None; @@ -350,13 +330,10 @@ impl RpcClient { return Err(RpcCallError::Closed); } self.write_tx - .send( - JSONRPCMessage::Notification(JSONRPCNotification { - method: method.to_string(), - params: Some(params), - }) - .into(), - ) + .send(JSONRPCMessage::Notification(JSONRPCNotification { + method: method.to_string(), + params: Some(params), + })) .await .map_err(|_| RpcCallError::Closed) } @@ -401,15 +378,12 @@ impl RpcClient { }; if self .write_tx - .send( - JSONRPCMessage::Request(JSONRPCRequest { - id: request_id.clone(), - method: method.to_string(), - params: Some(params), - trace: None, - }) - .into(), - ) + .send(JSONRPCMessage::Request(JSONRPCRequest { + id: request_id.clone(), + method: method.to_string(), + params: Some(params), + trace: None, + })) .await .is_err() { @@ -432,64 +406,6 @@ impl RpcClient { serde_json::from_value(response).map_err(RpcCallError::Json) } - pub(crate) async fn call_batch( - &self, - calls: Vec, - ) -> Result>, RpcCallError> { - if calls.is_empty() { - return Ok(Vec::new()); - } - if calls.len() > MAX_RPC_BATCH_REQUESTS { - return Err(RpcCallError::InvalidBatch(format!( - "JSON-RPC batch contains {} requests; maximum is {MAX_RPC_BATCH_REQUESTS}", - calls.len() - ))); - } - - let mut request_ids = Vec::with_capacity(calls.len()); - let mut requests = Vec::with_capacity(calls.len()); - let mut response_receivers = Vec::with_capacity(calls.len()); - { - let mut pending = self.pending.lock().await; - if self.closed.load(Ordering::Acquire) || *self.disconnected_rx.borrow() { - return Err(RpcCallError::Closed); - } - for RpcBatchCall { method, params } in calls { - let request_id = - RequestId::Integer(self.next_request_id.fetch_add(1, Ordering::SeqCst)); - let (response_tx, response_rx) = oneshot::channel(); - pending.insert(request_id.clone(), response_tx); - request_ids.push(request_id.clone()); - requests.push(JSONRPCMessage::Request(JSONRPCRequest { - id: request_id, - method, - params: Some(params), - trace: None, - })); - response_receivers.push(response_rx); - } - } - - if self - .write_tx - .send(JsonRpcWireMessage::Batch(requests)) - .await - .is_err() - { - let mut pending = self.pending.lock().await; - for request_id in request_ids { - pending.remove(&request_id); - } - return Err(RpcCallError::Closed); - } - - let mut results = Vec::with_capacity(response_receivers.len()); - for response_rx in response_receivers { - results.push(response_rx.await.map_err(|_| RpcCallError::Closed)?); - } - Ok(results) - } - #[cfg(test)] pub(crate) async fn pending_request_count(&self) -> usize { self.pending.lock().await.len() @@ -508,34 +424,22 @@ impl Drop for RpcClient { pub(crate) fn encode_server_message( message: RpcServerOutboundMessage, -) -> Result { +) -> Result { match message { RpcServerOutboundMessage::Response { request_id, result } => { Ok(JSONRPCMessage::Response(JSONRPCResponse { id: request_id, result, - }) - .into()) + })) } RpcServerOutboundMessage::Error { request_id, error } => { Ok(JSONRPCMessage::Error(JSONRPCError { id: request_id, error, - }) - .into()) + })) } RpcServerOutboundMessage::Notification(notification) => { - Ok(JSONRPCMessage::Notification(notification).into()) - } - RpcServerOutboundMessage::Batch(messages) => { - let mut encoded = Vec::with_capacity(messages.len()); - for message in messages { - match encode_server_message(message)? { - JsonRpcWireMessage::Single(message) => encoded.push(message), - JsonRpcWireMessage::Batch(messages) => encoded.extend(messages), - } - } - Ok(JsonRpcWireMessage::Batch(encoded)) + Ok(JSONRPCMessage::Notification(notification)) } } } diff --git a/codex-rs/exec-server/src/server/handler.rs b/codex-rs/exec-server/src/server/handler.rs index bb13e37847..e0ede94c77 100644 --- a/codex-rs/exec-server/src/server/handler.rs +++ b/codex-rs/exec-server/src/server/handler.rs @@ -17,7 +17,6 @@ use crate::client::http_client::ReqwestHttpRequestRunner; use crate::protocol::EnvironmentInfo; use crate::protocol::ExecParams; use crate::protocol::ExecResponse; -use crate::protocol::ExecServerCapabilities; use crate::protocol::FsCanonicalizeParams; use crate::protocol::FsCanonicalizeResponse; use crate::protocol::FsCloseParams; @@ -42,7 +41,7 @@ use crate::protocol::FsWriteFileParams; use crate::protocol::FsWriteFileResponse; use crate::protocol::HttpRequestParams; use crate::protocol::InitializeParams; -use crate::protocol::InitializeWireResponse; +use crate::protocol::InitializeResponse; use crate::protocol::ReadParams; use crate::protocol::ReadResponse; use crate::protocol::SignalParams; @@ -108,7 +107,7 @@ impl ExecServerHandler { pub(crate) async fn initialize( &self, params: InitializeParams, - ) -> Result { + ) -> Result { if self.initialize_requested.swap(true, Ordering::SeqCst) { return Err(invalid_request( "initialize may only be sent once per connection".to_string(), @@ -136,10 +135,7 @@ impl ExecServerHandler { .session .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(session); - Ok(InitializeWireResponse { - session_id, - capabilities: ExecServerCapabilities { rpc_batch: true }, - }) + Ok(InitializeResponse { session_id }) } pub(crate) fn initialized(&self) -> Result<(), String> { diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index ca0fa1a427..7d203dd495 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -1,7 +1,7 @@ use std::sync::Arc; -use futures::StreamExt; use tokio::sync::mpsc; +use tokio::task::JoinSet; use tracing::debug; use tracing::warn; @@ -9,7 +9,7 @@ use crate::ExecServerRuntimePaths; use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; -use crate::connection::MAX_RPC_BATCH_REQUESTS; +use crate::protocol::INITIALIZED_METHOD; use crate::rpc::RpcNotificationSender; use crate::rpc::RpcRouter; use crate::rpc::RpcServerOutboundMessage; @@ -20,7 +20,7 @@ use crate::server::ExecServerHandler; use crate::server::registry::build_router; use crate::server::session_registry::SessionRegistry; -const MAX_RPC_BATCH_CONCURRENCY: usize = 32; +const MAX_CONCURRENT_READ_ONLY_REQUESTS: usize = 32; #[derive(Clone)] pub(crate) struct ConnectionProcessor { @@ -67,6 +67,8 @@ async fn run_connection( notifications, runtime_paths, )); + let mut read_only_tasks = JoinSet::new(); + let mut read_only_requests_enabled = false; let outbound_task = tokio::spawn(async move { while let Some(message) = outgoing_rx.recv().await { @@ -84,13 +86,32 @@ async fn run_connection( }); // Process inbound events sequentially to preserve initialize/initialized ordering. - while let Some(event) = incoming_rx.recv().await { + loop { if !handler.is_session_attached() { debug!("exec-server connection evicted after session resume"); break; } + let event = tokio::select! { + result = read_only_tasks.join_next(), if !read_only_tasks.is_empty() => { + if !send_finished_request(result, &outgoing_tx).await { + break; + } + continue; + } + event = incoming_rx.recv() => { + let Some(event) = event else { + break; + }; + event + } + }; match event { JsonRpcConnectionEvent::MalformedMessage { reason } => { + if !drain_read_only_tasks(&mut read_only_tasks, &outgoing_tx, &mut disconnected_rx) + .await + { + break; + } warn!("ignoring malformed exec-server message: {reason}"); if outgoing_tx .send(RpcServerOutboundMessage::Error { @@ -105,6 +126,36 @@ async fn run_connection( } JsonRpcConnectionEvent::Message(message) => match message { codex_app_server_protocol::JSONRPCMessage::Request(request) => { + if read_only_requests_enabled + && router.is_request_concurrent(request.method.as_str()) + { + if !limit_read_only_tasks( + &mut read_only_tasks, + &outgoing_tx, + &mut disconnected_rx, + ) + .await + { + break; + } + spawn_read_only_request( + &mut read_only_tasks, + Arc::clone(&router), + Arc::clone(&handler), + request, + ); + continue; + } + + if !drain_read_only_tasks( + &mut read_only_tasks, + &outgoing_tx, + &mut disconnected_rx, + ) + .await + { + break; + } let message = tokio::select! { message = dispatch_request( Arc::clone(&router), @@ -123,6 +174,15 @@ async fn run_connection( } } codex_app_server_protocol::JSONRPCMessage::Notification(notification) => { + if !drain_read_only_tasks( + &mut read_only_tasks, + &outgoing_tx, + &mut disconnected_rx, + ) + .await + { + break; + } let Some(route) = router.notification_route(notification.method.as_str()) else { warn!( @@ -131,6 +191,7 @@ async fn run_connection( ); break; }; + let enables_read_only_requests = notification.method == INITIALIZED_METHOD; let result = tokio::select! { result = route(Arc::clone(&handler), notification) => result, _ = disconnected_rx.changed() => { @@ -144,8 +205,20 @@ async fn run_connection( warn!("closing exec-server connection after protocol error: {err}"); break; } + if enables_read_only_requests { + read_only_requests_enabled = true; + } } codex_app_server_protocol::JSONRPCMessage::Response(response) => { + if !drain_read_only_tasks( + &mut read_only_tasks, + &outgoing_tx, + &mut disconnected_rx, + ) + .await + { + break; + } warn!( "closing exec-server connection after unexpected client response: {:?}", response.id @@ -153,6 +226,15 @@ async fn run_connection( break; } codex_app_server_protocol::JSONRPCMessage::Error(error) => { + if !drain_read_only_tasks( + &mut read_only_tasks, + &outgoing_tx, + &mut disconnected_rx, + ) + .await + { + break; + } warn!( "closing exec-server connection after unexpected client error: {:?}", error.id @@ -160,99 +242,6 @@ async fn run_connection( break; } }, - JsonRpcConnectionEvent::Batch(messages) => { - if messages.is_empty() || messages.len() > MAX_RPC_BATCH_REQUESTS { - let message = if messages.is_empty() { - "JSON-RPC batch must not be empty".to_string() - } else { - format!( - "JSON-RPC batch contains {} requests; maximum is {MAX_RPC_BATCH_REQUESTS}", - messages.len() - ) - }; - if outgoing_tx - .send(RpcServerOutboundMessage::Error { - request_id: codex_app_server_protocol::RequestId::Integer(-1), - error: invalid_request(message), - }) - .await - .is_err() - { - break; - } - continue; - } - - let batch = futures::stream::iter(messages) - .map(|message| { - let router = Arc::clone(&router); - let handler = Arc::clone(&handler); - async move { - match message { - codex_app_server_protocol::JSONRPCMessage::Request(request) - if !router.is_request_batchable(request.method.as_str()) => - { - Some(RpcServerOutboundMessage::Error { - request_id: request.id, - error: invalid_request(format!( - "`{}` cannot be sent in a JSON-RPC batch", - request.method - )), - }) - } - codex_app_server_protocol::JSONRPCMessage::Request(request) => { - dispatch_request(router, handler, request).await - } - codex_app_server_protocol::JSONRPCMessage::Notification(_) => { - Some(RpcServerOutboundMessage::Error { - request_id: codex_app_server_protocol::RequestId::Integer( - -1, - ), - error: invalid_request( - "notifications cannot be sent in a JSON-RPC batch" - .to_string(), - ), - }) - } - codex_app_server_protocol::JSONRPCMessage::Response(response) => { - Some(RpcServerOutboundMessage::Error { - request_id: response.id, - error: invalid_request( - "client responses cannot be sent in a JSON-RPC batch" - .to_string(), - ), - }) - } - codex_app_server_protocol::JSONRPCMessage::Error(error) => { - Some(RpcServerOutboundMessage::Error { - request_id: error.id, - error: invalid_request( - "client errors cannot be sent in a JSON-RPC batch" - .to_string(), - ), - }) - } - } - } - }) - .buffered(MAX_RPC_BATCH_CONCURRENCY) - .filter_map(std::future::ready) - .collect::>(); - let messages = tokio::select! { - messages = batch => messages, - _ = disconnected_rx.changed() => { - debug!("exec-server transport disconnected while handling batch"); - break; - } - }; - if outgoing_tx - .send(RpcServerOutboundMessage::Batch(messages)) - .await - .is_err() - { - break; - } - } JsonRpcConnectionEvent::Disconnected { reason } => { if let Some(reason) = reason { debug!("exec-server connection disconnected: {reason}"); @@ -262,6 +251,7 @@ async fn run_connection( } } + read_only_tasks.abort_all(); handler.shutdown().await; drop(handler); drop(outgoing_tx); @@ -272,6 +262,70 @@ async fn run_connection( let _ = outbound_task.await; } +async fn limit_read_only_tasks( + tasks: &mut JoinSet>, + outgoing_tx: &mpsc::Sender, + disconnected_rx: &mut tokio::sync::watch::Receiver, +) -> bool { + while tasks.len() >= MAX_CONCURRENT_READ_ONLY_REQUESTS { + if !drain_one_read_only_task(tasks, outgoing_tx, disconnected_rx).await { + return false; + } + } + true +} + +async fn drain_read_only_tasks( + tasks: &mut JoinSet>, + outgoing_tx: &mpsc::Sender, + disconnected_rx: &mut tokio::sync::watch::Receiver, +) -> bool { + while !tasks.is_empty() { + if !drain_one_read_only_task(tasks, outgoing_tx, disconnected_rx).await { + return false; + } + } + true +} + +async fn drain_one_read_only_task( + tasks: &mut JoinSet>, + outgoing_tx: &mpsc::Sender, + disconnected_rx: &mut tokio::sync::watch::Receiver, +) -> bool { + let result = tokio::select! { + result = tasks.join_next() => result, + _ = disconnected_rx.changed() => { + debug!("exec-server transport disconnected while handling read-only request"); + return false; + } + }; + send_finished_request(result, outgoing_tx).await +} + +async fn send_finished_request( + result: Option, tokio::task::JoinError>>, + outgoing_tx: &mpsc::Sender, +) -> bool { + match result { + Some(Ok(Some(message))) => outgoing_tx.send(message).await.is_ok(), + Some(Ok(None)) | None => true, + Some(Err(err)) => { + warn!("read-only exec-server request task failed: {err}"); + true + } + } +} + +fn spawn_read_only_request( + tasks: &mut JoinSet>, + router: Arc>, + handler: Arc, + request: codex_app_server_protocol::JSONRPCRequest, +) { + tasks.spawn(dispatch_request(router, handler, request)); +} + async fn dispatch_request( router: Arc>, handler: Arc, @@ -301,7 +355,6 @@ mod tests { use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_utils_path_uri::PathUri; - use pretty_assertions::assert_eq; use serde::Serialize; use serde::de::DeserializeOwned; use tokio::io::AsyncBufReadExt; @@ -322,10 +375,6 @@ mod tests { use crate::protocol::EXEC_TERMINATE_METHOD; use crate::protocol::ExecParams; use crate::protocol::ExecResponse; - use crate::protocol::FS_READ_DIRECTORY_METHOD; - use crate::protocol::FS_WRITE_FILE_METHOD; - use crate::protocol::FsReadDirectoryParams; - use crate::protocol::FsReadDirectoryResponse; use crate::protocol::INITIALIZE_METHOD; use crate::protocol::INITIALIZED_METHOD; use crate::protocol::InitializeParams; @@ -335,82 +384,6 @@ mod tests { use crate::protocol::TerminateResponse; use crate::server::session_registry::SessionRegistry; - #[tokio::test] - async fn json_rpc_batch_allows_only_read_only_filesystem_requests() { - let registry = SessionRegistry::new(); - let (mut writer, mut lines, task) = spawn_test_connection(registry, "batch"); - - send_request( - &mut writer, - /*id*/ 1, - INITIALIZE_METHOD, - &InitializeParams { - client_name: "exec-server-test".to_string(), - resume_session_id: None, - }, - ) - .await; - let _: InitializeResponse = read_response(&mut lines, /*expected_id*/ 1).await; - send_notification(&mut writer, INITIALIZED_METHOD, &()).await; - - let cwd = PathUri::from_path(std::env::current_dir().expect("cwd")).expect("cwd URI"); - write_batch( - &mut writer, - &[ - JSONRPCMessage::Request(JSONRPCRequest { - id: RequestId::Integer(2), - method: FS_READ_DIRECTORY_METHOD.to_string(), - params: Some( - serde_json::to_value(FsReadDirectoryParams { - path: cwd, - sandbox: None, - }) - .expect("serialize fs/readDirectory params"), - ), - trace: None, - }), - JSONRPCMessage::Request(JSONRPCRequest { - id: RequestId::Integer(3), - method: FS_WRITE_FILE_METHOD.to_string(), - params: Some(serde_json::json!({})), - trace: None, - }), - ], - ) - .await; - - let messages = read_batch(&mut lines).await; - assert_eq!(messages.len(), 2); - match &messages[0] { - JSONRPCMessage::Response(JSONRPCResponse { id, result }) => { - assert_eq!(*id, RequestId::Integer(2)); - let _: FsReadDirectoryResponse = - serde_json::from_value(result.clone()).expect("decode fs/readDirectory result"); - } - other => panic!("expected fs/readDirectory response, got {other:?}"), - } - match &messages[1] { - JSONRPCMessage::Error(error) => { - assert_eq!(error.id, RequestId::Integer(3)); - assert_eq!(error.error.code, -32600); - assert!( - error - .error - .message - .contains("`fs/writeFile` cannot be sent in a JSON-RPC batch") - ); - } - other => panic!("expected fs/writeFile error, got {other:?}"), - } - - drop(writer); - drop(lines); - timeout(Duration::from_secs(1), task) - .await - .expect("processor should exit") - .expect("processor should join"); - } - #[tokio::test] async fn transport_disconnect_detaches_session_during_in_flight_read() { let registry = SessionRegistry::new(); @@ -556,12 +529,6 @@ mod tests { writer.write_all(b"\n").await.expect("write newline"); } - async fn write_batch(writer: &mut DuplexStream, messages: &[JSONRPCMessage]) { - let encoded = serde_json::to_vec(messages).expect("serialize JSON-RPC batch"); - writer.write_all(&encoded).await.expect("write batch"); - writer.write_all(b"\n").await.expect("write newline"); - } - async fn read_response( lines: &mut Lines>, expected_id: i64, @@ -581,15 +548,6 @@ mod tests { } } - async fn read_batch(lines: &mut Lines>) -> Vec { - let line = lines - .next_line() - .await - .expect("read batch") - .expect("batch line"); - serde_json::from_str(&line).expect("decode JSON-RPC batch") - } - fn exec_params(process_id: ProcessId) -> ExecParams { let mut env = HashMap::new(); if let Some(path) = std::env::var_os("PATH") { diff --git a/codex-rs/exec-server/src/server/registry.rs b/codex-rs/exec-server/src/server/registry.rs index 29dd7af6f4..891a0885de 100644 --- a/codex-rs/exec-server/src/server/registry.rs +++ b/codex-rs/exec-server/src/server/registry.rs @@ -93,7 +93,7 @@ pub(crate) fn build_router() -> RpcRouter { handler.terminate(params).await }, ); - router.batchable_request( + router.concurrent_request( FS_READ_FILE_METHOD, |handler: Arc, params: FsReadFileParams| async move { handler.fs_read_file(params).await @@ -129,19 +129,19 @@ pub(crate) fn build_router() -> RpcRouter { handler.fs_create_directory(params).await }, ); - router.batchable_request( + router.concurrent_request( FS_GET_METADATA_METHOD, |handler: Arc, params: FsGetMetadataParams| async move { handler.fs_get_metadata(params).await }, ); - router.batchable_request( + router.concurrent_request( FS_CANONICALIZE_METHOD, |handler: Arc, params: FsCanonicalizeParams| async move { handler.fs_canonicalize(params).await }, ); - router.batchable_request( + router.concurrent_request( FS_READ_DIRECTORY_METHOD, |handler: Arc, params: FsReadDirectoryParams| async move { handler.fs_read_directory(params).await