scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

perf(proxy): pipeline spend counter increments into one Redis call per request (#40371)

Browse source 
* 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 <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-09 10:35:50 -07:00 • committed by GitHub
parent ea0851de79
commit 996ee5635a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 609 additions and 372 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

331
litellm/proxy/proxy_server.py
View file

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

538
tests/test_litellm/proxy/proxy_server/test_spend_counters.py
View file

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

9
tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py
View file

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

9
tests/test_litellm/proxy/test_budget_reservation.py
View file

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

94
tests/test_litellm/proxy/test_proxy_server.py
View file

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