Merge branch 'main' into jif/investigate-shell-snapshot-issue-for

codex-rs/core/src/tools/runtimes/mod.rs

@@ -95,8 +95,26 @@ pub(crate) fn maybe_wrap_shell_lc_with_snapshot(
         .iter()
         .map(|arg| format!(" '{}'", shell_single_quote(arg)))
         .collect::<String>();
+    // Preserve command-process environment precedence:
+    // 1) Capture the current exported environment before snapshot sourcing.
+    // 2) Source the snapshot (best effort).
+    // 3) Re-apply the original exported environment so command/worktree values
+    //    win on conflicts, while snapshot-only vars remain present.
+    // 4) Exec the original shell command.
+    //
+    // We keep the capture in-memory (no temp files) to avoid leaking exported
+    // environment values to disk.
     let rewritten_script = format!(
-        "if . '{snapshot_path}' >/dev/null 2>&1; then :; fi; exec '{original_shell}' -c '{original_script}'{trailing_args}"
+        r#"__codex_restore_env="$(export -p 2>/dev/null || true)"
+
+if . '{snapshot_path}' >/dev/null 2>&1; then :; fi
+
+if [ -n "$__codex_restore_env" ]; then
+  eval "$__codex_restore_env" >/dev/null 2>&1 || true
+fi
+unset __codex_restore_env >/dev/null 2>&1 || true
+
+exec '{original_shell}' -c '{original_script}'{trailing_args}"#
     );
 
     vec![shell_path.to_string(), "-c".to_string(), rewritten_script]
@@ -113,6 +131,7 @@ mod tests {
     use crate::shell_snapshot::ShellSnapshot;
     use pretty_assertions::assert_eq;
     use std::path::PathBuf;
+    use std::process::Command;
     use std::sync::Arc;
     use tempfile::tempdir;
     use tokio::sync::watch;
@@ -306,4 +325,63 @@ mod tests {
         assert!(rewritten[2].contains("if . '"));
         assert!(rewritten[2].contains("exec '/bin/bash' -c 'echo hello'"));
     }
+
+    #[test]
+    fn maybe_wrap_shell_lc_with_snapshot_restores_original_environment_precedence() {
+        let dir = tempdir().expect("create temp dir");
+        let snapshot_path = dir.path().join("snapshot.sh");
+        std::fs::write(
+            &snapshot_path,
+            "# Snapshot file\nexport TEST_ENV_SNAPSHOT=global\nexport SNAPSHOT_ONLY=from_snapshot\n",
+        )
+        .expect("write snapshot");
+        let session_shell = shell_with_snapshot(
+            ShellType::Bash,
+            "/bin/bash",
+            snapshot_path,
+            dir.path().to_path_buf(),
+        );
+        let command = vec![
+            "/bin/bash".to_string(),
+            "-lc".to_string(),
+            "printf '%s|%s' \"$TEST_ENV_SNAPSHOT\" \"${SNAPSHOT_ONLY-unset}\"".to_string(),
+        ];
+
+        let rewritten = maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, dir.path());
+        let output = Command::new(&rewritten[0])
+            .args(&rewritten[1..])
+            .env("TEST_ENV_SNAPSHOT", "worktree")
+            .output()
+            .expect("run rewritten command");
+
+        assert!(output.status.success(), "command failed: {output:?}");
+        assert_eq!(
+            String::from_utf8_lossy(&output.stdout),
+            "worktree|from_snapshot"
+        );
+    }
+
+    #[test]
+    fn maybe_wrap_shell_lc_with_snapshot_avoids_temp_file_env_dump() {
+        let dir = tempdir().expect("create temp dir");
+        let snapshot_path = dir.path().join("snapshot.sh");
+        std::fs::write(&snapshot_path, "# Snapshot file\n").expect("write snapshot");
+        let session_shell = shell_with_snapshot(
+            ShellType::Bash,
+            "/bin/bash",
+            snapshot_path,
+            dir.path().to_path_buf(),
+        );
+        let command = vec![
+            "/bin/bash".to_string(),
+            "-lc".to_string(),
+            "echo hello".to_string(),
+        ];
+
+        let rewritten = maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, dir.path());
+
+        assert!(!rewritten[2].contains("codex-restore-env-"));
+        assert!(!rewritten[2].contains("export -p >"));
+        assert!(!rewritten[2].contains("rm -f \"$__codex_restore_env_file\""));
+    }
 }