diff --git a/codex-rs/codex-mcp/src/binding.rs b/codex-rs/codex-mcp/src/binding.rs index d1f1577a70..6cec4cb393 100644 --- a/codex-rs/codex-mcp/src/binding.rs +++ b/codex-rs/codex-mcp/src/binding.rs @@ -96,7 +96,11 @@ impl McpBinding { server: &str, params: Option, ) -> Result { - self.clients.list_resources(server, params).await + if self.clients.client(server).is_some() { + self.clients.list_resources(server, params).await + } else { + self.connections.list_resources(server, params).await + } } pub async fn list_all_resources( @@ -111,7 +115,13 @@ impl McpBinding { server: &str, params: Option, ) -> Result { - self.clients.list_resource_templates(server, params).await + if self.clients.client(server).is_some() { + self.clients.list_resource_templates(server, params).await + } else { + self.connections + .list_resource_templates(server, params) + .await + } } pub async fn list_all_resource_templates( @@ -128,7 +138,11 @@ impl McpBinding { server: &str, params: ReadResourceRequestParams, ) -> Result { - self.clients.read_resource(server, params).await + if self.clients.client(server).is_some() { + self.clients.read_resource(server, params).await + } else { + self.connections.read_resource(server, params).await + } } } diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 752b1a3ac8..34a0f35613 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -23,6 +23,7 @@ pub use tool_catalog::tool_is_model_visible; use std::collections::HashMap; use std::sync::Arc; +use std::sync::OnceLock; use std::sync::atomic::Ordering; use std::time::Duration; @@ -144,6 +145,7 @@ impl McpServerView { pub(crate) struct McpConnectionSet { servers: HashMap, required_servers: Vec, + optional_startup_deadline: OnceLock, tool_catalog_revision: Arc>, codex_apps_tools_override: RwLock>>, codex_apps_refresh_lock: Mutex<()>, @@ -483,6 +485,7 @@ impl McpConnectionSet { let manager = Self { servers, required_servers, + optional_startup_deadline: OnceLock::new(), tool_catalog_revision: Arc::new(RwLock::new(0)), codex_apps_tools_override: RwLock::new(None), codex_apps_refresh_lock: Mutex::new(()), @@ -540,6 +543,7 @@ impl McpConnectionSet { Self { servers: HashMap::new(), required_servers: Vec::new(), + optional_startup_deadline: OnceLock::new(), tool_catalog_revision: Arc::new(RwLock::new(0)), codex_apps_tools_override: RwLock::new(None), codex_apps_refresh_lock: Mutex::new(()), diff --git a/codex-rs/codex-mcp/src/connection_manager/resources.rs b/codex-rs/codex-mcp/src/connection_manager/resources.rs index 8a46971acb..06edde90b4 100644 --- a/codex-rs/codex-mcp/src/connection_manager/resources.rs +++ b/codex-rs/codex-mcp/src/connection_manager/resources.rs @@ -4,6 +4,7 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; use anyhow::anyhow; +use rmcp::model::ListResourceTemplatesResult; use rmcp::model::ListResourcesResult; use rmcp::model::PaginatedRequestParams; use rmcp::model::ReadResourceRequestParams; @@ -154,6 +155,19 @@ impl McpConnectionSet { .with_context(|| format!("resources/list failed for `{server}`")) } + pub async fn list_resource_templates( + &self, + server: &str, + params: Option, + ) -> Result { + let (managed, timeout) = self.client_by_name(server).await?; + managed + .client + .list_resource_templates(params, timeout) + .await + .with_context(|| format!("resources/templates/list failed for `{server}`")) + } + pub async fn read_resource( &self, server: &str, diff --git a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs index 1da3134648..47afb2242d 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -1,5 +1,6 @@ use std::sync::Arc; use std::sync::atomic::Ordering; +use std::time::Duration; use std::time::Instant; use anyhow::Context; @@ -30,6 +31,7 @@ use crate::tools::normalize_tools_for_model_with_prefix; const MCP_UI_META_KEY: &str = "ui"; const MCP_UI_VISIBILITY_META_KEY: &str = "visibility"; const MCP_UI_MODEL_VISIBILITY: &str = "model"; +const OPTIONAL_MCP_STARTUP_GRACE: Duration = Duration::from_secs(1); /// Returns whether a tool may be included in model-facing tool declarations. /// @@ -132,10 +134,14 @@ impl McpConnectionSet { self: &Arc, config: Arc, plugins_available: bool, + required_servers: &[String], ) -> McpBinding { let revision = self.tool_catalog_revision.read().await; let mut listed_tools = Vec::new(); let mut clients = std::collections::HashMap::new(); + let optional_startup_deadline = *self + .optional_startup_deadline + .get_or_init(|| tokio::time::Instant::now() + OPTIONAL_MCP_STARTUP_GRACE); for (server_name, view) in &self.servers { if !view .connection @@ -143,7 +149,15 @@ impl McpConnectionSet { .startup_complete .load(Ordering::Acquire) { - if view.connection.client.has_cached_tools() { + let required = self.required_servers.binary_search(server_name).is_ok(); + let has_cached_tools = view.connection.client.has_cached_tools(); + let must_wait_for_startup = required + || self.is_selected_plugin_mcp_server(server_name) + || required_servers + .iter() + .any(|required| required == server_name) + || (server_name == CODEX_APPS_MCP_SERVER_NAME && !has_cached_tools); + if !must_wait_for_startup && has_cached_tools { if let Some(server_tools) = view.listed_tools(&self.tool_plugin_provenance).await { @@ -156,6 +170,17 @@ impl McpConnectionSet { } continue; } + if !must_wait_for_startup + && tokio::time::timeout_at( + optional_startup_deadline, + view.connection.client.client(), + ) + .await + .is_err() + { + trace!(server_name = %server_name, "omitting pending optional MCP server"); + continue; + } let _ = view.connection.client.client().await; } view.connection.client.reconnect_failed_startup().await; diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index d78896d63b..ed8f226e1d 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -86,6 +86,7 @@ impl McpConnectionSet { Self { servers: HashMap::new(), required_servers: Vec::new(), + optional_startup_deadline: OnceLock::new(), tool_catalog_revision: Arc::new(RwLock::new(0)), codex_apps_tools_override: RwLock::new(None), codex_apps_refresh_lock: Mutex::new(()), @@ -194,6 +195,7 @@ async fn capture_binding(manager: &Arc) -> McpBinding { .capture_binding_with_metadata( Arc::new(crate::mcp::tests::test_mcp_config(std::env::temp_dir())), /*plugins_available*/ false, + /*required_servers*/ &[], ) .await } @@ -1666,6 +1668,87 @@ async fn capture_binding_exposes_cached_tools_before_startup() { ); } +#[tokio::test(start_paused = true)] +async fn capture_binding_skips_pending_optional_servers_after_one_shared_startup_grace() { + let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionSet::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + let mut plugin_config = crate::mcp::tests::test_mcp_config(std::env::temp_dir()); + let mut catalog = crate::ResolvedMcpCatalog::builder(); + catalog.register(crate::McpServerRegistration::from_plugin( + "pending-one".to_string(), + crate::McpPluginAttribution::new("optional-plugin".to_string(), "Optional".to_string()), + /*plugin_order*/ 0, + serde_json::from_value(serde_json::json!({ "command": "optional-plugin" })) + .expect("optional plugin MCP config"), + )); + plugin_config.mcp_server_catalog = catalog.build(); + manager.tool_plugin_provenance = Arc::new(crate::tool_plugin_provenance(&plugin_config)); + for server_name in ["pending-one", "pending-two"] { + manager.insert_test_client( + server_name.to_string(), + AsyncManagedClient { + client: futures::future::pending::>() + .boxed() + .shared(), + is_codex_apps_mcp_server: false, + cached_server_info: None, + codex_apps_tools_cache_context: None, + tool_catalog_cache_context: None, + startup_complete: Arc::new(AtomicBool::new(false)), + startup_reconnect: None, + cancel_token: CancellationToken::new(), + }, + ); + } + + let manager = Arc::new(manager); + let binding = tokio::time::timeout(Duration::from_millis(1500), capture_binding(&manager)) + .await + .expect("all optional servers should share a single startup grace"); + assert!(binding.tools().is_empty()); + + let binding = tokio::time::timeout(Duration::from_millis(1), capture_binding(&manager)) + .await + .expect("later bindings must not restart the optional startup grace"); + assert!(binding.tools().is_empty()); + + assert!( + tokio::time::timeout( + Duration::from_millis(1), + binding.list_resources("pending-one", /*params*/ None), + ) + .await + .is_err(), + "resources must wait for an omitted server instead of failing immediately" + ); + assert!( + tokio::time::timeout( + Duration::from_millis(1), + binding.list_all_resources(|server| server == "pending-one"), + ) + .await + .is_ok(), + "resource discovery must not wait for an omitted optional server" + ); + + let required_servers = vec!["pending-one".to_string()]; + let binding = tokio::time::timeout( + Duration::from_millis(1), + manager.capture_binding_with_metadata( + Arc::new(crate::mcp::tests::test_mcp_config(std::env::temp_dir())), + /*plugins_available*/ false, + &required_servers, + ), + ) + .await; + assert!(binding.is_err(), "explicitly requested servers must wait"); +} + #[tokio::test] async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { let managed_client = diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 27af83ef3f..c678d708a6 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -214,12 +214,20 @@ impl McpRuntime { /// Captures the latest published configuration and live client handles. pub async fn current_binding(&self) -> Option> { + self.current_binding_with_required_servers(&[]).await + } + + /// Captures the latest runtime, waiting for servers explicitly required by this turn. + pub async fn current_binding_with_required_servers( + &self, + required_servers: &[String], + ) -> Option> { let current = self.current.load_full(); let config = Arc::clone(current.config.as_ref()?); Some(Arc::new( current .connections - .capture_binding_with_metadata(config, current.plugins_available) + .capture_binding_with_metadata(config, current.plugins_available, required_servers) .await, )) } @@ -250,7 +258,7 @@ impl McpRuntime { Some(Arc::new( current .connections - .capture_binding_with_metadata(config, current.plugins_available) + .capture_binding_with_metadata(config, current.plugins_available, &[]) .await, )) } diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index eb57ce1d21..21c1293aae 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -260,6 +260,7 @@ impl Session { self: &Arc, turn_context: &TurnContext, selected_capability_roots: &[ResolvedSelectedCapabilityRoot], + required_servers: &[String], ) -> Arc { let ready_selected_capability_roots = Self::ready_selected_capability_roots(selected_capability_roots); @@ -272,7 +273,12 @@ impl Session { self.mark_mcp_runtime_dirty(); } self.refresh_mcp_if_dirty().await; - if let Some(binding) = self.services.mcp_runtime.current_binding().await { + if let Some(binding) = self + .services + .mcp_runtime + .current_binding_with_required_servers(required_servers) + .await + { return binding; } let config = Arc::new(self.runtime_mcp_config(&turn_context.config).await); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 4c5c7b1134..02625a9810 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3023,6 +3023,20 @@ impl Session { self: &Arc, turn_context: Arc, cancellation_token: &CancellationToken, + ) -> CodexResult> { + self.capture_step_context_with_required_mcp_servers( + turn_context, + cancellation_token, + /*required_servers*/ &[], + ) + .await + } + + pub(crate) async fn capture_step_context_with_required_mcp_servers( + self: &Arc, + turn_context: Arc, + cancellation_token: &CancellationToken, + required_servers: &[String], ) -> CodexResult> { // Keep selections fixed for the turn while allowing their startup work to finish. let environments = turn_context.environments.refresh_readiness(); @@ -3046,7 +3060,11 @@ impl Session { extension_data.insert(selected_capability_roots.clone()); let (mcp, prepared_recommendations) = async { tokio::join!( - self.mcp_runtime_for_step(turn_context.as_ref(), &selected_capability_roots), + self.mcp_runtime_for_step( + turn_context.as_ref(), + &selected_capability_roots, + required_servers, + ), turn::prepare_tool_recommendations(self.as_ref(), turn_context.as_ref()), ) } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 031ffcc7e4..cb3f85eb0d 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -8096,7 +8096,11 @@ async fn refresh_mcp_servers_uses_latest_state_for_existing_turns() { .await .expect("a fresh cancellation token cannot be cancelled"); let rematerialized_old = session - .mcp_runtime_for_step(&turn_context, /*selected_capability_roots*/ &[]) + .mcp_runtime_for_step( + &turn_context, + /*selected_capability_roots*/ &[], + /*required_servers*/ &[], + ) .await; let configured_servers = codex_mcp::configured_mcp_servers(new_step.mcp.config()); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 6c76a46bfa..6eb6ae4ba6 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -183,9 +183,26 @@ pub(crate) async fn run_turn( return Ok(None); } + let user_input = turn_user_input(&input); + let (required_servers, mentioned_plugins) = + match required_mcp_servers_for_input(&sess, turn_context.as_ref(), &user_input) + .or_cancel(&cancellation_token) + .await + { + Ok(requirements) => requirements, + Err(err) => { + run_hooks_and_record_inputs(&sess, &turn_context, &input).await; + return Err(err.into()); + } + }; + // run_turn owns the step used to seed context and make the first sampling request. let first_step_context = match sess - .capture_step_context(Arc::clone(&turn_context), &cancellation_token) + .capture_step_context_with_required_mcp_servers( + Arc::clone(&turn_context), + &cancellation_token, + &required_servers, + ) .await { Ok(step_context) => step_context, @@ -204,7 +221,8 @@ pub(crate) async fn run_turn( let Some((injection_items, explicitly_enabled_connectors)) = build_skills_and_plugins( &sess, first_step_context.as_ref(), - &input, + &user_input, + &mentioned_plugins, &cancellation_token, ) .await @@ -276,10 +294,26 @@ pub(crate) async fn run_turn( // Capture once so context, advertised tools, and tool calls share one request view. let step_context = match next_step_context.take() { Some(step_context) => step_context, - None => { + None if pending_input.is_empty() => { sess.capture_step_context(Arc::clone(&turn_context), &cancellation_token) .await? } + None => { + let pending_user_input = turn_user_input(&pending_input); + let (required_servers, _) = required_mcp_servers_for_input( + &sess, + turn_context.as_ref(), + &pending_user_input, + ) + .or_cancel(&cancellation_token) + .await?; + sess.capture_step_context_with_required_mcp_servers( + Arc::clone(&turn_context), + &cancellation_token, + &required_servers, + ) + .await? + } }; let sampling_request_result: CodexResult<_> = async { super::time_reminder::maybe_record_current_time_reminder( @@ -567,11 +601,127 @@ pub(crate) async fn run_hooks_and_record_inputs( blocked_input && !accepted_user_input } +fn turn_user_input(input: &[TurnInput]) -> Vec { + input + .iter() + .filter_map(|item| match item { + TurnInput::UserInput { content, .. } => Some(content.as_slice()), + TurnInput::ResponseItem(_) | TurnInput::InterAgentCommunication(_) => None, + }) + .flatten() + .cloned() + .collect() +} + +async fn required_mcp_servers_for_input( + sess: &Arc, + turn_context: &TurnContext, + user_input: &[UserInput], +) -> (Vec, Vec) { + if crate::guardian::is_guardian_reviewer_source(&turn_context.session_source) { + return (Vec::new(), Vec::new()); + } + + // Plugin capabilities depend on authentication, so project them only after + // the runtime has aligned the plugin manager with its current account. + sess.refresh_mcp_if_dirty().await; + let loaded_plugins = sess + .services + .plugins_manager + .plugins_for_config(&turn_context.config.plugins_config_input()) + .await; + let current_config = sess.services.mcp_runtime.current_config(); + let mentioned_plugins = + collect_explicit_plugin_mentions(user_input, loaded_plugins.capability_summaries()); + let mut required_servers = mentioned_plugins + .iter() + .flat_map(|plugin| plugin.mcp_server_names.iter().cloned()) + .collect::>(); + + let messages = user_input + .iter() + .filter_map(|input| match input { + UserInput::Text { text, .. } => Some(text.clone()), + _ => None, + }) + .collect::>(); + let mentions = collect_tool_mentions_from_messages(&messages); + let paths = user_input + .iter() + .filter_map(|input| match input { + UserInput::Mention { path, .. } => Some(path.clone()), + _ => None, + }) + .chain(mentions.paths); + required_servers.extend(paths.filter_map(|path| { + path.strip_prefix("mcp://") + .filter(|server| !server.is_empty()) + .map(str::to_string) + })); + + let connector_slug_counts = if turn_context.apps_enabled() && !mentions.plain_names.is_empty() { + let cached_connectors = + connectors::list_cached_accessible_connectors_from_mcp_tools(&turn_context.config) + .await; + let accessible_connectors = match cached_connectors { + Some(connectors) => connectors, + None => sess + .services + .mcp_runtime + .current_binding() + .await + .map(|binding| connectors::accessible_connectors_from_mcp_tools(binding.tools())) + .unwrap_or_default(), + }; + let connector_ids = current_config + .iter() + .flat_map(|config| config.connector_snapshot.connector_ids()) + .map(|connector_id| connector_id.0.clone()); + build_connector_slug_counts( + &codex_connectors::merge::merge_plugin_connectors_with_accessible( + connector_ids, + accessible_connectors, + ), + ) + } else { + HashMap::new() + }; + let skills_outcome = turn_context.turn_skills.snapshot.outcome(); + let mentioned_skills = collect_explicit_skill_mentions( + user_input, + &skills_outcome.skills, + &skills_outcome.disabled_paths, + &connector_slug_counts, + ); + for skill in mentioned_skills { + if let Some(dependencies) = skill.dependencies { + required_servers.extend( + dependencies + .tools + .into_iter() + .filter(|tool| tool.r#type.eq_ignore_ascii_case("mcp")) + .map(|tool| tool.value), + ); + } + if let Some(plugin_id) = skill.plugin_id.as_deref() + && let Some(plugin) = loaded_plugins + .capability_summaries() + .iter() + .find(|plugin| plugin.config_name == plugin_id) + { + required_servers.extend(plugin.mcp_server_names.iter().cloned()); + } + } + + (required_servers.into_iter().collect(), mentioned_plugins) +} + #[instrument(level = "trace", skip_all)] async fn build_skills_and_plugins( sess: &Arc, step_context: &StepContext, - input: &[TurnInput], + user_input: &[UserInput], + mentioned_plugins: &[crate::plugins::PluginCapabilitySummary], cancellation_token: &CancellationToken, ) -> Option<(Vec, HashSet)> { let turn_context = step_context.turn.as_ref(); @@ -581,30 +731,12 @@ async fn build_skills_and_plugins( return Some((Vec::new(), HashSet::new())); } - let user_input = input - .iter() - .filter_map(|item| match item { - TurnInput::UserInput { content, .. } => Some(content.as_slice()), - TurnInput::ResponseItem(_) | TurnInput::InterAgentCommunication(_) => None, - }) - .flatten() - .cloned() - .collect::>(); let tracking = build_track_events_context( turn_context.model_info.slug.clone(), sess.thread_id.to_string(), turn_context.sub_id.clone(), turn_context.originator.clone(), ); - let loaded_plugins = sess - .services - .plugins_manager - .plugins_for_config(&turn_context.config.plugins_config_input()) - .await; - // Structured plugin:// mentions are resolved from the current session's - // enabled plugins, then converted into turn-scoped guidance below. - let mentioned_plugins = - collect_explicit_plugin_mentions(&user_input, loaded_plugins.capability_summaries()); let connector_snapshot = step_context.mcp.config().connector_snapshot.clone(); let mcp_tools = if turn_context.apps_enabled() || !mentioned_plugins.is_empty() { // Plugin mentions need raw MCP/app inventory even when app tools @@ -629,12 +761,12 @@ async fn build_skills_and_plugins( let skills_outcome = turn_context.turn_skills.snapshot.outcome(); let connector_slug_counts = build_connector_slug_counts(&available_connectors); let extension_injection_items = - build_extension_turn_input_items(sess, step_context, &user_input, cancellation_token) + build_extension_turn_input_items(sess, step_context, user_input, cancellation_token) .await?; let skill_name_counts_lower = build_skill_name_counts(&skills_outcome.skills, &skills_outcome.disabled_paths).1; let mentioned_skills = collect_explicit_skill_mentions( - &user_input, + user_input, &skills_outcome.skills, &skills_outcome.disabled_paths, &connector_slug_counts, @@ -678,8 +810,8 @@ async fn build_skills_and_plugins( &skill_name_counts_lower, ); let plugin_items = - build_plugin_injections(&mentioned_plugins, &mcp_tools, &available_connectors); - let mut explicitly_enabled_connectors = collect_explicit_app_ids(&user_input); + build_plugin_injections(mentioned_plugins, &mcp_tools, &available_connectors); + 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 .iter() @@ -698,7 +830,7 @@ async fn build_skills_and_plugins( sess.services .analytics_events_client .track_app_mentioned(tracking.clone(), mentioned_app_invocations); - for summary in &mentioned_plugins { + for summary in mentioned_plugins { if let Some(plugin) = sess .services .plugins_manager diff --git a/codex-rs/core/tests/suite/plugins.rs b/codex-rs/core/tests/suite/plugins.rs index bfc9ff31c0..10d035a5a9 100644 --- a/codex-rs/core/tests/suite/plugins.rs +++ b/codex-rs/core/tests/suite/plugins.rs @@ -45,6 +45,7 @@ use core_test_support::wait_for_event_match; use core_test_support::wait_for_mcp_server; use pretty_assertions::assert_eq; use tempfile::TempDir; +use test_case::test_case; use wiremock::MockServer; const SAMPLE_PLUGIN_CONFIG_NAME: &str = "sample@test"; @@ -166,6 +167,28 @@ fn write_plugin_mcp_plugin(home: &TempDir, command: &str) { .expect("write plugin mcp config"); } +fn block_plugin_mcp_startup(home: &TempDir, command: &str) -> std::path::PathBuf { + let barrier = home.path().join("allow-plugin-initialize"); + std::fs::write( + sample_plugin_root(home).join(".mcp.json"), + serde_json::to_vec(&serde_json::json!({ + "mcpServers": { + "sample": { + "command": command, + "cwd": ".", + "env": { + "MCP_TEST_INITIALIZE_BARRIER_FILE": barrier, + }, + "startup_timeout_sec": 10, + }, + }, + })) + .expect("serialize blocked plugin MCP configuration"), + ) + .expect("write blocked plugin MCP configuration"); + barrier +} + fn write_plugin_app_plugin(home: &TempDir) { write_plugin_app_plugin_with_name(home, "sample"); } @@ -836,8 +859,20 @@ async fn explicit_plugin_mentions_keep_non_conflicting_mcp_for_chatgpt_auth() -> Ok(()) } +#[derive(Clone, Copy)] +enum ExplicitMcpRequest { + Plugin, + PluginSkill, + ServerMention, + LinkedServerMention, +} + +#[test_case(ExplicitMcpRequest::Plugin; "plugin mention")] +#[test_case(ExplicitMcpRequest::PluginSkill; "plugin skill")] +#[test_case(ExplicitMcpRequest::ServerMention; "MCP server mention")] +#[test_case(ExplicitMcpRequest::LinkedServerMention; "linked MCP server mention")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn explicit_plugin_mentions_use_mcp_for_api_key_dual_surface_plugins() -> Result<()> { +async fn explicitly_requested_mcp_waits_for_startup(request: ExplicitMcpRequest) -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; let mock = mount_plugin_tool_search_turn(&server).await; @@ -850,9 +885,10 @@ async fn explicit_plugin_mentions_use_mcp_for_api_key_dual_surface_plugins() -> return Ok(()); } }; - write_plugin_skill_plugin(codex_home.as_ref()); + let skill_path = std::fs::canonicalize(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 initialize_barrier = block_plugin_mcp_startup(codex_home.as_ref(), &rmcp_test_server_bin); let mut builder = test_codex() .with_home(codex_home) @@ -868,37 +904,68 @@ async fn explicit_plugin_mentions_use_mcp_for_api_key_dual_surface_plugins() -> .await .expect("create new conversation"); let codex = Arc::clone(&test_codex.codex); - wait_for_mcp_server(&codex, "sample").await?; + let input = match request { + ExplicitMcpRequest::Plugin => UserInput::Mention { + name: "sample".into(), + path: format!("plugin://{SAMPLE_PLUGIN_CONFIG_NAME}"), + }, + ExplicitMcpRequest::PluginSkill => UserInput::Skill { + name: "sample:sample-search".into(), + path: skill_path, + }, + ExplicitMcpRequest::ServerMention => UserInput::Mention { + name: "sample".into(), + path: "mcp://sample".into(), + }, + ExplicitMcpRequest::LinkedServerMention => UserInput::Text { + text: "use [$sample](mcp://sample)".to_string(), + text_elements: Vec::new(), + }, + }; codex .submit(Op::UserInput { - items: vec![codex_protocol::user_input::UserInput::Mention { - name: "sample".into(), - path: format!("plugin://{SAMPLE_PLUGIN_CONFIG_NAME}"), - }], + items: vec![input], final_output_json_schema: None, responsesapi_client_metadata: None, additional_context: Default::default(), thread_settings: Default::default(), }) .await?; + tokio::time::sleep(Duration::from_millis(1200)).await; + assert!( + mock.requests().is_empty(), + "an explicitly requested MCP should finish starting before inference" + ); + std::fs::write(initialize_barrier, "ready")?; wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; let requests = mock.requests(); - let request = &requests[0]; - let developer_messages = request.message_input_texts("developer"); - assert!( - developer_messages - .iter() - .any(|text| text.contains("Skills from this plugin")), - "expected plugin skills guidance: {developer_messages:?}" - ); - assert!( - developer_messages - .iter() - .any(|text| text.contains("MCP servers from this plugin")), - "expected visible plugin MCP guidance: {developer_messages:?}" - ); + let model_request = &requests[0]; + let developer_messages = model_request.message_input_texts("developer"); + if matches!(request, ExplicitMcpRequest::Plugin) { + assert!( + developer_messages + .iter() + .any(|text| text.contains("Skills from this plugin")), + "expected plugin skills guidance: {developer_messages:?}" + ); + assert!( + developer_messages + .iter() + .any(|text| text.contains("MCP servers from this plugin")), + "expected visible plugin MCP guidance: {developer_messages:?}" + ); + } + if matches!(request, ExplicitMcpRequest::PluginSkill) { + let user_messages = model_request.message_input_texts("user"); + assert!( + user_messages + .iter() + .any(|message| message.contains("sample:sample-search")), + "expected explicitly requested skill instructions: {user_messages:?}" + ); + } assert!( !developer_messages .iter() @@ -906,7 +973,7 @@ async fn explicit_plugin_mentions_use_mcp_for_api_key_dual_surface_plugins() -> "expected plugin app guidance to be suppressed for API-key auth: {developer_messages:?}" ); assert!( - request + model_request .tool_by_name(SAMPLE_PLUGIN_APP_NAMESPACE, SEARCH_CALENDAR_CREATE_TOOL) .is_none(), "plugin app tool should not leak into the request for API-key auth" diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index addaa54c70..2dba184f7e 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -26,7 +26,9 @@ use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::Environment; use codex_exec_server::HttpRedirectPolicy; use codex_exec_server::HttpRequestParams; +use codex_features::Feature; use codex_login::CodexAuth; +use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::MCP_SANDBOX_STATE_META_CAPABILITY; use codex_mcp::SandboxState; use codex_models_manager::manager::RefreshStrategy; @@ -46,6 +48,7 @@ use codex_protocol::openai_models::TruncationPolicyConfig; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::McpInvocation; +use codex_protocol::protocol::McpStartupStatus; use codex_protocol::protocol::McpToolCallBeginEvent; use codex_protocol::protocol::Op; use codex_protocol::protocol::RolloutItem; @@ -808,6 +811,89 @@ server_names = ["history", "notes"] Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn apps_enabled_turn_skips_pending_optional_mcp_without_cached_tools() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let apps_server = AppsTestServer::mount(&server).await?; + let apps_base_url = apps_server.chatgpt_base_url.clone(); + let response_mock = mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_assistant_message("msg-1", "done"), + responses::ev_completed("resp-1"), + ]), + ) + .await; + let pending_mcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let pending_mcp_url = format!("http://{}/mcp", pending_mcp_listener.local_addr()?); + + let fixture = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(move |config| { + config + .features + .enable(Feature::Apps) + .expect("test config should allow Apps override"); + config.chatgpt_base_url = apps_base_url; + insert_mcp_server( + config, + "pending_optional", + McpServerTransportConfig::StreamableHttp { + url: pending_mcp_url, + bearer_token_env_var: None, + http_headers: None, + env_http_headers: None, + }, + TestMcpServerOptions::default(), + ); + }) + .build_with_auto_env(&server) + .await?; + + let (_pending_mcp_connection, _) = + tokio::time::timeout(Duration::from_secs(5), pending_mcp_listener.accept()) + .await + .context("optional MCP startup should connect before the first turn")??; + tokio::time::timeout(Duration::from_secs(5), async { + loop { + let event = fixture + .codex + .next_event() + .await + .context("event stream ended before Codex Apps became ready")?; + if let EventMsg::McpStartupUpdate(update) = event.msg + && update.server == CODEX_APPS_MCP_SERVER_NAME + && matches!(update.status, McpStartupStatus::Ready) + { + break Ok::<(), anyhow::Error>(()); + } + } + }) + .await + .context("Codex Apps should finish starting before the first turn")??; + + tokio::time::timeout(Duration::from_secs(5), fixture.submit_turn("hello")) + .await + .context("a pending optional MCP must not block the first turn")??; + let body = response_mock.single_request().body_json(); + assert!(body["input"].to_string().contains("")); + let tools = body["tools"].as_array().expect("model request tools"); + assert!(tools.iter().all(|tool| { + tool.get("name") + .or_else(|| tool.get("type")) + .and_then(Value::as_str) + .is_none_or(|name| !name.starts_with("mcp__pending_optional")) + })); + + tokio::time::timeout(Duration::from_secs(2), fixture.codex.shutdown_and_wait()) + .await + .context("shutdown should cancel pending optional MCP startup")??; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn shutdown_cancels_startup_prewarm_waiting_for_mcp_startup() -> anyhow::Result<()> { skip_if_no_network!(Ok(()));