diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 4cc849762aa..c30f785d546 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1073,7 +1073,9 @@ async def _verify_generic_sso_or_fall_back_to_id_token( ) -async def _raise_if_sso_exceeds_free_user_limit(premium_user: bool, prisma_client: PrismaClient | None) -> None: +async def _raise_if_sso_exceeds_free_user_limit( + premium_user: bool, prisma_client: PrismaClient | None, *, adding_user: bool = False +) -> None: """Free tier allows SSO for up to 5 billable users; beyond that requires an Enterprise license.""" if premium_user is True: return @@ -1085,7 +1087,7 @@ async def _raise_if_sso_exceeds_free_user_limit(premium_user: bool, prisma_clien code=status.HTTP_403_FORBIDDEN, ) billable_users: Final = await UserRepository(prisma_client).count_billable_users() - if billable_users and billable_users > 5: + if billable_users + int(adding_user) > 5: raise ProxyException( message="You must be a LiteLLM Enterprise user to use SSO for more than 5 users. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You configured SSO (one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, `GENERIC_CLIENT_ID`, or SAML) in your env. Please unset it", type=ProxyErrorTypes.auth_error, @@ -1942,6 +1944,11 @@ async def get_user_info_from_db( if user_info is not None: break + if user_info is None: + from litellm.proxy.proxy_server import premium_user + + await _raise_if_sso_exceeds_free_user_limit(premium_user, prisma_client, adding_user=True) + verbose_proxy_logger.debug( "user_info: %s; litellm.default_internal_user_params: %s", user_info, litellm.default_internal_user_params ) @@ -1962,6 +1969,8 @@ async def get_user_info_from_db( ) return user_info + except ProxyException: + raise except Exception as e: if resolved_by_database_only and PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e): verbose_proxy_logger.exception("SSO login failed: the database could not resolve the account: %s", e) diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 873182a5145..a5c1b4fc86f 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -17,7 +17,7 @@ from fastapi import HTTPException, Request import litellm from litellm._uuid import uuid from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse, ProxyException +from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse, ProxyException, SSOUserDefinedValues from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO from litellm.proxy.management_endpoints.types import CustomOpenID @@ -1878,7 +1878,7 @@ async def test_get_user_info_from_db_user_not_exists_creates_user(): from litellm.proxy._types import NewUserResponse, SSOUserDefinedValues from litellm.proxy.management_endpoints.ui_sso import get_user_info_from_db - prisma_client = MagicMock() + prisma_client = MagicMock(**{"db.litellm_usertable.count": AsyncMock(return_value=0)}) user_api_key_cache = MagicMock() proxy_logging_obj = MagicMock() user_email = "newuser@example.com" @@ -2115,6 +2115,7 @@ async def test_get_user_info_from_db_fails_with_500_when_the_database_drops_whil **{ "db.litellm_usertable.find_unique": AsyncMock(return_value=None), "db.litellm_usertable.find_first": AsyncMock(return_value=None), + "db.litellm_usertable.count": AsyncMock(return_value=0), } ) @@ -5991,7 +5992,6 @@ class TestPKCEFunctionality: @pytest.mark.asyncio async def test_pkce_token_exchange_http200_with_error_body(self): """Provider returns HTTP 200 but with an error field instead of tokens.""" - from litellm.proxy._types import ProxyException error_body = { "error": "invalid_grant", @@ -6088,7 +6088,6 @@ class TestPKCEFunctionality: @pytest.mark.asyncio async def test_pkce_userinfo_raises_when_both_sources_unavailable(self): """When userinfo endpoint fails AND no id_token, raise ProxyException.""" - from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler with patch( @@ -6115,7 +6114,6 @@ class TestPKCEFunctionality: async def test_pkce_userinfo_http200_empty_body_no_id_token_raises(self): """When userinfo returns HTTP 200 with an empty/null body and no id_token is available, _get_pkce_userinfo raises ProxyException.""" - from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler mock_resp = MagicMock() @@ -6153,7 +6151,6 @@ class TestPKCEFunctionality: from starlette.requests import Request - from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler mock_cache = MagicMock(redis_cache=None) @@ -6278,7 +6275,6 @@ class TestPKCEFunctionality: from starlette.requests import Request - from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler # Cache returns an integer — unexpected format @@ -6366,7 +6362,6 @@ class TestPKCEFunctionality: returns a non-200 status (e.g. 401 Unauthorized from provider).""" from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler mock_response = MagicMock() @@ -6486,7 +6481,6 @@ class TestPKCEFunctionality: async def test_pkce_token_exchange_null_json_body_raises_proxy_exception(self): """HTTP 200 with JSON body `null` raises a clean ProxyException instead of AttributeError when .get() is called on the None return value.""" - from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler with patch( @@ -6520,7 +6514,6 @@ class TestPKCEFunctionality: async def test_pkce_token_exchange_http200_no_error_field_no_access_token(self): """HTTP 200 with no error field and no access_token raises ProxyException with a descriptive message showing the actual response keys.""" - from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler body_without_token = {"token_type": "Bearer", "scope": "openid"} @@ -7906,7 +7899,6 @@ class TestCliSsoAttributionMetadata: @pytest.mark.asyncio async def test_cli_sso_callback_rejects_restricted_sso_group(self): """CLI SSO must enforce restricted_sso_group before upserting the user.""" - from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints import ui_sso from litellm.proxy.management_endpoints.types import CustomOpenID @@ -8609,7 +8601,6 @@ class TestPKCEStateCookieBinding: """When PKCE is enabled and a code_verifier is in the cache, the callback must reject a request that has no ``litellm_oauth_state`` cookie (browser-to-server binding missing).""" - from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import ( SSOAuthenticationHandler, get_generic_sso_response, @@ -8662,7 +8653,6 @@ class TestPKCEStateCookieBinding: async def test_pkce_callback_rejects_state_cookie_mismatch(self): """The Login-CSRF shape: attacker mints state ``A``, victim's browser carries cookie state ``B``. The callback must reject.""" - from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import ( SSOAuthenticationHandler, get_generic_sso_response, @@ -9202,7 +9192,6 @@ async def test_saml_callback_enforces_free_sso_user_limit_after_validation(): """An IdP-initiated assertion must not bypass the >5 free-SSO-user Enterprise gate that /sso/key/generate enforces; the ACS re-checks it after validating the assertion, so the entitlement DB query never runs on unvalidated input.""" - from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.types import CustomOpenID from litellm.proxy.management_endpoints.ui_sso import saml_callback @@ -9249,6 +9238,110 @@ async def test_saml_callback_enforces_free_sso_user_limit_after_validation(): assert call_order == ["validate", "count"] +def _prisma_with_users(total_users: int) -> MagicMock: + prisma_client: Final = MagicMock() + user_table: Final = prisma_client.db.litellm_usertable + user_table.count = AsyncMock(side_effect=lambda where=None: total_users if where is None else 0) + user_table.find_unique = AsyncMock(return_value=None) + user_table.find_first = AsyncMock(return_value=None) + prisma_client.insert_data = AsyncMock(return_value=MagicMock()) + return prisma_client + + +def _inserted_tables(prisma_client: MagicMock) -> list[str]: + return [call.kwargs["table_name"] for call in prisma_client.insert_data.await_args_list] + + +def _sso_login(user_id: str) -> tuple[CustomOpenID, SSOUserDefinedValues]: + email: Final = f"{user_id}@litellm.ai" + openid: Final = CustomOpenID( + id=user_id, + email=email, + first_name=None, + last_name=None, + display_name=user_id, + picture=None, + provider="microsoft", + team_ids=[], + ) + user_defined_values: Final = SSOUserDefinedValues( + models=[], + user_id=user_id, + user_email=email, + user_role=None, + max_budget=None, + budget_duration=None, + ) + return openid, user_defined_values + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("premium_user", "existing_users", "refused"), + [(False, 5, True), (False, 4, False), (True, 5, False)], +) +async def test_free_tier_sso_refuses_the_sixth_user_instead_of_locking_everyone_out( + premium_user: bool, existing_users: int, refused: bool +) -> None: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.ui_sso import get_user_info_from_db + + openid, user_defined_values = _sso_login("new-hire") + prisma_client: Final = _prisma_with_users(existing_users) + with ( + patch("litellm.proxy.proxy_server.premium_user", premium_user), + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + ): + try: + await get_user_info_from_db( + result=openid, + prisma_client=prisma_client, + user_api_key_cache=UserApiKeyCache(), + proxy_logging_obj=MagicMock(), + user_email=openid.email, + user_defined_values=user_defined_values, + alternate_user_id="new-hire", + ) + error_code = None + except ProxyException as e: + error_code = str(e.code) + + expected: Final = ("403", []) if refused else (None, ["user"]) + assert (error_code, _inserted_tables(prisma_client)) == expected, ( + f"premium={premium_user}, {existing_users} existing users: expected (error code, inserted tables) {expected}" + ) + + +@pytest.mark.asyncio +async def test_free_tier_sso_at_the_user_limit_still_lets_existing_users_log_in() -> None: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.ui_sso import get_user_info_from_db + + openid, user_defined_values = _sso_login("existing-admin") + existing_row: Final = LiteLLM_UserTable( + user_id="existing-admin", user_email=openid.email, user_role="proxy_admin", max_budget=None + ) + cache: Final = UserApiKeyCache() + await cache.async_set_cache(key="existing-admin", value=existing_row) + + prisma_client: Final = _prisma_with_users(5) + with ( + patch("litellm.proxy.proxy_server.premium_user", False), + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + ): + user_info = await get_user_info_from_db( + result=openid, + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=MagicMock(), + user_email=openid.email, + user_defined_values=user_defined_values, + alternate_user_id="existing-admin", + ) + + assert (user_info, _inserted_tables(prisma_client)) == (existing_row, []) + + @pytest.mark.asyncio async def test_cli_poll_key_tolerates_missing_user_row(): """The CLI poll must still mint the JWT when the user lookup raises, @@ -9610,7 +9703,6 @@ async def test_non_json_userinfo_falls_back_to_id_token_claims(): async def test_non_json_userinfo_without_id_token_raises_401(): """Same unreadable userinfo but no decodable id_token in the token response: the login must fail closed with a 401 ProxyException naming the userinfo endpoint.""" - from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response request = _generic_sso_callback_request()