diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 400d30bc0a2..1a024b6b2b0 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1273,6 +1273,24 @@ class RedisCache(BaseCache): if len(self.redis_batch_writing_buffer) >= self.redis_flush_size: await self.flush_cache_buffer() # logging done in here + @staticmethod + async def _incrbyfloat_with_ttl( + _redis_client: "Redis", key: str, value: float, ttl: int | None, refresh_ttl: bool + ) -> float: + """INCRBYFLOAT plus its TTL command in one round trip; a third only when an unexpiring key needs an EXPIRE.""" + if ttl is None: + return await _redis_client.incrbyfloat(name=key, amount=value) + async with _redis_client.pipeline(transaction=False) as pipe: + pipe.incrbyfloat(name=key, amount=value) + if refresh_ttl: + pipe.expire(key, ttl) + else: + pipe.ttl(key) + result, ttl_or_expire = await pipe.execute() + if not refresh_ttl and ttl_or_expire == -1: + await _redis_client.expire(key, ttl) + return float(result) + @_redis_circuit_breaker_guard async def async_increment( self, @@ -1289,14 +1307,9 @@ class RedisCache(BaseCache): _used_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) try: - result: Final = await _redis_client.incrbyfloat(name=key, amount=value) - if _used_ttl is not None: - if refresh_ttl: - await _redis_client.expire(key, _used_ttl) - else: - current_ttl: Final = await _redis_client.ttl(key) - if current_ttl == -1: - await _redis_client.expire(key, _used_ttl) + result: Final = await self._incrbyfloat_with_ttl( + _redis_client, key=key, value=value, ttl=_used_ttl, refresh_ttl=refresh_ttl + ) ## LOGGING ## end_time = time.time() diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 540ce6738fc..2528f07f92c 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -52,6 +52,12 @@ from litellm.types.integrations.prometheus import ( _sanitize_prometheus_label_value, validate_prometheus_deployment_and_latency_caller_identity, ) +from litellm.types.proxy.carried_budget_state import ( + KeyBudgetSnapshot, + OrgBudgetSnapshot, + TeamBudgetSnapshot, + UserBudgetSnapshot, +) from litellm.types.utils import ( StandardLoggingGuardrailInformation, StandardLoggingPayload, @@ -1941,6 +1947,8 @@ class PrometheusLogger(CustomLogger): _user_spend: Final = _metadata.get("user_api_key_user_spend", None) _user_max_budget: Final = _metadata.get("user_api_key_user_max_budget", None) + _user_email: Final = _metadata.get("user_api_key_user_email", None) + _org_alias: Final = _metadata.get("user_api_key_org_alias", None) # Bound the per-request budget-metric emission so that slow Redis/DB # lookups under load cannot consume the whole LoggingWorker watchdog @@ -1957,6 +1965,7 @@ class PrometheusLogger(CustomLogger): response_cost=response_cost, key_max_budget=_api_key_max_budget, key_spend=_api_key_spend, + carried=KeyBudgetSnapshot.from_metadata(_metadata), ), self._set_team_budget_metrics_after_api_request( user_api_team=user_api_team, @@ -1964,16 +1973,21 @@ class PrometheusLogger(CustomLogger): team_spend=_team_spend, team_max_budget=_team_max_budget, response_cost=response_cost, + carried=TeamBudgetSnapshot.from_metadata(_metadata), ), self._set_user_budget_metrics_after_api_request( user_id=user_id, user_spend=_user_spend, user_max_budget=_user_max_budget, response_cost=response_cost, + carried=UserBudgetSnapshot.from_metadata(_metadata), + user_email=_user_email if isinstance(_user_email, str) else None, ), self._set_org_budget_metrics_after_api_request( org_id=user_api_key_org_id, response_cost=response_cost, + carried=OrgBudgetSnapshot.from_metadata(_metadata), + org_alias=_org_alias if isinstance(_org_alias, str) else None, ), return_exceptions=True, ) @@ -3821,6 +3835,7 @@ class PrometheusLogger(CustomLogger): team_spend: float | None, team_max_budget: float | None, response_cost: float, + carried: TeamBudgetSnapshot | None = None, ): """ Set team budget metrics after an LLM API request @@ -3839,6 +3854,7 @@ class PrometheusLogger(CustomLogger): spend=team_spend, max_budget=team_max_budget, response_cost=response_cost, + carried=carried, ) self._set_team_budget_metrics(team_object) @@ -3850,18 +3866,26 @@ class PrometheusLogger(CustomLogger): spend: float | None, max_budget: float | None, response_cost: float, + carried: TeamBudgetSnapshot | None = None, ) -> LiteLLM_TeamTable: """ Assemble a LiteLLM_TeamTable object - for fields not available in metadata, we fetch from db - Fields not available in metadata: - - `budget_reset_at` + ``budget_reset_at`` comes from the auth-carried snapshot when the request has one, + otherwise from the team lookup """ from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache _total_team_spend: Final = (spend or 0) + response_cost + if carried is not None: + return LiteLLM_TeamTable( + team_id=team_id, + team_alias=team_alias, + spend=_total_team_spend, + max_budget=max_budget if max_budget is not None else carried.max_budget, + budget_reset_at=carried.budget_reset_at, + ) team_object: Final = LiteLLM_TeamTable( team_id=team_id, team_alias=team_alias, @@ -3946,11 +3970,13 @@ class PrometheusLogger(CustomLogger): self, org_id: str | None, response_cost: float, + carried: OrgBudgetSnapshot | None = None, + org_alias: str | None = None, ): """ Set org budget metrics after an LLM API request - - Fetches org info via cache (get_org_object) + - Uses the auth-carried org budget when the request has one, else fetches via get_org_object - Sets org budget metrics """ if isinstance(self.litellm_remaining_org_budget_metric, NoOpMetric): @@ -3959,6 +3985,16 @@ class PrometheusLogger(CustomLogger): if not org_id: return + if carried is not None: + self._set_org_budget_metrics( + org_id=org_id, + org_alias=org_alias or "", + spend=carried.spend + response_cost, + max_budget=carried.max_budget, + budget_reset_at=None, + ) + return + from litellm.proxy.auth.auth_checks import get_org_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache @@ -3979,7 +4015,6 @@ class PrometheusLogger(CustomLogger): if org_info is None: return - org_alias: Final = org_info.organization_alias or "" _total_org_spend: Final = (org_info.spend or 0.0) + response_cost budget_table: Final = org_info.litellm_budget_table max_budget: Final = budget_table.max_budget if budget_table else None @@ -3987,7 +4022,7 @@ class PrometheusLogger(CustomLogger): self._set_org_budget_metrics( org_id=org_id, - org_alias=org_alias, + org_alias=org_info.organization_alias or "", spend=_total_org_spend, max_budget=max_budget, budget_reset_at=budget_reset_at, @@ -4084,6 +4119,7 @@ class PrometheusLogger(CustomLogger): response_cost: float, key_max_budget: float | None, key_spend: float | None, + carried: KeyBudgetSnapshot | None = None, ): if isinstance(self.litellm_remaining_api_key_budget_metric, NoOpMetric): return @@ -4095,6 +4131,7 @@ class PrometheusLogger(CustomLogger): key_max_budget=key_max_budget, key_spend=key_spend, response_cost=response_cost, + carried=carried, ) self._set_key_budget_metrics(user_api_key_dict) @@ -4105,6 +4142,7 @@ class PrometheusLogger(CustomLogger): key_max_budget: float | None, key_spend: float | None, response_cost: float, + carried: KeyBudgetSnapshot | None = None, ) -> UserAPIKeyAuth: """ Assemble a UserAPIKeyAuth object @@ -4113,6 +4151,14 @@ class PrometheusLogger(CustomLogger): from litellm.proxy.proxy_server import prisma_client, user_api_key_cache _total_key_spend: Final = (key_spend or 0) + response_cost + if carried is not None: + return UserAPIKeyAuth( + token=user_api_key, + key_alias=user_api_key_alias, + max_budget=key_max_budget, + spend=_total_key_spend, + budget_reset_at=carried.budget_reset_at, + ) user_api_key_dict: Final = UserAPIKeyAuth( token=user_api_key, key_alias=user_api_key_alias, @@ -4140,6 +4186,8 @@ class PrometheusLogger(CustomLogger): user_spend: float | None, user_max_budget: float | None, response_cost: float, + carried: UserBudgetSnapshot | None = None, + user_email: str | None = None, ): """ Set user budget metrics after an LLM API request @@ -4157,6 +4205,8 @@ class PrometheusLogger(CustomLogger): spend=user_spend, max_budget=user_max_budget, response_cost=response_cost, + carried=carried, + user_email=user_email, ) self._set_user_budget_metrics(user_object) @@ -4167,18 +4217,28 @@ class PrometheusLogger(CustomLogger): spend: float | None, max_budget: float | None, response_cost: float, + carried: UserBudgetSnapshot | None = None, + user_email: str | None = None, ) -> LiteLLM_UserTable: """ Assemble a LiteLLM_UserTable object - for fields not available in metadata, we fetch from db - Fields not available in metadata: - - `budget_reset_at` + ``budget_reset_at`` and ``user_alias`` come from the auth-carried snapshot when the + request has one, otherwise from the user lookup """ from litellm.proxy.auth.auth_checks import get_user_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache _total_user_spend: Final = (spend or 0) + response_cost + if carried is not None: + return LiteLLM_UserTable( + user_id=user_id, + spend=_total_user_spend, + max_budget=max_budget if max_budget is not None else carried.max_budget, + budget_reset_at=carried.budget_reset_at, + user_alias=carried.user_alias, + user_email=user_email, + ) user_object: Final = LiteLLM_UserTable( user_id=user_id, spend=_total_user_spend, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3d22923c0a8..85677f4eb3c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -40,6 +40,11 @@ from litellm.types.mcp import ( MCPTransportType, ) from litellm.types.mcp_server.mcp_server_manager import MCPInfo +from litellm.types.proxy.carried_budget_state import ( + OrgBudgetSnapshot, + TeamBudgetSnapshot, + UserBudgetSnapshot, +) from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry from litellm.types.router import RouterErrors, UpdateRouterConfig from litellm.types.secret_managers.main import KeyManagementSystem @@ -3106,6 +3111,9 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob ), ) budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True) + team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True) + user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True) + org_budget_snapshot: OrgBudgetSnapshot | None = Field(default=None, exclude=True) matched_model_access_groups: list[str] | None = Field(default=None, exclude=True) budget_throttle_pct: float | None = Field(default=None, exclude=True) user: Any | None = None # Expanded user object when expand=user is used diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 18095aaafb4..9c175242a9a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -102,6 +102,7 @@ from litellm.proxy.guardrails.tool_name_extraction import ( ) from litellm.proxy.route_llm_request import route_request from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start +from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.object_permission_repository import ObjectPermissionRepository @@ -5635,6 +5636,8 @@ async def _organization_max_budget_check( if org_table is None: return + carry_organization_budget_state(valid_token=valid_token, org_table=org_table) + # Get max_budget from organization's budget table org_max_budget: float | None = None if org_table.litellm_budget_table is not None: diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2481a7436a7..9828311112e 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -103,6 +103,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( ) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.spend_tracking.carried_budget_state import carry_team_and_user_budget_state from litellm.proxy.spend_tracking.spend_counter_batch import ( bind_admission_counter_keys, release_spend_counter_batch, @@ -2631,6 +2632,11 @@ async def _run_centralized_common_checks( None if isinstance(end_user_result, BaseException) else end_user_result ) global_proxy_spend: float | None = None if isinstance(global_spend_result, BaseException) else global_spend_result + carry_team_and_user_budget_state( + valid_token=user_api_key_auth_obj, + team_object=team_object, + user_object=user_object, + ) if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: user_api_key_auth_obj.org_id = team_object.organization_id diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index b3ddf9bb8cd..89a07234c6c 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -24,7 +24,7 @@ from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import Litellm_EntityType from litellm.proxy.db.db_lookup_gate import db_lookup_gate -from litellm.proxy.spend_tracking.spend_counter_batch import active_spend_counter_batch +from litellm.proxy.spend_tracking.spend_counter_batch import read_batched_spend_counter, record_spend_counter_value from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( BudgetWindowSpendRepository, @@ -197,9 +197,8 @@ class SpendCounterReseed: @staticmethod async def _read_active_batch(counter_key: str) -> tuple[float | None, bool] | None: - """The request's admission MGET already answered for this counter; a Redis miss there is authoritative.""" - batch: Final = active_spend_counter_batch() - return None if batch is None else await batch.read(counter_key) + """The request's MGET answers for this counter; a Redis miss there is authoritative.""" + return await read_batched_spend_counter(counter_key) @staticmethod async def coalesced( @@ -263,6 +262,7 @@ class SpendCounterReseed: key=counter_key, value=current_value, ) + record_spend_counter_value(counter_key, current_value) else: cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) seeded_spend: Final = max(db_spend, float(cached_spend)) if cached_spend is not None else db_spend @@ -414,8 +414,11 @@ class SpendCounterReseed: ) -> float | None: lock: Final = await SpendCounterReseed._get_lock(counter_key) async with lock: - redis_clean_miss = False - if spend_counter_cache.redis_cache is not None: + batched: Final = await SpendCounterReseed._read_active_batch(counter_key) + if batched is not None and batched[0] is not None: + return batched[0] + redis_clean_miss = batched is not None + if spend_counter_cache.redis_cache is not None and not redis_clean_miss: try: val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) if val is not None: @@ -459,6 +462,7 @@ class SpendCounterReseed: key=counter_key, value=current_value, ) + record_spend_counter_value(counter_key, float(current_value)) else: cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) seeded_spend: Final = ( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 8f7b515c22a..59971e54e46 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -64,6 +64,7 @@ from litellm.proxy.common_utils.callback_utils import ( strip_callback_config, ) from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers +from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY # Cache special headers as a frozenset for O(1) lookup performance @@ -2300,6 +2301,7 @@ async def add_litellm_data_to_request( data[_metadata_variable_name]["user_api_key_user_max_budget"] = user_api_key_dict.user_max_budget user_model_budget: Final = user_api_key_dict.user_model_max_budget data[_metadata_variable_name]["user_api_key_user_model_max_budget"] = user_model_budget # rebind-ok: out-param + data[_metadata_variable_name].update(carried_budget_metadata(user_api_key_dict)) data[_metadata_variable_name]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata) data[_metadata_variable_name]["user_api_key_team_metadata"] = strip_callback_config(user_api_key_dict.team_metadata) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3d751b2ad25..606a590c24b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -26,7 +26,6 @@ from collections.abc import ( MutableMapping, Sequence, ) -from dataclasses import dataclass from datetime import datetime, timedelta, timezone from types import MappingProxyType, UnionType from typing import ( @@ -659,7 +658,15 @@ from litellm.proxy.route_priority import hot_routes_first from litellm.proxy.search_endpoints.endpoints import router as search_router from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start -from litellm.proxy.spend_tracking.spend_counter_batch import active_spend_counter_batch +from litellm.proxy.spend_tracking.spend_counter_batch import ( + PendingSpendIncrement, + active_spend_counter_batch, + forget_spend_counter, + post_call_counter_keys, + read_batched_spend_counter, + record_spend_counter_value, + spend_counter_batch_scope, +) from litellm.proxy.spend_tracking.spend_management_endpoints import ( router as spend_management_router, ) @@ -2628,6 +2635,7 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None if needs_update: spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend) if spend_counter_cache.redis_cache is not None: + forget_spend_counter(counter_key) try: await spend_counter_cache.redis_cache.async_set_max(key=counter_key, value=db_spend) except Exception: @@ -2768,12 +2776,6 @@ async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) return fallback_spend, False -@dataclass(frozen=True, slots=True) -class _PendingSpendIncrement: - counter_key: str - increment: float - - async def increment_spend_counters( token: str | None, team_id: str | None, @@ -2796,6 +2798,45 @@ async def increment_spend_counters( Awaited (not create_task) in the cost callback, so the counter is updated before the next request's auth check runs. """ + with spend_counter_batch_scope( + spend_counter_cache.redis_cache, + counter_keys=post_call_counter_keys( + token=token, + team_id=team_id, + user_id=user_id, + org_id=org_id, + end_user_id=end_user_id, + tags=tags, + model_access_groups=model_access_groups, + ), + ): + await _increment_spend_counters_batched( + token=token, + team_id=team_id, + user_id=user_id, + response_cost=response_cost, + org_id=org_id, + budget_reservation=budget_reservation, + end_user_id=end_user_id, + tags=tags, + request_started_at=request_started_at, + model_access_groups=model_access_groups, + ) + + +async def _increment_spend_counters_batched( + token: str | None, + team_id: str | None, + user_id: str | None, + response_cost: float | None, + org_id: str | None, + budget_reservation: dict | None, + end_user_id: str | None, + tags: list[str] | None, + request_started_at: datetime | None, + model_access_groups: Sequence[str] | None, +): + """Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET.""" reserved_counter_keys: Final = await _reconcile_budget_reservation_for_counter_update( budget_reservation=budget_reservation, response_cost=response_cost, @@ -2808,7 +2849,7 @@ async def increment_spend_counters( cost: Final[float] = response_cost - async def _key_scope(key_token: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: + async def _key_scope(key_token: str) -> tuple[PendingSpendIncrement | BaseException, ...]: # key_token arrives pre-hashed from metadata["user_api_key"] (auth flow # hashes raw "sk-..." keys before they reach the callback). The # startswith("sk-") check is a safety net matching update_cache — @@ -2819,7 +2860,7 @@ async def increment_spend_counters( hash_token(token=key_token) if isinstance(key_token, str) and key_token.startswith("sk-") else key_token ) key_counter_key: Final = f"spend:key:{hashed_token}" - key_pending: Final[tuple[_PendingSpendIncrement, ...]] = ( + key_pending: Final[tuple[PendingSpendIncrement, ...]] = ( () if key_counter_key in reserved_counter_keys else ( @@ -2831,7 +2872,7 @@ async def increment_spend_counters( ) ) - async def _key_window_increment(window: object) -> _PendingSpendIncrement | None: + async def _key_window_increment(window: object) -> PendingSpendIncrement | None: duration = ( window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) ) @@ -2878,9 +2919,9 @@ async def increment_spend_counters( ) return key_pending + tuple(item for item in window_pending if item is not None) - async def _team_scope(scope_team_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: + async def _team_scope(scope_team_id: str) -> tuple[PendingSpendIncrement | BaseException, ...]: team_counter_key: Final = f"spend:team:{scope_team_id}" - team_pending: Final[tuple[_PendingSpendIncrement, ...]] = ( + team_pending: Final[tuple[PendingSpendIncrement, ...]] = ( () if team_counter_key in reserved_counter_keys else ( @@ -2892,7 +2933,7 @@ async def increment_spend_counters( ) ) - async def _team_window_increment(window: object) -> _PendingSpendIncrement | None: + async def _team_window_increment(window: object) -> PendingSpendIncrement | None: duration = ( window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) ) @@ -2941,7 +2982,7 @@ async def increment_spend_counters( async def _team_member_scope( scope_user_id: str, scope_team_id: str - ) -> tuple[_PendingSpendIncrement | BaseException, ...]: + ) -> tuple[PendingSpendIncrement | BaseException, ...]: team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}" if team_member_counter_key in reserved_counter_keys: return () @@ -2953,7 +2994,7 @@ async def increment_spend_counters( ), ) - async def _user_scope(scope_user_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: + async def _user_scope(scope_user_id: str) -> tuple[PendingSpendIncrement | BaseException, ...]: user_counter_key: Final = f"spend:user:{scope_user_id}" if user_counter_key in reserved_counter_keys: return () @@ -3063,7 +3104,7 @@ async def _prepare_end_user_and_tag_spend_increments( tags: list[str] | None, response_cost: float, reserved_counter_keys: set[str], -) -> tuple[_PendingSpendIncrement | BaseException, ...]: +) -> tuple[PendingSpendIncrement | BaseException, ...]: unique_tags: Final = ( tuple(dict.fromkeys(tag for tag in tags if tag and isinstance(tag, str))) if tags is not None else () ) @@ -3100,7 +3141,7 @@ async def _prepare_model_access_group_spend_increments( model_access_groups: Sequence[object], response_cost: float, reserved_counter_keys: set[str], -) -> tuple[_PendingSpendIncrement | BaseException, ...]: +) -> tuple[PendingSpendIncrement | BaseException, ...]: """Charge the model access groups that authorized this request. Without this the counter auth reads is written only by the reservation path, so @@ -3133,7 +3174,7 @@ async def _prepare_org_spend_increment( org_id: str | None, response_cost: float, reserved_counter_keys: set[str], -) -> tuple[_PendingSpendIncrement, ...]: +) -> tuple[PendingSpendIncrement, ...]: if org_id is None: return () @@ -3151,7 +3192,7 @@ async def _prepare_unreserved_spend_counter_increment( source_cache_key: str | list[str], increment: float, reserved_counter_keys: set[str], -) -> _PendingSpendIncrement | None: +) -> PendingSpendIncrement | None: if counter_key in reserved_counter_keys: return None @@ -3166,7 +3207,7 @@ async def _prepare_spend_counter_increment( counter_key: str, source_cache_key: str | list[str], increment: float, -) -> _PendingSpendIncrement: +) -> PendingSpendIncrement: """ Initialize counter from the authoritative DB spend value if not yet set, then return the pending increment for the caller to apply in one @@ -3188,7 +3229,7 @@ async def _prepare_spend_counter_increment( counter_key=counter_key, source_cache_key=source_cache_key, ) - return _PendingSpendIncrement(counter_key=counter_key, increment=increment) + return PendingSpendIncrement(counter_key=counter_key, increment=increment) async def _enqueue_window_spend_row_update( @@ -3247,7 +3288,7 @@ async def _prepare_window_spend_counter_increment( window_duration: str | None, window_start: datetime | None, increment: float, -) -> _PendingSpendIncrement | None: +) -> PendingSpendIncrement | None: if window_start is None: verbose_proxy_logger.warning( "Skipping spend counter increment for invalid budget window %s", @@ -3264,7 +3305,7 @@ async def _prepare_window_spend_counter_increment( ) if initialized is False: return None - return _PendingSpendIncrement(counter_key=counter_key, increment=increment) + return PendingSpendIncrement(counter_key=counter_key, increment=increment) async def _ensure_spend_counter_initialized( @@ -3331,6 +3372,14 @@ async def _ensure_window_spend_counter_initialized( async def _is_spend_counter_cache_warm(counter_key: str) -> bool: + batched: Final = await read_batched_spend_counter(counter_key) + if batched is not None: + batched_value, _ = batched + if batched_value is None: + return False + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=batched_value) + return True + if spend_counter_cache.redis_cache is not None: try: current_value: Final[object] = await spend_counter_cache.redis_cache.async_get_cache( @@ -3375,6 +3424,7 @@ async def _increment_spend_counter_cache(counter_key: str, increment: float): key=counter_key, value=current_value, ) + record_spend_counter_value(counter_key, float(current_value)) return current_value return await SpendCounterReseed.increment_in_memory( @@ -3383,6 +3433,7 @@ async def _increment_spend_counter_cache(counter_key: str, increment: float): async def _invalidate_spend_counter(counter_key: str): + forget_spend_counter(counter_key) spend_counter_cache.in_memory_cache.delete_cache(key=counter_key) if spend_counter_cache.redis_cache is not None: try: @@ -3395,7 +3446,16 @@ async def _invalidate_spend_counter(counter_key: str): ) -async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncrement]) -> None: +async def _apply_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> None: + try: + await increment_spend_counters_pipeline(pending=pending) + except RedisCircuitBreakerOpenError: + return + + +async def increment_spend_counters_pipeline(pending: Sequence[PendingSpendIncrement]) -> None: + """One INCRBYFLOAT+EXPIRE pipeline for every pending counter; on failure every counter is invalidated + before the error propagates, so no caller can read a half-applied batch.""" if not pending: return redis_cache: Final = spend_counter_cache.redis_cache @@ -3412,13 +3472,12 @@ async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncreme ] try: results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list) - except Exception as e: + except Exception: await asyncio.gather(*(_invalidate_spend_counter(counter_key=item.counter_key) for item in pending)) - if isinstance(e, RedisCircuitBreakerOpenError): - return raise for item, current_value in zip(pending, results or ()): spend_counter_cache.in_memory_cache.set_cache(key=item.counter_key, value=current_value) + record_spend_counter_value(item.counter_key, float(current_value)) async def update_cache( diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index eedec0619db..6074a50a69b 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -32,6 +32,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( tag_cache_key, team_membership_reservation_cache_key, ) +from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement, spend_counter_batch_scope from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router from litellm.rust_bridge.token_counter import RustTokenizer, count_input_tokens, rust_tokenizer @@ -257,46 +258,47 @@ async def reserve_budget_for_request( applied_entries: Final[list[dict[str, float | str]]] = [] try: - for counter in counters: - entry = _counter_to_reservation_entry( - counter=counter, - reserved_cost=reservation_cost, - ) - applied_entries.append(entry) - try: - reserved_value = await _reserve_counter( + with _counters_batch_scope(frozenset(counter.counter_key for counter in counters)): + for counter in counters: + entry = _counter_to_reservation_entry( counter=counter, - reservation_cost=reservation_cost, + reserved_cost=reservation_cost, ) - except _CounterReservationUnavailable as exc: - if exc.touched_counter and not exc.counter_invalidated: - await _release_applied_entries_best_effort( - entries=[entry], - default_reserved_cost=reservation_cost, + applied_entries.append(entry) + try: + reserved_value = await _reserve_counter( + counter=counter, + reservation_cost=reservation_cost, ) - applied_entries.remove(entry) - if fail_closed_budget_enforcement: - _raise_reservation_unavailable(counter_key=counter.counter_key) - continue + except _CounterReservationUnavailable as exc: + if exc.touched_counter and not exc.counter_invalidated: + await _release_applied_entries_best_effort( + entries=[entry], + default_reserved_cost=reservation_cost, + ) + applied_entries.remove(entry) + if fail_closed_budget_enforcement: + _raise_reservation_unavailable(counter_key=counter.counter_key) + continue - if reserved_value is not None: - current_spend = reserved_value - else: - cached_spend = current_spend_by_counter_key.get(counter.counter_key) - if cached_spend is None: - cached_spend = await _get_current_counter_value(counter=counter) - current_spend = cached_spend + reservation_cost - if current_spend > counter.max_budget: - reservation_cost = await _apply_over_budget_reservation_policy( - counter=counter, - valid_token=valid_token, - entry=entry, - applied_entries=applied_entries, - reservation_cost=reservation_cost, - current_spend=current_spend, - fail_closed_budget_enforcement=fail_closed_budget_enforcement, - ) - continue + if reserved_value is not None: + current_spend = reserved_value + else: + cached_spend = current_spend_by_counter_key.get(counter.counter_key) + if cached_spend is None: + cached_spend = await _get_current_counter_value(counter=counter) + current_spend = cached_spend + reservation_cost + if current_spend > counter.max_budget: + reservation_cost = await _apply_over_budget_reservation_policy( + counter=counter, + valid_token=valid_token, + entry=entry, + applied_entries=applied_entries, + reservation_cost=reservation_cost, + current_spend=current_spend, + fail_closed_budget_enforcement=fail_closed_budget_enforcement, + ) + continue except Exception: await _release_applied_entries_best_effort( entries=applied_entries, @@ -878,67 +880,92 @@ async def _get_current_counter_value(counter: _BudgetCounter) -> float: ) +def _counters_batch_scope(counter_keys: frozenset[str]) -> spend_counter_batch_scope: + """Each counter is read once, then written, so one MGET up front serves every read in the loop.""" + from litellm.proxy.proxy_server import spend_counter_cache + + return spend_counter_batch_scope(spend_counter_cache.redis_cache, counter_keys=counter_keys) + + +@dataclass(frozen=True, slots=True) +class _EntryAdjustment: + entry: dict[str, float | str] + counter_key: str + target_adjustment: float + adjustment: float + + +def _entry_adjustment( + entry: dict[str, float | str], actual_cost: float, default_reserved_cost: float +) -> _EntryAdjustment | None: + counter_key: Final = entry.get("counter_key") + if counter_key is None: + return None + target_adjustment: Final = actual_cost - _get_entry_reserved_cost( + entry=entry, default_reserved_cost=default_reserved_cost + ) + adjustment: Final = target_adjustment - float(entry.get("applied_adjustment") or 0.0) + if adjustment == 0: + return None + return _EntryAdjustment( + entry=entry, counter_key=str(counter_key), target_adjustment=target_adjustment, adjustment=adjustment + ) + + async def _set_reserved_entries_actual_cost( entries: list[dict], actual_cost: float, default_reserved_cost: float, reseed_on_inconsistent: bool = True, ) -> None: - for entry in entries: - await _set_reserved_entry_actual_cost( - entry=entry, - actual_cost=actual_cost, - default_reserved_cost=default_reserved_cost, - reseed_on_inconsistent=reseed_on_inconsistent, + """Every reserved counter is read from one MGET and the consistent adjustments go out in one pipeline. + A counter that was flushed or reseeded since reservation is settled on its own after the pipeline.""" + from litellm.proxy.proxy_server import increment_spend_counters_pipeline + + with _counters_batch_scope(frozenset(str(entry["counter_key"]) for entry in entries if "counter_key" in entry)): + adjustments: Final = tuple( + adjustment + for entry in entries + if (adjustment := _entry_adjustment(entry, actual_cost, default_reserved_cost)) is not None ) - - -async def _set_reserved_entry_actual_cost( - entry: dict, - actual_cost: float, - default_reserved_cost: float, - reseed_on_inconsistent: bool = True, -) -> None: - from litellm.proxy.proxy_server import ( - _increment_spend_counter_cache, - reseed_spend_counter_from_db, - ) - - counter_key: Final = entry.get("counter_key") - if counter_key is None: - return - reserved_cost: Final = _get_entry_reserved_cost( - entry=entry, - default_reserved_cost=default_reserved_cost, - ) - target_adjustment: Final = actual_cost - reserved_cost - applied_adjustment: Final = float(entry.get("applied_adjustment") or 0.0) - adjustment: Final = target_adjustment - applied_adjustment - if adjustment == 0: - return - if await _counter_can_apply_adjustment( - counter_key=counter_key, - adjustment=adjustment, - ): - await _increment_spend_counter_cache( - counter_key=counter_key, - increment=adjustment, + consistent: Final = tuple( + await asyncio.gather( + *( + _counter_can_apply_adjustment(counter_key=item.counter_key, adjustment=item.adjustment) + for item in adjustments + ) + ) ) - elif reseed_on_inconsistent: - # Post-call reconcile / release: the counter was flushed, expired or reseeded - # between reservation and reconcile, so the optimistic delta no longer applies. - # Reseed from the DB floor (which cannot include this request's cost yet) and - # add the settled cost, since increment_spend_counters skips reserved keys. - reseeded: Final = await reseed_spend_counter_from_db(counter_key=counter_key) - if reseeded and actual_cost > 0: - await _increment_spend_counter_cache(counter_key=counter_key, increment=actual_cost) - else: - # Pre-call admission resize: the in-flight reservation cost is not yet - # persisted, so the DB floor would discard it. Keep the original - # fail-closed behavior (raise -> reserve_budget_for_request releases and - # denies) rather than admitting against an inconsistent counter. - raise RuntimeError(f"Cannot resize budget reservation against inconsistent counter {counter_key}") - entry["applied_adjustment"] = target_adjustment + inconsistent: Final = tuple(item for item, ok in zip(adjustments, consistent) if not ok) + if inconsistent and not reseed_on_inconsistent: + # Pre-call admission resize: the in-flight reservation cost is not yet + # persisted, so the DB floor would discard it. Keep the original + # fail-closed behavior (raise -> reserve_budget_for_request releases and + # denies) rather than admitting against an inconsistent counter. + raise RuntimeError( + f"Cannot resize budget reservation against inconsistent counter {inconsistent[0].counter_key}" + ) + applicable: Final = tuple(item for item, ok in zip(adjustments, consistent) if ok) + await increment_spend_counters_pipeline( + pending=tuple( + PendingSpendIncrement(counter_key=item.counter_key, increment=item.adjustment) for item in applicable + ) + ) + for item in inconsistent: + await _reseed_reserved_entry(item=item, actual_cost=actual_cost) + for item in adjustments: + item.entry["applied_adjustment"] = item.target_adjustment + + +async def _reseed_reserved_entry(item: _EntryAdjustment, actual_cost: float) -> None: + """Post-call reconcile / release of a counter that was flushed, expired or reseeded between reservation and + reconcile: the optimistic delta no longer applies, so reseed from the DB floor (which cannot include this + request's cost yet) and add the settled cost, since increment_spend_counters skips reserved keys.""" + from litellm.proxy.proxy_server import _increment_spend_counter_cache, reseed_spend_counter_from_db + + reseeded: Final = await reseed_spend_counter_from_db(counter_key=item.counter_key) + if reseeded and actual_cost > 0: + await _increment_spend_counter_cache(counter_key=item.counter_key, increment=actual_cost) async def _counter_can_apply_adjustment( @@ -963,8 +990,8 @@ async def _release_applied_entries_best_effort( ) -> None: for entry in entries: try: - await _set_reserved_entry_actual_cost( - entry=entry, + await _set_reserved_entries_actual_cost( + entries=[entry], # mutable-ok: the reconcile takes the reservation's list of entries actual_cost=0.0, default_reserved_cost=default_reserved_cost, ) diff --git a/litellm/proxy/spend_tracking/carried_budget_state.py b/litellm/proxy/spend_tracking/carried_budget_state.py new file mode 100644 index 00000000000..efd3a78d211 --- /dev/null +++ b/litellm/proxy/spend_tracking/carried_budget_state.py @@ -0,0 +1,60 @@ +"""Pins the budget state auth resolved onto ``UserAPIKeyAuth`` and emits it as request metadata.""" + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from litellm.models.organization import LiteLLM_OrganizationTable +from litellm.models.team import LiteLLM_TeamTable +from litellm.models.user import LiteLLM_UserTable +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.proxy.carried_budget_state import ( + OrgBudgetSnapshot, + TeamBudgetSnapshot, + UserBudgetSnapshot, +) + + +def carry_team_and_user_budget_state( + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + user_object: LiteLLM_UserTable | None, +) -> None: + if team_object is not None: + valid_token.team_budget_snapshot = TeamBudgetSnapshot( # rebind-ok: the request credential is pinned in place + budget_reset_at=team_object.budget_reset_at, + max_budget=team_object.max_budget, + ) + if user_object is not None: + valid_token.user_budget_snapshot = UserBudgetSnapshot( # rebind-ok: same object the caller keeps using + budget_reset_at=user_object.budget_reset_at, + max_budget=user_object.max_budget, + user_alias=user_object.user_alias, + ) + + +def carry_organization_budget_state(valid_token: UserAPIKeyAuth, org_table: LiteLLM_OrganizationTable) -> None: + budget_table: Final = org_table.litellm_budget_table + valid_token.organization_alias = ( + org_table.organization_alias + ) # rebind-ok: the request credential is pinned in place + valid_token.org_budget_snapshot = OrgBudgetSnapshot( # rebind-ok: same object the caller keeps using + spend=org_table.spend, + max_budget=budget_table.max_budget if budget_table is not None else None, + ) + + +def carried_budget_metadata(valid_token: UserAPIKeyAuth) -> Mapping[str, object]: + snapshots: Final = ( + valid_token.team_budget_snapshot, + valid_token.user_budget_snapshot, + valid_token.org_budget_snapshot, + ) + return MappingProxyType( + { + key: value + for snapshot in snapshots + if snapshot is not None + for key, value in snapshot.metadata_entries().items() + } + ) diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index 52e6e98b500..7106d88c655 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/litellm/proxy/spend_tracking/spend_counter_batch.py @@ -1,8 +1,9 @@ -"""One Redis MGET for every spend counter the admission checks read, instead of one GET per counter.""" +"""One Redis MGET per phase (admission, reservation, post-call) for the spend counters it reads, not one GET each.""" import asyncio -from collections.abc import Iterator, Mapping +from collections.abc import Iterator, Mapping, Sequence from contextvars import ContextVar, Token +from dataclasses import dataclass from types import MappingProxyType, TracebackType from typing import Final @@ -11,11 +12,18 @@ from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import RedisCache from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import model_access_group_spend_counter_key _CounterValues: Final = TypeAdapter(dict[str, float | None]) _NO_VALUES: Final[Mapping[str, float | None]] = MappingProxyType({}) +@dataclass(frozen=True, slots=True) +class PendingSpendIncrement: + counter_key: str + increment: float + + class SpendCounterBatch: """Bound counters are read with one MGET on first use; counters bound later join the next MGET. ``async_batch_get_cache`` maps a clean miss to ``None`` and drops keys only when Redis failed, so an absent @@ -35,6 +43,10 @@ class SpendCounterBatch: def counter_keys(self) -> frozenset[str]: return self._keys + @property + def is_open(self) -> bool: + return self._open + def bind(self, counter_keys: frozenset[str]) -> None: if self._open: self._keys = self._keys | counter_keys @@ -52,12 +64,29 @@ class SpendCounterBatch: return None return loaded[counter_key], True + def record(self, counter_key: str, value: float) -> None: + """A write returned the counter's new value; later reads in this scope see it instead of the MGET value.""" + if not self._open: + return + key: Final = frozenset((counter_key,)) + self._keys = self._keys | key + self._fetched = self._fetched | key + self._loaded = MappingProxyType({**self._loaded, counter_key: value}) + + def forget(self, counter_key: str) -> None: + """A write left the counter's value unknown; later reads in this scope go to Redis.""" + key: Final = frozenset((counter_key,)) + self._keys = self._keys - key + self._fetched = self._fetched - key + self._loaded = MappingProxyType({k: v for k, v in self._loaded.items() if k != counter_key}) + async def _load(self) -> Mapping[str, float | None]: async with self._lock: pending: Final = self._keys - self._fetched if pending: self._fetched = self._fetched | pending - self._loaded = MappingProxyType({**self._loaded, **await self._fetch(pending)}) + fetched: Final = await self._fetch(pending) + self._loaded = MappingProxyType({**fetched, **self._loaded}) return self._loaded async def _fetch(self, keys: frozenset[str]) -> Mapping[str, float | None]: @@ -78,17 +107,26 @@ def active_spend_counter_batch() -> SpendCounterBatch | None: class spend_counter_batch_scope: - """Reads inside the scope share one MGET once ``bind_admission_counter_keys`` has run.""" + """Reads inside the scope share one MGET for the keys bound here or by ``bind_*`` calls inside it. + Opened inside a scope whose batch is still open, it binds into that batch so both phases share the MGET.""" - __slots__ = ("_redis_cache", "_token") + __slots__ = ("_counter_keys", "_redis_cache", "_token") - def __init__(self, redis_cache: RedisCache | None) -> None: + def __init__(self, redis_cache: RedisCache | None, counter_keys: frozenset[str] = frozenset()) -> None: self._redis_cache: Final = redis_cache + self._counter_keys: Final = counter_keys self._token: Token[SpendCounterBatch | None] | None = None def __enter__(self) -> None: - if self._redis_cache is not None: - self._token = _active_batch.set(SpendCounterBatch(self._redis_cache)) + if self._redis_cache is None: + return + outer: Final = _active_batch.get() + if outer is not None and outer.is_open: + outer.bind(self._counter_keys) + return + batch: Final = SpendCounterBatch(self._redis_cache) + batch.bind(self._counter_keys) + self._token = _active_batch.set(batch) def __exit__( self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None @@ -122,9 +160,58 @@ def admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> fr return frozenset(_iter_admission_counter_keys(token, end_user_id)) +def post_call_counter_keys( + token: str | None, + team_id: str | None, + user_id: str | None, + org_id: str | None, + end_user_id: str | None, + tags: Sequence[object] | None, + model_access_groups: Sequence[object] | None, +) -> frozenset[str]: + """Every counter ``increment_spend_counters`` warm-checks, except budget windows which bind on read.""" + entity_keys: Final = admission_counter_keys( + UserAPIKeyAuth(token=token, team_id=team_id, user_id=user_id, org_id=org_id), end_user_id + ) + tag_keys: Final = frozenset(f"spend:tag:{tag}" for tag in tags or () if tag and isinstance(tag, str)) + group_keys: Final = frozenset( + model_access_group_spend_counter_key(group) + for group in model_access_groups or () + if group and isinstance(group, str) + ) + return entity_keys | tag_keys | group_keys + + def bind_admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> None: """Idempotent: call again after the token gains ids (end user, team org) so those counters join the MGET.""" + bind_spend_counter_keys(admission_counter_keys(token, end_user_id)) + + +def bind_spend_counter_keys(counter_keys: frozenset[str]) -> None: batch: Final = _active_batch.get() if batch is None: return - batch.bind(admission_counter_keys(token, end_user_id)) + batch.bind(counter_keys) + + +def record_spend_counter_value(counter_key: str, value: float) -> None: + batch: Final = _active_batch.get() + if batch is None: + return + batch.record(counter_key, value) + + +def forget_spend_counter(counter_key: str) -> None: + batch: Final = _active_batch.get() + if batch is None: + return + batch.forget(counter_key) + + +async def read_batched_spend_counter(counter_key: str) -> tuple[float | None, bool] | None: + """Bind-on-read for counters only known at read time (budget windows); the first reader pays the MGET.""" + batch: Final = _active_batch.get() + if batch is None: + return None + batch.bind(frozenset((counter_key,))) + return await batch.read(counter_key) diff --git a/litellm/types/proxy/carried_budget_state.py b/litellm/types/proxy/carried_budget_state.py new file mode 100644 index 00000000000..0e64b91645d --- /dev/null +++ b/litellm/types/proxy/carried_budget_state.py @@ -0,0 +1,64 @@ +"""Budget fields auth already resolved, carried on the request so success logging does no object lookups. + +Auth pins one snapshot per entity on ``UserAPIKeyAuth`` (request-scoped, never cached), the pre-call +setup writes them into the request metadata under the aliased ``user_api_key_*`` names, and a logger +reads them back with ``from_metadata``. ``None`` means this request never carried that entity +(unauthenticated route, custom auth, budget check skipped) and the logger keeps its own lookup. +""" + +from collections.abc import Mapping +from datetime import datetime +from types import MappingProxyType + +from pydantic import BaseModel, ConfigDict, Field, ValidationError +from typing_extensions import Self + + +class _BudgetSnapshot(BaseModel): + model_config = ConfigDict(frozen=True, populate_by_name=True, extra="ignore") + + def metadata_entries(self) -> Mapping[str, object]: + return MappingProxyType(self.model_dump(by_alias=True, mode="json")) + + @classmethod + def from_metadata(cls, metadata: Mapping[str, object]) -> Self | None: + try: + return cls.model_validate(metadata) + except ValidationError: + return None + + +class KeyBudgetSnapshot(_BudgetSnapshot): + """Read-only view of the ``user_api_key_budget_reset_at`` entry ``add_user_api_key_auth_to_request_metadata`` writes.""" + + budget_reset_at: datetime | None = Field( + validation_alias="user_api_key_budget_reset_at", serialization_alias="user_api_key_budget_reset_at" + ) + + +class TeamBudgetSnapshot(_BudgetSnapshot): + budget_reset_at: datetime | None = Field( + validation_alias="user_api_key_team_budget_reset_at", serialization_alias="user_api_key_team_budget_reset_at" + ) + max_budget: float | None = Field( + validation_alias="user_api_key_team_table_max_budget", serialization_alias="user_api_key_team_table_max_budget" + ) + + +class UserBudgetSnapshot(_BudgetSnapshot): + budget_reset_at: datetime | None = Field( + validation_alias="user_api_key_user_budget_reset_at", serialization_alias="user_api_key_user_budget_reset_at" + ) + max_budget: float | None = Field( + validation_alias="user_api_key_user_table_max_budget", serialization_alias="user_api_key_user_table_max_budget" + ) + user_alias: str | None = Field( + validation_alias="user_api_key_user_alias", serialization_alias="user_api_key_user_alias" + ) + + +class OrgBudgetSnapshot(_BudgetSnapshot): + spend: float = Field(validation_alias="user_api_key_org_spend", serialization_alias="user_api_key_org_spend") + max_budget: float | None = Field( + validation_alias="user_api_key_org_max_budget", serialization_alias="user_api_key_org_max_budget" + ) diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 3c93687d467..4bf7894829e 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1205,6 +1205,92 @@ async def test_a_probe_overtaken_by_a_later_outage_leaves_the_breaker_to_the_new assert breaker._state == breaker.CLOSED +class _RoundTripCountingRedis: + """Fake redis.asyncio client: one round trip per awaited command or pipeline execute.""" + + def __init__(self, ttl: int) -> None: + self.values: dict[str, float] = {} + self.ttls: dict[str, int] = {} + self.round_trips = 0 + self._initial_ttl = ttl + + async def incrbyfloat(self, name: str, amount: float) -> float: + self.round_trips += 1 + return self._incr(name, amount) + + async def expire(self, name: str, time: int) -> bool: + self.round_trips += 1 + self.ttls[name] = time + return True + + def _incr(self, name: str, amount: float) -> float: + self.values[name] = self.values.get(name, 0.0) + amount + self.ttls.setdefault(name, self._initial_ttl) + return self.values[name] + + def pipeline(self, transaction: bool) -> "_RoundTripCountingRedis._Pipeline": + return _RoundTripCountingRedis._Pipeline(self) + + class _Pipeline: + def __init__(self, client: "_RoundTripCountingRedis") -> None: + self._client = client + self._commands: list[tuple[str, tuple[object, ...]]] = [] + + async def __aenter__(self) -> "_RoundTripCountingRedis._Pipeline": + return self + + async def __aexit__(self, *exc_info: object) -> None: + return None + + def incrbyfloat(self, name: str, amount: float) -> None: + self._commands.append(("incrbyfloat", (name, amount))) + + def expire(self, name: str, time: int) -> None: + self._commands.append(("expire", (name, time))) + + def ttl(self, name: str) -> None: + self._commands.append(("ttl", (name,))) + + async def execute(self) -> list[object]: + self._client.round_trips += 1 + results: list[object] = [] + for command, args in self._commands: + if command == "incrbyfloat": + results.append(self._client._incr(str(args[0]), float(args[1]))) # pyright: ignore[reportArgumentType] # fake stores str/float + elif command == "expire": + self._client.ttls[str(args[0])] = int(args[1]) # pyright: ignore[reportArgumentType] # fake stores int + results.append(True) + else: + results.append(self._client.ttls.get(str(args[0]), -2)) + return results + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("refresh_ttl", "existing_ttl", "expected_round_trips", "expected_ttl"), + [ + pytest.param(True, 100, 2, 60, id="refresh_ttl: INCRBYFLOAT+EXPIRE in one round trip each"), + pytest.param(False, 100, 2, 100, id="keep ttl: INCRBYFLOAT+TTL in one round trip each, no EXPIRE"), + pytest.param(False, -1, 3, 60, id="unexpiring key: INCRBYFLOAT+TTL then EXPIRE once, 1 trip after"), + ], +) +async def test_async_increment_pipelines_the_ttl_command( + monkeypatch, redis_no_ping, refresh_ttl, existing_ttl, expected_round_trips, expected_ttl +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace="ns") + client = _RoundTripCountingRedis(ttl=existing_ttl) + + with patch.object(redis_cache, "init_async_client", return_value=client): + first = await redis_cache.async_increment(key="spend:key:k", value=1.5, ttl=60, refresh_ttl=refresh_ttl) + second = await redis_cache.async_increment(key="spend:key:k", value=2.0, ttl=60, refresh_ttl=refresh_ttl) + + assert (first, second) == (1.5, 3.5) + assert client.values == {"ns:spend:key:k": 3.5} + assert client.ttls == {"ns:spend:key:k": expected_ttl} + assert client.round_trips == expected_round_trips + + class _SetRecordingPipeline: def __init__(self) -> None: self.sets: list[tuple[str, str, timedelta | None]] = [] diff --git a/tests/test_litellm/integrations/test_prometheus_carried_budget_state.py b/tests/test_litellm/integrations/test_prometheus_carried_budget_state.py new file mode 100644 index 00000000000..e49e6baa8b0 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_carried_budget_state.py @@ -0,0 +1,289 @@ +""" +Post-request budget gauges read the key/team/user/org state auth already resolved +from request metadata. get_*_object only runs when that state is missing (custom +auth, SDK callers, failure paths) +""" + +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from prometheus_client import REGISTRY + +from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_OrganizationTable, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.spend_tracking.carried_budget_state import ( + carried_budget_metadata, + carry_organization_budget_state, + carry_team_and_user_budget_state, +) +from litellm.types.proxy.carried_budget_state import ( + KeyBudgetSnapshot, + TeamBudgetSnapshot, + UserBudgetSnapshot, +) + +TEAM_RESET_AT = datetime(2026, 10, 1, tzinfo=timezone.utc) +USER_RESET_AT = datetime(2026, 11, 1, tzinfo=timezone.utc) +KEY_RESET_AT = datetime(2026, 12, 1, tzinfo=timezone.utc) + + +@pytest.fixture(autouse=True) +def cleanup_prometheus_registry(): + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + yield + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + + +@pytest.fixture +def prometheus_logger(): + return PrometheusLogger() + + +@pytest.fixture +def getters(): + """Every response-path object getter, patched where prometheus imports them from.""" + mocks = { + "get_key_object": AsyncMock(return_value=UserAPIKeyAuth(token="hashed", budget_reset_at=KEY_RESET_AT)), + "get_team_object": AsyncMock( + return_value=LiteLLM_TeamTable(team_id="t1", budget_reset_at=TEAM_RESET_AT, max_budget=300.0) + ), + "get_user_object": AsyncMock( + return_value=LiteLLM_UserTable( + user_id="u1", + budget_reset_at=USER_RESET_AT, + user_email="alice@example.com", + user_alias="Alice", + max_budget=50.0, + ) + ), + "get_org_object": AsyncMock( + return_value=LiteLLM_OrganizationTable( + organization_id="o1", + organization_alias="platform-org", + budget_id="b1", + created_by="admin", + updated_by="admin", + spend=40.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=500.0), + ) + ), + } + with ( + patch.multiple( # test-quality-ok: prometheus reads these proxy_server globals at call time, no injection seam + "litellm.proxy.proxy_server", prisma_client=MagicMock(), user_api_key_cache=MagicMock() + ), + patch.multiple( # test-quality-ok: the getters are the DB boundary this test counts calls to + "litellm.proxy.auth.auth_checks", **mocks + ), + ): + yield mocks + + +def _authed_token() -> UserAPIKeyAuth: + token = UserAPIKeyAuth( + token="hashed", + key_alias="key-alias", + team_id="t1", + team_alias="team-alias", + user_id="u1", + user_email="alice@example.com", + org_id="o1", + spend=1.0, + max_budget=10.0, + team_spend=20.0, + team_max_budget=300.0, + user_spend=5.0, + user_max_budget=50.0, + budget_reset_at=KEY_RESET_AT, + ) + carry_team_and_user_budget_state( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", budget_reset_at=TEAM_RESET_AT, max_budget=300.0), + user_object=LiteLLM_UserTable(user_id="u1", budget_reset_at=USER_RESET_AT, user_alias="Alice", max_budget=50.0), + ) + carry_organization_budget_state( + valid_token=token, + org_table=LiteLLM_OrganizationTable( + organization_id="o1", + organization_alias="platform-org", + budget_id="b1", + created_by="admin", + updated_by="admin", + spend=40.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=500.0), + ), + ) + return token + + +def _request_metadata(token: UserAPIKeyAuth) -> dict: + """What add_user_api_key_auth_to_request_metadata leaves in litellm_params["metadata"].""" + return { + **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(token), + **carried_budget_metadata(token), + } + + +def _stub_gauges(prometheus_logger: PrometheusLogger) -> None: + for name in ( + "litellm_remaining_api_key_budget_metric", + "litellm_api_key_max_budget_metric", + "litellm_api_key_budget_remaining_hours_metric", + "litellm_remaining_team_budget_metric", + "litellm_team_max_budget_metric", + "litellm_team_budget_remaining_hours_metric", + "litellm_remaining_user_budget_metric", + "litellm_user_max_budget_metric", + "litellm_user_budget_remaining_hours_metric", + "litellm_remaining_org_budget_metric", + "litellm_org_max_budget_metric", + "litellm_org_budget_remaining_hours_metric", + ): + setattr(prometheus_logger, name, MagicMock()) + + +async def _emit(prometheus_logger: PrometheusLogger, metadata: dict) -> None: + await prometheus_logger._increment_remaining_budget_metrics( + user_api_team="t1", + user_api_team_alias="team-alias", + user_api_key="hashed", + user_api_key_alias="key-alias", + litellm_params={"metadata": metadata}, + response_cost=2.0, + user_id="u1", + user_api_key_org_id="o1", + ) + + +@pytest.mark.asyncio +async def test_authed_request_sets_every_gauge_without_any_object_getter(prometheus_logger, getters): + _stub_gauges(prometheus_logger) + + await _emit(prometheus_logger, _request_metadata(_authed_token())) + + assert all(getter.await_count == 0 for getter in getters.values()), { + name: getter.await_count for name, getter in getters.items() + } + remaining = { + "key": prometheus_logger.litellm_remaining_api_key_budget_metric.labels().set.call_args[0][0], + "team": prometheus_logger.litellm_remaining_team_budget_metric.labels().set.call_args[0][0], + "user": prometheus_logger.litellm_remaining_user_budget_metric.labels().set.call_args[0][0], + "org": prometheus_logger.litellm_remaining_org_budget_metric.labels().set.call_args[0][0], + } + assert remaining == { + "key": pytest.approx(7.0), + "team": pytest.approx(278.0), + "user": pytest.approx(43.0), + "org": 458.0, + } + prometheus_logger.litellm_org_max_budget_metric.labels().set.assert_called_once_with(500.0) + prometheus_logger.litellm_api_key_budget_remaining_hours_metric.labels().set.assert_called_once() + prometheus_logger.litellm_team_budget_remaining_hours_metric.labels().set.assert_called_once() + prometheus_logger.litellm_user_budget_remaining_hours_metric.labels().set.assert_called_once() + + +@pytest.mark.asyncio +async def test_metadata_without_carried_state_still_fetches_each_object_once(prometheus_logger, getters): + _stub_gauges(prometheus_logger) + + await _emit(prometheus_logger, {"user_api_key_team_spend": 20.0, "user_api_key_team_max_budget": 300.0}) + + assert {name: getter.await_count for name, getter in getters.items()} == { + "get_key_object": 1, + "get_team_object": 1, + "get_user_object": 1, + "get_org_object": 1, + } + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(458.0) + prometheus_logger.litellm_team_budget_remaining_hours_metric.labels().set.assert_called_once() + + +@pytest.mark.asyncio +async def test_partial_carried_state_only_skips_the_carried_objects(prometheus_logger, getters): + _stub_gauges(prometheus_logger) + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1") + carry_team_and_user_budget_state( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", budget_reset_at=TEAM_RESET_AT), + user_object=None, + ) + + await _emit(prometheus_logger, dict(carried_budget_metadata(token))) + + assert {name: getter.await_count for name, getter in getters.items()} == { + "get_key_object": 1, + "get_team_object": 0, + "get_user_object": 1, + "get_org_object": 1, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_max_budget", [300.0, None], ids=["metadata has max_budget", "filled from object"]) +async def test_carried_objects_match_what_the_getters_would_have_produced( + prometheus_logger, getters, metadata_max_budget +): + metadata = _request_metadata(_authed_token()) + user_max_budget = 50.0 if metadata_max_budget is not None else None + + carried_team = await prometheus_logger._assemble_team_object( + team_id="t1", + team_alias="team-alias", + spend=20.0, + max_budget=metadata_max_budget, + response_cost=2.0, + carried=TeamBudgetSnapshot.from_metadata(metadata), + ) + fetched_team = await prometheus_logger._assemble_team_object( + team_id="t1", team_alias="team-alias", spend=20.0, max_budget=metadata_max_budget, response_cost=2.0 + ) + carried_user = await prometheus_logger._assemble_user_object( + user_id="u1", + spend=5.0, + max_budget=user_max_budget, + response_cost=2.0, + carried=UserBudgetSnapshot.from_metadata(metadata), + user_email="alice@example.com", + ) + fetched_user = await prometheus_logger._assemble_user_object( + user_id="u1", spend=5.0, max_budget=user_max_budget, response_cost=2.0 + ) + carried_key = await prometheus_logger._assemble_key_object( + user_api_key="hashed", + user_api_key_alias="key-alias", + key_max_budget=10.0, + key_spend=1.0, + response_cost=2.0, + carried=KeyBudgetSnapshot.from_metadata(metadata), + ) + fetched_key = await prometheus_logger._assemble_key_object( + user_api_key="hashed", user_api_key_alias="key-alias", key_max_budget=10.0, key_spend=1.0, response_cost=2.0 + ) + + assert carried_team == fetched_team + assert carried_team.max_budget == 300.0 + assert carried_user == fetched_user + assert carried_user.max_budget == 50.0 + assert carried_key == fetched_key + assert {name: getter.await_count for name, getter in getters.items()} == { + "get_key_object": 1, + "get_team_object": 1, + "get_user_object": 1, + "get_org_object": 0, + } diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index dc614d18662..d9a065db3cb 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5791,6 +5791,42 @@ async def test_budget_checks_only_run_on_llm_api_routes(scope, route, expect_blo assert await _run() is True +@pytest.mark.asyncio +async def test_organization_budget_check_carries_org_state_on_the_token(): + """The org row auth already fetched is pinned on the token so the response path + (Prometheus org budget gauges) reads it from request metadata instead of calling + get_org_object again.""" + from litellm.proxy._types import LiteLLM_OrganizationTable + from litellm.proxy.auth.auth_checks import _organization_max_budget_check + from litellm.types.proxy.carried_budget_state import OrgBudgetSnapshot + + org_table = LiteLLM_OrganizationTable( + organization_id="o1", + organization_alias="platform-org", + budget_id="b1", + created_by="admin", + updated_by="admin", + spend=12.5, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0), + ) + token = UserAPIKeyAuth(token="k1", org_id="o1") + user_api_key_cache = UserApiKeyCache() + await user_api_key_cache.async_set_cache( + key="org_id:o1:with_budget", value=org_table, model_type=LiteLLM_OrganizationTable + ) + + await _organization_max_budget_check( + valid_token=token, + team_object=None, + prisma_client=MagicMock(), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=MagicMock(), + ) + + assert token.organization_alias == "platform-org" + assert token.org_budget_snapshot == OrgBudgetSnapshot(spend=12.5, max_budget=100.0) + + @pytest.mark.parametrize("route", ["/health", "/health/services", "/health/test_connection"]) @pytest.mark.asyncio async def test_spend_capable_non_llm_routes_still_enforce_budget(route): diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index a257288ebe0..0cdbcde6abc 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -5,7 +5,7 @@ import os import subprocess import sys from contextlib import contextmanager -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from pathlib import Path from textwrap import dedent from types import SimpleNamespace @@ -23,6 +23,7 @@ from litellm.proxy._types import ( LiteLLM_JWTAuth, LiteLLM_BudgetTable, LiteLLM_EndUserTable, + LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LitellmUserRoles, ProxyErrorTypes, @@ -48,6 +49,7 @@ from litellm.proxy.auth.user_api_key_auth import ( get_api_key, user_api_key_auth, ) +from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata class _RoutingRequest: @@ -4019,6 +4021,65 @@ async def test_centralized_common_checks_routes_header_tags_to_litellm_metadata( assert "metadata" not in request_data +@pytest.mark.asyncio +async def test_centralized_common_checks_carries_team_and_user_budget_state_on_the_token(): + """The team and user objects auth resolves are pinned on the token so the + response path (Prometheus budget gauges) reads them from request metadata + instead of calling get_team_object / get_user_object again.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + + reset_at = datetime(2026, 10, 1, tzinfo=timezone.utc) + token = UserAPIKeyAuth(api_key="sk-test", token="hashed", team_id="t1", user_id="u1") + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache( + key="team_id:t1", + value=LiteLLM_TeamTableCachedObj(team_id="t1", budget_reset_at=reset_at, max_budget=300.0), + ) + await user_api_key_cache.async_set_cache( + key="u1", + value=LiteLLM_UserTable(user_id="u1", user_alias="Alice", budget_reset_at=None, max_budget=None), + ) + proxy_logging_obj = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock() + attrs = { + **_proxy_attrs_for_centralized_checks(user_custom_auth=None), + "prisma_client": MagicMock(), + "user_api_key_cache": user_api_key_cache, + "proxy_logging_obj": proxy_logging_obj, + } + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with patch( # test-quality-ok: the authz gate has its own tests above; this one checks the carry step before it + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-5.4-mini"}, + route="/chat/completions", + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + assert dict(carried_budget_metadata(token)) == { + "user_api_key_team_budget_reset_at": "2026-10-01T00:00:00Z", + "user_api_key_team_table_max_budget": 300.0, + "user_api_key_user_budget_reset_at": None, + "user_api_key_user_table_max_budget": None, + "user_api_key_user_alias": "Alice", + } + + @pytest.mark.asyncio async def test_centralized_common_checks_skipped_for_custom_auth_without_flag(): """Existing RPS guarantee: custom-auth deployments without diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index 69ec174903b..bca6344b3f7 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -350,7 +350,7 @@ async def test_cold_reseed_preserves_concurrent_local_increment( increment_task: Final = asyncio.create_task( proxy_server._apply_spend_counter_increments( - pending=(proxy_server._PendingSpendIncrement(counter_key=counter_key, increment=increment),) + pending=(proxy_server.PendingSpendIncrement(counter_key=counter_key, increment=increment),) ) if batch else proxy_server._increment_spend_counter_cache(counter_key=counter_key, increment=increment) diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 2f47736a398..6eac53df645 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -1151,10 +1151,10 @@ async def test_prepare_window_spend_counter_increment_missing_window_start_inval # --------------------------------------------------------------------------- -def _two_pending_increments() -> tuple[ps._PendingSpendIncrement, ...]: +def _two_pending_increments() -> tuple[ps.PendingSpendIncrement, ...]: return ( - ps._PendingSpendIncrement(counter_key="spend:key:k", increment=1.5), - ps._PendingSpendIncrement(counter_key="spend:team:t", increment=1.5), + ps.PendingSpendIncrement(counter_key="spend:key:k", increment=1.5), + ps.PendingSpendIncrement(counter_key="spend:team:t", increment=1.5), ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py b/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py new file mode 100644 index 00000000000..fe852be775c --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py @@ -0,0 +1,119 @@ +"""Auth-resolved budget state rides on ``UserAPIKeyAuth`` and round-trips through request metadata.""" + +from datetime import datetime, timezone + +from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_OrganizationTable, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.spend_tracking.carried_budget_state import ( + carried_budget_metadata, + carry_organization_budget_state, + carry_team_and_user_budget_state, +) +from litellm.types.proxy.carried_budget_state import ( + KeyBudgetSnapshot, + OrgBudgetSnapshot, + TeamBudgetSnapshot, + UserBudgetSnapshot, +) + +RESET_AT = datetime(2026, 10, 1, 12, 30, tzinfo=timezone.utc) + + +def test_team_and_user_state_round_trips_through_metadata(): + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1") + carry_team_and_user_budget_state( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", budget_reset_at=RESET_AT, max_budget=300.0), + user_object=LiteLLM_UserTable(user_id="u1", budget_reset_at=None, max_budget=None, user_alias="Alice"), + ) + + metadata = dict(carried_budget_metadata(token)) + + assert metadata == { + "user_api_key_team_budget_reset_at": RESET_AT.isoformat().replace("+00:00", "Z"), + "user_api_key_team_table_max_budget": 300.0, + "user_api_key_user_budget_reset_at": None, + "user_api_key_user_table_max_budget": None, + "user_api_key_user_alias": "Alice", + } + assert TeamBudgetSnapshot.from_metadata(metadata) == TeamBudgetSnapshot(budget_reset_at=RESET_AT, max_budget=300.0) + assert UserBudgetSnapshot.from_metadata(metadata) == UserBudgetSnapshot( + budget_reset_at=None, max_budget=None, user_alias="Alice" + ) + + +def test_missing_objects_leave_no_metadata_and_no_snapshot(): + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1") + carry_team_and_user_budget_state(valid_token=token, team_object=None, user_object=None) + + assert dict(carried_budget_metadata(token)) == {} + assert TeamBudgetSnapshot.from_metadata({}) is None + assert UserBudgetSnapshot.from_metadata({"user_api_key_user_alias": "Alice"}) is None + assert OrgBudgetSnapshot.from_metadata({"user_api_key_org_spend": 1.0}) is None + assert KeyBudgetSnapshot.from_metadata({}) is None + + +def test_organization_state_carries_alias_spend_and_max_budget(): + token = UserAPIKeyAuth(token="hashed", org_id="o1") + org = LiteLLM_OrganizationTable( + organization_id="o1", + organization_alias="platform-org", + budget_id="b1", + created_by="admin", + updated_by="admin", + spend=12.5, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0), + ) + + carry_organization_budget_state(valid_token=token, org_table=org) + + assert token.organization_alias == "platform-org" + assert OrgBudgetSnapshot.from_metadata(carried_budget_metadata(token)) == OrgBudgetSnapshot( + spend=12.5, max_budget=100.0 + ) + + +def test_organization_without_budget_table_carries_no_cap(): + token = UserAPIKeyAuth(token="hashed", org_id="o1") + org = LiteLLM_OrganizationTable( + organization_id="o1", + budget_id="b1", + created_by="admin", + updated_by="admin", + spend=3.0, + ) + + carry_organization_budget_state(valid_token=token, org_table=org) + + assert token.org_budget_snapshot == OrgBudgetSnapshot(spend=3.0, max_budget=None) + + +def test_key_snapshot_parses_the_iso_string_auth_metadata_writes(): + assert KeyBudgetSnapshot.from_metadata({"user_api_key_budget_reset_at": RESET_AT.isoformat()}) == KeyBudgetSnapshot( + budget_reset_at=RESET_AT + ) + assert KeyBudgetSnapshot.from_metadata({"user_api_key_budget_reset_at": None}) == KeyBudgetSnapshot( + budget_reset_at=None + ) + + +def test_snapshots_never_reach_the_serialized_token(): + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1") + carry_team_and_user_budget_state( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", budget_reset_at=RESET_AT), + user_object=LiteLLM_UserTable(user_id="u1", user_alias="Alice"), + ) + token.org_budget_snapshot = OrgBudgetSnapshot(spend=1.0, max_budget=2.0) + + dumped = token.model_dump() + + assert "team_budget_snapshot" not in dumped + assert "user_budget_snapshot" not in dumped + assert "org_budget_snapshot" not in dumped + assert UserAPIKeyAuth(**dumped).team_budget_snapshot is None diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_counter_batch.py b/tests/test_litellm/proxy/spend_tracking/test_spend_counter_batch.py index 288c3d8c7b4..3e4b817fab8 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_counter_batch.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_counter_batch.py @@ -3,7 +3,7 @@ from __future__ import annotations import asyncio -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from unittest.mock import AsyncMock, MagicMock import pytest @@ -52,6 +52,24 @@ class CountingRedis(RedisCache): self.commands.append(f"MGET {' '.join(key_list)}") return {key: self.store.get(key) for key in key_list} + def get_ttl(self, **kwargs: object) -> int | None: + return None + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + self.commands.append(f"INCRBYFLOAT {key} {value}") + return self._incr(key, value) + + async def async_increment_pipeline( + self, increment_list: Sequence[Mapping[str, object]], **kwargs: object + ) -> list[float]: + self.commands.append(f"PIPELINE {' '.join(str(op['key']) for op in increment_list)}") + return [self._incr(str(op["key"]), float(str(op["increment_value"]))) for op in increment_list] + + def _incr(self, key: str, value: float) -> float: + total = float(str(self.store.get(key, 0.0))) + value + self.store[key] = total + return total + def _spend_counter_cache(redis: RedisCache | None, in_memory: dict[str, float] | None = None) -> MagicMock: cache = MagicMock() @@ -119,6 +137,37 @@ async def test_keys_bound_after_the_first_read_join_one_more_mget_for_only_the_n assert redis.commands == ["MGET spend:key:hashed", "MGET spend:org:org"] +@pytest.mark.asyncio +async def test_a_recorded_write_result_answers_later_reads_without_another_redis_read(): + redis = CountingRedis({"spend:key:hashed": 1.0}) + batch = SpendCounterBatch(redis) + batch.bind(frozenset({"spend:key:hashed"})) + assert await batch.read("spend:key:hashed") == (1.0, True) + + batch.record("spend:key:hashed", 3.5) + batch.record("spend:org:org", 7.0) + + assert await batch.read("spend:key:hashed") == (3.5, True) + assert await batch.read("spend:org:org") == (7.0, True) + assert redis.commands == ["MGET spend:key:hashed"] + + +@pytest.mark.asyncio +async def test_a_forgotten_counter_is_read_fresh_from_redis_when_it_is_bound_again(): + redis = CountingRedis({"spend:key:hashed": 1.0}) + batch = SpendCounterBatch(redis) + batch.bind(frozenset({"spend:key:hashed"})) + batch.record("spend:key:hashed", 3.5) + + batch.forget("spend:key:hashed") + assert await batch.read("spend:key:hashed") is None + + redis.store["spend:key:hashed"] = 9.0 + batch.bind(frozenset({"spend:key:hashed"})) + assert await batch.read("spend:key:hashed") == (9.0, True) + assert redis.commands == ["MGET spend:key:hashed"] + + @pytest.mark.asyncio async def test_failed_mget_hands_every_counter_back_to_the_caller(): batch = SpendCounterBatch(CountingRedis(fail=True)) @@ -298,3 +347,181 @@ async def test_reseed_outside_the_scope_still_re_checks_redis_itself(): assert value == 4.0 assert redis.commands == ["GET spend:key:hashed"] + + +POST_CALL_KEYS = TOKEN_KEYS | {"spend:tag:prod", "spend:model_access_group:premium"} + + +@pytest.mark.asyncio +async def test_post_call_increment_for_every_entity_costs_one_mget_and_one_pipeline(monkeypatch): + redis = CountingRedis({key: 1.0 for key in POST_CALL_KEYS}) + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + monkeypatch.setattr(ps, "prisma_client", None) + + await ps.increment_spend_counters( + token="hashed", + team_id="team", + user_id="user", + org_id="org", + end_user_id="eu", + tags=["prod"], + model_access_groups=["premium"], + response_cost=0.5, + ) + + assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE"], redis.commands + assert set(redis.commands[0].split()[1:]) == POST_CALL_KEYS + assert set(redis.commands[1].split()[1:]) == POST_CALL_KEYS + assert {key: redis.store[key] for key in POST_CALL_KEYS} == {key: 1.5 for key in POST_CALL_KEYS} + + +@pytest.mark.asyncio +async def test_post_call_cold_counters_seed_from_the_mget_miss_without_a_second_read(monkeypatch): + redis = CountingRedis({"spend:key:hashed": 1.0}) + redis.async_set_cache = AsyncMock(return_value=True) + prisma = MagicMock() + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=MagicMock(spend=4.0)) + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + monkeypatch.setattr(ps, "prisma_client", prisma) + + await ps.increment_spend_counters(token="hashed", team_id="team", user_id=None, response_cost=0.5) + + assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE"], redis.commands + redis.async_set_cache.assert_awaited_once_with(key="spend:team:team", value=4.0, nx=True) + assert redis.store["spend:key:hashed"] == 1.5 + + +RESERVED_KEYS = frozenset( + {"spend:key:hashed", "spend:team:team", "spend:team_member:user:team", "spend:end_user:eu", "spend:org:org"} +) + + +def _reservation(reserved_cost: float, counter_keys: frozenset[str] = RESERVED_KEYS) -> dict[str, object]: + return { + "reserved_cost": reserved_cost, + "entries": [ + {"counter_key": key, "entity_type": "Key", "entity_id": key, "reserved_cost": reserved_cost} + for key in sorted(counter_keys) + ], + } + + +@pytest.mark.asyncio +async def test_post_call_with_a_reservation_costs_one_mget_one_reconcile_pipeline_one_increment_pipeline(monkeypatch): + redis = CountingRedis({key: 1.0 for key in POST_CALL_KEYS}) + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + monkeypatch.setattr(ps, "prisma_client", None) + reservation = _reservation(reserved_cost=0.4) + + await ps.increment_spend_counters( + token="hashed", + team_id="team", + user_id="user", + org_id="org", + end_user_id="eu", + tags=["prod"], + model_access_groups=["premium"], + response_cost=0.5, + budget_reservation=reservation, + ) + + assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE", "PIPELINE"], redis.commands + assert set(redis.commands[0].split()[1:]) == POST_CALL_KEYS, "reconcile and warm checks share the MGET" + assert set(redis.commands[1].split()[1:]) == RESERVED_KEYS + assert set(redis.commands[2].split()[1:]) == POST_CALL_KEYS - RESERVED_KEYS + assert {key: round(redis.store[key], 6) for key in POST_CALL_KEYS} == { + key: (1.1 if key in RESERVED_KEYS else 1.5) for key in POST_CALL_KEYS + } + assert [round(entry["applied_adjustment"], 6) for entry in reservation["entries"]] == [0.1] * len(RESERVED_KEYS) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_reconcile_settles_a_flushed_counter_on_its_own_after_the_shared_pipeline(monkeypatch): + from litellm.proxy.spend_tracking.budget_reservation import reconcile_budget_reservation + + redis = CountingRedis({key: 1.0 for key in RESERVED_KEYS - {"spend:team:team"}}) + redis.async_set_max = AsyncMock(return_value=4.0) + prisma = MagicMock() + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=MagicMock(spend=4.0)) + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + monkeypatch.setattr(ps, "prisma_client", prisma) + reservation = _reservation(reserved_cost=0.4) + + await reconcile_budget_reservation(budget_reservation=reservation, actual_cost=0.5) + + assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE", "INCRBYFLOAT"], redis.commands + assert set(redis.commands[1].split()[1:]) == RESERVED_KEYS - {"spend:team:team"} + assert redis.commands[2] == "INCRBYFLOAT spend:team:team 0.5" + redis.async_set_max.assert_awaited_once() + assert redis.async_set_max.await_args.kwargs["key"] == "spend:team:team" + assert all(round(entry["applied_adjustment"], 6) == 0.1 for entry in reservation["entries"]) + + +@pytest.mark.asyncio +async def test_pre_call_resize_against_an_inconsistent_counter_writes_nothing_and_denies(monkeypatch): + from litellm.proxy.spend_tracking.budget_reservation import _resize_applied_reservation + + redis = CountingRedis({"spend:key:hashed": 1.0}) + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + monkeypatch.setattr(ps, "prisma_client", None) + entries = _reservation(reserved_cost=0.4, counter_keys=frozenset({"spend:key:hashed", "spend:team:team"}))[ + "entries" + ] + + with pytest.raises(RuntimeError, match="spend:team:team"): + await _resize_applied_reservation(entries=entries, current_reserved_cost=0.4, new_reserved_cost=0.9) + + assert [c.split()[0] for c in redis.commands] == ["MGET"], redis.commands + assert redis.store["spend:key:hashed"] == 1.0 + assert all("applied_adjustment" not in entry for entry in entries) + + +@pytest.mark.asyncio +async def test_a_failed_reconcile_pipeline_invalidates_every_reserved_counter_and_falls_back(monkeypatch): + redis = CountingRedis({key: 1.0 for key in POST_CALL_KEYS}) + redis.async_delete_cache = AsyncMock() + reconcile_pipeline_failed = False + + async def _pipeline(increment_list: Sequence[Mapping[str, object]], **kwargs: object) -> list[float]: + nonlocal reconcile_pipeline_failed + if not reconcile_pipeline_failed: + reconcile_pipeline_failed = True + raise ConnectionError("redis down") + return await CountingRedis.async_increment_pipeline(redis, increment_list, **kwargs) + + redis.async_increment_pipeline = _pipeline # pyright: ignore[reportAttributeAccessIssue] # instance override + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + monkeypatch.setattr(ps, "prisma_client", None) + reservation = _reservation(reserved_cost=0.4) + + await ps.increment_spend_counters( + token="hashed", + team_id="team", + user_id="user", + org_id="org", + end_user_id="eu", + response_cost=0.5, + budget_reservation=reservation, + ) + + assert {call.kwargs["key"] for call in redis.async_delete_cache.await_args_list} == RESERVED_KEYS + assert all("applied_adjustment" not in entry for entry in reservation["entries"]) + assert redis.commands[-1].split()[0] == "PIPELINE" + assert set(redis.commands[-1].split()[1:]) == RESERVED_KEYS | {"spend:user:user"} + + +def test_a_scope_opened_inside_an_open_scope_joins_its_batch_and_a_closed_one_gets_its_own(): + redis = CountingRedis() + with spend_counter_batch_scope(redis, counter_keys=frozenset({"spend:key:a"})): + outer = active_spend_counter_batch() + assert outer is not None + with spend_counter_batch_scope(redis, counter_keys=frozenset({"spend:key:b"})): + assert active_spend_counter_batch() is outer + assert outer.counter_keys == {"spend:key:a", "spend:key:b"} + release_spend_counter_batch() + with spend_counter_batch_scope(redis, counter_keys=frozenset({"spend:key:c"})): + inner = active_spend_counter_batch() + assert inner is not outer + assert inner is not None and inner.counter_keys == {"spend:key:c"} + assert active_spend_counter_batch() is outer diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 3b3031647dc..03a24ec5e98 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -8419,7 +8419,7 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments(): async def assert_reservation_not_finalized_yet(**kwargs): assert budget_reservation["finalized"] is False incremented_counters.append(kwargs["counter_key"]) - return ps._PendingSpendIncrement( + return ps.PendingSpendIncrement( counter_key=kwargs["counter_key"], increment=kwargs["increment"] )