MacOS snapshots

codex-rs/core/src/codex.rs

@@ -286,6 +286,7 @@ pub(crate) struct Session {
     codex_linux_sandbox_exe: Option<PathBuf>,
     user_shell: shell::Shell,
     show_raw_agent_reasoning: bool,
+    shell_snapshot_path: Option<PathBuf>,
 }
 
 /// The context needed for a single turn of the conversation.
@@ -391,7 +392,7 @@ impl Session {
         let rollout_fut = RolloutRecorder::new(&config, session_id, user_instructions.clone());
 
         let mcp_fut = McpConnectionManager::new(config.mcp_servers.clone());
-        let default_shell_fut = shell::default_user_shell();
+        let default_shell_fut = shell::default_user_shell(session_id);
         let history_meta_fut = crate::message_history::history_metadata(&config);
 
         // Join all independent futures.
@@ -464,6 +465,12 @@ impl Session {
             cwd,
             disable_response_storage,
         };
+        // TODO(jif) add support for other shells.
+        let shell_snapshot_path = match &default_shell {
+            shell::Shell::Zsh(zsh) => zsh.snapshot_path.clone(),
+            _ => None,
+        };
+
         let sess = Arc::new(Session {
             session_id,
             tx_event: tx_event.clone(),
@@ -475,6 +482,7 @@ impl Session {
             codex_linux_sandbox_exe: config.codex_linux_sandbox_exe.clone(),
             user_shell: default_shell,
             show_raw_agent_reasoning: config.show_raw_agent_reasoning,
+            shell_snapshot_path,
         });
 
         // Dispatch the SessionConfiguredEvent first and then report any errors.
@@ -947,6 +955,9 @@ impl Session {
 impl Drop for Session {
     fn drop(&mut self) {
         self.interrupt_task();
+        if let Some(path) = &self.shell_snapshot_path {
+            shell::delete_shell_snapshot(path);
+        }
     }
 }
 
@@ -2293,13 +2304,25 @@ pub struct ExecInvokeArgs<'a> {
     pub stdout_stream: Option<StdoutStream>,
 }
 
+fn should_translate_shell_command(
+    shell: &crate::shell::Shell,
+    shell_policy: &ShellEnvironmentPolicy,
+) -> bool {
+    matches!(shell, crate::shell::Shell::PowerShell(_))
+        || shell_policy.use_profile
+        || matches!(
+            shell,
+            crate::shell::Shell::Zsh(zsh) if zsh.snapshot_path.is_some()
+        )
+}
+
 fn maybe_translate_shell_command(
     params: ExecParams,
     sess: &Session,
     turn_context: &TurnContext,
 ) -> ExecParams {
-    let should_translate = matches!(sess.user_shell, crate::shell::Shell::PowerShell(_))
-        || turn_context.shell_environment_policy.use_profile;
+    let should_translate =
+        should_translate_shell_command(&sess.user_shell, &turn_context.shell_environment_policy);
 
     if should_translate
         && let Some(command) = sess
@@ -2908,10 +2931,13 @@ fn convert_call_tool_result_to_function_call_output_payload(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::config_types::ShellEnvironmentPolicyInherit;
     use mcp_types::ContentBlock;
     use mcp_types::TextContent;
     use pretty_assertions::assert_eq;
     use serde_json::json;
+    use std::collections::HashMap;
+    use std::path::PathBuf;
     use std::time::Duration as StdDuration;
 
     fn text_block(s: &str) -> ContentBlock {
@@ -2922,6 +2948,46 @@ mod tests {
         })
     }
 
+    fn shell_policy_with_profile(use_profile: bool) -> ShellEnvironmentPolicy {
+        ShellEnvironmentPolicy {
+            inherit: ShellEnvironmentPolicyInherit::All,
+            ignore_default_excludes: false,
+            exclude: Vec::new(),
+            r#set: HashMap::new(),
+            include_only: Vec::new(),
+            use_profile,
+        }
+    }
+
+    fn zsh_shell(snapshot_path: Option<PathBuf>) -> shell::Shell {
+        shell::Shell::Zsh(shell::ZshShell {
+            shell_path: "/bin/zsh".to_string(),
+            zshrc_path: "/Users/example/.zshrc".to_string(),
+            snapshot_path,
+        })
+    }
+
+    #[test]
+    fn translates_commands_when_shell_policy_requests_profile() {
+        let policy = shell_policy_with_profile(true);
+        let shell = zsh_shell(None);
+        assert!(should_translate_shell_command(&shell, &policy));
+    }
+
+    #[test]
+    fn translates_commands_for_zsh_with_snapshot() {
+        let policy = shell_policy_with_profile(false);
+        let shell = zsh_shell(Some(PathBuf::from("/tmp/snapshot")));
+        assert!(should_translate_shell_command(&shell, &policy));
+    }
+
+    #[test]
+    fn bypasses_translation_for_zsh_without_snapshot_or_profile() {
+        let policy = shell_policy_with_profile(false);
+        let shell = zsh_shell(None);
+        assert!(!should_translate_shell_command(&shell, &policy));
+    }
+
     #[test]
     fn prefers_structured_content_when_present() {
         let ctr = CallToolResult {

codex-rs/core/src/shell.rs

@@ -1,12 +1,18 @@
 use serde::Deserialize;
 use serde::Serialize;
 use shlex;
+use std::path::Path;
 use std::path::PathBuf;
+use tokio::process::Command;
+use tracing::trace;
+use tracing::warn;
+use uuid::Uuid;
 
 #[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize)]
 pub struct ZshShell {
-    shell_path: String,
-    zshrc_path: String,
+    pub(crate) shell_path: String,
+    pub(crate) zshrc_path: String,
+    pub(crate) snapshot_path: Option<PathBuf>,
 }
 
 #[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize)]
@@ -26,22 +32,38 @@ impl Shell {
     pub fn format_default_shell_invocation(&self, command: Vec<String>) -> Option<Vec<String>> {
         match self {
             Shell::Zsh(zsh) => {
-                if !std::path::Path::new(&zsh.zshrc_path).exists() {
+                if !Path::new(&zsh.zshrc_path).exists() {
                     return None;
                 }
 
-                let mut result = vec![zsh.shell_path.clone()];
-                result.push("-lc".to_string());
-
                 let joined = strip_bash_lc(&command)
                     .or_else(|| shlex::try_join(command.iter().map(|s| s.as_str())).ok());
 
-                if let Some(joined) = joined {
-                    result.push(format!("source {} && ({joined})", zsh.zshrc_path));
-                } else {
-                    return None;
+                let joined = joined?;
+
+                if let Some(snapshot_path) = &zsh.snapshot_path
+                    && snapshot_path.exists()
+                {
+                    let snapshot_path_string = snapshot_path.to_string_lossy();
+                    trace!(
+                        snapshot_path = %snapshot_path_string,
+                        "using cached zsh snapshot"
+                    );
+                    return Some(vec![
+                        zsh.shell_path.clone(),
+                        "-c".to_string(),
+                        format!("source {snapshot_path_string} && ({joined})"),
+                    ]);
                 }
-                Some(result)
+
+                trace!("no snapshot available; falling back to zshrc");
+
+                let zshrc_path = &zsh.zshrc_path;
+                Some(vec![
+                    zsh.shell_path.clone(),
+                    "-lc".to_string(),
+                    format!("source {zshrc_path} && ({joined})"),
+                ])
             }
             Shell::PowerShell(ps) => {
                 // If model generated a bash command, prefer a detected bash fallback
@@ -117,14 +139,11 @@ fn strip_bash_lc(command: &Vec<String>) -> Option<String> {
 }
 
 #[cfg(target_os = "macos")]
-pub async fn default_user_shell() -> Shell {
-    use tokio::process::Command;
-    use whoami;
-
+pub async fn default_user_shell(session_id: Uuid) -> Shell {
     let user = whoami::username();
-    let home = format!("/Users/{user}");
+    let home = PathBuf::from(format!("/Users/{user}"));
     let output = Command::new("dscl")
-        .args([".", "-read", &home, "UserShell"])
+        .args([".", "-read", home.to_string_lossy().as_ref(), "UserShell"])
         .output()
         .await
         .ok();
@@ -138,9 +157,15 @@ pub async fn default_user_shell() -> Shell {
                 if let Some(shell_path) = line.strip_prefix("UserShell: ")
                     && shell_path.ends_with("/zsh")
                 {
+                    let snapshot_path = ensure_zsh_snapshot(shell_path, &home, session_id).await;
+                    if snapshot_path.is_none() {
+                        trace!("failed to prepare zsh snapshot; using live profile");
+                    }
+
                     return Shell::Zsh(ZshShell {
                         shell_path: shell_path.to_string(),
-                        zshrc_path: format!("{home}/.zshrc"),
+                        zshrc_path: home.join(".zshrc").to_string_lossy().to_string(),
+                        snapshot_path,
                     });
                 }
             }
@@ -152,12 +177,12 @@ pub async fn default_user_shell() -> Shell {
 }
 
 #[cfg(all(not(target_os = "macos"), not(target_os = "windows")))]
-pub async fn default_user_shell() -> Shell {
+pub async fn default_user_shell(_session_id: Uuid) -> Shell {
     Shell::Unknown
 }
 
 #[cfg(target_os = "windows")]
-pub async fn default_user_shell() -> Shell {
+pub async fn default_user_shell(_session_id: Uuid) -> Shell {
     use tokio::process::Command;
 
     // Prefer PowerShell 7+ (`pwsh`) if available, otherwise fall back to Windows PowerShell.
@@ -196,11 +221,131 @@ pub async fn default_user_shell() -> Shell {
     }
 }
 
+#[cfg(target_os = "macos")]
+fn zsh_profile_paths(home: &Path) -> Vec<PathBuf> {
+    [".zshrc", ".zprofile", ".zshenv", ".zlogin"]
+        .into_iter()
+        .map(|name| home.join(name))
+        .collect()
+}
+
+#[cfg(target_os = "macos")]
+async fn ensure_zsh_snapshot(shell_path: &str, home: &Path, session_id: Uuid) -> Option<PathBuf> {
+    let snapshot_path = home.join(format!(".codex_shell_snapshot_{session_id}.zsh"));
+
+    // Check if an update in the profile requires to re-generate the snapshot.
+    let snapshot_is_stale = match tokio::fs::metadata(&snapshot_path).await {
+        Ok(metadata) => match metadata.modified() {
+            Ok(snapshot_modified) => {
+                let mut stale = false;
+                for profile in zsh_profile_paths(home) {
+                    if let Ok(profile_metadata) = tokio::fs::metadata(&profile).await {
+                        match profile_metadata.modified() {
+                            Ok(profile_modified) => {
+                                if profile_modified > snapshot_modified {
+                                    stale = true;
+                                    break;
+                                }
+                            }
+                            Err(_) => {
+                                stale = true;
+                                break;
+                            }
+                        }
+                    }
+                }
+                stale
+            }
+            Err(_) => true,
+        },
+        Err(_) => true,
+    };
+
+    if !snapshot_is_stale {
+        return Some(snapshot_path);
+    }
+
+    match regenerate_zsh_snapshot(shell_path, home, &snapshot_path).await {
+        Ok(()) => Some(snapshot_path),
+        Err(err) => {
+            warn!("failed to generate zsh snapshot: {err}");
+            None
+        }
+    }
+}
+
+#[cfg(target_os = "macos")]
+async fn regenerate_zsh_snapshot(
+    shell_path: &str,
+    home: &Path,
+    snapshot_path: &Path,
+) -> std::io::Result<()> {
+    // Use `emulate -L sh` instead of `set -o posix` so we work on zsh builds
+    // that disable that option. Guard `alias -p` with `|| true` so the script
+    // keeps a zero exit status even if aliases are disabled.
+    let mut source_profiles = String::new();
+    for profile in zsh_profile_paths(home) {
+        let profile_string = profile.to_string_lossy().into_owned();
+        let quoted =
+            shlex::try_quote(&profile_string).unwrap_or_else(|_| profile_string.clone().into());
+        source_profiles.push_str(&format!("[ -f {quoted} ] && source {quoted}; "));
+    }
+
+    let capture_script = format!(
+        "{source_profiles}setopt posixbuiltins; export -p; {{ alias | sed 's/^/alias /'; }} 2>/dev/null || true"
+    );
+    let output = Command::new(shell_path)
+        .arg("-lc")
+        .arg(capture_script)
+        .env("HOME", home)
+        .output()
+        .await?;
+
+    if !output.status.success() {
+        return Err(std::io::Error::other(format!(
+            "snapshot capture exited with status {}",
+            output.status
+        )));
+    }
+
+    let mut contents = String::from("# Generated by Codex. Do not edit.\n");
+
+    contents.push_str(&String::from_utf8_lossy(&output.stdout));
+    contents.push('\n');
+
+    let tmp_path = snapshot_path.with_extension("tmp");
+    tokio::fs::write(&tmp_path, contents).await?;
+
+    #[cfg(unix)]
+    {
+        // Restrict the snapshot to user read/write so that environment variables or aliases
+        // that may contain secrets are not exposed to other users on the system.
+        use std::os::unix::fs::PermissionsExt;
+        let permissions = std::fs::Permissions::from_mode(0o600);
+        tokio::fs::set_permissions(&tmp_path, permissions).await?;
+    }
+
+    tokio::fs::rename(&tmp_path, snapshot_path).await?;
+    Ok(())
+}
+
+#[cfg(target_os = "macos")]
+pub(crate) fn delete_shell_snapshot(path: &Path) {
+    if let Err(err) = std::fs::remove_file(path) {
+        trace!(?path, %err, "failed to delete shell snapshot");
+    }
+}
+
+#[cfg(not(target_os = "macos"))]
+pub(crate) fn delete_shell_snapshot(_path: &Path) {}
+
 #[cfg(test)]
 #[cfg(target_os = "macos")]
 mod tests {
     use super::*;
+    use std::path::Path;
     use std::process::Command;
+    use uuid::Uuid;
 
     #[tokio::test]
     async fn test_current_shell_detects_zsh() {
@@ -213,13 +358,13 @@ mod tests {
         let home = std::env::var("HOME").unwrap();
         let shell_path = String::from_utf8_lossy(&shell.stdout).trim().to_string();
         if shell_path.ends_with("/zsh") {
-            assert_eq!(
-                default_user_shell().await,
-                Shell::Zsh(ZshShell {
-                    shell_path: shell_path.to_string(),
-                    zshrc_path: format!("{home}/.zshrc",),
-                })
-            );
+            match default_user_shell(Uuid::new_v4()).await {
+                Shell::Zsh(zsh) => {
+                    assert_eq!(zsh.shell_path, shell_path);
+                    assert_eq!(zsh.zshrc_path, format!("{home}/.zshrc"));
+                }
+                other => panic!("unexpected shell returned: {other:?}"),
+            }
         }
     }
 
@@ -228,11 +373,77 @@ mod tests {
         let shell = Shell::Zsh(ZshShell {
             shell_path: "/bin/zsh".to_string(),
             zshrc_path: "/does/not/exist/.zshrc".to_string(),
+            snapshot_path: None,
         });
         let actual_cmd = shell.format_default_shell_invocation(vec!["myecho".to_string()]);
         assert_eq!(actual_cmd, None);
     }
 
+    #[tokio::test]
+    async fn test_snapshot_generation_uses_session_id_and_cleanup() {
+        let shell_path = "/bin/zsh";
+        if !Path::new(shell_path).exists() {
+            return;
+        }
+
+        let temp_home = tempfile::tempdir().unwrap();
+        std::fs::write(
+            temp_home.path().join(".zshrc"),
+            "export SNAPSHOT_TEST_VAR=1\nalias snapshot_test_alias='echo hi'\n",
+        )
+        .unwrap();
+
+        let session_id = Uuid::new_v4();
+        let snapshot_path = ensure_zsh_snapshot(shell_path, temp_home.path(), session_id)
+            .await
+            .expect("snapshot path");
+
+        let filename = snapshot_path
+            .file_name()
+            .unwrap()
+            .to_string_lossy()
+            .to_string();
+        assert!(filename.contains(&session_id.to_string()));
+        assert!(snapshot_path.exists());
+
+        let snapshot_path_second = ensure_zsh_snapshot(shell_path, temp_home.path(), session_id)
+            .await
+            .expect("snapshot path");
+        assert_eq!(snapshot_path, snapshot_path_second);
+
+        let contents = std::fs::read_to_string(&snapshot_path).unwrap();
+        assert!(contents.contains("alias snapshot_test_alias='echo hi'"));
+        assert!(contents.contains("SNAPSHOT_TEST_VAR=1"));
+
+        delete_shell_snapshot(&snapshot_path);
+        assert!(!snapshot_path.exists());
+    }
+
+    #[test]
+    fn format_default_shell_invocation_prefers_snapshot_when_available() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let snapshot_path = temp_dir.path().join("snapshot.zsh");
+        std::fs::write(&snapshot_path, "export SNAPSHOT_READY=1").unwrap();
+
+        let shell = Shell::Zsh(ZshShell {
+            shell_path: "/bin/zsh".to_string(),
+            zshrc_path: {
+                let path = temp_dir.path().join(".zshrc");
+                std::fs::write(&path, "# test zshrc").unwrap();
+                path.to_string_lossy().to_string()
+            },
+            snapshot_path: Some(snapshot_path.clone()),
+        });
+
+        let invocation = shell.format_default_shell_invocation(vec!["echo".to_string()]);
+        let expected_command = vec!["/bin/zsh".to_string(), "-c".to_string(), {
+            let snapshot_path = snapshot_path.to_string_lossy();
+            format!("source {snapshot_path} && (echo)")
+        }];
+
+        assert_eq!(invocation, Some(expected_command));
+    }
+
     #[tokio::test]
     async fn test_run_with_profile_escaping_and_execution() {
         let shell_path = "/bin/zsh";
@@ -292,6 +503,7 @@ mod tests {
             let shell = Shell::Zsh(ZshShell {
                 shell_path: shell_path.to_string(),
                 zshrc_path: zshrc_path.to_str().unwrap().to_string(),
+                snapshot_path: None,
             });
 
             let actual_cmd = shell

codex-rs/core/tests/suite/prompt_caching.rs

@@ -17,6 +17,7 @@ use core_test_support::load_default_config_for_test;
 use core_test_support::load_sse_fixture_with_id;
 use core_test_support::wait_for_event;
 use tempfile::TempDir;
+use uuid::Uuid;
 use wiremock::Mock;
 use wiremock::MockServer;
 use wiremock::ResponseTemplate;
@@ -269,7 +270,7 @@ async fn prefixes_context_and_instructions_once_and_consistently_across_requests
     let requests = server.received_requests().await.unwrap();
     assert_eq!(requests.len(), 2, "expected two POST requests");
 
-    let shell = default_user_shell().await;
+    let shell = default_user_shell(Uuid::new_v4()).await;
 
     let expected_env_text = format!(
         r#"<environment_context>