mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Merge pull request #41588 from BerriAI/litellm_ui_session_personal_key_budget_ceiling
* fix(key_generate): use user's budget for UI session personal keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Update litellm/proxy/management_endpoints/key_management_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(user_update): evict cached user on max_budget change so the personal key ceiling refreshes on every worker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: jesus-berri <jesus@berri.ai> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: yucheng <yucheng@berri.ai>
This commit is contained in:
yucheng-berri
GitHub
commit
51ae7ab1fd
4 changed files with 129 additions and 8 deletions
|
|
@ -107,6 +107,7 @@ if TYPE_CHECKING:
|
|||
router: Final = APIRouter()
|
||||
_USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig])
|
||||
_USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE: Final = 50
|
||||
_USER_BUDGET_CACHE_FIELDS: Final = frozenset({"max_budget", "model_max_budget"})
|
||||
|
||||
|
||||
def _user_table(
|
||||
|
|
@ -1571,7 +1572,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 or "metadata" in data_json:
|
||||
if not _USER_BUDGET_CACHE_FIELDS.isdisjoint(non_default_values) or "metadata" in data_json:
|
||||
await evict_and_broadcast(
|
||||
cache_keys=(non_default_values["user_id"],),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
|
|
@ -1902,7 +1903,7 @@ async def bulk_user_update(
|
|||
),
|
||||
)
|
||||
|
||||
if "model_max_budget" in non_default_values:
|
||||
if not _USER_BUDGET_CACHE_FIELDS.isdisjoint(non_default_values):
|
||||
for start in range(0, len(all_users_in_db), _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE):
|
||||
await asyncio.gather(
|
||||
*(
|
||||
|
|
|
|||
|
|
@ -1257,11 +1257,9 @@ async def _common_key_generation_helper(
|
|||
|
||||
# Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller
|
||||
# cannot grant a key a higher budget than their own authority.
|
||||
is_ui_session_team_key = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID and _requested_team_id is not None
|
||||
# Session tokens (lite login) carry max_budget=None to avoid a per-session
|
||||
# LLM spend cap, but that None must not be read as "unlimited delegation
|
||||
# authority". A personal key (no team) has no team-budget enforcement at
|
||||
# request time, so a session token cannot delegate any budget for one.
|
||||
# UI session personal keys are capped by user_max_budget when it is available.
|
||||
is_ui_session_token: Final = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID
|
||||
is_ui_session_team_key = is_ui_session_token and _requested_team_id is not None
|
||||
if (
|
||||
user_api_key_dict.is_session_token
|
||||
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
|
||||
|
|
@ -1279,7 +1277,9 @@ async def _common_key_generation_helper(
|
|||
},
|
||||
)
|
||||
delegation_ceiling: Final = (
|
||||
user_api_key_dict.max_budget
|
||||
user_api_key_dict.user_max_budget
|
||||
if is_ui_session_token and user_api_key_dict.user_max_budget is not None
|
||||
else user_api_key_dict.max_budget
|
||||
if user_api_key_dict.max_budget is not None
|
||||
else (team_table.max_budget if user_api_key_dict.is_session_token and team_table is not None else None)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2317,6 +2317,49 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke
|
|||
broadcast.assert_awaited_once_with(cache_key=saved_user.user_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("all_users", [False, True], ids=["single-user", "bulk-all-users"])
|
||||
async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: MockerFixture, all_users: bool) -> 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, bulk_user_update
|
||||
from litellm.types.proxy.management_endpoints.internal_user_endpoints import BulkUpdateUserRequest
|
||||
|
||||
saved_user: Final = LiteLLM_UserTable(user_id="user-spruce", max_budget=500.0)
|
||||
prisma_client: Final = mocker.MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user)
|
||||
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)
|
||||
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,
|
||||
)
|
||||
admin: Final = UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
|
||||
if all_users:
|
||||
await bulk_user_update(
|
||||
data=BulkUpdateUserRequest(all_users=True, user_updates={"max_budget": 50.0}),
|
||||
user_api_key_dict=admin,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
prisma_client.db.litellm_usertable.update_many.assert_awaited_once_with(where={}, data={"max_budget": 50.0})
|
||||
else:
|
||||
await _update_single_user_helper(
|
||||
user_request=UpdateUserRequest(user_id=saved_user.user_id, max_budget=50.0),
|
||||
user_api_key_dict=admin,
|
||||
)
|
||||
assert prisma_client.update_data.call_args.kwargs["data"]["max_budget"] == 50.0
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -15831,6 +15831,83 @@ async def test_ghsa_q775_ui_session_token_personal_key_still_capped():
|
|||
assert "cannot exceed" in msg.lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_session_token_personal_key_ceiling_is_user_budget():
|
||||
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
|
||||
|
||||
data = GenerateKeyRequest(max_budget=100)
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-ui-session",
|
||||
user_id="user-1",
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID,
|
||||
max_budget=1.0,
|
||||
user_max_budget=500.0,
|
||||
)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()), # test-quality-ok: helper reads proxy_server.prisma_client directly
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: helper reads proxy_server.user_api_key_cache directly
|
||||
patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: helper reads proxy_server.llm_router directly
|
||||
patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: helper reads proxy_server.premium_user directly
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"), # test-quality-ok: helper reads proxy_server.litellm_proxy_admin_name directly
|
||||
patch( # test-quality-ok: helper has no dependency injection seam for key persistence
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn"
|
||||
) as mock_generate_key,
|
||||
):
|
||||
mock_generate_key.return_value = {"key": "sk-test-key", "token_id": "token-id"}
|
||||
try:
|
||||
await _common_key_generation_helper(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
team_table=None,
|
||||
)
|
||||
except (HTTPException, ProxyException) as err:
|
||||
msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", ""))
|
||||
assert "cannot exceed" not in msg.lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_session_token_personal_key_above_user_budget_rejected():
|
||||
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
|
||||
|
||||
data = GenerateKeyRequest(max_budget=600)
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-ui-session",
|
||||
user_id="user-1",
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID,
|
||||
max_budget=1.0,
|
||||
user_max_budget=500.0,
|
||||
)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()), # test-quality-ok: helper reads proxy_server.prisma_client directly
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: helper reads proxy_server.user_api_key_cache directly
|
||||
patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: helper reads proxy_server.llm_router directly
|
||||
patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: helper reads proxy_server.premium_user directly
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"), # test-quality-ok: helper reads proxy_server.litellm_proxy_admin_name directly
|
||||
patch( # test-quality-ok: helper has no dependency injection seam for key persistence
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn"
|
||||
) as mock_generate_key,
|
||||
):
|
||||
mock_generate_key.return_value = {"key": "sk-test-key", "token_id": "token-id"}
|
||||
with pytest.raises((HTTPException, ProxyException)) as exc_info:
|
||||
await _common_key_generation_helper(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
team_table=None,
|
||||
)
|
||||
err = exc_info.value
|
||||
code = getattr(err, "status_code", None) or getattr(err, "code", None)
|
||||
msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", ""))
|
||||
assert str(code) == "400"
|
||||
assert "cannot exceed" in msg.lower()
|
||||
assert "500.0" in msg
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue