diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 5590fd3b5eb5..2e2f4ce450e1 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -561,7 +561,7 @@ impl ExecServerClient { &self, process_id: &ProcessId, ) -> Result { - self.call( + self.call_for_cleanup( EXEC_TERMINATE_METHOD, &TerminateParams { process_id: process_id.clone(), @@ -592,7 +592,7 @@ impl ExecServerClient { &self, params: FsCloseParams, ) -> Result { - self.call(FS_CLOSE_METHOD, ¶ms).await + self.call_for_cleanup(FS_CLOSE_METHOD, ¶ms).await } pub async fn fs_write_file( @@ -788,6 +788,15 @@ impl ExecServerClient { { map_rpc_call_result(rpc_client.call(method, params).await) } + + async fn call_for_cleanup(&self, method: &str, params: &P) -> Result + where + P: serde::Serialize, + T: serde::de::DeserializeOwned, + { + let rpc_client = self.inner.rpc_client().await?; + map_rpc_call_result(rpc_client.call_for_cleanup(method, params).await) + } } fn map_rpc_call_result(result: Result) -> Result { @@ -830,6 +839,9 @@ impl From for ExecServerError { RpcCallError::TimedOut { method, timeout } => Self::Protocol(format!( "timed out waiting for exec-server `{method}` response after {timeout:?}" )), + RpcCallError::PendingRequestLimitExceeded { limit } => Self::Protocol(format!( + "exec-server has reached its limit of {limit} pending requests" + )), } } } diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index 252948cfc1da..54ccf05046bd 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/codex-rs/exec-server/src/client_recovery.rs @@ -445,7 +445,7 @@ impl Inner { Ok(false) => {} Err(error) => { let terminated: Result = rpc_client - .call( + .call_for_cleanup( EXEC_TERMINATE_METHOD, &TerminateParams { process_id: process_id.clone(), diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index e762fe8b66c3..9671f25d3933 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -18,6 +18,8 @@ use serde::Serialize; use serde::de::DeserializeOwned; use serde_json::Value; use tokio::sync::Mutex; +use tokio::sync::Semaphore; +use tokio::sync::SemaphorePermit; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::sync::watch; @@ -29,6 +31,8 @@ use crate::connection::JsonRpcConnectionEvent; use crate::connection::JsonRpcTransport; pub(crate) const SESSION_ALREADY_ATTACHED_ERROR_CODE: i64 = -32010; +const MAX_IN_FLIGHT_REGULAR_CALLS: usize = 1024; +const RESERVED_CLEANUP_CALLS: usize = 1; #[derive(Debug)] pub(crate) enum RpcCallError { @@ -40,6 +44,8 @@ pub(crate) enum RpcCallError { Server(JSONRPCErrorError), /// The executor did not return a response before the caller's deadline. TimedOut { method: String, timeout: Duration }, + /// The client already has the maximum number of regular RPC calls in flight. + PendingRequestLimitExceeded { limit: usize }, } type PendingRequest = oneshot::Sender>; @@ -240,6 +246,8 @@ pub(crate) struct RpcClient { // can be delivered for their request id. disconnected_rx: watch::Receiver, closed: Arc, + shared_call_slots: Semaphore, + cleanup_call_slots: Semaphore, next_request_id: AtomicI64, transport_tasks: Vec>, transport: JsonRpcTransport, @@ -302,6 +310,8 @@ impl RpcClient { pending, disconnected_rx, closed, + shared_call_slots: Semaphore::new(MAX_IN_FLIGHT_REGULAR_CALLS), + cleanup_call_slots: Semaphore::new(RESERVED_CLEANUP_CALLS), next_request_id: AtomicI64::new(1), transport_tasks, transport, @@ -342,6 +352,17 @@ impl RpcClient { drain_pending(&self.pending).await; } + // Callers keep this permit until `call_inner` returns, so an executor + // cannot free admission early by guessing a request id and replying before + // the request leaves the outbound queue. + fn acquire_regular_call_slot(&self) -> Result, RpcCallError> { + self.shared_call_slots.try_acquire().map_err(|_| { + RpcCallError::PendingRequestLimitExceeded { + limit: MAX_IN_FLIGHT_REGULAR_CALLS, + } + }) + } + #[tracing::instrument( name = "codex.exec_server.request", level = "info", @@ -357,6 +378,7 @@ impl RpcClient { P: Serialize, T: DeserializeOwned, { + let _call_slot = self.acquire_regular_call_slot()?; self.call_inner(method, params, RpcCallTimeout::None).await } @@ -370,10 +392,43 @@ impl RpcClient { P: Serialize, T: DeserializeOwned, { + let _call_slot = self.acquire_regular_call_slot()?; self.call_inner(method, params, RpcCallTimeout::After(call_timeout)) .await } + #[tracing::instrument( + name = "codex.exec_server.request", + level = "info", + skip_all, + fields( + otel.kind = "client", + otel.name = method, + method, + ) + )] + pub(crate) async fn call_for_cleanup( + &self, + method: &str, + params: &P, + ) -> Result + where + P: Serialize, + T: DeserializeOwned, + { + let _call_slot = match self.shared_call_slots.try_acquire() { + Ok(call_slot) => call_slot, + Err(_) => match self.cleanup_call_slots.try_acquire() { + Ok(call_slot) => call_slot, + Err(_) => { + self.close_transport().await; + return Err(RpcCallError::Closed); + } + }, + }; + self.call_inner(method, params, RpcCallTimeout::None).await + } + async fn call_inner( &self, method: &str, @@ -394,6 +449,7 @@ impl RpcClient { if self.closed.load(Ordering::Acquire) || *self.disconnected_rx.borrow() { return Err(RpcCallError::Closed); } + pending.retain(|_, response_tx| !response_tx.is_closed()); pending.insert(request_id.clone(), response_tx); } @@ -610,10 +666,13 @@ async fn drain_pending(pending: &Mutex>) { #[cfg(test)] mod tests { + use std::sync::Arc; use std::time::Duration; use codex_exec_server_protocol::JSONRPCMessage; + use codex_exec_server_protocol::JSONRPCNotification; use codex_exec_server_protocol::JSONRPCResponse; + use codex_exec_server_protocol::RequestId; use opentelemetry::trace::TracerProvider as _; use opentelemetry_sdk::trace::InMemorySpanExporter; use opentelemetry_sdk::trace::SdkTracerProvider; @@ -621,13 +680,18 @@ mod tests { use tokio::io::AsyncBufReadExt; use tokio::io::AsyncWriteExt; use tokio::io::BufReader; + use tokio::task::JoinSet; use tokio::time::timeout; use tracing::Instrument; use tracing_subscriber::filter::filter_fn; use tracing_subscriber::prelude::*; + use super::MAX_IN_FLIGHT_REGULAR_CALLS; + use super::RpcCallError; use super::RpcClient; use crate::connection::JsonRpcConnection; + use crate::connection::JsonRpcConnectionEvent; + use crate::connection::JsonRpcTransport; async fn read_jsonrpc_line(lines: &mut tokio::io::Lines>) -> JSONRPCMessage where @@ -731,6 +795,39 @@ mod tests { } } + #[tokio::test(start_paused = true)] + async fn rpc_client_call_has_no_implicit_deadline() { + let (client_stdin, server_reader) = tokio::io::duplex(4096); + let (mut server_writer, client_stdout) = tokio::io::duplex(4096); + let connection = + JsonRpcConnection::from_stdio(client_stdout, client_stdin, "test-rpc".to_string()); + let (client, _events_rx) = RpcClient::new(connection); + let mut lines = BufReader::new(server_reader).lines(); + + let params = serde_json::json!({}); + let call = client.call::<_, serde_json::Value>("slow", ¶ms); + tokio::pin!(call); + assert!(futures::poll!(call.as_mut()).is_pending()); + let request = match read_jsonrpc_line(&mut lines).await { + JSONRPCMessage::Request(request) => request, + other => panic!("expected JSON-RPC request, got {other:?}"), + }; + + tokio::time::advance(Duration::from_secs(61)).await; + assert!(futures::poll!(call.as_mut()).is_pending()); + + let expected = serde_json::json!({ "value": "done" }); + write_jsonrpc_line( + &mut server_writer, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: expected.clone(), + }), + ) + .await; + assert_eq!(call.await.expect("RPC response"), expected); + } + #[tokio::test] async fn rpc_client_timeout_removes_pending_request() { let (client_stdin, server_reader) = tokio::io::duplex(4096); @@ -765,6 +862,110 @@ mod tests { } } + #[tokio::test] + async fn rpc_client_bounds_in_flight_calls_and_preserves_cleanup() { + let (outgoing_tx, outgoing_rx) = tokio::sync::mpsc::channel(/*buffer*/ 1); + outgoing_tx + .send(JSONRPCMessage::Notification(JSONRPCNotification { + method: "blocker".to_string(), + params: None, + })) + .await + .expect("outbound queue should accept the blocker"); + let (incoming_tx, incoming_rx) = tokio::sync::mpsc::channel(MAX_IN_FLIGHT_REGULAR_CALLS); + let (_disconnected_tx, disconnected_rx) = tokio::sync::watch::channel(/*init*/ false); + let connection = JsonRpcConnection { + outgoing_tx, + incoming_rx, + disconnected_rx, + task_handles: Vec::new(), + transport: JsonRpcTransport::Plain, + }; + let (client, _events_rx) = RpcClient::new(connection); + let client = Arc::new(client); + let mut calls = JoinSet::new(); + + for index in 0..MAX_IN_FLIGHT_REGULAR_CALLS { + let client = Arc::clone(&client); + calls.spawn(async move { + client + .call::<_, serde_json::Value>("pending", &serde_json::json!({ "index": index })) + .await + }); + } + timeout(Duration::from_secs(1), async { + while client.pending_request_count().await < MAX_IN_FLIGHT_REGULAR_CALLS { + tokio::task::yield_now().await; + } + }) + .await + .expect("pending requests should reach the regular limit"); + + for request_id in 1..=MAX_IN_FLIGHT_REGULAR_CALLS { + incoming_tx + .send(JsonRpcConnectionEvent::Message(JSONRPCMessage::Response( + JSONRPCResponse { + id: RequestId::Integer( + i64::try_from(request_id).expect("request id should fit in i64"), + ), + result: serde_json::json!({}), + }, + ))) + .await + .expect("reader should accept the spoofed response"); + } + timeout(Duration::from_secs(1), async { + while client.pending_request_count().await != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("spoofed responses should drain response routing"); + + let params = serde_json::json!({}); + let overflow = client.call::<_, serde_json::Value>("overflow", ¶ms); + tokio::pin!(overflow); + assert!(matches!( + futures::poll!(overflow.as_mut()), + std::task::Poll::Ready(Err(RpcCallError::PendingRequestLimitExceeded { limit })) + if limit == MAX_IN_FLIGHT_REGULAR_CALLS + )); + + let cleanup_client = Arc::clone(&client); + calls.spawn(async move { + let params = serde_json::json!({}); + cleanup_client + .call_for_cleanup::<_, serde_json::Value>("cleanup", ¶ms) + .await + }); + timeout(Duration::from_secs(1), async { + while client.pending_request_count().await != 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("cleanup request should use the reserved capacity"); + + let cleanup_params = serde_json::json!({}); + let cleanup_overflow = timeout( + Duration::from_secs(1), + client.call_for_cleanup::<_, serde_json::Value>("cleanup-overflow", &cleanup_params), + ) + .await + .expect("cleanup circuit breaker should not block"); + assert!(matches!(cleanup_overflow, Err(RpcCallError::Closed))); + assert!(client.is_disconnected()); + assert_eq!(client.pending_request_count().await, 0); + + drop(outgoing_rx); + while let Some(call) = calls.join_next().await { + assert!(matches!( + call.expect("pending call task should join"), + Err(RpcCallError::Closed) + )); + } + } + #[tokio::test(flavor = "current_thread")] async fn rpc_client_propagates_current_trace_context() { let span_exporter = InMemorySpanExporter::default();