From 9a46fd33a0ac62e7d700f2667e1b643c50c4970a Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 30 Jul 2026 09:20:17 +0000 Subject: [PATCH] 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 | 15 +++- codex-rs/core/src/tools/parallel.rs | 13 ++- codex-rs/core/src/tools/registry.rs | 15 ++-- codex-rs/core/src/tools/registry_tests.rs | 105 ++++++++++++++++++++++ codex-rs/core/src/tools/router.rs | 4 +- codex-rs/core/src/tools/router_tests.rs | 33 ++++--- codex-rs/core/src/tools/spec_plan.rs | 5 ++ 7 files changed, 156 insertions(+), 34 deletions(-) diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 2bce827ff0..1ddb838b76 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/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) -> Option> { + 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 { diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index deb0ab85ac..0dd5bc6f45 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -98,7 +98,7 @@ impl ToolCallRuntime { cancellation_token: CancellationToken, ) -> impl std::future::Future> { 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> = 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 { diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 3927941abb..37a165e53b 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/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 { - /// 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) -> Option> { None } @@ -288,8 +289,8 @@ impl ToolExecutor 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) -> Option> { + 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)])) } - fn tool(&self, name: &ToolName) -> Option> { + pub(crate) fn tool(&self, name: &ToolName) -> Option> { 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 { let mut names = self.tools.keys().cloned().collect::>(); diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index 12ab6b3f28..eb60c67541 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/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 for TestHandler { impl CoreToolRuntime for TestHandler {} +struct ReadinessTestHandler { + handler: TestHandler, + readiness_waits: Arc, +} + +impl ToolExecutor 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) -> Option> { + 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; + 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; diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 1e604297af..8f63390b86 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/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> { + self.registry.tool(&call.tool_name) } pub fn tool_waits_for_runtime_cancellation(&self, call: &ToolCall) -> bool { diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index 3abe1cdb41..78b79d526b 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/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(()) diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 215ea66463..e0c616a6de 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/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 for MultiAgentV2NamespaceOverride { } impl CoreToolRuntime for MultiAgentV2NamespaceOverride { + fn wait_until_ready<'a>(&'a self, session: &'a Arc) -> Option> { + self.handler.wait_until_ready(session) + } + fn matches_kind(&self, payload: &crate::tools::context::ToolPayload) -> bool { self.handler.matches_kind(payload) }