Merge 4753536654 into 4b1928387f

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)

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()