mirror of
https://github.com/openai/codex.git
synced 2026-10-20 00:18:49 +00:00
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:
@@ -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,
|
||||
},
|
||||
)]),
|
||||
}),
|
||||
|
||||
@@ -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()
|
||||
},
|
||||
)]),
|
||||
}),
|
||||
|
||||
@@ -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]);
|
||||
|
||||
@@ -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;
|
||||
};
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
@@ -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)]
|
||||
|
||||
Reference in New Issue
Block a user