Retry failed shell snapshot captures (#40447)

## Why

A failed shell snapshot capture was cached for the lifetime of its cache entry, so later commands continued using the original command even after a transient profile error was fixed.

## What changed

- Retry failed captures after a one-second backoff, up to three total attempts.
- Share each attempt across concurrent requests and preserve command fallback during failures and backoff.
- Keep the final failure cached after the retry budget is exhausted.

## Testing

Cover recovery and retry exhaustion for concurrent requests, local and remote execution, pipes and TTYs, and Bash and zsh where available.

GitOrigin-RevId: b432ef7c9eb35a3a0a703ac675b4db6954c43a80

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<ExecEnvPolicy>,
     sandbox: Option<FileSystemSandboxContext>,
-    snapshot: Arc<OnceCell<Option<ShellSnapshot>>>,
+    attempts: usize,
+    // Failed captures store the earliest time another attempt may start.
+    snapshot: Arc<OnceCell<Result<ShellSnapshot, Instant>>>,
 }
 
 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 {

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(
+            &params,
+            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(
+            &params,
+            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(&params, &mut prepared),
+            cache.prepare(&params, &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(&params, &mut prepared)
+                .await
+                .expect("capture must stay cached during backoff");
+            assert_eq!(
+                (&prepared.command, &prepared.env),
+                (&params.argv, &params.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() {

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(())
 }