openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 17:38:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor shell snapshot store seam

Browse Source
This commit is contained in:
starr-openai
2026-06-03 00:19:03 -07:00
parent eeb7b76548
commit 8f4d1c2ce6
11 changed files with 280 additions and 143 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/codex_delegate.rs
+1
View File
@@ -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),
}))
codex-rs/core/src/lib.rs
+1 -1
View File
@@ -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;
codex-rs/core/src/session/mod.rs
+5 -11
View File
@@ -421,6 +421,7 @@ pub(crate) struct CodexSpawnArgs {
pub(crate) analytics_events_client: Option<AnalyticsEventsClient>,
pub(crate) thread_store: Arc<dyn ThreadStore>,
pub(crate) artifact_store: Arc<dyn ArtifactStore>,
pub(crate) shell_snapshot_store: Arc<dyn crate::shell_snapshot::ShellSnapshotStore>,
pub(crate) attestation_provider: Option<Arc<dyn AttestationProvider>>,
pub(crate) inherited_multi_agent_version: Option<MultiAgentVersion>,
}
@@ -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;
codex-rs/core/src/session/session.rs
+3 -1
View File
@@ -506,6 +506,7 @@ impl Session {
analytics_events_client: Option<AnalyticsEventsClient>,
thread_store: Arc<dyn ThreadStore>,
artifact_store: Arc<dyn ArtifactStore>,
shell_snapshot_store: Arc<dyn crate::shell_snapshot::ShellSnapshotStore>,
parent_rollout_thread_trace: ThreadTraceContext,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
multi_agent_version: Option<MultiAgentVersion>,
@@ -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),
codex-rs/core/src/session/tests.rs
+15
View File
@@ -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),
codex-rs/core/src/session/tests/guardian_tests.rs
+4
View File
@@ -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,
})
codex-rs/core/src/session/turn_context.rs
-4
View File
@@ -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,
);
codex-rs/core/src/shell_snapshot.rs
+160 -103
View File
@@ -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<StateDbHandle>,
) -> ShellSnapshotStoreFuture<'_>;
}
/// Future returned by shell snapshot store cleanup work.
pub type ShellSnapshotStoreFuture<'a> = Pin<Box<dyn Future<Output = Result<()>> + 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<StateDbHandle>,
) -> 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<StateDbHandle>,
) -> 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<dyn ShellSnapshotStore>,
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<dyn ShellSnapshotStore>,
session_id: ThreadId,
session_cwd: AbsolutePathBuf,
shell: Shell,
@@ -70,7 +207,7 @@ impl ShellSnapshot {
state_db: Option<StateDbHandle>,
) {
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<dyn ShellSnapshotStore>,
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<StateDbHandle>,
) -> std::result::Result<Self, &'static str> {
// 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<StateDbHandle>,
) -> 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);
codex-rs/core/src/shell_snapshot_tests.rs
+20 -23
View File
@@ -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(())
codex-rs/core/src/state/service.rs
+1
View File
@@ -52,6 +52,7 @@ pub(crate) struct SessionServices {
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_store: Arc<dyn crate::shell_snapshot::ShellSnapshotStore>,
pub(crate) show_raw_agent_reasoning: bool,
pub(crate) exec_policy: Arc<ExecPolicyManager>,
pub(crate) auth_manager: Arc<AuthManager>,
codex-rs/core/src/thread_manager.rs
+70
View File
@@ -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<ExtensionRegistry<Config>>,
thread_store: Arc<dyn ThreadStore>,
artifact_store: Arc<dyn ArtifactStore>,
shell_snapshot_store: Arc<dyn ShellSnapshotStore>,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
session_source: SessionSource,
installation_id: String,
@@ -295,6 +298,68 @@ impl ThreadManager {
state_db: Option<StateDbHandle>,
installation_id: String,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
) -> 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<AuthManager>,
session_source: SessionSource,
environment_manager: Arc<EnvironmentManager>,
extensions: Arc<ExtensionRegistry<Config>>,
analytics_events_client: Option<AnalyticsEventsClient>,
thread_store: Arc<dyn ThreadStore>,
shell_snapshot_store: Arc<dyn ShellSnapshotStore>,
state_db: Option<StateDbHandle>,
installation_id: String,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
) -> 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<AuthManager>,
session_source: SessionSource,
environment_manager: Arc<EnvironmentManager>,
extensions: Arc<ExtensionRegistry<Config>>,
analytics_events_client: Option<AnalyticsEventsClient>,
thread_store: Arc<dyn ThreadStore>,
artifact_store: Arc<dyn ArtifactStore>,
shell_snapshot_store: Arc<dyn ShellSnapshotStore>,
state_db: Option<StateDbHandle>,
installation_id: String,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
) -> 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<dyn ShellSnapshotStore> =
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,
}))