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