diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 89c04208d65..34fbe9d9247 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -181,13 +181,15 @@ quota_management... | team_multi_window | fallback | spend_counter 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 | retrieve_batch_cost_joins_retrieving_key + | poller_batch_cost_joins_creating_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] ``` diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index d0afcaca848..ad0914d455b 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -58,3 +58,8 @@ - {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, rust_control_plane], 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, batch create, and a replayed callback log 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, rust_control_plane], 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.retrieve_batch_cost_joins_retrieving_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [retrieve_batch_cost_joins_retrieving_key], exercised_on: [batches], source: "proxy/batches_endpoints/endpoints.py", rationale: "The retrieve that first sees a batch in a terminal state prices it inline and writes its {provider_batch_id}_batch_cost row against the retrieving key, so the batch each run creates is one OpenAI fails at validation within seconds and the test retrieves it by its raw provider id with the same key until it is failed; a raw id is never owned by the CheckBatchCost poller, and the row must carry that key's token hash and alias"} +- {id: quota_management.spend_tracking.key_attribution.poller_batch_cost_joins_creating_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [poller_batch_cost_joins_creating_key], exercised_on: [batches], source: "enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py", rationale: "The CheckBatchCost poller bills a completed, positive-cost batch created through a unified id against the key that created it, a different writer from the inline retrieve. No test claims this cell yet: OpenAI's completion window is 24h and both e2e stacks boot a fresh Postgres per build, so a completed batch is out of one run's reach and the managed list never shows an earlier run's batch; the cell stays visible as a gap until a run can hand a completed batch to the poller"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 62810e6cfd9..faf8557498b 100644 --- a/tests/e2e/models.py +++ b/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 diff --git a/tests/e2e/quota_management/spend_tracking/conftest.py b/tests/e2e/quota_management/spend_tracking/conftest.py index 0597c9af400..9c8ffd18144 100644 --- a/tests/e2e/quota_management/spend_tracking/conftest.py +++ b/tests/e2e/quota_management/spend_tracking/conftest.py @@ -36,6 +36,7 @@ DRIVER_MODELS: tuple[tuple[str, str, str], ...] = ( ("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"), ) 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 056799b8499..9ac97f57f47 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -15,9 +15,12 @@ import time from collections.abc import Callable from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from typing import Final from e2e_config import unique_marker from e2e_http import ( + FileUploadForm, + Headers, NoBody, ProbeResult, Result, @@ -35,6 +38,8 @@ from models import ( DateRangeParams, EmbedBody, EmbedResponse, + KeyGenerateBody, + KeyGenerateResponse, OpenAPISchema, SpendCalculateBody, SpendCalculateResponse, @@ -43,13 +48,27 @@ from models import ( SpendLogsPageParams, SpendTagsResponse, TagSpend, + UserDeleteBody, + UserDeleteResponse, + UserNewBody, + UserNewResponse, + UserRole, ) -from proxy_client import ProxyClient +from proxy_client import Converged, ProxyClient, await_converged +from pydantic import BaseModel, Field __all__ = [ + "BatchCreateBody", + "CallbackLogMetadata", + "CallbackLogPayload", + "BatchObject", + "DailyActivityKeyBreakdown", + "FileObject", "ProbeResult", + "ResponseIdentity", "SpendClient", "SpendLogRow", + "StreamingResponse", "build_client", "is_ok", "unique_marker", @@ -57,6 +76,139 @@ __all__ = [ ] +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 + + +class ProviderQuery(BaseModel): + provider: str + + +class CallbackLogMetadata(BaseModel): + user_api_key_hash: str + user_api_key_alias: str + user_api_key_user_id: str + + +class CallbackLogPayload(BaseModel): + id: str + litellm_call_id: str + model: str + call_type: str = "acompletion" + start_time: float = Field(serialization_alias="startTime") + end_time: float = Field(serialization_alias="endTime") + response_cost: float + prompt_tokens: int + completion_tokens: int + total_tokens: int + metadata: CallbackLogMetadata + + +class CallbackLogRecord(BaseModel): + status: str = "success" + standard_logging_payload: CallbackLogPayload + + +class CallbackLogsRequest(BaseModel): + records: list[CallbackLogRecord] + + +class CallbackLogsResponse(BaseModel): + processed: int + failed: 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 +359,166 @@ class SpendClient: 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: + _ = unwrap( + 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 retrieve_batch(self, key: str, batch_id: str, *, provider: str) -> BatchObject: + return unwrap( + self.proxy.transport.get( + f"/v1/batches/{batch_id}", + headers=self.proxy.transport.bearer(key), + params=ProviderQuery(provider=provider), + response_type=BatchObject, + ) + ) + + def replay_callback_log(self, key: str, payload: CallbackLogPayload) -> CallbackLogsResponse: + return unwrap( + self.proxy.transport.post( + "/v1/rust_control_plane/logs", + headers=self.proxy.transport.bearer(key), + json=CallbackLogsRequest(records=[CallbackLogRecord(standard_logging_payload=payload)]), + response_type=CallbackLogsResponse, + ) + ) + + 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( diff --git a/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py new file mode 100644 index 00000000000..4a2c23927c6 --- /dev/null +++ b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py @@ -0,0 +1,405 @@ +"""Every spend row a live proxy writes joins its virtual key (MAT-180). + +One virtual key with an alias, owned by a user with an email, drives every spend +write path a key can reach: /chat/completions, /queue/chat/completions, +/v1/messages, /v1/responses, /embeddings, the Gemini native passthrough, a batch +input file upload, a batch create, and a replayed callback log (POST +/v1/rust_control_plane/logs, the writer an external gateway feeds). Each row those calls write must carry +`api_key` equal to the key's LiteLLM_VerificationToken.token (the sha256 hash +/key/generate returns as `token`), which is the join /spend/logs?api_key= and +/user/daily/activity rely on to report key_alias and user_email. A row keyed by a +re-hashed token (v1.99.0's regression, #39568 and #39572) shows up as a +key-hash-* row with no alias and no email in the customer's usage exports. + +The health-check service account writes rows too; those must stay keyed by the +literal service-account name, never by a hash of it. A batch's cost row is +written by the retrieve that first sees the batch in a terminal state, so the +batch the run creates is one OpenAI fails at validation within seconds (its one +line targets /v1/embeddings under a /v1/chat/completions batch), and the test +retrieves it by its raw provider id with the same key until it is failed. A raw +id is never owned by the CheckBatchCost poller, so that retrieve prices the batch +inline against the retrieving key and its {provider_batch_id}_batch_cost row +must join the key's token with its alias. A completed batch with a positive +cost is out of a single run's reach (OpenAI's completion window is 24h, and a +stack booted fresh per run lists no earlier run's batches), so the poller's own +row is not asserted here. + +/spend/logs carries no email field, so the email assertion lives on +/user/daily/activity alone; /spend/logs is held to the alias in metadata. +""" + +import base64 +import time +from collections.abc import Iterator +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Final + +import pytest +from models import KeyGenerateBody +from proxy_client import Converged, await_converged +from pydantic import BaseModel +from spend_e2e_client import ( + BatchCreateBody, + BatchObject, + CallbackLogMetadata, + CallbackLogPayload, + DailyActivityKeyBreakdown, + ResponseIdentity, + SpendClient, + SpendLogRow, + StreamingResponse, + unique_marker, +) + +pytestmark = pytest.mark.e2e + +CHAT_MODEL: Final = "gemini-2.5-flash" +MESSAGES_MODEL: Final = "claude-haiku-4-5" +RESPONSES_MODEL: Final = "openai-responses-codex" +EMBED_MODEL: Final = "openai-text-embedding-3-small" +BATCH_MODEL: Final = "openai-gpt-4o-mini" +BATCH_BACKEND_MODEL: Final = "gpt-4o-mini" +BATCH_PROVIDER: Final = "openai" +HEALTH_SERVICE_ACCOUNT: Final = "litellm-internal-health-check" +BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "expired"}) +FAILED_BATCH_POLL_SECONDS: Final = 120.0 +FAILED_BATCH_POLL_INTERVAL_SECONDS: Final = 5.0 +MAX_TOKENS: Final = 8 +REPLAY_RESPONSE_COST: Final = 0.0001 +REPLAY_PROMPT_TOKENS: Final = 5 +REPLAY_COMPLETION_TOKENS: Final = 1 +WRITE_PATHS: Final = ( + "chat_completions", + "queue_chat_completions", + "messages", + "responses", + "embeddings", + "gemini_passthrough", + "batch_file_upload", + "batch_create", + "callback_replay", +) + + +class EmbeddingLineBody(BaseModel): + model: str + input: str + + +class EmbeddingLine(BaseModel): + custom_id: str + method: str = "POST" + url: str = "/v1/embeddings" + body: EmbeddingLineBody + + +@dataclass(frozen=True, slots=True) +class AttributedKey: + key: str + token: str + alias: str + email: str + user_id: str + + +@dataclass(frozen=True, slots=True) +class WritePath: + name: str + request_id: str + + +@dataclass(frozen=True, slots=True) +class DrivenKey: + identity: AttributedKey + paths: tuple[WritePath, ...] + started_at: datetime + + +def _body_id(name: str, sent: StreamingResponse) -> WritePath: + assert sent.ok, f"{name} failed with {sent.status_code}: {sent.body[:300]}" + response_id: Final = ResponseIdentity.model_validate_json(sent.body).id + assert response_id, f"{name} answered without a response id: {sent.body[:300]}" + return WritePath(name=name, request_id=response_id) + + +def _call_id(name: str, sent: StreamingResponse) -> WritePath: + assert sent.ok, f"{name} failed with {sent.status_code}: {sent.body[:300]}" + assert sent.call_id, f"{name} answered without an x-litellm-call-id header" + return WritePath(name=name, request_id=sent.call_id) + + +def _endpoint_mismatched_jsonl(marker: str) -> bytes: + line: Final = EmbeddingLine(custom_id=marker, body=EmbeddingLineBody(model=BATCH_BACKEND_MODEL, input=marker)) + return f"{line.model_dump_json()}\n".encode() + + +def _drive_batch(client: SpendClient, identity: AttributedKey, marker: str) -> tuple[WritePath, WritePath]: + uploaded: Final = client.upload_batch_file(identity.key, BATCH_MODEL, _endpoint_mismatched_jsonl(marker)) + created: Final = client.create_batch( + identity.key, + BatchCreateBody( + input_file_id=uploaded.id, + model=BATCH_MODEL, + metadata={"run": marker}, + ), + ) + return ( + WritePath(name="batch_file_upload", request_id=uploaded.id), + WritePath(name="batch_create", request_id=created.id), + ) + + +def _drive_callback_replay(client: SpendClient, identity: AttributedKey, marker: str) -> WritePath: + request_id: Final = f"callback-replay-{marker}" + finished_at: Final = time.time() + replayed: Final = client.replay_callback_log( + identity.key, + CallbackLogPayload( + id=request_id, + litellm_call_id=request_id, + model=CHAT_MODEL, + start_time=finished_at - 1, + end_time=finished_at, + response_cost=REPLAY_RESPONSE_COST, + prompt_tokens=REPLAY_PROMPT_TOKENS, + completion_tokens=REPLAY_COMPLETION_TOKENS, + total_tokens=REPLAY_PROMPT_TOKENS + REPLAY_COMPLETION_TOKENS, + metadata=CallbackLogMetadata( + user_api_key_hash=identity.token, + user_api_key_alias=identity.alias, + user_api_key_user_id=identity.user_id, + ), + ), + ) + assert replayed.processed == 1 and replayed.failed == 0, f"callback replay rejected the payload: {replayed}" + return WritePath(name="callback_replay", request_id=request_id) + + +def _drive_every_write_path(client: SpendClient, identity: AttributedKey) -> tuple[WritePath, ...]: + marker: Final = unique_marker() + prompt: Final = f"Reply with the word ok. {marker}" + key: Final = identity.key + return ( + _body_id("chat_completions", client.send_chat(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)), + _body_id("queue_chat_completions", client.send_queued_chat(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)), + _body_id("messages", client.send_messages(key, MESSAGES_MODEL, prompt, max_tokens=MAX_TOKENS)), + _body_id("responses", client.send_responses(key, RESPONSES_MODEL, prompt)), + _call_id("embeddings", client.send_embed(key, EMBED_MODEL, prompt)), + _call_id("gemini_passthrough", client.send_gemini_generate(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)), + *_drive_batch(client, identity, marker), + _drive_callback_replay(client, identity, marker), + ) + + +def _provider_batch_id(unified_batch_id: str) -> str: + encoded: Final = unified_batch_id.removeprefix("batch_") + decoded: Final = base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)).decode() + return decoded.removeprefix("litellm:").split(";", 1)[0] + + +def _driven_batch_id(driven: DrivenKey) -> str: + return next(path.request_id for path in driven.paths if path.name == "batch_create") + + +def _await_terminal_batch(client: SpendClient, key: str, provider_batch_id: str) -> BatchObject: + outcome: Final = await_converged( + lambda: client.retrieve_batch(key, provider_batch_id, provider=BATCH_PROVIDER), + converged=lambda batch: batch.status in BATCH_TERMINAL_STATUSES, + timeout=FAILED_BATCH_POLL_SECONDS, + interval=FAILED_BATCH_POLL_INTERVAL_SECONDS, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else outcome.last_result + + +def _health_rows_between(client: SpendClient, started_at: datetime) -> list[SpendLogRow]: + return [ + row + for row in client.proxy.spend_logs_window( + start=started_at - timedelta(minutes=1), end=datetime.now(timezone.utc) + timedelta(minutes=1) + ) + if HEALTH_SERVICE_ACCOUNT in (row.request_tags or []) + ] + + +def _health_rows_since(client: SpendClient, started_at: datetime) -> list[SpendLogRow]: + outcome: Final = await_converged( + lambda: _health_rows_between(client, started_at), + converged=lambda rows: bool(rows), + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else outcome.last_result + + +class TestKeyAttribution: + @pytest.fixture(scope="class") + def driven(self, client: SpendClient) -> Iterator[DrivenKey]: + marker: Final = unique_marker() + user_id: Final = client.create_user( + email=f"key-attribution-{marker}@example.com", + role="proxy_admin", + user_id=f"key-attribution-{marker}", + ) + record: Final = client.generate_key_record( + KeyGenerateBody(models=[], user_id=user_id, key_alias=f"key-attribution-{marker}") + ) + assert record.token, "/key/generate answered without the key's token hash" + assert record.key_alias, "/key/generate dropped the key alias" + identity: Final = AttributedKey( + key=record.key, + token=record.token, + alias=record.key_alias, + email=f"key-attribution-{marker}@example.com", + user_id=user_id, + ) + started_at: Final = datetime.now(timezone.utc) + try: + yield DrivenKey( + identity=identity, + paths=_drive_every_write_path(client, identity), + started_at=started_at, + ) + finally: + client.proxy.delete_key(identity.key) + client.delete_user(identity.user_id) + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.joins_key", + exercised_on=[ + "chat_completions", + "messages", + "responses", + "embeddings", + "batches", + "files", + "google_native", + "rust_control_plane", + ], + ) + def test_every_write_path_row_joins_the_key(self, client: SpendClient, driven: DrivenKey) -> None: + assert tuple(path.name for path in driven.paths) == WRITE_PATHS + found: Final = tuple((path, client.proxy.poll_logs_for_request_id(path.request_id)) for path in driven.paths) + unwritten: Final = [path.name for path, rows in found if not rows] + assert not unwritten, f"write paths that produced no spend row within the poll window: {unwritten}" + unjoined: Final = [ + (path.name, row.call_type, row.api_key) + for path, rows in found + for row in rows + if row.api_key != driven.identity.token + ] + assert not unjoined, ( + "spend rows whose api_key does not join LiteLLM_VerificationToken.token " + f"{driven.identity.token}: {unjoined}" + ) + unaliased: Final = [ + (path.name, row.call_type, row.metadata.user_api_key_alias if row.metadata else None) + for path, rows in found + for row in rows + if row.metadata is None or row.metadata.user_api_key_alias != driven.identity.alias + ] + assert not unaliased, f"spend rows written without key alias {driven.identity.alias!r}: {unaliased}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.reports_alias_and_email", + exercised_on=[ + "chat_completions", + "messages", + "responses", + "embeddings", + "batches", + "files", + "google_native", + "rust_control_plane", + ], + ) + def test_spend_logs_by_key_return_every_row_with_the_alias(self, client: SpendClient, driven: DrivenKey) -> None: + expected_ids: Final = frozenset(path.request_id for path in driven.paths) + rows: Final = client.poll_logs_for_key( + driven.identity.key, + min_rows=len(driven.paths), + predicate=lambda found: expected_ids <= frozenset(row.request_id or "" for row in found), + ) + missing: Final = expected_ids - frozenset(row.request_id or "" for row in rows) + assert not missing, ( + f"/spend/logs?api_key= does not return {len(missing)} of {len(expected_ids)} rows for the key: " + f"{sorted(path.name for path in driven.paths if path.request_id in missing)}" + ) + aliases: Final = frozenset(row.metadata.user_api_key_alias if row.metadata else None for row in rows) + assert aliases == {driven.identity.alias}, f"/spend/logs rows carry aliases {sorted(map(str, aliases))}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.reports_alias_and_email", + exercised_on=[ + "chat_completions", + "messages", + "responses", + "embeddings", + "batches", + "files", + "google_native", + "rust_control_plane", + ], + ) + def test_user_daily_activity_reports_alias_and_email(self, client: SpendClient, driven: DrivenKey) -> None: + breakdown: Final[DailyActivityKeyBreakdown | None] = client.poll_daily_activity_for_key( + driven.identity.token, + start=driven.started_at - timedelta(days=1), + end=datetime.now(timezone.utc) + timedelta(days=1), + min_requests=len(driven.paths), + ) + assert breakdown is not None, ( + f"/user/daily/activity?api_key={driven.identity.token} has no api_keys breakdown: " + "the key's rows did not aggregate under its token" + ) + assert breakdown.metrics.api_requests >= len(driven.paths), ( + f"/user/daily/activity counts {breakdown.metrics.api_requests} requests for the key, " + f"expected at least {len(driven.paths)}" + ) + assert breakdown.metadata.key_alias == driven.identity.alias, f"key_alias={breakdown.metadata.key_alias!r}" + assert breakdown.metadata.user_email == driven.identity.email, f"user_email={breakdown.metadata.user_email!r}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.health_rows_keep_service_account", + exercised_on=["chat_completions"], + ) + def test_health_check_rows_keep_the_service_account_key(self, client: SpendClient) -> None: + started_at: Final = datetime.now(timezone.utc) + probe: Final = client.health(CHAT_MODEL) + assert probe.healthy, f"/health?model={CHAT_MODEL} answered {probe.status_code}: {probe.body[:300]}" + rows: Final = _health_rows_since(client, started_at) + assert rows, f"/health?model={CHAT_MODEL} wrote no {HEALTH_SERVICE_ACCOUNT}-tagged spend row" + rehashed: Final = [(row.request_id, row.api_key) for row in rows if row.api_key != HEALTH_SERVICE_ACCOUNT] + assert not rehashed, f"health-check rows keyed by something other than {HEALTH_SERVICE_ACCOUNT!r}: {rehashed}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key", + exercised_on=["batches"], + ) + def test_terminal_batch_cost_row_joins_the_retrieving_key(self, client: SpendClient, driven: DrivenKey) -> None: + provider_batch_id: Final = _provider_batch_id(_driven_batch_id(driven)) + fetched: Final = _await_terminal_batch(client, driven.identity.key, provider_batch_id) + assert fetched.status == "failed", ( + f"endpoint-mismatched batch {provider_batch_id} is {fetched.status!r} after " + f"{FAILED_BATCH_POLL_SECONDS:.0f}s, so its terminal cost row cannot be asserted" + ) + cost_request_id: Final = f"{provider_batch_id}_batch_cost" + rows: Final = client.proxy.poll_logs_for_request_id(cost_request_id) + assert rows, f"retrieving failed batch {provider_batch_id} wrote no cost row under {cost_request_id}" + call_types: Final = tuple(sorted({row.call_type or "" for row in rows})) + assert call_types == ("aretrieve_batch",), f"cost rows under {cost_request_id} carry call types {call_types}" + unjoined: Final = [ + (row.call_type, row.api_key, row.metadata.user_api_key_alias if row.metadata else None) + for row in rows + if row.api_key != driven.identity.token + or row.metadata is None + or row.metadata.user_api_key_alias != driven.identity.alias + ] + assert not unjoined, ( + f"batch cost rows that do not join the retrieving key's token {driven.identity.token} " + f"with alias {driven.identity.alias!r}: {unjoined}" + )