fix(batch): track cost for managed batches with no attributable key/u… - #35468
Conversation
Greptile SummaryThe PR restores spend persistence for completed managed batches that have no attributable key, user, team, or end user.
Confidence Score: 5/5The PR appears safe to merge. No blocking failure remains.
|
| Filename | Overview |
|---|---|
| enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py | Propagates the batch’s creation-time team ID into the synthetic completed-batch logging event. |
| litellm/proxy/hooks/proxy_track_cost_callback.py | Extends unattributed spend tracking to completed managed-batch retrieval events while preserving the global spend-update disable switch. |
| tests/proxy_unit_tests/test_check_batch_cost.py | Adds an end-to-end mocked regression test proving an unattributed completed batch reaches the spend database writer with its calculated cost. |
| tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py | Extends callback guard and database-write tests to cover unattributed aretrieve_batch events. |
Reviews (2): Last reviewed commit: "Merge branch 'litellm_internal_staging' ..." | Re-trigger Greptile
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
…ser/team LiteLLM_ManagedObjectTable only stores created_by (user_id) and team_id, never the raw API key hash. A batch created with the master key or a team-less key has both null, so CheckBatchCost's synthetic logging_obj for the completed batch carried no attributable key/user/team/end-user. _should_track_cost_callback silently skipped the DB write in that case (by design, to avoid tracking truly anonymous requests), with no error or warning: batch_processed still became true, but no LiteLLM_SpendLogs row was ever written despite real, already-incurred provider cost. Extend the same allowance already made for unauthenticated pass-through requests to aretrieve_batch's cost event, and pass job.team_id through so a batch's team gets real attribution when one exists.
62478ea to
833670f
Compare
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
| # user_api_key_team_id (from .team_id) -- both are None for batches created with | ||
| # the master key or a team-less key, since the table never stores the raw key | ||
| # hash. The batch already incurred real provider cost, so track it regardless. | ||
| CallTypes.aretrieve_batch.value, |
There was a problem hiding this comment.
Low: Unauthenticated batch retrieval can inflate recorded spend
This allowlist entry also covers the public retrieve-batch route. With proxy authentication disabled, an attacker who repeatedly retrieves a completed batch causes async_success_handler to recompute and persist the entire batch cost on every GET, even though that cost was not incurred again; these writes also update the global proxy-budget row when a maximum budget is configured. Use a distinct internal call type or another non-client-controllable marker for the synthetic CheckBatchCost event rather than allowing all unattributed aretrieve_batch calls.
There was a problem hiding this comment.
Auth-disabled proxies already let anyone incur unlimited real spend, so inflating recorded spend adds no marginal risk. The aggregate re-increment predates this PR
PR overviewThis PR updates proxy cost tracking so managed batch costs are recorded even when they cannot be attributed to a specific key or user. It adds batch retrieval handling to the proxy cost callback. One security issue remains open: when proxy authentication is disabled, repeated retrieval of a completed batch can repeatedly record its full cost and inflate spend and global budget accounting. Exploitation depends on the public batch retrieval route being exposed without authentication, but it could disrupt accounting and budget enforcement in that configuration. Open issues (1)
Fixed/addressed: 0 · PR risk: 4/10 |
a0e627f
into
BerriAI:litellm_internal_staging
2 participants
TLDR
Problem this solves:
How it solves it:
User Flow
Before: the batch completes, but its spend never shows up anywhere on the gateway
After: the same batch shows up as billed spend shortly after it completes
Relevant issues
Fixes #35358
Linear ticket
Resolves LIT-5183
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
Live proxies against real OpenAI batch spend, one per build, each on its own fresh postgres database. Before runs at staging head 0acca3e, after runs at this PR's head e0c4c7c. Config pins
gpt-5.4-nano-batch(openai/gpt-5.4-nano, mode batch) andproxy_batch_polling_interval: 30Two repro details matter. Polling
GET /v1/batches/{id}to completion masks the bug on both builds, because the completion-detecting GET writes the cost itself attributed to the calling key, so after batch creation no further HTTP calls touch the batch and only the background CheckBatchCost job can observe completion. And the issue title's master-key case already bills at staging head, since a master-key batch storescreated_by = default_user_id, which passes the pre-fix guard; the surviving gap is a key with no user and no team, where the managed-object row hascreated_by NULL, team_id NULLSame steps on both proxies (before on 41783, after on 45291):
No further requests are made. Both proxy logs then show the poller picking the completed batch up, with the same all-None attribution reaching the track-cost guard:
Spend logs afterwards, queried several poller cycles later. The before database has no billed row for that batch, the after database has exactly one:
The billed amount is exact for gpt-5.4-nano batch pricing: 26 prompt tokens at 1e-07 plus 8 completion tokens at 6.25e-07 is 7.6e-06 USD. A master-key control batch billed the same 7.6e-06 on both builds, and no batch double-billed anywhere, including one that was both poller-reconciled and then retrieved over HTTP
Type
🐛 Bug Fix
Changes
_should_track_cost_callbackskips any cost event that carries no key, user, team, or end-user, with a carve-out for unauthenticated pass-through routes. CheckBatchCost attributes its synthetic logging object fromLiteLLM_ManagedObjectTable, which stores onlycreated_byandteam_idand never the raw key hash, so a batch created by a key with no user and no team reaches that guard all-None and its cost event was silently dropped after the batch was already marked processedThe carve-out set is renamed from
_PASS_THROUGH_CALL_TYPESto_UNATTRIBUTED_TRACKABLE_CALL_TYPESand now includesaretrieve_batch: the batch already incurred real provider cost, so it must be tracked regardless of attribution, the same reasoning as pass-through. The synthetic logging object also now carriesuser_api_key_team_idfrom the job, so a team-owned batch gets real team attribution instead of falling through to the track-anyway pathRegression tests:
test_completed_batch_with_no_attributable_owner_still_writes_spend_logruns the real logging pipeline through_ProxyDBLogger(the layer the bug hid behind) and asserts the spend write happens and the job is still marked processed, and the two parametrized guard tests intest_proxy_track_cost_callback.pypinaretrieve_batchas tracked without an owner while other unattributed call types stay skippedFinal Attestation