From f03cc6d88b826c2752b20bdce944f9ad1e507e83 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Tue, 4 Aug 2026 11:31:52 -0700 Subject: [PATCH] fix(router): stop counting an image's base64 payload as suffix tokens (#2778) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `_netcost_message_tokens` walked block-list content itself and fell back to `str(block)` for anything that wasn't `text` or `tool_result` — on the stated assumption that such blocks *"rarely dominate a suffix"*. An `image` block is the exception that breaks it: `str()` embeds the whole base64 payload. ```text counted real over 512x512 PNG 20,034 349 57x 1092x1092 screenshot 100,034 1,589 63x 1568x1568 233,367 1,600 146x ``` **Why this changes behaviour, not just a number.** S is the cache-bust cost — the tokens re-written if message *j* is mutated. `apply()` builds it as a running suffix sum: ```python for j in range(num_messages - 1, -1, -1): netcost_suffix_tokens[j] = netcost_suffix_tokens[j + 1] + _netcost_message_tokens(...) ``` So one image inflates S for **every message before it**, and the break-even gate then declines to compress any of them. A single screenshot could switch off net-cost-gated compression for the whole earlier conversation — and screenshots are routine in agent sessions. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made Delegate block-list content to `tokenizers.base.count_content_blocks`, deleting the local walk. That counter already guards exactly this case — its comment reads *"1MB image = ~330K fake tokens without this"* — so this walk simply predated it. Beyond the raw fix, this removes a **second pricing rule**: the gate now values images the same way the tokenizer that computes `tokens_before`/`tokens_after` does (a flat 1600, "max after auto-resize"). Pricing images one way for the gate and another for the savings math is the same class of problem as #2761. Verified byte-identical on the shapes the old walk handled correctly: ```text old walk canonical text only 101 101 tool_result str 81 81 tool_result list 61 61 image only 100,034 1,600 mixed 100,036 1,602 ``` ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check` + `ruff format`) - [x] Type checking passes (`mypy`) - [x] New tests added for new functionality ### Test Output ```text $ pytest tests/test_netcost_suffix_image_tokens.py -q 8 passed $ git stash push headroom/ && pytest tests/test_netcost_suffix_image_tokens.py -q 4 failed, 4 passed # the 4 failures are the payload-scaling assertions; the 4 passes are the # text/tool_result/string shapes, included to prove delegation is behaviour-preserving ``` All netcost + content-router suites: ```text $ pytest tests/test_netcost_gate.py tests/test_content_router_*.py \ tests/test_transforms_content_router.py tests/test_netcost_suffix_image_tokens.py -q 126 passed ``` ```text $ ruff check headroom/transforms/content_router.py tests/... All checks passed! $ mypy headroom/transforms/content_router.py no new errors ``` Deferring the full suite to CI — no maturin/Rust core in this environment. ## One existing test rewritten — please look at this bit `test_netcost_gate.py::TestNetCostHelpers::test_message_tokens_block_list_beats_repr` fails under the fix, and I want to be explicit that I changed a test rather than bury it. It built its image block as `{"type": "image", "source": {"data": "x" * 500}}`. A 500-char stub is **cheaper than a single image's real token cost**, so `str()` over it looked harmless (~130 tokens) and its assertion `abs(helper - text_only) < text_only * 0.5` held. That unrepresentative fixture is precisely why the payload-scaling bug survived — the test named "beats repr" was passing on the one payload size where repr happens not to be catastrophic. Rewritten to use a realistic 200KB payload and to assert what actually matters: ```python assert helper >= text_only # text still counted in full assert helper - text_only <= 2000 # image cost is bounded, not payload-scaled assert helper < count_text(str(content)) / 10 # ...and far below repr ``` I checked this both ways, so it is a real test and not a rubber stamp: ```text old test + fixed code -> FAILS (it was pinning the defect) new test + main -> FAILS (it catches the real bug) new test + fixed code -> passes ``` ## Known limitation The canonical estimate is a flat 1600 per image regardless of dimensions, so a small icon is now over-charged (~1600 vs ~13 real) where repr would have charged ~200. I kept the flat constant deliberately: it is the value every other counter in the codebase uses, and introducing a third rule here to shave small-icon cost would recreate the inconsistency this PR removes. The error is bounded at 1600 tokens and biases the gate conservative, versus an unbounded 100K+ error before. --- headroom/transforms/content_router.py | 48 +++++------- tests/test_netcost_gate.py | 21 +++-- tests/test_netcost_suffix_image_tokens.py | 93 +++++++++++++++++++++++ 3 files changed, 125 insertions(+), 37 deletions(-) create mode 100644 tests/test_netcost_suffix_image_tokens.py diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 1ce7b0fca..242e0aaa2 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -60,6 +60,7 @@ from ..config import ( ) from ..parser import CCR_RETRIEVAL_MARKER_RE from ..tokenizer import Tokenizer +from ..tokenizers.base import count_content_blocks from ..tokenizers.estimator import EstimatingTokenCounter from . import mixed_content as _mixed_content from .base import Transform @@ -1176,42 +1177,29 @@ def _gain_bucket(gain: float) -> str: def _netcost_message_tokens(message: dict[str, Any], tokenizer: Tokenizer) -> int: """Token count of a message for net-cost suffix (S) estimation. - String content is counted directly. Anthropic block-list content is - counted by summing the text-bearing fields (``text`` blocks and - ``tool_result`` content) rather than stringifying the whole list, which - would count Python ``repr`` punctuation and type names and badly - miscount S — the value that drives the break-even gate decision. + String content is counted directly. Block-list content is delegated to the + canonical block counter, which knows how to price non-text blocks. + + This function used to walk the list itself and fall back to + ``str(block)`` for anything that was not ``text`` or ``tool_result``, on the + stated assumption that such blocks "rarely dominate a suffix". An ``image`` + block is the exception that breaks it: ``str()`` embeds the whole base64 + payload, so one screenshot counted ~100,000 tokens instead of ~1,600 + (57x-146x over, growing with image size). + + That mattered because S is the cache-bust cost — the tokens re-written if + message *j* is mutated — so an image inflated S for **every message before + it**, and the break-even gate then refused to compress any of them. + ``BaseTokenizer._count_content_parts`` already solves this (see its "1MB + image = ~330K fake tokens without this" guard); this walk simply predated + it. Delegating also means new block types are priced in one place. """ content = message.get("content", "") if isinstance(content, str): return tokenizer.count_text(content) if not isinstance(content, list): return tokenizer.count_text(str(content)) - total = 0 - for block in content: - if not isinstance(block, dict): - total += tokenizer.count_text(str(block)) - continue - block_type = block.get("type") - if block_type == "text": - total += tokenizer.count_text(str(block.get("text", ""))) - elif block_type == "tool_result": - tc = block.get("content", "") - if isinstance(tc, str): - total += tokenizer.count_text(tc) - elif isinstance(tc, list): - for sub in tc: - if isinstance(sub, dict) and sub.get("type") == "text": - total += tokenizer.count_text(str(sub.get("text", ""))) - else: - total += tokenizer.count_text(str(sub)) - else: - total += tokenizer.count_text(str(tc)) - else: - # Other blocks (image, tool_use input, …) — repr is a rough proxy - # but bounded; these rarely dominate a suffix. - total += tokenizer.count_text(str(block)) - return total + return count_content_blocks(content, tokenizer.count_text) class CompressionCache: diff --git a/tests/test_netcost_gate.py b/tests/test_netcost_gate.py index 0bb71e8df..8a070eafe 100644 --- a/tests/test_netcost_gate.py +++ b/tests/test_netcost_gate.py @@ -148,8 +148,13 @@ class TestNetCostHelpers: assert _gain_bucket(float("inf")) == "nan" def test_message_tokens_block_list_beats_repr(self, tokenizer): - # str(content) over a block list counts repr punctuation/type names; - # the block-aware helper counts only the text-bearing payload. + # str(content) over a block list embeds the whole base64 payload; the + # block-aware helper prices the image at its pixel cost instead. + # + # This used to use a 500-char stub image, which is *smaller* than a + # single image's real token cost -- so repr looked cheap and the + # payload-scaling bug stayed invisible. Use a realistically sized + # payload, which is what actually occurs (screenshots). from headroom.transforms.content_router import _netcost_message_tokens text = "word " * 200 @@ -157,15 +162,17 @@ class TestNetCostHelpers: "role": "user", "content": [ {"type": "text", "text": text}, - {"type": "image", "source": {"data": "x" * 500}}, + {"type": "image", "source": {"data": "x" * 200_000}}, ], } helper = _netcost_message_tokens(block_msg, tokenizer) text_only = tokenizer.count_text(text) - # Helper tracks the text payload closely; the image block adds only a - # small repr proxy, far less than stringifying the whole list. - assert abs(helper - text_only) < text_only * 0.5 - assert helper < tokenizer.count_text(str(block_msg["content"])) + # The text payload is still counted in full, and the image adds a + # bounded pixel-based cost rather than a payload-scaled one. + assert helper >= text_only + assert helper - text_only <= 2000 + # ...which is dramatically less than stringifying the whole list. + assert helper < tokenizer.count_text(str(block_msg["content"])) / 10 def test_message_tokens_tool_result_blocks(self, tokenizer): from headroom.transforms.content_router import _netcost_message_tokens diff --git a/tests/test_netcost_suffix_image_tokens.py b/tests/test_netcost_suffix_image_tokens.py new file mode 100644 index 000000000..1355c309d --- /dev/null +++ b/tests/test_netcost_suffix_image_tokens.py @@ -0,0 +1,93 @@ +"""An image must not inflate the net-cost suffix (S). + +``_netcost_message_tokens`` used to walk block-list content itself and fall back +to ``str(block)`` for anything that was not ``text`` or ``tool_result``, on the +stated assumption that such blocks "rarely dominate a suffix". An ``image`` block +breaks that assumption completely: ``str()`` embeds the whole base64 payload. + + 512x512 PNG 20,034 counted vs ~349 real 57x + 1092x1092 shot 100,034 counted vs ~1,589 real 63x + 1568x1568 233,367 counted vs ~1,600 real 146x + +S is the cache-bust cost -- the tokens re-written if message *j* is mutated -- so +an image inflated S for **every message before it**, and the break-even gate then +declined to compress any of them. A single screenshot could switch off net-cost +compression for the whole earlier conversation. +""" + +from __future__ import annotations + +import base64 + +import pytest + +from headroom.tokenizer import Tokenizer +from headroom.tokenizers import get_tokenizer +from headroom.transforms.content_router import _netcost_message_tokens + + +@pytest.fixture +def tok() -> Tokenizer: + return Tokenizer(get_tokenizer("claude-sonnet-4-6"), "claude-sonnet-4-6") + + +def _image_block(payload_bytes: int) -> dict: + data = base64.b64encode(b"\x89PNG" + b"\x00" * payload_bytes).decode() + return { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": data}, + } + + +@pytest.mark.parametrize("payload_bytes", [120_000, 600_000, 1_400_000]) +def test_image_is_not_counted_as_its_base64_payload(tok: Tokenizer, payload_bytes: int) -> None: + """Cost must not scale with the base64 length.""" + message = {"role": "user", "content": [_image_block(payload_bytes)]} + + counted = _netcost_message_tokens(message, tok) + + # Anthropic caps image cost around 1600 tokens; anything in the tens of + # thousands means the payload is being counted as text. + assert counted <= 2000, f"image counted as {counted:,} tokens" + + +def test_image_cost_does_not_grow_with_payload_size(tok: Tokenizer) -> None: + """A 12x larger payload must not cost ~12x more.""" + small = _netcost_message_tokens({"role": "user", "content": [_image_block(120_000)]}, tok) + large = _netcost_message_tokens({"role": "user", "content": [_image_block(1_400_000)]}, tok) + + assert large == small + + +@pytest.mark.parametrize( + "content", + [ + [{"type": "text", "text": "hello world " * 50}], + [{"type": "tool_result", "content": "result text " * 40}], + [{"type": "tool_result", "content": [{"type": "text", "text": "x " * 60}]}], + ], + ids=["text", "tool_result_str", "tool_result_list"], +) +def test_text_bearing_blocks_are_unchanged(tok: Tokenizer, content: list) -> None: + """Delegation must be behaviour-preserving for what already worked. + + These are the shapes the old local walk handled correctly; pinning them + keeps the delegation from quietly changing suffix sizes on normal traffic. + """ + counted = _netcost_message_tokens({"role": "user", "content": content}, tok) + text = "".join( + block.get("text", "") + or (block.get("content") if isinstance(block.get("content"), str) else "") + or "".join( + sub.get("text", "") for sub in (block.get("content") or []) if isinstance(sub, dict) + ) + for block in content + ) + + assert counted == pytest.approx(tok.count_text(text), abs=2) + + +def test_plain_string_content_still_counted(tok: Tokenizer) -> None: + message = {"role": "user", "content": "plain " * 20} + + assert _netcost_message_tokens(message, tok) == tok.count_text(message["content"])