From 61a3dd4387dabbbc9725bfac0b99dec1d902cfb3 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Sat, 8 Aug 2026 04:08:25 +0000 Subject: [PATCH] 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 | 3 + codex-rs/code-mode-host/Cargo.toml | 7 +- .../code-mode-host/src/grpc/conversions.rs | 174 +++++ .../src/grpc/conversions_tests.rs | 113 +++ codex-rs/code-mode-host/src/grpc/delegate.rs | 180 +++++ codex-rs/code-mode-host/src/grpc/events.rs | 151 ++++ codex-rs/code-mode-host/src/grpc/mod.rs | 369 ++++++++++ .../src/grpc/robustness_tests.rs | 665 ++++++++++++++++++ codex-rs/code-mode-host/src/grpc/routing.rs | 307 ++++++++ .../code-mode-host/src/grpc/service_tests.rs | 348 +++++++++ codex-rs/code-mode-host/src/grpc/session.rs | 465 ++++++++++++ .../code-mode-host/src/grpc/validation.rs | 53 ++ codex-rs/code-mode-host/src/grpc/waits.rs | 92 +++ codex-rs/code-mode-host/src/host_tests.rs | 14 +- codex-rs/code-mode-host/src/lib.rs | 24 +- codex-rs/code-mode-host/src/transport.rs | 2 +- codex-rs/code-mode-host/tests/websocket.rs | 2 +- 17 files changed, 2950 insertions(+), 19 deletions(-) create mode 100644 codex-rs/code-mode-host/src/grpc/conversions.rs create mode 100644 codex-rs/code-mode-host/src/grpc/conversions_tests.rs create mode 100644 codex-rs/code-mode-host/src/grpc/delegate.rs create mode 100644 codex-rs/code-mode-host/src/grpc/events.rs create mode 100644 codex-rs/code-mode-host/src/grpc/mod.rs create mode 100644 codex-rs/code-mode-host/src/grpc/robustness_tests.rs create mode 100644 codex-rs/code-mode-host/src/grpc/routing.rs create mode 100644 codex-rs/code-mode-host/src/grpc/service_tests.rs create mode 100644 codex-rs/code-mode-host/src/grpc/session.rs create mode 100644 codex-rs/code-mode-host/src/grpc/validation.rs create mode 100644 codex-rs/code-mode-host/src/grpc/waits.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 4c0cc0e6eb..39654f1d7d 100644 --- a/codex-rs/Cargo.lock +++ b/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", diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index 2f1febc2fb..71c1237d20 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/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 } diff --git a/codex-rs/code-mode-host/src/grpc/conversions.rs b/codex-rs/code-mode-host/src/grpc/conversions.rs new file mode 100644 index 0000000000..7c6f6134bd --- /dev/null +++ b/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, +) -> Result { + 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 { + 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::, _>>()?, + 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 { + 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>, field: &str) -> Result, 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; diff --git a/codex-rs/code-mode-host/src/grpc/conversions_tests.rs b/codex-rs/code-mode-host/src/grpc/conversions_tests.rs new file mode 100644 index 0000000000..61d49156db --- /dev/null +++ b/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()), + }, + )), + } + ); +} diff --git a/codex-rs/code-mode-host/src/grpc/delegate.rs b/codex-rs/code-mode-host/src/grpc/delegate.rs new file mode 100644 index 0000000000..4b540f8c4b --- /dev/null +++ b/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, +} + +impl GrpcDelegate { + pub(super) fn new(session: Weak) -> 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, + id: Option, + kind: CallbackKind, +} + +impl PendingCallback { + fn tool(session: Arc, id: Uuid) -> Self { + Self { + session, + id: Some(id), + kind: CallbackKind::Tool, + } + } + + fn notification(session: Arc, 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), + } + } +} diff --git a/codex-rs/code-mode-host/src/grpc/events.rs b/codex-rs/code-mode-host/src/grpc/events.rs new file mode 100644 index 0000000000..affcd0ceee --- /dev/null +++ b/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, + permits: Arc, + closed: CancellationToken, + writer: TaskTracker, +} + +struct QueuedEvent { + message: proto::SessionEvent, + _queue_permit: OwnedSemaphorePermit, + _cell_permit: Option, +} + +impl EventSender { + pub(super) fn new( + output: mpsc::Sender>, + closed: CancellationToken, + ) -> Self { + let (sender, mut receiver) = mpsc::unbounded_channel::(); + 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, + ) -> 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, + ) -> 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 { + 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) +} diff --git a/codex-rs/code-mode-host/src/grpc/mod.rs b/codex-rs/code-mode-host/src/grpc/mod.rs new file mode 100644 index 0000000000..84145a4bfd --- /dev/null +++ b/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 = Pin> + Send + 'static>>; +type GrpcFuture<'a, T> = Pin, Status>> + Send + 'a>>; + +/// Serves transport-independent, leased code-mode sessions over gRPC. +#[derive(Clone)] +pub struct GrpcCodeModeHost { + state: Arc, +} + +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>, 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, 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>, 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, 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, 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>, 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, 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, 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, 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; + type SubscribeToToolCallsStream = GrpcStream; + type ExecuteStream = GrpcStream; + + fn open_session<'a, 'async_trait>( + &'a self, + request: Request, + ) -> 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, + ) -> 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, + ) -> 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, + ) -> 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, + ) -> 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, + ) -> 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, + ) -> 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, + ) -> 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, + ) -> GrpcFuture<'async_trait, proto::WaitResponse> + where + 'a: 'async_trait, + Self: 'async_trait, + { + Box::pin(self.terminate_request(request.into_inner())) + } +} + +struct ExecutionAdmission { + session: Arc, + execution_id: Option, +} + +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; diff --git a/codex-rs/code-mode-host/src/grpc/robustness_tests.rs b/codex-rs/code-mode-host/src/grpc/robustness_tests.rs new file mode 100644 index 0000000000..ced3290edd --- /dev/null +++ b/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(result: Result) { + 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::>(); + 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::>(); + 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::>(); + 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() + ); +} diff --git a/codex-rs/code-mode-host/src/grpc/routing.rs b/codex-rs/code-mode-host/src/grpc/routing.rs new file mode 100644 index 0000000000..e6a29efdce --- /dev/null +++ b/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, + filters: Vec, + ) -> Result, 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::>() + }; + 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>, + response: oneshot::Sender>, + 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::>(); + (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::>(); + 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, + ) -> 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(¬ification_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(¬ification_id) { + Some(response) => Some(response), + None if state.seen_notifications.contains(¬ification_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(¬ification_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, + ); + } + } +} diff --git a/codex-rs/code-mode-host/src/grpc/service_tests.rs b/codex-rs/code-mode-host/src/grpc/service_tests.rs new file mode 100644 index 0000000000..2ee9e7bcc9 --- /dev/null +++ b/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) { + 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) { + 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, + })), + } + ); +} diff --git a/codex-rs/code-mode-host/src/grpc/session.rs b/codex-rs/code-mode-host/src/grpc/session.rs new file mode 100644 index 0000000000..c88c6add86 --- /dev/null +++ b/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>>, + limits: HostLimits, + delegate_permits: Arc, + control_permits: Arc, +} + +pub(super) struct GrpcSession { + pub(super) id: Uuid, + pub(super) runtime: Arc, + pub(super) closed: CancellationToken, + pub(super) state: Mutex, + events: EventSender, + cells_changed: Notify, + delegate_permits: Arc, +} + +#[derive(Default)] +pub(super) struct SessionState { + pub(super) cells: HashMap, + pending_executions: HashSet, + pending_closures: HashSet, + seen_executions: BoundedIds, + pub(super) subscriptions: Vec, + pub(super) next_subscription: usize, + pub(super) pending_invocations: HashMap, + pub(super) seen_invocations: BoundedIds, + pub(super) pending_notifications: HashMap>, + pub(super) seen_notifications: BoundedIds, + pub(super) waits: HashMap, + 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, + pub(super) sender: mpsc::Sender>, +} + +pub(super) struct PendingInvocation { + pub(super) subscription_id: Uuid, + pub(super) response: oneshot::Sender>, +} + +#[derive(Default)] +pub(super) struct BoundedIds { + ids: HashSet, + order: VecDeque, +} + +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, + limits: CodeModeSessionCellExecutionLimits, + ) -> Result, 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, 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) { + 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 { + 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 { + self.limits + .cell_permit() + .map_err(|_| Status::resource_exhausted("code-mode host has too many active cells")) + } + + pub(super) fn control_permit(&self) -> Result { + 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, + limits: CodeModeSessionCellExecutionLimits, + ) -> Arc { + Arc::new_cyclic(|weak: &Weak| { + 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 { + 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, 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 { + 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 { + 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, + ) -> Result<(), String> { + self.events.send_now(event, cell_permit) + } +} + +impl BoundedIds +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(&self, id: &Q) -> bool + where + T: Borrow, + Q: Eq + Hash + ?Sized, + { + self.ids.contains(id) + } + + pub(super) fn remove(&mut self, id: &Q) -> bool + where + T: Borrow, + Q: Eq + Hash + ?Sized, + { + self.ids.remove(id) + } +} diff --git a/codex-rs/code-mode-host/src/grpc/validation.rs b/codex-rs/code-mode-host/src/grpc/validation.rs new file mode 100644 index 0000000000..7328a96f33 --- /dev/null +++ b/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 { + 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(()) +} diff --git a/codex-rs/code-mode-host/src/grpc/waits.rs b/codex-rs/code-mode-host/src/grpc/waits.rs new file mode 100644 index 0000000000..9f0b0ef477 --- /dev/null +++ b/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, + id: String, + cancellation: CancellationToken, + retired: CancellationToken, +} + +impl WaitRegistration { + pub(super) fn new(session: Arc, id: String) -> Result { + 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(()) + } +} diff --git a/codex-rs/code-mode-host/src/host_tests.rs b/codex-rs/code-mode-host/src/host_tests.rs index 07a2d4d717..f5ac27a5db 100644 --- a/codex-rs/code-mode-host/src/host_tests.rs +++ b/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), }; diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs index dd4cf24bb6..05b119d142 100644 --- a/codex-rs/code-mode-host/src/lib.rs +++ b/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 { + Arc::clone(&self.request_permits).try_acquire_owned() + } + + fn cell_permit(&self) -> Result { + 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>>, + limits: Arc, seen_session_ids: Mutex, requests: Mutex, request_tasks: TaskTracker, - request_permits: Arc, - active_cell_permits: Arc, closing: AtomicBool, peer: Arc, } @@ -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) = - Arc::clone(&self.active_cell_permits).try_acquire_owned() - else { + let Ok(active_cell_permit) = self.limits.cell_permit() else { self.respond( request_id, Err("code-mode host has too many active cells".to_string()), diff --git a/codex-rs/code-mode-host/src/transport.rs b/codex-rs/code-mode-host/src/transport.rs index 2f3ed0a7dc..d4dddd081a 100644 --- a/codex-rs/code-mode-host/src/transport.rs +++ b/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:#}"); } }) } diff --git a/codex-rs/code-mode-host/tests/websocket.rs b/codex-rs/code-mode-host/tests/websocket.rs index 9ff50e1b84..249f06af17 100644 --- a/codex-rs/code-mode-host/tests/websocket.rs +++ b/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); } }