From 12ed76b09bf311d8d9a6f25be2f9a8eb37daf879 Mon Sep 17 00:00:00 2001 From: ningyi-oai Date: Wed, 2 Sep 2026 00:37:11 +0000 Subject: [PATCH] Record result sources in app tool analytics (#42164) ## What changed - Add per-tool `analytics_result_source` requirements, with support for the `detailed_message_search_v1` format and normal configuration precedence. - Attach host-generated source IDs from accepted app tool results to recorded executed tool calls for both direct and Code Mode invocations. - Bound and deduplicate source metadata, reject caller-supplied values, and shed optional source data before dropping tool-call or completion metadata. - Preserve source updates across waits and retries while preventing stale or compacted records from overwriting accepted metadata. ## Testing Add coverage for requirement parsing and merging, trusted source snapshots, metadata budgets, retries, waits, and compaction. GitOrigin-RevId: 7d5e394219eb33e0ab7b6db3f039138a78569d17 --- codex-rs/config/src/config_requirements.rs | 101 +++++++++- codex-rs/config/src/lib.rs | 2 + .../connectors/src/app_tool_policy_tests.rs | 2 + .../core/src/tools/executed_tool_calls.rs | 94 +++++++++- .../src/tools/executed_tool_calls_tests.rs | 172 +++++++++++++++++- codex-rs/core/src/tools/handlers/mcp.rs | 6 + codex-rs/core/src/tools/registry.rs | 4 + codex-rs/protocol/src/models.rs | 3 + .../src/models/executed_tool_calls.rs | 92 ++++++++++ .../src/models/executed_tool_calls_tests.rs | 116 ++++++++++++ codex-rs/tools/src/tool_output.rs | 9 + 11 files changed, 584 insertions(+), 17 deletions(-) diff --git a/codex-rs/config/src/config_requirements.rs b/codex-rs/config/src/config_requirements.rs index 0b6785b594..8fbeed6b48 100644 --- a/codex-rs/config/src/config_requirements.rs +++ b/codex-rs/config/src/config_requirements.rs @@ -14,8 +14,10 @@ use serde::de::value::Error as ValueDeserializerError; use serde::de::value::StrDeserializer; use std::collections::BTreeMap; use std::collections::BTreeSet; +use std::convert::Infallible; use std::fmt; use std::path::PathBuf; +use std::str::FromStr; use wildmatch::WildMatchPattern; use super::requirements_exec_policy::RequirementsExecPolicyToml; @@ -822,11 +824,53 @@ impl FeatureRequirementsToml { #[derive(Deserialize, Debug, Clone, Default, PartialEq, Eq)] pub struct AppToolRequirementToml { pub approval_mode: Option, + /// Opt-in analytics extraction for this exact tool, not a tool argument. + /// The highest-priority rule wins as a whole, including unsupported formats. + pub analytics_result_source: Option, +} + +#[derive(Deserialize, Debug, Clone, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct AppToolResultSourceRequirementToml { + /// Result format to parse; currently only `detailed_message_search_v1` is supported. + pub format: AppToolResultSourceFormat, + /// Source kind emitted alongside each extracted ID. + #[serde(rename = "type")] + pub source_type: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AppToolResultSourceFormat { + DetailedMessageSearchV1, + /// Keep unknown formats so higher-priority rules still override lower ones. + Unknown(String), +} + +impl FromStr for AppToolResultSourceFormat { + type Err = Infallible; + + fn from_str(value: &str) -> Result { + Ok(match value { + "detailed_message_search_v1" => Self::DetailedMessageSearchV1, + _ => Self::Unknown(value.to_string()), + }) + } +} + +impl<'de> Deserialize<'de> for AppToolResultSourceFormat { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + String::deserialize(deserializer)? + .parse() + .map_err(D::Error::custom) + } } impl AppToolRequirementToml { pub fn is_empty(&self) -> bool { - self.approval_mode.is_none() + self.approval_mode.is_none() && self.analytics_result_source.is_none() } } @@ -897,6 +941,9 @@ pub(crate) fn merge_app_requirements_descending( if base_tool.approval_mode.is_none() { base_tool.approval_mode = incoming_tool.approval_mode; } + if base_tool.analytics_result_source.is_none() { + base_tool.analytics_result_source = incoming_tool.analytics_result_source; + } } } } @@ -3093,6 +3140,7 @@ allowed_approvals_reviewers = ["user"] "calendar/list_events".to_string(), AppToolRequirementToml { approval_mode: Some(AppToolApproval::Approve), + analytics_result_source: None, }, )]), }), @@ -3103,6 +3151,56 @@ allowed_approvals_reviewers = ["user"] Ok(()) } + #[test] + fn app_tool_result_source_requirements_parse_and_merge() -> Result<()> { + let rule = r#" + [apps.connector_123123.tools."messages/search"] + analytics_result_source = { format = "detailed_message_search_v1", type = "message_room" } + "#; + let requirements: ConfigRequirementsToml = from_str(rule)?; + assert!(!requirements.is_empty()); + let source = requirements.apps.expect("apps should be present"); + + for higher_rule in [ + None, + Some(("unsupported", "other_resource")), + Some(("detailed_message_search_v1", "other_resource")), + ] { + let mut merged = source.clone(); + merged + .apps + .get_mut("connector_123123") + .expect("app should be present") + .tools + .as_mut() + .expect("tools should be present") + .tools + .get_mut("messages/search") + .expect("tool should be present") + .analytics_result_source = higher_rule.map(|(format, source_type)| { + from_str(&format!("format = {format:?}\ntype = {source_type:?}")) + .expect("complete source rule should parse, including unknown formats") + }); + let expected = if higher_rule.is_none() { + source.clone() + } else { + merged.clone() + }; + + merge_app_requirements_descending(&mut merged, source.clone()); + + assert_eq!(merged, expected); + } + + for incomplete_rule in [ + r#"format = "detailed_message_search_v1""#, + r#"type = "message_room""#, + ] { + assert!(from_str::(incomplete_rule).is_err()); + } + Ok(()) + } + fn apps_requirements(entries: &[(&str, Option)]) -> AppsRequirementsToml { AppsRequirementsToml { apps: entries @@ -3135,6 +3233,7 @@ allowed_approvals_reviewers = ["user"] tool_name.to_string(), AppToolRequirementToml { approval_mode: Some(approval_mode), + analytics_result_source: None, }, )]), }), diff --git a/codex-rs/config/src/lib.rs b/codex-rs/config/src/lib.rs index fc5e05a687..0b41b2e04e 100644 --- a/codex-rs/config/src/lib.rs +++ b/codex-rs/config/src/lib.rs @@ -83,6 +83,8 @@ pub use config_layer_source::ConfigLayerSource; pub use config_layer_source::format_config_layer_source; pub use config_requirements::AppRequirementToml; pub use config_requirements::AppToolRequirementToml; +pub use config_requirements::AppToolResultSourceFormat; +pub use config_requirements::AppToolResultSourceRequirementToml; pub use config_requirements::AppToolsRequirementsToml; pub use config_requirements::AppsRequirementsToml; pub use config_requirements::AutoReviewRequirementsToml; diff --git a/codex-rs/connectors/src/app_tool_policy_tests.rs b/codex-rs/connectors/src/app_tool_policy_tests.rs index 8d2222fdaa..e255de4972 100644 --- a/codex-rs/connectors/src/app_tool_policy_tests.rs +++ b/codex-rs/connectors/src/app_tool_policy_tests.rs @@ -53,6 +53,7 @@ fn evaluator_reuses_one_snapshot_across_tools() { "events/create".to_string(), AppToolRequirementToml { approval_mode: Some(AppToolApproval::Approve), + ..Default::default() }, )]), }), @@ -937,6 +938,7 @@ fn app_tool_requirements( tool_name.to_string(), AppToolRequirementToml { approval_mode: Some(approval_mode), + ..Default::default() }, )]), }), diff --git a/codex-rs/core/src/tools/executed_tool_calls.rs b/codex-rs/core/src/tools/executed_tool_calls.rs index cda42c5a13..ec4843b5aa 100644 --- a/codex-rs/core/src/tools/executed_tool_calls.rs +++ b/codex-rs/core/src/tools/executed_tool_calls.rs @@ -8,6 +8,7 @@ use codex_protocol::models::ResponseItem; use codex_protocol::models::bound_executed_tool_calls_for_prompt_prioritizing_recent; use codex_protocol::models::executed_tool_call_metadata_bytes; use codex_protocol::openai_models::ToolMode; +use indexmap::IndexMap; use serde_json::Value as JsonValue; use crate::tools::context::ToolCallSource; @@ -42,6 +43,9 @@ struct RetainedToolCalls { complete: bool, cell_id: Option, runtime_cell_id: Option, + // Invocation IDs stay local and are retained only with their original output's calls. + call_index_by_id: HashMap, + sources_updated: bool, } #[derive(Default, PartialEq, Eq)] @@ -56,7 +60,7 @@ enum CellCompletion { #[derive(Default)] struct RecordedCell { - pending_calls: Vec, + pending_calls: IndexMap, pending_full_argument_bytes: usize, completion: CellCompletion, originating_call_id: Option, @@ -169,6 +173,7 @@ impl ExecutedToolCallRecorder { ToolCallSource::CodeMode { cell_id, .. } => { self.record_nested_tool_call( CellId::new(cell_id.clone()), + call.call_id.clone(), recorded_call, original_bytes, ); @@ -179,6 +184,7 @@ impl ExecutedToolCallRecorder { fn record_nested_tool_call( &self, cell_id: CellId, + call_id: String, call: ExecutedToolCall, original_bytes: usize, ) { @@ -222,10 +228,49 @@ impl ExecutedToolCallRecorder { } else { CellCompletion::Incomplete }; - cell.pending_calls.push(call); + cell.pending_calls.insert(call_id, call); state.pending_nested_calls += 1; } + pub(crate) fn record_tool_result_sources( + &self, + source: &ToolCallSource, + call_id: &str, + result_sources: codex_protocol::models::ToolResultSources, + ) -> bool { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let call = match source { + ToolCallSource::Direct | ToolCallSource::DirectPlaintextMessage => { + state.direct_calls.get_mut(call_id) + } + ToolCallSource::CodeMode { cell_id, .. } => state + .cells + .get_mut(&CellId::new(cell_id.clone())) + .and_then(|cell| cell.pending_calls.get_mut(call_id)), + }; + if let Some(call) = call { + return call.set_tool_result_sources(result_sources); + } + let ToolCallSource::CodeMode { cell_id, .. } = source else { + return false; + }; + let Some((retained, index)) = state.retained_calls.values_mut().find_map(|retained| { + if retained.runtime_cell_id.as_ref()?.as_str() != cell_id.as_str() { + return None; + } + let index = *retained.call_index_by_id.get(call_id)?; + Some((retained, index)) + }) else { + return false; + }; + // Older retry copies must not overwrite this output's accepted result metadata. + retained.sources_updated = true; + retained.calls[index].set_tool_result_sources(result_sources) + } + pub(crate) fn attach_pending_to_prompt( &self, items: &mut [ResponseItem], @@ -246,6 +291,13 @@ impl ExecutedToolCallRecorder { return false; } + // Updated records supersede older retry snapshots. + retry_cache.retain(|key, _| { + !state + .retained_calls + .get(key) + .is_some_and(|retained| retained.sources_updated) + }); let mut pending_retry_outputs = retry_cache.keys().cloned().collect::>(); let mut pending_retained_outputs = state.retained_calls.keys().cloned().collect::>(); @@ -294,13 +346,19 @@ impl ExecutedToolCallRecorder { .remove(call_id) .into_iter() .collect::>(); + let mut call_index_by_id = HashMap::new(); if let Some(output_cell_id) = state.output_cells.remove(call_id) && let Some(cell) = state.cells.get_mut(&output_cell_id) { cell_id = cell.originating_call_id.clone(); runtime_cell_id = Some(output_cell_id.clone()); let pending_calls = cell.pending_calls.len(); - calls.append(&mut cell.pending_calls); + for (call_id, call) in cell.pending_calls.drain(..) { + if matches!(call.arguments(), ExecutedToolCallArguments::Raw(_)) { + call_index_by_id.insert(call_id, calls.len()); + } + calls.push(call); + } cell.pending_full_argument_bytes = 0; complete = cell.completion == CellCompletion::Complete; if matches!( @@ -326,6 +384,8 @@ impl ExecutedToolCallRecorder { complete, cell_id: cell_id.clone(), runtime_cell_id, + call_index_by_id, + sources_updated: false, }, ); calls @@ -348,6 +408,7 @@ impl ExecutedToolCallRecorder { if retained_bytes > MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT { bound_executed_tool_calls_for_prompt_prioritizing_recent(items); let retained_before_bounding = std::mem::take(&mut state.retained_calls); + let mut bounded_outputs = HashSet::new(); for item in items { let call_id = match &*item { ResponseItem::FunctionCallOutput { @@ -363,10 +424,15 @@ impl ExecutedToolCallRecorder { }; let key = (std::mem::discriminant(&*item), call_id.clone()); let metadata = item.executed_tool_call_metadata(); - if let Some(retained) = retained_before_bounding.get(&key) + let unique_output = bounded_outputs.insert(key.clone()); + if !unique_output && let Some(retained) = state.retained_calls.get_mut(&key) { + retained.call_index_by_id.clear(); + } + let previous = retained_before_bounding.get(&key); + if let Some(retained) = previous && let Some(runtime_cell_id) = &retained.runtime_cell_id - && metadata.and_then(|metadata| metadata.executed_tool_calls.as_ref()) - != Some(&retained.calls) + && metadata + .is_none_or(|metadata| !metadata.has_same_tool_calls(&retained.calls)) && let Some(cell) = state.cells.get_mut(runtime_cell_id) { cell.completion = CellCompletion::Incomplete; @@ -378,9 +444,16 @@ impl ExecutedToolCallRecorder { .is_some_and(|calls| !calls.is_empty()) || metadata.tool_calls_complete.is_some()) { - let runtime_cell_id = retained_before_bounding - .get(&key) - .and_then(|retained| retained.runtime_cell_id.clone()); + let runtime_cell_id = + previous.and_then(|retained| retained.runtime_cell_id.clone()); + // Budget rewriting or duplicate outputs invalidate the original index slots. + let call_index_by_id = previous + .filter(|retained| { + unique_output + && metadata.executed_tool_calls.as_ref() == Some(&retained.calls) + }) + .map(|retained| retained.call_index_by_id.clone()) + .unwrap_or_default(); let retained = state.retained_calls.entry(key).or_default(); retained.runtime_cell_id = runtime_cell_id; retained.calls = metadata.executed_tool_calls.clone().unwrap_or_default(); @@ -388,6 +461,9 @@ impl ExecutedToolCallRecorder { retained.cell_id = Some(cell_id.clone()); } retained.complete |= metadata.tool_calls_complete == Some(true); + retained.call_index_by_id = call_index_by_id; + retained.sources_updated = + previous.is_some_and(|retained| retained.sources_updated); } } } diff --git a/codex-rs/core/src/tools/executed_tool_calls_tests.rs b/codex-rs/core/src/tools/executed_tool_calls_tests.rs index 3a2011aaad..441ac44dfb 100644 --- a/codex-rs/core/src/tools/executed_tool_calls_tests.rs +++ b/codex-rs/core/src/tools/executed_tool_calls_tests.rs @@ -1,5 +1,6 @@ use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseInputItem; +use codex_protocol::models::ToolResultSources; use pretty_assertions::assert_eq; use serde_json::json; @@ -30,12 +31,18 @@ fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() { ToolMode::Direct, ); } + assert!(recorder.record_tool_result_sources( + &ToolCallSource::Direct, + "direct-0", + ToolResultSources::new(Vec::new()) + )); let cell_id = CellId::new("bounded-cell".to_string()); recorder.start_cell(&cell_id, "bounded-output"); - for _ in 0..MAX_PENDING_EXECUTED_TOOL_CALLS + 2 { + for index in 0..MAX_PENDING_EXECUTED_TOOL_CALLS + 2 { recorder.record_nested_tool_call( cell_id.clone(), + format!("nested-{index}"), ExecutedToolCall::new("nested_tool".to_string(), json!({})), /*original_bytes*/ 2, ); @@ -251,6 +258,7 @@ fn tool_call_completeness_requires_finished_lossless_recording() { if scenario != "empty" { recorder.record_nested_tool_call( cell_id.clone(), + "nested-call".to_string(), ExecutedToolCall::new("nested_tool".to_string(), json!({})), /*original_bytes*/ 2, ); @@ -271,12 +279,22 @@ fn tool_call_completeness_requires_finished_lossless_recording() { #[test] fn tool_call_completeness_survives_waits_without_changing_deltas() { - for truncated in [false, true] { + let sources = ToolResultSources::new(vec![codex_protocol::models::ToolResultSource { + r#type: "test_channel".to_string(), + id: "CLATE".to_string(), + }]); + for (truncated, sources) in [ + (false, sources.clone()), + (true, sources), + (false, ToolResultSources::parse_failed()), + (true, ToolResultSources::parse_failed()), + ] { let recorder = ExecutedToolCallRecorder::default(); let cell_id = CellId::new("multi-wait".to_string()); recorder.start_cell(&cell_id, "exec"); let mut history = Vec::new(); let mut expected = Vec::new(); + let mut retry_cache = HashMap::new(); for (index, call_id) in ["exec", "wait-1", "wait-2", "wait-3"] .into_iter() .enumerate() @@ -296,7 +314,43 @@ fn tool_call_completeness_survives_waits_without_changing_deltas() { } else { ExecutedToolCall::new("nested_tool".to_string(), json!({})) }; - recorder.record_nested_tool_call(cell_id.clone(), call.clone(), original_bytes); + recorder.record_nested_tool_call( + cell_id.clone(), + format!("nested-{index}"), + call.clone(), + original_bytes, + ); + let mut call = call; + if index == 1 { + let source = |cell_id: &str| ToolCallSource::CodeMode { + cell_id: cell_id.to_string(), + runtime_tool_call_id: "runtime-call".to_string(), + }; + assert!(!recorder.record_tool_result_sources( + &source("other-cell"), + "nested-0", + ToolResultSources::new(Vec::new()), + )); + let source = source(cell_id.as_str()); + assert!(recorder.record_tool_result_sources( + &source, + "nested-0", + sources.clone(), + )); + let mut delayed_call = + ExecutedToolCall::new("nested_tool".to_string(), json!({})); + assert!(delayed_call.set_tool_result_sources(sources.clone())); + let mut delayed_output = output("exec"); + delayed_output.append_executed_tool_calls(vec![delayed_call]); + delayed_output.set_tool_call_cell_id("exec"); + expected[0] = delayed_output; + assert!(recorder.record_tool_result_sources( + &source, + "nested-1", + ToolResultSources::new(Vec::new()) + )); + assert!(call.set_tool_result_sources(ToolResultSources::new(Vec::new()))); + } expected_output.append_executed_tool_calls(vec![call]); } else if index == 3 { recorder.finish_cell_recording(&cell_id); @@ -308,7 +362,6 @@ fn tool_call_completeness_survives_waits_without_changing_deltas() { expected_output.mark_tool_calls_complete(); } expected.push(expected_output); - let mut retry_cache = HashMap::new(); for _ in 0..2 { let mut prompt = history.clone(); assert!(recorder.attach_pending_to_prompt(&mut prompt, &mut retry_cache)); @@ -330,6 +383,7 @@ fn cell_correlation_uses_originating_exec_across_runtime_restarts() { recorder.start_cell(&runtime_cell_id, originating_call_id); recorder.record_nested_tool_call( runtime_cell_id.clone(), + "nested-call".to_string(), ExecutedToolCall::new("nested_tool".to_string(), json!({})), /*original_bytes*/ 2, ); @@ -364,9 +418,10 @@ fn request_truncation_prevents_completion_after_compaction() { let arguments = json!({ "payload": "x".repeat(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES - r#"{"payload":""}"#.len()), }); - for _ in 0..4 { + for index in 0..4 { recorder.record_nested_tool_call( cell_id.clone(), + format!("nested-{index}"), ExecutedToolCall::new("nested_tool".to_string(), arguments.clone()), MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES, ); @@ -383,6 +438,17 @@ fn request_truncation_prevents_completion_after_compaction() { ExecutedToolCallArguments::Truncated { .. } ))) ); + let source = ToolCallSource::CodeMode { + cell_id: cell_id.as_str().to_string(), + runtime_tool_call_id: "runtime-call".to_string(), + }; + for index in 0..4 { + assert!(!recorder.record_tool_result_sources( + &source, + &format!("nested-{index}"), + ToolResultSources::new(Vec::new()), + )); + } recorder.attach_pending_to_prompt(&mut [], &mut HashMap::new()); recorder.register_cell(&cell_id, "wait"); @@ -397,6 +463,88 @@ fn request_truncation_prevents_completion_after_compaction() { ); } +#[test] +fn source_shedding_preserves_completion_after_compaction() { + let recorder = ExecutedToolCallRecorder::default(); + let cell_id = CellId::new("compacted-cell".to_string()); + recorder.start_cell(&cell_id, "exec"); + + let argument_bytes = MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES - 256; + let arguments = json!({ + "payload": "x".repeat(argument_bytes - r#"{"payload":""}"#.len()), + }); + for index in 0..4 { + recorder.record_nested_tool_call( + cell_id.clone(), + format!("nested-{index}"), + ExecutedToolCall::new("nested_tool".to_string(), arguments.clone()), + argument_bytes, + ); + } + let source = ToolCallSource::CodeMode { + cell_id: cell_id.as_str().to_string(), + runtime_tool_call_id: "runtime-call".to_string(), + }; + let mut retry_cache = HashMap::new(); + recorder.attach_pending_to_prompt(&mut [output("exec")], &mut retry_cache); + assert!( + recorder.record_tool_result_sources( + &source, + "nested-0", + ToolResultSources::new( + (0..32) + .map(|index| codex_protocol::models::ToolResultSource { + r#type: "test_channel".to_string(), + id: format!("C{index:063}"), + }) + .collect() + ), + ) + ); + + assert!(recorder.record_tool_result_sources( + &source, + "nested-1", + ToolResultSources::parse_failed() + )); + let mut initial = [output("exec")]; + assert!(recorder.attach_pending_to_prompt(&mut initial, &mut retry_cache)); + let calls = initial[0] + .executed_tool_call_metadata() + .and_then(|metadata| metadata.executed_tool_calls.as_ref()) + .expect("source shedding must preserve recorded calls"); + assert_eq!( + calls, + &vec![ExecutedToolCall::new("nested_tool".to_string(), arguments); 4] + ); + assert!(!recorder.record_tool_result_sources( + &source, + "nested-0", + ToolResultSources::new(Vec::new()) + )); + let mut retry = [output("exec")]; + assert!(recorder.attach_pending_to_prompt(&mut retry, &mut retry_cache)); + assert_eq!(retry, initial); + + recorder.attach_pending_to_prompt(&mut [], &mut retry_cache); + assert!(retry_cache.is_empty()); + assert!(!recorder.record_tool_result_sources( + &source, + "nested-1", + ToolResultSources::new(Vec::new()) + )); + recorder.register_cell(&cell_id, "wait"); + recorder.finish_cell_recording(&cell_id); + let mut final_output = [output("wait")]; + recorder.attach_pending_to_prompt(&mut final_output, &mut HashMap::new()); + assert_eq!( + final_output[0] + .executed_tool_call_metadata() + .and_then(|metadata| metadata.tool_calls_complete), + Some(true), + ); +} + #[test] fn finished_cells_without_more_waits_do_not_block_new_calls() { let recorder = ExecutedToolCallRecorder::default(); @@ -404,13 +552,23 @@ fn finished_cells_without_more_waits_do_not_block_new_calls() { for index in 0..MAX_PENDING_EXECUTED_TOOL_CALLS { let cell = CellId::new(format!("cell-{index}")); recorder.start_cell(&cell, cell.as_str()); - recorder.record_nested_tool_call(cell.clone(), call.clone(), /*original_bytes*/ 2); + recorder.record_nested_tool_call( + cell.clone(), + format!("nested-{index}"), + call.clone(), + /*original_bytes*/ 2, + ); recorder.attach_pending_to_prompt(&mut [output(cell.as_str())], &mut HashMap::new()); recorder.finish_cell_recording(&cell); } let fresh = CellId::new("fresh".to_string()); recorder.start_cell(&fresh, "fresh-output"); - recorder.record_nested_tool_call(fresh.clone(), call.clone(), /*original_bytes*/ 2); + recorder.record_nested_tool_call( + fresh.clone(), + "fresh-call".to_string(), + call.clone(), + /*original_bytes*/ 2, + ); recorder.finish_cell_recording(&fresh); let mut expected = output("fresh-output"); expected.append_executed_tool_calls(vec![call]); diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 4d45c64349..48e828fcfa 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -284,6 +284,12 @@ impl CoreToolRuntime for McpHandler { } fn on_tool_result_accepted(&self, invocation: &ToolInvocation, result: &dyn ToolOutput) { + // Direct calls also record sources, before the Code Mode-only evidence path below. + if let Some(recorder) = invocation.session.services.executed_tool_calls.as_ref() + && let Some(sources) = result.tool_result_sources() + { + recorder.record_tool_result_sources(&invocation.source, &invocation.call_id, sources); + } let ToolCallSource::CodeMode { cell_id, .. } = &invocation.source else { return; }; diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 8e19ece6bc..e73711f8e6 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -242,6 +242,10 @@ impl ToolOutput for PostToolUseFeedbackOutput { fn code_mode_result(&self, payload: &ToolPayload) -> Value { self.original.code_mode_result(payload) } + + fn tool_result_sources(&self) -> Option { + self.original.tool_result_sources() + } } #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index 9005f811a0..71846545fe 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -41,6 +41,9 @@ pub use crate::permission_profile_snapshot::PermissionProfileSnapshot; pub use executed_tool_calls::ExecutedToolCall; pub use executed_tool_calls::ExecutedToolCallArguments; pub use executed_tool_calls::ExecutedToolCallTruncation; +pub use executed_tool_calls::MAX_TOOL_RESULT_SOURCE_FIELD_BYTES; +pub use executed_tool_calls::ToolResultSource; +pub use executed_tool_calls::ToolResultSources; pub use executed_tool_calls::bound_executed_tool_calls_for_prompt; pub use executed_tool_calls::bound_executed_tool_calls_for_prompt_prioritizing_recent; pub use executed_tool_calls::executed_tool_call_metadata_bytes; diff --git a/codex-rs/protocol/src/models/executed_tool_calls.rs b/codex-rs/protocol/src/models/executed_tool_calls.rs index 8f504ff05c..b07ffba3b5 100644 --- a/codex-rs/protocol/src/models/executed_tool_calls.rs +++ b/codex-rs/protocol/src/models/executed_tool_calls.rs @@ -7,6 +7,10 @@ use super::InternalChatMessageMetadataPassthrough; use super::ResponseItem; const MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES: usize = 8 * 1024; +/// Maximum distinct result sources retained for one tool invocation. +const MAX_TOOL_RESULT_SOURCES: usize = 32; +/// Maximum UTF-8 bytes for each source's `type` and `id` separately, not the source list. +pub const MAX_TOOL_RESULT_SOURCE_FIELD_BYTES: usize = 128; /// Maximum serialized warehouse-only attempted-tool metadata in one request. const MAX_EXECUTED_TOOL_CALL_METADATA_BYTES: usize = 32 * 1024; const EXECUTED_TOOL_CALL_METADATA_FIELD_BYTES: usize = b"\"executed_tool_calls\":".len(); @@ -56,6 +60,18 @@ pub fn executed_tool_call_metadata_bytes(item: &ResponseItem) -> usize { .saturating_add(executed_tool_call_metadata_field_bytes(metadata)) } +impl InternalChatMessageMetadataPassthrough { + /// Compares call order, names and arguments, ignoring optional source metadata. + pub fn has_same_tool_calls(&self, calls: &[ExecutedToolCall]) -> bool { + self.executed_tool_calls.as_ref().is_some_and(|recorded| { + recorded.len() == calls.len() + && recorded.iter().zip(calls).all(|(recorded, call)| { + recorded.name == call.name && recorded.arguments() == call.arguments() + }) + }) + } +} + /// Bounds attempted-tool metadata fairly across the complete serialized request. pub fn bound_executed_tool_calls_for_prompt(items: &mut [ResponseItem]) { bound_executed_tool_calls_for_prompt_with_priority(items, /*prioritize_recent*/ false); @@ -112,6 +128,23 @@ fn bound_executed_tool_calls_for_prompt_with_priority( original_metadata_bytes.saturating_add(executed_tool_call_metadata_bytes(item)); } + // Source evidence is optional; dropping it must not discard calls or their completion proof. + if original_metadata_bytes > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { + original_metadata_bytes = 0; + for item in items.iter_mut() { + if let Some(metadata) = item + .internal_chat_message_metadata_passthrough_mut() + .and_then(Option::as_mut) + { + for call in metadata.executed_tool_calls.iter_mut().flatten() { + call.tool_result_sources = None; + } + } + original_metadata_bytes = + original_metadata_bytes.saturating_add(executed_tool_call_metadata_bytes(item)); + } + } + // A terminal marker can be on an empty wait output, separate from the lost calls. if truncated || original_metadata_bytes > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { for item in items.iter_mut() { @@ -298,6 +331,58 @@ pub struct ExecutedToolCall { pub name: String, #[ts(type = "unknown")] arguments: ExecutedToolCallArguments, + /// Host-generated analytics only: ignore input JSON rather than accepting caller-supplied + /// evidence, and keep this out of public schemas and generated clients. + #[serde(default, skip_deserializing, skip_serializing_if = "Option::is_none")] + #[schemars(skip)] + #[ts(skip)] + tool_result_sources: Option, +} + +/// A bounded capture update. Omitted updates still clear any previously recorded evidence. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ToolResultSources(Option); + +#[derive(Debug, Clone, Serialize, PartialEq, Eq)] +enum ToolResultSourcesValue { + #[serde(rename = "parse_failed")] + ParseFailed, + // Preserve the existing array shape for successful captures, including an empty array. + #[serde(untagged)] + Sources(Vec), +} + +impl ToolResultSources { + /// Deduplicates a complete source list, discarding all sources if it exceeds a limit. + pub fn new(sources: Vec) -> Self { + let mut unique_sources = Vec::new(); + for source in sources { + if unique_sources.contains(&source) { + continue; + } + if unique_sources.len() == MAX_TOOL_RESULT_SOURCES + || source.r#type.len() > MAX_TOOL_RESULT_SOURCE_FIELD_BYTES + || source.id.len() > MAX_TOOL_RESULT_SOURCE_FIELD_BYTES + { + return Self(None); + } + unique_sources.push(source); + } + Self(Some(ToolResultSourcesValue::Sources(unique_sources))) + } + + /// Records that source parsing was attempted but failed, not that a budget was exceeded. + pub fn parse_failed() -> Self { + Self(Some(ToolResultSourcesValue::ParseFailed)) + } +} + +/// A trusted source identity observed in an accepted tool result by the host. +#[derive(Debug, Clone, Serialize, PartialEq, Eq)] +pub struct ToolResultSource { + #[serde(rename = "type")] + pub r#type: String, + pub id: String, } /// Trusted truncation details generated locally for an oversized attempted tool call. @@ -325,6 +410,7 @@ impl ExecutedToolCall { Self { name, arguments: ExecutedToolCallArguments::Raw(arguments), + tool_result_sources: None, } } @@ -340,6 +426,12 @@ impl ExecutedToolCall { &self.arguments } + /// Replaces this invocation's capture outcome, including clearing omitted evidence. + pub fn set_tool_result_sources(&mut self, sources: ToolResultSources) -> bool { + self.tool_result_sources = sources.0; + self.tool_result_sources.is_some() + } + fn truncation(&self) -> Option<&ExecutedToolCallTruncation> { match &self.arguments { ExecutedToolCallArguments::Raw(_) => None, diff --git a/codex-rs/protocol/src/models/executed_tool_calls_tests.rs b/codex-rs/protocol/src/models/executed_tool_calls_tests.rs index b49443fd11..34f673a8ea 100644 --- a/codex-rs/protocol/src/models/executed_tool_calls_tests.rs +++ b/codex-rs/protocol/src/models/executed_tool_calls_tests.rs @@ -140,6 +140,42 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { assert_eq!(oversized_items, bounded_items); } + let mut items = ["exec", "wait"].map(|call_id| { + let mut item = ResponseItem::from(ResponseInputItem::FunctionCallOutput { + call_id: call_id.to_string(), + output: FunctionCallOutputPayload::from_text(String::new()), + }); + item.set_tool_call_cell_id("cell-\"\\"); + item + }); + items[0].append_executed_tool_calls(vec![ + ExecutedToolCall::new( + "test_tool".to_string(), + serde_json::json!({ "payload": "x".repeat(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES - 256) }), + ); + 4 + ]); + items[1].mark_tool_calls_complete(); + let expected = items.clone(); + assert!( + first_executed_tool_call(&mut items[0]) + .expect("recorded call should exist") + .set_tool_result_sources(ToolResultSources::new( + (0..MAX_TOOL_RESULT_SOURCES) + .map(|index| ToolResultSource { + r#type: "document".to_string(), + id: format!( + "R{index:0width$}", + width = MAX_TOOL_RESULT_SOURCE_FIELD_BYTES - 1 + ), + }) + .collect(), + )) + ); + assert!(metadata_bytes(&items)? > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); + bound_executed_tool_calls_for_prompt(&mut items); + assert_eq!(items, expected); + // Empty waits can carry only the marker; its bytes still count toward the budget. for metadata in [None, Some(passthrough_metadata("turn-1"))] { let mut item = ResponseItem::from(ResponseInputItem::FunctionCallOutput { @@ -229,6 +265,17 @@ fn tool_call_completeness_is_host_only_and_fail_closed() -> Result<()> { "tool_calls_complete": true, }))?; assert_eq!(untrusted, passthrough_metadata("turn-1")); + for sources in [ + serde_json::json!([{ "type": "test_resource", "id": "ATTACKER" }]), + serde_json::json!("parse_failed"), + ] { + let untrusted_call = serde_json::from_value::(serde_json::json!({ + "name": "test_tool", + "arguments": {}, + "tool_result_sources": sources, + }))?; + assert_eq!(untrusted_call, call); + } let mut item = ResponseItem::from(ResponseInputItem::FunctionCallOutput { call_id: "call-1".to_string(), @@ -272,3 +319,72 @@ fn tool_call_completeness_is_host_only_and_fail_closed() -> Result<()> { } Ok(()) } + +#[test] +fn tool_result_source_snapshots_replace_atomically() -> Result<()> { + let source = |kind: &str, id: &str| ToolResultSource { + r#type: kind.to_string(), + id: id.to_string(), + }; + let mut call = ExecutedToolCall::new("test_tool".to_string(), serde_json::json!({})); + let mut sources = (0..MAX_TOOL_RESULT_SOURCES - 1) + .map(|index| source("test_resource", &format!("R{index}"))) + .collect::>(); + sources.push(source("other_resource", "R0")); + sources.push(sources[0].clone()); + let capture = ToolResultSources::new(sources.clone()); + sources.truncate(MAX_TOOL_RESULT_SOURCES); + assert_eq!( + capture, + ToolResultSources(Some(ToolResultSourcesValue::Sources(sources.clone()))) + ); + assert!(call.set_tool_result_sources(capture)); + sources.push(source("test_resource", "OVERFLOW")); + let capture = ToolResultSources::new(sources); + assert_eq!(capture, ToolResultSources(None)); + assert!(!call.set_tool_result_sources(capture)); + assert!( + serde_json::to_value(&call)? + .get("tool_result_sources") + .is_none() + ); + + // Measure UTF-8 bytes, and clear old evidence instead of keeping a partial replacement. + let field = format!( + "é{}", + "x".repeat(MAX_TOOL_RESULT_SOURCE_FIELD_BYTES - "é".len()) + ); + let bounded = source(&field, &field); + let oversized = format!("{field}x"); + for invalid in [ + source(&oversized, "R1"), + source("test_resource", &oversized), + ] { + assert!(call.set_tool_result_sources(ToolResultSources::new(vec![bounded.clone()]))); + assert_eq!( + call.tool_result_sources, + Some(ToolResultSourcesValue::Sources(vec![bounded.clone()])) + ); + let capture = ToolResultSources::new(vec![source("test_resource", "R1"), invalid]); + assert_eq!(capture, ToolResultSources(None)); + assert!(!call.set_tool_result_sources(capture)); + assert_eq!(call.tool_result_sources, None); + } + + assert!(call.set_tool_result_sources(ToolResultSources::new(vec![bounded]))); + assert!(call.set_tool_result_sources(ToolResultSources::parse_failed())); + assert_eq!( + serde_json::to_value(&call)?, + serde_json::json!({ + "name": "test_tool", + "arguments": {}, + "tool_result_sources": "parse_failed", + }) + ); + assert!(call.set_tool_result_sources(ToolResultSources::new(Vec::new()))); + assert_eq!( + serde_json::to_value(&call)?["tool_result_sources"], + serde_json::json!([]) + ); + Ok(()) +} diff --git a/codex-rs/tools/src/tool_output.rs b/codex-rs/tools/src/tool_output.rs index 0aa22b626d..a33232e20d 100644 --- a/codex-rs/tools/src/tool_output.rs +++ b/codex-rs/tools/src/tool_output.rs @@ -55,6 +55,11 @@ pub trait ToolOutput: Send { fn code_mode_result(&self, payload: &ToolPayload) -> JsonValue { response_input_to_code_mode_result(self.to_response_item("", payload)) } + + /// Reports configured source capture only after acceptance; `None` means no capture attempt. + fn tool_result_sources(&self) -> Option { + None + } } impl ToolOutput for Box @@ -96,6 +101,10 @@ where fn code_mode_result(&self, payload: &ToolPayload) -> JsonValue { (**self).code_mode_result(payload) } + + fn tool_result_sources(&self) -> Option { + (**self).tool_result_sources() + } } #[derive(Clone, Debug, PartialEq)]