Let unrelated tools run while MCP servers start (#35937)

## Why

Cached MCP definitions can make a tool callable before its server has finished
starting. Waiting for that server while holding the tool execution gate blocks
unrelated calls in the same response.

## What changed

- Expose the owning MCP server through tool runtime metadata, including tools
  with exposure overrides.
- Refresh MCP state and wait for that specific server before entering the
  parallel execution gate.

## Testing

Extend the cached MCP tool regression test to verify that an unrelated plan
update completes while MCP startup remains pending.

GitOrigin-RevId: 884d154bc403dfc1d347d265da3d6a887948e7b6
This commit is contained in:
jif authored and copyberry committed 2026-07-29 11:14:16 +00:00
1 parent fe01054a28
commit fbf666fa98
7 files changed
+91 -6

No files matched your search

+4
View File
@@ -163,6 +163,10 @@ impl McpHandler {
}
impl CoreToolRuntime for McpHandler {
fn mcp_server_name(&self) -> Option<&str> {
Some(&self.tool_info.server_name)
}
fn telemetry_tags<'a>(
&'a self,
_invocation: &'a ToolInvocation,
+10
View File
@@ -98,6 +98,7 @@ impl ToolCallRuntime {
cancellation_token: CancellationToken,
) -> impl std::future::Future<Output = Result<AnyToolResult, FunctionCallError>> {
let supports_parallel = self.router.tool_supports_parallel(&call);
let mcp_server = self.router.mcp_server_name(&call).map(str::to_owned);
let router = Arc::clone(&self.router);
let session = Arc::clone(&self.session);
let step_context = Arc::clone(&self.step_context);
@@ -130,6 +131,15 @@ impl ToolCallRuntime {
let mut dispatch_handle: AbortOnDropHandle<Result<AnyToolResult, FunctionCallError>> =
AbortOnDropHandle::new(tokio::spawn(async move {
if let Some(server) = mcp_server {
session.refresh_mcp_if_dirty().await;
session
.services
.mcp_runtime
.wait_for_server_startup(&server)
.await;
}
let _guard = if supports_parallel {
Either::Left(lock.read().await)
} else {
+13
View File
@@ -51,6 +51,11 @@ pub use codex_tools::ToolExposure;
/// Implementers provide the shared `ToolExecutor` behavior plus optional
/// core-owned metadata for hooks, telemetry, tool search, and argument diffs.
pub(crate) trait CoreToolRuntime: ToolExecutor<ToolInvocation> {
/// Returns the MCP server that owns this runtime, when this is an MCP tool.
fn mcp_server_name(&self) -> Option<&str> {
None
}
fn matches_kind(&self, payload: &ToolPayload) -> bool {
matches!(
payload,
@@ -283,6 +288,10 @@ impl ToolExecutor<ToolInvocation> for ExposureOverride {
}
impl CoreToolRuntime for ExposureOverride {
fn mcp_server_name(&self) -> Option<&str> {
self.handler.mcp_server_name()
}
fn matches_kind(&self, payload: &ToolPayload) -> bool {
self.handler.matches_kind(payload)
}
@@ -392,6 +401,10 @@ impl ToolRegistry {
self.tools.get(name).map(Arc::clone)
}
pub(crate) fn mcp_server_name(&self, name: &ToolName) -> Option<&str> {
self.tools.get(name)?.mcp_server_name()
}
#[cfg(test)]
pub(crate) fn tool_names_for_test(&self) -> Vec<ToolName> {
let mut names = self.tools.keys().cloned().collect::<Vec<_>>();
+4
View File
@@ -149,6 +149,10 @@ impl ToolRouter {
.unwrap_or(false)
}
pub(crate) fn mcp_server_name(&self, call: &ToolCall) -> Option<&str> {
self.registry.mcp_server_name(&call.tool_name)
}
pub fn tool_waits_for_runtime_cancellation(&self, call: &ToolCall) -> bool {
self.registry
.waits_for_runtime_cancellation(&call.tool_name)
+27 -6
View File
@@ -7,6 +7,8 @@ use crate::session::tests::make_session_and_context;
use crate::tools::context::ToolPayload;
use crate::tools::handlers::McpHandler;
use crate::tools::registry::CoreToolRuntime;
use crate::tools::registry::ToolExposure;
use crate::tools::registry::override_tool_exposure;
use crate::turn_diff_tracker::TurnDiffTracker;
use codex_extension_api::ExtensionData;
use codex_extension_api::ExtensionRegistry;
@@ -156,6 +158,17 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow
})
.expect("test session should expose a parallel shell-like tool");
assert_eq!(
router.mcp_server_name(&ToolCall {
tool_name: ToolName::plain(parallel_tool_name),
call_id: "call-local-tool".to_string(),
payload: ToolPayload::Function {
arguments: "{}".to_string(),
},
}),
None
);
assert!(!router.tool_supports_parallel(&ToolCall {
tool_name: ToolName::namespaced("mcp__server__", parallel_tool_name),
call_id: "call-namespaced-tool".to_string(),
@@ -248,12 +261,15 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> {
"mcp__echo__",
"query_with_delay",
)),
mcp_runtime(mcp_tool_info(
"hello_echo",
/*supports_parallel_tool_calls*/ false,
"mcp__hello_echo__",
"query_with_delay",
)),
override_tool_exposure(
mcp_runtime(mcp_tool_info(
"hello_echo",
/*supports_parallel_tool_calls*/ false,
"mcp__hello_echo__",
"query_with_delay",
)),
ToolExposure::DirectModelOnly,
),
],
extension_tool_executors: Vec::new(),
wait_for_environment_tool_config: None,
@@ -271,6 +287,7 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> {
encrypted_function_args: None,
};
assert!(router.tool_supports_parallel(&call));
assert_eq!(router.mcp_server_name(&call), Some("echo"));
let different_server_call = ToolCall {
tool_name: ToolName::namespaced("mcp__hello_echo__", "query_with_delay"),
@@ -281,6 +298,10 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> {
encrypted_function_args: None,
};
assert!(!router.tool_supports_parallel(&different_server_call));
assert_eq!(
router.mcp_server_name(&different_server_call),
Some("hello_echo")
);
Ok(())
}