From c9dc0f6338450c2bc8b70b96852f7ac422d75bbf Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Fri, 29 May 2026 11:42:52 -0700 Subject: [PATCH] code-mode: introduce durable session interface (#24180) ## Summary Introduce a `CodeModeSession` interface for executing and managing code-mode cells. This moves cell lifecycle, callback delegation, termination, and shutdown behind a session abstraction, while continuing to use the existing in-process implementation, and the ability to implement an external process one behind this interface. A Codex session owns one `CodeModeSession`, which in turn owns its running cells and stored code-mode state. Each cell is represented to the caller as a `StartedCell`, exposing its cell ID and initial response. It also introduces a `CodeModeSessionDelegate` callback interface. A session uses the delegate to invoke nested host tools and emit notifications while a cell is running, allowing the runtime to communicate with its owning Codex session without depending directly on core turn handling. image --- codex-rs/Cargo.lock | 2 - codex-rs/code-mode/Cargo.toml | 2 - codex-rs/code-mode/src/lib.rs | 13 +- codex-rs/code-mode/src/runtime/mod.rs | 32 +- codex-rs/code-mode/src/service.rs | 952 ++++++++++++------ codex-rs/core/src/session/handlers.rs | 3 + codex-rs/core/src/session/turn.rs | 16 +- codex-rs/core/src/tools/code_mode/delegate.rs | 311 ++++++ .../src/tools/code_mode/execute_handler.rs | 38 +- codex-rs/core/src/tools/code_mode/mod.rs | 107 +- .../core/src/tools/code_mode/wait_handler.rs | 38 +- 11 files changed, 1056 insertions(+), 458 deletions(-) create mode 100644 codex-rs/core/src/tools/code_mode/delegate.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index d5434ad97..2eeafe30b 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2430,8 +2430,6 @@ dependencies = [ name = "codex-code-mode" version = "0.0.0" dependencies = [ - "async-channel", - "async-trait", "codex-protocol", "deno_core_icudata", "pretty_assertions", diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index 19d6c3ab4..879404d8a 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/codex-rs/code-mode/Cargo.toml @@ -16,8 +16,6 @@ sandbox = ["v8/v8_enable_sandbox"] workspace = true [dependencies] -async-channel = { workspace = true } -async-trait = { workspace = true } codex-protocol = { workspace = true } deno_core_icudata = { workspace = true } serde = { workspace = true, features = ["derive"] } diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index 3da8c7732..37ba1d2a9 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -29,9 +29,18 @@ pub use runtime::WaitOutcome; pub use runtime::WaitRequest; pub use runtime::WaitToPendingOutcome; pub use runtime::WaitToPendingRequest; +pub use service::CellId; pub use service::CodeModeService; -pub use service::CodeModeTurnHost; -pub use service::CodeModeTurnWorker; +pub use service::CodeModeSession; +pub use service::CodeModeSessionDelegate; +pub use service::CodeModeSessionProvider; +pub use service::CodeModeSessionProviderFuture; +pub use service::CodeModeSessionResultFuture; +pub use service::InProcessCodeModeSessionProvider; +pub use service::NoopCodeModeSessionDelegate; +pub use service::NotificationFuture; +pub use service::StartedCell; +pub use service::ToolInvocationFuture; pub const PUBLIC_TOOL_NAME: &str = "exec"; pub const WAIT_TOOL_NAME: &str = "wait"; diff --git a/codex-rs/code-mode/src/runtime/mod.rs b/codex-rs/code-mode/src/runtime/mod.rs index 89893d2b2..21757328e 100644 --- a/codex-rs/code-mode/src/runtime/mod.rs +++ b/codex-rs/code-mode/src/runtime/mod.rs @@ -19,6 +19,7 @@ use crate::description::EnabledToolMetadata; use crate::description::ToolDefinition; use crate::description::enabled_tool_metadata; use crate::response::FunctionCallOutputContentItem; +use crate::service::CellId; pub const DEFAULT_EXEC_YIELD_TIME_MS: u64 = 10_000; pub const DEFAULT_WAIT_YIELD_TIME_MS: u64 = 10_000; @@ -27,11 +28,6 @@ const EXIT_SENTINEL: &str = "__codex_code_mode_exit__"; #[derive(Clone, Debug)] pub struct ExecuteRequest { - /// Runtime cell id for this execution. - /// - /// Callers allocate this before execution so tracing, waits, and nested tool - /// calls can refer to the cell as soon as JavaScript starts. - pub cell_id: String, pub tool_call_id: String, pub enabled_tools: Vec, pub source: String, @@ -41,14 +37,13 @@ pub struct ExecuteRequest { #[derive(Clone, Debug)] pub struct WaitRequest { - pub cell_id: String, + pub cell_id: CellId, pub yield_time_ms: u64, - pub terminate: bool, } #[derive(Clone, Debug)] pub struct WaitToPendingRequest { - pub cell_id: String, + pub cell_id: CellId, } /// Result of waiting on a code-mode cell. @@ -73,7 +68,7 @@ pub enum ExecuteToPendingOutcome { /// The cell is waiting for more runtime input after draining the runtime /// input queue that was ready at the pending boundary. Pending { - cell_id: String, + cell_id: CellId, content_items: Vec, /// Runtime tool-call ids emitted before this paused execution frontier /// sealed. Hosts can use these ids to drain their tool-call transport @@ -105,15 +100,15 @@ impl From for RuntimeResponse { #[derive(Debug, PartialEq, Serialize)] pub enum RuntimeResponse { Yielded { - cell_id: String, + cell_id: CellId, content_items: Vec, }, Terminated { - cell_id: String, + cell_id: CellId, content_items: Vec, }, Result { - cell_id: String, + cell_id: CellId, content_items: Vec, error_text: Option, }, @@ -126,23 +121,13 @@ pub enum RuntimeResponse { /// if their tool-call graph requires globally unique ids. #[derive(Debug)] pub struct CodeModeNestedToolCall { - pub cell_id: String, + pub cell_id: CellId, pub runtime_tool_call_id: String, pub tool_name: ToolName, pub tool_kind: CodeModeToolKind, pub input: Option, } -#[derive(Debug)] -pub(crate) enum TurnMessage { - ToolCall(CodeModeNestedToolCall), - Notify { - cell_id: String, - call_id: String, - text: String, - }, -} - #[derive(Debug)] pub(crate) enum RuntimeCommand { ToolResponse { id: String, result: JsonValue }, @@ -460,7 +445,6 @@ mod tests { fn execute_request(source: &str) -> ExecuteRequest { ExecuteRequest { - cell_id: "1".to_string(), tool_call_id: "call_1".to_string(), enabled_tools: Vec::new(), source: source.to_string(), diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index d153f1b41..6589e6ec5 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -1,14 +1,20 @@ use std::collections::HashMap; +use std::fmt; +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use std::time::Duration; -use async_trait::async_trait; +use serde::Deserialize; +use serde::Serialize; use serde_json::Value as JsonValue; use tokio::sync::Mutex; use tokio::sync::mpsc; use tokio::sync::oneshot; +use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; use tracing::warn; @@ -22,35 +28,162 @@ use crate::runtime::RuntimeCommand; use crate::runtime::RuntimeControlCommand; use crate::runtime::RuntimeEvent; use crate::runtime::RuntimeResponse; -use crate::runtime::TurnMessage; use crate::runtime::WaitOutcome; use crate::runtime::WaitRequest; use crate::runtime::WaitToPendingOutcome; use crate::runtime::WaitToPendingRequest; use crate::runtime::spawn_runtime; -#[async_trait] -pub trait CodeModeTurnHost: Send + Sync { - async fn invoke_tool( - &self, +pub type CodeModeSessionResultFuture<'a, T> = + Pin> + Send + 'a>>; +pub type CodeModeSessionProviderFuture<'a> = + CodeModeSessionResultFuture<'a, Arc>; +pub type ToolInvocationFuture<'a> = + Pin> + Send + 'a>>; +pub type NotificationFuture<'a> = Pin> + Send + 'a>>; + +#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] +pub struct CellId(String); + +impl CellId { + pub fn new(value: String) -> Self { + Self(value) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl AsRef for CellId { + fn as_ref(&self) -> &str { + self.as_str() + } +} + +impl fmt::Display for CellId { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(self.as_str()) + } +} + +pub struct StartedCell { + pub cell_id: CellId, + initial_response_rx: oneshot::Receiver, +} + +impl StartedCell { + pub async fn initial_response(self) -> Result { + self.initial_response_rx + .await + .map_err(|_| "exec runtime ended unexpectedly".to_string()) + } +} + +/// Host callbacks used by a code-mode session while cells are executing. +pub trait CodeModeSessionDelegate: Send + Sync { + fn invoke_tool<'a>( + &'a self, invocation: CodeModeNestedToolCall, cancellation_token: CancellationToken, - ) -> Result; + ) -> ToolInvocationFuture<'a>; - async fn notify(&self, call_id: String, cell_id: String, text: String) -> Result<(), String>; + fn notify<'a>( + &'a self, + call_id: String, + cell_id: CellId, + text: String, + cancellation_token: CancellationToken, + ) -> NotificationFuture<'a>; + + /// Releases delegate state associated with a cell after it reaches a terminal state. + fn cell_closed(&self, cell_id: &CellId); +} + +pub struct NoopCodeModeSessionDelegate; + +impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate { + fn invoke_tool<'a>( + &'a self, + _invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async move { + cancellation_token.cancelled().await; + Err("code mode nested tools are unavailable".to_string()) + }) + } + + fn notify<'a>( + &'a self, + _call_id: String, + _cell_id: CellId, + _text: String, + _cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async { Ok(()) }) + } + + fn cell_closed(&self, _cell_id: &CellId) {} +} + +/// A durable code-mode session owned by one Codex thread. +/// +/// Cells executed in the same session share stored values. Separate sessions +/// must keep those values isolated. Implementations may execute cells +/// in-process or remotely. +pub trait CodeModeSession: Send + Sync { + fn execute<'a>( + &'a self, + request: ExecuteRequest, + ) -> CodeModeSessionResultFuture<'a, StartedCell>; + + fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome>; + + fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome>; + + fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()>; +} + +/// Creates code-mode sessions for one Codex thread. +/// +/// Providers choose where a session executes and receive the host delegate that +/// the session should use for nested tool calls and notifications. +pub trait CodeModeSessionProvider: Send + Sync { + fn create_session<'a>( + &'a self, + delegate: Arc, + ) -> CodeModeSessionProviderFuture<'a>; +} + +#[derive(Default)] +pub struct InProcessCodeModeSessionProvider; + +impl CodeModeSessionProvider for InProcessCodeModeSessionProvider { + fn create_session<'a>( + &'a self, + delegate: Arc, + ) -> CodeModeSessionProviderFuture<'a> { + Box::pin(async move { + let session: Arc = + Arc::new(CodeModeService::with_delegate(delegate)); + Ok(session) + }) + } } #[derive(Clone)] -struct SessionHandle { - control_tx: mpsc::UnboundedSender, +struct CellHandle { + control_tx: mpsc::UnboundedSender, runtime_tx: std::sync::mpsc::Sender, + cancellation_token: CancellationToken, } struct Inner { stored_values: Mutex>, - sessions: Mutex>, - turn_message_tx: async_channel::Sender, - turn_message_rx: async_channel::Receiver, + cells: Mutex>, + delegate: Arc, + shutting_down: AtomicBool, next_cell_id: AtomicU64, } @@ -60,50 +193,50 @@ pub struct CodeModeService { impl CodeModeService { pub fn new() -> Self { - let (turn_message_tx, turn_message_rx) = async_channel::unbounded(); + Self::with_delegate(Arc::new(NoopCodeModeSessionDelegate)) + } + pub fn with_delegate(delegate: Arc) -> Self { Self { inner: Arc::new(Inner { stored_values: Mutex::new(HashMap::new()), - sessions: Mutex::new(HashMap::new()), - turn_message_tx, - turn_message_rx, + cells: Mutex::new(HashMap::new()), + delegate, + shutting_down: AtomicBool::new(false), next_cell_id: AtomicU64::new(1), }), } } - async fn stored_values(&self) -> HashMap { - self.inner.stored_values.lock().await.clone() + fn allocate_cell_id(&self) -> CellId { + CellId::new( + self.inner + .next_cell_id + .fetch_add(1, Ordering::Relaxed) + .to_string(), + ) } - /// Reserves the runtime cell id for a future `execute` request. - /// - /// The runtime can issue nested tool calls before the first `execute` - /// response is returned. Hosts that need a parent trace object for those - /// nested calls should allocate the cell id up front and pass it back on the - /// `ExecuteRequest`. - pub fn allocate_cell_id(&self) -> String { - self.inner - .next_cell_id - .fetch_add(1, Ordering::Relaxed) - .to_string() - } - - pub async fn execute(&self, request: ExecuteRequest) -> Result { + pub async fn execute(&self, request: ExecuteRequest) -> Result { + if self.inner.shutting_down.load(Ordering::Acquire) { + return Err("code mode session is shutting down".to_string()); + } let initial_yield_time_ms = request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS); let (response_tx, response_rx) = oneshot::channel(); - self.start_session( + let cell_id = self.allocate_cell_id(); + self.start_cell( + cell_id.clone(), request, - SessionResponseSender::Runtime(response_tx), + CellResponseSender::Runtime(response_tx), Some(initial_yield_time_ms), PendingRuntimeMode::Continue, ) .await?; - response_rx - .await - .map_err(|_| "exec runtime ended unexpectedly".to_string()) + Ok(StartedCell { + cell_id, + initial_response_rx: response_rx, + }) } pub async fn execute_to_pending( @@ -111,9 +244,11 @@ impl CodeModeService { request: ExecuteRequest, ) -> Result { let (response_tx, response_rx) = oneshot::channel(); - self.start_session( + let cell_id = self.allocate_cell_id(); + self.start_cell( + cell_id, request, - SessionResponseSender::ExecuteToPending(response_tx), + CellResponseSender::ExecuteToPending(response_tx), /*initial_yield_time_ms*/ None, PendingRuntimeMode::PauseUntilResumed, ) @@ -124,47 +259,50 @@ impl CodeModeService { .map_err(|_| "exec runtime ended unexpectedly".to_string()) } - async fn start_session( + async fn start_cell( &self, + cell_id: CellId, request: ExecuteRequest, - initial_response_tx: SessionResponseSender, + initial_response_tx: CellResponseSender, initial_yield_time_ms: Option, pending_mode: PendingRuntimeMode, ) -> Result<(), String> { - let cell_id = request.cell_id.clone(); let (event_tx, event_rx) = mpsc::unbounded_channel(); let (control_tx, control_rx) = mpsc::unbounded_channel(); - let stored_values = self.stored_values().await; + let stored_values = self.inner.stored_values.lock().await.clone(); + let cancellation_token = CancellationToken::new(); let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = { - let mut sessions = self.inner.sessions.lock().await; - if sessions.contains_key(&cell_id) { + let mut cells = self.inner.cells.lock().await; + if self.inner.shutting_down.load(Ordering::Acquire) { + return Err("code mode session is shutting down".to_string()); + } + if cells.contains_key(&cell_id) { return Err(format!("exec cell {cell_id} already exists")); } let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = spawn_runtime(stored_values, request, event_tx, pending_mode)?; - // Keep the session registry locked through insertion so a - // caller-owned cell id cannot race with another execute and replace - // a live runtime. - sessions.insert( + cells.insert( cell_id.clone(), - SessionHandle { + CellHandle { control_tx, runtime_tx: runtime_tx.clone(), + cancellation_token: cancellation_token.clone(), }, ); (runtime_tx, runtime_control_tx, runtime_terminate_handle) }; - tokio::spawn(run_session_control( + tokio::spawn(run_cell_control( Arc::clone(&self.inner), - SessionControlContext { - cell_id: cell_id.clone(), + CellControlContext { + cell_id, runtime_tx, runtime_control_tx, pending_mode, runtime_terminate_handle, + cancellation_token, }, event_rx, control_rx, @@ -176,34 +314,44 @@ impl CodeModeService { } pub async fn wait(&self, request: WaitRequest) -> Result { - let cell_id = request.cell_id.clone(); - let handle = self - .inner - .sessions - .lock() - .await - .get(&request.cell_id) - .cloned(); + let WaitRequest { + cell_id, + yield_time_ms, + } = request; + let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); let Some(handle) = handle else { return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); }; let (response_tx, response_rx) = oneshot::channel(); - let control_message = if request.terminate { - SessionControlCommand::Terminate { response_tx } - } else { - SessionControlCommand::Poll { - yield_time_ms: request.yield_time_ms, - response_tx, - } + let control_message = CellControlCommand::Poll { + yield_time_ms, + response_tx, }; if handle.control_tx.send(control_message).is_err() { return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); } match response_rx.await { Ok(response) => Ok(WaitOutcome::LiveCell(response)), - Err(_) => Ok(WaitOutcome::MissingCell(missing_cell_response( - request.cell_id, - ))), + Err(_) => Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))), + } + } + + pub async fn terminate(&self, cell_id: CellId) -> Result { + let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); + let Some(handle) = handle else { + return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); + }; + let (response_tx, response_rx) = oneshot::channel(); + if handle + .control_tx + .send(CellControlCommand::Terminate { response_tx }) + .is_err() + { + return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); + } + match response_rx.await { + Ok(response) => Ok(WaitOutcome::LiveCell(response)), + Err(_) => Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))), } } @@ -211,14 +359,8 @@ impl CodeModeService { &self, request: WaitToPendingRequest, ) -> Result { - let cell_id = request.cell_id.clone(); - let handle = self - .inner - .sessions - .lock() - .await - .get(&request.cell_id) - .cloned(); + let cell_id = request.cell_id; + let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); let Some(handle) = handle else { return Ok(WaitToPendingOutcome::MissingCell(missing_cell_response( cell_id, @@ -227,7 +369,7 @@ impl CodeModeService { let (response_tx, response_rx) = oneshot::channel(); if handle .control_tx - .send(SessionControlCommand::PollToPending { response_tx }) + .send(CellControlCommand::PollToPending { response_tx }) .is_err() { return Ok(WaitToPendingOutcome::MissingCell(missing_cell_response( @@ -237,74 +379,33 @@ impl CodeModeService { match response_rx.await { Ok(response) => Ok(WaitToPendingOutcome::LiveCell(response)), Err(_) => Ok(WaitToPendingOutcome::MissingCell(missing_cell_response( - request.cell_id, + cell_id, ))), } } - pub fn start_turn_worker(&self, host: Arc) -> CodeModeTurnWorker { - let (shutdown_tx, mut shutdown_rx) = oneshot::channel(); - let inner = Arc::clone(&self.inner); - let turn_message_rx = self.inner.turn_message_rx.clone(); - - tokio::spawn(async move { - loop { - let next_message = tokio::select! { - _ = &mut shutdown_rx => break, - message = turn_message_rx.recv() => message.ok(), - }; - let Some(next_message) = next_message else { - break; - }; - match next_message { - TurnMessage::Notify { - cell_id, - call_id, - text, - } => { - if let Err(err) = host.notify(call_id, cell_id.clone(), text).await { - warn!( - "failed to deliver code mode notification for cell {cell_id}: {err}" - ); - } - } - TurnMessage::ToolCall(invocation) => { - let host = Arc::clone(&host); - let inner = Arc::clone(&inner); - tokio::spawn(async move { - let cell_id = invocation.cell_id.clone(); - let runtime_tool_call_id = invocation.runtime_tool_call_id.clone(); - let response = - host.invoke_tool(invocation, CancellationToken::new()).await; - let runtime_tx = inner - .sessions - .lock() - .await - .get(&cell_id) - .map(|handle| handle.runtime_tx.clone()); - let Some(runtime_tx) = runtime_tx else { - return; - }; - let command = match response { - Ok(result) => RuntimeCommand::ToolResponse { - id: runtime_tool_call_id, - result, - }, - Err(error_text) => RuntimeCommand::ToolError { - id: runtime_tool_call_id, - error_text, - }, - }; - let _ = runtime_tx.send(command); - }); - } - } - } - }); - - CodeModeTurnWorker { - shutdown_tx: Some(shutdown_tx), + pub async fn shutdown(&self) -> Result<(), String> { + self.inner.shutting_down.store(true, Ordering::Release); + let handles = self + .inner + .cells + .lock() + .await + .values() + .cloned() + .collect::>(); + for handle in handles { + handle.cancellation_token.cancel(); + let (response_tx, _response_rx) = oneshot::channel(); + let _ = handle + .control_tx + .send(CellControlCommand::Terminate { response_tx }); + let _ = handle.runtime_tx.send(RuntimeCommand::Terminate); } + while !self.inner.cells.lock().await.is_empty() { + tokio::task::yield_now().await; + } + Ok(()) } } @@ -314,19 +415,44 @@ impl Default for CodeModeService { } } -pub struct CodeModeTurnWorker { - shutdown_tx: Option>, -} - -impl Drop for CodeModeTurnWorker { +impl Drop for CodeModeService { fn drop(&mut self) { - if let Some(shutdown_tx) = self.shutdown_tx.take() { - let _ = shutdown_tx.send(()); + self.inner.shutting_down.store(true, Ordering::Release); + if let Ok(cells) = self.inner.cells.try_lock() { + for handle in cells.values() { + handle.cancellation_token.cancel(); + let (response_tx, _response_rx) = oneshot::channel(); + let _ = handle + .control_tx + .send(CellControlCommand::Terminate { response_tx }); + let _ = handle.runtime_tx.send(RuntimeCommand::Terminate); + } } } } -enum SessionControlCommand { +impl CodeModeSession for CodeModeService { + fn execute<'a>( + &'a self, + request: ExecuteRequest, + ) -> CodeModeSessionResultFuture<'a, StartedCell> { + Box::pin(CodeModeService::execute(self, request)) + } + + fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(CodeModeService::wait(self, request)) + } + + fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(CodeModeService::terminate(self, cell_id)) + } + + fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { + Box::pin(CodeModeService::shutdown(self)) + } +} + +enum CellControlCommand { Poll { yield_time_ms: u64, response_tx: oneshot::Sender, @@ -339,7 +465,7 @@ enum SessionControlCommand { }, } -enum SessionResponseSender { +enum CellResponseSender { Runtime(oneshot::Sender), ExecuteToPending(oneshot::Sender), } @@ -349,15 +475,16 @@ struct PendingResult { error_text: Option, } -struct SessionControlContext { - cell_id: String, +struct CellControlContext { + cell_id: CellId, runtime_tx: std::sync::mpsc::Sender, runtime_control_tx: std::sync::mpsc::Sender, pending_mode: PendingRuntimeMode, runtime_terminate_handle: v8::IsolateHandle, + cancellation_token: CancellationToken, } -fn missing_cell_response(cell_id: String) -> RuntimeResponse { +fn missing_cell_response(cell_id: CellId) -> RuntimeResponse { RuntimeResponse::Result { error_text: Some(format!("exec cell {cell_id} not found")), cell_id, @@ -365,29 +492,29 @@ fn missing_cell_response(cell_id: String) -> RuntimeResponse { } } -fn pending_result_response(cell_id: &str, result: PendingResult) -> RuntimeResponse { +fn pending_result_response(cell_id: &CellId, result: PendingResult) -> RuntimeResponse { RuntimeResponse::Result { - cell_id: cell_id.to_string(), + cell_id: cell_id.clone(), content_items: result.content_items, error_text: result.error_text, } } -fn send_terminal_response(response_tx: SessionResponseSender, response: RuntimeResponse) { +fn send_terminal_response(response_tx: CellResponseSender, response: RuntimeResponse) { match response_tx { - SessionResponseSender::Runtime(response_tx) => { + CellResponseSender::Runtime(response_tx) => { let _ = response_tx.send(response); } - SessionResponseSender::ExecuteToPending(response_tx) => { + CellResponseSender::ExecuteToPending(response_tx) => { let _ = response_tx.send(ExecuteToPendingOutcome::Completed(response)); } } } fn send_or_buffer_result( - cell_id: &str, + cell_id: &CellId, result: PendingResult, - response_tx: &mut Option, + response_tx: &mut Option, pending_result: &mut Option, ) -> bool { if let Some(response_tx) = response_tx.take() { @@ -401,42 +528,41 @@ fn send_or_buffer_result( } fn send_yield_response( - cell_id: &str, + cell_id: &CellId, content_items: &mut Vec, - response_tx: &mut Option, + response_tx: &mut Option, ) { let Some(current_response_tx) = response_tx.take() else { return; }; match current_response_tx { - SessionResponseSender::Runtime(response_tx) => { + CellResponseSender::Runtime(response_tx) => { let _ = response_tx.send(RuntimeResponse::Yielded { - cell_id: cell_id.to_string(), + cell_id: cell_id.clone(), content_items: std::mem::take(content_items), }); } - SessionResponseSender::ExecuteToPending(execute_to_pending_tx) => { - *response_tx = Some(SessionResponseSender::ExecuteToPending( - execute_to_pending_tx, - )); + CellResponseSender::ExecuteToPending(execute_to_pending_tx) => { + *response_tx = Some(CellResponseSender::ExecuteToPending(execute_to_pending_tx)); } } } -async fn run_session_control( +async fn run_cell_control( inner: Arc, - context: SessionControlContext, + context: CellControlContext, mut event_rx: mpsc::UnboundedReceiver, - mut control_rx: mpsc::UnboundedReceiver, - initial_response_tx: SessionResponseSender, + mut control_rx: mpsc::UnboundedReceiver, + initial_response_tx: CellResponseSender, initial_yield_time_ms: Option, ) { - let SessionControlContext { + let CellControlContext { cell_id, runtime_tx, runtime_control_tx, pending_mode, runtime_terminate_handle, + cancellation_token, } = context; let mut content_items = Vec::new(); let mut pending_tool_call_ids = Vec::new(); @@ -445,6 +571,7 @@ async fn run_session_control( let mut termination_requested = false; let mut runtime_closed = false; let mut yield_timer: Option>> = None; + let mut notification_tasks = JoinSet::new(); loop { tokio::select! { @@ -492,11 +619,11 @@ async fn run_session_control( RuntimeEvent::Pending => { if let Some(current_response_tx) = response_tx.take() { match current_response_tx { - SessionResponseSender::Runtime(runtime_response_tx) => { + CellResponseSender::Runtime(runtime_response_tx) => { response_tx = - Some(SessionResponseSender::Runtime(runtime_response_tx)); + Some(CellResponseSender::Runtime(runtime_response_tx)); } - SessionResponseSender::ExecuteToPending(response_tx) => { + CellResponseSender::ExecuteToPending(response_tx) => { let _ = response_tx.send(ExecuteToPendingOutcome::Pending { cell_id: cell_id.clone(), content_items: std::mem::take(&mut content_items), @@ -516,11 +643,26 @@ async fn run_session_control( send_yield_response(&cell_id, &mut content_items, &mut response_tx); } RuntimeEvent::Notify { call_id, text } => { - let _ = inner.turn_message_tx.send(TurnMessage::Notify { - cell_id: cell_id.clone(), - call_id, - text, - }).await; + let delegate = Arc::clone(&inner.delegate); + let cell_id = cell_id.clone(); + let cancellation_token = cancellation_token.child_token(); + notification_tasks.spawn(async move { + tokio::select! { + result = delegate.notify( + call_id, + cell_id.clone(), + text, + cancellation_token.clone(), + ) => { + if let Err(err) = result { + warn!( + "failed to deliver code mode notification for cell {cell_id}: {err}" + ); + } + } + _ = cancellation_token.cancelled() => {} + } + }); } RuntimeEvent::ToolCall { id, @@ -533,15 +675,25 @@ async fn run_session_control( } let tool_call = CodeModeNestedToolCall { cell_id: cell_id.clone(), - runtime_tool_call_id: id, + runtime_tool_call_id: id.clone(), tool_name: name, tool_kind: kind, input, }; - let _ = inner - .turn_message_tx - .send(TurnMessage::ToolCall(tool_call)) - .await; + let delegate = Arc::clone(&inner.delegate); + let runtime_tx = runtime_tx.clone(); + let cancellation_token = cancellation_token.child_token(); + tokio::spawn(async move { + let response = tokio::select! { + response = delegate.invoke_tool(tool_call, cancellation_token.clone()) => response, + _ = cancellation_token.cancelled() => return, + }; + let command = match response { + Ok(result) => RuntimeCommand::ToolResponse { id, result }, + Err(error_text) => RuntimeCommand::ToolError { id, error_text }, + }; + let _ = runtime_tx.send(command); + }); } RuntimeEvent::Result { stored_value_writes, @@ -558,6 +710,7 @@ async fn run_session_control( } break; } + drain_notification_tasks(&mut notification_tasks).await; inner .stored_values .lock() @@ -578,12 +731,19 @@ async fn run_session_control( } } } + task_result = notification_tasks.join_next(), if !notification_tasks.is_empty() => { + if let Some(Err(err)) = task_result + && !err.is_cancelled() + { + warn!("code mode notification task failed: {err}"); + } + } maybe_command = control_rx.recv() => { let Some(command) = maybe_command else { break; }; match command { - SessionControlCommand::Poll { + CellControlCommand::Poll { yield_time_ms, response_tx: next_response_tx, } => { @@ -591,11 +751,11 @@ async fn run_session_control( let _ = next_response_tx.send(pending_result_response(&cell_id, result)); break; } - response_tx = Some(SessionResponseSender::Runtime(next_response_tx)); + response_tx = Some(CellResponseSender::Runtime(next_response_tx)); yield_timer = Some(Box::pin(tokio::time::sleep(Duration::from_millis(yield_time_ms)))); resume_paused_runtime(&runtime_control_tx, pending_mode); } - SessionControlCommand::PollToPending { + CellControlCommand::PollToPending { response_tx: next_response_tx, } => { if let Some(result) = pending_result.take() { @@ -605,18 +765,19 @@ async fn run_session_control( break; } response_tx = - Some(SessionResponseSender::ExecuteToPending(next_response_tx)); + Some(CellResponseSender::ExecuteToPending(next_response_tx)); yield_timer = None; resume_paused_runtime(&runtime_control_tx, pending_mode); } - SessionControlCommand::Terminate { response_tx: next_response_tx } => { + CellControlCommand::Terminate { response_tx: next_response_tx } => { if let Some(result) = pending_result.take() { let _ = next_response_tx.send(pending_result_response(&cell_id, result)); break; } - response_tx = Some(SessionResponseSender::Runtime(next_response_tx)); + response_tx = Some(CellResponseSender::Runtime(next_response_tx)); termination_requested = true; + cancellation_token.cancel(); yield_timer = None; let _ = runtime_tx.send(RuntimeCommand::Terminate); terminate_paused_runtime(&runtime_control_tx, pending_mode); @@ -650,8 +811,21 @@ async fn run_session_control( } let _ = runtime_tx.send(RuntimeCommand::Terminate); + cancellation_token.cancel(); + drain_notification_tasks(&mut notification_tasks).await; terminate_paused_runtime(&runtime_control_tx, pending_mode); - inner.sessions.lock().await.remove(&cell_id); + inner.cells.lock().await.remove(&cell_id); + inner.delegate.cell_closed(&cell_id); +} + +async fn drain_notification_tasks(notification_tasks: &mut JoinSet<()>) { + while let Some(result) = notification_tasks.join_next().await { + if let Err(err) = result + && !err.is_cancelled() + { + warn!("code mode notification task failed: {err}"); + } + } } fn resume_paused_runtime( @@ -677,6 +851,7 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::AtomicU64; + use std::sync::atomic::Ordering; use std::time::Duration; use codex_protocol::ToolName; @@ -685,19 +860,21 @@ mod tests { use tokio::sync::mpsc; use tokio::sync::oneshot; + use super::CellControlCommand; + use super::CellControlContext; + use super::CellId; + use super::CellResponseSender; use super::CodeModeService; use super::Inner; + use super::NoopCodeModeSessionDelegate; use super::PendingRuntimeMode; use super::RuntimeCommand; use super::RuntimeResponse; - use super::SessionControlCommand; - use super::SessionControlContext; - use super::SessionResponseSender; use super::WaitOutcome; use super::WaitRequest; use super::WaitToPendingOutcome; use super::WaitToPendingRequest; - use super::run_session_control; + use super::run_cell_control; use crate::CodeModeToolKind; use crate::FunctionCallOutputContentItem; use crate::ToolDefinition; @@ -708,7 +885,6 @@ mod tests { fn execute_request(source: &str) -> ExecuteRequest { ExecuteRequest { - cell_id: "1".to_string(), tool_call_id: "call_1".to_string(), enabled_tools: Vec::new(), source: source.to_string(), @@ -717,13 +893,26 @@ mod tests { } } + fn cell_id(value: &str) -> CellId { + CellId::new(value.to_string()) + } + + async fn execute(service: &CodeModeService, request: ExecuteRequest) -> RuntimeResponse { + service + .execute(request) + .await + .unwrap() + .initial_response() + .await + .unwrap() + } + fn test_inner() -> Arc { - let (turn_message_tx, turn_message_rx) = async_channel::unbounded(); Arc::new(Inner { stored_values: Mutex::new(HashMap::new()), - sessions: Mutex::new(HashMap::new()), - turn_message_tx, - turn_message_rx, + cells: Mutex::new(HashMap::new()), + delegate: Arc::new(NoopCodeModeSessionDelegate), + shutting_down: std::sync::atomic::AtomicBool::new(false), next_cell_id: AtomicU64::new(1), }) } @@ -732,19 +921,20 @@ mod tests { async fn synchronous_exit_returns_successfully() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#"text("before"); exit(); text("after");"#.to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "before".to_string(), }], @@ -753,6 +943,116 @@ mod tests { ); } + #[tokio::test] + async fn stored_values_are_shared_between_cells_but_not_sessions() { + let first_session = CodeModeService::new(); + let second_session = CodeModeService::new(); + + let write_response = execute( + &first_session, + ExecuteRequest { + source: r#"store("key", "visible");"#.to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + let same_session = execute( + &first_session, + ExecuteRequest { + source: r#"text(String(load("key")));"#.to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + let other_session = execute( + &second_session, + ExecuteRequest { + source: r#"text(String(load("key")));"#.to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + write_response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: Vec::new(), + error_text: None, + } + ); + assert_eq!( + same_session, + RuntimeResponse::Result { + cell_id: cell_id("2"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "visible".to_string(), + }], + error_text: None, + } + ); + assert_eq!( + other_session, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "undefined".to_string(), + }], + error_text: None, + } + ); + } + + #[tokio::test] + async fn shutdown_interrupts_cpu_bound_cells() { + let service = CodeModeService::new(); + + let cell = service + .execute(ExecuteRequest { + source: "while (true) {}".to_string(), + ..execute_request("") + }) + .await + .unwrap(); + assert_eq!( + cell.initial_response().await.unwrap(), + RuntimeResponse::Yielded { + cell_id: cell_id("1"), + content_items: Vec::new(), + } + ); + + tokio::time::timeout(Duration::from_secs(1), service.shutdown()) + .await + .unwrap() + .unwrap(); + } + + #[tokio::test] + async fn start_cell_rejects_new_cell_after_shutdown_begins() { + let service = CodeModeService::new(); + service.inner.shutting_down.store(true, Ordering::Release); + let (response_tx, _response_rx) = oneshot::channel(); + + let error = service + .start_cell( + cell_id("late-cell"), + execute_request(""), + CellResponseSender::Runtime(response_tx), + Some(/*initial_yield_time_ms*/ 1), + PendingRuntimeMode::Continue, + ) + .await + .unwrap_err(); + + assert_eq!(error, "code mode session is shutting down".to_string()); + assert!(service.inner.cells.lock().await.is_empty()); + } + #[tokio::test] async fn execute_to_pending_returns_completed_for_synchronous_results() { let service = CodeModeService::new(); @@ -769,7 +1069,7 @@ mod tests { assert_eq!( response, ExecuteToPendingOutcome::Completed(RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), }], @@ -797,7 +1097,7 @@ mod tests { assert_eq!( response, ExecuteToPendingOutcome::Pending { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "before".to_string(), }], @@ -805,19 +1105,12 @@ mod tests { } ); - let termination = service - .wait(WaitRequest { - cell_id: "1".to_string(), - yield_time_ms: 1, - terminate: true, - }) - .await - .unwrap(); + let termination = service.terminate(cell_id("1")).await.unwrap(); assert_eq!( termination, WaitOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), }) ); @@ -853,25 +1146,18 @@ await Promise.all([ assert_eq!( response, ExecuteToPendingOutcome::Pending { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: vec!["tool-1".to_string(), "tool-2".to_string()], } ); - let termination = service - .wait(WaitRequest { - cell_id: "1".to_string(), - yield_time_ms: 1, - terminate: true, - }) - .await - .unwrap(); + let termination = service.terminate(cell_id("1")).await.unwrap(); assert_eq!( termination, WaitOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), }) ); @@ -910,7 +1196,7 @@ await Promise.all([ assert_eq!( initial_response, ExecuteToPendingOutcome::Pending { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: vec!["tool-1".to_string(), "tool-2".to_string()], } @@ -918,10 +1204,10 @@ await Promise.all([ let runtime_tx = service .inner - .sessions + .cells .lock() .await - .get("1") + .get(&cell_id("1")) .unwrap() .runtime_tx .clone(); @@ -932,7 +1218,7 @@ await Promise.all([ let resumed_response = tokio::time::timeout( Duration::from_secs(1), service.wait_to_pending(WaitToPendingRequest { - cell_id: "1".to_string(), + cell_id: cell_id("1"), }), ) .await @@ -942,25 +1228,18 @@ await Promise.all([ assert_eq!( resumed_response, WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Pending { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: vec!["tool-3".to_string()], }) ); - let termination = service - .wait(WaitRequest { - cell_id: "1".to_string(), - yield_time_ms: 1, - terminate: true, - }) - .await - .unwrap(); + let termination = service.terminate(cell_id("1")).await.unwrap(); assert_eq!( termination, WaitOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), }) ); @@ -987,7 +1266,7 @@ await new Promise(() => {}); assert_eq!( initial_response, ExecuteToPendingOutcome::Pending { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: Vec::new(), } @@ -995,10 +1274,10 @@ await new Promise(() => {}); let runtime_tx = service .inner - .sessions + .cells .lock() .await - .get("1") + .get(&cell_id("1")) .unwrap() .runtime_tx .clone(); @@ -1009,7 +1288,7 @@ await new Promise(() => {}); let resumed_response = tokio::time::timeout( Duration::from_secs(1), service.wait_to_pending(WaitToPendingRequest { - cell_id: "1".to_string(), + cell_id: cell_id("1"), }), ) .await @@ -1019,7 +1298,7 @@ await new Promise(() => {}); assert_eq!( resumed_response, WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Pending { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "after".to_string(), }], @@ -1027,19 +1306,12 @@ await new Promise(() => {}); }) ); - let termination = service - .wait(WaitRequest { - cell_id: "1".to_string(), - yield_time_ms: 1, - terminate: true, - }) - .await - .unwrap(); + let termination = service.terminate(cell_id("1")).await.unwrap(); assert_eq!( termination, WaitOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), }) ); @@ -1065,7 +1337,7 @@ text("done"); assert_eq!( initial_response, ExecuteToPendingOutcome::Pending { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: Vec::new(), } @@ -1073,10 +1345,10 @@ text("done"); let runtime_tx = service .inner - .sessions + .cells .lock() .await - .get("1") + .get(&cell_id("1")) .unwrap() .runtime_tx .clone(); @@ -1087,7 +1359,7 @@ text("done"); let resumed_response = tokio::time::timeout( Duration::from_secs(1), service.wait_to_pending(WaitToPendingRequest { - cell_id: "1".to_string(), + cell_id: cell_id("1"), }), ) .await @@ -1098,7 +1370,7 @@ text("done"); resumed_response, WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Completed( RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), }], @@ -1112,19 +1384,20 @@ text("done"); async fn v8_console_is_not_exposed_on_global_this() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#"text(String(Object.hasOwn(globalThis, "console")));"#.to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "false".to_string(), }], @@ -1137,8 +1410,9 @@ text("done"); async fn date_locale_string_formats_with_icu_data() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" const value = new Date("2025-01-02T03:04:05Z") .toLocaleString("fr-FR", { @@ -1156,14 +1430,14 @@ text(value); .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "jeudi 2 janvier \u{e0} 03:04:05".to_string(), }], @@ -1176,8 +1450,9 @@ text(value); async fn intl_date_time_format_formats_with_icu_data() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" const formatter = new Intl.DateTimeFormat("fr-FR", { weekday: "long", @@ -1194,14 +1469,14 @@ text(formatter.format(new Date("2025-01-02T03:04:05Z"))); .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "jeudi 2 janvier \u{e0} 03:04:05".to_string(), }], @@ -1214,8 +1489,9 @@ text(formatter.format(new Date("2025-01-02T03:04:05Z"))); async fn output_helpers_return_undefined() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" const returnsUndefined = [ text("first"), @@ -1227,14 +1503,14 @@ text(JSON.stringify(returnsUndefined)); .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![ FunctionCallOutputContentItem::InputText { text: "first".to_string(), @@ -1256,8 +1532,9 @@ text(JSON.stringify(returnsUndefined)); async fn image_helper_accepts_raw_mcp_image_block_with_original_detail() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" image({ type: "image", @@ -1269,14 +1546,14 @@ image({ .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputImage { image_url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==".to_string(), detail: Some(crate::ImageDetail::Original), @@ -1290,8 +1567,9 @@ image({ async fn image_helper_second_arg_overrides_explicit_object_detail() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" image( { @@ -1304,14 +1582,14 @@ image( .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputImage { image_url: "https://example.com/image.jpg".to_string(), detail: Some(crate::ImageDetail::Original), @@ -1325,8 +1603,9 @@ image( async fn image_helper_second_arg_overrides_raw_mcp_image_detail() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" image( { @@ -1341,14 +1620,14 @@ image( .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputImage { image_url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==".to_string(), detail: Some(crate::ImageDetail::High), @@ -1362,8 +1641,9 @@ image( async fn image_helper_accepts_low_detail() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" image({ image_url: "https://example.com/image.jpg", @@ -1373,14 +1653,14 @@ image({ .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputImage { image_url: "https://example.com/image.jpg".to_string(), detail: Some(crate::ImageDetail::Low), @@ -1394,8 +1674,9 @@ image({ async fn image_helper_rejects_unsupported_detail() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" image({ image_url: "https://example.com/image.jpg", @@ -1405,14 +1686,14 @@ image({ .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), error_text: Some( "image detail must be one of: auto, low, high, original".to_string() @@ -1425,8 +1706,9 @@ image({ async fn image_helper_rejects_raw_mcp_result_container() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" image({ content: [ @@ -1443,14 +1725,14 @@ image({ .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), error_text: Some( "image expects a non-empty image URL string, an object with image_url and optional detail, or a raw MCP image block".to_string(), @@ -1465,9 +1747,8 @@ image({ let response = service .wait(WaitRequest { - cell_id: "missing".to_string(), + cell_id: cell_id("missing"), yield_time_ms: 1, - terminate: false, }) .await .unwrap(); @@ -1475,7 +1756,7 @@ image({ assert_eq!( response, WaitOutcome::MissingCell(RuntimeResponse::Result { - cell_id: "missing".to_string(), + cell_id: cell_id("missing"), content_items: Vec::new(), error_text: Some("exec cell missing not found".to_string()), }) @@ -1501,18 +1782,19 @@ image({ ) .unwrap(); - tokio::spawn(run_session_control( + tokio::spawn(run_cell_control( inner, - SessionControlContext { - cell_id: "cell-1".to_string(), + CellControlContext { + cell_id: cell_id("cell-1"), runtime_tx: runtime_tx.clone(), runtime_control_tx, pending_mode: PendingRuntimeMode::Continue, runtime_terminate_handle, + cancellation_token: tokio_util::sync::CancellationToken::new(), }, event_rx, control_rx, - SessionResponseSender::Runtime(initial_response_tx), + CellResponseSender::Runtime(initial_response_tx), Some(/*initial_yield_time_ms*/ 60_000), )); @@ -1521,14 +1803,14 @@ image({ assert_eq!( initial_response_rx.await.unwrap(), RuntimeResponse::Yielded { - cell_id: "cell-1".to_string(), + cell_id: cell_id("cell-1"), content_items: Vec::new(), } ); let (terminate_response_tx, terminate_response_rx) = oneshot::channel(); control_tx - .send(SessionControlCommand::Terminate { + .send(CellControlCommand::Terminate { response_tx: terminate_response_tx, }) .unwrap(); @@ -1545,7 +1827,7 @@ image({ assert_eq!( terminate_response.await, RuntimeResponse::Terminated { - cell_id: "cell-1".to_string(), + cell_id: cell_id("cell-1"), content_items: Vec::new(), } ); diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 135ea6119..7730a30ad 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -630,6 +630,9 @@ async fn shutdown_session_runtime(sess: &Arc) { .unified_exec_manager .terminate_all_processes() .await; + if let Err(err) = sess.services.code_mode_service.shutdown().await { + warn!("failed to shutdown code mode session: {err}"); + } let mcp_shutdown = { let mut manager = sess.services.mcp_connection_manager.write().await; manager.begin_shutdown() diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 8219c1901..7944be57a 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -943,16 +943,12 @@ async fn run_sampling_request( Arc::clone(&turn_context), Arc::clone(&turn_diff_tracker), ); - let _code_mode_worker = sess - .services - .code_mode_service - .start_turn_worker( - &sess, - &turn_context, - Arc::clone(&router), - Arc::clone(&turn_diff_tracker), - ) - .await; + let _code_mode_worker = sess.services.code_mode_service.start_turn_worker( + &sess, + &turn_context, + Arc::clone(&router), + Arc::clone(&turn_diff_tracker), + ); let max_retries = turn_context.provider.info().stream_max_retries(); let mut retries = 0; let mut initial_input = Some(input); diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs new file mode 100644 index 000000000..06bb71a28 --- /dev/null +++ b/codex-rs/core/src/tools/code_mode/delegate.rs @@ -0,0 +1,311 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex; + +use codex_code_mode::CellId; +use codex_code_mode::CodeModeNestedToolCall; +use codex_code_mode::CodeModeSessionDelegate; +use codex_code_mode::NotificationFuture; +use codex_code_mode::ToolInvocationFuture; +use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::ResponseItem; +use serde_json::Value as JsonValue; +use tokio::sync::oneshot; +use tokio::sync::watch; +use tokio_util::sync::CancellationToken; + +use super::ExecContext; +use super::PUBLIC_TOOL_NAME; +use super::call_nested_tool; +use crate::tools::ToolRouter; +use crate::tools::context::SharedTurnDiffTracker; +use crate::tools::parallel::ToolCallRuntime; + +pub(super) struct CodeModeDispatchBroker { + dispatch_tx: async_channel::Sender, + dispatch_rx: async_channel::Receiver, + dispatch_gates: Arc>>>, +} + +impl CodeModeDispatchBroker { + pub(super) fn new() -> Self { + let (dispatch_tx, dispatch_rx) = async_channel::unbounded(); + Self { + dispatch_tx, + dispatch_rx, + dispatch_gates: Arc::new(Mutex::new(HashMap::new())), + } + } + + pub(super) fn mark_cell_ready_for_dispatch(&self, cell_id: &CellId) { + dispatch_gate(&self.dispatch_gates, cell_id).send_replace(true); + } + + pub(super) fn close_cell(&self, cell_id: &CellId) { + remove_dispatch_gate(&self.dispatch_gates, cell_id); + } + + pub(super) fn start_turn_worker( + &self, + exec: ExecContext, + router: Arc, + tracker: SharedTurnDiffTracker, + ) -> CodeModeDispatchWorker { + let tool_runtime = ToolCallRuntime::new( + router, + Arc::clone(&exec.session), + Arc::clone(&exec.turn), + tracker, + ); + let host = Arc::new(CoreTurnHost { exec, tool_runtime }); + let dispatch_rx = self.dispatch_rx.clone(); + let dispatch_gates = Arc::clone(&self.dispatch_gates); + let (shutdown_tx, mut shutdown_rx) = oneshot::channel(); + tokio::spawn(async move { + loop { + let message = tokio::select! { + _ = &mut shutdown_rx => break, + message = dispatch_rx.recv() => message.ok(), + }; + let Some(message) = message else { + break; + }; + match message { + DispatchMessage::Notify { + call_id, + cell_id, + text, + cancellation_token, + response_tx, + } => { + let response = if wait_until_cell_ready_for_dispatch( + &dispatch_gates, + &cell_id, + &cancellation_token, + ) + .await + { + host.notify(call_id, cell_id, text).await + } else { + remove_dispatch_gate(&dispatch_gates, &cell_id); + Err("code mode notification cancelled".to_string()) + }; + let _ = response_tx.send(response); + } + DispatchMessage::InvokeTool { + invocation, + cancellation_token, + response_tx, + } => { + let cell_id = invocation.cell_id.clone(); + if !wait_until_cell_ready_for_dispatch( + &dispatch_gates, + &cell_id, + &cancellation_token, + ) + .await + { + remove_dispatch_gate(&dispatch_gates, &cell_id); + continue; + } + let host = Arc::clone(&host); + tokio::spawn(async move { + let response = tokio::select! { + response = host.invoke_tool( + invocation, + cancellation_token.clone(), + ) => response, + _ = cancellation_token.cancelled() => return, + }; + let _ = response_tx.send(response); + }); + } + } + } + }); + CodeModeDispatchWorker { + shutdown_tx: Some(shutdown_tx), + } + } +} + +fn dispatch_gate( + dispatch_gates: &Mutex>>, + cell_id: &CellId, +) -> watch::Sender { + let mut dispatch_gates = match dispatch_gates.lock() { + Ok(dispatch_gates) => dispatch_gates, + Err(poisoned) => poisoned.into_inner(), + }; + dispatch_gates + .entry(cell_id.clone()) + .or_insert_with(|| watch::channel(false).0) + .clone() +} + +fn remove_dispatch_gate( + dispatch_gates: &Mutex>>, + cell_id: &CellId, +) { + let mut dispatch_gates = match dispatch_gates.lock() { + Ok(dispatch_gates) => dispatch_gates, + Err(poisoned) => poisoned.into_inner(), + }; + dispatch_gates.remove(cell_id); +} + +async fn wait_until_cell_ready_for_dispatch( + dispatch_gates: &Mutex>>, + cell_id: &CellId, + cancellation_token: &CancellationToken, +) -> bool { + if cancellation_token.is_cancelled() { + return false; + } + let mut ready_rx = dispatch_gate(dispatch_gates, cell_id).subscribe(); + loop { + if *ready_rx.borrow_and_update() { + return true; + } + tokio::select! { + changed = ready_rx.changed() => { + if changed.is_err() { + return false; + } + } + _ = cancellation_token.cancelled() => return false, + } + } +} + +impl CodeModeSessionDelegate for CodeModeDispatchBroker { + fn invoke_tool<'a>( + &'a self, + invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async move { + if cancellation_token.is_cancelled() { + return Err("code mode nested tool call cancelled".to_string()); + } + let (response_tx, response_rx) = oneshot::channel(); + self.dispatch_tx + .send(DispatchMessage::InvokeTool { + invocation, + cancellation_token: cancellation_token.clone(), + response_tx, + }) + .await + .map_err(|_| "code mode nested tool dispatcher is unavailable".to_string())?; + tokio::select! { + response = response_rx => response + .map_err(|_| "code mode nested tool dispatcher stopped".to_string())?, + _ = cancellation_token.cancelled() => { + Err("code mode nested tool call cancelled".to_string()) + } + } + }) + } + + fn notify<'a>( + &'a self, + call_id: String, + cell_id: CellId, + text: String, + cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async move { + if cancellation_token.is_cancelled() { + return Err("code mode notification cancelled".to_string()); + } + let (response_tx, response_rx) = oneshot::channel(); + self.dispatch_tx + .send(DispatchMessage::Notify { + call_id, + cell_id, + text, + cancellation_token: cancellation_token.clone(), + response_tx, + }) + .await + .map_err(|_| "code mode notification dispatcher is unavailable".to_string())?; + tokio::select! { + response = response_rx => response + .map_err(|_| "code mode notification dispatcher stopped".to_string())?, + _ = cancellation_token.cancelled() => { + Err("code mode notification cancelled".to_string()) + } + } + }) + } + + fn cell_closed(&self, cell_id: &CellId) { + self.close_cell(cell_id); + } +} + +enum DispatchMessage { + InvokeTool { + invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + response_tx: oneshot::Sender>, + }, + Notify { + call_id: String, + cell_id: CellId, + text: String, + cancellation_token: CancellationToken, + response_tx: oneshot::Sender>, + }, +} + +pub(crate) struct CodeModeDispatchWorker { + shutdown_tx: Option>, +} + +impl Drop for CodeModeDispatchWorker { + fn drop(&mut self) { + if let Some(shutdown_tx) = self.shutdown_tx.take() { + let _ = shutdown_tx.send(()); + } + } +} + +struct CoreTurnHost { + exec: ExecContext, + tool_runtime: ToolCallRuntime, +} + +impl CoreTurnHost { + async fn invoke_tool( + &self, + invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> Result { + call_nested_tool( + self.exec.clone(), + self.tool_runtime.clone(), + invocation, + cancellation_token, + ) + .await + .map_err(|error| error.to_string()) + } + + async fn notify(&self, call_id: String, cell_id: CellId, text: String) -> Result<(), String> { + if text.trim().is_empty() { + return Ok(()); + } + self.exec + .session + .inject_if_running(vec![ResponseItem::CustomToolCallOutput { + call_id, + name: Some(PUBLIC_TOOL_NAME.to_string()), + output: FunctionCallOutputPayload::from_text(text), + }]) + .await + .map_err(|_| { + format!("failed to inject exec notify message for cell {cell_id}: no active turn") + }) + } +} diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 29f562384..1d69c8f4c 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -38,9 +38,22 @@ impl CodeModeExecuteHandler { let exec = ExecContext { session, turn }; let enabled_tools = codex_tools::collect_code_mode_tool_definitions(&self.nested_tool_specs); - // Allocate before starting V8 so the trace can create the parent - // CodeCell before model-authored JavaScript issues nested tool calls. - let runtime_cell_id = exec.session.services.code_mode_service.allocate_cell_id(); + let started_at = std::time::Instant::now(); + let started_cell = exec + .session + .services + .code_mode_service + .execute(codex_code_mode::ExecuteRequest { + tool_call_id: call_id.clone(), + enabled_tools, + source: args.code.clone(), + yield_time_ms: args.yield_time_ms, + max_output_tokens: args.max_output_tokens, + }) + .await + .map_err(FunctionCallError::RespondToModel)?; + let cell_id = started_cell.cell_id.clone(); + let runtime_cell_id = cell_id.to_string(); let code_cell_trace = exec .session .services @@ -51,19 +64,12 @@ impl CodeModeExecuteHandler { call_id.as_str(), args.code.as_str(), ); - let started_at = std::time::Instant::now(); - let response = exec - .session + exec.session .services .code_mode_service - .execute(codex_code_mode::ExecuteRequest { - cell_id: runtime_cell_id, - tool_call_id: call_id, - enabled_tools, - source: args.code, - yield_time_ms: args.yield_time_ms, - max_output_tokens: args.max_output_tokens, - }) + .mark_cell_ready_for_dispatch(&cell_id); + let response = started_cell + .initial_response() .await .map_err(FunctionCallError::RespondToModel)?; // Record the raw runtime boundary. The model-visible custom-tool output @@ -74,6 +80,10 @@ impl CodeModeExecuteHandler { // here when the first response also ended the runtime. if !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) { code_cell_trace.record_ended(&response); + exec.session + .services + .code_mode_service + .finish_cell_dispatch(&cell_id); } handle_runtime_response(&exec, response, args.max_output_tokens, started_at) .await diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index ff005646c..ade23cf3b 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -1,3 +1,4 @@ +mod delegate; mod execute_handler; pub(crate) mod execute_spec; mod response_adapter; @@ -7,13 +8,12 @@ pub(crate) mod wait_spec; use std::sync::Arc; use std::time::Duration; +use codex_code_mode::CellId; use codex_code_mode::CodeModeNestedToolCall; +use codex_code_mode::CodeModeSession; use codex_code_mode::CodeModeToolKind; -use codex_code_mode::CodeModeTurnHost; use codex_code_mode::RuntimeResponse; use codex_protocol::models::FunctionCallOutputContentItem; -use codex_protocol::models::FunctionCallOutputPayload; -use codex_protocol::models::ResponseItem; use serde_json::Value as JsonValue; use tokio_util::sync::CancellationToken; @@ -36,6 +36,8 @@ use codex_utils_output_truncation::TruncationPolicy; use codex_utils_output_truncation::formatted_truncate_text_content_items_with_policy; use codex_utils_output_truncation::truncate_function_output_items_with_policy; +use delegate::CodeModeDispatchBroker; +use delegate::CodeModeDispatchWorker; pub(crate) use execute_handler::CodeModeExecuteHandler; use response_adapter::into_function_call_output_content_items; pub(crate) use wait_handler::CodeModeWaitHandler; @@ -56,42 +58,67 @@ pub(crate) struct ExecContext { } pub(crate) struct CodeModeService { - inner: codex_code_mode::CodeModeService, + session: Option>, + dispatch_broker: Arc, } impl CodeModeService { pub(crate) fn new() -> Self { + let dispatch_broker = Arc::new(CodeModeDispatchBroker::new()); Self { - inner: codex_code_mode::CodeModeService::new(), + session: Some(Arc::new(codex_code_mode::CodeModeService::with_delegate( + dispatch_broker.clone(), + ))), + dispatch_broker, } } - pub(crate) fn allocate_cell_id(&self) -> String { - self.inner.allocate_cell_id() - } - pub(crate) async fn execute( &self, request: codex_code_mode::ExecuteRequest, - ) -> Result { - self.inner.execute(request).await + ) -> Result { + self.session()?.execute(request).await } pub(crate) async fn wait( &self, request: codex_code_mode::WaitRequest, ) -> Result { - self.inner.wait(request).await + self.session()?.wait(request).await } - pub(crate) async fn start_turn_worker( + pub(crate) async fn terminate( + &self, + cell_id: CellId, + ) -> Result { + self.session()?.terminate(cell_id).await + } + + pub(crate) async fn shutdown(&self) -> Result<(), String> { + match &self.session { + Some(session) => session.shutdown().await, + None => Ok(()), + } + } + + pub(crate) fn mark_cell_ready_for_dispatch(&self, cell_id: &codex_code_mode::CellId) { + self.dispatch_broker.mark_cell_ready_for_dispatch(cell_id); + } + + pub(crate) fn finish_cell_dispatch(&self, cell_id: &CellId) { + self.dispatch_broker.close_cell(cell_id); + } + + pub(crate) fn start_turn_worker( &self, session: &Arc, turn: &Arc, router: Arc, tracker: SharedTurnDiffTracker, - ) -> Option { - if !matches!(turn.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) { + ) -> Option { + if !matches!(turn.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) + || self.session.is_none() + { return None; } @@ -99,50 +126,16 @@ impl CodeModeService { session: Arc::clone(session), turn: Arc::clone(turn), }; - let tool_runtime = - ToolCallRuntime::new(router, Arc::clone(session), Arc::clone(turn), tracker); - let host = Arc::new(CoreTurnHost { exec, tool_runtime }); - Some(self.inner.start_turn_worker(host)) - } -} - -struct CoreTurnHost { - exec: ExecContext, - tool_runtime: ToolCallRuntime, -} - -#[async_trait::async_trait] -impl CodeModeTurnHost for CoreTurnHost { - async fn invoke_tool( - &self, - invocation: CodeModeNestedToolCall, - cancellation_token: CancellationToken, - ) -> Result { - call_nested_tool( - self.exec.clone(), - self.tool_runtime.clone(), - invocation, - cancellation_token, + Some( + self.dispatch_broker + .start_turn_worker(exec, router, tracker), ) - .await - .map_err(|error| error.to_string()) } - async fn notify(&self, call_id: String, cell_id: String, text: String) -> Result<(), String> { - if text.trim().is_empty() { - return Ok(()); - } - self.exec - .session - .inject_if_running(vec![ResponseItem::CustomToolCallOutput { - call_id, - name: Some(PUBLIC_TOOL_NAME.to_string()), - output: FunctionCallOutputPayload::from_text(text), - }]) - .await - .map_err(|_| { - format!("failed to inject exec notify message for cell {cell_id}: no active turn") - }) + fn session(&self) -> Result<&Arc, String> { + self.session + .as_ref() + .ok_or_else(|| "code mode is unavailable".to_string()) } } @@ -273,7 +266,7 @@ async fn call_nested_tool( .handle_tool_call_with_source( call, ToolCallSource::CodeMode { - cell_id, + cell_id: cell_id.to_string(), runtime_tool_call_id, }, cancellation_token, diff --git a/codex-rs/core/src/tools/code_mode/wait_handler.rs b/codex-rs/core/src/tools/code_mode/wait_handler.rs index d0c0453df..74cc314a7 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -73,17 +73,24 @@ impl ToolExecutor for CodeModeWaitHandler { let args: ExecWaitArgs = parse_arguments(&arguments)?; let exec = ExecContext { session, turn }; let started_at = std::time::Instant::now(); - let wait_response = exec - .session - .services - .code_mode_service - .wait(codex_code_mode::WaitRequest { - cell_id: args.cell_id, - yield_time_ms: args.yield_time_ms, - terminate: args.terminate, - }) - .await - .map_err(FunctionCallError::RespondToModel)?; + let cell_id = codex_code_mode::CellId::new(args.cell_id); + let wait_response = if args.terminate { + exec.session + .services + .code_mode_service + .terminate(cell_id) + .await + } else { + exec.session + .services + .code_mode_service + .wait(codex_code_mode::WaitRequest { + cell_id, + yield_time_ms: args.yield_time_ms, + }) + .await + } + .map_err(FunctionCallError::RespondToModel)?; if let codex_code_mode::WaitOutcome::LiveCell(response) = &wait_response && !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) { @@ -98,8 +105,15 @@ impl ToolExecutor for CodeModeWaitHandler { exec.session .services .rollout_thread_trace - .code_cell_trace_context(exec.turn.sub_id.as_str(), runtime_cell_id) + .code_cell_trace_context( + exec.turn.sub_id.as_str(), + runtime_cell_id.as_str(), + ) .record_ended(response); + exec.session + .services + .code_mode_service + .finish_cell_dispatch(runtime_cell_id); } handle_runtime_response(&exec, wait_response.into(), args.max_tokens, started_at) .await