Use LRU cache for token counting (#9586)

Signed-off-by: jh-block <jhugo@block.xyz>

Cargo.lock

@@ -2668,20 +2668,6 @@ dependencies = [
  "parking_lot_core",
 ]
 
-[[package]]
-name = "dashmap"
-version = "6.2.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c"
-dependencies = [
- "cfg-if",
- "crossbeam-utils",
- "hashbrown 0.14.5",
- "lock_api",
- "once_cell",
- "parking_lot_core",
-]
-
 [[package]]
 name = "data-encoding"
 version = "2.11.0"
@@ -4449,7 +4435,6 @@ version = "1.37.0"
 dependencies = [
  "agent-client-protocol",
  "agent-client-protocol-schema",
- "ahash",
  "anyhow",
  "arboard",
  "async-stream",
@@ -4469,7 +4454,6 @@ dependencies = [
  "chrono",
  "clap",
  "ctor",
- "dashmap 6.2.1",
  "dirs 5.0.1",
  "dotenvy",
  "dtor",
@@ -6313,7 +6297,7 @@ dependencies = [
  "async-trait",
  "boxed_error",
  "capacity_builder",
- "dashmap 5.5.3",
+ "dashmap",
  "deno_error",
  "deno_maybe_sync",
  "deno_media_type",

crates/goose/Cargo.toml

@@ -152,8 +152,6 @@ blake3 = { version = "1", default-features = false, features = ["std"] }
 fs2 = { workspace = true }
 tokio-stream = { workspace = true, features = ["io-util"] }
 tempfile = { workspace = true }
-dashmap = { version = "6", default-features = false }
-ahash = { version = "0.8.11", default-features = false, features = ["std"] }
 tokio-util = { workspace = true, features = ["compat"] }
 agent-client-protocol-schema = { workspace = true }
 agent-client-protocol = { workspace = true, features = ["unstable"] }

crates/goose/src/token_counter.rs

@@ -1,8 +1,7 @@
-use ahash::AHasher;
-use dashmap::DashMap;
+use lru::LruCache;
 use rmcp::model::Tool;
-use std::hash::{Hash, Hasher};
-use std::sync::Arc;
+use std::num::NonZeroUsize;
+use std::sync::{Arc, Mutex};
 use tiktoken_rs::CoreBPE;
 use tokio::sync::OnceCell;
 
@@ -10,7 +9,7 @@ use crate::conversation::message::Message;
 
 static TOKENIZER: OnceCell<Arc<CoreBPE>> = OnceCell::const_new();
 
-const MAX_TOKEN_CACHE_SIZE: usize = 10_000;
+const MAX_TOKEN_CACHE_SIZE: usize = 1_024;
 
 // token use for various bits of a tool calls:
 const FUNC_INIT: usize = 7;
@@ -22,38 +21,54 @@ const FUNC_END: usize = 12;
 
 pub struct TokenCounter {
     tokenizer: Arc<CoreBPE>,
-    token_cache: Arc<DashMap<u64, usize>>,
+    token_cache: Mutex<LruCache<TokenCacheKey, usize>>,
+}
+
+#[derive(Clone, Copy, Eq, Hash, PartialEq)]
+struct TokenCacheKey {
+    len: usize,
+    hash: [u8; 32],
+}
+
+impl TokenCacheKey {
+    fn from_text(text: &str) -> Self {
+        Self {
+            len: text.len(),
+            hash: *blake3::hash(text.as_bytes()).as_bytes(),
+        }
+    }
 }
 
 impl TokenCounter {
     pub async fn new() -> Result<Self, String> {
         let tokenizer = get_tokenizer().await?;
+        let cache_capacity =
+            NonZeroUsize::new(MAX_TOKEN_CACHE_SIZE).expect("token cache capacity must be non-zero");
         Ok(Self {
             tokenizer,
-            token_cache: Arc::new(DashMap::new()),
+            token_cache: Mutex::new(LruCache::new(cache_capacity)),
         })
     }
 
     pub fn count_tokens(&self, text: &str) -> usize {
-        let mut hasher = AHasher::default();
-        text.hash(&mut hasher);
-        let hash = hasher.finish();
-
-        if let Some(count) = self.token_cache.get(&hash) {
-            return *count;
+        let cache_key = TokenCacheKey::from_text(text);
+        if let Some(count) = self
+            .token_cache
+            .lock()
+            .expect("token cache mutex poisoned")
+            .get(&cache_key)
+            .copied()
+        {
+            return count;
         }
 
         let tokens = self.tokenizer.encode_with_special_tokens(text);
         let count = tokens.len();
 
-        if self.token_cache.len() >= MAX_TOKEN_CACHE_SIZE {
-            if let Some(entry) = self.token_cache.iter().next() {
-                let old_hash = *entry.key();
-                self.token_cache.remove(&old_hash);
-            }
-        }
-
-        self.token_cache.insert(hash, count);
+        self.token_cache
+            .lock()
+            .expect("token cache mutex poisoned")
+            .put(cache_key, count);
         count
     }
 
@@ -173,11 +188,17 @@ impl TokenCounter {
     }
 
     pub fn clear_cache(&self) {
-        self.token_cache.clear();
+        self.token_cache
+            .lock()
+            .expect("token cache mutex poisoned")
+            .clear();
     }
 
     pub fn cache_size(&self) -> usize {
-        self.token_cache.len()
+        self.token_cache
+            .lock()
+            .expect("token cache mutex poisoned")
+            .len()
     }
 }
 
@@ -260,13 +281,13 @@ mod tests {
         let counter = create_token_counter().await.unwrap();
 
         let mut cached_texts = Vec::new();
-        for i in 0..50 {
+        for i in 0..=MAX_TOKEN_CACHE_SIZE {
             let text = format!("Test string number {}", i);
             counter.count_tokens(&text);
             cached_texts.push(text);
         }
 
-        assert!(counter.cache_size() <= MAX_TOKEN_CACHE_SIZE);
+        assert_eq!(counter.cache_size(), MAX_TOKEN_CACHE_SIZE);
 
         let recent_text = &cached_texts[cached_texts.len() - 1];
         let start_size = counter.cache_size();