scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

Merge pull request #40465 from BerriAI/litellm_e2e_spend_rows_join_virtual_key

Browse source 
test(e2e): every spend row a virtual key writes joins its token across all write paths
This commit is contained in:
Mateo Wang 2026-09-09 19:25:41 -07:00 • committed by GitHub
commit 9839bfcdbd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 730 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
tests/e2e/CLAUDE.md
View file

@ -181,13 +181,15 @@ 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 | 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]
```

5
tests/e2e/coverage_registry/quota_management.yaml
View file

@ -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"}

2
tests/e2e/models.py
View file

@ -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

1
tests/e2e/quota_management/spend_tracking/conftest.py
View file

@ -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"),
)

314
tests/e2e/quota_management/spend_tracking/spend_e2e_client.py
View file

@ -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(

405
tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py Normal file
View file

@ -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}"
)