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