fix(batches): attribute Vertex batch cost spend logs to key/team/tags

[marcus-linktree]

Relevant issues

Resolves #33316

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

Screenshots / Proof of Fix

Verified on a local proxy running the equivalent change. A Vertex batch submitted through the passthrough with a team-scoped, tagged virtual key now produces an attributed aretrieve_batch cost row with real spend, key alias, team, and tags, where before no row was written at all and the key showed $0 spend

Type

🐛 Bug Fix

Changes

Vertex batches created through the passthrough (/vertex_ai/.../batchPredictionJobs) logged no attributed cost row. At create time the managed object was stored with no request identity (api_key="", user_api_key_user_id unset), so when the CheckBatchCost poller later logged the cost it had a blank key, user, and team. In that state the DB spend-logger drops the row entirely rather than writing an unattributed one, so the batch cost was silently lost and the key showed $0 spend

The identity needed for attribution is already resolved in memory by auth at create time and copied into the passthrough request metadata, so this snapshots it onto the managed object file_object jsonb right then (the key hash, the owning user and team, their aliases, user email, and the request tags, which fall back to the key's own tags when the request sent none) and reads it back when the batch completes. CheckBatchCost issues no per-batch key, team, or user database lookup; batches created before the snapshot existed fall back to the stored row. This needs no schema migration

To make that snapshot safe to trust, the passthrough re-asserts the authenticated identity fields (user_api_key_user_id, user_api_key_team_id, and the aliases and email) in the request metadata after the client-supplied metadata is merged, mirroring how the key hash was already protected. Previously a request body could set user_api_key_user_id or user_api_key_team_id and have the completed batch charged to another user or team's spend

The snapshot is stored only for internal cost attribution and must never reach a caller. Because Batch inherits pydantic extra="allow", anything under the stored file_object round-trips into the LiteLLMBatch that the retrieve and list endpoints return, which would expose the creator's key hash, email, and aliases to anyone with team-level access to the batch. The two places that build a batch from a stored row now strip the internal key from a copy before returning, and the preserve-on-update path writes the snapshot into the serialized DB file_object rather than onto the returned response object, so it stays server-side

The snapshot also has to survive later writes to the same row. When a retrieve detects a status change, update_batch_in_database overwrites file_object with the raw provider batch, which does not carry litellm_batch_attribution, so without care the later cost poll would fall back to attribution with no key hash and the spend would skip the caller's key-level budget. The update path now reads the stored snapshot off the existing row and carries it onto the response before serializing, so it persists across every status update through completion

In pass_through_endpoints.py the identity re-assertion closes that spoofing gap. In vertex_passthrough_logging_handler.py the create path stashes the authenticated identity plus tags onto file_object. In check_batch_cost.py the poll path builds the spend-log metadata from that snapshot with no DB call, keeping the pre-existing user-table enrichment only for the legacy no-snapshot fallback

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
3. Submit a Vertex batch through the gateway with that key; wait for completion and one proxy_batch_polling_interval
4. GET /spend/logs: a call_type=aretrieve_batch row appears, attributed to the key alias, team alias, and tags, with non-zero spend

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

[CLAassistant]

All committers have signed the CLA.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR fixes Vertex batch jobs submitted through the passthrough endpoint producing no attributed spend row — the batch cost was silently dropped because the managed object was stored with a blank key, user, and team. The fix snapshots the authenticated identity and request tags onto the managed object's file_object jsonb at create time, reads them back at poll time in CheckBatchCost with no extra DB lookups, and preserves the snapshot across every subsequent status-update write.

• Attribution snapshot: _store_batch_managed_object writes litellm_batch_attribution into batch_object.metadata before calling store_unified_object_id; _get_batch_spend_attribution reads it back at poll time, falling back to the legacy user-table path for batches created before the snapshot existed.
• Identity re-assertion: pass_through_endpoints.py now overwrites user_id, team_id, team_alias, key_alias, and user_email in the merged metadata from the authenticated UserAPIKeyAuth after the client-metadata merge, preventing a request body from redirecting batch spend to another user or team.
• Snapshot scrubbing: strip_internal_batch_attribution is applied in both get_batch_from_database and managed_files.list_batches so the creator's key hash, email, and aliases never reach a caller; update_batch_in_database carries the snapshot forward from the stored row so it survives status updates.

Confidence Score: 5/5

Safe to merge — the changes are well-scoped, all edge cases (legacy batches, None fields, malformed JSON, stash failures) are handled gracefully, and the scrubbing of internal data from API responses is consistently applied across all read paths.

The snapshot/restore flow is correct end-to-end: identity re-assertion happens before the snapshot is taken, the snapshot survives status-update overwrites via explicit carry-forward, and the internal key is stripped before any batch object is returned to a caller. Test coverage is thorough, including the spoofing regression, DB-lookup elimination, legacy fallback, and the non-mutation of the raw DB row.

No files require special attention.

Important Files Changed

Filename	Overview
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py	Adds _get_batch_attribution, _get_batch_spend_attribution helpers; replaces ad-hoc user-table lookup with snapshot-based attribution, with clean legacy fallback. Logic is correct and handles all edge cases (None fields filtered, empty attribution dict, JSON/dict duality of file_object).
litellm/proxy/openai_files_endpoints/common_utils.py	Adds strip_internal_batch_attribution and read_stored_batch_attribution helpers, applies stripping in get_batch_from_database, and preserves the attribution snapshot across status updates in update_batch_in_database. Both string and dict forms of file_object are handled.
litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py	Snapshots the authenticated identity + request tags onto batch_object.metadata before storage; stash failures are caught and logged so they never abort batch creation. Tag resolution priority (request → litellm_metadata → key auth metadata) is consistent with the passthrough pattern.
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py	Re-asserts user_id, team_id, team_alias, key_alias, and user_email from the authenticated UserAPIKeyAuth after the client-metadata merge, preventing a request body from spoofing the batch-cost owner. Mirrors the existing api_key re-assertion pattern.
enterprise/litellm_enterprise/proxy/hooks/managed_files.py	Applies strip_internal_batch_attribution to the list-batches path, ensuring the attribution snapshot is scrubbed from the caller-facing response. One-liner change with correct placement.
tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py	Adds thorough regression tests: attribution preserved across status updates, stripping does not mutate the raw row, get_batch_from_database does not leak the snapshot, read_stored_batch_attribution handles all None/malformed inputs. Assertion polarity is correct.
tests/proxy_unit_tests/test_check_batch_cost.py	New TestBatchCostAttribution class tests snapshot recovery, zero-DB-lookup attribution, spoofing protection via snapshot > stored columns, missing-field filtering, and legacy fallback. All assertions are directionally correct.
tests/pass_through_unit_tests/test_pass_through_unit_tests.py	Adds spoofing-regression test that verifies client-supplied identity fields are overwritten by the authenticated key after the metadata merge.
tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py	Adds three tests for _store_batch_managed_object: key+tags stashed correctly, key-level tags used as fallback when no request tags, and stash failure is swallowed. Coverage is solid.

_{Reviews (6): Last reviewed commit: "fix(batches): keep the batch attribution..." | Re-trigger Greptile}

[veria-ai]

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

[marcus-linktree]

@greptileai

[marcus-linktree]

@greptileai

[marcus-linktree]

@greptileai

[marcus-linktree]

@greptileai

[marcus-linktree]

@veria-ai this is ready for another look

[marcus-linktree]

@veria-ai this is ready for another look

[marcus-linktree]

@greptileai

[yucheng-berri]

Thanks for the contribution. i adopted your PR and merged in #34456