Preserve more tool metadata within outgoing message budgets (#47758)

## Why

The 32 KiB per-result cap and 128 KiB aggregate budget discard tool metadata even when the outgoing request has room to retain it.

## What changed

- Remove the per-result capture cap and raise the aggregate executed-tool metadata budget to 2 MiB.
- Apply a 15 MiB soft budget to serialized HTTP requests and WebSocket messages, using the actual WebSocket delta. Trim optional result metadata on the wire copy while preserving ordinary content and logical history.
- Prioritize `openai/resource_access` metadata during trimming, and include `overage_bytes` in omission markers.
- Record bytes removed by metadata budgets in `codex.tool_metadata.shed_bytes`, tagged by stage.

## Testing

Add coverage for HTTP and WebSocket message boundaries, resource preservation, residual marker removal, unchanged history, and shedding metrics. Update direct-tool and code-mode tests to verify metadata survives beyond the previous limits.

GitOrigin-RevId: ed00996b77945300f60db1e52eff22b4de2707c1
This commit is contained in:
ningyi-oai
2026-09-24 05:48:45 +00:00
committed by copyberry
parent b0a9dcb843
commit 53446f90a5
15 changed files with 969 additions and 279 deletions
+12
View File
@@ -32,6 +32,9 @@ use std::sync::OnceLock;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
#[path = "client_tool_metadata.rs"]
mod tool_metadata;
use crate::CodexResponsesHeaders;
use async_channel::Sender;
use codex_api::AgentIdentityTelemetry;
@@ -1743,6 +1746,9 @@ impl ModelClientSession {
);
let inference_trace_attempt = inference_trace.start_attempt();
inference_trace_attempt.add_request_headers(&mut options.extra_headers);
if let Some(input) = tool_metadata::bounded_input(&request, &request.input) {
request.input = input;
}
inference_trace_attempt.record_started(&request);
let client = ApiResponsesClient::new(
transport,
@@ -2022,6 +2028,12 @@ impl ModelClientSession {
);
let mut ws_request = ResponsesWsRequest::ResponseCreate(ws_payload);
stamp_ws_stream_request_start_ms(&mut ws_request);
let ResponsesWsRequest::ResponseCreate(payload) = &ws_request;
let bounded_input = tool_metadata::bounded_input(&ws_request, payload.input);
if let Some(input) = bounded_input.as_deref() {
let ResponsesWsRequest::ResponseCreate(payload) = &mut ws_request;
payload.input = input;
}
if !previous_response_id_from_untraced_warmup {
inference_trace_attempt.record_started(&ws_request);
}
+5 -7
View File
@@ -611,8 +611,8 @@ fn responses_request_limits_internal_metadata_to_resolved_first_party_https_endp
}
#[test]
fn responses_request_preserves_smaller_result_metadata_under_aggregate_budget() -> anyhow::Result<()>
{
fn responses_request_preserves_result_metadata_above_previous_aggregate_budget()
-> anyhow::Result<()> {
let result_metadata = [
json!({ "payload": "l".repeat(31 * 1024) }),
json!({ "payload": "m".repeat(20 * 1024) }),
@@ -643,9 +643,6 @@ fn responses_request_preserves_smaller_result_metadata_under_aggregate_budget()
history.push(output);
}
let original_history = serde_json::to_value(&history)?;
let mut expected_input = original_history.clone();
expected_input[1]["internal_chat_message_metadata_passthrough"]["executed_tool_calls"][0]["tool_result_metadata"] =
json!("omitted_due_to_size_limit");
let mut features = codex_features::Features::default();
features.enable(codex_features::Feature::ExecutedToolCallMetadata);
@@ -682,7 +679,7 @@ fn responses_request_preserves_smaller_result_metadata_under_aggregate_budget()
)?;
let body = serde_json::to_value(&request)?;
// Whole-input equality covers bindings, arguments, results, sources and completion too.
assert_eq!(body["input"], expected_input);
assert_eq!(body["input"], original_history);
let mut without_metadata = body.clone();
for item in without_metadata["input"].as_array_mut().unwrap() {
item.as_object_mut()
@@ -691,7 +688,8 @@ fn responses_request_preserves_smaller_result_metadata_under_aggregate_budget()
}
let metadata_bytes =
serde_json::to_vec(&body)?.len() - serde_json::to_vec(&without_metadata)?.len();
assert!(metadata_bytes <= 128 * 1024);
assert!(metadata_bytes > 128 * 1024);
assert!(metadata_bytes <= 2 * 1024 * 1024);
assert_eq!(serde_json::to_value(&history)?, original_history);
Ok(())
}
+82
View File
@@ -0,0 +1,82 @@
//! Bounds optional tool-result metadata using the actual outgoing Responses message.
//! Only the wire copy changes; model-visible input and WebSocket continuation history stay intact.
use crate::tools::metadata_metrics;
use crate::utils::json::serialized_json_bytes;
use codex_protocol::models::ResponseItem;
use serde::Serialize;
const MAX_RESPONSE_MESSAGE_BYTES: usize = 15 * 1024 * 1024;
pub(super) fn bounded_input<T: Serialize>(
message: &T,
input: &[ResponseItem],
) -> Option<Vec<ResponseItem>> {
if !input.iter().any(ResponseItem::has_tool_result_metadata) {
return None;
}
// Serialization errors belong to the transport's existing error path, not this soft budget.
let message_bytes = serialized_json_bytes(message).ok()?;
if message_bytes <= MAX_RESPONSE_MESSAGE_BYTES {
return None;
}
let before = metadata_metrics::metadata_bytes(input);
let mut bounded = input.to_vec();
let overage = message_bytes - MAX_RESPONSE_MESSAGE_BYTES;
for item in &mut bounded {
item.retain_tool_resource_access_or_omit_metadata(overage);
}
let after = metadata_metrics::metadata_bytes(&bounded);
let mut remaining_message_bytes = message_bytes.saturating_sub(before.saturating_sub(after));
// Small values and omission markers can still cause the overage. Remove their
// fields before sacrificing resource evidence, even when it occurs earlier.
reduce_input_until_fit(
&mut bounded,
&mut remaining_message_bytes,
ResponseItem::retain_tool_resource_access,
);
let overage = remaining_message_bytes.saturating_sub(MAX_RESPONSE_MESSAGE_BYTES);
reduce_input_until_fit(&mut bounded, &mut remaining_message_bytes, |item| {
item.omit_tool_result_metadata(overage);
});
// Newly created markers are optional too. Clear those before any small
// resource-only value that was cheaper to retain than replace with a marker.
reduce_input_until_fit(
&mut bounded,
&mut remaining_message_bytes,
ResponseItem::retain_tool_resource_access,
);
reduce_input_until_fit(
&mut bounded,
&mut remaining_message_bytes,
ResponseItem::clear_tool_result_metadata,
);
let after = metadata_metrics::metadata_bytes(&bounded);
if before == after {
return None;
}
metadata_metrics::record_shedding("message", before, after, codex_otel::global().as_ref());
// Ordinary content alone can exceed the soft limit. Preserve it and existing error handling.
Some(bounded)
}
fn reduce_input_until_fit(
input: &mut [ResponseItem],
message_bytes: &mut usize,
reduce: impl Fn(&mut ResponseItem),
) {
for item in input {
if *message_bytes <= MAX_RESPONSE_MESSAGE_BYTES {
break;
}
let before = metadata_metrics::metadata_bytes(std::slice::from_ref(item));
reduce(item);
let after = metadata_metrics::metadata_bytes(std::slice::from_ref(item));
// Include the field name and comma when a result-metadata field vanishes.
*message_bytes = message_bytes.saturating_sub(before.saturating_sub(after));
}
}
#[cfg(test)]
#[path = "client_tool_metadata_tests.rs"]
mod tests;
@@ -0,0 +1,263 @@
use super::*;
use codex_api::ResponseCreateWsRequest;
use codex_api::ResponsesApiRequest;
use codex_api::ResponsesWsRequest;
use codex_protocol::models::ExecutedToolCall;
use codex_protocol::models::FunctionCallOutputPayload;
use codex_protocol::models::ResponseInputItem;
use codex_protocol::models::ToolResultMetadata;
use codex_protocol::models::ToolResultSource;
use codex_protocol::models::ToolResultSources;
use pretty_assertions::assert_eq;
use serde_json::json;
fn request() -> ResponsesApiRequest {
request_with_metadata(&json!({
"provider": "x".repeat(16 * 1024),
"openai/resource_access": {"resources": ["é".repeat(4 * 1024)]},
}))
}
fn request_with_metadata(metadata: &serde_json::Value) -> ResponsesApiRequest {
let mut output = ResponseItem::from(ResponseInputItem::FunctionCallOutput {
call_id: "call".to_string(),
output: FunctionCallOutputPayload::from_text("keep result".to_string()),
});
let mut call = ExecutedToolCall::new("mcp__apps__read".to_string(), json!({"query": "keep"}));
call.set_tool_result_metadata(ToolResultMetadata::new(metadata));
call.set_tool_result_sources(ToolResultSources::new(vec![ToolResultSource {
r#type: "test_resource".to_string(),
id: "resource".to_string(),
}]));
output.append_executed_tool_calls(vec![call]);
output.set_tool_call_cell_id("cell");
output.mark_tool_calls_complete();
ResponsesApiRequest {
model: "test".to_string(),
instructions: "x".to_string(),
input: vec![output],
tools: None,
tool_choice: "auto".to_string(),
parallel_tool_calls: true,
reasoning: None,
store: false,
stream: true,
stream_options: None,
include: Vec::new(),
service_tier: None,
prompt_cache_key: None,
text: None,
client_metadata: None,
access_programs: None,
}
}
#[test]
fn http_message_budget_preserves_resources_until_the_message_still_exceeds_limit() {
for overage in [0, 4 * 1024, 20 * 1024, 40 * 1024] {
let mut request = request();
let target_bytes = MAX_RESPONSE_MESSAGE_BYTES + overage;
request.instructions =
"x".repeat(target_bytes - serialized_json_bytes(&request).unwrap() + 1);
assert_eq!(serialized_json_bytes(&request).unwrap(), target_bytes);
assert!(metadata_metrics::metadata_bytes(&request.input) < 32 * 1024);
let original_input = request.input.clone();
let bounded = bounded_input(&request, &request.input);
if overage == 0 {
assert_eq!(bounded, None);
continue;
}
let bounded = bounded.unwrap();
let encoded = serde_json::to_value(&bounded[0]).unwrap();
let metadata = &encoded["internal_chat_message_metadata_passthrough"]["executed_tool_calls"]
[0]["tool_result_metadata"];
if overage == 4 * 1024 {
assert_eq!(
metadata,
&json!({
"openai/resource_access": {"resources": ["é".repeat(4 * 1024)]},
})
);
} else if overage == 20 * 1024 {
let mut resource_only = original_input.clone();
resource_only[0].retain_tool_resource_access();
let first_reduction = metadata_metrics::metadata_bytes(&original_input)
- metadata_metrics::metadata_bytes(&resource_only);
assert_eq!(
metadata,
&json!(format!(
"omitted_due_to_size_limit (overage_bytes={})",
overage - first_reduction,
))
);
} else {
assert!(
encoded["internal_chat_message_metadata_passthrough"]["executed_tool_calls"][0]
.get("tool_result_metadata")
.is_none()
);
}
assert_eq!(request.input, original_input);
let mut ordinary = original_input;
let mut bounded_ordinary = bounded.clone();
for items in [&mut ordinary, &mut bounded_ordinary] {
for item in items {
item.clear_tool_result_metadata();
}
}
assert_eq!(bounded_ordinary, ordinary);
request.input = bounded;
// The last case cannot fit even without raw metadata: ordinary content is never cut.
assert_eq!(
serialized_json_bytes(&request).unwrap() > MAX_RESPONSE_MESSAGE_BYTES,
overage == 40 * 1024,
);
assert_eq!(bounded_input(&request, &request.input), None);
}
}
#[test]
fn message_budget_removes_residual_values_and_markers_when_ordinary_input_fits() {
for metadata in [
json!({}),
json!(null),
json!("omitted_due_to_size_limit"),
json!("omitted_due_to_size_limit (overage_bytes=1)"),
json!({"provider": "x".repeat(1024)}),
json!({"openai/resource_access": {}}),
] {
for websocket in [false, true] {
let mut request = request_with_metadata(&metadata);
let mut ordinary = request.clone();
ordinary.input[0].clear_tool_result_metadata();
let ordinary_bytes = if websocket {
serialized_json_bytes(&ResponsesWsRequest::ResponseCreate(
ResponseCreateWsRequest::from(&ordinary),
))
.unwrap()
} else {
serialized_json_bytes(&ordinary).unwrap()
};
// The original instructions contain one byte, so leave exactly one
// byte below the limit after removing only result metadata.
ordinary.instructions = "x".repeat(MAX_RESPONSE_MESSAGE_BYTES - ordinary_bytes);
request.instructions = ordinary.instructions.clone();
let original_input = request.input.clone();
let bounded = if websocket {
let ordinary_message =
ResponsesWsRequest::ResponseCreate(ResponseCreateWsRequest::from(&ordinary));
assert_eq!(
serialized_json_bytes(&ordinary_message).unwrap(),
MAX_RESPONSE_MESSAGE_BYTES - 1,
);
let mut message =
ResponsesWsRequest::ResponseCreate(ResponseCreateWsRequest::from(&request));
assert!(serialized_json_bytes(&message).unwrap() > MAX_RESPONSE_MESSAGE_BYTES);
let bounded = bounded_input(&message, &request.input)
.expect("optional result metadata must not keep a fitting message oversized");
let ResponsesWsRequest::ResponseCreate(payload) = &mut message;
payload.input = &bounded;
assert_eq!(
serialized_json_bytes(&message).unwrap(),
MAX_RESPONSE_MESSAGE_BYTES - 1,
);
bounded
} else {
assert_eq!(
serialized_json_bytes(&ordinary).unwrap(),
MAX_RESPONSE_MESSAGE_BYTES - 1,
);
assert!(serialized_json_bytes(&request).unwrap() > MAX_RESPONSE_MESSAGE_BYTES);
let bounded = bounded_input(&request, &request.input)
.expect("optional result metadata must not keep a fitting message oversized");
let mut message = request.clone();
message.input = bounded.clone();
assert_eq!(message, ordinary);
assert_eq!(
serialized_json_bytes(&message).unwrap(),
MAX_RESPONSE_MESSAGE_BYTES - 1,
);
bounded
};
// Whole-input equality covers sources, calls, completion and ordinary output.
assert_eq!(bounded, ordinary.input);
assert_eq!(request.input, original_input);
}
}
}
#[test]
fn residual_marker_removal_preserves_earlier_resources_and_stops_when_message_fits() {
let mut request = request_with_metadata(&json!({
"openai/resource_access": {
"schema_version": 1,
"resource_coverage": "complete",
"resources": [],
},
}));
for call_id in ["marker", "later-marker"] {
let mut output = ResponseItem::from(ResponseInputItem::FunctionCallOutput {
call_id: call_id.to_string(),
output: FunctionCallOutputPayload::from_text("keep result".to_string()),
});
let mut call = ExecutedToolCall::new("mcp__apps__read".to_string(), json!({}));
call.set_tool_result_metadata(ToolResultMetadata::new(&json!(
"omitted_due_to_size_limit (overage_bytes=1)"
)));
output.append_executed_tool_calls(vec![call]);
output.mark_tool_calls_complete();
request.input.push(output);
}
let mut expected = request.clone();
expected.input[1].clear_tool_result_metadata();
expected.instructions =
"x".repeat(MAX_RESPONSE_MESSAGE_BYTES - serialized_json_bytes(&expected).unwrap() + 1);
request.instructions = expected.instructions.clone();
assert_eq!(
serialized_json_bytes(&expected).unwrap(),
MAX_RESPONSE_MESSAGE_BYTES
);
assert!(serialized_json_bytes(&request).unwrap() > MAX_RESPONSE_MESSAGE_BYTES);
let original_input = request.input.clone();
let bounded = bounded_input(&request, &request.input).unwrap();
assert_eq!(bounded, expected.input);
assert_eq!(request.input, original_input);
request.input = bounded;
assert_eq!(request, expected);
assert_eq!(
serialized_json_bytes(&request).unwrap(),
MAX_RESPONSE_MESSAGE_BYTES
);
}
#[test]
fn websocket_budget_uses_the_actual_delta_and_does_not_mutate_logical_history() {
let mut request = request();
let old_output = ResponseItem::from(ResponseInputItem::FunctionCallOutput {
call_id: "old".to_string(),
output: FunctionCallOutputPayload::from_text("x".repeat(MAX_RESPONSE_MESSAGE_BYTES)),
});
request.input.insert(0, old_output);
let delta = &request.input[1..];
let delta_message = ResponsesWsRequest::ResponseCreate(ResponseCreateWsRequest {
previous_response_id: Some("previous-response".to_string()),
input: delta,
..ResponseCreateWsRequest::from(&request)
});
assert!(serialized_json_bytes(&request).unwrap() > MAX_RESPONSE_MESSAGE_BYTES);
assert!(serialized_json_bytes(&delta_message).unwrap() < 32 * 1024);
assert_eq!(bounded_input(&delta_message, delta), None);
let mut full = ResponsesWsRequest::ResponseCreate(ResponseCreateWsRequest::from(&request));
let bounded = bounded_input(&full, &request.input).unwrap();
let ResponsesWsRequest::ResponseCreate(payload) = &mut full;
payload.input = &bounded;
assert_eq!(bounded[0], request.input[0]);
let mut ordinary = request.input[1].clone();
let mut bounded_ordinary = bounded[1].clone();
ordinary.clear_tool_result_metadata();
bounded_ordinary.clear_tool_result_metadata();
assert_eq!(bounded_ordinary, ordinary);
assert_eq!(request.input[1], self::request().input[0]);
assert_eq!(bounded_input(&full, &bounded), None);
}
+3 -3
View File
@@ -48,7 +48,7 @@ async fn local_compaction_respects_tool_metadata_state(
.await;
let mut items = vec![user_message("Update the plan")];
for index in 0..20 {
for index in 0..300 {
let call_id = format!("direct-{index}");
let arguments = json!({"plan": [{"step": "x".repeat(7 * 1024), "status": "completed"}]});
assert!(serde_json::to_vec(&arguments)?.len() < 8 * 1024);
@@ -129,7 +129,7 @@ async fn local_compaction_respects_tool_metadata_state(
_ => None,
})
.collect::<serde_json::Result<Vec<_>>>()?;
assert_eq!(outputs.len(), 20);
assert_eq!(outputs.len(), 300);
let metadata_bytes: usize = outputs
.iter()
.map(|item| {
@@ -140,7 +140,7 @@ async fn local_compaction_respects_tool_metadata_state(
.sum();
// Compaction does not rebudget source records as a normal inference request.
// Passthrough bytes are also excluded from model token estimates.
assert!(metadata_bytes > 128 * 1024);
assert!(metadata_bytes > 2 * 1024 * 1024);
if !metadata_enabled {
let mut config = (*session.get_config().await).clone();
+16 -1
View File
@@ -25,6 +25,7 @@ use codex_protocol::mcp::McpAttributionSource;
use codex_protocol::models::ExecutedToolCall;
use codex_protocol::models::ExecutedToolCallArguments;
use codex_protocol::models::ResponseItem;
use codex_protocol::models::ToolResultMetadata;
use codex_protocol::models::bound_executed_tool_calls_for_prompt;
use codex_protocol::models::bound_executed_tool_calls_for_prompt_prioritizing_recent;
use codex_protocol::models::executed_tool_call_metadata_bytes;
@@ -36,6 +37,7 @@ use crate::session::step_context::StepContext;
use crate::tools::context::ToolCallSource;
use crate::tools::context::ToolOutput;
use crate::tools::context::ToolPayload;
use crate::tools::metadata_metrics;
use crate::tools::router::ToolCall;
use crate::utils::json::serialized_json_bytes;
@@ -344,6 +346,7 @@ impl ExecutedToolCalls {
let retained = self.retained_direct_metadata_bytes.load(Ordering::Relaxed);
let available = MAX_RETAINED_DIRECT_METADATA_BYTES.saturating_sub(retained);
let mut bytes = executed_tool_call_metadata_bytes(item);
let original_bytes = bytes;
if bytes > available {
item.retain_tool_resource_access();
bytes = executed_tool_call_metadata_bytes(item);
@@ -354,11 +357,23 @@ impl ExecutedToolCalls {
}
if bytes > available {
item.clear_executed_tool_calls();
metadata_metrics::record_shedding(
"direct_retained",
original_bytes,
executed_tool_call_metadata_bytes(item),
codex_otel::global().as_ref(),
);
return;
}
// All Direct attachments share the recorder lock, including cloned handles.
self.retained_direct_metadata_bytes
.store(retained + bytes, Ordering::Relaxed);
metadata_metrics::record_shedding(
"direct_retained",
original_bytes,
bytes,
codex_otel::global().as_ref(),
);
}
/// Remember IDs from response items that bypass local tool dispatch.
@@ -487,7 +502,7 @@ impl ExecutedToolCalls {
let ToolCallSource::CodeMode { cell_id, .. } = source else {
return false;
};
let metadata = codex_protocol::models::ToolResultMetadata::new(metadata);
let metadata = ToolResultMetadata::new(metadata);
let has_metadata = metadata.is_some();
let mut state = self.lock_state();
let Some(state) = state.as_mut() else {
@@ -48,7 +48,12 @@ impl ExecutedToolCalls {
return;
};
if attached || items.iter().any(has_direct_call_metadata) {
bound_executed_tool_calls_for_prompt(items);
metadata_metrics::bound_prompt_metadata(
items,
bound_executed_tool_calls_for_prompt,
"request",
codex_otel::global().as_ref(),
);
}
}
@@ -292,7 +297,12 @@ impl ExecutedToolCalls {
if metadata_bytes > MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT {
match budget_scope {
MetadataBudgetScope::AllCalls => {
bound_executed_tool_calls_for_prompt_prioritizing_recent(items);
metadata_metrics::bound_prompt_metadata(
items,
bound_executed_tool_calls_for_prompt_prioritizing_recent,
"retained",
codex_otel::global().as_ref(),
);
}
MetadataBudgetScope::CodeModeOnly => {
// Validate bindings against the complete history above, but do not make
@@ -306,7 +316,12 @@ impl ExecutedToolCalls {
})
.map(|(index, item)| (index, item.clone()))
.unzip();
bound_executed_tool_calls_for_prompt_prioritizing_recent(&mut bounded);
metadata_metrics::bound_prompt_metadata(
&mut bounded,
bound_executed_tool_calls_for_prompt_prioritizing_recent,
"compaction",
codex_otel::global().as_ref(),
);
for (index, item) in indices.into_iter().zip(bounded) {
items[index] = item;
}
@@ -127,13 +127,14 @@ fn compaction_attaches_pending_and_retained_code_mode_metadata() {
#[test]
fn failed_compaction_preserves_observations_absent_from_trimmed_retries() {
let call_count = 100;
for rebudget in [false, true] {
let recorder = new_recorder(InitialHistory::New);
let mut history = Vec::new();
for origin in ["old", "recent"] {
let cell = CellId::new(origin.to_string());
recorder.start_cell(&cell, origin);
for index in 0..6 {
for index in 0..call_count {
record_nested_call(&recorder, &cell, &format!("nested-{index}"));
}
recorder.finish_cell_recording(&cell);
@@ -150,7 +151,7 @@ fn failed_compaction_preserves_observations_absent_from_trimmed_retries() {
cell_id: "recent".to_string(),
runtime_tool_call_id: "runtime".to_string(),
};
for index in 0..6 {
for index in 0..call_count {
assert!(recorder.record_tool_result_metadata(
&source,
&format!("nested-{index}"),
@@ -163,15 +164,24 @@ fn failed_compaction_preserves_observations_absent_from_trimmed_retries() {
assert_eq!(tool_calls_complete(&trimmed_attempt[1]), Some(true));
if rebudget {
let encoded = serde_json::to_value(&trimmed_attempt[1]).unwrap();
assert_eq!(
encoded["internal_chat_message_metadata_passthrough"]["executed_tool_calls"][0]["tool_result_metadata"],
json!("omitted_due_to_size_limit"),
assert!(
encoded["internal_chat_message_metadata_passthrough"]["executed_tool_calls"][0]["tool_result_metadata"]
.as_str()
.is_some_and(|marker| marker.starts_with("omitted_due_to_size_limit (overage_bytes=")),
);
}
// Exhausted retries leave the session's original history intact.
let mut expected = full_attempt;
expected[2..].clone_from_slice(&trimmed_attempt);
// The restored inventory must fit without causing another round of raw-result shedding.
assert!(
expected
.iter()
.map(executed_tool_call_metadata_bytes)
.sum::<usize>()
<= 2 * 1024 * 1024
);
let mut resumed = history.clone();
recorder.attach_to_prompt(&mut resumed, &mut HashMap::new());
assert_eq!(resumed, expected);
@@ -189,11 +199,14 @@ fn compaction_bounds_code_mode_metadata_without_rebudgeting_direct_history() {
let mut history = Vec::new();
let arguments = json!({"value": "x".repeat(7 * 1024)});
let argument_bytes = serialized_json_bytes(&arguments).unwrap();
for index in 0..24 {
for index in 0..96 {
let origin = format!("exec-{index}");
let cell = CellId::new(format!("cell-{index}"));
recorder.start_cell(&cell, &origin);
let call = ExecutedToolCall::new("nested_tool".to_string(), arguments.clone());
let mut call = ExecutedToolCall::new("nested_tool".to_string(), arguments.clone());
call.set_tool_result_metadata(ToolResultMetadata::new(&json!({
"provider": "x".repeat(24 * 1024),
})));
recorder.record_nested_tool_call(
cell.clone(),
format!("nested-{index}"),
@@ -222,25 +235,35 @@ fn compaction_bounds_code_mode_metadata_without_rebudgeting_direct_history() {
> 128 * 1024
);
let mut code_mode_only = history[..direct_start].to_vec();
recorder.attach_to_compaction_prompt(&mut code_mode_only);
let mut compact = history.clone();
recorder.attach_to_compaction_prompt(&mut compact);
assert_eq!(&compact[direct_start..], &history[direct_start..]);
assert_eq!(&compact[..direct_start], &code_mode_only);
assert!(
compact[..direct_start]
.iter()
.map(executed_tool_call_metadata_bytes)
.sum::<usize>()
<= 128 * 1024
<= 2 * 1024 * 1024
);
let mut newest = exec_output("exec-23");
newest.append_executed_tool_calls(vec![ExecutedToolCall::new(
"nested_tool".to_string(),
arguments,
)]);
newest.set_tool_call_cell_id("exec-23");
let mut newest = exec_output("exec-95");
let mut newest_call = ExecutedToolCall::new("nested_tool".to_string(), arguments);
newest_call.set_tool_result_metadata(ToolResultMetadata::new(&json!({
"provider": "x".repeat(24 * 1024),
})));
newest.append_executed_tool_calls(vec![newest_call]);
newest.set_tool_call_cell_id("exec-95");
newest.mark_tool_calls_complete();
assert_eq!(compact[direct_start - 1], newest);
assert_ne!(tool_calls_complete(&compact[1]), Some(true));
let oldest = serde_json::to_value(&compact[1]).unwrap();
assert!(
oldest["internal_chat_message_metadata_passthrough"]["executed_tool_calls"][0]["tool_result_metadata"]
.as_str()
.is_some_and(|marker| marker.starts_with("omitted_due_to_size_limit (overage_bytes=")),
);
assert_eq!(tool_calls_complete(&compact[1]), Some(true));
let mut retry = history;
recorder.attach_to_compaction_prompt(&mut retry);
assert_eq!(retry, compact);
@@ -609,42 +632,48 @@ fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() {
fn executed_tool_call_recorder_bounds_retained_history_and_keeps_latest_calls() {
let recorder = new_recorder(InitialHistory::Forked(Vec::new()));
let mut history = Vec::new();
let arguments = serde_json::to_string(&json!({ "payload": "x".repeat(1024) }))
.expect("tool arguments must serialize");
let arguments = json!({ "payload": "x".repeat(7 * 1024) });
let argument_bytes = serialized_json_bytes(&arguments).expect("tool arguments must serialize");
let cell_count = 100;
let calls_per_cell = 4;
assert!(argument_bytes < MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES);
assert!(
argument_bytes * calls_per_cell < MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT
);
assert!(argument_bytes * calls_per_cell * cell_count > 2 * 1024 * 1024);
let mut prompt = Vec::new();
for index in 0..512 {
for index in 0..cell_count {
let call_id = format!("retained-{index}");
let cell_id = CellId::new(format!("retained-cell-{index}"));
recorder.start_cell(&cell_id, &call_id);
recorder.record_nested_tool_call(
cell_id.clone(),
format!("nested-{index}"),
ExecutedToolCall::new(
format!("retained_tool_{index}"),
json!({ "payload": "x".repeat(1024) }),
),
arguments.len(),
);
let expected_calls = (0..calls_per_cell)
.map(|nested_index| {
ExecutedToolCall::new(
format!("retained_tool_{index}_{nested_index}"),
arguments.clone(),
)
})
.collect::<Vec<_>>();
for (nested_index, call) in expected_calls.iter().enumerate() {
recorder.record_nested_tool_call(
cell_id.clone(),
format!("nested-{index}-{nested_index}"),
call.clone(),
argument_bytes,
);
}
recorder.finish_cell_recording(&cell_id);
history.extend([exec_input(&call_id), exec_output(&call_id)]);
prompt = history.clone();
assert!(recorder.attach_pending_to_prompt(&mut prompt, &mut HashMap::new()));
codex_protocol::models::bound_executed_tool_calls_for_prompt(&mut prompt);
let latest_call = prompt
let latest_calls = prompt
.last()
.and_then(ResponseItem::executed_tool_call_metadata)
.and_then(|metadata| metadata.executed_tool_calls.as_ref())
.and_then(|calls| calls.first())
.map(serde_json::to_value)
.transpose()
.expect("latest tool call must serialize")
.expect("latest tool call must remain in retained metadata");
assert_eq!(latest_call["name"], format!("retained_tool_{index}"));
assert_eq!(
latest_call["arguments"],
json!({ "payload": "x".repeat(1024) }),
);
.expect("latest tool calls must remain in retained metadata");
assert_eq!(latest_calls, &expected_calls);
}
let state = recorder.lock_state();
@@ -655,7 +684,7 @@ fn executed_tool_call_recorder_bounds_retained_history_and_keeps_latest_calls()
.map(|retained| serialized_json_bytes(&retained.calls))
.sum::<serde_json::Result<usize>>()
.expect("retained calls must serialize");
assert!(retained_bytes <= 128 * 1024);
assert!(retained_bytes <= 2 * 1024 * 1024);
let metadata = prompt
.iter()
@@ -664,10 +693,11 @@ fn executed_tool_call_recorder_bounds_retained_history_and_keeps_latest_calls()
.flatten()
.map(|call| serde_json::to_value(call).expect("retained call must serialize"))
.collect::<Vec<_>>();
assert!(metadata.len() < 512);
assert!(metadata.len() < cell_count * calls_per_cell);
// Dropped outputs do not add their omission counts to unrelated surviving calls.
assert!(metadata.iter().all(|call| {
call["arguments"]["_codex_executed_tool_call_truncated"]["omitted_calls"].is_null()
let omitted = &call["arguments"]["_codex_executed_tool_call_truncated"]["omitted_calls"];
omitted.is_null() || *omitted == json!(calls_per_cell - 1)
}));
}
@@ -943,7 +973,7 @@ fn tool_call_completeness_survives_waits_without_changing_deltas() {
#[test]
fn unrelated_raw_overflow_retains_a_complete_exec_within_the_request_budget() {
let request_budget = 128 * 1024;
let request_budget = 2 * 1024 * 1024;
let recorder = new_recorder(InitialHistory::New);
let cell = CellId::new("clean-overflow-cell".to_string());
recorder.start_cell(&cell, "exec-clean");
@@ -954,7 +984,7 @@ fn unrelated_raw_overflow_retains_a_complete_exec_within_the_request_budget() {
let arguments = json!("x".repeat(7 * 1024));
assert!(serialized_json_bytes(&arguments).unwrap() <= MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES);
unrelated.append_executed_tool_calls(
(0..24)
(0..request_budget / (7 * 1024) + 1)
.map(|index| ExecutedToolCall::new(format!("direct-{index}"), arguments.clone()))
.collect(),
);
@@ -1210,7 +1240,7 @@ fn request_truncation_prevents_completion_after_compaction() {
}
let mut initial = vec![exec_input("exec"), exec_output("exec")];
for index in 0..3 {
for index in 0..63 {
let origin = format!("other-{index}");
let other_cell = CellId::new(origin.clone());
recorder.start_cell(&other_cell, &origin);
@@ -1261,7 +1291,7 @@ fn request_truncation_prevents_completion_after_compaction() {
}
#[test]
fn result_metadata_updates_the_exact_retained_call_and_marks_oversized_retries() {
fn result_metadata_updates_the_exact_retained_call_and_preserves_large_retries() {
let recorder = new_recorder(InitialHistory::Forked(Vec::new()));
let cell_id = CellId::new("metadata-cell".to_string());
recorder.start_cell(&cell_id, "exec");
@@ -1298,14 +1328,15 @@ fn result_metadata_updates_the_exact_retained_call_and_marks_oversized_retries()
assert!(recorder.attach_pending_to_prompt(&mut retry, &mut retry_cache));
assert_eq!(retry, [expected]);
let large_metadata = json!({ "oversized": "x".repeat(40 * 1024) });
assert!(recorder.record_tool_result_metadata(
&source(cell_id.as_str()),
"first",
&json!({ "oversized": "x".repeat(32 * 1024) }),
&large_metadata,
));
let mut expected = serde_json::to_value(&initial[1..]).unwrap();
expected[0]["internal_chat_message_metadata_passthrough"]["executed_tool_calls"][0]["tool_result_metadata"] =
json!("omitted_due_to_size_limit");
large_metadata;
let mut retry = [exec_output("exec")];
assert!(recorder.attach_pending_to_prompt(&mut retry, &mut retry_cache));
assert_eq!(serde_json::to_value(&retry).unwrap(), expected);
@@ -1323,12 +1354,22 @@ fn untracked_history_budget_loss_keeps_later_wait_incomplete() {
/*original_bytes*/ 8192,
);
let mut prompt = vec![exec_input("exec"), exec_output("exec")];
for index in 0..256 {
let untracked_count = 256;
let per_output_budget = 2 * 1024 * 1024 / untracked_count;
let mut empty_output = output("untracked");
empty_output.append_executed_tool_calls(vec![ExecutedToolCall::new(
"other_tool".to_string(),
json!(""),
)]);
// Fill the budget with intact unmanaged outputs, leaving no room for the tracked record.
let payload_bytes = per_output_budget - executed_tool_call_metadata_bytes(&empty_output);
for index in 0..untracked_count {
let mut item = output(&format!("untracked-{index}"));
item.append_executed_tool_calls(vec![ExecutedToolCall::new(
"other_tool".to_string(),
json!("x".repeat(1024)),
json!("x".repeat(payload_bytes)),
)]);
assert_eq!(executed_tool_call_metadata_bytes(&item), per_output_budget);
prompt.push(item);
}
recorder.attach_to_prompt(&mut prompt, &mut HashMap::new());
@@ -1348,11 +1389,12 @@ fn result_metadata_shedding_preserves_completion_after_compaction() {
let cell_id = CellId::new("compacted-cell".to_string());
recorder.start_cell(&cell_id, "exec");
let argument_bytes = MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT / 5 - 256;
let call_count = 79;
let argument_bytes = MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT / call_count - 64;
let arguments = json!({
"payload": "x".repeat(argument_bytes - r#"{"payload":""}"#.len()),
});
for index in 0..5 {
for index in 0..call_count {
recorder.record_nested_tool_call(
cell_id.clone(),
format!("nested-{index}"),
@@ -1371,7 +1413,7 @@ fn result_metadata_shedding_preserves_completion_after_compaction() {
);
// Each snapshot fits alone; together they push the calls over the aggregate budget.
let large_metadata = json!({ "provider_data": "x".repeat(26 * 1024) });
for index in [0, 2, 3, 4] {
for index in (0..call_count).filter(|index| *index != 1) {
assert!(recorder.record_tool_result_metadata(
&source,
&format!("nested-{index}"),
@@ -1390,14 +1432,20 @@ fn result_metadata_shedding_preserves_completion_after_compaction() {
.executed_tool_call_metadata()
.and_then(|metadata| metadata.executed_tool_calls.as_ref())
.expect("metadata shedding must preserve recorded calls");
let mut expected_calls = vec![ExecutedToolCall::new("nested_tool".to_string(), arguments); 5];
for call in &mut expected_calls[2..] {
let mut expected_calls =
vec![ExecutedToolCall::new("nested_tool".to_string(), arguments); call_count];
for call in &mut expected_calls {
call.set_tool_result_metadata(ToolResultMetadata::new(&large_metadata));
}
expected_calls[0]
.set_tool_result_metadata(ToolResultMetadata::new(&json!("omitted_due_to_size_limit")));
expected_calls[1]
.set_tool_result_metadata(ToolResultMetadata::new(&json!({ "status": "unavailable" })));
let mut unbounded = exec_output("exec");
unbounded.append_executed_tool_calls(expected_calls.clone());
unbounded.set_tool_call_cell_id("exec");
let overage = executed_tool_call_metadata_bytes(&unbounded) - 2 * 1024 * 1024;
expected_calls[0].set_tool_result_metadata(ToolResultMetadata::new(&json!(format!(
"omitted_due_to_size_limit (overage_bytes={overage})"
))));
assert_eq!(calls, &expected_calls);
let mut retry = [exec_output("exec")];
assert!(recorder.attach_pending_to_prompt(&mut retry, &mut retry_cache));
@@ -0,0 +1,61 @@
//! Measures bytes shed by existing metadata budgets without changing their policy.
//! Samples describe budget applications, not unique tool calls or full request sizes.
use codex_otel::MetricsClient;
use codex_protocol::models::ResponseItem;
use codex_protocol::models::executed_tool_call_metadata_bytes;
const SHED_BYTES_METRIC: &str = "codex.tool_metadata.shed_bytes";
const BYTE_BUCKETS: &[f64] = &[
1024.0, 4096.0, 16384.0, 32768.0, 65536.0, 131072.0, 262144.0, 524288.0, 1048576.0, 2097152.0,
4194304.0, 8388608.0, 16777216.0,
];
pub(super) fn bound_prompt_metadata(
items: &mut [ResponseItem],
bound: fn(&mut [ResponseItem]),
stage: &'static str,
metrics: Option<&MetricsClient>,
) {
let before = metrics.map(|_| metadata_bytes(items));
bound(items);
if let Some(before) = before {
// Measure only the budget operation: attaching observations can itself add bytes.
record_shedding(stage, before, metadata_bytes(items), metrics);
}
}
pub(crate) fn metadata_bytes(items: &[ResponseItem]) -> usize {
items.iter().fold(0_usize, |bytes, item| {
bytes.saturating_add(executed_tool_call_metadata_bytes(item))
})
}
pub(crate) fn record_shedding(
stage: &'static str,
before: usize,
after: usize,
metrics: Option<&MetricsClient>,
) {
let Some(metrics) = metrics else {
return;
};
let shed = before.saturating_sub(after);
if shed == 0 || before == usize::MAX || after == usize::MAX {
return;
}
// Count the complete attempted-tool metadata envelope, including argument and
// inventory truncation performed by retention/request budgets.
if let Err(error) = metrics.histogram_with_boundaries(
SHED_BYTES_METRIC,
i64::try_from(shed).unwrap_or(i64::MAX),
BYTE_BUCKETS,
&[("stage", stage)],
) {
tracing::warn!("tool metadata shedding metric failed: {error}");
}
}
#[cfg(test)]
#[path = "metadata_metrics_tests.rs"]
mod tests;
@@ -0,0 +1,131 @@
use super::*;
use codex_otel::MetricsConfig;
use codex_protocol::models::ExecutedToolCall;
use codex_protocol::models::FunctionCallOutputPayload;
use codex_protocol::models::ResponseInputItem;
use codex_protocol::models::ToolResultMetadata;
use codex_protocol::models::bound_executed_tool_calls_for_prompt;
use codex_protocol::models::bound_executed_tool_calls_for_prompt_prioritizing_recent;
use opentelemetry_sdk::metrics::InMemoryMetricExporter;
use opentelemetry_sdk::metrics::data::AggregatedMetrics;
use opentelemetry_sdk::metrics::data::MetricData;
use opentelemetry_sdk::metrics::data::ScopeMetrics;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::collections::BTreeMap;
fn metrics() -> MetricsClient {
MetricsClient::new(
MetricsConfig::in_memory(
"test",
"codex-core",
env!("CARGO_PKG_VERSION"),
InMemoryMetricExporter::default(),
)
.with_runtime_reader(),
)
.expect("in-memory metrics")
}
fn samples(metrics: &MetricsClient) -> BTreeMap<String, (u64, f64)> {
let snapshot = metrics.snapshot().expect("metrics snapshot");
let mut samples = BTreeMap::new();
for metric in snapshot.scope_metrics().flat_map(ScopeMetrics::metrics) {
assert_eq!(metric.name(), SHED_BYTES_METRIC);
let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = metric.data() else {
panic!("expected shed-byte histogram");
};
for point in histogram.data_points() {
let attributes = point.attributes().collect::<Vec<_>>();
assert_eq!(attributes.len(), 1);
assert_eq!(attributes[0].key.as_str(), "stage");
samples.insert(
attributes[0].value.as_str().to_string(),
(point.count(), point.sum()),
);
}
}
samples
}
#[test]
fn prompt_reports_each_budget_application_without_changing_output() {
let metrics = metrics();
let original = (0..96)
.map(|index| {
let mut item = ResponseItem::from(ResponseInputItem::FunctionCallOutput {
call_id: format!("call_{index}"),
output: FunctionCallOutputPayload::from_text("unchanged output".to_string()),
});
let mut call = ExecutedToolCall::new("mcp__apps__read".to_string(), json!({}));
call.set_tool_result_metadata(ToolResultMetadata::new(&json!({
"openai/resource_access": {"resources": ["x".repeat(24_000)]},
})));
item.append_executed_tool_calls(vec![call]);
item.mark_tool_calls_complete();
item
})
.collect::<Vec<_>>();
let mut expected_samples = BTreeMap::new();
for (stage, bound) in [
(
"request",
bound_executed_tool_calls_for_prompt as fn(&mut [ResponseItem]),
),
(
"retained",
bound_executed_tool_calls_for_prompt_prioritizing_recent,
),
(
"compaction",
bound_executed_tool_calls_for_prompt_prioritizing_recent,
),
] {
let mut expected = original.clone();
bound(&mut expected);
let shed = metadata_bytes(&original) - metadata_bytes(&expected);
assert!(shed > 0);
let mut actual = original.clone();
bound_prompt_metadata(&mut actual, bound, stage, Some(&metrics));
assert_eq!(actual, expected);
// Reapplying an already-satisfied budget must not add a zero-byte sample.
bound_prompt_metadata(&mut actual, bound, stage, Some(&metrics));
assert_eq!(actual, expected);
let mut without_metrics = original.clone();
bound_prompt_metadata(&mut without_metrics, bound, stage, /*metrics*/ None);
assert_eq!(without_metrics, expected);
expected_samples.insert(stage.to_string(), (1, shed as f64));
}
assert_eq!(samples(&metrics), expected_samples);
}
#[test]
fn direct_retention_reports_only_positive_known_shedding() {
let metrics = metrics();
for (before, after) in [
(0, 0),
(10, 10),
(10, 20),
(usize::MAX, 0),
(10, usize::MAX),
] {
record_shedding("direct_retained", before, after, Some(&metrics));
}
record_shedding(
"direct_retained",
/*before*/ 100,
/*after*/ 1,
/*metrics*/ None,
);
assert_eq!(samples(&metrics), BTreeMap::new());
record_shedding(
"direct_retained",
/*before*/ 100,
/*after*/ 27,
Some(&metrics),
);
assert_eq!(
samples(&metrics),
BTreeMap::from([("direct_retained".to_string(), (1, 73.0))]),
);
}
+1
View File
@@ -10,6 +10,7 @@ pub(crate) mod handlers;
pub(crate) mod hook_names;
pub(crate) mod hosted_spec;
pub(crate) mod lifecycle;
pub(crate) mod metadata_metrics;
mod multi_agent_tool;
pub(crate) mod network_approval;
pub(crate) mod orchestrator;
+15 -20
View File
@@ -2026,7 +2026,7 @@ async fn direct_result_metadata_retained_budget_preserves_resource_access() -> R
"openai/resource_access": resource_access,
});
let metadata_bytes = serde_json::to_vec(&result_metadata)?.len();
assert!(metadata_bytes < 32 * 1024);
assert!(metadata_bytes * 32 < 1024 * 1024);
assert!(metadata_bytes * 33 > 1024 * 1024);
let apps_server = mount_result_metadata_app(
&server,
@@ -2105,7 +2105,7 @@ async fn direct_result_metadata_retained_budget_preserves_resource_access() -> R
assert_eq!(call["params"]["name"], RESULT_METADATA_TOOL);
assert_eq!(call["params"]["arguments"], arguments);
}
// The request-metadata helper reapplies the smaller 128 KiB budget, hiding this boundary.
// Inspect retained history directly to isolate the Direct admission budget.
let history = test.codex.conversation_history_snapshot().await;
assert!(
history
@@ -2156,7 +2156,9 @@ async fn direct_result_metadata_retained_budget_preserves_resource_access() -> R
#[test_case(ToolMode::Direct; "direct")]
#[test_case(ToolMode::CodeModeOnly; "code_mode")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn result_metadata_budget_preserves_smaller_results(tool_mode: ToolMode) -> Result<()> {
async fn result_metadata_preserves_results_within_request_budget(
tool_mode: ToolMode,
) -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let arguments = [
@@ -2174,7 +2176,6 @@ async fn result_metadata_budget_preserves_smaller_results(tool_mode: ToolMode) -
"resource_coverage": "complete",
"resources": [],
});
let resource_only = serde_json::json!({ "openai/resource_access": resource_access });
let result_metadata = [
serde_json::json!({
"payload": "l".repeat(31 * 1024),
@@ -2187,7 +2188,11 @@ async fn result_metadata_budget_preserves_smaller_results(tool_mode: ToolMode) -
serde_json::json!({ "payload": "m".repeat(20 * 1024) }),
serde_json::json!({
"payload": "o".repeat(40 * 1024),
"openai/resource_access": resource_access,
"openai/resource_access": {
"schema_version": 1,
"resource_coverage": "complete",
"resources": ["r".repeat(40 * 1024)],
},
}),
serde_json::json!({ "status": "ok" }),
];
@@ -2195,10 +2200,9 @@ async fn result_metadata_budget_preserves_smaller_results(tool_mode: ToolMode) -
.iter()
.map(|metadata| serde_json::to_vec(metadata).unwrap().len())
.collect::<Vec<_>>();
assert!(metadata_sizes[..6].iter().all(|bytes| *bytes < 32 * 1024));
assert!(metadata_sizes[..6].iter().sum::<usize>() > 128 * 1024);
assert!(metadata_sizes[6] > 32 * 1024);
assert!(metadata_sizes[7] < 32 * 1024);
assert!(metadata_sizes.iter().sum::<usize>() > 128 * 1024);
assert!(metadata_sizes.iter().sum::<usize>() < 1024 * 1024);
assert!(serde_json::to_vec(&result_metadata[6]["openai/resource_access"])?.len() > 32 * 1024);
for (arguments, metadata) in arguments.iter().zip(&result_metadata) {
let result = serde_json::json!({
"content": [{ "type": "text", "text": RESULT_METADATA_PRIVATE_RESULT }],
@@ -2335,7 +2339,7 @@ async fn result_metadata_budget_preserves_smaller_results(tool_mode: ToolMode) -
.iter()
.map(codex_protocol::models::executed_tool_call_metadata_bytes)
.sum::<usize>()
<= 128 * 1024
<= 2 * 1024 * 1024
);
let captured = serde_json::to_value(captured)?;
for (input, expected_metadata) in [
@@ -2343,16 +2347,7 @@ async fn result_metadata_budget_preserves_smaller_results(tool_mode: ToolMode) -
(request.input(), None),
(
captured.as_array().unwrap().clone(),
Some([
resource_only.clone(),
result_metadata[1].clone(),
result_metadata[2].clone(),
result_metadata[3].clone(),
result_metadata[4].clone(),
result_metadata[5].clone(),
resource_only.clone(),
result_metadata[7].clone(),
]),
Some(result_metadata.clone()),
),
] {
let mut calls = Vec::new();
@@ -1,4 +1,4 @@
//! Direct-call metadata coverage, including malformed calls and request-budget pruning.
//! Direct-call metadata coverage, including malformed calls and metadata budgets.
use anyhow::Result;
use codex_features::Feature;
@@ -203,7 +203,7 @@ async fn direct_call_metadata_during_compaction_respects_provider_support(
}
#[test_case(false, 0; "metadata disabled")]
#[test_case(true, 24; "request budget exceeded")]
#[test_case(true, 24; "above previous request budget")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn direct_function_and_tool_search_mark_complete_attempts(
metadata_enabled: bool,
@@ -211,7 +211,7 @@ async fn direct_function_and_tool_search_mark_complete_attempts(
) -> Result<()> {
skip_if_no_network!(Ok(()));
let request_budget = 128 * 1024;
let request_budget = 2 * 1024 * 1024;
let server = start_mock_server().await;
let mut builder = test_codex().with_config(move |config| {
configure_search_capable_model(config);
@@ -226,13 +226,14 @@ async fn direct_function_and_tool_search_mark_complete_attempts(
let valid_arguments = json!({"plan": [{"step": "read", "status": "in_progress"}]});
let malformed_arguments = "{malformed";
let search_arguments = json!({"query": "nonexistent-completeness-proof-tool", "limit": null});
// The bulk calls fit the per-call limit; their combined metadata forces pruning.
// These calls exceed the old request budget but stay below Direct's retained-history limit.
let budget_arguments =
json!({"plan": [{"step": "x".repeat(7 * 1024), "status": "in_progress"}]});
let budget_arguments_json = budget_arguments.to_string();
assert!(budget_arguments_json.len() < 8 * 1024);
if budget_calls > 0 {
assert!(budget_calls * budget_arguments_json.len() > request_budget);
assert!(budget_calls * budget_arguments_json.len() > 128 * 1024);
assert!(budget_calls * budget_arguments_json.len() < 1024 * 1024);
}
let mut events = vec![ev_response_created("resp-1")];
if budget_calls > 0 {
@@ -322,36 +323,22 @@ async fn direct_function_and_tool_search_mark_complete_attempts(
);
assert!(metadata.get("tool_calls_complete").is_none());
}
let mut truncated = 0;
for index in 0..budget_calls {
let output = request.function_call_output(&format!("plan-budget-{index}"));
assert_eq!(output["output"], json!("Plan updated"));
let metadata = tool_call_metadata(output);
metadata_bytes += serde_json::to_vec(&metadata)?.len();
let calls = metadata["executed_tool_calls"]
.as_array()
.expect("recorded calls");
assert_eq!(calls.len(), 1);
if calls[0]["arguments"]
.get("_codex_executed_tool_call_truncated")
.is_some()
{
truncated += 1;
assert!(metadata.get("tool_calls_complete").is_none());
} else {
assert_eq!(
metadata,
json!({
"executed_tool_calls": [{"name": "update_plan", "arguments": budget_arguments}],
"tool_calls_complete": true,
})
);
}
assert_eq!(
metadata,
json!({
"executed_tool_calls": [{"name": "update_plan", "arguments": budget_arguments}],
"tool_calls_complete": true,
})
);
}
assert!(metadata_bytes <= request_budget);
assert!(
budget_calls == 0 || truncated > 0,
"request must exercise pruning"
);
if budget_calls > 0 {
assert!(metadata_bytes > 128 * 1024);
}
Ok(())
}
@@ -1,4 +1,3 @@
use std::collections::BTreeMap;
use std::collections::HashSet;
use schemars::JsonSchema;
@@ -15,9 +14,7 @@ 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 = 128 * 1024;
/// Maximum serialized raw metadata captured from one tool result.
const MAX_TOOL_RESULT_METADATA_BYTES: usize = 32 * 1024;
const MAX_EXECUTED_TOOL_CALL_METADATA_BYTES: usize = 2 * 1024 * 1024;
const RESOURCE_ACCESS_METADATA_KEY: &str = "openai/resource_access";
const EXECUTED_TOOL_CALL_METADATA_FIELD_BYTES: usize = b"\"executed_tool_calls\":".len();
const INTERNAL_CHAT_MESSAGE_METADATA_PASSTHROUGH_FIELD_BYTES: usize =
@@ -135,8 +132,6 @@ fn bound_executed_tool_calls_for_prompt_with_priority(
}
// Raw result metadata must not displace existing source evidence, calls or completion proof.
// Shed the largest values first so one large result does not discard unrelated small results.
let omitted = ToolResultMetadata::omitted_due_to_size_limit();
let omitted_bytes = serde_json::to_vec(&omitted).map_or(usize::MAX, |value| value.len());
let mut result_metadata = Vec::new();
for (item_index, item) in items.iter_mut().enumerate() {
if let Some(metadata) = item
@@ -182,10 +177,13 @@ fn bound_executed_tool_calls_for_prompt_with_priority(
total_metadata_bytes =
total_metadata_bytes.saturating_sub(bytes.saturating_sub(retained_bytes));
*bytes = retained_bytes;
} else if *bytes > omitted_bytes {
**metadata = omitted.clone();
total_metadata_bytes = total_metadata_bytes.saturating_sub(*bytes - omitted_bytes);
*bytes = omitted_bytes;
} else {
let retained_bytes = metadata.omit_if_smaller(
*bytes,
total_metadata_bytes - MAX_EXECUTED_TOOL_CALL_METADATA_BYTES,
);
total_metadata_bytes = total_metadata_bytes.saturating_sub(*bytes - retained_bytes);
*bytes = retained_bytes;
}
}
// Resource evidence is optional too; it must not displace the call inventory.
@@ -197,17 +195,18 @@ fn bound_executed_tool_calls_for_prompt_with_priority(
break;
}
// An omission marker can be larger than a small object, including an empty `_meta`.
if *bytes > omitted_bytes {
**metadata = omitted.clone();
total_metadata_bytes = total_metadata_bytes.saturating_sub(*bytes - omitted_bytes);
*bytes = omitted_bytes;
}
let retained_bytes = metadata.omit_if_smaller(
*bytes,
total_metadata_bytes - MAX_EXECUTED_TOOL_CALL_METADATA_BYTES,
);
total_metadata_bytes = total_metadata_bytes.saturating_sub(*bytes - retained_bytes);
*bytes = retained_bytes;
}
// Use the updated sizes: markers should be removed before smaller provider metadata.
result_metadata.sort_by_key(|(bytes, order, call_index, metadata)| {
(
std::cmp::Reverse(*bytes),
metadata.0 != omitted.0,
!metadata.is_omitted_due_to_size_limit(),
*order,
*call_index,
)
@@ -340,23 +339,9 @@ impl std::fmt::Debug for ToolResultMetadata {
}
impl ToolResultMetadata {
/// Bounds serialization before cloning, retaining resource evidence if the full result is too large.
/// Captures the complete result; retention and outgoing-request budgets apply later.
pub fn new(metadata: &serde_json::Value) -> Self {
let mut limit = MetadataSizeLimit(MAX_TOOL_RESULT_METADATA_BYTES);
if serde_json::to_writer(&mut limit, metadata).is_ok() {
return Self(Some(metadata.clone()));
}
if let Some(resource_access) = metadata.get(RESOURCE_ACCESS_METADATA_KEY) {
let resource_metadata =
BTreeMap::from([(RESOURCE_ACCESS_METADATA_KEY, resource_access)]);
let mut limit = MetadataSizeLimit(MAX_TOOL_RESULT_METADATA_BYTES);
if serde_json::to_writer(&mut limit, &resource_metadata).is_ok() {
return Self(Some(serde_json::json!({
(RESOURCE_ACCESS_METADATA_KEY): resource_access,
})));
}
}
Self::omitted_due_to_size_limit()
Self(Some(metadata.clone()))
}
/// Leaves generic metadata unchanged when it has no resource-access field.
@@ -371,39 +356,48 @@ impl ToolResultMetadata {
true
}
fn omitted_due_to_size_limit() -> Self {
fn omitted_due_to_size_limit(overage_bytes: usize) -> Self {
// MCP `_meta` is an object; this string is a harness omission status, not provider data.
Self(Some(serde_json::Value::String(
"omitted_due_to_size_limit".to_string(),
)))
Self(Some(serde_json::Value::String(format!(
"omitted_due_to_size_limit (overage_bytes={overage_bytes})"
))))
}
fn is_omitted_due_to_size_limit(&self) -> bool {
let Some(serde_json::Value::String(value)) = self.0.as_ref() else {
return false;
};
value == "omitted_due_to_size_limit"
|| value
.strip_prefix("omitted_due_to_size_limit (overage_bytes=")
.and_then(|value| value.strip_suffix(')'))
.is_some_and(|value| value.parse::<usize>().is_ok())
}
fn omit_if_smaller(&mut self, original_bytes: usize, overage_bytes: usize) -> usize {
if self.is_none() || self.is_omitted_due_to_size_limit() {
return original_bytes;
}
let omitted = Self::omitted_due_to_size_limit(overage_bytes);
let omitted_bytes = serde_json::to_vec(&omitted).map_or(usize::MAX, |value| value.len());
if omitted_bytes < original_bytes {
*self = omitted;
omitted_bytes
} else {
original_bytes
}
}
fn is_none(&self) -> bool {
self.0.is_none()
}
/// Whether the bounded snapshot contains metadata or an omission marker.
/// Whether the captured snapshot contains metadata or an omission marker.
pub fn is_some(&self) -> bool {
self.0.is_some()
}
}
struct MetadataSizeLimit(usize);
impl std::io::Write for MetadataSizeLimit {
fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
self.0 = self
.0
.checked_sub(bytes.len())
.ok_or_else(|| std::io::Error::other("tool result metadata exceeds the byte limit"))?;
Ok(bytes.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
/// A bounded capture update. Omitted updates still clear any previously recorded evidence.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolResultSources(Option<Vec<ToolResultSource>>);
@@ -494,7 +488,7 @@ impl ExecutedToolCall {
self.tool_result_sources.is_some()
}
/// Replaces the entire `_meta` snapshot, including with an omission marker if oversized.
/// Replaces the entire captured `_meta` snapshot.
pub fn set_tool_result_metadata(&mut self, metadata: ToolResultMetadata) {
self.tool_result_metadata = metadata;
}
@@ -614,6 +608,17 @@ impl ResponseItem {
result_metadata(self).eq(result_metadata(other))
}
/// Whether an output carries raw result metadata or a harness omission marker.
pub fn has_tool_result_metadata(&self) -> bool {
self.executed_tool_call_metadata().is_some_and(|metadata| {
metadata
.executed_tool_calls
.iter()
.flatten()
.any(|call| call.tool_result_metadata.is_some())
})
}
/// Omits raw tool results without changing existing call, source or completion metadata.
pub fn clear_tool_result_metadata(&mut self) {
if let Some(metadata) = self
@@ -626,6 +631,36 @@ impl ResponseItem {
}
}
/// Marks larger raw results with the full request's overage, preserving existing markers.
pub fn omit_tool_result_metadata(&mut self, overage_bytes: usize) {
if let Some(metadata) = self
.internal_chat_message_metadata_passthrough_mut()
.and_then(Option::as_mut)
{
for call in metadata.executed_tool_calls.iter_mut().flatten() {
let result = &mut call.tool_result_metadata;
let bytes = serde_json::to_vec(result).map_or(usize::MAX, |value| value.len());
result.omit_if_smaller(bytes, overage_bytes);
}
}
}
/// Preserves complete resource evidence, marking other large results with request overage.
pub fn retain_tool_resource_access_or_omit_metadata(&mut self, overage_bytes: usize) {
if let Some(metadata) = self
.internal_chat_message_metadata_passthrough_mut()
.and_then(Option::as_mut)
{
for call in metadata.executed_tool_calls.iter_mut().flatten() {
let result = &mut call.tool_result_metadata;
if !result.retain_resource_access() {
let bytes = serde_json::to_vec(result).map_or(usize::MAX, |value| value.len());
result.omit_if_smaller(bytes, overage_bytes);
}
}
}
}
/// Reduces optional results to resource evidence before considering call-inventory loss.
pub fn retain_tool_resource_access(&mut self) {
if let Some(metadata) = self
@@ -110,7 +110,7 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> {
String::new(),
serde_json::Value::Null,
)]);
item.set_tool_call_cell_id("cell-\"\\");
item.set_tool_call_cell_id(&format!("cell-\"\\{}", "x".repeat(512)));
item.mark_tool_calls_complete();
item
})
@@ -199,7 +199,8 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> {
let arguments = serde_json::json!({ "payload": "x".repeat(7 * 1024) });
let argument_bytes = serde_json::to_vec(&arguments)?.len();
assert!(argument_bytes < MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES);
assert!(24 * argument_bytes > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES);
let call_count = MAX_EXECUTED_TOOL_CALL_METADATA_BYTES / argument_bytes + 1;
assert!(call_count * argument_bytes > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES);
let mut items = ["A", "wait-A", "B"].map(|call_id| {
let mut item = output(call_id);
item.mark_tool_calls_complete();
@@ -208,7 +209,7 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> {
for item in &mut items[..2] {
item.set_tool_call_cell_id("cell-A");
}
let calls = vec![ExecutedToolCall::new("test_tool".to_string(), arguments); 24];
let calls = vec![ExecutedToolCall::new("test_tool".to_string(), arguments); call_count];
items[0].append_executed_tool_calls(calls.clone());
items[2].append_executed_tool_calls(vec![ExecutedToolCall::new(
"test_tool".to_string(),
@@ -230,7 +231,7 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> {
assert_eq!(recorded[0].name, "test_tool");
let truncation = recorded[0].truncation().unwrap();
assert_eq!(truncation.original_bytes, argument_bytes);
assert_eq!(truncation.omitted_calls, Some(23));
assert_eq!(truncation.omitted_calls, Some(call_count - 1));
assert_eq!(
bounded[1]
.executed_tool_call_metadata()
@@ -249,7 +250,7 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> {
.truncation()
.unwrap()
.omitted_calls,
Some(47),
Some(2 * call_count - 1),
);
}
@@ -266,6 +267,22 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> {
16
]);
items[1].mark_tool_calls_complete();
// Leave room for calls, but not their sources, without exceeding the argument limit.
let padding_bytes = MAX_EXECUTED_TOOL_CALL_METADATA_BYTES - metadata_bytes(&items)? - 1024;
let calls = items[0]
.ensure_tool_call_metadata()
.unwrap()
.executed_tool_calls
.as_mut()
.unwrap();
let call_count = calls.len();
for (index, call) in calls.iter_mut().enumerate() {
call.name.push_str(
&"x".repeat(
padding_bytes / call_count + usize::from(index < padding_bytes % call_count),
),
);
}
let expected = items.clone();
assert!(
first_executed_tool_call(&mut items[0])
@@ -293,7 +310,7 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> {
.internal_chat_message_metadata_passthrough_mut()
.unwrap() = metadata;
let without_marker = item.clone();
item.set_tool_call_cell_id("cell-\"\\");
item.set_tool_call_cell_id(&format!("cell-\"\\{}", "x".repeat(256)));
item.mark_tool_calls_complete();
assert_eq!(
metadata_bytes(std::slice::from_ref(&item))?,
@@ -501,7 +518,7 @@ fn tool_result_source_snapshots_replace_atomically() -> Result<()> {
}
#[test]
fn tool_result_metadata_is_host_only_bounded_and_redacted() -> Result<()> {
fn tool_result_metadata_is_host_only_and_redacted_without_a_per_result_limit() -> Result<()> {
let mut call = ExecutedToolCall::new("test_tool".to_string(), serde_json::json!({}));
let without_metadata = call.clone();
for metadata in [
@@ -510,9 +527,14 @@ fn tool_result_metadata_is_host_only_bounded_and_redacted() -> Result<()> {
"other": false,
}),
serde_json::json!({
"openai/resource_access": { "resources": ["not-for-debug"] },
"other": "preserved when the full metadata fits",
"openai/resource_access": {
"resource_coverage": "complete",
"resources": ["not-for-debug", "x".repeat(40 * 1024)],
},
"other": "preserved along with resource_access",
}),
// Capture preserves the original value even when serialization escapes increase its size.
serde_json::json!({ "value": "\0".repeat(32 * 1024) }),
serde_json::json!({}),
serde_json::json!(null),
] {
@@ -535,74 +557,80 @@ fn tool_result_metadata_is_host_only_bounded_and_redacted() -> Result<()> {
without_metadata
);
}
let exact_limit = serde_json::json!({
"value": "x".repeat(MAX_TOOL_RESULT_METADATA_BYTES - r#"{"value":""}"#.len()),
});
let capture = ToolResultMetadata::new(&exact_limit);
assert!(capture.is_some());
call.set_tool_result_metadata(capture);
assert_eq!(
serde_json::to_value(&call)?["tool_result_metadata"],
exact_limit
);
// JSON escaping counts toward the bound, not just the in-memory string length.
let oversized = serde_json::json!({ "value": "\0".repeat(MAX_TOOL_RESULT_METADATA_BYTES / 2) });
let capture = ToolResultMetadata::new(&oversized);
assert!(capture.is_some());
call.set_tool_result_metadata(capture);
let wire = serde_json::to_value(&call)?;
assert_eq!(
wire,
serde_json::json!({
"name": "test_tool",
"arguments": {},
"tool_result_metadata": "omitted_due_to_size_limit",
})
);
assert_eq!(
serde_json::from_value::<ExecutedToolCall>(wire)?,
without_metadata
);
let resource_access = serde_json::json!({
"schema_version": 1,
"resource_coverage": "incomplete",
"resources": [{ "id": "not-for-debug" }],
"reasons": ["tool_error"],
});
let oversized = serde_json::json!({
"openai/resource_access": resource_access,
"other": "x".repeat(MAX_TOOL_RESULT_METADATA_BYTES),
});
call.set_tool_result_metadata(ToolResultMetadata::new(&oversized));
let wire = serde_json::to_value(&call)?;
assert_eq!(
wire["tool_result_metadata"],
serde_json::json!({ "openai/resource_access": resource_access }),
);
assert!(!format!("{call:?}").contains("not-for-debug"));
assert_eq!(
serde_json::from_value::<ExecutedToolCall>(wire)?,
without_metadata
);
// Do not trim resource lists while preserving a provider's coverage claim.
let oversized_resources = serde_json::json!({
"openai/resource_access": {
"resource_coverage": "complete",
"resources": ["x".repeat(MAX_TOOL_RESULT_METADATA_BYTES)],
},
});
assert_eq!(
ToolResultMetadata::new(&oversized_resources),
ToolResultMetadata::omitted_due_to_size_limit(),
);
Ok(())
}
#[test]
fn wire_budget_helpers_preserve_resource_evidence_markers_and_small_values() {
let resource_only = serde_json::json!({
"openai/resource_access": { "resource_coverage": "complete", "resources": [] },
});
let mut item = output("exec");
assert!(!item.has_tool_result_metadata());
item.append_executed_tool_calls(
[
ToolResultMetadata::new(&serde_json::json!({
"openai/resource_access": resource_only["openai/resource_access"],
"payload": "x".repeat(512),
})),
ToolResultMetadata::new(&serde_json::json!({
"omitted_due_to_size_limit": { "overage_bytes": 1 },
"payload": "x".repeat(512),
})),
ToolResultMetadata::new(&serde_json::json!({})),
ToolResultMetadata::new(&serde_json::json!("omitted_due_to_size_limit")),
ToolResultMetadata::omitted_due_to_size_limit(/*overage_bytes*/ 123_456),
ToolResultMetadata::default(),
]
.into_iter()
.map(|metadata| {
let mut call = ExecutedToolCall::new("test_tool".to_string(), serde_json::json!({}));
call.set_tool_result_metadata(metadata);
call.set_tool_result_sources(ToolResultSources::parse_failed());
call
})
.collect(),
);
item.set_tool_call_cell_id("exec");
item.mark_tool_calls_complete();
assert!(item.has_tool_result_metadata());
let mut expected = item.clone();
let calls = expected
.ensure_tool_call_metadata()
.unwrap()
.executed_tool_calls
.as_mut()
.unwrap();
calls[0].set_tool_result_metadata(ToolResultMetadata::new(&resource_only));
calls[1].set_tool_result_metadata(ToolResultMetadata::omitted_due_to_size_limit(
/*overage_bytes*/ 77,
));
item.retain_tool_resource_access_or_omit_metadata(/*overage_bytes*/ 77);
assert_eq!(item, expected);
item.retain_tool_resource_access_or_omit_metadata(/*overage_bytes*/ 1_234);
assert_eq!(item, expected);
first_executed_tool_call(&mut expected)
.unwrap()
.set_tool_result_metadata(ToolResultMetadata::omitted_due_to_size_limit(
/*overage_bytes*/ 1,
));
item.omit_tool_result_metadata(/*overage_bytes*/ 1);
assert_eq!(item, expected);
item.omit_tool_result_metadata(/*overage_bytes*/ 99_999);
assert_eq!(item, expected);
item.clear_tool_result_metadata();
assert!(!item.has_tool_result_metadata());
}
#[test]
fn raw_result_metadata_is_shed_before_sources_calls_or_completion() -> Result<()> {
let mut call = ExecutedToolCall::new("test_tool".to_string(), serde_json::json!({}));
// Keep each result below its capture limit while filling the larger request budget.
let name = format!(
"test_tool{}",
"x".repeat(MAX_EXECUTED_TOOL_CALL_METADATA_BYTES / 8 - 16 * 1024)
);
let mut call = ExecutedToolCall::new(name, serde_json::json!({}));
call.set_tool_result_sources(ToolResultSources::new(vec![ToolResultSource {
r#type: "test_resource".to_string(),
id: "R1".to_string(),
@@ -633,17 +661,17 @@ fn raw_result_metadata_is_shed_before_sources_calls_or_completion() -> Result<()
.unwrap()
.set_tool_result_metadata(metadata.clone());
}
let original_bytes = items
.iter()
.map(executed_tool_call_metadata_bytes)
.sum::<usize>();
let mut expected = items.clone();
first_executed_tool_call(&mut expected[0])
.unwrap()
.set_tool_result_metadata(ToolResultMetadata::omitted_due_to_size_limit());
assert!(
items
.iter()
.map(executed_tool_call_metadata_bytes)
.sum::<usize>()
> MAX_EXECUTED_TOOL_CALL_METADATA_BYTES
);
.set_tool_result_metadata(ToolResultMetadata::omitted_due_to_size_limit(
original_bytes - MAX_EXECUTED_TOOL_CALL_METADATA_BYTES,
));
assert!(original_bytes > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES);
for bound in [
bound_executed_tool_calls_for_prompt,
bound_executed_tool_calls_for_prompt_prioritizing_recent,
@@ -677,7 +705,7 @@ fn result_metadata_shedding_keeps_smaller_results_within_one_output() {
serde_json::json!({ "value": "é".repeat(10_000) }),
],
1,
ToolResultMetadata::omitted_due_to_size_limit(),
None,
),
(
[
@@ -686,7 +714,7 @@ fn result_metadata_shedding_keeps_smaller_results_within_one_output() {
serde_json::json!({ "value": "\0".repeat(4_500) }),
],
2,
ToolResultMetadata::omitted_due_to_size_limit(),
None,
),
(
[
@@ -698,7 +726,7 @@ fn result_metadata_shedding_keeps_smaller_results_within_one_output() {
}),
],
2,
ToolResultMetadata::new(&resource_only),
Some(ToolResultMetadata::new(&resource_only)),
),
] {
let mut metadata = metadata.to_vec();
@@ -706,6 +734,7 @@ fn result_metadata_shedding_keeps_smaller_results_within_one_output() {
serde_json::json!({ "padding": "x".repeat(17_000) });
6
]);
let name_padding = (MAX_EXECUTED_TOOL_CALL_METADATA_BYTES - 128 * 1024) / metadata.len();
let mut item = output("exec");
item.append_executed_tool_calls(
metadata
@@ -713,7 +742,7 @@ fn result_metadata_shedding_keeps_smaller_results_within_one_output() {
.enumerate()
.map(|(index, metadata)| {
let mut call = ExecutedToolCall::new(
format!("tool_{index}"),
format!("tool_{index}{}", "x".repeat(name_padding)),
serde_json::json!({ "argument": index }),
);
call.set_tool_result_sources(ToolResultSources::new(vec![ToolResultSource {
@@ -728,6 +757,11 @@ fn result_metadata_shedding_keeps_smaller_results_within_one_output() {
item.set_tool_call_cell_id("exec");
item.mark_tool_calls_complete();
assert!(executed_tool_call_metadata_bytes(&item) > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES);
let replacement = replacement.unwrap_or_else(|| {
ToolResultMetadata::omitted_due_to_size_limit(
executed_tool_call_metadata_bytes(&item) - MAX_EXECUTED_TOOL_CALL_METADATA_BYTES,
)
});
let mut expected = item.clone();
expected
.ensure_tool_call_metadata()
@@ -762,7 +796,7 @@ fn result_metadata_shedding_removes_only_markers_needed_to_fit() {
id: "R1".to_string(),
}]));
let mut calls = vec![call; 20];
let omitted = ToolResultMetadata::omitted_due_to_size_limit();
let omitted = ToolResultMetadata::omitted_due_to_size_limit(/*overage_bytes*/ 1);
let marker_bytes = serde_json::to_vec(&omitted).unwrap().len();
// A real object tied with a marker must survive even though it is older.
let equal_sized = ToolResultMetadata::new(&serde_json::json!({
@@ -790,9 +824,10 @@ fn result_metadata_shedding_removes_only_markers_needed_to_fit() {
.unwrap();
let call_count = calls.len();
for (index, call) in calls.iter_mut().enumerate() {
call.arguments = ExecutedToolCallArguments::Raw(serde_json::json!("x".repeat(
remaining_bytes / call_count + usize::from(index < remaining_bytes % call_count)
)));
// Pad names so the per-argument limit cannot preempt result-metadata shedding.
call.name.push_str(&"x".repeat(
remaining_bytes / call_count + usize::from(index < remaining_bytes % call_count),
));
}
assert_eq!(
executed_tool_call_metadata_bytes(&item),
@@ -806,8 +841,9 @@ fn result_metadata_shedding_removes_only_markers_needed_to_fit() {
.as_mut()
.unwrap()[1]
.set_tool_result_metadata(ToolResultMetadata::default());
// First downgrade the larger, newer result, then remove the older marker. The
// removal order must use the updated sizes rather than the original raw size.
let tied_markers = item.clone();
// The new marker has a larger overage value and therefore more bytes. Remove
// that marker first, while tied smaller markers still precede provider data.
item.ensure_tool_call_metadata()
.unwrap()
.executed_tool_calls
@@ -816,18 +852,29 @@ fn result_metadata_shedding_removes_only_markers_needed_to_fit() {
.set_tool_result_metadata(ToolResultMetadata::new(&serde_json::json!({
"large": "x".repeat(4 * 1024),
})));
for bound in [
bound_executed_tool_calls_for_prompt,
bound_executed_tool_calls_for_prompt_prioritizing_recent,
] {
let mut bounded = item.clone();
bound(std::slice::from_mut(&mut bounded));
assert_eq!(bounded, expected);
assert!(
executed_tool_call_metadata_bytes(&bounded) <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES
);
bound(std::slice::from_mut(&mut bounded));
assert_eq!(bounded, expected);
let mut larger_marker_expected = item.clone();
larger_marker_expected
.ensure_tool_call_metadata()
.unwrap()
.executed_tool_calls
.as_mut()
.unwrap()[3]
.set_tool_result_metadata(ToolResultMetadata::default());
for (item, expected) in [(tied_markers, expected), (item, larger_marker_expected)] {
for bound in [
bound_executed_tool_calls_for_prompt,
bound_executed_tool_calls_for_prompt_prioritizing_recent,
] {
let mut bounded = item.clone();
bound(std::slice::from_mut(&mut bounded));
assert_eq!(bounded, expected);
assert!(
executed_tool_call_metadata_bytes(&bounded)
<= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES
);
bound(std::slice::from_mut(&mut bounded));
assert_eq!(bounded, expected);
}
}
}
@@ -855,10 +902,10 @@ fn result_metadata_markers_do_not_displace_existing_evidence() {
.unwrap();
let call_count = calls.len();
for (index, call) in calls.iter_mut().enumerate() {
// Fill the budget exactly, with every argument still below its individual limit.
call.arguments = ExecutedToolCallArguments::Raw(serde_json::json!("x".repeat(
remaining_bytes / call_count + usize::from(index < remaining_bytes % call_count)
)));
// Fill the budget exactly without increasing arguments past their individual limit.
call.name.push_str(&"x".repeat(
remaining_bytes / call_count + usize::from(index < remaining_bytes % call_count),
));
}
assert_eq!(
executed_tool_call_metadata_bytes(&item),