fix(batches): account a managed batch's cost exactly once

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -52,6 +52,33 @@ def __init__(
         # 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
+
+    @staticmethod
+    def _is_missing_batch_processed_column_error(err: Exception) -> bool:
+        message: Final = str(err).lower()
+        return "batch_processed" in message or "unknown column" in message or "does not exist" in message
+
+    async def confirm_batch_processed_support(self) -> None:
+        """
+        Probe the batch_processed column before the proxy serves traffic, so the retrieve
+        path never sees an unconfirmed poller on a schema that has the column and accounts
+        inline for a batch the first poll cycle then accounts again.
+        """
+        try:
+            await self.prisma_client.db.litellm_managedobjecttable.find_first(
+                where={"file_purpose": "batch", "batch_processed": False}
+            )
+        except Exception as probe_err:
+            if not self._is_missing_batch_processed_column_error(probe_err):
+                verbose_proxy_logger.debug(
+                    f"CheckBatchCost: batch_processed probe failed, the poll cycle will confirm support: {probe_err}"
+                )
+                return
+            self._has_batch_processed_column = False
+            verbose_proxy_logger.warning("CheckBatchCost: batch_processed column not found, querying without it")
+            return
+        self.batch_processed_support_confirmed = True
 
     async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> Dict[str, Any]:
         """
@@ -724,8 +751,9 @@ async def check_batch_cost(self):
                     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():
+                if not self._is_missing_batch_processed_column_error(query_err):
                     raise
                 # Permanent schema gap — cache the result so future cycles skip straight to fallback
                 self._has_batch_processed_column = False

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -54,6 +54,9 @@
     normalize_mime_type_for_provider,
     resolve_managed_output_file_model_name,
 )
+from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
+    request_tags_from_metadata,
+)
 from litellm.types.llms.openai import (  # pyright: ignore[reportAttributeAccessIssue]
     AllMessageValues,
     AsyncCursorPage,
@@ -1146,6 +1149,7 @@ async def async_post_call_success_hook(
             ## Check if unified_file_id is in the response
             unified_file_id = response._hidden_params.get("unified_file_id")  # managed file id
             unified_batch_id = response._hidden_params.get("unified_batch_id")  # managed batch id
+            is_batch_create: Final = unified_file_id is not None
             model_id = cast(Optional[str], response._hidden_params.get("model_id"))
             model_name = cast(Optional[str], response._hidden_params.get("model_name"))
 
@@ -1216,13 +1220,16 @@ async def async_post_call_success_hook(
                             model_mappings={model_id: provider_file_id},
                             user_api_key_dict=user_api_key_dict,
                         )
+            request_metadata: Final = data.get("litellm_metadata")
             await self.store_unified_object_id(
                 unified_object_id=response.id,
                 file_object=response,
                 litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
                 model_object_id=original_response_id,
                 file_purpose="batch",
                 user_api_key_dict=user_api_key_dict,
+                request_tags=request_tags_from_metadata(request_metadata if isinstance(request_metadata, dict) else {}),
+                persist_attribution=is_batch_create,
             )
 
             # Only record batch creation metric on actual create (not retrieve/cancel).

litellm/proxy/batches_endpoints/endpoints.py

@@ -25,6 +25,7 @@
     _is_base64_encoded_unified_file_id,
     add_internal_model_credentials,
     apply_team_provider_credentials,
+    batch_cost_poller_is_active,
     decode_model_from_file_id,
     encode_batch_response_ids,
     encode_file_id_with_model,
@@ -497,6 +498,14 @@ async def retrieve_batch(
                 "Batch %s is in non-terminal state %s, syncing with provider", batch_id, response.status
             )
 
+        poller_owns_accounting: Final = bool(unified_batch_id) and batch_cost_poller_is_active()
+        if poller_owns_accounting:
+            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
         if model_from_id is not None:
@@ -581,6 +590,7 @@ async def retrieve_batch(
             verbose_proxy_logger=verbose_proxy_logger,
             db_batch_object=db_batch_object,
             operation="retrieve",
+            poller_owns_accounting=poller_owns_accounting,
         )
 
         ### CALL HOOKS ### - modify outgoing data

litellm/proxy/openai_files_endpoints/common_utils.py

@@ -1256,6 +1256,38 @@ async def get_batch_from_database(
         return None, None
 
 
+def batch_cost_poller_is_active() -> bool:
+    """
+    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, 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
+
+    if not PROXY_BATCH_POLLING_ENABLED:
+        return False
+    try:
+        import litellm.proxy.proxy_server as proxy_server_module
+
+        scheduler = getattr(proxy_server_module, "scheduler", None)
+        if scheduler is None:
+            return False
+        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
+
+
 async def update_batch_in_database(
     batch_id: str,
     unified_batch_id: str | Literal[False],
@@ -1266,6 +1298,7 @@ async def update_batch_in_database(
     db_batch_object=None,
     operation: str = "update",
     user_api_key_dict=None,
+    poller_owns_accounting: bool | None = None,
 ):
     """
     Update batch status and object in ManagedObjectTable.
@@ -1280,6 +1313,12 @@ async def update_batch_in_database(
         db_batch_object: Optional existing database object; fetched by unified_object_id when omitted
         operation: Description of operation ("update", "cancel", etc.)
         user_api_key_dict: Optional auth context for creating managed file IDs
+        poller_owns_accounting: Whether the caller already decided that the cost poller
+            owns this batch's accounting. Callers that suppress their own inline
+            accounting must pass the same decision they acted on, because re-deciding
+            here can observe a poller that became usable in between and leave the batch
+            unmarked after it was already accounted for, billing it twice. Left None by
+            callers that record no cost themselves.
     """
     import litellm.utils
 
@@ -1329,15 +1368,8 @@ async def update_batch_in_database(
             "updated_at": litellm.utils.get_utc_datetime(),
         }
 
-        # When a batch reaches completion, also mark batch_processed=True.
-        # The cost callback is enqueued asynchronously during the
-        # aretrieve_batch call that detected completion (via the @client
-        # decorator).  It is not awaited, so there is a theoretical window
-        # where the callback hasn't executed yet.  In practice the callback
-        # completes reliably.  Setting the flag here unblocks file deletion
-        # which queries batch_processed=False.  CheckBatchCost acts as a
-        # safety net for the rare case where the callback fails.
-        if db_status == "complete":
+        poller_owns: Final = batch_cost_poller_is_active() if poller_owns_accounting is None else poller_owns_accounting
+        if db_status == "complete" and not poller_owns:
             update_data["batch_processed"] = True
 
         try:

litellm/proxy/proxy_server.py

@@ -8986,6 +8986,7 @@ async def _scheduled_ptu_rollup() -> None:
                     llm_router=llm_router,
                     track_unmanaged_batch_cost=general_settings.get("track_unmanaged_batch_cost", False),
                 )
+                await check_batch_cost_job.confirm_batch_processed_support()
                 scheduler.add_job(
                     check_batch_cost_job.check_batch_cost,
                     "interval",

tests/proxy_unit_tests/test_check_batch_cost.py

@@ -114,6 +114,45 @@ async def test_cleanup_scoped_to_batch_file_purpose(
         assert "stale_expired" in where["status"]["not_in"]
         assert "created_at" in where
 
+    @pytest.mark.asyncio
+    async def test_startup_probe_confirms_batch_processed_support(
+        self, check_batch_cost_instance, mock_prisma_client
+    ):
+        mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None)
+
+        await check_batch_cost_instance.confirm_batch_processed_support()
+
+        probe_where = mock_prisma_client.db.litellm_managedobjecttable.find_first.call_args[1]["where"]
+        assert probe_where["batch_processed"] is False
+        assert check_batch_cost_instance.batch_processed_support_confirmed is True
+        assert check_batch_cost_instance._has_batch_processed_column is True
+
+    @pytest.mark.asyncio
+    async def test_startup_probe_marks_column_absent(
+        self, check_batch_cost_instance, mock_prisma_client
+    ):
+        mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(
+            side_effect=Exception("column batch_processed does not exist")
+        )
+
+        await check_batch_cost_instance.confirm_batch_processed_support()
+
+        assert check_batch_cost_instance.batch_processed_support_confirmed is False
+        assert check_batch_cost_instance._has_batch_processed_column is False
+
+    @pytest.mark.asyncio
+    async def test_startup_probe_transient_error_defers_to_poll_cycle(
+        self, check_batch_cost_instance, mock_prisma_client
+    ):
+        mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(
+            side_effect=Exception("connection reset by peer")
+        )
+
+        await check_batch_cost_instance.confirm_batch_processed_support()
+
+        assert check_batch_cost_instance.batch_processed_support_confirmed is False
+        assert check_batch_cost_instance._has_batch_processed_column is True
+
     @pytest.mark.asyncio
     async def test_find_many_uses_pagination_and_excludes_stale(
         self, check_batch_cost_instance, mock_prisma_client
@@ -143,6 +182,7 @@ async def test_find_many_uses_pagination_and_excludes_stale(
         assert "complete" not in not_in
         assert "completed" not in not_in
         assert find_call[1]["where"]["batch_processed"] is False
+        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 +211,7 @@ async def test_fallback_query_used_when_batch_processed_missing(
         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/enterprise/proxy/test_managed_files_hook.py

@@ -828,3 +828,49 @@ async def test_cost_job_and_retrieve_paths_mint_identical_unified_output_file_id
         model_id="model-deploy-xyz",
         model_name=cost_job_model_name,
     )
+
+
+@pytest.mark.asyncio
+async def test_batch_create_hook_persists_creating_key_and_tags():
+    """Regression: the /v1/batches create hook must persist the creating key and the
+    request's tags on the managed object row. CheckBatchCost, which owns the batch's
+    accounting once the retrieve path defers to it, bills whatever the row carries, and
+    without these columns the cost lands on the user alone and the key's spend and
+    budget never see it."""
+    managed_files = _make_managed_files_instance()
+    creator = UserAPIKeyAuth(api_key="sk-the-creator", user_id="alice", parent_otel_span=None)
+    create_response = _make_batch_response(status="validating", output_file_id=None)
+
+    await managed_files.async_post_call_success_hook(
+        data={"litellm_metadata": {"tags": ["env:prod", "team:ml"], "user_api_key": creator.api_key}},
+        user_api_key_dict=creator,
+        response=create_response,
+    )
+
+    managed_files.store_unified_object_id.assert_awaited_once()
+    stored = managed_files.store_unified_object_id.await_args.kwargs
+    assert stored["persist_attribution"] is True
+    assert stored["request_tags"] == ("env:prod", "team:ml")
+    assert stored["user_api_key_dict"] is creator
+
+
+@pytest.mark.asyncio
+async def test_batch_retrieve_hook_does_not_claim_attribution():
+    """A retrieve carries unified_batch_id but no unified_file_id, so it must not rewrite
+    the row's paying key to whoever happens to poll the batch."""
+    managed_files = _make_managed_files_instance()
+    retrieve_response = _make_batch_response(status="in_progress", output_file_id=None)
+    retrieve_response._hidden_params = {
+        "unified_batch_id": "some-unified-batch-id",
+        "model_id": "model-deploy-xyz",
+        "model_name": "azure/gpt-4",
+    }
+
+    await managed_files.async_post_call_success_hook(
+        data={"litellm_metadata": {"tags": ["poller:tag"]}},
+        user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-poller", user_id="bob", parent_otel_span=None),
+        response=retrieve_response,
+    )
+
+    managed_files.store_unified_object_id.assert_awaited_once()
+    assert managed_files.store_unified_object_id.await_args.kwargs["persist_attribution"] is False

tests/test_litellm/proxy/batches_endpoints/test_endpoints.py

@@ -2410,3 +2410,34 @@ async def test_cancel__unified_batch_id_allowed_when_managed_files_required(canc
         await call_cancel(cancel_harness, _unified_batch_id())
 
     assert cancel_harness.router_acancel.call_count == 1
+
+
+
+
+@pytest.mark.asyncio
+async def test_retrieve__managed_batch_defers_cost_to_the_poller_when_it_is_running(retrieve_harness):
+    with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=True)):
+        await call_retrieve(retrieve_harness, _unified_batch_id())
+
+    assert retrieve_harness.router.aretrieve_batch.await_count == 1
+    metadata = retrieve_harness.router.aretrieve_batch.await_args.kwargs.get("litellm_metadata") or {}
+    assert metadata.get("batch_ignore_default_logging") is True
+
+
+@pytest.mark.asyncio
+async def test_retrieve__managed_batch_still_accounts_inline_without_a_poller(retrieve_harness):
+    with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=False)):
+        await call_retrieve(retrieve_harness, _unified_batch_id())
+
+    assert retrieve_harness.router.aretrieve_batch.await_count == 1
+    metadata = retrieve_harness.router.aretrieve_batch.await_args.kwargs.get("litellm_metadata") or {}
+    assert metadata.get("batch_ignore_default_logging") is None
+
+
+@pytest.mark.asyncio
+async def test_retrieve__raw_batch_id_is_untouched_by_the_poller_handoff(retrieve_harness):
+    with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=True)):
+        await call_retrieve(retrieve_harness, "batch-raw-xyz")
+
+    metadata = retrieve_harness.litellm_aretrieve.await_args.kwargs.get("litellm_metadata") or {}
+    assert metadata.get("batch_ignore_default_logging") is None