fix(core): wait for applicable shell snapshots before execution

Co-authored-by: Codex noreply@openai.com

codex-rs/core/src/session/session.rs

@@ -842,7 +842,8 @@ impl Session {
             // Create the mutable state for the Session.
             let shell_snapshot_tx = if config.features.enabled(Feature::ShellSnapshot) {
                 if let Some(snapshot) = session_configuration.inherited_shell_snapshot.clone() {
-                    let (tx, rx) = watch::channel(Some(snapshot));
+                    let (tx, rx) =
+                        watch::channel(crate::shell_snapshot::ShellSnapshotState::Ready(snapshot));
                     default_shell.shell_snapshot = rx;
                     tx
                 } else {
@@ -856,7 +857,8 @@ impl Session {
                     )
                 }
             } else {
-                let (tx, rx) = watch::channel(None);
+                let (tx, rx) =
+                    watch::channel(crate::shell_snapshot::ShellSnapshotState::Disabled);
                 default_shell.shell_snapshot = rx;
                 tx
             };

codex-rs/core/src/session/tests.rs

@@ -4539,7 +4539,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
         })),
         rollout_thread_trace: codex_rollout_trace::ThreadTraceContext::disabled(),
         user_shell: Arc::new(default_user_shell()),
-        shell_snapshot_tx: watch::channel(None).0,
+        shell_snapshot_tx: watch::channel(crate::shell_snapshot::ShellSnapshotState::Disabled).0,
         show_raw_agent_reasoning: config.show_raw_agent_reasoning,
         exec_policy,
         auth_manager: auth_manager.clone(),
@@ -6369,7 +6369,7 @@ where
         })),
         rollout_thread_trace: codex_rollout_trace::ThreadTraceContext::disabled(),
         user_shell: Arc::new(default_user_shell()),
-        shell_snapshot_tx: watch::channel(None).0,
+        shell_snapshot_tx: watch::channel(crate::shell_snapshot::ShellSnapshotState::Disabled).0,
         show_raw_agent_reasoning: config.show_raw_agent_reasoning,
         exec_policy,
         auth_manager: Arc::clone(&auth_manager),

codex-rs/core/src/shell.rs

@@ -1,5 +1,9 @@
+use crate::path_utils;
 use crate::shell_detect::detect_shell_type;
 use crate::shell_snapshot::ShellSnapshot;
+use crate::shell_snapshot::ShellSnapshotFailure;
+use crate::shell_snapshot::ShellSnapshotState;
+use codex_utils_absolute_path::AbsolutePathBuf;
 use serde::Deserialize;
 use serde::Serialize;
 use std::path::PathBuf;
@@ -24,7 +28,7 @@ pub struct Shell {
         skip_deserializing,
         default = "empty_shell_snapshot_receiver"
     )]
-    pub(crate) shell_snapshot: watch::Receiver<Option<Arc<ShellSnapshot>>>,
+    pub(crate) shell_snapshot: watch::Receiver<ShellSnapshotState>,
 }
 
 impl Shell {
@@ -71,12 +75,65 @@ impl Shell {
 
     /// Return the shell snapshot if existing.
     pub fn shell_snapshot(&self) -> Option<Arc<ShellSnapshot>> {
-        self.shell_snapshot.borrow().clone()
+        match self.shell_snapshot.borrow().clone() {
+            ShellSnapshotState::Ready(snapshot) => Some(snapshot),
+            ShellSnapshotState::Disabled
+            | ShellSnapshotState::Pending { .. }
+            | ShellSnapshotState::Failed { .. } => None,
+        }
+    }
+
+    pub(crate) async fn wait_for_shell_snapshot(
+        &self,
+        cwd: &AbsolutePathBuf,
+    ) -> Result<Option<Arc<ShellSnapshot>>, ShellSnapshotFailure> {
+        let mut receiver = self.shell_snapshot.clone();
+        loop {
+            match receiver.borrow().clone() {
+                ShellSnapshotState::Disabled => return Ok(None),
+                ShellSnapshotState::Pending { cwd: snapshot_cwd } => {
+                    if !path_utils::paths_match_after_normalization(
+                        snapshot_cwd.as_path(),
+                        cwd.as_path(),
+                    ) {
+                        return Ok(None);
+                    }
+                }
+                ShellSnapshotState::Ready(snapshot) => {
+                    return if path_utils::paths_match_after_normalization(
+                        snapshot.cwd.as_path(),
+                        cwd.as_path(),
+                    ) {
+                        Ok(Some(snapshot))
+                    } else {
+                        Ok(None)
+                    };
+                }
+                ShellSnapshotState::Failed {
+                    cwd: snapshot_cwd,
+                    failure_reason,
+                } => {
+                    return if path_utils::paths_match_after_normalization(
+                        snapshot_cwd.as_path(),
+                        cwd.as_path(),
+                    ) {
+                        Err(failure_reason)
+                    } else {
+                        Ok(None)
+                    };
+                }
+            }
+
+            receiver
+                .changed()
+                .await
+                .map_err(|_| ShellSnapshotFailure::TaskEnded)?;
+        }
     }
 }
 
-pub(crate) fn empty_shell_snapshot_receiver() -> watch::Receiver<Option<Arc<ShellSnapshot>>> {
-    let (_tx, rx) = watch::channel(None);
+pub(crate) fn empty_shell_snapshot_receiver() -> watch::Receiver<ShellSnapshotState> {
+    let (_tx, rx) = watch::channel(ShellSnapshotState::Disabled);
     rx
 }
 

codex-rs/core/src/shell_snapshot.rs

@@ -30,21 +30,65 @@ pub struct ShellSnapshot {
     pub cwd: AbsolutePathBuf,
 }
 
-const SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(10);
+/// Why a snapshot for a specific shell working directory could not be produced.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum ShellSnapshotFailure {
+    WriteFailed,
+    ValidationFailed,
+    TaskEnded,
+}
+
+impl ShellSnapshotFailure {
+    fn telemetry_value(self) -> &'static str {
+        match self {
+            Self::WriteFailed => "write_failed",
+            Self::ValidationFailed => "validation_failed",
+            Self::TaskEnded => "task_ended",
+        }
+    }
+}
+
+impl std::fmt::Display for ShellSnapshotFailure {
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::WriteFailed => formatter.write_str("shell environment capture failed"),
+            Self::ValidationFailed => formatter.write_str("captured shell environment was invalid"),
+            Self::TaskEnded => formatter.write_str("shell environment capture task ended"),
+        }
+    }
+}
+
+/// Snapshot readiness for the session shell, scoped to the working directory it captures.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub(crate) enum ShellSnapshotState {
+    Disabled,
+    Pending {
+        cwd: AbsolutePathBuf,
+    },
+    Ready(Arc<ShellSnapshot>),
+    Failed {
+        cwd: AbsolutePathBuf,
+        failure_reason: ShellSnapshotFailure,
+    },
+}
+
+const SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(60);
 const SNAPSHOT_RETENTION: Duration = Duration::from_secs(60 * 60 * 24 * 3); // 3 days retention.
 const SNAPSHOT_DIR: &str = "shell_snapshots";
 const EXCLUDED_EXPORT_VARS: &[&str] = &["PWD", "OLDPWD"];
 
 impl ShellSnapshot {
-    pub fn start_snapshotting(
+    pub(crate) fn start_snapshotting(
         codex_home: AbsolutePathBuf,
         session_id: ThreadId,
         session_cwd: AbsolutePathBuf,
         shell: &mut Shell,
         session_telemetry: SessionTelemetry,
         state_db: Option<StateDbHandle>,
-    ) -> watch::Sender<Option<Arc<ShellSnapshot>>> {
-        let (shell_snapshot_tx, shell_snapshot_rx) = watch::channel(None);
+    ) -> watch::Sender<ShellSnapshotState> {
+        let (shell_snapshot_tx, shell_snapshot_rx) = watch::channel(ShellSnapshotState::Pending {
+            cwd: session_cwd.clone(),
+        });
         shell.shell_snapshot = shell_snapshot_rx;
 
         Self::spawn_snapshot_task(
@@ -60,15 +104,18 @@ impl ShellSnapshot {
         shell_snapshot_tx
     }
 
-    pub fn refresh_snapshot(
+    pub(crate) fn refresh_snapshot(
         codex_home: AbsolutePathBuf,
         session_id: ThreadId,
         session_cwd: AbsolutePathBuf,
         shell: Shell,
-        shell_snapshot_tx: watch::Sender<Option<Arc<ShellSnapshot>>>,
+        shell_snapshot_tx: watch::Sender<ShellSnapshotState>,
         session_telemetry: SessionTelemetry,
         state_db: Option<StateDbHandle>,
     ) {
+        let _ = shell_snapshot_tx.send_replace(ShellSnapshotState::Pending {
+            cwd: session_cwd.clone(),
+        });
         Self::spawn_snapshot_task(
             codex_home,
             session_id,
@@ -85,7 +132,7 @@ impl ShellSnapshot {
         session_id: ThreadId,
         session_cwd: AbsolutePathBuf,
         snapshot_shell: Shell,
-        shell_snapshot_tx: watch::Sender<Option<Arc<ShellSnapshot>>>,
+        shell_snapshot_tx: watch::Sender<ShellSnapshotState>,
         session_telemetry: SessionTelemetry,
         state_db: Option<StateDbHandle>,
     ) {
@@ -107,10 +154,27 @@ impl ShellSnapshot {
                 let _ = timer.map(|timer| timer.record(&[("success", success_tag)]));
                 let mut counter_tags = vec![("success", success_tag)];
                 if let Some(failure_reason) = snapshot.as_ref().err() {
-                    counter_tags.push(("failure_reason", *failure_reason));
+                    counter_tags.push(("failure_reason", failure_reason.telemetry_value()));
                 }
                 session_telemetry.counter("codex.shell_snapshot", /*inc*/ 1, &counter_tags);
-                let _ = shell_snapshot_tx.send(snapshot.ok());
+                let completed_state = match snapshot {
+                    Ok(snapshot) => ShellSnapshotState::Ready(snapshot),
+                    Err(failure_reason) => ShellSnapshotState::Failed {
+                        cwd: session_cwd.clone(),
+                        failure_reason,
+                    },
+                };
+                let _ = shell_snapshot_tx.send_if_modified(|state| {
+                    if matches!(
+                        state,
+                        ShellSnapshotState::Pending { cwd } if cwd == &session_cwd
+                    ) {
+                        *state = completed_state;
+                        true
+                    } else {
+                        false
+                    }
+                });
             }
             .instrument(snapshot_span),
         );
@@ -122,7 +186,7 @@ impl ShellSnapshot {
         session_cwd: &AbsolutePathBuf,
         shell: &Shell,
         state_db: Option<StateDbHandle>,
-    ) -> std::result::Result<Self, &'static str> {
+    ) -> std::result::Result<Self, ShellSnapshotFailure> {
         // File to store the snapshot
         let extension = match shell.shell_type {
             ShellType::PowerShell => "ps1",
@@ -158,7 +222,7 @@ impl ShellSnapshot {
                 "Failed to create shell snapshot for {}: {err:?}",
                 shell.name()
             );
-            return Err("write_failed");
+            return Err(ShellSnapshotFailure::WriteFailed);
         }
         tracing::info!(
             "Shell snapshot successfully created: {}",
@@ -168,13 +232,13 @@ impl ShellSnapshot {
         if let Err(err) = validate_snapshot(shell, &temp_path, session_cwd).await {
             tracing::error!("Shell snapshot validation failed: {err:?}");
             remove_snapshot_file(&temp_path).await;
-            return Err("validation_failed");
+            return Err(ShellSnapshotFailure::ValidationFailed);
         }
 
         if let Err(err) = fs::rename(&temp_path, &path).await {
             tracing::warn!("Failed to finalize shell snapshot: {err:?}");
             remove_snapshot_file(&temp_path).await;
-            return Err("write_failed");
+            return Err(ShellSnapshotFailure::WriteFailed);
         }
 
         Ok(Self {

codex-rs/core/src/shell_snapshot_tests.rs

@@ -1,4 +1,5 @@
 use super::*;
+use codex_protocol::protocol::SessionSource;
 use core_test_support::PathBufExt;
 use core_test_support::PathExt;
 use pretty_assertions::assert_eq;
@@ -217,6 +218,50 @@ async fn try_new_creates_and_deletes_snapshot_file() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn refresh_snapshot_marks_new_cwd_pending_before_recapture() -> Result<()> {
+    let dir = tempdir()?;
+    let previous_cwd = dir.path().join("previous").abs();
+    let next_cwd = dir.path().join("next").abs();
+    let (shell_snapshot_tx, shell_snapshot_rx) =
+        watch::channel(ShellSnapshotState::Ready(Arc::new(ShellSnapshot {
+            path: dir.path().join("previous.sh").abs(),
+            cwd: previous_cwd,
+        })));
+    let shell = Shell {
+        shell_type: ShellType::Bash,
+        shell_path: PathBuf::from("/bin/bash"),
+        shell_snapshot: shell_snapshot_rx.clone(),
+    };
+
+    ShellSnapshot::refresh_snapshot(
+        dir.path().abs(),
+        ThreadId::new(),
+        next_cwd.clone(),
+        shell,
+        shell_snapshot_tx,
+        SessionTelemetry::new(
+            ThreadId::new(),
+            "gpt-test",
+            "gpt-test",
+            /*account_id*/ None,
+            /*account_email*/ None,
+            /*auth_mode*/ None,
+            "test-originator".to_string(),
+            /*log_user_prompts*/ false,
+            "test-terminal".to_string(),
+            SessionSource::Cli,
+        ),
+        /*state_db*/ None,
+    );
+
+    assert_eq!(
+        shell_snapshot_rx.borrow().clone(),
+        ShellSnapshotState::Pending { cwd: next_cwd }
+    );
+    Ok(())
+}
+
 #[cfg(unix)]
 #[tokio::test]
 async fn try_new_uses_distinct_generation_paths() -> Result<()> {

codex-rs/core/src/state/service.rs

@@ -50,7 +50,7 @@ pub(crate) struct SessionServices {
     pub(crate) hooks: ArcSwap<Hooks>,
     pub(crate) rollout_thread_trace: ThreadTraceContext,
     pub(crate) user_shell: Arc<crate::shell::Shell>,
-    pub(crate) shell_snapshot_tx: watch::Sender<Option<Arc<crate::shell_snapshot::ShellSnapshot>>>,
+    pub(crate) shell_snapshot_tx: watch::Sender<crate::shell_snapshot::ShellSnapshotState>,
     pub(crate) show_raw_agent_reasoning: bool,
     pub(crate) exec_policy: Arc<ExecPolicyManager>,
     pub(crate) auth_manager: Arc<AuthManager>,

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

@@ -20,9 +20,12 @@ use crate::session::TurnInput;
 use crate::session::turn_context::TurnContext;
 use crate::state::TaskKind;
 use crate::tools::format_exec_output_str;
+use crate::tools::runtimes::await_shell_snapshot_for_command;
 use crate::tools::runtimes::maybe_wrap_shell_lc_with_snapshot;
+use crate::tools::runtimes::shell_snapshot_failure_message;
 use crate::turn_timing::now_unix_timestamp_ms;
 use crate::user_shell_command::user_shell_command_record_item;
+use codex_protocol::error::CodexErr;
 use codex_protocol::exec_output::ExecToolCallOutput;
 use codex_protocol::exec_output::StreamOutput;
 use codex_protocol::protocol::EventMsg;
@@ -146,15 +149,6 @@ pub(crate) async fn execute_user_shell_command(
             exec_env_map.remove(PROXY_GIT_SSH_COMMAND_ENV_KEY);
         }
     }
-    let exec_command = maybe_wrap_shell_lc_with_snapshot(
-        &display_command,
-        session_shell.as_ref(),
-        #[allow(deprecated)]
-        &turn_context.cwd,
-        &turn_context.shell_environment_policy.r#set,
-        &exec_env_map,
-    );
-
     let call_id = Uuid::new_v4().to_string();
     let raw_command = command;
     #[allow(deprecated)]
@@ -178,42 +172,61 @@ pub(crate) async fn execute_user_shell_command(
         )
         .await;
 
-    let permission_profile = PermissionProfile::Disabled;
-    let exec_env = ExecRequest {
-        command: exec_command.clone(),
-        cwd: cwd.clone(),
-        env: exec_env_map,
-        exec_server_env_config: None,
-        // `/shell` is the explicit full-access escape hatch, so it must not
-        // inherit a managed proxy from the surrounding session or turn.
-        network: None,
-        // TODO(zhao-oai): Now that we have ExecExpiration::Cancellation, we
-        // should use that instead of an "arbitrarily large" timeout here.
-        expiration: USER_SHELL_TIMEOUT_MS.into(),
-        capture_policy: ExecCapturePolicy::ShellTool,
-        sandbox: SandboxType::None,
-        windows_sandbox_policy_cwd: cwd.clone(),
-        windows_sandbox_level: turn_context.windows_sandbox_level,
-        windows_sandbox_private_desktop: turn_context
-            .config
-            .permissions
-            .windows_sandbox_private_desktop,
-        permission_profile: permission_profile.clone(),
-        file_system_sandbox_policy: permission_profile.file_system_sandbox_policy(),
-        network_sandbox_policy: permission_profile.network_sandbox_policy(),
-        windows_sandbox_filesystem_overrides: None,
-        arg0: None,
-    };
+    let exec_result =
+        match await_shell_snapshot_for_command(&display_command, session_shell.as_ref(), &cwd)
+            .or_cancel(&cancellation_token)
+            .await
+        {
+            Ok(Ok(())) => {
+                let exec_command = maybe_wrap_shell_lc_with_snapshot(
+                    &display_command,
+                    session_shell.as_ref(),
+                    #[allow(deprecated)]
+                    &turn_context.cwd,
+                    &turn_context.shell_environment_policy.r#set,
+                    &exec_env_map,
+                );
+                let permission_profile = PermissionProfile::Disabled;
+                let exec_env = ExecRequest {
+                    command: exec_command,
+                    cwd: cwd.clone(),
+                    env: exec_env_map,
+                    exec_server_env_config: None,
+                    // `/shell` is the explicit full-access escape hatch, so it must not
+                    // inherit a managed proxy from the surrounding session or turn.
+                    network: None,
+                    // TODO(zhao-oai): Now that we have ExecExpiration::Cancellation, we
+                    // should use that instead of an "arbitrarily large" timeout here.
+                    expiration: USER_SHELL_TIMEOUT_MS.into(),
+                    capture_policy: ExecCapturePolicy::ShellTool,
+                    sandbox: SandboxType::None,
+                    windows_sandbox_policy_cwd: cwd.clone(),
+                    windows_sandbox_level: turn_context.windows_sandbox_level,
+                    windows_sandbox_private_desktop: turn_context
+                        .config
+                        .permissions
+                        .windows_sandbox_private_desktop,
+                    permission_profile: permission_profile.clone(),
+                    file_system_sandbox_policy: permission_profile.file_system_sandbox_policy(),
+                    network_sandbox_policy: permission_profile.network_sandbox_policy(),
+                    windows_sandbox_filesystem_overrides: None,
+                    arg0: None,
+                };
+                let stdout_stream = Some(StdoutStream {
+                    sub_id: turn_context.sub_id.clone(),
+                    call_id: call_id.clone(),
+                    tx_event: session.get_tx_event(),
+                });
 
-    let stdout_stream = Some(StdoutStream {
-        sub_id: turn_context.sub_id.clone(),
-        call_id: call_id.clone(),
-        tx_event: session.get_tx_event(),
-    });
-
-    let exec_result = execute_exec_request(exec_env, stdout_stream, /*after_spawn*/ None)
-        .or_cancel(&cancellation_token)
-        .await;
+                execute_exec_request(exec_env, stdout_stream, /*after_spawn*/ None)
+                    .or_cancel(&cancellation_token)
+                    .await
+            }
+            Ok(Err(failure_reason)) => Ok(Err(CodexErr::Io(std::io::Error::other(
+                shell_snapshot_failure_message(failure_reason),
+            )))),
+            Err(CancelErr::Cancelled) => Err(CancelErr::Cancelled),
+        };
 
     match exec_result {
         Err(CancelErr::Cancelled) => {

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

@@ -12,6 +12,7 @@ use crate::session::tests::make_session_and_context;
 use crate::shell::Shell;
 use crate::shell::ShellType;
 use crate::shell_snapshot::ShellSnapshot;
+use crate::shell_snapshot::ShellSnapshotState;
 use crate::tools::context::FunctionToolOutput;
 use crate::tools::context::ToolCallSource;
 use crate::tools::context::ToolInvocation;
@@ -128,10 +129,11 @@ async fn shell_command_handler_to_exec_params_uses_session_shell_and_turn_contex
 
 #[test]
 fn shell_command_handler_respects_explicit_login_flag() {
-    let (_tx, shell_snapshot) = watch::channel(Some(Arc::new(ShellSnapshot {
-        path: test_path_buf("/tmp/snapshot.sh").abs(),
-        cwd: test_path_buf("/tmp").abs(),
-    })));
+    let (_tx, shell_snapshot) =
+        watch::channel(ShellSnapshotState::Ready(Arc::new(ShellSnapshot {
+            path: test_path_buf("/tmp/snapshot.sh").abs(),
+            cwd: test_path_buf("/tmp").abs(),
+        })));
     let shell = Shell {
         shell_type: ShellType::Bash,
         shell_path: PathBuf::from("/bin/bash"),

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

@@ -9,6 +9,7 @@ use crate::path_utils;
 use crate::sandboxing::SandboxPermissions;
 use crate::shell::Shell;
 use crate::shell::ShellType;
+use crate::shell_snapshot::ShellSnapshotFailure;
 use crate::tools::sandboxing::ToolError;
 #[cfg(target_os = "macos")]
 use codex_network_proxy::CODEX_PROXY_GIT_SSH_COMMAND_MARKER;
@@ -99,6 +100,24 @@ pub(crate) fn disable_powershell_profile_for_elevated_windows_sandbox(
     command
 }
 
+pub(crate) async fn await_shell_snapshot_for_command(
+    command: &[String],
+    session_shell: &Shell,
+    cwd: &AbsolutePathBuf,
+) -> Result<(), ShellSnapshotFailure> {
+    if cfg!(windows) || command.get(1).map(String::as_str) != Some("-lc") {
+        return Ok(());
+    }
+
+    session_shell.wait_for_shell_snapshot(cwd).await.map(|_| ())
+}
+
+pub(crate) fn shell_snapshot_failure_message(failure_reason: ShellSnapshotFailure) -> String {
+    format!(
+        "failed to initialize the shell snapshot: {failure_reason}; refusing to execute without the expected shell environment"
+    )
+}
+
 /// POSIX-only helper: for commands produced by `Shell::derive_exec_args`
 /// for Bash/Zsh/sh of the form `[shell_path, "-lc", "<script>"]`, and
 /// when a snapshot is configured on the session shell, rewrite the argv

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

@@ -4,6 +4,8 @@ use crate::exec::ExecExpiration;
 use crate::sandboxing::ExecOptions;
 use crate::shell::ShellType;
 use crate::shell_snapshot::ShellSnapshot;
+use crate::shell_snapshot::ShellSnapshotFailure;
+use crate::shell_snapshot::ShellSnapshotState;
 use crate::tools::sandboxing::SandboxAttempt;
 use crate::tools::sandboxing::managed_network_for_sandbox_permissions;
 #[cfg(target_os = "macos")]
@@ -55,10 +57,11 @@ fn shell_with_snapshot(
     snapshot_path: AbsolutePathBuf,
     snapshot_cwd: AbsolutePathBuf,
 ) -> Shell {
-    let (_tx, shell_snapshot) = watch::channel(Some(Arc::new(ShellSnapshot {
-        path: snapshot_path,
-        cwd: snapshot_cwd,
-    })));
+    let (_tx, shell_snapshot) =
+        watch::channel(ShellSnapshotState::Ready(Arc::new(ShellSnapshot {
+            path: snapshot_path,
+            cwd: snapshot_cwd,
+        })));
     Shell {
         shell_type,
         shell_path: PathBuf::from(shell_path),
@@ -66,6 +69,100 @@ fn shell_with_snapshot(
     }
 }
 
+#[tokio::test]
+async fn await_shell_snapshot_for_command_waits_for_pending_snapshot() {
+    let dir = tempdir().expect("create temp dir");
+    let cwd = dir.path().abs();
+    let (sender, shell_snapshot) = watch::channel(ShellSnapshotState::Pending { cwd: cwd.clone() });
+    let session_shell = Shell {
+        shell_type: ShellType::Bash,
+        shell_path: PathBuf::from("/bin/bash"),
+        shell_snapshot,
+    };
+    let command = vec![
+        "/bin/bash".to_string(),
+        "-lc".to_string(),
+        "printf ready".to_string(),
+    ];
+
+    let waiter = tokio::spawn(async move {
+        await_shell_snapshot_for_command(&command, &session_shell, &cwd).await
+    });
+    tokio::task::yield_now().await;
+    assert_eq!(waiter.is_finished(), false);
+
+    sender
+        .send(ShellSnapshotState::Ready(Arc::new(ShellSnapshot {
+            path: dir.path().join("snapshot.sh").abs(),
+            cwd: dir.path().abs(),
+        })))
+        .expect("publish snapshot");
+
+    assert_eq!(waiter.await.expect("join snapshot waiter"), Ok(()));
+}
+
+#[tokio::test]
+async fn await_shell_snapshot_for_command_reports_failed_snapshot() {
+    let dir = tempdir().expect("create temp dir");
+    let cwd = dir.path().abs();
+    let (_sender, shell_snapshot) = watch::channel(ShellSnapshotState::Failed {
+        cwd: cwd.clone(),
+        failure_reason: ShellSnapshotFailure::WriteFailed,
+    });
+    let session_shell = Shell {
+        shell_type: ShellType::Bash,
+        shell_path: PathBuf::from("/bin/bash"),
+        shell_snapshot,
+    };
+    let command = vec![
+        "/bin/bash".to_string(),
+        "-lc".to_string(),
+        "printf never-runs".to_string(),
+    ];
+
+    assert_eq!(
+        await_shell_snapshot_for_command(&command, &session_shell, &cwd).await,
+        Err(ShellSnapshotFailure::WriteFailed)
+    );
+}
+
+#[tokio::test]
+async fn await_shell_snapshot_for_command_ignores_snapshot_for_other_cwd() {
+    let snapshot_dir = tempdir().expect("create snapshot temp dir");
+    let command_dir = tempdir().expect("create command temp dir");
+    let snapshot_cwd = snapshot_dir.path().abs();
+    let command_cwd = command_dir.path().abs();
+    let (sender, shell_snapshot) = watch::channel(ShellSnapshotState::Pending {
+        cwd: snapshot_cwd.clone(),
+    });
+    let session_shell = Shell {
+        shell_type: ShellType::Bash,
+        shell_path: PathBuf::from("/bin/bash"),
+        shell_snapshot,
+    };
+    let command = vec![
+        "/bin/bash".to_string(),
+        "-lc".to_string(),
+        "printf without-snapshot".to_string(),
+    ];
+
+    assert_eq!(
+        await_shell_snapshot_for_command(&command, &session_shell, &command_cwd).await,
+        Ok(())
+    );
+
+    sender
+        .send(ShellSnapshotState::Failed {
+            cwd: snapshot_cwd,
+            failure_reason: ShellSnapshotFailure::WriteFailed,
+        })
+        .expect("publish unrelated snapshot failure");
+    assert_eq!(
+        await_shell_snapshot_for_command(&command, &session_shell, &command_cwd).await,
+        Ok(())
+    );
+}
+
 async fn test_network_proxy() -> anyhow::Result<NetworkProxy> {
     let state = codex_network_proxy::build_config_state(
         NetworkProxyConfig::default(),

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

@@ -20,10 +20,12 @@ use crate::shell::ShellType;
 use crate::tools::flat_tool_name;
 use crate::tools::network_approval::NetworkApprovalMode;
 use crate::tools::network_approval::NetworkApprovalSpec;
+use crate::tools::runtimes::await_shell_snapshot_for_command;
 use crate::tools::runtimes::build_sandbox_command;
 use crate::tools::runtimes::disable_powershell_profile_for_elevated_windows_sandbox;
 use crate::tools::runtimes::exec_env_for_sandbox_permissions;
 use crate::tools::runtimes::maybe_wrap_shell_lc_with_snapshot;
+use crate::tools::runtimes::shell_snapshot_failure_message;
 use crate::tools::sandboxing::Approvable;
 use crate::tools::sandboxing::ApprovalCtx;
 use crate::tools::sandboxing::ExecApprovalRequirement;
@@ -232,6 +234,11 @@ impl ToolRuntime<ShellRequest, ExecToolCallOutput> for ShellRuntime {
         let managed_network =
             managed_network_for_sandbox_permissions(req.network.as_ref(), req.sandbox_permissions);
         let env = exec_env_for_sandbox_permissions(&req.env, req.sandbox_permissions);
+        await_shell_snapshot_for_command(&req.command, session_shell.as_ref(), &req.cwd)
+            .await
+            .map_err(|failure_reason| {
+                ToolError::Rejected(shell_snapshot_failure_message(failure_reason))
+            })?;
         let command = maybe_wrap_shell_lc_with_snapshot(
             &req.command,
             session_shell.as_ref(),

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

@@ -17,11 +17,13 @@ use crate::shell::ShellType;
 use crate::tools::flat_tool_name;
 use crate::tools::network_approval::NetworkApprovalMode;
 use crate::tools::network_approval::NetworkApprovalSpec;
+use crate::tools::runtimes::await_shell_snapshot_for_command;
 use crate::tools::runtimes::build_sandbox_command;
 use crate::tools::runtimes::disable_powershell_profile_for_elevated_windows_sandbox;
 use crate::tools::runtimes::exec_env_for_sandbox_permissions;
 use crate::tools::runtimes::maybe_wrap_shell_lc_with_snapshot;
 use crate::tools::runtimes::shell::zsh_fork_backend;
+use crate::tools::runtimes::shell_snapshot_failure_message;
 use crate::tools::sandboxing::Approvable;
 use crate::tools::sandboxing::ApprovalCtx;
 use crate::tools::sandboxing::ExecApprovalRequirement;
@@ -265,6 +267,11 @@ impl<'a> ToolRuntime<UnifiedExecRequest, UnifiedExecProcess> for UnifiedExecRunt
         let command = if environment_is_remote {
             base_command.to_vec()
         } else {
+            await_shell_snapshot_for_command(base_command, session_shell.as_ref(), &req.cwd)
+                .await
+                .map_err(|failure_reason| {
+                    ToolError::Rejected(shell_snapshot_failure_message(failure_reason))
+                })?;
             maybe_wrap_shell_lc_with_snapshot(
                 base_command,
                 session_shell.as_ref(),