From 703d253ed676c26e4b9708a374b7344b5c1bcc53 Mon Sep 17 00:00:00 2001 From: Anthony Tafoya Date: Tue, 13 Oct 2026 17:13:14 +0000 Subject: [PATCH] Wire V2 capability discovery through exec-server RPC (#53723) ## What changed - Register `capabilities/discoverV2` for initialized sessions and expose it through `ExecServerClient` and remote `Environment` APIs, checking executor support before calling it. - Reuse startup discovery locations and retry failed prewarming on demand. Replay discovery once after executor recovery while preserving the request parameters. - Track whether all required filesystem watches registered successfully and reload locations and metadata on each request when registration fails. ## Testing Add coverage for RPC metadata refresh, omission of skill instruction bodies, sandbox isolation, startup scan recovery, disconnect replay, and successful or failed watch registration. GitOrigin-RevId: 726e8f7e8b4d835b877f31d1d8a9f8eea1ad9c8b --- codex-rs/exec-server/src/client.rs | 10 + .../src/discoverV2/capability_discoveries.rs | 5 +- .../src/discoverV2/capability_manager.rs | 49 ++-- .../capability_manager/capability_watchers.rs | 8 +- codex-rs/exec-server/src/discoverV2/mod.rs | 3 - .../tests/capability_discoveries_tests.rs | 13 +- .../tests/capability_manager_tests.rs | 88 ++++--- .../tests/capability_watchers_tests.rs | 109 ++++++--- codex-rs/exec-server/src/environment.rs | 37 +-- .../src/environment/discovery_recovery.rs | 49 ++++ .../src/environment/discovery_v2.rs | 43 ++++ .../src/environment/discovery_v2_tests.rs | 100 ++++++++ codex-rs/exec-server/src/server/handler.rs | 17 ++ .../exec-server/src/server/handler/tests.rs | 7 + .../src/server/process_otel_tests.rs | 2 + codex-rs/exec-server/src/server/processor.rs | 131 ++--------- codex-rs/exec-server/src/server/registry.rs | 8 + .../src/server/request_dispatcher_tests.rs | 3 + codex-rs/exec-server/tests/capabilities_v2.rs | 214 ++++++++++++++++++ codex-rs/exec-server/tests/initialize.rs | 40 +++- .../file-watcher/src/file_watcher_tests.rs | 101 +++++++++ codex-rs/file-watcher/src/lib.rs | 39 +++- 22 files changed, 802 insertions(+), 274 deletions(-) create mode 100644 codex-rs/exec-server/src/environment/discovery_recovery.rs create mode 100644 codex-rs/exec-server/src/environment/discovery_v2.rs create mode 100644 codex-rs/exec-server/src/environment/discovery_v2_tests.rs create mode 100644 codex-rs/exec-server/tests/capabilities_v2.rs diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 368662ff1d..2ce504dabb 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -46,9 +46,12 @@ use crate::process::ExecProcessEvent; use crate::process::ExecProcessEventLog; use crate::process::ExecProcessEventReceiver; use crate::process_telemetry::trace_process_id; +use crate::protocol::CAPABILITIES_DISCOVER_V2_METHOD; use crate::protocol::CAPABILITY_ROOTS_DISCOVER_METHOD; use crate::protocol::CapabilityRootsDiscoverParams; use crate::protocol::CapabilityRootsDiscoverResponse; +use crate::protocol::DiscoverV2CapabilitiesRequest; +use crate::protocol::DiscoverV2CapabilitiesResponse; use crate::protocol::ENVIRONMENT_CONFIG_READ_METHOD; use crate::protocol::ENVIRONMENT_INFO_METHOD; use crate::protocol::ENVIRONMENT_STATUS_METHOD; @@ -893,6 +896,13 @@ impl ExecServerClient { self.call(CAPABILITY_ROOTS_DISCOVER_METHOD, ¶ms).await } + pub async fn discover_capabilities_v2( + &self, + params: DiscoverV2CapabilitiesRequest, + ) -> Result { + self.call(CAPABILITIES_DISCOVER_V2_METHOD, ¶ms).await + } + pub async fn read(&self, params: ReadParams) -> Result { self.call(EXEC_READ_METHOD, ¶ms).await } diff --git a/codex-rs/exec-server/src/discoverV2/capability_discoveries.rs b/codex-rs/exec-server/src/discoverV2/capability_discoveries.rs index bda6f20e67..b681a417c0 100644 --- a/codex-rs/exec-server/src/discoverV2/capability_discoveries.rs +++ b/codex-rs/exec-server/src/discoverV2/capability_discoveries.rs @@ -87,8 +87,9 @@ pub(crate) async fn load_capability_discovery_batch( }; let mut budget = ResponseBudget::default(); // Use V1 concurrency; admit results in location order. - let mut discoveries = futures::stream::iter(locations) - .map(|location| load_capability_discovery(file_system, location.clone())) + // Yield owned locations to satisfy the RPC router's Send future bound. + let mut discoveries = futures::stream::iter(locations.iter().cloned()) + .map(|location| load_capability_discovery(file_system, location)) .buffered(MAX_CONCURRENT_ROOTS) .enumerate(); while let Some((index, discovery)) = discoveries.next().await { diff --git a/codex-rs/exec-server/src/discoverV2/capability_manager.rs b/codex-rs/exec-server/src/discoverV2/capability_manager.rs index 2f3e37de57..037a1c6afb 100644 --- a/codex-rs/exec-server/src/discoverV2/capability_manager.rs +++ b/codex-rs/exec-server/src/discoverV2/capability_manager.rs @@ -1,5 +1,5 @@ //! Retains watched global locations and sandbox-scoped response snapshots. -//! Replacement snapshots isolate in-flight loads; unavailable watchers force request-time reloads. +//! Replacement snapshots isolate in-flight loads; failed watch registration forces request-time reloads. use crate::DiscoverV2CapabilitiesRequest; use crate::DiscoverV2CapabilitiesResponse; @@ -10,6 +10,7 @@ use crate::discover_v2::capability_locations::CapabilityLocationRequest; use crate::discover_v2::capability_locations::CapabilityLocations; use crate::discover_v2::capability_locations::prewarm_locations; use codex_file_watcher::WatchRegistration; +use codex_utils_path_uri::PathUri; use std::io; use std::sync::Arc; use std::sync::atomic::AtomicBool; @@ -35,7 +36,7 @@ struct LocationState { dirty: AtomicBool, locations: CapabilityLocations, sandbox_discoveries: Mutex>>, - _registration: WatchRegistration, + registration: WatchRegistration, } struct CachedSandboxDiscovery { @@ -54,24 +55,22 @@ impl CapabilityManager { }) } + /// Resolves home paths on first initialization and reuses the prewarmed locations. pub(crate) async fn prewarm_locations( self: &Arc, - request: CapabilityLocationRequest, ) -> io::Result<&(CapabilityLocationRequest, CapabilityLocations)> { let cached = self .prewarmed_locations .get_or_try_init(|| async { + let request = CapabilityLocationRequest { + codex_home: PathUri::from_abs_path(&codex_utils_home_dir::find_codex_home()?), + user_home: dirs::home_dir() + .and_then(|home| PathUri::from_host_native_path(home).ok()), + }; let locations = prewarm_locations(&self.file_system, &request).await?; - Ok::<_, io::Error>((request.clone(), locations)) + Ok::<_, io::Error>((request, locations)) }) .await?; - // Check the winning request even when another caller initialized the cell. - if cached.0 != request { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "capability home paths differ from the prewarmed configuration", - )); - } self.initialize_watchers(&cached.1).await; Ok(cached) } @@ -79,21 +78,9 @@ impl CapabilityManager { /// Reuses startup locations, retrying prewarming if the background scan failed. pub(crate) async fn get_or_refresh_discovery( self: &Arc, - location_request: CapabilityLocationRequest, params: DiscoverV2CapabilitiesRequest, ) -> io::Result> { - let (location_request, _) = match self.prewarm_locations(location_request).await { - Ok(prewarmed) => prewarmed, - Err(error) => { - return Ok(Arc::new(DiscoverV2CapabilitiesResponse { - warnings: vec![format!( - "could not prewarm capability discovery: {:?}", - error.kind() - )], - ..Default::default() - })); - } - }; + let (location_request, _) = self.prewarm_locations().await?; let location_state = self.get_or_refresh_locations(location_request).await?; let discovery = { let mut discoveries = location_state.sandbox_discoveries.lock().await; @@ -115,7 +102,7 @@ impl CapabilityManager { discovery } }; - if self.watchers.has_subscription() { + if location_state.registration.registration_succeeded() { return discovery .response .get_or_try_init(|| { @@ -140,8 +127,8 @@ impl CapabilityManager { &self, request: &CapabilityLocationRequest, ) -> io::Result> { - if self.watchers.has_subscription() - && let Some(location_state) = self.location_snapshot().await + if let Some(location_state) = self.location_snapshot().await + && location_state.registration.registration_succeeded() && !location_state.dirty.load(Ordering::Relaxed) { return Ok(location_state); @@ -151,13 +138,13 @@ impl CapabilityManager { .acquire() .await .map_err(io::Error::other)?; - if self.watchers.has_subscription() - && let Some(location_state) = self.location_snapshot().await + if let Some(location_state) = self.location_snapshot().await + && location_state.registration.registration_succeeded() && !location_state.dirty.load(Ordering::Relaxed) { return Ok(location_state); } - // Retry dirty locations; without a subscription, reload on every request. + // Retry dirty locations; failed watch registration reloads on every request. let locations = match prewarm_locations(&self.file_system, request).await { Ok(locations) => locations, Err(error) => { @@ -170,7 +157,7 @@ impl CapabilityManager { }; let next = Arc::new(LocationState { dirty: AtomicBool::new(false), - _registration: self.watchers.register_paths(&locations), + registration: self.watchers.register_paths(&locations), locations, sandbox_discoveries: Mutex::default(), }); diff --git a/codex-rs/exec-server/src/discoverV2/capability_manager/capability_watchers.rs b/codex-rs/exec-server/src/discoverV2/capability_manager/capability_watchers.rs index be92ca91a7..11af0eb074 100644 --- a/codex-rs/exec-server/src/discoverV2/capability_manager/capability_watchers.rs +++ b/codex-rs/exec-server/src/discoverV2/capability_manager/capability_watchers.rs @@ -47,10 +47,6 @@ impl CapabilityWatchers { } } - pub(super) fn has_subscription(&self) -> bool { - self.subscription.get().is_some() - } - /// Registers fixed roots, including missing ones, without separate watches for symlink targets. pub(super) fn register_paths(&self, locations: &CapabilityLocations) -> WatchRegistration { let Some(subscription) = self.subscription.get() else { @@ -83,7 +79,7 @@ impl CapabilityManager { dirty: AtomicBool::new(false), locations: prewarmed.clone(), sandbox_discoveries: Mutex::default(), - _registration: registration, + registration, }); let mut current_state = self.current_state.lock().await; let weak = Arc::downgrade(self); @@ -128,7 +124,7 @@ impl CapabilityManager { dirty: AtomicBool::new(false), locations, sandbox_discoveries: Mutex::default(), - _registration: registration, + registration, }); // Requests populate the new cache lazily; in-flight loads retain their old snapshot. *self.current_state.lock().await = Some(next); diff --git a/codex-rs/exec-server/src/discoverV2/mod.rs b/codex-rs/exec-server/src/discoverV2/mod.rs index faab1775e6..cfbe4c812d 100644 --- a/codex-rs/exec-server/src/discoverV2/mod.rs +++ b/codex-rs/exec-server/src/discoverV2/mod.rs @@ -1,9 +1,6 @@ //! Executor-side V2 capability discovery. //! Connection and RPC routing stay outside this module; legacy V1 scanning remains separate. -// Staged implementation; the RPC is registered in E04. -#![allow(dead_code)] - pub(crate) mod capability_discoveries; pub(crate) mod capability_locations; pub(crate) mod capability_manager; diff --git a/codex-rs/exec-server/src/discoverV2/tests/capability_discoveries_tests.rs b/codex-rs/exec-server/src/discoverV2/tests/capability_discoveries_tests.rs index 23f2a267ae..dd144081ec 100644 --- a/codex-rs/exec-server/src/discoverV2/tests/capability_discoveries_tests.rs +++ b/codex-rs/exec-server/src/discoverV2/tests/capability_discoveries_tests.rs @@ -87,15 +87,12 @@ async fn discovery_preserves_inventory_and_v1_metadata() -> anyhow::Result<()> { // V1 keeps undeclared/migrated skills and excludes the skill beyond its depth limit. assert_eq!(v1.skills.len(), 4); let manager = CapabilityManager::new(file_system); - manager.prewarm_locations(request.clone()).await?; + manager.seed_prewarmed_locations(request.clone()).await?; let response = manager - .get_or_refresh_discovery( - request.clone(), - DiscoverV2CapabilitiesRequest { - cwd: PathUri::from_host_native_path(temp.path().join("project"))?, - sandbox: None, - }, - ) + .get_or_refresh_discovery(DiscoverV2CapabilitiesRequest { + cwd: PathUri::from_host_native_path(temp.path().join("project"))?, + sandbox: None, + }) .await?; assert_eq!( response diff --git a/codex-rs/exec-server/src/discoverV2/tests/capability_manager_tests.rs b/codex-rs/exec-server/src/discoverV2/tests/capability_manager_tests.rs index ca300006e1..996492181e 100644 --- a/codex-rs/exec-server/src/discoverV2/tests/capability_manager_tests.rs +++ b/codex-rs/exec-server/src/discoverV2/tests/capability_manager_tests.rs @@ -7,11 +7,23 @@ use codex_protocol::protocol::SandboxPolicy; use codex_utils_path_uri::PathUri; use pretty_assertions::assert_eq; -async fn setup() -> anyhow::Result<( - tempfile::TempDir, - CapabilityLocationRequest, - Arc, -)> { +impl CapabilityManager { + // Supply fixture paths through the existing cell, without production override state. + pub(crate) async fn seed_prewarmed_locations( + self: &Arc, + request: CapabilityLocationRequest, + ) -> io::Result<()> { + self.prewarmed_locations + .get_or_try_init(|| async { + let locations = prewarm_locations(&self.file_system, &request).await?; + Ok::<_, io::Error>((request, locations)) + }) + .await?; + self.prewarm_locations().await.map(|_| ()) + } +} + +async fn setup() -> anyhow::Result<(tempfile::TempDir, Arc)> { let directory = tempfile::tempdir()?; write( directory.path(), @@ -21,36 +33,34 @@ async fn setup() -> anyhow::Result<( let request = request(&std::fs::canonicalize(directory.path())?)?; let file_system = LocalFileSystem::unsandboxed(); let manager = CapabilityManager::new(file_system); - manager.prewarm_locations(request.clone()).await?; + manager.seed_prewarmed_locations(request).await?; capability_watchers::tests::stop_listener(&manager).await?; - Ok((directory, request, manager)) + Ok((directory, manager)) } #[tokio::test] async fn discovery_responses_are_cached_by_sandbox_context() -> anyhow::Result<()> { - let (directory, request, manager) = setup().await?; + let (directory, manager) = setup().await?; + let (first_locations, second_locations) = + tokio::try_join!(manager.prewarm_locations(), manager.prewarm_locations())?; + assert!(std::ptr::eq(first_locations, second_locations)); let cwd = PathUri::from_host_native_path(directory.path().join("project"))?; let other_cwd = PathUri::from_host_native_path(directory.path().join("other"))?; let params = DiscoverV2CapabilitiesRequest { cwd: cwd.clone(), sandbox: None, }; - let first = manager - .get_or_refresh_discovery(request.clone(), params.clone()) - .await?; + let first = manager.get_or_refresh_discovery(params.clone()).await?; assert_eq!(first.skills.len(), 1); let snapshot = manager .location_snapshot() .await .ok_or_else(|| anyhow::anyhow!("missing snapshot"))?; let other = manager - .get_or_refresh_discovery( - request.clone(), - DiscoverV2CapabilitiesRequest { - cwd: other_cwd.clone(), - ..params.clone() - }, - ) + .get_or_refresh_discovery(DiscoverV2CapabilitiesRequest { + cwd: other_cwd.clone(), + ..params.clone() + }) .await?; assert!(Arc::ptr_eq(&first, &other)); @@ -65,16 +75,12 @@ async fn discovery_responses_are_cached_by_sandbox_context() -> anyhow::Result<( )?), ..params.clone() }; - let response = manager - .get_or_refresh_discovery(request.clone(), sandbox.clone()) - .await?; + let response = manager.get_or_refresh_discovery(sandbox.clone()).await?; assert_eq!(response, previous); assert!(!Arc::ptr_eq(&response, &previous)); assert!(Arc::ptr_eq( &response, - &manager - .get_or_refresh_discovery(request.clone(), sandbox) - .await? + &manager.get_or_refresh_discovery(sandbox).await? )); previous = response; } @@ -83,7 +89,7 @@ async fn discovery_responses_are_cached_by_sandbox_context() -> anyhow::Result<( MAX_CACHED_SANDBOX_CONTEXTS ); // The oldest context is evicted, but its previously returned response stays valid. - let reloaded = manager.get_or_refresh_discovery(request, params).await?; + let reloaded = manager.get_or_refresh_discovery(params).await?; assert_eq!(reloaded, first); assert!(!Arc::ptr_eq(&reloaded, &first)); assert_eq!( @@ -100,7 +106,7 @@ fn cancelled_initialization_retries_and_warning_responses_are_cached() -> anyhow .max_blocking_threads(1) .build()? .block_on(async { - let (directory, request, manager) = setup().await?; + let (directory, manager) = setup().await?; let params = DiscoverV2CapabilitiesRequest { cwd: PathUri::from_host_native_path(directory.path().join("project"))?, sandbox: None, @@ -118,8 +124,7 @@ fn cancelled_initialization_retries_and_warning_responses_are_cached() -> anyhow let _ = wait.recv(); }); ready.await?; - let mut first = - Box::pin(manager.get_or_refresh_discovery(request.clone(), params.clone())); + let mut first = Box::pin(manager.get_or_refresh_discovery(params.clone())); assert!(futures::poll!(&mut first).is_pending()); let discovery = snapshot .sandbox_discoveries @@ -128,8 +133,7 @@ fn cancelled_initialization_retries_and_warning_responses_are_cached() -> anyhow .first() .cloned() .ok_or_else(|| anyhow::anyhow!("missing sandbox discovery"))?; - let mut second = - Box::pin(manager.get_or_refresh_discovery(request.clone(), params.clone())); + let mut second = Box::pin(manager.get_or_refresh_discovery(params.clone())); assert!(futures::poll!(&mut second).is_pending()); assert!(discovery.response.get().is_none()); assert!(snapshot.sandbox_discoveries.try_lock().is_ok()); @@ -141,9 +145,7 @@ fn cancelled_initialization_retries_and_warning_responses_are_cached() -> anyhow assert_eq!(original.skills.len(), 1); assert!(Arc::ptr_eq( &original, - &manager - .get_or_refresh_discovery(request.clone(), params.clone()) - .await?, + &manager.get_or_refresh_discovery(params.clone()).await?, )); write( @@ -154,17 +156,15 @@ fn cancelled_initialization_retries_and_warning_responses_are_cached() -> anyhow // Invalidation drops entries instead of overwriting their initialized responses. snapshot.sandbox_discoveries.lock().await.clear(); let (partial, concurrent) = tokio::try_join!( - manager.get_or_refresh_discovery(request.clone(), params.clone()), - manager.get_or_refresh_discovery(request.clone(), params.clone()), + manager.get_or_refresh_discovery(params.clone()), + manager.get_or_refresh_discovery(params.clone()), )?; assert!(Arc::ptr_eq(&partial, &concurrent)); assert!(partial.skills.is_empty()); assert!(!partial.warnings.is_empty()); assert!(Arc::ptr_eq( &partial, - &manager - .get_or_refresh_discovery(request.clone(), params.clone()) - .await? + &manager.get_or_refresh_discovery(params.clone()).await? )); write( @@ -174,14 +174,10 @@ fn cancelled_initialization_retries_and_warning_responses_are_cached() -> anyhow )?; assert!(Arc::ptr_eq( &partial, - &manager - .get_or_refresh_discovery(request.clone(), params.clone()) - .await? + &manager.get_or_refresh_discovery(params.clone()).await? )); snapshot.sandbox_discoveries.lock().await.clear(); - let recovered = manager - .get_or_refresh_discovery(request.clone(), params.clone()) - .await?; + let recovered = manager.get_or_refresh_discovery(params.clone()).await?; assert!(recovered.warnings.is_empty()); assert_eq!( recovered @@ -193,9 +189,7 @@ fn cancelled_initialization_retries_and_warning_responses_are_cached() -> anyhow ); assert!(Arc::ptr_eq( &recovered, - &manager - .get_or_refresh_discovery(request.clone(), params) - .await? + &manager.get_or_refresh_discovery(params).await? )); Ok(()) }) diff --git a/codex-rs/exec-server/src/discoverV2/tests/capability_watchers_tests.rs b/codex-rs/exec-server/src/discoverV2/tests/capability_watchers_tests.rs index 4f3484bf6a..65ccf360be 100644 --- a/codex-rs/exec-server/src/discoverV2/tests/capability_watchers_tests.rs +++ b/codex-rs/exec-server/src/discoverV2/tests/capability_watchers_tests.rs @@ -38,13 +38,12 @@ async fn watch_events_invalidate_discovery_and_no_subscription_disables_caching( let request = request(&root)?; let demo_skill = request.codex_home.join("skills/demo/SKILL.md")?; let manager = CapabilityManager::new(LocalFileSystem::unsandboxed()); + manager.seed_prewarmed_locations(request.clone()).await?; let params = DiscoverV2CapabilitiesRequest { cwd: PathUri::from_host_native_path(&root)?, sandbox: None, }; - let original = manager - .get_or_refresh_discovery(request.clone(), params.clone()) - .await?; + let original = manager.get_or_refresh_discovery(params.clone()).await?; assert!( original .skills @@ -121,9 +120,7 @@ async fn watch_events_invalidate_discovery_and_no_subscription_disables_caching( .context("native watcher missed creation of a global root")??; // Drive subsequent refreshes explicitly so native delivery cannot race cache assertions. stop_listener(&manager).await?; - let refreshed = manager - .get_or_refresh_discovery(request.clone(), params.clone()) - .await?; + let refreshed = manager.get_or_refresh_discovery(params.clone()).await?; let added_skill = home_skills.join("added/SKILL.md")?; assert!( refreshed @@ -145,9 +142,7 @@ async fn watch_events_invalidate_discovery_and_no_subscription_disables_caching( )?), ..params.clone() }; - manager - .get_or_refresh_discovery(request.clone(), sandbox_params) - .await?; + manager.get_or_refresh_discovery(sandbox_params).await?; let retained = manager .location_snapshot() .await @@ -189,9 +184,7 @@ async fn watch_events_invalidate_discovery_and_no_subscription_disables_caching( .map(|plugin| plugin.version.as_str()), Some("1.0.0") ); - let after = manager - .get_or_refresh_discovery(request.clone(), params.clone()) - .await?; + let after = manager.get_or_refresh_discovery(params.clone()).await?; assert!( after .skills @@ -202,9 +195,7 @@ async fn watch_events_invalidate_discovery_and_no_subscription_disables_caching( assert_eq!(next.sandbox_discoveries.lock().await.len(), 1); assert!(Arc::ptr_eq( &after, - &manager - .get_or_refresh_discovery(request.clone(), params.clone()) - .await? + &manager.get_or_refresh_discovery(params.clone()).await? )); #[cfg(unix)] @@ -215,9 +206,7 @@ async fn watch_events_invalidate_discovery_and_no_subscription_disables_caching( std::fs::rename(&cache, &saved_cache)?; std::os::unix::fs::symlink(&cache, &cache)?; manager.refresh_locations().await; - let fallback = manager - .get_or_refresh_discovery(request.clone(), params.clone()) - .await?; + let fallback = manager.get_or_refresh_discovery(params.clone()).await?; assert!(Arc::ptr_eq(&after, &fallback)); std::fs::remove_file(&cache)?; std::fs::rename(&saved_cache, &cache)?; @@ -226,9 +215,7 @@ async fn watch_events_invalidate_discovery_and_no_subscription_disables_caching( "codex/skills/recovered/SKILL.md", "---\nname: recovered\ndescription: recovered after failed refresh\n---\n", )?; - let recovered = manager - .get_or_refresh_discovery(request.clone(), params.clone()) - .await?; + let recovered = manager.get_or_refresh_discovery(params.clone()).await?; assert!( recovered .skills @@ -238,9 +225,7 @@ async fn watch_events_invalidate_discovery_and_no_subscription_disables_caching( assert!(!Arc::ptr_eq(&fallback, &recovered)); assert!(Arc::ptr_eq( &recovered, - &manager - .get_or_refresh_discovery(request.clone(), params.clone()) - .await? + &manager.get_or_refresh_discovery(params.clone()).await? )); } @@ -250,9 +235,8 @@ async fn watch_events_invalidate_discovery_and_no_subscription_disables_caching( .context("new manager unexpectedly shared")? .watchers .file_watcher = None; - let before = uncached - .get_or_refresh_discovery(request.clone(), params.clone()) - .await?; + uncached.seed_prewarmed_locations(request.clone()).await?; + let before = uncached.get_or_refresh_discovery(params.clone()).await?; assert!(uncached.watchers.subscription.get().is_none()); write( &root, @@ -264,9 +248,7 @@ async fn watch_events_invalidate_discovery_and_no_subscription_disables_caching( "codex/plugins/cache/acme/calendar/1.0.0/.codex-plugin/plugin.json", r#"{"name":"calendar"}"#, )?; - let after = uncached - .get_or_refresh_discovery(request.clone(), params.clone()) - .await?; + let after = uncached.get_or_refresh_discovery(params.clone()).await?; assert_eq!(after.plugins.len(), before.plugins.len() + 1); assert!( after @@ -280,7 +262,72 @@ async fn watch_events_invalidate_discovery_and_no_subscription_disables_caching( .iter() .any(|skill| skill.path == demo_skill && skill.description == "uncached edit") ); - let repeated = uncached.get_or_refresh_discovery(request, params).await?; + let repeated = uncached.get_or_refresh_discovery(params).await?; + assert_eq!(after, repeated); + assert!(!Arc::ptr_eq(&after, &repeated)); + Ok(()) +} + +#[tokio::test] +async fn failed_watch_registration_reloads_locations_and_contents() -> anyhow::Result<()> { + let directory = tempfile::tempdir()?; + let root = std::fs::canonicalize(directory.path())?; + write( + &root, + "codex/skills/demo/SKILL.md", + "---\nname: demo\ndescription: original\n---\n", + )?; + let request = request(&root)?; + let params = DiscoverV2CapabilitiesRequest { + cwd: PathUri::from_host_native_path(&root)?, + sandbox: None, + }; + let mut manager = CapabilityManager::new(LocalFileSystem::unsandboxed()); + // A subscriber exists, but no OS watches can be registered. + Arc::get_mut(&mut manager) + .context("new manager unexpectedly shared")? + .watchers + .file_watcher = Some(Arc::new(FileWatcher::noop())); + manager.seed_prewarmed_locations(request.clone()).await?; + let before = manager.get_or_refresh_discovery(params.clone()).await?; + assert_eq!(before.skills.len(), 1); + assert!(before.plugins.is_empty()); + assert!(manager.watchers.subscription.get().is_some()); + let initial = manager + .location_snapshot() + .await + .context("missing snapshot")?; + assert!(!initial.registration.registration_succeeded()); + + write( + &root, + "codex/skills/demo/SKILL.md", + "---\nname: demo\ndescription: edited\n---\n", + )?; + write( + &root, + "codex/plugins/cache/acme/calendar/1.0.0/.codex-plugin/plugin.json", + r#"{"name":"calendar"}"#, + )?; + let after = manager.get_or_refresh_discovery(params.clone()).await?; + assert_eq!(after.skills.len(), 1); + assert_eq!(after.skills[0].description, "edited"); + assert_eq!(after.plugins.len(), 1); + assert_eq!(after.plugins[0].id, "calendar@acme"); + let snapshot = manager + .location_snapshot() + .await + .context("missing refreshed snapshot")?; + assert!(!Arc::ptr_eq(&initial, &snapshot)); + assert!( + snapshot + .sandbox_discoveries + .lock() + .await + .iter() + .all(|discovery| discovery.response.get().is_none()) + ); + let repeated = manager.get_or_refresh_discovery(params).await?; assert_eq!(after, repeated); assert!(!Arc::ptr_eq(&after, &repeated)); Ok(()) diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 277f218069..8c3d89b88d 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -1,4 +1,6 @@ mod connect_options; +mod discovery_recovery; +mod discovery_v2; use std::collections::HashMap; use std::collections::HashSet; @@ -1063,7 +1065,7 @@ impl Environment { ) -> Result { match &self.remote_client { Some(client) => { - let mut connection_state = client.subscribe_connection_state(); + let connection_state = client.subscribe_connection_state(); let client = client.get().await?; let discover = || async { if params.roots.iter().any(|root| { @@ -1083,37 +1085,8 @@ impl Environment { } client.discover_capability_roots(params.clone()).await }; - match discover().await { - Err(error) if crate::client::is_retryable_recovery_error(&error) => { - tracing::warn!(%error, "replaying capability discovery after executor recovery"); - let recovered = - tokio::time::timeout(std::time::Duration::from_secs(8), async { - loop { - match self.readiness_result() { - Some(Ok(())) => return Some(Ok(())), - Some(Err(error)) - if !crate::client::is_retryable_recovery_error( - &error, - ) => - { - return Some(Err(error)); - } - Some(Err(_)) | None => {} - } - if connection_state.changed().await.is_err() { - return None; - } - } - }) - .await; - match recovered { - Ok(Some(Ok(()))) => discover().await, - Ok(Some(Err(error))) => Err(error), - Ok(None) | Err(_) => Err(error), - } - } - response => response, - } + self.discover_with_recovery(connection_state, discover) + .await } None => crate::discover_capability_roots(self.filesystem.as_ref(), params) .await diff --git a/codex-rs/exec-server/src/environment/discovery_recovery.rs b/codex-rs/exec-server/src/environment/discovery_recovery.rs new file mode 100644 index 0000000000..a4f23a0045 --- /dev/null +++ b/codex-rs/exec-server/src/environment/discovery_recovery.rs @@ -0,0 +1,49 @@ +//! Shared recovery for idempotent V1 and V2 discovery requests. + +use super::Environment; +use super::EnvironmentConnectionState; +use crate::ExecServerError; +use crate::client::is_retryable_recovery_error; +use std::future::Future; +use std::time::Duration; +use tokio::sync::watch; + +impl Environment { + /// Replays discovery once after recovery; callers subscribe before obtaining the client. + pub(super) async fn discover_with_recovery( + &self, + mut connection_state: watch::Receiver, + discover: F, + ) -> Result + where + F: Fn() -> Fut, + Fut: Future>, + { + match discover().await { + Err(error) if is_retryable_recovery_error(&error) => { + tracing::warn!(%error, "replaying capability discovery after executor recovery"); + let recovered = tokio::time::timeout(Duration::from_secs(8), async { + loop { + match self.readiness_result() { + Some(Ok(())) => return Some(Ok(())), + Some(Err(error)) if !is_retryable_recovery_error(&error) => { + return Some(Err(error)); + } + Some(Err(_)) | None => {} + } + if connection_state.changed().await.is_err() { + return None; + } + } + }) + .await; + match recovered { + Ok(Some(Ok(()))) => discover().await, + Ok(Some(Err(error))) => Err(error), + Ok(None) | Err(_) => Err(error), + } + } + response => response, + } + } +} diff --git a/codex-rs/exec-server/src/environment/discovery_v2.rs b/codex-rs/exec-server/src/environment/discovery_v2.rs new file mode 100644 index 0000000000..ec927bae8c --- /dev/null +++ b/codex-rs/exec-server/src/environment/discovery_v2.rs @@ -0,0 +1,43 @@ +//! Acquires executor inventories without activation and replays safe reads after recovery. + +use super::Environment; +use crate::DiscoverV2CapabilitiesRequest; +use crate::DiscoverV2CapabilitiesResponse; +use crate::ExecServerError; + +impl Environment { + /// Returns metadata for the caller's cwd from the remote session's V2 lifecycle. + /// Local environments and older executors retain the V1 selected-root API. + pub async fn discover_capabilities_v2( + &self, + params: DiscoverV2CapabilitiesRequest, + ) -> Result { + let Some(client) = &self.remote_client else { + return Err(ExecServerError::Protocol( + "capabilities/discoverV2 requires a remote environment".to_string(), + )); + }; + let connection_state = client.subscribe_connection_state(); + let client = client.get().await?; + let discover = || async { + if !client + .environment_info() + .await? + .capabilities + .capability_discovery_v2 + { + return Err(ExecServerError::Protocol( + "capabilities/discoverV2 is unavailable in this environment session" + .to_string(), + )); + } + client.discover_capabilities_v2(params.clone()).await + }; + self.discover_with_recovery(connection_state, discover) + .await + } +} + +#[cfg(test)] +#[path = "discovery_v2_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/environment/discovery_v2_tests.rs b/codex-rs/exec-server/src/environment/discovery_v2_tests.rs new file mode 100644 index 0000000000..db280f09d9 --- /dev/null +++ b/codex-rs/exec-server/src/environment/discovery_v2_tests.rs @@ -0,0 +1,100 @@ +//! Real transport recovery preserves a discovery request's cwd and retries it once. + +use super::DiscoverV2CapabilitiesRequest; +use super::DiscoverV2CapabilitiesResponse; +use super::Environment; +use crate::CAPABILITIES_DISCOVER_V2_METHOD; +use crate::client_api::ExecServerTransportParams; +use codex_config::ScopedSkillsConfig; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; +use codex_utils_path_uri::PathUri; +use futures::SinkExt; +use futures::StreamExt; +use pretty_assertions::assert_eq; +use serde_json::json; +use std::sync::Arc; +use std::time::Duration; +use tokio::net::TcpListener; +use tokio::sync::Mutex; +use tokio_tungstenite::accept_async; +use tokio_tungstenite::tungstenite::Message; +use tokio_util::task::AbortOnDropHandle; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn discovery_retries_disconnect_within_the_same_request() -> anyhow::Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let url = format!("ws://{}", listener.local_addr()?); + let attempts = Arc::new(Mutex::new(Vec::new())); + let recorded = Arc::clone(&attempts); + let _server = AbortOnDropHandle::new(tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + let mut socket = accept_async(stream).await?; + while let Some(message) = socket.next().await { + let Message::Text(text) = message? else { + continue; + }; + let request: serde_json::Value = serde_json::from_str(text.as_ref())?; + if request.get("id").is_none() { + continue; + } + let info = json!({ + "shell":{"name":"bash","path":"/bin/bash"}, + "capabilities":{"capabilityDiscoveryV2":true} + }); + let result = match request["method"].as_str() { + Some("initialize") => { + // V2 support is advertised by the executor, not negotiated by the client. + assert!(request["params"].get("capabilityDiscoveryV2").is_none()); + json!({"sessionId":"discovery-recovery", "environmentInfo":info}) + } + Some("environment/info") => info, + Some(CAPABILITIES_DISCOVER_V2_METHOD) => { + let first = { + let mut attempts = recorded.lock().await; + attempts.push(request["params"].clone()); + attempts.len() == 1 + }; + if first { + break; + } + serde_json::to_value(DiscoverV2CapabilitiesResponse::default())? + } + method => anyhow::bail!("unexpected executor method: {method:?}"), + }; + socket + .send(Message::Text( + json!({"id":request["id"],"result":result}) + .to_string() + .into(), + )) + .await?; + } + } + anyhow::Ok(()) + })); + let environment = Environment::remote_with_transport( + ExecServerTransportParams::websocket_url(url, Duration::from_secs(5)), + /*local_runtime_paths*/ None, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ScopedSkillsConfig::default(), + ); + let params = DiscoverV2CapabilitiesRequest { + cwd: PathUri::parse("file:///caller/project-b")?, + sandbox: None, + }; + let response = tokio::time::timeout( + Duration::from_secs(12), + environment.discover_capabilities_v2(params.clone()), + ) + .await??; + assert_eq!(response, DiscoverV2CapabilitiesResponse::default()); + assert_eq!( + *attempts.lock().await, + vec![ + serde_json::to_value(¶ms)?, + serde_json::to_value(¶ms)? + ] + ); + Ok(()) +} diff --git a/codex-rs/exec-server/src/server/handler.rs b/codex-rs/exec-server/src/server/handler.rs index d4ecb1b07f..e7f41edcba 100644 --- a/codex-rs/exec-server/src/server/handler.rs +++ b/codex-rs/exec-server/src/server/handler.rs @@ -13,10 +13,13 @@ use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; +use crate::DiscoverV2CapabilitiesRequest; +use crate::DiscoverV2CapabilitiesResponse; use crate::ExecServerRuntimeOptions; use crate::client::http_client::PendingRouteAwareHttpBodyStream; use crate::client::http_client::RouteAwareHttpClient; use crate::client::http_client::RouteAwareHttpRequestRunner; +use crate::discover_v2::capability_manager::CapabilityManager; use crate::environment_config::ReadEnvironmentConfigError; use crate::environment_config::read_environment_config; use crate::protocol::CapabilityRootsDiscoverParams; @@ -79,6 +82,7 @@ use crate::telemetry::ExecutorRegistration; pub(crate) struct ExecServerHandler { pub(super) executor_registration: Option>, session_registry: Arc, + capability_manager: Arc, notifications: RpcNotificationSender, session: StdMutex>, active_body_stream_ids: Mutex>, @@ -94,6 +98,7 @@ pub(crate) struct ExecServerHandler { impl ExecServerHandler { pub(crate) fn new( session_registry: Arc, + capability_manager: Arc, notifications: RpcNotificationSender, runtime_paths: ExecServerRuntimeOptions, http_client_factory: HttpClientFactory, @@ -101,6 +106,7 @@ impl ExecServerHandler { Self { executor_registration: None, session_registry, + capability_manager, notifications, session: StdMutex::new(None), active_body_stream_ids: Mutex::new(HashSet::new()), @@ -330,6 +336,17 @@ impl ExecServerHandler { self.file_system.discover_capability_roots(params).await } + pub(crate) async fn discover_capabilities_v2( + &self, + params: DiscoverV2CapabilitiesRequest, + ) -> Result, JSONRPCErrorError> { + self.require_initialized_for("capability discovery")?; + self.capability_manager + .get_or_refresh_discovery(params) + .await + .map_err(|error| internal_error(error.to_string())) + } + pub(crate) async fn fs_open( &self, params: FsOpenParams, diff --git a/codex-rs/exec-server/src/server/handler/tests.rs b/codex-rs/exec-server/src/server/handler/tests.rs index 00de894176..bf42de8feb 100644 --- a/codex-rs/exec-server/src/server/handler/tests.rs +++ b/codex-rs/exec-server/src/server/handler/tests.rs @@ -12,6 +12,7 @@ use uuid::Uuid; use super::ExecServerHandler; use crate::ExecServerRuntimeOptions; use crate::ProcessId; +use crate::discover_v2::capability_manager::CapabilityManager; use crate::protocol::ExecParams; use crate::protocol::InitializeParams; use crate::protocol::ReadParams; @@ -94,6 +95,7 @@ async fn initialized_handler() -> Arc { let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); let handler = Arc::new(ExecServerHandler::new( registry, + CapabilityManager::new(crate::LocalFileSystem::default()), RpcNotificationSender::new(outgoing_tx), test_runtime_paths(), test_http_client_factory(), @@ -173,6 +175,7 @@ async fn long_poll_read_fails_after_session_resume() { let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); let first_handler = Arc::new(ExecServerHandler::new( Arc::clone(®istry), + CapabilityManager::new(crate::LocalFileSystem::default()), RpcNotificationSender::new(first_tx), test_runtime_paths(), test_http_client_factory(), @@ -217,6 +220,7 @@ async fn long_poll_read_fails_after_session_resume() { let (second_tx, _second_rx) = mpsc::channel(16); let second_handler = Arc::new(ExecServerHandler::new( registry, + CapabilityManager::new(crate::LocalFileSystem::default()), RpcNotificationSender::new(second_tx), test_runtime_paths(), test_http_client_factory(), @@ -251,6 +255,7 @@ async fn active_session_resume_is_rejected() { let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); let first_handler = Arc::new(ExecServerHandler::new( Arc::clone(®istry), + CapabilityManager::new(crate::LocalFileSystem::default()), RpcNotificationSender::new(first_tx), test_runtime_paths(), test_http_client_factory(), @@ -266,6 +271,7 @@ async fn active_session_resume_is_rejected() { let (second_tx, _second_rx) = mpsc::channel(16); let second_handler = Arc::new(ExecServerHandler::new( registry, + CapabilityManager::new(crate::LocalFileSystem::default()), RpcNotificationSender::new(second_tx), test_runtime_paths(), test_http_client_factory(), @@ -295,6 +301,7 @@ async fn output_and_exit_are_retained_after_notification_receiver_closes() { let (outgoing_tx, outgoing_rx) = mpsc::channel(16); let handler = Arc::new(ExecServerHandler::new( SessionRegistry::new(crate::ExecServerTelemetry::default()), + CapabilityManager::new(crate::LocalFileSystem::default()), RpcNotificationSender::new(outgoing_tx), test_runtime_paths(), test_http_client_factory(), diff --git a/codex-rs/exec-server/src/server/process_otel_tests.rs b/codex-rs/exec-server/src/server/process_otel_tests.rs index 93638142d8..e7142c3dcb 100644 --- a/codex-rs/exec-server/src/server/process_otel_tests.rs +++ b/codex-rs/exec-server/src/server/process_otel_tests.rs @@ -40,6 +40,7 @@ use super::session_registry::SessionRegistry; use crate::ExecServerRuntimeOptions; use crate::ExecServerTelemetry; use crate::connection::JsonRpcConnectionEvent; +use crate::discover_v2::capability_manager::CapabilityManager; use crate::protocol::EXEC_METHOD; use crate::protocol::EXEC_READ_METHOD; use crate::protocol::EXEC_TERMINATE_METHOD; @@ -738,6 +739,7 @@ fn new_handler( ) -> Arc { let mut handler = ExecServerHandler::new( Arc::clone(sessions), + CapabilityManager::new(crate::LocalFileSystem::default()), RpcNotificationSender::new(outgoing), ExecServerRuntimeOptions::new( std::env::current_exe().expect("test executable"), diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 81cf9c54bf..0db3bc38a3 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -1,4 +1,3 @@ -use std::io; use std::sync::Arc; use std::time::Instant; @@ -14,7 +13,6 @@ use crate::LocalFileSystem; use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; -use crate::discover_v2::capability_locations::CapabilityLocationRequest; use crate::discover_v2::capability_manager::CapabilityManager; use crate::rpc::RpcCallError; use crate::rpc::RpcNotificationSender; @@ -31,7 +29,6 @@ use crate::telemetry::ConnectionTransport; use crate::telemetry::ExecServerTelemetry; use crate::telemetry::ExecutorRegistration; use codex_http_client::HttpClientFactory; -use codex_utils_path_uri::PathUri; #[derive(Clone)] pub(crate) struct ConnectionProcessor { @@ -45,49 +42,11 @@ pub(crate) struct ConnectionProcessor { } impl ConnectionProcessor { - #[cfg(test)] - pub(crate) fn new(runtime_paths: ExecServerRuntimeOptions) -> Self { - Self::new_with_location_request( - runtime_paths, - ExecServerTelemetry::default(), - codex_http_client::HttpClientFactory::new( - codex_http_client::OutboundProxyPolicy::ReqwestDefault, - ), - RequestDispatchMode::Inline, - // Connection-only tests must not scan the real user home. - Err(io::Error::other( - "capability home paths not supplied by this test", - )), - ) - } - pub(crate) fn new_with_telemetry( runtime_paths: ExecServerRuntimeOptions, telemetry: ExecServerTelemetry, http_client_factory: HttpClientFactory, request_dispatch_mode: RequestDispatchMode, - ) -> Self { - let request = - codex_utils_home_dir::find_codex_home().map(|codex_home| CapabilityLocationRequest { - codex_home: PathUri::from_abs_path(&codex_home), - user_home: dirs::home_dir() - .and_then(|home| PathUri::from_host_native_path(home).ok()), - }); - Self::new_with_location_request( - runtime_paths, - telemetry, - http_client_factory, - request_dispatch_mode, - request, - ) - } - - fn new_with_location_request( - runtime_paths: ExecServerRuntimeOptions, - telemetry: ExecServerTelemetry, - http_client_factory: HttpClientFactory, - request_dispatch_mode: RequestDispatchMode, - request: io::Result, ) -> Self { // Library callers may bypass CLI startup. Capture the version before serving clients. let _ = BuildInfo::get(); @@ -97,11 +56,7 @@ impl ConnectionProcessor { // Start before any client connects; all sessions share these inert global locations. let prewarm = tokio::spawn(async move { // Home resolution failures are nonfatal, just like scan failures. - let result = match request { - Ok(request) => manager.prewarm_locations(request).await.map(|_| ()), - Err(error) => Err(error), - }; - if let Err(error) = result { + if let Err(error) = manager.prewarm_locations().await { tracing::warn!(error_kind = ?error.kind(), "capability location prewarming unavailable"); } }); @@ -157,7 +112,7 @@ async fn run_connection( executor_registration: Option>, ) { let ConnectionProcessor { - capability_manager: _capability_manager, + capability_manager, capability_prewarm: _capability_prewarm, session_registry, runtime_paths, @@ -179,6 +134,7 @@ async fn run_connection( let requests = notifications.request_sender(); let mut handler = ExecServerHandler::new( session_registry, + capability_manager, notifications, runtime_paths, http_client_factory, @@ -370,67 +326,26 @@ mod tests { use crate::rpc_server_requests::RpcServerRequestSender; use crate::server::session_registry::SessionRegistry; - #[tokio::test] - async fn startup_prewarms_without_a_client_connection() -> anyhow::Result<()> { - let paths = ExecServerRuntimeOptions::new( - std::env::current_exe()?, - /*codex_linux_sandbox_exe*/ None, - )?; - let directory = tempfile::tempdir()?; - let request = crate::discover_v2::capability_locations::CapabilityLocationRequest { - codex_home: PathUri::from_host_native_path(directory.path().join("codex"))?, - user_home: Some(PathUri::from_host_native_path( - directory.path().join("home"), - )?), - }; - let skill_root = directory.path().join("codex/skills"); - std::fs::create_dir_all(&skill_root)?; - let mut processor = super::ConnectionProcessor::new_with_location_request( - paths, - crate::ExecServerTelemetry::default(), - codex_http_client::HttpClientFactory::new( - codex_http_client::OutboundProxyPolicy::ReqwestDefault, - ), - crate::server::RequestDispatchMode::Inline, - Ok(request.clone()), - ); - Arc::get_mut(&mut processor.capability_prewarm) - .ok_or_else(|| anyhow::anyhow!("startup task unexpectedly shared"))? - .await?; - let cloned_processor = processor.clone(); - assert!(Arc::ptr_eq( - &processor.capability_manager, - &cloned_processor.capability_manager, - )); - let manager = &processor.capability_manager; - // Removing the fixture proves startup retained it before any later call. - std::fs::remove_dir(&skill_root)?; - let (first, second) = tokio::try_join!( - manager.prewarm_locations(request.clone()), - manager.prewarm_locations(request.clone()), - )?; - assert!(std::ptr::eq(first, second)); - assert_eq!(first.0, request); - let expected_root = request.codex_home.join("skills")?; - assert!( - first - .1 - .locations - .iter() - .any(|location| location.root == expected_root) - ); - let mut different_request = request.clone(); - different_request.codex_home = request.codex_home.join("different")?; - let Err(error) = manager.prewarm_locations(different_request).await else { - anyhow::bail!("different home paths must not reuse cached locations"); - }; - assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput); - assert!(std::ptr::eq( - first, - manager.prewarm_locations(request).await?, - )); - processor.shutdown().await; - Ok(()) + impl super::ConnectionProcessor { + // Connection-only tests do not prewarm or scan the real user home. + pub(crate) fn new(runtime_paths: ExecServerRuntimeOptions) -> Self { + let telemetry = crate::ExecServerTelemetry::default(); + Self { + capability_manager: crate::discover_v2::capability_manager::CapabilityManager::new( + crate::LocalFileSystem::with_runtime_paths(runtime_paths.clone()), + ), + capability_prewarm: Arc::new(tokio_util::task::AbortOnDropHandle::new( + tokio::spawn(async {}), + )), + session_registry: SessionRegistry::new(telemetry.clone()), + runtime_paths, + telemetry, + http_client_factory: codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + request_dispatch_mode: crate::server::RequestDispatchMode::Inline, + } + } } #[tokio::test] diff --git a/codex-rs/exec-server/src/server/registry.rs b/codex-rs/exec-server/src/server/registry.rs index 609a3a748e..842af18bee 100644 --- a/codex-rs/exec-server/src/server/registry.rs +++ b/codex-rs/exec-server/src/server/registry.rs @@ -5,7 +5,9 @@ use opentelemetry::trace::TraceContextExt; use std::sync::Arc; use crate::local_file_system::resolve_existing_path; +use crate::protocol::CAPABILITIES_DISCOVER_V2_METHOD; use crate::protocol::CAPABILITY_ROOTS_DISCOVER_METHOD; +use crate::protocol::DiscoverV2CapabilitiesRequest; use crate::protocol::ENVIRONMENT_CONFIG_READ_METHOD; use crate::protocol::ENVIRONMENT_INFO_METHOD; use crate::protocol::ENVIRONMENT_STATUS_METHOD; @@ -109,6 +111,12 @@ pub(crate) fn build_router() -> RpcRouter { handler.discover_capability_roots(params).await }, ); + router.request( + CAPABILITIES_DISCOVER_V2_METHOD, + |handler: Arc, params: DiscoverV2CapabilitiesRequest| async move { + handler.discover_capabilities_v2(params).await + }, + ); router.request( EXEC_READ_METHOD, |handler: Arc, params: ReadParams| async move { diff --git a/codex-rs/exec-server/src/server/request_dispatcher_tests.rs b/codex-rs/exec-server/src/server/request_dispatcher_tests.rs index 5da5ed0d2f..680a4bea22 100644 --- a/codex-rs/exec-server/src/server/request_dispatcher_tests.rs +++ b/codex-rs/exec-server/src/server/request_dispatcher_tests.rs @@ -33,6 +33,7 @@ use super::RequestDispatcher; use super::RequestTaskResult; use crate::ExecServerRuntimeOptions; use crate::connection::JsonRpcConnectionEvent; +use crate::discover_v2::capability_manager::CapabilityManager; use crate::rpc::RpcNotificationSender; use crate::rpc::RpcRouter; use crate::rpc::RpcServerOutboundMessage; @@ -177,6 +178,7 @@ async fn request_queue_waits_for_dispatcher_admission_before_recording_telemetry let requests = notifications.request_sender(); let handler = Arc::new(ExecServerHandler::new( SessionRegistry::new(telemetry.clone()), + CapabilityManager::new(crate::LocalFileSystem::default()), notifications, ExecServerRuntimeOptions::new( std::env::current_exe().expect("current executable"), @@ -386,6 +388,7 @@ fn request_telemetry_fixture(router: RpcRouter) -> RequestTel let requests = notifications.request_sender(); let handler = Arc::new(ExecServerHandler::new( SessionRegistry::new(telemetry.clone()), + CapabilityManager::new(crate::LocalFileSystem::default()), notifications, ExecServerRuntimeOptions::new( std::env::current_exe().expect("current executable"), diff --git a/codex-rs/exec-server/tests/capabilities_v2.rs b/codex-rs/exec-server/tests/capabilities_v2.rs new file mode 100644 index 0000000000..3eacf46f85 --- /dev/null +++ b/codex-rs/exec-server/tests/capabilities_v2.rs @@ -0,0 +1,214 @@ +//! V2 RPC metadata refresh and sandbox isolation over a real connection. +mod common; + +use anyhow::Context; +use codex_exec_server::CAPABILITIES_DISCOVER_V2_METHOD; +use codex_exec_server::DiscoverV2CapabilitiesRequest; +use codex_exec_server::DiscoverV2CapabilitiesResponse; +use codex_exec_server::InitializeResponse; +use codex_exec_server_protocol::JSONRPCMessage; +use codex_utils_path_uri::PathUri; +use common::exec_server::ExecServerHarness; +use common::exec_server::exec_server_with_env; + +#[cfg(any(target_os = "linux", target_os = "macos", windows))] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn v2_rpc_refreshes_metadata_and_enforces_sandbox() -> anyhow::Result<()> { + let home = tempfile::tempdir()?; + let home_path = std::fs::canonicalize(home.path())?; + let codex_home = tempfile::tempdir()?; + let codex_home_path = std::fs::canonicalize(codex_home.path())?; + let skills = codex_home_path.join("skills/demo/SKILL.md"); + let plugin = codex_home_path.join("plugins/cache/market/mail/1.0.0"); + write( + &plugin.join(".codex-plugin/plugin.json"), + r#"{"name":"mail"}"#, + )?; + write( + &plugin.join("skills/search/SKILL.md"), + "---\nname: search\ndescription: Search mail\n---\n", + )?; + write( + &skills, + "---\nname: demo\ndescription: original\n---\nPRIVATE INSTRUCTION BODY", + )?; + let mut server = exec_server_with_env( + [ + ("HOME", &home_path), + ("USERPROFILE", &home_path), + ("CODEX_HOME", &codex_home_path), + ], + &[], + ) + .await?; + // Existing clients need no discovery-specific initialize field. + let initialized: InitializeResponse = serde_json::from_value( + request( + &mut server, + "initialize", + serde_json::json!({"clientName": "v2-client"}), + ) + .await?, + )?; + server + .send_notification("initialized", serde_json::json!({})) + .await?; + let info = initialized + .environment_info + .context("initialize environment info")?; + assert!(info.capabilities.capability_discovery_v2); + // Keep the requested cwd isolated from the test runner's repository. + let project = tempfile::tempdir()?; + let cwd = PathUri::from_host_native_path(project.path())?; + let params = DiscoverV2CapabilitiesRequest { + cwd: cwd.clone(), + sandbox: None, + }; + let first = discover(&mut server, ¶ms).await?; + assert!( + first + .skills + .iter() + .any(|skill| skill.name == "demo" && skill.description == "original") + ); + assert!(!serde_json::to_string(&first)?.contains("PRIVATE INSTRUCTION BODY")); + + assert!( + first + .plugins + .iter() + .any(|plugin| plugin.id == "mail@market" && plugin.skills.len() == 1) + ); + + // Exercise one content refresh through the real RPC and native watcher. + write(&skills, "---\nname: demo\ndescription: updated\n---\n")?; + tokio::time::timeout(std::time::Duration::from_secs(30), async { + let mut interval = tokio::time::interval(std::time::Duration::from_millis(50)); + loop { + interval.tick().await; + if discover(&mut server, ¶ms) + .await? + .skills + .iter() + .any(|skill| skill.name == "demo" && skill.description == "updated") + { + break anyhow::Ok(()); + } + } + }) + .await + .context("waiting for discovery refresh")??; + + // A restricted request must not reuse the unrestricted context's cached plugin metadata. + { + use codex_exec_server::FileSystemSandboxContext; + use codex_protocol::models::PermissionProfile; + use codex_protocol::permissions::FileSystemAccessMode; + use codex_protocol::permissions::FileSystemSandboxEntry; + use codex_protocol::permissions::FileSystemSandboxPolicy; + use codex_protocol::permissions::NetworkSandboxPolicy; + use codex_utils_absolute_path::AbsolutePathBuf; + let allowed = codex_home_path.join("skills"); + let policy = FileSystemSandboxPolicy::restricted(vec![FileSystemSandboxEntry::new( + AbsolutePathBuf::from_absolute_path(&allowed)?.into(), + FileSystemAccessMode::Read, + )]); + let sandbox = FileSystemSandboxContext::from_permission_profile( + PermissionProfile::from_runtime_permissions(&policy, NetworkSandboxPolicy::Restricted), + cwd.clone(), + ); + #[cfg(windows)] + let sandbox = FileSystemSandboxContext { + windows_sandbox_selection: codex_exec_server::WindowsSandboxSelection::RestrictedToken, + ..sandbox + }; + let restricted = DiscoverV2CapabilitiesRequest { + cwd: cwd.clone(), + sandbox: Some(sandbox), + }; + // Linux (Bubblewrap) and macOS (Seatbelt) enforce the allowed directory directly. + #[cfg(any(target_os = "linux", target_os = "macos"))] + { + let response = discover(&mut server, &restricted).await?; + assert!(response.plugins.is_empty()); + assert!(response.skills.iter().any(|skill| skill.name == "demo")); + assert!(!serde_json::to_string(&response)?.contains("mail@market")); + } + // Windows restricted tokens reject split read access; a supported read-only policy succeeds. + #[cfg(windows)] + { + let Err(error) = discover(&mut server, &restricted).await else { + anyhow::bail!("unsupported read restrictions must not reuse cached metadata"); + }; + assert!( + error.chain().any(|cause| cause + .to_string() + .contains("cannot enforce split filesystem read restrictions")), + "{error:#}" + ); + let mut sandbox = FileSystemSandboxContext::from_legacy_sandbox_policy( + codex_protocol::protocol::SandboxPolicy::new_read_only_policy(), + cwd.clone(), + )?; + sandbox.windows_sandbox_selection = + codex_exec_server::WindowsSandboxSelection::RestrictedToken; + let permitted = DiscoverV2CapabilitiesRequest { + cwd: cwd.clone(), + sandbox: Some(sandbox), + }; + let response = discover(&mut server, &permitted).await?; + assert!( + response + .plugins + .iter() + .any(|plugin| plugin.id == "mail@market") + ); + assert!(response.skills.iter().any(|skill| skill.name == "demo")); + } + } + + server.shutdown().await?; + Ok(()) +} + +async fn discover( + server: &mut ExecServerHarness, + params: &DiscoverV2CapabilitiesRequest, +) -> anyhow::Result { + Ok(serde_json::from_value( + request( + server, + CAPABILITIES_DISCOVER_V2_METHOD, + serde_json::to_value(params)?, + ) + .await + .with_context(|| format!("discover V2 for {}", params.cwd))?, + )?) +} + +async fn request( + server: &mut ExecServerHarness, + method: &str, + params: serde_json::Value, +) -> anyhow::Result { + let id = server.send_request(method, params).await?; + let message = server + .wait_for_event(|event| { + matches!(event, + JSONRPCMessage::Response(response) if response.id == id + ) || matches!(event, JSONRPCMessage::Error(error) if error.id == id) + }) + .await + .with_context(|| format!("waiting for {method}"))?; + match message { + JSONRPCMessage::Response(response) => Ok(response.result), + other => anyhow::bail!("{method} failed: {other:?}"), + } +} + +fn write(path: &std::path::Path, contents: &str) -> std::io::Result<()> { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + std::fs::write(path, contents) +} diff --git a/codex-rs/exec-server/tests/initialize.rs b/codex-rs/exec-server/tests/initialize.rs index 53bb19d886..c9510011ac 100644 --- a/codex-rs/exec-server/tests/initialize.rs +++ b/codex-rs/exec-server/tests/initialize.rs @@ -3,6 +3,9 @@ mod common; use anyhow::Context; use codex_build_info::BuildInfo; use codex_build_info::build_id; +use codex_exec_server::CAPABILITIES_DISCOVER_V2_METHOD; +use codex_exec_server::DiscoverV2CapabilitiesRequest; +use codex_exec_server::DiscoverV2CapabilitiesResponse; use codex_exec_server::EnvironmentInfo; use codex_exec_server::InitializeParams; use codex_exec_server::InitializeResponse; @@ -102,7 +105,8 @@ async fn exec_server_accepts_initialize(version: Option<&str>) -> anyhow::Result } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn discovery_support_does_not_require_client_opt_in() -> anyhow::Result<()> { +async fn discovery_prewarms_before_initialize_and_recovers_without_client_opt_in() +-> anyhow::Result<()> { let codex_home = TempDir::new()?; #[cfg(not(windows))] let discovery_home = { @@ -121,11 +125,12 @@ async fn discovery_support_does_not_require_client_opt_in() -> anyhow::Result<() let helper_paths = common::exec_server::test_codex_helper_paths()?; let mut command = Command::new(helper_paths.codex_exe); command.args(["exec-server", "--listen", "ws://127.0.0.1:0"]); - command.env("CODEX_HOME", discovery_home); + command.env("CODEX_HOME", &discovery_home); command.env("HOME", &user_home); command.env("USERPROFILE", &user_home); command.env("RUST_LOG", "codex_exec_server=warn"); let mut server = ExecServerHarness::start_with_stderr(command, stderr.into()).await?; + // No initialize or discovery request has been sent: the scan must run at startup. let observed_failure = tokio::time::timeout(std::time::Duration::from_secs(5), async { loop { let logs = tokio::fs::read_to_string(&stderr_path).await?; @@ -172,10 +177,39 @@ async fn discovery_support_does_not_require_client_opt_in() -> anyhow::Result<() let environment_info_id = server .send_request("environment/info", serde_json::json!({})) .await?; - let JSONRPCMessage::Response(JSONRPCResponse { id, .. }) = server.next_event().await? else { + let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = server.next_event().await? + else { panic!("expected environment info response after degraded discovery startup"); }; assert_eq!(id, environment_info_id); + let info: EnvironmentInfo = serde_json::from_value(result)?; + // Repair the failed startup fixture; discovery must retry and find the new skill. + #[cfg(not(windows))] + std::fs::remove_file(discovery_home.join("plugins"))?; + let skill_root = discovery_home.join("skills/recovered"); + std::fs::create_dir_all(&skill_root)?; + std::fs::write( + skill_root.join("SKILL.md"), + "---\nname: recovered\ndescription: Found after startup recovery\n---\n", + )?; + let discovery_id = server + .send_request( + CAPABILITIES_DISCOVER_V2_METHOD, + serde_json::to_value(DiscoverV2CapabilitiesRequest { + cwd: info.cwd.context("startup cwd")?, + sandbox: None, + })?, + ) + .await?; + let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = server.next_event().await? + else { + anyhow::bail!("expected discovery to recover after repairing the startup fixture"); + }; + assert_eq!(id, discovery_id); + let discovery: DiscoverV2CapabilitiesResponse = serde_json::from_value(result)?; + assert!(discovery.skills.iter().any(|skill| { + skill.name == "recovered" && skill.description == "Found after startup recovery" + })); server.shutdown().await?; Ok(()) diff --git a/codex-rs/file-watcher/src/file_watcher_tests.rs b/codex-rs/file-watcher/src/file_watcher_tests.rs index 8bd7858c21..ecf54cabd1 100644 --- a/codex-rs/file-watcher/src/file_watcher_tests.rs +++ b/codex-rs/file-watcher/src/file_watcher_tests.rs @@ -591,3 +591,104 @@ async fn dropping_live_watcher_releases_inner_watcher() { assert_eq!(weak_inner.upgrade().is_none(), true); } + +#[tokio::test] +async fn registration_reports_success_for_new_and_shared_watches() +-> Result<(), Box> { + let directory = tempfile::tempdir()?; + let watcher = Arc::new(FileWatcher::new()?); + let (subscriber, _rx) = watcher.add_subscriber(); + let (other_subscriber, _other_rx) = watcher.add_subscriber(); + // Missing directories succeed through a non-recursive watch on their existing parent. + for path in [ + directory.path().to_path_buf(), + directory.path().join("missing"), + ] { + let first = subscriber.register_path(path.clone(), /*recursive*/ true); + let shared = other_subscriber.register_path(path.clone(), /*recursive*/ false); + assert!(first.registration_succeeded(), "new watch for {path:?}"); + assert!(shared.registration_succeeded(), "shared watch for {path:?}"); + drop(first); + assert_eq!( + watcher.watch_counts_for_test(directory.path()), + Some((1, 0)) + ); + drop(shared); + assert_eq!(watcher.watch_counts_for_test(directory.path()), None); + } + Ok(()) +} + +#[tokio::test] +async fn registration_reports_missing_os_watches() -> Result<(), Box> { + let directory = tempfile::tempdir()?; + let root = directory.path().canonicalize()?; + let missing = root.join("a-missing"); + let watched = root.join("b-watched"); + std::fs::create_dir(&watched)?; + let watcher = Arc::new(FileWatcher::new()?); + let (subscriber, _rx) = watcher.add_subscriber(); + let paths = vec![ + WatchPath { + path: missing.clone(), + recursive: true, + }, + WatchPath { + path: watched.clone(), + recursive: true, + }, + ]; + // Model a path disappearing after resolution but before the OS registration. + let registrations = paths + .iter() + .map(|watch| SubscriberWatchRegistration { + key: SubscriberWatchKey { + requested: watch.clone(), + matched: watch.clone(), + }, + actual: watch.clone(), + fallback: false, + }) + .collect::>(); + assert!(!watcher.register_paths(subscriber.id, ®istrations)); + + std::fs::create_dir(&missing)?; + // The directory now exists, but the unchanged mode does not retry its OS watch. + let repeated = subscriber.register_paths(paths); + assert!(!repeated.registration_succeeded()); + // The other path registered successfully despite the first path's failure. + let successful = subscriber.register_path(watched.clone(), /*recursive*/ true); + assert!(successful.registration_succeeded()); + { + let inner = watcher + .inner + .as_ref() + .ok_or("missing watcher backend")? + .lock() + .map_err(|_| "poisoned watcher lock")?; + assert_eq!( + inner.watched_paths, + HashMap::from([(watched.clone(), RecursiveMode::Recursive)]) + ); + } + drop(successful); + drop(repeated); + // Release the original registrations made directly through the subscriber ID. + drop(subscriber); + assert_eq!(watcher.watch_counts_for_test(&missing), None); + assert_eq!(watcher.watch_counts_for_test(&watched), None); + Ok(()) +} + +#[test] +fn noop_registration_reports_failure() -> Result<(), Box> { + let directory = tempfile::tempdir()?; + let watcher = Arc::new(FileWatcher::noop()); + let (subscriber, _rx) = watcher.add_subscriber(); + let registration = + subscriber.register_path(directory.path().to_path_buf(), /*recursive*/ true); + assert!(!registration.registration_succeeded()); + drop(registration); + assert_eq!(watcher.watch_counts_for_test(directory.path()), None); + Ok(()) +} diff --git a/codex-rs/file-watcher/src/lib.rs b/codex-rs/file-watcher/src/lib.rs index a0b1a42401..1e2e22f08d 100644 --- a/codex-rs/file-watcher/src/lib.rs +++ b/codex-rs/file-watcher/src/lib.rs @@ -318,9 +318,10 @@ impl FileWatcherSubscriber { } }) .collect::>(); - self.file_watcher.register_paths(self.id, &watched_paths); + let registration_succeeded = self.file_watcher.register_paths(self.id, &watched_paths); WatchRegistration { + registration_succeeded, file_watcher: Arc::downgrade(&self.file_watcher), subscriber_id: self.id, watched_paths: watched_paths @@ -344,14 +345,23 @@ impl Drop for FileWatcherSubscriber { /// RAII guard for a set of active path registrations. pub struct WatchRegistration { + registration_succeeded: bool, file_watcher: std::sync::Weak, subscriber_id: SubscriberId, watched_paths: Vec, } +impl WatchRegistration { + /// Whether all required OS watches registered successfully at creation time. + pub fn registration_succeeded(&self) -> bool { + self.registration_succeeded + } +} + impl Default for WatchRegistration { fn default() -> Self { Self { + registration_succeeded: false, file_watcher: std::sync::Weak::new(), subscriber_id: 0, watched_paths: Vec::new(), @@ -433,7 +443,7 @@ impl FileWatcher { &self, subscriber_id: SubscriberId, watched_paths: &[SubscriberWatchRegistration], - ) { + ) -> bool { let mut state = self .state .write() @@ -443,7 +453,7 @@ impl FileWatcher { for registration in watched_paths { let actual = { let Some(subscriber) = state.subscribers.get_mut(&subscriber_id) else { - return; + return false; }; match subscriber.watched_paths.entry(registration.key.clone()) { std::collections::hash_map::Entry::Occupied(mut entry) => { @@ -473,6 +483,29 @@ impl FileWatcher { self.reconfigure_watch(&actual.path, next_mode, &mut inner_guard); } } + + // Registration remains best-effort. Check that every requested path has an OS watch. + let Some(inner) = &self.inner else { + return false; + }; + let guard = inner_guard.get_or_insert_with(|| { + inner + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + }); + let Some(subscriber) = state.subscribers.get(&subscriber_id) else { + return false; + }; + watched_paths.iter().all(|registration| { + let Some(watch) = subscriber.watched_paths.get(®istration.key) else { + return false; + }; + match guard.watched_paths.get(&watch.actual.path) { + Some(RecursiveMode::Recursive) => true, + Some(RecursiveMode::NonRecursive) => !watch.actual.recursive, + None => false, + } + }) } fn unregister_paths(&self, subscriber_id: SubscriberId, watched_paths: &[SubscriberWatchKey]) {