From 9d0eae74cd3d543c40dcdfa4e4007b43445e6543 Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 1 Sep 2026 12:13:06 +0000 Subject: [PATCH] Include shared histories in rollout compression (#42039) ## What changed - Make `local_thread_store_compression` compress cold rollout files across shared and forked histories without a separate compression mode. - Retire `local_thread_store_shared_compression` while continuing to accept it in strict configuration without changing compression behavior. - Read rollout files through the compressed-rollout reader when `codex exec resume` determines the latest turn's working directory. ## Testing - Cover compression and restoration of an archived fork chain. - Cover cwd-based resume selection when only the compressed rollout remains. - Cover the removed compatibility key in strict configuration. GitOrigin-RevId: 8edebf61eccc1f79cabb8e641f8982c76f337bf1 --- codex-rs/Cargo.lock | 2 + .../core/src/config/config_loader_tests.rs | 31 ++++ codex-rs/core/src/thread_manager.rs | 19 +-- .../core/tests/suite/rollout_compression.rs | 3 - codex-rs/exec/Cargo.toml | 2 + codex-rs/exec/src/lib.rs | 30 ++-- codex-rs/exec/tests/suite/resume.rs | 8 ++ codex-rs/features/src/lib.rs | 7 +- codex-rs/rollout/src/compression.rs | 59 ++------ codex-rs/rollout/src/compression_tests.rs | 134 ++++-------------- codex-rs/rollout/src/lib.rs | 1 - .../rollout/src/rollout_reference_index.rs | 133 +++++------------ .../src/rollout_reference_index_tests.rs | 13 -- 13 files changed, 137 insertions(+), 305 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index b59de3da7f..f5bb83c900 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2999,6 +2999,7 @@ dependencies = [ "codex-model-provider-info", "codex-otel", "codex-protocol", + "codex-rollout", "codex-utils-absolute-path", "codex-utils-cargo-bin", "codex-utils-cli", @@ -3023,6 +3024,7 @@ dependencies = [ "uuid", "walkdir", "wiremock", + "zstd", ] [[package]] diff --git a/codex-rs/core/src/config/config_loader_tests.rs b/codex-rs/core/src/config/config_loader_tests.rs index 52f5f54fb1..bb25c2c543 100644 --- a/codex-rs/core/src/config/config_loader_tests.rs +++ b/codex-rs/core/src/config/config_loader_tests.rs @@ -779,6 +779,37 @@ foo = true"#; assert_eq!(config_error.range.start.column, 1); } +#[tokio::test] +async fn strict_config_accepts_removed_shared_compression_key_without_changing_compression() { + for enabled in [false, true] { + let tmp = tempdir().expect("tempdir"); + let removed = !enabled; + std::fs::write( + tmp.path().join(CONFIG_TOML_FILE), + format!( + "[features]\nlocal_thread_store_compression = {enabled}\nlocal_thread_store_shared_compression = {removed}\n" + ), + ) + .expect("write config"); + + let config = ConfigBuilder::default() + .codex_home(tmp.path().to_path_buf()) + .fallback_cwd(Some(tmp.path().to_path_buf())) + .loader_overrides(LoaderOverrides::without_managed_config_for_tests()) + .strict_config(/*strict_config*/ true) + .build() + .await + .expect("removed feature remains accepted"); + + assert_eq!( + config + .features + .enabled(Feature::LocalThreadStoreCompression), + enabled + ); + } +} + #[test] fn strict_config_points_to_unknown_nested_key() { let tmp = tempdir().expect("tempdir"); diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 7f6648bcb8..c316178bde 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -67,7 +67,6 @@ use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::W3cTraceContext; -use codex_rollout::RolloutCompressionMode; use codex_rollout::state_db::StateDbHandle; use codex_skills_extension::HostSkillsService; use codex_thread_store::InMemoryThreadStore; @@ -385,14 +384,6 @@ pub fn thread_store_from_config( let compression_enabled = config .features .enabled(Feature::LocalThreadStoreCompression); - let compression_mode = if config - .features - .enabled(Feature::LocalThreadStoreSharedCompression) - { - RolloutCompressionMode::IncludeShared - } else { - RolloutCompressionMode::Standalone - }; let background_migration_enabled = config .features .enabled(Feature::BackgroundPaginatedRolloutMigration); @@ -409,17 +400,11 @@ pub fn thread_store_from_config( warn!("failed to migrate legacy rollouts on startup: {err}"); } if compression_enabled { - codex_rollout::spawn_rollout_compression_worker( - codex_home, - compression_mode, - ); + codex_rollout::spawn_rollout_compression_worker(codex_home); } }); } else if compression_enabled { - codex_rollout::spawn_rollout_compression_worker( - config.codex_home.to_path_buf(), - compression_mode, - ); + codex_rollout::spawn_rollout_compression_worker(config.codex_home.to_path_buf()); } store } diff --git a/codex-rs/core/tests/suite/rollout_compression.rs b/codex-rs/core/tests/suite/rollout_compression.rs index 1bef3ab3b1..e361cd375c 100644 --- a/codex-rs/core/tests/suite/rollout_compression.rs +++ b/codex-rs/core/tests/suite/rollout_compression.rs @@ -140,9 +140,6 @@ async fn compressed_shared_fork_resume_preserves_checkpoint_and_frozen_history() config .features .enable(Feature::LocalThreadStoreCompression)?; - config - .features - .enable(Feature::LocalThreadStoreSharedCompression)?; // Use production feature wiring, rather than manually writing zstd files or calling the worker. let store = codex_core::thread_store_from_config(&config, codex_core::init_state_db(&config).await); diff --git a/codex-rs/exec/Cargo.toml b/codex-rs/exec/Cargo.toml index 9c3a236bf7..83da982309 100644 --- a/codex-rs/exec/Cargo.toml +++ b/codex-rs/exec/Cargo.toml @@ -37,6 +37,7 @@ codex-login = { workspace = true } codex-model-provider-info = { workspace = true } codex-otel = { workspace = true } codex-protocol = { workspace = true } +codex-rollout = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-cli = { workspace = true } codex-utils-oss = { workspace = true } @@ -77,3 +78,4 @@ tracing-opentelemetry = { workspace = true } uuid = { workspace = true } walkdir = { workspace = true } wiremock = { workspace = true } +zstd = { workspace = true } diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index 9552613cc7..68cde2368b 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -1563,20 +1563,24 @@ async fn latest_thread_cwd(thread: &AppServerThread) -> PathBuf { } async fn parse_latest_turn_context_cwd(path: &Path) -> Option { - let text = tokio::fs::read_to_string(path).await.ok()?; - for line in text.lines().rev() { - let trimmed = line.trim(); - if trimmed.is_empty() { - continue; + let path = path.to_path_buf(); + tokio::task::spawn_blocking(move || { + let reader = codex_rollout::open_rollout_seekable_reader(&path).ok()?; + let mut scanner = codex_rollout::ReverseJsonlScanner::new(reader).ok()?; + while let Some(outcome) = scanner.scan_next::().ok()? { + if let codex_rollout::ScanOutcome::Parsed(RolloutLine { + item: RolloutItem::TurnContext(item), + .. + }) = outcome + { + return Some(item.cwd.into_path_buf()); + } } - let Ok(rollout_line) = serde_json::from_str::(trimmed) else { - continue; - }; - if let RolloutItem::TurnContext(item) = rollout_line.item { - return Some(item.cwd.into_path_buf()); - } - } - None + None + }) + .await + .ok() + .flatten() } fn cwds_match(current_cwd: &Path, session_cwd: &Path) -> bool { diff --git a/codex-rs/exec/tests/suite/resume.rs b/codex-rs/exec/tests/suite/resume.rs index ead58a07d2..d318bcf737 100644 --- a/codex-rs/exec/tests/suite/resume.rs +++ b/codex-rs/exec/tests/suite/resume.rs @@ -564,6 +564,14 @@ async fn exec_resume_last_respects_cwd_filter_and_all_flag() -> anyhow::Result<( "resume --last --all should pick newest session" ); + // Selection must still use the latest turn's cwd when only the compressed rollout exists. + zstd::stream::copy_encode( + std::fs::File::open(&path_b)?, + std::fs::File::create(path_b.with_extension("jsonl.zst"))?, + /*level*/ 3, + )?; + std::fs::remove_file(&path_b)?; + let marker_a2 = format!("resume-cwd-a-2-{}", Uuid::new_v4()); let prompt_a2 = format!("echo {marker_a2}"); test.cmd_with_server(&server) diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index cd66f8cd3c..65c91e1f06 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -171,9 +171,10 @@ pub enum Feature { MemoryTool, /// Enable importing project-scoped memory from external agents. ExternalAgentMemoryImport, - /// Compress cold local thread-store rollout files. + /// Compress cold local thread-store rollout files, including shared histories. + /// Requires every reader of the Codex home to support compressed shared histories. LocalThreadStoreCompression, - /// Allow rollout compression on homes used exclusively by compressed-lineage-aware readers. + /// Removed compatibility flag; local_thread_store_compression controls all rollout files. LocalThreadStoreSharedCompression, /// Migrate legacy local rollout files to paginated history in the background. BackgroundPaginatedRolloutMigration, @@ -1081,7 +1082,7 @@ pub const FEATURES: &[FeatureSpec] = &[ FeatureSpec { id: Feature::LocalThreadStoreSharedCompression, key: "local_thread_store_shared_compression", - stage: Stage::UnderDevelopment, + stage: Stage::Removed, default_enabled: false, }, FeatureSpec { diff --git a/codex-rs/rollout/src/compression.rs b/codex-rs/rollout/src/compression.rs index 36336c91db..2e9ed39120 100644 --- a/codex-rs/rollout/src/compression.rs +++ b/codex-rs/rollout/src/compression.rs @@ -21,22 +21,13 @@ const OPEN_ROLLOUT_LINE_READER_RETRY_DELAY: Duration = Duration::from_millis(50) const TEMP_SUFFIX: &str = ".tmp"; static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0); -/// Whether cold files belonging to shared rollout lineages may be compressed. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum RolloutCompressionMode { - /// Preserve compatibility with readers that require shared lineages to remain plain JSONL. - Standalone, - /// Requires every reader of this Codex home to support compressed shared lineages. - IncludeShared, -} - /// Starts a best-effort background job that compresses cold local rollout files. /// /// The worker is fire-and-forget: failures are logged, startup is not blocked, /// and a run marker under `codex_home` prevents overlapping or too-frequent /// compression runs from the same local store. -pub fn spawn_rollout_compression_worker(codex_home: PathBuf, mode: RolloutCompressionMode) { - worker::spawn(codex_home, mode) +pub fn spawn_rollout_compression_worker(codex_home: PathBuf) { + worker::spawn(codex_home) } /// Returns the modified time for the existing plain or compressed rollout file. @@ -253,10 +244,8 @@ mod worker { use tokio::task::JoinSet; use crate::ARCHIVED_SESSIONS_SUBDIR; - use crate::RolloutReferenceIndex; use crate::SESSIONS_SUBDIR; - use super::RolloutCompressionMode; use super::RolloutFile; use super::metrics; use super::path; @@ -334,7 +323,7 @@ mod worker { } } - pub(super) fn spawn(codex_home: PathBuf, mode: RolloutCompressionMode) { + pub(super) fn spawn(codex_home: PathBuf) { let Ok(handle) = tokio::runtime::Handle::try_current() else { metrics::run("skipped_no_runtime"); warn!( @@ -344,7 +333,7 @@ mod worker { return; }; handle.spawn(async move { - if let Err(err) = run(codex_home.clone(), mode).await { + if let Err(err) = run(codex_home.clone()).await { warn!( "rollout compression worker failed for {}: {err}", codex_home.display() @@ -353,7 +342,7 @@ mod worker { }); } - pub(super) async fn run(codex_home: PathBuf, mode: RolloutCompressionMode) -> io::Result<()> { + pub(super) async fn run(codex_home: PathBuf) -> io::Result<()> { let Some(_maintenance_guard) = crate::try_acquire_rollout_maintenance_lock(codex_home.as_path())? else { @@ -384,15 +373,6 @@ mod worker { let started_at = Instant::now(); let result = async { cleanup_stale_temps(codex_home.as_path()).await?; - let Some(reference_index) = RolloutReferenceIndex::scan_until( - codex_home.as_path(), - started_at, - WORKER_MAX_RUNTIME, - ) - .await? - else { - return Ok(CompressionStats::default()); - }; let mut stats = CompressionStats::default(); for root in [ codex_home.join(ARCHIVED_SESSIONS_SUBDIR), @@ -401,14 +381,7 @@ mod worker { if started_at.elapsed() >= WORKER_MAX_RUNTIME { break; } - compress_rollouts_in_root( - root.as_path(), - started_at, - &reference_index, - mode, - &mut stats, - ) - .await?; + compress_rollouts_in_root(root.as_path(), started_at, &mut stats).await?; } Ok::<_, io::Error>(stats) } @@ -448,8 +421,6 @@ mod worker { async fn compress_rollouts_in_root( root: &Path, started_at: Instant, - reference_index: &RolloutReferenceIndex, - mode: RolloutCompressionMode, stats: &mut CompressionStats, ) -> io::Result<()> { if !tokio::fs::try_exists(root).await.unwrap_or(false) { @@ -508,26 +479,14 @@ mod worker { continue; } let path = rollout_file.into_path(); - let Some(rollout_id) = crate::rollout_id_from_path(path.as_path()) else { + if crate::rollout_id_from_path(path.as_path()).is_none() { stats.skipped = stats.skipped.saturating_add(1); metrics::file("skipped_unreadable_meta"); continue; - }; - let Ok(meta) = crate::read_session_meta_line(path.as_path()).await else { - stats.skipped = stats.skipped.saturating_add(1); - metrics::file("skipped_unreadable_meta"); - continue; - }; - if mode == RolloutCompressionMode::Standalone - && reference_index.reference_count(rollout_id) > 0 - { - stats.skipped = stats.skipped.saturating_add(1); - metrics::file("skipped_referenced"); - continue; } - if mode == RolloutCompressionMode::Standalone && meta.meta.history_base.is_some() { + if crate::read_session_meta_line(path.as_path()).await.is_err() { stats.skipped = stats.skipped.saturating_add(1); - metrics::file("skipped_fork_pointer"); + metrics::file("skipped_unreadable_meta"); continue; } stats.scanned = stats.scanned.saturating_add(1); diff --git a/codex-rs/rollout/src/compression_tests.rs b/codex-rs/rollout/src/compression_tests.rs index f9d7a74921..d2ac7450a6 100644 --- a/codex-rs/rollout/src/compression_tests.rs +++ b/codex-rs/rollout/src/compression_tests.rs @@ -289,11 +289,7 @@ async fn worker_compresses_old_active_and_archived_rollouts() -> anyhow::Result< let fresh_temp = active_path.with_file_name("rollout-fresh.jsonl.zst.tmp"); fs::write(&fresh_temp, "fresh temp")?; - worker::run( - home.path().to_path_buf(), - RolloutCompressionMode::Standalone, - ) - .await?; + worker::run(home.path().to_path_buf()).await?; assert!(!active_path.exists()); assert!(compressed_rollout_path(&active_path).exists()); @@ -323,105 +319,53 @@ async fn worker_waits_for_rollout_maintenance_before_compressing() -> anyhow::Re let guard = crate::try_acquire_rollout_maintenance_lock(home.path())? .expect("claim rollout maintenance lock"); - worker::run( - home.path().to_path_buf(), - RolloutCompressionMode::Standalone, - ) - .await?; + worker::run(home.path().to_path_buf()).await?; assert!(path.exists()); assert!(!compressed_rollout_path(&path).exists()); drop(guard); - worker::run( - home.path().to_path_buf(), - RolloutCompressionMode::Standalone, - ) - .await?; + worker::run(home.path().to_path_buf()).await?; assert!(!path.exists()); assert!(compressed_rollout_path(&path).exists()); Ok(()) } #[tokio::test] -async fn worker_compresses_archived_fork_chain_only_with_shared_mode() -> anyhow::Result<()> { - for mode in [ - RolloutCompressionMode::Standalone, - RolloutCompressionMode::IncludeShared, - ] { - // Each mode gets a fresh home, without bypassing the worker's maintenance cooldown. - let home = TempDir::new()?; - let thread_id = ThreadId::from_string(&Uuid::from_u128(15).to_string())?; - let source_uuid = Uuid::from_u128(16); - let source_rollout_id = ThreadId::from_string(&source_uuid.to_string())?; - let source_path = rollout_path(home.path(), "2025-01-03T12-00-00", source_uuid); - write_rollout(&source_path, thread_id, "referenced source")?; - set_old_mtime(&source_path)?; - - let child_uuid = Uuid::from_u128(17); - let child_path = archived_rollout_path(home.path(), "2025-01-03T12-00-01", child_uuid); - write_rollout(&child_path, thread_id, "fork child")?; - set_history_base( - child_path.as_path(), - HistoryPosition { - thread_id: source_rollout_id, - end_ordinal_exclusive: 2, - end_byte_offset: std::fs::metadata(source_path.as_path())?.len(), - }, - )?; - set_old_mtime(&child_path)?; - let original_source = fs::read(&source_path)?; - let original_child = fs::read(&child_path)?; - - worker::run(home.path().to_path_buf(), mode).await?; - - for (path, original) in [ - (&source_path, original_source), - (&child_path, original_child), - ] { - match mode { - RolloutCompressionMode::Standalone => assert_eq!(fs::read(path)?, original), - RolloutCompressionMode::IncludeShared => { - assert!(!path.exists(), "shared compression should replace {path:?}"); - let mut restored = Vec::new(); - crate::open_rollout_seekable_reader(path)?.read_to_end(&mut restored)?; - assert_eq!(restored, original); - } - } - } - } - Ok(()) -} - -#[tokio::test] -async fn worker_skips_source_referenced_by_archived_compressed_rollout() -> anyhow::Result<()> { +async fn worker_compresses_archived_fork_chain() -> anyhow::Result<()> { let home = TempDir::new()?; - let source_uuid = Uuid::from_u128(18); - let source_id = ThreadId::from_string(&source_uuid.to_string())?; + let thread_id = ThreadId::from_string(&Uuid::from_u128(15).to_string())?; + let source_uuid = Uuid::from_u128(16); + let source_rollout_id = ThreadId::from_string(&source_uuid.to_string())?; let source_path = rollout_path(home.path(), "2025-01-03T12-00-00", source_uuid); - write_rollout(&source_path, source_id, "referenced source")?; + write_rollout(&source_path, thread_id, "referenced source")?; set_old_mtime(&source_path)?; - let child_uuid = Uuid::from_u128(19); - let child_id = ThreadId::from_string(&child_uuid.to_string())?; + let child_uuid = Uuid::from_u128(17); let child_path = archived_rollout_path(home.path(), "2025-01-03T12-00-01", child_uuid); - write_rollout(&child_path, child_id, "fork child")?; + write_rollout(&child_path, thread_id, "fork child")?; set_history_base( child_path.as_path(), HistoryPosition { - thread_id: source_id, + thread_id: source_rollout_id, end_ordinal_exclusive: 2, end_byte_offset: std::fs::metadata(source_path.as_path())?.len(), }, )?; - compress_now(child_path.as_path())?; + set_old_mtime(&child_path)?; + let original_source = fs::read(&source_path)?; + let original_child = fs::read(&child_path)?; - worker::run( - home.path().to_path_buf(), - RolloutCompressionMode::Standalone, - ) - .await?; + worker::run(home.path().to_path_buf()).await?; - assert!(source_path.exists()); + for (path, original) in [ + (&source_path, original_source), + (&child_path, original_child), + ] { + assert!(!path.exists(), "compression should replace {path:?}"); + let mut restored = Vec::new(); + crate::open_rollout_seekable_reader(path)?.read_to_end(&mut restored)?; + assert_eq!(restored, original); + } Ok(()) } @@ -438,11 +382,7 @@ async fn worker_skips_unreadable_metadata_without_blocking_other_compression() - let unreadable_path = rollout_path(home.path(), "2025-01-03T12-00-01", Uuid::from_u128(21)); fs::write(unreadable_path.as_path(), "{not json}\n")?; - worker::run( - home.path().to_path_buf(), - RolloutCompressionMode::Standalone, - ) - .await?; + worker::run(home.path().to_path_buf()).await?; assert!(!source_path.exists()); assert!(compressed_rollout_path(&source_path).exists()); @@ -516,11 +456,7 @@ async fn compression_preserves_rollout_permissions() -> anyhow::Result<()> { fs::set_permissions(&rollout_path, fs::Permissions::from_mode(0o600))?; set_old_mtime(&rollout_path)?; - worker::run( - home.path().to_path_buf(), - RolloutCompressionMode::Standalone, - ) - .await?; + worker::run(home.path().to_path_buf()).await?; let compressed_path = compressed_rollout_path(&rollout_path); assert!(!rollout_path.exists()); @@ -602,11 +538,7 @@ async fn compression_preserves_read_only_rollout_permissions() -> anyhow::Result fs::set_permissions(&rollout_path, fs::Permissions::from_mode(0o400))?; let source_modified = fs::metadata(&rollout_path)?.modified()?; - worker::run( - home.path().to_path_buf(), - RolloutCompressionMode::Standalone, - ) - .await?; + worker::run(home.path().to_path_buf()).await?; let compressed_path = compressed_rollout_path(&rollout_path); let compressed_metadata = fs::metadata(&compressed_path)?; @@ -627,11 +559,7 @@ async fn worker_skips_existing_compressed_archived_rollouts() -> anyhow::Result< let compressed_path = compressed_rollout_path(&rollout_path); set_old_mtime(&compressed_path)?; - worker::run( - home.path().to_path_buf(), - RolloutCompressionMode::Standalone, - ) - .await?; + worker::run(home.path().to_path_buf()).await?; assert!(!rollout_path.exists()); assert!(compressed_path.exists()); @@ -655,11 +583,7 @@ async fn worker_skips_when_fresh_run_marker_exists() -> anyhow::Result<()> { fs::create_dir_all(marker_dir.as_path())?; fs::write(marker_dir.join("rollout-compression.lock"), "recent run")?; - worker::run( - home.path().to_path_buf(), - RolloutCompressionMode::Standalone, - ) - .await?; + worker::run(home.path().to_path_buf()).await?; assert!(rollout_path.exists()); assert!(!compressed_rollout_path(&rollout_path).exists()); diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index e319c933c1..20f966ae90 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -77,7 +77,6 @@ pub static INTERACTIVE_SESSION_SOURCES: LazyLock> = LazyLock: }); pub use codex_protocol::protocol::SessionMeta; -pub use compression::RolloutCompressionMode; pub use compression::RolloutLineReader; pub use compression::existing_rollout_path; pub use compression::open_rollout_line_reader; diff --git a/codex-rs/rollout/src/rollout_reference_index.rs b/codex-rs/rollout/src/rollout_reference_index.rs index 7a16082eb6..204e05f180 100644 --- a/codex-rs/rollout/src/rollout_reference_index.rs +++ b/codex-rs/rollout/src/rollout_reference_index.rs @@ -5,8 +5,6 @@ use std::collections::hash_map::Entry; use std::io; use std::path::Path; use std::path::PathBuf; -use std::time::Duration; -use std::time::Instant; use codex_protocol::RolloutId; use codex_protocol::ThreadId; @@ -34,22 +32,13 @@ struct IndexedRollout { } impl RolloutReferenceIndex { - /// Scans active and archived local rollout metadata without a deadline. + /// Scans active and archived local rollout metadata. pub async fn scan(codex_home: &Path) -> io::Result { - let Some(index) = Self::scan_with_deadline( - vec![ - codex_home.join(ARCHIVED_SESSIONS_SUBDIR), - codex_home.join(SESSIONS_SUBDIR), - ], - ScanDeadline::Unlimited, - ) - .await? - else { - return Err(io::Error::other( - "unlimited rollout reference scan exceeded a deadline", - )); - }; - Ok(index) + Self::scan_paths(vec![ + codex_home.join(ARCHIVED_SESSIONS_SUBDIR), + codex_home.join(SESSIONS_SUBDIR), + ]) + .await } /// Scans only unarchived rollouts to locate files that still need to be archived. @@ -57,79 +46,18 @@ impl RolloutReferenceIndex { /// Reference counts exclude archived history and must not be used to decide whether a /// rollout can be deleted or compressed. pub async fn scan_unarchived(codex_home: &Path) -> io::Result { - Self::scan_with_deadline( - vec![codex_home.join(SESSIONS_SUBDIR)], - ScanDeadline::Unlimited, - ) - .await? - .ok_or_else(|| io::Error::other("unlimited rollout reference scan exceeded a deadline")) + Self::scan_paths(vec![codex_home.join(SESSIONS_SUBDIR)]).await } - /// Scans active and archived local rollout metadata until the worker deadline expires. - /// - /// Returns None instead of a partial index when the deadline expires. - pub(crate) async fn scan_until( - codex_home: &Path, - started_at: Instant, - max_runtime: Duration, - ) -> io::Result> { - Self::scan_with_deadline( - vec![ - codex_home.join(ARCHIVED_SESSIONS_SUBDIR), - codex_home.join(SESSIONS_SUBDIR), - ], - ScanDeadline::Until { - started_at, - max_runtime, - }, - ) - .await - } - - /// Returns how many other discovered rollouts directly reference `rollout_id`. - pub fn reference_count(&self, rollout_id: RolloutId) -> usize { - self.reference_counts_by_rollout - .get(&rollout_id) - .copied() - .unwrap_or_default() - } - - /// Returns the direct history-base edge for `rollout_id`, if one was discovered. - pub fn history_base(&self, rollout_id: RolloutId) -> Option<&HistoryPosition> { - self.rollouts_by_id - .get(&rollout_id) - .and_then(|rollout| rollout.history_base.as_ref()) - } - - /// Returns rollout IDs and paths whose session metadata belongs to `thread_id`. - pub fn rollouts_for_thread( - &self, - thread_id: ThreadId, - ) -> impl Iterator { - self.rollouts_by_id - .iter() - .filter(move |(_, rollout)| rollout.thread_id == thread_id) - .map(|(rollout_id, rollout)| (*rollout_id, rollout.path.as_path())) - } - - async fn scan_with_deadline( - mut stack: Vec, - deadline: ScanDeadline, - ) -> io::Result> { + async fn scan_paths(mut stack: Vec) -> io::Result { let mut rollouts_by_id = HashMap::new(); while let Some(directory) = stack.pop() { - if deadline.expired() { - return Ok(None); - } let mut entries = match tokio::fs::read_dir(directory.as_path()).await { Ok(entries) => entries, Err(err) if err.kind() == io::ErrorKind::NotFound => continue, Err(err) => return Err(err), }; loop { - if deadline.expired() { - return Ok(None); - } let Some(entry) = entries.next_entry().await? else { break; }; @@ -173,31 +101,36 @@ impl RolloutReferenceIndex { .entry(history_base.thread_id) .or_default() += 1; } - Ok(Some(Self { + Ok(Self { rollouts_by_id, reference_counts_by_rollout, - })) + }) } -} -#[derive(Clone, Copy)] -enum ScanDeadline { - Unlimited, - Until { - started_at: Instant, - max_runtime: Duration, - }, -} + /// Returns how many other discovered rollouts directly reference `rollout_id`. + pub fn reference_count(&self, rollout_id: RolloutId) -> usize { + self.reference_counts_by_rollout + .get(&rollout_id) + .copied() + .unwrap_or_default() + } -impl ScanDeadline { - fn expired(self) -> bool { - match self { - Self::Unlimited => false, - Self::Until { - started_at, - max_runtime, - } => started_at.elapsed() >= max_runtime, - } + /// Returns the direct history-base edge for `rollout_id`, if one was discovered. + pub fn history_base(&self, rollout_id: RolloutId) -> Option<&HistoryPosition> { + self.rollouts_by_id + .get(&rollout_id) + .and_then(|rollout| rollout.history_base.as_ref()) + } + + /// Returns rollout IDs and paths whose session metadata belongs to `thread_id`. + pub fn rollouts_for_thread( + &self, + thread_id: ThreadId, + ) -> impl Iterator { + self.rollouts_by_id + .iter() + .filter(move |(_, rollout)| rollout.thread_id == thread_id) + .map(|(rollout_id, rollout)| (*rollout_id, rollout.path.as_path())) } } diff --git a/codex-rs/rollout/src/rollout_reference_index_tests.rs b/codex-rs/rollout/src/rollout_reference_index_tests.rs index b82fdfbea5..091399fa00 100644 --- a/codex-rs/rollout/src/rollout_reference_index_tests.rs +++ b/codex-rs/rollout/src/rollout_reference_index_tests.rs @@ -1,8 +1,6 @@ use std::fs; use std::path::Path; use std::path::PathBuf; -use std::time::Duration; -use std::time::Instant; use codex_protocol::ThreadId; use codex_protocol::protocol::HistoryPosition; @@ -171,17 +169,6 @@ async fn self_history_base_does_not_count_as_reference() -> anyhow::Result<()> { Ok(()) } -#[tokio::test] -async fn expired_deadline_returns_none() -> anyhow::Result<()> { - let home = TempDir::new()?; - - let index = - RolloutReferenceIndex::scan_until(home.path(), Instant::now(), Duration::ZERO).await?; - - assert!(index.is_none()); - Ok(()) -} - fn active_rollout_path(home: &Path, uuid: Uuid) -> PathBuf { home.join("sessions/2025/01/03") .join(format!("rollout-2025-01-03T12-00-00-{uuid}.jsonl"))