From d0a8dcd15721adb02a3227d08d3cde888e5546e5 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Mon, 7 Sep 2026 15:01:50 +0000 Subject: [PATCH] Limit archive rollout reads to requested threads (#43494) ## Why Archiving threads reads metadata from every active rollout, including unrelated compressed files. ## What changed Add `RolloutReferenceIndex::scan_unarchived_threads` and use it when archiving. Filter canonical filenames by the requested thread IDs before reading contents, while retaining metadata-based ownership checks. ## Testing Add tests covering reverted rollouts, requested descendants, compressed rollouts, exclusion of archived and unrelated files, and candidate ownership checks. GitOrigin-RevId: a64f7fa25e5493ed20312f1d8216422752272bae --- .../rollout/src/rollout_reference_index.rs | 43 +++++++-- .../src/rollout_reference_index_tests.rs | 96 +++++++++++++++++++ .../thread-store/src/local/archive_thread.rs | 15 +-- 3 files changed, 141 insertions(+), 13 deletions(-) diff --git a/codex-rs/rollout/src/rollout_reference_index.rs b/codex-rs/rollout/src/rollout_reference_index.rs index 204e05f180..a55d2344f9 100644 --- a/codex-rs/rollout/src/rollout_reference_index.rs +++ b/codex-rs/rollout/src/rollout_reference_index.rs @@ -1,6 +1,7 @@ //! Indexes direct fork references found in local rollout files. use std::collections::HashMap; +use std::collections::HashSet; use std::collections::hash_map::Entry; use std::io; use std::path::Path; @@ -13,6 +14,7 @@ use codex_protocol::protocol::HistoryPosition; use crate::ARCHIVED_SESSIONS_SUBDIR; use crate::SESSIONS_SUBDIR; use crate::compression::RolloutFile; +use crate::rollout_file_name::RolloutFileName; /// Direct history-base edges discovered from local rollout metadata. /// @@ -34,10 +36,13 @@ struct IndexedRollout { impl RolloutReferenceIndex { /// Scans active and archived local rollout metadata. pub async fn scan(codex_home: &Path) -> io::Result { - Self::scan_paths(vec![ - codex_home.join(ARCHIVED_SESSIONS_SUBDIR), - codex_home.join(SESSIONS_SUBDIR), - ]) + Self::scan_paths( + vec![ + codex_home.join(ARCHIVED_SESSIONS_SUBDIR), + codex_home.join(SESSIONS_SUBDIR), + ], + /*thread_ids*/ None, + ) .await } @@ -46,10 +51,30 @@ 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_paths(vec![codex_home.join(SESSIONS_SUBDIR)]).await + Self::scan_paths( + vec![codex_home.join(SESSIONS_SUBDIR)], + /*thread_ids*/ None, + ) + .await } - async fn scan_paths(mut stack: Vec) -> io::Result { + /// Scans unarchived files whose canonical filenames belong to the requested threads. + /// + /// Skips unrelated rollout contents, including compressed files. Metadata still determines + /// ownership among the candidates. Reference counts are partial and must not be used to + /// decide whether a rollout can be deleted or compressed. + pub async fn scan_unarchived_threads( + codex_home: &Path, + thread_ids: &[ThreadId], + ) -> io::Result { + let thread_ids = thread_ids.iter().copied().collect(); + Self::scan_paths(vec![codex_home.join(SESSIONS_SUBDIR)], Some(&thread_ids)).await + } + + async fn scan_paths( + mut stack: Vec, + thread_ids: Option<&HashSet>, + ) -> io::Result { let mut rollouts_by_id = HashMap::new(); while let Some(directory) = stack.pop() { let mut entries = match tokio::fs::read_dir(directory.as_path()).await { @@ -73,9 +98,13 @@ impl RolloutReferenceIndex { let Some(rollout_file) = RolloutFile::from_path(path) else { continue; }; - let Some(rollout_id) = crate::rollout_id_from_path(rollout_file.path()) else { + let Some(file_name) = RolloutFileName::parse(rollout_file.plain_file_name()) else { continue; }; + if thread_ids.is_some_and(|ids| !ids.contains(&file_name.thread_id())) { + continue; + } + let rollout_id = file_name.rollout_id(); let Ok(meta) = crate::read_session_meta_line(rollout_file.path()).await else { continue; }; diff --git a/codex-rs/rollout/src/rollout_reference_index_tests.rs b/codex-rs/rollout/src/rollout_reference_index_tests.rs index 091399fa00..036e870724 100644 --- a/codex-rs/rollout/src/rollout_reference_index_tests.rs +++ b/codex-rs/rollout/src/rollout_reference_index_tests.rs @@ -151,6 +151,102 @@ async fn unarchived_scan_finds_all_active_rollouts_owned_by_a_thread() -> anyhow Ok(()) } +#[tokio::test] +async fn filtered_unarchived_scan_includes_reverts_and_requested_descendants() -> anyhow::Result<()> +{ + let home = TempDir::new()?; + let owner_uuid = Uuid::from_u128(40); + let owner_id = thread_id(owner_uuid)?; + let replacement_id = thread_id(Uuid::from_u128(41))?; + let descendant_uuid = Uuid::from_u128(42); + let descendant_id = thread_id(descendant_uuid)?; + let unrelated_uuid = Uuid::from_u128(43); + let unrelated_id = thread_id(unrelated_uuid)?; + let original_path = active_rollout_path(home.path(), owner_uuid); + let replacement_path = original_path.with_file_name(format!( + "rollout-2025-01-03T12-00-00-{owner_id}_{replacement_id}.jsonl" + )); + let descendant_path = active_rollout_path(home.path(), descendant_uuid); + + write_rollout(original_path.clone(), owner_id, /*history_base*/ None)?; + compress_now(&original_path)?; + write_rollout( + replacement_path.clone(), + owner_id, + /*history_base*/ None, + )?; + // A leftover compressed sibling must not supersede the plain rollout. + compress_now(&replacement_path)?; + write_rollout( + replacement_path.clone(), + owner_id, + /*history_base*/ None, + )?; + write_rollout( + descendant_path.clone(), + descendant_id, + Some(history_position(owner_id)), + )?; + write_rollout( + archived_rollout_path(home.path(), owner_uuid), + owner_id, + /*history_base*/ None, + )?; + let unrelated_path = active_rollout_path(home.path(), unrelated_uuid); + write_rollout( + unrelated_path.clone(), + unrelated_id, + /*history_base*/ None, + )?; + compress_now(&unrelated_path)?; + + let index = + RolloutReferenceIndex::scan_unarchived_threads(home.path(), &[owner_id, descendant_id]) + .await?; + let mut owned: Vec<_> = index + .rollouts_for_thread(owner_id) + .map(|(id, path)| (id, path.to_path_buf())) + .collect(); + owned.sort_by_key(|(id, _)| id.to_string()); + assert_eq!( + owned, + vec![ + (owner_id, original_path.with_extension("jsonl.zst")), + (replacement_id, replacement_path), + ] + ); + assert_eq!( + index.rollouts_for_thread(descendant_id).collect::>(), + vec![(descendant_id, descendant_path.as_path())] + ); + assert_eq!( + index.rollouts_for_thread(unrelated_id).collect::>(), + Vec::new() + ); + Ok(()) +} + +#[tokio::test] +async fn filtered_unarchived_scan_validates_candidate_ownership() -> anyhow::Result<()> { + let home = TempDir::new()?; + let requested_uuid = Uuid::from_u128(50); + let requested_id = thread_id(requested_uuid)?; + let actual_owner = thread_id(Uuid::from_u128(51))?; + write_rollout( + active_rollout_path(home.path(), requested_uuid), + actual_owner, + /*history_base*/ None, + )?; + + let index = + RolloutReferenceIndex::scan_unarchived_threads(home.path(), &[requested_id]).await?; + assert_eq!( + index.rollouts_for_thread(requested_id).collect::>(), + Vec::new() + ); + Ok(()) +} + #[tokio::test] async fn self_history_base_does_not_count_as_reference() -> anyhow::Result<()> { let home = TempDir::new()?; diff --git a/codex-rs/thread-store/src/local/archive_thread.rs b/codex-rs/thread-store/src/local/archive_thread.rs index c1ac4a7946..075a0a873f 100644 --- a/codex-rs/thread-store/src/local/archive_thread.rs +++ b/codex-rs/thread-store/src/local/archive_thread.rs @@ -39,12 +39,15 @@ pub(super) async fn archive_threads( } } let _writer_guards = store.acquire_writer_locks(&lock_thread_ids).await?; - // Already-archived rollouts need no move. Avoid reading the entire archive on every request. - let reference_index = RolloutReferenceIndex::scan_unarchived(store.config.codex_home.as_path()) - .await - .map_err(|err| ThreadStoreError::Internal { - message: format!("failed to scan thread rollout files: {err}"), - })?; + // Only inspect active files whose names belong to the threads being archived. + let reference_index = RolloutReferenceIndex::scan_unarchived_threads( + store.config.codex_home.as_path(), + &thread_ids, + ) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to scan thread rollout files: {err}"), + })?; let parent_thread_id = thread_ids[0]; let mut archived_thread_ids = Vec::new();