diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst index 66e60eba98..79241e9314 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst differ diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index a55a8bbbc5..387efab65e 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -22,7 +22,10 @@ use crate::context::ManagedDeveloperInstructions; use crate::context::MultiAgentRoleInstructions; use crate::context::SubagentNotification; use crate::init_state_db; +use crate::responses_metadata::CodexResponsesRequestKind; use crate::session::SessionSettingsUpdate; +use crate::session::step_context::StepContext; +use crate::thread_manager::ForkSnapshot; use crate::thread_manager::StartThreadOptions; use crate::tools::handlers::multi_agents_common::thread_spawn_source; use assert_matches::assert_matches; @@ -46,11 +49,15 @@ use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::ModeKind; +use codex_protocol::config_types::ReasoningSummary; use codex_protocol::config_types::Settings; use codex_protocol::error::CodexErrorDetails; use codex_protocol::items::TurnItem; use codex_protocol::items::UserMessageItem; use codex_protocol::mcp::ClientMcpExtensions; +use codex_protocol::mcp::McpAttribution; +use codex_protocol::mcp::McpAttributionSource; +use codex_protocol::mcp::McpAttributionStatus; use codex_protocol::mcp::OPENAI_FORM_EXTENSION_ID; use codex_protocol::models::ContentItem; use codex_protocol::models::ContentItemKind; @@ -425,6 +432,48 @@ async fn persist_thread_for_tree_resume(thread: &Arc, message: &str .expect("test thread rollout should flush"); } +/// Builds a request entirely in process; it does not install a transport or send anything. +/// The destination filtering behavior is covered separately in `client_tests`. +async fn mcp_attribution_in_constructed_request(thread: &CodexThread) -> McpAttribution { + let turn_context = thread.session.new_default_turn().await; + let step_context = StepContext::for_test(Arc::clone(&turn_context)); + let prompt = crate::session::turn::build_prompt( + thread + .session + .clone_history() + .await + .for_prompt(&step_context.settings.model_info.input_modalities), + &step_context, + thread.session.get_prompt_base_instructions().await, + ); + let metadata = thread + .session + .responses_metadata(&step_context, CodexResponsesRequestKind::Turn) + .await; + let request = thread + .session + .services + .model_client + .build_responses_request( + &prompt, + &step_context.settings.model_info, + /*effort*/ None, + ReasoningSummary::None, + /*service_tier*/ None, + &metadata, + /*include_internal*/ true, + ) + .expect("construct request"); + serde_json::from_str( + request + .client_metadata + .as_ref() + .and_then(|metadata| metadata.get("mcp_attribution")) + .expect("request MCP attribution"), + ) + .expect("valid request MCP attribution") +} + async fn wait_for_live_thread_spawn_children( control: &LocalAgentControl, parent_thread_id: ThreadId, @@ -3009,6 +3058,103 @@ async fn spawn_agent_full_fork_legacy_compaction_rebuilds_child_instructions_onc } } +#[tokio::test] +async fn resume_and_cold_fork_restore_mcp_attribution_in_constructed_requests() { + let harness = AgentControlHarness::new().await; + let (source_thread_id, source_thread) = harness.start_thread().await; + let turn_context = source_thread.session.new_default_turn().await; + source_thread + .session + .record_conversation_items( + turn_context.as_ref(), + turn_context.model_info(), + &[user_message("source turn")], + ) + .await; + let source = McpAttributionSource { + connector_id: None, + plugin_id: Some("messages@openai-bundled".to_string()), + server_name: "messages".to_string(), + tool_name: "search".to_string(), + first_turn_id: turn_context.sub_id.clone(), + }; + source_thread + .session + .services + .executed_tool_calls + .record_mcp_source(source.clone()); + source_thread + .session + .record_conversation_items( + turn_context.as_ref(), + turn_context.model_info(), + &[assistant_message("source result", /*phase*/ None)], + ) + .await; + source_thread.ensure_rollout_materialized().await; + source_thread + .flush_rollout() + .await + .expect("flush source rollout"); + let rollout_path = source_thread.rollout_path().expect("source rollout path"); + source_thread + .shutdown_and_wait() + .await + .expect("shutdown source thread"); + assert!( + harness + .manager + .remove_thread(&source_thread_id) + .await + .is_some() + ); + drop(source_thread); + + let forked = harness + .manager + .fork_thread( + ForkSnapshot::Interrupted, + StartThreadOptions::new(harness.config.clone()), + rollout_path.clone(), + ) + .await + .expect("fork unloaded source thread"); + let resumed = harness + .manager + .resume_thread_from_rollout( + harness.config.clone(), + rollout_path, + AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")), + /*parent_trace*/ None, + ClientMcpExtensions::default(), + ) + .await + .expect("resume source thread"); + let expected = McpAttribution { + status: McpAttributionStatus::Complete, + sources: vec![source], + }; + assert_eq!( + mcp_attribution_in_constructed_request(&forked.thread).await, + expected + ); + assert_eq!( + mcp_attribution_in_constructed_request(&resumed.thread).await, + expected + ); + + forked + .thread + .shutdown_and_wait() + .await + .expect("shutdown forked thread"); + resumed + .thread + .shutdown_and_wait() + .await + .expect("shutdown resumed thread"); +} + #[tokio::test] async fn spawn_agent_fork_flushes_parent_rollout_before_loading_history() { let harness = AgentControlHarness::new().await; @@ -3075,11 +3221,35 @@ async fn spawn_agent_fork_flushes_parent_rollout_before_loading_history() { async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { let harness = AgentControlHarness::new().await; let (parent_thread_id, parent_thread) = harness.start_thread().await; - + let old_turn_context = parent_thread.session.new_default_turn().await; parent_thread - .inject_response_items(vec![user_message("old parent context")]) - .await - .expect("inject old parent context"); + .session + .record_conversation_items( + old_turn_context.as_ref(), + old_turn_context.model_info(), + &[user_message("old parent context")], + ) + .await; + let source = McpAttributionSource { + connector_id: None, + plugin_id: None, + server_name: "example".to_string(), + tool_name: "search".to_string(), + first_turn_id: old_turn_context.sub_id.clone(), + }; + parent_thread + .session + .services + .executed_tool_calls + .record_mcp_source(source.clone()); + parent_thread + .session + .record_conversation_items( + old_turn_context.as_ref(), + old_turn_context.model_info(), + &[assistant_message("old MCP result", /*phase*/ None)], + ) + .await; let queued_communication = InterAgentCommunication::new( AgentPath::root(), AgentPath::try_from("/root/worker").expect("agent path"), @@ -3176,6 +3346,10 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { !history_contains_text(history.raw_items(), "old parent context"), "forked child history should drop parent context outside the requested last-N turn window" ); + assert!( + !history_contains_text(history.raw_items(), "old MCP result"), + "forked child history should drop the old result while retaining its attribution" + ); assert!( !history_contains_text(history.raw_items(), "queued message"), "forked child history should drop queued inter-agent messages outside the requested last-N turn window" @@ -3196,6 +3370,13 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { .is_none(), "last-N forked child should rebuild context after truncating the cached prefix" ); + assert_eq!( + mcp_attribution_in_constructed_request(&child_thread).await, + McpAttribution { + status: McpAttributionStatus::Complete, + sources: vec![source], + } + ); let _ = harness .control diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 1c5b49353a..bc5f59e479 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -405,6 +405,16 @@ fn response_items_equal_ignoring_internal_metadata( previous == current } +/// Whether the resolved outbound Responses destination may receive internal tool metadata. +fn is_internal_metadata_destination(provider: &ApiProvider) -> bool { + url::Url::parse(&provider.base_url).ok().is_some_and(|url| { + url.scheme() == "https" + && url.host_str().is_some_and(|host| { + host == "api.openai.com" || codex_http_client::is_allowed_chatgpt_host(host) + }) + }) +} + impl WebsocketSession { fn reset(&mut self, reason: Option<&'static str>) { // Per-socket backend metrics call a resend after reconnect "initial". @@ -802,9 +812,10 @@ impl ModelClient { fn build_ws_client_metadata( &self, responses_metadata: &CodexResponsesMetadata, + include_internal: bool, use_responses_lite: bool, ) -> HashMap { - let mut client_metadata = responses_metadata.client_metadata(); + let mut client_metadata = responses_metadata.client_metadata(include_internal); if use_responses_lite { client_metadata.insert( WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY.to_string(), @@ -866,6 +877,7 @@ impl ModelClient { } } + #[allow(clippy::too_many_arguments)] pub(crate) fn build_responses_request( &self, prompt: &Prompt, @@ -874,6 +886,7 @@ impl ModelClient { summary: ReasoningSummaryConfig, service_tier: Option, responses_metadata: &CodexResponsesMetadata, + include_internal: bool, ) -> Result { let mut input = prompt.get_formatted_input_for_request(model_info); if !self.reasoning_effort_override_enabled(model_info) { @@ -963,6 +976,12 @@ impl ModelClient { } else { model_info.service_tier_for_request(service_tier) }; + if !include_internal { + for item in &mut input { + item.clear_tool_result_metadata(); + } + } + let client_metadata = responses_metadata.client_metadata(include_internal); let request = ResponsesApiRequest { model: model_info.slug.clone(), instructions, @@ -978,32 +997,12 @@ impl ModelClient { service_tier, prompt_cache_key, text, - client_metadata: Some(responses_metadata.client_metadata()), + client_metadata: Some(client_metadata), access_programs: None, }; Ok(request) } - fn filter_tool_result_metadata(input: &mut [ResponseItem], api_provider: &ApiProvider) { - // Check the resolved destination only when sending, not for local budget estimates. - // HTTP and WS (including v2 compaction) share this raw-metadata-only filter. - let result_metadata_allowed = - url::Url::parse(&api_provider.base_url) - .ok() - .is_some_and(|url| { - url.scheme() == "https" - && url.host_str().is_some_and(|host| { - host == "api.openai.com" - || codex_http_client::is_allowed_chatgpt_host(host) - }) - }); - if !result_metadata_allowed { - for item in input { - item.clear_tool_result_metadata(); - } - } - } - fn prepare_response_items_for_request(&self, input: &mut [ResponseItem]) { for item in input { if item.id().is_some_and(|id| !id.is_prefixed()) { @@ -1629,6 +1628,7 @@ impl ModelClientSession { .client .current_client_setup(ClientRouting::Workspace) .await?; + let include_internal = is_internal_metadata_destination(&client_setup.api_provider); let responses_headers = self .client .responses_headers(client_setup.auth.as_ref(), &model_info.slug); @@ -1666,11 +1666,8 @@ impl ModelClientSession { summary, service_tier.clone(), responses_metadata, + include_internal, )?; - ModelClient::filter_tool_result_metadata( - &mut request.input, - &client_setup.api_provider, - ); self.client.set_guardian_metadata( &mut request.client_metadata, responses_metadata.parent_response_id.as_deref(), @@ -1811,6 +1808,7 @@ impl ModelClientSession { .client .current_client_setup(ClientRouting::Workspace) .await?; + let include_internal = is_internal_metadata_destination(&client_setup.api_provider); let responses_headers = self .client .responses_headers(client_setup.auth.as_ref(), &model_info.slug); @@ -1828,11 +1826,8 @@ impl ModelClientSession { summary, service_tier.clone(), responses_metadata, + include_internal, )?; - ModelClient::filter_tool_result_metadata( - &mut request.input, - &client_setup.api_provider, - ); let guardian_reviewer = responses_headers .get("x-codex-guardian") .is_some_and(|value| value == "reviewer"); @@ -1905,9 +1900,11 @@ impl ModelClientSession { { crate::guardian::observe_guardian_request(session_telemetry, &request); } - let mut client_metadata = self - .client - .build_ws_client_metadata(responses_metadata, model_info.use_responses_lite); + let mut client_metadata = self.client.build_ws_client_metadata( + responses_metadata, + include_internal, + model_info.use_responses_lite, + ); if let Some(turn_state) = self.turn_state.get() { client_metadata.insert(X_CODEX_TURN_STATE_HEADER.to_string(), turn_state.clone()); } diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 24b785a6fa..0f62891632 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -12,6 +12,8 @@ use crate::AttestationContext; use crate::AttestationProvider; use crate::GenerateAttestationFuture; use crate::responses_metadata::CodexResponsesMetadata; +use crate::responses_metadata::MAX_MCP_ATTRIBUTION_BYTES; +use crate::responses_metadata::MCP_ATTRIBUTION_CLIENT_METADATA_KEY; use crate::test_support::TestCodexResponsesRequestKind; use crate::test_support::responses_metadata as test_responses_metadata; use base64::Engine; @@ -42,6 +44,9 @@ use codex_protocol::ThreadId; use codex_protocol::auth::AuthMode; use codex_protocol::error::CodexErr; use codex_protocol::error::CodexErrorDetails; +use codex_protocol::mcp::McpAttribution; +use codex_protocol::mcp::McpAttributionSource; +use codex_protocol::mcp::McpAttributionStatus; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; use codex_protocol::models::ExecutedToolCall; @@ -479,7 +484,7 @@ fn output_with_tool_result_metadata(metadata: ToolResultMetadata) -> ResponseIte } #[test] -fn responses_request_limits_raw_tool_metadata_to_resolved_first_party_https_endpoint() +fn responses_request_limits_internal_metadata_to_resolved_first_party_https_endpoint() -> anyhow::Result<()> { let provider = ModelProviderInfo::create_openai_provider(Some("https://api.openai.com/v1".to_string())); @@ -492,17 +497,28 @@ fn responses_request_limits_raw_tool_metadata_to_resolved_first_party_https_endp "private": { "resource": "raw-result-metadata" }, }))); let without_raw_metadata = output_with_tool_result_metadata(ToolResultMetadata::default()); + let attribution = McpAttribution { + status: McpAttributionStatus::Complete, + sources: vec![McpAttributionSource { + connector_id: Some("connector_example".to_string()), + plugin_id: Some("example@openai-bundled".to_string()), + server_name: "codex_apps".to_string(), + tool_name: "search".to_string(), + first_turn_id: "turn_123".to_string(), + }], + }; let prompt = Prompt { input: vec![output.clone()], ..Default::default() }; - let responses_metadata = test_responses_metadata_for_client( + let mut responses_metadata = test_responses_metadata_for_client( &client, /*turn_id*/ None, format!("{}:0", client.state.thread_id), /*parent_thread_id*/ None, TestCodexResponsesRequestKind::Turn, ); + responses_metadata.mcp_attribution = Some(attribution.clone()); for (base_url, allowed) in [ ("https://api.openai.com/v1", true), ("https://chatgpt.com/backend-api/codex", true), @@ -515,18 +531,19 @@ fn responses_request_limits_raw_tool_metadata_to_resolved_first_party_https_endp ("not a URL", false), ] { api_provider.base_url = base_url.to_string(); + let include_internal = super::is_internal_metadata_destination(&api_provider); for responses_lite in [false, true] { let mut model = test_model_info(); model.use_responses_lite = responses_lite; - let mut request = client.build_responses_request( + let request = client.build_responses_request( &prompt, &model, /*effort*/ None, codex_protocol::config_types::ReasoningSummary::None, /*service_tier*/ None, &responses_metadata, + include_internal, )?; - ModelClient::filter_tool_result_metadata(&mut request.input, &api_provider); assert_eq!( request.input.last(), Some(if allowed { @@ -536,9 +553,58 @@ fn responses_request_limits_raw_tool_metadata_to_resolved_first_party_https_endp }), "resolved endpoint: {base_url}, responses_lite: {responses_lite}", ); + let expected_json = allowed + .then(|| serde_json::to_string(&attribution).map(serde_json::Value::String)) + .transpose()?; + assert_eq!( + serde_json::to_value(&request)? + .get("client_metadata") + .and_then(|metadata| metadata.get(MCP_ATTRIBUTION_CLIENT_METADATA_KEY)), + expected_json.as_ref(), + ); + let ws_client_metadata = client.build_ws_client_metadata( + &responses_metadata, + include_internal, + responses_lite, + ); + let ws_request = codex_api::ResponseCreateWsRequest { + client_metadata: Some(ws_client_metadata), + ..codex_api::ResponseCreateWsRequest::from(&request) + }; + assert_eq!( + serde_json::to_value(ws_request)? + .get("client_metadata") + .and_then(|metadata| metadata.get(MCP_ATTRIBUTION_CLIENT_METADATA_KEY)), + expected_json.as_ref(), + ); + assert!( + serde_json::to_value(&request)? + .get("mcp_attribution") + .is_none() + ); assert_eq!(prompt.input, vec![output.clone()]); } } + let mut oversized_attribution = attribution; + oversized_attribution.sources[0].tool_name = "a".repeat(MAX_MCP_ATTRIBUTION_BYTES); + responses_metadata.mcp_attribution = Some(oversized_attribution); + let request = client.build_responses_request( + &prompt, + &test_model_info(), + /*effort*/ None, + codex_protocol::config_types::ReasoningSummary::None, + /*service_tier*/ None, + &responses_metadata, + /*include_internal*/ true, + )?; + assert_eq!( + request.client_metadata.as_ref().and_then(|metadata| { + metadata + .get(MCP_ATTRIBUTION_CLIENT_METADATA_KEY) + .map(String::as_str) + }), + Some(r#"{"status":"attribution_error"}"#), + ); Ok(()) } @@ -615,7 +681,9 @@ fn websocket_incremental_reuse_tracks_raw_result_metadata() -> anyhow::Result<() }); previous_output.set_turn_id_if_missing("previous-turn"); current_output.set_turn_id_if_missing("current-turn"); - let mut previous = client.build_responses_request( + api_provider.base_url = base_url.to_string(); + let include_internal = super::is_internal_metadata_destination(&api_provider); + let previous = client.build_responses_request( &Prompt { input: vec![previous_output], ..Default::default() @@ -625,16 +693,24 @@ fn websocket_incremental_reuse_tracks_raw_result_metadata() -> anyhow::Result<() codex_protocol::config_types::ReasoningSummary::None, /*service_tier*/ None, &responses_metadata, + include_internal, )?; let follow_up = ResponseItem::from(ResponseInputItem::FunctionCallOutput { call_id: "wait-call".to_string(), output: FunctionCallOutputPayload::from_text("done".to_string()), }); - let mut current = previous.clone(); - current.input = vec![current_output, follow_up.clone()]; - api_provider.base_url = base_url.to_string(); - ModelClient::filter_tool_result_metadata(&mut previous.input, &api_provider); - ModelClient::filter_tool_result_metadata(&mut current.input, &api_provider); + let current = client.build_responses_request( + &Prompt { + input: vec![current_output, follow_up.clone()], + ..Default::default() + }, + &test_model_info(), + /*effort*/ None, + codex_protocol::config_types::ReasoningSummary::None, + /*service_tier*/ None, + &responses_metadata, + include_internal, + )?; let mut session = client.new_session(); session.websocket_session.last_request = Some(previous); @@ -760,6 +836,7 @@ fn responses_lite_prefix_ids_track_thread_and_payload() -> anyhow::Result<()> { /*parent_thread_id*/ None, TestCodexResponsesRequestKind::Turn, ), + /*include_internal*/ true, ) }; @@ -834,6 +911,7 @@ fn websocket_continuation_reset_reason_survives_failed_reconnect_and_turn_bounda /*parent_thread_id*/ None, TestCodexResponsesRequestKind::Turn, ), + /*include_internal*/ true, ) .expect("build continuation request"); let mut session = client.new_session(); @@ -880,6 +958,7 @@ fn reasoning_effort_in_request( /*parent_thread_id*/ None, TestCodexResponsesRequestKind::Turn, ), + /*include_internal*/ true, ) .expect("build responses request") .reasoning @@ -1173,8 +1252,11 @@ fn build_ws_client_metadata_includes_window_lineage_and_turn_metadata() { Some(parent_thread_id), TestCodexResponsesRequestKind::Turn, ); - let client_metadata = - client.build_ws_client_metadata(&responses_metadata, /*use_responses_lite*/ false); + let client_metadata = client.build_ws_client_metadata( + &responses_metadata, + /*include_internal*/ true, + /*use_responses_lite*/ false, + ); let parent_thread_id = parent_thread_id.to_string(); let turn_metadata: serde_json::Value = serde_json::from_str( client_metadata diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index e102da28b4..c1e18ba7eb 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -267,14 +267,9 @@ async fn run_compact_task_inner_impl( let max_retries = turn_context.provider.info().stream_max_retries(); let mut retries = 0; + // Reuse one client session so turn-scoped state (sticky routing and websocket incremental + // request tracking) survives retries within this compact turn. let mut client_session = sess.services.model_client.new_session(); - // Reuse one client session so turn-scoped state (sticky routing, websocket incremental - // request tracking) - // survives retries within this compact turn. - let responses_metadata = sess - .compaction_responses_metadata(turn_context.as_ref(), compaction_metadata) - .await; - let compaction_response = loop { // Clone is required because of the loop let mut turn_input = history @@ -289,6 +284,9 @@ async fn run_compact_task_inner_impl( base_instructions: sess.get_prompt_base_instructions().await, ..Default::default() }; + let responses_metadata = sess + .compaction_responses_metadata(turn_context.as_ref(), compaction_metadata) + .await; let attempt_result = drain_to_completed( &sess, turn_context.as_ref(), diff --git a/codex-rs/core/src/guardian/input_budget.rs b/codex-rs/core/src/guardian/input_budget.rs index 9514eb31e8..2fce16ebaa 100644 --- a/codex-rs/core/src/guardian/input_budget.rs +++ b/codex-rs/core/src/guardian/input_budget.rs @@ -100,6 +100,7 @@ pub(crate) async fn finalize( &session .responses_metadata(step, CodexResponsesRequestKind::Turn) .await, + /*include_internal*/ true, )?; let mut existing = super::request_budget::estimate_request_tokens(&request) .max(usize::try_from(session.get_total_token_usage().await).unwrap_or(usize::MAX)); diff --git a/codex-rs/core/src/guardian/request_budget.rs b/codex-rs/core/src/guardian/request_budget.rs index 0f9fa0e0fe..8de9ff59ea 100644 --- a/codex-rs/core/src/guardian/request_budget.rs +++ b/codex-rs/core/src/guardian/request_budget.rs @@ -70,6 +70,7 @@ pub(crate) fn check_prompt( codex_protocol::config_types::ReasoningSummary::None, /*service_tier*/ None, metadata, + /*include_internal*/ true, )?; if estimate_request_tokens(&request) > effective_input_token_limit(model, config.model_context_window) diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 9c32660abb..ad8f420b2d 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -595,6 +595,21 @@ async fn handle_approved_mcp_tool_call( if let Some(elicitation_type) = elicitation_type { track_mcp_tool_call_elicitation(sess, turn_context, call_id, elicitation_type); } + // Direct and Code Mode calls share this boundary. Once an approved call enters + // it, conservatively attribute returned errors too: they may contain peer data. + let source_connector_id = prepared_call + .is_host_owned_apps() + .then(|| prepared_call.tool_info().connector_id.clone()) + .flatten(); + sess.services.executed_tool_calls.record_mcp_source( + codex_protocol::mcp::McpAttributionSource { + connector_id: source_connector_id, + plugin_id: prepared_call.plugin_id().map(str::to_string), + server_name: prepared_call.server_name().to_string(), + tool_name: prepared_call.tool_info().tool.name.to_string(), + first_turn_id: turn_context.sub_id.clone(), + }, + ); notify_mcp_tool_call_completed( sess, turn_context, diff --git a/codex-rs/core/src/responses_metadata.rs b/codex-rs/core/src/responses_metadata.rs index fbd3661c3b..dcc6f5e407 100644 --- a/codex-rs/core/src/responses_metadata.rs +++ b/codex-rs/core/src/responses_metadata.rs @@ -8,10 +8,12 @@ use codex_analytics::CompactionStrategy; use codex_analytics::CompactionTrigger; use codex_git_utils::SanitizedGitUrl; use codex_protocol::ThreadId; +use codex_protocol::mcp::McpAttribution; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadSource; use codex_utils_string::to_ascii_json_string; +use codex_utils_string::to_json_string_bounded; use http::HeaderMap as ApiHeaderMap; use http::HeaderValue; use serde::Serialize; @@ -40,6 +42,8 @@ pub(crate) const TOOL_NAMESPACES_INFO_KEY: &str = "tool_namespaces_info"; pub(crate) const TURN_STARTED_AT_UNIX_MS_KEY: &str = "turn_started_at_unix_ms"; pub(crate) const HISTORY_INGEST_REQUESTED_KEY: &str = "history_ingest_requested"; pub(crate) const ANALYTICS_ENABLED_KEY: &str = "analytics_enabled"; +pub(crate) const MCP_ATTRIBUTION_CLIENT_METADATA_KEY: &str = "mcp_attribution"; +pub(crate) const MAX_MCP_ATTRIBUTION_BYTES: usize = 16 * 1024; pub(crate) const FORKED_FROM_THREAD_ID_KEY: &str = "forked_from_thread_id"; pub(crate) const FORKED_FROM_ORDINAL_EXCLUSIVE_KEY: &str = "forked_from_ordinal_exclusive"; @@ -60,6 +64,7 @@ pub(crate) const WORKSPACES_KEY: &str = "workspaces"; // when submitting a turn, but they must not override fields owned by core. const RESERVED_METADATA_KEYS: &[&str] = &[ "guardian_credits_requested", + MCP_ATTRIBUTION_CLIENT_METADATA_KEY, INSTALLATION_ID_KEY, X_CODEX_INSTALLATION_ID_HEADER, SESSION_ID_KEY, @@ -95,6 +100,7 @@ const RESERVED_METADATA_KEYS: &[&str] = &[ NODE_REPL_DISABLED_KEY, WORKSPACES_KEY, ]; + // These keys were previously valid user configuration. Accept existing configs while filtering // their values before constructing Core-owned request metadata. const BACKWARD_COMPATIBLE_RESERVED_METADATA_KEYS: &[&str] = &[ @@ -255,6 +261,8 @@ pub struct CodexResponsesMetadata { /// Selected session analytics client's collection state, independent of event eligibility or delivery. /// Absent when the request has no initialized session analytics context. pub(crate) analytics_enabled: Option, + /// Cumulative MCP attribution for this logical request; body-only and model-invisible. + pub(crate) mcp_attribution: Option, pub(crate) extra: BTreeMap, } @@ -296,6 +304,7 @@ impl CodexResponsesMetadata { turn_started_at_unix_ms: None, history_ingest_requested: None, analytics_enabled: None, + mcp_attribution: None, extra: BTreeMap::new(), } } @@ -312,7 +321,7 @@ impl CodexResponsesMetadata { serde_json::to_value(self.turn_metadata_payload()).ok() } - pub(crate) fn client_metadata(&self) -> HashMap { + pub(crate) fn client_metadata(&self, include_internal: bool) -> HashMap { let mut client_metadata = HashMap::from([ ( X_CODEX_INSTALLATION_ID_HEADER.to_string(), @@ -348,6 +357,11 @@ impl CodexResponsesMetadata { { client_metadata.insert(X_CODEX_TURN_METADATA_HEADER.to_string(), turn_metadata_json); } + if include_internal && let Some(attribution) = &self.mcp_attribution { + let serialized = to_json_string_bounded(attribution, MAX_MCP_ATTRIBUTION_BYTES) + .unwrap_or_else(|_| r#"{"status":"attribution_error"}"#.to_string()); + client_metadata.insert(MCP_ATTRIBUTION_CLIENT_METADATA_KEY.to_string(), serialized); + } client_metadata } diff --git a/codex-rs/core/src/session/inject_tests.rs b/codex-rs/core/src/session/inject_tests.rs index 92446c1d0e..b8a67a986a 100644 --- a/codex-rs/core/src/session/inject_tests.rs +++ b/codex-rs/core/src/session/inject_tests.rs @@ -14,7 +14,7 @@ async fn harness_authored_configuration_updates_preserve_metadata_and_resume() { let (session, turn_context, rx_event) = make_session_and_context_with_rx().await; assert!(!session.enabled(Feature::RetainClientDeveloperMessages)); - let expected = ResponseItemEnvelope { + let mut expected = ResponseItemEnvelope { item: ResponseItem::ConfigurationUpdate { reasoning: ConfigurationReasoning { effort: ReasoningEffort::High, @@ -33,6 +33,13 @@ async fn harness_authored_configuration_updates_preserve_metadata_and_resume() { ) .await; + expected.metadata.as_mut().unwrap().mcp_attribution = Some( + session + .services + .executed_tool_calls + .mcp_attribution_snapshot(), + ); + let recorded = session.clone_history().await.into_annotated_items(); assert_eq!(recorded, vec![expected.clone()]); let mut raw_items = Vec::new(); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index a4b4ba0032..ac7a8d431d 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -158,6 +158,7 @@ use codex_protocol::request_permissions::RequestPermissionsResponse; use codex_protocol::request_user_input::RequestUserInputArgs; use codex_protocol::request_user_input::RequestUserInputResponse; use codex_rmcp_client::ElicitationResponse; +use codex_rollout::should_persist_response_item; use codex_rollout::state_db; use codex_rollout_trace::ThreadStartedTraceMetadata; use codex_rollout_trace::ThreadTraceContext; @@ -3554,6 +3555,32 @@ impl Session { .get_or_insert_with(|| with_serialization_allowance(policy).token_budget()); } } + // Last-N-turn forks retain a suffix starting at a user turn boundary. Repeat the + // cumulative checkpoint there so the suffix remains self-contained even when the fork + // happens mid-turn; dirty checkpoints preserve new sources between boundaries. + let force_mcp_checkpoint = items + .iter() + .any(|envelope| crate::context_manager::is_user_turn_boundary(&envelope.item)); + let mcp_revision = self + .services + .executed_tool_calls + .mcp_attribution_checkpoint(force_mcp_checkpoint) + .and_then(|(attribution, revision)| { + let first_persisted = items + .iter() + .position(|envelope| should_persist_response_item(&envelope.item))?; + for (index, envelope) in items.iter_mut().enumerate() { + // Rollout batches may be partially written. Checkpoint the first persisted + // item, and repeat the checkpoint at turn boundaries retained by forks. + if index == first_persisted + || crate::context_manager::is_user_turn_boundary(&envelope.item) + { + envelope.metadata.get_or_insert_default().mcp_attribution = + Some(attribution.clone()); + } + } + Some(revision) + }); let response_items = items .iter() .map(|envelope| envelope.item.clone()) @@ -3593,7 +3620,13 @@ impl Session { } let rollout_items: Vec = items.into_iter().map(RolloutItem::ResponseItem).collect(); - self.persist_rollout_items(&rollout_items).await; + if self.persist_rollout_items(&rollout_items).await + && let Some(revision) = mcp_revision + { + self.services + .executed_tool_calls + .mark_mcp_attribution_persisted(revision); + } if turn_context.config.memories.disable_on_external_context && let Some(item) = response_items .iter() @@ -3959,6 +3992,16 @@ impl Session { for envelope in &mut items { Self::assign_missing_response_item_id(&mut envelope.item); } + let mcp_revision = self + .services + .executed_tool_calls + .mcp_attribution_checkpoint(/*force*/ true) + .and_then(|(attribution, revision)| { + items.last_mut().map(|envelope| { + envelope.metadata.get_or_insert_default().mcp_attribution = Some(attribution); + revision + }) + }); if let Some(checkpoint) = items.iter_mut().rev().find(|envelope| { matches!( envelope.item, @@ -4022,7 +4065,13 @@ impl Session { rollout_items.push(RolloutItem::EventMsg( thread_settings::applied_event(self).await, )); - self.persist_rollout_items(&rollout_items).await; + if self.persist_rollout_items(&rollout_items).await + && let Some(revision) = mcp_revision + { + self.services + .executed_tool_calls + .mark_mcp_attribution_persisted(revision); + } { let mut state = self.state.lock().await; state.queue_pending_session_start_source(codex_hooks::SessionStartSource::Compact); @@ -4354,12 +4403,14 @@ impl Session { } #[tracing::instrument(level = "trace", skip_all, fields(item_count = items.len()))] - pub(crate) async fn persist_rollout_items(&self, items: &[RolloutItem]) { + pub(crate) async fn persist_rollout_items(&self, items: &[RolloutItem]) -> bool { if let Some(live_thread) = self.live_thread() && let Err(e) = live_thread.append_items(items).await { error!("failed to record rollout items: {e:#}"); + return false; } + true } pub(crate) async fn clone_history(&self) -> ContextManager { diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 92c1481375..0f729c2e21 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -716,6 +716,7 @@ impl Session { CodexResponsesMetadata { window_number: Some(window_number), context_window_id: Some(context_window_id), + mcp_attribution: Some(self.services.executed_tool_calls.mcp_attribution_snapshot()), analytics_enabled: Some(self.services.analytics_events_client.is_enabled()), history_ingest_requested: turn_context .config diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index fc1df5abf7..a4b3219867 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -33,8 +33,10 @@ use crate::shell_snapshot::ShellSnapshot; use crate::test_support::models_manager_with_provider; use crate::tools::format_exec_output_str; use crate::tools::registry::ToolRegistry; +use codex_analytics::CompactionImplementation; use codex_analytics::CompactionPhase; use codex_analytics::CompactionReason; +use codex_analytics::CompactionTrigger; use codex_config::ConfigLayerStack; use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; use codex_config::LoaderOverrides; @@ -107,6 +109,7 @@ use tracing::Span; use crate::connectors::AppInfo; use crate::responses_metadata::CodexResponsesRequestKind; +use crate::responses_metadata::CompactionTurnMetadata; use crate::rollout::recorder::RolloutRecorder; use crate::state::ActiveTurn; use crate::state::TaskKind; @@ -149,6 +152,9 @@ use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::Settings; use codex_protocol::items::HookPromptFragment; use codex_protocol::items::build_hook_prompt_message; +use codex_protocol::mcp::McpAttribution; +use codex_protocol::mcp::McpAttributionSource; +use codex_protocol::mcp::McpAttributionStatus; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; use codex_protocol::models::ContentItemKind; @@ -2662,6 +2668,12 @@ async fn annotated_history_uses_explicit_model_without_a_step( session .record_annotated_conversation_items(&turn_context, &model_info, expected.clone()) .await; + expected[0].metadata.get_or_insert_default().mcp_attribution = Some( + session + .services + .executed_tool_calls + .mcp_attribution_snapshot(), + ); for envelope in &mut expected { envelope .metadata @@ -3707,9 +3719,14 @@ async fn start_new_context_window_persists_checkpoint_state() { | RolloutItem::RealtimeItem(_) | RolloutItem::EventMsg(_) => None, }); + let mut expected_history = live_history.annotated_items().to_vec(); + // Compaction's parallel metadata vector represents absent entries as default metadata. + for envelope in &mut expected_history { + envelope.metadata.get_or_insert_default(); + } assert_eq!( persisted_compacted.and_then(|compacted| compacted.replacement_history.clone()), - Some(live_history.annotated_items().to_vec()) + Some(expected_history) ); assert_eq!( persisted_compacted.map(|compacted| { @@ -5603,6 +5620,102 @@ async fn settings_checkpoint_waits_for_accepted_settings_persistence() { ); } +#[tokio::test] +async fn mcp_attribution_checkpoints_cover_batch_prefixes_compaction_and_restore() { + let (mut session, turn_context) = make_session_and_context().await; + session.services.executed_tool_calls = + crate::state::ExecutedToolCalls::new(&turn_context.config.features, &InitialHistory::New); + let rollout_path = attach_thread_persistence(&mut session).await; + let source = McpAttributionSource { + connector_id: None, + plugin_id: None, + server_name: "example".to_string(), + tool_name: "search".to_string(), + first_turn_id: "turn_1".to_string(), + }; + session + .services + .executed_tool_calls + .record_mcp_source(source.clone()); + let expected = McpAttribution { + status: McpAttributionStatus::Complete, + sources: vec![source], + }; + session + .record_annotated_conversation_items( + &turn_context, + turn_context.model_info(), + vec![ + ResponseItemEnvelope::new(ResponseItem::FunctionCallOutput { + id: None, + call_id: Some("call_1".to_string()), + name: None, + namespace: None, + output: FunctionCallOutputPayload::from_text("result".to_string()), + internal_chat_message_metadata_passthrough: None, + }), + ResponseItemEnvelope::new(user_message("next turn")), + ], + ) + .await; + let checkpoints = session + .clone_history() + .await + .annotated_items() + .iter() + .map(|envelope| { + envelope + .metadata + .as_ref() + .and_then(|metadata| metadata.mcp_attribution.clone()) + }) + .collect::>(); + assert_eq!( + checkpoints, + vec![Some(expected.clone()), Some(expected.clone())] + ); + + let (window_number, window_ids) = session.advance_auto_compact_window().await; + session + .replace_compacted_history( + vec![ResponseItemEnvelope::new(user_message("compacted history"))], + /*reference_context_item*/ None, + /*world_state_baseline*/ None, + CompactedHistoryMetadata { + message: "summary".to_string(), + window_number, + window_ids, + compaction_response_id: None, + compaction_model_hash: None, + reviewer_compaction_hash: None, + }, + ) + .await; + session + .flush_rollout() + .await + .expect("flush compacted history"); + + let (items, _, _) = RolloutRecorder::load_rollout_items(&rollout_path) + .await + .expect("read compacted history"); + let checkpoint = items.iter().rev().find_map(|item| match item { + RolloutItem::Compacted(compacted) => compacted + .replacement_history + .as_ref()? + .iter() + .rev() + .find_map(|envelope| envelope.metadata.as_ref()?.mcp_attribution.as_ref()), + _ => None, + }); + assert_eq!(checkpoint, Some(&expected)); + let restored = crate::state::ExecutedToolCalls::new( + &turn_context.config.features, + &InitialHistory::Forked(items), + ); + assert_eq!(restored.mcp_attribution_snapshot(), expected); +} + #[tokio::test] async fn session_settings_commit_keeps_snapshot_across_postcommit_wait() { let (session, _turn_context, _rx) = make_session_and_context_with_auth_and_config_and_rx( @@ -6018,6 +6131,50 @@ pub(crate) async fn build_world_state_from_turn_context( .expect("world state should build") } +#[tokio::test] +async fn response_metadata_builders_capture_fresh_mcp_attribution() { + let (mut session, turn_context) = make_session_and_context().await; + session.services.executed_tool_calls = + crate::state::ExecutedToolCalls::new(&turn_context.config.features, &InitialHistory::New); + let turn_context = Arc::new(turn_context); + let step_context = StepContext::for_test(Arc::clone(&turn_context)); + let before = session + .responses_metadata(&step_context, CodexResponsesRequestKind::Turn) + .await; + let source = McpAttributionSource { + connector_id: None, + plugin_id: None, + server_name: "example".to_string(), + tool_name: "search".to_string(), + first_turn_id: turn_context.sub_id.clone(), + }; + session + .services + .executed_tool_calls + .record_mcp_source(source.clone()); + let after = session + .responses_metadata(&step_context, CodexResponsesRequestKind::Turn) + .await; + let compaction = session + .compaction_responses_metadata( + &turn_context, + CompactionTurnMetadata::new( + CompactionTrigger::Auto, + CompactionReason::ContextLimit, + CompactionImplementation::Responses, + CompactionPhase::MidTurn, + ), + ) + .await; + let expected = Some(McpAttribution { + status: McpAttributionStatus::Complete, + sources: vec![source], + }); + assert_eq!(before.mcp_attribution, Some(McpAttribution::default())); + assert_eq!(after.mcp_attribution, expected); + assert_eq!(compaction.mcp_attribution, expected); +} + #[tokio::test] async fn responses_metadata_uses_selected_harness_analytics_client() { for enabled in [true, false] { diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 360607f3f4..c17ff03a7e 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -518,16 +518,12 @@ pub(crate) async fn run_turn( .instrument(trace_span!("run_turn.prepare_sampling_request_input")) .await; - let responses_metadata = sess - .responses_metadata(step_context.as_ref(), CodexResponsesRequestKind::Turn) - .await; run_sampling_request( Arc::clone(&sess), Arc::clone(&step_context), Arc::clone(&turn_context.extension_data), Arc::clone(&turn_diff_tracker), &mut client_session, - &responses_metadata, sampling_request_input, cancellation_token.child_token(), ) @@ -1590,7 +1586,6 @@ async fn run_sampling_request( turn_store: Arc, turn_diff_tracker: SharedTurnDiffTracker, client_session: &mut ModelClientSession, - responses_metadata: &CodexResponsesMetadata, input: Vec, cancellation_token: CancellationToken, ) -> CodexResult<(SamplingRequestResult, Vec)> { @@ -1631,13 +1626,16 @@ async fn run_sampling_request( step_context.as_ref(), base_instructions.clone(), ); + let responses_metadata = sess + .responses_metadata(step_context.as_ref(), CodexResponsesRequestKind::Turn) + .await; if crate::guardian::is_basic_session_source(&turn_context.session_source) { crate::guardian::check_guardian_prompt_budget( &sess, &prompt, &turn_context.config, &step_context.settings.model_info, - responses_metadata, + &responses_metadata, )?; } let err = match try_run_sampling_request( @@ -1646,7 +1644,7 @@ async fn run_sampling_request( Arc::clone(&step_context), Arc::clone(&turn_store), client_session, - responses_metadata, + &responses_metadata, Arc::clone(&turn_diff_tracker), &prompt, cancellation_token.child_token(), diff --git a/codex-rs/core/src/test_support.rs b/codex-rs/core/src/test_support.rs index b93c2ba700..0cf04f0f4c 100644 --- a/codex-rs/core/src/test_support.rs +++ b/codex-rs/core/src/test_support.rs @@ -64,6 +64,18 @@ pub async fn history_with_tool_call_metadata( items } +/// Returns the recorder state used by the next session metadata snapshot. +/// This does not change destination filtering or issue a request. +pub fn mcp_attribution_snapshot( + thread: &crate::CodexThread, +) -> codex_protocol::mcp::McpAttribution { + thread + .session + .services + .executed_tool_calls + .mcp_attribution_snapshot() +} + /// Test-only provider that supplies no user instructions. #[derive(Debug, Default)] pub struct EmptyUserInstructionsProvider; diff --git a/codex-rs/core/src/tools/executed_tool_calls.rs b/codex-rs/core/src/tools/executed_tool_calls.rs index 43e8aa534e..aec694bd24 100644 --- a/codex-rs/core/src/tools/executed_tool_calls.rs +++ b/codex-rs/core/src/tools/executed_tool_calls.rs @@ -4,6 +4,7 @@ #[cfg(test)] #[path = "executed_tool_calls_direct_tests.rs"] mod direct_tests; +mod mcp_attribution; mod request_metadata; use std::collections::HashMap; @@ -19,6 +20,8 @@ use codex_code_mode::CellId; use codex_features::Feature; use codex_features::Features; use codex_history::InitialHistory; +use codex_protocol::mcp::McpAttribution; +use codex_protocol::mcp::McpAttributionSource; use codex_protocol::models::ExecutedToolCall; use codex_protocol::models::ExecutedToolCallArguments; use codex_protocol::models::ResponseItem; @@ -58,6 +61,7 @@ pub(crate) struct ExecutedToolCalls { state: Arc>>, retained_direct_metadata_bytes: Arc, pending_direct_calls: Arc, + mcp_attribution: mcp_attribution::McpAttributionRecorder, } // The tool future owns this reservation, so completion or cancellation releases it. @@ -220,9 +224,26 @@ impl ExecutedToolCalls { }))), retained_direct_metadata_bytes: Arc::new(AtomicUsize::new(0)), pending_direct_calls: Arc::new(AtomicUsize::new(0)), + mcp_attribution: mcp_attribution::McpAttributionRecorder::new(history), } } + pub(crate) fn record_mcp_source(&self, source: McpAttributionSource) { + self.mcp_attribution.record(source); + } + + pub(crate) fn mcp_attribution_snapshot(&self) -> McpAttribution { + self.mcp_attribution.snapshot() + } + + pub(crate) fn mcp_attribution_checkpoint(&self, force: bool) -> Option<(McpAttribution, u64)> { + self.mcp_attribution.checkpoint(force) + } + + pub(crate) fn mark_mcp_attribution_persisted(&self, revision: u64) { + self.mcp_attribution.mark_persisted(revision); + } + fn lock_state(&self) -> MutexGuard<'_, Option> { self.state .lock() diff --git a/codex-rs/core/src/tools/executed_tool_calls/mcp_attribution.rs b/codex-rs/core/src/tools/executed_tool_calls/mcp_attribution.rs new file mode 100644 index 0000000000..fce3f8f36d --- /dev/null +++ b/codex-rs/core/src/tools/executed_tool_calls/mcp_attribution.rs @@ -0,0 +1,169 @@ +//! Cumulative MCP result attribution, independent of the best-effort call recorder. + +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::MutexGuard; + +use codex_history::InitialHistory; +use codex_history::RolloutItem; +use codex_protocol::mcp::McpAttribution; +use codex_protocol::mcp::McpAttributionSource; +use codex_protocol::mcp::McpAttributionStatus; + +#[derive(Clone)] +pub(super) struct McpAttributionRecorder(Arc>); + +struct State { + attribution: McpAttribution, + revision: u64, + persisted_revision: u64, +} + +#[derive(PartialEq, Eq)] +struct SourceIdentity<'a> { + connector_id: Option<&'a str>, + plugin_id: Option<&'a str>, + server_name: &'a str, + tool_name: &'a str, +} + +impl<'a> From<&'a McpAttributionSource> for SourceIdentity<'a> { + fn from(source: &'a McpAttributionSource) -> Self { + Self { + connector_id: source.connector_id.as_deref(), + plugin_id: source.plugin_id.as_deref(), + server_name: &source.server_name, + tool_name: &source.tool_name, + } + } +} + +impl Default for McpAttributionRecorder { + fn default() -> Self { + Self::new(&InitialHistory::New) + } +} + +impl McpAttributionRecorder { + pub(super) fn new(history: &InitialHistory) -> Self { + let mut state = State { + attribution: McpAttribution::default(), + // Persist an initial checkpoint even when no MCP result has been recorded. + revision: 1, + persisted_revision: 0, + }; + let mut found_checkpoint = matches!(history, InitialHistory::New | InitialHistory::Cleared); + for item in history.get_rollout_items() { + match item { + RolloutItem::ResponseItem(envelope) => { + if let Some(checkpoint) = envelope + .metadata + .as_ref() + .and_then(|metadata| metadata.mcp_attribution.as_ref()) + { + found_checkpoint = true; + state.merge_checkpoint(checkpoint); + } + } + RolloutItem::Compacted(compacted) => { + for checkpoint in compacted + .replacement_history + .iter() + .flatten() + .filter_map(|envelope| envelope.metadata.as_ref()) + .filter_map(|metadata| metadata.mcp_attribution.as_ref()) + { + found_checkpoint = true; + state.merge_checkpoint(checkpoint); + } + } + _ => {} + } + } + if !found_checkpoint { + // Pre-attribution history cannot establish that earlier context was MCP-free. + state.mark_error(); + } + Self(Arc::new(Mutex::new(state))) + } + + fn lock(&self) -> MutexGuard<'_, State> { + self.0.lock().unwrap_or_else(|poisoned| { + let mut state = poisoned.into_inner(); + state.mark_error(); + state + }) + } + + pub(super) fn record(&self, source: McpAttributionSource) { + self.lock().record(source, /*restoring*/ false); + } + + pub(super) fn snapshot(&self) -> McpAttribution { + self.lock().attribution.clone() + } + + pub(super) fn checkpoint(&self, force: bool) -> Option<(McpAttribution, u64)> { + let state = self.lock(); + (force || state.revision != state.persisted_revision) + .then(|| (state.attribution.clone(), state.revision)) + } + + pub(super) fn mark_persisted(&self, revision: u64) { + let mut state = self.lock(); + state.persisted_revision = state.persisted_revision.max(revision.min(state.revision)); + } +} + +impl State { + fn mark_error(&mut self) { + if self.attribution.status != McpAttributionStatus::AttributionError { + self.attribution.status = McpAttributionStatus::AttributionError; + self.revision += 1; + } + } + + fn merge_checkpoint(&mut self, checkpoint: &McpAttribution) { + if checkpoint.status == McpAttributionStatus::AttributionError + || (checkpoint.status == McpAttributionStatus::None && !checkpoint.sources.is_empty()) + || (checkpoint.status == McpAttributionStatus::Complete + && checkpoint.sources.is_empty()) + { + self.mark_error(); + } + for source in &checkpoint.sources { + self.record(source.clone(), /*restoring*/ true); + } + } + + fn record(&mut self, source: McpAttributionSource, restoring: bool) { + let identity = SourceIdentity::from(&source); + if let Some(existing) = self + .attribution + .sources + .iter() + .find(|existing| SourceIdentity::from(*existing) == identity) + { + if restoring && existing.first_turn_id != source.first_turn_id { + self.mark_error(); + } + return; + } + if source.server_name.is_empty() + || source.tool_name.is_empty() + || source.first_turn_id.is_empty() + { + self.mark_error(); + return; + } + if self.attribution.status == McpAttributionStatus::None { + self.attribution.status = McpAttributionStatus::Complete; + } + self.attribution.sources.push(source); + self.revision += 1; + } +} + +#[cfg(test)] +#[path = "mcp_attribution_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/tools/executed_tool_calls/mcp_attribution_tests.rs b/codex-rs/core/src/tools/executed_tool_calls/mcp_attribution_tests.rs new file mode 100644 index 0000000000..87033720a2 --- /dev/null +++ b/codex-rs/core/src/tools/executed_tool_calls/mcp_attribution_tests.rs @@ -0,0 +1,129 @@ +//! Tests for cumulative MCP attribution restoration and checkpoint acknowledgements. + +use super::*; +use codex_history::CodexHarnessMetadata; +use codex_history::CompactedItem; +use codex_history::ResponseItemEnvelope; +use codex_protocol::models::ResponseItem; +use pretty_assertions::assert_eq; + +fn source(tool_name: &str, first_turn_id: &str) -> McpAttributionSource { + McpAttributionSource { + connector_id: None, + plugin_id: None, + server_name: "example".to_string(), + tool_name: tool_name.to_string(), + first_turn_id: first_turn_id.to_string(), + } +} + +fn envelope(attribution: Option) -> ResponseItemEnvelope { + ResponseItemEnvelope { + item: ResponseItem::Message { + id: None, + role: "user".to_string(), + content: Vec::new(), + phase: None, + internal_chat_message_metadata_passthrough: None, + }, + metadata: attribution.map(|mcp_attribution| CodexHarnessMetadata { + mcp_attribution: Some(mcp_attribution), + ..Default::default() + }), + } +} + +#[test] +fn records_the_first_turn_for_each_unique_source() { + let recorder = McpAttributionRecorder::default(); + recorder.record(source("search", "turn_1")); + recorder.record(source("search", "turn_2")); + recorder.record(source("fetch", "turn_2")); + + assert_eq!( + recorder.snapshot(), + McpAttribution { + status: McpAttributionStatus::Complete, + sources: vec![source("search", "turn_1"), source("fetch", "turn_2")], + } + ); +} + +#[test] +fn restores_cumulative_item_and_compaction_checkpoints() { + let initial = McpAttribution { + status: McpAttributionStatus::Complete, + sources: vec![source("search", "turn_1")], + }; + let cumulative = McpAttribution { + status: McpAttributionStatus::Complete, + sources: vec![source("search", "turn_1"), source("fetch", "turn_2")], + }; + let history = InitialHistory::Forked(vec![ + RolloutItem::ResponseItem(envelope(Some(initial))), + RolloutItem::Compacted(CompactedItem { + message: "summary".to_string(), + replacement_history: Some(vec![envelope(Some(cumulative.clone()))]), + guardian_history: None, + retained_context: None, + mcp_resource_origins: None, + window_number: None, + first_window_id: None, + previous_window_id: None, + window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, + }), + ]); + + assert_eq!(McpAttributionRecorder::new(&history).snapshot(), cumulative); +} + +#[test] +fn legacy_or_conflicting_history_is_not_complete() { + let legacy = InitialHistory::Forked(vec![RolloutItem::ResponseItem(envelope( + /*attribution*/ None, + ))]); + assert_eq!( + McpAttributionRecorder::new(&legacy).snapshot(), + McpAttribution { + status: McpAttributionStatus::AttributionError, + sources: Vec::new(), + } + ); + + let history = InitialHistory::Forked( + ["turn_1", "turn_2"] + .map(|turn_id| { + RolloutItem::ResponseItem(envelope(Some(McpAttribution { + status: McpAttributionStatus::Complete, + sources: vec![source("search", turn_id)], + }))) + }) + .to_vec(), + ); + assert_eq!( + McpAttributionRecorder::new(&history).snapshot(), + McpAttribution { + status: McpAttributionStatus::AttributionError, + sources: vec![source("search", "turn_1")], + } + ); +} + +#[test] +fn acknowledging_an_older_checkpoint_does_not_clear_newer_changes() { + let recorder = McpAttributionRecorder::default(); + let (_, initial_revision) = recorder + .checkpoint(/*force*/ false) + .expect("initial checkpoint"); + recorder.record(source("search", "turn_1")); + recorder.mark_persisted(initial_revision); + + let (_, latest_revision) = recorder + .checkpoint(/*force*/ false) + .expect("dirty checkpoint"); + recorder.mark_persisted(latest_revision); + assert_eq!(recorder.checkpoint(/*force*/ false), None); + assert!(recorder.checkpoint(/*force*/ true).is_some()); +} diff --git a/codex-rs/core/src/turn_metadata_tests.rs b/codex-rs/core/src/turn_metadata_tests.rs index c0316896be..a4ea4b3597 100644 --- a/codex-rs/core/src/turn_metadata_tests.rs +++ b/codex-rs/core/src/turn_metadata_tests.rs @@ -9,6 +9,7 @@ use crate::responses_metadata::CompactionTurnMetadata; use crate::responses_metadata::FORKED_FROM_ORDINAL_EXCLUSIVE_KEY; use crate::responses_metadata::INSTALLATION_ID_KEY; use crate::responses_metadata::LEGACY_CODE_MODE_TOOL_NAMES_KEY; +use crate::responses_metadata::MCP_ATTRIBUTION_CLIENT_METADATA_KEY; use crate::responses_metadata::NODE_REPL_AUTO_REVIEW_REQUIRED_KEY; use crate::responses_metadata::NODE_REPL_DISABLED_KEY; use crate::responses_metadata::PARENT_TURN_ID_KEY; @@ -181,7 +182,7 @@ async fn detached_memory_responses_metadata_starts_an_independent_root_turn() { let turn_id = parsed["turn_id"].as_str().expect("memory turn ID"); uuid::Uuid::parse_str(turn_id).expect("memory turn ID is a UUID"); assert_eq!(parsed[ROOT_TURN_ID_KEY], parsed["turn_id"]); - let client_metadata = metadata.client_metadata(); + let client_metadata = metadata.client_metadata(/*include_internal*/ true); assert_eq!( client_metadata.get("turn_id").map(String::as_str), Some(turn_id) @@ -1100,6 +1101,7 @@ fn responses_api_metadata_rejects_reserved_keys() { TURN_TRIGGER_KEY, WINDOW_ID_KEY, CONTEXT_WINDOW_ID_KEY, + MCP_ATTRIBUTION_CLIENT_METADATA_KEY, ] { assert_eq!( validate_extra_metadata( diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index d443287687..a2bebfe6d0 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -5939,6 +5939,41 @@ contentLength=0" Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn code_mode_records_mcp_completion_even_when_result_is_not_printed() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let (test, second_mock) = run_code_mode_turn_with_rmcp( + &server, + "call the rmcp echo tool without printing its result", + r#"await tools.mcp__rmcp__echo({ message: "discarded" }); text("done");"#, + ) + .await?; + let request = second_mock.single_request(); + assert_eq!( + custom_tool_output_body_and_success(&request, "call-1").0, + "done" + ); + + let first_turn_id = request.body_json()["client_metadata"]["turn_id"].clone(); + assert!(first_turn_id.is_string()); + assert_eq!( + serde_json::to_value(codex_core::test_support::mcp_attribution_snapshot( + &test.codex + ))?, + serde_json::json!({ + "status": "complete", + "sources": [{ + "server_name": "rmcp", + "tool_name": "echo", + "first_turn_id": first_turn_id, + }], + }) + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[test_case("node_repl"; "node_repl")] #[test_case("cua_repl"; "cua_repl")] diff --git a/codex-rs/core/tests/suite/mcp_turn_metadata.rs b/codex-rs/core/tests/suite/mcp_turn_metadata.rs index 7e01e1381b..75b3ba1dfc 100644 --- a/codex-rs/core/tests/suite/mcp_turn_metadata.rs +++ b/codex-rs/core/tests/suite/mcp_turn_metadata.rs @@ -1493,6 +1493,28 @@ async fn apps_default_writes_prompts_for_writes_but_not_reads() -> Result<()> { recorded_apps_tool_call_by_call_id(&server, read_call_id).await; recorded_apps_tool_call_by_call_id(&server, write_call_id).await; + let first_turn_id = responses.requests()[2].body_json()["client_metadata"]["turn_id"].clone(); + assert!(first_turn_id.is_string()); + assert_eq!( + serde_json::to_value(codex_core::test_support::mcp_attribution_snapshot( + &test.codex + ))?, + json!({ + "status": "complete", + "sources": [{ + "connector_id": "calendar", + "server_name": "codex_apps", + "tool_name": "calendar_list_events", + "first_turn_id": first_turn_id, + }, { + "connector_id": "calendar", + "server_name": "codex_apps", + "tool_name": "calendar_create_event", + "first_turn_id": first_turn_id, + }], + }) + ); + test.codex.ensure_rollout_materialized().await; test.codex.flush_rollout().await?; let rollout_path = test.codex.rollout_path().expect("rollout path"); diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index a1fd36fe2f..bad5c1135b 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -670,7 +670,8 @@ async fn text_only_mcp_content_uses_content_items() -> anyhow::Result<()> { }) .await; - let output = final_mock.single_request().function_call_output(call_id); + let request = final_mock.single_request(); + let output = request.function_call_output(call_id); let header = output["output"][0]["text"] .as_str() .expect("first content item should contain the wall-time header"); @@ -689,6 +690,22 @@ async fn text_only_mcp_content_uses_content_items() -> anyhow::Result<()> { ]) ); + let first_turn_id = request.body_json()["client_metadata"]["turn_id"].clone(); + assert!(first_turn_id.is_string()); + assert_eq!( + serde_json::to_value(codex_core::test_support::mcp_attribution_snapshot( + &fixture.codex + ))?, + json!({ + "status": "complete", + "sources": [{ + "server_name": "rmcp", + "tool_name": "image_scenario", + "first_turn_id": first_turn_id, + }], + }) + ); + server.verify().await; Ok(()) } diff --git a/codex-rs/history/src/lib.rs b/codex-rs/history/src/lib.rs index 278d8dc839..4b611539cf 100644 --- a/codex-rs/history/src/lib.rs +++ b/codex-rs/history/src/lib.rs @@ -12,6 +12,8 @@ use std::sync::Arc; use codex_protocol::ThreadId; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::dynamic_tools::DynamicToolSpec; +use codex_protocol::mcp::McpAttribution; +use codex_protocol::mcp::McpAttributionStatus; use codex_protocol::mcp::McpResourceOriginCheckpoint; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; @@ -83,11 +85,34 @@ pub struct CodexHarnessMetadata { #[serde(default, skip_serializing_if = "std::ops::Not::not")] pub inherited_user_message: bool, + /// Cumulative MCP tools/call attribution checkpoint, never model-visible. + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "deserialize_mcp_attribution_checkpoint" + )] + pub mcp_attribution: Option, + /// Sender context captured by the host when this task message was accepted. #[serde(default, skip_serializing_if = "Option::is_none")] pub sender_user_messages: Option>, } +fn deserialize_mcp_attribution_checkpoint<'de, D>( + deserializer: D, +) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + let value = serde_json::Value::deserialize(deserializer)?; + Ok(Some(serde_json::from_value(value).unwrap_or_else(|_| { + McpAttribution { + status: McpAttributionStatus::AttributionError, + sources: Vec::new(), + } + }))) +} + impl ResponseItemEnvelope { /// Wraps a raw Responses API item for persisted history. pub fn new(item: ResponseItem) -> Self { diff --git a/codex-rs/history/src/tests.rs b/codex-rs/history/src/tests.rs index b04d2448db..d4fcbfe2f8 100644 --- a/codex-rs/history/src/tests.rs +++ b/codex-rs/history/src/tests.rs @@ -161,6 +161,32 @@ fn response_item_envelope_accessors_preserve_item() { assert_eq!(envelope.into_item(), replacement_item); } +#[test] +fn mcp_checkpoint_handles_missing_or_unknown_identity() -> Result<()> { + let metadata: CodexHarnessMetadata = serde_json::from_value(json!({}))?; + assert_eq!(metadata.mcp_attribution, None); + + let restored: CodexHarnessMetadata = serde_json::from_value(json!({ + "mcp_attribution": { + "status": "complete", + "sources": [{ + "server_name": "example", + "tool_name": "search", + "first_turn_id": "turn_1", + "future_identity": "unknown to this reader" + }] + } + }))?; + assert_eq!( + restored.mcp_attribution, + Some(McpAttribution { + status: McpAttributionStatus::AttributionError, + sources: Vec::new(), + }) + ); + Ok(()) +} + #[test] /// Keeps legacy response-item rollout lines readable and byte-shape compatible. fn response_item_rollout_line_preserves_shape() -> Result<()> { diff --git a/codex-rs/protocol/src/mcp.rs b/codex-rs/protocol/src/mcp.rs index 317174a568..c898c2af4e 100644 --- a/codex-rs/protocol/src/mcp.rs +++ b/codex-rs/protocol/src/mcp.rs @@ -52,6 +52,42 @@ pub fn is_node_repl_backed_tool(name: &str, namespace: Option<&str>) -> bool { .is_some_and(|(server, _)| is_node_repl_backed_server(server)) } +/// Producer-reported completeness of the thread's MCP tools/call attribution. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum McpAttributionStatus { + #[default] + None, + Complete, + AttributionError, +} + +/// One source and the first runtime turn in which Codex recorded it. +/// +/// Server and tool names are retained when no stable identifier is available; +/// they are not guaranteed to be globally unique or independently authenticated. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct McpAttributionSource { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub connector_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub plugin_id: Option, + pub server_name: String, + pub tool_name: String, + pub first_turn_id: String, +} + +/// Cumulative, model-invisible attribution reported on Responses API requests. +/// This is provenance data, not a training-eligibility decision or attestation. +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct McpAttribution { + pub status: McpAttributionStatus, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub sources: Vec, +} + /// Bounded app-resource provenance retained across a compaction checkpoint. #[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] pub struct McpResourceOriginCheckpoint { diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index c74d53ee98..5f2eb7f84c 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -143,6 +143,7 @@ pub use persistence_metrics::RolloutPersistenceTelemetry; pub use persistence_metrics::measure_and_filter_rollout_items; pub use policy::is_persisted_rollout_item; pub use policy::persisted_rollout_items; +pub use policy::should_persist_response_item; pub use policy::should_persist_response_item_for_memories; pub use recorder::RolloutRecorder; pub use recorder::RolloutRecorderParams; diff --git a/codex-rs/utils/string/src/json.rs b/codex-rs/utils/string/src/json.rs index fd5e7d65bc..da69ed6318 100644 --- a/codex-rs/utils/string/src/json.rs +++ b/codex-rs/utils/string/src/json.rs @@ -1,5 +1,4 @@ -//! JSON serialization helpers for output that must remain parseable as JSON -//! while staying safe for ASCII-only transports. +//! JSON serialization helpers for ASCII-only transports and bounded output. use std::io; @@ -54,6 +53,43 @@ where .map_err(|err| serde_json::Error::io(io::Error::new(io::ErrorKind::InvalidData, err))) } +/// Serialize UTF-8 JSON while retaining at most `max_bytes` of output. +pub fn to_json_string_bounded(value: &T, max_bytes: usize) -> serde_json::Result +where + T: Serialize + ?Sized, +{ + struct BoundedBuffer { + bytes: Vec, + max_bytes: usize, + } + + impl io::Write for BoundedBuffer { + fn write(&mut self, bytes: &[u8]) -> io::Result { + if bytes.len() > self.max_bytes.saturating_sub(self.bytes.len()) { + return Err(io::Error::new( + io::ErrorKind::WriteZero, + "JSON byte limit exceeded", + )); + } + self.bytes.extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + let mut buffer = BoundedBuffer { + bytes: Vec::new(), + max_bytes, + }; + let mut serializer = serde_json::Serializer::new(&mut buffer); + value.serialize(&mut serializer)?; + String::from_utf8(buffer.bytes) + .map_err(|err| serde_json::Error::io(io::Error::new(io::ErrorKind::InvalidData, err))) +} + #[cfg(test)] mod tests { use std::collections::BTreeMap; @@ -65,6 +101,7 @@ mod tests { use serde_json::json; use super::to_ascii_json_string; + use super::to_json_string_bounded; #[test] fn to_ascii_json_string_escapes_non_ascii_strings() { @@ -119,4 +156,14 @@ mod tests { let parsed: Value = serde_json::from_str(&serialized).expect("serialized json"); assert_eq!(parsed, expected_value); } + + #[test] + fn bounded_json_counts_utf8_output_bytes() { + assert_eq!( + to_json_string_bounded(&"é", /*max_bytes*/ 4).unwrap(), + r#""é""# + ); + assert!(to_json_string_bounded(&"é", /*max_bytes*/ 3).is_err()); + assert!(to_json_string_bounded(&"a".repeat(100_000), 16 * 1024).is_err()); + } } diff --git a/codex-rs/utils/string/src/lib.rs b/codex-rs/utils/string/src/lib.rs index 2d081f79e2..1618739b08 100644 --- a/codex-rs/utils/string/src/lib.rs +++ b/codex-rs/utils/string/src/lib.rs @@ -2,6 +2,7 @@ mod json; mod truncate; pub use json::to_ascii_json_string; +pub use json::to_json_string_bounded; pub use truncate::approx_bytes_for_tokens; pub use truncate::approx_token_count; pub use truncate::approx_tokens_from_byte_count;