feat(cloud): add scoped support admin sessions (#2369)

* feat(cloud): add scoped support admin sessions * style(web): format support admin session changes * fix(cloud): isolate support adapter sessions * fix(cloud): authenticate plugin assets and report workspace resources --------- Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
2026-08-09 04:40:57 +00:00 · 2026-07-31 17:41:55 +08:00
parent 9df021eb8f
commit 404e3466d9
37 changed files with 1703 additions and 102 deletions
@@ -15,12 +15,15 @@ from cryptography.exceptions import InvalidSignature
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey

+from .support_admin import SupportAdminReplayError, SupportAdminSessionError, hash_grant_jti
+
 if typing.TYPE_CHECKING:
    from ..core.app import Application


 CONTROL_PLANE_TYP = 'langbot-control-plane+jwt'
 LAUNCH_KIND = 'workspace.launch'
+SUPPORT_ADMIN_LAUNCH_KIND = 'workspace.support_admin_launch'
 EXPECTED_ISSUER = 'langbot-space'
 EXPECTED_AUDIENCE = 'langbot-cloud-runtime'
 _CONSUMED_JTI_MAX_ENTRIES = 4096
@@ -123,15 +126,60 @@ class SpaceLaunchService:
        payload = claims.get('payload')
        if not isinstance(payload, dict):
            raise SpaceLaunchError('Launch assertion payload must be a JSON object')
-        account_uuid = _required_string(payload, 'account_uuid')
+        kind = _required_string(claims, 'kind')
        workspace_uuid = _required_string(payload, 'workspace_uuid')
        if expected_workspace_uuid is not None and workspace_uuid != expected_workspace_uuid:
            raise SpaceLaunchError('Launch assertion targets another Workspace')
-        await self._consume_jti(_required_string(claims, 'jti'), _required_int(claims, 'exp', minimum=1))
-        return {
+        if kind == SUPPORT_ADMIN_LAUNCH_KIND:
+            if 'account_uuid' in payload:
+                raise SpaceLaunchError('Admin launch assertion must not identify a customer Account')
+            if payload.get('launch_mode') != 'support_admin' or payload.get('principal_type') != 'support_admin':
+                raise SpaceLaunchError('Admin launch principal must be support_admin')
+            actor_account_uuid = _required_string(payload, 'actor_account_uuid')
+            if _required_string(payload, 'effective_role') != 'owner':
+                raise SpaceLaunchError('Admin launch effective role must be owner')
+            issued_at = _required_int(claims, 'iat')
+            expires_at = _required_int(claims, 'exp', minimum=1)
+            if expires_at - issued_at > 90:
+                raise SpaceLaunchError('Admin launch assertion lifetime exceeds 90 seconds')
+            grant_jti_hash = hash_grant_jti(_required_string(claims, 'jti'))
+            result = {
+                'workspace_uuid': workspace_uuid,
+                'launch_mode': 'support_admin',
+                'actor_account_uuid': actor_account_uuid,
+                'effective_role': 'owner',
+                'grant_jti_hash': grant_jti_hash,
+            }
+            support_service = getattr(self.ap, 'support_admin_session_service', None)
+            if support_service is None or not callable(getattr(support_service, 'consume_launch_grant', None)):
+                raise SpaceLaunchError('Durable support admin session service is unavailable')
+            try:
+                support_session = await support_service.consume_launch_grant(
+                    grant_jti_hash=grant_jti_hash,
+                    workspace_uuid=workspace_uuid,
+                    actor_account_uuid=actor_account_uuid,
+                )
+            except SupportAdminReplayError as exc:
+                raise SpaceLaunchError('Launch assertion has already been consumed') from exc
+            except SupportAdminSessionError as exc:
+                raise SpaceLaunchError(str(exc)) from exc
+            result['support_admin_token'] = support_session.token
+            self.ap.logger.info(
+                'cloud_support_admin_launch_consumed actor_account_uuid=%s workspace_uuid=%s',
+                result['actor_account_uuid'],
+                workspace_uuid,
+            )
+            return result
+
+        if payload.get('launch_mode') is not None:
+            raise SpaceLaunchError('Launch assertion mode is unsupported')
+        account_uuid = _required_string(payload, 'account_uuid')
+        result = {
            'account_uuid': account_uuid,
            'workspace_uuid': workspace_uuid,
        }
+        await self._consume_jti(_required_string(claims, 'jti'), _required_int(claims, 'exp', minimum=1))
+        return result

    def _verify_assertion(self, token: str) -> dict[str, typing.Any]:
        if not getattr(getattr(self.ap, 'deployment', None), 'multi_workspace_enabled', False):
@@ -169,8 +217,9 @@ class SpaceLaunchService:
            raise SpaceLaunchError('Launch assertion subject targets another instance')
        if _required_string(claims, 'instance_uuid') != instance_uuid:
            raise SpaceLaunchError('Launch assertion instance UUID does not match this Core')
-        if _required_string(claims, 'kind') != LAUNCH_KIND:
-            raise SpaceLaunchError('Launch assertion kind is not workspace.launch')
+        kind = _required_string(claims, 'kind')
+        if kind not in {LAUNCH_KIND, SUPPORT_ADMIN_LAUNCH_KIND}:
+            raise SpaceLaunchError('Launch assertion kind is not supported')

        issued_at = _required_int(claims, 'iat')
        not_before = _required_int(claims, 'nbf')
@@ -0,0 +1,248 @@
+from __future__ import annotations
+
+import dataclasses
+import datetime
+import hashlib
+import re
+import time
+import typing
+
+import jwt
+from sqlalchemy.exc import IntegrityError
+
+from ..entity.persistence.support_admin import SupportAdminTemporarySession
+from ..workspace.errors import WorkspaceError
+
+if typing.TYPE_CHECKING:
+    from ..core.app import Application
+
+
+SUPPORT_ADMIN_TOKEN_TYP = 'langbot-support-admin+jwt'
+SUPPORT_ADMIN_TOKEN_KIND = 'support_admin.session'
+SUPPORT_ADMIN_EFFECTIVE_ROLE = 'owner'
+SUPPORT_ADMIN_MAX_TOKEN_SECONDS = 300
+_SHA256_HEX = re.compile(r'^[0-9a-f]{64}$')
+
+
+class SupportAdminSessionError(ValueError):
+    """Raised when a support-admin session or token is not admissible."""
+
+
+class SupportAdminReplayError(SupportAdminSessionError):
+    """Raised when a launch grant JTI has already been consumed."""
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class IssuedSupportAdminSession:
+    token: str
+    grant_jti_hash: str
+    workspace_uuid: str
+    actor_account_uuid: str
+    issued_at: datetime.datetime
+    expires_at: datetime.datetime
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class SupportAdminSessionIdentity:
+    grant_jti_hash: str
+    workspace_uuid: str
+    actor_account_uuid: str
+    instance_uuid: str
+    placement_generation: int
+
+
+def hash_grant_jti(jti: str) -> str:
+    return hashlib.sha256(jti.encode('utf-8')).hexdigest()
+
+
+class SupportAdminSessionService:
+    """Issue and validate temporary Workspace-scoped support-admin sessions."""
+
+    def __init__(
+        self,
+        ap: Application,
+        *,
+        wall_time: typing.Callable[[], float] = time.time,
+    ) -> None:
+        self.ap = ap
+        self._wall_time = wall_time
+
+    async def consume_launch_grant(
+        self,
+        *,
+        grant_jti_hash: str,
+        workspace_uuid: str,
+        actor_account_uuid: str,
+    ) -> IssuedSupportAdminSession:
+        self._validate_grant_hash(grant_jti_hash)
+        if not workspace_uuid or not actor_account_uuid:
+            raise SupportAdminSessionError('Support admin session requires an actor and Workspace')
+
+        issued_at = self._utcnow()
+        expires_at = issued_at + datetime.timedelta(seconds=SUPPORT_ADMIN_MAX_TOKEN_SECONDS)
+        tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
+        if not callable(tenant_uow):
+            raise SupportAdminSessionError('Support admin sessions require tenant persistence')
+
+        try:
+            async with tenant_uow(workspace_uuid) as uow:
+                await self.ap.workspace_service.get_execution_binding(workspace_uuid, session=uow.session)
+                uow.session.add(
+                    SupportAdminTemporarySession(
+                        grant_jti_hash=grant_jti_hash,
+                        workspace_uuid=workspace_uuid,
+                        actor_account_uuid=actor_account_uuid,
+                        issued_at=issued_at,
+                        expires_at=expires_at,
+                    )
+                )
+                await uow.session.flush()
+        except IntegrityError as exc:
+            raise SupportAdminReplayError('Launch assertion has already been consumed') from exc
+        except WorkspaceError as exc:
+            raise SupportAdminSessionError('Workspace is unavailable for support access') from exc
+
+        return IssuedSupportAdminSession(
+            token=self._encode_token(
+                grant_jti_hash=grant_jti_hash,
+                workspace_uuid=workspace_uuid,
+                actor_account_uuid=actor_account_uuid,
+                issued_at=issued_at,
+                expires_at=expires_at,
+            ),
+            grant_jti_hash=grant_jti_hash,
+            workspace_uuid=workspace_uuid,
+            actor_account_uuid=actor_account_uuid,
+            issued_at=issued_at,
+            expires_at=expires_at,
+        )
+
+    def is_support_admin_token(self, token: str) -> bool:
+        """Return True only for compact JWTs marked as support-admin tokens."""
+
+        if not isinstance(token, str) or token.count('.') != 2:
+            return False
+        try:
+            header = jwt.get_unverified_header(token)
+        except jwt.PyJWTError:
+            return False
+        if header.get('typ') == SUPPORT_ADMIN_TOKEN_TYP:
+            return True
+        try:
+            payload = jwt.decode(token, options={'verify_signature': False})
+        except jwt.PyJWTError:
+            return False
+        return payload.get('kind') == SUPPORT_ADMIN_TOKEN_KIND
+
+    async def authenticate_token(
+        self,
+        token: str,
+        *,
+        requested_workspace_uuid: str | None,
+    ) -> SupportAdminSessionIdentity:
+        if not self.is_support_admin_token(token):
+            raise SupportAdminSessionError('Not a support admin token')
+        workspace_uuid = (requested_workspace_uuid or '').strip()
+        if not workspace_uuid:
+            raise SupportAdminSessionError('Support admin token requires an explicit Workspace selector')
+
+        jwt_secret = self.ap.instance_config.data['system']['jwt']['secret']
+        try:
+            payload = jwt.decode(
+                token,
+                jwt_secret,
+                algorithms=['HS256'],
+                issuer='langbot-core',
+                audience=self._audience(workspace_uuid),
+                options={'require': ['exp', 'iat', 'nbf', 'iss', 'aud']},
+            )
+        except jwt.PyJWTError as exc:
+            raise SupportAdminSessionError('Invalid support admin token') from exc
+        self._validate_payload(payload, workspace_uuid)
+        grant_jti_hash = payload['grant_jti_hash']
+        actor_account_uuid = payload['actor_account_uuid']
+
+        tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
+        if not callable(tenant_uow):
+            raise SupportAdminSessionError('Support admin sessions require tenant persistence')
+
+        now = self._utcnow()
+        async with tenant_uow(workspace_uuid) as uow:
+            session = await uow.session.get(SupportAdminTemporarySession, grant_jti_hash)
+            if (
+                session is None
+                or session.workspace_uuid != workspace_uuid
+                or session.actor_account_uuid != actor_account_uuid
+                or session.revoked_at is not None
+                or session.expires_at <= now
+            ):
+                raise SupportAdminSessionError('Support admin session is inactive')
+            binding = await self.ap.workspace_service.get_execution_binding(workspace_uuid, session=uow.session)
+            session.last_used_at = now
+            await uow.session.flush()
+
+        return SupportAdminSessionIdentity(
+            grant_jti_hash=grant_jti_hash,
+            workspace_uuid=workspace_uuid,
+            actor_account_uuid=actor_account_uuid,
+            instance_uuid=binding.instance_uuid,
+            placement_generation=binding.placement_generation,
+        )
+
+    async def revoke_session(self, grant_jti_hash: str, workspace_uuid: str) -> None:
+        self._validate_grant_hash(grant_jti_hash)
+        now = self._utcnow()
+        async with self.ap.persistence_mgr.tenant_uow(workspace_uuid) as uow:
+            row = await uow.session.get(SupportAdminTemporarySession, grant_jti_hash)
+            if row is not None and row.revoked_at is None:
+                row.revoked_at = now
+
+    def _encode_token(
+        self,
+        *,
+        grant_jti_hash: str,
+        workspace_uuid: str,
+        actor_account_uuid: str,
+        issued_at: datetime.datetime,
+        expires_at: datetime.datetime,
+    ) -> str:
+        jwt_secret = self.ap.instance_config.data['system']['jwt']['secret']
+        payload: dict[str, typing.Any] = {
+            'kind': SUPPORT_ADMIN_TOKEN_KIND,
+            'iss': 'langbot-core',
+            'aud': self._audience(workspace_uuid),
+            'sub': f'support-admin:{actor_account_uuid}',
+            'iat': issued_at,
+            'nbf': issued_at,
+            'exp': expires_at,
+            'actor_account_uuid': actor_account_uuid,
+            'workspace_uuid': workspace_uuid,
+            'effective_role': SUPPORT_ADMIN_EFFECTIVE_ROLE,
+            'grant_jti_hash': grant_jti_hash,
+        }
+        return jwt.encode(payload, jwt_secret, algorithm='HS256', headers={'typ': SUPPORT_ADMIN_TOKEN_TYP})
+
+    def _validate_payload(self, payload: dict[str, typing.Any], workspace_uuid: str) -> None:
+        if payload.get('kind') != SUPPORT_ADMIN_TOKEN_KIND:
+            raise SupportAdminSessionError('Invalid support admin token kind')
+        if payload.get('workspace_uuid') != workspace_uuid:
+            raise SupportAdminSessionError('Support admin session is scoped to another Workspace')
+        if payload.get('effective_role') != SUPPORT_ADMIN_EFFECTIVE_ROLE:
+            raise SupportAdminSessionError('Invalid support admin token role')
+        actor_account_uuid = payload.get('actor_account_uuid')
+        if not isinstance(actor_account_uuid, str) or not actor_account_uuid.strip():
+            raise SupportAdminSessionError('Invalid support admin actor')
+        grant_jti_hash = payload.get('grant_jti_hash')
+        if not isinstance(grant_jti_hash, str) or not _SHA256_HEX.match(grant_jti_hash):
+            raise SupportAdminSessionError('Invalid support admin grant')
+
+    def _audience(self, workspace_uuid: str) -> str:
+        return f'langbot-support-admin:{self.ap.workspace_service.instance_uuid}:{workspace_uuid}'
+
+    @staticmethod
+    def _validate_grant_hash(grant_jti_hash: str) -> None:
+        if not _SHA256_HEX.match(grant_jti_hash):
+            raise SupportAdminSessionError('Invalid support admin grant')
+
+    def _utcnow(self) -> datetime.datetime:
+        return datetime.datetime.fromtimestamp(self._wall_time(), datetime.UTC).replace(tzinfo=None)