Bound exec-server pending RPCs (#31578)

## Why An untrusted exec-server can stop reading requests, never answer them, or send guessed responses before queued requests are written. Without client-side admission, the orchestrator can retain unbounded RPC call futures and request payloads. We need a hard bound without adding a blanket timeout, because individual operations already own their timeout and cleanup semantics. ## What changed - hold one of 64 shared admission permits for the full lifetime of each regular RPC call, so an early response cannot free capacity while the call remains queued - allow `process/terminate` and `fs/close` to use one additional cleanup permit - close the transport and fail pending calls if the cleanup permit is also stuck, allowing teardown or recovery to release remote resources - leave the wire format and existing timeout behavior unchanged ## Tests - `rpc_client_call_has_no_implicit_deadline` verifies that ordinary calls remain untimed - `rpc_client_bounds_in_flight_calls_and_preserves_cleanup` covers the regular-call cap, guessed early responses, cleanup admission, and the cleanup circuit breaker
2026-10-18 17:18:52 +00:00 · 2026-07-08 19:46:18 +01:00
parent 8d83687aaa
commit a14b4c2d7f
3 changed files with 216 additions and 3 deletions
@@ -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"
+            )),
        }
    }
 }
@@ -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(),
@@ -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,10 +666,13 @@ 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;
@@ -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<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();