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

Track cumulative MCP attribution across requests and thread history (#47081)

Browse Source 
## What changed

- Record each MCP source and its first turn, including connector and plugin identifiers when available. Cover direct and Code Mode calls, including returned errors and results that are not printed.
- Persist cumulative attribution checkpoints and restore them across compaction, resume, and forks, including forks limited to recent turns.
- Attach fresh attribution to Responses API `client_metadata` under `mcp_attribution`, outside model-visible content. Restrict transmission to allowed OpenAI HTTPS destinations for both HTTP and WebSocket requests.
- Report `attribution_error` for legacy history without checkpoints, malformed or conflicting attribution, and serialized attribution exceeding 16 KiB.

## Testing

Add coverage for source deduplication, checkpoint persistence and restoration, recent-turn forks, direct and Code Mode calls, destination filtering, and the serialized size limit.

GitOrigin-RevId: fd285e152ca9e24bfa000debb20013709d67a86e
This commit is contained in:
peilin-openai
2026-09-21 18:19:07 +00:00
committed by copyberry
parent 7cec3e110c
commit ac98537678
28 changed files with 1119 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst
Generated
BIN
View File
Binary file not shown.
codex-rs/core/src/agent/control_tests.rs
+185 -4
View File
@@ -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<CodexThread>, 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
codex-rs/core/src/client.rs
+30 -33
View File
@@ -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<String, String> {
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<String>,
responses_metadata: &CodexResponsesMetadata,
include_internal: bool,
) -> Result<ResponsesApiRequest> {
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());
}
codex-rs/core/src/client_tests.rs
+94 -12
View File
@@ -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
codex-rs/core/src/compact.rs
+5 -7
View File
@@ -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(),
codex-rs/core/src/guardian/input_budget.rs
+1
View File
@@ -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));
codex-rs/core/src/guardian/request_budget.rs
+1
View File
@@ -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)
codex-rs/core/src/mcp_tool_call.rs
+15
View File
@@ -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,
codex-rs/core/src/responses_metadata.rs
+15 -1
View File
@@ -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<bool>,
/// Cumulative MCP attribution for this logical request; body-only and model-invisible.
pub(crate) mcp_attribution: Option<McpAttribution>,
pub(crate) extra: BTreeMap<String, String>,
}
@@ -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<String, String> {
pub(crate) fn client_metadata(&self, include_internal: bool) -> HashMap<String, String> {
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
}
codex-rs/core/src/session/inject_tests.rs
+8 -1
View File
@@ -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();
codex-rs/core/src/session/mod.rs
+54 -3
View File
@@ -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<RolloutItem> =
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 {
codex-rs/core/src/session/session.rs
+1
View File
@@ -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
codex-rs/core/src/session/tests.rs
+158 -1
View File
@@ -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::<Vec<_>>();
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] {
codex-rs/core/src/session/turn.rs
+5 -7
View File
@@ -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<codex_extension_api::ExtensionData>,
turn_diff_tracker: SharedTurnDiffTracker,
client_session: &mut ModelClientSession,
responses_metadata: &CodexResponsesMetadata,
input: Vec<ResponseItem>,
cancellation_token: CancellationToken,
) -> CodexResult<(SamplingRequestResult, Vec<ResponseItem>)> {
@@ -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(),
codex-rs/core/src/test_support.rs
+12
View File
@@ -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;
codex-rs/core/src/tools/executed_tool_calls.rs
+21
View File
@@ -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<Mutex<Option<ExecutedToolCallRecorderState>>>,
retained_direct_metadata_bytes: Arc<AtomicUsize>,
pending_direct_calls: Arc<AtomicUsize>,
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<ExecutedToolCallRecorderState>> {
self.state
.lock()
codex-rs/core/src/tools/executed_tool_calls/mcp_attribution.rs
+169
View File
@@ -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<Mutex<State>>);
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;
codex-rs/core/src/tools/executed_tool_calls/mcp_attribution_tests.rs
+129
View File
@@ -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<McpAttribution>) -> 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());
}
codex-rs/core/src/turn_metadata_tests.rs
+3 -1
View File
@@ -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(
codex-rs/core/tests/suite/code_mode.rs
+35
View File
@@ -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")]
codex-rs/core/tests/suite/mcp_turn_metadata.rs
+22
View File
@@ -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");
codex-rs/core/tests/suite/rmcp_client.rs
+18 -1
View File
@@ -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(())
}
codex-rs/history/src/lib.rs
+25
View File
@@ -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<McpAttribution>,
/// 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<Box<SenderUserMessages>>,
}
fn deserialize_mcp_attribution_checkpoint<'de, D>(
deserializer: D,
) -> Result<Option<McpAttribution>, 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 {
codex-rs/history/src/tests.rs
+26
View File
@@ -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<()> {
codex-rs/protocol/src/mcp.rs
+36
View File
@@ -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<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub plugin_id: Option<String>,
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<McpAttributionSource>,
}
/// Bounded app-resource provenance retained across a compaction checkpoint.
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
pub struct McpResourceOriginCheckpoint {
codex-rs/rollout/src/lib.rs
+1
View File
@@ -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;
codex-rs/utils/string/src/json.rs
+49 -2
View File
@@ -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<T>(value: &T, max_bytes: usize) -> serde_json::Result<String>
where
T: Serialize + ?Sized,
{
struct BoundedBuffer {
bytes: Vec<u8>,
max_bytes: usize,
}
impl io::Write for BoundedBuffer {
fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
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());
}
}
codex-rs/utils/string/src/lib.rs
+1
View File
@@ -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;