From 5433ae7e8caf725191bca6a579596b404f4acf72 Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Sat, 14 Feb 2026 00:30:35 -0500 Subject: [PATCH] Fix: bypass managed files access check in batch polling by calling afile_content directly --- .../proxy/common_utils/check_batch_cost.py | 39 ++-- .../proxy/test_managed_files_access_check.py | 200 ++++++++++++++++++ 2 files changed, 220 insertions(+), 19 deletions(-) create mode 100644 tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index bb25e4f062..807d237df5 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -4,7 +4,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t from litellm._uuid import uuid from datetime import datetime -from typing import TYPE_CHECKING, Optional, cast +from typing import TYPE_CHECKING, Optional from litellm._logging import verbose_proxy_logger @@ -35,14 +35,11 @@ class CheckBatchCost: - if not, return False - if so, return True """ - from litellm_enterprise.proxy.hooks.managed_files import ( - _PROXY_LiteLLMManagedFiles, - ) - from litellm.batches.batch_utils import ( _get_file_content_as_dictionary, calculate_batch_cost_and_usage, ) + from litellm.files.main import afile_content from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.proxy.openai_files_endpoints.common_utils import ( @@ -102,27 +99,31 @@ class CheckBatchCost: continue ## RETRIEVE THE BATCH JOB OUTPUT FILE - managed_files_obj = cast( - Optional[_PROXY_LiteLLMManagedFiles], - self.proxy_logging_obj.get_proxy_hook("managed_files"), - ) if ( response.status == "completed" and response.output_file_id is not None - and managed_files_obj is not None ): verbose_proxy_logger.info( f"Batch ID: {batch_id} is complete, tracking cost and usage" ) - # track cost - model_file_id_mapping = { - response.output_file_id: {model_id: response.output_file_id} - } - _file_content = await managed_files_obj.afile_content( - file_id=response.output_file_id, - litellm_parent_otel_span=None, - llm_router=self.llm_router, - model_file_id_mapping=model_file_id_mapping, + + # Extract raw provider file ID from the unified file ID + # (async_post_call_success_hook may have replaced output_file_id with a unified ID) + raw_output_file_id = response.output_file_id + decoded = _is_base64_encoded_unified_file_id(raw_output_file_id) + if decoded: + try: + raw_output_file_id = decoded.split("llm_output_file_id,")[1].split(";")[0] + except (IndexError, AttributeError): + pass + + # Call litellm.afile_content directly with deployment credentials, + # bypassing the managed files access-control hooks that would + # reject this background job's default_user_id identity + credentials = self.llm_router.get_deployment_credentials_with_provider(model_id) or {} + _file_content = await afile_content( + file_id=raw_output_file_id, + **credentials, ) file_content_as_dict = _get_file_content_as_dictionary( diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py new file mode 100644 index 0000000000..2db5a2214c --- /dev/null +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py @@ -0,0 +1,200 @@ +""" +Tests for managed files access control in batch polling context. + +Regression test for: batch polling job running as default_user_id gets 403 +when trying to access managed files created by a real user. + +The fix (Option C) makes check_batch_cost call litellm.afile_content directly +with deployment credentials, bypassing the managed files access-control hooks. +""" + +import base64 +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from fastapi import HTTPException + +from litellm.proxy._types import UserAPIKeyAuth + + +def _make_user_api_key_dict(user_id: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-test", + user_id=user_id, + parent_otel_span=None, + ) + + +def _make_unified_file_id() -> str: + """Create a base64-encoded unified file ID that passes _is_base64_encoded_unified_file_id.""" + raw = "litellm_proxy:application/octet-stream;unified_id,test-123;target_model_names,azure-gpt-4" + return base64.b64encode(raw.encode()).decode() + + +def _make_managed_files_instance(file_created_by: str, unified_file_id: str): + """Create a _PROXY_LiteLLMManagedFiles with a mocked DB that returns a file owned by file_created_by.""" + from litellm_enterprise.proxy.hooks.managed_files import ( + _PROXY_LiteLLMManagedFiles, + ) + + mock_db_record = MagicMock() + mock_db_record.created_by = file_created_by + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock( + return_value=mock_db_record + ) + + instance = _PROXY_LiteLLMManagedFiles( + internal_usage_cache=MagicMock(), + prisma_client=mock_prisma, + ) + return instance + + +# --- Access control unit tests (document existing behavior) --- + + +@pytest.mark.asyncio +async def test_should_allow_file_owner_access(): + """File owner can access their own file — baseline sanity check.""" + unified_file_id = _make_unified_file_id() + managed_files = _make_managed_files_instance( + file_created_by="user-A", + unified_file_id=unified_file_id, + ) + user = _make_user_api_key_dict("user-A") + data = {"file_id": unified_file_id} + + result = await managed_files.check_managed_file_id_access(data, user) + assert result is True + + +@pytest.mark.asyncio +async def test_should_block_different_user_access(): + """A different regular user cannot access another user's file — correct behavior.""" + unified_file_id = _make_unified_file_id() + managed_files = _make_managed_files_instance( + file_created_by="user-A", + unified_file_id=unified_file_id, + ) + user = _make_user_api_key_dict("user-B") + data = {"file_id": unified_file_id} + + with pytest.raises(HTTPException) as exc_info: + await managed_files.check_managed_file_id_access(data, user) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_should_block_default_user_id_access(): + """ + default_user_id is correctly blocked by the access check. + This documents the existing behavior that the Option C fix works around. + """ + unified_file_id = _make_unified_file_id() + managed_files = _make_managed_files_instance( + file_created_by="user-A", + unified_file_id=unified_file_id, + ) + system_user = _make_user_api_key_dict("default_user_id") + data = {"file_id": unified_file_id} + + with pytest.raises(HTTPException) as exc_info: + await managed_files.check_managed_file_id_access(data, system_user) + assert exc_info.value.status_code == 403 + + +# --- Option C fix test: check_batch_cost bypasses managed files hook --- + + +@pytest.mark.asyncio +async def test_check_batch_cost_should_call_afile_content_directly_with_credentials(): + """ + check_batch_cost should call litellm.afile_content directly with deployment + credentials, bypassing managed_files_obj.afile_content and its access-control + hooks. This avoids the 403 that occurs when the background job runs as + default_user_id. + """ + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + # Build a unified object ID in the expected format: + # litellm_proxy;model_id:{};llm_batch_id:{};llm_output_file_id:{} + unified_raw = "litellm_proxy;model_id:model-deploy-xyz;llm_batch_id:batch-123;llm_output_file_id:file-raw-output" + unified_object_id = base64.b64encode(unified_raw.encode()).decode() + + # Mock a pending job from the DB + mock_job = MagicMock() + mock_job.unified_object_id = unified_object_id + mock_job.created_by = "user-A" + mock_job.id = "job-1" + + # Mock prisma + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock() + + # Mock proxy_logging_obj — should NOT be called for file content + mock_proxy_logging = MagicMock() + mock_managed_files_hook = MagicMock() + mock_managed_files_hook.afile_content = AsyncMock() + mock_proxy_logging.get_proxy_hook = MagicMock(return_value=mock_managed_files_hook) + + # Mock the batch response (completed, with output file) + from litellm.types.utils import LiteLLMBatch + batch_response = LiteLLMBatch( + id="batch-123", + completion_window="24h", + created_at=1700000000, + endpoint="/v1/chat/completions", + input_file_id="file-input", + object="batch", + status="completed", + output_file_id="file-raw-output", + ) + + # Mock router + mock_router = MagicMock() + mock_router.aretrieve_batch = AsyncMock(return_value=batch_response) + mock_router.get_deployment_credentials_with_provider = MagicMock( + return_value={ + "api_key": "test-key", + "api_base": "https://test.azure.com/", + "custom_llm_provider": "azure", + } + ) + + mock_deployment = MagicMock() + mock_deployment.litellm_params.custom_llm_provider = "azure" + mock_deployment.litellm_params.model = "azure/gpt-4" + mock_router.get_deployment = MagicMock(return_value=mock_deployment) + + checker = CheckBatchCost( + proxy_logging_obj=mock_proxy_logging, + prisma_client=mock_prisma, + llm_router=mock_router, + ) + + mock_file_content = MagicMock() + mock_file_content.content = b'{"id":"req-1","response":{"status_code":200,"body":{"id":"cmpl-1","object":"chat.completion","created":1700000000,"model":"gpt-4","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}}}\n' + + with patch( + "litellm.files.main.afile_content", + new_callable=AsyncMock, + return_value=mock_file_content, + ) as mock_direct_afile_content: + await checker.check_batch_cost() + + # afile_content should be called directly (not through managed_files_obj) + mock_direct_afile_content.assert_called_once() + call_kwargs = mock_direct_afile_content.call_args.kwargs + + assert call_kwargs.get("api_key") == "test-key", ( + f"afile_content should receive api_key from deployment credentials. " + f"Got: {call_kwargs}" + ) + + # managed_files_obj.afile_content should NOT have been called + mock_managed_files_hook.afile_content.assert_not_called()