diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 15f61c4f07..89e75623fa 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -422,11 +422,14 @@ impl Session { }) .cloned() .collect::>(); - Some(Arc::new( - cache - .snapshot(&selected_capability_roots, &sandbox_contexts) - .await, - )) + let discovery = cache + .snapshot(&selected_capability_roots, &sandbox_contexts) + .await; + if cache.take_recovered_discovery() { + // Root selection is unchanged, but recovered manifests can change MCP servers. + self.mark_mcp_runtime_dirty(); + } + Some(Arc::new(discovery)) } pub(crate) async fn resolve_selected_capability_roots_for_step( diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 4caa44158c..6a666e0094 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3143,7 +3143,8 @@ impl Session { &environments, turn_context.windows_sandbox_level, ) - .await; + .or_cancel(cancellation_token) + .await?; let extension_data = codex_extension_api::ExtensionData::new(turn_context.sub_id.clone()); extension_data.insert(selected_capability_roots.clone()); if let Some(discovery) = &executor_capability_discovery { diff --git a/codex-rs/exec-server/src/capability_discovery_cache.rs b/codex-rs/exec-server/src/capability_discovery_cache.rs index b1e9fb5408..33bbcb8c1f 100644 --- a/codex-rs/exec-server/src/capability_discovery_cache.rs +++ b/codex-rs/exec-server/src/capability_discovery_cache.rs @@ -1,6 +1,8 @@ use std::collections::BTreeMap; use std::collections::HashMap; use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; @@ -15,12 +17,13 @@ use crate::FileSystemSandboxContext; /// Thread-scoped cache shared by capability consumers using the high-level executor API. /// -/// A single miss batches every requested root by environment. The cache deliberately has no -/// invalidation: selected roots are already treated as stable for the lifetime of a thread by the -/// existing plugin and skill providers. +/// A single miss batches every requested root by environment. Successful discoveries and +/// permanent failures remain cached by root and sandbox; transient failures are retried on the +/// next request. Recovery is reported so dependent MCP projections can be invalidated. pub struct ExecutorCapabilityDiscoveryCache { environment_manager: Arc, entries: Mutex>, + recovered_discovery: AtomicBool, } impl std::fmt::Debug for ExecutorCapabilityDiscoveryCache { @@ -34,9 +37,9 @@ impl std::fmt::Debug for ExecutorCapabilityDiscoveryCache { struct CachedRoot { selected_root: SelectedCapabilityRoot, sandbox: Option, - // Both successes and failures are memoized for the thread. Retrying a transient failure for - // the same stable selected root requires explicit invalidation or a new thread. result: Result, String>, + // Preserve transport classification after the public snapshot reduces errors to strings. + retryable: bool, } impl ExecutorCapabilityDiscoveryCache { @@ -44,9 +47,15 @@ impl ExecutorCapabilityDiscoveryCache { Self { environment_manager, entries: Mutex::new(Vec::new()), + recovered_discovery: AtomicBool::new(false), } } + /// Reports whether a previously failed root has recovered since the last observation. + pub fn take_recovered_discovery(&self) -> bool { + self.recovered_discovery.swap(false, Ordering::AcqRel) + } + /// Returns discoveries in the same order as `selected_roots`. #[tracing::instrument( name = "capability_roots.discovery_cache.resolve", @@ -69,6 +78,7 @@ impl ExecutorCapabilityDiscoveryCache { !entries.iter().any(|cached| { cached.selected_root == **selected_root && cached.sandbox.as_ref() == sandbox + && !cached.retryable }) }) .cloned() @@ -81,7 +91,10 @@ impl ExecutorCapabilityDiscoveryCache { .iter_mut() .find(|cached| cached.selected_root == discovered_root.selected_root) { - if cached.sandbox != discovered_root.sandbox { + if cached.sandbox != discovered_root.sandbox || cached.result.is_err() { + if cached.result.is_err() && discovered_root.result.is_ok() { + self.recovered_discovery.store(true, Ordering::Release); + } *cached = discovered_root; } } else { @@ -153,6 +166,7 @@ impl ExecutorCapabilityDiscoveryCache { selected_root, sandbox: sandbox.clone(), result: Err(error.clone()), + retryable: true, }) .collect::>(); }; @@ -173,6 +187,7 @@ impl ExecutorCapabilityDiscoveryCache { let response = match environment.discover_capability_roots(params).await { Ok(response) => response, Err(error) => { + let retryable = crate::client::is_retryable_recovery_error(&error); let error = error.to_string(); return selected_roots .into_iter() @@ -180,6 +195,7 @@ impl ExecutorCapabilityDiscoveryCache { selected_root, sandbox: sandbox.clone(), result: Err(error.clone()), + retryable, }) .collect(); } @@ -196,6 +212,7 @@ impl ExecutorCapabilityDiscoveryCache { selected_root, sandbox: sandbox.clone(), result: Err(error.clone()), + retryable: false, }) .collect(); } @@ -218,6 +235,7 @@ impl ExecutorCapabilityDiscoveryCache { selected_root, sandbox: sandbox.clone(), result, + retryable: false, } }) .collect() diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 1ebde3302f..bee9a2a56d 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -125,6 +125,7 @@ use codex_http_client::HttpClientFactory; pub(crate) mod http_client; #[path = "client_recovery.rs"] mod recovery; +pub(crate) use recovery::is_retryable_recovery_error; const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); const INITIALIZE_TIMEOUT: Duration = Duration::from_secs(10); diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index 6588de2109..606a81e3f6 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/codex-rs/exec-server/src/client_recovery.rs @@ -770,7 +770,10 @@ impl ExecServerClient { } } -pub(super) fn is_retryable_recovery_error(error: &ExecServerError) -> bool { +pub(crate) fn is_retryable_recovery_error(error: &ExecServerError) -> bool { + if let ExecServerError::ConnectionAttempt(error) = error { + return is_retryable_recovery_error(error.as_ref()); + } is_transport_closed_error(error) || matches!( error, diff --git a/codex-rs/exec-server/src/client_recovery_tests.rs b/codex-rs/exec-server/src/client_recovery_tests.rs index 3dfcbe0ffa..545154728b 100644 --- a/codex-rs/exec-server/src/client_recovery_tests.rs +++ b/codex-rs/exec-server/src/client_recovery_tests.rs @@ -41,6 +41,9 @@ fn recovery_retries_transient_registry_errors() { assert!(is_retryable_registry_error(&error)); assert!(is_retryable_recovery_error(&error)); + assert!(is_retryable_recovery_error( + &ExecServerError::ConnectionAttempt(Arc::new(error)) + )); } #[test] @@ -57,6 +60,9 @@ fn recovery_does_not_retry_other_registry_conflicts() { assert!(!is_retryable_registry_error(&error)); assert!(!is_retryable_recovery_error(&error)); + assert!(!is_retryable_recovery_error( + &ExecServerError::ConnectionAttempt(Arc::new(error)) + )); } #[test] diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 7616692486..df631145c1 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -926,22 +926,48 @@ impl Environment { ) -> Result { match &self.remote_client { Some(client) => { + let mut connection_state = client.subscribe_connection_state(); let client = client.get().await?; - if params.roots.iter().any(|root| { - root.sandbox - .as_ref() - .is_some_and(crate::FileSystemSandboxContext::should_run_in_sandbox) - }) && !client - .environment_info() - .await? - .capabilities - .capability_discovery_sandbox - { - return Err(ExecServerError::Protocol( - "exec-server does not support sandboxed capability discovery".to_string(), - )); + let discover = || async { + if params.roots.iter().any(|root| { + root.sandbox + .as_ref() + .is_some_and(crate::FileSystemSandboxContext::should_run_in_sandbox) + }) && !client + .environment_info() + .await? + .capabilities + .capability_discovery_sandbox + { + return Err(ExecServerError::Protocol( + "exec-server does not support sandboxed capability discovery" + .to_string(), + )); + } + 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 { + while self.readiness_result().is_none_or(|result| result.is_err()) { + if connection_state.changed().await.is_err() { + return false; + } + } + true + }) + .await + .unwrap_or(false); + if recovered { + discover().await + } else { + Err(error) + } + } + response => response, } - client.discover_capability_roots(params).await } None => crate::discover_capability_roots(self.filesystem.as_ref(), params) .await diff --git a/codex-rs/exec-server/tests/environment.rs b/codex-rs/exec-server/tests/environment.rs index 01cb6a16df..814ae7f2d2 100644 --- a/codex-rs/exec-server/tests/environment.rs +++ b/codex-rs/exec-server/tests/environment.rs @@ -1,11 +1,17 @@ mod common; +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; use anyhow::Context; use codex_exec_server::EnvironmentManager; +use codex_exec_server::ExecutorCapabilityDiscoveryCache; use codex_exec_server::REMOTE_ENVIRONMENT_ID; use codex_exec_server::SelectedCapabilityRootsStatus; +use codex_exec_server_protocol::CAPABILITY_ROOTS_DISCOVER_METHOD; use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; use codex_http_client::cache_system_proxy_route_for_test; @@ -13,6 +19,8 @@ use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_utils_path_uri::PathUri; use common::exec_server::exec_server; +use futures::SinkExt; +use futures::StreamExt; use pretty_assertions::assert_eq; use tokio::io::AsyncReadExt; use tokio::io::AsyncWriteExt; @@ -21,6 +29,9 @@ use tokio::net::TcpStream; use tokio::sync::oneshot; use tokio::time::sleep; use tokio::time::timeout; +use tokio_tungstenite::accept_async; +use tokio_tungstenite::connect_async; +use tokio_tungstenite::tungstenite::Message; use tokio_util::task::AbortOnDropHandle; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -164,3 +175,141 @@ async fn selected_capability_inspection_tracks_connection_recovery() -> anyhow:: Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(remote_exec_server)] +async fn capability_discovery_retries_executor_disconnect_within_same_request() -> anyhow::Result<()> +{ + let server = exec_server().await?; + let proxy_listener = TcpListener::bind("127.0.0.1:0").await?; + let proxy_websocket_url = format!("ws://{}", proxy_listener.local_addr()?); + let upstream_websocket_url = server.websocket_url().to_string(); + let discovery_attempts = Arc::new(AtomicUsize::new(0)); + let proxy_discovery_attempts = Arc::clone(&discovery_attempts); + let _proxy_task = AbortOnDropHandle::new(tokio::spawn(async move { + while let Ok((downstream, _)) = proxy_listener.accept().await { + let mut downstream = accept_async(downstream).await?; + let (mut upstream, _) = connect_async(&upstream_websocket_url).await?; + + loop { + tokio::select! { + message = downstream.next() => { + let Some(message) = message.transpose()? else { + break; + }; + if let Message::Text(message_text) = &message { + let request = serde_json::from_str::(message_text.as_ref())?; + if request.get("method").and_then(serde_json::Value::as_str) + == Some(CAPABILITY_ROOTS_DISCOVER_METHOD) + { + let attempt = proxy_discovery_attempts.fetch_add(1, Ordering::SeqCst); + if attempt == 0 { + break; + } + sleep(Duration::from_secs(9)).await; + } + } + upstream.send(message).await?; + } + message = upstream.next() => { + let Some(message) = message.transpose()? else { + break; + }; + downstream.send(message).await?; + } + } + } + } + Ok::<(), anyhow::Error>(()) + })); + let manager = Arc::new( + EnvironmentManager::create_for_tests( + Some(proxy_websocket_url), + /*local_runtime_paths*/ None, + ) + .await, + ); + manager + .default_environment() + .context("remote environment")? + .info() + .await?; + + let cache = Arc::new(ExecutorCapabilityDiscoveryCache::new(Arc::clone(&manager))); + let skill_root = tempfile::tempdir()?; + let selected_roots = vec![SelectedCapabilityRoot { + id: "recovering-skill".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + path: PathUri::from_host_native_path(skill_root.path())?, + }, + }]; + + let snapshot = timeout( + Duration::from_secs(12), + cache.snapshot(&selected_roots, &HashMap::new()), + ) + .await + .context("capability discovery did not retry within the same request")?; + let discovery = snapshot.roots()[0] + .result + .as_ref() + .map_err(|error| anyhow::anyhow!("{error}"))?; + + assert_eq!(discovery.id, "recovering-skill"); + assert_eq!( + 2, + discovery_attempts.load(Ordering::SeqCst), + "same-request retry must issue a second capability discovery RPC" + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn capability_discovery_retries_after_executor_reconnects() -> anyhow::Result<()> { + let server = exec_server().await?; + let manager = Arc::new(EnvironmentManager::default_for_tests()); + let cache = ExecutorCapabilityDiscoveryCache::new(Arc::clone(&manager)); + let skill_root = tempfile::tempdir()?; + let refused_listener = TcpListener::bind("127.0.0.1:0").await?; + let refused_address = refused_listener.local_addr()?; + drop(refused_listener); + manager.upsert_environment( + "recovering".to_string(), + format!("ws://{refused_address}"), + Some(Duration::from_millis(100)), + )?; + let selected_roots = vec![SelectedCapabilityRoot { + id: "recovering-skill".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "recovering".to_string(), + path: PathUri::from_host_native_path(skill_root.path())?, + }, + }]; + + let failed_snapshot = cache.snapshot(&selected_roots, &HashMap::new()).await; + assert!(failed_snapshot.roots()[0].result.is_err()); + assert!(!cache.take_recovered_discovery()); + + manager.upsert_environment( + "recovering".to_string(), + server.websocket_url().to_string(), + /*connect_timeout*/ None, + )?; + manager + .get_environment("recovering") + .context("recovered environment")? + .wait_until_ready() + .await?; + + let recovered_snapshot = cache.snapshot(&selected_roots, &HashMap::new()).await; + let discovery = recovered_snapshot.roots()[0] + .result + .as_ref() + .map_err(|error| anyhow::anyhow!("{error}"))?; + + assert_eq!(discovery.id, "recovering-skill"); + assert!(cache.take_recovered_discovery()); + assert!(!cache.take_recovered_discovery()); + Ok(()) +} diff --git a/codex-rs/ext/skills/src/state.rs b/codex-rs/ext/skills/src/state.rs index c3f0589a26..60f0148263 100644 --- a/codex-rs/ext/skills/src/state.rs +++ b/codex-rs/ext/skills/src/state.rs @@ -103,10 +103,9 @@ impl SkillsThreadState { /// Returns catalogs for stable selected roots. /// - /// The first catalog returned for a root remains cached until this thread state is dropped. - /// Environment availability only controls whether the root is projected into the current - /// step; it never invalidates the cache. There is intentionally no filesystem watcher or - /// content-based invalidation because selected environment roots are treated as stable. + /// Successful catalogs, including empty or warning-bearing catalogs, remain cached until + /// this thread state is dropped. Catalogs backed by failed discovery are not cached, so + /// later steps can recover. There is no filesystem watcher because selected roots are stable. #[tracing::instrument( name = "skills.executor.catalog_snapshot", level = "info", @@ -140,6 +139,10 @@ impl SkillsThreadState { providers: &SkillProviders, query: SkillListQuery, ) -> SkillCatalog { + let discovery_failed = query + .executor_capability_discovery + .as_ref() + .is_some_and(|discovery| discovery.roots().iter().any(|root| root.result.is_err())); let sandbox_contexts = query .executor_capability_discovery .as_ref() @@ -158,6 +161,9 @@ impl SkillsThreadState { } let roots = query.executor_roots.clone(); let discovered = providers.list_executor_for_turn(query).await; + if discovery_failed { + return discovered; + } let mut cache = self .executor_discovery_cache .lock() diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index f8c320c57a..9200307751 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -8,6 +8,8 @@ use codex_config::ConfigLayerEntry; use codex_config::ConfigLayerSource; use codex_config::ConfigLayerStack; use codex_config::ConfigRequirementsToml; +use codex_exec_server::CapabilityRootDiscovery; +use codex_exec_server::ExecutorCapabilityDiscoverySnapshot; use codex_exec_server::LOCAL_FS; use codex_extension_api::ConversationHistory; use codex_extension_api::ExtensionData; @@ -1125,6 +1127,45 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach assert!(restored_fragment.body().contains("lint-fix")); assert_eq!(1, list_calls.load(Ordering::Relaxed)); + let failed_discovery = ExecutorCapabilityDiscoverySnapshot::new( + &selected_roots, + vec![Err("exec-server transport disconnected".to_string())], + Default::default(), + ); + let recovered_discovery = ExecutorCapabilityDiscoverySnapshot::new( + &selected_roots, + vec![Ok(Arc::new(CapabilityRootDiscovery { + id: "lint-fix".to_string(), + path: PathUri::parse("file:///skills/lint-fix")?, + plugin: None, + skills: Vec::new(), + namespace_manifests: Vec::new(), + warnings: Vec::new(), + error: None, + }))], + Default::default(), + ); + for (turn_id, discovery, expected_list_calls) in [ + ("failed-discovery", &failed_discovery, 2), + ("recovered-discovery", &recovered_discovery, 3), + ("cached-discovery", &recovered_discovery, 3), + ] { + registry.context_contributors()[0] + .contribute_world_state(WorldStateContributionInput { + thread_id: codex_protocol::ThreadId::new(), + turn_id, + environments: &[], + ready_selected_capability_roots: &selected_roots, + executor_capability_discovery: Some(discovery), + extension_metrics: None, + session_store: &session_store, + thread_store: &thread_store, + turn_store: &ExtensionData::new(turn_id), + }) + .await; + assert_eq!(expected_list_calls, list_calls.load(Ordering::Relaxed)); + } + let mut listing_disabled_config = config.clone(); listing_disabled_config.include_instructions = false; registry.config_contributors()[0].on_config_changed(