openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-31 19:28:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
## 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
This commit is contained in:
rhan-oai
2026-10-30 17:32:09 +00:00
committed by copyberry
parent cbe16d525c
commit 5be57bedd0
4 changed files with 259 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/mcp_tool_call.rs
+17 -4
View File
@@ -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
+169 -8
View File
@@ -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| {