diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index a0fe245a8..314262d42 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -72,46 +72,9 @@ class AnthropicHandlerMixin: """Mixin providing Anthropic API handler methods for HeadroomProxy.""" async def _count_tokens_offloaded(self, model, messages): # noqa: ANN001, ANN201 - """Resolve a tokenizer and count messages off the event loop. + from headroom.proxy.token_counting import count_tokens_offloaded - Tokenizer resolution can be expensive on first use (HuggingFace - backends may download vocab files) and counting a full Claude Code - conversation is CPU-bound, so both run on the compression executor - bounded by ``COMPRESSION_TIMEOUT_SECONDS`` (GH #1701: an unbounded - on-loop load froze the whole server). On timeout or error this - fails open to character-based estimation. - - Returns: - Tuple of ``(tokenizer, token_count)``. The tokenizer is fully - initialized, so later ``count_messages`` calls on it are pure - CPU work. - """ - from headroom.proxy.helpers import COMPRESSION_TIMEOUT_SECONDS - from headroom.tokenizers import EstimatingTokenCounter, get_tokenizer - - def _resolve_and_count(): # noqa: ANN202 - tokenizer = get_tokenizer(model) - return tokenizer, tokenizer.count_messages(messages) - - try: - return await self._run_compression_in_executor( - _resolve_and_count, - timeout=float(COMPRESSION_TIMEOUT_SECONDS), - ) - except Exception as e: # fail open — includes asyncio.TimeoutError - # Log the downgrade once per model, not per request. - fallback_models = getattr(self, "_token_count_fallback_models", None) - if fallback_models is None: - fallback_models = set() - self._token_count_fallback_models = fallback_models - if model not in fallback_models: - fallback_models.add(model) - logger.warning( - f"Token counting for model {model} failed or timed out " - f"({e.__class__.__name__}); falling back to estimation" - ) - estimator = EstimatingTokenCounter() - return estimator, estimator.count_messages(messages) + return await count_tokens_offloaded(self, model, messages) @staticmethod def _resolve_ccr_workspace( diff --git a/headroom/proxy/handlers/gemini.py b/headroom/proxy/handlers/gemini.py index dba8bbbb1..0d176d129 100644 --- a/headroom/proxy/handlers/gemini.py +++ b/headroom/proxy/handlers/gemini.py @@ -38,6 +38,16 @@ def _usage_int(value: Any, default: int = 0) -> int: class GeminiHandlerMixin: """Mixin providing Gemini API handler methods for HeadroomProxy.""" + async def _count_tokens_offloaded(self, model, messages): # noqa: ANN001, ANN201 + from headroom.proxy.token_counting import count_tokens_offloaded + + return await count_tokens_offloaded(self, model, messages) + + async def _count_texts_offloaded(self, model, texts): # noqa: ANN001, ANN201 + from headroom.proxy.token_counting import count_texts_offloaded + + return await count_texts_offloaded(self, model, texts) + def _is_cloudcode_antigravity_request( self, body: dict[str, Any], headers: dict[str, str] ) -> bool: @@ -259,7 +269,6 @@ class GeminiHandlerMixin: from fastapi.responses import JSONResponse, Response from headroom.proxy.helpers import MAX_REQUEST_BODY_SIZE, _read_request_json - from headroom.tokenizers import get_tokenizer from headroom.utils import extract_user_query start_time = time.time() @@ -491,9 +500,8 @@ class GeminiHandlerMixin: headers=response_headers, ) - # Token counting - tokenizer = get_tokenizer(model) - original_tokens = tokenizer.count_messages(messages) + # Token counting (offloaded off the event loop — GH #1701) + tokenizer, original_tokens = await self._count_tokens_offloaded(model, messages) # Optimization transforms_applied: list[str] = [] @@ -816,7 +824,6 @@ class GeminiHandlerMixin: from fastapi.responses import JSONResponse from headroom.proxy.helpers import _read_request_json - from headroom.tokenizers import get_tokenizer from headroom.utils import extract_user_query start_time = time.time() @@ -880,8 +887,10 @@ class GeminiHandlerMixin: if isinstance(contents, list) and idx < len(contents) } - tokenizer = get_tokenizer(model) - original_tokens = tokenizer.count_messages(messages) if messages else 0 + # Token counting (offloaded off the event loop — GH #1701) + tokenizer, original_tokens = await self._count_tokens_offloaded(model, messages) + if not messages: + original_tokens = 0 optimized_messages = messages optimized_tokens = original_tokens transforms_applied: list[str] = [] @@ -981,7 +990,6 @@ class GeminiHandlerMixin: from fastapi.responses import JSONResponse from headroom.proxy.helpers import _read_request_json - from headroom.tokenizers import get_tokenizer start_time = time.time() request_id = await self._next_request_id() @@ -1019,14 +1027,17 @@ class GeminiHandlerMixin: request_id=request_id, ) - # Token counting - tokenizer = get_tokenizer(model) - original_tokens = 0 - for content in contents: - parts = content.get("parts", []) - for part in parts: - if "text" in part: - original_tokens += tokenizer.count_text(part["text"]) + # Token counting (offloaded off the event loop — GH #1701). Reuse the + # shared _dict_parts coercion and keep only str text values: count_text + # raises on a non-str part value and the fail-open path re-runs the same + # input, so a malformed part would otherwise 500 the streaming request. + text_parts = [ + part["text"] + for content in (contents if isinstance(contents, list) else []) + for part in self._dict_parts(content) + if isinstance(part.get("text"), str) + ] + _, original_tokens = await self._count_texts_offloaded(model, text_parts) optimization_latency = (time.time() - start_time) * 1000 @@ -1069,7 +1080,6 @@ class GeminiHandlerMixin: from fastapi.responses import JSONResponse, Response from headroom.proxy.helpers import _read_request_json - from headroom.tokenizers import get_tokenizer from headroom.utils import extract_user_query start_time = time.time() @@ -1140,9 +1150,8 @@ class GeminiHandlerMixin: headers=response_headers, ) - # Token counting (original) - tokenizer = get_tokenizer(model) - original_tokens = tokenizer.count_messages(messages) + # Token counting (original, offloaded off the event loop — GH #1701) + tokenizer, original_tokens = await self._count_tokens_offloaded(model, messages) # Apply compression using the same pipeline as generateContent transforms_applied: list[str] = [] diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 3637d7d63..3830d5a35 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -1316,6 +1316,11 @@ def _prefers_http1_passthrough(base_url: str) -> bool: class OpenAIHandlerMixin: """Mixin providing OpenAI API handler methods for HeadroomProxy.""" + async def _count_tokens_offloaded(self, model, messages): # noqa: ANN001, ANN201 + from headroom.proxy.token_counting import count_tokens_offloaded + + return await count_tokens_offloaded(self, model, messages) + OPENAI_RESPONSES_ROUTER_MIN_BYTES = 512 OPENAI_RESPONSES_OUTPUT_TYPES = _RESPONSES_OUTPUT_ITEM_TYPES @@ -2576,7 +2581,6 @@ class OpenAIHandlerMixin: _read_request_json, ) from headroom.proxy.modes import is_cache_mode, is_token_mode - from headroom.tokenizers import get_tokenizer from headroom.utils import extract_user_query start_time = time.time() @@ -2905,9 +2909,8 @@ class OpenAIHandlerMixin: return Response(content=cached.response_body, headers=response_headers) - # Token counting - tokenizer = get_tokenizer(model) - original_tokens = tokenizer.count_messages(messages) + # Token counting (offloaded off the event loop — GH #1701) + tokenizer, original_tokens = await self._count_tokens_offloaded(model, messages) # Hook: pre_compress _hook_biases = None @@ -4257,7 +4260,6 @@ class OpenAIHandlerMixin: MAX_REQUEST_BODY_SIZE, read_request_json_with_bytes, ) - from headroom.tokenizers import get_tokenizer from headroom.utils import extract_user_query start_time = time.time() @@ -4474,9 +4476,8 @@ class OpenAIHandlerMixin: detail=f"Rate limited. Retry after {wait_seconds:.1f}s", ) - # Token counting on converted messages - tokenizer = get_tokenizer(model) - original_tokens = tokenizer.count_messages(messages) + # Token counting on converted messages (offloaded off the event loop — GH #1701) + tokenizer, original_tokens = await self._count_tokens_offloaded(model, messages) # Defaults below feed downstream telemetry and memory injection. # If optimization remains enabled, the Responses payload is compressed diff --git a/headroom/proxy/token_counting.py b/headroom/proxy/token_counting.py new file mode 100644 index 000000000..fb9ca107b --- /dev/null +++ b/headroom/proxy/token_counting.py @@ -0,0 +1,79 @@ +"""Offloaded token-count helpers shared by proxy handlers. + +Tokenizer resolution can be expensive on first use (HuggingFace backends may +download vocab files) and counting a full Claude Code conversation is CPU-bound, +so both run on the caller's compression executor bounded by +``COMPRESSION_TIMEOUT_SECONDS`` (GH #1701: an unbounded on-loop load froze the +whole server). On timeout, error, or a missing executor this fails open to +character-based estimation. + +Shared by every provider handler mixin (Anthropic, OpenAI, Gemini): the OpenAI +``/v1/chat/completions`` and ``/v1/responses`` endpoints are multi-provider +passthroughs, so an HF-routed model (qwen, deepseek, llama, ...) can reach them +and trigger the same cold load. +""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from typing import Any, cast + +logger = logging.getLogger("headroom.proxy") + + +def _record_fallback_model(owner: Any, model: Any, message: str) -> None: + fallback_models = getattr(owner, "_token_count_fallback_models", None) + if fallback_models is None: + fallback_models = set() + owner._token_count_fallback_models = fallback_models + if model not in fallback_models: + fallback_models.add(model) + logger.warning(message) + + +async def _count_offloaded(owner: Any, model: Any, count: Callable[[Any], int]) -> tuple[Any, int]: + """Resolve a tokenizer and apply ``count`` off the event loop when possible. + + ``count`` maps a resolved tokenizer to a token total. Returns + ``(tokenizer, total)``; the returned tokenizer is fully initialized, so later + counts on it are pure CPU work. Fails open to ``EstimatingTokenCounter`` when + the owner has no compression executor, or on timeout/error. + """ + from headroom.proxy.helpers import COMPRESSION_TIMEOUT_SECONDS + from headroom.tokenizers import EstimatingTokenCounter, get_tokenizer + + runner = getattr(owner, "_run_compression_in_executor", None) + if runner is None: + estimator = EstimatingTokenCounter() + return estimator, count(estimator) + + def _resolve_and_count() -> tuple[Any, int]: + tokenizer = get_tokenizer(model) + return tokenizer, count(tokenizer) + + try: + result = await runner(_resolve_and_count, timeout=float(COMPRESSION_TIMEOUT_SECONDS)) + return cast(tuple[Any, int], result) + except Exception as e: # fail open — includes asyncio.TimeoutError + _record_fallback_model( + owner, + model, + f"Token counting for model {model} failed or timed out " + f"({e.__class__.__name__}); falling back to estimation", + ) + estimator = EstimatingTokenCounter() + return estimator, count(estimator) + + +async def count_tokens_offloaded(owner: Any, model: Any, messages: Any) -> tuple[Any, int]: + """Resolve a tokenizer and count ``messages`` off the event loop when possible.""" + return await _count_offloaded(owner, model, lambda counter: counter.count_messages(messages)) + + +async def count_texts_offloaded(owner: Any, model: Any, texts: Any) -> tuple[Any, int]: + """Resolve a tokenizer and count text fragments off the event loop when possible.""" + text_list = list(texts) + return await _count_offloaded( + owner, model, lambda counter: sum(counter.count_text(text) for text in text_list) + ) diff --git a/tests/test_gemini_nonjson_status.py b/tests/test_gemini_nonjson_status.py index fd8673936..341fbaab2 100644 --- a/tests/test_gemini_nonjson_status.py +++ b/tests/test_gemini_nonjson_status.py @@ -57,6 +57,14 @@ class _Handler(GeminiHandlerMixin): async def _record_request_outcome(self, outcome) -> None: # noqa: ANN001 self.outcomes.append(outcome) + async def _count_tokens_offloaded(self, model, messages): # noqa: ANN001, ANN201 + # Test stub for HeadroomProxy._count_tokens_offloaded: resolve the + # tokenizer and count inline (the real method offloads to the executor). + from headroom.tokenizers import get_tokenizer + + tokenizer = get_tokenizer(model) + return tokenizer, tokenizer.count_messages(messages) + @pytest.mark.asyncio async def test_generate_content_forwards_non_json_upstream_status( diff --git a/tests/test_openai_codex_routing.py b/tests/test_openai_codex_routing.py index 47839983d..16a44ea1b 100644 --- a/tests/test_openai_codex_routing.py +++ b/tests/test_openai_codex_routing.py @@ -205,6 +205,14 @@ class _DummyOpenAIHandler(OpenAIHandlerMixin): # synchronously so MagicMock call_count assertions fire. return fn() + async def _count_tokens_offloaded(self, model, messages): # noqa: ANN001, ANN201 + # Test stub for HeadroomProxy._count_tokens_offloaded: resolve the + # tokenizer and count inline (the real method offloads to the executor). + from headroom.tokenizers import get_tokenizer + + tokenizer = get_tokenizer(model) + return tokenizer, tokenizer.count_messages(messages) + async def _record_request_outcome(self, outcome) -> None: # Test stub: delegates to the production funnel so wire shape # matches HeadroomProxy._record_request_outcome. diff --git a/tests/test_tokenizer_count_offload.py b/tests/test_tokenizer_count_offload.py index 239d5e1ef..e4d77e4bc 100644 --- a/tests/test_tokenizer_count_offload.py +++ b/tests/test_tokenizer_count_offload.py @@ -4,8 +4,10 @@ handler. For HF-backed models (e.g. deepseek-*) first use triggers an unbounded network download, freezing the whole server (610s request, then /livez, /readyz and /health hang until kill). The fix routes resolution + counting through HeadroomProxy._count_tokens_offloaded (compression executor, bounded by -COMPRESSION_TIMEOUT_SECONDS, fail-open to estimation), and offloads the inline -batch pipeline.apply() calls the same way. +COMPRESSION_TIMEOUT_SECONDS, fail-open to estimation) — shared by every provider +handler (Anthropic, OpenAI, Gemini), since the OpenAI passthrough endpoints +receive the same HF-backed models — and offloads the inline batch +pipeline.apply() calls the same way. """ from __future__ import annotations @@ -17,7 +19,18 @@ import time from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin from headroom.proxy.handlers.batch import BatchHandlerMixin -from headroom.proxy.server import ProxyConfig, create_app +from headroom.proxy.handlers.gemini import GeminiHandlerMixin +from headroom.proxy.handlers.openai import OpenAIHandlerMixin +from headroom.proxy.server import ( + CompressionQuarantinedError, + ProxyConfig, + create_app, +) +from headroom.proxy.token_counting import ( + _count_offloaded, + count_texts_offloaded, + count_tokens_offloaded, +) from headroom.tokenizers import EstimatingTokenCounter @@ -36,11 +49,36 @@ def _make_proxy(): # noqa: ANN202 — returns the internal HeadroomProxy def test_handlers_offload_token_counting_and_batch_apply() -> None: """Wiring guard: the request paths must use the offloaded helpers, not inline get_tokenizer/count_messages or pipeline.apply on the event loop.""" - fn = AnthropicHandlerMixin.handle_anthropic_messages + # Every provider handler that counts the original conversation must route + # resolution + counting through the shared fail-open helper, never inline on + # the loop. OpenAI /chat + /responses are multi-provider passthroughs, so an + # HF-routed model (qwen, deepseek, llama, ...) can reach them and cold-load. + for mixin, method in ( + (AnthropicHandlerMixin, "handle_anthropic_messages"), + (OpenAIHandlerMixin, "handle_openai_chat"), + (OpenAIHandlerMixin, "handle_openai_responses"), + (GeminiHandlerMixin, "handle_gemini_generate_content"), + (GeminiHandlerMixin, "handle_google_cloudcode_stream"), + (GeminiHandlerMixin, "handle_gemini_count_tokens"), + ): + fn = getattr(mixin, method) + assert inspect.iscoroutinefunction(fn), f"{method} must be async" + src = inspect.getsource(fn) + assert "_count_tokens_offloaded(" in src, f"{method}: token counting not offloaded" + assert "tokenizer = get_tokenizer(" not in src, ( + f"{method}: tokenizer resolved inline on the loop" + ) + + fn = GeminiHandlerMixin.handle_gemini_stream_generate_content assert inspect.iscoroutinefunction(fn) src = inspect.getsource(fn) - assert "_count_tokens_offloaded(" in src, "token counting not offloaded" + assert "_count_texts_offloaded(" in src, "streaming Gemini text counting not offloaded" assert "tokenizer = get_tokenizer(" not in src, "tokenizer resolved inline on the loop" + assert "count_text(" not in src, "streaming Gemini count_text still runs on the loop" + assert "_dict_parts(" in src, "streaming Gemini must reuse the shared _dict_parts coercion" + assert 'isinstance(part.get("text"), str)' in src, ( + "streaming Gemini must skip non-str text so count_text can't 500" + ) for mixin, method in ( (AnthropicHandlerMixin, "handle_anthropic_batch_create"), @@ -54,7 +92,7 @@ def test_handlers_offload_token_counting_and_batch_apply() -> None: assert "_run_compression_in_executor(" in src, f"{method}: apply() not offloaded" assert "COMPRESSION_TIMEOUT_SECONDS" in src, f"{method}: offload missing timeout" - helper_src = inspect.getsource(AnthropicHandlerMixin._count_tokens_offloaded) + helper_src = inspect.getsource(_count_offloaded) assert "COMPRESSION_TIMEOUT_SECONDS" in helper_src assert "EstimatingTokenCounter" in helper_src, "helper must fail open to estimation" @@ -124,3 +162,115 @@ async def test_count_tokens_offloaded_fails_open(monkeypatch) -> None: # noqa: assert tokens > 0 # Logged-once bookkeeping records the downgraded model. assert "deepseek-chat" in proxy._token_count_fallback_models + + +async def test_count_tokens_offloaded_fails_open_on_executor_quarantine() -> None: + """Now that OpenAI/Gemini counting shares the compression executor, an + unrelated request's compression timeout can quarantine it — the next + ``_run_compression_in_executor`` call raises ``CompressionQuarantinedError`` + immediately (process-wide state). A request that is only counting tokens + must not 500 on that; it fails open to estimation like any other error.""" + # The executor's ``except Exception`` fail-open only catches the quarantine + # error because it subclasses Exception — pin that contract. + assert issubclass(CompressionQuarantinedError, Exception) + + proxy = _make_proxy() + # Record a concurrent compression as timed out so the real executor guard + # quarantines the next call — no mock of the helper itself. + proxy._compression_timed_out_in_flight = 1 + + tokenizer, tokens = await proxy._count_tokens_offloaded( + "qwen2.5-coder", [{"role": "user", "content": "hello world"}] + ) + + assert isinstance(tokenizer, EstimatingTokenCounter) + assert tokens > 0 + assert "qwen2.5-coder" in proxy._token_count_fallback_models + + +async def test_count_tokens_offloaded_returns_count_text_capable_tokenizer() -> None: + """The fail-open tokenizer should still support text counting for callers + that need per-fragment accounting.""" + proxy = _make_proxy() + # Quarantine forces the fail-open branch (an EstimatingTokenCounter). + proxy._compression_timed_out_in_flight = 1 + + # The empty-messages count is intentionally discarded by that handler + # (it sums text parts itself), so only the tokenizer matters here. + tokenizer, _ = await proxy._count_tokens_offloaded("qwen2.5-coder", []) + + assert isinstance(tokenizer, EstimatingTokenCounter) + # The streaming handler's per-part loop must not raise on the fallback. + assert tokenizer.count_text("hello world") > 0 + + +async def test_count_texts_offloaded_runs_on_worker_thread(monkeypatch) -> None: # noqa: ANN001 + proxy = _make_proxy() + loop_thread = threading.current_thread().name + seen: dict[str, str] = {} + + class _SpyTokenizer(EstimatingTokenCounter): + def count_text(self, text): # noqa: ANN001, ANN201 + seen["thread"] = threading.current_thread().name + return super().count_text(text) + + monkeypatch.setattr("headroom.tokenizers.get_tokenizer", lambda *a, **k: _SpyTokenizer()) + + _, tokens = await proxy._count_texts_offloaded("gemini-pro", ["hello", "world"]) + + assert tokens > 0 + assert seen["thread"].startswith("headroom-compress") + assert seen["thread"] != loop_thread + + +async def test_count_texts_offloaded_fails_open(monkeypatch) -> None: # noqa: ANN001 + """The texts variant downgrades to estimation on a resolution error, the same + as the messages variant (its fail-open branch was previously uncovered).""" + proxy = _make_proxy() + + def _boom(*a, **k): # noqa: ANN002, ANN003, ANN202 + raise RuntimeError("tokenizer backend exploded") + + monkeypatch.setattr("headroom.tokenizers.get_tokenizer", _boom) + + tokenizer, tokens = await proxy._count_texts_offloaded("deepseek-chat", ["hello", "world"]) + + assert isinstance(tokenizer, EstimatingTokenCounter) + assert tokens > 0 + assert "deepseek-chat" in proxy._token_count_fallback_models + + +async def test_count_offloaded_without_executor_estimates() -> None: + """An owner with no compression executor (a lightweight caller or test double) + fails open to estimation inline instead of crashing on the missing runner.""" + + class _NoExecutorOwner: + pass + + owner = _NoExecutorOwner() + + tok, n_msg = await count_tokens_offloaded( + owner, "gpt-4", [{"role": "user", "content": "hello world"}] + ) + assert isinstance(tok, EstimatingTokenCounter) + assert n_msg > 0 + + tok2, n_txt = await count_texts_offloaded(owner, "gemini-pro", ["hello", "world"]) + assert isinstance(tok2, EstimatingTokenCounter) + assert n_txt > 0 + + +async def test_count_texts_offloaded_sums_fragments(monkeypatch) -> None: # noqa: ANN001 + """The streaming rewrite sums per-fragment counts, matching the old per-part + count_text loop it replaced.""" + proxy = _make_proxy() + monkeypatch.setattr( + "headroom.tokenizers.get_tokenizer", lambda *a, **k: EstimatingTokenCounter() + ) + fragments = ["hello", "world", "foo"] + + _, total = await proxy._count_texts_offloaded("gemini-pro", fragments) + + est = EstimatingTokenCounter() + assert total == sum(est.count_text(f) for f in fragments) + assert total > 0