From 43c8d9858fdf613e80e5318104b794669ebabcb1 Mon Sep 17 00:00:00 2001 From: Matthew Zeng Date: Mon, 13 Apr 2026 22:16:36 -0700 Subject: [PATCH] Keep advertised MCP tools sticky across requests --- codex-rs/core/src/codex.rs | 87 ++++++++++++---- codex-rs/core/src/codex_tests.rs | 101 ++++++++++++++++++ codex-rs/core/src/connectors.rs | 13 --- codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/mcp_tool_catalog.rs | 145 ++++++++++++++++++++++++++ codex-rs/core/src/state/session.rs | 21 ++++ 6 files changed, 334 insertions(+), 34 deletions(-) create mode 100644 codex-rs/core/src/mcp_tool_catalog.rs diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 0516f0cda4..48296316ad 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -277,6 +277,8 @@ use crate::injection::tool_kind_for_path; use crate::instructions::UserInstructions; use crate::mcp::McpManager; use crate::mcp_skill_dependencies::maybe_prompt_and_install_mcp_dependencies; +use crate::mcp_tool_catalog::McpToolCatalogSnapshot; +use crate::mcp_tool_catalog::McpToolCatalogUpdate; use crate::memories; use crate::mentions::build_connector_slug_counts; use crate::mentions::build_skill_name_counts; @@ -2354,6 +2356,39 @@ impl Session { state.get_connector_selection() } + async fn merge_advertised_mcp_catalog_for_model( + &self, + update: McpToolCatalogUpdate, + ) -> McpToolCatalogSnapshot { + let mut state = self.state.lock().await; + state.merge_advertised_mcp_tools(update) + } + + async fn list_model_visible_mcp_catalog_for_turn( + &self, + cancellation_token: &CancellationToken, + ) -> CodexResult { + let manager = self.services.mcp_connection_manager.read().await; + let has_servers = manager.has_servers(); + let tools = manager + .list_all_tools() + .or_cancel(cancellation_token) + .await?; + drop(manager); + Ok(self + .merge_advertised_mcp_catalog_for_model(McpToolCatalogUpdate { has_servers, tools }) + .await) + } + + async fn list_model_visible_mcp_catalog(&self) -> McpToolCatalogSnapshot { + let manager = self.services.mcp_connection_manager.read().await; + let has_servers = manager.has_servers(); + let tools = manager.list_all_tools().await; + drop(manager); + self.merge_advertised_mcp_catalog_for_model(McpToolCatalogUpdate { has_servers, tools }) + .await + } + // Clears connector IDs that were accumulated for explicit selection. pub(crate) async fn clear_connector_selection(&self) { let mut state = self.state.lock().await; @@ -3831,13 +3866,14 @@ impl Session { } } if turn_context.config.include_apps_instructions && turn_context.apps_enabled() { - let mcp_connection_manager = self.services.mcp_connection_manager.read().await; - let accessible_and_enabled_connectors = - connectors::list_accessible_and_enabled_connectors_from_manager( - &mcp_connection_manager, - &turn_context.config, - ) - .await; + let mcp_catalog = self.list_model_visible_mcp_catalog().await; + let accessible_and_enabled_connectors = connectors::with_app_enabled_state( + connectors::accessible_connectors_from_mcp_tools(&mcp_catalog.tools), + &turn_context.config, + ) + .into_iter() + .filter(|connector| connector.is_accessible && connector.is_enabled) + .collect::>(); if let Some(apps_section) = render_apps_section(&accessible_and_enabled_connectors) { developer_sections.push(apps_section); } @@ -4447,12 +4483,19 @@ impl Session { name: &str, namespace: Option<&str>, ) -> Option { - self.services + if let Some(tool) = self + .services .mcp_connection_manager .read() .await .resolve_tool_info(name, namespace) .await + { + return Some(tool); + } + + let state = self.state.lock().await; + state.resolve_advertised_mcp_tool(name, namespace) } pub async fn interrupt_task(self: &Arc) { @@ -4514,6 +4557,15 @@ impl Session { sandbox_cwd: turn_context.cwd.to_path_buf(), use_legacy_landlock: turn_context.features.use_legacy_landlock(), }; + let manager = self.services.mcp_connection_manager.read().await; + let existing_has_servers = manager.has_servers(); + let existing_tools = manager.list_all_tools().await; + drop(manager); + self.merge_advertised_mcp_catalog_for_model(McpToolCatalogUpdate { + has_servers: existing_has_servers, + tools: existing_tools, + }) + .await; { let mut guard = self.services.mcp_startup_cancellation_token.lock().await; guard.cancel(); @@ -6148,15 +6200,10 @@ pub(crate) async fn run_turn( // are normally hidden so we can describe the plugin's currently // usable capabilities for this turn. match sess - .services - .mcp_connection_manager - .read() - .await - .list_all_tools() - .or_cancel(&cancellation_token) + .list_model_visible_mcp_catalog_for_turn(&cancellation_token) .await { - Ok(mcp_tools) => mcp_tools, + Ok(mcp_catalog) => mcp_catalog.tools, Err(_) if turn_context.apps_enabled() => return None, Err(_) => HashMap::new(), } @@ -7041,13 +7088,11 @@ pub(crate) async fn built_tools( skills_outcome: Option<&SkillLoadOutcome>, cancellation_token: &CancellationToken, ) -> CodexResult> { - let mcp_connection_manager = sess.services.mcp_connection_manager.read().await; - let has_mcp_servers = mcp_connection_manager.has_servers(); - let all_mcp_tools = mcp_connection_manager - .list_all_tools() - .or_cancel(cancellation_token) + let mcp_catalog = sess + .list_model_visible_mcp_catalog_for_turn(cancellation_token) .await?; - drop(mcp_connection_manager); + let has_mcp_servers = mcp_catalog.has_servers; + let all_mcp_tools = mcp_catalog.tools; let loaded_plugins = sess .services .plugins_manager diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 5648fffd7c..255a94f179 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -10,6 +10,7 @@ use crate::config_loader::RequirementSource; use crate::config_loader::Sourced; use crate::exec::ExecCapturePolicy; use crate::function_tool::FunctionCallError; +use crate::mcp_tool_catalog::McpToolCatalogUpdate; use crate::mcp_tool_exposure::DIRECT_MCP_TOOL_EXPOSURE_THRESHOLD; use crate::mcp_tool_exposure::build_mcp_tool_exposure; use crate::shell::default_user_shell; @@ -1010,6 +1011,106 @@ fn mcp_tool_exposure_directly_exposes_explicit_apps_in_large_search_sets() { assert!(deferred_tools.contains_key("mcp__rmcp__tool_0")); } +#[tokio::test] +async fn built_tools_keeps_advertised_mcp_tools_when_live_manager_is_empty() { + let (session, turn_context, _rx) = make_session_and_context_with_rx().await; + session + .merge_advertised_mcp_catalog_for_model(McpToolCatalogUpdate { + has_servers: true, + tools: numbered_mcp_tools(1), + }) + .await; + + let router = built_tools( + session.as_ref(), + turn_context.as_ref(), + &[], + &HashSet::new(), + /*skills_outcome*/ None, + &CancellationToken::new(), + ) + .await + .expect("build tools"); + + assert!( + router + .model_visible_specs() + .iter() + .any(|spec| spec.name() == "mcp__rmcp__tool_0"), + "previously advertised MCP tool should remain model-visible" + ); +} + +#[tokio::test] +async fn built_tools_keeps_mcp_resource_tools_when_server_catalog_was_advertised() { + let (session, turn_context, _rx) = make_session_and_context_with_rx().await; + session + .merge_advertised_mcp_catalog_for_model(McpToolCatalogUpdate { + has_servers: true, + tools: HashMap::new(), + }) + .await; + + let router = built_tools( + session.as_ref(), + turn_context.as_ref(), + &[], + &HashSet::new(), + /*skills_outcome*/ None, + &CancellationToken::new(), + ) + .await + .expect("build tools"); + + let specs = router.model_visible_specs(); + + assert!(specs.iter().any(|spec| spec.name() == "list_mcp_resources")); + assert!( + specs + .iter() + .any(|spec| spec.name() == "list_mcp_resource_templates") + ); + assert!(specs.iter().any(|spec| spec.name() == "read_mcp_resource")); +} + +#[tokio::test] +async fn build_tool_call_routes_advertised_mcp_tool_when_live_manager_is_empty() { + let (session, _turn_context, _rx) = make_session_and_context_with_rx().await; + session + .merge_advertised_mcp_catalog_for_model(McpToolCatalogUpdate { + has_servers: true, + tools: numbered_mcp_tools(1), + }) + .await; + + let call = ToolRouter::build_tool_call( + session.as_ref(), + ResponseItem::FunctionCall { + id: None, + name: "mcp__rmcp__tool_0".to_string(), + namespace: None, + arguments: "{}".to_string(), + call_id: "call-1".to_string(), + }, + ) + .await + .expect("build tool call") + .expect("function call should produce a tool call"); + + match call.payload { + ToolPayload::Mcp { + server, + tool, + raw_arguments, + } => { + assert_eq!(server, "rmcp"); + assert_eq!(tool, "tool_0"); + assert_eq!(raw_arguments, "{}"); + } + other => panic!("expected MCP payload, got {other:?}"), + } +} + #[tokio::test] async fn reconstruct_history_matches_live_compactions() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 14154cffa5..5edb061e70 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -104,19 +104,6 @@ pub async fn list_accessible_connectors_from_mcp_tools( ) } -pub(crate) async fn list_accessible_and_enabled_connectors_from_manager( - mcp_connection_manager: &McpConnectionManager, - config: &Config, -) -> Vec { - with_app_enabled_state( - accessible_connectors_from_mcp_tools(&mcp_connection_manager.list_all_tools().await), - config, - ) - .into_iter() - .filter(|connector| connector.is_accessible && connector.is_enabled) - .collect() -} - pub(crate) async fn list_tool_suggest_discoverable_tools_with_auth( config: &Config, auth: Option<&CodexAuth>, diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 5e62199d2d..34323a18a9 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -46,6 +46,7 @@ pub use landlock::spawn_command_under_linux_sandbox; pub(crate) mod mcp; mod mcp_skill_dependencies; mod mcp_tool_approval_templates; +mod mcp_tool_catalog; mod mcp_tool_exposure; mod network_policy_decision; pub(crate) mod network_proxy_loader; diff --git a/codex-rs/core/src/mcp_tool_catalog.rs b/codex-rs/core/src/mcp_tool_catalog.rs new file mode 100644 index 0000000000..dad9832477 --- /dev/null +++ b/codex-rs/core/src/mcp_tool_catalog.rs @@ -0,0 +1,145 @@ +use std::collections::HashMap; + +use codex_mcp::ToolInfo; + +#[derive(Debug, Default)] +pub(crate) struct McpToolCatalog { + has_servers: bool, + tools: HashMap, +} + +#[derive(Debug)] +pub(crate) struct McpToolCatalogUpdate { + pub(crate) has_servers: bool, + pub(crate) tools: HashMap, +} + +#[derive(Debug)] +pub(crate) struct McpToolCatalogSnapshot { + pub(crate) has_servers: bool, + pub(crate) tools: HashMap, +} + +impl McpToolCatalog { + pub(crate) fn merge(&mut self, update: McpToolCatalogUpdate) -> McpToolCatalogSnapshot { + let McpToolCatalogUpdate { has_servers, tools } = update; + self.has_servers |= has_servers; + for (name, tool) in tools { + self.tools.entry(name).or_insert(tool); + } + self.snapshot() + } + + pub(crate) fn snapshot(&self) -> McpToolCatalogSnapshot { + McpToolCatalogSnapshot { + has_servers: self.has_servers, + tools: self.tools.clone(), + } + } + + pub(crate) fn resolve(&self, name: &str, namespace: Option<&str>) -> Option { + let qualified_name = qualified_mcp_tool_name(name, namespace); + self.tools.get(&qualified_name).cloned() + } +} + +fn qualified_mcp_tool_name(name: &str, namespace: Option<&str>) -> String { + match namespace { + Some(namespace) if name.starts_with(namespace) => name.to_string(), + Some(namespace) => format!("{namespace}{name}"), + None => name.to_string(), + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use rmcp::model::Tool; + use serde_json::Map as JsonObject; + + use super::*; + + fn tool_info(server_name: &str, callable_name: &str) -> ToolInfo { + ToolInfo { + server_name: server_name.to_string(), + callable_name: callable_name.to_string(), + callable_namespace: format!("mcp__{server_name}__"), + server_instructions: None, + tool: Tool { + name: callable_name.to_string().into(), + title: None, + description: Some(format!("Test tool: {callable_name}").into()), + input_schema: Arc::new(JsonObject::default()), + output_schema: None, + annotations: None, + execution: None, + icons: None, + meta: None, + }, + connector_id: None, + connector_name: None, + plugin_display_names: Vec::new(), + connector_description: None, + } + } + + #[test] + fn merge_preserves_first_seen_tool_definition() { + let mut catalog = McpToolCatalog::default(); + catalog.merge(McpToolCatalogUpdate { + has_servers: true, + tools: HashMap::from([( + "mcp__observability__get_dashboard".to_string(), + tool_info("observability", "get_dashboard"), + )]), + }); + catalog.merge(McpToolCatalogUpdate { + has_servers: true, + tools: HashMap::from([( + "mcp__observability__get_dashboard".to_string(), + tool_info("other", "replacement"), + )]), + }); + + let resolved = catalog + .resolve("mcp__observability__get_dashboard", None) + .expect("tool should resolve"); + assert_eq!(resolved.server_name, "observability"); + assert_eq!(resolved.tool.name.as_ref(), "get_dashboard"); + } + + #[test] + fn resolve_matches_namespaced_tool_calls() { + let mut catalog = McpToolCatalog::default(); + catalog.merge(McpToolCatalogUpdate { + has_servers: true, + tools: HashMap::from([( + "mcp__observability__get_dashboard".to_string(), + tool_info("observability", "get_dashboard"), + )]), + }); + + let resolved = catalog + .resolve("get_dashboard", Some("mcp__observability__")) + .expect("tool should resolve"); + assert_eq!(resolved.server_name, "observability"); + } + + #[test] + fn merge_preserves_advertised_server_presence() { + let mut catalog = McpToolCatalog::default(); + let snapshot = catalog.merge(McpToolCatalogUpdate { + has_servers: true, + tools: HashMap::new(), + }); + assert!(snapshot.has_servers); + + let snapshot = catalog.merge(McpToolCatalogUpdate { + has_servers: false, + tools: HashMap::new(), + }); + assert!(snapshot.has_servers); + assert!(snapshot.tools.is_empty()); + } +} diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 206f75060c..24c1711dc2 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -9,7 +9,11 @@ use std::collections::HashSet; use crate::codex::PreviousTurnSettings; use crate::codex::SessionConfiguration; use crate::context_manager::ContextManager; +use crate::mcp_tool_catalog::McpToolCatalog; +use crate::mcp_tool_catalog::McpToolCatalogSnapshot; +use crate::mcp_tool_catalog::McpToolCatalogUpdate; use crate::session_startup_prewarm::SessionStartupPrewarmHandle; +use codex_mcp::ToolInfo; use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; @@ -28,6 +32,7 @@ pub(crate) struct SessionState { /// model/realtime handling on subsequent regular turns (including full-context /// reinjection after resume or `/compact`). previous_turn_settings: Option, + advertised_mcp_tools: McpToolCatalog, /// Startup prewarmed session prepared during session initialization. pub(crate) startup_prewarm: Option, pub(crate) active_connector_selection: HashSet, @@ -47,6 +52,7 @@ impl SessionState { dependency_env: HashMap::new(), mcp_dependency_prompted: HashSet::new(), previous_turn_settings: None, + advertised_mcp_tools: McpToolCatalog::default(), startup_prewarm: None, active_connector_selection: HashSet::new(), pending_session_start_source: None, @@ -73,6 +79,21 @@ impl SessionState { self.previous_turn_settings = previous_turn_settings; } + pub(crate) fn merge_advertised_mcp_tools( + &mut self, + update: McpToolCatalogUpdate, + ) -> McpToolCatalogSnapshot { + self.advertised_mcp_tools.merge(update) + } + + pub(crate) fn resolve_advertised_mcp_tool( + &self, + name: &str, + namespace: Option<&str>, + ) -> Option { + self.advertised_mcp_tools.resolve(name, namespace) + } + pub(crate) fn clone_history(&self) -> ContextManager { self.history.clone() }