diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index 35088efce..80112c61c 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/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(); diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index 67f4c0416..9093dfba4 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/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(), diff --git a/codex-rs/core/src/tools/runtimes/mod.rs b/codex-rs/core/src/tools/runtimes/mod.rs index b3d7b3292..92c8f3922 100644 --- a/codex-rs/core/src/tools/runtimes/mod.rs +++ b/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, ) -> Vec { + 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::(); - 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) { + let mut keys = explicit_env_overrides + .keys() + .filter(|key| is_valid_shell_variable_name(key)) + .collect::>(); + 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::>() + .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::>() + .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"); + } } diff --git a/codex-rs/core/src/tools/runtimes/shell.rs b/codex-rs/core/src/tools/runtimes/shell.rs index 3b7f970d7..3d7a61ef7 100644 --- a/codex-rs/core/src/tools/runtimes/shell.rs +++ b/codex-rs/core/src/tools/runtimes/shell.rs @@ -37,6 +37,7 @@ pub struct ShellRequest { pub cwd: PathBuf, pub timeout_ms: Option, pub env: std::collections::HashMap, + pub explicit_env_overrides: std::collections::HashMap, pub network: Option, pub sandbox_permissions: SandboxPermissions, pub justification: Option, @@ -166,8 +167,12 @@ impl ToolRuntime for ShellRuntime { ) -> Result { 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) { diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index 4c6c61107..89b159c5b 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -41,6 +41,7 @@ pub struct UnifiedExecRequest { pub command: Vec, pub cwd: PathBuf, pub env: HashMap, + pub explicit_env_overrides: HashMap, pub network: Option, pub tty: bool, pub sandbox_permissions: SandboxPermissions, @@ -170,8 +171,12 @@ impl<'a> ToolRuntime for UnifiedExecRunt ) -> Result { 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) { diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 197cda479..6896df14b 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/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, diff --git a/codex-rs/core/tests/suite/shell_snapshot.rs b/codex-rs/core/tests/suite/shell_snapshot.rs index 42a2107cd..b5007fd58 100644 --- a/codex-rs/core/tests/suite/shell_snapshot.rs +++ b/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, +} + async fn wait_for_snapshot(codex_home: &Path) -> Result { 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 { } } +async fn wait_for_file_contents(path: &Path) -> Result { + 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 { + 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 { - 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 { + 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 { #[allow(clippy::expect_used)] async fn run_shell_command_snapshot(command: &str) -> Result { - 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 { + 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 { }) } +#[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 { + 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?;