Merge pull request #41328 from BerriAI/litellm_e2e_cost_calculation_scripted_provider

test(integration): literal request/response cost tracking suite on a test-owned cost map

.circleci/config.yml

@@ -3015,7 +3015,7 @@ workflows:
           name: integration-<< matrix.suite >>
           matrix:
             parameters:
-              suite: [management, accounting, database, providers, extensions, sdk, browser]
+              suite: [management, accounting, database, providers, extensions, sdk, cost, browser]
           filters:
             branches:
               only:

.circleci/scripts/run_integration.sh

@@ -9,6 +9,7 @@ fi
 suite="${1:?integration suite required}"
 results="test-results/integration-${suite}"
 mkdir -p "$results"
+shard_timeout=11m
 integration_identity="$(.venv/bin/python -c 'import uuid; print(uuid.uuid4().hex)')"
 upstream_pid=""
 proxy_pid=""
@@ -108,13 +109,26 @@ awk '$3 == "REJECT" && $1 > 0 { rejected=1 } END { exit !rejected }' "$results/e
 setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
   .venv/bin/python -m integration._support.upstream > "$results/upstream.log" 2>&1 &
 upstream_pid=$!
+if [ "$suite" = cost ]; then
+  export INTEGRATION_WORKERS=8
+fi
 start_proxy() {
   local port="$1"
   local log_name="$2"
+  local -a cost_map_env
+  if [ "$suite" = cost ]; then
+    cost_map_env=(
+      "LITELLM_MODEL_COST_MAP_URL=$INTEGRATION_UPSTREAM_URL/_cost_map"
+      "MODEL_COST_MAP_MIN_MODEL_COUNT=1"
+      "MODEL_COST_MAP_MAX_SHRINK_RATIO=0"
+    )
+  else
+    cost_map_env=("LITELLM_LOCAL_MODEL_COST_MAP=True")
+  fi
   setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
     DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
     LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
-    LITELLM_MODE=PRODUCTION LITELLM_LOCAL_MODEL_COST_MAP=True STORE_MODEL_IN_DB=True \
+    LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True "${cost_map_env[@]}" \
     AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
     .venv/bin/python -m integration._support.proxy --config tests/integration/proxy_config.yaml \
     --host 127.0.0.1 --port "$port" --num_workers 1 --telemetry False \
@@ -158,11 +172,12 @@ if [ "$suite" = browser ]; then
   exit 0
 fi
 
-timeout --signal=TERM --kill-after=20s 11m env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
+timeout --signal=TERM --kill-after=20s "$shard_timeout" env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
   INTEGRATION_RUN_ID="$integration_identity" \
   DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
   INTEGRATION_PROXY_URL="$INTEGRATION_PROXY_URL" INTEGRATION_PEER_URL="$INTEGRATION_PEER_URL" \
   INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
+  INTEGRATION_WORKERS="${INTEGRATION_WORKERS:-1}" \
   INTEGRATION_MASTER_KEY="$INTEGRATION_MASTER_KEY" LITELLM_MODE=PRODUCTION \
   INTEGRATION_SEED="$INTEGRATION_SEED" \
   INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \

tests/integration/README.md

@@ -2,7 +2,9 @@
 
 These tests exercise a running gateway, PostgreSQL and Redis with an owned local upstream. CircleCI owns this suite. Tests are grouped by behavior, with no automatic test retries or fallback to paid provider calls
 
-Use `tests/integration/run.py management`, `accounting`, `database`, `providers`, `extensions` or `sdk` to run a selected group. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate
+The `cost` group is driven by `cost_tracking_cases.json`, which contains the cost map, literal requests, literal provider responses and expected accounting values. Each case has a name, contract ID, cost-map model, optional deployment overrides, request body, tagged response and exact or recount expectations. Request bodies use `$MODEL` for the registered proxy model, while responses use `$REQUEST_ID` for the per-run scenario ID. To add a case, add a cost-map entry when the model is new, add the request body and exact provider response data, add hand-computed expected values and register the node ID in `contracts.json`. The upstream serves each stored response for any path under `/<scenario_id>`, while the test-owned cost map is served over loopback through `LITELLM_MODEL_COST_MAP_URL`
+
+Use `tests/integration/run.py management`, `accounting`, `database`, `providers`, `extensions`, `sdk` or `cost` to run a selected group. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate
 
 Management also requires `INTEGRATION_PEER_URL`, `REDIS_HOST` and `REDIS_PORT`. CircleCI starts two directly addressed proxy processes sharing only that job's stores. The test-only CLI wrapper supplies enterprise route entitlement, following the existing behavior suite's convention. It does not qualify license validation; run it with one worker and no reload
 

tests/integration/_support/manifest.py

@@ -20,6 +20,7 @@ OWNED_DIRECTORIES: Final = frozenset(
         "observability",
         "compatibility",
         "sdk",
+        "cost_calculation",
     }
 )
 

tests/integration/_support/upstream.py

@@ -1,21 +1,35 @@
 from __future__ import annotations
 
 import argparse
-from dataclasses import dataclass, field
 from collections import deque
+from collections.abc import Mapping
+import json
+from dataclasses import dataclass, field
+import os
+from pathlib import Path
 from queue import SimpleQueue
-from typing import Final
+import struct
+from typing import Final, cast
+import zlib
 
+import httpx
 import uvicorn
-from pydantic import JsonValue, TypeAdapter
+from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
 from starlette.applications import Starlette
 from starlette.requests import Request
 from starlette.responses import JSONResponse, Response
 from starlette.routing import Route
 
 from _fake_openai_endpoint_server import chat_completions, completions, embeddings, health, moderations
+from integration.cost_calculation.cost_tracking_case import (
+    EventStreamResponse,
+    JsonResponse,
+    SseResponse,
+    StoredResponse,
+)
 
 JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+CASES_FILE: Final = Path(__file__).resolve().parents[1] / "cost_calculation" / "cost_tracking_cases.json"
 INTERNAL_FIELDS: Final = frozenset(
     {
         "litellm_params",
@@ -44,10 +58,58 @@ class Observation:
     body: dict[str, JsonValue]
 
 
+class _ScenarioRegistration(BaseModel):
+    scenario_id: str
+    response: StoredResponse
+
+
+def _aws_str_header(name: str, value: str) -> bytes:
+    name_bytes: Final = name.encode()
+    value_bytes: Final = value.encode()
+    return (
+        struct.pack("!B", len(name_bytes))
+        + name_bytes
+        + struct.pack("!B", 7)
+        + struct.pack("!H", len(value_bytes))
+        + value_bytes
+    )
+
+
+def _aws_event_frame(event_type: str, payload: Mapping[str, JsonValue], scenario_id: str) -> bytes:
+    payload_bytes: Final = json.dumps(payload, separators=(",", ":")).replace(
+        "$REQUEST_ID", scenario_id
+    ).encode()
+    headers_bytes: Final = (
+        _aws_str_header(":event-type", event_type)
+        + _aws_str_header(":content-type", "application/json")
+        + _aws_str_header(":message-type", "event")
+    )
+    total_length: Final = 12 + len(headers_bytes) + len(payload_bytes) + 4
+    prelude: Final = struct.pack("!II", total_length, len(headers_bytes))
+    prelude_crc: Final = struct.pack("!I", zlib.crc32(prelude) & 0xFFFFFFFF)
+    message: Final = prelude + prelude_crc + headers_bytes + payload_bytes
+    return message + struct.pack("!I", zlib.crc32(message) & 0xFFFFFFFF)
+
+
+class ScenarioStore:
+    def __init__(self) -> None:
+        self._scenarios: dict[str, StoredResponse] = {}
+
+    def put(self, scenario_id: str, response: StoredResponse) -> None:
+        self._scenarios[scenario_id] = response
+
+    def drop(self, scenario_id: str) -> bool:
+        return self._scenarios.pop(scenario_id, None) is not None
+
+    def get(self, scenario_id: str) -> StoredResponse | None:
+        return self._scenarios.get(scenario_id)
+
+
 @dataclass(frozen=True, slots=True)
 class Provider:
     observations: SimpleQueue[Observation] = field(default_factory=SimpleQueue)
     scripts: dict[str, deque[int]] = field(default_factory=dict)
+    scenario_store: ScenarioStore = field(default_factory=ScenarioStore)
 
     async def chat(self, request: Request) -> Response:
         body: Final = JSON_OBJECT.validate_json(await request.body())
@@ -78,7 +140,7 @@ class Provider:
         return await chat_completions(request)
 
     async def script(self, request: Request) -> Response:
-        name: Final = request.path_params["model"]
+        name: Final = cast(str, request.path_params["model"])
         if request.method in {"DELETE", "GET"} and name not in self.scripts:
             return JSONResponse({"error": "Script not found"}, status_code=404)
         if request.method == "GET":
@@ -103,25 +165,122 @@ class Provider:
             }
         )
 
+    async def register_scenario(self, request: Request) -> Response:
+        try:
+            registration: Final = _ScenarioRegistration.model_validate_json(await request.body())
+        except ValidationError as exc:
+            return JSONResponse({"error": str(exc)}, status_code=400)
+        self.scenario_store.put(registration.scenario_id, registration.response)
+        return JSONResponse({"scenario_id": registration.scenario_id})
+
+    async def delete_scenario(self, request: Request) -> Response:
+        scenario_id: Final = cast(str, request.path_params["scenario_id"])
+        deleted: Final = self.scenario_store.drop(scenario_id)
+        return JSONResponse({"deleted": deleted}, status_code=200 if deleted else 404)
+
+    async def cost_map(self, _request: Request) -> Response:
+        cases_file: Final = JSON_OBJECT.validate_json(CASES_FILE.read_bytes())
+        return JSONResponse(cases_file["cost_map"])
+
+    async def oauth_token(self, _request: Request) -> Response:
+        return JSONResponse(
+            {
+                "access_token": "scripted-token",
+                "token_type": "Bearer",
+                "expires_in": 3600,
+            }
+        )
+
+    async def scripted(self, request: Request) -> Response:
+        segments: Final = tuple(segment for segment in cast(str, request.path_params["path"]).split("/") if segment)
+        if not segments:
+            return JSONResponse({"error": "Unknown scenario"}, status_code=404)
+        scenario_id: Final = segments[0].split(":", 1)[0]
+        response: Final = self.scenario_store.get(scenario_id)
+        if response is None:
+            return JSONResponse({"error": "Unknown scenario"}, status_code=404)
+        return self._response(response, scenario_id)
+
+    @staticmethod
+    def _response(response: StoredResponse, scenario_id: str) -> Response:
+        match response:
+            case JsonResponse():
+                return Response(
+                    content=json.dumps(response.body, separators=(",", ":")).replace(
+                        "$REQUEST_ID", scenario_id
+                    ).encode(),
+                    media_type=response.content_type,
+                )
+            case SseResponse():
+                stream_body: Final = ("\n\n".join(response.frames) + "\n\n").replace(
+                    "$REQUEST_ID", scenario_id
+                )
+                return Response(content=stream_body.encode(), media_type=response.content_type)
+            case EventStreamResponse():
+                event_body: Final = b"".join(
+                    _aws_event_frame(event.event_type, event.payload, scenario_id) for event in response.events
+                )
+                return Response(content=event_body, media_type=response.content_type)
+
     def app(self) -> Starlette:
         return Starlette(
             routes=[
                 Route("/health", health),
                 Route("/__observations", self.observed),
                 Route("/__scripts/{model}", self.script, methods=["POST", "DELETE", "GET"]),
+                Route("/__scenarios", self.register_scenario, methods=["POST"]),
+                Route("/__scenarios/{scenario_id}", self.delete_scenario, methods=["DELETE"]),
+                Route("/_cost_map", self.cost_map, methods=["GET"]),
+                Route("/_oauth/token", self.oauth_token, methods=["POST"]),
                 Route("/v1/chat/completions", self.chat, methods=["POST"]),
                 Route("/v1/completions", completions, methods=["POST"]),
                 Route("/v1/embeddings", embeddings, methods=["POST"]),
                 Route("/v1/moderations", moderations, methods=["POST"]),
+                Route("/{path:path}", self.scripted, methods=["POST"]),
             ]
         )
 
 
+CONTROL_URL: Final = os.environ.get("INTEGRATION_UPSTREAM_URL", "http://127.0.0.1:8190").rstrip("/")
+
+
+@dataclass(frozen=True, slots=True)
+class ScenarioHandle:
+    scenario_id: str
+    control_url: str
+
+    def api_base(self) -> str:
+        return f"{self.control_url}/{self.scenario_id}"
+
+
+def register_scenario(scenario_id: str, response: StoredResponse) -> ScenarioHandle:
+    http_response: Final = httpx.post(
+        f"{CONTROL_URL}/__scenarios",
+        json={"scenario_id": scenario_id, "response": response.model_dump(mode="json")},
+        trust_env=False,
+        timeout=15,
+    )
+    http_response.raise_for_status()
+    return ScenarioHandle(
+        scenario_id=scenario_id,
+        control_url=CONTROL_URL,
+    )
+
+
+def delete_scenario(handle: ScenarioHandle) -> None:
+    response: Final = httpx.delete(
+        f"{CONTROL_URL}/__scenarios/{handle.scenario_id}",
+        trust_env=False,
+        timeout=15,
+    )
+    response.raise_for_status()
+
+
 def main() -> None:
     parser: Final = argparse.ArgumentParser()
     parser.add_argument("--port", type=int, default=8190)
     arguments: Final = parser.parse_args()
-    uvicorn.run(Provider().app(), host="127.0.0.1", port=arguments.port, access_log=False)
+    uvicorn.run(Provider().app(), host="127.0.0.1", port=cast(int, arguments.port), access_log=False)
 
 
 if __name__ == "__main__":

tests/integration/conftest.py

@@ -1,20 +1,20 @@
 from __future__ import annotations
 
+import hashlib
 import json
 import os
-import hashlib
+from collections.abc import Iterator, Sequence
 from importlib.metadata import version
-from collections.abc import Generator, Iterator
 from pathlib import Path
 from typing import Final
 
-import pytest
 import httpx
+import pytest
 from redis import Redis
 
 from tests.integration._support.client import Gateway, eventually, gateway_from_environment
-from tests.integration._support.manifest import OWNED_DIRECTORIES, contracts
 from tests.integration._support.generation import LIFECYCLE_SETTINGS
+from tests.integration._support.manifest import OWNED_DIRECTORIES, contracts
 
 COLLECTED: Final = pytest.StashKey[tuple[str, ...]]()
 REPORTS: Final = pytest.StashKey[list[pytest.TestReport]]()
@@ -28,6 +28,24 @@ def pytest_configure(config: pytest.Config) -> None:
     config.addinivalue_line("markers", "integration: owned real-service integration contracts")
     config.addinivalue_line("markers", "covers(*ids): independently asserted behavior contracts")
     config.stash[REPORTS] = []
+    config.pluginmanager.register(IntegrationReportPlugin(config))
+
+
+class IntegrationReportPlugin:
+    def __init__(self, config: pytest.Config) -> None:
+        self.config = config
+
+    def pytest_runtest_logreport(self, report: pytest.TestReport) -> None:
+        self.config.stash[REPORTS].append(report)
+
+    @pytest.hookimpl(optionalhook=True)
+    def pytest_xdist_node_collection_finished(self, node: object, ids: Sequence[str]) -> None:
+        self.config.stash[COLLECTED] = tuple(nodeid for nodeid in ids if _owned(nodeid))
+
+
+def _owned(nodeid: str) -> bool:
+    parts: Final = Path(nodeid.split("::", 1)[0]).parts
+    return parts[:2] == ("tests", "integration") and len(parts) > 3 and parts[2] in OWNED_DIRECTORIES
 
 
 def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
@@ -54,16 +72,9 @@ def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item
     config.stash[COLLECTED] = tuple(item.nodeid for item in owned)
 
 
-@pytest.hookimpl(wrapper=True)
-def pytest_runtest_makereport(
-    item: pytest.Item, call: pytest.CallInfo[None]
-) -> Generator[None, pytest.TestReport, pytest.TestReport]:
-    report: Final = yield
-    item.config.stash[REPORTS].append(report)
-    return report
-
-
 def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None:
+    if hasattr(session.config, "workerinput"):
+        return
     destination: Final = os.environ.get("INTEGRATION_RESULTS_DIR")
     if destination is None:
         return

tests/integration/cost_calculation/conftest.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import functools
+import json
+import os
+from collections.abc import Mapping
+from hashlib import sha256
+from typing import Final
+
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from pydantic import BaseModel, ConfigDict
+
+from integration._support.client import JSON_OBJECT, Scenario, eventually, object_value, string_value
+from integration._support.database import read_rows
+from integration._support.upstream import delete_scenario, register_scenario
+from integration.cost_calculation.cost_tracking_case import CostTrackingTestCase
+
+
+class CostBreakdown(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    input_cost: float | None = None
+    output_cost: float | None = None
+    cache_read_cost: float | None = None
+    cache_creation_cost: float | None = None
+    reasoning_cost: float | None = None
+    tool_usage_cost: float | None = None
+    total_cost: float | None = None
+    service_tier: str | None = None
+
+
+class CostMetadata(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    cost_breakdown: CostBreakdown | None = None
+
+
+class CostRow(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    spend: float | None = None
+    prompt_tokens: int | None = None
+    completion_tokens: int | None = None
+    metadata: CostMetadata | None = None
+
+    @property
+    def breakdown(self) -> CostBreakdown:
+        assert self.metadata is not None and self.metadata.cost_breakdown is not None
+        return self.metadata.cost_breakdown
+
+
+def approx_equal(actual: float, expected: float) -> bool:
+    return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2)
+
+
+def assert_total_is_sum_of_components(row: CostRow, context: str) -> None:
+    breakdown: Final = row.breakdown
+    total: Final = sum(
+        cost or 0.0
+        for cost in (breakdown.input_cost, breakdown.output_cost, breakdown.tool_usage_cost)
+    )
+    assert breakdown.total_cost is not None and approx_equal(breakdown.total_cost, total), (
+        f"{context}: total_cost {breakdown.total_cost} != input_cost {breakdown.input_cost} "
+        f"+ output_cost {breakdown.output_cost} + tool_usage_cost {breakdown.tool_usage_cost} "
+        f"(sum {total})"
+    )
+    assert row.spend is not None and approx_equal(row.spend, breakdown.total_cost), (
+        f"{context}: row spend {row.spend} != breakdown total_cost {breakdown.total_cost}"
+    )
+
+
+def _row(value: Mapping[str, object]) -> CostRow | None:
+    metadata_value: Final = value.get("metadata")
+    metadata: Final = json.loads(metadata_value) if isinstance(metadata_value, str) else metadata_value
+    parsed: Final = CostRow.model_validate({**value, "metadata": metadata})
+    return parsed if parsed.metadata and parsed.metadata.cost_breakdown else None
+
+
+def poll_cost_row(key: str) -> CostRow:
+    digest: Final = sha256(key.encode()).hexdigest()
+
+    def read() -> CostRow | None:
+        rows: Final = read_rows(
+            'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
+            (digest,),
+        )
+        return next((parsed for row in rows if (parsed := _row(row)) is not None), None)
+
+    result: Final = eventually(read, lambda row: row is not None, seconds=60)
+    assert result is not None
+    return result
+
+
+@functools.cache
+def _vertex_private_key_pem() -> str:
+    return rsa.generate_private_key(public_exponent=65537, key_size=2048).private_bytes(
+        serialization.Encoding.PEM,
+        serialization.PrivateFormat.PKCS8,
+        serialization.NoEncryption(),
+    ).decode()
+
+
+def _vertex_service_account_json(url: str) -> str:
+    return json.dumps(
+        {
+            "type": "service_account",
+            "project_id": "cc-scripted-project",
+            "private_key_id": "scripted",
+            "private_key": _vertex_private_key_pem(),
+            "client_email": "scripted@cc-scripted-project.iam.gserviceaccount.com",
+            "client_id": "0",
+            "auth_uri": f"{url}/_oauth/authorize",
+            "token_uri": f"{url}/_oauth/token",
+        }
+    )
+
+
+def register_scenario_deployment(
+    scenario: Scenario,
+    case: CostTrackingTestCase,
+    marker: str,
+    key: str,
+) -> str:
+    control_url: Final = os.environ["INTEGRATION_UPSTREAM_URL"].rstrip("/")
+    run_marker: Final = sha256(key.encode()).hexdigest()[:12]
+    handle: Final = register_scenario(f"sc-{marker}-{run_marker}", case.response)
+    scenario.cleanups.callback(delete_scenario, handle)
+    model_name: Final = f"cost-{marker}-{run_marker}"
+    parameters: Final = {
+        "model": case.litellm_model,
+        "api_key": case.api_key,
+        "api_base": handle.api_base(),
+        **case.litellm_params,
+        **(
+            {"vertex_credentials": _vertex_service_account_json(control_url)}
+            if case.rates.litellm_provider == "vertex_ai-language-models"
+            else {}
+        ),
+    }
+    created: Final = scenario.gateway.post(
+        "/model/new",
+        JSON_OBJECT.validate_python({
+            "model_name": model_name,
+            "litellm_params": parameters,
+            "model_info": (
+                {"base_model": case.base_model}
+                if case.base_model is not None
+                else {}
+            ),
+        }),
+    )
+    identity: Final = string_value(object_value(created["model_info"])["id"])
+    scenario.cleanups.callback(scenario.delete_model, identity)
+    return model_name

tests/integration/cost_calculation/cost_tracking_case.py

@@ -0,0 +1,253 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from pathlib import Path
+from types import MappingProxyType
+from typing import Annotated, Final, Literal, TypeAlias
+
+from pydantic import BaseModel, ConfigDict, Field, JsonValue
+
+CASES_PATH: Final = Path(__file__).resolve().parent / "cost_tracking_cases.json"
+
+
+class SearchContextCostPerQuery(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    search_context_size_low: float | None = None
+    search_context_size_medium: float | None = None
+    search_context_size_high: float | None = None
+
+
+class ProviderSpecificEntry(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    fast: float | None = None
+    us: float | None = None
+
+
+class CostMapEntry(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    litellm_provider: str
+    mode: str
+    max_tokens: int | None = None
+    max_input_tokens: int | None = None
+    max_output_tokens: int | None = None
+    supports_function_calling: bool | None = None
+    input_cost_per_token: float | None = None
+    output_cost_per_token: float | None = None
+    cache_read_input_token_cost: float | None = None
+    cache_creation_input_token_cost: float | None = None
+    cache_creation_input_token_cost_above_1hr: float | None = None
+    cache_read_input_token_cost_above_200k_tokens: float | None = None
+    cache_creation_input_token_cost_above_200k_tokens: float | None = None
+    output_cost_per_reasoning_token: float | None = None
+    input_cost_per_audio_token: float | None = None
+    output_cost_per_audio_token: float | None = None
+    input_cost_per_image_token: float | None = None
+    input_cost_per_video_token: float | None = None
+    input_cost_per_token_above_200k_tokens: float | None = None
+    output_cost_per_token_above_200k_tokens: float | None = None
+    input_cost_per_token_flex: float | None = None
+    output_cost_per_token_flex: float | None = None
+    input_cost_per_token_priority: float | None = None
+    output_cost_per_token_priority: float | None = None
+    search_context_cost_per_query: SearchContextCostPerQuery | None = None
+    web_search_billing_unit: str | None = None
+    google_maps_grounding_cost_per_query: float | None = None
+    file_search_cost_per_1k_calls: float | None = None
+    provider_specific_entry: ProviderSpecificEntry | None = None
+
+
+class Deployment(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    model: str | None = None
+    base_model: str | None = None
+
+
+class JsonResponse(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    content_type: Literal["application/json"]
+    body: dict[str, JsonValue]
+
+
+class SseResponse(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    content_type: Literal["text/event-stream"]
+    frames: tuple[str, ...]
+
+
+class EventStreamEvent(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    event_type: str
+    payload: dict[str, JsonValue]
+
+
+class EventStreamResponse(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    content_type: Literal["application/vnd.amazon.eventstream"]
+    events: tuple[EventStreamEvent, ...]
+
+
+StoredResponse: TypeAlias = Annotated[
+    JsonResponse | SseResponse | EventStreamResponse,
+    Field(discriminator="content_type"),
+]
+
+
+class ExactExpected(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    spend: float
+    input_cost: float
+    output_cost: float
+    prompt_tokens: int
+    completion_tokens: int
+
+
+class RecountRates(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    input_cost_per_token: float
+    output_cost_per_token: float
+
+
+class RecountExpected(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    recount: RecountRates
+
+
+Expected: TypeAlias = ExactExpected | RecountExpected
+
+
+class CostTrackingTestCase(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    name: str
+    covers: str
+    model: str
+    deployment: Deployment | None = None
+    request: dict[str, JsonValue]
+    response: StoredResponse
+    expected: Expected
+
+    @property
+    def rates(self) -> CostMapEntry:
+        return COST_MAP[self.model]
+
+    @property
+    def litellm_model(self) -> str:
+        provider: Final = self.rates.litellm_provider
+        prefix: Final = (
+            "openai"
+            if provider == "openai" and self.rates.mode == "chat"
+            else "openai/responses"
+            if provider == "openai"
+            else _PROVIDER_PREFIXES.get(provider)
+        )
+        if prefix is None:
+            raise ValueError(f"unsupported cost-map provider {provider} for {self.model}")
+        return self.deployment.model if self.deployment and self.deployment.model is not None else (
+            self.model if prefix == "" else f"{prefix}/{self.model}"
+        )
+
+    @property
+    def litellm_params(self) -> Mapping[str, str]:
+        return _LITELLM_PARAMS[self.rates.litellm_provider]
+
+    @property
+    def api_key(self) -> str:
+        return "sk-scripted-provider"
+
+    @property
+    def base_model(self) -> str | None:
+        return self.deployment.base_model if self.deployment else None
+
+
+class _CasesFile(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    cost_map: dict[str, CostMapEntry]
+    cases: tuple[CostTrackingTestCase, ...]
+
+
+_PROVIDER_PREFIXES: Final[Mapping[str, str]] = MappingProxyType(
+    {
+        "anthropic": "anthropic",
+        "bedrock_converse": "bedrock/converse",
+        "vertex_ai-language-models": "vertex_ai",
+        "gemini": "",
+        "together_ai": "",
+        "fireworks_ai": "",
+        "azure": "",
+    }
+)
+_LITELLM_PARAMS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType(
+    {
+        "anthropic": MappingProxyType({}),
+        "bedrock_converse": MappingProxyType(
+            {
+                "aws_access_key_id": "AKIASCRIPTEDPROVIDER",
+                "aws_secret_access_key": "scripted-secret",
+                "aws_region_name": "us-east-1",
+            }
+        ),
+        "vertex_ai-language-models": MappingProxyType(
+            {"vertex_project": "cc-scripted-project", "vertex_location": "us-central1"}
+        ),
+        "gemini": MappingProxyType({}),
+        "together_ai": MappingProxyType({}),
+        "fireworks_ai": MappingProxyType({}),
+        "azure": MappingProxyType({"api_version": "2025-04-01-preview"}),
+        "openai": MappingProxyType({}),
+    }
+)
+
+_LOADED: Final = _CasesFile.model_validate_json(CASES_PATH.read_bytes())
+COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType(dict(_LOADED.cost_map))
+CASES: Final[tuple[CostTrackingTestCase, ...]] = _LOADED.cases
+_LITELLM_MODELS: Final = tuple(case.litellm_model for case in CASES)
+
+
+def data_errors() -> tuple[str, ...]:
+    case_models: Final = frozenset(case.model for case in CASES)
+    unknown_models: Final = sorted(case.model for case in CASES if case.model not in COST_MAP)
+    missing_cases: Final = sorted(model for model in COST_MAP if model not in case_models)
+    duplicate_names: Final = sorted(
+        name for name in {case.name for case in CASES} if sum(case.name == name for case in CASES) > 1
+    )
+    input_rates: Final = tuple(
+        (entry.input_cost_per_token, model) for model, entry in COST_MAP.items()
+    )
+    shared_input_rates: Final = sorted(
+        f"{rate}: {tuple(model for value, model in input_rates if value == rate)}"
+        for rate in {value for value, _ in input_rates if value is not None}
+        if sum(value == rate for value, _ in input_rates) > 1
+    )
+    recount_mismatches: Final = sorted(
+        case.name
+        for case in CASES
+        if isinstance(case.expected, RecountExpected)
+        and case.model in COST_MAP
+        and (
+            case.expected.recount.input_cost_per_token != (COST_MAP[case.model].input_cost_per_token or 0.0)
+            or case.expected.recount.output_cost_per_token != (COST_MAP[case.model].output_cost_per_token or 0.0)
+        )
+    )
+    return tuple(
+        message
+        for message in (
+            f"case models absent from cost_map: {unknown_models}" if unknown_models else None,
+            f"cost-map entries without cases: {missing_cases}" if missing_cases else None,
+            f"duplicate case names: {duplicate_names}" if duplicate_names else None,
+            f"cost-map entries share input_cost_per_token: {shared_input_rates}" if shared_input_rates else None,
+            f"recount rates differ from cost-map rates: {recount_mismatches}" if recount_mismatches else None,
+        )
+        if message is not None
+    )

tests/integration/cost_calculation/test_cost_tracking.py

@@ -0,0 +1,101 @@
+"""Cost tracking coverage for literal integration request and response data."""
+
+from __future__ import annotations
+
+from hashlib import sha256
+from typing import Final, cast
+
+import pytest
+
+from integration._support.client import JSON_OBJECT, Gateway
+from integration.cost_calculation.conftest import (
+    approx_equal,
+    assert_total_is_sum_of_components,
+    poll_cost_row,
+    register_scenario_deployment,
+)
+from integration.cost_calculation.cost_tracking_case import (
+    CASES,
+    CostTrackingTestCase,
+    ExactExpected,
+    RecountExpected,
+    data_errors,
+)
+
+if _data_errors := data_errors():
+    raise ValueError("\n".join(_data_errors))
+
+
+_CASES: Final = tuple(
+    pytest.param(case, marks=pytest.mark.covers(case.covers), id=case.name)
+    for case in CASES
+)
+
+
+def _assert_stream_has_no_error(response_text: str) -> None:
+    for line in response_text.splitlines():
+        if not line.startswith("data:"):
+            continue
+        payload = line.removeprefix("data:").strip()
+        if payload == "[DONE]":
+            continue
+        parsed = JSON_OBJECT.validate_json(payload)
+        assert "error" not in parsed, f"stream carried an error event: {parsed}"
+
+
+@pytest.mark.parametrize("case", _CASES)
+def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) -> None:
+    marker: Final = sha256(case.name.encode()).hexdigest()[:12]
+    with gateway.scenario() as scenario:
+        key: Final = scenario.key()
+        model_name: Final = register_scenario_deployment(scenario, case, marker, key)
+        response: Final = gateway.request(
+            "POST",
+            "/v1/chat/completions",
+            {**case.request, "model": model_name},
+            key=key,
+        )
+        assert response.is_success, f"{case.name}: proxy returned {response.status_code}: {response.text[:400]}"
+        if case.response.content_type == "text/event-stream":
+            _assert_stream_has_no_error(response.text)
+        row: Final = poll_cost_row(key)
+        if isinstance(case.expected, RecountExpected):
+            assert row.prompt_tokens is not None and row.prompt_tokens > 0, (
+                f"{case.name}: recount case counted no input tokens: prompt_tokens={row.prompt_tokens}"
+            )
+            assert row.completion_tokens is not None and row.completion_tokens > 0, (
+                f"{case.name}: recount case counted no output tokens: completion_tokens={row.completion_tokens}"
+            )
+            recount: Final = row.prompt_tokens * case.expected.recount.input_cost_per_token + (
+                row.completion_tokens * case.expected.recount.output_cost_per_token
+            )
+            assert row.spend is not None and approx_equal(row.spend, recount), (
+                f"{case.name}: spend {row.spend} != recount {recount} at map rates"
+            )
+            assert_total_is_sum_of_components(row, case.name)
+            return
+        expected: Final = case.expected
+        assert isinstance(expected, ExactExpected)
+        if case.response.content_type == "application/json":
+            header: Final = cast(str | None, response.headers.get("x-litellm-response-cost"))
+            assert header is not None and approx_equal(float(header), expected.spend), (
+                f"{case.name}: x-litellm-response-cost {header} != expected {expected.spend}"
+            )
+        assert row.spend is not None and approx_equal(row.spend, expected.spend), (
+            f"{case.name}: spend {row.spend} != expected {expected.spend} "
+            f"(breakdown {row.breakdown.model_dump()})"
+        )
+        breakdown: Final = row.breakdown
+        assert breakdown.input_cost is not None and approx_equal(breakdown.input_cost, expected.input_cost), (
+            f"{case.name}: input_cost {breakdown.input_cost} != expected {expected.input_cost}"
+        )
+        assert breakdown.output_cost is not None and approx_equal(breakdown.output_cost, expected.output_cost), (
+            f"{case.name}: output_cost {breakdown.output_cost} != expected {expected.output_cost}"
+        )
+        assert row.prompt_tokens == expected.prompt_tokens, (
+            f"{case.name}: prompt_tokens {row.prompt_tokens} != expected {expected.prompt_tokens}"
+        )
+        assert row.completion_tokens == expected.completion_tokens, (
+            f"{case.name}: completion_tokens {row.completion_tokens} != expected {expected.completion_tokens}"
+        )
+        assert_total_is_sum_of_components(row, case.name)

tests/integration/run.py

@@ -18,6 +18,7 @@ def main() -> int:
     parser.add_argument("--results", type=Path, default=Path("test-results/integration"))
     parser.add_argument("--seed", type=int, default=int(os.environ.get("INTEGRATION_SEED", "4106601")))
     parser.add_argument("--order-seed", type=int, default=int(os.environ.get("INTEGRATION_ORDER_SEED", "0")))
+    parser.add_argument("--workers", type=int, default=int(os.environ.get("INTEGRATION_WORKERS", "1")))
     options: Final = parser.parse_args()
     root: Final = Path(__file__).resolve().parents[2]
     selected: Final = tuple(
@@ -56,6 +57,11 @@ def main() -> int:
             f"--hypothesis-seed={options.seed}",
             f"--integration-order-seed={options.order_seed}",
             f"--junitxml={output / 'junit.xml'}",
+            *(
+                ("-n", str(options.workers))
+                if options.workers > 1
+                else ()
+            ),
         ],
         cwd=root,
         env=environment,