Build MCP tool runtimes before router planning (#33149)

## What changed

- Convert filtered MCP tool metadata into `CoreToolRuntime` instances before constructing the tool router.
- Pass those runtimes through the shared tool-planning path instead of maintaining separate direct and deferred MCP tool lists.
- Preserve MCP visibility, app policy filtering, and direct versus deferred exposure behavior.

GitOrigin-RevId: 67f7fbeecc2992318c859e038cd97938833911db
This commit is contained in:
jif
2026-07-14 20:36:29 +00:00
committed by copyberry
parent eb80df0317
commit 1447cee36b
9 changed files with 141 additions and 190 deletions
+31 -20
View File
@@ -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> {
+35 -33
View File
@@ -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,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)
);
}
+2 -13
View File
@@ -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(),
},
+3 -8
View File
@@ -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(),
@@ -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(),
},
+2 -3
View File
@@ -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],
+17 -22
View File
@@ -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<ExtensionRegistry<Config>> {
#[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<dyn CoreToolRuntime> {
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(),
},
+8 -49
View File
@@ -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<dyn ToolExecutor<ExtensionToolCall>>],
dynamic_tools: &'a [DynamicToolSpec],
@@ -174,8 +171,7 @@ fn build_tool_specs_and_registry(
) -> (Vec<ToolSpec>, 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 {
+42 -40
View File
@@ -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<Vec<ToolInfo>>,
deferred_mcp_tools: Option<Vec<ToolInfo>>,
tool_runtimes: Vec<Arc<dyn CoreToolRuntime>>,
tool_suggest_candidates: Option<ToolSuggestCandidates>,
extension_tool_executors: Vec<Arc<dyn ToolExecutor<ExtensionToolCall>>>,
dynamic_tools: Vec<DynamicToolSpec>,
@@ -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<dyn CoreToolRuntime> {
let handler: Arc<dyn CoreToolRuntime> = 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] {