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/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_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 ------------------------------------------------------------------- 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()