From 8f4d1c2ce6d7439a410234b1f87fff6fb3362d94 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 2 Jun 2026 23:35:33 -0700 Subject: [PATCH] refactor shell snapshot store seam --- codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/lib.rs | 2 +- codex-rs/core/src/session/mod.rs | 16 +- codex-rs/core/src/session/session.rs | 4 +- codex-rs/core/src/session/tests.rs | 15 + .../core/src/session/tests/guardian_tests.rs | 4 + codex-rs/core/src/session/turn_context.rs | 4 - codex-rs/core/src/shell_snapshot.rs | 263 +++++++++++------- codex-rs/core/src/shell_snapshot_tests.rs | 43 ++- codex-rs/core/src/state/service.rs | 1 + codex-rs/core/src/thread_manager.rs | 70 +++++ 11 files changed, 280 insertions(+), 143 deletions(-) diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 69d37ad558..7f32130c98 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -106,6 +106,7 @@ pub(crate) async fn run_codex_thread_interactive( analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), thread_store: Arc::clone(&parent_session.services.thread_store), artifact_store: Arc::clone(&parent_session.services.artifact_store), + shell_snapshot_store: Arc::clone(&parent_session.services.shell_snapshot_store), attestation_provider: parent_session.services.attestation_provider.clone(), inherited_multi_agent_version: Some(MultiAgentVersion::Disabled), })) diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 64a6175f29..4e3082d9b0 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -133,7 +133,7 @@ mod rollout; pub(crate) mod safety; mod session_rollout_init_error; pub mod shell; -pub(crate) mod shell_snapshot; +pub mod shell_snapshot; pub mod spawn; pub(crate) mod state_db_bridge; pub use state_db_bridge::StateDbHandle; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index f15f8e8b32..6682c8b8a3 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -421,6 +421,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, pub(crate) artifact_store: Arc, + pub(crate) shell_snapshot_store: Arc, pub(crate) attestation_provider: Option>, pub(crate) inherited_multi_agent_version: Option, } @@ -502,6 +503,7 @@ impl Codex { analytics_events_client, thread_store, artifact_store, + shell_snapshot_store, attestation_provider, inherited_multi_agent_version, } = args; @@ -649,6 +651,7 @@ impl Codex { analytics_events_client, thread_store, artifact_store, + shell_snapshot_store, parent_rollout_thread_trace, attestation_provider, multi_agent_version, @@ -1353,7 +1356,6 @@ impl Session { &self, previous_cwd: &AbsolutePathBuf, next_cwd: &AbsolutePathBuf, - codex_home: &AbsolutePathBuf, session_source: &SessionSource, ) { if previous_cwd == next_cwd { @@ -1372,7 +1374,7 @@ impl Session { } ShellSnapshot::refresh_snapshot( - codex_home.clone(), + Arc::clone(&self.services.shell_snapshot_store), self.conversation_id, next_cwd.clone(), self.services.user_shell.as_ref().clone(), @@ -1393,7 +1395,6 @@ impl Session { previous_cwd, permission_profile_changed, next_cwd, - codex_home, session_source, ) = { let mut state = self.state.lock().await; @@ -1415,7 +1416,6 @@ impl Session { let permission_profile_changed = previous_permission_profile != updated_permission_profile; let next_cwd = updated.cwd.clone(); - let codex_home = updated.codex_home.clone(); let session_source = updated.session_source.clone(); state.session_configuration = updated; ( @@ -1424,18 +1424,12 @@ impl Session { previous_cwd, permission_profile_changed, next_cwd, - codex_home, session_source, ) }; self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref()); - self.maybe_refresh_shell_snapshot_for_cwd( - &previous_cwd, - &next_cwd, - &codex_home, - &session_source, - ); + self.maybe_refresh_shell_snapshot_for_cwd(&previous_cwd, &next_cwd, &session_source); if permission_profile_changed { self.refresh_managed_network_proxy_for_current_permission_profile() .await; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 7ef3899065..d24ffd2f12 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -506,6 +506,7 @@ impl Session { analytics_events_client: Option, thread_store: Arc, artifact_store: Arc, + shell_snapshot_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, attestation_provider: Option>, multi_agent_version: Option, @@ -861,7 +862,7 @@ impl Session { tx } else { ShellSnapshot::start_snapshotting( - config.codex_home.clone(), + Arc::clone(&shell_snapshot_store), thread_id, session_configuration.cwd.clone(), &mut default_shell, @@ -1008,6 +1009,7 @@ impl Session { rollout_thread_trace, user_shell: Arc::new(default_shell), shell_snapshot_tx, + shell_snapshot_store, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, auth_manager: Arc::clone(&auth_manager), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 93496e851c..d6afaefdf2 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4635,6 +4635,9 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { Arc::new(crate::artifact_store::LocalArtifactStore::from_codex_home( &config.codex_home, )), + Arc::new( + crate::shell_snapshot::LocalShellSnapshotStore::from_codex_home(&config.codex_home), + ), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, Some(config.multi_agent_version_from_features()), @@ -4767,6 +4770,9 @@ 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_store: Arc::new( + crate::shell_snapshot::LocalShellSnapshotStore::from_codex_home(&config.codex_home), + ), show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, auth_manager: auth_manager.clone(), @@ -4991,6 +4997,9 @@ async fn make_session_with_config_and_rx( Arc::new(crate::artifact_store::LocalArtifactStore::from_codex_home( &config.codex_home, )), + Arc::new( + crate::shell_snapshot::LocalShellSnapshotStore::from_codex_home(&config.codex_home), + ), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, Some(config.multi_agent_version_from_features()), @@ -5106,6 +5115,9 @@ async fn make_session_with_history_source_and_agent_control_and_rx( Arc::new(crate::artifact_store::LocalArtifactStore::from_codex_home( &config.codex_home, )), + Arc::new( + crate::shell_snapshot::LocalShellSnapshotStore::from_codex_home(&config.codex_home), + ), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, Some(config.multi_agent_version_from_features()), @@ -6866,6 +6878,9 @@ 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_store: Arc::new( + crate::shell_snapshot::LocalShellSnapshotStore::from_codex_home(&config.codex_home), + ), show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, auth_manager: Arc::clone(&auth_manager), diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 21ac52a195..980ac11faa 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -705,6 +705,9 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { let artifact_store = Arc::new(crate::artifact_store::LocalArtifactStore::from_codex_home( &config.codex_home, )); + let shell_snapshot_store = Arc::new( + crate::shell_snapshot::LocalShellSnapshotStore::from_codex_home(&config.codex_home), + ); let CodexSpawnOk { codex, .. } = Codex::spawn(CodexSpawnArgs { config, @@ -737,6 +740,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { analytics_events_client: None, thread_store, artifact_store, + shell_snapshot_store, attestation_provider: None, inherited_multi_agent_version: None, }) diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index d600368137..ac2c6aae45 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -609,7 +609,6 @@ impl Session { let next_permission_profile = next.permission_profile(); let permission_profile_changed = previous_permission_profile != next_permission_profile; - let codex_home = next.codex_home.clone(); let session_source = next.session_source.clone(); let previous_config = notify_config_contributors.then(|| { Self::build_effective_session_config(&state.session_configuration) @@ -622,7 +621,6 @@ impl Session { turn_environments, permission_profile_changed, previous_cwd, - codex_home, session_source, previous_config, new_config, @@ -637,7 +635,6 @@ impl Session { turn_environments, permission_profile_changed, previous_cwd, - codex_home, session_source, previous_config, new_config, @@ -661,7 +658,6 @@ impl Session { self.maybe_refresh_shell_snapshot_for_cwd( &previous_cwd, &session_configuration.cwd, - &codex_home, &session_source, ); diff --git a/codex-rs/core/src/shell_snapshot.rs b/codex-rs/core/src/shell_snapshot.rs index b328a977d7..b25a2c36cc 100644 --- a/codex-rs/core/src/shell_snapshot.rs +++ b/codex-rs/core/src/shell_snapshot.rs @@ -1,5 +1,7 @@ +use std::future::Future; use std::io::ErrorKind; use std::path::Path; +use std::pin::Pin; use std::process::Stdio; use std::sync::Arc; use std::time::Duration; @@ -35,9 +37,144 @@ const SNAPSHOT_RETENTION: Duration = Duration::from_secs(60 * 60 * 24 * 3); // 3 const SNAPSHOT_DIR: &str = "shell_snapshots"; const EXCLUDED_EXPORT_VARS: &[&str] = &["PWD", "OLDPWD"]; +/// Persists executor-local shell snapshot files and owns retention cleanup. +/// Implementations must return local readable paths because shell execution sources snapshots. +pub trait ShellSnapshotStore: Send + Sync + 'static { + /// Allocates the final and temporary paths for one snapshot generation. + fn snapshot_paths(&self, session_id: ThreadId, shell_type: ShellType) -> ShellSnapshotPaths; + + /// Removes stale inactive snapshots according to the store's retention policy. + fn cleanup_stale_snapshots( + &self, + active_session_id: ThreadId, + state_db: Option, + ) -> ShellSnapshotStoreFuture<'_>; +} + +/// Future returned by shell snapshot store cleanup work. +pub type ShellSnapshotStoreFuture<'a> = Pin> + Send + 'a>>; + +#[derive(Clone)] +/// Codex Home backed executor-local shell snapshot store. +pub struct LocalShellSnapshotStore { + codex_home: AbsolutePathBuf, +} + +impl LocalShellSnapshotStore { + /// Constructs a local shell snapshot store rooted at one Codex Home. + pub fn from_codex_home(codex_home: &AbsolutePathBuf) -> Self { + Self { + codex_home: codex_home.clone(), + } + } + + fn snapshot_dir(&self) -> AbsolutePathBuf { + self.codex_home.join(SNAPSHOT_DIR) + } + + async fn cleanup_stale_snapshots_impl( + &self, + active_session_id: ThreadId, + state_db: Option, + ) -> Result<()> { + let snapshot_dir = self.snapshot_dir(); + + let mut entries = match fs::read_dir(&snapshot_dir).await { + Ok(entries) => entries, + Err(err) if err.kind() == ErrorKind::NotFound => return Ok(()), + Err(err) => return Err(err.into()), + }; + + let now = SystemTime::now(); + let active_session_id = active_session_id.to_string(); + + while let Some(entry) = entries.next_entry().await? { + if !entry.file_type().await?.is_file() { + continue; + } + + let path = entry.path(); + + let file_name = entry.file_name(); + let file_name = file_name.to_string_lossy(); + let Some(session_id) = snapshot_session_id_from_file_name(&file_name) else { + remove_snapshot_file(&path).await; + continue; + }; + if session_id == active_session_id { + continue; + } + + let rollout_path = + find_thread_path_by_id_str(&self.codex_home, session_id, state_db.as_deref()) + .await?; + let Some(rollout_path) = rollout_path else { + remove_snapshot_file(&path).await; + continue; + }; + + let modified = match fs::metadata(&rollout_path).await.and_then(|m| m.modified()) { + Ok(modified) => modified, + Err(err) => { + tracing::warn!( + "Failed to check rollout age for snapshot {}: {err:?}", + path.display() + ); + continue; + } + }; + + if now + .duration_since(modified) + .ok() + .is_some_and(|age| age >= SNAPSHOT_RETENTION) + { + remove_snapshot_file(&path).await; + } + } + + Ok(()) + } +} + +impl ShellSnapshotStore for LocalShellSnapshotStore { + fn snapshot_paths(&self, session_id: ThreadId, shell_type: ShellType) -> ShellSnapshotPaths { + let extension = match shell_type { + ShellType::PowerShell => "ps1", + _ => "sh", + }; + let nonce = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .map(|duration| duration.as_nanos()) + .unwrap_or(0); + let snapshot_dir = self.snapshot_dir(); + + ShellSnapshotPaths { + path: snapshot_dir.join(format!("{session_id}.{nonce}.{extension}")), + temp_path: snapshot_dir.join(format!("{session_id}.tmp-{nonce}")), + } + } + + fn cleanup_stale_snapshots( + &self, + active_session_id: ThreadId, + state_db: Option, + ) -> ShellSnapshotStoreFuture<'_> { + Box::pin(self.cleanup_stale_snapshots_impl(active_session_id, state_db)) + } +} + +/// Final and temporary local paths for one shell snapshot generation. +pub struct ShellSnapshotPaths { + /// Final path sourced by later shell execution. + pub path: AbsolutePathBuf, + /// Temporary path populated before validation and rename. + pub temp_path: AbsolutePathBuf, +} + impl ShellSnapshot { - pub fn start_snapshotting( - codex_home: AbsolutePathBuf, + pub(crate) fn start_snapshotting( + shell_snapshot_store: Arc, session_id: ThreadId, session_cwd: AbsolutePathBuf, shell: &mut Shell, @@ -48,7 +185,7 @@ impl ShellSnapshot { shell.shell_snapshot = shell_snapshot_rx; Self::spawn_snapshot_task( - codex_home, + shell_snapshot_store, session_id, session_cwd, shell.clone(), @@ -60,8 +197,8 @@ impl ShellSnapshot { shell_snapshot_tx } - pub fn refresh_snapshot( - codex_home: AbsolutePathBuf, + pub(crate) fn refresh_snapshot( + shell_snapshot_store: Arc, session_id: ThreadId, session_cwd: AbsolutePathBuf, shell: Shell, @@ -70,7 +207,7 @@ impl ShellSnapshot { state_db: Option, ) { Self::spawn_snapshot_task( - codex_home, + shell_snapshot_store, session_id, session_cwd, shell, @@ -81,7 +218,7 @@ impl ShellSnapshot { } fn spawn_snapshot_task( - codex_home: AbsolutePathBuf, + shell_snapshot_store: Arc, session_id: ThreadId, session_cwd: AbsolutePathBuf, snapshot_shell: Shell, @@ -92,13 +229,24 @@ impl ShellSnapshot { let snapshot_span = info_span!("shell_snapshot", thread_id = %session_id); tokio::spawn( async move { + let cleanup_shell_snapshot_store = Arc::clone(&shell_snapshot_store); + let cleanup_session_id = session_id; + let cleanup_state_db = state_db.clone(); + tokio::spawn(async move { + if let Err(err) = cleanup_shell_snapshot_store + .cleanup_stale_snapshots(cleanup_session_id, cleanup_state_db) + .await + { + tracing::warn!("Failed to clean up shell snapshots: {err:?}"); + } + }); + let timer = session_telemetry.start_timer("codex.shell_snapshot.duration_ms", &[]); let snapshot = ShellSnapshot::try_new( - &codex_home, + shell_snapshot_store.as_ref(), session_id, &session_cwd, &snapshot_shell, - state_db, ) .await .map(Arc::new); @@ -117,38 +265,13 @@ impl ShellSnapshot { } async fn try_new( - codex_home: &AbsolutePathBuf, + shell_snapshot_store: &dyn ShellSnapshotStore, session_id: ThreadId, session_cwd: &AbsolutePathBuf, shell: &Shell, - state_db: Option, ) -> std::result::Result { - // File to store the snapshot - let extension = match shell.shell_type { - ShellType::PowerShell => "ps1", - _ => "sh", - }; - let nonce = SystemTime::now() - .duration_since(SystemTime::UNIX_EPOCH) - .map(|duration| duration.as_nanos()) - .unwrap_or(0); - let path = codex_home - .join(SNAPSHOT_DIR) - .join(format!("{session_id}.{nonce}.{extension}")); - let temp_path = codex_home - .join(SNAPSHOT_DIR) - .join(format!("{session_id}.tmp-{nonce}")); - - // Clean the (unlikely) leaked snapshot files. - let codex_home = codex_home.clone(); - let cleanup_session_id = session_id; - tokio::spawn(async move { - if let Err(err) = - cleanup_stale_snapshots(&codex_home, cleanup_session_id, state_db).await - { - tracing::warn!("Failed to clean up shell snapshots: {err:?}"); - } - }); + let ShellSnapshotPaths { path, temp_path } = + shell_snapshot_store.snapshot_paths(session_id, shell.shell_type.clone()); // Make the new snapshot. if let Err(err) = @@ -494,72 +617,6 @@ $envVars | ForEach-Object { "## } -/// Removes shell snapshots that either lack a matching session rollout file or -/// whose rollouts have not been updated within the retention window. -/// The active session id is exempt from cleanup. -pub async fn cleanup_stale_snapshots( - codex_home: &AbsolutePathBuf, - active_session_id: ThreadId, - state_db: Option, -) -> Result<()> { - let snapshot_dir = codex_home.join(SNAPSHOT_DIR); - - let mut entries = match fs::read_dir(&snapshot_dir).await { - Ok(entries) => entries, - Err(err) if err.kind() == ErrorKind::NotFound => return Ok(()), - Err(err) => return Err(err.into()), - }; - - let now = SystemTime::now(); - let active_session_id = active_session_id.to_string(); - - while let Some(entry) = entries.next_entry().await? { - if !entry.file_type().await?.is_file() { - continue; - } - - let path = entry.path(); - - let file_name = entry.file_name(); - let file_name = file_name.to_string_lossy(); - let Some(session_id) = snapshot_session_id_from_file_name(&file_name) else { - remove_snapshot_file(&path).await; - continue; - }; - if session_id == active_session_id { - continue; - } - - let rollout_path = - find_thread_path_by_id_str(codex_home, session_id, state_db.as_deref()).await?; - let Some(rollout_path) = rollout_path else { - remove_snapshot_file(&path).await; - continue; - }; - - let modified = match fs::metadata(&rollout_path).await.and_then(|m| m.modified()) { - Ok(modified) => modified, - Err(err) => { - tracing::warn!( - "Failed to check rollout age for snapshot {}: {err:?}", - path.display() - ); - continue; - } - }; - - if now - .duration_since(modified) - .ok() - .is_some_and(|age| age >= SNAPSHOT_RETENTION) - { - remove_snapshot_file(&path).await; - } - } - - Ok(()) -} - async fn remove_snapshot_file(path: &Path) { if let Err(err) = fs::remove_file(path).await { tracing::warn!("Failed to delete shell snapshot at {:?}: {err:?}", path); diff --git a/codex-rs/core/src/shell_snapshot_tests.rs b/codex-rs/core/src/shell_snapshot_tests.rs index 0199347b4e..7cf3e486d7 100644 --- a/codex-rs/core/src/shell_snapshot_tests.rs +++ b/codex-rs/core/src/shell_snapshot_tests.rs @@ -196,13 +196,13 @@ async fn try_new_creates_and_deletes_snapshot_file() -> Result<()> { shell_path: PathBuf::from("/bin/bash"), shell_snapshot: crate::shell::empty_shell_snapshot_receiver(), }; + let shell_snapshot_store = LocalShellSnapshotStore::from_codex_home(&dir.path().abs()); let snapshot = ShellSnapshot::try_new( - &dir.path().abs(), + &shell_snapshot_store, ThreadId::new(), &dir.path().abs(), &shell, - /*state_db*/ None, ) .await .expect("snapshot should be created"); @@ -227,25 +227,16 @@ async fn try_new_uses_distinct_generation_paths() -> Result<()> { shell_path: PathBuf::from("/bin/bash"), shell_snapshot: crate::shell::empty_shell_snapshot_receiver(), }; + let shell_snapshot_store = LocalShellSnapshotStore::from_codex_home(&dir.path().abs()); - let initial_snapshot = ShellSnapshot::try_new( - &dir.path().abs(), - session_id, - &dir.path().abs(), - &shell, - /*state_db*/ None, - ) - .await - .expect("initial snapshot should be created"); - let refreshed_snapshot = ShellSnapshot::try_new( - &dir.path().abs(), - session_id, - &dir.path().abs(), - &shell, - /*state_db*/ None, - ) - .await - .expect("refreshed snapshot should be created"); + let initial_snapshot = + ShellSnapshot::try_new(&shell_snapshot_store, session_id, &dir.path().abs(), &shell) + .await + .expect("initial snapshot should be created"); + let refreshed_snapshot = + ShellSnapshot::try_new(&shell_snapshot_store, session_id, &dir.path().abs(), &shell) + .await + .expect("refreshed snapshot should be created"); let initial_path = initial_snapshot.path.clone(); let refreshed_path = refreshed_snapshot.path.clone(); @@ -439,7 +430,9 @@ async fn cleanup_stale_snapshots_removes_orphans_and_keeps_live() -> Result<()> fs::write(&orphan_snapshot, "orphan").await?; fs::write(&invalid_snapshot, "invalid").await?; - cleanup_stale_snapshots(&codex_home, ThreadId::new(), /*state_db*/ None).await?; + LocalShellSnapshotStore::from_codex_home(&codex_home) + .cleanup_stale_snapshots(ThreadId::new(), /*state_db*/ None) + .await?; assert_eq!(live_snapshot.exists(), true); assert_eq!(orphan_snapshot.exists(), false); @@ -462,7 +455,9 @@ async fn cleanup_stale_snapshots_removes_stale_rollouts() -> Result<()> { set_file_mtime(&rollout_path, SNAPSHOT_RETENTION + Duration::from_secs(60))?; - cleanup_stale_snapshots(&codex_home, ThreadId::new(), /*state_db*/ None).await?; + LocalShellSnapshotStore::from_codex_home(&codex_home) + .cleanup_stale_snapshots(ThreadId::new(), /*state_db*/ None) + .await?; assert_eq!(stale_snapshot.exists(), false); Ok(()) @@ -483,7 +478,9 @@ async fn cleanup_stale_snapshots_skips_active_session() -> Result<()> { set_file_mtime(&rollout_path, SNAPSHOT_RETENTION + Duration::from_secs(60))?; - cleanup_stale_snapshots(&codex_home, active_session, /*state_db*/ None).await?; + LocalShellSnapshotStore::from_codex_home(&codex_home) + .cleanup_stale_snapshots(active_session, /*state_db*/ None) + .await?; assert_eq!(active_snapshot.exists(), true); Ok(()) diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 728d5dca8f..6bbe63999e 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -52,6 +52,7 @@ pub(crate) struct SessionServices { pub(crate) rollout_thread_trace: ThreadTraceContext, pub(crate) user_shell: Arc, pub(crate) shell_snapshot_tx: watch::Sender>>, + pub(crate) shell_snapshot_store: Arc, pub(crate) show_raw_agent_reasoning: bool, pub(crate) exec_policy: Arc, pub(crate) auth_manager: Arc, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 49c2299b3f..1ffc79ec4c 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -15,7 +15,9 @@ use crate::session::CodexSpawnArgs; use crate::session::CodexSpawnOk; use crate::session::INITIAL_SUBMIT_ID; use crate::session::resolve_multi_agent_version; +use crate::shell_snapshot::LocalShellSnapshotStore; use crate::shell_snapshot::ShellSnapshot; +use crate::shell_snapshot::ShellSnapshotStore; use crate::tasks::InterruptedTurnHistoryMarker; use crate::tasks::interrupted_turn_history_marker; use codex_analytics::AnalyticsEventsClient; @@ -211,6 +213,7 @@ pub(crate) struct ThreadManagerState { extensions: Arc>, thread_store: Arc, artifact_store: Arc, + shell_snapshot_store: Arc, attestation_provider: Option>, session_source: SessionSource, installation_id: String, @@ -295,6 +298,68 @@ impl ThreadManager { state_db: Option, installation_id: String, attestation_provider: Option>, + ) -> Self { + Self::new_with_storage_backends( + config, + auth_manager, + session_source, + environment_manager, + extensions, + analytics_events_client, + thread_store, + artifact_store, + Arc::new(LocalShellSnapshotStore::from_codex_home(&config.codex_home)), + state_db, + installation_id, + attestation_provider, + ) + } + + #[allow(clippy::too_many_arguments)] + /// Constructs a thread manager with an explicit shell snapshot backend. + pub fn new_with_shell_snapshot_store( + config: &Config, + auth_manager: Arc, + session_source: SessionSource, + environment_manager: Arc, + extensions: Arc>, + analytics_events_client: Option, + thread_store: Arc, + shell_snapshot_store: Arc, + state_db: Option, + installation_id: String, + attestation_provider: Option>, + ) -> Self { + Self::new_with_storage_backends( + config, + auth_manager, + session_source, + environment_manager, + extensions, + analytics_events_client, + thread_store, + Arc::new(LocalArtifactStore::from_codex_home(&config.codex_home)), + shell_snapshot_store, + state_db, + installation_id, + attestation_provider, + ) + } + + #[allow(clippy::too_many_arguments)] + fn new_with_storage_backends( + config: &Config, + auth_manager: Arc, + session_source: SessionSource, + environment_manager: Arc, + extensions: Arc>, + analytics_events_client: Option, + thread_store: Arc, + artifact_store: Arc, + shell_snapshot_store: Arc, + state_db: Option, + installation_id: String, + attestation_provider: Option>, ) -> Self { let codex_home = config.codex_home.clone(); let restriction_product = session_source.restriction_product(); @@ -321,6 +386,7 @@ impl ThreadManager { extensions, thread_store, artifact_store, + shell_snapshot_store, attestation_provider, auth_manager, session_source, @@ -395,6 +461,8 @@ impl ThreadManager { restriction_product, )); let mcp_manager = Arc::new(McpManager::new(Arc::clone(&plugins_manager))); + let shell_snapshot_store: Arc = + Arc::new(LocalShellSnapshotStore::from_codex_home(&skills_codex_home)); let skills_manager = Arc::new(SkillsManager::new_with_restriction_product( skills_codex_home.clone(), /*bundled_skills_enabled*/ true, @@ -425,6 +493,7 @@ impl ThreadManager { extensions: empty_extension_registry(), thread_store, artifact_store, + shell_snapshot_store, attestation_provider: None, auth_manager, session_source: SessionSource::Exec, @@ -1372,6 +1441,7 @@ impl ThreadManagerState { analytics_events_client: self.analytics_events_client.clone(), thread_store: Arc::clone(&self.thread_store), artifact_store: Arc::clone(&self.artifact_store), + shell_snapshot_store: Arc::clone(&self.shell_snapshot_store), attestation_provider: self.attestation_provider.clone(), inherited_multi_agent_version: multi_agent_version, }))