Files
codex/codex-rs/code-mode-host/src/lib.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

799 lines
27 KiB
Rust

use std::collections::HashMap;
use std::collections::HashSet;
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::PoisonError;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
use std::time::Instant;
use anyhow::Context;
use anyhow::Result;
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::TerminateRequest;
use codex_code_mode_protocol::WaitRequest;
use codex_code_mode_protocol::host::BOUNDED_OUTPUT_CAPABILITY;
use codex_code_mode_protocol::host::Capability;
use codex_code_mode_protocol::host::CapabilitySet;
use codex_code_mode_protocol::host::ClientToHost;
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::HandshakeRejectReason;
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;
use codex_code_mode_protocol::host::SessionId;
use codex_code_mode_protocol::host::SupportedProtocolVersions;
use codex_code_mode_protocol::host::WireWaitOutcome;
use codex_code_mode_protocol::host::YIELD_OBSERVATION_CAPABILITY;
use codex_code_mode_runtime::InProcessCodeModeSession;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::sync::OwnedSemaphorePermit;
use tokio::sync::Semaphore;
use tokio::sync::TryAcquireError;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
use self::delegate::RemoteDelegate;
use self::peer::HostPeer;
pub use self::grpc::GrpcCodeModeHost;
pub use self::process_metrics::CodeModeProcessMetricsCollector;
pub use self::transport::DEFAULT_LISTEN_URL;
pub use codex_code_mode_protocol::CodeModeProcessMetrics;
mod delegate;
mod grpc;
mod grpc_stdio;
mod grpc_transport;
mod peer;
mod process_metrics;
mod transport;
// 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);
// Allow a full callback-completion burst plus concurrent control messages.
const OUTGOING_CHANNEL_CAPACITY: usize = MAX_PENDING_DELEGATE_CALLS * 2;
enum NegotiatedConnection {
Rejected,
Accepted { max_pending_delegate_calls: usize },
}
struct HostLimits {
request_permits: Arc<Semaphore>,
active_cell_permits: Arc<Semaphore>,
}
impl HostLimits {
fn new() -> Self {
Self {
request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)),
active_cell_permits: Arc::new(Semaphore::new(MAX_ACTIVE_CELLS)),
}
}
fn request_permit(&self) -> Result<OwnedSemaphorePermit, TryAcquireError> {
Arc::clone(&self.request_permits).try_acquire_owned()
}
fn cell_permit(&self) -> Result<OwnedSemaphorePermit, TryAcquireError> {
Arc::clone(&self.active_cell_permits).try_acquire_owned()
}
}
/// Runs the code-mode host on its configured stdio or gRPC transport.
#[tracing::instrument(
name = "code_mode_host.run_main",
level = "info",
skip_all,
fields(otel.name = "code_mode_host.run_main")
)]
pub async fn run_main(listen_url: &str) -> Result<()> {
transport::run_transport(listen_url, GrpcCodeModeHost::new()).await
}
/// Runs with an optional gRPC metrics collector initialized before sandbox entry.
/// The legacy stdio transport remains unchanged.
pub async fn run_main_with_process_metrics(
listen_url: &str,
collector: CodeModeProcessMetricsCollector,
) -> Result<()> {
transport::run_transport(
listen_url,
GrpcCodeModeHost::with_process_metrics(collector),
)
.await
}
/// Runs one code-mode host connection over the process standard streams.
pub async fn run_stdio() -> Result<()> {
run(tokio::io::stdin(), tokio::io::stdout()).await
}
/// Runs one code-mode host connection over an ordered input/output pair.
async fn run<R, W>(reader: R, writer: W) -> Result<()>
where
R: AsyncRead + Send + Unpin + 'static,
W: AsyncWrite + Send + Unpin + 'static,
{
let mut reader = FramedReader::new(reader);
let mut writer = FramedWriter::new(writer);
let max_pending_delegate_calls = match negotiate(&mut reader, &mut writer).await? {
NegotiatedConnection::Rejected => return Ok(()),
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, max_pending_delegate_calls));
let state = Arc::new(HostState {
sessions: Mutex::new(HashMap::new()),
limits: Arc::new(HostLimits::new()),
seen_session_ids: Mutex::new(SeenSessionIds::default()),
requests: Mutex::new(RequestRegistry::default()),
request_tasks: TaskTracker::new(),
closing: AtomicBool::new(false),
peer: Arc::clone(&peer),
});
let writer_disconnected = peer.disconnection_token();
let writer_task =
tokio::spawn(async move { drive_writer(writer, outgoing_rx, writer_disconnected).await });
let writer_peer = Arc::clone(&peer);
let writer_supervisor = tokio::spawn(async move {
match writer_task.await {
Ok(Ok(())) if !writer_peer.is_disconnected() => {
writer_peer.fail("code-mode writer task exited unexpectedly".to_string());
}
Ok(Ok(())) => {}
Ok(Err(err)) => {
writer_peer.fail(format!("code-mode writer task failed: {err:#}"));
}
Err(err) => {
writer_peer.fail(format!("code-mode writer task failed: {err}"));
}
}
});
let input_result = async {
loop {
let message = tokio::select! {
biased;
_ = peer.disconnected() => break,
message = reader.read() => message.context("failed to read code-mode client message")?,
};
let Some(message) = message else {
break;
};
match message {
ClientToHost::ClientHello(_) => {
anyhow::bail!("received a second code-mode client hello");
}
ClientToHost::Request { id, request } => {
state.spawn_request(id, request)?;
}
ClientToHost::CancelRequest { id } => {
state.cancel_request(id);
}
ClientToHost::YieldRequest { id } => {
state.yield_request(id);
}
ClientToHost::DelegateResponse { id, result } => {
peer.complete(id, result.into_result()).await;
}
}
}
Ok::<(), anyhow::Error>(())
}
.await;
peer.disconnect();
if tokio::time::timeout(SHUTDOWN_TIMEOUT, state.disconnect())
.await
.is_err()
{
peer.fail("timed out shutting down code-mode host state".to_string());
}
drop(state);
tokio::time::timeout(SHUTDOWN_TIMEOUT, writer_supervisor)
.await
.context("timed out supervising code-mode writer task")?
.context("code-mode writer supervisor task failed")?;
let failure = peer.failure();
drop(peer);
input_result?;
if let Some(failure) = failure {
anyhow::bail!(failure);
}
Ok(())
}
async fn drive_writer<W: AsyncWrite + Unpin>(
mut writer: FramedWriter<W>,
mut outgoing: mpsc::Receiver<EncodedFrame>,
disconnected: CancellationToken,
) -> Result<()> {
loop {
tokio::select! {
_ = disconnected.cancelled() => return Ok(()),
frame = outgoing.recv() => {
let Some(frame) = frame else {
return Ok(());
};
tokio::select! {
_ = disconnected.cancelled() => return Ok(()),
result = writer.write_frame(&frame) => {
result.context("failed to write code-mode host message")?;
}
}
}
}
}
}
async fn negotiate<R: AsyncRead + Unpin, W: AsyncWrite + Unpin>(
reader: &mut FramedReader<R>,
writer: &mut FramedWriter<W>,
) -> Result<NegotiatedConnection> {
let Some(first_message) = reader
.read()
.await
.context("failed to read code-mode client hello")?
else {
return Ok(NegotiatedConnection::Rejected);
};
let ClientToHost::ClientHello(client_hello) = first_message else {
writer
.write(&HostToClient::HandshakeRejected {
reason: HandshakeRejectReason::InvalidHello {
message: "first message must be connection/hello".to_string(),
},
})
.await
.context("failed to reject invalid code-mode client hello")?;
return Ok(NegotiatedConnection::Rejected);
};
let supported_versions = SupportedProtocolVersions::try_new([ProtocolVersion::V1])?;
if !client_hello
.supported_versions()
.contains(ProtocolVersion::V1)
{
writer
.write(&HostToClient::HandshakeRejected {
reason: HandshakeRejectReason::NoCompatibleVersion { supported_versions },
})
.await
.context("failed to reject incompatible code-mode client")?;
return Ok(NegotiatedConnection::Rejected);
}
let host_capabilities = CapabilitySet::try_new(
[
SESSION_RESOURCE_LIMITS_CAPABILITY,
YIELD_OBSERVATION_CAPABILITY,
BOUNDED_OUTPUT_CAPABILITY,
LARGE_DELEGATE_CAPACITY_CAPABILITY,
]
.into_iter()
.map(Capability::new)
.collect::<Result<Vec<_>, _>>()?
.into_iter()
.filter(|capability| {
client_hello.required_capabilities().contains(capability)
|| client_hello.optional_capabilities().contains(capability)
}),
)?;
if let Some(capability) = client_hello
.required_capabilities()
.iter()
.find(|capability| !host_capabilities.contains(capability))
{
writer
.write(&HostToClient::HandshakeRejected {
reason: HandshakeRejectReason::MissingRequiredCapability {
capability: capability.clone(),
},
})
.await
.context("failed to reject unsupported code-mode capability")?;
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 {
max_pending_delegate_calls,
})
}
struct HostState {
sessions: Mutex<HashMap<SessionId, Arc<InProcessCodeModeSession>>>,
limits: Arc<HostLimits>,
seen_session_ids: Mutex<SeenSessionIds>,
requests: Mutex<RequestRegistry>,
request_tasks: TaskTracker,
closing: AtomicBool,
peer: Arc<HostPeer>,
}
impl HostState {
fn spawn_request(
self: &Arc<Self>,
request_id: RequestId,
request: HostRequest,
) -> Result<(), anyhow::Error> {
// Include host admission and scheduling before the runtime observes the request.
let received_at = Instant::now();
let request_kind = RequestKind::from(&request);
let (cancellation, yield_signal) = self
.requests
.lock()
.unwrap_or_else(PoisonError::into_inner)
.start(request_id, request_kind)?;
let Ok(permit) = self.limits.request_permit() else {
self.respond(
request_id,
Err("code-mode host has too many in-flight requests".to_string()),
);
self.finish_request(request_id);
return Ok(());
};
let state = Arc::clone(self);
let request_task = self.request_tasks.spawn(async move {
let _permit = permit;
state
.handle_request(request_id, request, cancellation, yield_signal, received_at)
.await;
state.finish_request(request_id);
});
self.supervise_request_task(request_task);
Ok(())
}
fn supervise_request_task(&self, task: tokio::task::JoinHandle<()>) {
let peer = Arc::clone(&self.peer);
tokio::spawn(async move {
if let Err(err) = task.await {
peer.fail(format!("code-mode request task failed: {err}"));
}
});
}
#[tracing::instrument(
name = "code_mode_host.request",
level = "info",
skip_all,
fields(
otel.name = "code_mode_host.request",
request.kind = RequestKind::from(&request).as_str(),
request.id = ?request_id,
call_id = tracing::field::Empty,
)
)]
async fn handle_request(
&self,
request_id: RequestId,
request: HostRequest,
cancellation: CancellationToken,
yield_signal: CancellationToken,
received_at: Instant,
) {
if let HostRequest::Execute { request, .. } = &request {
tracing::Span::current().record("call_id", request.tool_call_id.as_str());
}
if self.closing.load(Ordering::Acquire) {
self.respond(
request_id,
Err("code-mode host is shutting down".to_string()),
);
return;
}
match request {
HostRequest::OpenSession {
session_id,
cell_execution_limits,
output_truncation,
} => {
let result = CodeModeSessionCellExecutionLimits::try_from(
cell_execution_limits.unwrap_or_default(),
)
.map_err(|error| format!("invalid code-mode session execution limits: {error}"))
.and_then(|limits| {
self.open_session(
session_id.clone(),
CodeModeSessionOptions {
cell_execution_limits: limits,
output_truncation: output_truncation
.map(CodeModeOutputTruncation::from)
.unwrap_or_default(),
},
)
})
.map(|()| HostResponse::SessionReady { session_id });
self.respond(request_id, result);
}
HostRequest::Execute {
session_id,
request,
} => {
if cancellation.is_cancelled() {
self.respond(request_id, Err("code-mode request cancelled".to_string()));
return;
}
let request = match request.try_into() {
Ok(request) => request,
Err(err) => {
self.respond(
request_id,
Err(format!("invalid code-mode execute request: {err}")),
);
return;
}
};
let session = match self.session(&session_id) {
Ok(session) => session,
Err(err) => {
self.respond(request_id, Err(err));
return;
}
};
let Ok(active_cell_permit) = self.limits.cell_permit() else {
self.respond(
request_id,
Err("code-mode host has too many active cells".to_string()),
);
return;
};
let delegate = Arc::new(RemoteDelegate::new(
session_id.clone(),
Arc::clone(&self.peer),
));
let result = session.execute(request, delegate, Some(yield_signal)).await;
match result {
Ok(started) => {
let cell_id = started.cell_id.clone();
self.respond(
request_id,
Ok(HostResponse::ExecutionStarted {
cell_id: cell_id.into(),
}),
);
let initial_response_sent = self.peer.start_cell(
session_id,
request_id,
started,
active_cell_permit,
received_at,
);
let _ = initial_response_sent.await;
}
Err(err) => self.respond(request_id, Err(err)),
}
}
HostRequest::Wait {
session_id,
request,
max_output_tokens,
} => {
let result = match self.session(&session_id) {
Ok(session) => {
tokio::select! {
biased;
_ = cancellation.cancelled() => {
Err("code-mode request cancelled".to_string())
}
result = session.wait(WaitRequest {
max_output_tokens,
..request.into()
}, Some(yield_signal)) => result.and_then(|outcome| {
let outcome = outcome.with_code_mode_host_duration(received_at.elapsed());
Ok(HostResponse::WaitCompleted {
outcome: WireWaitOutcome::try_from(outcome)
.map_err(|error| error.to_string())?,
})
}),
}
}
Err(err) => Err(err),
};
self.respond(request_id, result);
}
HostRequest::Terminate {
session_id,
cell_id,
max_output_tokens,
} => {
let result = match self.session(&session_id) {
Ok(session) => session
.terminate(TerminateRequest {
cell_id: cell_id.into(),
max_output_tokens,
})
.await
.and_then(|outcome| {
let outcome =
outcome.with_code_mode_host_duration(received_at.elapsed());
Ok(HostResponse::WaitCompleted {
outcome: WireWaitOutcome::try_from(outcome)
.map_err(|error| error.to_string())?,
})
}),
Err(err) => Err(err),
};
self.respond(request_id, result);
}
HostRequest::ShutdownSession { session_id } => {
let session = self
.sessions
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remove(&session_id);
let result = match session {
Some(session) => match session.shutdown().await {
Ok(()) => {
self.peer.wait_for_session_cells(&session_id).await;
Ok(HostResponse::SessionClosed { session_id })
}
Err(err) => Err(err),
},
None => Err(format!("unknown code-mode session {session_id}")),
};
self.respond(request_id, result);
}
}
}
fn open_session(
&self,
session_id: SessionId,
options: CodeModeSessionOptions,
) -> Result<(), String> {
let mut sessions = self.sessions.lock().unwrap_or_else(PoisonError::into_inner);
if sessions.contains_key(&session_id) {
return Err(format!(
"code-mode session ID `{session_id}` is already open"
));
}
if self.closing.load(Ordering::Acquire) {
return Err("code-mode host is shutting down".to_string());
}
if !self
.seen_session_ids
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remember(session_id.clone())
{
return Err(format!("code-mode session ID `{session_id}` was reused"));
}
let peer = Arc::downgrade(&self.peer);
let task_failure_handler = Arc::new(move |reason| {
if let Some(peer) = peer.upgrade() {
peer.fail(reason);
}
});
sessions.insert(
session_id,
Arc::new(InProcessCodeModeSession::with_options(
options,
Some(task_failure_handler),
)),
);
Ok(())
}
fn session(&self, session_id: &SessionId) -> Result<Arc<InProcessCodeModeSession>, String> {
self.sessions
.lock()
.unwrap_or_else(PoisonError::into_inner)
.get(session_id)
.cloned()
.ok_or_else(|| format!("unknown code-mode session {session_id}"))
}
fn respond(&self, id: RequestId, result: Result<HostResponse, String>) {
self.peer.respond(id, result);
}
fn cancel_request(&self, request_id: RequestId) {
self.requests
.lock()
.unwrap_or_else(PoisonError::into_inner)
.cancel(request_id);
}
fn yield_request(&self, request_id: RequestId) {
self.requests
.lock()
.unwrap_or_else(PoisonError::into_inner)
.yield_observation(request_id);
}
fn finish_request(&self, request_id: RequestId) {
self.requests
.lock()
.unwrap_or_else(PoisonError::into_inner)
.finish(request_id);
}
async fn disconnect(&self) {
self.closing.store(true, Ordering::Release);
self.requests
.lock()
.unwrap_or_else(PoisonError::into_inner)
.cancel_all();
self.request_tasks.close();
self.request_tasks.wait().await;
let sessions = self
.sessions
.lock()
.unwrap_or_else(PoisonError::into_inner)
.drain()
.map(|(_, session)| session)
.collect::<Vec<_>>();
for session in sessions {
let _ = session.shutdown().await;
}
}
}
#[derive(Clone, Copy)]
enum RequestKind {
OpenSession,
Execute,
Wait,
Terminate,
ShutdownSession,
}
impl RequestKind {
fn from(request: &HostRequest) -> Self {
match request {
HostRequest::OpenSession { .. } => Self::OpenSession,
HostRequest::Execute { .. } => Self::Execute,
HostRequest::Wait { .. } => Self::Wait,
HostRequest::Terminate { .. } => Self::Terminate,
HostRequest::ShutdownSession { .. } => Self::ShutdownSession,
}
}
fn is_cancellable(self) -> bool {
matches!(self, Self::Execute | Self::Wait)
}
fn as_str(self) -> &'static str {
match self {
Self::OpenSession => "open_session",
Self::Execute => "execute",
Self::Wait => "wait",
Self::Terminate => "terminate",
Self::ShutdownSession => "shutdown_session",
}
}
}
struct ActiveRequest {
kind: RequestKind,
cancellation: CancellationToken,
yield_signal: CancellationToken,
}
#[derive(Default)]
struct RequestRegistry {
active: HashMap<RequestId, ActiveRequest>,
recent: HashSet<RequestId>,
recent_order: VecDeque<RequestId>,
}
impl RequestRegistry {
fn start(
&mut self,
request_id: RequestId,
kind: RequestKind,
) -> Result<(CancellationToken, CancellationToken), anyhow::Error> {
if self.active.contains_key(&request_id) || self.recent.contains(&request_id) {
anyhow::bail!("duplicate code-mode request ID {request_id:?}");
}
let cancellation = CancellationToken::new();
let yield_signal = CancellationToken::new();
self.active.insert(
request_id,
ActiveRequest {
kind,
cancellation: cancellation.clone(),
yield_signal: yield_signal.clone(),
},
);
Ok((cancellation, yield_signal))
}
fn cancel(&mut self, request_id: RequestId) {
if let Some(request) = self.active.get(&request_id)
&& request.kind.is_cancellable()
{
request.cancellation.cancel();
}
}
fn yield_observation(&self, request_id: RequestId) {
if let Some(request) = self.active.get(&request_id)
&& matches!(request.kind, RequestKind::Execute | RequestKind::Wait)
{
request.yield_signal.cancel();
}
}
fn finish(&mut self, request_id: RequestId) {
if self.active.remove(&request_id).is_none() {
return;
}
self.recent.insert(request_id);
self.recent_order.push_back(request_id);
while self.recent_order.len() > MAX_RECENT_REQUEST_IDS {
if let Some(expired) = self.recent_order.pop_front() {
self.recent.remove(&expired);
}
}
}
fn cancel_all(&self) {
for request in self.active.values() {
request.cancellation.cancel();
}
}
}
#[derive(Default)]
struct SeenSessionIds {
ids: HashSet<SessionId>,
order: VecDeque<SessionId>,
}
impl SeenSessionIds {
fn remember(&mut self, session_id: SessionId) -> bool {
if !self.ids.insert(session_id.clone()) {
return false;
}
self.order.push_back(session_id);
while self.order.len() > MAX_RECENT_SESSION_IDS {
if let Some(expired) = self.order.pop_front() {
self.ids.remove(&expired);
}
}
true
}
}
#[cfg(test)]
#[path = "host_tests.rs"]
mod tests;