mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
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:
committed by
GitHub
Unverified
parent
67b805fc11
commit
45912a6dc6
@@ -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()?;
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user