From dbb875d23cfd1a327cf6b51d7b50dad4b4601268 Mon Sep 17 00:00:00 2001 From: ningyi-oai Date: Thu, 24 Sep 2026 22:42:51 +0000 Subject: [PATCH] Bound tool-call observations to the outgoing Responses message budget (#47957) ## Why Removing tool-result metadata alone can leave requests above the 15 MiB message budget when recorded tool calls and arguments occupy the remaining space. If trimming loses calls or arguments from a Code Mode cell, later waits must not report its call inventory as complete. ## What changed - Bound all optional tool-call observations against the serialized HTTP or WebSocket message size, preserving ordinary tool outputs and continuation history. - Shed generic result metadata, sources, and recorded arguments before resource-access evidence, retaining tool names and smaller metadata where possible. - Invalidate completeness across affected cell outputs and later waits when the outgoing copy loses calls or arguments. - Preserve size-limit omission markers for direct tool-result metadata when they fit, while continuing to exclude that metadata from custom providers. ## Testing Add regression coverage for HTTP and WebSocket message budgets, unchanged tool results and history, resource-evidence retention, direct-call omission markers, and Code Mode waits after argument truncation. GitOrigin-RevId: 538dc9d8c11b0830ec70327422f98a1803122e85 --- codex-rs/core/src/client.rs | 14 + codex-rs/core/src/client_tests.rs | 94 +++- codex-rs/core/src/client_tool_metadata.rs | 57 +-- .../core/src/client_tool_metadata_tests.rs | 105 ++++- codex-rs/core/src/session/session.rs | 1 + .../core/src/tools/executed_tool_calls.rs | 50 +- .../request_metadata_tests.rs | 70 +++ .../tools/executed_tool_calls_direct_tests.rs | 45 ++ codex-rs/core/tests/suite/code_mode.rs | 139 ++++++ .../core/tests/suite/direct_tool_metadata.rs | 439 ++++++++++++++++++ codex-rs/protocol/src/models.rs | 1 + .../src/models/executed_tool_calls.rs | 275 +++++++++-- .../src/models/executed_tool_calls_tests.rs | 197 ++++++++ 13 files changed, 1369 insertions(+), 118 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index f8df44c924..f8faf45498 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -36,6 +36,7 @@ use std::sync::atomic::Ordering; mod tool_metadata; use crate::CodexResponsesHeaders; +use crate::tools::ExecutedToolCalls; use async_channel::Sender; use codex_api::AgentIdentityTelemetry; use codex_api::ApiError; @@ -270,6 +271,7 @@ pub struct ModelClient { http_client_factory: HttpClientFactory, restored_history: bool, request_contributors: Vec>, + executed_tool_calls: Option, } /// A turn-scoped streaming session created from a [`ModelClient`]. @@ -548,9 +550,15 @@ impl ModelClient { http_client_factory, restored_history: false, request_contributors, + executed_tool_calls: None, } } + pub(crate) fn with_executed_tool_calls(mut self, recorder: ExecutedToolCalls) -> Self { + self.executed_tool_calls = Some(recorder); + self + } + pub(crate) fn reasoning_effort_override_enabled(&self, model_info: &ModelInfo) -> bool { self.state.reasoning_effort_override_enabled && self.state.provider.info().is_openai() @@ -1747,6 +1755,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) { + if let Some(recorder) = &self.client.executed_tool_calls { + recorder.invalidate_wire_inventory_loss(&request.input, &input); + } request.input = input; } inference_trace_attempt.record_started(&request); @@ -2031,6 +2042,9 @@ impl ModelClientSession { 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() { + if let Some(recorder) = &self.client.executed_tool_calls { + recorder.invalidate_wire_inventory_loss(payload.input, input); + } let ResponsesWsRequest::ResponseCreate(payload) = &mut ws_request; payload.input = input; } diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 92c84d4b66..fc88e15ec8 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -485,9 +485,69 @@ fn output_with_tool_result_metadata(metadata: ToolResultMetadata) -> ResponseIte output } -#[test] -fn responses_request_limits_internal_metadata_to_resolved_first_party_https_endpoint() +#[tokio::test] +async fn responses_request_limits_internal_metadata_to_resolved_first_party_https_endpoint() -> anyhow::Result<()> { + use crate::session::step_context::StepContext; + use crate::tools::ExecutedToolCalls; + use crate::tools::context::ToolCallSource; + use crate::tools::context::ToolPayload; + use crate::tools::router::ToolCall; + use codex_features::Feature; + use codex_features::Features; + + let (_, turn) = crate::session::tests::make_session_and_context().await; + let step = StepContext::for_test(Arc::new(turn)); + let mut features = Features::default(); + features.enable(Feature::ExecutedToolCallMetadata); + let recorder = ExecutedToolCalls::new(&features, &codex_history::InitialHistory::New); + let resource_metadata = json!({"openai/resource_access": {"payload": "x".repeat(600 * 1024)}}); + let mut outputs = Vec::new(); + for id in ["first", "second"] { + let call = ToolCall { + tool_name: codex_tools::ToolName::plain("mcp__apps__read"), + call_id: id.to_string(), + payload: ToolPayload::Function { + arguments: json!({"query": id}).to_string(), + }, + encrypted_function_args: None, + }; + let (mut recorded, permit) = recorder + .prepare_direct_call(&call, &ToolCallSource::Direct, &step) + .expect("direct observation slot"); + recorded.set_tool_result_metadata(ToolResultMetadata::new(&resource_metadata)); + let mut item = ResponseItem::from(ResponseInputItem::FunctionCallOutput { + call_id: id.to_string(), + output: FunctionCallOutputPayload::from_text(format!("result for {id}")), + }); + recorder.attach_direct_call_to_output(&mut item, Some((recorded, permit))); + outputs.push(item); + } + let original_outputs = outputs.clone(); + recorder.attach_to_prompt(&mut outputs, &mut Default::default()); + assert_eq!(outputs, original_outputs); + let recorded = serde_json::to_value(&outputs)?; + assert_eq!( + recorded[0]["internal_chat_message_metadata_passthrough"]["executed_tool_calls"][0]["tool_result_metadata"], + resource_metadata, + ); + assert!( + recorded[1]["internal_chat_message_metadata_passthrough"]["executed_tool_calls"][0] + ["tool_result_metadata"] + .as_str() + .is_some_and(|value| value.starts_with("omitted_due_to_size_limit (overage_bytes=")) + ); + let omitted_output = outputs.pop().expect("second direct output"); + let mut without_omitted_metadata = ResponseItem::from(ResponseInputItem::FunctionCallOutput { + call_id: "second".to_string(), + output: FunctionCallOutputPayload::from_text("result for second".to_string()), + }); + without_omitted_metadata.append_executed_tool_calls(vec![ExecutedToolCall::new( + "mcp__apps__read".to_string(), + json!({"query": "second"}), + )]); + without_omitted_metadata.mark_tool_calls_complete(); + let provider = ModelProviderInfo::create_openai_provider(Some("https://api.openai.com/v1".to_string())); let mut api_provider = provider.to_api_provider(/*auth_mode*/ None)?; @@ -511,7 +571,7 @@ fn responses_request_limits_internal_metadata_to_resolved_first_party_https_endp }], }; let prompt = Prompt { - input: vec![output.clone()], + input: vec![output.clone(), omitted_output.clone()], ..Default::default() }; let mut responses_metadata = test_responses_metadata_for_client( @@ -547,15 +607,23 @@ fn responses_request_limits_internal_metadata_to_resolved_first_party_https_endp &responses_metadata, include_internal, )?; - assert_eq!( - request.input.last(), - Some(if allowed { - &output - } else { - &without_raw_metadata - }), + let expected_input = if allowed { + vec![output.clone(), omitted_output.clone()] + } else { + vec![ + without_raw_metadata.clone(), + without_omitted_metadata.clone(), + ] + }; + assert!( + request.input.ends_with(&expected_input), "resolved endpoint: {base_url}, responses_lite: {responses_lite}", ); + let expected_wire_input = serde_json::to_value(&request.input)?; + assert_eq!( + serde_json::to_value(&request)?["input"], + expected_wire_input + ); let expected_json = allowed .then(|| serde_json::to_string(&attribution).map(serde_json::Value::String)) .transpose()?; @@ -574,8 +642,10 @@ fn responses_request_limits_internal_metadata_to_resolved_first_party_https_endp client_metadata: Some(ws_client_metadata), ..codex_api::ResponseCreateWsRequest::from(&request) }; + let ws_request = serde_json::to_value(ws_request)?; + assert_eq!(ws_request["input"], expected_wire_input); assert_eq!( - serde_json::to_value(ws_request)? + ws_request .get("client_metadata") .and_then(|metadata| metadata.get(MCP_ATTRIBUTION_CLIENT_METADATA_KEY)), expected_json.as_ref(), @@ -585,7 +655,7 @@ fn responses_request_limits_internal_metadata_to_resolved_first_party_https_endp .get("mcp_attribution") .is_none() ); - assert_eq!(prompt.input, vec![output.clone()]); + assert_eq!(prompt.input, vec![output.clone(), omitted_output.clone()]); } } responses_metadata.mcp_attribution = Some(McpAttribution { diff --git a/codex-rs/core/src/client_tool_metadata.rs b/codex-rs/core/src/client_tool_metadata.rs index c546eaca9f..5fecb42a1b 100644 --- a/codex-rs/core/src/client_tool_metadata.rs +++ b/codex-rs/core/src/client_tool_metadata.rs @@ -1,9 +1,10 @@ -//! Bounds optional tool-result metadata using the actual outgoing Responses message. +//! Bounds optional tool observations 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 codex_protocol::models::bound_executed_tool_calls_for_message; use serde::Serialize; const MAX_RESPONSE_MESSAGE_BYTES: usize = 15 * 1024 * 1024; @@ -12,7 +13,13 @@ pub(super) fn bounded_input( message: &T, input: &[ResponseItem], ) -> Option> { - if !input.iter().any(ResponseItem::has_tool_result_metadata) { + if !input.iter().any(|item| { + item.executed_tool_call_metadata().is_some_and(|metadata| { + metadata.executed_tool_calls.is_some() + || metadata.cell_id.is_some() + || metadata.tool_calls_complete.is_some() + }) + }) { return None; } // Serialization errors belong to the transport's existing error path, not this soft budget. @@ -23,34 +30,7 @@ pub(super) fn bounded_input( 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, - ); + bound_executed_tool_calls_for_message(&mut bounded, before.saturating_sub(overage)); let after = metadata_metrics::metadata_bytes(&bounded); if before == after { return None; @@ -60,23 +40,6 @@ pub(super) fn bounded_input( 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 index 976846d95b..4d093231db 100644 --- a/codex-rs/core/src/client_tool_metadata_tests.rs +++ b/codex-rs/core/src/client_tool_metadata_tests.rs @@ -19,16 +19,25 @@ fn request() -> ResponsesApiRequest { } fn request_with_metadata(metadata: &serde_json::Value) -> ResponsesApiRequest { + request_with_metadata_and_source(metadata, /*with_source*/ true) +} + +fn request_with_metadata_and_source( + metadata: &serde_json::Value, + with_source: bool, +) -> 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(), - }])); + if with_source { + 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(); @@ -79,8 +88,11 @@ fn http_message_budget_preserves_resources_until_the_message_still_exceeds_limit }) ); } else if overage == 20 * 1024 { - let mut resource_only = original_input.clone(); - resource_only[0].retain_tool_resource_access(); + let resource_only = request_with_metadata_and_source( + &json!({"openai/resource_access": {"resources": ["é".repeat(4 * 1024)]}}), + /*with_source*/ false, + ) + .input; let first_reduction = metadata_metrics::metadata_bytes(&original_input) - metadata_metrics::metadata_bytes(&resource_only); assert_eq!( @@ -98,11 +110,19 @@ fn http_message_budget_preserves_resources_until_the_message_still_exceeds_limit ); } assert_eq!(request.input, original_input); - let mut ordinary = original_input; + let mut ordinary = if overage == 20 * 1024 { + request_with_metadata_and_source(&json!({}), /*with_source*/ false).input + } else { + 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(); + if overage == 40 * 1024 { + item.clear_executed_tool_calls(); + } else { + item.clear_tool_result_metadata(); + } } } assert_eq!(bounded_ordinary, ordinary); @@ -143,6 +163,11 @@ fn message_budget_removes_residual_values_and_markers_when_ordinary_input_fits() ordinary.instructions = "x".repeat(MAX_RESPONSE_MESSAGE_BYTES - ordinary_bytes); request.instructions = ordinary.instructions.clone(); let original_input = request.input.clone(); + let mut expected = ordinary.clone(); + if metadata.get("openai/resource_access").is_some() { + expected.input = + request_with_metadata_and_source(&metadata, /*with_source*/ false).input; + } let bounded = if websocket { let ordinary_message = ResponsesWsRequest::ResponseCreate(ResponseCreateWsRequest::from(&ordinary)); @@ -157,10 +182,13 @@ fn message_budget_removes_residual_values_and_markers_when_ordinary_input_fits() .expect("optional result metadata must not keep a fitting message oversized"); let ResponsesWsRequest::ResponseCreate(payload) = &mut message; payload.input = &bounded; + let expected_message = + ResponsesWsRequest::ResponseCreate(ResponseCreateWsRequest::from(&expected)); assert_eq!( serialized_json_bytes(&message).unwrap(), - MAX_RESPONSE_MESSAGE_BYTES - 1, + serialized_json_bytes(&expected_message).unwrap(), ); + assert!(serialized_json_bytes(&message).unwrap() <= MAX_RESPONSE_MESSAGE_BYTES); bounded } else { assert_eq!( @@ -172,15 +200,12 @@ fn message_budget_removes_residual_values_and_markers_when_ordinary_input_fits() .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, - ); + assert_eq!(message, expected); + assert!(serialized_json_bytes(&message).unwrap() <= MAX_RESPONSE_MESSAGE_BYTES); bounded }; // Whole-input equality covers sources, calls, completion and ordinary output. - assert_eq!(bounded, ordinary.input); + assert_eq!(bounded, expected.input); assert_eq!(request.input, original_input); } } @@ -255,9 +280,55 @@ fn websocket_budget_uses_the_actual_delta_and_does_not_mutate_logical_history() 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(); + ordinary.clear_executed_tool_calls(); + bounded_ordinary.clear_executed_tool_calls(); assert_eq!(bounded_ordinary, ordinary); assert_eq!(request.input[1], self::request().input[0]); assert_eq!(bounded_input(&full, &bounded), None); } + +#[test] +fn inventory_only_message_budget_counts_the_complete_encoded_envelope() { + for websocket in [false, true] { + let mut request = request(); + request.input[0].clear_tool_result_metadata(); + let original_input = request.input.clone(); + let mut ordinary = request.clone(); + ordinary.input[0].clear_executed_tool_calls(); + let ordinary_bytes = if websocket { + serialized_json_bytes(&ResponsesWsRequest::ResponseCreate( + ResponseCreateWsRequest::from(&ordinary), + )) + .unwrap() + } else { + serialized_json_bytes(&ordinary).unwrap() + }; + request.instructions = "x".repeat(MAX_RESPONSE_MESSAGE_BYTES - ordinary_bytes + 1); + ordinary.instructions = request.instructions.clone(); + if websocket { + let mut message = + ResponsesWsRequest::ResponseCreate(ResponseCreateWsRequest::from(&request)); + let bounded = bounded_input(&message, &request.input) + .expect("inventory without raw results must not keep the message oversized"); + assert_eq!(bounded, ordinary.input); + let ResponsesWsRequest::ResponseCreate(payload) = &mut message; + payload.input = &bounded; + assert_eq!( + serialized_json_bytes(&message).unwrap(), + MAX_RESPONSE_MESSAGE_BYTES + ); + } else { + let bounded = bounded_input(&request, &request.input) + .expect("inventory without raw results must not keep the message oversized"); + assert_eq!(bounded, ordinary.input); + let mut message = request.clone(); + message.input = bounded; + assert_eq!(message, ordinary); + assert_eq!( + serialized_json_bytes(&message).unwrap(), + MAX_RESPONSE_MESSAGE_BYTES + ); + } + assert_eq!(request.input, original_input); + } +} diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 3f9cd15720..b6f39ac38a 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1735,6 +1735,7 @@ impl Session { workspace_routing.as_ref().clone(), extensions.model_request_contributors().to_vec(), ) + .with_executed_tool_calls(executed_tool_calls.clone()) .with_restored_history(matches!( &initial_history, InitialHistory::Resumed(_) | InitialHistory::Forked(_) diff --git a/codex-rs/core/src/tools/executed_tool_calls.rs b/codex-rs/core/src/tools/executed_tool_calls.rs index f5952e4a80..99ca2028d2 100644 --- a/codex-rs/core/src/tools/executed_tool_calls.rs +++ b/codex-rs/core/src/tools/executed_tool_calls.rs @@ -66,6 +66,15 @@ pub(crate) struct ExecutedToolCalls { mcp_attribution: mcp_attribution::McpAttributionRecorder, } +// Avoid exposing recorded call arguments or result metadata through client Debug output. +impl std::fmt::Debug for ExecutedToolCalls { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("ExecutedToolCalls") + .finish_non_exhaustive() + } +} + // The tool future owns this reservation, so completion or cancellation releases it. pub(crate) struct DirectCallPermit { recording: Weak<()>, @@ -257,6 +266,41 @@ impl ExecutedToolCalls { .unwrap_or_else(std::sync::PoisonError::into_inner) } + /// A later wait cannot claim a complete inventory if an earlier wire copy + /// lost recorded calls or arguments from the same Code Mode cell. + pub(crate) fn invalidate_wire_inventory_loss( + &self, + original: &[ResponseItem], + bounded: &[ResponseItem], + ) { + let mut state = self.lock_state(); + let Some(state) = state.as_mut() else { + return; + }; + // Request bounding clones the same items in the same order; it only edits metadata. + for (original, bounded) in original.iter().zip(bounded) { + let Some(metadata) = original.executed_tool_call_metadata() else { + continue; + }; + let Some(origin) = metadata.cell_id.as_deref() else { + continue; + }; + let Some(calls) = metadata + .executed_tool_calls + .as_deref() + .filter(|calls| !calls.is_empty()) + else { + continue; + }; + if bounded + .executed_tool_call_metadata() + .is_none_or(|bounded| !bounded.has_same_tool_calls(calls)) + { + state.invalidate_origin(origin); + } + } + } + /// Called under the session config lock; this never changes execution features. pub(crate) fn refresh(&self, features: &Features) { let enabled = Self::is_enabled(features); @@ -348,7 +392,11 @@ impl ExecutedToolCalls { let mut bytes = executed_tool_call_metadata_bytes(item); let original_bytes = bytes; if bytes > available { - item.retain_tool_resource_access(); + item.retain_tool_resource_access_or_omit_metadata(bytes - available); + bytes = executed_tool_call_metadata_bytes(item); + } + if bytes > available { + item.omit_tool_result_metadata(bytes - available); bytes = executed_tool_call_metadata_bytes(item); } if bytes > available { 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 fe70a5cdd2..274ff6b7b7 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 @@ -1578,3 +1578,73 @@ fn finished_cells_without_more_waits_do_not_block_new_calls() { assert!(recorder.attach_pending_to_prompt(&mut items, &mut HashMap::new())); assert_eq!(items, [exec_input("fresh-output"), expected]); } + +#[test] +fn wire_inventory_loss_keeps_later_wait_incomplete() { + for (scenario, expect_complete) in [ + ("metadata_only", true), + ("arguments", false), + ("name", false), + ("removed", false), + ] { + let recorder = new_recorder(InitialHistory::New); + let cell = CellId::new("runtime-a".to_string()); + let unrelated = CellId::new("runtime-b".to_string()); + recorder.start_cell(&cell, "exec-a"); + record_nested_call(&recorder, &cell, "nested-a"); + recorder.start_cell(&unrelated, "exec-b"); + record_nested_call(&recorder, &unrelated, "nested-b"); + let mut original = [exec_input("exec-a"), exec_output("exec-a")]; + assert!(recorder.attach_pending_to_prompt(&mut original, &mut HashMap::new())); + let ResponseItem::CustomToolCallOutput { + internal_chat_message_metadata_passthrough: Some(metadata), + .. + } = &mut original[1] + else { + panic!("expected exec output metadata"); + }; + metadata.executed_tool_calls.as_mut().expect("calls")[0] + .set_tool_result_metadata(ToolResultMetadata::new(&json!({"value": "raw"}))); + let mut bounded = original.clone(); + match scenario { + "metadata_only" => bounded[1].clear_tool_result_metadata(), + "removed" => bounded[1].clear_executed_tool_calls(), + "arguments" | "name" => { + let ResponseItem::CustomToolCallOutput { + internal_chat_message_metadata_passthrough: Some(metadata), + .. + } = &mut bounded[1] + else { + panic!("expected exec output metadata"); + }; + let call = &mut metadata.executed_tool_calls.as_mut().expect("calls")[0]; + if scenario == "arguments" { + *call = ExecutedToolCall::truncated( + "nested_tool".to_string(), + /*original_bytes*/ 2, + /*max_bytes*/ 0, + ); + } else { + call.name = "another_tool".to_string(); + } + } + _ => unreachable!(), + } + recorder.invalidate_wire_inventory_loss(&original, &bounded); + recorder.register_cell(&cell, "wait-a"); + recorder.finish_cell_recording(&cell); + let mut terminal = [wait_input("wait-a", &cell), output("wait-a")]; + recorder.attach_pending_to_prompt(&mut terminal, &mut HashMap::new()); + assert_eq!( + tool_calls_complete(&terminal[1]), + expect_complete.then_some(true), + "{scenario}" + ); + + // A different cell must not lose its completion marker. + recorder.finish_cell_recording(&unrelated); + let mut other = [exec_input("exec-b"), exec_output("exec-b")]; + assert!(recorder.attach_pending_to_prompt(&mut other, &mut HashMap::new())); + assert_eq!(tool_calls_complete(&other[1]), Some(true), "{scenario}"); + } +} diff --git a/codex-rs/core/src/tools/executed_tool_calls_direct_tests.rs b/codex-rs/core/src/tools/executed_tool_calls_direct_tests.rs index 2e2a439afd..4d0f1ea7f1 100644 --- a/codex-rs/core/src/tools/executed_tool_calls_direct_tests.rs +++ b/codex-rs/core/src/tools/executed_tool_calls_direct_tests.rs @@ -15,6 +15,51 @@ fn output(call_id: &str) -> ResponseItem { }) } +#[test] +fn direct_budget_keeps_an_omission_marker_when_it_fits() { + for metadata in [ + json!({"provider": "x".repeat(1024)}), + json!({"openai/resource_access": "x".repeat(1024)}), + ] { + let mut features = Features::default(); + features.enable(Feature::ExecutedToolCallMetadata); + let recorder = ExecutedToolCalls::new(&features, &InitialHistory::New); + let mut call = ExecutedToolCall::new("test_tool".to_string(), json!({"argument": "kept"})); + call.set_tool_result_metadata(ToolResultMetadata::new(&metadata)); + let mut expected = output("direct"); + expected.append_executed_tool_calls(vec![call.clone()]); + expected.mark_tool_calls_complete(); + let available = 512; + let overage = executed_tool_call_metadata_bytes(&expected) - available; + let mut omitted = + ExecutedToolCall::new("test_tool".to_string(), json!({"argument": "kept"})); + omitted.set_tool_result_metadata(ToolResultMetadata::new(&json!(format!( + "omitted_due_to_size_limit (overage_bytes={overage})" + )))); + expected.clear_executed_tool_calls(); + expected.append_executed_tool_calls(vec![omitted]); + expected.mark_tool_calls_complete(); + let retained = MAX_RETAINED_DIRECT_METADATA_BYTES - available; + recorder + .retained_direct_metadata_bytes + .store(retained, Ordering::Relaxed); + + let mut item = output("direct"); + recorder.attach_direct_call_to_output( + &mut item, + Some((call, recorder.reserve_direct_call().unwrap())), + ); + + assert_eq!(item, expected); + assert_eq!( + recorder + .retained_direct_metadata_bytes + .load(Ordering::Relaxed), + retained + executed_tool_call_metadata_bytes(&expected), + ); + } +} + #[tokio::test] async fn direct_pending_limit_releases_on_completion_or_dropped_future() { let (_, turn) = crate::session::tests::make_session_and_context().await; diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 98a755931b..1c589fbfde 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -8842,3 +8842,142 @@ text(JSON.stringify({ Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn code_mode_oversized_websocket_yield_keeps_later_wait_incomplete() -> Result<()> { + skip_if_no_network!(Ok(())); + const LIMIT: usize = 15 * 1024 * 1024; + const PROMPT: &str = "Record a call, yield, then stop"; + + // Calibrate a first request with the same tools/features/turn prompt; its + // exact serialized overhead varies with the model catalog and headers. + let configure = |config: &mut Config, instructions: String| { + config.base_instructions = Some(instructions); + config.model_context_window = Some(20_000_000); + config.model_auto_compact_token_limit = Some(20_000_000); + config.features.disable(Feature::TokenBudget).unwrap(); + config.features.enable(Feature::CodeMode).unwrap(); + config + .features + .enable(Feature::ExecutedToolCallMetadata) + .unwrap(); + config + .features + .disable(Feature::RemoteCompactionV2) + .unwrap(); + config.model_provider.request_max_retries = Some(0); + config.model_provider.stream_max_retries = Some(0); + }; + let warmup = || vec![ev_response_created("warmup"), ev_completed("warmup")]; + let probe_server = responses::start_websocket_server(vec![vec![ + warmup(), + vec![ev_response_created("probe"), ev_completed("probe")], + ]]) + .await; + let mut probe_builder = test_codex() + .with_model("test-gpt-5.1-codex") + .with_config(move |config| configure(config, String::new())); + let probe = probe_builder + .build_with_websocket_server(&probe_server) + .await?; + probe.submit_turn(PROMPT).await?; + let probe_connection = probe_server.single_connection(); + assert_eq!(probe_connection.len(), 2); + let base_bytes = serde_json::to_vec(&probe_connection[1].body_json())?.len(); + assert!(base_bytes + 4 * 1024 < LIMIT); + probe.codex.shutdown_and_wait().await?; + probe_server.shutdown().await; + + // One 7 KiB invocation stays under the recorder's per-output argument + // budget. It pushes the yielded delta over the message budget only. + let instructions = "x".repeat(LIMIT - base_bytes - 4 * 1024); + let code = r#" +await tools.test_sync_tool({ barrier: { id: "x".repeat(7000), participants: 1 } }); +text("yielded"); +yield_control(); +await new Promise(() => {}); +"#; + let mut exec = ev_custom_tool_call("exec-a", "exec", code); + exec["item"]["id"] = serde_json::json!("ctc_exec_a"); + let mut wait = + responses::ev_function_call("wait-a", "wait", r#"{"cell_id":"1","terminate":true}"#); + wait["item"]["id"] = serde_json::json!("fc_wait_a"); + let server = responses::start_websocket_server(vec![vec![ + warmup(), + vec![ev_response_created("resp-1"), exec, ev_completed("resp-1")], + vec![ev_response_created("resp-2"), wait, ev_completed("resp-2")], + vec![ev_response_created("resp-3"), ev_completed("resp-3")], + ]]) + .await; + let mut builder = test_codex() + .with_model("test-gpt-5.1-codex") + .with_config(move |config| configure(config, instructions)); + let test = builder.build_with_websocket_server(&server).await?; + test.submit_turn(PROMPT).await?; + let connection = server.single_connection(); + assert_eq!(connection.len(), 4); + let first = connection[1].body_json(); + let yielded_request = connection[2].body_json(); + let terminal_request = connection[3].body_json(); + assert!(serde_json::to_vec(&first)?.len() <= LIMIT); + assert!(serde_json::to_vec(&yielded_request)?.len() <= LIMIT); + assert_eq!(yielded_request["previous_response_id"], "resp-1"); + assert_eq!(terminal_request["previous_response_id"], "resp-2"); + + let find_output = |request: &Value, call_id: &str| -> Value { + request["input"] + .as_array() + .unwrap() + .iter() + .find(|item| { + item["call_id"] == call_id + && matches!( + item["type"].as_str(), + Some("custom_tool_call_output" | "function_call_output") + ) + }) + .unwrap() + .clone() + }; + let yielded = find_output(&yielded_request, "exec-a"); + let observed = &yielded["internal_chat_message_metadata_passthrough"]; + let output_text = match &yielded["output"] { + Value::String(text) => text.clone(), + Value::Array(items) => items + .iter() + .filter_map(|item| item["text"].as_str()) + .collect::>() + .join("\n"), + _ => panic!("unexpected Code Mode output"), + }; + assert!(output_text.contains("yielded")); + assert!(output_text.contains("Script running with cell ID 1")); + // The request budget must actually trim a recorded call's arguments. + assert_eq!(observed["cell_id"], "exec-a"); + let calls = observed["executed_tool_calls"].as_array().unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0]["name"], "test_sync_tool"); + assert!( + calls[0]["arguments"] + .get("_codex_executed_tool_call_truncated") + .is_some() + ); + let terminal = find_output(&terminal_request, "wait-a"); + assert_eq!( + terminal["internal_chat_message_metadata_passthrough"].get("tool_calls_complete"), + None + ); + + // Compare ordinary outputs against the live session history. No actual + // invocation, output, or wait request is changed by metadata trimming. + let history = test.codex.conversation_history_snapshot().await; + let history = serde_json::to_value(history.items().collect::>())?; + let history_request = serde_json::json!({"input": history}); + for (call_id, actual) in [("exec-a", yielded), ("wait-a", terminal)] { + let original = find_output(&history_request, call_id); + assert_eq!(actual["output"], original["output"]); + } + test.codex.shutdown_and_wait().await?; + server.shutdown().await; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/direct_tool_metadata.rs b/codex-rs/core/tests/suite/direct_tool_metadata.rs index b968829db3..7c52d6dcaf 100644 --- a/codex-rs/core/tests/suite/direct_tool_metadata.rs +++ b/codex-rs/core/tests/suite/direct_tool_metadata.rs @@ -2,7 +2,23 @@ use anyhow::Result; use codex_features::Feature; +use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_model_provider::RemoteCompactionSupport; +use core_test_support::apps_test_server::AppsTestServer; +use core_test_support::apps_test_server::SEARCH_CALENDAR_LIST_TOOL; +use core_test_support::apps_test_server::SEARCH_CALENDAR_NAMESPACE; +use core_test_support::apps_test_server::recorded_apps_tool_calls; +use core_test_support::apps_test_server::search_capable_apps_builder; +use core_test_support::responses::ev_function_call_with_namespace; +use core_test_support::wait_for_mcp_server; +use wiremock::Mock; +use wiremock::Request; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_partial_json; +use wiremock::matchers::method; +use wiremock::matchers::path; +use wiremock::matchers::path_regex; + use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; @@ -16,6 +32,7 @@ use core_test_support::responses::ev_tool_search_call; use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; +use core_test_support::responses::start_websocket_server; use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; @@ -32,6 +49,269 @@ pub(super) fn tool_call_metadata(mut item: Value) -> Value { metadata } +#[test_case(false; "http")] +#[test_case(true; "websocket")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn message_budget_sheds_inventory_without_changing_tool_results_or_history( + websocket: bool, +) -> Result<()> { + skip_if_no_network!(Ok(())); + let message_limit = 15 * 1024 * 1024; + let call_count = 16; + // Each argument fits 8 KiB and all observations fit Direct's 1 MiB budget. + // Only the final message budget should shed this inventory. HTTP includes + // the original call arguments; a WebSocket delta already has those upstream. + let arguments = json!({"plan": [{"step": "x".repeat(7 * 1024), "status": "in_progress"}]}); + assert!(arguments.to_string().len() < 8 * 1024); + assert!(call_count * arguments.to_string().len() < 1024 * 1024); + let next_arguments = json!({"plan": [{"step": "done", "status": "completed"}]}); + let instruction_bytes = message_limit - if websocket { 64 * 1024 } else { 192 * 1024 }; + let instructions = "padding ".repeat(instruction_bytes / 8); + let mut builder = test_codex().with_config(move |config| { + config.base_instructions = Some(instructions); + config.model_context_window = Some(20_000_000); + config.model_auto_compact_token_limit = Some(20_000_000); + config + .features + .disable(Feature::TokenBudget) + .expect("disable token budget"); + config + .features + .disable(Feature::CodeModeOnly) + .expect("disable code-mode-only tools"); + config + .features + .disable(Feature::CodeMode) + .expect("disable code mode"); + config + .features + .enable(Feature::ExecutedToolCallMetadata) + .expect("enable tool-call metadata"); + config.update_plan_enabled = true; + config.model_provider.request_max_retries = Some(0); + config.model_provider.stream_max_retries = Some(0); + }); + let mut batch = vec![ev_response_created("resp-1")]; + for index in 0..call_count { + let mut call = ev_function_call( + &format!("plan-{index}"), + "update_plan", + &arguments.to_string(), + ); + // Stable server IDs keep the WebSocket continuation on the delta path. + call["item"]["id"] = json!(format!("fc_plan_{index}")); + batch.push(call); + } + batch.push(ev_completed("resp-1")); + let mut next_call = ev_function_call("plan-next", "update_plan", &next_arguments.to_string()); + next_call["item"]["id"] = json!("fc_plan_next"); + let events = vec![ + batch, + vec![ + ev_response_created("resp-2"), + next_call, + ev_completed("resp-2"), + ], + vec![ev_response_created("resp-3"), ev_completed("resp-3")], + ]; + let http_server = if websocket { + None + } else { + Some(start_mock_server().await) + }; + let http_mock = if let Some(server) = &http_server { + Some(mount_sse_sequence(server, events.iter().cloned().map(sse).collect()).await) + } else { + None + }; + let websocket_server = if websocket { + let mut responses = vec![vec![ev_response_created("warmup"), ev_completed("warmup")]]; + responses.extend(events); + Some(start_websocket_server(vec![responses]).await) + } else { + None + }; + let test = if let Some(server) = &websocket_server { + builder.build_with_websocket_server(server).await? + } else { + builder + .build_with_auto_env(http_server.as_ref().expect("HTTP test server")) + .await? + }; + test.submit_turn("Update the plan, then finish it").await?; + let requests = if let Some(server) = &websocket_server { + let connection = server.single_connection(); + assert_eq!(connection.len(), 4); + assert_eq!(connection[0].body_json()["generate"], false); + connection[1..] + .iter() + .map(core_test_support::responses::WebSocketRequest::body_json) + .collect::>() + } else { + http_mock + .as_ref() + .expect("HTTP response sequence") + .requests() + .iter() + .map(core_test_support::responses::ResponsesRequest::body_json) + .collect() + }; + assert_eq!(requests.len(), 3); + if websocket { + assert_eq!(requests[1]["previous_response_id"], "resp-1"); + assert_eq!(requests[2]["previous_response_id"], "resp-2"); + } + for (index, request) in requests.iter().enumerate() { + assert_eq!( + request["instructions"] + .as_str() + .expect("request instructions") + .len(), + instruction_bytes + ); + let request_bytes = serde_json::to_vec(request)?.len(); + assert!( + request_bytes <= message_limit, + "request {index} is {request_bytes} bytes (limit {message_limit})" + ); + } + + let history = test.codex.conversation_history_snapshot().await; + let history = serde_json::to_value(history.items().collect::>())?; + let outputs = history + .as_array() + .expect("serialized history items") + .iter() + .filter(|item| item["type"] == "function_call_output") + .collect::>(); + assert_eq!(outputs.len(), call_count + 1); + for (index, output) in outputs.iter().enumerate() { + let (call_id, expected_arguments) = if index < call_count { + (format!("plan-{index}"), &arguments) + } else { + ("plan-next".to_string(), &next_arguments) + }; + assert_eq!(output["call_id"], call_id); + assert_eq!(output["output"], "Plan updated"); + assert_eq!( + tool_call_metadata((**output).clone()), + json!({ + "executed_tool_calls": [{"name": "update_plan", "arguments": expected_arguments}], + "tool_calls_complete": true, + }) + ); + } + + let wire_calls = requests[1]["input"] + .as_array() + .expect("second request input") + .iter() + .filter(|item| item["type"] == "function_call") + .collect::>(); + assert_eq!(wire_calls.len(), if websocket { 0 } else { call_count }); + for (index, call) in wire_calls.iter().enumerate() { + assert_eq!(call["call_id"], format!("plan-{index}")); + assert_eq!(call["name"], "update_plan"); + assert_eq!(call["arguments"], arguments.to_string()); + } + // Inspect the bounded wire copy before rebuilding an unbounded comparison below. + let wire_outputs = requests[1]["input"] + .as_array() + .expect("second request input") + .iter() + .filter(|item| item["type"] == "function_call_output") + .collect::>(); + assert_eq!(wire_outputs.len(), call_count); + let argument_bytes = serde_json::to_vec(&arguments)?.len() as u64; + let mut retained_arguments = 0; + let mut truncated_arguments = 0; + for (index, output) in wire_outputs.iter().enumerate() { + assert_eq!(output["call_id"], format!("plan-{index}")); + assert_eq!(output["output"], "Plan updated"); + let metadata = tool_call_metadata((*output).clone()); + let calls = metadata["executed_tool_calls"] + .as_array() + .expect("wire call inventory"); + assert_eq!(calls.len(), 1); + let call = &calls[0]; + if call["arguments"] == arguments { + assert_eq!( + metadata, + json!({ + "executed_tool_calls": [{"name": "update_plan", "arguments": arguments}], + "tool_calls_complete": true, + }) + ); + retained_arguments += 1; + } else { + let truncation = &call["arguments"]["_codex_executed_tool_call_truncated"]; + let max_bytes = truncation["max_bytes"] + .as_u64() + .expect("wire truncation limit"); + assert!(max_bytes < argument_bytes); + assert_eq!( + metadata, + json!({ + "executed_tool_calls": [{ + "name": "update_plan", + "arguments": {"_codex_executed_tool_call_truncated": { + "original_bytes": argument_bytes, + "max_bytes": max_bytes, + }}, + }], + }) + ); + truncated_arguments += 1; + } + } + assert!( + retained_arguments > 0, + "wire request lost every full observation" + ); + assert!( + truncated_arguments > 0, + "wire request did not shed any arguments" + ); + + let mut unbounded = requests[1].clone(); + let mut matched_outputs = 0; + for item in unbounded["input"] + .as_array_mut() + .expect("unbounded request input") + { + if item["type"] != "function_call_output" { + continue; + } + let source = outputs + .iter() + .find(|output| output["call_id"] == item["call_id"]) + .expect("live history output matching the wire call ID"); + assert_eq!(item["output"], source["output"]); + item["internal_chat_message_metadata_passthrough"] = + source["internal_chat_message_metadata_passthrough"].clone(); + matched_outputs += 1; + } + assert_eq!(matched_outputs, call_count); + assert!(serde_json::to_vec(&unbounded)?.len() > message_limit); + if websocket { + let delta = requests[2]["input"] + .as_array() + .expect("WebSocket continuation input"); + assert_eq!(delta.len(), 1); + assert_eq!(delta[0]["call_id"], "plan-next"); + assert_eq!(delta[0]["output"], "Plan updated"); + assert_eq!( + tool_call_metadata(delta[0].clone()), + tool_call_metadata((*outputs[call_count]).clone()) + ); + } + test.codex.shutdown_and_wait().await?; + if let Some(server) = websocket_server { + server.shutdown().await; + } + Ok(()) +} + #[test_case(RemoteCompactionSupport::Unsupported, true; "local")] #[test_case(RemoteCompactionSupport::V2, true; "remote_v2")] #[test_case(RemoteCompactionSupport::V2, false; "remote_v2_disabled_after_capture")] @@ -342,3 +622,162 @@ async fn direct_function_and_tool_search_mark_complete_attempts( } Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn direct_metadata_limit_retains_marker_without_sending_it_to_custom_provider() -> Result<()> +{ + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let apps = AppsTestServer::mount(&server).await?; + let raw_metadata = json!({"openai/resource_access": {"payload": "x".repeat(600 * 1024)}}); + let raw_bytes = serde_json::to_vec(&raw_metadata)?.len(); + assert!(raw_bytes < 1024 * 1024); + assert!(raw_bytes * 2 > 1024 * 1024); + assert!(raw_bytes * 2 < 2 * 1024 * 1024); + let fixture_metadata = raw_metadata.clone(); + Mock::given(method("POST")) + .and(path_regex("^/api/codex/ps/mcp/?$")) + .and(body_partial_json(json!({"method": "tools/call"}))) + .respond_with(move |request: &Request| { + let body: Value = serde_json::from_slice(&request.body).expect("MCP request"); + let query = body["params"]["arguments"]["query"] + .as_str() + .expect("query"); + ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + "jsonrpc": "2.0", "id": body["id"], + "result": { + "content": [{"type": "text", "text": format!("result for {query}")}], + "_meta": fixture_metadata, + "isError": false + } + })) + }) + .with_priority(/*p*/ 1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/analytics-events/events")) + .respond_with(ResponseTemplate::new(/*s*/ 200)) + .mount(&server) + .await; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call_with_namespace( + "first", + SEARCH_CALENDAR_NAMESPACE, + SEARCH_CALENDAR_LIST_TOOL, + r#"{"query":"first"}"#, + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_function_call_with_namespace( + "second", + SEARCH_CALENDAR_NAMESPACE, + SEARCH_CALENDAR_LIST_TOOL, + r#"{"query":"second"}"#, + ), + ev_completed("resp-2"), + ]), + sse(vec![ev_response_created("resp-3"), ev_completed("resp-3")]), + ], + ) + .await; + let mut builder = search_capable_apps_builder(apps.chatgpt_base_url).with_config(|config| { + config.analytics_enabled = Some(true); + config + .features + .enable(Feature::ExecutedToolCallMetadata) + .expect("enable executed tool call metadata"); + config + .features + .disable(Feature::CodeMode) + .expect("disable code mode"); + config + .features + .disable(Feature::CodeModeOnly) + .expect("disable code-mode-only tools"); + }); + // build() always selects the local test environment; all network destinations are the mock. + let test = builder.build(&server).await?; + assert!(test.codex.analytics_enabled()); + wait_for_mcp_server(&test.codex, CODEX_APPS_MCP_SERVER_NAME).await?; + test.submit_turn_with_approval_and_permission_profile( + "Use [$calendar](app://calendar) to list events twice.", + AskForApproval::OnRequest, + PermissionProfile::read_only(), + ) + .await?; + + let calls = recorded_apps_tool_calls(&server).await; + assert_eq!(calls.len(), 2); + for (call, query) in calls.iter().zip(["first", "second"]) { + assert_eq!(call["params"]["name"], "calendar_list_events"); + assert_eq!(call["params"]["arguments"], json!({"query": query})); + } + let requests = responses.requests(); + assert_eq!(requests.len(), 3); + let first_output = requests[2].function_call_output("first"); + let second_output = requests[2].function_call_output("second"); + let history = test.codex.conversation_history_snapshot().await; + let history = serde_json::to_value(history.items().collect::>())?; + let recorded_outputs = [&first_output, &second_output].map(|output| { + history + .as_array() + .expect("history items") + .iter() + .find(|item| { + item["type"] == "function_call_output" && item["call_id"] == output["call_id"] + }) + .expect("original tool output in history") + }); + for ((wire, recorded), query) in [&first_output, &second_output] + .into_iter() + .zip(recorded_outputs) + .zip(["first", "second"]) + { + assert_eq!(wire["output"], recorded["output"]); + assert!( + wire["output"] + .to_string() + .contains(&format!("result for {query}")) + ); + let recorded_metadata = tool_call_metadata(recorded.clone()); + let wire_metadata = tool_call_metadata(wire.clone()); + assert_eq!(recorded_metadata["tool_calls_complete"], true); + assert_eq!( + recorded_metadata["executed_tool_calls"][0]["name"], + format!("{SEARCH_CALENDAR_NAMESPACE}{SEARCH_CALENDAR_LIST_TOOL}") + ); + assert_eq!( + recorded_metadata["executed_tool_calls"][0]["arguments"], + json!({"query": query}) + ); + // A custom Responses provider must never receive internal result metadata. + assert!( + wire_metadata["executed_tool_calls"][0] + .get("tool_result_metadata") + .is_none() + ); + } + let first_metadata = tool_call_metadata(recorded_outputs[0].clone()); + let second_metadata = tool_call_metadata(recorded_outputs[1].clone()); + assert_eq!( + first_metadata["executed_tool_calls"][0]["tool_result_metadata"], + raw_metadata + ); + let marker = second_metadata["executed_tool_calls"][0]["tool_result_metadata"] + .as_str() + .expect("second result should have an omission marker"); + let overage = marker + .strip_prefix("omitted_due_to_size_limit (overage_bytes=") + .and_then(|value| value.strip_suffix(')')) + .and_then(|value| value.parse::().ok()) + .expect("marker should report omitted bytes"); + assert!(overage > 0); + Ok(()) +} diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index d51ded354c..318a4faa3d 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -48,6 +48,7 @@ pub use executed_tool_calls::MAX_TOOL_RESULT_SOURCE_FIELD_BYTES; pub use executed_tool_calls::ToolResultMetadata; pub use executed_tool_calls::ToolResultSource; pub use executed_tool_calls::ToolResultSources; +pub use executed_tool_calls::bound_executed_tool_calls_for_message; pub use executed_tool_calls::bound_executed_tool_calls_for_prompt; pub use executed_tool_calls::bound_executed_tool_calls_for_prompt_prioritizing_recent; pub use executed_tool_calls::executed_tool_call_metadata_bytes; diff --git a/codex-rs/protocol/src/models/executed_tool_calls.rs b/codex-rs/protocol/src/models/executed_tool_calls.rs index 7ec8afa81f..88e7775a30 100644 --- a/codex-rs/protocol/src/models/executed_tool_calls.rs +++ b/codex-rs/protocol/src/models/executed_tool_calls.rs @@ -121,17 +121,113 @@ fn bound_executed_tool_calls_for_prompt_with_priority( } clear_damaged_cell_completeness(items, &damaged_cells); - let metadata_bytes = |items: &[ResponseItem]| { - items.iter().fold(0_usize, |bytes, item| { - bytes.saturating_add(executed_tool_call_metadata_bytes(item)) - }) - }; - let mut total_metadata_bytes = metadata_bytes(items); - if total_metadata_bytes <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { + bound_executed_tool_calls_with_metadata_budget( + items, + MAX_EXECUTED_TOOL_CALL_METADATA_BYTES, + prioritize_recent, + /*whole_message*/ false, + ); +} + +/// Bounds optional observations to the space left in the actual outgoing message. +/// Ordinary output and unrelated item metadata are never changed. +pub fn bound_executed_tool_calls_for_message( + items: &mut [ResponseItem], + max_metadata_bytes: usize, +) { + bound_executed_tool_calls_with_metadata_budget( + items, + max_metadata_bytes, + /*prioritize_recent*/ false, + /*whole_message*/ true, + ); +} + +fn bound_executed_tool_calls_with_metadata_budget( + items: &mut [ResponseItem], + max_metadata_bytes: usize, + prioritize_recent: bool, + whole_message: bool, +) { + let mut damaged_cells = HashSet::new(); + + let total_metadata_bytes = metadata_bytes(items); + if total_metadata_bytes <= max_metadata_bytes { return; } - // 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 mut remaining_bytes = shed_generic_result_metadata( + items, + max_metadata_bytes, + prioritize_recent, + whole_message, + total_metadata_bytes, + ); + + // At the message limit, shed optional sources and recorded arguments before + // resource-access evidence. The prompt budget keeps its existing order. + if whole_message && metadata_bytes(items) > max_metadata_bytes { + let overage_bytes = + shed_result_sources(items, max_metadata_bytes, /*whole_message*/ true); + if overage_bytes > 0 { + truncate_call_arguments_to_fit( + items, + max_metadata_bytes, + overage_bytes, + &mut damaged_cells, + ); + } + remaining_bytes = metadata_bytes(items); + } + shed_remaining_result_metadata( + items, + max_metadata_bytes, + prioritize_recent, + whole_message, + remaining_bytes, + ); + + // Recheck the serialized request rather than relying on the incremental size accounting. + if metadata_bytes(items) <= max_metadata_bytes { + return; + } + // Omission markers are optional too; keep the original call budget if they cannot fit. + for item in items.iter_mut() { + item.clear_tool_result_metadata(); + } + + if metadata_bytes(items) <= max_metadata_bytes { + return; + } + // The whole-message path already handled sources and arguments before resources. + if !whole_message { + shed_result_sources(items, max_metadata_bytes, /*whole_message*/ false); + } + + if metadata_bytes(items) <= max_metadata_bytes { + return; + } + + distribute_remaining_budget( + items, + max_metadata_bytes, + prioritize_recent, + &mut damaged_cells, + ); + clear_damaged_cell_completeness(items, &damaged_cells); +} + +fn metadata_bytes(items: &[ResponseItem]) -> usize { + items.iter().fold(0_usize, |bytes, item| { + bytes.saturating_add(executed_tool_call_metadata_bytes(item)) + }) +} + +type ResultMetadataEntry<'a> = (usize, usize, usize, &'a mut ToolResultMetadata); + +fn result_metadata_by_size( + items: &mut [ResponseItem], + prioritize_recent: bool, +) -> Vec> { let mut result_metadata = Vec::new(); for (item_index, item) in items.iter_mut().enumerate() { if let Some(metadata) = item @@ -167,8 +263,20 @@ fn bound_executed_tool_calls_for_prompt_with_priority( result_metadata.sort_by_key(|(bytes, order, call_index, _)| { (std::cmp::Reverse(*bytes), *order, *call_index) }); + result_metadata +} + +fn shed_generic_result_metadata( + items: &mut [ResponseItem], + max_metadata_bytes: usize, + prioritize_recent: bool, + whole_message: bool, + mut total_metadata_bytes: usize, +) -> usize { + // Shed the largest values first so one large result does not discard unrelated small results. + let mut result_metadata = result_metadata_by_size(items, prioritize_recent); for (bytes, _, _, metadata) in &mut result_metadata { - if total_metadata_bytes <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { + if total_metadata_bytes <= max_metadata_bytes { break; } if metadata.retain_resource_access() { @@ -178,33 +286,54 @@ fn bound_executed_tool_calls_for_prompt_with_priority( total_metadata_bytes.saturating_sub(bytes.saturating_sub(retained_bytes)); *bytes = retained_bytes; } else { - let retained_bytes = metadata.omit_if_smaller( - *bytes, - total_metadata_bytes - MAX_EXECUTED_TOOL_CALL_METADATA_BYTES, - ); + let retained_bytes = + metadata.omit_if_smaller(*bytes, total_metadata_bytes - max_metadata_bytes); total_metadata_bytes = total_metadata_bytes.saturating_sub(*bytes - retained_bytes); *bytes = retained_bytes; } } + // Tiny generic values and markers can be cheaper than replacing them with a + // new marker. Remove their complete fields before sacrificing resource evidence. + if whole_message { + for (bytes, _, _, metadata) in &mut result_metadata { + if total_metadata_bytes <= max_metadata_bytes { + break; + } + if metadata.retain_resource_access() { + continue; + } + **metadata = ToolResultMetadata::default(); + total_metadata_bytes = total_metadata_bytes + .saturating_sub(bytes.saturating_add(b",\"tool_result_metadata\":".len())); + *bytes = 0; + } + } + total_metadata_bytes +} + +fn shed_remaining_result_metadata( + items: &mut [ResponseItem], + max_metadata_bytes: usize, + prioritize_recent: bool, + whole_message: bool, + mut total_metadata_bytes: usize, +) { + let mut result_metadata = result_metadata_by_size(items, prioritize_recent); // Resource evidence is optional too; it must not displace the call inventory. - result_metadata.sort_by_key(|(bytes, order, call_index, _)| { - (std::cmp::Reverse(*bytes), *order, *call_index) - }); for (bytes, _, _, metadata) in &mut result_metadata { - if total_metadata_bytes <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { + if total_metadata_bytes <= max_metadata_bytes { break; } // An omission marker can be larger than a small object, including an empty `_meta`. - let retained_bytes = metadata.omit_if_smaller( - *bytes, - total_metadata_bytes - MAX_EXECUTED_TOOL_CALL_METADATA_BYTES, - ); + let retained_bytes = + metadata.omit_if_smaller(*bytes, total_metadata_bytes - max_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)| { ( + whole_message && !metadata.is_omitted_due_to_size_limit(), std::cmp::Reverse(*bytes), !metadata.is_omitted_due_to_size_limit(), *order, @@ -212,47 +341,112 @@ fn bound_executed_tool_calls_for_prompt_with_priority( ) }); for (bytes, _, _, metadata) in result_metadata { - if total_metadata_bytes <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { + if total_metadata_bytes <= max_metadata_bytes { break; } + if metadata.is_none() { + continue; + } *metadata = ToolResultMetadata::default(); // The required name and arguments always precede this optional serialized field. total_metadata_bytes = total_metadata_bytes .saturating_sub(bytes.saturating_add(b",\"tool_result_metadata\":".len())); } +} - // Recheck the serialized request rather than relying on the incremental size accounting. - if metadata_bytes(items) <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { - return; - } - // Omission markers are optional too; keep the original call budget if they cannot fit. - for item in items.iter_mut() { - item.clear_tool_result_metadata(); - } - - if metadata_bytes(items) <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { - return; - } +fn shed_result_sources( + items: &mut [ResponseItem], + max_metadata_bytes: usize, + whole_message: bool, +) -> usize { // Source evidence is optional; dropping it must not discard calls or their completion proof. + let mut overage_bytes = metadata_bytes(items).saturating_sub(max_metadata_bytes); for item in items.iter_mut() { if let Some(metadata) = item .internal_chat_message_metadata_passthrough_mut() .and_then(Option::as_mut) { for call in metadata.executed_tool_calls.iter_mut().flatten() { - call.tool_result_sources = None; + if whole_message && overage_bytes == 0 { + break; + } + if let Some(sources) = call.tool_result_sources.take() { + let bytes = + serde_json::to_vec(&sources).map_or(usize::MAX, |value| value.len()); + overage_bytes = overage_bytes + .saturating_sub(bytes.saturating_add(b",\"tool_result_sources\":".len())); + } } } } - if metadata_bytes(items) <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { - return; - } + overage_bytes +} +fn truncate_call_arguments_to_fit( + items: &mut [ResponseItem], + max_metadata_bytes: usize, + mut overage_bytes: usize, + damaged_cells: &mut HashSet, +) { + // A large outgoing message should not lose the other tool names in a + // cell when replacing one call's arguments is enough to make it fit. + for index in 0..items.len() { + while overage_bytes > 0 { + let Some(metadata) = items[index] + .internal_chat_message_metadata_passthrough_mut() + .and_then(Option::as_mut) + else { + break; + }; + let mut arguments_changed = false; + for call in metadata.executed_tool_calls.iter_mut().flatten() { + if call.truncation().is_some() { + continue; + } + let original_bytes = + serde_json::to_vec(&call.arguments).map_or(usize::MAX, |value| value.len()); + let truncated = ExecutedToolCallArguments::Truncated { + truncation: ExecutedToolCallTruncation { + original_bytes, + max_bytes: original_bytes + .saturating_sub(overage_bytes) + .min(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES), + omitted_calls: None, + original_name_bytes: None, + }, + }; + let truncated_bytes = + serde_json::to_vec(&truncated).map_or(usize::MAX, |value| value.len()); + if truncated_bytes < original_bytes { + call.arguments = truncated; + arguments_changed = true; + metadata.tool_calls_complete = None; + damaged_cells.extend(metadata.cell_id.clone()); + break; + } + } + if !arguments_changed { + break; + } + // Losing one call's arguments invalidates every output in that + // cell. Count those removed fields before trimming another call. + clear_damaged_cell_completeness(items, damaged_cells); + overage_bytes = metadata_bytes(items).saturating_sub(max_metadata_bytes); + } + } +} + +fn distribute_remaining_budget( + items: &mut [ResponseItem], + max_metadata_bytes: usize, + prioritize_recent: bool, + damaged_cells: &mut HashSet, +) { let mut remaining_items = items .iter() .filter(|item| executed_tool_call_metadata_bytes(item) > 0) .count(); - let mut remaining_bytes = MAX_EXECUTED_TOOL_CALL_METADATA_BYTES; + let mut remaining_bytes = max_metadata_bytes; for item in items.iter_mut() { let item_bytes = executed_tool_call_metadata_bytes(item); if item_bytes == 0 { @@ -275,7 +469,6 @@ fn bound_executed_tool_calls_for_prompt_with_priority( remaining_bytes = remaining_bytes.saturating_sub(executed_tool_call_metadata_bytes(item)); remaining_items -= 1; } - clear_damaged_cell_completeness(items, &damaged_cells); } fn clear_damaged_cell_completeness(items: &mut [ResponseItem], damaged_cells: &HashSet) { 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 9e79d43007..9195a693af 100644 --- a/codex-rs/protocol/src/models/executed_tool_calls_tests.rs +++ b/codex-rs/protocol/src/models/executed_tool_calls_tests.rs @@ -623,6 +623,200 @@ fn wire_budget_helpers_preserve_resource_evidence_markers_and_small_values() { assert!(!item.has_tool_result_metadata()); } +#[test] +fn message_budget_keeps_smaller_resource_evidence_when_larger_evidence_must_be_shed() { + for (sizes, larger_index) in [([1024, 512], 0), ([512, 1024], 1)] { + let mut items = sizes.map(|size| { + let mut item = ResponseItem::from(ResponseInputItem::FunctionCallOutput { + call_id: format!("call-{size}"), + output: FunctionCallOutputPayload::from_text("unchanged output".to_string()), + }); + let mut call = + ExecutedToolCall::new(format!("tool_{size}"), serde_json::json!({"query": "kept"})); + call.set_tool_result_sources(ToolResultSources::new(vec![ToolResultSource { + r#type: "test_resource".to_string(), + id: format!("resource-{size}"), + }])); + call.set_tool_result_metadata(ToolResultMetadata::new(&serde_json::json!({ + "openai/resource_access": {"opaque": "é".repeat(size)}, + }))); + item.append_executed_tool_calls(vec![call]); + item.set_tool_call_cell_id(&format!("cell-{size}")); + item.mark_tool_calls_complete(); + item + }); + let mut expected = items.clone(); + for item in &mut expected { + first_executed_tool_call(item).unwrap().tool_result_sources = None; + } + let budget = expected + .iter() + .map(executed_tool_call_metadata_bytes) + .sum::() + - 128; + first_executed_tool_call(&mut expected[larger_index]) + .unwrap() + .set_tool_result_metadata(ToolResultMetadata::omitted_due_to_size_limit( + /*overage_bytes*/ 128, + )); + + bound_executed_tool_calls_for_message(&mut items, budget); + assert_eq!(items, expected); + assert!( + items + .iter() + .map(executed_tool_call_metadata_bytes) + .sum::() + <= budget + ); + bound_executed_tool_calls_for_message(&mut items, budget); + assert_eq!(items, expected); + } +} + +#[test] +fn message_budget_drops_only_the_sources_needed_to_fit() { + let mut items = ["first", "second"].map(|id| { + let mut item = ResponseItem::from(ResponseInputItem::FunctionCallOutput { + call_id: id.to_string(), + output: FunctionCallOutputPayload::from_text("unchanged output".to_string()), + }); + let mut call = + ExecutedToolCall::new(format!("tool_{id}"), serde_json::json!({"query": "kept"})); + call.set_tool_result_sources(ToolResultSources::new(vec![ToolResultSource { + r#type: "test_resource".to_string(), + id: id.to_string(), + }])); + call.set_tool_result_metadata(ToolResultMetadata::new(&serde_json::json!({ + "openai/resource_access": {"resource_coverage": "complete", "resources": []}, + }))); + item.append_executed_tool_calls(vec![call]); + item.set_tool_call_cell_id(id); + item.mark_tool_calls_complete(); + item + }); + let mut expected = items.clone(); + first_executed_tool_call(&mut expected[0]) + .unwrap() + .tool_result_sources = None; + let budget = expected + .iter() + .map(executed_tool_call_metadata_bytes) + .sum::(); + assert!( + items + .iter() + .map(executed_tool_call_metadata_bytes) + .sum::() + > budget + ); + + bound_executed_tool_calls_for_message(&mut items, budget); + assert_eq!(items, expected); + bound_executed_tool_calls_for_message(&mut items, budget); + assert_eq!(items, expected); +} + +#[test] +fn message_budget_removes_generic_fields_before_resources_in_the_same_output() { + for generic in [ + serde_json::json!({}), + serde_json::json!(null), + serde_json::json!("omitted_due_to_size_limit"), + ] { + let mut item = output("exec"); + for metadata in [ + serde_json::json!({"openai/resource_access": {"opaque": "é".repeat(256)}}), + generic.clone(), + generic, + ] { + let mut call = ExecutedToolCall::new("read".to_string(), serde_json::json!({})); + call.set_tool_result_metadata(ToolResultMetadata::new(&metadata)); + item.append_executed_tool_calls(vec![call]); + } + item.set_tool_call_cell_id("cell"); + item.mark_tool_calls_complete(); + let mut expected = item.clone(); + expected + .ensure_tool_call_metadata() + .unwrap() + .executed_tool_calls + .as_mut() + .unwrap()[1] + .set_tool_result_metadata(ToolResultMetadata::default()); + let budget = executed_tool_call_metadata_bytes(&expected); + bound_executed_tool_calls_for_message(std::slice::from_mut(&mut item), budget); + assert_eq!(item, expected); + } +} + +#[test] +fn message_budget_preserves_names_and_turn_metadata_and_invalidates_shared_cell_completion() { + let arguments = serde_json::json!({"payload": "é".repeat(256)}); + let argument_bytes = serde_json::to_vec(&arguments).unwrap().len(); + let mut exec = output("exec"); + exec.append_executed_tool_calls(vec![ + ExecutedToolCall::new("first_é\"".to_string(), arguments.clone()), + ExecutedToolCall::new("second".to_string(), arguments), + ]); + first_executed_tool_call(&mut exec) + .unwrap() + .set_tool_result_metadata(ToolResultMetadata::new(&serde_json::json!({ + "openai/resource_access": {"resources": ["R1"]} + }))); + let mut wait = output("wait"); + for item in [&mut exec, &mut wait] { + item.set_tool_call_cell_id("cell_é\""); + item.mark_tool_calls_complete(); + let metadata = item.ensure_tool_call_metadata().unwrap(); + metadata.turn_id = Some("turn_é\"".to_string()); + metadata.create_time = Some(serde_json::Number::from(123)); + } + let original = vec![exec, wait]; + let original_metadata_bytes = original + .iter() + .map(executed_tool_call_metadata_bytes) + .sum::(); + // The middle case needs the sibling wait's completeness bytes to fit after + // truncating only the first call; the second call's arguments must survive. + for budget in [ + original_metadata_bytes - 64, + original_metadata_bytes - (argument_bytes - 32), + 0, + ] { + let mut items = original.clone(); + let mut expected = original.clone(); + if budget == 0 { + for item in &mut expected { + item.clear_executed_tool_calls(); + } + } else { + first_executed_tool_call(&mut expected[0]) + .unwrap() + .set_truncation( + argument_bytes, + argument_bytes - (original_metadata_bytes - budget), + /*omitted_calls*/ None, + ); + for item in &mut expected { + item.clear_tool_calls_complete(); + } + } + bound_executed_tool_calls_for_message(&mut items, budget); + assert_eq!(items, expected); + let bounded_metadata_bytes = items + .iter() + .map(executed_tool_call_metadata_bytes) + .sum::(); + assert!(bounded_metadata_bytes <= budget); + assert_eq!( + serde_json::to_vec(&original).unwrap().len() + - serde_json::to_vec(&items).unwrap().len(), + original_metadata_bytes - bounded_metadata_bytes, + ); + } +} + #[test] fn raw_result_metadata_is_shed_before_sources_calls_or_completion() -> Result<()> { // Keep each result below its capture limit while filling the larger request budget. @@ -773,6 +967,9 @@ fn result_metadata_shedding_keeps_smaller_results_within_one_output() { for bound in [ bound_executed_tool_calls_for_prompt, bound_executed_tool_calls_for_prompt_prioritizing_recent, + |items: &mut [ResponseItem]| { + bound_executed_tool_calls_for_message(items, MAX_EXECUTED_TOOL_CALL_METADATA_BYTES) + }, ] { let mut bounded = item.clone(); bound(std::slice::from_mut(&mut bounded));