fix(tokenizers): count HuggingFace chat templates, and resolve gpt-5 / gateway-wrapped names

Three selection defects, all measured against real counters on identical text.

1. HuggingFace-routed models counted a whole conversation as 2 tokens.

transformers >= 5 defaults `apply_chat_template(tokenize=True)` to
`return_dict=True` and returns a BatchEncoding, so `len(formatted)` counted DICT
KEYS — input_ids and attention_mask — instead of tokens. Measured on Qwen2.5-72B
with a 6,000-char message: count_messages returned 2 and count_message returned
-1 (BaseTokenizer subtracts a 3-token reply overhead from it), against a true
1,003. After the fix: 1,020 and 1,017.

That is a ~99.8% undercount on every HF-routed family whose resolved tokenizer
carries a chat template — llama, qwen, deepseek, phi, yi, falcon, starcoder.
pyproject pins transformers>=5.5.0,<6.0, so the affected version is the only
installable one, and no test covered count_messages.

It hid behind a second bug while I was reproducing it: DeepSeek-V3 mis-resolves to
deepseek-llm-7b-base, a 2023 model with no chat template, which falls back to the
estimator and looks correct. That mis-resolution is left for a follow-up.

2. The current OpenAI flagships had no pattern.

MODEL_PATTERNS stopped at ^gpt-4 / ^o1 / ^o3, so gpt-5, gpt-5.1, gpt-5-mini,
gpt-5.1-codex and o4-mini all fell through to the char estimator. Deviation vs the
correct o200k encoding: +20% English, -33% JSON, -44% logs. Added ^gpt-5 and ^o4.

3. Every pattern is ^-anchored, so gateway-wrapped ids matched nothing.

bedrock/anthropic.claude-*, vertex_ai/claude-*, openrouter/anthropic/claude-*,
anthropic/claude-*, azure/gpt-4o, us.anthropic.claude-*-v1:0 and friends all
resolved to the estimator instead of their family's counter. Deviation: +15%
English, -33% JSON, -38% logs. handlers/openai.py already documents that
LiteLLM's `headroom` guardrail passes exactly these forms.

`_detect_backend` now tries progressively-unwrapped candidates — path segments
stripped left to right, then Bedrock's dotted [region.]vendor.model — with the
FULL name first, so no currently-correct resolution can move and an unknown alias
still falls back to estimation rather than matching by accident.

20 new tests covering all three, plus the no-regression cases: bare names
unchanged, unknown aliases still estimated, wrapped Gemini matching its bare form
exactly, and candidate ordering.

ruff check + format clean (0.15.17). tests/test_huggingface_tokenizer_timeout.py
+ tests/test_tokenizers/: 12 passed on this branch and 12 on clean upstream/main.
tests/test_evals_cjk_tokenization.py cannot collect in this env
(ModuleNotFoundError: headroom._core, the compiled extension this machine cannot
build) — identical on baseline, so CI is the check there.

headroom/tokenizers/huggingface.py

@@ -325,11 +325,22 @@ class HuggingFaceTokenizer(BaseTokenizer):
         # Try to use chat template for accurate counting
         if hasattr(self.tokenizer, "apply_chat_template"):
             try:
-                # Apply chat template and count
+                # ``return_dict=False`` is load-bearing. transformers >= 5 defaults
+                # ``apply_chat_template(tokenize=True)`` to ``return_dict=True``,
+                # which hands back a BatchEncoding — so ``len(formatted)`` counted
+                # DICT KEYS (2: input_ids, attention_mask) instead of tokens.
+                # Measured on Qwen2.5-72B, a 6,000-char message: count_messages
+                # returned 2 and count_message returned -1 (base subtracts a
+                # 3-token reply overhead), against a true 1,003 tokens. That is a
+                # ~99.8% undercount on every HF-routed family whose resolved
+                # tokenizer carries a chat template — llama, qwen, deepseek, phi,
+                # yi, falcon, starcoder. pyproject pins transformers>=5.5.0,<6.0,
+                # so the affected version is the only installable one.
                 formatted = self.tokenizer.apply_chat_template(
                     messages,
                     tokenize=True,
                     add_generation_prompt=True,
+                    return_dict=False,
                 )
                 return len(formatted)
             except Exception:

headroom/tokenizers/registry.py

@@ -24,11 +24,13 @@ logger = logging.getLogger(__name__)
 # Order matters - more specific patterns first
 MODEL_PATTERNS: list[tuple[str, str]] = [
     # OpenAI models -> tiktoken
+    (r"^gpt-5", "tiktoken"),
     (r"^gpt-4o", "tiktoken"),
     (r"^gpt-4", "tiktoken"),
     (r"^gpt-3\.5", "tiktoken"),
     (r"^o1", "tiktoken"),
     (r"^o3", "tiktoken"),
+    (r"^o4", "tiktoken"),
     (r"^text-embedding", "tiktoken"),
     (r"^text-davinci", "tiktoken"),
     (r"^code-", "tiktoken"),
@@ -74,6 +76,42 @@ MODEL_PATTERNS: list[tuple[str, str]] = [
 ]
 
 
+def _name_candidates(model_lower: str) -> tuple[str, ...]:
+    """Progressively-unwrapped forms of a model name, most specific first.
+
+    Every entry in :data:`MODEL_PATTERNS` is anchored with ``^``, which is right
+    for a bare model id and wrong for the wrapped ids gateways actually send. A
+    name like ``bedrock/anthropic.claude-sonnet-4-6-v1:0`` matched nothing and
+    fell through to the char estimator instead of the Claude counter — measured
+    deviation on identical text: +15% English, -33% JSON, -38% logs. Affected
+    every ``bedrock/``, ``vertex_ai/``, ``openrouter/``, ``anthropic/``,
+    ``azure/``, ``groq/`` and ``litellm/`` form, plus Bedrock's bare
+    ``anthropic.claude-…`` and its ``us.``/``eu.``/``apac.`` region variants.
+
+    Yielding candidates rather than rewriting the name keeps the exact-match case
+    first, so no currently-correct resolution can change.
+    """
+    seen: list[str] = []
+
+    def add(name: str) -> None:
+        if name and name not in seen:
+            seen.append(name)
+
+    add(model_lower)
+    # Strip provider path segments left-to-right: openrouter/anthropic/claude-x
+    # yields anthropic/claude-x then claude-x.
+    rest = model_lower
+    while "/" in rest:
+        rest = rest.split("/", 1)[1]
+        add(rest)
+    # Bedrock dotted ids: [region.]vendor.model
+    for candidate in list(seen):
+        parts = candidate.split(".")
+        for i in range(1, len(parts)):
+            add(".".join(parts[i:]))
+    return tuple(seen)
+
+
 class TokenizerRegistry:
     """Registry for tokenizer instances and factories.
 
@@ -289,9 +327,10 @@ class TokenizerRegistry:
         """
         model_lower = model.lower()
 
-        for pattern, backend in MODEL_PATTERNS:
-            if re.match(pattern, model_lower):
-                return backend
+        for candidate in _name_candidates(model_lower):
+            for pattern, backend in MODEL_PATTERNS:
+                if re.match(pattern, candidate):
+                    return backend
 
         # Default to estimation for unknown models
         return "estimation"

tests/test_tokenizer_selection_coverage.py

@@ -0,0 +1,108 @@
+"""Model names must resolve to the tokenizer their model actually uses.
+
+Two selection gaps, both measured against real counters on identical text:
+
+1. ``MODEL_PATTERNS`` stopped at ``^gpt-4``/``^o1``/``^o3``, so the current
+   flagships — ``gpt-5``, ``gpt-5.1``, ``o4-mini`` — fell through to the char
+   estimator. Deviation vs the correct o200k encoding: +20% English, -33% JSON,
+   -44% logs.
+
+2. Every pattern is ``^``-anchored, which is right for a bare model id and wrong
+   for the wrapped ids gateways send. ``bedrock/anthropic.claude-3-5-sonnet``,
+   ``vertex_ai/claude-…``, ``openrouter/anthropic/claude-…``, ``azure/gpt-4o``
+   and Bedrock's ``us.anthropic.claude-…`` all matched nothing. LiteLLM's
+   ``headroom`` guardrail passes exactly these forms.
+
+The estimator is a legitimate FALLBACK; the bug is reaching it when a real
+tokenizer for that family exists.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from headroom.tokenizers import get_tokenizer
+from headroom.tokenizers.registry import _name_candidates
+
+_TIKTOKEN = "TiktokenCounter"
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        "gpt-5",
+        "gpt-5.1",
+        "gpt-5-mini",
+        "gpt-5.1-codex",
+        "o4-mini",
+    ],
+)
+def test_current_openai_flagships_get_a_real_tokenizer(model: str) -> None:
+    """These fell to EstimatingTokenCounter before ^gpt-5 / ^o4 were added."""
+    assert type(get_tokenizer(model)).__name__ == _TIKTOKEN
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        # gateway path prefixes
+        "bedrock/anthropic.claude-3-5-sonnet",
+        "vertex_ai/claude-sonnet-4-6",
+        "openrouter/anthropic/claude-sonnet-4-6",
+        "anthropic/claude-opus-4",
+        "litellm/claude-sonnet-4-6",
+        # Bedrock dotted ids, with and without a region segment
+        "anthropic.claude-3-5-sonnet-20241022-v2:0",
+        "us.anthropic.claude-sonnet-4-6-v1:0",
+        "eu.anthropic.claude-sonnet-4-6-v1:0",
+        # OpenAI behind a gateway
+        "azure/gpt-4o",
+        "openrouter/openai/gpt-4o",
+    ],
+)
+def test_gateway_wrapped_names_resolve_like_their_bare_form(model: str) -> None:
+    assert type(get_tokenizer(model)).__name__ == _TIKTOKEN
+
+
+def test_wrapped_gemini_matches_the_bare_form_exactly() -> None:
+    """Prefix stripping must reach the google backend, not the generic fallback."""
+    text = "hello world " * 200
+    assert get_tokenizer("vertex_ai/gemini-2.5-pro").count_text(text) == get_tokenizer(
+        "gemini-2.5-pro"
+    ).count_text(text)
+
+
+def test_bare_names_are_unaffected() -> None:
+    """The exact-match candidate is tried first, so nothing already-correct moves."""
+    for model, expected in (
+        ("gpt-4o", _TIKTOKEN),
+        ("gpt-3.5-turbo", _TIKTOKEN),
+        ("o1-preview", _TIKTOKEN),
+        ("o3-mini", _TIKTOKEN),
+        ("claude-sonnet-4-6", _TIKTOKEN),
+    ):
+        assert type(get_tokenizer(model)).__name__ == expected, model
+
+
+def test_unknown_alias_still_falls_back_to_estimation() -> None:
+    """Prefix stripping must not invent a match for a genuinely unknown model."""
+    assert type(get_tokenizer("my-gateway/big-model")).__name__ == "EstimatingTokenCounter"
+    assert type(get_tokenizer("totally-unknown-xyz")).__name__ == "EstimatingTokenCounter"
+
+
+def test_name_candidates_orders_most_specific_first() -> None:
+    """The full name must be candidate 0 so exact registrations always win."""
+    got = _name_candidates("openrouter/anthropic/claude-sonnet-4-6")
+    assert got[0] == "openrouter/anthropic/claude-sonnet-4-6"
+    assert "anthropic/claude-sonnet-4-6" in got
+    assert "claude-sonnet-4-6" in got
+
+    dotted = _name_candidates("us.anthropic.claude-sonnet-4-6-v1:0")
+    assert dotted[0] == "us.anthropic.claude-sonnet-4-6-v1:0"
+    assert "claude-sonnet-4-6-v1:0" in dotted
+
+
+def test_name_candidates_is_deduplicated_and_finite() -> None:
+    got = _name_candidates("a/b/c.d.e")
+    assert len(got) == len(set(got))
+    assert got[0] == "a/b/c.d.e"