scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 152

Merge pull request #40837 from BerriAI/litellm_user-model-budget-clear

Browse source 
fix(proxy): persist clearing user model budgets
This commit is contained in:
yuneng-jiang 2026-09-11 23:43:25 -07:00 • committed by GitHub
commit 9071ca503e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 169 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -22,6 +22,7 @@ from typing import Any, Final, Literal, Protocol, cast, overload
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
@ -30,6 +31,7 @@ from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import get_team_object, get_user_object
from litellm.proxy.auth.password_policy import validate_password_policy
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.common_utils.user_api_key_cache import (
object_permission_cache_key,
user_object_permission_id_cache_key,
@ -86,6 +88,7 @@ from litellm.types.proxy.management_endpoints.scim_v2 import (
SCIM_ENTITLEMENTS_METADATA_KEY,
SCIM_ROLES_METADATA_KEY,
)
from litellm.types.utils import BudgetConfig
if TYPE_CHECKING:
from prisma import models as prisma_models
@ -96,6 +99,8 @@ if TYPE_CHECKING:
from litellm.proxy.utils import ProxyLogging
router: Final = APIRouter()
_USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig])
_USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE: Final = 50
def _user_table(
@ -1252,6 +1257,13 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda
if k == "max_budget":
if "max_budget" in fields_set:
non_default_values[k] = v
elif k == "model_max_budget":
if k in fields_set:
try:
_USER_MODEL_BUDGET_ADAPTER.validate_python({} if v is None else v)
except ValidationError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
non_default_values[k] = {} if v is None else v
elif (
v is not None
and v
@ -1421,7 +1433,7 @@ async def _update_single_user_helper(
Returns the updated user data or raises an exception on failure.
"""
from litellm.proxy.proxy_server import general_settings, litellm_proxy_admin_name, prisma_client
from litellm.proxy.proxy_server import general_settings, litellm_proxy_admin_name, prisma_client, user_api_key_cache
if prisma_client is None:
raise Exception("Not connected to DB!")
@ -1464,7 +1476,7 @@ 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", "soft_budget", "spend", "object_permission")
_protected_fields: Final = ("max_budget", "model_max_budget", "soft_budget", "spend", "object_permission")
for _field in _protected_fields:
if _field in non_default_values or _field in _sent_fields:
raise HTTPException(
@ -1548,6 +1560,12 @@ async def _update_single_user_helper(
await _invalidate_user_spend_counter_if_changed(non_default_values)
if "model_max_budget" in non_default_values:
await evict_and_broadcast(
cache_keys=(non_default_values["user_id"],),
user_api_key_cache=user_api_key_cache,
)
if "object_permission_id" in non_default_values:
await _invalidate_cached_user_entitlement(
user_id=non_default_values.get("user_id"),
@ -1802,7 +1820,7 @@ async def bulk_user_update(
}'
```
"""
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client, user_api_key_cache
if prisma_client is None:
raise HTTPException(
@ -1867,9 +1885,22 @@ async def bulk_user_update(
# Perform bulk database update
await UserRepository(prisma_client).table.update_many(
where={},
data=non_default_values, # Update all users
data=(
{**non_default_values, "model_max_budget": json.dumps(non_default_values["model_max_budget"])}
if "model_max_budget" in non_default_values
else non_default_values
),
)
if "model_max_budget" in non_default_values:
for start in range(0, len(all_users_in_db), _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE):
await asyncio.gather(
*(
evict_and_broadcast(cache_keys=(user.user_id,), user_api_key_cache=user_api_key_cache)
for user in all_users_in_db[start : start + _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE]
)
)
# Create individual success results
for user in all_users_in_db:
results.append(

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

@ -1,9 +1,12 @@
import json
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Final
import pytest
from fastapi.testclient import TestClient
from fastapi import HTTPException
from pytest_mock import MockerFixture
from litellm.proxy._types import (
@ -2128,6 +2131,128 @@ def test_update_internal_user_params_keeps_original_max_budget_when_not_provided
assert "user_alias" in non_default_values
@pytest.mark.parametrize("cleared_budget", [{}, None], ids=["empty-map", "null"])
def test_update_internal_user_params_clears_model_budget(cleared_budget: dict[str, object] | None) -> None:
request: Final = UpdateUserRequest(user_id="user-spruce", model_max_budget=cleared_budget)
update: Final = _update_internal_user_params(data_json=request.model_dump(exclude_unset=True), data=request)
assert update == {"user_id": "user-spruce", "model_max_budget": {}}
def test_update_internal_user_params_preserves_model_budget_presence_and_neighbors() -> None:
omitted: Final = UpdateUserRequest(user_id="user-spruce", user_alias="Spruce")
assert _update_internal_user_params(data_json=omitted.model_dump(), data=omitted) == {
"user_id": "user-spruce",
"user_alias": "Spruce",
}
replacement: Final = {"model-spruce": {"budget_limit": 0, "time_period": "1d"}, "model-birch": 5.0, "model-cedar": 0}
request: Final = UpdateUserRequest(
user_id="user-spruce",
model_max_budget=replacement,
max_budget=50,
user_alias=None,
models=[],
allowed_cache_controls=[],
config={},
)
assert _update_internal_user_params(data_json=request.model_dump(exclude_unset=True), data=request) == {
"user_id": "user-spruce",
"model_max_budget": replacement,
"max_budget": 50,
}
@pytest.mark.parametrize("invalid_budget", [{"model-spruce": "invalid"}, {"model-spruce": {"budget_limit": "invalid"}}])
def test_update_internal_user_params_rejects_invalid_model_budget(invalid_budget: dict[str, object]) -> None:
request: Final = UpdateUserRequest(user_id="user-spruce", model_max_budget=invalid_budget)
with pytest.raises(HTTPException) as exc:
_update_internal_user_params(data_json=request.model_dump(exclude_unset=True), data=request)
assert exc.value.status_code == 400
@pytest.mark.asyncio
async def test_user_model_budget_update_by_email_refreshes_cached_user(mocker: MockerFixture) -> None:
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper
saved_user: Final = LiteLLM_UserTable(
user_id="user-spruce",
user_email="spruce@example.test",
model_max_budget={"model-spruce": {"budget_limit": 5, "time_period": "1d"}},
max_budget=50,
)
prisma_client: Final = mocker.MagicMock()
prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user)
prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user])
prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user})
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
cache: Final = UserApiKeyCache()
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new_callable=mocker.AsyncMock,
)
await _update_single_user_helper(
user_request=UpdateUserRequest(user_email=saved_user.user_email, model_max_budget={}),
user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert prisma_client.update_data.call_args.kwargs["data"]["model_max_budget"] == {}
assert "max_budget" not in prisma_client.update_data.call_args.kwargs["data"]
assert await cache.async_get_cache(key=saved_user.user_id, model_type=LiteLLM_UserTable) is None
broadcast.assert_awaited_once_with(cache_key=saved_user.user_id)
@pytest.mark.asyncio
async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocker: MockerFixture) -> None:
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.internal_user_endpoints import bulk_user_update
from litellm.types.proxy.management_endpoints.internal_user_endpoints import BulkUpdateUserRequest
saved_user: Final = LiteLLM_UserTable(user_id="user-spruce", model_max_budget={"model-spruce": {"budget_limit": 5}})
prisma_client: Final = mocker.MagicMock()
prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user])
prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1)
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
cache: Final = UserApiKeyCache()
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new_callable=mocker.AsyncMock,
)
with pytest.raises(HTTPException) as exc:
await bulk_user_update(
data=BulkUpdateUserRequest(all_users=True, user_updates={"model_max_budget": {"model-spruce": "invalid"}}),
user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN),
litellm_changed_by=None,
)
assert exc.value.status_code == 400
prisma_client.db.litellm_usertable.update_many.assert_not_called()
assert await cache.async_get_cache(key=saved_user.user_id, model_type=LiteLLM_UserTable) == saved_user
response: Final = await bulk_user_update(
data=BulkUpdateUserRequest(all_users=True, user_updates={"model_max_budget": None}),
user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN),
litellm_changed_by=None,
)
prisma_client.db.litellm_usertable.update_many.assert_awaited_once_with(where={}, data={"model_max_budget": "{}"})
prisma_client.update_data.assert_not_called()
assert response.successful_updates == 1
assert response.results[0].updated_user["model_max_budget"] == {}
assert await cache.async_get_cache(key=saved_user.user_id, model_type=LiteLLM_UserTable) is None
broadcast.assert_awaited_once_with(cache_key=saved_user.user_id)
def test_generate_request_base_validator():
"""
Test that GenerateRequestBase validator converts empty string to None for max_budget
@ -3498,7 +3623,11 @@ def test_enforce_user_info_access_blocks_cross_user_lookup():
@pytest.mark.asyncio
async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker):
@pytest.mark.parametrize(
("budget_field", "budget_value"),
[("max_budget", 999999), ("model_max_budget", {}), ("model_max_budget", None)],
)
async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker, budget_field, budget_value):
"""Non-admin updating their own record must be blocked from modifying
max_budget (self-escalation)."""
from fastapi import HTTPException
@ -3508,6 +3637,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker):
)
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "user-1", "data": {"user_id": "user-1"}})
existing_user = mocker.MagicMock()
existing_user.model_dump.return_value = {
"user_id": "user-1",
@ -3519,10 +3649,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker):
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
user_request = UpdateUserRequest(
user_id="user-1",
max_budget=999999,
)
user_request = UpdateUserRequest.model_validate({"user_id": "user-1", budget_field: budget_value})
caller = UserAPIKeyAuth(
user_id="user-1",
user_role=LitellmUserRoles.INTERNAL_USER,
@ -3533,7 +3660,8 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker):
user_request=user_request, user_api_key_dict=caller
)
assert exc.value.status_code == 403
assert "max_budget" in str(exc.value.detail)
assert budget_field in str(exc.value.detail)
mock_prisma_client.update_data.assert_not_called()
@pytest.mark.asyncio