Add regression coverage for timed-out zsh snapshots

codex-rs/core/src/shell_snapshot.rs

@@ -12,7 +12,6 @@ use crate::shell::ShellType;
 use crate::shell::get_shell;
 use anyhow::Context;
 use anyhow::Result;
-use anyhow::anyhow;
 use anyhow::bail;
 use codex_otel::SessionTelemetry;
 use codex_protocol::ThreadId;
@@ -35,6 +34,18 @@ const SNAPSHOT_RETENTION: Duration = Duration::from_secs(60 * 60 * 24 * 3); // 3
 const SNAPSHOT_DIR: &str = "shell_snapshots";
 const EXCLUDED_EXPORT_VARS: &[&str] = &["PWD", "OLDPWD"];
 
+#[derive(Debug, thiserror::Error)]
+#[error("Snapshot command timed out for {shell_name}")]
+struct SnapshotCommandTimeout {
+    shell_name: &'static str,
+}
+
+#[derive(Clone, Copy)]
+enum ZshSnapshotMode {
+    SourceUserRc,
+    SkipUserRc,
+}
+
 impl ShellSnapshot {
     pub fn start_snapshotting(
         codex_home: AbsolutePathBuf,
@@ -227,7 +238,7 @@ async fn write_shell_snapshot(
 async fn capture_snapshot(shell: &Shell, cwd: &AbsolutePathBuf) -> Result<String> {
     let shell_type = shell.shell_type.clone();
     match shell_type {
-        ShellType::Zsh => run_shell_script(shell, &zsh_snapshot_script(), cwd).await,
+        ShellType::Zsh => capture_zsh_snapshot(shell, cwd, SNAPSHOT_TIMEOUT).await,
         ShellType::Bash => run_shell_script(shell, &bash_snapshot_script(), cwd).await,
         ShellType::Sh => run_shell_script(shell, &sh_snapshot_script(), cwd).await,
         ShellType::PowerShell => run_shell_script(shell, powershell_snapshot_script(), cwd).await,
@@ -235,6 +246,39 @@ async fn capture_snapshot(shell: &Shell, cwd: &AbsolutePathBuf) -> Result<String
     }
 }
 
+async fn capture_zsh_snapshot(
+    shell: &Shell,
+    cwd: &AbsolutePathBuf,
+    snapshot_timeout: Duration,
+) -> Result<String> {
+    let snapshot = run_script_with_timeout(
+        shell,
+        &zsh_snapshot_script(ZshSnapshotMode::SourceUserRc),
+        snapshot_timeout,
+        /*use_login_shell*/ true,
+        cwd,
+    )
+    .await;
+
+    match snapshot {
+        Ok(snapshot) => Ok(snapshot),
+        Err(err) if err.is::<SnapshotCommandTimeout>() => {
+            tracing::warn!(
+                "Zsh shell snapshot timed out while sourcing user rc; retrying without user rc"
+            );
+            run_script_with_timeout(
+                shell,
+                &zsh_snapshot_script(ZshSnapshotMode::SkipUserRc),
+                snapshot_timeout,
+                /*use_login_shell*/ true,
+                cwd,
+            )
+            .await
+        }
+        Err(err) => Err(err),
+    }
+}
+
 fn strip_snapshot_preamble(snapshot: &str) -> Result<String> {
     let marker = "# Snapshot file";
     let Some(start) = snapshot.find(marker) else {
@@ -299,7 +343,7 @@ async fn run_script_with_timeout(
     handler.kill_on_drop(true);
     let output = timeout(snapshot_timeout, handler.output())
         .await
-        .map_err(|_| anyhow!("Snapshot command timed out for {shell_name}"))?
+        .map_err(|_| SnapshotCommandTimeout { shell_name })?
         .with_context(|| format!("Failed to execute {shell_name}"))?;
 
     if !output.status.success() {
@@ -315,15 +359,20 @@ fn excluded_exports_regex() -> String {
     EXCLUDED_EXPORT_VARS.join("|")
 }
 
-fn zsh_snapshot_script() -> String {
+fn zsh_snapshot_script(mode: ZshSnapshotMode) -> String {
     let excluded = excluded_exports_regex();
-    let script = r##"if [[ -n "$ZDOTDIR" ]]; then
+    let source_user_rc = match mode {
+        ZshSnapshotMode::SourceUserRc => {
+            r##"if [[ -n "$ZDOTDIR" ]]; then
   rc="$ZDOTDIR/.zshrc"
 else
   rc="$HOME/.zshrc"
 fi
-[[ -r "$rc" ]] && . "$rc"
-print '# Snapshot file'
+[[ -r "$rc" ]] && . "$rc""##
+        }
+        ZshSnapshotMode::SkipUserRc => "",
+    };
+    let script = r##"print '# Snapshot file'
 print '# Unset all aliases to avoid conflicts with functions'
 print 'unalias -a 2>/dev/null || true'
 print '# Functions'
@@ -356,7 +405,7 @@ if [[ -n "$export_lines" ]]; then
   print -r -- "$export_lines"
 fi
 "##;
-    script.replace("EXCLUDED_EXPORTS", &excluded)
+    format!("{source_user_rc}\n{script}").replace("EXCLUDED_EXPORTS", &excluded)
 }
 
 fn bash_snapshot_script() -> String {

codex-rs/core/src/shell_snapshot_tests.rs

@@ -1,9 +1,12 @@
 use super::*;
+use anyhow::anyhow;
 use core_test_support::PathBufExt;
 use core_test_support::PathExt;
 use pretty_assertions::assert_eq;
 #[cfg(unix)]
 use std::os::unix::ffi::OsStrExt;
+#[cfg(unix)]
+use std::os::unix::fs::PermissionsExt;
 use std::path::PathBuf;
 #[cfg(unix)]
 use std::process::Command;
@@ -374,6 +377,41 @@ async fn timed_out_snapshot_shell_is_terminated() -> Result<()> {
     Ok(())
 }
 
+#[cfg(unix)]
+#[tokio::test]
+async fn timed_out_zsh_user_rc_uses_login_shell_fallback_snapshot() -> Result<()> {
+    let dir = tempdir()?;
+    let shell_path = dir.path().join("zsh-wrapper.sh");
+    fs::write(
+        &shell_path,
+        r#"#!/bin/sh
+marker="$0.first-call"
+if [ ! -e "$marker" ]; then
+  : > "$marker"
+  while :; do :; done
+fi
+case "$*" in
+  *".zshrc"*) exit 2 ;;
+esac
+printf '# Snapshot file\nexport PATH=/fallback/bin\n'
+"#,
+    )
+    .await?;
+    std::fs::set_permissions(&shell_path, std::fs::Permissions::from_mode(0o755))?;
+
+    let shell = Shell {
+        shell_type: ShellType::Zsh,
+        shell_path,
+        shell_snapshot: crate::shell::empty_shell_snapshot_receiver(),
+    };
+
+    let snapshot = capture_zsh_snapshot(&shell, &dir.path().abs(), Duration::from_secs(1)).await?;
+
+    assert_eq!(snapshot, "# Snapshot file\nexport PATH=/fallback/bin\n");
+
+    Ok(())
+}
+
 #[cfg(target_os = "macos")]
 #[tokio::test]
 async fn macos_zsh_snapshot_includes_sections() -> Result<()> {