mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
fix(scim): use members_with_roles as the source of truth for group membership (#34162)
* fix(scim): use members_with_roles as the source of truth for group membership SCIM group provisioning tracked membership inconsistently. Team creation and the real team endpoints persist membership in members_with_roles (and each member's user.teams), but the SCIM group PATCH handler and the GET /Groups listing read the legacy team.members String[] column, which team creation never populates. Seeding a PATCH result from that empty column made an Okta "add member" operation recompute the member set from scratch and silently drop everyone already in the team, so users ended up missing from the groups they were provisioned into. Reading the same empty column on GET /Groups reported an empty member list back to the IdP, which drove repeated re-provisioning. Separately, add_new_member appended the team id to user.teams with an unconditional array push. Under the concurrent group PATCHes an IdP sends during a reconcile, each request passed the members_with_roles duplicate check and pushed, so user.teams accumulated duplicate ids for the same team. A duplicate also breaks auth logic that keys off the number of teams a user belongs to. Read current membership from members_with_roles in the SCIM group PATCH seed and the GET /Groups listing, and make the user.teams append idempotent via a filtered update that no-ops once the team is present. Resolves LIT-4283 * fix(scim): address review; atomic user-creation and stop writing legacy members Keep the concurrent-safe team append but create the user via an atomic upsert (create-or-update) instead of a check-then-create, so provisioning the same new user concurrently cannot race into a duplicate-key failure; the team is still appended idempotently by a filtered update so an existing user gets no duplicate team id. Stop writing the legacy team.members column in the group PATCH apply so the only membership record is the source of truth (members_with_roles plus each member's user.teams), reconciled by team_member_add/team_member_delete. Tests: existing add_new_member and team-creation mocks updated to the upsert plus filtered-append shape, and new tests cover atomic creation and that the PATCH apply does not write the legacy members column.
This commit is contained in:
Yassin Kortam
GitHub
parent
0fcaadf11c
commit
38467631b6
5 changed files with 347 additions and 37 deletions
|
|
@ -1654,8 +1654,11 @@ async def get_groups(
|
|||
# Convert to SCIM format
|
||||
scim_groups = []
|
||||
for team in teams:
|
||||
# Get team members with display names
|
||||
members = await _get_team_members_display(team.members or [])
|
||||
# Get team members with display names. members_with_roles is the
|
||||
# source of truth; the legacy `members` column is not populated by
|
||||
# team creation, so reading it here would report an empty member
|
||||
# list to the IdP and trigger repeated re-provisioning.
|
||||
members = await _get_team_members_display(await _get_team_member_user_ids_from_team(team))
|
||||
verbose_proxy_logger.debug(f"SCIM GET GROUPS members: {members}")
|
||||
team_alias = getattr(team, "team_alias", team.team_id)
|
||||
team_created_at = team.created_at.isoformat() if team.created_at else None
|
||||
|
|
@ -1885,8 +1888,12 @@ async def _process_group_patch_operations(
|
|||
existing_metadata = existing_team.metadata or {}
|
||||
metadata = dict(existing_metadata) if existing_metadata else {}
|
||||
|
||||
# Track member changes
|
||||
current_members = set(existing_team.members or [])
|
||||
# Track member changes. members_with_roles is the source of truth for team
|
||||
# membership; the legacy `members` column is not populated by team creation
|
||||
# or the real team endpoints, so seeding from it would make an `add`/`remove`
|
||||
# operation recompute the member set from an empty base and silently drop
|
||||
# everyone already in the team.
|
||||
current_members = set(await _get_team_member_user_ids_from_team(existing_team))
|
||||
final_members = current_members.copy()
|
||||
|
||||
# Process each patch operation
|
||||
|
|
@ -1963,24 +1970,24 @@ async def _process_group_patch_operations(
|
|||
return update_data, final_members
|
||||
|
||||
|
||||
async def _apply_group_patch_updates(
|
||||
group_id: str, update_data: Dict[str, Any], final_members: Set[str], prisma_client
|
||||
):
|
||||
"""Apply patch updates to the group in the database."""
|
||||
# Serialize metadata if present
|
||||
async def _apply_group_patch_updates(group_id: str, update_data: Dict[str, Any], prisma_client):
|
||||
"""Apply the group's metadata/displayName patch updates to the database.
|
||||
|
||||
Membership itself is not written here; it is reconciled onto the source of
|
||||
truth (members_with_roles and each member's user.teams) by
|
||||
_handle_group_membership_changes via team_member_add/team_member_delete.
|
||||
Writing the legacy `members` column here too would create a second, unread
|
||||
copy of membership that could drift from the source of truth.
|
||||
"""
|
||||
if "metadata" in update_data and isinstance(update_data["metadata"], dict):
|
||||
update_data["metadata"] = safe_dumps(update_data["metadata"])
|
||||
|
||||
# Update members list
|
||||
update_data["members"] = list(final_members)
|
||||
|
||||
# Update team in database
|
||||
updated_team = await TeamRepository(prisma_client).table.update(
|
||||
where={"team_id": group_id},
|
||||
data=update_data,
|
||||
)
|
||||
|
||||
return updated_team
|
||||
if update_data:
|
||||
return await TeamRepository(prisma_client).table.update(
|
||||
where={"team_id": group_id},
|
||||
data=update_data,
|
||||
)
|
||||
return await TeamRepository(prisma_client).table.find_unique(where={"team_id": group_id})
|
||||
|
||||
|
||||
async def _handle_group_membership_changes(group_id: str, current_members: Set[str], final_members: Set[str]):
|
||||
|
|
@ -2036,8 +2043,8 @@ async def patch_group(
|
|||
# Track current members BEFORE update for comparison
|
||||
current_members = set(await _get_team_member_user_ids_from_team(existing_team))
|
||||
|
||||
# Apply updates to the database
|
||||
updated_team = await _apply_group_patch_updates(group_id, update_data, final_members, prisma_client)
|
||||
# Apply the metadata/displayName updates to the database
|
||||
updated_team = await _apply_group_patch_updates(group_id, update_data, prisma_client)
|
||||
|
||||
# Refresh team data from database to get the latest state after concurrent updates
|
||||
# This prevents race conditions when multiple PATCH requests come in simultaneously
|
||||
|
|
|
|||
|
|
@ -252,6 +252,21 @@ async def _resolve_member_budget_id(
|
|||
return response.budget_id
|
||||
|
||||
|
||||
async def _append_team_id_if_absent(prisma_client: PrismaClient, user_id: str, team_id: str) -> None:
|
||||
"""Append team_id to a user's teams array, only if it is not already present.
|
||||
|
||||
The row-level filter makes the append a no-op once the team is present, so
|
||||
repeated or concurrent adds of the same team cannot accumulate duplicate
|
||||
team ids in user.teams (a duplicate also breaks auth logic that keys off the
|
||||
number of teams a user belongs to). Teams added concurrently for a different
|
||||
team id are unaffected, since each update filters on its own team id.
|
||||
"""
|
||||
await UserRepository(prisma_client).table.update_many(
|
||||
where={"user_id": user_id, "NOT": {"teams": {"has": team_id}}},
|
||||
data={"teams": {"push": [team_id]}},
|
||||
)
|
||||
|
||||
|
||||
async def add_new_member(
|
||||
new_member: Member,
|
||||
max_budget_in_team: Optional[float],
|
||||
|
|
@ -276,13 +291,16 @@ async def add_new_member(
|
|||
## ADD TEAM ID, to USER TABLE IF NEW ##
|
||||
if new_member.user_id is not None:
|
||||
new_user_defaults = get_new_internal_user_defaults(user_id=new_member.user_id)
|
||||
# Upsert ensures the user row exists atomically (no create race when the
|
||||
# same new user is provisioned concurrently), seeding teams on create.
|
||||
# The teams append lives in the filtered update below rather than the
|
||||
# upsert's update branch so an already-existing user does not get a
|
||||
# duplicate team id.
|
||||
_returned_user = await UserRepository(prisma_client).table.upsert(
|
||||
where={"user_id": new_member.user_id},
|
||||
data={
|
||||
"update": {"teams": {"push": [team_id]}},
|
||||
"create": {"teams": [team_id], **new_user_defaults}, # type: ignore
|
||||
},
|
||||
data={"create": {"teams": [team_id], **new_user_defaults}, "update": {}},
|
||||
)
|
||||
await _append_team_id_if_absent(prisma_client, new_member.user_id, team_id)
|
||||
if _returned_user is not None:
|
||||
returned_user = LiteLLM_UserTable(**_returned_user.model_dump())
|
||||
elif new_member.user_email is not None:
|
||||
|
|
@ -302,12 +320,8 @@ async def add_new_member(
|
|||
returned_user = LiteLLM_UserTable(**_returned_user.model_dump())
|
||||
elif len(existing_user_row) == 1:
|
||||
user_info = existing_user_row[0]
|
||||
_returned_user = await UserRepository(prisma_client).table.update(
|
||||
where={"user_id": user_info.user_id}, # type: ignore
|
||||
data={"teams": {"push": [team_id]}},
|
||||
)
|
||||
if _returned_user is not None:
|
||||
returned_user = LiteLLM_UserTable(**_returned_user.model_dump())
|
||||
await _append_team_id_if_absent(prisma_client, user_info.user_id, team_id)
|
||||
returned_user = LiteLLM_UserTable(**user_info.model_dump())
|
||||
elif len(existing_user_row) > 1:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
|
|
|
|||
|
|
@ -4,14 +4,17 @@ import pytest
|
|||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_TeamTable,
|
||||
LiteLLM_UserTable,
|
||||
LitellmUserRoles,
|
||||
Member,
|
||||
NewUserRequest,
|
||||
NewUserResponse,
|
||||
ProxyException,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.scim.scim_v2 import (
|
||||
UserProvisionerHelpers,
|
||||
_apply_group_patch_updates,
|
||||
_extract_group_member_ids,
|
||||
_handle_team_membership_changes,
|
||||
_process_group_patch_operations,
|
||||
|
|
@ -19,6 +22,7 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import (
|
|||
create_group,
|
||||
create_user,
|
||||
delete_group,
|
||||
get_groups,
|
||||
get_users,
|
||||
get_service_provider_config,
|
||||
patch_group,
|
||||
|
|
@ -2855,3 +2859,161 @@ async def test_patch_group_rename_recomputes_retained_members(mocker):
|
|||
|
||||
recompute_mock.assert_awaited_once()
|
||||
assert set(recompute_mock.call_args[0][1]) == {"user1"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_group_patch_operations_add_retains_existing_members(
|
||||
mocker, monkeypatch
|
||||
):
|
||||
"""A SCIM group ``add`` operation must not drop members already in the team.
|
||||
|
||||
Team membership lives in members_with_roles; team creation leaves the legacy
|
||||
``members`` column empty. Seeding the patch result from that empty column
|
||||
made an ``add`` recompute the member set from scratch and remove everyone
|
||||
already in the team. The result set must be seeded from members_with_roles so
|
||||
existing members survive an add of a new one.
|
||||
"""
|
||||
|
||||
async def mock_get_config():
|
||||
return {"litellm_settings": {"scim_upsert_user": True}}
|
||||
|
||||
from litellm.proxy.proxy_server import proxy_config
|
||||
|
||||
monkeypatch.setattr(proxy_config, "get_config", mock_get_config)
|
||||
|
||||
existing_team = LiteLLM_TeamTable(
|
||||
team_id="team-1",
|
||||
team_alias="Team One",
|
||||
members=[], # legacy column intentionally empty, as real teams leave it
|
||||
members_with_roles=[Member(user_id="existing-user", role="user")],
|
||||
)
|
||||
patch_ops = SCIMPatchOp(
|
||||
schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"],
|
||||
Operations=[
|
||||
SCIMPatchOperation(op="add", path="members", value=[{"value": "new-user"}])
|
||||
],
|
||||
)
|
||||
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
mock_prisma_client.db = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
|
||||
# new-user already exists in the DB
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=mocker.MagicMock(user_id="new-user")
|
||||
)
|
||||
|
||||
_, final_members = await _process_group_patch_operations(
|
||||
patch_ops=patch_ops,
|
||||
existing_team=existing_team,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
assert final_members == {"existing-user", "new-user"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_group_patch_operations_remove_uses_members_with_roles(
|
||||
mocker, monkeypatch
|
||||
):
|
||||
"""A ``remove`` op must diff against members_with_roles, so removing one
|
||||
member leaves the rest of the team intact rather than emptying it."""
|
||||
|
||||
async def mock_get_config():
|
||||
return {"litellm_settings": {"scim_upsert_user": True}}
|
||||
|
||||
from litellm.proxy.proxy_server import proxy_config
|
||||
|
||||
monkeypatch.setattr(proxy_config, "get_config", mock_get_config)
|
||||
|
||||
existing_team = LiteLLM_TeamTable(
|
||||
team_id="team-1",
|
||||
team_alias="Team One",
|
||||
members=[],
|
||||
members_with_roles=[
|
||||
Member(user_id="keep-user", role="user"),
|
||||
Member(user_id="drop-user", role="user"),
|
||||
],
|
||||
)
|
||||
patch_ops = SCIMPatchOp(
|
||||
schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"],
|
||||
Operations=[
|
||||
SCIMPatchOperation(
|
||||
op="remove", path="members", value=[{"value": "drop-user"}]
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
mock_prisma_client.db = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=mocker.MagicMock(user_id="drop-user")
|
||||
)
|
||||
|
||||
_, final_members = await _process_group_patch_operations(
|
||||
patch_ops=patch_ops,
|
||||
existing_team=existing_team,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
assert final_members == {"keep-user"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_groups_reports_members_from_members_with_roles(mocker):
|
||||
"""GET /Groups must report members from members_with_roles (the source of
|
||||
truth), not the legacy ``members`` column that team creation leaves empty.
|
||||
Reporting an empty member list makes the IdP repeatedly re-provision."""
|
||||
team = LiteLLM_TeamTable(
|
||||
team_id="team-1",
|
||||
team_alias="Team One",
|
||||
members=[], # legacy column empty
|
||||
members_with_roles=[Member(user_id="member-1", role="user")],
|
||||
)
|
||||
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
mock_prisma_client.db = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[team])
|
||||
mock_prisma_client.db.litellm_teamtable.count = AsyncMock(return_value=1)
|
||||
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=mocker.MagicMock(user_id="member-1", user_email="member-1@example.com")
|
||||
)
|
||||
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
|
||||
AsyncMock(return_value=mock_prisma_client),
|
||||
)
|
||||
|
||||
response = await get_groups(startIndex=1, count=10, filter=None)
|
||||
|
||||
assert [m.value for m in response.Resources[0].members] == ["member-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_group_patch_updates_does_not_write_legacy_members(mocker):
|
||||
"""The group PATCH apply must not write the legacy ``members`` column.
|
||||
|
||||
Membership is reconciled onto the source of truth (members_with_roles and
|
||||
each member's user.teams) separately; writing the legacy column here too
|
||||
would create a second, unread copy of membership that can drift from the
|
||||
source of truth, which is the inconsistency this PR removes.
|
||||
"""
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
mock_prisma_client.db = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
|
||||
updated = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=updated)
|
||||
|
||||
result = await _apply_group_patch_updates(
|
||||
group_id="team-1",
|
||||
update_data={"team_alias": "Renamed"},
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
assert result is updated
|
||||
mock_prisma_client.db.litellm_teamtable.update.assert_awaited_once()
|
||||
written = mock_prisma_client.db.litellm_teamtable.update.call_args.kwargs["data"]
|
||||
assert "members" not in written
|
||||
assert written["team_alias"] == "Renamed"
|
||||
|
|
|
|||
|
|
@ -4106,6 +4106,8 @@ async def test_new_team_max_budget_within_user_limit():
|
|||
}
|
||||
mock_prisma.db.litellm_usertable = MagicMock()
|
||||
mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user)
|
||||
mock_prisma.db.litellm_usertable.update_many = AsyncMock()
|
||||
mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user)
|
||||
mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user)
|
||||
|
||||
# Mock team membership table
|
||||
|
|
@ -4247,6 +4249,8 @@ async def test_new_team_org_scoped_budget_bypasses_user_limit():
|
|||
}
|
||||
mock_prisma.db.litellm_usertable = MagicMock()
|
||||
mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user)
|
||||
mock_prisma.db.litellm_usertable.update_many = AsyncMock()
|
||||
mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user)
|
||||
mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user)
|
||||
|
||||
# Mock team membership table
|
||||
|
|
@ -4393,6 +4397,8 @@ async def test_new_team_org_scoped_models_bypasses_user_limit():
|
|||
}
|
||||
mock_prisma.db.litellm_usertable = MagicMock()
|
||||
mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user)
|
||||
mock_prisma.db.litellm_usertable.update_many = AsyncMock()
|
||||
mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user)
|
||||
mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user)
|
||||
|
||||
# Mock team membership table
|
||||
|
|
@ -7245,6 +7251,8 @@ async def test_new_team_soft_budget_validation(
|
|||
}
|
||||
mock_prisma.db.litellm_usertable = MagicMock()
|
||||
mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user)
|
||||
mock_prisma.db.litellm_usertable.update_many = AsyncMock()
|
||||
mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user)
|
||||
mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user)
|
||||
|
||||
# Mock team membership table
|
||||
|
|
|
|||
|
|
@ -202,7 +202,8 @@ async def test_add_new_member_clones_default_team_budget_id():
|
|||
"teams": [test_team_id],
|
||||
"user_role": "internal_user",
|
||||
}
|
||||
mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(
|
||||
mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response)
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=mock_user_response
|
||||
)
|
||||
|
||||
|
|
@ -305,7 +306,8 @@ async def test_add_new_member_budget_duration_only_clones_default_max_budget():
|
|||
"teams": ["team-dc"],
|
||||
"user_role": "internal_user",
|
||||
}
|
||||
mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(
|
||||
mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response)
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=mock_user_response
|
||||
)
|
||||
mock_default_budget_row = MagicMock()
|
||||
|
|
@ -388,7 +390,8 @@ async def test_add_new_member_no_budget_when_no_default_and_no_max_budget():
|
|||
"teams": [test_team_id],
|
||||
"user_role": "internal_user",
|
||||
}
|
||||
mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(
|
||||
mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response)
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=mock_user_response
|
||||
)
|
||||
|
||||
|
|
@ -455,7 +458,8 @@ async def test_add_new_member_creates_new_budget_when_max_budget_provided():
|
|||
"teams": [test_team_id],
|
||||
"user_role": "internal_user",
|
||||
}
|
||||
mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(
|
||||
mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response)
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=mock_user_response
|
||||
)
|
||||
|
||||
|
|
@ -531,7 +535,8 @@ async def test_add_new_member_persists_budget_duration():
|
|||
"teams": ["team-dur"],
|
||||
"user_role": "internal_user",
|
||||
}
|
||||
mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(
|
||||
mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response)
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=mock_user_response
|
||||
)
|
||||
mock_budget_response = MagicMock()
|
||||
|
|
@ -594,7 +599,8 @@ async def test_add_new_member_persists_budget_duration_without_max_budget():
|
|||
"teams": ["team-dur2"],
|
||||
"user_role": "internal_user",
|
||||
}
|
||||
mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(
|
||||
mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response)
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=mock_user_response
|
||||
)
|
||||
mock_budget_response = MagicMock()
|
||||
|
|
@ -997,3 +1003,116 @@ async def test_attach_object_permission_to_dict_with_none_object_permission_id()
|
|||
|
||||
# Verify no database query was made
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_new_member_appends_team_only_if_absent_for_existing_user():
|
||||
"""Adding an existing user to a team must append the team id only if it is
|
||||
not already present.
|
||||
|
||||
add_new_member is the single writer of user.teams for every team add
|
||||
(/team/member_add, /user/new, SSO, SCIM). An unconditional append let
|
||||
repeated or concurrent adds accumulate duplicate team ids in user.teams,
|
||||
which also breaks auth logic that keys off the number of teams a user
|
||||
belongs to. The append must go through a filtered update that no-ops when
|
||||
the team is already present, and it must not fall through to creating a new
|
||||
user row for a user that already exists.
|
||||
"""
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
new_member = Member(user_id="existing-user", role="user")
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
mock_user_after = MagicMock()
|
||||
mock_user_after.model_dump.return_value = {
|
||||
"user_id": "existing-user",
|
||||
"user_email": None,
|
||||
"teams": ["team-1"],
|
||||
"user_role": "internal_user",
|
||||
}
|
||||
mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_after)
|
||||
mock_prisma_client.db.litellm_usertable.update_many = AsyncMock()
|
||||
# no team default budget and no explicit budget -> no team membership row
|
||||
mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
result_user, _ = await add_new_member(
|
||||
new_member=new_member,
|
||||
max_budget_in_team=None,
|
||||
prisma_client=mock_prisma_client,
|
||||
team_id="team-1",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name="admin",
|
||||
)
|
||||
|
||||
assert result_user is not None
|
||||
assert result_user.user_id == "existing-user"
|
||||
|
||||
# the append must be a filtered, idempotent update keyed off the team id, so
|
||||
# a repeated or concurrent add of a team the user already has is a no-op
|
||||
mock_prisma_client.db.litellm_usertable.update_many.assert_called_once()
|
||||
where = mock_prisma_client.db.litellm_usertable.update_many.call_args.kwargs["where"]
|
||||
assert where["user_id"] == "existing-user"
|
||||
assert where["NOT"] == {"teams": {"has": "team-1"}}
|
||||
data = mock_prisma_client.db.litellm_usertable.update_many.call_args.kwargs["data"]
|
||||
assert data == {"teams": {"push": ["team-1"]}}
|
||||
|
||||
# upsert (not an unconditional teams push) is what ensures the row exists, so
|
||||
# its update branch must not carry a teams push that would duplicate
|
||||
mock_prisma_client.db.litellm_usertable.upsert.assert_called_once()
|
||||
upsert_update = mock_prisma_client.db.litellm_usertable.upsert.call_args.kwargs["data"]["update"]
|
||||
assert "teams" not in upsert_update
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_new_member_creates_missing_user_atomically_via_upsert():
|
||||
"""A brand-new user added to a team must be created via an atomic upsert, not
|
||||
a separate existence check followed by create.
|
||||
|
||||
Concurrent provisioning of the same new user (which SCIM group reconciles do)
|
||||
would race a check-then-create into a duplicate-key failure. The upsert seeds
|
||||
teams on create, and the filtered append is a no-op because the team is
|
||||
already present on the freshly created row.
|
||||
"""
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
new_member = Member(user_id="brand-new-user", role="user")
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
mock_created = MagicMock()
|
||||
mock_created.model_dump.return_value = {
|
||||
"user_id": "brand-new-user",
|
||||
"user_email": None,
|
||||
"teams": ["team-1"],
|
||||
"user_role": "internal_user",
|
||||
}
|
||||
mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_created)
|
||||
mock_prisma_client.db.litellm_usertable.update_many = AsyncMock()
|
||||
mock_prisma_client.db.litellm_usertable.create = AsyncMock()
|
||||
mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
result_user, _ = await add_new_member(
|
||||
new_member=new_member,
|
||||
max_budget_in_team=None,
|
||||
prisma_client=mock_prisma_client,
|
||||
team_id="team-1",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name="admin",
|
||||
)
|
||||
|
||||
assert result_user is not None
|
||||
assert result_user.user_id == "brand-new-user"
|
||||
|
||||
# existence is established by an atomic upsert (create-or-update), never a
|
||||
# non-atomic standalone create that could race under concurrent provisioning
|
||||
mock_prisma_client.db.litellm_usertable.upsert.assert_called_once()
|
||||
mock_prisma_client.db.litellm_usertable.create.assert_not_called()
|
||||
create_data = mock_prisma_client.db.litellm_usertable.upsert.call_args.kwargs["data"]["create"]
|
||||
assert create_data["teams"] == ["team-1"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue