scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 137

fix(proxy): reset a key's budget-window counters on spend reset (#38686)

Browse source 
* fix(proxy): reset a key's budget-window counters and broadcast the reset cross-pod

/key/{id}/reset_spend already reset the key's lifetime spend counter in
Redis, but a key with its own budget_limits (an extra time-windowed cap,
e.g. a daily budget layered on top of the lifetime max_budget) kept its
window counter untouched, so the key stayed 429'd on
"ExceededBudget: Key over <duration> budget" even after the admin action
reported spend back to $0.

Force-expire each window on reset: zero its Redis counter and restart the
window from now (window_start is derived as reset_at - budget_duration,
so reset_at must float to now + duration, not the next calendar boundary
get_budget_reset_time gives key creation - that boundary can still be
in the past relative to the spend that triggered the block).

Also close a second, narrower race: _delete_cache_key_object evicted the
cached key object only on the handling pod, so another pod could keep
serving the stale pre-reset object (and re-derive the pre-reset spend
counter via its own floor-marker cache) until its own TTL expired. It now
broadcasts the eviction, matching the pattern already used for team,
team-member, customer, and tag caches.

* fix(proxy): evict the cached key object after every reset_spend DB write

Greptile P1: eviction ran before the window-reset DB write committed, so a
request racing the reset could re-fetch and re-cache the pre-write row,
pinning that pod to the stale budget_limits for the rest of its own cache
TTL even after the write went through. Move the eviction to run last.

* test: pin real cache state and satisfy the test-quality gate

test_delete_cache_key_object_broadcasts_invalidation now asserts a real
UserApiKeyCache no longer holds the evicted entry, rather than only
inspecting a mock's call args. Suppress test-quality-ok on the
hash_token/_check_proxy_or_team_admin_for_key/_delete_cache_key_object/
publish_auth_cache_invalidation patches: none has an HTTP boundary to
fake, matching the pattern the file already uses for these same targets.

* fix(proxy): narrow budget_limits by the str branch, not the list branch

isinstance(x, list) in the else branch still leaves Sequence[object] | str
(a tuple satisfies Sequence without being a list), so json.loads() saw a
possible non-str argument. Check isinstance(x, str) instead, which narrows
each branch to exactly the type it needs.

* fix(proxy): persist advanced budget-window boundaries before zeroing counters

Greptile P1: publishing a zeroed window counter before the new reset_at
committed let a request racing the write compute window_start from the
stale boundary, re-sum the historical spend log rows the reset was
clearing, and put the counter right back above budget. Compute every
window's new boundary, persist all of them in one DB write, then zero
each window's Redis counter only once that write has landed.
This commit is contained in:
Yassin Kortam 2026-08-28 15:13:03 -07:00 • committed by GitHub
parent 42d278cfad
commit 671f89d8bc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 399 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/proxy/auth/auth_checks.py
View file

@ -2596,6 +2596,11 @@ async def _delete_cache_key_object(
dropped before the Redis round trip. Letting a cache-backend error raise here therefore reports
failure for work that succeeded without making the cache any less stale; the leftover Redis
entry expires at its TTL either way.
Also broadcasts the eviction to every other worker (LIT-3803): auth serves this object
cache-first with no freshness check, so a worker that never receives the broadcast keeps
admitting requests against the pre-mutation object (e.g. a just-reset spend) until its own
copy's TTL expires.
"""
key: Final = hashed_token
@ -2612,6 +2617,8 @@ async def _delete_cache_key_object(
e,
)
await publish_auth_cache_invalidation(cache_key=key)
async def delete_cache_key_objects(
hashed_tokens: Sequence[str],
@ -2623,8 +2630,9 @@ async def delete_cache_key_objects(
`/key/delete`. Auth resolves a cached key object without re-reading its team, so a key left
cached after its row is gone keeps buying access until its TTL expires.
Evicting locally only reaches this worker, so each token is also broadcast: a deleted key left
in a peer worker's in-memory cache still authenticates there until its TTL expires.
Evicting locally only reaches this worker; `_delete_cache_key_object` itself broadcasts each
token, so a deleted key left in a peer worker's in-memory cache still authenticates there until
its TTL expires.
Best-effort per key: the rows are already deleted by the time this runs, so an unreachable
cache backend must not abort the caller partway through its own cascade.
@ -2648,7 +2656,6 @@ async def delete_cache_key_objects(
hashed_token,
result,
)
await publish_auth_cache_invalidation(cache_key=hashed_token)
class _TeamNotFoundDetail(TypedDict):

158
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -62,6 +62,9 @@ from litellm.proxy.auth.auth_utils import (
enforce_output_token_estimates_are_admin_only,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
publish_auth_cache_invalidation,
)
from litellm.proxy.common_utils.callback_config_validation import logging_metadata_config_error
from litellm.proxy.common_utils.callback_utils import (
decrypt_callback_vars,
@ -5171,6 +5174,125 @@ def _validate_reset_spend_value(reset_to: object, key_in_db: LiteLLM_Verificatio
return reset_to
async def _set_spend_counter_with_floor_and_broadcast(counter_key: str, value: float) -> None:
"""
Set a Redis-backed spend counter to `value`, mirror it into the short-lived
spend_db_floor marker `_authoritative_floor_spend` reads, and broadcast both
to every worker (LIT-3803 pattern: setting, not deleting, means a worker's
own self-delivered broadcast still carries the reset value forward).
Without the floor marker, `_authoritative_floor_spend` can re-derive a
stale, pre-reset value from a marker another worker cached moments earlier
and raise the just-reset counter right back up via `_repair_stale_spend_counter`.
Without the broadcast, a worker that already cached the pre-reset key object
or floor marker keeps enforcing against it until its own TTL expires.
"""
from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=value, ttl=60)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=value, ttl=60)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to update spend counter %s in Redis: %s. "
"Budget checks may use stale value until counter expires.",
counter_key,
redis_err,
)
floor_key: Final = f"spend_db_floor:{counter_key}"
spend_counter_cache.in_memory_cache.set_cache(key=floor_key, value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS)
await publish_auth_cache_invalidation(cache_key=counter_key, new_value=value, ttl=60)
await publish_auth_cache_invalidation(cache_key=floor_key, new_value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS)
def _budget_limit_windows(budget_limits: Sequence[object] | str | None) -> tuple[Mapping[str, object], ...]:
"""Coerce a key's stored `budget_limits` into a tuple of plain window dicts.
It is a DB Json column, so a caller reading it straight off `find_unique`
gets an already-parsed list; one reading it off `json.dumps`'d text (or a
raw SQL row) gets the string form. Either way each entry is a plain dict,
except wherever a caller already validated the field through a pydantic
model (e.g. `UserAPIKeyAuth.budget_limits`), which yields `BudgetLimitEntry`
objects instead -- coerced here via `model_dump()`, matching
`_set_budget_reset_at`'s identical coercion in team_endpoints.py.
"""
if not budget_limits:
return ()
raw_windows: Final = json.loads(budget_limits) if isinstance(budget_limits, str) else budget_limits
return tuple(raw_window if isinstance(raw_window, dict) else raw_window.model_dump() for raw_window in raw_windows)
def _advance_one_key_budget_window(window: Mapping[str, object]) -> Mapping[str, object]:
"""Restart one budget window from now, by advancing its `reset_at`.
`window_start` is derived elsewhere as `reset_at - budget_duration`
(`get_budget_window_start`), so `reset_at` must be set to `now +
budget_duration` -- a window floating from THIS moment -- to make
`window_start` land at `now` and exclude the historical spend that
triggered the block. Reusing `get_budget_reset_time`/
`ResetBudgetJob._reset_expired_window`'s calendar-standardized boundary
(e.g. "next midnight") would not do that: for a "1d" window `next
midnight - 1d` is simply the START of the calendar day already in
progress, which still covers that spend. That reuse is only safe for the
scheduled job, which runs right as `reset_at` naturally elapses, so the
elapsed boundary it computes is already close to "now". A manual reset
can happen at any point mid-window, so it needs the floating form
instead. A window with no `budget_duration` is returned unchanged.
"""
duration = window.get("budget_duration")
if not isinstance(duration, str) or not duration:
return window
new_reset_at: Final = datetime.now(timezone.utc) + timedelta(seconds=duration_in_seconds(duration))
return { # mutable-ok: this is the JSON payload persisted to budget_limits' Json column, which requires a plain dict
**window,
"reset_at": new_reset_at.isoformat(),
}
async def _reset_key_budget_windows(
prisma_client: PrismaClient,
hashed_api_key: str,
budget_limits: Sequence[object] | str | None,
) -> None:
"""Force-expire every one of a key's own `budget_limits` windows (extra
time-windowed caps layered on top of the lifetime max_budget, e.g. a daily
limit) so a manual spend reset also clears them, not just the lifetime
counter.
Persists the advanced `reset_at` boundaries BEFORE zeroing any window's
Redis counter, not after: a window counter reading zero is only durable
once every reader recomputing its floor from the DB sees the new
boundary too (`get_current_spend` re-derives a window counter from real
`LiteLLM_SpendLogs` rows inside `[window_start, now)` on every read below
max_budget, see its `is_window` branch). Zeroing first would let a
request racing the DB write compute `window_start` from the stale
pre-reset boundary, re-sum the unchanged historical spend, and put the
counter right back where it was before the write ever landed.
"""
windows: Final = _budget_limit_windows(budget_limits)
if not windows:
return
reset_windows: Final = tuple(_advance_one_key_budget_window(w) for w in windows)
# prisma-client-py's typed update() takes plain dict literals for `where`/`data`; there is no
# frozen-mapping equivalent to pass instead.
reset_payload: Final = {"budget_limits": json.dumps(reset_windows, default=str)} # mutable-ok: prisma data kwarg
await VerificationTokenRepository(prisma_client).table.update(
where={"token": hashed_api_key}, # mutable-ok: prisma where kwarg
data=reset_payload,
)
for window in reset_windows:
duration = window.get("budget_duration")
if isinstance(duration, str) and duration:
counter_key = f"spend:key:{hashed_api_key}:window:{duration}"
await _set_spend_counter_with_floor_and_broadcast(counter_key=counter_key, value=0.0)
@router.post(
"/key/{key:path}/reset_spend",
tags=["key management"],
@ -5236,30 +5358,30 @@ async def reset_key_spend_fn(
detail={"error": "Failed to update key spend"},
)
# Reset the lifetime spend counter to the new value (not 0.0, so partial
# resets are reflected correctly), and force-expire any of the key's own
# budget_limits windows, so get_current_spend() returns the correct
# amount for every enforcement check immediately instead of the stale
# pre-reset value.
_counter_key: Final = f"spend:key:{hashed_api_key}"
await _set_spend_counter_with_floor_and_broadcast(counter_key=_counter_key, value=reset_to)
await _reset_key_budget_windows(
prisma_client=prisma_client,
hashed_api_key=hashed_api_key,
budget_limits=_key_in_db.budget_limits,
)
# Evicting the cached key object LAST (after every DB write above has
# committed) matters: a request landing between an earlier eviction and
# a later write would re-fetch and re-cache the pre-write row, pinning
# that pod to the stale budget_limits/spend for the rest of its own
# cache TTL even though the DB is already correct.
await _delete_cache_key_object(
hashed_token=hashed_api_key,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
# Set Redis spend counter to the new value so get_current_spend()
# returns the correct amount immediately instead of the stale pre-reset value.
# We use reset_to (not 0.0) so partial resets are reflected correctly.
from litellm.proxy.proxy_server import spend_counter_cache
_counter_key: Final = f"spend:key:{hashed_api_key}"
spend_counter_cache.in_memory_cache.set_cache(key=_counter_key, value=reset_to, ttl=60)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(key=_counter_key, value=reset_to, ttl=60)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to update spend counter %s in Redis: %s. "
"Budget checks may use stale value until counter expires.",
_counter_key,
redis_err,
)
max_budget: Final = updated_key.max_budget
budget_reset_at: Final = updated_key.budget_reset_at

251
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -1,4 +1,5 @@
import json
from datetime import datetime, timedelta, timezone
import litellm
import pytest
@ -26,7 +27,7 @@ from litellm.proxy._types import (
ResetSpendRequest,
UpdateKeyRequest,
)
from litellm.proxy.auth.auth_checks import _project_cache_key
from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _project_cache_key
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.key_management_endpoints import (
@ -7222,9 +7223,255 @@ async def test_reset_key_spend_success(monkeypatch):
assert response["max_budget"] == 200.0
mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once()
mock_delete_cache.assert_awaited_once()
mock_spend_counter_cache.in_memory_cache.set_cache.assert_called_once_with(
mock_spend_counter_cache.in_memory_cache.set_cache.assert_any_call(
key=f"spend:key:{hashed_key}", value=50.0, ttl=60
)
# spend_db_floor marker is also set to the reset value (LIT-3803 pattern),
# so a request landing on a pod with a warm pre-reset floor marker cannot
# re-derive and re-apply the stale spend.
mock_spend_counter_cache.in_memory_cache.set_cache.assert_any_call(
key=f"spend_db_floor:spend:key:{hashed_key}", value=50.0, ttl=5
)
@pytest.mark.asyncio
async def test_reset_key_spend_resets_budget_windows(monkeypatch):
"""
Regression test: a key with an extra time-windowed budget (`budget_limits`,
e.g. a daily cap layered on top of the lifetime max_budget) must have that
window's own Redis counter reset too, and its `reset_at` advanced, not just
the lifetime spend/counter.
Before the fix, reset_key_spend_fn only reset spend:key:{hash}, leaving
spend:key:{hash}:window:{duration} at its pre-reset value. Since
get_current_spend always re-derives a window counter from real
LiteLLM_SpendLogs rows inside the still-open window, merely zeroing that
counter without also advancing reset_at is not durable either: the very
next request would re-sum the unchanged historical spend and put the
counter right back above the window's max_budget, so
_virtual_key_multi_budget_check kept raising BudgetExceededError (429) on
every request even though the key's own reported spend read $0.
"""
mock_prisma_client = MagicMock()
mock_user_api_key_cache = MagicMock()
mock_proxy_logging_obj = MagicMock()
hashed_key = "hashed-window-budget-key"
key_in_db = LiteLLM_VerificationToken(
token=hashed_key,
user_id="test-user",
spend=80.0,
max_budget=1000.0,
litellm_budget_table=None,
budget_limits=[
{
"budget_duration": "1d",
"max_budget": 50.0,
"reset_at": "2020-01-01T00:00:00+00:00",
}
],
)
updated_key = LiteLLM_VerificationToken(
token=hashed_key,
user_id="test-user",
spend=0.0,
max_budget=1000.0,
budget_reset_at=None,
)
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=key_in_db
)
mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
return_value=updated_key
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
)
mock_spend_counter_cache = MagicMock()
mock_spend_counter_cache.redis_cache = MagicMock()
mock_spend_counter_cache.redis_cache.async_set_cache = AsyncMock()
monkeypatch.setattr(
"litellm.proxy.proxy_server.spend_counter_cache",
mock_spend_counter_cache,
)
with (
patch("litellm.proxy.proxy_server.hash_token") as mock_hash_token, # test-quality-ok: no HTTP boundary; same pattern as test_reset_key_spend_success
patch( # test-quality-ok: no HTTP boundary; same pattern as test_reset_key_spend_success
"litellm.proxy.management_endpoints.key_management_endpoints._check_proxy_or_team_admin_for_key"
) as mock_check_admin,
patch( # test-quality-ok: no HTTP boundary; same pattern as test_reset_key_spend_success
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object"
) as mock_delete_cache,
):
mock_hash_token.return_value = hashed_key
mock_check_admin.return_value = None
mock_delete_cache.return_value = None
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
)
before_call = datetime.now(timezone.utc)
response = await reset_key_spend_fn(
key="sk-test-key",
data=ResetSpendRequest(reset_to=0.0),
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
after_call = datetime.now(timezone.utc)
assert response["spend"] == 0.0
window_counter_key = f"spend:key:{hashed_key}:window:1d"
mock_spend_counter_cache.in_memory_cache.set_cache.assert_any_call(
key=window_counter_key, value=0.0, ttl=60
)
mock_spend_counter_cache.redis_cache.async_set_cache.assert_any_call(
key=window_counter_key, value=0.0, ttl=60
)
mock_spend_counter_cache.in_memory_cache.set_cache.assert_any_call(
key=f"spend_db_floor:{window_counter_key}", value=0.0, ttl=5
)
# The window's DB row must be advanced past the historical spend that
# triggered the block, or the next authoritative-floor recompute re-sums
# the still-open window's spend logs and silently re-inflates the counter.
# reset_at must land at (roughly) now + 1 day: get_budget_window_start
# derives window_start as reset_at - budget_duration, so this is what
# makes window_start land at "now" and exclude the historical spend that
# triggered the block. The next *calendar-aligned* midnight (what a naive
# get_budget_reset_time("1d") call would give) is the wrong value here --
# it would put window_start at the start of the day already in progress,
# which still covers that spend.
assert mock_prisma_client.db.litellm_verificationtoken.update.call_count == 2
window_update_call = mock_prisma_client.db.litellm_verificationtoken.update.call_args_list[1]
assert window_update_call.kwargs["where"] == {"token": hashed_key}
persisted_windows = json.loads(window_update_call.kwargs["data"]["budget_limits"])
assert len(persisted_windows) == 1
assert persisted_windows[0]["budget_duration"] == "1d"
assert persisted_windows[0]["max_budget"] == 50.0
persisted_reset_at = datetime.fromisoformat(persisted_windows[0]["reset_at"])
assert before_call + timedelta(days=1) <= persisted_reset_at <= after_call + timedelta(days=1)
@pytest.mark.asyncio
async def test_reset_key_spend_no_budget_limits_skips_window_reset(monkeypatch):
"""A key with no budget_limits must not trigger any extra DB write beyond
the lifetime spend update; _reset_key_budget_windows should be a no-op."""
mock_prisma_client = MagicMock()
mock_user_api_key_cache = MagicMock()
mock_proxy_logging_obj = MagicMock()
hashed_key = "hashed-no-window-key"
key_in_db = LiteLLM_VerificationToken(
token=hashed_key,
user_id="test-user",
spend=100.0,
max_budget=200.0,
litellm_budget_table=None,
budget_limits=None,
)
updated_key = LiteLLM_VerificationToken(
token=hashed_key,
user_id="test-user",
spend=0.0,
max_budget=200.0,
budget_reset_at=None,
)
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=key_in_db
)
mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
return_value=updated_key
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
)
mock_spend_counter_cache = MagicMock()
mock_spend_counter_cache.redis_cache = None
monkeypatch.setattr(
"litellm.proxy.proxy_server.spend_counter_cache",
mock_spend_counter_cache,
)
with (
patch("litellm.proxy.proxy_server.hash_token") as mock_hash_token, # test-quality-ok: no HTTP boundary; same pattern as test_reset_key_spend_success
patch( # test-quality-ok: no HTTP boundary; same pattern as test_reset_key_spend_success
"litellm.proxy.management_endpoints.key_management_endpoints._check_proxy_or_team_admin_for_key"
) as mock_check_admin,
patch( # test-quality-ok: no HTTP boundary; same pattern as test_reset_key_spend_success
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object"
) as mock_delete_cache,
):
mock_hash_token.return_value = hashed_key
mock_check_admin.return_value = None
mock_delete_cache.return_value = None
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
)
response = await reset_key_spend_fn(
key="sk-test-key",
data=ResetSpendRequest(reset_to=0.0),
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
assert response["spend"] == 0.0
mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once()
@pytest.mark.asyncio
async def test_delete_cache_key_object_broadcasts_invalidation(monkeypatch):
"""
Regression test (LIT-3803 pattern applied to keys): evicting a key's
cached auth object must broadcast the invalidation to every other worker,
or a worker that already cached the pre-mutation object (e.g. pre-reset
spend) keeps serving it until its own local TTL expires, even though this
worker's own cache and the DB have already moved on.
"""
real_user_api_key_cache = UserApiKeyCache()
await real_user_api_key_cache.async_set_cache(
key="hashed-broadcast-key",
value=UserAPIKeyAuth(api_key="sk-broadcast", spend=100.0),
model_type=UserAPIKeyAuth,
)
mock_proxy_logging_obj = MagicMock()
mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock()
with patch( # test-quality-ok: pub/sub broadcast to other workers has no HTTP boundary to fake
"litellm.proxy.auth.auth_checks.publish_auth_cache_invalidation"
) as mock_publish:
mock_publish.return_value = None
await _delete_cache_key_object(
hashed_token="hashed-broadcast-key",
user_api_key_cache=real_user_api_key_cache,
proxy_logging_obj=mock_proxy_logging_obj,
)
# Real, observable state: the cache object itself no longer holds the entry.
assert real_user_api_key_cache.get_cache(key="hashed-broadcast-key") is None
mock_publish.assert_awaited_once_with(cache_key="hashed-broadcast-key")
@pytest.mark.asyncio