Keep deferred MCP discovery when flattening tools

This commit is contained in:
Mitsuhiro Kotake
2026-06-17 15:08:58 +09:00
parent 9b1033c43d
commit 6a8098abf1
7 changed files with 125 additions and 37 deletions
+1
View File
@@ -69,6 +69,7 @@ pub use request_user_input::RequestUserInputHandler;
pub use shell::ShellCommandHandler;
pub(crate) use shell::ShellCommandHandlerOptions;
pub use test_sync::TestSyncHandler;
pub(crate) use tool_search::NamespaceToolSpecMode;
pub use tool_search::ToolSearchHandler;
pub use unified_exec::ExecCommandHandler;
pub(crate) use unified_exec::ExecCommandHandlerOptions;
@@ -19,15 +19,26 @@ use codex_tools::ToolSearchInfo;
use codex_tools::ToolSearchSourceInfo;
use codex_tools::ToolSpec;
use codex_tools::coalesce_loadable_tool_specs;
use codex_tools::flatten_responses_api_namespace;
pub struct ToolSearchHandler {
entries: Vec<ToolSearchEntry>,
namespace_tool_spec_mode: NamespaceToolSpecMode,
search_source_infos: Vec<ToolSearchSourceInfo>,
search_engine: SearchEngine<usize>,
}
#[derive(Clone, Copy)]
pub(crate) enum NamespaceToolSpecMode {
Preserve,
Flatten,
}
impl ToolSearchHandler {
pub(crate) fn new(search_infos: Vec<ToolSearchInfo>) -> Self {
pub(crate) fn new(
search_infos: Vec<ToolSearchInfo>,
namespace_tool_spec_mode: NamespaceToolSpecMode,
) -> Self {
let mut entries = Vec::with_capacity(search_infos.len());
let mut search_source_infos = Vec::new();
for search_info in search_infos {
@@ -47,6 +58,7 @@ impl ToolSearchHandler {
Self {
entries,
namespace_tool_spec_mode,
search_source_infos,
search_engine,
}
@@ -132,9 +144,27 @@ impl ToolSearchHandler {
&self,
results: impl IntoIterator<Item = &'a ToolSearchEntry>,
) -> Result<Vec<LoadableToolSpec>, FunctionCallError> {
Ok(coalesce_loadable_tool_specs(
results.into_iter().map(|entry| entry.output.clone()),
))
let tools =
coalesce_loadable_tool_specs(results.into_iter().map(|entry| entry.output.clone()));
if matches!(
self.namespace_tool_spec_mode,
NamespaceToolSpecMode::Preserve
) {
return Ok(tools);
}
Ok(tools
.into_iter()
.flat_map(|tool| match tool {
LoadableToolSpec::Function(tool) => vec![LoadableToolSpec::Function(tool)],
LoadableToolSpec::Namespace(namespace) => {
flatten_responses_api_namespace(namespace)
.into_iter()
.map(LoadableToolSpec::Function)
.collect()
}
})
.collect())
}
}
@@ -187,7 +217,7 @@ mod tests {
.search_info()
.expect("dynamic handler should return search info")
}));
let handler = ToolSearchHandler::new(search_infos);
let handler = ToolSearchHandler::new(search_infos, NamespaceToolSpecMode::Preserve);
let results = [
&handler.entries[0],
&handler.entries[2],
@@ -255,6 +285,35 @@ mod tests {
);
}
#[test]
fn search_results_flatten_mcp_namespaces_without_namespace_tools() {
let search_info = McpHandler::new(tool_info("calendar", "create_event", "Create events"))
.expect("MCP tool should convert")
.search_info()
.expect("MCP handler should return search info");
let handler = ToolSearchHandler::new(vec![search_info], NamespaceToolSpecMode::Flatten);
let tools = handler
.search_output_tools([&handler.entries[0]])
.expect("MCP search output should serialize");
assert_eq!(
tools,
vec![LoadableToolSpec::Function(ResponsesApiTool {
name: "mcp__calendar__create_event".to_string(),
description: "Create events desktop tool".to_string(),
strict: false,
defer_loading: Some(true),
parameters: codex_tools::JsonSchema::object(
Default::default(),
/*required*/ None,
Some(false.into()),
),
output_schema: None,
})],
);
}
fn tool_info(server_name: &str, tool_name: &str, description_prefix: &str) -> ToolInfo {
ToolInfo {
server_name: server_name.to_string(),
+21 -26
View File
@@ -14,6 +14,7 @@ use crate::tools::handlers::ListAvailablePluginsToInstallHandler;
use crate::tools::handlers::ListMcpResourceTemplatesHandler;
use crate::tools::handlers::ListMcpResourcesHandler;
use crate::tools::handlers::McpHandler;
use crate::tools::handlers::NamespaceToolSpecMode;
use crate::tools::handlers::NewContextWindowHandler;
use crate::tools::handlers::PlanHandler;
use crate::tools::handlers::ReadMcpResourceHandler;
@@ -81,6 +82,7 @@ use codex_tools::can_request_original_image_detail;
use codex_tools::collect_code_mode_exec_prompt_tool_definitions;
use codex_tools::collect_request_plugin_install_entries;
use codex_tools::default_namespace_description;
use codex_tools::flatten_responses_api_namespace;
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;
@@ -230,7 +232,10 @@ fn build_model_visible_specs_and_registry(
merged_specs
.into_iter()
.flat_map(|spec| match spec {
ToolSpec::Namespace(namespace) => flatten_namespace_spec(namespace),
ToolSpec::Namespace(namespace) => flatten_responses_api_namespace(namespace)
.into_iter()
.map(ToolSpec::Function)
.collect(),
other => vec![other],
})
.collect()
@@ -239,22 +244,6 @@ fn build_model_visible_specs_and_registry(
(model_visible_specs, registry)
}
/// Flattens a namespace into plain `function` specs named `<namespace>__<tool>`.
/// The canonical names match hooks and proxies, so the registry resolves them.
fn flatten_namespace_spec(namespace: ResponsesApiNamespace) -> Vec<ToolSpec> {
namespace
.tools
.into_iter()
.map(|tool| match tool {
ResponsesApiNamespaceTool::Function(mut function) => {
let tool_name = ToolName::namespaced(namespace.name.as_str(), function.name);
function.name = tool_name.canonical_flat_name().into_owned();
ToolSpec::Function(function)
}
})
.collect()
}
fn spec_for_model_request(
turn_context: &TurnContext,
exposure: ToolExposure,
@@ -778,12 +767,11 @@ fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mu
} else {
let agent_type_description =
agent_type_description(turn_context, context.default_agent_type_description);
let exposure =
if search_tool_enabled(turn_context) && namespace_tools_enabled(turn_context) {
ToolExposure::Deferred
} else {
ToolExposure::Direct
};
let exposure = if search_tool_enabled(turn_context) {
ToolExposure::Deferred
} else {
ToolExposure::Direct
};
planned_tools.add_with_exposure(
SpawnAgentHandler::new(SpawnAgentToolOptions {
available_models: turn_context.available_models.clone(),
@@ -866,7 +854,7 @@ fn append_tool_search_executor(
planned_tools: &mut PlannedTools,
) {
let turn_context = context.turn_context;
if !(search_tool_enabled(turn_context) && namespace_tools_enabled(turn_context)) {
if !search_tool_enabled(turn_context) {
return;
}
@@ -880,7 +868,15 @@ fn append_tool_search_executor(
return;
}
planned_tools.add(ToolSearchHandler::new(search_infos));
let namespace_tool_spec_mode = if namespace_tools_enabled(turn_context) {
NamespaceToolSpecMode::Preserve
} else {
NamespaceToolSpecMode::Flatten
};
planned_tools.add(ToolSearchHandler::new(
search_infos,
namespace_tool_spec_mode,
));
}
fn prepend_code_mode_executors(
@@ -914,7 +910,6 @@ fn append_extension_tool_executors(
reserved_tool_names.insert(ToolName::plain(codex_code_mode::WAIT_TOOL_NAME));
}
if search_tool_enabled(turn_context)
&& namespace_tools_enabled(turn_context)
&& planned_tools
.runtimes()
.iter()
+20 -2
View File
@@ -692,13 +692,14 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() {
deferred_mcp_tools: Some(vec![mcp_tool("searchable", "mcp__searchable", "lookup")]),
..ToolPlanInputs::default()
};
let searchable_deferred_mcp_tools = searchable_mcp.deferred_mcp_tools.clone();
let missing_model_capability = probe_with(
|turn| {
turn.model_info.supports_search_tool = false;
},
ToolPlanInputs {
deferred_mcp_tools: searchable_mcp.deferred_mcp_tools.clone(),
deferred_mcp_tools: searchable_deferred_mcp_tools.clone(),
..ToolPlanInputs::default()
},
)
@@ -723,7 +724,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(),
deferred_mcp_tools: searchable_deferred_mcp_tools.clone(),
..ToolPlanInputs::default()
},
)
@@ -742,6 +743,23 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() {
"tool_search",
&ToolName::namespaced("mcp__searchable", "lookup").to_string(),
]);
let namespace_tools_disabled = probe_with(
|turn| {
turn.model_info.supports_search_tool = true;
disable_namespace_tools(turn);
},
ToolPlanInputs {
deferred_mcp_tools: searchable_deferred_mcp_tools,
..ToolPlanInputs::default()
},
)
.await;
namespace_tools_disabled.assert_visible_contains(&["tool_search"]);
namespace_tools_disabled.assert_registered_contains(&[
"tool_search",
&ToolName::namespaced("mcp__searchable", "lookup").to_string(),
]);
}
#[tokio::test]
+2 -4
View File
@@ -894,16 +894,14 @@ async fn stdio_mcp_tool_call_flattens_namespace_when_provider_disables_it() -> a
tool.get("type").and_then(Value::as_str) == Some("function")
&& tool.get("name").and_then(Value::as_str) == Some(flat_tool_name.as_str())
}),
"MCP tool should be sent as a flat function: {:?}",
request_body
"MCP tool should be sent as a flat function: {request_body:?}"
);
assert!(
tools.iter().all(|tool| {
tool.get("type").and_then(Value::as_str) != Some("namespace")
|| tool.get("name").and_then(Value::as_str) != Some("mcp__rmcp")
}),
"MCP namespace wrapper should be omitted: {:?}",
request_body
"MCP namespace wrapper should be omitted: {request_body:?}"
);
let output = final_mock
+1
View File
@@ -59,6 +59,7 @@ pub use responses_api::ResponsesApiTool;
pub use responses_api::coalesce_loadable_tool_specs;
pub use responses_api::default_namespace_description;
pub use responses_api::dynamic_tool_to_responses_api_tool;
pub use responses_api::flatten_responses_api_namespace;
pub use responses_api::mcp_tool_to_deferred_responses_api_tool;
pub use responses_api::mcp_tool_to_responses_api_tool;
pub use responses_api::tool_definition_to_responses_api_tool;
+16
View File
@@ -104,6 +104,22 @@ pub fn coalesce_loadable_tool_specs(
coalesced_specs
}
/// Flattens a namespace into plain Responses API functions named
/// `<namespace>__<tool>`.
pub fn flatten_responses_api_namespace(namespace: ResponsesApiNamespace) -> Vec<ResponsesApiTool> {
namespace
.tools
.into_iter()
.map(|tool| match tool {
ResponsesApiNamespaceTool::Function(mut function) => {
let tool_name = ToolName::namespaced(namespace.name.as_str(), function.name);
function.name = tool_name.canonical_flat_name().into_owned();
function
}
})
.collect()
}
pub fn mcp_tool_to_responses_api_tool(
tool_name: &ToolName,
tool: &rmcp::model::Tool,