agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: search_tool (#10657)

Browse Source 
**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
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/codex.rs
+199 -1
View File
@@ -139,6 +139,8 @@ use crate::mcp::effective_mcp_servers;
use crate::mcp::maybe_prompt_and_install_mcp_dependencies;
use crate::mcp::with_codex_apps_mcp;
use crate::mcp_connection_manager::McpConnectionManager;
use crate::mcp_connection_manager::filter_codex_apps_mcp_tools_only;
use crate::mcp_connection_manager::filter_mcp_tools_by_name;
use crate::mentions::build_connector_slug_counts;
use crate::mentions::build_skill_name_counts;
use crate::mentions::collect_explicit_app_paths;
@@ -504,6 +506,9 @@ pub(crate) struct Session {
next_internal_sub_id: AtomicU64,
}
const SEARCH_TOOL_DEVELOPER_INSTRUCTIONS: &str =
include_str!("../templates/search_tool/developer_instructions.md");
/// The context needed for a single turn of the thread.
#[derive(Debug)]
pub(crate) struct TurnContext {
@@ -1257,6 +1262,21 @@ impl Session {
}
}
pub(crate) async fn merge_mcp_tool_selection(&self, tool_names: Vec<String>) -> Vec<String> {
let mut state = self.state.lock().await;
state.merge_mcp_tool_selection(tool_names)
}
pub(crate) async fn get_mcp_tool_selection(&self) -> Option<Vec<String>> {
let state = self.state.lock().await;
state.get_mcp_tool_selection()
}
pub(crate) async fn clear_mcp_tool_selection(&self) {
let mut state = self.state.lock().await;
state.clear_mcp_tool_selection();
}
async fn record_initial_history(&self, conversation_history: InitialHistory) {
let turn_context = self.new_default_turn().await;
match conversation_history {
@@ -2182,6 +2202,11 @@ impl Session {
if let Some(developer_instructions) = turn_context.developer_instructions.as_deref() {
items.push(DeveloperInstructions::new(developer_instructions.to_string()).into());
}
if turn_context.tools_config.search_tool {
items.push(
DeveloperInstructions::new(SEARCH_TOOL_DEVELOPER_INSTRUCTIONS.to_string()).into(),
);
}
// Add developer instructions from collaboration_mode if they exist and are non-empty
let (collaboration_mode, base_instructions) = {
let state = self.state.lock().await;
@@ -4119,6 +4144,7 @@ async fn run_sampling_request(
.list_all_tools()
.or_cancel(&cancellation_token)
.await?;
let connectors_for_tools = if turn_context.config.features.enabled(Feature::Apps) {
let connectors = connectors::accessible_connectors_from_mcp_tools(&mcp_tools);
Some(filter_connectors_for_input(
@@ -4130,9 +4156,25 @@ async fn run_sampling_request(
} else {
None
};
if let Some(connectors) = connectors_for_tools.as_ref() {
if turn_context.config.features.enabled(Feature::SearchTool) {
let mut selected_mcp_tools =
if let Some(selected_tools) = sess.get_mcp_tool_selection().await {
filter_mcp_tools_by_name(mcp_tools.clone(), &selected_tools)
} else {
HashMap::new()
};
if let Some(connectors) = connectors_for_tools.as_ref() {
let apps_mcp_tools = filter_codex_apps_mcp_tools_only(mcp_tools, connectors);
selected_mcp_tools.extend(apps_mcp_tools);
}
mcp_tools = selected_mcp_tools;
} else if let Some(connectors) = connectors_for_tools.as_ref() {
mcp_tools = filter_codex_apps_mcp_tools(mcp_tools, connectors);
}
let router = Arc::new(ToolRouter::from_config(
&turn_context.tools_config,
Some(
@@ -4958,6 +5000,8 @@ pub(super) fn get_last_assistant_message_from_turn(responses: &[ResponseItem]) -
pub(crate) use tests::make_session_and_context;
#[cfg(test)]
pub(crate) use tests::make_session_and_context_with_rx;
#[cfg(test)]
pub(crate) use tests::make_session_configuration_for_tests;
#[cfg(test)]
mod tests {
@@ -4967,6 +5011,7 @@ mod tests {
use crate::config::test_config;
use crate::exec::ExecToolCallOutput;
use crate::function_tool::FunctionCallError;
use crate::mcp_connection_manager::ToolInfo;
use crate::shell::default_user_shell;
use crate::tools::format_exec_output_str;
@@ -5006,6 +5051,8 @@ mod tests {
use codex_protocol::mcp::CallToolResult as McpCallToolResult;
use pretty_assertions::assert_eq;
use rmcp::model::JsonObject;
use rmcp::model::Tool;
use serde::Deserialize;
use serde_json::json;
use std::path::PathBuf;
@@ -5042,6 +5089,30 @@ mod tests {
}
}
fn make_mcp_tool(
server_name: &str,
tool_name: &str,
connector_id: Option<&str>,
connector_name: Option<&str>,
) -> ToolInfo {
ToolInfo {
server_name: server_name.to_string(),
tool_name: tool_name.to_string(),
tool: Tool {
name: tool_name.to_string().into(),
title: None,
description: Some(format!("Test tool: {tool_name}").into()),
input_schema: Arc::new(JsonObject::default()),
output_schema: None,
annotations: None,
icons: None,
meta: None,
},
connector_id: connector_id.map(str::to_string),
connector_name: connector_name.map(str::to_string),
}
}
#[tokio::test]
async fn get_base_instructions_no_user_content() {
let prompt_with_apply_patch_instructions =
@@ -5145,6 +5216,93 @@ mod tests {
assert_eq!(selected, Vec::new());
}
#[test]
fn search_tool_selection_keeps_codex_apps_tools_without_mentions() {
let selected_tool_names = vec![
"mcp__codex_apps__calendar_create_event".to_string(),
"mcp__rmcp__echo".to_string(),
];
let mcp_tools = HashMap::from([
(
"mcp__codex_apps__calendar_create_event".to_string(),
make_mcp_tool(
CODEX_APPS_MCP_SERVER_NAME,
"calendar_create_event",
Some("calendar"),
Some("Calendar"),
),
),
(
"mcp__rmcp__echo".to_string(),
make_mcp_tool("rmcp", "echo", None, None),
),
]);
let mut selected_mcp_tools =
filter_mcp_tools_by_name(mcp_tools.clone(), &selected_tool_names);
let connectors = connectors::accessible_connectors_from_mcp_tools(&mcp_tools);
let connectors = filter_connectors_for_input(
connectors,
&[user_message("run the selected tools")],
&[],
&HashMap::new(),
);
let apps_mcp_tools = filter_codex_apps_mcp_tools_only(mcp_tools, &connectors);
selected_mcp_tools.extend(apps_mcp_tools);
let mut tool_names: Vec<String> = selected_mcp_tools.into_keys().collect();
tool_names.sort();
assert_eq!(
tool_names,
vec![
"mcp__codex_apps__calendar_create_event".to_string(),
"mcp__rmcp__echo".to_string(),
]
);
}
#[test]
fn apps_mentions_add_codex_apps_tools_to_search_selected_set() {
let selected_tool_names = vec!["mcp__rmcp__echo".to_string()];
let mcp_tools = HashMap::from([
(
"mcp__codex_apps__calendar_create_event".to_string(),
make_mcp_tool(
CODEX_APPS_MCP_SERVER_NAME,
"calendar_create_event",
Some("calendar"),
Some("Calendar"),
),
),
(
"mcp__rmcp__echo".to_string(),
make_mcp_tool("rmcp", "echo", None, None),
),
]);
let mut selected_mcp_tools =
filter_mcp_tools_by_name(mcp_tools.clone(), &selected_tool_names);
let connectors = connectors::accessible_connectors_from_mcp_tools(&mcp_tools);
let connectors = filter_connectors_for_input(
connectors,
&[user_message("use $calendar and then echo the response")],
&[],
&HashMap::new(),
);
let apps_mcp_tools = filter_codex_apps_mcp_tools_only(mcp_tools, &connectors);
selected_mcp_tools.extend(apps_mcp_tools);
let mut tool_names: Vec<String> = selected_mcp_tools.into_keys().collect();
tool_names.sort();
assert_eq!(
tool_names,
vec![
"mcp__codex_apps__calendar_create_event".to_string(),
"mcp__rmcp__echo".to_string(),
]
);
}
#[tokio::test]
async fn reconstruct_history_matches_live_compactions() {
let (session, turn_context) = make_session_and_context().await;
@@ -5849,6 +6007,46 @@ mod tests {
)
}
pub(crate) async fn make_session_configuration_for_tests() -> SessionConfiguration {
let codex_home = tempfile::tempdir().expect("create temp dir");
let config = build_test_config(codex_home.path()).await;
let config = Arc::new(config);
let model = ModelsManager::get_model_offline(config.model.as_deref());
let model_info = ModelsManager::construct_model_info_offline(model.as_str(), &config);
let reasoning_effort = config.model_reasoning_effort;
let collaboration_mode = CollaborationMode {
mode: ModeKind::Default,
settings: Settings {
model,
reasoning_effort,
developer_instructions: None,
},
};
SessionConfiguration {
provider: config.model_provider.clone(),
collaboration_mode,
model_reasoning_summary: config.model_reasoning_summary,
developer_instructions: config.developer_instructions.clone(),
user_instructions: config.user_instructions.clone(),
personality: config.personality,
base_instructions: config
.base_instructions
.clone()
.unwrap_or_else(|| model_info.get_model_instructions(config.personality)),
compact_prompt: config.compact_prompt.clone(),
approval_policy: config.approval_policy.clone(),
sandbox_policy: config.sandbox_policy.clone(),
windows_sandbox_level: WindowsSandboxLevel::from_config(&config),
cwd: config.cwd.clone(),
codex_home: config.codex_home.clone(),
thread_name: None,
original_config_do_not_use: Arc::clone(&config),
session_source: SessionSource::Exec,
dynamic_tools: Vec::new(),
}
}
pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
let (tx_event, _rx_event) = async_channel::unbounded();
let codex_home = tempfile::tempdir().expect("create temp dir");