From 6ad0e943cc727dc836d7c671f3377db30107f4d9 Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 10 Jul 2026 09:34:16 +0000 Subject: [PATCH] Reduce startup latency for ancestor discovery (#32106) ## Why Serial metadata probes add remote round trips while Codex discovers project instructions, repository skills, and project-root markers during startup. ## What changed - Probe ancestor directories concurrently while preserving discovery order. - Raise the shared find-up probe window to cover typical project hierarchies. - Refresh project instructions in parallel with plugin and skill warmup. ## Testing Add coverage that blocks one metadata request and verifies all expected project-root and `AGENTS.md` ancestor probes have started before it is released. GitOrigin-RevId: 11449837c7a212ec8ed385e7290476e926d93531 --- codex-rs/core-skills/src/loader.rs | 48 +++++++++++---- codex-rs/core/src/agents_md.rs | 38 ++++++++---- codex-rs/core/src/agents_md_tests.rs | 92 ++++++++++++++++++++++++++-- codex-rs/core/src/session/session.rs | 12 ++-- codex-rs/file-system/src/find_up.rs | 4 +- 5 files changed, 157 insertions(+), 37 deletions(-) diff --git a/codex-rs/core-skills/src/loader.rs b/codex-rs/core-skills/src/loader.rs index aa663e3e6007..9f496d60d7fa 100644 --- a/codex-rs/core-skills/src/loader.rs +++ b/codex-rs/core-skills/src/loader.rs @@ -146,6 +146,9 @@ const MAX_DEPENDENCY_URL_LEN: usize = MAX_DESCRIPTION_LEN; // Traversal depth from the skills root. const MAX_SCAN_DEPTH: usize = 6; const MAX_SKILLS_DIRS_PER_ROOT: usize = 2000; +// Keep ancestor metadata probes within one remote round trip for typical project hierarchies while +// leaving room for other startup discovery on the shared exec-server transport. +const MAX_CONCURRENT_ANCESTOR_PROBES: usize = 256; struct ResolvedDiscoveredSkill { skill: DiscoveredSkill, @@ -387,10 +390,19 @@ async fn repo_agents_skill_roots( let project_root = find_project_root(fs.as_ref(), cwd, &project_root_markers).await; let dirs = dirs_between_project_root_and_cwd(cwd, &project_root); let mut roots = Vec::new(); - for dir in dirs { - let agents_skills = dir.join(AGENTS_DIR_NAME).join(SKILLS_DIR_NAME); - let agents_skills_uri = PathUri::from_abs_path(&agents_skills); - match fs.get_metadata(&agents_skills_uri, /*sandbox*/ None).await { + let mut results = futures::stream::iter(dirs) + .map(|dir| { + let fs = Arc::clone(&fs); + async move { + let agents_skills = dir.join(AGENTS_DIR_NAME).join(SKILLS_DIR_NAME); + let agents_skills_uri = PathUri::from_abs_path(&agents_skills); + let result = fs.get_metadata(&agents_skills_uri, /*sandbox*/ None).await; + (agents_skills, result) + } + }) + .buffered(MAX_CONCURRENT_ANCESTOR_PROBES); + while let Some((agents_skills, result)) = results.next().await { + match result { Ok(metadata) if metadata.is_directory => roots.push(SkillRoot { path: agents_skills, scope: SkillScope::Repo, @@ -443,19 +455,29 @@ async fn find_project_root( return cwd.clone(); } + let mut probes = Vec::new(); for ancestor in cwd.ancestors() { for marker in project_root_markers { let marker_path = ancestor.join(marker); + probes.push((ancestor.clone(), marker_path)); + } + } + let mut results = futures::stream::iter(probes) + .map(|(ancestor, marker_path)| async move { let marker_path_uri = PathUri::from_abs_path(&marker_path); - match fs.get_metadata(&marker_path_uri, /*sandbox*/ None).await { - Ok(_) => return ancestor, - Err(err) if err.kind() == io::ErrorKind::NotFound => {} - Err(err) => { - tracing::warn!( - "failed to stat project root marker {}: {err:#}", - marker_path.display() - ); - } + let result = fs.get_metadata(&marker_path_uri, /*sandbox*/ None).await; + (ancestor, marker_path, result) + }) + .buffered(MAX_CONCURRENT_ANCESTOR_PROBES); + while let Some((ancestor, marker_path, result)) = results.next().await { + match result { + Ok(_) => return ancestor, + Err(err) if err.kind() == io::ErrorKind::NotFound => {} + Err(err) => { + tracing::warn!( + "failed to stat project root marker {}: {err:#}", + marker_path.display() + ); } } } diff --git a/codex-rs/core/src/agents_md.rs b/codex-rs/core/src/agents_md.rs index 1c01c8873ff2..e0b80b2f61a4 100644 --- a/codex-rs/core/src/agents_md.rs +++ b/codex-rs/core/src/agents_md.rs @@ -29,6 +29,7 @@ use codex_file_system::FindUpErrorPolicy; use codex_file_system::find_nearest_ancestor_with_markers; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; +use futures::StreamExt; use std::io; use toml::Value as TomlValue; use tracing::error; @@ -42,6 +43,11 @@ pub const LOCAL_AGENTS_MD_FILENAME: &str = "AGENTS.override.md"; /// concatenated with the following separator. const AGENTS_MD_SEPARATOR: &str = "\n\n--- project-doc ---\n\n"; +// Metadata probes are cheap and the exec-server transport already bounds total in-flight calls. +// This covers typical project hierarchies in one remote round trip without monopolizing that +// transport when independent startup discovery runs concurrently. +const MAX_CONCURRENT_ANCESTOR_PROBES: usize = 256; + /// Loads project AGENTS.md content and combines it with host-provided user /// instructions. pub(crate) async fn load_project_instructions( @@ -198,22 +204,28 @@ async fn agents_md_paths( vec![dir] }; - let mut found = Vec::new(); let candidate_filenames = candidate_filenames(config); - for directory in search_dirs { - for name in &candidate_filenames { - let candidate = directory - .join(name) - .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?; - match fs.get_metadata(&candidate, /*sandbox*/ None).await { - Ok(metadata) if metadata.is_file => { - found.push(candidate); - break; + let candidate_filenames = &candidate_filenames; + let mut results = futures::stream::iter(search_dirs) + .map(|directory| async move { + for name in candidate_filenames { + let candidate = directory + .join(name) + .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?; + match fs.get_metadata(&candidate, /*sandbox*/ None).await { + Ok(metadata) if metadata.is_file => return Ok(Some(candidate)), + Ok(_) => {} + Err(err) if err.kind() == io::ErrorKind::NotFound => {} + Err(err) => return Err(err), } - Ok(_) => {} - Err(err) if err.kind() == io::ErrorKind::NotFound => {} - Err(err) => return Err(err), } + Ok(None) + }) + .buffered(MAX_CONCURRENT_ANCESTOR_PROBES); + let mut found = Vec::new(); + while let Some(result) = results.next().await { + if let Some(candidate) = result? { + found.push(candidate); } } Ok(found) diff --git a/codex-rs/core/src/agents_md_tests.rs b/codex-rs/core/src/agents_md_tests.rs index d93e4634dbb6..89bc5414f4e4 100644 --- a/codex-rs/core/src/agents_md_tests.rs +++ b/codex-rs/core/src/agents_md_tests.rs @@ -719,9 +719,8 @@ async fn marker_search_does_not_wait_for_a_higher_ancestor() { } #[tokio::test] -async fn project_root_marker_search_pipelines_bounded_window_and_continues() { +async fn project_root_marker_search_starts_all_ancestor_probes() { const NESTING_DEPTH: usize = 9; - const CONCURRENCY_LIMIT: usize = 8; let tmp = tempfile::tempdir().expect("tempdir"); fs::write(tmp.path().join(".git"), "").unwrap(); @@ -734,6 +733,7 @@ async fn project_root_marker_search_pipelines_bounded_window_and_continues() { let mut config = make_config(&tmp, /*limit*/ 4096, /*instructions*/ None).await; config.cwd = nested.abs(); + let expected_probe_count = config.cwd.ancestors().count(); let cwd = PathUri::from_abs_path(&config.cwd); let metadata_calls = Arc::new(MetadataCallCounts::default()); let fs = FailingFileSystem { @@ -752,7 +752,7 @@ async fn project_root_marker_search_pipelines_bounded_window_and_continues() { .lock() .expect("metadata paths lock") .len() - >= CONCURRENCY_LIMIT + >= expected_probe_count { break; } @@ -767,7 +767,7 @@ async fn project_root_marker_search_pipelines_bounded_window_and_continues() { .lock() .expect("metadata paths lock") .len(), - CONCURRENCY_LIMIT + expected_probe_count ); metadata_calls.release.notify_one(); @@ -785,6 +785,90 @@ async fn project_root_marker_search_pipelines_bounded_window_and_continues() { ); } +#[tokio::test] +async fn agents_md_search_starts_all_directory_probes() { + const NESTING_DEPTH: usize = 9; + + let tmp = tempfile::tempdir().expect("tempdir"); + fs::write(tmp.path().join(".git"), "").unwrap(); + fs::write(tmp.path().join("AGENTS.md"), "project doc").unwrap(); + let mut nested = tmp.path().to_path_buf(); + for depth in 0..NESTING_DEPTH { + nested.push(format!("nested-{depth}")); + } + fs::create_dir_all(&nested).unwrap(); + + let mut config = make_config(&tmp, /*limit*/ 4096, /*instructions*/ None).await; + config.cwd = nested.abs(); + let cwd = PathUri::from_abs_path(&config.cwd); + let mut search_dirs = config + .cwd + .ancestors() + .take(NESTING_DEPTH + 1) + .collect::>(); + search_dirs.reverse(); + let expected_probes = search_dirs + .into_iter() + .map(|directory| PathUri::from_abs_path(&directory.join(LOCAL_AGENTS_MD_FILENAME))) + .collect::>(); + let metadata_calls = Arc::new(MetadataCallCounts::default()); + let fs = FailingFileSystem { + path: tmp.path().join(LOCAL_AGENTS_MD_FILENAME).abs(), + failure: InjectedFailure::MetadataBlocked, + metadata_calls: Arc::clone(&metadata_calls), + }; + + let search = + tokio::spawn(async move { super::agents_md_paths(&config.config, &cwd, &fs).await }); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + let started = metadata_calls.started.notified(); + if expected_probes.iter().all(|candidate| { + metadata_calls + .paths + .lock() + .expect("metadata paths lock") + .contains(candidate) + }) { + break; + } + started.await; + } + }) + .await + .expect("all directory probes should start"); + + let mut actual_probes = metadata_calls + .paths + .lock() + .expect("metadata paths lock") + .iter() + .filter(|path| expected_probes.contains(path)) + .map(ToString::to_string) + .collect::>(); + actual_probes.sort(); + let mut expected_probes = expected_probes + .into_iter() + .map(|path| path.to_string()) + .collect::>(); + expected_probes.sort(); + assert_eq!(actual_probes, expected_probes); + + metadata_calls.release.notify_one(); + let paths = tokio::time::timeout(std::time::Duration::from_secs(5), search) + .await + .expect("AGENTS.md search should complete") + .expect("AGENTS.md search task") + .expect("AGENTS.md discovery"); + + assert_eq!( + paths, + vec![PathUri::from_abs_path( + &tmp.path().join(DEFAULT_AGENTS_MD_FILENAME).abs() + )] + ); +} + #[tokio::test] async fn empty_project_root_markers_only_probe_cwd_candidates() { let tmp = tempfile::tempdir().expect("tempdir"); diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 7c59ec22b7b2..af3e4c512b8f 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -907,10 +907,7 @@ impl Session { turn_environments.update_selections(session_configuration.environment_selections()); let resolved_environments = turn_environments.snapshot().await; let agents_md_manager = Arc::new(AgentsMdManager::new(user_instructions)); - agents_md_manager - .refresh(config.as_ref(), &resolved_environments) - .await; - let plugin_skill_errors = warm_plugins_and_skills_for_session_init( + let plugin_skill_warmup = warm_plugins_and_skills_for_session_init( Arc::clone(&config), Arc::clone(&plugins_manager), Arc::clone(&skills_service), @@ -919,8 +916,11 @@ impl Session { .instrument(info_span!( "session_init.plugin_skill_warmup", otel.name = "session_init.plugin_skill_warmup", - )) - .await; + )); + let ((), plugin_skill_errors) = tokio::join!( + agents_md_manager.refresh(config.as_ref(), &resolved_environments), + plugin_skill_warmup, + ); for err in &plugin_skill_errors { error!( "failed to load skill {}: {}", diff --git a/codex-rs/file-system/src/find_up.rs b/codex-rs/file-system/src/find_up.rs index c0278569f7a1..b28ca504a6ce 100644 --- a/codex-rs/file-system/src/find_up.rs +++ b/codex-rs/file-system/src/find_up.rs @@ -6,7 +6,9 @@ use codex_utils_path_uri::PathUri; use futures::StreamExt; use std::io; -const MAX_CONCURRENT_PROBES: usize = 8; +// Keep enough ordinary metadata calls in flight to cover typical ancestor chains in one remote +// round trip, while leaving room for independent startup discovery to run at the same time. +const MAX_CONCURRENT_PROBES: usize = 256; /// Controls how an upward marker search handles metadata errors other than `NotFound`. #[derive(Clone, Copy, Debug, Eq, PartialEq)]