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
This commit is contained in:
Anthony Tafoya
2026-10-13 17:22:34 +00:00
committed by copyberry
parent 6cc7d9cc5c
commit 703d253ed6
22 changed files with 802 additions and 274 deletions
+10
View File
@@ -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, &params).await
}
pub async fn discover_capabilities_v2(
&self,
params: DiscoverV2CapabilitiesRequest,
) -> Result<DiscoverV2CapabilitiesResponse, ExecServerError> {
self.call(CAPABILITIES_DISCOVER_V2_METHOD, &params).await
}
pub async fn read(&self, params: ReadParams) -> Result<ReadResponse, ExecServerError> {
self.call(EXEC_READ_METHOD, &params).await
}
@@ -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 {
@@ -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<Vec<Arc<CachedSandboxDiscovery>>>,
_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<Self>,
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<Self>,
location_request: CapabilityLocationRequest,
params: DiscoverV2CapabilitiesRequest,
) -> io::Result<Arc<DiscoverV2CapabilitiesResponse>> {
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<Arc<LocationState>> {
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(),
});
@@ -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);
@@ -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;
@@ -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
@@ -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<CapabilityManager>,
)> {
impl CapabilityManager {
// Supply fixture paths through the existing cell, without production override state.
pub(crate) async fn seed_prewarmed_locations(
self: &Arc<Self>,
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<CapabilityManager>)> {
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(())
})
@@ -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(())
+5 -32
View File
@@ -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<CapabilityRootsDiscoverResponse, ExecServerError> {
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
@@ -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<T, F, Fut>(
&self,
mut connection_state: watch::Receiver<EnvironmentConnectionState>,
discover: F,
) -> Result<T, ExecServerError>
where
F: Fn() -> Fut,
Fut: Future<Output = Result<T, ExecServerError>>,
{
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,
}
}
}
@@ -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<DiscoverV2CapabilitiesResponse, ExecServerError> {
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;
@@ -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(&params)?,
serde_json::to_value(&params)?
]
);
Ok(())
}
@@ -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<Arc<ExecutorRegistration>>,
session_registry: Arc<SessionRegistry>,
capability_manager: Arc<CapabilityManager>,
notifications: RpcNotificationSender,
session: StdMutex<Option<SessionHandle>>,
active_body_stream_ids: Mutex<HashSet<String>>,
@@ -94,6 +98,7 @@ pub(crate) struct ExecServerHandler {
impl ExecServerHandler {
pub(crate) fn new(
session_registry: Arc<SessionRegistry>,
capability_manager: Arc<CapabilityManager>,
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<Arc<DiscoverV2CapabilitiesResponse>, 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,
@@ -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<ExecServerHandler> {
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(&registry),
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(&registry),
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(),
@@ -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<ExecServerHandler> {
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"),
+23 -108
View File
@@ -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<CapabilityLocationRequest>,
) -> 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<Arc<ExecutorRegistration>>,
) {
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]
@@ -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<ExecServerHandler> {
handler.discover_capability_roots(params).await
},
);
router.request(
CAPABILITIES_DISCOVER_V2_METHOD,
|handler: Arc<ExecServerHandler>, params: DiscoverV2CapabilitiesRequest| async move {
handler.discover_capabilities_v2(params).await
},
);
router.request(
EXEC_READ_METHOD,
|handler: Arc<ExecServerHandler>, params: ReadParams| async move {
@@ -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<ExecServerHandler>) -> 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"),
@@ -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, &params).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, &params)
.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<DiscoverV2CapabilitiesResponse> {
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<serde_json::Value> {
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)
}
+37 -3
View File
@@ -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(())
@@ -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<dyn std::error::Error>> {
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<dyn std::error::Error>> {
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::<Vec<_>>();
assert!(!watcher.register_paths(subscriber.id, &registrations));
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<dyn std::error::Error>> {
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(())
}
+36 -3
View File
@@ -318,9 +318,10 @@ impl FileWatcherSubscriber {
}
})
.collect::<Vec<_>>();
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<FileWatcher>,
subscriber_id: SubscriberId,
watched_paths: Vec<SubscriberWatchKey>,
}
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(&registration.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]) {