From 109f00da0a338b309f6751f4166bff1267ebd3bb Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Thu, 13 Nov 2025 16:59:00 -0800 Subject: [PATCH] start --- codex-rs/core/src/codex.rs | 10 +++++++- codex-rs/core/src/context_manager/history.rs | 5 ++-- codex-rs/core/src/truncate.rs | 4 ++-- codex-rs/utils/tokenizer/src/lib.rs | 25 ++++++++++++++++++++ 4 files changed, 38 insertions(+), 6 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 273f7a4975..dfe4f899de 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -133,6 +133,7 @@ use codex_protocol::protocol::InitialHistory; use codex_protocol::user_input::UserInput; use codex_utils_readiness::Readiness; use codex_utils_readiness::ReadinessFlag; +use codex_utils_tokenizer::shared_default_tokenizer; /// The high-level interface to the Codex system. /// It operates as a queue pair where you send submissions and receive events. @@ -247,6 +248,10 @@ impl Codex { } } +async fn warm_up_tokenizer() { + let _ = tokio::task::spawn_blocking(shared_default_tokenizer).await; +} + /// Context for an initialized model agent /// /// A session has at most 1 running task at a time, and can be interrupted by user input. @@ -484,6 +489,7 @@ impl Session { // - spin up MCP connection manager // - perform default shell discovery // - load history metadata + // - warm up the shared tokenizer let rollout_fut = RolloutRecorder::new(&config, rollout_params); let mcp_fut = McpConnectionManager::new( @@ -496,6 +502,7 @@ impl Session { config.mcp_servers.iter(), config.mcp_oauth_credentials_store_mode, ); + let tokenizer_warmup_fut = warm_up_tokenizer(); // Join all independent futures. let ( @@ -509,7 +516,8 @@ impl Session { mcp_fut, default_shell_fut, history_meta_fut, - auth_statuses_fut + auth_statuses_fut, + tokenizer_warmup_fut ); let rollout_recorder = rollout_recorder.map_err(|e| { diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 07ff54ad1c..729e3d48a4 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -6,7 +6,7 @@ use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; -use codex_utils_tokenizer::Tokenizer; +use codex_utils_tokenizer::shared_default_tokenizer; use std::ops::Deref; /// Transcript of conversation history @@ -78,8 +78,7 @@ impl ContextManager { // /!\ The value is a lower bound estimate and does not represent the exact // context length. pub(crate) fn estimate_token_count(&self, turn_context: &TurnContext) -> Option { - let model = turn_context.client.get_model(); - let tokenizer = Tokenizer::for_model(model.as_str()).ok()?; + let tokenizer = shared_default_tokenizer()?; let model_family = turn_context.client.get_model_family(); Some( diff --git a/codex-rs/core/src/truncate.rs b/codex-rs/core/src/truncate.rs index 3f0be8fcaf..b53fed90d8 100644 --- a/codex-rs/core/src/truncate.rs +++ b/codex-rs/core/src/truncate.rs @@ -1,7 +1,7 @@ //! Utilities for truncating large chunks of output while preserving a prefix //! and suffix on UTF-8 boundaries. -use codex_utils_tokenizer::Tokenizer; +use codex_utils_tokenizer::shared_default_tokenizer; /// Truncate the middle of a UTF-8 string to at most `max_bytes` bytes, /// preserving the beginning and the end. Returns the possibly truncated @@ -15,7 +15,7 @@ pub(crate) fn truncate_middle(s: &str, max_bytes: usize) -> (String, Option // Build a tokenizer for counting (default to o200k_base; fall back to cl100k_base). // If both fail, fall back to a 4-bytes-per-token estimate. - let tok = Tokenizer::try_default().ok(); + let tok = shared_default_tokenizer(); let token_count = |text: &str| -> u64 { if let Some(ref t) = tok { t.count(text) as u64 diff --git a/codex-rs/utils/tokenizer/src/lib.rs b/codex-rs/utils/tokenizer/src/lib.rs index 6cda6e635f..bc9bf0146f 100644 --- a/codex-rs/utils/tokenizer/src/lib.rs +++ b/codex-rs/utils/tokenizer/src/lib.rs @@ -1,4 +1,6 @@ use std::fmt; +use std::sync::Arc; +use std::sync::OnceLock; use anyhow::Context; use anyhow::Error as AnyhowError; @@ -107,6 +109,19 @@ impl Tokenizer { } } +static DEFAULT_TOKENIZER: OnceLock, TokenizerError>> = OnceLock::new(); + +/// Return a shared default tokenizer (`O200kBase`), loading it once per process. +/// Returns `None` if initialization fails. +#[must_use] +pub fn shared_default_tokenizer() -> Option> { + DEFAULT_TOKENIZER + .get_or_init(|| Tokenizer::try_default().map(Arc::new)) + .as_ref() + .ok() + .cloned() +} + #[cfg(test)] mod tests { use super::*; @@ -158,4 +173,14 @@ mod tests { assert_eq!(tok.encode(text, false), fallback.encode(text, false)); Ok(()) } + + #[test] + fn shared_default_tokenizer_is_cached() { + let first = shared_default_tokenizer().expect("default tokenizer"); + let second = shared_default_tokenizer().expect("default tokenizer reused"); + + let ptr1 = Arc::as_ptr(&first); + let ptr2 = Arc::as_ptr(&second); + assert_eq!(ptr1, ptr2); + } }