fix: reconcile post-rebase core fallout

Author: ericjuta

codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 use crate::function_tool::FunctionCallError;
 use crate::maybe_emit_implicit_skill_invocation;
 use crate::tools::context::ExecCommandToolOutput;
+use crate::tools::context::ToolCallSource;
 use crate::tools::context::ToolInvocation;
 use crate::tools::context::ToolPayload;
 use crate::tools::context::boxed_tool_output;
@@ -98,6 +99,7 @@ impl ToolExecutor<ToolInvocation> for ExecCommandHandler {
             cancellation_token,
             tracker,
             call_id,
+            source,
             payload,
             ..
         } = invocation;
@@ -112,6 +114,10 @@ impl ToolExecutor<ToolInvocation> for ExecCommandHandler {
         };
 
         let manager: &UnifiedExecProcessManager = &session.services.unified_exec_manager;
+        let cancellation_token = match source {
+            ToolCallSource::CodeMode { .. } => cancellation_token,
+            ToolCallSource::Direct => tokio_util::sync::CancellationToken::new(),
+        };
         let context = UnifiedExecContext::new_with_cancellation(
             session.clone(),
             turn.clone(),

codex-rs/core/src/tools/parallel.rs

@@ -156,11 +156,21 @@ impl ToolCallRuntime {
                                 Err(err) => return Err(Self::tool_task_join_error(err)),
                             }
                         } else {
-                            handle.abort();
-                            match handle.await {
-                                Ok(result) => return result,
-                                Err(err) if err.is_cancelled() => {}
-                                Err(err) => return Err(Self::tool_task_join_error(err)),
+                            match tokio::time::timeout(
+                                TOOL_CANCELLATION_CLEANUP_GRACE,
+                                &mut handle,
+                            )
+                            .await
+                            {
+                                Ok(result) => return result.map_err(Self::tool_task_join_error)?,
+                                Err(_) => {
+                                    handle.abort();
+                                    match handle.await {
+                                        Ok(result) => return result,
+                                        Err(err) if err.is_cancelled() => {}
+                                        Err(err) => return Err(Self::tool_task_join_error(err)),
+                                    }
+                                }
                             }
                         }
                         let response = Self::aborted_response(&call, secs);

codex-rs/core/tests/suite/compact_remote.rs

@@ -3431,6 +3431,7 @@ async fn remote_sampling_context_window_exceeded_compacts_and_retries() -> Resul
             }],
             final_output_json_schema: None,
             responsesapi_client_metadata: None,
+            additional_context: Default::default(),
             thread_settings: Default::default(),
         })
         .await?;
@@ -3445,6 +3446,7 @@ async fn remote_sampling_context_window_exceeded_compacts_and_retries() -> Resul
             }],
             final_output_json_schema: None,
             responsesapi_client_metadata: None,
+            additional_context: Default::default(),
             thread_settings: Default::default(),
         })
         .await?;

codex-rs/core/tests/suite/hooks_operator_smoke.rs

@@ -280,6 +280,7 @@ async fn submit_smoke_turn(
             }],
             final_output_json_schema: None,
             responsesapi_client_metadata: None,
+            additional_context: Default::default(),
             thread_settings: codex_protocol::protocol::ThreadSettingsOverrides {
                 cwd: Some(test.config.cwd.to_path_buf()),
                 approval_policy: Some(approval_policy),