feat(BA-4865): Apply RBAC validator infrastructure to Session actions (#9624)

Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>

Author: fregataa

changes/9624.feature.md

@@ -0,0 +1 @@
+Add RBAC validator infrastructure to Session actions following BEP-1048 patterns

src/ai/backend/manager/api/gql/data_loader/kernel/loader.py

@@ -2,8 +2,10 @@
 
 from collections.abc import Sequence
 
+from ai.backend.common.contexts.user import current_user
 from ai.backend.common.types import KernelId
 from ai.backend.manager.data.kernel.types import KernelInfo
+from ai.backend.manager.errors.user import UserNotFound
 from ai.backend.manager.repositories.base import BatchQuerier, NoPagination
 from ai.backend.manager.repositories.scheduler.options import KernelConditions
 from ai.backend.manager.services.session.actions.search_kernel import SearchKernelsAction
@@ -26,13 +28,17 @@ async def load_kernels_by_ids(
     if not kernel_ids:
         return []
 
+    user = current_user()
+    if user is None:
+        raise UserNotFound("User not found in context")
+
     querier = BatchQuerier(
         pagination=NoPagination(),
         conditions=[KernelConditions.by_ids(kernel_ids)],
     )
 
     action_result = await processor.search_kernels.wait_for_complete(
-        SearchKernelsAction(querier=querier)
+        SearchKernelsAction(querier=querier, user_id=user.user_id)
     )
 
     kernel_map: dict[KernelId, KernelInfo] = {kernel.id: kernel for kernel in action_result.data}

src/ai/backend/manager/api/gql/data_loader/session/loader.py

@@ -2,8 +2,10 @@
 
 from collections.abc import Sequence
 
+from ai.backend.common.contexts.user import current_user
 from ai.backend.common.types import SessionId
 from ai.backend.manager.data.session.types import SessionData
+from ai.backend.manager.errors.user import UserNotFound
 from ai.backend.manager.repositories.base import BatchQuerier, NoPagination
 from ai.backend.manager.repositories.scheduler.options import SessionConditions
 from ai.backend.manager.services.session.actions.search import SearchSessionsAction
@@ -17,13 +19,17 @@ async def load_sessions_by_ids(
     if not session_ids:
         return []
 
+    user = current_user()
+    if user is None:
+        raise UserNotFound("User not found in context")
+
     querier = BatchQuerier(
         pagination=NoPagination(),
         conditions=[SessionConditions.by_ids(session_ids)],
     )
 
     action_result = await processor.search_sessions.wait_for_complete(
-        SearchSessionsAction(querier=querier)
+        SearchSessionsAction(querier=querier, user_id=user.user_id)
     )
 
     session_map: dict[SessionId, SessionData] = {

src/ai/backend/manager/api/gql/kernel/fetcher/kernel.py

@@ -5,6 +5,7 @@
 import strawberry
 from strawberry import Info
 
+from ai.backend.common.contexts.user import current_user
 from ai.backend.common.types import KernelId
 from ai.backend.manager.api.gql.adapter import PaginationOptions, PaginationSpec
 from ai.backend.manager.api.gql.base import encode_cursor
@@ -16,6 +17,7 @@
     KernelV2OrderByGQL,
 )
 from ai.backend.manager.api.gql.types import StrawberryGQLContext
+from ai.backend.manager.errors.user import UserNotFound
 from ai.backend.manager.models.kernel import KernelRow
 from ai.backend.manager.repositories.base import QueryCondition
 from ai.backend.manager.repositories.scheduler.options import KernelConditions
@@ -45,6 +47,10 @@ async def fetch_kernels(
     offset: int | None = None,
     base_conditions: list[QueryCondition] | None = None,
 ) -> KernelV2ConnectionGQL:
+    user = current_user()
+    if user is None:
+        raise UserNotFound("User not found in context")
+
     querier = info.context.gql_adapter.build_querier(
         PaginationOptions(
             first=first,
@@ -61,7 +67,7 @@ async def fetch_kernels(
     )
 
     action_result = await info.context.processors.session.search_kernels.wait_for_complete(
-        SearchKernelsAction(querier=querier)
+        SearchKernelsAction(querier=querier, user_id=user.user_id)
     )
     nodes = [KernelV2GQL.from_kernel_info(kernel_info) for kernel_info in action_result.data]
     edges = [KernelV2EdgeGQL(node=node, cursor=encode_cursor(node.id)) for node in nodes]

src/ai/backend/manager/api/gql/session/fetcher/session.py

@@ -5,6 +5,7 @@
 import strawberry
 from strawberry import Info
 
+from ai.backend.common.contexts.user import current_user
 from ai.backend.manager.api.gql.adapter import PaginationOptions, PaginationSpec
 from ai.backend.manager.api.gql.base import encode_cursor
 from ai.backend.manager.api.gql.session.types import (
@@ -15,6 +16,7 @@
     SessionV2OrderByGQL,
 )
 from ai.backend.manager.api.gql.types import StrawberryGQLContext
+from ai.backend.manager.errors.user import UserNotFound
 from ai.backend.manager.models.session import SessionRow
 from ai.backend.manager.repositories.base import QueryCondition
 from ai.backend.manager.repositories.scheduler.options import SessionConditions, SessionOrders
@@ -44,6 +46,10 @@ async def fetch_sessions(
     offset: int | None = None,
     base_conditions: list[QueryCondition] | None = None,
 ) -> SessionV2ConnectionGQL:
+    user = current_user()
+    if user is None:
+        raise UserNotFound("User not found in context")
+
     querier = info.context.gql_adapter.build_querier(
         PaginationOptions(
             first=first,
@@ -60,7 +66,7 @@ async def fetch_sessions(
     )
 
     action_result = await info.context.processors.session.search_sessions.wait_for_complete(
-        SearchSessionsAction(querier=querier)
+        SearchSessionsAction(querier=querier, user_id=user.user_id)
     )
 
     nodes = [SessionV2GQL.from_data(session_data) for session_data in action_result.data]

src/ai/backend/manager/api/gql/session/types.py

@@ -12,6 +12,7 @@
 from strawberry import ID, Info
 from strawberry.relay import Connection, Edge, Node, NodeID
 
+from ai.backend.common.contexts.user import current_user
 from ai.backend.common.types import SessionId
 from ai.backend.manager.api.gql.base import OrderDirection, StringFilter, UUIDFilter, encode_cursor
 from ai.backend.manager.api.gql.common.types import (
@@ -41,6 +42,7 @@
 from ai.backend.manager.api.gql.types import GQLFilter, GQLOrderBy, StrawberryGQLContext
 from ai.backend.manager.api.gql.user.types.node import UserV2GQL
 from ai.backend.manager.data.session.types import SessionData, SessionStatus
+from ai.backend.manager.errors.user import UserNotFound
 from ai.backend.manager.repositories.base import (
     BatchQuerier,
     NoPagination,
@@ -468,13 +470,17 @@ async def images(self) -> ImageV2ConnectionGQL:
         description="Added in 26.3.0. The kernels belonging to this session."
     )
     async def kernels(self, info: Info[StrawberryGQLContext]) -> KernelV2ConnectionGQL:
+        user = current_user()
+        if user is None:
+            raise UserNotFound("User not found in context")
+
         session_id = SessionId(UUID(str(self.id)))
         querier = BatchQuerier(
             pagination=NoPagination(),
             conditions=[KernelConditions.by_session_ids([session_id])],
         )
         action_result = await info.context.processors.session.search_kernels.wait_for_complete(
-            SearchKernelsAction(querier=querier)
+            SearchKernelsAction(querier=querier, user_id=user.user_id)
         )
         nodes = [KernelV2GQL.from_kernel_info(kernel) for kernel in action_result.data]
         edges = [KernelV2EdgeGQL(node=node, cursor=encode_cursor(node.id)) for node in nodes]

src/ai/backend/manager/api/rest/compute_sessions/handler.py

@@ -7,6 +7,7 @@
 from typing import Final
 
 from ai.backend.common.api_handlers import APIResponse, BodyParam
+from ai.backend.common.contexts.user import current_user
 from ai.backend.common.dto.manager.compute_session import (
     PaginationInfo,
     SearchComputeSessionsRequest,
@@ -15,6 +16,7 @@
 from ai.backend.common.types import SessionId
 from ai.backend.logging import BraceStyleAdapter
 from ai.backend.manager.dto.context import UserContext
+from ai.backend.manager.errors.user import UserNotFound
 from ai.backend.manager.services.session.actions.search import SearchSessionsAction
 from ai.backend.manager.services.session.actions.search_kernel import SearchKernelsAction
 from ai.backend.manager.services.session.processors import SessionProcessors
@@ -39,10 +41,14 @@ async def search_sessions(
         """Search compute sessions with nested container data."""
         log.info("SEARCH_SESSIONS (ak:{})", ctx.access_key)
 
+        user = current_user()
+        if user is None:
+            raise UserNotFound("User not found in context")
+
         # Step 1: Search sessions
         session_querier = self._adapter.build_session_querier(body.parsed)
         session_result = await self._session.search_sessions.wait_for_complete(
-            SearchSessionsAction(querier=session_querier)
+            SearchSessionsAction(querier=session_querier, user_id=user.user_id)
         )
 
         # Step 2: Fetch kernels for found sessions
@@ -51,7 +57,7 @@ async def search_sessions(
         if session_ids:
             kernel_querier = self._adapter.build_kernel_querier_for_sessions(session_ids)
             kernel_result = await self._session.search_kernels.wait_for_complete(
-                SearchKernelsAction(querier=kernel_querier)
+                SearchKernelsAction(querier=kernel_querier, user_id=user.user_id)
             )
             kernels_by_session = self._adapter.group_kernels_by_session(kernel_result.data)
 

src/ai/backend/manager/api/rest/session/handler.py

@@ -20,6 +20,7 @@
 from pydantic import BaseModel
 
 from ai.backend.common.api_handlers import APIResponse, BaseResponseModel, BodyParam, QueryParam
+from ai.backend.common.contexts.user import current_user
 from ai.backend.common.dto.manager.session.request import (
     CommitSessionRequest,
     CompleteRequest,
@@ -91,6 +92,7 @@
 from ai.backend.manager.errors.api import InvalidAPIParameters
 from ai.backend.manager.errors.auth import InsufficientPrivilege
 from ai.backend.manager.errors.resource import NoCurrentTaskContext
+from ai.backend.manager.errors.user import UserNotFound
 from ai.backend.manager.models.user import UserRole
 from ai.backend.manager.services.agent.actions.sync_agent_registry import (
     SyncAgentRegistryAction,
@@ -515,6 +517,9 @@ async def match_sessions(
             )
         )
         requester_access_key, owner_access_key = scope.requester_access_key, scope.owner_access_key
+        user = current_user()
+        if user is None:
+            raise UserNotFound("User not found in context")
         log.info(
             "MATCH_SESSIONS(ak:{0}/{1}, prefix:{2})",
             requester_access_key,
@@ -525,6 +530,7 @@ async def match_sessions(
             MatchSessionsAction(
                 id_or_name_prefix=params.id,
                 owner_access_key=owner_access_key,
+                user_id=user.user_id,
             )
         )
         return APIResponse.build(HTTPStatus.OK, MatchSessionsResponse(matches=result.result))

src/ai/backend/manager/services/session/actions/create_cluster.py

@@ -3,15 +3,23 @@
 from dataclasses import dataclass
 from typing import Any, override
 
+from ai.backend.common.data.permission.types import RBACElementType, ScopeType
 from ai.backend.common.types import AccessKey, SessionTypes
 from ai.backend.manager.actions.action import BaseActionResult
 from ai.backend.manager.actions.types import ActionOperationType
+from ai.backend.manager.data.permission.types import RBACElementRef
 from ai.backend.manager.models.user import UserRole
-from ai.backend.manager.services.session.base import SessionAction
+from ai.backend.manager.services.session.base import SessionScopeAction
 
 
 @dataclass
-class CreateClusterAction(SessionAction):
+class CreateClusterAction(SessionScopeAction):
+    """Create a new cluster session.
+
+    RBAC validation checks if the user has CREATE permission in USER scope.
+    Scope is always USER scope with user_id.
+    """
+
     session_name: str
     user_id: uuid.UUID
     user_role: UserRole
@@ -37,6 +45,21 @@ def entity_id(self) -> str | None:
     def operation_type(cls) -> ActionOperationType:
         return ActionOperationType.CREATE
 
+    @override
+    def scope_type(self) -> ScopeType:
+        return ScopeType.USER
+
+    @override
+    def scope_id(self) -> str:
+        return str(self.user_id)
+
+    @override
+    def target_element(self) -> RBACElementRef:
+        return RBACElementRef(
+            element_type=RBACElementType.USER,
+            element_id=str(self.user_id),
+        )
+
 
 @dataclass
 class CreateClusterActionResult(BaseActionResult):

src/ai/backend/manager/services/session/actions/create_from_params.py

@@ -6,11 +6,13 @@
 
 import yarl
 
+from ai.backend.common.data.permission.types import RBACElementType, ScopeType
 from ai.backend.common.types import AccessKey, ClusterMode, SessionTypes
 from ai.backend.manager.actions.action import BaseActionResult
 from ai.backend.manager.actions.types import ActionOperationType
+from ai.backend.manager.data.permission.types import RBACElementRef
 from ai.backend.manager.models.user import UserRole
-from ai.backend.manager.services.session.base import SessionAction
+from ai.backend.manager.services.session.base import SessionScopeAction
 
 
 # TODO: Idea: Refactor this type using pydantic and utilize as API model
@@ -41,7 +43,13 @@ class CreateFromParamsActionParams:
 
 
 @dataclass
-class CreateFromParamsAction(SessionAction):
+class CreateFromParamsAction(SessionScopeAction):
+    """Create a new session from parameters.
+
+    RBAC validation checks if the user has CREATE permission in USER scope.
+    Scope is always USER scope with user_id.
+    """
+
     params: CreateFromParamsActionParams
     user_id: uuid.UUID
     user_role: UserRole
@@ -58,6 +66,21 @@ def entity_id(self) -> str | None:
     def operation_type(cls) -> ActionOperationType:
         return ActionOperationType.CREATE
 
+    @override
+    def scope_type(self) -> ScopeType:
+        return ScopeType.USER
+
+    @override
+    def scope_id(self) -> str:
+        return str(self.user_id)
+
+    @override
+    def target_element(self) -> RBACElementRef:
+        return RBACElementRef(
+            element_type=RBACElementType.USER,
+            element_id=str(self.user_id),
+        )
+
 
 @dataclass
 class CreateFromParamsActionResult(BaseActionResult):