merge: token-count memo rejects cross-tokenizer counts under concurrent rebind (headroomlabs-ai#2710)

Author: nangsontay

headroom/cache/token_count_memo.py

@@ -46,6 +46,17 @@ class TokenCountMemo:
     to a fresh ``EstimatingTokenCounter`` mid-session — clears the memo and
     rebinds. The registry caches tokenizer instances per model, so on the
     normal path the binding is stable across a session's requests.
+
+    The store is per-session and a session's requests run concurrently, so
+    binding alone is not enough: request A can bind tokenizer A, release the
+    lock, and still be mid-count when request B rebinds to tokenizer B (the
+    fail-open path builds a *fresh* estimator per request — see
+    ``proxy/token_counting._count_offloaded``). ``get``/``put`` therefore
+    re-check tokenizer identity under the same lock that owns the binding, so
+    an obsolete binding's reads miss and its writes are dropped rather than
+    leaking counts across tokenizers. That makes each lookup/populate atomic
+    with respect to tokenizer identity without holding the lock across the
+    (slow) tokenizer call itself.
     """
 
     def __init__(self, max_entries: int = 10000) -> None:
@@ -83,8 +94,16 @@ def message_hash(message: dict[str, Any]) -> str:
         raw = TokenCountMemo.canonical_message(message)
         return hashlib.sha256(raw.encode("utf-8")).hexdigest()
 
-    def get(self, key: str, canonical: str | None = None) -> int | None:
+    def get(self, key: str, canonical: str | None = None, *, tokenizer: object) -> int | None:
+        """Read a count, but only for the tokenizer currently bound.
+
+        ``tokenizer`` is required so no caller can silently opt out of the
+        identity check: a stale binding must miss, not serve another
+        tokenizer's numbers.
+        """
         with self._lock:
+            if self._bound_tokenizer is not tokenizer:
+                return None
             entries = self._counts.get(key)
             if entries is None:
                 return None
@@ -93,8 +112,16 @@ def get(self, key: str, canonical: str | None = None) -> int | None:
                 return entries[0][1]
             return next((count for stored, count in entries if stored == canonical), None)
 
-    def put(self, key: str, count: int, canonical: str | None = None) -> None:
+    def put(self, key: str, count: int, canonical: str | None = None, *, tokenizer: object) -> None:
+        """Store a count, dropping writes from an obsolete tokenizer binding.
+
+        A concurrent rebind (another request's tokenizer swap) already cleared
+        the store for its own tokenizer; letting this write land would mark a
+        stale count as belonging to the new binding.
+        """
         with self._lock:
+            if self._bound_tokenizer is not tokenizer:
+                return
             canonical = key if canonical is None else canonical
             entries = self._counts.pop(key, [])
             entries = [(stored, value) for stored, value in entries if stored != canonical]
@@ -121,6 +148,11 @@ def count_messages_memoized(
     ``count_messages`` call instead. The memo is also bound to the tokenizer
     instance — a swap (count fail-open downgrade) clears stale counts rather
     than serving numbers computed under a different tokenizer.
+
+    Every ``get``/``put`` re-asserts that identity under the memo lock, so a
+    concurrent rebind on the shared per-session memo degrades this call to
+    uncached counting (all values still produced by ``tokenizer``) instead of
+    trading counts with the other request's tokenizer.
     """
     if not getattr(tokenizer, "ADDITIVE_COUNTS", False):
         return tokenizer.count_messages(messages)
@@ -129,9 +161,9 @@ def count_messages_memoized(
     for msg in messages:
         canonical = TokenCountMemo.canonical_message(msg)
         key = TokenCountMemo.message_hash(msg)
-        count = memo.get(key, canonical)
+        count = memo.get(key, canonical, tokenizer=tokenizer)
         if count is None:
             count = tokenizer.count_message(msg)
-            memo.put(key, count, canonical)
+            memo.put(key, count, canonical, tokenizer=tokenizer)
         total += count
     return total + tokenizer.REPLY_OVERHEAD

tests/test_token_count_memo.py

@@ -11,6 +11,8 @@
 
 from __future__ import annotations
 
+import threading
+
 import pytest
 
 from headroom.cache.token_count_memo import TokenCountMemo, count_messages_memoized
@@ -133,7 +135,7 @@ def test_memoized_count_handles_append_only_delta() -> None:
 
     count_messages_memoized(memo, tokenizer, turn1)
     prefix_key = TokenCountMemo.message_hash(turn1[0])
-    assert memo.get(prefix_key) is not None
+    assert memo.get(prefix_key, tokenizer=tokenizer) is not None
 
     memoized_turn2 = count_messages_memoized(memo, tokenizer, turn2)
     assert memoized_turn2 == tokenizer.count_messages(turn2)
@@ -174,21 +176,25 @@ def count_messages(self, messages) -> int:  # noqa: ANN001
 class TestTokenCountMemoEviction:
     def test_eviction_at_max_entries(self) -> None:
         memo = TokenCountMemo(max_entries=3)
-        memo.put("a", 1)
-        memo.put("b", 2)
-        memo.put("c", 3)
-        memo.get("a")  # touch "a" so it's not the least-recently-used
-        memo.put("d", 4)  # should evict "b" (oldest untouched)
-
-        assert memo.get("a") == 1
-        assert memo.get("b") is None
-        assert memo.get("c") == 3
-        assert memo.get("d") == 4
+        tok = object()  # get/put require the bound tokenizer's identity
+        memo.bind_or_reset(tok)
+        memo.put("a", 1, tokenizer=tok)
+        memo.put("b", 2, tokenizer=tok)
+        memo.put("c", 3, tokenizer=tok)
+        memo.get("a", tokenizer=tok)  # touch "a" so it's not the least-recently-used
+        memo.put("d", 4, tokenizer=tok)  # should evict "b" (oldest untouched)
+
+        assert memo.get("a", tokenizer=tok) == 1
+        assert memo.get("b", tokenizer=tok) is None
+        assert memo.get("c", tokenizer=tok) == 3
+        assert memo.get("d", tokenizer=tok) == 4
 
     def test_get_stats_reports_entry_count(self) -> None:
         memo = TokenCountMemo()
-        memo.put("a", 1)
-        memo.put("b", 2)
+        tok = object()
+        memo.bind_or_reset(tok)
+        memo.put("a", 1, tokenizer=tok)
+        memo.put("b", 2, tokenizer=tok)
         assert memo.get_stats()["entries"] == 2
 
 
@@ -258,3 +264,94 @@ def test_same_tokenizer_keeps_memo_bound(self) -> None:
         entries = memo.get_stats()["entries"]
         count_messages_memoized(memo, tokenizer, messages)
         assert memo.get_stats()["entries"] == entries  # no clear, warm hits
+
+
+class _GatedTokenizer:
+    """Additive tokenizer with a fixed per-message count and an interleaving hook.
+
+    ``per_message`` differs between the two instances so a cross-tokenizer leak
+    shows up as a wrong number rather than a coincidence.
+    """
+
+    ADDITIVE_COUNTS = True
+    REPLY_OVERHEAD = 0
+
+    def __init__(self, per_message: int, before_count: object = None) -> None:
+        self.per_message = per_message
+        self._before_count = before_count
+
+    def count_message(self, message: dict) -> int:  # noqa: ARG002
+        if self._before_count is not None:
+            self._before_count()
+        return self.per_message
+
+    def count_messages(self, messages: list[dict]) -> int:
+        return sum(self.count_message(m) for m in messages) + self.REPLY_OVERHEAD
+
+
+class TestCrossTokenizerRebindRace:
+    """The per-session memo is shared by concurrent requests, and the count
+    fail-open path (``proxy/token_counting._count_offloaded``) hands each
+    request a *fresh* ``EstimatingTokenCounter``. So request A can be mid-count
+    under tokenizer A while request B rebinds the memo to tokenizer B. A's
+    counts must never reach B.
+    """
+
+    def test_put_after_concurrent_rebind_never_leaks_to_new_binding(self) -> None:
+        messages = [{"role": "user", "content": "shared prefix message"}]
+        key = TokenCountMemo.message_hash(messages[0])
+        canonical = TokenCountMemo.canonical_message(messages[0])
+
+        a_is_counting = threading.Event()
+        b_has_rebound = threading.Event()
+
+        def _pause_a() -> None:
+            # A has already called bind_or_reset and missed the cache; hold it
+            # here so B's whole rebind+count lands before A's put.
+            a_is_counting.set()
+            assert b_has_rebound.wait(timeout=10), "B never rebound"
+
+        a = _GatedTokenizer(per_message=7, before_count=_pause_a)
+        b = _GatedTokenizer(per_message=31)
+        memo = TokenCountMemo()
+
+        a_total: list[int] = []
+        a_thread = threading.Thread(
+            target=lambda: a_total.append(count_messages_memoized(memo, a, messages)),
+            daemon=True,
+        )
+        a_thread.start()
+        assert a_is_counting.wait(timeout=10), "A never started counting"
+
+        # B rebinds (clearing A's era) and populates its own count.
+        b_total = count_messages_memoized(memo, b, messages)
+        b_has_rebound.set()
+        a_thread.join(timeout=10)
+        assert not a_thread.is_alive()
+
+        # A's own total stays exact under tokenizer A — the fix degrades A to
+        # uncached counting, it does not hand A B's numbers.
+        assert a_total == [7]
+        assert b_total == 31
+
+        # The proof: A's post-rebind put was dropped, so the memo still holds
+        # only B's count for the shared message and a fresh B read agrees.
+        assert memo.get(key, canonical, tokenizer=b) == 31
+        assert count_messages_memoized(memo, b, messages) == 31
+
+    def test_get_from_obsolete_binding_misses_instead_of_reading_new_counts(self) -> None:
+        """Mirror direction: a stale binding must not read the rebinder's
+        counts either (the swap guarantee stays symmetric)."""
+        messages = [{"role": "user", "content": "shared prefix message"}]
+        key = TokenCountMemo.message_hash(messages[0])
+        canonical = TokenCountMemo.canonical_message(messages[0])
+
+        a = _GatedTokenizer(per_message=7)
+        b = _GatedTokenizer(per_message=31)
+        memo = TokenCountMemo()
+
+        count_messages_memoized(memo, a, messages)
+        count_messages_memoized(memo, b, messages)  # rebinds to b
+
+        assert memo.get(key, canonical, tokenizer=a) is None
+        assert memo.get(key, canonical, tokenizer=b) == 31