diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 66bbda1ca4e..c1f9407cdad 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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): diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 97999cbb6d7..0d12b012c18 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 6045b64023d..42e56ceabd3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -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