From 944c9cf7947a7e6f17b1e3cb850688ecab9e9bc1 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 10 Jul 2026 00:16:17 +0100 Subject: [PATCH] Cache watched rg file inventories --- codex-rs/Cargo.lock | 5 + codex-rs/Cargo.toml | 1 + codex-rs/rg-shim/Cargo.toml | 22 +++ codex-rs/rg-shim/src/lib.rs | 74 +++++++ codex-rs/rg-shim/src/lib_tests.rs | 52 +++++ codex-rs/rg-shim/src/main.rs | 47 +++-- codex-rs/rg-shim/src/main_tests.rs | 12 ++ codex-rs/rg-shim/src/manager.rs | 275 ++++++++++++++++++++++++++ codex-rs/rg-shim/src/manager_tests.rs | 46 +++++ 9 files changed, 523 insertions(+), 11 deletions(-) create mode 100644 codex-rs/rg-shim/src/lib.rs create mode 100644 codex-rs/rg-shim/src/lib_tests.rs create mode 100644 codex-rs/rg-shim/src/manager.rs create mode 100644 codex-rs/rg-shim/src/manager_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 17513928ab..192e365b54 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3816,8 +3816,13 @@ dependencies = [ name = "codex-rg" version = "0.0.0" dependencies = [ + "notify", "pretty_assertions", + "sha2 0.10.9", "tempfile", + "tokio", + "tracing", + "uuid", ] [[package]] diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 6da5891dc5..546a3e78b7 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -223,6 +223,7 @@ codex-prompts = { path = "prompts" } codex-responses-api-proxy = { path = "responses-api-proxy" } codex-response-debug-context = { path = "response-debug-context" } codex-rmcp-client = { path = "rmcp-client" } +codex-rg = { path = "rg-shim" } codex-rollout = { path = "rollout" } codex-rollout-trace = { path = "rollout-trace" } codex-sandboxing = { path = "sandboxing" } diff --git a/codex-rs/rg-shim/Cargo.toml b/codex-rs/rg-shim/Cargo.toml index d8f35cdb41..2c106c853a 100644 --- a/codex-rs/rg-shim/Cargo.toml +++ b/codex-rs/rg-shim/Cargo.toml @@ -8,9 +8,31 @@ license.workspace = true name = "codex-rg" path = "src/main.rs" +[lib] +name = "codex_rg" +path = "src/lib.rs" +doctest = false + +[features] +manager = [ + "dep:notify", + "dep:tempfile", + "dep:tokio", + "dep:tracing", + "dep:uuid", +] + [lints] workspace = true +[dependencies] +notify = { workspace = true, optional = true } +sha2 = { workspace = true } +tempfile = { workspace = true, optional = true } +tokio = { workspace = true, features = ["macros", "rt", "sync"], optional = true } +tracing = { workspace = true, optional = true } +uuid = { workspace = true, features = ["v4"], optional = true } + [dev-dependencies] pretty_assertions = { workspace = true } tempfile = { workspace = true } diff --git a/codex-rs/rg-shim/src/lib.rs b/codex-rs/rg-shim/src/lib.rs new file mode 100644 index 0000000000..a77d574d99 --- /dev/null +++ b/codex-rs/rg-shim/src/lib.rs @@ -0,0 +1,74 @@ +use sha2::Digest; +use sha2::Sha256; +use std::fs; +use std::fs::File; +use std::path::Path; +use std::path::PathBuf; + +/// Private environment key used to point the packaged shim at an +/// exec-server-owned cache generation. +pub const CACHE_ROOT_ENV: &str = "CODEX_INTERNAL_RG_CACHE_ROOT"; + +const FILES_FILENAME: &str = "files"; +const READY_FILENAME: &str = "ready"; +const ROOT_FILENAME: &str = "root"; + +#[derive(Clone, Debug, Eq, PartialEq)] +struct InventoryCache { + directory: PathBuf, + files: PathBuf, + ready: PathBuf, + root: PathBuf, + root_text: String, +} + +impl InventoryCache { + fn new(cache_root: &Path, repository_root: &Path) -> Option { + let root_text = repository_root.to_str()?.to_string(); + let key = format!("{:x}", Sha256::digest(root_text.as_bytes())); + let directory = cache_root.join(key); + Some(Self { + files: directory.join(FILES_FILENAME), + ready: directory.join(READY_FILENAME), + root: directory.join(ROOT_FILENAME), + directory, + root_text, + }) + } + + fn open(&self) -> Option { + let generation_before = fs::read(&self.ready).ok()?; + if generation_before.is_empty() || fs::read_to_string(&self.root).ok()? != self.root_text { + return None; + } + let files = File::open(&self.files).ok()?; + let generation_after = fs::read(&self.ready).ok()?; + (generation_before == generation_after).then_some(files) + } +} + +/// Opens an exact cached `rg --files` result for `cwd` when the executor has +/// published a live generation for the containing Git worktree. +pub fn open_file_inventory(cache_root: &Path, cwd: &Path) -> Option { + let repository_root = find_repository_root(cwd)?; + if fs::canonicalize(cwd).ok()? != repository_root { + return None; + } + InventoryCache::new(cache_root, &repository_root)?.open() +} + +pub(crate) fn find_repository_root(cwd: &Path) -> Option { + let cwd = fs::canonicalize(cwd).ok()?; + cwd.ancestors() + .find(|ancestor| ancestor.join(".git").exists()) + .map(Path::to_path_buf) +} + +#[cfg(feature = "manager")] +mod manager; +#[cfg(feature = "manager")] +pub use manager::RgCacheManager; + +#[cfg(test)] +#[path = "lib_tests.rs"] +mod tests; diff --git a/codex-rs/rg-shim/src/lib_tests.rs b/codex-rs/rg-shim/src/lib_tests.rs new file mode 100644 index 0000000000..8f28510142 --- /dev/null +++ b/codex-rs/rg-shim/src/lib_tests.rs @@ -0,0 +1,52 @@ +use super::InventoryCache; +use super::find_repository_root; +use super::open_file_inventory; +use pretty_assertions::assert_eq; +use std::fs; +use std::io::Read; +use tempfile::TempDir; + +#[test] +fn finds_nearest_git_worktree_root() { + let temp_dir = TempDir::new().expect("temp dir"); + let outer = temp_dir.path().join("outer"); + let inner = outer.join("nested"); + let child = inner.join("src"); + fs::create_dir_all(outer.join(".git")).expect("outer git dir"); + fs::create_dir_all(inner.join(".git")).expect("inner git dir"); + fs::create_dir_all(&child).expect("child dir"); + + assert_eq!( + find_repository_root(&child), + Some(fs::canonicalize(inner).expect("canonical inner root")) + ); +} + +#[test] +fn opens_only_matching_ready_inventory_at_repository_root() { + let temp_dir = TempDir::new().expect("temp dir"); + let repository = temp_dir.path().join("repository"); + let child = repository.join("src"); + let cache_root = temp_dir.path().join("cache"); + fs::create_dir_all(repository.join(".git")).expect("git dir"); + fs::create_dir_all(&child).expect("child dir"); + let repository = fs::canonicalize(repository).expect("canonical repository"); + let cache = InventoryCache::new(&cache_root, &repository).expect("cache layout"); + fs::create_dir_all(&cache.directory).expect("cache dir"); + fs::write(&cache.root, repository.to_string_lossy().as_bytes()).expect("root marker"); + fs::write(&cache.files, "src/lib.rs\n").expect("inventory"); + + assert!(open_file_inventory(&cache_root, &repository).is_none()); + fs::write(&cache.ready, "generation-1").expect("ready marker"); + assert!(open_file_inventory(&cache_root, &child).is_none()); + + let mut inventory = String::new(); + open_file_inventory(&cache_root, &repository) + .expect("ready inventory") + .read_to_string(&mut inventory) + .expect("read inventory"); + assert_eq!(inventory, "src/lib.rs\n"); + + fs::write(&cache.root, "/different/repository").expect("mismatched root marker"); + assert!(open_file_inventory(&cache_root, &repository).is_none()); +} diff --git a/codex-rs/rg-shim/src/main.rs b/codex-rs/rg-shim/src/main.rs index bd3c27af9c..7a5ddb8c69 100644 --- a/codex-rs/rg-shim/src/main.rs +++ b/codex-rs/rg-shim/src/main.rs @@ -1,5 +1,9 @@ +use codex_rg::CACHE_ROOT_ENV; +use codex_rg::open_file_inventory; use std::ffi::OsStr; +use std::ffi::OsString; use std::io; +use std::io::Write; use std::path::Path; use std::path::PathBuf; use std::process::Command; @@ -9,9 +13,7 @@ const PATH_DIRNAME: &str = "codex-path"; const RESOURCES_DIRNAME: &str = "codex-resources"; fn main() { - let result = std::env::current_exe() - .and_then(|current_exe| real_rg_path(¤t_exe)) - .and_then(run_real_rg); + let result = run(); match result { Ok(exit_code) => std::process::exit(exit_code), Err(error) => { @@ -21,6 +23,33 @@ fn main() { } } +fn run() -> io::Result { + let current_exe = std::env::current_exe()?; + let real_rg = real_rg_path(¤t_exe)?; + let args = std::env::args_os().skip(1).collect::>(); + if is_default_files_request(&args) + && std::env::var_os("RIPGREP_CONFIG_PATH").is_none() + && let Some(cache_root) = std::env::var_os(CACHE_ROOT_ENV) + && let Some(mut inventory) = + open_file_inventory(Path::new(&cache_root), &std::env::current_dir()?) + { + let mut stdout = io::stdout().lock(); + return match io::copy(&mut inventory, &mut stdout) { + Ok(_) => { + stdout.flush()?; + Ok(0) + } + Err(error) if error.kind() == io::ErrorKind::BrokenPipe => Ok(0), + Err(error) => Err(error), + }; + } + run_real_rg(real_rg, args) +} + +fn is_default_files_request(args: &[OsString]) -> bool { + args == [OsStr::new("--files")] +} + fn real_rg_path(current_exe: &Path) -> io::Result { let path_dir = current_exe.parent().ok_or_else(|| { io::Error::new( @@ -60,19 +89,15 @@ fn real_rg_path(current_exe: &Path) -> io::Result { } #[cfg(unix)] -fn run_real_rg(real_rg: PathBuf) -> io::Result { +fn run_real_rg(real_rg: PathBuf, args: Vec) -> io::Result { use std::os::unix::process::CommandExt; - Err(Command::new(real_rg) - .args(std::env::args_os().skip(1)) - .exec()) + Err(Command::new(real_rg).args(args).exec()) } #[cfg(windows)] -fn run_real_rg(real_rg: PathBuf) -> io::Result { - let status = Command::new(real_rg) - .args(std::env::args_os().skip(1)) - .status()?; +fn run_real_rg(real_rg: PathBuf, args: Vec) -> io::Result { + let status = Command::new(real_rg).args(args).status()?; Ok(status.code().unwrap_or(1)) } diff --git a/codex-rs/rg-shim/src/main_tests.rs b/codex-rs/rg-shim/src/main_tests.rs index 9f9e30e5d1..b589b1127a 100644 --- a/codex-rs/rg-shim/src/main_tests.rs +++ b/codex-rs/rg-shim/src/main_tests.rs @@ -1,8 +1,10 @@ use super::PACKAGE_METADATA_FILENAME; use super::PATH_DIRNAME; use super::RESOURCES_DIRNAME; +use super::is_default_files_request; use super::real_rg_path; use pretty_assertions::assert_eq; +use std::ffi::OsString; use std::fs; use tempfile::TempDir; @@ -35,3 +37,13 @@ fn rejects_executable_outside_package_path() { assert_eq!(error.kind(), std::io::ErrorKind::NotFound); } + +#[test] +fn only_exact_default_files_request_uses_inventory() { + assert!(is_default_files_request(&[OsString::from("--files")])); + assert!(!is_default_files_request(&[ + OsString::from("--files"), + OsString::from("-0") + ])); + assert!(!is_default_files_request(&[OsString::from("--hidden")])); +} diff --git a/codex-rs/rg-shim/src/manager.rs b/codex-rs/rg-shim/src/manager.rs new file mode 100644 index 0000000000..dda0b89a72 --- /dev/null +++ b/codex-rs/rg-shim/src/manager.rs @@ -0,0 +1,275 @@ +use crate::InventoryCache; +use crate::find_repository_root; +use notify::Event; +use notify::EventKind; +use notify::RecommendedWatcher; +use notify::RecursiveMode; +use notify::Watcher; +use std::collections::HashSet; +use std::fmt; +use std::fs; +use std::io; +use std::io::Write; +use std::path::Path; +use std::path::PathBuf; +use std::process::Command; +use std::process::Stdio; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Instant; +use tempfile::NamedTempFile; +use tempfile::TempDir; +use tokio::sync::mpsc; +use tracing::debug; +use tracing::warn; +use uuid::Uuid; + +/// Owns ephemeral, watched ripgrep cache generations for one exec-server. +#[derive(Clone)] +pub struct RgCacheManager { + inner: Arc, +} + +struct Inner { + cache_root: TempDir, + real_rg: PathBuf, + repositories: Mutex>, +} + +impl fmt::Debug for RgCacheManager { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("RgCacheManager") + .field("cache_root", &self.inner.cache_root.path()) + .field("real_rg", &self.inner.real_rg) + .finish_non_exhaustive() + } +} + +impl RgCacheManager { + /// Creates a process-scoped cache generation below `cache_parent`. + pub fn new(real_rg: PathBuf, cache_parent: &Path) -> io::Result { + fs::create_dir_all(cache_parent)?; + let cache_root = tempfile::Builder::new() + .prefix("generation-") + .tempdir_in(cache_parent)?; + Ok(Self { + inner: Arc::new(Inner { + cache_root, + real_rg, + repositories: Mutex::new(HashSet::new()), + }), + }) + } + + /// Returns the private cache root exposed read-only to eligible children. + pub fn cache_root(&self) -> &Path { + self.inner.cache_root.path() + } + + /// Returns the canonical Git worktree containing `cwd`. + pub fn repository_root(cwd: &Path) -> Option { + find_repository_root(cwd) + } + + /// Starts watching and building an inventory for `repository_root`. + pub fn observe_repository(&self, repository_root: PathBuf) { + let mut repositories = self + .inner + .repositories + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if !repositories.insert(repository_root.clone()) { + return; + } + drop(repositories); + + let manager = self.clone(); + let worker_root = repository_root.clone(); + tokio::spawn(async move { + if let Err(error) = manager.run_repository(worker_root.clone()).await { + warn!( + repository = %worker_root.display(), + %error, + "ripgrep inventory worker stopped" + ); + } + }); + } + + async fn run_repository(&self, repository_root: PathBuf) -> io::Result<()> { + let cache = InventoryCache::new(self.cache_root(), &repository_root) + .ok_or_else(|| io::Error::other("repository path is not valid UTF-8"))?; + fs::create_dir_all(&cache.directory)?; + fs::write( + &cache.root, + repository_root + .to_str() + .ok_or_else(|| io::Error::other("repository path is not valid UTF-8"))?, + )?; + + let (event_tx, mut event_rx) = mpsc::channel(1); + let generation = Arc::new(AtomicU64::new(0)); + let callback_generation = Arc::clone(&generation); + let watch_failed = Arc::new(AtomicBool::new(false)); + let callback_watch_failed = Arc::clone(&watch_failed); + let watch_root = repository_root.clone(); + let _watcher = tokio::task::spawn_blocking(move || { + let mut watcher = + notify::recommended_watcher(move |event: notify::Result| match event { + Ok(event) if event_is_mutating(&event) => { + callback_generation.fetch_add(1, Ordering::Release); + let _ = event_tx.try_send(()); + } + Ok(_) => {} + Err(error) => { + warn!(%error, "ripgrep inventory watcher failed"); + callback_watch_failed.store(true, Ordering::Release); + callback_generation.fetch_add(1, Ordering::Release); + let _ = event_tx.try_send(()); + } + }) + .map_err(io::Error::other)?; + watcher + .watch(&watch_root, RecursiveMode::Recursive) + .map_err(io::Error::other)?; + Ok::(watcher) + }) + .await + .map_err(io::Error::other)??; + + loop { + if !self + .build_clean_generation( + &repository_root, + &cache, + &generation, + &watch_failed, + &mut event_rx, + ) + .await? + { + return Ok(()); + } + match event_rx.recv().await { + Some(()) if !watch_failed.load(Ordering::Acquire) => invalidate(&cache), + Some(()) | None => { + invalidate(&cache); + return Ok(()); + } + } + } + } + + async fn build_clean_generation( + &self, + repository_root: &Path, + cache: &InventoryCache, + generation: &AtomicU64, + watch_failed: &AtomicBool, + event_rx: &mut mpsc::Receiver<()>, + ) -> io::Result { + loop { + while event_rx.try_recv().is_ok() {} + if event_rx.is_closed() || watch_failed.load(Ordering::Acquire) { + return Ok(false); + } + invalidate(cache); + let started_at = Instant::now(); + let generation_before = generation.load(Ordering::Acquire); + let real_rg = self.inner.real_rg.clone(); + let build_root = repository_root.to_path_buf(); + let cache_directory = cache.directory.clone(); + let mut build = tokio::task::spawn_blocking(move || { + build_inventory(&real_rg, &build_root, &cache_directory) + }); + let mut dirty = false; + let inventory = loop { + tokio::select! { + result = &mut build => { + break result.map_err(io::Error::other)??; + } + event = event_rx.recv() => { + match event { + Some(()) if !watch_failed.load(Ordering::Acquire) => dirty = true, + Some(()) | None => return Ok(false), + } + } + } + }; + if watch_failed.load(Ordering::Acquire) { + return Ok(false); + } + if dirty || generation.load(Ordering::Acquire) != generation_before { + continue; + } + + publish(cache, inventory)?; + debug!( + repository = %repository_root.display(), + elapsed_ms = started_at.elapsed().as_millis(), + "ripgrep file inventory is ready" + ); + return Ok(true); + } + } +} + +fn build_inventory( + real_rg: &Path, + repository_root: &Path, + cache_directory: &Path, +) -> io::Result { + let inventory = NamedTempFile::new_in(cache_directory)?; + let output = inventory.reopen()?; + let status = Command::new(real_rg) + .arg("--files") + .current_dir(repository_root) + .env_remove(crate::CACHE_ROOT_ENV) + .env_remove("RIPGREP_CONFIG_PATH") + .stdin(Stdio::null()) + .stdout(Stdio::from(output)) + .stderr(Stdio::null()) + .status()?; + if !status.success() { + return Err(io::Error::other(format!( + "ripgrep inventory exited with {status}" + ))); + } + Ok(inventory) +} + +fn publish(cache: &InventoryCache, inventory: NamedTempFile) -> io::Result<()> { + match fs::remove_file(&cache.files) { + Ok(()) => {} + Err(error) if error.kind() == io::ErrorKind::NotFound => {} + Err(error) => return Err(error), + } + inventory + .persist(&cache.files) + .map_err(|error| error.error)?; + let mut ready = NamedTempFile::new_in(&cache.directory)?; + writeln!(ready, "{}", Uuid::new_v4())?; + ready.flush()?; + ready.persist(&cache.ready).map_err(|error| error.error)?; + Ok(()) +} + +fn invalidate(cache: &InventoryCache) { + match fs::remove_file(&cache.ready) { + Ok(()) => {} + Err(error) if error.kind() == io::ErrorKind::NotFound => {} + Err(error) => warn!(%error, "failed to invalidate ripgrep inventory"), + } +} + +fn event_is_mutating(event: &Event) -> bool { + !matches!(event.kind, EventKind::Access(_)) +} + +#[cfg(all(test, unix))] +#[path = "manager_tests.rs"] +mod tests; diff --git a/codex-rs/rg-shim/src/manager_tests.rs b/codex-rs/rg-shim/src/manager_tests.rs new file mode 100644 index 0000000000..c908dd0329 --- /dev/null +++ b/codex-rs/rg-shim/src/manager_tests.rs @@ -0,0 +1,46 @@ +use super::build_inventory; +use super::invalidate; +use super::publish; +use crate::InventoryCache; +use pretty_assertions::assert_eq; +use std::fs; +use std::io::Read; +use std::os::unix::fs::PermissionsExt; +use tempfile::TempDir; + +#[test] +fn built_inventory_is_visible_only_while_ready() { + let temp_dir = TempDir::new().expect("temp dir"); + let repository = temp_dir.path().join("repository"); + let cache_root = temp_dir.path().join("cache"); + fs::create_dir_all(repository.join(".git")).expect("git dir"); + fs::create_dir_all(&cache_root).expect("cache root"); + let repository = fs::canonicalize(repository).expect("canonical repository"); + let cache = InventoryCache::new(&cache_root, &repository).expect("cache layout"); + fs::create_dir_all(&cache.directory).expect("cache dir"); + fs::write(&cache.root, repository.to_string_lossy().as_bytes()).expect("root marker"); + + let fake_rg = temp_dir.path().join("rg"); + fs::write(&fake_rg, "#!/bin/sh\nprintf 'src/lib.rs\\nREADME.md\\n'\n").expect("fake rg"); + let mut permissions = fs::metadata(&fake_rg) + .expect("fake rg metadata") + .permissions(); + permissions.set_mode(0o755); + fs::set_permissions(&fake_rg, permissions).expect("fake rg permissions"); + + let inventory = + build_inventory(&fake_rg, &repository, &cache.directory).expect("build inventory"); + assert!(cache.open().is_none()); + + publish(&cache, inventory).expect("publish inventory"); + let mut output = String::new(); + cache + .open() + .expect("ready inventory") + .read_to_string(&mut output) + .expect("read inventory"); + assert_eq!(output, "src/lib.rs\nREADME.md\n"); + + invalidate(&cache); + assert!(cache.open().is_none()); +}