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 75584ce4f65..976048b9521 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2909,6 +2909,129 @@ class TestCLIKeyRegenerationFlow: # Verify session was deleted after JWT generation mock_cache.delete_cache.assert_called_once() + @pytest.mark.asyncio + async def test_cli_poll_key_does_not_cap_session_when_user_has_budget(self): + """A user with a configured budget must not get the max_ui_session_budget fallback cap.""" + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + cli_poll_key, + ) + + session_data = { + "user_id": "budgeted-user", + "user_role": "internal_user", + "teams": [], + "team_details": [], + "models": ["gpt-4"], + "user_email": "budgeted@example.com", + } + mock_user_info = LiteLLM_UserTable( + user_id="budgeted-user", + user_role="internal_user", + teams=[], + models=["gpt-4"], + max_budget=100.0, + ) + mock_cache = MagicMock() + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "sso_complete": True, + "user_code_verified": True, + "session_data": session_data, + } + mock_jwt_token = "eyJhbGciOiJIUzI1NiJ9.budgeted.token" + + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.prisma_client"), + patch( + "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", + return_value=mock_jwt_token, + ) as mock_get_jwt, + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(return_value=mock_user_info), + ), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + new=AsyncMock( + side_effect=AssertionError("team lookup must be skipped") + ), + ), + ): + result = await cli_poll_key( + key_id="cli-session-budgeted", + team_id=None, + x_litellm_cli_poll_secret="poll-secret", + ) + + assert result["status"] == "ready" + mock_get_jwt.assert_called_once() + assert mock_get_jwt.call_args.kwargs["max_budget"] is None + + @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.""" + from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + cli_poll_key, + ) + + session_data = { + "user_id": "unbudgeted-user", + "user_role": "internal_user", + "teams": ["team-x"], + "team_details": [{"team_id": "team-x", "team_alias": "Team X"}], + "models": ["gpt-4"], + "user_email": "unbudgeted@example.com", + } + mock_user_info = LiteLLM_UserTable( + user_id="unbudgeted-user", + user_role="internal_user", + teams=["team-x"], + models=["gpt-4"], + max_budget=None, + ) + mock_team = LiteLLM_TeamTableCachedObj(team_id="team-x", max_budget=None) + mock_cache = MagicMock() + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "sso_complete": True, + "user_code_verified": True, + "session_data": session_data, + } + mock_jwt_token = "eyJhbGciOiJIUzI1NiJ9.unbudgeted.token" + + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.prisma_client"), + patch( + "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", + return_value=mock_jwt_token, + ) as mock_get_jwt, + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(return_value=mock_user_info), + ), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + new=AsyncMock(return_value=mock_team), + ), + ): + result = await cli_poll_key( + key_id="cli-session-unbudgeted", + team_id="team-x", + x_litellm_cli_poll_secret="poll-secret", + ) + + assert result["status"] == "ready" + mock_get_jwt.assert_called_once() + assert ( + mock_get_jwt.call_args.kwargs["max_budget"] == litellm.max_ui_session_budget + ) + class TestGetAppRolesFromIdToken: """Test the get_app_roles_from_id_token method"""