Add scoped host metadata hooks for MCP tool requests (#58395)

## What changed

Expose `McpToolRequestMetaHook` and `McpToolRequestMetaProvider` so hosts can supply a per-invocation `_meta` value using the thread, turn, and call IDs. Apply the hook before dispatch only when the tool, server name, source, and full transport configuration match the selected local MCP registration.

Reject metadata key collisions and propagate provider errors to the tool caller without dispatching the call.

## Testing

Add integration coverage for matching registrations, different contributors and routing headers, unrelated tools and servers, provider failures, and existing metadata keys.

GitOrigin-RevId: 2bf51c375a04617d61b24ca1cbbcaa810fb06720

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;

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<Value>>;
+}
+
+/// 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<dyn McpToolRequestMetaProvider>,
+}
+
+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<Value>,
+    ) -> 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(())
+    }
+}

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::<McpToolRequestMetaHook>()
+                    {
+                        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| {

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<McpServerContribution>);
 
-impl McpServerContributor<Config> for AppsExtensionEndpoint {
+impl McpServerContributor<Config> for ExtensionServers {
     fn id(&self) -> &'static str {
         "apps_protocol_test"
     }
@@ -51,7 +59,7 @@ impl McpServerContributor<Config> for AppsExtensionEndpoint {
         &'a self,
         _context: McpServerContributionContext<'a, Config>,
     ) -> ExtensionFuture<'a, Vec<McpServerContribution>> {
-        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<Value>> {
+        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<McpServerConfig> {
+        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| {