From 3ee91e52208480a491248d4227a070d9dc69921c Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Sat, 17 Oct 2026 15:24:36 -0700
Subject: [PATCH 1/8] fix(batch): preserve recoverable costs and rotate polling
 pages

---
 .../proxy/common_utils/check_batch_cost.py    |  94 +++++-----
 .../managed_object_attribution.py             |   7 +
 .../spend/test_batch_poll_starvation.py       |  71 ++++++--
 .../common_utils/test_check_batch_cost.py     | 169 +++++++++++++-----
 4 files changed, 241 insertions(+), 100 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 2d2bc92dbd2..c73ef351d96 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -109,6 +109,8 @@ class CheckBatchCost:
         # the guaranteed-failing primary query on every subsequent cycle.
         self._has_batch_processed_column: bool = True
         self.batch_processed_support_confirmed: bool = False
+        self._poll_cursor: tuple[datetime, str] | None = None
+        self._poll_before: datetime | None = None
         self._line_item_stores: Final[set[asyncio.Task[None]]] = set()
         self._line_item_store_turn: Final = asyncio.Lock()
         self._poll_idle: Final = asyncio.Event()
@@ -161,29 +163,9 @@ class CheckBatchCost:
                 f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
             )
 
-        if not self._has_batch_processed_column:
-            return
-
-        # A row already in a terminal status is never rewritten by the sweep above, so
-        # without this it keeps a poll-page slot forever and starves newer batches.
-        retired: Final = await _managed_object_table(self.prisma_client).update_many(
-            where={
-                "file_purpose": "batch",
-                "batch_processed": False,
-                "status": {"in": ["complete", "completed"]},
-                "created_at": {"lt": cutoff},
-            },
-            data={"batch_processed": True},
-        )
-        if retired > 0:
-            verbose_proxy_logger.warning(
-                f"CheckBatchCost: gave up on {retired} completed managed objects older than "
-                f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
-            )
-
     async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]":
         """Query batch jobs without the batch_processed filter (for older schemas)."""
-        return await _managed_object_table(self.prisma_client).find_many(
+        return await self._find_jobs(
             where={
                 "file_purpose": "batch",
                 "status": {
@@ -197,10 +179,40 @@ class CheckBatchCost:
                     ]
                 },
             },
-            take=MAX_OBJECTS_PER_POLL_CYCLE,
-            order={"created_at": "asc"},
         )
 
+    async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[ManagedObjectRow]:
+        cursor: Final = self._poll_cursor
+        if cursor is None:
+            self._poll_before = datetime.now(timezone.utc)
+        page: Final = await _managed_object_table(self.prisma_client).find_many(
+            where={
+                **where,
+                "created_at": {"lte": self._poll_before},
+                **(
+                    {
+                        "OR": [
+                            {"created_at": {"gt": cursor[0]}},
+                            {"created_at": cursor[0], "id": {"gt": cursor[1]}},
+                        ]
+                    }
+                    if cursor is not None
+                    else {}
+                ),
+            },
+            take=MAX_OBJECTS_PER_POLL_CYCLE,
+            order=[{"created_at": "asc"}, {"id": "asc"}],
+        )
+        self._poll_cursor = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
+        return page
+
+    def _unprocessed_snapshot(self, job: ManagedObjectRow) -> Mapping[str, object]:
+        return {
+            "id": job.id,
+            "updated_at": job.updated_at,
+            **({"batch_processed": False} if self._has_batch_processed_column else {}),
+        }
+
     async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None:
         """
         Take a row that can never be costed out of the poll page. Leaving it selectable
@@ -208,24 +220,19 @@ class CheckBatchCost:
         once enough such rows accumulate no newer batch is ever reached. Older schemas
         without batch_processed can only be excluded through the status filter.
         """
-        data: Final = (
-            {"batch_processed": True}
-            if self._has_batch_processed_column
-            else {"status": "stale_expired"}
-        )
+        data: Final = {"batch_processed": True} if self._has_batch_processed_column else {"status": "stale_expired"}
         try:
-            await _managed_object_table(self.prisma_client).update(
-                where={"id": job.id},
+            retired: Final = await _managed_object_table(self.prisma_client).update_many(
+                where=self._unprocessed_snapshot(job),
                 data=data,
             )
         except Exception as db_err:
-            verbose_proxy_logger.error(
-                f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}"
-            )
+            verbose_proxy_logger.error(f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}")
+            return
+        if retired == 0:
             return
         verbose_proxy_logger.warning(
-            f"CheckBatchCost: job {job.id} can never be costed ({reason}), "
-            "so it will no longer be polled"
+            f"CheckBatchCost: job {job.id} can never be costed ({reason}), so it will no longer be polled"
         )
 
     async def _claim_job_for_costing(self, job: "ManagedObjectRow", response: "LiteLLMBatch") -> bool:
@@ -355,18 +362,18 @@ class CheckBatchCost:
                 db_batch_object=job,
                 unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id),
             )
-            update_data: Final[dict] = {
+            update_data: Final[dict[str, object]] = {
                 "status": response.status,
                 "file_object": response.model_dump_json(),
                 **({"batch_processed": True} if self._has_batch_processed_column else {}),
             }
-            await _managed_object_table(self.prisma_client).update(
-                where={"id": job.id},
+            finalized: Final = await _managed_object_table(self.prisma_client).update_many(
+                where=self._unprocessed_snapshot(job),
                 data=update_data,
             )
-            verbose_proxy_logger.info(
-                f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
-            )
+            if finalized == 0:
+                return False
+            verbose_proxy_logger.info(f"CheckBatchCost: marked job {job.id} as {response.status} in DB")
             return True
         except Exception as db_err:
             verbose_proxy_logger.error(
@@ -1028,6 +1035,7 @@ class CheckBatchCost:
         """
         try:
             from litellm.integrations.prometheus import PrometheusLogger
+
             prom_logger = PrometheusLogger.get_instance()
         except Exception as e:
             verbose_proxy_logger.error(f"CheckBatchCost: could not get Prometheus logger: {e}")
@@ -1052,7 +1060,7 @@ class CheckBatchCost:
                 # endpoint may transition a batch to "complete" before
                 # CheckBatchCost runs.  The batch_processed=False filter
                 # already prevents reprocessing finished batches.
-                jobs = await _managed_object_table(self.prisma_client).find_many(
+                jobs = await self._find_jobs(
                     where={
                         "file_purpose": "batch",
                         "batch_processed": False,
@@ -1065,8 +1073,6 @@ class CheckBatchCost:
                             ]
                         },
                     },
-                    take=MAX_OBJECTS_PER_POLL_CYCLE,
-                    order={"created_at": "asc"},
                 )
                 self.batch_processed_support_confirmed = True
             except Exception as query_err:
diff --git a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py
index 3c7735b53c0..a079fcf0b45 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py
@@ -9,6 +9,7 @@ from __future__ import annotations
 
 from collections.abc import Mapping
 from dataclasses import dataclass
+from datetime import datetime
 from typing import TYPE_CHECKING, Final, Protocol, cast
 
 from pydantic import TypeAdapter
@@ -27,6 +28,12 @@ if TYPE_CHECKING:
 
 
 class ManagedObjectRow(Protocol):
+    @property
+    def created_at(self) -> datetime: ...
+
+    @property
+    def updated_at(self) -> datetime: ...
+
     @property
     def id(self) -> str: ...
 
diff --git a/tests/integration/spend/test_batch_poll_starvation.py b/tests/integration/spend/test_batch_poll_starvation.py
index f68f10e3d8a..4039bd80a8b 100644
--- a/tests/integration/spend/test_batch_poll_starvation.py
+++ b/tests/integration/spend/test_batch_poll_starvation.py
@@ -1,12 +1,13 @@
 import json
 import os
+from functools import partial
 from hashlib import sha256
 from typing import Final
 
 import pytest
 from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value
-from integration._support.database import read_rows
-from integration._support.upstream import delete_scenario, register_scenario
+from integration._support.database import read_rows, write_rows
+from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
 from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse
 from pydantic import JsonValue
 
@@ -59,17 +60,17 @@ def _accepting_routes() -> dict[str, JsonResponse | TextResponse]:
     }
 
 
-def _gone_at_provider_routes() -> RoutedResponse:
+def _gone_at_provider_routes(*, retryable: bool = False) -> RoutedResponse:
     return RoutedResponse(
         content_type="application/x-routed",
         routes={
             **_accepting_routes(),
             "GET /batches/batch-$REQUEST_ID": JsonResponse(
                 content_type="application/json",
-                status=404,
+                status=503 if retryable else 404,
                 body={
                     "error": {
-                        "message": "No batch found with id 'batch-$REQUEST_ID'.",
+                        "message": "Temporary failure" if retryable else "No batch found with id 'batch-$REQUEST_ID'.",
                         "type": "invalid_request_error",
                         "param": "id",
                         "code": "batch_not_found",
@@ -80,7 +81,7 @@ def _gone_at_provider_routes() -> RoutedResponse:
     )
 
 
-def _completed_routes() -> RoutedResponse:
+def _completed_routes(*, ready: bool = True) -> RoutedResponse:
     output_line: Final = {
         "id": "batch_req_1",
         "custom_id": "r1",
@@ -106,7 +107,8 @@ def _completed_routes() -> RoutedResponse:
         routes={
             **_accepting_routes(),
             "GET /batches/batch-$REQUEST_ID": JsonResponse(
-                content_type="application/json", body=_batch("completed", "file-out-$REQUEST_ID")
+                content_type="application/json",
+                body=_batch("completed", "file-out-$REQUEST_ID") if ready else _batch("in_progress", None),
             ),
             "GET /files/file-out-$REQUEST_ID/content": TextResponse(
                 content_type="application/jsonl", body=json.dumps(output_line, separators=(",", ":")) + "\n"
@@ -115,7 +117,7 @@ def _completed_routes() -> RoutedResponse:
     )
 
 
-def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse) -> str:
+def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse) -> tuple[str, ScenarioHandle]:
     scenario_id: Final = f"poll-{marker}-{sha256(os.urandom(16)).hexdigest()[:12]}"
     handle: Final = register_scenario(scenario_id, routes)
     scenario.cleanups.callback(delete_scenario, handle)
@@ -133,7 +135,7 @@ def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse
         },
     )
     scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"]))
-    return string_value(created["model_name"])
+    return string_value(created["model_name"]), handle
 
 
 def _submitted_batch_id(gateway: Gateway, key: str, model_name: str) -> str:
@@ -173,20 +175,39 @@ def _managed_rows(batch_ids: tuple[str, ...]) -> list[dict[str, JsonValue]]:
     )
 
 
+@pytest.mark.parametrize("retryable", [False, True])
+@pytest.mark.parametrize("old_completed", [False, True])
 @pytest.mark.timeout(180)
 @pytest.mark.covers("quota_management.spend_tracking.batch_costs.uncostable_rows_retire_so_newer_batches_are_costed")
-def test_batches_gone_at_provider_do_not_starve_a_newer_batch_out_of_cost_polling(gateway: Gateway) -> None:
+def test_uncostable_batches_do_not_starve_billable_batches(
+    gateway: Gateway, retryable: bool, old_completed: bool
+) -> None:
     with gateway.scenario() as scenario:
-        key: Final = scenario.key()
+        team_id: Final = scenario.team()
+        key: Final = scenario.key(team_id=team_id)
         gone_batch_ids: Final = tuple(
             _submitted_batch_id(
-                gateway, key, _scripted_deployment(scenario, f"gone{index}", _gone_at_provider_routes())
+                gateway,
+                key,
+                _scripted_deployment(scenario, f"gone{index}", _gone_at_provider_routes(retryable=retryable))[0],
             )
             for index in range(MAX_OBJECTS_PER_POLL_CYCLE)
         )
-        costable_batch_id: Final = _submitted_batch_id(
-            gateway, key, _scripted_deployment(scenario, "costable", _completed_routes())
-        )
+        model, upstream = _scripted_deployment(scenario, "costable", _completed_routes(ready=False))
+        costable_batch_id: Final = _submitted_batch_id(gateway, key, model)
+        for batch_id in (*gone_batch_ids, costable_batch_id):
+            scenario.cleanups.callback(
+                write_rows,
+                'DELETE FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s',
+                (batch_id,),
+            )
+        if old_completed:
+            write_rows(
+                'UPDATE "LiteLLM_ManagedObjectTable" SET status = %s, '
+                "created_at = NOW() - INTERVAL '30 days' WHERE unified_object_id = %s",
+                ("completed", costable_batch_id),
+            )
+        register_scenario(upstream.scenario_id, _completed_routes(), control_url=upstream.control_url)
         spend_rows: Final = eventually(
             lambda: read_rows(
                 'SELECT call_type, status, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" '
@@ -209,4 +230,22 @@ def test_batches_gone_at_provider_do_not_starve_a_newer_batch_out_of_cost_pollin
             }
         ]
         assert _managed_rows((costable_batch_id,)) == [{"batch_processed": True}]
-        assert _managed_rows(gone_batch_ids) == [{"batch_processed": True}] * MAX_OBJECTS_PER_POLL_CYCLE
+        expected_blockers: Final = [{"batch_processed": not retryable}] * MAX_OBJECTS_PER_POLL_CYCLE
+        blockers: Final = eventually(
+            lambda: _managed_rows(gone_batch_ids), lambda rows: rows == expected_blockers, seconds=120
+        )
+        assert blockers == expected_blockers
+
+        expected_spend: Final = BATCH_COST_SHARE * (
+            PROMPT_TOKENS * INPUT_COST_PER_TOKEN + COMPLETION_TOKENS * OUTPUT_COST_PER_TOKEN
+        )
+        for table, column, identity in (
+            ("LiteLLM_VerificationToken", "token", sha256(key.encode()).hexdigest()),
+            ("LiteLLM_TeamTable", "team_id", team_id),
+        ):
+            balance: Final = eventually(
+                partial(read_rows, f'SELECT spend FROM "{table}" WHERE {column} = %s', (identity,)),
+                lambda rows: len(rows) == 1 and rows[0]["spend"] == pytest.approx(expected_spend),
+                seconds=70,
+            )
+            assert balance == [{"spend": pytest.approx(expected_spend)}]
diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py
index eb2a5b7f66c..fc5c9192c62 100644
--- a/tests/unit/proxy/common_utils/test_check_batch_cost.py
+++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py
@@ -10,6 +10,8 @@ import asyncio
 import json
 import logging
 from contextlib import contextmanager
+from copy import copy
+from datetime import datetime, timezone
 from hashlib import sha256
 from typing import TYPE_CHECKING, Final
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -211,7 +213,7 @@ class TestCheckBatchCost:
 
         find_call = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args
         assert find_call[1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE
-        assert find_call[1]["order"] == {"created_at": "asc"}
+        assert find_call[1]["order"] == [{"created_at": "asc"}, {"id": "asc"}]
         not_in = find_call[1]["where"]["status"]["not_in"]
         assert "stale_expired" in not_in
         # "complete"/"completed" are intentionally NOT excluded from the
@@ -973,10 +975,10 @@ class TestCheckBatchCost:
 
         await check_batch_cost_instance.check_batch_cost()
 
-        assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
+        assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, (
             f"Expected update() to be called exactly once for a {terminal_status} job"
         )
-        update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"]
+        update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
         assert update_data["status"] == terminal_status
         assert update_data["batch_processed"] is True, (
             "terminal-status update() must set batch_processed=True so polling stops"
@@ -1079,9 +1081,13 @@ class TestCheckBatchCost:
         assert owner.user_id == "user-1"
         assert owner.team_id == "team-1"
 
-        assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1
-        update_call = mock_prisma_client.db.litellm_managedobjecttable.update.call_args
-        assert update_call.kwargs["where"] == {"id": "job-terminal-mint-1"}
+        assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2
+        update_call = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args
+        assert update_call.kwargs["where"] == {
+            "id": "job-terminal-mint-1",
+            "batch_processed": False,
+            "updated_at": mock_job.updated_at,
+        }
         update_data = update_call.kwargs["data"]
         assert update_data["status"] == terminal_status
         assert update_data["batch_processed"] is True
@@ -1146,10 +1152,10 @@ class TestCheckBatchCost:
         ) as mock_afile_content:
             await check_batch_cost_instance.check_batch_cost()
 
-        assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
+        assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, (
             "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"]
+        update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.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"
@@ -1538,10 +1544,10 @@ class TestCheckBatchCost:
 
         assert mock_afile_content.await_count == 1
         mock_calculate.assert_not_awaited()
-        assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
+        assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, (
             "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"]
+        update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
         assert update_data["status"] == "failed"
         assert update_data["batch_processed"] is True
 
@@ -2737,9 +2743,10 @@ class TestPollPageStarvation:
         await self._instance(prisma, llm_router).check_batch_cost()
 
         llm_router.aretrieve_batch.assert_not_awaited()
-        prisma.db.litellm_managedobjecttable.update.assert_awaited_once()
-        call = prisma.db.litellm_managedobjecttable.update.call_args[1]
-        assert call["where"] == {"id": "job-no-model"}
+        prisma.db.litellm_managedobjecttable.update_many.assert_awaited()
+        call = prisma.db.litellm_managedobjecttable.update_many.call_args[1]
+        assert call["where"]["id"] == "job-no-model"
+        assert call["where"]["batch_processed"] is False
         assert call["data"] == {"batch_processed": True}
 
     @pytest.mark.asyncio
@@ -2767,8 +2774,8 @@ class TestPollPageStarvation:
 
         await self._instance(prisma, llm_router).check_batch_cost()
 
-        prisma.db.litellm_managedobjecttable.update.assert_awaited_once()
-        assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"batch_processed": True}
+        prisma.db.litellm_managedobjecttable.update_many.assert_awaited()
+        assert prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] == {"batch_processed": True}
 
     @pytest.mark.asyncio
     async def test_provider_404_with_deployment_gone_keeps_job(self):
@@ -2827,24 +2834,7 @@ class TestPollPageStarvation:
 
         await instance.check_batch_cost()
 
-        assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"status": "stale_expired"}
-
-    @pytest.mark.asyncio
-    async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self):
-        """A row already in a terminal status is never rewritten by the staleness sweep, so
-        it needs its own bound or it starves newer batches indefinitely."""
-        prisma = self._prisma([])
-
-        await self._instance(prisma, MagicMock()).check_batch_cost()
-
-        calls = prisma.db.litellm_managedobjecttable.update_many.call_args_list
-        assert len(calls) == 2, "expected the staleness sweep plus the never-costed sweep"
-        where = calls[1][1]["where"]
-        assert where["file_purpose"] == "batch"
-        assert where["batch_processed"] is False
-        assert where["status"] == {"in": ["complete", "completed"]}
-        assert "created_at" in where
-        assert calls[1][1]["data"] == {"batch_processed": True}
+        assert prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] == {"status": "stale_expired"}
 
     @pytest.mark.asyncio
     async def test_newer_batch_is_polled_once_dead_rows_are_retired(self):
@@ -2882,7 +2872,11 @@ class TestPollPageStarvation:
 
         await self._instance(prisma, llm_router).check_batch_cost()
 
-        retired = [call[1]["where"]["id"] for call in prisma.db.litellm_managedobjecttable.update.call_args_list]
+        retired = [
+            call[1]["where"]["id"]
+            for call in prisma.db.litellm_managedobjecttable.update_many.call_args_list
+            if "id" in call[1]["where"]
+        ]
         assert retired == ["job-no-model", "job-gone"]
         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"
@@ -2930,7 +2924,8 @@ class _FakeManagedObjectRow:
         self.team_id = None
         self.api_key = None
         self.request_tags = None
-        self.created_at = 1700000000
+        self.created_at = datetime(9999, 1, 1, tzinfo=timezone.utc)
+        self.updated_at = datetime(2000, 1, 1, tzinfo=timezone.utc)
         self.file_object = json.dumps(
             {
                 "id": "batch-456",
@@ -2960,7 +2955,9 @@ class _FakeManagedObjectTable:
     def _matches(self, where: dict) -> bool:
         for key, value in where.items():
             if key == "created_at":
-                return False
+                if "lt" in value and not self.row.created_at < value["lt"]:
+                    return False
+                continue
             if key == "status":
                 if self.row.status in value.get("not_in", []):
                     return False
@@ -2973,7 +2970,9 @@ class _FakeManagedObjectTable:
     async def _update_many(self, *, where: dict, data: dict) -> int:
         if not self._matches(where):
             return 0
-        if "batch_processed" in where:
+        if "status" in data and "file_object" in data:
+            self.journal.append("finalize")
+        elif "batch_processed" in where:
             self.journal.append("claim" if data.get("batch_processed") else "release")
         for key, value in data.items():
             setattr(self.row, key, value)
@@ -3693,12 +3692,15 @@ class TestBatchLineItemsFromThePoller:
         assert logger.payloads == []
 
     @pytest.mark.asyncio
-    async def test_a_batch_that_could_not_be_retired_sends_no_lines_until_it_is(self, monkeypatch):
+    @pytest.mark.parametrize("database_error", [False, True])
+    async def test_a_batch_that_could_not_be_retired_sends_no_lines_until_it_is(self, monkeypatch, database_error):
         import litellm
 
         journal: list = []
         row, instance = self._setup(journal, status="expired", output_file_id=None, error_file_id="file-err-1")
-        instance.prisma_client.db.litellm_managedobjecttable.update = AsyncMock(side_effect=RuntimeError("db down"))
+        instance.prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            side_effect=RuntimeError("db down") if database_error else None, return_value=0
+        )
         logger = self._recording_logger(journal)
         monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,))
 
@@ -3834,7 +3836,6 @@ class TestBatchLineItemsFromThePoller:
         assert row.batch_processed is True
         assert f"so the line items of batch batch-456 (job {row.id}) were not stored" in caplog.text
 
-
     @pytest.mark.asyncio
     async def test_batches_store_their_lines_one_at_a_time_in_the_order_they_were_retired(self, monkeypatch):
         import litellm
@@ -4109,3 +4110,91 @@ class TestCreatorTurnedOffMessageLogging:
 
         assert await instance._creator_turned_off_message_logging(self._job(api_key=api_key)) is False
         instance.prisma_client.db.litellm_verificationtoken.find_unique.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_old_completed_batch_remains_available_for_costing() -> None:
+    row: Final = _FakeManagedObjectRow()
+    row.status = "completed"
+    row.created_at = datetime(2000, 1, 1, tzinfo=timezone.utc)
+    prisma: Final = TestMultiPodBatchCostClaim._prisma(row, [])
+    poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router())
+
+    await poller._cleanup_stale_managed_objects()
+
+    assert row.batch_processed is False
+    assert row.status == "completed"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("progress", ["claimed", "refreshed"])
+async def test_stale_retirement_preserves_concurrent_progress(progress: str) -> None:
+    row: Final = _FakeManagedObjectRow()
+    selected: Final = copy(row)
+    if progress == "claimed":
+        row.batch_processed = True
+    row.updated_at = datetime(2000, 1, 2, tzinfo=timezone.utc)
+    prisma: Final = TestMultiPodBatchCostClaim._prisma(row, [])
+    poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router())
+
+    await poller._retire_job(selected, "provider batch disappeared")
+
+    assert row.batch_processed is (progress == "claimed")
+    assert row.updated_at > selected.updated_at
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("progress", ["claimed", "refreshed"])
+async def test_stale_terminal_response_preserves_concurrent_progress(progress: str) -> None:
+    from litellm.types.utils import LiteLLMBatch
+
+    row: Final = _FakeManagedObjectRow()
+    selected: Final = copy(row)
+    if progress == "claimed":
+        row.batch_processed = True
+    row.updated_at = datetime(2000, 1, 2, tzinfo=timezone.utc)
+    row.file_object = '{"revision":"fresh"}'
+    prisma: Final = TestMultiPodBatchCostClaim._prisma(row, [])
+    poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router())
+    response: Final = LiteLLMBatch(
+        id="batch-456",
+        completion_window="24h",
+        created_at=1,
+        endpoint="/v1/chat/completions",
+        input_file_id="file-input-1",
+        object="batch",
+        status="failed",
+    )
+
+    assert await poller._finalize_unbilled_terminal_job(selected, response) is False
+
+    assert row.file_object == '{"revision":"fresh"}'
+    assert row.status == "in_progress"
+    assert row.batch_processed is (progress == "claimed")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("legacy", [False, True])
+async def test_permanently_lost_batch_retires_unchanged_row(legacy: bool) -> None:
+    row: Final = _FakeManagedObjectRow()
+    prisma: Final = TestMultiPodBatchCostClaim._prisma(row, [])
+    poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router())
+    poller._has_batch_processed_column = not legacy
+
+    await poller._retire_job(row, "provider batch disappeared")
+
+    assert row.batch_processed is (not legacy)
+    assert row.status == ("stale_expired" if legacy else "in_progress")
+
+
+@pytest.mark.asyncio
+async def test_failed_retirement_keeps_batch_available() -> None:
+    row: Final = _FakeManagedObjectRow()
+    prisma: Final = TestMultiPodBatchCostClaim._prisma(row, [])
+    prisma.db.litellm_managedobjecttable.update_many = AsyncMock(side_effect=RuntimeError("database unavailable"))
+    poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router())
+
+    await poller._retire_job(row, "provider batch disappeared")
+
+    assert row.batch_processed is False
+    assert row.status == "in_progress"

From abddf0363fd390eda33503308b2e0c1a94f0f00a Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Sat, 17 Oct 2026 15:35:59 -0700
Subject: [PATCH 2/8] test(batch): type retirement failure regression

---
 tests/unit/proxy/common_utils/test_check_batch_cost.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py
index fc5c9192c62..5d28abd3c31 100644
--- a/tests/unit/proxy/common_utils/test_check_batch_cost.py
+++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py
@@ -3693,7 +3693,9 @@ class TestBatchLineItemsFromThePoller:
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize("database_error", [False, True])
-    async def test_a_batch_that_could_not_be_retired_sends_no_lines_until_it_is(self, monkeypatch, database_error):
+    async def test_a_batch_that_could_not_be_retired_sends_no_lines_until_it_is(
+        self, monkeypatch: pytest.MonkeyPatch, database_error: bool
+    ) -> None:
         import litellm
 
         journal: list = []

From aab0323a4dc65b6ce6ae3e1cb6b21a1154a07e2d Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Sat, 17 Oct 2026 16:46:35 -0700
Subject: [PATCH 3/8] fix(batch): alternate new arrivals with retry polling

---
 .../proxy/common_utils/check_batch_cost.py    | 24 ++++--
 .../common_utils/test_check_batch_cost.py     | 82 +++++++++++++++++--
 2 files changed, 96 insertions(+), 10 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 c73ef351d96..81f4aac4885 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -111,6 +111,8 @@ class CheckBatchCost:
         self.batch_processed_support_confirmed: bool = False
         self._poll_cursor: tuple[datetime, str] | None = None
         self._poll_before: datetime | None = None
+        self._arrival_cursor: tuple[datetime, str] = (datetime.now(timezone.utc), "")
+        self._poll_arrivals: bool = False
         self._line_item_stores: Final[set[asyncio.Task[None]]] = set()
         self._line_item_store_turn: Final = asyncio.Lock()
         self._poll_idle: Final = asyncio.Event()
@@ -182,13 +184,21 @@ class CheckBatchCost:
         )
 
     async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[ManagedObjectRow]:
-        cursor: Final = self._poll_cursor
-        if cursor is None:
+        arrivals: Final = self._poll_arrivals
+        page: Final = await self._find_job_page(where, arrivals=arrivals)
+        self._poll_arrivals = not arrivals if page else arrivals
+        if page:
+            return page
+        return await self._find_job_page(where, arrivals=not arrivals)
+
+    async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]:
+        cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor
+        if not arrivals and cursor is None:
             self._poll_before = datetime.now(timezone.utc)
         page: Final = await _managed_object_table(self.prisma_client).find_many(
             where={
                 **where,
-                "created_at": {"lte": self._poll_before},
+                "created_at": {"lte": datetime.now(timezone.utc) if arrivals else self._poll_before},
                 **(
                     {
                         "OR": [
@@ -203,7 +213,11 @@ class CheckBatchCost:
             take=MAX_OBJECTS_PER_POLL_CYCLE,
             order=[{"created_at": "asc"}, {"id": "asc"}],
         )
-        self._poll_cursor = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
+        if arrivals:
+            if page:
+                self._arrival_cursor = (page[-1].created_at, page[-1].id)
+        else:
+            self._poll_cursor = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
         return page
 
     def _unprocessed_snapshot(self, job: ManagedObjectRow) -> Mapping[str, object]:
@@ -325,7 +339,7 @@ class CheckBatchCost:
         """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."""
+        404 must not retire a potentially recoverable row."""
         return self.llm_router.get_deployment(model_id=model_id) is not None
 
     @staticmethod
diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py
index 5d28abd3c31..4cf7910cbb0 100644
--- a/tests/unit/proxy/common_utils/test_check_batch_cost.py
+++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py
@@ -9,11 +9,14 @@ ARN unified_object_id) batches with no managed unified id.
 import asyncio
 import json
 import logging
+from collections.abc import Mapping, Sequence
 from contextlib import contextmanager
 from copy import copy
+from dataclasses import dataclass
 from datetime import datetime, timezone
 from hashlib import sha256
-from typing import TYPE_CHECKING, Final
+from itertools import chain
+from typing import TYPE_CHECKING, Final, cast
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -23,6 +26,7 @@ from litellm.types.utils import LiteLLMBatch
 
 if TYPE_CHECKING:
     from litellm.batches.batch_utils import BatchCostUsageResult
+    from litellm_enterprise.proxy.common_utils.managed_object_attribution import ManagedObjectRow
 
 _IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id"
 _CLAIM_UNIFIED_BATCH_ID = "dW5pZmllZF9iYXRjaF9pZA=="
@@ -236,13 +240,13 @@ class TestCheckBatchCost:
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
         # First find_many (primary query) raises with a schema error; second (fallback) returns empty
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
-            side_effect=[Exception("column batch_processed does not exist"), []]
+            side_effect=[Exception("column batch_processed does not exist"), [], []]
         )
 
         await check_batch_cost_instance.check_batch_cost()
 
         calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list
-        assert len(calls) == 2
+        assert len(calls) == 3
         fallback_where = calls[1][1]["where"]
         assert "batch_processed" not in fallback_where
         assert "stale_expired" in fallback_where["status"]["not_in"]
@@ -263,8 +267,7 @@ class TestCheckBatchCost:
 
         await check_batch_cost_instance.check_batch_cost()
 
-        # Only one find_many call — the fallback directly, no primary query attempt
-        assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1
+        assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 2
         fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"]
         assert "batch_processed" not in fallback_where
 
@@ -4200,3 +4203,72 @@ async def test_failed_retirement_keeps_batch_available() -> None:
 
     assert row.batch_processed is False
     assert row.status == "in_progress"
+
+
+@dataclass(frozen=True, slots=True)
+class _PollingBatch:
+    id: str
+    created_at: datetime
+
+
+def _matches_poll_filter(row: _PollingBatch, where: Mapping[str, object]) -> bool:
+    for key, value in where.items():
+        if key == "OR":
+            if not any(_matches_poll_filter(row, branch) for branch in cast(Sequence[Mapping[str, object]], value)):
+                return False
+            continue
+        actual: Final = row.created_at if key == "created_at" else row.id
+        if not isinstance(value, Mapping):
+            if actual != value:
+                return False
+            continue
+        for operator, bound in cast(Mapping[str, datetime | str], value).items():
+            if isinstance(actual, datetime) and isinstance(bound, datetime):
+                if (operator == "gt" and actual <= bound) or (operator == "lte" and actual > bound):
+                    return False
+            elif isinstance(actual, str) and isinstance(bound, str) and operator == "gt" and actual <= bound:
+                return False
+    return True
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("page_size", [1, 3])
+async def test_new_batches_progress_without_waiting_for_retry_scan(
+    monkeypatch: pytest.MonkeyPatch, page_size: int
+) -> None:
+    from litellm_enterprise.proxy.common_utils import check_batch_cost as module
+
+    monkeypatch.setattr(module, "MAX_OBJECTS_PER_POLL_CYCLE", page_size)
+    backlog: Final = tuple(
+        _PollingBatch(f"old-{index:03}", datetime(2000, 1, 1, tzinfo=timezone.utc)) for index in range(page_size * 5)
+    )
+    arrivals: Final = tuple(
+        _PollingBatch(f"new-{index:03}", datetime(2020, index // page_size + 1, 1, tzinfo=timezone.utc))
+        for index in range(page_size * 3)
+    )
+
+    async def select_page(*, where: Mapping[str, object], take: int, order: object) -> tuple[_PollingBatch, ...]:
+        return tuple(row for row in (*backlog, *arrivals) if _matches_poll_filter(row, where))[:take]
+
+    prisma: Final = MagicMock()
+    prisma.db.litellm_managedobjecttable.find_many = AsyncMock(side_effect=select_page)
+    with patch.object(module, "datetime") as clock:
+        clock.now.return_value = datetime(2010, 1, 1, tzinfo=timezone.utc)
+        poller: Final = module.CheckBatchCost(MagicMock(), prisma, MagicMock())
+        first: Final = await poller._find_jobs({})
+        assert tuple(row.id for row in first) == tuple(row.id for row in backlog[:page_size])
+
+        async def poll_next_turn(turn: int) -> Sequence[ManagedObjectRow]:
+            clock.now.return_value = datetime(2020, turn // 2 + 1, 2, tzinfo=timezone.utc)
+            return await poller._find_jobs({})
+
+        pages: Final = tuple([await poll_next_turn(turn) for turn in range(5)])
+        selected: Final = frozenset(row.id for row in chain.from_iterable(pages))
+        assert all(len(page) <= page_size for page in pages)
+        assert all(row.id in selected for row in arrivals)
+        assert all(row.id in selected for row in backlog[page_size : 3 * page_size])
+        clock.now.return_value = datetime(2021, 1, 1, tzinfo=timezone.utc)
+        restarted: Final = module.CheckBatchCost(MagicMock(), prisma, MagicMock())
+        recovered: Final = tuple([await restarted._find_jobs({}) for _ in range(10)])
+        recovered_ids: Final = frozenset(row.id for row in chain.from_iterable(recovered))
+        assert recovered_ids == frozenset(row.id for row in (*backlog, *arrivals))

From ebd6eb84e9df394176e627ce4bcf6b063c8cfe85 Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Sat, 17 Oct 2026 17:00:46 -0700
Subject: [PATCH 4/8] test(batch): defer polling row annotation on Python 3.12

---
 tests/unit/proxy/common_utils/test_check_batch_cost.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py
index 4cf7910cbb0..11b84dc0fdb 100644
--- a/tests/unit/proxy/common_utils/test_check_batch_cost.py
+++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py
@@ -4258,7 +4258,7 @@ async def test_new_batches_progress_without_waiting_for_retry_scan(
         first: Final = await poller._find_jobs({})
         assert tuple(row.id for row in first) == tuple(row.id for row in backlog[:page_size])
 
-        async def poll_next_turn(turn: int) -> Sequence[ManagedObjectRow]:
+        async def poll_next_turn(turn: int) -> Sequence["ManagedObjectRow"]:
             clock.now.return_value = datetime(2020, turn // 2 + 1, 2, tzinfo=timezone.utc)
             return await poller._find_jobs({})
 

From 0d2177598ee59502c509c8829139aaac0319e28c Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Sat, 17 Oct 2026 17:16:20 -0700
Subject: [PATCH 5/8] test(batch): reconcile fixtures with atomic metadata
 claims

---
 .../unit/proxy/common_utils/test_check_batch_cost.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py
index 11b84dc0fdb..8f694465984 100644
--- a/tests/unit/proxy/common_utils/test_check_batch_cost.py
+++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py
@@ -743,7 +743,7 @@ class TestCheckBatchCost:
 
             await check_batch_cost_instance.check_batch_cost()
 
-        assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3, (
+        assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, (
             "Expected update() to be called exactly once for the completed job"
         )
         update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
@@ -876,7 +876,7 @@ class TestCheckBatchCost:
 
         mock_update_database.assert_awaited_once()
         assert mock_update_database.call_args.kwargs["response_cost"] == 0.01
-        assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3, (
+        assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, (
             "the job must still be marked processed once cost tracking succeeds"
         )
 
@@ -1374,7 +1374,7 @@ class TestCheckBatchCost:
             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_many.call_count >= 3
+        assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2
         update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
         assert update_data["batch_processed"] is True
         assert update_data["status"] == terminal_status, (
@@ -1946,7 +1946,7 @@ class TestUnmanagedVertexRouting:
         mock_logging_obj.async_success_handler.assert_awaited_once()
         assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.01
 
-        assert prisma.db.litellm_managedobjecttable.update_many.call_count >= 3
+        assert prisma.db.litellm_managedobjecttable.update_many.call_count == 2
         update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
         assert update_data["batch_processed"] is True
         assert update_data["status"] == "complete"
@@ -2163,7 +2163,7 @@ class TestUnmanagedBedrockRouting:
         mock_logging_obj.async_success_handler.assert_awaited_once()
         assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.02
 
-        assert prisma.db.litellm_managedobjecttable.update_many.call_count >= 3
+        assert prisma.db.litellm_managedobjecttable.update_many.call_count == 2
         update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
         assert update_data["batch_processed"] is True
         assert update_data["status"] == "complete"
@@ -2973,7 +2973,7 @@ class _FakeManagedObjectTable:
     async def _update_many(self, *, where: dict, data: dict) -> int:
         if not self._matches(where):
             return 0
-        if "status" in data and "file_object" in data:
+        if "updated_at" in where:
             self.journal.append("finalize")
         elif "batch_processed" in where:
             self.journal.append("claim" if data.get("batch_processed") else "release")

From cb562f16a0c181f001169d93ca99759efaabb888 Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Sat, 17 Oct 2026 23:48:11 -0700
Subject: [PATCH 6/8] refactor(batch): simplify reconciliation to one bounded
 scan

---
 .../proxy/common_utils/check_batch_cost.py    | 22 +----
 .../spend/test_batch_poll_starvation.py       | 58 +++++++++++++
 .../common_utils/test_check_batch_cost.py     | 81 +------------------
 3 files changed, 66 insertions(+), 95 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 81f4aac4885..e8addb4f586 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -111,8 +111,6 @@ class CheckBatchCost:
         self.batch_processed_support_confirmed: bool = False
         self._poll_cursor: tuple[datetime, str] | None = None
         self._poll_before: datetime | None = None
-        self._arrival_cursor: tuple[datetime, str] = (datetime.now(timezone.utc), "")
-        self._poll_arrivals: bool = False
         self._line_item_stores: Final[set[asyncio.Task[None]]] = set()
         self._line_item_store_turn: Final = asyncio.Lock()
         self._poll_idle: Final = asyncio.Event()
@@ -184,21 +182,13 @@ class CheckBatchCost:
         )
 
     async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[ManagedObjectRow]:
-        arrivals: Final = self._poll_arrivals
-        page: Final = await self._find_job_page(where, arrivals=arrivals)
-        self._poll_arrivals = not arrivals if page else arrivals
-        if page:
-            return page
-        return await self._find_job_page(where, arrivals=not arrivals)
-
-    async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]:
-        cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor
-        if not arrivals and cursor is None:
+        cursor: Final = self._poll_cursor
+        if cursor is None:
             self._poll_before = datetime.now(timezone.utc)
         page: Final = await _managed_object_table(self.prisma_client).find_many(
             where={
                 **where,
-                "created_at": {"lte": datetime.now(timezone.utc) if arrivals else self._poll_before},
+                "created_at": {"lte": self._poll_before},
                 **(
                     {
                         "OR": [
@@ -213,11 +203,7 @@ class CheckBatchCost:
             take=MAX_OBJECTS_PER_POLL_CYCLE,
             order=[{"created_at": "asc"}, {"id": "asc"}],
         )
-        if arrivals:
-            if page:
-                self._arrival_cursor = (page[-1].created_at, page[-1].id)
-        else:
-            self._poll_cursor = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
+        self._poll_cursor = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
         return page
 
     def _unprocessed_snapshot(self, job: ManagedObjectRow) -> Mapping[str, object]:
diff --git a/tests/integration/spend/test_batch_poll_starvation.py b/tests/integration/spend/test_batch_poll_starvation.py
index 4039bd80a8b..5981c3bc425 100644
--- a/tests/integration/spend/test_batch_poll_starvation.py
+++ b/tests/integration/spend/test_batch_poll_starvation.py
@@ -1,12 +1,16 @@
 import json
 import os
+from datetime import datetime, timezone
 from functools import partial
 from hashlib import sha256
+from itertools import chain
 from typing import Final
+from uuid import uuid4
 
 import pytest
 from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value
 from integration._support.database import read_rows, write_rows
+from integration._support.prisma_clients import connected_proxy_prisma
 from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
 from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse
 from pydantic import JsonValue
@@ -249,3 +253,57 @@ def test_uncostable_batches_do_not_starve_billable_batches(
                 seconds=70,
             )
             assert balance == [{"spend": pytest.approx(expected_spend)}]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("tail", [0, 1])
+async def test_batch_scan_finishes_and_revisits_retries_despite_new_arrivals(tail: int) -> None:
+    from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+    from litellm.proxy.utils import ProxyLogging
+    from litellm.router import Router
+    from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
+
+    owner: Final = f"batch-scan-{uuid4().hex}"
+    count: Final = 2 * MAX_OBJECTS_PER_POLL_CYCLE + tail
+
+    def insert_rows(group: str, size: int, created_at: datetime) -> tuple[str, ...]:
+        prefix: Final = f"{owner}-{group}-"
+        write_rows(
+            'INSERT INTO "LiteLLM_ManagedObjectTable" '
+            '(id, unified_object_id, model_object_id, file_object, file_purpose, created_by, created_at, updated_at) '
+            "SELECT %s || lpad(n::text, 8, '0'), %s || n, %s || n, '{}'::jsonb, 'batch-scan-test', %s, "
+            "%s::timestamptz, NOW() FROM generate_series(0, %s::int - 1) n",
+            (prefix, prefix, prefix, owner, created_at.isoformat(), str(size)),
+        )
+        return tuple(f"{prefix}{index:08}" for index in range(size))
+
+    try:
+        old_ids: Final = insert_rows("old", count, datetime(2000, 1, 1, tzinfo=timezone.utc))
+        async with connected_proxy_prisma() as prisma:
+            logging: Final = ProxyLogging(UserApiKeyCache())
+            router: Final = Router()
+            poller: Final = CheckBatchCost(logging, prisma, router)
+            where: Final = {"created_by": owner, "batch_processed": False}
+            first: Final = await poller._find_jobs(where)
+            assert tuple(row.id for row in first) == old_ids[:MAX_OBJECTS_PER_POLL_CYCLE]
+
+            new_ids: Final = insert_rows("new", MAX_OBJECTS_PER_POLL_CYCLE, datetime.now(timezone.utc))
+            second: Final = await poller._find_jobs(where)
+            assert tuple(row.id for row in second) == old_ids[MAX_OBJECTS_PER_POLL_CYCLE : 2 * MAX_OBJECTS_PER_POLL_CYCLE]
+            insert_rows("later", MAX_OBJECTS_PER_POLL_CYCLE, datetime.now(timezone.utc))
+            last: Final = await poller._find_jobs(where)
+            assert tuple(row.id for row in last) == old_ids[2 * MAX_OBJECTS_PER_POLL_CYCLE :]
+
+            revisited: Final = await poller._find_jobs(where)
+            assert tuple(row.id for row in revisited) == old_ids[:MAX_OBJECTS_PER_POLL_CYCLE]
+            insert_rows("during-next-scan", MAX_OBJECTS_PER_POLL_CYCLE, datetime.now(timezone.utc))
+            pages: Final = tuple([await poller._find_jobs(where) for _ in range(3)])
+            assert all(len(page) <= MAX_OBJECTS_PER_POLL_CYCLE for page in pages)
+            selected: Final = frozenset(row.id for row in chain.from_iterable(pages))
+            assert frozenset(new_ids) <= selected
+
+            restarted: Final = CheckBatchCost(logging, prisma, router)
+            restart_page: Final = await restarted._find_jobs(where)
+            assert tuple(row.id for row in restart_page) == old_ids[:MAX_OBJECTS_PER_POLL_CYCLE]
+    finally:
+        write_rows('DELETE FROM "LiteLLM_ManagedObjectTable" WHERE created_by = %s', (owner,))
diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py
index 8f694465984..350606f7892 100644
--- a/tests/unit/proxy/common_utils/test_check_batch_cost.py
+++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py
@@ -9,14 +9,11 @@ ARN unified_object_id) batches with no managed unified id.
 import asyncio
 import json
 import logging
-from collections.abc import Mapping, Sequence
 from contextlib import contextmanager
 from copy import copy
-from dataclasses import dataclass
 from datetime import datetime, timezone
 from hashlib import sha256
-from itertools import chain
-from typing import TYPE_CHECKING, Final, cast
+from typing import TYPE_CHECKING, Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -26,7 +23,6 @@ from litellm.types.utils import LiteLLMBatch
 
 if TYPE_CHECKING:
     from litellm.batches.batch_utils import BatchCostUsageResult
-    from litellm_enterprise.proxy.common_utils.managed_object_attribution import ManagedObjectRow
 
 _IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id"
 _CLAIM_UNIFIED_BATCH_ID = "dW5pZmllZF9iYXRjaF9pZA=="
@@ -240,13 +236,13 @@ class TestCheckBatchCost:
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
         # First find_many (primary query) raises with a schema error; second (fallback) returns empty
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
-            side_effect=[Exception("column batch_processed does not exist"), [], []]
+            side_effect=[Exception("column batch_processed does not exist"), []]
         )
 
         await check_batch_cost_instance.check_batch_cost()
 
         calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list
-        assert len(calls) == 3
+        assert len(calls) == 2
         fallback_where = calls[1][1]["where"]
         assert "batch_processed" not in fallback_where
         assert "stale_expired" in fallback_where["status"]["not_in"]
@@ -267,7 +263,7 @@ class TestCheckBatchCost:
 
         await check_batch_cost_instance.check_batch_cost()
 
-        assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 2
+        assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1
         fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"]
         assert "batch_processed" not in fallback_where
 
@@ -4203,72 +4199,3 @@ async def test_failed_retirement_keeps_batch_available() -> None:
 
     assert row.batch_processed is False
     assert row.status == "in_progress"
-
-
-@dataclass(frozen=True, slots=True)
-class _PollingBatch:
-    id: str
-    created_at: datetime
-
-
-def _matches_poll_filter(row: _PollingBatch, where: Mapping[str, object]) -> bool:
-    for key, value in where.items():
-        if key == "OR":
-            if not any(_matches_poll_filter(row, branch) for branch in cast(Sequence[Mapping[str, object]], value)):
-                return False
-            continue
-        actual: Final = row.created_at if key == "created_at" else row.id
-        if not isinstance(value, Mapping):
-            if actual != value:
-                return False
-            continue
-        for operator, bound in cast(Mapping[str, datetime | str], value).items():
-            if isinstance(actual, datetime) and isinstance(bound, datetime):
-                if (operator == "gt" and actual <= bound) or (operator == "lte" and actual > bound):
-                    return False
-            elif isinstance(actual, str) and isinstance(bound, str) and operator == "gt" and actual <= bound:
-                return False
-    return True
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("page_size", [1, 3])
-async def test_new_batches_progress_without_waiting_for_retry_scan(
-    monkeypatch: pytest.MonkeyPatch, page_size: int
-) -> None:
-    from litellm_enterprise.proxy.common_utils import check_batch_cost as module
-
-    monkeypatch.setattr(module, "MAX_OBJECTS_PER_POLL_CYCLE", page_size)
-    backlog: Final = tuple(
-        _PollingBatch(f"old-{index:03}", datetime(2000, 1, 1, tzinfo=timezone.utc)) for index in range(page_size * 5)
-    )
-    arrivals: Final = tuple(
-        _PollingBatch(f"new-{index:03}", datetime(2020, index // page_size + 1, 1, tzinfo=timezone.utc))
-        for index in range(page_size * 3)
-    )
-
-    async def select_page(*, where: Mapping[str, object], take: int, order: object) -> tuple[_PollingBatch, ...]:
-        return tuple(row for row in (*backlog, *arrivals) if _matches_poll_filter(row, where))[:take]
-
-    prisma: Final = MagicMock()
-    prisma.db.litellm_managedobjecttable.find_many = AsyncMock(side_effect=select_page)
-    with patch.object(module, "datetime") as clock:
-        clock.now.return_value = datetime(2010, 1, 1, tzinfo=timezone.utc)
-        poller: Final = module.CheckBatchCost(MagicMock(), prisma, MagicMock())
-        first: Final = await poller._find_jobs({})
-        assert tuple(row.id for row in first) == tuple(row.id for row in backlog[:page_size])
-
-        async def poll_next_turn(turn: int) -> Sequence["ManagedObjectRow"]:
-            clock.now.return_value = datetime(2020, turn // 2 + 1, 2, tzinfo=timezone.utc)
-            return await poller._find_jobs({})
-
-        pages: Final = tuple([await poll_next_turn(turn) for turn in range(5)])
-        selected: Final = frozenset(row.id for row in chain.from_iterable(pages))
-        assert all(len(page) <= page_size for page in pages)
-        assert all(row.id in selected for row in arrivals)
-        assert all(row.id in selected for row in backlog[page_size : 3 * page_size])
-        clock.now.return_value = datetime(2021, 1, 1, tzinfo=timezone.utc)
-        restarted: Final = module.CheckBatchCost(MagicMock(), prisma, MagicMock())
-        recovered: Final = tuple([await restarted._find_jobs({}) for _ in range(10)])
-        recovered_ids: Final = frozenset(row.id for row in chain.from_iterable(recovered))
-        assert recovered_ids == frozenset(row.id for row in (*backlog, *arrivals))

From 235696d95ab1762e0b2f9d95f194368e4bd9760c Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Sun, 18 Oct 2026 00:22:56 -0700
Subject: [PATCH 7/8] fix(batch): restore arrival priority during retry scans

---
 .../proxy/common_utils/check_batch_cost.py    | 22 +++++++++++++---
 .../spend/test_batch_poll_starvation.py       | 26 ++++++++-----------
 .../common_utils/test_check_batch_cost.py     |  6 ++---
 3 files changed, 32 insertions(+), 22 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 e8addb4f586..81f4aac4885 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -111,6 +111,8 @@ class CheckBatchCost:
         self.batch_processed_support_confirmed: bool = False
         self._poll_cursor: tuple[datetime, str] | None = None
         self._poll_before: datetime | None = None
+        self._arrival_cursor: tuple[datetime, str] = (datetime.now(timezone.utc), "")
+        self._poll_arrivals: bool = False
         self._line_item_stores: Final[set[asyncio.Task[None]]] = set()
         self._line_item_store_turn: Final = asyncio.Lock()
         self._poll_idle: Final = asyncio.Event()
@@ -182,13 +184,21 @@ class CheckBatchCost:
         )
 
     async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[ManagedObjectRow]:
-        cursor: Final = self._poll_cursor
-        if cursor is None:
+        arrivals: Final = self._poll_arrivals
+        page: Final = await self._find_job_page(where, arrivals=arrivals)
+        self._poll_arrivals = not arrivals if page else arrivals
+        if page:
+            return page
+        return await self._find_job_page(where, arrivals=not arrivals)
+
+    async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]:
+        cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor
+        if not arrivals and cursor is None:
             self._poll_before = datetime.now(timezone.utc)
         page: Final = await _managed_object_table(self.prisma_client).find_many(
             where={
                 **where,
-                "created_at": {"lte": self._poll_before},
+                "created_at": {"lte": datetime.now(timezone.utc) if arrivals else self._poll_before},
                 **(
                     {
                         "OR": [
@@ -203,7 +213,11 @@ class CheckBatchCost:
             take=MAX_OBJECTS_PER_POLL_CYCLE,
             order=[{"created_at": "asc"}, {"id": "asc"}],
         )
-        self._poll_cursor = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
+        if arrivals:
+            if page:
+                self._arrival_cursor = (page[-1].created_at, page[-1].id)
+        else:
+            self._poll_cursor = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
         return page
 
     def _unprocessed_snapshot(self, job: ManagedObjectRow) -> Mapping[str, object]:
diff --git a/tests/integration/spend/test_batch_poll_starvation.py b/tests/integration/spend/test_batch_poll_starvation.py
index 5981c3bc425..f8f4898fa91 100644
--- a/tests/integration/spend/test_batch_poll_starvation.py
+++ b/tests/integration/spend/test_batch_poll_starvation.py
@@ -3,7 +3,6 @@ import os
 from datetime import datetime, timezone
 from functools import partial
 from hashlib import sha256
-from itertools import chain
 from typing import Final
 from uuid import uuid4
 
@@ -256,15 +255,14 @@ def test_uncostable_batches_do_not_starve_billable_batches(
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("tail", [0, 1])
-async def test_batch_scan_finishes_and_revisits_retries_despite_new_arrivals(tail: int) -> None:
+@pytest.mark.parametrize("count", [2 * MAX_OBJECTS_PER_POLL_CYCLE, 2 * MAX_OBJECTS_PER_POLL_CYCLE + 1, 1000])
+async def test_arrivals_progress_while_retry_scan_advances_and_wraps(count: int) -> None:
     from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
     from litellm.proxy.utils import ProxyLogging
     from litellm.router import Router
     from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
 
     owner: Final = f"batch-scan-{uuid4().hex}"
-    count: Final = 2 * MAX_OBJECTS_PER_POLL_CYCLE + tail
 
     def insert_rows(group: str, size: int, created_at: datetime) -> tuple[str, ...]:
         prefix: Final = f"{owner}-{group}-"
@@ -287,20 +285,18 @@ async def test_batch_scan_finishes_and_revisits_retries_despite_new_arrivals(tai
             first: Final = await poller._find_jobs(where)
             assert tuple(row.id for row in first) == old_ids[:MAX_OBJECTS_PER_POLL_CYCLE]
 
-            new_ids: Final = insert_rows("new", MAX_OBJECTS_PER_POLL_CYCLE, datetime.now(timezone.utc))
-            second: Final = await poller._find_jobs(where)
-            assert tuple(row.id for row in second) == old_ids[MAX_OBJECTS_PER_POLL_CYCLE : 2 * MAX_OBJECTS_PER_POLL_CYCLE]
-            insert_rows("later", MAX_OBJECTS_PER_POLL_CYCLE, datetime.now(timezone.utc))
-            last: Final = await poller._find_jobs(where)
-            assert tuple(row.id for row in last) == old_ids[2 * MAX_OBJECTS_PER_POLL_CYCLE :]
+            for start in range(MAX_OBJECTS_PER_POLL_CYCLE, count + 1, MAX_OBJECTS_PER_POLL_CYCLE):
+                new_ids: Final = insert_rows(str(start), MAX_OBJECTS_PER_POLL_CYCLE, datetime.now(timezone.utc))
+                arrivals: Final = await poller._find_jobs(where)
+                assert tuple(row.id for row in arrivals) == new_ids
+                retries: Final = await poller._find_jobs(where)
+                assert tuple(row.id for row in retries) == old_ids[start : start + MAX_OBJECTS_PER_POLL_CYCLE]
 
             revisited: Final = await poller._find_jobs(where)
             assert tuple(row.id for row in revisited) == old_ids[:MAX_OBJECTS_PER_POLL_CYCLE]
-            insert_rows("during-next-scan", MAX_OBJECTS_PER_POLL_CYCLE, datetime.now(timezone.utc))
-            pages: Final = tuple([await poller._find_jobs(where) for _ in range(3)])
-            assert all(len(page) <= MAX_OBJECTS_PER_POLL_CYCLE for page in pages)
-            selected: Final = frozenset(row.id for row in chain.from_iterable(pages))
-            assert frozenset(new_ids) <= selected
+            latest_ids: Final = insert_rows("latest", MAX_OBJECTS_PER_POLL_CYCLE, datetime.now(timezone.utc))
+            latest: Final = await poller._find_jobs(where)
+            assert tuple(row.id for row in latest) == latest_ids
 
             restarted: Final = CheckBatchCost(logging, prisma, router)
             restart_page: Final = await restarted._find_jobs(where)
diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py
index 350606f7892..683423829e5 100644
--- a/tests/unit/proxy/common_utils/test_check_batch_cost.py
+++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py
@@ -236,13 +236,13 @@ class TestCheckBatchCost:
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
         # First find_many (primary query) raises with a schema error; second (fallback) returns empty
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
-            side_effect=[Exception("column batch_processed does not exist"), []]
+            side_effect=[Exception("column batch_processed does not exist"), [], []]
         )
 
         await check_batch_cost_instance.check_batch_cost()
 
         calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list
-        assert len(calls) == 2
+        assert len(calls) == 3
         fallback_where = calls[1][1]["where"]
         assert "batch_processed" not in fallback_where
         assert "stale_expired" in fallback_where["status"]["not_in"]
@@ -263,7 +263,7 @@ class TestCheckBatchCost:
 
         await check_batch_cost_instance.check_batch_cost()
 
-        assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1
+        assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 2
         fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"]
         assert "batch_processed" not in fallback_where
 

From ae74c537076ac51caa6c4138f649167bee392d2a Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Mon, 19 Oct 2026 19:33:08 -0700
Subject: [PATCH 8/8] fix(batch): preserve recoverable jobs on unrelated
 resource errors

---
 .../proxy/common_utils/check_batch_cost.py    |   9 +-
 .../spend/test_batch_poll_starvation.py       |  44 ++++++-
 .../common_utils/test_check_batch_cost.py     | 116 ++++++++++++++++++
 3 files changed, 165 insertions(+), 4 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 81f4aac4885..7e4fb8c8247 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -4,6 +4,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
 
 import asyncio
 import json
+import re
 from collections.abc import Mapping, Sequence
 from dataclasses import replace as dataclasses_replace
 from datetime import datetime, timedelta, timezone
@@ -39,6 +40,10 @@ if TYPE_CHECKING:
     from litellm.types.utils import LiteLLMBatch
 
 
+def _error_names_resource(error: Exception, resource_id: str) -> bool:
+    return resource_id in re.findall(r"""[^\s"'`]+""", str(error))
+
+
 CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
 
 PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
@@ -333,7 +338,7 @@ class CheckBatchCost:
 
         from litellm.exceptions import NotFoundError
 
-        return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error)
+        return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, batch_id)
 
     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
@@ -354,7 +359,7 @@ class CheckBatchCost:
 
         if not output_file_id:
             return False
-        return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
+        return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id)
 
     async def _finalize_unbilled_terminal_job(
         self, job: "ManagedObjectRow", response: "LiteLLMBatch"
diff --git a/tests/integration/spend/test_batch_poll_starvation.py b/tests/integration/spend/test_batch_poll_starvation.py
index f8f4898fa91..cdecabbbc9c 100644
--- a/tests/integration/spend/test_batch_poll_starvation.py
+++ b/tests/integration/spend/test_batch_poll_starvation.py
@@ -6,13 +6,14 @@ from hashlib import sha256
 from typing import Final
 from uuid import uuid4
 
+import httpx
 import pytest
 from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value
 from integration._support.database import read_rows, write_rows
 from integration._support.prisma_clients import connected_proxy_prisma
 from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
 from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse
-from pydantic import JsonValue
+from pydantic import JsonValue, TypeAdapter
 
 from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
 
@@ -120,6 +121,36 @@ def _completed_routes(*, ready: bool = True) -> RoutedResponse:
     )
 
 
+def _batch_was_polled(upstream: ScenarioHandle) -> bool:
+    response: Final = httpx.get(f"{upstream.control_url}/__observations", timeout=15, trust_env=False)
+    response.raise_for_status()
+    observations: Final = TypeAdapter(dict[str, list[dict[str, JsonValue]]]).validate_json(response.content)
+    return any(
+        item.get("method") == "GET" and string_value(item["path"]).startswith(f"/{upstream.scenario_id}/batches/")
+        for item in observations["requests"]
+    )
+
+
+def _missing_output_routes() -> RoutedResponse:
+    return RoutedResponse(
+        content_type="application/x-routed",
+        routes={
+            **_completed_routes().routes,
+            "GET /files/file-out-$REQUEST_ID/content": JsonResponse(
+                content_type="application/json",
+                status=404,
+                body={
+                    "error": {
+                        "message": "No file found with id 'file-out-$REQUEST_ID'.",
+                        "type": "invalid_request_error",
+                        "code": "file_not_found",
+                    }
+                },
+            ),
+        },
+    )
+
+
 def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse) -> tuple[str, ScenarioHandle]:
     scenario_id: Final = f"poll-{marker}-{sha256(os.urandom(16)).hexdigest()[:12]}"
     handle: Final = register_scenario(scenario_id, routes)
@@ -196,9 +227,12 @@ def test_uncostable_batches_do_not_starve_billable_batches(
             )
             for index in range(MAX_OBJECTS_PER_POLL_CYCLE)
         )
+        missing_output_id: Final = _submitted_batch_id(
+            gateway, key, _scripted_deployment(scenario, "missing-output", _missing_output_routes())[0]
+        )
         model, upstream = _scripted_deployment(scenario, "costable", _completed_routes(ready=False))
         costable_batch_id: Final = _submitted_batch_id(gateway, key, model)
-        for batch_id in (*gone_batch_ids, costable_batch_id):
+        for batch_id in (*gone_batch_ids, missing_output_id, costable_batch_id):
             scenario.cleanups.callback(
                 write_rows,
                 'DELETE FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s',
@@ -210,6 +244,7 @@ def test_uncostable_batches_do_not_starve_billable_batches(
                 "created_at = NOW() - INTERVAL '30 days' WHERE unified_object_id = %s",
                 ("completed", costable_batch_id),
             )
+        assert eventually(lambda: _batch_was_polled(upstream), bool, seconds=120)
         register_scenario(upstream.scenario_id, _completed_routes(), control_url=upstream.control_url)
         spend_rows: Final = eventually(
             lambda: read_rows(
@@ -233,6 +268,11 @@ def test_uncostable_batches_do_not_starve_billable_batches(
             }
         ]
         assert _managed_rows((costable_batch_id,)) == [{"batch_processed": True}]
+        assert eventually(
+            lambda: _managed_rows((missing_output_id,)),
+            lambda rows: rows == [{"batch_processed": True}],
+            seconds=120,
+        ) == [{"batch_processed": True}]
         expected_blockers: Final = [{"batch_processed": not retryable}] * MAX_OBJECTS_PER_POLL_CYCLE
         blockers: Final = eventually(
             lambda: _managed_rows(gone_batch_ids), lambda rows: rows == expected_blockers, seconds=120
diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py
index 683423829e5..1704fa58342 100644
--- a/tests/unit/proxy/common_utils/test_check_batch_cost.py
+++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py
@@ -3219,6 +3219,89 @@ class TestMultiPodBatchCostClaim:
             ["batch-test-model"] if deployment_model == "openai/batch-test-model" else ["batch-test-model"] * successes
         )
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("error_source", ("batch", "output"))
+    async def test_a_404_for_another_identifier_preserves_recovery(self, error_source: str) -> None:
+        import base64
+
+        import httpx
+        import litellm
+        import respx
+
+        from litellm.integrations.custom_logger import CustomLogger
+        from litellm.types.utils import SpecialEnums
+
+        class SpendRecorder(CustomLogger):
+            def __init__(self) -> None:
+                super().__init__()
+                self.costs: tuple[float, ...] = ()
+
+            async def async_log_success_event(
+                self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
+            ) -> None:
+                if kwargs["call_type"] == "aretrieve_batch":
+                    cost: Final = kwargs["response_cost"]
+                    assert isinstance(cost, float)
+                    self.costs += (cost,)
+
+        recorder: Final = SpendRecorder()
+        litellm.logging_callback_manager.add_litellm_async_success_callback(recorder)
+        litellm.register_model(
+            {
+                "batch-recovery-model": {
+                    "litellm_provider": "openai",
+                    "mode": "chat",
+                    "input_cost_per_token_batches": 0.01,
+                    "output_cost_per_token_batches": 0.02,
+                }
+            }
+        )
+        row: Final = _FakeManagedObjectRow()
+        row.unified_object_id = (
+            base64.urlsafe_b64encode(
+                SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format("model-123", "batch-456").encode()
+            )
+            .decode()
+            .rstrip("=")
+        )
+        router: Final = self._router()
+        response: Final = router.aretrieve_batch.return_value
+        error: Final = litellm.NotFoundError(
+            message=f"No object found with id '{'batch-456' if error_source == 'batch' else _CLAIM_OUTPUT_FILE_ID}0'.",
+            model="gpt-4",
+            llm_provider="openai",
+        )
+        router.aretrieve_batch.side_effect = (error, response) if error_source == "batch" else (response, response)
+        prisma: Final = self._prisma(row, [])
+        output: Final = json.dumps(
+            {
+                "response": {
+                    "status_code": 200,
+                    "body": {
+                        "model": "batch-recovery-model",
+                        "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
+                    },
+                },
+                "error": None,
+            }
+        ).encode()
+        with respx.mock:
+            route: Final = respx.get(f"https://api.openai.com/v1/files/{_CLAIM_OUTPUT_FILE_ID}/content")
+            route.mock(
+                side_effect=error if error_source == "output" else None,
+                return_value=None if error_source == "output" else httpx.Response(200, content=output),
+            )
+            await self._instance(prisma, router).check_batch_cost()
+            assert row.batch_processed is False
+            assert recorder.costs == ()
+            route.mock(return_value=httpx.Response(200, content=output))
+            await self._instance(prisma, router).check_batch_cost()
+            await self._instance(prisma, router).check_batch_cost()
+
+        assert row.batch_processed is True
+        assert row.status == "complete"
+        assert recorder.costs == (pytest.approx(10 * 0.01 + 5 * 0.02),)
+
     @staticmethod
     def _claim_calls(prisma) -> list:
         return [
@@ -4199,3 +4282,36 @@ async def test_failed_retirement_keeps_batch_available() -> None:
 
     assert row.batch_processed is False
     assert row.status == "in_progress"
+
+
+@pytest.mark.parametrize("resource_id", ("batch-123", "file-123", "gs://bucket/job/predictions.jsonl"))
+@pytest.mark.parametrize("other_id", ("{id}4", "prefix-{id}", "{id}/child", "{id}.backup", "{id}?generation=2"))
+def test_another_resource_in_a_404_is_not_permanent_loss(resource_id: str, other_id: str) -> None:
+    from litellm.exceptions import NotFoundError
+    from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
+
+    error: Final = NotFoundError(
+        message=f"No object found with id '{other_id.format(id=resource_id)}'.", model="test", llm_provider="openai"
+    )
+    assert CheckBatchCost._is_batch_gone_at_provider(error, resource_id) is False
+    assert CheckBatchCost._is_output_file_gone_at_provider(error, resource_id) is False
+
+
+@pytest.mark.parametrize("resource_id", ("batch-123", "file-123", "gs://bucket/job/predictions.jsonl"))
+@pytest.mark.parametrize("message", ("No object found with id '{id}'.", 'File "{id}" does not exist', "{id} not found"))
+def test_a_404_naming_the_exact_resource_is_permanent_loss(resource_id: str, message: str) -> None:
+    from litellm.exceptions import NotFoundError
+    from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
+
+    error: Final = NotFoundError(message=message.format(id=resource_id), model="test", llm_provider="openai")
+    assert CheckBatchCost._is_batch_gone_at_provider(error, resource_id) is True
+    assert CheckBatchCost._is_output_file_gone_at_provider(error, resource_id) is True
+
+
+@pytest.mark.parametrize("output_file_id", (None, ""))
+def test_a_404_without_an_output_identity_remains_retryable(output_file_id: str | None) -> None:
+    from litellm.exceptions import NotFoundError
+    from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
+
+    error: Final = NotFoundError(message="Output file not found", model="test", llm_provider="openai")
+    assert CheckBatchCost._is_output_file_gone_at_provider(error, output_file_id) is False