fix(batches): attribute Vertex passthrough batch cost to key/team/tags - #34456
Conversation
Greptile SummaryAdds creator attribution for Vertex passthrough batch costs.
Confidence Score: 5/5The PR appears safe to merge. No blocking failures remain within the scope of the previous review threads.
|
| Filename | Overview |
|---|---|
| enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py | Reconstructs delayed batch-cost attribution from the creator identity persisted on the managed-object row. |
| enterprise/litellm_enterprise/proxy/hooks/managed_files.py | Adds write-once key and tag attribution during creation and an update-only path for batch observations. |
| litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py | Distinguishes Vertex batch creation from polling and schedules creator-attributed registration with explicit failure logging. |
| litellm/proxy/hooks/proxy_track_cost_callback.py | Prevents delayed batch spend from replacing captured identity with the virtual key’s current ownership. |
| litellm-proxy-extras/litellm_proxy_extras/migrations/20260730000000_add_api_key_and_request_tags_to_managed_object_table/migration.sql | Adds nullable key and JSON tag columns used to retain batch creator attribution. |
Reviews (21): Last reviewed commit: "fix(batches): keep batch state in sync o..." | Re-trigger Greptile
Greptile SummaryThis PR adds create-time identity and tag snapshots for Vertex passthrough batch cost attribution, strips those internal snapshots from API responses, and preserves them across most managed-object rewrites Confidence Score: 3/5The PR should not merge until initial snapshot creation is made atomic so concurrent batch activity cannot replace the creator attribution Sequential writers preserve attribution correctly, but concurrent first-time stores can both observe no existing row and allow the later upsert to persist a different caller's identity; terminal polling also drops the snapshot after spend is logged enterprise/litellm_enterprise/proxy/hooks/managed_files.py, enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
|
| Filename | Overview |
|---|---|
| enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py | Reads snapshot identity for spend logging, but terminal status writes still remove the persisted snapshot |
| enterprise/litellm_enterprise/proxy/hooks/managed_files.py | Adds response stripping and snapshot carry-forward, but the read-before-upsert sequence does not make first-writer attribution immutable under concurrency |
| litellm/proxy/openai_files_endpoints/common_utils.py | Adds focused helpers for reading, stripping, and merging internal attribution and applies preservation during batch updates |
| litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py | Captures authenticated identity and tags on Vertex batch objects before asynchronous managed-object storage |
| litellm/proxy/pass_through_endpoints/managed_id_rewriter.py | Preserves the stored attribution while refreshing managed batch snapshots |
| litellm/proxy/pass_through_endpoints/pass_through_endpoints.py | Reasserts authenticated identity after merging client metadata |
| tests/pass_through_unit_tests/test_pass_through_unit_tests.py | Verifies client metadata cannot replace authenticated passthrough identity |
| tests/proxy_unit_tests/test_check_batch_cost.py | Covers snapshot extraction and spend metadata but not preservation by terminal poller writes |
| tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py | Covers sequential snapshot preservation and response stripping but does not exercise concurrent initial stores |
| tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py | Covers identity and tag snapshot construction for Vertex batch responses |
Comments Outside Diff (1)
-
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py, line 663-672 (link)Terminal writes discard attribution
The terminal poller writes
response.model_dump_json()without the preservation helper, permanently removing the create-time attribution snapshot after normal processing and leaving later audit or reprocessing paths without the captured identityKnowledge Base Used: Cost Tracking and Budget Enforcement
Reviews (2): Last reviewed commit: "fix(batches): make batch-cost attributio..." | Re-trigger Greptile
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 3 · PR risk: 0/10 |
|
bugbot run |
|
The concurrent-initial-store race is closed in 54acf65. The passthrough handlers now register the managed object only on the create route (POST |
|
The fix looks correct. Both handlers now gate registration on the collection endpoint:
This fully closes the race: since only the creating call ever reaches The two original concerns (initial attribution race and non-atomic initialization) are addressed. No remaining issues from my earlier comments. |
|
On the remaining concurrent-initial-insert concern: after 54acf65 there is no code path that inserts a batch row with a non-creator identity, so the write-once columns cannot be set to a competing caller. There are four So while If you'd still prefer a belt-and-suspenders change, the passthrough create store could be awaited rather than dispatched via |
|
Verified live against the real Anthropic Message Batches API (no mocks) that a poll cannot take over a batch's attribution. Setup: local proxy on :4000, Postgres, key A (
So no poll registers or rewrites the row; only the create does. Combined with the columns being written solely by the upsert create branch, a competing caller cannot own the batch spend. This is the concurrent-insert scenario the review flagged, exercised end to end. Separately, this live run caught a real bug the mocked tests missed: a batch created with no tags passed |
|
Good end-to-end validation. The three points all check out cleanly:
The fix is solid. Live verification covers the concurrent-attribution scenario, and the Prisma-rejection bug caught in that run is now fixed with a test. LGTM. |
|
Live e2e verification and regression check on the current head (e5a3d46), real Anthropic API, no mocks, enterprise license enabled with a real team and tagged virtual keys. Feature works:
No regression, checked live and at the unit level:
No backward-incompatible changes:
The terminal batch-cost spend row could not be exercised in the local single-instance setup because the pre-existing CheckBatchCost AsyncIOScheduler does not fire without the surrounding infra; that scheduler is not changed by this PR. The poller's attribution build from the columns is covered by unit tests. Note on CI: the osv-scan check is red from a newly published dependency advisory, not from this branch. No dependency or lockfile changed since the last green osv-scan on 6b77e62; the only files changed since are Python handlers and tests. |
bb2f2f1 to
cd15cb4
Compare
| api_key = getattr(job, "api_key", None) | ||
| team_id = getattr(job, "team_id", None) | ||
| request_tags = getattr(job, "request_tags", None) | ||
|
|
||
| metadata: Dict[str, Any] = { | ||
| "user_api_key_user_id": job.created_by, | ||
| "user_api_key": api_key, | ||
| "user_api_key_team_id": team_id, | ||
| **(await self._get_user_info(batch_id, job.created_by)), | ||
| } | ||
|
|
||
| if api_key is not None: | ||
| metadata["user_api_key_alias"] = await self._get_key_alias(batch_id, api_key) | ||
| team_alias = await self._get_team_alias(team_id) | ||
| if team_alias is not None: | ||
| metadata["user_api_key_team_alias"] = team_alias | ||
| if isinstance(request_tags, list) and request_tags: | ||
| metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)] | ||
|
|
||
| return metadata |
There was a problem hiding this comment.
🟡 New attribution code builds and mutates a dictionary instead of the required immutable style
CLAUDE.md requires new code to avoid mutation and to annotate every variable with : Final, but the new attribution builder seeds a dictionary and then mutates it three times (metadata[...] = ... at enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py:110-123) with plain non-final locals.
Impact: The new code does not follow the repository's mandated coding conventions and adds to the lint budgets it is supposed to ratchet down.
Rule source
CLAUDE.md: "No mutation; don't reassign variables... Instead of mutable lists and dicts, prefer tuples, frozen dataclasses..." and "Annotate every variable with : Final (LIT010)", plus LIT001/LIT002 guidance to build values in one shot with comprehensions or a single expression rather than seeding an empty container and mutating it. api_key, team_id, request_tags, metadata, and team_alias are all unannotated, and metadata is mutated after construction. The same pattern appears in enterprise/litellm_enterprise/proxy/hooks/managed_files.py:198-206.
Prompt for agents
_build_creator_attribution_metadata in enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py builds a dict and then mutates it for alias, team alias and tags, with unannotated locals. Repo conventions (CLAUDE.md) require Final-annotated locals and constructing values in one shot instead of seeding and mutating. Restructure it so the alias/team-alias/tag entries are resolved first into Final locals and the metadata mapping is built in a single expression (conditional dict-unpacking works here), and apply the same treatment to the attribution_columns construction in enterprise/litellm_enterprise/proxy/hooks/managed_files.py.
Was this helpful? React with 👍 or 👎 to provide feedback.
|
|
||
| decoded = _is_base64_encoded_unified_file_id(output_file_id) | ||
| assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] | ||
| class TestBatchCostAttribution: |
There was a problem hiding this comment.
🟡 New test class is glued to the previous test without blank lines, failing the mandatory formatting check
The new attribution test class is placed immediately after the last line of the preceding test (class TestBatchCostAttribution: at tests/proxy_unit_tests/test_check_batch_cost.py:1644) with no blank lines, which the repository's mandatory Black formatting check rejects.
Impact: The pull request fails the required formatting/lint stage in CI until the spacing is fixed.
Black / PEP8 two-blank-line rule before a top-level class
CONTRIBUTING.md states "Black formatting is enforced in CI. All PRs must pass the Black formatting check." Line 1643 is the last assert of test_falls_back_to_deployment_model_group_without_managed_input_file, and line 1644 starts a new top-level class with zero separating blank lines. Black requires two blank lines between a module-level definition and preceding code.
| class TestBatchCostAttribution: | |
| class TestBatchCostAttribution: |
Was this helpful? React with 👍 or 👎 to provide feedback.
|
@greptileai review latest head |
|
bugbot run |
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 81ee8bc. Configure here.
|
Live proxy test on real Vertex batches, polled by a different key: batch cost now bills the creating key, team and tags, no regressions Top row ($0.001723, blank key and team) is a row hand blanked to the pre-PR shape, so it is what used to be lost. Same rows straight out of PostgresRecording: https://app.devin.ai/attachments/605c7f05-9f80-48a9-81c0-b613cd2bbb18/pr34456-ui-edited.mp4 |
81ee8bc to
f38ef4f
Compare
|
@greptileai please review the current head f38ef4f |
b09cca9 to
16a82d0
Compare
|
@greptileai please re-review the current head 16a82d0. The branch was rebased onto |
|
Re-tested live against the new head Batch cost bills the creating key, the team it had at creation, and its tags.
Rollups: the creating team is billed, the poller and the moved-to team are not
New in this head: a poll refreshes state without claiming the row
One thing to be aware of rather than a bug: the poll rewrites Regression checkChat completions, background Responses, Not covered this round: the OpenAI managed batch was still The same rows straight out of Postgres |
…ribution A poll of a Vertex passthrough batch wrote nothing to the managed-object row, so status and file_object stayed frozen at the create-time snapshot and GET /v1/batches served a stale status and an empty output file id for the life of the batch. Only the create may claim a batch, but every observation of one may refresh its state. store_unified_object_id takes create_if_missing, which the poll clears: it refreshes status and file_object through update_many, and leaves a row that is absent absent rather than creating one owned by the observer, since created_by and team_id are written by whoever reaches the create branch. The update payload is now shared with the upsert so it cannot drift into writing api_key, request_tags, created_by or team_id. The passthrough identity re-assertion that was previously part of this PR ships separately in #36121, so this PR keeps only the batch attribution work. The creating key owns user_api_key_alias only when it actually has one. Guarding the overwrite on the presence of a key rather than on a resolved alias nulled the field out for every key generated without key_alias, and for any key rotated or deleted before its batch finished, losing the creating user's alias that the spend row previously carried. The guard now matches the team-alias line below it.
16a82d0 to
255677d
Compare
|
@greptileai please re-review the current head 255677d. Since the last review this adds one fix: the creating key owns |
|
Bugbot run |
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 255677d. Configure here.
mateo-berri
left a comment
There was a problem hiding this comment.
q: why isn't the team id enough to figure out the team associated with the spend row?
|
Closing the one gap from the previous comment: the OpenAI managed batch finished on OpenAI's side 11h22m after creation and the poller billed it 62s later, so the OpenAI managed batch path is now covered end to end at head Its managed-object row is unchanged from pre-PR behavior: |
TLDR
Problem this solves:
How it solves it:
Relationship to other PRs:
LiteLLM_ManagedObjectTablenever stored the key hash; this PR stores itRelevant issues
This supersedes #33319 by @marcus-linktree, which reported the same failure and fixed it by storing the attribution as a
litellm_batch_attributionsnapshot on the batch file object. This PR started from that approach and moved to two dedicated columns onLiteLLM_ManagedObjectTable, so the values are queryable, cannot be echoed back by a client, and are write-once at the DB level rather than inside a JSON blob. The problem statement and the first working fix are his.Linear ticket
Pre-Submission checklist
Type
🐛 Bug Fix
Changes
At batch-create time the Vertex passthrough stored the managed object with
api_key=""and no user, so whenCheckBatchCostlater logged the batch cost it had a blank key, user, and team. The DB spend logger drops a blank-identity row rather than writing an unattributed one, so the cost was silently lost and the creating key showed no spend.The identity is already resolved in memory by auth and copied into the passthrough request metadata. This persists the hashed key and the request tags on the managed-object row in two new columns (
api_key,request_tags, added by a migration), andCheckBatchCost._build_creator_attribution_metadatarebuilds the spend metadata from the row at cost time, resolving the key alias from the hashed token and the team alias from the team id. Rows created before the columns existed fall back tocreated_byandteam_id.Both columns are written only in the upsert create branch and only when the caller is the batch create, so a later status update, retrieve, or poll can neither reassign nor clear them. Every other caller of
store_unified_object_id(the managed-batch and fine-tune post-call hooks, the background Responses API, the Anthropic passthrough) leaves the flag at its default and writes no new columns, so those flows are byte-for-byte unchanged.The batch is claimed from its create route only. The Vertex handler is reached for both the create and a poll, because dispatch is a substring match on the route; claiming the batch from a poll would race the create's own detached write and could leave the batch owned by the polling caller. A poll therefore passes
create_if_missing=Falseandpersist_attribution=False, which refreshesstatusandfile_objectthrough anupdate_manyand leaves a row that is absent absent, rather than creating one owned by the observer. Claiming and observing are separate concerns: only the create may claim a batch, but every observation of one may refresh its state, which is what keepsGET /v1/batchesfrom serving the create-time snapshot for the life of the batch. The registration's failure is reported explicitly instead of surfacing as an unretrieved task exception.The metadata a batch is billed against is the passthrough request metadata, which a client body could previously overwrite with another user, team, org, or end user. That is fixed in #36121 and is already on
litellm_internal_staging, so this PR inherits it rather than carrying it.Behavior changes
aretrieve_batchspend row attributed to the creating key and its tags. fix(batch): track cost for managed batches with no attributable key/u… #35468 made the row itself appear for a batch with no attributable owner; this gives that row an ownerstatusandfile_objectwhen the row exists, so this is a change to which key owns a batch, not to whether its state is kept currentuser_api_key_aliason a batch-cost row is the creating key's alias when the key has one. Previously it carried the creating user'suser_alias, which is a different entity. A key generated withoutkey_alias, and a key rotated or deleted before its batch finished, both keep the creating user's alias rather than emitting a null, so the field never becomes less resolvable than it was_PROXY_track_cost_callbackreads the key as it stands at the moment of logging, which for a batch is whenever the poller sees it finish; a key moved between users or teams while its batch ran would otherwise have billed the key's current owner rather than the batch's creator. That backfill is now skipped foraretrieve_batchand unchanged for every other call typeLiteLLM_VerificationToken.max_budget) now see batch spend. Attributing the batch cost to the creating key is the point of this PR, and the spend increment that follows is what makes a key with a budget able to trip on batch spend that previously escaped it. An operator with key-level caps can see a key exceed immediately after upgrading, on cost that was silently untracked beforeLiteLLM_TeamMembership) now see batch spend, because a batch-cost row carries a team id. This lands with the team id itself, which is already onlitellm_internal_staging, not with this PR; it is called out here because an operator with per-member caps configured can see a member trip immediately after upgrading, on spend that previously escaped the capQA runbook
vertex_aimodel withuse_in_pass_through: true, a Postgres-backed proxy, and an enterprise licenseapi_key, its team id, and its tags inrequest_tagsapi_keyandteam_idstill name key Aproxy_batch_polling_intervalGET /spend/logs: acall_type=aretrieve_batchrow appears, attributed to key A's alias, team alias, and tags, with non-zero spend, and no spend attributed to key BTesting
Every new test was confirmed to fail on the parent commit and pass with the fix. Coverage: the create persists the key hash and tags; a non-create caller writes neither; the columns are write-once across a later store; unset columns are omitted rather than passed as
None(prisma rejectsNonefor the Json column); the batch is registered from the create route only, across collection, trailing-slash, query-string, and id-scoped routes; request-tag precedence over key tags with non-string tags dropped; the attribution metadata carries key, team, both aliases and tags, tolerates legacy rows, keeps the key when a team key has no user, and survives an alias lookup failure. A poll refreshesstatusandfile_objectwhile writing none ofapi_key,request_tags,created_byorteam_idand creating no row when none exists; the create still upserts and claims attribution; and the callers that do not passcreate_if_missing(the fine-tune hook, the Responses API background path, the Anthropic passthrough) still create their rows unchanged.Screenshots / Proof of Fix
A poll of a Vertex batch has to keep the managed-object row current without becoming its owner, so this is an A/B where the only variable is that behavior: the control is this same tree with the
if is_batch_create:gate restored, so the poll writes nothing. Real handler, real managed-files hook, real Postgres, andCheckBatchCostis deliberately never started, because a pure-passthrough model is not a routable deployment and the poller could not heal the row in that config anyway.Key A creates the batch, Vertex finishes it, and key B polls it. Without the poll write the row is frozen at the create-time snapshot and
GET /v1/batchesreportsvalidatingfor a batch that finished; with it the row readscompleted. Attribution is byte-identical on both sides, so the poller refreshes the state without becoming the payer.output_file_idis the same on both sides and is not part of the defect: Vertex echoes the caller'soutputConfigon the create, so the destination is known before the job runs. The value that actually goes stale isstatus.The rest of the evidence below is the attribution half, re-run with this PR merged into
litellm_internal_stagingat8db2fbaad0. Real proxy, real Postgres, real Gemini calls over the passthrough; nothing in the transport or storage layer is stubbed. 210 of 211 cases pass.The four
aretrieve_batchrows at the top of the request log are the batch-cost rows written by theCheckBatchCostpoller. Each carries the creating key hash and its tags, which is the attribution this PR adds; before it these rows had a blank key and were dropped by the DB spend logger. The rows beneath them are real Gemini passthrough calls billed to the owning team, which is the regression belt for the surfaces this PR does not intend to change.The
hash-MOVEDrow is the create-time identity case. Its key was reassigned to another team and another user while the batch was running, and the row is still attributed to the key that created the batch, with neither the new team nor the new user appearing anywhere on it. Opening the row showsCall Type: aretrieve_batchand theLiteLLM Proxy/CheckBatchCosttag, confirming the row came from the deferred poller rather than the request path.Scope of the run: 91 DB cases against real Postgres, 75 HTTP cases against the running proxy, 31 end-to-end spend-row cases through the real logging pipeline, and 14 budget-enforcement cases, split across happy, bad, edge, pre-existing, and new paths. The batch response object is built locally rather than fetched from Vertex, because a real Vertex batch needs GCS and hours of wall clock; everything from the managed-object row through attribution, the callback, the DB writer and the resulting spend row is real code against a real database.
The single non-pass is pre-existing and untouched here: a passthrough body whose
litellm_metadatais a JSON string rather than an object returns 500. The_metadata.update(litellm_metadata)line responsible dates to June 2025 and is byte-identical on the merge base.Final Attestation
Note
Medium Risk
Changes batch spend attribution, managed-object persistence, and poll-vs-create behavior on a billing path; mis-attribution or stale rows would affect key/team budgets and spend logs, though the PR adds extensive tests and write-once guards.
Overview
Vertex passthrough batches previously stored managed objects with an empty key and no tags, so CheckBatchCost logged batch spend with blank identity and the DB dropped those rows. This PR persists the creating key hash and request tags on
LiteLLM_ManagedObjectTable(migration addsapi_key,request_tags) and rebuilds spend metadata at completion via_build_creator_attribution_metadata(key/team aliases, tags, safe handling of nullcreated_byand legacy rows).Managed object writes are split by intent: only the batch create path sets
persist_attribution=True(write-once on upsert create); polls/retrieves usecreate_if_missing=Falseandupdate_manyto refreshstatus/file_objectwithout claiming ownership or creating rows. The Vertex handler distinguishes create vs id-scoped routes, passes real key hash and tag precedence (request tags over key metadata), and surfaces async registration failures via a done callback.Spend callback skips re-reading key identity from the DB for
aretrieve_batchso deferred batch-cost logs keep create-time attribution instead of the key’s current owner.Reviewed by Cursor Bugbot for commit 255677d. Bugbot is set up for automated code reviews on this repo. Configure here.