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,