diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index a3e2549171..87c9e526d3 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -37,6 +37,8 @@ pub use sandbox_metadata::MCP_SANDBOX_STATE_CAPABILITY; pub use sandbox_metadata::MCP_SANDBOX_STATE_META_CAPABILITY; pub use sandbox_metadata::SandboxState; pub use tool_catalog_cache::McpToolCatalogCache; +pub use tool_request_meta::McpToolRequestMetaHook; +pub use tool_request_meta::McpToolRequestMetaProvider; pub use tools::ToolInfo; pub use trusted_access::TrustedAccessContext; @@ -139,5 +141,6 @@ mod sandbox_executable; mod sandbox_metadata; pub(crate) mod server; mod tool_catalog_cache; +mod tool_request_meta; pub(crate) mod tools; mod trusted_access; diff --git a/codex-rs/codex-mcp/src/tool_request_meta.rs b/codex-rs/codex-mcp/src/tool_request_meta.rs new file mode 100644 index 0000000000..7c327e1a6c --- /dev/null +++ b/codex-rs/codex-mcp/src/tool_request_meta.rs @@ -0,0 +1,70 @@ +//! Host-provided metadata for a specific tool on a specific MCP registration. + +use std::sync::Arc; + +use crate::McpServerSource; +use crate::PreparedMcpCall; +use anyhow::Context; +use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; +use codex_config::McpServerTransportConfig; +use codex_protocol::ThreadId; +use futures::future::BoxFuture; +use serde_json::Value; + +/// Supplies one MCP request metadata value for a live tool invocation. An error +/// fails the tool call before dispatch and is reported to the tool caller. +pub trait McpToolRequestMetaProvider: Send + Sync { + fn value<'a>( + &'a self, + thread_id: ThreadId, + turn_id: &'a str, + call_id: &'a str, + ) -> BoxFuture<'a, anyhow::Result>; +} + +/// Install in thread extension data to add a private `_meta` field to one remote +/// tool on the selected local MCP registration. Bind the source and full transport +/// supplied by the trusted host, including routing headers for a proxied endpoint. +pub struct McpToolRequestMetaHook { + pub server_name: String, + pub source: McpServerSource, + pub transport: McpServerTransportConfig, + pub tool_name: String, + /// A metadata key owned by the provider. Existing metadata cannot be overwritten. + pub key: String, + pub provider: Arc, +} + +impl McpToolRequestMetaHook { + /// Adds the provider’s value only when the immutable call matches the selected registration. + pub async fn apply( + &self, + call: &PreparedMcpCall, + thread_id: ThreadId, + turn_id: &str, + call_id: &str, + meta: &mut Option, + ) -> anyhow::Result<()> { + if call.server_name() == self.server_name + && call.tool_info().tool.name.as_ref() == self.tool_name + && call.server_environment_id() == DEFAULT_MCP_SERVER_ENVIRONMENT_ID + && let Some(server) = call.config().mcp_server_catalog.server(call.server_name()) + && server.config().is_local_environment() + && server.source() == &self.source + && server.config().transport == self.transport + { + let meta = meta + .as_mut() + .and_then(Value::as_object_mut) + .context("missing MCP request metadata object")?; + anyhow::ensure!( + !meta.contains_key(&self.key), + "MCP request metadata already contains {}", + self.key + ); + let value = self.provider.value(thread_id, turn_id, call_id).await?; + meta.insert(self.key.clone(), value); + } + Ok(()) + } +} diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index c67d062548..5abbbce2a3 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -40,6 +40,7 @@ use codex_features::Feature; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::MCP_TOOL_CODEX_APPS_META_KEY; use codex_mcp::McpPermissionPromptAutoApproveContext; +use codex_mcp::McpToolRequestMetaHook; use codex_mcp::PreparedMcpCall; use codex_mcp::SandboxState; use codex_mcp::ToolInfo; @@ -584,10 +585,22 @@ async fn handle_approved_mcp_tool_call( .services .rollout_thread_trace .start_mcp_call_trace(call_id); - Ok(( - rewritten_arguments, - mcp_call_trace.add_request_meta(request_meta), - )) + let mut request_meta = mcp_call_trace.add_request_meta(request_meta); + if let Some(hook) = sess + .services + .thread_extension_data + .get::() + { + hook.apply( + &prepared_call, + sess.thread_id, + &turn_context.sub_id, + call_id, + &mut request_meta, + ) + .await?; + } + Ok((rewritten_arguments, request_meta)) }) .await .map_err(|error| { diff --git a/codex-rs/core/tests/suite/mcp_extension_protocol.rs b/codex-rs/core/tests/suite/mcp_extension_protocol.rs index 81711e9a9d..2a8b9ad32c 100644 --- a/codex-rs/core/tests/suite/mcp_extension_protocol.rs +++ b/codex-rs/core/tests/suite/mcp_extension_protocol.rs @@ -1,7 +1,8 @@ -//! Verifies that an extension can choose MCP protocol mode for its own HTTP server. +//! Verifies protocol selection and request metadata for extension-contributed MCP servers. use std::collections::HashMap; use std::sync::Arc; +use std::sync::Mutex; use std::time::Duration; use codex_config::McpServerConfig; @@ -15,14 +16,20 @@ use codex_extension_api::McpServerContributionContext; use codex_extension_api::McpServerContributor; use codex_features::Feature; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_mcp::McpServerSource; +use codex_mcp::McpToolRequestMetaHook; +use codex_mcp::McpToolRequestMetaProvider; +use codex_protocol::ThreadId; use codex_protocol::approvals::ElicitationRequest; use codex_protocol::mcp::ClientMcpExtensions; use codex_protocol::mcp::OPENAI_ELICITATION_EXTENSION_ID; +use codex_protocol::models::FunctionCallOutputBody; use codex_protocol::protocol::ElicitationAction; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use core_test_support::apps_test_server::AppsTestServer; use core_test_support::apps_test_server::apps_enabled_builder; +use core_test_support::apps_test_server::recorded_apps_tool_calls; use core_test_support::responses; use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; @@ -31,6 +38,7 @@ use core_test_support::wait_for_mcp_server; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; +use test_case::test_case; use tokio::time::timeout; use tokio_util::task::AbortOnDropHandle; use wiremock::Mock; @@ -40,9 +48,9 @@ use wiremock::ResponseTemplate; use wiremock::matchers::method; use wiremock::matchers::path; -struct AppsExtensionEndpoint(McpServerContribution); +struct ExtensionServers(Vec); -impl McpServerContributor for AppsExtensionEndpoint { +impl McpServerContributor for ExtensionServers { fn id(&self) -> &'static str { "apps_protocol_test" } @@ -51,7 +59,7 @@ impl McpServerContributor for AppsExtensionEndpoint { &'a self, _context: McpServerContributionContext<'a, Config>, ) -> ExtensionFuture<'a, Vec> { - Box::pin(async move { vec![self.0.clone()] }) + Box::pin(async move { self.0.clone() }) } } @@ -93,13 +101,13 @@ async fn extension_protocol_mode_does_not_change_other_http_servers() -> anyhow: .chatgpt_base_url ); let mut extensions = ExtensionRegistryBuilder::new(); - extensions.mcp_server_contributor(Arc::new(AppsExtensionEndpoint( + extensions.mcp_server_contributor(Arc::new(ExtensionServers(vec![ McpServerContribution::SetWithProtocolMode { name: "extension_apps".to_string(), config: Box::new(serde_json::from_value(json!({ "url": extension_url }))?), protocol_mode: extension_mode, }, - ))); + ]))); let fixture = test_codex() .with_extensions(Arc::new(extensions.build())) @@ -159,6 +167,159 @@ async fn extension_protocol_mode_does_not_change_other_http_servers() -> anyhow: Ok(()) } +struct InvocationMetadata; + +impl McpToolRequestMetaProvider for InvocationMetadata { + fn value<'a>( + &'a self, + thread_id: ThreadId, + turn_id: &'a str, + call_id: &'a str, + ) -> ExtensionFuture<'a, anyhow::Result> { + Box::pin(async move { + anyhow::ensure!(call_id == "prepared", "cannot prepare call {call_id}"); + Ok(json!([thread_id, turn_id, call_id])) + }) + } +} + +enum HookOutcome { + Applies, + Skipped, + KeyConflict, +} + +#[test_case("apps_protocol_test", "approved", HookOutcome::Applies; "matching registration")] +#[test_case("other_contributor", "approved", HookOutcome::Skipped; "different contributor")] +#[test_case("apps_protocol_test", "other", HookOutcome::Skipped; "different routing header")] +#[test_case("apps_protocol_test", "approved", HookOutcome::KeyConflict; "existing metadata key")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn request_metadata_is_scoped_to_the_selected_tool_and_registration( + contributor: &str, + target: &str, + outcome: HookOutcome, +) -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let model = responses::start_mock_server().await; + let gateway = responses::start_mock_server().await; + let tools = ["prepare", "read"].map(|name| { + json!({ + "name": name, + "inputSchema": {"type": "object"}, + "annotations": {"readOnlyHint": true}, + }) + }); + AppsTestServer::mount_with_tools(&gateway, Arc::new(Mutex::new(Vec::from(tools)))).await?; + let config_for = |target: &str| -> serde_json::Result { + serde_json::from_value(json!({ + "url": format!("{}/api/codex/ps/mcp", gateway.uri()), + "http_headers": {"x-test-target": target}, + })) + }; + // The two registrations differ only by name, as when a gateway hosts several tools. + let config = config_for("approved")?; + let servers = ["selected", "other"].map(|name| McpServerContribution::Set { + name: name.into(), + config: Box::new(config.clone()), + }); + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.mcp_server_contributor(Arc::new(ExtensionServers(Vec::from(servers)))); + let fixture = test_codex() + .with_extensions(Arc::new(extensions.build())) + .build_with_auto_env(&model) + .await?; + fixture + .codex + .thread_extension_data() + .insert(McpToolRequestMetaHook { + server_name: "selected".into(), + source: McpServerSource::Extension { + id: contributor.into(), + host_owned_apps: false, + }, + transport: config_for(target)?.transport, + tool_name: "prepare".into(), + key: if matches!(outcome, HookOutcome::KeyConflict) { + "callId" + } else { + "x-test-context" + } + .into(), + provider: Arc::new(InvocationMetadata), + }); + wait_for_mcp_server(&fixture.codex, "selected").await?; + + let mut events = vec![responses::ev_response_created("calls")]; + events.extend( + [ + ("prepared", "mcp__selected", "prepare"), + ("refused", "mcp__selected", "prepare"), + ("other-tool", "mcp__selected", "read"), + ("other-server", "mcp__other", "prepare"), + ] + .map(|(id, ns, tool)| responses::ev_function_call_with_namespace(id, ns, tool, "{}")), + ); + events.push(responses::ev_completed("calls")); + let mock = responses::mount_sse_sequence( + &model, + vec![ + responses::sse(events), + responses::sse(vec![ + responses::ev_assistant_message("done", "done"), + responses::ev_completed("end"), + ]), + ], + ) + .await; + fixture.submit_turn("Call the tools").await?; + + let calls = recorded_apps_tool_calls(&gateway).await; + let mut dispatched: Vec<_> = calls + .iter() + .map(|call| { + let meta = &call["params"]["_meta"]; + ( + meta["callId"].as_str().expect("dispatch call ID"), + meta.get("x-test-context").cloned(), + ) + }) + .collect(); + dispatched.sort_unstable_by_key(|(id, _)| *id); + let mut expected = vec![("other-server", None), ("other-tool", None)]; + let request = mock.last_request().expect("tool outputs"); + let error = match outcome { + HookOutcome::Applies => { + let body = request.body_json(); + let turn_id = body["client_metadata"]["turn_id"] + .as_str() + .expect("model turn ID"); + expected.push(( + "prepared", + Some(json!([ + fixture.session_configured.thread_id, + turn_id, + "prepared" + ])), + )); + Some("cannot prepare call refused") + } + HookOutcome::Skipped => { + expected.extend([("prepared", None), ("refused", None)]); + None + } + HookOutcome::KeyConflict => Some("MCP request metadata already contains callId"), + }; + if let Some(message) = error { + let item = request.function_call_output("refused"); + let body: FunctionCallOutputBody = serde_json::from_value(item["output"].clone())?; + let output = body.to_text().expect("tool error"); + assert!(output.contains(message), "{output}"); + } + assert_eq!(dispatched, expected); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn hosted_apps_protocol_override_preserves_native_verification() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); @@ -222,14 +383,14 @@ async fn hosted_apps_protocol_override_preserves_native_verification() -> anyhow .mount(&apps_server) .await; let mut extensions = ExtensionRegistryBuilder::new(); - extensions.mcp_server_contributor(Arc::new(AppsExtensionEndpoint( + extensions.mcp_server_contributor(Arc::new(ExtensionServers(vec![ McpServerContribution::HostedApps { config: Box::new(serde_json::from_value(json!({ "url": format!("{}/api/codex/ps/mcp", apps_server.uri()) }))?), protocol_mode: Some(McpProtocolMode::V20260728), }, - ))); + ]))); let fixture = apps_enabled_builder(apps_server.uri()) .with_extensions(Arc::new(extensions.build())) .with_config(|config| {