feat: async shell snapshot (#9600)

codex-rs/core/src/codex.rs

@@ -723,18 +723,12 @@ impl Session {
         let mut default_shell = shell::default_user_shell();
         // Create the mutable state for the Session.
         if config.features.enabled(Feature::ShellSnapshot) {
-            let timer = otel_manager.start_timer("codex.shell_snapshot.duration_ms", &[]);
-            default_shell.shell_snapshot =
-                ShellSnapshot::try_new(&config.codex_home, conversation_id, &default_shell)
-                    .await
-                    .map(Arc::new);
-            let success = if default_shell.shell_snapshot.is_some() {
-                "true"
-            } else {
-                "false"
-            };
-            let _ = timer.map(|timer| timer.record(&[("success", success)]));
-            otel_manager.counter("codex.shell_snapshot", 1, &[("success", success)])
+            ShellSnapshot::start_snapshotting(
+                config.codex_home.clone(),
+                conversation_id,
+                &mut default_shell,
+                otel_manager.clone(),
+            );
         }
         let state = SessionState::new(session_configuration.clone());
 

codex-rs/core/src/environment_context.rs

@@ -95,7 +95,7 @@ mod tests {
         Shell {
             shell_type: ShellType::Bash,
             shell_path: PathBuf::from("/bin/bash"),
-            shell_snapshot: None,
+            shell_snapshot: crate::shell::empty_shell_snapshot_receiver(),
         }
     }
 
@@ -189,7 +189,7 @@ mod tests {
             Shell {
                 shell_type: ShellType::Bash,
                 shell_path: "/bin/bash".into(),
-                shell_snapshot: None,
+                shell_snapshot: crate::shell::empty_shell_snapshot_receiver(),
             },
         );
         let context2 = EnvironmentContext::new(
@@ -197,7 +197,7 @@ mod tests {
             Shell {
                 shell_type: ShellType::Zsh,
                 shell_path: "/bin/zsh".into(),
-                shell_snapshot: None,
+                shell_snapshot: crate::shell::empty_shell_snapshot_receiver(),
             },
         );
 

codex-rs/core/src/shell.rs

@@ -1,9 +1,9 @@
+use crate::shell_snapshot::ShellSnapshot;
 use serde::Deserialize;
 use serde::Serialize;
 use std::path::PathBuf;
 use std::sync::Arc;
-
-use crate::shell_snapshot::ShellSnapshot;
+use tokio::sync::watch;
 
 #[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize)]
 pub enum ShellType {
@@ -14,12 +14,16 @@ pub enum ShellType {
     Cmd,
 }
 
-#[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Shell {
     pub(crate) shell_type: ShellType,
     pub(crate) shell_path: PathBuf,
-    #[serde(skip_serializing, skip_deserializing, default)]
-    pub(crate) shell_snapshot: Option<Arc<ShellSnapshot>>,
+    #[serde(
+        skip_serializing,
+        skip_deserializing,
+        default = "empty_shell_snapshot_receiver"
+    )]
+    pub(crate) shell_snapshot: watch::Receiver<Option<Arc<ShellSnapshot>>>,
 }
 
 impl Shell {
@@ -63,8 +67,26 @@ impl Shell {
             }
         }
     }
+
+    /// Return the shell snapshot if existing.
+    pub fn shell_snapshot(&self) -> Option<Arc<ShellSnapshot>> {
+        self.shell_snapshot.borrow().clone()
+    }
 }
 
+pub(crate) fn empty_shell_snapshot_receiver() -> watch::Receiver<Option<Arc<ShellSnapshot>>> {
+    let (_tx, rx) = watch::channel(None);
+    rx
+}
+
+impl PartialEq for Shell {
+    fn eq(&self, other: &Self) -> bool {
+        self.shell_type == other.shell_type && self.shell_path == other.shell_path
+    }
+}
+
+impl Eq for Shell {}
+
 #[cfg(unix)]
 fn get_user_shell_path() -> Option<PathBuf> {
     use libc::getpwuid;
@@ -139,7 +161,7 @@ fn get_zsh_shell(path: Option<&PathBuf>) -> Option<Shell> {
     shell_path.map(|shell_path| Shell {
         shell_type: ShellType::Zsh,
         shell_path,
-        shell_snapshot: None,
+        shell_snapshot: empty_shell_snapshot_receiver(),
     })
 }
 
@@ -149,7 +171,7 @@ fn get_bash_shell(path: Option<&PathBuf>) -> Option<Shell> {
     shell_path.map(|shell_path| Shell {
         shell_type: ShellType::Bash,
         shell_path,
-        shell_snapshot: None,
+        shell_snapshot: empty_shell_snapshot_receiver(),
     })
 }
 
@@ -159,7 +181,7 @@ fn get_sh_shell(path: Option<&PathBuf>) -> Option<Shell> {
     shell_path.map(|shell_path| Shell {
         shell_type: ShellType::Sh,
         shell_path,
-        shell_snapshot: None,
+        shell_snapshot: empty_shell_snapshot_receiver(),
     })
 }
 
@@ -175,7 +197,7 @@ fn get_powershell_shell(path: Option<&PathBuf>) -> Option<Shell> {
     shell_path.map(|shell_path| Shell {
         shell_type: ShellType::PowerShell,
         shell_path,
-        shell_snapshot: None,
+        shell_snapshot: empty_shell_snapshot_receiver(),
     })
 }
 
@@ -185,7 +207,7 @@ fn get_cmd_shell(path: Option<&PathBuf>) -> Option<Shell> {
     shell_path.map(|shell_path| Shell {
         shell_type: ShellType::Cmd,
         shell_path,
-        shell_snapshot: None,
+        shell_snapshot: empty_shell_snapshot_receiver(),
     })
 }
 
@@ -194,13 +216,13 @@ fn ultimate_fallback_shell() -> Shell {
         Shell {
             shell_type: ShellType::Cmd,
             shell_path: PathBuf::from("cmd.exe"),
-            shell_snapshot: None,
+            shell_snapshot: empty_shell_snapshot_receiver(),
         }
     } else {
         Shell {
             shell_type: ShellType::Sh,
             shell_path: PathBuf::from("/bin/sh"),
-            shell_snapshot: None,
+            shell_snapshot: empty_shell_snapshot_receiver(),
         }
     }
 }
@@ -426,7 +448,7 @@ mod tests {
         let test_bash_shell = Shell {
             shell_type: ShellType::Bash,
             shell_path: PathBuf::from("/bin/bash"),
-            shell_snapshot: None,
+            shell_snapshot: empty_shell_snapshot_receiver(),
         };
         assert_eq!(
             test_bash_shell.derive_exec_args("echo hello", false),
@@ -440,7 +462,7 @@ mod tests {
         let test_zsh_shell = Shell {
             shell_type: ShellType::Zsh,
             shell_path: PathBuf::from("/bin/zsh"),
-            shell_snapshot: None,
+            shell_snapshot: empty_shell_snapshot_receiver(),
         };
         assert_eq!(
             test_zsh_shell.derive_exec_args("echo hello", false),
@@ -454,7 +476,7 @@ mod tests {
         let test_powershell_shell = Shell {
             shell_type: ShellType::PowerShell,
             shell_path: PathBuf::from("pwsh.exe"),
-            shell_snapshot: None,
+            shell_snapshot: empty_shell_snapshot_receiver(),
         };
         assert_eq!(
             test_powershell_shell.derive_exec_args("echo hello", false),
@@ -481,7 +503,7 @@ mod tests {
                 Shell {
                     shell_type: ShellType::Zsh,
                     shell_path: PathBuf::from(shell_path),
-                    shell_snapshot: None,
+                    shell_snapshot: empty_shell_snapshot_receiver(),
                 }
             );
         }

codex-rs/core/src/shell_snapshot.rs

@@ -1,6 +1,7 @@
 use std::io::ErrorKind;
 use std::path::Path;
 use std::path::PathBuf;
+use std::sync::Arc;
 use std::time::Duration;
 use std::time::SystemTime;
 
@@ -12,9 +13,11 @@ use anyhow::Context;
 use anyhow::Result;
 use anyhow::anyhow;
 use anyhow::bail;
+use codex_otel::OtelManager;
 use codex_protocol::ThreadId;
 use tokio::fs;
 use tokio::process::Command;
+use tokio::sync::watch;
 use tokio::time::timeout;
 
 #[derive(Clone, Debug, PartialEq, Eq)]
@@ -27,7 +30,31 @@ const SNAPSHOT_RETENTION: Duration = Duration::from_secs(60 * 60 * 24 * 7); // 7
 const SNAPSHOT_DIR: &str = "shell_snapshots";
 
 impl ShellSnapshot {
-    pub async fn try_new(codex_home: &Path, session_id: ThreadId, shell: &Shell) -> Option<Self> {
+    pub fn start_snapshotting(
+        codex_home: PathBuf,
+        session_id: ThreadId,
+        shell: &mut Shell,
+        otel_manager: OtelManager,
+    ) {
+        let (shell_snapshot_tx, shell_snapshot_rx) = watch::channel(None);
+        shell.shell_snapshot = shell_snapshot_rx;
+
+        let snapshot_shell = shell.clone();
+        let snapshot_session_id = session_id;
+        tokio::spawn(async move {
+            let timer = otel_manager.start_timer("codex.shell_snapshot.duration_ms", &[]);
+            let snapshot =
+                ShellSnapshot::try_new(&codex_home, snapshot_session_id, &snapshot_shell)
+                    .await
+                    .map(Arc::new);
+            let success = if snapshot.is_some() { "true" } else { "false" };
+            let _ = timer.map(|timer| timer.record(&[("success", success)]));
+            otel_manager.counter("codex.shell_snapshot", 1, &[("success", success)]);
+            let _ = shell_snapshot_tx.send(snapshot);
+        });
+    }
+
+    async fn try_new(codex_home: &Path, session_id: ThreadId, shell: &Shell) -> Option<Self> {
         // File to store the snapshot
         let extension = match shell.shell_type {
             ShellType::PowerShell => "ps1",
@@ -74,7 +101,7 @@ impl Drop for ShellSnapshot {
     }
 }
 
-pub async fn write_shell_snapshot(shell_type: ShellType, output_path: &Path) -> Result<PathBuf> {
+async fn write_shell_snapshot(shell_type: ShellType, output_path: &Path) -> Result<PathBuf> {
     if shell_type == ShellType::PowerShell || shell_type == ShellType::Cmd {
         bail!("Shell snapshot not supported yet for {shell_type:?}");
     }
@@ -407,7 +434,7 @@ mod tests {
         let shell = Shell {
             shell_type: ShellType::Bash,
             shell_path: PathBuf::from("/bin/bash"),
-            shell_snapshot: None,
+            shell_snapshot: crate::shell::empty_shell_snapshot_receiver(),
         };
 
         let snapshot = ShellSnapshot::try_new(dir.path(), ThreadId::new(), &shell)
@@ -449,7 +476,7 @@ mod tests {
         let shell = Shell {
             shell_type: ShellType::Sh,
             shell_path,
-            shell_snapshot: None,
+            shell_snapshot: crate::shell::empty_shell_snapshot_receiver(),
         };
 
         let err = run_shell_script_with_timeout(&shell, "ignored", Duration::from_millis(500))

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

@@ -305,6 +305,7 @@ mod tests {
     use crate::shell::ShellType;
     use crate::shell_snapshot::ShellSnapshot;
     use crate::tools::handlers::ShellCommandHandler;
+    use tokio::sync::watch;
 
     /// The logic for is_known_safe_command() has heuristics for known shells,
     /// so we must ensure the commands generated by [ShellCommandHandler] can be
@@ -314,14 +315,14 @@ mod tests {
         let bash_shell = Shell {
             shell_type: ShellType::Bash,
             shell_path: PathBuf::from("/bin/bash"),
-            shell_snapshot: None,
+            shell_snapshot: crate::shell::empty_shell_snapshot_receiver(),
         };
         assert_safe(&bash_shell, "ls -la");
 
         let zsh_shell = Shell {
             shell_type: ShellType::Zsh,
             shell_path: PathBuf::from("/bin/zsh"),
-            shell_snapshot: None,
+            shell_snapshot: crate::shell::empty_shell_snapshot_receiver(),
         };
         assert_safe(&zsh_shell, "ls -la");
 
@@ -329,7 +330,7 @@ mod tests {
             let powershell = Shell {
                 shell_type: ShellType::PowerShell,
                 shell_path: path.to_path_buf(),
-                shell_snapshot: None,
+                shell_snapshot: crate::shell::empty_shell_snapshot_receiver(),
             };
             assert_safe(&powershell, "ls -Name");
         }
@@ -338,7 +339,7 @@ mod tests {
             let pwsh = Shell {
                 shell_type: ShellType::PowerShell,
                 shell_path: path.to_path_buf(),
-                shell_snapshot: None,
+                shell_snapshot: crate::shell::empty_shell_snapshot_receiver(),
             };
             assert_safe(&pwsh, "ls -Name");
         }
@@ -391,12 +392,13 @@ mod tests {
 
     #[test]
     fn shell_command_handler_respects_explicit_login_flag() {
+        let (_tx, shell_snapshot) = watch::channel(Some(Arc::new(ShellSnapshot {
+            path: PathBuf::from("/tmp/snapshot.sh"),
+        })));
         let shell = Shell {
             shell_type: ShellType::Bash,
             shell_path: PathBuf::from("/bin/bash"),
-            shell_snapshot: Some(Arc::new(ShellSnapshot {
-                path: PathBuf::from("/tmp/snapshot.sh"),
-            })),
+            shell_snapshot,
         };
 
         let login_command =

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

@@ -236,7 +236,7 @@ impl ToolHandler for UnifiedExecHandler {
 fn get_command(args: &ExecCommandArgs, session_shell: Arc<Shell>) -> Vec<String> {
     let model_shell = args.shell.as_ref().map(|shell_str| {
         let mut shell = get_shell_by_model_provided_path(&PathBuf::from(shell_str));
-        shell.shell_snapshot = None;
+        shell.shell_snapshot = crate::shell::empty_shell_snapshot_receiver();
         shell
     });
 

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

@@ -54,7 +54,7 @@ pub(crate) fn maybe_wrap_shell_lc_with_snapshot(
     command: &[String],
     session_shell: &Shell,
 ) -> Vec<String> {
-    let Some(snapshot) = &session_shell.shell_snapshot else {
+    let Some(snapshot) = session_shell.shell_snapshot() else {
         return command.to_vec();
     };
 

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

@@ -20,9 +20,11 @@ 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::path::Path;
 use std::path::PathBuf;
 use tokio::fs;
 use tokio::time::Duration;
+use tokio::time::Instant;
 use tokio::time::sleep;
 
 #[derive(Debug)]
@@ -34,6 +36,24 @@ struct SnapshotRun {
     codex_home: PathBuf,
 }
 
+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
+            && let Some(entry) = entries.next_entry().await?
+        {
+            return Ok(entry.path());
+        }
+
+        if Instant::now() >= deadline {
+            anyhow::bail!("timed out waiting for shell snapshot");
+        }
+
+        sleep(Duration::from_millis(25)).await;
+    }
+}
+
 #[allow(clippy::expect_used)]
 async fn run_snapshot_command(command: &str) -> Result<SnapshotRun> {
     let builder = test_codex().with_config(|config| {
@@ -89,12 +109,7 @@ async fn run_snapshot_command(command: &str) -> Result<SnapshotRun> {
         _ => None,
     })
     .await;
-    let mut entries = fs::read_dir(codex_home.join("shell_snapshots")).await?;
-    let snapshot_path = entries
-        .next_entry()
-        .await?
-        .map(|entry| entry.path())
-        .expect("shell snapshot created");
+    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 {
@@ -167,12 +182,7 @@ async fn run_shell_command_snapshot(command: &str) -> Result<SnapshotRun> {
         _ => None,
     })
     .await;
-    let mut entries = fs::read_dir(codex_home.join("shell_snapshots")).await?;
-    let snapshot_path = entries
-        .next_entry()
-        .await?
-        .map(|entry| entry.path())
-        .expect("shell snapshot created");
+    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 {
@@ -305,12 +315,7 @@ async fn shell_command_snapshot_still_intercepts_apply_patch() -> Result<()> {
 
     assert_eq!(fs::read_to_string(&target).await?, "hello from snapshot\n");
 
-    let mut entries = fs::read_dir(codex_home.join("shell_snapshots")).await?;
-    let snapshot_path = entries
-        .next_entry()
-        .await?
-        .map(|entry| entry.path())
-        .expect("shell snapshot created");
+    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);
 
@@ -328,12 +333,7 @@ async fn shell_snapshot_deleted_after_shutdown_with_skills() -> Result<()> {
     let codex_home = home.path().to_path_buf();
     let codex = harness.test().codex.clone();
 
-    let mut entries = fs::read_dir(codex_home.join("shell_snapshots")).await?;
-    let snapshot_path = entries
-        .next_entry()
-        .await?
-        .map(|entry| entry.path())
-        .expect("shell snapshot created");
+    let snapshot_path = wait_for_snapshot(&codex_home).await?;
     assert!(snapshot_path.exists());
 
     codex.submit(Op::Shutdown {}).await?;