From 8e9e1254bfd63d339e28d417eae40542eb158800 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Fri, 24 Jul 2026 01:31:51 -0700 Subject: [PATCH] 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 | 46 ++++++++----- .../proxy/batches_endpoints/test_endpoints.py | 69 +++++++++++++++---- 2 files changed, 83 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index ed2913fdf14..0e483f5820c 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/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, diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 3092da9e1e5..50b02601f28 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/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"