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:
Anton Panasenko
2026-02-09 12:53:50 -08:00
committed by GitHub
Unverified
parent 9450cd9ce5
commit becc3a0424
15 changed files with 1238 additions and 1 deletions
+104
View File
@@ -29,6 +29,7 @@ pub(crate) struct SessionState {
pub(crate) pending_resume_previous_model: Option<String>,
/// Startup regular task pre-created during session initialization.
pub(crate) startup_regular_task: Option<RegularTask>,
pub(crate) active_mcp_tool_selection: Option<Vec<String>>,
}
impl SessionState {
@@ -45,6 +46,7 @@ impl SessionState {
initial_context_seeded: false,
pending_resume_previous_model: None,
startup_regular_task: None,
active_mcp_tool_selection: None,
}
}
@@ -140,6 +142,32 @@ impl SessionState {
pub(crate) fn take_startup_regular_task(&mut self) -> Option<RegularTask> {
self.startup_regular_task.take()
}
pub(crate) fn merge_mcp_tool_selection(&mut self, tool_names: Vec<String>) -> Vec<String> {
if tool_names.is_empty() {
return self.active_mcp_tool_selection.clone().unwrap_or_default();
}
let mut merged = self.active_mcp_tool_selection.take().unwrap_or_default();
let mut seen: HashSet<String> = merged.iter().cloned().collect();
for tool_name in tool_names {
if seen.insert(tool_name.clone()) {
merged.push(tool_name);
}
}
self.active_mcp_tool_selection = Some(merged.clone());
merged
}
pub(crate) fn get_mcp_tool_selection(&self) -> Option<Vec<String>> {
self.active_mcp_tool_selection.clone()
}
pub(crate) fn clear_mcp_tool_selection(&mut self) {
self.active_mcp_tool_selection = None;
}
}
// Sometimes new snapshots don't include credits or plan information.
@@ -155,3 +183,79 @@ fn merge_rate_limit_fields(
}
snapshot
}
#[cfg(test)]
mod tests {
use super::*;
use crate::codex::make_session_configuration_for_tests;
use pretty_assertions::assert_eq;
#[tokio::test]
async fn merge_mcp_tool_selection_deduplicates_and_preserves_order() {
let session_configuration = make_session_configuration_for_tests().await;
let mut state = SessionState::new(session_configuration);
let merged = state.merge_mcp_tool_selection(vec![
"mcp__rmcp__echo".to_string(),
"mcp__rmcp__image".to_string(),
"mcp__rmcp__echo".to_string(),
]);
assert_eq!(
merged,
vec![
"mcp__rmcp__echo".to_string(),
"mcp__rmcp__image".to_string(),
]
);
let merged = state.merge_mcp_tool_selection(vec![
"mcp__rmcp__image".to_string(),
"mcp__rmcp__search".to_string(),
]);
assert_eq!(
merged,
vec![
"mcp__rmcp__echo".to_string(),
"mcp__rmcp__image".to_string(),
"mcp__rmcp__search".to_string(),
]
);
}
#[tokio::test]
async fn merge_mcp_tool_selection_empty_input_is_noop() {
let session_configuration = make_session_configuration_for_tests().await;
let mut state = SessionState::new(session_configuration);
state.merge_mcp_tool_selection(vec![
"mcp__rmcp__echo".to_string(),
"mcp__rmcp__image".to_string(),
]);
let merged = state.merge_mcp_tool_selection(Vec::new());
assert_eq!(
merged,
vec![
"mcp__rmcp__echo".to_string(),
"mcp__rmcp__image".to_string(),
]
);
assert_eq!(
state.get_mcp_tool_selection(),
Some(vec![
"mcp__rmcp__echo".to_string(),
"mcp__rmcp__image".to_string(),
])
);
}
#[tokio::test]
async fn clear_mcp_tool_selection_removes_selection() {
let session_configuration = make_session_configuration_for_tests().await;
let mut state = SessionState::new(session_configuration);
state.merge_mcp_tool_selection(vec!["mcp__rmcp__echo".to_string()]);
state.clear_mcp_tool_selection();
assert_eq!(state.get_mcp_tool_selection(), None);
}
}