From bb51c2b8458c3148ff8ac0e3b5bc4d3ce907a3d2 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 2 Jun 2026 10:10:56 +0100 Subject: [PATCH] reuse compressed search snippets --- codex-rs/rollout/src/compression_tests.rs | 26 +++++++ codex-rs/rollout/src/lib.rs | 1 + codex-rs/rollout/src/search.rs | 75 ++++++++++++++----- .../thread-store/src/local/search_threads.rs | 23 +++--- 4 files changed, 93 insertions(+), 32 deletions(-) diff --git a/codex-rs/rollout/src/compression_tests.rs b/codex-rs/rollout/src/compression_tests.rs index 8eb238053656..52094c2b8a73 100644 --- a/codex-rs/rollout/src/compression_tests.rs +++ b/codex-rs/rollout/src/compression_tests.rs @@ -23,6 +23,7 @@ use crate::RolloutConfig; use crate::RolloutRecorder; use crate::RolloutRecorderParams; use crate::append_rollout_item_to_path; +use crate::search_rollout_matches; #[tokio::test] async fn load_rollout_items_reads_compressed_rollout() -> anyhow::Result<()> { @@ -104,6 +105,31 @@ async fn append_rollout_item_materializes_compressed_rollout() -> anyhow::Result Ok(()) } +#[tokio::test] +async fn search_rollout_matches_returns_compressed_snippet() -> anyhow::Result<()> { + let home = TempDir::new()?; + let uuid = Uuid::from_u128(15); + let thread_id = ThreadId::from_string(&uuid.to_string())?; + let rollout_path = rollout_path(home.path(), "2025-01-03T12-00-00", uuid); + write_rollout(&rollout_path, thread_id, "targeted search term")?; + compress_now(&rollout_path)?; + let compressed_path = compressed_rollout_path(&rollout_path); + + let matches = search_rollout_matches( + std::path::Path::new("missing-rg-for-test"), + home.path(), + /*archived*/ false, + "search term", + ) + .await?; + + assert_eq!( + matches.get(compressed_path.as_path()), + Some(&Some("targeted search term".to_string())) + ); + Ok(()) +} + #[tokio::test] async fn worker_compresses_old_archived_rollouts_only() -> anyhow::Result<()> { let home = TempDir::new()?; diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index ec9fc27f34e8..9e6db0fb8db7 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -67,6 +67,7 @@ pub use recorder::RolloutRecorder; pub use recorder::RolloutRecorderParams; pub use recorder::append_rollout_item_to_path; pub use search::first_rollout_content_match_snippet; +pub use search::search_rollout_matches; pub use search::search_rollout_paths; pub use session_index::append_thread_name; pub use session_index::find_thread_meta_by_name_str; diff --git a/codex-rs/rollout/src/search.rs b/codex-rs/rollout/src/search.rs index 2270913548f5..705d24e7b4fb 100644 --- a/codex-rs/rollout/src/search.rs +++ b/codex-rs/rollout/src/search.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::collections::HashSet; use std::io; use std::path::Path; @@ -20,21 +21,42 @@ use super::compression; const MATCH_CONTEXT_BEFORE_CHARS: usize = 48; const MATCH_CONTEXT_AFTER_CHARS: usize = 96; +pub type RolloutSearchMatches = HashMap>; + pub async fn search_rollout_paths( rg_command: &Path, codex_home: &Path, archived: bool, search_term: &str, ) -> io::Result> { + Ok( + search_rollout_matches(rg_command, codex_home, archived, search_term) + .await? + .into_keys() + .collect(), + ) +} + +pub async fn search_rollout_matches( + rg_command: &Path, + codex_home: &Path, + archived: bool, + search_term: &str, +) -> io::Result { let root = codex_home.join(if archived { ARCHIVED_SESSIONS_SUBDIR } else { SESSIONS_SUBDIR }); - let search_term = json_escaped_search_term(search_term)?; - let mut matches = - ripgrep_rollout_paths(rg_command, root.as_path(), search_term.as_str()).await?; - matches.extend(scan_compressed_rollout_paths(root.as_path(), search_term.as_str()).await?); + let json_search_term = json_escaped_search_term(search_term)?; + let Some(plain_matches) = + ripgrep_rollout_paths(rg_command, root.as_path(), json_search_term.as_str()).await? + else { + return scan_rollout_matches(root.as_path(), json_search_term.as_str(), search_term).await; + }; + let mut matches: RolloutSearchMatches = + plain_matches.into_iter().map(|path| (path, None)).collect(); + matches.extend(scan_compressed_rollout_matches(root.as_path(), search_term).await?); Ok(matches) } @@ -42,9 +64,9 @@ async fn ripgrep_rollout_paths( rg_command: &Path, root: &Path, search_term: &str, -) -> io::Result> { +) -> io::Result>> { if !tokio::fs::try_exists(root).await.unwrap_or(false) { - return Ok(HashSet::new()); + return Ok(Some(HashSet::new())); } let output = match Command::new(rg_command) @@ -62,13 +84,13 @@ async fn ripgrep_rollout_paths( { Ok(output) => output, Err(err) if err.kind() == io::ErrorKind::NotFound => { - return scan_rollout_paths(root, search_term).await; + return Ok(None); } Err(err) => return Err(err), }; if !output.status.success() { if output.status.code() == Some(1) && output.stderr.is_empty() { - return Ok(HashSet::new()); + return Ok(Some(HashSet::new())); } return Err(io::Error::other(format!( @@ -88,13 +110,17 @@ async fn ripgrep_rollout_paths( matches.insert(path); } - Ok(matches) + Ok(Some(matches)) } -async fn scan_rollout_paths(root: &Path, search_term: &str) -> io::Result> { - let mut matches = HashSet::new(); +async fn scan_rollout_matches( + root: &Path, + json_search_term: &str, + search_term: &str, +) -> io::Result { + let mut matches = HashMap::new(); let mut dirs = vec![root.to_path_buf()]; - let search_term = case_insensitive_literal_regex(search_term)?; + let json_search_term = case_insensitive_literal_regex(json_search_term)?; while let Some(dir) = dirs.pop() { let mut entries = match tokio::fs::read_dir(dir).await { @@ -115,8 +141,16 @@ async fn scan_rollout_paths(root: &Path, search_term: &str) -> io::Result io::Result> { - let mut matches = HashSet::new(); +) -> io::Result { + let mut matches = HashMap::new(); let mut dirs = vec![root.to_path_buf()]; - let search_term = case_insensitive_literal_regex(search_term)?; while let Some(dir) = dirs.pop() { let mut entries = match tokio::fs::read_dir(dir).await { @@ -181,8 +214,10 @@ async fn scan_compressed_rollout_paths( if !rollout_file.is_compressed() { continue; } - if rollout_contains(rollout_file.path(), &search_term).await? { - matches.insert(rollout_file.into_path()); + if let Some(snippet) = + first_rollout_content_match_snippet(rollout_file.path(), search_term).await? + { + matches.insert(rollout_file.into_path(), Some(snippet)); } } } diff --git a/codex-rs/thread-store/src/local/search_threads.rs b/codex-rs/thread-store/src/local/search_threads.rs index ed9ad3a0a233..20f5da824335 100644 --- a/codex-rs/thread-store/src/local/search_threads.rs +++ b/codex-rs/thread-store/src/local/search_threads.rs @@ -7,7 +7,7 @@ use codex_rollout::RolloutConfig; use codex_rollout::find_thread_names_by_ids; use codex_rollout::first_rollout_content_match_snippet; use codex_rollout::parse_cursor; -use codex_rollout::search_rollout_paths; +use codex_rollout::search_rollout_matches; use super::LocalThreadStore; use super::helpers::distinct_thread_metadata_title; @@ -64,7 +64,7 @@ pub(super) async fn search_threads( generate_memories: false, }; let rg_command = InstallContext::current().rg_command(); - let matching_paths = search_rollout_paths( + let matching_rollouts = search_rollout_matches( rg_command.as_path(), store.config.codex_home.as_path(), params.archived, @@ -74,7 +74,7 @@ pub(super) async fn search_threads( .map_err(|err| ThreadStoreError::Internal { message: format!("failed to search rollout contents: {err}"), })?; - if matching_paths.is_empty() { + if matching_rollouts.is_empty() { return Ok(ThreadSearchPage { items: Vec::new(), next_cursor: None, @@ -95,7 +95,7 @@ pub(super) async fn search_threads( search_term: None, use_state_db_only: state_db.is_some(), }; - let mut remaining_paths = matching_paths; + let mut remaining_rollouts = matching_rollouts; loop { let page = list_rollout_threads( @@ -109,16 +109,15 @@ pub(super) async fn search_threads( ) .await?; for item in page.items { - if !remaining_paths.remove(item.path.as_path()) { - continue; - } - let Some(snippet) = - first_rollout_content_match_snippet(item.path.as_path(), search_term) + let Some(snippet) = (match remaining_rollouts.remove(item.path.as_path()) { + Some(Some(snippet)) => Some(snippet), + Some(None) => first_rollout_content_match_snippet(item.path.as_path(), search_term) .await .map_err(|err| ThreadStoreError::Internal { message: format!("failed to read rollout search match: {err}"), - })? - else { + })?, + None => None, + }) else { continue; }; matching_items.push(ThreadSearchItem { item, snippet }); @@ -128,7 +127,7 @@ pub(super) async fn search_threads( } page_cursor = page.next_cursor; if matching_items.len() > params.page_size - || remaining_paths.is_empty() + || remaining_rollouts.is_empty() || page_cursor.is_none() { break;