scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

fix: remove CLI poll session budget cap

Browse source
This commit is contained in:
Cursor Agent 2026-06-25 17:58:40 +00:00
parent b3f0085d1f
commit 60b81fb8b6
No known key found for this signature in database
2 changed files with 8 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -2245,10 +2245,8 @@ async def cli_poll_key(
"""
from litellm.proxy.auth.auth_checks import (
ExperimentalUIJWTToken,
get_team_object,
get_user_object,
)
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
from litellm.proxy.proxy_server import user_api_key_cache
try:
flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=user_api_key_cache)
@ -2330,40 +2328,10 @@ async def cli_poll_key(
models=session_data.get("models", []),
)
user_db_obj = await get_user_object(
user_id=user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
)
user_budget = user_db_obj.max_budget if user_db_obj is not None else None
team_budget: Optional[float] = None
team_budget_resolved = False
if team_id is not None:
try:
team_obj = await get_team_object(
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
team_budget = team_obj.max_budget
team_budget_resolved = True
except Exception:
pass
session_max_budget = (
litellm.max_ui_session_budget
if user_budget is None
and (team_id is None or (team_budget_resolved and team_budget is None))
else None
)
jwt_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(
user_info=user_info,
team_id=team_id,
team_alias=team_alias,
max_budget=session_max_budget,
)
# Delete cache entry (single-use)

14
tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
View file

@ -2904,7 +2904,7 @@ class TestCLIKeyRegenerationFlow:
jwt_call_args = mock_get_jwt.call_args
assert jwt_call_args.kwargs["team_id"] == selected_team
assert jwt_call_args.kwargs["team_alias"] == "Team B"
assert jwt_call_args.kwargs["max_budget"] is None
assert "max_budget" not in jwt_call_args.kwargs
# Verify session was deleted after JWT generation
mock_cache.delete_cache.assert_called_once()
@ -2968,11 +2968,13 @@ class TestCLIKeyRegenerationFlow:
assert result["status"] == "ready"
mock_get_jwt.assert_called_once()
assert mock_get_jwt.call_args.kwargs["max_budget"] is None
assert "max_budget" not in mock_get_jwt.call_args.kwargs
@pytest.mark.asyncio
async def test_cli_poll_key_caps_session_when_user_and_team_have_no_budget(self):
"""With no user and no team budget, the session falls back to max_ui_session_budget."""
async def test_cli_poll_key_does_not_cap_session_when_user_and_team_have_no_budget(
self,
):
"""With no user and no team budget, the session has no per-key budget cap."""
from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable
from litellm.proxy.management_endpoints.ui_sso import (
_hash_cli_sso_secret,
@ -3028,9 +3030,7 @@ class TestCLIKeyRegenerationFlow:
assert result["status"] == "ready"
mock_get_jwt.assert_called_once()
assert (
mock_get_jwt.call_args.kwargs["max_budget"] == litellm.max_ui_session_budget
)
assert "max_budget" not in mock_get_jwt.call_args.kwargs
class TestGetAppRolesFromIdToken: