From 5339ec50e788a6bb9090380d1e43060e13bcdb97 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:42:20 -0700 Subject: [PATCH] fix(batches): persist managed file ids for cancelled/failed/expired batches When the batch cost poller found a batch in a terminal failed, expired, or cancelled state it wrote the provider response straight to the managed object table, so the stored blob kept raw provider file ids and a raw batch id. Since the row is final after batch_processed=True and the read paths only resolve existing managed ids, every later GET /batches/{id} and GET /batches leaked raw provider output and error file ids that clients cannot fetch through the proxy. The terminal branch now normalizes the response with ensure_batch_response_managed_file_ids before persisting, minting managed ids under the batch owner's identity POST /batches/{id}/cancel had the same gap: it called update_batch_in_database without the caller's auth context, so a cancel response that already carried provider file ids could never mint managed ids. The endpoint now forwards user_api_key_dict --- .../proxy/common_utils/check_batch_cost.py | 14 ++ litellm/proxy/batches_endpoints/endpoints.py | 1 + .../proxy_unit_tests/test_check_batch_cost.py | 151 +++++++++++++++--- .../proxy/batches_endpoints/test_endpoints.py | 11 ++ 4 files changed, 158 insertions(+), 19 deletions(-) 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 3ed63b0d9ee..a12b0bb7170 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -656,6 +656,20 @@ class CheckBatchCost: elif response.status in ("failed", "expired", "cancelled"): try: + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + ensure_batch_response_managed_file_ids, + ) + + response.id = job.unified_object_id + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=self.proxy_logging_obj.get_proxy_hook("managed_files"), + prisma_client=self.prisma_client, + verbose_proxy_logger=verbose_proxy_logger, + db_batch_object=job, + unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id), + ) update_data = { "status": response.status, "file_object": response.model_dump_json(), diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index deed665d3f2..7b5af6c1068 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -961,6 +961,7 @@ async def cancel_batch( prisma_client=prisma_client, verbose_proxy_logger=verbose_proxy_logger, operation="cancel", + user_api_key_dict=user_api_key_dict, ) ### CALL HOOKS ### - modify outgoing data diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index a15abd023d8..d6c5d8fc809 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -499,7 +499,7 @@ class TestCheckBatchCost: must be written back with that status and batch_processed=True so it stops being polled forever. """ - from unittest.mock import patch + import base64 mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( return_value=0 @@ -511,7 +511,9 @@ class TestCheckBatchCost: mock_job = MagicMock() mock_job.id = "job-terminal-1" - mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" + mock_job.unified_object_id = base64.urlsafe_b64encode( + b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456" + ).decode() mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True @@ -527,23 +529,7 @@ class TestCheckBatchCost: mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;" - - with ( - patch( - "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id", - side_effect=[decoded_id, None], - ), - patch( - "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id", - return_value="model-123", - ), - patch( - "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id", - return_value="batch-456", - ), - ): - await check_batch_cost_instance.check_batch_cost() + await check_batch_cost_instance.check_batch_cost() assert ( mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 @@ -556,6 +542,133 @@ class TestCheckBatchCost: update_data["batch_processed"] is True ), "terminal-status update() must set batch_processed=True so polling stops" + @pytest.mark.asyncio + @pytest.mark.parametrize("terminal_status", ["failed", "expired", "cancelled"]) + async def test_terminal_status_persists_managed_output_file_ids( + self, + check_batch_cost_instance, + mock_prisma_client, + mock_llm_router, + terminal_status, + ): + """A cancelled/failed/expired batch with provider output files must be persisted + with unified managed file IDs, never raw provider IDs. Raw IDs written here leak + to every later GET /batches/{id} and GET /batches because the terminal row is + final (batch_processed=True) and read paths only resolve, never mint. + """ + import base64 + import json + + from litellm.types.utils import LiteLLMBatch + + unified_batch_uid = base64.urlsafe_b64encode( + b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456" + ).decode() + raw_output_file_id = "file-terminal-out-abc" + raw_error_file_id = "file-terminal-err-xyz" + raw_input_file_id = "file-terminal-in-123" + unified_input_file_id = base64.urlsafe_b64encode( + b"litellm_proxy:application/octet-stream;unified_id,in-1;target_model_names,gpt-5-batch" + ).decode() + unified_output_file_id = base64.urlsafe_b64encode( + f"litellm_proxy:application/octet-stream;unified_id,u-1;llm_output_file_id,{raw_output_file_id}".encode() + ).decode() + unified_error_file_id = base64.urlsafe_b64encode( + f"litellm_proxy:application/octet-stream;unified_id,u-2;llm_output_file_id,{raw_error_file_id}".encode() + ).decode() + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) + + input_file_row = MagicMock() + input_file_row.unified_file_id = unified_input_file_id + + def find_managed_file(where): + if where["flat_model_file_ids"]["has"] == raw_input_file_id: + return input_file_row + return None + + mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( + side_effect=find_managed_file + ) + + mock_job = MagicMock() + mock_job.id = "job-terminal-mint-1" + mock_job.unified_object_id = unified_batch_uid + mock_job.created_by = "user-1" + mock_job.team_id = "team-1" + + check_batch_cost_instance._has_batch_processed_column = True + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + + response = LiteLLMBatch( + id="batch-456", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id=raw_input_file_id, + object="batch", + status=terminal_status, + output_file_id=raw_output_file_id, + error_file_id=raw_error_file_id, + ) + mock_llm_router.aretrieve_batch = AsyncMock(return_value=response) + + mock_hook = MagicMock() + mock_hook.get_unified_output_file_id.side_effect = [ + unified_output_file_id, + unified_error_file_id, + ] + mock_hook.store_unified_file_id = AsyncMock() + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( + mock_hook + ) + + await check_batch_cost_instance.check_batch_cost() + + mock_hook.get_unified_output_file_id.assert_any_call( + output_file_id=raw_output_file_id, + model_id="model-123", + model_name="gpt-5-batch", + ) + mock_hook.get_unified_output_file_id.assert_any_call( + output_file_id=raw_error_file_id, + model_id="model-123", + model_name="gpt-5-batch", + ) + stored = { + next(iter(c.kwargs["model_mappings"].values())): c.kwargs["file_id"] + for c in mock_hook.store_unified_file_id.call_args_list + } + assert stored == { + raw_output_file_id: unified_output_file_id, + raw_error_file_id: unified_error_file_id, + } + for store_call in mock_hook.store_unified_file_id.call_args_list: + assert store_call.kwargs["user_api_key_dict"].user_id == "user-1" + assert store_call.kwargs["user_api_key_dict"].team_id == "team-1" + + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + update_call = mock_prisma_client.db.litellm_managedobjecttable.update.call_args + assert update_call.kwargs["where"] == {"id": "job-terminal-mint-1"} + update_data = update_call.kwargs["data"] + assert update_data["status"] == terminal_status + assert update_data["batch_processed"] is True + persisted = json.loads(update_data["file_object"]) + assert persisted["id"] == unified_batch_uid + assert persisted["input_file_id"] == unified_input_file_id + assert persisted["output_file_id"] == unified_output_file_id + assert persisted["error_file_id"] == unified_error_file_id + assert raw_output_file_id not in update_data["file_object"] + assert raw_error_file_id not in update_data["file_object"] + @pytest.mark.asyncio async def test_raw_output_file_id_converted_to_managed_id( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index f2d37fbe842..64e1dcda5c9 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -1896,6 +1896,17 @@ async def test_cancel__unified_batch_id_routes_to_router(cancel_harness): assert cancel_harness.update_batch_in_db.call_args.kwargs["operation"] == "cancel" +@pytest.mark.asyncio +async def test_cancel__db_write_receives_caller_auth(cancel_harness): + """update_batch_in_database can only mint managed IDs for a cancelled batch's + output files when it has an auth context, so cancel must forward the caller's.""" + caller = UserAPIKeyAuth(api_key="sk-test", user_id="user-cancel-1") + with patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID): + await call_cancel(cancel_harness, "batch-unified-blob", user=caller) + + assert cancel_harness.update_batch_in_db.call_args.kwargs["user_api_key_dict"] is caller + + @pytest.mark.asyncio async def test_cancel__unified_missing_model_id_400(cancel_harness): # unified id with no model_id segment -> get_model_id returns None -> 400.