diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index 6a1a7c40ec..5b70b15147 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -250,6 +250,7 @@ fn build_compacted_history_preserves_user_message_passthrough_metadata() { internal_chat_message_metadata_passthrough: Some( InternalChatMessageMetadataPassthrough { turn_id: Some("turn-1".to_string()), + ..Default::default() }, ), }], diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 37d0cda789..97d65fc396 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -1201,6 +1201,7 @@ fn record_items_truncates_function_call_output_content() { }, internal_chat_message_metadata_passthrough: Some(InternalChatMessageMetadataPassthrough { turn_id: Some("turn-1".to_string()), + ..Default::default() }), }; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 88b9345f70..031ffcc7e4 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -2038,6 +2038,7 @@ async fn reconstruct_history_uses_replacement_history_verbatim() { phase: None, internal_chat_message_metadata_passthrough: Some(InternalChatMessageMetadataPassthrough { turn_id: Some("compact-turn".to_string()), + ..Default::default() }), }; let replacement_history = vec![ diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index 51e22d10fd..dc9d9438a7 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -28,6 +28,13 @@ use schemars::JsonSchema; use crate::ResponseItemId; use crate::mcp::CallToolResult; +mod executed_tool_calls; + +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; + /// Controls the per-command sandbox override requested by a shell-like tool call. #[derive( Debug, Clone, Copy, Default, Eq, Hash, PartialEq, Serialize, Deserialize, JsonSchema, TS, @@ -778,6 +785,11 @@ pub struct InternalChatMessageMetadataPassthrough { #[serde(default, skip_serializing_if = "Option::is_none")] #[ts(optional)] pub turn_id: Option, + /// Warehouse-only Responses metadata, not part of the public app-server protocol. + #[serde(default, skip_deserializing, skip_serializing_if = "Option::is_none")] + #[schemars(skip)] + #[ts(skip)] + pub executed_tool_calls: Option>, } impl InternalChatMessageMetadataPassthrough { @@ -2283,6 +2295,25 @@ mod tests { }))?; assert_eq!(unknown_metadata, item); + item.append_executed_tool_calls(vec![ExecutedToolCall::new( + "test_tool".to_string(), + serde_json::json!({"value": 1}), + )]); + let serialized = serde_json::to_value(&item)?; + assert_eq!( + serialized["internal_chat_message_metadata_passthrough"]["executed_tool_calls"], + serde_json::json!([{"name": "test_tool", "arguments": {"value": 1}}]), + ); + let deserialized = serde_json::from_value::(serialized)?; + assert_eq!(deserialized.turn_id(), Some("turn-1")); + assert!( + deserialized + .executed_tool_call_metadata() + .and_then(|metadata| metadata.executed_tool_calls.as_ref()) + .is_none(), + "provider and rollout input cannot forge local attempted-tool metadata", + ); + item.clear_executed_tool_calls(); item.set_turn_id_if_missing("turn-2"); assert_eq!(item.turn_id(), Some("turn-1")); @@ -2348,6 +2379,7 @@ mod tests { fn passthrough_metadata(turn_id: &str) -> InternalChatMessageMetadataPassthrough { InternalChatMessageMetadataPassthrough { turn_id: Some(turn_id.to_string()), + ..Default::default() } } diff --git a/codex-rs/protocol/src/models/executed_tool_calls.rs b/codex-rs/protocol/src/models/executed_tool_calls.rs new file mode 100644 index 0000000000..65f1e438c0 --- /dev/null +++ b/codex-rs/protocol/src/models/executed_tool_calls.rs @@ -0,0 +1,405 @@ +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; +use ts_rs::TS; + +use super::InternalChatMessageMetadataPassthrough; +use super::ResponseItem; + +const MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES: usize = 8 * 1024; +/// Maximum serialized warehouse-only attempted-tool metadata in one request. +const MAX_EXECUTED_TOOL_CALL_METADATA_BYTES: usize = 32 * 1024; +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(); + +fn executed_tool_call_metadata_field_bytes( + metadata: &InternalChatMessageMetadataPassthrough, +) -> usize { + EXECUTED_TOOL_CALL_METADATA_FIELD_BYTES + + if metadata.turn_id.is_some() { + 1 + } else { + INTERNAL_CHAT_MESSAGE_METADATA_PASSTHROUGH_FIELD_BYTES + 3 + } +} + +fn executed_tool_call_metadata_bytes(item: &ResponseItem) -> usize { + let Some(metadata) = item.executed_tool_call_metadata() else { + return 0; + }; + let Some(calls) = metadata + .executed_tool_calls + .as_ref() + .filter(|calls| !calls.is_empty()) + else { + return 0; + }; + serde_json::to_vec(calls) + .map(|bytes| { + bytes + .len() + .saturating_add(executed_tool_call_metadata_field_bytes(metadata)) + }) + .unwrap_or(usize::MAX) +} + +/// Bounds attempted-tool metadata fairly across the complete serialized request. +pub fn bound_executed_tool_calls_for_prompt(items: &mut [ResponseItem]) { + let mut remaining_items = 0_usize; + let mut original_calls = 0_usize; + let mut original_metadata_bytes = 0_usize; + + for item in items.iter_mut() { + let Some(calls) = item + .internal_chat_message_metadata_passthrough_mut() + .and_then(Option::as_mut) + .and_then(|metadata| metadata.executed_tool_calls.as_mut()) + .filter(|calls| !calls.is_empty()) + else { + continue; + }; + for call in calls { + let argument_bytes = serde_json::to_vec(&call.arguments) + .map(|bytes| bytes.len()) + .unwrap_or(usize::MAX); + if call.truncation().is_none() && argument_bytes > MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES + { + call.set_truncation( + argument_bytes, + MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES, + /*omitted_calls*/ None, + ); + } + original_calls = original_calls.saturating_add(1).saturating_add( + call.truncation() + .and_then(|truncation| truncation.omitted_calls) + .unwrap_or_default(), + ); + } + remaining_items += 1; + original_metadata_bytes = + original_metadata_bytes.saturating_add(executed_tool_call_metadata_bytes(item)); + } + + if original_metadata_bytes <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { + return; + } + + let overflow_fallback = items.iter().enumerate().find_map(|(index, item)| { + item.executed_tool_call_metadata() + .and_then(|metadata| metadata.executed_tool_calls.as_ref()) + .and_then(|calls| calls.first()) + .map(|call| (index, call.clone())) + }); + + let omitted_call_reservation = serde_json::to_vec(&serde_json::json!({ + "_codex_executed_tool_call_truncated": ExecutedToolCallTruncation { + original_bytes: usize::MAX, + max_bytes: usize::MAX, + omitted_calls: Some(usize::MAX), + original_name_bytes: Some(usize::MAX), + }, + })) + .map(|bytes| bytes.len()) + .unwrap_or(usize::MAX); + let mut remaining_bytes = + MAX_EXECUTED_TOOL_CALL_METADATA_BYTES.saturating_sub(omitted_call_reservation); + for item in items.iter_mut() { + let Some(metadata) = item.executed_tool_call_metadata() else { + continue; + }; + if metadata + .executed_tool_calls + .as_ref() + .is_none_or(Vec::is_empty) + { + continue; + } + + let metadata_field_bytes = executed_tool_call_metadata_field_bytes(metadata); + item.bound_executed_tool_calls_with_budget( + (remaining_bytes / remaining_items).saturating_sub(metadata_field_bytes), + ); + remaining_bytes = remaining_bytes.saturating_sub(executed_tool_call_metadata_bytes(item)); + remaining_items -= 1; + } + + let represented_calls = items + .iter() + .filter_map(ResponseItem::executed_tool_call_metadata) + .filter_map(|metadata| metadata.executed_tool_calls.as_ref()) + .flatten() + .map(|call| { + 1 + call + .truncation() + .and_then(|truncation| truncation.omitted_calls) + .unwrap_or_default() + }) + .sum::(); + if represented_calls == original_calls { + return; + } + + if represented_calls == 0 { + if let Some((index, mut call)) = overflow_fallback { + let original_bytes = call + .truncation() + .map(|truncation| truncation.original_bytes) + .unwrap_or_else(|| { + serde_json::to_vec(&call.arguments) + .map(|bytes| bytes.len()) + .unwrap_or(usize::MAX) + }); + let original_name_bytes = call.name.len(); + let name_boundary = call.name.floor_char_boundary( + original_name_bytes.min(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES / 2), + ); + call.name.truncate(name_boundary); + call.set_truncation_with_name( + original_bytes, + /*max_bytes*/ 0, + Some(original_calls.saturating_sub(1)), + (name_boundary < original_name_bytes).then_some(original_name_bytes), + ); + items[index].append_executed_tool_calls(vec![call]); + } + 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 original_bytes = call + .truncation() + .map(|truncation| truncation.original_bytes) + .unwrap_or_else(|| { + serde_json::to_vec(&call.arguments) + .map(|bytes| bytes.len()) + .unwrap_or(usize::MAX) + }); + let max_bytes = call + .truncation() + .map(|truncation| truncation.max_bytes) + .unwrap_or_default(); + let previous_omissions = call + .truncation() + .and_then(|truncation| truncation.omitted_calls) + .unwrap_or_default(); + call.set_truncation( + original_bytes, + max_bytes, + Some( + previous_omissions.saturating_add(original_calls.saturating_sub(represented_calls)), + ), + ); + } +} + +/// Raw model arguments or trusted truncation metadata for an attempted tool call. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonSchema, TS)] +#[serde(untagged)] +pub enum ExecutedToolCallArguments { + Raw(serde_json::Value), + #[serde(skip_deserializing)] + Truncated { + #[serde(rename = "_codex_executed_tool_call_truncated")] + truncation: ExecutedToolCallTruncation, + }, +} + +/// A model-attempted Codex tool invocation captured at the shared runtime boundary. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonSchema, TS)] +pub struct ExecutedToolCall { + pub name: String, + #[ts(type = "unknown")] + arguments: ExecutedToolCallArguments, +} + +/// Trusted truncation details generated locally for an oversized attempted tool call. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonSchema, TS)] +pub struct ExecutedToolCallTruncation { + original_bytes: usize, + max_bytes: usize, + #[serde(default, skip_serializing_if = "Option::is_none")] + omitted_calls: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + original_name_bytes: Option, +} + +impl ExecutedToolCall { + /// Creates a recorded call without treating model-provided JSON as trusted metadata. + pub fn new(name: String, arguments: serde_json::Value) -> Self { + let arguments = if arguments + .as_object() + .is_some_and(|object| object.contains_key("_codex_executed_tool_call_truncated")) + { + serde_json::json!({ "_codex_executed_tool_call_raw": arguments }) + } else { + arguments + }; + Self { + name, + arguments: ExecutedToolCallArguments::Raw(arguments), + } + } + + /// Replaces oversized arguments with internally generated truncation metadata. + pub fn truncated(name: String, original_bytes: usize, max_bytes: usize) -> Self { + let mut call = Self::new(name, serde_json::Value::Null); + call.set_truncation(original_bytes, max_bytes, /*omitted_calls*/ None); + call + } + + /// Returns the raw arguments or locally generated truncation payload. + pub fn arguments(&self) -> &ExecutedToolCallArguments { + &self.arguments + } + + fn truncation(&self) -> Option<&ExecutedToolCallTruncation> { + match &self.arguments { + ExecutedToolCallArguments::Raw(_) => None, + ExecutedToolCallArguments::Truncated { truncation } => Some(truncation), + } + } + + fn set_truncation( + &mut self, + original_bytes: usize, + max_bytes: usize, + omitted_calls: Option, + ) { + self.set_truncation_with_name( + original_bytes, + max_bytes, + omitted_calls, + /*original_name_bytes*/ None, + ); + } + + fn set_truncation_with_name( + &mut self, + original_bytes: usize, + max_bytes: usize, + omitted_calls: Option, + original_name_bytes: Option, + ) { + self.arguments = ExecutedToolCallArguments::Truncated { + truncation: ExecutedToolCallTruncation { + original_bytes, + max_bytes, + omitted_calls, + original_name_bytes, + }, + }; + } +} + +impl ResponseItem { + /// Attaches model-attempted tool invocations without replacing existing item metadata. + pub fn append_executed_tool_calls(&mut self, calls: Vec) { + if calls.is_empty() { + return; + } + let Some(metadata) = self.internal_chat_message_metadata_passthrough_mut() else { + return; + }; + metadata + .get_or_insert_with(InternalChatMessageMetadataPassthrough::default) + .executed_tool_calls + .get_or_insert_with(Vec::new) + .extend(calls); + } + + /// Returns warehouse-only attempted-tool metadata for any supported item variant. + pub fn executed_tool_call_metadata(&self) -> Option<&InternalChatMessageMetadataPassthrough> { + self.internal_chat_message_metadata_passthrough() + } + + /// Bounds one request item's attempted calls to its share of the prompt budget. + fn bound_executed_tool_calls_with_budget(&mut self, max_metadata_bytes: usize) { + let Some(metadata) = self.internal_chat_message_metadata_passthrough_mut() else { + return; + }; + let Some(passthrough) = metadata.as_mut() else { + return; + }; + let Some(calls) = passthrough.executed_tool_calls.as_mut() else { + return; + }; + + let mut serialized_bytes = 2; + let mut retained_calls = 0; + calls.retain_mut(|call| { + let separator_bytes = usize::from(retained_calls > 0); + let remaining_bytes = max_metadata_bytes + .saturating_sub(serialized_bytes) + .saturating_sub(separator_bytes); + let argument_bytes = serde_json::to_vec(&call.arguments) + .map(|bytes| bytes.len()) + .unwrap_or(usize::MAX); + let call_bytes = serde_json::to_vec(&*call) + .map(|bytes| bytes.len()) + .unwrap_or(usize::MAX); + + if call_bytes > remaining_bytes + || argument_bytes > MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES + { + let original_bytes = call + .truncation() + .map(|truncation| truncation.original_bytes) + .unwrap_or(argument_bytes); + let omitted_calls = call + .truncation() + .and_then(|truncation| truncation.omitted_calls); + call.set_truncation( + original_bytes, + remaining_bytes.min(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES), + omitted_calls, + ); + } + + let call_bytes = serde_json::to_vec(&*call) + .map(|bytes| bytes.len()) + .unwrap_or(usize::MAX); + if call_bytes > remaining_bytes { + return false; + } + + serialized_bytes = serialized_bytes + .saturating_add(separator_bytes) + .saturating_add(call_bytes); + retained_calls += 1; + true + }); + + if calls.is_empty() { + passthrough.executed_tool_calls = None; + if *passthrough == InternalChatMessageMetadataPassthrough::default() { + *metadata = None; + } + } + } + + /// Removes untrusted warehouse-only tool records without changing the turn ID. + pub fn clear_executed_tool_calls(&mut self) { + let Some(metadata) = self.internal_chat_message_metadata_passthrough_mut() else { + return; + }; + let Some(passthrough) = metadata.as_mut() else { + return; + }; + passthrough.executed_tool_calls = None; + if *passthrough == InternalChatMessageMetadataPassthrough::default() { + *metadata = None; + } + } +} + +#[cfg(test)] +#[path = "executed_tool_calls_tests.rs"] +mod tests; diff --git a/codex-rs/protocol/src/models/executed_tool_calls_tests.rs b/codex-rs/protocol/src/models/executed_tool_calls_tests.rs new file mode 100644 index 0000000000..5409b7359a --- /dev/null +++ b/codex-rs/protocol/src/models/executed_tool_calls_tests.rs @@ -0,0 +1,173 @@ +use anyhow::Result; +use pretty_assertions::assert_eq; + +use super::super::FunctionCallOutputBody; +use super::super::FunctionCallOutputPayload; +use super::*; + +fn passthrough_metadata(turn_id: &str) -> InternalChatMessageMetadataPassthrough { + InternalChatMessageMetadataPassthrough { + turn_id: Some(turn_id.to_string()), + ..Default::default() + } +} + +#[test] +fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { + let metadata_bytes = |items: &[ResponseItem]| -> Result { + items.iter().try_fold(0_usize, |bytes, item| { + let with_metadata = serde_json::to_vec(item)?.len(); + let mut without_metadata = item.clone(); + without_metadata.clear_executed_tool_calls(); + let without_metadata = serde_json::to_vec(&without_metadata)?.len(); + Ok(bytes + with_metadata.saturating_sub(without_metadata)) + }) + }; + + for with_turn_id in [true, false] { + let mut items = (0..1_000) + .map(|index| { + let mut item = ResponseItem::FunctionCallOutput { + id: None, + call_id: index.to_string(), + output: FunctionCallOutputPayload { + body: FunctionCallOutputBody::Text(String::new()), + success: None, + }, + internal_chat_message_metadata_passthrough: with_turn_id + .then(|| passthrough_metadata("turn-1")), + }; + item.append_executed_tool_calls(vec![ExecutedToolCall::new( + String::new(), + serde_json::Value::Null, + )]); + item + }) + .collect::>(); + + assert!(metadata_bytes(&items)? > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); + bound_executed_tool_calls_for_prompt(&mut items); + assert!(metadata_bytes(&items)? <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); + + let calls = items + .iter() + .filter_map(ResponseItem::executed_tool_call_metadata) + .filter_map(|metadata| metadata.executed_tool_calls.as_ref()) + .flatten() + .collect::>(); + let omitted_calls = calls + .iter() + .filter_map(|call| call.truncation()) + .filter_map(|truncation| truncation.omitted_calls) + .sum::(); + assert!(omitted_calls > 0); + assert_eq!(calls.len() + omitted_calls, 1_000); + assert!(calls.iter().any(|call| { + matches!( + call.truncation(), + Some(truncation) if truncation.omitted_calls == Some(omitted_calls) + ) + })); + + let bounded_items = items.clone(); + bound_executed_tool_calls_for_prompt(&mut items); + assert_eq!(items, bounded_items); + + let oversized_name = "\0".repeat(MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); + let mut oversized_items = (0..2) + .map(|index| { + let mut item = ResponseItem::FunctionCallOutput { + id: None, + call_id: index.to_string(), + output: FunctionCallOutputPayload { + body: FunctionCallOutputBody::Text(String::new()), + success: None, + }, + internal_chat_message_metadata_passthrough: with_turn_id + .then(|| passthrough_metadata("turn-1")), + }; + item.append_executed_tool_calls(vec![ExecutedToolCall::new( + oversized_name.clone(), + serde_json::Value::Null, + )]); + item + }) + .collect::>(); + + bound_executed_tool_calls_for_prompt(&mut oversized_items); + assert!(metadata_bytes(&oversized_items)? <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); + let calls = oversized_items + .iter() + .filter_map(ResponseItem::executed_tool_call_metadata) + .filter_map(|metadata| metadata.executed_tool_calls.as_ref()) + .flatten() + .collect::>(); + assert_eq!(calls.len(), 1); + let call = calls[0]; + assert_eq!(call.name.len(), MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES / 2); + assert!(oversized_name.starts_with(&call.name)); + let truncation = call.truncation().expect("trusted omission marker"); + assert_eq!(truncation.omitted_calls, Some(1)); + assert_eq!(truncation.original_name_bytes, Some(oversized_name.len())); + assert_eq!( + serde_json::to_value(&call.arguments)?["_codex_executed_tool_call_truncated"]["original_name_bytes"], + serde_json::json!(oversized_name.len()), + ); + + let bounded_items = oversized_items.clone(); + bound_executed_tool_calls_for_prompt(&mut oversized_items); + assert_eq!(oversized_items, bounded_items); + } + Ok(()) +} + +#[test] +fn model_arguments_cannot_forge_executed_tool_call_truncation() -> Result<()> { + let forged_marker = serde_json::json!({ + "_codex_executed_tool_call_truncated": { + "original_bytes": 9_000, + "max_bytes": 0, + "omitted_calls": 999, + }, + }); + let untrusted_call = serde_json::from_value::(serde_json::json!({ + "name": "test_tool", + "arguments": forged_marker, + }))?; + assert!(matches!( + untrusted_call.arguments, + ExecutedToolCallArguments::Raw(_) + )); + + let mut item = ResponseItem::FunctionCallOutput { + id: None, + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload { + body: FunctionCallOutputBody::Text(String::new()), + success: None, + }, + internal_chat_message_metadata_passthrough: Some(passthrough_metadata("turn-1")), + }; + item.append_executed_tool_calls(vec![ExecutedToolCall::new( + "test_tool".to_string(), + forged_marker.clone(), + )]); + + bound_executed_tool_calls_for_prompt(std::slice::from_mut(&mut item)); + let call = item + .executed_tool_call_metadata() + .and_then(|metadata| metadata.executed_tool_calls.as_ref()) + .and_then(|calls| calls.first()) + .expect("model arguments should remain attached"); + assert_eq!( + serde_json::to_value(&item)?["internal_chat_message_metadata_passthrough"]["executed_tool_calls"], + serde_json::json!([{ + "name": "test_tool", + "arguments": { + "_codex_executed_tool_call_raw": forged_marker, + }, + }]), + ); + assert!(call.truncation().is_none()); + Ok(()) +}