From 35176fa64aec30de01a9b650bb3fc88da71b9244 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:05:50 -0700 Subject: [PATCH 1/2] fix(guardrails): retry usage upserts only on connection errors The daily guardrail metrics and usage-unit upserts are non-idempotent increments, but the retry loop re-sent every failed row on any exception. An ambiguous post-send failure such as a read timeout after the write had already committed therefore stacked a second increment and inflated the billable unit totals served by the guardrail usage endpoints. Retry only DB_RETRY_SAFE_ERROR_TYPES (httpx.ConnectError), the same rule the spend writer and autorouter rollup use for increment upserts, and log any other failure once as terminal for that row while the rest of the batch still lands. Follows up #37225 --- litellm/proxy/guardrails/usage_tracking.py | 19 +++-- .../proxy/guardrails/test_usage_tracking.py | 71 +++++++++++++++++-- 2 files changed, 79 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 727f767469f..0e2d37c8d57 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/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: diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 62a4bbbbe6e..693a73b152d 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/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() From 15823b1be345daf7e0ffaa2191ffe524feaa80f7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:38:43 -0700 Subject: [PATCH 2/2] fix(guardrails): degrade usage units to empty when the units table is missing GET /guardrails/usage/overview and GET /guardrails/usage/detail/{id} 500ed on a database that has not applied 20260817143646_add_daily_guardrail_usage_units yet (pip installs on litellm-proxy-extras 0.4.86 with DISABLE_SCHEMA_UPDATE=true). Both endpoints now return their metrics with empty units and log one warning until the migration lands. --- litellm/proxy/guardrails/usage_endpoints.py | 12 +++++- .../proxy/guardrails/test_usage_endpoints.py | 39 +++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index a73efed30ad..ca89c7587ba 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/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: diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py index b0f98b81c05..f63c08a2c39 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py +++ b/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 -------------------------------------------------------------------