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

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> {
-        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> {
-        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<PathBuf>) -> io::Result<Self> {
+    /// 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<Self> {
+        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<PathBuf>,
+        thread_ids: Option<&HashSet<ThreadId>>,
+    ) -> io::Result<Self> {
         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;
                 };

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<_>>(),
+        vec![(descendant_id, descendant_path.as_path())]
+    );
+    assert_eq!(
+        index.rollouts_for_thread(unrelated_id).collect::<Vec<_>>(),
+        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<_>>(),
+        Vec::new()
+    );
+    Ok(())
+}
+
 #[tokio::test]
 async fn self_history_base_does_not_count_as_reference() -> anyhow::Result<()> {
     let home = TempDir::new()?;

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();