scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165

refactor(proxy/batches): scope managed-file handling to resolution, drop ownership check

Browse source 
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
This commit is contained in:
Yucheng Zhu 2026-07-24 16:45:04 -07:00
parent df7fa74c93
commit 18dd984902
2 changed files with 28 additions and 103 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
litellm/proxy/batches_endpoints/endpoints.py
View file

@ -12,7 +12,6 @@ from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest
from litellm.llms.base_llm.managed_resources.isolation import can_access_resource
from litellm.proxy._types import *
from litellm.proxy.common_utils.callback_utils import sanitize_openai_provider_metadata
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -45,19 +44,16 @@ from litellm.types.llms.openai import LiteLLMBatchCreateRequest
router = APIRouter()
async def _resolve_managed_input_file_storage_url(
input_file_id: str,
user_api_key_dict: UserAPIKeyAuth,
) -> "str | None":
async def _resolve_managed_input_file_storage_url(input_file_id: str) -> "str | None":
"""Resolve a managed (unified) input_file_id to its backend storage_url.
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.
Provider batch handlers (e.g. Vertex AI, which parses a `publishers/`
segment out of the file URI) need a real storage location; the opaque
unified token crashes them. Returns None when there is no database, the
lookup fails, or the row has no storage_url yet, so callers fall back to
the original id (which the managed-files deployment hook can still map).
Raises a 404 when no managed-file row exists, since the token cannot be
resolved and dispatching it would crash the provider.
"""
from litellm.proxy.proxy_server import prisma_client
@ -67,18 +63,11 @@ async def _resolve_managed_input_file_storage_url(
db_file = await ManagedFileRepository(prisma_client).table.find_first(where={"unified_file_id": input_file_id})
except Exception as e:
verbose_proxy_logger.warning("create_batch: managed file lookup failed for %s: %s", input_file_id, e)
raise HTTPException(
status_code=503,
detail={"error": "Unable to verify managed file access; please retry"},
)
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,
):
return None
if db_file is None:
raise HTTPException(
status_code=404,
detail={"error": f"File not found: {input_file_id}"},
detail={"error": f"Managed file not found: {input_file_id}"},
)
return db_file.storage_url or None
@ -265,15 +254,7 @@ 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,
)
resolved_storage_url = await _resolve_managed_input_file_storage_url(input_file_id)
if resolved_storage_url is not None:
_create_batch_data["input_file_id"] = resolved_storage_url

88
tests/test_litellm/proxy/batches_endpoints/test_endpoints.py
View file

@ -198,9 +198,6 @@ def harness():
stack.enter_context(patch.object(proxy_server, "general_settings", {}))
stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock()))
stack.enter_context(patch.object(proxy_server, "version", "test-version"))
# Default to a database-less proxy so managed-file resolution is a no-op
# unless a test opts in; keeps the routing rows that do not exercise it
# deterministic regardless of cross-file prisma_client pollution.
stack.enter_context(patch.object(proxy_server, "prisma_client", None))
h = Harness(
@ -543,8 +540,6 @@ async def test_create__unified_file_id_resolves_real_storage_url(harness):
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()
@ -557,64 +552,19 @@ async def test_create__unified_file_id_resolves_real_storage_url(harness):
patch.object(proxy_server, "prisma_client", MagicMock()),
patch.object(endpoints, "ManagedFileRepository", fake_repo_cls),
):
resp = await call_create(harness, user=UserAPIKeyAuth(api_key="sk-test", user_id="user-1"))
resp = await call_create(harness)
# The real storage_url - not the opaque unified id - must be what's
# forwarded to the router/provider.
assert harness.router_kwargs()["input_file_id"] == fake_db_file.storage_url
# The lookup key is the RAW base64 id from the request, not the decoded string.
find_first.assert_awaited_once_with(where={"unified_file_id": "litellm_proxy_unified_id"})
# The unified id is still what's returned to the client.
assert resp.input_file_id == "litellm_proxy_unified_id"
assert resp._hidden_params["unified_file_id"] == "unified-xyz"
@pytest.mark.asyncio
async def test_create__unified_file_id_other_tenant_gets_404(harness):
"""Ownership gate: a caller who does not own the managed file (different
user, different team, not an admin) must get a 404, and nothing may be
dispatched. Uses the real can_access_resource so the semantics cannot
drift from the files retrieve/download endpoints."""
set_body(
harness,
{
"input_file_id": "litellm_proxy_unified_id",
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
},
)
fake_db_file = MagicMock(
storage_url="gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.0/abc",
created_by="owner-user",
team_id="owner-team",
)
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),
):
with pytest.raises(ProxyException) as exc:
await call_create(harness, user=UserAPIKeyAuth(api_key="sk-test", user_id="intruder"))
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_db_error_fails_closed_503(harness):
"""A lookup failure must fail closed: the ownership gate could not run, and
the managed-files deployment hook maps unified ids from cache without
re-checking ownership, so dispatching the unverified id would let a caller
use another tenant's file during a database outage. Expect a clear 503 and
no dispatch."""
async def test_create__unified_file_id_db_error_falls_back_to_raw_id(harness):
"""A lookup failure must not abort batch creation. Resolution is best-effort
crash prevention, so on a database error it falls back to dispatching the
original id, which the managed-files deployment hook can still map."""
set_body(
harness,
{
@ -635,12 +585,9 @@ async def test_create__unified_file_id_db_error_fails_closed_503(harness):
patch.object(proxy_server, "prisma_client", MagicMock()),
patch.object(endpoints, "ManagedFileRepository", fake_repo_cls),
):
with pytest.raises(ProxyException) as exc:
await call_create(harness)
await call_create(harness)
assert exc.value.code == "503"
harness.router_acreate.assert_not_called()
harness.litellm_acreate.assert_not_called()
assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id"
@pytest.mark.asyncio
@ -670,8 +617,6 @@ async def test_create__multi_model_unified_file_with_loadbalancing_keeps_router_
):
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
assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id"
harness.litellm_acreate.assert_not_called()
@ -680,10 +625,9 @@ async def test_create__multi_model_unified_file_with_loadbalancing_keeps_router_
@pytest.mark.asyncio
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."""
resolved to a real storage location, so it fails closed with a 404 rather
than dispatching the opaque token, which would hit the Vertex
publishers-segment IndexError this PR exists to prevent."""
set_body(
harness,
{
@ -713,12 +657,12 @@ async def test_create__unified_file_id_missing_row_fails_closed_404(harness):
@pytest.mark.asyncio
async def test_create__unified_file_id_owned_legacy_row_without_storage_url_dispatches_raw(
async def test_create__unified_file_id_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."""
"""A managed file whose row predates the storage_url column still dispatches
the original id (the managed-files deployment hook maps it); the row exists,
so this is not the missing-row fail-closed case."""
set_body(
harness,
{
@ -728,7 +672,7 @@ async def test_create__unified_file_id_owned_legacy_row_without_storage_url_disp
},
)
fake_db_file = MagicMock(storage_url=None, created_by="user-1", team_id=None)
fake_db_file = MagicMock(storage_url=None)
find_first = AsyncMock(return_value=fake_db_file)
fake_repo_instance = MagicMock()
fake_repo_instance.table.find_first = find_first
@ -740,7 +684,7 @@ async def test_create__unified_file_id_owned_legacy_row_without_storage_url_disp
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"))
await call_create(harness)
assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id"