fix: avoid hardcoded shell cancellation check

Co-authored-by: Codex noreply@openai.com

codex-rs/core/src/exec.rs

@@ -56,6 +56,7 @@ const SIGKILL_CODE: i32 = 9;
 const TIMEOUT_CODE: i32 = 64;
 const EXIT_CODE_SIGNAL_BASE: i32 = 128; // conventional shell: 128 + signal
 const EXEC_TIMEOUT_EXIT_CODE: i32 = 124; // conventional timeout exit code
+const CANCELLATION_TERMINATION_GRACE_PERIOD: Duration = Duration::from_millis(50);
 
 // I/O buffer sizing
 const READ_CHUNK_SIZE: usize = 8192; // bytes per read
@@ -1354,15 +1355,47 @@ async fn consume_output(
             (exit_status, false)
         }
         outcome = &mut expiration_wait => {
-            kill_child_process_group(&mut child)?;
-            child.start_kill()?;
-            let timed_out = matches!(outcome, Some(ExecExpirationOutcome::TimedOut));
-            let exit_status = if timed_out {
-                synthetic_exit_status(EXIT_CODE_SIGNAL_BASE + TIMEOUT_CODE)
-            } else {
-                synthetic_exit_status_for_code(/*code*/ 1)
-            };
-            (exit_status, timed_out)
+            match outcome {
+                Some(ExecExpirationOutcome::TimedOut) => {
+                    kill_child_process_group(&mut child)?;
+                    child.start_kill()?;
+                    (
+                        synthetic_exit_status(EXIT_CODE_SIGNAL_BASE + TIMEOUT_CODE),
+                        true,
+                    )
+                }
+                Some(ExecExpirationOutcome::Cancelled) => {
+                    let process_group_id = child.id();
+                    let should_escalate = if let Some(process_group_id) = process_group_id {
+                        codex_utils_pty::process_group::terminate_process_group(process_group_id)?
+                    } else {
+                        false
+                    };
+                    match tokio::time::timeout(
+                        CANCELLATION_TERMINATION_GRACE_PERIOD,
+                        child.wait(),
+                    )
+                    .await
+                    {
+                        Ok(status) => {
+                            status?;
+                            if should_escalate
+                                && let Some(process_group_id) = process_group_id
+                            {
+                                codex_utils_pty::process_group::kill_process_group(
+                                    process_group_id,
+                                )?;
+                            }
+                        }
+                        Err(_) => {
+                            kill_child_process_group(&mut child)?;
+                            child.start_kill()?;
+                        }
+                    }
+                    (synthetic_exit_status_for_code(/*code*/ 1), false)
+                }
+                None => unreachable!("expiration wait only resolves while expiration is active"),
+            }
         }
         _ = tokio::signal::ctrl_c() => {
             kill_child_process_group(&mut child)?;

codex-rs/core/src/exec_tests.rs

@@ -1128,6 +1128,115 @@ async fn process_exec_tool_call_respects_cancellation_token() -> Result<()> {
     Ok(())
 }
 
+#[cfg(unix)]
+#[tokio::test]
+async fn process_exec_tool_call_cancellation_allows_sigterm_cleanup() -> Result<()> {
+    let temp_dir = tempfile::TempDir::new()?;
+    let ready_marker = temp_dir.path().join("ready");
+    let cleanup_marker = temp_dir.path().join("cleanup");
+    let descendant_pid_marker = temp_dir.path().join("descendant-pid");
+    let command = vec![
+        "/bin/sh".to_string(),
+        "-c".to_string(),
+        concat!(
+            "(trap '' TERM; sleep 60) &\n",
+            "printf '%s' \"$!\" > \"$DESCENDANT_PID_MARKER\"\n",
+            "trap 'printf cleaned > \"$CLEANUP_MARKER\"; exit 0' TERM\n",
+            "printf ready > \"$READY_MARKER\"\n",
+            "while :; do sleep 1; done"
+        )
+        .to_string(),
+    ];
+    let cwd = codex_utils_absolute_path::AbsolutePathBuf::current_dir()?;
+    let mut env: HashMap<String, String> = std::env::vars().collect();
+    env.insert(
+        "READY_MARKER".to_string(),
+        ready_marker.to_string_lossy().into_owned(),
+    );
+    env.insert(
+        "CLEANUP_MARKER".to_string(),
+        cleanup_marker.to_string_lossy().into_owned(),
+    );
+    env.insert(
+        "DESCENDANT_PID_MARKER".to_string(),
+        descendant_pid_marker.to_string_lossy().into_owned(),
+    );
+    let cancel_token = CancellationToken::new();
+    let cancel_tx = cancel_token.clone();
+    tokio::spawn(async move {
+        for _ in 0..50 {
+            if ready_marker.exists() {
+                cancel_tx.cancel();
+                return;
+            }
+            tokio::time::sleep(Duration::from_millis(20)).await;
+        }
+        cancel_tx.cancel();
+    });
+    let params = ExecParams {
+        command,
+        cwd: cwd.clone(),
+        expiration: ExecExpiration::DefaultTimeout.with_cancellation(cancel_token),
+        capture_policy: ExecCapturePolicy::ShellTool,
+        env,
+        network: None,
+        sandbox_permissions: SandboxPermissions::UseDefault,
+        windows_sandbox_level: codex_protocol::config_types::WindowsSandboxLevel::Disabled,
+        windows_sandbox_private_desktop: false,
+        justification: None,
+        arg0: None,
+    };
+
+    let result = timeout(
+        Duration::from_secs(5),
+        process_exec_tool_call(
+            params,
+            &PermissionProfile::Disabled,
+            &cwd,
+            &None,
+            /*use_legacy_landlock*/ false,
+            /*stdout_stream*/ None,
+        ),
+    )
+    .await
+    .expect("cancellation should stop the process promptly");
+    let output = result.expect("cancellation should return a non-timeout exec result");
+    assert!(!output.timed_out);
+    assert_eq!(
+        std::fs::read_to_string(cleanup_marker)?,
+        "cleaned",
+        "SIGTERM cleanup trap should run before cancellation falls back to a hard kill"
+    );
+    let descendant_pid = std::fs::read_to_string(descendant_pid_marker)?
+        .parse::<i32>()
+        .map_err(|error| {
+            io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!("failed to parse descendant pid: {error}"),
+            )
+        })?;
+    let mut killed = false;
+    for _ in 0..20 {
+        if unsafe { libc::kill(descendant_pid, 0) } == -1
+            && let Some(libc::ESRCH) = std::io::Error::last_os_error().raw_os_error()
+        {
+            killed = true;
+            break;
+        }
+        tokio::time::sleep(Duration::from_millis(100)).await;
+    }
+    if !killed {
+        unsafe {
+            libc::kill(descendant_pid, libc::SIGKILL);
+        }
+    }
+    assert!(
+        killed,
+        "TERM-ignoring descendant process with pid {descendant_pid} is still alive"
+    );
+    Ok(())
+}
+
 #[cfg(unix)]
 fn long_running_command() -> Vec<String> {
     vec![

codex-rs/core/src/sandboxing/mod.rs

@@ -28,7 +28,6 @@ use codex_sandboxing::SandboxType;
 use codex_sandboxing::compatibility_sandbox_policy_for_permission_profile;
 use codex_utils_absolute_path::AbsolutePathBuf;
 use std::collections::HashMap;
-
 #[derive(Debug)]
 pub(crate) struct ExecOptions {
     pub(crate) expiration: ExecExpiration,

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

@@ -9909,6 +9909,74 @@ async fn rejects_escalated_permissions_when_policy_not_on_request() {
         ExecApprovalRequirement::Skip { .. }
     ));
 }
+
+#[cfg(unix)]
+#[tokio::test]
+async fn shell_tool_cancellation_waits_for_runtime_cleanup() -> anyhow::Result<()> {
+    let session = make_session_with_config(|config| {
+        let cwd = config.cwd.clone();
+        config
+            .permissions
+            .set_legacy_sandbox_policy(SandboxPolicy::DangerFullAccess, cwd.as_path())
+            .expect("test setup should allow sandbox policy");
+    })
+    .await?;
+    let turn_context = session.new_default_turn().await;
+    let session = Arc::new(session);
+    let turn_context = Arc::new(turn_context);
+    let temp_dir = tempfile::TempDir::new()?;
+    let ready_marker = temp_dir.path().join("ready");
+    let cleanup_marker = temp_dir.path().join("cleanup");
+    let command = format!(
+        "trap 'printf cleaned > \"{}\"; exit 0' TERM\nprintf ready > \"{}\"\nwhile :; do sleep 1; done",
+        cleanup_marker.display(),
+        ready_marker.display(),
+    );
+    let item = ResponseItem::FunctionCall {
+        id: None,
+        name: "shell_command".to_string(),
+        namespace: None,
+        arguments: serde_json::json!({
+            "command": command,
+            "workdir": ".",
+            "timeout_ms": 60_000,
+        })
+        .to_string(),
+        call_id: "shell-cleanup-call".to_string(),
+    };
+    let call = ToolRouter::build_tool_call(item)?
+        .expect("shell command response item should build a tool call");
+    let cancellation_token = CancellationToken::new();
+    let cancellation_tx = cancellation_token.clone();
+    let handle = tokio::spawn(
+        test_tool_runtime(Arc::clone(&session), Arc::clone(&turn_context))
+            .handle_tool_call(call, cancellation_token),
+    );
+
+    let mut ready = false;
+    for _ in 0..50 {
+        if ready_marker.exists() {
+            ready = true;
+            break;
+        }
+        tokio::time::sleep(Duration::from_millis(20)).await;
+    }
+    if !ready {
+        cancellation_tx.cancel();
+        let _ = timeout(Duration::from_secs(5), handle).await;
+        anyhow::bail!("shell command should reach the ready marker");
+    }
+
+    cancellation_tx.cancel();
+    timeout(Duration::from_secs(5), handle)
+        .await
+        .expect("cancelled shell tool should finish promptly")
+        .expect("shell tool task should join")
+        .expect("cancelled shell tool should return a response item");
+    assert_eq!(std::fs::read_to_string(cleanup_marker)?, "cleaned");
+    Ok(())
+}
+
 #[tokio::test]
 async fn unified_exec_rejects_escalated_permissions_when_policy_not_on_request() {
     use crate::sandboxing::SandboxPermissions;

codex-rs/core/src/tools/handlers/shell.rs

@@ -2,6 +2,7 @@ use codex_features::Feature;
 use codex_protocol::models::ShellCommandToolCallParams;
 use serde_json::Value as JsonValue;
 use std::sync::Arc;
+use tokio_util::sync::CancellationToken;
 
 use crate::exec::ExecParams;
 use crate::exec_policy::ExecApprovalRequest;
@@ -44,6 +45,7 @@ fn shell_command_payload_command(payload: &ToolPayload) -> Option<String> {
 struct RunExecLikeArgs {
     tool_name: ToolName,
     exec_params: ExecParams,
+    cancellation_token: CancellationToken,
     hook_command: String,
     shell_type: Option<ShellType>,
     additional_permissions: Option<AdditionalPermissionProfile>,
@@ -59,6 +61,7 @@ async fn run_exec_like(args: RunExecLikeArgs) -> Result<FunctionToolOutput, Func
     let RunExecLikeArgs {
         tool_name,
         exec_params,
+        cancellation_token,
         hook_command,
         shell_type,
         additional_permissions,
@@ -183,6 +186,7 @@ async fn run_exec_like(args: RunExecLikeArgs) -> Result<FunctionToolOutput, Func
         hook_command,
         cwd: exec_params.cwd.clone(),
         timeout_ms: exec_params.expiration.timeout_ms(),
+        cancellation_token,
         env: exec_params.env.clone(),
         explicit_env_overrides,
         network: exec_params.network.clone(),

codex-rs/core/src/tools/handlers/shell/shell_command.rs

@@ -152,6 +152,7 @@ impl ToolExecutor<ToolInvocation> for ShellCommandHandler {
         let ToolInvocation {
             session,
             turn,
+            cancellation_token,
             tracker,
             call_id,
             payload,
@@ -189,6 +190,7 @@ impl ToolExecutor<ToolInvocation> for ShellCommandHandler {
         run_exec_like(RunExecLikeArgs {
             tool_name,
             exec_params,
+            cancellation_token,
             hook_command: params.command,
             shell_type,
             additional_permissions: params.additional_permissions.clone(),
@@ -209,6 +211,10 @@ impl CoreToolRuntime for ShellCommandHandler {
         matches!(payload, ToolPayload::Function { .. })
     }
 
+    fn waits_for_runtime_cancellation(&self) -> bool {
+        true
+    }
+
     fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
         shell_command_payload_command(&invocation.payload).map(|command| PreToolUsePayload {
             tool_name: HookToolName::bash(),

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

@@ -92,6 +92,7 @@ impl ToolCallRuntime {
         let tracker = Arc::clone(&self.tracker);
         let lock = Arc::clone(&self.parallel_execution);
         let invocation_cancellation_token = cancellation_token.clone();
+        let wait_for_runtime_cancellation = self.router.tool_waits_for_runtime_cancellation(&call);
         let started = Instant::now();
         let abort_session = Arc::clone(&session);
         let abort_source = source.clone();
@@ -140,23 +141,30 @@ impl ToolCallRuntime {
                     } else {
                         let secs = started.elapsed().as_secs_f32().max(0.1);
                         abort_dispatch_span.record("aborted", true);
-                        handle.abort();
-                        match handle.await {
-                            Ok(result) => result,
-                            Err(err) if err.is_cancelled() => {
-                                let response = Self::aborted_response(&call, secs);
-                                notify_tool_aborted(
-                                    abort_session.as_ref(),
-                                    abort_turn.as_ref(),
-                                    call.call_id.as_str(),
-                                    &call.tool_name,
-                                    abort_source,
-                                )
-                                .await;
-                                Ok(response)
+                        if wait_for_runtime_cancellation {
+                            match handle.await {
+                                Ok(_) => {}
+                                Err(err) if err.is_cancelled() => {}
+                                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)),
                             }
-                            Err(err) => Err(Self::tool_task_join_error(err)),
                         }
+                        let response = Self::aborted_response(&call, secs);
+                        notify_tool_aborted(
+                            abort_session.as_ref(),
+                            abort_turn.as_ref(),
+                            call.call_id.as_str(),
+                            &call.tool_name,
+                            abort_source,
+                        )
+                        .await;
+                        Ok(response)
                     }
                 },
             }

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

@@ -55,6 +55,10 @@ pub(crate) trait CoreToolRuntime: ToolExecutor<ToolInvocation> {
         )
     }
 
+    fn waits_for_runtime_cancellation(&self) -> bool {
+        false
+    }
+
     fn telemetry_tags<'a>(
         &'a self,
         _invocation: &'a ToolInvocation,
@@ -213,6 +217,10 @@ impl CoreToolRuntime for ExposureOverride {
         self.handler.matches_kind(payload)
     }
 
+    fn waits_for_runtime_cancellation(&self) -> bool {
+        self.handler.waits_for_runtime_cancellation()
+    }
+
     fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
         self.handler.pre_tool_use_payload(invocation)
     }
@@ -303,6 +311,11 @@ impl ToolRegistry {
         Some(tool.supports_parallel_tool_calls())
     }
 
+    pub(crate) fn waits_for_runtime_cancellation(&self, name: &ToolName) -> Option<bool> {
+        let tool = self.tool(name)?;
+        Some(tool.waits_for_runtime_cancellation())
+    }
+
     #[allow(dead_code)]
     pub(crate) async fn dispatch_any(
         &self,

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

@@ -74,6 +74,12 @@ impl ToolRouter {
             .unwrap_or(false)
     }
 
+    pub fn tool_waits_for_runtime_cancellation(&self, call: &ToolCall) -> bool {
+        self.registry
+            .waits_for_runtime_cancellation(&call.tool_name)
+            .unwrap_or(false)
+    }
+
     #[instrument(level = "trace", skip_all, err)]
     pub fn build_tool_call(item: ResponseItem) -> Result<Option<ToolCall>, FunctionCallError> {
         match item {

codex-rs/core/src/tools/runtimes/shell.rs

@@ -44,6 +44,7 @@ use codex_shell_command::powershell::prefix_powershell_script_with_utf8;
 use codex_utils_absolute_path::AbsolutePathBuf;
 use futures::future::BoxFuture;
 use std::collections::HashMap;
+use tokio_util::sync::CancellationToken;
 
 #[derive(Clone, Debug)]
 pub struct ShellRequest {
@@ -52,6 +53,7 @@ pub struct ShellRequest {
     pub hook_command: String,
     pub cwd: AbsolutePathBuf,
     pub timeout_ms: Option<u64>,
+    pub cancellation_token: CancellationToken,
     pub env: HashMap<String, String>,
     pub explicit_env_overrides: HashMap<String, String>,
     pub network: Option<NetworkProxy>,
@@ -265,6 +267,7 @@ impl ToolRuntime<ShellRequest, ExecToolCallOutput> for ShellRuntime {
         let command =
             build_sandbox_command(&command, &req.cwd, &env, req.additional_permissions.clone())?;
         let mut expiration: crate::exec::ExecExpiration = req.timeout_ms.into();
+        expiration = expiration.with_cancellation(req.cancellation_token.clone());
         if let Some(cancellation) = attempt.network_denial_cancellation_token.clone() {
             expiration = expiration.with_cancellation(cancellation);
         }

codex-rs/utils/pty/src/process_group.rs

@@ -94,7 +94,7 @@ pub fn kill_process_group_by_pid(pid: u32) -> io::Result<()> {
     let pgid = unsafe { libc::getpgid(pid) };
     if pgid == -1 {
         let err = io::Error::last_os_error();
-        if err.kind() != ErrorKind::NotFound {
+        if err.kind() != ErrorKind::NotFound && err.raw_os_error() != Some(libc::ESRCH) {
             return Err(err);
         }
         return Ok(());
@@ -103,7 +103,7 @@ pub fn kill_process_group_by_pid(pid: u32) -> io::Result<()> {
     let result = unsafe { libc::killpg(pgid, libc::SIGKILL) };
     if result == -1 {
         let err = io::Error::last_os_error();
-        if err.kind() != ErrorKind::NotFound {
+        if err.kind() != ErrorKind::NotFound && err.raw_os_error() != Some(libc::ESRCH) {
             return Err(err);
         }
     }
@@ -129,7 +129,7 @@ pub fn terminate_process_group(process_group_id: u32) -> io::Result<bool> {
     let result = unsafe { libc::killpg(pgid, libc::SIGTERM) };
     if result == -1 {
         let err = io::Error::last_os_error();
-        if err.kind() == ErrorKind::NotFound {
+        if err.kind() == ErrorKind::NotFound || err.raw_os_error() == Some(libc::ESRCH) {
             return Ok(false);
         }
         return Err(err);
@@ -153,7 +153,7 @@ pub fn kill_process_group(process_group_id: u32) -> io::Result<()> {
     let result = unsafe { libc::killpg(pgid, libc::SIGKILL) };
     if result == -1 {
         let err = io::Error::last_os_error();
-        if err.kind() != ErrorKind::NotFound {
+        if err.kind() != ErrorKind::NotFound && err.raw_os_error() != Some(libc::ESRCH) {
             return Err(err);
         }
     }