From b44b8f474d4bdc5f99a85a5994b5e8bb4412d5bd Mon Sep 17 00:00:00 2001 From: Hyu Date: Thu, 3 Sep 2026 23:14:44 +0800 Subject: [PATCH] fix(cloud): provision login workspace just in time (#2505) * fix(cloud): provision login workspace just in time * fix(oauth): send callback URI during code exchange * fix(oauth): preserve callback URI through browser exchange * fix(oauth): negotiate redirect-bound codes --------- Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com> --- .../pkg/api/http/controller/groups/user.py | 91 +++---- src/langbot/pkg/api/http/service/space.py | 10 +- src/langbot/pkg/api/http/service/user.py | 5 +- src/langbot/pkg/cloud/directory_projection.py | 85 ++++++- .../integration/api/test_user_space_oauth.py | 233 ++++++++++++------ .../api/service/test_space_service.py | 11 +- .../cloud/test_directory_projection.py | 84 ++++++- web/src/app/auth/space/callback/page.tsx | 18 +- web/src/app/infra/http/BackendClient.ts | 15 +- web/src/app/login/page.tsx | 8 +- .../unit/cloud-new-account-entry.test.mjs | 12 +- 11 files changed, 407 insertions(+), 165 deletions(-) diff --git a/src/langbot/pkg/api/http/controller/groups/user.py b/src/langbot/pkg/api/http/controller/groups/user.py index af37172db..03c84bf47 100644 --- a/src/langbot/pkg/api/http/controller/groups/user.py +++ b/src/langbot/pkg/api/http/controller/groups/user.py @@ -9,7 +9,6 @@ from .. import group from .....entity.errors import account as account_errors from ...context import RequestContext from .....cloud.launch import SpaceLaunchError -from .....workspace.errors import WorkspaceNotFoundError from ...service.user import ControlPlaneDirectoryRequiredError, PublicRegistrationClosedError @@ -144,13 +143,6 @@ class UserRouterGroup(group.RouterGroup): try: redirect_uri = self._validate_space_redirect_uri(redirect_uri, bind=False) launch_workspace_uuid = quart.request.args.get('launch_workspace_uuid') - cloud_entry = quart.request.args.get('cloud_entry') == '1' - if ( - cloud_entry - and not launch_workspace_uuid - and getattr(getattr(self.ap, 'deployment', None), 'mode', 'oss') == 'cloud' - ): - return self.success(data={'authorize_url': self.ap.space_service.get_cloud_entry_url()}) if launch_workspace_uuid: if not getattr(getattr(self.ap, 'deployment', None), 'multi_workspace_enabled', False): return self.fail(1, 'Space launch requires Cloud mode') @@ -194,6 +186,9 @@ class UserRouterGroup(group.RouterGroup): json_data = await quart.request.json code = json_data.get('code') state = json_data.get('state') + redirect_uri = json_data.get('redirect_uri') or ( + quart.request.url_root.rstrip('/') + '/auth/space/callback' + ) launch_assertion = json_data.get('launch_assertion') workspace_uuid = json_data.get('workspace_uuid') @@ -207,8 +202,11 @@ class UserRouterGroup(group.RouterGroup): return self.fail(1, 'Missing authorization code') if not state: return self.fail(1, 'Missing state parameter') + if not str(code).startswith('v4_'): + return self.fail(1, 'Unsupported Space OAuth code contract') try: + redirect_uri = self._validate_space_redirect_uri(str(redirect_uri), bind=False) consumed_state = await self.ap.user_service.consume_space_oauth_state_details(state, 'login') # Exchange code for tokens launch_workspace_uuid = consumed_state.launch_workspace_uuid @@ -226,24 +224,36 @@ class UserRouterGroup(group.RouterGroup): code, workspace_uuids, workspace_created_ats, + redirect_uri=redirect_uri, ) access_token = token_data.get('access_token') refresh_token = token_data.get('refresh_token') expires_in = token_data.get('expires_in', 0) + cloud_workspace_uuid = token_data.get('cloud_workspace_uuid') if not access_token: return self.fail(1, 'Failed to get access token from Space') - # Authenticate and create/update local user + cloud_mode = getattr(getattr(self.ap, 'deployment', None), 'mode', 'oss') == 'cloud' + if cloud_mode and launch_workspace_uuid and launch_workspace_uuid != cloud_workspace_uuid: + return self.fail(1, 'Space OAuth Workspace binding mismatch') + target_workspace_uuid = launch_workspace_uuid or cloud_workspace_uuid + if cloud_mode: + if not target_workspace_uuid: + return self.fail(1, 'Space OAuth response is missing the Cloud Workspace binding') + await self.ap.directory_projection_service.reconcile_workspaces((target_workspace_uuid,)) + + # Authenticate only after the signed, exact Workspace delta has + # established the Account and membership runtime shadow rows. jwt_token, user_obj = await self.ap.user_service.authenticate_space_user( access_token, refresh_token, expires_in ) - if launch_workspace_uuid: + if target_workspace_uuid: try: access = await self.ap.workspace_collaboration_service.resolve_account_workspace( user_obj.uuid, - launch_workspace_uuid, + target_workspace_uuid, ) except Exception: self.ap.logger.warning('Rejected Space OAuth launch for unauthorized Workspace') @@ -375,12 +385,17 @@ class UserRouterGroup(group.RouterGroup): json_data = await quart.request.json code = json_data.get('code') state = json_data.get('state') + redirect_uri = json_data.get('redirect_uri') or ( + quart.request.url_root.rstrip('/') + '/auth/space/callback?mode=bind' + ) if not code: return self.http_status(400, -1, 'Missing authorization code') if not state: return self.http_status(400, -1, 'Missing state parameter') + if not str(code).startswith('v4_'): + return self.http_status(400, -1, 'Unsupported Space OAuth code contract') try: user_obj = await self.ap.user_service.consume_space_oauth_state(state, 'bind') @@ -393,7 +408,10 @@ class UserRouterGroup(group.RouterGroup): return self.http_status(400, -1, 'Only local accounts can bind to Space') try: - updated_user = await self.ap.user_service.bind_space_account(user_obj.user, code) + redirect_uri = self._validate_space_redirect_uri(str(redirect_uri), bind=True) + updated_user = await self.ap.user_service.bind_space_account( + user_obj.user, code, redirect_uri=redirect_uri + ) jwt_token = await self.ap.user_service.generate_jwt_token(updated_user) return self.success( data={ @@ -436,49 +454,18 @@ class UserRouterGroup(group.RouterGroup): } ) - account = await self.ap.user_service.get_user_by_uuid(launch['account_uuid']) projection_service = self.ap.directory_projection_service - access = None - # A first Cloud launch creates the personal Workspace immediately - # before redirecting here. Pull a bounded number of signed event - # pages until both the Account and its target Workspace membership - # are visible instead of rejecting during the background-sync window. - for attempt in range(4): - if account is not None: - self.ap.user_service._require_active_account(account) - try: - access = await self.ap.workspace_collaboration_service.resolve_account_workspace( - account.uuid, - launch['workspace_uuid'], - ) - break - except WorkspaceNotFoundError: - if projection_service is None: - raise - elif projection_service is None: - break - - if attempt == 3: - break - await projection_service.sync_once() - account = await self.ap.user_service.get_user_by_uuid(launch['account_uuid']) - - if access is None and projection_service is not None: - # The target event may be deeper than the bounded incremental - # page budget. One authoritative signed snapshot catches this - # process up without turning the callback into unbounded polling. - await projection_service.refresh_snapshot() - account = await self.ap.user_service.get_user_by_uuid(launch['account_uuid']) - if account is not None: - self.ap.user_service._require_active_account(account) - access = await self.ap.workspace_collaboration_service.resolve_account_workspace( - account.uuid, - launch['workspace_uuid'], - ) + if projection_service is None: + raise SpaceLaunchError('Cloud directory projection is unavailable') + await projection_service.reconcile_workspaces((launch['workspace_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') - if access is None: # pragma: no cover - bounded loop resolves or raises. - raise SpaceLaunchError('Launch Workspace is not projected into Core') + self.ap.user_service._require_active_account(account) + access = await self.ap.workspace_collaboration_service.resolve_account_workspace( + account.uuid, + launch['workspace_uuid'], + ) token = await self.ap.user_service.generate_jwt_token(account) return self.success( data={ diff --git a/src/langbot/pkg/api/http/service/space.py b/src/langbot/pkg/api/http/service/space.py index 174cc965c..101eafdbd 100644 --- a/src/langbot/pkg/api/http/service/space.py +++ b/src/langbot/pkg/api/http/service/space.py @@ -119,21 +119,18 @@ class SpaceService: space_config = self._get_space_config() authorize_url = space_config['oauth_authorize_url'] - params = {'redirect_uri': redirect_uri} + params = {'redirect_uri': redirect_uri, 'code_contract': 'redirect-v1'} if state: params['state'] = state return f'{authorize_url}?{urlencode(params)}' - def get_cloud_entry_url(self) -> str: - """Return the Space-owned Cloud selector for a Cloud Account login.""" - - return f'{self._get_space_config()["url"].rstrip("/")}/cloud?environment=beta' - async def exchange_oauth_code( self, code: str, workspace_uuids: list[str] | None = None, workspace_created_ats: dict[str, int] | None = None, + *, + redirect_uri: str = '', ) -> typing.Dict: """Exchange OAuth authorization code for tokens""" from langbot.pkg.utils import constants @@ -146,6 +143,7 @@ class SpaceService: f'{space_url}/api/v1/accounts/oauth/token', json={ 'code': code, + 'redirect_uri': redirect_uri, 'instance_id': constants.instance_id, # Sending an explicit empty list tells new Space servers not to # synthesize a legacy instance-derived Workspace binding. diff --git a/src/langbot/pkg/api/http/service/user.py b/src/langbot/pkg/api/http/service/user.py index 5e3ca2082..7fadccf6f 100644 --- a/src/langbot/pkg/api/http/service/user.py +++ b/src/langbot/pkg/api/http/service/user.py @@ -774,7 +774,7 @@ class UserService: f'email:{normalized_email}', ) - async def bind_space_account(self, user_email: str, code: str) -> user.User: + async def bind_space_account(self, user_email: str, code: str, *, redirect_uri: str = '') -> user.User: """Bind Space account to existing local account""" local_account = await self.get_user_by_email(user_email) if local_account is None: @@ -794,12 +794,13 @@ class UserService: code, [binding.workspace_uuid], {binding.workspace_uuid: created_ts}, + redirect_uri=redirect_uri, ) else: # Compatibility for early/bootstrap call sites that have not wired # WorkspaceService yet; old Space servers still derive the legacy # Workspace identity from instance_id when the field is omitted. - token_data = await self.ap.space_service.exchange_oauth_code(code) + token_data = await self.ap.space_service.exchange_oauth_code(code, redirect_uri=redirect_uri) access_token = token_data.get('access_token') refresh_token = token_data.get('refresh_token') expires_in = token_data.get('expires_in', 0) diff --git a/src/langbot/pkg/cloud/directory_projection.py b/src/langbot/pkg/cloud/directory_projection.py index c5fd2147d..35a46d960 100644 --- a/src/langbot/pkg/cloud/directory_projection.py +++ b/src/langbot/pkg/cloud/directory_projection.py @@ -173,6 +173,77 @@ class DirectoryProjectionService: async with self._sync_lock: await self._sync_once() + async def reconcile_workspaces(self, workspace_uuids: Iterable[str]) -> None: + """Synchronously project an exact Workspace set without moving the event cursor.""" + + requested = tuple(sorted({str(value).strip() for value in workspace_uuids if str(value).strip()})) + if not requested: + raise DirectoryProjectionUnavailableError('Targeted directory reconciliation requires a Workspace') + if len(requested) > self.event_limit: + raise DirectoryProjectionUnavailableError('Targeted directory reconciliation exceeds the batch limit') + async with self._sync_lock: + delta = await self.provider.fetch_workspaces(self.instance_uuid, requested) + await self._apply_targeted_delta(delta, requested) + + async def _apply_targeted_delta( + self, + delta: DirectoryDelta, + requested_workspace_uuids: tuple[str, ...], + ) -> None: + if not isinstance(delta, DirectoryDelta): + raise DirectoryProjectionUnavailableError('Directory provider returned an invalid delta') + workspace_count, membership_count = self._validate_batch_capacity( + delta.workspaces, + full_snapshot=False, + ) + delta = DirectoryDelta.model_validate(delta.model_dump()) + if delta.instance_uuid != self.instance_uuid: + raise DirectoryProjectionUnavailableError('Directory delta targets another LangBot instance') + requested = set(requested_workspace_uuids) + if set(delta.requested_workspace_uuids) != requested: + raise DirectoryProjectionUnavailableError('Directory delta does not match the requested Workspaces') + if {workspace.uuid for workspace in delta.workspaces} != requested: + raise DirectoryProjectionUnavailableError('Directory delta omitted a requested Workspace') + + directory_uow = getattr(self.ap.persistence_mgr, 'directory_projection_uow', None) + if not callable(directory_uow): + raise DirectoryProjectionUnavailableError('Directory projection persistence scope is unavailable') + + async with directory_uow(self.instance_uuid) as uow: + session = uow.session + state = await session.scalar( + sqlalchemy.select(DirectoryProjectionState) + .where(DirectoryProjectionState.instance_uuid == self.instance_uuid) + .with_for_update() + ) + if state is None: + raise DirectoryProjectionUnavailableError('Directory projection is not initialized') + snapshot = DirectorySnapshot( + instance_uuid=self.instance_uuid, + cursor=state.cursor, + generated_at=delta.generated_at, + workspaces=delta.workspaces, + ) + accounts_by_uuid = await self._apply_accounts(session, snapshot, preserve_existing=True) + await self._apply_workspaces(session, snapshot, accounts_by_uuid=accounts_by_uuid) + active_workspace_count = await self._enforce_active_workspace_capacity(session) + await session.flush() + + await self._update_entitlement_workspace_activity( + snapshot.workspaces, + requested_workspace_uuids=requested, + ) + self._publish_runtime_execution_projection( + snapshot.workspaces, + affected_workspace_uuids=requested, + ) + self._request_model_catalog_sync() + self._record_batch_cardinality( + active_workspaces=active_workspace_count, + workspaces=workspace_count, + memberships=membership_count, + ) + async def _sync_once(self) -> None: cursor = self._consumer_cursor if cursor is None: @@ -723,7 +794,13 @@ class DirectoryProjectionService: for row in inbox_rows: row.applied_at = now - async def _apply_accounts(self, session: Any, snapshot: DirectorySnapshot) -> dict[str, User]: + async def _apply_accounts( + self, + session: Any, + snapshot: DirectorySnapshot, + *, + preserve_existing: bool = False, + ) -> dict[str, User]: selected: dict[str, DirectoryMember] = {} emails: dict[str, str] = {} for workspace in snapshot.workspaces: @@ -788,6 +865,12 @@ class DirectoryProjectionService: continue if account.source != AccountSource.CLOUD_PROJECTION.value: raise DirectoryProjectionUnavailableError('Directory account UUID collides with a local Core account') + if preserve_existing: + # A targeted Workspace fetch has no independently monotonic + # Account revision. It may create a missing runtime shadow, but + # ordered event/snapshot projection remains the only updater of + # existing Account identity and status fields. + continue if account.projection_revision > snapshot.cursor: raise DirectoryProjectionUnavailableError('Directory account revision rolled back') projected_account = self._account_projection(member) diff --git a/tests/integration/api/test_user_space_oauth.py b/tests/integration/api/test_user_space_oauth.py index 34c09451b..af2bab140 100644 --- a/tests/integration/api/test_user_space_oauth.py +++ b/tests/integration/api/test_user_space_oauth.py @@ -11,7 +11,6 @@ import pytest import quart from langbot.pkg.api.http.controller.groups.user import UserRouterGroup -from langbot.pkg.workspace.errors import WorkspaceNotFoundError pytestmark = pytest.mark.integration @@ -71,7 +70,6 @@ async def space_oauth_api(): application.space_service.get_oauth_authorize_url = Mock( side_effect=lambda redirect_uri, state: f'https://space.example/authorize?state={state}' ) - application.space_service.get_cloud_entry_url = Mock(return_value='https://space.example/cloud?environment=beta') application.space_service.exchange_oauth_code = AsyncMock( return_value={ 'access_token': 'space-access-token', @@ -129,7 +127,7 @@ async def test_cloud_launch_state_is_server_issued_and_workspace_bound(space_oau @pytest.mark.asyncio -async def test_cloud_login_entry_redirects_to_space_workspace_launcher(space_oauth_api): +async def test_cloud_login_entry_uses_normal_stateful_oauth(space_oauth_api): application, client = space_oauth_api application.deployment.mode = 'cloud' @@ -143,9 +141,9 @@ async def test_cloud_login_entry_redirects_to_space_workspace_launcher(space_oau ) assert response.status_code == 200 - assert (await response.get_json())['data']['authorize_url'] == ('https://space.example/cloud?environment=beta') - application.space_service.get_cloud_entry_url.assert_called_once_with() - application.user_service.issue_space_oauth_state.assert_not_awaited() + authorize_url = (await response.get_json())['data']['authorize_url'] + assert authorize_url.startswith('https://space.example/authorize?state=') + application.user_service.issue_space_oauth_state.assert_awaited_once_with('login') @pytest.mark.asyncio @@ -272,10 +270,14 @@ async def test_server_side_webhook_origin_supports_bundled_ui(space_oauth_api): async def test_login_callback_requires_and_consumes_server_state(space_oauth_api): application, client = space_oauth_api - missing = await client.post('/api/v1/user/space/callback', json={'code': 'oauth-code'}) + missing = await client.post('/api/v1/user/space/callback', json={'code': 'v4_oauth-code'}) response = await client.post( '/api/v1/user/space/callback', - json={'code': 'oauth-code', 'state': 'opaque-login-state'}, + json={ + 'code': 'v4_oauth-code', + 'state': 'opaque-login-state', + 'redirect_uri': 'https://oss.example/auth/space/callback', + }, ) assert (await missing.get_json())['code'] == 1 @@ -283,12 +285,146 @@ async def test_login_callback_requires_and_consumes_server_state(space_oauth_api assert (await response.get_json())['data']['token'] == 'space-login-token' application.user_service.consume_space_oauth_state_details.assert_awaited_once_with('opaque-login-state', 'login') application.space_service.exchange_oauth_code.assert_awaited_once_with( - 'oauth-code', + 'v4_oauth-code', [WORKSPACE_UUID], {WORKSPACE_UUID: int(WORKSPACE_CREATED_AT.timestamp())}, + redirect_uri='https://oss.example/auth/space/callback', ) +@pytest.mark.asyncio +async def test_login_callback_rejects_downgraded_legacy_code(space_oauth_api): + application, client = space_oauth_api + + response = await client.post( + '/api/v1/user/space/callback', + json={'code': 'v2_legacy-code', 'state': 'opaque-login-state'}, + ) + + payload = await response.get_json() + assert response.status_code == 200 + assert payload['code'] == 1 + assert 'code contract' in payload['msg'] + application.space_service.exchange_oauth_code.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_cloud_login_callback_reconciles_authorized_workspace_before_local_authentication(space_oauth_api): + application, client = space_oauth_api + application.deployment.mode = 'cloud' + calls: list[str] = [] + application.directory_projection_service = SimpleNamespace( + reconcile_workspaces=AsyncMock(side_effect=lambda _workspace_uuids: calls.append('reconcile')) + ) + application.space_service.exchange_oauth_code.return_value = { + 'access_token': 'space-access-token', + 'refresh_token': 'space-refresh-token', + 'expires_in': 3600, + 'cloud_workspace_uuid': WORKSPACE_UUID, + } + + authenticated_account = application.user_service.authenticate_space_user.return_value[1] + + async def authenticate(*_args): + calls.append('authenticate') + return 'space-login-token', authenticated_account + + application.user_service.authenticate_space_user.side_effect = authenticate + + response = await client.post( + '/api/v1/user/space/callback', + json={'code': 'v4_oauth-code', 'state': 'opaque-login-state'}, + ) + + assert response.status_code == 200 + assert (await response.get_json())['data']['workspace_uuid'] == WORKSPACE_UUID + assert calls == ['reconcile', 'authenticate'] + application.directory_projection_service.reconcile_workspaces.assert_awaited_once_with((WORKSPACE_UUID,)) + + +@pytest.mark.asyncio +async def test_cloud_login_callback_fails_closed_without_workspace_binding(space_oauth_api): + application, client = space_oauth_api + application.deployment.mode = 'cloud' + application.directory_projection_service = SimpleNamespace(reconcile_workspaces=AsyncMock()) + + response = await client.post( + '/api/v1/user/space/callback', + json={'code': 'v4_oauth-code', 'state': 'opaque-login-state'}, + ) + + payload = await response.get_json() + assert response.status_code == 200 + assert payload['code'] == 1 + assert 'Cloud Workspace binding' in payload['msg'] + application.directory_projection_service.reconcile_workspaces.assert_not_awaited() + application.user_service.authenticate_space_user.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_cloud_login_callback_requires_code_binding_for_launch_state(space_oauth_api): + application, client = space_oauth_api + application.deployment.mode = 'cloud' + application.directory_projection_service = SimpleNamespace(reconcile_workspaces=AsyncMock()) + application.user_service.consume_space_oauth_state_details.return_value = SimpleNamespace( + launch_workspace_uuid=WORKSPACE_UUID + ) + + response = await client.post( + '/api/v1/user/space/callback', + json={'code': 'v4_oauth-code', 'state': 'opaque-login-state'}, + ) + + payload = await response.get_json() + assert response.status_code == 200 + assert payload['code'] == 1 + assert 'Workspace binding' in payload['msg'] + application.directory_projection_service.reconcile_workspaces.assert_not_awaited() + application.user_service.authenticate_space_user.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_cloud_login_callback_rejects_conflicting_state_and_code_workspace_bindings(space_oauth_api): + application, client = space_oauth_api + application.deployment.mode = 'cloud' + application.directory_projection_service = SimpleNamespace(reconcile_workspaces=AsyncMock()) + application.user_service.consume_space_oauth_state_details.return_value = SimpleNamespace( + launch_workspace_uuid=WORKSPACE_UUID + ) + application.space_service.exchange_oauth_code.return_value = { + 'access_token': 'space-access-token', + 'refresh_token': 'space-refresh-token', + 'expires_in': 3600, + 'cloud_workspace_uuid': 'workspace-from-another-flow', + } + + response = await client.post( + '/api/v1/user/space/callback', + json={'code': 'v4_oauth-code', 'state': 'opaque-login-state'}, + ) + + payload = await response.get_json() + assert response.status_code == 200 + assert payload['code'] == 1 + assert 'Workspace binding' in payload['msg'] + application.directory_projection_service.reconcile_workspaces.assert_not_awaited() + application.user_service.authenticate_space_user.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_oss_login_callback_does_not_request_cloud_reconciliation(space_oauth_api): + application, client = space_oauth_api + application.directory_projection_service = SimpleNamespace(reconcile_workspaces=AsyncMock()) + + response = await client.post( + '/api/v1/user/space/callback', + json={'code': 'v4_oauth-code', 'state': 'opaque-login-state'}, + ) + + assert response.status_code == 200 + application.directory_projection_service.reconcile_workspaces.assert_not_awaited() + + @pytest.mark.asyncio async def test_login_callback_launch_state_selects_asserted_workspace(space_oauth_api): application, client = space_oauth_api @@ -299,7 +435,7 @@ async def test_login_callback_launch_state_selects_asserted_workspace(space_oaut response = await client.post( '/api/v1/user/space/callback', - json={'code': 'oauth-code', 'state': 'opaque-login-state'}, + json={'code': 'v4_oauth-code', 'state': 'opaque-login-state'}, ) assert response.status_code == 200 @@ -398,18 +534,22 @@ async def test_bind_callback_uses_opaque_state_and_never_treats_it_as_jwt(space_ rejected = await client.post( '/api/v1/user/bind-space', - json={'code': 'attacker-code', 'state': 'jwt.must-not-be-used'}, + json={'code': 'v4_attacker-code', 'state': 'jwt.must-not-be-used'}, ) response = await client.post( '/api/v1/user/bind-space', - json={'code': 'oauth-code', 'state': 'opaque-bind-state'}, + json={'code': 'v4_oauth-code', 'state': 'opaque-bind-state'}, ) assert rejected.status_code == 401 assert response.status_code == 200 assert (await response.get_json())['data']['token'] == 'rotated-account-token' application.user_service.verify_jwt_token.assert_not_awaited() - application.user_service.bind_space_account.assert_awaited_once_with('owner@example.com', 'oauth-code') + application.user_service.bind_space_account.assert_awaited_once_with( + 'owner@example.com', + 'v4_oauth-code', + redirect_uri='http://localhost/auth/space/callback?mode=bind', + ) @pytest.mark.asyncio @@ -417,6 +557,7 @@ async def test_direct_launch_assertion_does_not_consume_normal_oauth_state(space application, client = space_oauth_api application.user_service.consume_space_oauth_state.reset_mock() application.space_service.exchange_oauth_code.reset_mock() + application.directory_projection_service = SimpleNamespace(reconcile_workspaces=AsyncMock()) response = await client.post( '/api/v1/user/space/callback', @@ -440,7 +581,7 @@ async def test_direct_launch_assertion_does_not_consume_normal_oauth_state(space @pytest.mark.asyncio -async def test_direct_launch_refreshes_new_workspace_projection_before_rejecting_account(space_oauth_api): +async def test_direct_launch_reconciles_exact_workspace_before_resolving_access(space_oauth_api): application, client = space_oauth_api projected_account = SimpleNamespace( uuid='account-a', @@ -448,8 +589,8 @@ async def test_direct_launch_refreshes_new_workspace_projection_before_rejecting account_type='space', status='active', ) - application.user_service.get_user_by_uuid = AsyncMock(side_effect=[None, None, projected_account]) - application.directory_projection_service = SimpleNamespace(sync_once=AsyncMock()) + application.user_service.get_user_by_uuid = AsyncMock(return_value=projected_account) + application.directory_projection_service = SimpleNamespace(reconcile_workspaces=AsyncMock()) response = await client.post( '/api/v1/user/space/callback', @@ -461,61 +602,5 @@ async def test_direct_launch_refreshes_new_workspace_projection_before_rejecting assert response.status_code == 200 assert (await response.get_json())['data']['workspace_uuid'] == WORKSPACE_UUID - assert application.directory_projection_service.sync_once.await_count == 2 - assert application.user_service.get_user_by_uuid.await_count == 3 - - -@pytest.mark.asyncio -async def test_direct_launch_refreshes_projection_when_account_exists_before_workspace(space_oauth_api): - application, client = space_oauth_api - projected_access = application.workspace_collaboration_service.resolve_account_workspace.return_value - application.workspace_collaboration_service.resolve_account_workspace = AsyncMock( - side_effect=[WorkspaceNotFoundError('Workspace not found'), projected_access] - ) - application.directory_projection_service = SimpleNamespace(sync_once=AsyncMock()) - - response = await client.post( - '/api/v1/user/space/callback', - json={ - 'workspace_uuid': WORKSPACE_UUID, - 'launch_assertion': 'signed-launch-token', - }, - ) - - assert response.status_code == 200 - assert (await response.get_json())['data']['workspace_uuid'] == WORKSPACE_UUID - application.directory_projection_service.sync_once.assert_awaited_once_with() - assert application.workspace_collaboration_service.resolve_account_workspace.await_count == 2 - - -@pytest.mark.asyncio -async def test_direct_launch_falls_back_to_snapshot_when_event_backlog_exceeds_page_budget(space_oauth_api): - application, client = space_oauth_api - projected_access = application.workspace_collaboration_service.resolve_account_workspace.return_value - application.workspace_collaboration_service.resolve_account_workspace = AsyncMock( - side_effect=[ - WorkspaceNotFoundError('Workspace not found'), - WorkspaceNotFoundError('Workspace not found'), - WorkspaceNotFoundError('Workspace not found'), - WorkspaceNotFoundError('Workspace not found'), - projected_access, - ] - ) - application.directory_projection_service = SimpleNamespace( - sync_once=AsyncMock(), - refresh_snapshot=AsyncMock(), - ) - - response = await client.post( - '/api/v1/user/space/callback', - json={ - 'workspace_uuid': WORKSPACE_UUID, - 'launch_assertion': 'signed-launch-token', - }, - ) - - assert response.status_code == 200 - assert (await response.get_json())['data']['workspace_uuid'] == WORKSPACE_UUID - assert application.directory_projection_service.sync_once.await_count == 3 - application.directory_projection_service.refresh_snapshot.assert_awaited_once_with() - assert application.workspace_collaboration_service.resolve_account_workspace.await_count == 5 + application.directory_projection_service.reconcile_workspaces.assert_awaited_once_with((WORKSPACE_UUID,)) + application.user_service.get_user_by_uuid.assert_awaited_once_with('account-a') diff --git a/tests/unit_tests/api/service/test_space_service.py b/tests/unit_tests/api/service/test_space_service.py index a77dd14a1..731244561 100644 --- a/tests/unit_tests/api/service/test_space_service.py +++ b/tests/unit_tests/api/service/test_space_service.py @@ -95,7 +95,9 @@ class TestSpaceServiceGetOAuthAuthorizeUrl: result = service.get_oauth_authorize_url('http://localhost/callback') # Verify - assert parse_qs(urlsplit(result).query)['redirect_uri'] == ['http://localhost/callback'] + query = parse_qs(urlsplit(result).query) + assert query['redirect_uri'] == ['http://localhost/callback'] + assert query['code_contract'] == ['redirect-v1'] assert 'https://space.langbot.app/auth/authorize' in result def test_get_oauth_authorize_url_with_state(self): @@ -578,12 +580,14 @@ class TestSpaceServiceExchangeOAuthCode: 'auth_code', ['workspace-1'], {'workspace-1': 1_700_000_000}, + redirect_uri='https://oss.example/auth/space/callback', ) # Verify assert result['access_token'] == 'new_access_token' assert mock_session_obj.post.call_args.kwargs['json'] == { 'code': 'auth_code', + 'redirect_uri': 'https://oss.example/auth/space/callback', 'instance_id': constants.instance_id, 'workspace_uuids': ['workspace-1'], 'workspace_created_ats': {'workspace-1': 1_700_000_000}, @@ -846,10 +850,7 @@ class TestSpaceServiceGetModelSelection: if response_shape == 'models-envelope': data = {'models': models} elif response_shape == 'availability-wrapper': - data = [ - {'model': model, 'latency_ms': index + 10, 'http_code': 200} - for index, model in enumerate(models) - ] + data = [{'model': model, 'latency_ms': index + 10, 'http_code': 200} for index, model in enumerate(models)] else: data = models payload = {'code': 0, 'data': data} diff --git a/tests/unit_tests/cloud/test_directory_projection.py b/tests/unit_tests/cloud/test_directory_projection.py index 1320a9abc..756baa7be 100644 --- a/tests/unit_tests/cloud/test_directory_projection.py +++ b/tests/unit_tests/cloud/test_directory_projection.py @@ -4,7 +4,7 @@ import asyncio import datetime import logging from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import AsyncMock, Mock import pytest import sqlalchemy @@ -215,6 +215,88 @@ async def test_directory_delta_requests_model_catalog_sync_after_commit(projecti request_sync.assert_called_once_with() +async def test_targeted_reconciliation_projects_new_workspace_without_advancing_event_cursor(projection_context): + application, session_factory = projection_context + provider = _Provider( + [_snapshot(7, workspaces=[])], + deltas=[_delta(workspaces=[_workspace(revision=8, name='JIT Workspace')])], + ) + service = DirectoryProjectionService(application, provider, INSTANCE_UUID) + await service.initialize() + + await service.reconcile_workspaces((WORKSPACE_UUID,)) + + async with session_factory() as session: + account = await session.scalar(sqlalchemy.select(User).where(User.uuid == ACCOUNT_UUID)) + workspace = await session.get(Workspace, WORKSPACE_UUID) + membership = await session.scalar( + sqlalchemy.select(WorkspaceMembership).where( + WorkspaceMembership.workspace_uuid == WORKSPACE_UUID, + WorkspaceMembership.account_uuid == ACCOUNT_UUID, + ) + ) + state = await session.get(DirectoryProjectionState, INSTANCE_UUID) + assert account is not None + assert workspace is not None and workspace.name == 'JIT Workspace' + assert membership is not None and membership.status == 'active' + assert state is not None and state.cursor == 7 + assert provider.delta_calls == 1 + assert provider.after_cursors == [] + + +async def test_targeted_reconciliation_preserves_existing_account_until_ordered_event_projection(projection_context): + application, session_factory = projection_context + targeted_workspace = _workspace(revision=8, name='Renamed Workspace').model_copy( + update={ + 'members': [ + _member(revision=8).model_copy(update={'display_name': 'Changed Account Name'}) + ] + } + ) + provider = _Provider( + [_snapshot(7)], + deltas=[_delta(workspaces=[targeted_workspace])], + ) + service = DirectoryProjectionService(application, provider, INSTANCE_UUID) + await service.initialize() + + await service.reconcile_workspaces((WORKSPACE_UUID,)) + + async with session_factory() as session: + account = await session.scalar(sqlalchemy.select(User).where(User.uuid == ACCOUNT_UUID)) + workspace = await session.get(Workspace, WORKSPACE_UUID) + state = await session.get(DirectoryProjectionState, INSTANCE_UUID) + assert account is not None and account.user == 'Workspace Owner' + assert account.projection_revision == 7 + assert workspace is not None and workspace.name == 'Renamed Workspace' + assert state is not None and state.cursor == 7 + + +async def test_targeted_reconciliation_only_updates_requested_workspace_side_effects(projection_context): + application, _session_factory = projection_context + provider = _Provider( + [_snapshot(7, workspaces=[])], + deltas=[_delta(workspaces=[_workspace(revision=8, name='JIT Workspace')])], + ) + service = DirectoryProjectionService(application, provider, INSTANCE_UUID) + await service.initialize() + service._reconcile_entitlement_snapshot_set = AsyncMock() + service._update_entitlement_workspace_activity = AsyncMock() + service._publish_runtime_execution_projection = Mock() + + await service.reconcile_workspaces((WORKSPACE_UUID,)) + + service._reconcile_entitlement_snapshot_set.assert_not_awaited() + service._update_entitlement_workspace_activity.assert_awaited_once() + assert service._update_entitlement_workspace_activity.await_args.kwargs == { + 'requested_workspace_uuids': {WORKSPACE_UUID}, + } + service._publish_runtime_execution_projection.assert_called_once() + assert service._publish_runtime_execution_projection.call_args.kwargs == { + 'affected_workspace_uuids': {WORKSPACE_UUID}, + } + + async def test_initial_snapshot_projects_core_owned_rows(projection_context): application, session_factory = projection_context reconcile_execution_projection = Mock() diff --git a/web/src/app/auth/space/callback/page.tsx b/web/src/app/auth/space/callback/page.tsx index c2c142589..ea41ded6c 100644 --- a/web/src/app/auth/space/callback/page.tsx +++ b/web/src/app/auth/space/callback/page.tsx @@ -43,17 +43,24 @@ const pendingSpaceOAuthLogins = new Map< function getOrCreateSpaceOAuthLoginPromise( authCode: string, state: string, + redirectUri: string, workspaceUuid?: string, launchAssertion?: string, ): Promise { - const requestKey = `${authCode}:${state}:${workspaceUuid ?? ''}:${launchAssertion ?? ''}`; + const requestKey = `${authCode}:${state}:${redirectUri}:${workspaceUuid ?? ''}:${launchAssertion ?? ''}`; const pendingRequest = pendingSpaceOAuthLogins.get(requestKey); if (pendingRequest) { return pendingRequest; } const requestPromise = httpClient - .exchangeSpaceOAuthCode(authCode, state, workspaceUuid, launchAssertion) + .exchangeSpaceOAuthCode( + authCode, + state, + redirectUri, + workspaceUuid, + launchAssertion, + ) .finally(() => { pendingSpaceOAuthLogins.delete(requestKey); }); @@ -95,6 +102,7 @@ function SpaceOAuthCallbackContent() { const response = await getOrCreateSpaceOAuthLoginPromise( authCode, state, + `${window.location.origin}/auth/space/callback`, workspaceUuid, launchAssertion, ); @@ -195,7 +203,11 @@ function SpaceOAuthCallbackContent() { async (authCode: string, state: string) => { setIsProcessing(true); try { - const response = await httpClient.bindSpaceAccount(authCode, state); + const response = await httpClient.bindSpaceAccount( + authCode, + state, + `${window.location.origin}/auth/space/callback?mode=bind`, + ); if (!isMountedRef.current) { return; } diff --git a/web/src/app/infra/http/BackendClient.ts b/web/src/app/infra/http/BackendClient.ts index 41aac21cb..c9d742c4c 100644 --- a/web/src/app/infra/http/BackendClient.ts +++ b/web/src/app/infra/http/BackendClient.ts @@ -1365,6 +1365,7 @@ export class BackendClient extends BaseHttpClient { public async bindSpaceAccount( code: string, state: string, + redirectUri: string, ): Promise<{ token: string; user: string; @@ -1372,7 +1373,7 @@ export class BackendClient extends BaseHttpClient { }> { const response = await this.instance.post( '/api/v1/user/bind-space', - { code, state }, + { code, state, redirect_uri: redirectUri }, { skipWorkspace: true } as RequestConfig, ); if (response.data.code !== 0) { @@ -1385,18 +1386,12 @@ export class BackendClient extends BaseHttpClient { } // ============ Space OAuth API (Redirect Flow) ============ - public getSpaceAuthorizeUrl( - redirectUri: string, - options?: { cloudEntry?: boolean }, - ): Promise<{ + public getSpaceAuthorizeUrl(redirectUri: string): Promise<{ authorize_url: string; }> { return this.get( '/api/v1/user/space/authorize-url', - { - redirect_uri: redirectUri, - ...(options?.cloudEntry ? { cloud_entry: '1' } : {}), - }, + { redirect_uri: redirectUri }, { skipWorkspace: true }, ); } @@ -1414,6 +1409,7 @@ export class BackendClient extends BaseHttpClient { public async exchangeSpaceOAuthCode( code: string, state: string, + redirectUri: string, workspaceUuid?: string, launchAssertion?: string, ): Promise<{ @@ -1428,6 +1424,7 @@ export class BackendClient extends BaseHttpClient { { code, state, + redirect_uri: redirectUri, workspace_uuid: workspaceUuid, launch_assertion: launchAssertion, }, diff --git a/web/src/app/login/page.tsx b/web/src/app/login/page.tsx index 061a21988..48436017f 100644 --- a/web/src/app/login/page.tsx +++ b/web/src/app/login/page.tsx @@ -202,13 +202,7 @@ export default function Login() { try { const currentOrigin = window.location.origin; const redirectUri = `${currentOrigin}/auth/space/callback`; - const response = await httpClient.getSpaceAuthorizeUrl(redirectUri, { - // Cloud Accounts must be launched from Space so a first visit can - // lazily create and project the personal Workspace. Invitation login - // remains on the OAuth callback path because it targets the invited - // Workspace instead. - cloudEntry: !getPendingInvitationToken(), - }); + const response = await httpClient.getSpaceAuthorizeUrl(redirectUri); window.location.href = response.authorize_url; } catch { toast.error(t('common.spaceLoginFailed')); diff --git a/web/tests/unit/cloud-new-account-entry.test.mjs b/web/tests/unit/cloud-new-account-entry.test.mjs index 5420da59e..8a096be33 100644 --- a/web/tests/unit/cloud-new-account-entry.test.mjs +++ b/web/tests/unit/cloud-new-account-entry.test.mjs @@ -7,11 +7,13 @@ const source = fs.readFileSync( 'utf8', ); -test('normal Cloud login enters through the Space Workspace launcher', () => { - assert.match(source, /cloudEntry:\s*!getPendingInvitationToken\(\)/); - assert.match(source, /getSpaceAuthorizeUrl\(redirectUri,\s*\{/); +test('normal Cloud login uses the standard Space OAuth callback path', () => { + assert.doesNotMatch(source, /cloudEntry/); + assert.match(source, /getSpaceAuthorizeUrl\(redirectUri\)/); }); -test('invitation login remains on the OAuth callback path', () => { - assert.match(source, /cloudEntry:\s*!getPendingInvitationToken\(\)/); +test('invitation login uses the same OAuth callback before accepting the invitation', () => { + assert.doesNotMatch(source, /cloudEntry/); + assert.match(source, /const invitationToken = getPendingInvitationToken\(\)/); + assert.match(source, /acceptWorkspaceInvitation\(invitationToken\)/); });