From df7fa74c932ad213ac7056d032dd6f5f3dbad98f Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Fri, 24 Jul 2026 15:44:05 -0700 Subject: [PATCH] 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 --- litellm/proxy/batches_endpoints/endpoints.py | 7 +--- .../proxy/batches_endpoints/test_endpoints.py | 39 +++++++------------ 2 files changed, 14 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 0e483f5820c..33b2927b580 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -244,12 +244,7 @@ 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 - and not unified_file_id - ): + elif litellm.enable_loadbalancing_on_batch_endpoints is True and is_router_model and router_model is not None: if llm_router is None: raise HTTPException( status_code=500, diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 50b02601f28..42f90977bc2 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -644,12 +644,14 @@ async def test_create__unified_file_id_db_error_fails_closed_503(harness): @pytest.mark.asyncio -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.""" +async def test_create__multi_model_unified_file_with_loadbalancing_keeps_router_branch(harness): + """Regression guard: a multi-model managed file dispatched with an explicit + router model under load balancing must keep taking the load-balanced router + branch, exactly as on the base revision, where the managed-files deployment + hook remaps the unified id per model. Routing it into the unified branch + instead would trip that branch's "exactly one model" 400 and break a path + that works today, so the unified-file resolution must not steal the + load-balanced branch.""" set_body( harness, { @@ -661,32 +663,17 @@ async def test_create__unified_file_id_with_loadbalancing_uses_resolving_branch( ) harness.is_known_model.return_value = True - fake_db_file = MagicMock( - storage_url="gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.0/abc", - 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(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), + patch.object(endpoints, "get_models_from_unified_file_id", return_value=["model-a", "model-b"]), ): - resp = await call_create(harness, user=UserAPIKeyAuth(api_key="sk-test", user_id="user-1")) + await call_create(harness) + # Load-balanced router branch fired with the request unchanged; the unified + # branch (and its single-model 400) was not reached. 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" + assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id" harness.litellm_acreate.assert_not_called()