From f17b3924707b756eb70999fe848afe31043a789f Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 14 Nov 2025 17:05:00 +0100 Subject: [PATCH] feat: cache tokenizer (#6609) --- codex-rs/Cargo.lock | 6 +- codex-rs/Cargo.toml | 1 + codex-rs/core/src/codex.rs | 4 ++ codex-rs/core/src/config/edit.rs | 4 ++ codex-rs/utils/cache/Cargo.toml | 2 +- codex-rs/utils/cache/src/lib.rs | 88 ++++++++++++++++++++--------- codex-rs/utils/tokenizer/Cargo.toml | 4 +- codex-rs/utils/tokenizer/src/lib.rs | 48 +++++++++++----- 8 files changed, 112 insertions(+), 45 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index d288c0f66..281f274c6 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1550,9 +1550,11 @@ name = "codex-utils-tokenizer" version = "0.0.0" dependencies = [ "anyhow", + "codex-utils-cache", "pretty_assertions", "thiserror 2.0.17", "tiktoken-rs", + "tokio", ] [[package]] @@ -6302,9 +6304,9 @@ dependencies = [ [[package]] name = "tiktoken-rs" -version = "0.7.0" +version = "0.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25563eeba904d770acf527e8b370fe9a5547bacd20ff84a0b6c3bc41288e5625" +checksum = "3a19830747d9034cd9da43a60eaa8e552dfda7712424aebf187b7a60126bae0d" dependencies = [ "anyhow", "base64", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index c50c69aa3..80bae6550 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -182,6 +182,7 @@ tempfile = "3.23.0" test-log = "0.2.18" textwrap = "0.16.2" thiserror = "2.0.17" +tiktoken-rs = "0.9" time = "0.3" tiny_http = "0.12" tokio = "1" diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 4422b3884..a13c7a0c4 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::warm_model_cache; /// The high-level interface to the Codex system. /// It operates as a queue pair where you send submissions and receive events. @@ -590,6 +591,9 @@ impl Session { // Create the mutable state for the Session. let state = SessionState::new(session_configuration.clone()); + // Warm the tokenizer cache for the session model without blocking startup. + warm_model_cache(&session_configuration.model); + let services = SessionServices { mcp_connection_manager, unified_exec_manager: UnifiedExecSessionManager::default(), diff --git a/codex-rs/core/src/config/edit.rs b/codex-rs/core/src/config/edit.rs index b4bfdf7aa..dc5ff1114 100644 --- a/codex-rs/core/src/config/edit.rs +++ b/codex-rs/core/src/config/edit.rs @@ -3,6 +3,7 @@ use crate::config::types::McpServerConfig; use crate::config::types::Notice; use anyhow::Context; use codex_protocol::config_types::ReasoningEffort; +use codex_utils_tokenizer::warm_model_cache; use std::collections::BTreeMap; use std::path::Path; use std::path::PathBuf; @@ -229,6 +230,9 @@ impl ConfigDocument { fn apply(&mut self, edit: &ConfigEdit) -> anyhow::Result { match edit { ConfigEdit::SetModel { model, effort } => Ok({ + if let Some(model) = &model { + warm_model_cache(model) + } let mut mutated = false; mutated |= self.write_profile_value( &["model"], diff --git a/codex-rs/utils/cache/Cargo.toml b/codex-rs/utils/cache/Cargo.toml index d10007153..e3397bfe1 100644 --- a/codex-rs/utils/cache/Cargo.toml +++ b/codex-rs/utils/cache/Cargo.toml @@ -9,7 +9,7 @@ workspace = true [dependencies] lru = { workspace = true } sha1 = { workspace = true } -tokio = { workspace = true, features = ["sync", "rt"] } +tokio = { workspace = true, features = ["sync", "rt", "rt-multi-thread"] } [dev-dependencies] tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread"] } diff --git a/codex-rs/utils/cache/src/lib.rs b/codex-rs/utils/cache/src/lib.rs index 38d07e52e..743c289ff 100644 --- a/codex-rs/utils/cache/src/lib.rs +++ b/codex-rs/utils/cache/src/lib.rs @@ -9,6 +9,7 @@ use tokio::sync::Mutex; use tokio::sync::MutexGuard; /// A minimal LRU cache protected by a Tokio mutex. +/// Calls outside a Tokio runtime are no-ops. pub struct BlockingLruCache { inner: Mutex>, } @@ -30,14 +31,16 @@ where where V: Clone, { - let mut guard = lock_blocking(&self.inner); - if let Some(v) = guard.get(&key) { - return v.clone(); + if let Some(mut guard) = lock_if_runtime(&self.inner) { + if let Some(v) = guard.get(&key) { + return v.clone(); + } + let v = value(); + // Insert and return a clone to keep ownership in the cache. + guard.put(key, v.clone()); + return v; } - let v = value(); - // Insert and return a clone to keep ownership in the cache. - guard.put(key, v.clone()); - v + value() } /// Like `get_or_insert_with`, but the value factory may fail. @@ -49,13 +52,15 @@ where where V: Clone, { - let mut guard = lock_blocking(&self.inner); - if let Some(v) = guard.get(&key) { - return Ok(v.clone()); + if let Some(mut guard) = lock_if_runtime(&self.inner) { + if let Some(v) = guard.get(&key) { + return Ok(v.clone()); + } + let v = value()?; + guard.put(key, v.clone()); + return Ok(v); } - let v = value()?; - guard.put(key, v.clone()); - Ok(v) + value() } /// Builds a cache if `capacity` is non-zero, returning `None` otherwise. @@ -71,12 +76,14 @@ where Q: Hash + Eq + ?Sized, V: Clone, { - lock_blocking(&self.inner).get(key).cloned() + let mut guard = lock_if_runtime(&self.inner)?; + guard.get(key).cloned() } /// Inserts `value` for `key`, returning the previous entry if it existed. pub fn insert(&self, key: K, value: V) -> Option { - lock_blocking(&self.inner).put(key, value) + let mut guard = lock_if_runtime(&self.inner)?; + guard.put(key, value) } /// Removes the entry for `key` if it exists, returning it. @@ -85,34 +92,39 @@ where K: Borrow, Q: Hash + Eq + ?Sized, { - lock_blocking(&self.inner).pop(key) + let mut guard = lock_if_runtime(&self.inner)?; + guard.pop(key) } /// Clears all entries from the cache. pub fn clear(&self) { - lock_blocking(&self.inner).clear(); + if let Some(mut guard) = lock_if_runtime(&self.inner) { + guard.clear(); + } } /// Executes `callback` with a mutable reference to the underlying cache. pub fn with_mut(&self, callback: impl FnOnce(&mut LruCache) -> R) -> R { - let mut guard = lock_blocking(&self.inner); - callback(&mut guard) + if let Some(mut guard) = lock_if_runtime(&self.inner) { + callback(&mut guard) + } else { + let mut disabled = LruCache::unbounded(); + callback(&mut disabled) + } } - /// Provides direct access to the cache guard for advanced use cases. - pub fn blocking_lock(&self) -> MutexGuard<'_, LruCache> { - lock_blocking(&self.inner) + /// Provides direct access to the cache guard when a Tokio runtime is available. + pub fn blocking_lock(&self) -> Option>> { + lock_if_runtime(&self.inner) } } -fn lock_blocking(m: &Mutex>) -> MutexGuard<'_, LruCache> +fn lock_if_runtime(m: &Mutex>) -> Option>> where K: Eq + Hash, { - match tokio::runtime::Handle::try_current() { - Ok(_) => tokio::task::block_in_place(|| m.blocking_lock()), - Err(_) => m.blocking_lock(), - } + tokio::runtime::Handle::try_current().ok()?; + Some(tokio::task::block_in_place(|| m.blocking_lock())) } /// Computes the SHA-1 digest of `bytes`. @@ -156,4 +168,26 @@ mod tests { assert_eq!(cache.get(&"a"), Some(1)); assert_eq!(cache.get(&"c"), Some(3)); } + + #[test] + fn disabled_without_runtime() { + let cache = BlockingLruCache::new(NonZeroUsize::new(2).expect("capacity")); + cache.insert("first", 1); + assert!(cache.get(&"first").is_none()); + + assert_eq!(cache.get_or_insert_with("first", || 2), 2); + assert!(cache.get(&"first").is_none()); + + assert!(cache.remove(&"first").is_none()); + cache.clear(); + + let result = cache.with_mut(|inner| { + inner.put("tmp", 3); + inner.get(&"tmp").cloned() + }); + assert_eq!(result, Some(3)); + assert!(cache.get(&"tmp").is_none()); + + assert!(cache.blocking_lock().is_none()); + } } diff --git a/codex-rs/utils/tokenizer/Cargo.toml b/codex-rs/utils/tokenizer/Cargo.toml index 6f6b4decf..7669d6b98 100644 --- a/codex-rs/utils/tokenizer/Cargo.toml +++ b/codex-rs/utils/tokenizer/Cargo.toml @@ -8,8 +8,10 @@ workspace = true [dependencies] anyhow = { workspace = true } +codex-utils-cache = { workspace = true } thiserror = { workspace = true } -tiktoken-rs = "0.7" +tiktoken-rs = { workspace = true } +tokio = { workspace = true } [dev-dependencies] pretty_assertions = { workspace = true } diff --git a/codex-rs/utils/tokenizer/src/lib.rs b/codex-rs/utils/tokenizer/src/lib.rs index 6cda6e635..1c343e439 100644 --- a/codex-rs/utils/tokenizer/src/lib.rs +++ b/codex-rs/utils/tokenizer/src/lib.rs @@ -1,7 +1,9 @@ use std::fmt; +use std::num::NonZeroUsize; +use std::sync::OnceLock; -use anyhow::Context; use anyhow::Error as AnyhowError; +use codex_utils_cache::BlockingLruCache; use thiserror::Error; use tiktoken_rs::CoreBPE; @@ -37,6 +39,26 @@ pub enum TokenizerError { }, } +fn model_cache() -> &'static BlockingLruCache { + static MODEL_CACHE: OnceLock> = OnceLock::new(); + MODEL_CACHE + .get_or_init(|| BlockingLruCache::new(NonZeroUsize::new(64).unwrap_or(NonZeroUsize::MIN))) +} + +/// Fire-and-forget function used to pre-warm model tokenizer loading. This is done +/// on a best-effort basis, without any guarantee about the state of the cache +/// before or after. +/// Only working in Tokio runtimes +pub fn warm_model_cache(model: &str) { + if tokio::runtime::Handle::try_current().is_err() { + return; + } + let model = model.to_string(); + tokio::spawn(async move { + let _ = Tokenizer::for_model(&model); + }); +} + /// Thin wrapper around a `tiktoken_rs::CoreBPE` tokenizer. #[derive(Clone)] pub struct Tokenizer { @@ -63,20 +85,13 @@ impl Tokenizer { /// Build a tokenizer using an `OpenAI` model name (maps to an encoding). /// Falls back to the `O200kBase` encoding when the model is unknown. pub fn for_model(model: &str) -> Result { - match tiktoken_rs::get_bpe_from_model(model) { - Ok(inner) => Ok(Self { inner }), - Err(model_error) => { - let inner = tiktoken_rs::o200k_base() - .with_context(|| { - format!("fallback after model lookup failure for {model}: {model_error}") - }) - .map_err(|source| TokenizerError::LoadEncoding { - kind: EncodingKind::O200kBase, - source, - })?; - Ok(Self { inner }) + let inner = model_cache().get_or_try_insert_with(model.to_owned(), || { + match tiktoken_rs::get_bpe_from_model(model) { + Ok(inner) => Ok(inner), + Err(_model_error) => Tokenizer::new(EncodingKind::O200kBase).map(|e| e.inner), } - } + })?; + Ok(Self { inner }) } /// Encode text to token IDs. If `with_special_tokens` is true, special @@ -158,4 +173,9 @@ mod tests { assert_eq!(tok.encode(text, false), fallback.encode(text, false)); Ok(()) } + + #[test] + fn warm_model_cache_without_runtime_is_noop() { + warm_model_cache("gpt-5"); + } }