From e3bc51703469bbb26e8d1e9b7b5ec6e26176c364 Mon Sep 17 00:00:00 2001 From: Israel Lot Date: Fri, 31 Jul 2026 00:34:47 +0000 Subject: [PATCH] fix(agent): stop double-counting api_content in the token estimator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `api_content` is a SUBSTITUTE for `content`, not an addition to it. `turn_context.substitute_api_content()` pops the sidecar and overwrites `content` at every API-bound message-build site (the `api_messages` build in `conversation_loop`, the max-iterations summary in `chat_completion_helpers`, the chat-completions transport), so exactly one of the two is ever sent to the provider. The preflight estimator counted both, because both `_estimate_message_chars` and `_estimate_message_tokens_without_images` walked every key of the persisted dict with a single-entry denylist (`_anthropic_content_blocks`). Any message whose sidecar differs from its clean stored content was counted twice — exactly 2.00x on a 40KB sidecar. The sidecar exists to keep the provider prompt-cache prefix byte-stable, so it is written on precisely the long, cache-pinned messages where the doubling hurts most. Because `estimate_messages_tokens_rough()` also feeds the compaction threshold via `context_compressor` and `conversation_loop`, the inflated estimate makes compression fire on phantom bytes. Fix: substitute rather than sum, mirroring the wire. The two estimator helpers had drifted into near-identical copies of the same shadow-building loop, so this factors the shared logic into `_wire_message_shadow()` and fixes the class once instead of patching one site and leaving the other. Image accounting is unchanged: base64 payloads are still replaced with a placeholder and charged at the flat `_count_image_tokens` rate, and the `_multimodal` text_summary path is preserved. Tests: three cases in `TestEstimateMessagesTokensRough` — sidecar equal to content is counted once, a sidecar that DIFFERS is still counted (a lower bound, so it fails if the field were dropped rather than substituted, which would undercount the real request), and a sidecar cannot smuggle raw base64 past the flat image rate. Verified on Linux (Python 3.11): 53 passed in tests/agent/test_model_metadata.py, 57 passed with tests/agent/test_context_breakdown.py, 656 passed / 3 skipped across the compression/context/token/estimate/prune surface of tests/agent. Mutation-tested: reverting the substitution fails the new equality test. `scripts/check-windows-footguns.py` is not applicable — no file I/O, process management, terminal handling, subprocesses, or signals. --- agent/model_metadata.py | 94 +++++++++++++++--------------- tests/agent/test_model_metadata.py | 37 ++++++++++++ 2 files changed, 85 insertions(+), 46 deletions(-) diff --git a/agent/model_metadata.py b/agent/model_metadata.py index 9fc6150fdc0b0..ea8dc220b3137 100644 --- a/agent/model_metadata.py +++ b/agent/model_metadata.py @@ -2982,6 +2982,52 @@ def _count_image_tokens(msg: Dict[str, Any], cost_per_image: int) -> int: return count * cost_per_image +def _wire_message_shadow(msg: Dict[str, Any]) -> Dict[str, Any]: + """Shadow of a message holding only what the provider actually receives. + + Two adjustments to the raw persisted dict: + + * ``api_content`` is a SUBSTITUTE for ``content``, not an addition to it. + ``turn_context.substitute_api_content()`` pops the sidecar and overwrites + ``content`` at every API-bound build site, so exactly one of the two is + ever sent. Counting both double-counts any message whose sidecar differs + from its clean stored content (2.00x on a 40KB sidecar). + * Base64 image payloads are replaced with a placeholder; they are charged + separately at a flat rate by ``_count_image_tokens``, and counting their + raw chars here would massively overestimate usage. + """ + shadow: Dict[str, Any] = {} + for k, v in msg.items(): + if k == "_anthropic_content_blocks": + continue + if k == "api_content": + shadow["content"] = v + continue + if k == "content": + if "api_content" in msg: + # The sidecar wins on the wire; skip the clean copy so the + # same logical content is not counted twice. + continue + if isinstance(v, list): + cleaned = [] + for part in v: + if isinstance(part, dict): + if part.get("type") in {"image", "image_url", "input_image"}: + cleaned.append({"type": part.get("type"), "image": "[stripped]"}) + else: + cleaned.append(part) + else: + cleaned.append(part) + shadow[k] = cleaned + elif isinstance(v, dict) and v.get("_multimodal"): + shadow[k] = v.get("text_summary", "") + else: + shadow[k] = v + else: + shadow[k] = v + return shadow + + def _estimate_message_chars(msg: Dict[str, Any]) -> int: """Char count for token estimation, excluding base64 image data. @@ -2990,58 +3036,14 @@ def _estimate_message_chars(msg: Dict[str, Any]) -> int: """ if not isinstance(msg, dict): return len(str(msg)) - shadow: Dict[str, Any] = {} - for k, v in msg.items(): - if k == "_anthropic_content_blocks": - continue - if k == "content": - if isinstance(v, list): - cleaned = [] - for part in v: - if isinstance(part, dict): - if part.get("type") in {"image", "image_url", "input_image"}: - cleaned.append({"type": part.get("type"), "image": "[stripped]"}) - else: - cleaned.append(part) - else: - cleaned.append(part) - shadow[k] = cleaned - elif isinstance(v, dict) and v.get("_multimodal"): - shadow[k] = v.get("text_summary", "") - else: - shadow[k] = v - else: - shadow[k] = v - return len(str(shadow)) + return len(str(_wire_message_shadow(msg))) def _estimate_message_tokens_without_images(msg: Dict[str, Any]) -> int: """Token estimate for a message shadow with image payloads stripped.""" if not isinstance(msg, dict): return estimate_tokens_rough(str(msg)) - shadow: Dict[str, Any] = {} - for k, v in msg.items(): - if k == "_anthropic_content_blocks": - continue - if k == "content": - if isinstance(v, list): - cleaned = [] - for part in v: - if isinstance(part, dict): - if part.get("type") in {"image", "image_url", "input_image"}: - cleaned.append({"type": part.get("type"), "image": "[stripped]"}) - else: - cleaned.append(part) - else: - cleaned.append(part) - shadow[k] = cleaned - elif isinstance(v, dict) and v.get("_multimodal"): - shadow[k] = v.get("text_summary", "") - else: - shadow[k] = v - else: - shadow[k] = v - return estimate_tokens_rough(str(shadow)) + return estimate_tokens_rough(str(_wire_message_shadow(msg))) def estimate_request_tokens_rough( diff --git a/tests/agent/test_model_metadata.py b/tests/agent/test_model_metadata.py index c4442ebe0abda..febcd24106f95 100644 --- a/tests/agent/test_model_metadata.py +++ b/tests/agent/test_model_metadata.py @@ -79,6 +79,43 @@ class TestEstimateMessagesTokensRough: # string representation. assert 1500 <= result < 2000 + def test_api_content_substitutes_for_content_not_added_to_it(self): + """``api_content`` replaces ``content`` on the wire, so count one. + + ``turn_context.substitute_api_content()`` pops the sidecar and + overwrites ``content`` at every API-bound build site. Counting both + doubled the estimate for any message carrying a sidecar. + """ + body = "cached prompt bytes " * 2000 + wire_shape = {"role": "user", "content": body} + persisted_shape = {"role": "user", "content": body, "api_content": body} + + assert estimate_messages_tokens_rough([persisted_shape]) == \ + estimate_messages_tokens_rough([wire_shape]) + + def test_api_content_is_counted_when_it_differs_from_content(self): + """The sidecar is what's sent, so its size is the one that matters.""" + big_sidecar = "cached prompt bytes " * 2000 + msg = {"role": "user", "content": "short", "api_content": big_sidecar} + + result = estimate_messages_tokens_rough([msg]) + + # Lower bound: fails if the sidecar were dropped rather than + # substituted (which would undercount the real request). + assert result >= (len(big_sidecar) // 4) * 0.9 + + def test_api_content_does_not_defeat_image_stripping(self): + """A sidecar must not smuggle raw base64 past the flat image rate.""" + import base64 + import os + + payload = "data:image/png;base64," + base64.b64encode(os.urandom(300_000)).decode() + msg = {"role": "user", + "content": [{"type": "image_url", "image_url": {"url": payload}}]} + + # Raw base64 would be ~100K tokens; the flat per-image model is ~1.5K. + assert estimate_messages_tokens_rough([msg]) < 5_000 + class TestEstimateRequestTokensRough: