From c34f07b186c3c14c123251a609510cf2a7372930 Mon Sep 17 00:00:00 2001 From: dadachann <185672915+dadachann@users.noreply.github.com> Date: Fri, 31 Jul 2026 07:25:33 +0000 Subject: [PATCH] feat(cloud): add scoped support admin sessions --- src/langbot/pkg/api/http/authz.py | 5 + src/langbot/pkg/api/http/context.py | 2 + src/langbot/pkg/api/http/controller/group.py | 125 ++++- .../groups/pipelines/websocket_chat.py | 30 +- .../pkg/api/http/controller/groups/user.py | 13 + .../api/http/controller/groups/workspaces.py | 27 +- src/langbot/pkg/api/http/service/user.py | 7 +- src/langbot/pkg/cloud/launch.py | 59 ++- src/langbot/pkg/cloud/support_admin.py | 248 ++++++++++ src/langbot/pkg/core/app.py | 3 + src/langbot/pkg/core/stages/build_app.py | 2 + .../pkg/entity/persistence/support_admin.py | 35 ++ .../versions/0016_support_admin_sessions.py | 88 ++++ src/langbot/pkg/persistence/mgr.py | 1 + src/langbot/pkg/persistence/tenant_uow.py | 1 + .../pkg/platform/sources/websocket_manager.py | 8 +- .../api/test_support_admin_launch.py | 453 ++++++++++++++++++ .../integration/api/test_user_space_oauth.py | 2 +- .../api/service/test_user_service.py | 1 + tests/unit_tests/cloud/test_space_launch.py | 88 ++++ web/src/app/auth/space/callback/page.tsx | 15 +- web/src/app/home/layout.tsx | 6 +- web/src/app/infra/http/BackendClient.ts | 4 +- web/src/app/infra/http/index.ts | 48 ++ web/tests/unit/support-admin-session.test.mjs | 39 ++ 25 files changed, 1285 insertions(+), 25 deletions(-) create mode 100644 src/langbot/pkg/cloud/support_admin.py create mode 100644 src/langbot/pkg/entity/persistence/support_admin.py create mode 100644 src/langbot/pkg/persistence/alembic/versions/0016_support_admin_sessions.py create mode 100644 tests/integration/api/test_support_admin_launch.py create mode 100644 web/tests/unit/support-admin-session.test.mjs diff --git a/src/langbot/pkg/api/http/authz.py b/src/langbot/pkg/api/http/authz.py index 36b796fe7..225e30aa8 100644 --- a/src/langbot/pkg/api/http/authz.py +++ b/src/langbot/pkg/api/http/authz.py @@ -74,6 +74,11 @@ class AuthorizationError(Exception): error_code = 'forbidden' +class AuthenticationDeniedError(AuthorizationError): + status_code = 401 + error_code = 'invalid_authentication' + + class WorkspaceRequiredError(AuthorizationError): status_code = 400 error_code = 'workspace_required' diff --git a/src/langbot/pkg/api/http/context.py b/src/langbot/pkg/api/http/context.py index 7880a1e15..eb5c8ac3b 100644 --- a/src/langbot/pkg/api/http/context.py +++ b/src/langbot/pkg/api/http/context.py @@ -9,6 +9,7 @@ class PrincipalType(enum.StrEnum): ACCOUNT = 'account' API_KEY = 'api_key' + SUPPORT_ADMIN = 'support_admin' SYSTEM = 'system' PUBLIC_BOT = 'public_bot' @@ -19,6 +20,7 @@ class PrincipalContext: principal_type: PrincipalType account_uuid: str | None = None + actor_account_uuid: str | None = None api_key_uuid: str | None = None diff --git a/src/langbot/pkg/api/http/controller/group.py b/src/langbot/pkg/api/http/controller/group.py index 1fc09a0b2..946eec78c 100644 --- a/src/langbot/pkg/api/http/controller/group.py +++ b/src/langbot/pkg/api/http/controller/group.py @@ -15,8 +15,17 @@ from ....workspace.collaboration import MembershipPermissionError, WorkspaceColl from ....workspace.errors import WorkspaceNotFoundError from ....cloud.entitlements import EntitlementUnavailableError from ....core.errors import TaskCapacityError -from ..authz import AuthorizationError, Permission, permissions_for_role, require_permission +from ..authz import ( + AuthenticationDeniedError, + AuthorizationError, + Permission, + PermissionDeniedError, + WorkspaceRequiredError, + permissions_for_role, + require_permission, +) from ..context import PrincipalContext, PrincipalType, RequestContext, WorkspaceContext +from ....cloud.support_admin import SupportAdminSessionError if typing.TYPE_CHECKING: from ....core.app import Application @@ -51,6 +60,17 @@ class AuthType(enum.Enum): USER_TOKEN_OR_API_KEY = 'user-token-or-api-key' +_SUPPORT_ADMIN_DENIED_PERMISSIONS = frozenset( + { + Permission.OWNER_TRANSFER.value, + Permission.MEMBER_VIEW.value, + Permission.MEMBER_INVITE.value, + Permission.MEMBER_UPDATE_ROLE.value, + Permission.MEMBER_REMOVE.value, + } +) + + class RouterGroup(abc.ABC): name: str @@ -95,6 +115,10 @@ class RouterGroup(abc.ABC): return self.http_status(401, -1, 'No valid user token provided') try: + if self._is_support_admin_token(token): + raise AuthenticationDeniedError( + 'Support admin tokens cannot be refreshed or used on account endpoints' + ) account, user_email = await self._authenticate_account(token) # Account-token routes deliberately stop before Workspace # selection. They may bootstrap a selector, but cannot @@ -111,8 +135,13 @@ class RouterGroup(abc.ABC): return self.http_status(401, -1, 'No valid user token provided') try: - account, user_email = await self._authenticate_account(token) - request_context = await self._resolve_account_context(account, auth_type) + request_context = await self._authenticate_support_admin(token, auth_type) + if request_context is not None: + self._require_support_admin_route_allowed(rule, f, permission) + user_email = None + else: + account, user_email = await self._authenticate_account(token) + request_context = await self._resolve_account_context(account, auth_type) if permission is not None: if request_context is None: raise AuthorizationError('Workspace authorization is unavailable') @@ -141,10 +170,20 @@ class RouterGroup(abc.ABC): return self._auth_error_response(e) elif auth_type == AuthType.USER_TOKEN_OR_API_KEY: + token = quart.request.headers.get('Authorization', '').replace('Bearer ', '') + if token and self._is_support_admin_token(token): + try: + request_context = await self._authenticate_support_admin(token, auth_type) + if request_context is None: + raise AuthenticationDeniedError('Invalid support admin token') + self._require_support_admin_route_allowed(rule, f, permission) + if permission is not None: + require_permission(request_context, permission) + self._inject_handler_context(f, kwargs, None, request_context) + except Exception as e: + return self._auth_error_response(e) # Try API key first (check X-API-Key header) - api_key = quart.request.headers.get('X-API-Key', '') - - if api_key: + elif api_key := quart.request.headers.get('X-API-Key', ''): # API key authentication try: request_context = await self._authenticate_api_key(api_key, auth_type) @@ -155,8 +194,6 @@ class RouterGroup(abc.ABC): return self._auth_error_response(e) else: # Try user token authentication (Authorization header) - token = quart.request.headers.get('Authorization', '').replace('Bearer ', '') - if not token: return self.http_status( 401, -1, 'No valid authentication provided (user token or API key required)' @@ -268,10 +305,82 @@ class RouterGroup(abc.ABC): raise ValueError('User not found') return account, account.user + def _is_support_admin_token(self, token: str) -> bool: + service = getattr(self.ap, 'support_admin_session_service', None) + detector = getattr(service, 'is_support_admin_token', None) + return callable(detector) and detector(token) is True + + async def _authenticate_support_admin( + self, + token: str, + auth_type: AuthType, + *, + workspace_uuid: str | None = None, + request_id: str | None = None, + ) -> RequestContext | None: + service = getattr(self.ap, 'support_admin_session_service', None) + detector = getattr(service, 'is_support_admin_token', None) + if service is None or not callable(detector) or detector(token) is not True: + return None + + requested_workspace_uuid = ( + workspace_uuid if workspace_uuid is not None else quart.request.headers.get('X-Workspace-Id') + ) + if not requested_workspace_uuid: + raise WorkspaceRequiredError('Support admin token requires an explicit Workspace selector') + try: + identity = await service.authenticate_token( + token, + requested_workspace_uuid=requested_workspace_uuid, + ) + except SupportAdminSessionError as exc: + raise AuthenticationDeniedError(str(exc)) from exc + + entitlement_revision = await self._resolve_entitlement_revision( + identity.instance_uuid, + identity.workspace_uuid, + ) + request_context = RequestContext( + instance_uuid=identity.instance_uuid, + placement_generation=identity.placement_generation, + request_id=request_id or self.request_id(), + auth_type=auth_type.value, + principal=PrincipalContext( + principal_type=PrincipalType.SUPPORT_ADMIN, + actor_account_uuid=identity.actor_account_uuid, + ), + workspace=WorkspaceContext( + workspace_uuid=identity.workspace_uuid, + membership_uuid=None, + role='owner', + permissions=permissions_for_role('owner') - _SUPPORT_ADMIN_DENIED_PERMISSIONS, + membership_revision=0, + ), + entitlement_revision=entitlement_revision, + ) + quart.g.request_context = request_context + quart.g.workspace_membership = None + return request_context + + @staticmethod + def _require_support_admin_route_allowed( + rule: str, + handler: RouteCallable, + permission: Permission | str | None, + ) -> None: + parameters = inspect.signature(handler).parameters + if rule.startswith('/api/v1/user/') or 'account' in parameters or 'user_email' in parameters: + raise AuthenticationDeniedError('Support admin tokens are not permitted on account endpoints') + permission_value = permission.value if isinstance(permission, Permission) else permission + if permission_value in _SUPPORT_ADMIN_DENIED_PERMISSIONS: + raise PermissionDeniedError(permission_value) + async def _resolve_account_context( self, account: typing.Any, auth_type: AuthType, + *, + token: str | None = None, ) -> RequestContext | None: collaboration_service = getattr(self.ap, 'workspace_collaboration_service', None) account_uuid = getattr(account, 'uuid', None) diff --git a/src/langbot/pkg/api/http/controller/groups/pipelines/websocket_chat.py b/src/langbot/pkg/api/http/controller/groups/pipelines/websocket_chat.py index 565107117..056217348 100644 --- a/src/langbot/pkg/api/http/controller/groups/pipelines/websocket_chat.py +++ b/src/langbot/pkg/api/http/controller/groups/pipelines/websocket_chat.py @@ -97,6 +97,16 @@ class WebSocketChatRouterGroup(group.RouterGroup): if not token or not workspace_uuid: raise ValueError('Authentication is required') + support_context = await self._authenticate_support_admin( + token, + group.AuthType.USER_TOKEN, + workspace_uuid=workspace_uuid, + request_id=quart.websocket.headers.get('X-Request-Id') or str(uuid.uuid4()), + ) + if support_context is not None: + require_permission(support_context, Permission.RUNTIME_OPERATE) + return support_context, token + account, _ = await self._authenticate_account(token) account_uuid = getattr(account, 'uuid', None) collaboration_service = getattr(self.ap, 'workspace_collaboration_service', None) @@ -131,6 +141,23 @@ class WebSocketChatRouterGroup(group.RouterGroup): ) -> RequestContext: """Recheck revocable account, membership, permission, and placement state.""" + if request_context.principal.principal_type == PrincipalType.SUPPORT_ADMIN: + current_context = await self._authenticate_support_admin( + token, + group.AuthType.USER_TOKEN, + workspace_uuid=request_context.workspace_uuid, + request_id=request_context.request_id, + ) + if current_context is None or current_context.principal != request_context.principal: + raise ValueError('WebSocket support admin session changed') + if ( + current_context.instance_uuid != request_context.instance_uuid + or current_context.placement_generation != request_context.placement_generation + ): + raise ValueError('WebSocket authorization changed') + require_permission(current_context, Permission.RUNTIME_OPERATE) + return current_context + account, _ = await self._authenticate_account(token) account_uuid = getattr(account, 'uuid', None) if account_uuid != request_context.account_uuid: @@ -211,6 +238,7 @@ class WebSocketChatRouterGroup(group.RouterGroup): scope=WebSocketScope.from_context(request_context), pipeline_uuid=pipeline_uuid, session_type=session_type, + trigger_principal=request_context.principal, metadata={'user_agent': quart.websocket.headers.get('User-Agent', '')}, send_queue_size=( self.ap.instance_config.data.get('system', {}) @@ -391,7 +419,7 @@ class WebSocketChatRouterGroup(group.RouterGroup): ) elif message_type == 'message': try: - await self._revalidate_websocket_authorization(request_context, token) + request_context = await self._revalidate_websocket_authorization(request_context, token) except Exception: await connection.send_queue.put({'type': 'error', 'message': 'Unauthorized'}) break diff --git a/src/langbot/pkg/api/http/controller/groups/user.py b/src/langbot/pkg/api/http/controller/groups/user.py index 8fa38bf73..8b86ba361 100644 --- a/src/langbot/pkg/api/http/controller/groups/user.py +++ b/src/langbot/pkg/api/http/controller/groups/user.py @@ -396,6 +396,19 @@ class UserRouterGroup(group.RouterGroup): launch_assertion, expected_workspace_uuid=workspace_uuid, ) + if launch.get('launch_mode') == 'support_admin': + token = launch.get('support_admin_token') + if not token: + raise SpaceLaunchError('Support admin launch session was not issued') + return self.success( + data={ + 'token': token, + 'workspace_uuid': launch['workspace_uuid'], + 'principal_type': 'support_admin', + 'actor_account_uuid': launch['actor_account_uuid'], + } + ) + account = await self.ap.user_service.get_user_by_uuid(launch['account_uuid']) if account is None: raise SpaceLaunchError('Launch Account is not projected into Core') diff --git a/src/langbot/pkg/api/http/controller/groups/workspaces.py b/src/langbot/pkg/api/http/controller/groups/workspaces.py index 5c43cf64b..6b6817b07 100644 --- a/src/langbot/pkg/api/http/controller/groups/workspaces.py +++ b/src/langbot/pkg/api/http/controller/groups/workspaces.py @@ -5,7 +5,7 @@ import typing import quart from ...authz import Permission, permissions_for_role -from ...context import RequestContext +from ...context import PrincipalType, RequestContext from ...service.user import AccountExistsLoginRequiredError, ControlPlaneDirectoryRequiredError from .....entity.persistence.workspace import Workspace, WorkspaceInvitation, WorkspaceMembership from .....entity.persistence.workspace import WorkspaceSource @@ -120,9 +120,6 @@ class WorkspacesRouterGroup(group.RouterGroup): @self.route('/current', methods=['GET'], permission=Permission.WORKSPACE_VIEW) async def _(request_context: RequestContext) -> typing.Any: membership = quart.g.workspace_membership - account = await self.ap.user_service.get_user_by_uuid(request_context.account_uuid) - if account is None: - return self.http_status(401, 'invalid_authentication', 'Account not found') workspace = await self.ap.workspace_service.get_workspace(request_context.workspace_uuid) plan_name: str | None = None resolver = getattr(self.ap, 'entitlement_resolver', None) @@ -132,6 +129,28 @@ class WorkspacesRouterGroup(group.RouterGroup): minimum_revision=request_context.entitlement_revision, ) plan_name = entitlement.plan_name + if request_context.principal.principal_type == PrincipalType.SUPPORT_ADMIN: + return self.success( + data={ + 'workspace': _workspace_payload(workspace), + 'membership': { + 'uuid': None, + 'workspace_uuid': request_context.workspace_uuid, + 'account_uuid': None, + 'email': None, + 'role': 'owner', + 'status': 'active', + 'joined_at': None, + 'created_at': None, + }, + 'permissions': sorted(request_context.workspace.permissions), + 'placement_generation': request_context.placement_generation, + 'plan_name': plan_name, + } + ) + account = await self.ap.user_service.get_user_by_uuid(request_context.account_uuid) + if account is None: + return self.http_status(401, 'invalid_authentication', 'Account not found') return self.success( data={ 'workspace': _workspace_payload(workspace), diff --git a/src/langbot/pkg/api/http/service/user.py b/src/langbot/pkg/api/http/service/user.py index 93a3441ad..d1ae53608 100644 --- a/src/langbot/pkg/api/http/service/user.py +++ b/src/langbot/pkg/api/http/service/user.py @@ -400,7 +400,10 @@ class UserService: return await self.generate_jwt_token(user_obj) - async def generate_jwt_token(self, account: user.User | str) -> str: + async def generate_jwt_token( + self, + account: user.User | str, + ) -> str: jwt_secret = self.ap.instance_config.data['system']['jwt']['secret'] jwt_expire = self.ap.instance_config.data['system']['jwt']['expire'] @@ -413,7 +416,7 @@ class UserService: # Lightweight unit-test and bootstrap callers may not have persistence wired. account_obj = None - payload = { + payload: dict[str, typing.Any] = { 'user': user_email, 'iss': self._jwt_identity()[0], 'aud': self._jwt_identity()[1], diff --git a/src/langbot/pkg/cloud/launch.py b/src/langbot/pkg/cloud/launch.py index ac5006870..ddc808beb 100644 --- a/src/langbot/pkg/cloud/launch.py +++ b/src/langbot/pkg/cloud/launch.py @@ -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') diff --git a/src/langbot/pkg/cloud/support_admin.py b/src/langbot/pkg/cloud/support_admin.py new file mode 100644 index 000000000..bd0607a20 --- /dev/null +++ b/src/langbot/pkg/cloud/support_admin.py @@ -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) diff --git a/src/langbot/pkg/core/app.py b/src/langbot/pkg/core/app.py index bb8d76a2f..bfd01387d 100644 --- a/src/langbot/pkg/core/app.py +++ b/src/langbot/pkg/core/app.py @@ -51,6 +51,7 @@ from ..workspace import collaboration as workspace_collaboration_module from ..workspace import invitation_delivery as invitation_delivery_module from ..cloud import bootstrap as cloud_bootstrap_module from ..cloud import launch as cloud_launch_module +from ..cloud import support_admin as cloud_support_admin_module from ..cloud import directory_projection as cloud_directory_projection_module from ..cloud import entitlements as cloud_entitlements_module from ..api.http.context import ExecutionContext, PrincipalContext, PrincipalType @@ -136,6 +137,8 @@ class Application: space_launch_service: cloud_launch_module.SpaceLaunchService = None + support_admin_session_service: cloud_support_admin_module.SupportAdminSessionService = None + deployment: cloud_bootstrap_module.OpenSourceDeployment | cloud_bootstrap_module.VerifiedCloudDeployment = None deployment_admission: cloud_bootstrap_module.DeploymentAdmissionGuard = None diff --git a/src/langbot/pkg/core/stages/build_app.py b/src/langbot/pkg/core/stages/build_app.py index 813a6c4e6..b08fdfaf7 100644 --- a/src/langbot/pkg/core/stages/build_app.py +++ b/src/langbot/pkg/core/stages/build_app.py @@ -42,6 +42,7 @@ from ...workspace import collaboration as workspace_collaboration_module from ...workspace import invitation_delivery as invitation_delivery_module from ...cloud import bootstrap as cloud_bootstrap from ...cloud import launch as cloud_launch_module +from ...cloud import support_admin as cloud_support_admin_module from ...cloud.directory import directory_projection_limits_from_config from ...cloud.directory_projection import DirectoryProjectionService from ...cloud.entitlements import EntitlementResolver @@ -180,6 +181,7 @@ class BuildAppStage(stage.BootingStage): workspace_service_inst, ) ap.invitation_delivery_service = invitation_delivery_module.InvitationDeliveryService(ap) + ap.support_admin_session_service = cloud_support_admin_module.SupportAdminSessionService(ap) ap.space_launch_service = cloud_launch_module.SpaceLaunchService(ap) user_service_inst = user_service.UserService(ap) diff --git a/src/langbot/pkg/entity/persistence/support_admin.py b/src/langbot/pkg/entity/persistence/support_admin.py new file mode 100644 index 000000000..3575d0135 --- /dev/null +++ b/src/langbot/pkg/entity/persistence/support_admin.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +import sqlalchemy + +from .base import Base + + +class SupportAdminTemporarySession(Base): + """Temporary support-admin Workspace access session.""" + + __tablename__ = 'support_admin_temporary_sessions' + + grant_jti_hash = sqlalchemy.Column(sqlalchemy.String(64), primary_key=True) + workspace_uuid = sqlalchemy.Column( + sqlalchemy.String(36), + sqlalchemy.ForeignKey('workspaces.uuid', ondelete='CASCADE'), + nullable=False, + ) + actor_account_uuid = sqlalchemy.Column(sqlalchemy.String(36), nullable=False) + issued_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False) + expires_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False) + revoked_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True) + last_used_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True) + + __table_args__ = ( + sqlalchemy.Index( + 'ix_support_admin_sessions_workspace_expiry', + 'workspace_uuid', + 'expires_at', + ), + sqlalchemy.CheckConstraint( + 'length(grant_jti_hash) = 64', + name='ck_support_admin_sessions_grant_jti_hash', + ), + ) diff --git a/src/langbot/pkg/persistence/alembic/versions/0016_support_admin_sessions.py b/src/langbot/pkg/persistence/alembic/versions/0016_support_admin_sessions.py new file mode 100644 index 000000000..28a1b34d6 --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/versions/0016_support_admin_sessions.py @@ -0,0 +1,88 @@ +"""add temporary support-admin sessions + +Revision ID: 0016_support_admin_sessions +Revises: 0015_cloud_core_collab +Create Date: 2026-07-31 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + + +revision = '0016_support_admin_sessions' +down_revision = '0015_cloud_core_collab' +branch_labels = None +depends_on = None + + +_TABLE_NAME = 'support_admin_temporary_sessions' +_POLICY_NAME = 'langbot_workspace_isolation' +_TENANT_SETTING = 'langbot.workspace_uuid' + + +def _setting(name: str) -> str: + return f"NULLIF(current_setting('{name}', true), '')" + + +def _quote(conn: sa.Connection, identifier: str) -> str: + return conn.dialect.identifier_preparer.quote(identifier) + + +def upgrade() -> None: + conn = op.get_bind() + existing_tables = set(sa.inspect(conn).get_table_names()) + if _TABLE_NAME not in existing_tables: + op.create_table( + _TABLE_NAME, + sa.Column('grant_jti_hash', sa.String(64), nullable=False), + sa.Column( + 'workspace_uuid', + sa.String(36), + sa.ForeignKey('workspaces.uuid', ondelete='CASCADE'), + nullable=False, + ), + sa.Column('actor_account_uuid', sa.String(36), nullable=False), + sa.Column('issued_at', sa.DateTime(), nullable=False), + sa.Column('expires_at', sa.DateTime(), nullable=False), + sa.Column('revoked_at', sa.DateTime(), nullable=True), + sa.Column('last_used_at', sa.DateTime(), nullable=True), + sa.CheckConstraint( + 'length(grant_jti_hash) = 64', + name='ck_support_admin_sessions_grant_jti_hash', + ), + sa.PrimaryKeyConstraint('grant_jti_hash'), + ) + op.create_index( + 'ix_support_admin_sessions_workspace_expiry', + _TABLE_NAME, + ['workspace_uuid', 'expires_at'], + unique=False, + ) + + if conn.dialect.name != 'postgresql': + return + + table = _quote(conn, _TABLE_NAME) + policy = _quote(conn, _POLICY_NAME) + expression = f'workspace_uuid::text = {_setting(_TENANT_SETTING)}' + op.execute(sa.text(f'ALTER TABLE {table} ENABLE ROW LEVEL SECURITY')) + op.execute(sa.text(f'ALTER TABLE {table} FORCE ROW LEVEL SECURITY')) + op.execute(sa.text(f'DROP POLICY IF EXISTS {policy} ON {table}')) + op.execute( + sa.text( + f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR ALL TO PUBLIC ' + f'USING ({expression}) WITH CHECK ({expression})' + ) + ) + + +def downgrade() -> None: + conn = op.get_bind() + if conn.dialect.name == 'postgresql': + table = _quote(conn, _TABLE_NAME) + policy = _quote(conn, _POLICY_NAME) + op.execute(sa.text(f'DROP POLICY IF EXISTS {policy} ON {table}')) + op.drop_index('ix_support_admin_sessions_workspace_expiry', table_name=_TABLE_NAME) + op.drop_table(_TABLE_NAME) diff --git a/src/langbot/pkg/persistence/mgr.py b/src/langbot/pkg/persistence/mgr.py index 9565dfdd0..121f19de5 100644 --- a/src/langbot/pkg/persistence/mgr.py +++ b/src/langbot/pkg/persistence/mgr.py @@ -54,6 +54,7 @@ _ALEMBIC_TENANT_TABLES = { 'workspace_memberships', 'workspace_invitations', 'workspace_execution_states', + 'support_admin_temporary_sessions', 'workspace_metadata', 'api_keys', 'bots', diff --git a/src/langbot/pkg/persistence/tenant_uow.py b/src/langbot/pkg/persistence/tenant_uow.py index 27db2d64a..53be402fa 100644 --- a/src/langbot/pkg/persistence/tenant_uow.py +++ b/src/langbot/pkg/persistence/tenant_uow.py @@ -43,6 +43,7 @@ TENANT_TABLE_COLUMNS: dict[str, str] = { 'workspace_memberships': 'workspace_uuid', 'workspace_invitations': 'workspace_uuid', 'workspace_execution_states': 'workspace_uuid', + 'support_admin_temporary_sessions': 'workspace_uuid', 'workspace_metadata': 'workspace_uuid', 'api_keys': 'workspace_uuid', 'bots': 'workspace_uuid', diff --git a/src/langbot/pkg/platform/sources/websocket_manager.py b/src/langbot/pkg/platform/sources/websocket_manager.py index 99ef98bba..03f94e517 100644 --- a/src/langbot/pkg/platform/sources/websocket_manager.py +++ b/src/langbot/pkg/platform/sources/websocket_manager.py @@ -9,7 +9,7 @@ from datetime import datetime import pydantic -from ...api.http.context import ExecutionContext +from ...api.http.context import ExecutionContext, PrincipalContext logger = logging.getLogger(__name__) _SESSION_FILTER_UNSET = object() @@ -95,6 +95,9 @@ class WebSocketConnection(pydantic.BaseModel): metadata: dict = pydantic.Field(default_factory=dict) """连接元数据（可存储额外信息）""" + trigger_principal: PrincipalContext | None = None + """Authenticated principal that opened this dashboard connection.""" + @property def scope(self) -> WebSocketScope: return WebSocketScope( @@ -112,6 +115,7 @@ class WebSocketConnection(pydantic.BaseModel): workspace_uuid=self.workspace_uuid, placement_generation=self.placement_generation, pipeline_uuid=self.pipeline_uuid, + trigger_principal=self.trigger_principal, ) @@ -138,6 +142,7 @@ class WebSocketConnectionManager: pipeline_uuid: str, session_type: str, metadata: dict | None = None, + trigger_principal: PrincipalContext | None = None, session_id: str | None = None, send_queue_size: int = _DEFAULT_SEND_QUEUE_SIZE, max_connections: int = 1024, @@ -174,6 +179,7 @@ class WebSocketConnectionManager: session_id=session_id, websocket=websocket, metadata=metadata or {}, + trigger_principal=trigger_principal, send_queue=asyncio.Queue(maxsize=send_queue_size), ) diff --git a/tests/integration/api/test_support_admin_launch.py b/tests/integration/api/test_support_admin_launch.py new file mode 100644 index 000000000..94647a7f1 --- /dev/null +++ b/tests/integration/api/test_support_admin_launch.py @@ -0,0 +1,453 @@ +from __future__ import annotations + +import base64 +import datetime +import json +import logging +import time +import uuid +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +import sqlalchemy +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey +from quart import Quart +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +from langbot.pkg.api.http.authz import Permission +from langbot.pkg.api.http.context import PrincipalType, RequestContext +from langbot.pkg.api.http.controller import group +from langbot.pkg.api.http.controller.groups.pipelines.websocket_chat import WebSocketChatRouterGroup +from langbot.pkg.api.http.controller.groups.user import UserRouterGroup +from langbot.pkg.cloud.launch import SpaceLaunchError, SpaceLaunchService +from langbot.pkg.cloud.support_admin import SupportAdminSessionService +from langbot.pkg.entity.persistence.base import Base +from langbot.pkg.entity.persistence.support_admin import SupportAdminTemporarySession +from langbot.pkg.entity.persistence.user import User +from langbot.pkg.entity.persistence.workspace import ( + Workspace, + WorkspaceExecutionState, + WorkspaceMembership, +) +from langbot.pkg.workspace.service import WorkspaceService + + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio] + + +INSTANCE_UUID = 'instance-support-admin' +WORKSPACE_UUID = '10000000-0000-4000-8000-000000000001' +OTHER_WORKSPACE_UUID = '10000000-0000-4000-8000-000000000002' +ACTOR_ACCOUNT_UUID = '20000000-0000-4000-8000-000000000001' +KEY_ID = 'support-admin-key-1' + + +def _base64url(raw: bytes) -> str: + return base64.urlsafe_b64encode(raw).rstrip(b'=').decode('ascii') + + +def _sign(private_key: Ed25519PrivateKey, claims: dict, *, key_id: str = KEY_ID) -> str: + header = {'alg': 'EdDSA', 'kid': key_id, 'typ': 'langbot-control-plane+jwt'} + encoded_header = _base64url(json.dumps(header, separators=(',', ':')).encode('utf-8')) + encoded_claims = _base64url(json.dumps(claims, separators=(',', ':')).encode('utf-8')) + signing_input = f'{encoded_header}.{encoded_claims}' + return f'{signing_input}.{_base64url(private_key.sign(signing_input.encode("ascii")))}' + + +def _admin_claims(*, now: int, jti: str | None = None, workspace_uuid: str = WORKSPACE_UUID) -> dict: + return { + 'iss': 'langbot-space', + 'aud': 'langbot-cloud-runtime', + 'sub': f'langbot-instance:{INSTANCE_UUID}', + 'jti': jti or str(uuid.uuid4()), + 'iat': now, + 'nbf': now - 5, + 'exp': now + 90, + 'instance_uuid': INSTANCE_UUID, + 'kind': 'workspace.support_admin_launch', + 'payload': { + 'workspace_uuid': workspace_uuid, + 'launch_mode': 'support_admin', + 'principal_type': 'support_admin', + 'actor_account_uuid': ACTOR_ACCOUNT_UUID, + 'effective_role': 'owner', + }, + } + + +@group.group_class('support_admin_probe', '/api/v1/support-admin-probe') +class SupportAdminProbeGroup(group.RouterGroup): + async def initialize(self) -> None: + @self.route('/user-token', auth_type=group.AuthType.USER_TOKEN, permission=Permission.WORKSPACE_VIEW) + async def _(request_context: RequestContext) -> str: + return self.success(data=_context_payload(request_context)) + + @self.route( + '/member-operation', + auth_type=group.AuthType.USER_TOKEN, + permission=Permission.MEMBER_VIEW, + ) + async def member_operation(request_context: RequestContext) -> str: + return self.success(data=_context_payload(request_context)) + + @self.route( + '/user-token-or-api-key', + auth_type=group.AuthType.USER_TOKEN_OR_API_KEY, + permission=Permission.WORKSPACE_VIEW, + ) + async def _(request_context: RequestContext) -> str: + return self.success(data=_context_payload(request_context)) + + +def _context_payload(request_context: RequestContext) -> dict: + return { + 'principal_type': request_context.principal.principal_type.value, + 'actor_account_uuid': request_context.principal.actor_account_uuid, + 'account_uuid': request_context.principal.account_uuid, + 'role': request_context.workspace.role, + 'membership_uuid': request_context.workspace.membership_uuid, + 'permissions': sorted(request_context.workspace.permissions), + } + + +class _TenantUow: + def __init__(self, engine): + self._engine = engine + self.session = None + self._transaction = None + + async def __aenter__(self): + session_factory = async_sessionmaker(self._engine, expire_on_commit=False) + self.session = session_factory() + self._transaction = await self.session.begin() + return self + + async def __aexit__(self, exc_type, exc, traceback): + try: + if exc_type is None: + await self._transaction.commit() + else: + await self._transaction.rollback() + finally: + await self.session.close() + + +class _TenantScope: + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback): + return False + + +class _PersistenceManager: + def __init__(self, engine): + self._engine = engine + self.mode = SimpleNamespace(value='oss_compat') + + def get_db_engine(self): + return self._engine + + def tenant_uow(self, workspace_uuid: str): + del workspace_uuid + return _TenantUow(self._engine) + + def tenant_scope(self, workspace_uuid: str): + del workspace_uuid + return _TenantScope() + + +@pytest.fixture +async def support_admin_api(tmp_path): + private_key = Ed25519PrivateKey.generate() + public_key = private_key.public_key().public_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PublicFormat.Raw, + ) + engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "support-admin.db"}') + async with engine.begin() as connection: + await connection.run_sync( + Base.metadata.create_all, + tables=[ + User.__table__, + Workspace.__table__, + WorkspaceExecutionState.__table__, + WorkspaceMembership.__table__, + SupportAdminTemporarySession.__table__, + ], + ) + for workspace_uuid, slug in ( + (WORKSPACE_UUID, 'support-admin-a'), + (OTHER_WORKSPACE_UUID, 'support-admin-b'), + ): + await connection.execute( + sqlalchemy.insert(Workspace).values( + uuid=workspace_uuid, + instance_uuid=INSTANCE_UUID, + name=slug, + slug=slug, + type='team', + status='active', + source='cloud_projection', + projection_revision=1, + ) + ) + await connection.execute( + sqlalchemy.insert(WorkspaceExecutionState).values( + workspace_uuid=workspace_uuid, + instance_uuid=INSTANCE_UUID, + active_generation=1, + state='active', + write_fenced=False, + source='cloud', + desired_state_revision=1, + ) + ) + + app = SimpleNamespace() + app.persistence_mgr = _PersistenceManager(engine) + app.instance_config = SimpleNamespace( + data={ + 'system': { + 'jwt': {'secret': 'support-admin-secret', 'expire': 3600}, + 'websocket_retention': {}, + }, + 'space': { + 'launch': { + 'control_plane_public_key': _base64url(public_key), + } + }, + 'api': {'global_api_key': ''}, + } + ) + app.logger = logging.getLogger('support-admin-test') + app.deployment = SimpleNamespace(mode='cloud', multi_workspace_enabled=True, verification_key_id=KEY_ID) + app.directory_projection_service = SimpleNamespace(require_ready=lambda: None) + app.workspace_service = WorkspaceService(app, instance_uuid=INSTANCE_UUID) + app.entitlement_resolver = SimpleNamespace( + instance_uuid=INSTANCE_UUID, + resolve=AsyncMock(return_value=SimpleNamespace(entitlement_revision=7)), + ) + app.support_admin_session_service = SupportAdminSessionService(app) + app.space_launch_service = SpaceLaunchService(app) + app.user_service = SimpleNamespace() + app.user_service.get_authenticated_account = AsyncMock(side_effect=AssertionError('normal account auth used')) + app.user_service.verify_jwt_token = AsyncMock(side_effect=AssertionError('normal token verification used')) + app.user_service.get_user_by_email = AsyncMock(side_effect=AssertionError('user lookup used')) + app.apikey_service = SimpleNamespace() + app.apikey_service.authenticate_api_key = AsyncMock( + return_value=SimpleNamespace( + instance_uuid=INSTANCE_UUID, + workspace_uuid=OTHER_WORKSPACE_UUID, + placement_generation=1, + api_key_uuid='api-key', + permissions=frozenset(permission.value for permission in Permission), + ) + ) + + quart_app = Quart(__name__) + await UserRouterGroup(app, quart_app).initialize() + await SupportAdminProbeGroup(app, quart_app).initialize() + + yield app, quart_app.test_client(), engine, private_key + await engine.dispose() + + +async def _issue_support_token(app, private_key: Ed25519PrivateKey, *, jti: str | None = None) -> dict[str, str]: + launch = await app.space_launch_service.consume_assertion( + _sign(private_key, _admin_claims(now=int(time.time()), jti=jti)), + expected_workspace_uuid=WORKSPACE_UUID, + ) + return launch + + +def _auth(token: str, workspace_uuid: str = WORKSPACE_UUID) -> dict[str, str]: + return {'Authorization': f'Bearer {token}', 'X-Workspace-Id': workspace_uuid} + + +async def test_support_admin_membership_only_routes_are_denied(support_admin_api): + app, client, _engine, private_key = support_admin_api + launch = await _issue_support_token(app, private_key) + + response = await client.get( + '/api/v1/support-admin-probe/member-operation', + headers=_auth(launch['support_admin_token']), + ) + + assert response.status_code == 403 + assert (await response.get_json())['code'] == 'permission_denied' + + +async def test_support_admin_check_token_is_rejected(support_admin_api): + app, client, _engine, private_key = support_admin_api + launch = await _issue_support_token(app, private_key) + + response = await client.get('/api/v1/user/check-token', headers=_auth(launch['support_admin_token'])) + + assert response.status_code == 401 + assert (await response.get_json())['code'] == 'invalid_authentication' + + +async def test_support_admin_cross_workspace_denied_for_user_token_and_or_api_key(support_admin_api): + app, client, _engine, private_key = support_admin_api + launch = await _issue_support_token(app, private_key) + + missing_selector = await client.get( + '/api/v1/support-admin-probe/user-token', + headers={'Authorization': f'Bearer {launch["support_admin_token"]}'}, + ) + user_response = await client.get( + '/api/v1/support-admin-probe/user-token', + headers=_auth(launch['support_admin_token'], OTHER_WORKSPACE_UUID), + ) + either_response = await client.get( + '/api/v1/support-admin-probe/user-token-or-api-key', + headers={ + **_auth(launch['support_admin_token'], OTHER_WORKSPACE_UUID), + 'X-API-Key': 'valid-api-key', + }, + ) + + assert missing_selector.status_code == 400 + assert user_response.status_code == 401 + assert either_response.status_code == 401 + app.apikey_service.authenticate_api_key.assert_not_awaited() + + +async def test_support_admin_request_context_has_actor_owner_and_no_membership(support_admin_api): + app, client, engine, private_key = support_admin_api + before_count = await _membership_count(engine) + launch = await _issue_support_token(app, private_key) + + response = await client.get( + '/api/v1/support-admin-probe/user-token', + headers=_auth(launch['support_admin_token']), + ) + + assert response.status_code == 200 + data = (await response.get_json())['data'] + permissions = set(data.pop('permissions')) + assert Permission.WORKSPACE_VIEW.value in permissions + assert Permission.RESOURCE_MANAGE.value in permissions + assert not permissions.intersection( + { + Permission.OWNER_TRANSFER.value, + Permission.MEMBER_VIEW.value, + Permission.MEMBER_INVITE.value, + Permission.MEMBER_UPDATE_ROLE.value, + Permission.MEMBER_REMOVE.value, + } + ) + assert data == { + 'principal_type': PrincipalType.SUPPORT_ADMIN.value, + 'actor_account_uuid': ACTOR_ACCOUNT_UUID, + 'account_uuid': None, + 'role': 'owner', + 'membership_uuid': None, + } + assert await _membership_count(engine) == before_count + + +async def test_support_admin_missing_workspace_is_controlled_launch_failure(support_admin_api): + app, _client, engine, private_key = support_admin_api + async with engine.begin() as connection: + await connection.execute( + sqlalchemy.delete(WorkspaceExecutionState).where(WorkspaceExecutionState.workspace_uuid == WORKSPACE_UUID) + ) + + with pytest.raises(SpaceLaunchError, match='unavailable'): + await _issue_support_token(app, private_key) + + +async def test_support_admin_launch_replay_is_durable_across_service_instances(support_admin_api): + app, _client, _engine, private_key = support_admin_api + jti = str(uuid.uuid4()) + + await _issue_support_token(app, private_key, jti=jti) + second_service = SpaceLaunchService(app) + + with pytest.raises(SpaceLaunchError, match='already been consumed'): + await second_service.consume_assertion( + _sign(private_key, _admin_claims(now=int(time.time()), jti=jti)), + expected_workspace_uuid=WORKSPACE_UUID, + ) + + +async def test_support_admin_persisted_expiry_and_revocation_are_enforced(support_admin_api): + app, client, engine, private_key = support_admin_api + launch = await _issue_support_token(app, private_key) + token = launch['support_admin_token'] + + async with engine.begin() as connection: + await connection.execute( + sqlalchemy.update(SupportAdminTemporarySession) + .where(SupportAdminTemporarySession.grant_jti_hash == launch['grant_jti_hash']) + .values(expires_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None) - datetime.timedelta(minutes=1)) + ) + expired = await client.get('/api/v1/support-admin-probe/user-token', headers=_auth(token)) + assert expired.status_code == 401 + + second = await _issue_support_token(app, private_key) + async with engine.begin() as connection: + await connection.execute( + sqlalchemy.update(SupportAdminTemporarySession) + .where(SupportAdminTemporarySession.grant_jti_hash == second['grant_jti_hash']) + .values(revoked_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None)) + ) + revoked = await client.get('/api/v1/support-admin-probe/user-token', headers=_auth(second['support_admin_token'])) + assert revoked.status_code == 401 + + +async def test_support_admin_websocket_preserves_actor_and_revalidates(support_admin_api): + app, _client, _engine, private_key = support_admin_api + launch = await _issue_support_token(app, private_key) + captured_contexts = [] + + class Adapter: + async def handle_websocket_message(self, connection, data): + del data + captured_contexts.append(connection.execution_context) + await connection.send_queue.put({'type': 'handled'}) + connection.is_active = False + + app.pipeline_service = SimpleNamespace(get_pipeline=AsyncMock(return_value=SimpleNamespace(uuid='pipeline-1'))) + app.platform_mgr = SimpleNamespace( + get_websocket_proxy_bot=AsyncMock(return_value=SimpleNamespace(adapter=Adapter())) + ) + + quart_app = Quart(__name__) + await WebSocketChatRouterGroup(app, quart_app).initialize() + + async with quart_app.test_client().websocket('/api/v1/pipelines/pipeline-1/ws/connect') as websocket: + await websocket.send( + json.dumps( + { + 'type': 'authenticate', + 'token': launch['support_admin_token'], + 'workspace_uuid': WORKSPACE_UUID, + } + ) + ) + connected = json.loads(await websocket.receive()) + assert connected['type'] == 'connected' + await websocket.send(json.dumps({'type': 'message', 'message': [{'type': 'text', 'text': 'hi'}]})) + handled = json.loads(await websocket.receive()) + assert handled['type'] == 'handled' + + assert captured_contexts + principal = captured_contexts[0].trigger_principal + assert principal is not None + assert principal.principal_type == PrincipalType.SUPPORT_ADMIN + assert principal.actor_account_uuid == ACTOR_ACCOUNT_UUID + + +async def _membership_count(engine) -> int: + async with engine.connect() as connection: + return int( + await connection.scalar( + sqlalchemy.select(sqlalchemy.func.count()).select_from(WorkspaceMembership), + ) + or 0 + ) diff --git a/tests/integration/api/test_user_space_oauth.py b/tests/integration/api/test_user_space_oauth.py index 17569707c..deb1e93e2 100644 --- a/tests/integration/api/test_user_space_oauth.py +++ b/tests/integration/api/test_user_space_oauth.py @@ -270,7 +270,7 @@ async def test_space_credits_are_resolved_from_workspace_owner(space_oauth_api): response = await client.get( '/api/v1/user/space-credits', - headers={'Authorization': 'Bearer account-token', 'X-Workspace-UUID': WORKSPACE_UUID}, + headers={'Authorization': 'Bearer account-token', 'X-Workspace-Id': WORKSPACE_UUID}, ) assert response.status_code == 200 diff --git a/tests/unit_tests/api/service/test_user_service.py b/tests/unit_tests/api/service/test_user_service.py index 96d501b95..a3479725e 100644 --- a/tests/unit_tests/api/service/test_user_service.py +++ b/tests/unit_tests/api/service/test_user_service.py @@ -418,6 +418,7 @@ class TestUserServiceGenerateJwtToken: assert token is not None + class TestUserServiceVerifyJwtToken: """Tests for verify_jwt_token method.""" diff --git a/tests/unit_tests/cloud/test_space_launch.py b/tests/unit_tests/cloud/test_space_launch.py index 14b6b7f17..be9c7c11a 100644 --- a/tests/unit_tests/cloud/test_space_launch.py +++ b/tests/unit_tests/cloud/test_space_launch.py @@ -11,6 +11,7 @@ from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey from langbot.pkg.cloud.launch import SpaceLaunchError, SpaceLaunchService +from langbot.pkg.cloud.support_admin import SupportAdminReplayError pytestmark = pytest.mark.asyncio @@ -57,9 +58,21 @@ def _service(private_key: Ed25519PrivateKey, *, now: int) -> SpaceLaunchService: encoding=serialization.Encoding.Raw, format=serialization.PublicFormat.Raw, ) + consumed: set[str] = set() + + class DurableSupportAdminService: + async def consume_launch_grant(self, **kwargs): + grant_hash = kwargs['grant_jti_hash'] + if grant_hash in consumed: + raise SupportAdminReplayError('already consumed') + consumed.add(grant_hash) + return SimpleNamespace(token='support-admin-token') + app = SimpleNamespace( deployment=SimpleNamespace(multi_workspace_enabled=True, verification_key_id=KEY_ID), workspace_service=SimpleNamespace(instance_uuid=INSTANCE_UUID), + logger=SimpleNamespace(info=lambda *args, **kwargs: None), + support_admin_session_service=DurableSupportAdminService(), instance_config=SimpleNamespace( data={ 'space': { @@ -86,6 +99,81 @@ async def test_consumes_valid_workspace_launch_assertion_once(): await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID) + + +async def test_consumes_admin_owner_launch_once_and_validates_claims(): + private_key = Ed25519PrivateKey.generate() + now = int(time.time()) + service = _service(private_key, now=now) + claims = _claims(now=now) + claims['kind'] = 'workspace.support_admin_launch' + claims['payload'].update( + { + 'launch_mode': 'support_admin', + 'principal_type': 'support_admin', + 'actor_account_uuid': '33333333-3333-4333-8333-333333333333', + 'effective_role': 'owner', + } + ) + claims['payload'].pop('account_uuid') + token = _sign(private_key, claims) + + launch = await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID) + + assert launch == { + 'workspace_uuid': WORKSPACE_UUID, + 'launch_mode': 'support_admin', + 'actor_account_uuid': '33333333-3333-4333-8333-333333333333', + 'effective_role': 'owner', + 'grant_jti_hash': launch['grant_jti_hash'], + 'support_admin_token': 'support-admin-token', + } + with pytest.raises(SpaceLaunchError, match='already been consumed'): + await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID) + + invalid = _claims(now=now) + invalid['kind'] = 'workspace.support_admin_launch' + invalid['payload'].update( + { + 'launch_mode': 'support_admin', + 'principal_type': 'support_admin', + 'actor_account_uuid': '33333333-3333-4333-8333-333333333333', + 'effective_role': 'member', + } + ) + invalid['payload'].pop('account_uuid') + with pytest.raises(SpaceLaunchError, match='effective role'): + await service.consume_assertion(_sign(private_key, invalid), expected_workspace_uuid=WORKSPACE_UUID) + + too_long = _claims(now=now) + too_long['kind'] = 'workspace.support_admin_launch' + too_long['exp'] = now + 91 + too_long['payload'].update( + { + 'launch_mode': 'support_admin', + 'principal_type': 'support_admin', + 'actor_account_uuid': '33333333-3333-4333-8333-333333333333', + 'effective_role': 'owner', + } + ) + too_long['payload'].pop('account_uuid') + with pytest.raises(SpaceLaunchError, match='lifetime exceeds 90 seconds'): + await service.consume_assertion(_sign(private_key, too_long), expected_workspace_uuid=WORKSPACE_UUID) + + impersonating = _claims(now=now) + impersonating['kind'] = 'workspace.support_admin_launch' + impersonating['payload'].update( + { + 'launch_mode': 'support_admin', + 'principal_type': 'support_admin', + 'actor_account_uuid': '33333333-3333-4333-8333-333333333333', + 'effective_role': 'owner', + } + ) + with pytest.raises(SpaceLaunchError, match='customer Account'): + await service.consume_assertion(_sign(private_key, impersonating), expected_workspace_uuid=WORKSPACE_UUID) + + async def test_replay_cache_does_not_scan_all_live_assertions(monkeypatch): private_key = Ed25519PrivateKey.generate() now = int(time.time()) diff --git a/web/src/app/auth/space/callback/page.tsx b/web/src/app/auth/space/callback/page.tsx index 33f3a4846..1028de6ac 100644 --- a/web/src/app/auth/space/callback/page.tsx +++ b/web/src/app/auth/space/callback/page.tsx @@ -3,6 +3,7 @@ import { useNavigate, useSearchParams } from 'react-router-dom'; import { httpClient } from '@/app/infra/http/HttpClient'; import { beginAuthenticatedSession, + beginSupportAdminSession, bootstrapWorkspaceSession, getPendingInvitationToken, } from '@/app/infra/http'; @@ -27,8 +28,10 @@ import langbotIcon from '@/app/assets/langbot-logo.webp'; type SpaceOAuthLoginResult = { token: string; - user: string; + user?: string; workspace_uuid?: string; + principal_type?: 'account' | 'support_admin'; + actor_account_uuid?: string; }; const pendingSpaceOAuthLogins = new Map< @@ -94,6 +97,16 @@ function SpaceOAuthCallbackContent() { return; } + if (response.principal_type === 'support_admin') { + if (!response.workspace_uuid) { + throw new Error('Support admin launch did not include a Workspace'); + } + beginSupportAdminSession(response.token, response.workspace_uuid); + await bootstrapWorkspaceSession(); + navigate('/home', { replace: true }); + return; + } + beginAuthenticatedSession(response.token, response.user); if (getPendingInvitationToken()) { navigate('/invitations/accept', { replace: true }); diff --git a/web/src/app/home/layout.tsx b/web/src/app/home/layout.tsx index fc91f7805..49b3969a2 100644 --- a/web/src/app/home/layout.tsx +++ b/web/src/app/home/layout.tsx @@ -17,6 +17,7 @@ import { bootstrapWorkspaceSession, systemInfo, initializeSystemInfo, + isSupportAdminSession, useCurrentWorkspace, } from '@/app/infra/http'; import { useNavigate, useLocation } from 'react-router-dom'; @@ -156,7 +157,10 @@ export default function HomeLayout({ // selected Workspace's wizard state. useEffect(() => { if (!identityReady) return; - if (systemInfo.wizard_status === 'none') { + if ( + systemInfo?.wizard_status === 'none' && + !isSupportAdminSession() + ) { navigate('/wizard', { replace: true }); } }, [identityReady, navigate]); diff --git a/web/src/app/infra/http/BackendClient.ts b/web/src/app/infra/http/BackendClient.ts index e1733eaff..91b9ab0fe 100644 --- a/web/src/app/infra/http/BackendClient.ts +++ b/web/src/app/infra/http/BackendClient.ts @@ -1327,8 +1327,10 @@ export class BackendClient extends BaseHttpClient { launchAssertion?: string, ): Promise<{ token: string; - user: string; + user?: string; workspace_uuid?: string; + principal_type?: 'account' | 'support_admin'; + actor_account_uuid?: string; }> { const response = await this.instance.post( '/api/v1/user/space/callback', diff --git a/web/src/app/infra/http/index.ts b/web/src/app/infra/http/index.ts index 6223144bb..0014bc272 100644 --- a/web/src/app/infra/http/index.ts +++ b/web/src/app/infra/http/index.ts @@ -217,10 +217,34 @@ export function beginAuthenticatedSession( if (typeof window === 'undefined') return; localStorage.removeItem('token'); localStorage.removeItem('userEmail'); + localStorage.removeItem('authPrincipalType'); localStorage.setItem('token', token); if (userEmail) localStorage.setItem('userEmail', userEmail); } +export function beginSupportAdminSession( + token: string, + workspaceUuid: string, +): void { + userInfo = null; + clearWorkspaceSelection(); + clearWorkspaceBootstrapSnapshot(); + + if (typeof window === 'undefined') return; + localStorage.removeItem('token'); + localStorage.removeItem('userEmail'); + localStorage.setItem('token', token); + localStorage.setItem('authPrincipalType', 'support_admin'); + setActiveWorkspaceUuid(workspaceUuid); +} + +export function isSupportAdminSession(): boolean { + return ( + typeof window !== 'undefined' && + localStorage.getItem('authPrincipalType') === 'support_admin' + ); +} + async function initializeSelectedWorkspace( workspaceUuid: string, workspaces: WorkspaceBootstrapEntry[], @@ -252,6 +276,27 @@ async function initializeSelectedWorkspace( export async function bootstrapWorkspaceSession( options: WorkspaceBootstrapOptions = {}, ): Promise { + if (isSupportAdminSession()) { + const selectedWorkspaceUuid = getActiveWorkspaceUuid(); + if (!selectedWorkspaceUuid) { + throw new Error('Support admin session is missing its Workspace scope'); + } + if ( + options.preferredWorkspaceUuid && + options.preferredWorkspaceUuid !== selectedWorkspaceUuid + ) { + throw new Error('Support admin session cannot change Workspace scope'); + } + await initializeWorkspaceInfo(); + const workspace = getCurrentWorkspaceSnapshot(); + if (!workspace || workspace.workspace.uuid !== selectedWorkspaceUuid) { + clearWorkspaceSelection(); + throw new Error('Support admin Workspace scope could not be initialized'); + } + clearWorkspaceBootstrapSnapshot(); + return { status: 'ready', workspace, workspaces: [] }; + } + if (options.resetSelection) { clearWorkspaceSelection(); clearWorkspaceBootstrapSnapshot(); @@ -339,6 +384,9 @@ export const clearUserInfo = (): void => { userInfo = null; clearWorkspaceSelection(); clearWorkspaceBootstrapSnapshot(); + if (typeof window !== 'undefined') { + localStorage.removeItem('authPrincipalType'); + } }; export { diff --git a/web/tests/unit/support-admin-session.test.mjs b/web/tests/unit/support-admin-session.test.mjs new file mode 100644 index 000000000..b687457ed --- /dev/null +++ b/web/tests/unit/support-admin-session.test.mjs @@ -0,0 +1,39 @@ +import assert from 'node:assert/strict'; +import fs from 'node:fs'; +import path from 'node:path'; +import test from 'node:test'; +import { fileURLToPath } from 'node:url'; + +const root = path.resolve(path.dirname(fileURLToPath(import.meta.url)), '../..'); +const read = (file) => fs.readFileSync(path.join(root, file), 'utf8'); + +test('support-admin launch stores a scoped principal instead of starting an Account session', () => { + const callback = read('src/app/auth/space/callback/page.tsx'); + assert.match(callback, /response\.principal_type === 'support_admin'/); + assert.match(callback, /beginSupportAdminSession\(response\.token, response\.workspace_uuid\)/); +}); + +test('support-admin workspace bootstrap never calls Account bootstrap', () => { + const source = read('src/app/infra/http/index.ts'); + assert.match(source, /export function beginSupportAdminSession/); + assert.match( + source, + /export function isSupportAdminSession\(\): boolean/, + ); + const supportBranch = source.indexOf('if (isSupportAdminSession())'); + const accountBootstrap = source.indexOf( + 'backendClient.getWorkspaceBootstrap()', + ); + assert.ok(supportBranch >= 0); + assert.ok(accountBootstrap > supportBranch); + assert.match( + source.slice(supportBranch, accountBootstrap), + /initializeWorkspaceInfo\([\s\S]*status: 'ready'/, + ); + assert.match( + source, + /localStorage\.setItem\('authPrincipalType', 'support_admin'\)/, + ); + const homeLayout = read('src/app/home/layout.tsx'); + assert.match(homeLayout, /!isSupportAdminSession\(\)/); +});