diff --git a/.github/screenshots/after_org_assigned.png b/.github/screenshots/after_org_assigned.png new file mode 100644 index 0000000000..75c6a8ed5f Binary files /dev/null and b/.github/screenshots/after_org_assigned.png differ diff --git a/.github/screenshots/after_org_detail.png b/.github/screenshots/after_org_detail.png new file mode 100644 index 0000000000..2b4d23e254 Binary files /dev/null and b/.github/screenshots/after_org_detail.png differ diff --git a/.github/screenshots/before_403_error.png b/.github/screenshots/before_403_error.png new file mode 100644 index 0000000000..686c2fe573 Binary files /dev/null and b/.github/screenshots/before_403_error.png differ diff --git a/.github/screenshots/before_no_org.png b/.github/screenshots/before_no_org.png new file mode 100644 index 0000000000..9a4cfaa1d8 Binary files /dev/null and b/.github/screenshots/before_no_org.png differ diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 8357b1c0fe..1ef40f0685 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -43,6 +43,7 @@ from litellm.proxy._types import ( LitellmUserRoles, Member, NewTeamRequest, + OrgMember, ProxyErrorTypes, ProxyException, SpecialManagementEndpointEnums, @@ -78,6 +79,9 @@ from litellm.proxy.management_endpoints.common_utils import ( _upsert_budget_and_membership, _user_has_admin_view, ) +from litellm.proxy.management_endpoints.organization_endpoints import ( + add_member_to_organization, +) from litellm.proxy.management_endpoints.tag_management_endpoints import ( get_daily_activity, ) @@ -1255,11 +1259,53 @@ async def _update_model_table( return _model_id +async def _auto_add_team_members_to_organization( + team: LiteLLM_TeamTable, + organization: LiteLLM_OrganizationTableWithMembers, + prisma_client: Any, +) -> None: + """ + When moving a team to an org, ensure all team members are also org members. + + For SSO/Entra setups without SCIM, users join teams automatically on login but + are never explicitly added to organizations. This silently upserts missing members + rather than blocking the team move. + """ + org_member_ids = ( + {m.user_id for m in organization.members} if organization.members else set() + ) + for member in team.members_with_roles: + if member.user_id is None: + continue + if member.user_id == SpecialProxyStrings.default_user_id.value: + continue + if member.user_id in org_member_ids: + continue + if organization.organization_id is None: + continue + try: + await add_member_to_organization( + member=OrgMember( + user_id=member.user_id, + role=LitellmUserRoles.INTERNAL_USER, + ), + organization_id=organization.organization_id, + prisma_client=prisma_client, + ) + except Exception as e: + verbose_proxy_logger.debug( + "_auto_add_team_members_to_organization: skipping user_id=%s - %s", + member.user_id, + e, + ) + + async def fetch_and_validate_organization( organization_id: str, existing_team_row: Any, llm_router: Optional[Router], prisma_client: Any, + user_api_key_dict: Optional[UserAPIKeyAuth] = None, ) -> Any: """ Fetch and validate an organization for team update operations. @@ -1294,14 +1340,25 @@ async def fetch_and_validate_organization( }, ) + is_proxy_admin = ( + user_api_key_dict is not None + and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + ) + organization = LiteLLM_OrganizationTableWithMembers(**organization_row.model_dump()) validate_team_org_change( team=LiteLLM_TeamTable(**existing_team_row.model_dump()), - organization=LiteLLM_OrganizationTableWithMembers( - **organization_row.model_dump() - ), + organization=organization, llm_router=llm_router, + is_proxy_admin=is_proxy_admin, ) + if is_proxy_admin: + await _auto_add_team_members_to_organization( + team=LiteLLM_TeamTable(**existing_team_row.model_dump()), + organization=organization, + prisma_client=prisma_client, + ) + return organization_row @@ -1309,14 +1366,20 @@ def validate_team_org_change( team: LiteLLM_TeamTable, organization: LiteLLM_OrganizationTableWithMembers, llm_router: Router, + is_proxy_admin: bool = False, ) -> bool: """ Validate that a team can be moved to an organization. - The org must have access to the team's models - The team budget cannot be greater than the org max_budget - - The team's user_id must be a member of the org + - For non-proxy-admins: all team members must already be org members - The team's tpm/rpm limit must be less than the org's tpm/rpm limit + + Proxy admins bypass the membership check and instead trigger auto-add of + missing members (handled by the caller). This supports SSO/Entra setups + where org membership tables are empty but proxy admins still need to group + teams under orgs for budget/model governance. """ # If the team's organization is the same as the new organization, return True @@ -1357,23 +1420,26 @@ def validate_team_org_change( }, ) - # Check if the team's user_id is a member of the org - team_members = [m.user_id for m in team.members_with_roles] - org_members = ( - [m.user_id for m in organization.members] if organization.members else [] - ) - not_in_org = [ - m - for m in team_members - if m not in org_members and m != SpecialProxyStrings.default_user_id.value - ] - if len(not_in_org) > 0: - raise HTTPException( - status_code=403, - detail={ - "error": f"Cannot move team to organization. Team has user_id {not_in_org} that is not a member of the organization." - }, + # For non-proxy-admins, require all team members to already be org members. + # This prevents a team admin from moving their team into an arbitrary org and + # thereby injecting members into that org without org admin approval. + if not is_proxy_admin: + team_members = [m.user_id for m in team.members_with_roles] + org_members = ( + [m.user_id for m in organization.members] if organization.members else [] ) + not_in_org = [ + m + for m in team_members + if m not in org_members and m != SpecialProxyStrings.default_user_id.value + ] + if len(not_in_org) > 0: + raise HTTPException( + status_code=403, + detail={ + "error": f"Cannot move team to organization. Team has user_id {not_in_org} that is not a member of the organization." + }, + ) # Check if the team's tpm/rpm limit is less than the org's tpm/rpm limit if ( @@ -1617,6 +1683,7 @@ async def update_team( # noqa: PLR0915 existing_team_row=existing_team_row, llm_router=llm_router, prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, ) elif data.organization_id is not None and len(data.organization_id) == 0: # unsetting the organization_id diff --git a/tests/test_litellm/proxy/test_team_org_move.py b/tests/test_litellm/proxy/test_team_org_move.py new file mode 100644 index 0000000000..2dc961bec8 --- /dev/null +++ b/tests/test_litellm/proxy/test_team_org_move.py @@ -0,0 +1,232 @@ +""" +Tests for moving teams to organizations. + +Covers the SSO/Entra scenario where: +- Proxy admins can move teams freely; missing members are auto-added to the org. +- Non-proxy-admins (team admins) must have all team members pre-added to the org, + preserving the original security model (no privilege escalation via team move). +""" +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy._types import ( + LiteLLM_OrganizationTableWithMembers, + LiteLLM_OrganizationMembershipTable, + LiteLLM_TeamTable, + LitellmUserRoles, + Member, + OrgMember, + SpecialProxyStrings, +) +from litellm.proxy.management_endpoints.team_endpoints import ( + _auto_add_team_members_to_organization, + validate_team_org_change, +) +from litellm.router import Router + + +def _make_org(organization_id="org-1", members=None, models=None): + from datetime import datetime + + return LiteLLM_OrganizationTableWithMembers( + organization_id=organization_id, + organization_alias="test-org", + budget_id="budget-test", + spend=0.0, + metadata={}, + models=models or [], + created_by="default_user_id", + updated_by="default_user_id", + members=members or [], + teams=[], + litellm_budget_table=None, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + +def _make_team(team_id="team-1", member_ids=None, organization_id=None): + members = [ + Member(user_id=uid, role="user") for uid in (member_ids or []) + ] + members.append(Member(user_id=SpecialProxyStrings.default_user_id.value, role="admin")) + return LiteLLM_TeamTable( + team_id=team_id, + team_alias="test-team", + organization_id=organization_id, + admins=[], + members=[], + members_with_roles=members, + metadata={}, + models=[], + blocked=False, + spend=0.0, + ) + + +def _make_org_membership(user_id): + from datetime import datetime + + return LiteLLM_OrganizationMembershipTable( + user_id=user_id, + organization_id="org-1", + user_role=LitellmUserRoles.INTERNAL_USER, + spend=0.0, + budget_id=None, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + +class TestValidateTeamOrgChange: + def test_proxy_admin_not_blocked_when_members_not_in_org(self): + """Proxy admins bypass the membership check — auto-add handles it instead.""" + router = MagicMock(spec=Router) + team = _make_team(member_ids=["sso-user-001", "sso-user-002"]) + org = _make_org(members=[]) + + result = validate_team_org_change( + team=team, organization=org, llm_router=router, is_proxy_admin=True + ) + assert result is True + + def test_non_admin_blocked_when_members_not_in_org(self): + """Team admins (non-proxy-admin) must have all members pre-added to the org.""" + router = MagicMock(spec=Router) + team = _make_team(member_ids=["sso-user-001"]) + org = _make_org(members=[]) + + with pytest.raises(Exception) as exc_info: + validate_team_org_change( + team=team, organization=org, llm_router=router, is_proxy_admin=False + ) + assert "403" in str(exc_info.value) or "not a member" in str(exc_info.value) + + def test_non_admin_passes_when_all_members_in_org(self): + """Team admin move succeeds when all team members are already org members.""" + router = MagicMock(spec=Router) + team = _make_team(member_ids=["u1"]) + org = _make_org(members=[_make_org_membership("u1")]) + + result = validate_team_org_change( + team=team, organization=org, llm_router=router, is_proxy_admin=False + ) + assert result is True + + def test_same_org_short_circuits(self): + """Moving to the same org is always a no-op, regardless of role.""" + router = MagicMock(spec=Router) + team = _make_team(member_ids=["u1"], organization_id="org-1") + org = _make_org(organization_id="org-1") + + assert validate_team_org_change( + team=team, organization=org, llm_router=router, is_proxy_admin=False + ) is True + assert validate_team_org_change( + team=team, organization=org, llm_router=router, is_proxy_admin=True + ) is True + + def test_default_user_excluded_from_membership_check(self): + """default_user_id is never checked for org membership.""" + router = MagicMock(spec=Router) + # Team has only default_user_id (added by _make_team) + team = _make_team(member_ids=[]) + org = _make_org(members=[]) + + # Should not raise even for non-proxy-admin + result = validate_team_org_change( + team=team, organization=org, llm_router=router, is_proxy_admin=False + ) + assert result is True + + +class TestAutoAddTeamMembersToOrg: + @pytest.mark.asyncio + async def test_adds_missing_members(self): + team = _make_team(member_ids=["sso-user-001", "sso-user-002"]) + org = _make_org(members=[]) + + mock_add = AsyncMock() + import litellm.proxy.management_endpoints.team_endpoints as te + original = te.add_member_to_organization + te.add_member_to_organization = mock_add + + try: + await _auto_add_team_members_to_organization( + team=team, + organization=org, + prisma_client=MagicMock(), + ) + finally: + te.add_member_to_organization = original + + assert mock_add.call_count == 2 + called_user_ids = { + call.kwargs["member"].user_id for call in mock_add.call_args_list + } + assert called_user_ids == {"sso-user-001", "sso-user-002"} + + @pytest.mark.asyncio + async def test_skips_existing_org_members(self): + team = _make_team(member_ids=["u1", "u2"]) + org = _make_org(members=[_make_org_membership("u1")]) + + mock_add = AsyncMock() + import litellm.proxy.management_endpoints.team_endpoints as te + original = te.add_member_to_organization + te.add_member_to_organization = mock_add + + try: + await _auto_add_team_members_to_organization( + team=team, + organization=org, + prisma_client=MagicMock(), + ) + finally: + te.add_member_to_organization = original + + assert mock_add.call_count == 1 + assert mock_add.call_args.kwargs["member"].user_id == "u2" + + @pytest.mark.asyncio + async def test_skips_default_user(self): + """default_user_id should never be added as an org member.""" + team = _make_team(member_ids=[]) + org = _make_org(members=[]) + + mock_add = AsyncMock() + import litellm.proxy.management_endpoints.team_endpoints as te + original = te.add_member_to_organization + te.add_member_to_organization = mock_add + + try: + await _auto_add_team_members_to_organization( + team=team, + organization=org, + prisma_client=MagicMock(), + ) + finally: + te.add_member_to_organization = original + + assert mock_add.call_count == 0 + + @pytest.mark.asyncio + async def test_logs_and_continues_on_error(self): + """Errors must not propagate — they are logged at DEBUG and skipped.""" + team = _make_team(member_ids=["u1"]) + org = _make_org(members=[]) + + mock_add = AsyncMock(side_effect=Exception("duplicate key")) + import litellm.proxy.management_endpoints.team_endpoints as te + original = te.add_member_to_organization + te.add_member_to_organization = mock_add + + try: + await _auto_add_team_members_to_organization( + team=team, + organization=org, + prisma_client=MagicMock(), + ) + finally: + te.add_member_to_organization = original