test(e2e): one request lands the same spend on every surface (#42540)

* test(e2e): one request lands the same spend on every surface

One priced chat request must show the same response_cost on the spend log row, /key/info, /team/info, the usage export's /user/daily/activity/aggregated row, and the litellm_spend_metric Prometheus sample; each is a separate writer, so the test fails naming the surface that drifted

* fix(e2e): scrape every replica's /metrics/ and enable prometheus in the replay lane

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

tests/e2e/coverage_registry/quota_management.yaml

@@ -46,6 +46,7 @@
 - {id: quota_management.spend_tracking.cache_hit.zero_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cache_hit, assertions: [zero_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A response-cache hit logs at zero cost with the cache-hit marker"}
 - {id: quota_management.spend_tracking.key_rollup.matches_sum_of_logs, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_rollup, assertions: [matches_sum_of_logs], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A key's rolled-up spend equals the sum of its log rows"}
 - {id: quota_management.spend_tracking.concurrent_burst.loses_no_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: concurrent_burst, assertions: [loses_no_spend], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "Concurrent calls all land as spend; no row lost to write contention"}
+- {id: quota_management.spend_tracking.surface_consistency.matches_every_surface, module: quota_management, tier: P1, behavior: spend_tracking, variant: surface_consistency, assertions: [matches_every_surface], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "One priced request lands the same response_cost on the spend log row, /key/info, /team/info, the usage export's /user/daily/activity/aggregated row, and the litellm_spend_metric Prometheus sample; each is a separate writer, so a rounding, dropped, or double-counted write on one drifts it from the rest (LIT-3620, LIT-5045)"}
 - {id: quota_management.spend_tracking.tags.attributes_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: tags, assertions: [attributes_spend], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Request tags round-trip to spend rows and tag rollups match tagged logs"}
 - {id: quota_management.spend_tracking.end_user.attributes_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: end_user, assertions: [attributes_spend], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "user= attribution lands the end-user id on the spend row"}
 - {id: quota_management.spend_tracking.per_model.writes_own_rows, module: quota_management, tier: P2, behavior: spend_tracking, variant: per_model, assertions: [writes_own_rows], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Each model on a shared key gets its own spend row"}

tests/e2e/gateway/record_replay_ci_config.yml

@@ -2,3 +2,6 @@ general_settings:
   master_key: os.environ/LITELLM_MASTER_KEY
   store_model_in_db: true
   disable_model_info_refresh: true
+
+litellm_settings:
+  callbacks: ["prometheus"]

tests/e2e/quota_management/spend_tracking/spend_e2e_client.py

@@ -12,9 +12,10 @@ helpers from one place.
 from __future__ import annotations
 
 import time
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
+from types import MappingProxyType
 from typing import Final
 
 from e2e_config import unique_marker
@@ -48,6 +49,7 @@ from models import (
     SpendLogsPageParams,
     SpendTagsResponse,
     TagSpend,
+    TeamInfoParams,
     UserDeleteBody,
     UserDeleteResponse,
     UserNewBody,
@@ -57,6 +59,8 @@ from models import (
 from proxy_client import Converged, ProxyClient, await_converged
 from pydantic import BaseModel, Field
 
+METRICS_PATH: Final = "/metrics/"
+
 __all__ = [
     "BatchCreateBody",
     "CallbackLogMetadata",
@@ -189,6 +193,7 @@ class DailyActivityKeyMetadata(BaseModel):
 
 class DailyActivityKeyMetrics(BaseModel):
     api_requests: int = 0
+    spend: float = 0.0
 
 
 class DailyActivityKeyBreakdown(BaseModel):
@@ -209,6 +214,14 @@ class DailyActivityResponse(BaseModel):
     results: list[DailyActivityRow] = []
 
 
+class TeamInfoSpend(BaseModel):
+    spend: float | None = None
+
+
+class TeamInfoSpendResponse(BaseModel):
+    team_info: TeamInfoSpend
+
+
 def _chat_body(
     model: str,
     content: str,
@@ -334,6 +347,42 @@ class SpendClient:
             time.sleep(self.proxy.poll_interval)
         return spend
 
+    def team_spend(self, team_id: str) -> float:
+        return (
+            unwrap(
+                self.proxy.transport.get(
+                    "/team/info",
+                    headers=self.proxy.transport.master,
+                    params=TeamInfoParams(team_id=team_id),
+                    response_type=TeamInfoSpendResponse,
+                )
+            ).team_info.spend
+            or 0.0
+        )
+
+    def poll_team_spend(self, team_id: str, *, minimum: float = 0.0) -> float:
+        outcome: Final = await_converged(
+            lambda: self.team_spend(team_id),
+            converged=lambda spend: spend > minimum,
+            timeout=self.proxy.poll_timeout,
+            interval=self.proxy.poll_interval,
+            now=time.monotonic,
+            sleep=time.sleep,
+        )
+        return outcome.result if isinstance(outcome, Converged) else outcome.last_result
+
+    def scrape_metrics(self) -> Mapping[str, ProbeResult]:
+        """GET /metrics/ on every replica in PROXY_REPLICA_URLS, keyed by replica. The
+        counter is per pod, so the union of the replicas is the fleet's exposition; the
+        trailing slash is the mounted app's own path, since bare /metrics answers a 307
+        whose Location drops the port behind a Host-rewriting balancer."""
+        return MappingProxyType(
+            {
+                replica: transport.probe(METRICS_PATH, params=NoBody())
+                for replica, transport in self.proxy.replicas.items()
+            }
+        )
+
     def spend_logs_page(
         self, *, api_key: str | None, page: int, page_size: int
     ) -> SpendLogsPage:
@@ -500,9 +549,21 @@ class SpendClient:
         return self.proxy.transport.probe("/health", params=HealthParams(model=model))
 
     def daily_activity_for_key(self, token: str, *, start: datetime, end: datetime) -> DailyActivityKeyBreakdown | None:
+        return self._key_breakdown("/user/daily/activity", token, start=start, end=end)
+
+    def usage_export_row_for_key(
+        self, token: str, *, start: datetime, end: datetime
+    ) -> DailyActivityKeyBreakdown | None:
+        """The key's row on /user/daily/activity/aggregated, the response the
+        dashboard's Export Usage Data CSV serializes."""
+        return self._key_breakdown("/user/daily/activity/aggregated", token, start=start, end=end)
+
+    def _key_breakdown(
+        self, route: str, token: str, *, start: datetime, end: datetime
+    ) -> DailyActivityKeyBreakdown | None:
         response: Final = unwrap(
             self.proxy.transport.get(
-                "/user/daily/activity",
+                route,
                 headers=self.proxy.transport.master,
                 params=DailyActivityParams(
                     start_date=start.strftime("%Y-%m-%d"),
@@ -519,9 +580,21 @@ class SpendClient:
 
     def poll_daily_activity_for_key(
         self, token: str, *, start: datetime, end: datetime, min_requests: int
+    ) -> DailyActivityKeyBreakdown | None:
+        return self._poll_key_breakdown(lambda: self.daily_activity_for_key(token, start=start, end=end), min_requests)
+
+    def poll_usage_export_row_for_key(
+        self, token: str, *, start: datetime, end: datetime, min_requests: int
+    ) -> DailyActivityKeyBreakdown | None:
+        return self._poll_key_breakdown(
+            lambda: self.usage_export_row_for_key(token, start=start, end=end), min_requests
+        )
+
+    def _poll_key_breakdown(
+        self, fetch: Callable[[], DailyActivityKeyBreakdown | None], min_requests: int
     ) -> DailyActivityKeyBreakdown | None:
         outcome: Final = await_converged(
-            lambda: self.daily_activity_for_key(token, start=start, end=end),
+            fetch,
             converged=lambda found: found is not None and found.metrics.api_requests >= min_requests,
             timeout=self.proxy.poll_timeout,
             interval=self.proxy.poll_interval,

tests/e2e/quota_management/spend_tracking/test_spend_surface_consistency_e2e.py

@@ -0,0 +1,161 @@
+"""One priced request must land the same response_cost on every spend surface.
+
+A customer reconciles the bill from whichever surface they look at: the spend
+log row, the key's and the team's rolled-up spend on /key/info and /team/info,
+the usage page's Export Usage Data CSV (the dashboard serializes the
+/user/daily/activity/aggregated rows it already holds; there is no server-side
+CSV endpoint), and the litellm_spend_metric counter Prometheus scrapes. Each is
+written by a different writer (the spend log insert, the key and team rollups in
+db_spend_update_writer, the daily spend tables, the Prometheus success callback),
+so one of them can drift without the others noticing: the cause of the
+key-versus-log mismatch in LIT-3620 and the export-versus-console mismatch in
+LIT-5045. The deployment carries its own per-token rates, so the expected cost
+is computed from the returned usage rather than read off any one surface, and
+every surface is held to that number.
+
+/metrics is per pod, so every replica the stack exports (PROXY_REPLICA_URLS) is
+scraped directly and the samples merged; a stack that exports only its balancer
+is scraped there until the pod that served the call answers. The request itself
+is sent once.
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Mapping
+from datetime import datetime, timedelta, timezone
+from itertools import groupby
+from math import isclose
+from types import MappingProxyType
+from typing import Final
+
+import pytest
+from e2e_config import provider_edge_base, unique_marker
+from e2e_http import ProbeResult
+from lifecycle import ResourceManager
+from models import ChatBody, ChatMessage, KeyGenerateBody, LiteLLMParamsBody, TeamNewBody
+from prometheus_client.parser import text_string_to_metric_families
+from proxy_client import Converged, await_converged
+from spend_e2e_client import SpendClient, unwrap
+from spend_reconciliation import INPUT_RATE, OUTPUT_RATE
+
+pytestmark = pytest.mark.e2e
+
+SPEND_METRIC: Final = "litellm_spend_metric_total"
+KEY_HASH_LABEL: Final = "hashed_api_key"
+TEAM_LABEL: Final = "team"
+
+SeriesLabels = tuple[tuple[str, str], ...]
+
+
+def _spend_series_for_key(scrapes: Mapping[str, ProbeResult], token: str) -> Mapping[SeriesLabels, float]:
+    samples: Final = sorted(
+        (tuple(sorted(sample.labels.items())), sample.value)
+        for scrape in scrapes.values()
+        if scrape.status_code == 200
+        for family in text_string_to_metric_families(scrape.body)
+        for sample in family.samples
+        if sample.name == SPEND_METRIC and sample.labels.get(KEY_HASH_LABEL) == token
+    )
+    return MappingProxyType(
+        {labels: sum(value for _, value in group) for labels, group in groupby(samples, key=lambda sample: sample[0])}
+    )
+
+
+def _poll_spend_series_for_key(client: SpendClient, token: str) -> Mapping[SeriesLabels, float]:
+    outcome: Final = await_converged(
+        client.scrape_metrics,
+        converged=lambda scrapes: bool(_spend_series_for_key(scrapes, token)),
+        timeout=client.proxy.poll_timeout,
+        interval=client.proxy.poll_interval,
+        now=time.monotonic,
+        sleep=time.sleep,
+    )
+    scrapes: Final = outcome.result if isinstance(outcome, Converged) else outcome.last_result
+    assert _spend_series_for_key(scrapes, token), (
+        f"{SPEND_METRIC} never exposed a series for {KEY_HASH_LABEL}={token} on any replica; "
+        f"last scrape status per replica: {({replica: scrape.status_code for replica, scrape in scrapes.items()})}"
+    )
+    return _spend_series_for_key(scrapes, token)
+
+
+def _same_spend(actual: float | None, expected: float) -> bool:
+    return actual is not None and isclose(actual, expected, rel_tol=1e-6, abs_tol=1e-9)
+
+
+class TestSpendSurfaceConsistency:
+    @pytest.mark.replayable
+    @pytest.mark.covers("quota_management.spend_tracking.surface_consistency.matches_every_surface")
+    def test_one_request_lands_the_same_spend_on_every_surface(
+        self, client: SpendClient, resources: ResourceManager
+    ) -> None:
+        started: Final = datetime.now(timezone.utc)
+        marker: Final = unique_marker()
+        base: Final = provider_edge_base("openai")
+        model: Final = f"e2e-spend-surfaces-{marker}"
+        model_id: Final = client.proxy.create_model(
+            model,
+            LiteLLMParamsBody(
+                model="openai/gpt-5.6-luna",
+                api_key="os.environ/OPENAI_API_KEY",
+                api_base=None if base is None else f"{base}/v1",
+                input_cost_per_token=INPUT_RATE,
+                output_cost_per_token=OUTPUT_RATE,
+            ),
+        )
+        resources.defer(lambda: client.proxy.delete_model(model_id))
+        team_id: Final = client.proxy.create_team(TeamNewBody(team_alias=f"e2e-spend-surfaces-{marker}"))
+        resources.defer(lambda: client.proxy.delete_team(team_id))
+        record: Final = client.generate_key_record(
+            KeyGenerateBody(team_id=team_id, models=[model], key_alias=f"e2e-spend-surfaces-{marker}")
+        )
+        resources.defer(lambda: client.proxy.delete_key(record.key))
+        assert record.token, "/key/generate answered without the key's token hash"
+        token: Final = record.token
+
+        response: Final = unwrap(
+            client.proxy.chat(
+                record.key,
+                ChatBody(
+                    model=model,
+                    messages=[ChatMessage(role="user", content=f"Reply with one word. {marker}")],
+                    max_completion_tokens=128,
+                ),
+            )
+        )
+        usage: Final = response.usage
+        assert response.id, "successful response must have an ID"
+        assert usage is not None and usage.prompt_tokens and usage.completion_tokens, f"no billable usage: {usage}"
+        expected: Final = usage.prompt_tokens * INPUT_RATE + usage.completion_tokens * OUTPUT_RATE
+
+        rows: Final = client.proxy.poll_logs_for_request_id(response.id)
+        assert len(rows) == 1, f"expected one spend row for {response.id}, saw {len(rows)}: {rows}"
+        row: Final = rows[0]
+        assert row.api_key == token, f"spend row keyed by {row.api_key}, not the key's token hash {token}"
+        assert row.team_id == team_id, f"spend row attributed to team {row.team_id}, not {team_id}"
+        assert row.status == "success", f"spend row status {row.status}"
+
+        key_spend: Final = client.poll_key_spend(record.key, minimum=expected * 0.999999)
+        team_spend: Final = client.poll_team_spend(team_id, minimum=expected * 0.999999)
+        export_row: Final = client.poll_usage_export_row_for_key(
+            token, start=started - timedelta(days=1), end=datetime.now(timezone.utc), min_requests=1
+        )
+        assert export_row is not None, f"/user/daily/activity/aggregated never listed key {token} under api_keys"
+        series: Final = _poll_spend_series_for_key(client, token)
+        off_team: Final = tuple(labels for labels in series if dict(labels).get(TEAM_LABEL) != team_id)
+        assert not off_team, f"{SPEND_METRIC} series for the key carry a team other than {team_id}: {off_team}"
+
+        observed: Final = MappingProxyType(
+            {
+                "/spend/logs row": row.spend,
+                "/key/info spend": key_spend,
+                "/team/info spend": team_spend,
+                "usage export row (/user/daily/activity/aggregated)": export_row.metrics.spend,
+                SPEND_METRIC: sum(series.values()),
+            }
+        )
+        drifted: Final = tuple(surface for surface, spend in observed.items() if not _same_spend(spend, expected))
+        assert not drifted, (
+            f"response_cost {expected} (usage {usage.prompt_tokens}x{INPUT_RATE} + "
+            f"{usage.completion_tokens}x{OUTPUT_RATE}) drifted on {drifted}; every surface: {dict(observed)}"
+        )