feat: update process_exec_tool_call() to take a cancellation token

codex-rs/app-server/src/codex_message_processor.rs

@@ -1171,7 +1171,7 @@ impl CodexMessageProcessor {
         let exec_params = ExecParams {
             command: params.command,
             cwd,
-            timeout_ms,
+            expiration: timeout_ms.into(),
             env,
             with_escalated_permissions: None,
             justification: None,

codex-rs/core/src/codex.rs

@@ -3051,6 +3051,7 @@ mod tests {
         let session = Arc::new(session);
         let mut turn_context = Arc::new(turn_context_raw);
 
+        let timeout_ms = 1000;
         let params = ExecParams {
             command: if cfg!(windows) {
                 vec![
@@ -3066,7 +3067,7 @@ mod tests {
                 ]
             },
             cwd: turn_context.cwd.clone(),
-            timeout_ms: Some(1000),
+            expiration: timeout_ms.into(),
             env: HashMap::new(),
             with_escalated_permissions: Some(true),
             justification: Some("test".to_string()),
@@ -3075,7 +3076,12 @@ mod tests {
 
         let params2 = ExecParams {
             with_escalated_permissions: Some(false),
-            ..params.clone()
+            command: params.command.clone(),
+            cwd: params.cwd.clone(),
+            expiration: timeout_ms.into(),
+            env: HashMap::new(),
+            justification: params.justification.clone(),
+            arg0: None,
         };
 
         let turn_diff_tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
@@ -3095,7 +3101,7 @@ mod tests {
                     arguments: serde_json::json!({
                         "command": params.command.clone(),
                         "workdir": Some(turn_context.cwd.to_string_lossy().to_string()),
-                        "timeout_ms": params.timeout_ms,
+                        "timeout_ms": params.expiration.timeout_ms(),
                         "with_escalated_permissions": params.with_escalated_permissions,
                         "justification": params.justification.clone(),
                     })
@@ -3132,7 +3138,7 @@ mod tests {
                     arguments: serde_json::json!({
                         "command": params2.command.clone(),
                         "workdir": Some(turn_context.cwd.to_string_lossy().to_string()),
-                        "timeout_ms": params2.timeout_ms,
+                        "timeout_ms": params2.expiration.timeout_ms(),
                         "with_escalated_permissions": params2.with_escalated_permissions,
                         "justification": params2.justification.clone(),
                     })

codex-rs/core/src/exec.rs

@@ -14,6 +14,7 @@ use tokio::io::AsyncRead;
 use tokio::io::AsyncReadExt;
 use tokio::io::BufReader;
 use tokio::process::Child;
+use tokio_util::sync::CancellationToken;
 
 use crate::error::CodexErr;
 use crate::error::Result;
@@ -47,20 +48,59 @@ const AGGREGATE_BUFFER_INITIAL_CAPACITY: usize = 8 * 1024; // 8 KiB
 /// Aggregation still collects full output; only the live event stream is capped.
 pub(crate) const MAX_EXEC_OUTPUT_DELTAS_PER_CALL: usize = 10_000;
 
-#[derive(Clone, Debug)]
+#[derive(Debug)]
 pub struct ExecParams {
     pub command: Vec<String>,
     pub cwd: PathBuf,
-    pub timeout_ms: Option<u64>,
+    pub expiration: ExecExpiration,
     pub env: HashMap<String, String>,
     pub with_escalated_permissions: Option<bool>,
     pub justification: Option<String>,
     pub arg0: Option<String>,
 }
 
-impl ExecParams {
-    pub fn timeout_duration(&self) -> Duration {
-        Duration::from_millis(self.timeout_ms.unwrap_or(DEFAULT_TIMEOUT_MS))
+/// Mechanism to terminate an exec invocation before it finishes naturally.
+#[derive(Debug)]
+pub enum ExecExpiration {
+    Timeout(Duration),
+    DefaultTimeout,
+    Cancellation(CancellationToken),
+}
+
+impl From<Option<u64>> for ExecExpiration {
+    fn from(timeout_ms: Option<u64>) -> Self {
+        timeout_ms.map_or(ExecExpiration::DefaultTimeout, |timeout_ms| {
+            ExecExpiration::Timeout(Duration::from_millis(timeout_ms))
+        })
+    }
+}
+
+impl From<u64> for ExecExpiration {
+    fn from(timeout_ms: u64) -> Self {
+        ExecExpiration::Timeout(Duration::from_millis(timeout_ms))
+    }
+}
+
+impl ExecExpiration {
+    async fn wait(self) {
+        match self {
+            ExecExpiration::Timeout(duration) => tokio::time::sleep(duration).await,
+            ExecExpiration::DefaultTimeout => {
+                tokio::time::sleep(Duration::from_millis(DEFAULT_TIMEOUT_MS)).await
+            }
+            ExecExpiration::Cancellation(cancel) => {
+                cancel.cancelled().await;
+            }
+        }
+    }
+
+    /// If ExecExpiration is a timeout, returns the timeout in milliseconds.
+    pub(crate) fn timeout_ms(&self) -> Option<u64> {
+        match self {
+            ExecExpiration::Timeout(duration) => Some(duration.as_millis() as u64),
+            ExecExpiration::DefaultTimeout => Some(DEFAULT_TIMEOUT_MS),
+            ExecExpiration::Cancellation(_) => None,
+        }
     }
 }
 
@@ -96,7 +136,7 @@ pub async fn process_exec_tool_call(
     let ExecParams {
         command,
         cwd,
-        timeout_ms,
+        expiration,
         env,
         with_escalated_permissions,
         justification,
@@ -115,15 +155,15 @@ pub async fn process_exec_tool_call(
         args: args.to_vec(),
         cwd,
         env,
-        timeout_ms,
+        expiration,
         with_escalated_permissions,
         justification,
     };
 
     let manager = SandboxManager::new();
     let exec_env = manager
         .transform(
-            &spec,
+            spec,
             sandbox_policy,
             sandbox_type,
             sandbox_cwd,
@@ -132,7 +172,7 @@ pub async fn process_exec_tool_call(
         .map_err(CodexErr::from)?;
 
     // Route through the sandboxing module for a single, unified execution path.
-    crate::sandboxing::execute_env(&exec_env, sandbox_policy, stdout_stream).await
+    crate::sandboxing::execute_env(exec_env, sandbox_policy, stdout_stream).await
 }
 
 pub(crate) async fn execute_exec_env(
@@ -144,7 +184,7 @@ pub(crate) async fn execute_exec_env(
         command,
         cwd,
         env,
-        timeout_ms,
+        expiration,
         sandbox,
         with_escalated_permissions,
         justification,
@@ -154,7 +194,7 @@ pub(crate) async fn execute_exec_env(
     let params = ExecParams {
         command,
         cwd,
-        timeout_ms,
+        expiration,
         env,
         with_escalated_permissions,
         justification,
@@ -179,9 +219,12 @@ async fn exec_windows_sandbox(
         command,
         cwd,
         env,
-        timeout_ms,
+        expiration,
         ..
     } = params;
+    // TODO(iceweasel-oai): run_windows_sandbox_capture should support all
+    // variants of ExecExpiration, not just timeout.
+    let timeout_ms = expiration.timeout_ms();
 
     let policy_str = serde_json::to_string(sandbox_policy).map_err(|err| {
         CodexErr::Io(io::Error::other(format!(
@@ -449,12 +492,12 @@ async fn exec(
     {
         return exec_windows_sandbox(params, sandbox_policy).await;
     }
-    let timeout = params.timeout_duration();
     let ExecParams {
         command,
         cwd,
         env,
         arg0,
+        expiration,
         ..
     } = params;
 
@@ -475,14 +518,14 @@ async fn exec(
         env,
     )
     .await?;
-    consume_truncated_output(child, timeout, stdout_stream).await
+    consume_truncated_output(child, expiration, stdout_stream).await
 }
 
 /// Consumes the output of a child process, truncating it so it is suitable for
 /// use as the output of a `shell` tool call. Also enforces specified timeout.
 async fn consume_truncated_output(
     mut child: Child,
-    timeout: Duration,
+    expiration: ExecExpiration,
     stdout_stream: Option<StdoutStream>,
 ) -> Result<RawExecToolCallOutput> {
     // Both stdout and stderr were configured with `Stdio::piped()`
@@ -516,20 +559,14 @@ async fn consume_truncated_output(
     ));
 
     let (exit_status, timed_out) = tokio::select! {
-        result = tokio::time::timeout(timeout, child.wait()) => {
-            match result {
-                Ok(status_result) => {
-                    let exit_status = status_result?;
-                    (exit_status, false)
-                }
-                Err(_) => {
-                    // timeout
-                    kill_child_process_group(&mut child)?;
-                    child.start_kill()?;
-                    // Debatable whether `child.wait().await` should be called here.
-                    (synthetic_exit_status(EXIT_CODE_SIGNAL_BASE + TIMEOUT_CODE), true)
-                }
-            }
+        status_result = child.wait() => {
+            let exit_status = status_result?;
+            (exit_status, false)
+        }
+        _ = expiration.wait() => {
+            kill_child_process_group(&mut child)?;
+            child.start_kill()?;
+            (synthetic_exit_status(EXIT_CODE_SIGNAL_BASE + TIMEOUT_CODE), true)
         }
         _ = tokio::signal::ctrl_c() => {
             kill_child_process_group(&mut child)?;
@@ -790,7 +827,7 @@ mod tests {
         let params = ExecParams {
             command,
             cwd: std::env::current_dir()?,
-            timeout_ms: Some(500),
+            expiration: 500.into(),
             env,
             with_escalated_permissions: None,
             justification: None,
@@ -824,4 +861,62 @@ mod tests {
         assert!(killed, "grandchild process with pid {pid} is still alive");
         Ok(())
     }
+
+    #[tokio::test]
+    async fn process_exec_tool_call_respects_cancellation_token() -> Result<()> {
+        let command = long_running_command();
+        let cwd = std::env::current_dir()?;
+        let env: HashMap<String, String> = std::env::vars().collect();
+        let cancel_token = CancellationToken::new();
+        let cancel_tx = cancel_token.clone();
+        let params = ExecParams {
+            command,
+            cwd: cwd.clone(),
+            expiration: ExecExpiration::Cancellation(cancel_token),
+            env,
+            with_escalated_permissions: None,
+            justification: None,
+            arg0: None,
+        };
+        tokio::spawn(async move {
+            tokio::time::sleep(Duration::from_millis(1_000)).await;
+            cancel_tx.cancel();
+        });
+        let result = process_exec_tool_call(
+            params,
+            SandboxType::None,
+            &SandboxPolicy::DangerFullAccess,
+            cwd.as_path(),
+            &None,
+            None,
+        )
+        .await;
+        let output = match result {
+            Err(CodexErr::Sandbox(SandboxErr::Timeout { output })) => output,
+            other => panic!("expected timeout error, got {other:?}"),
+        };
+        assert!(output.timed_out);
+        assert_eq!(output.exit_code, EXEC_TIMEOUT_EXIT_CODE);
+        Ok(())
+    }
+
+    #[cfg(unix)]
+    fn long_running_command() -> Vec<String> {
+        vec![
+            "/bin/sh".to_string(),
+            "-c".to_string(),
+            "sleep 30".to_string(),
+        ]
+    }
+
+    #[cfg(windows)]
+    fn long_running_command() -> Vec<String> {
+        vec![
+            "powershell.exe".to_string(),
+            "-NonInteractive".to_string(),
+            "-NoLogo".to_string(),
+            "-Command".to_string(),
+            "Start-Sleep -Seconds 30".to_string(),
+        ]
+    }
 }