scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-05 02:41:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 48

Merge 8c79bbc1eb into c2c2a623c0

Browse source
This commit is contained in:
mubashir1osmani 2026-09-13 00:02:45 -07:00 • committed by GitHub
commit f7a1a1593d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 84 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -3370,6 +3370,7 @@ class ExperimentalUIJWTToken:
models=user_info.models,
max_parallel_requests=None,
user_role=LitellmUserRoles(user_info.user_role),
user_max_budget=user_info.max_budget,
)
return encrypt_value_helper(valid_token.model_dump_json(exclude_none=True))

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

@ -1068,7 +1068,8 @@ 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
is_ui_session_token = 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
# 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
@ -1089,11 +1090,16 @@ async def _common_key_generation_helper(
)
},
)
delegation_ceiling: Final = (
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)
)
# The dashboard key's max_budget only caps playground spend. Its owner's
# personal budget is the authority for creating personal keys.
if is_ui_session_token and user_api_key_dict.user_max_budget is not None:
delegation_ceiling = user_api_key_dict.user_max_budget
elif user_api_key_dict.max_budget is not None:
delegation_ceiling = user_api_key_dict.max_budget
else:
delegation_ceiling = (
team_table.max_budget if user_api_key_dict.is_session_token and team_table is not None else None
)
if (
user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
and not is_ui_session_team_key

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

@ -3646,7 +3646,7 @@ class SSOAuthenticationHandler:
user_id=user_defined_values["user_id"],
user_role=user_defined_values["user_role"] or user_role,
models=[],
max_budget=litellm.max_ui_session_budget,
max_budget=user_defined_values["max_budget"],
)
if _user_info is None:
raise HTTPException(

2
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -142,6 +142,7 @@ def test_get_experimental_ui_login_jwt_auth_token_valid(valid_sso_user_defined_v
assert token_data["user_role"] == LitellmUserRoles.PROXY_ADMIN.value
assert token_data["models"] == ["gpt-3.5-turbo"]
assert token_data["max_budget"] == litellm.max_ui_session_budget
assert token_data["user_max_budget"] == 100.0
# Verify expiration time is set and valid (Experimental UI uses fixed 10-min expiry)
assert "expires" in token_data
@ -270,6 +271,7 @@ def test_get_key_object_from_ui_hash_key_valid(valid_sso_user_defined_values, mo
assert key_object.user_role == LitellmUserRoles.PROXY_ADMIN
assert key_object.models == ["gpt-3.5-turbo"]
assert key_object.max_budget == litellm.max_ui_session_budget
assert key_object.user_max_budget == 100.0
def test_get_key_object_from_ui_hash_key_invalid():

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

@ -14101,6 +14101,74 @@ 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_personal_key_uses_user_budget_as_delegation_ceiling():
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
data = GenerateKeyRequest(max_budget=20)
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=30.0,
)
generated_key = {"key": "sk-generated"}
with (
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.premium_user", False),
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
new=AsyncMock(return_value=generated_key),
) as mock_generate_key,
):
result = await _common_key_generation_helper(
data=data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
team_table=None,
)
assert result.key == "sk-generated"
assert mock_generate_key.await_args.kwargs["key_max_budget"] == 20
@pytest.mark.asyncio
async def test_ui_session_personal_key_cannot_exceed_user_budget():
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
data = GenerateKeyRequest(max_budget=31)
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=30.0,
)
with (
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.premium_user", False),
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"),
pytest.raises(HTTPException) 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,
)
assert exc_info.value.status_code == 400
assert "own max_budget (30.0)" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption():
"""