mirror of
https://github.com/openai/codex.git
synced 2026-09-04 15:08:45 +00:00
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:
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()),
|
||||
|
||||
@@ -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));
|
||||
|
||||
@@ -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(()));
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user