From 432e24a11fcceaa6884ac3f501440dee147e6a22 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 18 Sep 2026 01:01:32 +0000 Subject: [PATCH] feat(proxy): support multiple budget windows on internal users Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration.sql | 1 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/models/user.py | 6 + litellm/proxy/_types.py | 1 + litellm/proxy/auth/auth_checks.py | 51 ++++++ litellm/proxy/auth/handle_jwt.py | 1 + litellm/proxy/auth/user_api_key_auth.py | 6 + .../proxy/common_utils/reset_budget_job.py | 17 ++ .../proxy/db/budget_window_spend_writer.py | 15 ++ litellm/proxy/db/spend_counter_reseed.py | 8 +- .../management_endpoints/common_utils.py | 37 +++- .../internal_user_endpoints.py | 30 +++- .../key_management_endpoints.py | 2 + litellm/proxy/proxy_server.py | 65 ++++++- litellm/proxy/schema.prisma | 1 + schema.prisma | 1 + .../proxy/auth/test_handle_jwt.py | 62 +++++++ .../proxy/auth/test_multi_budget_windows.py | 132 +++++++++++++- .../common_utils/test_reset_budget_job.py | 48 ++++- .../db/test_budget_window_spend_writer.py | 10 +- .../proxy/db/test_spend_counter_reseed.py | 37 +++- .../test_internal_user_endpoints.py | 164 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 19 ++ 23 files changed, 684 insertions(+), 31 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260918004206_add_user_budget_limits/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260918004206_add_user_budget_limits/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260918004206_add_user_budget_limits/migration.sql new file mode 100644 index 00000000000..15117fd66eb --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260918004206_add_user_budget_limits/migration.sql @@ -0,0 +1 @@ +ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "budget_limits" JSONB; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 1894518e51d..3a57c3a1d4e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -260,6 +260,7 @@ model LiteLLM_UserTable { policies String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") + budget_limits Json? // multiple concurrent budget windows for the user created_at DateTime? @default(now()) @map("created_at") updated_at DateTime? @default(now()) @updatedAt @map("updated_at") diff --git a/litellm/models/user.py b/litellm/models/user.py index 82f78c28078..dc790865e0e 100644 --- a/litellm/models/user.py +++ b/litellm/models/user.py @@ -5,6 +5,7 @@ Canonical definition for ``litellm_usertable``. Re-exported from ``litellm.proxy._types`` for backwards compatibility. """ +import json from datetime import datetime from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -13,6 +14,7 @@ from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.models.organization_membership import ( LiteLLM_OrganizationMembershipTable, ) +from litellm.models.team import BudgetLimitEntry from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -40,6 +42,7 @@ class LiteLLM_UserTable(LiteLLMPydanticObjectBase): policies: list[str] = [] model_spend: dict | None = {} model_max_budget: dict | None = {} + budget_limits: list[BudgetLimitEntry] | None = None created_at: datetime | None = None updated_at: datetime | None = None organization_memberships: list[LiteLLM_OrganizationMembershipTable] | None = None @@ -56,6 +59,9 @@ class LiteLLM_UserTable(LiteLLMPydanticObjectBase): values.update({"models": []}) if values.get("teams") is None: values.update({"teams": []}) + raw_budget_limits = values.get("budget_limits") + if isinstance(raw_budget_limits, str): + values["budget_limits"] = json.loads(raw_budget_limits) return values def is_over_budget(self) -> bool: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b0d31df92ce..2c883e833e2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3133,6 +3133,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob # and validating it here would make one malformed row fail auth outright. # resolve_model_budget validates the single entry a request actually needs. user_model_max_budget: Mapping[str, object] | None = None + user_budget_limits: list[BudgetLimitEntry] | None = None request_route: str | None = None is_session_token: bool = False # Server-only marker set exclusively by the MCP gateway admission path diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3dd2e2d8eb2..36c284ae3e2 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1109,6 +1109,11 @@ async def common_checks( valid_token=valid_token, ), _team_multi_budget_check(team_object=team_object), + _user_multi_budget_check( + valid_token=valid_token, + team_object=team_object, + general_settings=general_settings, + ), _virtual_key_multi_budget_check(valid_token=valid_token) if valid_token is not None else None, _team_soft_budget_check( team_object=team_object, @@ -5520,6 +5525,52 @@ async def _team_multi_budget_check( ) +async def _user_multi_budget_check( + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + general_settings: dict, +): + """ + Raises BudgetExceededError if any budget window in valid_token.user_budget_limits is exceeded. + + Each window has its own Redis counter keyed by spend:user:{user_id}:window:{budget_duration}. + Using budget_duration (not list index) keeps counters stable when windows are reordered + or removed during a user update. Skipped for keys owned by a team unless + apply_user_budget_to_team_keys is enabled, matching the flat user budget check. + """ + if valid_token is None or not valid_token.user_budget_limits or valid_token.user_id is None: + return + is_team_key: Final = team_object is not None and team_object.team_id is not None + if is_team_key and general_settings.get("apply_user_budget_to_team_keys") is not True: + return + + from litellm.proxy.proxy_server import get_current_spend + + for window in valid_token.user_budget_limits: + w: dict = window if isinstance(window, dict) else window.model_dump() + counter_key = f"spend:user:{valid_token.user_id}:window:{w['budget_duration']}" + window_spend = await get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + max_budget=w["max_budget"], + window_entity_type="User", + window_entity_id=valid_token.user_id, + window_duration=str(w["budget_duration"]), + window_start=get_budget_window_start(w), + ) + if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: + raise litellm.BudgetExceededError( + current_cost=window_spend, + max_budget=w["max_budget"], + message=( + f"ExceededBudget: User={valid_token.user_id} over {w['budget_duration']} budget. " + f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" + ), + entity_type=Litellm_EntityType.USER.value, + entity_id=valid_token.user_id, + ) + + async def _team_soft_budget_check( team_object: LiteLLM_TeamTable | None, valid_token: UserAPIKeyAuth | None, diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 6a28cd7ff99..36eda157760 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -2734,6 +2734,7 @@ class JWTAuthManager: user_tpm_limit=user.tpm_limit if user is not None and not admin else None, user_rpm_limit=user.rpm_limit if user is not None and not admin else None, user_model_max_budget=user.model_max_budget if user is not None and not admin else None, + user_budget_limits=user.budget_limits if user is not None and not admin else None, **team_grants( team_object=result["team_object"], team_membership=result.get("team_membership"), diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4cbd4213463..c19c255f192 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1420,6 +1420,7 @@ async def _refresh_session_token_grants( **team_grants(team_object, team_membership, user_object.user_id), "user_role": _get_user_role(user_object), "models": () if team_object is not None else user_models(user_object), + "user_budget_limits": user_object.budget_limits, } ) ) @@ -1757,6 +1758,9 @@ async def _user_api_key_auth_builder( auto_registered.user_model_max_budget = ( user_object.model_max_budget if user_object is not None else None ) + auto_registered.user_budget_limits = ( + user_object.budget_limits if user_object is not None else None + ) valid_token = auto_registered api_key = valid_token.token or "" @@ -2192,6 +2196,7 @@ async def _user_api_key_auth_builder( # user's own per-model budget reaches enforcement and the post-call # increment through the row fetched here. valid_token.user_model_max_budget = user_obj.model_max_budget + valid_token.user_budget_limits = user_obj.budget_limits if ( user_obj is not None @@ -3294,6 +3299,7 @@ async def _return_user_api_key_auth_obj( user_spend=getattr(user_obj, "spend", None), user_max_budget=getattr(user_obj, "max_budget", None), user_model_max_budget=getattr(user_obj, "model_max_budget", None), + user_budget_limits=getattr(user_obj, "budget_limits", None), ) if user_obj is not None and _is_user_proxy_admin(user_obj=user_obj): user_api_key_kwargs.update( diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 1299a4df243..ef4c87c2ed4 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -60,6 +60,7 @@ from litellm.repositories.unit_of_work import ( budget_cascade_unit_of_work, spend_reset_unit_of_work, ) +from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) @@ -387,6 +388,13 @@ async def _write_team_windows(prisma_client: PrismaClient, row_id: str, payload: ) +async def _write_user_windows(prisma_client: PrismaClient, row_id: str, payload: str) -> None: + await UserRepository(prisma_client).table.update( + where={"user_id": row_id}, + data={"budget_limits": payload}, + ) + + @dataclass(frozen=True, slots=True) class _WindowSource: """A table whose rows carry their own per-window budget limits.""" @@ -433,6 +441,15 @@ _WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = ( retry_subject="team", write=_write_team_windows, ), + _WindowSource( + table="LiteLLM_UserTable", + id_column="user_id", + entity_type=Litellm_EntityType.USER, + counter_prefix="spend:user", + log_subject="users", + retry_subject="user", + write=_write_user_windows, + ), ) diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py index 8cf2f737063..57dd3f0c686 100644 --- a/litellm/proxy/db/budget_window_spend_writer.py +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -87,6 +87,19 @@ _SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL: Final = ( "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" ) +_SEED_FROM_SPEND_LOGS_USER_SQL: Final = ( + "SELECT COALESCE(SUM(spend), 0.0) AS total, " + "COALESCE(SUM(spend) FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) AS before_batch " + 'FROM "LiteLLM_SpendLogs" ' + 'WHERE "user" = $1 AND "startTime" >= ($2::timestamptz AT TIME ZONE \'UTC\')' +) + +_SEED_FROM_SPEND_LOGS_USER_UNBOUNDED_SQL: Final = ( + "SELECT COALESCE(SUM(spend), 0.0) AS total, COALESCE(SUM(spend), 0.0) AS before_batch " + 'FROM "LiteLLM_SpendLogs" ' + 'WHERE "user" = $1 AND "startTime" >= ($2::timestamptz AT TIME ZONE \'UTC\')' +) + _UPSERT_TRANSACTION_TIMEOUT: Final = timedelta(seconds=60) @@ -144,6 +157,8 @@ async def spend_logs_seed_totals( bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_KEY_SQL, _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL elif entity_type == Litellm_EntityType.TEAM.value: bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_TEAM_SQL, _SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL + elif entity_type == Litellm_EntityType.USER.value: + bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_USER_SQL, _SEED_FROM_SPEND_LOGS_USER_UNBOUNDED_SQL else: return None rows: Final = ( diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 89a07234c6c..0176d96fab4 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -49,6 +49,7 @@ _WINDOW_SPEND_ENTITY_TYPES: Final[Mapping[str, str]] = MappingProxyType( { "Key": Litellm_EntityType.KEY.value, "Team": Litellm_EntityType.TEAM.value, + "User": Litellm_EntityType.USER.value, } ) @@ -58,6 +59,7 @@ _WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType( { "Key": "api_key", "Team": "team_id", + "User": "user", } ) @@ -129,7 +131,7 @@ class SpendCounterReseed: # Per-window key/team counters share prefixes with primary counters # but don't correspond to a DB row. Do not reject arbitrary entity IDs # or tag names that merely contain ":window:". - if SpendCounterReseed._is_key_or_team_window_counter(counter_key): + if SpendCounterReseed._is_entity_window_counter(counter_key): return None try: async with db_lookup_gate.current(): @@ -181,8 +183,8 @@ class SpendCounterReseed: return float(row.spend or 0.0) @staticmethod - def _is_key_or_team_window_counter(counter_key: str) -> bool: - for prefix in ("spend:key:", "spend:team:"): + def _is_entity_window_counter(counter_key: str) -> bool: + for prefix in ("spend:key:", "spend:team:", "spend:user:"): if not counter_key.startswith(prefix): continue _, separator, duration = counter_key.rpartition(":window:") diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 29f24d2465f..b03a48a596e 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,5 +1,5 @@ import math -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Optional, Union @@ -41,6 +41,41 @@ def validate_budget_duration(budget_duration: str | None, status_code: int = 400 raise HTTPException(status_code=status_code, detail={"error": error}) +def validate_budget_limits(budget_limits: Sequence[object] | None, status_code: int = 400) -> None: + """Reject malformed budget windows before they are persisted: each entry + needs a valid duration, a positive finite cap, and a unique budget_duration. + Duplicate durations collide on the (entity, window) spend row, and a + non-positive cap can never be meaningful spend headroom. + """ + from litellm.models.team import BudgetLimitEntry + from litellm.proxy.common_utils.timezone_utils import budget_duration_error + + if not budget_limits: + return + windows: Final[tuple[BudgetLimitEntry, ...]] = tuple( + entry if isinstance(entry, BudgetLimitEntry) else BudgetLimitEntry.model_validate(entry) + for entry in budget_limits + ) + for window in windows: + error: Final = budget_duration_error(window.budget_duration) + if error is not None: + raise HTTPException(status_code=status_code, detail={"error": error}) + if not math.isfinite(window.max_budget) or window.max_budget <= 0: + raise HTTPException( + status_code=status_code, + detail={ + "error": f"budget_limits entry max_budget ({window.max_budget}) must be a positive finite number." + }, + ) + durations: Final[tuple[str, ...]] = tuple(window.budget_duration for window in windows) + duplicate: Final = next((d for d in durations if durations.count(d) > 1), None) + if duplicate is not None: + raise HTTPException( + status_code=status_code, + detail={"error": f"budget_limits has a duplicate budget_duration '{duplicate}'."}, + ) + + from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.proxy._types import ( diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index ba7a3309a90..440a1a0fb57 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -48,6 +48,7 @@ from litellm.proxy.management_endpoints.common_utils import ( _user_has_admin_view, require_caller_user_id_for_non_admin, validate_budget_duration, + validate_budget_limits, validate_finite_spend, ) from litellm.proxy.management_endpoints.key_management_endpoints import ( @@ -524,6 +525,7 @@ async def new_user( detail=CommonProxyErrors.db_not_connected_error.value, ) validate_budget_duration(data.budget_duration) + validate_budget_limits(data.budget_limits) # Check for duplicate user_id or email await _check_duplicate_user_id(data.user_id, prisma_client) @@ -1248,6 +1250,19 @@ def _process_keys_for_user_info( return returned_keys +def _prepare_user_budget_limits(value: object) -> str: + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + if not value: + return json.dumps(None) + initialized_windows: Final = [] + for window in cast(Sequence[object], value): + w = window if isinstance(window, dict) else window.model_dump() # pyright: ignore[reportAttributeAccessIssue] # BudgetLimitEntry or its JSON dict + w["reset_at"] = get_budget_reset_time(budget_duration=w["budget_duration"]).isoformat() + initialized_windows.append(w) + return json.dumps(initialized_windows) + + def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | UpdateUserRequestNoUserIDorEmail) -> dict: non_default_values: Final = {} fields_set: Final = data.fields_set() if hasattr(data, "fields_set") else set() @@ -1256,6 +1271,10 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda if k in ("max_budget", "budget_duration"): if k in fields_set: non_default_values[k] = v + elif k == "budget_limits": + if k in fields_set: + validate_budget_limits(v) + non_default_values[k] = _prepare_user_budget_limits(v) elif k == "model_max_budget": if k in fields_set: try: @@ -1477,7 +1496,14 @@ async def _update_single_user_helper( # because `_update_internal_user_params` drops empty values, and `object_permission: {}` is # precisely the clear-my-own-ceiling case this must refuse. _sent_fields: Final = user_request.fields_set() if hasattr(user_request, "fields_set") else set() - _protected_fields: Final = ("max_budget", "model_max_budget", "soft_budget", "spend", "object_permission") + _protected_fields: Final = ( + "max_budget", + "model_max_budget", + "budget_limits", + "soft_budget", + "spend", + "object_permission", + ) for _field in _protected_fields: if _field in non_default_values or _field in _sent_fields: raise HTTPException( @@ -1561,7 +1587,7 @@ async def _update_single_user_helper( await _invalidate_user_spend_counter_if_changed(non_default_values) - if "model_max_budget" in non_default_values: + if "model_max_budget" in non_default_values or "budget_limits" in non_default_values: await evict_and_broadcast( cache_keys=(non_default_values["user_id"],), user_api_key_cache=user_api_key_cache, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 802a7c3e469..9a755f4396a 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4507,6 +4507,8 @@ async def generate_key_helper_fn( # Only when supplied: the SSO and default-key callers reach this with the # empty default, and writing that would clear an existing user's budgets. user_data["model_max_budget"] = model_max_budget_json + if budget_limits_json is not None: + user_data["budget_limits"] = budget_limits_json key_data: Final = { "token": token, "key_alias": key_alias, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d7d8413d2ce..ce5b942f3a6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3024,16 +3024,65 @@ async def _increment_spend_counters_batched( 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 () - return ( - await _prepare_spend_counter_increment( - counter_key=user_counter_key, - source_cache_key=scope_user_id, - increment=cost, - ), + user_pending: Final[tuple[PendingSpendIncrement, ...]] = ( + () + if user_counter_key in reserved_counter_keys + else ( + await _prepare_spend_counter_increment( + counter_key=user_counter_key, + source_cache_key=scope_user_id, + increment=cost, + ), + ) ) + async def _user_window_increment(window: object) -> PendingSpendIncrement | None: + duration = ( + window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) + ) + user_window_reset_at = ( + window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) + ) + user_window_counter: Final = f"spend:user:{scope_user_id}:window:{duration}" + user_window_start = get_budget_window_start(window) + pending_window: Final = ( + await _prepare_window_spend_counter_increment( + counter_key=user_window_counter, + entity_type="User", + entity_id=scope_user_id, + window_duration=duration, + window_start=user_window_start, + increment=cost, + ) + if user_window_counter not in reserved_counter_keys + else None + ) + await _enqueue_window_spend_row_update( + entity_type=Litellm_EntityType.USER, + entity_id=scope_user_id, + reset_at=user_window_reset_at, + window_duration=duration, + window_start=user_window_start, + increment=cost, + request_started_at=request_started_at, + ) + return pending_window + + user_obj: Final[object] = await user_api_key_cache.async_get_cache(key=scope_user_id) + if user_obj is None: + return user_pending + user_budget_limits = getattr(user_obj, "budget_limits", None) or ( + user_obj.get("budget_limits") if isinstance(user_obj, dict) else None + ) + if isinstance(user_budget_limits, str): + user_budget_limits = json.loads(user_budget_limits) + if not isinstance(user_budget_limits, list): + return user_pending + window_pending: Final = await asyncio.gather( + *(_user_window_increment(window) for window in user_budget_limits), return_exceptions=True + ) + return user_pending + tuple(item for item in window_pending if item is not None) + scope_coros: Final = tuple( coro for coro in ( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 1894518e51d..3a57c3a1d4e 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -260,6 +260,7 @@ model LiteLLM_UserTable { policies String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") + budget_limits Json? // multiple concurrent budget windows for the user created_at DateTime? @default(now()) @map("created_at") updated_at DateTime? @default(now()) @updatedAt @map("updated_at") diff --git a/schema.prisma b/schema.prisma index 1894518e51d..3a57c3a1d4e 100644 --- a/schema.prisma +++ b/schema.prisma @@ -260,6 +260,7 @@ model LiteLLM_UserTable { policies String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") + budget_limits Json? // multiple concurrent budget windows for the user created_at DateTime? @default(now()) @map("created_at") updated_at DateTime? @default(now()) @updatedAt @map("updated_at") diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 15defb196af..f9cb65a734d 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -13,6 +13,7 @@ import litellm from litellm.proxy._types import ( DEFAULT_JWKS_STALE_TTL, + JWTAuthBuilderResult, JWTLiteLLMRoleMap, LiteLLM_JWTAuth, LiteLLM_ModelTable, @@ -7144,3 +7145,64 @@ async def test_admin_jwt_team_header_only_provisions_during_admission(monkeypatc else: create_team.assert_not_awaited() assert result["team_id"] is None + + +def test_jwt_built_user_api_key_auth_carries_user_budget_limits(): + """JWT-authenticated requests have no key row, so the user's budget windows + must ride UserAPIKeyAuth.user_budget_limits from the loaded user object.""" + user = LiteLLM_UserTable( + user_id="jwt-user", + user_role=LitellmUserRoles.INTERNAL_USER.value, + budget_limits=[ + {"budget_duration": "1d", "max_budget": 10.0}, + {"budget_duration": "30d", "max_budget": 100.0}, + ], + ) + result = JWTAuthBuilderResult( + is_proxy_admin=False, + team_object=None, + user_object=user, + end_user_object=None, + org_object=None, + token="jwt", + team_id=None, + user_id="jwt-user", + user_email="jwt@example.com", + end_user_id=None, + org_id=None, + team_membership=None, + jwt_claims={}, + agent_id=None, + ) + + auth = JWTAuthManager.user_api_key_auth_from_result(result=result) + + windows = [w.model_dump() if not isinstance(w, dict) else w for w in auth.user_budget_limits or []] + assert [(w["budget_duration"], w["max_budget"]) for w in windows] == [("1d", 10.0), ("30d", 100.0)] + + +def test_jwt_admin_does_not_inherit_user_budget_limits(): + user = LiteLLM_UserTable( + user_id="admin-user", + budget_limits=[{"budget_duration": "1d", "max_budget": 10.0}], + ) + result = JWTAuthBuilderResult( + is_proxy_admin=True, + team_object=None, + user_object=user, + end_user_object=None, + org_object=None, + token="jwt", + team_id=None, + user_id="admin-user", + user_email=None, + end_user_id=None, + org_id=None, + team_membership=None, + jwt_claims={}, + agent_id=None, + ) + + auth = JWTAuthManager.user_api_key_auth_from_result(result=result) + + assert auth.user_budget_limits is None diff --git a/tests/test_litellm/proxy/auth/test_multi_budget_windows.py b/tests/test_litellm/proxy/auth/test_multi_budget_windows.py index 0f01391b2f5..7628733754a 100644 --- a/tests/test_litellm/proxy/auth/test_multi_budget_windows.py +++ b/tests/test_litellm/proxy/auth/test_multi_budget_windows.py @@ -8,7 +8,10 @@ import pytest import litellm from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.auth.auth_checks import _virtual_key_multi_budget_check +from litellm.proxy.auth.auth_checks import ( + _user_multi_budget_check, + _virtual_key_multi_budget_check, +) def _make_valid_token(**kwargs) -> UserAPIKeyAuth: @@ -68,9 +71,7 @@ async def test_over_first_window_raises(): call_count += 1 return val - with patch( - "litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend - ): + with patch("litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend): with pytest.raises(litellm.BudgetExceededError) as exc_info: await _virtual_key_multi_budget_check(valid_token=token) @@ -100,9 +101,7 @@ async def test_over_second_window_raises(): call_count += 1 return val - with patch( - "litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend - ): + with patch("litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend): with pytest.raises(litellm.BudgetExceededError) as exc_info: await _virtual_key_multi_budget_check(valid_token=token) @@ -135,3 +134,122 @@ async def test_budget_limit_entry_objects_coerced(): ): # Should not raise TypeError / KeyError — model_dump() coerces the object await _virtual_key_multi_budget_check(valid_token=token) + + +def _make_user_token(**kwargs) -> UserAPIKeyAuth: + defaults = dict( + user_id="user-1", + spend=0.0, + user_budget_limits=None, + ) + defaults.update(kwargs) + return UserAPIKeyAuth(**defaults) + + +@pytest.mark.asyncio +async def test_user_with_no_windows_passes(): + await _user_multi_budget_check(valid_token=_make_user_token(), team_object=None, general_settings={}) + + +@pytest.mark.asyncio +async def test_user_under_all_windows_passes(): + token = _make_user_token( + user_budget_limits=[ + {"budget_duration": "24h", "max_budget": 10.0, "reset_at": None}, + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": None}, + ] + ) + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new_callable=AsyncMock, + return_value=1.0, + ) as spend_mock: + await _user_multi_budget_check(valid_token=token, team_object=None, general_settings={}) + + counter_keys = [call.kwargs["counter_key"] for call in spend_mock.await_args_list] + assert counter_keys == [ + "spend:user:user-1:window:24h", + "spend:user:user-1:window:30d", + ] + assert all(call.kwargs["window_entity_type"] == "User" for call in spend_mock.await_args_list) + + +@pytest.mark.asyncio +async def test_user_over_any_window_raises(): + token = _make_user_token( + user_budget_limits=[ + {"budget_duration": "24h", "max_budget": 50.0, "reset_at": None}, + {"budget_duration": "30d", "max_budget": 5.0, "reset_at": None}, + ] + ) + + spend_by_window = [1.0, 10.0] + call_count = 0 + + async def fake_get_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + nonlocal call_count + val = spend_by_window[call_count] + call_count += 1 + return val + + with patch("litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _user_multi_budget_check(valid_token=token, team_object=None, general_settings={}) + + err = exc_info.value + assert err.status_code == 429 + assert "30d" in str(err) + assert "User=user-1" in str(err) + + +@pytest.mark.asyncio +async def test_jwt_built_token_carries_user_budget_limits_and_is_blocked(): + """JWT auth has no key; user windows must still be enforced through the + UserAPIKeyAuth.user_budget_limits field populated from the user row.""" + token = _make_user_token( + api_key=None, + user_budget_limits=[ + {"budget_duration": "1d", "max_budget": 2.0, "reset_at": None}, + ], + ) + assert token.user_budget_limits[0].max_budget == 2.0 + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new_callable=AsyncMock, + return_value=5.0, + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _user_multi_budget_check(valid_token=token, team_object=None, general_settings={}) + + assert "user-1" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_user_windows_skipped_for_team_key_unless_flag_set(): + """Matches _user_max_budget_check: keys owned by a team don't inherit the + user's windows unless apply_user_budget_to_team_keys is enabled.""" + from litellm.proxy._types import LiteLLM_TeamTable + + token = _make_user_token(user_budget_limits=[{"budget_duration": "1d", "max_budget": 2.0, "reset_at": None}]) + team = LiteLLM_TeamTable(team_id="team-1") + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new_callable=AsyncMock, + return_value=100.0, + ) as spend_mock: + await _user_multi_budget_check(valid_token=token, team_object=team, general_settings={}) + spend_mock.assert_not_awaited() + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new_callable=AsyncMock, + return_value=100.0, + ): + with pytest.raises(litellm.BudgetExceededError): + await _user_multi_budget_check( + valid_token=token, + team_object=team, + general_settings={"apply_user_budget_to_team_keys": True}, + ) diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index e96069ffa99..e0e53b055af 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -924,6 +924,7 @@ def _make_reset_budget_windows_job( monkeypatch, key_rows: List[Dict[str, Any]], team_rows: List[Dict[str, Any]], + user_rows: List[Dict[str, Any]] | None = None, ): """Build a ResetBudgetJob with a fully-mocked prisma client and a fake `litellm.proxy.proxy_server` module exposing a stub `spend_counter_cache`. @@ -933,17 +934,20 @@ def _make_reset_budget_windows_job( prisma_client = MagicMock() async def fake_query_raw(query: str, *args, **kwargs): - # Dispatch by table name in the SQL so a single stub covers both calls. + # Dispatch by table name in the SQL so a single stub covers all calls. if '"LiteLLM_VerificationToken"' in query: return key_rows if '"LiteLLM_TeamTable"' in query: return team_rows + if '"LiteLLM_UserTable"' in query: + return user_rows or [] raise AssertionError(f"Unexpected query_raw call: {query}") prisma_client.db.query_raw = AsyncMock(side_effect=fake_query_raw) prisma_client.db.execute_raw = AsyncMock(return_value=1) prisma_client.db.litellm_verificationtoken.update = AsyncMock(return_value=None) prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None) + prisma_client.db.litellm_usertable.update = AsyncMock(return_value=None) # Stub out litellm.proxy.proxy_server so the in-function # `from litellm.proxy.proxy_server import spend_counter_cache` resolves @@ -971,13 +975,15 @@ def test_reset_budget_windows_uses_is_not_null_filter(monkeypatch): asyncio.run(job.reset_budget_windows()) queries = [call.args[0] for call in prisma_client.db.query_raw.await_args_list] - assert len(queries) == 2, queries - key_query, team_query = queries + assert len(queries) == 3, queries + key_query, team_query, user_query = queries assert '"LiteLLM_VerificationToken"' in key_query assert "budget_limits IS NOT NULL" in key_query assert '"LiteLLM_TeamTable"' in team_query assert "budget_limits IS NOT NULL" in team_query + assert '"LiteLLM_UserTable"' in user_query + assert "budget_limits IS NOT NULL" in user_query def test_reset_budget_windows_resets_expired_key_window(monkeypatch): @@ -1093,6 +1099,40 @@ def test_reset_budget_windows_rolls_the_team_window_spend_row(monkeypatch): assert rolls[0][1:4] == ("team", "team-expired", "30d") +def test_reset_budget_windows_rolls_the_user_window_spend_row(monkeypatch): + """A user whose window's `reset_at` has passed gets a rolled + LiteLLM_BudgetWindowSpend row, a bumped `reset_at`, and a cleared + `spend:user:{id}:window:{duration}` counter, matching team behavior.""" + now = datetime.utcnow() + expired = (now - timedelta(minutes=5)).isoformat() + "Z" + + user_rows = [ + { + "user_id": "user-expired", + "budget_limits": [{"budget_duration": "1d", "reset_at": expired}], + } + ] + job, prisma_client, spend_counter_cache = _make_reset_budget_windows_job( + monkeypatch, key_rows=[], team_rows=[], user_rows=user_rows + ) + + asyncio.run(job.reset_budget_windows()) + + prisma_client.db.litellm_usertable.update.assert_awaited_once() + call_kwargs = prisma_client.db.litellm_usertable.update.await_args.kwargs + assert call_kwargs["where"] == {"user_id": "user-expired"} + written_windows = json.loads(call_kwargs["data"]["budget_limits"]) + assert len(written_windows) == 1 + new_reset_at = datetime.fromisoformat(written_windows[0]["reset_at"].replace("Z", "+00:00")).replace(tzinfo=None) + assert new_reset_at > now + + rolls = _window_spend_rolls(prisma_client) + assert len(rolls) == 1 + assert rolls[0][1:4] == ("user", "user-expired", "1d") + + spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:user:user-expired:window:1d", value=0.0) + + def test_reset_budget_windows_does_not_roll_an_unexpired_window(monkeypatch): now = datetime.utcnow() future = (now + timedelta(hours=1)).isoformat() + "Z" @@ -2736,7 +2776,7 @@ def _cursor_paginating_window_job(monkeypatch, key_rows: List[Dict[str, Any]]): visited: List[str] = [] async def fake_query_raw(query: str, *args, **kwargs): - if '"LiteLLM_TeamTable"' in query: + if '"LiteLLM_TeamTable"' in query or '"LiteLLM_UserTable"' in query: return [] cursor, limit = args[0], args[1] page = [row for row in ordered if row["token"] > cursor][:limit] diff --git a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py index 130f0c56ccf..cfe7493a8d0 100644 --- a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py +++ b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py @@ -344,8 +344,8 @@ async def test_all_upserts_are_committed_in_one_transaction(): @pytest.mark.asyncio async def test_unknown_entity_type_contributes_no_seed(): - """Only key and team windows have a LiteLLM_SpendLogs column to aggregate; - anything else starts from its increment alone.""" + """Only key, team, and user windows have a LiteLLM_SpendLogs column to + aggregate; anything else starts from its increment alone.""" db = _FakeDB(existing_rows=[]) async def no_such_column(prisma_client, entity_type, entity_id, window_start, batch_started_at): @@ -353,7 +353,7 @@ async def test_unknown_entity_type_contributes_no_seed(): await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=(build_window_spend_transaction("user", "u1", "30d", WINDOW_A, 1.0),), + transactions=(build_window_spend_transaction("organization", "o1", "30d", WINDOW_A, 1.0),), spend_logs_aggregate=no_such_column, ) @@ -515,7 +515,7 @@ async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet(): @pytest.mark.asyncio @pytest.mark.parametrize( "entity_type, expected_column", - [("key", "api_key = $1"), ("team", "team_id = $1")], + [("key", "api_key = $1"), ("team", "team_id = $1"), ("user", '"user" = $1')], ) async def test_seed_aggregate_sql_splits_the_window_at_the_batch_start(entity_type, expected_column): db = _FakeDB(existing_rows=[{"total": 1.25, "before_batch": 0.75}]) @@ -569,7 +569,7 @@ async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs totals = await spend_logs_seed_totals( prisma_client=_FakePrismaClient(db), - entity_type="user", + entity_type="organization", entity_id="u1", window_start=WINDOW_A, batch_started_at=None, 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 bca6344b3f7..0a2785a6302 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -156,6 +156,41 @@ async def test_window_from_table_maps_team_entity_type(): assert inner["entity_type"] == "team" +@pytest.mark.asyncio +async def test_window_from_table_maps_user_entity_type(): + prisma = _FakePrismaClient(row=_row(WINDOW_START, 7.5)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="User", + entity_id="user-1", + window_duration="1d", + expected_window_start=WINDOW_START, + ) + + assert result == 7.5 + inner = prisma.db.litellm_budgetwindowspend.where_clauses[0]["entity_type_entity_id_window_duration"] + assert inner["entity_type"] == "user" + + +@pytest.mark.asyncio +async def test_user_window_falls_back_to_spend_logs_aggregate_on_user_column(): + """With no maintained row, a user window must aggregate LiteLLM_SpendLogs + grouped by the ``user`` column, like key/team windows do on theirs.""" + prisma = _FakePrismaClient(row=None, spend_logs_total=6.25) + + result = await SpendCounterReseed.window_from_db( + prisma_client=prisma, + entity_type="User", + entity_id="user-1", + window_duration="1d", + window_start=WINDOW_START, + ) + + assert result == 6.25 + assert prisma.db.litellm_spendlogs.call_count == 1 + + @pytest.mark.asyncio async def test_window_from_table_trusts_row_newer_than_expected_window(): """Regression: a pod holding a stale ``reset_at`` computes an expected start @@ -211,7 +246,7 @@ async def test_window_from_table_treats_naive_row_timestamp_as_utc(): "prisma, entity_type", [ (_FakePrismaClient(row=None), "Key"), - (_FakePrismaClient(row=_row(WINDOW_START, 1.0)), "User"), + (_FakePrismaClient(row=_row(WINDOW_START, 1.0)), "Organization"), (_FakePrismaClient(error=RuntimeError("connection reset")), "Key"), (None, "Key"), ], diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 0d8b19345f1..df911223d50 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -4519,3 +4519,167 @@ async def test_user_update_hashes_and_persists_strong_password(_admin_prisma, mo written_data = mock_prisma_client.update_data.call_args.kwargs["data"] assert written_data.get("password") is not None assert written_data["password"] != strong_password + + +@pytest.mark.asyncio +async def test_new_user_forwards_budget_limits_into_user_persistence(mocker): + """/user/new must pass the requested windows down to generate_key_helper_fn + so they land on the user row (the helper used to drop them).""" + from litellm.proxy.management_endpoints.internal_user_endpoints import new_user + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.count = mocker.AsyncMock(return_value=5) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_id", + new=mocker.AsyncMock(), + ) + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_email", + new=mocker.AsyncMock(), + ) + mock_license = mocker.MagicMock() + mock_license.is_over_limit.return_value = False + mocker.patch("litellm.proxy.proxy_server._license_check", mock_license) + + helper = mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints.generate_key_helper_fn", + new=mocker.AsyncMock(return_value={"user_id": "u-1", "key": "sk-1", "expires": None}), + ) + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + windows = [{"budget_duration": "1d", "max_budget": 10.0}] + await new_user( + data=NewUserRequest(user_email="w@example.com", budget_limits=windows), + user_api_key_dict=admin, + ) + + helper.assert_awaited_once() + forwarded = helper.await_args.kwargs["budget_limits"] + assert [(w["budget_duration"], w["max_budget"]) for w in forwarded] == [("1d", 10.0)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "windows", + [ + [{"budget_duration": "1d", "max_budget": 10.0}, {"budget_duration": "1d", "max_budget": 5.0}], + [{"budget_duration": "1d", "max_budget": -3.0}], + [{"budget_duration": "not-a-duration", "max_budget": 10.0}], + ], + ids=["duplicate_window", "non_positive_cap", "invalid_duration"], +) +async def test_new_user_rejects_malformed_budget_limits(mocker, windows): + from litellm.proxy.management_endpoints.internal_user_endpoints import new_user + + mock_prisma_client = mocker.MagicMock() + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + duplicate_check = mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_id", + new=mocker.AsyncMock(), + ) + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with pytest.raises(ProxyException) as exc_info: + await new_user( + data=NewUserRequest(user_email="w@example.com", budget_limits=windows), + user_api_key_dict=admin, + ) + + assert str(exc_info.value.code) == "400" + duplicate_check.assert_not_awaited() + + +def test_update_internal_user_params_writes_budget_limits_with_initialized_reset_at(): + from litellm.proxy._types import UpdateUserRequest + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_internal_user_params, + ) + + data = UpdateUserRequest( + user_id="u-1", + budget_limits=[ + {"budget_duration": "1d", "max_budget": 10.0}, + {"budget_duration": "30d", "max_budget": 100.0}, + ], + ) + + non_default_values = _update_internal_user_params(data_json=data.model_dump(exclude_unset=True), data=data) + + written = json.loads(non_default_values["budget_limits"]) + assert len(written) == 2 + for window in written: + assert window["reset_at"] is not None + + +@pytest.mark.asyncio +async def test_update_user_replaces_budget_limits(_admin_prisma, mocker): + """/user/update persists the replacement list into the user row as JSON.""" + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = {"user_id": "target-user"} + existing_user.user_id = "target-user" + _admin_prisma.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + _admin_prisma.update_data = mocker.AsyncMock(return_value={"user_id": "target-user"}) + _admin_prisma.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) + + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + await _update_single_user_helper( + user_request=UpdateUserRequest( + user_id="target-user", + budget_limits=[{"budget_duration": "7d", "max_budget": 50.0}], + ), + user_api_key_dict=admin_caller, + ) + + written = json.loads(_admin_prisma.update_data.call_args.kwargs["data"]["budget_limits"]) + assert [(w["budget_duration"], w["max_budget"]) for w in written] == [("7d", 50.0)] + assert written[0]["reset_at"] is not None + + +@pytest.mark.asyncio +async def test_update_user_clears_budget_limits_with_empty_list(_admin_prisma, mocker): + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = {"user_id": "target-user"} + existing_user.user_id = "target-user" + _admin_prisma.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + _admin_prisma.update_data = mocker.AsyncMock(return_value={"user_id": "target-user"}) + _admin_prisma.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) + + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + await _update_single_user_helper( + user_request=UpdateUserRequest(user_id="target-user", budget_limits=[]), + user_api_key_dict=admin_caller, + ) + + assert json.loads(_admin_prisma.update_data.call_args.kwargs["data"]["budget_limits"]) is None + + +@pytest.mark.asyncio +async def test_update_user_rejects_duplicate_budget_window(_admin_prisma, mocker): + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + + with pytest.raises(HTTPException) as exc_info: + await _update_single_user_helper( + user_request=UpdateUserRequest( + user_id="target-user", + budget_limits=[ + {"budget_duration": "1d", "max_budget": 10.0}, + {"budget_duration": "1d", "max_budget": 5.0}, + ], + ), + user_api_key_dict=admin_caller, + ) + + assert exc_info.value.status_code == 400 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 41c4956dba6..567a479530c 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -12592,6 +12592,25 @@ async def test_team_window_spend_row_is_enqueued(): assert enqueued[0]["spend"] == pytest.approx(1.5) +@pytest.mark.asyncio +async def test_user_window_spend_row_is_enqueued(): + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=3) + user_obj = MagicMock() + user_obj.budget_limits = [{"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()}] + + with _window_spend_enqueue_env({"user-1": user_obj}) as queue: + await increment_spend_counters(token=None, team_id=None, user_id="user-1", response_cost=1.5) + enqueued = await _drain(queue) + + assert len(enqueued) == 1 + assert enqueued[0]["entity_type"] == "user" + assert enqueued[0]["entity_id"] == "user-1" + assert enqueued[0]["window_duration"] == "7d" + assert enqueued[0]["spend"] == pytest.approx(1.5) + + @pytest.mark.asyncio async def test_window_spend_row_is_enqueued_even_when_the_counter_was_reserved(): """A reservation only pre-charged the cache counter with an estimate; the