Merge pull request #41999 from BerriAI/litellm_cost_shard_harness_extensions

test(integration): endpoint, breakdown component and failure support in the cost harness

tests/integration/_support/upstream.py

@@ -210,6 +210,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

@@ -397,6 +397,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"
     ],
@@ -1327,6 +1336,69 @@
     "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[us.anthropic.claude-opus-5-v1:0-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_cache_read]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_reasoning]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_stream]": [
+      "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_stream_cache_read]": [
+      "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_incomplete]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_previous_response_id]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_web_search_medium]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.3-codex-responses_file_search]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_service_tier_flex]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_service_tier_priority]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_input_text]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_cache_read]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_cache_write_5m]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_cache_write_1h]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_web_search]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_stream]": [
+      "quota_management.spend_tracking.scripted_wire.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_stream_cache_read]": [
+      "quota_management.spend_tracking.scripted_wire.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_tiered_input_above_200k]": [
+      "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-messages_input_text]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[us.anthropic.claude-opus-5-v1:0-messages_input_text]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[anthropic.claude-sonnet-5-v1:0-messages_cache_read]": [
+      "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
     "tests/integration/mcp/test_mcp_lifecycle.py::test_health_intersects_route_restricted_key_grants_in_both_management_modes": [
       "other.mcp.health.restricted_keys_intersect_grants_in_both_modes"
     ],

tests/integration/cost_calculation/conftest.py

@@ -50,6 +50,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 +101,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
@@ -146,7 +172,7 @@ class CostTrackingTestCase(BaseModel):
         provider: Final = self.rates.litellm_provider
         prefix: Final = (
             "openai"
-            if provider == "openai" and self.rates.mode == "chat"
+            if provider == "openai" and (self.rates.mode == "chat" or self.endpoint == "/v1/responses")
             else "openai/responses"
             if provider == "openai"
             else _PROVIDER_PREFIXES.get(provider)
@@ -240,6 +266,48 @@ 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 not 400 <= case.response.status <= 599
+                or not 400 <= case.expected.failure.status <= 599
+            )
+        )
+        or (
+            not isinstance(case.expected, FailureExpected)
+            and isinstance(case.response, JsonResponse)
+            and case.response.status != 200
+        )
+    )
     return tuple(
         message
         for message in (
@@ -248,6 +316,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/test_cost_tracking.py

@@ -12,12 +12,14 @@ 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,
 )
@@ -40,7 +42,9 @@ def _assert_stream_has_no_error(response_text: str) -> None:
         if payload == "[DONE]":
             continue
         parsed = JSON_OBJECT.validate_json(payload)
-        assert "error" not in parsed, f"stream carried an error event: {parsed}"
+        assert (
+            "error" not in parsed and parsed.get("type") not in {"error", "response.failed"}
+        ), f"stream carried an error event: {parsed}"
 
 
 @pytest.mark.parametrize("case", _CASES)
@@ -51,10 +55,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 +108,62 @@ 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, actual_component, expected_component in (
+            (
+                "cache_read_cost",
+                "x-litellm-response-cost-cache-read",
+                breakdown.cache_read_cost,
+                expected.cache_read_cost,
+            ),
+            (
+                "cache_creation_cost",
+                "x-litellm-response-cost-cache-creation",
+                breakdown.cache_creation_cost,
+                expected.cache_creation_cost,
+            ),
+            (
+                "reasoning_cost",
+                "x-litellm-response-cost-reasoning",
+                breakdown.reasoning_cost,
+                expected.reasoning_cost,
+            ),
+            (
+                "tool_usage_cost",
+                "x-litellm-response-cost-tool-usage",
+                breakdown.tool_usage_cost,
+                expected.tool_usage_cost,
+            ),
+        ):
+            if expected_component is None:
+                continue
+            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}"
         )