feat: search_tool

Author: apanasenko-oai

codex-rs/Cargo.lock

@@ -129,6 +129,7 @@ async-stream = "0.3.6"
 async-trait = "0.1.89"
 axum = { version = "0.8", default-features = false }
 base64 = "0.22.1"
+bm25 = "2.3.2"
 bytes = "1.10.1"
 chardetng = "0.1.17"
 chrono = "0.4.43"

codex-rs/Cargo.toml

@@ -23,6 +23,7 @@ arc-swap = "1.8.0"
 async-channel = { workspace = true }
 async-trait = { workspace = true }
 base64 = { workspace = true }
+bm25 = { workspace = true }
 chardetng = { workspace = true }
 chrono = { workspace = true, features = ["serde"] }
 clap = { workspace = true, features = ["derive"] }

codex-rs/core/Cargo.toml

@@ -211,6 +211,9 @@
             "runtime_metrics": {
               "type": "boolean"
             },
+            "search_tool": {
+              "type": "boolean"
+            },
             "shell_snapshot": {
               "type": "boolean"
             },
@@ -1229,6 +1232,9 @@
         "runtime_metrics": {
           "type": "boolean"
         },
+        "search_tool": {
+          "type": "boolean"
+        },
         "shell_snapshot": {
           "type": "boolean"
         },

codex-rs/core/config.schema.json

@@ -480,6 +480,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 {
@@ -1104,6 +1107,16 @@ impl Session {
         }
     }
 
+    pub(crate) async fn set_next_mcp_tool_selection(&self, tool_names: Vec<String>) {
+        let mut state = self.state.lock().await;
+        state.set_next_mcp_tool_selection(tool_names);
+    }
+
+    pub(crate) async fn take_next_mcp_tool_selection(&self) -> Option<Vec<String>> {
+        let mut state = self.state.lock().await;
+        state.take_next_mcp_tool_selection()
+    }
+
     async fn record_initial_history(&self, conversation_history: InitialHistory) {
         let turn_context = self.new_default_turn().await;
         match conversation_history {
@@ -1945,7 +1958,7 @@ impl Session {
         &self,
         turn_context: &TurnContext,
     ) -> Vec<ResponseItem> {
-        let mut items = Vec::<ResponseItem>::with_capacity(4);
+        let mut items = Vec::<ResponseItem>::with_capacity(6);
         let shell = self.user_shell();
         items.push(
             DeveloperInstructions::from_policy(
@@ -1960,6 +1973,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;
@@ -3779,6 +3797,15 @@ fn filter_codex_apps_mcp_tools(
     mcp_tools
 }
 
+fn filter_mcp_tools_by_name(
+    mut mcp_tools: HashMap<String, crate::mcp_connection_manager::ToolInfo>,
+    selected_tools: &[String],
+) -> HashMap<String, crate::mcp_connection_manager::ToolInfo> {
+    let allowed: HashSet<&str> = selected_tools.iter().map(String::as_str).collect();
+    mcp_tools.retain(|name, _| allowed.contains(name.as_str()));
+    mcp_tools
+}
+
 fn codex_apps_connector_id(tool: &crate::mcp_connection_manager::ToolInfo) -> Option<&str> {
     tool.connector_id.as_deref()
 }
@@ -3813,19 +3840,31 @@ 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(
-            connectors,
-            &input,
-            tool_selection.explicit_app_paths,
-            tool_selection.skill_name_counts_lower,
-        ))
+    let search_tool_enabled = turn_context
+        .config
+        .features
+        .enabled(Feature::SearchTool);
+    if search_tool_enabled {
+        if let Some(selected_tools) = sess.take_next_mcp_tool_selection().await {
+            mcp_tools = filter_mcp_tools_by_name(mcp_tools, &selected_tools);
+        } else {
+            mcp_tools.clear();
+        }
     } else {
-        None
-    };
-    if let Some(connectors) = connectors_for_tools.as_ref() {
-        mcp_tools = filter_codex_apps_mcp_tools(mcp_tools, connectors);
+        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(
+                connectors,
+                &input,
+                tool_selection.explicit_app_paths,
+                tool_selection.skill_name_counts_lower,
+            ))
+        } else {
+            None
+        };
+        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,

codex-rs/core/src/codex.rs

@@ -87,6 +87,8 @@ pub enum Feature {
     /// Allow the model to request web searches that fetch cached content.
     /// Takes precedence over `WebSearchRequest`.
     WebSearchCached,
+    /// Allow the model to search MCP tools via BM25 before exposing them.
+    SearchTool,
     /// Gate the execpolicy enforcement for shell/unified exec.
     ExecPolicy,
     /// Allow the model to request approval and propose exec rules.
@@ -424,6 +426,16 @@ pub const FEATURES: &[FeatureSpec] = &[
         stage: Stage::Deprecated,
         default_enabled: false,
     },
+    FeatureSpec {
+        id: Feature::SearchTool,
+        key: "search_tool",
+        stage: Stage::Experimental {
+            name: "Tool search",
+            menu_description: "Search MCP tools with BM25 before exposing them.",
+            announcement: "NEW! Try tool search to reduce MCP tool noise. Enable in /experimental!",
+        },
+        default_enabled: false,
+    },
     // Experimental program. Rendered in the `/experimental` menu for users.
     FeatureSpec {
         id: Feature::ShellSnapshot,

codex-rs/core/src/features.rs

@@ -24,6 +24,7 @@ pub(crate) struct SessionState {
     /// TODO(owen): This is a temporary solution to avoid updating a thread's updated_at
     /// timestamp when resuming a session. Remove this once SQLite is in place.
     pub(crate) initial_context_seeded: bool,
+    pub(crate) next_mcp_tool_selection: Option<Vec<String>>,
 }
 
 impl SessionState {
@@ -38,6 +39,7 @@ impl SessionState {
             dependency_env: HashMap::new(),
             mcp_dependency_prompted: HashSet::new(),
             initial_context_seeded: false,
+            next_mcp_tool_selection: None,
         }
     }
 
@@ -125,6 +127,14 @@ impl SessionState {
     pub(crate) fn dependency_env(&self) -> HashMap<String, String> {
         self.dependency_env.clone()
     }
+
+    pub(crate) fn set_next_mcp_tool_selection(&mut self, tool_names: Vec<String>) {
+        self.next_mcp_tool_selection = Some(tool_names);
+    }
+
+    pub(crate) fn take_next_mcp_tool_selection(&mut self) -> Option<Vec<String>> {
+        self.next_mcp_tool_selection.take()
+    }
 }
 
 // Sometimes new snapshots don't include credits or plan information.

codex-rs/core/src/state/session.rs

@@ -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,7 @@ 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 use search_tool_bm25::SearchToolBm25Handler;
 pub use shell::ShellCommandHandler;
 pub use shell::ShellHandler;
 pub use test_sync::TestSyncHandler;