test: extend cost tracking integration harness

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

tests/integration/_support/upstream.py

@@ -1,25 +1,19 @@
 from __future__ import annotations
 
 import argparse
+import json
+import os
+import struct
+import zlib
 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
-import struct
 from typing import Final, cast
-import zlib
 
 import httpx
 import uvicorn
-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,
@@ -27,6 +21,11 @@ from integration.cost_calculation.cost_tracking_case import (
     SseResponse,
     StoredResponse,
 )
+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
 
 JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
 CASES_FILE: Final = Path(__file__).resolve().parents[1] / "cost_calculation" / "cost_tracking_cases.json"
@@ -210,6 +209,7 @@ class Provider:
                         "$REQUEST_ID", scenario_id
                     ).encode(),
                     media_type=response.content_type,
+                    status_code=response.status,
                 )
             case SseResponse():
                 stream_body: Final = ("\n\n".join(response.frames) + "\n\n").replace(

tests/integration/contracts.json

@@ -382,6 +382,15 @@
     "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-gpt-5.6-stream_full_usage]": [
       "quota_management.spend_tracking.scripted_wire.logs_cost"
     ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_native_json]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-upstream_500_zero_spend]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-upstream_429_zero_spend]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
     "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-haiku-4-5-input_text]": [
       "quota_management.spend_tracking.cost_matrix.logs_cost"
     ],

tests/integration/cost_calculation/conftest.py

@@ -9,12 +9,11 @@ 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
+from pydantic import BaseModel, ConfigDict
 
 
 class CostBreakdown(BaseModel):
@@ -50,6 +49,15 @@ class CostRow(BaseModel):
         return self.metadata.cost_breakdown
 
 
+class FailureRow(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    spend: float
+    status: str
+    prompt_tokens: int | None = None
+    completion_tokens: int | None = None
+
+
 def approx_equal(actual: float, expected: float) -> bool:
     return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2)
 
@@ -92,6 +100,28 @@ def poll_cost_row(key: str) -> CostRow:
     return result
 
 
+def poll_failure_row(key: str) -> FailureRow:
+    digest: Final = sha256(key.encode()).hexdigest()
+
+    def read() -> FailureRow | None:
+        rows: Final = read_rows(
+            'SELECT spend, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
+            (digest,),
+        )
+        return next(
+            (
+                parsed
+                for row in rows
+                if (parsed := FailureRow.model_validate(row)).status == "failure"
+            ),
+            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(

tests/integration/cost_calculation/cost_tracking_case.py

@@ -71,6 +71,7 @@ class JsonResponse(BaseModel):
 
     content_type: Literal["application/json"]
     body: dict[str, JsonValue]
+    status: int = 200
 
 
 class SseResponse(BaseModel):
@@ -108,6 +109,10 @@ class ExactExpected(BaseModel):
     output_cost: float
     prompt_tokens: int
     completion_tokens: int
+    cache_read_cost: float | None = None
+    cache_creation_cost: float | None = None
+    reasoning_cost: float | None = None
+    tool_usage_cost: float | None = None
 
 
 class RecountRates(BaseModel):
@@ -123,7 +128,19 @@ class RecountExpected(BaseModel):
     recount: RecountRates
 
 
-Expected: TypeAlias = ExactExpected | RecountExpected
+class FailureDetails(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    status: int
+
+
+class FailureExpected(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    failure: FailureDetails
+
+
+Expected: TypeAlias = ExactExpected | RecountExpected | FailureExpected
 
 
 class CostTrackingTestCase(BaseModel):
@@ -132,6 +149,15 @@ class CostTrackingTestCase(BaseModel):
     name: str
     covers: str
     model: str
+    endpoint: Literal[
+        "/v1/chat/completions",
+        "/v1/responses",
+        "/v1/messages",
+        "/v1/embeddings",
+        "/v1/rerank",
+        "/v1/completions",
+        "/v1/moderations",
+    ] = "/v1/chat/completions"
     deployment: Deployment | None = None
     request: dict[str, JsonValue]
     response: StoredResponse
@@ -240,6 +266,44 @@ def data_errors() -> tuple[str, ...]:
             or case.expected.recount.output_cost_per_token != (COST_MAP[case.model].output_cost_per_token or 0.0)
         )
     )
+    component_mismatches: Final = sorted(
+        case.name
+        for case in CASES
+        if isinstance(case.expected, ExactExpected)
+        and any(
+            component is not None
+            for component in (
+                case.expected.cache_read_cost,
+                case.expected.cache_creation_cost,
+                case.expected.reasoning_cost,
+                case.expected.tool_usage_cost,
+            )
+        )
+        and (
+            (case.expected.cache_read_cost or 0.0) + (case.expected.cache_creation_cost or 0.0)
+            > case.expected.input_cost
+            or (case.expected.reasoning_cost or 0.0) > case.expected.output_cost
+            or not _approx_equal(
+                case.expected.input_cost
+                + case.expected.output_cost
+                + (case.expected.tool_usage_cost or 0.0),
+                case.expected.spend,
+            )
+        )
+    )
+    failure_response_mismatches: Final = sorted(
+        case.name
+        for case in CASES
+        if (
+            isinstance(case.expected, FailureExpected)
+            and (not isinstance(case.response, JsonResponse) or case.response.status < 400)
+        )
+        or (
+            not isinstance(case.expected, FailureExpected)
+            and isinstance(case.response, JsonResponse)
+            and case.response.status != 200
+        )
+    )
     return tuple(
         message
         for message in (
@@ -248,6 +312,14 @@ def data_errors() -> tuple[str, ...]:
             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,
+            f"breakdown components are inconsistent: {component_mismatches}" if component_mismatches else None,
+            f"failure response statuses are inconsistent: {failure_response_mismatches}"
+            if failure_response_mismatches
+            else None,
         )
         if message is not None
     )
+
+
+def _approx_equal(actual: float, expected: float) -> bool:
+    return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2)

tests/integration/cost_calculation/cost_tracking_cases.json

@@ -534,7 +534,8 @@
         "input_cost": 0.00616704,
         "output_cost": 0.00627,
         "prompt_tokens": 12928,
-        "completion_tokens": 380
+        "completion_tokens": 380,
+        "cache_read_cost": 0.00405504
       }
     },
     {
@@ -605,7 +606,8 @@
         "input_cost": 0.0397056,
         "output_cost": 0.005775,
         "prompt_tokens": 9728,
-        "completion_tokens": 350
+        "completion_tokens": 350,
+        "cache_creation_cost": 0.038016
       }
     },
     {
@@ -681,7 +683,8 @@
         "input_cost": 0.0574464,
         "output_cost": 0.005775,
         "prompt_tokens": 9728,
-        "completion_tokens": 350
+        "completion_tokens": 350,
+        "cache_creation_cost": 0.0557568
       }
     },
     {
@@ -3417,7 +3420,8 @@
         "input_cost": 0.002232,
         "output_cost": 0.065484,
         "prompt_tokens": 1240,
-        "completion_tokens": 4040
+        "completion_tokens": 4040,
+        "reasoning_cost": 0.05742
       }
     },
     {
@@ -3638,7 +3642,8 @@
         "input_cost": 0.003312,
         "output_cost": 0.0059328,
         "prompt_tokens": 1840,
-        "completion_tokens": 412
+        "completion_tokens": 412,
+        "tool_usage_cost": 0.0125
       }
     },
     {
@@ -4494,7 +4499,8 @@
         "input_cost": 0.0018688,
         "output_cost": 0.0019,
         "prompt_tokens": 12928,
-        "completion_tokens": 380
+        "completion_tokens": 380,
+        "cache_read_cost": 0.0012288
       }
     },
     {
@@ -16955,7 +16961,8 @@
         "input_cost": 0.00276,
         "output_cost": 0.004944,
         "prompt_tokens": 1840,
-        "completion_tokens": 412
+        "completion_tokens": 412,
+        "tool_usage_cost": 0.0025
       }
     },
     {
@@ -21797,6 +21804,120 @@
         "completion_tokens": 1592
       }
     },
+    {
+      "name": "gpt-5.6-responses_native_json",
+      "covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
+      "model": "gpt-5.6",
+      "endpoint": "/v1/responses",
+      "request": {
+        "model": "$MODEL",
+        "input": "responses native fixture",
+        "stream": false
+      },
+      "response": {
+        "content_type": "application/json",
+        "body": {
+          "object": "response",
+          "status": "completed",
+          "model": "gpt-5.6",
+          "output": [
+            {
+              "type": "message",
+              "id": "msg_$REQUEST_ID",
+              "status": "completed",
+              "role": "assistant",
+              "content": [
+                {
+                  "type": "output_text",
+                  "text": "scripted response",
+                  "annotations": []
+                }
+              ]
+            }
+          ],
+          "usage": {
+            "input_tokens": 11,
+            "output_tokens": 7,
+            "total_tokens": 18,
+            "input_tokens_details": {
+              "cached_tokens": 0
+            },
+            "output_tokens_details": {
+              "reasoning_tokens": 0
+            }
+          }
+        }
+      },
+      "expected": {
+        "spend": 0.00011725,
+        "input_cost": 1.925e-05,
+        "output_cost": 9.8e-05,
+        "prompt_tokens": 11,
+        "completion_tokens": 7
+      }
+    },
+    {
+      "name": "gpt-5.6-upstream_500_zero_spend",
+      "covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
+      "model": "gpt-5.6",
+      "request": {
+        "model": "$MODEL",
+        "messages": [
+          {
+            "role": "user",
+            "content": "scripted upstream failure 500"
+          }
+        ],
+        "stream": false
+      },
+      "response": {
+        "content_type": "application/json",
+        "status": 500,
+        "body": {
+          "error": {
+            "message": "scripted upstream failure",
+            "type": "server_error",
+            "code": "500"
+          }
+        }
+      },
+      "expected": {
+        "failure": {
+          "status": 500
+        }
+      }
+    },
+    {
+      "name": "gpt-5.6-upstream_429_zero_spend",
+      "covers": "quota_management.spend_tracking.cost_matrix.logs_cost",
+      "model": "gpt-5.6",
+      "request": {
+        "model": "$MODEL",
+        "messages": [
+          {
+            "role": "user",
+            "content": "scripted upstream failure 429"
+          }
+        ],
+        "stream": false
+      },
+      "response": {
+        "content_type": "application/json",
+        "status": 429,
+        "body": {
+          "error": {
+            "message": "scripted upstream failure",
+            "type": "rate_limit_error",
+            "code": "429"
+          }
+        }
+      },
+      "expected": {
+        "failure": {
+          "status": 429
+        }
+      }
+    },
     {
       "name": "meta.llama4-maverick-17b-instruct-v1:0-input_text",
       "covers": "quota_management.spend_tracking.cost_matrix.logs_cost",

tests/integration/cost_calculation/test_cost_tracking.py

@@ -6,18 +6,19 @@ 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,
+    poll_failure_row,
     register_scenario_deployment,
 )
 from integration.cost_calculation.cost_tracking_case import (
     CASES,
     CostTrackingTestCase,
     ExactExpected,
+    FailureExpected,
     RecountExpected,
     data_errors,
 )
@@ -51,10 +52,22 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
         model_name: Final = register_scenario_deployment(scenario, case, marker, key)
         response: Final = gateway.request(
             "POST",
-            "/v1/chat/completions",
+            case.endpoint,
             {**case.request, "model": model_name},
             key=key,
         )
+        if isinstance(case.expected, FailureExpected):
+            assert response.status_code == case.expected.failure.status, (
+                f"{case.name}: proxy returned {response.status_code}, expected {case.expected.failure.status}: "
+                f"{response.text[:400]}"
+            )
+            response_cost: Final = response.headers.get("x-litellm-response-cost")
+            assert response_cost is None or approx_equal(float(response_cost), 0.0), (
+                f"{case.name}: failure response cost was {response_cost}"
+            )
+            row: Final = poll_failure_row(key)
+            assert row.spend == 0, f"{case.name}: failure spend was {row.spend}"
+            return
         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)
@@ -92,6 +105,43 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
         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}"
         )
+        for field, header_name, expected_component in (
+            ("cache_read_cost", "x-litellm-response-cost-cache-read", expected.cache_read_cost),
+            ("cache_creation_cost", "x-litellm-response-cost-cache-creation", expected.cache_creation_cost),
+            ("reasoning_cost", "x-litellm-response-cost-reasoning", expected.reasoning_cost),
+            ("tool_usage_cost", "x-litellm-response-cost-tool-usage", expected.tool_usage_cost),
+        ):
+            if expected_component is None:
+                continue
+            actual_component: Final = getattr(breakdown, field)
+            assert actual_component is not None and approx_equal(actual_component, expected_component), (
+                f"{case.name}: {field} {actual_component} != expected {expected_component}"
+            )
+            if case.response.content_type == "application/json":
+                header: Final = response.headers.get(header_name)
+                assert header is not None and approx_equal(float(header), expected_component), (
+                    f"{case.name}: {header_name} {header} != expected {expected_component}"
+                )
+        if case.response.content_type == "application/json" and any(
+            component is not None
+            for component in (
+                expected.cache_read_cost,
+                expected.cache_creation_cost,
+                expected.reasoning_cost,
+                expected.tool_usage_cost,
+            )
+        ):
+            input_header: Final = response.headers.get("x-litellm-response-cost-input")
+            output_header: Final = response.headers.get("x-litellm-response-cost-output")
+            expected_input_header: Final = expected.input_cost - (
+                expected.cache_read_cost or 0.0
+            ) - (expected.cache_creation_cost or 0.0)
+            assert input_header is not None and approx_equal(float(input_header), expected_input_header), (
+                f"{case.name}: x-litellm-response-cost-input {input_header} != expected {expected_input_header}"
+            )
+            assert output_header is not None and approx_equal(float(output_header), expected.output_cost), (
+                f"{case.name}: x-litellm-response-cost-output {output_header} != expected {expected.output_cost}"
+            )
         assert row.prompt_tokens == expected.prompt_tokens, (
             f"{case.name}: prompt_tokens {row.prompt_tokens} != expected {expected.prompt_tokens}"
         )