Support tool search for dynamic tools

This commit is contained in:
Sayan Sisodiya
2026-04-16 23:44:07 +08:00
parent 54af119326
commit f1e0ee2de5
3 changed files with 361 additions and 80 deletions

View File

@@ -4,6 +4,7 @@ use app_test_support::McpProcess;
use app_test_support::create_final_assistant_message_sse_response;
use app_test_support::create_mock_responses_server_sequence_unchecked;
use app_test_support::to_response;
use app_test_support::write_models_cache_with_models;
use codex_app_server_protocol::DynamicToolCallOutputContentItem;
use codex_app_server_protocol::DynamicToolCallParams;
use codex_app_server_protocol::DynamicToolCallResponse;
@@ -21,6 +22,7 @@ use codex_app_server_protocol::ThreadStartResponse;
use codex_app_server_protocol::TurnStartParams;
use codex_app_server_protocol::TurnStartResponse;
use codex_app_server_protocol::UserInput as V2UserInput;
use codex_models_manager::bundled_models_response;
use codex_protocol::models::FunctionCallOutputBody;
use codex_protocol::models::FunctionCallOutputContentItem;
use codex_protocol::models::FunctionCallOutputPayload;
@@ -28,6 +30,7 @@ use core_test_support::responses;
use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;
use std::io::Write;
use std::path::Path;
use std::time::Duration;
use tempfile::TempDir;
@@ -196,6 +199,178 @@ async fn thread_start_keeps_hidden_dynamic_tools_out_of_model_requests() -> Resu
Ok(())
}
/// Exercises deferred dynamic tool discovery, the follow-up tool call, and the tool response.
#[tokio::test]
async fn deferred_dynamic_tool_can_be_discovered_and_called_through_tool_search() -> Result<()> {
let search_call_id = "tool-search-1";
let dynamic_call_id = "dyn-search-call-1";
let tool_name = "automation_update";
let tool_description = "Create, update, view, or delete recurring automations.";
let tool_args = json!({ "mode": "create" });
let tool_call_arguments = serde_json::to_string(&tool_args)?;
let responses = vec![
responses::sse(vec![
responses::ev_response_created("resp-1"),
responses::ev_tool_search_call(
search_call_id,
&json!({
"query": "recurring automations",
"limit": 8,
}),
),
responses::ev_completed("resp-1"),
]),
responses::sse(vec![
responses::ev_response_created("resp-2"),
responses::ev_function_call(dynamic_call_id, tool_name, &tool_call_arguments),
responses::ev_completed("resp-2"),
]),
create_final_assistant_message_sse_response("Done")?,
];
let server = create_mock_responses_server_sequence_unchecked(responses).await;
let codex_home = TempDir::new()?;
write_search_capable_models_cache(codex_home.path())?;
create_config_toml(codex_home.path(), &server.uri())?;
enable_tool_search_feature(codex_home.path())?;
let mut mcp = McpProcess::new(codex_home.path()).await?;
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
let input_schema = json!({
"type": "object",
"properties": {
"mode": { "type": "string" }
},
"required": ["mode"],
"additionalProperties": false,
});
let dynamic_tool = DynamicToolSpec {
name: tool_name.to_string(),
description: tool_description.to_string(),
input_schema: input_schema.clone(),
defer_loading: true,
};
let thread_req = mcp
.send_thread_start_request(ThreadStartParams {
dynamic_tools: Some(vec![dynamic_tool]),
..Default::default()
})
.await?;
let thread_resp: JSONRPCResponse = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(thread_req)),
)
.await??;
let ThreadStartResponse { thread, .. } = to_response::<ThreadStartResponse>(thread_resp)?;
let thread_id = thread.id.clone();
let turn_req = mcp
.send_turn_start_request(TurnStartParams {
thread_id: thread_id.clone(),
input: vec![V2UserInput::Text {
text: "Use the automation tool".to_string(),
text_elements: Vec::new(),
}],
..Default::default()
})
.await?;
let turn_resp: JSONRPCResponse = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(turn_req)),
)
.await??;
let TurnStartResponse { turn } = to_response::<TurnStartResponse>(turn_resp)?;
let turn_id = turn.id.clone();
let started = wait_for_dynamic_tool_started(&mut mcp, dynamic_call_id).await?;
assert_eq!(started.thread_id, thread_id.clone());
assert_eq!(started.turn_id, turn_id.clone());
let request = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_request_message(),
)
.await??;
let (request_id, params) = match request {
ServerRequest::DynamicToolCall { request_id, params } => (request_id, params),
other => panic!("expected DynamicToolCall request, got {other:?}"),
};
assert_eq!(
params,
DynamicToolCallParams {
thread_id: thread_id.clone(),
turn_id: turn_id.clone(),
call_id: dynamic_call_id.to_string(),
tool: tool_name.to_string(),
arguments: tool_args.clone(),
}
);
mcp.send_response(
request_id,
serde_json::to_value(DynamicToolCallResponse {
content_items: vec![DynamicToolCallOutputContentItem::InputText {
text: "dynamic-search-ok".to_string(),
}],
success: true,
})?,
)
.await?;
let completed = wait_for_dynamic_tool_completed(&mut mcp, dynamic_call_id).await?;
assert_eq!(completed.thread_id, thread_id);
assert_eq!(completed.turn_id, turn_id);
timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_notification_message("turn/completed"),
)
.await??;
let bodies = responses_bodies(&server).await?;
let first_request = bodies
.first()
.context("expected an initial responses request")?;
assert!(
find_tool_by_type(first_request, "tool_search").is_some(),
"initial request should advertise tool_search: {first_request:?}"
);
assert!(
find_tool(first_request, tool_name).is_none(),
"deferred dynamic tool should not be directly advertised before search"
);
let search_tools = bodies
.iter()
.find_map(|body| tool_search_output_tools(body, search_call_id))
.context("expected tool_search_output in follow-up request")?;
assert_eq!(
search_tools,
vec![json!({
"type": "function",
"name": tool_name,
"description": tool_description,
"strict": false,
"defer_loading": true,
"parameters": input_schema,
})]
);
let payload = bodies
.iter()
.find_map(|body| function_call_output_payload(body, dynamic_call_id))
.context("expected function_call_output in post-tool follow-up request")?;
assert_eq!(
payload,
FunctionCallOutputPayload::from_text("dynamic-search-ok".to_string())
);
Ok(())
}
/// Exercises the full dynamic tool call path (server request, client response, model output).
#[tokio::test]
async fn dynamic_tool_call_round_trip_sends_text_content_items_to_model() -> Result<()> {
@@ -583,6 +758,30 @@ fn find_tool<'a>(body: &'a Value, name: &str) -> Option<&'a Value> {
})
}
fn find_tool_by_type<'a>(body: &'a Value, tool_type: &str) -> Option<&'a Value> {
body.get("tools")
.and_then(Value::as_array)
.and_then(|tools| {
tools
.iter()
.find(|tool| tool.get("type").and_then(Value::as_str) == Some(tool_type))
})
}
fn tool_search_output_tools(body: &Value, call_id: &str) -> Option<Vec<Value>> {
body.get("input")
.and_then(Value::as_array)
.and_then(|items| {
items.iter().find(|item| {
item.get("type").and_then(Value::as_str) == Some("tool_search_output")
&& item.get("call_id").and_then(Value::as_str) == Some(call_id)
})
})
.and_then(|item| item.get("tools"))
.and_then(Value::as_array)
.cloned()
}
fn function_call_output_payload(body: &Value, call_id: &str) -> Option<FunctionCallOutputPayload> {
function_call_output_raw_output(body, call_id)
.and_then(|output| serde_json::from_value(output).ok())
@@ -663,3 +862,24 @@ stream_max_retries = 0
),
)
}
fn enable_tool_search_feature(codex_home: &Path) -> std::io::Result<()> {
let mut config_toml = std::fs::OpenOptions::new()
.append(true)
.open(codex_home.join("config.toml"))?;
config_toml.write_all(b"\n[features]\ntool_search = true\n")
}
fn write_search_capable_models_cache(codex_home: &Path) -> Result<()> {
let mut model = bundled_models_response()
.context("bundled models should parse")?
.models
.into_iter()
.find(|model| model.slug == "gpt-5.4")
.context("expected bundled gpt-5.4 model")?;
model.slug = "mock-model".to_string();
model.display_name = "mock-model".to_string();
model.supports_search_tool = true;
write_models_cache_with_models(codex_home, vec![model])?;
Ok(())
}

View File

@@ -10,12 +10,13 @@ use bm25::SearchEngine;
use bm25::SearchEngineBuilder;
use codex_mcp::ToolInfo;
use codex_protocol::dynamic_tools::DynamicToolSpec;
use codex_tools::ResponsesApiNamespace;
use codex_tools::ResponsesApiNamespaceTool;
use codex_tools::TOOL_SEARCH_DEFAULT_LIMIT;
use codex_tools::TOOL_SEARCH_TOOL_NAME;
use codex_tools::ToolSearchOutputTool;
use codex_tools::ToolSearchResultSource;
use codex_tools::collect_tool_search_output_tools;
use codex_tools::dynamic_tool_to_responses_api_tool;
use codex_tools::mcp_tool_to_deferred_responses_api_tool;
const COMPUTER_USE_MCP_SERVER_NAME: &str = "computer-use";
const COMPUTER_USE_TOOL_SEARCH_LIMIT: usize = 20;
@@ -31,9 +32,9 @@ impl ToolSearchHandler {
dynamic_tools: Vec<DynamicToolSpec>,
) -> Self {
let mut mcp_entries: Vec<ToolSearchEntry> = mcp_tools
.into_iter()
.map(|(name, info)| ToolSearchEntry::Mcp {
name,
.into_values()
.map(|info| ToolSearchEntry::Mcp {
name: info.canonical_tool_name().display(),
info: Box::new(info),
})
.collect();
@@ -118,39 +119,7 @@ impl ToolSearchHandler {
use_default_limit: bool,
) -> Result<Vec<ToolSearchOutputTool>, FunctionCallError> {
let results = self.search_result_entries(query, limit, use_default_limit);
let mut tools = Vec::new();
let mut pending_mcp_sources = Vec::new();
for entry in results {
match entry {
ToolSearchEntry::Mcp { info, .. } => {
pending_mcp_sources.push(ToolSearchResultSource {
server_name: info.server_name.as_str(),
tool_namespace: info.callable_namespace.as_str(),
tool_name: info.callable_name.as_str(),
tool: &info.tool,
connector_name: info.connector_name.as_deref(),
connector_description: info.connector_description.as_deref(),
});
}
ToolSearchEntry::Dynamic { tool } => {
tools.extend(
collect_tool_search_output_tools(pending_mcp_sources.drain(..))
.map_err(tool_search_output_error)?,
);
tools.push(ToolSearchOutputTool::Function(
dynamic_tool_to_responses_api_tool(tool)
.map_err(tool_search_output_error)?,
));
}
}
}
tools.extend(
collect_tool_search_output_tools(pending_mcp_sources)
.map_err(tool_search_output_error)?,
);
Ok(tools)
search_output_tools(results)
}
fn search_result_entries(
@@ -184,6 +153,60 @@ impl ToolSearchHandler {
}
}
fn search_output_tools<'a>(
results: impl IntoIterator<Item = &'a ToolSearchEntry>,
) -> Result<Vec<ToolSearchOutputTool>, FunctionCallError> {
let mut tools = Vec::new();
// Preserve search order: group MCP tools under namespaces, emit dynamic tools directly.
for entry in results {
match entry {
ToolSearchEntry::Mcp { info, .. } => {
let tool_name = info.canonical_tool_name();
let namespace = info.callable_namespace.as_str();
let namespace_tool =
mcp_tool_to_deferred_responses_api_tool(&tool_name, &info.tool)
.map(ResponsesApiNamespaceTool::Function)
.map_err(tool_search_output_error)?;
if let Some(output) = tools.iter_mut().find_map(|tool| match tool {
ToolSearchOutputTool::Namespace(output) if output.name == namespace => {
Some(output)
}
ToolSearchOutputTool::Namespace(_) | ToolSearchOutputTool::Function(_) => None,
}) {
output.tools.push(namespace_tool);
} else {
tools.push(ToolSearchOutputTool::Namespace(ResponsesApiNamespace {
name: namespace.to_string(),
description: mcp_namespace_description(info),
tools: vec![namespace_tool],
}));
}
}
ToolSearchEntry::Dynamic { tool } => {
tools.push(ToolSearchOutputTool::Function(
dynamic_tool_to_responses_api_tool(tool).map_err(tool_search_output_error)?,
));
}
}
}
Ok(tools)
}
fn mcp_namespace_description(info: &ToolInfo) -> String {
info.connector_description
.clone()
.or_else(|| {
info.connector_name
.as_deref()
.map(str::trim)
.filter(|connector_name| !connector_name.is_empty())
.map(|connector_name| format!("Tools for working with {connector_name}."))
})
.unwrap_or_else(|| format!("Tools from the {} MCP server.", info.server_name))
}
fn limit_results_per_server(results: Vec<&ToolSearchEntry>) -> Vec<&ToolSearchEntry> {
results
.into_iter()
@@ -322,57 +345,95 @@ fn build_dynamic_search_text(tool: &DynamicToolSpec) -> String {
#[cfg(test)]
mod tests {
use super::*;
use codex_tools::JsonSchema;
use codex_tools::ResponsesApiNamespace;
use codex_tools::ResponsesApiNamespaceTool;
use codex_tools::ResponsesApiTool;
use pretty_assertions::assert_eq;
use rmcp::model::Tool;
use std::sync::Arc;
#[test]
fn search_returns_deferred_dynamic_tools() {
let handler = ToolSearchHandler::new(
std::collections::HashMap::new(),
vec![DynamicToolSpec {
name: "automation_update".to_string(),
description: "Create, update, view, or delete recurring automations.".to_string(),
input_schema: serde_json::json!({
"type": "object",
"properties": {
"mode": { "type": "string" },
"kind": { "type": "string" }
},
"required": ["mode"],
"additionalProperties": false,
}),
defer_loading: true,
}],
);
fn mixed_search_results_coalesce_mcp_namespaces() {
let entries = [
ToolSearchEntry::Mcp {
name: "mcp__calendar__create_event".to_string(),
info: Box::new(tool_info("calendar", "create_event", "Create events")),
},
ToolSearchEntry::Dynamic {
tool: DynamicToolSpec {
name: "automation_update".to_string(),
description: "Create, update, view, or delete recurring automations."
.to_string(),
input_schema: serde_json::json!({
"type": "object",
"properties": {
"mode": { "type": "string" },
},
"required": ["mode"],
"additionalProperties": false,
}),
defer_loading: true,
},
},
ToolSearchEntry::Mcp {
name: "mcp__calendar__list_events".to_string(),
info: Box::new(tool_info("calendar", "list_events", "List events")),
},
];
let tools = handler
.search(
"automation_update",
TOOL_SEARCH_DEFAULT_LIMIT,
/*use_default_limit*/ true,
)
.expect("search deferred dynamic tools");
let tools =
search_output_tools(entries.iter()).expect("mixed search output should serialize");
assert_eq!(
tools,
vec![ToolSearchOutputTool::Function(ResponsesApiTool {
name: "automation_update".to_string(),
description: "Create, update, view, or delete recurring automations.".to_string(),
strict: false,
defer_loading: Some(true),
parameters: JsonSchema::object(
std::collections::BTreeMap::from([
("kind".to_string(), JsonSchema::string(None)),
("mode".to_string(), JsonSchema::string(None)),
]),
Some(vec!["mode".to_string()]),
Some(false.into()),
),
output_schema: None,
})],
vec![
ToolSearchOutputTool::Namespace(ResponsesApiNamespace {
name: "mcp__calendar__".to_string(),
description: "Tools from the calendar MCP server.".to_string(),
tools: vec![
ResponsesApiNamespaceTool::Function(ResponsesApiTool {
name: "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,
}),
ResponsesApiNamespaceTool::Function(ResponsesApiTool {
name: "list_events".to_string(),
description: "List 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,
}),
],
}),
ToolSearchOutputTool::Function(ResponsesApiTool {
name: "automation_update".to_string(),
description: "Create, update, view, or delete recurring automations."
.to_string(),
strict: false,
defer_loading: Some(true),
parameters: codex_tools::JsonSchema::object(
std::collections::BTreeMap::from([(
"mode".to_string(),
codex_tools::JsonSchema::string(None),
)]),
Some(vec!["mode".to_string()]),
Some(false.into()),
),
output_schema: None,
}),
],
);
}

View File

@@ -177,7 +177,7 @@ async fn search_tool_enabled_by_default_adds_tool_search() -> Result<()> {
"parameters": {
"type": "object",
"properties": {
"query": {"type": "string", "description": "Search query for MCP tools."},
"query": {"type": "string", "description": "Search query for deferred tools."},
"limit": {"type": "number", "description": "Maximum number of tools to return (defaults to 8)."},
},
"required": ["query"],