Repository navigation
test(e2e): every spend row a virtual key writes joins its token across all write paths - #40465
Conversation
…s 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
🤖 Devin AI EngineerI'll be helping with this pull request! Here's what you should know: ✅ I will automatically:
Note: I can only respond to comments from users who have write access to this repository. ⚙️ Control Options:
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR adds live end-to-end regression coverage ensuring spend rows retain virtual-key attribution across nine write paths and usage-reporting APIs.
Confidence Score: 5/5The PR appears safe to merge; no actionable defect or outstanding previous finding remains. The recent cleanup only removes unused helpers and makes teardown enforce the existing integer
|
| Filename | Overview |
|---|---|
| tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py | Adds comprehensive live spend-attribution checks across supported write paths, reporting APIs, health checks, and terminal batch retrieval. |
| tests/e2e/quota_management/spend_tracking/spend_e2e_client.py | Adds typed helpers for the new attribution scenarios, validates user deletion during teardown, and removes unused batch-list helpers. |
| tests/e2e/coverage_registry/quota_management.yaml | Registers the new attribution coverage and accurately leaves poller-based completed-batch accounting unclaimed. |
| tests/e2e/models.py | Extends E2E response models with the generated token and spend-row key alias required by the assertions. |
| tests/e2e/quota_management/spend_tracking/conftest.py | Registers the OpenAI batch model needed by the new live batch scenario. |
| tests/e2e/CLAUDE.md | Documents the new key-attribution coverage vocabulary and separate batch-accounting assertions. |
Reviews (5): Last reviewed commit: "test(e2e): drop the unused batch list cl..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
|
bugbot run |
Devin e2e QATested at: 3020a13 Setup shared by both legs: a proxy on a random free port (20656) with a fresh Postgres DB per leg, Read-back after the driver: Before (096984b)Every write path by hand, rows read back
The new module against the proxy
After (3020a13)Every write path by hand, rows read back
The new module against the proxy
|
…ve the callback replay writer
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
GitHub Actions
Error
…d of a cross-run baton
The batch list is served from LiteLLM_ManagedObjectTable whenever the managed
files hook is loaded, and the Buildkite e2e stacks bundle a fresh Postgres per
build, so a prior run's marker batch is never listed and the baton could only
ever pass vacuously. Each run now creates a batch OpenAI fails at validation
within seconds, retrieves it by its raw provider id with the same key until it
is failed, and asserts the {provider_batch_id}_batch_cost row that retrieve
writes joins the key's token hash and alias
GitHub Actions
Error
| metadata=CallbackLogMetadata( | ||
| user_api_key_hash=identity.token, | ||
| user_api_key_alias=identity.alias, | ||
| user_api_key_user_id=identity.user_id, | ||
| ), |
There was a problem hiding this comment.
The callback replay supplies the expected token hash and alias in its payload, then checks that those same values appear in the spend row. The endpoint copies these payload fields into spend attribution instead of deriving them from the authenticated key, so this test stays green even if authenticated-key attribution is broken. As a result, it does not provide the claimed regression coverage for this write path. Exercise a payload/auth mismatch or assert behavior that the request did not provide verbatim.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
There was a problem hiding this comment.
The proxy's hash handling still owns this row's api_key: the mutated proxy rehashes it and this leg goes red, so it is not self-fulfilling
There was a problem hiding this comment.
You're right for the regression this test is intended to guard. The callback payload supplies the expected token, but the persisted spend row still passes through the proxy's normal key-handling/redaction path; a mutation that re-hashes api_key changes identity.token into a different value, so test_every_write_path_row_joins_the_key fails. The alias assertion also verifies the callback metadata survives persistence. My original comment overstated the issue: this test does not prove the endpoint derives attribution from authentication rather than trusting payload metadata, but it does provide valid coverage for the row-key rehash regression described by MAT-180. I’m withdrawing the finding.
| 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) |
There was a problem hiding this comment.
This test now creates a batch that must fail and checks inline accounting from a raw-ID retrieve. That is a different path from completed, positive-cost batch accounting, so a regression limited to completed-batch or poller attribution can pass this P1 coverage cell. Keep separate completed-batch coverage, or register this narrower failed-retrieve scenario without using it as a replacement.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
The retrieve writer and the CheckBatchCost poller are two different spend writers, and the in-run failed batch only proves the first. Split the key_attribution batch cell in two: retrieve_batch_cost_joins_retrieving_key, which test_terminal_batch_cost_row_joins_the_retrieving_key claims, and poller_batch_cost_joins_creating_key, which no test claims yet and so shows up on the coverage dashboard as a P1 gap instead of hiding behind the retrieve leg. The rationale records why one run cannot hand the poller a completed batch on a stack that boots a fresh Postgres per build.
GitHub Actions
Error
…r teardown miss The failed-batch redesign left list_batches, BatchList, BatchListQuery and the batch object's metadata and created_at fields with no caller, and they duplicated the batches suite's own client. delete_user discarded its result, so a user that outlived the class fixture went unnoticed; unwrap turns that into a teardown error like delete_key already does.
GitHub Actions
Failure
|
bugbot run |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit 71cf1c7. Configure here.
This branch had an error being deployed
1 participant
TLDR
Problem this solves:
How it solves it:
User Flow
Before: an admin exporting usage per key after a proxy upgrade finds every request under a scrambled key with no alias or owner, and no CI stage had flagged it
{"key_alias": "team-search", "user_id": "sam"}(sam hasuser_emailset) and get backkeyand a 64-charactertoken[]api_keysbreakdown has no entry for the tokenkey_aliasmissing, so the export attributes the spend to nobodyAfter: the same upgrade turns the Buildkite e2e stage red before it ships, naming the write path whose rows stopped joining the key
{"key_alias": "team-search", "user_id": "sam"}and get backkeyand a 64-charactertokenapi_keyequal totokenandmetadata.user_api_key_aliasequal toteam-searchkey_alias: team-searchanduser_email: sam@...tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.pywith the write path and the offendingapi_keyin the assertion message, so the release never reaches themRelevant issues
Guards against the shape of #39568 and #39572
Linear ticket
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
uv run pytest tests/test_litellm/<your_test_file>.py -v. Leave the suites (make test-unit-*,make test-unit) to CI: it finishes in ~15 minutes where a laptop takes an hour or more@greptileaito re-request a review after pushing changes)Screenshots / Proof of Fix
Setup shared by both legs: a proxy on a random free port (42270) with a fresh Postgres DB,
store_model_in_db: true,PROXY_BATCH_POLLING_ENABLED=true, and config modelsgemini-2.5-flash,claude-haiku-4-5,openai-responses-codex(openai/gpt-5.3-codex),openai-text-embedding-3-small,gpt-4o-mini. Real provider spend on every leg.$Mis the master key,$Bishttp://localhost:42270,$JisContent-Type: application/json. The "broken proxy" in both legs is the same 2-line mutation of the spend-row key redaction that v1.99.0 shipped in effect: every key, including the health-check service account, gets hashed once more before the row is writtenCreate the user with an email and the key with an alias:
Drive every write path with that key (the same commands the new test runs):
Then read the rows back:
GET $B/spend/logs?api_key=$KEY,GET $B/spend/logsfiltered to rows carrying the alias or the token (printed ascall_type, whetherapi_keyequals the key'stoken,metadata.user_api_key_alias), the newest health-check row, andGET $B/user/daily/activity?start_date=2026-09-08&end_date=2026-09-10&api_key=<sha256 of $KEY>Before (096984b)
Every write path on the broken proxy, existing suite green
tests/e2e/quota_management/spend_tracking/at this commit has no test that asserts a spend row'sapi_keyequals the key's token, so the e2e stage passes with every row unattributedAfter (654afb5; 5db99ca only renames the coverage cell the batch test claims, 71cf1c7 only drops unused client code and makes the user teardown assert its response; the module re-ran against the head proxy at 71cf1c7, 5 passed in 91.55s)
Every write path on the broken proxy, new test red
/spend/logs?api_key=<key> rows: 0, sameapi_key == tokenFalse on every row, same missing daily activity entryEvery write path on the proxy at head, new test green
The two writers added after review, on the proxy at head
Replay one finished call through the callback log writer with the key (the payload names the key's own token hash, the way an external gateway reports it), and retrieve a batch OpenAI failed at validation by its raw provider id with the same key.
$TOKENis the sha256 of$KEY, the value /key/generate returns astokenOutput (rows printed as
call_type, whetherapi_keyequals$TOKEN,metadata.user_api_key_alias,spend)The validating retrieve wrote nothing; the first retrieve that saw the batch failed wrote the single
<provider batch id>_batch_costrow, keyed by the retrieving key's token with its alias. A completed batch retrieved the same way on this proxy bookedspend 2.1e-06under the same shape, so the writer is one and the same for every terminal statusType
✅ Test
Caveats (if any)
The batch cost leg asserts on a batch OpenAI fails at validation, so the row it checks is a $0
aretrieve_batchrow. The completed, positive-cost case runs through the same retrieve writer (shown above on this proxy), but a completed batch is out of a single run's reach: OpenAI's completion window is 24h, andGET /v1/batchesis served fromLiteLLM_ManagedObjectTablewhenever the managed files hook is loaded, so a stack booted with a fresh Postgres per build (both Buildkite e2e stacks) never lists an earlier run's batch. That also means the CheckBatchCost poller's own row, which bills a unified batch id against the key that created it, has no batch to bill in-run within the poll budget, so it is not asserted; it is registered as its own P1 cell,quota_management.spend_tracking.key_attribution.poller_batch_cost_joins_creating_key, which no test claims, so the coverage dashboard shows the gap instead of hiding it behind the retrieve leg (follow-up ticket filed). The batches suite'sTestBatchTerminalStatebaton relies on the same list and has the same blind spot; that is a separate ticket. Until that follow-up lands, a regression confined to the poller writer (CheckBatchCost, which only picks up batches that reached a completed status) leaves this suite green; the proof's mutation re-hashes every row, which is wider than a poller-only breakThe nine paths are driven once by a class-scoped fixture, so a provider failure on any one of them errors the four tests that read its rows instead of failing only that path, and a row the proxy never writes costs the 120s per-request-id poll budget before the join assertion names it
Spend writers this stack cannot drive stay uncovered: MCP tool-call rows (needs an MCP server registered), web search interception rows (needs the interception config), semantic text index rows, and realtime sessions. The health-check leg covers the service-account writer, and the callback replay covers the third v1.99.0 writer
QA runbook
Prerequisites: a proxy with a database,
store_model_in_db: true,PROXY_BATCH_POLLING_ENABLED=true, and credentials for Gemini (gemini-2.5-flash), Anthropic (claude-haiku-4-5), and OpenAI (openai-responses-codex,openai-text-embedding-3-small,openai-gpt-4o-mini; the suite registers missing driver models itself)./queue/chat/completionsonly accepts a key whose user isproxy_adminand needspriorityin the body, and the Gemini passthrough takes the key inx-goog-api-key, notAuthorizationtests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py::TestKeyAttribution::test_every_write_path_row_joins_the_key - every spend row from every write path carries the key's own token hash and alias
{"user_id":"qa-attr","user_email":"qa-attr@example.com","user_role":"proxy_admin","auto_create_key":false}{"key_alias":"qa-attr","user_id":"qa-attr"}and keep bothkeyandtokenfrom the response"priority": 0), POST /v1/messages (claude-haiku-4-5), POST /v1/responses (openai-responses-codex), POST /embeddings (openai-text-embedding-3-small), POST /gemini/v1beta/models/gemini-2.5-flash:generateContent withx-goog-api-key: <key>, POST /v1/files?model=openai-gpt-4o-mini with a one-line batch jsonl, and POST /v1/batches with that file id; expect 200 on all eight and keep each response'sid(embeddings and passthrough: thex-litellm-call-idresponse header)records[]entry whosestandard_logging_payloadnames a freshid,model,response_cost, andmetadata.user_api_key_hashequal totoken; expect{"processed":1,"failed":0}api_keyequalstokenexactly and whosemetadata.user_api_key_aliasisqa-attrtests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py::TestKeyAttribution::test_spend_logs_by_key_return_every_row_with_the_alias - filtering spend logs by the raw key returns all nine rows, every one with the alias
metadata.user_api_key_aliasequal toqa-attron every row (nonull, no other alias)tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py::TestKeyAttribution::test_user_daily_activity_reports_alias_and_email - the daily activity report lists the key under its token with the alias and the owner's email
results[].breakdown.api_keys[<token>]present withmetrics.api_requestsof at least 9,metadata.key_aliasequal toqa-attr, andmetadata.user_emailequal toqa-attr@example.comtests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py::TestKeyAttribution::test_health_check_rows_keep_the_service_account_key - a health check's spend row is keyed by the literal service account name, never by a hash of it
request_tagscontainslitellm-internal-health-checkapi_keyto be exactlylitellm-internal-health-checktests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py::TestKeyAttribution::test_terminal_batch_cost_row_joins_the_retrieving_key - the cost row written when a batch is retrieved in a terminal state joins the retrieving key
litellm:batch_<provider id>;model,<name>after thebatch_prefix)statusisfailed(the test allows 120s)aretrieve_batchrows, every one withapi_keyequal totokenandmetadata.user_api_key_aliasequal toqa-attrFinal Attestation
Link to Devin session: https://app.devin.ai/sessions/5b21193c1cd64d7f9d94b0033ca087fe
Open in Devin Desktop: https://app.devin.ai/desktop/session/5b21193c1cd64d7f9d94b0033ca087fe?variant=devin
Requested by: @mateo-berri
Note
Low Risk
Changes are limited to e2e tests, client models, and coverage docs; no proxy or spend-tracking production code is modified.
Overview
Adds e2e coverage for virtual-key spend attribution, guarding regressions like #39568/#39572 where spend rows stop joining
LiteLLM_VerificationToken.tokenand usage exports show unattributedkey-hash-*rows.A new
test_key_attribution_e2e.pydrives one aliased key (with an owning user email) through nine spend write paths—chat, queued chat, messages, responses, embeddings, Gemini passthrough, batch file/create, and rust control-plane callback replay—and asserts each row’sapi_keyequals/key/generate’stokenand carriesmetadata.user_api_key_alias. Separate cases cover/spend/logs?api_key=,/user/daily/activity(alias + email), health rows staying onlitellm-internal-health-check, and inline batch cost on terminal retrieve ({provider_id}_batch_cost).Harness updates:
SpendClientgains typed helpers for those routes and daily-activity polling;models.pyaddsKeyGenerateResponse.tokenand spend-log alias metadata; spend-tracking coverage registry documents fivekey_attributioncells (poller batch cost remains an explicit uncovered gap).Reviewed by Cursor Bugbot for commit 71cf1c7. Bugbot is set up for automated code reviews on this repo. Configure here.