fix(proxy): bill cancelled and failed batches that still produced an output file

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -24,12 +24,16 @@
 
 CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
 
-TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
+PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
     "completed",
     "complete",
     "failed",
     "expired",
     "cancelled",
+)
+
+TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
+    *PROVIDER_TERMINAL_BATCH_STATUSES,
     "stale_expired",
 )
 
@@ -286,6 +290,57 @@ def _batch_deployment_exists(self, model_id: str) -> bool:
         404 must not retire the row; the staleness sweep bounds it instead."""
         return self.llm_router.get_deployment(model_id=model_id) is not None
 
+    @staticmethod
+    def _is_output_file_gone_at_provider(error: Exception, output_file_id: Optional[str]) -> bool:
+        """A 404 naming the output file means there is nothing to fetch on this or any
+        later poll: providers like Vertex AI advertise an output path for every batch,
+        including terminal ones that never wrote it. Any other failure may be
+        transient, so it keeps retrying until the staleness sweep bounds it."""
+        import openai
+
+        from litellm.exceptions import NotFoundError
+
+        if not output_file_id:
+            return False
+        return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
+
+    async def _finalize_unbilled_terminal_job(
+        self, job: "LiteLLM_ManagedObjectTable", response: "LiteLLMBatch"
+    ) -> None:
+        """Persist a terminal batch that has nothing billable, converting any raw
+        provider file ids to managed ids, and take it out of the poll page."""
+        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: Final[dict] = {
+                "status": response.status,
+                "file_object": response.model_dump_json(),
+                **({"batch_processed": True} if self._has_batch_processed_column else {}),
+            }
+            await self.prisma_client.db.litellm_managedobjecttable.update(
+                where={"id": job.id},
+                data=update_data,
+            )
+            verbose_proxy_logger.info(
+                f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
+            )
+        except Exception as db_err:
+            verbose_proxy_logger.error(
+                f"CheckBatchCost: failed to mark job {job.id} as {response.status} in DB: {db_err}"
+            )
+
     @staticmethod
     def _record_error(
         prom_logger: Optional["PrometheusLogger"], error_type: str
@@ -796,7 +851,7 @@ async def check_batch_cost(self):
 
             ## RETRIEVE THE BATCH JOB OUTPUT FILE
             if (
-                response.status in ("completed", "complete", "expired")
+                response.status in PROVIDER_TERMINAL_BATCH_STATUSES
                 and response.output_file_id is not None
             ):
                 try:
@@ -808,6 +863,15 @@ async def check_batch_cost(self):
                         prom_logger=prom_logger,
                     )
                 except Exception as tracking_err:
+                    if self._is_output_file_gone_at_provider(
+                        tracking_err, response.output_file_id
+                    ) and self._batch_deployment_exists(model_id):
+                        verbose_proxy_logger.warning(
+                            f"CheckBatchCost: output file {response.output_file_id} of batch {batch_id} "
+                            f"does not exist at the provider; retiring job {job.id} unbilled"
+                        )
+                        await self._finalize_unbilled_terminal_job(job, response)
+                        continue
                     verbose_proxy_logger.error(
                         f"CheckBatchCost: failed to track cost for batch {batch_id} "
                         f"(job {job.id}); leaving it unprocessed so the next poll retries: {tracking_err}"
@@ -837,45 +901,8 @@ async def check_batch_cost(self):
                         f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
                     )
 
-            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,
-                        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(),
-                    }
-                    if self._has_batch_processed_column:
-                        update_data["batch_processed"] = True
-                    await self.prisma_client.db.litellm_managedobjecttable.update(
-                        where={"id": job.id},
-                        data=update_data,
-                    )
-                    verbose_proxy_logger.info(
-                        f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
-                    )
-                except Exception as db_err:
-                    verbose_proxy_logger.error(
-                        f"CheckBatchCost: failed to mark job {job.id} as {response.status} in DB: {db_err}"
-                    )
+            elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES:
+                await self._finalize_unbilled_terminal_job(job, response)
 
         # Record polling run metrics (always, even if nothing was processed)
         if prom_logger:

tests/proxy_unit_tests/test_check_batch_cost.py

@@ -817,12 +817,12 @@ async def test_terminal_status_persists_managed_output_file_ids(
         mock_llm_router,
         terminal_status,
     ):
-        """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.)
+        """A cancelled/failed batch with a provider error file (and no output file) 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. (Any terminal status with an output file is billed through the completed
+        path instead, covered by test_terminal_status_with_output_file_is_billed.)
         """
         import base64
         import json
@@ -832,15 +832,11 @@ async def test_terminal_status_persists_managed_output_file_ids(
         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()
@@ -884,29 +880,21 @@ def find_managed_file(where):
             input_file_id=raw_input_file_id,
             object="batch",
             status=terminal_status,
-            output_file_id=raw_output_file_id,
+            output_file_id=None,
             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.get_unified_output_file_id.side_effect = [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(
+        mock_hook.get_unified_output_file_id.assert_called_once_with(
             output_file_id=raw_error_file_id,
             model_id="model-123",
             model_name="gpt-5-batch",
@@ -915,10 +903,7 @@ def find_managed_file(where):
             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,
-        }
+        assert stored == {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"
@@ -932,9 +917,8 @@ def find_managed_file(where):
         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["output_file_id"] is None
         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
@@ -1067,12 +1051,17 @@ async def test_non_terminal_status_left_unprocessed(
         ), "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
+    @pytest.mark.parametrize("terminal_status", ["expired", "cancelled", "failed"])
+    async def test_terminal_status_with_output_file_is_billed(
+        self,
+        check_batch_cost_instance,
+        mock_prisma_client,
+        mock_llm_router,
+        terminal_status,
     ):
-        """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.
+        """A terminal (expired/cancelled/failed) 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
 
@@ -1085,7 +1074,7 @@ async def test_expired_with_output_file_is_billed(
         )
 
         mock_job = MagicMock()
-        mock_job.id = "job-expired-with-output-1"
+        mock_job.id = "job-terminal-with-output-1"
         mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA=="
         mock_job.created_by = "user-1"
 
@@ -1095,10 +1084,10 @@ async def test_expired_with_output_file_is_billed(
         )
 
         mock_response = MagicMock()
-        mock_response.status = "expired"
+        mock_response.status = terminal_status
         mock_response.output_file_id = "file-output-123"
         mock_response.model_dump_json.return_value = (
-            '{"id":"batch-1","status":"expired"}'
+            f'{{"id":"batch-1","status":"{terminal_status}"}}'
         )
 
         mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
@@ -1164,7 +1153,7 @@ async def test_expired_with_output_file_is_billed(
 
         assert (
             mock_afile_content.await_count == 1
-        ), "expired batch with an output file must fetch results and be billed"
+        ), f"{terminal_status} 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
@@ -1174,8 +1163,85 @@ async def test_expired_with_output_file_is_billed(
         ]["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"
+            update_data["status"] == terminal_status
+        ), f"billed {terminal_status} batch must keep its real terminal status in the DB"
+
+    @pytest.mark.asyncio
+    async def test_terminal_batch_with_missing_output_file_is_retired_unbilled(
+        self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
+    ):
+        """A terminal batch whose advertised output file 404s at the provider has
+        nothing to fetch on this or any later poll (Vertex AI advertises an output
+        path for every batch, even ones that never wrote it), so the job must be
+        retired as terminal on the first cycle instead of retrying until the
+        staleness sweep gives up on it.
+        """
+        import base64
+        from unittest.mock import patch
+
+        from litellm.exceptions import NotFoundError
+
+        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-output-gone-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]
+        )
+
+        missing_output_file_id = "gs://batch-out/job-1/predictions.jsonl"
+        mock_response = MagicMock()
+        mock_response.status = "failed"
+        mock_response.output_file_id = missing_output_file_id
+        mock_response.error_file_id = None
+        mock_response.model_dump_json.return_value = (
+            '{"id":"batch-1","status":"failed"}'
+        )
+
+        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"}
+        )
+
+        with (
+            patch(
+                "litellm.files.main.afile_content",
+                new_callable=AsyncMock,
+                side_effect=NotFoundError(
+                    message=f"404: output file {missing_output_file_id} does not exist",
+                    model="gemini-2.5-pro",
+                    llm_provider="vertex_ai",
+                ),
+            ) as mock_afile_content,
+            patch(
+                "litellm.batches.batch_utils.calculate_batch_cost_and_usage",
+                new_callable=AsyncMock,
+            ) as mock_calculate,
+        ):
+            await check_batch_cost_instance.check_batch_cost()
+
+        assert mock_afile_content.await_count == 1
+        mock_calculate.assert_not_awaited()
+        assert (
+            mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1
+        ), "a terminal batch with a 404ing output file must be retired, not retried forever"
+        update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[
+            1
+        ]["data"]
+        assert update_data["status"] == "failed"
+        assert update_data["batch_processed"] is True
 
     @pytest.mark.asyncio
     async def test_raw_output_file_id_converted_to_managed_id(