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.

<img width="2121" height="1001" alt="image"
src="https://github.com/user-attachments/assets/c349a819-2a59-485c-bda4-2caf68ac4c31"
/>

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",

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"] }

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";

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<ToolDefinition>,
     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<FunctionCallOutputContentItem>,
         /// 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<WaitOutcome> for RuntimeResponse {
 #[derive(Debug, PartialEq, Serialize)]
 pub enum RuntimeResponse {
     Yielded {
-        cell_id: String,
+        cell_id: CellId,
         content_items: Vec<FunctionCallOutputContentItem>,
     },
     Terminated {
-        cell_id: String,
+        cell_id: CellId,
         content_items: Vec<FunctionCallOutputContentItem>,
     },
     Result {
-        cell_id: String,
+        cell_id: CellId,
         content_items: Vec<FunctionCallOutputContentItem>,
         error_text: Option<String>,
     },
@@ -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<JsonValue>,
 }
 
-#[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(),

codex-rs/core/src/session/handlers.rs

@@ -630,6 +630,9 @@ async fn shutdown_session_runtime(sess: &Arc<Session>) {
         .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()

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);

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<DispatchMessage>,
+    dispatch_rx: async_channel::Receiver<DispatchMessage>,
+    dispatch_gates: Arc<Mutex<HashMap<CellId, watch::Sender<bool>>>>,
+}
+
+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<ToolRouter>,
+        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<HashMap<CellId, watch::Sender<bool>>>,
+    cell_id: &CellId,
+) -> watch::Sender<bool> {
+    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<HashMap<CellId, watch::Sender<bool>>>,
+    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<HashMap<CellId, watch::Sender<bool>>>,
+    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<Result<JsonValue, String>>,
+    },
+    Notify {
+        call_id: String,
+        cell_id: CellId,
+        text: String,
+        cancellation_token: CancellationToken,
+        response_tx: oneshot::Sender<Result<(), String>>,
+    },
+}
+
+pub(crate) struct CodeModeDispatchWorker {
+    shutdown_tx: Option<oneshot::Sender<()>>,
+}
+
+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<JsonValue, String> {
+        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")
+            })
+    }
+}

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

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<Arc<dyn CodeModeSession>>,
+    dispatch_broker: Arc<CodeModeDispatchBroker>,
 }
 
 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<RuntimeResponse, String> {
-        self.inner.execute(request).await
+    ) -> Result<codex_code_mode::StartedCell, String> {
+        self.session()?.execute(request).await
     }
 
     pub(crate) async fn wait(
         &self,
         request: codex_code_mode::WaitRequest,
     ) -> Result<codex_code_mode::WaitOutcome, String> {
-        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<codex_code_mode::WaitOutcome, String> {
+        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<Session>,
         turn: &Arc<TurnContext>,
         router: Arc<ToolRouter>,
         tracker: SharedTurnDiffTracker,
-    ) -> Option<codex_code_mode::CodeModeTurnWorker> {
-        if !matches!(turn.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) {
+    ) -> Option<CodeModeDispatchWorker> {
+        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<JsonValue, String> {
-        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<dyn CodeModeSession>, 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,

codex-rs/core/src/tools/code_mode/wait_handler.rs

@@ -73,17 +73,24 @@ impl ToolExecutor<ToolInvocation> 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<ToolInvocation> 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