From 3020a13e24f56a4068055eaadce039cd527f7faa Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:34:12 -0700 Subject: [PATCH] test(e2e): every spend row a virtual key writes joins its token across all write paths One aliased key owned by a user with an email drives chat, queued chat, messages, responses, embeddings, the Gemini passthrough, a batch file upload, and a batch create against a live proxy. Each row must carry api_key equal to the key's LiteLLM_VerificationToken.token and the alias in metadata, and /spend/logs?api_key= and /user/daily/activity must report the key with its alias and email. Health-check rows must keep the literal service-account key, and the batch cost row for a completed marker batch must join the key that created it. A re-hashed api_key (the v1.99.0 regression fixed by #39568 and #39572) now fails the Buildkite e2e stage naming the write path Resolves MAT-180 --- tests/e2e/CLAUDE.md | 5 +- .../coverage_registry/quota_management.yaml | 4 + tests/e2e/models.py | 2 + .../spend_tracking/conftest.py | 1 + .../spend_tracking/spend_e2e_client.py | 284 +++++++++++++- .../test_key_attribution_e2e.py | 349 ++++++++++++++++++ 6 files changed, 642 insertions(+), 3 deletions(-) create mode 100644 tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 89c04208d65..b92d84b8b25 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -181,13 +181,14 @@ 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 | 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] ``` diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index d0afcaca848..7013b59b83d 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/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"} 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..42ebfdb98af 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,14 @@ import time 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 @@ from models import ( DateRangeParams, EmbedBody, EmbedResponse, + KeyGenerateBody, + KeyGenerateResponse, OpenAPISchema, SpendCalculateBody, SpendCalculateResponse, @@ -43,13 +50,24 @@ from models import ( 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", @@ -57,6 +75,112 @@ __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 + 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 @@ 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: + _ = 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( 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..986033dbbde --- /dev/null +++ b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py @@ -0,0 +1,349 @@ +"""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, and a batch create. 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. The batch cost row is +written by the CheckBatchCost poller once the batch completes, up to an hour +later, so it rides a cross-run baton like the batches suite: each run submits a +one-line marker batch whose metadata records the token it expects on the cost +row, and asserts on the newest completed marker from any run (a cold start with +no completed marker is a documented vacuous pass, never a skip). + +/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 time +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Final, Iterator + +import pytest + +from models import ChatMessage, KeyGenerateBody +from proxy_client import Converged, await_converged +from spend_e2e_client import ( + BatchCreateBody, + BatchObject, + DailyActivityKeyBreakdown, + ResponseIdentity, + SpendClient, + SpendLogRow, + StreamingResponse, + unique_marker, +) +from pydantic import BaseModel + +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" +HEALTH_SERVICE_ACCOUNT: Final = "litellm-internal-health-check" +BATON_MARKER_KEY: Final = "litellm_e2e_suite" +BATON_MARKER_VALUE: Final = "key-attribution-baton" +BATON_EXPECTED_TOKEN_KEY: Final = "expected_api_key" +BATON_POLL_SECONDS: Final = 300.0 +BATON_POLL_INTERVAL_SECONDS: Final = 10.0 +BATON_LIST_LIMIT: Final = 100 +MAX_TOKENS: Final = 8 +WRITE_PATHS: Final = ( + "chat_completions", + "queue_chat_completions", + "messages", + "responses", + "embeddings", + "gemini_passthrough", + "batch_file_upload", + "batch_create", +) + + +class BatchLineBody(BaseModel): + model: str + messages: list[ChatMessage] + max_tokens: int + + +class BatchLine(BaseModel): + custom_id: str + method: str = "POST" + url: str = "/v1/chat/completions" + body: BatchLineBody + + +@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 _batch_jsonl(marker: str) -> bytes: + line: Final = BatchLine( + custom_id=marker, + body=BatchLineBody( + model=BATCH_BACKEND_MODEL, + messages=[ChatMessage(role="user", content=f"Reply with the word ok. {marker}")], + max_tokens=MAX_TOKENS, + ), + ) + 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, _batch_jsonl(marker)) + created: Final = client.create_batch( + identity.key, + BatchCreateBody( + input_file_id=uploaded.id, + model=BATCH_MODEL, + metadata={ + BATON_MARKER_KEY: BATON_MARKER_VALUE, + BATON_EXPECTED_TOKEN_KEY: identity.token, + "run": marker, + }, + ), + ) + return ( + WritePath(name="batch_file_upload", request_id=uploaded.id), + WritePath(name="batch_create", request_id=created.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), + ) + + +def _completed_baton(listed: list[BatchObject]) -> BatchObject | None: + return max( + ( + batch + for batch in listed + if batch.status == "completed" and (batch.metadata or {}).get(BATON_MARKER_KEY) == BATON_MARKER_VALUE + ), + key=lambda batch: batch.created_at or 0, + default=None, + ) + + +def _newest_completed_baton(client: SpendClient, key: str) -> BatchObject | None: + outcome: Final = await_converged( + lambda: _completed_baton(client.list_batches(key, BATCH_MODEL, limit=BATON_LIST_LIMIT)), + converged=lambda completed: completed is not None, + timeout=BATON_POLL_SECONDS, + interval=BATON_POLL_INTERVAL_SECONDS, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else None + + +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"], + ) + 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"], + ) + 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"], + ) + 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.batch_cost_joins_key", + exercised_on=["batches"], + ) + def test_completed_batch_cost_row_joins_the_key(self, client: SpendClient, driven: DrivenKey) -> None: + completed: Final = _newest_completed_baton(client, driven.identity.key) + if completed is None: + return + expected_token: Final = (completed.metadata or {}).get(BATON_EXPECTED_TOKEN_KEY) + assert expected_token, f"marker batch {completed.id} lost its expected token metadata: {completed.metadata}" + fetched: Final = client.retrieve_batch(driven.identity.key, completed.id) + assert fetched.status == "completed", f"listed-completed marker retrieved as {fetched.status!r}" + rows: Final = client.proxy.poll_logs_for_request_id( + f"{fetched.id}_batch_cost", + predicate=lambda found: any((row.spend or 0) > 0 for row in found), + ) + priced: Final = [row for row in rows if (row.spend or 0) > 0] + assert priced, f"completed batch {fetched.id} has no positive-cost spend row under {fetched.id}_batch_cost" + unjoined: Final = [(row.call_type, row.api_key) for row in priced if row.api_key != expected_token] + assert not unjoined, f"batch cost rows whose api_key does not join the key's token {expected_token}: {unjoined}"