diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 88c0c80342..31f2249d85 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -87,6 +87,7 @@ 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; @@ -97,6 +98,7 @@ 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; @@ -199,6 +201,7 @@ struct Inner { http_body_streams_write_lock: Mutex<()>, http_body_stream_next_id: AtomicU64, session_id: OnceLock, + rpc_batch_supported: AtomicBool, reconnect_strategy: Option, } @@ -471,7 +474,7 @@ impl ExecServerClient { } = options; timeout(initialize_timeout, async { - let response: InitializeResponse = rpc_client + let response: InitializeWireResponse = rpc_client .call( INITIALIZE_METHOD, &InitializeParams { @@ -484,6 +487,9 @@ 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 {}", @@ -493,7 +499,9 @@ impl ExecServerClient { rpc_client .notify(INITIALIZED_METHOD, &serde_json::json!({})) .await?; - Ok(response) + Ok(InitializeResponse { + session_id: response.session_id, + }) }) .await .map_err(|_| ExecServerError::InitializeTimedOut { @@ -708,6 +716,10 @@ 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() } @@ -741,6 +753,7 @@ 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 }; @@ -761,6 +774,29 @@ 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, @@ -813,6 +849,7 @@ 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 2d2e5afe06..255eb26ec8 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -12,6 +12,8 @@ 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; @@ -29,6 +31,7 @@ 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); @@ -38,10 +41,37 @@ 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. @@ -221,7 +251,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>, @@ -249,10 +279,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(JsonRpcConnectionEvent::Message(message)) + .send(message.into_connection_event()) .await .is_err() { @@ -395,7 +425,7 @@ impl JsonRpcConnection { Some(Ok(message)) => match message.parse_jsonrpc_frame() { Ok(JsonRpcWebSocketFrame::Message(message)) => { if incoming_tx - .send(JsonRpcConnectionEvent::Message(message)) + .send(message.into_connection_event()) .await .is_err() { @@ -464,7 +494,7 @@ impl JsonRpcConnection { } enum JsonRpcWebSocketFrame { - Message(JSONRPCMessage), + Message(JsonRpcWireMessage), Close, Ignore, } @@ -549,7 +579,7 @@ async fn send_malformed_message( async fn write_jsonrpc_line_message( writer: &mut BufWriter, - message: &JSONRPCMessage, + message: &JsonRpcWireMessage, ) -> std::io::Result<()> where W: AsyncWrite + Unpin, @@ -564,7 +594,7 @@ where async fn send_websocket_jsonrpc_message( websocket_writer: &mut W, connection_label: &str, - message: &JSONRPCMessage, + message: &JsonRpcWireMessage, ) -> Result<(), String> where W: Sink + Unpin, @@ -584,7 +614,7 @@ where } } -fn serialize_jsonrpc_message(message: &JSONRPCMessage) -> Result { +fn serialize_jsonrpc_message(message: &JsonRpcWireMessage) -> Result { serde_json::to_string(message) } @@ -694,7 +724,7 @@ mod tests { ); let message = test_jsonrpc_message(); - connection.outgoing_tx.send(message.clone()).await?; + connection.outgoing_tx.send(message.clone().into()).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/lib.rs b/codex-rs/exec-server/src/lib.rs index c96c443447..d66d2f9b7e 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -44,6 +44,9 @@ pub use codex_file_system::ExecutorFileSystem; pub use codex_file_system::ExecutorFileSystemFuture; pub use codex_file_system::FILE_READ_CHUNK_SIZE; pub use codex_file_system::FileMetadata; +pub use codex_file_system::FileSystemOperation; +pub use codex_file_system::FileSystemOperationOutput; +pub use codex_file_system::FileSystemOperationResult; pub use codex_file_system::FileSystemReadStream; pub use codex_file_system::FileSystemResult; pub use codex_file_system::FileSystemSandboxContext; 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 57a05579dc..04f25dc446 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(JsonRpcConnectionEvent::Message(message)) + .try_send(message.into_connection_event()) .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 9119236cfd..d757a12393 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)?)?; + let ciphertext = harness_transport.encrypt(&frame_jsonrpc_message(&message.clone().into())?)?; 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 1989525224..62e82e608c 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(JsonRpcConnectionEvent::Message(message)) + .send(message.into_connection_event()) .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 24ee48173e..b84253695c 100644 --- a/codex-rs/exec-server/src/noise_relay/message_framing.rs +++ b/codex-rs/exec-server/src/noise_relay/message_framing.rs @@ -1,6 +1,5 @@ -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; @@ -12,7 +11,9 @@ 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: &JSONRPCMessage) -> Result, ExecServerError> { +pub(crate) fn frame_jsonrpc_message( + message: &JsonRpcWireMessage, +) -> 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; @@ -39,7 +40,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 c3df14bf56..5cc3ed10f1 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,6 +7,7 @@ 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() { @@ -16,7 +17,7 @@ fn fragments_and_reassembles_large_jsonrpc_message() { "data": "x".repeat(128 * 1024), })), }); - let framed = frame_jsonrpc_message(&message).unwrap(); + let framed = frame_jsonrpc_message(&message.clone().into()).unwrap(); assert!(framed.len() > 128 * 1024); let mut decoder = JsonRpcMessageDecoder::default(); @@ -25,7 +26,7 @@ fn fragments_and_reassembles_large_jsonrpc_message() { decoded.extend(decoder.push(record).unwrap()); } - assert_eq!(decoded, vec![message]); + assert_eq!(decoded, vec![JsonRpcWireMessage::Single(message)]); } #[test] diff --git a/codex-rs/exec-server/src/protocol.rs b/codex-rs/exec-server/src/protocol.rs index 97ef86aad2..cde670a2ed 100644 --- a/codex-rs/exec-server/src/protocol.rs +++ b/codex-rs/exec-server/src/protocol.rs @@ -66,6 +66,21 @@ 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 033da4b2a7..ca467d0d0f 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -1,7 +1,6 @@ use std::collections::HashMap; use std::time::Duration; -use codex_app_server_protocol::JSONRPCMessage; use futures::Sink; use futures::SinkExt; use futures::Stream; @@ -25,6 +24,7 @@ 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: &JSONRPCMessage) -> Result, ExecServerError> { +pub(crate) fn jsonrpc_payload(message: &JsonRpcWireMessage) -> Result, ExecServerError> { serde_json::to_vec(message).map_err(ExecServerError::Json) } @@ -346,7 +346,7 @@ where &mut websocket, &mut keepalive, &incoming_tx, - JsonRpcConnectionEvent::Message(message), + message.into_connection_event(), ) .await { @@ -847,6 +847,7 @@ 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; @@ -880,7 +881,7 @@ mod tests { encode_relay_message_frame(&RelayMessageFrame::data( stream_id, /*seq*/ 0, - jsonrpc_payload(&message)?, + jsonrpc_payload(&message.clone().into())?, )) .into(), )) @@ -1029,7 +1030,7 @@ mod tests { let message = test_jsonrpc_message(); control.set_write_blocked(); - connection.outgoing_tx.send(message.clone()).await?; + connection.outgoing_tx.send(message.clone().into()).await?; control.wait_for_blocked_write().await?; control.send_inbound(Message::Pong(b"check".to_vec().into()))?; assert!( @@ -1047,7 +1048,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); + assert_eq!(frame.into_jsonrpc_message()?, message.into()); 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 f9136e1461..eadabba5d9 100644 --- a/codex-rs/exec-server/src/remote_file_system.rs +++ b/codex-rs/exec-server/src/remote_file_system.rs @@ -10,20 +10,31 @@ 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; @@ -223,6 +234,68 @@ 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 { @@ -310,6 +383,118 @@ 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 82cf2b200a..a6d3df0fef 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -25,6 +25,8 @@ 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; @@ -36,6 +38,13 @@ 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>; @@ -63,6 +72,7 @@ pub(crate) enum RpcServerOutboundMessage { error: JSONRPCErrorError, }, Notification(JSONRPCNotification), + Batch(Vec), } #[derive(Clone)] @@ -222,7 +232,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 @@ -252,7 +262,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 = loop { + let disconnect_reason = 'reader: loop { let Some(event) = incoming_rx.recv().await else { break None; }; @@ -265,6 +275,16 @@ 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; @@ -310,10 +330,13 @@ impl RpcClient { return Err(RpcCallError::Closed); } self.write_tx - .send(JSONRPCMessage::Notification(JSONRPCNotification { - method: method.to_string(), - params: Some(params), - })) + .send( + JSONRPCMessage::Notification(JSONRPCNotification { + method: method.to_string(), + params: Some(params), + }) + .into(), + ) .await .map_err(|_| RpcCallError::Closed) } @@ -358,12 +381,15 @@ impl RpcClient { }; if self .write_tx - .send(JSONRPCMessage::Request(JSONRPCRequest { - id: request_id.clone(), - method: method.to_string(), - params: Some(params), - trace: None, - })) + .send( + JSONRPCMessage::Request(JSONRPCRequest { + id: request_id.clone(), + method: method.to_string(), + params: Some(params), + trace: None, + }) + .into(), + ) .await .is_err() { @@ -386,6 +412,64 @@ 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() @@ -404,22 +488,34 @@ 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)) + 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)) } } } diff --git a/codex-rs/exec-server/src/server/handler.rs b/codex-rs/exec-server/src/server/handler.rs index e0ede94c77..bb13e37847 100644 --- a/codex-rs/exec-server/src/server/handler.rs +++ b/codex-rs/exec-server/src/server/handler.rs @@ -17,6 +17,7 @@ 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; @@ -41,7 +42,7 @@ use crate::protocol::FsWriteFileParams; use crate::protocol::FsWriteFileResponse; use crate::protocol::HttpRequestParams; use crate::protocol::InitializeParams; -use crate::protocol::InitializeResponse; +use crate::protocol::InitializeWireResponse; use crate::protocol::ReadParams; use crate::protocol::ReadResponse; use crate::protocol::SignalParams; @@ -107,7 +108,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(), @@ -135,7 +136,10 @@ impl ExecServerHandler { .session .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(session); - Ok(InitializeResponse { session_id }) + Ok(InitializeWireResponse { + session_id, + capabilities: ExecServerCapabilities { rpc_batch: true }, + }) } 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 952f414213..d1ccd57463 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -1,5 +1,6 @@ use std::sync::Arc; +use futures::StreamExt; use tokio::sync::mpsc; use tracing::debug; use tracing::warn; @@ -8,7 +9,13 @@ 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::FS_CANONICALIZE_METHOD; +use crate::protocol::FS_GET_METADATA_METHOD; +use crate::protocol::FS_READ_DIRECTORY_METHOD; +use crate::protocol::FS_READ_FILE_METHOD; use crate::rpc::RpcNotificationSender; +use crate::rpc::RpcRouter; use crate::rpc::RpcServerOutboundMessage; use crate::rpc::encode_server_message; use crate::rpc::invalid_request; @@ -17,6 +24,8 @@ use crate::server::ExecServerHandler; use crate::server::registry::build_router; use crate::server::session_registry::SessionRegistry; +const MAX_RPC_BATCH_CONCURRENCY: usize = 32; + #[derive(Clone)] pub(crate) struct ConnectionProcessor { session_registry: Arc, @@ -100,29 +109,19 @@ async fn run_connection( } JsonRpcConnectionEvent::Message(message) => match message { codex_app_server_protocol::JSONRPCMessage::Request(request) => { - if let Some(route) = router.request_route(request.method.as_str()) { - let message = tokio::select! { - message = route(Arc::clone(&handler), request) => message, - _ = disconnected_rx.changed() => { - debug!("exec-server transport disconnected while handling request"); - break; - } - }; - if let Some(message) = message - && outgoing_tx.send(message).await.is_err() - { + let message = tokio::select! { + message = dispatch_request( + Arc::clone(&router), + Arc::clone(&handler), + request, + ) => message, + _ = disconnected_rx.changed() => { + debug!("exec-server transport disconnected while handling request"); break; } - } else if outgoing_tx - .send(RpcServerOutboundMessage::Error { - request_id: request.id, - error: method_not_found(format!( - "exec-server stub does not implement `{}` yet", - request.method - )), - }) - .await - .is_err() + }; + if let Some(message) = message + && outgoing_tx.send(message).await.is_err() { break; } @@ -165,6 +164,99 @@ 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 !is_batchable_request_method(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}"); @@ -184,6 +276,35 @@ async fn run_connection( let _ = outbound_task.await; } +fn is_batchable_request_method(method: &str) -> bool { + // Batch handling is only for remote skill discovery lookups. Keep it read-only so concurrent + // execution cannot reorder mutations, process I/O, HTTP side effects, or file handle lifetimes. + matches!( + method, + FS_CANONICALIZE_METHOD + | FS_GET_METADATA_METHOD + | FS_READ_DIRECTORY_METHOD + | FS_READ_FILE_METHOD + ) +} + +async fn dispatch_request( + router: Arc>, + handler: Arc, + request: codex_app_server_protocol::JSONRPCRequest, +) -> Option { + let Some(route) = router.request_route(request.method.as_str()) else { + return Some(RpcServerOutboundMessage::Error { + request_id: request.id, + error: method_not_found(format!( + "exec-server stub does not implement `{}` yet", + request.method + )), + }); + }; + route(handler, request).await +} + #[cfg(test)] mod tests { use std::collections::HashMap; @@ -196,6 +317,7 @@ 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; @@ -216,6 +338,10 @@ 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; @@ -225,6 +351,82 @@ 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(); @@ -370,6 +572,12 @@ 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, @@ -389,6 +597,15 @@ 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/file-system/src/lib.rs b/codex-rs/file-system/src/lib.rs index 3ba7874404..1fd1c3ebf3 100644 --- a/codex-rs/file-system/src/lib.rs +++ b/codex-rs/file-system/src/lib.rs @@ -12,6 +12,7 @@ use codex_protocol::protocol::SandboxPolicy; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; use futures::Stream; +use futures::StreamExt; use std::future::Future; use std::io; use std::path::Path; @@ -21,6 +22,7 @@ use std::task::Poll; /// Maximum chunk size returned by [`ExecutorFileSystem::read_file_stream`]. pub const FILE_READ_CHUNK_SIZE: usize = 1024 * 1024; +const FILE_SYSTEM_BATCH_CONCURRENCY: usize = 32; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub struct CreateDirectoryOptions { @@ -56,6 +58,26 @@ pub struct ReadDirectoryEntry { pub is_file: bool, } +/// One independent, read-only filesystem operation eligible for batching. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum FileSystemOperation { + Canonicalize { path: PathUri }, + ReadFile { path: PathUri }, + GetMetadata { path: PathUri }, + ReadDirectory { path: PathUri }, +} + +/// Typed output for a [`FileSystemOperation`]. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum FileSystemOperationOutput { + Canonicalize(PathUri), + ReadFile(Vec), + GetMetadata(FileMetadata), + ReadDirectory(Vec), +} + +pub type FileSystemOperationResult = FileSystemResult; + #[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)] #[serde(rename_all = "camelCase")] pub struct FileSystemSandboxContext { @@ -263,4 +285,62 @@ pub trait ExecutorFileSystem: Send + Sync { copy_options: CopyOptions, sandbox: Option<&'a FileSystemSandboxContext>, ) -> ExecutorFileSystemFuture<'a, ()>; + + /// Executes independent read-only filesystem operations concurrently. + /// + /// Results retain request order. Operations have no ordering or transactional relationship; + /// callers must put dependent operations in separate batches. + fn execute_batch<'a>( + &'a self, + operations: Vec, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, Vec> { + Box::pin(execute_batch_with_scalar_operations( + self, operations, sandbox, + )) + } +} + +/// Default batch implementation used by local filesystems and by remote fallbacks. +pub async fn execute_batch_with_scalar_operations( + file_system: &F, + operations: Vec, + sandbox: Option<&FileSystemSandboxContext>, +) -> FileSystemResult> +where + F: ExecutorFileSystem + ?Sized, +{ + Ok(futures::stream::iter(operations) + .map(|operation| execute_scalar_operation(file_system, operation, sandbox)) + .buffered(FILE_SYSTEM_BATCH_CONCURRENCY) + .collect() + .await) +} + +async fn execute_scalar_operation( + file_system: &F, + operation: FileSystemOperation, + sandbox: Option<&FileSystemSandboxContext>, +) -> FileSystemOperationResult +where + F: ExecutorFileSystem + ?Sized, +{ + match operation { + FileSystemOperation::Canonicalize { path } => file_system + .canonicalize(&path, sandbox) + .await + .map(FileSystemOperationOutput::Canonicalize), + FileSystemOperation::ReadFile { path } => file_system + .read_file(&path, sandbox) + .await + .map(FileSystemOperationOutput::ReadFile), + FileSystemOperation::GetMetadata { path } => file_system + .get_metadata(&path, sandbox) + .await + .map(FileSystemOperationOutput::GetMetadata), + FileSystemOperation::ReadDirectory { path } => file_system + .read_directory(&path, sandbox) + .await + .map(FileSystemOperationOutput::ReadDirectory), + } }