feat: cache tokenizer (#6609)

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",

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"

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(),

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<bool> {
         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"],

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"] }

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<K, V> {
     inner: Mutex<LruCache<K, V>>,
 }
@@ -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<V> {
-        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>,
         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<R>(&self, callback: impl FnOnce(&mut LruCache<K, V>) -> 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<K, V>> {
-        lock_blocking(&self.inner)
+    /// Provides direct access to the cache guard when a Tokio runtime is available.
+    pub fn blocking_lock(&self) -> Option<MutexGuard<'_, LruCache<K, V>>> {
+        lock_if_runtime(&self.inner)
     }
 }
 
-fn lock_blocking<K, V>(m: &Mutex<LruCache<K, V>>) -> MutexGuard<'_, LruCache<K, V>>
+fn lock_if_runtime<K, V>(m: &Mutex<LruCache<K, V>>) -> Option<MutexGuard<'_, LruCache<K, V>>>
 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());
+    }
 }

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 }

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<String, CoreBPE> {
+    static MODEL_CACHE: OnceLock<BlockingLruCache<String, CoreBPE>> = 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<Self, TokenizerError> {
-        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");
+    }
 }