From 5920bb69da757387c99f380699e016788cb083aa Mon Sep 17 00:00:00 2001 From: redshift <213178690+1ftredsh@users.noreply.github.com> Date: Sun, 6 Sep 2026 12:15:13 +0200 Subject: [PATCH] fix(payment): wire Responses input images into reservation, fix original-detail fallbacks Address review on PR #680: - calculate_discounted_max_cost now calls estimate_image_tokens_from_input on body["input"], so a Responses request carrying an input_image reserves image tokens (previously the estimator was defined but never called, reserving 0). - _estimate_input_image_tokens now fetches and measures remote original-detail images instead of blindly reserving the 36,000-token worst case (~117x over-reserve for a 512x512 image), falling back to the worst case only when the fetch fails or the file is a file_id reference. - Broken/undecodable data URLs now reserve the declared detail's worst case (36,000 for original) rather than the 85-token low-detail floor: base64 is validated and PIL failures fall back explicitly instead of silently using _get_image_dimensions' 512x512 default. - Restore test_discount_cannot_be_dodged_by_hiding_prompt_in_tools (dropped in the rebase) and add regression tests for the Responses input-image wiring, remote original-detail fetch, and broken-data-URL fallback. --- routstr/payment/helpers.py | 99 ++++++++++----- tests/unit/test_payment_helpers.py | 198 +++++++++++++++++++++++++++++ 2 files changed, 264 insertions(+), 33 deletions(-) diff --git a/routstr/payment/helpers.py b/routstr/payment/helpers.py index 9cc34e56..05304684 100644 --- a/routstr/payment/helpers.py +++ b/routstr/payment/helpers.py @@ -230,17 +230,24 @@ async def calculate_discounted_max_cost( # for work the reservation never covered. prompt_tokens = estimate_prompt_tokens(body) + # Images are billed as tokens by the upstream but carry no text for + # ``estimate_prompt_tokens`` to count, so they are estimated separately and + # added on both the chat (``messages``) and Responses (``input``) paths. + image_tokens = 0 if isinstance(messages, list): - image_tokens = await estimate_image_tokens_in_messages(messages) - if image_tokens > 0: - logger.debug( - "Found images in request", - extra={ - "model": model, - "image_tokens": image_tokens, - }, - ) - prompt_tokens += image_tokens + image_tokens += await estimate_image_tokens_in_messages(messages) + input_data = body.get("input") + if input_data is not None: + image_tokens += await estimate_image_tokens_from_input(input_data) + if image_tokens > 0: + logger.debug( + "Found images in request", + extra={ + "model": model, + "image_tokens": image_tokens, + }, + ) + prompt_tokens += image_tokens if prompt_tokens > 0: estimated_prompt_delta_sats = ( @@ -629,34 +636,60 @@ async def estimate_image_tokens_in_messages(messages: list) -> int: async def _estimate_input_image_tokens(item: dict) -> int: """Estimate tokens for a Responses API ``input_image`` item. - Honors the item-level ``detail``. The dimensions of ``file_id`` - references can't be fetched here, so they get conservative - estimates: the max-size tile math for high/auto and the 30,000-patch - worst case (36,000 tokens) for original, so we never under-reserve. + Honors the item-level ``detail``. Data-URL images are measured from their + decoded bytes; remote URLs are fetched and measured like the chat path. + Only ``file_id`` references (whose dimensions cannot be fetched here) and + unfetchable/broken images fall back to conservative estimates: the + max-size tile math for high/auto and the 30,000-patch worst case (36,000 + tokens) for original, so we never under-reserve. """ detail = item.get("detail") or "auto" - if image_url := item.get("image_url"): - if isinstance(image_url, dict): - image_url = image_url.get("url", "") - if isinstance(image_url, str) and image_url.startswith("data:image/"): - try: - _, base64_data = image_url.split(",", 1) - image_bytes = base64.b64decode(base64_data) - width, height = _get_image_dimensions(image_bytes) - return _calculate_image_tokens(width, height, detail) - except Exception as e: - logger.warning( - "Failed to process base64 image", extra={"error": str(e)} - ) - return 85 - # Remote URLs and file_id both have unfetchable dimensions here; fall - # through to the conservative estimates below. - if item.get("file_id") or item.get("image_url"): + image_url = item.get("image_url") + if isinstance(image_url, dict): + image_url = image_url.get("url", "") + + def _worst_case() -> int: + # Dimensions unknown: reserve the worst case for the declared detail so + # a broken/unreadable image still covers what the upstream could bill. if detail == "original": return _MAX_ORIGINAL_IMAGE_TOKENS - # We can't fetch an uploaded file's dimensions here; assume the - # largest vision image so we don't under-reserve. return _calculate_image_tokens(2048, 2048, detail) + + if isinstance(image_url, str) and image_url: + if image_url.startswith("data:image/"): + image_bytes = None + try: + _, base64_data = image_url.split(",", 1) + image_bytes = base64.b64decode(base64_data, validate=True) + except Exception as e: + logger.warning( + "Failed to decode base64 image", extra={"error": str(e)} + ) + if image_bytes is not None: + try: + img = Image.open(BytesIO(image_bytes)) + return _calculate_image_tokens(img.size[0], img.size[1], detail) + except Exception as e: + logger.warning( + "Failed to read image dimensions", extra={"error": str(e)} + ) + # Undecodable / unreadable data URL: reserve the worst case. + return _worst_case() + # Remote URL: fetch and measure like the chat path so a small image + # does not reserve the original-detail worst case. + image_bytes = await _fetch_image_from_url(image_url) + if image_bytes: + try: + img = Image.open(BytesIO(image_bytes)) + return _calculate_image_tokens(img.size[0], img.size[1], detail) + except Exception as e: + logger.warning( + "Failed to read image dimensions", extra={"error": str(e)} + ) + # Unfetchable or unreadable: fall through to the conservative estimate. + + if item.get("file_id") or image_url: + return _worst_case() return 0 diff --git a/tests/unit/test_payment_helpers.py b/tests/unit/test_payment_helpers.py index 391f8f1f..84ab8959 100644 --- a/tests/unit/test_payment_helpers.py +++ b/tests/unit/test_payment_helpers.py @@ -211,6 +211,204 @@ async def test_discount_counts_legacy_token_id_prompt() -> None: assert cost == 50_000 +async def test_discount_cannot_be_dodged_by_hiding_prompt_in_tools() -> None: + """A large prompt moved from messages into tool schemas must reserve the + same cost — otherwise a caller undercharges by hiding weight from the + estimator.""" + from routstr.payment.helpers import calculate_discounted_max_cost + + pricing = Mock() + pricing.prompt = 0.5 + pricing.completion = 0.01 + pricing.max_prompt_cost = 100.0 + pricing.max_completion_cost = 100.0 + + model_obj = Mock() + model_obj.sats_pricing = pricing + model_obj.top_provider = None + model_obj.context_length = None + + big_text = "word " * 2_000 + base = {"model": "test-model", "max_tokens": 10} + in_messages = { + **base, + "messages": [{"role": "user", "content": big_text}], + } + hiding_places = { + "tools": { + **base, + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + {"type": "function", "function": {"name": "f", "description": big_text}} + ], + }, + # Anthropic forwards a top-level system prompt; it is billed like any other. + "system": { + **base, + "messages": [{"role": "user", "content": "hi"}], + "system": big_text, + }, + # A key named like an image field must not win an image exclusion. + "image-named key": { + **base, + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"function": {"parameters": {"data": big_text}}}], + }, + # Nor may a caller-chosen "data:" prefix, in any field the body allows. + "data-prefixed content": { + **base, + "messages": [{"role": "user", "content": "data:" + big_text}], + }, + "data-prefixed text block": { + **base, + "messages": [ + { + "role": "user", + "content": [{"type": "text", "text": "data:" + big_text}], + } + ], + }, + "data-prefixed system": { + **base, + "messages": [{"role": "user", "content": "hi"}], + "system": "data:" + big_text, + }, + } + + with ( + patch.object(settings, "fixed_pricing", False), + patch.object(settings, "tolerance_percentage", 0), + patch.object(settings, "min_request_msat", 1000), + ): + cost_messages = await calculate_discounted_max_cost( + 150_000, in_messages, model_obj + ) + for where, body in hiding_places.items(): + cost = await calculate_discounted_max_cost(150_000, body, model_obj) + # Same prompt weight → at least the same reservation, never the floor. + assert cost >= cost_messages, where + assert cost > 1000, where + + +async def test_discounted_max_cost_counts_responses_input_images() -> None: + """A Responses ``input_image`` must add image tokens to the reservation. + + Regression for the review finding that ``estimate_image_tokens_from_input`` + was defined but never called: a Responses body carries ``input``, not + ``messages``, so its images were previously reserved at zero tokens. + """ + import base64 + from io import BytesIO + + from PIL import Image + + from routstr.payment.helpers import calculate_discounted_max_cost + + pricing = Mock() + pricing.prompt = 0.001 + pricing.completion = 0.001 + pricing.max_prompt_cost = 100.0 + pricing.max_completion_cost = 0.0 + + model_obj = Mock() + model_obj.sats_pricing = pricing + model_obj.top_provider = None + model_obj.context_length = None + + image = Image.new("RGB", (512, 512), "red") + buffer = BytesIO() + image.save(buffer, format="JPEG") + data_url = "data:image/jpeg;base64," + base64.b64encode(buffer.getvalue()).decode() + + no_image = { + "model": "test-model", + "input": [{"role": "user", "content": "hi"}], + } + with_image = { + "model": "test-model", + "input": [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "hi"}, + {"type": "input_image", "image_url": data_url, "detail": "high"}, + ], + } + ], + } + + with ( + patch.object(settings, "fixed_pricing", False), + patch.object(settings, "tolerance_percentage", 0), + patch.object(settings, "min_request_msat", 1000), + ): + cost_no_image = await calculate_discounted_max_cost(100_000, no_image, model_obj) + cost_with_image = await calculate_discounted_max_cost( + 100_000, with_image, model_obj + ) + + # The 512x512 high-detail image (85 + 170 = 255 tokens) is billed as prompt + # weight, so it reserves strictly more than the identical text-only body. + assert cost_with_image > cost_no_image + + +async def test_estimate_input_image_tokens_remote_original_fetches() -> None: + """A remote ``original`` image is fetched and measured, not worst-cased. + + Regression for the review finding that any non-data URL with + ``detail: \"original\"`` reserved the 36,000-token worst case without + trying to fetch — a 512x512 image reserved ~117x its real cost. + """ + from io import BytesIO + from unittest.mock import patch as mock_patch + + from PIL import Image + + from routstr.payment.helpers import _estimate_input_image_tokens + + image = Image.new("RGB", (512, 512), "red") + buffer = BytesIO() + image.save(buffer, format="JPEG") + image_bytes = buffer.getvalue() + + with mock_patch( + "routstr.payment.helpers._fetch_image_from_url", + new=AsyncMock(return_value=image_bytes), + ): + # 512x512 original -> 16x16 = 256 patches -> ceil(256 * 1.2) = 308 tokens, + # far below the 36,000 worst case a blind fallback would reserve. + assert await _estimate_input_image_tokens( + {"type": "input_image", "image_url": "https://x.test/i.jpg", "detail": "original"} + ) == 308 + + # When the fetch fails, fall back to the original-detail worst case. + with mock_patch( + "routstr.payment.helpers._fetch_image_from_url", + new=AsyncMock(return_value=None), + ): + assert await _estimate_input_image_tokens( + {"type": "input_image", "image_url": "https://x.test/i.jpg", "detail": "original"} + ) == 36_000 + + +async def test_estimate_input_image_tokens_broken_original_data_url() -> None: + """A broken data URL with ``detail: \"original\"`` reserves the worst case. + + Regression for the review finding that the ``except`` branch returned 85 + (low-detail) regardless of the declared detail. + """ + from routstr.payment.helpers import _estimate_input_image_tokens + + # "!!!" is not valid base64, so decoding raises before any dimension read. + assert await _estimate_input_image_tokens( + {"type": "input_image", "image_url": "data:image/jpeg;base64,!!!", "detail": "original"} + ) == 36_000 + # Non-original details fall back to the max-size tile math, not the 85 floor. + assert await _estimate_input_image_tokens( + {"type": "input_image", "image_url": "data:image/jpeg;base64,!!!", "detail": "high"} + ) == 85 + (170 * 4) + + async def test_discounted_max_cost_body_max_output_tokens_fallback() -> None: """Body ``max_output_tokens`` (Responses API) is honored as a completion cap.""" from routstr.payment.helpers import calculate_discounted_max_cost