feat: reuse compressed rollout search snippets

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

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;

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,31 +21,52 @@ use super::compression;
 const MATCH_CONTEXT_BEFORE_CHARS: usize = 48;
 const MATCH_CONTEXT_AFTER_CHARS: usize = 96;
 
+pub type RolloutSearchMatches = HashMap<PathBuf, Option<String>>;
+
 pub async fn search_rollout_paths(
     rg_command: &Path,
     codex_home: &Path,
     archived: bool,
     search_term: &str,
 ) -> io::Result<HashSet<PathBuf>> {
+    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<RolloutSearchMatches> {
     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)
 }
 
 async fn ripgrep_rollout_paths(
     rg_command: &Path,
     root: &Path,
     search_term: &str,
-) -> io::Result<HashSet<PathBuf>> {
+) -> io::Result<Option<HashSet<PathBuf>>> {
     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<HashSet<PathBuf>> {
-    let mut matches = HashSet::new();
+async fn scan_rollout_matches(
+    root: &Path,
+    json_search_term: &str,
+    search_term: &str,
+) -> io::Result<RolloutSearchMatches> {
+    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<HashSe
             let Some(rollout_file) = compression::RolloutFile::from_path(path) else {
                 continue;
             };
-            if rollout_contains(rollout_file.path(), &search_term).await? {
-                matches.insert(rollout_file.into_path());
+            if rollout_file.is_compressed() {
+                if let Some(snippet) =
+                    first_rollout_content_match_snippet(rollout_file.path(), search_term).await?
+                {
+                    matches.insert(rollout_file.into_path(), Some(snippet));
+                }
+                continue;
+            }
+            if rollout_contains(rollout_file.path(), &json_search_term).await? {
+                matches.insert(rollout_file.into_path(), None);
             }
         }
     }
@@ -151,13 +185,12 @@ pub async fn first_rollout_content_match_snippet(
     Ok(None)
 }
 
-async fn scan_compressed_rollout_paths(
+async fn scan_compressed_rollout_matches(
     root: &Path,
     search_term: &str,
-) -> io::Result<HashSet<PathBuf>> {
-    let mut matches = HashSet::new();
+) -> io::Result<RolloutSearchMatches> {
+    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));
             }
         }
     }

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;