fix(managed_files): check ownership before the batch rate limiter reads the input file and keep key-owned resources across key regenerate

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/hooks/batch_rate_limiter.py

@@ -1128,6 +1128,24 @@ class _PROXY_BatchRateLimiter(CustomLogger):
                     },
                 )
 
+    @staticmethod
+    async def _check_managed_file_ownership(
+        managed_files_obj: object,
+        file_id: str,
+        user_api_key_dict: UserAPIKeyAuth,
+    ) -> None:
+        """Deny before the input file is read, so a caller cannot make the proxy
+        download, count, or charge its own limits for another tenant's file."""
+        from litellm.proxy.openai_files_endpoints.common_utils import ManagedResourceAccessChecker
+
+        if not isinstance(managed_files_obj, ManagedResourceAccessChecker):
+            raise HTTPException(status_code=500, detail="Managed resource ownership validation is unavailable.")
+        if not await managed_files_obj.can_user_call_unified_file_id(file_id, user_api_key_dict):
+            raise HTTPException(
+                status_code=403,
+                detail=f"User {user_api_key_dict.user_id} does not have access to the file {file_id}",
+            )
+
     async def _fetch_managed_file_content(
         self,
         file_id: str,
@@ -1170,6 +1188,8 @@ class _PROXY_BatchRateLimiter(CustomLogger):
         if llm_router is None:
             raise ValueError("llm_router not available. Cannot access managed files.")
 
+        await self._check_managed_file_ownership(managed_files_obj, file_id, user_api_key_dict)
+
         # Use the managed files hook to get file content
         # This properly handles user permissions and file ownership
         file_content: Final = await managed_files_obj.afile_content(

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -122,6 +122,7 @@ from litellm.proxy.management_helpers.access_group_key_sync import (
     sync_key_update_access_group_membership,
 )
 from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
+from litellm.proxy.management_helpers.managed_resource_owner_sync import repoint_key_owned_managed_resources
 from litellm.proxy.management_helpers.object_permission_utils import (  # noqa: F401  # legacy module exports
     ObjectPermissionUpsert,
     _set_object_permission,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
@@ -5957,6 +5958,13 @@ async def _execute_virtual_key_regeneration(
 
     await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
 
+    await repoint_key_owned_managed_resources(
+        prisma_client=prisma_client,
+        existing_key_row=key_in_db,
+        previous_key_token=hashed_api_key,
+        new_key_token=new_token_hash,
+    )
+
     # After credential invalidation, so a failure here can never keep the old key alive.
     await sync_key_regeneration_access_group_membership(
         prisma_client=prisma_client,

litellm/proxy/management_helpers/managed_resource_owner_sync.py

@@ -0,0 +1,44 @@
+"""Carry key-owned managed resources over to a regenerated key's new hash."""
+
+from __future__ import annotations
+
+from typing import Final, Protocol
+
+from litellm.proxy._types import LiteLLM_VerificationToken
+from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper
+from litellm.repositories.table_repositories import ManagedObjectRepository
+
+_KEY_OWNED_TABLES: Final = (
+    "LiteLLM_ManagedFileTable",
+    "LiteLLM_ManagedObjectTable",
+    "LiteLLM_ManagedVectorStoreTable",
+)
+
+
+class _RawWriter(Protocol):
+    async def execute_raw(self, query: str, *args: str) -> int: ...
+
+
+def _raw_writer(prisma_client: object) -> _RawWriter:
+    db: Final = ManagedObjectRepository(prisma_client).prisma_client.db  # pyright: ignore[reportAny]  # untyped Prisma client
+    return writer_wrapper(db)  # pyright: ignore[reportAny, reportReturnType]  # untyped Prisma client behind the pin
+
+
+async def repoint_key_owned_managed_resources(
+    prisma_client: object,
+    existing_key_row: LiteLLM_VerificationToken,
+    previous_key_token: str,
+    new_key_token: str,
+) -> None:
+    """A key with neither a user_id nor a team_id owns its managed files, batches and
+    vector stores as ``key:<hashed token>`` (``resolve_resource_owner_id``), so a
+    regenerated key would otherwise be denied its own resources."""
+    if existing_key_row.user_id is not None or existing_key_row.team_id is not None:
+        return
+    writer: Final = _raw_writer(prisma_client)
+    for table in _KEY_OWNED_TABLES:
+        await writer.execute_raw(
+            f'UPDATE "{table}" SET "created_by" = $2 WHERE "created_by" = $1',
+            f"key:{previous_key_token}",
+            f"key:{new_key_token}",
+        )

tests/integration/management/test_managed_file_create_ownership.py

@@ -10,7 +10,7 @@ import httpx
 import pytest
 from integration._support.client import JSON_OBJECT, Gateway, Scenario, string_value
 from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
-from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse
+from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse
 from pydantic import JsonValue
 
 _FILE: Final = {
@@ -198,3 +198,74 @@ def test_fine_tuning_create_from_own_managed_files_still_dispatches(gateway: Gat
         )
         assert created.status_code == 200, created.text
         assert _upstream_creates(gateway, tenants.handle, "/fine_tuning/jobs") == 1
+
+
+def _upstream_calls(gateway: Gateway, handle: ScenarioHandle, method: str, suffix: str) -> int:
+    response: Final = httpx.get(f"{gateway.upstream_url}/__observations", timeout=15, trust_env=False)
+    assert response.status_code == 200, response.text
+    requests: Final = JSON_OBJECT.validate_json(response.content)["requests"]
+    assert isinstance(requests, list), response.text
+    return sum(
+        1
+        for request in requests
+        if isinstance(request, dict)
+        and request.get("method") == method
+        and handle.scenario_id in str(request.get("path"))
+        and str(request.get("path")).endswith(suffix)
+    )
+
+
+def test_batch_create_from_another_teams_file_is_403_before_the_rate_limiter_reads_it(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario:
+        scenario_id: Final = f"managed-file-create-ownership-{uuid.uuid4().hex}"
+        handle: Final = register_scenario(scenario_id, _routes())
+        scenario.cleanups.callback(delete_scenario, handle)
+        model: Final = scenario.model(api_base=handle.api_base())
+        line: Final = json.dumps(
+            {
+                "custom_id": "req-1",
+                "method": "POST",
+                "url": "/v1/chat/completions",
+                "body": {"model": model, "messages": [{"role": "user", "content": "ping"}], "max_tokens": 5},
+            }
+        )
+        register_scenario(
+            scenario_id,
+            RoutedResponse(
+                content_type="application/x-routed",
+                routes={
+                    **_routes().routes,
+                    "GET /files/file-in-$REQUEST_ID/content": TextResponse(
+                        content_type="application/jsonl", body=line + "\n"
+                    ),
+                },
+            ),
+        )
+        team_a: Final = scenario.team(models=[model])
+        team_b: Final = scenario.team(models=[model])
+        owner: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_a)
+        outsider: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_b, tpm_limit=1)
+        owner_file_id: Final = _upload(scenario, owner, model, "batch")
+        denied: Final = _create_batch(gateway, outsider, model, owner_file_id)
+        assert denied.status_code == 403, denied.text
+        assert owner_file_id in denied.text, denied.text
+        assert _upstream_calls(gateway, handle, "GET", "/content") == 0
+
+
+def test_regenerated_key_without_user_or_team_keeps_its_managed_files(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario:
+        tenants: Final = _tenants(scenario, "batch")
+        original: Final = string_value(gateway.post("/key/generate", {})["key"])
+        file_id: Final = _upload(scenario, original, tenants.model, "batch")
+        regenerated: Final = string_value(gateway.post("/key/regenerate", {"key": original})["key"])
+        scenario.cleanups.callback(scenario.delete_key, regenerated)
+        batch: Final = _create_batch(gateway, regenerated, tenants.model, file_id)
+        retrieved: Final = gateway.request("GET", f"/v1/files/{file_id}", key=regenerated)
+        stranger: Final = _create_batch(
+            gateway, string_value(gateway.post("/key/generate", {})["key"]), tenants.model, file_id
+        )
+        assert (batch.status_code, retrieved.status_code, stranger.status_code) == (200, 200, 403), (
+            batch.text,
+            retrieved.text,
+            stranger.text,
+        )

tests/unit/proxy/hooks/test_batch_rate_limiter.py

@@ -11,6 +11,7 @@ import time
 from collections.abc import Iterator, Sequence
 from datetime import datetime, timezone
 from typing import Final, Literal
+from unittest.mock import AsyncMock, MagicMock
 
 import httpx
 import pytest
@@ -481,10 +482,9 @@ async def test_batch_rate_limiter_without_user_context(openai_files: respx.MockR
     assert usage_without_context.total_tokens == usage_with_context.total_tokens == _token_counter_total(rows)
 
 
-@pytest.mark.asyncio
-async def test_managed_file_is_read_through_the_managed_files_hook_with_user_context(
-    openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
-):
+async def _install_managed_file_hook(
+    openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rows: str
+) -> tuple[respx.Route, respx.Route]:
     from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
 
     from litellm import Router
@@ -494,24 +494,22 @@ async def test_managed_file_is_read_through_the_managed_files_hook_with_user_con
     from litellm.proxy.utils import ProxyLogging
 
     assert is_base64_encoded_unified_file_id(_MANAGED_FILE_ID)
-    rows: Final = _batch_rows(("Test message for regression",))
     provider_route: Final = _serve_file(openai_files, "file-provider-1", rows)
-    standard_route: Final = _serve_file(openai_files, "file-abc123", rows)
     unrouted_managed_read: Final = _serve_file(openai_files, _MANAGED_FILE_ID, rows)
-    file_cache: Final = InternalUsageCache(dual_cache=DualCache())
-    await file_cache.async_set_cache(
-        key=_MANAGED_FILE_ID,
-        value=LiteLLM_ManagedFileTable(
-            unified_file_id=_MANAGED_FILE_ID,
-            model_mappings={"deployment-1": "file-provider-1"},
-            flat_model_file_ids=["file-provider-1"],
-            created_by="test-user-regression",
-        ).model_dump(),
-        litellm_parent_otel_span=None,
+    row: Final = LiteLLM_ManagedFileTable(
+        unified_file_id=_MANAGED_FILE_ID,
+        model_mappings={"deployment-1": "file-provider-1"},
+        flat_model_file_ids=["file-provider-1"],
+        created_by="test-user-regression",
+        team_id="team-owner",
     )
+    file_cache: Final = InternalUsageCache(dual_cache=DualCache())
+    await file_cache.async_set_cache(key=_MANAGED_FILE_ID, value=row.model_dump(), litellm_parent_otel_span=None)
+    prisma_client: Final = MagicMock()
+    prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=row)
     proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache())
     proxy_logging.proxy_hook_mapping["managed_files"] = _PROXY_LiteLLMManagedFiles(
-        internal_usage_cache=file_cache, prisma_client=None
+        internal_usage_cache=file_cache, prisma_client=prisma_client
     )
     router: Final = Router(
         model_list=[
@@ -524,6 +522,16 @@ async def test_managed_file_is_read_through_the_managed_files_hook_with_user_con
     )
     monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging)
     monkeypatch.setattr(proxy_server, "llm_router", router)
+    return provider_route, unrouted_managed_read
+
+
+@pytest.mark.asyncio
+async def test_managed_file_is_read_through_the_managed_files_hook_with_user_context(
+    openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
+):
+    rows: Final = _batch_rows(("Test message for regression",))
+    provider_route, unrouted_managed_read = await _install_managed_file_hook(openai_files, monkeypatch, rows)
+    standard_route: Final = _serve_file(openai_files, "file-abc123", rows)
     _, _, batch_limiter = _make_limiters()
     user_api_key_dict: Final = UserAPIKeyAuth(
         api_key="test-key-regression", user_id="test-user-regression", tpm_limit=1000, rpm_limit=10
@@ -541,3 +549,25 @@ async def test_managed_file_is_read_through_the_managed_files_hook_with_user_con
     assert not unrouted_managed_read.called
     assert managed_usage.request_count == standard_usage.request_count == 1
     assert managed_usage.total_tokens == standard_usage.total_tokens == _token_counter_total(rows)
+
+
+@pytest.mark.asyncio
+async def test_managed_file_of_another_tenant_is_denied_before_it_is_read(
+    openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
+):
+    rows: Final = _batch_rows(("Test message for regression",))
+    provider_route, unrouted_managed_read = await _install_managed_file_hook(openai_files, monkeypatch, rows)
+    _, _, batch_limiter = _make_limiters()
+    outsider: Final = UserAPIKeyAuth(api_key="test-key-outsider", user_id="outsider", team_id="team-other", tpm_limit=1)
+
+    with pytest.raises(HTTPException) as denied:
+        await batch_limiter.async_pre_call_hook(
+            user_api_key_dict=outsider,
+            cache=DualCache(),
+            data=dict(_create_batch_data(_MANAGED_FILE_ID)),
+            call_type="acreate_batch",
+        )
+
+    assert denied.value.status_code == 403
+    assert not provider_route.called
+    assert not unrouted_managed_read.called