fix(batches): only hand accounting to the poller once it can mark batches done

The handoff asked whether the poller was running, when what matters is whether it
will actually account for the batch. Those differ on a schema without the
batch_processed column: the poller cannot filter on it, so it falls back to a
query that excludes complete and completed rows, and it cannot set it either. A
caller retrieving a provider-completed batch before the poller saw it therefore
suppressed inline accounting, then marked the row complete, and the fallback query
could never find it again. Nobody accounted for that batch, so its cost escaped
the caller's budget entirely.

The poller now publishes batch_processed_support_confirmed, set only once a
filtered query has actually succeeded, and the handoff requires it. Defaulting to
unconfirmed keeps accounting on the retrieve path in exactly the cases the poller
would drop the batch, including the window before the poller's first cycle. All
four combinations account exactly once: unconfirmed leaves the retrieve
accounting and setting the marker, whether or not the column exists, and
confirmed is only reachable when the column is present, where the poller accounts
and sets it.

A scheduler that hands back something other than a bound method leaves no poller
to interrogate, which reads as unconfirmed rather than as working.

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -51,6 +51,7 @@ class CheckBatchCost:
         # Cached after the first poll cycle. Once we know the column is absent we skip
         # the guaranteed-failing primary query on every subsequent cycle.
         self._has_batch_processed_column: bool = True
+        self.batch_processed_support_confirmed: bool = False
 
     async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> Dict[str, Any]:
         """
@@ -722,6 +723,7 @@ class CheckBatchCost:
                     take=MAX_OBJECTS_PER_POLL_CYCLE,
                     order={"created_at": "asc"},
                 )
+                self.batch_processed_support_confirmed = True
             except Exception as query_err:
                 if "batch_processed" not in str(query_err).lower() and "unknown column" not in str(query_err).lower() and "does not exist" not in str(query_err).lower():
                     raise

litellm/proxy/batches_endpoints/endpoints.py

@@ -498,10 +498,11 @@ async def retrieve_batch(
             )
 
         if unified_batch_id and batch_cost_poller_is_active():
-            data["litellm_metadata"] = {
-                **(data.get("litellm_metadata") or {}),
-                "batch_ignore_default_logging": True,
-            }
+            litellm_metadata = data.get("litellm_metadata")
+            if not isinstance(litellm_metadata, dict):
+                litellm_metadata = {}  # mutable-ok: the suppression flag must live inside litellm_metadata for the success handler to read it, and this request carried no mapping to extend
+                data["litellm_metadata"] = litellm_metadata
+            litellm_metadata["batch_ignore_default_logging"] = True
 
         # Retrieve from provider (for non-terminal states or if DB lookup failed)
         # SCENARIO 1: Batch ID is encoded with model info

litellm/proxy/openai_files_endpoints/common_utils.py

@@ -1233,11 +1233,16 @@ async def get_batch_from_database(
 
 def batch_cost_poller_is_active() -> bool:
     """
-    Whether the CheckBatchCost poller is running and will therefore account for a
-    managed batch's cost itself.
+    Whether the CheckBatchCost poller will account for a managed batch's cost itself.
 
-    False whenever the poller cannot be relied on: polling disabled by config, or the
-    job absent from the scheduler because the enterprise import failed.
+    False whenever the poller cannot be relied on: polling disabled by config, the job
+    absent from the scheduler because the enterprise import failed, or the poller not
+    yet having confirmed that the batch_processed column exists. That last condition
+    matters because the poller needs the column both to find outstanding batches and to
+    mark them accounted; without it the poller falls back to a query that excludes
+    terminal statuses, so a batch the retrieve path has already marked complete becomes
+    invisible to it. Defaulting to False until the poller confirms support keeps the
+    retrieve path accounting in exactly the cases the poller would drop the batch.
     """
     from litellm.constants import PROXY_BATCH_POLLING_ENABLED
 
@@ -1249,7 +1254,11 @@ def batch_cost_poller_is_active() -> bool:
         scheduler = getattr(proxy_server_module, "scheduler", None)
         if scheduler is None:
             return False
-        return scheduler.get_job("check_batch_cost_job") is not None
+        job = scheduler.get_job("check_batch_cost_job")
+        if job is None:
+            return False
+        poller = getattr(getattr(job, "func", None), "__self__", None)
+        return getattr(poller, "batch_processed_support_confirmed", False) is True
     except Exception:  # noqa: BLE001  # scheduler backends raise varied types from get_job; an unreadable scheduler means the poller cannot be relied on
         return False
 

tests/proxy_unit_tests/test_check_batch_cost.py

@@ -143,6 +143,11 @@ class TestCheckBatchCost:
         assert "complete" not in not_in
         assert "completed" not in not_in
         assert find_call[1]["where"]["batch_processed"] is False
+        # A successful filtered query is the only proof the column exists. The retrieve
+        # path reads this to decide whether handing accounting to the poller is safe:
+        # without the column the poller's fallback query excludes complete/completed, so
+        # a batch already marked complete would never be accounted by anyone.
+        assert check_batch_cost_instance.batch_processed_support_confirmed is True
 
     @pytest.mark.asyncio
     async def test_fallback_query_used_when_batch_processed_missing(
@@ -171,6 +176,7 @@ class TestCheckBatchCost:
         assert calls[1][1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE
         # Column absence is now cached — next call should go straight to fallback
         assert check_batch_cost_instance._has_batch_processed_column is False
+        assert check_batch_cost_instance.batch_processed_support_confirmed is False
 
     @pytest.mark.asyncio
     async def test_column_absence_cached_across_cycles(

tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py

@@ -106,19 +106,44 @@ class _FakeScheduler:
         return self._job
 
 
+class _FakePoller:
+    def __init__(self, confirmed):
+        self.batch_processed_support_confirmed = confirmed
+
+    def check_batch_cost(self):
+        return None
+
+
+def _job_for(poller):
+    if poller is None:
+        return None
+    job = MagicMock()
+    job.func = poller.check_batch_cost
+    return job
+
+
 @pytest.mark.parametrize(
     "polling_enabled, job, expected",
     [
-        (True, object(), True),
+        (True, _job_for(_FakePoller(confirmed=True)), True),
+        (True, _job_for(_FakePoller(confirmed=False)), False),
         (True, None, False),
-        (False, object(), False),
+        (False, _job_for(_FakePoller(confirmed=True)), False),
+    ],
+    ids=[
+        "poller-running-and-column-confirmed",
+        "poller-running-but-column-unconfirmed",
+        "job-absent-enterprise-import-failed",
+        "polling-disabled-by-config",
     ],
-    ids=["poller-running", "job-absent-enterprise-import-failed", "polling-disabled-by-config"],
 )
 def test_batch_cost_poller_is_active(monkeypatch, polling_enabled, job, expected):
     """The predicate must only claim the poller when it can actually be relied on, so a
-    proxy with polling switched off or without the enterprise job keeps accounting for
-    batch cost on the retrieve path."""
+    proxy with polling switched off, without the enterprise job, or whose poller has not
+    confirmed batch_processed support keeps accounting for batch cost on the retrieve
+    path. The unconfirmed case is the one that matters for legacy schemas: without the
+    column the poller falls back to a query excluding terminal statuses, so a batch the
+    retrieve path already marked complete would never be accounted by anyone."""
     import litellm.constants
     import litellm.proxy.proxy_server as proxy_server_module
     from litellm.proxy.openai_files_endpoints.common_utils import (
@@ -206,3 +231,98 @@ async def test_retrieving_a_completed_batch_still_marks_processed_without_a_cost
 
     assert data["batch_processed"] is True
     assert data["status"] == "complete"
+
+
+def test_batch_cost_poller_is_active_is_false_when_the_job_has_no_bound_poller(monkeypatch):
+    """A scheduler that hands back a plain function rather than a bound method leaves no
+    poller to interrogate, so the predicate stays conservative instead of assuming the
+    column is supported."""
+    import litellm.constants
+    import litellm.proxy.proxy_server as proxy_server_module
+    from litellm.proxy.openai_files_endpoints.common_utils import (
+        batch_cost_poller_is_active,
+    )
+
+    def unbound_check_batch_cost():
+        return None
+
+    job = MagicMock()
+    job.func = unbound_check_batch_cost
+
+    monkeypatch.setattr(litellm.constants, "PROXY_BATCH_POLLING_ENABLED", True, raising=False)
+    monkeypatch.setattr(proxy_server_module, "scheduler", _FakeScheduler(job), raising=False)
+
+    assert batch_cost_poller_is_active() is False
+
+
+def test_batch_cost_poller_is_active_is_false_when_get_job_raises(monkeypatch):
+    """Scheduler backends raise varied types; an unreadable scheduler must not be read as
+    a working poller."""
+    import litellm.constants
+    import litellm.proxy.proxy_server as proxy_server_module
+    from litellm.proxy.openai_files_endpoints.common_utils import (
+        batch_cost_poller_is_active,
+    )
+
+    class _ExplodingScheduler:
+        def get_job(self, job_id):
+            raise RuntimeError("scheduler not started")
+
+    monkeypatch.setattr(litellm.constants, "PROXY_BATCH_POLLING_ENABLED", True, raising=False)
+    monkeypatch.setattr(proxy_server_module, "scheduler", _ExplodingScheduler(), raising=False)
+
+    assert batch_cost_poller_is_active() is False
+
+
+
+@pytest.mark.asyncio
+async def test_retrieving_a_batch_whose_status_is_unchanged_writes_nothing(monkeypatch):
+    """A caller polling an already-complete batch must not write at all, so repeated polls
+    cannot flip batch_processed or disturb whichever component owns accounting."""
+    import litellm.proxy.openai_files_endpoints.common_utils as cu
+
+    monkeypatch.setattr(cu, "batch_cost_poller_is_active", lambda: False)
+    monkeypatch.setattr(cu, "ensure_batch_response_managed_file_ids", AsyncMock())
+
+    prisma_client = MagicMock()
+    update_mock = AsyncMock()
+    prisma_client.db.litellm_managedobjecttable.update = update_mock
+
+    db_batch_object = MagicMock()
+    db_batch_object.status = "completed"
+
+    await cu.update_batch_in_database(
+        batch_id="unified-batch-id",
+        unified_batch_id="unified-batch-id",
+        response=_completed_batch(),
+        managed_files_obj=MagicMock(),
+        prisma_client=prisma_client,
+        verbose_proxy_logger=MagicMock(),
+        db_batch_object=db_batch_object,
+        operation="retrieve",
+    )
+
+    update_mock.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_update_batch_in_database_is_a_noop_for_unmanaged_batches(monkeypatch):
+    """Batches with no managed object row have neither the flag nor a poller queue entry, so
+    this path must leave them alone entirely."""
+    import litellm.proxy.openai_files_endpoints.common_utils as cu
+
+    prisma_client = MagicMock()
+    update_mock = AsyncMock()
+    prisma_client.db.litellm_managedobjecttable.update = update_mock
+
+    await cu.update_batch_in_database(
+        batch_id="batch-raw-xyz",
+        unified_batch_id=False,
+        response=_completed_batch(),
+        managed_files_obj=MagicMock(),
+        prisma_client=prisma_client,
+        verbose_proxy_logger=MagicMock(),
+        operation="retrieve",
+    )
+
+    update_mock.assert_not_awaited()