diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 15d05633679..8d0d05c3858 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -1628,6 +1628,19 @@ async def organization_member_delete( existing_user_email_row: Final = await find_member_if_email(data.user_email, prisma_client) data.user_id = existing_user_email_row.user_id + if data.user_id is None: + raise HTTPException( + status_code=400, + detail={"error": "Either user_id or user_email needs to be passed in"}, + ) + + await _remove_user_from_organization_teams( + user_id=data.user_id, + organization_id=data.organization_id, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) + member_to_delete: Final = await _table(OrganizationMembershipRepository(prisma_client)).delete( where={ "user_id_organization_id": { @@ -1643,6 +1656,31 @@ async def organization_member_delete( raise e +async def _remove_user_from_organization_teams( + user_id: str, + organization_id: str, + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient, +) -> None: + """Membership in an organization's teams is granted through the organization, so leaving the + organization has to take the user off every one of its teams, through the same per-team path as + /team/member_delete (advisory lock, roster rewrite, LiteLLM_TeamMembership and team-key cleanup, + cache eviction). Teams outside the organization and the user's personal keys are left alone.""" + from litellm.proxy.management_endpoints.team_endpoints import team_member_delete + + organization_teams: Final = await TeamRepository(prisma_client).find_by_organization_id(organization_id) + member_team_ids: Final = tuple( + team.team_id + for team in organization_teams + if any(member.user_id == user_id for member in team.members_with_roles) + ) + for team_id in member_team_ids: + await team_member_delete( + data=TeamMemberDeleteRequest(team_id=team_id, user_id=user_id), + user_api_key_dict=user_api_key_dict, + ) + + async def add_member_to_organization( member: OrgMember, organization_id: str, diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index 8d2c6c0054c..6cf45a6d047 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -2,7 +2,7 @@ import asyncio import json from contextlib import asynccontextmanager from types import MappingProxyType, SimpleNamespace -from typing import Final, Mapping, Optional, cast +from typing import TYPE_CHECKING, AsyncIterator, Final, Mapping, Optional, Sequence, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -15,6 +15,9 @@ from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import ( JWTMappingRow, ) +if TYPE_CHECKING: + from litellm.proxy._types import LiteLLM_VerificationToken + @pytest.mark.asyncio async def test_organization_update_object_permissions_existing_permission(monkeypatch): @@ -1718,3 +1721,197 @@ def test_organization_info_response_keeps_budget_reset_at(): ) assert response.model_dump()["litellm_budget_table"]["budget_reset_at"] == reset_at + + +class _Row(SimpleNamespace): + def model_dump(self) -> dict[str, object]: + return dict(vars(self)) + + +class _TeamRows: + def __init__(self, rows: Sequence[Mapping[str, object]]) -> None: + self.rows: dict[str, dict[str, object]] = {str(row["team_id"]): dict(row) for row in rows} + + async def find_many(self, where: Mapping[str, object]) -> list[_Row]: + return [_Row(**row) for row in self.rows.values() if row["organization_id"] == where["organization_id"]] + + async def find_unique(self, where: Mapping[str, str]) -> _Row | None: + row = self.rows.get(where["team_id"]) + return None if row is None else _Row(**row) + + async def update(self, where: Mapping[str, str], data: Mapping[str, str]) -> _Row: + self.rows[where["team_id"]]["members_with_roles"] = json.loads(data["members_with_roles"]) + return _Row(**self.rows[where["team_id"]]) + + +class _UserRows: + def __init__(self, rows: Sequence[Mapping[str, object]]) -> None: + self.rows: dict[str, dict[str, object]] = {str(row["user_id"]): dict(row) for row in rows} + + async def find_many(self, where: Mapping[str, Mapping[str, Sequence[str]]]) -> list[_Row]: + return [_Row(**row) for user_id, row in self.rows.items() if user_id in where["user_id"]["in"]] + + async def update(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, Sequence[str]]]) -> _Row: + self.rows[where["user_id"]]["teams"] = list(data["teams"]["set"]) + return _Row(**self.rows[where["user_id"]]) + + +class _TeamMembershipRows: + def __init__(self, rows: Sequence[tuple[str, str]]) -> None: + self.rows: set[tuple[str, str]] = set(rows) + + async def delete_many(self, where: Mapping[str, str]) -> int: + doomed = {(where["team_id"], where["user_id"])} & self.rows + self.rows -= doomed + return len(doomed) + + +class _KeyRows: + def __init__(self, rows: Sequence["LiteLLM_VerificationToken"]) -> None: + self.rows: dict[str, "LiteLLM_VerificationToken"] = {key.token: key for key in rows} + + def _matching(self, where: Mapping[str, object]) -> list["LiteLLM_VerificationToken"]: + user_filter = cast(Mapping[str, Sequence[str]], where["user_id"]) + return [ + key for key in self.rows.values() if key.user_id in user_filter["in"] and key.team_id == where["team_id"] + ] + + async def find_many(self, where: Mapping[str, object]) -> list["LiteLLM_VerificationToken"]: + return self._matching(where) + + async def delete_many(self, where: Mapping[str, object]) -> int: + doomed = self._matching(where) + for key in doomed: + del self.rows[key.token] + return len(doomed) + + +class _AppendOnlyRows: + def __init__(self) -> None: + self.records: list[Mapping[str, object]] = [] + + async def create_many(self, data: Sequence[Mapping[str, object]]) -> int: + self.records.extend(data) + return len(data) + + +class _NoRows: + async def find_many(self, where: Mapping[str, object] | None = None) -> list[object]: + return [] + + +class _OrgMembershipRows: + def __init__(self, rows: Sequence[tuple[str, str]]) -> None: + self.rows: set[tuple[str, str]] = set(rows) + + async def delete(self, where: Mapping[str, Mapping[str, str]]) -> _Row: + member = (where["user_id_organization_id"]["user_id"], where["user_id_organization_id"]["organization_id"]) + self.rows.remove(member) + return _Row(user_id=member[0], organization_id=member[1]) + + +class _OrgMemberDeleteTx: + def __init__(self, db: SimpleNamespace) -> None: + self._db = db + + def __getattr__(self, table_name: str) -> object: + return getattr(self._db, table_name) + + async def query_raw(self, sql: str, team_id: str) -> list[Mapping[str, object]]: + from litellm.proxy.management_endpoints.team_endpoints import TEAM_ADVISORY_LOCK_SQL + + if sql == TEAM_ADVISORY_LOCK_SQL: + return [] + team = self._db.litellm_teamtable.rows.get(team_id) + return [] if team is None else [{"members_with_roles": team["members_with_roles"]}] + + +class _OrgMemberDeletePrisma: + """In-memory stand-in for every table /organization/member_delete and the per-team cleanup it + cascades into read or write, so the test asserts the rows left behind rather than the calls made.""" + + def __init__( + self, + *, + teams: Sequence[Mapping[str, object]], + users: Sequence[Mapping[str, object]], + team_memberships: Sequence[tuple[str, str]], + keys: Sequence["LiteLLM_VerificationToken"], + org_memberships: Sequence[tuple[str, str]], + ) -> None: + self.db = SimpleNamespace( + litellm_teamtable=_TeamRows(teams), + litellm_usertable=_UserRows(users), + litellm_teammembership=_TeamMembershipRows(team_memberships), + litellm_verificationtoken=_KeyRows(keys), + litellm_deletedverificationtoken=_AppendOnlyRows(), + litellm_auditlog=_AppendOnlyRows(), + litellm_jwtkeymapping=_NoRows(), + litellm_organizationmembership=_OrgMembershipRows(org_memberships), + ) + + @asynccontextmanager + async def tx(self, timeout: object = None) -> AsyncIterator[_OrgMemberDeleteTx]: + yield _OrgMemberDeleteTx(self.db) + + +@pytest.mark.asyncio +async def test_organization_member_delete_removes_the_user_from_every_team_of_that_organization(monkeypatch): + """Team membership inside an organization is granted through the organization, so removing the + user from the organization has to take them off its teams the way /team/member_delete would, or + they keep their team admin role, their team keys and the ability to mint new ones (LIT-9248). + Teams in other organizations and the user's personal key are not the organization's to touch.""" + from litellm.proxy._types import ( + LiteLLM_VerificationToken, + LitellmUserRoles, + OrganizationMemberDeleteRequest, + UserAPIKeyAuth, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.organization_endpoints import organization_member_delete + + lead: Final = {"user_id": "lead", "user_email": None, "role": "admin"} + dev: Final = {"user_id": "dev", "user_email": None, "role": "admin"} + org_team_key: Final = LiteLLM_VerificationToken(token="hash-org-team-key", user_id="dev", team_id="team-eng") + other_org_team_key: Final = LiteLLM_VerificationToken(token="hash-other-team-key", user_id="dev", team_id="team-ext") + personal_key: Final = LiteLLM_VerificationToken(token="hash-personal-key", user_id="dev", team_id=None) + prisma: Final = _OrgMemberDeletePrisma( + teams=( + {"team_id": "team-eng", "organization_id": "org-eng", "members_with_roles": [lead, dev]}, + {"team_id": "team-eng-2", "organization_id": "org-eng", "members_with_roles": [lead]}, + {"team_id": "team-ext", "organization_id": "org-ext", "members_with_roles": [dev]}, + ), + users=( + {"user_id": "dev", "teams": ["team-eng", "team-ext"]}, + {"user_id": "lead", "teams": ["team-eng", "team-eng-2"]}, + ), + team_memberships=(("team-eng", "dev"), ("team-eng", "lead"), ("team-eng-2", "lead"), ("team-ext", "dev")), + keys=(org_team_key, other_org_team_key, personal_key), + org_memberships=(("dev", "org-eng"), ("lead", "org-eng"), ("dev", "org-ext")), + ) + cache: Final = UserApiKeyCache() + for cache_key in ("team_id:team-eng", "team_id:team-ext", org_team_key.token, personal_key.token): + cache.set_cache(key=cache_key, value={"cached": True}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None) + + removed: Final = await organization_member_delete( + data=OrganizationMemberDeleteRequest(organization_id="org-eng", user_id="dev"), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert (removed.user_id, removed.organization_id) == ("dev", "org-eng") + assert prisma.db.litellm_organizationmembership.rows == {("lead", "org-eng"), ("dev", "org-ext")} + assert prisma.db.litellm_teamtable.rows["team-eng"]["members_with_roles"] == [lead] + assert prisma.db.litellm_teamtable.rows["team-eng-2"]["members_with_roles"] == [lead] + assert prisma.db.litellm_teamtable.rows["team-ext"]["members_with_roles"] == [dev] + assert prisma.db.litellm_usertable.rows["dev"]["teams"] == ["team-ext"] + assert prisma.db.litellm_usertable.rows["lead"]["teams"] == ["team-eng", "team-eng-2"] + assert prisma.db.litellm_teammembership.rows == {("team-eng", "lead"), ("team-eng-2", "lead"), ("team-ext", "dev")} + assert set(prisma.db.litellm_verificationtoken.rows) == {other_org_team_key.token, personal_key.token} + assert [record["token"] for record in prisma.db.litellm_deletedverificationtoken.records] == [org_team_key.token] + assert cache.get_cache(key="team_id:team-eng") is None + assert cache.get_cache(key=org_team_key.token) is None + assert cache.get_cache(key="team_id:team-ext") == {"cached": True} + assert cache.get_cache(key=personal_key.token) == {"cached": True}