Skip to content

fix(batches): attribute Vertex passthrough batch cost to key/team/tags - #34456

Merged
yucheng-berri merged 1 commit into
litellm_internal_stagingfrom
litellm_vertex_batch_cost_attribution_oss
Aug 8, 2026
Merged

fix(batches): attribute Vertex passthrough batch cost to key/team/tags#34456
yucheng-berri merged 1 commit into
litellm_internal_stagingfrom
litellm_vertex_batch_cost_attribution_oss

Conversation

@yucheng-berri

@yucheng-berri yucheng-berri commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

TLDR

Problem this solves:

  • A Vertex batch created through the passthrough stored its managed object with no request identity, so the batch cost was logged with a blank key, user, and team
  • The batch cost row carried no key and no tags, so the creating key showed no spend and the cost could not be attributed to a key, a team, or a tag

How it solves it:

  • Persist the creating key hash and its tags on the managed-object row in two new columns, written only in the upsert create branch
  • Rebuild the spend metadata from those columns when the batch completes, resolving the key and team aliases, so the batch-cost row is attributed like a non-batch request

Relationship to other PRs:

Relevant issues

This supersedes #33319 by @marcus-linktree, which reported the same failure and fixed it by storing the attribution as a litellm_batch_attribution snapshot on the batch file object. This PR started from that approach and moved to two dedicated columns on LiteLLM_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

  • I have added meaningful tests
  • My PR passes all CI/CD checks (e.g., lint, format, unit tests)
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review

Type

🐛 Bug Fix

Changes

At batch-create time the Vertex passthrough stored the managed object with api_key="" and no user, so when CheckBatchCost later 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), and CheckBatchCost._build_creator_attribution_metadata rebuilds 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 to created_by and team_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=False and persist_attribution=False, which refreshes status and file_object through an update_many and 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 keeps GET /v1/batches from 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

  • A completed Vertex passthrough batch now writes an aretrieve_batch spend 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 owner
  • Polling or retrieving a batch cannot change or clear which key, team, and tags the batch cost is attributed to
  • A Vertex batch is no longer claimed from a poll. Previously a poll would register a batch whose create-time registration was lost, attributed to the polling caller; such a batch is now not cost-tracked at all, since billing the wrong key is worse than not billing. A poll still refreshes the row's status and file_object when the row exists, so this is a change to which key owns a batch, not to whether its state is kept current
  • user_api_key_alias on a batch-cost row is the creating key's alias when the key has one. Previously it carried the creating user's user_alias, which is a different entity. A key generated without key_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
  • A batch-cost row keeps the identity captured when the batch was created. The key backfill in _PROXY_track_cost_callback reads 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 for aretrieve_batch and unchanged for every other call type
  • Key budgets (LiteLLM_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 before
  • Per-member team budgets (LiteLLM_TeamMembership) now see batch spend, because a batch-cost row carries a team id. This lands with the team id itself, which is already on litellm_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 cap

QA runbook

  1. Configure a vertex_ai model with use_in_pass_through: true, a Postgres-backed proxy, and an enterprise license
  2. Create a team-scoped virtual key with tags (call it key A) and a second key in a different team (key B)
  3. Submit a batch through the passthrough with key A, then confirm the managed-object row carries key A's hashed key in api_key, its team id, and its tags in request_tags
  4. Poll the batch through the passthrough with key B, then confirm the row's api_key and team_id still name key A
  5. Wait for completion and one proxy_batch_polling_interval
  6. GET /spend/logs: a call_type=aretrieve_batch row appears, attributed to key A's alias, team alias, and tags, with non-zero spend, and no spend attributed to key B

Testing

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 rejects None for 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 refreshes status and file_object while writing none of api_key, request_tags, created_by or team_id and creating no row when none exists; the create still upserts and claims attribution; and the callers that do not pass create_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, and CheckBatchCost is 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/batches reports validating for a batch that finished; with it the row reads completed. Attribution is byte-identical on both sides, so the poller refreshes the state without becoming the payer.

Poll state sync A/B

output_file_id is the same on both sides and is not part of the defect: Vertex echoes the caller's outputConfig on the create, so the destination is known before the job runs. The value that actually goes stale is status.

The rest of the evidence below is the attribution half, re-run with this PR merged into litellm_internal_staging at 8db2fbaad0. Real proxy, real Postgres, real Gemini calls over the passthrough; nothing in the transport or storage layer is stubbed. 210 of 211 cases pass.

Live proxy matrix summary

The four aretrieve_batch rows at the top of the request log are the batch-cost rows written by the CheckBatchCost poller. 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.

Request logs showing batch cost attributed to the creating key

The hash-MOVED row 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 shows Call Type: aretrieve_batch and the LiteLLM Proxy/CheckBatchCost tag, confirming the row came from the deferred poller rather than the request path.

Batch cost row detail

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_metadata is 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

  • The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

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 adds api_key, request_tags) and rebuilds spend metadata at completion via _build_creator_attribution_metadata (key/team aliases, tags, safe handling of null created_by and 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 use create_if_missing=False and update_many to refresh status/file_object without 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_batch so 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.

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@greptile-apps

greptile-apps Bot commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

Adds creator attribution for Vertex passthrough batch costs.

  • Persists the creating key hash and request tags on managed-object rows.
  • Reconstructs creator, key, team, alias, and tag metadata when completed batch spend is logged.
  • Makes Vertex polling update existing batch state without creating or claiming missing rows.
  • Preserves captured batch identity instead of backfilling delayed spend from the key’s current owner.
  • Adds the corresponding Prisma schema migration and regression coverage.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failures remain within the scope of the previous review threads.

Important Files Changed

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-apps[bot]

This comment was marked as resolved.

@greptile-apps

greptile-apps Bot commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This 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/5

The 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

Security Review

A concurrent initial store can still replace the creator attribution because preservation is implemented as a non-atomic read followed by an upsert

Important Files Changed

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)

  1. enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py, line 663-672 (link)

    P2 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 identity

    Knowledge Base Used: Cost Tracking and Budget Enforcement

Reviews (2): Last reviewed commit: "fix(batches): make batch-cost attributio..." | Re-trigger Greptile

greptile-apps[bot]

This comment was marked as resolved.

devin-ai-integration[bot]

This comment was marked as resolved.

@codecov

codecov Bot commented Jul 24, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 78.12500% with 7 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...der_handlers/vertex_passthrough_logging_handler.py 75.00% 7 Missing ⚠️

📢 Thoughts on this report? Let us know!

@codspeed-hq

codspeed-hq Bot commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing litellm_vertex_batch_cost_attribution_oss (255677d) with litellm_internal_staging (97a59c8)

Open in CodSpeed

@CLAassistant

CLAassistant commented Jul 24, 2026

Copy link
Copy Markdown

CLA assistant check
All committers have signed the CLA.

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

veria-ai[bot]

This comment was marked as resolved.

@veria-ai

veria-ai Bot commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

PR overview

All previously flagged issues have been addressed. No open security concerns remain on this pull request.

Security review

No open security issues remain on this pull request.

Fixed/addressed: 3 · PR risk: 0/10

devin-ai-integration[bot]

This comment was marked as resolved.

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

cursor[bot]

This comment was marked as resolved.

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

greptile-apps[bot]

This comment was marked as resolved.

devin-ai-integration[bot]

This comment was marked as resolved.

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

The concurrent-initial-store race is closed in 54acf65. The passthrough handlers now register the managed object only on the create route (POST .../batchPredictionJobs and .../messages/batches), never on the .../{id} poll/retrieve path, for both Vertex and Anthropic. Since only the creating request ever calls store, no poll can win the insert and stamp its own identity into the write-once columns. Added a regression test asserting a retrieve-shaped route does not register (fails when the gate is reverted).

@greptileai

@greptile-apps

greptile-apps Bot commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

The fix looks correct. Both handlers now gate registration on the collection endpoint:

  • Vertex: url_route.split("?")[0].rstrip("/").endswith("batchPredictionJobs") — only a POST to .../batchPredictionJobs passes; a poll/retrieve at .../batchPredictionJobs/{id} does not.
  • Anthropic: url_route.split("?")[0].rstrip("/").endswith("batches") — same logic for .../messages/batches vs .../messages/batches/{id}.

This fully closes the race: since only the creating call ever reaches _store_batch_managed_object, there's no competing inserter to win the write-once columns. The regression test asserting that a retrieve-shaped route does not register is the right coverage for this.

The two original concerns (initial attribution race and non-atomic initialization) are addressed. No remaining issues from my earlier comments.

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

greptile-apps[bot]

This comment was marked as resolved.

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

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 store_unified_object_id call sites. For a passthrough batch, only the create route reaches store_batch_managed_object (the .../{id} poll and retrieve paths are gated out), so exactly one call registers that batch's unified_object_id, once, with the creating key. The managed-files post-call hook (managed_files.py:1192/:1241) is the only other batch caller, and it serves the separate unified /v1/batches flow, where it awaits the store inside the create request itself, so the row exists with the creator's identity before any retrieve can run; a later retrieve by a different key takes the update branch, which never touches user_api_key/request_tags. It does not fire with a passthrough batch's id. The response endpoint caller is file_purpose=response, not a batch.

So while store_unified_object_id in isolation looks racy, no reachable pair of callers writes the same batch row with different identities. The create store being a background task is the pre-existing logging-path pattern and is benign here because it is the sole writer of that id.

If you'd still prefer a belt-and-suspenders change, the passthrough create store could be awaited rather than dispatched via asyncio.create_task, but that alters the non-blocking logging path for a case that is not reachable today.

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

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 (batch-creator-A) and key B (poller-B).

  1. Key A creates a batch through the passthrough (POST /anthropic/v1/messages/batches, real msgbatch_...). The managed-object row is written with user_api_key equal to key A's hashed token.
  2. Key B polls the same batch three times (GET /anthropic/v1/messages/batches/{id}, all 200). The row's user_api_key is unchanged, still key A's hash, never key B's.
  3. Proxy logs confirm the create-only gate: one Stored batch managed object call for the POST, zero for the three GET polls.

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 request_tags=None, which prisma rejects for the Json column, so the row was never registered. Fixed in e5a3d46 by omitting unset optional columns, with a regression test.

@greptileai

@greptile-apps

greptile-apps Bot commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

Good end-to-end validation. The three points all check out cleanly:

  1. Poll gate — the create-only registration confirmed in logs is exactly what the "only upsert on POST" guard was meant to enforce. Seeing zero Stored batch managed object calls for the GETs is the right signal.

  2. Attribution stability under concurrent reads — Key A's hash holding across three Key B polls is the concurrent-insert scenario that needed live confirmation. The upsert update branch not touching user_api_key/team_id/created_by is what makes this safe, and the log confirms it held.

  3. The request_tags=None Prisma rejection — this is a real common-path bug (most batches have no tags) and the fix in e5a3d46 is correct: build the create dict conditionally so nullable Json columns are omitted rather than set to None. The regression test for the no-tags path is the right coverage here.

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.

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

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:

  • Create through the Anthropic passthrough with key A (team-alpha, tags [prod, batch]) writes the managed object row with user_api_key = key A's hashed token, team_id = team-alpha, request_tags = [prod, batch]
  • Key B polling the same batch leaves the row attributed to key A, never key B
  • A request body claiming user_api_key_team_id = SPOOFED-TEAM and user_api_key_user_id = attacker is ignored; the row uses the authenticated key A and team-alpha

No regression, checked live and at the unit level:

  • 373 unit tests pass across managed_files, the unified /v1/batches endpoint, the files endpoint, check_batch_cost, and the Vertex and Anthropic passthrough suites
  • A normal chat completion still works and logs real spend for the key
  • Retrieving a batch returns a response with no user_api_key and no litellm_batch_attribution; the internal identity never reaches the caller
  • The create-only registration leaves exactly one managed object row per batch after a poll (no duplicate or churned rows)

No backward-incompatible changes:

  • The two new columns are nullable and additive; older rows read as NULL and fall back to created_by/team_id (covered by test_attribution_metadata_tolerates_legacy_rows)
  • store_unified_object_id gained one optional parameter with a default, so every existing caller keeps working unchanged
  • Pre-PR released behavior had no batch attribution at all, so populating the columns is purely additive

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.

@greptileai

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@yucheng-berri
yucheng-berri requested a review from a team July 25, 2026 02:21
@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

greptile-apps[bot]

This comment was marked as resolved.

@yucheng-berri
yucheng-berri force-pushed the litellm_vertex_batch_cost_attribution_oss branch from bb2f2f1 to cd15cb4 Compare August 6, 2026 20:04

@devin-ai-integration devin-ai-integration Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Devin Review found 2 new potential issues.

View 2 additional findings in Devin Review.

Open in Devin Review

Comment thread litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
Comment on lines +106 to +125
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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.
Open in Devin Review

Was this helpful? React with 👍 or 👎 to provide feedback.

@devin-ai-integration devin-ai-integration Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Devin Review found 1 new potential issue.

View 2 additional findings in Devin Review.

Open in Devin Review


decoded = _is_base64_encoded_unified_file_id(output_file_id)
assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP]
class TestBatchCostAttribution:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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.

Suggested change
class TestBatchCostAttribution:
class TestBatchCostAttribution:
Open in Devin Review

Was this helpful? React with 👍 or 👎 to provide feedback.

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai review latest head

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

greptile-apps[bot]

This comment was marked as resolved.

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ 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.

@devin-ai-integration

Copy link
Copy Markdown
Contributor

Live proxy test on real Vertex batches, polled by a different key: batch cost now bills the creating key, team and tags, no regressions

logs

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. key-b-poller polled every batch and was billed nothing

Row detail, request-tag override, key moved to another team mid flight

drawer
tag override
moved key

Rollups: creating team billed, poller team and moved-to team not

teams
keys
poller

Regression check: non-batch traffic unchanged

regressions

Same rows straight out of Postgres
$ bash 03_show_managed_rows.sh   # managed-object row per real Vertex batch
=== [A] vertex batch 958926705248960512
 api_key_resolves_to |   created_by_user    |     team     |           request_tags            |  status  | batch_processed 
---------------------+----------------------+--------------+-----------------------------------+----------+-----------------
 key-a-batch-creator | creator-a-user-alias | team-batch-a | ["batch-tag-1", "cost-center-42"] | complete | t
(1 row)

=== [L] vertex batch 6151577075607142400
 api_key_resolves_to |   created_by_user    |     team     | request_tags |  status  | batch_processed 
---------------------+----------------------+--------------+--------------+----------+-----------------
 <none>              | creator-a-user-alias | team-batch-a |              | complete | t
(1 row)

=== [M] vertex batch 9216276622032764928
 api_key_resolves_to |   created_by_user    |     team     |   request_tags    |  status  | batch_processed 
---------------------+----------------------+--------------+-------------------+----------+-----------------
 key-m-will-move     | creator-a-user-alias | team-batch-a | ["moved-key-tag"] | complete | t
(1 row)

=== [OLD] vertex batch 8610542472151433216
 api_key_resolves_to | created_by_user |  team  | request_tags |  status  | batch_processed 
---------------------+-----------------+--------+--------------+----------+-----------------
 <none>              | <null>          | <null> |              | complete | t
(1 row)

=== [P] vertex batch 5595382521626886144
 api_key_resolves_to | created_by_user | team | request_tags | status | batch_processed 
---------------------+-----------------+------+--------------+--------+-----------------
(0 rows)

=== [R] vertex batch 3510215894154346496
  api_key_resolves_to  | created_by_user |  team  | request_tags |  status  | batch_processed 
-----------------------+-----------------+--------+--------------+----------+-----------------
 key-r-no-team-no-tags | <null>          | <null> | []           | complete | t
(1 row)

=== [TAGREQ] vertex batch 7835923336243707904
 api_key_resolves_to |   created_by_user    |     team     |               request_tags                |  status  | batch_processed 
---------------------+----------------------+--------------+-------------------------------------------+----------+-----------------
 key-a-batch-creator | creator-a-user-alias | team-batch-a | ["request-tag-alpha", "request-tag-beta"] | complete | t
(1 row)


$ bash 07_show_batch_cost_rows.sh   # batch-cost rows written by CheckBatchCost
      model       | billed_key_hash_resolves_to |   logged_key_alias    | logged_team_alias |     billed_user      |                                                        tags                                                        |         spend         
------------------+-----------------------------+-----------------------+-------------------+----------------------+--------------------------------------------------------------------------------------------------------------------+-----------------------
 gemini-2.5-flash | <NO KEY>                    | <null>                | <null>            | <null>               | ["User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                                          |             0.0017226
 gemini-2.5-flash | key-a-batch-creator         | key-a-batch-creator   | team-batch-a      | creator-a-user-alias | ["batch-tag-1", "cost-center-42", "User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]         | 0.0005426000000000001
 gemini-2.5-flash | <NO KEY>                    | creator-a-user-alias  | team-batch-a      | creator-a-user-alias | ["User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                                          | 0.0004913500000000001
 gemini-2.5-flash | key-m-will-move             | key-m-will-move       | team-batch-a      | creator-a-user-alias | ["moved-key-tag", "User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                         | 0.0002876000000000001
 gemini-2.5-flash | key-r-no-team-no-tags       | key-r-no-team-no-tags | <null>            | <null>               | ["User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                                          | 7.135000000000001e-05
 gemini-2.5-flash | key-a-batch-creator         | key-a-batch-creator   | team-batch-a      | creator-a-user-alias | ["request-tag-alpha", "request-tag-beta", "User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"] |              5.76e-05
 gpt-4o-mini      | <NO KEY>                    | creator-a-user-alias  | team-batch-a      | creator-a-user-alias | ["User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                                          |             1.275e-06
(7 rows)

== spend attributed to the poller key (key-b-poller) by the batch poller: must be none ==
 batch_cost_rows_billed_to_key_b 
---------------------------------
                               0
(1 row)


$ bash 06_show_spend_logs.sh   # every spend row, incl. regression traffic
       call_type       |                   model                   |      billed_key       |   logged_key_alias    |     team     |         user         |                                                        tags                                                        |   spend    | status  
-----------------------+-------------------------------------------+-----------------------+-----------------------+--------------+----------------------+--------------------------------------------------------------------------------------------------------------------+------------+---------
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-m-will-move       | key-m-will-move       | team-batch-a | creator-a-user-alias | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-l-legacy-row      | key-l-legacy-row      | team-batch-a | creator-a-user-alias | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-r-no-team-no-tags | key-r-no-team-no-tags | <null>       | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-p-poll-register   | key-p-poll-register   | team-batch-a | creator-a-user-alias | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["request-tag-alpha", "request-tag-beta"]                                                                          | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-r-no-team-no-tags | key-r-no-team-no-tags | <null>       | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | unknown                                   | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | gemini-2.5-flash                          | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | []                                                                                                                 | 0.00005240 | success
 acompletion           | openai/gpt-4o-mini                        | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["batch-tag-1", "cost-center-42", "User-Agent: curl", "User-Agent: curl/7.81.0"]                                   | 0.00000255 | success
 aresponses            | openai/gpt-4o-mini                        | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["batch-tag-1", "cost-center-42", "User-Agent: curl", "User-Agent: curl/7.81.0"]                                   | 0.00000000 | success
 acreate_file          | openai/gpt-4o-mini                        | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["batch-tag-1", "cost-center-42", "User-Agent: curl", "User-Agent: curl/7.81.0"]                                   | 0.00000000 | success
 acreate_batch         | openai/gpt-4o-mini                        | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["batch-tag-1", "cost-center-42", "User-Agent: curl", "User-Agent: curl/7.81.0"]                                   | 0.00000000 | success
                       | openai/gpt-4o-mini                        | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["batch-tag-1", "cost-center-42"]                                                                                  | 0.00000000 | failure
 aget_responses        | openai/gpt-4o-mini                        | <NO KEY>              | <null>                | <null>       | creator-a-user-alias | []                                                                                                                 | 0.00000315 | success
 aretrieve_batch       | gpt-4o-mini                               | <NO KEY>              | creator-a-user-alias  | team-batch-a | creator-a-user-alias | ["User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                                          | 0.00000128 | success
 aretrieve_batch       | gemini-2.5-flash                          | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["batch-tag-1", "cost-center-42", "User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]         | 0.00054260 | success
 aretrieve_batch       | gemini-2.5-flash                          | key-m-will-move       | key-m-will-move       | team-batch-a | creator-a-user-alias | ["moved-key-tag", "User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                         | 0.00028760 | success
 aretrieve_batch       | gemini-2.5-flash                          | <NO KEY>              | creator-a-user-alias  | team-batch-a | creator-a-user-alias | ["User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                                          | 0.00049135 | success
 aretrieve_batch       | gemini-2.5-flash                          | key-r-no-team-no-tags | key-r-no-team-no-tags | <null>       | <null>               | ["User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                                          | 0.00007135 | success
 aretrieve_batch       | gemini-2.5-flash                          | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["request-tag-alpha", "request-tag-beta", "User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"] | 0.00005760 | success
 aretrieve_batch       | gemini-2.5-flash                          | <NO KEY>              | <null>                | <null>       | <null>               | ["User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                                          | 0.00172260 | success
(29 rows)

Recording: https://app.devin.ai/attachments/605c7f05-9f80-48a9-81c0-b613cd2bbb18/pr34456-ui-edited.mp4

@yucheng-berri
yucheng-berri force-pushed the litellm_vertex_batch_cost_attribution_oss branch from 81ee8bc to f38ef4f Compare August 8, 2026 17:17
@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai please review the current head f38ef4f

greptile-apps[bot]

This comment was marked as resolved.

@yucheng-berri
yucheng-berri force-pushed the litellm_vertex_batch_cost_attribution_oss branch 2 times, most recently from b09cca9 to 16a82d0 Compare August 8, 2026 18:17
@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai please re-review the current head 16a82d0. The branch was rebased onto litellm_internal_staging and squashed to a single commit, and I replied to the failed-registration finding with the pre-PR behavior it would restore.

@devin-ai-integration

Copy link
Copy Markdown
Contributor

Re-tested live against the new head 16a82d0de7 (the earlier comment was 81ee8bcfb4, before the poll-refresh commit). Everything from scratch: fresh proxy, fresh keys/teams, 7 new real Vertex batchPredictionJobs through the passthrough with real GCS input and output, left to actually finish, billed by the real CheckBatchCost poller, and every batch polled in between by key-b-poller from a different team.

Batch cost bills the creating key, the team it had at creation, and its tags.

logs

$0.000798 carries request-tag-alpha/beta only, so request tags win over key tags. $0.000075 is key-m-will-move, moved to team-batch-c-moved while its batch was running, still billed to team-batch-a. $0.000925 is a row hand-blanked to the pre-migration shape and it falls back to user plus team. $0.000068 is a row blanked to the full pre-PR shape, which is what every Vertex batch used to look like before this PR.

Row detail, request-tag override, moved key

drawer

tag override

moved key

Rollups: the creating team is billed, the poller and the moved-to team are not

teams

keys

poller

team-batch-a $0.0041, team-batch-b (poller) $0.00, team-batch-c-moved $0.00. Filtering the logs by key-b-poller gives 13 rows, all zero-cost polls, no aretrieve_batch.

New in this head: a poll refreshes state without claiming the row
== batch A row, as the create left it ==
  status   |  api_key  |  request_tags  |  created_by  |  team_id  |  updated_by  | file_object_status |  output_file_id  
-----------+---+---+---+---+---+---+---
 completed | 4d834de6710354348de28c13bf9c45c05866f08d2baf8f9dc24ccf88dc390ece | ["batch-tag-1", "cost-center-42"] | 1397d717-41ef-4be0-a363-e0cb5b29588b | e2fa582d-9784-48be-9be2-e72d0fdc14d6 | 1397d717-41ef-4be0-a363-e0cb5b29588b | completed  | <managed base64 file id, decodes to gs://litellm_bucket-16/litellm-vertex-files/pr34456/out-A/.../predictions.jsonl>
(1 row)

== vertex's own view of batch A ==
   state: JOB_STATE_SUCCEEDED
== batch A row after that poll by key-b-poller (other team) ==
  status   |  api_key  |  request_tags  |  created_by  |  team_id  | updated_by | file_object_status |  output_file_id  
-----------+---+---+---+---+------------+---+---
 completed | 4d834de6710354348de28c13bf9c45c05866f08d2baf8f9dc24ccf88dc390ece | ["batch-tag-1", "cost-center-42"] | 1397d717-41ef-4be0-a363-e0cb5b29588b | e2fa582d-9784-48be-9be2-e72d0fdc14d6 | <null>     | completed  | gs://litellm_bucket-16/litellm-vertex-files/pr34456/out-A/prediction-model-2026-08-08T18:45:27.897977Z/predictions.jsonl
(1 row)

status and file_object refresh, while api_key, request_tags, created_by and team_id stay byte identical. Batch 1761886852874240000 had its managed row deleted and was then polled by the observer key: still (0 rows), so a poll does not resurrect a missing row and cannot claim one.

One thing to be aware of rather than a bug: the poll rewrites file_object.output_file_id from the managed base64 id to the raw provider gs:// path, as the diff above shows. GET /v1/batches/<unified id> still returns the managed base64 id before and after the poll, so nothing user visible changes, but the column value does differ depending on who touched the row last.

Regression check

regressions

Chat completions, background Responses, /v1/files, /v1/batches, a 400 failure row and non-batch Vertex generateContent all still attribute to key-a-batch-creator / team-batch-a with the key's tags, and the OpenAI managed-object row still has api_key NULL as before.

Not covered this round: the OpenAI managed batch was still validating upstream when I wrapped up, so this data set has no gpt-4o-mini aretrieve_batch row. That path was exercised in the previous round and is untouched by the new commit.

The same rows straight out of Postgres
# Live proxy test of PR #34456 at head 16a82d0de7, merged into litellm_internal_staging (b0fd3e1e30)

$ bash 03_show_managed_rows.sh
=== [A] vertex batch 2885534959903178752
 api_key_resolves_to |   created_by_user    |     team     |           request_tags            |  status  | batch_processed 
---------------------+----------------------+--------------+-----------------------------------+----------+-----------------
 key-a-batch-creator | creator-a-user-alias | team-batch-a | ["batch-tag-1", "cost-center-42"] | complete | t
(1 row)

=== [L] vertex batch 3200786933819113472
 api_key_resolves_to |   created_by_user    |     team     | request_tags |  status  | batch_processed 
---------------------+----------------------+--------------+--------------+----------+-----------------
 <none>              | creator-a-user-alias | team-batch-a |              | complete | t
(1 row)

=== [M] vertex batch 6610011851738578944
 api_key_resolves_to |   created_by_user    |     team     |   request_tags    |  status  | batch_processed 
---------------------+----------------------+--------------+-------------------+----------+-----------------
 key-m-will-move     | creator-a-user-alias | team-batch-a | ["moved-key-tag"] | complete | t
(1 row)

=== [OLD] vertex batch 8770613772969574400
 api_key_resolves_to | created_by_user |  team  | request_tags |  status  | batch_processed 
---------------------+-----------------+--------+--------------+----------+-----------------
 <none>              | <null>          | <null> |              | complete | t
(1 row)

=== [P] vertex batch 1761886852874240000
 api_key_resolves_to | created_by_user | team | request_tags | status | batch_processed 
---------------------+-----------------+------+--------------+--------+-----------------
(0 rows)

=== [R] vertex batch 6673062246521765888
  api_key_resolves_to  | created_by_user |  team  | request_tags |  status  | batch_processed 
-----------------------+-----------------+--------+--------------+----------+-----------------
 key-r-no-team-no-tags | <null>          | <null> | []           | complete | t
(1 row)

=== [TAGREQ] vertex batch 1939779038155374592
 api_key_resolves_to |   created_by_user    |     team     |               request_tags                |  status  | batch_processed 
---------------------+----------------------+--------------+-------------------------------------------+----------+-----------------
 key-a-batch-creator | creator-a-user-alias | team-batch-a | ["request-tag-alpha", "request-tag-beta"] | complete | t
(1 row)


$ bash 08_poll_refresh.sh   # a poll refreshes state without claiming the row
(captured live earlier this run, see below)

$ bash 07_show_batch_cost_rows.sh
      model       | billed_key_hash_resolves_to |   logged_key_alias    | logged_team_alias |     billed_user      |                                                        tags                                                        |         spend         
------------------+-----------------------------+-----------------------+-------------------+----------------------+--------------------------------------------------------------------------------------------------------------------+-----------------------
 gemini-2.5-flash | <NO KEY>                    | creator-a-user-alias  | team-batch-a      | creator-a-user-alias | ["User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                                          |             0.0009251
 gemini-2.5-flash | key-a-batch-creator         | key-a-batch-creator   | team-batch-a      | creator-a-user-alias | ["request-tag-alpha", "request-tag-beta", "User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"] | 0.0007976000000000001
 gemini-2.5-flash | key-a-batch-creator         | key-a-batch-creator   | team-batch-a      | creator-a-user-alias | ["batch-tag-1", "cost-center-42", "User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]         |             0.0002251
 gemini-2.5-flash | key-r-no-team-no-tags       | key-r-no-team-no-tags | <null>            | <null>               | ["User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                                          |            0.00011885
 gemini-2.5-flash | key-m-will-move             | key-m-will-move       | team-batch-a      | creator-a-user-alias | ["moved-key-tag", "User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                         | 7.510000000000001e-05
 gemini-2.5-flash | <NO KEY>                    | <null>                | <null>            | <null>               | ["User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                                          | 6.760000000000002e-05
(6 rows)

== spend attributed to the poller key (key-b-poller) by the batch poller: must be none ==
 batch_cost_rows_billed_to_key_b 
---------------------------------
                               0
(1 row)


$ bash 06_show_spend_logs.sh
       call_type       |                   model                   |      billed_key       |   logged_key_alias    |     team     |         user         |                                                        tags                                                        |   spend    | status  
-----------------------+-------------------------------------------+-----------------------+-----------------------+--------------+----------------------+--------------------------------------------------------------------------------------------------------------------+------------+---------
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-m-will-move       | key-m-will-move       | team-batch-a | creator-a-user-alias | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-l-legacy-row      | key-l-legacy-row      | team-batch-a | creator-a-user-alias | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-r-no-team-no-tags | key-r-no-team-no-tags | <null>       | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-p-poll-register   | key-p-poll-register   | team-batch-a | creator-a-user-alias | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["request-tag-alpha", "request-tag-beta"]                                                                          | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-r-no-team-no-tags | key-r-no-team-no-tags | <null>       | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | unknown                                   | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | publishers/google/models/gemini-2.5-flash | key-b-poller          | key-b-poller          | team-batch-b | <null>               | []                                                                                                                 | 0.00000000 | success
 pass_through_endpoint | gemini-2.5-flash                          | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | []                                                                                                                 | 0.00004990 | success
 acompletion           | openai/gpt-4o-mini                        | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["batch-tag-1", "cost-center-42", "User-Agent: curl", "User-Agent: curl/7.81.0"]                                   | 0.00000255 | success
 aresponses            | openai/gpt-4o-mini                        | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["batch-tag-1", "cost-center-42", "User-Agent: curl", "User-Agent: curl/7.81.0"]                                   | 0.00000000 | success
 acreate_file          | openai/gpt-4o-mini                        | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["batch-tag-1", "cost-center-42", "User-Agent: curl", "User-Agent: curl/7.81.0"]                                   | 0.00000000 | success
 acreate_batch         | openai/gpt-4o-mini                        | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["batch-tag-1", "cost-center-42", "User-Agent: curl", "User-Agent: curl/7.81.0"]                                   | 0.00000000 | success
                       | openai/gpt-4o-mini                        | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["batch-tag-1", "cost-center-42"]                                                                                  | 0.00000000 | failure
 aget_responses        | openai/gpt-4o-mini                        | <NO KEY>              | <null>                | <null>       | creator-a-user-alias | []                                                                                                                 | 0.00000315 | success
 aretrieve_batch       | gemini-2.5-flash                          | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["batch-tag-1", "cost-center-42", "User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]         | 0.00022510 | success
 aretrieve_batch       | gemini-2.5-flash                          | key-m-will-move       | key-m-will-move       | team-batch-a | creator-a-user-alias | ["moved-key-tag", "User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                         | 0.00007510 | success
 aretrieve_batch       | gemini-2.5-flash                          | <NO KEY>              | creator-a-user-alias  | team-batch-a | creator-a-user-alias | ["User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                                          | 0.00092510 | success
 aretrieve_batch       | gemini-2.5-flash                          | key-r-no-team-no-tags | key-r-no-team-no-tags | <null>       | <null>               | ["User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                                          | 0.00011885 | success
 aretrieve_batch       | gemini-2.5-flash                          | key-a-batch-creator   | key-a-batch-creator   | team-batch-a | creator-a-user-alias | ["request-tag-alpha", "request-tag-beta", "User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"] | 0.00079760 | success
 aretrieve_batch       | gemini-2.5-flash                          | <NO KEY>              | <null>                | <null>       | <null>               | ["User-Agent: LiteLLM Proxy", "User-Agent: LiteLLM Proxy/CheckBatchCost"]                                          | 0.00006760 | success
(31 rows)

…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.
@yucheng-berri
yucheng-berri force-pushed the litellm_vertex_batch_cost_attribution_oss branch from 16a82d0 to 255677d Compare August 8, 2026 20:44
@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai please re-review the current head 255677d. Since the last review this adds one fix: the creating key owns user_api_key_alias only when it actually has one, so a key generated without key_alias (or rotated before its batch finished) keeps the creating user's alias instead of nulling the field out. The guard now matches the team-alias line below it. The PR description also gained the key-budget behavior-change bullet.

@mateo-berri

Copy link
Copy Markdown
Contributor

Bugbot run

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ 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.

@yucheng-berri
yucheng-berri enabled auto-merge (squash) August 8, 2026 22:14

@mateo-berri mateo-berri left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

q: why isn't the team id enough to figure out the team associated with the spend row?

@mateo-berri mateo-berri left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Thanks!

@yucheng-berri
yucheng-berri merged commit efc4e6f into litellm_internal_staging Aug 8, 2026
93 checks passed
@yucheng-berri
yucheng-berri deleted the litellm_vertex_batch_cost_attribution_oss branch August 8, 2026 23:01
@devin-ai-integration

Copy link
Copy Markdown
Contributor

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 16a82d0de7 too.

model        | spend     | call_type       | key_alias | team_alias   | startTime
gpt-4o-mini  | 1.275e-06 | aretrieve_batch | <none>    | team-batch-a | 2026-08-09 06:13:01

Its managed-object row is unchanged from pre-PR behavior: api_key NULL, attribution resolved through created_by and team_id. Only the Vertex passthrough create route sets the new columns, which is the intended scope.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants