Reuse tool runtime for code mode worker (#14496)

## Summary
- create the turn-scoped `ToolCallRuntime` before starting the code mode
worker so the worker reuses the same runtime and router
- thread the shared runtime through the code mode service/worker path
and use it for nested tool calls
- model aborted tool calls as a concrete `ToolOutput` so aborted
responses still produce valid tool output shapes

## Testing
- `just fmt`
- `cargo test -p codex-core` (still running locally)

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

@@ -4,7 +4,6 @@ use crate::codex::Session;
 use crate::codex::TurnContext;
 use crate::function_tool::FunctionCallError;
 use crate::tools::context::FunctionToolOutput;
-use crate::tools::context::SharedTurnDiffTracker;
 use crate::tools::context::ToolInvocation;
 use crate::tools::context::ToolPayload;
 use crate::tools::registry::ToolHandler;
@@ -25,14 +24,9 @@ impl CodeModeExecuteHandler {
         &self,
         session: std::sync::Arc<Session>,
         turn: std::sync::Arc<TurnContext>,
-        tracker: SharedTurnDiffTracker,
         code: String,
     ) -> Result<FunctionToolOutput, FunctionCallError> {
-        let exec = ExecContext {
-            session,
-            turn,
-            tracker,
-        };
+        let exec = ExecContext { session, turn };
         let enabled_tools = build_enabled_tools(&exec).await;
         let service = &exec.session.services.code_mode_service;
         let stored_values = service.stored_values().await;
@@ -94,7 +88,6 @@ impl ToolHandler for CodeModeExecuteHandler {
         let ToolInvocation {
             session,
             turn,
-            tracker,
             tool_name,
             payload,
             ..
@@ -102,7 +95,7 @@ impl ToolHandler for CodeModeExecuteHandler {
 
         match payload {
             ToolPayload::Custom { input } if tool_name == PUBLIC_TOOL_NAME => {
-                self.execute(session, turn, tracker, input).await
+                self.execute(session, turn, input).await
             }
             _ => Err(FunctionCallError::RespondToModel(format!(
                 "{PUBLIC_TOOL_NAME} expects raw JavaScript source text"

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

@@ -18,8 +18,8 @@ use crate::tools::ToolRouter;
 use crate::tools::code_mode_description::augment_tool_spec_for_code_mode;
 use crate::tools::code_mode_description::code_mode_tool_reference;
 use crate::tools::context::FunctionToolOutput;
-use crate::tools::context::SharedTurnDiffTracker;
 use crate::tools::context::ToolPayload;
+use crate::tools::parallel::ToolCallRuntime;
 use crate::tools::router::ToolCall;
 use crate::tools::router::ToolCallSource;
 use crate::tools::router::ToolRouterParams;
@@ -42,7 +42,6 @@ pub(crate) const DEFAULT_WAIT_YIELD_TIME_MS: u64 = 10_000;
 pub(super) struct ExecContext {
     pub(super) session: Arc<Session>,
     pub(super) turn: Arc<TurnContext>,
-    pub(super) tracker: SharedTurnDiffTracker,
 }
 
 pub(crate) use execute_handler::CodeModeExecuteHandler;
@@ -270,8 +269,10 @@ async fn build_nested_router(exec: &ExecContext) -> ToolRouter {
 
 async fn call_nested_tool(
     exec: ExecContext,
+    tool_runtime: ToolCallRuntime,
     tool_name: String,
     input: Option<JsonValue>,
+    cancellation_token: tokio_util::sync::CancellationToken,
 ) -> JsonValue {
     if tool_name == PUBLIC_TOOL_NAME {
         return JsonValue::String(format!("{PUBLIC_TOOL_NAME} cannot invoke itself"));
@@ -302,14 +303,8 @@ async fn call_nested_tool(
         tool_namespace: None,
         payload,
     };
-    let result = router
-        .dispatch_tool_call_with_code_mode_result(
-            exec.session.clone(),
-            exec.turn.clone(),
-            exec.tracker.clone(),
-            call,
-            ToolCallSource::CodeMode,
-        )
+    let result = tool_runtime
+        .handle_tool_call_with_source(call, ToolCallSource::CodeMode, cancellation_token)
         .await;
 
     match result {

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

@@ -9,8 +9,10 @@ use tracing::warn;
 use crate::codex::Session;
 use crate::codex::TurnContext;
 use crate::features::Feature;
+use crate::tools::ToolRouter;
 use crate::tools::context::SharedTurnDiffTracker;
 use crate::tools::js_repl::resolve_compatible_node;
+use crate::tools::parallel::ToolCallRuntime;
 
 use super::ExecContext;
 use super::PUBLIC_TOOL_NAME;
@@ -65,7 +67,8 @@ impl CodeModeService {
         &self,
         session: &Arc<Session>,
         turn: &Arc<TurnContext>,
-        tracker: &SharedTurnDiffTracker,
+        router: Arc<ToolRouter>,
+        tracker: SharedTurnDiffTracker,
     ) -> Option<CodeModeWorker> {
         if !turn.features.enabled(Feature::CodeMode) {
             return None;
@@ -73,8 +76,9 @@ impl CodeModeService {
         let exec = ExecContext {
             session: Arc::clone(session),
             turn: Arc::clone(turn),
-            tracker: Arc::clone(tracker),
         };
+        let tool_runtime =
+            ToolCallRuntime::new(router, Arc::clone(session), Arc::clone(turn), tracker);
         let mut process_slot = match self.ensure_started().await {
             Ok(process_slot) => process_slot,
             Err(err) => {
@@ -88,7 +92,7 @@ impl CodeModeService {
             );
             return None;
         };
-        Some(process.worker(exec))
+        Some(process.worker(exec, tool_runtime))
     }
 
     pub(crate) async fn allocate_session_id(&self) -> i32 {

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

@@ -54,7 +54,6 @@ impl ToolHandler for CodeModeWaitHandler {
         let ToolInvocation {
             session,
             turn,
-            tracker,
             tool_name,
             payload,
             ..
@@ -63,11 +62,7 @@ impl ToolHandler for CodeModeWaitHandler {
         match payload {
             ToolPayload::Function { arguments } if tool_name == WAIT_TOOL_NAME => {
                 let args: ExecWaitArgs = parse_arguments(&arguments)?;
-                let exec = ExecContext {
-                    session,
-                    turn,
-                    tracker,
-                };
+                let exec = ExecContext { session, turn };
                 let request_id = exec
                     .session
                     .services

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

@@ -1,4 +1,5 @@
 use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
 use tracing::warn;
 
 use super::ExecContext;
@@ -7,6 +8,7 @@ use super::call_nested_tool;
 use super::process::CodeModeProcess;
 use super::process::write_message;
 use super::protocol::HostToNodeMessage;
+use crate::tools::parallel::ToolCallRuntime;
 pub(crate) struct CodeModeWorker {
     shutdown_tx: Option<oneshot::Sender<()>>,
 }
@@ -20,7 +22,11 @@ impl Drop for CodeModeWorker {
 }
 
 impl CodeModeProcess {
-    pub(super) fn worker(&self, exec: ExecContext) -> CodeModeWorker {
+    pub(super) fn worker(
+        &self,
+        exec: ExecContext,
+        tool_runtime: ToolCallRuntime,
+    ) -> CodeModeWorker {
         let (shutdown_tx, mut shutdown_rx) = oneshot::channel();
         let stdin = self.stdin.clone();
         let tool_call_rx = self.tool_call_rx.clone();
@@ -37,13 +43,20 @@ impl CodeModeProcess {
                     break;
                 };
                 let exec = exec.clone();
+                let tool_runtime = tool_runtime.clone();
                 let stdin = stdin.clone();
                 tokio::spawn(async move {
                     let response = HostToNodeMessage::Response {
                         request_id: tool_call.request_id,
                         id: tool_call.id,
-                        code_mode_result: call_nested_tool(exec, tool_call.name, tool_call.input)
-                            .await,
+                        code_mode_result: call_nested_tool(
+                            exec,
+                            tool_runtime,
+                            tool_call.name,
+                            tool_call.input,
+                            CancellationToken::new(),
+                        )
+                        .await,
                     };
                     if let Err(err) = write_message(&stdin, &response).await {
                         warn!("failed to write {PUBLIC_TOOL_NAME} tool response: {err}");