diff --git a/codex-rs/core/src/mcp_tool_exposure.rs b/codex-rs/core/src/mcp_tool_exposure.rs index fda2b04c9b97..1bbd02860820 100644 --- a/codex-rs/core/src/mcp_tool_exposure.rs +++ b/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, - pub(crate) deferred_tools: Option>, -} +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> { + 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 = 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 { diff --git a/codex-rs/core/src/mcp_tool_exposure_test.rs b/codex-rs/core/src/mcp_tool_exposure_test.rs index 8e9cc2256fba..605ee763f0d3 100644 --- a/codex-rs/core/src/mcp_tool_exposure_test.rs +++ b/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 { .collect() } -fn tool_names(tools: &[ToolInfo]) -> HashSet { +fn expected_runtimes( + tools: &[ToolInfo], + exposure: ToolExposure, +) -> HashMap { tools .iter() - .map(codex_mcp::ToolInfo::canonical_tool_name) + .map(|tool| (tool.canonical_tool_name(), exposure)) + .collect() +} + +fn runtimes_by_name(runtimes: &[Arc]) -> HashMap { + 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,7 +186,7 @@ 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, @@ -181,10 +194,9 @@ async fn excludes_tools_hidden_from_model_exposure() { ); 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,7 +236,7 @@ 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, @@ -232,10 +244,9 @@ enabled = true ); 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] @@ -243,16 +254,14 @@ 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) + ); } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 53314ea9ae4b..b27a280eac93 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -654,8 +654,7 @@ fn test_tool_runtime(session: Arc, turn_context: Arc) -> 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(), }, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 2ecdaacc55ea..24e9458c4fa6 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/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> { 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(), diff --git a/codex-rs/core/src/stream_events_utils_tests.rs b/codex-rs/core/src/stream_events_utils_tests.rs index eb04e25942a8..02d73cd2a04d 100644 --- a/codex-rs/core/src/stream_events_utils_tests.rs +++ b/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(), }, diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index aca86b124b3a..32e139cae734 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/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>, - pub(crate) deferred_mcp_tools: Option>, + pub(crate) tool_runtimes: Vec>, pub(crate) tool_suggest_candidates: Option, pub(crate) extension_tool_executors: Vec>>, pub(crate) dynamic_tools: &'a [DynamicToolSpec], diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index 1f38cf7580a8..b918035ccc87 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -4,6 +4,8 @@ use crate::config::Config; use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::tools::context::ToolPayload; +use crate::tools::handlers::McpHandler; +use crate::tools::registry::CoreToolRuntime; use crate::turn_diff_tracker::TurnDiffTracker; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionRegistry; @@ -105,21 +107,14 @@ fn extension_tool_test_registry() -> Arc> { #[tokio::test] async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow::Result<()> { - let (session, turn) = make_session_and_context().await; + let (_, turn) = make_session_and_context().await; let turn = Arc::new(turn); let step_context = StepContext::for_test(Arc::clone(&turn)); - let mcp_tools = session - .services - .latest_mcp_runtime() - .manager() - .list_all_tools() - .await; let router = ToolRouter::from_context( step_context.as_ref(), ToolRouterParams { tool_suggest_candidates: None, - deferred_mcp_tools: None, - mcp_tools: Some(mcp_tools), + tool_runtimes: Vec::new(), extension_tool_executors: Vec::new(), dynamic_tools: turn.dynamic_tools.as_slice(), }, @@ -217,21 +212,20 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { step_context.as_ref(), ToolRouterParams { tool_suggest_candidates: None, - deferred_mcp_tools: None, - mcp_tools: Some(vec![ - mcp_tool_info( + tool_runtimes: vec![ + mcp_runtime(mcp_tool_info( "echo", /*supports_parallel_tool_calls*/ true, "mcp__echo__", "query_with_delay", - ), - mcp_tool_info( + )), + mcp_runtime(mcp_tool_info( "hello_echo", /*supports_parallel_tool_calls*/ false, "mcp__hello_echo__", "query_with_delay", - ), - ]), + )), + ], extension_tool_executors: Vec::new(), dynamic_tools: turn.dynamic_tools.as_slice(), }, @@ -268,8 +262,7 @@ async fn tools_without_handlers_do_not_support_parallel() -> anyhow::Result<()> step_context.as_ref(), ToolRouterParams { tool_suggest_candidates: None, - deferred_mcp_tools: None, - mcp_tools: None, + tool_runtimes: Vec::new(), extension_tool_executors: Vec::new(), dynamic_tools: turn.dynamic_tools.as_slice(), }, @@ -325,8 +318,7 @@ async fn specs_filter_deferred_dynamic_tools() -> anyhow::Result<()> { step_context.as_ref(), ToolRouterParams { tool_suggest_candidates: None, - deferred_mcp_tools: None, - mcp_tools: None, + tool_runtimes: Vec::new(), extension_tool_executors: Vec::new(), dynamic_tools: &dynamic_tools, }, @@ -368,6 +360,10 @@ fn mcp_tool_info( } } +fn mcp_runtime(tool_info: codex_mcp::ToolInfo) -> Arc { + Arc::new(McpHandler::new(tool_info).expect("MCP tool spec should build")) +} + #[tokio::test] async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow::Result<()> { let (mut session, turn) = make_session_and_context().await; @@ -393,8 +389,7 @@ async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow step_context.as_ref(), ToolRouterParams { tool_suggest_candidates: None, - deferred_mcp_tools: None, - mcp_tools: None, + tool_runtimes: Vec::new(), extension_tool_executors: extension_tool_executors(&session), dynamic_tools: turn.dynamic_tools.as_slice(), }, diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index ace4c85fd1f6..83169918e0b3 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -16,7 +16,6 @@ use crate::tools::handlers::GetContextRemainingHandler; use crate::tools::handlers::ListAvailablePluginsToInstallHandler; use crate::tools::handlers::ListMcpResourceTemplatesHandler; use crate::tools::handlers::ListMcpResourcesHandler; -use crate::tools::handlers::McpHandler; use crate::tools::handlers::NewContextWindowHandler; use crate::tools::handlers::PlanHandler; use crate::tools::handlers::ReadMcpResourceHandler; @@ -61,7 +60,6 @@ use crate::tools::router::ToolRouter; use crate::tools::router::ToolRouterParams; use codex_features::Feature; use codex_login::AuthManager; -use codex_mcp::ToolInfo; use codex_protocol::config_types::WebSearchMode; use codex_protocol::dynamic_tools::DynamicToolNamespaceTool; use codex_protocol::dynamic_tools::DynamicToolSpec; @@ -145,8 +143,7 @@ impl PlannedTools { #[derive(Clone, Copy)] struct CoreToolPlanContext<'a> { step_context: &'a StepContext, - mcp_tools: Option<&'a [ToolInfo]>, - deferred_mcp_tools: Option<&'a [ToolInfo]>, + tool_runtimes: &'a [PlannedRuntime], tool_suggest_candidates: Option<&'a crate::tools::router::ToolSuggestCandidates>, extension_tool_executors: &'a [Arc>], dynamic_tools: &'a [DynamicToolSpec], @@ -174,8 +171,7 @@ fn build_tool_specs_and_registry( ) -> (Vec, ToolRegistry) { let turn_context = step_context.turn.as_ref(); let ToolRouterParams { - mcp_tools, - deferred_mcp_tools, + tool_runtimes, tool_suggest_candidates, extension_tool_executors, dynamic_tools, @@ -184,8 +180,7 @@ fn build_tool_specs_and_registry( crate::agent::role::spawn_tool_spec::build(&std::collections::BTreeMap::new()); let context = CoreToolPlanContext { step_context, - mcp_tools: mcp_tools.as_deref(), - deferred_mcp_tools: deferred_mcp_tools.as_deref(), + tool_runtimes: &tool_runtimes, tool_suggest_candidates: tool_suggest_candidates.as_ref(), extension_tool_executors: &extension_tool_executors, dynamic_tools, @@ -613,7 +608,9 @@ fn add_tool_sources(context: &CoreToolPlanContext<'_>, planned_tools: &mut Plann add_mcp_resource_tools(context, planned_tools); add_core_utility_tools(context, planned_tools); add_collaboration_tools(context, planned_tools); - add_mcp_runtime_tools(context, planned_tools); + for runtime in context.tool_runtimes { + planned_tools.add_arc(Arc::clone(runtime)); + } add_extension_tools(context, planned_tools); add_dynamic_tools(context, planned_tools); for spec in hosted_model_tool_specs(context) { @@ -695,7 +692,7 @@ fn unified_exec_should_include_shell_parameter( #[instrument(level = "trace", skip_all)] fn add_mcp_resource_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) { - if context.mcp_tools.is_some() { + if context.step_context.mcp.manager().has_servers() { planned_tools.add(ListMcpResourcesHandler); planned_tools.add(ListMcpResourceTemplatesHandler); planned_tools.add(ReadMcpResourceHandler); @@ -879,45 +876,7 @@ fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mu } } -#[instrument( - level = "trace", - skip_all, - fields( - direct_mcp_tool_count = context.mcp_tools.map_or(0, <[ToolInfo]>::len), - deferred_mcp_tool_count = context.deferred_mcp_tools.map_or(0, <[ToolInfo]>::len) - ) -)] -fn add_mcp_runtime_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) { - if let Some(mcp_tools) = context.mcp_tools { - for tool in mcp_tools { - match McpHandler::new(tool.clone()) { - Ok(handler) => planned_tools.add(handler), - Err(err) => warn!( - "Skipping MCP tool `{}`: failed to build tool spec: {err}", - tool.canonical_tool_name() - ), - } - } - } - - if let Some(deferred_mcp_tools) = context.deferred_mcp_tools { - for tool in deferred_mcp_tools { - match McpHandler::new(tool.clone()) { - Ok(handler) => planned_tools.add_with_exposure(handler, ToolExposure::Deferred), - Err(err) => warn!( - "Skipping deferred MCP tool `{}`: failed to build tool spec: {err}", - tool.canonical_tool_name() - ), - } - } - } -} - -#[instrument( - level = "trace", - skip_all, - fields(dynamic_tool_count = context.dynamic_tools.len()) -)] +#[instrument(level = "trace", skip_all, fields(dynamic_tool_count = context.dynamic_tools.len()))] fn add_dynamic_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) { for spec in context.dynamic_tools { match spec { diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index b7fb9070379a..c084e30ac11e 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -34,8 +34,11 @@ use crate::config::CurrentTimeReminderConfig; use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::session::turn_context::TurnContext; +use crate::tools::handlers::McpHandler; use crate::tools::handlers::ToolSearchHandlerCache; use crate::tools::handlers::multi_agents_spec::MULTI_AGENT_V1_NAMESPACE; +use crate::tools::registry::CoreToolRuntime; +use crate::tools::registry::override_tool_exposure; use crate::tools::router::ToolRouter; use crate::tools::router::ToolRouterParams; use crate::tools::router::ToolSuggestCandidates; @@ -45,8 +48,7 @@ const MULTI_AGENT_V2_NAMESPACE: &str = "collaboration"; #[derive(Default)] struct ToolPlanInputs { - mcp_tools: Option>, - deferred_mcp_tools: Option>, + tool_runtimes: Vec>, tool_suggest_candidates: Option, extension_tool_executors: Vec>>, dynamic_tools: Vec, @@ -187,8 +189,7 @@ async fn probe_with( step_context.as_ref(), ToolRouterParams { tool_suggest_candidates: inputs.tool_suggest_candidates, - mcp_tools: inputs.mcp_tools, - deferred_mcp_tools: inputs.deferred_mcp_tools, + tool_runtimes: inputs.tool_runtimes, extension_tool_executors: inputs.extension_tool_executors, dynamic_tools: inputs.dynamic_tools.as_slice(), }, @@ -369,12 +370,16 @@ fn mcp_tool(server: &str, namespace: &str, name: &str) -> ToolInfo { } } -fn invalid_mcp_tool(server: &str, namespace: &str, name: &str) -> ToolInfo { - let mut tool = mcp_tool(server, namespace, name); - tool.tool.input_schema = Arc::new(rmcp::model::object(json!({ - "type": "null", - }))); - tool +fn mcp_runtime( + server: &str, + namespace: &str, + name: &str, + exposure: ToolExposure, +) -> Arc { + let handler: Arc = Arc::new( + McpHandler::new(mcp_tool(server, namespace, name)).expect("MCP tool spec should build"), + ); + override_tool_exposure(handler, exposure) } fn dynamic_tool(namespace: Option<&str>, name: &str, defer_loading: bool) -> DynamicToolSpec { @@ -678,8 +683,7 @@ async fn environment_tools_follow_the_step_context() { let plan = ToolPlanProbe::from_router(ToolRouter::from_context( step_context.as_ref(), ToolRouterParams { - mcp_tools: None, - deferred_mcp_tools: None, + tool_runtimes: Vec::new(), tool_suggest_candidates: None, extension_tool_executors: Vec::new(), dynamic_tools: &[], @@ -737,23 +741,28 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() { let direct_mcp = probe_with( |_| {}, ToolPlanInputs { - mcp_tools: Some(vec![mcp_tool("direct", "mcp__direct", "lookup")]), + tool_runtimes: vec![mcp_runtime( + "direct", + "mcp__direct", + "lookup", + ToolExposure::Direct, + )], ..ToolPlanInputs::default() }, ) .await; - direct_mcp.assert_visible_contains(&[ - "list_mcp_resources", - "list_mcp_resource_templates", - "read_mcp_resource", - ]); assert_eq!( direct_mcp.namespace_function_names("mcp__direct"), &["lookup".to_string()] ); let searchable_mcp = ToolPlanInputs { - deferred_mcp_tools: Some(vec![mcp_tool("searchable", "mcp__searchable", "lookup")]), + tool_runtimes: vec![mcp_runtime( + "searchable", + "mcp__searchable", + "lookup", + ToolExposure::Deferred, + )], ..ToolPlanInputs::default() }; @@ -762,7 +771,7 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() { turn.model_info.supports_search_tool = false; }, ToolPlanInputs { - deferred_mcp_tools: searchable_mcp.deferred_mcp_tools.clone(), + tool_runtimes: searchable_mcp.tool_runtimes.clone(), ..ToolPlanInputs::default() }, ) @@ -787,7 +796,7 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() { use_bedrock_provider(turn); }, ToolPlanInputs { - deferred_mcp_tools: searchable_mcp.deferred_mcp_tools.clone(), + tool_runtimes: searchable_mcp.tool_runtimes.clone(), ..ToolPlanInputs::default() }, ) @@ -838,8 +847,12 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() { let first_router = ToolRouter::from_context( first_step_context.as_ref(), ToolRouterParams { - mcp_tools: None, - deferred_mcp_tools: Some(vec![mcp_tool("first", "mcp__first", "lookup")]), + tool_runtimes: vec![mcp_runtime( + "first", + "mcp__first", + "lookup", + ToolExposure::Deferred, + )], tool_suggest_candidates: None, extension_tool_executors: Vec::new(), dynamic_tools: &[], @@ -855,8 +868,12 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() { let second_router = ToolRouter::from_context( second_step_context.as_ref(), ToolRouterParams { - mcp_tools: None, - deferred_mcp_tools: Some(vec![mcp_tool("second", "mcp__second", "lookup")]), + tool_runtimes: vec![mcp_runtime( + "second", + "mcp__second", + "lookup", + ToolExposure::Deferred, + )], tool_suggest_candidates: None, extension_tool_executors: Vec::new(), dynamic_tools: &[], @@ -886,21 +903,6 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() { assert!(!second_description.contains("- first: Tools from first.")); } -#[tokio::test] -async fn invalid_mcp_tools_are_not_registered() { - let plan = probe_with( - |_| {}, - ToolPlanInputs { - mcp_tools: Some(vec![invalid_mcp_tool("invalid", "mcp__invalid", "lookup")]), - ..ToolPlanInputs::default() - }, - ) - .await; - - plan.assert_visible_lacks(&["mcp__invalid"]); - plan.assert_registered_lacks(&[&ToolName::namespaced("mcp__invalid", "lookup").to_string()]); -} - #[tokio::test] async fn request_plugin_install_requires_all_discovery_features() { for disabled_feature in [Feature::ToolSuggest, Feature::Apps, Feature::Plugins] {