diff --git a/codex-rs/code-mode-host/src/grpc/capacity_tests.rs b/codex-rs/code-mode-host/src/grpc/capacity_tests.rs new file mode 100644 index 0000000000..87c45812f5 --- /dev/null +++ b/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::>(); + 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); +} diff --git a/codex-rs/code-mode-host/src/grpc/mod.rs b/codex-rs/code-mode-host/src/grpc/mod.rs index 7a4a0c0d27..ea7b9a6d5f 100644 --- a/codex-rs/code-mode-host/src/grpc/mod.rs +++ b/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>, 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; diff --git a/codex-rs/code-mode-host/src/grpc/routing.rs b/codex-rs/code-mode-host/src/grpc/routing.rs index 3a16dc7792..6fe2f76fdf 100644 --- a/codex-rs/code-mode-host/src/grpc/routing.rs +++ b/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( diff --git a/codex-rs/code-mode-host/src/grpc/service_tests.rs b/codex-rs/code-mode-host/src/grpc/service_tests.rs index d3cb8a70c4..48888e738e 100644 --- a/codex-rs/code-mode-host/src/grpc/service_tests.rs +++ b/codex-rs/code-mode-host/src/grpc/service_tests.rs @@ -21,6 +21,7 @@ async fn open_session_with_truncation( ) -> (proto::SessionOpened, GrpcStream) { 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), })) diff --git a/codex-rs/code-mode-host/src/grpc/session.rs b/codex-rs/code-mode-host/src/grpc/session.rs index c1988ce8a3..5e76337435 100644 --- a/codex-rs/code-mode-host/src/grpc/session.rs +++ b/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, + client_delegate_permits: Arc, } #[derive(Default)] @@ -111,6 +112,7 @@ impl GrpcHostState { pub(super) fn open_session( self: &Arc, options: CodeModeSessionOptions, + max_pending_delegate_calls: usize, ) -> Result, 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, options: CodeModeSessionOptions, + max_pending_delegate_calls: usize, ) -> Arc { Arc::new_cyclic(|weak: &Weak| { 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 { - 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( diff --git a/codex-rs/code-mode-host/src/host_tests.rs b/codex-rs/code-mode-host/src/host_tests.rs index 097fa58be9..cf55ff9ce4 100644 --- a/codex-rs/code-mode-host/src/host_tests.rs +++ b/codex-rs/code-mode-host/src/host_tests.rs @@ -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"); }); diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs index aa6d99f3cd..d1a801a6e6 100644 --- a/codex-rs/code-mode-host/src/lib.rs +++ b/codex-rs/code-mode-host/src/lib.rs @@ -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::(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( 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( 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 { diff --git a/codex-rs/code-mode-host/src/peer.rs b/codex-rs/code-mode-host/src/peer.rs index d39682ee22..2980dbebf6 100644 --- a/codex-rs/code-mode-host/src/peer.rs +++ b/codex-rs/code-mode-host/src/peer.rs @@ -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) -> Self { + pub(super) fn new( + outgoing_tx: mpsc::Sender, + 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), diff --git a/codex-rs/code-mode-host/src/peer_tests.rs b/codex-rs/code-mode-host/src/peer_tests.rs index 6aa9ddfbbd..a05d1df4d2 100644 --- a/codex-rs/code-mode-host/src/peer_tests.rs +++ b/codex-rs/code-mode-host/src/peer_tests.rs @@ -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 diff --git a/codex-rs/code-mode-host/tests/capacity.rs b/codex-rs/code-mode-host/tests/capacity.rs new file mode 100644 index 0000000000..4b628fb410 --- /dev/null +++ b/codex-rs/code-mode-host/tests/capacity.rs @@ -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; +} diff --git a/codex-rs/code-mode-host/tests/grpc.rs b/codex-rs/code-mode-host/tests/grpc.rs index 69aef1441f..d3994545af 100644 --- a/codex-rs/code-mode-host/tests/grpc.rs +++ b/codex-rs/code-mode-host/tests/grpc.rs @@ -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, }) diff --git a/codex-rs/code-mode-host/tests/grpc_tcp.rs b/codex-rs/code-mode-host/tests/grpc_tcp.rs index 4daceb2d79..eb81054e46 100644 --- a/codex-rs/code-mode-host/tests/grpc_tcp.rs +++ b/codex-rs/code-mode-host/tests/grpc_tcp.rs @@ -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, }) diff --git a/codex-rs/code-mode-host/tests/grpc_timing.rs b/codex-rs/code-mode-host/tests/grpc_timing.rs index ffe59d784f..e97e3c26aa 100644 --- a/codex-rs/code-mode-host/tests/grpc_timing.rs +++ b/codex-rs/code-mode-host/tests/grpc_timing.rs @@ -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, }) diff --git a/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto b/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto index 043b7f8cd1..80c212421b 100644 --- a/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto +++ b/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto @@ -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 { diff --git a/codex-rs/code-mode-protocol/src/host/mod.rs b/codex-rs/code-mode-protocol/src/host/mod.rs index f3cd0286f0..817f11dba4 100644 --- a/codex-rs/code-mode-protocol/src/host/mod.rs +++ b/codex-rs/code-mode-protocol/src/host/mod.rs @@ -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. diff --git a/codex-rs/code-mode/src/grpc_session/callbacks.rs b/codex-rs/code-mode/src/grpc_session/callbacks.rs index e3df7b0515..7800b5fc41 100644 --- a/codex-rs/code-mode/src/grpc_session/callbacks.rs +++ b/codex-rs/code-mode/src/grpc_session/callbacks.rs @@ -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() diff --git a/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs b/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs index d32bc1efcd..a30c4ee724 100644 --- a/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs +++ b/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs @@ -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 { 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()); diff --git a/codex-rs/code-mode/src/grpc_session/mod.rs b/codex-rs/code-mode/src/grpc_session/mod.rs index e805e24666..61f8b3324d 100644 --- a/codex-rs/code-mode/src/grpc_session/mod.rs +++ b/codex-rs/code-mode/src/grpc_session/mod.rs @@ -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; 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(request: &mut tonic::Request, span: &tracing::Span) { if let Some(traceparent) = @@ -76,6 +80,7 @@ fn inject_span_traceparent(request: &mut tonic::Request, span: &tracing::S #[derive(Clone)] pub struct GrpcCodeModeSessionProvider { transport: Arc, + completion_permits: Arc, // Reloaded threads share this counter rather than reusing a cell namespace. next_generation: Arc, } @@ -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>, pub(super) stopped: CancellationToken, stream_tasks: TaskTracker, + completion_permits: Arc, _transport: Arc, } diff --git a/codex-rs/code-mode/src/remote_session/connection.rs b/codex-rs/code-mode/src/remote_session/connection.rs index 504c0fa80a..6b59d07cdb 100644 --- a/codex-rs/code-mode/src/remote_session/connection.rs +++ b/codex-rs/code-mode/src/remote_session/connection.rs @@ -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(