From 46aa019e805352c9d7fd9a740cbf7f8b9aeb162d Mon Sep 17 00:00:00 2001 From: Tamir Duberstein Date: Tue, 25 Aug 2026 13:04:44 +0000 Subject: [PATCH] Count serialized JSON bytes without allocating strings (#40604) ## Why History estimation and executed-tool-call limits only need the size of serialized JSON, not the serialized output itself. ## What changed - Add a shared `serialized_json_bytes` helper that streams JSON into a byte counter. - Use it for response-item size estimates and executed-tool-call argument limits. - Preserve serialization errors so callers can apply their existing fallback behavior. ## Testing - Cover response-item sizing with escaped characters, newlines, and multibyte text. GitOrigin-RevId: a4417996df96df5c51cae62830adc496912f69c7 --- codex-rs/core/src/context_manager/history.rs | 5 ++-- .../core/src/context_manager/history_tests.rs | 2 +- .../core/src/tools/executed_tool_calls.rs | 29 ++++--------------- .../src/tools/executed_tool_calls_tests.rs | 3 +- codex-rs/core/src/utils/json.rs | 22 ++++++++++++++ codex-rs/core/src/utils/mod.rs | 1 + 6 files changed, 34 insertions(+), 28 deletions(-) create mode 100644 codex-rs/core/src/utils/json.rs diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 98254ebc56..94ee21ae96 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -7,6 +7,7 @@ use crate::event_mapping::has_non_contextual_dev_message_content; use crate::event_mapping::is_contextual_dev_message_content; use crate::event_mapping::is_contextual_user_message_content; use crate::session::turn_context::TurnContext; +use crate::utils::json::serialized_json_bytes; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_context_fragments::set_annotated_content; @@ -666,8 +667,8 @@ fn estimate_response_item_model_visible_bytes(item: &ResponseItem) -> i64 { .. } => i64::try_from(estimate_reasoning_length(content.len())).unwrap_or(i64::MAX), item => { - let raw = serde_json::to_string(item) - .map(|serialized| i64::try_from(serialized.len()).unwrap_or(i64::MAX)) + let raw = serialized_json_bytes(item) + .map(|len| i64::try_from(len).unwrap_or(i64::MAX)) .unwrap_or_default(); let (image_payload_bytes, image_replacement_bytes) = image_data_url_estimate_adjustment(item); diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 06fa8b4b08..a81085e2e4 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -2724,7 +2724,7 @@ fn text_only_items_unchanged() { id: None, role: "assistant".to_string(), content: vec![ContentItem::OutputText { - text: "Hello world, this is a response.".to_string(), + text: "Hello, \"world\"!\nこんにちは".to_string(), }], phase: None, internal_chat_message_metadata_passthrough: None, diff --git a/codex-rs/core/src/tools/executed_tool_calls.rs b/codex-rs/core/src/tools/executed_tool_calls.rs index 7434728e7e..f16363fe83 100644 --- a/codex-rs/core/src/tools/executed_tool_calls.rs +++ b/codex-rs/core/src/tools/executed_tool_calls.rs @@ -12,6 +12,7 @@ use serde_json::Value as JsonValue; use crate::tools::context::ToolCallSource; use crate::tools::context::ToolPayload; use crate::tools::router::ToolCall; +use crate::utils::json::serialized_json_bytes; const MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES: usize = 8 * 1024; const MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT: usize = 32 * 1024; @@ -39,28 +40,6 @@ struct RecordedCell { pending_full_argument_bytes: usize, } -#[derive(Default)] -struct JsonByteCounter(usize); - -impl std::io::Write for JsonByteCounter { - fn write(&mut self, bytes: &[u8]) -> std::io::Result { - self.0 = self.0.saturating_add(bytes.len()); - Ok(bytes.len()) - } - - fn flush(&mut self) -> std::io::Result<()> { - Ok(()) - } -} - -fn serialized_json_bytes(value: &T) -> usize { - let mut counter = JsonByteCounter::default(); - if serde_json::to_writer(&mut counter, value).is_err() { - return usize::MAX; - } - counter.0 -} - impl ExecutedToolCallRecorder { pub(crate) fn record_tool_call( &self, @@ -87,8 +66,10 @@ impl ExecutedToolCallRecorder { let original_bytes = match &call.payload { ToolPayload::Function { arguments } => arguments.len(), - ToolPayload::Custom { input } => serialized_json_bytes(input), - ToolPayload::ToolSearch { arguments } => serialized_json_bytes(arguments), + ToolPayload::Custom { input } => serialized_json_bytes(input).unwrap_or(usize::MAX), + ToolPayload::ToolSearch { arguments } => { + serialized_json_bytes(arguments).unwrap_or(usize::MAX) + } }; let name = codex_tools::code_mode_name_for_tool_name(&call.tool_name); let recorded_call = if original_bytes > MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES { diff --git a/codex-rs/core/src/tools/executed_tool_calls_tests.rs b/codex-rs/core/src/tools/executed_tool_calls_tests.rs index c266d7c390..6a4a8d6c8a 100644 --- a/codex-rs/core/src/tools/executed_tool_calls_tests.rs +++ b/codex-rs/core/src/tools/executed_tool_calls_tests.rs @@ -199,7 +199,8 @@ fn executed_tool_call_recorder_bounds_retained_history_and_reports_omissions() { .retained_calls .values() .map(serialized_json_bytes) - .sum::(); + .sum::>() + .expect("retained calls must serialize"); assert!(retained_bytes <= MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT); let metadata = prompt diff --git a/codex-rs/core/src/utils/json.rs b/codex-rs/core/src/utils/json.rs new file mode 100644 index 0000000000..54b2303f1f --- /dev/null +++ b/codex-rs/core/src/utils/json.rs @@ -0,0 +1,22 @@ +#[derive(Default)] +struct JsonByteCounter(usize); + +impl std::io::Write for JsonByteCounter { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0 = self.0.saturating_add(bytes.len()); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +/// Counts serialized JSON bytes without retaining the serialized output. +pub(crate) fn serialized_json_bytes( + value: &T, +) -> serde_json::Result { + let mut counter = JsonByteCounter::default(); + serde_json::to_writer(&mut counter, value)?; + Ok(counter.0) +} diff --git a/codex-rs/core/src/utils/mod.rs b/codex-rs/core/src/utils/mod.rs index adbdcef12b..b4bc359f3a 100644 --- a/codex-rs/core/src/utils/mod.rs +++ b/codex-rs/core/src/utils/mod.rs @@ -1 +1,2 @@ +pub(crate) mod json; pub mod path_utils;