diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index e79f11978a4..4a3b1050e4b 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -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) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 976048b9521..0a4810a638a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -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: