Expose MCP provenance to tool lifecycle extensions (#40976)

## What changed

- Add optional `McpToolContext` metadata to `ToolStartInput`, exposing the
  model-visible MCP tool details and its source classification without exposing
  the executable client.
- Classify MCP calls as connectors, configured servers, plugin servers,
  executor-selected plugins, or other registrations based on the prepared call.
- Prepare each MCP call before notifying tool lifecycle contributors and reuse
  that same call for execution so the callback describes the call that runs.

## Testing

- Cover host-owned connector calls and extension-owned Apps server calls,
  including their distinct provenance and executed tool names.

GitOrigin-RevId: efd23f511b1045ffdd96349d621ad62365569b2b
This commit is contained in:
felixxia-oai
2026-08-26 21:01:32 +00:00
committed by copyberry
parent daa3eaf10f
commit 21ff2e802c
10 changed files with 264 additions and 3 deletions

View File

@@ -236,6 +236,18 @@ impl PreparedMcpCall {
&self.server_name
}
/// Returns whether this call is bound to the host-owned Codex Apps server.
pub fn is_host_owned_apps(&self) -> bool {
self.config
.mcp_server_catalog
.server(&self.server_name)
.is_some_and(|registration| {
registration
.source()
.is_host_owned_apps(&self.server_name, registration.config())
})
}
pub fn server_origin(&self) -> Option<&str> {
self.server_metadata
.origin

View File

@@ -119,6 +119,7 @@ pub(crate) async fn handle_mcp_tool_call(
call_id: String,
originating_item_id: Option<ResponseItemId>,
tool_info: &ToolInfo,
prepared_call: Option<PreparedMcpCall>,
hook_tool_name: HookToolName,
invocation_tool_name: ToolName,
arguments: String,
@@ -149,7 +150,7 @@ pub(crate) async fn handle_mcp_tool_call(
arguments: arguments_value.clone(),
};
let Some(prepared_call) = sess.prepare_mcp_call(&server, &tool_name).await else {
let Some(prepared_call) = prepared_call else {
let item_metadata =
McpToolCallItemMetadata::from_tool_metadata(&server, /*metadata*/ None);
let result = notify_mcp_tool_call_skip(

View File

@@ -17,11 +17,13 @@ use crate::tools::context::ToolPayload;
use crate::tools::context::boxed_tool_output;
use crate::tools::flat_tool_name;
use crate::tools::hook_names::HookToolName;
use crate::tools::lifecycle::notify_tool_start;
use crate::tools::registry::CoreToolRuntime;
use crate::tools::registry::PostToolUsePayload;
use crate::tools::registry::PreToolUsePayload;
use crate::tools::registry::ToolExecutor;
use crate::tools::registry::ToolTelemetryTags;
use codex_extension_api::McpToolContext;
use codex_mcp::ToolInfo;
use codex_protocol::mcp::is_node_repl_backed_server;
use codex_protocol::user_input::UserInput;
@@ -171,6 +173,26 @@ impl McpHandler {
&self,
invocation: ToolInvocation,
) -> Result<Box<dyn crate::tools::context::ToolOutput>, FunctionCallError> {
let prepared_mcp_call = invocation
.session
.prepare_mcp_call(
&self.tool_info.server_name,
self.tool_info.tool.name.as_ref(),
)
.await;
let mcp_tool = prepared_mcp_call.as_ref().map(|call| {
McpToolContext::from_prepared_call(
call,
invocation
.turn
.config
.mcp_servers
.get()
.get(call.server_name()),
)
});
notify_tool_start(&invocation, mcp_tool.as_ref()).await;
let originating_item_id = invocation.originating_item_id().await;
let ToolInvocation {
session,
@@ -200,6 +222,7 @@ impl McpHandler {
call_id.clone(),
originating_item_id,
&self.tool_info,
prepared_mcp_call,
self.hook_tool_name(),
tool_name,
payload,

View File

@@ -1,5 +1,6 @@
use std::sync::Arc;
use codex_extension_api::McpToolContext;
use codex_extension_api::ToolCallOutcome;
use codex_extension_api::ToolCallSource as ExtensionToolCallSource;
use codex_extension_api::ToolFinishInput;
@@ -11,7 +12,10 @@ use crate::session::turn_context::TurnContext;
use crate::tools::context::ToolCallSource;
use crate::tools::context::ToolInvocation;
pub(crate) async fn notify_tool_start(invocation: &ToolInvocation) {
pub(crate) async fn notify_tool_start(
invocation: &ToolInvocation,
mcp_tool: Option<&McpToolContext>,
) {
let contributors = invocation
.session
.services
@@ -32,6 +36,7 @@ pub(crate) async fn notify_tool_start(invocation: &ToolInvocation) {
turn_id: invocation.turn.sub_id.as_str(),
call_id: invocation.call_id.as_str(),
tool_name: &invocation.tool_name,
mcp_tool,
payload: &invocation.payload,
conversation_history: Arc::clone(&conversation_history),
source: extension_tool_call_source(invocation.source.clone()),

View File

@@ -619,7 +619,9 @@ impl ToolRegistry {
}
}
notify_tool_start(&invocation).await;
if tool.mcp_server_name().is_none() {
notify_tool_start(&invocation, /*mcp_tool*/ None).await;
}
let mut control_tool_analytics = tool
.is_builtin_control_tool()
.then(|| ControlToolCallGuard::new(&invocation));

View File

@@ -6,24 +6,39 @@ use std::sync::Mutex;
use anyhow::Context;
use anyhow::Result;
use codex_core::config::Config;
use codex_extension_api::ExtensionFuture;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::McpServerContribution;
use codex_extension_api::McpServerContributionContext;
use codex_extension_api::McpServerContributor;
use codex_extension_api::McpToolSource;
use codex_extension_api::ResponseItem;
use codex_extension_api::ToolLifecycleContributor;
use codex_extension_api::ToolLifecycleFuture;
use codex_extension_api::ToolStartInput;
use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
use codex_protocol::models::ContentItem;
use core_test_support::apps_test_server::AppsTestServer;
use core_test_support::apps_test_server::SEARCH_CALENDAR_LIST_TOOL;
use core_test_support::apps_test_server::SEARCH_CALENDAR_NAMESPACE;
use core_test_support::apps_test_server::apps_enabled_builder;
use core_test_support::apps_test_server::recorded_apps_tool_call_by_call_id;
use core_test_support::hooks::trust_discovered_hooks;
use core_test_support::responses;
use core_test_support::skip_if_no_network;
use core_test_support::skip_if_wine_exec;
use core_test_support::test_codex::test_codex;
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;
struct RecordedHistory {
call_id: String,
arguments: String,
items: Vec<ResponseItem>,
mcp_tool: Option<(String, Option<String>, McpToolSource)>,
}
#[derive(Default)]
@@ -31,6 +46,36 @@ struct ConversationHistoryRecorder {
histories: Mutex<Vec<RecordedHistory>>,
}
#[derive(Clone, Copy)]
enum AppsServerOwner {
Host,
Extension,
}
struct ExtensionOwnedAppsServer {
url: String,
}
impl McpServerContributor<Config> for ExtensionOwnedAppsServer {
fn id(&self) -> &'static str {
"extension_owned_apps_lifecycle_test"
}
fn contribute<'a>(
&'a self,
_context: McpServerContributionContext<'a, Config>,
) -> ExtensionFuture<'a, Vec<McpServerContribution>> {
Box::pin(async move {
let config = serde_json::from_value(json!({ "url": self.url }))
.expect("test Apps MCP server config should be valid");
vec![McpServerContribution::Set {
name: CODEX_APPS_MCP_SERVER_NAME.to_string(),
config: Box::new(config),
}]
})
}
}
impl ToolLifecycleContributor for ConversationHistoryRecorder {
fn on_tool_start<'a>(&'a self, input: ToolStartInput<'a>) -> ToolLifecycleFuture<'a> {
Box::pin(async move {
@@ -41,6 +86,13 @@ impl ToolLifecycleContributor for ConversationHistoryRecorder {
call_id: input.call_id.to_owned(),
arguments: input.payload.log_payload().into_owned(),
items: input.conversation_history.items().cloned().collect(),
mcp_tool: input.mcp_tool.map(|tool| {
(
tool.tool_info().server_name.clone(),
tool.tool_info().connector_id.clone(),
tool.source().clone(),
)
}),
});
})
}
@@ -155,6 +207,84 @@ async fn tool_start_receives_conversation_history() -> Result<()> {
Ok(())
}
#[test_case(AppsServerOwner::Host, McpToolSource::Connector; "host_owned_apps")]
#[test_case(AppsServerOwner::Extension, McpToolSource::Other; "extension_owned_apps")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn tool_start_receives_executed_mcp_call_for_connector(
owner: AppsServerOwner,
expected_source: McpToolSource,
) -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let apps_server = AppsTestServer::mount(&server).await?;
let call_id = "calendar-lifecycle-call";
responses::mount_sse_sequence(
&server,
vec![
responses::sse(vec![
responses::ev_function_call_with_namespace(
call_id,
SEARCH_CALENDAR_NAMESPACE,
SEARCH_CALENDAR_LIST_TOOL,
"{}",
),
responses::ev_completed("first-response"),
]),
responses::sse(vec![
responses::ev_assistant_message("assistant-1", "done"),
responses::ev_completed("second-response"),
]),
],
)
.await;
let recorder = Arc::new(ConversationHistoryRecorder::default());
let mut extensions = ExtensionRegistryBuilder::<Config>::new();
extensions.tool_lifecycle_contributor(recorder.clone());
if matches!(owner, AppsServerOwner::Extension) {
extensions.mcp_server_contributor(Arc::new(ExtensionOwnedAppsServer {
url: format!("{}/api/codex/ps/mcp", apps_server.chatgpt_base_url),
}));
}
let test = apps_enabled_builder(apps_server.chatgpt_base_url)
.with_extensions(Arc::new(extensions.build()))
.build_with_auto_env(&server)
.await?;
wait_for_mcp_server(&test.codex, CODEX_APPS_MCP_SERVER_NAME).await?;
test.submit_text_turn("List my calendar events.").await?;
{
let histories = recorder
.histories
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let [history] = histories.as_slice() else {
panic!("expected one tool start, got {}", histories.len());
};
assert_eq!(history.call_id, call_id);
assert_eq!(
history.mcp_tool,
Some((
CODEX_APPS_MCP_SERVER_NAME.to_string(),
Some("calendar".to_string()),
expected_source,
)),
);
}
let executed_call = recorded_apps_tool_call_by_call_id(&server, call_id).await;
assert_eq!(
executed_call
.pointer("/params/name")
.and_then(Value::as_str),
Some("calendar_list_events")
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn tool_start_receives_rewritten_payload_and_post_hook_history() -> Result<()> {
skip_if_no_network!(Ok(()));

View File

@@ -42,6 +42,8 @@ pub use thread_lifecycle::ThreadReadyInput;
pub use thread_lifecycle::ThreadResumeInput;
pub use thread_lifecycle::ThreadStartInput;
pub use thread_lifecycle::ThreadStopInput;
pub use tool_lifecycle::McpToolContext;
pub use tool_lifecycle::McpToolSource;
pub use tool_lifecycle::ToolCallOutcome;
pub use tool_lifecycle::ToolFinishInput;
pub use tool_lifecycle::ToolLifecycleFuture;

View File

@@ -2,6 +2,9 @@ use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use codex_config::McpServerConfig;
use codex_mcp::McpServerSource;
use codex_mcp::PreparedMcpCall;
use codex_tools::ToolCallSource;
use codex_tools::ToolName;
use codex_tools::ToolPayload;
@@ -33,6 +36,75 @@ pub enum ToolCallOutcome {
Aborted,
}
/// Provenance captured from the immutable MCP call selected for one tool invocation.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum McpToolSource {
/// A connector routed through the host-owned Codex Apps MCP server.
Connector,
/// An MCP server whose frozen registration matches the active Codex configuration.
Config,
/// An MCP server registered by a locally loaded plugin.
Plugin {
/// Identifier of the plugin that owns this MCP server.
id: String,
},
/// An executor-selected plugin whose root has not been attested by the host.
SelectedPlugin,
/// A compatibility or extension registration without user-owned provenance.
Other,
}
/// Read-only metadata and provenance captured from the MCP call that will execute.
#[derive(Clone, Debug)]
pub struct McpToolContext {
tool: crate::McpToolInfo,
source: McpToolSource,
}
impl McpToolContext {
/// Snapshots a prepared call without exposing its executable client to extensions.
///
/// Configured servers retain their provenance only when their captured connection
/// still matches the host configuration for the current tool invocation.
pub fn from_prepared_call(
call: &PreparedMcpCall,
configured_server: Option<&McpServerConfig>,
) -> Self {
let tool = call.tool_info().clone();
let source = if tool.connector_id.is_some() && call.is_host_owned_apps() {
McpToolSource::Connector
} else if call.is_selected_plugin_server() {
McpToolSource::SelectedPlugin
} else if let Some(id) = call.plugin_id() {
McpToolSource::Plugin { id: id.to_owned() }
} else if call
.config()
.mcp_server_catalog
.server(call.server_name())
.is_some_and(|server| {
matches!(server.source(), McpServerSource::Config)
&& configured_server.is_some_and(|configured| server.config() == configured)
})
{
McpToolSource::Config
} else {
McpToolSource::Other
};
Self { tool, source }
}
/// Returns frozen metadata for the exact model-visible MCP tool being executed.
pub fn tool_info(&self) -> &crate::McpToolInfo {
&self.tool
}
/// Returns the registration source captured with the executable call.
pub fn source(&self) -> &McpToolSource {
&self.source
}
}
/// Input supplied when the host starts executing one tool call.
pub struct ToolStartInput<'a> {
/// Store scoped to the host session runtime.
@@ -47,6 +119,8 @@ pub struct ToolStartInput<'a> {
pub call_id: &'a str,
/// Tool name as routed by the host.
pub tool_name: &'a ToolName,
/// Read-only metadata and provenance from the exact MCP call that will execute.
pub mcp_tool: Option<&'a McpToolContext>,
/// Finalized tool arguments, including any pre-tool-use hook rewrites.
///
/// Payloads can contain sensitive plaintext and must not be logged.

View File

@@ -17,6 +17,7 @@ pub use capabilities::NoopResponseItemInjector;
pub use capabilities::ResponseItemInjectionFuture;
pub use capabilities::ResponseItemInjector;
pub use codex_context_fragments::ContextualUserFragment;
pub use codex_mcp::ToolInfo as McpToolInfo;
pub use codex_protocol::models::ContentItemKind;
pub use codex_protocol::models::ResponseItem;
pub use codex_protocol::security_risk::SecurityRiskScore;
@@ -49,6 +50,8 @@ pub use contributors::ExtensionFuture;
pub use contributors::McpServerContribution;
pub use contributors::McpServerContributionContext;
pub use contributors::McpServerContributor;
pub use contributors::McpToolContext;
pub use contributors::McpToolSource;
pub use contributors::PreviousWorldStateSection;
pub use contributors::PromptFragment;
pub use contributors::PromptSlot;

View File

@@ -160,6 +160,7 @@ async fn installed_extension_reconnects_after_auth_refresh() -> Result<()> {
turn_id: "turn-1",
call_id,
tool_name: &tool_name,
mcp_tool: None,
payload: &payload,
conversation_history: Arc::new(TestConversationHistory(Vec::new())),
source: ToolCallSource::Direct,
@@ -360,6 +361,7 @@ async fn sandboxed_shell_classification_respects_review_scope() -> Result<()> {
turn_id: "turn-1",
call_id: "call-2",
tool_name: &tool_name,
mcp_tool: None,
payload: &payload,
conversation_history: Arc::new(TestConversationHistory(Vec::new())),
source: ToolCallSource::Direct,
@@ -437,6 +439,7 @@ async fn computer_use_only_scores_cannot_approve_other_actions() -> Result<()> {
turn_id: "turn-1",
call_id: "ordinary-call",
tool_name: &ordinary_tool,
mcp_tool: None,
payload: &payload,
conversation_history: Arc::new(TestConversationHistory(Vec::new())),
source: ToolCallSource::CodeMode {
@@ -812,6 +815,7 @@ async fn sample_configured_conversation_history_with_source(
turn_id: "turn-1",
call_id: "call-1",
tool_name: &tool_name,
mcp_tool: None,
payload: &tool_payload,
conversation_history: Arc::new(conversation_history),
source,
@@ -878,6 +882,7 @@ impl GuardianFailureFixture {
turn_id: "turn-1",
call_id: "call-1",
tool_name: &tool_name,
mcp_tool: None,
payload: &payload,
conversation_history: Arc::new(TestConversationHistory(Vec::new())),
source: ToolCallSource::Direct,
@@ -2015,6 +2020,7 @@ async fn contributor_skips_required_models_in_standard_scope() -> Result<()> {
turn_id: "turn-1",
call_id: "protected.md",
tool_name: &tool_name,
mcp_tool: None,
payload: &payload,
conversation_history: Arc::new(TestConversationHistory(vec![oversized_compaction])),
source: ToolCallSource::Direct,
@@ -2095,6 +2101,7 @@ async fn contributor_counts_failed_thread_lookups_toward_score_lag() -> Result<(
turn_id: "turn-1",
call_id: "missing.md",
tool_name: &tool_name,
mcp_tool: None,
payload: &payload,
conversation_history: Arc::new(TestConversationHistory(Vec::new())),
source: ToolCallSource::Direct,
@@ -2469,6 +2476,7 @@ async fn contributor_reuses_the_latest_compatible_parent_compaction() -> Result<
turn_id: "turn-1",
call_id: "call-1",
tool_name: &tool_name,
mcp_tool: None,
payload: &tool_payload,
conversation_history: Arc::new(conversation_history),
source: ToolCallSource::Direct,
@@ -2538,6 +2546,7 @@ async fn contributor_reuses_the_latest_compatible_parent_compaction() -> Result<
turn_id: "turn-1",
call_id: "call-2",
tool_name: &tool_name,
mcp_tool: None,
payload: &tool_payload,
conversation_history: Arc::new(TestConversationHistory(vec![
latest_compaction,