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,

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

@@ -20,6 +20,7 @@ use core_test_support::wait_for_event;
 use core_test_support::wait_for_event_match;
 use pretty_assertions::assert_eq;
 use serde_json::json;
+use std::collections::HashMap;
 use std::path::Path;
 use std::path::PathBuf;
 use tokio::fs;
@@ -36,6 +37,17 @@ struct SnapshotRun {
     codex_home: PathBuf,
 }
 
+const POLICY_PATH_FOR_TEST: &str = "/codex/policy/path";
+const SNAPSHOT_PATH_FOR_TEST: &str = "/codex/snapshot/path";
+const SNAPSHOT_MARKER_VAR: &str = "CODEX_SNAPSHOT_POLICY_MARKER";
+const SNAPSHOT_MARKER_VALUE: &str = "from_snapshot";
+const POLICY_SUCCESS_OUTPUT: &str = "policy-after-snapshot";
+
+#[derive(Debug, Default)]
+struct SnapshotRunOptions {
+    shell_environment_set: HashMap<String, String>,
+}
+
 async fn wait_for_snapshot(codex_home: &Path) -> Result<PathBuf> {
     let snapshot_dir = codex_home.join("shell_snapshots");
     let deadline = Instant::now() + Duration::from_secs(5);
@@ -54,12 +66,57 @@ async fn wait_for_snapshot(codex_home: &Path) -> Result<PathBuf> {
     }
 }
 
+async fn wait_for_file_contents(path: &Path) -> Result<String> {
+    let deadline = Instant::now() + Duration::from_secs(5);
+    loop {
+        match fs::read_to_string(path).await {
+            Ok(contents) => return Ok(contents),
+            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
+            Err(err) => return Err(err.into()),
+        }
+
+        if Instant::now() >= deadline {
+            anyhow::bail!("timed out waiting for file {}", path.display());
+        }
+
+        sleep(Duration::from_millis(25)).await;
+    }
+}
+
+fn policy_set_path_for_test() -> HashMap<String, String> {
+    HashMap::from([("PATH".to_string(), POLICY_PATH_FOR_TEST.to_string())])
+}
+
+fn snapshot_override_content_for_policy_test() -> String {
+    format!(
+        "# Snapshot file\nexport PATH='{SNAPSHOT_PATH_FOR_TEST}'\nexport {SNAPSHOT_MARKER_VAR}='{SNAPSHOT_MARKER_VALUE}'\n"
+    )
+}
+
+fn command_asserting_policy_after_snapshot() -> String {
+    format!(
+        "if [ \"${{{SNAPSHOT_MARKER_VAR}:-}}\" = \"{SNAPSHOT_MARKER_VALUE}\" ] && [ \"$PATH\" != \"{SNAPSHOT_PATH_FOR_TEST}\" ]; then case \":$PATH:\" in *\":{POLICY_PATH_FOR_TEST}:\"*) printf \"{POLICY_SUCCESS_OUTPUT}\" ;; *) printf \"path=%s marker=%s\" \"$PATH\" \"${{{SNAPSHOT_MARKER_VAR}:-missing}}\" ;; esac; else printf \"path=%s marker=%s\" \"$PATH\" \"${{{SNAPSHOT_MARKER_VAR}:-missing}}\"; fi"
+    )
+}
+
 #[allow(clippy::expect_used)]
 async fn run_snapshot_command(command: &str) -> Result<SnapshotRun> {
-    let builder = test_codex().with_config(|config| {
+    run_snapshot_command_with_options(command, SnapshotRunOptions::default()).await
+}
+
+#[allow(clippy::expect_used)]
+async fn run_snapshot_command_with_options(
+    command: &str,
+    options: SnapshotRunOptions,
+) -> Result<SnapshotRun> {
+    let SnapshotRunOptions {
+        shell_environment_set,
+    } = options;
+    let builder = test_codex().with_config(move |config| {
         config.use_experimental_unified_exec_tool = true;
         config.features.enable(Feature::UnifiedExec);
         config.features.enable(Feature::ShellSnapshot);
+        config.permissions.shell_environment_policy.r#set = shell_environment_set;
     });
     let harness = TestCodexHarness::with_builder(builder).await?;
     let args = json!({
@@ -132,8 +189,20 @@ async fn run_snapshot_command(command: &str) -> Result<SnapshotRun> {
 
 #[allow(clippy::expect_used)]
 async fn run_shell_command_snapshot(command: &str) -> Result<SnapshotRun> {
-    let builder = test_codex().with_config(|config| {
+    run_shell_command_snapshot_with_options(command, SnapshotRunOptions::default()).await
+}
+
+#[allow(clippy::expect_used)]
+async fn run_shell_command_snapshot_with_options(
+    command: &str,
+    options: SnapshotRunOptions,
+) -> Result<SnapshotRun> {
+    let SnapshotRunOptions {
+        shell_environment_set,
+    } = options;
+    let builder = test_codex().with_config(move |config| {
         config.features.enable(Feature::ShellSnapshot);
+        config.permissions.shell_environment_policy.r#set = shell_environment_set;
     });
     let harness = TestCodexHarness::with_builder(builder).await?;
     let args = json!({
@@ -204,6 +273,64 @@ async fn run_shell_command_snapshot(command: &str) -> Result<SnapshotRun> {
     })
 }
 
+#[allow(clippy::expect_used)]
+async fn run_tool_turn_on_harness(
+    harness: &TestCodexHarness,
+    prompt: &str,
+    call_id: &str,
+    tool_name: &str,
+    args: serde_json::Value,
+) -> Result<ExecCommandEndEvent> {
+    let responses = vec![
+        sse(vec![
+            ev_response_created("resp-1"),
+            ev_function_call(call_id, tool_name, &serde_json::to_string(&args)?),
+            ev_completed("resp-1"),
+        ]),
+        sse(vec![
+            ev_response_created("resp-2"),
+            ev_assistant_message("msg-1", "done"),
+            ev_completed("resp-2"),
+        ]),
+    ];
+    mount_sse_sequence(harness.server(), responses).await;
+
+    let test = harness.test();
+    let codex = test.codex.clone();
+    let session_model = test.session_configured.model.clone();
+    let cwd = test.cwd_path().to_path_buf();
+    codex
+        .submit(Op::UserTurn {
+            items: vec![UserInput::Text {
+                text: prompt.into(),
+                text_elements: Vec::new(),
+            }],
+            final_output_json_schema: None,
+            cwd,
+            approval_policy: AskForApproval::Never,
+            sandbox_policy: SandboxPolicy::DangerFullAccess,
+            model: session_model,
+            effort: None,
+            summary: ReasoningSummary::Auto,
+            collaboration_mode: None,
+            personality: None,
+        })
+        .await?;
+
+    wait_for_event_match(&codex, |ev| match ev {
+        EventMsg::ExecCommandBegin(ev) if ev.call_id == call_id => Some(ev.clone()),
+        _ => None,
+    })
+    .await;
+    let end = wait_for_event_match(&codex, |ev| match ev {
+        EventMsg::ExecCommandEnd(ev) if ev.call_id == call_id => Some(ev.clone()),
+        _ => None,
+    })
+    .await;
+    wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
+    Ok(end)
+}
+
 fn normalize_newlines(text: &str) -> String {
     text.replace("\r\n", "\n")
 }
@@ -260,6 +387,100 @@ async fn linux_shell_command_uses_shell_snapshot() -> Result<()> {
     Ok(())
 }
 
+#[cfg_attr(target_os = "windows", ignore)]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn shell_command_snapshot_preserves_shell_environment_policy_set() -> Result<()> {
+    let builder = test_codex().with_config(|config| {
+        config.features.enable(Feature::ShellSnapshot);
+        config.permissions.shell_environment_policy.r#set = policy_set_path_for_test();
+    });
+    let harness = TestCodexHarness::with_builder(builder).await?;
+    let codex_home = harness.test().home.path().to_path_buf();
+    run_tool_turn_on_harness(
+        &harness,
+        "warm up shell snapshot",
+        "shell-snapshot-policy-warmup",
+        "shell_command",
+        json!({
+            "command": "printf warmup",
+            "timeout_ms": 1_000,
+        }),
+    )
+    .await?;
+    let snapshot_path = wait_for_snapshot(&codex_home).await?;
+    fs::write(&snapshot_path, snapshot_override_content_for_policy_test()).await?;
+
+    let command = command_asserting_policy_after_snapshot();
+    let end = run_tool_turn_on_harness(
+        &harness,
+        "verify shell policy after snapshot",
+        "shell-snapshot-policy-assert",
+        "shell_command",
+        json!({
+            "command": command,
+            "timeout_ms": 1_000,
+        }),
+    )
+    .await?;
+
+    assert_eq!(
+        normalize_newlines(&end.stdout).trim(),
+        POLICY_SUCCESS_OUTPUT
+    );
+    assert_eq!(end.exit_code, 0);
+    assert!(snapshot_path.starts_with(codex_home));
+
+    Ok(())
+}
+
+#[cfg_attr(not(target_os = "linux"), ignore)]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn linux_unified_exec_snapshot_preserves_shell_environment_policy_set() -> Result<()> {
+    let builder = test_codex().with_config(|config| {
+        config.use_experimental_unified_exec_tool = true;
+        config.features.enable(Feature::UnifiedExec);
+        config.features.enable(Feature::ShellSnapshot);
+        config.permissions.shell_environment_policy.r#set = policy_set_path_for_test();
+    });
+    let harness = TestCodexHarness::with_builder(builder).await?;
+    let codex_home = harness.test().home.path().to_path_buf();
+    run_tool_turn_on_harness(
+        &harness,
+        "warm up unified exec shell snapshot",
+        "shell-snapshot-policy-warmup-exec",
+        "exec_command",
+        json!({
+            "cmd": "printf warmup",
+            "yield_time_ms": 1_000,
+        }),
+    )
+    .await?;
+    let snapshot_path = wait_for_snapshot(&codex_home).await?;
+    fs::write(&snapshot_path, snapshot_override_content_for_policy_test()).await?;
+
+    let command = command_asserting_policy_after_snapshot();
+    let end = run_tool_turn_on_harness(
+        &harness,
+        "verify unified exec policy after snapshot",
+        "shell-snapshot-policy-assert-exec",
+        "exec_command",
+        json!({
+            "cmd": command,
+            "yield_time_ms": 1_000,
+        }),
+    )
+    .await?;
+
+    assert_eq!(
+        normalize_newlines(&end.stdout).trim(),
+        POLICY_SUCCESS_OUTPUT
+    );
+    assert_eq!(end.exit_code, 0);
+    assert!(snapshot_path.starts_with(codex_home));
+
+    Ok(())
+}
+
 #[cfg_attr(target_os = "windows", ignore)]
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn shell_command_snapshot_still_intercepts_apply_patch() -> Result<()> {
@@ -316,7 +537,10 @@ async fn shell_command_snapshot_still_intercepts_apply_patch() -> Result<()> {
 
     wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
 
-    assert_eq!(fs::read_to_string(&target).await?, "hello from snapshot\n");
+    assert_eq!(
+        wait_for_file_contents(&target).await?,
+        "hello from snapshot\n"
+    );
 
     let snapshot_path = wait_for_snapshot(&codex_home).await?;
     let snapshot_content = fs::read_to_string(&snapshot_path).await?;