From daced81f20a6b98c01beecf66fa997310d90bfb4 Mon Sep 17 00:00:00 2001 From: amasen02 Date: Fri, 4 Sep 2026 15:37:53 +0530 Subject: [PATCH 01/10] fix(proxy): invalidate end-user spend counter and cache on budget reset (#39726) Signed-off-by: amasen02 --- .../proxy/common_utils/reset_budget_job.py | 24 ++++++++++++++++- .../common_utils/test_reset_budget_job.py | 26 +++++++++++++++++++ 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 47f69732e95..12ba75aea24 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -38,6 +38,7 @@ from litellm.proxy.common_utils.timezone_utils import ( get_budget_reset_settings, ) from litellm.proxy.common_utils.user_api_key_cache import ( + end_user_cache_key, model_access_group_cache_key, model_access_group_spend_counter_key, tag_cache_key, @@ -177,6 +178,21 @@ def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...] return (model_access_group_cache_key(row.access_group_name),) +def _enduser_counter_key(row: _EndUserRow) -> str: + return f"spend:end_user:{row.user_id}" + + +def _enduser_cache_keys(row: _EndUserRow) -> tuple[str, ...]: + return (end_user_cache_key(row.user_id),) + + +def _enduser_carried_spend(row: _EndUserRow, caps: Mapping[str, float]) -> float: + if not caps: + return 0.0 + effective_budget_id = row.budget_id or litellm.max_end_user_budget_id + return _carried_spend(row.spend, caps.get(effective_budget_id) if effective_budget_id is not None else None) + + def _budget_link_where( budget_ids: Sequence[str], extra: Mapping[str, object] = MappingProxyType({}), @@ -650,6 +666,7 @@ class ResetBudgetJob: if _rollover_enabled() else {} # mutable-ok: empty sentinel immediately frozen by MappingProxyType ) + endusers: Final[tuple[_EndUserRow, ...]] = await self._collect_endusers_to_reset(budget_ids) return _BudgetCascade( budgets=tuple(budgets_to_reset), budget_ids=budget_ids, @@ -661,7 +678,7 @@ class ResetBudgetJob: for b in budgets_to_reset if b.budget_id is not None and b.budget_duration is not None ), - endusers=await self._collect_endusers_to_reset(budget_ids), + endusers=endusers, counter_resets=( *( (_team_membership_counter_key(row), _row_carried_spend(row, rollover_caps)) @@ -674,6 +691,10 @@ class ResetBudgetJob: (_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in model_access_groups ), + *( + (_enduser_counter_key(row), _enduser_carried_spend(row, rollover_caps)) + for row in endusers + ), ), rollover_caps=rollover_caps, cache_keys=( @@ -682,6 +703,7 @@ class ResetBudgetJob: *(key for row in orgs for key in _org_cache_keys(row)), *(key for row in tags for key in _tag_cache_keys(row)), *(key for row in model_access_groups for key in _model_access_group_cache_keys(row)), + *(key for row in endusers for key in _enduser_cache_keys(row)), ), ) diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 03b05bd9d87..e6bbfd8c7d4 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1495,6 +1495,32 @@ def test_budget_table_reset_invalidates_every_tag_not_just_the_first(reset_budge assert deleted == {"tag:tenant-a", "tag:tenant-b", "tag:tenant-c"} +def test_budget_table_reset_invalidates_enduser_counter_and_cache(reset_budget_job, mock_prisma_client, monkeypatch): + """When an end user's budget resets, its Redis spend counter is zeroed and its management cache is evicted.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + budget = _budget_row(budget_id="budget-1") + mock_prisma_client.data["budget"] = [budget] + test_enduser = type( + "LiteLLM_EndUserTable", + (), + { + "spend": 20.0, + "litellm_budget_table": budget, + "budget_id": "budget-1", + "user_id": "customer-42", + }, + ) + mock_prisma_client.data["enduser"] = [test_enduser] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:end_user:customer-42", value=0.0, ttl=60) + counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:end_user:customer-42", value=0.0, ttl=60) + deleted = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} + assert "end_user_id:customer-42" in deleted + + + def test_budget_table_reset_commits_even_when_cache_eviction_fails(reset_budget_job, mock_prisma_client, monkeypatch): """Eviction runs after the commit, so a broken cache cannot undo the write.""" counter_cache = _make_counter_invalidation_job(monkeypatch) From 3623aecc6419ed5442bea4efd435cdca3246101a Mon Sep 17 00:00:00 2001 From: amasen02 Date: Fri, 4 Sep 2026 16:33:57 +0530 Subject: [PATCH 02/10] style(proxy): add Final type annotations to enduser budget reset variables --- litellm/proxy/common_utils/reset_budget_job.py | 2 +- .../proxy/common_utils/test_reset_budget_job.py | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 12ba75aea24..4fb544cbb15 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -189,7 +189,7 @@ def _enduser_cache_keys(row: _EndUserRow) -> tuple[str, ...]: def _enduser_carried_spend(row: _EndUserRow, caps: Mapping[str, float]) -> float: if not caps: return 0.0 - effective_budget_id = row.budget_id or litellm.max_end_user_budget_id + effective_budget_id: Final[str | None] = row.budget_id or litellm.max_end_user_budget_id return _carried_spend(row.spend, caps.get(effective_budget_id) if effective_budget_id is not None else None) diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index e6bbfd8c7d4..bc9926a314f 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -4,7 +4,7 @@ import sys import types from datetime import datetime, timedelta, timezone from datetime import time as dt_time -from typing import Any, Dict, List +from typing import Any, Dict, Final, List from unittest.mock import AsyncMock, MagicMock import httpx @@ -1497,10 +1497,10 @@ def test_budget_table_reset_invalidates_every_tag_not_just_the_first(reset_budge def test_budget_table_reset_invalidates_enduser_counter_and_cache(reset_budget_job, mock_prisma_client, monkeypatch): """When an end user's budget resets, its Redis spend counter is zeroed and its management cache is evicted.""" - counter_cache = _make_counter_invalidation_job(monkeypatch) - budget = _budget_row(budget_id="budget-1") + counter_cache: Final = _make_counter_invalidation_job(monkeypatch) + budget: Final = _budget_row(budget_id="budget-1") mock_prisma_client.data["budget"] = [budget] - test_enduser = type( + test_enduser: Final = type( "LiteLLM_EndUserTable", (), { @@ -1516,7 +1516,7 @@ def test_budget_table_reset_invalidates_enduser_counter_and_cache(reset_budget_j counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:end_user:customer-42", value=0.0, ttl=60) counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:end_user:customer-42", value=0.0, ttl=60) - deleted = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} + deleted: Final = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} assert "end_user_id:customer-42" in deleted From 9c8594c7b8c4c948abb85507e7611762b0b0f70f Mon Sep 17 00:00:00 2001 From: amasen02 Date: Fri, 4 Sep 2026 16:37:58 +0530 Subject: [PATCH 03/10] style(proxy): format reset_budget_job with ruff --- litellm/proxy/common_utils/reset_budget_job.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 4fb544cbb15..f2648c8466e 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -691,10 +691,7 @@ class ResetBudgetJob: (_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in model_access_groups ), - *( - (_enduser_counter_key(row), _enduser_carried_spend(row, rollover_caps)) - for row in endusers - ), + *((_enduser_counter_key(row), _enduser_carried_spend(row, rollover_caps)) for row in endusers), ), rollover_caps=rollover_caps, cache_keys=( From 2042364fc2976ea735ab3d8c77dd4f4b27df3b84 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 4 Sep 2026 11:50:28 -0700 Subject: [PATCH 04/10] fix(proxy): strip every TypedDict qualifier before numeric form-field detection _numeric_form_type only peeled a single ReadOnly layer, so a field still wrapped in Required/NotRequired was read as non-numeric and dropped from the mapping. Which qualifiers survive get_type_hints varies by interpreter version and by include_extras, so on Python 3.10 NotRequired[ReadOnly[int]] reached the check intact and the field was silently skipped, which is what turns the mapped test red on the 3.10 leg only. Peel Required/NotRequired/ReadOnly/Annotated in any order and nesting instead. The one production caller feeds a schema with no qualifiers, so the resulting mapping is unchanged on every interpreter in the matrix, but a field written the house-convention way stops being dropped. --- .../proxy/common_utils/http_parsing_utils.py | 16 +++++++++++++--- .../common_utils/test_http_parsing_utils.py | 18 ++++++++++++++++++ 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 552d1ea434f..a396e543e94 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -2,11 +2,11 @@ import json import re from collections.abc import Collection, Mapping from types import MappingProxyType, UnionType -from typing import Any, Final, Union, get_args, get_origin +from typing import Annotated, Any, Final, Union, get_args, get_origin import orjson from fastapi import Request, UploadFile, status -from typing_extensions import ReadOnly +from typing_extensions import NotRequired, ReadOnly, Required from litellm._logging import verbose_proxy_logger from litellm.constants import MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB @@ -18,6 +18,8 @@ from litellm.types.router import Deployment _FORM_CONTENT_TYPES: Final[frozenset[str]] = frozenset({"application/x-www-form-urlencoded", "multipart/form-data"}) +_ANNOTATION_QUALIFIERS: Final[frozenset[object]] = frozenset({Annotated, NotRequired, ReadOnly, Required}) + def _normalize_media_type(content_type: str) -> str: """Return the bare media type per RFC 7231: strip params, trim, lowercase.""" @@ -42,9 +44,17 @@ def _is_json_content_type(content_type: str) -> bool: return _normalize_media_type(content_type) == "application/json" +def _unqualified(annotation: object) -> object: + """Which qualifiers ``get_type_hints`` already stripped varies by interpreter version, so peel them all.""" + if get_origin(annotation) not in _ANNOTATION_QUALIFIERS: + return annotation + qualified: Final[tuple[object, ...]] = get_args(annotation) + return _unqualified(qualified[0]) + + def _numeric_form_type(annotation: object) -> type[int] | type[float] | None: """The scalar to parse an ``int``/``float``-typed field as, else ``None``.""" - unwrapped: Final = get_args(annotation)[0] if get_origin(annotation) is ReadOnly else annotation + unwrapped: Final = _unqualified(annotation) candidates: Final = ( tuple(arg for arg in get_args(unwrapped) if arg is not type(None)) if get_origin(unwrapped) in (Union, UnionType) diff --git a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py index fcfb9342176..011571a37e0 100644 --- a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py @@ -1053,6 +1053,8 @@ class TestNumericFormFields: read_only: ReadOnly[int | None] not_required: NotRequired[ReadOnly[int]] required: Required[ReadOnly[Annotated[float, "meta"]]] + read_only_not_required: ReadOnly[NotRequired[int]] + read_only_required: ReadOnly[Required[float]] assert dict(numeric_form_fields(get_type_hints(Schema))) == { "plain": int, @@ -1061,6 +1063,22 @@ class TestNumericFormFields: "read_only": int, "not_required": int, "required": float, + "read_only_not_required": int, + "read_only_required": float, + } + + def test_qualifiers_are_unwrapped_when_get_type_hints_keeps_extras(self): + from typing_extensions import Annotated, NotRequired, ReadOnly, Required, TypedDict + + class Schema(TypedDict, total=False): + annotated: ReadOnly[Annotated[int, "meta"]] + not_required: NotRequired[ReadOnly[int]] + required: Required[ReadOnly[Annotated[float, "meta"]]] + + assert dict(numeric_form_fields(get_type_hints(Schema, include_extras=True))) == { + "annotated": int, + "not_required": int, + "required": float, } def test_non_scalar_and_bool_fields_are_skipped(self): From 976f8625f34c9c0eb7ac5e5976493dde2c4cd997 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 16:42:29 -0700 Subject: [PATCH 05/10] test(proxy): cover default-tier end-user counter reset with rollover --- .../common_utils/test_reset_budget_job.py | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index bc9926a314f..56c0efb41d2 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -3054,6 +3054,38 @@ def test_budget_cascade_carries_enduser_overage_when_rollover_enabled( } in enduser_writes +def test_budget_cascade_carries_default_tier_enduser_counter_when_rollover_enabled( + rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch +): + """An end user on the default budget (no budget_id on its row) 5 over the cap + keeps a counter of 5 in the next window and loses its cached object.""" + import litellm + + counter_cache: Final = _make_counter_invalidation_job(monkeypatch) + monkeypatch.setattr(litellm, "max_end_user_budget_id", "default-enduser-budget") + mock_prisma_client.data["budget"] = [ + _budget_row(budget_id="default-enduser-budget", budget_duration="1d", max_budget=10.0) + ] + implicit_enduser: Final = type( + "EndUserRow", + (), + { + "spend": 15.0, + "user_id": "enduser-implicit", + "budget_id": None, + "model_dump": lambda self=None: {"spend": 15.0, "user_id": "enduser-implicit", "budget_id": None, "blocked": False}, + }, + ) + mock_prisma_client.db.litellm_endusertable.set_find_many_results([implicit_enduser]) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:end_user:enduser-implicit", value=5.0, ttl=60) + counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:end_user:enduser-implicit", value=5.0, ttl=60) + deleted: Final = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} + assert "end_user_id:enduser-implicit" in deleted + + def _replay_spend_writes(writes, spend): """Apply the queued update_many statements in order, the way the DB transaction executes them, and return the row's final spend.""" From 89086db28233514c3cc07333fcffdcd974cc8573 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:39:46 -0700 Subject: [PATCH 06/10] fix(proxy): floor end-user budget checks on the DB row after a reset The reset job evicts the cached end-user object only from its own worker's in-memory cache (plus Redis), so every other uvicorn worker and replica keeps the pre-reset spend for up to user_api_key_cache_ttl (60s by default). Those workers pass that stale spend as fallback_spend, and since the authoritative floor read returned None for spend:end_user: keys, get_current_spend handed the stale value straight back and the end user kept getting 429 after the rollover on every worker but the one that ran the reset. The floor read now consults LiteLLM_EndUserTable.spend for end-user counters, the same way keys, teams, users, and orgs already read their rows. It runs only when the shared counter sits below the cached spend (a reset or a Redis restart) and stays behind the existing 5s in-process marker, so the normal request path still does no DB read. Cold end-user counters keep seeding from the cached object rather than the row, so from_db is unchanged for them. --- litellm/proxy/db/spend_counter_reseed.py | 25 +++++- litellm/proxy/proxy_server.py | 46 +++++++---- .../proxy/db/test_spend_counter_reseed.py | 69 +++++++++++++++- .../proxy/proxy_server/test_spend_counters.py | 79 +++++++++++++++++-- 4 files changed, 196 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 7b3c261036e..a38b8a47dbd 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -26,6 +26,7 @@ from litellm.proxy._types import Litellm_EntityType from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( BudgetWindowSpendRepository, + EndUserRepository, SpendLogsRepository, TeamMembershipRepository, ) @@ -36,6 +37,8 @@ from litellm.repositories.verification_token_repository import ( ) if TYPE_CHECKING: + from prisma.types import LiteLLM_EndUserTableWhereUniqueInput + from litellm.caching.dual_cache import DualCache from litellm.proxy.utils import PrismaClient @@ -47,6 +50,8 @@ _WINDOW_SPEND_ENTITY_TYPES: Final[Mapping[str, str]] = MappingProxyType( } ) +END_USER_COUNTER_PREFIX: Final = "spend:end_user:" + _WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType( { "Key": "api_key", @@ -74,6 +79,10 @@ class SpendCounterReseed: End-user and tag spend counters intentionally do not reseed here. Their auth paths already load the corresponding objects via get_end_user_object() and get_tag_objects_batch(); callers pass those values as fallback_spend. + end_user_from_db is the one end-user read, used only as the budget floor when + a counter sits below that cached spend: a worker that did not run the budget + reset still caches the pre-reset end-user object, and LiteLLM_EndUserTable + is the row the reset zeroed. """ _locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict() @@ -129,7 +138,7 @@ class SpendCounterReseed: elif counter_key.startswith("spend:user:"): user_id = counter_key[len("spend:user:") :] row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) - elif counter_key.startswith("spend:end_user:") or counter_key.startswith("spend:tag:"): + elif counter_key.startswith(END_USER_COUNTER_PREFIX) or counter_key.startswith("spend:tag:"): return None elif counter_key.startswith("spend:org:"): org_id: Final = counter_key[len("spend:org:") :] @@ -143,6 +152,20 @@ class SpendCounterReseed: return None return float(getattr(row, "spend", 0.0) or 0.0) + @staticmethod + async def end_user_from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None: + if prisma_client is None or not counter_key.startswith(END_USER_COUNTER_PREFIX): + return None + where: Final[LiteLLM_EndUserTableWhereUniqueInput] = {"user_id": counter_key[len(END_USER_COUNTER_PREFIX) :]} + try: + row: Final = await EndUserRepository(prisma_client).table.find_unique(where=where) + except Exception: # noqa: BLE001 # a failed floor read falls back to the cached spend, like from_db + verbose_proxy_logger.exception("SpendCounterReseed.end_user_from_db: failed for %s", counter_key) + return None + if row is None: + return None + return float(row.spend or 0.0) + @staticmethod def _is_key_or_team_window_counter(counter_key: str) -> bool: for prefix in ("spend:key:", "spend:team:"): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1f39a78e12a..47c0811d903 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -423,7 +423,7 @@ from litellm.proxy.db.proxy_worker_heartbeat import ( PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS, ProxyWorkerHeartbeat, ) -from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed +from litellm.proxy.db.spend_counter_reseed import END_USER_COUNTER_PREFIX, SpendCounterReseed from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config @@ -2580,6 +2580,29 @@ async def reseed_spend_counter_from_db(counter_key: str) -> None: await _repair_stale_spend_counter(counter_key=counter_key, db_spend=db_spend) +async def _floor_spend_from_db( + counter_key: str, + window_entity_type: str | None, + window_entity_id: str | None, + window_duration: str | None, + window_start: datetime | None, +) -> float | None: + if counter_key.startswith(END_USER_COUNTER_PREFIX): + return await SpendCounterReseed.end_user_from_db(prisma_client=prisma_client, counter_key=counter_key) + entity_spend: Final = await SpendCounterReseed.from_db(prisma_client=prisma_client, counter_key=counter_key) + if entity_spend is not None: + return entity_spend + if window_entity_type is None or window_entity_id is None or window_start is None: + return None + return await SpendCounterReseed.window_from_db( + prisma_client=prisma_client, + entity_type=window_entity_type, + entity_id=window_entity_id, + window_duration=window_duration, + window_start=window_start, + ) + + async def _authoritative_floor_spend( counter_key: str, window_entity_type: str | None = None, @@ -2592,20 +2615,13 @@ async def _authoritative_floor_spend( if cached is not None: return float(cached) - db_spend = await SpendCounterReseed.from_db(prisma_client=prisma_client, counter_key=counter_key) - if ( - db_spend is None - and window_entity_type is not None - and window_entity_id is not None - and window_start is not None - ): - db_spend = await SpendCounterReseed.window_from_db( - prisma_client=prisma_client, - entity_type=window_entity_type, - entity_id=window_entity_id, - window_duration=window_duration, - window_start=window_start, - ) + db_spend: Final = await _floor_spend_from_db( + counter_key=counter_key, + window_entity_type=window_entity_type, + window_entity_id=window_entity_id, + window_duration=window_duration, + window_start=window_start, + ) if db_spend is None: return None diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index 816f9ae72f4..3bd6d93328d 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -18,7 +18,7 @@ from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc) -class _FakeWindowSpendTable: +class _FakeFindUniqueTable: def __init__(self, row: SimpleNamespace | None, error: Exception | None = None) -> None: self._row = row self._error = error @@ -47,10 +47,13 @@ class _FakePrismaClient: row: SimpleNamespace | None = None, spend_logs_total: float = 0.0, error: Exception | None = None, + end_user_row: SimpleNamespace | None = None, + end_user_error: Exception | None = None, ) -> None: self.db = SimpleNamespace( - litellm_budgetwindowspend=_FakeWindowSpendTable(row=row, error=error), + litellm_budgetwindowspend=_FakeFindUniqueTable(row=row, error=error), litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total), + litellm_endusertable=_FakeFindUniqueTable(row=end_user_row, error=end_user_error), ) @@ -248,3 +251,65 @@ async def test_coalesced_window_seeds_a_cold_counter_from_the_row(): assert result == 4.5 assert cache.in_memory_cache.get_cache(key=counter_key) == 4.5 assert prisma.db.litellm_spendlogs.call_count == 0 + + +@pytest.mark.asyncio +async def test_end_user_from_db_reads_the_end_user_row_by_user_id(): + prisma = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=0.0)) + + result = await SpendCounterReseed.end_user_from_db( + prisma_client=prisma, counter_key="spend:end_user:customer-42" + ) + + assert result == 0.0 + assert prisma.db.litellm_endusertable.where_clauses == [{"user_id": "customer-42"}] + + +@pytest.mark.asyncio +async def test_end_user_from_db_returns_the_recorded_spend(): + prisma = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=12.5)) + + assert ( + await SpendCounterReseed.end_user_from_db(prisma_client=prisma, counter_key="spend:end_user:customer-42") + == 12.5 + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("counter_key", ["spend:key:hashed", "spend:team:t1", "spend:tag:t1"]) +async def test_end_user_from_db_ignores_other_counter_kinds_without_touching_the_db(counter_key): + prisma = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="x", spend=5.0)) + + assert await SpendCounterReseed.end_user_from_db(prisma_client=prisma, counter_key=counter_key) is None + assert prisma.db.litellm_endusertable.where_clauses == [] + + +@pytest.mark.asyncio +async def test_end_user_from_db_returns_none_without_a_row_a_client_or_on_db_error(): + assert ( + await SpendCounterReseed.end_user_from_db(prisma_client=None, counter_key="spend:end_user:customer-42") + is None + ) + assert ( + await SpendCounterReseed.end_user_from_db( + prisma_client=_FakePrismaClient(end_user_row=None), counter_key="spend:end_user:customer-42" + ) + is None + ) + assert ( + await SpendCounterReseed.end_user_from_db( + prisma_client=_FakePrismaClient(end_user_error=RuntimeError("db down")), + counter_key="spend:end_user:customer-42", + ) + is None + ) + + +@pytest.mark.asyncio +async def test_from_db_still_never_reads_the_end_user_row(): + """A cold end-user counter keeps seeding from the cached end-user object the auth + path already loaded; the row is read only as the budget floor.""" + prisma = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=5.0)) + + assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:end_user:customer-42") is None + assert prisma.db.litellm_endusertable.where_clauses == [] diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index fb3de990deb..7cc1390fcbd 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -223,21 +223,90 @@ async def test_get_current_spend_floor_caches_db_read(monkeypatch): @pytest.mark.asyncio async def test_get_current_spend_floors_end_user_tag_against_fallback(monkeypatch): - """End-user and tag counters have no DB row (from_db returns None). When the - counter is stale-low, enforcement falls back to the caller's recorded spend - (loaded fresh in auth) instead of trusting the stale counter.""" + """Tag counters have no DB row (from_db returns None), and an end-user counter has + none to read without a DB client. When such a counter is stale-low, enforcement + falls back to the caller's recorded spend (loaded fresh in auth) instead of + trusting the stale counter.""" fake_cache = _make_spend_counter_cache(redis_get_value=2.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", AsyncMock(return_value=None)) + for counter_key in ("spend:end_user:e1", "spend:tag:t1"): + result = await ps.get_current_spend( + counter_key=counter_key, + fallback_spend=20.0, + max_budget=10.0, + ) + + assert result == 20.0 + # no DB row to repair against, so the shared counter is left untouched + fake_cache.redis_cache.async_set_max.assert_not_called() + + +def _make_prisma_with_end_user_row(spend: float | None): + prisma = MagicMock() + prisma.db.litellm_endusertable.find_unique = AsyncMock( + return_value=None if spend is None else MagicMock(spend=spend) + ) + return prisma + + +@pytest.mark.asyncio +async def test_get_current_spend_end_user_floor_admits_after_a_reset_on_a_stale_worker(monkeypatch): + """The reset job zeroes LiteLLM_EndUserTable.spend and the shared counter, but it + evicts the cached end-user object only on the worker that ran the reset. Every + other worker still passes the pre-reset spend as fallback_spend, and that stale + copy must not out-vote the reset row.""" + fake_cache = _make_spend_counter_cache(redis_get_value=0.0) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + prisma = _make_prisma_with_end_user_row(spend=0.0) + monkeypatch.setattr(ps, "prisma_client", prisma) + result = await ps.get_current_spend( - counter_key="spend:end_user:e1", + counter_key="spend:end_user:customer-42", + fallback_spend=0.000032, + max_budget=0.00003, + fallback_authoritative=True, + ) + + assert result == 0.0 + prisma.db.litellm_endusertable.find_unique.assert_awaited_once_with(where={"user_id": "customer-42"}) + fake_cache.redis_cache.async_set_max.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_current_spend_end_user_floor_repairs_a_stale_low_counter(monkeypatch): + """After a Redis restart the end-user counter can sit below the recorded spend; + the row wins and the shared counter is raised so other workers stop admitting on + the stale value.""" + fake_cache = _make_spend_counter_cache(redis_get_value=2.0) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "prisma_client", _make_prisma_with_end_user_row(spend=12.0)) + + result = await ps.get_current_spend( + counter_key="spend:end_user:customer-42", + fallback_spend=12.0, + max_budget=10.0, + ) + + assert result == 12.0 + fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key="spend:end_user:customer-42", value=12.0) + + +@pytest.mark.asyncio +async def test_get_current_spend_end_user_without_a_row_keeps_the_cached_spend(monkeypatch): + fake_cache = _make_spend_counter_cache(redis_get_value=0.0) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "prisma_client", _make_prisma_with_end_user_row(spend=None)) + + result = await ps.get_current_spend( + counter_key="spend:end_user:customer-42", fallback_spend=20.0, max_budget=10.0, ) assert result == 20.0 - # no DB row to repair against, so the shared counter is left untouched fake_cache.redis_cache.async_set_max.assert_not_called() From b4fd63f621cd8dd98755944aeca0e9dfaa78daa1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:53:45 -0700 Subject: [PATCH 07/10] chore(proxy): annotate the new spend-counter test locals and correct the floor comments --- litellm/proxy/proxy_server.py | 7 ++++--- .../proxy/db/test_spend_counter_reseed.py | 11 ++++++----- .../proxy/proxy_server/test_spend_counters.py | 17 +++++++++-------- 3 files changed, 19 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 47c0811d903..9ad3c910f58 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2477,7 +2477,8 @@ async def get_current_spend( authoritative source depends on the counter: primary key/team/user/org counters read the DB row; per-window counters (``window_start`` supplied) read the maintained window-spend row and only aggregate spend logs when - that row is missing or stale; end-user/tag counters have no DB row, so the caller's + that row is missing or stale; end-user counters read ``LiteLLM_EndUserTable``, the + row the budget reset zeroes; tag counters have no DB row, so the caller's ``fallback_spend`` (loaded fresh in auth) is authoritative. The DB read is skipped for healthy primary counters (counter at or above recorded spend) and cached in-process for a few seconds, so a persistently stale counter @@ -2511,8 +2512,8 @@ async def get_current_spend( await _repair_stale_spend_counter(counter_key=counter_key, db_spend=authoritative) return authoritative elif fallback_spend > current: - # end-user / tag counters have no DB row; fallback_spend is the - # authoritative recorded value loaded in auth. + # nothing to read (tag counters, an end user without a row or a DB client, a + # failed read); fallback_spend is the authoritative recorded value loaded in auth. return fallback_spend # Opt-in hard guarantee: when the spend backing this admit decision came diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index 3bd6d93328d..8cb3fc665eb 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -9,6 +9,7 @@ from __future__ import annotations from datetime import datetime, timedelta, timezone from types import SimpleNamespace +from typing import Final import pytest @@ -255,9 +256,9 @@ async def test_coalesced_window_seeds_a_cold_counter_from_the_row(): @pytest.mark.asyncio async def test_end_user_from_db_reads_the_end_user_row_by_user_id(): - prisma = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=0.0)) + prisma: Final = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=0.0)) - result = await SpendCounterReseed.end_user_from_db( + result: Final = await SpendCounterReseed.end_user_from_db( prisma_client=prisma, counter_key="spend:end_user:customer-42" ) @@ -267,7 +268,7 @@ async def test_end_user_from_db_reads_the_end_user_row_by_user_id(): @pytest.mark.asyncio async def test_end_user_from_db_returns_the_recorded_spend(): - prisma = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=12.5)) + prisma: Final = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=12.5)) assert ( await SpendCounterReseed.end_user_from_db(prisma_client=prisma, counter_key="spend:end_user:customer-42") @@ -278,7 +279,7 @@ async def test_end_user_from_db_returns_the_recorded_spend(): @pytest.mark.asyncio @pytest.mark.parametrize("counter_key", ["spend:key:hashed", "spend:team:t1", "spend:tag:t1"]) async def test_end_user_from_db_ignores_other_counter_kinds_without_touching_the_db(counter_key): - prisma = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="x", spend=5.0)) + prisma: Final = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="x", spend=5.0)) assert await SpendCounterReseed.end_user_from_db(prisma_client=prisma, counter_key=counter_key) is None assert prisma.db.litellm_endusertable.where_clauses == [] @@ -309,7 +310,7 @@ async def test_end_user_from_db_returns_none_without_a_row_a_client_or_on_db_err async def test_from_db_still_never_reads_the_end_user_row(): """A cold end-user counter keeps seeding from the cached end-user object the auth path already loaded; the row is read only as the budget floor.""" - prisma = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=5.0)) + prisma: Final = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=5.0)) assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:end_user:customer-42") is None assert prisma.db.litellm_endusertable.where_clauses == [] diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 7cc1390fcbd..ef6f8120c82 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -22,6 +22,7 @@ from __future__ import annotations import asyncio from datetime import datetime +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -233,7 +234,7 @@ async def test_get_current_spend_floors_end_user_tag_against_fallback(monkeypatc monkeypatch.setattr(ps.SpendCounterReseed, "from_db", AsyncMock(return_value=None)) for counter_key in ("spend:end_user:e1", "spend:tag:t1"): - result = await ps.get_current_spend( + result: Final = await ps.get_current_spend( counter_key=counter_key, fallback_spend=20.0, max_budget=10.0, @@ -245,7 +246,7 @@ async def test_get_current_spend_floors_end_user_tag_against_fallback(monkeypatc def _make_prisma_with_end_user_row(spend: float | None): - prisma = MagicMock() + prisma: Final = MagicMock() prisma.db.litellm_endusertable.find_unique = AsyncMock( return_value=None if spend is None else MagicMock(spend=spend) ) @@ -258,9 +259,9 @@ async def test_get_current_spend_end_user_floor_admits_after_a_reset_on_a_stale_ evicts the cached end-user object only on the worker that ran the reset. Every other worker still passes the pre-reset spend as fallback_spend, and that stale copy must not out-vote the reset row.""" - fake_cache = _make_spend_counter_cache(redis_get_value=0.0) + fake_cache: Final = _make_spend_counter_cache(redis_get_value=0.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - prisma = _make_prisma_with_end_user_row(spend=0.0) + prisma: Final = _make_prisma_with_end_user_row(spend=0.0) monkeypatch.setattr(ps, "prisma_client", prisma) result = await ps.get_current_spend( @@ -280,11 +281,11 @@ async def test_get_current_spend_end_user_floor_repairs_a_stale_low_counter(monk """After a Redis restart the end-user counter can sit below the recorded spend; the row wins and the shared counter is raised so other workers stop admitting on the stale value.""" - fake_cache = _make_spend_counter_cache(redis_get_value=2.0) + fake_cache: Final = _make_spend_counter_cache(redis_get_value=2.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "prisma_client", _make_prisma_with_end_user_row(spend=12.0)) - result = await ps.get_current_spend( + result: Final = await ps.get_current_spend( counter_key="spend:end_user:customer-42", fallback_spend=12.0, max_budget=10.0, @@ -296,11 +297,11 @@ async def test_get_current_spend_end_user_floor_repairs_a_stale_low_counter(monk @pytest.mark.asyncio async def test_get_current_spend_end_user_without_a_row_keeps_the_cached_spend(monkeypatch): - fake_cache = _make_spend_counter_cache(redis_get_value=0.0) + fake_cache: Final = _make_spend_counter_cache(redis_get_value=0.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "prisma_client", _make_prisma_with_end_user_row(spend=None)) - result = await ps.get_current_spend( + result: Final = await ps.get_current_spend( counter_key="spend:end_user:customer-42", fallback_spend=20.0, max_budget=10.0, From 4ffd2ffb25017c861a350f24a27b6441ceeb2fd1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:00:40 -0700 Subject: [PATCH 08/10] test(proxy): parametrize the stale end-user counter case so no Final local sits in a loop --- .../proxy/proxy_server/test_spend_counters.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index ef6f8120c82..86e97a334df 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -223,24 +223,24 @@ async def test_get_current_spend_floor_caches_db_read(monkeypatch): @pytest.mark.asyncio -async def test_get_current_spend_floors_end_user_tag_against_fallback(monkeypatch): +@pytest.mark.parametrize("counter_key", ("spend:end_user:e1", "spend:tag:t1")) +async def test_get_current_spend_floors_end_user_tag_against_fallback(monkeypatch, counter_key): """Tag counters have no DB row (from_db returns None), and an end-user counter has none to read without a DB client. When such a counter is stale-low, enforcement falls back to the caller's recorded spend (loaded fresh in auth) instead of trusting the stale counter.""" - fake_cache = _make_spend_counter_cache(redis_get_value=2.0) + fake_cache: Final = _make_spend_counter_cache(redis_get_value=2.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", AsyncMock(return_value=None)) - for counter_key in ("spend:end_user:e1", "spend:tag:t1"): - result: Final = await ps.get_current_spend( - counter_key=counter_key, - fallback_spend=20.0, - max_budget=10.0, - ) + result: Final = await ps.get_current_spend( + counter_key=counter_key, + fallback_spend=20.0, + max_budget=10.0, + ) - assert result == 20.0 + assert result == 20.0 # no DB row to repair against, so the shared counter is left untouched fake_cache.redis_cache.async_set_max.assert_not_called() From cd113c3a2e46bd8d468ea48219d861c107e06260 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:13:51 -0700 Subject: [PATCH 09/10] ci: allowlist the bounded _unqualified qualifier peel in the recursion detector --- tests/code_coverage_tests/recursive_detector.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 0578dc60119..e9f87ba6cae 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -66,6 +66,7 @@ IGNORE_FUNCTIONS = [ "_json_safe", # max depth set (_MAX_DEPTH) plus a seen-ids cycle guard for self-referential input. "_redact_agent_params_tree", # max depth set (default 10), same shape as _redact_sensitive_litellm_params. "_restore_redacted_nested_value", # max depth set (default 10), mirrors _redact_agent_params_tree on the write side. + "_unqualified", # bounded by the qualifier depth of a static TypedDict annotation (Annotated, Required/NotRequired, ReadOnly around one type, no cycles possible). ] From e2741b564376f8b30ec93848cf2d9c3e94aa32af Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 4 Sep 2026 18:24:16 -0700 Subject: [PATCH 10/10] fix(datadog_llm_obs): keep the guardrail audit record under message redaction (#39702) * fix(datadog_llm_obs): keep the guardrail audit record under message redaction Redaction nulled `guardrail_information` on the span whole, so an operator running `turn_off_message_logging` (or a caller sending `x-litellm-enable-message-redaction`) lost the record of which guardrails ran, what they returned, and what they masked. Four of the record's fields can quote the prompt; the rest report what the guardrail decided without reproducing it. Replace only those four, the way `_sanitize_guardrail_information_for_spend_logs` already does for spend logs, and declare the field list once in `litellm/types/utils.py` so both readers share it. * fix(datadog_llm_obs): keep a lone guardrail record, and test through the span Review round 1. A guardrail that writes the metadata key itself leaves a single record where the type says list, which Prometheus already normalizes at `_guardrail_overhead_seconds`. Redaction dropped that shape and the latency extraction raised on it, so the span was lost outright. Normalize once and use it in both places. The new tests now drive `create_llm_obs_payload` instead of reading the module's private helpers and the record's declared field names. --- .../integrations/datadog/datadog_llm_obs.py | 64 +++++++-- .../spend_tracking/spend_tracking_utils.py | 10 +- litellm/types/utils.py | 44 ++++++ .../datadog/test_datadog_llm_obs.py | 129 +++++++++++++++++- 4 files changed, 228 insertions(+), 19 deletions(-) diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index ec86c0ae1d9..728bf41856f 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -19,7 +19,7 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid -from litellm.constants import REDACTED_BY_LITELLM +from litellm.constants import REDACTED_BY_LITELLM, REDACTED_BY_LITELM_STRING from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.integrations.datadog.datadog_handler import ( get_datadog_base_url_from_env, @@ -46,9 +46,10 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens from litellm.types.integrations.datadog_llm_obs import * from litellm.types.utils import ( + AUDIT_GUARDRAIL_FIELDS, + PROMPT_CARRYING_GUARDRAIL_FIELDS, PROMPT_QUOTING_ROUTING_DECISION_FIELDS, CallTypes, - StandardLoggingGuardrailInformation, StandardLoggingPayload, StandardLoggingPayloadErrorInformation, ) @@ -60,6 +61,8 @@ _SAFE_REDACTED_MESSAGE_ROLES: Final = frozenset( {"agent", "assistant", "developer", "function", "model", "system", "tool", "user"} ) +_CLASSIFIED_GUARDRAIL_FIELDS: Final = AUDIT_GUARDRAIL_FIELDS | PROMPT_CARRYING_GUARDRAIL_FIELDS + _PROMPT_CARRYING_METADATA_FIELDS: Final = frozenset( { "routing_decision", @@ -108,6 +111,49 @@ def _router_span_fields( ) +def _guardrail_entries(guardrail_information: object) -> tuple[Mapping[str, object], ...]: + """The guardrail records as a sequence, whatever shape the payload carries. + + `guardrail_information` is typed as a list, but a guardrail that writes the metadata key itself + can leave a single record there; Prometheus normalizes the same shape at + `_guardrail_overhead_seconds`. + """ + if isinstance(guardrail_information, Mapping): + return (guardrail_information,) + if isinstance(guardrail_information, (list, tuple)): + return tuple(entry for entry in guardrail_information if isinstance(entry, Mapping)) + return () + + +def _guardrail_entry_without_prompt_carriers(entry: Mapping[str, object]) -> Mapping[str, object]: + """One guardrail record kept as its audit fields, with the prompt-quoting ones marked redacted. + + Built as an allow-list rather than a deny-list: a key neither set classifies is dropped, so a + guardrail that records its own extra detail cannot put the caller's prompt on a redacted span. + """ + return { # mutable-ok: a fresh record built per entry, handed straight to the span serializer + field: REDACTED_BY_LITELM_STRING if field in PROMPT_CARRYING_GUARDRAIL_FIELDS else value + for field, value in entry.items() + if field in _CLASSIFIED_GUARDRAIL_FIELDS + } + + +def _guardrail_information_without_prompt_carriers( + guardrail_information: object, +) -> tuple[Mapping[str, object], ...] | None: + """The guardrail records reduced to what a redacted span may carry. + + Redaction removes the prompt, not the record that a guardrail ran: the name, mode, status, + timings and masked-entity counts are what an operator reads to answer whether a guardrail + caught anything on a request, and none of them reproduce the prompt. Field-level rather than + dropping the list, which is what `_sanitize_guardrail_information_for_spend_logs` already does + for spend logs. + """ + if guardrail_information is None: + return None + return tuple(_guardrail_entry_without_prompt_carriers(entry) for entry in _guardrail_entries(guardrail_information)) + + def _metadata_without_prompt_carriers(standard_logging_metadata: Mapping[str, Any]) -> Mapping[str, Any]: """The metadata minus the records that quote prompts, tool arguments, tool results, or retrieved text.""" return MappingProxyType( @@ -872,7 +918,9 @@ class DataDogLLMObsLogger(CustomBatchLogger): "cache_key": standard_logging_payload.get("cache_key", "unknown"), "saved_cache_cost": standard_logging_payload.get("saved_cache_cost", 0), "guardrail_information": ( - None if redact_prompt_text else standard_logging_payload.get("guardrail_information", None) + _guardrail_information_without_prompt_carriers(standard_logging_payload.get("guardrail_information")) + if redact_prompt_text + else standard_logging_payload.get("guardrail_information", None) ), "is_streamed_request": self._get_stream_value_from_payload(standard_logging_payload), "latency_metrics": dict(self._get_latency_metrics(standard_logging_payload)), @@ -904,14 +952,12 @@ class DataDogLLMObsLogger(CustomBatchLogger): latency_metrics["litellm_overhead_time_ms"] = litellm_overhead_ms # Guardrail overhead latency - guardrail_info: Final[list[StandardLoggingGuardrailInformation] | None] = standard_logging_payload.get( - "guardrail_information" - ) - if guardrail_info is not None: + guardrail_info: Final = _guardrail_entries(standard_logging_payload.get("guardrail_information")) + if guardrail_info: total_duration = 0.0 for info in guardrail_info: - _guardrail_duration_seconds: float | None = info.get("duration") - if _guardrail_duration_seconds is not None: + _guardrail_duration_seconds = info.get("duration") + if isinstance(_guardrail_duration_seconds, (int, float, str)): total_duration += float(_guardrail_duration_seconds) if total_duration > 0: diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 9ea2170d6ab..8a06bf68b81 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -37,6 +37,7 @@ from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload, SpendLogsR from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.proxy.utils import PrismaClient, hash_token from litellm.types.utils import ( + PROMPT_CARRYING_GUARDRAIL_FIELDS, CallTypes, CostBreakdown, StandardLoggingGuardrailInformation, @@ -1073,13 +1074,6 @@ def _sanitize_guardrail_information_for_spend_logs( return [_redact_prompt_fields_in_guardrail_entry(entry) for entry in entries if isinstance(entry, dict)] -_PROMPT_CARRYING_GUARDRAIL_FIELDS: Final = ( - "guardrail_request", - "guardrail_response", - "match_details", - "classification", -) - _NUMERIC_COMPRESSION_STAT_KEYS: Final = ( "tokens_before", "tokens_after", @@ -1114,7 +1108,7 @@ def _redact_prompt_fields_in_guardrail_entry( preserved_stats: Final = _numeric_compression_stats_from_guardrail_response(entry.get("guardrail_response")) redacted: Final[StandardLoggingGuardrailInformation] = { **entry, - **{key: REDACTED_BY_LITELM_STRING for key in _PROMPT_CARRYING_GUARDRAIL_FIELDS if key in entry}, + **{key: REDACTED_BY_LITELM_STRING for key in PROMPT_CARRYING_GUARDRAIL_FIELDS if key in entry}, } if preserved_stats is None: return redacted diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6745e65f81e..238986f86a6 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3080,6 +3080,50 @@ class GuardrailMode(TypedDict, total=False): GuardrailStatus = Literal["success", "guardrail_intervened", "guardrail_failed_to_respond", "not_run"] +# Fields on a guardrail record whose values can quote the caller's prompt: the payload sent to the +# guardrail, the provider response that echoes it back, and the two first-party hooks that inline +# prompt substrings (``block_code_execution`` and ``litellm_content_filter``). Every other field +# reports what the guardrail decided without reproducing the prompt, so redaction replaces these +# four and keeps the rest of the record. +PROMPT_CARRYING_GUARDRAIL_FIELDS: Final[frozenset[str]] = frozenset( + { + "guardrail_request", + "guardrail_response", + "match_details", + "classification", + } +) + +# The rest of the record: what the guardrail is, what it decided, how long it took and what it cost. +# None of these reproduce the prompt, so a redacted record keeps them and stays explainable. +# `test_every_guardrail_field_is_classified` fails if a field is added to the record without being +# placed in one set or the other, so a new field is dropped from redacted records rather than +# shipped unexamined. +AUDIT_GUARDRAIL_FIELDS: Final[frozenset[str]] = frozenset( + { + "guardrail_name", + "guardrail_provider", + "guardrail_mode", + "guardrail_status", + "start_time", + "end_time", + "duration", + "masked_entity_count", + "guardrail_id", + "policy_template", + "detection_method", + "confidence_score", + "patterns_checked", + "alert_recipients", + "risk_score", + "violation_categories", + "guardrail_action", + "guardrail_usage", + "guardrail_cost", + "guardrail_cost_in_spend", + } +) + class StandardLoggingGuardrailInformation(TypedDict, total=False): guardrail_name: str | None diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py index 34c62864c4e..0555447e34f 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py @@ -12,7 +12,7 @@ spelling of prompt-cache counts (`prompt_tokens_details.cached_tokens`). import json import os from datetime import datetime, timedelta -from typing import Any +from typing import Any, Final from unittest.mock import patch import pytest @@ -20,6 +20,8 @@ import pytest import litellm from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import StandardLoggingGuardrailInformation TOOL_DEFINITION: dict[str, Any] = { "type": "function", @@ -631,10 +633,133 @@ def test_redaction_drops_every_prompt_carrying_metadata_record(logger: DataDogLL for record in sensitive_metadata: assert record not in redacted["meta"]["metadata"] assert record in unredacted["meta"]["metadata"] - assert redacted["meta"]["metadata"]["guardrail_information"] is None + assert redacted["meta"]["metadata"]["guardrail_information"] == [ + {"guardrail_name": "g", "guardrail_request": "REDACTED_BY_LITELM"} + ] # the record survives; only the field quoting the prompt is replaced assert unredacted["meta"]["metadata"]["guardrail_information"] is not None +_AUDIT_RECORD: Final[StandardLoggingGuardrailInformation] = StandardLoggingGuardrailInformation( + guardrail_name="bedrock-pii", + guardrail_provider="bedrock", + guardrail_mode=GuardrailEventHooks.pre_call, + guardrail_status="guardrail_intervened", + guardrail_response={"action": "MASK", "match": "alice@acme.com"}, + match_details=[{"pattern": "email", "match": "alice@acme.com"}], + classification="the user asked for alice@acme.com", + masked_entity_count={"EMAIL": 2}, + violation_categories=["pii"], + duration=0.01, +) + + +def _payload_with_guardrail_record(guardrail_information: object) -> dict[str, Any]: + payload = build_payload() + payload["standard_logging_object"]["guardrail_information"] = guardrail_information + return payload + + +def test_redaction_keeps_the_guardrail_audit_record(logger: DataDogLLMObsLogger) -> None: + """Redaction removes the prompt, not the operator's record that a guardrail intervened.""" + redacted = _span_json( + _redacting_logger(turn_off_message_logging=True), + _payload_with_guardrail_record([dict(_AUDIT_RECORD)]), + ) + record = redacted["meta"]["metadata"]["guardrail_information"][0] + + for field in ("guardrail_request", "guardrail_response", "match_details", "classification"): + assert record.get(field, "REDACTED_BY_LITELM") == "REDACTED_BY_LITELM" + assert record["guardrail_name"] == "bedrock-pii" + assert record["guardrail_provider"] == "bedrock" + assert record["guardrail_mode"] == "pre_call" + assert record["guardrail_status"] == "guardrail_intervened" + assert record["masked_entity_count"] == {"EMAIL": 2} + assert record["violation_categories"] == ["pii"] + assert record["duration"] == 0.01 + assert "alice@acme.com" not in safe_dumps(redacted["meta"]["metadata"]) + + +def test_a_caller_supplied_redaction_header_cannot_blank_the_guardrail_record( + logger: DataDogLLMObsLogger, +) -> None: + """Any key may redact its own prompts with the header; none may erase what a guardrail caught.""" + payload = _payload_with_guardrail_record([dict(_AUDIT_RECORD)]) + payload["litellm_params"] = {"metadata": {"headers": {"x-litellm-enable-message-redaction": "true"}}} + + span = _span_json(logger, payload) + record = span["meta"]["metadata"]["guardrail_information"][0] + + assert span["meta"]["input"]["messages"] == [{"role": "user", "content": "redacted-by-litellm"}] + assert record["guardrail_status"] == "guardrail_intervened" + assert record["masked_entity_count"] == {"EMAIL": 2} + assert "alice@acme.com" not in safe_dumps(span["meta"]["metadata"]) + + +def test_a_guardrails_own_extra_field_never_reaches_a_redacted_span(logger: DataDogLLMObsLogger) -> None: + """A guardrail may record whatever it likes; only classified fields survive redaction.""" + span = _span_json( + _redacting_logger(turn_off_message_logging=True), + _payload_with_guardrail_record([{**_AUDIT_RECORD, "matched_text": "the caller asked about alice@acme.com"}]), + ) + record = span["meta"]["metadata"]["guardrail_information"][0] + + assert "matched_text" not in record + assert record["guardrail_status"] == "guardrail_intervened" + assert "alice@acme.com" not in safe_dumps(span["meta"]["metadata"]) + + +def test_a_lone_guardrail_record_survives_redaction(logger: DataDogLLMObsLogger) -> None: + """A guardrail that writes the metadata key itself leaves one record, not a list of them.""" + span = _span_json( + _redacting_logger(turn_off_message_logging=True), + _payload_with_guardrail_record(dict(_AUDIT_RECORD)), + ) + metadata = span["meta"]["metadata"] + + assert metadata["guardrail_information"] == [ + { + "guardrail_name": "bedrock-pii", + "guardrail_provider": "bedrock", + "guardrail_mode": "pre_call", + "guardrail_status": "guardrail_intervened", + "guardrail_response": "REDACTED_BY_LITELM", + "match_details": "REDACTED_BY_LITELM", + "classification": "REDACTED_BY_LITELM", + "masked_entity_count": {"EMAIL": 2}, + "violation_categories": ["pii"], + "duration": 0.01, + } + ] + assert metadata["latency_metrics"]["guardrail_overhead_time_ms"] == 10.0 + + +@pytest.mark.parametrize("guardrail_information", [None, [], 5, "abc", [None, "x"], {}]) +def test_odd_guardrail_shapes_still_produce_a_span( + guardrail_information: object, +) -> None: + """The redacted branch replaced an expression that could not fail, so it must not start failing.""" + span = _span_json( + _redacting_logger(turn_off_message_logging=True), + _payload_with_guardrail_record(guardrail_information), + ) + + assert span["meta"]["input"]["messages"] == [{"role": "user", "content": "redacted-by-litellm"}] + assert span["meta"]["metadata"]["guardrail_information"] in (None, [], [{}]) + + +def test_a_redacted_span_carries_every_declared_guardrail_field() -> None: + """A field added to the record without a redaction decision would be dropped, so it fails here.""" + declared = dict.fromkeys(StandardLoggingGuardrailInformation.__annotations__, "alice@acme.com") + payload = _payload_with_guardrail_record([{**declared, "duration": 0.01}]) + + span = _span_json(_redacting_logger(turn_off_message_logging=True), payload) + record = span["meta"]["metadata"]["guardrail_information"][0] + + assert set(record) == set(declared) + for field in ("guardrail_request", "guardrail_response", "match_details", "classification"): + assert record[field] == "REDACTED_BY_LITELM" + + def test_tool_definitions_accept_the_bare_anthropic_shape(logger: DataDogLLMObsLogger) -> None: """The Anthropic surface declares tools unwrapped, with input_schema instead of parameters.""" payload = build(