diff --git a/codex-rs/exec-server/src/shell_snapshot.rs b/codex-rs/exec-server/src/shell_snapshot.rs index 5967e94be5..e352d63641 100644 --- a/codex-rs/exec-server/src/shell_snapshot.rs +++ b/codex-rs/exec-server/src/shell_snapshot.rs @@ -16,6 +16,7 @@ use tokio::io::AsyncReadExt; use tokio::process::Command; use tokio::sync::Mutex; use tokio::sync::OnceCell; +use tokio::time::Instant; use crate::FileSystemSandboxContext; use crate::local_process::shell_environment_policy; @@ -31,6 +32,8 @@ const MAX_SNAPSHOT_BYTES: usize = 512 * 1024; const MAX_SNAPSHOT_ENV_VALUE_BYTES: usize = 60 * 1024; const MAX_SNAPSHOT_SCOPE_BYTES: usize = 256; const SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(10); +const SNAPSHOT_RETRY_BACKOFF: Duration = Duration::from_secs(1); +const MAX_SNAPSHOT_ATTEMPTS: usize = 3; #[derive(Default)] pub(crate) struct ShellSnapshotCache { @@ -42,7 +45,9 @@ struct CachedShellSnapshot { cwd: PathUri, env_policy: Option, sandbox: Option, - snapshot: Arc>>, + attempts: usize, + // Failed captures store the earliest time another attempt may start. + snapshot: Arc>>, } struct ShellSnapshot { @@ -93,7 +98,16 @@ impl ShellSnapshotCache { && entry.sandbox == params.sandbox }); let cached = position.and_then(|position| { - let entry = entries.remove(position)?; + let mut entry = entries.remove(position)?; + // Share each failed attempt during backoff. After the retry + // budget is exhausted, keep falling back until eviction. + if entry.attempts < MAX_SNAPSHOT_ATTEMPTS + && let Some(Err(retry_at)) = entry.snapshot.get() + && Instant::now() >= *retry_at + { + entry.attempts += 1; + entry.snapshot = Arc::new(OnceCell::new()); + } let snapshot = Arc::clone(&entry.snapshot); entries.push_back(entry); Some(snapshot) @@ -107,6 +121,7 @@ impl ShellSnapshotCache { cwd: params.cwd.clone(), env_policy: params.env_policy.clone(), sandbox: params.sandbox.clone(), + attempts: 1, snapshot: Arc::clone(&snapshot), }; entries.push_back(entry); @@ -117,15 +132,14 @@ impl ShellSnapshotCache { snapshot } }; - let Some(snapshot) = snapshot + let Ok(snapshot) = snapshot .get_or_init(|| async { - match capture_snapshot(params, prepared, shell_type).await { - Ok(snapshot) => Some(snapshot), - Err(err) => { + capture_snapshot(params, prepared, shell_type) + .await + .map_err(|err| { tracing::warn!("failed to capture shell snapshot: {err:?}"); - None - } - } + Instant::now() + SNAPSHOT_RETRY_BACKOFF + }) }) .await else { diff --git a/codex-rs/exec-server/src/shell_snapshot_tests.rs b/codex-rs/exec-server/src/shell_snapshot_tests.rs index 8e6aeda771..1f24ef5de8 100644 --- a/codex-rs/exec-server/src/shell_snapshot_tests.rs +++ b/codex-rs/exec-server/src/shell_snapshot_tests.rs @@ -2,9 +2,120 @@ use std::collections::HashMap; use codex_protocol::config_types::ShellEnvironmentPolicyInherit; use pretty_assertions::assert_eq; +use test_case::test_case; +use super::MAX_SNAPSHOT_ATTEMPTS; +use super::SNAPSHOT_RETRY_BACKOFF; +use super::ShellSnapshotCache; use super::parse_snapshot; +use crate::process_sandbox::prepare_exec_request; use crate::protocol::ExecEnvPolicy; +use crate::protocol::ExecParams; +use crate::protocol::ProcessId; +use crate::protocol::ShellInfo; +use crate::protocol::ShellSnapshotRequest; + +#[test_case(2; "recovers_on_second_attempt")] +#[test_case(3; "recovers_on_last_attempt")] +#[test_case(4; "stops_after_three_failures")] +#[tokio::test] +async fn snapshot_failure_retries_are_bounded_and_single_flight( + recovery_attempt: usize, +) -> anyhow::Result<()> { + let home = tempfile::TempDir::new()?; + let profile = home.path().join(".bashrc"); + std::fs::write(&profile, "printf x >> \"$HOME/captures\"\nexit 7\n")?; + let params = ExecParams { + process_id: ProcessId::from("snapshot-retry"), + argv: vec![ + "/bin/bash".to_string(), + "-lc".to_string(), + "true".to_string(), + ], + cwd: codex_utils_path_uri::PathUri::from_host_native_path(home.path())?, + env: HashMap::from([ + ( + "HOME".to_string(), + home.path().to_string_lossy().into_owned(), + ), + ("PATH".to_string(), "/usr/bin:/bin".to_string()), + ]), + env_policy: None, + shell_snapshot: Some(ShellSnapshotRequest { + scope_id: "attachment-1".to_string(), + shell: ShellInfo { + name: "bash".to_string(), + path: "/bin/bash".to_string(), + }, + }), + tty: false, + pipe_stdin: false, + arg0: None, + sandbox: None, + enforce_managed_network: false, + managed_network: None, + network_proxy: None, + }; + let cache = ShellSnapshotCache::default(); + + for attempt in 1..=5 { + if attempt == recovery_attempt { + std::fs::write( + &profile, + "printf x >> \"$HOME/captures\"\nprofile_helper() { printf recovered; }\n", + )?; + } + let mut prepared = prepare_exec_request( + ¶ms, + params.env.clone(), + /*runtime_paths*/ None, + /*network_policy_decider*/ None, + /*network_policy_audit_observer*/ None, + ) + .await + .expect("prepare capture"); + let mut concurrent = prepare_exec_request( + ¶ms, + params.env.clone(), + /*runtime_paths*/ None, + /*network_policy_decider*/ None, + /*network_policy_audit_observer*/ None, + ) + .await + .expect("prepare concurrent capture"); + let (first, second) = tokio::join!( + cache.prepare(¶ms, &mut prepared), + cache.prepare(¶ms, &mut concurrent), + ); + first.expect("capture failure must preserve command fallback"); + second.expect("concurrent request must share the capture attempt"); + assert_eq!( + (&prepared.command, &prepared.env), + (&concurrent.command, &concurrent.env) + ); + + tokio::time::pause(); + if attempt < recovery_attempt || recovery_attempt > MAX_SNAPSHOT_ATTEMPTS { + cache + .prepare(¶ms, &mut prepared) + .await + .expect("capture must stay cached during backoff"); + assert_eq!( + (&prepared.command, &prepared.env), + (¶ms.argv, ¶ms.env) + ); + } else { + assert_ne!(prepared.command, params.argv); + } + assert_eq!( + std::fs::read_to_string(home.path().join("captures"))?, + "x".repeat(attempt.min(recovery_attempt).min(MAX_SNAPSHOT_ATTEMPTS)) + ); + tokio::time::advance(SNAPSHOT_RETRY_BACKOFF).await; + tokio::time::resume(); + } + Ok(()) +} #[test] fn snapshot_filters_profile_exports_after_capture() { diff --git a/codex-rs/exec-server/tests/exec_process.rs b/codex-rs/exec-server/tests/exec_process.rs index 5b107dd245..82fc6590f5 100644 --- a/codex-rs/exec-server/tests/exec_process.rs +++ b/codex-rs/exec-server/tests/exec_process.rs @@ -351,13 +351,39 @@ async fn shell_snapshot_v2_remote_managed_proxy_uses_prepared_execution_context( } #[cfg(unix)] +#[test_case(false, false, "bash", 1; "local_pipe_recovery")] +#[test_case(false, true, "bash", 1; "local_tty_recovery")] +#[test_case(true, false, "bash", 1; "remote_pipe_recovery")] +#[test_case(true, true, "bash", 1; "remote_tty_recovery")] +#[test_case(false, false, "bash", 3; "local_retry_budget_exhausted")] +#[test_case(true, false, "bash", 3; "remote_retry_budget_exhausted")] +#[cfg_attr(target_os = "macos", test_case(false, false, "zsh", 1; "local_zsh_recovery"))] +#[cfg_attr(target_os = "macos", test_case(true, false, "zsh", 1; "remote_zsh_recovery"))] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn shell_snapshot_v2_capture_failure_falls_back_to_original_command() -> Result<()> { - let context = create_process_context(/*use_remote*/ false).await?; +#[serial_test::serial(remote_exec_server)] +async fn shell_snapshot_v2_capture_failure_falls_back_and_retries( + use_remote: bool, + tty: bool, + shell_name: &str, + failures_before_repair: usize, +) -> Result<()> { + if use_remote + && let Some(warning) = + codex_sandboxing::system_bwrap_warning(&PermissionProfile::workspace_write()) + { + eprintln!("skipping sandbox test: {warning}"); + return Ok(()); + } + let context = create_process_context(use_remote).await?; let home = TempDir::new()?; let cwd = PathUri::from_host_native_path(home.path())?; + let (shell_path, profile_name) = match shell_name { + "bash" => ("/bin/bash", ".bashrc"), + "zsh" => ("/bin/zsh", ".zshrc"), + name => anyhow::bail!("unsupported test shell {name}"), + }; std::fs::write( - home.path().join(".bashrc"), + home.path().join(profile_name), "printf x >> \"$HOME/captures\"\nexit 7\n", )?; let policy = ExecEnvPolicy { @@ -373,30 +399,35 @@ async fn shell_snapshot_v2_capture_failure_falls_back_to_original_command() -> R let mut params = ExecParams { process_id: ProcessId::from("snapshot-first"), argv: vec![ - "/bin/bash".to_string(), + shell_path.to_string(), "-lc".to_string(), - "printf original".to_string(), + "if command -v profile_helper >/dev/null; then profile_helper; else printf original; fi".to_string(), ], - cwd, + cwd: cwd.clone(), env_policy: Some(policy), shell_snapshot: Some(ShellSnapshotRequest { scope_id: "attachment-1".to_string(), shell: ShellInfo { - name: "bash".to_string(), - path: "/bin/bash".to_string(), + name: shell_name.to_string(), + path: shell_path.to_string(), }, }), env: HashMap::new(), - tty: false, + tty, pipe_stdin: false, arg0: None, - sandbox: None, + sandbox: use_remote.then(|| { + FileSystemSandboxContext::from_permission_profile_with_cwd( + PermissionProfile::workspace_write(), + cwd, + ) + }), enforce_managed_network: false, managed_network: None, network_proxy: None, }; - for attempt in 0..2 { + for attempt in 0..failures_before_repair { params.process_id = ProcessId::from(format!("snapshot-fallback-{attempt}")); let fallback = context.backend.start(params.clone()).await?; let fallback_output = collect_process_output_from_events(fallback.process).await?; @@ -404,8 +435,36 @@ async fn shell_snapshot_v2_capture_failure_falls_back_to_original_command() -> R fallback_output, ("original".to_string(), String::new(), Some(0), true) ); + // A real remote executor has its own clock; the unit test uses a + // paused clock to check requests made during the one-second backoff. + sleep(Duration::from_millis(1100)).await; } - assert_eq!(std::fs::read_to_string(home.path().join("captures"))?, "x"); + assert_eq!( + std::fs::read_to_string(home.path().join("captures"))?, + "x".repeat(failures_before_repair) + ); + + std::fs::write( + home.path().join(profile_name), + "printf x >> \"$HOME/captures\"\nprofile_helper() { printf recovered; }\n", + )?; + let (expected_output, expected_captures) = if failures_before_repair == 3 { + ("original", "xxx") + } else { + ("recovered", "xx") + }; + for attempt in 0..2 { + params.process_id = ProcessId::from(format!("snapshot-after-repair-{attempt}")); + let started = context.backend.start(params.clone()).await?; + assert_eq!( + collect_process_output_from_events(started.process).await?, + (expected_output.to_string(), String::new(), Some(0), true) + ); + } + assert_eq!( + std::fs::read_to_string(home.path().join("captures"))?, + expected_captures + ); Ok(()) }