diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 4f809b0a1d..12c341b5a9 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -610,6 +610,92 @@ fn responses_request_limits_internal_metadata_to_resolved_first_party_https_endp Ok(()) } +#[test] +fn responses_request_preserves_smaller_result_metadata_under_aggregate_budget() -> anyhow::Result<()> +{ + let result_metadata = [ + json!({ "payload": "l".repeat(31 * 1024) }), + json!({ "payload": "m".repeat(20 * 1024) }), + json!({ "payload": "m".repeat(20 * 1024) }), + json!({ "payload": "m".repeat(20 * 1024) }), + json!({ "payload": "m".repeat(20 * 1024) }), + json!({ "payload": "m".repeat(20 * 1024) }), + json!({ "status": "ok" }), + ]; + let sizes = result_metadata + .iter() + .map(|metadata| serde_json::to_vec(metadata).unwrap().len()) + .collect::>(); + assert!(sizes.iter().all(|bytes| *bytes < 32 * 1024)); + assert!(sizes.iter().sum::() > 128 * 1024); + let mut history = Vec::new(); + for (index, metadata) in result_metadata.iter().enumerate() { + let id = format!("tool-call-{index}"); + history.push(serde_json::from_value(json!({ + "type": "function_call", "call_id": id, "name": "test_tool", + "arguments": json!({ "query": "keep" }).to_string(), + }))?); + let mut output = output_with_tool_result_metadata(ToolResultMetadata::new(metadata)); + let ResponseItem::FunctionCallOutput { call_id, .. } = &mut output else { + unreachable!("helper returns a function call output"); + }; + *call_id = Some(id); + 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); + let recorder = + crate::tools::ExecutedToolCalls::new(&features, &codex_history::InitialHistory::New); + let mut prompt = Prompt { + input: history.clone(), + ..Default::default() + }; + // Follow the sampling path: budget the request copy before client serialization. + recorder.attach_to_prompt(&mut prompt.input, &mut Default::default()); + let provider = + ModelProviderInfo::create_openai_provider(Some("https://api.openai.com/v1".to_string())); + let api_provider = provider.to_api_provider(/*auth_mode*/ None)?; + let mut client = test_model_client(SessionSource::Cli); + Arc::get_mut(&mut client.state) + .expect("test client should have unique session state") + .provider = create_model_provider(provider, /*auth_manager*/ None); + let responses_metadata = test_responses_metadata_for_client( + &client, + /*turn_id*/ None, + format!("{}:0", client.state.thread_id), + /*parent_thread_id*/ None, + TestCodexResponsesRequestKind::Turn, + ); + let request = client.build_responses_request( + &prompt, + &test_model_info(), + /*effort*/ None, + codex_protocol::config_types::ReasoningSummary::None, + /*service_tier*/ None, + &responses_metadata, + super::is_internal_metadata_destination(&api_provider), + )?; + let body = serde_json::to_value(&request)?; + // Whole-input equality covers bindings, arguments, results, sources and completion too. + assert_eq!(body["input"], expected_input); + let mut without_metadata = body.clone(); + for item in without_metadata["input"].as_array_mut().unwrap() { + item.as_object_mut() + .unwrap() + .remove("internal_chat_message_metadata_passthrough"); + } + let metadata_bytes = + serde_json::to_vec(&body)?.len() - serde_json::to_vec(&without_metadata)?.len(); + assert!(metadata_bytes <= 128 * 1024); + assert_eq!(serde_json::to_value(&history)?, original_history); + Ok(()) +} + #[test] fn websocket_incremental_reuse_tracks_raw_result_metadata() -> anyhow::Result<()> { let provider = diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index 088d1234cf..bbefb411b8 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..5 { + for index in 0..20 { 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(), 5); + assert_eq!(outputs.len(), 20); 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 > 32 * 1024); + assert!(metadata_bytes > 128 * 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 db8f8e6aa6..418d399999 100644 --- a/codex-rs/core/src/tools/executed_tool_calls.rs +++ b/codex-rs/core/src/tools/executed_tool_calls.rs @@ -344,6 +344,10 @@ 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); + if bytes > available { + item.retain_tool_resource_access(); + bytes = executed_tool_call_metadata_bytes(item); + } if bytes > available { item.clear_tool_result_metadata(); bytes = executed_tool_call_metadata_bytes(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 7f79e1d227..f4d3176161 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 @@ -154,7 +154,7 @@ fn failed_compaction_preserves_observations_absent_from_trimmed_retries() { assert!(recorder.record_tool_result_metadata( &source, &format!("nested-{index}"), - &json!({"provider": "x".repeat(7 * 1024)}), + &json!({"provider": "x".repeat(24 * 1024)}), )); } } @@ -189,7 +189,7 @@ 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..6 { + for index in 0..24 { let origin = format!("exec-{index}"); let cell = CellId::new(format!("cell-{index}")); recorder.start_cell(&cell, &origin); @@ -204,7 +204,7 @@ fn compaction_bounds_code_mode_metadata_without_rebudgeting_direct_history() { history.extend([exec_input(&origin), exec_output(&origin)]); } let direct_start = history.len(); - for index in 0..6 { + for index in 0..24 { let mut item = output(&format!("direct-{index}")); let mut call = ExecutedToolCall::new("direct_tool".to_string(), arguments.clone()); call.set_tool_result_metadata(ToolResultMetadata::new(&json!({"id": index}))); @@ -219,7 +219,7 @@ fn compaction_bounds_code_mode_metadata_without_rebudgeting_direct_history() { .iter() .map(executed_tool_call_metadata_bytes) .sum::() - > MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT + > 128 * 1024 ); let mut compact = history.clone(); @@ -230,14 +230,14 @@ fn compaction_bounds_code_mode_metadata_without_rebudgeting_direct_history() { .iter() .map(executed_tool_call_metadata_bytes) .sum::() - <= MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT + <= 128 * 1024 ); - let mut newest = exec_output("exec-5"); + 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-5"); + newest.set_tool_call_cell_id("exec-23"); newest.mark_tool_calls_complete(); assert_eq!(compact[direct_start - 1], newest); assert_ne!(tool_calls_complete(&compact[1]), Some(true)); @@ -269,15 +269,21 @@ fn compaction_keeps_direct_outputs_in_code_mode_binding_validation() { #[test] fn direct_retained_metadata_budget_sheds_results_then_calls_across_refresh() { let recorder = new_recorder(InitialHistory::New); + let resource_access = json!({"resource_coverage": "complete", "resources": []}); let mut call = ExecutedToolCall::new("test_tool".to_string(), json!({"argument": "kept"})); call.set_tool_result_metadata(ToolResultMetadata::new(&json!({ "provider": "x".repeat(1024), + "openai/resource_access": resource_access, }))); let mut first = output("first"); recorder.attach_direct_call_to_output( &mut first, Some((call.clone(), recorder.reserve_direct_call().unwrap())), ); + let mut expected_first = output("first"); + expected_first.append_executed_tool_calls(vec![call.clone()]); + expected_first.mark_tool_calls_complete(); + assert_eq!(first, expected_first); let full_bytes = executed_tool_call_metadata_bytes(&first); assert!(full_bytes > 0); assert_eq!( @@ -287,10 +293,39 @@ fn direct_retained_metadata_budget_sheds_results_then_calls_across_refresh() { full_bytes, ); + let mut resource_only_call = call.clone(); + resource_only_call.set_tool_result_metadata(ToolResultMetadata::new(&json!({ + "openai/resource_access": resource_access, + }))); + let mut resource_only_output = ResponseItem::from(ResponseInputItem::FunctionCallOutput { + call_id: "resource-only".to_string(), + output: FunctionCallOutputPayload::from_text("ordinary result".to_string()), + }); + let mut expected_resource_only = resource_only_output.clone(); + expected_resource_only.append_executed_tool_calls(vec![resource_only_call]); + expected_resource_only.mark_tool_calls_complete(); + let resource_only_bytes = executed_tool_call_metadata_bytes(&expected_resource_only); + assert!(full_bytes > resource_only_bytes); + recorder.retained_direct_metadata_bytes.store( + MAX_RETAINED_DIRECT_METADATA_BYTES - resource_only_bytes, + Ordering::Relaxed, + ); + recorder.clone().attach_direct_call_to_output( + &mut resource_only_output, + Some((call.clone(), recorder.reserve_direct_call().unwrap())), + ); + assert_eq!(resource_only_output, expected_resource_only); + assert_eq!( + recorder + .retained_direct_metadata_bytes + .load(Ordering::Relaxed), + MAX_RETAINED_DIRECT_METADATA_BYTES, + ); + let mut without_result = first.clone(); without_result.clear_tool_result_metadata(); let call_bytes = executed_tool_call_metadata_bytes(&without_result); - assert!(full_bytes > call_bytes); + assert!(resource_only_bytes > call_bytes); recorder.retained_direct_metadata_bytes.store( MAX_RETAINED_DIRECT_METADATA_BYTES - call_bytes, Ordering::Relaxed, @@ -302,6 +337,13 @@ fn direct_retained_metadata_budget_sheds_results_then_calls_across_refresh() { &mut second, Some((call.clone(), recorder.reserve_direct_call().unwrap())), ); + let mut expected_second = output("second"); + expected_second.append_executed_tool_calls(vec![ExecutedToolCall::new( + "test_tool".to_string(), + json!({"argument": "kept"}), + )]); + expected_second.mark_tool_calls_complete(); + assert_eq!(second, expected_second); let second_metadata = second .executed_tool_call_metadata() .expect("call remains when only its result metadata exceeds the budget"); @@ -613,7 +655,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 <= MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT); + assert!(retained_bytes <= 128 * 1024); let metadata = prompt .iter() @@ -901,7 +943,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 = 32 * 1024; + let request_budget = 128 * 1024; let recorder = new_recorder(InitialHistory::New); let cell = CellId::new("clean-overflow-cell".to_string()); recorder.start_cell(&cell, "exec-clean"); @@ -912,7 +954,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..6) + (0..24) .map(|index| ExecutedToolCall::new(format!("direct-{index}"), arguments.clone())) .collect(), ); @@ -1167,7 +1209,22 @@ fn request_truncation_prevents_completion_after_compaction() { ); } - let mut initial = [exec_input("exec"), exec_output("exec")]; + let mut initial = vec![exec_input("exec"), exec_output("exec")]; + for index in 0..3 { + let origin = format!("other-{index}"); + let other_cell = CellId::new(origin.clone()); + recorder.start_cell(&other_cell, &origin); + for nested in 0..4 { + recorder.record_nested_tool_call( + other_cell.clone(), + format!("other-{index}-{nested}"), + ExecutedToolCall::new("nested_tool".to_string(), arguments.clone()), + MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES, + ); + } + recorder.finish_cell_recording(&other_cell); + initial.extend([exec_input(&origin), exec_output(&origin)]); + } assert!(recorder.attach_pending_to_prompt(&mut initial, &mut HashMap::new())); assert!( initial[1] @@ -1266,7 +1323,7 @@ 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..64 { + for index in 0..256 { let mut item = output(&format!("untracked-{index}")); item.append_executed_tool_calls(vec![ExecutedToolCall::new( "other_tool".to_string(), @@ -1291,11 +1348,11 @@ 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_ARGUMENT_BYTES - 256; + let argument_bytes = MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT / 5 - 256; let arguments = json!({ "payload": "x".repeat(argument_bytes - r#"{"payload":""}"#.len()), }); - for index in 0..4 { + for index in 0..5 { recorder.record_nested_tool_call( cell_id.clone(), format!("nested-{index}"), @@ -1312,12 +1369,15 @@ fn result_metadata_shedding_preserves_completion_after_compaction() { &mut [exec_input("exec"), exec_output("exec")], &mut retry_cache, ); - // This snapshot fits alone, but pushes the four calls over the aggregate budget. - assert!(recorder.record_tool_result_metadata( - &source, - "nested-0", - &json!({ "provider_data": "x".repeat(4 * 1024) }), - )); + // 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] { + assert!(recorder.record_tool_result_metadata( + &source, + &format!("nested-{index}"), + &large_metadata, + )); + } assert!(recorder.record_tool_result_metadata( &source, @@ -1330,10 +1390,14 @@ 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); 4]; - for call in &mut expected_calls[..2] { - call.set_tool_result_metadata(ToolResultMetadata::new(&json!("omitted_due_to_size_limit"))); + let mut expected_calls = vec![ExecutedToolCall::new("nested_tool".to_string(), arguments); 5]; + for call in &mut expected_calls[2..] { + 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" }))); 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/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index f06e0c2246..b753df9674 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -1402,7 +1402,7 @@ async fn code_mode_tool_call_completeness_is_private_and_opt_in( r#" const args = { barrier: { id: "", participants: 1 } }; args.barrier.id = "x".repeat(8192 - JSON.stringify(args).length); -for (let index = 0; index < 4; index++) await tools.test_sync_tool(args); +for (let index = 0; index < 17; index++) await tools.test_sync_tool(args); text("done"); yield_control(); await new Promise(() => {}); @@ -2010,6 +2010,379 @@ enum ResultMetadataAnalytics { HostDisabled(Option), } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn direct_result_metadata_retained_budget_preserves_resource_access() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = responses::start_mock_server().await; + let arguments = serde_json::json!({ "search": "retained-budget" }); + let resource_access = serde_json::json!({ + "schema_version": 1, + "resource_coverage": "complete", + "resources": [], + }); + let resource_only = serde_json::json!({ "openai/resource_access": resource_access }); + let result_metadata = serde_json::json!({ + "payload": "x".repeat(31 * 1024), + "openai/resource_access": resource_access, + }); + let metadata_bytes = serde_json::to_vec(&result_metadata)?.len(); + assert!(metadata_bytes < 32 * 1024); + assert!(metadata_bytes * 33 > 1024 * 1024); + let apps_server = mount_result_metadata_app( + &server, + Some(result_metadata.clone()), + /*is_error*/ false, + ) + .await?; + let mut builder = + result_metadata_apps_builder(apps_server.chatgpt_base_url, "user@example.com").with_config( + |config| { + config.features.disable(Feature::CodeModeOnly).unwrap(); + config.features.disable(Feature::CodeMode).unwrap(); + }, + ); + let test = builder.build_with_auto_env(&server).await?; + let call_ids = (0..33) + .map(|index| format!("call-{index}")) + .collect::>(); + let mut events = call_ids[..32] + .iter() + .map(|call_id| { + responses::ev_function_call_with_namespace( + call_id, + "mcp__codex_apps__messagesearch", + RESULT_METADATA_TOOL, + &arguments.to_string(), + ) + }) + .collect::>(); + events.push(ev_completed("resp-1")); + let initial = responses::mount_sse_once(&server, sse(events)).await; + // Cross the retained-history limit after the first batch has finished, regardless of order. + let overflow = responses::mount_sse_once( + &server, + sse(vec![ + responses::ev_function_call_with_namespace( + &call_ids[32], + "mcp__codex_apps__messagesearch", + RESULT_METADATA_TOOL, + &arguments.to_string(), + ), + ev_completed("resp-2"), + ]), + ) + .await; + let completion = responses::mount_sse_once(&server, sse(vec![ev_completed("resp-3")])).await; + test.submit_turn("Search the connected app 33 times") + .await?; + let initial = initial.single_request(); + assert!(initial.has_message_with_input_texts("user", |texts| { + texts == ["Search the connected app 33 times"] + })); + assert!( + !tool_names(&initial.body_json()) + .iter() + .any(|name| name == "exec") + ); + for (request, expected_calls) in [ + (overflow.single_request(), &call_ids[..32]), + (completion.single_request(), &call_ids[..]), + ] { + for call_id in expected_calls { + let output = request.function_call_output(call_id); + assert_eq!( + &output["output"].as_array().expect("direct output content")[1..], + &[serde_json::json!({ + "type": "input_text", + "text": RESULT_METADATA_PRIVATE_RESULT, + })], + ); + } + } + let calls = recorded_apps_tool_calls(&server).await; + assert_eq!(calls.len(), call_ids.len()); + for call in calls { + 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. + let history = test.codex.conversation_history_snapshot().await; + assert!( + history + .items() + .map(codex_protocol::models::executed_tool_call_metadata_bytes) + .sum::() + <= 1024 * 1024 + ); + let outputs = history + .items() + .filter(|item| matches!(item, ResponseItem::FunctionCallOutput { .. })) + .collect::>(); + assert_eq!(outputs.len(), call_ids.len()); + for (index, call_id) in call_ids.iter().enumerate() { + let item = outputs + .iter() + .find(|item| matches!(item, ResponseItem::FunctionCallOutput { call_id: Some(id), .. } if id == call_id)) + .expect("retained Direct output"); + if let ResponseItem::FunctionCallOutput { output, .. } = item { + assert_eq!(output.success, Some(true)); + } + let output = serde_json::to_value(item)?; + let metadata = &output["internal_chat_message_metadata_passthrough"]; + assert_eq!(metadata["tool_calls_complete"], true); + let calls = metadata["executed_tool_calls"].as_array().unwrap(); + assert_eq!(calls.len(), 1); + assert!( + calls[0]["name"] + .as_str() + .unwrap() + .ends_with(RESULT_METADATA_TOOL) + ); + assert_eq!(calls[0]["arguments"], arguments); + assert_eq!( + calls[0]["tool_result_metadata"], + if index < 32 { + &result_metadata + } else { + &resource_only + } + .clone(), + ); + } + test.codex.shutdown_and_wait().await?; + Ok(()) +} + +#[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<()> { + skip_if_no_network!(Ok(())); + let server = responses::start_mock_server().await; + let arguments = [ + serde_json::json!({ "search": "large" }), + serde_json::json!({ "search": "medium-1" }), + serde_json::json!({ "search": "medium-2" }), + serde_json::json!({ "search": "medium-3" }), + serde_json::json!({ "search": "medium-4" }), + serde_json::json!({ "search": "medium-5" }), + serde_json::json!({ "search": "oversized" }), + serde_json::json!({ "search": "small" }), + ]; + let resource_access = serde_json::json!({ + "schema_version": 1, + "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), + "openai/resource_access": resource_access, + }), + serde_json::json!({ "payload": "m".repeat(20 * 1024) }), + serde_json::json!({ "payload": "m".repeat(20 * 1024) }), + serde_json::json!({ "payload": "m".repeat(20 * 1024) }), + serde_json::json!({ "payload": "m".repeat(20 * 1024) }), + serde_json::json!({ "payload": "m".repeat(20 * 1024) }), + serde_json::json!({ + "payload": "o".repeat(40 * 1024), + "openai/resource_access": resource_access, + }), + serde_json::json!({ "status": "ok" }), + ]; + let metadata_sizes = result_metadata + .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); + for (arguments, metadata) in arguments.iter().zip(&result_metadata) { + let result = serde_json::json!({ + "content": [{ "type": "text", "text": RESULT_METADATA_PRIVATE_RESULT }], + "isError": false, + "_meta": metadata, + }); + Mock::given(method("POST")) + .and(path("/api/codex/ps/mcp")) + .and(body_partial_json(serde_json::json!({ + "method": "tools/call", + "params": { "arguments": arguments }, + }))) + .respond_with(move |request: &Request| { + let request: Value = serde_json::from_slice(&request.body).unwrap(); + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "jsonrpc": "2.0", + "id": request["id"], + "result": result, + })) + }) + .with_priority(/*p*/ 1) + .mount(&server) + .await; + } + // Exact argument matches take precedence over the shared fallback mounted afterward. + let apps_server = mount_result_metadata_app( + &server, /*result_metadata*/ None, /*is_error*/ false, + ) + .await?; + let direct = tool_mode == ToolMode::Direct; + let mut builder = + result_metadata_apps_builder(apps_server.chatgpt_base_url, "user@example.com").with_config( + move |config| { + if direct { + config.features.disable(Feature::CodeModeOnly).unwrap(); + config.features.disable(Feature::CodeMode).unwrap(); + } + }, + ); + let test = builder.build_with_auto_env(&server).await?; + let call_ids = [ + "call-1", "call-2", "call-3", "call-4", "call-5", "call-6", "call-7", "call-8", + ]; + let mut events = if direct { + arguments + .iter() + .zip(call_ids) + .map(|(arguments, call_id)| { + responses::ev_function_call_with_namespace( + call_id, + "mcp__codex_apps__messagesearch", + RESULT_METADATA_TOOL, + &arguments.to_string(), + ) + }) + .collect::>() + } else { + let arguments = serde_json::to_string(&arguments)?; + let code = format!( + "const tool = ALL_TOOLS.find(({{ name }}) => name.endsWith(\"{RESULT_METADATA_TOOL}\")); \ + const results = []; \ + for (const args of {arguments}) {{ \ + const result = await tools[tool.name](args); \ + results.push({{ content: result.content, isError: Boolean(result.isError), \ + hasMeta: Object.hasOwn(result, \"_meta\") }}); \ + }} \ + text(JSON.stringify(results));" + ); + vec![ev_custom_tool_call("call-1", "exec", &code)] + }; + events.push(ev_completed("resp-1")); + let initial = responses::mount_sse_once(&server, sse(events)).await; + let follow_up = responses::mount_sse_once(&server, sse(vec![ev_completed("resp-2")])).await; + test.submit_turn("Search the connected app eight times") + .await?; + let initial = initial.single_request(); + assert!(initial.has_message_with_input_texts("user", |texts| { + texts == ["Search the connected app eight times"] + })); + assert_eq!( + tool_names(&initial.body_json()) + .iter() + .any(|name| name == "exec"), + !direct, + ); + let request = follow_up.single_request(); + let mut actual_arguments = recorded_apps_tool_calls(&server) + .await + .into_iter() + .map(|call| { + assert_eq!(call["params"]["name"], RESULT_METADATA_TOOL); + call["params"]["arguments"].clone() + }) + .collect::>(); + actual_arguments.sort_by_key(Value::to_string); + assert_eq!(actual_arguments, arguments); + if direct { + for call_id in call_ids { + let output = request.function_call_output(call_id); + assert_eq!( + &output["output"].as_array().expect("direct output content")[1..], + &[serde_json::json!({ + "type": "input_text", + "text": RESULT_METADATA_PRIVATE_RESULT, + })], + ); + assert!(!output["output"].to_string().contains("payload")); + } + } else { + let (body, success) = custom_tool_output_body_and_success(&request, "call-1"); + assert_ne!(success, Some(false), "Code Mode failed: {body}"); + assert_eq!( + serde_json::from_str::(&body)?, + serde_json::json!(vec![ + serde_json::json!({ + "content": [{ "type": "text", "text": RESULT_METADATA_PRIVATE_RESULT }], + "isError": false, + "hasMeta": false, + }); + arguments.len() + ]), + ); + } + let captured = codex_core::test_support::history_with_tool_call_metadata(&test.codex).await; + if direct { + for item in &captured { + if let ResponseItem::FunctionCallOutput { output, .. } = item { + assert_eq!(output.success, Some(true)); + } + } + } + assert!( + captured + .iter() + .map(codex_protocol::models::executed_tool_call_metadata_bytes) + .sum::() + <= 128 * 1024 + ); + let captured = serde_json::to_value(captured)?; + for (input, expected_metadata) in [ + // Custom inference endpoints must strip raw metadata, including omission markers. + (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(), + ]), + ), + ] { + let mut calls = Vec::new(); + for output in &input { + let metadata = &output["internal_chat_message_metadata_passthrough"]; + if let Some(recorded) = metadata["executed_tool_calls"].as_array() { + assert_eq!(metadata["tool_calls_complete"], true); + calls.extend(recorded.iter().cloned()); + } + } + assert_eq!(calls.len(), arguments.len()); + let expected_calls = arguments + .iter() + .enumerate() + .map(|(index, arguments)| { + let name = calls[index]["name"].as_str().unwrap(); + assert!(name.ends_with(RESULT_METADATA_TOOL)); + let mut call = serde_json::json!({ "name": name, "arguments": arguments }); + if let Some(metadata) = &expected_metadata { + call["tool_result_metadata"] = metadata[index].clone(); + } + call + }) + .collect::>(); + assert_eq!(calls, expected_calls); + } + test.codex.shutdown_and_wait().await?; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[test_case(true, true, true, false, "user@example.com", ToolMode::CodeModeOnly, ResultMetadataAnalytics::Config(None); "copies_full_metadata_without_rules")] #[test_case(true, true, true, true, "user@example.com", ToolMode::CodeModeOnly, ResultMetadataAnalytics::Config(None); "accepted_error_keeps_metadata")] @@ -2067,8 +2440,8 @@ async fn result_metadata_follows_call_binding( .with_config(move |config| { config.analytics_enabled = analytics_enabled; if direct { - config.features.disable(Feature::CodeMode).unwrap(); config.features.disable(Feature::CodeModeOnly).unwrap(); + config.features.disable(Feature::CodeMode).unwrap(); } if !metadata_enabled { config diff --git a/codex-rs/core/tests/suite/direct_tool_metadata.rs b/codex-rs/core/tests/suite/direct_tool_metadata.rs index 9c3bffaf71..a6bb8059e4 100644 --- a/codex-rs/core/tests/suite/direct_tool_metadata.rs +++ b/codex-rs/core/tests/suite/direct_tool_metadata.rs @@ -203,7 +203,7 @@ async fn direct_call_metadata_during_compaction_respects_provider_support( } #[test_case(false, 0; "metadata disabled")] -#[test_case(true, 6; "request budget exceeded")] +#[test_case(true, 24; "request budget exceeded")] #[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 = 32 * 1024; + let request_budget = 128 * 1024; let server = start_mock_server().await; let mut builder = test_codex().with_config(move |config| { configure_search_capable_model(config); diff --git a/codex-rs/protocol/src/models/executed_tool_calls.rs b/codex-rs/protocol/src/models/executed_tool_calls.rs index 227042c596..3fe38189be 100644 --- a/codex-rs/protocol/src/models/executed_tool_calls.rs +++ b/codex-rs/protocol/src/models/executed_tool_calls.rs @@ -1,3 +1,4 @@ +use std::collections::BTreeMap; use std::collections::HashSet; use schemars::JsonSchema; @@ -14,7 +15,10 @@ const MAX_TOOL_RESULT_SOURCES: usize = 32; /// Maximum UTF-8 bytes for each source's `type` and `id` separately, not the source list. pub const MAX_TOOL_RESULT_SOURCE_FIELD_BYTES: usize = 128; /// Maximum serialized warehouse-only attempted-tool metadata in one request. -const MAX_EXECUTED_TOOL_CALL_METADATA_BYTES: usize = 32 * 1024; +const 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 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 = b"\"internal_chat_message_metadata_passthrough\":".len(); @@ -125,23 +129,100 @@ fn bound_executed_tool_calls_for_prompt_with_priority( bytes.saturating_add(executed_tool_call_metadata_bytes(item)) }) }; - if metadata_bytes(items) <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { + let mut total_metadata_bytes = metadata_bytes(items); + if total_metadata_bytes <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { return; } // Raw result metadata must not displace existing source evidence, calls or completion proof. - for item in items.iter_mut() { + // 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 .internal_chat_message_metadata_passthrough_mut() .and_then(Option::as_mut) { - for call in metadata.executed_tool_calls.iter_mut().flatten() { + // The retained-history entry point has already reversed the items. Break size ties + // by evicting older outputs, while keeping the original call order within an output. + let eviction_order = if prioritize_recent { + usize::MAX - item_index + } else { + item_index + }; + for (call_index, call) in metadata + .executed_tool_calls + .iter_mut() + .flatten() + .enumerate() + { if call.tool_result_metadata.is_some() { - call.tool_result_metadata = ToolResultMetadata::omitted_due_to_size_limit(); + let bytes = serde_json::to_vec(&call.tool_result_metadata) + .map_or(usize::MAX, |value| value.len()); + result_metadata.push(( + bytes, + eviction_order, + call_index, + &mut call.tool_result_metadata, + )); } } } } + 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 { + break; + } + if metadata.retain_resource_access() { + let retained_bytes = + serde_json::to_vec(&**metadata).map_or(usize::MAX, |value| value.len()); + 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; + } + } + // 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 { + 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; + } + } + // 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, + *order, + *call_index, + ) + }); + for (bytes, _, _, metadata) in result_metadata { + if total_metadata_bytes <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { + break; + } + *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; } @@ -247,7 +328,7 @@ pub struct ExecutedToolCall { tool_result_metadata: ToolResultMetadata, } -/// An entire MCP result's `_meta`, or a string marker when omitted due to the size limit. +/// MCP result metadata, optionally reduced to resource evidence, or a size-omission marker. #[derive(Clone, Default, Serialize, PartialEq, Eq)] #[serde(transparent)] pub struct ToolResultMetadata(Option); @@ -259,14 +340,35 @@ impl std::fmt::Debug for ToolResultMetadata { } impl ToolResultMetadata { - /// Bounds serialization before cloning arbitrary MCP metadata; no keys are filtered. + /// Bounds serialization before cloning, retaining resource evidence if the full result is too large. pub fn new(metadata: &serde_json::Value) -> Self { - let mut limit = MetadataSizeLimit(MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); + let mut limit = MetadataSizeLimit(MAX_TOOL_RESULT_METADATA_BYTES); if serde_json::to_writer(&mut limit, metadata).is_ok() { - Self(Some(metadata.clone())) - } else { - Self::omitted_due_to_size_limit() + 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() + } + + /// Leaves generic metadata unchanged when it has no resource-access field. + fn retain_resource_access(&mut self) -> bool { + let Some(serde_json::Value::Object(metadata)) = self.0.as_mut() else { + return false; + }; + if !metadata.contains_key(RESOURCE_ACCESS_METADATA_KEY) { + return false; + } + metadata.retain(|key, _| key == RESOURCE_ACCESS_METADATA_KEY); + true } fn omitted_due_to_size_limit() -> Self { @@ -524,6 +626,20 @@ impl ResponseItem { } } + /// Reduces optional results to resource evidence before considering call-inventory loss. + pub fn retain_tool_resource_access(&mut self) { + 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() { + if !call.tool_result_metadata.retain_resource_access() { + call.tool_result_metadata = ToolResultMetadata::default(); + } + } + } + } + /// Replaces an over-budget output's calls with its own omission marker. fn bound_executed_tool_calls_with_budget(&mut self, max_metadata_bytes: usize) { let Some(metadata) = self.executed_tool_call_metadata() else { 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 5747aabab1..7d033a27a1 100644 --- a/codex-rs/protocol/src/models/executed_tool_calls_tests.rs +++ b/codex-rs/protocol/src/models/executed_tool_calls_tests.rs @@ -100,7 +100,7 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { }; for with_turn_id in [true, false] { - let mut items = (0..1_000) + let mut items = (0..4_000) .map(|index| { let mut item = output(&index.to_string()); *item @@ -133,7 +133,7 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { .flatten() .collect::>(); // Shares smaller than a local marker omit that output; surviving calls stay unchanged. - assert!(!calls.is_empty() && calls.len() < 1_000); + assert!(!calls.is_empty() && calls.len() < 4_000); assert!(calls.iter().all(|call| { **call == ExecutedToolCall::new(String::new(), serde_json::Value::Null) })); @@ -199,7 +199,7 @@ 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!(6 * argument_bytes > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); + assert!(24 * 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 +208,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); 6]; + let calls = vec![ExecutedToolCall::new("test_tool".to_string(), arguments); 24]; items[0].append_executed_tool_calls(calls.clone()); items[2].append_executed_tool_calls(vec![ExecutedToolCall::new( "test_tool".to_string(), @@ -230,7 +230,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(5)); + assert_eq!(truncation.omitted_calls, Some(23)); assert_eq!( bounded[1] .executed_tool_call_metadata() @@ -249,7 +249,7 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { .truncation() .unwrap() .omitted_calls, - Some(11), + Some(47), ); } @@ -263,7 +263,7 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { "test_tool".to_string(), serde_json::json!({ "payload": "x".repeat(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES - 256) }), ); - 4 + 16 ]); items[1].mark_tool_calls_complete(); let expected = items.clone(); @@ -299,7 +299,7 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { metadata_bytes(std::slice::from_ref(&item))?, executed_tool_call_metadata_bytes(&item), ); - let mut items = vec![item; 2_000]; + let mut items = vec![item; 8_000]; assert!(metadata_bytes(&items)? > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); bound_executed_tool_calls_for_prompt_prioritizing_recent(&mut items); assert!(metadata_bytes(&items)? <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); @@ -307,7 +307,7 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { .iter() .filter(|item| executed_tool_call_metadata_bytes(item) > 0) .count(); - assert!(retained > 0 && retained < 2_000); + assert!(retained > 0 && retained < 8_000); for item in &mut items { assert!( item.executed_tool_call_metadata() @@ -315,7 +315,7 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { ); item.clear_executed_tool_calls(); } - assert_eq!(items, vec![without_marker; 2_000]); + assert_eq!(items, vec![without_marker; 8_000]); } Ok(()) @@ -509,6 +509,10 @@ fn tool_result_metadata_is_host_only_bounded_and_redacted() -> Result<()> { "arbitrary-key": { "secret": "not-for-debug", "ids": [1, "é", null] }, "other": false, }), + serde_json::json!({ + "openai/resource_access": { "resources": ["not-for-debug"] }, + "other": "preserved when the full metadata fits", + }), serde_json::json!({}), serde_json::json!(null), ] { @@ -533,7 +537,7 @@ fn tool_result_metadata_is_host_only_bounded_and_redacted() -> Result<()> { } let exact_limit = serde_json::json!({ - "value": "x".repeat(MAX_EXECUTED_TOOL_CALL_METADATA_BYTES - r#"{"value":""}"#.len()), + "value": "x".repeat(MAX_TOOL_RESULT_METADATA_BYTES - r#"{"value":""}"#.len()), }); let capture = ToolResultMetadata::new(&exact_limit); assert!(capture.is_some()); @@ -543,8 +547,7 @@ fn tool_result_metadata_is_host_only_bounded_and_redacted() -> Result<()> { exact_limit ); // JSON escaping counts toward the bound, not just the in-memory string length. - let oversized = - serde_json::json!({ "value": "\0".repeat(MAX_EXECUTED_TOOL_CALL_METADATA_BYTES / 2) }); + 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); @@ -561,6 +564,39 @@ fn tool_result_metadata_is_host_only_bounded_and_redacted() -> Result<()> { 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(()) } @@ -571,7 +607,17 @@ fn raw_result_metadata_is_shed_before_sources_calls_or_completion() -> Result<() r#type: "test_resource".to_string(), id: "R1".to_string(), }])); - let mut items = ["first", "second", "without-metadata"].map(|call_id| { + let mut items = [ + "first", + "second", + "third", + "fourth", + "fifth", + "sixth", + "seventh", + "without-metadata", + ] + .map(|call_id| { let mut item = ResponseItem::from(ResponseInputItem::FunctionCallOutput { call_id: call_id.to_string(), output: FunctionCallOutputPayload::from_text("unchanged output".to_string()), @@ -581,20 +627,16 @@ fn raw_result_metadata_is_shed_before_sources_calls_or_completion() -> Result<() item.mark_tool_calls_complete(); item }); - let mut expected = items.clone(); let metadata = ToolResultMetadata::new(&serde_json::json!({ "large": "x".repeat(20_000) })); - for item in &mut items[..2] { + for item in &mut items[..7] { first_executed_tool_call(item) .unwrap() .set_tool_result_metadata(metadata.clone()); } - for item in &mut expected[..2] { - first_executed_tool_call(item) - .unwrap() - .set_tool_result_metadata(ToolResultMetadata::new(&serde_json::json!( - "omitted_due_to_size_limit" - ))); - } + 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() @@ -602,20 +644,193 @@ fn raw_result_metadata_is_shed_before_sources_calls_or_completion() -> Result<() .sum::() > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES ); - bound_executed_tool_calls_for_prompt(&mut items); - assert_eq!(items, expected); - assert!( - items - .iter() - .map(executed_tool_call_metadata_bytes) - .sum::() - <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES - ); - bound_executed_tool_calls_for_prompt(&mut items); - assert_eq!(items, expected); + for bound in [ + bound_executed_tool_calls_for_prompt, + bound_executed_tool_calls_for_prompt_prioritizing_recent, + ] { + let mut bounded = items.clone(); + bound(&mut bounded); + assert_eq!(bounded, expected); + assert!( + bounded + .iter() + .map(executed_tool_call_metadata_bytes) + .sum::() + <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES + ); + bound(&mut bounded); + assert_eq!(bounded, expected); + } Ok(()) } +#[test] +fn result_metadata_shedding_keeps_smaller_results_within_one_output() { + let resource_only = serde_json::json!({ + "openai/resource_access": { "resource_coverage": "complete", "resources": [] }, + }); + for (metadata, reduced_index, replacement) in [ + ( + [ + serde_json::json!({}), + serde_json::json!({ "value": "é".repeat(10_000) }), + serde_json::json!({ "value": "é".repeat(10_000) }), + ], + 1, + ToolResultMetadata::omitted_due_to_size_limit(), + ), + ( + [ + serde_json::json!({}), + serde_json::json!({ "value": "x".repeat(8_000) }), + serde_json::json!({ "value": "\0".repeat(4_500) }), + ], + 2, + ToolResultMetadata::omitted_due_to_size_limit(), + ), + ( + [ + serde_json::json!({}), + serde_json::json!({ "value": "x".repeat(8_000) }), + serde_json::json!({ + "value": "\0".repeat(4_500), + "openai/resource_access": resource_only["openai/resource_access"], + }), + ], + 2, + ToolResultMetadata::new(&resource_only), + ), + ] { + let mut metadata = metadata.to_vec(); + metadata.extend(vec![ + serde_json::json!({ "padding": "x".repeat(17_000) }); + 6 + ]); + let mut item = output("exec"); + item.append_executed_tool_calls( + metadata + .iter() + .enumerate() + .map(|(index, metadata)| { + let mut call = ExecutedToolCall::new( + format!("tool_{index}"), + serde_json::json!({ "argument": index }), + ); + call.set_tool_result_sources(ToolResultSources::new(vec![ToolResultSource { + r#type: "test_resource".to_string(), + id: format!("R{index}"), + }])); + call.set_tool_result_metadata(ToolResultMetadata::new(metadata)); + call + }) + .collect(), + ); + 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 mut expected = item.clone(); + expected + .ensure_tool_call_metadata() + .unwrap() + .executed_tool_calls + .as_mut() + .unwrap()[reduced_index] + .set_tool_result_metadata(replacement); + 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); + } + } +} + +#[test] +fn result_metadata_shedding_removes_only_markers_needed_to_fit() { + let mut item = output("exec"); + let mut call = ExecutedToolCall::new("test_tool".to_string(), serde_json::json!({})); + call.set_tool_result_sources(ToolResultSources::new(vec![ToolResultSource { + r#type: "test_resource".to_string(), + id: "R1".to_string(), + }])); + let mut calls = vec![call; 20]; + let omitted = ToolResultMetadata::omitted_due_to_size_limit(); + 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!({ + "value": "x".repeat(marker_bytes - r#"{"value":""}"#.len()), + })); + assert_eq!( + serde_json::to_vec(&equal_sized).unwrap().len(), + marker_bytes + ); + calls[0].set_tool_result_metadata(equal_sized); + for index in [1, 3] { + calls[index].set_tool_result_metadata(omitted.clone()); + } + calls[2].set_tool_result_metadata(ToolResultMetadata::new(&serde_json::json!({}))); + item.append_executed_tool_calls(calls); + item.set_tool_call_cell_id("exec"); + item.mark_tool_calls_complete(); + let remaining_bytes = + MAX_EXECUTED_TOOL_CALL_METADATA_BYTES + 1 - executed_tool_call_metadata_bytes(&item); + let calls = item + .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.arguments = ExecutedToolCallArguments::Raw(serde_json::json!("x".repeat( + remaining_bytes / call_count + usize::from(index < remaining_bytes % call_count) + ))); + } + assert_eq!( + executed_tool_call_metadata_bytes(&item), + MAX_EXECUTED_TOOL_CALL_METADATA_BYTES + 1 + ); + let mut expected = item.clone(); + expected + .ensure_tool_call_metadata() + .unwrap() + .executed_tool_calls + .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. + item.ensure_tool_call_metadata() + .unwrap() + .executed_tool_calls + .as_mut() + .unwrap()[3] + .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); + } +} + #[test] fn result_metadata_markers_do_not_displace_existing_evidence() { let mut item = ResponseItem::from(ResponseInputItem::FunctionCallOutput { @@ -627,7 +842,7 @@ fn result_metadata_markers_do_not_displace_existing_evidence() { r#type: "test_resource".to_string(), id: "R1".to_string(), }])); - item.append_executed_tool_calls(vec![call; 4]); + item.append_executed_tool_calls(vec![call; 20]); item.set_tool_call_cell_id("exec"); item.mark_tool_calls_complete(); let remaining_bytes = @@ -638,22 +853,29 @@ fn result_metadata_markers_do_not_displace_existing_evidence() { .executed_tool_calls .as_mut() .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 / 4 + usize::from(index < remaining_bytes % 4)) - )); + call.arguments = ExecutedToolCallArguments::Raw(serde_json::json!("x".repeat( + remaining_bytes / call_count + usize::from(index < remaining_bytes % call_count) + ))); } assert_eq!( executed_tool_call_metadata_bytes(&item), MAX_EXECUTED_TOOL_CALL_METADATA_BYTES ); let expected = item.clone(); - first_executed_tool_call(&mut item) - .unwrap() - .set_tool_result_metadata(ToolResultMetadata::new( - &serde_json::json!({ "too-large": "x".repeat(32 * 1024) }), - )); - bound_executed_tool_calls_for_prompt(std::slice::from_mut(&mut item)); - assert_eq!(item, expected); + for metadata in [ + serde_json::json!({ "too-large": "x".repeat(32 * 1024) }), + serde_json::json!({ + "openai/resource_access": { "resource_coverage": "complete", "resources": [] }, + }), + ] { + let mut item = expected.clone(); + first_executed_tool_call(&mut item) + .unwrap() + .set_tool_result_metadata(ToolResultMetadata::new(&metadata)); + bound_executed_tool_calls_for_prompt(std::slice::from_mut(&mut item)); + assert_eq!(item, expected); + } }