fix(BA-5650-G): align scheduler db_source and tests with renamed fields

Update remaining call sites in scheduler db_source and test fixtures
to use the renamed dataclass fields:

- ``main_access_key`` (previously ``access_key``) on
  ``ScheduledSessionData``, ``SessionDataForPull``, ``SessionDataForStart``,
  ``SessionWorkload``, ``SweptSessionInfo``, ``TerminatingSessionData``,
  ``PendingSessionData``.
- ``owner_id`` (previously ``user_uuid``) on
  ``SessionEnqueueData``, ``KernelEnqueueData``, ``SessionDataForStart``,
  ``SessionWorkload``.
- ``pending_sessions.owner_ids`` (previously ``user_uuids``) on
  ``PendingSessions``.

Joined ``users`` for SweptSessionInfo and ScheduledSessionData queries
so ``main_access_key`` is sourced from the user record rather than the
session column being dropped in slice K.

Drop unused ``target_main_access_key`` argument from
``EndpointRow.delegate_endpoint_ownership`` call (already unused by the
ORM helper); keep the parameter on the repository facade for caller
compatibility.

Inject ``user_repository=MagicMock()`` into ``AgentRegistry`` test
construction in ``test_reconcile_agent_resources``.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: jopemachine

src/ai/backend/manager/repositories/scheduler/db_source/db_source.py

@@ -700,7 +700,7 @@ async def _fetch_user_policies(
         """Fetch user resource policies for users in pending sessions."""
         user_policies: dict[UUID, UserResourcePolicy] = {}
 
-        if not pending_sessions.user_uuids:
+        if not pending_sessions.owner_ids:
             return user_policies
 
         user_policy_result = await db_sess.execute(
@@ -716,7 +716,7 @@ async def _fetch_user_policies(
                 KeyPairResourcePolicyRow,
                 KeyPairRow.resource_policy == KeyPairResourcePolicyRow.name,
             )
-            .where(UserRow.uuid.in_(pending_sessions.user_uuids))
+            .where(UserRow.uuid.in_(pending_sessions.owner_ids))
         )
 
         for row in user_policy_result:
@@ -1146,11 +1146,14 @@ async def get_terminating_sessions_by_ids(
                     for kernel in session_row.kernels
                 ]
 
+                owner_main_ak = (
+                    session_row.user.main_access_key if session_row.user else None
+                )
                 terminating_sessions.append(
                     TerminatingSessionData(
                         session_id=session_row.id,
-                        access_key=AccessKey(session_row.access_key)
-                        if session_row.access_key
+                        main_access_key=AccessKey(owner_main_ak)
+                        if owner_main_ak
                         else AccessKey(""),
                         creation_id=session_row.creation_id or "",
                         status=session_row.status,
@@ -1183,11 +1186,12 @@ async def get_pending_timeout_sessions_by_ids(
                 sa.select(
                     SessionRow.id,
                     SessionRow.creation_id,
-                    SessionRow.access_key,
+                    UserRow.main_access_key,
                     SessionRow.created_at,
                     ScalingGroupRow.scheduler_opts,
                 )
                 .select_from(SessionRow)
+                .join(UserRow, SessionRow.user_uuid == UserRow.uuid)
                 .join(ScalingGroupRow, SessionRow.scaling_group_name == ScalingGroupRow.name)
                 .where(
                     SessionRow.id.in_(session_ids),
@@ -1213,7 +1217,7 @@ async def get_pending_timeout_sessions_by_ids(
                         SweptSessionInfo(
                             session_id=row.id,
                             creation_id=row.creation_id,
-                            access_key=row.access_key,
+                            main_access_key=row.main_access_key,
                         )
                     )
 
@@ -1302,8 +1306,8 @@ async def enqueue_session(
                 id=session_data.id,
                 creation_id=session_data.creation_id,
                 name=session_data.name,
-                access_key=session_data.access_key,
-                user_uuid=session_data.user_uuid,
+                access_key=session_data.main_access_key,
+                user_uuid=session_data.owner_id,
                 group_id=session_data.group_id,
                 domain_name=session_data.domain_name,
                 scaling_group_name=session_data.scaling_group_name,
@@ -1349,8 +1353,8 @@ async def enqueue_session(
                     scaling_group=kernel.scaling_group,
                     domain_name=kernel.domain_name,
                     group_id=kernel.group_id,
-                    user_uuid=kernel.user_uuid,
-                    access_key=kernel.access_key,
+                    user_uuid=kernel.owner_id,
+                    access_key=kernel.main_access_key,
                     image=kernel.image,
                     architecture=kernel.architecture,
                     registry=kernel.registry,
@@ -1387,7 +1391,7 @@ async def enqueue_session(
                 element_type=RBACElementType.SESSION,
                 scope_ref=RBACElementRef(
                     element_type=RBACElementType.USER,
-                    element_id=str(session_data.user_uuid),
+                    element_id=str(session_data.owner_id),
                 ),
                 additional_scope_refs=[
                     RBACElementRef(
@@ -1843,21 +1847,28 @@ async def allocate_sessions(
             # First, fetch session data to get creation_id and access_key
             session_ids = {alloc.session_id for alloc in allocation_batch.allocations}
             if session_ids:
-                query = sa.select(
-                    SessionRow.id, SessionRow.creation_id, SessionRow.access_key
-                ).where(SessionRow.id.in_(session_ids))
+                query = (
+                    sa.select(
+                        SessionRow.id, SessionRow.creation_id, UserRow.main_access_key
+                    )
+                    .select_from(SessionRow)
+                    .join(UserRow, SessionRow.user_uuid == UserRow.uuid)
+                    .where(SessionRow.id.in_(session_ids))
+                )
                 result = await db_sess.execute(query)
-                session_data_map = {row.id: (row.creation_id, row.access_key) for row in result}
+                session_data_map = {
+                    row.id: (row.creation_id, row.main_access_key) for row in result
+                }
 
                 # Create SessionEventData for each allocated session
                 for allocation in allocation_batch.allocations:
                     if session_data := session_data_map.get(allocation.session_id):
-                        creation_id, access_key = session_data
+                        creation_id, main_access_key = session_data
                         scheduled_sessions.append(
                             ScheduledSessionData(
                                 session_id=allocation.session_id,
                                 creation_id=creation_id,
-                                access_key=access_key,
+                                main_access_key=main_access_key,
                                 reason="triggered-by-scheduler",
                             )
                         )
@@ -2917,7 +2928,9 @@ async def _get_sessions_by_statuses(
             scheduled_session = ScheduledSessionData(
                 session_id=session.id,
                 creation_id=session.creation_id or "",
-                access_key=AccessKey(session.access_key) if session.access_key else AccessKey(""),
+                main_access_key=AccessKey(session.access_key)
+                if session.access_key
+                else AccessKey(""),
                 reason="triggered-by-scheduler",
             )
             scheduled_sessions.append(scheduled_session)
@@ -2962,7 +2975,7 @@ async def _get_scheduled_sessions(self, db_sess: SASession) -> list[ScheduledSes
                 ScheduledSessionData(
                     session_id=session.id,
                     creation_id=session.creation_id or "",
-                    access_key=AccessKey(session.access_key)
+                    main_access_key=AccessKey(session.access_key)
                     if session.access_key
                     else AccessKey(""),
                     reason="triggered-by-scheduler",
@@ -3102,7 +3115,7 @@ async def _get_sessions_for_pull(
                 sessions_map[session_id] = SessionDataForPull(
                     session_id=session_id,
                     creation_id=row.creation_id,
-                    access_key=row.access_key,
+                    main_access_key=row.access_key,
                     kernels=[],
                 )
 
@@ -3294,13 +3307,13 @@ async def _get_sessions_for_start(
                 SessionDataForStart(
                     session_id=session_info["id"],
                     creation_id=session_info["creation_id"],
-                    access_key=session_info["access_key"],
+                    main_access_key=session_info["access_key"],
                     session_type=session_info["session_type"],
                     name=session_info["name"],
                     cluster_mode=session_info["cluster_mode"],
                     kernels=kernel_bindings,
                     environ=session_info.get("environ", {}),
-                    user_uuid=session_info["user_uuid"],
+                    owner_id=session_info["user_uuid"],
                     user_email=user_info.email,
                     user_name=user_info.username,
                 )
@@ -4074,7 +4087,7 @@ async def _fetch_sessions_for_pull_by_ids(
                 sessions_map[session_id] = SessionDataForPull(
                     session_id=session_id,
                     creation_id=row.creation_id,
-                    access_key=row.access_key,
+                    main_access_key=row.access_key,
                     kernels=[],
                 )
 
@@ -4293,13 +4306,13 @@ async def _fetch_sessions_for_start_by_ids(
                 SessionDataForStart(
                     session_id=session_info["id"],
                     creation_id=session_info["creation_id"],
-                    access_key=session_info["access_key"],
+                    main_access_key=session_info["access_key"],
                     session_type=session_info["session_type"],
                     name=session_info["name"],
                     cluster_mode=session_info["cluster_mode"],
                     kernels=kernel_bindings,
                     environ=session_info.get("environ", {}),
-                    user_uuid=session_info["user_uuid"],
+                    owner_id=session_info["user_uuid"],
                     user_email=user_info.email,
                     user_name=user_info.username,
                 )
@@ -4369,7 +4382,7 @@ async def search_sessions_with_kernels(
                 sessions_map[row.id] = SessionDataForPull(
                     session_id=row.id,
                     creation_id=row.creation_id,
-                    access_key=row.access_key,
+                    main_access_key=row.access_key,
                     kernels=[],
                 )
 
@@ -4625,13 +4638,13 @@ async def search_sessions_with_kernels_and_user(
                     SessionDataForStart(
                         session_id=session_info["id"],
                         creation_id=session_info["creation_id"],
-                        access_key=session_info["access_key"],
+                        main_access_key=session_info["access_key"],
                         session_type=session_info["session_type"],
                         name=session_info["name"],
                         cluster_mode=session_info["cluster_mode"],
                         kernels=session_info["kernels"],
                         environ=session_info.get("environ") or {},
-                        user_uuid=session_info["user_uuid"],
+                        owner_id=session_info["user_uuid"],
                         user_email=user_info.email,
                         user_name=user_info.username,
                     )

src/ai/backend/manager/repositories/user/db_source/db_source.py

@@ -674,10 +674,16 @@ async def delegate_endpoint_ownership(
         target_user_uuid: UUID,
         target_main_access_key: AccessKey,
     ) -> None:
-        """Delegate endpoint ownership to another user."""
+        """Delegate endpoint ownership to another user.
+
+        ``target_main_access_key`` is accepted for caller compatibility but is
+        no longer needed by ``EndpointRow.delegate_endpoint_ownership`` — the
+        downstream ``delegate_ownership`` calls only need ``target_user_uuid``.
+        """
+        del target_main_access_key  # unused, kept for caller signature compatibility
         async with self._db.begin_session() as session:
             await EndpointRow.delegate_endpoint_ownership(
-                session, user_uuid, target_user_uuid, target_main_access_key
+                session, user_uuid, target_user_uuid
             )
 
     async def delete_endpoints(

src/ai/backend/manager/sokovan/scheduling_controller/preparers/preparer.py

@@ -125,7 +125,7 @@ async def prepare(
             id=session_id,
             creation_id=spec.session_creation_id,
             name=spec.session_name,
-            main_access_key=spec.main_access_key,
+            main_access_key=spec.access_key,
             owner_id=spec.user_scope.user_uuid,
             group_id=spec.user_scope.group_id,
             domain_name=spec.user_scope.domain_name,
@@ -255,7 +255,7 @@ async def _prepare_kernels(
                 domain_name=spec.user_scope.domain_name,
                 group_id=spec.user_scope.group_id,
                 owner_id=spec.user_scope.user_uuid,
-                main_access_key=spec.main_access_key,
+                main_access_key=spec.access_key,
                 image=image_info.canonical if image_info else self.DEFAULT_IMAGE_NAME,
                 architecture=image_info.architecture if image_info else self.DEFAULT_ARCHITECTURE,
                 registry=image_info.registry if image_info else self.DEFAULT_REGISTRY,