Preserve large arguments in executed tool call metadata (#52268)

## Why

Fixed argument limits truncated recorded tool calls even when their metadata fit within the available message budget.

## What changed

- Remove the 8 KiB per-call and 32 KiB per-output argument limits for direct calls and Code Mode.
- Preserve full recorded arguments while retaining message metadata budgets, pending-call limits, and the direct metadata retention limit.
- Rename the normalization helper to `normalize_executed_tool_call_completeness`; it now only clears completeness when calls or tool names were lost.

## Testing

Update regression coverage for large arguments, escaped inputs, reserved-key wrappers, and argument and result metadata preservation across Code Mode waits.

GitOrigin-RevId: 8c199ccf39191d5d44eef02d28b1773654735b43
This commit is contained in:
ningyi-oai
2026-10-08 20:40:47 +00:00
committed by copyberry
parent 3b6ab3471a
commit 515c291d87
10 changed files with 97 additions and 168 deletions
@@ -8,7 +8,6 @@ use ts_rs::TS;
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.
@@ -86,12 +85,10 @@ impl InternalChatMessageMetadataPassthrough {
}
}
/// Bounds recorded arguments, preserving completion unless calls or tool names were lost.
/// Returns cells whose arguments were newly truncated. Ordinary tool-call arguments and
/// outputs are not changed.
pub fn normalize_executed_tool_call_arguments(items: &mut [ResponseItem]) -> HashSet<String> {
/// Clears completion across a cell when recorded calls or tool names were lost.
/// Argument truncation preserves inventory completeness; arguments and outputs are unchanged.
pub fn normalize_executed_tool_call_completeness(items: &mut [ResponseItem]) {
let mut damaged_cells = HashSet::new();
let mut newly_truncated_cells = HashSet::new();
for item in items.iter_mut() {
let Some(metadata) = item
.internal_chat_message_metadata_passthrough_mut()
@@ -99,33 +96,17 @@ pub fn normalize_executed_tool_call_arguments(items: &mut [ResponseItem]) -> Has
else {
continue;
};
let mut inventory_lost = false;
let mut newly_truncated = false;
for call in metadata.executed_tool_calls.iter_mut().flatten() {
let argument_bytes = serde_json::to_vec(&call.arguments)
.map(|bytes| bytes.len())
.unwrap_or(usize::MAX);
if call.truncation().is_none() && argument_bytes > MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES
{
call.set_truncation(
argument_bytes,
MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES,
/*omitted_calls*/ None,
);
newly_truncated = true;
}
inventory_lost |= !call.has_complete_inventory();
}
if inventory_lost {
if metadata
.executed_tool_calls
.iter()
.flatten()
.any(|call| !call.has_complete_inventory())
{
metadata.tool_calls_complete = None;
damaged_cells.extend(metadata.cell_id.clone());
}
if newly_truncated {
newly_truncated_cells.extend(metadata.cell_id.clone());
}
}
clear_damaged_cell_completeness(items, &damaged_cells);
newly_truncated_cells
}
/// Bounds optional observations to the space left in the actual outgoing message.
@@ -339,9 +320,7 @@ fn truncate_call_arguments_to_fit(
let truncated = ExecutedToolCallArguments::Truncated {
truncation: ExecutedToolCallTruncation {
original_bytes,
max_bytes: original_bytes
.saturating_sub(overage_bytes)
.min(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES),
max_bytes: original_bytes.saturating_sub(overage_bytes),
omitted_calls: None,
original_name_bytes: None,
},
@@ -849,7 +828,7 @@ impl ResponseItem {
let omitted_calls = (represented_calls > 1).then_some(represented_calls - 1);
call.set_truncation_with_name(
original_bytes,
max_call_bytes.min(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES),
max_call_bytes,
omitted_calls,
original_name_bytes,
);
@@ -862,7 +841,7 @@ impl ResponseItem {
let call = &mut calls[0];
call.set_truncation_with_name(
original_bytes,
max_call_bytes.min(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES),
max_call_bytes,
omitted_calls,
Some(original_name_bytes.unwrap_or(call.name.len())),
);
@@ -132,7 +132,7 @@ fn model_arguments_cannot_forge_executed_tool_call_truncation() -> Result<()> {
forged_marker.clone(),
)]);
normalize_executed_tool_call_arguments(std::slice::from_mut(&mut item));
normalize_executed_tool_call_completeness(std::slice::from_mut(&mut item));
let call = item
.executed_tool_call_metadata()
.and_then(|metadata| metadata.executed_tool_calls.as_ref())
@@ -177,7 +177,7 @@ fn tool_call_completeness_is_host_only_and_fail_closed() -> Result<()> {
let mut item = output("call-1");
item.set_tool_call_cell_id("cell-1");
item.mark_tool_calls_complete();
normalize_executed_tool_call_arguments(std::slice::from_mut(&mut item));
normalize_executed_tool_call_completeness(std::slice::from_mut(&mut item));
assert_eq!(
serde_json::to_value(&item)?["internal_chat_message_metadata_passthrough"],
serde_json::json!({
@@ -188,21 +188,15 @@ fn tool_call_completeness_is_host_only_and_fail_closed() -> Result<()> {
item.clear_executed_tool_calls();
assert!(item.executed_tool_call_metadata().is_none());
for (call, newly_truncated) in [
(
ExecutedToolCall::new(
"test_tool".to_string(),
serde_json::json!({ "payload": "x".repeat(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES + 1) }),
),
true,
for call in [
ExecutedToolCall::new(
"test_tool".to_string(),
serde_json::json!({ "payload": "x".repeat(9_000) }),
),
(
ExecutedToolCall::truncated(
"test_tool".to_string(),
/*original_bytes*/ 9_000,
/*max_bytes*/ 0,
),
false,
ExecutedToolCall::truncated(
"test_tool".to_string(),
/*original_bytes*/ 9_000,
/*max_bytes*/ 0,
),
] {
for same_cell in [false, true] {
@@ -214,12 +208,11 @@ fn tool_call_completeness_is_host_only_and_fail_closed() -> Result<()> {
}
item.mark_tool_calls_complete();
}
let changed_cells = normalize_executed_tool_call_arguments(&mut items);
assert_eq!(
changed_cells.contains("cell-1"),
same_cell && newly_truncated
);
assert!(normalize_executed_tool_call_arguments(&mut items).is_empty());
let expected = items.clone();
normalize_executed_tool_call_completeness(&mut items);
assert_eq!(items, expected);
normalize_executed_tool_call_completeness(&mut items);
assert_eq!(items, expected);
assert_eq!(
items.map(|item| item
.executed_tool_call_metadata()
@@ -605,7 +598,7 @@ fn inventory_loss_still_clears_completeness() {
let mut item = output("call");
item.append_executed_tool_calls(vec![call]);
item.mark_tool_calls_complete();
normalize_executed_tool_call_arguments(std::slice::from_mut(&mut item));
normalize_executed_tool_call_completeness(std::slice::from_mut(&mut item));
assert_eq!(
item.executed_tool_call_metadata()
.unwrap()