From 257d5eceff1079f1c8f6c593611325873b7ced96 Mon Sep 17 00:00:00 2001
From: Ahmed Ibrahim <aibrahim@openai.com>
Date: Fri, 24 Oct 2025 12:31:08 -0700
Subject: [PATCH] use shared tokenizer

---
 codex-rs/core/src/codex.rs                |  3 ++
 codex-rs/core/src/conversation_history.rs | 48 +++++++++++++++--------
 2 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs
index 1a3be60534..ec2caa8a67 100644
--- a/codex-rs/core/src/codex.rs
+++ b/codex-rs/core/src/codex.rs
@@ -59,6 +59,7 @@ use crate::config::Config;
 use crate::config_types::McpServerTransportConfig;
 use crate::config_types::ShellEnvironmentPolicy;
 use crate::conversation_history::ConversationHistory;
+use crate::conversation_history::prefetch_tokenizer_in_background;
 use crate::environment_context::EnvironmentContext;
 use crate::error::CodexErr;
 use crate::error::Result as CodexResult;
@@ -159,6 +160,8 @@ impl Codex {
         conversation_history: InitialHistory,
         session_source: SessionSource,
     ) -> CodexResult<CodexSpawnOk> {
+        // Start loading the tokenizer in the background so we don't block later.
+        prefetch_tokenizer_in_background();
         let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY);
         let (tx_event, rx_event) = async_channel::unbounded();
 
diff --git a/codex-rs/core/src/conversation_history.rs b/codex-rs/core/src/conversation_history.rs
index 4ffe03bdfc..4c91a3278e 100644
--- a/codex-rs/core/src/conversation_history.rs
+++ b/codex-rs/core/src/conversation_history.rs
@@ -6,8 +6,11 @@ use codex_protocol::models::ResponseItem;
 use codex_protocol::protocol::TokenUsage;
 use codex_protocol::protocol::TokenUsageInfo;
 use codex_utils_tokenizer::Tokenizer;
+use tokio::task;
 use tracing::error;
 
+static TOKENIZER: OnceLock<Option<Arc<Tokenizer>>> = OnceLock::new();
+
 use crate::error::CodexErr;
 
 /// Transcript of conversation history
@@ -16,17 +19,13 @@ pub(crate) struct ConversationHistory {
     /// The oldest items are at the beginning of the vector.
     items: Vec<ResponseItem>,
     token_info: Option<TokenUsageInfo>,
-    tokenizer: Option<Arc<Tokenizer>>,
 }
 
 impl ConversationHistory {
     pub(crate) fn new() -> Self {
-        // load the tokenizer once and share it across all instances. It takes ~500 ms to load.
-        let tokenizer = shared_tokenizer();
         Self {
             items: Vec::new(),
             token_info: None,
-            tokenizer,
         }
     }
 
@@ -121,8 +120,9 @@ impl ConversationHistory {
         let Some(context_window) = info.model_context_window else {
             return Ok(());
         };
-        let Some(tokenizer) = self.tokenizer.as_deref() else {
-            return Ok(());
+        let tokenizer = match shared_tokenizer() {
+            Some(t) => t,
+            None => return Ok(()),
         };
 
         let mut input_tokens: i64 = 0;
@@ -413,17 +413,33 @@ fn error_or_panic(message: String) {
 }
 
 fn shared_tokenizer() -> Option<Arc<Tokenizer>> {
-    static TOKENIZER: OnceLock<Option<Arc<Tokenizer>>> = OnceLock::new();
-    TOKENIZER
-        .get_or_init(|| match Tokenizer::try_default() {
-            Ok(tokenizer) => Some(Arc::new(tokenizer)),
-            Err(error) => {
-                error!("failed to create tokenizer: {error}");
-                None
+    TOKENIZER.get().and_then(|opt| opt.as_ref().map(Arc::clone))
+}
+
+/// Kick off background initialization of the shared tokenizer without blocking the caller.
+pub(crate) fn prefetch_tokenizer_in_background() {
+    if TOKENIZER.get().is_some() {
+        return;
+    }
+
+    // Spawn a background task to initialize the tokenizer. Use spawn_blocking in case
+    // initialization performs CPU-heavy work or file I/O.
+    tokio::spawn(async {
+        let result = task::spawn_blocking(Tokenizer::try_default).await;
+        match result {
+            Ok(Ok(tokenizer)) => {
+                let _ = TOKENIZER.set(Some(Arc::new(tokenizer)));
             }
-        })
-        .as_ref()
-        .map(Arc::clone)
+            Ok(Err(error)) => {
+                error!("failed to create tokenizer: {error}");
+                let _ = TOKENIZER.set(None);
+            }
+            Err(join_error) => {
+                error!("failed to join tokenizer init task: {join_error}");
+                let _ = TOKENIZER.set(None);
+            }
+        }
+    });
 }
 
 /// Anything that is not a system message or "reasoning" message is considered