openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-02 19:43:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

Share MCP tool specs until search results are selected (#45509)

Browse Source 
## Why

Building MCP search entries eagerly cloned tool specs and normalized schemas even for tools that were never selected.

## What changed

Store search specs in `Arc<ToolSpec>` and let MCP search entries share the handler's spec. Materialize and normalize loadable specs only for selected results, preserving existing result formatting and dynamic-tool cache equality behavior.

## Testing

Add coverage for function, freeform, and namespace specs that verifies shared specs produce equivalent results, retain the source while needed, and release it when the search entry is dropped.

GitOrigin-RevId: 4260f2e2527834d8a856b1528654c6951aab8a7d
This commit is contained in:
Owen Lin
2026-09-14 20:20:29 +00:00
committed by copyberry
parent b9bfc0aff8
commit ea3c4848d8
5 changed files with 145 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/tools/handlers/mcp.rs
+2 -2
View File
@@ -157,9 +157,9 @@ impl ToolExecutor<ToolInvocation> for McpHandler {
.map(str::to_string),
});
ToolSearchInfo::from_spec(
ToolSearchInfo::from_shared_spec(
build_mcp_search_text(&self.tool_info),
self.spec(),
Arc::clone(&self.spec),
source_info,
)
}
codex-rs/core/src/tools/handlers/mcp_search_tests.rs
+4 -4
View File
@@ -29,7 +29,7 @@ fn search_info_uses_connector_name_for_output_namespace_description() {
let handler = McpHandler::new(tool_info).expect("MCP tool spec should build");
let search_info = handler.search_info().expect("MCP search info");
let LoadableToolSpec::Namespace(namespace) = search_info.entry.output else {
let LoadableToolSpec::Namespace(namespace) = search_info.entry.to_loadable_spec() else {
panic!("expected namespace search output");
};
assert_eq!(namespace.description, "Tools for working with Calendar.");
@@ -57,7 +57,7 @@ fn mcp_namespace_descriptions_preserve_complete_metadata() {
description: Some(full_description.clone()),
})
);
let LoadableToolSpec::Namespace(namespace) = search_info.entry.output else {
let LoadableToolSpec::Namespace(namespace) = search_info.entry.to_loadable_spec() else {
panic!("expected namespace search output");
};
assert_eq!(namespace.description, full_description);
@@ -83,7 +83,7 @@ fn mcp_namespace_descriptions_are_bounded_at_512_kib() {
description: Some(full_description),
})
);
let LoadableToolSpec::Namespace(namespace) = search_info.entry.output else {
let LoadableToolSpec::Namespace(namespace) = search_info.entry.to_loadable_spec() else {
panic!("expected namespace search output");
};
assert_eq!(namespace.description, expected_description);
@@ -104,7 +104,7 @@ fn agent_plugin_namespace_descriptions_use_the_stricter_bound() {
description: Some(expected_description.clone()),
})
);
let LoadableToolSpec::Namespace(namespace) = search_info.entry.output else {
let LoadableToolSpec::Namespace(namespace) = search_info.entry.to_loadable_spec() else {
panic!("expected namespace search output");
};
assert_eq!(namespace.description, expected_description);
codex-rs/core/src/tools/handlers/tool_search.rs
+1 -1
View File
@@ -250,7 +250,7 @@ impl ToolSearchHandler {
results: impl IntoIterator<Item = &'a ToolSearchEntry>,
) -> Result<Vec<LoadableToolSpec>, FunctionCallError> {
Ok(coalesce_loadable_tool_specs(
results.into_iter().map(|entry| entry.output.clone()),
results.into_iter().map(ToolSearchEntry::to_loadable_spec),
))
}
}
codex-rs/tools/src/tool_search.rs
+66 -42
View File
@@ -7,11 +7,22 @@ use crate::ToolSearchSourceInfo;
use crate::ToolSpec;
use crate::default_namespace_description;
use codex_protocol::DEFAULT_FUNCTION_NAMESPACE;
use std::sync::Arc;
#[derive(Clone, PartialEq)]
pub struct ToolSearchEntry {
pub search_text: String,
pub output: LoadableToolSpec,
spec: Arc<ToolSpec>,
}
impl ToolSearchEntry {
/// Materialize only selected results; output schemas remain shared until discarded.
pub fn to_loadable_spec(&self) -> LoadableToolSpec {
let Some(output) = normalize_search_spec(self.spec.as_ref().clone()) else {
unreachable!("search entries contain only loadable tools");
};
output
}
}
#[derive(Clone, PartialEq)]
@@ -21,6 +32,21 @@ pub struct ToolSearchInfo {
}
impl ToolSearchInfo {
/// Keep the immutable catalog snapshot alive without duplicating its schemas.
pub fn from_shared_spec(
search_text: String,
spec: Arc<ToolSpec>,
source_info: Option<ToolSearchSourceInfo>,
) -> Option<Self> {
match spec.as_ref() {
ToolSpec::Function(_) | ToolSpec::Freeform(_) | ToolSpec::Namespace(_) => Some(Self {
entry: ToolSearchEntry { search_text, spec },
source_info,
}),
ToolSpec::ToolSearch { .. } | ToolSpec::WebSearch { .. } => None,
}
}
pub fn from_tool_spec(
spec: ToolSpec,
source_info: Option<ToolSearchSourceInfo>,
@@ -34,56 +60,54 @@ impl ToolSearchInfo {
spec: ToolSpec,
source_info: Option<ToolSearchSourceInfo>,
) -> Option<Self> {
let output = match spec {
ToolSpec::Function(mut tool) => {
tool.defer_loading = Some(true);
tool.output_schema = None;
LoadableToolSpec::Namespace(ResponsesApiNamespace {
name: DEFAULT_FUNCTION_NAMESPACE.to_string(),
description: default_namespace_description(DEFAULT_FUNCTION_NAMESPACE),
tools: vec![ResponsesApiNamespaceTool::Function(tool)],
})
}
ToolSpec::Freeform(mut tool) => {
tool.defer_loading = Some(true);
LoadableToolSpec::Namespace(ResponsesApiNamespace {
name: DEFAULT_FUNCTION_NAMESPACE.to_string(),
description: default_namespace_description(DEFAULT_FUNCTION_NAMESPACE),
tools: vec![ResponsesApiNamespaceTool::Custom(tool)],
})
}
ToolSpec::Namespace(mut namespace) => {
if namespace.description.trim().is_empty() {
namespace.description = default_namespace_description(&namespace.name);
}
for tool in &mut namespace.tools {
match tool {
ResponsesApiNamespaceTool::Function(tool) => {
tool.defer_loading = Some(true);
tool.output_schema = None;
}
ResponsesApiNamespaceTool::Custom(tool) => {
tool.defer_loading = Some(true);
}
}
}
LoadableToolSpec::Namespace(namespace)
}
ToolSpec::ToolSearch { .. } | ToolSpec::WebSearch { .. } => {
return None;
}
};
// Dynamic-tool cache entries compare normalized specs by value.
// Preserve that behavior; shared MCP specs normalize only when selected.
let output = normalize_search_spec(spec)?;
Some(Self {
entry: ToolSearchEntry {
search_text,
output,
spec: Arc::new(output.into()),
},
source_info,
})
}
}
fn normalize_search_spec(spec: ToolSpec) -> Option<LoadableToolSpec> {
let mut namespace = match spec {
ToolSpec::Function(tool) => ResponsesApiNamespace {
name: DEFAULT_FUNCTION_NAMESPACE.to_string(),
description: default_namespace_description(DEFAULT_FUNCTION_NAMESPACE),
tools: vec![ResponsesApiNamespaceTool::Function(tool)],
},
ToolSpec::Freeform(tool) => ResponsesApiNamespace {
name: DEFAULT_FUNCTION_NAMESPACE.to_string(),
description: default_namespace_description(DEFAULT_FUNCTION_NAMESPACE),
tools: vec![ResponsesApiNamespaceTool::Custom(tool)],
},
ToolSpec::Namespace(namespace) => namespace,
ToolSpec::ToolSearch { .. } | ToolSpec::WebSearch { .. } => {
return None;
}
};
if namespace.description.trim().is_empty() {
namespace.description = default_namespace_description(&namespace.name);
}
for tool in &mut namespace.tools {
match tool {
ResponsesApiNamespaceTool::Function(tool) => {
tool.defer_loading = Some(true);
tool.output_schema = None;
}
ResponsesApiNamespaceTool::Custom(tool) => {
tool.defer_loading = Some(true);
}
}
}
Some(LoadableToolSpec::Namespace(namespace))
}
fn default_tool_search_text(spec: &ToolSpec) -> String {
let mut parts = Vec::new();
codex-rs/tools/src/tool_search_tests.rs
+72 -3
View File
@@ -2,6 +2,69 @@ use super::*;
use pretty_assertions::assert_eq;
use std::collections::BTreeMap;
#[test]
fn shared_search_specs_preserve_results_and_release_the_source() {
let function = ResponsesApiTool {
name: "lookup".to_string(),
description: "Look up an item".to_string(),
strict: true,
defer_loading: None,
parameters: JsonSchema::object(
BTreeMap::new(),
/*required*/ None,
/*additional_properties*/ None,
),
output_schema: Some(serde_json::json!({"type": "object"}).into()),
};
let custom = crate::FreeformTool {
name: "patch".to_string(),
description: "Apply a patch".to_string(),
defer_loading: None,
format: crate::FreeformToolFormat {
r#type: "grammar".to_string(),
syntax: "lark".to_string(),
definition: "start: \"patch\"".to_string(),
},
};
for spec in [
ToolSpec::Function(function.clone()),
ToolSpec::Freeform(custom.clone()),
ToolSpec::Namespace(ResponsesApiNamespace {
name: "example".to_string(),
description: String::new(),
tools: vec![
ResponsesApiNamespaceTool::Function(function),
ResponsesApiNamespaceTool::Custom(custom),
],
}),
] {
let expected =
ToolSearchInfo::from_spec("query".to_string(), spec.clone(), /*source_info*/ None)
.unwrap()
.entry
.to_loadable_spec();
let spec = Arc::new(spec);
let weak = Arc::downgrade(&spec);
let info = ToolSearchInfo::from_shared_spec(
"query".to_string(),
Arc::clone(&spec),
/*source_info*/ None,
)
.unwrap();
drop(spec);
assert!(
weak.upgrade().is_some(),
"search retains the original spec rather than a deep copy"
);
assert_eq!(info.entry.to_loadable_spec(), expected);
drop(info);
assert!(
weak.upgrade().is_none(),
"search does not keep discarded catalogs alive"
);
}
}
#[test]
fn top_level_function_search_results_use_the_default_namespace() {
let function_tool = ResponsesApiTool {
@@ -23,7 +86,10 @@ fn top_level_function_search_results_use_the_default_namespace() {
.expect("top-level function should be searchable");
assert_eq!(
(search_info.entry.search_text, search_info.entry.output),
(
search_info.entry.search_text.clone(),
search_info.entry.to_loadable_spec()
),
(
"lookup_order lookup order Look up an order".to_string(),
LoadableToolSpec::Namespace(ResponsesApiNamespace {
@@ -58,7 +124,10 @@ fn top_level_custom_tools_are_searchable() {
.expect("top-level custom tool should be searchable");
assert_eq!(
(search_info.entry.search_text, search_info.entry.output),
(
search_info.entry.search_text.clone(),
search_info.entry.to_loadable_spec()
),
(
"apply_patch Apply a patch lark".to_string(),
LoadableToolSpec::Namespace(ResponsesApiNamespace {
@@ -159,7 +228,7 @@ fn mixed_namespaced_function_and_custom_tools_are_searchable() {
"editor Editing tools lookup_order lookup order Look up an order apply_patch Apply a patch lark"
);
assert_eq!(
search_info.entry.output,
search_info.entry.to_loadable_spec(),
LoadableToolSpec::Namespace(crate::ResponsesApiNamespace {
name: "editor".to_string(),
description: "Editing tools".to_string(),