fix(batches): mark terminal batch with no output file as processed in CheckBatchCost

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:
@@ -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:
@@ -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,

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,18 +672,20 @@ 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,
         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
+        """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,246 @@ 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
+        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(
         self, check_batch_cost_instance, mock_prisma_client, mock_llm_router