agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Simplify MCP tool handler plumbing (#21595)

Browse Source 
## Why
The MCP tool path had accumulated a few core-owned special cases: a
dedicated payload variant, resolver plumbing, a legacy `AfterToolUse`
translation path, and a side channel for parallel-call metadata. That
made `ToolRegistry` and the spec builder know more about MCP than they
needed to.

This change moves MCP-specific execution details back onto `ToolInfo`
and `McpHandler` so `codex-core` can treat MCP calls like normal
function calls while still preserving MCP-specific dispatch and
telemetry behavior where it belongs.

## What changed
- removed `resolve_mcp_tool_info`, `ToolPayload::Mcp`, `ToolKind`, and
the remaining registry-side MCP resolver path
- stored MCP routing metadata directly on `McpHandler` and `ToolInfo`,
including `supports_parallel_tool_calls`
- deleted the legacy `AfterToolUse` consumer in `core`, which removes
the need for handler-specific `after_tool_use_payload` implementations
- switched tool-result telemetry to handler-provided tags and kept
MCP-specific dispatch payload construction inside the handler
- simplified tool spec planning/building by passing `ToolInfo` directly
and dropping the direct/deferred MCP wrapper structs and the
parallel-server side table

## Testing
- `cargo check -p codex-core -p codex-mcp -p codex-otel`
- `cargo test -p codex-core
mcp_parallel_support_uses_exact_payload_server`
- `cargo test -p codex-core
direct_mcp_tools_register_namespaced_handlers`
- `cargo test -p codex-core
search_tool_description_lists_each_mcp_source_once`
- `cargo test -p codex-mcp
list_all_tools_uses_startup_snapshot_while_client_is_pending`
- `just fix -p codex-core -p codex-mcp -p codex-otel`
This commit is contained in:
pakrym-oai
2026-05-12 00:11:31 +00:00
committed by GitHub
parent e16b4e46d4
commit ed5944ba1d
70 changed files with 412 additions and 638 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-mcp/src/connection_manager.rs
+21 -22
View File
@@ -7,7 +7,6 @@
//! `codex-core`.
use std::collections::HashMap;
use std::collections::HashSet;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
@@ -44,7 +43,6 @@ use codex_config::Constrained;
use codex_config::McpServerTransportConfig;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_login::CodexAuth;
use codex_protocol::ToolName;
use codex_protocol::mcp::CallToolResult;
use codex_protocol::models::PermissionProfile;
use codex_protocol::protocol::AskForApproval;
@@ -131,17 +129,6 @@ impl McpConnectionManager {
.is_none_or(|metadata| metadata.pollutes_memory)
}
pub fn parallel_tool_call_server_names(&self) -> HashSet<String> {
self.server_metadata
.iter()
.filter_map(|(name, metadata)| {
metadata
.supports_parallel_tool_calls
.then_some(name.clone())
})
.collect()
}
pub fn is_host_owned_codex_apps_server(&self, server_name: &str) -> bool {
self.host_owned_codex_apps_enabled && server_name == CODEX_APPS_MCP_SERVER_NAME
}
@@ -373,7 +360,11 @@ impl McpConnectionManager {
let Some(server_tools) = managed_client.listed_tools().await else {
continue;
};
tools.extend(server_tools);
tools.extend(
server_tools
.into_iter()
.map(|tool| self.with_server_metadata(tool)),
);
}
normalize_tools_for_model(tools)
}
@@ -424,11 +415,26 @@ impl McpConnectionManager {
.into_iter()
.map(|mut tool| {
tool.tool = tool_with_model_visible_input_schema(&tool.tool);
tool
self.with_server_metadata(tool)
});
Ok(normalize_tools_for_model(tools))
}
fn with_server_metadata(&self, mut tool: ToolInfo) -> ToolInfo {
let Some(metadata) = self.server_metadata.get(&tool.server_name) else {
tool.supports_parallel_tool_calls = false;
tool.server_origin = None;
return tool;
};
tool.supports_parallel_tool_calls = metadata.supports_parallel_tool_calls;
tool.server_origin = metadata
.origin
.as_ref()
.map(|origin| origin.as_str().to_string());
tool
}
/// Returns a single map that contains all resources. Each key is the
/// server name and the value is a vector of resources.
pub async fn list_all_resources(&self) -> HashMap<String, Vec<Resource>> {
@@ -662,13 +668,6 @@ impl McpConnectionManager {
.with_context(|| format!("resources/read failed for `{server}` ({uri})"))
}
pub async fn resolve_tool_info(&self, tool_name: &ToolName) -> Option<ToolInfo> {
let all_tools = self.list_all_tools().await;
all_tools
.into_iter()
.find(|tool| tool.canonical_tool_name() == *tool_name)
}
async fn client_by_name(&self, name: &str) -> Result<ManagedClient> {
self.clients
.get(name)
codex-rs/codex-mcp/src/connection_manager_tests.rs
+43 -38
View File
@@ -10,6 +10,7 @@ use crate::elicitation::elicitation_is_rejected_by_policy;
use crate::rmcp_client::AsyncManagedClient;
use crate::rmcp_client::ManagedClient;
use crate::rmcp_client::StartupOutcomeError;
use crate::server::McpServerOrigin;
use crate::tools::ToolFilter;
use crate::tools::ToolInfo;
use crate::tools::filter_tools;
@@ -38,6 +39,8 @@ fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo {
let tool_namespace = format!("mcp__{server_name}__");
ToolInfo {
server_name: server_name.to_string(),
supports_parallel_tool_calls: false,
server_origin: None,
callable_name: tool_name.to_string(),
callable_namespace: tool_namespace,
namespace_description: None,
@@ -712,44 +715,6 @@ async fn list_all_tools_uses_startup_snapshot_while_client_is_pending() {
assert_eq!(tool.callable_name, "calendar_create_event");
}
#[tokio::test]
async fn resolve_tool_info_accepts_canonical_namespaced_tool_names() {
let startup_tools = vec![create_test_tool("rmcp", "echo")];
let pending_client = futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
.boxed()
.shared();
let approval_policy = Constrained::allow_any(AskForApproval::OnFailure);
let permission_profile = Constrained::allow_any(PermissionProfile::default());
let mut manager =
McpConnectionManager::new_uninitialized(&approval_policy, &permission_profile);
manager.clients.insert(
"rmcp".to_string(),
AsyncManagedClient {
client: pending_client,
startup_snapshot: Some(startup_tools),
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
},
);
let tool = manager
.resolve_tool_info(&ToolName::namespaced("mcp__rmcp__", "echo"))
.await
.expect("split MCP tool namespace and name should resolve");
let expected = ("rmcp", "mcp__rmcp__", "echo", "echo");
assert_eq!(
(
tool.server_name.as_str(),
tool.callable_namespace.as_str(),
tool.callable_name.as_str(),
tool.tool.name.as_ref(),
),
expected
);
}
#[tokio::test]
async fn list_all_tools_blocks_while_client_is_pending_without_startup_snapshot() {
let pending_client = futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
@@ -842,6 +807,46 @@ async fn list_all_tools_uses_startup_snapshot_when_client_startup_fails() {
assert_eq!(tool.callable_name, "calendar_create_event");
}
#[tokio::test]
async fn list_all_tools_adds_server_metadata_to_cached_tools() {
let server_name = "docs";
let startup_tools = vec![create_test_tool(server_name, "search")];
let pending_client = futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
.boxed()
.shared();
let approval_policy = Constrained::allow_any(AskForApproval::OnFailure);
let permission_profile = Constrained::allow_any(PermissionProfile::default());
let mut manager =
McpConnectionManager::new_uninitialized(&approval_policy, &permission_profile);
manager.server_metadata.insert(
server_name.to_string(),
McpServerMetadata {
pollutes_memory: true,
origin: Some(McpServerOrigin::StreamableHttp(
"https://docs.example".to_string(),
)),
supports_parallel_tool_calls: true,
},
);
manager.clients.insert(
server_name.to_string(),
AsyncManagedClient {
client: pending_client,
startup_snapshot: Some(startup_tools),
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
},
);
let tools = manager.list_all_tools().await;
assert_eq!(tools.len(), 1);
let tool = &tools[0];
assert_eq!(tool.server_name, server_name);
assert!(tool.supports_parallel_tool_calls);
assert_eq!(tool.server_origin.as_deref(), Some("https://docs.example"));
}
#[test]
fn elicitation_capability_uses_2025_06_18_shape_for_form_only_support() {
let capability = Some(ElicitationCapability::default());
codex-rs/codex-mcp/src/rmcp_client.rs
+2
View File
@@ -375,6 +375,8 @@ pub(crate) async fn list_tools_for_client_uncached(
};
ToolInfo {
server_name: server_name.to_owned(),
supports_parallel_tool_calls: false,
server_origin: None,
callable_name,
callable_namespace,
namespace_description,
codex-rs/codex-mcp/src/tools.rs
+6
View File
@@ -29,6 +29,12 @@ pub(crate) const MCP_TOOLS_CACHE_WRITE_DURATION_METRIC: &str =
pub struct ToolInfo {
/// Raw MCP server name used for routing the tool call.
pub server_name: String,
/// Whether calls routed to this server may run in parallel.
#[serde(default)]
pub supports_parallel_tool_calls: bool,
/// MCP server origin used for telemetry and diagnostics, when known.
#[serde(default)]
pub server_origin: Option<String>,
/// Model-visible tool name used in Responses API tool declarations.
#[serde(rename = "tool_name", alias = "callable_name")]
pub callable_name: String,