From 6f4200ee36d3e745228315f87de04881db13caca Mon Sep 17 00:00:00 2001 From: dadachann <185672915+dadachann@users.noreply.github.com> Date: Fri, 31 Jul 2026 04:26:15 +0000 Subject: [PATCH] feat(cloud): support scoped admin owner sessions --- src/langbot/pkg/api/http/controller/group.py | 18 ++++++- .../groups/pipelines/websocket_chat.py | 11 ++++ .../pkg/api/http/controller/groups/user.py | 14 ++++- src/langbot/pkg/api/http/service/user.py | 42 ++++++++++++++- src/langbot/pkg/cloud/launch.py | 26 ++++++++- .../api/service/test_user_service.py | 47 ++++++++++++++++ tests/unit_tests/cloud/test_space_launch.py | 53 +++++++++++++++++++ 7 files changed, 205 insertions(+), 6 deletions(-) diff --git a/src/langbot/pkg/api/http/controller/group.py b/src/langbot/pkg/api/http/controller/group.py index 1fc09a0b2..7afce94c9 100644 --- a/src/langbot/pkg/api/http/controller/group.py +++ b/src/langbot/pkg/api/http/controller/group.py @@ -164,7 +164,7 @@ class RouterGroup(abc.ABC): try: account, user_email = await self._authenticate_account(token) - request_context = await self._resolve_account_context(account, auth_type) + request_context = await self._resolve_account_context(account, auth_type, token=token) if permission is not None: if request_context is None: raise AuthorizationError('Workspace authorization is unavailable') @@ -272,6 +272,8 @@ class RouterGroup(abc.ABC): 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) @@ -280,6 +282,20 @@ class RouterGroup(abc.ABC): return None requested_workspace_uuid = quart.request.headers.get('X-Workspace-Id') + scope_resolver = getattr(self.ap.user_service, 'get_admin_owner_scope', None) + admin_owner_scope = typing.cast( + dict[str, str] | None, + scope_resolver(token) if token and callable(scope_resolver) else None, + ) + if admin_owner_scope is not None: + scoped_workspace_uuid = admin_owner_scope['workspace_uuid'] + if requested_workspace_uuid and requested_workspace_uuid != scoped_workspace_uuid: + self.ap.logger.warning( + 'cloud_admin_owner_scope_rejected actor_account_uuid=%s target_workspace_uuid=%s requested_workspace_uuid=%s', + admin_owner_scope['actor_account_uuid'], scoped_workspace_uuid, requested_workspace_uuid, + ) + raise MembershipPermissionError('Admin owner session is scoped to another Workspace') + requested_workspace_uuid = scoped_workspace_uuid access = await collaboration_service.resolve_account_workspace(account_uuid, requested_workspace_uuid) entitlement_revision = await self._resolve_entitlement_revision( access.execution.instance_uuid, 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..53f84a95b 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 @@ -98,6 +98,17 @@ class WebSocketChatRouterGroup(group.RouterGroup): raise ValueError('Authentication is required') account, _ = await self._authenticate_account(token) + scope_resolver = getattr(self.ap.user_service, 'get_admin_owner_scope', None) + admin_owner_scope = typing.cast( + dict[str, str] | None, + scope_resolver(token) if callable(scope_resolver) else None, + ) + if admin_owner_scope is not None and workspace_uuid != admin_owner_scope['workspace_uuid']: + self.ap.logger.warning( + 'cloud_admin_owner_scope_rejected actor_account_uuid=%s target_workspace_uuid=%s requested_workspace_uuid=%s', + admin_owner_scope['actor_account_uuid'], admin_owner_scope['workspace_uuid'], workspace_uuid, + ) + raise ValueError('Admin owner session is scoped to another Workspace') account_uuid = getattr(account, 'uuid', None) collaboration_service = getattr(self.ap, 'workspace_collaboration_service', None) if not isinstance(account_uuid, str) or collaboration_service is None: diff --git a/src/langbot/pkg/api/http/controller/groups/user.py b/src/langbot/pkg/api/http/controller/groups/user.py index 8fa38bf73..6bd055275 100644 --- a/src/langbot/pkg/api/http/controller/groups/user.py +++ b/src/langbot/pkg/api/http/controller/groups/user.py @@ -404,7 +404,19 @@ class UserRouterGroup(group.RouterGroup): account.uuid, launch['workspace_uuid'], ) - token = await self.ap.user_service.generate_jwt_token(account) + admin_owner_scope = None + if launch.get('launch_mode') == 'admin_owner': + if access.membership.role != 'owner': + raise SpaceLaunchError('Admin launch principal is not an active Workspace owner') + admin_owner_scope = { + 'actor_account_uuid': launch['actor_account_uuid'], + 'workspace_uuid': launch['workspace_uuid'], + 'effective_role': 'owner', + } + token = await self.ap.user_service.generate_jwt_token( + account, + admin_owner_scope=admin_owner_scope, + ) return self.success( data={ 'token': token, diff --git a/src/langbot/pkg/api/http/service/user.py b/src/langbot/pkg/api/http/service/user.py index 93a3441ad..fa1a68ad9 100644 --- a/src/langbot/pkg/api/http/service/user.py +++ b/src/langbot/pkg/api/http/service/user.py @@ -400,9 +400,21 @@ 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, + *, + admin_owner_scope: dict[str, str] | None = None, + ) -> str: jwt_secret = self.ap.instance_config.data['system']['jwt']['secret'] jwt_expire = self.ap.instance_config.data['system']['jwt']['expire'] + if admin_owner_scope is not None: + expected = {'actor_account_uuid', 'workspace_uuid', 'effective_role'} + if set(admin_owner_scope) != expected or admin_owner_scope.get('effective_role') != 'owner': + raise ValueError('Invalid admin owner token scope') + if not admin_owner_scope.get('actor_account_uuid') or not admin_owner_scope.get('workspace_uuid'): + raise ValueError('Invalid admin owner token scope') + jwt_expire = min(int(jwt_expire), 300) account_obj: user.User | None = account if not isinstance(account, str) and hasattr(account, 'user') else None user_email = account_obj.user if account_obj is not None else account @@ -413,7 +425,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], @@ -427,9 +439,35 @@ class UserService: 'account_revision': account_obj.projection_revision, } ) + if admin_owner_scope is not None: + payload['admin_owner_scope'] = dict(admin_owner_scope) return jwt.encode(payload, jwt_secret, algorithm='HS256') + def get_admin_owner_scope(self, token: str) -> dict[str, str] | None: + jwt_secret = self.ap.instance_config.data['system']['jwt']['secret'] + issuer, audience = self._jwt_identity() + payload = jwt.decode( + token, + jwt_secret, + algorithms=['HS256'], + issuer=issuer, + audience=audience, + options={'require': ['exp', 'iss', 'aud']}, + ) + scope = payload.get('admin_owner_scope') + if scope is None: + return None + if not isinstance(scope, dict) or set(scope) != {'actor_account_uuid', 'workspace_uuid', 'effective_role'}: + raise ValueError('Invalid admin owner token scope') + if scope.get('effective_role') != 'owner': + raise ValueError('Invalid admin owner token scope') + actor = scope.get('actor_account_uuid') + workspace = scope.get('workspace_uuid') + if not isinstance(actor, str) or not actor or not isinstance(workspace, str) or not workspace: + raise ValueError('Invalid admin owner token scope') + return {'actor_account_uuid': actor, 'workspace_uuid': workspace, 'effective_role': 'owner'} + async def verify_jwt_token(self, token: str) -> str: account = await self.get_authenticated_account(token, allow_unresolved_legacy=True) if isinstance(account, str): diff --git a/src/langbot/pkg/cloud/launch.py b/src/langbot/pkg/cloud/launch.py index ac5006870..f69de42f3 100644 --- a/src/langbot/pkg/cloud/launch.py +++ b/src/langbot/pkg/cloud/launch.py @@ -127,11 +127,33 @@ class SpaceLaunchService: 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 { + launch_mode = payload.get('launch_mode') + result = { 'account_uuid': account_uuid, 'workspace_uuid': workspace_uuid, } + if launch_mode is not None: + if launch_mode != 'admin_owner': + raise SpaceLaunchError('Launch assertion mode is unsupported') + 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') + result.update({ + 'launch_mode': 'admin_owner', + 'actor_account_uuid': actor_account_uuid, + 'effective_role': 'owner', + }) + await self._consume_jti(_required_string(claims, 'jti'), _required_int(claims, 'exp', minimum=1)) + if launch_mode == 'admin_owner': + self.ap.logger.info( + 'cloud_admin_owner_launch_consumed actor_account_uuid=%s workspace_uuid=%s owner_account_uuid=%s', + result['actor_account_uuid'], workspace_uuid, account_uuid, + ) + return result def _verify_assertion(self, token: str) -> dict[str, typing.Any]: if not getattr(getattr(self.ap, 'deployment', None), 'multi_workspace_enabled', False): diff --git a/tests/unit_tests/api/service/test_user_service.py b/tests/unit_tests/api/service/test_user_service.py index 96d501b95..458e0644e 100644 --- a/tests/unit_tests/api/service/test_user_service.py +++ b/tests/unit_tests/api/service/test_user_service.py @@ -418,6 +418,53 @@ class TestUserServiceGenerateJwtToken: assert token is not None + + async def test_admin_owner_token_is_scoped_and_capped_at_five_minutes(self): + ap = SimpleNamespace() + ap.instance_config = SimpleNamespace() + ap.instance_config.data = {'system': {'jwt': {'secret': 'test_secret', 'expire': 7200}}} + service = UserService(ap) + before = datetime.datetime.now(datetime.timezone.utc).timestamp() + + token = await service.generate_jwt_token( + 'owner@example.com', + admin_owner_scope={ + 'actor_account_uuid': 'admin-uuid', + 'workspace_uuid': 'workspace-uuid', + 'effective_role': 'owner', + }, + ) + payload = jwt.decode( + token, + 'test_secret', + algorithms=['HS256'], + options={'verify_aud': False}, + ) + + assert payload['exp'] <= before + 301 + assert service.get_admin_owner_scope(token) == { + 'actor_account_uuid': 'admin-uuid', + 'workspace_uuid': 'workspace-uuid', + 'effective_role': 'owner', + } + + async def test_admin_owner_token_rejects_invalid_scope(self): + ap = SimpleNamespace() + ap.instance_config = SimpleNamespace() + ap.instance_config.data = {'system': {'jwt': {'secret': 'test_secret', 'expire': 7200}}} + service = UserService(ap) + + with pytest.raises(ValueError, match='Invalid admin owner token scope'): + await service.generate_jwt_token( + 'owner@example.com', + admin_owner_scope={ + 'actor_account_uuid': 'admin-uuid', + 'workspace_uuid': 'workspace-uuid', + 'effective_role': 'member', + }, + ) + + 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..443ec986b 100644 --- a/tests/unit_tests/cloud/test_space_launch.py +++ b/tests/unit_tests/cloud/test_space_launch.py @@ -60,6 +60,7 @@ def _service(private_key: Ed25519PrivateKey, *, now: int) -> SpaceLaunchService: 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), instance_config=SimpleNamespace( data={ 'space': { @@ -86,6 +87,58 @@ 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['payload'].update( + { + 'launch_mode': 'admin_owner', + 'actor_account_uuid': '33333333-3333-4333-8333-333333333333', + 'effective_role': 'owner', + } + ) + token = _sign(private_key, claims) + + launch = await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID) + + assert launch == { + 'account_uuid': ACCOUNT_UUID, + 'workspace_uuid': WORKSPACE_UUID, + 'launch_mode': 'admin_owner', + 'actor_account_uuid': '33333333-3333-4333-8333-333333333333', + 'effective_role': 'owner', + } + with pytest.raises(SpaceLaunchError, match='already been consumed'): + await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID) + + invalid = _claims(now=now) + invalid['payload'].update( + { + 'launch_mode': 'admin_owner', + 'actor_account_uuid': '33333333-3333-4333-8333-333333333333', + 'effective_role': 'member', + } + ) + 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['exp'] = now + 91 + too_long['payload'].update( + { + 'launch_mode': 'admin_owner', + 'actor_account_uuid': '33333333-3333-4333-8333-333333333333', + 'effective_role': 'owner', + } + ) + with pytest.raises(SpaceLaunchError, match='lifetime exceeds 90 seconds'): + await service.consume_assertion(_sign(private_key, too_long), 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())