mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
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:
@@ -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> {
|
||||
|
||||
@@ -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)
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
},
|
||||
|
||||
@@ -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(),
|
||||
},
|
||||
|
||||
@@ -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],
|
||||
|
||||
@@ -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(),
|
||||
},
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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] {
|
||||
|
||||
Reference in New Issue
Block a user