From 689e9d862f0ee8b1580ab0ef848ca9e66aaa30ca Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Thu, 25 Jun 2026 12:35:04 +0530
Subject: [PATCH] fix(ui_sso): add _enforce_free_sso_user_limit and propagate
 ProxyException
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Three fixes for tests added in this PR:

1. Extract _enforce_free_sso_user_limit() — inline free-tier user
   limit check was only in google_login; now available as a standalone
   function so insert_sso_user and login flows can share the logic.

2. insert_sso_user now accepts an optional prisma_client param and
   calls _enforce_free_sso_user_limit (block_at_limit=True) before
   inserting, so a non-premium proxy cannot onboard a 6th SSO user.

3. get_user_info_from_db re-raises ProxyException instead of swallowing
   it — a 403 from the user limit check was being caught by the bare
   except and silently returning None.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 litellm/proxy/management_endpoints/ui_sso.py | 39 ++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py
index 199de54ff09..a32ae21d8c1 100644
--- a/litellm/proxy/management_endpoints/ui_sso.py
+++ b/litellm/proxy/management_endpoints/ui_sso.py
@@ -1734,6 +1734,8 @@ async def get_user_info_from_db(
         )
 
         return user_info
+    except ProxyException:
+        raise
     except Exception as e:
         verbose_proxy_logger.exception(
             f"[Non-Blocking] Error trying to add sso user to db: {e}"
@@ -2368,9 +2370,38 @@ async def cli_poll_key(
         )
 
 
+async def _enforce_free_sso_user_limit(
+    prisma_client: PrismaClient,
+    premium_user: bool,
+    block_at_limit: bool,
+) -> None:
+    """Raise ProxyException if the free-tier SSO user limit (5) is exceeded.
+
+    block_at_limit=True  → block when count >= 5 (for inserts: inserting would exceed)
+    block_at_limit=False → block when count >  5 (for logins: existing user over limit)
+    """
+    if premium_user:
+        return
+    FREE_SSO_USER_LIMIT = 5
+    total_users = await prisma_client.db.litellm_usertable.count()
+    threshold = FREE_SSO_USER_LIMIT if block_at_limit else FREE_SSO_USER_LIMIT + 1
+    if total_users is not None and total_users >= threshold:
+        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"
+            ),
+            type=ProxyErrorTypes.auth_error,
+            param="premium_user",
+            code=status.HTTP_403_FORBIDDEN,
+        )
+
+
 async def insert_sso_user(
     result_openid: Optional[Union[OpenID, dict]],
     user_defined_values: Optional[SSOUserDefinedValues] = None,
+    prisma_client: Optional[PrismaClient] = None,
 ) -> NewUserResponse:
     """
     Helper function to create a New User in LiteLLM DB after a successful SSO login
@@ -2385,6 +2416,14 @@ async def insert_sso_user(
     verbose_proxy_logger.debug(
         f"Inserting SSO user into DB. User values: {user_defined_values}"
     )
+    if prisma_client is not None:
+        from litellm.proxy.proxy_server import premium_user as _premium_user
+
+        await _enforce_free_sso_user_limit(
+            prisma_client=prisma_client,
+            premium_user=_premium_user,
+            block_at_limit=True,
+        )
     if result_openid is None:
         raise ValueError("result_openid is None")
     if isinstance(result_openid, dict):