diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 5740a878608..6c34e5daa5c 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -49,6 +49,7 @@ - {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.end_user.attributes_responses_header, module: quota_management, tier: P1, behavior: spend_tracking, variant: end_user, assertions: [attributes_responses_header], exercised_on: [responses], source: "proxy/auth/auth_utils.py", rationale: "A /v1/responses call carrying x-litellm-customer-id or x-litellm-end-user-id plus x-litellm-tags, the headers Codex CLI attaches through its config.toml http_headers because it has no body field for the end user, lands the end user and the tags on a costed aresponses spend row whose spend the customer's /customer/info total matches (LIT-8575)"} - {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"} - {id: quota_management.spend_tracking.failure.writes_failure_row, module: quota_management, tier: P1, behavior: spend_tracking, variant: failure, assertions: [writes_failure_row], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_log_error_logger.py", rationale: "A failed call writes a failure-status spend row"} - {id: quota_management.spend_tracking.failure.writes_normalized_error, module: quota_management, tier: P1, behavior: spend_tracking, variant: failure, assertions: [writes_normalized_error], exercised_on: [chat_completions], source: "litellm_core_utils/error_normalization.py", rationale: "Failure rows carry a stable metadata.error_information.normalized_error key next to the unchanged error_message, so two upstream auth failures with different provider wording share one cluster key a dashboard can group by"} diff --git a/tests/e2e/quota_management/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md b/tests/e2e/quota_management/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md index 6baebc4c28c..32dc0c47dda 100644 --- a/tests/e2e/quota_management/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md +++ b/tests/e2e/quota_management/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md @@ -23,7 +23,7 @@ proxy + SpendLogs rows. Status: `covered` / `partial` / `gap`. | per-model / per-provider attribution | `test_spend_tracking_utils.py` | unit | covered | yes (`test_each_model_on_a_shared_key_gets_its_own_row`) | | field population (model/tokens/api_key/team/org) | `test_spend_tracking_utils.py` | unit | partial | yes (asserts real values) | | `request_tags` propagation | `test_db_spend_update_writer.py` | unit | partial | yes (`test_request_tags_round_trip`) | -| `end_user` attribution | unit | unit | partial | yes (`test_end_user_spend_attributed_on_row`) | +| `end_user` attribution | unit | unit | partial | yes (`test_end_user_spend_attributed_on_row`, `test_end_user_header_attributes_responses_row`) | ## Cost calculation by modality @@ -67,6 +67,7 @@ proxy + SpendLogs rows. Status: `covered` / `partial` / `gap`. | `test_request_tags_round_trip` | tags persist onto the row | | `test_tag_spend_matches_sum_of_tagged_logs` | `/spend/tags` SUM/COUNT == tagged rows | | `test_end_user_spend_attributed_on_row` | `end_user` attributed + costed | +| `test_end_user_header_attributes_responses_row` | `x-litellm-customer-id` / `x-litellm-end-user-id` + `x-litellm-tags` headers on `/v1/responses` (the Codex CLI `http_headers` shape) attributed + tagged + costed, and `/customer/info` spend equals the row | | `test_each_model_on_a_shared_key_gets_its_own_row` | per-model/provider rows, correct model + cost, distinct request_ids matching response id | | `test_failure_call_writes_failure_status_row` | failed call -> `status=failure`, `spend=0` | | `test_spend_calculate_returns_nonzero_cost` | cost-map smoke (no batch wait) | diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index 8b63b063e14..e607c12b731 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -20,6 +20,7 @@ from typing import Final from e2e_config import unique_marker from e2e_http import ( + AuthHeaders, FileUploadForm, Headers, NoBody, @@ -36,6 +37,7 @@ from models import ( ChatMessage, ChatMetadata, ChatResponse, + CustomerInfoParams, DateRangeParams, EmbedBody, EmbedResponse, @@ -63,9 +65,10 @@ METRICS_PATH: Final = "/metrics/" __all__ = [ "BatchCreateBody", + "BatchObject", "CallbackLogMetadata", "CallbackLogPayload", - "BatchObject", + "ClientAttributionHeaders", "DailyActivityKeyBreakdown", "FileObject", "ProbeResult", @@ -80,6 +83,16 @@ __all__ = [ ] +class ClientAttributionHeaders(AuthHeaders): + """The attribution headers a coding agent attaches to every call from its own + config (Codex CLI's config.toml ``http_headers``, Claude Code's + ``ANTHROPIC_CUSTOM_HEADERS``) because it has no body field for the end user.""" + + x_litellm_customer_id: str | None = Field(default=None, alias="x-litellm-customer-id") + x_litellm_end_user_id: str | None = Field(default=None, alias="x-litellm-end-user-id") + x_litellm_tags: str | None = Field(default=None, alias="x-litellm-tags") + + class GeminiApiKeyHeaders(Headers): x_goog_api_key: str = Field(serialization_alias="x-goog-api-key") content_type: str = Field(default="application/json", serialization_alias="Content-Type") @@ -222,6 +235,10 @@ class TeamInfoSpendResponse(BaseModel): team_info: TeamInfoSpend +class CustomerSpendResponse(BaseModel): + spend: float | None = None + + def _chat_body( model: str, content: str, @@ -371,6 +388,31 @@ class SpendClient: ) return outcome.result if isinstance(outcome, Converged) else outcome.last_result + def customer_spend(self, customer_id: str) -> float: + """0.0 until the spend writer has upserted the end-user row, which /customer/info 404s before.""" + looked_up: Final = self.proxy.transport.get( + "/customer/info", + headers=self.proxy.transport.master, + params=CustomerInfoParams(end_user_id=customer_id), + response_type=CustomerSpendResponse, + ) + match looked_up: + case Success(data=data): + return data.spend or 0.0 + case _: + return 0.0 + + def poll_customer_spend(self, customer_id: str, *, minimum: float = 0.0) -> float: + outcome: Final = await_converged( + lambda: self.customer_spend(customer_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 @@ -479,9 +521,12 @@ class SpendClient: ) def send_responses(self, key: str, model: str, content: str) -> StreamingResponse: + return self.send_responses_with_headers(self.proxy.transport.bearer(key), model, content) + + def send_responses_with_headers(self, headers: AuthHeaders, model: str, content: str) -> StreamingResponse: return self.proxy.transport.send( "/v1/responses", - headers=self.proxy.transport.bearer(key), + headers=headers, json=ResponsesBody(model=model, input=content), ) diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py index 286421e2e3f..6633396b538 100644 --- a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py @@ -24,7 +24,14 @@ import pytest from e2e_http import RateLimitedError, Success from lifecycle import ResourceManager from models import KeyGenerateBody, LiteLLMParamsBody, SpendLogs, SpendLogsParams -from spend_e2e_client import SpendClient, SpendLogRow, is_ok, unique_marker, unwrap +from spend_e2e_client import ( + ClientAttributionHeaders, + SpendClient, + SpendLogRow, + is_ok, + unique_marker, + unwrap, +) pytestmark = pytest.mark.e2e @@ -425,6 +432,41 @@ def test_end_user_spend_attributed_on_row( assert (row.spend or 0) > 0, f"end-user row should cost > 0: {_summarize(rows)}" +@pytest.mark.covers("quota_management.spend_tracking.end_user.attributes_responses_header") +@pytest.mark.parametrize("header", ["x-litellm-customer-id", "x-litellm-end-user-id"]) +def test_end_user_header_attributes_responses_row( + client: SpendClient, scoped_key: str, resources: ResourceManager, header: str +) -> None: + """Codex CLI has no body field for the end user, so its config.toml http_headers + attach the customer header (and x-litellm-tags) to every /v1/responses call. + A regression that stops reading either header on the Responses route, drops the + tags, costs the row at zero, or leaves the customer's own spend total behind the + row fails here.""" + customer = resources.customer(f"e2e-codex-{unique_marker()}") + tag = f"codex-{unique_marker()}" + headers = ClientAttributionHeaders.model_validate( + {"authorization": f"Bearer {scoped_key}", header: customer, "x-litellm-tags": tag} + ) + sent = client.send_responses_with_headers( + headers, "openai-responses-codex", f"one word {unique_marker()}" + ) + assert sent.ok, f"/v1/responses failed with {sent.status_code}: {sent.body[:300]}" + + rows = client.poll_logs_for_key( + scoped_key, predicate=lambda rs: any(r.end_user == customer for r in rs) + ) + row = _require_row( + rows, lambda r: r.end_user == customer, f"attributed to end_user {customer!r} via {header}" + ) + assert row.call_type == "aresponses", f"row is not a Responses row: {_summarize(rows)}" + assert tag in (row.request_tags or []), f"tag {tag!r} missing from {row.request_tags}" + assert (row.spend or 0) > 0, f"end-user row should cost > 0: {_summarize(rows)}" + customer_total = client.poll_customer_spend(customer) + assert _approx_equal(customer_total, row.spend or 0), ( + f"/customer/info spend {customer_total} != the row's {row.spend}: {_summarize(rows)}" + ) + + @pytest.mark.covers("quota_management.spend_tracking.per_model.writes_own_rows") def test_each_model_on_a_shared_key_gets_its_own_row( client: SpendClient, scoped_key: str