diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index d4969b2d03..e5341d35f6 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -564,10 +564,12 @@ 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 - if user_info is None or user_info.user_role is None: - user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY - else: - user_defined_values["user_role"] = user_info.user_role + # Check if user_role already exists in user_defined_values (from JWT/SSO response) + if user_defined_values.get("user_role") is None: + if user_info is None or user_info.user_role is None: + user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + else: + user_defined_values["user_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: @@ -1581,7 +1583,7 @@ class SSOAuthenticationHandler: user_id=user_id, user_email=user_email, max_budget=max_internal_user_budget, - user_role=None, + user_role=user_role, budget_duration=internal_user_budget_duration, ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py b/tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py index f6248d3662..5ecf076757 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py +++ b/tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py @@ -1,55 +1,92 @@ -""" -Unit tests for EntraID app roles JWT claim extraction. - -This module tests the get_app_roles_from_id_token method to ensure it correctly -extracts app roles from Microsoft EntraID JWT tokens and prevents regressions. -""" - -import pytest import jwt from litellm.proxy.management_endpoints.ui_sso import MicrosoftSSOHandler +from litellm.proxy.management_endpoints.types import get_litellm_user_role +from litellm.proxy._types import LitellmUserRoles -class TestEntraIDAppRoles: - """Test EntraID app roles extraction from JWT tokens""" +def test_extracts_proxy_admin_role_from_jwt(): + """Ensure supported app roles like 'proxy_admin' are extracted from the id_token.""" + payload = { + "sub": "user123", + "email": "admin@company.com", + "app_roles": ["proxy_admin"], + "aud": "litellm-app", + "iss": "https://login.microsoftonline.com/tenant-id/v2.0", + "exp": 9999999999, + } - def test_get_app_roles_from_id_token_works_without_roles(self): - """Test that JWT token works fine without app_roles claim""" - # Arrange - Token without app_roles (normal user) - payload = { - "sub": "user123", - "email": "user@company.com", - "aud": "litellm-app", - "iss": "https://login.microsoftonline.com/tenant-id/v2.0", - "exp": 9999999999, - } - no_roles_token = jwt.encode(payload, "secret", algorithm="HS256") + token = jwt.encode(payload, "secret", algorithm="HS256") + roles = MicrosoftSSOHandler.get_app_roles_from_id_token(token) - # Act - result = MicrosoftSSOHandler.get_app_roles_from_id_token(no_roles_token) + assert roles == ["proxy_admin"] - # Assert - Should return empty list, not error - assert result == [] - assert len(result) == 0 - def test_get_app_roles_from_id_token_assigns_roles_when_present(self): - """Test that valid app roles are properly assigned when present""" - # Arrange - Token with valid roles - payload = { - "sub": "user123", - "email": "admin@company.com", - "app_roles": ["proxy_admin"], - "aud": "litellm-app", - "iss": "https://login.microsoftonline.com/tenant-id/v2.0", - "exp": 9999999999, - } - valid_roles_token = jwt.encode(payload, "secret", algorithm="HS256") +def test_maps_internal_user_role(): + """Ensure internal_user role is correctly mapped to LitellmUserRoles.""" + payload = { + "sub": "user456", + "email": "user@company.com", + "app_roles": ["internal_user"], + "aud": "litellm-app", + "iss": "https://login.microsoftonline.com/tenant-id/v2.0", + "exp": 9999999999, + } - # Act - result = MicrosoftSSOHandler.get_app_roles_from_id_token(valid_roles_token) + token = jwt.encode(payload, "secret", algorithm="HS256") + roles = MicrosoftSSOHandler.get_app_roles_from_id_token(token) + + # Map to LitellmUserRoles + chosen = None + for r in roles: + mapped = get_litellm_user_role(r) + if mapped is not None: + chosen = mapped + break + + assert chosen == LitellmUserRoles.INTERNAL_USER + + +def test_maps_proxy_admin_viewer_role(): + """Ensure proxy_admin_viewer role is correctly mapped.""" + payload = { + "sub": "user789", + "email": "viewer@company.com", + "app_roles": ["proxy_admin_viewer"], + "aud": "litellm-app", + "iss": "https://login.microsoftonline.com/tenant-id/v2.0", + "exp": 9999999999, + } + + token = jwt.encode(payload, "secret", algorithm="HS256") + roles = MicrosoftSSOHandler.get_app_roles_from_id_token(token) + + chosen = None + for r in roles: + mapped = get_litellm_user_role(r) + if mapped is not None: + chosen = mapped + break + + assert chosen == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + + +def test_defaults_to_internal_user_viewer_when_no_role(): + """Ensure default role is internal_user_viewer when no app role is present.""" + payload = { + "sub": "user_no_role", + "email": "noRole@company.com", + "aud": "litellm-app", + "iss": "https://login.microsoftonline.com/tenant-id/v2.0", + "exp": 9999999999, + } + + token = jwt.encode(payload, "secret", algorithm="HS256") + roles = MicrosoftSSOHandler.get_app_roles_from_id_token(token) + + assert roles == [] + + # Default role would be internal_user_viewer + default_role = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + assert default_role.value == "internal_user_viewer" - # Assert - Should extract the role - assert result == ["proxy_admin"] - assert len(result) == 1 - assert "proxy_admin" in result