openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-08-23 13:09:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Retain attempted tool metadata across prompts (#36507)

Browse Source 
## 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
This commit is contained in:
ningyi-oai
2026-08-01 17:29:46 +00:00
committed by copyberry
parent feee0b07c7
commit a1dd74b535
5 changed files with 223 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
codex-rs/core/src/tools/executed_tool_calls.rs
View File

@@ -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<String, ExecutedToolCall>,
cells: HashMap<CellId, RecordedCell>,
output_cells: HashMap<String, CellId>,
retained_calls: HashMap<(std::mem::Discriminant<ResponseItem>, String), Vec<ExecutedToolCall>>,
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::<HashSet<_>>();
let mut pending_retained_outputs =
state.retained_calls.keys().cloned().collect::<HashSet<_>>();
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
}

107
codex-rs/core/src/tools/executed_tool_calls_tests.rs
View File

@@ -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::<usize>();
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::<Vec<_>>();
let omitted_calls = metadata
.iter()
.filter_map(|call| {
call["arguments"]["_codex_executed_tool_call_truncated"]["omitted_calls"].as_u64()
})
.sum::<u64>();
assert!(omitted_calls > 0);
assert_eq!(metadata.len() as u64 + omitted_calls, 512);
}

44
codex-rs/core/tests/suite/tools.rs
View File

@@ -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"),

2
codex-rs/protocol/src/models.rs
View File

@@ -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(

44
codex-rs/protocol/src/models/executed_tool_calls.rs
View File

@@ -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)]