From 8947008fd2a9fbb461e18fbfbe2bd7c072c832f2 Mon Sep 17 00:00:00 2001
From: mateo <mateo@berri.ai>
Date: Thu, 13 Aug 2026 04:22:23 +0000
Subject: [PATCH] fix(batches): only retire on a 404 that names the batch

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../proxy/common_utils/check_batch_cost.py    | 14 ++++++---
 .../proxy_unit_tests/test_check_batch_cost.py | 31 ++++++++++++++++++-
 2 files changed, 39 insertions(+), 6 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 44371c8299a..818c44d2039 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -235,14 +235,18 @@ class CheckBatchCost:
         return bool(decoded) and get_model_id_from_unified_batch_id(decoded) is None
 
     @staticmethod
-    def _is_batch_gone_at_provider(error: Exception) -> bool:
-        """A 404 from the provider means it dropped its record of the batch, so no later
-        retrieve can ever succeed."""
+    def _is_batch_gone_at_provider(error: Exception, batch_id: str) -> bool:
+        """
+        A 404 naming the batch means the provider dropped its record of it, so no later
+        retrieve can ever succeed. A 404 about anything else, a renamed Azure deployment
+        or a fallback deployment that never saw this batch, is still fixable in config, so
+        it keeps retrying.
+        """
         import openai
 
         from litellm.exceptions import NotFoundError
 
-        return isinstance(error, (NotFoundError, openai.NotFoundError))
+        return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error)
 
     def _batch_deployment_exists(self, model_id: str) -> bool:
         """A 404 only proves the batch is gone when it came from the batch's own
@@ -753,7 +757,7 @@ class CheckBatchCost:
                 )
                 if prom_logger:
                     prom_logger.record_check_batch_cost_error("provider_retrieval_error")
-                if self._is_batch_gone_at_provider(e) and self._batch_deployment_exists(model_id):
+                if self._is_batch_gone_at_provider(e, batch_id) and self._batch_deployment_exists(model_id):
                     await self._retire_job(job, f"batch {batch_id} no longer exists at the provider")
                 continue
 
diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py
index fc3743a0490..fa274324fd6 100644
--- a/tests/proxy_unit_tests/test_check_batch_cost.py
+++ b/tests/proxy_unit_tests/test_check_batch_cost.py
@@ -1983,7 +1983,9 @@ class TestPollPageStarvation:
         async def _retrieve(model, batch_id, litellm_metadata):
             if batch_id == "batch_deadbeef":
                 raise litellm.NotFoundError(
-                    message="No batch found", model=model, llm_provider="openai"
+                    message=f"No batch found with id '{batch_id}'.",
+                    model=model,
+                    llm_provider="openai",
                 )
             return in_progress
 
@@ -2000,3 +2002,30 @@ class TestPollPageStarvation:
         assert (
             llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live"
         ), "the newer healthy batch must still be polled in the same cycle"
+
+    @pytest.mark.asyncio
+    async def test_404_that_does_not_name_the_batch_keeps_job_for_retry(self):
+        """A 404 about something other than the batch, e.g. a renamed Azure deployment, is
+        fixable in config, so the row must survive to be costed after the fix."""
+        import litellm
+
+        prisma = self._prisma(
+            [
+                self._job(
+                    "job-bad-deployment",
+                    self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_real"),
+                )
+            ]
+        )
+        llm_router = MagicMock()
+        llm_router.aretrieve_batch = AsyncMock(
+            side_effect=litellm.NotFoundError(
+                message="Error code: 404 - DeploymentNotFound",
+                model="model-123",
+                llm_provider="azure",
+            )
+        )
+
+        await self._instance(prisma, llm_router).check_batch_cost()
+
+        prisma.db.litellm_managedobjecttable.update.assert_not_awaited()