test(e2e): derive goldens from first-principles rate selection

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/e2e/cost_calculation/cases.json

@@ -62,7 +62,15 @@
       "name": "web_search",
       "usage": {"fresh_input_tokens": 100, "output_tokens": 30, "web_search_calls": 3},
       "requires_rates": ["search_context_cost_per_query"],
-      "requires_caps": ["web_search"]
+      "requires_caps": ["web_search"],
+      "wires": ["openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate"]
+    },
+    {
+      "name": "web_search_single",
+      "usage": {"fresh_input_tokens": 100, "output_tokens": 30, "web_search_calls": 1},
+      "requires_rates": ["search_context_cost_per_query"],
+      "requires_caps": ["web_search"],
+      "wires": ["openai_chat", "together_chat", "fireworks_chat", "azure_chat"]
     },
     {
       "name": "stream",

tests/e2e/cost_calculation/expected.json

@@ -160,7 +160,7 @@
     "prompt_tokens": 120,
     "spend": 0.032
   },
-  "azure/gpt-5.4-mini|web_search": {
+  "azure/gpt-5.4-mini|web_search_single": {
     "completion_tokens": 30,
     "input_cost": 0.016,
     "output_cost": 0.009600000000000001,
@@ -272,7 +272,7 @@
     "prompt_tokens": 120,
     "spend": 0.03
   },
-  "azure/gpt-5.6|web_search": {
+  "azure/gpt-5.6|web_search_single": {
     "completion_tokens": 30,
     "input_cost": 0.015,
     "output_cost": 0.009,
@@ -615,7 +615,7 @@
     "prompt_tokens": 120,
     "spend": 0.028000000000000004
   },
-  "fireworks_ai/deepseek-v4p1-flash|web_search": {
+  "fireworks_ai/deepseek-v4p1-flash|web_search_single": {
     "completion_tokens": 30,
     "input_cost": 0.014000000000000002,
     "output_cost": 0.008400000000000001,
@@ -706,7 +706,7 @@
     "prompt_tokens": 120,
     "spend": 0.024
   },
-  "fireworks_ai/kimi-k3|web_search": {
+  "fireworks_ai/kimi-k3|web_search_single": {
     "completion_tokens": 30,
     "input_cost": 0.012000000000000002,
     "output_cost": 0.007200000000000001,
@@ -797,7 +797,7 @@
     "prompt_tokens": 120,
     "spend": 0.026000000000000002
   },
-  "fireworks_ai/qwen3p8-max|web_search": {
+  "fireworks_ai/qwen3p8-max|web_search_single": {
     "completion_tokens": 30,
     "input_cost": 0.013000000000000001,
     "output_cost": 0.007800000000000001,
@@ -1462,7 +1462,7 @@
     "prompt_tokens": 120,
     "spend": 0.008
   },
-  "gpt-5.4-mini|web_search": {
+  "gpt-5.4-mini|web_search_single": {
     "completion_tokens": 30,
     "input_cost": 0.004,
     "output_cost": 0.0024000000000000002,
@@ -1679,7 +1679,7 @@
     "prompt_tokens": 120,
     "spend": 0.002
   },
-  "gpt-5.6|web_search": {
+  "gpt-5.6|web_search_single": {
     "completion_tokens": 30,
     "input_cost": 0.001,
     "output_cost": 0.0006000000000000001,
@@ -1826,7 +1826,7 @@
     "prompt_tokens": 120,
     "spend": 0.02
   },
-  "together_ai/moonshotai/Kimi-K3|web_search": {
+  "together_ai/moonshotai/Kimi-K3|web_search_single": {
     "completion_tokens": 30,
     "input_cost": 0.01,
     "output_cost": 0.006,
@@ -1938,7 +1938,7 @@
     "prompt_tokens": 120,
     "spend": 0.022
   },
-  "together_ai/zai-org/GLM-5.3|web_search": {
+  "together_ai/zai-org/GLM-5.3|web_search_single": {
     "completion_tokens": 30,
     "input_cost": 0.011000000000000001,
     "output_cost": 0.0066,

tests/e2e/cost_calculation/generate_expected.py

@@ -19,8 +19,6 @@ from dataclasses import dataclass
 from types import MappingProxyType
 from typing import Final
 
-from pydantic import TypeAdapter
-
 from cost_matrix import (
     EXPECTED_PATH,
     FRONTIER_MODELS,
@@ -32,19 +30,11 @@ from cost_matrix import (
     cases_for,
     expected_key,
 )
-
-# Wires whose response surface reports a real web-search call count; the
-# chat-completions wires only expose url_citation annotations, so their billed
-# count floors to one.
-_EXACT_WEB_SEARCH_WIRES: Final = frozenset(
-    {"openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate"}
-)
+from pydantic import TypeAdapter
 
 
-def billed_web_search_calls(model: FrontierModel, case: Case) -> int:
-    if case.usage.web_search_calls == 0:
-        return 0
-    return case.usage.web_search_calls if model.wire in _EXACT_WEB_SEARCH_WIRES else 1
+def _first_present(*rates: float | None) -> float | None:
+    return next((rate for rate in rates if rate is not None), None)
 
 
 @dataclass(frozen=True, slots=True)
@@ -67,11 +57,14 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost:
 
     Input = fresh*in + read*read + 5m*create + 1h*create_1h + audio_in*audio_in;
     output = text*out + reasoning*reasoning + audio_out*audio_out; plus the
-    billed web-search calls at the medium search-context rate. Above-threshold
-    swaps every input/output rate to its ``_above_200k_tokens`` variant when
-    total prompt tokens exceed the threshold; a service tier swaps input/output
-    to the tier's variants, falling back to the base rate when a variant is
-    unset -- mirroring _get_token_base_cost in litellm's cost calculator.
+    billed web-search calls at the medium search-context rate. Every billed
+    token is a token the provider charged for: a component whose entry has no
+    dedicated rate bills at the ordinary input or output rate, and a present
+    rate (including an explicit 0.0) is authoritative. When the total prompt
+    tokens exceed the threshold, input/output rates come from the
+    ``_above_200k_tokens`` variants; a service tier takes its ``_priority`` or
+    ``_flex`` variant when the entry carries one, and otherwise bills at the
+    base rate.
     """
     rates: Final[CostMapEntry] = model.override_rates if case.response_model_override else model.rates
     u: Final = case.usage
@@ -81,46 +74,53 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost:
     )
     tiered: Final = prompt_tokens > TIER_THRESHOLD_TOKENS
     in_rate: Final = (
-        (rates.input_cost_per_token_above_200k_tokens if tiered else None)
-        or (rates.input_cost_per_token_priority if case.service_tier == "priority" else None)
-        or (rates.input_cost_per_token_flex if case.service_tier == "flex" else None)
-        or rates.input_cost_per_token
+        _first_present(
+            rates.input_cost_per_token_above_200k_tokens if tiered else None,
+            rates.input_cost_per_token_priority if case.service_tier == "priority" else None,
+            rates.input_cost_per_token_flex if case.service_tier == "flex" else None,
+            rates.input_cost_per_token,
+        )
         or 0.0
     )
     out_rate: Final = (
-        (rates.output_cost_per_token_above_200k_tokens if tiered else None)
-        or (rates.output_cost_per_token_priority if case.service_tier == "priority" else None)
-        or (rates.output_cost_per_token_flex if case.service_tier == "flex" else None)
-        or rates.output_cost_per_token
+        _first_present(
+            rates.output_cost_per_token_above_200k_tokens if tiered else None,
+            rates.output_cost_per_token_priority if case.service_tier == "priority" else None,
+            rates.output_cost_per_token_flex if case.service_tier == "flex" else None,
+            rates.output_cost_per_token,
+        )
         or 0.0
     )
-    write_rate: Final = (
-        rates.cache_creation_input_token_cost
-        if rates.cache_creation_input_token_cost is not None
-        else in_rate
+    read_rate: Final = _first_present(rates.cache_read_input_token_cost, in_rate) or 0.0
+    write_rate: Final = _first_present(rates.cache_creation_input_token_cost, in_rate) or 0.0
+    write_1h_rate: Final = (
+        _first_present(rates.cache_creation_input_token_cost_above_1hr, write_rate) or 0.0
     )
+    audio_in_rate: Final = _first_present(rates.input_cost_per_audio_token, in_rate) or 0.0
+    reasoning_rate: Final = _first_present(rates.output_cost_per_reasoning_token, out_rate) or 0.0
+    audio_out_rate: Final = _first_present(rates.output_cost_per_audio_token, out_rate) or 0.0
     input_cost: Final = (
         u.fresh_input_tokens * in_rate
-        + u.cache_read_tokens
-        * (rates.cache_read_input_token_cost if rates.cache_read_input_token_cost is not None else in_rate)
+        + u.cache_read_tokens * read_rate
         + u.cache_write_5m_tokens * write_rate
-        + u.cache_write_1h_tokens
-        * (
-            rates.cache_creation_input_token_cost_above_1hr
-            if rates.cache_creation_input_token_cost_above_1hr is not None
-            else write_rate
-        )
-        + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0)
+        + u.cache_write_1h_tokens * write_1h_rate
+        + u.audio_input_tokens * audio_in_rate
     )
     output_cost: Final = (
         u.output_tokens * out_rate
-        + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate)
-        + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate)
+        + u.reasoning_tokens * reasoning_rate
+        + u.audio_output_tokens * audio_out_rate
     )
     search: Final = rates.search_context_cost_per_query
-    tool_cost: Final = billed_web_search_calls(model, case) * (
-        search.search_context_size_medium if search and search.search_context_size_medium else 0.0
+    medium_rate: Final = (
+        search.search_context_size_medium if search is not None else None
     )
+    if u.web_search_calls and medium_rate is None:
+        raise ValueError(
+            f"{model.map_key}: case {case.name} bills {u.web_search_calls} web-search "
+            "calls but the entry has no search_context_cost_per_query medium rate"
+        )
+    tool_cost: Final = u.web_search_calls * (medium_rate if medium_rate is not None else 0.0)
     return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost)