fix(proxy): charge sessions a poll's guardrail-inclusive cost, count a job cost row's guardrail run, and keep batch_cost_is_final importable from litellm.batches

litellm/batches/batch_utils.py

@@ -10,6 +10,7 @@ from pydantic import JsonValue, TypeAdapter
 
 import litellm
 from litellm._logging import verbose_logger
+from litellm.litellm_core_utils import batch_utils as core_batch_utils
 from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
 from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
 from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo
@@ -39,6 +40,7 @@ class BatchCostUsageResult:
 
 
 _DECODED_BATCH_LINE: Final = TypeAdapter(object)
+batch_cost_is_final: Final = core_batch_utils.batch_cost_is_final
 
 
 def _uses_native_vertex_output(

litellm/proxy/db/db_spend_update_writer.py

@@ -160,6 +160,22 @@ def _org_member_transaction_key(org_id: str, user_id: str) -> str:
     return f"organization_id::{quote(org_id, safe='')}::user_id::{quote(user_id, safe='')}"
 
 
+async def _track_guardrail_usage_of_cost_row(
+    payload: SpendLogsPayload, prisma_client: "PrismaClient", disable_spend_logs: bool
+) -> None:
+    """A job's cost row skips the spend log queue, so its guardrail runs are counted here instead of at the flush."""
+    if disable_spend_logs:
+        return
+    from litellm.proxy.guardrails.usage_tracking import process_spend_logs_guardrail_usage
+
+    try:
+        await process_spend_logs_guardrail_usage(prisma_client=prisma_client, logs_to_process=(payload,))
+    except Exception as e:  # noqa: BLE001  # guardrail usage is best effort, as at the spend log flush
+        verbose_proxy_logger.warning(
+            "Guardrail usage tracking failed for spend row %s (non-fatal): %s", payload["request_id"], e
+        )
+
+
 def _is_settle_once_cost_row(payload: SpendLogsPayload) -> bool:
     """A poll of a job whose cost is charged once, by the first poll that sees it, however often it is polled.
 
@@ -807,15 +823,16 @@ class DBSpendUpdateWriter:
                 data=[prisma_client.jsonify_object(row)],
                 skip_duplicates=True,
             )
-            if claimed == 1:
-                return True
-            existing: Final = await spend_logs.find_unique(where={"request_id": request_id})
+            existing: Final = None if claimed == 1 else await spend_logs.find_unique(where={"request_id": request_id})
         except Exception as e:  # noqa: BLE001  # prisma raises its own hierarchy; an unreachable DB queues the row like any other spend log
             verbose_proxy_logger.warning(
                 "Could not claim spend row %s for a job's cost, queueing it: %s", request_id, e
             )
             await self._insert_spend_log_to_db(payload=prisma_client.jsonify_object(row), prisma_client=prisma_client)
             return True
+        if claimed == 1:
+            await _track_guardrail_usage_of_cost_row(payload, prisma_client, disable_spend_logs)
+            return True
         if existing is None or existing.call_type != payload["call_type"] or existing.status != "success":
             verbose_proxy_logger.warning(
                 "Spend row %s belongs to a %s request, so this job's cost is charged without a row of its own",
@@ -826,7 +843,12 @@ class DBSpendUpdateWriter:
         if existing.spend > 0:
             verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this job", request_id)
             return False
-        return await self._take_over_uncharged_cost_row(payload=payload, prisma_client=prisma_client, row=row)
+        taken_over: Final = await self._take_over_uncharged_cost_row(
+            payload=payload, prisma_client=prisma_client, row=row
+        )
+        if taken_over:
+            await _track_guardrail_usage_of_cost_row(payload, prisma_client, disable_spend_logs)
+        return taken_over
 
     async def _take_over_uncharged_cost_row(
         self, payload: SpendLogsPayload, prisma_client: "PrismaClient", row: Mapping[str, object]

litellm/proxy/hooks/max_budget_per_session_limiter.py

@@ -28,6 +28,7 @@ from litellm.exceptions import RateLimitType
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.batch_utils import (
     ChargeOnceJob,
+    charged_poll_cost,
     defers_batch_cost,
     job_charged_once_marker_key,
     job_to_charge_once,
@@ -177,7 +178,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger):
                 return
 
             callback_fields: Final = _CALLBACK_FIELDS.validate_python(kwargs)
-            response_cost: Final = finite_cost(callback_fields.get("response_cost")) or 0.0
+            response_cost: Final = finite_cost(charged_poll_cost(callback_fields)) or 0.0
             if response_cost <= 0:
                 return
 
@@ -196,7 +197,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger):
             verbose_proxy_logger.debug(
                 "MaxBudgetPerSessionHandler: incremented session %s spend by %.6f",
                 session_id,
-                response_cost,
+                charge,
             )
         except Exception as e:
             verbose_proxy_logger.warning(

tests/unit/proxy/db/test_db_spend_update_writer.py

@@ -1,6 +1,7 @@
 import asyncio
 import copy
 import functools
+import itertools
 import json
 import logging
 import re
@@ -4787,6 +4788,52 @@ async def test_update_database_charges_a_video_once_from_the_poll_that_sees_it_c
     }
 
 
+def _indexed_guardrail_request_ids(prisma: MagicMock, index_table: str) -> tuple[str, ...]:
+    calls: Final = getattr(prisma.db, index_table).create_many.await_args_list
+    batches: Final = tuple(call.kwargs["data"] for call in calls)
+    return tuple(row["request_id"] for row in itertools.chain.from_iterable(batches))
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("disable_spend_logs", [False, True], ids=["spend_logs_on", "spend_logs_off"])
+@pytest.mark.parametrize("left_uncharged", [False, True], ids=["claimed_row", "taken_over_row"])
+async def test_a_job_cost_row_counts_its_polls_guardrail_run_in_guardrail_usage_once(
+    disable_spend_logs: bool, left_uncharged: bool
+):
+    """The job's cost row is written by its claim, not the spend log queue whose flush counts guardrail runs, so a
+    finished video's guardrail run must still reach guardrail usage, once, and not when spend logs are off."""
+    from litellm.repositories.table_repositories import (
+        DailyGuardrailMetricsRepository,
+        SpendLogGuardrailIndexRepository,
+    )
+
+    db_writer = DBSpendUpdateWriter()
+    db_writer._batch_database_updates = AsyncMock()
+    prisma = _spend_logs_table_prisma()
+    getattr(prisma.db, SpendLogGuardrailIndexRepository.table_name).create_many = AsyncMock(return_value=1)
+    getattr(prisma.db, DailyGuardrailMetricsRepository.table_name).upsert = AsyncMock()
+    if left_uncharged:
+        prisma.db.litellm_spendlogs.rows["video_abc_video_cost"] = {**_video_status_payload(), "spend": 0.0}
+    guarded_poll: Final = {
+        **_video_status_payload(),
+        "metadata": json.dumps(
+            {"guardrail_information": [{"guardrail_name": "priced", "guardrail_status": "success"}]}
+        ),
+    }
+
+    charged = [
+        await _update_database_with(
+            db_writer, prisma, guarded_poll, disable_spend_logs=disable_spend_logs, response_cost=0.43
+        )
+        for _ in range(2)
+    ]
+
+    assert charged == [True, False]
+    assert _indexed_guardrail_request_ids(prisma, SpendLogGuardrailIndexRepository.table_name) == (
+        () if disable_spend_logs else ("video_abc_video_cost",)
+    )
+
+
 def _job_poll(
     litellm_call_id: str,
     response_cost: float,

tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py

@@ -222,14 +222,14 @@ async def test_a_batch_retrieve_before_its_cost_is_final_charges_the_session_not
 
 
 def _video_poll(
-    cost: float, guardrail_cost: float = 0.0, session_id: str = "session-jobs", status: str = "completed"
+    video_cost: float, guardrail_cost: float = 0.0, session_id: str = "session-jobs", status: str = "completed"
 ) -> dict[str, object]:
     return {
         "call_type": "avideo_status",
-        "response_cost": cost,
+        "response_cost": video_cost,
         POLLED_VIDEO_STATUS_KWARG: status,
         "standard_logging_object": {
-            "response_cost": cost,
+            "response_cost": video_cost + guardrail_cost,
             "guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": guardrail_cost}],
         },
         "litellm_params": {"metadata": {"session_id": session_id, "agent_id": "agent-budget-123"}},
@@ -237,28 +237,39 @@ def _video_poll(
 
 
 @pytest.mark.asyncio
-async def test_every_video_poll_pays_its_guardrail_and_only_the_finished_poll_pays_the_session_for_the_video():
+@pytest.mark.parametrize(
+    ("video_cost", "guardrail_cost"),
+    [(0.40, 0.01), (0.05, 0.10)],
+    ids=["video_costs_more_than_its_guardrail", "video_costs_no_more_than_its_guardrail"],
+)
+async def test_every_video_poll_pays_its_guardrail_and_only_the_finished_poll_pays_the_session_for_the_video(
+    video_cost: float, guardrail_cost: float
+):
     from litellm.types.videos.main import VideoObject
 
     handler = PROXY_MaxBudgetPerSessionHandler(internal_usage_cache=InternalUsageCache(DualCache()))
+    completed_polls: Final = 3
 
     with patch(
         "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry"
     ) as mock_registry:
         mock_registry.get_agent_by_id.return_value = _make_mock_agent(max_budget_per_session=5.0)
         await handler.async_log_success_event(
-            _video_poll(0.01, guardrail_cost=0.01, status="in_progress"),
+            _video_poll(0.0, guardrail_cost=guardrail_cost, status="in_progress"),
             VideoObject(id="job-1", object="video", status="in_progress"),
             None,
             None,
         )
-        for _ in range(3):
+        for _ in range(completed_polls):
             await handler.async_log_success_event(
-                _video_poll(0.41, guardrail_cost=0.01), VideoObject(id="job-1", object="video", status="completed"), None, None
+                _video_poll(video_cost, guardrail_cost=guardrail_cost),
+                VideoObject(id="job-1", object="video", status="completed"),
+                None,
+                None,
             )
 
     spend = await handler._get_current_spend(handler._make_cache_key("session-jobs"))
-    assert spend == pytest.approx(0.01 + 0.41 + 2 * 0.01)
+    assert spend == pytest.approx(video_cost + (1 + completed_polls) * guardrail_cost)
 
 
 class _ClockedRedis: