From f18ba2858eb80861777262941dfd820461d69653 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Tue, 3 Mar 2026 12:49:02 -0800 Subject: [PATCH] refactor: prepare unified exec for zsh-fork backend --- codex-rs/core/src/tools/runtimes/shell.rs | 2 +- .../tools/runtimes/shell/unix_escalation.rs | 123 +++++++++++++++++- .../runtimes/shell/unix_escalation_tests.rs | 4 + .../core/src/tools/runtimes/unified_exec.rs | 58 ++++++++- codex-rs/core/src/tools/spec.rs | 18 +++ codex-rs/core/src/unified_exec/process.rs | 7 +- .../core/src/unified_exec/process_manager.rs | 7 +- codex-rs/shell-escalation/src/lib.rs | 2 + .../src/unix/escalate_server.rs | 101 +++++++++++--- codex-rs/shell-escalation/src/unix/mod.rs | 1 + .../shell-escalation/src/unix/stopwatch.rs | 32 ++++- 11 files changed, 322 insertions(+), 33 deletions(-) diff --git a/codex-rs/core/src/tools/runtimes/shell.rs b/codex-rs/core/src/tools/runtimes/shell.rs index 101d9309fd..c4d15273a6 100644 --- a/codex-rs/core/src/tools/runtimes/shell.rs +++ b/codex-rs/core/src/tools/runtimes/shell.rs @@ -5,7 +5,7 @@ Executes shell requests under the orchestrator: asks for approval when needed, builds a CommandSpec, and runs it under the current SandboxAttempt. */ #[cfg(unix)] -mod unix_escalation; +pub(crate) mod unix_escalation; use crate::command_canonicalization::canonicalize_command_for_approval; use crate::exec::ExecToolCallOutput; diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs index 36f63be4fa..4fc22209d9 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs @@ -6,6 +6,7 @@ use crate::exec::ExecToolCallOutput; use crate::exec::SandboxType; use crate::exec::is_likely_sandbox_denied; use crate::features::Feature; +use crate::sandboxing::ExecRequest; use crate::sandboxing::SandboxPermissions; use crate::shell::ShellType; use crate::skills::SkillMetadata; @@ -36,6 +37,7 @@ use codex_shell_escalation::EscalationDecision; use codex_shell_escalation::EscalationExecution; use codex_shell_escalation::EscalationPermissions; use codex_shell_escalation::EscalationPolicy; +use codex_shell_escalation::EscalationSession; use codex_shell_escalation::ExecParams; use codex_shell_escalation::ExecResult; use codex_shell_escalation::Permissions as EscalatedPermissions; @@ -51,6 +53,11 @@ use tokio::sync::RwLock; use tokio_util::sync::CancellationToken; use uuid::Uuid; +pub(crate) struct PreparedUnifiedExecZshFork { + pub(crate) exec_request: ExecRequest, + pub(crate) escalation_session: EscalationSession, +} + pub(super) async fn try_run_zsh_fork( req: &ShellRequest, attempt: &SandboxAttempt<'_>, @@ -95,7 +102,7 @@ pub(super) async fn try_run_zsh_fork( justification, arg0, } = sandbox_exec_request; - let ParsedShellCommand { script, login } = extract_shell_script(&command)?; + let ParsedShellCommand { script, login, .. } = extract_shell_script(&command)?; let effective_timeout = Duration::from_millis( req.timeout_ms .unwrap_or(crate::exec::DEFAULT_EXEC_COMMAND_TIMEOUT_MS), @@ -172,6 +179,103 @@ pub(super) async fn try_run_zsh_fork( map_exec_result(attempt.sandbox, exec_result).map(Some) } +pub(crate) async fn prepare_unified_exec_zsh_fork( + req: &crate::tools::runtimes::unified_exec::UnifiedExecRequest, + attempt: &SandboxAttempt<'_>, + ctx: &ToolCtx, + exec_request: ExecRequest, +) -> Result, ToolError> { + let Some(shell_zsh_path) = ctx.session.services.shell_zsh_path.as_ref() else { + tracing::warn!("ZshFork backend specified, but shell_zsh_path is not configured."); + return Ok(None); + }; + if !ctx.session.features().enabled(Feature::ShellZshFork) { + tracing::warn!("ZshFork backend specified, but ShellZshFork feature is not enabled."); + return Ok(None); + } + if !matches!(ctx.session.user_shell().shell_type, ShellType::Zsh) { + tracing::warn!("ZshFork backend specified, but user shell is not Zsh."); + return Ok(None); + } + + let parsed = match extract_shell_script(&exec_request.command) { + Ok(parsed) => parsed, + Err(err) => { + tracing::warn!("ZshFork unified exec fallback: {err:?}"); + return Ok(None); + } + }; + if parsed.program != shell_zsh_path.to_string_lossy() { + tracing::warn!( + "ZshFork backend specified, but unified exec command targets `{}` instead of `{}`.", + parsed.program, + shell_zsh_path.display(), + ); + return Ok(None); + } + + let exec_policy = Arc::new(RwLock::new( + ctx.session.services.exec_policy.current().as_ref().clone(), + )); + let command_executor = CoreShellCommandExecutor { + command: exec_request.command.clone(), + cwd: exec_request.cwd.clone(), + sandbox_policy: exec_request.sandbox_policy.clone(), + sandbox: exec_request.sandbox, + env: exec_request.env.clone(), + network: exec_request.network.clone(), + windows_sandbox_level: exec_request.windows_sandbox_level, + sandbox_permissions: exec_request.sandbox_permissions, + justification: exec_request.justification.clone(), + arg0: exec_request.arg0.clone(), + sandbox_policy_cwd: ctx.turn.cwd.clone(), + macos_seatbelt_profile_extensions: ctx + .turn + .config + .permissions + .macos_seatbelt_profile_extensions + .clone(), + codex_linux_sandbox_exe: ctx.turn.codex_linux_sandbox_exe.clone(), + use_linux_sandbox_bwrap: ctx.turn.features.enabled(Feature::UseLinuxSandboxBwrap), + }; + let main_execve_wrapper_exe = ctx + .session + .services + .main_execve_wrapper_exe + .clone() + .ok_or_else(|| { + ToolError::Rejected( + "zsh fork feature enabled, but execve wrapper is not configured".to_string(), + ) + })?; + let escalation_policy = CoreShellActionProvider { + policy: Arc::clone(&exec_policy), + session: Arc::clone(&ctx.session), + turn: Arc::clone(&ctx.turn), + call_id: ctx.call_id.clone(), + approval_policy: ctx.turn.approval_policy.value(), + sandbox_policy: attempt.policy.clone(), + sandbox_permissions: req.sandbox_permissions, + prompt_permissions: req.additional_permissions.clone(), + stopwatch: Stopwatch::unlimited(), + }; + + let escalate_server = EscalateServer::new( + shell_zsh_path.clone(), + main_execve_wrapper_exe, + escalation_policy, + ); + let escalation_session = escalate_server + .start_session(Arc::new(command_executor)) + .map_err(|err| ToolError::Rejected(err.to_string()))?; + let mut exec_request = exec_request; + exec_request.env.extend(escalation_session.env().clone()); + Ok(Some(PreparedUnifiedExecZshFork { + exec_request, + escalation_session, + })) +} + struct CoreShellActionProvider { policy: Arc>, session: Arc, @@ -809,6 +913,7 @@ impl CoreShellCommandExecutor { #[derive(Debug, Eq, PartialEq)] struct ParsedShellCommand { + program: String, script: String, login: bool, } @@ -817,12 +922,20 @@ fn extract_shell_script(command: &[String]) -> Result Some((script.to_owned(), false)), - [_, flag, script] if flag == "-lc" => Some((script.to_owned(), true)), + if let Some((program, script, login)) = command.windows(3).find_map(|parts| match parts { + [program, flag, script] if flag == "-c" => { + Some((program.to_owned(), script.to_owned(), false)) + } + [program, flag, script] if flag == "-lc" => { + Some((program.to_owned(), script.to_owned(), true)) + } _ => None, }) { - return Ok(ParsedShellCommand { script, login }); + return Ok(ParsedShellCommand { + program, + script, + login, + }); } Err(ToolError::Rejected( diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs index b88f5b0542..ad663a3fed 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs @@ -64,6 +64,7 @@ fn extract_shell_script_preserves_login_flag() { assert_eq!( extract_shell_script(&["/bin/zsh".into(), "-lc".into(), "echo hi".into()]).unwrap(), ParsedShellCommand { + program: "/bin/zsh".to_string(), script: "echo hi".to_string(), login: true, } @@ -71,6 +72,7 @@ fn extract_shell_script_preserves_login_flag() { assert_eq!( extract_shell_script(&["/bin/zsh".into(), "-c".into(), "echo hi".into()]).unwrap(), ParsedShellCommand { + program: "/bin/zsh".to_string(), script: "echo hi".to_string(), login: false, } @@ -89,6 +91,7 @@ fn extract_shell_script_supports_wrapped_command_prefixes() { ]) .unwrap(), ParsedShellCommand { + program: "/bin/zsh".to_string(), script: "echo hello".to_string(), login: true, } @@ -105,6 +108,7 @@ fn extract_shell_script_supports_wrapped_command_prefixes() { ]) .unwrap(), ParsedShellCommand { + program: "/bin/zsh".to_string(), script: "pwd".to_string(), login: false, } diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index f4d823bf0e..4ae2276d8c 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -16,6 +16,8 @@ use crate::tools::network_approval::NetworkApprovalMode; use crate::tools::network_approval::NetworkApprovalSpec; use crate::tools::runtimes::build_command_spec; use crate::tools::runtimes::maybe_wrap_shell_lc_with_snapshot; +#[cfg(unix)] +use crate::tools::runtimes::shell::unix_escalation; use crate::tools::sandboxing::Approvable; use crate::tools::sandboxing::ApprovalCtx; use crate::tools::sandboxing::ExecApprovalRequirement; @@ -28,6 +30,7 @@ use crate::tools::sandboxing::ToolError; use crate::tools::sandboxing::ToolRuntime; use crate::tools::sandboxing::sandbox_override_for_first_attempt; use crate::tools::sandboxing::with_cached_approval; +use crate::tools::spec::UnifiedExecBackendConfig; use crate::unified_exec::UnifiedExecError; use crate::unified_exec::UnifiedExecProcess; use crate::unified_exec::UnifiedExecProcessManager; @@ -63,11 +66,12 @@ pub struct UnifiedExecApprovalKey { pub struct UnifiedExecRuntime<'a> { manager: &'a UnifiedExecProcessManager, + backend: UnifiedExecBackendConfig, } impl<'a> UnifiedExecRuntime<'a> { - pub fn new(manager: &'a UnifiedExecProcessManager) -> Self { - Self { manager } + pub fn new(manager: &'a UnifiedExecProcessManager, backend: UnifiedExecBackendConfig) -> Self { + Self { manager, backend } } } @@ -194,11 +198,59 @@ impl<'a> ToolRuntime for UnifiedExecRunt req.justification.clone(), ) .map_err(|_| ToolError::Rejected("missing command line for PTY".to_string()))?; + #[cfg(unix)] + if self.backend == UnifiedExecBackendConfig::ZshFork { + let exec_env = attempt + .env_for(spec, req.network.as_ref()) + .map_err(|err| ToolError::Codex(err.into()))?; + match unix_escalation::prepare_unified_exec_zsh_fork(req, attempt, ctx, exec_env) + .await? + { + Some(prepared) => { + let unix_escalation::PreparedUnifiedExecZshFork { + exec_request, + escalation_session, + } = prepared; + return self + .manager + .open_session_with_exec_env( + &exec_request, + req.tty, + Some(escalation_session), + ) + .await + .map_err(|err| match err { + UnifiedExecError::SandboxDenied { output, .. } => { + ToolError::Codex(CodexErr::Sandbox(SandboxErr::Denied { + output: Box::new(output), + network_policy_decision: None, + })) + } + other => ToolError::Rejected(other.to_string()), + }); + } + None => { + tracing::warn!( + "UnifiedExec ZshFork backend specified, but conditions for using it were not met, falling back to direct execution", + ); + } + } + } + let spec = build_command_spec( + &command, + &req.cwd, + &env, + ExecExpiration::DefaultTimeout, + req.sandbox_permissions, + req.additional_permissions.clone(), + req.justification.clone(), + ) + .map_err(|_| ToolError::Rejected("missing command line for PTY".to_string()))?; let exec_env = attempt .env_for(spec, req.network.as_ref()) .map_err(|err| ToolError::Codex(err.into()))?; self.manager - .open_session_with_exec_env(&exec_env, req.tty) + .open_session_with_exec_env(&exec_env, req.tty, None) .await .map_err(|err| match err { UnifiedExecError::SandboxDenied { output, .. } => { diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 0e53670750..2e1dacfa9a 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -42,10 +42,17 @@ pub enum ShellCommandBackendConfig { ZshFork, } +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub enum UnifiedExecBackendConfig { + Direct, + ZshFork, +} + #[derive(Debug, Clone)] pub(crate) struct ToolsConfig { pub shell_type: ConfigShellToolType, shell_command_backend: ShellCommandBackendConfig, + pub unified_exec_backend: UnifiedExecBackendConfig, pub allow_login_shell: bool, pub apply_patch_tool_type: Option, pub web_search_mode: Option, @@ -94,6 +101,12 @@ impl ToolsConfig { } else { ShellCommandBackendConfig::Classic }; + let unified_exec_backend = + if features.enabled(Feature::ShellTool) && features.enabled(Feature::ShellZshFork) { + UnifiedExecBackendConfig::ZshFork + } else { + UnifiedExecBackendConfig::Direct + }; let shell_type = if !features.enabled(Feature::ShellTool) { ConfigShellToolType::Disabled @@ -132,6 +145,7 @@ impl ToolsConfig { Self { shell_type, shell_command_backend, + unified_exec_backend, allow_login_shell: true, apply_patch_tool_type, web_search_mode: *web_search_mode, @@ -2727,6 +2741,10 @@ mod tests { tools_config.shell_command_backend, ShellCommandBackendConfig::ZshFork ); + assert_eq!( + tools_config.unified_exec_backend, + UnifiedExecBackendConfig::ZshFork + ); } #[test] diff --git a/codex-rs/core/src/unified_exec/process.rs b/codex-rs/core/src/unified_exec/process.rs index 69790dbad2..a2728c390b 100644 --- a/codex-rs/core/src/unified_exec/process.rs +++ b/codex-rs/core/src/unified_exec/process.rs @@ -17,6 +17,7 @@ use crate::exec::StreamOutput; use crate::exec::is_likely_sandbox_denied; use crate::truncate::TruncationPolicy; use crate::truncate::formatted_truncate_text; +use codex_shell_escalation::EscalationSession; use codex_utils_pty::ExecCommandSession; use codex_utils_pty::SpawnedPty; @@ -44,6 +45,7 @@ pub(crate) struct UnifiedExecProcess { output_drained: Arc, output_task: JoinHandle<()>, sandbox_type: SandboxType, + _escalation_session: Option, } impl UnifiedExecProcess { @@ -51,6 +53,7 @@ impl UnifiedExecProcess { process_handle: ExecCommandSession, initial_output_rx: tokio::sync::broadcast::Receiver>, sandbox_type: SandboxType, + escalation_session: Option, ) -> Self { let output_buffer = Arc::new(Mutex::new(HeadTailBuffer::default())); let output_notify = Arc::new(Notify::new()); @@ -92,6 +95,7 @@ impl UnifiedExecProcess { output_drained, output_task, sandbox_type, + _escalation_session: escalation_session, } } @@ -196,13 +200,14 @@ impl UnifiedExecProcess { pub(super) async fn from_spawned( spawned: SpawnedPty, sandbox_type: SandboxType, + escalation_session: Option, ) -> Result { let SpawnedPty { session: process_handle, output_rx, mut exit_rx, } = spawned; - let managed = Self::new(process_handle, output_rx, sandbox_type); + let managed = Self::new(process_handle, output_rx, sandbox_type, escalation_session); let exit_ready = matches!(exit_rx.try_recv(), Ok(_) | Err(TryRecvError::Closed)); diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 821abc683e..54eb9fefdf 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -51,6 +51,7 @@ use crate::unified_exec::process::OutputBuffer; use crate::unified_exec::process::OutputHandles; use crate::unified_exec::process::UnifiedExecProcess; use crate::unified_exec::resolve_max_tokens; +use codex_shell_escalation::EscalationSession; const UNIFIED_EXEC_ENV: [(&str, &str); 10] = [ ("NO_COLOR", "1"), @@ -528,6 +529,7 @@ impl UnifiedExecProcessManager { &self, env: &ExecRequest, tty: bool, + escalation_session: Option, ) -> Result { let (program, args) = env .command @@ -555,7 +557,7 @@ impl UnifiedExecProcessManager { }; let spawned = spawn_result.map_err(|err| UnifiedExecError::create_process(err.to_string()))?; - UnifiedExecProcess::from_spawned(spawned, env.sandbox).await + UnifiedExecProcess::from_spawned(spawned, env.sandbox, escalation_session).await } pub(super) async fn open_session_with_sandbox( @@ -569,7 +571,8 @@ impl UnifiedExecProcessManager { Some(context.session.conversation_id), )); let mut orchestrator = ToolOrchestrator::new(); - let mut runtime = UnifiedExecRuntime::new(self); + let mut runtime = + UnifiedExecRuntime::new(self, context.turn.tools_config.unified_exec_backend); let exec_approval_requirement = context .session .services diff --git a/codex-rs/shell-escalation/src/lib.rs b/codex-rs/shell-escalation/src/lib.rs index 1cc42a46db..6d08a78716 100644 --- a/codex-rs/shell-escalation/src/lib.rs +++ b/codex-rs/shell-escalation/src/lib.rs @@ -14,6 +14,8 @@ pub use unix::EscalationPermissions; #[cfg(unix)] pub use unix::EscalationPolicy; #[cfg(unix)] +pub use unix::EscalationSession; +#[cfg(unix)] pub use unix::ExecParams; #[cfg(unix)] pub use unix::ExecResult; diff --git a/codex-rs/shell-escalation/src/unix/escalate_server.rs b/codex-rs/shell-escalation/src/unix/escalate_server.rs index bb45c3ed71..75cda09ad3 100644 --- a/codex-rs/shell-escalation/src/unix/escalate_server.rs +++ b/codex-rs/shell-escalation/src/unix/escalate_server.rs @@ -7,7 +7,9 @@ use std::time::Duration; use anyhow::Context as _; use codex_utils_absolute_path::AbsolutePathBuf; +use socket2::Socket; use tokio::process::Command; +use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; use crate::unix::escalate_protocol::ESCALATE_SOCKET_ENV_VAR; @@ -82,6 +84,25 @@ pub struct PreparedExec { pub arg0: Option, } +#[derive(Debug)] +pub struct EscalationSession { + env: HashMap, + task: JoinHandle>, + _client_socket: Socket, +} + +impl EscalationSession { + pub fn env(&self) -> &HashMap { + &self.env + } +} + +impl Drop for EscalationSession { + fn drop(&mut self) { + self.task.abort(); + } +} + pub struct EscalateServer { bash_path: PathBuf, execve_wrapper: PathBuf, @@ -106,16 +127,42 @@ impl EscalateServer { cancel_rx: CancellationToken, command_executor: Arc, ) -> anyhow::Result { + let session = self.start_session(Arc::clone(&command_executor))?; + let command = vec![ + self.bash_path.to_string_lossy().to_string(), + if params.login == Some(false) { + "-c".to_string() + } else { + "-lc".to_string() + }, + params.command, + ]; + let workdir = AbsolutePathBuf::try_from(params.workdir)?; + let result = command_executor + .run( + command, + workdir.to_path_buf(), + session.env().clone(), + cancel_rx, + ) + .await?; + Ok(result) + } + + pub fn start_session( + &self, + command_executor: Arc, + ) -> anyhow::Result { let (escalate_server, escalate_client) = AsyncDatagramSocket::pair()?; let client_socket = escalate_client.into_inner(); // Only the client endpoint should cross exec into the wrapper process. client_socket.set_cloexec(false)?; - let escalate_task = tokio::spawn(escalate_task( + let task = tokio::spawn(escalate_task( escalate_server, Arc::clone(&self.policy), Arc::clone(&command_executor), )); - let mut env = std::env::vars().collect::>(); + let mut env = HashMap::new(); env.insert( ESCALATE_SOCKET_ENV_VAR.to_string(), client_socket.as_raw_fd().to_string(), @@ -128,22 +175,11 @@ impl EscalateServer { LEGACY_BASH_EXEC_WRAPPER_ENV_VAR.to_string(), self.execve_wrapper.to_string_lossy().to_string(), ); - - let command = vec![ - self.bash_path.to_string_lossy().to_string(), - if params.login == Some(false) { - "-c".to_string() - } else { - "-lc".to_string() - }, - params.command, - ]; - let workdir = AbsolutePathBuf::try_from(params.workdir)?; - let result = command_executor - .run(command, workdir.to_path_buf(), env, cancel_rx) - .await?; - escalate_task.abort(); - Ok(result) + Ok(EscalationSession { + env, + task, + _client_socket: client_socket, + }) } } @@ -390,6 +426,35 @@ mod tests { } } + #[tokio::test] + async fn start_session_exposes_wrapper_env_overlay() -> anyhow::Result<()> { + let execve_wrapper = PathBuf::from("/tmp/codex-execve-wrapper"); + let execve_wrapper_str = execve_wrapper.to_string_lossy().to_string(); + let server = EscalateServer::new( + PathBuf::from("/bin/bash"), + execve_wrapper.clone(), + DeterministicEscalationPolicy { + decision: EscalationDecision::run(), + }, + ); + + let session = server.start_session(Arc::new(ForwardingShellCommandExecutor))?; + let env = session.env(); + assert_eq!(env.get(EXEC_WRAPPER_ENV_VAR), Some(&execve_wrapper_str)); + assert_eq!( + env.get(LEGACY_BASH_EXEC_WRAPPER_ENV_VAR), + Some(&execve_wrapper_str) + ); + let socket_fd = env + .get(ESCALATE_SOCKET_ENV_VAR) + .expect("session should export shell escalation socket"); + let socket_fd = socket_fd.parse::()?; + assert!(socket_fd >= 0); + assert_ne!(unsafe { libc::fcntl(socket_fd, libc::F_GETFD) }, -1); + + Ok(()) + } + #[tokio::test] async fn handle_escalate_session_respects_run_in_sandbox_decision() -> anyhow::Result<()> { let (server, client) = AsyncSocket::pair()?; diff --git a/codex-rs/shell-escalation/src/unix/mod.rs b/codex-rs/shell-escalation/src/unix/mod.rs index 6de12297a4..13b7146ab4 100644 --- a/codex-rs/shell-escalation/src/unix/mod.rs +++ b/codex-rs/shell-escalation/src/unix/mod.rs @@ -66,6 +66,7 @@ pub use self::escalate_protocol::EscalateAction; pub use self::escalate_protocol::EscalationDecision; pub use self::escalate_protocol::EscalationExecution; pub use self::escalate_server::EscalateServer; +pub use self::escalate_server::EscalationSession; pub use self::escalate_server::ExecParams; pub use self::escalate_server::ExecResult; pub use self::escalate_server::PreparedExec; diff --git a/codex-rs/shell-escalation/src/unix/stopwatch.rs b/codex-rs/shell-escalation/src/unix/stopwatch.rs index f564434813..19821d8be2 100644 --- a/codex-rs/shell-escalation/src/unix/stopwatch.rs +++ b/codex-rs/shell-escalation/src/unix/stopwatch.rs @@ -9,7 +9,7 @@ use tokio_util::sync::CancellationToken; #[derive(Clone, Debug)] pub struct Stopwatch { - limit: Duration, + limit: Option, inner: Arc>, notify: Arc, } @@ -30,13 +30,27 @@ impl Stopwatch { active_pauses: 0, })), notify: Arc::new(Notify::new()), - limit, + limit: Some(limit), + } + } + + pub fn unlimited() -> Self { + Self { + inner: Arc::new(Mutex::new(StopwatchState { + elapsed: Duration::ZERO, + running_since: Some(Instant::now()), + active_pauses: 0, + })), + notify: Arc::new(Notify::new()), + limit: None, } } pub fn cancellation_token(&self) -> CancellationToken { - let limit = self.limit; let token = CancellationToken::new(); + let Some(limit) = self.limit else { + return token; + }; let cancel = token.clone(); let inner = Arc::clone(&self.inner); let notify = Arc::clone(&self.notify); @@ -208,4 +222,16 @@ mod tests { // Now the stopwatch should resume and hit the limit shortly after. token.cancelled().await; } + + #[tokio::test] + async fn unlimited_stopwatch_never_cancels() { + let stopwatch = Stopwatch::unlimited(); + let token = stopwatch.cancellation_token(); + + assert!( + timeout(Duration::from_millis(30), token.cancelled()) + .await + .is_err() + ); + } }