Increase shared code-mode host capacity safely (#59590)

## Why

Shared hosts serve independent chats and their subagents. Larger callback workloads need room for unrelated work and control messages, while gRPC completion bursts can exhaust HTTP/2's small-DATA-frame budget and disconnect shared sessions.

## What changed

- Raise the active-cell limit from 128 to 2,048 and the pending delegate callback limit from 1,024 to 4,096. Expand request and transport queues to leave room for completion bursts and control traffic.
- Negotiate larger callback budgets through the stdio `delegate-capacity-4096` capability and gRPC `max_pending_delegate_calls` field. Preserve the 1,024-callback budget for legacy clients and enforce both session and shared host limits on gRPC.
- Limit concurrent gRPC completion RPCs to 64 across a provider's sessions, allowing queued completions to be canceled without waiting for delivery capacity.

## Testing

Add stdio and gRPC integration tests covering 256 cells and 4,096 callbacks across shared sessions while unrelated work completes. Add tests for callback budget isolation and reuse, rejection of zero capacity, and cancellation of queued completions. Include an ignored resource-intensive test for full host capacity.

GitOrigin-RevId: e85bfa24cc46cc8a735b464d80adf8066a6a6d3a

codex-rs/code-mode-host/src/grpc/capacity_tests.rs

@@ -0,0 +1,78 @@
+#![cfg(test)]
+//! Negotiated callback budgets protect older clients on a shared host.
+
+use codex_code_mode_protocol::grpc as proto;
+use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHost;
+use codex_code_mode_protocol::host::LEGACY_MAX_PENDING_DELEGATE_CALLS;
+use futures::StreamExt;
+use pretty_assertions::assert_eq;
+use tonic::Code;
+use tonic::Request;
+
+use super::GrpcCodeModeHost;
+use super::tests::open_session;
+
+#[tokio::test]
+async fn client_callback_capacity_is_isolated_and_reusable() {
+    let host = GrpcCodeModeHost::new();
+    let (legacy_id, _legacy_lease) = open_session(&host).await;
+    let legacy = host.state.session(&legacy_id).expect("legacy session");
+    let legacy_permits = (0..LEGACY_MAX_PENDING_DELEGATE_CALLS)
+        .map(|_| legacy.delegate_permit().expect("legacy callback"))
+        .collect::<Vec<_>>();
+    assert_eq!(
+        legacy.delegate_permit().unwrap_err(),
+        "code-mode session has too many pending delegate calls"
+    );
+
+    let mut lease = host
+        .open_session(Request::new(proto::OpenSessionRequest {
+            max_pending_delegate_calls: Some(2),
+            ..Default::default()
+        }))
+        .await
+        .expect("open capacity-limited session")
+        .into_inner();
+    let event = lease.next().await.expect("opened event").expect("event");
+    let Some(proto::session_event::Event::Opened(opened)) = event.event else {
+        panic!("expected session opening");
+    };
+    let session = host.state.session(&opened.session_id).expect("session");
+    let first = session.delegate_permit().expect("first callback");
+    let second = session.delegate_permit().expect("second callback");
+    assert_eq!(
+        session.delegate_permit().unwrap_err(),
+        "code-mode session has too many pending delegate calls"
+    );
+
+    drop(first);
+    let replacement = session.delegate_permit().expect("reused callback capacity");
+    drop(legacy_permits);
+    let legacy_replacement = legacy.delegate_permit().expect("reused legacy capacity");
+    drop((second, replacement, legacy_replacement));
+    host.close_session(Request::new(proto::CloseSessionRequest {
+        session_id: opened.session_id,
+    }))
+    .await
+    .expect("close new session");
+    host.close_session(Request::new(proto::CloseSessionRequest {
+        session_id: legacy_id,
+    }))
+    .await
+    .expect("close legacy session");
+}
+
+#[tokio::test]
+async fn zero_callback_capacity_is_rejected_before_opening_a_session() {
+    let host = GrpcCodeModeHost::new();
+    let result = host
+        .open_session(Request::new(proto::OpenSessionRequest {
+            max_pending_delegate_calls: Some(0),
+            ..Default::default()
+        }))
+        .await;
+    let Err(error) = result else {
+        panic!("zero callback capacity must be rejected");
+    };
+    assert_eq!(error.code(), Code::InvalidArgument);
+}

codex-rs/code-mode-host/src/grpc/mod.rs

@@ -19,6 +19,8 @@ use codex_code_mode_protocol::TerminateRequest;
 use codex_code_mode_protocol::WaitRequest;
 use codex_code_mode_protocol::grpc as proto;
 use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHost;
+use codex_code_mode_protocol::host::LEGACY_MAX_PENDING_DELEGATE_CALLS;
+use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
 use codex_protocol::protocol::W3cTraceContext;
 use futures::Stream;
 use futures::StreamExt;
@@ -78,6 +80,16 @@ impl GrpcCodeModeHost {
         request: proto::OpenSessionRequest,
     ) -> Result<Response<GrpcStream<proto::SessionEvent>>, Status> {
         let _permit = self.state.request_permit()?;
+        let max_pending_delegate_calls = request
+            .max_pending_delegate_calls
+            .map(|limit| limit as usize)
+            .unwrap_or(LEGACY_MAX_PENDING_DELEGATE_CALLS)
+            .min(MAX_PENDING_DELEGATE_CALLS);
+        if max_pending_delegate_calls == 0 {
+            return Err(Status::invalid_argument(
+                "callback capacity must be positive",
+            ));
+        }
         let limits = conversions::session_limits(request.cell_execution_limits)?;
         let output_truncation = match request
             .output_truncation
@@ -95,6 +107,7 @@ impl GrpcCodeModeHost {
                 cell_execution_limits: limits,
                 output_truncation,
             },
+            max_pending_delegate_calls,
         )?))
     }
 
@@ -533,3 +546,7 @@ mod tests;
 #[cfg(test)]
 #[path = "robustness_tests.rs"]
 mod robustness_tests;
+
+#[cfg(test)]
+#[path = "capacity_tests.rs"]
+mod capacity_tests;

codex-rs/code-mode-host/src/grpc/routing.rs

@@ -24,7 +24,7 @@ use super::session::ToolSubscription;
 use super::validation;
 use crate::OUTGOING_CHANNEL_CAPACITY;
 
-const MAX_SUBSCRIPTIONS: usize = OUTGOING_CHANNEL_CAPACITY;
+const MAX_SUBSCRIPTIONS: usize = 128;
 
 impl GrpcSession {
     pub(super) fn subscribe(

codex-rs/code-mode-host/src/grpc/service_tests.rs

@@ -21,6 +21,7 @@ async fn open_session_with_truncation(
 ) -> (proto::SessionOpened, GrpcStream<proto::SessionEvent>) {
     let mut stream = host
         .open_session(Request::new(proto::OpenSessionRequest {
+            max_pending_delegate_calls: None,
             cell_execution_limits: None,
             output_truncation,
         }))
@@ -605,6 +606,7 @@ async fn unknown_output_truncation_is_rejected_before_opening_a_session() {
     let host = GrpcCodeModeHost::new();
     let result = host
         .open_session(Request::new(proto::OpenSessionRequest {
+            max_pending_delegate_calls: None,
             cell_execution_limits: None,
             output_truncation: Some(42),
         }))

codex-rs/code-mode-host/src/grpc/session.rs

@@ -54,6 +54,7 @@ pub(super) struct GrpcSession {
     events: EventSender,
     cells_changed: Notify,
     delegate_permits: Arc<Semaphore>,
+    client_delegate_permits: Arc<Semaphore>,
 }
 
 #[derive(Default)]
@@ -111,6 +112,7 @@ impl GrpcHostState {
     pub(super) fn open_session(
         self: &Arc<Self>,
         options: CodeModeSessionOptions,
+        max_pending_delegate_calls: usize,
     ) -> Result<GrpcStream<proto::SessionEvent>, Status> {
         let id = Uuid::new_v4();
         let (events, receiver) = mpsc::channel(OUTGOING_CHANNEL_CAPACITY);
@@ -122,6 +124,7 @@ impl GrpcHostState {
             closed,
             Arc::clone(&self.delegate_permits),
             options,
+            max_pending_delegate_calls,
         );
         events
             .try_send(Ok(proto::SessionEvent {
@@ -222,6 +225,7 @@ impl GrpcSession {
         closed: CancellationToken,
         delegate_permits: Arc<Semaphore>,
         options: CodeModeSessionOptions,
+        max_pending_delegate_calls: usize,
     ) -> Arc<Self> {
         Arc::new_cyclic(|weak: &Weak<Self>| {
             let failure_session = weak.clone();
@@ -242,6 +246,7 @@ impl GrpcSession {
                 events,
                 cells_changed: Notify::new(),
                 delegate_permits,
+                client_delegate_permits: Arc::new(Semaphore::new(max_pending_delegate_calls)),
             }
         })
     }
@@ -420,10 +425,14 @@ impl GrpcSession {
         );
     }
 
-    pub(super) fn delegate_permit(&self) -> Result<OwnedSemaphorePermit, String> {
-        Arc::clone(&self.delegate_permits)
+    pub(super) fn delegate_permit(&self) -> Result<[OwnedSemaphorePermit; 2], String> {
+        let client = Arc::clone(&self.client_delegate_permits)
             .try_acquire_owned()
-            .map_err(|_| "code-mode host has too many pending delegate calls".to_string())
+            .map_err(|_| "code-mode session has too many pending delegate calls".to_string())?;
+        let host = Arc::clone(&self.delegate_permits)
+            .try_acquire_owned()
+            .map_err(|_| "code-mode host has too many pending delegate calls".to_string())?;
+        Ok([client, host])
     }
 
     pub(super) async fn send_event(