test(integration): assert the streamed usage.cost matches the spend row on every cost case (#46762)

* test(integration): assert the streamed usage.cost matches the spend row on every cost case

* test(integration): type the streamed usage frames as JsonValue

* test(integration): only require streamed usage.cost when the client received a usage frame

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/integration/cost_calculation/assertions.py

@@ -1,6 +1,10 @@
 from __future__ import annotations
 
+from collections.abc import Iterator, Mapping
+from typing import Final
+
 import httpx
+from integration._support.client import JSON_OBJECT
 from integration.cost_calculation.conftest import (
     CostBreakdown,
     CostRow,
@@ -8,6 +12,47 @@ from integration.cost_calculation.conftest import (
     assert_total_is_sum_of_components,
 )
 from integration.cost_calculation.cost_tracking_case import ExactExpected, RecountExpected
+from pydantic import JsonValue
+
+STREAMED_COST_ENDPOINTS: Final = frozenset({"/v1/chat/completions", "/v1/completions", "/v1/responses"})
+
+
+def _data_frames(response_text: str) -> Iterator[Mapping[str, JsonValue]]:
+    for line in response_text.splitlines():
+        if not line.startswith("data:"):
+            continue
+        payload = line.removeprefix("data:").strip()
+        if payload != "[DONE]":
+            yield JSON_OBJECT.validate_json(payload)
+
+
+def _frame_usage(frame: Mapping[str, JsonValue]) -> Mapping[str, JsonValue] | None:
+    usage: Final = frame.get("usage")
+    if isinstance(usage, dict):
+        return usage
+    nested: Final = frame.get("response")
+    nested_usage: Final = nested.get("usage") if isinstance(nested, dict) else None
+    return nested_usage if isinstance(nested_usage, dict) else None
+
+
+def streamed_usage(response_text: str) -> Mapping[str, JsonValue] | None:
+    """The last streamed frame's usage object, the only place a streaming client can read a cost from."""
+    usages: Final = tuple(usage for usage in map(_frame_usage, _data_frames(response_text)) if usage is not None)
+    return usages[-1] if usages else None
+
+
+def assert_streamed_usage_cost(case_name: str, endpoint: str, response_text: str, spend: float) -> None:
+    usage: Final = streamed_usage(response_text)
+    if usage is None:
+        return
+    cost: Final = usage.get("cost")
+    seen: Final = float(cost) if isinstance(cost, (int, float)) else None
+    if endpoint in STREAMED_COST_ENDPOINTS and spend != 0:
+        assert seen is not None and approx_equal(seen, spend), (
+            f"{case_name}: streamed usage.cost {seen} != spend {spend}"
+        )
+    elif seen is not None:
+        assert approx_equal(seen, spend), f"{case_name}: streamed usage.cost {seen} != spend {spend}"
 
 
 def assert_breakdown(
@@ -113,9 +158,17 @@ def assert_exact(
     )
     if breakdown is not None:
         assert_total_is_sum_of_components(row, breakdown, case_name)
+    if response is not None and response_content_type == "text/event-stream":
+        assert_streamed_usage_cost(case_name, response.request.url.path, response.text, expected.spend)
 
 
-def assert_recount(case_name: str, expected: RecountExpected, row: CostRow) -> None:
+def assert_recount(
+    case_name: str,
+    response_content_type: str,
+    expected: RecountExpected,
+    row: CostRow,
+    response: httpx.Response | None,
+) -> None:
     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}"
     )
@@ -146,3 +199,5 @@ def assert_recount(case_name: str, expected: RecountExpected, row: CostRow) -> N
     )
     assert row.breakdown is not None, f"{case_name}: no cost_breakdown persisted"
     assert_total_is_sum_of_components(row, row.breakdown, case_name)
+    if response is not None and response_content_type == "text/event-stream":
+        assert_streamed_usage_cost(case_name, response.request.url.path, response.text, row.spend)

tests/integration/cost_calculation/test_cost_tracking.py

@@ -242,7 +242,7 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
             row: Final = poll_cost_row(key)
             assert isinstance(expected, RecountExpected)
             assert row.status == "success", f"{case.name}: disconnect row status was {row.status}"
-            assert_recount(case.name, expected, row)
+            assert_recount(case.name, case.response.content_type, expected, row, None)
             _settled_rows(scenario, gateway, key, 1)
             return
         responses: Final = tuple(
@@ -288,7 +288,7 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase)
         rows: Final = poll_rows(key, row_count)
         if isinstance(expected, RecountExpected):
             row: Final = rows[0]
-            assert_recount(case.name, expected, row)
+            assert_recount(case.name, case.response.content_type, expected, row, response)
             _settled_rows(scenario, gateway, key, row_count)
             return
         assert isinstance(expected, ExactExpected)