mirror of
https://github.com/openai/codex.git
synced 2026-10-30 19:18:49 +00:00
Use generic RPC batch requests
This commit is contained in:
Generated
+1
@@ -3049,6 +3049,7 @@ dependencies = [
|
||||
"codex-utils-path-uri",
|
||||
"futures",
|
||||
"serde",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(())
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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
@@ -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))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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") {
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user