mirror of
https://github.com/openai/codex.git
synced 2026-11-13 22:47:20 +00: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:
1 parent
67b805fc11
commit
45912a6dc6
4 files changed
+93
-32
No files matched your search
@@ -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()?;
|
||||||
|
|||||||
@@ -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;
|
||||||
|
|||||||
@@ -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;
|
||||||
|
|||||||
Reference in new issue
Block a user