From b75024c465280ad1ee9f5f07728619840b4c58ca Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 21 Jan 2026 10:41:13 +0000 Subject: [PATCH] feat: async shell snapshot (#9600) --- codex-rs/core/src/codex.rs | 18 +++---- codex-rs/core/src/environment_context.rs | 6 +-- codex-rs/core/src/shell.rs | 54 +++++++++++++------ codex-rs/core/src/shell_snapshot.rs | 35 ++++++++++-- codex-rs/core/src/tools/handlers/shell.rs | 16 +++--- .../core/src/tools/handlers/unified_exec.rs | 2 +- codex-rs/core/src/tools/runtimes/mod.rs | 2 +- codex-rs/core/tests/suite/shell_snapshot.rs | 48 ++++++++--------- 8 files changed, 113 insertions(+), 68 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 83fd5a98c..a05521e80 100644 --- a/codex-rs/core/src/codex.rs +++ b/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()); diff --git a/codex-rs/core/src/environment_context.rs b/codex-rs/core/src/environment_context.rs index 3e340ebbd..28b5032fa 100644 --- a/codex-rs/core/src/environment_context.rs +++ b/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(), }, ); diff --git a/codex-rs/core/src/shell.rs b/codex-rs/core/src/shell.rs index b6adf861d..c866afd3b 100644 --- a/codex-rs/core/src/shell.rs +++ b/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>, + #[serde( + skip_serializing, + skip_deserializing, + default = "empty_shell_snapshot_receiver" + )] + pub(crate) shell_snapshot: watch::Receiver>>, } impl Shell { @@ -63,8 +67,26 @@ impl Shell { } } } + + /// Return the shell snapshot if existing. + pub fn shell_snapshot(&self) -> Option> { + self.shell_snapshot.borrow().clone() + } } +pub(crate) fn empty_shell_snapshot_receiver() -> watch::Receiver>> { + 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 { use libc::getpwuid; @@ -139,7 +161,7 @@ fn get_zsh_shell(path: Option<&PathBuf>) -> Option { 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_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_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_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_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(), } ); } diff --git a/codex-rs/core/src/shell_snapshot.rs b/codex-rs/core/src/shell_snapshot.rs index be316afa5..81ac293b9 100644 --- a/codex-rs/core/src/shell_snapshot.rs +++ b/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 { + 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 { // 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 { +async fn write_shell_snapshot(shell_type: ShellType, output_path: &Path) -> Result { 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)) diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index 0e14da68f..dc9f198cc 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/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 = diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index daa39bc35..c9c5a3a71 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/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) -> Vec { 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 }); diff --git a/codex-rs/core/src/tools/runtimes/mod.rs b/codex-rs/core/src/tools/runtimes/mod.rs index 044173bfb..ddf367178 100644 --- a/codex-rs/core/src/tools/runtimes/mod.rs +++ b/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 { - let Some(snapshot) = &session_shell.shell_snapshot else { + let Some(snapshot) = session_shell.shell_snapshot() else { return command.to_vec(); }; diff --git a/codex-rs/core/tests/suite/shell_snapshot.rs b/codex-rs/core/tests/suite/shell_snapshot.rs index 443a41aba..2177ab9bd 100644 --- a/codex-rs/core/tests/suite/shell_snapshot.rs +++ b/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 { + 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 { let builder = test_codex().with_config(|config| { @@ -89,12 +109,7 @@ async fn run_snapshot_command(command: &str) -> Result { _ => 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 { _ => 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?;