fix(proxy): claim batch cost rows atomically so multi-pod polling can't double-bill (#37685)

Every pod and uvicorn worker schedules its own CheckBatchCost poller against the
shared managed-object table, so two of them can select the same completed batch in
one polling window and both write an aretrieve_batch spend log for it, counting
that batch's cost twice.

Claim the row with a compare-and-swap on batch_processed, and skip the batch when
another pod already holds it. The claim sits immediately before the spend log is
written rather than before the results fetch, because batch_processed is also what
blocks deletion of the files the fetch reads and what keeps an unbilled row
selectable by later poll cycles, so claiming up front would strand the spend of any
worker that died mid-fetch. A failed spend log write hands the row back.

Co-authored-by: Yassin Kortam <yassin@berri.ai>

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -255,6 +255,52 @@ class CheckBatchCost:
             "so it will no longer be polled"
         )
 
+    async def _claim_job_for_costing(self, job: "LiteLLM_ManagedObjectTable") -> bool:
+        """
+        Atomically flip batch_processed from false to true, returning whether this pod won
+        the row. Every pod and uvicorn worker schedules its own poller against the shared
+        table, so without this compare-and-swap two of them can select the same completed
+        batch in one window and both emit an aretrieve_batch spend log for it. Schemas
+        without the column can't be claimed, so they keep the pre-existing behavior.
+
+        Called immediately before the spend log is written rather than before the results
+        fetch, because batch_processed is also what holds off deletion of the files that
+        fetch reads and what keeps an unbilled row selectable by the next poll cycle.
+        """
+        if not self._has_batch_processed_column:
+            return True
+        try:
+            claimed: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
+                where={"id": job.id, "batch_processed": False},
+                data={"batch_processed": True},
+            )
+        except Exception as db_err:
+            verbose_proxy_logger.error(
+                f"CheckBatchCost: failed to claim job {job.id} for cost tracking: {db_err}"
+            )
+            return False
+        return claimed > 0
+
+    async def _release_job_claim(self, job: "LiteLLM_ManagedObjectTable") -> None:
+        """Give a claimed row back once billing it failed, so a later poll cycle retries it.
+
+        Safe to match on batch_processed=True: while this poller is active the retrieve
+        path leaves the column alone (batch_cost_poller_is_active), so a true value here
+        is always this pod's own claim.
+        """
+        if not self._has_batch_processed_column:
+            return
+        try:
+            await self.prisma_client.db.litellm_managedobjecttable.update_many(
+                where={"id": job.id, "batch_processed": True},
+                data={"batch_processed": False},
+            )
+        except Exception as db_err:
+            verbose_proxy_logger.error(
+                f"CheckBatchCost: failed to release the claim on job {job.id}, "
+                f"so its cost will not be retried: {db_err}"
+            )
+
     @staticmethod
     def _has_unified_id_without_model(job: "LiteLLM_ManagedObjectTable") -> bool:
         """A unified id that decodes but carries no model_id can never be routed."""
@@ -572,9 +618,10 @@ class CheckBatchCost:
         """
         Fetch a completed batch's results, compute cost/usage, and emit the
         aretrieve_batch spend log. Returns (model_name, llm_provider) on
-        success, None when the job can't be routed to a deployment. Raises on
-        results-fetch or cost-computation failures so the caller can leave the
-        job unprocessed and retry it on a later poll.
+        success, None when the job can't be routed to a deployment or when
+        another pod claimed it. Raises on results-fetch or cost-computation
+        failures so the caller can leave the job unprocessed and retry it on a
+        later poll.
         """
         from litellm.batches.batch_utils import (
             _get_file_content_as_dictionary,
@@ -743,12 +790,23 @@ class CheckBatchCost:
             optional_params={},
         )
 
-        await logging_obj.async_success_handler(
-            result=response,
-            batch_cost=batch_cost,
-            batch_usage=batch_usage,
-            batch_models=batch_models,
-        )
+        if not await self._claim_job_for_costing(job):
+            verbose_proxy_logger.info(
+                f"CheckBatchCost: batch {batch_id} (job {job.id}) was claimed by another pod "
+                "in this window, so its cost is already being tracked there"
+            )
+            return None
+
+        try:
+            await logging_obj.async_success_handler(
+                result=response,
+                batch_cost=batch_cost,
+                batch_usage=batch_usage,
+                batch_models=batch_models,
+            )
+        except Exception:
+            await self._release_job_claim(job)
+            raise
 
         # Record batch duration (completed_at - created_at)
         if prom_logger and response.completed_at and response.created_at:

tests/proxy_unit_tests/test_check_batch_cost.py

@@ -6,11 +6,17 @@ Vertex (raw gs:// input_file_id) and Bedrock (raw s3:// input_file_id,
 ARN unified_object_id) batches with no managed unified id.
 """
 
+import asyncio
+import json
+from contextlib import contextmanager
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from fastapi import HTTPException
 
 _IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id"
+_CLAIM_UNIFIED_BATCH_ID = "dW5pZmllZF9iYXRjaF9pZA=="
+_CLAIM_OUTPUT_FILE_ID = "file-output-123"
 
 
 def _unmanaged_vertex_file_object(
@@ -95,7 +101,7 @@ class TestCheckBatchCost:
     ):
         """_cleanup_stale_managed_objects scopes its update to file_purpose='batch' only."""
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
-            return_value=0
+            return_value=1
         )
         # Return empty so the main poll loop exits immediately
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
@@ -161,7 +167,7 @@ class TestCheckBatchCost:
         from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
 
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
-            return_value=0
+            return_value=1
         )
         mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
             return_value=[]
@@ -192,7 +198,7 @@ class TestCheckBatchCost:
         from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
 
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
-            return_value=0
+            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(
@@ -221,7 +227,7 @@ class TestCheckBatchCost:
         from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
 
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
-            return_value=0
+            return_value=1
         )
         # Simulate column already known absent from a previous cycle
         check_batch_cost_instance._has_batch_processed_column = False
@@ -254,7 +260,7 @@ class TestCheckBatchCost:
         from unittest.mock import patch
 
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
-            return_value=0
+            return_value=1
         )
         mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
         mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
@@ -563,7 +569,7 @@ class TestCheckBatchCost:
         from unittest.mock import patch
 
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
-            return_value=0
+            return_value=1
         )
         mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
         mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
@@ -679,7 +685,7 @@ class TestCheckBatchCost:
         import litellm
         from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger
 
-        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0)
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
         mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
         mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
 
@@ -801,7 +807,7 @@ class TestCheckBatchCost:
         from unittest.mock import patch
 
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
-            return_value=0
+            return_value=1
         )
         mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
         mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
@@ -869,7 +875,7 @@ class TestCheckBatchCost:
         import base64
 
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
-            return_value=0
+            return_value=1
         )
         mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
         mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
@@ -944,7 +950,7 @@ class TestCheckBatchCost:
         ).decode()
 
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
-            return_value=0
+            return_value=1
         )
         mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
         mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
@@ -1044,7 +1050,7 @@ class TestCheckBatchCost:
         from unittest.mock import patch
 
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
-            return_value=0
+            return_value=1
         )
         mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
         mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
@@ -1111,7 +1117,7 @@ class TestCheckBatchCost:
         from unittest.mock import patch
 
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
-            return_value=0
+            return_value=1
         )
         mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
 
@@ -1168,7 +1174,7 @@ class TestCheckBatchCost:
         from unittest.mock import patch
 
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
-            return_value=0
+            return_value=1
         )
         mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
         mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
@@ -1284,7 +1290,7 @@ class TestCheckBatchCost:
         from litellm.exceptions import NotFoundError
 
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
-            return_value=0
+            return_value=1
         )
         mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
         mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
@@ -1355,7 +1361,7 @@ class TestCheckBatchCost:
         through the proxy, causing API_KEY errors when clients call GET /files/{id}/content.
         """
         mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
-            return_value=0
+            return_value=1
         )
         mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
         mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
@@ -1672,7 +1678,7 @@ class TestUnmanagedVertexRouting:
         prisma = instance.prisma_client
         prisma.db = MagicMock()
         prisma.db.litellm_managedobjecttable = MagicMock()
-        prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0)
+        prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
         prisma.db.litellm_managedobjecttable.update = AsyncMock()
         prisma.db.litellm_managedobjecttable.find_many = AsyncMock(
             return_value=[self._job()]
@@ -1902,7 +1908,7 @@ class TestUnmanagedBedrockRouting:
         prisma = instance.prisma_client
         prisma.db = MagicMock()
         prisma.db.litellm_managedobjecttable = MagicMock()
-        prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0)
+        prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
         prisma.db.litellm_managedobjecttable.update = AsyncMock()
         prisma.db.litellm_managedobjecttable.find_many = AsyncMock(
             return_value=[self._job()]
@@ -2577,3 +2583,353 @@ class TestPollPageStarvation:
         await self._instance(prisma, llm_router).check_batch_cost()
 
         prisma.db.litellm_managedobjecttable.update.assert_not_awaited()
+
+class _FakeManagedObjectRow:
+    """One managed batch row the provider has finished but nothing has costed yet."""
+
+    def __init__(self):
+        self.id = "job-claim-1"
+        self.unified_object_id = _CLAIM_UNIFIED_BATCH_ID
+        self.model_object_id = "batch-456"
+        self.file_purpose = "batch"
+        self.status = "in_progress"
+        self.batch_processed = False
+        self.created_by = "user-1"
+        self.team_id = None
+        self.api_key = None
+        self.request_tags = None
+        self.created_at = 1700000000
+        self.file_object = json.dumps(
+            {"id": "batch-456", "status": "in_progress", "input_file_id": "file-input-1",
+             "output_file_id": _CLAIM_OUTPUT_FILE_ID}
+        )
+
+
+class _FakeManagedObjectTable:
+    """A LiteLLM_ManagedObjectTable double backed by one real, mutable row.
+
+    It honours the batch_processed and status filters, so the poller's compare-and-swap
+    and the managed-files deletion guard both read the same state a shared Postgres row
+    would give them. Staleness sweeps (the only queries scoped by created_at) never match.
+    """
+
+    def __init__(self, row: _FakeManagedObjectRow, journal: list):
+        self.row = row
+        self.journal = journal
+        self.update_many = AsyncMock(side_effect=self._update_many)
+        self.update = AsyncMock(side_effect=self._update)
+        self.find_many = AsyncMock(side_effect=self._find_many)
+        self.find_first = AsyncMock(return_value=None)
+
+    def _matches(self, where: dict) -> bool:
+        for key, value in where.items():
+            if key == "created_at":
+                return False
+            if key == "status":
+                if self.row.status in value.get("not_in", []):
+                    return False
+                if "in" in value and self.row.status not in value["in"]:
+                    return False
+            elif getattr(self.row, key) != value:
+                return False
+        return True
+
+    async def _update_many(self, *, where: dict, data: dict) -> int:
+        if not self._matches(where):
+            return 0
+        if "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)
+        return 1
+
+    async def _update(self, *, where: dict, data: dict) -> None:
+        self.journal.append("finalize")
+        for key, value in data.items():
+            setattr(self.row, key, value)
+
+    async def _find_many(self, *, where: dict, take=None, order=None) -> list:
+        return [self.row] if self._matches(where) else []
+
+
+class TestMultiPodBatchCostClaim:
+    """LIT-4827 regression: every pod and uvicorn worker schedules its own poller against
+    the shared LiteLLM_ManagedObjectTable, so a completed batch must be claimed atomically
+    before its cost is logged. Without the claim two pods select the same row in one window
+    and both write an aretrieve_batch spend log for it, double counting the spend.
+
+    The claim sits immediately before the spend-log write rather than before the results
+    fetch, because batch_processed is also what keeps an unbilled row selectable by later
+    poll cycles and what blocks deletion of the files the fetch reads."""
+
+    @staticmethod
+    def _instance(prisma, llm_router):
+        from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
+
+        proxy_logging_obj = MagicMock()
+        proxy_logging_obj.get_proxy_hook.return_value = None
+        return CheckBatchCost(
+            proxy_logging_obj=proxy_logging_obj,
+            prisma_client=prisma,
+            llm_router=llm_router,
+        )
+
+    @staticmethod
+    def _prisma(row: _FakeManagedObjectRow, journal: list):
+        prisma = MagicMock()
+        prisma.db.litellm_managedobjecttable = _FakeManagedObjectTable(row, journal)
+        prisma.db.litellm_managedfiletable.find_many = AsyncMock(return_value=[])
+        prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None)
+        prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
+        prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None)
+        return prisma
+
+    @staticmethod
+    def _router():
+        response = MagicMock()
+        response.status = "completed"
+        response.output_file_id = _CLAIM_OUTPUT_FILE_ID
+        response.error_file_id = None
+        response.created_at = 1
+        response.completed_at = 2
+        response.model_dump_json.return_value = '{"id":"batch-456","status":"completed"}'
+
+        deployment = MagicMock()
+        deployment.litellm_params.custom_llm_provider = "openai"
+        deployment.litellm_params.model = "gpt-4"
+        deployment.model_info.model_dump.return_value = {}
+
+        router = MagicMock()
+        router.aretrieve_batch = AsyncMock(return_value=response)
+        router.get_deployment_credentials_with_provider = MagicMock(
+            return_value={"api_key": "sk-test"}
+        )
+        router.get_deployment = MagicMock(return_value=deployment)
+        return router
+
+    @staticmethod
+    @contextmanager
+    def _billing_patches(journal: list, during_fetch=None, bill_error=None):
+        """Patch the cost path a batch runs through, journalling the results fetch and the
+        spend-log write. during_fetch runs while the output file is being read, which is
+        the window an interrupted worker or a concurrent file deletion lands in."""
+        file_content = MagicMock()
+        file_content.content = b'{"id":"req-1"}'
+
+        async def _afile_content(**kwargs):
+            journal.append("fetch")
+            if during_fetch is not None:
+                await during_fetch()
+            return file_content
+
+        async def _bill(**kwargs):
+            journal.append("bill")
+            if bill_error is not None:
+                raise bill_error
+
+        def _is_b64(file_id):
+            if file_id == _CLAIM_UNIFIED_BATCH_ID:
+                return "llm_model_id,model-123;llm_batch_id,batch-456;"
+            return False
+
+        logging_obj = MagicMock()
+        logging_obj.async_success_handler = AsyncMock(side_effect=_bill)
+
+        with (
+            patch(_IS_B64, side_effect=_is_b64),
+            patch(
+                "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id",
+                return_value="model-123",
+            ),
+            patch(
+                "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id",
+                return_value="batch-456",
+            ),
+            patch("litellm.files.main.afile_content", new=AsyncMock(side_effect=_afile_content)),
+            patch(
+                "litellm.batches.batch_utils._get_file_content_as_dictionary",
+                return_value=[{"id": "req-1"}],
+            ),
+            patch(
+                "litellm.batches.batch_utils.calculate_batch_cost_and_usage",
+                new_callable=AsyncMock,
+                return_value=(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"]),
+            ),
+            patch(
+                "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
+                return_value=("gpt-4", "openai", None, None),
+            ),
+            patch("litellm.litellm_core_utils.litellm_logging.Logging", return_value=logging_obj),
+        ):
+            yield logging_obj
+
+    @staticmethod
+    def _claim_calls(prisma) -> list:
+        return [
+            call.kwargs
+            for call in prisma.db.litellm_managedobjecttable.update_many.call_args_list
+            if "id" in call.kwargs["where"]
+        ]
+
+    @staticmethod
+    async def _run_deletion_guard(prisma, file_id: str) -> None:
+        """Run the real managed-files deletion guard against the row the poller is costing."""
+        from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
+
+        cache = MagicMock()
+        cache.async_get_cache = AsyncMock(return_value=None)
+        cache.async_set_cache = AsyncMock()
+        guard = _PROXY_LiteLLMManagedFiles(internal_usage_cache=cache, prisma_client=prisma)
+
+        scheduler = MagicMock()
+        scheduler.get_job.return_value = MagicMock()
+        with patch("litellm.proxy.proxy_server.scheduler", scheduler):
+            await guard._check_file_deletion_allowed(file_id)
+
+    @pytest.mark.asyncio
+    async def test_winning_pod_claims_the_row_between_fetching_and_billing(self):
+        """The claim flips batch_processed false -> true after the results are in hand and
+        before the spend log is written, so a concurrent pod's claim finds no matching row."""
+        row = _FakeManagedObjectRow()
+        journal = []
+        prisma = self._prisma(row, journal)
+
+        with self._billing_patches(journal) as logging_obj:
+            await self._instance(prisma, self._router()).check_batch_cost()
+
+        assert journal == ["fetch", "claim", "bill", "finalize"]
+        assert self._claim_calls(prisma) == [
+            {
+                "where": {"id": "job-claim-1", "batch_processed": False},
+                "data": {"batch_processed": True},
+            }
+        ]
+        logging_obj.async_success_handler.assert_awaited_once()
+        assert row.batch_processed is True
+
+    @pytest.mark.asyncio
+    async def test_a_pod_that_loses_the_claim_after_fetching_does_not_bill(self):
+        """Both pods select the row and fetch its results in the same window. The one whose
+        compare-and-swap finds the row already taken must not write a second spend log."""
+        row = _FakeManagedObjectRow()
+        journal = []
+        prisma = self._prisma(row, journal)
+
+        async def _other_pod_wins_the_row():
+            row.batch_processed = True
+
+        with self._billing_patches(journal, during_fetch=_other_pod_wins_the_row) as logging_obj:
+            await self._instance(prisma, self._router()).check_batch_cost()
+
+        assert journal == ["fetch"]
+        logging_obj.async_success_handler.assert_not_awaited()
+        assert self._claim_calls(prisma) == [
+            {
+                "where": {"id": "job-claim-1", "batch_processed": False},
+                "data": {"batch_processed": True},
+            }
+        ]
+
+    @pytest.mark.asyncio
+    async def test_a_failed_spend_log_write_releases_the_claim(self):
+        """A transient failure while billing a claimed batch must hand the row back, or its
+        spend is silently lost instead of being retried on the next cycle."""
+        row = _FakeManagedObjectRow()
+        journal = []
+        prisma = self._prisma(row, journal)
+
+        with self._billing_patches(journal, bill_error=Exception("spend log write failed")):
+            await self._instance(prisma, self._router()).check_batch_cost()
+
+        assert journal == ["fetch", "claim", "bill", "release"]
+        assert row.batch_processed is False
+        assert self._claim_calls(prisma)[-1] == {
+            "where": {"id": "job-claim-1", "batch_processed": True},
+            "data": {"batch_processed": False},
+        }
+
+    @pytest.mark.asyncio
+    async def test_a_worker_interrupted_mid_costing_leaves_the_batch_billable(self):
+        """A pod killed while reading a batch's results must leave the row for a later
+        cycle. Claiming before the fetch marked the batch processed for good, so the pod
+        that died took that batch's spend with it and no other pod ever selected it."""
+        row = _FakeManagedObjectRow()
+        journal = []
+        prisma = self._prisma(row, journal)
+        reached_fetch = asyncio.Event()
+
+        async def _never_returns():
+            reached_fetch.set()
+            await asyncio.Event().wait()
+
+        with self._billing_patches(journal, during_fetch=_never_returns) as logging_obj:
+            interrupted = asyncio.create_task(
+                self._instance(prisma, self._router()).check_batch_cost()
+            )
+            await asyncio.wait_for(reached_fetch.wait(), timeout=5)
+            assert row.batch_processed is False, "an in-flight costing must not mark the row processed"
+            interrupted.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await interrupted
+
+        assert journal == ["fetch"]
+        logging_obj.async_success_handler.assert_not_awaited()
+
+        survivor_journal = []
+        survivor_prisma = self._prisma(row, survivor_journal)
+        with self._billing_patches(survivor_journal) as survivor_logging:
+            await self._instance(survivor_prisma, self._router()).check_batch_cost()
+
+        assert survivor_journal == ["fetch", "claim", "bill", "finalize"]
+        survivor_logging.async_success_handler.assert_awaited_once()
+        assert row.batch_processed is True
+
+    @pytest.mark.asyncio
+    async def test_costing_in_flight_keeps_the_referenced_file_undeletable(self):
+        """The deletion guard only holds files whose batch still has batch_processed false,
+        so claiming the row before the fetch let a concurrent delete remove the very output
+        file the in-flight costing was about to read."""
+        row = _FakeManagedObjectRow()
+        journal = []
+        prisma = self._prisma(row, journal)
+        reached_fetch = asyncio.Event()
+        finish_fetch = asyncio.Event()
+
+        async def _wait_for_the_delete_attempt():
+            reached_fetch.set()
+            await finish_fetch.wait()
+
+        with self._billing_patches(journal, during_fetch=_wait_for_the_delete_attempt):
+            costing = asyncio.create_task(
+                self._instance(prisma, self._router()).check_batch_cost()
+            )
+            await asyncio.wait_for(reached_fetch.wait(), timeout=5)
+
+            with pytest.raises(HTTPException) as blocked:
+                await self._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID)
+            assert blocked.value.status_code == 400
+            assert _CLAIM_OUTPUT_FILE_ID in blocked.value.detail
+
+            finish_fetch.set()
+            await asyncio.wait_for(costing, timeout=5)
+
+        assert journal == ["fetch", "claim", "bill", "finalize"]
+        assert row.batch_processed is True
+        await self._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID)
+
+    @pytest.mark.asyncio
+    async def test_schema_without_batch_processed_still_bills(self):
+        """Older schemas have no column to claim, so they keep the pre-fix behavior instead
+        of losing every batch's cost."""
+        row = _FakeManagedObjectRow()
+        journal = []
+        prisma = self._prisma(row, journal)
+        instance = self._instance(prisma, self._router())
+        instance._has_batch_processed_column = False
+
+        with self._billing_patches(journal) as logging_obj:
+            await instance.check_batch_cost()
+
+        assert self._claim_calls(prisma) == []
+        assert journal == ["fetch", "bill", "finalize"]
+        logging_obj.async_success_handler.assert_awaited_once()

tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py

@@ -192,6 +192,7 @@ async def test_check_batch_cost_should_call_afile_content_directly_with_credenti
         return_value=[mock_job]
     )
     mock_prisma.db.litellm_managedobjecttable.update = AsyncMock()
+    mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
 
     # Mock proxy_logging_obj — should NOT be called for file content
     mock_proxy_logging = MagicMock()