openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-27 18:48:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Honor explicit plugin mentions during MCP startup (#42406)

Browse Source 
## 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
This commit is contained in:
Matthew Zeng
2026-09-03 01:33:32 +00:00
committed by copyberry
parent b7f710273e
commit 460b63e5f4
14 changed files with 617 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs
+370 -54
View File
@@ -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(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn selected_capabilities_become_available_between_samples_in_one_turn() -> Result<()> {
const USER_INPUT_CALL_ID: &str = "pause-for-environment";
#[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(
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<std::path::PathBuf> {
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(())
}
codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs
+6 -1
View File
@@ -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<crate::McpConfig>,
plugins_available: bool,
required_servers: &[String],
required_plugins: &HashSet<String>,
) -> 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);
codex-rs/codex-mcp/src/connection_manager_tests.rs
+73 -12
View File
@@ -231,6 +231,7 @@ async fn capture_binding(manager: &Arc<McpConnectionSet>) -> 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
codex-rs/codex-mcp/src/runtime.rs
+35 -10
View File
@@ -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<Arc<McpBinding>> {
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<String>,
) -> Option<Arc<McpBinding>> {
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<PublishedMcpRuntime>,
required_servers: &[String],
required_plugins: &HashSet<String>,
) -> Option<Arc<McpBinding>> {
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));
}
codex-rs/core/src/plugins/mentions.rs
+14 -11
View File
@@ -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 `<plugin>@<marketplace>` IDs; selected roots may supply opaque IDs.
pub(crate) fn collect_explicit_plugin_ids(input: &[UserInput]) -> HashSet<String> {
let messages = input
.iter()
.filter_map(|item| match item {
@@ -76,7 +89,7 @@ pub(crate) fn collect_explicit_plugin_mentions(
})
.collect::<Vec<String>>();
let mentioned_config_names: HashSet<String> = 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()
}
codex-rs/core/src/plugins/mentions_tests.rs
+29 -8
View File
@@ -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::<PluginCapabilitySummary>::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()]),
);
}
codex-rs/core/src/plugins/mod.rs
+1
View File
@@ -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;
codex-rs/core/src/session/mcp.rs
+2 -1
View File
@@ -329,6 +329,7 @@ impl Session {
turn_context: &TurnContext,
selected_capability_roots: &[ResolvedSelectedCapabilityRoot],
required_servers: &[String],
required_plugins: &HashSet<String>,
) -> Arc<codex_mcp::McpBinding> {
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;
codex-rs/core/src/session/mod.rs
+11 -1
View File
@@ -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<TurnContext>,
cancellation_token: &CancellationToken,
required_servers: &[String],
required_plugins: &HashSet<String>,
) -> CodexResult<Arc<StepContext>> {
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<TurnContext>,
cancellation_token: &CancellationToken,
required_servers: &[String],
required_plugins: &HashSet<String>,
) -> CodexResult<Arc<StepContext>> {
// 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()),
)
codex-rs/core/src/session/tests.rs
+1
View File
@@ -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;
codex-rs/core/src/session/turn.rs
+44 -8
View File
@@ -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<String>,
required_plugins: HashSet<String>,
}
/// 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<Session>,
turn_context: Arc<TurnContext>,
input: Vec<TurnInput>,
mcp_startup_requirements: &mut McpStartupRequirements,
prewarmed_client_session: Option<ModelClientSession>,
cancellation_token: CancellationToken,
) -> CodexResult<Option<String>> {
@@ -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?
}
codex-rs/core/src/tasks/regular.rs
+3
View File
@@ -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(),
)
codex-rs/skills/src/mentions.rs
+3
View File
@@ -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())
}
codex-rs/skills/src/mentions_tests.rs
+25
View File
@@ -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]);
}