diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 7013b59b83d..d3966eed48f 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -58,7 +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.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.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"} +- {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: "proxy/batches_endpoints/endpoints.py", rationale: "Retrieving a completed batch by its raw provider id prices it inline against the retrieving key, so the {provider_batch_id}_batch_cost row must carry that key's token hash and alias; the target is the newest completed cross-run marker batch this proxy has not billed yet, since completion lags by up to 24h, and a cold start with no completed marker is a documented vacuous pass. The CheckBatchCost poller's own unified-id row needs the batch create row in the same database, which a stack booted fresh per run never holds for a completed marker"} 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 42ebfdb98af..18b714d8aeb 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -17,8 +17,6 @@ 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, @@ -57,9 +55,12 @@ from models import ( UserRole, ) from proxy_client import Converged, ProxyClient, await_converged +from pydantic import BaseModel, Field __all__ = [ "BatchCreateBody", + "CallbackLogMetadata", + "CallbackLogPayload", "BatchObject", "DailyActivityKeyBreakdown", "FileObject", @@ -147,6 +148,44 @@ class BatchListQuery(BaseModel): limit: int +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 @@ -445,16 +484,26 @@ class SpendClient: ) ).data - def retrieve_batch(self, key: str, batch_id: str) -> 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=NoBody(), + 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)) 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 index 986033dbbde..cac4fd7fb3c 100644 --- a/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py @@ -3,7 +3,8 @@ 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 +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 @@ -11,29 +12,40 @@ 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). +literal service-account name, never by a hash of it. A batch's cost row lands +only once the batch completes, up to 24h later, so the batch cost path rides a +cross-run baton like the batches suite: each run submits a one-line marker batch +and never cancels it, and the newest completed marker from any run that this +proxy has not billed yet is retrieved by its raw provider id with this run's +key. That retrieve is the writer under test: a raw id is never poller-owned, so +the proxy prices it inline against the retrieving key, and the +{provider_batch_id}_batch_cost row must join this run's token with its alias. +The CheckBatchCost poller's own row (the unified id, billed against the +submitting key) needs the batch's create row in the same database, which a stack +booted fresh per run never holds for a completed marker, so that writer is out +of this test's reach. 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 base64 import time +from collections.abc import Iterator from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Final, Iterator +from typing import Final import pytest - -from models import ChatMessage, KeyGenerateBody +from models import ChatMessage, KeyGenerateBody, SpendLogsParams from proxy_client import Converged, await_converged +from pydantic import BaseModel from spend_e2e_client import ( BatchCreateBody, BatchObject, + CallbackLogMetadata, + CallbackLogPayload, DailyActivityKeyBreakdown, ResponseIdentity, SpendClient, @@ -41,7 +53,6 @@ from spend_e2e_client import ( StreamingResponse, unique_marker, ) -from pydantic import BaseModel pytestmark = pytest.mark.e2e @@ -51,14 +62,17 @@ 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" 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_SECONDS: Final = 30.0 BATON_POLL_INTERVAL_SECONDS: Final = 10.0 BATON_LIST_LIMIT: Final = 100 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", @@ -68,6 +82,7 @@ WRITE_PATHS: Final = ( "gemini_passthrough", "batch_file_upload", "batch_create", + "callback_replay", ) @@ -138,11 +153,7 @@ def _drive_batch(client: SpendClient, identity: AttributedKey, marker: str) -> t BatchCreateBody( input_file_id=uploaded.id, model=BATCH_MODEL, - metadata={ - BATON_MARKER_KEY: BATON_MARKER_VALUE, - BATON_EXPECTED_TOKEN_KEY: identity.token, - "run": marker, - }, + metadata={BATON_MARKER_KEY: BATON_MARKER_VALUE, "run": marker}, ), ) return ( @@ -151,6 +162,32 @@ def _drive_batch(client: SpendClient, identity: AttributedKey, marker: str) -> t ) +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}" @@ -163,31 +200,52 @@ def _drive_every_write_path(client: SpendClient, identity: AttributedKey) -> tup _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 _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 _completed_batons(listed: list[BatchObject]) -> tuple[BatchObject, ...]: + return tuple( + sorted( + ( + 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, + reverse=True, + ) ) -def _newest_completed_baton(client: SpendClient, key: str) -> BatchObject | None: +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 _batch_cost_request_id(unified_batch_id: str) -> str: + return f"{_provider_batch_id(unified_batch_id)}_batch_cost" + + +def _newest_unbilled_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, + lambda: _completed_batons(client.list_batches(key, BATCH_MODEL, limit=BATON_LIST_LIMIT)), + converged=lambda completed: bool(completed), timeout=BATON_POLL_SECONDS, interval=BATON_POLL_INTERVAL_SECONDS, now=time.monotonic, sleep=time.sleep, ) - return outcome.result if isinstance(outcome, Converged) else None + completed: Final = outcome.result if isinstance(outcome, Converged) else () + return next( + ( + batch + for batch in completed + if not client.proxy.spend_logs(SpendLogsParams(request_id=_batch_cost_request_id(batch.id))) + ), + None, + ) def _health_rows_between(client: SpendClient, started_at: datetime) -> list[SpendLogRow]: @@ -246,7 +304,16 @@ class TestKeyAttribution: @pytest.mark.covers( "quota_management.spend_tracking.key_attribution.joins_key", - exercised_on=["chat_completions", "messages", "responses", "embeddings", "batches", "files", "google_native"], + 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 @@ -275,7 +342,16 @@ class TestKeyAttribution: @pytest.mark.covers( "quota_management.spend_tracking.key_attribution.reports_alias_and_email", - exercised_on=["chat_completions", "messages", "responses", "embeddings", "batches", "files", "google_native"], + 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) @@ -294,7 +370,16 @@ class TestKeyAttribution: @pytest.mark.covers( "quota_management.spend_tracking.key_attribution.reports_alias_and_email", - exercised_on=["chat_completions", "messages", "responses", "embeddings", "batches", "files", "google_native"], + 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( @@ -332,18 +417,27 @@ class TestKeyAttribution: 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) + completed: Final = _newest_unbilled_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) + provider_batch_id: Final = _provider_batch_id(completed.id) + fetched: Final = client.retrieve_batch(driven.identity.key, provider_batch_id, provider=BATCH_PROVIDER) assert fetched.status == "completed", f"listed-completed marker retrieved as {fetched.status!r}" + cost_request_id: Final = _batch_cost_request_id(completed.id) rows: Final = client.proxy.poll_logs_for_request_id( - f"{fetched.id}_batch_cost", + cost_request_id, 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}" + assert priced, f"retrieving completed batch {provider_batch_id} wrote no positive-cost row under {cost_request_id}" + unjoined: Final = [ + (row.call_type, row.api_key, row.metadata.user_api_key_alias if row.metadata else None) + for row in priced + 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}" + )