From a76e723b26c44d362c7d22ea2809b325fd856888 Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Thu, 15 Oct 2026 19:56:45 +0200 Subject: [PATCH] feat(proxy): report when each per-model budget window resets (#45797) * feat(proxy): report when each per-model budget window resets model_max_budget_usage on /key/info, /user/info and /team/info now carries budget_reset_at per model, the open window's start plus its period, or null when no window is open. The window start is read in the same batched lookup as the spend counter * fix(proxy): count a monthly budget window from its own start when reporting its reset duration_in_seconds measures Nmo from now, so a window opened on January 15 and read on February 10 reported February 12, then null while the counter still refused requests until February 15. duration_in_seconds_from measures the month from the window start, and duration_in_seconds now shares that month arithmetic * fix(proxy): measure months from timezone-aware datetimes and assert whole usage values datetime.fromtimestamp without a tz tripped DTZ006. The month arithmetic now runs on an aware local datetime, which keeps the same wall-clock result. The duration tests move to parametrized functions, since subTest payloads carrying datetimes cannot cross the xdist worker boundary, and the new usage tests compare the whole entry * fix(budgets): end a monthly budget window a calendar month after it opens Each charge recounted the month from the charge date, so a monthly window opened in a month of another length expired days before or after the reset /key/info reports. The window is now measured from its start, by the same helper the usage report uses. A period ending after year 9999 reports no reset instead of failing the read. * fix(budgets): count monthly budget windows in UTC for every router budget Months were counted on the host's local calendar, so replicas with different TZ settings disagreed on when a monthly window ends. Provider, deployment, tag and per-model budgets now all measure a window from its start in UTC, and a test covers the router budgets. A reset past year 9999 is documented as null. * test(budgets): cover monthly tag budgets in the router window test * test(budgets): add integration cells for budget_reset_at and monthly budget windows * test(budgets): find a monthly window in another span on every calendar day --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/duration_parser.py | 54 +- .../proxy/hooks/model_max_budget_limiter.py | 104 +- .../key_management_endpoints.py | 3 +- litellm/router_strategy/budget_limiter.py | 58 +- .../management/test_budget_reset_windows.py | 1020 +++++++++++++++++ .../test_duration_parser.py | 32 + ...test_unit_test_max_model_budget_limiter.py | 416 ++++++- .../test_internal_user_endpoints.py | 2 +- .../router_strategy/test_budget_limiter.py | 128 ++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 3 +- 10 files changed, 1677 insertions(+), 143 deletions(-) create mode 100644 tests/integration/management/test_budget_reset_windows.py 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