fix(ui_sso): add _enforce_free_sso_user_limit and propagate ProxyException

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>
This commit is contained in:
Sameer Kankute 2026-06-25 12:35:04 +05:30
parent ca1d2d01e4
commit 689e9d862f
No known key found for this signature in database

View file

@ -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):