mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
feat: search_tool (#10657)
**Why We Did This** - The goal is to reduce MCP tool context pollution by not exposing the full MCP tool list up front - It forces an explicit discovery step (`search_tool_bm25`) so the model narrows tool scope before making MCP calls, which helps relevance and lowers prompt/tool clutter. **What It Changed** - Added a new experimental feature flag `search_tool` in `core/src/features.rs:90` and `core/src/features.rs:430`. - Added config/schema support for that flag in `core/config.schema.json:214` and `core/config.schema.json:1235`. - Added BM25 dependency (`bm25`) in `Cargo.toml:129` and `core/Cargo.toml:23`. - Added new tool handler `search_tool_bm25` in `core/src/tools/handlers/search_tool_bm25.rs:18`. - Registered the handler and tool spec in `core/src/tools/handlers/mod.rs:11` and `core/src/tools/spec.rs:780` and `core/src/tools/spec.rs:1344`. - Extended `ToolsConfig` to carry `search_tool` enablement in `core/src/tools/spec.rs:32` and `core/src/tools/spec.rs:56`. - Injected dedicated developer instructions for tool-discovery workflow in `core/src/codex.rs:483` and `core/src/codex.rs:1976`, using `core/templates/search_tool/developer_instructions.md:1`. - Added session state to store one-shot selected MCP tools in `core/src/state/session.rs:27` and `core/src/state/session.rs:131`. - Added filtering so when feature is enabled, only selected MCP tools are exposed on the next request (then consumed) in `core/src/codex.rs:3800` and `core/src/codex.rs:3843`. - Added E2E suite coverage for enablement/instructions/hide-until-search/one-turn-selection in `core/tests/suite/search_tool.rs:72`, `core/tests/suite/search_tool.rs:109`, `core/tests/suite/search_tool.rs:147`, and `core/tests/suite/search_tool.rs:218`. - Refactored test helper utilities to support config-driven tool collection in `core/tests/suite/tools.rs:281`. **Net Behavioral Effect** - With `search_tool` **off**: existing MCP behavior (tools exposed normally). - With `search_tool` **on**: MCP tools start hidden, model must call `search_tool_bm25`, and only returned `selected_tools` are available for the next model call.
This commit is contained in:
committed by
GitHub
Unverified
parent
9450cd9ce5
commit
becc3a0424
@@ -8,6 +8,7 @@ mod mcp_resource;
|
||||
mod plan;
|
||||
mod read_file;
|
||||
mod request_user_input;
|
||||
mod search_tool_bm25;
|
||||
mod shell;
|
||||
mod test_sync;
|
||||
mod unified_exec;
|
||||
@@ -28,6 +29,8 @@ pub use plan::PlanHandler;
|
||||
pub use read_file::ReadFileHandler;
|
||||
pub use request_user_input::RequestUserInputHandler;
|
||||
pub(crate) use request_user_input::request_user_input_tool_description;
|
||||
pub(crate) use search_tool_bm25::DEFAULT_LIMIT as SEARCH_TOOL_BM25_DEFAULT_LIMIT;
|
||||
pub use search_tool_bm25::SearchToolBm25Handler;
|
||||
pub use shell::ShellCommandHandler;
|
||||
pub use shell::ShellHandler;
|
||||
pub use test_sync::TestSyncHandler;
|
||||
|
||||
@@ -0,0 +1,217 @@
|
||||
use async_trait::async_trait;
|
||||
use bm25::Document;
|
||||
use bm25::Language;
|
||||
use bm25::SearchEngineBuilder;
|
||||
use codex_protocol::models::FunctionCallOutputBody;
|
||||
use serde::Deserialize;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::function_tool::FunctionCallError;
|
||||
use crate::mcp_connection_manager::ToolInfo;
|
||||
use crate::tools::context::ToolInvocation;
|
||||
use crate::tools::context::ToolOutput;
|
||||
use crate::tools::context::ToolPayload;
|
||||
use crate::tools::handlers::parse_arguments;
|
||||
use crate::tools::registry::ToolHandler;
|
||||
use crate::tools::registry::ToolKind;
|
||||
|
||||
pub struct SearchToolBm25Handler;
|
||||
|
||||
pub(crate) const DEFAULT_LIMIT: usize = 8;
|
||||
|
||||
fn default_limit() -> usize {
|
||||
DEFAULT_LIMIT
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct SearchToolBm25Args {
|
||||
query: String,
|
||||
#[serde(default = "default_limit")]
|
||||
limit: usize,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct ToolEntry {
|
||||
name: String,
|
||||
server_name: String,
|
||||
title: Option<String>,
|
||||
description: Option<String>,
|
||||
connector_id: Option<String>,
|
||||
connector_name: Option<String>,
|
||||
input_keys: Vec<String>,
|
||||
search_text: String,
|
||||
}
|
||||
|
||||
impl ToolEntry {
|
||||
fn new(name: String, info: ToolInfo) -> Self {
|
||||
let input_keys = info
|
||||
.tool
|
||||
.input_schema
|
||||
.get("properties")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.map(|map| map.keys().cloned().collect::<Vec<_>>())
|
||||
.unwrap_or_default();
|
||||
let search_text = build_search_text(&name, &info, &input_keys);
|
||||
Self {
|
||||
name,
|
||||
server_name: info.server_name,
|
||||
title: info.tool.title,
|
||||
description: info
|
||||
.tool
|
||||
.description
|
||||
.map(|description| description.to_string()),
|
||||
connector_id: info.connector_id,
|
||||
connector_name: info.connector_name,
|
||||
input_keys,
|
||||
search_text,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolHandler for SearchToolBm25Handler {
|
||||
fn kind(&self) -> ToolKind {
|
||||
ToolKind::Function
|
||||
}
|
||||
|
||||
async fn handle(&self, invocation: ToolInvocation) -> Result<ToolOutput, FunctionCallError> {
|
||||
let ToolInvocation {
|
||||
payload, session, ..
|
||||
} = invocation;
|
||||
|
||||
let arguments = match payload {
|
||||
ToolPayload::Function { arguments } => arguments,
|
||||
_ => {
|
||||
return Err(FunctionCallError::Fatal(
|
||||
"search_tool_bm25 handler received unsupported payload".to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
let args: SearchToolBm25Args = parse_arguments(&arguments)?;
|
||||
let query = args.query.trim();
|
||||
if query.is_empty() {
|
||||
return Err(FunctionCallError::RespondToModel(
|
||||
"query must not be empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
if args.limit == 0 {
|
||||
return Err(FunctionCallError::RespondToModel(
|
||||
"limit must be greater than zero".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let limit = args.limit;
|
||||
|
||||
let mcp_tools = session
|
||||
.services
|
||||
.mcp_connection_manager
|
||||
.read()
|
||||
.await
|
||||
.list_all_tools()
|
||||
.await;
|
||||
|
||||
let mut entries: Vec<ToolEntry> = mcp_tools
|
||||
.into_iter()
|
||||
.map(|(name, info)| ToolEntry::new(name, info))
|
||||
.collect();
|
||||
entries.sort_by(|a, b| a.name.cmp(&b.name));
|
||||
|
||||
if entries.is_empty() {
|
||||
let active_selected_tools = session.get_mcp_tool_selection().await.unwrap_or_default();
|
||||
let content = json!({
|
||||
"query": query,
|
||||
"total_tools": 0,
|
||||
"active_selected_tools": active_selected_tools,
|
||||
"tools": [],
|
||||
})
|
||||
.to_string();
|
||||
return Ok(ToolOutput::Function {
|
||||
body: FunctionCallOutputBody::Text(content),
|
||||
success: Some(true),
|
||||
});
|
||||
}
|
||||
|
||||
let documents: Vec<Document<usize>> = entries
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(idx, entry)| Document::new(idx, entry.search_text.clone()))
|
||||
.collect();
|
||||
let search_engine =
|
||||
SearchEngineBuilder::<usize>::with_documents(Language::English, documents).build();
|
||||
let results = search_engine.search(query, limit);
|
||||
|
||||
let mut selected_tools = Vec::new();
|
||||
let mut result_payloads = Vec::new();
|
||||
for result in results {
|
||||
let Some(entry) = entries.get(result.document.id) else {
|
||||
continue;
|
||||
};
|
||||
selected_tools.push(entry.name.clone());
|
||||
result_payloads.push(json!({
|
||||
"name": entry.name.clone(),
|
||||
"server": entry.server_name.clone(),
|
||||
"title": entry.title.clone(),
|
||||
"description": entry.description.clone(),
|
||||
"connector_id": entry.connector_id.clone(),
|
||||
"connector_name": entry.connector_name.clone(),
|
||||
"input_keys": entry.input_keys.clone(),
|
||||
"score": result.score,
|
||||
}));
|
||||
}
|
||||
|
||||
let active_selected_tools = session.merge_mcp_tool_selection(selected_tools).await;
|
||||
|
||||
let content = json!({
|
||||
"query": query,
|
||||
"total_tools": entries.len(),
|
||||
"active_selected_tools": active_selected_tools,
|
||||
"tools": result_payloads,
|
||||
})
|
||||
.to_string();
|
||||
|
||||
Ok(ToolOutput::Function {
|
||||
body: FunctionCallOutputBody::Text(content),
|
||||
success: Some(true),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn build_search_text(name: &str, info: &ToolInfo, input_keys: &[String]) -> String {
|
||||
let mut parts = vec![
|
||||
name.to_string(),
|
||||
info.tool_name.clone(),
|
||||
info.server_name.clone(),
|
||||
];
|
||||
|
||||
if let Some(title) = info.tool.title.as_deref()
|
||||
&& !title.trim().is_empty()
|
||||
{
|
||||
parts.push(title.to_string());
|
||||
}
|
||||
|
||||
if let Some(description) = info.tool.description.as_deref()
|
||||
&& !description.trim().is_empty()
|
||||
{
|
||||
parts.push(description.to_string());
|
||||
}
|
||||
|
||||
if let Some(connector_name) = info.connector_name.as_deref()
|
||||
&& !connector_name.trim().is_empty()
|
||||
{
|
||||
parts.push(connector_name.to_string());
|
||||
}
|
||||
|
||||
if let Some(connector_id) = info.connector_id.as_deref()
|
||||
&& !connector_id.trim().is_empty()
|
||||
{
|
||||
parts.push(connector_id.to_string());
|
||||
}
|
||||
|
||||
if !input_keys.is_empty() {
|
||||
parts.extend(input_keys.iter().cloned());
|
||||
}
|
||||
|
||||
parts.join(" ")
|
||||
}
|
||||
@@ -4,6 +4,7 @@ use crate::client_common::tools::ToolSpec;
|
||||
use crate::features::Feature;
|
||||
use crate::features::Features;
|
||||
use crate::tools::handlers::PLAN_TOOL;
|
||||
use crate::tools::handlers::SEARCH_TOOL_BM25_DEFAULT_LIMIT;
|
||||
use crate::tools::handlers::apply_patch::create_apply_patch_freeform_tool;
|
||||
use crate::tools::handlers::apply_patch::create_apply_patch_json_tool;
|
||||
use crate::tools::handlers::collab::DEFAULT_WAIT_TIMEOUT_MS;
|
||||
@@ -31,6 +32,7 @@ pub(crate) struct ToolsConfig {
|
||||
pub apply_patch_tool_type: Option<ApplyPatchToolType>,
|
||||
pub web_search_mode: Option<WebSearchMode>,
|
||||
pub supports_image_input: bool,
|
||||
pub search_tool: bool,
|
||||
pub collab_tools: bool,
|
||||
pub collaboration_modes_tools: bool,
|
||||
pub request_rule_enabled: bool,
|
||||
@@ -54,6 +56,7 @@ impl ToolsConfig {
|
||||
let include_collab_tools = features.enabled(Feature::Collab);
|
||||
let include_collaboration_modes_tools = features.enabled(Feature::CollaborationModes);
|
||||
let request_rule_enabled = features.enabled(Feature::RequestRule);
|
||||
let include_search_tool = features.enabled(Feature::SearchTool);
|
||||
|
||||
let shell_type = if !features.enabled(Feature::ShellTool) {
|
||||
ConfigShellToolType::Disabled
|
||||
@@ -85,6 +88,7 @@ impl ToolsConfig {
|
||||
apply_patch_tool_type,
|
||||
web_search_mode: *web_search_mode,
|
||||
supports_image_input: model_info.input_modalities.contains(&InputModality::Image),
|
||||
search_tool: include_search_tool,
|
||||
collab_tools: include_collab_tools,
|
||||
collaboration_modes_tools: include_collaboration_modes_tools,
|
||||
request_rule_enabled,
|
||||
@@ -800,6 +804,36 @@ fn create_grep_files_tool() -> ToolSpec {
|
||||
})
|
||||
}
|
||||
|
||||
fn create_search_tool_bm25_tool() -> ToolSpec {
|
||||
let properties = BTreeMap::from([
|
||||
(
|
||||
"query".to_string(),
|
||||
JsonSchema::String {
|
||||
description: Some("Search query for MCP tools.".to_string()),
|
||||
},
|
||||
),
|
||||
(
|
||||
"limit".to_string(),
|
||||
JsonSchema::Number {
|
||||
description: Some(format!(
|
||||
"Maximum number of tools to return (defaults to {SEARCH_TOOL_BM25_DEFAULT_LIMIT})."
|
||||
)),
|
||||
},
|
||||
),
|
||||
]);
|
||||
|
||||
ToolSpec::Function(ResponsesApiTool {
|
||||
name: "search_tool_bm25".to_string(),
|
||||
description: "Searches MCP tool metadata with BM25 and exposes matching tools for the next model call.".to_string(),
|
||||
strict: false,
|
||||
parameters: JsonSchema::Object {
|
||||
properties,
|
||||
required: Some(vec!["query".to_string()]),
|
||||
additional_properties: Some(false.into()),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
fn create_read_file_tool() -> ToolSpec {
|
||||
let indentation_properties = BTreeMap::from([
|
||||
(
|
||||
@@ -1261,6 +1295,7 @@ pub(crate) fn build_specs(
|
||||
use crate::tools::handlers::PlanHandler;
|
||||
use crate::tools::handlers::ReadFileHandler;
|
||||
use crate::tools::handlers::RequestUserInputHandler;
|
||||
use crate::tools::handlers::SearchToolBm25Handler;
|
||||
use crate::tools::handlers::ShellCommandHandler;
|
||||
use crate::tools::handlers::ShellHandler;
|
||||
use crate::tools::handlers::TestSyncHandler;
|
||||
@@ -1280,6 +1315,7 @@ pub(crate) fn build_specs(
|
||||
let mcp_resource_handler = Arc::new(McpResourceHandler);
|
||||
let shell_command_handler = Arc::new(ShellCommandHandler);
|
||||
let request_user_input_handler = Arc::new(RequestUserInputHandler);
|
||||
let search_tool_handler = Arc::new(SearchToolBm25Handler);
|
||||
|
||||
match &config.shell_type {
|
||||
ConfigShellToolType::Default => {
|
||||
@@ -1334,6 +1370,11 @@ pub(crate) fn build_specs(
|
||||
builder.register_handler("request_user_input", request_user_input_handler);
|
||||
}
|
||||
|
||||
if config.search_tool {
|
||||
builder.push_spec_with_parallel_support(create_search_tool_bm25_tool(), true);
|
||||
builder.register_handler("search_tool_bm25", search_tool_handler);
|
||||
}
|
||||
|
||||
if let Some(apply_patch_tool_type) = &config.apply_patch_tool_type {
|
||||
match apply_patch_tool_type {
|
||||
ApplyPatchToolType::Freeform => {
|
||||
|
||||
Reference in New Issue
Block a user