From 19184694f59eb1934f3d550cae932d9f432f82a0 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:09:16 +0000 Subject: [PATCH 1/2] fix(batches): mark terminal batch with no output file as processed in CheckBatchCost A managed batch whose request lines all failed can reach a terminal provider status (completed) with output_file_id=None and only an error_file_id. Such a row matched neither the completed-with-output billing branch nor the failed/expired/cancelled branch, so batch_processed stayed False and the poller re-selected it on every cycle for the lifetime of the deployment; output/error file deletion is also gated on batch_processed, so those files could never be deleted. Broaden the terminal handling so a completed/complete/expired batch with an output file is billed, and any terminal batch with nothing to bill (failed/cancelled, or completed/expired with no output) is marked terminal exactly once. Non-terminal statuses (validating/in_progress) are still left for the next poll, and an expired batch that did produce output is now billed. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 10 +- .../proxy_unit_tests/test_check_batch_cost.py | 252 +++++++++++++++++- 2 files changed, 254 insertions(+), 8 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 dc8f17fb665f..00cc184a515c 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -671,7 +671,7 @@ async def check_batch_cost(self): ## RETRIEVE THE BATCH JOB OUTPUT FILE if ( - response.status == "completed" + response.status in ("completed", "complete", "expired") and response.output_file_id is not None ): try: @@ -712,7 +712,13 @@ async def check_batch_cost(self): f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}" ) - elif response.status in ("failed", "expired", "cancelled"): + elif response.status in ( + "completed", + "complete", + "failed", + "expired", + "cancelled", + ): try: from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 6d7ada17ec52..7616a1d5ddc3 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -623,9 +623,9 @@ async def test_terminal_status_marks_job_processed( mock_llm_router, terminal_status, ): - """When the provider reports a terminal status (failed/expired/cancelled), the row - must be written back with that status and batch_processed=True so it stops being - polled forever. + """When the provider reports a terminal status with nothing to bill + (failed/cancelled, or expired with no output file), the row must be written back + with that status and batch_processed=True so it stops being polled forever. """ import base64 @@ -651,6 +651,7 @@ async def test_terminal_status_marks_job_processed( mock_response = MagicMock() mock_response.status = terminal_status + mock_response.output_file_id = None mock_response.model_dump_json.return_value = ( f'{{"id":"batch-1","status":"{terminal_status}"}}' ) @@ -671,7 +672,7 @@ async def test_terminal_status_marks_job_processed( ), "terminal-status update() must set batch_processed=True so polling stops" @pytest.mark.asyncio - @pytest.mark.parametrize("terminal_status", ["failed", "expired", "cancelled"]) + @pytest.mark.parametrize("terminal_status", ["failed", "cancelled"]) async def test_terminal_status_persists_managed_output_file_ids( self, check_batch_cost_instance, @@ -679,10 +680,12 @@ async def test_terminal_status_persists_managed_output_file_ids( 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 + """A cancelled/failed 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. + (Expired with an output file is billed through the completed path instead, + covered by test_expired_with_output_file_is_billed.) """ import base64 import json @@ -797,6 +800,243 @@ def find_managed_file(where): 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 + @pytest.mark.parametrize("completed_status", ["completed", "complete"]) + async def test_completed_without_output_file_marked_processed_without_billing( + self, + check_batch_cost_instance, + mock_prisma_client, + mock_llm_router, + completed_status, + ): + """#35354 regression: a terminal completed batch whose request lines all failed + reaches `completed` with output_file_id=None (only an error_file_id). + + Pre-fix it matched neither the completed-with-output branch nor the + failed/expired/cancelled branch, so batch_processed stayed False and the row + was re-selected on every poll cycle forever. It must now be marked terminal + exactly once, without being billed (no output means nothing to bill). + """ + import base64 + from unittest.mock import patch + + 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 + ) + + mock_job = MagicMock() + mock_job.id = "job-completed-no-output-1" + 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 + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + + mock_response = MagicMock() + mock_response.status = completed_status + mock_response.output_file_id = None + mock_response.error_file_id = "file-error-123" + mock_response.model_dump_json.return_value = ( + f'{{"id":"batch-1","status":"{completed_status}"}}' + ) + + mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) + # Billing reads credentials off the router; if it is touched we billed a batch + # that has no output, which is the behaviour this test guards against. + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) + + with patch( + "litellm.files.main.afile_content", + new_callable=AsyncMock, + ) as mock_afile_content: + await check_batch_cost_instance.check_batch_cost() + + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "a completed batch with no output file must be marked processed exactly once" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] + assert update_data["status"] == completed_status + assert ( + update_data["batch_processed"] is True + ), "completed-without-output update() must set batch_processed=True so polling stops" + assert ( + mock_afile_content.await_count == 0 + ), "a batch with no output file must not be billed" + assert ( + mock_llm_router.get_deployment_credentials_with_provider.call_count == 0 + ), "a batch with no output file must not enter the cost-tracking path" + + @pytest.mark.asyncio + async def test_non_terminal_status_left_unprocessed( + self, check_batch_cost_instance, mock_prisma_client, mock_llm_router + ): + """A batch still validating/in_progress must NOT be treated as terminal: no DB + write, so it keeps being polled until it actually reaches a terminal status. + """ + from unittest.mock import patch + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + + mock_job = MagicMock() + mock_job.id = "job-in-progress-1" + mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" + mock_job.created_by = "user-1" + + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + + mock_response = MagicMock() + mock_response.status = "in_progress" + mock_response.output_file_id = None + + 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() + + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 + ), "a non-terminal batch must not be written back (would stop polling prematurely)" + + @pytest.mark.asyncio + async def test_expired_with_output_file_is_billed( + self, check_batch_cost_instance, mock_prisma_client, mock_llm_router + ): + """An expired batch that still produced an output file served real request lines, + so it must be billed (cost tracked) and then marked processed, not silently + marked terminal without billing. + """ + from unittest.mock import patch + + 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 + ) + + mock_job = MagicMock() + mock_job.id = "job-expired-with-output-1" + mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" + mock_job.created_by = "user-1" + + assert check_batch_cost_instance._has_batch_processed_column is True + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + + mock_response = MagicMock() + mock_response.status = "expired" + mock_response.output_file_id = "file-output-123" + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"expired"}' + ) + + mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) + + mock_deployment = MagicMock() + mock_deployment.litellm_params.custom_llm_provider = "openai" + mock_deployment.litellm_params.model = "gpt-4" + mock_deployment.model_info.model_dump.return_value = {} + mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment) + + mock_file_content = MagicMock() + mock_file_content.content = b'{"id":"req-1"}' + + 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", + ), + patch( + "litellm.files.main.afile_content", + new_callable=AsyncMock, + return_value=mock_file_content, + ) as mock_afile_content, + patch( + "litellm.batches.batch_utils._get_file_content_as_dictionary", + return_value=[{"id": "req-1"}], + ), + patch( + "litellm.batches.batch_utils.calculate_batch_cost_and_usage", + new_callable=AsyncMock, + return_value=( + 0.01, + {"prompt_tokens": 10, "completion_tokens": 5}, + ["gpt-4"], + ), + ), + patch( + "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", + return_value=("gpt-4", "openai", None, None), + ), + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, + ): + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + mock_logging_cls.return_value = mock_logging_obj + + await check_batch_cost_instance.check_batch_cost() + + assert ( + mock_afile_content.await_count == 1 + ), "expired batch with an output file must fetch results and be billed" + mock_logging_obj.async_success_handler.assert_awaited_once() + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] + assert update_data["batch_processed"] is True + @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 From eacea13a257d934627714b554dd1fd2c2b44b261 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 13 Aug 2026 21:24:54 -0700 Subject: [PATCH 2/2] fix(batches): persist real terminal status when billing expired batches --- .../litellm_enterprise/proxy/common_utils/check_batch_cost.py | 2 +- tests/proxy_unit_tests/test_check_batch_cost.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) 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 00cc184a515c..38266f6c3ea4 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -698,7 +698,7 @@ async def check_batch_cost(self): # mark the job as complete try: update_data: dict = { - "status": "complete", + "status": response.status if response.status != "completed" else "complete", "file_object": response.model_dump_json(), } if self._has_batch_processed_column: diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 7616a1d5ddc3..ca9d5f7f7d46 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1036,6 +1036,9 @@ async def test_expired_with_output_file_is_billed( 1 ]["data"] assert update_data["batch_processed"] is True + assert ( + update_data["status"] == "expired" + ), "billed expired batch must keep its real terminal status in the DB" @pytest.mark.asyncio async def test_raw_output_file_id_converted_to_managed_id(