[Feat] SSO - Ensure role from SSO provider is used when a user is inserted onto LiteLLM (#16794)

* test_apply_user_info_values_sso_role_takes_precedence

* fix SSO
This commit is contained in:
Ishaan Jaff
2025-11-18 15:35:38 -08:00
committed by GitHub
parent 97bb899f52
commit e2fc225201
3 changed files with 112 additions and 6 deletions
@@ -13084,7 +13084,7 @@
"supports_audio_output": false,
"supports_function_calling": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_system_messages": false,
"supports_tool_choice": true,
"supports_vision": true
},
+41 -5
View File
@@ -62,7 +62,11 @@ from litellm.proxy.management_endpoints.sso_helper_utils import (
has_admin_ui_access,
)
from litellm.proxy.management_endpoints.team_endpoints import new_team, team_member_add
from litellm.proxy.management_endpoints.types import CustomOpenID, get_litellm_user_role
from litellm.proxy.management_endpoints.types import (
CustomOpenID,
get_litellm_user_role,
is_valid_litellm_user_role,
)
from litellm.proxy.utils import (
PrismaClient,
ProxyLogging,
@@ -554,6 +558,20 @@ async def get_user_info_from_db(
return None
def _should_use_role_from_sso_response(sso_role: Optional[str]) -> bool:
"""returns true if SSO upsert should use the 'role' defined on the SSO response"""
if sso_role is None:
return False
if not is_valid_litellm_user_role(sso_role):
verbose_proxy_logger.debug(
f"SSO role '{sso_role}' is not a valid LiteLLM user role. "
"Ignoring role from SSO response. See LitellmUserRoles enum for valid roles."
)
return False
return True
def apply_user_info_values_to_sso_user_defined_values(
user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]],
@@ -564,12 +582,22 @@ def apply_user_info_values_to_sso_user_defined_values(
if user_info is not None and user_info.user_id is not None:
user_defined_values["user_id"] = user_info.user_id
# Check if user_role already exists in user_defined_values (from JWT/SSO response)
if user_defined_values.get("user_role") is None:
# SSO role takes precedence - only use DB role if SSO didn't provide one
# This ensures SSO is the authoritative source for user roles
sso_role = user_defined_values.get("user_role")
db_role = user_info.user_role if user_info else None
if _should_use_role_from_sso_response(sso_role):
# SSO provided a valid role, keep it and log that we're using it
verbose_proxy_logger.info(f"Using SSO role: {sso_role} (DB role was: {db_role})")
else:
# SSO didn't provide a valid role, fall back to DB role or default
if user_info is None or user_info.user_role is None:
user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value
verbose_proxy_logger.debug("No SSO or DB role found, using default: INTERNAL_USER_VIEW_ONLY")
else:
user_defined_values["user_role"] = user_info.user_role
verbose_proxy_logger.debug(f"Using DB role: {user_info.user_role}")
# Preserve the user's existing models from the database
if user_info is not None and hasattr(user_info, "models") and user_info.models:
@@ -1540,7 +1568,15 @@ class SSOAuthenticationHandler:
},
)
# generic client id
# Extract user_role from result (works for all SSO providers)
if result is not None:
_user_role = getattr(result, "user_role", None)
if _user_role is not None:
# Convert enum to string if needed
user_role = _user_role.value if isinstance(_user_role, LitellmUserRoles) else _user_role
verbose_proxy_logger.debug(f"Extracted user_role from SSO result: {user_role}")
# generic client id - override with custom attribute name if specified
if generic_client_id is not None and result is not None:
generic_user_role_attribute_name = os.getenv(
"GENERIC_USER_ROLE_ATTRIBUTE", "role"
@@ -533,6 +533,75 @@ def test_apply_user_info_values_to_sso_user_defined_values_with_models():
assert sso_user_defined_values["models"] == ["no-default-models"]
def test_apply_user_info_values_sso_role_takes_precedence():
"""
Test that SSO role takes precedence over DB role.
When Microsoft SSO returns a user_role, it should be used instead of the role stored in the database.
This ensures SSO is the authoritative source for user roles.
"""
from litellm.proxy._types import LiteLLM_UserTable, SSOUserDefinedValues
from litellm.proxy.management_endpoints.ui_sso import (
apply_user_info_values_to_sso_user_defined_values,
)
user_info = LiteLLM_UserTable(
user_id="123",
user_email="test@example.com",
user_role="internal_user_viewer",
models=["model-1"],
)
user_defined_values: SSOUserDefinedValues = {
"models": [],
"user_id": "456",
"user_email": "test@example.com",
"user_role": "proxy_admin_viewer",
"max_budget": None,
"budget_duration": None,
}
sso_user_defined_values = apply_user_info_values_to_sso_user_defined_values(
user_info=user_info,
user_defined_values=user_defined_values,
)
assert sso_user_defined_values is not None
assert sso_user_defined_values["user_id"] == "123"
assert sso_user_defined_values["user_role"] == "proxy_admin_viewer"
assert sso_user_defined_values["models"] == ["model-1"]
def test_get_user_email_and_id_extracts_microsoft_role():
"""
Test that _get_user_email_and_id_from_result extracts user_role from Microsoft SSO.
This ensures Microsoft SSO roles (from app_roles in id_token) are properly
extracted and converted from enum to string.
"""
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.management_endpoints.types import CustomOpenID
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
result = CustomOpenID(
id="test-user-id",
email="test@example.com",
display_name="Test User",
provider="microsoft",
team_ids=["team-1"],
user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
)
parsed = SSOAuthenticationHandler._get_user_email_and_id_from_result(
result=result,
generic_client_id=None,
)
assert parsed.get("user_email") == "test@example.com"
assert parsed.get("user_id") == "test-user-id"
assert parsed.get("user_role") == "proxy_admin_viewer"
@pytest.mark.asyncio
async def test_get_user_info_from_db():
"""
@@ -2068,6 +2137,7 @@ class TestProcessSSOJWTAccessToken:
async def test_get_ui_settings_includes_api_doc_base_url():
"""Ensure the UI settings endpoint surfaces the optional API doc override."""
from fastapi import Request
from litellm.proxy.management_endpoints.ui_sso import get_ui_settings
mock_request = Request(