feat: extend active tools to lifecycle of turn

Author: apanasenko-oai

codex-rs/core/src/codex.rs

@@ -1107,14 +1107,19 @@ impl Session {
         }
     }
 
-    pub(crate) async fn set_next_mcp_tool_selection(&self, tool_names: Vec<String>) {
+    pub(crate) async fn merge_mcp_tool_selection(&self, tool_names: Vec<String>) -> Vec<String> {
         let mut state = self.state.lock().await;
-        state.set_next_mcp_tool_selection(tool_names);
+        state.merge_mcp_tool_selection(tool_names)
     }
 
-    pub(crate) async fn take_next_mcp_tool_selection(&self) -> Option<Vec<String>> {
+    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.take_next_mcp_tool_selection()
+        state.clear_mcp_tool_selection();
     }
 
     async fn record_initial_history(&self, conversation_history: InitialHistory) {
@@ -3842,7 +3847,7 @@ async fn run_sampling_request(
         .await?;
     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 {
+        if let Some(selected_tools) = sess.get_mcp_tool_selection().await {
             mcp_tools = filter_mcp_tools_by_name(mcp_tools, &selected_tools);
         } else {
             mcp_tools.clear();
@@ -4676,6 +4681,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 {
@@ -5484,6 +5491,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");

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

@@ -24,7 +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>>,
+    pub(crate) active_mcp_tool_selection: Option<Vec<String>>,
 }
 
 impl SessionState {
@@ -39,7 +39,7 @@ impl SessionState {
             dependency_env: HashMap::new(),
             mcp_dependency_prompted: HashSet::new(),
             initial_context_seeded: false,
-            next_mcp_tool_selection: None,
+            active_mcp_tool_selection: None,
         }
     }
 
@@ -128,12 +128,30 @@ impl SessionState {
         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 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 take_next_mcp_tool_selection(&mut self) -> Option<Vec<String>> {
-        self.next_mcp_tool_selection.take()
+    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;
     }
 }
 
@@ -150,3 +168,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);
+    }
+}

codex-rs/core/src/tasks/mod.rs

@@ -117,6 +117,7 @@ impl Session {
         task: T,
     ) {
         self.abort_all_tasks(TurnAbortReason::Replaced).await;
+        self.clear_mcp_tool_selection().await;
         self.seed_initial_context_if_needed(turn_context.as_ref())
             .await;
 

codex-rs/core/src/tools/handlers/search_tool_bm25.rs

@@ -114,11 +114,11 @@ impl ToolHandler for SearchToolBm25Handler {
         entries.sort_by(|a, b| a.name.cmp(&b.name));
 
         if entries.is_empty() {
-            session.set_next_mcp_tool_selection(Vec::new()).await;
+            let active_selected_tools = session.get_mcp_tool_selection().await.unwrap_or_default();
             let content = json!({
                 "query": query,
                 "total_tools": 0,
-                "selected_tools": [],
+                "active_selected_tools": active_selected_tools,
                 "tools": [],
             })
             .to_string();
@@ -157,14 +157,12 @@ impl ToolHandler for SearchToolBm25Handler {
             }));
         }
 
-        session
-            .set_next_mcp_tool_selection(selected_tools.clone())
-            .await;
+        let active_selected_tools = session.merge_mcp_tool_selection(selected_tools).await;
 
         let content = json!({
             "query": query,
             "total_tools": entries.len(),
-            "selected_tools": selected_tools,
+            "active_selected_tools": active_selected_tools,
             "tools": result_payloads,
         })
         .to_string();

codex-rs/core/templates/search_tool/developer_instructions.md

@@ -8,8 +8,9 @@ Follow this workflow:
    - `query` (required): focused terms that describe the capability you need.
    - `limit` (optional): maximum number of tools to return (default `8`, max `50`).
 2. Use the returned `tools` list to decide which MCP tools are relevant.
-3. On the next request, only the returned `selected_tools` will be available. Invoke the MCP tool(s) you need there.
-4. MCP tool selections are consumed after that next request. Search again if you need different MCP tools.
+3. Matching tools are added to `active_selected_tools`. Only tools in `active_selected_tools` are available for the remainder of the current turn.
+4. Repeated searches in the same turn are additive: new matches are unioned into `active_selected_tools`.
+5. `active_selected_tools` resets at the start of the next turn.
 
 Notes:
 - Core tools remain available without searching.