refactor(BA-5650): unify session ownership to owner_id and main_access_key [1/2]

changes/11046.enhance.md

@@ -0,0 +1 @@
+Collapse `SessionRepository` / `SessionDBSource` signatures to take `owner_id: UUID` instead of `owner_access_key: AccessKey`. No external behavior change.

changes/11048.enhance.md

@@ -0,0 +1 @@
+Propagate the `owner_id` / `main_access_key` signature rename into the sokovan data classes, scheduler handlers, provisioner validators, launcher, scheduling controller, and sequencers.

changes/11050.breaking.md

@@ -0,0 +1 @@
+Remove `owner_access_key` from REST v1 session API; resolve `owner_id` via `current_user()` in the service layer. Clients must migrate to `owner_id` (user UUID) for delegation.

src/ai/backend/common/dto/manager/session/request.py

@@ -135,7 +135,7 @@ class CreateFromTemplateRequest(BaseRequestModel):
         default=None,
         validation_alias=AliasChoices("callback_url", "callbackUrl", "callbackURL"),
     )
-    owner_access_key: str | None = None
+    owner_id: UUID | None = None
 
 
 class CreateFromParamsRequest(BaseRequestModel):
@@ -214,7 +214,7 @@ class CreateFromParamsRequest(BaseRequestModel):
         default=None,
         validation_alias=AliasChoices("callback_url", "callbackUrl", "callbackURL"),
     )
-    owner_access_key: str | None = None
+    owner_id: UUID | None = None
 
 
 class CreateClusterRequest(BaseRequestModel):
@@ -252,7 +252,7 @@ class CreateClusterRequest(BaseRequestModel):
         ge=0,
         validation_alias=AliasChoices("max_wait_seconds", "maxWaitSeconds"),
     )
-    owner_access_key: str | None = None
+    owner_id: UUID | None = None
 
 
 # ---------------------------------------------------------------------------
@@ -352,14 +352,11 @@ class DestroySessionRequest(BaseRequestModel):
 
     forced: bool = False
     recursive: bool = False
-    owner_access_key: str | None = None
 
 
 class RestartSessionRequest(BaseRequestModel):
     """PATCH ``/{session_name}``"""
 
-    owner_access_key: str | None = None
-
 
 class MatchSessionsRequest(BaseRequestModel):
     """GET ``/_/match``"""
@@ -419,10 +416,6 @@ class ListFilesRequest(BaseRequestModel):
 class GetContainerLogsRequest(BaseRequestModel):
     """GET ``/{session_name}/logs``"""
 
-    owner_access_key: str | None = Field(
-        default=None,
-        validation_alias=AliasChoices("owner_access_key", "ownerAccessKey"),
-    )
     kernel_id: UUID | None = Field(
         default=None,
         validation_alias=AliasChoices("kernel_id", "kernelId"),
@@ -441,5 +434,3 @@ class GetTaskLogsRequest(BaseRequestModel):
 
 class GetStatusHistoryRequest(BaseRequestModel):
     """GET ``/{session_name}/status-history``"""
-
-    owner_access_key: str | None = None

src/ai/backend/manager/api/adapters/session.py

@@ -223,8 +223,8 @@ async def enqueue(
 
         When ``input.owner_id`` is set, the session is created on behalf of the
         target user: their main access key, role, and domain are used in place
-        of the caller's. Resolution and authorization of the delegated user
-        are handled by the downstream session service, not by this adapter.
+        of the caller's. The target user must be loadable via the user
+        processor (RBAC enforced).
         """
         batch_spec: SessionBatchSpec | None = None
         if input.batch is not None:
@@ -849,12 +849,10 @@ async def shutdown_service(
         self,
         session_id: UUID,
         input: ShutdownSessionServiceInput,
-        access_key: str,
     ) -> None:
         """Shut down a service in a session."""
         action = ShutdownServiceAction(
             session_name=str(session_id),
-            owner_access_key=AccessKey(access_key),
             service_name=input.service,
         )
         await self._processors.session.shutdown_service.wait_for_complete(action)
@@ -866,13 +864,11 @@ async def shutdown_service(
     async def get_logs(
         self,
         session_id: UUID,
-        access_key: str,
         kernel_id: UUID | None = None,
     ) -> SessionLogsPayload:
         """Get container logs for a session."""
         action = GetContainerLogsAction(
             session_name=str(session_id),
-            owner_access_key=AccessKey(access_key),
             kernel_id=KernelId(kernel_id) if kernel_id else None,
         )
         result = await self._processors.session.get_container_logs.wait_for_complete(action)
@@ -887,14 +883,12 @@ async def update(
         self,
         session_id: UUID,
         input: UpdateSessionInput,
-        access_key: str,
     ) -> UpdateSessionPayload:
         """Update session fields (currently supports rename only)."""
         if input.name is not None:
             action = RenameSessionAction(
                 session_name=str(session_id),
                 new_name=input.name,
-                owner_access_key=AccessKey(access_key),
             )
             result = await self._processors.session.rename_session.wait_for_complete(action)
             return UpdateSessionPayload(session=self._session_data_to_node(result.session_data))
@@ -933,13 +927,13 @@ def _session_data_to_node(data: SessionData) -> SessionNode:
         return SessionNode(
             id=data.id,
             domain_name=data.domain_name,
-            user_id=data.user_uuid,
+            user_id=data.owner_id,
             project_id=data.group_id,
             metadata=SessionMetadataInfoGQLDTO(
                 creation_id=data.creation_id or "",
                 name=data.name or "",
                 session_type=data.session_type.value,
-                access_key=str(data.access_key) if data.access_key else "",
+                access_key="",
                 cluster_mode=data.cluster_mode.name,
                 cluster_size=data.cluster_size,
                 priority=data.priority,
@@ -1011,8 +1005,8 @@ def _kernel_info_to_node(info: KernelInfo) -> KernelNode:
                 session_type=info.session.session_type.value,
             ),
             user_info=KernelUserInfoGQLDTO(
-                user_id=info.user_permission.user_uuid,
-                access_key=info.user_permission.access_key,
+                user_id=info.user_permission.owner_id,
+                access_key=info.user_permission.main_access_key,
                 domain_name=info.user_permission.domain_name,
                 group_id=info.user_permission.group_id,
             ),

src/ai/backend/manager/api/gql_legacy/session.py

@@ -23,7 +23,6 @@
 
 from ai.backend.common import validators as tx
 from ai.backend.common.defs.session import SESSION_PRIORITY_MAX, SESSION_PRIORITY_MIN
-from ai.backend.common.exception import SessionWithInvalidStateError
 from ai.backend.common.types import (
     ClusterMode,
     KernelId,
@@ -395,19 +394,25 @@ def from_dataclass(
         cls,
         ctx: GraphQueryContext,
         session_data: SessionData,
+        main_access_key: str | None,
         *,
         permissions: Iterable[ComputeSessionPermission] | None = None,
     ) -> Self:
+        """Build a ``ComputeSessionNode`` from session data.
+
+        ``main_access_key`` must be pre-resolved by the caller (typically
+        via ``UserRepository.get_main_access_key_by_id(session_data.owner_id)``
+        or by eagerly loading ``session_data.owner``). Keeping the helper
+        synchronous avoids a hidden per-session DB query and lets the
+        caller batch the lookup across nodes.
+        """
         status_history = session_data.status_history or {}
         raw_scheduled_at = status_history.get(SessionStatus.SCHEDULED.name)
         if not session_data.vfolder_mounts:
             vfolder_mounts = []
         else:
             vfolder_mounts = [vf.vfid.folder_id for vf in session_data.vfolder_mounts]
 
-        if session_data.owner is None:
-            raise SessionWithInvalidStateError()
-
         result = cls(
             # identity
             id=session_data.id,  # auto-converted to Relay global ID
@@ -422,9 +427,9 @@ def from_dataclass(
             # ownership
             domain_name=session_data.domain_name,
             project_id=session_data.group_id,
-            user_id=session_data.user_uuid,
-            access_key=session_data.access_key,
-            owner=UserNode.from_dataclass(ctx, session_data.owner),
+            user_id=session_data.owner_id,
+            access_key=main_access_key,
+            owner=UserNode.from_dataclass(ctx, session_data.owner) if session_data.owner else None,
             # status
             status=session_data.status.name,
             # status_changed=row.status_changed,  # FIXME: generated attribute
@@ -918,8 +923,14 @@ async def mutate_and_get_payload(
             )
         )
 
+        session_data = result.session_data
+        main_access_key = (
+            session_data.owner.main_access_key
+            if session_data.owner
+            else await graph_ctx.user_repository.get_main_access_key_by_id(session_data.owner_id)
+        )
         return ModifyComputeSession(
-            ComputeSessionNode.from_dataclass(graph_ctx, result.session_data),
+            ComputeSessionNode.from_dataclass(graph_ctx, session_data, main_access_key),
             input.get("client_mutation_id"),
         )
 
@@ -969,8 +980,16 @@ async def mutate(
                     )
                 )
             )
+            session_data = action_result.session_data
+            main_access_key = (
+                session_data.owner.main_access_key
+                if session_data.owner
+                else await graph_ctx.user_repository.get_main_access_key_by_id(
+                    session_data.owner_id
+                )
+            )
             session_nodes.append(
-                ComputeSessionNode.from_dataclass(graph_ctx, action_result.session_data)
+                ComputeSessionNode.from_dataclass(graph_ctx, session_data, main_access_key)
             )
 
         return CheckAndTransitStatus(session_nodes, input.get("client_mutation_id"))