From 06add9e9d833783c1144316c0cb3cb142377d897 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Mon, 3 Aug 2026 22:40:38 -0700 Subject: [PATCH] fix(providers): stop pricing modern content blocks at zero (#2760) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Each token counter in `headroom/providers/` had grown its own shortened content-block walker, handling only the shapes its provider was expected to send. Everything else fell through and contributed **nothing**. Measured on one 6,800-char block, via `count_messages` of a single-block message — so 7–8 is message overhead alone: | block type | OpenAI ctr | Anthropic ctr | |---|---|---| | `text` (control) | 3409 | 3748 | | `tool_result` | **8** | 3748 | | `thinking` | **8** | **7** | | `document` | **8** | **7** | | `mcp_tool_result` | **8** | **7** | | `output_text` | **8** | **7** | | `refusal` | **8** | **7** | Two things make this worse than a coverage gap: 1. **Each counter zeroed blocks from its own provider.** `output_text` and `refusal` are OpenAI Responses shapes; `thinking` and `document` are Anthropic's. 2. **These are the counters the live pipelines use.** `proxy/server.py` builds them with `AnthropicProvider` / `OpenAIProvider`, so this is the main request path — not an edge case. #2743 fixed this for `/v1/compress` only, by routing that route to the registry tokenizers, whose `BaseTokenizer` walker is complete. Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made Rather than add a **fifth** partial walker, the counters now delegate to the audited one: - `tokenizers/base.py` — new `count_content_blocks(parts, count_text_fn)` plus a thin `_DelegatingBlockCounter` adapter, since the provider counters are not `BaseTokenizer` subclasses. `BaseTokenizer` itself is untouched. - `providers/openai.py`, `providers/anthropic.py`, `providers/openai_compatible.py` — list-content branches delegate. **Why delegate instead of adding a `count_text(str(block))` catch-all:** that would serialize a base64 blob and price it as text. `tiktoken_counter.py` already documents the failure — a 1MB image becomes ~330K phantom tokens. The shared walker gives media a pixel/byte-based estimate. **Scope:** the three counters that accumulate token counts. `google.py` and `cohere.py` extract a *text string* first and count that, so the same defect there needs a differently-shaped fix — left as a follow-up. ## Testing - [x] Unit tests pass - [x] Linting passes (`ruff check` + `format --check`, pinned 0.15.17) - [x] New tests added - [x] Manual testing performed ### Test Output ```text $ pytest tests/test_provider_counter_content_blocks.py -q 12 passed in 0.52s $ uvx ruff@0.15.17 check headroom/ tests/... --exclude headroom/dashboard/templates All checks passed! ``` **After the fix**, every shape lands within ~1% of the equivalent plain text, and media stays bounded: ```text block OpenAI Anthropic text (control) 3409 3748 tool_result 3409 3748 thinking 3419 3759 document 3423 3763 mcp_tool_result 3422 3762 output_text 3420 3760 refusal 3421 3761 image b64 200KB 1608 1607 <- pixel estimate, not ~50K as text ``` ## Real Behavior Proof — including a regression I caught **This change flipped an existing test**, and I only found it because every suite was run against clean `upstream/main` in the same environment with the failure sets diffed: ```text before the test rewrite: upstream/main : 1 failed, 104 passed this branch : 2 failed, 103 passed <- regression diff : + test_openai_compatible_token_counter_ignores_unhandled_content_shapes ``` That test asserted `content: [{"type": "image"}, 123] == 8` — i.e. it **pinned the defect**, that unhandled shapes contribute nothing. Rewritten as `..._prices_declared_media`: a declared image is now priced (1608) while a bare int is still correctly ignored (8), with the rationale in the docstring. ```text after the rewrite: upstream/main : 1 failed, 104 passed, 10 skipped, 25 errors this branch : 1 failed, 104 passed, 10 skipped, 25 errors failure sets : IDENTICAL ``` - **Pre-existing, not from this change:** the 1 failure and all 25 errors. The errors are all in `test_compress_route_tokenizer_by_model.py`, whose loopback `TestClient` fixture this throwaway env cannot satisfy. - **Environment note:** `content_router` and several suites need the compiled `headroom._core`, which isn't in a fresh worktree (gitignored, built in-place). I copied the built `.so` in to run these and removed it before committing. - **Not tested:** no live provider call, so the *absolute* accuracy of the 1600 image estimate against a real Anthropic/OpenAI bill is unverified — it is the value `BaseTokenizer` already used, and this PR only changes which blocks reach it. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective - [x] I did **not** edit `CHANGELOG.md` ## Related Fourth PR from one tokenizer-consistency audit: #2757 (litellm total prompt / `--budget`), #2758 (HuggingFace chat templates, `gpt-5`, gateway-wrapped names), #2759 (router token units). Plus #2756, which splits the local/provider token scales in `RequestOutcome`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- headroom/providers/anthropic.py | 20 ++-- headroom/providers/openai.py | 19 ++-- headroom/providers/openai_compatible.py | 17 ++- headroom/tokenizers/base.py | 37 ++++++ tests/test_provider_counter_content_blocks.py | 106 ++++++++++++++++++ tests/test_providers/test_universal.py | 19 +++- 6 files changed, 193 insertions(+), 25 deletions(-) create mode 100644 tests/test_provider_counter_content_blocks.py diff --git a/headroom/providers/anthropic.py b/headroom/providers/anthropic.py index 107fc53c5..0bca5be84 100644 --- a/headroom/providers/anthropic.py +++ b/headroom/providers/anthropic.py @@ -24,6 +24,7 @@ import warnings from typing import Any, cast from headroom import paths as _paths +from headroom.tokenizers.base import count_content_blocks from .base import Provider, TokenCounter @@ -395,15 +396,16 @@ class AnthropicTokenCounter(TokenCounter): if isinstance(content, str): tokens += self.count_text(content) elif isinstance(content, list): - for block in content: - if isinstance(block, dict): - if block.get("type") == "text": - tokens += self.count_text(block.get("text", "")) - elif block.get("type") == "tool_use": - tokens += self.count_text(block.get("name", "")) - tokens += self.count_text(str(block.get("input", {}))) - elif block.get("type") == "tool_result": - tokens += self.count_text(str(block.get("content", ""))) + # Delegate to the audited shared walker instead of a partial + # per-provider one. Each provider counter had grown its own + # shortened branch list, so every modern block priced at ~0: + # measured on a 6,800-char block this returned 8 tokens for + # tool_result, thinking, document, mcp_tool_result — and for + # output_text / refusal, which are OpenAI's OWN Responses shapes. + # The shared walker is also image-safe: a 200KB base64 image gets + # a pixel-based 1600, not the ~50K phantom text tokens a naive + # str(block) catch-all would produce. + tokens += count_content_blocks(content, self.count_text) # OpenAI format tool calls if "tool_calls" in message: diff --git a/headroom/providers/openai.py b/headroom/providers/openai.py index cc91d2b71..73a953c21 100644 --- a/headroom/providers/openai.py +++ b/headroom/providers/openai.py @@ -16,6 +16,7 @@ from functools import lru_cache from typing import Any, cast from headroom import paths as _paths +from headroom.tokenizers.base import count_content_blocks from .base import Provider, TokenCounter @@ -331,14 +332,16 @@ class OpenAITokenCounter: if isinstance(content, str): tokens += self.count_text(content) elif isinstance(content, list): - for part in content: - if isinstance(part, dict): - if part.get("type") == "text": - tokens += self.count_text(part.get("text", "")) - elif part.get("type") == "image_url": - tokens += 85 # Low detail image estimate - elif isinstance(part, str): - tokens += self.count_text(part) + # Delegate to the audited shared walker instead of a partial + # per-provider one. Each provider counter had grown its own + # shortened branch list, so every modern block priced at ~0: + # measured on a 6,800-char block this returned 8 tokens for + # tool_result, thinking, document, mcp_tool_result — and for + # output_text / refusal, which are OpenAI's OWN Responses shapes. + # The shared walker is also image-safe: a 200KB base64 image gets + # a pixel-based 1600, not the ~50K phantom text tokens a naive + # str(block) catch-all would produce. + tokens += count_content_blocks(content, self.count_text) # Name field name = message.get("name") diff --git a/headroom/providers/openai_compatible.py b/headroom/providers/openai_compatible.py index 32c88db70..da0ebd3f4 100644 --- a/headroom/providers/openai_compatible.py +++ b/headroom/providers/openai_compatible.py @@ -24,6 +24,7 @@ from dataclasses import dataclass from typing import Any from headroom.tokenizers import get_tokenizer +from headroom.tokenizers.base import count_content_blocks from .base import Provider @@ -159,12 +160,16 @@ class OpenAICompatibleTokenCounter: if isinstance(content, str): tokens += self.count_text(content) elif isinstance(content, list): - for part in content: - if isinstance(part, dict): - if part.get("type") == "text": - tokens += self.count_text(part.get("text", "")) - elif isinstance(part, str): - tokens += self.count_text(part) + # Delegate to the audited shared walker instead of a partial + # per-provider one. Each provider counter had grown its own + # shortened branch list, so every modern block priced at ~0: + # measured on a 6,800-char block this returned 8 tokens for + # tool_result, thinking, document, mcp_tool_result — and for + # output_text / refusal, which are OpenAI's OWN Responses shapes. + # The shared walker is also image-safe: a 200KB base64 image gets + # a pixel-based 1600, not the ~50K phantom text tokens a naive + # str(block) catch-all would produce. + tokens += count_content_blocks(content, self.count_text) name = message.get("name") if name: diff --git a/headroom/tokenizers/base.py b/headroom/tokenizers/base.py index e90c5534c..fcd38ad13 100644 --- a/headroom/tokenizers/base.py +++ b/headroom/tokenizers/base.py @@ -8,6 +8,7 @@ from __future__ import annotations import json from abc import ABC, abstractmethod +from collections.abc import Callable from typing import Any, Protocol, runtime_checkable @@ -359,3 +360,39 @@ class BaseTokenizer(ABC): NotImplementedError: If decoding is not supported. """ raise NotImplementedError(f"{self.__class__.__name__} does not support decoding") + + +class _DelegatingBlockCounter(BaseTokenizer): + """Adapter exposing :meth:`BaseTokenizer._count_content_parts` to non-subclasses. + + The provider token counters in ``headroom/providers/`` are not + ``BaseTokenizer`` subclasses, and each grew its own shortened content-block + walker that handled only the shapes its provider was expected to send. The + result was that every one of them priced most modern blocks at ~0: measured + on a 6,800-char block, ``OpenAITokenCounter`` returned 8 tokens for + ``tool_result``/``thinking``/``document``/``mcp_tool_result`` and — its own + Responses shapes — ``output_text``/``refusal``; ``AnthropicTokenCounter`` + returned 7 for ``thinking``/``document``, which are Anthropic's own. + + Rather than add a fifth partial walker, this lets them borrow the audited one. + It is image-safe (base64 blobs get a pixel-based estimate instead of being + serialized and priced as text) and bounds oversized blobs, which a naive + ``count_text(str(block))`` catch-all does not. + """ + + __slots__ = ("_count_text_fn",) + + def __init__(self, count_text_fn: Callable[[str], int]) -> None: + self._count_text_fn = count_text_fn + + def count_text(self, text: str) -> int: + return self._count_text_fn(text) + + +def count_content_blocks(parts: list[Any], count_text_fn: Callable[[str], int]) -> int: + """Count a multi-part content list using *count_text_fn* for text. + + Shared entry point for provider counters. See + :class:`_DelegatingBlockCounter` for why this exists. + """ + return _DelegatingBlockCounter(count_text_fn)._count_content_parts(parts) diff --git a/tests/test_provider_counter_content_blocks.py b/tests/test_provider_counter_content_blocks.py new file mode 100644 index 000000000..c4ad57088 --- /dev/null +++ b/tests/test_provider_counter_content_blocks.py @@ -0,0 +1,106 @@ +"""Provider token counters must price every content block, not just ``text``. + +Each counter in ``headroom/providers/`` had grown its own shortened content-block +walker handling only the shapes its provider was expected to send. Everything else +fell through and contributed nothing. Measured on one 6,800-char block +(``count_messages`` of a single-block message, so 7-8 is message overhead alone): + + block type OpenAI ctr Anthropic ctr + text (control) 3409 3748 + tool_result 8 3748 + thinking 8 7 + document 8 7 + mcp_tool_result 8 7 + output_text 8 7 + refusal 8 7 + +Note each counter zeroed blocks from its OWN provider — ``output_text`` and +``refusal`` are OpenAI Responses shapes, ``thinking`` and ``document`` are +Anthropic's. And these counters are what the LIVE proxy pipelines use +(``proxy/server.py`` builds them with ``AnthropicProvider`` / ``OpenAIProvider``), +so this was the main request path, not an edge case. + +The counters now delegate to ``count_content_blocks``, which reuses +``BaseTokenizer._count_content_parts``. That matters beyond coverage: a naive +``count_text(str(block))`` catch-all would serialize a base64 image and price it +as text — a 1MB screenshot reads as ~330K phantom tokens. The shared walker gives +media a pixel/byte-based estimate instead. +""" + +from __future__ import annotations + +import pytest + +from headroom.providers.anthropic import AnthropicProvider +from headroom.providers.openai import OpenAIProvider +from headroom.tokenizers.base import count_content_blocks + +_BIG = "x " * 3400 # ~6,800 chars + + +def _counters(): + return { + "openai": OpenAIProvider().get_token_counter("gpt-4o"), + "anthropic": AnthropicProvider(warn=False).get_token_counter("claude-sonnet-4-6"), + } + + +@pytest.mark.parametrize( + "block", + [ + {"type": "tool_result", "tool_use_id": "t", "content": _BIG}, + {"type": "tool_use", "id": "t", "name": "grep", "input": {"pattern": _BIG}}, + {"type": "thinking", "thinking": _BIG}, + {"type": "document", "source": {"data": _BIG}}, + {"type": "mcp_tool_result", "content": _BIG}, + {"type": "output_text", "text": _BIG}, + {"type": "refusal", "refusal": _BIG}, + {"type": "search_result", "content": _BIG}, + ], + ids=lambda b: str(b.get("type")), +) +def test_no_provider_counter_prices_a_large_block_at_zero(block: dict) -> None: + """Every one of these returned 7-8 tokens — message overhead only.""" + message = {"role": "user", "content": [block]} + for name, counter in _counters().items(): + got = counter.count_messages([message]) + assert got > 1_000, f"{name} priced a ~6,800-char {block['type']} block at {got}" + + +def test_base64_media_is_not_priced_as_text() -> None: + """The reason a str(block) catch-all would have been the wrong fix.""" + image = {"type": "image", "source": {"type": "base64", "data": "A" * 200_000}} + message = {"role": "user", "content": [image]} + + for name, counter in _counters().items(): + got = counter.count_messages([message]) + # 200KB of base64 as text would be ~50,000 tokens; the pixel estimate is 1600. + assert got < 5_000, f"{name} priced a 200KB base64 image as text: {got}" + assert got > 1_000, f"{name} ignored a declared image entirely: {got}" + + +def test_plain_text_blocks_are_unchanged() -> None: + """The control: the shape both counters already handled must not move.""" + message = {"role": "user", "content": [{"type": "text", "text": _BIG}]} + for name, counter in _counters().items(): + text_form = {"role": "user", "content": _BIG} + block_form = counter.count_messages([message]) + # A text block and the equivalent string should agree closely. + assert abs(block_form - counter.count_messages([text_form])) <= 5, name + + +def test_shared_walker_ignores_non_block_parts() -> None: + """A bare int is not a block and must contribute nothing.""" + assert count_content_blocks([123], len) == 0 + assert count_content_blocks([], len) == 0 + + +def test_shared_walker_counts_nested_tool_result_blocks() -> None: + """A tool that returns blocks nests them; they must be walked, not serialized.""" + nested = { + "type": "tool_result", + "tool_use_id": "t", + "content": [{"type": "text", "text": _BIG}], + } + flat = {"type": "tool_result", "tool_use_id": "t", "content": _BIG} + assert abs(count_content_blocks([nested], len) - count_content_blocks([flat], len)) < 100 diff --git a/tests/test_providers/test_universal.py b/tests/test_providers/test_universal.py index 26190ed46..b840837fc 100644 --- a/tests/test_providers/test_universal.py +++ b/tests/test_providers/test_universal.py @@ -200,7 +200,18 @@ class TestOpenAICompatibleProvider: assert tokens == 55 assert total == 34 - def test_openai_compatible_token_counter_ignores_unhandled_content_shapes(self, monkeypatch): + def test_openai_compatible_token_counter_prices_declared_media(self, monkeypatch): + """An image block costs tokens; a non dict/str part still contributes none. + + This previously asserted that BOTH contribute 0 — i.e. it pinned the + defect. The counter handled only ``type == "text"``, so every other block + priced at ~0: measured on a 6,800-char block, tool_result / thinking / + document / mcp_tool_result all returned 8 tokens, overhead only. Counters + now delegate to the shared walker, which prices a declared image with the + pixel-based estimate (1600, the max after provider auto-resize) rather + than either ignoring it or serializing its base64 as text. + """ + class DummyTokenizer: def count_text(self, text: str) -> int: return len(text) @@ -211,8 +222,12 @@ class TestOpenAICompatibleProvider: ) counter = OpenAICompatibleProvider().get_token_counter("demo-model") + # Non-list, non-str content is still ignored. assert counter.count_message({"role": "user", "content": {}}) == 8 - assert counter.count_message({"role": "user", "content": [{"type": "image"}, 123]}) == 8 + # A bare int is not a block and still contributes nothing. + assert counter.count_message({"role": "user", "content": [123]}) == 8 + # A declared image is now priced instead of silently free. + assert counter.count_message({"role": "user", "content": [{"type": "image"}, 123]}) == 1608 def test_get_context_limit_prefix_output_buffer_and_partial_pricing(self): provider = OpenAICompatibleProvider(