Fix PromptTrie.pop_prefixes() off-by-one when pruning immediate prefixes (#1078)

Signed-off-by: Yuan Lik Xun <lxyuan0420@gmail.com>
Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

mlx_lm/models/cache.py

@@ -1434,18 +1434,22 @@ class PromptTrie:
     def pop_prefixes(self, model: Any, tokens: List[int]):
         values = []
         current = self._trie[model]
-        for i in range(len(tokens) - 1):
+        for i, tok in enumerate(tokens):
             if "__value__" in current:
                 values.append((i, current.pop("__value__")))
-            current = current[tokens[i]]
+            current = current[tok]
         return values
 
     def search(self, model: Any, tokens: List[int]) -> PromptTrieResult:
         if model not in self._trie:
             return PromptTrieResult(model, None, None, None, 0)
 
-        # Walk the tokens as far as we can
         current = self._trie[model]
+
+        if not tokens and "__value__" in current:
+            return PromptTrieResult(model, [], None, None, 0)
+
+        # Walk the tokens as far as we can
         last_index = -1
         index = 0
         while index < len(tokens) and tokens[index] in current:
@@ -1455,7 +1459,7 @@ class PromptTrie:
             index += 1
 
         # Got an exact match
-        if last_index == len(tokens) - 1:
+        if last_index == len(tokens) - 1 >= 0:
             return PromptTrieResult(model, tokens, None, None, 0)
 
         # Check if we found a prefix at any point

tests/test_server.py

@@ -537,6 +537,41 @@ class TestLRUPromptCache(unittest.TestCase):
         self.assertEqual(c, [MockCache("test4")])
         self.assertEqual(t, [])
 
+    def test_insert_trimmable_cache_removes_immediate_prefix(self):
+        cache = LRUPromptCache(max_size=10)
+        model = ("test", None, None)
+
+        cache.insert_cache(model, [1, 2], [MockCache("ab")])
+        self.assertEqual(len(cache), 1)
+        self.assertEqual(cache.nbytes, 2)
+
+        cache.insert_cache(model, [1, 2, 3], [MockCache("abc")])
+        self.assertEqual(len(cache), 1)
+        self.assertEqual(cache.nbytes, 3)
+
+    def test_insert_empty_tokens_does_not_self_destruct(self):
+        cache = LRUPromptCache(max_size=10)
+        model = ("test", None, None)
+
+        cache.insert_cache(model, [], [MockCache("root")])
+        self.assertEqual(len(cache), 1)
+        self.assertEqual(cache.nbytes, 4)
+
+        c, t = cache.fetch_nearest_cache(model, [])
+        self.assertIsNotNone(c)
+        self.assertEqual(t, [])
+
+    def test_fetch_empty_tokens_after_root_eviction(self):
+        cache = LRUPromptCache(max_size=10)
+        model = ("test", None, None)
+
+        cache.insert_cache(model, [], [MockCache("root")])
+        cache.insert_cache(model, [1], [MockCache("a")])
+
+        c, t = cache.fetch_nearest_cache(model, [])
+        self.assertIsNone(c)
+        self.assertEqual(t, [])
+
     def test_lru_bytes(self):
         cache = LRUPromptCache(max_size=100, max_bytes=10)
         model = ("test", None, None)