use anyhow::Result;use codex_features::Feature;use codex_protocol::models::PermissionProfile;use codex_protocol::protocol::AskForApproval;use codex_protocol::protocol::EventMsg;use codex_protocol::protocol::ExecCommandBeginEvent;use codex_protocol::protocol::ExecCommandEndEvent;use codex_protocol::protocol::Op;use codex_protocol::user_input::UserInput;use core_test_support::responses::ev_assistant_message;use core_test_support::responses::ev_completed;use core_test_support::responses::ev_function_call;use core_test_support::responses::ev_response_created;use core_test_support::responses::mount_sse_sequence;use core_test_support::responses::sse;use core_test_support::test_codex::TestCodexHarness;use core_test_support::test_codex::local_selections;use core_test_support::test_codex::test_codex;use core_test_support::test_codex::turn_permission_fields;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;use tokio::time::Duration;use tokio::time::Instant;use tokio::time::sleep;#[derive(Debug)]struct SnapshotRun {    begin: ExecCommandBeginEvent,    end: ExecCommandEndEvent,    snapshot_path: PathBuf,    snapshot_content: String,    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);    loop {        if let Ok(mut entries) = fs::read_dir(&snapshot_dir).await {            while let Some(entry) = entries.next_entry().await? {                let path = entry.path();                let Some(extension) = path.extension().and_then(|ext| ext.to_str()) else {                    continue;                };                if extension == "sh" || extension == "ps1" {                    return Ok(path);                }            }        }        if Instant::now() >= deadline {            anyhow::bail!("timed out waiting for shell snapshot");        }        sleep(Duration::from_millis(25)).await;    }}async fn wait_for_file_contents(path: &Path) -> Result<String> {    let deadline = Instant::now() + Duration::from_secs(15);    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"    )}async fn run_snapshot_command(command: &str) -> Result<SnapshotRun> {    run_snapshot_command_with_options(command, SnapshotRunOptions::default()).await}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)            .expect("test config should allow feature update");        config            .features            .enable(Feature::ShellSnapshot)            .expect("test config should allow feature update");        config.permissions.shell_environment_policy.r#set = shell_environment_set;    });    let harness = TestCodexHarness::with_builder(builder).await?;    let args = json!({        "cmd": command,        "yield_time_ms": 1000,    });    let call_id = "shell-snapshot-exec";    let responses = vec![        sse(vec![            ev_response_created("resp-1"),            ev_function_call(call_id, "exec_command", &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 codex_home = test.home.path().to_path_buf();    let session_model = test.session_configured.model.clone();    let cwd = test.config.cwd.clone();    let (sandbox_policy, permission_profile) =        turn_permission_fields(PermissionProfile::Disabled, cwd.as_path());    codex        .submit(Op::UserInput {            items: vec![UserInput::Text {                text: "run unified exec with shell snapshot".into(),                text_elements: Vec::new(),            }],            final_output_json_schema: None,            responsesapi_client_metadata: None,            additional_context: Default::default(),            thread_settings: codex_protocol::protocol::ThreadSettingsOverrides {                environments: Some(local_selections(cwd)),                approval_policy: Some(AskForApproval::Never),                sandbox_policy: Some(sandbox_policy),                permission_profile,                collaboration_mode: Some(codex_protocol::config_types::CollaborationMode {                    mode: codex_protocol::config_types::ModeKind::Default,                    settings: codex_protocol::config_types::Settings {                        model: session_model,                        reasoning_effort: None,                        developer_instructions: None,                    },                }),                ..Default::default()            },        })        .await?;    let begin = wait_for_event_match(&codex, |ev| match ev {        EventMsg::ExecCommandBegin(ev) if ev.call_id == call_id => Some(ev.clone()),        _ => None,    })    .await;    let snapshot_path = wait_for_snapshot(&codex_home).await?;    let snapshot_content = fs::read_to_string(&snapshot_path).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(SnapshotRun {        begin,        end,        snapshot_path,        snapshot_content,        codex_home,    })}async fn run_shell_command_snapshot(command: &str) -> Result<SnapshotRun> {    run_shell_command_snapshot_with_options(command, SnapshotRunOptions::default()).await}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)            .expect("test config should allow feature update");        config.permissions.shell_environment_policy.r#set = shell_environment_set;    });    let harness = TestCodexHarness::with_builder(builder).await?;    let args = json!({        "command": command,        "timeout_ms": 1000,    });    let call_id = "shell-snapshot-command";    let responses = vec![        sse(vec![            ev_response_created("resp-1"),            ev_function_call(call_id, "shell_command", &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 codex_home = test.home.path().to_path_buf();    let session_model = test.session_configured.model.clone();    let cwd = test.config.cwd.clone();    let (sandbox_policy, permission_profile) =        turn_permission_fields(PermissionProfile::Disabled, cwd.as_path());    codex        .submit(Op::UserInput {            items: vec![UserInput::Text {                text: "run shell_command with shell snapshot".into(),                text_elements: Vec::new(),            }],            final_output_json_schema: None,            responsesapi_client_metadata: None,            additional_context: Default::default(),            thread_settings: codex_protocol::protocol::ThreadSettingsOverrides {                environments: Some(local_selections(cwd)),                approval_policy: Some(AskForApproval::Never),                sandbox_policy: Some(sandbox_policy),                permission_profile,                collaboration_mode: Some(codex_protocol::config_types::CollaborationMode {                    mode: codex_protocol::config_types::ModeKind::Default,                    settings: codex_protocol::config_types::Settings {                        model: session_model,                        reasoning_effort: None,                        developer_instructions: None,                    },                }),                ..Default::default()            },        })        .await?;    let begin = wait_for_event_match(&codex, |ev| match ev {        EventMsg::ExecCommandBegin(ev) if ev.call_id == call_id => Some(ev.clone()),        _ => None,    })    .await;    let snapshot_path = wait_for_snapshot(&codex_home).await?;    let snapshot_content = fs::read_to_string(&snapshot_path).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(SnapshotRun {        begin,        end,        snapshot_path,        snapshot_content,        codex_home,    })}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.config.cwd.clone();    let (sandbox_policy, permission_profile) =        turn_permission_fields(PermissionProfile::Disabled, cwd.as_path());    codex        .submit(Op::UserInput {            items: vec![UserInput::Text {                text: prompt.into(),                text_elements: Vec::new(),            }],            final_output_json_schema: None,            responsesapi_client_metadata: None,            additional_context: Default::default(),            thread_settings: codex_protocol::protocol::ThreadSettingsOverrides {                environments: Some(local_selections(cwd)),                approval_policy: Some(AskForApproval::Never),                sandbox_policy: Some(sandbox_policy),                permission_profile,                collaboration_mode: Some(codex_protocol::config_types::CollaborationMode {                    mode: codex_protocol::config_types::ModeKind::Default,                    settings: codex_protocol::config_types::Settings {                        model: session_model,                        reasoning_effort: None,                        developer_instructions: None,                    },                }),                ..Default::default()            },        })        .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")}fn assert_posix_snapshot_sections(snapshot: &str) {    assert!(snapshot.contains("# Snapshot file"));    assert!(snapshot.contains("aliases "));    assert!(snapshot.contains("exports "));    assert!(snapshot.contains("setopts "));    assert!(        snapshot.contains("PATH"),        "snapshot should include PATH exports; snapshot={snapshot:?}"    );}#[cfg_attr(not(target_os = "linux"), ignore)]#[tokio::test(flavor = "multi_thread", worker_threads = 2)]async fn linux_unified_exec_uses_shell_snapshot() -> Result<()> {    let command = "echo snapshot-linux";    let run = run_snapshot_command(command).await?;    let stdout = normalize_newlines(&run.end.stdout);    assert_eq!(run.begin.command.get(1).map(String::as_str), Some("-lc"));    assert_eq!(run.begin.command.get(2).map(String::as_str), Some(command));    assert_eq!(run.begin.command.len(), 3);    assert!(run.snapshot_path.starts_with(&run.codex_home));    assert_posix_snapshot_sections(&run.snapshot_content);    assert_eq!(run.end.exit_code, 0);    assert!(        stdout.contains("snapshot-linux"),        "stdout should contain snapshot marker; stdout={stdout:?}"    );    Ok(())}#[cfg_attr(target_os = "windows", ignore)]#[tokio::test(flavor = "multi_thread", worker_threads = 2)]async fn linux_shell_command_uses_shell_snapshot() -> Result<()> {    let command = "echo shell-command-snapshot-linux";    let run = run_shell_command_snapshot(command).await?;    assert_eq!(run.begin.command.get(1).map(String::as_str), Some("-lc"));    assert_eq!(run.begin.command.get(2).map(String::as_str), Some(command));    assert_eq!(run.begin.command.len(), 3);    assert!(run.snapshot_path.starts_with(&run.codex_home));    assert_posix_snapshot_sections(&run.snapshot_content);    assert_eq!(        normalize_newlines(&run.end.stdout).trim(),        "shell-command-snapshot-linux"    );    assert_eq!(run.end.exit_code, 0);    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)            .expect("test config should allow feature update");        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)            .expect("test config should allow feature update");        config            .features            .enable(Feature::ShellSnapshot)            .expect("test config should allow feature update");        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<()> {    let builder = test_codex().with_config(|config| {        config            .features            .enable(Feature::ShellSnapshot)            .expect("test config should allow feature update");    });    let harness = TestCodexHarness::with_builder(builder).await?;    let test = harness.test();    let codex = test.codex.clone();    let cwd = test.config.cwd.clone();    let codex_home = test.home.path().to_path_buf();    let target = cwd.join("snapshot-apply.txt");    let script = "apply_patch <<'EOF'\n*** Begin Patch\n*** Add File: snapshot-apply.txt\n+hello from snapshot\n*** End Patch\nEOF\n";    let args = json!({        "command": script,        // Keep this above the default because intercepted apply_patch still        // performs filesystem work that can be slow in Bazel macOS test        // environments.        "timeout_ms": 5_000,    });    let call_id = "shell-snapshot-apply-patch";    let responses = vec![        sse(vec![            ev_response_created("resp-1"),            ev_function_call(call_id, "shell_command", &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 model = test.session_configured.model.clone();    let (sandbox_policy, permission_profile) =        turn_permission_fields(PermissionProfile::Disabled, cwd.as_path());    codex        .submit(Op::UserInput {            items: vec![UserInput::Text {                text: "apply patch via shell_command with snapshot".into(),                text_elements: Vec::new(),            }],            final_output_json_schema: None,            responsesapi_client_metadata: None,            additional_context: Default::default(),            thread_settings: codex_protocol::protocol::ThreadSettingsOverrides {                environments: Some(local_selections(cwd.clone())),                approval_policy: Some(AskForApproval::Never),                sandbox_policy: Some(sandbox_policy),                permission_profile,                collaboration_mode: Some(codex_protocol::config_types::CollaborationMode {                    mode: codex_protocol::config_types::ModeKind::Default,                    settings: codex_protocol::config_types::Settings {                        model,                        reasoning_effort: None,                        developer_instructions: None,                    },                }),                ..Default::default()            },        })        .await?;    let snapshot_path = wait_for_snapshot(&codex_home).await?;    let snapshot_content = fs::read_to_string(&snapshot_path).await?;    assert_posix_snapshot_sections(&snapshot_content);    let mut saw_patch_begin = false;    let mut patch_end = None;    wait_for_event(&codex, |ev| match ev {        EventMsg::PatchApplyBegin(begin) if begin.call_id == call_id => {            saw_patch_begin = true;            false        }        EventMsg::PatchApplyEnd(end) if end.call_id == call_id => {            patch_end = Some(end.clone());            false        }        EventMsg::TurnComplete(_) => true,        _ => false,    })    .await;    assert!(        saw_patch_begin,        "expected apply_patch to emit PatchApplyBegin"    );    let patch_end = patch_end.expect("expected apply_patch to emit PatchApplyEnd");    assert!(        patch_end.success,        "expected apply_patch to finish successfully: stdout={:?} stderr={:?}",        patch_end.stdout, patch_end.stderr,    );    assert_eq!(        wait_for_file_contents(&target).await?,        "hello from snapshot\n"    );    Ok(())}#[cfg_attr(target_os = "windows", ignore)]#[tokio::test(flavor = "multi_thread", worker_threads = 2)]async fn shell_snapshot_deleted_after_shutdown_with_skills() -> Result<()> {    let builder = test_codex().with_config(|config| {        config            .features            .enable(Feature::ShellSnapshot)            .expect("test config should allow feature update");    });    let harness = TestCodexHarness::with_builder(builder).await?;    let home = harness.test().home.clone();    let codex_home = home.path().to_path_buf();    let codex = harness.test().codex.clone();    let snapshot_path = wait_for_snapshot(&codex_home).await?;    assert!(snapshot_path.exists());    codex.submit(Op::Shutdown {}).await?;    wait_for_event(&codex, |ev| matches!(ev, EventMsg::ShutdownComplete)).await;    drop(codex);    drop(harness);    sleep(Duration::from_millis(150)).await;    assert_eq!(        snapshot_path.exists(),        false,        "snapshot should be removed after shutdown"    );    Ok(())}#[cfg_attr(not(target_os = "macos"), ignore)]#[cfg_attr(    target_os = "macos",    ignore = "requires unrestricted networking on macOS")]#[tokio::test(flavor = "multi_thread", worker_threads = 2)]async fn macos_unified_exec_uses_shell_snapshot() -> Result<()> {    let command = "echo snapshot-macos";    let run = run_snapshot_command(command).await?;    let shell_path = run        .begin        .command        .first()        .expect("shell path recorded")        .clone();    assert_eq!(run.begin.command.get(1).map(String::as_str), Some("-c"));    assert_eq!(        run.begin.command.get(2).map(String::as_str),        Some(". \"$0\" && exec \"$@\"")    );    assert_eq!(run.begin.command.get(4), Some(&shell_path));    assert_eq!(run.begin.command.get(5).map(String::as_str), Some("-c"));    assert_eq!(run.begin.command.last(), Some(&command.to_string()));    assert!(run.snapshot_path.starts_with(&run.codex_home));    assert_posix_snapshot_sections(&run.snapshot_content);    assert_eq!(normalize_newlines(&run.end.stdout).trim(), "snapshot-macos");    assert_eq!(run.end.exit_code, 0);    Ok(())}// #[cfg_attr(not(target_os = "windows"), ignore)]#[ignore]#[tokio::test(flavor = "multi_thread", worker_threads = 2)]async fn windows_unified_exec_uses_shell_snapshot() -> Result<()> {    let command = "Write-Output snapshot-windows";    let run = run_snapshot_command(command).await?;    let snapshot_index = run        .begin        .command        .iter()        .position(|arg| arg.contains("shell_snapshots"))        .expect("snapshot argument exists");    assert!(run.begin.command.iter().any(|arg| arg == "-NoProfile"));    assert!(        run.begin            .command            .iter()            .any(|arg| arg == "param($snapshot) . $snapshot; & @args")    );    assert!(snapshot_index > 0);    assert_eq!(run.begin.command.last(), Some(&command.to_string()));    assert!(run.snapshot_path.starts_with(&run.codex_home));    assert!(run.snapshot_content.contains("# Snapshot file"));    assert!(run.snapshot_content.contains("# aliases "));    assert!(run.snapshot_content.contains("# exports "));    assert_eq!(        normalize_newlines(&run.end.stdout).trim(),        "snapshot-windows"    );    assert_eq!(run.end.exit_code, 0);    Ok(())}