fix(tokenizers): bound tiktoken vocab load so a stalled download cannot hang requests (#956) (#994)

## Description

The #956 symptom — `compression_first_stage` always timing out at ~30s
with 0 tokens removed — is not a PyO3/event-loop issue (compression
itself runs ~120ms on Python 3.14). Root cause: `tiktoken` downloads its
BPE vocab via `requests.get(...)` with no timeout, loaded lazily inside
the compression worker (`TiktokenCounter.encoding`,
`AnthropicProvider.__init__`). On a firewalled network that blocks
indefinitely, so the worker hangs and `asyncio.wait_for` trips at 30s on
every request (the hung download never caches, so it repeats). Refs #956
(runtime half).

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)

## Changes Made

- Load the encoding on a worker thread bounded by
`HEADROOM_TIKTOKEN_LOAD_TIMEOUT_SECONDS` (default 10s); on timeout raise
`TiktokenLoadError` and fall back to estimation (registry ->
EstimatingTokenCounter; Anthropic provider -> character estimate).
- Remember the first timed-out encoding so later requests fail fast
instead of re-blocking.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] New tests added for new functionality

### Test Output

```text
$ pytest tests/test_tokenizers/test_tiktoken_load_timeout.py -q
4 passed in 0.69s
```

## Real Behavior Proof

- Environment: Linux, Python 3.14 and 3.13
- Exact command / steps: timed the Rust compression call sync / via
run_in_executor / 2x concurrent on both interpreters; ran the
bounded-loader tests against a simulated stalled get_encoding
- Observed result: compression ~118-120ms on both 3.14 and 3.13 (no
event-loop block); the bounded loader raises/falls back within the
timeout instead of hanging
- Not tested: the real firewalled-network stall (could not reproduce on
an unfirewalled host); the no-timeout requests.get is confirmed in
tiktoken's source

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

headroom/providers/anthropic.py

@@ -274,11 +274,18 @@ class AnthropicTokenCounter(TokenCounter):
             )
             _FALLBACK_WARNING_SHOWN = True
 
-        # Load tiktoken as fallback
+        # Load tiktoken as fallback — bounded, so a stalled vocab download can't
+        # hang token counting inside a request (tiktoken's downloader has no
+        # network timeout); on timeout we estimate by characters instead (GH #956).
         try:
-            import tiktoken
+            from headroom.tokenizers.tiktoken_counter import (
+                TiktokenLoadError,
+                load_encoding,
+            )
 
-            self._encoding = tiktoken.get_encoding("cl100k_base")
+            self._encoding = load_encoding("cl100k_base")
+        except TiktokenLoadError:
+            self._encoding = None  # count_text() falls back to a character estimate
         except ImportError:
             if not self._use_api:
                 warnings.warn(

headroom/tokenizers/registry.py

@@ -287,10 +287,24 @@ class TokenizerRegistry:
         return "estimation"
 
     def _create_tiktoken(self, model: str) -> TokenCounter:
-        """Create tiktoken-based tokenizer."""
-        try:
-            from .tiktoken_counter import TiktokenCounter
+        """Create tiktoken-based tokenizer.
 
+        Forces the (bounded) encoding load up front so a stalled vocab download
+        falls back to estimation instead of hanging later inside a request (GH #956).
+        """
+        try:
+            from .tiktoken_counter import (
+                TiktokenCounter,
+                TiktokenLoadError,
+                get_encoding_for_model,
+                load_encoding,
+            )
+
+            try:
+                load_encoding(get_encoding_for_model(model))
+            except TiktokenLoadError as exc:
+                logger.warning("tiktoken unavailable (%s); using estimation.", exc)
+                return EstimatingTokenCounter()
             return TiktokenCounter(model)
         except ImportError:
             logger.warning("tiktoken not installed. Install with: pip install tiktoken")

headroom/tokenizers/tiktoken_counter.py

@@ -10,11 +10,38 @@ It supports multiple encodings:
 
 from __future__ import annotations
 
+import logging
+import os
+import threading
 from functools import lru_cache
 from typing import Any
 
 from .base import BaseTokenizer
 
+logger = logging.getLogger(__name__)
+
+
+class TiktokenLoadError(RuntimeError):
+    """Raised when a tiktoken encoding can't be loaded in time.
+
+    tiktoken downloads its BPE vocab on first use via ``requests.get`` with no
+    timeout, so a stalled/firewalled connection can block indefinitely. We bound
+    that load and raise this instead, so callers fall back to estimation rather
+    than hanging the request (see GH #956).
+    """
+
+
+# Encoding names whose bounded load already timed out — don't block on them again.
+_load_failed: set[str] = set()
+
+
+def _load_timeout_seconds() -> float:
+    try:
+        return float(os.environ.get("HEADROOM_TIKTOKEN_LOAD_TIMEOUT_SECONDS", "10"))
+    except (TypeError, ValueError):
+        return 10.0
+
+
 # Model to encoding mapping
 MODEL_TO_ENCODING = {
     # GPT-4o family (o200k_base)
@@ -78,10 +105,54 @@ DEFAULT_ENCODING = "cl100k_base"
 
 @lru_cache(maxsize=8)
 def _get_encoding(encoding_name: str):
-    """Get tiktoken encoding, cached for performance."""
+    """Get a tiktoken encoding, cached for performance.
+
+    Bounded by ``HEADROOM_TIKTOKEN_LOAD_TIMEOUT_SECONDS`` (default 10s): tiktoken's
+    vocab download has no network timeout, so we run the load on a worker thread
+    and raise :class:`TiktokenLoadError` if it doesn't finish in time, letting
+    callers fall back to estimation rather than hang the request (GH #956). The
+    first timed-out encoding is remembered so later calls fail fast instead of
+    re-blocking on every request.
+    """
     import tiktoken
 
-    return tiktoken.get_encoding(encoding_name)
+    if encoding_name in _load_failed:
+        raise TiktokenLoadError(f"tiktoken encoding {encoding_name!r} previously failed to load")
+
+    box: dict[str, Any] = {}
+
+    def _load() -> None:
+        try:
+            box["enc"] = tiktoken.get_encoding(encoding_name)
+        except BaseException as exc:  # noqa: BLE001 - re-raised in the calling thread
+            box["err"] = exc
+
+    worker = threading.Thread(target=_load, name=f"tiktoken-load-{encoding_name}", daemon=True)
+    worker.start()
+    worker.join(_load_timeout_seconds())
+
+    if worker.is_alive():
+        _load_failed.add(encoding_name)
+        logger.warning(
+            "tiktoken encoding %r did not load within %.1fs (likely a stalled vocab "
+            "download); falling back to token estimation. Pre-populate TIKTOKEN_CACHE_DIR "
+            "or tune HEADROOM_TIKTOKEN_LOAD_TIMEOUT_SECONDS.",
+            encoding_name,
+            _load_timeout_seconds(),
+        )
+        raise TiktokenLoadError(f"tiktoken encoding {encoding_name!r} load timed out")
+    if "err" in box:
+        raise box["err"]
+    return box["enc"]
+
+
+def load_encoding(encoding_name: str) -> Any:
+    """Public, bounded tiktoken-encoding loader.
+
+    Returns the tiktoken encoding, or raises :class:`TiktokenLoadError` if the
+    vocab can't be loaded within the timeout (see :func:`_get_encoding`, GH #956).
+    """
+    return _get_encoding(encoding_name)
 
 
 def get_encoding_for_model(model: str) -> str:

tests/test_tokenizers/test_tiktoken_load_timeout.py

@@ -0,0 +1,81 @@
+"""tiktoken vocab loading must be bounded (GH #956).
+
+tiktoken downloads its BPE vocab via ``requests.get`` with no timeout, so a
+stalled/firewalled connection blocks indefinitely. The proxy calls this lazily
+inside a request worker, so the only bound was the 30s compression timeout —
+yielding "every request times out, 0 compression". The bounded loader caps the
+wait and falls back to estimation instead.
+"""
+
+from __future__ import annotations
+
+import time
+
+import pytest
+
+from headroom.tokenizers import tiktoken_counter as tc
+from headroom.tokenizers.estimator import EstimatingTokenCounter
+from headroom.tokenizers.registry import TokenizerRegistry
+
+
+@pytest.fixture(autouse=True)
+def _reset_encoding_state():
+    tc._get_encoding.cache_clear()
+    tc._load_failed.clear()
+    yield
+    tc._get_encoding.cache_clear()
+    tc._load_failed.clear()
+
+
+def _stalled_get_encoding(_name: str):
+    # Simulates tiktoken's unbounded network download stalling.
+    time.sleep(2.0)
+    return object()
+
+
+def test_load_encoding_is_bounded_on_stall(monkeypatch: pytest.MonkeyPatch) -> None:
+    import tiktoken
+
+    monkeypatch.setattr(tiktoken, "get_encoding", _stalled_get_encoding)
+    monkeypatch.setenv("HEADROOM_TIKTOKEN_LOAD_TIMEOUT_SECONDS", "0.2")
+
+    start = time.perf_counter()
+    with pytest.raises(tc.TiktokenLoadError):
+        tc.load_encoding("stall-enc")
+    elapsed = time.perf_counter() - start
+    assert elapsed < 1.5, f"load was not bounded (took {elapsed:.2f}s vs the 2s stall)"
+
+
+def test_failed_encoding_short_circuits(monkeypatch: pytest.MonkeyPatch) -> None:
+    import tiktoken
+
+    monkeypatch.setattr(tiktoken, "get_encoding", _stalled_get_encoding)
+    monkeypatch.setenv("HEADROOM_TIKTOKEN_LOAD_TIMEOUT_SECONDS", "0.2")
+
+    with pytest.raises(tc.TiktokenLoadError):
+        tc.load_encoding("stall-enc-2")
+
+    # A second request must fail instantly via the _load_failed short-circuit,
+    # not wait out the timeout again (this is what makes it not "every request").
+    start = time.perf_counter()
+    with pytest.raises(tc.TiktokenLoadError):
+        tc.load_encoding("stall-enc-2")
+    assert time.perf_counter() - start < 0.1
+
+
+def test_fast_load_returns_encoding(monkeypatch: pytest.MonkeyPatch) -> None:
+    import tiktoken
+
+    sentinel = object()
+    monkeypatch.setattr(tiktoken, "get_encoding", lambda _name: sentinel)
+    assert tc.load_encoding("fast-enc") is sentinel
+
+
+def test_registry_falls_back_to_estimator_on_stall(monkeypatch: pytest.MonkeyPatch) -> None:
+    import tiktoken
+
+    monkeypatch.setattr(tiktoken, "get_encoding", _stalled_get_encoding)
+    monkeypatch.setenv("HEADROOM_TIKTOKEN_LOAD_TIMEOUT_SECONDS", "0.2")
+
+    counter = TokenizerRegistry()._create_tiktoken("gpt-4")
+    assert isinstance(counter, EstimatingTokenCounter)