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
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user