Files
codex/codex-rs/code-mode-host/src/peer_tests.rs
T
pash-openai 56032e9732 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
2026-11-05 07:13:37 +00:00

100 lines
3.3 KiB
Rust

use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::RuntimeResponse;
use codex_code_mode_protocol::StartedCell;
use codex_code_mode_protocol::host::DelegateRequest;
use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
use codex_code_mode_protocol::host::RequestId;
use codex_code_mode_protocol::host::SessionId;
use pretty_assertions::assert_eq;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::sync::oneshot::error::TryRecvError;
use tokio_util::sync::CancellationToken;
use super::HostPeer;
fn session_id(value: &str) -> SessionId {
SessionId::new(value).expect("session ID")
}
#[tokio::test]
async fn start_cell_reports_when_initial_response_is_enqueued() {
let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*max_capacity*/ 4);
let peer = Arc::new(HostPeer::new(outgoing_tx, MAX_PENDING_DELEGATE_CALLS));
let cell_id = CellId::new("cell-1".to_string());
let (response_tx, response_rx) = oneshot::channel();
let started = StartedCell::new(cell_id.clone(), response_rx);
let active_cell_permits = Arc::new(Semaphore::new(/*permits*/ 1));
let active_cell_permit = Arc::clone(&active_cell_permits)
.try_acquire_owned()
.expect("active cell permit");
let mut initial_response_sent = peer.start_cell(
session_id("session-1"),
RequestId::new(/*value*/ 1),
started,
active_cell_permit,
Instant::now(),
);
assert_eq!(initial_response_sent.try_recv(), Err(TryRecvError::Empty));
response_tx
.send(RuntimeResponse::Result {
code_mode_host_duration: None,
cell_id: cell_id.clone(),
content_items: Vec::new(),
result_buffered_for: Duration::ZERO,
error_text: None,
})
.expect("initial response receiver");
initial_response_sent
.await
.expect("initial response completion");
outgoing_rx.recv().await.expect("initial response frame");
assert_eq!(active_cell_permits.available_permits(), 0);
peer.close_cell(session_id("session-1"), cell_id);
let permit = tokio::time::timeout(
Duration::from_secs(1),
Arc::clone(&active_cell_permits).acquire_owned(),
)
.await
.expect("cell permit should be released")
.expect("cell permit semaphore should remain open");
drop(permit);
}
#[tokio::test]
async fn pending_delegate_limit_rejects_call_without_disconnecting() {
let (outgoing_tx, _outgoing_rx) = mpsc::channel(/*max_capacity*/ 1);
let peer = Arc::new(HostPeer::new(outgoing_tx, MAX_PENDING_DELEGATE_CALLS));
let permits = Arc::clone(&peer.delegate_permits)
.acquire_many_owned(MAX_PENDING_DELEGATE_CALLS as u32)
.await
.expect("delegate permits");
let result = peer
.call(
session_id("session-1"),
DelegateRequest::Notify {
call_id: "call-1".to_string(),
cell_id: CellId::new("cell-1".to_string()).into(),
text: "hello".to_string(),
},
CancellationToken::new(),
)
.await;
assert_eq!(
result,
Err("code-mode host has too many pending delegate calls".to_string())
);
assert!(!peer.is_disconnected());
drop(permits);
}