fix(proxy/batches): resolve managed unified input_file_id to storage_url with ownership check before dispatch by yucheng-berri · Pull Request #34474 · BerriAI/litellm

yucheng-berri · 2026-07-24T06:51:23Z

Relevant issues

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

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 (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

Screenshots / Proof of Fix

Verified end to end on a live proxy against real infrastructure at the current HEAD, no mocks: a real Vertex Gemini batchPredictionJobs dispatch, a Postgres-backed proxy, and the enterprise managed-files flow. The batch input was uploaded as a managed file with target_model_names, then batches.create was called with the returned unified id

Owner path, the request succeeds and dispatches a real Vertex job:

$ curl -s http://localhost:4021/v1/files -H "Authorization: Bearer $OWNER" \
    -F purpose=batch -F file=@input.jsonl -F target_model_names=gemini-2.5-flash
# -> {"id":"<base64 unified_file_id>", ...}

$ curl -s -w "\n%{http_code}\n" http://localhost:4021/v1/batches -H "Authorization: Bearer $OWNER" \
    -H "Content-Type: application/json" \
    -d '{"input_file_id":"<unified_file_id>","endpoint":"/v1/chat/completions","completion_window":"24h"}'
# -> {"id":"...","object":"batch","status":"validating","input_file_id":"<unified_file_id>", ...}
# 200

The response input_file_id is the original unified id, not the internal storage_url. Vertex confirms a real job was created (queried with the same service-account credentials the proxy used):

batchPredictionJobs/2335769800166342656 -> JOB_STATE_SUCCEEDED  model=publishers/google/models/gemini-2.5-flash

and the proxy log shows the opaque unified id was resolved to the real backend path before dispatch:

input_file_id='gs://.../litellm-vertex-files/publishers/google/models/gemini-2.5-flash/4f81513a-...'

The resolution fails closed. With the managed-file row removed (database present), the owner's own retry returns 404 rather than dispatching the opaque token, which the provider cannot parse into a real file:

$ # DELETE FROM "LiteLLM_ManagedFileTable"; then owner retries the create above
# 404

The owner run created a real Vertex job; the missing-row denial created none. Correctness is also locked by the routing-contract tests in tests/test_litellm/proxy/batches_endpoints/test_endpoints.py (78 passed, 2 xfailed). The behavior-defining tests were mutation-checked: reverting the missing-row 404 to a raw-id fallback, dropping the response restore, and reverting the lookup key to the decoded string each make the corresponding test fail

Regression safety, base vs this branch, live

The same four batch requests were run against a proxy on litellm_internal_staging (base) and a proxy on this branch, same config, same Postgres, real Vertex Gemini batchPredictionJobs, no mocks. Every path a normal caller exercises is byte-identical; the only difference is a crash that is now a clean error

Case	Base (`litellm_internal_staging`)	This branch	Meaning
Non-managed raw `gs://` batch (`custom_llm_provider: vertex_ai`)	200, real Vertex job	200, real Vertex job	identical, untouched path
Managed unified batch, owner key	200, real Vertex job	200, real Vertex job	identical, happy path preserved
Managed unified batch, managed-file row missing	500 `list index out of range`	404	crash becomes a clean, fail-closed error
Multi-model managed file, load balancing on, explicit `model`	200, routed by load balancing	200, routed by load balancing	identical, verified after removing an over-broad branch guard

No legitimate workflow that worked on staging behaves differently here. The only backward-incompatible surface is that a managed-file id with no backing row, which previously crashed the provider, now returns a clean 404

Type

🐛 Bug Fix

Changes

Adopts the storage_url resolution from the original PR, scoped to that one problem. The managed unified input_file_id is resolved to its backend storage_url before dispatch, so provider batch handlers that parse a real path (Vertex splits on publishers/) receive a gs:// URI instead of the opaque token. The substitution stays on the single-model unified branch, so the load-balanced branch keeps the original id for the deployment hook's model_file_id_mapping. When a managed id has no backing row the request fails closed with a 404, so the token is never dispatched into the provider crash; database lookup errors and legacy rows without a storage_url fall back to the original id

Two things surfaced during adoption. The lookup key was wrong: LiteLLM_ManagedFileTable.unified_file_id stores the raw base64 id (see schema.prisma and the enterprise hook), but the original change queried with the decoded litellm_proxy:... string, so it never matched in production; it now queries with the raw id, locked by a regression test. The inspect.isawaitable null-guard is also removed, with the tests using AsyncMock so production code carries no test-harness accommodations

Cross-tenant ownership is deliberately out of scope. Batch create had no ownership check before this PR and the gap spans every managed-file call type, so it is being fixed at the source in the enterprise managed-files pre-call hook (tracked in LIT-4797) rather than partially patched in this one endpoint

Behavior changes

When a managed file row with a storage_url exists, the unified batch branch dispatches the resolved storage_url instead of the opaque token; the response still returns the original unified input_file_id
A managed unified id with no backing row (database present) now returns 404 instead of dispatching a token that crashes the provider
No behavior change for non-managed batches, the load-balanced path, when the proxy has no database, when the row has no storage_url, or on a lookup error; all fall back to prior behavior

QA runbook

Upload a batch input file with target_model_names set (managed files, enterprise) against a vertex_ai model
client.batches.create(input_file_id=<unified id>, endpoint="/v1/chat/completions", completion_window="24h") with the owning key succeeds and the Vertex job receives the real gs:// storage path
Delete the LiteLLM_ManagedFileTable row and repeat step 2: expect a clean 404 instead of a 500 list index out of range

Credit

Adopted from #34260 by @htourinho-clgx. Mirrored onto a litellm_ branch so CircleCI and the internal lint workflow run; the original commits are preserved with their authorship

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-create auth and file dispatch for managed unified IDs (tenant isolation and fail-closed DB errors); scope is limited to the unified single-model branch with strong test coverage.

Overview
Batch create for managed (unified) input_file_id now looks up LiteLLM_ManagedFileTable by the raw base64 id, enforces can_access_resource (404 if missing row or wrong tenant), and substitutes storage_url before router dispatch so providers like Vertex get a real gs:// URI instead of the opaque token. DB lookup errors return 503 with no dispatch; legacy rows without storage_url still send the original id after ownership passes.

Resolution runs only on the single-model unified branch (not load-balanced multi-model paths, which keep the unified id for deployment-hook mapping). Responses still return the client’s unified input_file_id.

Tests add routing-contract coverage for resolution, cross-tenant denial, fail-closed 404/503, load-balancing precedence, and default prisma_client=None in the harness.

^{Reviewed by Cursor Bugbot for commit df7fa74. Bugbot is set up for automated code reviews on this repo. Configure here.}

…atch create litellm.create_batch() against a Vertex AI-backed model crashes with an opaque error when the input file was uploaded as a LiteLLM-managed 'unified file' (multi-model file upload). The base64-encoded unified_file_id token is a LiteLLM-internal identifier, not a real provider-side file reference, but the batches_endpoints create_batch handler forwards it unchanged to llm_router.acreate_batch() / litellm.acreate_batch() for the unified_file_id branch. Provider-specific code that expects a real file location (e.g. Vertex AI's batch transformation, which parses a 'publishers/' segment out of the GCS URI) then fails on the opaque token. Resolve the unified_file_id to its real backend location (LiteLLM_ManagedFileTable.storage_url) before dispatch, mirroring the same lookup already used by the files retrieve/download endpoints for managed files. Falls back to the previous (unchanged) behavior if no managed-file record exists. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

… test harnesses

…olving managed input_file_id The adopted resolution queried LiteLLM_ManagedFileTable with the decoded litellm_proxy string, but the unified_file_id column stores the raw base64 file id (see schema.prisma and the enterprise managed-files hook), so the lookup never matched in production and silently fell back to the opaque id. Query with the raw id instead and lock the key with a regression test. Move the resolution above the dispatch branches so the load-balanced router path receives the resolved storage_url too, enforce managed-file ownership with the same can_access_resource semantics the files retrieve and download endpoints use (404 on denial), and downgrade database failures to a logged fallback instead of aborting batch creation. Unresolved ids still dispatch unchanged because the managed-files deployment hook can map them via model_file_id_mapping

yucheng-berri · 2026-07-24T06:51:36Z

@greptileai

greptile-apps · 2026-07-24T06:53:35Z

Greptile Summary

Resolves managed unified batch input IDs to backend storage URLs before single-model dispatch.

Returns 404 when the managed-file row is missing and 503 when its database lookup fails.
Preserves the original unified ID in batch responses and leaves load-balanced multi-model routing unchanged.
Adds regression tests for resolution, failure handling, legacy rows, response restoration, and routing precedence.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failures remain.

Important Files Changed

Filename	Overview
litellm/proxy/batches_endpoints/endpoints.py	Adds fail-closed managed-file lookup and storage URL substitution before single-model batch dispatch.
tests/test_litellm/proxy/batches_endpoints/test_endpoints.py	Adds focused regression coverage for managed unified-file resolution and preserves existing routing behavior.

_{Reviews (8): Last reviewed commit: "fix(proxy/batches): fail closed with 503..." | Re-trigger Greptile}

devin-ai-integration

✅ Devin Review: No Issues Found

Devin Review analyzed this PR and found no bugs or issues to report.

veria-ai · 2026-07-24T06:55:46Z

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: 1 · PR risk: 0/10

codecov · 2026-07-24T06:57:06Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

…p errors A lookup exception previously fell back to dispatching the original unified id with the ownership gate unexecuted; the managed-files deployment hook maps unified ids from cache without re-checking ownership, so a database outage let a caller dispatch another tenant's file. Raise a clear 503 instead and lock the behavior with a regression test. No-database and no-row cases still fall back unchanged

yucheng-berri · 2026-07-24T07:00:16Z

@greptileai

yucheng-berri · 2026-07-24T07:13:17Z

bugbot run

codspeed-hq · 2026-07-24T07:41:39Z

Merging this PR will not alter performance

✅ 31 untouched benchmarks

_{Comparing litellm_batch_unified_file_resolution (891c6e4) with litellm_internal_staging (7263aa0)¹}

No successful run was found on litellm_internal_staging (842f32d) during the generation of this report, so 7263aa0 was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩

The batch routing harness left proxy_server.prisma_client at its module global, which a sibling test in the same shard can leave as a MagicMock. The unified-file rows that do not opt into managed-file resolution then entered the resolver and awaited a non-awaitable mock, surfacing as a 503. Patch prisma_client to None by default so those rows stay a no-op; resolution tests still override it explicitly

yucheng-berri · 2026-07-24T08:15:09Z

@greptileai

…l closed on missing row Cursor flagged that hoisting the storage_url substitution above the load-balanced dispatch branch broke two things on that path: the model_file_id_mapping deployment filter keys on the original unified id, and the response returned the internal storage_url instead of the unified id. Move the resolution back inside the unified branch and exclude unified ids from the load-balanced branch so a managed file always takes the resolving path (which restores input_file_id and the unified_file_id hidden param on the response), and a load-balanced batch keeps the original id for deployment filtering. Also fail closed with a 404 when a unified id has no managed-file row while a database is present: the id cannot be ownership-verified, and dispatching it would both bypass the gate and hit the Vertex publishers-segment IndexError. Owned rows without a storage_url (legacy) still dispatch the original id

yucheng-berri · 2026-07-24T08:32:32Z

@greptileai

… branch Excluding unified ids from the load-balanced branch (and not unified_file_id) regressed a path that works on the base revision: a multi-model managed file dispatched with an explicit router model under load balancing was routed into the unified branch, which raises a 400 for anything other than exactly one target model. Verified live against base (200, managed-files deployment hook remaps the unified id per model) versus the guarded branch (400 Expected 1 model, got 2). Restore the original three-condition load-balanced branch so that path keeps working unchanged. Unified-file storage_url resolution and the ownership 404 still apply on the non-load-balanced unified branch, which is the common managed-batch flow; the load-balanced managed path retains its existing behavior and its pre-existing enterprise-hook ownership gap, unchanged from base

yucheng-berri · 2026-07-24T22:44:28Z

@greptileai

yucheng-berri · 2026-07-24T22:47:32Z

bugbot run

cursor

Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit df7fa74. Configure here.}

cursor · 2026-07-24T22:51:59Z

                )
            model = target_model_names[0]
            _create_batch_data["model"] = model
+


LB path skips file ownership

High Severity

Moving the _resolve_managed_input_file_storage_url call into the unified-file branch bypasses ownership checks for load-balanced batch creations. When enable_loadbalancing_on_batch_endpoints is active and a model is specified, the ownership check is skipped, allowing cross-tenant access to managed files.

^{Reviewed by Cursor Bugbot for commit df7fa74. Configure here.}

Superseded: as of the latest commit the resolve call does no ownership check on any branch, by design. Ownership for batch create is out of scope for this resolution PR (a pre-existing gap across all managed-file call types) and is tracked for a source fix in the enterprise pre-call hook in LIT-4797. The PR description was updated to say so

…rop ownership check Narrow this PR to its one problem: resolving a managed unified input_file_id to its backend storage_url so provider batch handlers (Vertex parses a publishers/ segment) receive a real location instead of the opaque token, and failing closed with a 404 when the token has no backing row so it is never dispatched into the provider crash. Remove the cross-tenant ownership check (can_access_resource) added earlier. Batch-create had no ownership enforcement before this PR, and the gap spans every managed-file call type, so it belongs in the enterprise managed-files pre-call hook (its acreate_batch branch) where files, batches and fine-tuning are covered uniformly, not partially in this one endpoint. Filed as a follow-up. This also removes the load-balanced-path ownership inconsistency the bots flagged, since there is no ownership branch to skip. Drop the inline comments flagged against the no-comments rule; behavior is documented in the helper docstring and the test docstrings

yucheng-berri · 2026-07-24T23:50:15Z

@greptileai

… errors A lookup exception previously fell back to dispatching the unresolved unified token, which defeats the fail-closed guarantee: the token still reaches the provider and can hit the same publishers-segment IndexError the resolution prevents. Treat a lookup error like the missing-row case and fail closed, but with a retryable 503 since the condition is transient. No-database and no-storage_url rows still fall back unchanged

yucheng-berri · 2026-07-25T00:07:36Z

@greptileai

yucheng-berri · 2026-07-25T01:20:49Z

Follow-up: #34584 restores the pre-existing fall-back behavior for managed batch input ids that cannot be resolved (missing row / lookup error), removing the 404 and 503 this PR added. Those two responses were backward-incompatible against the behavior before this PR (the same requests previously dispatched the original id via the managed-files deployment hook). #34584 keeps the storage_url resolution and makes the managed-file handling strictly additive. A deliberate fail-closed at the source (managed-files pre-call hook, all call types) is tracked separately in LIT-4797

…url with ownership check before dispatch (BerriAI#34474) * fix: resolve unified_file_id to real storage_url before dispatching batch create litellm.create_batch() against a Vertex AI-backed model crashes with an opaque error when the input file was uploaded as a LiteLLM-managed 'unified file' (multi-model file upload). The base64-encoded unified_file_id token is a LiteLLM-internal identifier, not a real provider-side file reference, but the batches_endpoints create_batch handler forwards it unchanged to llm_router.acreate_batch() / litellm.acreate_batch() for the unified_file_id branch. Provider-specific code that expects a real file location (e.g. Vertex AI's batch transformation, which parses a 'publishers/' segment out of the GCS URI) then fails on the opaque token. Resolve the unified_file_id to its real backend location (LiteLLM_ManagedFileTable.storage_url) before dispatch, mirroring the same lookup already used by the files retrieve/download endpoints for managed files. Falls back to the previous (unchanged) behavior if no managed-file record exists. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix(proxy/batches): null-guard await on find_first for sync MagicMock test harnesses * fix(proxy/batches): enforce ownership and correct lookup key when resolving managed input_file_id The adopted resolution queried LiteLLM_ManagedFileTable with the decoded litellm_proxy string, but the unified_file_id column stores the raw base64 file id (see schema.prisma and the enterprise managed-files hook), so the lookup never matched in production and silently fell back to the opaque id. Query with the raw id instead and lock the key with a regression test. Move the resolution above the dispatch branches so the load-balanced router path receives the resolved storage_url too, enforce managed-file ownership with the same can_access_resource semantics the files retrieve and download endpoints use (404 on denial), and downgrade database failures to a logged fallback instead of aborting batch creation. Unresolved ids still dispatch unchanged because the managed-files deployment hook can map them via model_file_id_mapping * fix(proxy/batches): fail closed when the managed file ownership lookup errors A lookup exception previously fell back to dispatching the original unified id with the ownership gate unexecuted; the managed-files deployment hook maps unified ids from cache without re-checking ownership, so a database outage let a caller dispatch another tenant's file. Raise a clear 503 instead and lock the behavior with a regression test. No-database and no-row cases still fall back unchanged * test(proxy/batches): default harness prisma_client to None The batch routing harness left proxy_server.prisma_client at its module global, which a sibling test in the same shard can leave as a MagicMock. The unified-file rows that do not opt into managed-file resolution then entered the resolver and awaited a non-awaitable mock, surfacing as a 503. Patch prisma_client to None by default so those rows stay a no-op; resolution tests still override it explicitly * fix(proxy/batches): keep unified resolution in its own branch and fail closed on missing row Cursor flagged that hoisting the storage_url substitution above the load-balanced dispatch branch broke two things on that path: the model_file_id_mapping deployment filter keys on the original unified id, and the response returned the internal storage_url instead of the unified id. Move the resolution back inside the unified branch and exclude unified ids from the load-balanced branch so a managed file always takes the resolving path (which restores input_file_id and the unified_file_id hidden param on the response), and a load-balanced batch keeps the original id for deployment filtering. Also fail closed with a 404 when a unified id has no managed-file row while a database is present: the id cannot be ownership-verified, and dispatching it would both bypass the gate and hit the Vertex publishers-segment IndexError. Owned rows without a storage_url (legacy) still dispatch the original id * fix(proxy/batches): do not divert unified files off the load-balanced branch Excluding unified ids from the load-balanced branch (and not unified_file_id) regressed a path that works on the base revision: a multi-model managed file dispatched with an explicit router model under load balancing was routed into the unified branch, which raises a 400 for anything other than exactly one target model. Verified live against base (200, managed-files deployment hook remaps the unified id per model) versus the guarded branch (400 Expected 1 model, got 2). Restore the original three-condition load-balanced branch so that path keeps working unchanged. Unified-file storage_url resolution and the ownership 404 still apply on the non-load-balanced unified branch, which is the common managed-batch flow; the load-balanced managed path retains its existing behavior and its pre-existing enterprise-hook ownership gap, unchanged from base * refactor(proxy/batches): scope managed-file handling to resolution, drop ownership check Narrow this PR to its one problem: resolving a managed unified input_file_id to its backend storage_url so provider batch handlers (Vertex parses a publishers/ segment) receive a real location instead of the opaque token, and failing closed with a 404 when the token has no backing row so it is never dispatched into the provider crash. Remove the cross-tenant ownership check (can_access_resource) added earlier. Batch-create had no ownership enforcement before this PR, and the gap spans every managed-file call type, so it belongs in the enterprise managed-files pre-call hook (its acreate_batch branch) where files, batches and fine-tuning are covered uniformly, not partially in this one endpoint. Filed as a follow-up. This also removes the load-balanced-path ownership inconsistency the bots flagged, since there is no ownership branch to skip. Drop the inline comments flagged against the no-comments rule; behavior is documented in the helper docstring and the test docstrings * fix(proxy/batches): fail closed with 503 when the managed-file lookup errors A lookup exception previously fell back to dispatching the unresolved unified token, which defeats the fail-closed guarantee: the token still reaches the provider and can hit the same publishers-segment IndexError the resolution prevents. Treat a lookup error like the missing-row case and fail closed, but with a retryable 503 since the condition is transient. No-database and no-storage_url rows still fall back unchanged --------- Co-authored-by: htourinho-clgx <htourinho@cotality.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

htourinho-clgx and others added 3 commits July 23, 2026 23:32

fix(proxy/batches): null-guard await on find_first for sync MagicMock…

411b525

… test harnesses