diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index a91b29002e3..b360ce40db1 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -27,6 +27,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( decode_model_from_file_id, encode_batch_response_ids, encode_file_id_with_model, + ensure_batch_response_managed_file_ids, get_batch_id_from_unified_batch_id, get_batch_from_database, get_credentials_for_model, @@ -440,10 +441,19 @@ async def retrieve_batch( ) # The DB may store raw provider file IDs (before hooks translate them). - # Resolve any raw input/output/error file IDs to unified IDs. + # Register any missing managed-file rows and rewrite raw IDs to unified + # IDs so terminal rows written by the poller don't leak raw provider + # file IDs that bypass the /v1/files ownership check. if unified_batch_id: - await resolve_input_file_id_to_unified(response, prisma_client) - await resolve_output_file_ids_to_unified(response, prisma_client) + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=managed_files_obj, + prisma_client=prisma_client, + verbose_proxy_logger=verbose_proxy_logger, + user_api_key_dict=user_api_key_dict, + db_batch_object=db_batch_object, + unified_batch_id=unified_batch_id, + ) asyncio.create_task( proxy_logging_obj.update_request_status( diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 2960b031cd3..c514857e0a6 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -933,6 +933,7 @@ async def ensure_batch_response_managed_file_ids( verbose_proxy_logger, user_api_key_dict=None, db_batch_object=None, + unified_batch_id: Optional[str] = None, ) -> None: """Normalize batch file IDs to managed unified IDs before DB persistence.""" await resolve_input_file_id_to_unified(response, prisma_client) @@ -943,6 +944,8 @@ async def ensure_batch_response_managed_file_ids( hidden_params = getattr(response, "_hidden_params", None) or {} model_id = hidden_params.get("model_id") + if not model_id and unified_batch_id: + model_id = get_model_id_from_unified_batch_id(unified_batch_id) if not model_id: return diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py index d8669960674..cc2a6c6cb21 100644 --- a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py +++ b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py @@ -200,6 +200,40 @@ async def test_ensure_batch_response_resolves_model_name_from_unified_file_id(): ) +@pytest.mark.asyncio +async def test_ensure_batch_response_derives_model_id_from_unified_batch_id(): + """Terminal DB rows have empty _hidden_params (rehydrated from stored JSON). + + Regression for GH #33989: the helper must still register the raw output/error + ids as managed files by deriving model_id from the unified batch id, otherwise + the terminal-retrieve short-circuit leaks raw provider file ids. + """ + unified_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ=" + response = _build_batch_response( + output_file_id="file-raw-output", + error_file_id="file-raw-error", + hidden_params={}, + ) + mock_managed_files = _build_managed_files_mock(unified_id=unified_id) + + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=mock_managed_files, + prisma_client=_build_prisma_mock(), + verbose_proxy_logger=MagicMock(), + db_batch_object=SimpleNamespace(created_by="user-from-db", team_id=None, status="failed"), + unified_batch_id="litellm_proxy;model_id:deployment-42;llm_batch_id:batch_abc", + ) + + assert response.output_file_id == unified_id + assert response.error_file_id == unified_id + assert mock_managed_files.get_unified_output_file_id.call_count == 2 + assert ( + mock_managed_files.get_unified_output_file_id.call_args.kwargs["model_id"] + == "deployment-42" + ) + + @pytest.mark.asyncio async def test_ensure_batch_response_returns_early_without_managed_files_obj(): """Without managed_files_obj, the helper is a no-op (no conversion attempted).""" diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 6a185988c9b..6884712a440 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -937,6 +937,7 @@ class RetrieveHarness: update_batch_in_db: AsyncMock resolve_input: AsyncMock resolve_output: AsyncMock + ensure_managed: AsyncMock @property def router_aretrieve(self) -> AsyncMock: @@ -975,6 +976,7 @@ def retrieve_harness(): update_batch_in_db = AsyncMock(return_value=None) resolve_input = AsyncMock(return_value=None) resolve_output = AsyncMock(return_value=None) + ensure_managed = AsyncMock(return_value=None) with ExitStack() as stack: stack.enter_context( @@ -1003,6 +1005,7 @@ def retrieve_harness(): stack.enter_context(patch.object(endpoints, "update_batch_in_database", update_batch_in_db)) stack.enter_context(patch.object(endpoints, "resolve_input_file_id_to_unified", resolve_input)) stack.enter_context(patch.object(endpoints, "resolve_output_file_ids_to_unified", resolve_output)) + stack.enter_context(patch.object(endpoints, "ensure_batch_response_managed_file_ids", ensure_managed)) stack.enter_context(patch.object(litellm, "aretrieve_batch", litellm_aretrieve)) stack.enter_context(patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", False)) stack.enter_context(patch.object(proxy_server, "llm_router", router)) @@ -1026,6 +1029,7 @@ def retrieve_harness(): update_batch_in_db=update_batch_in_db, resolve_input=resolve_input, resolve_output=resolve_output, + ensure_managed=ensure_managed, ) @@ -1260,15 +1264,20 @@ async def test_retrieve__db_terminal_state_short_circuits(retrieve_harness, stat @pytest.mark.asyncio async def test_retrieve__db_terminal_unified_resolves_file_ids(retrieve_harness): + db_batch_object = MagicMock() db_response = make_batch(id="batch-from-db", status="completed") - retrieve_harness.get_batch_from_db.return_value = (MagicMock(), db_response) + retrieve_harness.get_batch_from_db.return_value = (db_batch_object, db_response) with patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID): await call_retrieve(retrieve_harness, "batch-unified-blob") - # Terminal short-circuit still resolves raw provider file ids to unified. - retrieve_harness.resolve_input.assert_called_once() - retrieve_harness.resolve_output.assert_called_once() + # Regression GH #33989: the terminal short-circuit must register missing + # managed-file rows (owned by the batch creator) and rewrite raw provider + # file ids, not merely look them up, so raw ids never leak past the ACL. + retrieve_harness.ensure_managed.assert_called_once() + ensure_kwargs = retrieve_harness.ensure_managed.call_args.kwargs + assert ensure_kwargs["unified_batch_id"] == UNIFIED_BATCH_ID + assert ensure_kwargs["db_batch_object"] is db_batch_object retrieve_harness.litellm_aretrieve.assert_not_called() retrieve_harness.router_aretrieve.assert_not_called()