From 062fa7a2bb96b7b6f5fd0bc1edb7c48d75024114 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Tue, 24 Mar 2026 15:45:40 -0700 Subject: [PATCH] Move string truncation helpers into codex-utils-string (#15572) - move the shared byte-based middle truncation logic from `core` into `codex-utils-string` - keep token-specific truncation in `codex-core` so rollout can reuse the shared helper in the next stacked PR --------- Co-authored-by: Codex --- codex-rs/Cargo.lock | 11 + codex-rs/Cargo.toml | 2 + codex-rs/core/Cargo.toml | 1 + codex-rs/core/src/codex.rs | 4 +- .../src/codex/rollout_reconstruction_tests.rs | 16 +- codex-rs/core/src/codex_tests.rs | 2 +- codex-rs/core/src/compact.rs | 6 +- codex-rs/core/src/context_manager/history.rs | 12 +- .../core/src/context_manager/history_tests.rs | 6 +- codex-rs/core/src/error.rs | 4 +- codex-rs/core/src/guardian/prompt.rs | 6 +- codex-rs/core/src/lib.rs | 1 - codex-rs/core/src/memories/prompts.rs | 4 +- .../core/src/models_manager/model_info.rs | 2 +- codex-rs/core/src/realtime_context.rs | 4 +- codex-rs/core/src/rollout/list.rs | 2 +- codex-rs/core/src/rollout/recorder.rs | 8 +- codex-rs/core/src/rollout/recorder_tests.rs | 2 +- codex-rs/core/src/rollout/tests.rs | 2 +- codex-rs/core/src/state/session.rs | 2 +- codex-rs/core/src/tools/code_mode/mod.rs | 6 +- codex-rs/core/src/tools/context.rs | 4 +- codex-rs/core/src/tools/js_repl/mod.rs | 4 +- codex-rs/core/src/tools/mod.rs | 6 +- codex-rs/core/src/truncate.rs | 363 ------------------ codex-rs/core/src/unified_exec/process.rs | 4 +- .../core/src/unified_exec/process_manager.rs | 2 +- codex-rs/protocol/Cargo.toml | 1 + codex-rs/protocol/src/protocol.rs | 46 +++ codex-rs/utils/output-truncation/BUILD.bazel | 6 + codex-rs/utils/output-truncation/Cargo.toml | 15 + codex-rs/utils/output-truncation/src/lib.rs | 142 +++++++ .../output-truncation}/src/truncate_tests.rs | 106 ++--- codex-rs/utils/string/src/lib.rs | 9 + codex-rs/utils/string/src/truncate.rs | 156 ++++++++ codex-rs/utils/string/src/truncate/tests.rs | 71 ++++ 36 files changed, 551 insertions(+), 487 deletions(-) delete mode 100644 codex-rs/core/src/truncate.rs create mode 100644 codex-rs/utils/output-truncation/BUILD.bazel create mode 100644 codex-rs/utils/output-truncation/Cargo.toml create mode 100644 codex-rs/utils/output-truncation/src/lib.rs rename codex-rs/{core => utils/output-truncation}/src/truncate_tests.rs (77%) create mode 100644 codex-rs/utils/string/src/truncate.rs create mode 100644 codex-rs/utils/string/src/truncate/tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 935fc413a..697b8f5d3 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1902,6 +1902,7 @@ dependencies = [ "codex-utils-cargo-bin", "codex-utils-home-dir", "codex-utils-image", + "codex-utils-output-truncation", "codex-utils-pty", "codex-utils-readiness", "codex-utils-stream-parser", @@ -2399,6 +2400,7 @@ dependencies = [ "codex-git-utils", "codex-utils-absolute-path", "codex-utils-image", + "codex-utils-string", "icu_decimal", "icu_locale_core", "icu_provider", @@ -2891,6 +2893,15 @@ dependencies = [ "codex-ollama", ] +[[package]] +name = "codex-utils-output-truncation" +version = "0.0.0" +dependencies = [ + "codex-protocol", + "codex-utils-string", + "pretty_assertions", +] + [[package]] name = "codex-utils-pty" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index e071ec807..ebd88c945 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -65,6 +65,7 @@ members = [ "utils/sleep-inhibitor", "utils/approval-presets", "utils/oss", + "utils/output-truncation", "utils/fuzzy-match", "utils/stream-parser", "codex-client", @@ -154,6 +155,7 @@ codex-utils-home-dir = { path = "utils/home-dir" } codex-utils-image = { path = "utils/image" } codex-utils-json-to-toml = { path = "utils/json-to-toml" } codex-utils-oss = { path = "utils/oss" } +codex-utils-output-truncation = { path = "utils/output-truncation" } codex-utils-pty = { path = "utils/pty" } codex-utils-readiness = { path = "utils/readiness" } codex-utils-rustls-provider = { path = "utils/rustls-provider" } diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 82ea26005..51dd305b3 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -55,6 +55,7 @@ codex-utils-absolute-path = { workspace = true } codex-utils-cache = { workspace = true } codex-utils-image = { workspace = true } codex-utils-home-dir = { workspace = true } +codex-utils-output-truncation = { workspace = true } codex-utils-pty = { workspace = true } codex-utils-readiness = { workspace = true } codex-secrets = { workspace = true } diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index b299bec36..2f8c2877f 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -46,7 +46,6 @@ use crate::stream_events_utils::handle_output_item_done; use crate::stream_events_utils::last_assistant_message_from_item; use crate::stream_events_utils::raw_assistant_output_text_from_item; use crate::stream_events_utils::record_completed_response_item; -use crate::truncate::TruncationPolicy; use crate::turn_metadata::TurnMetadataState; use crate::util::error_or_panic; use async_channel::Receiver; @@ -117,6 +116,7 @@ use codex_protocol::request_user_input::RequestUserInputResponse; use codex_rmcp_client::ElicitationResponse; use codex_rmcp_client::OAuthCredentialsStoreMode; use codex_terminal_detection::user_agent; +use codex_utils_output_truncation::TruncationPolicy; use codex_utils_stream_parser::AssistantTextChunk; use codex_utils_stream_parser::AssistantTextStreamParser; use codex_utils_stream_parser::ProposedPlanSegment; @@ -1007,7 +1007,7 @@ impl TurnContext { user_instructions: self.user_instructions.clone(), developer_instructions: self.developer_instructions.clone(), final_output_json_schema: self.final_output_json_schema.clone(), - truncation_policy: Some(self.truncation_policy.into()), + truncation_policy: Some(self.truncation_policy), } } diff --git a/codex-rs/core/src/codex/rollout_reconstruction_tests.rs b/codex-rs/core/src/codex/rollout_reconstruction_tests.rs index 906c72c92..1ed668d56 100644 --- a/codex-rs/core/src/codex/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/codex/rollout_reconstruction_tests.rs @@ -77,7 +77,7 @@ async fn record_initial_history_resumed_bare_turn_context_does_not_hydrate_previ user_instructions: None, developer_instructions: None, final_output_json_schema: None, - truncation_policy: Some(turn_context.truncation_policy.into()), + truncation_policy: Some(turn_context.truncation_policy), }; let rollout_items = vec![RolloutItem::TurnContext(previous_context_item)]; @@ -116,7 +116,7 @@ async fn record_initial_history_resumed_hydrates_previous_turn_settings_from_lif user_instructions: None, developer_instructions: None, final_output_json_schema: None, - truncation_policy: Some(turn_context.truncation_policy.into()), + truncation_policy: Some(turn_context.truncation_policy), }; let turn_id = previous_context_item .turn_id @@ -866,7 +866,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis user_instructions: None, developer_instructions: None, final_output_json_schema: None, - truncation_policy: Some(turn_context.truncation_policy.into()), + truncation_policy: Some(turn_context.truncation_policy), }; let previous_turn_id = previous_context_item .turn_id @@ -938,7 +938,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis user_instructions: None, developer_instructions: None, final_output_json_schema: None, - truncation_policy: Some(turn_context.truncation_policy.into()), + truncation_policy: Some(turn_context.truncation_policy), })) .expect("serialize expected reference context item") ); @@ -967,7 +967,7 @@ async fn record_initial_history_resumed_aborted_turn_without_id_clears_active_tu user_instructions: None, developer_instructions: None, final_output_json_schema: None, - truncation_policy: Some(turn_context.truncation_policy.into()), + truncation_policy: Some(turn_context.truncation_policy), }; let previous_turn_id = previous_context_item .turn_id @@ -1073,7 +1073,7 @@ async fn record_initial_history_resumed_unmatched_abort_preserves_active_turn_fo user_instructions: None, developer_instructions: None, final_output_json_schema: None, - truncation_policy: Some(turn_context.truncation_policy.into()), + truncation_policy: Some(turn_context.truncation_policy), }; let rollout_items = vec![ @@ -1175,7 +1175,7 @@ async fn record_initial_history_resumed_trailing_incomplete_turn_compaction_clea user_instructions: None, developer_instructions: None, final_output_json_schema: None, - truncation_policy: Some(turn_context.truncation_policy.into()), + truncation_policy: Some(turn_context.truncation_policy), }; let previous_turn_id = previous_context_item .turn_id @@ -1319,7 +1319,7 @@ async fn record_initial_history_resumed_replaced_incomplete_compacted_turn_clear user_instructions: None, developer_instructions: None, final_output_json_schema: None, - truncation_policy: Some(turn_context.truncation_policy.into()), + truncation_policy: Some(turn_context.truncation_policy), }; let previous_turn_id = previous_context_item .turn_id diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 8846f0a64..0ba0a1beb 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -1287,7 +1287,7 @@ async fn record_initial_history_forked_hydrates_previous_turn_settings() { user_instructions: None, developer_instructions: None, final_output_json_schema: None, - truncation_policy: Some(turn_context.truncation_policy.into()), + truncation_policy: Some(turn_context.truncation_policy), }; let turn_id = previous_context_item .turn_id diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 7686c5b65..82fd5c15e 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -15,9 +15,6 @@ use crate::protocol::CompactedItem; use crate::protocol::EventMsg; use crate::protocol::TurnStartedEvent; use crate::protocol::WarningEvent; -use crate::truncate::TruncationPolicy; -use crate::truncate::approx_token_count; -use crate::truncate::truncate_text; use crate::util::backoff; use codex_protocol::items::ContextCompactionItem; use codex_protocol::items::TurnItem; @@ -25,6 +22,9 @@ use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; use codex_protocol::user_input::UserInput; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::approx_token_count; +use codex_utils_output_truncation::truncate_text; use futures::prelude::*; use tracing::error; diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index b518d8ce4..ec2df30d3 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -3,12 +3,6 @@ use crate::context_manager::normalize; 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::truncate::TruncationPolicy; -use crate::truncate::approx_bytes_for_tokens; -use crate::truncate::approx_token_count; -use crate::truncate::approx_tokens_from_byte_count_i64; -use crate::truncate::truncate_function_output_items_with_policy; -use crate::truncate::truncate_text; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_protocol::models::BaseInstructions; @@ -25,6 +19,12 @@ use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::TurnContextItem; use codex_utils_cache::BlockingLruCache; use codex_utils_cache::sha1_digest; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::approx_bytes_for_tokens; +use codex_utils_output_truncation::approx_token_count; +use codex_utils_output_truncation::approx_tokens_from_byte_count_i64; +use codex_utils_output_truncation::truncate_function_output_items_with_policy; +use codex_utils_output_truncation::truncate_text; use std::num::NonZeroUsize; use std::ops::Deref; use std::sync::LazyLock; diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 0c92edcfb..3c508e05f 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -1,6 +1,4 @@ use super::*; -use crate::truncate; -use crate::truncate::TruncationPolicy; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_git_utils::GhostCommit; @@ -23,6 +21,8 @@ use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::TurnContextItem; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::truncate_text; use image::ImageBuffer; use image::ImageFormat; use image::Rgba; @@ -179,7 +179,7 @@ fn reasoning_with_encrypted_content(len: usize) -> ResponseItem { } fn truncate_exec_output(content: &str) -> String { - truncate::truncate_text(content, TruncationPolicy::Tokens(EXEC_FORMAT_MAX_TOKENS)) + truncate_text(content, TruncationPolicy::Tokens(EXEC_FORMAT_MAX_TOKENS)) } fn approx_token_count_for_text(text: &str) -> i64 { diff --git a/codex-rs/core/src/error.rs b/codex-rs/core/src/error.rs index b84180d16..b0c55a55c 100644 --- a/codex-rs/core/src/error.rs +++ b/codex-rs/core/src/error.rs @@ -2,8 +2,6 @@ use crate::exec::ExecToolCallOutput; use crate::network_policy_decision::NetworkPolicyDecisionPayload; use crate::token_data::KnownPlan; use crate::token_data::PlanType; -use crate::truncate::TruncationPolicy; -use crate::truncate::truncate_text; use chrono::DateTime; use chrono::Datelike; use chrono::Local; @@ -15,6 +13,8 @@ use codex_protocol::ThreadId; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::ErrorEvent; use codex_protocol::protocol::RateLimitSnapshot; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::truncate_text; use reqwest::StatusCode; use serde_json; use std::io; diff --git a/codex-rs/core/src/guardian/prompt.rs b/codex-rs/core/src/guardian/prompt.rs index be0291648..5741a65fc 100644 --- a/codex-rs/core/src/guardian/prompt.rs +++ b/codex-rs/core/src/guardian/prompt.rs @@ -7,9 +7,9 @@ use serde_json::Value; use crate::codex::Session; use crate::compact::content_items_to_text; use crate::event_mapping::is_contextual_user_message_content; -use crate::truncate::approx_bytes_for_tokens; -use crate::truncate::approx_token_count; -use crate::truncate::approx_tokens_from_byte_count; +use codex_utils_output_truncation::approx_bytes_for_tokens; +use codex_utils_output_truncation::approx_token_count; +use codex_utils_output_truncation::approx_tokens_from_byte_count; use super::GUARDIAN_MAX_MESSAGE_ENTRY_TOKENS; use super::GUARDIAN_MAX_MESSAGE_TRANSCRIPT_TOKENS; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index decebc84e..d4a8ae93a 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -78,7 +78,6 @@ mod stream_events_utils; pub mod test_support; mod text_encoding; pub use codex_login::token_data; -mod truncate; mod unified_exec; pub mod windows_sandbox; pub use client::X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER; diff --git a/codex-rs/core/src/memories/prompts.rs b/codex-rs/core/src/memories/prompts.rs index 1659e1c1c..851179426 100644 --- a/codex-rs/core/src/memories/prompts.rs +++ b/codex-rs/core/src/memories/prompts.rs @@ -1,13 +1,13 @@ use crate::memories::memory_root; use crate::memories::phase_one; use crate::memories::storage::rollout_summary_file_stem_from_parts; -use crate::truncate::TruncationPolicy; -use crate::truncate::truncate_text; use askama::Template; use codex_protocol::openai_models::ModelInfo; use codex_state::Phase2InputSelection; use codex_state::Stage1Output; use codex_state::Stage1OutputRef; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::truncate_text; use std::path::Path; use tokio::fs; use tracing::warn; diff --git a/codex-rs/core/src/models_manager/model_info.rs b/codex-rs/core/src/models_manager/model_info.rs index 7d3c6e9d1..055a6459d 100644 --- a/codex-rs/core/src/models_manager/model_info.rs +++ b/codex-rs/core/src/models_manager/model_info.rs @@ -10,8 +10,8 @@ use codex_protocol::openai_models::WebSearchToolType; use codex_protocol::openai_models::default_input_modalities; use crate::config::Config; -use crate::truncate::approx_bytes_for_tokens; use codex_features::Feature; +use codex_utils_output_truncation::approx_bytes_for_tokens; use tracing::warn; pub const BASE_INSTRUCTIONS: &str = include_str!("../../prompt.md"); diff --git a/codex-rs/core/src/realtime_context.rs b/codex-rs/core/src/realtime_context.rs index bcb24c567..1b845db1e 100644 --- a/codex-rs/core/src/realtime_context.rs +++ b/codex-rs/core/src/realtime_context.rs @@ -1,13 +1,13 @@ use crate::codex::Session; use crate::compact::content_items_to_text; use crate::event_mapping::is_contextual_user_message_content; -use crate::truncate::TruncationPolicy; -use crate::truncate::truncate_text; use chrono::Utc; use codex_git_utils::resolve_root_git_project_for_trust; use codex_protocol::models::ResponseItem; use codex_state::SortKey; use codex_state::ThreadMetadata; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::truncate_text; use dirs::home_dir; use std::cmp::Reverse; use std::collections::HashMap; diff --git a/codex-rs/core/src/rollout/list.rs b/codex-rs/core/src/rollout/list.rs index bf32e23c8..6ae5908d8 100644 --- a/codex-rs/core/src/rollout/list.rs +++ b/codex-rs/core/src/rollout/list.rs @@ -1,7 +1,7 @@ use async_trait::async_trait; use std::cmp::Reverse; use std::ffi::OsStr; -use std::io::{self}; +use std::io; use std::num::NonZero; use std::ops::ControlFlow; use std::path::Path; diff --git a/codex-rs/core/src/rollout/recorder.rs b/codex-rs/core/src/rollout/recorder.rs index 25e7ab203..30b0734df 100644 --- a/codex-rs/core/src/rollout/recorder.rs +++ b/codex-rs/core/src/rollout/recorder.rs @@ -1,7 +1,7 @@ //! Persist Codex session rollouts (.jsonl) so sessions can be replayed or inspected later. +use std::fs; use std::fs::File; -use std::fs::{self}; use std::io::Error as IoError; use std::path::Path; use std::path::PathBuf; @@ -17,8 +17,8 @@ use time::format_description::FormatItem; use time::format_description::well_known::Rfc3339; use time::macros::format_description; use tokio::io::AsyncWriteExt; +use tokio::sync::mpsc; use tokio::sync::mpsc::Sender; -use tokio::sync::mpsc::{self}; use tokio::sync::oneshot; use tracing::info; use tracing::trace; @@ -44,8 +44,6 @@ use crate::default_client::originator; use crate::path_utils; use crate::state_db; use crate::state_db::StateDbHandle; -use crate::truncate::TruncationPolicy; -use crate::truncate::truncate_text; use codex_git_utils::collect_git_info; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::GitInfo as ProtocolGitInfo; @@ -58,6 +56,8 @@ use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_state::StateRuntime; use codex_state::ThreadMetadataBuilder; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::truncate_text; /// Records all [`ResponseItem`]s for a session and flushes them to disk after /// every update. diff --git a/codex-rs/core/src/rollout/recorder_tests.rs b/codex-rs/core/src/rollout/recorder_tests.rs index 8ca7b58a6..9443c877b 100644 --- a/codex-rs/core/src/rollout/recorder_tests.rs +++ b/codex-rs/core/src/rollout/recorder_tests.rs @@ -10,8 +10,8 @@ use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::TurnContextItem; use codex_protocol::protocol::UserMessageEvent; use pretty_assertions::assert_eq; +use std::fs; use std::fs::File; -use std::fs::{self}; use std::io::Write; use std::path::Path; use std::path::PathBuf; diff --git a/codex-rs/core/src/rollout/tests.rs b/codex-rs/core/src/rollout/tests.rs index 44e536e50..17c1792d3 100644 --- a/codex-rs/core/src/rollout/tests.rs +++ b/codex-rs/core/src/rollout/tests.rs @@ -1,9 +1,9 @@ #![allow(clippy::unwrap_used, clippy::expect_used)] use std::ffi::OsStr; +use std::fs; use std::fs::File; use std::fs::FileTimes; -use std::fs::{self}; use std::io::Write; use std::path::Path; diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 81978160a..1a1423616 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -13,8 +13,8 @@ use crate::protocol::RateLimitSnapshot; use crate::protocol::TokenUsage; use crate::protocol::TokenUsageInfo; use crate::session_startup_prewarm::SessionStartupPrewarmHandle; -use crate::truncate::TruncationPolicy; use codex_protocol::protocol::TurnContextItem; +use codex_utils_output_truncation::TruncationPolicy; /// Persistent, session-scoped state previously stored directly on `Session`. pub(crate) struct SessionState { diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index a4838d246..a1f21a622 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -27,11 +27,11 @@ use crate::tools::parallel::ToolCallRuntime; use crate::tools::router::ToolCall; use crate::tools::router::ToolCallSource; use crate::tools::router::ToolRouterParams; -use crate::truncate::TruncationPolicy; -use crate::truncate::formatted_truncate_text_content_items_with_policy; -use crate::truncate::truncate_function_output_items_with_policy; use crate::unified_exec::resolve_max_tokens; use codex_features::Feature; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::formatted_truncate_text_content_items_with_policy; +use codex_utils_output_truncation::truncate_function_output_items_with_policy; pub(crate) use execute_handler::CodeModeExecuteHandler; use response_adapter::into_function_call_output_content_items; diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 74efb0989..12db45102 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -4,8 +4,6 @@ use crate::codex::TurnContext; use crate::tools::TELEMETRY_PREVIEW_MAX_BYTES; use crate::tools::TELEMETRY_PREVIEW_MAX_LINES; use crate::tools::TELEMETRY_PREVIEW_TRUNCATION_NOTICE; -use crate::truncate::TruncationPolicy; -use crate::truncate::formatted_truncate_text; use crate::turn_diff_tracker::TurnDiffTracker; use crate::unified_exec::resolve_max_tokens; use codex_protocol::mcp::CallToolResult; @@ -16,6 +14,8 @@ use codex_protocol::models::ResponseInputItem; use codex_protocol::models::SearchToolCallParams; use codex_protocol::models::ShellToolCallParams; use codex_protocol::models::function_call_output_content_items_to_text; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::formatted_truncate_text; use codex_utils_string::take_bytes_at_char_boundary; use serde::Serialize; use serde_json::Value as JsonValue; diff --git a/codex-rs/core/src/tools/js_repl/mod.rs b/codex-rs/core/src/tools/js_repl/mod.rs index 27d8cb9e6..0a29fb120 100644 --- a/codex-rs/core/src/tools/js_repl/mod.rs +++ b/codex-rs/core/src/tools/js_repl/mod.rs @@ -42,12 +42,12 @@ use crate::original_image_detail::normalize_output_image_detail; use crate::sandboxing::ExecOptions; use crate::tools::ToolRouter; use crate::tools::context::SharedTurnDiffTracker; -use crate::truncate::TruncationPolicy; -use crate::truncate::truncate_text; use codex_sandboxing::SandboxCommand; use codex_sandboxing::SandboxManager; use codex_sandboxing::SandboxTransformRequest; use codex_sandboxing::SandboxablePreference; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::truncate_text; pub(crate) const JS_REPL_PRAGMA_PREFIX: &str = "// codex-js-repl:"; const KERNEL_SOURCE: &str = include_str!("kernel.js"); diff --git a/codex-rs/core/src/tools/mod.rs b/codex-rs/core/src/tools/mod.rs index 4e495190e..87a01c8e3 100644 --- a/codex-rs/core/src/tools/mod.rs +++ b/codex-rs/core/src/tools/mod.rs @@ -15,9 +15,9 @@ pub mod sandboxing; pub mod spec; use crate::exec::ExecToolCallOutput; -use crate::truncate::TruncationPolicy; -use crate::truncate::formatted_truncate_text; -use crate::truncate::truncate_text; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::formatted_truncate_text; +use codex_utils_output_truncation::truncate_text; pub use router::ToolRouter; use serde::Serialize; diff --git a/codex-rs/core/src/truncate.rs b/codex-rs/core/src/truncate.rs deleted file mode 100644 index 707fbe22e..000000000 --- a/codex-rs/core/src/truncate.rs +++ /dev/null @@ -1,363 +0,0 @@ -//! Utilities for truncating large chunks of output while preserving a prefix -//! and suffix on UTF-8 boundaries, and helpers for line/token‑based truncation -//! used across the core crate. - -use codex_protocol::models::FunctionCallOutputContentItem; -use codex_protocol::openai_models::TruncationMode; -use codex_protocol::openai_models::TruncationPolicyConfig; -use codex_protocol::protocol::TruncationPolicy as ProtocolTruncationPolicy; - -const APPROX_BYTES_PER_TOKEN: usize = 4; - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub enum TruncationPolicy { - Bytes(usize), - Tokens(usize), -} - -impl From for ProtocolTruncationPolicy { - fn from(value: TruncationPolicy) -> Self { - match value { - TruncationPolicy::Bytes(bytes) => Self::Bytes(bytes), - TruncationPolicy::Tokens(tokens) => Self::Tokens(tokens), - } - } -} - -impl From for TruncationPolicy { - fn from(config: TruncationPolicyConfig) -> Self { - match config.mode { - TruncationMode::Bytes => Self::Bytes(config.limit as usize), - TruncationMode::Tokens => Self::Tokens(config.limit as usize), - } - } -} - -impl TruncationPolicy { - /// Returns a token budget derived from this policy. - /// - /// - For `Tokens`, this is the explicit token limit. - /// - For `Bytes`, this is an approximate token budget using the global - /// bytes-per-token heuristic. - pub fn token_budget(&self) -> usize { - match self { - TruncationPolicy::Bytes(bytes) => { - usize::try_from(approx_tokens_from_byte_count(*bytes)).unwrap_or(usize::MAX) - } - TruncationPolicy::Tokens(tokens) => *tokens, - } - } - - /// Returns a byte budget derived from this policy. - /// - /// - For `Bytes`, this is the explicit byte limit. - /// - For `Tokens`, this is an approximate byte budget using the global - /// bytes-per-token heuristic. - pub fn byte_budget(&self) -> usize { - match self { - TruncationPolicy::Bytes(bytes) => *bytes, - TruncationPolicy::Tokens(tokens) => approx_bytes_for_tokens(*tokens), - } - } -} - -impl std::ops::Mul for TruncationPolicy { - type Output = Self; - - fn mul(self, multiplier: f64) -> Self::Output { - match self { - TruncationPolicy::Bytes(bytes) => { - TruncationPolicy::Bytes((bytes as f64 * multiplier).ceil() as usize) - } - TruncationPolicy::Tokens(tokens) => { - TruncationPolicy::Tokens((tokens as f64 * multiplier).ceil() as usize) - } - } - } -} - -pub(crate) fn formatted_truncate_text(content: &str, policy: TruncationPolicy) -> String { - if content.len() <= policy.byte_budget() { - return content.to_string(); - } - let total_lines = content.lines().count(); - let result = truncate_text(content, policy); - format!("Total output lines: {total_lines}\n\n{result}") -} - -pub(crate) fn truncate_text(content: &str, policy: TruncationPolicy) -> String { - match policy { - TruncationPolicy::Bytes(_) => truncate_with_byte_estimate(content, policy), - TruncationPolicy::Tokens(_) => { - let (truncated, _) = truncate_with_token_budget(content, policy); - truncated - } - } -} - -pub(crate) fn formatted_truncate_text_content_items_with_policy( - items: &[FunctionCallOutputContentItem], - policy: TruncationPolicy, -) -> (Vec, Option) { - let text_segments = items - .iter() - .filter_map(|item| match item { - FunctionCallOutputContentItem::InputText { text } => Some(text.as_str()), - FunctionCallOutputContentItem::InputImage { .. } => None, - }) - .collect::>(); - - if text_segments.is_empty() { - return (items.to_vec(), None); - } - - let mut combined = String::new(); - for text in &text_segments { - if !combined.is_empty() { - combined.push('\n'); - } - combined.push_str(text); - } - - if combined.len() <= policy.byte_budget() { - return (items.to_vec(), None); - } - - let mut out = vec![FunctionCallOutputContentItem::InputText { - text: formatted_truncate_text(&combined, policy), - }]; - out.extend(items.iter().filter_map(|item| match item { - FunctionCallOutputContentItem::InputImage { image_url, detail } => { - Some(FunctionCallOutputContentItem::InputImage { - image_url: image_url.clone(), - detail: *detail, - }) - } - FunctionCallOutputContentItem::InputText { .. } => None, - })); - - (out, Some(approx_token_count(&combined))) -} - -/// Globally truncate function output items to fit within the given -/// truncation policy's budget, preserving as many text/image items as -/// possible and appending a summary for any omitted text items. -pub(crate) fn truncate_function_output_items_with_policy( - items: &[FunctionCallOutputContentItem], - policy: TruncationPolicy, -) -> Vec { - let mut out: Vec = Vec::with_capacity(items.len()); - let mut remaining_budget = match policy { - TruncationPolicy::Bytes(_) => policy.byte_budget(), - TruncationPolicy::Tokens(_) => policy.token_budget(), - }; - let mut omitted_text_items = 0usize; - - for it in items { - match it { - FunctionCallOutputContentItem::InputText { text } => { - if remaining_budget == 0 { - omitted_text_items += 1; - continue; - } - - let cost = match policy { - TruncationPolicy::Bytes(_) => text.len(), - TruncationPolicy::Tokens(_) => approx_token_count(text), - }; - - if cost <= remaining_budget { - out.push(FunctionCallOutputContentItem::InputText { text: text.clone() }); - remaining_budget = remaining_budget.saturating_sub(cost); - } else { - let snippet_policy = match policy { - TruncationPolicy::Bytes(_) => TruncationPolicy::Bytes(remaining_budget), - TruncationPolicy::Tokens(_) => TruncationPolicy::Tokens(remaining_budget), - }; - let snippet = truncate_text(text, snippet_policy); - if snippet.is_empty() { - omitted_text_items += 1; - } else { - out.push(FunctionCallOutputContentItem::InputText { text: snippet }); - } - remaining_budget = 0; - } - } - FunctionCallOutputContentItem::InputImage { image_url, detail } => { - out.push(FunctionCallOutputContentItem::InputImage { - image_url: image_url.clone(), - detail: *detail, - }); - } - } - } - - if omitted_text_items > 0 { - out.push(FunctionCallOutputContentItem::InputText { - text: format!("[omitted {omitted_text_items} text items ...]"), - }); - } - - out -} - -/// Truncate the middle of a UTF-8 string to at most `max_tokens` tokens, -/// preserving the beginning and the end. Returns the possibly truncated string -/// and `Some(original_token_count)` if truncation occurred; otherwise returns -/// the original string and `None`. -fn truncate_with_token_budget(s: &str, policy: TruncationPolicy) -> (String, Option) { - if s.is_empty() { - return (String::new(), None); - } - let max_tokens = policy.token_budget(); - - let byte_len = s.len(); - if max_tokens > 0 && byte_len <= approx_bytes_for_tokens(max_tokens) { - return (s.to_string(), None); - } - - let truncated = truncate_with_byte_estimate(s, policy); - let approx_total_usize = approx_token_count(s); - let approx_total = u64::try_from(approx_total_usize).unwrap_or(u64::MAX); - if truncated == s { - (truncated, None) - } else { - (truncated, Some(approx_total)) - } -} - -/// Truncate a string using a byte budget derived from the token budget, without -/// performing any real tokenization. This keeps the logic purely byte-based and -/// uses a bytes placeholder in the truncated output. -fn truncate_with_byte_estimate(s: &str, policy: TruncationPolicy) -> String { - if s.is_empty() { - return String::new(); - } - - let total_chars = s.chars().count(); - let max_bytes = policy.byte_budget(); - - if max_bytes == 0 { - // No budget to show content; just report that everything was truncated. - let marker = format_truncation_marker( - policy, - removed_units_for_source(policy, s.len(), total_chars), - ); - return marker; - } - - if s.len() <= max_bytes { - return s.to_string(); - } - - let total_bytes = s.len(); - - let (left_budget, right_budget) = split_budget(max_bytes); - - let (removed_chars, left, right) = split_string(s, left_budget, right_budget); - - let marker = format_truncation_marker( - policy, - removed_units_for_source(policy, total_bytes.saturating_sub(max_bytes), removed_chars), - ); - - assemble_truncated_output(left, right, &marker) -} - -fn split_string(s: &str, beginning_bytes: usize, end_bytes: usize) -> (usize, &str, &str) { - if s.is_empty() { - return (0, "", ""); - } - - let len = s.len(); - let tail_start_target = len.saturating_sub(end_bytes); - let mut prefix_end = 0usize; - let mut suffix_start = len; - let mut removed_chars = 0usize; - let mut suffix_started = false; - - for (idx, ch) in s.char_indices() { - let char_end = idx + ch.len_utf8(); - if char_end <= beginning_bytes { - prefix_end = char_end; - continue; - } - - if idx >= tail_start_target { - if !suffix_started { - suffix_start = idx; - suffix_started = true; - } - continue; - } - - removed_chars = removed_chars.saturating_add(1); - } - - if suffix_start < prefix_end { - suffix_start = prefix_end; - } - - let before = &s[..prefix_end]; - let after = &s[suffix_start..]; - - (removed_chars, before, after) -} - -fn format_truncation_marker(policy: TruncationPolicy, removed_count: u64) -> String { - match policy { - TruncationPolicy::Tokens(_) => format!("…{removed_count} tokens truncated…"), - TruncationPolicy::Bytes(_) => format!("…{removed_count} chars truncated…"), - } -} - -fn split_budget(budget: usize) -> (usize, usize) { - let left = budget / 2; - (left, budget - left) -} - -fn removed_units_for_source( - policy: TruncationPolicy, - removed_bytes: usize, - removed_chars: usize, -) -> u64 { - match policy { - TruncationPolicy::Tokens(_) => approx_tokens_from_byte_count(removed_bytes), - TruncationPolicy::Bytes(_) => u64::try_from(removed_chars).unwrap_or(u64::MAX), - } -} - -fn assemble_truncated_output(prefix: &str, suffix: &str, marker: &str) -> String { - let mut out = String::with_capacity(prefix.len() + marker.len() + suffix.len() + 1); - out.push_str(prefix); - out.push_str(marker); - out.push_str(suffix); - out -} - -pub(crate) fn approx_token_count(text: &str) -> usize { - let len = text.len(); - len.saturating_add(APPROX_BYTES_PER_TOKEN.saturating_sub(1)) / APPROX_BYTES_PER_TOKEN -} - -pub(crate) fn approx_bytes_for_tokens(tokens: usize) -> usize { - tokens.saturating_mul(APPROX_BYTES_PER_TOKEN) -} - -pub(crate) fn approx_tokens_from_byte_count(bytes: usize) -> u64 { - let bytes_u64 = bytes as u64; - bytes_u64.saturating_add((APPROX_BYTES_PER_TOKEN as u64).saturating_sub(1)) - / (APPROX_BYTES_PER_TOKEN as u64) -} - -pub(crate) fn approx_tokens_from_byte_count_i64(bytes: i64) -> i64 { - if bytes <= 0 { - return 0; - } - let bytes = usize::try_from(bytes).unwrap_or(usize::MAX); - i64::try_from(approx_tokens_from_byte_count(bytes)).unwrap_or(i64::MAX) -} - -#[cfg(test)] -#[path = "truncate_tests.rs"] -mod tests; diff --git a/codex-rs/core/src/unified_exec/process.rs b/codex-rs/core/src/unified_exec/process.rs index 58faed27d..68cf18bce 100644 --- a/codex-rs/core/src/unified_exec/process.rs +++ b/codex-rs/core/src/unified_exec/process.rs @@ -15,9 +15,9 @@ use tokio_util::sync::CancellationToken; use crate::exec::ExecToolCallOutput; use crate::exec::StreamOutput; use crate::exec::is_likely_sandbox_denied; -use crate::truncate::TruncationPolicy; -use crate::truncate::formatted_truncate_text; use codex_sandboxing::SandboxType; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::formatted_truncate_text; use codex_utils_pty::ExecCommandSession; use codex_utils_pty::SpawnedPty; diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 52d668c00..a760249d6 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -27,7 +27,6 @@ use crate::tools::orchestrator::ToolOrchestrator; use crate::tools::runtimes::unified_exec::UnifiedExecRequest as UnifiedExecToolRequest; use crate::tools::runtimes::unified_exec::UnifiedExecRuntime; use crate::tools::sandboxing::ToolCtx; -use crate::truncate::approx_token_count; use crate::unified_exec::ExecCommandRequest; use crate::unified_exec::MAX_UNIFIED_EXEC_PROCESSES; use crate::unified_exec::MAX_YIELD_TIME_MS; @@ -50,6 +49,7 @@ use crate::unified_exec::process::OutputBuffer; use crate::unified_exec::process::OutputHandles; use crate::unified_exec::process::SpawnLifecycleHandle; use crate::unified_exec::process::UnifiedExecProcess; +use codex_utils_output_truncation::approx_token_count; const UNIFIED_EXEC_ENV: [(&str, &str); 10] = [ ("NO_COLOR", "1"), diff --git a/codex-rs/protocol/Cargo.toml b/codex-rs/protocol/Cargo.toml index eaaf791bb..ebfdaddab 100644 --- a/codex-rs/protocol/Cargo.toml +++ b/codex-rs/protocol/Cargo.toml @@ -16,6 +16,7 @@ codex-execpolicy = { workspace = true } codex-git-utils = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-image = { workspace = true } +codex-utils-string = { workspace = true } icu_decimal = { workspace = true } icu_locale_core = { workspace = true } icu_provider = { workspace = true, features = ["sync"] } diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index c40514e3d..6856f1941 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -7,6 +7,7 @@ use std::collections::HashMap; use std::collections::HashSet; use std::ffi::OsStr; use std::fmt; +use std::ops::Mul; use std::path::Path; use std::path::PathBuf; use std::str::FromStr; @@ -2629,6 +2630,51 @@ pub enum TruncationPolicy { Tokens(usize), } +impl From for TruncationPolicy { + fn from(config: crate::openai_models::TruncationPolicyConfig) -> Self { + match config.mode { + crate::openai_models::TruncationMode::Bytes => Self::Bytes(config.limit as usize), + crate::openai_models::TruncationMode::Tokens => Self::Tokens(config.limit as usize), + } + } +} + +impl TruncationPolicy { + pub fn token_budget(&self) -> usize { + match self { + TruncationPolicy::Bytes(bytes) => { + usize::try_from(codex_utils_string::approx_tokens_from_byte_count(*bytes)) + .unwrap_or(usize::MAX) + } + TruncationPolicy::Tokens(tokens) => *tokens, + } + } + + pub fn byte_budget(&self) -> usize { + match self { + TruncationPolicy::Bytes(bytes) => *bytes, + TruncationPolicy::Tokens(tokens) => { + codex_utils_string::approx_bytes_for_tokens(*tokens) + } + } + } +} + +impl Mul for TruncationPolicy { + type Output = Self; + + fn mul(self, multiplier: f64) -> Self::Output { + match self { + TruncationPolicy::Bytes(bytes) => { + TruncationPolicy::Bytes((bytes as f64 * multiplier).ceil() as usize) + } + TruncationPolicy::Tokens(tokens) => { + TruncationPolicy::Tokens((tokens as f64 * multiplier).ceil() as usize) + } + } + } +} + #[derive(Serialize, Deserialize, Clone, JsonSchema)] pub struct RolloutLine { pub timestamp: String, diff --git a/codex-rs/utils/output-truncation/BUILD.bazel b/codex-rs/utils/output-truncation/BUILD.bazel new file mode 100644 index 000000000..eefd22ddb --- /dev/null +++ b/codex-rs/utils/output-truncation/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "output-truncation", + crate_name = "codex_utils_output_truncation", +) diff --git a/codex-rs/utils/output-truncation/Cargo.toml b/codex-rs/utils/output-truncation/Cargo.toml new file mode 100644 index 000000000..7ad0ccfd4 --- /dev/null +++ b/codex-rs/utils/output-truncation/Cargo.toml @@ -0,0 +1,15 @@ +[package] +edition.workspace = true +license.workspace = true +name = "codex-utils-output-truncation" +version.workspace = true + +[lints] +workspace = true + +[dependencies] +codex-protocol = { workspace = true } +codex-utils-string = { workspace = true } + +[dev-dependencies] +pretty_assertions = { workspace = true } diff --git a/codex-rs/utils/output-truncation/src/lib.rs b/codex-rs/utils/output-truncation/src/lib.rs new file mode 100644 index 000000000..24b1630da --- /dev/null +++ b/codex-rs/utils/output-truncation/src/lib.rs @@ -0,0 +1,142 @@ +//! Helpers for truncating tool and exec output using [`TruncationPolicy`](codex_protocol::protocol::TruncationPolicy). + +use codex_protocol::models::FunctionCallOutputContentItem; +pub use codex_utils_string::approx_bytes_for_tokens; +pub use codex_utils_string::approx_token_count; +pub use codex_utils_string::approx_tokens_from_byte_count; +use codex_utils_string::truncate_middle_chars; +use codex_utils_string::truncate_middle_with_token_budget; + +pub use codex_protocol::protocol::TruncationPolicy; + +pub fn formatted_truncate_text(content: &str, policy: TruncationPolicy) -> String { + if content.len() <= policy.byte_budget() { + return content.to_string(); + } + + let total_lines = content.lines().count(); + let result = truncate_text(content, policy); + format!("Total output lines: {total_lines}\n\n{result}") +} + +pub fn truncate_text(content: &str, policy: TruncationPolicy) -> String { + match policy { + TruncationPolicy::Bytes(bytes) => truncate_middle_chars(content, bytes), + TruncationPolicy::Tokens(tokens) => truncate_middle_with_token_budget(content, tokens).0, + } +} + +pub fn formatted_truncate_text_content_items_with_policy( + items: &[FunctionCallOutputContentItem], + policy: TruncationPolicy, +) -> (Vec, Option) { + let text_segments = items + .iter() + .filter_map(|item| match item { + FunctionCallOutputContentItem::InputText { text } => Some(text.as_str()), + FunctionCallOutputContentItem::InputImage { .. } => None, + }) + .collect::>(); + + if text_segments.is_empty() { + return (items.to_vec(), None); + } + + let mut combined = String::new(); + for text in &text_segments { + if !combined.is_empty() { + combined.push('\n'); + } + combined.push_str(text); + } + + if combined.len() <= policy.byte_budget() { + return (items.to_vec(), None); + } + + let mut out = vec![FunctionCallOutputContentItem::InputText { + text: formatted_truncate_text(&combined, policy), + }]; + out.extend(items.iter().filter_map(|item| match item { + FunctionCallOutputContentItem::InputImage { image_url, detail } => { + Some(FunctionCallOutputContentItem::InputImage { + image_url: image_url.clone(), + detail: *detail, + }) + } + FunctionCallOutputContentItem::InputText { .. } => None, + })); + + (out, Some(approx_token_count(&combined))) +} + +pub fn truncate_function_output_items_with_policy( + items: &[FunctionCallOutputContentItem], + policy: TruncationPolicy, +) -> Vec { + let mut out: Vec = Vec::with_capacity(items.len()); + let mut remaining_budget = match policy { + TruncationPolicy::Bytes(_) => policy.byte_budget(), + TruncationPolicy::Tokens(_) => policy.token_budget(), + }; + let mut omitted_text_items = 0usize; + + for item in items { + match item { + FunctionCallOutputContentItem::InputText { text } => { + if remaining_budget == 0 { + omitted_text_items += 1; + continue; + } + + let cost = match policy { + TruncationPolicy::Bytes(_) => text.len(), + TruncationPolicy::Tokens(_) => approx_token_count(text), + }; + + if cost <= remaining_budget { + out.push(FunctionCallOutputContentItem::InputText { text: text.clone() }); + remaining_budget = remaining_budget.saturating_sub(cost); + } else { + let snippet_policy = match policy { + TruncationPolicy::Bytes(_) => TruncationPolicy::Bytes(remaining_budget), + TruncationPolicy::Tokens(_) => TruncationPolicy::Tokens(remaining_budget), + }; + let snippet = truncate_text(text, snippet_policy); + if snippet.is_empty() { + omitted_text_items += 1; + } else { + out.push(FunctionCallOutputContentItem::InputText { text: snippet }); + } + remaining_budget = 0; + } + } + FunctionCallOutputContentItem::InputImage { image_url, detail } => { + out.push(FunctionCallOutputContentItem::InputImage { + image_url: image_url.clone(), + detail: *detail, + }); + } + } + } + + if omitted_text_items > 0 { + out.push(FunctionCallOutputContentItem::InputText { + text: format!("[omitted {omitted_text_items} text items ...]"), + }); + } + + out +} + +pub fn approx_tokens_from_byte_count_i64(bytes: i64) -> i64 { + if bytes <= 0 { + return 0; + } + + let bytes = usize::try_from(bytes).unwrap_or(usize::MAX); + i64::try_from(approx_tokens_from_byte_count(bytes)).unwrap_or(i64::MAX) +} + +#[cfg(test)] +mod truncate_tests; diff --git a/codex-rs/core/src/truncate_tests.rs b/codex-rs/utils/output-truncation/src/truncate_tests.rs similarity index 77% rename from codex-rs/core/src/truncate_tests.rs rename to codex-rs/utils/output-truncation/src/truncate_tests.rs index 5a61a9a26..f159a6b62 100644 --- a/codex-rs/core/src/truncate_tests.rs +++ b/codex-rs/utils/output-truncation/src/truncate_tests.rs @@ -1,49 +1,13 @@ -use super::TruncationPolicy; -use super::approx_token_count; -use super::formatted_truncate_text; -use super::formatted_truncate_text_content_items_with_policy; -use super::split_string; -use super::truncate_function_output_items_with_policy; -use super::truncate_text; -use super::truncate_with_token_budget; +use crate::TruncationPolicy; +use crate::approx_token_count; +use crate::approx_tokens_from_byte_count_i64; +use crate::formatted_truncate_text; +use crate::formatted_truncate_text_content_items_with_policy; +use crate::truncate_function_output_items_with_policy; +use crate::truncate_text; use codex_protocol::models::FunctionCallOutputContentItem; use pretty_assertions::assert_eq; -#[test] -fn split_string_works() { - assert_eq!(split_string("hello world", 5, 5), (1, "hello", "world")); - assert_eq!(split_string("abc", 0, 0), (3, "", "")); -} - -#[test] -fn split_string_handles_empty_string() { - assert_eq!(split_string("", 4, 4), (0, "", "")); -} - -#[test] -fn split_string_only_keeps_prefix_when_tail_budget_is_zero() { - assert_eq!(split_string("abcdef", 3, 0), (3, "abc", "")); -} - -#[test] -fn split_string_only_keeps_suffix_when_prefix_budget_is_zero() { - assert_eq!(split_string("abcdef", 0, 3), (3, "", "def")); -} - -#[test] -fn split_string_handles_overlapping_budgets_without_removal() { - assert_eq!(split_string("abcdef", 4, 4), (0, "abcd", "ef")); -} - -#[test] -fn split_string_respects_utf8_boundaries() { - assert_eq!(split_string("😀abc😀", 5, 5), (1, "😀a", "c😀")); - - assert_eq!(split_string("😀😀😀😀😀", 1, 1), (5, "", "")); - assert_eq!(split_string("😀😀😀😀😀", 7, 7), (3, "😀", "😀")); - assert_eq!(split_string("😀😀😀😀😀", 8, 8), (1, "😀😀", "😀😀")); -} - #[test] fn truncate_bytes_less_than_placeholder_returns_placeholder() { let content = "example output"; @@ -126,31 +90,6 @@ fn truncate_tokens_reports_original_line_count_when_truncated() { ); } -#[test] -fn truncate_with_token_budget_returns_original_when_under_limit() { - let s = "short output"; - let limit = 100; - let (out, original) = truncate_with_token_budget(s, TruncationPolicy::Tokens(limit)); - assert_eq!(out, s); - assert_eq!(original, None); -} - -#[test] -fn truncate_with_token_budget_reports_truncation_at_zero_limit() { - let s = "abcdef"; - let (out, original) = truncate_with_token_budget(s, TruncationPolicy::Tokens(0)); - assert_eq!(out, "…2 tokens truncated…"); - assert_eq!(original, Some(2)); -} - -#[test] -fn truncate_middle_tokens_handles_utf8_content() { - let s = "😀😀😀😀😀😀😀😀😀😀\nsecond line with text\n"; - let (out, tokens) = truncate_with_token_budget(s, TruncationPolicy::Tokens(8)); - assert_eq!(out, "😀😀😀😀…8 tokens truncated… line with text\n"); - assert_eq!(tokens, Some(16)); -} - #[test] fn truncate_middle_bytes_handles_utf8_content() { let s = "😀😀😀😀😀😀😀😀😀😀\nsecond line with text\n"; @@ -185,7 +124,6 @@ fn truncates_across_multiple_under_limit_texts_and_reports_omitted() { let output = truncate_function_output_items_with_policy(&items, TruncationPolicy::Tokens(limit)); - // Expect: t1 (full), t2 (full), image, t3 (truncated), summary mentioning 2 omitted. assert_eq!(output.len(), 5); let first_text = match &output[0] { @@ -245,6 +183,29 @@ fn formatted_truncate_text_content_items_with_policy_returns_original_under_limi assert_eq!(original_token_count, None); } +#[test] +fn formatted_truncate_text_content_items_with_policy_preserves_empty_leading_text_behavior() { + let items = vec![ + FunctionCallOutputContentItem::InputText { + text: String::new(), + }, + FunctionCallOutputContentItem::InputText { + text: "abc".to_string(), + }, + ]; + + let (output, original_token_count) = + formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Bytes(0)); + + assert_eq!( + output, + vec![FunctionCallOutputContentItem::InputText { + text: "Total output lines: 1\n\n…3 chars truncated…".to_string(), + }] + ); + assert_eq!(original_token_count, Some(1)); +} + #[test] fn formatted_truncate_text_content_items_with_policy_merges_text_and_appends_images() { let items = vec![ @@ -311,3 +272,10 @@ fn formatted_truncate_text_content_items_with_policy_merges_all_text_for_token_b ); assert_eq!(original_token_count, Some(5)); } + +#[test] +fn byte_count_conversion_clamps_non_positive_values() { + assert_eq!(approx_tokens_from_byte_count_i64(/*bytes*/ -1), 0); + assert_eq!(approx_tokens_from_byte_count_i64(/*bytes*/ 0), 0); + assert_eq!(approx_tokens_from_byte_count_i64(/*bytes*/ 5), 2); +} diff --git a/codex-rs/utils/string/src/lib.rs b/codex-rs/utils/string/src/lib.rs index 0cb218f34..a667a517b 100644 --- a/codex-rs/utils/string/src/lib.rs +++ b/codex-rs/utils/string/src/lib.rs @@ -1,3 +1,11 @@ +mod truncate; + +pub use truncate::approx_bytes_for_tokens; +pub use truncate::approx_token_count; +pub use truncate::approx_tokens_from_byte_count; +pub use truncate::truncate_middle_chars; +pub use truncate::truncate_middle_with_token_budget; + // Truncate a &str to a byte budget at a char boundary (prefix) #[inline] pub fn take_bytes_at_char_boundary(s: &str, maxb: usize) -> &str { @@ -112,6 +120,7 @@ fn parse_markdown_hash_location_point(point: &str) -> Option<(&str, Option<&str> } #[cfg(test)] +#[allow(warnings, clippy::all)] mod tests { use super::find_uuids; use super::normalize_markdown_hash_location_suffix; diff --git a/codex-rs/utils/string/src/truncate.rs b/codex-rs/utils/string/src/truncate.rs new file mode 100644 index 000000000..fa5b87177 --- /dev/null +++ b/codex-rs/utils/string/src/truncate.rs @@ -0,0 +1,156 @@ +//! Utilities for truncating large chunks of output while preserving a prefix +//! and suffix on UTF-8 boundaries. + +const APPROX_BYTES_PER_TOKEN: usize = 4; + +/// Truncate a string to `max_bytes` using a character-count marker. +pub fn truncate_middle_chars(s: &str, max_bytes: usize) -> String { + truncate_with_byte_estimate(s, max_bytes, /*use_tokens*/ false) +} + +/// Truncate the middle of a UTF-8 string to at most `max_tokens` approximate +/// tokens, preserving the beginning and the end. Returns the possibly +/// truncated string and `Some(original_token_count)` if truncation occurred; +/// otherwise returns the original string and `None`. +pub fn truncate_middle_with_token_budget(s: &str, max_tokens: usize) -> (String, Option) { + if s.is_empty() { + return (String::new(), None); + } + + if max_tokens > 0 && s.len() <= approx_bytes_for_tokens(max_tokens) { + return (s.to_string(), None); + } + + let truncated = truncate_with_byte_estimate( + s, + approx_bytes_for_tokens(max_tokens), + /*use_tokens*/ true, + ); + let total_tokens = u64::try_from(approx_token_count(s)).unwrap_or(u64::MAX); + + if truncated == s { + (truncated, None) + } else { + (truncated, Some(total_tokens)) + } +} + +fn truncate_with_byte_estimate(s: &str, max_bytes: usize, use_tokens: bool) -> String { + if s.is_empty() { + return String::new(); + } + + let total_chars = s.chars().count(); + + if max_bytes == 0 { + return format_truncation_marker( + use_tokens, + removed_units(use_tokens, s.len(), total_chars), + ); + } + + if s.len() <= max_bytes { + return s.to_string(); + } + + let total_bytes = s.len(); + let (left_budget, right_budget) = split_budget(max_bytes); + let (removed_chars, left, right) = split_string(s, left_budget, right_budget); + let marker = format_truncation_marker( + use_tokens, + removed_units( + use_tokens, + total_bytes.saturating_sub(max_bytes), + removed_chars, + ), + ); + + assemble_truncated_output(left, right, &marker) +} + +pub fn approx_token_count(text: &str) -> usize { + let len = text.len(); + len.saturating_add(APPROX_BYTES_PER_TOKEN.saturating_sub(1)) / APPROX_BYTES_PER_TOKEN +} + +pub fn approx_bytes_for_tokens(tokens: usize) -> usize { + tokens.saturating_mul(APPROX_BYTES_PER_TOKEN) +} + +pub fn approx_tokens_from_byte_count(bytes: usize) -> u64 { + let bytes_u64 = bytes as u64; + bytes_u64.saturating_add((APPROX_BYTES_PER_TOKEN as u64).saturating_sub(1)) + / (APPROX_BYTES_PER_TOKEN as u64) +} + +fn split_string(s: &str, beginning_bytes: usize, end_bytes: usize) -> (usize, &str, &str) { + if s.is_empty() { + return (0, "", ""); + } + + let len = s.len(); + let tail_start_target = len.saturating_sub(end_bytes); + let mut prefix_end = 0usize; + let mut suffix_start = len; + let mut removed_chars = 0usize; + let mut suffix_started = false; + + for (idx, ch) in s.char_indices() { + let char_end = idx + ch.len_utf8(); + if char_end <= beginning_bytes { + prefix_end = char_end; + continue; + } + + if idx >= tail_start_target { + if !suffix_started { + suffix_start = idx; + suffix_started = true; + } + continue; + } + + removed_chars = removed_chars.saturating_add(1); + } + + if suffix_start < prefix_end { + suffix_start = prefix_end; + } + + let before = &s[..prefix_end]; + let after = &s[suffix_start..]; + + (removed_chars, before, after) +} + +fn split_budget(budget: usize) -> (usize, usize) { + let left = budget / 2; + (left, budget - left) +} + +fn format_truncation_marker(use_tokens: bool, removed_count: u64) -> String { + if use_tokens { + format!("…{removed_count} tokens truncated…") + } else { + format!("…{removed_count} chars truncated…") + } +} + +fn removed_units(use_tokens: bool, removed_bytes: usize, removed_chars: usize) -> u64 { + if use_tokens { + approx_tokens_from_byte_count(removed_bytes) + } else { + u64::try_from(removed_chars).unwrap_or(u64::MAX) + } +} + +fn assemble_truncated_output(prefix: &str, suffix: &str, marker: &str) -> String { + let mut out = String::with_capacity(prefix.len() + marker.len() + suffix.len() + 1); + out.push_str(prefix); + out.push_str(marker); + out.push_str(suffix); + out +} + +#[cfg(test)] +mod tests; diff --git a/codex-rs/utils/string/src/truncate/tests.rs b/codex-rs/utils/string/src/truncate/tests.rs new file mode 100644 index 000000000..5d517e86d --- /dev/null +++ b/codex-rs/utils/string/src/truncate/tests.rs @@ -0,0 +1,71 @@ +use super::split_string; +use super::truncate_middle_chars; +use super::truncate_middle_with_token_budget; +use pretty_assertions::assert_eq; + +#[test] +fn split_string_works() { + assert_eq!(split_string("hello world", 5, 5), (1, "hello", "world")); + assert_eq!(split_string("abc", 0, 0), (3, "", "")); +} + +#[test] +fn split_string_handles_empty_string() { + assert_eq!(split_string("", 4, 4), (0, "", "")); +} + +#[test] +fn split_string_only_keeps_prefix_when_tail_budget_is_zero() { + assert_eq!(split_string("abcdef", 3, 0), (3, "abc", "")); +} + +#[test] +fn split_string_only_keeps_suffix_when_prefix_budget_is_zero() { + assert_eq!(split_string("abcdef", 0, 3), (3, "", "def")); +} + +#[test] +fn split_string_handles_overlapping_budgets_without_removal() { + assert_eq!(split_string("abcdef", 4, 4), (0, "abcd", "ef")); +} + +#[test] +fn split_string_respects_utf8_boundaries() { + assert_eq!(split_string("😀abc😀", 5, 5), (1, "😀a", "c😀")); + + assert_eq!(split_string("😀😀😀😀😀", 1, 1), (5, "", "")); + assert_eq!(split_string("😀😀😀😀😀", 7, 7), (3, "😀", "😀")); + assert_eq!(split_string("😀😀😀😀😀", 8, 8), (1, "😀😀", "😀😀")); +} + +#[test] +fn truncate_with_token_budget_returns_original_when_under_limit() { + let s = "short output"; + let limit = 100; + let (out, original) = truncate_middle_with_token_budget(s, limit); + assert_eq!(out, s); + assert_eq!(original, None); +} + +#[test] +fn truncate_with_token_budget_reports_truncation_at_zero_limit() { + let s = "abcdef"; + let (out, original) = truncate_middle_with_token_budget(s, 0); + assert_eq!(out, "…2 tokens truncated…"); + assert_eq!(original, Some(2)); +} + +#[test] +fn truncate_middle_tokens_handles_utf8_content() { + let s = "😀😀😀😀😀😀😀😀😀😀\nsecond line with text\n"; + let (out, tokens) = truncate_middle_with_token_budget(s, 8); + assert_eq!(out, "😀😀😀😀…8 tokens truncated… line with text\n"); + assert_eq!(tokens, Some(16)); +} + +#[test] +fn truncate_middle_bytes_handles_utf8_content() { + let s = "😀😀😀😀😀😀😀😀😀😀\nsecond line with text\n"; + let out = truncate_middle_chars(s, 20); + assert_eq!(out, "😀😀…21 chars truncated…with text\n"); +}