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
2026-06-02 11:32:36 +02:00
committed by GitHub
Unverified
parent 67b805fc11
commit 45912a6dc6
4 changed files with 93 additions and 32 deletions
+26
View File
@@ -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()?;
+1
View File
@@ -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;
+55 -20
View File
@@ -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<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)
}
@@ -42,9 +64,9 @@ 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));
}
}
}
@@ -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;