scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

Merge 874819590c into fca3d4f513

Browse source
This commit is contained in:
mubashir1osmani 2026-11-05 21:51:48 +00:00 • committed by GitHub
commit e0548169ab
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
29 changed files with 2755 additions and 107 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

80
litellm/caching/redis_cache.py
View file

@ -20,7 +20,7 @@ import threading
import time
import weakref
from collections import Counter
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Iterator, Mapping, Sequence
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import timedelta
@ -468,6 +468,14 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep
_swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1)
async def _first_keys(scan: AsyncIterator[object], count: int | None) -> AsyncIterator[str]:
taken: Final = itertools.count(1)
async for key in scan:
yield key.decode() if isinstance(key, bytes) else str(key)
if count is not None and next(taken) >= count:
return
_LIVE_BREAKERS: Final[weakref.WeakSet[RedisCircuitBreaker]] = weakref.WeakSet()
@ -991,10 +999,9 @@ class RedisCache(BaseCache):
return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys))
@_redis_circuit_breaker_guard
async def async_scan_iter(self, pattern: str, count: int = 100) -> list:
async def async_scan_iter(self, pattern: str, count: int = 100, *, unbounded: bool = False) -> list[str]:
start_time: Final = time.time()
try:
keys: Final = []
_redis_client: Final = self.init_async_client()
if not hasattr(_redis_client, "scan_iter"):
verbose_logger.debug(
@ -1002,11 +1009,13 @@ class RedisCache(BaseCache):
)
return []
pattern = self.check_and_fix_namespace(key=pattern)
async for key in _redis_client.scan_iter(match=pattern + "*", count=count):
keys.append(key)
if len(keys) >= count:
break
namespaced_pattern: Final = self.check_and_fix_namespace(key=pattern)
keys: Final = [
key
async for key in _first_keys(
_redis_client.scan_iter(match=namespaced_pattern + "*", count=count), None if unbounded else count
)
]
## LOGGING ##
end_time = time.time()
@ -1309,10 +1318,15 @@ class RedisCache(BaseCache):
_record_swallowed_redis_failure(self._circuit_breaker, e)
@_redis_circuit_breaker_guard
async def async_set_cache_pipeline_with_ttls(self, cache_list: Sequence[tuple[str, object, float | None]]) -> None:
"""One round trip for writes whose TTLs differ; a ``None`` TTL falls back to the default TTL."""
async def async_set_cache_pipeline_with_ttls(
self, cache_list: Sequence[tuple[str, object, float | None]]
) -> Exception | None:
"""One round trip for writes whose TTLs differ; a ``None`` TTL falls back to the default TTL.
Returns the Redis error when the writes were lost, so a caller that must not report
success on a lost write can tell. An open circuit breaker still raises."""
if len(cache_list) == 0:
return
return None
commands: Final = tuple(
(self.check_and_fix_namespace(key=cache_key), json.dumps(cache_value), self.get_ttl(ttl=ttl))
for cache_key, cache_value, ttl in cache_list
@ -1323,16 +1337,6 @@ class RedisCache(BaseCache):
for cache_key, json_cache_value, ttl in commands:
pipe.set(name=cache_key, value=json_cache_value, ex=None if ttl is None else timedelta(seconds=ttl))
await pipe.execute()
asyncio.create_task(
self.service_logger_obj.async_service_success_hook(
service=ServiceTypes.REDIS,
duration=time.time() - start_time,
call_type="async_set_cache_pipeline_with_ttls",
caller=_get_call_stack_info(),
start_time=start_time,
end_time=time.time(),
)
)
except Exception as e:
asyncio.create_task(
self.service_logger_obj.async_service_failure_hook(
@ -1352,6 +1356,19 @@ class RedisCache(BaseCache):
e,
)
_record_swallowed_redis_failure(self._circuit_breaker, e)
return e
else:
asyncio.create_task(
self.service_logger_obj.async_service_success_hook(
service=ServiceTypes.REDIS,
duration=time.time() - start_time,
call_type="async_set_cache_pipeline_with_ttls",
caller=_get_call_stack_info(),
start_time=start_time,
end_time=time.time(),
)
)
return None
async def _set_cache_sadd_helper(
self,
@ -1528,12 +1545,18 @@ class RedisCache(BaseCache):
)
raise e
def reset_fence_key(self, key: str) -> str:
"""Hash-tagged with the namespaced ``key`` so a cluster keeps the fence and the key on one slot."""
return f"{{{self.check_and_fix_namespace(key=key)}}}:reset_fence"
@_redis_circuit_breaker_guard
async def async_set_max(
self,
key: str,
value: float,
ttl: int | None = None,
*,
fence_key: str | None = None,
) -> float | None:
"""Atomically set ``key`` to ``value`` only when ``value`` is greater
than the stored value (or the key is unset), refreshing the TTL.
@ -1543,11 +1566,17 @@ class RedisCache(BaseCache):
concurrent increment that has already pushed the counter higher. The
GET/compare/SET runs in a single Lua call, so it is also atomic across
racing callers and pods. Returns the resulting value.
While ``fence_key`` exists the write is held and the stored value returned
as is: a spend reset sets one (``reset_fence_key``) so a repair still
carrying the pre-reset total cannot raise the counter the reset zeroed.
"""
_redis_client: Final = self.init_async_client()
_used_ttl: Final = self.get_ttl(ttl=ttl)
key = self.check_and_fix_namespace(key=key)
keys: Final = (key,) if fence_key is None else (key, self.check_and_fix_namespace(key=fence_key))
lua: Final = (
"if KEYS[2] ~= nil and redis.call('EXISTS', KEYS[2]) == 1 then return redis.call('GET', KEYS[1]) end "
"local cur = redis.call('GET', KEYS[1]) "
"if cur == false or tonumber(cur) < tonumber(ARGV[1]) then "
"redis.call('SET', KEYS[1], ARGV[1]) "
@ -1557,7 +1586,7 @@ class RedisCache(BaseCache):
)
result = cast(
"str | bytes | int | float | None",
await _redis_client.eval(lua, 1, key, str(value), str(int(_used_ttl or 0))),
await _redis_client.eval(lua, len(keys), *keys, str(value), str(int(_used_ttl or 0))),
)
if result is None:
return None
@ -1934,11 +1963,10 @@ class RedisCache(BaseCache):
raise e
@_redis_circuit_breaker_guard
async def delete_cache_keys(self, keys):
async def delete_cache_keys(self, keys: Iterable[str]) -> None:
_redis_client: Final = self._async_commands()
keys = [self.check_and_fix_namespace(key=key) for key in keys]
# keys is a list, unpack it so it gets passed as individual elements to delete
await _redis_client.delete(*keys)
namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in keys]
await _redis_client.delete(*namespaced_keys)
def client_list(self) -> list:
client_list: Final[list] = self.redis_client.client_list()

1
litellm/constants.py
View file

@ -1987,6 +1987,7 @@ _background_interaction_cost_polling_env: Final = os.getenv(
BACKGROUND_INTERACTION_COST_POLLING_ENABLED: Final = _background_interaction_cost_polling_env == "true"
BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS: Final = 300.0
PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = get_env_int("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)
SPEND_COUNTER_REDIS_TTL_SECONDS: Final = 3600
PROXY_BATCH_WRITE_AT: Final = get_env_int("PROXY_BATCH_WRITE_AT", 10) # in seconds, increased from 10
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS: Final = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30)

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

@ -3328,7 +3328,11 @@ async def invalidate_team_member_spend_state(
await await_inflight_team_membership_load(user_id=user_id, team_id=team_id)
if new_spend is not None:
from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
from litellm.proxy.proxy_server import (
SPEND_DB_FLOOR_CACHE_TTL_SECONDS,
spend_counter_cache,
spend_counter_redis_ttl,
)
spend_counter_key: Final = f"spend:team_member:{user_id}:{team_id}"
spend_db_floor_key: Final = f"spend_db_floor:{spend_counter_key}"
@ -3336,7 +3340,9 @@ async def invalidate_team_member_spend_state(
spend_counter_cache.in_memory_cache.set_cache(key=spend_counter_key, value=new_spend, ttl=60)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(key=spend_counter_key, value=new_spend, ttl=60)
await spend_counter_cache.redis_cache.async_set_cache(
key=spend_counter_key, value=new_spend, ttl=spend_counter_redis_ttl()
)
except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up
verbose_proxy_logger.warning(
"Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next "

100
litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py
View file

@ -1,15 +1,13 @@
import asyncio
import json
from collections.abc import Mapping, Sequence
from dataclasses import asdict, dataclass
from types import MappingProxyType
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
from pydantic import TypeAdapter
from pydantic import TypeAdapter, ValidationError
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.proxy.common_utils.config_sync_pubsub import ( # noqa: F401 # legacy module exports
ConfigSyncPubSub,
_ConfigSyncPubSub, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
@ -29,12 +27,14 @@ AUTH_CACHE_INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation"
_POLL_TIMEOUT_SECONDS: Final = 1.0
_MAX_PENDING_PUBLISHES: Final = 1024
_MAX_IN_FLIGHT_PUBLISHES: Final = 16
_KEYS_PER_MESSAGE: Final = 500
_MESSAGE_CACHE_KEYS: Final = TypeAdapter(tuple[str, ...])
_MESSAGE_FIELDS: Final = TypeAdapter(Mapping[str, object])
_BULK_EVICTION_WAIT_SECONDS: Final = 2.0
_pending_publishes: Final[set[asyncio.Task[None]]] = set() # mutable-ok: strong refs keep background publishes alive
_in_flight_publishes: Final = asyncio.Semaphore(_MAX_IN_FLIGHT_PUBLISHES)
_BACKOFF_INITIAL_SECONDS: Final = 5.0
_BACKOFF_MAX_SECONDS: Final = 60.0
_DECODED_JSON: Final = TypeAdapter(object)
def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str:
@ -45,16 +45,21 @@ def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str:
@dataclass(frozen=True, slots=True)
class _CacheInvalidationMessage:
cache_key: str
cache_keys: tuple[str, ...]
new_value: float | None = None
ttl: float | None = None
def _cache_invalidation_message_json(cache_key: str, new_value: float | None = None, ttl: float | None = None) -> str:
message: Final[Mapping[str, object]] = MappingProxyType(
asdict(_CacheInvalidationMessage(cache_key=cache_key, new_value=new_value, ttl=ttl))
)
return json.dumps({field: value for field, value in message.items() if value is not None})
def _cache_invalidation_message_json(
cache_keys: Sequence[str], new_value: float | None = None, ttl: float | None = None
) -> str:
fields: Final = {
"cache_key": cache_keys[0] if len(cache_keys) == 1 else None,
"cache_keys": list(cache_keys) if len(cache_keys) > 1 else None,
"new_value": new_value,
"ttl": ttl,
}
return json.dumps({field: value for field, value in fields.items() if value is not None})
def _finite_number_or_none(value: object) -> float | None:
@ -63,22 +68,29 @@ def _finite_number_or_none(value: object) -> float | None:
return float(value)
def _message_cache_keys(single: object, many: object) -> tuple[str, ...]:
if isinstance(single, str):
return (single,)
try:
return _MESSAGE_CACHE_KEYS.validate_python(many)
except ValidationError:
return ()
def _message_from_data(data: object) -> _CacheInvalidationMessage | None:
if isinstance(data, bytes):
data = data.decode("utf-8", errors="replace") # rebind-ok: normalizing the wire payload to str
if not isinstance(data, str):
return None
try:
parsed: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(data)))
except json.JSONDecodeError:
parsed: Final = _MESSAGE_FIELDS.validate_json(data)
except ValidationError:
return None
if parsed is None:
return None
cache_key: Final = parsed.get("cache_key")
if not isinstance(cache_key, str):
cache_keys: Final = _message_cache_keys(parsed.get("cache_key"), parsed.get("cache_keys"))
if len(cache_keys) == 0:
return None
return _CacheInvalidationMessage(
cache_key=cache_key,
cache_keys=cache_keys,
new_value=_finite_number_or_none(parsed.get("new_value")),
ttl=_finite_number_or_none(parsed.get("ttl")),
)
@ -99,7 +111,7 @@ async def _publish_batch_to_redis(redis_cache: "RedisCache", cache_keys: tuple[s
channel: Final = auth_cache_invalidation_channel(redis_cache)
pipeline: Final = client.pipeline(transaction=False)
for cache_key in cache_keys:
_ = pipeline.publish(channel, _cache_invalidation_message_json(cache_key))
_ = pipeline.publish(channel, _cache_invalidation_message_json((cache_key,)))
async with _in_flight_publishes:
_ = await pipeline.execute()
except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors
@ -119,6 +131,19 @@ def _has_room_for_publish(key_count: int) -> bool:
return False
def _queue_publish(
redis_cache: "RedisCache", cache_keys: tuple[str, ...], new_value: float | None, ttl: float | None
) -> None:
if not _has_room_for_publish(len(cache_keys)):
return
task: Final = asyncio.create_task(
_publish_to_redis(
redis_cache, cache_keys[0], _cache_invalidation_message_json(cache_keys, new_value=new_value, ttl=ttl)
)
)
_pending_publishes.add(task)
async def publish_auth_cache_invalidation(
cache_key: str, new_value: float | None = None, ttl: float | None = None
) -> None:
@ -141,15 +166,17 @@ async def publish_auth_cache_invalidation(
Redis connection at once; the rest wait in the task set, so a wedge cannot
drain the shared connection pool.
"""
await publish_auth_cache_invalidations(cache_keys=(cache_key,), new_value=new_value, ttl=ttl)
async def publish_auth_cache_invalidations(
cache_keys: Sequence[str], new_value: float | None = None, ttl: float | None = None
) -> None:
redis_cache: Final = coordination_redis_cache()
if redis_cache is None or not _has_room_for_publish(1):
if redis_cache is None:
return
task: Final = asyncio.create_task(
_publish_to_redis(
redis_cache, cache_key, _cache_invalidation_message_json(cache_key, new_value=new_value, ttl=ttl)
)
)
_pending_publishes.add(task)
for start in range(0, len(cache_keys), _KEYS_PER_MESSAGE):
_queue_publish(redis_cache, tuple(cache_keys[start : start + _KEYS_PER_MESSAGE]), new_value, ttl)
await asyncio.sleep(0)
@ -232,17 +259,23 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us
class AuthCacheInvalidationSubscriber:
__slots__ = ("_additional_in_memory_caches", "_redis_cache", "_task", "_user_api_key_cache")
"""``counter_cache`` is the one cache a value-carrying message (a spend reset) is written into;
a delete reaches every additional cache. Writing a reset's thousands of counters into every
additional cache would evict what the others hold, the BYOK credentials among them."""
__slots__ = ("_additional_in_memory_caches", "_counter_cache", "_redis_cache", "_task", "_user_api_key_cache")
def __init__(
self,
redis_cache: "RedisCache",
user_api_key_cache: "UserApiKeyCache",
additional_in_memory_caches: Sequence["InMemoryCache"] = (),
counter_cache: "InMemoryCache | None" = None,
) -> None:
self._redis_cache = redis_cache
self._user_api_key_cache = user_api_key_cache
self._additional_in_memory_caches = tuple(additional_in_memory_caches)
self._counter_cache = counter_cache
self._task: asyncio.Task[None] | None = None
def start(self) -> None:
@ -298,12 +331,15 @@ class AuthCacheInvalidationSubscriber:
if parsed is None:
return
if parsed.new_value is not None:
for additional_cache in self._additional_in_memory_caches:
additional_cache.set_cache(parsed.cache_key, parsed.new_value, ttl=parsed.ttl)
if self._counter_cache is None:
return
for cache_key in parsed.cache_keys:
self._counter_cache.set_cache(cache_key, parsed.new_value, ttl=parsed.ttl)
return
_drop_from_memory(parsed.cache_key, self._user_api_key_cache)
for additional_cache in self._additional_in_memory_caches:
additional_cache.delete_cache(parsed.cache_key)
for cache_key in parsed.cache_keys:
_drop_from_memory(cache_key, self._user_api_key_cache)
for additional_cache in self._additional_in_memory_caches:
additional_cache.delete_cache(cache_key)
@staticmethod
async def _close_pubsub(pubsub: ConfigSyncPubSub) -> None:

6
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -1482,8 +1482,12 @@ class ResetBudgetJob:
new_value: Final = await ResetBudgetJob._window_carried_spend(window, counter_key, spend_counter_cache)
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=new_value)
if spend_counter_cache.redis_cache is not None:
from litellm.proxy.proxy_server import spend_counter_redis_ttl
try:
await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=new_value)
await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key, value=new_value, ttl=spend_counter_redis_ttl()
)
except Exception as redis_err:
verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err)
budget_duration: Final = window["budget_duration"]

11
litellm/proxy/db/spend_counter_reseed.py
View file

@ -137,7 +137,7 @@ class SpendCounterReseed:
# Per-window key/team counters share prefixes with primary counters
# but don't correspond to a DB row. Do not reject arbitrary entity IDs
# or tag names that merely contain ":window:".
if SpendCounterReseed._is_key_or_team_window_counter(counter_key):
if SpendCounterReseed.is_key_or_team_window_counter(counter_key):
return None
try:
row: Final = await bounded_db_lookup(
@ -199,7 +199,7 @@ class SpendCounterReseed:
return float(row.spend or 0.0)
@staticmethod
def _is_key_or_team_window_counter(counter_key: str) -> bool:
def is_key_or_team_window_counter(counter_key: str) -> bool:
for prefix in ("spend:key:", "spend:team:"):
if not counter_key.startswith(prefix):
continue
@ -267,10 +267,13 @@ class SpendCounterReseed:
current_value: float = float(db_spend)
try:
if spend_counter_cache.redis_cache is not None:
from litellm.proxy.proxy_server import spend_counter_redis_ttl
seeded: Final = await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key,
value=db_spend,
nx=True,
ttl=spend_counter_redis_ttl(),
)
if seeded:
current_value = float(db_spend)
@ -462,10 +465,13 @@ class SpendCounterReseed:
return None
try:
if spend_counter_cache.redis_cache is not None:
from litellm.proxy.proxy_server import spend_counter_redis_ttl
seeded: Final = await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key,
value=window_spend,
nx=True,
ttl=spend_counter_redis_ttl(),
)
if seeded:
current_value = window_spend
@ -475,6 +481,7 @@ class SpendCounterReseed:
current_value = await spend_counter_cache.redis_cache.async_increment(
key=counter_key,
value=window_spend,
ttl=spend_counter_redis_ttl(),
)
else:
current_value = float(current_cached_value)

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

@ -4263,14 +4263,16 @@ async def update_key_fn(
)
if data.spend is not None:
from litellm.proxy.proxy_server import spend_counter_cache
from litellm.proxy.proxy_server import spend_counter_cache, spend_counter_redis_ttl
counter_key: Final = f"spend:key:{hash_token_if_needed(key)}"
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=data.spend, ttl=60)
if spend_counter_cache.redis_cache is not None:
try:
with service_target(SPEND_COUNTERS_TARGET):
await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=data.spend, ttl=60)
await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key, value=data.spend, ttl=spend_counter_redis_ttl()
)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to update spend counter %s in Redis after key spend update: %s. "
@ -6820,12 +6822,18 @@ async def _set_spend_counter_with_floor_and_broadcast(counter_key: str, value: f
Without the broadcast, a worker that already cached the pre-reset key object
or floor marker keeps enforcing against it until its own TTL expires.
"""
from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
from litellm.proxy.proxy_server import (
SPEND_DB_FLOOR_CACHE_TTL_SECONDS,
spend_counter_cache,
spend_counter_redis_ttl,
)
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=value, ttl=60)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=value, ttl=60)
await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key, value=value, ttl=spend_counter_redis_ttl()
)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to update spend counter %s in Redis: %s. "

11
litellm/proxy/management_endpoints/organization_endpoints.py
View file

@ -817,6 +817,17 @@ async def update_organization(
include={"members": True, "teams": True, "litellm_budget_table": True},
)
if data.spend is not None:
from litellm.proxy.proxy_server import (
raise_if_spend_counters_unwritten,
set_spend_counters_with_floor_and_broadcast,
)
lost: Final = await set_spend_counters_with_floor_and_broadcast(
counter_keys=(f"spend:org:{data.organization_id}",), value=data.spend
)
raise_if_spend_counters_unwritten(lost, saved="The organization spend was saved")
return response

148
litellm/proxy/proxy_server.py
View file

@ -66,6 +66,7 @@ from pydantic import BaseModel, ConfigDict, Json, JsonValue, TypeAdapter, Valida
from pydantic.fields import FieldInfo, PydanticUndefined
from pydantic.main import IncEx
from pydantic_core import PydanticSerializationError, to_jsonable_python
from redis.exceptions import RedisError
from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never
from litellm._uuid import uuid
@ -81,6 +82,7 @@ from litellm.constants import (
BASE_MCP_ROUTE,
DAILY_TAG_SPEND_BATCH_MULTIPLIER,
DEFAULT_MAX_RECURSE_DEPTH,
DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE,
DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL,
DEFAULT_SHARED_HEALTH_CHECK_TTL,
DEFAULT_SLACK_ALERTING_THRESHOLD,
@ -366,6 +368,7 @@ from litellm.constants import (
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG,
SPEND_CAPTURE_RATE_CHECK_JOB_ID,
SPEND_COUNTER_REDIS_TTL_SECONDS,
SPEND_LOG_SIZE_CHECK_INTERVAL,
USER_SPEND_ALERTS_JOB_ID,
WEEKLY_SPEND_REPORT_JOB_ID,
@ -516,6 +519,7 @@ from litellm.proxy.common_utils.anthropic_error_payload import (
)
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
AuthCacheInvalidationSubscriber,
publish_auth_cache_invalidations,
)
from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy
from litellm.proxy.common_utils.codex_model_catalog import codex_model_list_body
@ -3338,6 +3342,23 @@ async def _drain_spend_event_producer_on_shutdown() -> None:
# Bounds authoritative DB re-reads when enforcing a budget against a
# stale-low spend counter: at most one DB read per counter per window.
SPEND_DB_FLOOR_CACHE_TTL_SECONDS: Final = 5
SPEND_RESET_FENCE_TTL_SECONDS: Final = 2 * SPEND_DB_FLOOR_CACHE_TTL_SECONDS
SPEND_DB_FLOOR_MARKER_PREFIX: Final = "spend_db_floor:"
_SPEND_COUNTER_TTL_BATCH_WRITE_MULTIPLE: Final = 10
def spend_counter_redis_ttl() -> int:
"""Redis TTL for spend counters. A counter that expires before its spend is flushed to the DB reseeds from a
stale-low row and re-admits a key that is over budget, so it must outlive many batch-write intervals."""
batch_write_at: Final = (
proxy_batch_write_at if isinstance(proxy_batch_write_at, (int, float)) else PROXY_BATCH_WRITE_AT
)
configured_redis_ttl: Final = int(litellm.default_redis_ttl) if litellm.default_redis_ttl is not None else 0
return max(
SPEND_COUNTER_REDIS_TTL_SECONDS,
int(_SPEND_COUNTER_TTL_BATCH_WRITE_MULTIPLE * batch_write_at),
configured_redis_ttl,
)
def _fail_closed_budget_enforcement() -> bool:
@ -3470,7 +3491,12 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None
if spend_counter_cache.redis_cache is not None:
forget_spend_counter(counter_key)
try:
repaired: Final = await spend_counter_cache.redis_cache.async_set_max(key=counter_key, value=db_spend)
repaired: Final = await spend_counter_cache.redis_cache.async_set_max(
key=counter_key,
value=db_spend,
ttl=spend_counter_redis_ttl(),
fence_key=spend_counter_cache.redis_cache.reset_fence_key(counter_key),
)
except Exception:
verbose_proxy_logger.debug(
"Unable to repair stale spend counter %s in Redis",
@ -4406,7 +4432,7 @@ async def refresh_spend_counter_ttl(counter_key: str) -> bool:
if spend_counter_cache.redis_cache is None:
return False
try:
return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key)
return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key, ttl=spend_counter_redis_ttl())
except Exception as e:
verbose_proxy_logger.debug("spend counter TTL refresh skipped for %s: %s", counter_key, e)
return False
@ -4419,6 +4445,7 @@ async def increment_spend_counter_cache(counter_key: str, increment: float) -> f
current_value: Final = await spend_counter_cache.redis_cache.async_increment(
key=counter_key,
value=increment,
ttl=spend_counter_redis_ttl(),
refresh_ttl=True,
)
except Exception:
@ -4450,6 +4477,118 @@ async def invalidate_spend_counter(counter_key: str) -> None:
_invalidate_spend_counter: Final = invalidate_spend_counter
_IN_MEMORY_COUNTER_KEYS: Final = TypeAdapter(tuple[str, ...])
_SPEND_COUNTER_STORE_FAILURES: Final = (RedisError, RedisCircuitBreakerOpenError, OSError, ValidationError)
def _spend_counter_keys_in_memory(prefix: str) -> frozenset[str]:
floor_prefix: Final = f"{SPEND_DB_FLOOR_MARKER_PREFIX}{prefix}"
cached_keys: Final = _IN_MEMORY_COUNTER_KEYS.validate_python(tuple(spend_counter_cache.in_memory_cache.cache_dict))
return frozenset(
key.removeprefix(SPEND_DB_FLOOR_MARKER_PREFIX) for key in cached_keys if key.startswith((prefix, floor_prefix))
)
async def _spend_counter_keys_in_redis(redis_cache: RedisCache, prefix: str) -> frozenset[str] | Exception:
namespace_prefix: Final = f"{redis_cache.namespace}:" if redis_cache.namespace else ""
try:
scanned: Final = await redis_cache.async_scan_iter(pattern=prefix, count=500, unbounded=True)
except _SPEND_COUNTER_STORE_FAILURES as e:
return e
return frozenset(key.removeprefix(namespace_prefix) for key in scanned)
def _keys_per_reset_pipeline(max_batch_size: int) -> int:
return max(1, max_batch_size // 2)
_KEYS_PER_RESET_PIPELINE: Final = _keys_per_reset_pipeline(DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE)
async def _write_reset_pipeline(
redis_cache: RedisCache, batch: Sequence[str], value: float, ttl: int
) -> Exception | None:
"""The fences go first, so a repair that lands between a fence and its counter already finds the fence."""
fences: Final = tuple((redis_cache.reset_fence_key(key), 1, SPEND_RESET_FENCE_TTL_SECONDS) for key in batch)
try:
return await redis_cache.async_set_cache_pipeline_with_ttls((*fences, *((key, value, ttl) for key in batch)))
except RedisCircuitBreakerOpenError as e:
return e
async def _write_spend_counters_to_redis(
redis_cache: RedisCache, keys: Sequence[str], value: float
) -> Exception | None:
ttl: Final = spend_counter_redis_ttl()
for start in range(0, len(keys), _KEYS_PER_RESET_PIPELINE):
if (
lost := await _write_reset_pipeline(redis_cache, keys[start : start + _KEYS_PER_RESET_PIPELINE], value, ttl)
) is not None:
return lost
return None
@with_service_target(SPEND_COUNTERS_TARGET)
async def set_spend_counters_with_floor_and_broadcast(counter_keys: Collection[str], value: float) -> Exception | None:
"""A spend reset SETs each counter and its spend_db_floor marker on every pod, never
deletes them: a pod that cached the pre-reset floor within SPEND_DB_FLOOR_CACHE_TTL_SECONDS
would otherwise repair the reseeded counter right back to the old total. The floors are
broadcast before the counters so a subscriber never holds the new counter under the old
floor, and a reset fence in Redis holds any repair that still races the broadcast.
Returns the Redis error when the counters could not be written there: this pod and the
broadcast still carry the new value, but the other pods fall back to the stale Redis
counters once their copies expire, so the caller must not report success."""
keys: Final = tuple(sorted(counter_keys))
if len(keys) == 0:
return None
floor_keys: Final = tuple(f"{SPEND_DB_FLOOR_MARKER_PREFIX}{key}" for key in keys)
for key, floor_key in zip(keys, floor_keys, strict=True):
spend_counter_cache.in_memory_cache.set_cache(key=key, value=value, ttl=60)
spend_counter_cache.in_memory_cache.set_cache(key=floor_key, value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS)
redis_cache: Final = spend_counter_cache.redis_cache
lost: Final = None if redis_cache is None else await _write_spend_counters_to_redis(redis_cache, keys, value)
await publish_auth_cache_invalidations(floor_keys, new_value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS)
await publish_auth_cache_invalidations(keys, new_value=value, ttl=60)
return lost
@with_service_target(SPEND_COUNTERS_TARGET)
async def reset_spend_counters_with_prefix(prefix: str) -> Exception | None:
"""Resets every counter under ``prefix`` this pod or Redis knows; returns the Redis error
when the counters there could not be listed or written, with this pod already reset."""
redis_cache: Final = spend_counter_cache.redis_cache
in_redis: Final[frozenset[str] | Exception] = (
frozenset() if redis_cache is None else await _spend_counter_keys_in_redis(redis_cache, prefix)
)
listed: Final[frozenset[str]] = frozenset() if isinstance(in_redis, Exception) else in_redis
counters: Final = frozenset(
key
for key in _spend_counter_keys_in_memory(prefix) | listed
if not SpendCounterReseed.is_key_or_team_window_counter(key)
)
lost: Final = await set_spend_counters_with_floor_and_broadcast(counters, value=0.0)
return in_redis if isinstance(in_redis, Exception) else lost
def raise_if_spend_counters_unwritten(lost: Exception | None, saved: str) -> None:
"""Maps a lost Redis counter write to the public error: the database write stands, so the
caller retries the same request once Redis accepts writes rather than undoing anything."""
if lost is None:
return
raise ProxyException(
message=(
f"{saved}, but the cached spend counters could not be reset in Redis ({lost}). Budget checks on other "
"instances keep the previous spend until those counters expire, so retry this request once Redis "
"accepts writes."
),
type=ProxyErrorTypes.internal_server_error,
param=None,
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
async def _apply_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> None:
if _defer_spend_counter_increments(pending):
@ -4472,7 +4611,7 @@ def _defer_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) ->
batch: Final = active_post_call_redis_batch(redis_cache)
if batch is None:
return False
ttl: Final = redis_cache.get_ttl()
ttl: Final = spend_counter_redis_ttl()
with service_target(SPEND_COUNTERS_TARGET):
for item in pending:
batch.increment(item.counter_key, item.increment, ttl).on_settled(_settle_spend_counter_increment(item))
@ -4526,7 +4665,7 @@ async def run_spend_counter_pipeline(pending: Sequence[PendingSpendIncrement]) -
for item in pending
]
)
ttl: Final = redis_cache.get_ttl()
ttl: Final = spend_counter_redis_ttl()
increment_list: Final = [
RedisPipelineIncrementOperation(key=item.counter_key, increment_value=item.increment, ttl=ttl)
for item in pending
@ -9054,6 +9193,7 @@ class ProxyConfig:
redis_cache=redis_cache,
user_api_key_cache=user_api_key_cache,
additional_in_memory_caches=(spend_counter_cache.in_memory_cache, byok_credential_cache),
counter_cache=spend_counter_cache.in_memory_cache,
)
self.auth_cache_invalidation_subscriber = subscriber
subscriber.start()

9
litellm/proxy/spend_tracking/spend_management_endpoints.py
View file

@ -3611,7 +3611,11 @@ async def global_spend_reset():
3. LiteLLM_TeamTable spend will be set = 0
"""
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.proxy_server import (
prisma_client,
raise_if_spend_counters_unwritten,
reset_spend_counters_with_prefix,
)
if prisma_client is None:
raise ProxyException(
@ -3626,6 +3630,9 @@ async def global_spend_reset():
_ = await transaction.litellm_teamtable.update_many(data={"spend": 0.0}, where={})
_ = await transaction.query_raw(_LOCK_ALL_KEYS_SQL)
_ = await transaction.litellm_verificationtoken.update_many(data={"spend": 0.0}, where={})
keys_lost: Final = await reset_spend_counters_with_prefix(prefix="spend:key:")
teams_lost: Final = await reset_spend_counters_with_prefix(prefix="spend:team:")
raise_if_spend_counters_unwritten(keys_lost if keys_lost is not None else teams_lost, saved="Spend was reset")
return {
"message": "Spend for all API Keys and Teams reset successfully",

33
tests/integration/_support/spend_counters.py Normal file
View file

@ -0,0 +1,33 @@
import os
from collections.abc import Iterator
from contextlib import contextmanager
from hashlib import sha256
from typing import Final
from redis import Redis
SPEND_COUNTER_TTL_SECONDS: Final = 3600
TTL_SLACK_SECONDS: Final = 600
@contextmanager
def rig_cache() -> Iterator[Redis]:
with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache:
yield cache
def key_counter(key: str) -> str:
return f"spend:key:{sha256(key.encode()).hexdigest()}"
def counter_ttl(cache: Redis, counter_key: str) -> int:
return int(str(cache.ttl(counter_key)))
def counter_value(cache: Redis, counter_key: str) -> float | None:
raw: Final = cache.get(counter_key)
return None if raw is None else float(str(raw, "utf-8") if isinstance(raw, bytes) else str(raw))
def outlives(ttl: int, expected: int = SPEND_COUNTER_TTL_SECONDS) -> bool:
return expected - TTL_SLACK_SECONDS < ttl <= expected

91
tests/integration/_support/upstream.py
View file

@ -6,6 +6,7 @@ import base64
import json
import os
import struct
import time
import uuid
import zlib
from collections import deque
@ -60,6 +61,46 @@ INTERNAL_FIELDS: Final = frozenset(
)
RESPONSES_USAGE: Final = {"input_tokens": 20, "output_tokens": 20, "total_tokens": 40}
def responses_body(model: str) -> dict[str, JsonValue]:
return {
"id": f"resp_{uuid.uuid4().hex[:24]}",
"object": "response",
"created_at": int(time.time()),
"status": "completed",
"model": model,
"output": [
{
"id": f"msg_{uuid.uuid4().hex[:24]}",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "Hello.", "annotations": []}],
}
],
"usage": dict(RESPONSES_USAGE),
}
async def responses_stream(completed: Mapping[str, JsonValue]) -> AsyncIterator[str]:
item: Final = cast(dict[str, JsonValue], cast(list[JsonValue], completed["output"])[0])
events: Final = (
("response.created", {"response": {**completed, "status": "in_progress", "output": [], "usage": None}}),
("response.output_item.added", {"output_index": 0, "item": {**item, "status": "in_progress", "content": []}}),
(
"response.output_text.delta",
{"output_index": 0, "content_index": 0, "item_id": item["id"], "delta": "Hello."},
),
("response.output_item.done", {"output_index": 0, "item": item}),
("response.completed", {"response": dict(completed)}),
)
for sequence_number, (event_type, payload) in enumerate(events):
data: Final = json.dumps({"type": event_type, "sequence_number": sequence_number, **payload})
yield f"event: {event_type}\ndata: {data}\n\n"
def error_type(status: int) -> str:
if status == 429:
return "rate_limit_error"
@ -222,24 +263,41 @@ class Provider:
for message in messages
):
return JSONResponse({"error": {"message": "Invalid selected message contract"}}, status_code=400)
script: Final = self.scripts.get(str(body["model"]))
if script is not None:
if not script:
return JSONResponse({"error": {"message": "Script exhausted", "type": "api_error"}}, status_code=500)
status: Final = script.popleft()
if status != 200:
return JSONResponse(
{
"error": {
"message": "Controlled provider failure",
"type": error_type(status),
"code": str(status),
}
},
status_code=status,
)
failure: Final = self._scripted_failure(str(body["model"]))
if failure is not None:
return failure
return await chat_completions(request)
def _scripted_failure(self, model: str) -> Response | None:
script: Final = self.scripts.get(model)
if script is None:
return None
if not script:
return JSONResponse({"error": {"message": "Script exhausted", "type": "api_error"}}, status_code=500)
status: Final = script.popleft()
if status == 200:
return None
return JSONResponse(
{"error": {"message": "Controlled provider failure", "type": error_type(status), "code": str(status)}},
status_code=status,
)
async def responses(self, request: Request) -> Response:
body: Final = JSON_OBJECT.validate_json(await request.body())
self.observations.put(
Observation(request.url.path, request.headers.get("authorization", ""), body, request.method)
)
model: Final = body.get("model")
if not isinstance(model, str) or "input" not in body:
return JSONResponse({"error": {"message": "model and input are required"}}, status_code=400)
failure: Final = self._scripted_failure(model)
if failure is not None:
return failure
completed: Final = responses_body(model)
if body.get("stream") is True:
return StreamingResponse(responses_stream(completed), media_type="text/event-stream")
return JSONResponse(completed)
async def vector_store_search(self, request: Request) -> Response:
body: Final = JSON_OBJECT.validate_json(await request.body())
self.observations.put(
@ -645,6 +703,7 @@ class Provider:
Route("/_cost_map", self.cost_map, methods=["GET"]),
Route("/_oauth/token", self.oauth_token, methods=["POST"]),
Route("/v1/chat/completions", self.chat, methods=["POST"]),
Route("/v1/responses", self.responses, methods=["POST"]),
Route("/v1/completions", completions, methods=["POST"]),
Route("/v1/embeddings", embeddings, methods=["POST"]),
Route("/v1/moderations", moderations, methods=["POST"]),

88
tests/integration/caching/test_batch_redis_requests_prefetch.py Normal file
View file

@ -0,0 +1,88 @@
"""The ``batch_redis_requests`` callback prefetches a key's response cache entries
into memory with one Redis scan, so the next identical request is served without a
Redis round trip for the entry."""
import os
import threading
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
from hashlib import sha256
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway, eventually, object_value
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.redis_process import OwnedRedis, owned_redis
from redis import Redis
from redis.exceptions import TimeoutError as RedisTimeoutError
ROOT: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3])
CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120)
def _config_with_batch_redis_requests(directory: Path) -> Path:
shared: Final = object_value(yaml.safe_load((ROOT / "tests/integration/proxy_config.yaml").read_text()))
litellm_settings: Final = {**object_value(shared["litellm_settings"]), "callbacks": ["batch_redis_requests"]}
path: Final = directory / f"proxy_config_{uuid.uuid4().hex}.yaml"
path.write_text(yaml.safe_dump({**shared, "litellm_settings": litellm_settings}, sort_keys=False))
return path
@contextmanager
def _monitored(redis: OwnedRedis) -> Iterator[list[str]]:
commands: Final[list[str]] = []
stopped: Final = threading.Event()
with Redis(host=redis.host, port=redis.port, socket_timeout=0.3) as client, client.monitor() as monitor:
def pump() -> None:
while not stopped.is_set():
try:
commands.append(str(monitor.next_command()["command"]))
except RedisTimeoutError:
continue
thread: Final = threading.Thread(target=pump, daemon=True)
thread.start()
try:
yield commands
finally:
stopped.set()
thread.join(timeout=5)
@pytest.mark.timeout(CELL_TIMEOUT)
def test_the_prefetched_response_cache_entry_is_served_from_memory_on_the_next_request(
gateway: Gateway, tmp_path: Path
) -> None:
with owned_redis(tmp_path) as redis:
overrides: Final = {"REDIS_HOST": redis.host, "REDIS_PORT": str(redis.port)}
with (
owned_proxy_process(
gateway, tmp_path, overrides, config=_config_with_batch_redis_requests(tmp_path)
) as owned,
owned.gateway.scenario() as scenario,
Redis(host=redis.host, port=redis.port) as cache,
httpx.Client(base_url=gateway.upstream_url, timeout=30) as upstream,
):
model: Final = scenario.model()
key: Final = scenario.key(models=[model])
body: Final = {"model": model, "messages": [{"role": "user", "content": f"prefetch {uuid.uuid4().hex}"}]}
namespace: Final = f"litellm:{sha256(key.encode()).hexdigest()}:acompletion"
first: Final = owned.gateway.request("POST", "/v1/chat/completions", body, key=key)
assert first.status_code == 200, first.text
entries: Final = eventually(lambda: cache.keys(f"{namespace}:*"), lambda keys: len(keys) == 1, seconds=10)
upstream.delete("/__observations").raise_for_status()
with _monitored(redis) as commands:
second: Final = owned.gateway.request("POST", "/v1/chat/completions", body, key=key)
assert second.status_code == 200, second.text
assert second.json()["id"] == first.json()["id"], (first.text, second.text)
assert upstream.get("/__observations").json()["requests"] == []
fetched: Final = [command for command in commands if command.startswith(f"GET {namespace}:")]
assert fetched == [], (entries, commands)
third: Final = owned.gateway.request("POST", "/v1/chat/completions", body, key=key)
assert third.status_code == 200 and third.json()["id"] == first.json()["id"], third.text
assert upstream.get("/__observations").json()["requests"] == []

488
tests/integration/spend/test_spend_counter_redis_ttl.py Normal file
View file

@ -0,0 +1,488 @@
import asyncio
import json
import uuid
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from hashlib import sha256
from typing import Final
import anthropic
import httpx
import openai
import pytest
from integration._support.client import Gateway, Scenario, eventually
from integration._support.database import read_rows
from integration._support.spend_counters import (
counter_ttl,
counter_value,
key_counter,
outlives,
rig_cache,
)
from pydantic import JsonValue
from redis import Redis
RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}
ONE_CALL_COST: Final = 0.06
ANTHROPIC_HEADERS: Final = {"anthropic-version": "2023-06-01"}
@dataclass(frozen=True, slots=True)
class Entities:
model: str
key: str
team: str
user: str
organization: str
end_user: str
tag: str
def counters(self) -> Mapping[str, str]:
return {
"key": key_counter(self.key),
"team": f"spend:team:{self.team}",
"team_member": f"spend:team_member:{self.user}:{self.team}",
"user": f"spend:user:{self.user}",
"end_user": f"spend:end_user:{self.end_user}",
"tag": f"spend:tag:{self.tag}",
"org": f"spend:org:{self.organization}",
}
def _entities(gateway: Gateway, scenario: Scenario) -> Entities:
model: Final = scenario.model(**RATES)
models: Final = [model]
organization: Final = scenario.organization(max_budget=50.0)
team: Final = scenario.team(models=models, max_budget=20.0, organization_id=organization)
user: Final = scenario.user(max_budget=20.0)
added: Final = gateway.request(
"POST",
"/team/member_add",
{"team_id": team, "member": {"user_id": user, "role": "user"}, "max_budget_in_team": 20.0},
)
assert added.status_code == 200, added.text
tag: Final = f"integration-tag-{uuid.uuid4().hex}"
key: Final = scenario.key(team_id=team, user_id=user, models=models, max_budget=5.0, tags=[tag])
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
return Entities(model, key, team, user, organization, end_user, tag)
def _base(gateway: Gateway) -> str:
return str(gateway.client.base_url).rstrip("/")
def _chat_openai_sync(gateway: Gateway, entities: Entities, text: str) -> str:
with openai.OpenAI(base_url=f"{_base(gateway)}/v1", api_key=entities.key, max_retries=0, timeout=60) as client:
response: Final = client.chat.completions.create(
model=entities.model, messages=[{"role": "user", "content": text}], user=entities.end_user
)
assert response.usage is not None and response.usage.total_tokens == 40, response
return response.id
def _chat_openai_async_stream(gateway: Gateway, entities: Entities, text: str) -> str:
async def run() -> str:
async with openai.AsyncOpenAI(
base_url=f"{_base(gateway)}/v1", api_key=entities.key, max_retries=0, timeout=60
) as client:
stream: Final = await client.chat.completions.create(
model=entities.model,
messages=[{"role": "user", "content": text}],
user=entities.end_user,
stream=True,
stream_options={"include_usage": True},
)
ids: Final = [chunk.id async for chunk in stream]
assert len(set(ids)) == 1, ids
return ids[0]
return asyncio.run(run())
def _chat_httpx_stream(gateway: Gateway, entities: Entities, text: str) -> str:
with gateway.client.stream(
"POST",
"/v1/chat/completions",
json={
"model": entities.model,
"messages": [{"role": "user", "content": text}],
"user": entities.end_user,
"stream": True,
"stream_options": {"include_usage": True},
},
headers={"Authorization": f"Bearer {entities.key}"},
) as response:
assert response.status_code == 200, response.read()
frames: Final = [line.removeprefix("data: ") for line in response.iter_lines() if line.startswith("data: ")]
ids: Final = {json.loads(frame)["id"] for frame in frames if frame != "[DONE]"}
assert len(ids) == 1, frames
return ids.pop()
def _messages_anthropic_sync(gateway: Gateway, entities: Entities, text: str) -> str:
with anthropic.Anthropic(base_url=_base(gateway), api_key=entities.key, max_retries=0, timeout=60) as client:
message: Final = client.messages.create(
model=entities.model,
max_tokens=32,
messages=[{"role": "user", "content": text}],
metadata={"user_id": entities.end_user},
)
assert message.usage.output_tokens > 0, message
return message.id
def _messages_anthropic_async_stream(gateway: Gateway, entities: Entities, text: str) -> str:
async def run() -> str:
async with anthropic.AsyncAnthropic(
base_url=_base(gateway), api_key=entities.key, max_retries=0, timeout=60
) as client:
async with client.messages.stream(
model=entities.model,
max_tokens=32,
messages=[{"role": "user", "content": text}],
metadata={"user_id": entities.end_user},
) as stream:
async for _ in stream.text_stream:
pass
final: Final = await stream.get_final_message()
assert final.usage.output_tokens > 0, final
return final.id
return asyncio.run(run())
def _messages_httpx(gateway: Gateway, entities: Entities, text: str) -> str:
response: Final = gateway.request(
"POST",
"/v1/messages",
{
"model": entities.model,
"max_tokens": 32,
"messages": [{"role": "user", "content": text}],
"metadata": {"user_id": entities.end_user},
},
key=entities.key,
headers=ANTHROPIC_HEADERS,
)
assert response.status_code == 200, response.text
return str(response.json()["id"])
def _responses_httpx(gateway: Gateway, entities: Entities, text: str) -> str:
response: Final = gateway.request(
"POST",
"/v1/responses",
{"model": entities.model, "input": text, "user": entities.end_user},
key=entities.key,
)
assert response.status_code == 200, response.text
return str(response.json()["id"])
def _responses_httpx_stream(gateway: Gateway, entities: Entities, text: str) -> str:
with gateway.client.stream(
"POST",
"/v1/responses",
json={"model": entities.model, "input": text, "user": entities.end_user, "stream": True},
headers={"Authorization": f"Bearer {entities.key}"},
) as response:
assert response.status_code == 200, response.read()
lines: Final = [line for line in response.iter_lines() if line.startswith("data: ")]
events: Final = [json.loads(line.removeprefix("data: ")) for line in lines if line != "data: [DONE]"]
completed: Final = [event for event in events if event.get("type") == "response.completed"]
assert len(completed) == 1, [event.get("type") for event in events]
return str(completed[0]["response"]["id"])
@dataclass(frozen=True, slots=True)
class Caller:
call: Callable[[Gateway, Entities, str], str]
upstream_path: str
exact_cost: float | None
CALLERS: Final[Mapping[str, Caller]] = {
"chat-openai-sync": Caller(_chat_openai_sync, "/v1/chat/completions", ONE_CALL_COST),
"chat-openai-async-stream": Caller(_chat_openai_async_stream, "/v1/chat/completions", ONE_CALL_COST),
"chat-httpx-stream": Caller(_chat_httpx_stream, "/v1/chat/completions", ONE_CALL_COST),
"messages-anthropic-sync": Caller(_messages_anthropic_sync, "/v1/responses", ONE_CALL_COST),
"messages-anthropic-async-stream": Caller(_messages_anthropic_async_stream, "/v1/responses", ONE_CALL_COST),
"messages-httpx": Caller(_messages_httpx, "/v1/responses", ONE_CALL_COST),
"responses-httpx": Caller(_responses_httpx, "/v1/responses", ONE_CALL_COST),
"responses-httpx-stream": Caller(_responses_httpx_stream, "/v1/responses", ONE_CALL_COST),
}
def _spend_rows(key: str) -> list[dict[str, JsonValue]]:
return read_rows(
'SELECT request_id, spend, cache_hit FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"',
(sha256(key.encode()).hexdigest(),),
)
def _one_spend_row(key: str) -> dict[str, JsonValue]:
rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 1, seconds=70)
return rows[0]
def _observed(cache: Redis, counters: Mapping[str, str]) -> Mapping[str, tuple[float | None, int]]:
return {name: (counter_value(cache, counter), counter_ttl(cache, counter)) for name, counter in counters.items()}
def _assert_every_counter_outlives(
cache: Redis, counters: Mapping[str, str], spend: float
) -> Mapping[str, tuple[float | None, int]]:
observed: Final = eventually(
lambda: _observed(cache, counters),
lambda states: all(
value is not None and value == pytest.approx(spend, abs=1e-9) for value, _ in states.values()
),
seconds=20,
)
for name, (_, ttl) in observed.items():
assert outlives(ttl), (name, observed)
return observed
@pytest.mark.parametrize("caller", sorted(CALLERS))
def test_every_counter_scope_outlives_the_batch_write(gateway: Gateway, caller: str) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
rig_cache() as cache,
):
entities: Final = _entities(gateway, scenario)
text: Final = f"counter ttl {caller} {uuid.uuid4().hex}"
upstream.delete("/__observations").raise_for_status()
response_id: Final = CALLERS[caller].call(gateway, entities, text)
received: Final = upstream.get("/__observations").json()["requests"]
paths: Final = [request["path"] for request in received if text in json.dumps(request["body"])]
assert len(paths) == 1 and paths[0].endswith(CALLERS[caller].upstream_path), received
row: Final = _one_spend_row(entities.key)
spend: Final = float(str(row["spend"]))
assert spend == pytest.approx(CALLERS[caller].exact_cost or spend) and spend > 0, row
assert response_id, row
_assert_every_counter_outlives(cache, entities.counters(), spend)
def test_a_response_cache_hit_keeps_the_counter_and_its_ttl(gateway: Gateway) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
rig_cache() as cache,
):
entities: Final = _entities(gateway, scenario)
body: Final = {
"model": entities.model,
"messages": [{"role": "user", "content": f"cache twin {uuid.uuid4().hex}"}],
"user": entities.end_user,
}
first: Final = gateway.request("POST", "/v1/chat/completions", body, key=entities.key)
assert first.status_code == 200, first.text
row: Final = _one_spend_row(entities.key)
spend: Final = float(str(row["spend"]))
before: Final = _assert_every_counter_outlives(cache, entities.counters(), spend)
upstream.delete("/__observations").raise_for_status()
second: Final = gateway.request("POST", "/v1/chat/completions", body, key=entities.key)
assert second.status_code == 200, second.text
assert second.json()["id"] == first.json()["id"], (first.text, second.text)
assert upstream.get("/__observations").json()["requests"] == []
rows: Final = eventually(lambda: _spend_rows(entities.key), lambda values: len(values) == 2, seconds=70)
assert str(rows[1]["cache_hit"]) == "True" and float(str(rows[1]["spend"])) == 0.0, rows
after: Final = _assert_every_counter_outlives(cache, entities.counters(), spend)
assert {name: value for name, (value, _) in after.items()} == pytest.approx(
{name: value for name, (value, _) in before.items()}, abs=1e-9
), (before, after)
def _exhausted_key(gateway: Gateway, scenario: Scenario) -> tuple[str, str]:
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST)
served: Final = gateway.chat(model, key=key, text=f"exhaust {uuid.uuid4().hex}")
assert served["usage"]["total_tokens"] == 40, served
eventually(
lambda: read_rows(
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),)
),
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= ONE_CALL_COST,
seconds=70,
)
return model, key
def _chat(gateway: Gateway, model: str, key: str) -> httpx.Response:
return gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"after {uuid.uuid4().hex}"}]},
key=key,
)
def _assert_denied(response: httpx.Response) -> None:
assert response.status_code == 422, response.text
assert response.json()["error"]["type"] == "budget_exceeded", response.text
def test_a_stale_low_counter_is_repaired_with_the_long_ttl(gateway: Gateway) -> None:
with gateway.scenario() as scenario, rig_cache() as cache:
model: Final = scenario.model(**RATES)
key: Final = scenario.key(
models=[model], budget_limits=[{"max_budget": ONE_CALL_COST, "budget_duration": "1d"}]
)
counter: Final = f"{key_counter(key)}:window:1d"
assert gateway.chat(model, key=key, text=f"exhaust window {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40
_one_spend_row(key)
assert counter_value(cache, counter) == pytest.approx(ONE_CALL_COST)
served: Final[list[int]] = []
def stale_then_request() -> int:
assert cache.set(counter, "0.01")
code: Final = _chat(gateway, model, key).status_code
served.append(code)
return code
assert eventually(stale_then_request, lambda code: code == 422, seconds=30) == 422
assert set(served) == {200, 422}, served
repaired: Final = (counter_value(cache, counter), counter_ttl(cache, counter))
floor: Final = repaired[0]
assert floor is not None and ONE_CALL_COST - 1e-9 <= floor <= served.count(200) * ONE_CALL_COST + 1e-9, (
repaired,
served,
)
assert outlives(repaired[1]), repaired
def test_a_lost_counter_is_reseeded_from_the_db_with_the_long_ttl(gateway: Gateway) -> None:
with gateway.scenario() as scenario, rig_cache() as cache:
model, key = _exhausted_key(gateway, scenario)
counter: Final = key_counter(key)
assert cache.delete(counter) == 1
_assert_denied(_chat(gateway, model, key))
reseeded: Final = eventually(
lambda: (counter_value(cache, counter), counter_ttl(cache, counter)),
lambda state: state[0] is not None and state[0] >= ONE_CALL_COST and outlives(state[1]),
seconds=15,
)
assert reseeded[0] == pytest.approx(ONE_CALL_COST), reseeded
def test_key_update_spend_sets_the_counter_with_the_long_ttl(gateway: Gateway) -> None:
with gateway.scenario() as scenario, rig_cache() as cache:
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=5.0)
counter: Final = key_counter(key)
assert gateway.chat(model, key=key, text=f"before update {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40
_one_spend_row(key)
updated: Final = gateway.request("POST", "/key/update", {"key": key, "spend": 1.25})
assert updated.status_code == 200, updated.text
assert (counter_value(cache, counter), outlives(counter_ttl(cache, counter))) == (1.25, True), updated.text
served: Final = _chat(gateway, model, key)
assert served.status_code == 200, served.text
eventually(
lambda: counter_value(cache, counter),
lambda value: value is not None and value == pytest.approx(1.25 + ONE_CALL_COST),
seconds=15,
)
assert outlives(counter_ttl(cache, counter))
def test_key_reset_spend_sets_the_counter_with_the_long_ttl(gateway: Gateway) -> None:
with gateway.scenario() as scenario, rig_cache() as cache:
model, key = _exhausted_key(gateway, scenario)
counter: Final = key_counter(key)
_assert_denied(_chat(gateway, model, key))
reset: Final = gateway.request("POST", f"/key/{key}/reset_spend", {"reset_to": 0.0})
assert reset.status_code == 200, reset.text
assert reset.json()["key_hash"] == sha256(key.encode()).hexdigest(), reset.text
assert (counter_value(cache, counter), outlives(counter_ttl(cache, counter))) == (0.0, True), reset.text
served: Final = _chat(gateway, model, key)
assert served.status_code == 200, served.text
def test_team_member_reset_spend_sets_the_counter_with_the_long_ttl(gateway: Gateway) -> None:
with gateway.scenario() as scenario, rig_cache() as cache:
model: Final = scenario.model(**RATES)
team: Final = scenario.team(models=[model])
user: Final = scenario.user()
added: Final = gateway.request(
"POST",
"/team/member_add",
{"team_id": team, "member": {"user_id": user, "role": "user"}, "max_budget_in_team": ONE_CALL_COST},
)
assert added.status_code == 200, added.text
key: Final = scenario.key(team_id=team, user_id=user, models=[model])
counter: Final = f"spend:team_member:{user}:{team}"
assert gateway.chat(model, key=key, text=f"member {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40
eventually(
lambda: read_rows(
'SELECT spend FROM "LiteLLM_TeamMembership" WHERE team_id=%s AND user_id=%s', (team, user)
),
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= ONE_CALL_COST,
seconds=70,
)
_assert_denied(_chat(gateway, model, key))
reset: Final = gateway.request("POST", f"/team/{team}/member/{user}/reset_spend", {"reset_to": 0.0})
assert reset.status_code == 200, reset.text
assert reset.json()["previous_spend"] == pytest.approx(ONE_CALL_COST), reset.text
assert (counter_value(cache, counter), outlives(counter_ttl(cache, counter))) == (0.0, True), reset.text
served: Final = _chat(gateway, model, key)
assert served.status_code == 200, served.text
def test_malformed_spend_writes_are_rejected_and_leave_the_counter_alone(gateway: Gateway) -> None:
with gateway.scenario() as scenario, rig_cache() as cache:
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=5.0)
bystander: Final = scenario.key(models=[model], max_budget=5.0)
counter: Final = key_counter(key)
assert gateway.chat(model, key=key, text=f"before bad writes {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40
_one_spend_row(key)
before: Final = (counter_value(cache, counter), counter_ttl(cache, counter))
assert before[0] == pytest.approx(ONE_CALL_COST) and outlives(before[1]), before
rejected: Final = tuple(
gateway.request(*attempt)
for attempt in (
("POST", "/key/update", {"key": key, "spend": "a" * 5000}),
("POST", "/key/update", {"key": key, "spend": [1.0]}),
("POST", "/key/update", {"key": key, "spend": {"spend": 1.0}}),
("POST", f"/key/{key}/reset_spend", {"reset_to": "abc"}),
("POST", f"/key/{key}/reset_spend", {"reset_to": True}),
("POST", f"/key/{key}/reset_spend", {"reset_to": [0.0]}),
("POST", f"/key/{key}/reset_spend", {"reset_to": ONE_CALL_COST * 10}),
)
)
for response in rejected:
assert response.status_code in (400, 422), response.text
assert "error" in response.json() or "detail" in response.json(), response.text
assert (counter_value(cache, counter), counter_ttl(cache, counter) <= before[1]) == (before[0], True)
assert outlives(counter_ttl(cache, counter))
served: Final = _chat(gateway, model, bystander)
assert served.status_code == 200, served.text
def _landed_chat(gateway: Gateway, model: str, key: str, index: int) -> str:
served: Final = gateway.chat(model, key=key, text=f"repeat {index} {uuid.uuid4().hex}")
eventually(lambda: _spend_rows(key), lambda rows: len(rows) == index + 1, seconds=70)
return str(served["id"])
def test_repeated_requests_log_once_each_and_the_counter_matches_their_sum(gateway: Gateway) -> None:
with gateway.scenario() as scenario, rig_cache() as cache:
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=5.0)
counter: Final = key_counter(key)
ids: Final = tuple(_landed_chat(gateway, model, key, index) for index in range(3))
assert len(set(ids)) == 3, ids
rows: Final = _spend_rows(key)
assert len(rows) == 3, rows
assert sorted(str(row["request_id"]) for row in rows) == sorted(ids), rows
total: Final = sum(float(str(row["spend"])) for row in rows)
assert total == pytest.approx(3 * ONE_CALL_COST), rows
eventually(
lambda: counter_value(cache, counter),
lambda value: value is not None and value == pytest.approx(total),
seconds=15,
)
assert outlives(counter_ttl(cache, counter))

26
tests/integration/spend/test_spend_counter_reset_fence.py Normal file
View file

@ -0,0 +1,26 @@
"""The reset fence: a spend reset SETs it in the same pipeline as the zeroed counter, and the stale-floor
repair's Lua script leaves a fenced counter alone, so a pod still holding the pre-reset floor cannot raise
the counter the reset zeroed for the whole counter TTL."""
import os
import uuid
from typing import Final
from integration._support.spend_counters import rig_cache
from litellm.caching.redis_cache import RedisCache
async def test_set_max_holds_while_the_reset_fence_exists() -> None:
cache: Final = RedisCache(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), namespace="fence")
counter: Final = f"spend:key:{uuid.uuid4().hex}"
fence: Final = cache.reset_fence_key(counter)
assert await cache.async_set_cache_pipeline_with_ttls(((fence, 1, 10), (counter, 0.0, 3600))) is None
assert await cache.async_set_max(key=counter, value=12.0, ttl=3600, fence_key=fence) == 0.0
with rig_cache() as raw:
assert raw.delete(cache.check_and_fix_namespace(key=fence)) == 1
assert await cache.async_set_max(key=counter, value=12.0, ttl=3600, fence_key=fence) == 12.0
with rig_cache() as raw:
raw.delete(cache.check_and_fix_namespace(key=counter))

667
tests/integration/spend/test_spend_counter_reset_paths.py Normal file
View file

@ -0,0 +1,667 @@
"""The two spend reset paths that touch the Redis spend counters.
``POST /global/spend/reset`` zeroes every key and team spend in the database. The
counters that enforce those budgets must follow it on every pod, or a key the
operator just reset stays blocked until the counter's TTL runs out. ``PATCH
/organization/update`` with ``spend`` sets one organization's counter to the given
value. Both write a reset fence next to each zeroed counter so a pod still holding
the pre-reset floor cannot repair the counter back up, and both answer 500 when the
counters could not be written to Redis, so the operator retries once it is back.
Every cell drives a real proxy: the cells that reset every key or stop Redis boot
their own proxy on a scratch database and an owned Redis.
"""
import asyncio
import os
import threading
import uuid
from collections.abc import Callable, Iterator, Mapping, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from hashlib import sha256
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway, Scenario, eventually, object_value
from integration._support.database import read_rows, scratch_database
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.redis_process import OwnedRedis, owned_redis
from integration._support.spend_counters import (
SPEND_COUNTER_TTL_SECONDS,
counter_ttl,
counter_value,
key_counter,
outlives,
rig_cache,
)
from pydantic import BaseModel, ConfigDict
from redis import Redis
from litellm.caching.redis_cache import RedisCache
ROOT: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3])
RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}
ONE_CALL_COST: Final = 0.06
INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation"
REDIS_AUTH_CACHE_SETTING: Final = frozenset({"enable_redis_auth_cache"})
REDIS_SETTINGS: Final = frozenset({"cache", "cache_params", *REDIS_AUTH_CACHE_SETTING})
SEEDED_COUNTERS: Final = 1200
SETTLE_SECONDS: Final = 20
BROADCAST_SECONDS: Final = 10
RACE_WINDOW_SECONDS: Final = 2
DATABASE_SECONDS: Final = 70
FENCE_SECONDS: Final = 10
RECOVERY_SECONDS: Final = 30
LOST_COUNTERS: Final = "but the cached spend counters could not be reset in Redis"
RETRY_HINT: Final = "retry this request once Redis accepts writes."
CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120)
class _BroadcastMessage(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
cache_key: str | None = None
cache_keys: tuple[str, ...] = ()
new_value: float | None = None
ttl: float | None = None
def names(self) -> tuple[str, ...]:
return (self.cache_key,) if self.cache_key is not None else self.cache_keys
@dataclass(frozen=True, slots=True)
class Broadcast:
key: str
value: float | None
ttl: float | None
class _Broadcasts:
def __init__(self, cache: Redis) -> None:
self._pubsub: Final = cache.pubsub()
self._pubsub.subscribe(INVALIDATION_CHANNEL)
confirmation: Final = self._pubsub.get_message(timeout=5)
assert confirmation is not None and confirmation["type"] == "subscribe", confirmation
self._condition: Final = threading.Condition()
self._received: Final[list[Broadcast]] = []
self._stopped: Final = threading.Event()
self._thread: Final = threading.Thread(target=self._pump, daemon=True)
self._thread.start()
def _pump(self) -> None:
while not self._stopped.is_set():
self._record(self._pubsub.get_message(timeout=0.2))
def _record(self, message: Mapping[str, object] | None) -> None:
if message is None or message["type"] != "message":
return
data: Final = message["data"]
assert isinstance(data, bytes | str), message
parsed: Final = _BroadcastMessage.model_validate_json(data)
with self._condition:
self._received.extend(Broadcast(key, parsed.new_value, parsed.ttl) for key in parsed.names())
self._condition.notify_all()
def wait_for(self, satisfied: Callable[[Sequence[Broadcast]], bool], seconds: float) -> bool:
with self._condition:
return self._condition.wait_for(lambda: satisfied(tuple(self._received)), timeout=seconds)
def received(self) -> tuple[Broadcast, ...]:
with self._condition:
return tuple(self._received)
def close(self) -> None:
self._stopped.set()
self._thread.join(timeout=5)
self._pubsub.close()
@contextmanager
def _broadcasts(cache: Redis) -> Iterator[_Broadcasts]:
broadcasts: Final = _Broadcasts(cache)
try:
yield broadcasts
finally:
broadcasts.close()
def _zeroed(counter: str) -> Callable[[Sequence[Broadcast]], bool]:
return lambda received: any(broadcast.key == counter and broadcast.value == 0.0 for broadcast in received)
def _zeroed_with_its_floor(counter: str) -> Callable[[Sequence[Broadcast]], bool]:
return lambda received: (
sum(1 for broadcast in received if broadcast.key.endswith(counter) and broadcast.value == 0.0) >= 2
)
@dataclass(frozen=True, slots=True)
class Isolated:
gateway: Gateway
cache: Redis
redis: OwnedRedis
database_url: str
log: Path
def _pod_environment(database_url: str, redis: OwnedRedis) -> Mapping[str, str]:
return {
"DATABASE_URL": database_url,
"REDIS_HOST": redis.host,
"REDIS_PORT": str(redis.port),
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "2",
}
@contextmanager
def _isolated(
gateway: Gateway,
directory: Path,
*,
workers: int = 1,
settings: tuple[tuple[str, str], ...] = (),
config: Path | None = None,
) -> Iterator[Isolated]:
with scratch_database() as database_url, owned_redis(directory) as redis:
overrides: Final[Mapping[str, str]] = {**_pod_environment(database_url, redis), **dict(settings)}
with (
owned_proxy_process(
gateway,
directory,
overrides,
workers=workers,
config=config,
remove_environment=("DATABASE_URL_READ_REPLICA",),
) as owned,
Redis(host=redis.host, port=redis.port) as cache,
):
yield Isolated(owned.gateway, cache, redis, database_url, owned.log)
@contextmanager
def _second_pod(gateway: Gateway, directory: Path, isolated: Isolated, config: Path) -> Iterator[Gateway]:
with owned_proxy_process(
gateway,
directory,
_pod_environment(isolated.database_url, isolated.redis),
config=config,
database_setup=(),
remove_environment=("DATABASE_URL_READ_REPLICA",),
) as owned:
yield owned.gateway
def _config_without(directory: Path, dropped: frozenset[str]) -> Path:
shared: Final = object_value(yaml.safe_load((ROOT / "tests/integration/proxy_config.yaml").read_text()))
litellm_settings: Final = {
name: value for name, value in object_value(shared["litellm_settings"]).items() if name not in dropped
}
path: Final = directory / f"proxy_config_{uuid.uuid4().hex}.yaml"
path.write_text(yaml.safe_dump({**shared, "litellm_settings": litellm_settings}, sort_keys=False))
return path
@contextmanager
def _isolated_without_redis(gateway: Gateway, directory: Path) -> Iterator[tuple[Gateway, str]]:
with (
scratch_database() as database_url,
owned_proxy_process(
gateway,
directory,
{"DATABASE_URL": database_url},
config=_config_without(directory, REDIS_SETTINGS),
remove_environment=("REDIS_HOST", "REDIS_PORT", "DATABASE_URL_READ_REPLICA"),
) as owned,
):
yield owned.gateway, database_url
def _one_spend(rows: Sequence[Mapping[str, object]]) -> float:
assert len(rows) == 1, rows
return float(str(rows[0]["spend"]))
def _key_spend(key: str, database_url: str | None = None) -> float:
return _one_spend(
read_rows(
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s',
(sha256(key.encode()).hexdigest(),),
database_url=database_url,
)
)
def _team_spend(team: str, database_url: str | None = None) -> float:
return _one_spend(
read_rows('SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id=%s', (team,), database_url=database_url)
)
def _user_spend(user: str, database_url: str | None = None) -> float:
return _one_spend(
read_rows('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=%s', (user,), database_url=database_url)
)
def _organization_spend(organization: str, database_url: str | None = None) -> float:
return _one_spend(
read_rows(
'SELECT spend FROM "LiteLLM_OrganizationTable" WHERE organization_id=%s',
(organization,),
database_url=database_url,
)
)
def _landed(read: Callable[[], float], least: float) -> float:
return eventually(read, lambda spend: spend >= least - 1e-9, seconds=DATABASE_SECONDS)
def _chat(client: httpx.Client, model: str, key: str, *, end_user: str | None = None) -> httpx.Response:
return client.post(
"/v1/chat/completions",
json={
"model": model,
"messages": [{"role": "user", "content": f"reset path {uuid.uuid4().hex}"}],
**({} if end_user is None else {"user": end_user}),
},
headers={"Authorization": f"Bearer {key}"},
)
def _assert_served(response: httpx.Response) -> str:
assert response.status_code == 200, response.text
body: Final = response.json()
assert body["usage"]["total_tokens"] == 40, body
return str(body["id"])
def _assert_denied(response: httpx.Response) -> None:
assert response.status_code == 422, response.text
assert response.json()["error"]["type"] == "budget_exceeded", response.text
def _assert_counters_unwritten(response: httpx.Response, saved: str) -> None:
assert response.status_code == 500, response.text
error: Final = response.json()["error"]
assert error["type"] == "internal_server_error", response.text
message: Final = str(error["message"])
assert message.startswith(f"{saved}, {LOST_COUNTERS}"), message
assert message.endswith(RETRY_HINT), message
def _accepted_once_redis_is_back(request: Callable[[], httpx.Response]) -> None:
def attempt() -> int:
response: Final = request()
assert response.status_code in {200, 500}, response.text
return response.status_code
eventually(attempt, lambda code: code == 200, seconds=RECOVERY_SECONDS)
def _fence(counter: str) -> str:
return f"{{{counter}}}:reset_fence"
async def _repair_as_a_pod_would(redis: OwnedRedis, counter: str, value: float) -> float | None:
cache: Final = RedisCache(host=redis.host, port=redis.port)
assert cache.reset_fence_key(counter) == _fence(counter)
return await cache.async_set_max(
key=counter, value=value, ttl=SPEND_COUNTER_TTL_SECONDS, fence_key=cache.reset_fence_key(counter)
)
def _state(cache: Redis, counter: str) -> tuple[float | None, int]:
return counter_value(cache, counter), counter_ttl(cache, counter)
def _settled(cache: Redis, counter: str, spend: float) -> int:
_, ttl = eventually(
lambda: _state(cache, counter),
lambda state: state[0] is not None and state[0] == pytest.approx(spend, abs=1e-9),
seconds=SETTLE_SECONDS,
)
return ttl
def _values(cache: Redis, counters: Sequence[str]) -> tuple[float | None, ...]:
return tuple(None if raw is None else float(raw) for raw in cache.mget(counters))
def _ttls(cache: Redis, counters: Sequence[str]) -> tuple[int, ...]:
with cache.pipeline() as pipeline:
for counter in counters:
pipeline.ttl(counter)
return tuple(int(ttl) for ttl in pipeline.execute())
def _exhaust(gateway: Gateway, cache: Redis, database_url: str, model: str, key: str) -> str:
counter: Final = key_counter(key)
_assert_served(_chat(gateway.client, model, key))
_settled(cache, counter, ONE_CALL_COST)
_landed(lambda: _key_spend(key, database_url), ONE_CALL_COST)
_assert_denied(_chat(gateway.client, model, key))
return counter
def _global_reset(gateway: Gateway) -> None:
reset: Final = gateway.request("POST", "/global/spend/reset")
assert reset.status_code == 200, reset.text
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_global_spend_reset_admits_a_blocked_key_on_the_next_request(gateway: Gateway, tmp_path: Path) -> None:
with (
_isolated(gateway, tmp_path, workers=2) as isolated,
isolated.gateway.scenario() as scenario,
httpx.Client(base_url=str(isolated.gateway.client.base_url), timeout=60) as other_connection,
):
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST)
counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key)
_global_reset(isolated.gateway)
assert outlives(_settled(isolated.cache, counter, 0.0))
assert _key_spend(key, isolated.database_url) == 0.0
_assert_served(_chat(other_connection, model, key))
assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST))
_assert_denied(_chat(other_connection, model, key))
_landed(lambda: _key_spend(key, isolated.database_url), ONE_CALL_COST)
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_global_spend_reset_admits_a_key_blocked_by_its_team_budget(gateway: Gateway, tmp_path: Path) -> None:
with _isolated(gateway, tmp_path) as isolated, isolated.gateway.scenario() as scenario:
model: Final = scenario.model(**RATES)
team: Final = scenario.team(models=[model], max_budget=ONE_CALL_COST)
key: Final = scenario.key(team_id=team, models=[model])
counter: Final = f"spend:team:{team}"
_assert_served(_chat(isolated.gateway.client, model, key))
_settled(isolated.cache, counter, ONE_CALL_COST)
_landed(lambda: _team_spend(team, isolated.database_url), ONE_CALL_COST)
_assert_denied(_chat(isolated.gateway.client, model, key))
_global_reset(isolated.gateway)
assert outlives(_settled(isolated.cache, counter, 0.0))
assert _team_spend(team, isolated.database_url) == 0.0
_assert_served(_chat(isolated.gateway.client, model, key))
assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST))
_assert_denied(_chat(isolated.gateway.client, model, key))
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_global_spend_reset_leaves_the_other_scopes_and_a_non_master_caller_alone(
gateway: Gateway, tmp_path: Path
) -> None:
with _isolated(gateway, tmp_path) as isolated, isolated.gateway.scenario() as scenario:
model: Final = scenario.model(**RATES)
organization: Final = scenario.organization(max_budget=50.0)
team: Final = scenario.team(models=[model], max_budget=20.0, organization_id=organization)
user: Final = scenario.user(max_budget=20.0)
added: Final = isolated.gateway.request(
"POST",
"/team/member_add",
{"team_id": team, "member": {"user_id": user, "role": "user"}, "max_budget_in_team": 20.0},
)
assert added.status_code == 200, added.text
tag: Final = f"integration-tag-{uuid.uuid4().hex}"
key: Final = scenario.key(team_id=team, user_id=user, models=[model], max_budget=5.0, tags=[tag])
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
reset_counters: Final = (key_counter(key), f"spend:team:{team}")
kept_counters: Final = (
f"spend:team_member:{user}:{team}",
f"spend:user:{user}",
f"spend:end_user:{end_user}",
f"spend:tag:{tag}",
f"spend:org:{organization}",
)
_assert_served(_chat(isolated.gateway.client, model, key, end_user=end_user))
for counter in (*reset_counters, *kept_counters):
assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)), counter
_landed(lambda: _user_spend(user, isolated.database_url), ONE_CALL_COST)
_landed(lambda: _organization_spend(organization, isolated.database_url), ONE_CALL_COST)
forbidden: Final = isolated.gateway.request("POST", "/global/spend/reset", key=key)
assert forbidden.status_code == 401, forbidden.text
assert forbidden.json()["error"]["type"] == "auth_error", forbidden.text
assert _values(isolated.cache, (*reset_counters, *kept_counters)) == pytest.approx(
(ONE_CALL_COST,) * 7, abs=1e-9
)
_global_reset(isolated.gateway)
for counter in reset_counters:
assert outlives(_settled(isolated.cache, counter, 0.0)), counter
assert _values(isolated.cache, kept_counters) == pytest.approx((ONE_CALL_COST,) * 5, abs=1e-9)
assert all(outlives(ttl) for ttl in _ttls(isolated.cache, kept_counters)), _ttls(isolated.cache, kept_counters)
assert _key_spend(key, isolated.database_url) == 0.0
assert _team_spend(team, isolated.database_url) == 0.0
assert _user_spend(user, isolated.database_url) >= ONE_CALL_COST - 1e-9
assert _organization_spend(organization, isolated.database_url) >= ONE_CALL_COST - 1e-9
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_global_spend_reset_zeroes_every_key_counter_past_one_scan_page_and_broadcasts_them(
gateway: Gateway, tmp_path: Path
) -> None:
with _isolated(gateway, tmp_path) as isolated:
stem: Final = f"spend:key:audit-{uuid.uuid4().hex}"
team_named_like_a_window: Final = f"spend:team:audit-{uuid.uuid4().hex}:window:prod"
counters: Final = (*(f"{stem}-{index}" for index in range(SEEDED_COUNTERS)), team_named_like_a_window)
windows: Final = tuple(f"{stem}-w{index}:window:1d" for index in range(3))
with isolated.cache.pipeline() as pipeline:
for name in (*counters, *windows):
pipeline.set(name, "1.0", ex=3600)
pipeline.execute()
with _broadcasts(isolated.cache) as broadcasts:
_global_reset(isolated.gateway)
eventually(
lambda: _values(isolated.cache, counters),
lambda values: all(value == 0.0 for value in values),
seconds=SETTLE_SECONDS,
)
ttls: Final = _ttls(isolated.cache, counters)
assert all(outlives(ttl) for ttl in ttls), sorted(ttls)[:5]
assert _values(isolated.cache, windows) == (1.0, 1.0, 1.0)
assert broadcasts.wait_for(
lambda received: {broadcast.key for broadcast in received if broadcast.value == 0.0} >= set(counters),
BROADCAST_SECONDS,
), len(broadcasts.received())
assert not any(broadcast.key in windows for broadcast in broadcasts.received())
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_global_spend_reset_with_a_redis_batch_size_of_one_resets_and_fences_every_counter(
gateway: Gateway, tmp_path: Path
) -> None:
with (
_isolated(gateway, tmp_path, settings=(("DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE", "1"),)) as isolated,
isolated.gateway.scenario() as scenario,
):
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST)
counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key)
seeded: Final = tuple(f"spend:key:audit-{uuid.uuid4().hex}" for _ in range(2))
with isolated.cache.pipeline() as pipeline:
for name in seeded:
pipeline.set(name, "1.0", ex=3600)
pipeline.execute()
counters: Final = (counter, *seeded)
_global_reset(isolated.gateway)
assert isolated.cache.exists(*(_fence(name) for name in counters)) == len(counters)
assert _values(isolated.cache, counters) == (0.0, 0.0, 0.0)
assert all(outlives(ttl) for ttl in _ttls(isolated.cache, counters)), _ttls(isolated.cache, counters)
_assert_served(_chat(isolated.gateway.client, model, key))
assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST))
_assert_denied(_chat(isolated.gateway.client, model, key))
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_pod_that_cached_the_pre_reset_floor_admits_the_next_request_after_a_reset_on_another_pod(
gateway: Gateway, tmp_path: Path
) -> None:
with (
_isolated(gateway, tmp_path, config=_config_without(tmp_path, REDIS_AUTH_CACHE_SETTING)) as isolated,
_second_pod(gateway, tmp_path, isolated, _config_without(tmp_path, REDIS_AUTH_CACHE_SETTING)) as other_pod,
isolated.gateway.scenario() as scenario,
):
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST)
counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key)
with _broadcasts(isolated.cache) as broadcasts:
assert isolated.cache.set(counter, "0.0", ex=3600)
_assert_denied(_chat(other_pod.client, model, key))
_settled(isolated.cache, counter, ONE_CALL_COST)
_global_reset(isolated.gateway)
broadcasts.wait_for(_zeroed_with_its_floor(counter), RACE_WINDOW_SECONDS)
assert other_pod.client.get("/health/liveliness").status_code == 200
_assert_served(_chat(other_pod.client, model, key))
assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST))
_assert_denied(_chat(other_pod.client, model, key))
assert broadcasts.wait_for(_zeroed(counter), BROADCAST_SECONDS), broadcasts.received()
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_global_spend_reset_fences_the_counter_against_a_stale_floor_repair(gateway: Gateway, tmp_path: Path) -> None:
with _isolated(gateway, tmp_path) as isolated, isolated.gateway.scenario() as scenario:
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST)
counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key)
_global_reset(isolated.gateway)
assert outlives(_settled(isolated.cache, counter, 0.0))
assert isolated.cache.get(_fence(counter)) == b"1"
assert 0 < counter_ttl(isolated.cache, _fence(counter)) <= FENCE_SECONDS
assert asyncio.run(_repair_as_a_pod_would(isolated.redis, counter, ONE_CALL_COST)) == 0.0
value, ttl = _state(isolated.cache, counter)
assert value == 0.0 and outlives(ttl), (value, ttl)
_assert_served(_chat(isolated.gateway.client, model, key))
assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST))
_assert_denied(_chat(isolated.gateway.client, model, key))
eventually(lambda: isolated.cache.exists(_fence(counter)), lambda found: found == 0, seconds=FENCE_SECONDS + 5)
late_total: Final = 2 * ONE_CALL_COST
assert asyncio.run(_repair_as_a_pod_would(isolated.redis, counter, late_total)) == pytest.approx(late_total)
value_after, ttl_after = _state(isolated.cache, counter)
assert value_after == pytest.approx(late_total) and outlives(ttl_after), (value_after, ttl_after)
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_global_spend_reset_while_redis_is_down_answers_500_and_the_retry_resets_the_counters(
gateway: Gateway, tmp_path: Path
) -> None:
with _isolated(gateway, tmp_path) as isolated, isolated.gateway.scenario() as scenario:
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST)
counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key)
isolated.redis.stop()
_assert_counters_unwritten(isolated.gateway.request("POST", "/global/spend/reset"), "Spend was reset")
assert _key_spend(key, isolated.database_url) == 0.0
_assert_served(_chat(isolated.gateway.client, model, key))
isolated.redis.start()
assert isolated.cache.ping()
_landed(lambda: _key_spend(key, isolated.database_url), ONE_CALL_COST)
_accepted_once_redis_is_back(lambda: isolated.gateway.request("POST", "/global/spend/reset"))
assert _key_spend(key, isolated.database_url) == 0.0
assert outlives(_settled(isolated.cache, counter, 0.0))
_assert_served(_chat(isolated.gateway.client, model, key))
assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST))
_assert_denied(_chat(isolated.gateway.client, model, key))
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_global_spend_reset_without_redis_admits_the_next_request(gateway: Gateway, tmp_path: Path) -> None:
with _isolated_without_redis(gateway, tmp_path) as (owned, database_url), owned.scenario() as scenario:
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST)
_assert_served(_chat(owned.client, model, key))
_landed(lambda: _key_spend(key, database_url), ONE_CALL_COST)
_assert_denied(_chat(owned.client, model, key))
_global_reset(owned)
assert _key_spend(key, database_url) == 0.0
_assert_served(_chat(owned.client, model, key))
_landed(lambda: _key_spend(key, database_url), ONE_CALL_COST)
_assert_denied(_chat(owned.client, model, key))
def _organization_entities(
gateway: Gateway, scenario: Scenario, database_url: str | None = None
) -> tuple[str, str, str]:
model: Final = scenario.model(**RATES)
organization: Final = scenario.organization(max_budget=ONE_CALL_COST)
team: Final = scenario.team(models=[model], organization_id=organization)
key: Final = scenario.key(team_id=team, models=[model])
counter: Final = f"spend:org:{organization}"
_assert_served(_chat(gateway.client, model, key))
_landed(lambda: _organization_spend(organization, database_url), ONE_CALL_COST)
_assert_denied(_chat(gateway.client, model, key))
return model, key, counter
def _organization_update(gateway: Gateway, counter: str, **fields: object) -> httpx.Response:
return gateway.request(
"PATCH", "/organization/update", {"organization_id": counter.removeprefix("spend:org:"), **fields}
)
def _update_organization(gateway: Gateway, counter: str, **fields: object) -> None:
updated: Final = _organization_update(gateway, counter, **fields)
assert updated.status_code == 200, updated.text
def test_an_organization_spend_set_to_zero_admits_a_key_blocked_by_the_organization_budget(gateway: Gateway) -> None:
with gateway.scenario() as scenario, rig_cache() as cache:
model, key, counter = _organization_entities(gateway, scenario)
_settled(cache, counter, ONE_CALL_COST)
with _broadcasts(cache) as broadcasts:
_update_organization(gateway, counter, spend=0.0)
assert outlives(_settled(cache, counter, 0.0))
assert _organization_spend(counter.removeprefix("spend:org:")) == 0.0
assert broadcasts.wait_for(_zeroed(counter), BROADCAST_SECONDS), broadcasts.received()
_assert_served(_chat(gateway.client, model, key))
assert outlives(_settled(cache, counter, ONE_CALL_COST))
_assert_denied(_chat(gateway.client, model, key))
def test_an_organization_update_without_spend_keeps_the_counter_and_the_block(gateway: Gateway) -> None:
with gateway.scenario() as scenario, rig_cache() as cache:
model, key, counter = _organization_entities(gateway, scenario)
_settled(cache, counter, ONE_CALL_COST)
_update_organization(gateway, counter, metadata={"note": f"audit {uuid.uuid4().hex}"})
value, ttl = _state(cache, counter)
assert value == pytest.approx(ONE_CALL_COST, abs=1e-9) and outlives(ttl), (value, ttl)
_assert_denied(_chat(gateway.client, model, key))
def test_an_organization_spend_set_above_the_budget_keeps_the_block_until_set_to_zero(gateway: Gateway) -> None:
with gateway.scenario() as scenario, rig_cache() as cache:
model, key, counter = _organization_entities(gateway, scenario)
_settled(cache, counter, ONE_CALL_COST)
_update_organization(gateway, counter, spend=1.5)
assert outlives(_settled(cache, counter, 1.5))
_assert_denied(_chat(gateway.client, model, key))
_update_organization(gateway, counter, spend=0.0)
assert outlives(_settled(cache, counter, 0.0))
_assert_served(_chat(gateway.client, model, key))
@pytest.mark.timeout(CELL_TIMEOUT)
def test_an_organization_spend_set_while_redis_is_down_answers_500_and_the_retry_resets_the_counter(
gateway: Gateway, tmp_path: Path
) -> None:
with _isolated(gateway, tmp_path) as isolated, isolated.gateway.scenario() as scenario:
model, key, counter = _organization_entities(isolated.gateway, scenario, isolated.database_url)
organization: Final = counter.removeprefix("spend:org:")
_settled(isolated.cache, counter, ONE_CALL_COST)
isolated.redis.stop()
_assert_counters_unwritten(
_organization_update(isolated.gateway, counter, spend=0.0), "The organization spend was saved"
)
assert _organization_spend(organization, isolated.database_url) == 0.0
_assert_served(_chat(isolated.gateway.client, model, key))
isolated.redis.start()
assert isolated.cache.ping()
_accepted_once_redis_is_back(lambda: _organization_update(isolated.gateway, counter, spend=0.0))
assert _organization_spend(organization, isolated.database_url) == 0.0
assert outlives(_settled(isolated.cache, counter, 0.0))
_assert_served(_chat(isolated.gateway.client, model, key))
assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST))
_assert_denied(_chat(isolated.gateway.client, model, key))

292
tests/integration/spend/test_spend_counter_ttl_chaos.py Normal file
View file

@ -0,0 +1,292 @@
import uuid
from collections.abc import Callable, Mapping
from concurrent.futures import ThreadPoolExecutor
from hashlib import sha256
from pathlib import Path
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
from integration._support.redis_process import owned_redis
from integration._support.spend_counters import (
counter_ttl,
counter_value,
key_counter,
outlives,
rig_cache,
)
from pydantic import JsonValue
from redis import Redis
pytestmark: Final = pytest.mark.shares_worker_state
RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}
ANTHROPIC_HEADERS: Final = {"anthropic-version": "2023-06-01"}
BURST: Final = 24
UNBOUNDED_BUDGET: Final = 100000.0
RESERVATION_ESTIMATE_CEILING: Final = 40.0
SLOW_UPSTREAM_SECONDS: Final = 3.0
LEAK_KEY_BUDGET: Final = 20.0
CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120)
RESPAWN_SECONDS: Final = graceful_stop_seconds() + 60
def _chat(gateway: Gateway, model: str, key: str, index: int, stream: bool) -> httpx.Response:
return gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": f"burst chat {index} {uuid.uuid4().hex}"}],
"stream": stream,
**({"stream_options": {"include_usage": True}} if stream else {}),
},
key=key,
)
def _messages(gateway: Gateway, model: str, key: str, index: int, stream: bool) -> httpx.Response:
return gateway.request(
"POST",
"/v1/messages",
{
"model": model,
"max_tokens": 32,
"messages": [{"role": "user", "content": f"burst messages {index} {uuid.uuid4().hex}"}],
"stream": stream,
},
key=key,
headers=ANTHROPIC_HEADERS,
)
def _responses(gateway: Gateway, model: str, key: str, index: int, stream: bool) -> httpx.Response:
return gateway.request(
"POST",
"/v1/responses",
{"model": model, "input": f"burst responses {index} {uuid.uuid4().hex}", "stream": stream},
key=key,
)
SURFACES: Final[tuple[Callable[[Gateway, str, str, int, bool], httpx.Response], ...]] = (
_chat,
_messages,
_responses,
)
def _spend_rows(key: str) -> list[dict[str, JsonValue]]:
return read_rows(
'SELECT request_id, spend, status FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
(sha256(key.encode()).hexdigest(),),
)
def _successes(key: str, count: int) -> list[dict[str, JsonValue]]:
rows: Final = eventually(
lambda: [row for row in _spend_rows(key) if str(row["status"]) == "success"],
lambda values: len(values) >= count,
seconds=90,
)
assert len(rows) == count, rows
assert len({str(row["request_id"]) for row in rows}) == count, rows
return rows
def _counter_matches(cache: Redis, counter: str, total: float) -> tuple[float | None, int]:
state: Final = eventually(
lambda: (counter_value(cache, counter), counter_ttl(cache, counter)),
lambda values: values[0] is not None and values[0] == pytest.approx(total, abs=1e-9),
seconds=20,
)
assert outlives(state[1]), state
return state
def test_a_burst_with_upstream_failures_lands_each_success_once_with_a_long_ttl(gateway: Gateway) -> None:
with (
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
gateway.scenario() as scenario,
rig_cache() as cache,
):
provider_model: Final = f"burst-{uuid.uuid4().hex}"
model: Final = scenario.model(model=f"openai/{provider_model}", **RATES)
key: Final = scenario.key(models=[model], max_budget=UNBOUNDED_BUDGET)
statuses: Final = [500] + [200, 200, 200] * (BURST // 4 + 2)
def remove_script() -> None:
response: Final = upstream.delete(f"/__scripts/{provider_model}")
assert response.status_code in (200, 404), response.text
scenario.cleanups.callback(remove_script)
configured: Final = upstream.post(f"/__scripts/{provider_model}", json={"statuses": statuses})
assert configured.status_code == 200, configured.text
def attempt(index: int) -> httpx.Response:
stream: Final = index % 2 == 1
return SURFACES[index % len(SURFACES)](gateway, model, key, index, stream)
with ThreadPoolExecutor(max_workers=BURST) as pool:
responses: Final = tuple(pool.map(attempt, range(BURST)))
codes: Final = [response.status_code for response in responses]
assert all(code in (200, 500) for code in codes), codes
succeeded: Final = codes.count(200)
assert 0 < succeeded < BURST, codes
rows: Final = _successes(key, succeeded)
total: Final = sum(float(str(row["spend"])) for row in rows)
assert total == pytest.approx(succeeded * 0.06), rows
_counter_matches(cache, key_counter(key), total)
def _workers(owned: OwnedProxy) -> tuple[int, ...]:
return tuple(sorted(owned.worker_pids()))
def _attempt(gateway: Gateway, model: str, key: str, index: int) -> int:
try:
return _chat(gateway, model, key, index, False).status_code
except httpx.HTTPError:
return 0
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_killed_worker_does_not_shorten_or_lose_the_counter(gateway: Gateway, tmp_path: Path) -> None:
with (
owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned,
owned.gateway.scenario() as scenario,
rig_cache() as cache,
):
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=UNBOUNDED_BUDGET)
before: Final = _workers(owned)
assert len(before) == 2, before
with ThreadPoolExecutor(max_workers=12) as pool:
first: Final = pool.submit(_attempt, owned.gateway, model, key, 0)
assert first.result(timeout=60) == 200
in_flight: Final = [pool.submit(_attempt, owned.gateway, model, key, index) for index in range(1, 12)]
with owned.replacing_worker(before[0]) as kill:
kill()
survivor: Final = _attempt(owned.gateway, model, key, 99)
codes: Final = [future.result(timeout=120) for future in in_flight]
assert survivor == 200, survivor
assert all(code in (0, 200) for code in codes), codes
replaced: Final = eventually(
lambda: _workers(owned),
lambda pids: len(pids) == 2 and before[0] not in pids and before[1] in pids,
seconds=RESPAWN_SECONDS,
)
assert replaced
after: Final = [_attempt(owned.gateway, model, key, index) for index in range(12, 18)]
assert after == [200] * 6, after
answered: Final = 1 + 1 + codes.count(200) + 6
landed: Final = eventually(
lambda: [row for row in _spend_rows(key) if str(row["status"]) == "success"],
lambda rows: len(rows) >= 1 + 1 + 6,
seconds=90,
)
assert len({str(row["request_id"]) for row in landed}) == len(landed) <= answered, landed
ceiling: Final = answered * 0.06 + len(in_flight) * RESERVATION_ESTIMATE_CEILING
state: Final = eventually(
lambda: (counter_value(cache, key_counter(key)), counter_ttl(cache, key_counter(key))),
lambda values: values[0] is not None and 8 * 0.06 - 1e-9 <= values[0] <= ceiling,
seconds=20,
)
assert outlives(state[1]), state
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_reservation_leaked_by_a_killed_worker_stays_for_the_counter_ttl(gateway: Gateway, tmp_path: Path) -> None:
with (
owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned,
owned.gateway.scenario() as scenario,
rig_cache() as cache,
):
slow: Final = scenario.model(model="openai/slow-endpoint", **RATES)
key: Final = scenario.key(models=[slow], max_budget=LEAK_KEY_BUDGET)
counter: Final = key_counter(key)
before: Final = _workers(owned)
assert len(before) == 2, before
with ThreadPoolExecutor(max_workers=1) as pool:
in_flight: Final = pool.submit(_attempt, owned.gateway, slow, key, 0)
reserved: Final = eventually(
lambda: counter_value(cache, counter),
lambda value: value is not None and value >= LEAK_KEY_BUDGET - 1e-9,
seconds=SLOW_UPSTREAM_SECONDS,
)
with owned.replacing_worker(before[0]) as first, owned.replacing_worker(before[1]) as second:
first.freeze()
second.freeze()
first()
second()
interrupted: Final = in_flight.result(timeout=60)
assert interrupted != 200, interrupted
leaked: Final = (counter_value(cache, counter), counter_ttl(cache, counter))
assert leaked[0] == pytest.approx(reserved), (reserved, leaked)
respawned: Final = eventually(
lambda: _workers(owned),
lambda pids: len(pids) == 2 and not set(pids) & set(before),
seconds=RESPAWN_SECONDS,
)
assert respawned
blocked: Final = eventually(
lambda: _attempt(owned.gateway, slow, key, 1), lambda code: code == 422, seconds=RESPAWN_SECONDS
)
assert blocked == 422, (blocked, leaked)
held: Final = (counter_value(cache, counter), counter_ttl(cache, counter))
assert held[0] == pytest.approx(reserved), (reserved, held)
assert outlives(held[1]), held
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_redis_outage_mid_burst_keeps_serving_and_rebuilds_the_counter_long(gateway: Gateway, tmp_path: Path) -> None:
with owned_redis(tmp_path) as redis:
overrides: Final[Mapping[str, str]] = {
"REDIS_HOST": redis.host,
"REDIS_PORT": str(redis.port),
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "2",
}
with (
owned_proxy_process(gateway, tmp_path, overrides) as owned,
owned.gateway.scenario() as scenario,
Redis(host=redis.host, port=redis.port) as cache,
):
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=UNBOUNDED_BUDGET)
counter: Final = key_counter(key)
def burst(start: int) -> list[int]:
with ThreadPoolExecutor(max_workers=6) as pool:
return list(
pool.map(lambda index: _attempt(owned.gateway, model, key, index), range(start, start + 6))
)
assert burst(0) == [200] * 6
_counter_matches(cache, counter, 6 * 0.06)
redis.stop()
during: Final = burst(6)
assert during == [200] * 6, during
redis.start()
assert cache.ping()
probes: Final[list[int]] = []
def probe_until_the_counter_is_back() -> float | None:
probes.append(_attempt(owned.gateway, model, key, 100 + len(probes)))
return counter_value(cache, counter)
eventually(probe_until_the_counter_is_back, lambda value: value is not None, seconds=30)
assert set(probes) == {200}, probes
after: Final = burst(12)
assert after == [200] * 6, after
rows: Final = _successes(key, 18 + len(probes))
total: Final = sum(float(str(row["spend"])) for row in rows)
assert total == pytest.approx((18 + len(probes)) * 0.06), rows
rebuilt: Final = eventually(
lambda: (counter_value(cache, counter), counter_ttl(cache, counter)),
lambda state: state[0] is not None and 7 * 0.06 - 1e-9 <= state[0] <= total + 1e-9,
seconds=20,
)
assert outlives(rebuilt[1]), rebuilt

230
tests/integration/spend/test_spend_counter_ttl_settings.py Normal file
View file

@ -0,0 +1,230 @@
import json
import os
import time
import uuid
from collections.abc import Mapping
from concurrent.futures import ThreadPoolExecutor
from datetime import UTC, datetime, timedelta
from hashlib import sha256
from pathlib import Path
from types import MappingProxyType
from typing import Final
import psycopg
import pytest
import yaml
from integration._support.client import Gateway, eventually, object_value
from integration._support.database import read_rows
from integration._support.process import graceful_stop_seconds, owned_proxy
from integration._support.spend_counters import counter_ttl, counter_value, key_counter, outlives, rig_cache
from pydantic import JsonValue
pytestmark: Final = pytest.mark.shares_worker_state
ROOT: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3])
RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}
ONE_CALL_COST: Final = 0.06
SLOW_UPSTREAM_SECONDS: Final = 3.0
SHORT_REDIS_TTL: Final = 4
LONG_REDIS_TTL: Final = 86400
ENV_REDIS_TTL: Final = "7200"
LONG_BATCH_WRITE_SECONDS: Final = 900
STRICT_KEY_BUDGET: Final = 50.0
CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120)
FAST_RESET_TICK: Final = MappingProxyType(
{"PROXY_BUDGET_RESCHEDULER_MIN_TIME": "2", "PROXY_BUDGET_RESCHEDULER_MAX_TIME": "2"}
)
def _config(
directory: Path,
*,
litellm_settings: Mapping[str, JsonValue] = MappingProxyType({}),
general_settings: Mapping[str, JsonValue] = MappingProxyType({}),
) -> Path:
shared: Final = object_value(yaml.safe_load((ROOT / "tests/integration/proxy_config.yaml").read_text()))
configured: Final = {
**shared,
"litellm_settings": {**object_value(shared["litellm_settings"]), **litellm_settings},
"general_settings": {**object_value(shared["general_settings"]), **general_settings},
}
path: Final = directory / f"proxy_config_{uuid.uuid4().hex}.yaml"
path.write_text(yaml.safe_dump(configured, sort_keys=False))
return path
def _chat(gateway: Gateway, model: str, key: str) -> str:
served: Final = gateway.chat(model, key=key, text=f"settings {uuid.uuid4().hex}")
assert served["usage"]["total_tokens"] == 40, served
return str(served["id"])
def _counter_state(counter: str) -> tuple[float | None, int]:
with rig_cache() as cache:
return counter_value(cache, counter), counter_ttl(cache, counter)
def _settled(counter: str, spend: float, expected_ttl: int) -> tuple[float | None, int]:
return eventually(
lambda: _counter_state(counter),
lambda state: state[0] is not None and state[0] == pytest.approx(spend) and outlives(state[1], expected_ttl),
seconds=15,
)
def _spend_rows(key: str) -> list[dict[str, JsonValue]]:
return read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),))
@pytest.mark.timeout(CELL_TIMEOUT)
def test_null_default_redis_ttl_falls_back_to_the_hour_floor(gateway: Gateway, tmp_path: Path) -> None:
config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": None})
with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario:
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=5.0)
assert _chat(owned, model, key)
assert _settled(key_counter(key), ONE_CALL_COST, 3600)
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_short_default_redis_ttl_never_shortens_the_counter_and_the_lease_refresh_keeps_it_long(
gateway: Gateway, tmp_path: Path
) -> None:
config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": SHORT_REDIS_TTL})
with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario, rig_cache() as cache:
model: Final = scenario.model(model="openai/slow-endpoint", **RATES)
key: Final = scenario.key(models=[model], max_budget=5.0)
counter: Final = key_counter(key)
with ThreadPoolExecutor(max_workers=1) as pool:
slow: Final = pool.submit(lambda: (_chat(owned, model, key), time.monotonic())[1])
eventually(lambda: cache.exists(counter), lambda present: present == 1, seconds=SLOW_UPSTREAM_SECONDS)
assert cache.expire(counter, 1000)
refreshed_at: Final = eventually(
lambda: (time.monotonic(), counter_ttl(cache, counter)),
lambda state: outlives(state[1]),
seconds=SLOW_UPSTREAM_SECONDS + 10,
)[0]
finished_at: Final = slow.result(timeout=60)
assert refreshed_at < finished_at, (refreshed_at, finished_at)
assert _settled(counter, ONE_CALL_COST, 3600)
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_long_default_redis_ttl_wins_on_every_write_path(gateway: Gateway, tmp_path: Path) -> None:
config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": LONG_REDIS_TTL})
with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario:
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=5.0)
counter: Final = key_counter(key)
assert _chat(owned, model, key)
assert _settled(counter, ONE_CALL_COST, LONG_REDIS_TTL)
updated: Final = owned.request("POST", "/key/update", {"key": key, "spend": 2.5})
assert updated.status_code == 200, updated.text
assert _settled(counter, 2.5, LONG_REDIS_TTL)
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_default_redis_ttl_given_as_an_environment_string_is_read_as_seconds(
gateway: Gateway, tmp_path: Path
) -> None:
config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": "os.environ/SPEND_COUNTER_REDIS_TTL"})
overrides: Final = {"SPEND_COUNTER_REDIS_TTL": ENV_REDIS_TTL}
with owned_proxy(gateway, tmp_path, overrides, config=config) as owned, owned.scenario() as scenario:
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=5.0)
counter: Final = key_counter(key)
assert _chat(owned, model, key)
assert _settled(counter, ONE_CALL_COST, int(ENV_REDIS_TTL))
reset: Final = owned.request("POST", f"/key/{key}/reset_spend", {"reset_to": 0.0})
assert reset.status_code == 200, reset.text
assert _settled(counter, 0.0, int(ENV_REDIS_TTL))
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_long_batch_write_interval_lifts_the_counter_ttl(gateway: Gateway, tmp_path: Path) -> None:
config: Final = _config(tmp_path, general_settings={"proxy_batch_write_at": LONG_BATCH_WRITE_SECONDS})
with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario:
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=5.0)
assert _chat(owned, model, key)
assert _settled(key_counter(key), ONE_CALL_COST, 10 * LONG_BATCH_WRITE_SECONDS)
@pytest.mark.timeout(CELL_TIMEOUT)
def test_strict_enforcement_repairs_a_stale_low_key_counter_with_the_long_ttl(gateway: Gateway, tmp_path: Path) -> None:
config: Final = _config(tmp_path, general_settings={"fail_closed_budget_enforcement": True})
with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario, rig_cache() as cache:
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], max_budget=STRICT_KEY_BUDGET)
counter: Final = key_counter(key)
assert _chat(owned, model, key)
eventually(
lambda: read_rows(
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),)
),
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= ONE_CALL_COST,
seconds=70,
)
attempts: Final[list[str]] = []
def stale_then_request() -> tuple[float | None, int]:
assert cache.set(counter, "0.01")
attempts.append(_chat(owned, model, key))
return eventually(
lambda: (counter_value(cache, counter), counter_ttl(cache, counter)),
lambda state: state[0] is not None and state[0] < STRICT_KEY_BUDGET / 2,
seconds=15,
)
repaired: Final = eventually(
stale_then_request,
lambda state: state[0] is not None and state[0] > 0.01 + ONE_CALL_COST + 1e-6,
seconds=30,
)
lifted: Final = repaired[0]
assert (
lifted is not None and 2 * ONE_CALL_COST - 1e-9 <= lifted <= (len(attempts) + 1) * ONE_CALL_COST + 1e-9
), (
repaired,
attempts,
)
assert outlives(repaired[1]), repaired
def _window_reset_at(key: str) -> datetime:
rows: Final = read_rows(
'SELECT budget_limits FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),)
)
assert len(rows) == 1, rows
raw: Final = rows[0]["budget_limits"]
windows: Final = json.loads(raw) if isinstance(raw, str) else raw
assert isinstance(windows, list) and len(windows) == 1, raw
return datetime.fromisoformat(str(object_value(windows[0])["reset_at"]).replace("Z", "+00:00"))
def _make_window_due(key: str) -> None:
due: Final = (datetime.now(UTC) - timedelta(days=1)).replace(tzinfo=None).isoformat()
windows: Final = [{"max_budget": 1.0, "budget_duration": "1d", "reset_at": due}]
with psycopg.connect(os.environ["DATABASE_URL"]) as connection:
connection.execute(
'UPDATE "LiteLLM_VerificationToken" SET budget_limits = %s::jsonb WHERE token = %s',
(json.dumps(windows), sha256(key.encode()).hexdigest()),
)
@pytest.mark.timeout(CELL_TIMEOUT)
def test_a_reset_budget_window_keeps_the_long_ttl(gateway: Gateway, tmp_path: Path) -> None:
with owned_proxy(gateway, tmp_path, FAST_RESET_TICK) as owned, owned.scenario() as scenario, rig_cache() as cache:
model: Final = scenario.model(**RATES)
key: Final = scenario.key(models=[model], budget_limits=[{"max_budget": 1.0, "budget_duration": "1d"}])
window_counter: Final = f"{key_counter(key)}:window:1d"
assert _chat(owned, model, key)
assert _settled(window_counter, ONE_CALL_COST, 3600)
_make_window_due(key)
rolled: Final = eventually(
lambda: (_window_reset_at(key), counter_value(cache, window_counter), counter_ttl(cache, window_counter)),
lambda state: state[0].replace(tzinfo=None) > datetime.now(UTC).replace(tzinfo=None),
seconds=45,
)
assert rolled[1] is not None and rolled[1] < ONE_CALL_COST, rolled
assert outlives(rolled[2]), rolled

97
tests/unit/caching/test_redis_cache.py
View file

@ -403,6 +403,44 @@ async def test_async_scan_iter_namespaces_pattern(namespace, expected_match, mon
assert captured["match"] == expected_match
@pytest.mark.asyncio
async def test_async_scan_iter_decodes_bytes_keys_and_stops_at_count(monkeypatch, redis_no_ping):
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache()
def scan_iter(match, count):
async def gen():
for key in (b"spend:key:a", "spend:key:b", b"spend:key:c", b"spend:key:d"):
yield key
return gen()
mock_redis_instance = MagicMock()
mock_redis_instance.scan_iter = scan_iter
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
keys = await redis_cache.async_scan_iter(pattern="spend:key:", count=3)
assert keys == ["spend:key:a", "spend:key:b", "spend:key:c"]
@pytest.mark.asyncio
async def test_async_scan_iter_unbounded_returns_every_matching_key(monkeypatch, redis_no_ping):
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache()
def scan_iter(match, count):
async def gen():
for key in (b"spend:key:a", "spend:key:b", b"spend:key:c", b"spend:key:d"):
yield key
return gen()
mock_redis_instance = MagicMock()
mock_redis_instance.scan_iter = scan_iter
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
keys = await redis_cache.async_scan_iter(pattern="spend:key:", count=3, unbounded=True)
assert keys == ["spend:key:a", "spend:key:b", "spend:key:c", "spend:key:d"]
@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")])
def test_increment_cache_namespaces_key(namespace, expected, monkeypatch, redis_no_ping):
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
@ -1516,6 +1554,65 @@ async def test_async_set_cache_pipeline_with_ttls_keeps_each_entry_ttl(monkeypat
]
class _LostPipeline(_SetRecordingPipeline):
async def execute(self) -> list[bool]:
from redis.exceptions import ConnectionError as RedisConnectionError
raise RedisConnectionError("redis down")
@pytest.mark.asyncio
async def test_async_set_cache_pipeline_with_ttls_returns_the_error_of_a_lost_write(monkeypatch, redis_no_ping):
"""A spend reset must not report success on a write Redis never applied, so the lost write comes back
as a value while every other caller keeps ignoring it."""
from redis.exceptions import ConnectionError as RedisConnectionError
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache()
client = MagicMock()
client.pipeline = MagicMock(side_effect=(_LostPipeline(), _SetRecordingPipeline()))
with patch.object(redis_cache, "init_async_client", return_value=client):
lost = await redis_cache.async_set_cache_pipeline_with_ttls((("k", 1, 60),))
kept = await redis_cache.async_set_cache_pipeline_with_ttls((("k", 1, 60),))
assert isinstance(lost, RedisConnectionError)
assert kept is None
@pytest.mark.parametrize("namespace", [None, "ns"])
def test_reset_fence_key_shares_the_cluster_slot_of_its_counter(namespace, monkeypatch, redis_no_ping):
"""The repair checks the fence and the counter in one Lua call, which a cluster only runs on one slot."""
from redis.crc import key_slot
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache(namespace=namespace)
counter = redis_cache.check_and_fix_namespace(key="spend:key:abc")
fence = redis_cache.check_and_fix_namespace(key=redis_cache.reset_fence_key("spend:key:abc"))
assert fence != counter
assert key_slot(fence.encode()) == key_slot(counter.encode())
@pytest.mark.asyncio
async def test_async_set_max_declares_the_reset_fence_as_a_second_key(monkeypatch, redis_no_ping):
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache(namespace="ns")
client = MagicMock()
client.eval = AsyncMock(return_value=b"3.0")
with patch.object(redis_cache, "init_async_client", return_value=client):
held = await redis_cache.async_set_max(
key="spend:key:abc", value=12.0, ttl=60, fence_key=redis_cache.reset_fence_key("spend:key:abc")
)
unfenced = await redis_cache.async_set_max(key="spend:key:abc", value=12.0, ttl=60)
assert (held, unfenced) == (3.0, 3.0)
fenced_call, plain_call = client.eval.await_args_list
assert fenced_call.args[1:4] == (2, "ns:spend:key:abc", "ns:{ns:spend:key:abc}:reset_fence")
assert plain_call.args[1:3] == (1, "ns:spend:key:abc")
class _ListPipeline:
def __init__(self, rows: list[str]) -> None:
self.rows = rows

6
tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py
View file

@ -10655,6 +10655,7 @@ async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_dele
caller passing new_spend must match that precedent, not merely delete the counter."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.proxy_server import spend_counter_redis_ttl
real_cache = UserApiKeyCache()
real_spend_counter_cache = DualCache()
@ -10673,7 +10674,9 @@ async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_dele
)
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 2.5
fake_redis_cache.async_set_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1", value=2.5, ttl=60)
fake_redis_cache.async_set_cache.assert_awaited_once_with(
key="spend:team_member:user-1:team-1", value=2.5, ttl=spend_counter_redis_ttl()
)
@pytest.mark.asyncio
@ -10804,6 +10807,7 @@ async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_t
redis_cache=_FakeRedisCache(),
user_api_key_cache=remote_user_api_key_cache,
additional_in_memory_caches=(remote_spend_counter_in_memory_cache,),
counter_cache=remote_spend_counter_in_memory_cache,
)
for cache_key in ("spend:team_member:user-1:team-1", "spend_db_floor:spend:team_member:user-1:team-1"):
remote_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API

58
tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py
View file

@ -19,6 +19,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
broadcast_evictions,
evict_and_broadcast,
publish_auth_cache_invalidation,
publish_auth_cache_invalidations,
)
from litellm.proxy.common_utils.user_api_key_cache import (
AUTH_OBJECTS_TARGET,
@ -298,6 +299,63 @@ async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message()
assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None
@pytest.mark.asyncio
async def test_publish_many_sends_one_message_per_five_hundred_keys() -> None:
client = _RecordingRedisClient()
keys = tuple(f"spend:key:{index}" for index in range(1001))
with patch(
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
return_value=_FakeRedisCache(client=client),
):
await publish_auth_cache_invalidations(cache_keys=keys, new_value=0.0, ttl=5)
await asyncio.gather(*pubsub_module._pending_publishes)
assert [channel for channel, _ in client.published] == [AUTH_CACHE_INVALIDATION_CHANNEL] * 3
assert [json.loads(message) for _, message in client.published] == [
{"cache_keys": list(keys[:500]), "new_value": 0.0, "ttl": 5},
{"cache_keys": list(keys[500:1000]), "new_value": 0.0, "ttl": 5},
{"cache_key": keys[1000], "new_value": 0.0, "ttl": 5},
]
@pytest.mark.asyncio
async def test_subscriber_sets_every_key_of_a_multi_key_message_in_the_counter_cache_only() -> None:
"""A global reset broadcasts thousands of counters; written into every additional cache they would evict
the BYOK credentials cached next to them."""
cache = UserApiKeyCache()
spend_counter_in_memory_cache = InMemoryCache()
spend_counter_in_memory_cache.set_cache("spend:key:a", 999.0)
spend_counter_in_memory_cache.set_cache("spend_db_floor:spend:key:a", 999.0)
byok_cache = InMemoryCache()
byok_cache.set_cache("byok:cred-1", {"api_key": "sk-cred"})
message = {
"type": "message",
"data": json.dumps({"cache_keys": ["spend:key:a", "spend:key:b"], "new_value": 0.0, "ttl": 60}).encode(),
}
pubsub = _QueuePubSub(initial_messages=[message])
subscriber = AuthCacheInvalidationSubscriber(
redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])),
user_api_key_cache=cache,
additional_in_memory_caches=(spend_counter_in_memory_cache, byok_cache),
counter_cache=spend_counter_in_memory_cache,
)
subscriber.start()
try:
for _ in range(200):
if spend_counter_in_memory_cache.get_cache("spend:key:b") == 0.0:
break
await asyncio.sleep(0.01)
finally:
await subscriber.stop()
assert spend_counter_in_memory_cache.get_cache("spend:key:a") == 0.0
assert spend_counter_in_memory_cache.get_cache("spend:key:b") == 0.0
assert spend_counter_in_memory_cache.get_cache("spend_db_floor:spend:key:a") == 999.0
assert byok_cache.get_cache("byok:cred-1") == {"api_key": "sk-cred"}
assert byok_cache.get_cache("spend:key:a") is None
@pytest.mark.asyncio
async def test_subscriber_ignores_malformed_messages() -> None:
cache = UserApiKeyCache()

32
tests/unit/proxy/common_utils/test_reset_budget_job.py
View file

@ -4845,3 +4845,35 @@ async def test_reset_budget_for_litellm_team_members_called():
assert len(team_member_writes) == 1
assert team_member_writes[0]["where"]["budget_id"]["in"] == ["budget1"]
assert team_member_writes[0]["data"] == {"spend": 0}
@pytest.mark.asyncio
async def test_reset_expired_window_writes_the_redis_counter_with_the_spend_counter_ttl(monkeypatch):
"""The window reset is a counter write like every other, so it must carry the spend counter TTL, not the 60s cache default."""
from datetime import datetime
from litellm.proxy.common_utils.reset_budget_job import BudgetResetSettings, Litellm_EntityType
fake_module = types.ModuleType("litellm.proxy.proxy_server")
fake_module.spend_counter_redis_ttl = lambda: 4321
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
spend_counter_cache = MagicMock()
spend_counter_cache.redis_cache.async_set_cache = AsyncMock()
window = {"budget_duration": "1d", "max_budget": 1.0, "reset_at": "2026-01-01T00:00:00Z"}
reset = await ResetBudgetJob._reset_expired_window(
window=window,
counter_key="spend:key:abc:window:1d",
spend_counter_cache=spend_counter_cache,
now=datetime(2026, 1, 2),
reset_settings=BudgetResetSettings(),
prisma_client=AsyncMock(),
entity_type=Litellm_EntityType.KEY,
entity_id="abc",
)
assert reset is True
spend_counter_cache.redis_cache.async_set_cache.assert_awaited_once_with(
key="spend:key:abc:window:1d", value=0.0, ttl=4321
)
spend_counter_cache.in_memory_cache.set_cache.assert_called_once_with(key="spend:key:abc:window:1d", value=0.0)

6
tests/unit/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -106,7 +106,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
validate_key_list_check,
validate_key_team_change,
)
from litellm.proxy.proxy_server import app
from litellm.proxy.proxy_server import app, spend_counter_redis_ttl
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateKeyRequest,
BulkUpdateKeyResponse,
@ -8781,7 +8781,7 @@ async def test_reset_key_spend_resets_budget_windows(monkeypatch):
key=window_counter_key, value=0.0, ttl=60
)
mock_spend_counter_cache.redis_cache.async_set_cache.assert_any_call(
key=window_counter_key, value=0.0, ttl=60
key=window_counter_key, value=0.0, ttl=spend_counter_redis_ttl()
)
mock_spend_counter_cache.in_memory_cache.set_cache.assert_any_call(
key=f"spend_db_floor:{window_counter_key}", value=0.0, ttl=5
@ -8993,7 +8993,7 @@ async def test_update_key_spend_updates_counter(monkeypatch):
key=f"spend:key:{hashed_key}", value=0.0, ttl=60
)
mock_spend_counter_cache.redis_cache.async_set_cache.assert_awaited_once_with(
key=f"spend:key:{hashed_key}", value=0.0, ttl=60
key=f"spend:key:{hashed_key}", value=0.0, ttl=spend_counter_redis_ttl()
)

79
tests/unit/proxy/management_endpoints/test_organization_endpoints.py
View file

@ -3,7 +3,7 @@ import json
from contextlib import asynccontextmanager
from types import MappingProxyType, SimpleNamespace
from typing import Final, Mapping, Optional, cast
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, call, patch
import pytest
from fastapi import HTTPException
@ -648,6 +648,7 @@ async def test_organization_member_add_budget_omission_and_null_leave_budget_uns
from litellm.proxy.management_endpoints.organization_endpoints import organization_member_add
user = LiteLLM_UserTable(user_id="user-1", user_role="internal_user")
async def create_membership(data):
return LiteLLM_OrganizationMembershipTable(
user_id="user-1",
@ -1582,6 +1583,82 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke
assert jwt_table.rows == (kept_row,)
def _org_spend_counter_cache(monkeypatch: pytest.MonkeyPatch, *, lost_write: Exception | None = None) -> MagicMock:
from litellm.caching.in_memory_cache import InMemoryCache
spend_counter_cache = MagicMock()
spend_counter_cache.in_memory_cache = InMemoryCache()
spend_counter_cache.in_memory_cache.set_cache(key="spend:org:org-1", value=4.2)
spend_counter_cache.in_memory_cache.set_cache(key="spend_db_floor:spend:org:org-1", value=4.2)
spend_counter_cache.redis_cache.reset_fence_key = lambda key: f"{{{key}}}:reset_fence"
spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls = AsyncMock(return_value=lost_write)
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache)
monkeypatch.setattr("litellm.proxy.proxy_server.publish_auth_cache_invalidations", AsyncMock())
return spend_counter_cache
@pytest.mark.asyncio
async def test_legacy_update_with_spend_sets_the_org_counter_and_its_floor_on_every_pod(monkeypatch):
"""A direct spend write must SET the cross-pod counter and its floor marker: a pod that cached the old floor
would otherwise repair a deleted-and-reseeded counter back up, and the org stays at the stale value until
the counter expires."""
from litellm.proxy import proxy_server
spend_counter_cache = _org_spend_counter_cache(monkeypatch)
await _run_legacy_update_organization(
monkeypatch, body={"organization_id": "org-1", "spend": 1.5}, existing_budget_id="budget-1"
)
spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.assert_awaited_once_with(
(
("{spend:org:org-1}:reset_fence", 1, proxy_server.SPEND_RESET_FENCE_TTL_SECONDS),
("spend:org:org-1", 1.5, proxy_server.spend_counter_redis_ttl()),
)
)
assert spend_counter_cache.in_memory_cache.get_cache(key="spend:org:org-1") == 1.5
assert spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:org:org-1") == 1.5
assert proxy_server.publish_auth_cache_invalidations.await_args_list == [
call(("spend_db_floor:spend:org:org-1",), new_value=1.5, ttl=proxy_server.SPEND_DB_FLOOR_CACHE_TTL_SECONDS),
call(("spend:org:org-1",), new_value=1.5, ttl=60),
]
@pytest.mark.asyncio
async def test_legacy_update_with_spend_answers_500_when_the_redis_counter_write_is_lost(monkeypatch):
"""The row is saved either way, but a lost counter write leaves the other pods enforcing the old spend
for up to the counter TTL, so the admin hears about it and retries instead of reading success."""
from redis.exceptions import ConnectionError as RedisConnectionError
from litellm.proxy import proxy_server
from litellm.proxy._types import ProxyException
spend_counter_cache = _org_spend_counter_cache(monkeypatch, lost_write=RedisConnectionError("redis down"))
with pytest.raises(ProxyException) as raised:
await _run_legacy_update_organization(
monkeypatch, body={"organization_id": "org-1", "spend": 1.5}, existing_budget_id="budget-1"
)
assert raised.value.code == "500"
assert raised.value.message.startswith("The organization spend was saved, but the cached spend counters")
proxy_server.prisma_client.db.litellm_organizationtable.update.assert_awaited_once()
assert spend_counter_cache.in_memory_cache.get_cache(key="spend:org:org-1") == 1.5
assert proxy_server.publish_auth_cache_invalidations.await_count == 2
@pytest.mark.asyncio
async def test_legacy_update_without_spend_keeps_the_org_spend_counter(monkeypatch):
spend_counter_cache = _org_spend_counter_cache(monkeypatch)
await _run_legacy_update_organization(
monkeypatch, body={"organization_id": "org-1", "organization_alias": "renamed"}, existing_budget_id="budget-1"
)
spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.assert_not_awaited()
assert spend_counter_cache.in_memory_cache.get_cache(key="spend:org:org-1") == 4.2
@pytest.mark.asyncio
async def test_delete_organization_is_refused_while_one_of_its_teams_holds_a_ptu_share(monkeypatch):
"""Deleting an organization deletes its teams, so a team still holding a PTU share stops it the

62
tests/unit/proxy/proxy_server/test_spend_counters.py
View file

@ -78,6 +78,7 @@ 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.reset_fence_key = lambda key: f"{{{key}}}:reset_fence"
cache.redis_cache.async_increment_pipeline = AsyncMock(return_value=None)
cache.redis_cache.get_ttl = MagicMock(return_value=None)
else:
@ -155,7 +156,9 @@ 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, ttl=ps.spend_counter_redis_ttl(), fence_key="{spend:key:abc}:reset_fence"
)
@pytest.mark.asyncio
@ -303,7 +306,12 @@ async def test_get_current_spend_end_user_floor_repairs_a_stale_low_counter(monk
)
assert result == 12.0
fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key="spend:end_user:customer-42", value=12.0)
fake_cache.redis_cache.async_set_max.assert_awaited_once_with(
key="spend:end_user:customer-42",
value=12.0,
ttl=ps.spend_counter_redis_ttl(),
fence_key="{spend:end_user:customer-42}:reset_fence",
)
@pytest.mark.asyncio
@ -347,7 +355,9 @@ 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, ttl=ps.spend_counter_redis_ttl(), fence_key=f"{{{counter_key}}}:reset_fence"
)
def _make_window_spend_prisma(row=None, spend_logs_total=0.0):
@ -388,7 +398,9 @@ 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, ttl=ps.spend_counter_redis_ttl(), fence_key=f"{{{counter_key}}}:reset_fence"
)
@pytest.mark.asyncio
@ -851,6 +863,40 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call(
assert counter_cache.in_memory_cache.get_cache(key=key) == 1.5
def test_spend_counter_redis_ttl_never_shortens_configured_redis_ttl(monkeypatch):
"""Deployments that configured a longer Redis default must not see spend counters expire sooner"""
monkeypatch.setattr(ps.litellm, "default_redis_ttl", 86400)
monkeypatch.setattr(ps, "proxy_batch_write_at", 60)
assert ps.spend_counter_redis_ttl() == 86400
@pytest.mark.asyncio
@pytest.mark.parametrize("batch_write_at", [60, 900])
async def test_spend_counter_redis_ttl_outlives_db_batch_write_interval(monkeypatch, batch_write_at):
"""A counter that expires before its buffered spend is written to the DB reseeds from a stale-low
row and re-admits an over-budget key, so every counter write must outlive the batch interval"""
counter_cache = ps.DualCache()
fake_redis = AsyncMock()
fake_redis.async_increment = AsyncMock(return_value=1.0)
fake_redis.async_increment_pipeline = AsyncMock(return_value=[1.0])
fake_redis.async_refresh_ttl = AsyncMock(return_value=True)
fake_redis.get_ttl = MagicMock(return_value=60)
fake_redis.default_ttl = 60
counter_cache.redis_cache = fake_redis
monkeypatch.setattr(ps, "spend_counter_cache", counter_cache)
monkeypatch.setattr(ps, "proxy_batch_write_at", batch_write_at)
max_flush_delay: Final = batch_write_at + 5
await ps.increment_spend_counter_cache(counter_key="spend:key:direct", increment=0.5)
await ps.run_spend_counter_pipeline((ps.PendingSpendIncrement(counter_key="spend:key:batched", increment=0.5),))
await ps.refresh_spend_counter_ttl(counter_key="spend:key:reserved")
assert fake_redis.async_increment.await_args.kwargs["ttl"] > max_flush_delay
assert fake_redis.async_increment_pipeline.await_args.kwargs["increment_list"][0]["ttl"] > max_flush_delay
assert fake_redis.async_refresh_ttl.await_args.kwargs["ttl"] > max_flush_delay
_PENDING_SPEND_COUNTER_KEYS: tuple[str, ...] = (
"spend:key:hashed",
"spend:team:team-1",
@ -1597,6 +1643,14 @@ async def test_update_cache_user_cache_failure_invalid_state_is_swallowed(monkey
assert result is None
@pytest.mark.parametrize(("max_batch_size", "keys_per_pipeline"), [(1, 1), (2, 1), (3, 1), (1000, 500)])
def test_a_reset_pipeline_takes_half_the_redis_batch_size_and_never_less_than_one_key(
max_batch_size: int, keys_per_pipeline: int
) -> None:
"""Each key rides with its fence, so half the batch size; a step of zero would make every reset raise."""
assert ps._keys_per_reset_pipeline(max_batch_size) == keys_per_pipeline
async def _drain_spend_cache_writes() -> None:
for _ in range(3):
await asyncio.sleep(0)

9
tests/unit/proxy/spend_tracking/test_spend_counter_batch.py
View file

@ -41,6 +41,7 @@ class CountingRedis(RedisCache):
self.store: dict[str, object] = dict(store or {})
self.fail = fail
self.commands: list[str] = []
self.namespace = None
async def async_get_cache(self, key: str, **kwargs: object) -> object:
if self.fail:
@ -351,7 +352,9 @@ async def test_reseed_treats_a_batched_clean_miss_as_authoritative_and_seeds_fro
assert value == 2.25
assert [c for c in redis.commands if c.startswith("GET")] == []
redis.async_set_cache.assert_awaited_once_with(key="spend:key:hashed", value=2.25, nx=True)
redis.async_set_cache.assert_awaited_once_with(
key="spend:key:hashed", value=2.25, nx=True, ttl=ps.spend_counter_redis_ttl()
)
@pytest.mark.asyncio
@ -404,7 +407,9 @@ async def test_post_call_cold_counters_seed_from_the_mget_miss_without_a_second_
await ps.increment_spend_counters(token="hashed", team_id="team", user_id=None, response_cost=0.5)
assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE"], redis.commands
redis.async_set_cache.assert_awaited_once_with(key="spend:team:team", value=4.0, nx=True)
redis.async_set_cache.assert_awaited_once_with(
key="spend:team:team", value=4.0, nx=True, ttl=ps.spend_counter_redis_ttl()
)
assert redis.store["spend:key:hashed"] == 1.5

154
tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py
View file

@ -5,10 +5,11 @@ import hashlib
import json
import re
import sqlite3
from collections.abc import Mapping, Sequence
from collections.abc import AsyncIterator, Mapping, Sequence
from contextlib import asynccontextmanager
from datetime import timezone
from typing import Final, Literal
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, call, patch
import prisma
import pytest
@ -8336,6 +8337,155 @@ async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader
]
def _spend_counter_cache_for_reset(
monkeypatch: pytest.MonkeyPatch, *, scan_results: object, lost_write: Exception | None = None
) -> MagicMock:
from litellm.caching.in_memory_cache import InMemoryCache
mock_prisma_client: Final = MagicMock()
mock_prisma_client.db.query_raw = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.update_many = AsyncMock()
mock_prisma_client.db.litellm_teamtable.update_many = AsyncMock()
@asynccontextmanager
async def transaction() -> AsyncIterator[MagicMock]:
yield mock_prisma_client.db
mock_prisma_client.tx = transaction
monkeypatch.setattr(ps, "prisma_client", mock_prisma_client)
spend_counter_cache: Final = MagicMock()
spend_counter_cache.in_memory_cache = InMemoryCache()
spend_counter_cache.in_memory_cache.set_cache(key="spend:key:aaa", value=0.5)
spend_counter_cache.in_memory_cache.set_cache(key="spend_db_floor:spend:key:aaa", value=0.5)
spend_counter_cache.in_memory_cache.set_cache(key="spend_db_floor:spend:key:floored", value=0.9)
spend_counter_cache.in_memory_cache.set_cache(key="spend:key:aaa:window:1d", value=0.3)
spend_counter_cache.in_memory_cache.set_cache(key="spend:team:t1", value=0.5)
spend_counter_cache.in_memory_cache.set_cache(key="spend:user:u1", value=0.5)
spend_counter_cache.redis_cache.namespace = "ns"
spend_counter_cache.redis_cache.reset_fence_key = lambda key: f"{{ns:{key}}}:reset_fence"
spend_counter_cache.redis_cache.async_scan_iter = AsyncMock(side_effect=scan_results)
spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls = AsyncMock(return_value=lost_write)
monkeypatch.setattr(ps, "spend_counter_cache", spend_counter_cache)
monkeypatch.setattr(ps, "publish_auth_cache_invalidations", AsyncMock())
return spend_counter_cache
_RESET_KEY_COUNTERS: Final = ("spend:key:aaa", "spend:key:bbb", "spend:key:floored")
_RESET_TEAM_COUNTERS: Final = ("spend:team:sales:window:east", "spend:team:t1")
def _reset_pipeline(counters: tuple[str, ...]) -> tuple[tuple[str, float, int], ...]:
"""Every fence is SET before any counter, so a repair racing the pipeline already finds its fence."""
return (
*((f"{{ns:{counter}}}:reset_fence", 1, ps.SPEND_RESET_FENCE_TTL_SECONDS) for counter in counters),
*((counter, 0.0, ps.spend_counter_redis_ttl()) for counter in counters),
)
def _reset_broadcasts(counters: tuple[str, ...]) -> tuple[object, object]:
"""The floors go out first: a subscriber that applied the new counter under the old floor would repair it."""
return (
call(
tuple(f"spend_db_floor:{counter}" for counter in counters),
new_value=0.0,
ttl=ps.SPEND_DB_FLOOR_CACHE_TTL_SECONDS,
),
call(counters, new_value=0.0, ttl=60),
)
@pytest.mark.asyncio
async def test_global_spend_reset_sets_every_key_and_team_counter_and_its_floor_to_zero(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Deleting the counters is not enough: a pod that cached the pre-reset floor within the last 5s would repair
the reseeded counter back to the old total for the whole counter TTL, so the reset SETs counter and floor
everywhere, fences each counter against that repair, and leaves budget-window counters, which the reset
does not touch in the DB, alone. A window counter is told by its trailing duration, so a team whose id
merely contains ``:window:`` is still reset."""
spend_counter_cache: Final = _spend_counter_cache_for_reset(
monkeypatch,
scan_results=[
["ns:spend:key:aaa", "ns:spend:key:bbb", "ns:spend:key:bbb:window:1d"],
["ns:spend:team:t1", "ns:spend:team:t1:window:1d", "ns:spend:team:sales:window:east"],
],
)
response: Final = await spend_management_endpoints.global_spend_reset()
assert response["status"] == "success"
assert [c.kwargs for c in spend_counter_cache.redis_cache.async_scan_iter.await_args_list] == [
{"pattern": "spend:key:", "count": 500, "unbounded": True},
{"pattern": "spend:team:", "count": 500, "unbounded": True},
]
assert [c.args for c in spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.await_args_list] == [
(_reset_pipeline(_RESET_KEY_COUNTERS),),
(_reset_pipeline(_RESET_TEAM_COUNTERS),),
]
in_memory: Final = spend_counter_cache.in_memory_cache
reset_counters: Final = (*_RESET_KEY_COUNTERS, *_RESET_TEAM_COUNTERS)
assert [in_memory.get_cache(key=counter) for counter in reset_counters] == [0.0] * 5
assert [in_memory.get_cache(key=f"spend_db_floor:{counter}") for counter in reset_counters] == [0.0] * 5
assert in_memory.get_cache(key="spend:key:aaa:window:1d") == 0.3
assert in_memory.get_cache(key="spend:team:t1:window:1d") is None
assert in_memory.get_cache(key="spend:user:u1") == 0.5
assert ps.publish_auth_cache_invalidations.await_args_list == [
*_reset_broadcasts(_RESET_KEY_COUNTERS),
*_reset_broadcasts(_RESET_TEAM_COUNTERS),
]
@pytest.mark.asyncio
async def test_global_spend_reset_answers_500_after_resetting_the_rows_and_this_pod_when_redis_cannot_be_listed(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""The counters Redis holds for the other pods stay at the old spend for the whole counter TTL when the
scan fails, so the reset says so instead of reporting success, after resetting what this pod knows."""
spend_counter_cache: Final = _spend_counter_cache_for_reset(monkeypatch, scan_results=ConnectionError("redis down"))
with pytest.raises(ProxyException) as raised:
await spend_management_endpoints.global_spend_reset()
assert raised.value.code == "500"
assert "retry this request once Redis accepts writes" in raised.value.message
ps.prisma_client.db.litellm_verificationtoken.update_many.assert_awaited_once_with(data={"spend": 0.0}, where={})
ps.prisma_client.db.litellm_teamtable.update_many.assert_awaited_once_with(data={"spend": 0.0}, where={})
assert [c.args for c in spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.await_args_list] == [
(_reset_pipeline(("spend:key:aaa", "spend:key:floored")),),
(_reset_pipeline(("spend:team:t1",)),),
]
assert spend_counter_cache.in_memory_cache.get_cache(key="spend:key:aaa") == 0.0
assert spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:key:floored") == 0.0
@pytest.mark.asyncio
async def test_global_spend_reset_answers_500_when_the_redis_counter_write_is_lost(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A lost pipeline leaves the other pods on the old counters for up to the counter TTL, so the reset still
resets the teams, this pod, and the broadcast, then reports the loss so the admin retries."""
from redis.exceptions import ConnectionError as RedisConnectionError
spend_counter_cache: Final = _spend_counter_cache_for_reset(
monkeypatch,
scan_results=[["ns:spend:key:aaa"], ["ns:spend:team:t1"]],
lost_write=RedisConnectionError("redis down"),
)
with pytest.raises(ProxyException) as raised:
await spend_management_endpoints.global_spend_reset()
assert raised.value.code == "500"
assert raised.value.message.startswith("Spend was reset, but the cached spend counters could not be reset in Redis")
assert "redis down" in raised.value.message
assert spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.await_count == 2
assert spend_counter_cache.in_memory_cache.get_cache(key="spend:team:t1") == 0.0
assert ps.publish_auth_cache_invalidations.await_args_list == [
*_reset_broadcasts(("spend:key:aaa", "spend:key:floored")),
*_reset_broadcasts(("spend:team:t1",)),
]
def test_spend_calculate_rejection_body_is_openai_shaped(client):
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user", api_key="hashed-k"

47
tests/unit/proxy/test_budget_reservation.py
View file

@ -41,6 +41,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
)
from litellm.proxy.spend_tracking.budget_reservation import (
_get_model_access_group_budget_counters,
_lease_renewals,
estimate_request_max_cost,
get_budget_window_start,
invalidate_budget_reservation_counters,
@ -2348,6 +2349,9 @@ class _ExpiringRedisCache:
self.refresh_count = 0
self.fail_first_refresh = fail_first_refresh
def _expiry(self, ttl: object) -> float:
return time.monotonic() + (float(ttl) if isinstance(ttl, (int, float)) else self.default_ttl)
def _evict_expired(self, key: str) -> None:
if self.expires_at.get(key, float("inf")) <= time.monotonic():
self.store.pop(key, None)
@ -2360,18 +2364,21 @@ class _ExpiringRedisCache:
async def async_increment(self, key: str, value: float, **kwargs: object) -> float:
self._evict_expired(key)
self.store[key] = self.store.get(key, 0.0) + float(value)
self.expires_at[key] = time.monotonic() + self.default_ttl
self.expires_at[key] = self._expiry(kwargs.get("ttl"))
return self.store[key]
def reset_fence_key(self, key: str) -> str:
return f"{{{key}}}:reset_fence"
async def async_set_max(self, key: str, value: float, **kwargs: object) -> float:
self._evict_expired(key)
self.store[key] = max(self.store.get(key, float("-inf")), float(value))
self.expires_at[key] = time.monotonic() + self.default_ttl
self.expires_at[key] = self._expiry(kwargs.get("ttl"))
return self.store[key]
async def async_set_cache(self, key: str, value: float, *args: object, **kwargs: object) -> bool:
self.store[key] = float(value)
self.expires_at[key] = time.monotonic() + self.default_ttl
self.expires_at[key] = self._expiry(kwargs.get("ttl"))
return True
async def async_delete_cache(self, key: str, *args: object, **kwargs: object) -> None:
@ -2386,13 +2393,15 @@ class _ExpiringRedisCache:
if key not in self.store:
return False
self.refresh_count += 1
self.expires_at[key] = time.monotonic() + (ttl if ttl is not None else self.default_ttl)
self.expires_at[key] = self._expiry(ttl)
return True
async def async_increment_pipeline(
self, increment_list: Sequence[RedisPipelineIncrementOperation], **kwargs: object
) -> list[float]:
return [await self.async_increment(op["key"], op["increment_value"]) for op in increment_list]
return [
await self.async_increment(op["key"], op["increment_value"], ttl=op.get("ttl")) for op in increment_list
]
def get_ttl(self, **kwargs: object) -> int | None:
return int(self.default_ttl)
@ -2401,11 +2410,13 @@ class _ExpiringRedisCache:
@pytest.mark.asyncio
async def test_reservation_survives_redis_counter_ttl_while_request_in_flight(
spend_counter_state,
monkeypatch,
):
"""A request that runs longer than the counter TTL must keep its reservation in Redis
(so a concurrent request on any worker still sees it), and renewal must stop once the
reservation is reconciled so an idle counter still expires on its own."""
counter_cache, key_cache = spend_counter_state
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_redis_ttl", lambda: 0.2)
redis_cache = _ExpiringRedisCache(default_ttl=0.2)
counter_cache.redis_cache = redis_cache
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
@ -2430,13 +2441,37 @@ async def test_reservation_survives_redis_counter_ttl_while_request_in_flight(
assert await redis_cache.async_get_cache(key=counter_key) is None
@pytest.mark.asyncio
async def test_reservation_lease_renewal_stops_promptly_with_long_spend_counter_ttl(
spend_counter_state,
monkeypatch,
):
"""A long spend-counter TTL must not stretch how often renewal checks for completion, or finished
requests keep their renewal tasks alive for half that TTL"""
counter_cache, key_cache = spend_counter_state
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_redis_ttl", lambda: 3600)
counter_cache.redis_cache = _ExpiringRedisCache(default_ttl=0.2)
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
valid_token = UserAPIKeyAuth(token="key-lease-long-ttl", spend=0.0, max_budget=1.0)
reservation = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj)
assert reservation is not None
assert len(_lease_renewals) == 1
await release_budget_reservation(reservation)
await asyncio.sleep(0.3)
assert len(_lease_renewals) == 0
@pytest.mark.asyncio
async def test_reservation_lease_keeps_renewing_after_transient_redis_failure(
spend_counter_state,
monkeypatch,
):
"""One failed EXPIRE (Redis blip, open circuit breaker) must not end renewal for the
rest of the request."""
counter_cache, key_cache = spend_counter_state
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_redis_ttl", lambda: 0.2)
redis_cache = _ExpiringRedisCache(default_ttl=0.2, fail_first_refresh=True)
counter_cache.redis_cache = redis_cache
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
@ -2453,11 +2488,13 @@ async def test_reservation_lease_keeps_renewing_after_transient_redis_failure(
@pytest.mark.asyncio
async def test_reservation_lease_stops_when_request_task_ends_without_reconciling(
spend_counter_state,
monkeypatch,
):
"""A request whose task ends without reconciling (client disconnect path that skips the
cost callbacks) must not keep renewing: the counter falls back to its plain TTL instead of
pinning the reservation until the request timeout."""
counter_cache, key_cache = spend_counter_state
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_redis_ttl", lambda: 0.2)
redis_cache = _ExpiringRedisCache(default_ttl=0.2)
counter_cache.redis_cache = redis_cache
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)

5
tests/unit/proxy/test_proxy_server_endpoints_and_startup.py
View file

@ -9093,7 +9093,9 @@ async def test_prepare_spend_counter_increment_reseeds_from_db_on_counter_miss()
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.
# Only the per-request delta (1.5) goes through INCRBYFLOAT.
fake_redis.async_set_cache.assert_awaited_once_with(key="spend:team:team-9", value=42.0, nx=True)
fake_redis.async_set_cache.assert_awaited_once_with(
key="spend:team:team-9", value=42.0, nx=True, ttl=proxy_server_module.spend_counter_redis_ttl()
)
writes = [(c["key"], c["value"]) for c in recorded_increments]
assert writes == [("spend:team:team-9", 1.5)]
finally:
@ -9522,6 +9524,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed()
key=counter_key,
value=2.25,
nx=True,
ttl=proxy_server_module.spend_counter_redis_ttl(),
)
assert redis_store[counter_key] == pytest.approx(3.25)
assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(3.25)