From 460b63e5f43cb1cc8e992ae9c3235ec22d918613 Mon Sep 17 00:00:00 2001 From: Matthew Zeng Date: Thu, 3 Sep 2026 00:07:26 +0000 Subject: [PATCH] Honor explicit plugin mentions during MCP startup (#42406) ## Why Selected plugin MCP servers may be skipped after the optional startup grace period. When a user explicitly mentions a plugin or MCP server, the turn should continue waiting so the requested tools are available to the model. ## What changed - Track explicitly mentioned plugin IDs and MCP servers across later samples, steered input, and same-turn restarts. - Require matching selected-plugin MCP servers to finish startup before capturing the tool catalog, while leaving unmentioned servers optional. - Match `plugin://` identities independently of display names and ignore targeting query parameters such as `app` and `browserFamily`. ## Testing Added coverage for linked and structured plugin mentions, direct and batched capability discovery, delayed environment attachment, steering, and same-turn restarts. GitOrigin-RevId: c509ef8df45df1a7edb0da6afdb2682dc53a4de8 --- .../suite/v2/selected_capability_stack.rs | 420 +++++++++++++++--- .../src/connection_manager/tool_catalog.rs | 7 +- .../codex-mcp/src/connection_manager_tests.rs | 85 +++- codex-rs/codex-mcp/src/runtime.rs | 45 +- codex-rs/core/src/plugins/mentions.rs | 25 +- codex-rs/core/src/plugins/mentions_tests.rs | 37 +- codex-rs/core/src/plugins/mod.rs | 1 + codex-rs/core/src/session/mcp.rs | 3 +- codex-rs/core/src/session/mod.rs | 12 +- codex-rs/core/src/session/tests.rs | 1 + codex-rs/core/src/session/turn.rs | 52 ++- codex-rs/core/src/tasks/regular.rs | 3 + codex-rs/skills/src/mentions.rs | 3 + codex-rs/skills/src/mentions_tests.rs | 25 ++ 14 files changed, 615 insertions(+), 104 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs b/codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs index 79c92d860058..cff3e7673323 100644 --- a/codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs +++ b/codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs @@ -5,12 +5,15 @@ use anyhow::Context; use anyhow::Result; use app_test_support::ChatGptAuthFixture; use app_test_support::TestAppServer; +use app_test_support::create_final_assistant_message_sse_response; +use app_test_support::create_request_user_input_sse_response; use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use app_test_support::write_mock_responses_config_toml_with_chatgpt_base_url; use codex_app_server_protocol::AppInfo; use codex_app_server_protocol::CapabilityRootLocation; use codex_app_server_protocol::EnvironmentAddResponse; +use codex_app_server_protocol::EnvironmentInfoResponse; use codex_app_server_protocol::ListMcpServerStatusParams; use codex_app_server_protocol::ListMcpServerStatusResponse; use codex_app_server_protocol::RequestId; @@ -22,6 +25,9 @@ use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnEnvironmentParams; use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; +use codex_app_server_protocol::TurnSteerParams; +use codex_app_server_protocol::TurnSteerResponse; use codex_app_server_protocol::UserInput; use codex_config::types::AuthCredentialsStoreMode; use codex_exec_server::LOCAL_ENVIRONMENT_ID; @@ -64,6 +70,122 @@ const MCP_SERVER_NAME: &str = "executor_probe"; const MCP_CALL_ID: &str = "selected-executor-mcp-call"; const CONNECTOR_ID: &str = "calendar"; +#[derive(Clone, Copy)] +enum PluginMention { + Unmentioned, + Link, + Structured, +} + +#[test_case(PluginMention::Unmentioned, false; "optional direct")] +#[test_case(PluginMention::Unmentioned, true; "optional batched")] +#[test_case(PluginMention::Link, false; "plugin link direct")] +#[test_case(PluginMention::Link, true; "plugin link batched")] +#[test_case(PluginMention::Structured, false; "structured mention direct")] +#[test_case(PluginMention::Structured, true; "structured mention batched")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn selected_plugin_mcp_startup_respects_explicit_mentions( + mention: PluginMention, + executor_capability_discovery: bool, +) -> Result<()> { + let explicitly_mentioned = !matches!(mention, PluginMention::Unmentioned); + let responses_server = responses::start_mock_server().await; + let (apps_url, apps_server_handle) = + start_apps_server_with_delays(Vec::new(), Vec::new(), Duration::ZERO, Duration::ZERO) + .await?; + let fixture = selected_capability_fixture(&responses_server.uri(), &apps_url)?; + let config_path = fixture.codex_home.path().join("config.toml"); + let config = std::fs::read_to_string(&config_path)?.replace( + "executor_capability_discovery = true", + &format!("executor_capability_discovery = {executor_capability_discovery}"), + ); + std::fs::write(config_path, config)?; + let initialize_barrier = fixture.block_mcp_startup()?; + let response_mock = responses::mount_sse_once( + &responses_server, + create_final_assistant_message_sse_response("Done")?, + ) + .await; + let mut app_server = TestAppServer::builder() + .with_codex_home(fixture.codex_home.path()) + // This fixture owns environments.toml and selects its environments explicitly. + .without_auto_env() + .build() + .await?; + timeout(READ_TIMEOUT, app_server.initialize()).await??; + let thread_id = start_thread( + &mut app_server, + fixture.selected_root, + fixture.environment_cwd.clone(), + ) + .await?; + let mut exec_server = + spawn_exec_server(fixture.codex_home.path(), &fixture.exec_server_url).await?; + add_environment(&mut app_server, &fixture.exec_server_url).await?; + + let text_input = |text| UserInput::Text { + text, + text_elements: Vec::new(), + }; + let input = match mention { + PluginMention::Unmentioned => text_input("Answer without using tools".to_string()), + PluginMention::Link => { + text_input(format!("Use [@executor-demo](plugin://{PLUGIN_ID}) now")) + } + PluginMention::Structured => UserInput::Mention { + name: PLUGIN_DISPLAY_NAME.to_string(), + path: format!("plugin://{PLUGIN_ID}?app=com.example.editor"), + }, + }; + let request_id = app_server + .send_turn_start_request(TurnStartParams { + thread_id, + input: vec![input], + environments: Some(vec![TurnEnvironmentParams { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: fixture.environment_cwd.into(), + runtime_workspace_roots: None, + }]), + ..Default::default() + }) + .await?; + timeout( + READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + wait_for_pid_file(&fixture.pid_file).await?; + if explicitly_mentioned { + // An explicit mention must outwait the optional one-second grace. + tokio::time::sleep(Duration::from_secs(3)).await; + assert!(response_mock.requests().is_empty()); + std::fs::write(&initialize_barrier, "ready")?; + } + let turn_timeout = if explicitly_mentioned { + READ_TIMEOUT + } else { + Duration::from_secs(5) + }; + timeout( + turn_timeout, + app_server.read_stream_until_notification_message("turn/completed"), + ) + .await??; + let request = response_mock.single_request(); + assert_selected_skill_catalog_available(&request); + assert_eq!( + request + .tool_by_name(&format!("mcp__{MCP_SERVER_NAME}"), "echo") + .is_some(), + explicitly_mentioned, + ); + + exec_server.kill().await?; + apps_server_handle.abort(); + let _ = apps_server_handle.await; + Ok(()) +} + #[test_case(false; "direct selected root discovery")] #[test_case(true; "batched executor capability discovery")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -257,10 +379,11 @@ async fn selected_capability_stack_tracks_environment_availability_and_resume() add_environment(&mut app_server, &fixture.exec_server_url).await?; wait_for_selected_mcp_server(&mut app_server, &thread_id).await?; + // A skill mention alone does not wait for MCP startup. run_turn( &mut app_server, &thread_id, - &format!("Use ${SKILL_NAME} and call its selected executor MCP"), + &format!("Use ${SKILL_NAME} and call [${MCP_SERVER_NAME}](mcp://{MCP_SERVER_NAME})"), fixture.environment_cwd.clone(), ) .await?; @@ -322,7 +445,9 @@ async fn selected_capability_stack_tracks_environment_availability_and_resume() run_turn( &mut app_server, &thread_id, - &format!("Use ${SKILL_NAME} after reattaching the selected executor"), + &format!( + "Use ${SKILL_NAME} with [${MCP_SERVER_NAME}](mcp://{MCP_SERVER_NAME}) after reattaching the selected executor" + ), fixture.environment_cwd, ) .await?; @@ -352,10 +477,44 @@ async fn selected_capability_stack_tracks_environment_availability_and_resume() Ok(()) } +#[derive(Clone, Copy)] +enum MentionTiming { + Unmentioned, + Initial, + InitialBeforeRestart, + Steered, +} + +#[derive(Clone, Copy)] +enum MentionTarget { + Plugin, + Server, +} + +#[test_case(MentionTiming::Unmentioned, MentionTarget::Plugin; "without a mention")] +#[test_case(MentionTiming::Initial, MentionTarget::Plugin; "with an initial plugin mention")] +#[test_case(MentionTiming::InitialBeforeRestart, MentionTarget::Plugin; "with an initial plugin mention across a restart")] +#[test_case(MentionTiming::Steered, MentionTarget::Plugin; "with a steered plugin mention")] +#[test_case(MentionTiming::Initial, MentionTarget::Server; "with an initial server mention")] +#[test_case(MentionTiming::Steered, MentionTarget::Server; "with a steered server mention")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn selected_capabilities_become_available_between_samples_in_one_turn() -> Result<()> { +async fn selected_capabilities_become_available_between_samples_in_one_turn( + mention_timing: MentionTiming, + mention_target: MentionTarget, +) -> Result<()> { const USER_INPUT_CALL_ID: &str = "pause-for-environment"; + const MCP_USER_INPUT_CALL_ID: &str = "pause-for-mcp"; + let mention_link = match mention_target { + MentionTarget::Plugin => format!("[@executor-demo](plugin://{PLUGIN_ID})"), + MentionTarget::Server => format!("[${MCP_SERVER_NAME}](mcp://{MCP_SERVER_NAME})"), + }; + let samples_before_attach = if matches!(mention_timing, MentionTiming::Steered) { + 2 + } else { + 1 + }; + let continue_answers = json!({ "answers": { "confirm_path": { "answers": ["yes"] } } }); let responses_server = responses::start_mock_server().await; let (apps_url, apps_server_handle) = start_apps_server_with_delays( vec![AppInfo { @@ -381,33 +540,43 @@ async fn selected_capabilities_become_available_between_samples_in_one_turn() -> ) .await?; let fixture = selected_capability_fixture(&responses_server.uri(), &apps_url)?; + let initialize_barrier = fixture.block_mcp_startup()?; + let stop_hook_barrier = if matches!(mention_timing, MentionTiming::InitialBeforeRestart) { + let barrier = fixture.codex_home.path().join("allow-stop-hook"); + let hook_path = fixture.codex_home.path().join("wait-for-stop.py"); + std::fs::write( + &hook_path, + "import json\nimport sys\nimport time\nfrom pathlib import Path\n\n\ + json.load(sys.stdin)\nwhile not Path(sys.argv[1]).exists():\n time.sleep(0.01)\n", + )?; + let command = toml::Value::String(format!( + "python3 \"{}\" \"{}\"", + hook_path.display(), + barrier.display() + )); + std::fs::write( + fixture.codex_home.path().join("requirements.toml"), + format!( + "[[hooks.Stop]]\n[[hooks.Stop.hooks]]\ntype = 'command'\ncommand = {command}\ntimeout = 60\n" + ), + )?; + let config_path = fixture.codex_home.path().join("config.toml"); + let config = std::fs::read_to_string(&config_path)? + .replace("[features]\n", "[features]\nhooks = true\n"); + std::fs::write(config_path, config)?; + Some(barrier) + } else { + None + }; let response_mock = responses::mount_sse_sequence( &responses_server, vec![ - responses::sse(vec![ - responses::ev_response_created("environment-pending"), - responses::ev_function_call( - USER_INPUT_CALL_ID, - "request_user_input", - &json!({ - "questions": [{ - "id": "continue", - "header": "Continue", - "question": "Continue after the executor is attached?", - "options": [{ - "label": "Yes (Recommended)", - "description": "Continue the same turn." - }, { - "label": "No", - "description": "Stop here." - }] - }], - "autoResolutionMs": 60_000 - }) - .to_string(), - ), - responses::ev_completed("environment-pending"), - ]), + if stop_hook_barrier.is_some() { + create_final_assistant_message_sse_response("Waiting for the executor")? + } else { + create_request_user_input_sse_response(USER_INPUT_CALL_ID)? + }, + create_request_user_input_sse_response(MCP_USER_INPUT_CALL_ID)?, responses::sse(vec![ responses::ev_response_created("environment-ready-call"), responses::ev_function_call_with_namespace( @@ -446,9 +615,16 @@ async fn selected_capabilities_become_available_between_samples_in_one_turn() -> .await?; let turn_start_id = app_server .send_turn_start_request(TurnStartParams { - thread_id, + thread_id: thread_id.clone(), input: vec![UserInput::Text { - text: "Use the executor when it becomes ready.".to_string(), + text: match mention_timing { + MentionTiming::Unmentioned | MentionTiming::Steered => { + "Use the executor when it becomes ready.".to_string() + } + MentionTiming::Initial | MentionTiming::InitialBeforeRestart => { + format!("Use {mention_link} when its executor becomes ready.") + } + }, text_elements: Vec::new(), }], environments: Some(vec![TurnEnvironmentParams { @@ -467,34 +643,145 @@ async fn selected_capabilities_become_available_between_samples_in_one_turn() -> ..Default::default() }) .await?; - timeout( + let response = timeout( READ_TIMEOUT, app_server.read_stream_until_response_message(RequestId::Integer(turn_start_id)), ) .await??; + let TurnStartResponse { turn } = to_response(response)?; - let request = timeout(READ_TIMEOUT, app_server.read_stream_until_request_message()).await??; - let ServerRequest::ToolRequestUserInput { request_id, .. } = request else { - panic!("expected request_user_input, got {request:?}"); + let request_id = if stop_hook_barrier.is_some() { + // Steering during a successful Stop hook restarts run_turn in the same RegularTask. + timeout( + READ_TIMEOUT, + app_server.read_stream_until_matching_notification( + "Stop hook started", + |notification| { + notification.method == "hook/started" + && notification.params.as_ref().is_some_and(|params| { + params["turnId"] == turn.id && params["run"]["eventName"] == "stop" + }) + }, + ), + ) + .await??; + None + } else { + let request = + timeout(READ_TIMEOUT, app_server.read_stream_until_request_message()).await??; + let ServerRequest::ToolRequestUserInput { request_id, .. } = request else { + panic!("expected request_user_input, got {request:?}"); + }; + Some(request_id) }; let requests = response_mock.requests(); assert_eq!(1, requests.len()); assert_selected_capabilities_absent(&requests[0]); + if !matches!(mention_timing, MentionTiming::Unmentioned) { + let steer_request_id = app_server + .send_turn_steer_request(TurnSteerParams { + thread_id, + input: vec![UserInput::Text { + text: if matches!(mention_timing, MentionTiming::Steered) { + format!("Use {mention_link} now.") + } else { + "Continue with the requested tool when it becomes available.".to_string() + }, + text_elements: Vec::new(), + }], + expected_turn_id: turn.id.clone(), + ..Default::default() + }) + .await?; + let response = timeout( + READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(steer_request_id)), + ) + .await??; + let response: TurnSteerResponse = to_response(response)?; + assert_eq!(response, TurnSteerResponse { turn_id: turn.id }); + } + let request_id = if matches!(mention_timing, MentionTiming::Steered) { + // Consume the mention while the executor is absent, then sample once more. + app_server + .send_response( + request_id.expect("steering pauses before attaching the executor"), + continue_answers.clone(), + ) + .await?; + let request = + timeout(READ_TIMEOUT, app_server.read_stream_until_request_message()).await??; + let ServerRequest::ToolRequestUserInput { request_id, .. } = request else { + panic!("expected request_user_input, got {request:?}"); + }; + let requests = response_mock.requests(); + assert_eq!(2, requests.len()); + assert_selected_capabilities_absent(&requests[1]); + Some(request_id) + } else { + request_id + }; let mut exec_server = spawn_exec_server(fixture.codex_home.path(), &fixture.exec_server_url).await?; add_environment(&mut app_server, &fixture.exec_server_url).await?; - tokio::time::sleep(Duration::from_millis(200)).await; - app_server - .send_response( - request_id, - json!({ - "answers": { - "continue": { "answers": ["yes"] } - } - }), - ) - .await?; + if let Some(request_id) = request_id { + app_server + .send_response(request_id, continue_answers.clone()) + .await?; + } + if let Some(stop_hook_barrier) = stop_hook_barrier { + std::fs::write(stop_hook_barrier, "ready")?; + } + let mcp_pid = wait_for_pid_file(&fixture.pid_file).await?; + if !matches!(mention_timing, MentionTiming::Unmentioned) { + // An explicit mention must still require startup after the executor attaches. + tokio::time::sleep(Duration::from_secs(3)).await; + assert_eq!( + response_mock.requests().len(), + samples_before_attach, + "the explicit mention must keep later samples waiting for MCP startup" + ); + std::fs::write(&initialize_barrier, "ready")?; + } + let request_id = if matches!(mention_timing, MentionTiming::Steered) { + None + } else { + let request = + timeout(READ_TIMEOUT, app_server.read_stream_until_request_message()).await??; + let ServerRequest::ToolRequestUserInput { request_id, .. } = request else { + panic!("expected request_user_input, got {request:?}"); + }; + let requests = response_mock.requests(); + assert_eq!(2, requests.len()); + assert_selected_skill_catalog_available(&requests[1]); + if matches!(mention_timing, MentionTiming::Unmentioned) { + assert!( + requests[1] + .tool_by_name(&format!("mcp__{MCP_SERVER_NAME}"), "echo") + .is_none() + ); + std::fs::write(&initialize_barrier, "ready")?; + } else { + assert_selected_plugin_tools(&requests[1]); + } + Some(request_id) + }; + timeout( + READ_TIMEOUT, + app_server.read_stream_until_matching_notification("selected MCP ready", |notification| { + notification.method == "mcpServer/startupStatus/updated" + && notification.params.as_ref().is_some_and(|params| { + params["name"] == MCP_SERVER_NAME && params["status"] == "ready" + }) + }), + ) + .await??; + if let Some(request_id) = request_id { + app_server + .send_response(request_id, continue_answers) + .await?; + } timeout( READ_TIMEOUT, app_server.read_stream_until_notification_message("turn/completed"), @@ -502,19 +789,21 @@ async fn selected_capabilities_become_available_between_samples_in_one_turn() -> .await??; let requests = response_mock.requests(); - assert_eq!(3, requests.len()); - assert_selected_skill_catalog_available(&requests[1]); - assert_selected_plugin_tools(&requests[1]); - assert_plugin_guidance_count(&requests[1], /*expected_count*/ 0); - assert_selected_plugin_tools(&requests[2]); - assert_plugin_guidance_count(&requests[2], /*expected_count*/ 0); - let output = requests[2].function_call_output(MCP_CALL_ID); + assert_eq!(4, requests.len()); + for request in &requests[samples_before_attach..] { + assert_selected_skill_catalog_available(request); + assert_plugin_guidance_count(request, /*expected_count*/ 0); + } + for request in &requests[2..] { + assert_selected_plugin_tools(request); + } + let output = requests[3].function_call_output(MCP_CALL_ID); let output = output["output"] .as_str() .expect("MCP function output should be text"); assert!(output.contains("ECHOING: same turn")); assert!(output.contains(EXECUTOR_ENV_VALUE)); - wait_for_pid_file(&fixture.pid_file).await?; + assert_eq!(mcp_pid, wait_for_pid_file(&fixture.pid_file).await?); exec_server.kill().await?; apps_server_handle.abort(); @@ -531,6 +820,20 @@ struct SelectedCapabilityFixture { environment_cwd: AbsolutePathBuf, } +impl SelectedCapabilityFixture { + fn block_mcp_startup(&self) -> Result { + let initialize_barrier = self._plugin.path().join("allow-mcp-initialize"); + let mcp_config_path = self._plugin.path().join(".mcp.json"); + let mut mcp_config: serde_json::Value = + serde_json::from_slice(&std::fs::read(&mcp_config_path)?)?; + mcp_config["mcpServers"][MCP_SERVER_NAME]["env"]["MCP_TEST_INITIALIZE_BARRIER_FILE"] = + json!(initialize_barrier); + mcp_config["mcpServers"][MCP_SERVER_NAME]["startup_timeout_sec"] = json!(30); + std::fs::write(mcp_config_path, serde_json::to_vec(&mcp_config)?)?; + Ok(initialize_barrier) + } +} + fn selected_capability_fixture( responses_server_uri: &str, apps_url: &str, @@ -799,6 +1102,19 @@ async fn add_environment(app_server: &mut TestAppServer, exec_server_url: &str) ) .await??; let _: EnvironmentAddResponse = to_response(response)?; + // Wait for executor attachment, independently of plugin MCP startup. + let request_id = app_server + .send_raw_request( + "environment/info", + Some(json!({ "environmentId": EXECUTOR_ID })), + ) + .await?; + let response = timeout( + READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let _: EnvironmentInfoResponse = to_response(response)?; Ok(()) } 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 36ad04e1730f..616ec1729f74 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -1,3 +1,4 @@ +use std::collections::HashSet; use std::sync::Arc; use std::sync::atomic::Ordering; use std::time::Instant; @@ -170,6 +171,7 @@ impl McpConnectionSet { config: Arc, plugins_available: bool, required_servers: &[String], + required_plugins: &HashSet, ) -> McpBinding { let revision = self.tool_catalog_revision.read().await; let mut listed_tools = Vec::new(); @@ -190,10 +192,13 @@ impl McpConnectionSet { let has_cached_tools = cached_tools.is_some(); let must_wait_for_startup = (required && (!view.connection.startup_is_dormant() || !has_cached_tools)) - || self.is_selected_plugin_mcp_server(server_name) || required_servers .iter() .any(|required| required == server_name) + || (self.is_selected_plugin_mcp_server(server_name) + && self + .plugin_id_for_mcp_server_name(server_name) + .is_some_and(|plugin_id| required_plugins.contains(plugin_id))) || (server_name == CODEX_APPS_MCP_SERVER_NAME && !has_cached_tools); if !must_wait_for_startup && has_cached_tools { return (server_name, view, cached_tools); diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 9c974ac4665d..0c0e57306e51 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -231,6 +231,7 @@ async fn capture_binding(manager: &Arc) -> McpBinding { Arc::new(config), /*plugins_available*/ false, /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), ) .await } @@ -2592,10 +2593,17 @@ async fn capture_binding_skips_pending_optional_servers_after_configured_shared_ serde_json::from_value(serde_json::json!({ "command": "optional-plugin" })) .expect("optional plugin MCP config"), )); + catalog.register(crate::McpServerRegistration::from_selected_plugin( + "pending-selected".to_string(), + crate::McpPluginAttribution::new("selected-plugin".to_string(), "Selected".to_string()), + /*selection_order*/ 0, + serde_json::from_value(serde_json::json!({ "command": "selected-plugin" })) + .expect("selected plugin MCP config"), + )); plugin_config.mcp_server_catalog = catalog.build(); plugin_config.optional_mcp_startup_grace = Duration::from_millis(250); manager.tool_plugin_provenance = Arc::new(crate::tool_plugin_provenance(&plugin_config)); - for server_name in ["pending-one", "pending-two"] { + for server_name in ["pending-one", "pending-two", "pending-selected"] { manager.insert_test_client( server_name.to_string(), AsyncManagedClient { @@ -2613,19 +2621,34 @@ async fn capture_binding_skips_pending_optional_servers_after_configured_shared_ ); } + let mut required_manager = McpConnectionSet::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + required_manager.tool_plugin_provenance = Arc::clone(&manager.tool_plugin_provenance); + required_manager.insert_test_client( + "pending-selected", + manager.test_client("pending-selected").clone(), + ); + required_manager.required_servers = vec!["pending-selected".to_string()]; + let manager = Arc::new(manager); assert_eq!(manager.stable_catalog_revision().await, None); + let started = tokio::time::Instant::now(); let binding = tokio::time::timeout( Duration::from_millis(500), manager.capture_binding_with_metadata( Arc::new(plugin_config), /*plugins_available*/ false, /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), ), ) .await .expect("all optional servers should share the configured startup grace"); assert!(binding.tools().is_empty()); + assert_eq!(started.elapsed(), Duration::from_millis(250)); let binding = tokio::time::timeout(Duration::from_millis(1), capture_binding(&manager)) .await @@ -2651,17 +2674,52 @@ async fn capture_binding_skips_pending_optional_servers_after_configured_shared_ "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"); + for server_name in ["pending-one", "pending-selected"] { + let required_servers = vec![server_name.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, + /*required_plugins*/ &HashSet::new(), + ), + ) + .await; + assert!(binding.is_err(), "explicitly requested servers must wait"); + } + // A plugin mention must still require startup after the optional grace has elapsed. + for (plugin_id, must_wait) in [ + ("selected-plugin", true), + ("optional-plugin", false), + ("selected-plugin-other", false), + ] { + let required_plugins = HashSet::from([plugin_id.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*/ &[], + &required_plugins, + ), + ) + .await; + assert_eq!( + binding.is_err(), + must_wait, + "plugin requirement {plugin_id}" + ); + } + assert!( + tokio::time::timeout( + Duration::from_millis(1500), + capture_binding(&Arc::new(required_manager)), + ) + .await + .is_err(), + "configured-required selected plugin servers must wait beyond the optional grace" + ); } #[tokio::test(start_paused = true)] @@ -2690,6 +2748,7 @@ async fn capture_binding_waits_for_optional_startup_when_shared_grace_is_disable Arc::new(config), /*plugins_available*/ false, /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), ) .await }); @@ -2818,6 +2877,7 @@ async fn capture_binding_shares_optional_startup_grace_across_connection_sets() Arc::new(disabled_config), /*plugins_available*/ false, /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), ), ) .await @@ -2845,6 +2905,7 @@ async fn capture_binding_shares_optional_startup_grace_across_connection_sets() Arc::new(updated_config), /*plugins_available*/ false, /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), ), ) .await diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index dc30070e9092..7a55d2cc526d 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -6,6 +6,7 @@ //! [`crate::connection_manager`]. use std::collections::HashMap; +use std::collections::HashSet; use std::path::PathBuf; use std::sync::Arc; use std::sync::Mutex; @@ -335,20 +336,29 @@ 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 + self.current_binding_with_requirements(&[], &HashSet::new()) + .await } - /// Captures the latest runtime, waiting for servers explicitly required by this turn. - pub async fn current_binding_with_required_servers( + /// Captures one runtime, waiting for explicitly required servers and selected plugins. + /// Plugin IDs are resolved by the captured connection set, even if a refresh publishes later. + pub async fn current_binding_with_requirements( &self, required_servers: &[String], + required_plugins: &HashSet, ) -> Option> { - Self::binding_from_published_runtime(self.current.load_full(), required_servers).await + Self::binding_from_published_runtime( + self.current.load_full(), + required_servers, + required_plugins, + ) + .await } async fn binding_from_published_runtime( current: Arc, required_servers: &[String], + required_plugins: &HashSet, ) -> Option> { let config = Arc::clone(current.config.as_ref()?); let stable_catalog_revision = current.connections.stable_catalog_revision().await; @@ -367,7 +377,12 @@ impl McpRuntime { let binding = Arc::new( current .connections - .capture_binding_with_metadata(config, current.plugins_available, required_servers) + .capture_binding_with_metadata( + config, + current.plugins_available, + required_servers, + required_plugins, + ) .await, ); if let Some(catalog_revision) = stable_catalog_revision @@ -445,7 +460,12 @@ impl McpRuntime { if !current.connections.wait_for_server_startup(server).await { return None; } - Self::binding_from_published_runtime(current, /*required_servers*/ &[]).await + Self::binding_from_published_runtime( + current, + /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), + ) + .await } /// Returns the latest published configuration without waiting for clients. @@ -877,12 +897,14 @@ mod tests { let first = McpRuntime::binding_from_published_runtime( Arc::clone(&published), /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), ) .await .expect("first binding"); let repeated = McpRuntime::binding_from_published_runtime( Arc::clone(&published), /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), ) .await .expect("repeated binding"); @@ -893,10 +915,13 @@ mod tests { cached_binding: Mutex::new(None), ..previous }); - let refreshed = - McpRuntime::binding_from_published_runtime(republished, /*required_servers*/ &[]) - .await - .expect("republished binding"); + let refreshed = McpRuntime::binding_from_published_runtime( + republished, + /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), + ) + .await + .expect("republished binding"); assert!(!Arc::ptr_eq(&first, &refreshed)); } diff --git a/codex-rs/core/src/plugins/mentions.rs b/codex-rs/core/src/plugins/mentions.rs index 13563806dad8..fb43b1313096 100644 --- a/codex-rs/core/src/plugins/mentions.rs +++ b/codex-rs/core/src/plugins/mentions.rs @@ -68,6 +68,19 @@ pub(crate) fn collect_explicit_plugin_mentions( return Vec::new(); } + // `config_name` stores the full plugin ID, not its display or mention name. + let mentioned_plugin_ids = collect_explicit_plugin_ids(input); + plugins + .iter() + .filter(|plugin| mentioned_plugin_ids.contains(plugin.config_name.as_str())) + .cloned() + .collect() +} + +/// Collect exact IDs from explicit `plugin://` references, independently of display names. +/// +/// Host plugins use `@` IDs; selected roots may supply opaque IDs. +pub(crate) fn collect_explicit_plugin_ids(input: &[UserInput]) -> HashSet { let messages = input .iter() .filter_map(|item| match item { @@ -76,7 +89,7 @@ pub(crate) fn collect_explicit_plugin_mentions( }) .collect::>(); - let mentioned_config_names: HashSet = input + input .iter() .filter_map(|item| match item { UserInput::Mention { path, .. } => Some(path.clone()), @@ -89,16 +102,6 @@ pub(crate) fn collect_explicit_plugin_mentions( ) .filter(|path| tool_kind_for_path(path.as_str()) == ToolMentionKind::Plugin) .filter_map(|path| plugin_config_name_from_path(path.as_str()).map(str::to_string)) - .collect(); - - if mentioned_config_names.is_empty() { - return Vec::new(); - } - - plugins - .iter() - .filter(|plugin| mentioned_config_names.contains(plugin.config_name.as_str())) - .cloned() .collect() } diff --git a/codex-rs/core/src/plugins/mentions_tests.rs b/codex-rs/core/src/plugins/mentions_tests.rs index 8184a4194568..ee0747eeb50b 100644 --- a/codex-rs/core/src/plugins/mentions_tests.rs +++ b/codex-rs/core/src/plugins/mentions_tests.rs @@ -4,6 +4,7 @@ use codex_protocol::user_input::UserInput; use pretty_assertions::assert_eq; use super::collect_explicit_app_ids; +use super::collect_explicit_plugin_ids; use super::collect_explicit_plugin_mentions; use crate::plugins::PluginCapabilitySummary; @@ -77,35 +78,39 @@ fn collect_explicit_app_ids_ignores_non_app_paths() { #[test] fn collect_explicit_plugin_mentions_from_structured_paths() { + let expected = plugin("sample@test", "Sample"); let plugins = vec![ - plugin("sample@test", "sample"), - plugin("other@test", "other"), + expected.clone(), + plugin("sample@other-marketplace", "Sample"), + plugin("other@test", "alias"), ]; let mentioned = collect_explicit_plugin_mentions( &[UserInput::Mention { - name: "sample".to_string(), + name: "alias".to_string(), path: "plugin://sample@test".to_string(), }], &plugins, ); - assert_eq!(mentioned, vec![plugin("sample@test", "sample")]); + assert_eq!(mentioned, vec![expected]); } #[test] fn collect_explicit_plugin_mentions_from_linked_text_mentions() { + let expected = plugin("sample@test", "Sample"); let plugins = vec![ - plugin("sample@test", "sample"), - plugin("other@test", "other"), + expected.clone(), + plugin("sample@other-marketplace", "Sample"), + plugin("other@test", "alias"), ]; let mentioned = collect_explicit_plugin_mentions( - &[text_input("use [@sample](plugin://sample@test)")], + &[text_input("use [@alias](plugin://sample@test)")], &plugins, ); - assert_eq!(mentioned, vec![plugin("sample@test", "sample")]); + assert_eq!(mentioned, vec![expected]); } #[test] @@ -154,3 +159,19 @@ fn collect_explicit_plugin_mentions_ignores_dollar_linked_plugin_mentions() { assert_eq!(mentioned, Vec::::new()); } + +#[test] +fn collect_explicit_plugin_ids_preserves_opaque_ids() { + let input = [ + text_input("use [@same](plugin://selected-one)"), + UserInput::Mention { + name: "same".to_string(), + path: "plugin://selected-two".to_string(), + }, + ]; + + assert_eq!( + collect_explicit_plugin_ids(&input), + HashSet::from(["selected-one".to_string(), "selected-two".to_string()]), + ); +} diff --git a/codex-rs/core/src/plugins/mod.rs b/codex-rs/core/src/plugins/mod.rs index ca0e71ce4b81..e47c8e28e666 100644 --- a/codex-rs/core/src/plugins/mod.rs +++ b/codex-rs/core/src/plugins/mod.rs @@ -23,6 +23,7 @@ pub(crate) use render::render_explicit_plugin_instructions; pub(crate) use mentions::build_connector_slug_counts; pub(crate) use mentions::collect_explicit_app_ids; +pub(crate) use mentions::collect_explicit_plugin_ids; pub(crate) use mentions::collect_explicit_plugin_mentions; pub(crate) use mentions::collect_tool_mentions_from_messages; diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 46af2d336e27..fb79de7d7b32 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -329,6 +329,7 @@ impl Session { turn_context: &TurnContext, selected_capability_roots: &[ResolvedSelectedCapabilityRoot], required_servers: &[String], + required_plugins: &HashSet, ) -> Arc { let ready_selected_capability_roots = Self::ready_selected_capability_roots(selected_capability_roots); @@ -365,7 +366,7 @@ impl Session { if let Some(binding) = self .services .mcp_runtime - .current_binding_with_required_servers(&required_servers) + .current_binding_with_requirements(&required_servers, required_plugins) .await { return binding; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 90fe56ec364d..0cf192be8182 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3459,6 +3459,7 @@ impl Session { turn_context, cancellation_token, /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), ) .await } @@ -3468,9 +3469,15 @@ impl Session { turn_context: Arc, cancellation_token: &CancellationToken, required_servers: &[String], + required_plugins: &HashSet, ) -> CodexResult> { let step_context = self - .capture_step_context_inner(turn_context, cancellation_token, required_servers) + .capture_step_context_inner( + turn_context, + cancellation_token, + required_servers, + required_plugins, + ) .await?; self.set_last_known_step_context(&step_context).await; Ok(step_context) @@ -3487,6 +3494,7 @@ impl Session { turn_context, cancellation_token, /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), ) .await } @@ -3497,6 +3505,7 @@ impl Session { turn_context: Arc, cancellation_token: &CancellationToken, required_servers: &[String], + required_plugins: &HashSet, ) -> CodexResult> { // Capture once before asynchronous planning; all request consumers // retain this immutable settings version even if the turn is updated. @@ -3574,6 +3583,7 @@ impl Session { turn_context.as_ref(), &selected_capability_roots, required_servers, + required_plugins, ), 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 c1b00dfa262a..ffac44c011ac 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -9072,6 +9072,7 @@ async fn refresh_mcp_servers_uses_latest_state_for_existing_turns() { &turn_context, /*selected_capability_roots*/ &[], /*required_servers*/ &[], + /*required_plugins*/ &HashSet::new(), ) .await; diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 72d97c082a72..c49579cb2253 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -138,6 +138,13 @@ use tracing::warn; const POST_SAMPLING_TOKEN_ESTIMATE_TARGET: &str = "codex_core::post_sampling_token_estimate"; +/// Explicit MCP startup requirements retained across restarts within one user turn. +#[derive(Default)] +pub(crate) struct McpStartupRequirements { + required_servers: Vec, + required_plugins: HashSet, +} + /// Takes initial turn input and runs a loop where, at each sampling request, /// the model replies with either: /// @@ -156,6 +163,7 @@ pub(crate) async fn run_turn( sess: Arc, turn_context: Arc, input: Vec, + mcp_startup_requirements: &mut McpStartupRequirements, prewarmed_client_session: Option, cancellation_token: CancellationToken, ) -> CodexResult> { @@ -192,7 +200,16 @@ pub(crate) async fn run_turn( } let user_input = turn_user_input(&input); - let (required_servers, mentioned_plugins) = + let allow_plugin_mentions = + !crate::guardian::is_basic_session_source(&turn_context.session_source); + let McpStartupRequirements { + required_servers, + required_plugins, + } = mcp_startup_requirements; + if allow_plugin_mentions { + required_plugins.extend(crate::plugins::collect_explicit_plugin_ids(&user_input)); + } + let (input_required_servers, mentioned_plugins) = match required_mcp_servers_for_input(&sess, turn_context.as_ref(), &user_input) .or_cancel(&cancellation_token) .await @@ -205,12 +222,17 @@ pub(crate) async fn run_turn( } }; + required_servers.extend(input_required_servers); + required_servers.sort_unstable(); + required_servers.dedup(); + // run_turn owns the step used to seed context and make the first sampling request. let first_step_context = match sess .capture_step_context_with_required_mcp_servers( Arc::clone(&turn_context), &cancellation_token, - &required_servers, + required_servers, + required_plugins, ) .await { @@ -343,24 +365,38 @@ 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, + Some(step_context) if pending_input.is_empty() => step_context, None if pending_input.is_empty() => { - sess.capture_step_context(Arc::clone(&turn_context), &cancellation_token) - .await? + sess.capture_step_context_with_required_mcp_servers( + Arc::clone(&turn_context), + &cancellation_token, + required_servers, + required_plugins, + ) + .await? } - None => { + Some(_) | None => { let pending_user_input = turn_user_input(&pending_input); - let (required_servers, _) = required_mcp_servers_for_input( + if allow_plugin_mentions { + required_plugins.extend(crate::plugins::collect_explicit_plugin_ids( + &pending_user_input, + )); + } + let (pending_required_servers, _) = required_mcp_servers_for_input( &sess, turn_context.as_ref(), &pending_user_input, ) .or_cancel(&cancellation_token) .await?; + required_servers.extend(pending_required_servers); + required_servers.sort_unstable(); + required_servers.dedup(); sess.capture_step_context_with_required_mcp_servers( Arc::clone(&turn_context), &cancellation_token, - &required_servers, + required_servers, + required_plugins, ) .await? } diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index 7b999967c3ac..fa5734218505 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -4,6 +4,7 @@ use tokio_util::sync::CancellationToken; use crate::session::TurnInput; use crate::session::session::Session; +use crate::session::turn::McpStartupRequirements; use crate::session::turn::run_hooks_and_record_inputs; use crate::session::turn::run_turn; use crate::session::turn_context::TurnContext; @@ -73,11 +74,13 @@ impl SessionTask for RegularTask { }; let mut next_input = input; let mut prewarmed_client_session = prewarmed_client_session; + let mut mcp_startup_requirements = McpStartupRequirements::default(); loop { let last_agent_message = run_turn( Arc::clone(&sess), Arc::clone(&ctx), next_input, + &mut mcp_startup_requirements, prewarmed_client_session.take(), cancellation_token.child_token(), ) diff --git a/codex-rs/skills/src/mentions.rs b/codex-rs/skills/src/mentions.rs index 507a8480dfaa..1f391235d1ed 100644 --- a/codex-rs/skills/src/mentions.rs +++ b/codex-rs/skills/src/mentions.rs @@ -64,8 +64,11 @@ pub fn app_id_from_path(path: &str) -> Option<&str> { .filter(|value| !value.is_empty()) } +/// Desktop app/browser mentions append `?app=...` or `?browserFamily=...`. +/// Ignore these targeting parameters when matching the plugin ID. pub fn plugin_config_name_from_path(path: &str) -> Option<&str> { path.strip_prefix(PLUGIN_PATH_PREFIX) + .and_then(|value| value.split('?').next()) .filter(|value| !value.is_empty()) } diff --git a/codex-rs/skills/src/mentions_tests.rs b/codex-rs/skills/src/mentions_tests.rs index b4f90e1d234f..bdf88748212d 100644 --- a/codex-rs/skills/src/mentions_tests.rs +++ b/codex-rs/skills/src/mentions_tests.rs @@ -78,3 +78,28 @@ fn handles_many_sigils_without_looping() { let prefix = "$".repeat(256); assert_mentions(&format!("{prefix} not-a-mention"), &[], &[]); } + +#[test] +fn plugin_config_names_ignore_mention_query_parameters() { + let paths = [ + "plugin://sample@test", + "plugin://sample@test?app=com.example.editor", + "plugin://sample@test?browserFamily=chrome", + ]; + + assert_eq!( + paths.map(plugin_config_name_from_path), + [Some("sample@test"); 3], + ); +} + +#[test] +fn plugin_config_names_require_a_plugin_identity() { + let paths = [ + "plugin://", + "plugin://?app=com.example.editor", + "app://sample@test?app=com.example.editor", + ]; + + assert_eq!(paths.map(plugin_config_name_from_path), [None; 3]); +}