headroom/tests/test_content_router_token_units.py
Tejas Chopra ad56dd382b
fix(router): compare token quantities in one unit (#2759)
## Description

Two places compared a token quantity against something measured in a
**different unit**. Both changed compression **behaviour**, not just
reporting — which is the worse class.

### 1. The CONFIG branch put a word count in a token ratio

`compressed_tokens = len(compressed.split())` was divided by
`original_tokens`, which comes from `_estimate_tokens(content)`. Words
run ~2.8× fewer than estimator tokens on config text, so a compressor
that returned its input **byte-identically** scored ~0.36.

`min_ratio` is 1.0 — accept any real shrink — so the router **accepted
the no-op**: cached the result, pinned a frozen "compress" verdict,
emitted a `router:config_compressor` label into `transforms_applied`,
and recorded a fabricated saving to TOIN.

```text
mkdocs.yml, compressor returns its input unchanged
  denominator (_estimate_tokens)   = 936
  OLD numerator len(split())       = 334  -> ratio 0.357  claims 64% saved  ACCEPTED
  NEW numerator (_estimate_tokens) = 936  -> ratio 1.000  correctly rejected
```

The sibling TABULAR branch already used `_estimate_tokens`; CONFIG was
the outlier.

### 2. The Kompress size gate tested a token cap in chars/4

`len(text_to_compress) > self._kompress_max_tokens * 4` under-counts
anything denser than 4 chars/token, and compact JSON runs ~3.2. Against
the 50,000-token default there is a band where an oversized payload
passes:

```text
records  chars     old_gate(len/4)  new_gate(tokens)
  2700   119,281   False            False
  4000   177,781   False            True    <- 44,445 vs 55,557 tokens, 11% over
  4400   195,781   False            True    <- 48,945 vs 61,182 tokens, 22% over
  5000   222,781   True             True
```

Those payloads entered ONNX inference — exactly the >30s non-preemptible
worker stall the gate exists to prevent (#1171).

Closes #

## Type of Change

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

## Changes Made

- `content_router.py` CONFIG branch — `compressed_tokens` now from
`_estimate_tokens(compressed)`, matching its denominator and every
sibling branch.
- `content_router.py` Kompress gate — compared with `_estimate_tokens`,
the unit the cap is actually expressed in. The extra O(n) char scan is
negligible against the inference it guards.

## 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_content_router_token_units.py -q
4 passed in 0.33s

$ uvx ruff@0.15.17 check headroom/transforms/content_router.py tests/test_content_router_token_units.py
All checks passed!
```

**Regression check against clean `upstream/main` in the same
environment:**

```text
tests/test_transforms/ + test_transforms_content_router.py + kompress suites
  upstream/main : 2 failed, 468 passed, 78 skipped
  this branch   : 2 failed, 468 passed, 78 skipped
  failure sets  : identical
```

The 2 failures are `test_kompress_failsafe`'s artifact-selection tests,
which need a real `onnxruntime` this throwaway env lacks. Unrelated to
this change.

The 4 new tests pin the unit contract *and* the bounds of the
disagreement band — including the cases where both formulations agree,
so the band is demonstrated rather than assumed.

## Real Behavior Proof

- **Environment:** macOS 26.4 arm64, isolated worktree off
`upstream/main`. `content_router` needs the compiled `headroom._core`,
which isn't in a fresh worktree (gitignored, built in-place); I copied
the built `.so` in to run these, then removed it before committing.
- **Observed:** both tables above are from running the real
`_estimate_tokens` against the real thresholds, not reconstructed
arithmetic.

- **Not tested:** no live ONNX inference — the >30s stall the gate
prevents is cited from #1171, not reproduced. The CONFIG no-op was
demonstrated at the ratio level rather than by driving a stubbed
compressor through `apply()`.

## 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

Third of three PRs from one tokenizer-consistency audit — see #2757
(litellm total-prompt / `--budget`) and #2758 (HuggingFace chat
templates, `gpt-5`, gateway-wrapped names). Separate subsystems,
separate risk.

Known remaining from the same audit, not in any of the three:
`_netcost_message_tokens` pricing an image by Python `repr` (34×
over-count, flag-gated), three transforms reporting via
`count_text(str(content))` where the pipeline uses `count_messages` (19%
apart in one log file), `frozen_message_count` walking a chars/3.5
estimate against provider-reported cached tokens, and `target_ratio`
honoured in words while documented as tokens.

🤖 Generated with [Claude Code](https://claude.com/claude-code)
2026-08-03 22:49:47 -07:00

84 lines
3.2 KiB
Python

"""Token counts compared inside ContentRouter must share one unit.
Two places measured a token quantity in a different unit from the thing it was
compared against. Both changed COMPRESSION BEHAVIOUR, not just reporting:
1. The CONFIG branch built its accept ratio as
``len(compressed.split()) / _estimate_tokens(content)`` — a WORD count over a
TOKEN count. Words run ~2.8x fewer than estimator tokens on config text, so a
compressor returning its input byte-identically scored ~0.36. ``min_ratio`` is
1.0 (accept any real shrink), so the router ACCEPTED the no-op, cached it,
froze the verdict, emitted a ``router:config_compressor`` label and recorded a
fabricated ~64% saving.
2. The Kompress size gate tested ``len(text) > max_tokens * 4`` — a TOKEN cap
evaluated in chars/4. Dense payloads run denser than 4 chars/token (compact
JSON ~3.2), so a band existed where an oversized payload passed the gate into
the >30s non-preemptible ONNX inference the gate exists to prevent (#1171).
"""
from __future__ import annotations
import json
from headroom.transforms.content_router import _estimate_tokens
_CAP = 50_000
def _compact_json(records: int) -> str:
return json.dumps(
[{"id": i, "status": "ok", "msg": f"value_{i}"} for i in range(records)],
separators=(",", ":"),
)
def test_dense_payload_over_the_cap_is_caught_by_the_token_unit() -> None:
"""chars/4 waves through a payload the token cap should stop.
177,781 chars of compact JSON: 44,445 by chars/4 (under the cap, silent) but
55,557 estimator tokens — 11% over.
"""
payload = _compact_json(4_000)
assert len(payload) // 4 <= _CAP, "premise: the old chars/4 test passes this"
assert _estimate_tokens(payload) > _CAP, "the token cap must be exceeded"
def test_the_two_units_agree_below_and_above_the_band() -> None:
"""Outside the disagreement band both formulations reach the same verdict."""
small = _compact_json(2_700)
assert len(small) // 4 <= _CAP and _estimate_tokens(small) <= _CAP
huge = _compact_json(5_000)
assert len(huge) // 4 > _CAP and _estimate_tokens(huge) > _CAP
def test_chars_over_four_understates_dense_content() -> None:
"""The mechanism, stated as a property rather than a magic number."""
payload = _compact_json(4_000)
assert _estimate_tokens(payload) > len(payload) // 4, (
"compact JSON is denser than 4 chars/token, which is why the unit matters"
)
def test_word_count_is_not_interchangeable_with_a_token_count() -> None:
"""Why the CONFIG numerator had to change.
A byte-identical no-op must score 1.0. Measured with a word count it scored
~0.36 on real config text and was accepted as a 64% saving.
"""
config_text = "\n".join(
f"key_{i}: value_{i} # inline comment explaining key_{i}" for i in range(200)
)
tokens = _estimate_tokens(config_text)
words = len(config_text.split())
# Same unit on both sides: a no-op is correctly a no-op.
assert tokens / tokens == 1.0
# Mixed units: the same no-op looks like a large saving.
assert words / tokens < 0.75, (
"if words and estimator tokens were interchangeable this bug could not exist"
)