scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

feat(proxy): support multiple budget windows on internal users

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
jesus 2026-09-18 01:01:32 +00:00
parent 2fb502a556
commit 432e24a11f
23 changed files with 684 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20260918004206_add_user_budget_limits/migration.sql Normal file
View file

@ -0,0 +1 @@
ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "budget_limits" JSONB;

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -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")

6
litellm/models/user.py
View file

@ -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:

1
litellm/proxy/_types.py
View file

@ -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

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

@ -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,

1
litellm/proxy/auth/handle_jwt.py
View file

@ -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"),

6
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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(

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

@ -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,
),
)

15
litellm/proxy/db/budget_window_spend_writer.py
View file

@ -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 = (

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

@ -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:")

37
litellm/proxy/management_endpoints/common_utils.py
View file

@ -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 (

30
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -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,

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

@ -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,

65
litellm/proxy/proxy_server.py
View file

@ -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 (

1
litellm/proxy/schema.prisma
View file

@ -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")

1
schema.prisma
View file

@ -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")

62
tests/test_litellm/proxy/auth/test_handle_jwt.py
View file

@ -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

132
tests/test_litellm/proxy/auth/test_multi_budget_windows.py
View file

@ -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},
)

48
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
View file

@ -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]

10
tests/test_litellm/proxy/db/test_budget_window_spend_writer.py
View file

@ -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,

37
tests/test_litellm/proxy/db/test_spend_counter_reseed.py
View file

@ -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"),
],

164
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -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

19
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -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