Avoid blocking turns on optional MCP startup (#35742)

## Why

A pending optional MCP server can delay the first model request even when the
turn does not need that server.

## What changed

- Give optional MCP servers a shared one-second startup grace, then omit servers
  that are still pending from the captured tool catalog.
- Continue waiting when the turn explicitly requires a server through a plugin,
  skill dependency, or `mcp://` mention, and preserve this behavior for input
  received between sampling steps.
- Route direct resource requests to the live connection set when a pending
  server was omitted from the binding, while keeping all-server resource
  discovery non-blocking.

## Testing

- Cover the shared grace period and resource behavior for pending optional
  servers.
- Verify plugin, plugin-skill, and direct MCP mentions wait for startup.
- Verify an Apps-enabled turn proceeds without tools from an unrelated pending
  optional server.

GitOrigin-RevId: b5f895c5a5362fe73f7d33250367662d4a217e4d
This commit is contained in:
jif authored and copyberry committed 2026-07-28 09:32:48 +00:00
1 parent 8e271dc02b
commit d9e1c9cd55
12 files changed
+519 -58

No files matched your search

+17 -3
View File
@@ -96,7 +96,11 @@ impl McpBinding {
server: &str, server: &str,
params: Option<PaginatedRequestParams>, params: Option<PaginatedRequestParams>,
) -> Result<ListResourcesResult> { ) -> Result<ListResourcesResult> {
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( pub async fn list_all_resources(
@@ -111,7 +115,13 @@ impl McpBinding {
server: &str, server: &str,
params: Option<PaginatedRequestParams>, params: Option<PaginatedRequestParams>,
) -> Result<ListResourceTemplatesResult> { ) -> Result<ListResourceTemplatesResult> {
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( pub async fn list_all_resource_templates(
@@ -128,7 +138,11 @@ impl McpBinding {
server: &str, server: &str,
params: ReadResourceRequestParams, params: ReadResourceRequestParams,
) -> Result<ReadResourceResult> { ) -> Result<ReadResourceResult> {
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
}
} }
} }
@@ -23,6 +23,7 @@ pub use tool_catalog::tool_is_model_visible;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::Arc; use std::sync::Arc;
use std::sync::OnceLock;
use std::sync::atomic::Ordering; use std::sync::atomic::Ordering;
use std::time::Duration; use std::time::Duration;
@@ -144,6 +145,7 @@ impl McpServerView {
pub(crate) struct McpConnectionSet { pub(crate) struct McpConnectionSet {
servers: HashMap<String, McpServerView>, servers: HashMap<String, McpServerView>,
required_servers: Vec<String>, required_servers: Vec<String>,
optional_startup_deadline: OnceLock<tokio::time::Instant>,
tool_catalog_revision: Arc<RwLock<u64>>, tool_catalog_revision: Arc<RwLock<u64>>,
codex_apps_tools_override: RwLock<Option<Vec<ToolInfo>>>, codex_apps_tools_override: RwLock<Option<Vec<ToolInfo>>>,
codex_apps_refresh_lock: Mutex<()>, codex_apps_refresh_lock: Mutex<()>,
@@ -483,6 +485,7 @@ impl McpConnectionSet {
let manager = Self { let manager = Self {
servers, servers,
required_servers, required_servers,
optional_startup_deadline: OnceLock::new(),
tool_catalog_revision: Arc::new(RwLock::new(0)), tool_catalog_revision: Arc::new(RwLock::new(0)),
codex_apps_tools_override: RwLock::new(None), codex_apps_tools_override: RwLock::new(None),
codex_apps_refresh_lock: Mutex::new(()), codex_apps_refresh_lock: Mutex::new(()),
@@ -540,6 +543,7 @@ impl McpConnectionSet {
Self { Self {
servers: HashMap::new(), servers: HashMap::new(),
required_servers: Vec::new(), required_servers: Vec::new(),
optional_startup_deadline: OnceLock::new(),
tool_catalog_revision: Arc::new(RwLock::new(0)), tool_catalog_revision: Arc::new(RwLock::new(0)),
codex_apps_tools_override: RwLock::new(None), codex_apps_tools_override: RwLock::new(None),
codex_apps_refresh_lock: Mutex::new(()), codex_apps_refresh_lock: Mutex::new(()),
@@ -4,6 +4,7 @@ use std::time::Duration;
use anyhow::Context; use anyhow::Context;
use anyhow::Result; use anyhow::Result;
use anyhow::anyhow; use anyhow::anyhow;
use rmcp::model::ListResourceTemplatesResult;
use rmcp::model::ListResourcesResult; use rmcp::model::ListResourcesResult;
use rmcp::model::PaginatedRequestParams; use rmcp::model::PaginatedRequestParams;
use rmcp::model::ReadResourceRequestParams; use rmcp::model::ReadResourceRequestParams;
@@ -154,6 +155,19 @@ impl McpConnectionSet {
.with_context(|| format!("resources/list failed for `{server}`")) .with_context(|| format!("resources/list failed for `{server}`"))
} }
pub async fn list_resource_templates(
&self,
server: &str,
params: Option<PaginatedRequestParams>,
) -> Result<ListResourceTemplatesResult> {
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( pub async fn read_resource(
&self, &self,
server: &str, server: &str,
@@ -1,5 +1,6 @@
use std::sync::Arc; use std::sync::Arc;
use std::sync::atomic::Ordering; use std::sync::atomic::Ordering;
use std::time::Duration;
use std::time::Instant; use std::time::Instant;
use anyhow::Context; 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_META_KEY: &str = "ui";
const MCP_UI_VISIBILITY_META_KEY: &str = "visibility"; const MCP_UI_VISIBILITY_META_KEY: &str = "visibility";
const MCP_UI_MODEL_VISIBILITY: &str = "model"; 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. /// Returns whether a tool may be included in model-facing tool declarations.
/// ///
@@ -132,10 +134,14 @@ impl McpConnectionSet {
self: &Arc<Self>, self: &Arc<Self>,
config: Arc<crate::McpConfig>, config: Arc<crate::McpConfig>,
plugins_available: bool, plugins_available: bool,
required_servers: &[String],
) -> McpBinding { ) -> McpBinding {
let revision = self.tool_catalog_revision.read().await; let revision = self.tool_catalog_revision.read().await;
let mut listed_tools = Vec::new(); let mut listed_tools = Vec::new();
let mut clients = std::collections::HashMap::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 { for (server_name, view) in &self.servers {
if !view if !view
.connection .connection
@@ -143,7 +149,15 @@ impl McpConnectionSet {
.startup_complete .startup_complete
.load(Ordering::Acquire) .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) = if let Some(server_tools) =
view.listed_tools(&self.tool_plugin_provenance).await view.listed_tools(&self.tool_plugin_provenance).await
{ {
@@ -156,6 +170,17 @@ impl McpConnectionSet {
} }
continue; 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; let _ = view.connection.client.client().await;
} }
view.connection.client.reconnect_failed_startup().await; view.connection.client.reconnect_failed_startup().await;
@@ -86,6 +86,7 @@ impl McpConnectionSet {
Self { Self {
servers: HashMap::new(), servers: HashMap::new(),
required_servers: Vec::new(), required_servers: Vec::new(),
optional_startup_deadline: OnceLock::new(),
tool_catalog_revision: Arc::new(RwLock::new(0)), tool_catalog_revision: Arc::new(RwLock::new(0)),
codex_apps_tools_override: RwLock::new(None), codex_apps_tools_override: RwLock::new(None),
codex_apps_refresh_lock: Mutex::new(()), codex_apps_refresh_lock: Mutex::new(()),
@@ -194,6 +195,7 @@ async fn capture_binding(manager: &Arc<McpConnectionSet>) -> McpBinding {
.capture_binding_with_metadata( .capture_binding_with_metadata(
Arc::new(crate::mcp::tests::test_mcp_config(std::env::temp_dir())), Arc::new(crate::mcp::tests::test_mcp_config(std::env::temp_dir())),
/*plugins_available*/ false, /*plugins_available*/ false,
/*required_servers*/ &[],
) )
.await .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::<Result<ManagedClient, StartupOutcomeError>>()
.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] #[tokio::test]
async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { async fn list_all_tools_applies_legacy_mcp_prefix_by_default() {
let managed_client = let managed_client =
+10 -2
View File
@@ -214,12 +214,20 @@ impl McpRuntime {
/// Captures the latest published configuration and live client handles. /// Captures the latest published configuration and live client handles.
pub async fn current_binding(&self) -> Option<Arc<McpBinding>> { pub async fn current_binding(&self) -> Option<Arc<McpBinding>> {
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<Arc<McpBinding>> {
let current = self.current.load_full(); let current = self.current.load_full();
let config = Arc::clone(current.config.as_ref()?); let config = Arc::clone(current.config.as_ref()?);
Some(Arc::new( Some(Arc::new(
current current
.connections .connections
.capture_binding_with_metadata(config, current.plugins_available) .capture_binding_with_metadata(config, current.plugins_available, required_servers)
.await, .await,
)) ))
} }
@@ -250,7 +258,7 @@ impl McpRuntime {
Some(Arc::new( Some(Arc::new(
current current
.connections .connections
.capture_binding_with_metadata(config, current.plugins_available) .capture_binding_with_metadata(config, current.plugins_available, &[])
.await, .await,
)) ))
} }
+7 -1
View File
@@ -260,6 +260,7 @@ impl Session {
self: &Arc<Self>, self: &Arc<Self>,
turn_context: &TurnContext, turn_context: &TurnContext,
selected_capability_roots: &[ResolvedSelectedCapabilityRoot], selected_capability_roots: &[ResolvedSelectedCapabilityRoot],
required_servers: &[String],
) -> Arc<codex_mcp::McpBinding> { ) -> Arc<codex_mcp::McpBinding> {
let ready_selected_capability_roots = let ready_selected_capability_roots =
Self::ready_selected_capability_roots(selected_capability_roots); Self::ready_selected_capability_roots(selected_capability_roots);
@@ -272,7 +273,12 @@ impl Session {
self.mark_mcp_runtime_dirty(); self.mark_mcp_runtime_dirty();
} }
self.refresh_mcp_if_dirty().await; 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; return binding;
} }
let config = Arc::new(self.runtime_mcp_config(&turn_context.config).await); let config = Arc::new(self.runtime_mcp_config(&turn_context.config).await);
+19 -1
View File
@@ -3023,6 +3023,20 @@ impl Session {
self: &Arc<Self>, self: &Arc<Self>,
turn_context: Arc<TurnContext>, turn_context: Arc<TurnContext>,
cancellation_token: &CancellationToken, cancellation_token: &CancellationToken,
) -> CodexResult<Arc<StepContext>> {
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<Self>,
turn_context: Arc<TurnContext>,
cancellation_token: &CancellationToken,
required_servers: &[String],
) -> CodexResult<Arc<StepContext>> { ) -> CodexResult<Arc<StepContext>> {
// Keep selections fixed for the turn while allowing their startup work to finish. // Keep selections fixed for the turn while allowing their startup work to finish.
let environments = turn_context.environments.refresh_readiness(); let environments = turn_context.environments.refresh_readiness();
@@ -3046,7 +3060,11 @@ impl Session {
extension_data.insert(selected_capability_roots.clone()); extension_data.insert(selected_capability_roots.clone());
let (mcp, prepared_recommendations) = async { let (mcp, prepared_recommendations) = async {
tokio::join!( 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()), turn::prepare_tool_recommendations(self.as_ref(), turn_context.as_ref()),
) )
} }
+5 -1
View File
@@ -8096,7 +8096,11 @@ async fn refresh_mcp_servers_uses_latest_state_for_existing_turns() {
.await .await
.expect("a fresh cancellation token cannot be cancelled"); .expect("a fresh cancellation token cannot be cancelled");
let rematerialized_old = session 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; .await;
let configured_servers = codex_mcp::configured_mcp_servers(new_step.mcp.config()); let configured_servers = codex_mcp::configured_mcp_servers(new_step.mcp.config());
+159 -27
View File
@@ -183,9 +183,26 @@ pub(crate) async fn run_turn(
return Ok(None); 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. // run_turn owns the step used to seed context and make the first sampling request.
let first_step_context = match sess 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 .await
{ {
Ok(step_context) => step_context, 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( let Some((injection_items, explicitly_enabled_connectors)) = build_skills_and_plugins(
&sess, &sess,
first_step_context.as_ref(), first_step_context.as_ref(),
&input, &user_input,
&mentioned_plugins,
&cancellation_token, &cancellation_token,
) )
.await .await
@@ -276,10 +294,26 @@ pub(crate) async fn run_turn(
// Capture once so context, advertised tools, and tool calls share one request view. // Capture once so context, advertised tools, and tool calls share one request view.
let step_context = match next_step_context.take() { let step_context = match next_step_context.take() {
Some(step_context) => step_context, Some(step_context) => step_context,
None => { None if pending_input.is_empty() => {
sess.capture_step_context(Arc::clone(&turn_context), &cancellation_token) sess.capture_step_context(Arc::clone(&turn_context), &cancellation_token)
.await? .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 { let sampling_request_result: CodexResult<_> = async {
super::time_reminder::maybe_record_current_time_reminder( 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 blocked_input && !accepted_user_input
} }
fn turn_user_input(input: &[TurnInput]) -> Vec<UserInput> {
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<Session>,
turn_context: &TurnContext,
user_input: &[UserInput],
) -> (Vec<String>, Vec<crate::plugins::PluginCapabilitySummary>) {
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::<HashSet<_>>();
let messages = user_input
.iter()
.filter_map(|input| match input {
UserInput::Text { text, .. } => Some(text.clone()),
_ => None,
})
.collect::<Vec<_>>();
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)] #[instrument(level = "trace", skip_all)]
async fn build_skills_and_plugins( async fn build_skills_and_plugins(
sess: &Arc<Session>, sess: &Arc<Session>,
step_context: &StepContext, step_context: &StepContext,
input: &[TurnInput], user_input: &[UserInput],
mentioned_plugins: &[crate::plugins::PluginCapabilitySummary],
cancellation_token: &CancellationToken, cancellation_token: &CancellationToken,
) -> Option<(Vec<ResponseItem>, HashSet<String>)> { ) -> Option<(Vec<ResponseItem>, HashSet<String>)> {
let turn_context = step_context.turn.as_ref(); let turn_context = step_context.turn.as_ref();
@@ -581,30 +731,12 @@ async fn build_skills_and_plugins(
return Some((Vec::new(), HashSet::new())); 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::<Vec<_>>();
let tracking = build_track_events_context( let tracking = build_track_events_context(
turn_context.model_info.slug.clone(), turn_context.model_info.slug.clone(),
sess.thread_id.to_string(), sess.thread_id.to_string(),
turn_context.sub_id.clone(), turn_context.sub_id.clone(),
turn_context.originator.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 connector_snapshot = step_context.mcp.config().connector_snapshot.clone();
let mcp_tools = if turn_context.apps_enabled() || !mentioned_plugins.is_empty() { let mcp_tools = if turn_context.apps_enabled() || !mentioned_plugins.is_empty() {
// Plugin mentions need raw MCP/app inventory even when app tools // 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 skills_outcome = turn_context.turn_skills.snapshot.outcome();
let connector_slug_counts = build_connector_slug_counts(&available_connectors); let connector_slug_counts = build_connector_slug_counts(&available_connectors);
let extension_injection_items = 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?; .await?;
let skill_name_counts_lower = let skill_name_counts_lower =
build_skill_name_counts(&skills_outcome.skills, &skills_outcome.disabled_paths).1; build_skill_name_counts(&skills_outcome.skills, &skills_outcome.disabled_paths).1;
let mentioned_skills = collect_explicit_skill_mentions( let mentioned_skills = collect_explicit_skill_mentions(
&user_input, user_input,
&skills_outcome.skills, &skills_outcome.skills,
&skills_outcome.disabled_paths, &skills_outcome.disabled_paths,
&connector_slug_counts, &connector_slug_counts,
@@ -678,8 +810,8 @@ async fn build_skills_and_plugins(
&skill_name_counts_lower, &skill_name_counts_lower,
); );
let plugin_items = let plugin_items =
build_plugin_injections(&mentioned_plugins, &mcp_tools, &available_connectors); build_plugin_injections(mentioned_plugins, &mcp_tools, &available_connectors);
let mut explicitly_enabled_connectors = collect_explicit_app_ids(&user_input); let mut explicitly_enabled_connectors = collect_explicit_app_ids(user_input);
explicitly_enabled_connectors.extend(skill_connector_ids); explicitly_enabled_connectors.extend(skill_connector_ids);
let connector_names_by_id = available_connectors let connector_names_by_id = available_connectors
.iter() .iter()
@@ -698,7 +830,7 @@ async fn build_skills_and_plugins(
sess.services sess.services
.analytics_events_client .analytics_events_client
.track_app_mentioned(tracking.clone(), mentioned_app_invocations); .track_app_mentioned(tracking.clone(), mentioned_app_invocations);
for summary in &mentioned_plugins { for summary in mentioned_plugins {
if let Some(plugin) = sess if let Some(plugin) = sess
.services .services
.plugins_manager .plugins_manager
+89 -22
View File
@@ -45,6 +45,7 @@ use core_test_support::wait_for_event_match;
use core_test_support::wait_for_mcp_server; use core_test_support::wait_for_mcp_server;
use pretty_assertions::assert_eq; use pretty_assertions::assert_eq;
use tempfile::TempDir; use tempfile::TempDir;
use test_case::test_case;
use wiremock::MockServer; use wiremock::MockServer;
const SAMPLE_PLUGIN_CONFIG_NAME: &str = "sample@test"; 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"); .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) { fn write_plugin_app_plugin(home: &TempDir) {
write_plugin_app_plugin_with_name(home, "sample"); 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(()) 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)] #[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(())); skip_if_no_network!(Ok(()));
let server = start_mock_server().await; let server = start_mock_server().await;
let mock = mount_plugin_tool_search_turn(&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(()); 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_mcp_plugin(codex_home.as_ref(), &rmcp_test_server_bin);
write_plugin_app_plugin(codex_home.as_ref()); 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() let mut builder = test_codex()
.with_home(codex_home) .with_home(codex_home)
@@ -868,37 +904,68 @@ async fn explicit_plugin_mentions_use_mcp_for_api_key_dual_surface_plugins() ->
.await .await
.expect("create new conversation"); .expect("create new conversation");
let codex = Arc::clone(&test_codex.codex); 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 codex
.submit(Op::UserInput { .submit(Op::UserInput {
items: vec![codex_protocol::user_input::UserInput::Mention { items: vec![input],
name: "sample".into(),
path: format!("plugin://{SAMPLE_PLUGIN_CONFIG_NAME}"),
}],
final_output_json_schema: None, final_output_json_schema: None,
responsesapi_client_metadata: None, responsesapi_client_metadata: None,
additional_context: Default::default(), additional_context: Default::default(),
thread_settings: Default::default(), thread_settings: Default::default(),
}) })
.await?; .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; wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
let requests = mock.requests(); let requests = mock.requests();
let request = &requests[0]; let model_request = &requests[0];
let developer_messages = request.message_input_texts("developer"); let developer_messages = model_request.message_input_texts("developer");
assert!( if matches!(request, ExplicitMcpRequest::Plugin) {
developer_messages assert!(
.iter() developer_messages
.any(|text| text.contains("Skills from this plugin")), .iter()
"expected plugin skills guidance: {developer_messages:?}" .any(|text| text.contains("Skills from this plugin")),
); "expected plugin skills guidance: {developer_messages:?}"
assert!( );
developer_messages assert!(
.iter() developer_messages
.any(|text| text.contains("MCP servers from this plugin")), .iter()
"expected visible plugin MCP guidance: {developer_messages:?}" .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!( assert!(
!developer_messages !developer_messages
.iter() .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:?}" "expected plugin app guidance to be suppressed for API-key auth: {developer_messages:?}"
); );
assert!( assert!(
request model_request
.tool_by_name(SAMPLE_PLUGIN_APP_NAMESPACE, SEARCH_CALENDAR_CREATE_TOOL) .tool_by_name(SAMPLE_PLUGIN_APP_NAMESPACE, SEARCH_CALENDAR_CREATE_TOOL)
.is_none(), .is_none(),
"plugin app tool should not leak into the request for API-key auth" "plugin app tool should not leak into the request for API-key auth"
+86
View File
@@ -26,7 +26,9 @@ use codex_exec_server::CreateDirectoryOptions;
use codex_exec_server::Environment; use codex_exec_server::Environment;
use codex_exec_server::HttpRedirectPolicy; use codex_exec_server::HttpRedirectPolicy;
use codex_exec_server::HttpRequestParams; use codex_exec_server::HttpRequestParams;
use codex_features::Feature;
use codex_login::CodexAuth; use codex_login::CodexAuth;
use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
use codex_mcp::MCP_SANDBOX_STATE_META_CAPABILITY; use codex_mcp::MCP_SANDBOX_STATE_META_CAPABILITY;
use codex_mcp::SandboxState; use codex_mcp::SandboxState;
use codex_models_manager::manager::RefreshStrategy; 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::AskForApproval;
use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::McpInvocation; use codex_protocol::protocol::McpInvocation;
use codex_protocol::protocol::McpStartupStatus;
use codex_protocol::protocol::McpToolCallBeginEvent; use codex_protocol::protocol::McpToolCallBeginEvent;
use codex_protocol::protocol::Op; use codex_protocol::protocol::Op;
use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::RolloutItem;
@@ -808,6 +811,89 @@ server_names = ["history", "notes"]
Ok(()) 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("<apps_instructions>"));
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)] #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn shutdown_cancels_startup_prewarm_waiting_for_mcp_startup() -> anyhow::Result<()> { async fn shutdown_cancels_startup_prewarm_waiting_for_mcp_startup() -> anyhow::Result<()> {
skip_if_no_network!(Ok(())); skip_if_no_network!(Ok(()));