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