Merge 9995e2fd99 into 0e7c4908bb

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, ...]] = (
@@ -109,6 +114,10 @@ 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._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()
@@ -161,29 +170,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 +186,52 @@ class CheckBatchCost:
                     ]
                 },
             },
-            take=MAX_OBJECTS_PER_POLL_CYCLE,
-            order={"created_at": "asc"},
         )
 
+    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:
+            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},
+                **(
+                    {
+                        "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"}],
+        )
+        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]:
+        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 +239,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:
@@ -312,13 +338,13 @@ 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
         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
@@ -333,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"
@@ -355,18 +381,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(
@@ -1047,6 +1073,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}")
@@ -1071,7 +1098,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,
@@ -1084,8 +1111,6 @@ class CheckBatchCost:
                             ]
                         },
                     },
-                    take=MAX_OBJECTS_PER_POLL_CYCLE,
-                    order={"created_at": "asc"},
                 )
                 self.batch_processed_support_confirmed = True
             except Exception as query_err:

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: ...
 

tests/integration/spend/test_batch_poll_starvation.py

@@ -1,14 +1,19 @@
 import json
 import os
+from datetime import datetime, timezone
+from functools import partial
 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
-from integration._support.upstream import delete_scenario, register_scenario
+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
 
@@ -59,17 +64,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 +85,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 +111,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 +121,37 @@ def _completed_routes() -> RoutedResponse:
     )
 
 
-def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse) -> str:
+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)
     scenario.cleanups.callback(delete_scenario, handle)
@@ -133,7 +169,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 +209,43 @@ 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())
+        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, missing_output_id, 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),
+            )
+        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(
                 'SELECT call_type, status, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" '
@@ -209,4 +268,78 @@ 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
+        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
+        )
+        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)}]
+
+
+@pytest.mark.asyncio
+@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}"
+
+    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]
+
+            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]
+            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)
+            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,))

tests/unit/proxy/common_utils/test_check_batch_cost.py

@@ -11,7 +11,8 @@ import json
 import logging
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
-from datetime import datetime
+from copy import copy
+from datetime import datetime, timezone
 from hashlib import sha256
 from queue import SimpleQueue
 from types import SimpleNamespace
@@ -217,7 +218,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
@@ -240,13 +241,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"]
@@ -267,8 +268,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
 
@@ -744,7 +744,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"]
@@ -877,7 +877,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"
         )
 
@@ -979,10 +979,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"
@@ -1085,9 +1085,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
@@ -1152,10 +1156,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"
@@ -1371,7 +1375,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, (
@@ -1544,10 +1548,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
 
@@ -1942,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"
@@ -2159,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"
@@ -2740,9 +2744,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
@@ -2770,8 +2775,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):
@@ -2830,24 +2835,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):
@@ -2885,7 +2873,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"
@@ -2933,7 +2925,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",
@@ -2963,7 +2956,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
@@ -2976,7 +2971,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 "updated_at" in where:
+            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)
@@ -3224,6 +3221,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 [
@@ -3696,12 +3776,17 @@ 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: pytest.MonkeyPatch, database_error: bool
+    ) -> None:
         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,))
 
@@ -4111,6 +4196,127 @@ class TestCreatorTurnedOffMessageLogging:
         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"
+
+
+@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
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "kind", ["sync", "async", "custom", "overlapping", "general", "mixed", "sync_only", "executor_shutdown"]