openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-05 15:18:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add exec-server JSON-RPC batching

Browse Source
This commit is contained in:
jif-oai
2026-06-19 11:58:39 +02:00
parent dac588f413
commit df2bb95976
15 changed files with 737 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
codex-rs/exec-server/src/client.rs
View File

@@ -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<String>,
rpc_batch_supported: AtomicBool,
reconnect_strategy: Option<ExecServerReconnectStrategy>,
}
@@ -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<RpcBatchCall>,
) -> Result<Vec<Result<Value, ExecServerError>>, 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<P, T>(
&self,
rpc_client: &Arc<RpcClient>,
@@ -813,6 +849,7 @@ impl From<RpcCallError> for ExecServerError {
code: error.code,
message: error.message,
},
RpcCallError::InvalidBatch(message) => Self::Protocol(message),
}
}
}

48
codex-rs/exec-server/src/connection.rs
View File

@@ -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<JSONRPCMessage>),
MalformedMessage { reason: String },
Disconnected { reason: Option<String> },
}
/// 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<JSONRPCMessage>),
}
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<JSONRPCMessage> 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<std::process::ExitStatus>
}
pub(crate) struct JsonRpcConnection {
pub(crate) outgoing_tx: mpsc::Sender<JSONRPCMessage>,
pub(crate) outgoing_tx: mpsc::Sender<JsonRpcWireMessage>,
pub(crate) incoming_rx: mpsc::Receiver<JsonRpcConnectionEvent>,
pub(crate) disconnected_rx: watch::Receiver<bool>,
pub(crate) task_handles: Vec<tokio::task::JoinHandle<()>>,
@@ -249,10 +279,10 @@ impl JsonRpcConnection {
if line.trim().is_empty() {
continue;
}
match serde_json::from_str::<JSONRPCMessage>(&line) {
match serde_json::from_str::<JsonRpcWireMessage>(&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<W>(
writer: &mut BufWriter<W>,
message: &JSONRPCMessage,
message: &JsonRpcWireMessage,
) -> std::io::Result<()>
where
W: AsyncWrite + Unpin,
@@ -564,7 +594,7 @@ where
async fn send_websocket_jsonrpc_message<W, M, E>(
websocket_writer: &mut W,
connection_label: &str,
message: &JSONRPCMessage,
message: &JsonRpcWireMessage,
) -> Result<(), String>
where
W: Sink<M, Error = E> + Unpin,
@@ -584,7 +614,7 @@ where
}
}
fn serialize_jsonrpc_message(message: &JSONRPCMessage) -> Result<String, serde_json::Error> {
fn serialize_jsonrpc_message(message: &JsonRpcWireMessage) -> Result<String, serde_json::Error> {
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!(

3
codex-rs/exec-server/src/lib.rs
View File

@@ -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;

2
codex-rs/exec-server/src/noise_relay/executor_stream.rs
View File

@@ -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(),

2
codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs
View File

@@ -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,

2
codex-rs/exec-server/src/noise_relay/harness.rs
View File

@@ -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)?;
}

9
codex-rs/exec-server/src/noise_relay/message_framing.rs
View File

@@ -1,6 +1,5 @@
use codex_app_server_protocol::JSONRPCMessage;
use crate::ExecServerError;
use crate::connection::JsonRpcWireMessage;
const LENGTH_PREFIX_BYTES: usize = size_of::<u32>();
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<Vec<u8>, ExecServerError> {
pub(crate) fn frame_jsonrpc_message(
message: &JsonRpcWireMessage,
) -> Result<Vec<u8>, 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<Vec<JSONRPCMessage>, ExecServerError> {
) -> Result<Vec<JsonRpcWireMessage>, ExecServerError> {
if plaintext_record.len() > NOISE_RECORD_PLAINTEXT_LEN {
return Err(ExecServerError::Protocol(
"Noise relay plaintext record exceeds maximum length".to_string(),

5
codex-rs/exec-server/src/noise_relay/message_framing_tests.rs
View File

@@ -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]

15
codex-rs/exec-server/src/protocol.rs
View File

@@ -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")]

15
codex-rs/exec-server/src/relay.rs
View File

@@ -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<JSONRPCMessage, ExecServerError> {
fn into_jsonrpc_message(self) -> Result<JsonRpcWireMessage, ExecServerError> {
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<Vec<u8>, ExecServerError> {
pub(crate) fn jsonrpc_payload(message: &JsonRpcWireMessage) -> Result<Vec<u8>, 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(())
}

185
codex-rs/exec-server/src/remote_file_system.rs
View File

@@ -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<FileSystemOperation>,
sandbox: Option<&FileSystemSandboxContext>,
) -> FileSystemResult<Vec<FileSystemOperationResult>> {
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::<Vec<_>>();
if output_kind_chunk.is_empty() {
break;
}
let call_chunk = calls
.by_ref()
.take(MAX_RPC_BATCH_REQUESTS)
.collect::<Vec<_>>();
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<FileSystemOperation>,
sandbox: Option<&'a FileSystemSandboxContext>,
) -> ExecutorFileSystemFuture<'a, Vec<FileSystemOperationResult>> {
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<FileSystemSandboxContext>,
) -> 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<RpcBatchCall> {
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<serde_json::Value, ExecServerError>,
) -> 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<T: DeserializeOwned>(value: serde_json::Value) -> io::Result<T> {
serde_json::from_value(value).map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))
}
fn remote_sandbox_context(

128
codex-rs/exec-server/src/rpc.rs
View File

@@ -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<Result<Value, RpcCallError>>;
@@ -63,6 +72,7 @@ pub(crate) enum RpcServerOutboundMessage {
error: JSONRPCErrorError,
},
Notification(JSONRPCNotification),
Batch(Vec<RpcServerOutboundMessage>),
}
#[derive(Clone)]
@@ -222,7 +232,7 @@ where
}
pub(crate) struct RpcClient {
write_tx: mpsc::Sender<JSONRPCMessage>,
write_tx: mpsc::Sender<JsonRpcWireMessage>,
pending: Arc<Mutex<HashMap<RequestId, PendingRequest>>>,
// 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<RpcBatchCall>,
) -> Result<Vec<Result<Value, RpcCallError>>, 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<JSONRPCMessage, serde_json::Error> {
) -> Result<JsonRpcWireMessage, serde_json::Error> {
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))
}
}
}

10
codex-rs/exec-server/src/server/handler.rs
View File

@@ -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<InitializeResponse, JSONRPCErrorError> {
) -> Result<InitializeWireResponse, JSONRPCErrorError> {
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> {

259
codex-rs/exec-server/src/server/processor.rs
View File

@@ -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<SessionRegistry>,
@@ -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::<Vec<_>>();
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<RpcRouter<ExecServerHandler>>,
handler: Arc<ExecServerHandler>,
request: codex_app_server_protocol::JSONRPCRequest,
) -> Option<RpcServerOutboundMessage> {
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<T: DeserializeOwned>(
lines: &mut Lines<BufReader<DuplexStream>>,
expected_id: i64,
@@ -389,6 +597,15 @@ mod tests {
}
}
async fn read_batch(lines: &mut Lines<BufReader<DuplexStream>>) -> Vec<JSONRPCMessage> {
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") {

80
codex-rs/file-system/src/lib.rs
View File

@@ -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<u8>),
GetMetadata(FileMetadata),
ReadDirectory(Vec<ReadDirectoryEntry>),
}
pub type FileSystemOperationResult = FileSystemResult<FileSystemOperationOutput>;
#[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<FileSystemOperation>,
sandbox: Option<&'a FileSystemSandboxContext>,
) -> ExecutorFileSystemFuture<'a, Vec<FileSystemOperationResult>> {
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<F>(
file_system: &F,
operations: Vec<FileSystemOperation>,
sandbox: Option<&FileSystemSandboxContext>,
) -> FileSystemResult<Vec<FileSystemOperationResult>>
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<F>(
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),
}
}