fix: retrack PR2 based on PR1

haydnli-shopify · haydnli-shopify · commit e89a64858da2 · 2025-06-10T14:29:16.000-04:00
diff --git a/src/dstack/_internal/server/routers/projects.py b/src/dstack/_internal/server/routers/projects.py
@@ -7,13 +7,17 @@
 from dstack._internal.server.db import get_session
 from dstack._internal.server.models import ProjectModel, UserModel
 from dstack._internal.server.schemas.projects import (
+    AddProjectMemberRequest,
     CreateProjectRequest,
     DeleteProjectsRequest,
+    RemoveProjectMemberRequest,
     SetProjectMembersRequest,
 )
 from dstack._internal.server.security.permissions import (
     Authenticated,
     ProjectManager,
+    ProjectManagerOrPublicJoin,
+    ProjectManagerOrSelfLeave,
     ProjectMember,
 )
 from dstack._internal.server.services import projects
@@ -92,3 +96,41 @@ async def set_project_members(
     )
     await session.refresh(project)
     return projects.project_model_to_project(project)
+
+
+@router.post(
+    "/{project_name}/add_members",
+)
+async def add_project_members(
+    body: AddProjectMemberRequest,
+    session: AsyncSession = Depends(get_session),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectManagerOrPublicJoin()),
+) -> Project:
+    user, project = user_project
+    await projects.add_project_members(
+        session=session,
+        user=user,
+        project=project,
+        members=body.members,
+    )
+    await session.refresh(project)
+    return projects.project_model_to_project(project)
+
+
+@router.post(
+    "/{project_name}/remove_members",
+)
+async def remove_project_members(
+    body: RemoveProjectMemberRequest,
+    session: AsyncSession = Depends(get_session),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectManagerOrSelfLeave()),
+) -> Project:
+    user, project = user_project
+    await projects.remove_project_members(
+        session=session,
+        user=user,
+        project=project,
+        usernames=body.usernames,
+    )
+    await session.refresh(project)
+    return projects.project_model_to_project(project)
diff --git a/src/dstack/_internal/server/schemas/projects.py b/src/dstack/_internal/server/schemas/projects.py
@@ -25,3 +25,13 @@ class MemberSetting(CoreModel):
 
 class SetProjectMembersRequest(CoreModel):
     members: List[MemberSetting]
+
+
+class AddProjectMemberRequest(CoreModel):
+    # Always accept a list of members for cleaner API design
+    members: List[MemberSetting]
+
+
+class RemoveProjectMemberRequest(CoreModel):
+    # Always accept a list of usernames for cleaner API design
+    usernames: List[str]
diff --git a/src/dstack/_internal/server/security/permissions.py b/src/dstack/_internal/server/security/permissions.py
@@ -99,6 +99,76 @@ async def __call__(
         return await get_project_member(session, project_name, token.credentials)
 
 
+class ProjectManagerOrPublicJoin:
+    """
+    Allows:
+    1. Project managers to add any members
+    2. Any authenticated user to join public projects themselves
+    """
+    async def __call__(
+        self,
+        project_name: str,
+        session: AsyncSession = Depends(get_session),
+        token: HTTPAuthorizationCredentials = Security(HTTPBearer()),
+    ) -> Tuple[UserModel, ProjectModel]:
+        user = await log_in_with_token(session=session, token=token.credentials)
+        if user is None:
+            raise error_invalid_token()
+        project = await get_project_model_by_name(session=session, project_name=project_name)
+        if project is None:
+            raise error_not_found()
+        
+        # Global admin can always manage projects
+        if user.global_role == GlobalRole.ADMIN:
+            return user, project
+            
+        # Project managers can add members
+        project_role = get_user_project_role(user=user, project=project)
+        if project_role in [ProjectRole.ADMIN, ProjectRole.MANAGER]:
+            return user, project
+            
+        # For public projects, any authenticated user can join (will be validated in service layer)
+        if project.is_public:
+            return user, project
+            
+        raise error_forbidden()
+
+
+class ProjectManagerOrSelfLeave:
+    """
+    Allows:
+    1. Project managers to remove any members
+    2. Any project member to leave (remove themselves)
+    """
+    async def __call__(
+        self,
+        project_name: str,
+        session: AsyncSession = Depends(get_session),
+        token: HTTPAuthorizationCredentials = Security(HTTPBearer()),
+    ) -> Tuple[UserModel, ProjectModel]:
+        user = await log_in_with_token(session=session, token=token.credentials)
+        if user is None:
+            raise error_invalid_token()
+        project = await get_project_model_by_name(session=session, project_name=project_name)
+        if project is None:
+            raise error_not_found()
+        
+        # Global admin can always manage projects
+        if user.global_role == GlobalRole.ADMIN:
+            return user, project
+            
+        # Project managers can remove members
+        project_role = get_user_project_role(user=user, project=project)
+        if project_role in [ProjectRole.ADMIN, ProjectRole.MANAGER]:
+            return user, project
+            
+        # Any project member can leave (will be validated in service layer)
+        if project_role is not None:
+            return user, project
+            
+        raise error_forbidden()
+
+
 class OptionalServiceAccount:
     def __init__(self, token: Optional[str]) -> None:
         self._token = token
diff --git a/src/dstack/_internal/server/services/projects.py b/src/dstack/_internal/server/services/projects.py
@@ -1,6 +1,6 @@
 import uuid
 from datetime import timezone
-from typing import Awaitable, Callable, List, Optional, Tuple
+from typing import Awaitable, Callable, List, Optional, Tuple, Union
 
 from sqlalchemy import delete, select, update
 from sqlalchemy import func as safunc
@@ -75,8 +75,8 @@ async def list_user_accessible_projects(
 ) -> List[Project]:
     """
     Returns all projects accessible to the user:
-    - For global admins: ALL projects in the system
-    - For regular users: Projects where user is a member + public projects where user is NOT a member
+    - Projects where user is a member (public or private)
+    - Public projects where user is NOT a member
     """
     if user.global_role == GlobalRole.ADMIN:
         projects = await list_project_models(session=session)
@@ -231,6 +231,99 @@ async def set_project_members(
     await session.commit()
 
 
+async def add_project_members(
+    session: AsyncSession,
+    user: UserModel,
+    project: ProjectModel,
+    members: List[MemberSetting],
+):
+    """Add multiple members to a project."""
+    # reload with members
+    project = await get_project_model_by_name_or_error(
+        session=session,
+        project_name=project.name,
+    )
+    requesting_user_role = get_user_project_role(user=user, project=project)
+
+    # Check if this is a self-join to public project
+    is_self_join_to_public = (
+        len(members) == 1 and
+        project.is_public and
+        (members[0].username == user.name or members[0].username == user.email) and
+        requesting_user_role is None  # User is not already a member
+    )
+
+    # Check permissions: only managers/admins can add members, EXCEPT for self-join to public projects
+    if not is_self_join_to_public:
+        if requesting_user_role not in [ProjectRole.ADMIN, ProjectRole.MANAGER]:
+            raise ForbiddenError("Access denied: insufficient permissions to add members")
+
+        # For project managers, check if they're trying to add admins
+        if user.global_role != GlobalRole.ADMIN and requesting_user_role == ProjectRole.MANAGER:
+            for member in members:
+                if member.project_role == ProjectRole.ADMIN:
+                    raise ForbiddenError("Access denied: only global admins can add project admins")
+    else:
+        # For self-join to public project, only allow USER role
+        if members[0].project_role != ProjectRole.USER:
+            raise ForbiddenError("Access denied: can only join public projects as user role")
+
+    # Collect all usernames to query
+    usernames = [member.username for member in members]
+    
+    # Find all users (by username or email)
+    res = await session.execute(
+        select(UserModel).where(
+            (UserModel.name.in_(usernames)) | (UserModel.email.in_(usernames))
+        )
+    )
+    users_found = res.scalars().all()
+    
+    # Create lookup maps for both username and email
+    username_to_user = {user.name: user for user in users_found}
+    email_to_user = {user.email: user for user in users_found if user.email}
+    
+    # Build a set of current member IDs so we can quickly check if users are already members
+    member_ids = {m.user_id for m in project.members}
+    # Build a map from user_id to member for efficient existing member updates
+    member_by_user_id = {m.user_id: m for m in project.members}
+    
+    # Get current max member_num
+    max_member_num = 0
+    for member in project.members:
+        if member.member_num is not None and member.member_num > max_member_num:
+            max_member_num = member.member_num
+
+    # Process each member to add
+    for member_setting in members:
+        user_to_add = username_to_user.get(member_setting.username) or email_to_user.get(member_setting.username)
+        if user_to_add is None:
+            # Error on adding non-existing users instead of silently skipping
+            raise ServerClientError(f"User not found: {member_setting.username}")
+
+        # Check if user is already a member using our fast lookup set
+        if user_to_add.id in member_ids:
+            # Update existing member role if different using our efficient lookup
+            existing_member = member_by_user_id[user_to_add.id]
+            if existing_member.project_role != member_setting.project_role:
+                existing_member.project_role = member_setting.project_role
+        else:
+            # Add new member
+            max_member_num += 1
+            await add_project_member(
+                session=session,
+                project=project,
+                user=user_to_add,
+                project_role=member_setting.project_role,
+                member_num=max_member_num,
+                commit=False,
+            )
+            # Keep track of newly added members for subsequent iterations
+            member_ids.add(user_to_add.id)
+
+    await session.commit()
+
+
 async def add_project_member(
     session: AsyncSession,
     project: ProjectModel,
@@ -511,3 +604,92 @@ def _is_project_admin(
             if m.project_role == ProjectRole.ADMIN:
                 return True
     return False
+
+
+async def remove_project_members(
+    session: AsyncSession,
+    user: UserModel,
+    project: ProjectModel,
+    usernames: List[str],
+):
+    """Remove multiple members from a project."""
+    # reload with members
+    project = await get_project_model_by_name_or_error(
+        session=session,
+        project_name=project.name,
+    )
+    requesting_user_role = get_user_project_role(user=user, project=project)
+
+    # Check if this is a self-leave (user removing themselves)
+    is_self_leave = (
+        len(usernames) == 1 and
+        (usernames[0] == user.name or usernames[0] == user.email) and
+        requesting_user_role is not None  # User is actually a member
+    )
+
+    # Check basic permissions: only managers/admins can remove members, EXCEPT for self-leave
+    if not is_self_leave:
+        if requesting_user_role not in [ProjectRole.ADMIN, ProjectRole.MANAGER]:
+            raise ForbiddenError("Access denied: insufficient permissions to remove members")
+
+    # Find all users to remove (by username or email)
+    res = await session.execute(
+        select(UserModel).where(
+            (UserModel.name.in_(usernames)) | (UserModel.email.in_(usernames))
+        )
+    )
+    users_found = res.scalars().all()
+    
+    # Create lookup maps
+    username_to_user = {user.name: user for user in users_found}
+    email_to_user = {user.email: user for user in users_found if user.email}
+    
+    # Build a set of member IDs for faster membership checks
+    member_user_ids = {m.user_id for m in project.members}
+    # Build a map from user_id to member for efficient member lookups
+    member_by_user_id = {m.user_id: m for m in project.members}
+
+    # Find members to remove and validate permissions
+    members_to_remove = []
+    admin_removals = 0
+    
+    for username in usernames:
+        user_to_remove = username_to_user.get(username) or email_to_user.get(username)
+        if user_to_remove is None:
+            # Error on removing non-existing users instead of silently skipping
+            raise ServerClientError(f"User not found: {username}")
+
+        # Check if user is actually a member before trying to remove them
+        if user_to_remove.id not in member_user_ids:
+            # Error on removing non-members instead of silently skipping
+            raise ServerClientError(f"User is not a member of this project: {username}")
+
+        # Get the member to remove using our efficient lookup
+        member_to_remove = member_by_user_id[user_to_remove.id]
+
+        # Check if trying to remove project admin
+        if member_to_remove.project_role == ProjectRole.ADMIN:
+            if is_self_leave:
+                # For self-leave, check if user is the last admin
+                total_admins = sum(1 for member in project.members if member.project_role == ProjectRole.ADMIN)
+                if total_admins <= 1:
+                    raise ServerClientError("Cannot leave project: you are the last admin")
+            else:
+                # For manager/admin removing other admins, only global admins can do this
+                if user.global_role != GlobalRole.ADMIN:
+                    raise ForbiddenError(f"Access denied: only global admins can remove project admins (user: {username})")
+            admin_removals += 1
+
+        members_to_remove.append(member_to_remove)
+
+    # Check we're not removing all admins (for non-self-leave operations)
+    if not is_self_leave:
+        total_admins = sum(1 for member in project.members if member.project_role == ProjectRole.ADMIN)
+        if admin_removals >= total_admins:
+            raise ServerClientError("Cannot remove all project admins")
+
+    # Remove all members
+    for member in members_to_remove:
+        await session.delete(member)
+    
+    await session.commit()
diff --git a/src/dstack/api/server/_projects.py b/src/dstack/api/server/_projects.py
diff --git a/src/tests/_internal/server/routers/test_projects.py b/src/tests/_internal/server/routers/test_projects.py
