From 9274860aa28c5d69d696077b92283dfcad90e0a5 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 16 Dec 2025 14:54:26 -0800 Subject: [PATCH 1/6] Base commit --- litellm/proxy/management_endpoints/ui_sso.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index d1db21a270..5094fc5de9 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -594,7 +594,7 @@ def _build_sso_user_update_data( user_id: Optional[str], ) -> dict: """ - Build the update data dictionary for SSO user upsert. + Build the update data dictionary for SSO user upsert Args: result: The SSO response containing user information From 42d7d757a3bb1e5da02ed53c494ce6abca2f7e61 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 16 Dec 2025 15:48:26 -0800 Subject: [PATCH 2/6] Adding role mappings to SSOConfig DB --- litellm/proxy/management_endpoints/ui_sso.py | 2 +- litellm/proxy/proxy_server.py | 1 + .../proxy_setting_endpoints.py | 12 ++ .../proxy/management_endpoints/ui_sso.py | 34 ++++- .../test_proxy_setting_endpoints.py | 135 ++++++++++++++++++ 5 files changed, 182 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 5094fc5de9..d1db21a270 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -594,7 +594,7 @@ def _build_sso_user_update_data( user_id: Optional[str], ) -> dict: """ - Build the update data dictionary for SSO user upsert + Build the update data dictionary for SSO user upsert. Args: result: The SSO response containing user information diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8fdea95d7a..dfadab1d53 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3578,6 +3578,7 @@ class ProxyConfig: ) if sso_settings is not None: # Capitalize all keys in sso_settings dictionary + sso_settings.sso_settings.pop("role_mappings", None) uppercase_sso_settings = { key.upper(): value for key, value in sso_settings.sso_settings.items() diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 9c99b625e9..d9a41d38b2 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -433,10 +433,21 @@ async def get_sso_settings(): if sso_db_record and sso_db_record.sso_settings: # Load settings from database sso_settings_dict = dict(sso_db_record.sso_settings) + + # Extract role_mappings before removing it (it's a dict, not an env variable) + role_mappings_data = sso_settings_dict.pop("role_mappings", None) + role_mappings = None + if role_mappings_data: + from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings + if isinstance(role_mappings_data, dict): + role_mappings = RoleMappings(**role_mappings_data) + elif isinstance(role_mappings_data, RoleMappings): + role_mappings = role_mappings_data decrypted_sso_settings_dict = proxy_config._decrypt_and_set_db_env_variables(environment_variables=sso_settings_dict) # Build SSO config with database values or environment fallback + sso_config = SSOConfig( google_client_id=decrypted_sso_settings_dict.get("google_client_id", None), google_client_secret=decrypted_sso_settings_dict.get("google_client_secret", None), @@ -451,6 +462,7 @@ async def get_sso_settings(): proxy_base_url=decrypted_sso_settings_dict.get("proxy_base_url", None), user_email=decrypted_sso_settings_dict.get("user_email"), ui_access_mode=decrypted_sso_settings_dict.get("ui_access_mode"), + role_mappings=role_mappings, ) # Get the schema for UI display diff --git a/litellm/types/proxy/management_endpoints/ui_sso.py b/litellm/types/proxy/management_endpoints/ui_sso.py index 820b016440..187d8c97c0 100644 --- a/litellm/types/proxy/management_endpoints/ui_sso.py +++ b/litellm/types/proxy/management_endpoints/ui_sso.py @@ -1,10 +1,12 @@ -from typing import List, Literal, Optional, Union +from typing import Dict, List, Literal, Optional, Union from pydantic import Field from typing_extensions import TypedDict from litellm.types.utils import LiteLLMPydanticObjectBase +from litellm.proxy._types import LitellmUserRoles + class LiteLLM_UpperboundKeyGenerateParams(LiteLLMPydanticObjectBase): """ @@ -60,6 +62,30 @@ class AccessControl_UI_AccessMode(LiteLLMPydanticObjectBase): sso_group_jwt_field: str +class RoleMappings(LiteLLMPydanticObjectBase): + """ + Configuration for mapping SSO groups to LiteLLM roles. + + The system will look at the group_claim field in the SSO token to determine + which role to assign the user based on the roles mapping. + """ + + provider: str = Field( + description="SSO Provider name (e.g., 'google', 'microsoft', 'generic')" + ) + group_claim: str = Field( + description="The field name in the SSO token that contains the groups array (e.g., 'groups', 'roles')" + ) + default_role: Optional[LitellmUserRoles] = Field( + default=None, + description="Default role to assign if user's groups don't match any role mappings. Must be a valid LitellmUserRoles value (e.g., 'proxy_admin', 'internal_user', 'proxy_admin_viewer')" + ) + roles: Dict[LitellmUserRoles, List[str]] = Field( + default_factory=dict, + description="Mapping of LiteLLM role names to arrays of SSO group names. Example: {'proxy_admin': ['group-1', 'group-2'], 'proxy_admin_viewer': ['group-3']}" + ) + + class SSOConfig(LiteLLMPydanticObjectBase): """ Configuration for SSO environment variables and settings @@ -127,6 +153,12 @@ class SSOConfig(LiteLLMPydanticObjectBase): description="Access mode for the UI", ) + # Role Mappings + role_mappings: Optional[RoleMappings] = Field( + default=None, + description="Configuration for mapping SSO groups to LiteLLM roles based on group claims in the SSO token", + ) + class DefaultTeamSSOParams(LiteLLMPydanticObjectBase): """ diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index d3c9915119..8fdfd6897a 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -290,6 +290,10 @@ class TestProxySettingEndpoints: assert "google_client_id" in data["field_schema"]["properties"] assert "description" in data["field_schema"]["properties"]["google_client_id"] + # Verify role_mappings is present in response (can be None if not set) + assert "role_mappings" in values + assert values["role_mappings"] is None + # Verify find_unique was called with correct parameters mock_prisma.db.litellm_ssoconfig.find_unique.assert_called_once() call_args = mock_prisma.db.litellm_ssoconfig.find_unique.call_args @@ -863,6 +867,10 @@ class TestProxySettingEndpoints: assert values["google_client_secret"] == "decrypted_google_secret" assert values["microsoft_client_id"] == "decrypted_microsoft_id" assert values["proxy_base_url"] == "https://decrypted.example.com" + + # Verify role_mappings is present in response (can be None if not set) + assert "role_mappings" in values + assert values["role_mappings"] is None def test_update_sso_settings_to_database(self, mock_proxy_config, mock_auth, monkeypatch): """Test updating SSO settings saves to the dedicated database table""" @@ -1062,6 +1070,7 @@ class TestProxySettingEndpoints: assert values.get("google_client_id") is None assert values.get("google_client_secret") is None assert values.get("microsoft_client_id") is None + assert values.get("role_mappings") is None def test_update_sso_settings_no_database_connection(self, mock_proxy_config, mock_auth, monkeypatch): """Test updating SSO settings when database is not connected""" @@ -1088,3 +1097,129 @@ class TestProxySettingEndpoints: data = response.json() assert "error" in data["detail"] assert "Database not connected" in data["detail"]["error"] + + def test_get_sso_settings_with_role_mappings(self, mock_proxy_config, mock_auth, monkeypatch): + """Test getting SSO settings when role_mappings is present in database""" + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy._types import LitellmUserRoles + + # Mock the prisma client with database record containing role_mappings + mock_prisma = MagicMock() + mock_db_record = MagicMock() + mock_db_record.sso_settings = { + "google_client_id": "test_google_client_id", + "role_mappings": { + "provider": "google", + "group_claim": "groups", + "default_role": LitellmUserRoles.INTERNAL_USER, + "roles": { + LitellmUserRoles.PROXY_ADMIN: ["admin-group"], + }, + }, + } + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_db_record) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # Mock decryption to return the values as-is (role_mappings should not be passed to decryption) + from litellm.proxy.proxy_server import proxy_config + def mock_decrypt(environment_variables): + # role_mappings should not be in environment_variables since it's extracted before decryption + assert "role_mappings" not in environment_variables + return environment_variables + + monkeypatch.setattr( + proxy_config, "_decrypt_and_set_db_env_variables", mock_decrypt + ) + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + data = response.json() + + # Verify role_mappings is returned correctly + values = data["values"] + assert "role_mappings" in values + assert values["role_mappings"] is not None + assert values["role_mappings"]["provider"] == "google" + assert values["role_mappings"]["group_claim"] == "groups" + assert values["role_mappings"]["default_role"] == LitellmUserRoles.INTERNAL_USER + assert values["role_mappings"]["roles"][LitellmUserRoles.PROXY_ADMIN] == ["admin-group"] + + def test_role_mappings_stored_and_retrieved(self, mock_proxy_config, mock_auth, monkeypatch): + """Test that role_mappings is properly stored and retrieved from SSO settings""" + import json + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + # Mock the prisma client + mock_prisma = MagicMock() + mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() + mock_prisma.db.litellm_config = MagicMock() + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.update = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # Mock encryption to return values as-is + from litellm.proxy.proxy_server import proxy_config + monkeypatch.setattr(proxy_config, "_encrypt_env_variables", lambda environment_variables: environment_variables) + + # SSO settings with role_mappings + role_mappings_data = { + "provider": "google", + "group_claim": "groups", + "default_role": LitellmUserRoles.INTERNAL_USER, + "roles": { + LitellmUserRoles.PROXY_ADMIN: ["admin-group"], + LitellmUserRoles.INTERNAL_USER: ["user-group"], + }, + } + + new_sso_settings = { + "google_client_id": "test_google_id", + "role_mappings": role_mappings_data, + } + + response = client.patch("/update/sso_settings", json=new_sso_settings) + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "success" + assert "role_mappings" in data["settings"] + + # Verify role_mappings structure in response + returned_role_mappings = data["settings"]["role_mappings"] + assert returned_role_mappings["provider"] == "google" + assert returned_role_mappings["group_claim"] == "groups" + assert returned_role_mappings["default_role"] == LitellmUserRoles.INTERNAL_USER + assert returned_role_mappings["roles"][LitellmUserRoles.PROXY_ADMIN] == ["admin-group"] + + # Verify upsert was called with role_mappings in the data + assert mock_prisma.db.litellm_ssoconfig.upsert.called + call_args = mock_prisma.db.litellm_ssoconfig.upsert.call_args + create_data = call_args.kwargs["data"]["create"] + stored_sso_settings = json.loads(create_data["sso_settings"]) + assert "role_mappings" in stored_sso_settings + assert stored_sso_settings["role_mappings"]["provider"] == "google" + + # Now test retrieving role_mappings + mock_db_record = MagicMock() + mock_db_record.sso_settings = stored_sso_settings + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_db_record) + monkeypatch.setattr( + proxy_config, "_decrypt_and_set_db_env_variables", lambda environment_variables: environment_variables + ) + + get_response = client.get("/get/sso_settings") + assert get_response.status_code == 200 + get_data = get_response.json() + + # Verify role_mappings is returned correctly + assert "role_mappings" in get_data["values"] + retrieved_role_mappings = get_data["values"]["role_mappings"] + assert retrieved_role_mappings is not None + assert retrieved_role_mappings["provider"] == "google" + assert retrieved_role_mappings["group_claim"] == "groups" + assert retrieved_role_mappings["default_role"] == LitellmUserRoles.INTERNAL_USER From 58330f852d82bd42f02f2e920cdb14ae0b927024 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 16 Dec 2025 17:36:08 -0800 Subject: [PATCH 3/6] WIP waiting for okta --- litellm/proxy/management_endpoints/ui_sso.py | 134 ++++++++++++++++++- 1 file changed, 127 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index d1db21a270..7a404bcda8 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -85,6 +85,58 @@ else: router = APIRouter() +def determine_role_from_groups( + user_groups: List[str], + role_mappings: "RoleMappings", +) -> Optional[LitellmUserRoles]: + """ + Determine the highest privilege role for a user based on their groups. + + Role hierarchy (highest to lowest): + - proxy_admin + - proxy_admin_viewer + - internal_user + - internal_user_viewer + + Args: + user_groups: List of group names from the SSO token + role_mappings: RoleMappings configuration object + + Returns: + The highest privilege role found, or default_role if no matches, or None + """ + if not role_mappings.roles: + # No role mappings configured, return default_role + return role_mappings.default_role + + # Role hierarchy (highest to lowest) + role_hierarchy = [ + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + ] + + # Convert user_groups to a set for efficient lookup + user_groups_set = set(user_groups) if isinstance(user_groups, list) else set() + + # Find the highest privilege role the user belongs to + for role in role_hierarchy: + if role in role_mappings.roles: + role_groups = role_mappings.roles[role] + if isinstance(role_groups, list) and user_groups_set.intersection(set(role_groups)): + verbose_proxy_logger.debug( + f"User groups {user_groups} matched role '{role.value}' via groups: {role_groups}" + ) + return role + + # No matching groups found, return default_role + verbose_proxy_logger.debug( + f"User groups {user_groups} did not match any role mappings, using default_role: {role_mappings.default_role}" + ) + return role_mappings.default_role + + def process_sso_jwt_access_token( access_token_str: Optional[str], sso_jwt_handler: Optional[JWTHandler], @@ -243,6 +295,7 @@ def generic_response_convertor( response, jwt_handler: JWTHandler, sso_jwt_handler: Optional[JWTHandler] = None, + role_mappings: Optional["RoleMappings"] = None, ) -> CustomOpenID: generic_user_id_attribute_name = os.getenv( "GENERIC_USER_ID_ATTRIBUTE", "preferred_username" @@ -281,16 +334,48 @@ def generic_response_convertor( team_ids = jwt_handler.get_team_ids_from_jwt(cast(dict, response)) all_teams.extend(team_ids) - # Extract user role from SSO response - user_role_from_sso = get_nested_value(response, generic_user_role_attribute_name) + # Determine user role based on role_mappings if available + # Only apply role_mappings for GENERIC SSO provider user_role: Optional[LitellmUserRoles] = None - if user_role_from_sso is not None: - role = get_litellm_user_role(user_role_from_sso) - if role is not None: - user_role = role + + if role_mappings is not None and role_mappings.provider.lower() in ["generic", "okta"]: + # Use role_mappings to determine role from groups + group_claim = role_mappings.group_claim + user_groups_raw = get_nested_value(response, group_claim) + + # Handle different formats: could be a list, string (comma-separated), or single value + user_groups: List[str] = [] + if isinstance(user_groups_raw, list): + user_groups = [str(g) for g in user_groups_raw] + elif isinstance(user_groups_raw, str): + # Handle comma-separated string + user_groups = [g.strip() for g in user_groups_raw.split(",") if g.strip()] + elif user_groups_raw is not None: + # Single value + user_groups = [str(user_groups_raw)] + + if user_groups: + user_role = determine_role_from_groups(user_groups, role_mappings) verbose_proxy_logger.debug( - f"Found valid LitellmUserRoles '{role.value}' from SSO attribute '{generic_user_role_attribute_name}'" + f"Determined role '{user_role.value if user_role else None}' from groups '{user_groups}' using role_mappings" ) + else: + # No groups found, use default_role + user_role = role_mappings.default_role + verbose_proxy_logger.debug( + f"No groups found in '{group_claim}', using default_role: {role_mappings.default_role}" + ) + + # Fallback to existing logic if role_mappings not used + if user_role is None: + user_role_from_sso = get_nested_value(response, generic_user_role_attribute_name) + if user_role_from_sso is not None: + role = get_litellm_user_role(user_role_from_sso) + if role is not None: + user_role = role + verbose_proxy_logger.debug( + f"Found valid LitellmUserRoles '{role.value}' from SSO attribute '{generic_user_role_attribute_name}'" + ) return CustomOpenID( id=get_nested_value(response, generic_user_id_attribute_name), @@ -369,6 +454,40 @@ async def get_generic_sso_response( userinfo_endpoint=generic_userinfo_endpoint, ) + # Get role_mappings from SSO settings if available + role_mappings: Optional["RoleMappings"] = None + try: + from litellm.proxy.utils import get_prisma_client_or_throw + + prisma_client = get_prisma_client_or_throw( + "Prisma client is None, connect a database to your proxy" + ) + + # Get SSO config from dedicated table + sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( + where={"id": "sso_config"} + ) + + if sso_db_record and sso_db_record.sso_settings: + sso_settings_dict = dict(sso_db_record.sso_settings) + role_mappings_data = sso_settings_dict.get("role_mappings") + + if role_mappings_data: + from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings + if isinstance(role_mappings_data, dict): + role_mappings = RoleMappings(**role_mappings_data) + elif isinstance(role_mappings_data, RoleMappings): + role_mappings = role_mappings_data + + verbose_proxy_logger.debug( + f"Loaded role_mappings for provider '{role_mappings.provider}'" + ) + except Exception as e: + # If we can't load role_mappings, continue with existing logic + verbose_proxy_logger.debug( + f"Could not load role_mappings from database: {e}. Continuing with existing role logic." + ) + def response_convertor(response, client): nonlocal received_response # return for user debugging received_response = response @@ -376,6 +495,7 @@ async def get_generic_sso_response( response=response, jwt_handler=jwt_handler, sso_jwt_handler=sso_jwt_handler, + role_mappings=role_mappings, ) SSOProvider = create_provider( From bf76e66d2ccaace63eb6353ab92e37b57d99b349 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 19 Dec 2025 09:42:59 -0800 Subject: [PATCH 4/6] Working SSO Mapping new user, overrides default user settings --- litellm/proxy/management_endpoints/ui_sso.py | 48 +++++++++++++++++++- 1 file changed, 46 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 7a404bcda8..5afccc6fe5 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1173,8 +1173,44 @@ async def insert_sso_user( if user_defined_values is None: raise ValueError("user_defined_values is None") + # Check if role_mappings is configured in SSO settings + role_mappings_configured = False + try: + from litellm.proxy.utils import get_prisma_client_or_throw + + prisma_client = get_prisma_client_or_throw( + "Prisma client is None, connect a database to your proxy" + ) + + # Get SSO config from dedicated table + sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( + where={"id": "sso_config"} + ) + + if sso_db_record and sso_db_record.sso_settings: + sso_settings_dict = dict(sso_db_record.sso_settings) + role_mappings_data = sso_settings_dict.get("role_mappings") + role_mappings_configured = role_mappings_data is not None + except Exception as e: + # If we can't check role_mappings, continue with existing logic + verbose_proxy_logger.debug( + f"Could not check role_mappings configuration: {e}. Using default behavior." + ) + + # Apply default_internal_user_params if litellm.default_internal_user_params: - user_defined_values.update(litellm.default_internal_user_params) # type: ignore + # If role_mappings is configured and user_role is already set from SSO, preserve it + if role_mappings_configured and user_defined_values.get("user_role") is not None: + # Preserve the SSO-extracted role, but apply other defaults + preserved_role = user_defined_values.get("user_role") + user_defined_values.update(litellm.default_internal_user_params) # type: ignore + user_defined_values["user_role"] = preserved_role # Restore preserved role + verbose_proxy_logger.debug( + f"Preserved SSO-extracted role '{preserved_role}' (role_mappings configured)" + ) + else: + # Default behavior: update all values including role + user_defined_values.update(litellm.default_internal_user_params) # type: ignore # Set budget for internal users if user_defined_values.get("user_role") == LitellmUserRoles.INTERNAL_USER.value: @@ -1812,7 +1848,15 @@ class SSOAuthenticationHandler: ) user_id = getattr(result, "id", None) user_email = getattr(result, "email", None) - user_role = getattr(result, generic_user_role_attribute_name, None) # type: ignore + if user_role is None: + _role_from_attr = getattr(result, generic_user_role_attribute_name, None) # type: ignore + if _role_from_attr is not None: + # Convert enum to string if needed + user_role = ( + _role_from_attr.value + if isinstance(_role_from_attr, LitellmUserRoles) + else _role_from_attr + ) if user_id is None and result is not None: _first_name = getattr(result, "first_name", "") or "" From 313a613a13e2756308ce828924864815ac36bd38 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 19 Dec 2025 10:56:11 -0800 Subject: [PATCH 5/6] Adding tests --- .../proxy/management_endpoints/test_ui_sso.py | 105 ++++++++++++++++++ 1 file changed, 105 insertions(+) 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 500fc67de8..2082946657 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -3043,3 +3043,108 @@ class TestAddMissingTeamMember: assert set(added_teams) == set( expected_teams_added ), f"Expected teams {expected_teams_added}, but got {added_teams}" + + +@pytest.mark.asyncio +async def test_role_mappings_override_default_internal_user_params(): + """ + Test that when role_mappings is configured in SSO settings, + the SSO-extracted role overrides default_internal_user_params role. + """ + from litellm.proxy._types import NewUserResponse, SSOUserDefinedValues + from litellm.proxy.management_endpoints.ui_sso import insert_sso_user + + # Save original default_internal_user_params + original_default_params = getattr(litellm, "default_internal_user_params", None) + + try: + # Set default_internal_user_params with a role that should be overridden + litellm.default_internal_user_params = { + "user_role": "internal_user", + "max_budget": 100, + "budget_duration": "30d", + "models": ["gpt-3.5-turbo"], + } + + # Mock SSO result + mock_result_openid = CustomOpenID( + id="test-user-123", + email="test@example.com", + display_name="Test User", + provider="microsoft", + team_ids=[], + ) + + # User defined values with SSO-extracted role (from role_mappings) + user_defined_values: SSOUserDefinedValues = { + "user_id": "test-user-123", + "user_email": "test@example.com", + "user_role": "proxy_admin", # Role from SSO role_mappings + "max_budget": None, + "budget_duration": None, + "models": [], + } + + # Mock Prisma client with SSO config that has role_mappings configured + mock_prisma = MagicMock() + mock_sso_config = MagicMock() + mock_sso_config.sso_settings = { + "role_mappings": { + "Admin": "proxy_admin", + "User": "internal_user", + } + } + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock( + return_value=mock_sso_config + ) + + # Mock new_user function + mock_new_user_response = NewUserResponse( + user_id="test-user-123", + key="sk-xxxxx", + teams=None, + ) + + with patch( + "litellm.proxy.utils.get_prisma_client_or_throw", + return_value=mock_prisma, + ), patch( + "litellm.proxy.management_endpoints.ui_sso.new_user", + return_value=mock_new_user_response, + ) as mock_new_user: + # Act + result = await insert_sso_user( + result_openid=mock_result_openid, + user_defined_values=user_defined_values, + ) + + # Assert - verify new_user was called with preserved SSO role + mock_new_user.assert_called_once() + call_args = mock_new_user.call_args + new_user_request = call_args.kwargs["data"] + + # The role from SSO should be preserved, not overridden by default_internal_user_params + assert ( + new_user_request.user_role == "proxy_admin" + ), "SSO-extracted role should override default_internal_user_params role" + + # Other default params should still be applied + assert ( + new_user_request.max_budget == 100 + ), "max_budget from default_internal_user_params should be applied" + assert ( + new_user_request.budget_duration == "30d" + ), "budget_duration from default_internal_user_params should be applied" + + # Note: models are applied via _update_internal_new_user_params inside new_user, + # not in insert_sso_user, so we verify user_defined_values was updated correctly + # by checking that the function completed successfully and other defaults were applied + # The models will be applied when new_user processes the request + + finally: + # Restore original default_internal_user_params + if original_default_params is not None: + litellm.default_internal_user_params = original_default_params + else: + if hasattr(litellm, "default_internal_user_params"): + delattr(litellm, "default_internal_user_params") From 2afa5fc9fb72984d34f8ed32c296fc8b00303d48 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 22 Dec 2025 11:03:19 -0800 Subject: [PATCH 6/6] ruff check and mypy linting --- litellm/proxy/management_endpoints/ui_sso.py | 126 ++++++++++++------- 1 file changed, 80 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 013859e1d5..dc976e1ce6 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -391,20 +391,8 @@ def generic_response_convertor( ) -async def get_generic_sso_response( - request: Request, - jwt_handler: JWTHandler, - sso_jwt_handler: Optional[ - JWTHandler - ], # sso specific jwt handler - used for restricted sso group access control - generic_client_id: str, - redirect_url: str, -) -> Tuple[Union[OpenID, dict], Optional[dict]]: # return received response - # make generic sso provider - from fastapi_sso.sso.base import DiscoveryDocument - from fastapi_sso.sso.generic import create_provider - - received_response: Optional[dict] = None +def _setup_generic_sso_env_vars(generic_client_id: str, redirect_url: str) -> Tuple[str, List[str], str, str, str, bool]: + """Setup and validate Generic SSO environment variables.""" generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None) generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split(" ") generic_authorization_endpoint = os.getenv("GENERIC_AUTHORIZATION_ENDPOINT", None) @@ -413,6 +401,8 @@ async def get_generic_sso_response( generic_include_client_id = ( os.getenv("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true" ) + + # Validate required environment variables if generic_client_secret is None: raise ProxyException( message="GENERIC_CLIENT_SECRET not set. Set it in .env file", @@ -441,6 +431,7 @@ async def get_generic_sso_response( param="GENERIC_USERINFO_ENDPOINT", code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) + verbose_proxy_logger.debug( f"authorization_endpoint: {generic_authorization_endpoint}\ntoken_endpoint: {generic_token_endpoint}\nuserinfo_endpoint: {generic_userinfo_endpoint}" ) @@ -448,6 +439,80 @@ async def get_generic_sso_response( f"GENERIC_REDIRECT_URI: {redirect_url}\nGENERIC_CLIENT_ID: {generic_client_id}\n" ) + return ( + generic_client_secret, + generic_scope, + generic_authorization_endpoint, + generic_token_endpoint, + generic_userinfo_endpoint, + generic_include_client_id, + ) + + +async def _setup_role_mappings() -> Optional["RoleMappings"]: + """Setup role mappings from SSO database settings.""" + role_mappings: Optional["RoleMappings"] = None + try: + from litellm.proxy.utils import get_prisma_client_or_throw + + prisma_client = get_prisma_client_or_throw( + "Prisma client is None, connect a database to your proxy" + ) + + # Get SSO config from dedicated table + sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( + where={"id": "sso_config"} + ) + + if sso_db_record and sso_db_record.sso_settings: + sso_settings_dict = dict(sso_db_record.sso_settings) + role_mappings_data = sso_settings_dict.get("role_mappings") + + if role_mappings_data: + from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings + if isinstance(role_mappings_data, dict): + role_mappings = RoleMappings(**role_mappings_data) + elif isinstance(role_mappings_data, RoleMappings): + role_mappings = role_mappings_data + + if role_mappings: + verbose_proxy_logger.debug( + f"Loaded role_mappings for provider '{role_mappings.provider}'" + ) + except Exception as e: + # If we can't load role_mappings, continue with existing logic + verbose_proxy_logger.debug( + f"Could not load role_mappings from database: {e}. Continuing with existing role logic." + ) + + return role_mappings + + +async def get_generic_sso_response( + request: Request, + jwt_handler: JWTHandler, + sso_jwt_handler: Optional[ + JWTHandler + ], # sso specific jwt handler - used for restricted sso group access control + generic_client_id: str, + redirect_url: str, +) -> Tuple[Union[OpenID, dict], Optional[dict]]: # return received response + # make generic sso provider + from fastapi_sso.sso.base import DiscoveryDocument + from fastapi_sso.sso.generic import create_provider + + received_response: Optional[dict] = None + + # Setup environment variables + ( + generic_client_secret, + generic_scope, + generic_authorization_endpoint, + generic_token_endpoint, + generic_userinfo_endpoint, + generic_include_client_id, + ) = _setup_generic_sso_env_vars(generic_client_id, redirect_url) + discovery = DiscoveryDocument( authorization_endpoint=generic_authorization_endpoint, token_endpoint=generic_token_endpoint, @@ -455,38 +520,7 @@ async def get_generic_sso_response( ) # Get role_mappings from SSO settings if available - role_mappings: Optional["RoleMappings"] = None - try: - from litellm.proxy.utils import get_prisma_client_or_throw - - prisma_client = get_prisma_client_or_throw( - "Prisma client is None, connect a database to your proxy" - ) - - # Get SSO config from dedicated table - sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( - where={"id": "sso_config"} - ) - - if sso_db_record and sso_db_record.sso_settings: - sso_settings_dict = dict(sso_db_record.sso_settings) - role_mappings_data = sso_settings_dict.get("role_mappings") - - if role_mappings_data: - from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings - if isinstance(role_mappings_data, dict): - role_mappings = RoleMappings(**role_mappings_data) - elif isinstance(role_mappings_data, RoleMappings): - role_mappings = role_mappings_data - - verbose_proxy_logger.debug( - f"Loaded role_mappings for provider '{role_mappings.provider}'" - ) - except Exception as e: - # If we can't load role_mappings, continue with existing logic - verbose_proxy_logger.debug( - f"Could not load role_mappings from database: {e}. Continuing with existing role logic." - ) + role_mappings = await _setup_role_mappings() def response_convertor(response, client): nonlocal received_response # return for user debugging