feat: use shell policy in shell snapshot (#11759)

Honor `shell_environment_policy.set` even after a shell snapshot

codex-rs/core/src/tasks/user_shell.rs

@@ -119,6 +119,7 @@ pub(crate) async fn execute_user_shell_command(
         &display_command,
         session_shell.as_ref(),
         turn_context.cwd.as_path(),
+        &turn_context.shell_environment_policy.r#set,
     );
 
     let call_id = Uuid::new_v4().to_string();

codex-rs/core/src/tools/handlers/shell.rs

@@ -254,7 +254,14 @@ impl ShellHandler {
         let mut exec_params = exec_params;
         let dependency_env = session.dependency_env().await;
         if !dependency_env.is_empty() {
-            exec_params.env.extend(dependency_env);
+            exec_params.env.extend(dependency_env.clone());
+        }
+
+        let mut explicit_env_overrides = turn.shell_environment_policy.r#set.clone();
+        for key in dependency_env.keys() {
+            if let Some(value) = exec_params.env.get(key) {
+                explicit_env_overrides.insert(key.clone(), value.clone());
+            }
         }
 
         // Approval policy guard for explicit escalation in non-OnRequest modes.
@@ -315,6 +322,7 @@ impl ShellHandler {
             cwd: exec_params.cwd.clone(),
             timeout_ms: exec_params.expiration.timeout_ms(),
             env: exec_params.env.clone(),
+            explicit_env_overrides,
             network: exec_params.network.clone(),
             sandbox_permissions: exec_params.sandbox_permissions,
             justification: exec_params.justification.clone(),

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

@@ -57,7 +57,12 @@ pub(crate) fn maybe_wrap_shell_lc_with_snapshot(
     command: &[String],
     session_shell: &Shell,
     cwd: &Path,
+    explicit_env_overrides: &HashMap<String, String>,
 ) -> Vec<String> {
+    if cfg!(windows) {
+        return command.to_vec();
+    }
+
     let Some(snapshot) = session_shell.shell_snapshot() else {
         return command.to_vec();
     };
@@ -95,24 +100,78 @@ pub(crate) fn maybe_wrap_shell_lc_with_snapshot(
         .iter()
         .map(|arg| format!(" '{}'", shell_single_quote(arg)))
         .collect::<String>();
-    let rewritten_script = format!(
-        "if . '{snapshot_path}' >/dev/null 2>&1; then :; fi; exec '{original_shell}' -c '{original_script}'{trailing_args}"
-    );
+    let (override_captures, override_exports) = build_override_exports(explicit_env_overrides);
+    let rewritten_script = if override_exports.is_empty() {
+        format!(
+            "if . '{snapshot_path}' >/dev/null 2>&1; then :; fi\n\nexec '{original_shell}' -c '{original_script}'{trailing_args}"
+        )
+    } else {
+        format!(
+            "{override_captures}\n\nif . '{snapshot_path}' >/dev/null 2>&1; then :; fi\n\n{override_exports}\n\nexec '{original_shell}' -c '{original_script}'{trailing_args}"
+        )
+    };
 
     vec![shell_path.to_string(), "-c".to_string(), rewritten_script]
 }
 
+fn build_override_exports(explicit_env_overrides: &HashMap<String, String>) -> (String, String) {
+    let mut keys = explicit_env_overrides
+        .keys()
+        .filter(|key| is_valid_shell_variable_name(key))
+        .collect::<Vec<_>>();
+    keys.sort_unstable();
+
+    if keys.is_empty() {
+        return (String::new(), String::new());
+    }
+
+    let captures = keys
+        .iter()
+        .enumerate()
+        .map(|(idx, key)| {
+            format!(
+                "__CODEX_SNAPSHOT_OVERRIDE_SET_{idx}=\"${{{key}+x}}\"\n__CODEX_SNAPSHOT_OVERRIDE_{idx}=\"${{{key}-}}\""
+            )
+        })
+        .collect::<Vec<_>>()
+        .join("\n");
+    let restores = keys
+        .iter()
+        .enumerate()
+        .map(|(idx, key)| {
+            format!(
+                "if [ -n \"${{__CODEX_SNAPSHOT_OVERRIDE_SET_{idx}}}\" ]; then export {key}=\"${{__CODEX_SNAPSHOT_OVERRIDE_{idx}}}\"; else unset {key}; fi"
+            )
+        })
+        .collect::<Vec<_>>()
+        .join("\n");
+
+    (captures, restores)
+}
+
+fn is_valid_shell_variable_name(name: &str) -> bool {
+    let mut chars = name.chars();
+    let Some(first) = chars.next() else {
+        return false;
+    };
+    if !(first == '_' || first.is_ascii_alphabetic()) {
+        return false;
+    }
+    chars.all(|c| c == '_' || c.is_ascii_alphanumeric())
+}
+
 fn shell_single_quote(input: &str) -> String {
     input.replace('\'', r#"'"'"'"#)
 }
 
-#[cfg(test)]
+#[cfg(all(test, unix))]
 mod tests {
     use super::*;
     use crate::shell::ShellType;
     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;
@@ -151,7 +210,12 @@ mod tests {
             "echo hello".to_string(),
         ];
 
-        let rewritten = maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, dir.path());
+        let rewritten = maybe_wrap_shell_lc_with_snapshot(
+            &command,
+            &session_shell,
+            dir.path(),
+            &HashMap::new(),
+        );
 
         assert_eq!(rewritten[0], "/bin/zsh");
         assert_eq!(rewritten[1], "-c");
@@ -176,7 +240,12 @@ mod tests {
             "echo 'hello'".to_string(),
         ];
 
-        let rewritten = maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, dir.path());
+        let rewritten = maybe_wrap_shell_lc_with_snapshot(
+            &command,
+            &session_shell,
+            dir.path(),
+            &HashMap::new(),
+        );
 
         assert!(rewritten[2].contains(r#"exec '/bin/bash' -c 'echo '"'"'hello'"'"''"#));
     }
@@ -198,7 +267,12 @@ mod tests {
             "echo hello".to_string(),
         ];
 
-        let rewritten = maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, dir.path());
+        let rewritten = maybe_wrap_shell_lc_with_snapshot(
+            &command,
+            &session_shell,
+            dir.path(),
+            &HashMap::new(),
+        );
 
         assert_eq!(rewritten[0], "/bin/bash");
         assert_eq!(rewritten[1], "-c");
@@ -223,7 +297,12 @@ mod tests {
             "echo hello".to_string(),
         ];
 
-        let rewritten = maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, dir.path());
+        let rewritten = maybe_wrap_shell_lc_with_snapshot(
+            &command,
+            &session_shell,
+            dir.path(),
+            &HashMap::new(),
+        );
 
         assert_eq!(rewritten[0], "/bin/sh");
         assert_eq!(rewritten[1], "-c");
@@ -250,7 +329,12 @@ mod tests {
             "arg1".to_string(),
         ];
 
-        let rewritten = maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, dir.path());
+        let rewritten = maybe_wrap_shell_lc_with_snapshot(
+            &command,
+            &session_shell,
+            dir.path(),
+            &HashMap::new(),
+        );
 
         assert!(
             rewritten[2].contains(
@@ -276,7 +360,12 @@ mod tests {
             "echo hello".to_string(),
         ];
 
-        let rewritten = maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, &command_cwd);
+        let rewritten = maybe_wrap_shell_lc_with_snapshot(
+            &command,
+            &session_shell,
+            &command_cwd,
+            &HashMap::new(),
+        );
 
         assert_eq!(rewritten, command);
     }
@@ -299,11 +388,214 @@ mod tests {
         ];
         let command_cwd = dir.path().join(".");
 
-        let rewritten = maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, &command_cwd);
+        let rewritten = maybe_wrap_shell_lc_with_snapshot(
+            &command,
+            &session_shell,
+            &command_cwd,
+            &HashMap::new(),
+        );
 
         assert_eq!(rewritten[0], "/bin/zsh");
         assert_eq!(rewritten[1], "-c");
         assert!(rewritten[2].contains("if . '"));
         assert!(rewritten[2].contains("exec '/bin/bash' -c 'echo hello'"));
     }
+
+    #[test]
+    fn maybe_wrap_shell_lc_with_snapshot_restores_explicit_override_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 explicit_env_overrides =
+            HashMap::from([("TEST_ENV_SNAPSHOT".to_string(), "worktree".to_string())]);
+        let rewritten = maybe_wrap_shell_lc_with_snapshot(
+            &command,
+            &session_shell,
+            dir.path(),
+            &explicit_env_overrides,
+        );
+        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_keeps_snapshot_path_without_override() {
+        let dir = tempdir().expect("create temp dir");
+        let snapshot_path = dir.path().join("snapshot.sh");
+        std::fs::write(
+            &snapshot_path,
+            "# Snapshot file\nexport PATH='/snapshot/bin'\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' \"$PATH\"".to_string(),
+        ];
+        let rewritten = maybe_wrap_shell_lc_with_snapshot(
+            &command,
+            &session_shell,
+            dir.path(),
+            &HashMap::new(),
+        );
+        let output = Command::new(&rewritten[0])
+            .args(&rewritten[1..])
+            .output()
+            .expect("run rewritten command");
+
+        assert!(output.status.success(), "command failed: {output:?}");
+        assert_eq!(String::from_utf8_lossy(&output.stdout), "/snapshot/bin");
+    }
+
+    #[test]
+    fn maybe_wrap_shell_lc_with_snapshot_applies_explicit_path_override() {
+        let dir = tempdir().expect("create temp dir");
+        let snapshot_path = dir.path().join("snapshot.sh");
+        std::fs::write(
+            &snapshot_path,
+            "# Snapshot file\nexport PATH='/snapshot/bin'\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' \"$PATH\"".to_string(),
+        ];
+        let explicit_env_overrides =
+            HashMap::from([("PATH".to_string(), "/worktree/bin".to_string())]);
+        let rewritten = maybe_wrap_shell_lc_with_snapshot(
+            &command,
+            &session_shell,
+            dir.path(),
+            &explicit_env_overrides,
+        );
+        let output = Command::new(&rewritten[0])
+            .args(&rewritten[1..])
+            .env("PATH", "/worktree/bin")
+            .output()
+            .expect("run rewritten command");
+
+        assert!(output.status.success(), "command failed: {output:?}");
+        assert_eq!(String::from_utf8_lossy(&output.stdout), "/worktree/bin");
+    }
+
+    #[test]
+    fn maybe_wrap_shell_lc_with_snapshot_does_not_embed_override_values_in_argv() {
+        let dir = tempdir().expect("create temp dir");
+        let snapshot_path = dir.path().join("snapshot.sh");
+        std::fs::write(
+            &snapshot_path,
+            "# Snapshot file\nexport OPENAI_API_KEY='snapshot-value'\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' \"$OPENAI_API_KEY\"".to_string(),
+        ];
+        let explicit_env_overrides = HashMap::from([(
+            "OPENAI_API_KEY".to_string(),
+            "super-secret-value".to_string(),
+        )]);
+        let rewritten = maybe_wrap_shell_lc_with_snapshot(
+            &command,
+            &session_shell,
+            dir.path(),
+            &explicit_env_overrides,
+        );
+
+        assert!(!rewritten[2].contains("super-secret-value"));
+        let output = Command::new(&rewritten[0])
+            .args(&rewritten[1..])
+            .env("OPENAI_API_KEY", "super-secret-value")
+            .output()
+            .expect("run rewritten command");
+        assert!(output.status.success(), "command failed: {output:?}");
+        assert_eq!(
+            String::from_utf8_lossy(&output.stdout),
+            "super-secret-value"
+        );
+    }
+
+    #[test]
+    fn maybe_wrap_shell_lc_with_snapshot_preserves_unset_override_variables() {
+        let dir = tempdir().expect("create temp dir");
+        let snapshot_path = dir.path().join("snapshot.sh");
+        std::fs::write(
+            &snapshot_path,
+            "# Snapshot file\nexport CODEX_TEST_UNSET_OVERRIDE='snapshot-value'\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(),
+            "if [ \"${CODEX_TEST_UNSET_OVERRIDE+x}\" = x ]; then printf 'set:%s' \"$CODEX_TEST_UNSET_OVERRIDE\"; else printf 'unset'; fi".to_string(),
+        ];
+        let explicit_env_overrides = HashMap::from([(
+            "CODEX_TEST_UNSET_OVERRIDE".to_string(),
+            "worktree-value".to_string(),
+        )]);
+        let rewritten = maybe_wrap_shell_lc_with_snapshot(
+            &command,
+            &session_shell,
+            dir.path(),
+            &explicit_env_overrides,
+        );
+
+        let output = Command::new(&rewritten[0])
+            .args(&rewritten[1..])
+            .env_remove("CODEX_TEST_UNSET_OVERRIDE")
+            .output()
+            .expect("run rewritten command");
+        assert!(output.status.success(), "command failed: {output:?}");
+        assert_eq!(String::from_utf8_lossy(&output.stdout), "unset");
+    }
 }

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

@@ -37,6 +37,7 @@ pub struct ShellRequest {
     pub cwd: PathBuf,
     pub timeout_ms: Option<u64>,
     pub env: std::collections::HashMap<String, String>,
+    pub explicit_env_overrides: std::collections::HashMap<String, String>,
     pub network: Option<NetworkProxy>,
     pub sandbox_permissions: SandboxPermissions,
     pub justification: Option<String>,
@@ -166,8 +167,12 @@ impl ToolRuntime<ShellRequest, ExecToolCallOutput> for ShellRuntime {
     ) -> Result<ExecToolCallOutput, ToolError> {
         let base_command = &req.command;
         let session_shell = ctx.session.user_shell();
-        let command =
-            maybe_wrap_shell_lc_with_snapshot(base_command, session_shell.as_ref(), &req.cwd);
+        let command = maybe_wrap_shell_lc_with_snapshot(
+            base_command,
+            session_shell.as_ref(),
+            &req.cwd,
+            &req.explicit_env_overrides,
+        );
         let command = if matches!(session_shell.shell_type, ShellType::PowerShell)
             && ctx.session.features().enabled(Feature::PowershellUtf8)
         {

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

@@ -41,6 +41,7 @@ pub struct UnifiedExecRequest {
     pub command: Vec<String>,
     pub cwd: PathBuf,
     pub env: HashMap<String, String>,
+    pub explicit_env_overrides: HashMap<String, String>,
     pub network: Option<NetworkProxy>,
     pub tty: bool,
     pub sandbox_permissions: SandboxPermissions,
@@ -170,8 +171,12 @@ impl<'a> ToolRuntime<UnifiedExecRequest, UnifiedExecProcess> for UnifiedExecRunt
     ) -> Result<UnifiedExecProcess, ToolError> {
         let base_command = &req.command;
         let session_shell = ctx.session.user_shell();
-        let command =
-            maybe_wrap_shell_lc_with_snapshot(base_command, session_shell.as_ref(), &req.cwd);
+        let command = maybe_wrap_shell_lc_with_snapshot(
+            base_command,
+            session_shell.as_ref(),
+            &req.cwd,
+            &req.explicit_env_overrides,
+        );
         let command = if matches!(session_shell.shell_type, ShellType::PowerShell)
             && ctx.session.features().enabled(Feature::PowershellUtf8)
         {

codex-rs/core/src/unified_exec/process_manager.rs

@@ -622,6 +622,7 @@ impl UnifiedExecProcessManager {
             command: request.command.clone(),
             cwd,
             env,
+            explicit_env_overrides: context.turn.shell_environment_policy.r#set.clone(),
             network: request.network.clone(),
             tty: request.tty,
             sandbox_permissions: request.sandbox_permissions,