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(