Delegate readiness waits to tool runtimes (#36120)

## What changed

- Add an optional `wait_until_ready` hook to `CoreToolRuntime` and invoke it
  before acquiring the tool execution gate.
- Move MCP refresh and server-startup waiting into the MCP runtime's hook.
- Forward readiness waits through tool exposure and namespace wrappers.
- Test exact runtime selection for plain and namespaced tools, including
  exposure overrides and missing tools.

GitOrigin-RevId: 8718589fa06f14d053a35a288186613f649d25ff

codex-rs/core/src/tools/handlers/mcp.rs

@@ -4,6 +4,7 @@ use std::time::Instant;
 use crate::function_tool::FunctionCallError;
 use crate::mcp_tool_call::handle_mcp_tool_call;
 use crate::original_image_detail::can_request_original_image_detail;
+use crate::session::session::Session;
 use crate::tools::context::McpToolOutput;
 use crate::tools::context::ToolInvocation;
 use crate::tools::context::ToolPayload;
@@ -24,6 +25,7 @@ use codex_tools::ToolSearchSourceInfo;
 use codex_tools::ToolSpec;
 use codex_tools::mcp_tool_to_responses_api_tool;
 use codex_utils_string::take_bytes_at_char_boundary;
+use futures::future::BoxFuture;
 use serde_json::Map;
 use serde_json::Value;
 
@@ -166,14 +168,21 @@ impl McpHandler {
 }
 
 impl CoreToolRuntime for McpHandler {
-    fn mcp_server_name(&self) -> Option<&str> {
-        Some(&self.tool_info.server_name)
+    fn wait_until_ready<'a>(&'a self, session: &'a Arc<Session>) -> Option<BoxFuture<'a, ()>> {
+        Some(Box::pin(async move {
+            session.refresh_mcp_if_dirty().await;
+            session
+                .services
+                .mcp_runtime
+                .wait_for_server_startup(&self.tool_info.server_name)
+                .await;
+        }))
     }
 
     fn telemetry_tags<'a>(
         &'a self,
         _invocation: &'a ToolInvocation,
-    ) -> futures::future::BoxFuture<'a, ToolTelemetryTags> {
+    ) -> BoxFuture<'a, ToolTelemetryTags> {
         Box::pin(async {
             let mut tags = vec![("mcp_server", self.tool_info.server_name.clone())];
             if let Some(origin) = &self.tool_info.server_origin {

codex-rs/core/src/tools/parallel.rs

@@ -98,7 +98,7 @@ impl ToolCallRuntime {
         cancellation_token: CancellationToken,
     ) -> impl std::future::Future<Output = Result<AnyToolResult, FunctionCallError>> {
         let supports_parallel = self.router.tool_supports_parallel(&call);
-        let mcp_server = self.router.mcp_server_name(&call).map(str::to_owned);
+        let tool_runtime = self.router.tool_runtime(&call);
         let router = Arc::clone(&self.router);
         let session = Arc::clone(&self.session);
         let step_context = Arc::clone(&self.step_context);
@@ -131,13 +131,10 @@ impl ToolCallRuntime {
 
         let mut dispatch_handle: AbortOnDropHandle<Result<AnyToolResult, FunctionCallError>> =
             AbortOnDropHandle::new(tokio::spawn(async move {
-                if let Some(server) = mcp_server {
-                    session.refresh_mcp_if_dirty().await;
-                    session
-                        .services
-                        .mcp_runtime
-                        .wait_for_server_startup(&server)
-                        .await;
+                if let Some(tool_runtime) = tool_runtime
+                    && let Some(readiness) = tool_runtime.wait_until_ready(&session)
+                {
+                    readiness.await;
                 }
 
                 let _guard = if supports_parallel {

codex-rs/core/src/tools/registry.rs

@@ -14,6 +14,7 @@ use crate::memory_usage::emit_metric_for_tool_read;
 use crate::memory_usage::shell_script_for_invocation;
 use crate::sandbox_tags::permission_profile_policy_tag;
 use crate::sandbox_tags::permission_profile_sandbox_tag;
+use crate::session::session::Session;
 use crate::session::turn_context::TurnContext;
 use crate::tools::context::FunctionToolOutput;
 use crate::tools::context::ToolInvocation;
@@ -51,8 +52,8 @@ pub use codex_tools::ToolExposure;
 /// Implementers provide the shared `ToolExecutor` behavior plus optional
 /// core-owned metadata for hooks, telemetry, tool search, and argument diffs.
 pub(crate) trait CoreToolRuntime: ToolExecutor<ToolInvocation> {
-    /// Returns the MCP server that owns this runtime, when this is an MCP tool.
-    fn mcp_server_name(&self) -> Option<&str> {
+    /// Returns a readiness wait for this exact tool before taking the execution gate.
+    fn wait_until_ready<'a>(&'a self, _session: &'a Arc<Session>) -> Option<BoxFuture<'a, ()>> {
         None
     }
 
@@ -288,8 +289,8 @@ impl ToolExecutor<ToolInvocation> for ExposureOverride {
 }
 
 impl CoreToolRuntime for ExposureOverride {
-    fn mcp_server_name(&self) -> Option<&str> {
-        self.handler.mcp_server_name()
+    fn wait_until_ready<'a>(&'a self, session: &'a Arc<Session>) -> Option<BoxFuture<'a, ()>> {
+        self.handler.wait_until_ready(session)
     }
 
     fn matches_kind(&self, payload: &ToolPayload) -> bool {
@@ -397,14 +398,10 @@ impl ToolRegistry {
         Self::new(HashMap::from([(name, handler as Arc<dyn CoreToolRuntime>)]))
     }
 
-    fn tool(&self, name: &ToolName) -> Option<Arc<dyn CoreToolRuntime>> {
+    pub(crate) fn tool(&self, name: &ToolName) -> Option<Arc<dyn CoreToolRuntime>> {
         self.tools.get(name).map(Arc::clone)
     }
 
-    pub(crate) fn mcp_server_name(&self, name: &ToolName) -> Option<&str> {
-        self.tools.get(name)?.mcp_server_name()
-    }
-
     #[cfg(test)]
     pub(crate) fn tool_names_for_test(&self) -> Vec<ToolName> {
         let mut names = self.tools.keys().cloned().collect::<Vec<_>>();

codex-rs/core/src/tools/registry_tests.rs

@@ -1,6 +1,10 @@
 use super::*;
+use crate::session::session::Session;
 use crate::session::step_context::StepContext;
+use futures::future::BoxFuture;
 use pretty_assertions::assert_eq;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
 
 struct TestHandler {
     tool_name: codex_tools::ToolName,
@@ -29,6 +33,33 @@ impl ToolExecutor<ToolInvocation> for TestHandler {
 
 impl CoreToolRuntime for TestHandler {}
 
+struct ReadinessTestHandler {
+    handler: TestHandler,
+    readiness_waits: Arc<AtomicUsize>,
+}
+
+impl ToolExecutor<ToolInvocation> for ReadinessTestHandler {
+    fn tool_name(&self) -> codex_tools::ToolName {
+        self.handler.tool_name()
+    }
+
+    fn spec(&self) -> codex_tools::ToolSpec {
+        self.handler.spec()
+    }
+
+    fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> {
+        self.handler.handle(invocation)
+    }
+}
+
+impl CoreToolRuntime for ReadinessTestHandler {
+    fn wait_until_ready<'a>(&'a self, _session: &'a Arc<Session>) -> Option<BoxFuture<'a, ()>> {
+        Some(Box::pin(async {
+            self.readiness_waits.fetch_add(1, Ordering::Relaxed);
+        }))
+    }
+}
+
 #[derive(Clone)]
 enum LifecycleTestResult {
     Ok { success: bool },
@@ -179,6 +210,80 @@ fn handler_looks_up_namespaced_aliases_explicitly() {
     );
 }
 
+#[tokio::test]
+async fn readiness_selects_exact_tool_and_forwards_exposure_overrides() {
+    let (session, _turn) = crate::session::tests::make_session_and_context().await;
+    let session = Arc::new(session);
+    let plain_name = codex_tools::ToolName::plain("echo");
+    let namespaced_name = codex_tools::ToolName::namespaced("mcp__server__", "echo");
+    assert!(
+        TestHandler {
+            tool_name: plain_name.clone(),
+        }
+        .wait_until_ready(&session)
+        .is_none()
+    );
+    let plain_readiness_waits = Arc::new(AtomicUsize::new(0));
+    let namespaced_readiness_waits = Arc::new(AtomicUsize::new(0));
+    let plain_handler = Arc::new(ReadinessTestHandler {
+        handler: TestHandler {
+            tool_name: plain_name.clone(),
+        },
+        readiness_waits: Arc::clone(&plain_readiness_waits),
+    }) as Arc<dyn CoreToolRuntime>;
+    let namespaced_handler = override_tool_exposure(
+        Arc::new(ReadinessTestHandler {
+            handler: TestHandler {
+                tool_name: namespaced_name.clone(),
+            },
+            readiness_waits: Arc::clone(&namespaced_readiness_waits),
+        }),
+        ToolExposure::DirectModelOnly,
+    );
+    let registry = ToolRegistry::from_tools([plain_handler, namespaced_handler]);
+
+    registry
+        .tool(&plain_name)
+        .expect("plain runtime should be registered")
+        .wait_until_ready(&session)
+        .expect("plain runtime should provide a readiness wait")
+        .await;
+    assert_eq!(
+        [
+            plain_readiness_waits.load(Ordering::Relaxed),
+            namespaced_readiness_waits.load(Ordering::Relaxed),
+        ],
+        [1, 0]
+    );
+
+    registry
+        .tool(&namespaced_name)
+        .expect("namespaced runtime should be registered")
+        .wait_until_ready(&session)
+        .expect("namespaced runtime should forward its readiness wait")
+        .await;
+    assert_eq!(
+        [
+            plain_readiness_waits.load(Ordering::Relaxed),
+            namespaced_readiness_waits.load(Ordering::Relaxed),
+        ],
+        [1, 1]
+    );
+
+    assert!(
+        registry
+            .tool(&codex_tools::ToolName::namespaced("mcp__missing__", "echo"))
+            .is_none()
+    );
+    assert_eq!(
+        [
+            plain_readiness_waits.load(Ordering::Relaxed),
+            namespaced_readiness_waits.load(Ordering::Relaxed),
+        ],
+        [1, 1]
+    );
+}
+
 #[tokio::test]
 async fn function_tools_expose_default_hook_payloads_and_rewrites() -> anyhow::Result<()> {
     let (session, turn) = crate::session::tests::make_session_and_context().await;

codex-rs/core/src/tools/router.rs

@@ -135,8 +135,8 @@ impl ToolRouter {
             .unwrap_or(false)
     }
 
-    pub(crate) fn mcp_server_name(&self, call: &ToolCall) -> Option<&str> {
-        self.registry.mcp_server_name(&call.tool_name)
+    pub(crate) fn tool_runtime(&self, call: &ToolCall) -> Option<Arc<dyn CoreToolRuntime>> {
+        self.registry.tool(&call.tool_name)
     }
 
     pub fn tool_waits_for_runtime_cancellation(&self, call: &ToolCall) -> bool {

codex-rs/core/src/tools/router_tests.rs

@@ -181,15 +181,17 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow
         .expect("test session should expose a parallel shell-like tool");
 
     assert_eq!(
-        router.mcp_server_name(&ToolCall {
-            tool_name: ToolName::plain(parallel_tool_name),
-            call_id: "call-local-tool".to_string(),
-            payload: ToolPayload::Function {
-                arguments: "{}".to_string(),
-            },
-            encrypted_function_args: None,
-        }),
-        None
+        router
+            .tool_runtime(&ToolCall {
+                tool_name: ToolName::plain(parallel_tool_name),
+                call_id: "call-local-tool".to_string(),
+                payload: ToolPayload::Function {
+                    arguments: "{}".to_string(),
+                },
+                encrypted_function_args: None,
+            })
+            .map(|runtime| runtime.tool_name()),
+        Some(ToolName::plain(parallel_tool_name))
     );
 
     assert!(!router.tool_supports_parallel(&ToolCall {
@@ -303,7 +305,12 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> {
         encrypted_function_args: None,
     };
     assert!(router.tool_supports_parallel(&call));
-    assert_eq!(router.mcp_server_name(&call), Some("echo"));
+    assert_eq!(
+        router
+            .tool_runtime(&call)
+            .map(|runtime| runtime.tool_name()),
+        Some(call.tool_name.clone())
+    );
 
     let different_server_call = ToolCall {
         tool_name: ToolName::namespaced("mcp__hello_echo__", "query_with_delay"),
@@ -315,8 +322,10 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> {
     };
     assert!(!router.tool_supports_parallel(&different_server_call));
     assert_eq!(
-        router.mcp_server_name(&different_server_call),
-        Some("hello_echo")
+        router
+            .tool_runtime(&different_server_call)
+            .map(|runtime| runtime.tool_name()),
+        Some(different_server_call.tool_name.clone())
     );
 
     Ok(())

codex-rs/core/src/tools/spec_plan.rs

@@ -86,6 +86,7 @@ use codex_tools::default_namespace_description;
 use codex_tools::request_user_input_available_modes;
 use codex_tools::shell_command_backend_for_features;
 use codex_tools::shell_type_for_model_and_features;
+use futures::future::BoxFuture;
 use std::collections::BTreeMap;
 use std::collections::HashSet;
 use std::sync::Arc;
@@ -1089,6 +1090,10 @@ impl ToolExecutor<ToolInvocation> for MultiAgentV2NamespaceOverride {
 }
 
 impl CoreToolRuntime for MultiAgentV2NamespaceOverride {
+    fn wait_until_ready<'a>(&'a self, session: &'a Arc<Session>) -> Option<BoxFuture<'a, ()>> {
+        self.handler.wait_until_ready(session)
+    }
+
     fn matches_kind(&self, payload: &crate::tools::context::ToolPayload) -> bool {
         self.handler.matches_kind(payload)
     }