feat: reuse compressed rollout search snippets (#25814)

## Summary
- teach rollout search to return precomputed snippets for compressed
rollouts
- reuse those snippets in local thread search instead of reopening
matching compressed files
- keep the no-`rg` fallback single-pass and add regression coverage for
the compressed path

## Why
`thread/search` currently decodes matching compressed rollouts twice:
once to discover the matching path and again to extract the snippet
shown in results. That defeats a meaningful part of the compressed-read
optimization work.

## Impact
Compressed rollout hits now pay one decode pass on the search path while
plain `.jsonl` hits keep the existing ripgrep-driven flow.

## Validation
- `just test -p codex-rollout`
- `just test -p codex-thread-store`
- `just fix -p codex-rollout`
- `just fix -p codex-thread-store`
- `just fmt`
This commit is contained in:
jif-oai authored and GitHub committed 2026-06-02 11:32:36 +02:00
1 parent 67b805fc11
commit 45912a6dc6
4 files changed
+93 -32

No files matched your search

+26
View File
@@ -23,6 +23,7 @@ use crate::RolloutConfig;
use crate::RolloutRecorder; use crate::RolloutRecorder;
use crate::RolloutRecorderParams; use crate::RolloutRecorderParams;
use crate::append_rollout_item_to_path; use crate::append_rollout_item_to_path;
use crate::search_rollout_matches;
#[tokio::test] #[tokio::test]
async fn load_rollout_items_reads_compressed_rollout() -> anyhow::Result<()> { 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(()) 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] #[tokio::test]
async fn worker_compresses_old_archived_rollouts_only() -> anyhow::Result<()> { async fn worker_compresses_old_archived_rollouts_only() -> anyhow::Result<()> {
let home = TempDir::new()?; let home = TempDir::new()?;
+1
View File
@@ -67,6 +67,7 @@ pub use recorder::RolloutRecorder;
pub use recorder::RolloutRecorderParams; pub use recorder::RolloutRecorderParams;
pub use recorder::append_rollout_item_to_path; pub use recorder::append_rollout_item_to_path;
pub use search::first_rollout_content_match_snippet; pub use search::first_rollout_content_match_snippet;
pub use search::search_rollout_matches;
pub use search::search_rollout_paths; pub use search::search_rollout_paths;
pub use session_index::append_thread_name; pub use session_index::append_thread_name;
pub use session_index::find_thread_meta_by_name_str; pub use session_index::find_thread_meta_by_name_str;
+55 -20
View File
@@ -1,3 +1,4 @@
use std::collections::HashMap;
use std::collections::HashSet; use std::collections::HashSet;
use std::io; use std::io;
use std::path::Path; use std::path::Path;
@@ -20,21 +21,42 @@ use super::compression;
const MATCH_CONTEXT_BEFORE_CHARS: usize = 48; const MATCH_CONTEXT_BEFORE_CHARS: usize = 48;
const MATCH_CONTEXT_AFTER_CHARS: usize = 96; const MATCH_CONTEXT_AFTER_CHARS: usize = 96;
pub type RolloutSearchMatches = HashMap<PathBuf, Option<String>>;
pub async fn search_rollout_paths( pub async fn search_rollout_paths(
rg_command: &Path, rg_command: &Path,
codex_home: &Path, codex_home: &Path,
archived: bool, archived: bool,
search_term: &str, search_term: &str,
) -> io::Result<HashSet<PathBuf>> { ) -> 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 { let root = codex_home.join(if archived {
ARCHIVED_SESSIONS_SUBDIR ARCHIVED_SESSIONS_SUBDIR
} else { } else {
SESSIONS_SUBDIR SESSIONS_SUBDIR
}); });
let search_term = json_escaped_search_term(search_term)?; let json_search_term = json_escaped_search_term(search_term)?;
let mut matches = let Some(plain_matches) =
ripgrep_rollout_paths(rg_command, root.as_path(), search_term.as_str()).await?; ripgrep_rollout_paths(rg_command, root.as_path(), json_search_term.as_str()).await?
matches.extend(scan_compressed_rollout_paths(root.as_path(), 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) Ok(matches)
} }
@@ -42,9 +64,9 @@ async fn ripgrep_rollout_paths(
rg_command: &Path, rg_command: &Path,
root: &Path, root: &Path,
search_term: &str, search_term: &str,
) -> io::Result<HashSet<PathBuf>> { ) -> io::Result<Option<HashSet<PathBuf>>> {
if !tokio::fs::try_exists(root).await.unwrap_or(false) { 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) let output = match Command::new(rg_command)
@@ -62,13 +84,13 @@ async fn ripgrep_rollout_paths(
{ {
Ok(output) => output, Ok(output) => output,
Err(err) if err.kind() == io::ErrorKind::NotFound => { Err(err) if err.kind() == io::ErrorKind::NotFound => {
return scan_rollout_paths(root, search_term).await; return Ok(None);
} }
Err(err) => return Err(err), Err(err) => return Err(err),
}; };
if !output.status.success() { if !output.status.success() {
if output.status.code() == Some(1) && output.stderr.is_empty() { 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!( return Err(io::Error::other(format!(
@@ -88,13 +110,17 @@ async fn ripgrep_rollout_paths(
matches.insert(path); matches.insert(path);
} }
Ok(matches) Ok(Some(matches))
} }
async fn scan_rollout_paths(root: &Path, search_term: &str) -> io::Result<HashSet<PathBuf>> { async fn scan_rollout_matches(
let mut matches = HashSet::new(); 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 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() { while let Some(dir) = dirs.pop() {
let mut entries = match tokio::fs::read_dir(dir).await { 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 { let Some(rollout_file) = compression::RolloutFile::from_path(path) else {
continue; continue;
}; };
if rollout_contains(rollout_file.path(), &search_term).await? { if rollout_file.is_compressed() {
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));
}
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) Ok(None)
} }
async fn scan_compressed_rollout_paths( async fn scan_compressed_rollout_matches(
root: &Path, root: &Path,
search_term: &str, search_term: &str,
) -> io::Result<HashSet<PathBuf>> { ) -> io::Result<RolloutSearchMatches> {
let mut matches = HashSet::new(); let mut matches = HashMap::new();
let mut dirs = vec![root.to_path_buf()]; let mut dirs = vec![root.to_path_buf()];
let search_term = case_insensitive_literal_regex(search_term)?;
while let Some(dir) = dirs.pop() { while let Some(dir) = dirs.pop() {
let mut entries = match tokio::fs::read_dir(dir).await { 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() { if !rollout_file.is_compressed() {
continue; continue;
} }
if rollout_contains(rollout_file.path(), &search_term).await? { if let Some(snippet) =
matches.insert(rollout_file.into_path()); first_rollout_content_match_snippet(rollout_file.path(), search_term).await?
{
matches.insert(rollout_file.into_path(), Some(snippet));
} }
} }
} }
@@ -7,7 +7,7 @@ use codex_rollout::RolloutConfig;
use codex_rollout::find_thread_names_by_ids; use codex_rollout::find_thread_names_by_ids;
use codex_rollout::first_rollout_content_match_snippet; use codex_rollout::first_rollout_content_match_snippet;
use codex_rollout::parse_cursor; use codex_rollout::parse_cursor;
use codex_rollout::search_rollout_paths; use codex_rollout::search_rollout_matches;
use super::LocalThreadStore; use super::LocalThreadStore;
use super::helpers::distinct_thread_metadata_title; use super::helpers::distinct_thread_metadata_title;
@@ -64,7 +64,7 @@ pub(super) async fn search_threads(
generate_memories: false, generate_memories: false,
}; };
let rg_command = InstallContext::current().rg_command(); let rg_command = InstallContext::current().rg_command();
let matching_paths = search_rollout_paths( let matching_rollouts = search_rollout_matches(
rg_command.as_path(), rg_command.as_path(),
store.config.codex_home.as_path(), store.config.codex_home.as_path(),
params.archived, params.archived,
@@ -74,7 +74,7 @@ pub(super) async fn search_threads(
.map_err(|err| ThreadStoreError::Internal { .map_err(|err| ThreadStoreError::Internal {
message: format!("failed to search rollout contents: {err}"), message: format!("failed to search rollout contents: {err}"),
})?; })?;
if matching_paths.is_empty() { if matching_rollouts.is_empty() {
return Ok(ThreadSearchPage { return Ok(ThreadSearchPage {
items: Vec::new(), items: Vec::new(),
next_cursor: None, next_cursor: None,
@@ -95,7 +95,7 @@ pub(super) async fn search_threads(
search_term: None, search_term: None,
use_state_db_only: state_db.is_some(), use_state_db_only: state_db.is_some(),
}; };
let mut remaining_paths = matching_paths; let mut remaining_rollouts = matching_rollouts;
loop { loop {
let page = list_rollout_threads( let page = list_rollout_threads(
@@ -109,16 +109,15 @@ pub(super) async fn search_threads(
) )
.await?; .await?;
for item in page.items { for item in page.items {
if !remaining_paths.remove(item.path.as_path()) { let Some(snippet) = (match remaining_rollouts.remove(item.path.as_path()) {
continue; Some(Some(snippet)) => Some(snippet),
} Some(None) => first_rollout_content_match_snippet(item.path.as_path(), search_term)
let Some(snippet) =
first_rollout_content_match_snippet(item.path.as_path(), search_term)
.await .await
.map_err(|err| ThreadStoreError::Internal { .map_err(|err| ThreadStoreError::Internal {
message: format!("failed to read rollout search match: {err}"), message: format!("failed to read rollout search match: {err}"),
})? })?,
else { None => None,
}) else {
continue; continue;
}; };
matching_items.push(ThreadSearchItem { item, snippet }); matching_items.push(ThreadSearchItem { item, snippet });
@@ -128,7 +127,7 @@ pub(super) async fn search_threads(
} }
page_cursor = page.next_cursor; page_cursor = page.next_cursor;
if matching_items.len() > params.page_size if matching_items.len() > params.page_size
|| remaining_paths.is_empty() || remaining_rollouts.is_empty()
|| page_cursor.is_none() || page_cursor.is_none()
{ {
break; break;