diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 80319fa8a..2abb94165 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use std::collections::HashSet; use std::fmt::Debug; +use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::AtomicU64; @@ -87,6 +88,7 @@ use tokio::sync::Mutex; use tokio::sync::OnceCell; use tokio::sync::RwLock; use tokio::sync::oneshot; +use tokio::sync::watch; use tokio_util::sync::CancellationToken; use tracing::Instrument; use tracing::debug; @@ -234,8 +236,6 @@ use codex_protocol::protocol::InitialHistory; use codex_protocol::user_input::UserInput; use codex_utils_readiness::Readiness; use codex_utils_readiness::ReadinessFlag; -use tokio::sync::watch; - /// The high-level interface to the Codex system. /// It operates as a queue pair where you send submissions and receive events. pub struct Codex { @@ -1033,14 +1033,19 @@ impl Session { let mut default_shell = shell::default_user_shell(); // Create the mutable state for the Session. - if config.features.enabled(Feature::ShellSnapshot) { + let shell_snapshot_tx = if config.features.enabled(Feature::ShellSnapshot) { ShellSnapshot::start_snapshotting( config.codex_home.clone(), conversation_id, + session_configuration.cwd.clone(), &mut default_shell, otel_manager.clone(), - ); - } + ) + } else { + let (tx, rx) = watch::channel(None); + default_shell.shell_snapshot = rx; + tx + }; let thread_name = match session_index::find_thread_name_by_id(&config.codex_home, &conversation_id).await { @@ -1064,6 +1069,7 @@ impl Session { hooks: Hooks::new(config.as_ref()), rollout: Mutex::new(rollout_recorder), user_shell: Arc::new(default_shell), + shell_snapshot_tx, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, auth_manager: Arc::clone(&auth_manager), @@ -1405,6 +1411,30 @@ impl Session { state.pending_resume_previous_model.take() } + fn maybe_refresh_shell_snapshot_for_cwd( + &self, + previous_cwd: &Path, + next_cwd: &Path, + codex_home: &Path, + ) { + if previous_cwd == next_cwd { + return; + } + + if !self.features.enabled(Feature::ShellSnapshot) { + return; + } + + ShellSnapshot::refresh_snapshot( + codex_home.to_path_buf(), + self.conversation_id, + next_cwd.to_path_buf(), + self.services.user_shell.as_ref().clone(), + self.services.shell_snapshot_tx.clone(), + self.services.otel_manager.clone(), + ); + } + pub(crate) async fn update_settings( &self, updates: SessionSettingsUpdate, @@ -1413,7 +1443,14 @@ impl Session { match state.session_configuration.apply(&updates) { Ok(updated) => { + let previous_cwd = state.session_configuration.cwd.clone(); + let next_cwd = updated.cwd.clone(); + let codex_home = updated.codex_home.clone(); state.session_configuration = updated; + drop(state); + + self.maybe_refresh_shell_snapshot_for_cwd(&previous_cwd, &next_cwd, &codex_home); + Ok(()) } Err(err) => { @@ -1428,14 +1465,16 @@ impl Session { sub_id: String, updates: SessionSettingsUpdate, ) -> ConstraintResult> { - let (session_configuration, sandbox_policy_changed) = { + let (session_configuration, sandbox_policy_changed, previous_cwd, codex_home) = { let mut state = self.state.lock().await; match state.session_configuration.clone().apply(&updates) { Ok(next) => { + let previous_cwd = state.session_configuration.cwd.clone(); let sandbox_policy_changed = state.session_configuration.sandbox_policy != next.sandbox_policy; + let codex_home = next.codex_home.clone(); state.session_configuration = next.clone(); - (next, sandbox_policy_changed) + (next, sandbox_policy_changed, previous_cwd, codex_home) } Err(err) => { drop(state); @@ -1452,6 +1491,12 @@ impl Session { } }; + self.maybe_refresh_shell_snapshot_for_cwd( + &previous_cwd, + &session_configuration.cwd, + &codex_home, + ); + Ok(self .new_turn_from_configuration( sub_id, @@ -6123,6 +6168,7 @@ mod tests { hooks: Hooks::new(&config), rollout: Mutex::new(None), user_shell: Arc::new(default_user_shell()), + shell_snapshot_tx: watch::channel(None).0, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, auth_manager: auth_manager.clone(), @@ -6255,6 +6301,7 @@ mod tests { hooks: Hooks::new(&config), rollout: Mutex::new(None), user_shell: Arc::new(default_user_shell()), + shell_snapshot_tx: watch::channel(None).0, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, auth_manager: Arc::clone(&auth_manager), diff --git a/codex-rs/core/src/shell_snapshot.rs b/codex-rs/core/src/shell_snapshot.rs index 127732824..6eb1ef71d 100644 --- a/codex-rs/core/src/shell_snapshot.rs +++ b/codex-rs/core/src/shell_snapshot.rs @@ -26,6 +26,7 @@ use tracing::info_span; #[derive(Clone, Debug, PartialEq, Eq)] pub struct ShellSnapshot { pub path: PathBuf, + pub cwd: PathBuf, } const SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(10); @@ -37,22 +38,63 @@ impl ShellSnapshot { pub fn start_snapshotting( codex_home: PathBuf, session_id: ThreadId, + session_cwd: PathBuf, shell: &mut Shell, otel_manager: OtelManager, - ) { + ) -> watch::Sender>> { 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; - let snapshot_span = info_span!("shell_snapshot", thread_id = %snapshot_session_id); + Self::spawn_snapshot_task( + codex_home, + session_id, + session_cwd, + shell.clone(), + shell_snapshot_tx.clone(), + otel_manager, + ); + + shell_snapshot_tx + } + + pub fn refresh_snapshot( + codex_home: PathBuf, + session_id: ThreadId, + session_cwd: PathBuf, + shell: Shell, + shell_snapshot_tx: watch::Sender>>, + otel_manager: OtelManager, + ) { + Self::spawn_snapshot_task( + codex_home, + session_id, + session_cwd, + shell, + shell_snapshot_tx, + otel_manager, + ); + } + + fn spawn_snapshot_task( + codex_home: PathBuf, + session_id: ThreadId, + session_cwd: PathBuf, + snapshot_shell: Shell, + shell_snapshot_tx: watch::Sender>>, + otel_manager: OtelManager, + ) { + let snapshot_span = info_span!("shell_snapshot", thread_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 snapshot = ShellSnapshot::try_new( + &codex_home, + session_id, + session_cwd.as_path(), + &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)]); @@ -62,7 +104,12 @@ impl ShellSnapshot { ); } - async fn try_new(codex_home: &Path, session_id: ThreadId, shell: &Shell) -> Option { + async fn try_new( + codex_home: &Path, + session_id: ThreadId, + session_cwd: &Path, + shell: &Shell, + ) -> Option { // File to store the snapshot let extension = match shell.shell_type { ShellType::PowerShell => "ps1", @@ -82,22 +129,26 @@ impl ShellSnapshot { }); // Make the new snapshot. - let snapshot = match write_shell_snapshot(shell.shell_type.clone(), &path).await { - Ok(path) => { - tracing::info!("Shell snapshot successfully created: {}", path.display()); - Some(Self { path }) - } - Err(err) => { - tracing::warn!( - "Failed to create shell snapshot for {}: {err:?}", - shell.name() - ); - None - } - }; + let snapshot = + match write_shell_snapshot(shell.shell_type.clone(), &path, session_cwd).await { + Ok(path) => { + tracing::info!("Shell snapshot successfully created: {}", path.display()); + Some(Self { + path, + cwd: session_cwd.to_path_buf(), + }) + } + Err(err) => { + tracing::warn!( + "Failed to create shell snapshot for {}: {err:?}", + shell.name() + ); + None + } + }; if let Some(snapshot) = snapshot.as_ref() - && let Err(err) = validate_snapshot(shell, &snapshot.path).await + && let Err(err) = validate_snapshot(shell, &snapshot.path, session_cwd).await { tracing::error!("Shell snapshot validation failed: {err:?}"); return None; @@ -118,14 +169,18 @@ impl Drop for ShellSnapshot { } } -async fn write_shell_snapshot(shell_type: ShellType, output_path: &Path) -> Result { +async fn write_shell_snapshot( + shell_type: ShellType, + output_path: &Path, + cwd: &Path, +) -> Result { if shell_type == ShellType::PowerShell || shell_type == ShellType::Cmd { bail!("Shell snapshot not supported yet for {shell_type:?}"); } let shell = get_shell(shell_type.clone(), None) .with_context(|| format!("No available shell for {shell_type:?}"))?; - let raw_snapshot = capture_snapshot(&shell).await?; + let raw_snapshot = capture_snapshot(&shell, cwd).await?; let snapshot = strip_snapshot_preamble(&raw_snapshot)?; if let Some(parent) = output_path.parent() { @@ -143,13 +198,13 @@ async fn write_shell_snapshot(shell_type: ShellType, output_path: &Path) -> Resu Ok(output_path.to_path_buf()) } -async fn capture_snapshot(shell: &Shell) -> Result { +async fn capture_snapshot(shell: &Shell, cwd: &Path) -> Result { let shell_type = shell.shell_type.clone(); match shell_type { - ShellType::Zsh => run_shell_script(shell, &zsh_snapshot_script()).await, - ShellType::Bash => run_shell_script(shell, &bash_snapshot_script()).await, - ShellType::Sh => run_shell_script(shell, &sh_snapshot_script()).await, - ShellType::PowerShell => run_shell_script(shell, powershell_snapshot_script()).await, + ShellType::Zsh => run_shell_script(shell, &zsh_snapshot_script(), cwd).await, + ShellType::Bash => run_shell_script(shell, &bash_snapshot_script(), cwd).await, + ShellType::Sh => run_shell_script(shell, &sh_snapshot_script(), cwd).await, + ShellType::PowerShell => run_shell_script(shell, powershell_snapshot_script(), cwd).await, ShellType::Cmd => bail!("Shell snapshotting is not yet supported for {shell_type:?}"), } } @@ -163,16 +218,16 @@ fn strip_snapshot_preamble(snapshot: &str) -> Result { Ok(snapshot[start..].to_string()) } -async fn validate_snapshot(shell: &Shell, snapshot_path: &Path) -> Result<()> { +async fn validate_snapshot(shell: &Shell, snapshot_path: &Path, cwd: &Path) -> Result<()> { let snapshot_path_display = snapshot_path.display(); let script = format!("set -e; . \"{snapshot_path_display}\""); - run_script_with_timeout(shell, &script, SNAPSHOT_TIMEOUT, false) + run_script_with_timeout(shell, &script, SNAPSHOT_TIMEOUT, false, cwd) .await .map(|_| ()) } -async fn run_shell_script(shell: &Shell, script: &str) -> Result { - run_script_with_timeout(shell, script, SNAPSHOT_TIMEOUT, true).await +async fn run_shell_script(shell: &Shell, script: &str, cwd: &Path) -> Result { + run_script_with_timeout(shell, script, SNAPSHOT_TIMEOUT, true, cwd).await } async fn run_script_with_timeout( @@ -180,6 +235,7 @@ async fn run_script_with_timeout( script: &str, snapshot_timeout: Duration, use_login_shell: bool, + cwd: &Path, ) -> Result { let args = shell.derive_exec_args(script, use_login_shell); let shell_name = shell.name(); @@ -189,6 +245,7 @@ async fn run_script_with_timeout( let mut handler = Command::new(&args[0]); handler.args(&args[1..]); handler.stdin(Stdio::null()); + handler.current_dir(cwd); #[cfg(unix)] unsafe { handler.pre_exec(|| { @@ -550,7 +607,7 @@ mod tests { async fn get_snapshot(shell_type: ShellType) -> Result { let dir = tempdir()?; let path = dir.path().join("snapshot.sh"); - write_shell_snapshot(shell_type, &path).await?; + write_shell_snapshot(shell_type, &path, dir.path()).await?; let content = fs::read_to_string(&path).await?; Ok(content) } @@ -602,11 +659,12 @@ mod tests { shell_snapshot: crate::shell::empty_shell_snapshot_receiver(), }; - let snapshot = ShellSnapshot::try_new(dir.path(), ThreadId::new(), &shell) + let snapshot = ShellSnapshot::try_new(dir.path(), ThreadId::new(), dir.path(), &shell) .await .expect("snapshot should be created"); let path = snapshot.path.clone(); assert!(path.exists()); + assert_eq!(snapshot.cwd, dir.path().to_path_buf()); drop(snapshot); @@ -635,9 +693,10 @@ mod tests { "HOME=\"{home_display}\"; export HOME; {}", bash_snapshot_script() ); - let output = run_script_with_timeout(&shell, &script, Duration::from_millis(500), true) - .await - .context("run snapshot command")?; + let output = + run_script_with_timeout(&shell, &script, Duration::from_millis(500), true, home) + .await + .context("run snapshot command")?; assert!( output.contains("# Snapshot file"), @@ -665,9 +724,10 @@ mod tests { shell_snapshot: crate::shell::empty_shell_snapshot_receiver(), }; - let err = run_script_with_timeout(&shell, &script, Duration::from_secs(1), true) - .await - .expect_err("snapshot shell should time out"); + let err = + run_script_with_timeout(&shell, &script, Duration::from_secs(1), true, dir.path()) + .await + .expect_err("snapshot shell should time out"); assert!( err.to_string().contains("timed out"), "expected timeout error, got {err:?}" diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index d9f57cb5f..0438119d4 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -17,6 +17,7 @@ use crate::unified_exec::UnifiedExecProcessManager; use codex_otel::OtelManager; use tokio::sync::Mutex; use tokio::sync::RwLock; +use tokio::sync::watch; use tokio_util::sync::CancellationToken; pub(crate) struct SessionServices { @@ -27,6 +28,7 @@ pub(crate) struct SessionServices { pub(crate) hooks: Hooks, pub(crate) rollout: Mutex>, pub(crate) user_shell: Arc, + pub(crate) shell_snapshot_tx: watch::Sender>>, pub(crate) show_raw_agent_reasoning: bool, pub(crate) exec_policy: ExecPolicyManager, pub(crate) auth_manager: Arc, diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index 0ee65eace..c034e26f4 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -113,7 +113,11 @@ pub(crate) async fn execute_user_shell_command( let use_login_shell = true; let session_shell = session.user_shell(); let display_command = session_shell.derive_exec_args(&command, use_login_shell); - let exec_command = maybe_wrap_shell_lc_with_snapshot(&display_command, session_shell.as_ref()); + let exec_command = maybe_wrap_shell_lc_with_snapshot( + &display_command, + session_shell.as_ref(), + turn_context.cwd.as_path(), + ); let call_id = Uuid::new_v4().to_string(); let raw_command = command; diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index fd678c3a3..18c594d99 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -455,6 +455,7 @@ mod tests { fn shell_command_handler_respects_explicit_login_flag() { let (_tx, shell_snapshot) = watch::channel(Some(Arc::new(ShellSnapshot { path: PathBuf::from("/tmp/snapshot.sh"), + cwd: PathBuf::from("/tmp"), }))); let shell = Shell { shell_type: ShellType::Bash, diff --git a/codex-rs/core/src/tools/runtimes/mod.rs b/codex-rs/core/src/tools/runtimes/mod.rs index 8e4c47722..b3d7b3292 100644 --- a/codex-rs/core/src/tools/runtimes/mod.rs +++ b/codex-rs/core/src/tools/runtimes/mod.rs @@ -5,6 +5,7 @@ Concrete ToolRuntime implementations for specific tools. Each runtime stays small and focused and reuses the orchestrator for approvals + sandbox + retry. */ use crate::exec::ExecExpiration; +use crate::path_utils; use crate::sandboxing::CommandSpec; use crate::sandboxing::SandboxPermissions; use crate::shell::Shell; @@ -50,10 +51,12 @@ pub(crate) fn build_command_spec( /// => user_shell -c ". SNAPSHOT (best effort); exec shell -c