agent/codex
1
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Move string truncation helpers into codex-utils-string (#15572)

Browse Source 
- 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 <noreply@openai.com>
This commit is contained in:
Ahmed Ibrahim
2026-03-24 15:45:40 -07:00
committed by GitHub
Unverified
parent 0b619afc87
commit 062fa7a2bb
36 changed files with 551 additions and 487 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+11
View File
@@ -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"
codex-rs/Cargo.toml
+2
View File
@@ -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" }
codex-rs/core/Cargo.toml
+1
View File
@@ -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 }
codex-rs/core/src/codex.rs
+2 -2
View File
@@ -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),
}
}
codex-rs/core/src/codex/rollout_reconstruction_tests.rs
+8 -8
View File
@@ -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
codex-rs/core/src/codex_tests.rs
+1 -1
View File
@@ -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
codex-rs/core/src/compact.rs
+3 -3
View File
@@ -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;
codex-rs/core/src/context_manager/history.rs
+6 -6
View File
@@ -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;
codex-rs/core/src/context_manager/history_tests.rs
+3 -3
View File
@@ -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 {
codex-rs/core/src/error.rs
+2 -2
View File
@@ -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;
codex-rs/core/src/guardian/prompt.rs
+3 -3
View File
@@ -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;
codex-rs/core/src/lib.rs
-1
View File
@@ -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;
codex-rs/core/src/memories/prompts.rs
+2 -2
View File
@@ -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;
codex-rs/core/src/models_manager/model_info.rs
+1 -1
View File
@@ -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");
codex-rs/core/src/realtime_context.rs
+2 -2
View File
@@ -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;
codex-rs/core/src/rollout/list.rs
+1 -1
View File
@@ -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;
codex-rs/core/src/rollout/recorder.rs
+4 -4
View File
@@ -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.
codex-rs/core/src/rollout/recorder_tests.rs
+1 -1
View File
@@ -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;
codex-rs/core/src/rollout/tests.rs
+1 -1
View File
@@ -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;
codex-rs/core/src/state/session.rs
+1 -1
View File
@@ -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 {
codex-rs/core/src/tools/code_mode/mod.rs
+3 -3
View File
@@ -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;
codex-rs/core/src/tools/context.rs
+2 -2
View File
@@ -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;
codex-rs/core/src/tools/js_repl/mod.rs
+2 -2
View File
@@ -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");
codex-rs/core/src/tools/mod.rs
+3 -3
View File
@@ -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;
codex-rs/core/src/truncate.rs
-363
View File
@@ -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<TruncationPolicy> for ProtocolTruncationPolicy {
fn from(value: TruncationPolicy) -> Self {
match value {
TruncationPolicy::Bytes(bytes) => Self::Bytes(bytes),
TruncationPolicy::Tokens(tokens) => Self::Tokens(tokens),
}
}
}
impl From<TruncationPolicyConfig> 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<f64> 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<FunctionCallOutputContentItem>, Option<usize>) {
let text_segments = items
.iter()
.filter_map(|item| match item {
FunctionCallOutputContentItem::InputText { text } => Some(text.as_str()),
FunctionCallOutputContentItem::InputImage { .. } => None,
})
.collect::<Vec<_>>();
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<FunctionCallOutputContentItem> {
let mut out: Vec<FunctionCallOutputContentItem> = 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<u64>) {
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;
codex-rs/core/src/unified_exec/process.rs
+2 -2
View File
@@ -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;
codex-rs/core/src/unified_exec/process_manager.rs
+1 -1
View File
@@ -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"),
codex-rs/protocol/Cargo.toml
+1
View File
@@ -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"] }
codex-rs/protocol/src/protocol.rs
+46
View File
@@ -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<crate::openai_models::TruncationPolicyConfig> 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<f64> 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,
codex-rs/utils/output-truncation/BUILD.bazel
+6
View File
@@ -0,0 +1,6 @@
load("//:defs.bzl", "codex_rust_crate")
codex_rust_crate(
name = "output-truncation",
crate_name = "codex_utils_output_truncation",
)
codex-rs/utils/output-truncation/Cargo.toml
+15
View File
@@ -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 }
codex-rs/utils/output-truncation/src/lib.rs
+142
View File
@@ -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<FunctionCallOutputContentItem>, Option<usize>) {
let text_segments = items
.iter()
.filter_map(|item| match item {
FunctionCallOutputContentItem::InputText { text } => Some(text.as_str()),
FunctionCallOutputContentItem::InputImage { .. } => None,
})
.collect::<Vec<_>>();
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<FunctionCallOutputContentItem> {
let mut out: Vec<FunctionCallOutputContentItem> = 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;
codex-rs/core/src/truncate_tests.rs β†’ codex-rs/utils/output-truncation/src/truncate_tests.rs
+37 -69
View File
@@ -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);
}
codex-rs/utils/string/src/lib.rs
+9
View File
@@ -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;
codex-rs/utils/string/src/truncate.rs
+156
View File
@@ -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<u64>) {
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;
codex-rs/utils/string/src/truncate/tests.rs
+71
View File
@@ -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");
}