fix(agent): stop double-counting api_content in the token estimator

`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.
This commit is contained in:
Israel Lot 2026-07-31 00:34:47 +00:00 committed by kshitij
parent fae0c4f5f4
commit e3bc517034
2 changed files with 85 additions and 46 deletions

View File

@ -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(

View File

@ -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: