From 9f5148aafa1b1c96d83c060f6e5fb60d06415a47 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Sat, 10 Oct 2026 11:22:47 -0700 Subject: [PATCH 1/5] test(managed_files): pin ownership checks on batch and fine-tuning create Tenant B can create a batch or fine-tuning job from tenant A's managed file id today. These fail on main for acreate_batch input_file_id and acreate_fine_tuning_job training_file and validation_file --- .../proxy/hooks/test_managed_files.py | 92 +++++++++++++++++++ 1 file changed, 92 insertions(+) diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index ed80bd9b928..ad655114c4e 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -5264,3 +5264,95 @@ async def test_list_user_batches_target_model_names_filter_rejected_with_400(): assert exc.value.type == "invalid_request_error" assert exc.value.param == "target_model_names" assert exc.value.message == "Filtering by 'target_model_names' is not supported when using managed batches." + + +def _tenant_a_input_file_id() -> str: + raw: Final = ( + "litellm_proxy:application/octet-stream;unified_id,7d1c4b0e-5a7f-4a43-9a0e-3f8e2a1b6c11;" + "target_model_names,gpt-4o-mini;llm_output_file_id,file-tenant-a;llm_output_file_model_id,model-a" + ) + return base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=") + + +def _tenant_a_input_file_row(unified_file_id: str) -> LiteLLM_ManagedFileTable: + return LiteLLM_ManagedFileTable( + unified_file_id=unified_file_id, + file_object=None, + model_mappings={"model-a": "file-tenant-a"}, + flat_model_file_ids=["file-tenant-a"], + created_by="user_a", + team_id="team_a", + ) + + +_INPUT_FILE_CREATE_CALLS: Final = ( + ("acreate_batch", "input_file_id"), + ("acreate_fine_tuning_job", "training_file"), + ("acreate_fine_tuning_job", "validation_file"), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("call_type", "file_field"), _INPUT_FILE_CREATE_CALLS) +async def test_create_from_another_tenants_managed_file_is_denied(call_type: str, file_field: str) -> None: + from litellm.proxy._types import UserAPIKeyAuth + + unified_file_id: Final = _tenant_a_input_file_id() + managed_files, _ = _managed_files_with_fake_prisma(_tenant_a_input_file_row(unified_file_id)) + + with pytest.raises(HTTPException) as exc_info: + await managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id="user_b", team_id="team_b", parent_otel_span=MagicMock()), + cache=MagicMock(), + data={"model": "gpt-4o-mini", "training_file": "file-own-training", file_field: unified_file_id}, + call_type=call_type, + ) + + assert exc_info.value.status_code == 403, f"{call_type}.{file_field} let tenant B use tenant A's file" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("call_type", "file_field"), _INPUT_FILE_CREATE_CALLS) +@pytest.mark.parametrize( + "caller_kwargs", + [ + {"user_id": "user_a", "team_id": "team_a"}, + {"user_id": "teammate_of_a", "team_id": "team_a"}, + {"user_id": "admin", "team_id": "team_b", "user_role": "proxy_admin"}, + ], + ids=["owner", "same_team", "proxy_admin"], +) +async def test_create_from_accessible_managed_file_is_allowed( + call_type: str, file_field: str, caller_kwargs: dict[str, str] +) -> None: + from litellm.proxy._types import UserAPIKeyAuth + + unified_file_id: Final = _tenant_a_input_file_id() + managed_files, _ = _managed_files_with_fake_prisma(_tenant_a_input_file_row(unified_file_id)) + + result: Final = await managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(**caller_kwargs, parent_otel_span=MagicMock()), + cache=MagicMock(), + data={"model": "gpt-4o-mini", file_field: unified_file_id}, + call_type=call_type, + ) + + assert isinstance(result, dict) and result[file_field] == unified_file_id, result + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("call_type", "file_field"), _INPUT_FILE_CREATE_CALLS) +async def test_create_from_raw_provider_file_id_skips_managed_ownership_lookup(call_type: str, file_field: str) -> None: + from litellm.proxy._types import UserAPIKeyAuth + + managed_files, table = _managed_files_with_fake_prisma() + + result: Final = await managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id="user_b", team_id="team_b", parent_otel_span=MagicMock()), + cache=MagicMock(), + data={"model": "gpt-4o-mini", file_field: "file-raw-provider-id"}, + call_type=call_type, + ) + + assert isinstance(result, dict) and result[file_field] == "file-raw-provider-id", result + assert not any("unified_file_id" in where for where in table.find_first_calls), table.find_first_calls From f673922070582e8b1df6451f4f95c6ffbcf9f5d2 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Sat, 10 Oct 2026 11:36:46 -0700 Subject: [PATCH 2/5] fix(managed_files): check input file ownership on batch and fine-tuning create Batch create and fine-tuning create only mapped the managed input file to provider ids and never checked who owned it, so any key with access to the model could build a batch or job on another tenant's file and then own the outputs. Both now run check_file_ids_access like chat and responses do, which 403s a non-owner and 404s a managed id with no ownership row. Fine-tuning also checks validation_file, which was never looked at --- .../proxy/hooks/managed_files.py | 14 ++++++++------ .../proxy/hooks/test_managed_files.py | 18 ++++++++++++++++++ 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 4c06f0b9802..67894fbbf3c 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1162,7 +1162,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) return False - async def check_file_ids_access(self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth) -> None: + async def check_file_ids_access(self, file_ids: Sequence[str], user_api_key_dict: UserAPIKeyAuth) -> None: """ Check if the user has access to a list of file IDs. Only checks managed (unified) file IDs. @@ -1269,6 +1269,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): elif call_type == CallTypes.acreate_batch.value: input_file_id = cast(Optional[str], data.get("input_file_id")) if input_file_id: + await self.check_file_ids_access((input_file_id,), user_api_key_dict) model_file_id_mapping = await self.get_model_file_id_mapping( [input_file_id], user_api_key_dict.parent_otel_span ) @@ -1315,11 +1316,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): elif retrieve_object_id and accessor_key == "batch_id": await self.enforce_batch_object_access(retrieve_object_id, user_api_key_dict) elif call_type == CallTypes.acreate_fine_tuning_job.value: - input_file_id = cast(Optional[str], data.get("training_file")) - if input_file_id: - model_file_id_mapping = await self.get_model_file_id_mapping( - [input_file_id], user_api_key_dict.parent_otel_span - ) + fine_tuning_file_ids: Final = tuple( + file_id + for file_id in (data.get("training_file"), data.get("validation_file")) + if isinstance(file_id, str) and file_id + ) + await self.check_file_ids_access(fine_tuning_file_ids, user_api_key_dict) return data diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index ad655114c4e..0b22d07ee00 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -5311,6 +5311,24 @@ async def test_create_from_another_tenants_managed_file_is_denied(call_type: str assert exc_info.value.status_code == 403, f"{call_type}.{file_field} let tenant B use tenant A's file" +@pytest.mark.asyncio +@pytest.mark.parametrize(("call_type", "file_field"), _INPUT_FILE_CREATE_CALLS) +async def test_create_from_managed_file_id_with_no_db_row_is_404(call_type: str, file_field: str) -> None: + from litellm.proxy._types import UserAPIKeyAuth + + managed_files, _ = _managed_files_with_fake_prisma() + + with pytest.raises(HTTPException) as exc_info: + await managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id="user_a", team_id="team_a", parent_otel_span=MagicMock()), + cache=MagicMock(), + data={"model": "gpt-4o-mini", file_field: _tenant_a_input_file_id()}, + call_type=call_type, + ) + + assert exc_info.value.status_code == 404, f"{call_type}.{file_field} dispatched a managed id with no owner row" + + @pytest.mark.asyncio @pytest.mark.parametrize(("call_type", "file_field"), _INPUT_FILE_CREATE_CALLS) @pytest.mark.parametrize( From e4e1eb9939264b752722a1bbb48864ca61368e02 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Sat, 10 Oct 2026 12:53:07 -0700 Subject: [PATCH 3/5] test(batches): pin that a pre-call ownership denial stops every batch create branch LIT-4797 noted an endpoint-level check missed the load-balanced branch because it dispatches first. The ownership check lives in the pre-call hook, which runs before branch selection, so a 403 there must keep both the load-balanced and unified-file branches from reaching the router --- .../proxy/batches_endpoints/test_endpoints.py | 30 ++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/tests/unit/proxy/batches_endpoints/test_endpoints.py b/tests/unit/proxy/batches_endpoints/test_endpoints.py index 3469dbee748..aab2074e49f 100644 --- a/tests/unit/proxy/batches_endpoints/test_endpoints.py +++ b/tests/unit/proxy/batches_endpoints/test_endpoints.py @@ -57,7 +57,7 @@ from litellm.types.llms.openai import BatchJobStatus from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from litellm.types.utils import CredentialItem, LiteLLMBatch, SpecialEnums -from fastapi import Request, Response +from fastapi import HTTPException, Request, Response # --------------------------------------------------------------------------- # # Fixtures: distinguishable credentials per model so a wrong/hardcoded model_id @@ -1164,6 +1164,34 @@ async def test_create__uses_acreate_batch_route_type(harness, openai_env_creds): assert harness.pre_call.call_args.kwargs["route_type"] == "acreate_batch" +@pytest.mark.asyncio +@pytest.mark.parametrize("loadbalancing", [True, False], ids=["load_balanced", "unified_branch"]) +async def test_create__pre_call_ownership_denial_stops_every_dispatch_branch(harness, loadbalancing: bool): + set_body( + harness, + { + "input_file_id": "litellm_proxy_unified_id", + "model": "vertex-model", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + harness.is_known_model.return_value = True + harness.pre_call.side_effect = HTTPException(status_code=403, detail="does not have access to the file") + + with ( + patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", loadbalancing), + patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value="unified-xyz"), + patch.object(endpoints, "get_models_from_unified_file_id", return_value=["vertex-model"]), + pytest.raises(ProxyException) as exc, + ): + await call_create(harness) + + assert exc.value.code == "403", exc.value + harness.router_acreate.assert_not_called() + harness.litellm_acreate.assert_not_called() + + def install_managed_files_hook(harness: Harness) -> AsyncMock: prisma_client = AsyncMock() managed_files = PROXY_LiteLLMManagedFiles(MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client) From 5761c86af8938061c3eed8a91bccf1316c3a3ec3 Mon Sep 17 00:00:00 2001 From: mrinal Date: Sat, 10 Oct 2026 20:26:57 +0000 Subject: [PATCH 4/5] test(integration): pin managed-file ownership on batch and fine-tuning create Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_managed_file_create_ownership.py | 200 ++++++++++++++++++ 1 file changed, 200 insertions(+) create mode 100644 tests/integration/management/test_managed_file_create_ownership.py diff --git a/tests/integration/management/test_managed_file_create_ownership.py b/tests/integration/management/test_managed_file_create_ownership.py new file mode 100644 index 00000000000..3cce34b892c --- /dev/null +++ b/tests/integration/management/test_managed_file_create_ownership.py @@ -0,0 +1,200 @@ +from __future__ import annotations + +import base64 +import json +import uuid +from dataclasses import dataclass +from typing import Final + +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 pydantic import JsonValue + +_FILE: Final = { + "id": "file-in-$REQUEST_ID", + "object": "file", + "purpose": "batch", + "bytes": 100, + "created_at": 1, + "filename": "input.jsonl", + "status": "processed", +} +_BATCH: Final = { + "id": "batch-$REQUEST_ID", + "object": "batch", + "endpoint": "/v1/chat/completions", + "errors": None, + "input_file_id": "file-in-$REQUEST_ID", + "completion_window": "24h", + "status": "validating", + "output_file_id": None, + "error_file_id": None, + "created_at": 1, + "expires_at": 1, + "request_counts": {"total": 1, "completed": 0, "failed": 0}, + "metadata": None, +} +_JOB: Final = { + "id": "ftjob-$REQUEST_ID", + "object": "fine_tuning.job", + "model": "gpt-4o-mini-2024-07-18", + "created_at": 1, + "fine_tuned_model": None, + "organization_id": "org-integration", + "result_files": [], + "status": "validating_files", + "validation_file": None, + "training_file": "file-in-$REQUEST_ID", + "hyperparameters": {"n_epochs": "auto"}, + "seed": 1, +} + + +@dataclass(frozen=True, slots=True) +class _Tenants: + model: str + handle: ScenarioHandle + owner_key: str + teammate_key: str + outsider_key: str + owner_file_id: str + + +def _routes() -> RoutedResponse: + return RoutedResponse( + content_type="application/x-routed", + routes={ + "POST /files": JsonResponse(content_type="application/json", body=_FILE), + "POST /batches": JsonResponse(content_type="application/json", body=_BATCH), + "POST /fine_tuning/jobs": JsonResponse(content_type="application/json", body=_JOB), + }, + ) + + +def _upload(scenario: Scenario, key: str, model: str, purpose: str) -> str: + line: Final = json.dumps( + { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": model, "messages": [{"role": "user", "content": "ping"}]}, + } + ).encode() + uploaded: Final = scenario.gateway.request_multipart( + "/v1/files", + {"purpose": purpose, "target_model_names": model}, + {"file": ("input.jsonl", line + b"\n", "application/jsonl")}, + key=key, + ) + assert uploaded.status_code == 200, uploaded.text + return string_value(JSON_OBJECT.validate_json(uploaded.content)["id"]) + + +def _tenants(scenario: Scenario, purpose: str) -> _Tenants: + handle: Final = register_scenario(f"managed-file-create-ownership-{uuid.uuid4().hex}", _routes()) + scenario.cleanups.callback(delete_scenario, handle) + model: Final = scenario.model(api_base=handle.api_base()) + team_a: Final = scenario.team(models=[model]) + team_b: Final = scenario.team(models=[model]) + owner_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_a, models=[model]) + teammate_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_a, models=[model]) + outsider_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_b, models=[model]) + return _Tenants(model, handle, owner_key, teammate_key, outsider_key, _upload(scenario, owner_key, model, purpose)) + + +def _upstream_creates(gateway: Gateway, handle: ScenarioHandle, path: 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") == "POST" + and handle.scenario_id in str(request.get("path")) + and str(request.get("path")).endswith(path) + ) + + +def _create_batch(gateway: Gateway, key: str, model: str, input_file_id: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/batches", + { + "input_file_id": input_file_id, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": model, + }, + key=key, + ) + + +def _create_job(gateway: Gateway, key: str, model: str, files: dict[str, JsonValue]) -> httpx.Response: + return gateway.request("POST", "/v1/fine_tuning/jobs", {"model": model, **files}, key=key) + + +def test_batch_create_from_another_teams_managed_file_is_403_and_never_reaches_the_provider(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + tenants: Final = _tenants(scenario, "batch") + denied: Final = _create_batch(gateway, tenants.outsider_key, tenants.model, tenants.owner_file_id) + assert denied.status_code == 403, denied.text + assert tenants.owner_file_id in denied.text, denied.text + assert _upstream_creates(gateway, tenants.handle, "/batches") == 0 + + +def test_batch_create_from_own_or_teammates_managed_file_still_dispatches(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + tenants: Final = _tenants(scenario, "batch") + owner: Final = _create_batch(gateway, tenants.owner_key, tenants.model, tenants.owner_file_id) + teammate: Final = _create_batch(gateway, tenants.teammate_key, tenants.model, tenants.owner_file_id) + admin: Final = _create_batch(gateway, gateway.key, tenants.model, tenants.owner_file_id) + assert (owner.status_code, teammate.status_code, admin.status_code) == (200, 200, 200), ( + owner.text, + teammate.text, + admin.text, + ) + assert JSON_OBJECT.validate_json(owner.content)["input_file_id"] == tenants.owner_file_id, owner.text + assert _upstream_creates(gateway, tenants.handle, "/batches") == 3 + + +def test_batch_create_from_a_managed_file_id_with_no_owner_row_is_404(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + tenants: Final = _tenants(scenario, "batch") + forged_raw: Final = ( + f"litellm_proxy:application/octet-stream;unified_id,{uuid.uuid4()};target_model_names,{tenants.model};" + "llm_output_file_id,file-forged;llm_output_file_model_id,model-forged" + ) + forged: Final = base64.urlsafe_b64encode(forged_raw.encode()).decode().rstrip("=") + missing: Final = _create_batch(gateway, tenants.outsider_key, tenants.model, forged) + assert missing.status_code == 404, missing.text + assert _upstream_creates(gateway, tenants.handle, "/batches") == 0 + + +@pytest.mark.parametrize("field", ["training_file", "validation_file"]) +def test_fine_tuning_create_from_another_teams_managed_file_is_403(gateway: Gateway, field: str) -> None: + with gateway.scenario() as scenario: + tenants: Final = _tenants(scenario, "fine-tune") + own_training: Final = _upload(scenario, tenants.outsider_key, tenants.model, "fine-tune") + files: Final[dict[str, JsonValue]] = {"training_file": own_training, field: tenants.owner_file_id} + denied: Final = _create_job(gateway, tenants.outsider_key, tenants.model, files) + assert denied.status_code == 403, denied.text + assert _upstream_creates(gateway, tenants.handle, "/fine_tuning/jobs") == 0 + + +def test_fine_tuning_create_from_own_managed_files_still_dispatches(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + tenants: Final = _tenants(scenario, "fine-tune") + validation: Final = _upload(scenario, tenants.teammate_key, tenants.model, "fine-tune") + created: Final = _create_job( + gateway, + tenants.owner_key, + tenants.model, + {"training_file": tenants.owner_file_id, "validation_file": validation}, + ) + assert created.status_code == 200, created.text + assert _upstream_creates(gateway, tenants.handle, "/fine_tuning/jobs") == 1 From 51d8cbbfd013854ba1ab0efaf3b39c46ca8745a1 Mon Sep 17 00:00:00 2001 From: mrinal Date: Sat, 10 Oct 2026 22:47:28 +0000 Subject: [PATCH 5/5] 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 | 20 +++++ .../key_management_endpoints.py | 8 ++ .../managed_resource_owner_sync.py | 44 +++++++++++ .../test_managed_file_create_ownership.py | 73 ++++++++++++++++++- .../proxy/hooks/test_batch_rate_limiter.py | 64 +++++++++++----- 5 files changed, 191 insertions(+), 18 deletions(-) create mode 100644 litellm/proxy/management_helpers/managed_resource_owner_sync.py diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 17f8bd277dd..5f214816578 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/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( diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 6e226868918..a6e9bb7b3d5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/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, diff --git a/litellm/proxy/management_helpers/managed_resource_owner_sync.py b/litellm/proxy/management_helpers/managed_resource_owner_sync.py new file mode 100644 index 00000000000..4171205b52b --- /dev/null +++ b/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:`` (``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}", + ) diff --git a/tests/integration/management/test_managed_file_create_ownership.py b/tests/integration/management/test_managed_file_create_ownership.py index 3cce34b892c..a14f1fced19 100644 --- a/tests/integration/management/test_managed_file_create_ownership.py +++ b/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, + ) diff --git a/tests/unit/proxy/hooks/test_batch_rate_limiter.py b/tests/unit/proxy/hooks/test_batch_rate_limiter.py index 2f71b177705..b4e972ea34f 100644 --- a/tests/unit/proxy/hooks/test_batch_rate_limiter.py +++ b/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