Build MCP tool runtimes before router planning

codex-rs/core/src/mcp_tool_exposure.rs

@@ -1,47 +1,58 @@
 use std::collections::HashSet;
+use std::sync::Arc;
 
 use codex_connectors::AppToolPolicyEvaluator;
 use codex_connectors::AppToolPolicyInput;
 use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
 use codex_mcp::ToolInfo as McpToolInfo;
 use codex_mcp::tool_is_model_visible;
+use codex_tools::ToolExposure;
 use tracing::instrument;
+use tracing::warn;
 
 use crate::config::Config;
 use crate::connectors;
-
-pub(crate) struct McpToolExposure {
-    pub(crate) direct_tools: Vec<McpToolInfo>,
-    pub(crate) deferred_tools: Option<Vec<McpToolInfo>>,
-}
+use crate::tools::handlers::McpHandler;
+use crate::tools::registry::CoreToolRuntime;
+use crate::tools::registry::override_tool_exposure;
 
 #[instrument(level = "trace", skip_all)]
-pub(crate) fn build_mcp_tool_exposure(
+pub(crate) fn build_mcp_tool_runtimes(
     all_mcp_tools: &[McpToolInfo],
     connectors: Option<&[connectors::AppInfo]>,
     config: &Config,
     search_tool_enabled: bool,
-) -> McpToolExposure {
-    let mut deferred_tools = filter_non_codex_apps_mcp_tools_only(all_mcp_tools);
+) -> Vec<Arc<dyn CoreToolRuntime>> {
+    let mut exposed_tools = filter_non_codex_apps_mcp_tools_only(all_mcp_tools);
     if let Some(connectors) = connectors {
-        deferred_tools.extend(filter_codex_apps_mcp_tools(
+        exposed_tools.extend(filter_codex_apps_mcp_tools(
             all_mcp_tools,
             connectors,
             config,
         ));
     }
 
-    if !search_tool_enabled {
-        return McpToolExposure {
-            direct_tools: deferred_tools,
-            deferred_tools: None,
-        };
-    }
-
-    McpToolExposure {
-        direct_tools: Vec::new(),
-        deferred_tools: (!deferred_tools.is_empty()).then_some(deferred_tools),
-    }
+    let exposure = if search_tool_enabled {
+        ToolExposure::Deferred
+    } else {
+        ToolExposure::Direct
+    };
+    exposed_tools
+        .into_iter()
+        .filter_map(|tool| {
+            let tool_name = tool.canonical_tool_name();
+            match McpHandler::new(tool) {
+                Ok(handler) => {
+                    let handler: Arc<dyn CoreToolRuntime> = Arc::new(handler);
+                    Some(override_tool_exposure(handler, exposure))
+                }
+                Err(err) => {
+                    warn!("Skipping MCP tool `{tool_name}`: failed to build tool spec: {err}");
+                    None
+                }
+            }
+        })
+        .collect()
 }
 
 fn filter_non_codex_apps_mcp_tools_only(mcp_tools: &[McpToolInfo]) -> Vec<McpToolInfo> {

codex-rs/core/src/mcp_tool_exposure_test.rs

@@ -1,8 +1,9 @@
-use std::collections::HashSet;
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
 use codex_mcp::ToolInfo;
+use codex_tools::ToolExposure;
 use codex_tools::ToolName;
 use pretty_assertions::assert_eq;
 use rmcp::model::JsonObject;
@@ -79,10 +80,20 @@ fn numbered_mcp_tools(count: usize) -> Vec<ToolInfo> {
         .collect()
 }
 
-fn tool_names(tools: &[ToolInfo]) -> HashSet<ToolName> {
+fn expected_runtimes(
+    tools: &[ToolInfo],
+    exposure: ToolExposure,
+) -> HashMap<ToolName, ToolExposure> {
     tools
         .iter()
-        .map(codex_mcp::ToolInfo::canonical_tool_name)
+        .map(|tool| (tool.canonical_tool_name(), exposure))
+        .collect()
+}
+
+fn runtimes_by_name(runtimes: &[Arc<dyn CoreToolRuntime>]) -> HashMap<ToolName, ToolExposure> {
+    runtimes
+        .iter()
+        .map(|runtime| (runtime.tool_name(), runtime.exposure()))
         .collect()
 }
 
@@ -101,12 +112,14 @@ async fn directly_exposes_effective_tool_sets_when_search_is_unavailable() {
     let config = test_config().await;
     let mcp_tools = numbered_mcp_tools(/*count*/ 2);
 
-    let exposure = build_mcp_tool_exposure(
+    let runtimes = build_mcp_tool_runtimes(
         &mcp_tools, /*connectors*/ None, &config, /*search_tool_enabled*/ false,
     );
 
-    assert_eq!(tool_names(&exposure.direct_tools), tool_names(&mcp_tools));
-    assert!(exposure.deferred_tools.is_none());
+    assert_eq!(
+        runtimes_by_name(&runtimes),
+        expected_runtimes(&mcp_tools, ToolExposure::Direct)
+    );
 }
 
 #[tokio::test]
@@ -173,18 +186,17 @@ async fn excludes_tools_hidden_from_model_exposure() {
     ];
     let connectors = vec![make_connector("calendar", "Calendar")];
 
-    let exposure = build_mcp_tool_exposure(
+    let runtimes = build_mcp_tool_runtimes(
         &mcp_tools,
         Some(connectors.as_slice()),
         &config,
         /*search_tool_enabled*/ false,
     );
 
     assert_eq!(
-        tool_names(&exposure.direct_tools),
-        tool_names(&[visible_tool, visible_app_tool])
+        runtimes_by_name(&runtimes),
+        expected_runtimes(&[visible_tool, visible_app_tool], ToolExposure::Direct)
     );
-    assert!(exposure.deferred_tools.is_none());
 }
 
 #[tokio::test]
@@ -224,35 +236,32 @@ enabled = true
     );
     let connectors = vec![make_connector("calendar", "Calendar")];
 
-    let exposure = build_mcp_tool_exposure(
+    let runtimes = build_mcp_tool_runtimes(
         &[enabled_tool.clone(), disabled_tool],
         Some(connectors.as_slice()),
         &config,
         /*search_tool_enabled*/ false,
     );
 
     assert_eq!(
-        tool_names(&exposure.direct_tools),
-        tool_names(&[enabled_tool])
+        runtimes_by_name(&runtimes),
+        expected_runtimes(&[enabled_tool], ToolExposure::Direct)
     );
-    assert!(exposure.deferred_tools.is_none());
 }
 
 #[tokio::test]
 async fn defers_effective_tool_sets_when_search_is_available() {
     let config = test_config().await;
     let mcp_tools = numbered_mcp_tools(/*count*/ 2);
 
-    let exposure = build_mcp_tool_exposure(
+    let runtimes = build_mcp_tool_runtimes(
         &mcp_tools, /*connectors*/ None, &config, /*search_tool_enabled*/ true,
     );
 
-    assert!(exposure.direct_tools.is_empty());
-    let deferred_tools = exposure
-        .deferred_tools
-        .as_ref()
-        .expect("MCP tools should be discoverable through tool_search");
-    assert_eq!(tool_names(deferred_tools), tool_names(&mcp_tools));
+    assert_eq!(
+        runtimes_by_name(&runtimes),
+        expected_runtimes(&mcp_tools, ToolExposure::Deferred)
+    );
 }
 
 #[tokio::test]
@@ -278,22 +287,15 @@ async fn defers_apps_and_non_app_mcp_tools() {
     ];
     let connectors = vec![make_connector("calendar", "Calendar")];
 
-    let exposure = build_mcp_tool_exposure(
+    let runtimes = build_mcp_tool_runtimes(
         &mcp_tools,
         Some(connectors.as_slice()),
         &config,
         /*search_tool_enabled*/ true,
     );
 
-    assert!(exposure.direct_tools.is_empty());
-    let deferred_tools = exposure
-        .deferred_tools
-        .as_ref()
-        .expect("MCP tools should be discoverable through tool_search");
-    let deferred_tool_names = tool_names(deferred_tools);
-    assert!(deferred_tool_names.contains(&ToolName::namespaced("mcp__rmcp", "tool")));
-    assert!(deferred_tool_names.contains(&ToolName::namespaced(
-        "mcp__codex_apps__calendar",
-        "_create_event"
-    )));
+    assert_eq!(
+        runtimes_by_name(&runtimes),
+        expected_runtimes(&mcp_tools, ToolExposure::Deferred)
+    );
 }

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

@@ -654,8 +654,7 @@ fn test_tool_runtime(session: Arc<Session>, turn_context: Arc<TurnContext>) -> T
         step_context.as_ref(),
         crate::tools::router::ToolRouterParams {
             tool_suggest_candidates: None,
-            mcp_tools: None,
-            deferred_mcp_tools: None,
+            tool_runtimes: Vec::new(),
             extension_tool_executors: Vec::new(),
             dynamic_tools: turn_context.dynamic_tools.as_slice(),
         },
@@ -10389,22 +10388,12 @@ async fn abort_review_task_emits_exited_then_aborted_and_records_history() {
 #[tokio::test]
 async fn fatal_tool_error_stops_turn_and_reports_error() {
     let (session, turn_context, _rx) = make_session_and_context_with_rx().await;
-    let tools = {
-        session
-            .services
-            .latest_mcp_runtime()
-            .manager()
-            .list_all_tools()
-            .await
-    };
-    let deferred_mcp_tools = Some(tools.clone());
     let step_context = StepContext::for_test(Arc::clone(&turn_context));
     let router = ToolRouter::from_context(
         step_context.as_ref(),
         crate::tools::router::ToolRouterParams {
             tool_suggest_candidates: None,
-            deferred_mcp_tools,
-            mcp_tools: Some(tools),
+            tool_runtimes: Vec::new(),
             extension_tool_executors: Vec::new(),
             dynamic_tools: turn_context.dynamic_tools.as_slice(),
         },

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

@@ -28,7 +28,7 @@ use crate::injection::ToolMentionKind;
 use crate::injection::app_id_from_path;
 use crate::injection::tool_kind_for_path;
 use crate::mcp_skill_dependencies::maybe_prompt_and_install_mcp_dependencies;
-use crate::mcp_tool_exposure::build_mcp_tool_exposure;
+use crate::mcp_tool_exposure::build_mcp_tool_runtimes;
 use crate::mentions::build_connector_slug_counts;
 use crate::mentions::build_skill_name_counts;
 use crate::mentions::collect_explicit_app_ids;
@@ -1222,8 +1222,6 @@ pub(crate) async fn built_tools(
     cancellation_token: &CancellationToken,
 ) -> CodexResult<Arc<ToolRouter>> {
     let turn_context = step_context.turn.as_ref();
-    let mcp_connection_manager = step_context.mcp.manager();
-    let has_mcp_servers = mcp_connection_manager.has_servers();
     let all_mcp_tools = step_context
         .mcp_tools()
         .or_cancel(cancellation_token)
@@ -1330,19 +1328,16 @@ pub(crate) async fn built_tools(
             .instrument(trace_span!("built_tools.load_discoverable_tools"))
             .await
         };
-    let mcp_tool_exposure = build_mcp_tool_exposure(
+    let mcp_tool_runtimes = build_mcp_tool_runtimes(
         all_mcp_tools,
         connectors.as_deref(),
         &turn_context.config,
         search_tool_enabled(turn_context),
     );
-    let mcp_tools = has_mcp_servers.then_some(mcp_tool_exposure.direct_tools);
-    let deferred_mcp_tools = mcp_tool_exposure.deferred_tools;
     Ok(Arc::new(ToolRouter::from_context(
         step_context,
         ToolRouterParams {
-            mcp_tools,
-            deferred_mcp_tools,
+            tool_runtimes: mcp_tool_runtimes,
             tool_suggest_candidates,
             extension_tool_executors: extension_tool_executors(sess),
             dynamic_tools: turn_context.dynamic_tools.as_slice(),

codex-rs/core/src/stream_events_utils_tests.rs

@@ -280,8 +280,7 @@ async fn handle_output_item_done_returns_contributed_last_agent_message() {
         step_context.as_ref(),
         crate::tools::router::ToolRouterParams {
             tool_suggest_candidates: None,
-            mcp_tools: None,
-            deferred_mcp_tools: None,
+            tool_runtimes: Vec::new(),
             extension_tool_executors: Vec::new(),
             dynamic_tools: turn_context.dynamic_tools.as_slice(),
         },

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

@@ -6,10 +6,10 @@ use crate::tools::context::ToolInvocation;
 use crate::tools::context::ToolPayload;
 use crate::tools::handlers::ToolSearchHandlerCache;
 use crate::tools::registry::AnyToolResult;
+use crate::tools::registry::CoreToolRuntime;
 use crate::tools::registry::ToolArgumentDiffConsumer;
 use crate::tools::registry::ToolRegistry;
 use crate::tools::spec_plan::build_tool_router;
-use codex_mcp::ToolInfo;
 use codex_protocol::dynamic_tools::DynamicToolSpec;
 use codex_protocol::models::ResponseItem;
 use codex_protocol::models::SearchToolCallParams;
@@ -38,8 +38,7 @@ pub struct ToolRouter {
 }
 
 pub(crate) struct ToolRouterParams<'a> {
-    pub(crate) mcp_tools: Option<Vec<ToolInfo>>,
-    pub(crate) deferred_mcp_tools: Option<Vec<ToolInfo>>,
+    pub(crate) tool_runtimes: Vec<Arc<dyn CoreToolRuntime>>,
     pub(crate) tool_suggest_candidates: Option<ToolSuggestCandidates>,
     pub(crate) extension_tool_executors: Vec<Arc<dyn ToolExecutor<ExtensionToolCall>>>,
     pub(crate) dynamic_tools: &'a [DynamicToolSpec],