Bound exec-server pending RPCs

codex-rs/exec-server/src/client.rs

@@ -561,7 +561,7 @@ impl ExecServerClient {
         &self,
         process_id: &ProcessId,
     ) -> Result<TerminateResponse, ExecServerError> {
-        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<FsCloseResponse, ExecServerError> {
-        self.call(FS_CLOSE_METHOD, &params).await
+        self.call_for_cleanup(FS_CLOSE_METHOD, &params).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<P, T>(&self, method: &str, params: &P) -> Result<T, ExecServerError>
+    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<T>(result: Result<T, RpcCallError>) -> Result<T, ExecServerError> {
@@ -830,6 +839,9 @@ impl From<RpcCallError> 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"
+            )),
         }
     }
 }

codex-rs/exec-server/src/client_recovery.rs

@@ -445,7 +445,7 @@ impl Inner {
                 Ok(false) => {}
                 Err(error) => {
                     let terminated: Result<TerminateResponse, ExecServerError> = rpc_client
-                        .call(
+                        .call_for_cleanup(
                             EXEC_TERMINATE_METHOD,
                             &TerminateParams {
                                 process_id: process_id.clone(),

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<Result<Value, RpcCallError>>;
@@ -240,6 +246,8 @@ pub(crate) struct RpcClient {
     // can be delivered for their request id.
     disconnected_rx: watch::Receiver<bool>,
     closed: Arc<AtomicBool>,
+    shared_call_slots: Semaphore,
+    cleanup_call_slots: Semaphore,
     next_request_id: AtomicI64,
     transport_tasks: Vec<JoinHandle<()>>,
     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<SemaphorePermit<'_>, 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<P, T>(
+        &self,
+        method: &str,
+        params: &P,
+    ) -> Result<T, RpcCallError>
+    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<P, T>(
         &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,24 +666,32 @@ async fn drain_pending(pending: &Mutex<HashMap<RequestId, PendingRequest>>) {
 
 #[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;
     use pretty_assertions::assert_eq;
     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<R>(lines: &mut tokio::io::Lines<BufReader<R>>) -> 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", &params);
+        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", &params);
+        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", &params)
+                .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();