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.
This commit is contained in:
mateo-berri 2026-09-04 17:39:46 -07:00
parent 976f8625f3
commit 89086db282
4 changed files with 196 additions and 23 deletions

View file

@ -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:"):

View file

@ -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

View file

@ -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 == []

View file

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