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] 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()