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: