test(e2e): every spend row a virtual key writes joins its token across all write paths

tests/e2e/CLAUDE.md

@@ -181,13 +181,14 @@ quota_management.<behavior>.<variant>.<assertion>
                                | team_multi_window | fallback | spend_counter
               <spend_tracking> chat_completions | stream | messages_bridge | embeddings
                                | cache_hit | key_rollup | concurrent_burst | tags | end_user
-                               | per_model | failure | spend_calculate | pagination
+                               | per_model | failure | spend_calculate | pagination | key_attribution
   assertion : blocks_over_limit | resets_after_window | headers_report_remaining | picks_under_tpm
               | blocks_then_resets | resets_windows_independently | alerts_without_blocking
               | isolates_per_model | isolates_per_member | isolates_per_group | enforced_across_keys
               | routes_to_fallback | reseed_matches_db | reports_spend | logs_cost | zero_cost
               | matches_sum_of_logs | loses_no_spend | attributes_spend | writes_own_rows
-              | writes_failure_row | returns_cost | keeps_total
+              | writes_failure_row | returns_cost | keeps_total | joins_key | reports_alias_and_email
+              | health_rows_keep_service_account | batch_cost_joins_key
   e.g.  quota_management.ratelimit.rpm.blocks_over_limit           exercised_on=[chat_completions, messages]
         quota_management.budget.key.blocks_over_limit              exercised_on=[chat_completions]
 ```

tests/e2e/coverage_registry/quota_management.yaml

@@ -58,3 +58,7 @@
 - {id: quota_management.spend_tracking.service_tier.bills_tier_rates, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier, assertions: [bills_tier_rates], exercised_on: [chat_completions], source: "cost_calculator.py", rationale: "A priority service_tier call bills input, output, and reasoning at the deployment's *_priority rates and records the tier on the row (#35923, #35925)"}
 - {id: quota_management.spend_tracking.cost_headers.additive_components, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_headers, assertions: [additive_components], exercised_on: [chat_completions], source: "proxy/common_request_processing.py", rationale: "The x-litellm-response-cost-* component headers sum to the total, input covers only fresh tokens, and reasoning stays a subset of output (#36965)"}
 - {id: quota_management.spend_tracking.passthrough_stream.injects_usage_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: passthrough_stream, assertions: [injects_usage_cost], exercised_on: [openai_passthrough], source: "proxy/pass_through_endpoints/streaming_handler.py", rationale: "With include_cost_in_streaming_usage on, the /openai passthrough's final streaming usage frame carries the proxy-computed cost (#36503). Uncovered: the flag is only settable in litellm_settings, and the shared e2e stack does not turn it on yet"}
+- {id: quota_management.spend_tracking.key_attribution.joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [joins_key], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Every spend row a virtual key writes across chat, queued chat, messages, responses, embeddings, the Gemini passthrough, file upload, and batch create carries api_key equal to the key's token hash and the key alias, the join the usage APIs depend on; a re-hashed token shows up as an unattributed key-hash-* row (#39568, #39572)"}
+- {id: quota_management.spend_tracking.key_attribution.reports_alias_and_email, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [reports_alias_and_email], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native], source: "proxy/management_endpoints/internal_user_endpoints.py", rationale: "/spend/logs?api_key= returns every one of the key's rows with its alias and /user/daily/activity aggregates them under the key's token with key_alias and user_email; /spend/logs carries no email field, so the email is asserted on daily activity only"}
+- {id: quota_management.spend_tracking.key_attribution.health_rows_keep_service_account, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [health_rows_keep_service_account], exercised_on: [chat_completions], source: "proxy/health_check.py", rationale: "A /health probe's spend row stays keyed by the literal litellm-internal-health-check service account rather than a hash of it, so health spend never appears as an unattributed key"}
+- {id: quota_management.spend_tracking.key_attribution.batch_cost_joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [batch_cost_joins_key], exercised_on: [batches], source: "enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py", rationale: "The _batch_cost row the batch cost poller writes once a batch completes carries the submitting key's token hash; it rides a cross-run marker batch whose metadata records the expected hash, since completion can lag by up to the polling interval. Cold start with no completed marker is a documented vacuous pass"}

tests/e2e/models.py

@@ -84,6 +84,7 @@ class KeyGenerateBody(BaseModel):
 
 class KeyGenerateResponse(BaseModel):
     key: str
+    token: str | None = None
     key_alias: str | None = None
     models: list[str] = []
     max_budget: float | None = None
@@ -672,6 +673,7 @@ class GuardrailRunRecord(BaseModel):
 
 
 class SpendLogMetadata(BaseModel):
+    user_api_key_alias: str | None = None
     applied_guardrails: list[str] | None = None
     guardrail_information: list[GuardrailRunRecord] | None = None
 

tests/e2e/quota_management/spend_tracking/conftest.py

@@ -36,6 +36,7 @@ def _driver_params(provider_model: str, env_var: str) -> LiteLLMParamsBody:
     ("claude-haiku-4-5", "anthropic/claude-haiku-4-5", "ANTHROPIC_API_KEY"),
     ("openai-text-embedding-3-small", "openai/text-embedding-3-small", "OPENAI_API_KEY"),
     ("openai-responses-codex", "openai/gpt-5.3-codex", "OPENAI_API_KEY"),
+    ("openai-gpt-4o-mini", "openai/gpt-4o-mini", "OPENAI_API_KEY"),
 )
 
 

tests/e2e/quota_management/spend_tracking/spend_e2e_client.py

@@ -15,9 +15,14 @@
 from collections.abc import Callable
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
+from typing import Final
+
+from pydantic import BaseModel, Field
 
 from e2e_config import unique_marker
 from e2e_http import (
+    FileUploadForm,
+    Headers,
     NoBody,
     ProbeResult,
     Result,
@@ -35,6 +40,8 @@
     DateRangeParams,
     EmbedBody,
     EmbedResponse,
+    KeyGenerateBody,
+    KeyGenerateResponse,
     OpenAPISchema,
     SpendCalculateBody,
     SpendCalculateResponse,
@@ -43,20 +50,137 @@
     SpendLogsPageParams,
     SpendTagsResponse,
     TagSpend,
+    UserDeleteBody,
+    UserDeleteResponse,
+    UserNewBody,
+    UserNewResponse,
+    UserRole,
 )
-from proxy_client import ProxyClient
+from proxy_client import Converged, ProxyClient, await_converged
 
 __all__ = [
+    "BatchCreateBody",
+    "BatchObject",
+    "DailyActivityKeyBreakdown",
+    "FileObject",
     "ProbeResult",
+    "ResponseIdentity",
     "SpendClient",
     "SpendLogRow",
+    "StreamingResponse",
     "build_client",
     "is_ok",
     "unique_marker",
     "unwrap",
 ]
 
 
+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")
+
+
+class GeminiPart(BaseModel):
+    text: str
+
+
+class GeminiContent(BaseModel):
+    parts: list[GeminiPart]
+
+
+class GeminiGenerationConfig(BaseModel):
+    maxOutputTokens: int
+
+
+class GeminiGenerateBody(BaseModel):
+    contents: list[GeminiContent]
+    generationConfig: GeminiGenerationConfig
+
+
+class ResponsesBody(BaseModel):
+    model: str
+    input: str
+    cache: dict[str, bool] | None = {"no-cache": True}
+
+
+class QueuedChatBody(ChatBody):
+    priority: int = 0
+
+
+class ResponseIdentity(BaseModel):
+    id: str | None = None
+
+
+class HealthParams(BaseModel):
+    model: str
+
+
+class ModelQuery(BaseModel):
+    model: str
+
+
+class FileObject(BaseModel):
+    id: str
+
+
+class BatchCreateBody(BaseModel):
+    input_file_id: str
+    endpoint: str = "/v1/chat/completions"
+    completion_window: str = "24h"
+    model: str
+    metadata: dict[str, str]
+
+
+class BatchObject(BaseModel):
+    id: str
+    status: str
+    metadata: dict[str, str] | None = None
+    created_at: int | None = None
+
+
+class BatchList(BaseModel):
+    data: list[BatchObject] = []
+
+
+class BatchListQuery(BaseModel):
+    model: str
+    limit: int
+
+
+class DailyActivityParams(BaseModel):
+    start_date: str
+    end_date: str
+    api_key: str
+
+
+class DailyActivityKeyMetadata(BaseModel):
+    key_alias: str | None = None
+    team_id: str | None = None
+    user_email: str | None = None
+
+
+class DailyActivityKeyMetrics(BaseModel):
+    api_requests: int = 0
+
+
+class DailyActivityKeyBreakdown(BaseModel):
+    metrics: DailyActivityKeyMetrics
+    metadata: DailyActivityKeyMetadata
+
+
+class DailyActivityBreakdown(BaseModel):
+    api_keys: dict[str, DailyActivityKeyBreakdown] = {}
+
+
+class DailyActivityRow(BaseModel):
+    date: str
+    breakdown: DailyActivityBreakdown
+
+
+class DailyActivityResponse(BaseModel):
+    results: list[DailyActivityRow] = []
+
+
 def _chat_body(
     model: str,
     content: str,
@@ -207,6 +331,164 @@ def spend_logs_page(
     def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult:
         return self.proxy.transport.probe(path, params=params)
 
+    def create_user(self, *, email: str, role: UserRole, user_id: str) -> str:
+        return unwrap(
+            self.proxy.transport.post(
+                "/user/new",
+                headers=self.proxy.transport.master,
+                json=UserNewBody(user_email=email, user_role=role, user_id=user_id),
+                response_type=UserNewResponse,
+            )
+        ).user_id
+
+    def delete_user(self, user_id: str) -> None:
+        _ = self.proxy.transport.post(
+            "/user/delete",
+            headers=self.proxy.transport.master,
+            json=UserDeleteBody(user_ids=[user_id]),
+            response_type=UserDeleteResponse,
+        )
+
+    def generate_key_record(self, body: KeyGenerateBody) -> KeyGenerateResponse:
+        return unwrap(
+            self.proxy.transport.post(
+                "/key/generate",
+                headers=self.proxy.transport.master,
+                json=body,
+                response_type=KeyGenerateResponse,
+            )
+        )
+
+    def send_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
+        return self.proxy.transport.send(
+            "/chat/completions",
+            headers=self.proxy.transport.bearer(key),
+            json=_chat_body(model, content, max_tokens=max_tokens),
+        )
+
+    def send_queued_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
+        return self.proxy.transport.send(
+            "/queue/chat/completions",
+            headers=self.proxy.transport.bearer(key),
+            json=QueuedChatBody(
+                model=model,
+                messages=[ChatMessage(role="user", content=content)],
+                max_tokens=max_tokens,
+            ),
+        )
+
+    def send_messages(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
+        return self.proxy.transport.send(
+            "/v1/messages",
+            headers=self.proxy.transport.bearer(key),
+            json=AnthropicMessagesBody(
+                model=model,
+                messages=[ChatMessage(role="user", content=content)],
+                max_tokens=max_tokens,
+            ),
+        )
+
+    def send_responses(self, key: str, model: str, content: str) -> StreamingResponse:
+        return self.proxy.transport.send(
+            "/v1/responses",
+            headers=self.proxy.transport.bearer(key),
+            json=ResponsesBody(model=model, input=content),
+        )
+
+    def send_embed(self, key: str, model: str, content: str) -> StreamingResponse:
+        return self.proxy.transport.send(
+            "/embeddings",
+            headers=self.proxy.transport.bearer(key),
+            json=EmbedBody(model=model, input=content),
+        )
+
+    def send_gemini_generate(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
+        return self.proxy.transport.send(
+            f"/gemini/v1beta/models/{model}:generateContent",
+            headers=GeminiApiKeyHeaders(x_goog_api_key=key),
+            json=GeminiGenerateBody(
+                contents=[GeminiContent(parts=[GeminiPart(text=content)])],
+                generationConfig=GeminiGenerationConfig(maxOutputTokens=max_tokens),
+            ),
+        )
+
+    def upload_batch_file(self, key: str, model: str, content: bytes) -> FileObject:
+        return unwrap(
+            self.proxy.transport.upload(
+                "/v1/files",
+                headers=self.proxy.transport.bearer(key),
+                form=FileUploadForm(purpose="batch"),
+                filename="key_attribution.jsonl",
+                content=content,
+                params=ModelQuery(model=model),
+                response_type=FileObject,
+            )
+        )
+
+    def create_batch(self, key: str, body: BatchCreateBody) -> BatchObject:
+        return unwrap(
+            self.proxy.transport.post(
+                "/v1/batches",
+                headers=self.proxy.transport.bearer(key),
+                json=body,
+                response_type=BatchObject,
+            )
+        )
+
+    def list_batches(self, key: str, model: str, *, limit: int) -> list[BatchObject]:
+        return unwrap(
+            self.proxy.transport.get(
+                "/v1/batches",
+                headers=self.proxy.transport.bearer(key),
+                params=BatchListQuery(model=model, limit=limit),
+                response_type=BatchList,
+            )
+        ).data
+
+    def retrieve_batch(self, key: str, batch_id: str) -> BatchObject:
+        return unwrap(
+            self.proxy.transport.get(
+                f"/v1/batches/{batch_id}",
+                headers=self.proxy.transport.bearer(key),
+                params=NoBody(),
+                response_type=BatchObject,
+            )
+        )
+
+    def health(self, model: str) -> ProbeResult:
+        return self.proxy.transport.probe("/health", params=HealthParams(model=model))
+
+    def daily_activity_for_key(self, token: str, *, start: datetime, end: datetime) -> DailyActivityKeyBreakdown | None:
+        response: Final = unwrap(
+            self.proxy.transport.get(
+                "/user/daily/activity",
+                headers=self.proxy.transport.master,
+                params=DailyActivityParams(
+                    start_date=start.strftime("%Y-%m-%d"),
+                    end_date=end.strftime("%Y-%m-%d"),
+                    api_key=token,
+                ),
+                response_type=DailyActivityResponse,
+            )
+        )
+        return next(
+            (row.breakdown.api_keys[token] for row in response.results if token in row.breakdown.api_keys),
+            None,
+        )
+
+    def poll_daily_activity_for_key(
+        self, token: str, *, start: datetime, end: datetime, min_requests: int
+    ) -> DailyActivityKeyBreakdown | None:
+        outcome: Final = await_converged(
+            lambda: self.daily_activity_for_key(token, start=start, end=end),
+            converged=lambda found: found is not None and found.metrics.api_requests >= min_requests,
+            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 openapi(self) -> OpenAPISchema:
         return unwrap(
             self.proxy.transport.get(