From 18dd539c486288132f3777a38d4e08831e36fff0 Mon Sep 17 00:00:00 2001 From: jimmyfraiture Date: Thu, 4 Sep 2025 12:14:17 -0700 Subject: [PATCH] MacOS snapshots --- codex-rs/core/src/codex.rs | 72 +++++- codex-rs/core/src/shell.rs | 266 ++++++++++++++++++-- codex-rs/core/tests/suite/prompt_caching.rs | 3 +- 3 files changed, 310 insertions(+), 31 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 2028c4240e..4fe19e9c7d 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -286,6 +286,7 @@ pub(crate) struct Session { codex_linux_sandbox_exe: Option, user_shell: shell::Shell, show_raw_agent_reasoning: bool, + shell_snapshot_path: Option, } /// 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, } +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) -> 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 { diff --git a/codex-rs/core/src/shell.rs b/codex-rs/core/src/shell.rs index 3a69874eb2..bb30055cc0 100644 --- a/codex-rs/core/src/shell.rs +++ b/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, } #[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize)] @@ -26,22 +32,38 @@ impl Shell { pub fn format_default_shell_invocation(&self, command: Vec) -> Option> { 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) -> Option { } #[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 { + [".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 { + 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 diff --git a/codex-rs/core/tests/suite/prompt_caching.rs b/codex-rs/core/tests/suite/prompt_caching.rs index 11c6ea9acf..78fcc88fc5 100644 --- a/codex-rs/core/tests/suite/prompt_caching.rs +++ b/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#"