Wire V2 capability discovery through exec-server RPC (#53723)

## What changed

- Register `capabilities/discoverV2` for initialized sessions and expose it through `ExecServerClient` and remote `Environment` APIs, checking executor support before calling it.
- Reuse startup discovery locations and retry failed prewarming on demand. Replay discovery once after executor recovery while preserving the request parameters.
- Track whether all required filesystem watches registered successfully and reload locations and metadata on each request when registration fails.

## Testing

Add coverage for RPC metadata refresh, omission of skill instruction bodies, sandbox isolation, startup scan recovery, disconnect replay, and successful or failed watch registration.

GitOrigin-RevId: 726e8f7e8b4d835b877f31d1d8a9f8eea1ad9c8b

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

@@ -13,10 +13,13 @@ use tokio::sync::Mutex;
 use tokio_util::sync::CancellationToken;
 use tokio_util::task::TaskTracker;
 
+use crate::DiscoverV2CapabilitiesRequest;
+use crate::DiscoverV2CapabilitiesResponse;
 use crate::ExecServerRuntimeOptions;
 use crate::client::http_client::PendingRouteAwareHttpBodyStream;
 use crate::client::http_client::RouteAwareHttpClient;
 use crate::client::http_client::RouteAwareHttpRequestRunner;
+use crate::discover_v2::capability_manager::CapabilityManager;
 use crate::environment_config::ReadEnvironmentConfigError;
 use crate::environment_config::read_environment_config;
 use crate::protocol::CapabilityRootsDiscoverParams;
@@ -79,6 +82,7 @@ use crate::telemetry::ExecutorRegistration;
 pub(crate) struct ExecServerHandler {
     pub(super) executor_registration: Option<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,

codex-rs/exec-server/src/server/handler/tests.rs

@@ -12,6 +12,7 @@ use uuid::Uuid;
 use super::ExecServerHandler;
 use crate::ExecServerRuntimeOptions;
 use crate::ProcessId;
+use crate::discover_v2::capability_manager::CapabilityManager;
 use crate::protocol::ExecParams;
 use crate::protocol::InitializeParams;
 use crate::protocol::ReadParams;
@@ -94,6 +95,7 @@ async fn initialized_handler() -> Arc<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(),

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

@@ -40,6 +40,7 @@ use super::session_registry::SessionRegistry;
 use crate::ExecServerRuntimeOptions;
 use crate::ExecServerTelemetry;
 use crate::connection::JsonRpcConnectionEvent;
+use crate::discover_v2::capability_manager::CapabilityManager;
 use crate::protocol::EXEC_METHOD;
 use crate::protocol::EXEC_READ_METHOD;
 use crate::protocol::EXEC_TERMINATE_METHOD;
@@ -738,6 +739,7 @@ fn new_handler(
 ) -> Arc<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"),

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

@@ -1,4 +1,3 @@
-use std::io;
 use std::sync::Arc;
 use std::time::Instant;
 
@@ -14,7 +13,6 @@ use crate::LocalFileSystem;
 use crate::connection::CHANNEL_CAPACITY;
 use crate::connection::JsonRpcConnection;
 use crate::connection::JsonRpcConnectionEvent;
-use crate::discover_v2::capability_locations::CapabilityLocationRequest;
 use crate::discover_v2::capability_manager::CapabilityManager;
 use crate::rpc::RpcCallError;
 use crate::rpc::RpcNotificationSender;
@@ -31,7 +29,6 @@ use crate::telemetry::ConnectionTransport;
 use crate::telemetry::ExecServerTelemetry;
 use crate::telemetry::ExecutorRegistration;
 use codex_http_client::HttpClientFactory;
-use codex_utils_path_uri::PathUri;
 
 #[derive(Clone)]
 pub(crate) struct ConnectionProcessor {
@@ -45,49 +42,11 @@ pub(crate) struct ConnectionProcessor {
 }
 
 impl ConnectionProcessor {
-    #[cfg(test)]
-    pub(crate) fn new(runtime_paths: ExecServerRuntimeOptions) -> Self {
-        Self::new_with_location_request(
-            runtime_paths,
-            ExecServerTelemetry::default(),
-            codex_http_client::HttpClientFactory::new(
-                codex_http_client::OutboundProxyPolicy::ReqwestDefault,
-            ),
-            RequestDispatchMode::Inline,
-            // Connection-only tests must not scan the real user home.
-            Err(io::Error::other(
-                "capability home paths not supplied by this test",
-            )),
-        )
-    }
-
     pub(crate) fn new_with_telemetry(
         runtime_paths: ExecServerRuntimeOptions,
         telemetry: ExecServerTelemetry,
         http_client_factory: HttpClientFactory,
         request_dispatch_mode: RequestDispatchMode,
-    ) -> Self {
-        let request =
-            codex_utils_home_dir::find_codex_home().map(|codex_home| CapabilityLocationRequest {
-                codex_home: PathUri::from_abs_path(&codex_home),
-                user_home: dirs::home_dir()
-                    .and_then(|home| PathUri::from_host_native_path(home).ok()),
-            });
-        Self::new_with_location_request(
-            runtime_paths,
-            telemetry,
-            http_client_factory,
-            request_dispatch_mode,
-            request,
-        )
-    }
-
-    fn new_with_location_request(
-        runtime_paths: ExecServerRuntimeOptions,
-        telemetry: ExecServerTelemetry,
-        http_client_factory: HttpClientFactory,
-        request_dispatch_mode: RequestDispatchMode,
-        request: io::Result<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]

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

@@ -5,7 +5,9 @@ use opentelemetry::trace::TraceContextExt;
 use std::sync::Arc;
 
 use crate::local_file_system::resolve_existing_path;
+use crate::protocol::CAPABILITIES_DISCOVER_V2_METHOD;
 use crate::protocol::CAPABILITY_ROOTS_DISCOVER_METHOD;
+use crate::protocol::DiscoverV2CapabilitiesRequest;
 use crate::protocol::ENVIRONMENT_CONFIG_READ_METHOD;
 use crate::protocol::ENVIRONMENT_INFO_METHOD;
 use crate::protocol::ENVIRONMENT_STATUS_METHOD;
@@ -109,6 +111,12 @@ pub(crate) fn build_router() -> RpcRouter<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 {

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

@@ -33,6 +33,7 @@ use super::RequestDispatcher;
 use super::RequestTaskResult;
 use crate::ExecServerRuntimeOptions;
 use crate::connection::JsonRpcConnectionEvent;
+use crate::discover_v2::capability_manager::CapabilityManager;
 use crate::rpc::RpcNotificationSender;
 use crate::rpc::RpcRouter;
 use crate::rpc::RpcServerOutboundMessage;
@@ -177,6 +178,7 @@ async fn request_queue_waits_for_dispatcher_admission_before_recording_telemetry
     let requests = notifications.request_sender();
     let handler = Arc::new(ExecServerHandler::new(
         SessionRegistry::new(telemetry.clone()),
+        CapabilityManager::new(crate::LocalFileSystem::default()),
         notifications,
         ExecServerRuntimeOptions::new(
             std::env::current_exe().expect("current executable"),
@@ -386,6 +388,7 @@ fn request_telemetry_fixture(router: RpcRouter<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"),