From da8414228831c36605bee568f3c90cc0d1c7b3f9 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 12 Aug 2026 21:19:37 -0700
Subject: [PATCH] fix(batches): only trust a 404 from the batch's own
 deployment

---
 .../proxy/common_utils/check_batch_cost.py    |  9 +++++-
 .../proxy_unit_tests/test_check_batch_cost.py | 29 +++++++++++++++++++
 2 files changed, 37 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 73d5fef08e0..44371c8299a 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -244,6 +244,13 @@ class CheckBatchCost:
 
         return isinstance(error, (NotFoundError, openai.NotFoundError))
 
+    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
+        deployment. Once that deployment leaves the router, default fallbacks can
+        silently send the retrieve to a provider that never saw the batch, so its
+        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 _record_error(
         prom_logger: Optional["PrometheusLogger"], error_type: str
@@ -746,7 +753,7 @@ class CheckBatchCost:
                 )
                 if prom_logger:
                     prom_logger.record_check_batch_cost_error("provider_retrieval_error")
-                if self._is_batch_gone_at_provider(e):
+                if self._is_batch_gone_at_provider(e) 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 0418bff0cc3..fc3743a0490 100644
--- a/tests/proxy_unit_tests/test_check_batch_cost.py
+++ b/tests/proxy_unit_tests/test_check_batch_cost.py
@@ -1878,6 +1878,35 @@ class TestPollPageStarvation:
             "batch_processed": True
         }
 
+    @pytest.mark.asyncio
+    async def test_provider_404_with_deployment_gone_keeps_job(self):
+        """With the batch's own deployment removed from the router, default fallbacks can
+        send the retrieve to a provider that never saw the batch. That 404 proves nothing,
+        so the row must stay unprocessed instead of losing its spend forever."""
+        import litellm
+
+        prisma = self._prisma(
+            [
+                self._job(
+                    "job-misrouted",
+                    self._encode("litellm_proxy;model_id:model-gone;llm_batch_id:batch_alive"),
+                )
+            ]
+        )
+        llm_router = MagicMock()
+        llm_router.get_deployment = MagicMock(return_value=None)
+        llm_router.aretrieve_batch = AsyncMock(
+            side_effect=litellm.NotFoundError(
+                message="No batch found with id 'batch_alive'.",
+                model="model-gone",
+                llm_provider="openai",
+            )
+        )
+
+        await self._instance(prisma, llm_router).check_batch_cost()
+
+        prisma.db.litellm_managedobjecttable.update.assert_not_awaited()
+
     @pytest.mark.asyncio
     async def test_transient_provider_error_keeps_job_for_retry(self):
         """A failure that may clear up (timeout, 5xx) must still leave the row unprocessed."""