Use generic RPC batch requests

This commit is contained in:
jif-oai
2026-06-19 18:30:54 +02:00
parent d761119c24
commit baa2809950
17 changed files with 458 additions and 345 deletions
+1
View File
@@ -3049,6 +3049,7 @@ dependencies = [
"codex-utils-path-uri",
"futures",
"serde",
"serde_json",
]
[[package]]
+25
View File
@@ -97,6 +97,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;
@@ -761,6 +762,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 +837,7 @@ impl From<RpcCallError> for ExecServerError {
code: error.code,
message: error.message,
},
RpcCallError::InvalidBatch(message) => Self::Protocol(message),
}
}
}
+38 -10
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,33 @@ 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> },
}
#[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 +247,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 +275,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 +421,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 +490,7 @@ impl JsonRpcConnection {
}
enum JsonRpcWebSocketFrame {
Message(JSONRPCMessage),
Message(JsonRpcWireMessage),
Close,
Ignore,
}
@@ -549,7 +575,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 +590,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 +610,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 +720,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!(
@@ -706,7 +732,9 @@ mod tests {
control.set_write_ready();
assert!(matches!(
timeout(Duration::from_secs(1), outbound_rx.next()).await?,
Some(Message::Text(text)) if serde_json::from_str::<JSONRPCMessage>(&text)? == message
Some(Message::Text(text))
if serde_json::from_str::<JsonRpcWireMessage>(&text)?
== JsonRpcWireMessage::Single(message)
));
drop(connection);
Ok(())
+2 -3
View File
@@ -42,11 +42,10 @@ pub use codex_file_system::CopyOptions;
pub use codex_file_system::CreateDirectoryOptions;
pub use codex_file_system::ExecutorFileSystem;
pub use codex_file_system::ExecutorFileSystemFuture;
pub use codex_file_system::ExecutorRpcBatchCall;
pub use codex_file_system::ExecutorRpcBatchResult;
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;
@@ -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(),
@@ -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,
@@ -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)?;
}
@@ -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(),
@@ -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]
+11 -7
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,10 @@ 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()?,
JsonRpcWireMessage::Single(message)
);
drop(connection);
Ok(())
}
@@ -9,6 +9,8 @@ use crate::CreateDirectoryOptions;
use crate::ExecServerError;
use crate::ExecutorFileSystem;
use crate::ExecutorFileSystemFuture;
use crate::ExecutorRpcBatchCall;
use crate::ExecutorRpcBatchResult;
use crate::FileMetadata;
use crate::FileSystemReadStream;
use crate::FileSystemResult;
@@ -16,6 +18,7 @@ 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::FsCopyParams;
use crate::protocol::FsCreateDirectoryParams;
@@ -24,6 +27,7 @@ use crate::protocol::FsReadDirectoryParams;
use crate::protocol::FsReadFileParams;
use crate::protocol::FsRemoveParams;
use crate::protocol::FsWriteFileParams;
use crate::rpc::RpcBatchCall;
const INVALID_REQUEST_ERROR_CODE: i64 = -32600;
const NOT_FOUND_ERROR_CODE: i64 = -32004;
@@ -223,6 +227,43 @@ impl RemoteFileSystem {
.map_err(map_remote_error)?;
Ok(())
}
async fn execute_rpc_batch(
&self,
calls: Vec<ExecutorRpcBatchCall>,
) -> FileSystemResult<Vec<ExecutorRpcBatchResult>> {
if calls.is_empty() {
return Ok(Vec::new());
}
let client = self.client.get().await.map_err(map_remote_error)?;
let mut decoded = Vec::with_capacity(calls.len());
let mut calls = calls.into_iter();
loop {
let call_chunk = calls
.by_ref()
.take(MAX_RPC_BATCH_REQUESTS)
.map(|call| RpcBatchCall {
method: call.method,
params: call.params,
})
.collect::<Vec<_>>();
if call_chunk.is_empty() {
break;
}
let results = client
.call_batch(call_chunk)
.await
.map_err(map_remote_error)?;
decoded.extend(
results
.into_iter()
.map(|result| result.map_err(map_remote_error)),
);
}
Ok(decoded)
}
}
impl ExecutorFileSystem for RemoteFileSystem {
@@ -310,6 +351,13 @@ impl ExecutorFileSystem for RemoteFileSystem {
sandbox,
))
}
fn execute_rpc_batch<'a>(
&'a self,
calls: Vec<ExecutorRpcBatchCall>,
) -> ExecutorFileSystemFuture<'a, Vec<ExecutorRpcBatchResult>> {
Box::pin(RemoteFileSystem::execute_rpc_batch(self, calls))
}
}
fn remote_sandbox_context(
+118 -36
View File
@@ -1,5 +1,4 @@
use std::collections::HashMap;
use std::collections::HashSet;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
@@ -26,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;
@@ -37,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>>;
@@ -64,6 +72,7 @@ pub(crate) enum RpcServerOutboundMessage {
error: JSONRPCErrorError,
},
Notification(JSONRPCNotification),
Batch(Vec<RpcServerOutboundMessage>),
}
#[derive(Clone)]
@@ -107,7 +116,6 @@ impl RpcNotificationSender {
pub(crate) struct RpcRouter<S> {
request_routes: HashMap<&'static str, RequestRoute<S>>,
concurrent_request_methods: HashSet<&'static str>,
notification_routes: HashMap<&'static str, NotificationRoute<S>>,
}
@@ -115,7 +123,6 @@ impl<S> Default for RpcRouter<S> {
fn default() -> Self {
Self {
request_routes: HashMap::new(),
concurrent_request_methods: HashSet::new(),
notification_routes: HashMap::new(),
}
}
@@ -136,7 +143,6 @@ where
F: Fn(Arc<S>, P) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<R, JSONRPCErrorError>> + Send + 'static,
{
self.concurrent_request_methods.remove(method);
self.request_routes.insert(
method,
Box::new(move |state, request| {
@@ -166,24 +172,12 @@ where
);
}
pub(crate) fn concurrent_request<P, R, F, Fut>(&mut self, method: &'static str, handler: F)
where
P: DeserializeOwned + Send + 'static,
R: Serialize + Send + 'static,
F: Fn(Arc<S>, P) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<R, JSONRPCErrorError>> + Send + 'static,
{
self.request(method, handler);
self.concurrent_request_methods.insert(method);
}
pub(crate) fn request_with_id<P, F, Fut>(&mut self, method: &'static str, handler: F)
where
P: DeserializeOwned + Send + 'static,
F: Fn(Arc<S>, RequestId, P) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<(), JSONRPCErrorError>> + Send + 'static,
{
self.concurrent_request_methods.remove(method);
self.request_routes.insert(
method,
Box::new(move |state, request| {
@@ -232,17 +226,13 @@ where
self.request_routes.get(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<S>> {
self.notification_routes.get(method)
}
}
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
@@ -272,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;
};
@@ -285,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;
@@ -330,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)
}
@@ -378,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()
{
@@ -406,6 +412,70 @@ impl RpcClient {
serde_json::from_value(response).map_err(RpcCallError::Json)
}
/// Sends independent RPC requests as one JSON-RPC batch.
///
/// The server may execute batch entries concurrently and does not guarantee execution order.
/// Callers must only batch requests that are safe to run in any order, and must send dependent
/// requests separately. Results are returned in the same order as `calls` by matching request
/// ids in the responses.
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()
@@ -424,22 +494,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))
}
}
}
+181 -152
View File
@@ -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::protocol::INITIALIZED_METHOD;
use crate::connection::MAX_RPC_BATCH_REQUESTS;
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_CONCURRENT_READ_ONLY_REQUESTS: usize = 32;
const MAX_RPC_BATCH_CONCURRENCY: usize = 32;
#[derive(Clone)]
pub(crate) struct ConnectionProcessor {
@@ -67,8 +67,6 @@ 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 {
@@ -85,33 +83,13 @@ async fn run_connection(
}
});
// Process inbound events sequentially to preserve initialize/initialized ordering.
loop {
while let Some(event) = incoming_rx.recv().await {
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 {
@@ -126,36 +104,6 @@ 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),
@@ -174,15 +122,6 @@ 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!(
@@ -191,7 +130,6 @@ 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() => {
@@ -205,20 +143,8 @@ 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
@@ -226,15 +152,6 @@ 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
@@ -242,6 +159,88 @@ 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) => {
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}");
@@ -251,7 +250,6 @@ async fn run_connection(
}
}
read_only_tasks.abort_all();
handler.shutdown().await;
drop(handler);
drop(outgoing_tx);
@@ -262,70 +260,6 @@ async fn run_connection(
let _ = outbound_task.await;
}
async fn limit_read_only_tasks(
tasks: &mut JoinSet<Option<RpcServerOutboundMessage>>,
outgoing_tx: &mpsc::Sender<RpcServerOutboundMessage>,
disconnected_rx: &mut tokio::sync::watch::Receiver<bool>,
) -> 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<Option<RpcServerOutboundMessage>>,
outgoing_tx: &mpsc::Sender<RpcServerOutboundMessage>,
disconnected_rx: &mut tokio::sync::watch::Receiver<bool>,
) -> 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<Option<RpcServerOutboundMessage>>,
outgoing_tx: &mpsc::Sender<RpcServerOutboundMessage>,
disconnected_rx: &mut tokio::sync::watch::Receiver<bool>,
) -> 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<Result<Option<RpcServerOutboundMessage>, tokio::task::JoinError>>,
outgoing_tx: &mpsc::Sender<RpcServerOutboundMessage>,
) -> 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<Option<RpcServerOutboundMessage>>,
router: Arc<RpcRouter<ExecServerHandler>>,
handler: Arc<ExecServerHandler>,
request: codex_app_server_protocol::JSONRPCRequest,
) {
tasks.spawn(dispatch_request(router, handler, request));
}
async fn dispatch_request(
router: Arc<RpcRouter<ExecServerHandler>>,
handler: Arc<ExecServerHandler>,
@@ -370,11 +304,16 @@ mod tests {
use crate::ExecServerRuntimePaths;
use crate::ProcessId;
use crate::connection::JsonRpcConnection;
use crate::protocol::ENVIRONMENT_INFO_METHOD;
use crate::protocol::EXEC_METHOD;
use crate::protocol::EXEC_READ_METHOD;
use crate::protocol::EXEC_TERMINATE_METHOD;
use crate::protocol::EnvironmentInfo;
use crate::protocol::ExecParams;
use crate::protocol::ExecResponse;
use crate::protocol::FS_READ_DIRECTORY_METHOD;
use crate::protocol::FsReadDirectoryParams;
use crate::protocol::FsReadDirectoryResponse;
use crate::protocol::INITIALIZE_METHOD;
use crate::protocol::INITIALIZED_METHOD;
use crate::protocol::InitializeParams;
@@ -384,6 +323,81 @@ mod tests {
use crate::protocol::TerminateResponse;
use crate::server::session_registry::SessionRegistry;
#[tokio::test]
async fn json_rpc_batch_dispatches_generic_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: ENVIRONMENT_INFO_METHOD.to_string(),
params: Some(serde_json::to_value(()).expect("serialize params")),
trace: None,
}),
JSONRPCMessage::Request(JSONRPCRequest {
id: RequestId::Integer(3),
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,
}),
],
)
.await;
let messages = read_batch(&mut lines).await;
let mut responses = HashMap::new();
for message in messages {
match message {
JSONRPCMessage::Response(JSONRPCResponse { id, result }) => {
responses.insert(id, result);
}
JSONRPCMessage::Error(error) => panic!("unexpected JSON-RPC error: {error:?}"),
other => panic!("expected JSON-RPC response, got {other:?}"),
}
}
let environment_info = responses
.remove(&RequestId::Integer(2))
.expect("environment/info response");
let _: EnvironmentInfo =
serde_json::from_value(environment_info).expect("decode environment/info response");
let read_directory = responses
.remove(&RequestId::Integer(3))
.expect("fs/readDirectory response");
let _: FsReadDirectoryResponse =
serde_json::from_value(read_directory).expect("decode fs/readDirectory response");
assert!(responses.is_empty());
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();
@@ -529,6 +543,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,
@@ -548,6 +568,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") {
+4 -4
View File
@@ -93,7 +93,7 @@ pub(crate) fn build_router() -> RpcRouter<ExecServerHandler> {
handler.terminate(params).await
},
);
router.concurrent_request(
router.request(
FS_READ_FILE_METHOD,
|handler: Arc<ExecServerHandler>, params: FsReadFileParams| async move {
handler.fs_read_file(params).await
@@ -129,19 +129,19 @@ pub(crate) fn build_router() -> RpcRouter<ExecServerHandler> {
handler.fs_create_directory(params).await
},
);
router.concurrent_request(
router.request(
FS_GET_METADATA_METHOD,
|handler: Arc<ExecServerHandler>, params: FsGetMetadataParams| async move {
handler.fs_get_metadata(params).await
},
);
router.concurrent_request(
router.request(
FS_CANONICALIZE_METHOD,
|handler: Arc<ExecServerHandler>, params: FsCanonicalizeParams| async move {
handler.fs_canonicalize(params).await
},
);
router.concurrent_request(
router.request(
FS_READ_DIRECTORY_METHOD,
|handler: Arc<ExecServerHandler>, params: FsReadDirectoryParams| async move {
handler.fs_read_directory(params).await
@@ -4,8 +4,6 @@ use codex_exec_server::CopyOptions;
use codex_exec_server::CreateDirectoryOptions;
use codex_exec_server::FILE_READ_CHUNK_SIZE;
use codex_exec_server::FileMetadata;
use codex_exec_server::FileSystemOperation;
use codex_exec_server::FileSystemOperationOutput;
use codex_exec_server::ReadDirectoryEntry;
use codex_exec_server::RemoveOptions;
use codex_protocol::models::AdditionalPermissionProfile;
@@ -198,55 +196,6 @@ async fn file_system_read_file_returns_bytes(
Ok(())
}
#[test_case(FileSystemImplementation::Local ; "local")]
#[test_case(FileSystemImplementation::Remote ; "remote")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn file_system_batch_preserves_operation_order(
implementation: FileSystemImplementation,
) -> Result<()> {
let context = create_file_system_context(implementation).await?;
let file_system = context.file_system;
let tmp = TempDir::new()?;
let first_path = tmp.path().join("first.txt");
let second_path = tmp.path().join("second.txt");
std::fs::write(&first_path, "first")?;
std::fs::write(&second_path, "second")?;
let first_uri = PathUri::from_path(&first_path)?;
let second_uri = PathUri::from_path(&second_path)?;
let outputs = file_system
.execute_batch(
vec![
FileSystemOperation::ReadFile {
path: first_uri.clone(),
},
FileSystemOperation::Canonicalize {
path: second_uri.clone(),
},
FileSystemOperation::ReadFile { path: second_uri },
],
/*sandbox*/ None,
)
.await
.with_context(|| format!("mode={implementation}"))?
.into_iter()
.collect::<std::io::Result<Vec<_>>>()?;
assert_eq!(
outputs,
vec![
FileSystemOperationOutput::ReadFile(b"first".to_vec()),
FileSystemOperationOutput::Canonicalize(PathUri::from_path(std::fs::canonicalize(
second_path
)?)?),
FileSystemOperationOutput::ReadFile(b"second".to_vec()),
]
);
Ok(())
}
#[test_case(FileSystemImplementation::Local ; "local")]
#[test_case(FileSystemImplementation::Remote ; "remote")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+1
View File
@@ -14,6 +14,7 @@ codex-utils-absolute-path = { workspace = true }
codex-utils-path-uri = { workspace = true }
futures = { workspace = true }
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
[lib]
test = false
+18 -73
View File
@@ -12,7 +12,6 @@ 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;
@@ -22,7 +21,6 @@ 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 {
@@ -58,25 +56,13 @@ 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 },
#[derive(Clone, Debug, PartialEq)]
pub struct ExecutorRpcBatchCall {
pub method: String,
pub params: serde_json::Value,
}
/// 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>;
pub type ExecutorRpcBatchResult = FileSystemResult<serde_json::Value>;
#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
@@ -286,61 +272,20 @@ pub trait ExecutorFileSystem: Send + Sync {
sandbox: Option<&'a FileSystemSandboxContext>,
) -> ExecutorFileSystemFuture<'a, ()>;
/// Executes independent read-only filesystem operations concurrently.
/// Sends independent executor RPC requests as one batch.
///
/// Results retain request order. Operations have no ordering or transactional relationship;
/// callers must put dependent operations in separate batches.
fn execute_batch<'a>(
/// The executor may execute entries concurrently and does not guarantee execution order.
/// Callers must only batch requests that are safe to run in any order, and must send dependent
/// requests separately. Results are matched back to the input order by request id.
fn execute_rpc_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),
_calls: Vec<ExecutorRpcBatchCall>,
) -> ExecutorFileSystemFuture<'a, Vec<ExecutorRpcBatchResult>> {
Box::pin(async move {
Err(io::Error::new(
io::ErrorKind::Unsupported,
"filesystem does not support RPC batching",
))
})
}
}