mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
fix(proxy): persist SSO display name as user_alias on login (#44065)
* fix(proxy): persist SSO display name as user_alias on login
Generic/Microsoft SSO already parsed the IdP display_name, first_name and last_name into the SSO result, but the user upsert only wrote user_email and user_role, so the Users table never showed a name. Store the display name (first + last as fallback) as user_alias on first login and on later logins of users whose alias is still empty; never overwrite an alias already set. Whitespace-only names are treated as missing.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(proxy): keep stored user_email when SSO login carries no email claim
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* revert: keep stored user_email change, login writes the IdP email as before
Reverts 47c65cecff. A stored email staying eligible for email-based account linking after the IdP stops sending it is not wanted; the PR goes back to the user_alias fix only
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
3a91029328
commit
9351463755
2 changed files with 196 additions and 5 deletions
|
|
@ -1853,10 +1853,36 @@ def _should_use_role_from_sso_response(sso_role: str | None) -> bool:
|
|||
return True
|
||||
|
||||
|
||||
class _SsoUserNames(Protocol):
|
||||
id: str | None
|
||||
display_name: str | None
|
||||
first_name: str | None
|
||||
last_name: str | None
|
||||
|
||||
|
||||
def _get_sso_user_alias(result: _SsoUserNames | Mapping[str, object] | None) -> str | None:
|
||||
"""Display name the IdP sent for the user, falling back to the joined first/last name."""
|
||||
if result is None:
|
||||
return None
|
||||
if isinstance(result, Mapping):
|
||||
raw_names: tuple[object, ...] = tuple(
|
||||
result.get(key) for key in ("id", "display_name", "first_name", "last_name")
|
||||
)
|
||||
else:
|
||||
raw_names = (result.id, result.display_name, result.first_name, result.last_name)
|
||||
user_id, display_name, first_name, last_name = (
|
||||
name.strip() or None if isinstance(name, str) else None for name in raw_names
|
||||
)
|
||||
if display_name and display_name != user_id:
|
||||
return display_name
|
||||
return " ".join(part for part in (first_name, last_name) if part) or None
|
||||
|
||||
|
||||
def _build_sso_user_update_data(
|
||||
result: Union["CustomOpenID", OpenID, dict] | None,
|
||||
result: Union["CustomOpenID", OpenID, Mapping[str, object]] | None,
|
||||
user_email: str | None,
|
||||
user_id: str | None,
|
||||
existing_user_alias: str | None = None,
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Build the update data dictionary for SSO user upsert.
|
||||
|
|
@ -1865,14 +1891,19 @@ def _build_sso_user_update_data(
|
|||
result: The SSO response containing user information
|
||||
user_email: The user's email from SSO
|
||||
user_id: The user's ID for logging purposes
|
||||
existing_user_alias: The user's current alias in the DB; only an empty alias is filled from SSO
|
||||
|
||||
Returns:
|
||||
dict: Update data containing user_email and optionally user_role if valid
|
||||
dict: Update data containing user_email, user_alias when newly available, and user_role if valid
|
||||
"""
|
||||
update_data: Final[dict[str, object]] = {"user_email": normalize_email(user_email)}
|
||||
sso_user_alias: Final = None if existing_user_alias else _get_sso_user_alias(result)
|
||||
update_data: Final[dict[str, object]] = {
|
||||
"user_email": normalize_email(user_email),
|
||||
**({"user_alias": sso_user_alias} if sso_user_alias is not None else {}),
|
||||
}
|
||||
|
||||
# Get SSO role from result and include if valid
|
||||
sso_role: Final = getattr(result, "user_role", None)
|
||||
sso_role: Final = result.user_role if isinstance(result, CustomOpenID) else None
|
||||
if sso_role is not None:
|
||||
# Convert enum to string if needed
|
||||
sso_role_str: Final = sso_role.value if isinstance(sso_role, LitellmUserRoles) else sso_role
|
||||
|
|
@ -2616,6 +2647,7 @@ async def insert_sso_user(
|
|||
new_user_request: Final = NewUserRequest(
|
||||
user_id=user_defined_values["user_id"],
|
||||
user_email=normalize_email(user_defined_values["user_email"]),
|
||||
user_alias=_get_sso_user_alias(result_openid),
|
||||
user_role=user_defined_values["user_role"],
|
||||
max_budget=user_defined_values["max_budget"],
|
||||
budget_duration=user_defined_values["budget_duration"],
|
||||
|
|
@ -3249,6 +3281,7 @@ class SSOAuthenticationHandler:
|
|||
result=result,
|
||||
user_email=user_email,
|
||||
user_id=user_id,
|
||||
existing_user_alias=user_info.user_alias if isinstance(user_info, LiteLLM_UserTable) else None,
|
||||
)
|
||||
|
||||
await _user_meta_db(UserRepository(prisma_client)).update_many(
|
||||
|
|
@ -3280,7 +3313,7 @@ class SSOAuthenticationHandler:
|
|||
if user_info is None:
|
||||
verbose_proxy_logger.debug("User not found in LiteLLM DB, skipping team member addition")
|
||||
return
|
||||
sso_teams: Final = getattr(result, "team_ids", [])
|
||||
sso_teams: Final = result.team_ids if isinstance(result, CustomOpenID) else []
|
||||
await add_missing_team_member(user_info=user_info, sso_teams=sso_teams)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -939,6 +939,83 @@ def test_build_sso_user_update_data_normalizes_email():
|
|||
assert "user_role" not in update_data
|
||||
|
||||
|
||||
def test_build_sso_user_update_data_fills_empty_user_alias_from_display_name():
|
||||
"""
|
||||
An existing SSO user with no alias gets the IdP display name on login.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.types import CustomOpenID
|
||||
from litellm.proxy.management_endpoints.ui_sso import _build_sso_user_update_data
|
||||
|
||||
sso_result = CustomOpenID(
|
||||
id="S-1-5-21-adfs-user",
|
||||
email="jane.doe@example.com",
|
||||
first_name="Jane",
|
||||
last_name="Doe",
|
||||
display_name="Doe, Jane",
|
||||
provider="generic",
|
||||
team_ids=[],
|
||||
)
|
||||
|
||||
update_data = _build_sso_user_update_data(
|
||||
result=sso_result,
|
||||
user_email="jane.doe@example.com",
|
||||
user_id="S-1-5-21-adfs-user",
|
||||
existing_user_alias=None,
|
||||
)
|
||||
|
||||
assert update_data == {"user_email": "jane.doe@example.com", "user_alias": "Doe, Jane"}
|
||||
|
||||
|
||||
def test_build_sso_user_update_data_keeps_existing_user_alias():
|
||||
"""
|
||||
An alias already stored for the user is never overwritten by the IdP display name.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.types import CustomOpenID
|
||||
from litellm.proxy.management_endpoints.ui_sso import _build_sso_user_update_data
|
||||
|
||||
sso_result = CustomOpenID(
|
||||
id="S-1-5-21-adfs-user",
|
||||
email="jane.doe@example.com",
|
||||
display_name="Doe, Jane",
|
||||
provider="generic",
|
||||
team_ids=[],
|
||||
)
|
||||
|
||||
update_data = _build_sso_user_update_data(
|
||||
result=sso_result,
|
||||
user_email="jane.doe@example.com",
|
||||
user_id="S-1-5-21-adfs-user",
|
||||
existing_user_alias="Admin-set alias",
|
||||
)
|
||||
|
||||
assert update_data == {"user_email": "jane.doe@example.com"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"result, expected_alias",
|
||||
[
|
||||
(
|
||||
CustomOpenID(id="user-1", display_name="Doe, Jane", first_name="Jane", last_name="Doe", team_ids=[]),
|
||||
"Doe, Jane",
|
||||
),
|
||||
(CustomOpenID(id="user-1", first_name="Jane", last_name="Doe", team_ids=[]), "Jane Doe"),
|
||||
(CustomOpenID(id="user-1", display_name="user-1", first_name="Jane", team_ids=[]), "Jane"),
|
||||
(CustomOpenID(id="user-1", display_name="user-1", team_ids=[]), None),
|
||||
(CustomOpenID(id="user-1", display_name=" ", first_name=" Jane ", last_name="Doe", team_ids=[]), "Jane Doe"),
|
||||
(CustomOpenID(id="user-1", display_name=" ", first_name=" ", team_ids=[]), None),
|
||||
({"id": "user-1", "display_name": "Dict User", "first_name": None, "last_name": None}, "Dict User"),
|
||||
(None, None),
|
||||
],
|
||||
)
|
||||
def test_get_sso_user_alias(result: CustomOpenID | dict[str, str | None] | None, expected_alias: str | None):
|
||||
"""
|
||||
The alias is the IdP display name unless it is just the user id, then the joined first/last name.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.ui_sso import _get_sso_user_alias
|
||||
|
||||
assert _get_sso_user_alias(result) == expected_alias
|
||||
|
||||
|
||||
def test_generic_response_convertor_normalizes_email():
|
||||
"""
|
||||
Test that generic_response_convertor normalizes email addresses.
|
||||
|
|
@ -1022,6 +1099,87 @@ async def test_upsert_sso_user_updates_role_for_existing_user():
|
|||
assert call_args.kwargs["data"]["user_role"] == "proxy_admin"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_sso_user_fills_user_alias_for_existing_user():
|
||||
"""
|
||||
An existing user row without an alias is updated with the SSO display name on login.
|
||||
"""
|
||||
from litellm.proxy._types import LiteLLM_UserTable
|
||||
from litellm.proxy.management_endpoints.types import CustomOpenID
|
||||
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_usertable.update_many = AsyncMock()
|
||||
|
||||
existing_user = LiteLLM_UserTable(
|
||||
user_id="S-1-5-21-adfs-user",
|
||||
user_email="jane.doe@example.com",
|
||||
user_role="internal_user",
|
||||
user_alias=None,
|
||||
)
|
||||
sso_result = CustomOpenID(
|
||||
id="S-1-5-21-adfs-user",
|
||||
email="jane.doe@example.com",
|
||||
first_name="Jane",
|
||||
last_name="Doe",
|
||||
display_name="Doe, Jane",
|
||||
provider="generic",
|
||||
team_ids=[],
|
||||
)
|
||||
|
||||
await SSOAuthenticationHandler.upsert_sso_user(
|
||||
result=sso_result,
|
||||
user_info=existing_user,
|
||||
user_email="jane.doe@example.com",
|
||||
user_defined_values=None,
|
||||
prisma_client=mock_prisma,
|
||||
)
|
||||
|
||||
mock_prisma.db.litellm_usertable.update_many.assert_called_once_with(
|
||||
where={"user_id": "S-1-5-21-adfs-user"},
|
||||
data={"user_email": "jane.doe@example.com", "user_alias": "Doe, Jane"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_insert_sso_user_sets_user_alias_from_display_name():
|
||||
"""
|
||||
A newly created SSO user is inserted with the IdP display name as user_alias.
|
||||
"""
|
||||
from litellm.proxy._types import NewUserResponse, SSOUserDefinedValues
|
||||
from litellm.proxy.management_endpoints.types import CustomOpenID
|
||||
from litellm.proxy.management_endpoints.ui_sso import insert_sso_user
|
||||
|
||||
sso_result = CustomOpenID(
|
||||
id="S-1-5-21-adfs-user",
|
||||
email="jane.doe@example.com",
|
||||
first_name="Jane",
|
||||
last_name="Doe",
|
||||
display_name="Doe, Jane",
|
||||
provider="generic",
|
||||
team_ids=[],
|
||||
)
|
||||
user_defined_values: SSOUserDefinedValues = {
|
||||
"models": [],
|
||||
"user_id": "S-1-5-21-adfs-user",
|
||||
"user_email": "jane.doe@example.com",
|
||||
"max_budget": None,
|
||||
"user_role": "internal_user",
|
||||
"budget_duration": None,
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.ui_sso.new_user",
|
||||
return_value=NewUserResponse(user_id="S-1-5-21-adfs-user", key="sk-xxxxx", teams=None),
|
||||
) as mock_new_user:
|
||||
await insert_sso_user(result_openid=sso_result, user_defined_values=user_defined_values)
|
||||
|
||||
new_user_request = mock_new_user.call_args.kwargs["data"]
|
||||
assert new_user_request.user_id == "S-1-5-21-adfs-user"
|
||||
assert new_user_request.user_email == "jane.doe@example.com"
|
||||
assert new_user_request.user_alias == "Doe, Jane"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_sso_user_does_not_update_invalid_role():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue