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

litellm/proxy/batches_endpoints/endpoints.py

@@ -51,13 +51,13 @@ async def _resolve_managed_input_file_storage_url(
 ) -> "str | None":
     """Resolve a managed (unified) input_file_id to its backend storage_url.
 
-    Returns None when the proxy has no database, no managed file row exists,
-    or the row has no storage_url; callers fall back to dispatching the
-    original id so the managed-files deployment hook can still map it via
-    model_file_id_mapping. Raises a 404 when the caller does not own the
-    managed file and a 503 when the lookup fails, so an unverifiable id is
-    never dispatched (the deployment hook maps ids from cache without
-    re-checking ownership).
+    Returns None only when the proxy has no database (nothing to resolve or
+    enforce against) or when an owned managed file has no storage_url yet
+    (legacy rows); callers fall back to dispatching the original id. Raises a
+    404 when the caller does not own the file or no row exists (an id that
+    cannot be verified must not be dispatched, since the managed-files
+    deployment hook maps ids from cache without re-checking ownership) and a
+    503 when the lookup itself fails.
     """
     from litellm.proxy.proxy_server import prisma_client
 
@@ -71,9 +71,7 @@ async def _resolve_managed_input_file_storage_url(
             status_code=503,
             detail={"error": "Unable to verify managed file access; please retry"},
         )
-    if db_file is None:
-        return None
-    if not can_access_resource(
+    if db_file is None or not can_access_resource(
         user_api_key_dict=user_api_key_dict,
         created_by=db_file.created_by,
         resource_team_id=db_file.team_id,
@@ -198,14 +196,6 @@ async def create_batch(
             model_from_file_id = decode_model_from_file_id(input_file_id)
             unified_file_id = _is_base64_encoded_unified_file_id(input_file_id)
 
-        if model_from_file_id is None and unified_file_id and input_file_id:
-            resolved_storage_url = await _resolve_managed_input_file_storage_url(
-                input_file_id=input_file_id,
-                user_api_key_dict=user_api_key_dict,
-            )
-            if resolved_storage_url is not None:
-                _create_batch_data["input_file_id"] = resolved_storage_url
-
         # SCENARIO 1: File ID is encoded with model info
         if model_from_file_id is not None and input_file_id:
             credentials = get_credentials_for_model(
@@ -254,7 +244,12 @@ async def create_batch(
 
             response.input_file_id = input_file_id
 
-        elif litellm.enable_loadbalancing_on_batch_endpoints is True and is_router_model and router_model is not None:
+        elif (
+            litellm.enable_loadbalancing_on_batch_endpoints is True
+            and is_router_model
+            and router_model is not None
+            and not unified_file_id
+        ):
             if llm_router is None:
                 raise HTTPException(
                     status_code=500,
@@ -274,6 +269,19 @@ async def create_batch(
                 )
             model = target_model_names[0]
             _create_batch_data["model"] = model
+
+            # Resolve the opaque unified id to its real backend storage_url and
+            # enforce ownership before dispatch. Kept inside this branch (not
+            # hoisted above load balancing) so the load-balanced path keeps the
+            # original id for model_file_id_mapping deployment filtering, and so
+            # the response restore below still returns the unified id.
+            resolved_storage_url = await _resolve_managed_input_file_storage_url(
+                input_file_id=input_file_id,
+                user_api_key_dict=user_api_key_dict,
+            )
+            if resolved_storage_url is not None:
+                _create_batch_data["input_file_id"] = resolved_storage_url
+
             if llm_router is None:
                 raise HTTPException(
                     status_code=500,

tests/test_litellm/proxy/batches_endpoints/test_endpoints.py

@@ -644,10 +644,12 @@ async def test_create__unified_file_id_db_error_fails_closed_503(harness):
 
 
 @pytest.mark.asyncio
-async def test_create__unified_file_id_loadbalanced_path_gets_resolved_id(harness):
-    """Resolution runs before dispatch branching, so the load-balanced router
-    branch also receives the resolved storage_url instead of the opaque
-    unified id."""
+async def test_create__unified_file_id_with_loadbalancing_uses_resolving_branch(harness):
+    """A unified file must not take the raw load-balanced dispatch branch even
+    when load balancing is enabled and a router model is present; it takes the
+    unified branch that resolves the storage_url and restores the response, so
+    the opaque id never reaches the provider and model_file_id_mapping (keyed on
+    the original id) is not clobbered on the load-balanced path."""
     set_body(
         harness,
         {
@@ -672,24 +674,29 @@ async def test_create__unified_file_id_loadbalanced_path_gets_resolved_id(harnes
     with (
         patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True),
         patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"),
+        patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gemini-2.0"]),
         patch.object(proxy_server, "prisma_client", MagicMock()),
         patch.object(endpoints, "ManagedFileRepository", fake_repo_cls),
     ):
-        await call_create(harness, user=UserAPIKeyAuth(api_key="sk-test", user_id="user-1"))
+        resp = await call_create(harness, user=UserAPIKeyAuth(api_key="sk-test", user_id="user-1"))
 
     assert harness.router_acreate.call_count == 1
+    # The resolving branch fired: model injected from the unified id, storage_url
+    # forwarded, response restored to the unified id (not the internal storage_url).
+    assert harness.router_kwargs()["model"] == "gemini-2.0"
     assert harness.router_kwargs()["input_file_id"] == fake_db_file.storage_url
+    assert resp.input_file_id == "litellm_proxy_unified_id"
+    assert resp._hidden_params["unified_file_id"] == "unified-xyz"
     harness.litellm_acreate.assert_not_called()
 
 
 @pytest.mark.asyncio
-async def test_create__unified_file_id_no_managed_file_record_falls_back_to_raw_id(
-    harness,
-):
-    """If there's no LiteLLM_ManagedFileTable row (or it has no storage_url),
-    fall back to dispatching the original id instead of raising: the
-    managed-files deployment hook can still map it via model_file_id_mapping,
-    and legacy rows without storage_url keep working."""
+async def test_create__unified_file_id_missing_row_fails_closed_404(harness):
+    """With a database present, a managed unified id that has no row cannot be
+    ownership-verified, so it fails closed with a 404 rather than dispatching
+    the opaque id. Dispatching it would both bypass the ownership gate (the
+    deployment hook maps cache-resident ids without re-checking) and hit the
+    Vertex publishers-segment IndexError this PR exists to prevent."""
     set_body(
         harness,
         {
@@ -710,7 +717,43 @@ async def test_create__unified_file_id_no_managed_file_record_falls_back_to_raw_
         patch.object(proxy_server, "prisma_client", MagicMock()),
         patch.object(endpoints, "ManagedFileRepository", fake_repo_cls),
     ):
-        await call_create(harness)
+        with pytest.raises(ProxyException) as exc:
+            await call_create(harness)
+
+    assert exc.value.code == "404"
+    harness.router_acreate.assert_not_called()
+    harness.litellm_acreate.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_create__unified_file_id_owned_legacy_row_without_storage_url_dispatches_raw(
+    harness,
+):
+    """An owned managed file whose row predates storage_url still dispatches the
+    original id (the managed-files deployment hook maps it); ownership is
+    verified, so this is not the fail-closed case."""
+    set_body(
+        harness,
+        {
+            "input_file_id": "litellm_proxy_unified_id",
+            "endpoint": "/v1/chat/completions",
+            "completion_window": "24h",
+        },
+    )
+
+    fake_db_file = MagicMock(storage_url=None, created_by="user-1", team_id=None)
+    find_first = AsyncMock(return_value=fake_db_file)
+    fake_repo_instance = MagicMock()
+    fake_repo_instance.table.find_first = find_first
+    fake_repo_cls = MagicMock(return_value=fake_repo_instance)
+
+    with (
+        patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"),
+        patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gemini-2.0"]),
+        patch.object(proxy_server, "prisma_client", MagicMock()),
+        patch.object(endpoints, "ManagedFileRepository", fake_repo_cls),
+    ):
+        await call_create(harness, user=UserAPIKeyAuth(api_key="sk-test", user_id="user-1"))
 
     assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id"