diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 66d8338993..f8df44c924 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -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); } diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 12c341b5a9..318486aacc 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -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(()) } diff --git a/codex-rs/core/src/client_tool_metadata.rs b/codex-rs/core/src/client_tool_metadata.rs new file mode 100644 index 0000000000..c546eaca9f --- /dev/null +++ b/codex-rs/core/src/client_tool_metadata.rs @@ -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( + message: &T, + input: &[ResponseItem], +) -> Option> { + 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; diff --git a/codex-rs/core/src/client_tool_metadata_tests.rs b/codex-rs/core/src/client_tool_metadata_tests.rs new file mode 100644 index 0000000000..976846d95b --- /dev/null +++ b/codex-rs/core/src/client_tool_metadata_tests.rs @@ -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); +} diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index bbefb411b8..320b114817 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -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::>>()?; - 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(); diff --git a/codex-rs/core/src/tools/executed_tool_calls.rs b/codex-rs/core/src/tools/executed_tool_calls.rs index 418d399999..f5952e4a80 100644 --- a/codex-rs/core/src/tools/executed_tool_calls.rs +++ b/codex-rs/core/src/tools/executed_tool_calls.rs @@ -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 { diff --git a/codex-rs/core/src/tools/executed_tool_calls/request_metadata.rs b/codex-rs/core/src/tools/executed_tool_calls/request_metadata.rs index 9515c137fd..053c4a490a 100644 --- a/codex-rs/core/src/tools/executed_tool_calls/request_metadata.rs +++ b/codex-rs/core/src/tools/executed_tool_calls/request_metadata.rs @@ -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; } diff --git a/codex-rs/core/src/tools/executed_tool_calls/request_metadata_tests.rs b/codex-rs/core/src/tools/executed_tool_calls/request_metadata_tests.rs index f4d3176161..fe70a5cdd2 100644 --- a/codex-rs/core/src/tools/executed_tool_calls/request_metadata_tests.rs +++ b/codex-rs/core/src/tools/executed_tool_calls/request_metadata_tests.rs @@ -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::() + <= 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::() - <= 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::>(); + 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::>() .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::>(); - 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)); diff --git a/codex-rs/core/src/tools/metadata_metrics.rs b/codex-rs/core/src/tools/metadata_metrics.rs new file mode 100644 index 0000000000..f4fee32b1f --- /dev/null +++ b/codex-rs/core/src/tools/metadata_metrics.rs @@ -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; diff --git a/codex-rs/core/src/tools/metadata_metrics_tests.rs b/codex-rs/core/src/tools/metadata_metrics_tests.rs new file mode 100644 index 0000000000..253f2d16fb --- /dev/null +++ b/codex-rs/core/src/tools/metadata_metrics_tests.rs @@ -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 { + 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::>(); + 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::>(); + 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))]), + ); +} diff --git a/codex-rs/core/src/tools/mod.rs b/codex-rs/core/src/tools/mod.rs index 464996b061..ce81b20491 100644 --- a/codex-rs/core/src/tools/mod.rs +++ b/codex-rs/core/src/tools/mod.rs @@ -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; diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index b753df9674..a8745323cb 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -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::>(); - assert!(metadata_sizes[..6].iter().all(|bytes| *bytes < 32 * 1024)); - assert!(metadata_sizes[..6].iter().sum::() > 128 * 1024); - assert!(metadata_sizes[6] > 32 * 1024); - assert!(metadata_sizes[7] < 32 * 1024); + assert!(metadata_sizes.iter().sum::() > 128 * 1024); + assert!(metadata_sizes.iter().sum::() < 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::() - <= 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(); diff --git a/codex-rs/core/tests/suite/direct_tool_metadata.rs b/codex-rs/core/tests/suite/direct_tool_metadata.rs index a6bb8059e4..b968829db3 100644 --- a/codex-rs/core/tests/suite/direct_tool_metadata.rs +++ b/codex-rs/core/tests/suite/direct_tool_metadata.rs @@ -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(()) } diff --git a/codex-rs/protocol/src/models/executed_tool_calls.rs b/codex-rs/protocol/src/models/executed_tool_calls.rs index 3fe38189be..7ec8afa81f 100644 --- a/codex-rs/protocol/src/models/executed_tool_calls.rs +++ b/codex-rs/protocol/src/models/executed_tool_calls.rs @@ -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::().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 { - 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>); @@ -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 diff --git a/codex-rs/protocol/src/models/executed_tool_calls_tests.rs b/codex-rs/protocol/src/models/executed_tool_calls_tests.rs index 7d033a27a1..9e79d43007 100644 --- a/codex-rs/protocol/src/models/executed_tool_calls_tests.rs +++ b/codex-rs/protocol/src/models/executed_tool_calls_tests.rs @@ -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::(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::(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::(); 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::() - > 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),