diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 8a1b9c112dc..e9a4fd718f3 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -123,6 +123,8 @@ class _ObjectPermissionRow(Protocol): class _UserTableClient(Protocol): async def find_unique(self, where: Mapping[str, object]) -> "PrismaUserTable | None": ... + async def find_many(self, where: Mapping[str, object], take: int | None = None) -> "Sequence[PrismaUserTable]": ... + class _BudgetTableClient(Protocol): async def create(self, data: Mapping[str, object]) -> "PrismaBudgetTable": ... @@ -1413,16 +1415,12 @@ async def find_member_if_email(user_email: str, prisma_client: PrismaClient) -> "error": f"Unique user not found for user_email={user_email}. Potential duplicate OR non-existent user_email in LiteLLM_UserTable. Use 'user_id' instead." }, ) - try: - existing_user_email_row: Final = await UserRepository(prisma_client).table.find_unique( - where={"user_email": user_email} - ) - except Exception: + existing_user_email_rows: Final = await _table(UserRepository(prisma_client)).find_many( + where={"user_email": user_email}, take=2 + ) + if len(existing_user_email_rows) != 1: raise not_unique_user_email_error - if existing_user_email_row is None: - raise not_unique_user_email_error - existing_user_email_row_pydantic: Final = LiteLLM_UserTable.model_validate(existing_user_email_row.model_dump()) - return existing_user_email_row_pydantic + return LiteLLM_UserTable.model_validate(existing_user_email_rows[0].model_dump()) @router.patch( @@ -1642,25 +1640,20 @@ async def add_member_to_organization( try: user_object: LiteLLM_UserTable | None = None existing_user_id_row = None - existing_user_email_row = None ## Check if user exists in LiteLLM_UserTable - user exists - either the user_id or user_email is in LiteLLM_UserTable if member.user_id is not None: existing_user_id_row = await _table(UserRepository(prisma_client)).find_unique( where={"user_id": member.user_id} ) - if existing_user_id_row is None and member.user_email is not None: - try: - existing_user_email_row = await UserRepository(prisma_client).table.find_unique( - where={"user_email": member.user_email} - ) - except Exception as e: - raise ValueError( - f"Potential NON-Existent or Duplicate user email in DB: Error finding a unique instance of user_email={member.user_email} in LiteLLM_UserTable.: {e}" - ) + existing_user_email_rows: Final = ( + await _table(UserRepository(prisma_client)).find_many(where={"user_email": member.user_email}, take=2) + if existing_user_id_row is None and member.user_email is not None + else () + ) ## If user does not exist, create a new user - if existing_user_id_row is None and existing_user_email_row is None: + if existing_user_id_row is None and not existing_user_email_rows: # Create a new user - since user does not exist user_id: Final[str] = member.user_id or str(uuid.uuid4()) new_user_defaults: Final = get_new_internal_user_defaults( @@ -1671,16 +1664,20 @@ async def add_member_to_organization( _returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user") if _returned_user is not None: user_object = LiteLLM_UserTable.model_validate(_returned_user.model_dump()) - elif existing_user_email_row is not None and ( - len(existing_user_email_row) # pyright: ignore[reportArgumentType] # find_unique yields a row, not a list - > 1 - ): + elif len(existing_user_email_rows) > 1: raise HTTPException( status_code=400, detail={"error": "Multiple users with this email found in db. Please use 'user_id' instead."}, ) - elif existing_user_email_row is not None: - user_object = LiteLLM_UserTable.model_validate(existing_user_email_row.model_dump()) + elif existing_user_email_rows and member.user_id is not None: + raise HTTPException( + status_code=400, + detail={ + "error": f"user_email '{member.user_email}' and user_id '{member.user_id}' do not belong to the same user." + }, + ) + elif existing_user_email_rows: + user_object = LiteLLM_UserTable.model_validate(existing_user_email_rows[0].model_dump()) elif existing_user_id_row is not None: user_object = LiteLLM_UserTable.model_validate(existing_user_id_row.model_dump()) else: @@ -1707,5 +1704,7 @@ async def add_member_to_organization( ) return user_object, organization_membership + except HTTPException: + raise except Exception as e: raise ValueError(f"Error adding member={member} to organization={organization_id}: {e}") diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index 5e06756f3a9..0d526aefb06 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -1575,6 +1575,158 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke assert jwt_table.rows == (kept_row,) +class _PrismaLikeUserTable: + """Rejects find_unique on user_email like Prisma does, because the column is not @unique""" + + def __init__(self, rows: tuple) -> None: + self.rows = rows + self.most_rows_read = 0 + + async def find_unique(self, where: Mapping[str, object]): + if set(where) != {"user_id"}: + raise ValueError(f"Could not find field at `findUniqueLiteLLM_UserTable.where.{next(iter(where))}`") + return next((row for row in self.rows if row.user_id == where["user_id"]), None) + + async def find_many(self, where: Mapping[str, object], take: int | None = None): + matches = [row for row in self.rows if all(getattr(row, k) == v for k, v in where.items())] + read = matches if take is None else matches[:take] + self.most_rows_read = max(self.most_rows_read, len(read)) + return read + + +def _org_prisma_with_users(monkeypatch, users: tuple): + from datetime import datetime + + from litellm.proxy._types import LiteLLM_OrganizationMembershipTable, LiteLLM_UserTable + + async def create_membership(data): + return LiteLLM_OrganizationMembershipTable( + user_id=data["user_id"], + organization_id=data["organization_id"], + user_role=data["user_role"], + created_at=datetime(2024, 1, 1), + updated_at=datetime(2024, 1, 1), + ) + + async def insert_user(data, table_name): + return LiteLLM_UserTable(user_id=data["user_id"], user_email=data.get("user_email"), user_role="internal_user") + + mock_prisma = SimpleNamespace( + db=SimpleNamespace( + litellm_organizationtable=SimpleNamespace(find_unique=AsyncMock(return_value=SimpleNamespace())), + litellm_usertable=_PrismaLikeUserTable(users), + litellm_organizationmembership=SimpleNamespace(create=create_membership), + ), + insert_data=insert_user, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.organization_endpoints._verify_org_access", + AsyncMock(), + ) + return mock_prisma + + +async def _add_org_member(member: dict): + from litellm.proxy._types import LitellmUserRoles, OrganizationMemberAddRequest, UserAPIKeyAuth + from litellm.proxy.management_endpoints.organization_endpoints import organization_member_add + + return await organization_member_add( + data=OrganizationMemberAddRequest(organization_id="org-1", member=member), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "member, expected_user_id, expected_email", + [ + ({"user_email": "alice@example.com", "role": "internal_user"}, "alice-id", "alice@example.com"), + ({"user_id": "new-id", "user_email": "bob@example.com", "role": "org_admin"}, "new-id", "bob@example.com"), + ({"user_email": "bob@example.com", "role": "internal_user"}, None, "bob@example.com"), + ], + ids=["email-of-existing-user", "unknown-id-with-new-email", "new-email-only"], +) +async def test_organization_member_add_resolves_member_by_email(member, expected_user_id, expected_email, monkeypatch): + from litellm.proxy._types import LiteLLM_UserTable + + alice: Final = LiteLLM_UserTable(user_id="alice-id", user_email="alice@example.com", user_role="internal_user") + _org_prisma_with_users(monkeypatch, (alice,)) + + response = await _add_org_member(member) + + user = response.updated_users[0] + membership = response.updated_organization_memberships[0] + assert user.user_email == expected_email, response + assert membership.user_id == user.user_id, response + assert membership.user_role == member["role"], response + if expected_user_id is None: + assert user.user_id != alice.user_id, response + else: + assert user.user_id == expected_user_id, response + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "member, expected_error", + [ + ({"user_email": "dup@example.com", "role": "internal_user"}, "Multiple users with this email"), + ( + {"user_id": "new-id", "user_email": "alice@example.com", "role": "internal_user"}, + "user_email 'alice@example.com' and user_id 'new-id' do not belong to the same user", + ), + ], + ids=["email-shared-by-several-users", "unknown-id-with-another-users-email"], +) +async def test_organization_member_add_rejects_an_ambiguous_member_with_400(member, expected_error, monkeypatch): + from litellm.proxy._types import LiteLLM_UserTable, ProxyException + + users: Final = ( + LiteLLM_UserTable(user_id="alice-id", user_email="alice@example.com", user_role="internal_user"), + *( + LiteLLM_UserTable(user_id=user_id, user_email="dup@example.com", user_role="internal_user") + for user_id in ("dup-1", "dup-2", "dup-3") + ), + ) + prisma = _org_prisma_with_users(monkeypatch, users) + + with pytest.raises(ProxyException) as exc: + await _add_org_member(member) + assert int(exc.value.code) == 400, exc.value.message + assert expected_error in str(exc.value.message) + assert prisma.db.litellm_usertable.most_rows_read <= 2, "an email lookup must not load every duplicate row" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "rows, expected_user_id", + [ + ((("alice-id", "alice@example.com"),), "alice-id"), + ((), None), + ((("a", "alice@example.com"), ("b", "alice@example.com")), None), + ], + ids=["unique", "missing", "duplicate"], +) +async def test_find_member_if_email_resolves_only_a_unique_email(rows, expected_user_id, monkeypatch): + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.management_endpoints.organization_endpoints import find_member_if_email + + prisma = _org_prisma_with_users( + monkeypatch, + tuple( + LiteLLM_UserTable(user_id=user_id, user_email=email, user_role="internal_user") for user_id, email in rows + ), + ) + + if expected_user_id is None: + with pytest.raises(HTTPException) as exc: + await find_member_if_email("alice@example.com", prisma) + assert exc.value.status_code == 400 + else: + assert (await find_member_if_email("alice@example.com", prisma)).user_id == expected_user_id + + @pytest.mark.asyncio async def test_delete_organization_is_refused_while_one_of_its_teams_holds_a_ptu_share(monkeypatch): """Deleting an organization deletes its teams, so a team still holding a PTU share stops it the