From 996ee5635ab79854992661279c2e8db793b3dda4 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 10:35:50 -0700 Subject: [PATCH] perf(proxy): pipeline spend counter increments into one Redis call per request (#40371) * perf(proxy): pipeline spend counter increments into one redis call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): apply surviving spend increments before raising scope error Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): ruff format spend counter helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): settle inner spend counter gathers and fall back per key on pipeline failure Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): suppress BLE001 on pipeline fallback catch Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): invalidate all batched spend counters on pipeline failure Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 331 +++++++---- .../proxy/proxy_server/test_spend_counters.py | 538 ++++++++++-------- .../test_budget_reservation_redis_failure.py | 9 + .../proxy/test_budget_reservation.py | 9 + tests/test_litellm/proxy/test_proxy_server.py | 94 ++- 5 files changed, 609 insertions(+), 372 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7aed0553b4b..c241e66049b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -17,6 +17,7 @@ import time import traceback import warnings from collections.abc import AsyncGenerator, AsyncIterator, Callable, Collection, Mapping, MutableMapping, Sequence +from dataclasses import dataclass from datetime import datetime, timedelta, timezone from types import MappingProxyType, UnionType from typing import ( @@ -131,6 +132,7 @@ from litellm.router_utils.auto_router_tuning_baseline import ( snapshot_tuning_baselines, tuning_limit_violation, ) +from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -2703,6 +2705,12 @@ async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) return fallback_spend, False +@dataclass(frozen=True, slots=True) +class _PendingSpendIncrement: + counter_key: str + increment: float + + async def increment_spend_counters( token: str | None, team_id: str | None, @@ -2737,7 +2745,7 @@ async def increment_spend_counters( cost: Final[float] = response_cost - async def _key_scope(key_token: str) -> None: + async def _key_scope(key_token: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: # key_token arrives pre-hashed from metadata["user_api_key"] (auth flow # hashes raw "sk-..." keys before they reach the callback). The # startswith("sk-") check is a safety net matching update_cache — @@ -2748,30 +2756,29 @@ async def increment_spend_counters( hash_token(token=key_token) if isinstance(key_token, str) and key_token.startswith("sk-") else key_token ) key_counter_key: Final = f"spend:key:{hashed_token}" - if key_counter_key not in reserved_counter_keys: - await _init_and_increment_spend_counter( - counter_key=key_counter_key, - source_cache_key=hashed_token, - increment=cost, + key_pending: Final[tuple[_PendingSpendIncrement, ...]] = ( + () + if key_counter_key in reserved_counter_keys + else ( + await _prepare_spend_counter_increment( + counter_key=key_counter_key, + source_cache_key=hashed_token, + increment=cost, + ), ) - - key_obj: Final[object] = await user_api_key_cache.async_get_cache(key=hashed_token) - if key_obj is None: - return - key_budget_limits = getattr(key_obj, "budget_limits", None) or ( - key_obj.get("budget_limits") if isinstance(key_obj, dict) else None ) - if isinstance(key_budget_limits, str): - key_budget_limits = json.loads(key_budget_limits) - if not isinstance(key_budget_limits, list): - return - for window in key_budget_limits: - duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration - key_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at - key_window_counter = f"spend:key:{hashed_token}:window:{duration}" + + async def _key_window_increment(window: object) -> _PendingSpendIncrement | None: + duration = ( + window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) + ) + key_window_reset_at = ( + window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) + ) + key_window_counter: Final = f"spend:key:{hashed_token}:window:{duration}" key_window_start = get_budget_window_start(window) - if key_window_counter not in reserved_counter_keys: - await _init_and_increment_window_spend_counter( + pending_window: Final = ( + await _prepare_window_spend_counter_increment( counter_key=key_window_counter, entity_type="Key", entity_id=hashed_token, @@ -2779,6 +2786,9 @@ async def increment_spend_counters( window_start=key_window_start, increment=cost, ) + if key_window_counter not in reserved_counter_keys + else None + ) await _enqueue_window_spend_row_update( entity_type=Litellm_EntityType.KEY, entity_id=hashed_token, @@ -2788,33 +2798,48 @@ async def increment_spend_counters( increment=cost, request_started_at=request_started_at, ) + return pending_window - async def _team_scope(scope_team_id: str) -> None: - team_counter_key: Final = f"spend:team:{scope_team_id}" - if team_counter_key not in reserved_counter_keys: - await _init_and_increment_spend_counter( - counter_key=team_counter_key, - source_cache_key=f"team_id:{scope_team_id}", - increment=cost, - ) - - team_obj: Final[object] = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}") - if team_obj is None: - return - team_budget_limits = getattr(team_obj, "budget_limits", None) or ( - team_obj.get("budget_limits") if isinstance(team_obj, dict) else None + key_obj: Final[object] = await user_api_key_cache.async_get_cache(key=hashed_token) + if key_obj is None: + return key_pending + key_budget_limits = getattr(key_obj, "budget_limits", None) or ( + key_obj.get("budget_limits") if isinstance(key_obj, dict) else None ) - if isinstance(team_budget_limits, str): - team_budget_limits = json.loads(team_budget_limits) - if not isinstance(team_budget_limits, list): - return - for window in team_budget_limits: - duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration - team_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at - team_window_counter = f"spend:team:{scope_team_id}:window:{duration}" + if isinstance(key_budget_limits, str): + key_budget_limits = json.loads(key_budget_limits) + if not isinstance(key_budget_limits, list): + return key_pending + window_pending: Final = await asyncio.gather( + *(_key_window_increment(window) for window in key_budget_limits), return_exceptions=True + ) + return key_pending + tuple(item for item in window_pending if item is not None) + + async def _team_scope(scope_team_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: + team_counter_key: Final = f"spend:team:{scope_team_id}" + team_pending: Final[tuple[_PendingSpendIncrement, ...]] = ( + () + if team_counter_key in reserved_counter_keys + else ( + await _prepare_spend_counter_increment( + counter_key=team_counter_key, + source_cache_key=f"team_id:{scope_team_id}", + increment=cost, + ), + ) + ) + + async def _team_window_increment(window: object) -> _PendingSpendIncrement | None: + duration = ( + window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) + ) + team_window_reset_at = ( + window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) + ) + team_window_counter: Final = f"spend:team:{scope_team_id}:window:{duration}" team_window_start = get_budget_window_start(window) - if team_window_counter not in reserved_counter_keys: - await _init_and_increment_window_spend_counter( + pending_window: Final = ( + await _prepare_window_spend_counter_increment( counter_key=team_window_counter, entity_type="Team", entity_id=scope_team_id, @@ -2822,6 +2847,9 @@ async def increment_spend_counters( window_start=team_window_start, increment=cost, ) + if team_window_counter not in reserved_counter_keys + else None + ) await _enqueue_window_spend_row_update( entity_type=Litellm_EntityType.TEAM, entity_id=scope_team_id, @@ -2831,25 +2859,47 @@ async def increment_spend_counters( increment=cost, request_started_at=request_started_at, ) + return pending_window - async def _team_member_scope(scope_user_id: str, scope_team_id: str) -> None: + team_obj: Final[object] = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}") + if team_obj is None: + return team_pending + team_budget_limits = getattr(team_obj, "budget_limits", None) or ( + team_obj.get("budget_limits") if isinstance(team_obj, dict) else None + ) + if isinstance(team_budget_limits, str): + team_budget_limits = json.loads(team_budget_limits) + if not isinstance(team_budget_limits, list): + return team_pending + window_pending: Final = await asyncio.gather( + *(_team_window_increment(window) for window in team_budget_limits), return_exceptions=True + ) + return team_pending + tuple(item for item in window_pending if item is not None) + + async def _team_member_scope( + scope_user_id: str, scope_team_id: str + ) -> tuple[_PendingSpendIncrement | BaseException, ...]: team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}" if team_member_counter_key in reserved_counter_keys: - return - await _init_and_increment_spend_counter( - counter_key=team_member_counter_key, - source_cache_key=f"team_membership:{scope_user_id}:{scope_team_id}", - increment=cost, + return () + return ( + await _prepare_spend_counter_increment( + counter_key=team_member_counter_key, + source_cache_key=f"team_membership:{scope_user_id}:{scope_team_id}", + increment=cost, + ), ) - async def _user_scope(scope_user_id: str) -> None: + async def _user_scope(scope_user_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: user_counter_key: Final = f"spend:user:{scope_user_id}" if user_counter_key in reserved_counter_keys: - return - await _init_and_increment_spend_counter( - counter_key=user_counter_key, - source_cache_key=scope_user_id, - increment=cost, + return () + return ( + await _prepare_spend_counter_increment( + counter_key=user_counter_key, + source_cache_key=scope_user_id, + increment=cost, + ), ) scope_coros: Final = tuple( @@ -2859,7 +2909,7 @@ async def increment_spend_counters( _team_scope(team_id) if team_id is not None else None, _team_member_scope(user_id, team_id) if user_id is not None and team_id is not None else None, _user_scope(user_id) if user_id is not None else None, - _increment_end_user_and_tag_spend_counters( + _prepare_end_user_and_tag_spend_increments( end_user_id=end_user_id, tags=tags, response_cost=cost, @@ -2867,14 +2917,14 @@ async def increment_spend_counters( ) if end_user_id is not None or tags is not None else None, - _increment_model_access_group_spend_counters( + _prepare_model_access_group_spend_increments( model_access_groups=model_access_groups, response_cost=cost, reserved_counter_keys=reserved_counter_keys, ) if model_access_groups else None, - _increment_org_spend_counter( + _prepare_org_spend_increment( org_id=org_id, response_cost=cost, reserved_counter_keys=reserved_counter_keys, @@ -2889,7 +2939,20 @@ async def increment_spend_counters( # as orphaned tasks that race the caller's reservation-counter invalidation; # all scopes settle, then the first error propagates as before. scope_results: Final = await asyncio.gather(*scope_coros, return_exceptions=True) - scope_errors: Final = [r for r in scope_results if isinstance(r, BaseException)] + scope_errors: Final = tuple( + item + for scope in scope_results + for item in (scope if isinstance(scope, tuple) else (scope,)) + if isinstance(item, BaseException) + ) + pending: Final = tuple( + item + for scope in scope_results + if not isinstance(scope, BaseException) + for item in scope + if not isinstance(item, BaseException) + ) + await _apply_spend_counter_increments(pending=pending) if scope_errors: raise scope_errors[0] @@ -2932,41 +2995,49 @@ async def _reconcile_budget_reservation_for_counter_update( return reserved_counter_keys -async def _increment_end_user_and_tag_spend_counters( +async def _prepare_end_user_and_tag_spend_increments( end_user_id: str | None, tags: list[str] | None, response_cost: float, reserved_counter_keys: set[str], -) -> None: - if end_user_id is not None: - await _init_and_increment_unreserved_spend_counter( - counter_key=f"spend:end_user:{end_user_id}", - source_cache_key=end_user_cache_key(end_user_id), - increment=response_cost, - reserved_counter_keys=reserved_counter_keys, - ) - - if tags is None: - return - - seen_tags: Final[set[str]] = set() - for tag_name in tags: - if not tag_name or not isinstance(tag_name, str) or tag_name in seen_tags: - continue - seen_tags.add(tag_name) - await _init_and_increment_unreserved_spend_counter( - counter_key=f"spend:tag:{tag_name}", - source_cache_key=tag_cache_key(tag_name), - increment=response_cost, - reserved_counter_keys=reserved_counter_keys, - ) +) -> tuple[_PendingSpendIncrement | BaseException, ...]: + unique_tags: Final = ( + tuple(dict.fromkeys(tag for tag in tags if tag and isinstance(tag, str))) if tags is not None else () + ) + results: Final = await asyncio.gather( + *( + coro + for coro in ( + _prepare_unreserved_spend_counter_increment( + counter_key=f"spend:end_user:{end_user_id}", + source_cache_key=end_user_cache_key(end_user_id), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + if end_user_id is not None + else None, + *( + _prepare_unreserved_spend_counter_increment( + counter_key=f"spend:tag:{tag_name}", + source_cache_key=tag_cache_key(tag_name), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + for tag_name in unique_tags + ), + ) + if coro is not None + ), + return_exceptions=True, + ) + return tuple(item for item in results if item is not None) -async def _increment_model_access_group_spend_counters( +async def _prepare_model_access_group_spend_increments( model_access_groups: Sequence[object], response_cost: float, reserved_counter_keys: set[str], -) -> None: +) -> tuple[_PendingSpendIncrement | BaseException, ...]: """Charge the model access groups that authorized this request. Without this the counter auth reads is written only by the reservation path, so @@ -2980,55 +3051,63 @@ async def _increment_model_access_group_spend_counters( unique_groups: Final = tuple( dict.fromkeys(group for group in model_access_groups if group and isinstance(group, str)) ) - for group in unique_groups: - await _init_and_increment_unreserved_spend_counter( - counter_key=model_access_group_spend_counter_key(group), - source_cache_key=model_access_group_cache_key(group), - increment=response_cost, - reserved_counter_keys=reserved_counter_keys, - ) + results: Final = await asyncio.gather( + *( + _prepare_unreserved_spend_counter_increment( + counter_key=model_access_group_spend_counter_key(group), + source_cache_key=model_access_group_cache_key(group), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + for group in unique_groups + ), + return_exceptions=True, + ) + return tuple(item for item in results if item is not None) -async def _increment_org_spend_counter( +async def _prepare_org_spend_increment( org_id: str | None, response_cost: float, reserved_counter_keys: set[str], -) -> None: +) -> tuple[_PendingSpendIncrement, ...]: if org_id is None: - return + return () - await _init_and_increment_unreserved_spend_counter( + pending: Final = await _prepare_unreserved_spend_counter_increment( counter_key=f"spend:org:{org_id}", source_cache_key=[f"org_id:{org_id}:with_budget", f"org_id:{org_id}"], increment=response_cost, reserved_counter_keys=reserved_counter_keys, ) + return (pending,) if pending is not None else () -async def _init_and_increment_unreserved_spend_counter( +async def _prepare_unreserved_spend_counter_increment( counter_key: str, source_cache_key: str | list[str], increment: float, reserved_counter_keys: set[str], -) -> None: +) -> _PendingSpendIncrement | None: if counter_key in reserved_counter_keys: - return + return None - await _init_and_increment_spend_counter( + return await _prepare_spend_counter_increment( counter_key=counter_key, source_cache_key=source_cache_key, increment=increment, ) -async def _init_and_increment_spend_counter( +async def _prepare_spend_counter_increment( counter_key: str, source_cache_key: str | list[str], increment: float, -): +) -> _PendingSpendIncrement: """ Initialize counter from the authoritative DB spend value if not yet - set, then atomically increment in both in-memory and Redis. + set, then return the pending increment for the caller to apply in one + pipelined Redis call. On first access per pod: 1. Check spend_counter_cache (in-memory -> Redis via DualCache) @@ -3040,13 +3119,13 @@ async def _init_and_increment_spend_counter( the counter as absent and seed it. Using increment means the worst case is over-counting (conservative, blocks slightly early) rather than under-counting (would allow overspend). - 4. Increment atomically (both in-memory + Redis) + 4. Increment is returned for the caller to apply via pipeline """ await _ensure_spend_counter_initialized( counter_key=counter_key, source_cache_key=source_cache_key, ) - await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) + return _PendingSpendIncrement(counter_key=counter_key, increment=increment) async def _enqueue_window_spend_row_update( @@ -3098,20 +3177,20 @@ async def _enqueue_window_spend_row_update( ) -async def _init_and_increment_window_spend_counter( +async def _prepare_window_spend_counter_increment( counter_key: str, entity_type: str, entity_id: str, window_duration: str | None, window_start: datetime | None, increment: float, -): +) -> _PendingSpendIncrement | None: if window_start is None: verbose_proxy_logger.warning( "Skipping spend counter increment for invalid budget window %s", counter_key, ) - return + return None initialized: Final = await _ensure_window_spend_counter_initialized( counter_key=counter_key, @@ -3121,8 +3200,8 @@ async def _init_and_increment_window_spend_counter( window_start=window_start, ) if initialized is False: - return - await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) + return None + return _PendingSpendIncrement(counter_key=counter_key, increment=increment) async def _ensure_spend_counter_initialized( @@ -3255,6 +3334,32 @@ async def _invalidate_spend_counter(counter_key: str): ) +async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncrement]) -> None: + if not pending: + return + redis_cache: Final = spend_counter_cache.redis_cache + if redis_cache is None: + for item in pending: + await spend_counter_cache.async_increment_cache( + key=item.counter_key, + value=item.increment, + refresh_ttl=True, + ) + return + ttl: Final = redis_cache.get_ttl() + increment_list: Final = [ # mutable-ok: async_increment_pipeline signature requires list[RedisPipelineIncrementOperation] + RedisPipelineIncrementOperation(key=item.counter_key, increment_value=item.increment, ttl=ttl) + for item in pending + ] + try: + results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list) + except Exception: + await asyncio.gather(*(_invalidate_spend_counter(counter_key=item.counter_key) for item in pending)) + raise + for item, current_value in zip(pending, results or ()): + spend_counter_cache.in_memory_cache.set_cache(key=item.counter_key, value=current_value) + + async def update_cache( token: str | None, user_id: str | None, 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 86e97a334df..c343652efd9 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -4,11 +4,12 @@ Pins covered: - ``get_current_spend`` - ``increment_spend_counters`` - ``_reconcile_budget_reservation_for_counter_update`` -- ``_increment_end_user_and_tag_spend_counters`` -- ``_increment_org_spend_counter`` -- ``_init_and_increment_unreserved_spend_counter`` -- ``_init_and_increment_spend_counter`` -- ``_init_and_increment_window_spend_counter`` +- ``_prepare_end_user_and_tag_spend_increments`` +- ``_prepare_org_spend_increment`` +- ``_prepare_unreserved_spend_counter_increment`` +- ``_prepare_spend_counter_increment`` +- ``_prepare_window_spend_counter_increment`` +- ``_apply_spend_counter_increments`` - ``_ensure_spend_counter_initialized`` - ``_get_source_cache_base_spend`` - ``_ensure_window_spend_counter_initialized`` @@ -48,9 +49,7 @@ def _make_spend_counter_cache( cache.in_memory_cache.delete_cache = MagicMock() if with_redis: cache.redis_cache = MagicMock() - cache.redis_cache.async_get_cache = AsyncMock( - return_value=redis_get_value, side_effect=redis_get_side_effect - ) + cache.redis_cache.async_get_cache = AsyncMock(return_value=redis_get_value, side_effect=redis_get_side_effect) cache.redis_cache.async_increment = AsyncMock( return_value=redis_increment_value, side_effect=redis_increment_side_effect, @@ -58,6 +57,8 @@ def _make_spend_counter_cache( cache.redis_cache.async_delete_cache = AsyncMock() cache.redis_cache.async_set_cache = AsyncMock() cache.redis_cache.async_set_max = AsyncMock() + cache.redis_cache.async_increment_pipeline = AsyncMock(return_value=None) + cache.redis_cache.get_ttl = MagicMock(return_value=None) else: cache.redis_cache = None cache.async_increment_cache = AsyncMock(return_value=redis_increment_value) @@ -70,9 +71,7 @@ def _make_spend_counter_cache( def _make_user_api_key_cache(get_value=None, get_side_effect=None): cache = MagicMock() - cache.async_get_cache = AsyncMock( - return_value=get_value, side_effect=get_side_effect - ) + cache.async_get_cache = AsyncMock(return_value=get_value, side_effect=get_side_effect) cache.async_set_cache_pipeline = AsyncMock() return cache @@ -109,9 +108,7 @@ async def test_get_current_spend_redis_error_falls_back_to_in_memory(monkeypatch ) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=99.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=99.0) assert result == 17.0 @@ -136,9 +133,7 @@ async def test_get_current_spend_floors_stale_low_counter_against_db(monkeypatch # the stale counter is repaired up to the authoritative DB value via a # monotonic set-max so other workers read the corrected total, and a # concurrent increment cannot be clobbered - fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key="spend:key:abc", value=12.0 - ) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key="spend:key:abc", value=12.0) @pytest.mark.asyncio @@ -169,9 +164,7 @@ async def test_get_current_spend_no_floor_without_max_budget(monkeypatch): from_db = AsyncMock(return_value=12.0) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=12.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0) assert result == 2.0 assert from_db.await_count == 0 @@ -210,12 +203,8 @@ async def test_get_current_spend_floor_caches_db_read(monkeypatch): from_db = AsyncMock(return_value=12.0) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) - first = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0 - ) - second = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0 - ) + first = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0) + second = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0) assert first == 12.0 assert second == 12.0 @@ -336,9 +325,7 @@ async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch): assert result == 15.0 assert wfsl.await_count == 1 - fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key=counter_key, value=15.0 - ) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key=counter_key, value=15.0) def _make_window_spend_prisma(row=None, spend_logs_total=0.0): @@ -379,9 +366,7 @@ async def test_get_current_spend_floors_window_against_maintained_row(monkeypatc assert result == 15.0 fake_prisma.db.litellm_spendlogs.group_by.assert_not_awaited() - fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key=counter_key, value=15.0 - ) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key=counter_key, value=15.0) @pytest.mark.asyncio @@ -393,9 +378,7 @@ async def test_get_current_spend_floors_window_against_logs_when_row_stale(monke window_start = datetime(2026, 1, 8, tzinfo=timezone.utc) fake_prisma = _make_window_spend_prisma( - row=SimpleNamespace( - window_start=window_start - timedelta(days=7), spend=999.0 - ), + row=SimpleNamespace(window_start=window_start - timedelta(days=7), spend=999.0), spend_logs_total=15.0, ) fake_cache = _make_spend_counter_cache(redis_get_value=2.0) @@ -423,21 +406,13 @@ async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypat rather than admitted on an unverifiable budget.""" from fastapi import HTTPException - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down") - ) + fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) with pytest.raises(HTTPException) as exc: - await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0 - ) + await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) assert exc.value.status_code == 503 @@ -445,18 +420,12 @@ async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypat async def test_get_current_spend_fail_closed_off_admits_when_unverifiable(monkeypatch): """Default (flag off): an unverifiable read keeps the existing behavior and admits using the cached fallback — no new rejection.""" - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down") - ) + fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "general_settings", {}) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) assert result == 1.0 @@ -466,13 +435,9 @@ async def test_get_current_spend_fail_closed_admits_when_redis_verified(monkeypa authoritative, so an under-budget request is admitted normally.""" fake_cache = _make_spend_counter_cache(redis_get_value=1.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) assert result == 1.0 @@ -481,16 +446,10 @@ async def test_get_current_spend_fail_closed_allows_authoritative_fallback(monke """End-user/tag callers pass fallback_authoritative=True (their spend is loaded fresh from the DB in auth), so fail-closed does not reject them even when the counter path is unreadable.""" - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down") - ) + fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) result = await ps.get_current_spend( counter_key="spend:end_user:e1", @@ -508,9 +467,7 @@ async def test_get_current_spend_strict_floors_when_fallback_also_stale(monkeypa re-checks the authoritative DB and enforces against it.""" fake_cache = _make_spend_counter_cache(redis_get_value=0.00001) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) from_db = AsyncMock(return_value=0.5) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) @@ -532,9 +489,7 @@ async def test_get_current_spend_strict_floors_when_fallback_also_stale(monkeypa @pytest.mark.asyncio async def test_increment_spend_counters_increments_all_buckets(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=5.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=5.0) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) @@ -543,9 +498,7 @@ async def test_increment_spend_counters_increments_all_buckets(monkeypatch): async def _fake_coalesced(**kwargs): return None - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_fake_coalesced) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_fake_coalesced)) await ps.increment_spend_counters( token="hashed-tok", @@ -554,25 +507,36 @@ async def test_increment_spend_counters_increments_all_buckets(monkeypatch): response_cost=5.0, ) + pipeline = fake_cache.redis_cache.async_increment_pipeline + pipeline.assert_awaited_once() + increment_list = pipeline.await_args.kwargs["increment_list"] + assert {op["key"] for op in increment_list} == { + "spend:key:hashed-tok", + "spend:team:t1", + "spend:team_member:u1:t1", + "spend:user:u1", + } + assert all(op["increment_value"] == 5.0 for op in increment_list) observed = { "redis_increment_called": fake_cache.redis_cache.async_increment.called, - "increment_calls": fake_cache.redis_cache.async_increment.call_count, + "pipeline_calls": pipeline.await_count, "user_cache_used": fake_user_cache.async_get_cache.called, } assert normalize(observed) == { - "redis_increment_called": True, - "increment_calls": 4, + "redis_increment_called": False, + "pipeline_calls": 1, "user_cache_used": True, } class _ConcurrencyProbe: - """Stand-in for redis_cache.async_increment that pins concurrency. + """Stand-in for redis_cache.async_get_cache that pins concurrency. - Each call registers itself as in-flight and blocks on ``release`` until the - test lets it proceed. ``all_arrived`` fires once ``expected`` distinct scope - increments are simultaneously suspended here, which can only happen if the - per-scope increments are gathered rather than awaited one after another. + Each warm-check read registers itself as in-flight and blocks on ``release`` + until the test lets it proceed. ``all_arrived`` fires once ``expected`` + distinct scope warm-checks are simultaneously suspended here, which can only + happen if the per-scope prepares are gathered rather than awaited one after + another. """ def __init__(self, expected_concurrency: int): @@ -581,36 +545,45 @@ class _ConcurrencyProbe: self.max_in_flight = 0 self.all_arrived = asyncio.Event() self.release = asyncio.Event() - self.values: dict[str, float] = {} + self.keys: list[str] = [] - async def async_increment(self, *, key, value, refresh_ttl=True): + async def async_get_cache(self, *, key, **kwargs): self.in_flight += 1 self.max_in_flight = max(self.max_in_flight, self.in_flight) + self.keys.append(key) if self.in_flight >= self.expected: self.all_arrived.set() if not self.release.is_set(): await self.release.wait() self.in_flight -= 1 - self.values[key] = self.values.get(key, 0.0) + value - return self.values[key] + return 1.0 @pytest.mark.asyncio async def test_increment_spend_counters_runs_scopes_concurrently(monkeypatch): """The six independent scopes (key, team, team_member, user, end_user+tags, - org) must be incremented concurrently. The probe only fires once all six are - suspended in async_increment at the same time, which is impossible if the + org) must prepare their increments concurrently. The probe only fires once + all eight warm-check reads (one per counter: 6 scopes + 2 tags) are + suspended in async_get_cache at the same time, which is impossible if the awaits are chained sequentially.""" - probe = _ConcurrencyProbe(expected_concurrency=6) - fake_cache = _make_spend_counter_cache(redis_get_value=None) - fake_cache.redis_cache.async_increment = probe.async_increment + probe = _ConcurrencyProbe(expected_concurrency=8) + fake_cache = _make_spend_counter_cache() + fake_cache.redis_cache.async_get_cache = probe.async_get_cache + recorded: dict[str, float] = {} + + async def _record_pipeline(increment_list, **_): + results = [] + for op in increment_list: + recorded[op["key"]] = recorded.get(op["key"], 0.0) + op["increment_value"] + results.append(recorded[op["key"]]) + return results + + fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) task = asyncio.create_task( ps.increment_spend_counters( @@ -630,16 +603,16 @@ async def test_increment_spend_counters_runs_scopes_concurrently(monkeypatch): probe.release.set() await task pytest.fail( - "scope increments did not run concurrently; sequential awaits " - f"detected (peak in-flight was {probe.max_in_flight}, expected 6)" + "scope prepares did not run concurrently; sequential awaits " + f"detected (peak in-flight was {probe.max_in_flight}, expected 8)" ) - assert probe.in_flight == 6 - assert probe.max_in_flight == 6 + assert probe.in_flight == 8 + assert probe.max_in_flight == 8 probe.release.set() await task - assert probe.values == { + assert recorded == { "spend:key:hashed-tok": 5.0, "spend:team:t1": 5.0, "spend:team_member:u1:t1": 5.0, @@ -659,26 +632,25 @@ async def test_increment_spend_counters_skips_reserved_counter_keys(monkeypatch) import litellm.proxy.spend_tracking.budget_reservation as br reserved = {"spend:key:hashed-tok", "spend:org:org1"} - monkeypatch.setattr( - br, "get_reserved_counter_keys", MagicMock(return_value=set(reserved)) - ) + monkeypatch.setattr(br, "get_reserved_counter_keys", MagicMock(return_value=set(reserved))) monkeypatch.setattr(br, "reconcile_budget_reservation", AsyncMock()) recorded: dict[str, float] = {} - async def _record_increment(*, key, value, refresh_ttl=True): - recorded[key] = recorded.get(key, 0.0) + value - return recorded[key] + async def _record_pipeline(increment_list, **_): + results = [] + for op in increment_list: + recorded[op["key"]] = recorded.get(op["key"], 0.0) + op["increment_value"] + results.append(recorded[op["key"]]) + return results fake_cache = _make_spend_counter_cache(redis_get_value=None) - fake_cache.redis_cache.async_increment = _record_increment + fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) reservation = {"finalized": False} await ps.increment_spend_counters( @@ -708,27 +680,46 @@ async def test_increment_spend_counters_failing_scope_propagates_after_siblings_ ): """A failure in one scope must propagate to the caller (so it can invalidate reserved counters) while every other scope still settles rather than being - left as an orphaned background task, and the reservation is not finalized.""" - recorded: dict[str, float] = {} + left as an orphaned background task, and the reservation is not finalized. + The surviving scopes' increments are still applied in the single pipeline: + dropping them would under-count spend, the unsafe direction for budget + enforcement.""" + warmed_keys: list[str] = [] - async def _increment(*, key, value, refresh_ttl=True): + async def _warm_check(*, key, **kwargs): + warmed_keys.append(key) if key == "spend:team:t1": - raise RuntimeError("redis increment failed") - recorded[key] = recorded.get(key, 0.0) + value - return recorded[key] + raise RuntimeError("redis get failed") + return 1.0 - fake_cache = _make_spend_counter_cache(redis_get_value=None) - fake_cache.redis_cache.async_increment = _increment + async def _reseed_fails(*, counter_key, **kwargs): + if counter_key == "spend:team:t1": + raise RuntimeError("reseed failed") + + applied: dict[str, float] = {} + + async def _record_pipeline(increment_list, **_): + results = [] + for op in increment_list: + applied[op["key"]] = op["increment_value"] + results.append(op["increment_value"]) + return results + + fake_cache = _make_spend_counter_cache() + fake_cache.redis_cache.async_get_cache = AsyncMock(side_effect=_warm_check) + fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) + ps.SpendCounterReseed, + "coalesced", + AsyncMock(side_effect=_reseed_fails), ) reservation = {"finalized": False} - with pytest.raises(RuntimeError, match="redis increment failed"): + with pytest.raises(RuntimeError, match="reseed failed"): await ps.increment_spend_counters( token="hashed-tok", team_id="t1", @@ -741,7 +732,19 @@ async def test_increment_spend_counters_failing_scope_propagates_after_siblings_ ) assert reservation["finalized"] is False - assert recorded == { + # every sibling scope settled (its warm-check ran) before the error propagated + assert set(warmed_keys) == { + "spend:key:hashed-tok", + "spend:team:t1", + "spend:team_member:u1:t1", + "spend:user:u1", + "spend:end_user:eu1", + "spend:tag:a", + "spend:org:org1", + } + # the surviving scopes' increments were still applied, in one pipeline call + fake_cache.redis_cache.async_increment_pipeline.assert_awaited_once() + assert applied == { "spend:key:hashed-tok": 5.0, "spend:team_member:u1:t1": 5.0, "spend:user:u1": 5.0, @@ -749,6 +752,7 @@ async def test_increment_spend_counters_failing_scope_propagates_after_siblings_ "spend:tag:a": 5.0, "spend:org:org1": 5.0, } + fake_cache.redis_cache.async_increment.assert_not_awaited() @pytest.mark.asyncio @@ -772,6 +776,108 @@ async def test_increment_spend_counters_zero_cost_is_noop_finalizes_reservation( assert reservation == {"finalized": True} assert fake_cache.redis_cache.async_increment.called is False + fake_cache.redis_cache.async_increment_pipeline.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( + monkeypatch, +): + """Every scope's increment must go out in a single async_increment_pipeline + call, not one INCRBYFLOAT round-trip per scope.""" + counter_cache = ps.DualCache() + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=1.0) # counters warm + + async def _pipeline(increment_list, **_): + return [1.5] * len(increment_list) + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=_pipeline) + fake_redis.async_increment = AsyncMock() + fake_redis.get_ttl = MagicMock(return_value=None) + counter_cache.redis_cache = fake_redis + monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) + monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "prisma_client", None) + + await ps.increment_spend_counters( + token="hashed", + team_id="team-1", + user_id="user-1", + response_cost=0.5, + org_id="org-1", + end_user_id="eu-1", + tags=["tag-a", "tag-b"], + ) + + fake_redis.async_increment_pipeline.assert_awaited_once() + assert fake_redis.async_increment.await_count == 0 + increment_list = fake_redis.async_increment_pipeline.await_args.kwargs["increment_list"] + expected_keys = { + "spend:key:hashed", + "spend:team:team-1", + "spend:team_member:user-1:team-1", + "spend:user:user-1", + "spend:end_user:eu-1", + "spend:tag:tag-a", + "spend:tag:tag-b", + "spend:org:org-1", + } + assert {op["key"] for op in increment_list} == expected_keys + assert all(op["increment_value"] == 0.5 for op in increment_list) + for key in expected_keys: + assert counter_cache.in_memory_cache.get_cache(key=key) == 1.5 + + +@pytest.mark.asyncio +async def test_increment_spend_counters_pipeline_failure_invalidates_all_counters( + monkeypatch, +): + """A failing pipeline must invalidate every pending counter so the next + request reseeds from the DB (which already holds this request's cost) + instead of trusting a value the write may have partially applied.""" + from redis.exceptions import MaxConnectionsError + + counter_cache = ps.DualCache() + pending_keys = ( + "spend:key:hashed", + "spend:team:team-1", + "spend:team_member:user-1:team-1", + "spend:user:user-1", + "spend:end_user:eu-1", + "spend:tag:tag-a", + "spend:tag:tag-b", + "spend:org:org-1", + ) + for key in pending_keys: + counter_cache.in_memory_cache.set_cache(key=key, value=1.0) + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=1.0) # counters warm + fake_redis.async_increment_pipeline = AsyncMock(side_effect=MaxConnectionsError()) + fake_redis.async_increment = AsyncMock() + fake_redis.async_delete_cache = AsyncMock() + fake_redis.get_ttl = MagicMock(return_value=None) + counter_cache.redis_cache = fake_redis + monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) + monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "prisma_client", None) + + with pytest.raises(MaxConnectionsError): + await ps.increment_spend_counters( + token="hashed", + team_id="team-1", + user_id="user-1", + response_cost=0.5, + org_id="org-1", + end_user_id="eu-1", + tags=["tag-a", "tag-b"], + ) + + assert fake_redis.async_increment.await_count == 0 + deleted_keys = {call.kwargs["key"] for call in fake_redis.async_delete_cache.await_args_list} + assert deleted_keys == set(pending_keys) + for key in pending_keys: + assert counter_cache.in_memory_cache.get_cache(key=key) is None # --------------------------------------------------------------------------- @@ -781,9 +887,7 @@ async def test_increment_spend_counters_zero_cost_is_noop_finalizes_reservation( @pytest.mark.asyncio async def test_reconcile_budget_reservation_for_counter_update_returns_empty_set_when_none(): - result = await ps._reconcile_budget_reservation_for_counter_update( - budget_reservation=None, response_cost=1.0 - ) + result = await ps._reconcile_budget_reservation_for_counter_update(budget_reservation=None, response_cost=1.0) assert result == set() @@ -818,179 +922,151 @@ async def test_reconcile_budget_reservation_for_counter_update_failure_invalidat # --------------------------------------------------------------------------- -# _increment_end_user_and_tag_spend_counters +# _prepare_end_user_and_tag_spend_increments # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_increment_end_user_and_tag_spend_counters_increments_each_unique_tag( +async def test_prepare_end_user_and_tag_spend_increments_returns_each_unique_tag( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=3.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=1.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) - await ps._increment_end_user_and_tag_spend_counters( + pending = await ps._prepare_end_user_and_tag_spend_increments( end_user_id="eu1", tags=["a", "b", "a", "", None], response_cost=3.0, reserved_counter_keys=set(), ) - observed = { - "increment_calls": fake_cache.redis_cache.async_increment.call_count, - "in_memory_set_calls": fake_cache.in_memory_cache.set_cache.call_count, - "called": fake_cache.redis_cache.async_increment.called, - } - assert normalize(observed) == { - "increment_calls": 3, - "in_memory_set_calls": 3, - "called": True, + assert {item.counter_key for item in pending} == { + "spend:end_user:eu1", + "spend:tag:a", + "spend:tag:b", } + assert all(item.increment == 3.0 for item in pending) @pytest.mark.asyncio -async def test_increment_end_user_and_tag_spend_counters_no_end_user_no_tags_invalid_input_noop( +async def test_prepare_end_user_and_tag_spend_increments_no_end_user_no_tags_invalid_input_noop( monkeypatch, ): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._increment_end_user_and_tag_spend_counters( + pending = await ps._prepare_end_user_and_tag_spend_increments( end_user_id=None, tags=None, response_cost=1.0, reserved_counter_keys=set(), ) + assert pending == () assert fake_cache.redis_cache.async_increment.called is False # --------------------------------------------------------------------------- -# _increment_org_spend_counter +# _prepare_org_spend_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_increment_org_spend_counter_increments_when_org_present(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=10.0 - ) +async def test_prepare_org_spend_increment_returns_pending_when_org_present(monkeypatch): + fake_cache = _make_spend_counter_cache(redis_get_value=1.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) - await ps._increment_org_spend_counter( + pending = await ps._prepare_org_spend_increment( org_id="org-1", response_cost=10.0, reserved_counter_keys=set(), ) - observed = { - "increment_called": fake_cache.redis_cache.async_increment.called, - "increment_calls": fake_cache.redis_cache.async_increment.call_count, - "counter_key_arg": fake_cache.redis_cache.async_increment.call_args.kwargs[ - "key" - ], - } - assert normalize(observed) == { - "increment_called": True, - "increment_calls": 1, - "counter_key_arg": "spend:org:org-1", - } + assert len(pending) == 1 + assert pending[0].counter_key == "spend:org:org-1" + assert pending[0].increment == 10.0 @pytest.mark.asyncio -async def test_increment_org_spend_counter_no_org_is_noop_invalid_id(monkeypatch): +async def test_prepare_org_spend_increment_no_org_is_noop_invalid_id(monkeypatch): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._increment_org_spend_counter( + pending = await ps._prepare_org_spend_increment( org_id=None, response_cost=1.0, reserved_counter_keys=set(), ) + assert pending == () assert fake_cache.redis_cache.async_increment.called is False # --------------------------------------------------------------------------- -# _init_and_increment_unreserved_spend_counter +# _prepare_unreserved_spend_counter_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_init_and_increment_unreserved_spend_counter_skips_reserved_keys( +async def test_prepare_unreserved_spend_counter_increment_skips_reserved_keys( monkeypatch, ): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._init_and_increment_unreserved_spend_counter( + pending = await ps._prepare_unreserved_spend_counter_increment( counter_key="spend:tag:x", source_cache_key="tag:x", increment=1.0, reserved_counter_keys={"spend:tag:x"}, ) + assert pending is None assert fake_cache.redis_cache.async_increment.called is False @pytest.mark.asyncio -async def test_init_and_increment_unreserved_spend_counter_proceeds_when_not_reserved( +async def test_prepare_unreserved_spend_counter_increment_proceeds_when_not_reserved( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=2.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None) fake_user_cache = _make_user_api_key_cache() + reseed = AsyncMock(return_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", reseed) - await ps._init_and_increment_unreserved_spend_counter( + pending = await ps._prepare_unreserved_spend_counter_increment( counter_key="spend:tag:y", source_cache_key="tag:y", increment=2.0, reserved_counter_keys=set(), ) - observed = { - "increment_called": fake_cache.redis_cache.async_increment.called, - "redis_get_called": fake_cache.redis_cache.async_get_cache.called, - "reseed_consulted": True, - } - assert observed == { - "increment_called": True, - "redis_get_called": True, - "reseed_consulted": True, - } + assert pending is not None + assert pending.counter_key == "spend:tag:y" + assert pending.increment == 2.0 + assert fake_cache.redis_cache.async_get_cache.called is True + assert reseed.called is True # --------------------------------------------------------------------------- -# _init_and_increment_spend_counter +# _prepare_spend_counter_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_init_and_increment_spend_counter_warm_cache_skips_reseed(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_value=11.0, redis_increment_value=14.0 - ) +async def test_prepare_spend_counter_increment_warm_cache_skips_reseed(monkeypatch): + fake_cache = _make_spend_counter_cache(redis_get_value=11.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) @@ -998,12 +1074,14 @@ async def test_init_and_increment_spend_counter_warm_cache_skips_reseed(monkeypa reseed = AsyncMock(return_value=None) monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", reseed) - await ps._init_and_increment_spend_counter( + pending = await ps._prepare_spend_counter_increment( counter_key="spend:key:k", source_cache_key="k", increment=3.0, ) + assert pending.counter_key == "spend:key:k" + assert pending.increment == 3.0 observed = { "reseed_called": reseed.called, "increment_called": fake_cache.redis_cache.async_increment.called, @@ -1011,23 +1089,21 @@ async def test_init_and_increment_spend_counter_warm_cache_skips_reseed(monkeypa } assert normalize(observed) == { "reseed_called": False, - "increment_called": True, + "increment_called": False, "in_memory_seeded_from_redis": True, } # --------------------------------------------------------------------------- -# _init_and_increment_window_spend_counter +# _prepare_window_spend_counter_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_init_and_increment_window_spend_counter_increments_when_initialized( +async def test_prepare_window_spend_counter_increment_returns_pending_when_initialized( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=0.0, redis_increment_value=5.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=0.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr( @@ -1036,7 +1112,7 @@ async def test_init_and_increment_window_spend_counter_increments_when_initializ AsyncMock(return_value=0.0), ) - await ps._init_and_increment_window_spend_counter( + pending = await ps._prepare_window_spend_counter_increment( counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", @@ -1045,26 +1121,19 @@ async def test_init_and_increment_window_spend_counter_increments_when_initializ increment=5.0, ) - observed = { - "redis_increment_called": fake_cache.redis_cache.async_increment.called, - "increment_calls": fake_cache.redis_cache.async_increment.call_count, - "in_memory_set_calls": fake_cache.in_memory_cache.set_cache.call_count, - } - assert normalize(observed) == { - "redis_increment_called": True, - "increment_calls": 1, - "in_memory_set_calls": 2, - } + assert pending is not None + assert pending.counter_key == "spend:key:k:window:1d" + assert pending.increment == 5.0 @pytest.mark.asyncio -async def test_init_and_increment_window_spend_counter_missing_window_start_invalid_skips( +async def test_prepare_window_spend_counter_increment_missing_window_start_invalid_skips( monkeypatch, ): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._init_and_increment_window_spend_counter( + pending = await ps._prepare_window_spend_counter_increment( counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", @@ -1073,6 +1142,7 @@ async def test_init_and_increment_window_spend_counter_missing_window_start_inva increment=5.0, ) + assert pending is None assert fake_cache.redis_cache.async_increment.called is False @@ -1114,16 +1184,12 @@ async def test_ensure_spend_counter_initialized_warm_skips_reseed_and_source( async def test_ensure_spend_counter_initialized_cold_seeds_from_source_cache( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=7.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=7.0) fake_user_cache = _make_user_api_key_cache(get_value={"spend": 7.0}) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) await ps._ensure_spend_counter_initialized( counter_key="spend:user:u", @@ -1163,9 +1229,7 @@ async def test_get_source_cache_base_spend_reads_first_hit_from_list(monkeypatch fake_user_cache.async_get_cache = AsyncMock(side_effect=_get) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) - result = await ps._get_source_cache_base_spend( - source_cache_key=["miss", "hit-obj", "miss2"] - ) + result = await ps._get_source_cache_base_spend(source_cache_key=["miss", "hit-obj", "miss2"]) observed = { "result": result, @@ -1294,9 +1358,7 @@ async def test_increment_spend_counter_cache_redis_path_returns_new_value(monkey fake_cache = _make_spend_counter_cache(redis_increment_value=44.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - result = await ps._increment_spend_counter_cache( - counter_key="spend:key:k", increment=4.0 - ) + result = await ps._increment_spend_counter_cache(counter_key="spend:key:k", increment=4.0) observed = { "result": result, @@ -1314,15 +1376,11 @@ async def test_increment_spend_counter_cache_redis_path_returns_new_value(monkey async def test_increment_spend_counter_cache_redis_error_raises_and_invalidates( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_increment_side_effect=RuntimeError("incr fail") - ) + fake_cache = _make_spend_counter_cache(redis_increment_side_effect=RuntimeError("incr fail")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) with pytest.raises(RuntimeError): - await ps._increment_spend_counter_cache( - counter_key="spend:key:k", increment=1.0 - ) + await ps._increment_spend_counter_cache(counter_key="spend:key:k", increment=1.0) assert fake_cache.in_memory_cache.delete_cache.called is True assert fake_cache.redis_cache.async_delete_cache.called is True @@ -1343,9 +1401,7 @@ async def test_invalidate_spend_counter_deletes_in_memory_and_redis(monkeypatch) observed = { "in_memory_delete_called": fake_cache.in_memory_cache.delete_cache.called, "redis_delete_called": fake_cache.redis_cache.async_delete_cache.called, - "delete_args_key": fake_cache.redis_cache.async_delete_cache.call_args.kwargs[ - "key" - ], + "delete_args_key": fake_cache.redis_cache.async_delete_cache.call_args.kwargs["key"], } assert normalize(observed) == { "in_memory_delete_called": True, @@ -1357,9 +1413,7 @@ async def test_invalidate_spend_counter_deletes_in_memory_and_redis(monkeypatch) @pytest.mark.asyncio async def test_invalidate_spend_counter_swallows_redis_failure_no_raise(monkeypatch): fake_cache = _make_spend_counter_cache() - fake_cache.redis_cache.async_delete_cache = AsyncMock( - side_effect=RuntimeError("redis down") - ) + fake_cache.redis_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) await ps._invalidate_spend_counter(counter_key="spend:key:k") diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py index c123eeeed36..6165af4920d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py @@ -41,6 +41,15 @@ class _FlakyRedisCache: self._store[key] = float(value) return True + async def async_increment_pipeline(self, increment_list, **kwargs): + results = [] + for op in increment_list: + results.append(await self.async_increment(op["key"], op["increment_value"])) + return results + + def get_ttl(self, **kwargs): + return None + @pytest.mark.asyncio async def test_direct_increment_runs_when_reservation_reconcile_hits_redis_failure( diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 6dab054d8ea..40ebc03781c 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2220,6 +2220,15 @@ class _ExpiringRedisCache: async def async_delete_cache(self, key: str, *args: object, **kwargs: object) -> None: self.store.pop(key, None) + async def async_increment_pipeline(self, increment_list, **kwargs): + results = [] + for op in increment_list: + results.append(await self.async_increment(op["key"], op["increment_value"])) + return results + + def get_ttl(self, **kwargs) -> None: + return None + @pytest.mark.asyncio async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 54579e6cb7c..b0f38978727 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7749,7 +7749,7 @@ async def test_increment_spend_counters_team_and_member(): @pytest.mark.asyncio -async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss(): +async def test_prepare_spend_counter_increment_reseeds_from_db_on_counter_miss(): """When the Redis counter is missing, the reseed path reads the authoritative spend from the DB (not a stale cache), so the next increment continues from the correct base value.""" @@ -7762,8 +7762,17 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( recorded_increments.append({"key": key, "value": value, "ttl": ttl}) return value + async def record_pipeline(increment_list, **kwargs): + results = [] + for op in increment_list: + await record_increment(key=op["key"], value=op["increment_value"], ttl=op["ttl"]) + results.append(op["increment_value"]) + return results + fake_redis = AsyncMock() fake_redis.async_increment = AsyncMock(side_effect=record_increment) + fake_redis.async_increment_pipeline = AsyncMock(side_effect=record_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) fake_redis.async_get_cache = AsyncMock(return_value=None) # counter missing fake_redis.async_set_cache = AsyncMock(return_value=True) # SET NX wins counter_cache.redis_cache = fake_redis @@ -7782,7 +7791,10 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( stale_cache.in_memory_cache.set_cache(key="team_id:team-9", value=stale_team) import litellm.proxy.proxy_server as ps - from litellm.proxy.proxy_server import _init_and_increment_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_spend_counter_increment, + ) orig_user, orig_counter, orig_prisma = ( ps.user_api_key_cache, @@ -7793,11 +7805,12 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_spend_counter( + pending = await _prepare_spend_counter_increment( counter_key="spend:team:team-9", source_cache_key="team_id:team-9", increment=1.5, ) + await _apply_spend_counter_increments(pending=(pending,)) fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-9"}) # Seed uses SET NX with db_spend (42) — cross-pod safe, no INCR of 42. @@ -7976,7 +7989,10 @@ async def test_reseed_spend_from_db_skips_window_variant_keys(): @pytest.mark.asyncio async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_window_spend_counter_increment, + ) counter_cache = DualCache() window_start = datetime.now(timezone.utc) - timedelta(hours=1) @@ -7992,7 +8008,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key="spend:key:key-window:window:1h", entity_type="Key", entity_id="key-window", @@ -8000,6 +8016,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): window_start=window_start, increment=0.5, ) + await _apply_spend_counter_increments(pending=(pending,) if pending is not None else ()) fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once_with( by=["api_key"], @@ -8015,7 +8032,10 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): @pytest.mark.asyncio async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_spend_counter_increment, + ) counter_cache = DualCache() counter_key = "spend:team:team-stale-local" @@ -8037,6 +8057,15 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): fake_redis.async_get_cache = AsyncMock(return_value=None) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) + + async def redis_increment_pipeline(increment_list, **_): + results = [] + for op in increment_list: + results.append(await redis_increment(key=op["key"], value=op["increment_value"])) + return results + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=redis_increment_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) counter_cache.redis_cache = fake_redis db_row = MagicMock() @@ -8055,11 +8084,12 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): ps.prisma_client = fake_prisma ps.user_api_key_cache = DualCache() try: - await _init_and_increment_spend_counter( + pending = await _prepare_spend_counter_increment( counter_key=counter_key, source_cache_key="team_id:team-stale-local", increment=1.5, ) + await _apply_spend_counter_increments(pending=(pending,)) fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-stale-local"}) # Seed via SET NX (42) + delta via INCRBYFLOAT (1.5) = 43.5. @@ -8074,7 +8104,10 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): @pytest.mark.asyncio async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_window_spend_counter_increment, + ) counter_cache = DualCache() counter_key = "spend:key:key-window-stale-local:window:1h" @@ -8097,6 +8130,15 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): fake_redis.async_get_cache = AsyncMock(return_value=None) fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + + async def redis_increment_pipeline(increment_list, **_): + results = [] + for op in increment_list: + results.append(await redis_increment(key=op["key"], value=op["increment_value"])) + return results + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=redis_increment_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() @@ -8111,7 +8153,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key=counter_key, entity_type="Key", entity_id="key-window-stale-local", @@ -8119,6 +8161,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): window_start=window_start, increment=0.5, ) + await _apply_spend_counter_increments(pending=(pending,) if pending is not None else ()) fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once_with( by=["api_key"], @@ -8138,7 +8181,10 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): @pytest.mark.asyncio async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_window_spend_counter_increment, + ) counter_cache = DualCache() counter_key = "spend:key:key-window-concurrent-seed:window:1h" @@ -8161,6 +8207,15 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() fake_redis.async_get_cache = AsyncMock(side_effect=redis_get_cache) fake_redis.async_set_cache = AsyncMock(return_value=False) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + + async def redis_increment_pipeline(increment_list, **_): + results = [] + for op in increment_list: + results.append(await redis_increment(key=op["key"], value=op["increment_value"])) + return results + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=redis_increment_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() @@ -8175,7 +8230,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key=counter_key, entity_type="Key", entity_id="key-window-concurrent-seed", @@ -8183,6 +8238,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() window_start=window_start, increment=0.5, ) + await _apply_spend_counter_increments(pending=(pending,) if pending is not None else ()) fake_redis.async_set_cache.assert_awaited_once_with( key=counter_key, @@ -8199,7 +8255,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() @pytest.mark.asyncio async def test_window_spend_counter_skips_invalid_window_start(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import _prepare_window_spend_counter_increment counter_cache = DualCache() @@ -8208,7 +8264,7 @@ async def test_window_spend_counter_skips_invalid_window_start(): orig_counter = ps.spend_counter_cache ps.spend_counter_cache = counter_cache try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key="spend:key:key-invalid-window:window:not-a-duration", entity_type="Key", entity_id="key-invalid-window", @@ -8216,6 +8272,7 @@ async def test_window_spend_counter_skips_invalid_window_start(): window_start=None, increment=0.5, ) + assert pending is None assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-invalid-window:window:not-a-duration") is None finally: @@ -8279,6 +8336,9 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments(): async def assert_reservation_not_finalized_yet(**kwargs): assert budget_reservation["finalized"] is False incremented_counters.append(kwargs["counter_key"]) + return ps._PendingSpendIncrement( + counter_key=kwargs["counter_key"], increment=kwargs["increment"] + ) import litellm.proxy.proxy_server as ps @@ -8287,7 +8347,7 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments(): ps.user_api_key_cache = DualCache() try: with patch( - "litellm.proxy.proxy_server._init_and_increment_spend_counter", + "litellm.proxy.proxy_server._prepare_spend_counter_increment", new=AsyncMock(side_effect=assert_reservation_not_finalized_yet), ): await increment_spend_counters( @@ -8620,7 +8680,7 @@ async def test_get_current_spend_uses_db_zero_over_stale_fallback(): async def test_concurrent_read_and_write_paths_share_one_db_query(): """ The read path (`get_current_spend`) and the write path - (`_init_and_increment_spend_counter`) both reseed cold counters from + (`_prepare_spend_counter_increment`) both reseed cold counters from the DB. They must share the per-counter lock so a concurrent pre-call enforcement read and post-call increment for the same counter collapse to one DB query, not two. @@ -8629,7 +8689,7 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): from litellm.caching.dual_cache import DualCache from litellm.proxy.proxy_server import ( - _init_and_increment_spend_counter, + _prepare_spend_counter_increment, get_current_spend, ) @@ -8683,7 +8743,7 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): try: results = await _asyncio.gather( get_current_spend(counter_key=counter_key, fallback_spend=0.0), - _init_and_increment_spend_counter( + _prepare_spend_counter_increment( counter_key=counter_key, source_cache_key="ignored", increment=1.5,