From a1dd74b53548ac48af59841ea9b8eac06b78d546 Mon Sep 17 00:00:00 2001 From: ningyi-oai Date: Sat, 1 Aug 2026 17:29:46 +0000 Subject: [PATCH] Retain attempted tool metadata across prompts (#36507) ## What changed - Reattach recorded `executed_tool_calls` metadata when an output is included in a subsequent prompt. - Bound retained metadata to 32 KiB, prioritizing recent calls and reporting omitted calls in truncation metadata. - Drop retained entries after their corresponding outputs leave the prompt history. ## Testing - Cover metadata replay, cleanup after compaction, bounded retained history, and propagation through later tool requests. GitOrigin-RevId: 1c23a26123be3b7ad51c61f4ad522139b71bb773 --- .../core/src/tools/executed_tool_calls.rs | 52 ++++++++- .../src/tools/executed_tool_calls_tests.rs | 107 +++++++++++++++++- codex-rs/core/tests/suite/tools.rs | 44 +++++-- codex-rs/protocol/src/models.rs | 2 + .../src/models/executed_tool_calls.rs | 44 +++++-- 5 files changed, 223 insertions(+), 26 deletions(-) diff --git a/codex-rs/core/src/tools/executed_tool_calls.rs b/codex-rs/core/src/tools/executed_tool_calls.rs index 4465626db0..5819cee083 100644 --- a/codex-rs/core/src/tools/executed_tool_calls.rs +++ b/codex-rs/core/src/tools/executed_tool_calls.rs @@ -4,6 +4,8 @@ use std::collections::HashSet; use codex_code_mode::CellId; use codex_protocol::models::ExecutedToolCall; use codex_protocol::models::ResponseItem; +use codex_protocol::models::bound_executed_tool_calls_for_prompt_prioritizing_recent; +use codex_protocol::models::executed_tool_call_metadata_bytes; use codex_protocol::openai_models::ToolMode; use serde_json::Value as JsonValue; @@ -27,6 +29,7 @@ struct ExecutedToolCallRecorderState { direct_calls: HashMap, cells: HashMap, output_cells: HashMap, + retained_calls: HashMap<(std::mem::Discriminant, String), Vec>, pending_nested_calls: usize, } @@ -183,17 +186,24 @@ impl ExecutedToolCallRecorder { .state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - if state.direct_calls.is_empty() && state.output_cells.is_empty() && retry_cache.is_empty() + if state.direct_calls.is_empty() + && state.output_cells.is_empty() + && state.retained_calls.is_empty() + && retry_cache.is_empty() { return false; } let mut pending_retry_outputs = retry_cache.keys().cloned().collect::>(); + let mut pending_retained_outputs = + state.retained_calls.keys().cloned().collect::>(); let mut attached = false; + let mut retained_bytes = 0_usize; for item in items.iter_mut().rev() { if state.direct_calls.is_empty() && state.output_cells.is_empty() && pending_retry_outputs.is_empty() + && pending_retained_outputs.is_empty() { break; } @@ -211,7 +221,13 @@ impl ExecutedToolCallRecorder { if !pending_retry_outputs.remove(&key) { continue; } + pending_retained_outputs.remove(&key); cached.clone() + } else if let Some(retained) = state.retained_calls.get(&key) { + if !pending_retained_outputs.remove(&key) { + continue; + } + retained.clone() } else { let mut calls = state .direct_calls @@ -232,12 +248,44 @@ impl ExecutedToolCallRecorder { if calls.is_empty() { continue; } - retry_cache.insert(key, calls.clone()); + retry_cache.insert(key.clone(), calls.clone()); + state.retained_calls.insert(key, calls.clone()); calls }; item.append_executed_tool_calls(calls); + retained_bytes = retained_bytes.saturating_add(executed_tool_call_metadata_bytes(item)); attached = true; } + if !pending_retained_outputs.is_empty() { + state + .retained_calls + .retain(|key, _| !pending_retained_outputs.contains(key)); + } + if retained_bytes > MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT { + bound_executed_tool_calls_for_prompt_prioritizing_recent(items); + state.retained_calls.clear(); + for item in items { + let call_id = match &*item { + ResponseItem::FunctionCallOutput { call_id, .. } + | ResponseItem::CustomToolCallOutput { call_id, .. } + | ResponseItem::ToolSearchOutput { + call_id: Some(call_id), + .. + } => call_id, + _ => continue, + }; + if let Some(calls) = item + .executed_tool_call_metadata() + .and_then(|metadata| metadata.executed_tool_calls.as_ref()) + .filter(|calls| !calls.is_empty()) + { + state.retained_calls.insert( + (std::mem::discriminant(&*item), call_id.clone()), + calls.clone(), + ); + } + } + } attached } diff --git a/codex-rs/core/src/tools/executed_tool_calls_tests.rs b/codex-rs/core/src/tools/executed_tool_calls_tests.rs index d22508178e..36a6adcd8c 100644 --- a/codex-rs/core/src/tools/executed_tool_calls_tests.rs +++ b/codex-rs/core/src/tools/executed_tool_calls_tests.rs @@ -102,12 +102,113 @@ fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() { }, }), ); + let expected_calls = calls.clone(); + + { + let state = recorder + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + assert_eq!(state.pending_nested_calls, 0); + assert!(!state.cells.contains_key(&cell_id)); + assert_eq!(retry_cache.len(), 1); + } + + let mut replayed_items = [ResponseItem::FunctionCallOutput { + id: None, + call_id: "bounded-output".to_string(), + output: FunctionCallOutputPayload::from_text(String::new()), + internal_chat_message_metadata_passthrough: None, + }]; + let mut replay_retry_cache = HashMap::new(); + assert!(recorder.attach_pending_to_prompt(&mut replayed_items, &mut replay_retry_cache)); + assert_eq!( + replayed_items[0] + .executed_tool_call_metadata() + .and_then(|metadata| metadata.executed_tool_calls.as_ref()), + Some(&expected_calls), + ); + + let mut compacted_retry_cache = HashMap::new(); + assert!(!recorder.attach_pending_to_prompt(&mut [], &mut compacted_retry_cache)); + let state = recorder + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + assert!(state.retained_calls.is_empty()); +} + +#[test] +fn executed_tool_call_recorder_bounds_retained_history_and_reports_omissions() { + let recorder = ExecutedToolCallRecorder::default(); + let mut history = Vec::new(); + let arguments = serde_json::to_string(&json!({ "payload": "x".repeat(1024) })) + .expect("tool arguments must serialize"); + let mut prompt = Vec::new(); + + for index in 0..512 { + let call_id = format!("retained-{index}"); + recorder.record_tool_call( + &ToolCall { + tool_name: codex_tools::ToolName::plain(format!("retained_tool_{index}")), + call_id: call_id.clone(), + payload: ToolPayload::Function { + arguments: arguments.clone(), + }, + encrypted_function_args: None, + }, + &ToolCallSource::Direct, + ToolMode::Direct, + ); + history.push(ResponseItem::FunctionCallOutput { + id: None, + call_id, + output: FunctionCallOutputPayload::from_text(String::new()), + internal_chat_message_metadata_passthrough: None, + }); + prompt = history.clone(); + assert!(recorder.attach_pending_to_prompt(&mut prompt, &mut HashMap::new())); + codex_protocol::models::bound_executed_tool_calls_for_prompt(&mut prompt); + let latest_call = prompt + .last() + .and_then(ResponseItem::executed_tool_call_metadata) + .and_then(|metadata| metadata.executed_tool_calls.as_ref()) + .and_then(|calls| calls.first()) + .map(serde_json::to_value) + .transpose() + .expect("latest tool call must serialize") + .expect("latest tool call must remain in retained metadata"); + assert_eq!(latest_call["name"], format!("retained_tool_{index}")); + assert_eq!( + latest_call["arguments"], + json!({ "payload": "x".repeat(1024) }), + ); + } let state = recorder .state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - assert_eq!(state.pending_nested_calls, 0); - assert!(!state.cells.contains_key(&cell_id)); - assert_eq!(retry_cache.len(), 1); + let retained_bytes = state + .retained_calls + .values() + .map(serialized_json_bytes) + .sum::(); + assert!(retained_bytes <= MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT); + + let metadata = prompt + .iter() + .filter_map(ResponseItem::executed_tool_call_metadata) + .filter_map(|metadata| metadata.executed_tool_calls.as_ref()) + .flatten() + .map(|call| serde_json::to_value(call).expect("retained call must serialize")) + .collect::>(); + let omitted_calls = metadata + .iter() + .filter_map(|call| { + call["arguments"]["_codex_executed_tool_call_truncated"]["omitted_calls"].as_u64() + }) + .sum::(); + assert!(omitted_calls > 0); + assert_eq!(metadata.len() as u64 + omitted_calls, 512); } diff --git a/codex-rs/core/tests/suite/tools.rs b/codex-rs/core/tests/suite/tools.rs index f47fb8fc07..4911aeea3e 100644 --- a/codex-rs/core/tests/suite/tools.rs +++ b/codex-rs/core/tests/suite/tools.rs @@ -294,21 +294,29 @@ async fn namespaced_custom_tool_call_preserves_namespace_through_dispatch_and_re PermissionProfile::Disabled, ) .await?; + let escaped_request = escaped_mock.single_request(); assert_eq!( - escaped_mock - .single_request() - .custom_tool_call_output(escaped_call_id)["internal_chat_message_metadata_passthrough"] + escaped_request.custom_tool_call_output(call_id)["internal_chat_message_metadata_passthrough"] ["executed_tool_calls"], json!([{ "name": format!("{namespace}__{tool_name}"), - "arguments": { - "_codex_executed_tool_call_truncated": { - "original_bytes": serde_json::to_vec(&escaped_input)?.len(), - "max_bytes": 8 * 1024, - }, - }, + "arguments": input, }]), ); + let expected_escaped_calls = json!([{ + "name": format!("{namespace}__{tool_name}"), + "arguments": { + "_codex_executed_tool_call_truncated": { + "original_bytes": serde_json::to_vec(&escaped_input)?.len(), + "max_bytes": 8 * 1024, + }, + }, + }]); + assert_eq!( + escaped_request.custom_tool_call_output(escaped_call_id)["internal_chat_message_metadata_passthrough"] + ["executed_tool_calls"], + expected_escaped_calls, + ); let direct_exec_call_id = "custom-direct-exec"; mount_sse_once( @@ -340,9 +348,21 @@ async fn namespaced_custom_tool_call_preserves_namespace_through_dispatch_and_re ) .await?; - let direct_exec_output = direct_exec_mock - .single_request() - .custom_tool_call_output(direct_exec_call_id); + let direct_exec_request = direct_exec_mock.single_request(); + assert_eq!( + direct_exec_request.custom_tool_call_output(call_id)["internal_chat_message_metadata_passthrough"] + ["executed_tool_calls"], + json!([{ + "name": format!("{namespace}__{tool_name}"), + "arguments": input, + }]), + ); + assert_eq!( + direct_exec_request.custom_tool_call_output(escaped_call_id)["internal_chat_message_metadata_passthrough"] + ["executed_tool_calls"], + expected_escaped_calls, + ); + let direct_exec_output = direct_exec_request.custom_tool_call_output(direct_exec_call_id); assert_eq!( direct_exec_output["output"], json!("unsupported custom tool call: exec"), diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index 29e3e859e1..2c2a6d6b7a 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -34,6 +34,8 @@ pub use executed_tool_calls::ExecutedToolCall; pub use executed_tool_calls::ExecutedToolCallArguments; pub use executed_tool_calls::ExecutedToolCallTruncation; 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; /// Controls the per-command sandbox override requested by a shell-like tool call. #[derive( diff --git a/codex-rs/protocol/src/models/executed_tool_calls.rs b/codex-rs/protocol/src/models/executed_tool_calls.rs index 65f1e438c0..668b06e11d 100644 --- a/codex-rs/protocol/src/models/executed_tool_calls.rs +++ b/codex-rs/protocol/src/models/executed_tool_calls.rs @@ -24,7 +24,8 @@ fn executed_tool_call_metadata_field_bytes( } } -fn executed_tool_call_metadata_bytes(item: &ResponseItem) -> usize { +/// Returns the exact serialized wire size of an item's attempted-tool metadata. +pub fn executed_tool_call_metadata_bytes(item: &ResponseItem) -> usize { let Some(metadata) = item.executed_tool_call_metadata() else { return 0; }; @@ -46,6 +47,20 @@ fn executed_tool_call_metadata_bytes(item: &ResponseItem) -> usize { /// Bounds attempted-tool metadata fairly across the complete serialized request. pub fn bound_executed_tool_calls_for_prompt(items: &mut [ResponseItem]) { + bound_executed_tool_calls_for_prompt_with_priority(items, /*prioritize_recent*/ false); +} + +/// Bounds retained history without letting older calls displace the newest calls. +pub fn bound_executed_tool_calls_for_prompt_prioritizing_recent(items: &mut [ResponseItem]) { + items.reverse(); + bound_executed_tool_calls_for_prompt_with_priority(items, /*prioritize_recent*/ true); + items.reverse(); +} + +fn bound_executed_tool_calls_for_prompt_with_priority( + items: &mut [ResponseItem], + prioritize_recent: bool, +) { let mut remaining_items = 0_usize; let mut original_calls = 0_usize; let mut original_metadata_bytes = 0_usize; @@ -118,8 +133,13 @@ pub fn bound_executed_tool_calls_for_prompt(items: &mut [ResponseItem]) { } let metadata_field_bytes = executed_tool_call_metadata_field_bytes(metadata); + let item_budget = if prioritize_recent { + remaining_bytes + } else { + remaining_bytes / remaining_items + }; item.bound_executed_tool_calls_with_budget( - (remaining_bytes / remaining_items).saturating_sub(metadata_field_bytes), + item_budget.saturating_sub(metadata_field_bytes), ); remaining_bytes = remaining_bytes.saturating_sub(executed_tool_call_metadata_bytes(item)); remaining_items -= 1; @@ -167,13 +187,12 @@ pub fn bound_executed_tool_calls_for_prompt(items: &mut [ResponseItem]) { return; } - if let Some(call) = items - .iter_mut() - .filter_map(ResponseItem::internal_chat_message_metadata_passthrough_mut) - .filter_map(Option::as_mut) - .filter_map(|metadata| metadata.executed_tool_calls.as_mut()) - .find_map(|calls| calls.first_mut()) - { + let omission_call = if prioritize_recent { + items.iter_mut().rev().find_map(first_executed_tool_call) + } else { + items.iter_mut().find_map(first_executed_tool_call) + }; + if let Some(call) = omission_call { let original_bytes = call .truncation() .map(|truncation| truncation.original_bytes) @@ -200,6 +219,13 @@ pub fn bound_executed_tool_calls_for_prompt(items: &mut [ResponseItem]) { } } +fn first_executed_tool_call(item: &mut ResponseItem) -> Option<&mut ExecutedToolCall> { + item.internal_chat_message_metadata_passthrough_mut() + .and_then(Option::as_mut) + .and_then(|metadata| metadata.executed_tool_calls.as_mut()) + .and_then(|calls| calls.first_mut()) +} + /// Raw model arguments or trusted truncation metadata for an attempted tool call. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonSchema, TS)] #[serde(untagged)]