Recover capability discovery after executor disconnects (#38420)

## Why

Transient executor disconnects could leave capability discovery and skill catalogs
stuck on a cached failure for the rest of a thread, even after the executor
reconnected.

## What changed

- Replay capability discovery after executor recovery and retry transient failures
  on later requests while continuing to cache permanent failures.
- Avoid caching skill catalogs produced from failed discovery so a later step can
  load the recovered catalog.
- Mark the MCP runtime dirty when recovered manifests change the projected MCP
  servers, and allow discovery to be cancelled with the turn.

## Testing

- Cover same-request recovery after a disconnect and recovery on a later request.
- Cover retry classification through connection-attempt errors and skill catalog
  caching after discovery recovers.

GitOrigin-RevId: a57f90844351e73ea831931f72a9ddc4e4f3335c

codex-rs/core/src/session/mcp.rs

@@ -422,11 +422,14 @@ impl Session {
             })
             .cloned()
             .collect::<Vec<_>>();
-        Some(Arc::new(
+        let discovery = cache
-            cache
             .snapshot(&selected_capability_roots, &sandbox_contexts)
-                .await,
+            .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(

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 {

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
+/// A single miss batches every requested root by environment. Successful discoveries and
-/// invalidation: selected roots are already treated as stable for the lifetime of a thread by the
+/// permanent failures remain cached by root and sandbox; transient failures are retried on the
-/// existing plugin and skill providers.
+/// next request. Recovery is reported so dependent MCP projections can be invalidated.
 pub struct ExecutorCapabilityDiscoveryCache {
     environment_manager: Arc<EnvironmentManager>,
     entries: Mutex<Vec<CachedRoot>>,
+    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<FileSystemSandboxContext>,
-    // 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<Arc<CapabilityRootDiscovery>, 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::<Vec<_>>();
                 };
@@ -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()

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);

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,

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]

codex-rs/exec-server/src/environment.rs

@@ -926,7 +926,9 @@ impl Environment {
     ) -> Result<CapabilityRootsDiscoverResponse, ExecServerError> {
         match &self.remote_client {
             Some(client) => {
+                let mut connection_state = client.subscribe_connection_state();
                 let client = client.get().await?;
+                let discover = || async {
                     if params.roots.iter().any(|root| {
                         root.sandbox
                             .as_ref()
@@ -938,10 +940,34 @@ impl Environment {
                         .capability_discovery_sandbox
                     {
                         return Err(ExecServerError::Protocol(
-                        "exec-server does not support sandboxed capability discovery".to_string(),
+                            "exec-server does not support sandboxed capability discovery"
+                                .to_string(),
                         ));
                     }
-                client.discover_capability_roots(params).await
+                    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,
+                }
             }
             None => crate::discover_capability_roots(self.filesystem.as_ref(), params)
                 .await

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::<serde_json::Value>(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(())
+}

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.
+    /// Successful catalogs, including empty or warning-bearing catalogs, remain cached until
-    /// Environment availability only controls whether the root is projected into the current
+    /// this thread state is dropped. Catalogs backed by failed discovery are not cached, so
-    /// step; it never invalidates the cache. There is intentionally no filesystem watcher or
+    /// later steps can recover. There is no filesystem watcher because selected roots are stable.
-    /// content-based invalidation because selected environment roots are treated as 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()

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(