openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-05 20:10:33 +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-host/src/grpc/capacity_tests.rs
+78
View File
@@ -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
+17
View File
@@ -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
+1 -1
View File
@@ -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
+2
View File
@@ -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
+12 -3
View File
@@ -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(
codex-rs/code-mode-host/src/host_tests.rs
+7 -4
View File
@@ -24,6 +24,8 @@ use codex_code_mode_protocol::host::HostHello;
use codex_code_mode_protocol::host::HostRequest;
use codex_code_mode_protocol::host::HostResponse;
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;
@@ -284,6 +286,7 @@ async fn supported_capabilities_are_negotiated_when_optional_or_required() {
for capability in [
SESSION_RESOURCE_LIMITS_CAPABILITY,
BOUNDED_OUTPUT_CAPABILITY,
LARGE_DELEGATE_CAPACITY_CAPABILITY,
] {
let capability = Capability::new(capability).expect("supported capability");
let host_capabilities =
@@ -660,7 +663,7 @@ fn request_history_is_bounded() {
#[tokio::test]
async fn request_task_panic_disconnects_host() {
let (outgoing_tx, _outgoing_rx) = mpsc::channel(/*max_capacity*/ 1);
let peer = Arc::new(HostPeer::new(outgoing_tx));
let peer = Arc::new(HostPeer::new(outgoing_tx, MAX_PENDING_DELEGATE_CALLS));
let state = HostState {
sessions: Mutex::new(HashMap::new()),
limits: Arc::new(HostLimits::new()),
@@ -688,7 +691,7 @@ async fn request_task_panic_disconnects_host() {
#[tokio::test]
async fn execute_request_id_remains_active_until_initial_response() {
let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*max_capacity*/ 4);
let peer = Arc::new(HostPeer::new(outgoing_tx));
let peer = Arc::new(HostPeer::new(outgoing_tx, MAX_PENDING_DELEGATE_CALLS));
let state = Arc::new(HostState {
sessions: Mutex::new(HashMap::new()),
limits: Arc::new(HostLimits::new()),
@@ -749,7 +752,7 @@ async fn execute_request_id_remains_active_until_initial_response() {
#[tokio::test]
async fn active_cell_limit_rejects_execute_without_disconnecting() {
let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*max_capacity*/ 1);
let peer = Arc::new(HostPeer::new(outgoing_tx));
let peer = Arc::new(HostPeer::new(outgoing_tx, MAX_PENDING_DELEGATE_CALLS));
let state = HostState {
sessions: Mutex::new(HashMap::new()),
limits: Arc::new(HostLimits {
@@ -797,7 +800,7 @@ async fn active_cell_limit_rejects_execute_without_disconnecting() {
#[tokio::test]
async fn cell_forwarding_panic_disconnects_host() {
let (outgoing_tx, _outgoing_rx) = mpsc::channel(/*max_capacity*/ 1);
let peer = Arc::new(HostPeer::new(outgoing_tx));
let peer = Arc::new(HostPeer::new(outgoing_tx, MAX_PENDING_DELEGATE_CALLS));
peer.spawn_critical("cell forwarding", async {
panic!("cell forwarding panic probe");
});
codex-rs/code-mode-host/src/lib.rs
+27 -9
View File
@@ -28,6 +28,9 @@ use codex_code_mode_protocol::host::HostHello;
use codex_code_mode_protocol::host::HostRequest;
use codex_code_mode_protocol::host::HostResponse;
use codex_code_mode_protocol::host::HostToClient;
use codex_code_mode_protocol::host::LARGE_DELEGATE_CAPACITY_CAPABILITY;
use codex_code_mode_protocol::host::LEGACY_MAX_PENDING_DELEGATE_CALLS;
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;
@@ -61,16 +64,20 @@ mod peer;
mod process_metrics;
mod transport;
const MAX_IN_FLIGHT_REQUESTS: usize = 256;
const MAX_ACTIVE_CELLS: usize = 128;
// A host is shared by independent chats and all their subagents.
// Test live V8 isolates before increasing this, not just cell permit capacity.
const MAX_ACTIVE_CELLS: usize = 2_048;
// Leave observation and control headroom during a full callback-completion burst.
const MAX_IN_FLIGHT_REQUESTS: usize = MAX_PENDING_DELEGATE_CALLS * 2;
const MAX_RECENT_REQUEST_IDS: usize = 4096;
const MAX_RECENT_SESSION_IDS: usize = 4096;
const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
const OUTGOING_CHANNEL_CAPACITY: usize = 128;
// Allow a full callback-completion burst plus concurrent control messages.
const OUTGOING_CHANNEL_CAPACITY: usize = MAX_PENDING_DELEGATE_CALLS * 2;
enum NegotiatedConnection {
Rejected,
Accepted,
Accepted { max_pending_delegate_calls: usize },
}
struct HostLimits {
@@ -132,12 +139,14 @@ where
{
let mut reader = FramedReader::new(reader);
let mut writer = FramedWriter::new(writer);
match negotiate(&mut reader, &mut writer).await? {
let max_pending_delegate_calls = match negotiate(&mut reader, &mut writer).await? {
NegotiatedConnection::Rejected => return Ok(()),
NegotiatedConnection::Accepted => {}
}
NegotiatedConnection::Accepted {
max_pending_delegate_calls,
} => max_pending_delegate_calls,
};
let (outgoing_tx, outgoing_rx) = mpsc::channel::<EncodedFrame>(OUTGOING_CHANNEL_CAPACITY);
let peer = Arc::new(HostPeer::new(outgoing_tx));
let peer = Arc::new(HostPeer::new(outgoing_tx, max_pending_delegate_calls));
let state = Arc::new(HostState {
sessions: Mutex::new(HashMap::new()),
limits: Arc::new(HostLimits::new()),
@@ -284,6 +293,7 @@ async fn negotiate<R: AsyncRead + Unpin, W: AsyncWrite + Unpin>(
SESSION_RESOURCE_LIMITS_CAPABILITY,
YIELD_OBSERVATION_CAPABILITY,
BOUNDED_OUTPUT_CAPABILITY,
LARGE_DELEGATE_CAPACITY_CAPABILITY,
]
.into_iter()
.map(Capability::new)
@@ -310,12 +320,20 @@ async fn negotiate<R: AsyncRead + Unpin, W: AsyncWrite + Unpin>(
return Ok(NegotiatedConnection::Rejected);
}
let max_pending_delegate_calls =
if host_capabilities.contains(&Capability::new(LARGE_DELEGATE_CAPACITY_CAPABILITY)?) {
MAX_PENDING_DELEGATE_CALLS
} else {
LEGACY_MAX_PENDING_DELEGATE_CALLS
};
let hello = HostHello::new(ProtocolVersion::V1, host_capabilities);
writer
.write(&HostToClient::HostHello(hello))
.await
.context("failed to write code-mode host hello")?;
Ok(NegotiatedConnection::Accepted)
Ok(NegotiatedConnection::Accepted {
max_pending_delegate_calls,
})
}
struct HostState {
codex-rs/code-mode-host/src/peer.rs
+5 -3
View File
@@ -14,7 +14,6 @@ use codex_code_mode_protocol::host::DelegateRequestId;
use codex_code_mode_protocol::host::DelegateResponse;
use codex_code_mode_protocol::host::EncodedFrame;
use codex_code_mode_protocol::host::HostToClient;
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 codex_code_mode_protocol::host::WireResult;
@@ -66,11 +65,14 @@ struct InitialRequest {
}
impl HostPeer {
pub(super) fn new(outgoing_tx: mpsc::Sender<EncodedFrame>) -> Self {
pub(super) fn new(
outgoing_tx: mpsc::Sender<EncodedFrame>,
max_pending_delegate_calls: usize,
) -> Self {
Self {
outgoing_tx,
pending: Mutex::new(HashMap::new()),
delegate_permits: Arc::new(Semaphore::new(MAX_PENDING_DELEGATE_CALLS)),
delegate_permits: Arc::new(Semaphore::new(max_pending_delegate_calls)),
cell_routes: StdMutex::new(HashMap::new()),
cell_routes_changed: Notify::new(),
next_request_id: AtomicI64::new(1),
codex-rs/code-mode-host/src/peer_tests.rs
+3 -3
View File
@@ -6,6 +6,7 @@ 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;
@@ -16,7 +17,6 @@ use tokio::sync::oneshot::error::TryRecvError;
use tokio_util::sync::CancellationToken;
use super::HostPeer;
use super::MAX_PENDING_DELEGATE_CALLS;
fn session_id(value: &str) -> SessionId {
SessionId::new(value).expect("session ID")
@@ -25,7 +25,7 @@ fn session_id(value: &str) -> SessionId {
#[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));
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);
@@ -72,7 +72,7 @@ async fn start_cell_reports_when_initial_response_is_enqueued() {
#[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));
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
codex-rs/code-mode-host/tests/capacity.rs
+213
View File
@@ -0,0 +1,213 @@
//! Shared-host capacity, callback completion, and unrelated-session progress.
use std::sync::Arc;
use codex_code_mode::CellId;
use codex_code_mode::CodeModeNestedToolCall;
use codex_code_mode::CodeModeSessionDelegate;
use codex_code_mode::CodeModeSessionProvider;
use codex_code_mode::CodeModeToolKind;
use codex_code_mode::ExecuteRequest;
use codex_code_mode::FunctionCallOutputContentItem;
use codex_code_mode::GrpcCodeModeSessionProvider;
use codex_code_mode::NotificationFuture;
use codex_code_mode::ProcessOwnedCodeModeSessionProvider;
use codex_code_mode::RuntimeResponse;
use codex_code_mode::ToolDefinition;
use codex_code_mode::ToolInvocationFuture;
use codex_code_mode::WaitOutcome;
use codex_code_mode::WaitRequest;
use codex_protocol::ToolName;
use pretty_assertions::assert_eq;
use serde_json::json;
use tokio::sync::Semaphore;
use tokio_util::sync::CancellationToken;
#[path = "support/host.rs"]
mod host;
struct BlockingDelegate {
started: Semaphore,
release: Semaphore,
}
impl CodeModeSessionDelegate for BlockingDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
cancellation: CancellationToken,
) -> ToolInvocationFuture<'a> {
Box::pin(async move {
self.started.add_permits(1);
tokio::select! {
permit = self.release.acquire() => {
permit.map_err(|error| error.to_string())?.forget();
Ok(json!("done"))
}
_ = cancellation.cancelled() => Err("cancelled".to_string()),
}
})
}
fn notify<'a>(
&'a self,
_call_id: String,
_cell_id: CellId,
_text: String,
_cancellation: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async { Ok(()) })
}
fn cell_closed(&self, _cell_id: &CellId) {}
}
#[allow(clippy::expect_used)]
async fn assert_shared_capacity(
provider: &dyn CodeModeSessionProvider,
cells: usize,
calls_per_cell: usize,
) {
let delegate = Arc::new(BlockingDelegate {
started: Semaphore::new(/*permits*/ 0),
release: Semaphore::new(/*permits*/ 0),
});
let mut sessions = Vec::new();
for _ in 0..2 {
sessions.push(
provider
.create_session()
.await
.expect("create shared session"),
);
}
let mut pending = Vec::new();
for index in 0..cells {
let session = &sessions[index % sessions.len()];
let request = ExecuteRequest {
tool_call_id: format!("capacity-{index}"),
enabled_tools: vec![ToolDefinition {
name: "block".to_string(),
tool_name: ToolName::plain("block"),
description: String::new(),
kind: CodeModeToolKind::Function,
input_schema: None,
input_schema_max_bytes: None,
output_schema: None,
}],
source: format!(
"const results = await Promise.all(Array.from({{length: {calls_per_cell}}}, () => tools.block({{}}))); text(results.length);"
),
yield_time_ms: Some(1),
max_output_tokens: None,
};
let started = session
.execute(request, delegate.clone(), /*preempt*/ None)
.await
.unwrap_or_else(|error| panic!("admit cell {index}: {error}"));
let cell_id = started.cell_id.clone();
let response = started.initial_response().await.expect("initial response");
assert!(
matches!(response, RuntimeResponse::Yielded { .. }),
"{response:?}"
);
pending.push((Arc::clone(session), cell_id));
}
let callbacks = cells * calls_per_cell;
delegate
.started
.acquire_many(callbacks as u32)
.await
.expect("all callbacks started")
.forget();
let probe = sessions[0]
.execute(
ExecuteRequest {
tool_call_id: "unrelated-probe".to_string(),
enabled_tools: Vec::new(),
source: "text('ready');".to_string(),
yield_time_ms: Some(60_000),
max_output_tokens: None,
},
delegate.clone(),
/*preempt*/ None,
)
.await
.expect("unrelated work admitted")
.initial_response()
.await
.expect("unrelated work completed");
assert!(
matches!(
probe,
RuntimeResponse::Result {
error_text: None,
..
}
),
"{probe:?}"
);
delegate.release.add_permits(callbacks);
for (session, cell_id) in pending {
let result = session
.wait(
WaitRequest {
cell_id,
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
.await
.expect("collect callback results");
let WaitOutcome::LiveCell(RuntimeResponse::Result {
content_items,
error_text,
..
}) = result
else {
panic!("expected completed cell: {result:?}");
};
assert_eq!(
(content_items, error_text),
(
vec![FunctionCallOutputContentItem::InputText {
text: calls_per_cell.to_string()
}],
None
)
);
}
for session in sessions {
session.shutdown().await.expect("close shared session");
}
}
#[tokio::test]
async fn stdio_shared_sessions_complete_thousands_of_callbacks() {
let provider = ProcessOwnedCodeModeSessionProvider::with_host_program(
codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"),
);
assert_shared_capacity(&provider, /*cells*/ 256, /*calls_per_cell*/ 16).await;
}
#[tokio::test]
async fn grpc_shared_sessions_complete_thousands_of_callbacks() {
let host = host::HostHarness::start("grpc://127.0.0.1:0")
.await
.expect("host");
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
assert_shared_capacity(&provider, /*cells*/ 256, /*calls_per_cell*/ 16).await;
}
#[tokio::test]
#[ignore = "requires resources for 2,048 live V8 isolates and OS threads"]
async fn grpc_full_host_capacity() {
let host = host::HostHarness::start("grpc://127.0.0.1:0")
.await
.expect("host");
let provider = GrpcCodeModeSessionProvider::new(host.endpoint);
assert_shared_capacity(&provider, /*cells*/ 2_047, /*calls_per_cell*/ 2).await;
}
codex-rs/code-mode-host/tests/grpc.rs
+1
View File
@@ -1171,6 +1171,7 @@ async fn dropping_a_grpc_lease_retires_its_server_session() -> Result<()> {
.context("connect raw gRPC client")?;
let mut lease = client
.open_session(grpc::OpenSessionRequest {
max_pending_delegate_calls: None,
cell_execution_limits: None,
output_truncation: None,
})
codex-rs/code-mode-host/tests/grpc_tcp.rs
+1
View File
@@ -149,6 +149,7 @@ async fn tcp_listener_opens_a_grpc_session() -> Result<()> {
}
let response = client
.open_session(grpc::OpenSessionRequest {
max_pending_delegate_calls: None,
cell_execution_limits: None,
output_truncation: None,
})
codex-rs/code-mode-host/tests/grpc_timing.rs
+1
View File
@@ -39,6 +39,7 @@ async fn execution_timing_includes_javascript_but_excludes_delayed_reads() -> Re
let mut client = CodeModeHostClient::connect(host.endpoint).await?;
let mut events = client
.open_session(grpc::OpenSessionRequest {
max_pending_delegate_calls: None,
cell_execution_limits: None,
output_truncation: None,
})
codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto
+4
View File
@@ -60,6 +60,10 @@ message OpenSessionRequest {
optional SessionCellExecutionLimits cell_execution_limits = 1;
// Absent preserves the legacy client-side truncation behavior.
optional OutputTruncation output_truncation = 2;
// Maximum pending tool callbacks the client can accept.
// Absent preserves the legacy 1,024 callback budget. Hosts may impose a lower
// shared limit, and older hosts ignore this field and retain their own limit.
optional uint32 max_pending_delegate_calls = 3;
}
enum OutputTruncation {
codex-rs/code-mode-protocol/src/host/mod.rs
+7 -1
View File
@@ -10,7 +10,13 @@ mod payload;
mod types;
/// Maximum number of unresolved delegate callbacks allowed per host connection.
pub const MAX_PENDING_DELEGATE_CALLS: usize = 1_024;
pub const MAX_PENDING_DELEGATE_CALLS: usize = 4_096;
/// Callback capacity of clients that do not advertise the larger budget.
pub const LEGACY_MAX_PENDING_DELEGATE_CALLS: usize = 1_024;
/// Stdio clients that can accept up to 4,096 unresolved delegate callbacks.
pub const LARGE_DELEGATE_CAPACITY_CAPABILITY: &str = "delegate-capacity-4096";
/// Stdio connection support for the optional `session/open.outputTruncation` field.
/// Each session selects its own setting after the connection hello.
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(