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/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<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,22 +926,48 @@ impl Environment {
     ) -> Result<CapabilityRootsDiscoverResponse, ExecServerError> {
         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

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