Merge pull request #37247 from BerriAI/litellm_guardrail_usage_retry_safe_errors

fix(guardrails): retry usage upserts only on connection errors

litellm/proxy/guardrails/usage_endpoints.py

@@ -14,6 +14,7 @@ from fastapi import APIRouter, Depends, Query
 from pydantic import BaseModel
 from typing_extensions import NotRequired, ReadOnly, TypedDict
 
+from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.repositories.table_repositories import (
@@ -111,7 +112,16 @@ async def _find_daily_guardrail_usage_units(
     prisma_client: "PrismaClient",
     where: "prisma_types.LiteLLM_DailyGuardrailUsageUnitsWhereInput",
 ) -> "Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]":
-    return await _daily_guardrail_usage_units_table(prisma_client).find_many(where=where)
+    from prisma.errors import TableNotFoundError
+
+    try:
+        return await _daily_guardrail_usage_units_table(prisma_client).find_many(where=where)
+    except TableNotFoundError as e:
+        verbose_proxy_logger.warning(
+            "Guardrail usage units are unavailable until the LiteLLM_DailyGuardrailUsageUnits migration is applied: %s",
+            e,
+        )
+        return ()
 
 
 def _counter_name(row: "prisma_models.LiteLLM_DailyGuardrailUsageUnits") -> str:

litellm/proxy/guardrails/usage_tracking.py

@@ -15,6 +15,7 @@ from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeVar
 
 from litellm._logging import verbose_proxy_logger
+from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES
 from litellm.proxy.utils import PrismaClient
 from litellm.repositories.table_repositories import (
     DailyGuardrailMetricsRepository,
@@ -63,11 +64,21 @@ async def _upsert_rows_with_retry(
     retries_left: int = _UPSERT_RETRY_TIMES,
 ) -> None:
     outcomes: Final = {key: await _attempt_upsert(upsert_row, key, value) for key, value in rows.items()}
-    failed: Final = MappingProxyType({key: rows[key] for key, error in outcomes.items() if error is not None})
-    if not failed:
+    for key, error in outcomes.items():
+        if error is not None and not isinstance(error, DB_RETRY_SAFE_ERROR_TYPES):
+            verbose_proxy_logger.warning(
+                "Guardrail usage tracking: %s upsert failed for %s and is not safe to retry (non-fatal): %s",
+                label,
+                key,
+                error,
+            )
+    retryable: Final = MappingProxyType(
+        {key: rows[key] for key, error in outcomes.items() if isinstance(error, DB_RETRY_SAFE_ERROR_TYPES)}
+    )
+    if not retryable:
         return
     if retries_left == 0:
-        for key in failed:
+        for key in retryable:
             verbose_proxy_logger.warning(
                 "Guardrail usage tracking: %s upsert failed for %s after %d retries (non-fatal): %s",
                 label,
@@ -77,7 +88,7 @@ async def _upsert_rows_with_retry(
             )
         return
     await sleep(2 ** (_UPSERT_RETRY_TIMES - retries_left))
-    await _upsert_rows_with_retry(failed, upsert_row, label, sleep, retries_left - 1)
+    await _upsert_rows_with_retry(retryable, upsert_row, label, sleep, retries_left - 1)
 
 
 def _guardrail_status_to_action(status: str | None) -> str:

tests/test_litellm/proxy/guardrails/test_usage_endpoints.py

@@ -19,6 +19,7 @@ import pytest
 sys.path.insert(0, os.path.abspath("../../.."))
 
 from fastapi import HTTPException
+from prisma.errors import TableNotFoundError
 
 from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
 from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler
@@ -295,6 +296,44 @@ async def test_detail_breaks_units_down_by_day_team_and_key():
     assert units_where == {"guardrail_id": {"in": ["yaml-pii", "yaml-1"]}, "date": {"gte": START, "lte": END}}
 
 
+def _units_table_missing() -> TableNotFoundError:
+    return TableNotFoundError(
+        data={"user_facing_error": {"meta": {"table": "public.LiteLLM_DailyGuardrailUsageUnits"}}}
+    )
+
+
+@pytest.mark.asyncio
+async def test_overview_degrades_units_to_empty_when_units_table_is_missing():
+    prisma = _prisma(metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)])
+    prisma.db.litellm_dailyguardrailusageunits.find_many = AsyncMock(side_effect=_units_table_missing())
+    handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"))
+    p1, p2 = _patches(prisma, handler)
+    with p1, p2:
+        resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
+    row = next(r for r in resp.rows if r.id == "yaml-uuid")
+    assert (row.requestsEvaluated, row.usageUnits) == (4, {})
+    assert (resp.totalRequests, resp.totalBlocked, resp.totalUsageUnits) == (4, 1, {})
+
+
+@pytest.mark.asyncio
+async def test_detail_degrades_units_to_empty_when_units_table_is_missing():
+    prisma = _prisma(metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)])
+    prisma.db.litellm_dailyguardrailusageunits.find_many = AsyncMock(side_effect=_units_table_missing())
+    handler = _config_handler(_yaml_guardrail())
+    p1, p2 = _patches(prisma, handler)
+    with p1, p2:
+        resp = await guardrails_usage_detail(
+            guardrail_id="yaml-1", start_date=START, end_date=END, user_api_key_dict=ADMIN
+        )
+    assert (resp.requestsEvaluated, resp.failRate) == (4, 25.0)
+    assert (resp.usage_units, list(resp.usage_units_daily), resp.usage_units_by_team, resp.usage_units_by_key) == (
+        {},
+        [],
+        {},
+        {},
+    )
+
+
 # ---- logs -------------------------------------------------------------------
 
 

tests/test_litellm/proxy/guardrails/test_usage_tracking.py

@@ -3,6 +3,7 @@ from datetime import datetime, timezone
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock
 
+import httpx
 import pytest
 
 from litellm.proxy.guardrails.usage_tracking import process_spend_logs_guardrail_usage
@@ -94,8 +95,8 @@ async def test_one_failing_upsert_does_not_drop_remaining_writes():
     permanently under-report billable counters.
     """
     prisma = _prisma()
-    prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = RuntimeError("db down")
-    prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [RuntimeError("db down"), None, None]
+    prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down")
+    prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [httpx.ConnectError("db down"), None, None]
     sleep, _ = _fake_sleep()
     logs = [
         _payload("r1", usage={"topicPolicyUnits": 1}),
@@ -113,12 +114,13 @@ async def test_one_failing_upsert_does_not_drop_remaining_writes():
 @pytest.mark.asyncio
 async def test_transient_upsert_failure_is_retried_with_backoff_for_failed_rows_only():
     """
-    A transient DB error must not permanently drop billed units from the
-    aggregates: only the rows that failed are re-sent, after exponential
-    backoff, and the batch ends once every row has landed.
+    A connection error (the write provably never reached the database) must
+    not permanently drop billed units from the aggregates: only the rows that
+    failed are re-sent, after exponential backoff, and the batch ends once
+    every row has landed.
     """
     prisma = _prisma()
-    prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [RuntimeError("blip"), None, None]
+    prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [httpx.ConnectError("blip"), None, None]
     sleep, delays = _fake_sleep()
     logs = [
         _payload("r1", usage={"topicPolicyUnits": 1}),
@@ -136,7 +138,7 @@ async def test_transient_upsert_failure_is_retried_with_backoff_for_failed_rows_
 @pytest.mark.asyncio
 async def test_persistent_upsert_failure_stops_after_three_retries():
     prisma = _prisma()
-    prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = RuntimeError("db down")
+    prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down")
     sleep, delays = _fake_sleep()
 
     await process_spend_logs_guardrail_usage(prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep)
@@ -146,6 +148,61 @@ async def test_persistent_upsert_failure_stops_after_three_retries():
     assert prisma.db.litellm_dailyguardrailusageunits.upsert.call_count == 1
 
 
+def _units_upsert_wheres(prisma: MagicMock) -> list[tuple]:
+    return [
+        tuple(
+            c.kwargs["where"]["guardrail_id_date_team_id_api_key_usage_unit"][k]
+            for k in ("guardrail_id", "date", "team_id", "api_key", "usage_unit")
+        )
+        for c in prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list
+    ]
+
+
+@pytest.mark.asyncio
+async def test_post_send_failure_is_never_retried_so_increments_cannot_double_count():
+    """
+    Follow-up to #37225: the units upsert is a non-idempotent increment, so an
+    ambiguous post-send failure (read timeout after the statement may have
+    committed) must be attempted exactly once. Re-sending it stacks a second
+    increment and inflates billable unit totals. Only a connection error proves
+    the write never reached the database and may be retried; the other rows in
+    the batch still land either way.
+    """
+    prisma = _prisma()
+    prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [
+        httpx.ReadTimeout("read timed out"),
+        httpx.ConnectError("refused"),
+        None,
+    ]
+    sleep, delays = _fake_sleep()
+    logs = [
+        _payload("r1", usage={"topicPolicyUnits": 1}),
+        _payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}),
+    ]
+
+    await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep)
+
+    timed_out_row = ("bedrock-guard", "2026-08-17", "", "hashed-key-2", "topicPolicyUnits")
+    refused_row = ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits")
+    assert _units_upsert_wheres(prisma) == [timed_out_row, refused_row, refused_row]
+    assert delays == [1]
+
+
+@pytest.mark.asyncio
+async def test_generic_upsert_exception_is_terminal_for_that_row_only():
+    prisma = _prisma()
+    prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = RuntimeError("constraint violation")
+    sleep, delays = _fake_sleep()
+
+    await process_spend_logs_guardrail_usage(prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep)
+
+    assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 1
+    assert delays == []
+    assert _units_upserts(prisma) == {
+        ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1,
+    }
+
+
 @pytest.mark.asyncio
 async def test_zero_and_non_int_usage_counters_are_skipped():
     prisma = _prisma()