From 229a73bb92086f83ae6be2e264baeb7e576e12ec Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Wed, 19 Nov 2025 22:43:05 -0800 Subject: [PATCH] feat: update process_exec_tool_call() to take a cancellation token --- .../app-server/src/codex_message_processor.rs | 1 + codex-rs/core/src/exec.rs | 136 ++++++++++++++++-- codex-rs/core/src/sandboxing/mod.rs | 4 +- codex-rs/core/src/tasks/user_shell.rs | 2 +- .../core/src/tools/runtimes/apply_patch.rs | 2 +- codex-rs/core/src/tools/runtimes/shell.rs | 2 +- codex-rs/core/tests/suite/exec.rs | 2 +- .../exec-server/src/posix/escalate_server.rs | 1 + .../linux-sandbox/tests/suite/landlock.rs | 2 + 9 files changed, 137 insertions(+), 15 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index bf9f0b9403..bea967e2ba 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -1201,6 +1201,7 @@ impl CodexMessageProcessor { sandbox_cwd.as_path(), &codex_linux_sandbox_exe, None, + None, ) .await { diff --git a/codex-rs/core/src/exec.rs b/codex-rs/core/src/exec.rs index 583c2233a2..ffa625e5e1 100644 --- a/codex-rs/core/src/exec.rs +++ b/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::sync::oneshot; use crate::error::CodexErr; use crate::error::Result; @@ -91,6 +92,7 @@ pub async fn process_exec_tool_call( sandbox_cwd: &Path, codex_linux_sandbox_exe: &Option, stdout_stream: Option, + cancel_rx: Option>, ) -> Result { let ExecParams { command, @@ -131,13 +133,14 @@ 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, cancel_rx).await } pub(crate) async fn execute_exec_env( env: ExecEnv, sandbox_policy: &SandboxPolicy, stdout_stream: Option, + cancel_rx: Option>, ) -> Result { let ExecEnv { command, @@ -161,7 +164,7 @@ pub(crate) async fn execute_exec_env( }; let start = Instant::now(); - let raw_output_result = exec(params, sandbox, sandbox_policy, stdout_stream).await; + let raw_output_result = exec(params, sandbox, sandbox_policy, stdout_stream, cancel_rx).await; let duration = start.elapsed(); finalize_exec_result(raw_output_result, sandbox, duration) } @@ -170,7 +173,9 @@ pub(crate) async fn execute_exec_env( async fn exec_windows_sandbox( params: ExecParams, sandbox_policy: &SandboxPolicy, + cancel_rx: Option>, ) -> Result { + let _ = cancel_rx; use crate::config::find_codex_home; use codex_windows_sandbox::run_windows_sandbox_capture; @@ -441,14 +446,16 @@ async fn exec( sandbox: SandboxType, sandbox_policy: &SandboxPolicy, stdout_stream: Option, + cancel_rx: Option>, ) -> Result { #[cfg(target_os = "windows")] if sandbox == SandboxType::WindowsRestrictedToken && !matches!(sandbox_policy, SandboxPolicy::DangerFullAccess) { - return exec_windows_sandbox(params, sandbox_policy).await; + return exec_windows_sandbox(params, sandbox_policy, cancel_rx).await; } let timeout = params.timeout_duration(); + let mut cancel_rx = cancel_rx; let ExecParams { command, cwd, @@ -464,7 +471,7 @@ async fn exec( )) })?; let arg0_ref = arg0.as_deref(); - let child = spawn_child_async( + let spawn_future = spawn_child_async( PathBuf::from(program), args.into(), arg0_ref, @@ -472,9 +479,20 @@ async fn exec( sandbox_policy, StdioPolicy::RedirectForShellTool, env, - ) - .await?; - consume_truncated_output(child, timeout, stdout_stream).await + ); + let child = { + let cancel_wait = cancel_future(cancel_rx.as_mut()); + tokio::pin!(cancel_wait); + tokio::select! { + result = spawn_future => { + result? + } + _ = &mut cancel_wait => { + return Ok(synthetic_timeout_output()); + } + } + }; + consume_truncated_output(child, timeout, cancel_rx, stdout_stream).await } /// Consumes the output of a child process, truncating it so it is suitable for @@ -482,6 +500,7 @@ async fn exec( async fn consume_truncated_output( mut child: Child, timeout: Duration, + cancel_rx: Option>, stdout_stream: Option, ) -> Result { // Both stdout and stderr were configured with `Stdio::piped()` @@ -514,6 +533,9 @@ async fn consume_truncated_output( Some(agg_tx.clone()), )); + let mut cancel_rx = cancel_rx; + let cancel_wait = cancel_future(cancel_rx.as_mut()); + tokio::pin!(cancel_wait); let (exit_status, timed_out) = tokio::select! { result = tokio::time::timeout(timeout, child.wait()) => { match result { @@ -522,10 +544,8 @@ async fn consume_truncated_output( (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) } } @@ -535,6 +555,11 @@ async fn consume_truncated_output( child.start_kill()?; (synthetic_exit_status(EXIT_CODE_SIGNAL_BASE + SIGKILL_CODE), false) } + _ = &mut cancel_wait => { + kill_child_process_group(&mut child)?; + child.start_kill()?; + (synthetic_exit_status(EXIT_CODE_SIGNAL_BASE + TIMEOUT_CODE), true) + } }; // Wait for the stdout/stderr collection tasks but guard against them @@ -604,6 +629,34 @@ async fn consume_truncated_output( }) } +fn synthetic_timeout_output() -> RawExecToolCallOutput { + RawExecToolCallOutput { + exit_status: synthetic_exit_status(EXIT_CODE_SIGNAL_BASE + TIMEOUT_CODE), + stdout: StreamOutput { + text: Vec::new(), + truncated_after_lines: None, + }, + stderr: StreamOutput { + text: Vec::new(), + truncated_after_lines: None, + }, + aggregated_output: StreamOutput { + text: Vec::new(), + truncated_after_lines: None, + }, + timed_out: true, + } +} + +async fn cancel_future( + cancel: Option<&mut oneshot::Receiver<()>>, +) -> std::result::Result<(), oneshot::error::RecvError> { + match cancel { + Some(cancel) => cancel.await, + None => std::future::pending::>().await, + } +} + async fn read_capped( mut reader: R, stream: Option, @@ -796,7 +849,14 @@ mod tests { arg0: None, }; - let output = exec(params, SandboxType::None, &SandboxPolicy::ReadOnly, None).await?; + let output = exec( + params, + SandboxType::None, + &SandboxPolicy::ReadOnly, + None, + None, + ) + .await?; assert!(output.timed_out); let stdout = output.stdout.from_utf8_lossy().text; @@ -823,4 +883,60 @@ 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 = std::env::vars().collect(); + let params = ExecParams { + command, + cwd: cwd.clone(), + timeout_ms: Some(30_000), + env, + with_escalated_permissions: None, + justification: None, + arg0: None, + }; + let (cancel_tx, cancel_rx) = oneshot::channel(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(1_000)).await; + let _ = cancel_tx.send(()); + }); + let result = process_exec_tool_call( + params, + SandboxType::None, + &SandboxPolicy::DangerFullAccess, + cwd.as_path(), + &None, + None, + Some(cancel_rx), + ) + .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 { + vec![ + "/bin/sh".to_string(), + "-c".to_string(), + "sleep 30".to_string(), + ] + } + + #[cfg(windows)] + fn long_running_command() -> Vec { + vec![ + "cmd.exe".to_string(), + "/C".to_string(), + "timeout /T 30 /NOBREAK > NUL".to_string(), + ] + } } diff --git a/codex-rs/core/src/sandboxing/mod.rs b/codex-rs/core/src/sandboxing/mod.rs index 4ecb2a8c12..26370e6bd6 100644 --- a/codex-rs/core/src/sandboxing/mod.rs +++ b/codex-rs/core/src/sandboxing/mod.rs @@ -25,6 +25,7 @@ use crate::tools::sandboxing::SandboxablePreference; use std::collections::HashMap; use std::path::Path; use std::path::PathBuf; +use tokio::sync::oneshot; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum SandboxPermissions { @@ -195,6 +196,7 @@ pub async fn execute_env( env: &ExecEnv, policy: &SandboxPolicy, stdout_stream: Option, + cancel_rx: Option>, ) -> crate::error::Result { - execute_exec_env(env.clone(), policy, stdout_stream).await + execute_exec_env(env.clone(), policy, stdout_stream, cancel_rx).await } diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index 9644d3678e..df5c7ef706 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -107,7 +107,7 @@ impl SessionTask for UserShellCommandTask { }); let sandbox_policy = SandboxPolicy::DangerFullAccess; - let exec_result = execute_exec_env(exec_env, &sandbox_policy, stdout_stream) + let exec_result = execute_exec_env(exec_env, &sandbox_policy, stdout_stream, None) .or_cancel(&cancellation_token) .await; diff --git a/codex-rs/core/src/tools/runtimes/apply_patch.rs b/codex-rs/core/src/tools/runtimes/apply_patch.rs index 0cdddd5087..0115316012 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch.rs @@ -155,7 +155,7 @@ impl ToolRuntime for ApplyPatchRuntime { let env = attempt .env_for(&spec) .map_err(|err| ToolError::Codex(err.into()))?; - let out = execute_env(&env, attempt.policy, Self::stdout_stream(ctx)) + let out = execute_env(&env, attempt.policy, Self::stdout_stream(ctx), None) .await .map_err(ToolError::Codex)?; Ok(out) diff --git a/codex-rs/core/src/tools/runtimes/shell.rs b/codex-rs/core/src/tools/runtimes/shell.rs index d71c4498e6..461695cde8 100644 --- a/codex-rs/core/src/tools/runtimes/shell.rs +++ b/codex-rs/core/src/tools/runtimes/shell.rs @@ -140,7 +140,7 @@ impl ToolRuntime for ShellRuntime { let env = attempt .env_for(&spec) .map_err(|err| ToolError::Codex(err.into()))?; - let out = execute_env(&env, attempt.policy, Self::stdout_stream(ctx)) + let out = execute_env(&env, attempt.policy, Self::stdout_stream(ctx), None) .await .map_err(ToolError::Codex)?; Ok(out) diff --git a/codex-rs/core/tests/suite/exec.rs b/codex-rs/core/tests/suite/exec.rs index ea5ab84879..83b859fc2f 100644 --- a/codex-rs/core/tests/suite/exec.rs +++ b/codex-rs/core/tests/suite/exec.rs @@ -41,7 +41,7 @@ async fn run_test_cmd(tmp: TempDir, cmd: Vec<&str>) -> Result