feat: persist and restore codex app's tools after search (#11780)

### What changed
1. Removed per-turn MCP selection reset in `core/src/tasks/mod.rs`.
2. Added `SessionState::set_mcp_tool_selection(Vec<String>)` in
`core/src/state/session.rs` for authoritative restore behavior (deduped,
order-preserving, empty clears).
3. Added rollout parsing in `core/src/codex.rs` to recover
`active_selected_tools` from prior `search_tool_bm25` outputs:
   - tracks matching `call_id`s
   - parses function output text JSON
   - extracts `active_selected_tools`
   - latest valid payload wins
   - malformed/non-matching payloads are ignored
4. Applied restore logic to resumed and forked startup paths in
`core/src/codex.rs`.
5. Updated instruction text to session/thread scope in
`core/templates/search_tool/tool_description.md`.
6. Expanded tests in `core/tests/suite/search_tool.rs`, plus unit
coverage in:
   - `core/src/codex.rs`
   - `core/src/state/session.rs`

### Behavior after change
1. Search activates matched tools.
2. Additional searches union into active selection.
3. Selection survives new turns in the same thread.
4. Resume/fork restores selection from rollout history.
5. Separate threads do not inherit selection unless forked.
This commit is contained in:
Anton Panasenko
2026-02-15 19:18:41 -08:00
committed by GitHub
parent 060a320e7d
commit 02abd9a8ea
12 changed files with 1201 additions and 201 deletions

View File

@@ -33,6 +33,7 @@ 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(crate) use search_tool_bm25::SEARCH_TOOL_BM25_TOOL_NAME;
pub use search_tool_bm25::SearchToolBm25Handler;
pub use shell::ShellCommandHandler;
pub use shell::ShellHandler;

View File

@@ -10,7 +10,6 @@ use std::collections::HashMap;
use std::collections::HashSet;
use crate::connectors;
use crate::features::Feature;
use crate::function_tool::FunctionCallError;
use crate::mcp::CODEX_APPS_MCP_SERVER_NAME;
use crate::mcp_connection_manager::ToolInfo;
@@ -23,6 +22,7 @@ use crate::tools::registry::ToolKind;
pub struct SearchToolBm25Handler;
pub(crate) const SEARCH_TOOL_BM25_TOOL_NAME: &str = "search_tool_bm25";
pub(crate) const DEFAULT_LIMIT: usize = 8;
fn default_limit() -> usize {
@@ -42,7 +42,6 @@ struct ToolEntry {
server_name: String,
title: Option<String>,
description: Option<String>,
connector_id: Option<String>,
connector_name: Option<String>,
input_keys: Vec<String>,
search_text: String,
@@ -66,7 +65,6 @@ impl ToolEntry {
.tool
.description
.map(|description| description.to_string()),
connector_id: info.connector_id,
connector_name: info.connector_name,
input_keys,
search_text,
@@ -91,9 +89,9 @@ impl ToolHandler for SearchToolBm25Handler {
let arguments = match payload {
ToolPayload::Function { arguments } => arguments,
_ => {
return Err(FunctionCallError::Fatal(
"search_tool_bm25 handler received unsupported payload".to_string(),
));
return Err(FunctionCallError::Fatal(format!(
"{SEARCH_TOOL_BM25_TOOL_NAME} handler received unsupported payload"
)));
}
};
@@ -120,15 +118,12 @@ impl ToolHandler for SearchToolBm25Handler {
.await
.list_all_tools()
.await;
let mcp_tools = if turn.config.features.enabled(Feature::Apps) {
let connectors = connectors::with_app_enabled_state(
connectors::accessible_connectors_from_mcp_tools(&mcp_tools),
&turn.config,
);
filter_codex_apps_mcp_tools(mcp_tools, &connectors)
} else {
mcp_tools
};
let connectors = connectors::with_app_enabled_state(
connectors::accessible_connectors_from_mcp_tools(&mcp_tools),
&turn.config,
);
let mcp_tools = filter_codex_apps_mcp_tools(mcp_tools, &connectors);
let mut entries: Vec<ToolEntry> = mcp_tools
.into_iter()
@@ -172,7 +167,6 @@ impl ToolHandler for SearchToolBm25Handler {
"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,
@@ -243,12 +237,6 @@ fn build_search_text(name: &str, info: &ToolInfo, input_keys: &[String]) -> Stri
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());
}

View File

@@ -8,6 +8,7 @@ use crate::features::Features;
use crate::mcp_connection_manager::ToolInfo;
use crate::tools::handlers::PLAN_TOOL;
use crate::tools::handlers::SEARCH_TOOL_BM25_DEFAULT_LIMIT;
use crate::tools::handlers::SEARCH_TOOL_BM25_TOOL_NAME;
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;
@@ -913,7 +914,7 @@ fn create_search_tool_bm25_tool(app_tools: &HashMap<String, ToolInfo>) -> ToolSp
SEARCH_TOOL_BM25_DESCRIPTION_TEMPLATE.replace("{{app_names}}", app_names.as_str());
ToolSpec::Function(ResponsesApiTool {
name: "search_tool_bm25".to_string(),
name: SEARCH_TOOL_BM25_TOOL_NAME.to_string(),
description,
strict: false,
parameters: JsonSchema::Object {
@@ -1507,7 +1508,7 @@ pub(crate) fn build_specs(
&& let Some(app_tools) = app_tools
{
builder.push_spec_with_parallel_support(create_search_tool_bm25_tool(&app_tools), true);
builder.register_handler("search_tool_bm25", search_tool_handler);
builder.register_handler(SEARCH_TOOL_BM25_TOOL_NAME, search_tool_handler);
}
if let Some(apply_patch_tool_type) = &config.apply_patch_tool_type {
@@ -2579,7 +2580,7 @@ mod tests {
)
.build();
let search_tool = find_tool(&tools, "search_tool_bm25");
let search_tool = find_tool(&tools, SEARCH_TOOL_BM25_TOOL_NAME);
let ToolSpec::Function(ResponsesApiTool { description, .. }) = &search_tool.spec else {
panic!("expected function tool");
};