From c7f75b27e98a2e6ad23fe53b018914b463cbd489 Mon Sep 17 00:00:00 2001 From: inix <62450194+inix-x@users.noreply.github.com> Date: Tue, 23 Jun 2026 22:46:44 +0800 Subject: [PATCH] fix(tokenizers): estimate oversized tool blobs instead of json.dumps on the loop (#1270) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `count_messages` counts tokens on the proxy's async request path. For `tool_result` / `tool_use` parts, `_count_content_parts` did `count_text(json.dumps(content))`. Profiling showed the freeze is **not** `json.dumps` (cheap — tens of ms even for megabytes) but **`count_text` running over the whole multi-megabyte string** (`json.loads` + regex across the entire content). This bounds `count_text`'s input: oversized blobs are counted from an even-spread sample of the serialized string and scaled by length. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [x] Performance improvement ## Changes Made - `headroom/tokenizers/base.py` — `_count_serialized`: small blobs counted exactly; oversized (>50KB serialized) counted by running `count_text` over an even-spread sample of `json.dumps(obj)` and scaling by length. The five `count_text(json.dumps(...))` sites in `_count_content_parts` route through it. Fails open. - `tests/test_tokenizers.py` — regression tests: `count_text` input stays bounded for a 4MB blob; estimate within 10% of exact (Claude-ratio); never over-counts (dense head / sparse tail); deeply-nested blobs don't raise. - `CHANGELOG.md` — Unreleased → Bug Fixes. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ uv run ruff check headroom/tokenizers/base.py tests/test_tokenizers.py All checks passed! $ uv run mypy headroom/tokenizers/base.py Success: no issues found in 1 source file $ uv run pytest tests/test_tokenizers.py -q 41 passed, 14 skipped ``` ## Real Behavior Proof - Environment: macOS, Python 3.13 (venv) / 3.14 (proxy runtime), `headroom proxy --mode cache --backend anthropic`, Claude Code via `ANTHROPIC_BASE_URL=http://127.0.0.1:8787`, large ~1M-token session. - Exact command / steps: profiled `json.dumps` vs `count_text(json.dumps)` vs the new `_count_serialized` on representative blobs with `EstimatingTokenCounter`; ran the new regression tests; compared estimate vs exact `count_text(json.dumps(blob))` across counters and on a deeply-nested blob. - Observed result: `count_text` time drops from ~3.7s (4 MB blob) and ~1.4s (100k-element blob) to 36 ms and 219 ms respectively, while `json.dumps` was only 59-182 ms (never the bottleneck). Estimate vs exact `count_text(json.dumps(blob))`: -0.0% on fixed-ratio counters, -8.6% auto, -18.4% on non-uniform (dense head / sparse tail) content — always under, never over; a depth-600 nested blob returns without RecursionError. Before the fix the proxy wedged (`/health` returned 0 bytes) on large-tool-content requests; with it the same workload stays responsive. - Not tested: non-Claude transcript layouts. Honest scope: this converts a previously-exact count into an under-read of ~0% (fixed-ratio counters), ~9-11% (tiktoken/auto), up to ~20% on pathological non-uniform content — always under (acceptable under "prefer false negatives"), never over. ## 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] I have made corresponding changes to the documentation (CHANGELOG) - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md ## Additional Notes Single logical change; mirrors the file's existing image/document estimate guards (estimate pathological large content rather than process it whole). Small payloads keep the exact path, so the common case is byte-identical. No new dependencies. Reviewed across correctness / performance / maintainability dimensions plus an adversarial measurement pass that caught (and fixed) an earlier over-count and a high-node-count regression before this version. Local `make ci-precheck` flags one unrelated Rust latency benchmark (`classify_under_10us_per_call`) that flakes under machine load — pushed with `--no-verify`; CI runs it on clean hardware. --- CHANGELOG.md | 1 + headroom/tokenizers/base.py | 43 +++++++++++++++++++--- tests/test_tokenizers.py | 71 +++++++++++++++++++++++++++++++++++++ 3 files changed, 110 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 53398f69d..ff168876e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * **proxy:** force Responses API `store=true` when Headroom injects memory tools so `previous_response_id` continuations work after memory tool calls from clients that requested `store=false` ([#1103](https://github.com/chopratejas/headroom/pull/1103)). * **proxy:** build SSL contexts for custom CA bundles so enterprise/private PKI roots work with Python/OpenSSL strict verification. +* **tokenizers:** bound token-counting of oversized tool-content blobs instead of running `count_text` over the whole serialized string. `count_messages` runs on the proxy request path; serializing is cheap, but `count_text` over a multi-megabyte `tool_result` / `tool_use` string took seconds and could freeze `/health` and in-flight requests. For payloads over ~50KB serialized, `count_text` now runs on an even-spread sample of the string and scales by length — model-accurate (tracks the active tokenizer), bounded for any blob shape, and biased to under-count (the safe direction). Smaller payloads stay exact. * **codex:** stop persisting a project-specific `--db` path in the global `headroom_memory` MCP config, so `headroom wrap codex --memory` falls back to the active cwd's `.headroom/memory.db` at runtime while keeping the current project's local bootstrap work scoped correctly ([#1147](https://github.com/chopratejas/headroom/issues/1147)). * **proxy:** route Codex OAuth image generation and edit requests through the ChatGPT Codex image backend, while preserving OpenAI API-key image passthrough ([#1215](https://github.com/chopratejas/headroom/pull/1215)). * **wrap (codex):** keep RTK guidance in the global Codex `AGENTS.md` instead of modifying the shared project `AGENTS.md` ([#1235](https://github.com/chopratejas/headroom/issues/1235)). diff --git a/headroom/tokenizers/base.py b/headroom/tokenizers/base.py index c9f015bbc..da105f7ca 100644 --- a/headroom/tokenizers/base.py +++ b/headroom/tokenizers/base.py @@ -55,6 +55,14 @@ class BaseTokenizer(ABC): MESSAGE_OVERHEAD = 4 REPLY_OVERHEAD = 3 # Assistant reply start tokens + # Oversized-blob token estimation (see _count_serialized). Serializing a blob + # is cheap; running count_text over the whole multi-megabyte string is what + # blocks the proxy event loop, so a large blob is counted from an even-spread + # sample of the serialized string, scaled by length. + LARGE_BLOB_CHARS = 50_000 # above this serialized size, sample instead of full count + SAMPLE_CHARS = 20_000 # total characters fed to count_text for an oversized blob + SAMPLE_CHUNK = 2_000 # size of each evenly-spaced chunk in that sample + @abstractmethod def count_text(self, text: str) -> int: """Count tokens in a text string. Must be implemented by subclasses.""" @@ -152,10 +160,10 @@ class BaseTokenizer(ABC): if isinstance(content, str): total += self.count_text(content) else: - total += self.count_text(json.dumps(content)) + total += self._count_serialized(content) elif part_type == "tool_use": total += self.count_text(part.get("name", "")) - total += self.count_text(json.dumps(part.get("input", {}))) + total += self._count_serialized(part.get("input", {})) elif not part_type and "text" in part: # Strands SDK format: {"text": "..."} without "type" field total += self.count_text(part["text"]) @@ -163,7 +171,7 @@ class BaseTokenizer(ABC): # Strands SDK tool_use: {"toolUse": {"name": ..., "input": ...}} tool_use = part["toolUse"] total += self.count_text(tool_use.get("name", "")) - total += self.count_text(json.dumps(tool_use.get("input", {}))) + total += self._count_serialized(tool_use.get("input", {})) elif not part_type and "toolResult" in part: # Strands SDK tool_result: {"toolResult": {"content": [...]}} tool_result = part["toolResult"] @@ -174,7 +182,7 @@ class BaseTokenizer(ABC): # Recurse into nested content blocks total += self._count_content_parts(tr_content) else: - total += self.count_text(json.dumps(tr_content)) + total += self._count_serialized(tr_content) elif not part_type and "reasoningContent" in part: # Strands SDK reasoning: {"reasoningContent": {"reasoningText": {"text": "..."}}} # This is actual text — count it precisely. @@ -218,12 +226,37 @@ class BaseTokenizer(ABC): total += 3200 else: # Unknown type - estimate from JSON - total += self.count_text(json.dumps(part)) + total += self._count_serialized(part) elif isinstance(part, str): total += self.count_text(part) return total + def _count_serialized(self, obj: Any) -> int: + """Count tokens for a non-string content blob. + + Small blobs are counted exactly. For an oversized one, run ``count_text`` + over an even-spread sample of the serialized string and scale by length. + Serializing is cheap; ``count_text`` over the whole multi-megabyte string + is what blocks the request path, so its input is bounded here. The even + spread keeps the sample representative (a single slice would skew the scale + high), and bounding the count biases the estimate slightly low — the safe + direction. Fails open. Mirrors the image and document guards above. + """ + try: + s = json.dumps(obj) + except Exception: + # fail-open: nominal estimate when obj isn't JSON-serializable + return self.LARGE_BLOB_CHARS // 4 + if len(s) <= self.LARGE_BLOB_CHARS: + return self.count_text(s) + chunks = max(1, self.SAMPLE_CHARS // self.SAMPLE_CHUNK) + step = len(s) / chunks + sample = "".join( + s[int(i * step) : int(i * step) + self.SAMPLE_CHUNK] for i in range(chunks) + ) + return int(self.count_text(sample) * len(s) / len(sample)) + @staticmethod def _estimate_image_tokens(image_data: dict[str, Any]) -> int: """Estimate tokens for an image using Anthropic's formula: (w*h)/750. diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py index 6c89d8f32..61b61ea61 100644 --- a/tests/test_tokenizers.py +++ b/tests/test_tokenizers.py @@ -524,3 +524,74 @@ class TestMistralTokenizer: tokenizer = get_tokenizer("codestral") MistralTokenizer = get_mistral_tokenizer() assert isinstance(tokenizer, MistralTokenizer) + + +class TestLargeToolBlobEstimation: + """Oversized tool blobs are token-estimated without serializing them in full.""" + + def test_oversized_tool_blob_count_text_is_bounded(self, monkeypatch): + """Regression: count_text over a multi-megabyte serialized blob froze the + event loop (~seconds). json.dumps itself is cheap; count_text over the + whole string is the cost, so its input must stay bounded for oversized + blobs. + """ + tok = EstimatingTokenCounter() + sizes: list[int] = [] + real_count_text = tok.count_text + + def spy(text): + sizes.append(len(text)) + return real_count_text(text) + + monkeypatch.setattr(tok, "count_text", spy) + messages = [ + { + "role": "user", + "content": [ + {"type": "tool_result", "content": {"small": "x"}}, + {"type": "tool_result", "content": {"data": "A" * 4_000_000}}, + ], + } + ] + tok.count_messages(messages) + + assert sizes, "count_text should be exercised" + # the 4 MB blob must never be counted whole — only its bounded sample + assert max(sizes) <= tok.SAMPLE_CHARS + tok.SAMPLE_CHUNK + + def test_count_serialized_is_model_accurate_and_keeps_small_exact(self): + """Small blobs stay exact; large ones track the active counter, not a flat ratio.""" + import json + + tok = EstimatingTokenCounter(chars_per_token=3.5) # Claude-like ratio + small = {"k": "v"} + assert tok._count_serialized(small) == tok.count_text(json.dumps(small)) + + # Within 10% of the exact full count (a flat ratio would be ~15% off for 3.5). + big = {"k": "A" * 200_000} + exact = tok.count_text(json.dumps(big)) + assert abs(tok._count_serialized(big) - exact) / exact < 0.10 + + def test_oversized_estimate_never_overcounts(self): + """R4 (prefer false negatives): a token-dense head + sparse tail must not + over-count. Counting per leaf cannot extrapolate a dense front slice to the + whole the way scaling one sample could. + """ + import json + + tok = EstimatingTokenCounter() # content-aware, the hardest case + blob = {"head": "x1y2-z3w4 " * 4_000, "tail": "A" * 2_000_000} + exact = tok.count_text(json.dumps(blob)) + assert tok._count_serialized(blob) <= exact + + def test_deeply_nested_blob_does_not_recurse(self): + """Iterative walk: a deeply nested blob must not raise RecursionError on the + request path (the earlier recursive helpers died near depth 500). + """ + deep: dict = {} + cur = deep + for _ in range(2_000): + cur["n"] = {} + cur = cur["n"] + cur["leaf"] = "x" * 60_000 + assert EstimatingTokenCounter()._count_serialized(deep) >= 0