From e2fc22520146fb80547606a12fe9ca38f2f3dfcd Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 18 Nov 2025 15:35:38 -0800 Subject: [PATCH] [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 --- ...odel_prices_and_context_window_backup.json | 2 +- litellm/proxy/management_endpoints/ui_sso.py | 46 ++++++++++-- .../proxy/management_endpoints/test_ui_sso.py | 70 +++++++++++++++++++ 3 files changed, 112 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index fa13c2f6a0..351e01bc45 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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 }, diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 0add5c02cf..44b593efea 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -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" diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index ccf8bf13e0..4f5f4e0d85 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -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(