From 801bd4d773ab374959344c36640ddbdec270ddb7 Mon Sep 17 00:00:00 2001
From: Dan Gillespie <169189814+ethernetdan-openai@users.noreply.github.com>
Date: Tue, 20 Oct 2026 18:28:10 +0000
Subject: [PATCH] Search compressed rollouts concurrently (#55652)

## What changed

Collect compressed rollout paths and scan them in up to four concurrent batches, reusing compiled search patterns within each batch. Preserve match snippets, plain rollout path keys, and per-file error handling.

## Testing

Extend the regression test to verify matches across seven compressed rollouts while skipping a corrupt file, and retain coverage for oversized search patterns and missing files.

GitOrigin-RevId: 522403654ab8577499f8f8043ccc017590c4833e
---
 codex-rs/Cargo.lock                  |  1 +
 codex-rs/rollout/Cargo.toml          |  1 +
 codex-rs/rollout/src/search.rs       | 24 +++++++++++++++++++-----
 codex-rs/rollout/src/search_tests.rs | 15 +++++++++++++--
 4 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock
index 91f53aa98132..65f02416533f 100644
--- a/codex-rs/Cargo.lock
+++ b/codex-rs/Cargo.lock
@@ -4545,6 +4545,7 @@ dependencies = [
  "codex-utils-path",
  "codex-utils-process",
  "codex-utils-string",
+ "futures",
  "libc",
  "memchr",
  "pretty_assertions",
diff --git a/codex-rs/rollout/Cargo.toml b/codex-rs/rollout/Cargo.toml
index 2bd21c925112..54100b5dd47f 100644
--- a/codex-rs/rollout/Cargo.toml
+++ b/codex-rs/rollout/Cargo.toml
@@ -27,6 +27,7 @@ codex-utils-absolute-path = { workspace = true }
 codex-utils-output-truncation = { workspace = true }
 codex-utils-path = { workspace = true }
 codex-utils-string = { workspace = true }
+futures = { workspace = true, features = ["alloc"] }
 memchr = { workspace = true }
 regex = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
diff --git a/codex-rs/rollout/src/search.rs b/codex-rs/rollout/src/search.rs
index c64f2a331025..439271830ae9 100644
--- a/codex-rs/rollout/src/search.rs
+++ b/codex-rs/rollout/src/search.rs
@@ -9,6 +9,7 @@ use codex_protocol::models::ContentItem;
 use codex_protocol::models::ResponseItem;
 use codex_protocol::protocol::EventMsg;
 use codex_protocol::protocol::strip_user_message_prefix;
+use futures::future::join_all;
 use regex::Regex;
 use regex::RegexBuilder;
 use tokio::process::Command;
@@ -223,8 +224,8 @@ async fn scan_compressed_rollout_matches(
     search_term: &str,
 ) -> io::Result<RolloutSearchMatches> {
     let mut matches = HashMap::new();
+    let mut paths = Vec::new();
     let mut dirs = vec![root.to_path_buf()];
-    let mut patterns = None;
 
     while let Some(dir) = dirs.pop() {
         let mut entries = match tokio::fs::read_dir(dir).await {
@@ -248,12 +249,21 @@ async fn scan_compressed_rollout_matches(
             if !rollout_file.is_compressed() {
                 continue;
             }
+            paths.push(rollout_file.into_path());
+        }
+    }
+
+    const COMPRESSED_SEARCH_WORKERS: usize = 4;
+    let scan = async |paths: &[PathBuf]| {
+        let mut matches = HashMap::new();
+        let mut patterns = None;
+        for path in paths {
             if let Some(snippet) =
-                first_content_match_snippet(rollout_file.path(), search_term, &mut patterns)
+                first_content_match_snippet(path.as_path(), search_term, &mut patterns)
                     .await
                     .unwrap_or_else(|err| {
                         tracing::warn!(
-                            path = %rollout_file.path().display(),
+                            path = %path.as_path().display(),
                             %err,
                             "Failed to search compressed rollout"
                         );
@@ -261,13 +271,17 @@ async fn scan_compressed_rollout_matches(
                     })
             {
                 matches.insert(
-                    compression::plain_rollout_path(rollout_file.path()),
+                    compression::plain_rollout_path(path.as_path()),
                     Some(snippet),
                 );
             }
         }
+        matches
+    };
+    let chunk_size = paths.len().div_ceil(COMPRESSED_SEARCH_WORKERS).max(1);
+    for batch in join_all(paths.chunks(chunk_size).map(scan)).await {
+        matches.extend(batch);
     }
-
     Ok(matches)
 }
 
diff --git a/codex-rs/rollout/src/search_tests.rs b/codex-rs/rollout/src/search_tests.rs
index c9925e55e9be..5ad05090ff0e 100644
--- a/codex-rs/rollout/src/search_tests.rs
+++ b/codex-rs/rollout/src/search_tests.rs
@@ -7,8 +7,19 @@ use tempfile::TempDir;
 #[tokio::test]
 async fn compressed_search_preserves_error_handling() -> io::Result<()> {
     let root = TempDir::new()?;
-    let path = root.path().join("rollout-0001.jsonl.zst");
-    std::fs::write(&path, zstd::stream::encode_all(&b""[..], 3)?)?;
+    let line = br#"{"timestamp":"2025-01-01T00:00:00Z","type":"event_msg","payload":{"type":"agent_message","message":"needle"}}"#;
+    let compressed = zstd::stream::encode_all(&line[..], 3)?;
+    let mut expected = RolloutSearchMatches::new();
+    for index in 0..7 {
+        let path = root.path().join(format!("rollout-{index:04}.jsonl"));
+        std::fs::write(path.with_extension("jsonl.zst"), &compressed)?;
+        expected.insert(path, Some("needle".to_string()));
+    }
+    std::fs::write(root.path().join("rollout-broken.jsonl.zst"), b"not zstd")?;
+    assert_eq!(
+        scan_compressed_rollout_matches(root.path(), "needle").await?,
+        expected
+    );
     let query = "a".repeat(500_000);
 
     assert!(case_insensitive_literal_regex(&query).is_err());