diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index b9a8d4ff08..d2227b411e 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -681,6 +681,7 @@ async fn refresh_codex_apps_after_connector_auth(sess: &Session, turn_context: & match mcp_tools_result { Ok(mcp_tools) => { + sess.late_mcp_tools.clear(); let auth = sess.services.auth_manager.auth().await; connectors::refresh_accessible_connectors_cache_from_mcp_tools( &turn_context.config, diff --git a/codex-rs/core/src/plugins/injection.rs b/codex-rs/core/src/plugins/injection.rs index 48e15247b5..ee7a9a3037 100644 --- a/codex-rs/core/src/plugins/injection.rs +++ b/codex-rs/core/src/plugins/injection.rs @@ -11,10 +11,17 @@ use crate::plugins::render_explicit_plugin_instructions; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::ToolInfo; +#[derive(Clone, Copy)] +pub(crate) enum PluginCapabilityResolution { + Runtime, + Declared, +} + pub(crate) fn build_plugin_injections( mentioned_plugins: &[PluginCapabilitySummary], mcp_tools: &[ToolInfo], available_connectors: &[connectors::AppInfo], + resolution: PluginCapabilityResolution, ) -> Vec { if mentioned_plugins.is_empty() { return Vec::new(); @@ -25,27 +32,36 @@ pub(crate) fn build_plugin_injections( mentioned_plugins .iter() .filter_map(|plugin| { - let available_mcp_servers = mcp_tools - .iter() - .filter(|tool| { - tool.server_name != CODEX_APPS_MCP_SERVER_NAME - && tool - .plugin_display_names - .iter() - .any(|plugin_name| plugin_name == &plugin.display_name) - }) - .map(|tool| tool.server_name.clone()) - .collect::>() - .into_iter() - .collect::>(); + let available_mcp_servers = match resolution { + PluginCapabilityResolution::Runtime => mcp_tools + .iter() + .filter(|tool| { + tool.server_name != CODEX_APPS_MCP_SERVER_NAME + && tool + .plugin_display_names + .iter() + .any(|plugin_name| plugin_name == &plugin.display_name) + }) + .map(|tool| tool.server_name.clone()) + .collect::>() + .into_iter() + .collect::>(), + PluginCapabilityResolution::Declared => plugin.mcp_server_names.clone(), + }; let available_apps = available_connectors .iter() .filter(|connector| { connector.is_enabled - && connector - .plugin_display_names - .iter() - .any(|plugin_name| plugin_name == &plugin.display_name) + && match resolution { + PluginCapabilityResolution::Runtime => connector + .plugin_display_names + .iter() + .any(|plugin_name| plugin_name == &plugin.display_name), + PluginCapabilityResolution::Declared => plugin + .app_connector_ids + .iter() + .any(|connector_id| connector_id.0 == connector.id), + } }) .map(connector_display_label) .collect::>() diff --git a/codex-rs/core/src/plugins/mod.rs b/codex-rs/core/src/plugins/mod.rs index 11e2c338bc..df8bb0c2ef 100644 --- a/codex-rs/core/src/plugins/mod.rs +++ b/codex-rs/core/src/plugins/mod.rs @@ -8,6 +8,7 @@ pub(crate) mod test_support; pub(crate) use codex_plugin::PluginCapabilitySummary; pub(crate) use discoverable::list_tool_suggest_discoverable_plugins; +pub(crate) use injection::PluginCapabilityResolution; pub(crate) use injection::build_plugin_injections; pub(crate) use render::render_explicit_plugin_instructions; diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 0fc971e985..955b13ccf0 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -373,6 +373,7 @@ impl Session { let mut manager = self.services.mcp_connection_manager.write().await; std::mem::replace(&mut *manager, refreshed_manager) }; + self.late_mcp_tools.clear(); old_manager.shutdown().await; } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index ed3cbf168c..9da6b60544 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1522,6 +1522,7 @@ impl Session { self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref()); self.services.skills_manager.clear_cache(); self.services.plugins_manager.clear_cache(); + self.late_mcp_tools.clear(); let hooks = build_hooks_for_config( config.as_ref(), self.services.plugins_manager.as_ref(), @@ -2846,7 +2847,10 @@ impl Session { .push(PersonalitySpecInstructions::new(personality_message).render()); } } - if turn_context.config.include_apps_instructions && turn_context.apps_enabled() { + if turn_context.config.include_apps_instructions + && turn_context.apps_enabled() + && !crate::tools::lazy_mcp::enabled(turn_context) + { let mcp_connection_manager = self.services.mcp_connection_manager.read().await; let accessible_and_enabled_connectors = connectors::list_accessible_and_enabled_connectors_from_manager( diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 86ec8f7631..a013332810 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -4,6 +4,7 @@ use crate::agents_md::LoadedAgentsMd; use crate::config::ConstraintError; use crate::skills::SkillError; use crate::state::ActiveTurn; +use crate::tools::registry::LateToolRegistry; use codex_protocol::SessionId; use codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE; use codex_protocol::config_types::ServiceTier; @@ -40,6 +41,7 @@ pub(crate) struct Session { pub(super) features: ManagedFeatures, pub(super) multi_agent_version: OnceLock, pub(super) pending_mcp_server_refresh_config: Mutex>, + pub(crate) late_mcp_tools: LateToolRegistry, pub(crate) conversation: Arc, pub(crate) active_turn: Mutex>, pub(crate) input_queue: InputQueue, @@ -1068,6 +1070,7 @@ impl Session { features: config.features.clone(), multi_agent_version, pending_mcp_server_refresh_config: Mutex::new(None), + late_mcp_tools: LateToolRegistry::default(), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), input_queue: InputQueue::new(), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index febd9f00b2..e094791647 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -572,6 +572,7 @@ fn test_tool_runtime(session: Arc, turn_context: Arc) -> T discoverable_tools: None, extension_tool_executors: Vec::new(), dynamic_tools: turn_context.dynamic_tools.as_slice(), + late_mcp_tools: None, }, )); let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); @@ -1440,10 +1441,12 @@ async fn refresh_runtime_config_refreshes_hooks() -> anyhow::Result<()> { }; assert!(session.hooks().preview_session_start(&request).is_empty()); + let late_mcp_generation = session.late_mcp_tools.generation(); let next_config = load_latest_config_for_session(&session).await; session.refresh_runtime_config(next_config).await; assert_eq!(session.hooks().preview_session_start(&request).len(), 1); + assert_ne!(session.late_mcp_tools.generation(), late_mcp_generation); Ok(()) } @@ -4920,6 +4923,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { features: config.features.clone(), multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()), pending_mcp_server_refresh_config: Mutex::new(None), + late_mcp_tools: crate::tools::registry::LateToolRegistry::default(), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), input_queue: super::input_queue::InputQueue::new(), @@ -6986,6 +6990,7 @@ where features: config.features.clone(), multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()), pending_mcp_server_refresh_config: Mutex::new(None), + late_mcp_tools: crate::tools::registry::LateToolRegistry::default(), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), input_queue: super::input_queue::InputQueue::new(), @@ -9429,6 +9434,7 @@ async fn fatal_tool_error_stops_turn_and_reports_error() { discoverable_tools: None, extension_tool_executors: Vec::new(), dynamic_tools: turn_context.dynamic_tools.as_slice(), + late_mcp_tools: None, }, ); let item = ResponseItem::CustomToolCall { diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 9d4e7d1e82..f2f9581d45 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -34,6 +34,7 @@ use crate::mentions::build_skill_name_counts; use crate::mentions::collect_explicit_app_ids; use crate::mentions::collect_explicit_plugin_mentions; use crate::mentions::collect_tool_mentions_from_messages; +use crate::plugins::PluginCapabilityResolution; use crate::plugins::build_plugin_injections; use crate::responses_retry::ResponsesStreamRequest; use crate::responses_retry::handle_retryable_response_stream_error; @@ -53,6 +54,7 @@ use crate::stream_events_utils::record_completed_response_item_with_finalized_fa use crate::tasks::emit_compact_metric; use crate::tools::ToolRouter; use crate::tools::context::SharedTurnDiffTracker; +use crate::tools::lazy_mcp; use crate::tools::parallel::ToolCallRuntime; use crate::tools::registry::ToolArgumentDiffConsumer; use crate::tools::router::ToolRouterParams; @@ -457,34 +459,41 @@ async fn build_skills_and_plugins( // enabled plugins, then converted into turn-scoped guidance below. let mentioned_plugins = collect_explicit_plugin_mentions(&user_input, loaded_plugins.capability_summaries()); - let mcp_tools = if turn_context.apps_enabled() || !mentioned_plugins.is_empty() { - // Plugin mentions need raw MCP/app inventory even when app tools - // 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) - .await - { - Ok(mcp_tools) => mcp_tools, - Err(_) if turn_context.apps_enabled() => return None, - Err(_) => Vec::new(), - } - } else { - Vec::new() - }; + let lazy_mcp_enabled = lazy_mcp::enabled(turn_context); + let mcp_tools = + if !lazy_mcp_enabled && (turn_context.apps_enabled() || !mentioned_plugins.is_empty()) { + // Plugin mentions need raw MCP/app inventory even when app tools + // 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) + .await + { + Ok(mcp_tools) => mcp_tools, + Err(_) if turn_context.apps_enabled() => return None, + Err(_) => Vec::new(), + } + } else { + Vec::new() + }; let available_connectors = if turn_context.apps_enabled() { - let connectors = codex_connectors::merge::merge_plugin_connectors_with_accessible( - loaded_plugins - .effective_apps() - .into_iter() - .map(|connector_id| connector_id.0), - connectors::accessible_connectors_from_mcp_tools(&mcp_tools), - ); + let plugin_app_ids = loaded_plugins + .effective_apps() + .into_iter() + .map(|connector_id| connector_id.0); + let connectors = if lazy_mcp_enabled { + codex_connectors::merge::merge_plugin_connectors(Vec::new(), plugin_app_ids) + } else { + codex_connectors::merge::merge_plugin_connectors_with_accessible( + plugin_app_ids, + connectors::accessible_connectors_from_mcp_tools(&mcp_tools), + ) + }; connectors::with_app_enabled_state(connectors, &turn_context.config) } else { Vec::new() @@ -540,8 +549,16 @@ async fn build_skills_and_plugins( &available_connectors, &skill_name_counts_lower, ); - let plugin_items = - build_plugin_injections(&mentioned_plugins, &mcp_tools, &available_connectors); + let plugin_items = build_plugin_injections( + &mentioned_plugins, + &mcp_tools, + &available_connectors, + if lazy_mcp_enabled { + PluginCapabilityResolution::Declared + } else { + PluginCapabilityResolution::Runtime + }, + ); let mut explicitly_enabled_connectors = collect_explicit_app_ids(&user_input); explicitly_enabled_connectors.extend(skill_connector_ids); let connector_names_by_id = available_connectors @@ -1086,10 +1103,14 @@ pub(crate) async fn built_tools( .instrument(trace_span!("read_mcp_connection_manager")) .await; let has_mcp_servers = mcp_connection_manager.has_servers(); - let all_mcp_tools = mcp_connection_manager - .list_all_tools() - .or_cancel(cancellation_token) - .await?; + let all_mcp_tools = if lazy_mcp::enabled(turn_context) { + Vec::new() + } else { + mcp_connection_manager + .list_all_tools() + .or_cancel(cancellation_token) + .await? + }; drop(mcp_connection_manager); let loaded_plugins = sess .services @@ -1171,6 +1192,7 @@ pub(crate) async fn built_tools( discoverable_tools, extension_tool_executors: extension_tool_executors(sess), dynamic_tools: turn_context.dynamic_tools.as_slice(), + late_mcp_tools: lazy_mcp::enabled(turn_context).then(|| sess.late_mcp_tools.clone()), }, ))) } diff --git a/codex-rs/core/src/stream_events_utils_tests.rs b/codex-rs/core/src/stream_events_utils_tests.rs index 9c4be94576..3b19b57de7 100644 --- a/codex-rs/core/src/stream_events_utils_tests.rs +++ b/codex-rs/core/src/stream_events_utils_tests.rs @@ -273,6 +273,7 @@ async fn handle_output_item_done_returns_contributed_last_agent_message() { discoverable_tools: None, extension_tool_executors: Vec::new(), dynamic_tools: turn_context.dynamic_tools.as_slice(), + late_mcp_tools: None, }, )); let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); diff --git a/codex-rs/core/src/tools/handlers/request_plugin_install.rs b/codex-rs/core/src/tools/handlers/request_plugin_install.rs index 3f7bd92d25..3c96c3b689 100644 --- a/codex-rs/core/src/tools/handlers/request_plugin_install.rs +++ b/codex-rs/core/src/tools/handlers/request_plugin_install.rs @@ -331,6 +331,7 @@ async fn refresh_missing_requested_connectors( match manager.hard_refresh_codex_apps_tools_cache().await { Ok(mcp_tools) => { + session.late_mcp_tools.clear(); let accessible_connectors = connectors::with_app_enabled_state( connectors::accessible_connectors_from_mcp_tools(&mcp_tools), &turn.config, diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index 2de1e20dd4..bf97dd3355 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -4,7 +4,9 @@ use crate::tools::context::ToolPayload; use crate::tools::context::ToolSearchOutput; use crate::tools::context::boxed_tool_output; use crate::tools::handlers::tool_search_spec::create_tool_search_tool; +use crate::tools::lazy_mcp; use crate::tools::registry::CoreToolRuntime; +use crate::tools::registry::LateToolRegistry; use crate::tools::registry::ToolExecutor; use bm25::Document; use bm25::Language; @@ -24,10 +26,19 @@ pub struct ToolSearchHandler { entries: Vec, search_source_infos: Vec, search_engine: SearchEngine, + late_mcp_tools: Option, } impl ToolSearchHandler { + #[cfg(test)] pub(crate) fn new(search_infos: Vec) -> Self { + Self::new_with_late_mcp_tools(search_infos, /*late_mcp_tools*/ None) + } + + pub(crate) fn new_with_late_mcp_tools( + search_infos: Vec, + late_mcp_tools: Option, + ) -> Self { let mut entries = Vec::with_capacity(search_infos.len()); let mut search_source_infos = Vec::new(); for search_info in search_infos { @@ -36,19 +47,19 @@ impl ToolSearchHandler { search_source_infos.push(source_info); } } - let documents: Vec> = entries - .iter() - .map(|entry| entry.search_text.clone()) - .enumerate() - .map(|(idx, search_text)| Document::new(idx, search_text)) - .collect(); - let search_engine = - SearchEngineBuilder::::with_documents(Language::English, documents).build(); + if late_mcp_tools.is_some() { + search_source_infos.push(ToolSearchSourceInfo { + name: "MCP servers".to_string(), + description: Some("Tools provided by configured MCP servers.".to_string()), + }); + } + let search_engine = Self::build_search_engine(&entries); Self { entries, search_source_infos, search_engine, + late_mcp_tools, } } } @@ -71,7 +82,12 @@ impl ToolExecutor for ToolSearchHandler { &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { - let ToolInvocation { payload, .. } = invocation; + let ToolInvocation { + session, + turn, + payload, + .. + } = invocation; let args = match payload { ToolPayload::ToolSearch { arguments } => arguments, @@ -96,11 +112,27 @@ impl ToolExecutor for ToolSearchHandler { )); } - if self.entries.is_empty() { + let mut late_entries = Vec::new(); + if let Some(late_mcp_tools) = &self.late_mcp_tools { + late_entries.extend( + lazy_mcp::register_eligible_tools(&session, &turn, late_mcp_tools) + .await + .into_iter() + .map(|search_info| search_info.entry), + ); + } + + if self.entries.is_empty() && late_entries.is_empty() { return Ok(boxed_tool_output(ToolSearchOutput { tools: Vec::new() })); } - let tools = self.search(query, limit)?; + let tools = if late_entries.is_empty() { + self.search(query, limit)? + } else { + let mut entries = self.entries.clone(); + entries.extend(late_entries); + Self::search_entries(&entries, query, limit)? + }; Ok(boxed_tool_output(ToolSearchOutput { tools })) } @@ -109,6 +141,16 @@ impl ToolExecutor for ToolSearchHandler { impl CoreToolRuntime for ToolSearchHandler {} impl ToolSearchHandler { + fn build_search_engine(entries: &[ToolSearchEntry]) -> SearchEngine { + let documents: Vec> = entries + .iter() + .map(|entry| entry.search_text.clone()) + .enumerate() + .map(|(idx, search_text)| Document::new(idx, search_text)) + .collect(); + SearchEngineBuilder::::with_documents(Language::English, documents).build() + } + fn search( &self, query: &str, @@ -123,6 +165,22 @@ impl ToolSearchHandler { self.search_output_tools(results) } + fn search_entries( + entries: &[ToolSearchEntry], + query: &str, + limit: usize, + ) -> Result, FunctionCallError> { + let search_engine = Self::build_search_engine(entries); + let results = search_engine + .search(query, limit) + .into_iter() + .map(|result| result.document.id) + .filter_map(|id| entries.get(id)); + Ok(coalesce_loadable_tool_specs( + results.map(|entry| entry.output.clone()), + )) + } + fn search_output_tools<'a>( &self, results: impl IntoIterator, diff --git a/codex-rs/core/src/tools/lazy_mcp.rs b/codex-rs/core/src/tools/lazy_mcp.rs new file mode 100644 index 0000000000..f3bf35ea77 --- /dev/null +++ b/codex-rs/core/src/tools/lazy_mcp.rs @@ -0,0 +1,86 @@ +use crate::connectors; +use crate::mcp_tool_exposure::build_mcp_tool_exposure; +use crate::session::session::Session; +use crate::session::turn_context::TurnContext; +use crate::tools::handlers::McpHandler; +use crate::tools::registry::CoreToolRuntime; +use crate::tools::registry::LateToolRegistry; +use crate::tools::registry::ToolExecutor; +use codex_features::Feature; +use codex_tools::ToolSearchInfo; +use std::sync::Arc; +use tracing::warn; + +pub(crate) fn enabled(turn_context: &TurnContext) -> bool { + turn_context + .config + .features + .enabled(Feature::ToolSearchAlwaysDeferMcpTools) + && turn_context.model_info.supports_search_tool + && turn_context.provider.capabilities().namespace_tools +} + +#[expect( + clippy::await_holding_invalid_type, + reason = "MCP tool listing borrows the read guard across the asynchronous inventory read" +)] +pub(crate) async fn register_eligible_tools( + session: &Arc, + turn_context: &TurnContext, + late_tools: &LateToolRegistry, +) -> Vec { + let registry_generation = late_tools.generation(); + let all_mcp_tools = session + .services + .mcp_connection_manager + .read() + .await + .list_all_tools() + .await; + let connectors = if turn_context.apps_enabled() { + let loaded_plugins = session + .services + .plugins_manager + .plugins_for_config(&turn_context.config.plugins_config_input()) + .await; + let connectors = codex_connectors::merge::merge_plugin_connectors_with_accessible( + loaded_plugins + .effective_apps() + .into_iter() + .map(|connector_id| connector_id.0), + connectors::accessible_connectors_from_mcp_tools(&all_mcp_tools), + ); + Some(connectors::with_app_enabled_state( + connectors, + &turn_context.config, + )) + } else { + None + }; + + let exposure = build_mcp_tool_exposure( + &all_mcp_tools, + connectors.as_deref(), + &turn_context.config, + /*search_tool_enabled*/ true, + ); + let mut search_infos = Vec::new(); + let mut handlers: Vec> = Vec::new(); + for tool_info in exposure.deferred_tools.unwrap_or(exposure.direct_tools) { + match McpHandler::new(tool_info) { + Ok(handler) => { + if let Some(search_info) = handler.search_info() { + search_infos.push(search_info); + } + handlers.push(Arc::new(handler)); + } + Err(err) => warn!("Skipping deferred MCP tool: failed to build tool spec: {err}"), + } + } + if late_tools.replace_if_generation(registry_generation, handlers) { + search_infos + } else { + warn!("Skipping deferred MCP tools because the MCP inventory changed during tool_search"); + Vec::new() + } +} diff --git a/codex-rs/core/src/tools/mod.rs b/codex-rs/core/src/tools/mod.rs index cd10d246db..b5dd9f5b43 100644 --- a/codex-rs/core/src/tools/mod.rs +++ b/codex-rs/core/src/tools/mod.rs @@ -4,6 +4,7 @@ pub(crate) mod events; pub(crate) mod handlers; pub(crate) mod hook_names; pub(crate) mod hosted_spec; +pub(crate) mod lazy_mcp; pub(crate) mod lifecycle; pub(crate) mod network_approval; pub(crate) mod orchestrator; diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 20e08e9f7d..5c678d9ba3 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::RwLock; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; @@ -324,11 +325,84 @@ impl CoreToolRuntime for ExposureOverride { pub struct ToolRegistry { tools: HashMap>, + late_tools: Option, +} + +#[derive(Clone, Default)] +pub(crate) struct LateToolRegistry { + state: Arc>, +} + +#[derive(Default)] +struct LateToolRegistryState { + generation: u64, + tools: HashMap>, +} + +impl LateToolRegistry { + pub(crate) fn generation(&self) -> u64 { + self.state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .generation + } + + pub(crate) fn replace_if_generation( + &self, + generation: u64, + tools: impl IntoIterator>, + ) -> bool { + let tools = tools + .into_iter() + .map(|tool| (tool.tool_name(), tool)) + .collect(); + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.generation != generation { + return false; + } + state.tools = tools; + true + } + + pub(crate) fn clear(&self) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.generation = state.generation.wrapping_add(1); + state.tools.clear(); + } + + fn tool(&self, name: &ToolName) -> Option> { + self.state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .tools + .get(name) + .map(Arc::clone) + } + + #[cfg(test)] + fn tool_names(&self) -> Vec { + self.state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .tools + .keys() + .cloned() + .collect() + } } impl ToolRegistry { fn new(tools: HashMap>) -> Self { - Self { tools } + Self { + tools, + late_tools: None, + } } pub(crate) fn from_tools(tools: impl IntoIterator>) -> Self { @@ -344,6 +418,15 @@ impl ToolRegistry { Self::new(tools_by_name) } + pub(crate) fn from_tools_with_late_registry( + tools: impl IntoIterator>, + late_tools: LateToolRegistry, + ) -> Self { + let mut registry = Self::from_tools(tools); + registry.late_tools = Some(late_tools); + registry + } + #[cfg(test)] pub(crate) fn empty_for_test() -> Self { Self::new(HashMap::new()) @@ -359,19 +442,26 @@ impl ToolRegistry { } fn tool(&self, name: &ToolName) -> Option> { - self.tools.get(name).map(Arc::clone) + self.tools + .get(name) + .map(Arc::clone) + .or_else(|| self.late_tools.as_ref()?.tool(name)) } #[cfg(test)] pub(crate) fn tool_names_for_test(&self) -> Vec { let mut names = self.tools.keys().cloned().collect::>(); + if let Some(late_tools) = &self.late_tools { + names.extend(late_tools.tool_names()); + } names.sort(); + names.dedup(); names } #[cfg(test)] pub(crate) fn tool_exposure(&self, name: &ToolName) -> Option { - self.tools.get(name).map(|tool| tool.exposure()) + self.tool(name).map(|tool| tool.exposure()) } pub(crate) fn create_diff_consumer( diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index f2566251c5..ced5f61313 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -27,6 +27,48 @@ impl ToolExecutor for TestHandler { impl CoreToolRuntime for TestHandler {} +#[test] +fn late_tool_registry_adds_tools_after_static_registry_creation() { + let static_name = codex_tools::ToolName::plain("static"); + let late_name = codex_tools::ToolName::namespaced("mcp__late", "lookup"); + let late_tools = LateToolRegistry::default(); + let registry = ToolRegistry::from_tools_with_late_registry( + [Arc::new(TestHandler { + tool_name: static_name.clone(), + }) as Arc], + late_tools.clone(), + ); + + assert!(registry.tool(&static_name).is_some()); + assert!(registry.tool(&late_name).is_none()); + + let generation = late_tools.generation(); + assert!(late_tools.replace_if_generation( + generation, + [Arc::new(TestHandler { + tool_name: late_name.clone(), + }) as Arc], + )); + + assert!(registry.tool(&late_name).is_some()); + assert_eq!( + registry.tool_names_for_test(), + vec![late_name.clone(), static_name.clone()] + ); + + late_tools.clear(); + assert!(registry.tool(&late_name).is_none()); + assert_eq!(registry.tool_names_for_test(), vec![static_name]); + + assert!(!late_tools.replace_if_generation( + generation, + [Arc::new(TestHandler { + tool_name: late_name.clone(), + }) as Arc], + )); + assert!(registry.tool(&late_name).is_none()); +} + #[derive(Clone)] enum LifecycleTestResult { Ok { success: bool }, diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index a095499ea0..aba84f6b74 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -5,6 +5,7 @@ use crate::tools::context::SharedTurnDiffTracker; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::registry::AnyToolResult; +use crate::tools::registry::LateToolRegistry; use crate::tools::registry::ToolArgumentDiffConsumer; use crate::tools::registry::ToolRegistry; use crate::tools::spec_plan::build_tool_router; @@ -42,6 +43,7 @@ pub(crate) struct ToolRouterParams<'a> { pub(crate) discoverable_tools: Option>, pub(crate) extension_tool_executors: Vec>>, pub(crate) dynamic_tools: &'a [DynamicToolSpec], + pub(crate) late_mcp_tools: Option, } impl ToolRouter { diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index cece7fc171..7b763ef351 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -116,6 +116,7 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow discoverable_tools: None, extension_tool_executors: Vec::new(), dynamic_tools: turn.dynamic_tools.as_slice(), + late_mcp_tools: None, }, ); @@ -195,6 +196,7 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { discoverable_tools: None, extension_tool_executors: Vec::new(), dynamic_tools: turn.dynamic_tools.as_slice(), + late_mcp_tools: None, }, ); @@ -230,6 +232,7 @@ async fn tools_without_handlers_do_not_support_parallel() -> anyhow::Result<()> discoverable_tools: None, extension_tool_executors: Vec::new(), dynamic_tools: turn.dynamic_tools.as_slice(), + late_mcp_tools: None, }, ); @@ -282,6 +285,7 @@ async fn specs_filter_deferred_dynamic_tools() -> anyhow::Result<()> { discoverable_tools: None, extension_tool_executors: Vec::new(), dynamic_tools: &dynamic_tools, + late_mcp_tools: None, }, ); @@ -343,6 +347,7 @@ async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow discoverable_tools: None, extension_tool_executors: extension_tool_executors(&session), dynamic_tools: turn.dynamic_tools.as_slice(), + late_mcp_tools: None, }, ); diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index fed62a0ee1..21e49975fa 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -48,6 +48,7 @@ use crate::tools::hosted_spec::WebSearchToolOptions; use crate::tools::hosted_spec::create_image_generation_tool; use crate::tools::hosted_spec::create_web_search_tool; use crate::tools::registry::CoreToolRuntime; +use crate::tools::registry::LateToolRegistry; use crate::tools::registry::ToolExposure; use crate::tools::registry::ToolRegistry; use crate::tools::registry::override_tool_exposure; @@ -146,6 +147,7 @@ struct CoreToolPlanContext<'a> { dynamic_tools: &'a [DynamicToolSpec], default_agent_type_description: &'a str, wait_agent_timeouts: WaitAgentTimeoutOptions, + late_mcp_tools: Option<&'a LateToolRegistry>, } pub(crate) fn build_tool_router( @@ -166,6 +168,7 @@ fn build_tool_specs_and_registry( discoverable_tools, extension_tool_executors, dynamic_tools, + late_mcp_tools, } = params; let default_agent_type_description = crate::agent::role::spawn_tool_spec::build(&std::collections::BTreeMap::new()); @@ -178,17 +181,19 @@ fn build_tool_specs_and_registry( dynamic_tools, default_agent_type_description: &default_agent_type_description, wait_agent_timeouts: wait_agent_timeout_options(turn_context), + late_mcp_tools: late_mcp_tools.as_ref(), }; let mut planned_tools = PlannedTools::default(); add_tool_sources(&context, &mut planned_tools); append_tool_search_executor(&context, &mut planned_tools); prepend_code_mode_executors(&context, &mut planned_tools); - build_model_visible_specs_and_registry(turn_context, planned_tools) + build_model_visible_specs_and_registry(turn_context, planned_tools, late_mcp_tools) } fn build_model_visible_specs_and_registry( turn_context: &TurnContext, planned_tools: PlannedTools, + late_mcp_tools: Option, ) -> (Vec, ToolRegistry) { let PlannedTools { runtimes, @@ -215,7 +220,12 @@ fn build_model_visible_specs_and_registry( } specs.extend(hosted_specs); - let registry = ToolRegistry::from_tools(runtimes); + let registry = match late_mcp_tools { + Some(late_mcp_tools) => { + ToolRegistry::from_tools_with_late_registry(runtimes, late_mcp_tools) + } + None => ToolRegistry::from_tools(runtimes), + }; let model_visible_specs = merge_into_namespaces(specs) .into_iter() .filter(|spec| { @@ -852,11 +862,14 @@ fn append_tool_search_executor( .filter(|executor| executor.exposure() == ToolExposure::Deferred) .filter_map(|executor| executor.search_info()) .collect::>(); - if search_infos.is_empty() { + if search_infos.is_empty() && context.late_mcp_tools.is_none() { return; } - planned_tools.add(ToolSearchHandler::new(search_infos)); + planned_tools.add(ToolSearchHandler::new_with_late_mcp_tools( + search_infos, + context.late_mcp_tools.cloned(), + )); } fn prepend_code_mode_executors( diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index f96f41f0f2..c95a8e4581 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -33,6 +33,7 @@ use serde_json::json; use crate::session::tests::make_session_and_context; use crate::session::turn_context::TurnContext; use crate::tools::handlers::multi_agents_spec::MULTI_AGENT_V1_NAMESPACE; +use crate::tools::registry::LateToolRegistry; use crate::tools::router::ToolRouter; use crate::tools::router::ToolRouterParams; @@ -43,6 +44,7 @@ struct ToolPlanInputs { discoverable_tools: Option>, extension_tool_executors: Vec>>, dynamic_tools: Vec, + late_mcp_tools: Option, } struct ToolPlanProbe { @@ -183,6 +185,7 @@ async fn probe_with( discoverable_tools: inputs.discoverable_tools, extension_tool_executors: inputs.extension_tool_executors, dynamic_tools: inputs.dynamic_tools.as_slice(), + late_mcp_tools: inputs.late_mcp_tools, }, ); ToolPlanProbe::from_router(router) @@ -685,6 +688,24 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() { "read_mcp_resource", ]); + let lazy_mcp = probe_with( + |turn| { + turn.model_info.supports_search_tool = true; + set_feature( + turn, + Feature::ToolSearchAlwaysDeferMcpTools, + /*enabled*/ true, + ); + }, + ToolPlanInputs { + late_mcp_tools: Some(LateToolRegistry::default()), + ..ToolPlanInputs::default() + }, + ) + .await; + lazy_mcp.assert_visible_contains(&["tool_search"]); + lazy_mcp.assert_registered_contains(&["tool_search"]); + let bedrock_namespace_capability = probe_with( |turn| { turn.model_info.supports_search_tool = true; diff --git a/codex-rs/core/tests/common/apps_test_server.rs b/codex-rs/core/tests/common/apps_test_server.rs index 7f7873e30b..c7aa2d6e15 100644 --- a/codex-rs/core/tests/common/apps_test_server.rs +++ b/codex-rs/core/tests/common/apps_test_server.rs @@ -7,6 +7,13 @@ use codex_login::CodexAuth; use codex_models_manager::bundled_models_response; use serde_json::Value; use serde_json::json; +use std::sync::Arc; +use std::sync::Condvar; +use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::time::Duration; +use tokio::sync::Notify; use wiremock::Mock; use wiremock::MockServer; use wiremock::Request; @@ -52,6 +59,69 @@ pub struct AppsTestServer { pub chatgpt_base_url: String, } +#[derive(Clone)] +pub struct ToolsListGate { + inner: Arc, +} + +struct ToolsListGateInner { + started: AtomicBool, + started_notify: Notify, + released: Mutex, + released_notify: Condvar, +} + +impl ToolsListGate { + fn new() -> Self { + Self { + inner: Arc::new(ToolsListGateInner { + started: AtomicBool::new(false), + started_notify: Notify::new(), + released: Mutex::new(false), + released_notify: Condvar::new(), + }), + } + } + + pub async fn wait_until_started(&self) { + while !self.inner.started.load(Ordering::Acquire) { + self.inner.started_notify.notified().await; + } + } + + pub fn release(&self) { + *self + .inner + .released + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = true; + self.inner.released_notify.notify_all(); + } + + fn block(&self) { + self.inner.started.store(true, Ordering::Release); + self.inner.started_notify.notify_waiters(); + let mut released = self + .inner + .released + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + while !*released { + released = self + .inner + .released_notify + .wait(released) + .unwrap_or_else(std::sync::PoisonError::into_inner); + } + } +} + +impl Drop for ToolsListGate { + fn drop(&mut self) { + self.release(); + } +} + #[derive(Clone, Copy)] pub enum AppsTestToolLoading { Direct, @@ -63,7 +133,47 @@ impl AppsTestServer { Self::mount_with_connector_name(server, CONNECTOR_NAME).await } + pub async fn mount_with_tools_delay( + server: &MockServer, + tools_list_delay: Duration, + ) -> Result { + mount_oauth_metadata(server).await; + mount_connectors_directory(server).await; + mount_streamable_http_json_rpc( + server, + CONNECTOR_NAME.to_string(), + CONNECTOR_DESCRIPTION.to_string(), + /*searchable*/ false, + /*include_app_only_tool*/ false, + tools_list_delay, + ) + .await; + Ok(Self { + chatgpt_base_url: server.uri(), + }) + } + + pub async fn mount_with_tools_gate(server: &MockServer) -> Result<(Self, ToolsListGate)> { + let gate = ToolsListGate::new(); + mount_oauth_metadata(server).await; + mount_connectors_directory(server).await; + mount_streamable_http_json_rpc_with_gate(server, gate.clone()).await; + Ok(( + Self { + chatgpt_base_url: server.uri(), + }, + gate, + )) + } + pub async fn mount_searchable(server: &MockServer) -> Result { + Self::mount_searchable_with_tools_delay(server, Duration::ZERO).await + } + + pub async fn mount_searchable_with_tools_delay( + server: &MockServer, + tools_list_delay: Duration, + ) -> Result { mount_oauth_metadata(server).await; mount_connectors_directory(server).await; mount_streamable_http_json_rpc( @@ -72,6 +182,7 @@ impl AppsTestServer { CONNECTOR_DESCRIPTION.to_string(), /*searchable*/ true, /*include_app_only_tool*/ false, + tools_list_delay, ) .await; Ok(Self { @@ -91,6 +202,7 @@ impl AppsTestServer { CONNECTOR_DESCRIPTION.to_string(), /*searchable*/ false, /*include_app_only_tool*/ false, + Duration::ZERO, ) .await; Ok(Self { @@ -110,6 +222,7 @@ impl AppsTestServer { CONNECTOR_DESCRIPTION.to_string(), matches!(tool_loading, AppsTestToolLoading::Searchable), /*include_app_only_tool*/ true, + Duration::ZERO, ) .await; Ok(Self { @@ -264,6 +377,7 @@ async fn mount_streamable_http_json_rpc( connector_description: String, searchable: bool, include_app_only_tool: bool, + tools_list_delay: Duration, ) { Mock::given(method("POST")) .and(path_regex("^/api/codex/apps/?$")) @@ -272,6 +386,23 @@ async fn mount_streamable_http_json_rpc( connector_description, searchable, include_app_only_tool, + tools_list_delay, + tools_list_gate: None, + }) + .mount(server) + .await; +} + +async fn mount_streamable_http_json_rpc_with_gate(server: &MockServer, gate: ToolsListGate) { + Mock::given(method("POST")) + .and(path_regex("^/api/codex/apps/?$")) + .respond_with(CodexAppsJsonRpcResponder { + connector_name: CONNECTOR_NAME.to_string(), + connector_description: CONNECTOR_DESCRIPTION.to_string(), + searchable: false, + include_app_only_tool: false, + tools_list_delay: Duration::ZERO, + tools_list_gate: Some(gate), }) .mount(server) .await; @@ -282,6 +413,8 @@ struct CodexAppsJsonRpcResponder { connector_description: String, searchable: bool, include_app_only_tool: bool, + tools_list_delay: Duration, + tools_list_gate: Option, } impl Respond for CodexAppsJsonRpcResponder { @@ -327,6 +460,9 @@ impl Respond for CodexAppsJsonRpcResponder { } "notifications/initialized" => ResponseTemplate::new(202), "tools/list" => { + if let Some(gate) = &self.tools_list_gate { + gate.block(); + } let id = body.get("id").cloned().unwrap_or(Value::Null); let mut response = json!({ "jsonrpc": "2.0", @@ -475,7 +611,9 @@ impl Respond for CodexAppsJsonRpcResponder { } })); } - ResponseTemplate::new(200).set_body_json(response) + ResponseTemplate::new(200) + .set_body_json(response) + .set_delay(self.tools_list_delay) } "tools/call" => { let id = body.get("id").cloned().unwrap_or(Value::Null); diff --git a/codex-rs/core/tests/suite/plugins.rs b/codex-rs/core/tests/suite/plugins.rs index 412b5edbbc..db28521010 100644 --- a/codex-rs/core/tests/suite/plugins.rs +++ b/codex-rs/core/tests/suite/plugins.rs @@ -11,6 +11,7 @@ use codex_login::CodexAuth; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use core_test_support::apps_test_server::AppsTestServer; +use core_test_support::apps_test_server::configure_search_capable_model; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_response_created; use core_test_support::responses::mount_sse_once; @@ -23,6 +24,7 @@ use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use core_test_support::wait_for_mcp_server; use tempfile::TempDir; +use tokio::time::timeout; use wiremock::MockServer; const SAMPLE_PLUGIN_CONFIG_NAME: &str = "sample@test"; @@ -326,6 +328,96 @@ async fn explicit_plugin_mentions_inject_plugin_guidance() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn lazy_explicit_plugin_mentions_use_declared_capabilities() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let apps_mock_server = MockServer::start().await; + let (apps_server, tools_list_gate) = + AppsTestServer::mount_with_tools_gate(&apps_mock_server).await?; + let mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + + let codex_home = Arc::new(TempDir::new()?); + let rmcp_test_server_bin = match stdio_server_bin() { + Ok(bin) => bin, + Err(err) => { + eprintln!("test_stdio_server binary not available, skipping test: {err}"); + return Ok(()); + } + }; + write_plugin_skill_plugin(codex_home.as_ref()); + write_plugin_mcp_plugin(codex_home.as_ref(), &rmcp_test_server_bin); + write_plugin_app_plugin(codex_home.as_ref()); + + let chatgpt_base_url = apps_server.chatgpt_base_url; + let mut builder = test_codex() + .with_home(codex_home) + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(move |config| { + config + .features + .enable(Feature::Apps) + .expect("test config should allow feature update"); + config + .features + .enable(Feature::ToolSearchAlwaysDeferMcpTools) + .expect("test config should allow feature update"); + configure_search_capable_model(config); + config.chatgpt_base_url = chatgpt_base_url; + }); + let test = builder.build(&server).await?; + + timeout(Duration::from_secs(/*secs*/ 10), async { + test.codex + .submit(Op::UserInput { + items: vec![codex_protocol::user_input::UserInput::Mention { + name: "sample".into(), + path: format!("plugin://{SAMPLE_PLUGIN_CONFIG_NAME}"), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + wait_for_event(&test.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + Result::<()>::Ok(()) + }) + .await + .expect("plugin mention should not wait for uncached MCP tools")?; + timeout( + Duration::from_secs(/*secs*/ 10), + tools_list_gate.wait_until_started(), + ) + .await + .expect("background MCP tools/list should reach the deliberate gate"); + tools_list_gate.release(); + + let request = mock.single_request(); + let developer_messages = request.message_input_texts("developer"); + assert!( + developer_messages + .iter() + .any(|text| text.contains("MCP servers from this plugin")), + "expected declared plugin MCP guidance: {developer_messages:?}" + ); + assert!( + developer_messages + .iter() + .any(|text| text.contains("Apps from this plugin")), + "expected declared plugin app guidance: {developer_messages:?}" + ); + let request_tools = tool_names(&request.body_json()); + assert!(request_tools.iter().any(|name| name == "tool_search")); + assert!(!request_tools.iter().any(|name| name == "mcp__sample")); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn explicit_plugin_mentions_track_plugin_used_analytics() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/tests/suite/search_tool.rs b/codex-rs/core/tests/suite/search_tool.rs index 686b9b9e42..8506d0a445 100644 --- a/codex-rs/core/tests/suite/search_tool.rs +++ b/codex-rs/core/tests/suite/search_tool.rs @@ -53,6 +53,8 @@ use serde_json::Value; use serde_json::json; use std::collections::HashMap; use std::time::Duration; +use tokio::time::timeout; +use wiremock::MockServer; const SEARCH_TOOL_DESCRIPTION_SNIPPETS: [&str; 2] = [ "You have access to tools from the following sources", @@ -180,7 +182,9 @@ async fn always_defer_feature_hides_small_app_tool_sets() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let apps_server = AppsTestServer::mount(&server).await?; + let apps_mock_server = MockServer::start().await; + let (apps_server, tools_list_gate) = + AppsTestServer::mount_with_tools_gate(&apps_mock_server).await?; let mock = mount_sse_once( &server, sse(vec![ @@ -200,12 +204,23 @@ async fn always_defer_feature_hides_small_app_tool_sets() -> Result<()> { }); let test = builder.build(&server).await?; - test.submit_turn_with_approval_and_permission_profile( - "list tools", - AskForApproval::Never, - PermissionProfile::Disabled, + timeout( + Duration::from_secs(/*secs*/ 10), + test.submit_turn_with_approval_and_permission_profile( + "list tools", + AskForApproval::Never, + PermissionProfile::Disabled, + ), ) - .await?; + .await + .expect("first turn should not wait for uncached MCP tools")?; + timeout( + Duration::from_secs(/*secs*/ 10), + tools_list_gate.wait_until_started(), + ) + .await + .expect("background MCP tools/list should reach the deliberate gate"); + tools_list_gate.release(); let body = mock.single_request().body_json(); let tools = tool_names(&body); @@ -1239,7 +1254,6 @@ async fn tool_search_surfaced_mcp_tool_errors_are_returned_to_model() -> Result< .expect("test mcp servers should accept any configuration"); }); let test = builder.build(&server).await?; - wait_for_mcp_server(&test.codex, "rmcp").await?; test.codex .submit(Op::UserInput {