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 123

fix(e2e): bill the batch cost baton through a raw-id retrieve and drive the callback replay writer

Browse source 
The completed marker batch is now retrieved by its raw provider id with this
run's key, so the proxy prices it inline and the {provider_batch_id}_batch_cost
row must join that key's token and alias. The CheckBatchCost poller only bills
batches it created in the same database, which a stack booted fresh per run
never holds for a completed marker, so the old unified-id assertion had no row
to find. Every run also replays one callback log through
POST /v1/rust_control_plane/logs, the third spend writer, and asserts its row
joins the key like the eight request paths
This commit is contained in:
mateo-berri 2026-09-09 15:54:11 -07:00
parent 3020a13e24
commit 07a0ca1074
3 changed files with 192 additions and 49 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

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

178
tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py
View file

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