Implement the gRPC code-mode host service (#37530)

## What changed

- Export `GrpcCodeModeHost` as a transport-independent implementation of the
  code-mode gRPC API.
- Support leased sessions, execution and wait lifecycle operations, filtered
  nested tool-call subscriptions, tool completions, and notification
  acknowledgements.
- Share host-wide request and active-cell limits across the existing and gRPC
  transports, and bound identifiers, metadata, subscriptions, and pending
  callbacks.

## Testing

- Add coverage for request conversion, ordered callback routing, cancellation,
  session cleanup, backpressure, malformed input, and resource-limit handling.

GitOrigin-RevId: f146ba7e6fe4e4aa02f25dd3f961120980516d0e

codex-rs/Cargo.lock

@@ -2569,11 +2569,14 @@ dependencies = [
  "codex-utils-cargo-bin",
  "futures",
  "pretty_assertions",
+ "prost",
  "serde_json",
  "tempfile",
  "tokio",
+ "tokio-stream",
  "tokio-tungstenite",
  "tokio-util",
+ "tonic",
  "tracing",
  "tracing-subscriber",
  "uuid",

codex-rs/code-mode-host/Cargo.toml

@@ -22,18 +22,21 @@ axum = { workspace = true, features = ["http1", "tokio", "ws"] }
 clap = { workspace = true, features = ["derive"] }
 codex-code-mode-protocol = { workspace = true }
 codex-code-mode-runtime = { workspace = true }
+codex-protocol = { workspace = true }
 futures = { workspace = true }
+prost = "0.14.3"
+serde_json = { workspace = true }
 tokio = { workspace = true, features = ["io-std", "io-util", "macros", "net", "process", "rt", "sync", "time"] }
+tokio-stream = { workspace = true }
 tokio-util = { workspace = true, features = ["rt"] }
+tonic = { workspace = true }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true }
 uuid = { workspace = true, features = ["v4"] }
 
 [dev-dependencies]
 codex-code-mode = { workspace = true }
-codex-protocol = { workspace = true }
 codex-utils-cargo-bin = { workspace = true }
 pretty_assertions = { workspace = true }
-serde_json = { workspace = true }
 tempfile = { workspace = true }
 tokio-tungstenite = { workspace = true }

codex-rs/code-mode-host/src/grpc/conversions.rs

@@ -0,0 +1,174 @@
+use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
+use codex_code_mode_protocol::CodeModeToolKind;
+use codex_code_mode_protocol::ExecuteRequest;
+use codex_code_mode_protocol::FunctionCallOutputContentItem;
+use codex_code_mode_protocol::ImageDetail;
+use codex_code_mode_protocol::RuntimeResponse;
+use codex_code_mode_protocol::ToolDefinition;
+use codex_code_mode_protocol::WaitOutcome;
+use codex_code_mode_protocol::grpc as proto;
+use codex_protocol::ToolName;
+use serde_json::Value as JsonValue;
+use tonic::Status;
+
+use super::validation;
+
+pub(super) fn session_limits(
+    limits: Option<proto::SessionCellExecutionLimits>,
+) -> Result<CodeModeSessionCellExecutionLimits, Status> {
+    let limits = limits.unwrap_or_default();
+    Ok(CodeModeSessionCellExecutionLimits {
+        max_yield_time_ms: limits.max_yield_time_ms,
+        max_heap_size_bytes: limits
+            .max_heap_size_bytes
+            .map(usize::try_from)
+            .transpose()
+            .map_err(|_| Status::invalid_argument("maximum heap size exceeds this platform"))?,
+    })
+}
+
+pub(super) fn execute_request(request: proto::ExecuteRequest) -> Result<ExecuteRequest, Status> {
+    validation::identifier(&request.tool_call_id, "tool call ID")?;
+    if request.enabled_tools.len() > validation::MAX_TOOL_DEFINITIONS {
+        return Err(Status::invalid_argument(format!(
+            "code-mode execution exceeds {} enabled tools",
+            validation::MAX_TOOL_DEFINITIONS
+        )));
+    }
+    Ok(ExecuteRequest {
+        tool_call_id: request.tool_call_id,
+        source: request.source,
+        enabled_tools: request
+            .enabled_tools
+            .into_iter()
+            .map(tool_definition)
+            .collect::<Result<Vec<_>, _>>()?,
+        yield_time_ms: request.yield_time_ms,
+        max_output_tokens: request
+            .max_output_tokens
+            .map(usize::try_from)
+            .transpose()
+            .map_err(|_| Status::invalid_argument("maximum output tokens exceeds this platform"))?,
+    })
+}
+
+fn tool_definition(definition: proto::ToolDefinition) -> Result<ToolDefinition, Status> {
+    validation::identifier(&definition.name, "tool definition name")?;
+    validation::bounded(
+        &definition.description,
+        validation::MAX_TOOL_DESCRIPTION_BYTES,
+        "tool description",
+    )?;
+    let name = definition
+        .tool_name
+        .ok_or_else(|| Status::invalid_argument("tool definition is missing its tool name"))?;
+    validation::tool_name(&name)?;
+    Ok(ToolDefinition {
+        name: definition.name,
+        tool_name: ToolName::new(name.namespace, name.name),
+        description: definition.description,
+        kind: match proto::ToolKind::try_from(definition.kind) {
+            Ok(proto::ToolKind::Function) => CodeModeToolKind::Function,
+            Ok(proto::ToolKind::Freeform) => CodeModeToolKind::Freeform,
+            Ok(proto::ToolKind::Unspecified) | Err(_) => {
+                return Err(Status::invalid_argument(
+                    "tool definition has an invalid kind",
+                ));
+            }
+        },
+        input_schema: json_field(definition.input_schema_json, "input schema")?,
+        output_schema: json_field(definition.output_schema_json, "output schema")?,
+    })
+}
+
+fn json_field(value: Option<Vec<u8>>, field: &str) -> Result<Option<JsonValue>, Status> {
+    value
+        .map(|value| {
+            serde_json::from_slice(&value)
+                .map_err(|error| Status::invalid_argument(format!("invalid tool {field}: {error}")))
+        })
+        .transpose()
+}
+
+pub(super) fn execution_outcome(response: RuntimeResponse) -> proto::ExecutionOutcome {
+    let (cell_id, content_items, outcome) = match response {
+        RuntimeResponse::Yielded {
+            cell_id,
+            content_items,
+        } => (
+            cell_id,
+            content_items,
+            proto::execution_outcome::Outcome::Yielded(proto::ExecutionYielded {}),
+        ),
+        RuntimeResponse::Terminated {
+            cell_id,
+            content_items,
+        } => (
+            cell_id,
+            content_items,
+            proto::execution_outcome::Outcome::Terminated(proto::ExecutionTerminated {}),
+        ),
+        RuntimeResponse::Result {
+            cell_id,
+            content_items,
+            error_text,
+        } => (
+            cell_id,
+            content_items,
+            proto::execution_outcome::Outcome::Completed(proto::ExecutionCompleted { error_text }),
+        ),
+    };
+    proto::ExecutionOutcome {
+        cell_id: cell_id.to_string(),
+        content_items: content_items.into_iter().map(content_item).collect(),
+        outcome: Some(outcome),
+    }
+}
+
+pub(super) fn wait_response(outcome: WaitOutcome) -> proto::WaitResponse {
+    let state = match outcome {
+        WaitOutcome::LiveCell(response) => {
+            proto::wait_response::State::LiveCell(execution_outcome(response))
+        }
+        WaitOutcome::MissingCell(response) => {
+            proto::wait_response::State::MissingCell(execution_outcome(response))
+        }
+    };
+    proto::WaitResponse { state: Some(state) }
+}
+
+fn content_item(item: FunctionCallOutputContentItem) -> proto::ContentItem {
+    let item = match item {
+        FunctionCallOutputContentItem::InputText { text } => {
+            proto::content_item::Item::Text(proto::TextContent { text })
+        }
+        FunctionCallOutputContentItem::InputImage { image_url, detail } => {
+            proto::content_item::Item::Image(proto::ImageContent {
+                image_url,
+                detail: detail.map(|detail| {
+                    (match detail {
+                        ImageDetail::Auto => proto::ImageDetail::Auto,
+                        ImageDetail::Low => proto::ImageDetail::Low,
+                        ImageDetail::High => proto::ImageDetail::High,
+                        ImageDetail::Original => proto::ImageDetail::Original,
+                    }) as i32
+                }),
+            })
+        }
+        FunctionCallOutputContentItem::InputAudio { audio_url } => {
+            proto::content_item::Item::Audio(proto::AudioContent { audio_url })
+        }
+    };
+    proto::ContentItem { item: Some(item) }
+}
+
+pub(super) fn tool_kind(kind: CodeModeToolKind) -> i32 {
+    match kind {
+        CodeModeToolKind::Function => proto::ToolKind::Function as i32,
+        CodeModeToolKind::Freeform => proto::ToolKind::Freeform as i32,
+    }
+}
+
+#[cfg(test)]
+#[path = "conversions_tests.rs"]
+mod tests;

codex-rs/code-mode-host/src/grpc/conversions_tests.rs

@@ -0,0 +1,113 @@
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::FunctionCallOutputContentItem;
+use codex_code_mode_protocol::ImageDetail;
+use codex_code_mode_protocol::RuntimeResponse;
+use codex_code_mode_protocol::grpc as proto;
+use pretty_assertions::assert_eq;
+use tonic::Code;
+
+use super::execute_request;
+use super::execution_outcome;
+
+#[test]
+fn rejects_missing_names_unknown_tool_kinds_and_invalid_json_schemas() {
+    let definition = proto::ToolDefinition {
+        name: "echo".to_string(),
+        tool_name: None,
+        description: String::new(),
+        kind: proto::ToolKind::Function as i32,
+        input_schema_json: None,
+        output_schema_json: None,
+    };
+    let request = |definition| proto::ExecuteRequest {
+        session_id: "session".to_string(),
+        execution_id: "execution".to_string(),
+        tool_call_id: "call".to_string(),
+        source: String::new(),
+        enabled_tools: vec![definition],
+        yield_time_ms: None,
+        max_output_tokens: None,
+    };
+
+    assert_eq!(
+        execute_request(request(definition.clone()))
+            .unwrap_err()
+            .code(),
+        Code::InvalidArgument
+    );
+    let definition = proto::ToolDefinition {
+        tool_name: Some(proto::ToolName {
+            name: "echo".to_string(),
+            namespace: Some("tools".to_string()),
+        }),
+        ..definition
+    };
+    assert_eq!(
+        execute_request(request(proto::ToolDefinition {
+            kind: proto::ToolKind::Unspecified as i32,
+            ..definition.clone()
+        }))
+        .unwrap_err()
+        .code(),
+        Code::InvalidArgument
+    );
+    assert_eq!(
+        execute_request(request(proto::ToolDefinition {
+            input_schema_json: Some(b"not-json".to_vec()),
+            ..definition
+        }))
+        .unwrap_err()
+        .code(),
+        Code::InvalidArgument
+    );
+}
+
+#[test]
+fn maps_text_image_audio_and_terminal_error_without_losing_details() {
+    let outcome = execution_outcome(RuntimeResponse::Result {
+        cell_id: CellId::new("cell".to_string()),
+        content_items: vec![
+            FunctionCallOutputContentItem::InputText {
+                text: "hello".to_string(),
+            },
+            FunctionCallOutputContentItem::InputImage {
+                image_url: "data:image/png;base64,YQ==".to_string(),
+                detail: Some(ImageDetail::Original),
+            },
+            FunctionCallOutputContentItem::InputAudio {
+                audio_url: "data:audio/wav;base64,YQ==".to_string(),
+            },
+        ],
+        error_text: Some("failed".to_string()),
+    });
+
+    assert_eq!(
+        outcome,
+        proto::ExecutionOutcome {
+            cell_id: "cell".to_string(),
+            content_items: vec![
+                proto::ContentItem {
+                    item: Some(proto::content_item::Item::Text(proto::TextContent {
+                        text: "hello".to_string(),
+                    })),
+                },
+                proto::ContentItem {
+                    item: Some(proto::content_item::Item::Image(proto::ImageContent {
+                        image_url: "data:image/png;base64,YQ==".to_string(),
+                        detail: Some(proto::ImageDetail::Original as i32),
+                    })),
+                },
+                proto::ContentItem {
+                    item: Some(proto::content_item::Item::Audio(proto::AudioContent {
+                        audio_url: "data:audio/wav;base64,YQ==".to_string(),
+                    })),
+                },
+            ],
+            outcome: Some(proto::execution_outcome::Outcome::Completed(
+                proto::ExecutionCompleted {
+                    error_text: Some("failed".to_string()),
+                },
+            )),
+        }
+    );
+}

codex-rs/code-mode-host/src/grpc/delegate.rs

@@ -0,0 +1,180 @@
+use std::sync::Arc;
+use std::sync::Weak;
+
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::CodeModeNestedToolCall;
+use codex_code_mode_protocol::CodeModeSessionDelegate;
+use codex_code_mode_protocol::NotificationFuture;
+use codex_code_mode_protocol::ToolInvocationFuture;
+use codex_code_mode_protocol::grpc as proto;
+use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+use super::session::GrpcSession;
+
+pub(super) struct GrpcDelegate {
+    session: Weak<GrpcSession>,
+}
+
+impl GrpcDelegate {
+    pub(super) fn new(session: Weak<GrpcSession>) -> Self {
+        Self { session }
+    }
+}
+
+impl CodeModeSessionDelegate for GrpcDelegate {
+    fn invoke_tool<'a>(
+        &'a self,
+        invocation: CodeModeNestedToolCall,
+        cancellation: CancellationToken,
+    ) -> ToolInvocationFuture<'a> {
+        Box::pin(async move {
+            let session = self
+                .session
+                .upgrade()
+                .ok_or_else(|| "code-mode session is closed".to_string())?;
+            let _permit = session.delegate_permit()?;
+            let execution_id = session
+                .execution_id(invocation.cell_id.as_str(), &cancellation)
+                .await?;
+            let input_json = invocation
+                .input
+                .as_ref()
+                .map(serde_json::to_vec)
+                .transpose()
+                .map_err(|error| format!("failed to encode code-mode tool input: {error}"))?;
+            let invocation_id = Uuid::new_v4();
+            let (response, receiver) = oneshot::channel();
+            session
+                .dispatch_tool(
+                    invocation,
+                    execution_id,
+                    invocation_id,
+                    input_json,
+                    response,
+                    &cancellation,
+                )
+                .await?;
+            let mut pending = PendingCallback::tool(Arc::clone(&session), invocation_id);
+            let result = tokio::select! {
+                biased;
+                result = receiver => result
+                    .map_err(|_| "code-mode client closed before returning tool output".to_string())?,
+                _ = cancellation.cancelled() => {
+                    Err("code mode delegate request cancelled".to_string())
+                }
+                _ = session.closed.cancelled() => {
+                    Err("code-mode session closed before returning tool output".to_string())
+                }
+            };
+            if result.is_ok() {
+                pending.disarm();
+            }
+            result
+        })
+    }
+
+    fn notify<'a>(
+        &'a self,
+        call_id: String,
+        cell_id: CellId,
+        text: String,
+        cancellation: CancellationToken,
+    ) -> NotificationFuture<'a> {
+        Box::pin(async move {
+            let session = self
+                .session
+                .upgrade()
+                .ok_or_else(|| "code-mode session is closed".to_string())?;
+            let _permit = session.delegate_permit()?;
+            let execution_id = session
+                .execution_id(cell_id.as_str(), &cancellation)
+                .await?;
+            let notification_id = Uuid::new_v4();
+            let (response, receiver) = oneshot::channel();
+            session
+                .begin_notification(
+                    notification_id,
+                    proto::Notification {
+                        notification_id: notification_id.to_string(),
+                        execution_id,
+                        cell_id: cell_id.to_string(),
+                        call_id,
+                        text,
+                    },
+                    response,
+                    &cancellation,
+                )
+                .await?;
+            let mut pending = PendingCallback::notification(Arc::clone(&session), notification_id);
+            tokio::select! {
+                biased;
+                result = receiver => {
+                    result.map_err(|_| {
+                        "code-mode client closed before acknowledging notification".to_string()
+                    })?;
+                    pending.disarm();
+                    Ok(())
+                }
+                _ = cancellation.cancelled() => {
+                    Err("code mode notification cancelled".to_string())
+                }
+                _ = session.closed.cancelled() => {
+                    Err("code-mode session closed before acknowledging notification".to_string())
+                }
+            }
+        })
+    }
+
+    fn cell_closed(&self, cell_id: &CellId) {
+        if let Some(session) = self.session.upgrade() {
+            session.close_cell(cell_id.as_str());
+        }
+    }
+}
+
+enum CallbackKind {
+    Tool,
+    Notification,
+}
+
+struct PendingCallback {
+    session: Arc<GrpcSession>,
+    id: Option<Uuid>,
+    kind: CallbackKind,
+}
+
+impl PendingCallback {
+    fn tool(session: Arc<GrpcSession>, id: Uuid) -> Self {
+        Self {
+            session,
+            id: Some(id),
+            kind: CallbackKind::Tool,
+        }
+    }
+
+    fn notification(session: Arc<GrpcSession>, id: Uuid) -> Self {
+        Self {
+            session,
+            id: Some(id),
+            kind: CallbackKind::Notification,
+        }
+    }
+
+    fn disarm(&mut self) {
+        self.id = None;
+    }
+}
+
+impl Drop for PendingCallback {
+    fn drop(&mut self) {
+        let Some(id) = self.id.take() else {
+            return;
+        };
+        match self.kind {
+            CallbackKind::Tool => self.session.cancel_invocation(id),
+            CallbackKind::Notification => self.session.cancel_notification(id),
+        }
+    }
+}

codex-rs/code-mode-host/src/grpc/events.rs

@@ -0,0 +1,151 @@
+use std::sync::Arc;
+
+use codex_code_mode_protocol::grpc as proto;
+use codex_code_mode_protocol::host::MAX_FRAME_BYTES;
+use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
+use prost::Message;
+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 tonic::Status;
+
+use crate::MAX_ACTIVE_CELLS;
+
+const MAX_BUFFERED_CONTROL_EVENTS: usize = MAX_PENDING_DELEGATE_CALLS * 2 + MAX_ACTIVE_CELLS;
+
+#[derive(Clone)]
+pub(super) struct EventSender {
+    sender: mpsc::UnboundedSender<QueuedEvent>,
+    permits: Arc<Semaphore>,
+    closed: CancellationToken,
+    writer: TaskTracker,
+}
+
+struct QueuedEvent {
+    message: proto::SessionEvent,
+    _queue_permit: OwnedSemaphorePermit,
+    _cell_permit: Option<OwnedSemaphorePermit>,
+}
+
+impl EventSender {
+    pub(super) fn new(
+        output: mpsc::Sender<Result<proto::SessionEvent, Status>>,
+        closed: CancellationToken,
+    ) -> Self {
+        let (sender, mut receiver) = mpsc::unbounded_channel::<QueuedEvent>();
+        let writer_closed = closed.clone();
+        let writer = TaskTracker::new();
+        writer.spawn(async move {
+            loop {
+                let event = tokio::select! {
+                    _ = writer_closed.cancelled() => return,
+                    event = receiver.recv() => match event {
+                        Some(event) => event,
+                        None => return,
+                    },
+                };
+                let QueuedEvent {
+                    message,
+                    _queue_permit,
+                    _cell_permit,
+                } = event;
+                tokio::select! {
+                    _ = writer_closed.cancelled() => return,
+                    result = output.send(Ok(message)) => {
+                        if result.is_err() {
+                            writer_closed.cancel();
+                            return;
+                        }
+                    }
+                }
+                drop(_queue_permit);
+                drop(_cell_permit);
+            }
+        });
+        writer.close();
+        Self {
+            sender,
+            permits: Arc::new(Semaphore::new(MAX_BUFFERED_CONTROL_EVENTS)),
+            closed,
+            writer,
+        }
+    }
+
+    pub(super) async fn shutdown(&self) {
+        self.closed.cancel();
+        self.writer.wait().await;
+    }
+
+    pub(super) async fn send(
+        &self,
+        event: proto::session_event::Event,
+        cancellation: &CancellationToken,
+    ) -> Result<(), String> {
+        let message = validate_event(event)?;
+        let permit = tokio::select! {
+            biased;
+            _ = self.closed.cancelled() => {
+                return Err("code-mode session event stream is closed".to_string());
+            }
+            _ = cancellation.cancelled() => {
+                return Err("code-mode session event was cancelled".to_string());
+            }
+            permit = Arc::clone(&self.permits).acquire_owned() => permit
+                .map_err(|_| "code-mode session event queue is closed".to_string())?,
+        };
+        self.enqueue(message, permit, /*cell_permit*/ None)
+    }
+
+    pub(super) fn send_now(
+        &self,
+        event: proto::session_event::Event,
+        cell_permit: Option<OwnedSemaphorePermit>,
+    ) -> Result<(), String> {
+        let message = validate_event(event)?;
+        if self.closed.is_cancelled() {
+            return Err("code-mode session event stream is closed".to_string());
+        }
+        match Arc::clone(&self.permits).try_acquire_owned() {
+            Ok(permit) => self.enqueue(message, permit, cell_permit),
+            Err(TryAcquireError::NoPermits) => {
+                self.closed.cancel();
+                Err("code-mode session event queue is full".to_string())
+            }
+            Err(TryAcquireError::Closed) => {
+                self.closed.cancel();
+                Err("code-mode session event queue is closed".to_string())
+            }
+        }
+    }
+
+    fn enqueue(
+        &self,
+        message: proto::SessionEvent,
+        permit: OwnedSemaphorePermit,
+        cell_permit: Option<OwnedSemaphorePermit>,
+    ) -> Result<(), String> {
+        self.sender
+            .send(QueuedEvent {
+                message,
+                _queue_permit: permit,
+                _cell_permit: cell_permit,
+            })
+            .map_err(|_| {
+                self.closed.cancel();
+                "code-mode session event stream is closed".to_string()
+            })
+    }
+}
+
+fn validate_event(event: proto::session_event::Event) -> Result<proto::SessionEvent, String> {
+    let message = proto::SessionEvent { event: Some(event) };
+    if message.encoded_len() > MAX_FRAME_BYTES {
+        return Err(format!(
+            "code-mode session event exceeds the {MAX_FRAME_BYTES}-byte gRPC message limit"
+        ));
+    }
+    Ok(message)
+}

codex-rs/code-mode-host/src/grpc/mod.rs

@@ -0,0 +1,369 @@
+mod conversions;
+mod delegate;
+mod events;
+mod routing;
+mod session;
+mod validation;
+mod waits;
+
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+
+use codex_code_mode_protocol::CellId;
+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 futures::Stream;
+use futures::StreamExt;
+use tokio::sync::mpsc;
+use tokio_stream::wrappers::ReceiverStream;
+use tonic::Request;
+use tonic::Response;
+use tonic::Status;
+
+use self::session::GrpcHostState;
+use self::session::GrpcSession;
+use self::waits::WaitRegistration;
+
+type GrpcStream<T> = Pin<Box<dyn Stream<Item = Result<T, Status>> + Send + 'static>>;
+type GrpcFuture<'a, T> = Pin<Box<dyn Future<Output = Result<Response<T>, Status>> + Send + 'a>>;
+
+/// Serves transport-independent, leased code-mode sessions over gRPC.
+#[derive(Clone)]
+pub struct GrpcCodeModeHost {
+    state: Arc<GrpcHostState>,
+}
+
+impl GrpcCodeModeHost {
+    /// Creates a host with independent session, execution, and callback limits.
+    pub fn new() -> Self {
+        Self {
+            state: Arc::new(GrpcHostState::new()),
+        }
+    }
+
+    async fn open_session_request(
+        &self,
+        request: proto::OpenSessionRequest,
+    ) -> Result<Response<GrpcStream<proto::SessionEvent>>, Status> {
+        let _permit = self.state.request_permit()?;
+        let limits = conversions::session_limits(request.cell_execution_limits)?;
+        Ok(Response::new(self.state.open_session(limits)?))
+    }
+
+    async fn close_session_request(
+        &self,
+        request: proto::CloseSessionRequest,
+    ) -> Result<Response<proto::CloseSessionResponse>, Status> {
+        let _permit = self.state.control_permit()?;
+        self.state.close_session(&request.session_id).await?;
+        Ok(Response::new(proto::CloseSessionResponse {}))
+    }
+
+    async fn subscribe_request(
+        &self,
+        request: proto::SubscribeToToolCallsRequest,
+    ) -> Result<Response<GrpcStream<proto::ToolCall>>, Status> {
+        let _permit = self.state.request_permit()?;
+        let session = self.state.session(&request.session_id)?;
+        Ok(Response::new(session.subscribe(request.tool_names)?))
+    }
+
+    async fn complete_tool_request(
+        &self,
+        request: proto::CompleteToolCallRequest,
+    ) -> Result<Response<proto::CompleteToolCallResponse>, Status> {
+        let _permit = self.state.control_permit()?;
+        let session = self.state.session(&request.session_id)?;
+        let invocation_id = validation::uuid(&request.invocation_id, "tool invocation ID")?;
+        let result = match request.outcome {
+            Some(proto::complete_tool_call_request::Outcome::Succeeded(result)) => Ok(
+                serde_json::from_slice(&result.output_json).map_err(|error| {
+                    Status::invalid_argument(format!("invalid code-mode tool output JSON: {error}"))
+                })?,
+            ),
+            Some(proto::complete_tool_call_request::Outcome::Failed(error)) => {
+                validation::bounded(
+                    &error.message,
+                    validation::MAX_TOOL_ERROR_BYTES,
+                    "tool error message",
+                )?;
+                Err(error.message)
+            }
+            None => {
+                return Err(Status::invalid_argument(
+                    "tool completion is missing its outcome",
+                ));
+            }
+        };
+        session.complete_invocation(invocation_id, result)?;
+        Ok(Response::new(proto::CompleteToolCallResponse {}))
+    }
+
+    async fn acknowledge_notification_request(
+        &self,
+        request: proto::AcknowledgeNotificationRequest,
+    ) -> Result<Response<proto::AcknowledgeNotificationResponse>, Status> {
+        let _permit = self.state.control_permit()?;
+        let session = self.state.session(&request.session_id)?;
+        let notification_id = validation::uuid(&request.notification_id, "notification ID")?;
+        session.acknowledge_notification(notification_id)?;
+        Ok(Response::new(proto::AcknowledgeNotificationResponse {}))
+    }
+
+    async fn execute_request(
+        &self,
+        request: proto::ExecuteRequest,
+    ) -> Result<Response<GrpcStream<proto::ExecuteEvent>>, Status> {
+        let session = self.state.session(&request.session_id)?;
+        validation::identifier(&request.execution_id, "execution ID")?;
+        let request_permit = self.state.request_permit()?;
+        let execution_id = request.execution_id.clone();
+        let request = conversions::execute_request(request)?;
+        let cell_permit = self.state.cell_permit()?;
+        session.reserve_execution(&execution_id)?;
+        let mut admission = ExecutionAdmission {
+            session: Arc::clone(&session),
+            execution_id: Some(execution_id.clone()),
+        };
+        let started = tokio::select! {
+            _ = session.closed.cancelled() => {
+                return Err(Status::cancelled("code-mode session is closed"));
+            }
+            result = session.runtime.execute(request) => {
+                result.map_err(Status::failed_precondition)?
+            }
+        };
+        let cell_id = started.cell_id.clone();
+        session.admit_execution(execution_id.clone(), cell_id.to_string(), cell_permit)?;
+
+        let (sender, receiver) = mpsc::channel(/*buffer*/ 2);
+        sender
+            .try_send(Ok(proto::ExecuteEvent {
+                event: Some(proto::execute_event::Event::Started(
+                    proto::ExecutionStarted {
+                        execution_id,
+                        cell_id: cell_id.to_string(),
+                    },
+                )),
+            }))
+            .map_err(|_| Status::internal("failed to publish code-mode execution admission"))?;
+        tokio::spawn(async move {
+            let _request_permit = request_permit;
+            tokio::select! {
+                biased;
+                _ = sender.closed() => {}
+                response = started.initial_response() => {
+                    let event = response.map(|response| proto::ExecuteEvent {
+                        event: Some(proto::execute_event::Event::Outcome(
+                            conversions::execution_outcome(response),
+                        )),
+                    }).map_err(Status::internal);
+                    let _ = sender.send(event).await;
+                }
+                _ = session.closed.cancelled() => {}
+            }
+        });
+
+        let stream = ReceiverStream::new(receiver).inspect(move |event| {
+            if matches!(
+                event,
+                Ok(proto::ExecuteEvent {
+                    event: Some(proto::execute_event::Event::Outcome(_)),
+                })
+            ) {
+                admission.disarm();
+            }
+        });
+        Ok(Response::new(Box::pin(stream)))
+    }
+
+    async fn wait_request(
+        &self,
+        request: proto::WaitRequest,
+    ) -> Result<Response<proto::WaitResponse>, Status> {
+        let session = self.state.session(&request.session_id)?;
+        validation::identifier(&request.cell_id, "cell ID")?;
+        validation::identifier(&request.wait_id, "wait ID")?;
+        let _permit = self.state.request_permit()?;
+        let registration = WaitRegistration::new(Arc::clone(&session), request.wait_id)?;
+        let request = WaitRequest {
+            cell_id: CellId::new(request.cell_id),
+            yield_time_ms: request.yield_time_ms,
+        };
+        let outcome = tokio::select! {
+            biased;
+            _ = registration.cancellation().cancelled() => {
+                return Err(Status::cancelled("code-mode wait was cancelled"));
+            }
+            _ = session.closed.cancelled() => {
+                return Err(Status::cancelled("code-mode session is closed"));
+            }
+            outcome = session.runtime.wait(request) => {
+                outcome.map_err(Status::failed_precondition)?
+            }
+        };
+        Ok(Response::new(conversions::wait_response(outcome)))
+    }
+
+    async fn cancel_wait_request(
+        &self,
+        request: proto::CancelWaitRequest,
+    ) -> Result<Response<proto::CancelWaitResponse>, Status> {
+        let _permit = self.state.control_permit()?;
+        let session = self.state.session(&request.session_id)?;
+        validation::identifier(&request.wait_id, "wait ID")?;
+        session.cancel_wait(&request.wait_id).await?;
+        Ok(Response::new(proto::CancelWaitResponse {}))
+    }
+
+    async fn terminate_request(
+        &self,
+        request: proto::TerminateRequest,
+    ) -> Result<Response<proto::WaitResponse>, Status> {
+        let session = self.state.session(&request.session_id)?;
+        validation::identifier(&request.cell_id, "cell ID")?;
+        let _permit = self.state.request_permit()?;
+        let result = session.terminate(CellId::new(request.cell_id)).await?;
+        Ok(Response::new(conversions::wait_response(result)))
+    }
+}
+
+impl Default for GrpcCodeModeHost {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl CodeModeHost for GrpcCodeModeHost {
+    type OpenSessionStream = GrpcStream<proto::SessionEvent>;
+    type SubscribeToToolCallsStream = GrpcStream<proto::ToolCall>;
+    type ExecuteStream = GrpcStream<proto::ExecuteEvent>;
+
+    fn open_session<'a, 'async_trait>(
+        &'a self,
+        request: Request<proto::OpenSessionRequest>,
+    ) -> GrpcFuture<'async_trait, Self::OpenSessionStream>
+    where
+        'a: 'async_trait,
+        Self: 'async_trait,
+    {
+        Box::pin(self.open_session_request(request.into_inner()))
+    }
+
+    fn close_session<'a, 'async_trait>(
+        &'a self,
+        request: Request<proto::CloseSessionRequest>,
+    ) -> GrpcFuture<'async_trait, proto::CloseSessionResponse>
+    where
+        'a: 'async_trait,
+        Self: 'async_trait,
+    {
+        Box::pin(self.close_session_request(request.into_inner()))
+    }
+
+    fn subscribe_to_tool_calls<'a, 'async_trait>(
+        &'a self,
+        request: Request<proto::SubscribeToToolCallsRequest>,
+    ) -> GrpcFuture<'async_trait, Self::SubscribeToToolCallsStream>
+    where
+        'a: 'async_trait,
+        Self: 'async_trait,
+    {
+        Box::pin(self.subscribe_request(request.into_inner()))
+    }
+
+    fn complete_tool_call<'a, 'async_trait>(
+        &'a self,
+        request: Request<proto::CompleteToolCallRequest>,
+    ) -> GrpcFuture<'async_trait, proto::CompleteToolCallResponse>
+    where
+        'a: 'async_trait,
+        Self: 'async_trait,
+    {
+        Box::pin(self.complete_tool_request(request.into_inner()))
+    }
+
+    fn acknowledge_notification<'a, 'async_trait>(
+        &'a self,
+        request: Request<proto::AcknowledgeNotificationRequest>,
+    ) -> GrpcFuture<'async_trait, proto::AcknowledgeNotificationResponse>
+    where
+        'a: 'async_trait,
+        Self: 'async_trait,
+    {
+        Box::pin(self.acknowledge_notification_request(request.into_inner()))
+    }
+
+    fn execute<'a, 'async_trait>(
+        &'a self,
+        request: Request<proto::ExecuteRequest>,
+    ) -> GrpcFuture<'async_trait, Self::ExecuteStream>
+    where
+        'a: 'async_trait,
+        Self: 'async_trait,
+    {
+        Box::pin(self.execute_request(request.into_inner()))
+    }
+
+    fn wait<'a, 'async_trait>(
+        &'a self,
+        request: Request<proto::WaitRequest>,
+    ) -> GrpcFuture<'async_trait, proto::WaitResponse>
+    where
+        'a: 'async_trait,
+        Self: 'async_trait,
+    {
+        Box::pin(self.wait_request(request.into_inner()))
+    }
+
+    fn cancel_wait<'a, 'async_trait>(
+        &'a self,
+        request: Request<proto::CancelWaitRequest>,
+    ) -> GrpcFuture<'async_trait, proto::CancelWaitResponse>
+    where
+        'a: 'async_trait,
+        Self: 'async_trait,
+    {
+        Box::pin(self.cancel_wait_request(request.into_inner()))
+    }
+
+    fn terminate<'a, 'async_trait>(
+        &'a self,
+        request: Request<proto::TerminateRequest>,
+    ) -> GrpcFuture<'async_trait, proto::WaitResponse>
+    where
+        'a: 'async_trait,
+        Self: 'async_trait,
+    {
+        Box::pin(self.terminate_request(request.into_inner()))
+    }
+}
+
+struct ExecutionAdmission {
+    session: Arc<GrpcSession>,
+    execution_id: Option<String>,
+}
+
+impl ExecutionAdmission {
+    fn disarm(&mut self) {
+        self.execution_id = None;
+    }
+}
+
+impl Drop for ExecutionAdmission {
+    fn drop(&mut self) {
+        if let Some(execution_id) = self.execution_id.take() {
+            self.session.abandon_execution(&execution_id);
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "service_tests.rs"]
+mod tests;
+
+#[cfg(test)]
+#[path = "robustness_tests.rs"]
+mod robustness_tests;

codex-rs/code-mode-host/src/grpc/robustness_tests.rs

@@ -0,0 +1,665 @@
+use std::sync::Arc;
+use std::time::Duration;
+
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::CodeModeNestedToolCall;
+use codex_code_mode_protocol::CodeModeToolKind;
+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::MAX_FRAME_BYTES;
+use codex_protocol::ToolName;
+use futures::FutureExt;
+use futures::StreamExt;
+use pretty_assertions::assert_eq;
+use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
+use tonic::Code;
+use tonic::Request;
+use tonic::Status;
+use uuid::Uuid;
+
+use super::GrpcCodeModeHost;
+use super::tests::execute_events;
+use super::tests::execute_request;
+use super::tests::open_session;
+use super::tests::tool;
+use super::validation::MAX_IDENTIFIER_BYTES;
+use super::validation::MAX_TOOL_DESCRIPTION_BYTES;
+use super::validation::MAX_TOOL_FILTERS;
+use crate::MAX_ACTIVE_CELLS;
+use crate::MAX_IN_FLIGHT_REQUESTS;
+use crate::OUTGOING_CHANNEL_CAPACITY;
+
+fn assert_invalid<T>(result: Result<T, Status>) {
+    match result {
+        Ok(_) => panic!("expected an oversized gRPC field to be rejected"),
+        Err(error) => assert_eq!(error.code(), Code::InvalidArgument),
+    }
+}
+
+fn invocation(cell_id: &str, name: &str) -> CodeModeNestedToolCall {
+    CodeModeNestedToolCall {
+        cell_id: CellId::new(cell_id.to_string()),
+        runtime_tool_call_id: "runtime-call".to_string(),
+        tool_name: ToolName::plain(name),
+        tool_kind: CodeModeToolKind::Function,
+        input: None,
+    }
+}
+
+#[tokio::test]
+async fn rejects_oversized_identifiers_tool_metadata_and_subscription_filters() {
+    let host = GrpcCodeModeHost::new();
+    let (session_id, _events) = open_session(&host).await;
+    let oversized_id = "x".repeat(MAX_IDENTIFIER_BYTES + 1);
+
+    assert_invalid(
+        host.close_session(Request::new(proto::CloseSessionRequest {
+            session_id: oversized_id.clone(),
+        }))
+        .await,
+    );
+    assert_invalid(
+        host.close_session(Request::new(proto::CloseSessionRequest {
+            session_id: "not-a-uuid".to_string(),
+        }))
+        .await,
+    );
+    assert_invalid(
+        host.complete_tool_call(Request::new(proto::CompleteToolCallRequest {
+            session_id: session_id.clone(),
+            invocation_id: "not-a-uuid".to_string(),
+            outcome: Some(proto::complete_tool_call_request::Outcome::Succeeded(
+                proto::ToolCallSucceeded {
+                    output_json: b"null".to_vec(),
+                },
+            )),
+        }))
+        .await,
+    );
+    assert_invalid(
+        host.acknowledge_notification(Request::new(proto::AcknowledgeNotificationRequest {
+            session_id: session_id.clone(),
+            notification_id: "not-a-uuid".to_string(),
+        }))
+        .await,
+    );
+    assert_invalid(
+        host.execute(Request::new(execute_request(
+            &session_id,
+            &oversized_id,
+            "text(\"hello\");",
+        )))
+        .await,
+    );
+    assert_invalid(
+        host.subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest {
+            session_id: session_id.clone(),
+            tool_names: vec![proto::ToolName {
+                name: oversized_id,
+                namespace: None,
+            }],
+        }))
+        .await,
+    );
+    assert_invalid(
+        host.subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest {
+            session_id: session_id.clone(),
+            tool_names: vec![
+                proto::ToolName {
+                    name: "echo".to_string(),
+                    namespace: None,
+                };
+                MAX_TOOL_FILTERS + 1
+            ],
+        }))
+        .await,
+    );
+
+    let mut request = execute_request(&session_id, "metadata", "text(\"hello\");");
+    request.enabled_tools = vec![proto::ToolDefinition {
+        description: "x".repeat(MAX_TOOL_DESCRIPTION_BYTES + 1),
+        ..tool("echo")
+    }];
+    assert_invalid(host.execute(Request::new(request)).await);
+    assert!(host.state.session(&session_id).is_ok());
+}
+
+#[tokio::test]
+async fn dropping_an_unread_buffered_execution_outcome_retires_its_cell() {
+    let host = GrpcCodeModeHost::new();
+    let (session_id, mut events) = open_session(&host).await;
+    let execution = host
+        .execute(Request::new(execute_request(
+            &session_id,
+            "execution-abandoned",
+            "await new Promise(() => {});",
+        )))
+        .await
+        .expect("start execution")
+        .into_inner();
+    let _reserved_permits = (0..MAX_IN_FLIGHT_REQUESTS - 1)
+        .map(|_| host.state.request_permit().expect("reserve request permit"))
+        .collect::<Vec<_>>();
+    let _execution_permit = tokio::time::timeout(Duration::from_secs(/*secs*/ 2), async {
+        loop {
+            if let Ok(permit) = host.state.request_permit() {
+                return permit;
+            }
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .expect("execution outcome should be buffered before dropping its unread stream");
+
+    drop(execution);
+
+    assert_eq!(
+        tokio::time::timeout(Duration::from_secs(/*secs*/ 2), events.next())
+            .await
+            .expect("abandoned cell should close")
+            .expect("cell closed event")
+            .expect("session event"),
+        proto::SessionEvent {
+            event: Some(proto::session_event::Event::CellClosed(proto::CellClosed {
+                execution_id: "execution-abandoned".to_string(),
+                cell_id: "1".to_string(),
+                final_tool_call_sequence: 0,
+            })),
+        }
+    );
+}
+
+#[tokio::test]
+async fn closing_a_session_releases_buffered_cell_permits() {
+    let host = GrpcCodeModeHost::new();
+    let (session_id, _events) = open_session(&host).await;
+    let session = host.state.session(&session_id).expect("open session");
+    let mut permits = (0..MAX_ACTIVE_CELLS)
+        .map(|_| {
+            host.state
+                .cell_permit()
+                .expect("reserve active-cell permit")
+        })
+        .collect::<Vec<_>>();
+    session
+        .send_event_now(
+            proto::session_event::Event::CellClosed(proto::CellClosed {
+                execution_id: "execution-queued".to_string(),
+                cell_id: "1".to_string(),
+                final_tool_call_sequence: 0,
+            }),
+            permits.pop(),
+        )
+        .expect("queue cell closure");
+
+    host.close_session(Request::new(proto::CloseSessionRequest { session_id }))
+        .await
+        .expect("close session");
+
+    assert!(
+        host.state.cell_permit().is_ok(),
+        "closing a session must release its queued active-cell permits"
+    );
+}
+
+#[tokio::test]
+async fn dropping_a_lease_after_its_host_shuts_down_closes_the_session() {
+    let host = GrpcCodeModeHost::new();
+    let (session_id, lease) = open_session(&host).await;
+    let session = host.state.session(&session_id).expect("open session");
+
+    drop(host);
+    drop(lease);
+
+    tokio::time::timeout(Duration::from_secs(/*secs*/ 2), session.closed.cancelled())
+        .await
+        .expect("dropping a lease must close its session after the host disappears");
+}
+
+#[tokio::test]
+async fn session_closure_cancels_pending_termination() {
+    let host = GrpcCodeModeHost::new();
+    let (session_id, _events) = open_session(&host).await;
+    let (cell_id, mut execution) = execute_events(
+        &host,
+        execute_request(
+            &session_id,
+            "execution-termination",
+            "await new Promise(() => {});",
+        ),
+    )
+    .await;
+    execution.next().await.expect("execution outcome").unwrap();
+    let session = host.state.session(&session_id).expect("open session");
+    let termination = host.terminate(Request::new(proto::TerminateRequest {
+        session_id,
+        cell_id,
+    }));
+    tokio::pin!(termination);
+    assert!(termination.as_mut().now_or_never().is_none());
+
+    session.closed.cancel();
+    let error = termination
+        .await
+        .expect_err("closed sessions must bound termination");
+
+    assert_eq!(error.code(), Code::Cancelled);
+}
+
+#[tokio::test]
+async fn oversized_encoded_tool_invocation_fails_without_closing_its_session() {
+    let host = GrpcCodeModeHost::new();
+    let (session_id, _events) = open_session(&host).await;
+    let mut subscription = host
+        .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest {
+            session_id: session_id.clone(),
+            tool_names: Vec::new(),
+        }))
+        .await
+        .unwrap()
+        .into_inner();
+    let (cell_id, mut execution) = execute_events(
+        &host,
+        execute_request(
+            &session_id,
+            "execution-oversized",
+            "await new Promise(() => {});",
+        ),
+    )
+    .await;
+    execution.next().await.unwrap().unwrap();
+    let session = host.state.session(&session_id).unwrap();
+    let cancellation = CancellationToken::new();
+    let (response, _receiver) = oneshot::channel();
+
+    let error = session
+        .dispatch_tool(
+            invocation(&cell_id, "echo"),
+            "execution-oversized".to_string(),
+            Uuid::new_v4(),
+            Some(vec![0; MAX_FRAME_BYTES]),
+            response,
+            &cancellation,
+        )
+        .await
+        .unwrap_err();
+    assert!(error.contains("gRPC message limit"));
+    assert!(host.state.session(&session_id).is_ok());
+    assert!(!session.closed.is_cancelled());
+
+    let (response, _receiver) = oneshot::channel();
+    let invocation_id = Uuid::new_v4();
+    session
+        .dispatch_tool(
+            invocation(&cell_id, "echo"),
+            "execution-oversized".to_string(),
+            invocation_id,
+            Some(b"{}".to_vec()),
+            response,
+            &cancellation,
+        )
+        .await
+        .unwrap();
+    let delivered = subscription.next().await.unwrap().unwrap();
+    assert_eq!(delivered.invocation_id, invocation_id.to_string());
+    assert_eq!(delivered.sequence, 1);
+}
+
+#[tokio::test]
+async fn unmatched_tool_subscription_fails_without_closing_its_session() {
+    let host = GrpcCodeModeHost::new();
+    let (session_id, _events) = open_session(&host).await;
+    let _unrelated = host
+        .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest {
+            session_id: session_id.clone(),
+            tool_names: vec![proto::ToolName {
+                name: "other".to_string(),
+                namespace: None,
+            }],
+        }))
+        .await
+        .unwrap()
+        .into_inner();
+    let (cell_id, mut execution) = execute_events(
+        &host,
+        execute_request(
+            &session_id,
+            "execution-unmatched",
+            "await new Promise(() => {});",
+        ),
+    )
+    .await;
+    execution.next().await.unwrap().unwrap();
+    let session = host.state.session(&session_id).unwrap();
+    let cancellation = CancellationToken::new();
+    let (response, _receiver) = oneshot::channel();
+
+    assert!(
+        session
+            .dispatch_tool(
+                invocation(&cell_id, "echo"),
+                "execution-unmatched".to_string(),
+                Uuid::new_v4(),
+                Some(b"{}".to_vec()),
+                response,
+                &cancellation,
+            )
+            .await
+            .is_err()
+    );
+    assert!(host.state.session(&session_id).is_ok());
+    assert!(!session.closed.is_cancelled());
+
+    let mut matching = host
+        .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest {
+            session_id,
+            tool_names: vec![proto::ToolName {
+                name: "echo".to_string(),
+                namespace: None,
+            }],
+        }))
+        .await
+        .unwrap()
+        .into_inner();
+    let (response, _receiver) = oneshot::channel();
+    let invocation_id = Uuid::new_v4();
+    session
+        .dispatch_tool(
+            invocation(&cell_id, "echo"),
+            "execution-unmatched".to_string(),
+            invocation_id,
+            Some(b"{}".to_vec()),
+            response,
+            &cancellation,
+        )
+        .await
+        .unwrap();
+    let delivered = matching.next().await.unwrap().unwrap();
+    assert_eq!(delivered.invocation_id, invocation_id.to_string());
+    assert_eq!(delivered.sequence, 1);
+}
+
+#[tokio::test]
+async fn missing_selected_subscription_retries_another_matching_subscription() {
+    let host = GrpcCodeModeHost::new();
+    let (session_id, _events) = open_session(&host).await;
+    let mut first = host
+        .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest {
+            session_id: session_id.clone(),
+            tool_names: Vec::new(),
+        }))
+        .await
+        .expect("subscribe first tool stream")
+        .into_inner();
+    let mut second = host
+        .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest {
+            session_id: session_id.clone(),
+            tool_names: Vec::new(),
+        }))
+        .await
+        .expect("subscribe second tool stream")
+        .into_inner();
+    let (cell_id, mut execution) = execute_events(
+        &host,
+        execute_request(
+            &session_id,
+            "execution-retry",
+            "await new Promise(() => {});",
+        ),
+    )
+    .await;
+    execution.next().await.expect("execution outcome").unwrap();
+    let session = host.state.session(&session_id).expect("open session");
+    let subscriptions = session
+        .state
+        .lock()
+        .unwrap()
+        .subscriptions
+        .iter()
+        .map(|subscription| (subscription.id, subscription.sender.clone()))
+        .collect::<Vec<_>>();
+    for (_, sender) in &subscriptions {
+        for _ in 0..OUTGOING_CHANNEL_CAPACITY {
+            sender
+                .try_send(Ok(proto::ToolCall::default()))
+                .expect("fill subscription queue");
+        }
+    }
+    let cancellation = CancellationToken::new();
+    let (response, _receiver) = oneshot::channel();
+    let invocation_id = Uuid::new_v4();
+    let dispatch = session.dispatch_tool(
+        invocation(&cell_id, "echo"),
+        "execution-retry".to_string(),
+        invocation_id,
+        /*input_json*/ None,
+        response,
+        &cancellation,
+    );
+    tokio::pin!(dispatch);
+    assert!(dispatch.as_mut().now_or_never().is_none());
+    session
+        .state
+        .lock()
+        .unwrap()
+        .subscriptions
+        .retain(|subscription| subscription.id != subscriptions[0].0);
+
+    first.next().await.expect("free first reservation").unwrap();
+    assert!(dispatch.as_mut().now_or_never().is_none());
+    second
+        .next()
+        .await
+        .expect("free surviving reservation")
+        .unwrap();
+    dispatch.await.expect("retry surviving subscription");
+
+    for _ in 1..OUTGOING_CHANNEL_CAPACITY {
+        second.next().await.expect("drain buffered call").unwrap();
+    }
+    assert_eq!(
+        second
+            .next()
+            .await
+            .expect("retried invocation")
+            .unwrap()
+            .invocation_id,
+        invocation_id.to_string()
+    );
+}
+
+#[tokio::test]
+async fn saturated_subscription_does_not_block_independently_filtered_tools() {
+    let host = GrpcCodeModeHost::new();
+    let (session_id, _events) = open_session(&host).await;
+    let mut slow = host
+        .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest {
+            session_id: session_id.clone(),
+            tool_names: vec![proto::ToolName {
+                name: "slow".to_string(),
+                namespace: None,
+            }],
+        }))
+        .await
+        .unwrap()
+        .into_inner();
+    let mut fast = host
+        .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest {
+            session_id: session_id.clone(),
+            tool_names: vec![proto::ToolName {
+                name: "fast".to_string(),
+                namespace: None,
+            }],
+        }))
+        .await
+        .unwrap()
+        .into_inner();
+    let (cell_id, mut execution) = execute_events(
+        &host,
+        execute_request(
+            &session_id,
+            "execution-backpressure",
+            "await new Promise(() => {});",
+        ),
+    )
+    .await;
+    execution.next().await.unwrap().unwrap();
+    let session = host.state.session(&session_id).unwrap();
+    let cancellation = CancellationToken::new();
+    let mut responses = Vec::new();
+
+    for _ in 0..OUTGOING_CHANNEL_CAPACITY {
+        let (response, receiver) = oneshot::channel();
+        responses.push(receiver);
+        session
+            .dispatch_tool(
+                invocation(&cell_id, "slow"),
+                "execution-backpressure".to_string(),
+                Uuid::new_v4(),
+                /*input_json*/ None,
+                response,
+                &cancellation,
+            )
+            .await
+            .unwrap();
+    }
+
+    let blocked_session = Arc::clone(&session);
+    let blocked_cell = cell_id.clone();
+    let blocked_cancellation = cancellation.clone();
+    let (response, receiver) = oneshot::channel();
+    responses.push(receiver);
+    let blocked = tokio::spawn(async move {
+        blocked_session
+            .dispatch_tool(
+                invocation(&blocked_cell, "slow"),
+                "execution-backpressure".to_string(),
+                Uuid::new_v4(),
+                /*input_json*/ None,
+                response,
+                &blocked_cancellation,
+            )
+            .await
+    });
+    tokio::task::yield_now().await;
+    assert!(!blocked.is_finished());
+
+    let (response, receiver) = oneshot::channel();
+    responses.push(receiver);
+    let invocation_id = Uuid::new_v4();
+    tokio::time::timeout(
+        Duration::from_secs(1),
+        session.dispatch_tool(
+            invocation(&cell_id, "fast"),
+            "execution-backpressure".to_string(),
+            invocation_id,
+            /*input_json*/ None,
+            response,
+            &cancellation,
+        ),
+    )
+    .await
+    .expect("saturated subscription must not block another tool")
+    .unwrap();
+    assert_eq!(
+        fast.next().await.unwrap().unwrap().invocation_id,
+        invocation_id.to_string()
+    );
+
+    slow.next().await.unwrap().unwrap();
+    tokio::time::timeout(Duration::from_secs(1), blocked)
+        .await
+        .expect("draining the subscription should release its blocked invocation")
+        .unwrap()
+        .unwrap();
+}
+
+#[tokio::test]
+async fn dropping_subscriptions_only_retires_sessions_with_unread_calls() {
+    let host = GrpcCodeModeHost::new();
+    let (session_id, _events) = open_session(&host).await;
+    let idle = host
+        .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest {
+            session_id: session_id.clone(),
+            tool_names: Vec::new(),
+        }))
+        .await
+        .unwrap()
+        .into_inner();
+    drop(idle);
+
+    let session = host.state.session(&session_id).unwrap();
+    tokio::time::timeout(Duration::from_secs(/*secs*/ 2), async {
+        while !session.state.lock().unwrap().subscriptions.is_empty() {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .expect("idle subscription should be removed without retiring its session");
+    assert!(host.state.session(&session_id).is_ok());
+
+    let first = host
+        .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest {
+            session_id: session_id.clone(),
+            tool_names: Vec::new(),
+        }))
+        .await
+        .unwrap()
+        .into_inner();
+    let mut second = host
+        .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest {
+            session_id: session_id.clone(),
+            tool_names: Vec::new(),
+        }))
+        .await
+        .unwrap()
+        .into_inner();
+    let mut request = execute_request(
+        &session_id,
+        "execution-subscription-drop",
+        r#"await tools.echo({attempt: 1});"#,
+    );
+    request.yield_time_ms = Some(/*value*/ 10_000);
+    request.enabled_tools = vec![tool("echo")];
+    let (_cell_id, mut execution) = execute_events(&host, request).await;
+    let first_subscription_id = session.state.lock().unwrap().subscriptions[0].id;
+    tokio::time::timeout(Duration::from_secs(/*secs*/ 2), async {
+        loop {
+            let owned = session
+                .state
+                .lock()
+                .unwrap()
+                .pending_invocations
+                .values()
+                .any(|invocation| invocation.subscription_id == first_subscription_id);
+            if owned {
+                break;
+            }
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .expect("first subscription should own an unread tool call");
+    drop(first);
+
+    tokio::time::timeout(Duration::from_secs(/*secs*/ 2), async {
+        while host.state.session(&session_id).is_ok() {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .expect("losing an unread call must retire its lease without leaving a sequence gap");
+    assert!(
+        tokio::time::timeout(Duration::from_secs(/*secs*/ 2), second.next())
+            .await
+            .expect("other subscriptions should close with their lease")
+            .is_none()
+    );
+    assert!(
+        tokio::time::timeout(Duration::from_secs(/*secs*/ 2), execution.next())
+            .await
+            .expect("execution should retire with its lease")
+            .is_none()
+    );
+}

codex-rs/code-mode-host/src/grpc/routing.rs

@@ -0,0 +1,307 @@
+use std::sync::Arc;
+use std::sync::PoisonError;
+
+use codex_code_mode_protocol::CodeModeNestedToolCall;
+use codex_code_mode_protocol::grpc as proto;
+use codex_code_mode_protocol::host::MAX_FRAME_BYTES;
+use futures::StreamExt;
+use futures::stream::FuturesUnordered;
+use prost::Message;
+use serde_json::Value as JsonValue;
+use tokio::sync::mpsc;
+use tokio::sync::oneshot;
+use tokio_stream::wrappers::ReceiverStream;
+use tokio_util::sync::CancellationToken;
+use tonic::Status;
+use uuid::Uuid;
+
+use super::GrpcStream;
+use super::conversions;
+use super::session::GrpcSession;
+use super::session::PendingInvocation;
+use super::session::ToolSubscription;
+use super::validation;
+use crate::OUTGOING_CHANNEL_CAPACITY;
+
+const MAX_SUBSCRIPTIONS: usize = OUTGOING_CHANNEL_CAPACITY;
+
+impl GrpcSession {
+    pub(super) fn subscribe(
+        self: &Arc<Self>,
+        filters: Vec<proto::ToolName>,
+    ) -> Result<GrpcStream<proto::ToolCall>, Status> {
+        validation::tool_filters(&filters)?;
+        let id = Uuid::new_v4();
+        let (sender, receiver) = mpsc::channel(OUTGOING_CHANNEL_CAPACITY);
+        let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+        if self.closed.is_cancelled() {
+            return Err(Status::cancelled("code-mode session is closed"));
+        }
+        if state.subscriptions.len() >= MAX_SUBSCRIPTIONS {
+            return Err(Status::resource_exhausted(
+                "code-mode session has too many tool subscriptions",
+            ));
+        }
+        state.subscriptions.push(ToolSubscription {
+            id,
+            filters,
+            sender: sender.clone(),
+        });
+        drop(state);
+
+        let session = Arc::downgrade(self);
+        let closed = self.closed.clone();
+        tokio::spawn(async move {
+            tokio::select! {
+                _ = sender.closed() => {}
+                _ = closed.cancelled() => return,
+            }
+            if let Some(session) = session.upgrade() {
+                let abandoned = {
+                    let mut state = session.state.lock().unwrap_or_else(PoisonError::into_inner);
+                    state
+                        .subscriptions
+                        .retain(|subscription| subscription.id != id);
+                    state
+                        .pending_invocations
+                        .extract_if(|_, invocation| invocation.subscription_id == id)
+                        .map(|(_, invocation)| invocation)
+                        .collect::<Vec<_>>()
+                };
+                if !abandoned.is_empty() {
+                    // A buffered call may never have reached the client. Closing
+                    // the lease avoids replaying a possibly delivered invocation
+                    // or leaving an unfillable execution-sequence gap.
+                    session.closed.cancel();
+                }
+                for invocation in abandoned {
+                    let _ = invocation.response.send(Err(
+                        "code-mode tool subscription closed before returning tool output"
+                            .to_string(),
+                    ));
+                }
+            }
+        });
+        Ok(Box::pin(ReceiverStream::new(receiver)))
+    }
+
+    pub(super) async fn dispatch_tool(
+        &self,
+        invocation: CodeModeNestedToolCall,
+        execution_id: String,
+        invocation_id: Uuid,
+        input_json: Option<Vec<u8>>,
+        response: oneshot::Sender<Result<JsonValue, String>>,
+        cancellation: &CancellationToken,
+    ) -> Result<(), String> {
+        let cell_id = invocation.cell_id.to_string();
+        let tool_name = proto::ToolName {
+            name: invocation.tool_name.name,
+            namespace: invocation.tool_name.namespace,
+        };
+        let (sequence, subscriptions) = {
+            let state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+            let Some(execution) = state.cells.get(&cell_id) else {
+                return Err("code-mode cell closed before dispatching its tool call".to_string());
+            };
+            let sequence = execution.tool_call_sequence.checked_add(1).ok_or_else(|| {
+                "code-mode execution tool-call sequence was exhausted".to_string()
+            })?;
+            let count = state.subscriptions.len();
+            let subscriptions = (0..count)
+                .map(|offset| &state.subscriptions[(state.next_subscription + offset) % count])
+                .filter(|subscription| {
+                    subscription.filters.is_empty()
+                        || subscription.filters.iter().any(|filter| {
+                            filter.name == tool_name.name && filter.namespace == tool_name.namespace
+                        })
+                })
+                .map(|subscription| (subscription.id, subscription.sender.clone()))
+                .collect::<Vec<_>>();
+            (sequence, subscriptions)
+        };
+        if subscriptions.is_empty() {
+            return Err("no code-mode tool subscription matches the requested tool".to_string());
+        }
+        let mut message = proto::ToolCall {
+            session_id: self.id.to_string(),
+            execution_id,
+            cell_id: cell_id.clone(),
+            invocation_id: invocation_id.to_string(),
+            runtime_tool_call_id: invocation.runtime_tool_call_id,
+            tool_name: Some(tool_name.clone()),
+            tool_kind: conversions::tool_kind(invocation.tool_kind),
+            input_json,
+            sequence,
+        };
+        if message.encoded_len() > MAX_FRAME_BYTES {
+            return Err(format!(
+                "code-mode tool invocation exceeds the {MAX_FRAME_BYTES}-byte gRPC message limit"
+            ));
+        }
+
+        let mut reservations =
+            subscriptions
+                .into_iter()
+                .map(|(id, sender)| async move {
+                    sender.reserve_owned().await.map(|permit| (id, permit))
+                })
+                .collect::<FuturesUnordered<_>>();
+        loop {
+            let (subscription_id, permit) = tokio::select! {
+                biased;
+                _ = cancellation.cancelled() => {
+                    return Err("code mode delegate request cancelled".to_string());
+                }
+                _ = self.closed.cancelled() => {
+                    return Err("code-mode session closed before dispatching its tool call".to_string());
+                }
+                reservation = reservations.next() => match reservation {
+                    Some(Ok(reservation)) => reservation,
+                    Some(Err(_)) => continue,
+                    None => {
+                        return Err("matching code-mode tool subscriptions are unavailable".to_string());
+                    }
+                },
+            };
+
+            let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+            let Some(subscription_index) = state.subscriptions.iter().position(|subscription| {
+                subscription.id == subscription_id && !subscription.sender.is_closed()
+            }) else {
+                continue;
+            };
+            let Some(execution) = state.cells.get_mut(&cell_id) else {
+                return Err("code-mode cell closed before dispatching its tool call".to_string());
+            };
+            let sequence = execution.tool_call_sequence.checked_add(1).ok_or_else(|| {
+                "code-mode execution tool-call sequence was exhausted".to_string()
+            })?;
+            message.sequence = sequence;
+            if message.encoded_len() > MAX_FRAME_BYTES {
+                return Err(format!(
+                    "code-mode tool invocation exceeds the {MAX_FRAME_BYTES}-byte gRPC message limit"
+                ));
+            }
+            execution.tool_call_sequence = sequence;
+            state.pending_invocations.insert(
+                invocation_id,
+                PendingInvocation {
+                    subscription_id,
+                    response,
+                },
+            );
+            state.seen_invocations.remember(invocation_id);
+            state.next_subscription = (subscription_index + 1) % state.subscriptions.len();
+            permit.send(Ok(message));
+            return Ok(());
+        }
+    }
+
+    pub(super) fn complete_invocation(
+        &self,
+        invocation_id: Uuid,
+        result: Result<JsonValue, String>,
+    ) -> Result<(), Status> {
+        let response = {
+            let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+            match state.pending_invocations.remove(&invocation_id) {
+                Some(invocation) => Some(invocation.response),
+                None if state.seen_invocations.contains(&invocation_id) => None,
+                None => {
+                    return Err(Status::not_found(format!(
+                        "unknown code-mode tool invocation {invocation_id}"
+                    )));
+                }
+            }
+        };
+        if let Some(response) = response {
+            let _ = response.send(result);
+        }
+        Ok(())
+    }
+
+    pub(super) fn cancel_invocation(&self, invocation_id: Uuid) {
+        let pending = self
+            .state
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .pending_invocations
+            .remove(&invocation_id);
+        if pending.is_some() {
+            let _ = self.send_event_now(
+                proto::session_event::Event::ToolCallCancelled(proto::ToolCallCancelled {
+                    invocation_id: invocation_id.to_string(),
+                }),
+                /*cell_permit*/ None,
+            );
+        }
+    }
+
+    pub(super) async fn begin_notification(
+        &self,
+        notification_id: Uuid,
+        notification: proto::Notification,
+        response: oneshot::Sender<()>,
+        cancellation: &CancellationToken,
+    ) -> Result<(), String> {
+        {
+            let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+            state
+                .pending_notifications
+                .insert(notification_id, response);
+            state.seen_notifications.remember(notification_id);
+        }
+        if let Err(error) = self
+            .send_event(
+                proto::session_event::Event::Notification(notification),
+                cancellation,
+            )
+            .await
+        {
+            self.state
+                .lock()
+                .unwrap_or_else(PoisonError::into_inner)
+                .pending_notifications
+                .remove(&notification_id);
+            return Err(error);
+        }
+        Ok(())
+    }
+
+    pub(super) fn acknowledge_notification(&self, notification_id: Uuid) -> Result<(), Status> {
+        let response = {
+            let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+            match state.pending_notifications.remove(&notification_id) {
+                Some(response) => Some(response),
+                None if state.seen_notifications.contains(&notification_id) => None,
+                None => {
+                    return Err(Status::not_found(format!(
+                        "unknown code-mode notification {notification_id}"
+                    )));
+                }
+            }
+        };
+        if let Some(response) = response {
+            let _ = response.send(());
+        }
+        Ok(())
+    }
+
+    pub(super) fn cancel_notification(&self, notification_id: Uuid) {
+        let pending = self
+            .state
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .pending_notifications
+            .remove(&notification_id);
+        if pending.is_some() {
+            let _ = self.send_event_now(
+                proto::session_event::Event::NotificationCancelled(proto::NotificationCancelled {
+                    notification_id: notification_id.to_string(),
+                }),
+                /*cell_permit*/ None,
+            );
+        }
+    }
+}

codex-rs/code-mode-host/src/grpc/service_tests.rs

@@ -0,0 +1,348 @@
+use super::GrpcCodeModeHost;
+use super::GrpcStream;
+use codex_code_mode_protocol::grpc as proto;
+use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHost;
+use futures::FutureExt;
+use futures::StreamExt;
+use pretty_assertions::assert_eq;
+use tonic::Code;
+use tonic::Request;
+
+pub(super) async fn open_session(
+    host: &GrpcCodeModeHost,
+) -> (String, GrpcStream<proto::SessionEvent>) {
+    let mut stream = host
+        .open_session(Request::new(proto::OpenSessionRequest {
+            cell_execution_limits: None,
+        }))
+        .await
+        .expect("open code-mode session")
+        .into_inner();
+    let event = stream
+        .next()
+        .await
+        .expect("session opened event")
+        .expect("session event");
+    let Some(proto::session_event::Event::Opened(opened)) = event.event else {
+        panic!("expected the first session event to open its lease");
+    };
+    (opened.session_id, stream)
+}
+
+pub(super) fn execute_request(
+    session_id: &str,
+    execution_id: &str,
+    source: &str,
+) -> proto::ExecuteRequest {
+    proto::ExecuteRequest {
+        session_id: session_id.to_string(),
+        execution_id: execution_id.to_string(),
+        tool_call_id: "outer-call".to_string(),
+        source: source.to_string(),
+        enabled_tools: Vec::new(),
+        yield_time_ms: Some(1),
+        max_output_tokens: None,
+    }
+}
+
+pub(super) fn tool(name: &str) -> proto::ToolDefinition {
+    proto::ToolDefinition {
+        name: name.to_string(),
+        tool_name: Some(proto::ToolName {
+            name: name.to_string(),
+            namespace: None,
+        }),
+        description: String::new(),
+        kind: proto::ToolKind::Function as i32,
+        input_schema_json: None,
+        output_schema_json: None,
+    }
+}
+
+pub(super) async fn execute_events(
+    host: &GrpcCodeModeHost,
+    request: proto::ExecuteRequest,
+) -> (String, GrpcStream<proto::ExecuteEvent>) {
+    let mut stream = host
+        .execute(Request::new(request))
+        .await
+        .expect("execute cell")
+        .into_inner();
+    let event = stream
+        .next()
+        .await
+        .expect("execution started event")
+        .expect("execution event");
+    let Some(proto::execute_event::Event::Started(started)) = event.event else {
+        panic!("expected execution admission before its outcome");
+    };
+    (started.cell_id, stream)
+}
+
+#[tokio::test]
+async fn execute_stream_starts_immediately_and_wait_preserves_missing_cells() {
+    let host = GrpcCodeModeHost::new();
+    let (session_id, mut session_events) = open_session(&host).await;
+    let mut request = execute_request(
+        &session_id,
+        "execution-1",
+        r#"text("before"); yield_control(); text("after");"#,
+    );
+    request.yield_time_ms = Some(60_000);
+    let (cell_id, mut execution) = execute_events(&host, request).await;
+
+    assert_eq!(
+        execution.next().await.unwrap().unwrap(),
+        proto::ExecuteEvent {
+            event: Some(proto::execute_event::Event::Outcome(
+                proto::ExecutionOutcome {
+                    cell_id: cell_id.clone(),
+                    content_items: vec![proto::ContentItem {
+                        item: Some(proto::content_item::Item::Text(proto::TextContent {
+                            text: "before".to_string(),
+                        })),
+                    }],
+                    outcome: Some(proto::execution_outcome::Outcome::Yielded(
+                        proto::ExecutionYielded {},
+                    )),
+                }
+            )),
+        }
+    );
+    let completed = host
+        .wait(Request::new(proto::WaitRequest {
+            session_id: session_id.clone(),
+            cell_id: cell_id.clone(),
+            wait_id: "wait-1".to_string(),
+            yield_time_ms: 60_000,
+        }))
+        .await
+        .expect("wait for completion")
+        .into_inner();
+    assert_eq!(
+        completed,
+        proto::WaitResponse {
+            state: Some(proto::wait_response::State::LiveCell(
+                proto::ExecutionOutcome {
+                    cell_id: cell_id.clone(),
+                    content_items: vec![proto::ContentItem {
+                        item: Some(proto::content_item::Item::Text(proto::TextContent {
+                            text: "after".to_string(),
+                        })),
+                    }],
+                    outcome: Some(proto::execution_outcome::Outcome::Completed(
+                        proto::ExecutionCompleted { error_text: None },
+                    )),
+                }
+            )),
+        }
+    );
+    assert_eq!(
+        session_events.next().await.unwrap().unwrap(),
+        proto::SessionEvent {
+            event: Some(proto::session_event::Event::CellClosed(proto::CellClosed {
+                execution_id: "execution-1".to_string(),
+                cell_id: cell_id.clone(),
+                final_tool_call_sequence: 0,
+            })),
+        }
+    );
+
+    let missing = host
+        .wait(Request::new(proto::WaitRequest {
+            session_id,
+            cell_id: "missing-cell".to_string(),
+            wait_id: "wait-missing".to_string(),
+            yield_time_ms: 1,
+        }))
+        .await
+        .expect("missing cells remain successful wait outcomes")
+        .into_inner();
+    assert!(matches!(
+        missing.state,
+        Some(proto::wait_response::State::MissingCell(_))
+    ));
+}
+
+#[tokio::test]
+async fn filtered_subscriptions_receive_ordered_calls_and_unary_completions() {
+    let host = GrpcCodeModeHost::new();
+    let (session_id, mut session_events) = open_session(&host).await;
+    let mut matching = host
+        .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest {
+            session_id: session_id.clone(),
+            tool_names: vec![proto::ToolName {
+                name: "echo".to_string(),
+                namespace: None,
+            }],
+        }))
+        .await
+        .unwrap()
+        .into_inner();
+    let mut unrelated = host
+        .subscribe_to_tool_calls(Request::new(proto::SubscribeToToolCallsRequest {
+            session_id: session_id.clone(),
+            tool_names: vec![proto::ToolName {
+                name: "other".to_string(),
+                namespace: None,
+            }],
+        }))
+        .await
+        .unwrap()
+        .into_inner();
+    let mut request = execute_request(
+        &session_id,
+        "execution-tools",
+        r#"text(await tools.echo({value: 1})); text(await tools.echo({value: 2}));"#,
+    );
+    request.yield_time_ms = Some(60_000);
+    request.enabled_tools = vec![tool("echo")];
+    let (cell_id, mut execution) = execute_events(&host, request).await;
+
+    for sequence in [1, 2] {
+        let invocation = matching.next().await.unwrap().unwrap();
+        assert_eq!(invocation.session_id, session_id);
+        assert_eq!(invocation.execution_id, "execution-tools");
+        assert_eq!(invocation.cell_id, cell_id);
+        assert_eq!(invocation.sequence, sequence);
+        assert_eq!(
+            invocation.input_json,
+            Some(format!(r#"{{"value":{sequence}}}"#).into_bytes())
+        );
+        assert!(unrelated.next().now_or_never().is_none());
+        host.complete_tool_call(Request::new(proto::CompleteToolCallRequest {
+            session_id: session_id.clone(),
+            invocation_id: invocation.invocation_id,
+            outcome: Some(proto::complete_tool_call_request::Outcome::Succeeded(
+                proto::ToolCallSucceeded {
+                    output_json: format!(r#""result-{sequence}""#).into_bytes(),
+                },
+            )),
+        }))
+        .await
+        .expect("complete delegated tool");
+    }
+
+    let outcome = execution.next().await.unwrap().unwrap();
+    assert!(matches!(
+        outcome.event,
+        Some(proto::execute_event::Event::Outcome(
+            proto::ExecutionOutcome {
+                outcome: Some(proto::execution_outcome::Outcome::Completed(_)),
+                ..
+            }
+        ))
+    ));
+    assert_eq!(
+        session_events.next().await.unwrap().unwrap(),
+        proto::SessionEvent {
+            event: Some(proto::session_event::Event::CellClosed(proto::CellClosed {
+                execution_id: "execution-tools".to_string(),
+                cell_id,
+                final_tool_call_sequence: 2,
+            })),
+        }
+    );
+}
+
+#[tokio::test]
+async fn cancellation_before_wait_admission_is_preserved() {
+    let host = GrpcCodeModeHost::new();
+    let (session_id, _events) = open_session(&host).await;
+    let (cell_id, mut execution) = execute_events(
+        &host,
+        execute_request(
+            &session_id,
+            "execution-wait",
+            "await new Promise(() => {});",
+        ),
+    )
+    .await;
+    execution.next().await.unwrap().unwrap();
+
+    host.cancel_wait(Request::new(proto::CancelWaitRequest {
+        session_id: session_id.clone(),
+        wait_id: "pre-cancelled".to_string(),
+    }))
+    .await
+    .expect("record cancellation before wait admission");
+    let cancelled = host
+        .wait(Request::new(proto::WaitRequest {
+            session_id: session_id.clone(),
+            cell_id: cell_id.clone(),
+            wait_id: "pre-cancelled".to_string(),
+            yield_time_ms: 60_000,
+        }))
+        .await
+        .unwrap_err();
+    assert_eq!(cancelled.code(), Code::Cancelled);
+
+    host.terminate(Request::new(proto::TerminateRequest {
+        session_id,
+        cell_id,
+    }))
+    .await
+    .unwrap();
+}
+
+#[tokio::test]
+async fn terminating_a_cell_cancels_unacknowledged_notifications() {
+    let host = GrpcCodeModeHost::new();
+    let (session_id, mut session_events) = open_session(&host).await;
+    let (cell_id, mut execution) = execute_events(
+        &host,
+        execute_request(
+            &session_id,
+            "execution-notify",
+            r#"notify("pending"); await new Promise(() => {});"#,
+        ),
+    )
+    .await;
+    let event = session_events.next().await.unwrap().unwrap();
+    let Some(proto::session_event::Event::Notification(notification)) = event.event else {
+        panic!("expected pending notification");
+    };
+    assert_eq!(notification.execution_id, "execution-notify");
+    assert_eq!(notification.cell_id, cell_id);
+    assert_eq!(notification.call_id, "outer-call");
+    assert_eq!(notification.text, "pending");
+    execution.next().await.unwrap().unwrap();
+
+    let terminated = host
+        .terminate(Request::new(proto::TerminateRequest {
+            session_id,
+            cell_id: cell_id.clone(),
+        }))
+        .await
+        .unwrap()
+        .into_inner();
+    assert!(matches!(
+        terminated.state,
+        Some(proto::wait_response::State::LiveCell(
+            proto::ExecutionOutcome {
+                outcome: Some(proto::execution_outcome::Outcome::Terminated(_)),
+                ..
+            }
+        ))
+    ));
+    assert_eq!(
+        session_events.next().await.unwrap().unwrap(),
+        proto::SessionEvent {
+            event: Some(proto::session_event::Event::NotificationCancelled(
+                proto::NotificationCancelled {
+                    notification_id: notification.notification_id,
+                },
+            )),
+        }
+    );
+    assert_eq!(
+        session_events.next().await.unwrap().unwrap(),
+        proto::SessionEvent {
+            event: Some(proto::session_event::Event::CellClosed(proto::CellClosed {
+                execution_id: "execution-notify".to_string(),
+                cell_id,
+                final_tool_call_sequence: 0,
+            })),
+        }
+    );
+}

codex-rs/code-mode-host/src/grpc/session.rs

@@ -0,0 +1,465 @@
+use std::borrow::Borrow;
+use std::collections::HashMap;
+use std::collections::HashSet;
+use std::collections::VecDeque;
+use std::collections::hash_map::Entry;
+use std::hash::Hash;
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::PoisonError;
+use std::sync::Weak;
+
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
+use codex_code_mode_protocol::WaitOutcome;
+use codex_code_mode_protocol::grpc as proto;
+use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
+use codex_code_mode_runtime::InProcessCodeModeSession;
+use serde_json::Value as JsonValue;
+use tokio::sync::Notify;
+use tokio::sync::OwnedSemaphorePermit;
+use tokio::sync::Semaphore;
+use tokio::sync::mpsc;
+use tokio::sync::oneshot;
+use tokio_stream::wrappers::ReceiverStream;
+use tokio_util::sync::CancellationToken;
+use tonic::Status;
+use uuid::Uuid;
+
+use super::GrpcStream;
+use super::delegate::GrpcDelegate;
+use super::events::EventSender;
+use super::validation;
+use super::waits::ActiveWait;
+use crate::HostLimits;
+use crate::MAX_ACTIVE_CELLS;
+use crate::MAX_IN_FLIGHT_REQUESTS;
+use crate::MAX_RECENT_REQUEST_IDS;
+use crate::OUTGOING_CHANNEL_CAPACITY;
+
+pub(super) struct GrpcHostState {
+    sessions: Mutex<HashMap<Uuid, Arc<GrpcSession>>>,
+    limits: HostLimits,
+    delegate_permits: Arc<Semaphore>,
+    control_permits: Arc<Semaphore>,
+}
+
+pub(super) struct GrpcSession {
+    pub(super) id: Uuid,
+    pub(super) runtime: Arc<InProcessCodeModeSession>,
+    pub(super) closed: CancellationToken,
+    pub(super) state: Mutex<SessionState>,
+    events: EventSender,
+    cells_changed: Notify,
+    delegate_permits: Arc<Semaphore>,
+}
+
+#[derive(Default)]
+pub(super) struct SessionState {
+    pub(super) cells: HashMap<String, ExecutionState>,
+    pending_executions: HashSet<String>,
+    pending_closures: HashSet<String>,
+    seen_executions: BoundedIds,
+    pub(super) subscriptions: Vec<ToolSubscription>,
+    pub(super) next_subscription: usize,
+    pub(super) pending_invocations: HashMap<Uuid, PendingInvocation>,
+    pub(super) seen_invocations: BoundedIds<Uuid>,
+    pub(super) pending_notifications: HashMap<Uuid, oneshot::Sender<()>>,
+    pub(super) seen_notifications: BoundedIds<Uuid>,
+    pub(super) waits: HashMap<String, ActiveWait>,
+    pub(super) seen_waits: BoundedIds,
+    pub(super) cancelled_waits: BoundedIds,
+}
+
+pub(super) struct ExecutionState {
+    pub(super) execution_id: String,
+    pub(super) tool_call_sequence: u64,
+    permit: OwnedSemaphorePermit,
+}
+
+pub(super) struct ToolSubscription {
+    pub(super) id: Uuid,
+    pub(super) filters: Vec<proto::ToolName>,
+    pub(super) sender: mpsc::Sender<Result<proto::ToolCall, Status>>,
+}
+
+pub(super) struct PendingInvocation {
+    pub(super) subscription_id: Uuid,
+    pub(super) response: oneshot::Sender<Result<JsonValue, String>>,
+}
+
+#[derive(Default)]
+pub(super) struct BoundedIds<T = String> {
+    ids: HashSet<T>,
+    order: VecDeque<T>,
+}
+
+impl GrpcHostState {
+    pub(super) fn new() -> Self {
+        Self {
+            sessions: Mutex::new(HashMap::new()),
+            limits: HostLimits::new(),
+            delegate_permits: Arc::new(Semaphore::new(MAX_PENDING_DELEGATE_CALLS)),
+            control_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)),
+        }
+    }
+
+    pub(super) fn open_session(
+        self: &Arc<Self>,
+        limits: CodeModeSessionCellExecutionLimits,
+    ) -> Result<GrpcStream<proto::SessionEvent>, Status> {
+        let id = Uuid::new_v4();
+        let (events, receiver) = mpsc::channel(OUTGOING_CHANNEL_CAPACITY);
+        let closed = CancellationToken::new();
+        let event_sender = EventSender::new(events.clone(), closed.clone());
+        let session = GrpcSession::new(
+            id,
+            event_sender,
+            closed,
+            Arc::clone(&self.delegate_permits),
+            limits,
+        );
+        events
+            .try_send(Ok(proto::SessionEvent {
+                event: Some(proto::session_event::Event::Opened(proto::SessionOpened {
+                    session_id: id.to_string(),
+                })),
+            }))
+            .map_err(|_| Status::internal("failed to publish the opened code-mode session"))?;
+        let mut sessions = self.sessions.lock().unwrap_or_else(PoisonError::into_inner);
+        if sessions.len() >= MAX_IN_FLIGHT_REQUESTS {
+            return Err(Status::resource_exhausted(
+                "code-mode host has too many open sessions",
+            ));
+        }
+        sessions.insert(id, Arc::clone(&session));
+        drop(sessions);
+
+        let host = Arc::downgrade(self);
+        tokio::spawn(async move {
+            tokio::select! {
+                _ = events.closed() => {}
+                _ = session.closed.cancelled() => {}
+            }
+            if let Some(host) = host.upgrade() {
+                host.close_lease(id, &session).await;
+            } else {
+                let _ = session.shutdown().await;
+            }
+        });
+
+        Ok(Box::pin(ReceiverStream::new(receiver)))
+    }
+
+    pub(super) fn session(&self, id: &str) -> Result<Arc<GrpcSession>, Status> {
+        let session_id = validation::uuid(id, "session ID")?;
+        self.sessions
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .get(&session_id)
+            .cloned()
+            .ok_or_else(|| Status::not_found(format!("unknown code-mode session {id}")))
+    }
+
+    pub(super) async fn close_session(&self, id: &str) -> Result<(), Status> {
+        let session_id = validation::uuid(id, "session ID")?;
+        let session = self
+            .sessions
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .remove(&session_id)
+            .ok_or_else(|| Status::not_found(format!("unknown code-mode session {id}")))?;
+        session.shutdown().await
+    }
+
+    async fn close_lease(&self, id: Uuid, expected: &Arc<GrpcSession>) {
+        let session = {
+            let mut sessions = self.sessions.lock().unwrap_or_else(PoisonError::into_inner);
+            if sessions
+                .get(&id)
+                .is_some_and(|session| Arc::ptr_eq(session, expected))
+            {
+                sessions.remove(&id)
+            } else {
+                None
+            }
+        };
+        if let Some(session) = session {
+            let _ = session.shutdown().await;
+        }
+    }
+
+    pub(super) fn request_permit(&self) -> Result<OwnedSemaphorePermit, Status> {
+        self.limits.request_permit().map_err(|_| {
+            Status::resource_exhausted("code-mode host has too many in-flight requests")
+        })
+    }
+
+    pub(super) fn cell_permit(&self) -> Result<OwnedSemaphorePermit, Status> {
+        self.limits
+            .cell_permit()
+            .map_err(|_| Status::resource_exhausted("code-mode host has too many active cells"))
+    }
+
+    pub(super) fn control_permit(&self) -> Result<OwnedSemaphorePermit, Status> {
+        Arc::clone(&self.control_permits)
+            .try_acquire_owned()
+            .map_err(|_| {
+                Status::resource_exhausted("code-mode host has too many in-flight control requests")
+            })
+    }
+}
+
+impl GrpcSession {
+    fn new(
+        id: Uuid,
+        events: EventSender,
+        closed: CancellationToken,
+        delegate_permits: Arc<Semaphore>,
+        limits: CodeModeSessionCellExecutionLimits,
+    ) -> Arc<Self> {
+        Arc::new_cyclic(|weak: &Weak<Self>| {
+            let delegate = Arc::new(GrpcDelegate::new(weak.clone()));
+            let failure_session = weak.clone();
+            let failure_handler = Arc::new(move |reason: String| {
+                if let Some(session) = failure_session.upgrade() {
+                    tracing::warn!(session_id = %session.id, "code-mode host session failed: {reason}");
+                    session.closed.cancel();
+                }
+            });
+            Self {
+                id,
+                runtime: Arc::new(
+                    InProcessCodeModeSession::with_delegate_and_task_failure_handler(
+                        delegate,
+                        failure_handler,
+                        limits,
+                    ),
+                ),
+                closed,
+                state: Mutex::new(SessionState::default()),
+                events,
+                cells_changed: Notify::new(),
+                delegate_permits,
+            }
+        })
+    }
+
+    async fn shutdown(&self) -> Result<(), Status> {
+        self.closed.cancel();
+        {
+            let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+            for wait in state.waits.values() {
+                wait.cancellation.cancel();
+            }
+            state.pending_invocations.clear();
+            state.pending_notifications.clear();
+            state.subscriptions.clear();
+        }
+        let result = self.runtime.shutdown().await.map_err(Status::internal);
+        self.events.shutdown().await;
+        self.state
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .cells
+            .clear();
+        result
+    }
+
+    pub(super) async fn terminate(&self, cell_id: CellId) -> Result<WaitOutcome, Status> {
+        tokio::select! {
+            biased;
+            _ = self.closed.cancelled() => {
+                Err(Status::cancelled("code-mode session is closed"))
+            }
+            result = self.runtime.terminate(cell_id) => {
+                result.map_err(Status::failed_precondition)
+            }
+        }
+    }
+
+    pub(super) fn reserve_execution(&self, execution_id: &str) -> Result<(), Status> {
+        validation::identifier(execution_id, "execution ID")?;
+        let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+        if self.closed.is_cancelled() {
+            return Err(Status::cancelled("code-mode session is closed"));
+        }
+        if state.pending_executions.contains(execution_id)
+            || state
+                .cells
+                .values()
+                .any(|execution| execution.execution_id == execution_id)
+            || !state.seen_executions.remember(execution_id.to_string())
+        {
+            return Err(Status::already_exists(format!(
+                "code-mode execution ID `{execution_id}` was reused"
+            )));
+        }
+        state.pending_executions.insert(execution_id.to_string());
+        Ok(())
+    }
+
+    pub(super) fn admit_execution(
+        &self,
+        execution_id: String,
+        cell_id: String,
+        permit: OwnedSemaphorePermit,
+    ) -> Result<(), Status> {
+        let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+        if !state.pending_executions.remove(&execution_id) {
+            return Err(Status::cancelled("code-mode execution was abandoned"));
+        }
+        let Entry::Vacant(entry) = state.cells.entry(cell_id.clone()) else {
+            return Err(Status::internal(
+                "code-mode runtime reused an active cell ID",
+            ));
+        };
+        entry.insert(ExecutionState {
+            execution_id,
+            tool_call_sequence: 0,
+            permit,
+        });
+        let closed = state.pending_closures.remove(&cell_id);
+        let closed_execution = closed.then(|| state.cells.remove(&cell_id)).flatten();
+        drop(state);
+        self.cells_changed.notify_waiters();
+        if let Some(execution) = closed_execution {
+            self.send_cell_closed(&cell_id, execution);
+        }
+        Ok(())
+    }
+
+    pub(super) fn abandon_execution(self: &Arc<Self>, execution_id: &str) {
+        let cell_id = {
+            let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+            state.pending_executions.remove(execution_id);
+            state
+                .cells
+                .iter()
+                .find(|(_, execution)| execution.execution_id == execution_id)
+                .map(|(cell_id, _)| cell_id.clone())
+        };
+        if let Some(cell_id) = cell_id {
+            let session = Arc::clone(self);
+            tokio::spawn(async move {
+                let _ = session.terminate(CellId::new(cell_id)).await;
+            });
+        }
+    }
+
+    pub(super) async fn execution_id(
+        &self,
+        cell_id: &str,
+        cancellation: &CancellationToken,
+    ) -> Result<String, String> {
+        loop {
+            let changed = self.cells_changed.notified();
+            if let Some(execution_id) = self
+                .state
+                .lock()
+                .unwrap_or_else(PoisonError::into_inner)
+                .cells
+                .get(cell_id)
+                .map(|execution| execution.execution_id.clone())
+            {
+                return Ok(execution_id);
+            }
+            tokio::select! {
+                _ = self.closed.cancelled() => {
+                    return Err("code-mode session closed before cell admission".to_string());
+                }
+                _ = cancellation.cancelled() => {
+                    return Err("code-mode callback was cancelled before cell admission".to_string());
+                }
+                _ = changed => {}
+            }
+        }
+    }
+
+    pub(super) fn close_cell(&self, cell_id: &str) {
+        let execution = {
+            let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+            match state.cells.remove(cell_id) {
+                Some(execution) => Some(execution),
+                None if state.pending_closures.len() < MAX_ACTIVE_CELLS => {
+                    state.pending_closures.insert(cell_id.to_string());
+                    None
+                }
+                None => {
+                    self.closed.cancel();
+                    None
+                }
+            }
+        };
+        if let Some(execution) = execution {
+            self.send_cell_closed(cell_id, execution);
+        }
+    }
+
+    fn send_cell_closed(&self, cell_id: &str, execution: ExecutionState) {
+        let _ = self.send_event_now(
+            proto::session_event::Event::CellClosed(proto::CellClosed {
+                execution_id: execution.execution_id,
+                cell_id: cell_id.to_string(),
+                final_tool_call_sequence: execution.tool_call_sequence,
+            }),
+            Some(execution.permit),
+        );
+    }
+
+    pub(super) fn delegate_permit(&self) -> Result<OwnedSemaphorePermit, String> {
+        Arc::clone(&self.delegate_permits)
+            .try_acquire_owned()
+            .map_err(|_| "code-mode host has too many pending delegate calls".to_string())
+    }
+
+    pub(super) async fn send_event(
+        &self,
+        event: proto::session_event::Event,
+        cancellation: &CancellationToken,
+    ) -> Result<(), String> {
+        self.events.send(event, cancellation).await
+    }
+
+    pub(super) fn send_event_now(
+        &self,
+        event: proto::session_event::Event,
+        cell_permit: Option<OwnedSemaphorePermit>,
+    ) -> Result<(), String> {
+        self.events.send_now(event, cell_permit)
+    }
+}
+
+impl<T> BoundedIds<T>
+where
+    T: Clone + Eq + Hash,
+{
+    pub(super) fn remember(&mut self, id: T) -> bool {
+        if !self.ids.insert(id.clone()) {
+            return false;
+        }
+        self.order.push_back(id);
+        while self.order.len() > MAX_RECENT_REQUEST_IDS {
+            if let Some(expired) = self.order.pop_front() {
+                self.ids.remove(&expired);
+            }
+        }
+        true
+    }
+
+    pub(super) fn contains<Q>(&self, id: &Q) -> bool
+    where
+        T: Borrow<Q>,
+        Q: Eq + Hash + ?Sized,
+    {
+        self.ids.contains(id)
+    }
+
+    pub(super) fn remove<Q>(&mut self, id: &Q) -> bool
+    where
+        T: Borrow<Q>,
+        Q: Eq + Hash + ?Sized,
+    {
+        self.ids.remove(id)
+    }
+}

codex-rs/code-mode-host/src/grpc/validation.rs

@@ -0,0 +1,53 @@
+use codex_code_mode_protocol::grpc as proto;
+use tonic::Status;
+use uuid::Uuid;
+
+pub(super) const MAX_IDENTIFIER_BYTES: usize = 256;
+pub(super) const MAX_TOOL_FILTERS: usize = 64;
+pub(super) const MAX_TOOL_DEFINITIONS: usize = 1_024;
+pub(super) const MAX_TOOL_DESCRIPTION_BYTES: usize = 16 * 1_024;
+pub(super) const MAX_TOOL_ERROR_BYTES: usize = 64 * 1_024;
+
+pub(super) fn identifier(value: &str, field: &str) -> Result<(), Status> {
+    if value.is_empty() {
+        return Err(Status::invalid_argument(format!(
+            "code-mode {field} must not be empty"
+        )));
+    }
+    bounded(value, MAX_IDENTIFIER_BYTES, field)
+}
+
+pub(super) fn uuid(value: &str, field: &str) -> Result<Uuid, Status> {
+    identifier(value, field)?;
+    Uuid::parse_str(value)
+        .map_err(|_| Status::invalid_argument(format!("code-mode {field} must be a UUID")))
+}
+
+pub(super) fn bounded(value: &str, maximum: usize, field: &str) -> Result<(), Status> {
+    if value.len() > maximum {
+        return Err(Status::invalid_argument(format!(
+            "code-mode {field} exceeds {maximum} bytes"
+        )));
+    }
+    Ok(())
+}
+
+pub(super) fn tool_name(name: &proto::ToolName) -> Result<(), Status> {
+    identifier(&name.name, "tool name")?;
+    if let Some(namespace) = name.namespace.as_deref() {
+        bounded(namespace, MAX_IDENTIFIER_BYTES, "tool namespace")?;
+    }
+    Ok(())
+}
+
+pub(super) fn tool_filters(filters: &[proto::ToolName]) -> Result<(), Status> {
+    if filters.len() > MAX_TOOL_FILTERS {
+        return Err(Status::invalid_argument(format!(
+            "code-mode tool subscription exceeds {MAX_TOOL_FILTERS} filters"
+        )));
+    }
+    for filter in filters {
+        tool_name(filter)?;
+    }
+    Ok(())
+}

codex-rs/code-mode-host/src/grpc/waits.rs

@@ -0,0 +1,92 @@
+use std::sync::Arc;
+use std::sync::PoisonError;
+
+use tokio_util::sync::CancellationToken;
+use tonic::Status;
+
+use super::session::GrpcSession;
+use super::validation;
+
+pub(super) struct ActiveWait {
+    pub(super) cancellation: CancellationToken,
+    retired: CancellationToken,
+}
+
+pub(super) struct WaitRegistration {
+    session: Arc<GrpcSession>,
+    id: String,
+    cancellation: CancellationToken,
+    retired: CancellationToken,
+}
+
+impl WaitRegistration {
+    pub(super) fn new(session: Arc<GrpcSession>, id: String) -> Result<Self, Status> {
+        validation::identifier(&id, "wait ID")?;
+        let cancellation = CancellationToken::new();
+        let retired = CancellationToken::new();
+        let mut state = session.state.lock().unwrap_or_else(PoisonError::into_inner);
+        if session.closed.is_cancelled() {
+            return Err(Status::cancelled("code-mode session is closed"));
+        }
+        if state.waits.contains_key(&id) || !state.seen_waits.remember(id.clone()) {
+            return Err(Status::already_exists(format!(
+                "code-mode wait ID `{id}` was reused"
+            )));
+        }
+        if state.cancelled_waits.remove(&id) {
+            return Err(Status::cancelled("code-mode wait was cancelled"));
+        }
+        state.waits.insert(
+            id.clone(),
+            ActiveWait {
+                cancellation: cancellation.clone(),
+                retired: retired.clone(),
+            },
+        );
+        drop(state);
+        Ok(Self {
+            session,
+            id,
+            cancellation,
+            retired,
+        })
+    }
+
+    pub(super) fn cancellation(&self) -> &CancellationToken {
+        &self.cancellation
+    }
+}
+
+impl Drop for WaitRegistration {
+    fn drop(&mut self) {
+        self.session
+            .state
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .waits
+            .remove(&self.id);
+        self.retired.cancel();
+    }
+}
+
+impl GrpcSession {
+    pub(super) async fn cancel_wait(&self, id: &str) -> Result<(), Status> {
+        validation::identifier(id, "wait ID")?;
+        let active = {
+            let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+            if let Some(wait) = state.waits.get(id) {
+                Some((wait.cancellation.clone(), wait.retired.clone()))
+            } else {
+                if !state.seen_waits.contains(id) {
+                    state.cancelled_waits.remember(id.to_string());
+                }
+                None
+            }
+        };
+        if let Some((cancellation, retired)) = active {
+            cancellation.cancel();
+            retired.cancelled().await;
+        }
+        Ok(())
+    }
+}

codex-rs/code-mode-host/src/host_tests.rs

@@ -38,8 +38,8 @@ use tokio_util::sync::CancellationToken;
 use tokio_util::task::TaskTracker;
 use uuid::Uuid;
 
+use super::HostLimits;
 use super::HostState;
-use super::MAX_ACTIVE_CELLS;
 use super::MAX_IN_FLIGHT_REQUESTS;
 use super::MAX_RECENT_REQUEST_IDS;
 use super::RequestKind;
@@ -588,11 +588,10 @@ async fn request_task_panic_disconnects_host() {
     let peer = Arc::new(HostPeer::new(outgoing_tx));
     let state = 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(),
-        request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)),
-        active_cell_permits: Arc::new(Semaphore::new(MAX_ACTIVE_CELLS)),
         closing: AtomicBool::new(false),
         peer: Arc::clone(&peer),
     };
@@ -617,11 +616,10 @@ async fn execute_request_id_remains_active_until_initial_response() {
     let peer = Arc::new(HostPeer::new(outgoing_tx));
     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(),
-        request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)),
-        active_cell_permits: Arc::new(Semaphore::new(MAX_ACTIVE_CELLS)),
         closing: AtomicBool::new(false),
         peer,
     });
@@ -679,11 +677,13 @@ async fn active_cell_limit_rejects_execute_without_disconnecting() {
     let peer = Arc::new(HostPeer::new(outgoing_tx));
     let state = HostState {
         sessions: Mutex::new(HashMap::new()),
+        limits: Arc::new(HostLimits {
+            request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)),
+            active_cell_permits: Arc::new(Semaphore::new(/*permits*/ 0)),
+        }),
         seen_session_ids: Mutex::new(SeenSessionIds::default()),
         requests: Mutex::new(RequestRegistry::default()),
         request_tasks: TaskTracker::new(),
-        request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)),
-        active_cell_permits: Arc::new(Semaphore::new(/*permits*/ 0)),
         closing: AtomicBool::new(false),
         peer: Arc::clone(&peer),
     };

codex-rs/code-mode-host/src/lib.rs

@@ -31,7 +31,9 @@ use codex_code_mode_protocol::host::TransportLane;
 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;
@@ -43,9 +45,11 @@ use self::transport::BulkConnectionRegistry;
 use self::transport::ConnectionReader;
 use self::transport::ConnectionWriter;
 
+pub use self::grpc::GrpcCodeModeHost;
 pub use self::transport::DEFAULT_LISTEN_URL;
 
 mod delegate;
+mod grpc;
 mod peer;
 mod transport;
 
@@ -75,6 +79,14 @@ impl HostLimits {
             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 WebSocket transport.
@@ -141,11 +153,10 @@ async fn run_connection(
     };
     let state = Arc::new(HostState {
         sessions: Mutex::new(HashMap::new()),
+        limits,
         seen_session_ids: Mutex::new(SeenSessionIds::default()),
         requests: Mutex::new(RequestRegistry::default()),
         request_tasks: TaskTracker::new(),
-        request_permits: Arc::clone(&limits.request_permits),
-        active_cell_permits: Arc::clone(&limits.active_cell_permits),
         closing: AtomicBool::new(false),
         peer: Arc::clone(&peer),
     });
@@ -368,11 +379,10 @@ async fn negotiate(
 
 struct HostState {
     sessions: Mutex<HashMap<SessionId, Arc<InProcessCodeModeSession>>>,
+    limits: Arc<HostLimits>,
     seen_session_ids: Mutex<SeenSessionIds>,
     requests: Mutex<RequestRegistry>,
     request_tasks: TaskTracker,
-    request_permits: Arc<Semaphore>,
-    active_cell_permits: Arc<Semaphore>,
     closing: AtomicBool,
     peer: Arc<HostPeer>,
 }
@@ -388,7 +398,7 @@ impl HostState {
             .lock()
             .unwrap_or_else(PoisonError::into_inner)
             .start(request_id, RequestKind::from(&request))?;
-        let Ok(permit) = Arc::clone(&self.request_permits).try_acquire_owned() else {
+        let Ok(permit) = self.limits.request_permit() else {
             self.respond(
                 request_id,
                 Err("code-mode host has too many in-flight requests".to_string()),
@@ -468,9 +478,7 @@ impl HostState {
                         return;
                     }
                 };
-                let Ok(active_cell_permit) =
+                let Ok(active_cell_permit) = self.limits.cell_permit() else {
-                    Arc::clone(&self.active_cell_permits).try_acquire_owned()
-                else {
                     self.respond(
                         request_id,
                         Err("code-mode host has too many active cells".to_string()),

codex-rs/code-mode-host/src/transport.rs

@@ -308,7 +308,7 @@ async fn websocket_upgrade_handler(
             )
             .await
             {
-                warn!(%peer_addr, "code-mode host websocket connection failed: {err:#}");
+                warn!(%peer_addr, "code-mode host session failed: {err:#}");
             }
         })
 }

codex-rs/code-mode-host/tests/websocket.rs

@@ -877,7 +877,7 @@ async fn malformed_websocket_frame_does_not_stop_the_listener() -> Result<()> {
                 .next_line()
                 .await?
                 .context("code-mode host exited before reporting the malformed frame")?;
-            if line.contains("code-mode host websocket connection failed") {
+            if line.contains("code-mode host session failed") {
                 return Ok::<_, anyhow::Error>(line);
             }
         }