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
This commit is contained in:
ningyi-oai
2026-09-02 01:06:18 +00:00
committed by copyberry
parent 8971fc25a9
commit 12ed76b09b
11 changed files with 584 additions and 17 deletions
+100 -1
View File
@@ -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<AppToolApproval>,
/// 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<AppToolResultSourceRequirementToml>,
}
#[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<Self, Self::Err> {
Ok(match value {
"detailed_message_search_v1" => Self::DetailedMessageSearchV1,
_ => Self::Unknown(value.to_string()),
})
}
}
impl<'de> Deserialize<'de> for AppToolResultSourceFormat {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
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::<AppToolResultSourceRequirementToml>(incomplete_rule).is_err());
}
Ok(())
}
fn apps_requirements(entries: &[(&str, Option<bool>)]) -> 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,
},
)]),
}),
+2
View File
@@ -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;
@@ -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()
},
)]),
}),
+85 -9
View File
@@ -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<String>,
runtime_cell_id: Option<CellId>,
// Invocation IDs stay local and are retained only with their original output's calls.
call_index_by_id: HashMap<String, usize>,
sources_updated: bool,
}
#[derive(Default, PartialEq, Eq)]
@@ -56,7 +60,7 @@ enum CellCompletion {
#[derive(Default)]
struct RecordedCell {
pending_calls: Vec<ExecutedToolCall>,
pending_calls: IndexMap<String, ExecutedToolCall>,
pending_full_argument_bytes: usize,
completion: CellCompletion,
originating_call_id: Option<String>,
@@ -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::<HashSet<_>>();
let mut pending_retained_outputs =
state.retained_calls.keys().cloned().collect::<HashSet<_>>();
@@ -294,13 +346,19 @@ impl ExecutedToolCallRecorder {
.remove(call_id)
.into_iter()
.collect::<Vec<_>>();
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);
}
}
}
@@ -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]);
+6
View File
@@ -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;
};
+4
View File
@@ -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<codex_protocol::models::ToolResultSources> {
self.original.tool_result_sources()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
+3
View File
@@ -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;
@@ -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<ToolResultSourcesValue>,
}
/// A bounded capture update. Omitted updates still clear any previously recorded evidence.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolResultSources(Option<ToolResultSourcesValue>);
#[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<ToolResultSource>),
}
impl ToolResultSources {
/// Deduplicates a complete source list, discarding all sources if it exceeds a limit.
pub fn new(sources: Vec<ToolResultSource>) -> 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,
@@ -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::<ExecutedToolCall>(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::<Vec<_>>();
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(())
}
+9
View File
@@ -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<codex_protocol::models::ToolResultSources> {
None
}
}
impl<T> ToolOutput for Box<T>
@@ -96,6 +101,10 @@ where
fn code_mode_result(&self, payload: &ToolPayload) -> JsonValue {
(**self).code_mode_result(payload)
}
fn tool_result_sources(&self) -> Option<codex_protocol::models::ToolResultSources> {
(**self).tool_result_sources()
}
}
#[derive(Clone, Debug, PartialEq)]