openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-06 20:20:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
## 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
This commit is contained in:
pash-openai
2026-11-05 07:13:37 +00:00
committed by copyberry
parent 695961ee88
commit 56032e9732
19 changed files with 435 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode/src/grpc_session/callbacks.rs
+10 -6
View File
@@ -190,12 +190,16 @@ impl SessionInner {
_ = cancellation.cancelled() => {}
// NotFound can mean an invocation in a live session was canceled.
// Check cancellation before treating a completion error as fatal.
result = deadline::request(
self,
"tool invocation completion",
Duration::ZERO,
client.complete_tool_call(self.request(request)),
) => {
result = async {
let _permit = self.completion_permits.acquire().await
.map_err(|_| "code-mode completion queue is closed".to_string())?;
deadline::request(
self,
"tool invocation completion",
Duration::ZERO,
client.complete_tool_call(self.request(request)),
).await
} => {
if let Err(error) = result
&& !cancellation.is_cancelled()
&& !self.stopped.is_cancelled()
codex-rs/code-mode/src/grpc_session/callbacks_tests.rs
+28
View File
@@ -10,6 +10,7 @@ use codex_code_mode_protocol::CodeModeOutputTruncation;
use http_body_util::Full;
use pretty_assertions::assert_eq;
use tokio::sync::Notify;
use tokio::sync::Semaphore;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
use tonic::body::Body;
@@ -19,6 +20,7 @@ use tower::service_fn;
use tower::util::BoxCloneSyncService;
use super::super::GrpcClient;
use super::super::MAX_IN_FLIGHT_COMPLETIONS;
use super::super::state::SessionState;
use super::super::transport::SharedTransport;
use super::SessionInner;
@@ -59,10 +61,36 @@ fn session(client: GrpcClient) -> Arc<SessionInner> {
shutdown_result: Mutex::new(/*t*/ None),
stopped: CancellationToken::new(),
stream_tasks: TaskTracker::new(),
completion_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_COMPLETIONS)),
_transport: Arc::new(SharedTransport::Connected(client)),
})
}
#[tokio::test]
async fn cancelled_completion_does_not_wait_for_delivery_capacity() {
let client = GrpcClient::new(BoxCloneSyncService::new(service_fn(|_request| async {
Ok(tonic::Status::internal("a queued completion must not send an RPC").into_http())
})));
let session = session(client);
let _capacity = session
.completion_permits
.acquire_many(MAX_IN_FLIGHT_COMPLETIONS as u32)
.await
.expect("fill delivery window");
let cancellation = CancellationToken::new();
let completion = session.complete_tool_call(
"queued".to_string(),
cancellation.clone(),
Ok(serde_json::Value::Null),
);
tokio::pin!(completion);
assert!(futures::poll!(completion.as_mut()).is_pending());
cancellation.cancel();
completion.await;
assert_eq!(session.require_open(), Ok(()));
}
#[tokio::test(start_paused = true)]
async fn shutdown_is_ordered_with_session_failure() {
let close_started = Arc::new(Notify::new());
codex-rs/code-mode/src/grpc_session/mod.rs
+11
View File
@@ -28,6 +28,7 @@ use codex_code_mode_protocol::grpc::code_mode_host_client::CodeModeHostClient;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use rand::Rng;
use tokio::sync::Semaphore;
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
@@ -57,6 +58,9 @@ type GrpcClient = CodeModeHostClient<GrpcTransport>;
const SHUTDOWN_ERROR: &str = "code mode session is shutting down";
const ROUTE_HEADER: &str = "x-code-mode-route";
// Bound result-delivery RPCs, not running tools. A full completion burst can
// exhaust HTTP/2's small-DATA-frame budget and disconnect every shared session.
const MAX_IN_FLIGHT_COMPLETIONS: usize = 64;
fn inject_span_traceparent<T>(request: &mut tonic::Request<T>, span: &tracing::Span) {
if let Some(traceparent) =
@@ -76,6 +80,7 @@ fn inject_span_traceparent<T>(request: &mut tonic::Request<T>, span: &tracing::S
#[derive(Clone)]
pub struct GrpcCodeModeSessionProvider {
transport: Arc<SharedTransport>,
completion_permits: Arc<Semaphore>,
// Reloaded threads share this counter rather than reusing a cell namespace.
next_generation: Arc<AtomicU64>,
}
@@ -124,6 +129,7 @@ impl GrpcCodeModeSessionProvider {
fn from_transport(transport: SharedTransport) -> Self {
Self {
transport: Arc::new(transport),
completion_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_COMPLETIONS)),
next_generation: Arc::new(AtomicU64::new(1)),
}
}
@@ -153,6 +159,9 @@ impl GrpcCodeModeSessionProvider {
let mut retries = 0;
loop {
let mut request = tonic::Request::new(grpc::OpenSessionRequest {
max_pending_delegate_calls: Some(
codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS as u32,
),
cell_execution_limits,
output_truncation: match requested_mode {
CodeModeOutputTruncation::Client => None,
@@ -243,6 +252,7 @@ impl GrpcCodeModeSessionProvider {
shutdown_result: Mutex::new(None),
stopped: CancellationToken::new(),
stream_tasks: TaskTracker::new(),
completion_permits: Arc::clone(&self.completion_permits),
_transport: Arc::clone(&self.transport),
});
let mut opening = OpeningSession {
@@ -387,6 +397,7 @@ pub(super) struct SessionInner {
shutdown_result: Mutex<Option<ShutdownResultReceiver>>,
pub(super) stopped: CancellationToken,
stream_tasks: TaskTracker,
completion_permits: Arc<Semaphore>,
_transport: Arc<SharedTransport>,
}
codex-rs/code-mode/src/remote_session/connection.rs
+7 -1
View File
@@ -28,6 +28,8 @@ use codex_code_mode_protocol::host::EncodedFrame;
use codex_code_mode_protocol::host::FramedReader;
use codex_code_mode_protocol::host::FramedWriter;
use codex_code_mode_protocol::host::HostToClient;
use codex_code_mode_protocol::host::LARGE_DELEGATE_CAPACITY_CAPABILITY;
use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
use codex_code_mode_protocol::host::ProtocolVersion;
use codex_code_mode_protocol::host::RequestId;
use codex_code_mode_protocol::host::SESSION_RESOURCE_LIMITS_CAPABILITY;
@@ -58,7 +60,8 @@ use self::reader::drive_reader;
mod driver;
mod reader;
const IPC_CHANNEL_CAPACITY: usize = 128;
// Accommodate a completion burst and concurrent control messages from a full host.
const IPC_CHANNEL_CAPACITY: usize = MAX_PENDING_DELEGATE_CALLS * 2;
const LOCAL_HOST_STARTUP_TIMEOUT: Duration = Duration::from_secs(30);
// TODO(anp) make this timeout configurable if 60 seconds is insufficient.
const DEFAULT_HOST_WAIT_TRANSPORT_TIMEOUT: Duration = Duration::from_secs(60);
@@ -215,10 +218,13 @@ impl Connection {
Capability::new(YIELD_OBSERVATION_CAPABILITY).map_err(|error| error.to_string())?;
let output_capability =
Capability::new(BOUNDED_OUTPUT_CAPABILITY).map_err(|error| error.to_string())?;
let delegate_capacity = Capability::new(LARGE_DELEGATE_CAPACITY_CAPABILITY)
.map_err(|error| error.to_string())?;
let optional_capabilities = CapabilitySet::try_new([
session_limits_capability,
yield_capability,
output_capability,
delegate_capacity,
])
.map_err(|error| error.to_string())?;
let hello = ClientHello::new(