scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #21503 from BerriAI/litellm_sso_role_map

Browse source 
fix(sso): preserve SSO role regardless of role_mappings config
This commit is contained in:
yuneng-jiang 2026-02-18 17:31:49 -08:00
parent a09d3e9162
commit f49e4bf299
2 changed files with 99 additions and 57 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -1375,50 +1375,22 @@ 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
generic_user_role_mappings = os.getenv("GENERIC_USER_ROLE_MAPPINGS", None)
if generic_user_role_mappings is not None:
role_mappings_configured = True
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:
# 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 if it's a valid LiteLLM role,
# regardless of how it was determined (role_mappings, Microsoft app_roles,
# GENERIC_USER_ROLE_ATTRIBUTE, custom SSO handler, etc.)
sso_role = user_defined_values.get("user_role")
if _should_use_role_from_sso_response(sso_role):
# Preserve the SSO-extracted role, but apply other defaults
preserved_role = user_defined_values.get("user_role")
preserved_role = sso_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)"
f"Preserved SSO-extracted role '{preserved_role}'"
)
else:
# Default behavior: update all values including role
# SSO didn't provide a valid role, apply all defaults including role
user_defined_values.update(litellm.default_internal_user_params) # type: ignore
# Set budget for internal users

112
tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
View file

@ -1532,10 +1532,16 @@ class TestSSOHandlerIntegration:
SSOAuthenticationHandler.should_use_sso_handler(None, None, None) is False
)
@patch.dict(os.environ, {}, clear=False)
def test_get_redirect_url_for_sso(self):
"""Test the redirect URL generation for SSO"""
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
# Remove env vars that override request base_url so the test is
# isolated from local settings.
os.environ.pop("PROXY_BASE_URL", None)
os.environ.pop("SERVER_ROOT_PATH", None)
# Mock request object
mock_request = MagicMock()
mock_request.base_url = "https://test.litellm.ai/"
@ -3400,19 +3406,6 @@ async def test_role_mappings_override_default_internal_user_params():
"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",
@ -3421,9 +3414,6 @@ async def test_role_mappings_override_default_internal_user_params():
)
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:
@ -3450,11 +3440,6 @@ async def test_role_mappings_override_default_internal_user_params():
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
@ -3465,6 +3450,91 @@ async def test_role_mappings_override_default_internal_user_params():
delattr(litellm, "default_internal_user_params")
@pytest.mark.asyncio
async def test_sso_role_preserved_without_role_mappings():
"""
Test that SSO-extracted role is preserved even when role_mappings is NOT configured.
This covers the case where the role comes from Microsoft app_roles or
GENERIC_USER_ROLE_ATTRIBUTE (not from LiteLLM's role_mappings feature).
Previously, the role was only preserved when role_mappings was configured,
causing admin users to be downgraded to internal_user.
"""
from litellm.proxy._types import NewUserResponse, SSOUserDefinedValues
from litellm.proxy.management_endpoints.ui_sso import insert_sso_user
original_default_params = getattr(litellm, "default_internal_user_params", None)
try:
# Set default_internal_user_params (as most deployments do)
litellm.default_internal_user_params = {
"user_role": "internal_user",
"max_budget": 50,
}
# Mock SSO result from Microsoft with app_roles-derived admin role
mock_result_openid = CustomOpenID(
id="msft-user-456",
email="admin@company.com",
display_name="Admin User",
provider="microsoft",
team_ids=["group-1"],
user_role=None, # role is in user_defined_values, not on the OpenID result
)
# User defined values with role from Microsoft app_roles (NOT role_mappings)
user_defined_values: SSOUserDefinedValues = {
"user_id": "msft-user-456",
"user_email": "admin@company.com",
"user_role": "proxy_admin", # Role from Microsoft app_roles
"max_budget": None,
"budget_duration": None,
"models": [],
}
mock_new_user_response = NewUserResponse(
user_id="msft-user-456",
key="sk-xxxxx",
teams=None,
)
# No role_mappings configured anywhere - the role came from app_roles
with patch(
"litellm.proxy.management_endpoints.ui_sso.new_user",
return_value=mock_new_user_response,
) as mock_new_user:
_ = await insert_sso_user(
result_openid=mock_result_openid,
user_defined_values=user_defined_values,
)
mock_new_user.assert_called_once()
call_args = mock_new_user.call_args
new_user_request = call_args.kwargs["data"]
# SSO role should be preserved even without role_mappings configured
assert (
new_user_request.user_role == "proxy_admin"
), "SSO role from app_roles should not be overwritten by default_internal_user_params"
# Other defaults should still apply
assert (
new_user_request.max_budget == 50
), "max_budget from default_internal_user_params should be applied"
# Verify user_defined_values was also updated (it's mutated in-place)
assert (
user_defined_values["user_role"] == "proxy_admin"
), "user_defined_values should retain the SSO role after insert_sso_user"
finally:
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")
class TestSSOReadinessEndpoint:
"""Test the /sso/readiness endpoint"""