diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index fe8b71e7623..d52716951df 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -74,38 +74,31 @@ def duration_in_seconds(duration: str) -> int: elif unit == "w": return value * 604800 elif unit == "mo": - now: Final = time_module.time() - current_time: Final = datetime.fromtimestamp(now) - - # Calculate target month and year, handling overflow past December - total_months: Final = current_time.month - 1 + value # 0-indexed months - target_year: Final = current_time.year + total_months // 12 - target_month: Final = total_months % 12 + 1 # back to 1-indexed - - # Determine the day to set for next month - target_day = current_time.day - last_day_of_target_month: Final = get_last_day_of_month(target_year, target_month) - - target_day = min(target_day, last_day_of_target_month) - - next_month: Final = datetime( - year=target_year, - month=target_month, - day=target_day, - hour=current_time.hour, - minute=current_time.minute, - second=current_time.second, - microsecond=current_time.microsecond, + return _months_in_seconds( + start=datetime.fromtimestamp(time_module.time(), tz=timezone.utc).astimezone(), months=value ) - # Calculate the duration until the first day of the next month - duration_until_next_month: Final = next_month - current_time - return int(duration_until_next_month.total_seconds()) - else: raise ValueError(f"Unsupported duration unit, passed duration: {duration}") +def duration_in_seconds_from(start: datetime, duration: str) -> int: + """How long a period of `duration` lasts when it begins at `start`. + + `duration_in_seconds` measures `Nmo` from now, so a monthly window that opened in a + month of a different length needs its own start to know when it ends. Every other + unit is a fixed length and ignores `start`. + """ + value, unit = _extract_from_regex(duration=_normalize_duration(duration)) + if unit == "mo": + return _months_in_seconds(start=start, months=value) + return duration_in_seconds(duration) + + +def _months_in_seconds(start: datetime, months: int) -> int: + return int((_shift_months(start, months) - start).total_seconds()) + + def get_next_standardized_reset_time( duration: str, current_time: datetime, @@ -166,8 +159,9 @@ def get_next_standardized_reset_time( return base_midnight + timedelta(days=1) -def _subtract_months(moment: datetime, months: int) -> datetime: - total_months: Final = moment.year * 12 + moment.month - 1 - months +def _shift_months(moment: datetime, months: int) -> datetime: + """`moment` moved by `months` calendar months, back when negative, clamped to the target month's last day.""" + total_months: Final = moment.year * 12 + moment.month - 1 + months year, month_index = divmod(total_months, 12) month: Final = month_index + 1 return moment.replace(year=year, month=month, day=min(moment.day, get_last_day_of_month(year, month))) @@ -182,9 +176,9 @@ def get_budget_window_start(duration: str, reset_at: datetime) -> datetime: return reset_at - timedelta(days=1) match unit: case "mo": - return _subtract_months(reset_at, value) + return _shift_months(reset_at, -value) case "d": - return _subtract_months(reset_at, 1) if value == 30 else reset_at - timedelta(days=value) + return _shift_months(reset_at, -1) if value == 30 else reset_at - timedelta(days=value) case "w": return reset_at - timedelta(weeks=value) case "h": diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index d05a62d6d1b..00b5f9cd31c 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -1,8 +1,9 @@ import asyncio import json import time -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass +from datetime import datetime, timezone from types import MappingProxyType from typing import Final, cast # noqa: TID251 # untyped callback payload requires explicit narrowing @@ -17,7 +18,7 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth -from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting, budget_window_seconds from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import BudgetConfig, StandardLoggingPayload @@ -42,6 +43,8 @@ _LEGACY_REQUEST_MODEL_SCOPES: Final = frozenset({Litellm_EntityType.KEY, Litellm _PROCESS_STARTED_AT: Final = time.monotonic() +_LATEST_RESET_AT: Final = datetime(9999, 12, 31, tzinfo=timezone.utc).timestamp() + _BUDGET_START_TIME_KEY_PREFIXES: Final = MappingProxyType( { Litellm_EntityType.KEY: "virtual_key_budget_start_time", @@ -206,17 +209,29 @@ def _bedrock_candidates(model: str) -> tuple[str, ...]: return (base_model, without_vendor) if without_vendor else (base_model,) +@dataclass(frozen=True, slots=True) +class _WindowReading: + current_spend: float + started_at: float | None + + async def build_model_max_budget_usage( entity_type: Litellm_EntityType, entity_id: str | None, model_max_budget: Mapping[str, object] | None, cache: DualCache | None, + now: datetime | None = None, ) -> dict[str, dict[str, object]]: """Current-window spend per configured budget model, as `/key/info` and `/user/info` report it. `cache` must be the DualCache the limiter writes the counters to; callers read it off the limiter rather than re-deriving it, so a scope that is being blocked can never report zero usage. + + `budget_reset_at` is when the open window ends and its spend counter + expires, which is the window's start plus its period. It is None when no + window is open, because nothing has been charged since the last one ended, + and when the window ends after year 9999, so it never resets. """ if cache is None or entity_id is None or not model_max_budget: return {} @@ -229,41 +244,80 @@ async def build_model_max_budget_usage( ) if not budgets: return {} - spend_keys: Final = tuple( - model_budget_spend_cache_key( - entity_type=entity_type, - entity_id=entity_id, - budget_model=budget_model, - budget_duration=budget_config.budget_duration, - ) - for budget_model, budget_config in budgets + windows: Final = await _read_windows( + cache=cache, + spend_keys=tuple( + model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=budget_model, + budget_duration=budget_config.budget_duration, + ) + for budget_model, budget_config in budgets + ), + start_time_keys=tuple( + model_budget_start_time_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=budget_model, + budget_duration=budget_config.budget_duration, + ) + for budget_model, budget_config in budgets + ), ) - current_spends: Final = await _current_window_spends(cache=cache, spend_keys=spend_keys) + current_time: Final = (now or datetime.now(timezone.utc)).timestamp() return { budget_model: { - "current_spend": round(current_spend, 4), + "current_spend": round(window.current_spend, 4), "budget_limit": budget_config.max_budget, "time_period": budget_config.budget_duration, + "budget_reset_at": _window_reset_at( + started_at=window.started_at, + budget_duration=budget_config.budget_duration, + now=current_time, + ), } - for (budget_model, budget_config), current_spend in zip(budgets, current_spends, strict=True) + for (budget_model, budget_config), window in zip(budgets, windows, strict=True) } @with_service_target("model_budgets") -async def _current_window_spends(cache: DualCache, spend_keys: Sequence[str]) -> tuple[float, ...]: - """Redis holds the window total across replicas; the in-memory copy is one replica's share.""" - keys: Final = list(spend_keys) +async def _read_windows( + cache: DualCache, spend_keys: Sequence[str], start_time_keys: Sequence[str] +) -> tuple[_WindowReading, ...]: + """Every window's spend and start time in one batched read, however many models are budgeted. + + Redis holds the window total across replicas; the in-memory copy is one replica's share. + """ + keys: Final = [*spend_keys, *start_time_keys] redis_cache: Final = cache.redis_cache if redis_cache is not None: shared: Final = await redis_cache.async_batch_get_cache(key_list=keys) - return tuple(_as_spend(shared.get(key)) for key in keys) + return _readings(tuple(_as_float(shared.get(key)) for key in keys), window_count=len(spend_keys)) # async_batch_get_cache returns None if it fails internally, and its result is # index-aligned with `keys` otherwise. An unusable result reads as a miss, # which is what a never-written counter already reads as. batched: Final = await cache.async_batch_get_cache(keys=keys) if not isinstance(batched, list) or len(batched) != len(keys): - return (0.0,) * len(keys) - return tuple(_as_spend(current_spend) for current_spend in batched) + return tuple(_WindowReading(current_spend=0.0, started_at=None) for _ in spend_keys) + return _readings(tuple(_as_float(value) for value in batched), window_count=len(spend_keys)) + + +def _readings(values: Sequence[float | None], window_count: int) -> tuple[_WindowReading, ...]: + """`values` holds every window's spend followed by every window's start, in the same order.""" + return tuple( + _WindowReading(current_spend=spend or 0.0, started_at=started_at) + for spend, started_at in zip(values[:window_count], values[window_count:], strict=True) + ) + + +def _window_reset_at(started_at: float | None, budget_duration: str | None, now: float) -> str | None: + if started_at is None or budget_duration is None: + return None + ends_at: Final = started_at + budget_window_seconds(started_at, budget_duration) + if ends_at <= now or ends_at > _LATEST_RESET_AT: + return None + return datetime.fromtimestamp(ends_at, tz=timezone.utc).isoformat() def _usable_budget_config(raw_budget_config: object) -> BudgetConfig | None: @@ -284,6 +338,15 @@ def _as_spend(current_spend: object) -> float: return 0.0 +def _as_float(value: object) -> float | None: + if value is None: + return None + try: + return float(value) # pyright: ignore[reportArgumentType] # non-numeric falls to the except + except (TypeError, ValueError): + return None + + def _resolve_entity_model_budgets( model: str, entity_budgets: Iterable[tuple[Litellm_EntityType, str | None, object]], @@ -317,8 +380,9 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): Example: key=$LITELLM_API_KEY, model=gpt-4o, max_budget=100, time_period=1d """ - def __init__(self, dual_cache: DualCache): + def __init__(self, dual_cache: DualCache, clock: Callable[[], float] = time.time): self.dual_cache = dual_cache + self._clock = clock self.redis_increment_operation_queue = [] self._redis_increment_queue_lock = asyncio.Lock() self._redis_increment_flush_lock = asyncio.Lock() diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 896f12cb0c7..048bc41496f 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4667,7 +4667,8 @@ async def info_key_fn( hour), so subtracting budget_duration from it does not give the window's start - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} - model_max_budget_usage: dict | None - Current-window spend per model, present only when - the key has per-model budgets + the key has per-model budgets. Each entry's budget_reset_at is when that model's open + window ends and its spend resets, or None when no window is open or it ends after year 9999 - budget_limits: list | None - Concurrent budget windows, exactly as stored - budget_limits_usage: dict | None - Current-window spend per budget window, e.g. {"1h": {"current_spend": 0.0009}}, present only when the key has budget windows diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 8ba1da824c6..b20f71cd51f 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -21,7 +21,8 @@ anthropic: import asyncio import builtins import logging -from collections.abc import Mapping, Sequence +import time +from collections.abc import Callable, Mapping, Sequence from datetime import datetime, timedelta, timezone from itertools import groupby from types import MappingProxyType @@ -36,7 +37,7 @@ from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) -from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.litellm_core_utils.duration_parser import duration_in_seconds_from from litellm.router_strategy.tag_based_routing import get_tags_from_request_kwargs from litellm.router_utils.cooldown_callbacks import ( get_prometheus_logger_from_callbacks, @@ -104,14 +105,20 @@ def _sum_increments_by_key(operations: Sequence[RedisPipelineIncrementOperation] return MappingProxyType({key: sum(operation["increment_value"] for operation in group) for key, group in by_key}) +def budget_window_seconds(started_at: float, budget_duration: str) -> int: + return duration_in_seconds_from(start=datetime.fromtimestamp(started_at, tz=timezone.utc), duration=budget_duration) + + class RouterBudgetLimiting(CustomLogger): def __init__( self, dual_cache: DualCache, provider_budget_config: dict | None, model_list: list[DeploymentTypedDict | dict[str, Any]] | None = None, + clock: Callable[[], float] = time.time, ): self.dual_cache = dual_cache + self._clock = clock self.redis_increment_operation_queue: list[RedisPipelineIncrementOperation] = [] self._redis_increment_queue_lock = asyncio.Lock() self._redis_increment_flush_lock = asyncio.Lock() @@ -559,41 +566,29 @@ class RouterBudgetLimiting(CustomLogger): if budget_config.budget_duration is None: return - current_time: Final = datetime.now(timezone.utc).timestamp() - ttl_seconds: Final = duration_in_seconds(budget_config.budget_duration) + current_time: Final = self._clock() + new_window_seconds: Final = budget_window_seconds(current_time, budget_config.budget_duration) - budget_start = await self._get_or_set_budget_start_time( + budget_start: Final = await self._get_or_set_budget_start_time( start_time_key=start_time_key, current_time=current_time, - ttl_seconds=ttl_seconds, + ttl_seconds=new_window_seconds, ) + open_window_seconds: Final = budget_window_seconds(budget_start, budget_config.budget_duration) + elapsed: Final = current_time - budget_start - if budget_start is None: - # First spend for this provider - budget_start = await self._handle_new_budget_window( - spend_key=spend_key, - start_time_key=start_time_key, - current_time=current_time, - response_cost=response_cost, - ttl_seconds=ttl_seconds, - ) - elif (current_time - budget_start) > ttl_seconds: - # Budget window expired - reset everything + if elapsed > open_window_seconds: verbose_router_logger.debug("Budget window expired - resetting everything") - budget_start = await self._handle_new_budget_window( + await self._handle_new_budget_window( spend_key=spend_key, start_time_key=start_time_key, current_time=current_time, response_cost=response_cost, - ttl_seconds=ttl_seconds, + ttl_seconds=new_window_seconds, ) else: - # Within existing window - increment spend - remaining_time: Final = ttl_seconds - (current_time - budget_start) - ttl_for_increment: Final = int(remaining_time) - await self._increment_spend_in_current_window( - spend_key=spend_key, response_cost=response_cost, ttl=ttl_for_increment + spend_key=spend_key, response_cost=response_cost, ttl=int(open_window_seconds - elapsed) ) verbose_router_logger.debug("Incremented spend for %s by %s", spend_key, response_cost) @@ -809,14 +804,15 @@ class RouterBudgetLimiting(CustomLogger): spend_key: Final = f"provider_spend:{provider}:{budget_config.budget_duration}" start_time_key: Final = f"provider_budget_start_time:{provider}" - ttl_seconds: int | None = None - if budget_config.budget_duration is not None: - ttl_seconds = duration_in_seconds(budget_config.budget_duration) + current_time: Final = self._clock() + ttl_seconds: Final = ( + budget_window_seconds(current_time, budget_config.budget_duration) + if budget_config.budget_duration is not None + else None + ) - budget_start = await self.dual_cache.async_get_cache(start_time_key) - if budget_start is None: - budget_start = datetime.now(timezone.utc).timestamp() - await self.dual_cache.async_set_cache(key=start_time_key, value=budget_start, ttl=ttl_seconds) + if await self.dual_cache.async_get_cache(start_time_key) is None: + await self.dual_cache.async_set_cache(key=start_time_key, value=current_time, ttl=ttl_seconds) _spend_key: Final = await self.dual_cache.async_get_cache(spend_key) if _spend_key is None: diff --git a/tests/integration/management/test_budget_reset_windows.py b/tests/integration/management/test_budget_reset_windows.py new file mode 100644 index 00000000000..34a0823cea2 --- /dev/null +++ b/tests/integration/management/test_budget_reset_windows.py @@ -0,0 +1,1020 @@ +import asyncio +import calendar +import json +import math +import os +import time +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from hashlib import sha256 +from itertools import product +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + held, + object_value, +) +from integration._support.database import write_rows +from integration._support.process import graceful_stop_seconds, owned_proxy, owned_proxy_process +from integration._support.redis_process import owned_redis +from integration._support.upstream import delete_scenario, register_scenario +from integration.authorization._hidden_alias_budget import ( + BUDGET_EXCEEDED, + CHAT_REPLY, + anthropic_client, + async_anthropic_client, + async_openai_client, + base_url, + fresh_chat, + landed_all_once, + landed_once, + openai_client, + settle, + spend_marker, + upstream_hits, + upstream_requests, +) +from integration.cost_calculation.cost_tracking_case import JsonResponse, SseResponse +from pydantic import JsonValue, TypeAdapter +from redis import Redis + +pytestmark: Final = pytest.mark.timeout(240) + +DAY: Final = 86400 +HOUR: Final = 3600 +REPLY: Final = "window reply" +PRICES: Final[Mapping[str, JsonValue]] = {"input_cost_per_token": 0.0001, "output_cost_per_token": 0.0002} +OWNED_PROXY_SECONDS: Final = 2 * graceful_stop_seconds() + 120 +REDIS_TEXT: Final = TypeAdapter(str | None) +REDIS_INT: Final = TypeAdapter(int) + +_RESPONSE: Final[dict[str, JsonValue]] = { + "id": "resp_$UNIQUE_ID", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "id": "msg_$UNIQUE_ID", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": REPLY, "annotations": []}], + } + ], + "usage": {"input_tokens": 20, "output_tokens": 20, "total_tokens": 40}, +} +_RESPONSE_EVENTS: Final[tuple[Mapping[str, JsonValue], ...]] = ( + { + "type": "response.created", + "sequence_number": 0, + "response": {**_RESPONSE, "status": "in_progress", "output": []}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_$UNIQUE_ID", + "output_index": 0, + "content_index": 0, + "delta": REPLY, + }, + {"type": "response.completed", "sequence_number": 2, "response": _RESPONSE}, +) + + +@dataclass(frozen=True, slots=True) +class BudgetRig: + gateway: Gateway + peer: Gateway + redis: Redis + chat: str + responses: str + responses_stream: str + + @property + def replicas(self) -> tuple[Gateway, Gateway]: + return (self.gateway, self.peer) + + +@dataclass(frozen=True, slots=True) +class Window: + start_key: str + counter_key: str + + +@dataclass(frozen=True, slots=True) +class MonthlyWindow: + months: int + started: datetime + + @property + def period(self) -> str: + return f"{self.months}mo" + + @property + def ends(self) -> datetime: + return months_later(self.started, self.months) + + def differs_from_a_window_opened_at(self, now: datetime) -> bool: + return self.ends - self.started != months_later(now, self.months) - now + + def open_at(self, now: datetime) -> bool: + return self.ends > now + timedelta(hours=1) + + +def months_later(moment: datetime, months: int) -> datetime: + year, month_index = divmod(moment.year * 12 + moment.month - 1 + months, 12) + return moment.replace( + year=year, month=month_index + 1, day=min(moment.day, calendar.monthrange(year, month_index + 1)[1]) + ) + + +def monthly_window_in_another_span() -> MonthlyWindow: + now: Final = datetime.fromtimestamp(math.floor(time.time()), tz=timezone.utc) + candidates: Final = ( + MonthlyWindow(months, now - timedelta(days=days_back)) for months, days_back in product((1, 2, 3), range(1, 93)) + ) + return next(window for window in candidates if window.differs_from_a_window_opened_at(now) and window.open_at(now)) + + +def hashed(key: str) -> str: + return sha256(key.encode()).hexdigest() + + +def key_window(key: str, model: str, period: str) -> Window: + return Window( + start_key=f"virtual_key_budget_start_time:{hashed(key)}:{model}:{period}", + counter_key=f"virtual_key_spend:{hashed(key)}:{model}:{period}", + ) + + +def budget(model: str, period: str, limit: float = 100) -> dict[str, JsonValue]: + return {model: {"budget_limit": limit, "time_period": period}} + + +def unused(limit: float, period: str) -> dict[str, JsonValue]: + return {"current_spend": 0.0, "budget_limit": limit, "time_period": period, "budget_reset_at": None} + + +def reset_at(started_at: float, seconds: float) -> str: + return datetime.fromtimestamp(started_at + seconds, tz=timezone.utc).isoformat() + + +def redis_text(redis: Redis, key: str) -> str | None: + return REDIS_TEXT.validate_python(redis.get(key)) + + +def redis_float(redis: Redis, key: str) -> float: + return float(str(redis_text(redis, key))) + + +def opened_at(redis: Redis, window: Window) -> float: + started: Final = eventually( + lambda: (redis_text(redis, window.start_key), redis_text(redis, window.counter_key)), + lambda found: found[0] is not None and found[1] is not None and float(found[1]) > 0, + seconds=30, + ) + return float(str(started[0])) + + +def expires_at(redis: Redis, key: str) -> float: + remaining: Final = REDIS_INT.validate_python(redis.pttl(key)) + assert remaining > 0, f"{key} has no expiry: {remaining}" + return time.time() + remaining / 1000 + + +def key_usage(candidate: Gateway, key: str) -> Mapping[str, JsonValue]: + return object_value(object_value(candidate.get("/key/info", {"key": key})["info"])["model_max_budget_usage"]) + + +def model_usage(candidate: Gateway, key: str, model: str) -> Mapping[str, JsonValue]: + return object_value(key_usage(candidate, key)[model]) + + +def resets_on(candidates: tuple[Gateway, ...], key: str, model: str) -> tuple[JsonValue, ...]: + return tuple(model_usage(candidate, key, model)["budget_reset_at"] for candidate in candidates) + + +def reported_everywhere(candidates: tuple[Gateway, ...], key: str, model: str, expected: str | None) -> None: + eventually( + lambda: resets_on(candidates, key, model), + lambda found: found == (expected,) * len(candidates), + seconds=30, + ) + + +def assert_spend_matches_counter(rig: BudgetRig, key: str, model: str, window: Window) -> None: + eventually( + lambda: ( + tuple(model_usage(candidate, key, model)["current_spend"] for candidate in rig.replicas), + round(redis_float(rig.redis, window.counter_key), 4), + ), + lambda found: found[0] == (found[1],) * len(rig.replicas) and found[1] > 0, + seconds=30, + ) + + +def assert_daily_window_reported(rig: BudgetRig, key: str, model: str) -> float: + window: Final = key_window(key, model, "1d") + started_at: Final = opened_at(rig.redis, window) + reported_everywhere(rig.replicas, key, model, reset_at(started_at, DAY)) + assert_spend_matches_counter(rig, key, model, window) + return started_at + + +def assert_window_kept(rig: BudgetRig, key: str, model: str, started_at: float) -> None: + window: Final = key_window(key, model, "1d") + assert redis_text(rig.redis, window.start_key) == json.dumps(started_at), model + reported_everywhere(rig.replicas, key, model, reset_at(started_at, DAY)) + assert_spend_matches_counter(rig, key, model, window) + + +def assert_charged_once(rig: BudgetRig, key: str, marker: str, response_id: str | None) -> None: + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1 + row: Final = landed_once(key, marker) + assert response_id is None or row["request_id"] == response_id, row + assert float(str(row["spend"])) > 0, row + + +def responses_group(scenario: Scenario, gateway: Gateway, name: str, response: JsonResponse | SseResponse) -> str: + handle: Final = register_scenario(name, response, control_url=gateway.upstream_url) + scenario.cleanups.callback(delete_scenario, handle) + return scenario.model(api_base=f"{handle.api_base()}/v1", **PRICES) + + +def fresh_responses(candidate: Gateway, model: str, key: str, marker: str, stream: bool) -> httpx.Response: + return httpx.post( + f"{base_url(candidate)}/v1/responses", + json={"model": model, "input": marker, "stream": stream}, + headers={"Authorization": f"Bearer {key}", **spend_marker(marker)}, + timeout=60, + trust_env=False, + ) + + +def settle_served_on(rig: BudgetRig, candidate: Gateway) -> None: + admin: Final = rig.gateway.key + settle(lambda marker: fresh_chat(candidate, rig.chat, admin, marker), 200, burst=8) + settle(lambda marker: fresh_responses(candidate, rig.responses, admin, marker, False), 200, burst=8) + settle(lambda marker: fresh_responses(candidate, rig.responses_stream, admin, marker, True), 200, burst=8) + + +@pytest.fixture(scope="module") +def rig() -> Iterator[BudgetRig]: + suffix: Final = uuid.uuid4().hex[:12] + with ( + gateway_from_environment() as gateway, + httpx.Client(base_url=os.environ["INTEGRATION_PEER_URL"], timeout=15, trust_env=False) as peer_client, + Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), decode_responses=True) as redis, + gateway.scenario() as scenario, + ): + built: Final = BudgetRig( + gateway=gateway, + peer=Gateway(peer_client, gateway.key, gateway.upstream_url), + redis=redis, + chat=scenario.model(**PRICES), + responses=responses_group( + scenario, + gateway, + f"reset-at-json-{suffix}", + JsonResponse(content_type="application/json", body=_RESPONSE), + ), + responses_stream=responses_group( + scenario, + gateway, + f"reset-at-sse-{suffix}", + SseResponse( + content_type="text/event-stream", + frames=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}" for event in _RESPONSE_EVENTS), + ), + ), + ) + for candidate in built.replicas: + settle_served_on(built, candidate) + yield built + + +def test_reset_is_null_on_both_replicas_until_the_first_charge(rig: BudgetRig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=budget(rig.chat, "1d")) + assert [model_usage(candidate, key, rig.chat) for candidate in rig.replicas] == [unused(100.0, "1d")] * 2 + + +def test_openai_chat_charge_reports_the_reset_on_key_info_and_v2_key_info(rig: BudgetRig) -> None: + marker: Final = "chat-sync-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=budget(rig.chat, "1d")) + with openai_client(rig.gateway, key) as client: + completion: Final = client.chat.completions.create( + model=rig.chat, messages=[{"role": "user", "content": marker}], extra_headers=spend_marker(marker) + ) + assert completion.choices[0].message.content == CHAT_REPLY, completion + assert_charged_once(rig, key, marker, completion.id) + started_at: Final = assert_daily_window_reported(rig, key, rig.chat) + listed: Final = rig.peer.post("/v2/key/info", {"keys": [key]})["info"] + assert isinstance(listed, list) and len(listed) == 1, listed + entry: Final = object_value(object_value(object_value(listed[0])["model_max_budget_usage"])[rig.chat]) + assert entry["budget_reset_at"] == reset_at(started_at, DAY), entry + + +async def test_streamed_openai_chat_charge_reports_the_reset(rig: BudgetRig) -> None: + marker: Final = "chat-stream-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=budget(rig.chat, "1d")) + async with async_openai_client(rig.gateway, key) as client: + stream: Final = await client.chat.completions.create( + model=rig.chat, + messages=[{"role": "user", "content": marker}], + stream=True, + stream_options={"include_usage": True}, + extra_headers=spend_marker(marker), + ) + chunks: Final = tuple([chunk async for chunk in stream]) + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == CHAT_REPLY + assert_charged_once(rig, key, marker, chunks[0].id) + assert_daily_window_reported(rig, key, rig.chat) + + +def test_anthropic_messages_charge_reports_the_reset(rig: BudgetRig) -> None: + marker: Final = "messages-sync-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=budget(rig.responses, "1d")) + with anthropic_client(rig.gateway, key) as client: + message: Final = client.messages.create( + model=rig.responses, + max_tokens=16, + messages=[{"role": "user", "content": marker}], + extra_headers=spend_marker(marker), + ) + assert [block.text for block in message.content if block.type == "text"] == [REPLY], message + assert_charged_once(rig, key, marker, message.id) + assert_daily_window_reported(rig, key, rig.responses) + + +async def test_streamed_anthropic_messages_charge_reports_the_reset(rig: BudgetRig) -> None: + marker: Final = "messages-stream-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=budget(rig.responses_stream, "1d")) + async with ( + async_anthropic_client(rig.gateway, key) as client, + client.messages.stream( + model=rig.responses_stream, + max_tokens=16, + messages=[{"role": "user", "content": marker}], + extra_headers=spend_marker(marker), + ) as stream, + ): + text: Final = "".join([piece async for piece in stream.text_stream]) + final: Final = await stream.get_final_message() + assert text == REPLY, final + assert_charged_once(rig, key, marker, final.id) + assert_daily_window_reported(rig, key, rig.responses_stream) + + +def test_openai_responses_charge_reports_the_reset(rig: BudgetRig) -> None: + marker: Final = "responses-sync-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=budget(rig.responses, "1d")) + with openai_client(rig.gateway, key) as client: + response: Final = client.responses.create( + model=rig.responses, input=marker, extra_headers=spend_marker(marker) + ) + assert response.output_text == REPLY, response + assert_charged_once(rig, key, marker, response.id) + assert_daily_window_reported(rig, key, rig.responses) + + +async def test_streamed_openai_responses_charge_reports_the_reset(rig: BudgetRig) -> None: + marker: Final = "responses-stream-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=budget(rig.responses_stream, "1d")) + async with async_openai_client(rig.gateway, key) as client: + stream: Final = await client.responses.create( + model=rig.responses_stream, input=marker, stream=True, extra_headers=spend_marker(marker) + ) + events: Final = tuple([event async for event in stream]) + assert events[-1].type == "response.completed", events + assert "".join(event.delta for event in events if event.type == "response.output_text.delta") == REPLY + assert_charged_once(rig, key, marker, None) + assert_daily_window_reported(rig, key, rig.responses_stream) + + +def test_raw_chat_on_the_peer_reports_the_reset_on_both_replicas(rig: BudgetRig) -> None: + marker: Final = "raw-peer-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=budget(rig.chat, "1d")) + response: Final = fresh_chat(rig.peer, rig.chat, key, marker) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == CHAT_REPLY, response.text + assert_charged_once(rig, key, marker, response.json()["id"]) + assert_daily_window_reported(rig, key, rig.chat) + + +def user_resets(rig: BudgetRig, user: str) -> tuple[JsonValue, ...]: + def entry(usage: JsonValue) -> JsonValue: + return object_value(object_value(usage)[rig.chat])["budget_reset_at"] + + return ( + *( + entry(object_value(candidate.get("/user/info", {"user_id": user})["user_info"])["model_max_budget_usage"]) + for candidate in rig.replicas + ), + entry(rig.peer.get("/v2/user/info", {"user_id": user})["model_max_budget_usage"]), + ) + + +def test_user_info_and_v2_user_info_report_the_user_window_reset(rig: BudgetRig) -> None: + marker: Final = "user-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + user: Final = scenario.user(model_max_budget=budget(rig.chat, "1d")) + key: Final = scenario.key(user_id=user) + response: Final = fresh_chat(rig.gateway, rig.chat, key, marker) + assert response.status_code == 200, response.text + assert_charged_once(rig, key, marker, response.json()["id"]) + started_at: Final = opened_at( + rig.redis, + Window( + start_key=f"user_model_budget_start_time:{user}:{rig.chat}:1d", + counter_key=f"user_model_spend:{user}:{rig.chat}:1d", + ), + ) + expected: Final = reset_at(started_at, DAY) + eventually(lambda: user_resets(rig, user), lambda found: found == (expected,) * 3, seconds=30) + + +def team_resets(rig: BudgetRig, team: str) -> tuple[JsonValue, ...]: + return tuple( + object_value( + object_value( + object_value(candidate.get("/team/info", {"team_id": team})["team_info"])["model_max_budget_usage"] + )[rig.chat] + )["budget_reset_at"] + for candidate in rig.replicas + ) + + +def test_team_info_reports_the_team_window_reset(rig: BudgetRig) -> None: + marker: Final = "team-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(model_max_budget=budget(rig.chat, "1h")) + key: Final = scenario.key(team_id=team) + response: Final = fresh_chat(rig.gateway, rig.chat, key, marker) + assert response.status_code == 200, response.text + assert_charged_once(rig, key, marker, response.json()["id"]) + started_at: Final = opened_at( + rig.redis, + Window( + start_key=f"team_model_budget_start_time:{team}:{rig.chat}:1h", + counter_key=f"team_model_spend:{team}:{rig.chat}:1h", + ), + ) + expected: Final = reset_at(started_at, HOUR) + eventually(lambda: team_resets(rig, team), lambda found: found == (expected,) * 2, seconds=30) + + +@dataclass(frozen=True, slots=True) +class Attempt: + at: float + status: int + + +def attempt(candidate: Gateway, model: str, key: str) -> Attempt: + response: Final = fresh_chat(candidate, model, key, "after-reset-" + uuid.uuid4().hex) + return Attempt(at=time.time(), status=response.status_code) + + +def test_a_used_up_minute_budget_is_refused_until_the_reported_reset(rig: BudgetRig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=budget(rig.chat, "1m", limit=0.000000001)) + first: Final = fresh_chat(rig.gateway, rig.chat, key, "spend-" + uuid.uuid4().hex) + assert first.status_code == 200, first.text + started_at: Final = opened_at(rig.redis, key_window(key, rig.chat, "1m")) + resets: Final = started_at + 60 + reported_everywhere(rig.replicas, key, rig.chat, reset_at(started_at, 60)) + refused: Final = eventually( + lambda: tuple(attempt(candidate, rig.chat, key) for candidate in rig.replicas), + lambda found: all(tried.status == BUDGET_EXCEEDED for tried in found) and found[0].at >= resets - 2, + seconds=90, + ) + allowed: Final = eventually( + lambda: attempt(rig.gateway, rig.chat, key), + lambda found: found.status != BUDGET_EXCEEDED, + seconds=30, + ) + assert allowed.status == 200, allowed + assert resets - 1.5 <= allowed.at <= resets + 5, (allowed, resets, refused) + + +def config_with(directory: Path, name: str, **sections: JsonValue) -> Path: + loaded: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + path: Final = directory / name + path.write_text(yaml.safe_dump({**loaded, **sections})) + return path + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS) +def test_without_redis_the_reset_comes_from_the_worker_memory(rig: BudgetRig, tmp_path: Path) -> None: + with ( + rig.gateway.scenario() as scenario, + owned_proxy( + rig.gateway, + tmp_path, + {}, + config=config_with(tmp_path, "no_redis_config.yaml", litellm_settings={}), + remove_environment=("REDIS_HOST", "REDIS_PORT"), + ) as owned, + ): + settle(lambda marker: fresh_chat(owned, rig.chat, owned.key, marker), 200, burst=4) + key: Final = scenario.key(model_max_budget=budget(rig.chat, "1d")) + before: Final = time.time() + response: Final = fresh_chat(owned, rig.chat, key, "no-redis-" + uuid.uuid4().hex) + assert response.status_code == 200, response.text + reported: Final = eventually( + lambda: model_usage(owned, key, rig.chat)["budget_reset_at"], + lambda found: found is not None, + seconds=30, + ) + resets: Final = datetime.fromisoformat(str(reported)).timestamp() + assert before + DAY <= resets <= time.time() + DAY, (before, reported) + assert redis_text(rig.redis, key_window(key, rig.chat, "1d").start_key) is None + + +def assert_cache_hit_kept_the_window(rig: BudgetRig, key: str, window: Window, started_at: float) -> None: + spent: Final = redis_float(rig.redis, window.counter_key) + held(lambda: redis_float(rig.redis, window.counter_key), lambda found: found == spent, holding=3) + assert redis_text(rig.redis, window.start_key) == json.dumps(started_at) + assert resets_on(rig.replicas, key, rig.chat) == (reset_at(started_at, DAY),) * 2 + + +def test_a_cached_chat_reply_keeps_the_window_start_and_reset(rig: BudgetRig) -> None: + content: Final = "cache-twin-" + uuid.uuid4().hex + first_marker: Final = "first-" + uuid.uuid4().hex + second_marker: Final = "second-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=budget(rig.chat, "1d")) + with openai_client(rig.gateway, key) as client: + first: Final = client.chat.completions.create( + model=rig.chat, + messages=[{"role": "user", "content": content}], + extra_headers=spend_marker(first_marker), + ) + landed_once(key, first_marker) + started_at: Final = assert_daily_window_reported(rig, key, rig.chat) + with openai_client(rig.peer, key) as client: + second: Final = client.chat.completions.with_raw_response.create( + model=rig.chat, + messages=[{"role": "user", "content": content}], + extra_headers=spend_marker(second_marker), + ) + assert second.parse().id == first.id, second.parse() + assert "x-litellm-cache-key" in second.headers, dict(second.headers) + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), content) == 1 + assert float(str(landed_once(key, second_marker)["spend"])) == 0.0 + assert_cache_hit_kept_the_window(rig, key, key_window(key, rig.chat, "1d"), started_at) + + +async def streamed_chat(candidate: Gateway, model: str, key: str, content: str, marker: str) -> tuple[str, str, bool]: + async with async_openai_client(candidate, key) as client: + raw: Final = await client.chat.completions.with_raw_response.create( + model=model, + messages=[{"role": "user", "content": content}], + stream=True, + extra_headers=spend_marker(marker), + ) + chunks: Final = tuple([chunk async for chunk in raw.parse()]) + text: Final = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) + return chunks[0].id, text, "x-litellm-cache-key" in raw.headers + + +async def test_a_cached_streamed_chat_reply_keeps_the_window_start_and_reset(rig: BudgetRig) -> None: + content: Final = "stream-cache-twin-" + uuid.uuid4().hex + first_marker: Final = "first-" + uuid.uuid4().hex + second_marker: Final = "second-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=budget(rig.chat, "1d")) + first: Final = await streamed_chat(rig.gateway, rig.chat, key, content, first_marker) + assert first[1] == CHAT_REPLY, first + landed_once(key, first_marker) + started_at: Final = assert_daily_window_reported(rig, key, rig.chat) + second: Final = await streamed_chat(rig.peer, rig.chat, key, content, second_marker) + assert second == (first[0], CHAT_REPLY, True), (first, second) + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), content) == 1 + assert float(str(landed_once(key, second_marker)["spend"])) == 0.0 + assert_cache_hit_kept_the_window(rig, key, key_window(key, rig.chat, "1d"), started_at) + + +MALFORMED: Final = ( + pytest.param("time_period", 5, id="period-int"), + pytest.param("time_period", ["1d"], id="period-list"), + pytest.param("time_period", "", id="period-empty"), + pytest.param("time_period", "9" * 5119 + "h", id="period-5kb-hours"), + pytest.param("time_period", "x" * 5120, id="period-5kb-text"), + pytest.param("time_period", None, id="period-null"), + pytest.param("budget_limit", ["1"], id="limit-list"), + pytest.param("budget_limit", "", id="limit-empty"), + pytest.param("budget_limit", "x" * 5120, id="limit-5kb"), +) + + +def reported_reset(info: httpx.Response, model: str) -> JsonValue: + usage: Final = object_value(info.json()["info"]).get("model_max_budget_usage") + entry: Final = usage.get(model) if isinstance(usage, dict) else None + return entry.get("budget_reset_at") if isinstance(entry, dict) else None + + +def reported_resets(rig: BudgetRig, key: str) -> tuple[tuple[int, JsonValue], ...]: + infos: Final = tuple(candidate.request("GET", "/key/info", params={"key": key}) for candidate in rig.replicas) + return tuple( + (info.status_code, reported_reset(info, rig.chat) if info.status_code == 200 else None) for info in infos + ) + + +def store_budget(key: str, model_max_budget: Mapping[str, JsonValue]) -> None: + write_rows( + 'UPDATE "LiteLLM_VerificationToken" SET model_max_budget = %s::jsonb WHERE token = %s', + (json.dumps(model_max_budget), hashed(key)), + ) + + +def charged_with_stored_entry(rig: BudgetRig, scenario: Scenario, entry: Mapping[str, JsonValue]) -> str: + key: Final = scenario.key(model_max_budget=budget(rig.chat, "1d")) + store_budget(key, {rig.chat: dict(entry)}) + marker: Final = "stored-entry-" + uuid.uuid4().hex + response: Final = fresh_chat(rig.gateway, rig.chat, key, marker) + assert response.status_code == 200, response.text + landed_once(key, marker) + return key + + +@pytest.mark.parametrize(("field", "value"), MALFORMED) +def test_a_stored_malformed_budget_entry_never_breaks_reporting_or_other_keys( + rig: BudgetRig, field: str, value: JsonValue +) -> None: + with rig.gateway.scenario() as scenario: + control: Final = scenario.key() + key: Final = charged_with_stored_entry(rig, scenario, {"budget_limit": 100, "time_period": "1d", field: value}) + held(lambda: reported_resets(rig, key), lambda found: found == ((200, None),) * 2, holding=3) + listed: Final = rig.peer.request("POST", "/v2/key/info", {"keys": [key]}) + assert listed.status_code == 200, listed.text + served: Final = fresh_chat(rig.peer, rig.chat, control, "control-" + uuid.uuid4().hex) + assert served.status_code == 200, served.text + + +def test_a_stored_entry_with_a_null_limit_still_reports_its_window_reset(rig: BudgetRig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = charged_with_stored_entry(rig, scenario, {"budget_limit": None, "time_period": "1d"}) + assert_daily_window_reported(rig, key, rig.chat) + + +def test_the_last_of_two_entries_for_one_model_wins(rig: BudgetRig) -> None: + first: Final = json.dumps({"budget_limit": 1, "time_period": "1h"}) + last: Final = json.dumps({"budget_limit": 2, "time_period": "1d"}) + model: Final = json.dumps(rig.chat) + with rig.gateway.scenario() as scenario: + created: Final = rig.gateway.client.post( + "/key/generate", + content=f'{{"model_max_budget": {{{model}: {first}, {model}: {last}}}}}'.encode(), + headers={"Authorization": f"Bearer {rig.gateway.key}", "Content-Type": "application/json"}, + ) + assert created.status_code == 200, created.text + key: Final = str(created.json()["key"]) + scenario.cleanups.callback(scenario.delete_key, key) + assert [model_usage(candidate, key, rig.chat) for candidate in rig.replicas] == [unused(2.0, "1d")] * 2 + response: Final = fresh_chat(rig.gateway, rig.chat, key, "duplicate-" + uuid.uuid4().hex) + assert response.status_code == 200, response.text + assert_daily_window_reported(rig, key, rig.chat) + + +def test_a_window_ending_after_year_9999_reports_no_reset(rig: BudgetRig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=budget(rig.chat, "99999999h")) + response: Final = fresh_chat(rig.gateway, rig.chat, key, "far-" + uuid.uuid4().hex) + assert response.status_code == 200, response.text + window: Final = key_window(key, rig.chat, "99999999h") + opened_at(rig.redis, window) + assert_spend_matches_counter(rig, key, rig.chat, window) + assert resets_on(rig.replicas, key, rig.chat) == (None, None) + + +def test_a_request_to_another_model_opens_no_window(rig: BudgetRig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=budget(rig.responses, "1d")) + marker: Final = "other-model-" + uuid.uuid4().hex + response: Final = fresh_chat(rig.gateway, rig.chat, key, marker) + assert response.status_code == 200, response.text + landed_once(key, marker) + assert [model_usage(candidate, key, rig.responses) for candidate in rig.replicas] == [unused(100.0, "1d")] * 2 + assert redis_text(rig.redis, key_window(key, rig.responses, "1d").start_key) is None + + +def seed_start(redis: Redis, start_key: str, started_at: float, ends_at: float) -> None: + assert redis.set(start_key, json.dumps(started_at), ex=math.ceil(ends_at - time.time())) + + +def charge_on_both_replicas(rig: BudgetRig, key: str, markers: tuple[str, ...]) -> None: + statuses: Final = tuple( + fresh_chat(rig.replicas[index % 2], rig.chat, key, marker).status_code for index, marker in enumerate(markers) + ) + assert statuses == (200,) * len(markers), statuses + + +def test_a_monthly_window_ends_one_calendar_month_after_its_seeded_start(rig: BudgetRig) -> None: + monthly: Final = monthly_window_in_another_span() + started_at: Final = monthly.started.timestamp() + ends_at: Final = monthly.ends.timestamp() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=budget(rig.chat, monthly.period)) + window: Final = key_window(key, rig.chat, monthly.period) + seed_start(rig.redis, window.start_key, started_at, ends_at) + first: Final = "monthly-" + uuid.uuid4().hex + charge_on_both_replicas(rig, key, (first,)) + spent: Final = float(str(landed_once(key, first)["spend"])) + eventually( + lambda: redis_float(rig.redis, window.counter_key), + lambda found: math.isclose(found, spent, rel_tol=1e-9), + seconds=30, + ) + assert abs(expires_at(rig.redis, window.counter_key) - ends_at) < 5, (monthly, ends_at) + assert resets_on(rig.replicas, key, rig.chat) == (monthly.ends.isoformat(),) * 2 + markers: Final = tuple(f"monthly-again-{index}-{uuid.uuid4().hex}" for index in range(3)) + charge_on_both_replicas(rig, key, markers) + total: Final = spent + sum(float(str(row["spend"])) for row in landed_all_once(key, frozenset(markers))) + eventually( + lambda: redis_float(rig.redis, window.counter_key), + lambda found: math.isclose(found, total, rel_tol=1e-9), + seconds=30, + ) + assert redis_text(rig.redis, window.start_key) == json.dumps(started_at) + assert abs(expires_at(rig.redis, window.counter_key) - ends_at) < 5, (monthly, ends_at) + reported_everywhere(rig.replicas, key, rig.chat, monthly.ends.isoformat()) + + +def test_two_budgets_on_one_key_each_report_their_own_reset(rig: BudgetRig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget={**budget(rig.chat, "1d"), **budget(rig.responses, "1h")}) + chat: Final = fresh_chat(rig.gateway, rig.chat, key, "daily-" + uuid.uuid4().hex) + assert chat.status_code == 200, chat.text + hourly: Final = fresh_responses(rig.peer, rig.responses, key, "hourly-" + uuid.uuid4().hex, False) + assert hourly.status_code == 200, hourly.text + daily_start: Final = opened_at(rig.redis, key_window(key, rig.chat, "1d")) + hourly_start: Final = opened_at(rig.redis, key_window(key, rig.responses, "1h")) + reported_everywhere(rig.replicas, key, rig.chat, reset_at(daily_start, DAY)) + reported_everywhere(rig.replicas, key, rig.responses, reset_at(hourly_start, HOUR)) + + +def test_an_ended_window_reports_no_reset_until_the_next_charge_opens_one(rig: BudgetRig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=budget(rig.chat, "1m")) + window: Final = key_window(key, rig.chat, "1m") + ended_start: Final = math.floor(time.time()) - 120.0 + seed_start(rig.redis, window.start_key, ended_start, time.time() + HOUR) + assert resets_on(rig.replicas, key, rig.chat) == (None, None) + response: Final = fresh_chat(rig.gateway, rig.chat, key, "reopen-" + uuid.uuid4().hex) + assert response.status_code == 200, response.text + reopened: Final = eventually( + lambda: float(redis_text(rig.redis, window.start_key) or ended_start), + lambda found: found > ended_start, + seconds=30, + ) + reported_everywhere(rig.replicas, key, rig.chat, reset_at(reopened, 60)) + + +BURST_KINDS: Final[tuple[str, ...]] = ( + "chat", + "chat-stream", + "responses", + "responses-stream", + "messages", + "messages-stream", +) + + +def burst_body(rig: BudgetRig, kind: str, marker: str) -> tuple[str, dict[str, JsonValue]]: + stream: Final = kind.endswith("-stream") + match kind.removesuffix("-stream"): + case "chat": + return "/v1/chat/completions", { + "model": rig.chat, + "messages": [{"role": "user", "content": marker}], + "stream": stream, + } + case "responses": + return "/v1/responses", { + "model": rig.responses_stream if stream else rig.responses, + "input": marker, + "stream": stream, + } + case _: + return "/v1/messages", { + "model": rig.responses_stream if stream else rig.responses, + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + "stream": stream, + } + + +async def burst_request(client: httpx.AsyncClient, rig: BudgetRig, kind: str, key: str, marker: str) -> int: + path, body = burst_body(rig, kind, marker) + async with client.stream( + "POST", path, json=body, headers={"Authorization": f"Bearer {key}", **spend_marker(marker)} + ) as response: + await response.aread() + return response.status_code + + +async def mixed_burst(rig: BudgetRig, key: str, markers: tuple[str, ...]) -> tuple[int, ...]: + async with ( + httpx.AsyncClient(base_url=base_url(rig.gateway), timeout=90, trust_env=False) as gateway, + httpx.AsyncClient(base_url=base_url(rig.peer), timeout=90, trust_env=False) as peer, + ): + return tuple( + await asyncio.gather( + *( + burst_request((gateway, peer)[index % 2], rig, BURST_KINDS[index % len(BURST_KINDS)], key, marker) + for index, marker in enumerate(markers) + ) + ) + ) + + +async def test_a_mixed_burst_over_both_replicas_lands_once_and_keeps_one_reset(rig: BudgetRig) -> None: + with rig.gateway.scenario() as scenario: + models: Final = (rig.chat, rig.responses, rig.responses_stream) + key: Final = scenario.key(model_max_budget={model: budget(model, "1d")[model] for model in models}) + opening: Final = ( + fresh_chat(rig.gateway, rig.chat, key, "open-" + uuid.uuid4().hex), + fresh_responses(rig.peer, rig.responses, key, "open-" + uuid.uuid4().hex, False), + fresh_responses(rig.gateway, rig.responses_stream, key, "open-" + uuid.uuid4().hex, True), + ) + assert [response.status_code for response in opening] == [200] * 3, [response.text for response in opening] + starts: Final = {model: opened_at(rig.redis, key_window(key, model, "1d")) for model in models} + markers: Final = tuple(f"burst-{index}-{uuid.uuid4().hex}" for index in range(30)) + statuses: Final = await mixed_burst(rig, key, markers) + assert statuses == (200,) * len(markers), statuses + landed_all_once(key, frozenset(markers)) + for model in models: + assert_window_kept(rig, key, model, starts[model]) + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS) +def test_a_killed_worker_leaves_the_reset_and_start_unchanged(rig: BudgetRig, tmp_path: Path) -> None: + with ( + rig.gateway.scenario() as scenario, + owned_proxy_process(rig.gateway, tmp_path, {}, workers=2) as owned, + ): + settle(lambda marker: fresh_chat(owned.gateway, rig.chat, owned.gateway.key, marker), 200, burst=8) + workers: Final = eventually(owned.worker_pids, lambda found: len(found) == 2, seconds=30) + key: Final = scenario.key(model_max_budget=budget(rig.chat, "1d")) + first: Final = "before-kill-" + uuid.uuid4().hex + assert fresh_chat(owned.gateway, rig.chat, key, first).status_code == 200 + landed_once(key, first) + started_at: Final = opened_at(rig.redis, key_window(key, rig.chat, "1d")) + expected: Final = reset_at(started_at, DAY) + reported_everywhere((owned.gateway, *rig.replicas), key, rig.chat, expected) + markers: Final = tuple(f"during-kill-{index}-{uuid.uuid4().hex}" for index in range(6)) + with owned.replacing_worker(min(workers)) as kill: + kill() + statuses: Final = tuple(fresh_chat(owned.gateway, rig.chat, key, marker).status_code for marker in markers) + assert statuses == (200,) * len(markers), statuses + landed_all_once(key, frozenset(markers)) + assert_window_kept(rig, key, rig.chat, started_at) + assert {resets_on((owned.gateway,), key, rig.chat) for _ in range(8)} == {(expected,)} + + +ROUTER_TAG: Final = "monthly-tag-" + uuid.uuid4().hex[:12] +ROUTER_DEPLOYMENT: Final = "monthly-deployment-" + uuid.uuid4().hex[:12] +ROUTER_MODEL: Final = "monthly-router-" + uuid.uuid4().hex[:12] +ROUTER_START_KEYS: Final[tuple[str, ...]] = ( + "provider_budget_start_time:openai", + f"deployment_budget_start_time:{ROUTER_DEPLOYMENT}", + f"tag_budget_start_time:{ROUTER_TAG}", +) + + +@dataclass(frozen=True, slots=True) +class RouterRig: + gateway: Gateway + redis: Redis + monthly: MonthlyWindow + + @property + def counters(self) -> tuple[str, ...]: + period: Final = self.monthly.period + return ( + f"provider_spend:openai:{period}", + f"deployment_spend:{ROUTER_DEPLOYMENT}:{period}", + f"tag_spend:{ROUTER_TAG}:{period}", + ) + + +def router_config(directory: Path, upstream_url: str, period: str) -> Path: + limits: Final[dict[str, JsonValue]] = {"budget_limit": 100, "time_period": period} + return config_with( + directory, + "router_budget_config.yaml", + model_list=[ + { + "model_name": ROUTER_MODEL, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{upstream_url}/v1", + "api_key": "integration-provider-key", + "max_budget": 100, + "budget_duration": period, + **PRICES, + }, + "model_info": {"id": ROUTER_DEPLOYMENT}, + } + ], + router_settings={"disable_cooldowns": True, "provider_budget_config": {"openai": limits}}, + litellm_settings={"tag_budget_config": {ROUTER_TAG: limits}}, + general_settings={ + **object_value( + JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))[ + "general_settings" + ] + ), + "coordination_redis": {"host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"}, + }, + ) + + +def tagged_chat(candidate: Gateway, key: str, marker: str) -> httpx.Response: + return fresh_chat(candidate, ROUTER_MODEL, key, marker, {"metadata": {"tags": [ROUTER_TAG]}}) + + +@pytest.fixture(scope="module") +def router_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[RouterRig]: + directory: Final = tmp_path_factory.mktemp("router-budgets") + monthly: Final = monthly_window_in_another_span() + with ( + gateway_from_environment() as gateway, + owned_redis(directory / "redis") as served, + Redis(host=served.host, port=served.port, decode_responses=True) as redis, + ): + for start_key in ROUTER_START_KEYS: + seed_start(redis, start_key, monthly.started.timestamp(), monthly.ends.timestamp()) + with owned_proxy( + gateway, + directory, + {"REDIS_HOST": served.host, "REDIS_PORT": str(served.port)}, + config=router_config(directory, gateway.upstream_url, monthly.period), + ) as owned: + settle(lambda marker: tagged_chat(owned, owned.key, marker), 200, burst=2) + yield RouterRig(gateway=owned, redis=redis, monthly=monthly) + + +def router_counters_charged(rig: RouterRig) -> tuple[float, ...]: + return eventually( + lambda: tuple(float(redis_text(rig.redis, counter) or 0) for counter in rig.counters), + lambda found: all(spent > 0 for spent in found), + seconds=30, + ) + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS) +def test_router_budget_counters_expire_one_calendar_month_after_their_seeded_start(router_rig: RouterRig) -> None: + router_counters_charged(router_rig) + ends_at: Final = router_rig.monthly.ends.timestamp() + expiries: Final = tuple(expires_at(router_rig.redis, counter) for counter in router_rig.counters) + assert all(abs(expiry - ends_at) < 5 for expiry in expiries), (router_rig.monthly, expiries) + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS) +def test_provider_budgets_report_the_reset_one_calendar_month_after_the_seeded_start(router_rig: RouterRig) -> None: + router_counters_charged(router_rig) + provider: Final = object_value(object_value(router_rig.gateway.get("/provider/budgets")["providers"])["openai"]) + reported: Final = datetime.fromisoformat(str(provider["budget_reset_at"])).timestamp() + assert abs(reported - router_rig.monthly.ends.timestamp()) < 5, (router_rig.monthly, provider) + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS) +def test_router_budget_starts_stay_put_across_more_charges(router_rig: RouterRig) -> None: + before: Final = router_counters_charged(router_rig) + markers: Final = tuple(f"router-again-{index}-{uuid.uuid4().hex}" for index in range(3)) + statuses: Final = tuple( + tagged_chat(router_rig.gateway, router_rig.gateway.key, marker).status_code for marker in markers + ) + assert statuses == (200,) * 3, statuses + eventually( + lambda: tuple(float(redis_text(router_rig.redis, counter) or 0) for counter in router_rig.counters), + lambda found: all(spent > earlier for spent, earlier in zip(found, before, strict=True)), + seconds=30, + ) + started: Final = json.dumps(router_rig.monthly.started.timestamp()) + assert tuple(redis_text(router_rig.redis, start_key) for start_key in ROUTER_START_KEYS) == (started,) * 3 + ends_at: Final = router_rig.monthly.ends.timestamp() + expiries: Final = tuple(expires_at(router_rig.redis, counter) for counter in router_rig.counters) + assert all(abs(expiry - ends_at) < 5 for expiry in expiries), (router_rig.monthly, expiries) diff --git a/tests/unit/litellm_core_utils/test_duration_parser.py b/tests/unit/litellm_core_utils/test_duration_parser.py index f5e50d67e71..f14015effe6 100644 --- a/tests/unit/litellm_core_utils/test_duration_parser.py +++ b/tests/unit/litellm_core_utils/test_duration_parser.py @@ -4,9 +4,12 @@ from typing import Final from unittest.mock import patch from zoneinfo import ZoneInfo +import pytest + import litellm.litellm_core_utils.duration_parser as duration_parser from litellm.litellm_core_utils.duration_parser import ( duration_in_seconds, + duration_in_seconds_from, get_budget_window_start, get_next_standardized_reset_time, ) @@ -373,5 +376,34 @@ class TestGetBudgetWindowStart(unittest.TestCase): ) +@pytest.mark.parametrize( + "start,duration,days", + [ + (datetime(2026, 1, 15, 12, tzinfo=timezone.utc), "1mo", 31), + (datetime(2026, 2, 10, 12, tzinfo=timezone.utc), "1mo", 28), + (datetime(2026, 1, 31, 12, tzinfo=timezone.utc), "1mo", 28), + (datetime(2028, 1, 31, 12, tzinfo=timezone.utc), "1mo", 29), + (datetime(2026, 12, 15, 12, tzinfo=timezone.utc), "1mo", 31), + (datetime(2026, 1, 15, 12, tzinfo=timezone.utc), "2mo", 59), + ], + ids=["january", "february", "clamped_to_february", "leap_february", "across_new_year", "two_months"], +) +def test_a_month_lasts_as_long_as_the_month_it_starts_in(start, duration, days): + assert duration_in_seconds_from(start, duration) == days * 86400 + + +@pytest.mark.parametrize("duration", ["45s", "30m", "1h", "1d", "7d", "2w", "30d", "daily", "monthly"]) +def test_fixed_length_units_do_not_depend_on_the_start(duration): + assert duration_in_seconds_from(datetime(2026, 2, 10, 12, tzinfo=timezone.utc), duration) == duration_in_seconds( + duration + ) + + +def test_a_month_starting_now_matches_duration_in_seconds(): + now: Final = datetime(2026, 1, 31, 9, 30, 0, 250000, tzinfo=timezone.utc).astimezone() + with patch.object(duration_parser.time_module, "time", return_value=now.timestamp()): + assert duration_in_seconds("1mo") == duration_in_seconds_from(now, "1mo") + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index dc9dac9a57b..33904637366 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py @@ -1,26 +1,37 @@ import asyncio +import time +from datetime import datetime, timedelta, timezone from types import MappingProxyType +from typing import Final from unittest.mock import AsyncMock, patch - import pytest import litellm from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache -from datetime import datetime, timezone - from litellm.litellm_core_utils.duration_parser import duration_in_seconds -from litellm.proxy._types import Litellm_EntityType +from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth from litellm.proxy.hooks.model_max_budget_limiter import ( - _budget_model_candidates, PROXY_VirtualKeyModelMaxBudgetLimiter, + _budget_model_candidates, build_model_max_budget_usage, + model_budget_spend_cache_key, + model_budget_start_time_cache_key, resolve_model_budget, ) -from litellm.proxy._types import UserAPIKeyAuth from litellm.types.caching import RedisPipelineIncrementOperation -from litellm.types.utils import BudgetConfig as GenericBudgetInfo + + +async def _reset_at_of_open_window( + cache: DualCache, entity_type: Litellm_EntityType, entity_id: str, budget_model: str, period: str +) -> str: + started_at = await cache.async_get_cache( + key=model_budget_start_time_cache_key( + entity_type=entity_type, entity_id=entity_id, budget_model=budget_model, budget_duration=period + ) + ) + return datetime.fromtimestamp(started_at + duration_in_seconds(period), tz=timezone.utc).isoformat() # Test class setup @@ -494,14 +505,10 @@ async def test_model_budget_limiter_initializes_redis_increment_queue_lock(): limiter = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) spend_key = "virtual_key_spend:test-key:gpt-4:1d" - await limiter._increment_spend_in_current_window( - spend_key=spend_key, response_cost=0.01, ttl=86400 - ) + await limiter._increment_spend_in_current_window(spend_key=spend_key, response_cost=0.01, ttl=86400) assert limiter.redis_increment_operation_queue == [ - RedisPipelineIncrementOperation( - key=spend_key, increment_value=0.01, ttl=86400 - ) + RedisPipelineIncrementOperation(key=spend_key, increment_value=0.01, ttl=86400) ] @@ -680,6 +687,9 @@ async def test_logged_spend_is_visible_to_key_info_usage_and_enforcement(request "current_spend": 0.75, "budget_limit": 1.0, "time_period": "1d", + "budget_reset_at": await _reset_at_of_open_window( + dual_cache, Litellm_EntityType.KEY, key_hash, "gpt-4", "1d" + ), } } @@ -747,7 +757,16 @@ async def test_user_model_budget_is_tracked_and_enforced(): entity_id=user_id, model_max_budget=user_model_max_budget, cache=dual_cache, - ) == {"gpt-4": {"current_spend": 1.5, "budget_limit": 1.0, "time_period": "1mo"}} + ) == { + "gpt-4": { + "current_spend": 1.5, + "budget_limit": 1.0, + "time_period": "1mo", + "budget_reset_at": await _reset_at_of_open_window( + dual_cache, Litellm_EntityType.USER, user_id, "gpt-4", "1mo" + ), + } + } with pytest.raises(litellm.BudgetExceededError) as exc: await limiter.is_user_within_model_budget( @@ -857,7 +876,9 @@ async def test_build_model_max_budget_usage_skips_unusable_entries(): }, cache=dual_cache, ) - assert usage == {"gpt-4": {"current_spend": 3.0, "budget_limit": 10.0, "time_period": "1d"}} + assert usage == { + "gpt-4": {"current_spend": 3.0, "budget_limit": 10.0, "time_period": "1d", "budget_reset_at": None} + } @pytest.mark.asyncio @@ -896,6 +917,9 @@ async def test_bedrock_traffic_charges_the_bare_family_name_budget(): "current_spend": 1.5, "budget_limit": 1.0, "time_period": "18h", + "budget_reset_at": await _reset_at_of_open_window( + dual_cache, Litellm_EntityType.KEY, key_hash, "claude-opus-4-8", "18h" + ), } } @@ -963,7 +987,16 @@ async def test_user_model_budget_window_resets_when_the_period_elapses(): entity_id=user_id, model_max_budget=user_model_max_budget, cache=dual_cache, - ) == {"gpt-4": {"current_spend": 1.5, "budget_limit": 1.0, "time_period": "1mo"}} + ) == { + "gpt-4": { + "current_spend": 1.5, + "budget_limit": 1.0, + "time_period": "1mo", + "budget_reset_at": await _reset_at_of_open_window( + dual_cache, Litellm_EntityType.USER, user_id, "gpt-4", "1mo" + ), + } + } @pytest.mark.asyncio @@ -995,7 +1028,7 @@ async def test_a_zero_dollar_cap_is_reported_as_a_cap_not_as_absent(): entity_id="hash-zero", model_max_budget={"gpt-4": {"budget_limit": 0, "time_period": "1d"}}, cache=DualCache(), - ) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 0.0, "time_period": "1d"}} + ) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 0.0, "time_period": "1d", "budget_reset_at": None}} @pytest.mark.asyncio @@ -1033,7 +1066,7 @@ async def test_usage_report_reads_every_counter_in_one_batched_lookup(): budget = {f"model-{i}": {"budget_limit": 1.0, "time_period": "1d"} for i in range(50)} with ( - patch.object(dual_cache, "async_batch_get_cache", new=AsyncMock(return_value=[None] * 50)) as batched, + patch.object(dual_cache, "async_batch_get_cache", new=AsyncMock(return_value=[None] * 100)) as batched, patch.object(dual_cache, "async_get_cache", new=AsyncMock()) as single, ): usage = await build_model_max_budget_usage( @@ -1044,7 +1077,8 @@ async def test_usage_report_reads_every_counter_in_one_batched_lookup(): ) assert batched.await_count == 1 - assert len(batched.await_args.kwargs["keys"]) == 50 + # One spend counter and one window start per budgeted model, all in the same round trip. + assert len(batched.await_args.kwargs["keys"]) == 100 assert single.await_count == 0 assert len(usage) == 50 @@ -1063,7 +1097,161 @@ async def test_usage_report_survives_a_batch_lookup_that_returns_nothing(): entity_id="hash-none", model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}}, cache=dual_cache, - ) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1d"}} + ) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1d", "budget_reset_at": None}} + + +_WINDOW_OPENED_AT: Final = datetime(2026, 10, 10, 12, 0, tzinfo=timezone.utc) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "seconds_since_window_opened,expected_reset_at", + [ + (0, "2026-10-10T13:00:00+00:00"), + (600, "2026-10-10T13:00:00+00:00"), + (3599, "2026-10-10T13:00:00+00:00"), + (3600, None), + (7200, None), + ], + ids=["just_opened", "mid_window", "last_second", "ended_exactly_now", "long_ended"], +) +async def test_usage_reports_when_the_open_window_resets(seconds_since_window_opened, expected_reset_at): + """ + The counter expires at the window's start plus its period, and that is when the + model is admitted again. A window that has already run its period is not open, + even if its start time is still sitting in the cache, so no reset is reported. + """ + dual_cache = DualCache() + key = dict(entity_type=Litellm_EntityType.KEY, entity_id="vk-hash", budget_model="gpt-4", budget_duration="1h") + await dual_cache.async_set_cache(key=model_budget_spend_cache_key(**key), value=0.4) + await dual_cache.async_set_cache(key=model_budget_start_time_cache_key(**key), value=_WINDOW_OPENED_AT.timestamp()) + + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="vk-hash", + model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1h"}}, + cache=dual_cache, + now=_WINDOW_OPENED_AT + timedelta(seconds=seconds_since_window_opened), + ) + + assert usage == { + "gpt-4": {"current_spend": 0.4, "budget_limit": 1.0, "time_period": "1h", "budget_reset_at": expected_reset_at} + } + + +_DAY: Final = 86400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "opened_at,read_after_days,window_days", + [ + (datetime(2026, 1, 15, 12, tzinfo=timezone.utc), 26, 31), + (datetime(2026, 1, 15, 12, tzinfo=timezone.utc), 29, 31), + (datetime(2026, 1, 31, 12, tzinfo=timezone.utc), 20, 28), + (datetime(2026, 12, 15, 12, tzinfo=timezone.utc), 30, 31), + ], + ids=["read_in_a_shorter_month", "read_past_the_shorter_month_length", "clamped_to_february", "across_new_year"], +) +async def test_a_monthly_window_resets_a_calendar_month_after_it_opened(opened_at, read_after_days, window_days): + """ + A month is only as long as the month the window opened in, which is the lifetime the + counter was given when it opened. Measuring the month from the read instead made a + window opened on January 15 and read on February 10 report February 12, and then null + while the counter still refused requests until February 15. + """ + dual_cache = DualCache() + started_at: Final = opened_at.timestamp() + await dual_cache.async_set_cache( + key=model_budget_start_time_cache_key( + entity_type=Litellm_EntityType.USER, entity_id="user-1", budget_model="gpt-4", budget_duration="1mo" + ), + value=started_at, + ) + + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.USER, + entity_id="user-1", + model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, + cache=dual_cache, + now=opened_at + timedelta(days=read_after_days), + ) + + assert usage == { + "gpt-4": { + "current_spend": 0.0, + "budget_limit": 1.0, + "time_period": "1mo", + "budget_reset_at": (opened_at + timedelta(days=window_days)).isoformat(), + } + } + + +@pytest.mark.asyncio +async def test_a_monthly_window_reports_no_reset_once_its_month_has_passed(): + dual_cache = DualCache() + opened_at: Final = datetime(2026, 1, 15, 12, tzinfo=timezone.utc) + await dual_cache.async_set_cache( + key=model_budget_start_time_cache_key( + entity_type=Litellm_EntityType.USER, entity_id="user-1", budget_model="gpt-4", budget_duration="1mo" + ), + value=opened_at.timestamp(), + ) + + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.USER, + entity_id="user-1", + model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, + cache=dual_cache, + now=opened_at + timedelta(days=31), + ) + + assert usage == { + "gpt-4": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1mo", "budget_reset_at": None} + } + + +@pytest.mark.asyncio +async def test_each_budgeted_model_reports_its_own_window_reset(): + """Two models on one key own separate windows, so one model's reset must never be reported for the other.""" + dual_cache = DualCache() + hourly = dict(entity_type=Litellm_EntityType.KEY, entity_id="vk-hash", budget_model="gpt-4", budget_duration="1h") + daily = dict(entity_type=Litellm_EntityType.KEY, entity_id="vk-hash", budget_model="claude-3", budget_duration="1d") + await dual_cache.async_set_cache( + key=model_budget_start_time_cache_key(**hourly), value=_WINDOW_OPENED_AT.timestamp() + ) + await dual_cache.async_set_cache( + key=model_budget_start_time_cache_key(**daily), + value=(_WINDOW_OPENED_AT - timedelta(hours=5)).timestamp(), + ) + + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="vk-hash", + model_max_budget={ + "gpt-4": {"budget_limit": 1.0, "time_period": "1h"}, + "claude-3": {"budget_limit": 1.0, "time_period": "1d"}, + "never-used": {"budget_limit": 1.0, "time_period": "1d"}, + }, + cache=dual_cache, + now=_WINDOW_OPENED_AT, + ) + + assert usage == { + "gpt-4": { + "current_spend": 0.0, + "budget_limit": 1.0, + "time_period": "1h", + "budget_reset_at": "2026-10-10T13:00:00+00:00", + }, + "claude-3": { + "current_spend": 0.0, + "budget_limit": 1.0, + "time_period": "1d", + "budget_reset_at": "2026-10-11T07:00:00+00:00", + }, + "never-used": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1d", "budget_reset_at": None}, + } @pytest.mark.asyncio @@ -1097,7 +1285,16 @@ async def test_one_malformed_scope_does_not_abort_the_other_scopes(): entity_id="hash-mixed", model_max_budget=key_budget, cache=dual_cache, - ) == {"gpt-4": {"current_spend": 0.25, "budget_limit": 10.0, "time_period": "1d"}} + ) == { + "gpt-4": { + "current_spend": 0.25, + "budget_limit": 10.0, + "time_period": "1d", + "budget_reset_at": await _reset_at_of_open_window( + dual_cache, Litellm_EntityType.KEY, "hash-mixed", "gpt-4", "1d" + ), + } + } @pytest.mark.asyncio @@ -1368,6 +1565,7 @@ class _SharedFakeRedis(RedisCache): def __init__(self): self._store = MappingProxyType({}) + self.expiry_seconds = MappingProxyType({}) async def async_set_cache(self, key, value, **kwargs): self._store = MappingProxyType({**self._store, key: value}) @@ -1382,6 +1580,7 @@ class _SharedFakeRedis(RedisCache): for op in increment_list: total = self._store.get(op["key"], 0.0) + op["increment_value"] self._store = MappingProxyType({**self._store, op["key"]: total}) + self.expiry_seconds = MappingProxyType({**self.expiry_seconds, op["key"]: op["ttl"]}) return [self._store[op["key"]] for op in increment_list] @@ -1441,10 +1640,168 @@ async def test_spend_logged_on_one_replica_is_enforced_and_reported_on_another() await replica_c.is_key_within_model_budget(user_api_key, "gpt-4") -def _log_success(limiter, **kwargs): - return limiter.async_log_success_event( - _success_kwargs(**kwargs), response_obj=None, start_time=None, end_time=None +@pytest.mark.asyncio +async def test_a_replica_reports_the_reset_of_a_window_another_replica_opened(): + """ + The window start is written by whichever replica served the first charge, so a + replica that has never seen this key must read it from the shared Redis, or it + would report no reset for a model it is actively refusing. + """ + shared_redis = _SharedFakeRedis() + opening_replica = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis)) + other_replica = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis)) + model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1h"}} + + await _log_spend(opening_replica, key_hash="vk-shared", model_max_budget=model_max_budget, response_cost=0.5) + + usage_on_other: Final = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="vk-shared", + model_max_budget=model_max_budget, + cache=other_replica.dual_cache, ) + started_at: Final = shared_redis._store[ + model_budget_start_time_cache_key( + entity_type=Litellm_EntityType.KEY, entity_id="vk-shared", budget_model="gpt-4", budget_duration="1h" + ) + ] + assert usage_on_other == { + "gpt-4": { + "current_spend": 0.5, + "budget_limit": 1.0, + "time_period": "1h", + "budget_reset_at": (datetime.fromtimestamp(started_at, tz=timezone.utc) + timedelta(hours=1)).isoformat(), + } + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("opened_at", "charged_at"), + [ + (datetime(2026, 1, 15, 12, tzinfo=timezone.utc), datetime(2026, 2, 1, 12, tzinfo=timezone.utc)), + (datetime(2026, 2, 15, 12, tzinfo=timezone.utc), datetime(2026, 3, 1, 12, tzinfo=timezone.utc)), + ], + ids=["opened_in_a_long_month", "opened_in_february"], +) +async def test_a_monthly_window_is_enforced_until_the_reset_it_reports(opened_at, charged_at): + """ + Each charge sets when the shared spend counter expires, which is when the model is + allowed again, so it must match the reset /key/info reports. + + Regression: the expiry was a month counted from the charge instead of the window + start, so a window opened in a month of another length stayed blocked past its + reported reset, or reported no reset while still blocked. + """ + shared_redis = _SharedFakeRedis() + limiter = PROXY_VirtualKeyModelMaxBudgetLimiter( + dual_cache=DualCache(redis_cache=shared_redis), clock=charged_at.timestamp + ) + model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}} + window: Final = { + "entity_type": Litellm_EntityType.KEY, + "entity_id": "vk-monthly", + "budget_model": "gpt-4", + "budget_duration": "1mo", + } + await shared_redis.async_set_cache(model_budget_start_time_cache_key(**window), opened_at.timestamp()) + + await _log_spend(limiter, key_hash="vk-monthly", model_max_budget=model_max_budget, response_cost=0.5) + + usage: Final = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="vk-monthly", + model_max_budget=model_max_budget, + cache=limiter.dual_cache, + now=charged_at, + ) + reported_reset: Final = datetime.fromisoformat(usage["gpt-4"]["budget_reset_at"]) + assert usage["gpt-4"]["current_spend"] == 0.5 + assert shared_redis.expiry_seconds[model_budget_spend_cache_key(**window)] == int( + (reported_reset - charged_at).total_seconds() + ) + + +@pytest.mark.asyncio +async def test_a_window_that_ends_past_the_calendar_reports_no_reset(): + """ + Regression: a period long enough to end after year 9999 raised while formatting the + reset, so /key/info, /user/info and /team/info answered 500. + """ + dual_cache = DualCache() + opened_at: Final = datetime(2026, 1, 15, 12, tzinfo=timezone.utc) + model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "99999999h"}} + window: Final = { + "entity_type": Litellm_EntityType.KEY, + "entity_id": "vk-forever", + "budget_model": "gpt-4", + "budget_duration": "99999999h", + } + await dual_cache.async_set_cache(model_budget_spend_cache_key(**window), 0.25) + await dual_cache.async_set_cache(model_budget_start_time_cache_key(**window), opened_at.timestamp()) + + usage: Final = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="vk-forever", + model_max_budget=model_max_budget, + cache=dual_cache, + now=opened_at + timedelta(days=1), + ) + + assert usage == { + "gpt-4": { + "current_spend": 0.25, + "budget_limit": 1.0, + "time_period": "99999999h", + "budget_reset_at": None, + } + } + + +@pytest.fixture +def host_timezone(request, monkeypatch): + monkeypatch.setenv("TZ", request.param) + time.tzset() + yield request.param + monkeypatch.undo() + time.tzset() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("host_timezone", ["UTC", "Europe/Amsterdam", "America/Los_Angeles"], indirect=True) +async def test_a_monthly_reset_is_the_same_on_every_host_timezone(host_timezone): + """ + Replicas behind one load balancer can run with different TZ settings, and they + must agree on when a monthly window ends. + + Regression: months were counted on the host's local calendar, so a window opened + late on January 30 UTC ended a day earlier on a host in Amsterdam. + """ + dual_cache = DualCache() + opened_at: Final = datetime(2027, 1, 30, 23, 30, tzinfo=timezone.utc) + model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}} + window: Final = { + "entity_type": Litellm_EntityType.KEY, + "entity_id": "vk-anywhere", + "budget_model": "gpt-4", + "budget_duration": "1mo", + } + await dual_cache.async_set_cache(model_budget_spend_cache_key(**window), 0.25) + await dual_cache.async_set_cache(model_budget_start_time_cache_key(**window), opened_at.timestamp()) + + usage: Final = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="vk-anywhere", + model_max_budget=model_max_budget, + cache=dual_cache, + now=datetime(2027, 2, 10, tzinfo=timezone.utc), + ) + + assert usage["gpt-4"]["budget_reset_at"] == "2027-02-28T23:30:00+00:00" + + +def _log_success(limiter, **kwargs): + return limiter.async_log_success_event(_success_kwargs(**kwargs), response_obj=None, start_time=None, end_time=None) @pytest.mark.asyncio @@ -1496,7 +1853,16 @@ async def test_team_model_budget_is_shared_by_every_key_without_an_override(requ entity_id="team-1", model_max_budget=team_model_max_budget, cache=dual_cache, - ) == {"gpt-4": {"current_spend": pytest.approx(1.2), "budget_limit": 1.0, "time_period": "1d"}} + ) == { + "gpt-4": { + "current_spend": pytest.approx(1.2), + "budget_limit": 1.0, + "time_period": "1d", + "budget_reset_at": await _reset_at_of_open_window( + dual_cache, Litellm_EntityType.TEAM, "team-1", "gpt-4", "1d" + ), + } + } @pytest.mark.asyncio diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index 84c4608d29b..fb463186d8c 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -3522,7 +3522,7 @@ async def test_user_info_v2_response_shape(mocker): # these two, so dropping them makes a save replace the user's budgets. assert response_dict["model_max_budget"] == {"gpt-3.5-turbo": {"budget_limit": 5.0, "time_period": "30d"}} assert response_dict["model_max_budget_usage"] == { - "gpt-3.5-turbo": {"current_spend": 0.0, "budget_limit": 5.0, "time_period": "30d"} + "gpt-3.5-turbo": {"current_spend": 0.0, "budget_limit": 5.0, "time_period": "30d", "budget_reset_at": None} } # Verify teams is a list of strings (team IDs), not team objects diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index cce61e443a9..685ea46547e 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -6,6 +6,7 @@ anthropic_messages, embedding and rerank surfaces leave it unset, which used to the callback raise before any spend was recorded, so those budgets never moved. """ +from datetime import datetime, timezone from typing import Final import litellm @@ -163,6 +164,90 @@ async def test_deployment_budget_tracked_when_provider_is_unresolvable(disable_b assert await limiter.dual_cache.async_get_cache("deployment_spend:deployment-1:1d") == 0.25 +_MONTHLY_DEPLOYMENT: Final = { + "model_name": "some-model", + "litellm_params": {"model": "openai/gpt-4o", "max_budget": 10.0, "budget_duration": "1mo"}, + "model_info": {"id": "deployment-1"}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("limiter_config", "tag_budget_config", "spend_key", "start_time_key"), + [ + ( + {"provider_budget_config": {"openai": {"budget_limit": 10.0, "time_period": "1mo"}}}, + None, + "provider_spend:openai:1mo", + "provider_budget_start_time:openai", + ), + ( + {"provider_budget_config": None, "model_list": [_MONTHLY_DEPLOYMENT]}, + None, + "deployment_spend:deployment-1:1mo", + "deployment_budget_start_time:deployment-1", + ), + ( + {"provider_budget_config": None}, + {"release": {"budget_limit": 10.0, "time_period": "1mo"}}, + "tag_spend:release:1mo", + "tag_budget_start_time:release", + ), + ], + ids=["provider", "deployment", "tag"], +) +@pytest.mark.parametrize( + ("opened_at", "charged_at", "closes_at"), + [ + ( + datetime(2026, 1, 15, 12, tzinfo=timezone.utc), + datetime(2026, 2, 1, 12, tzinfo=timezone.utc), + datetime(2026, 2, 15, 12, tzinfo=timezone.utc), + ), + ( + datetime(2026, 2, 15, 12, tzinfo=timezone.utc), + datetime(2026, 3, 1, 12, tzinfo=timezone.utc), + datetime(2026, 3, 15, 12, tzinfo=timezone.utc), + ), + ], + ids=["opened_in_a_long_month", "opened_in_february"], +) +async def test_a_monthly_router_budget_ends_a_calendar_month_after_it_opens( + disable_budget_sync, + monkeypatch, + limiter_config, + tag_budget_config, + spend_key, + start_time_key, + opened_at, + charged_at, + closes_at, +): + """ + The spend counter's Redis expiry is when the router lets the provider, + deployment or tag take traffic again. + + Regression: it was a month counted from the charge, so a window opened on January 15 + and charged on February 1 unblocked on February 12 instead of February 15. + """ + monkeypatch.setattr(litellm, "tag_budget_config", tag_budget_config) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + limiter = RouterBudgetLimiting(dual_cache=DualCache(), clock=charged_at.timestamp, **limiter_config) + await limiter.dual_cache.async_set_cache(key=start_time_key, value=opened_at.timestamp()) + + await _log_success( + limiter, + { + **_success_kwargs(provider_in_litellm_params="openai", provider_in_payload="openai"), + "metadata": {"tags": ["release"]}, + }, + ) + + assert limiter.redis_increment_operation_queue == [ + {"key": spend_key, "increment_value": 0.25, "ttl": int((closes_at - charged_at).total_seconds())} + ] + + @pytest.mark.asyncio async def test_get_budget_config_for_provider(): """ @@ -175,9 +260,7 @@ async def test_get_budget_config_for_provider(): "anthropic": BudgetConfig(budget_duration="7d", max_budget=500), } - provider_budget = RouterBudgetLimiting( - dual_cache=DualCache(), provider_budget_config=config - ) + provider_budget = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config=config) # Test existing providers openai_config = provider_budget._get_budget_config_for_provider("openai") @@ -230,16 +313,9 @@ async def test_get_current_provider_spend(): @pytest.mark.asyncio async def test_get_llm_provider_for_deployment_resolves_provider_prefixes() -> None: - provider_budget: Final = RouterBudgetLimiting( - dual_cache=DualCache(), provider_budget_config={} - ) + provider_budget: Final = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={}) - assert ( - provider_budget._get_llm_provider_for_deployment( - {"litellm_params": {"model": "openai/gpt-4o"}} - ) - == "openai" - ) + assert provider_budget._get_llm_provider_for_deployment({"litellm_params": {"model": "openai/gpt-4o"}}) == "openai" assert ( provider_budget._get_llm_provider_for_deployment( {"litellm_params": {"model": "azure/gpt-4o", "api_base": "https://example.azure.com"}} @@ -251,9 +327,7 @@ async def test_get_llm_provider_for_deployment_resolves_provider_prefixes() -> N @pytest.mark.asyncio async def test_budget_start_time_is_created_once() -> None: - provider_budget: Final = RouterBudgetLimiting( - dual_cache=DualCache(), provider_budget_config={} - ) + provider_budget: Final = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={}) first_start: Final = await provider_budget._get_or_set_budget_start_time( start_time_key="provider_budget_start_time:openai", @@ -272,9 +346,7 @@ async def test_budget_start_time_is_created_once() -> None: @pytest.mark.asyncio async def test_new_budget_window_sets_spend_and_start_time() -> None: - provider_budget: Final = RouterBudgetLimiting( - dual_cache=DualCache(), provider_budget_config={} - ) + provider_budget: Final = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={}) start_time: Final = await provider_budget._handle_new_budget_window( spend_key="provider_spend:openai:1d", @@ -286,22 +358,13 @@ async def test_new_budget_window_sets_spend_and_start_time() -> None: assert start_time == 1000.0 assert await provider_budget.dual_cache.async_get_cache("provider_spend:openai:1d") == 0.5 - assert ( - await provider_budget.dual_cache.async_get_cache( - "provider_budget_start_time:openai" - ) - == 1000.0 - ) + assert await provider_budget.dual_cache.async_get_cache("provider_budget_start_time:openai") == 1000.0 @pytest.mark.asyncio async def test_spend_increment_updates_memory_and_queues_redis_operation() -> None: - provider_budget: Final = RouterBudgetLimiting( - dual_cache=DualCache(), provider_budget_config={} - ) - await provider_budget.dual_cache.async_set_cache( - key="provider_spend:openai:1d", value=1.0, ttl=86400 - ) + provider_budget: Final = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={}) + await provider_budget.dual_cache.async_set_cache(key="provider_spend:openai:1d", value=1.0, ttl=86400) await provider_budget._increment_spend_in_current_window( spend_key="provider_spend:openai:1d", @@ -309,10 +372,7 @@ async def test_spend_increment_updates_memory_and_queues_redis_operation() -> No ttl=86400, ) - assert ( - await provider_budget.dual_cache.async_get_cache("provider_spend:openai:1d") - == 1.5 - ) + assert await provider_budget.dual_cache.async_get_cache("provider_spend:openai:1d") == 1.5 assert provider_budget.redis_increment_operation_queue == [ { "key": "provider_spend:openai:1d", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 73cd88d2e50..88d3233598b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8565,7 +8565,8 @@ export interface paths { * hour), so subtracting budget_duration from it does not give the window's start * - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} * - model_max_budget_usage: dict | None - Current-window spend per model, present only when - * the key has per-model budgets + * the key has per-model budgets. Each entry's budget_reset_at is when that model's open + * window ends and its spend resets, or None when no window is open or it ends after year 9999 * - budget_limits: list | None - Concurrent budget windows, exactly as stored * - budget_limits_usage: dict | None - Current-window spend per budget window, e.g. * {"1h": {"current_spend": 0.0009}}, present only when the key has budget windows