scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-27 01:22:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 82

fix(key_generate): use user's budget for UI session personal keys

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-09-17 09:47:14 +00:00
parent 4b368bf066
commit 48bde68781
2 changed files with 84 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1233,11 +1233,10 @@ 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.
# Session tokens (lite login) use their session max_budget for team keys, but
# 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
@ -1255,7 +1254,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)
)

77
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -15509,6 +15509,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():
"""