fix(tokenizer): coerce non-string tool_call fields before counting (#2801)

## Description `/v1/compress` returned HTTP 503 with an unhandled `TypeError` when a message carried a `tool_calls[].function.arguments` value that was not a string. `arguments` is a JSON *string* per the OpenAI spec, but OpenAI-compatible upstreams do emit `None` or a raw object there, and every token counter passed the value straight to `tiktoken.encode()`. Because the malformed message persists in conversation history, the failure was sticky: every later request replaying that history failed too, regardless of destination provider. Reported in #2782. The exact repro in that issue (`arguments: null`) no longer raises — `count_text` grew a falsy guard since 0.33.0 — but the root cause is still live for any *truthy* non-string, which I reproduced against all four counters on `main` before the fix. ## Type of Change - [x] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Documentation update - [ ] Refactor / internal cleanup ## Changes Made - `headroom/tokenizers/base.py`: new `coerce_countable_text()`. Strings pass through untouched, `None` counts as nothing, dict/list/tuple are JSON-serialized, anything else falls back to `str()`. The serialized form is capped at 200K chars so a malformed upstream can't turn a token *estimate* into a multi-megabyte encode. - Applied at the tool-call field sites (`function.name`, `function.arguments`, `id`, and the legacy `function_call`) in `tokenizers/base.py`, `tokenizers/tiktoken_counter.py`, `providers/openai.py`, `providers/openai_compatible.py`, `providers/anthropic.py`. - Guarded `{"function": null}` / `{"id": null}`, which reach the same encode path. - New test file `tests/test_tool_call_arguments_not_a_string.py` (11 cases). Serializing dicts rather than the one-liner suggested in the issue (`str(func.get("arguments") or "")`) is deliberate: `str()` on a dict yields Python repr with single quotes, which is not what the upstream would have billed, and it is unbounded. ## Testing - [x] Existing tests pass - [x] New tests added for the fix - [ ] Manual testing performed New tests: ``` $ python -m pytest tests/test_tool_call_arguments_not_a_string.py -q tests\test_tool_call_arguments_not_a_string.py ........... [100%] ============================= 11 passed in 0.40s ============================== ``` Surrounding tokenizer/provider suites, unchanged: ``` $ python -m pytest tests/test_tool_call_arguments_not_a_string.py tests/test_tokenizer.py \ tests/test_tokenizers.py tests/test_tokenizers \ tests/test_provider_counter_content_blocks.py tests/test_provider_tokenizer_one_ruler.py -q tests\test_provider_counter_content_blocks.py ............ [ 91%] tests\test_provider_tokenizer_one_ruler.py ......... [100%] ======================= 91 passed, 14 skipped in 1.50s ======================== ``` Lint/format on the touched files: ``` $ ruff check <touched files> && ruff format --check <touched files> All checks passed! 6 files already formatted ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.13.11, pytest 9.1.1, repo at `upstream/main` (d0a86d40) with this branch applied; `headroom._core` built locally. - Exact command / steps: ran the same script before and after the change, driving the four counters directly (the crash site the proxy 503 unwinds to): ``` python -c " from headroom.providers.openai import OpenAITokenCounter from headroom.providers.openai_compatible import OpenAICompatibleTokenCounter from headroom.providers.anthropic import AnthropicTokenCounter from headroom.tokenizers.tiktoken_counter import TiktokenCounter def mk(a): return [{'role':'assistant','content':None,'tool_calls':[ {'id':'c1','type':'function','function':{'name':'read_file','arguments':a}}]}] for name,c in [('openai',OpenAITokenCounter('gpt-4o')), ('compat',OpenAICompatibleTokenCounter('gpt-4o')), ('anthropic',AnthropicTokenCounter('claude-sonnet-4')), ('tiktoken',TiktokenCounter('gpt-4o'))]: for a in [None, {'path':'x'}, 5]: try: print(name, repr(a), c.count_messages(mk(a))) except Exception as e: print(name, repr(a), 'ERR', type(e).__name__, e) " ``` - Observed result: before the change, all four counters raised on every truthy non-string; after, all return finite counts and an object `arguments` prices within 5 tokens of its JSON string form. ``` BEFORE openai None 21 openai {'path': 'x'} ERR TypeError expected string or buffer openai 5 ERR TypeError expected string or buffer compat {'path': 'x'} ERR TypeError expected string or buffer anthropic {'path': 'x'} ERR TypeError expected string or buffer tiktoken {'path': 'x'} ERR TypeError expected string or buffer AFTER openai None 21 {'path': 'x'} 27 5 22 '{"path":"x"}' 26 compat None 20 {'path': 'x'} 26 5 21 '{"path":"x"}' 25 anthropic None 9 {'path': 'x'} 15 5 10 '{"path":"x"}' 14 tiktoken None 14 {'path': 'x'} 20 5 15 '{"path":"x"}' 19 ``` - Not tested: I did not exercise a live `headroom proxy --mode cache` + `curl /v1/compress` round trip, nor a real OpenAI-compatible upstream that emits object `arguments`. The proxy path was verified only down to the counters that its traceback terminates in, plus the automated tests above. Also untested: `providers/google.py`, `cohere.py`, `litellm.py`, which have no tool-call counting branch and so were left alone. ## Review Readiness - [x] I have performed a self-review - [x] I have commented my code where the reasoning is not obvious - [x] This PR is ready for human review Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-08-05 17:32:14 +02:00
parent d0a86d409f
commit b6f9877c78
6 changed files with 162 additions and 30 deletions
@@ -24,7 +24,7 @@ import warnings
 from typing import Any, cast

 from headroom import paths as _paths
-from headroom.tokenizers.base import count_content_blocks
+from headroom.tokenizers.base import coerce_countable_text, count_content_blocks

 from .base import Provider, TokenCounter

@@ -411,9 +411,9 @@ class AnthropicTokenCounter(TokenCounter):
        if "tool_calls" in message:
            for tool_call in message.get("tool_calls", []):
                if isinstance(tool_call, dict):
-                    func = tool_call.get("function", {})
-                    tokens += self.count_text(func.get("name", ""))
-                    tokens += self.count_text(func.get("arguments", ""))
+                    func = tool_call.get("function") or {}
+                    tokens += self.count_text(coerce_countable_text(func.get("name")))
+                    tokens += self.count_text(coerce_countable_text(func.get("arguments")))

        return tokens

@@ -16,7 +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 headroom.tokenizers.base import coerce_countable_text, count_content_blocks

 from .base import Provider, TokenCounter

@@ -403,10 +403,10 @@ class OpenAITokenCounter:
        tool_calls = message.get("tool_calls")
        if tool_calls:
            for tc in tool_calls:
-                func = tc.get("function", {})
-                tokens += self.count_text(func.get("name", ""))
-                tokens += self.count_text(func.get("arguments", ""))
-                tokens += self.count_text(tc.get("id", ""))
+                func = tc.get("function") or {}
+                tokens += self.count_text(coerce_countable_text(func.get("name")))
+                tokens += self.count_text(coerce_countable_text(func.get("arguments")))
+                tokens += self.count_text(coerce_countable_text(tc.get("id")))
                tokens += 10  # Structural overhead

        # Tool call ID for tool responses
@@ -24,7 +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 headroom.tokenizers.base import coerce_countable_text, count_content_blocks

 from .base import Provider

@@ -178,9 +178,9 @@ class OpenAICompatibleTokenCounter:
        tool_calls = message.get("tool_calls")
        if tool_calls:
            for tc in tool_calls:
-                func = tc.get("function", {})
-                tokens += self.count_text(func.get("name", ""))
-                tokens += self.count_text(func.get("arguments", ""))
+                func = tc.get("function") or {}
+                tokens += self.count_text(coerce_countable_text(func.get("name")))
+                tokens += self.count_text(coerce_countable_text(func.get("arguments")))
                tokens += 10

        tool_call_id = message.get("tool_call_id")
@@ -11,6 +11,40 @@ from abc import ABC, abstractmethod
 from collections.abc import Callable
 from typing import Any, Protocol, runtime_checkable

+#: Cap on the serialized form of a non-string tool-call field. A malformed
+#: upstream can put an arbitrarily large object where a JSON string belongs;
+#: serializing it unbounded would turn a token *estimate* into a multi-megabyte
+#: encode. Truncating keeps the estimate finite and the request alive.
+_MAX_COERCED_FIELD_CHARS = 200_000
+
+
+def coerce_countable_text(value: Any) -> str:
+    """Return *value* as text safe to pass to ``count_text``.
+
+    Tool-call fields (``function.name``, ``function.arguments``, ``id``) are
+    strings per the OpenAI spec, but OpenAI-compatible upstreams do emit
+    ``None`` or a raw object there. Passing those straight to
+    ``tiktoken.encode()`` raises ``TypeError: expected string or buffer``, which
+    failed the whole compression with a 503 — and kept failing on every later
+    request, because the malformed message stays in conversation history.
+
+    ``None`` counts as nothing; dicts/lists are JSON-serialized so an object
+    ``arguments`` is priced roughly like the JSON string it should have been;
+    anything else falls back to ``str()``.
+    """
+    if isinstance(value, str):
+        return value
+    if value is None:
+        return ""
+    if isinstance(value, dict | list | tuple):
+        try:
+            text = json.dumps(value, ensure_ascii=False, default=str)
+        except (TypeError, ValueError):
+            text = str(value)
+    else:
+        text = str(value)
+    return text[:_MAX_COERCED_FIELD_CHARS]
+

@runtime_checkable
 class TokenCounter(Protocol):
@@ -311,20 +345,20 @@ class BaseTokenizer(ABC):
            total += 4  # Tool call overhead

            if "function" in call:
-                func = call["function"]
-                total += self.count_text(func.get("name", ""))
-                total += self.count_text(func.get("arguments", ""))
+                func = call["function"] or {}
+                total += self.count_text(coerce_countable_text(func.get("name")))
+                total += self.count_text(coerce_countable_text(func.get("arguments")))

            if "id" in call:
-                total += self.count_text(call["id"])
+                total += self.count_text(coerce_countable_text(call["id"]))

        return total

    def _count_function_call(self, function_call: dict[str, Any]) -> int:
        """Count tokens in legacy function call."""
        total = 4  # Function call overhead
-        total += self.count_text(function_call.get("name", ""))
-        total += self.count_text(function_call.get("arguments", ""))
+        total += self.count_text(coerce_countable_text(function_call.get("name")))
+        total += self.count_text(coerce_countable_text(function_call.get("arguments")))
        return total

    def encode(self, text: str) -> list[int]:
@@ -16,7 +16,7 @@ import threading
 from functools import lru_cache
 from typing import Any

-from .base import BaseTokenizer
+from .base import BaseTokenizer, coerce_countable_text

 logger = logging.getLogger(__name__)

@@ -316,24 +316,25 @@ class TiktokenCounter(BaseTokenizer):
                            elif isinstance(part, str):
                                total += self.count_text(part)
                elif key == "role":
-                    total += self.count_text(value)
+                    total += self.count_text(coerce_countable_text(value))
                elif key == "name":
-                    total += self.count_text(value)
+                    total += self.count_text(coerce_countable_text(value))
                    total += 1  # Name adds 1 token
                elif key == "tool_calls":
-                    for tool_call in value:
+                    for tool_call in value or []:
                        total += 3  # Tool call overhead
                        if "function" in tool_call:
-                            func = tool_call["function"]
-                            total += self.count_text(func.get("name", ""))
-                            total += self.count_text(func.get("arguments", ""))
+                            func = tool_call["function"] or {}
+                            total += self.count_text(coerce_countable_text(func.get("name")))
+                            total += self.count_text(coerce_countable_text(func.get("arguments")))
                        if "id" in tool_call:
-                            total += self.count_text(tool_call["id"])
+                            total += self.count_text(coerce_countable_text(tool_call["id"]))
                elif key == "tool_call_id":
-                    total += self.count_text(value)
+                    total += self.count_text(coerce_countable_text(value))
                elif key == "function_call":
-                    total += self.count_text(value.get("name", ""))
-                    total += self.count_text(value.get("arguments", ""))
+                    value = value or {}
+                    total += self.count_text(coerce_countable_text(value.get("name")))
+                    total += self.count_text(coerce_countable_text(value.get("arguments")))

        # Every reply is primed with assistant
        total += self.REPLY_OVERHEAD
@@ -0,0 +1,97 @@
+"""Token counters must survive a tool_call whose fields aren't strings (GH #2782).
+
+``function.arguments`` is a JSON *string* per the OpenAI spec, but
+OpenAI-compatible upstreams do emit ``None`` or a raw object there. Every counter
+passed the value straight to ``tiktoken.encode()``, which raises
+``TypeError: expected string or buffer`` — so ``/v1/compress`` failed the whole
+request with a 503. Worse, the malformed message stays in conversation history,
+so every later request replaying that history failed too, regardless of provider.
+
+``arguments: None`` stopped raising once ``count_text`` grew its falsy guard, but
+any *truthy* non-string (``{"path": "x"}``, ``5``) still crashed all four
+counters. The fix is ``coerce_countable_text`` at the tool-call field sites, so a
+dict is priced roughly like the JSON string it should have been rather than
+either crashing or silently counting as zero.
+"""
+
+from __future__ import annotations
+
+import json
+
+import pytest
+
+from headroom.providers.anthropic import AnthropicProvider
+from headroom.providers.openai import OpenAITokenCounter
+from headroom.providers.openai_compatible import OpenAICompatibleTokenCounter
+from headroom.tokenizers.base import coerce_countable_text
+from headroom.tokenizers.tiktoken_counter import TiktokenCounter
+
+
+def _counters():
+    return {
+        "openai": OpenAITokenCounter("gpt-4o"),
+        "openai_compatible": OpenAICompatibleTokenCounter("gpt-4o"),
+        "anthropic": AnthropicProvider(warn=False).get_token_counter("claude-sonnet-4-6"),
+        "tiktoken": TiktokenCounter("gpt-4o"),
+    }
+
+
+def _message(arguments: object) -> dict:
+    return {
+        "role": "assistant",
+        "content": None,
+        "tool_calls": [
+            {
+                "id": "call_1",
+                "type": "function",
+                "function": {"name": "read_file", "arguments": arguments},
+            }
+        ],
+    }
+
+
+@pytest.mark.parametrize(
+    "arguments",
+    [None, {"path": "x"}, ["a", "b"], 5, True, 1.5],
+    ids=["none", "dict", "list", "int", "bool", "float"],
+)
+def test_non_string_arguments_do_not_raise(arguments: object) -> None:
+    """The reported crash: 503 + TypeError out of tiktoken.encode."""
+    for name, counter in _counters().items():
+        got = counter.count_messages([_message(arguments)])
+        assert got > 0, f"{name} priced the whole message at {got}"
+
+
+def test_object_arguments_are_priced_like_their_json_form() -> None:
+    """Not just non-crashing: a dict must not silently count as zero."""
+    payload = {"path": "src/very/long/path/to/a/file.py", "start": 1, "end": 400}
+    for name, counter in _counters().items():
+        as_object = counter.count_messages([_message(payload)])
+        as_json = counter.count_messages([_message(json.dumps(payload))])
+        assert abs(as_object - as_json) <= 5, f"{name}: {as_object} vs {as_json}"
+
+
+def test_null_function_and_id_do_not_raise() -> None:
+    """``{"function": null}`` / ``{"id": null}`` reach the same encode path."""
+    message = {"role": "assistant", "tool_calls": [{"id": None, "function": None}]}
+    for name, counter in _counters().items():
+        assert counter.count_messages([message]) > 0, name
+
+
+def test_legacy_function_call_with_object_arguments_does_not_raise() -> None:
+    message = {"role": "assistant", "function_call": {"name": "f", "arguments": {"a": 1}}}
+    for name, counter in _counters().items():
+        assert counter.count_messages([message]) > 0, name
+
+
+def test_oversized_object_arguments_are_bounded() -> None:
+    """A malformed upstream must not turn an estimate into a megabyte encode."""
+    huge = {"blob": "x" * 5_000_000}
+    assert len(coerce_countable_text(huge)) <= 200_000
+
+
+def test_string_arguments_are_untouched() -> None:
+    """The control: the spec-compliant shape must not move."""
+    args = json.dumps({"path": "a.py"})
+    assert coerce_countable_text(args) == args
+    assert coerce_countable_text(None) == ""