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>
This commit is contained in:
Hyu
2026-09-03 23:14:44 +08:00
committed by GitHub
parent ab52684a01
commit b44b8f474d
11 changed files with 407 additions and 165 deletions
@@ -9,7 +9,6 @@ from .. import group
from .....entity.errors import account as account_errors from .....entity.errors import account as account_errors
from ...context import RequestContext from ...context import RequestContext
from .....cloud.launch import SpaceLaunchError from .....cloud.launch import SpaceLaunchError
from .....workspace.errors import WorkspaceNotFoundError
from ...service.user import ControlPlaneDirectoryRequiredError, PublicRegistrationClosedError from ...service.user import ControlPlaneDirectoryRequiredError, PublicRegistrationClosedError
@@ -144,13 +143,6 @@ class UserRouterGroup(group.RouterGroup):
try: try:
redirect_uri = self._validate_space_redirect_uri(redirect_uri, bind=False) redirect_uri = self._validate_space_redirect_uri(redirect_uri, bind=False)
launch_workspace_uuid = quart.request.args.get('launch_workspace_uuid') 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 launch_workspace_uuid:
if not getattr(getattr(self.ap, 'deployment', None), 'multi_workspace_enabled', False): if not getattr(getattr(self.ap, 'deployment', None), 'multi_workspace_enabled', False):
return self.fail(1, 'Space launch requires Cloud mode') return self.fail(1, 'Space launch requires Cloud mode')
@@ -194,6 +186,9 @@ class UserRouterGroup(group.RouterGroup):
json_data = await quart.request.json json_data = await quart.request.json
code = json_data.get('code') code = json_data.get('code')
state = json_data.get('state') 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') launch_assertion = json_data.get('launch_assertion')
workspace_uuid = json_data.get('workspace_uuid') workspace_uuid = json_data.get('workspace_uuid')
@@ -207,8 +202,11 @@ class UserRouterGroup(group.RouterGroup):
return self.fail(1, 'Missing authorization code') return self.fail(1, 'Missing authorization code')
if not state: if not state:
return self.fail(1, 'Missing state parameter') return self.fail(1, 'Missing state parameter')
if not str(code).startswith('v4_'):
return self.fail(1, 'Unsupported Space OAuth code contract')
try: 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') consumed_state = await self.ap.user_service.consume_space_oauth_state_details(state, 'login')
# Exchange code for tokens # Exchange code for tokens
launch_workspace_uuid = consumed_state.launch_workspace_uuid launch_workspace_uuid = consumed_state.launch_workspace_uuid
@@ -226,24 +224,36 @@ class UserRouterGroup(group.RouterGroup):
code, code,
workspace_uuids, workspace_uuids,
workspace_created_ats, workspace_created_ats,
redirect_uri=redirect_uri,
) )
access_token = token_data.get('access_token') access_token = token_data.get('access_token')
refresh_token = token_data.get('refresh_token') refresh_token = token_data.get('refresh_token')
expires_in = token_data.get('expires_in', 0) expires_in = token_data.get('expires_in', 0)
cloud_workspace_uuid = token_data.get('cloud_workspace_uuid')
if not access_token: if not access_token:
return self.fail(1, 'Failed to get access token from Space') 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( jwt_token, user_obj = await self.ap.user_service.authenticate_space_user(
access_token, refresh_token, expires_in access_token, refresh_token, expires_in
) )
if launch_workspace_uuid: if target_workspace_uuid:
try: try:
access = await self.ap.workspace_collaboration_service.resolve_account_workspace( access = await self.ap.workspace_collaboration_service.resolve_account_workspace(
user_obj.uuid, user_obj.uuid,
launch_workspace_uuid, target_workspace_uuid,
) )
except Exception: except Exception:
self.ap.logger.warning('Rejected Space OAuth launch for unauthorized Workspace') 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 json_data = await quart.request.json
code = json_data.get('code') code = json_data.get('code')
state = json_data.get('state') 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: if not code:
return self.http_status(400, -1, 'Missing authorization code') return self.http_status(400, -1, 'Missing authorization code')
if not state: if not state:
return self.http_status(400, -1, 'Missing state parameter') 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: try:
user_obj = await self.ap.user_service.consume_space_oauth_state(state, 'bind') 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') return self.http_status(400, -1, 'Only local accounts can bind to Space')
try: 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) jwt_token = await self.ap.user_service.generate_jwt_token(updated_user)
return self.success( return self.success(
data={ 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 projection_service = self.ap.directory_projection_service
access = None if projection_service is None:
# A first Cloud launch creates the personal Workspace immediately raise SpaceLaunchError('Cloud directory projection is unavailable')
# before redirecting here. Pull a bounded number of signed event await projection_service.reconcile_workspaces((launch['workspace_uuid'],))
# pages until both the Account and its target Workspace membership account = await self.ap.user_service.get_user_by_uuid(launch['account_uuid'])
# 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 account is None: if account is None:
raise SpaceLaunchError('Launch Account is not projected into Core') raise SpaceLaunchError('Launch Account is not projected into Core')
if access is None: # pragma: no cover - bounded loop resolves or raises. self.ap.user_service._require_active_account(account)
raise SpaceLaunchError('Launch Workspace is not projected into Core') 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) token = await self.ap.user_service.generate_jwt_token(account)
return self.success( return self.success(
data={ data={
+4 -6
View File
@@ -119,21 +119,18 @@ class SpaceService:
space_config = self._get_space_config() space_config = self._get_space_config()
authorize_url = space_config['oauth_authorize_url'] authorize_url = space_config['oauth_authorize_url']
params = {'redirect_uri': redirect_uri} params = {'redirect_uri': redirect_uri, 'code_contract': 'redirect-v1'}
if state: if state:
params['state'] = state params['state'] = state
return f'{authorize_url}?{urlencode(params)}' 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( async def exchange_oauth_code(
self, self,
code: str, code: str,
workspace_uuids: list[str] | None = None, workspace_uuids: list[str] | None = None,
workspace_created_ats: dict[str, int] | None = None, workspace_created_ats: dict[str, int] | None = None,
*,
redirect_uri: str = '',
) -> typing.Dict: ) -> typing.Dict:
"""Exchange OAuth authorization code for tokens""" """Exchange OAuth authorization code for tokens"""
from langbot.pkg.utils import constants from langbot.pkg.utils import constants
@@ -146,6 +143,7 @@ class SpaceService:
f'{space_url}/api/v1/accounts/oauth/token', f'{space_url}/api/v1/accounts/oauth/token',
json={ json={
'code': code, 'code': code,
'redirect_uri': redirect_uri,
'instance_id': constants.instance_id, 'instance_id': constants.instance_id,
# Sending an explicit empty list tells new Space servers not to # Sending an explicit empty list tells new Space servers not to
# synthesize a legacy instance-derived Workspace binding. # synthesize a legacy instance-derived Workspace binding.
+3 -2
View File
@@ -774,7 +774,7 @@ class UserService:
f'email:{normalized_email}', 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""" """Bind Space account to existing local account"""
local_account = await self.get_user_by_email(user_email) local_account = await self.get_user_by_email(user_email)
if local_account is None: if local_account is None:
@@ -794,12 +794,13 @@ class UserService:
code, code,
[binding.workspace_uuid], [binding.workspace_uuid],
{binding.workspace_uuid: created_ts}, {binding.workspace_uuid: created_ts},
redirect_uri=redirect_uri,
) )
else: else:
# Compatibility for early/bootstrap call sites that have not wired # Compatibility for early/bootstrap call sites that have not wired
# WorkspaceService yet; old Space servers still derive the legacy # WorkspaceService yet; old Space servers still derive the legacy
# Workspace identity from instance_id when the field is omitted. # 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') access_token = token_data.get('access_token')
refresh_token = token_data.get('refresh_token') refresh_token = token_data.get('refresh_token')
expires_in = token_data.get('expires_in', 0) expires_in = token_data.get('expires_in', 0)
+84 -1
View File
@@ -173,6 +173,77 @@ class DirectoryProjectionService:
async with self._sync_lock: async with self._sync_lock:
await self._sync_once() 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: async def _sync_once(self) -> None:
cursor = self._consumer_cursor cursor = self._consumer_cursor
if cursor is None: if cursor is None:
@@ -723,7 +794,13 @@ class DirectoryProjectionService:
for row in inbox_rows: for row in inbox_rows:
row.applied_at = now 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] = {} selected: dict[str, DirectoryMember] = {}
emails: dict[str, str] = {} emails: dict[str, str] = {}
for workspace in snapshot.workspaces: for workspace in snapshot.workspaces:
@@ -788,6 +865,12 @@ class DirectoryProjectionService:
continue continue
if account.source != AccountSource.CLOUD_PROJECTION.value: if account.source != AccountSource.CLOUD_PROJECTION.value:
raise DirectoryProjectionUnavailableError('Directory account UUID collides with a local Core account') 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: if account.projection_revision > snapshot.cursor:
raise DirectoryProjectionUnavailableError('Directory account revision rolled back') raise DirectoryProjectionUnavailableError('Directory account revision rolled back')
projected_account = self._account_projection(member) projected_account = self._account_projection(member)
+159 -74
View File
@@ -11,7 +11,6 @@ import pytest
import quart import quart
from langbot.pkg.api.http.controller.groups.user import UserRouterGroup from langbot.pkg.api.http.controller.groups.user import UserRouterGroup
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
pytestmark = pytest.mark.integration pytestmark = pytest.mark.integration
@@ -71,7 +70,6 @@ async def space_oauth_api():
application.space_service.get_oauth_authorize_url = Mock( application.space_service.get_oauth_authorize_url = Mock(
side_effect=lambda redirect_uri, state: f'https://space.example/authorize?state={state}' 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( application.space_service.exchange_oauth_code = AsyncMock(
return_value={ return_value={
'access_token': 'space-access-token', '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 @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, client = space_oauth_api
application.deployment.mode = 'cloud' 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 response.status_code == 200
assert (await response.get_json())['data']['authorize_url'] == ('https://space.example/cloud?environment=beta') authorize_url = (await response.get_json())['data']['authorize_url']
application.space_service.get_cloud_entry_url.assert_called_once_with() assert authorize_url.startswith('https://space.example/authorize?state=')
application.user_service.issue_space_oauth_state.assert_not_awaited() application.user_service.issue_space_oauth_state.assert_awaited_once_with('login')
@pytest.mark.asyncio @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): async def test_login_callback_requires_and_consumes_server_state(space_oauth_api):
application, client = 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( response = await client.post(
'/api/v1/user/space/callback', '/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 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' 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.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( application.space_service.exchange_oauth_code.assert_awaited_once_with(
'oauth-code', 'v4_oauth-code',
[WORKSPACE_UUID], [WORKSPACE_UUID],
{WORKSPACE_UUID: int(WORKSPACE_CREATED_AT.timestamp())}, {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 @pytest.mark.asyncio
async def test_login_callback_launch_state_selects_asserted_workspace(space_oauth_api): async def test_login_callback_launch_state_selects_asserted_workspace(space_oauth_api):
application, client = 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( response = await client.post(
'/api/v1/user/space/callback', '/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 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( rejected = await client.post(
'/api/v1/user/bind-space', '/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( response = await client.post(
'/api/v1/user/bind-space', '/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 rejected.status_code == 401
assert response.status_code == 200 assert response.status_code == 200
assert (await response.get_json())['data']['token'] == 'rotated-account-token' assert (await response.get_json())['data']['token'] == 'rotated-account-token'
application.user_service.verify_jwt_token.assert_not_awaited() 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 @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, client = space_oauth_api
application.user_service.consume_space_oauth_state.reset_mock() application.user_service.consume_space_oauth_state.reset_mock()
application.space_service.exchange_oauth_code.reset_mock() application.space_service.exchange_oauth_code.reset_mock()
application.directory_projection_service = SimpleNamespace(reconcile_workspaces=AsyncMock())
response = await client.post( response = await client.post(
'/api/v1/user/space/callback', '/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 @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 application, client = space_oauth_api
projected_account = SimpleNamespace( projected_account = SimpleNamespace(
uuid='account-a', uuid='account-a',
@@ -448,8 +589,8 @@ async def test_direct_launch_refreshes_new_workspace_projection_before_rejecting
account_type='space', account_type='space',
status='active', status='active',
) )
application.user_service.get_user_by_uuid = AsyncMock(side_effect=[None, None, projected_account]) application.user_service.get_user_by_uuid = AsyncMock(return_value=projected_account)
application.directory_projection_service = SimpleNamespace(sync_once=AsyncMock()) application.directory_projection_service = SimpleNamespace(reconcile_workspaces=AsyncMock())
response = await client.post( response = await client.post(
'/api/v1/user/space/callback', '/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 response.status_code == 200
assert (await response.get_json())['data']['workspace_uuid'] == WORKSPACE_UUID assert (await response.get_json())['data']['workspace_uuid'] == WORKSPACE_UUID
assert application.directory_projection_service.sync_once.await_count == 2 application.directory_projection_service.reconcile_workspaces.assert_awaited_once_with((WORKSPACE_UUID,))
assert application.user_service.get_user_by_uuid.await_count == 3 application.user_service.get_user_by_uuid.assert_awaited_once_with('account-a')
@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
@@ -95,7 +95,9 @@ class TestSpaceServiceGetOAuthAuthorizeUrl:
result = service.get_oauth_authorize_url('http://localhost/callback') result = service.get_oauth_authorize_url('http://localhost/callback')
# Verify # 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 assert 'https://space.langbot.app/auth/authorize' in result
def test_get_oauth_authorize_url_with_state(self): def test_get_oauth_authorize_url_with_state(self):
@@ -578,12 +580,14 @@ class TestSpaceServiceExchangeOAuthCode:
'auth_code', 'auth_code',
['workspace-1'], ['workspace-1'],
{'workspace-1': 1_700_000_000}, {'workspace-1': 1_700_000_000},
redirect_uri='https://oss.example/auth/space/callback',
) )
# Verify # Verify
assert result['access_token'] == 'new_access_token' assert result['access_token'] == 'new_access_token'
assert mock_session_obj.post.call_args.kwargs['json'] == { assert mock_session_obj.post.call_args.kwargs['json'] == {
'code': 'auth_code', 'code': 'auth_code',
'redirect_uri': 'https://oss.example/auth/space/callback',
'instance_id': constants.instance_id, 'instance_id': constants.instance_id,
'workspace_uuids': ['workspace-1'], 'workspace_uuids': ['workspace-1'],
'workspace_created_ats': {'workspace-1': 1_700_000_000}, 'workspace_created_ats': {'workspace-1': 1_700_000_000},
@@ -846,10 +850,7 @@ class TestSpaceServiceGetModelSelection:
if response_shape == 'models-envelope': if response_shape == 'models-envelope':
data = {'models': models} data = {'models': models}
elif response_shape == 'availability-wrapper': elif response_shape == 'availability-wrapper':
data = [ data = [{'model': model, 'latency_ms': index + 10, 'http_code': 200} for index, model in enumerate(models)]
{'model': model, 'latency_ms': index + 10, 'http_code': 200}
for index, model in enumerate(models)
]
else: else:
data = models data = models
payload = {'code': 0, 'data': data} payload = {'code': 0, 'data': data}
@@ -4,7 +4,7 @@ import asyncio
import datetime import datetime
import logging import logging
from types import SimpleNamespace from types import SimpleNamespace
from unittest.mock import Mock from unittest.mock import AsyncMock, Mock
import pytest import pytest
import sqlalchemy import sqlalchemy
@@ -215,6 +215,88 @@ async def test_directory_delta_requests_model_catalog_sync_after_commit(projecti
request_sync.assert_called_once_with() 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): async def test_initial_snapshot_projects_core_owned_rows(projection_context):
application, session_factory = projection_context application, session_factory = projection_context
reconcile_execution_projection = Mock() reconcile_execution_projection = Mock()
+15 -3
View File
@@ -43,17 +43,24 @@ const pendingSpaceOAuthLogins = new Map<
function getOrCreateSpaceOAuthLoginPromise( function getOrCreateSpaceOAuthLoginPromise(
authCode: string, authCode: string,
state: string, state: string,
redirectUri: string,
workspaceUuid?: string, workspaceUuid?: string,
launchAssertion?: string, launchAssertion?: string,
): Promise<SpaceOAuthLoginResult> { ): Promise<SpaceOAuthLoginResult> {
const requestKey = `${authCode}:${state}:${workspaceUuid ?? ''}:${launchAssertion ?? ''}`; const requestKey = `${authCode}:${state}:${redirectUri}:${workspaceUuid ?? ''}:${launchAssertion ?? ''}`;
const pendingRequest = pendingSpaceOAuthLogins.get(requestKey); const pendingRequest = pendingSpaceOAuthLogins.get(requestKey);
if (pendingRequest) { if (pendingRequest) {
return pendingRequest; return pendingRequest;
} }
const requestPromise = httpClient const requestPromise = httpClient
.exchangeSpaceOAuthCode(authCode, state, workspaceUuid, launchAssertion) .exchangeSpaceOAuthCode(
authCode,
state,
redirectUri,
workspaceUuid,
launchAssertion,
)
.finally(() => { .finally(() => {
pendingSpaceOAuthLogins.delete(requestKey); pendingSpaceOAuthLogins.delete(requestKey);
}); });
@@ -95,6 +102,7 @@ function SpaceOAuthCallbackContent() {
const response = await getOrCreateSpaceOAuthLoginPromise( const response = await getOrCreateSpaceOAuthLoginPromise(
authCode, authCode,
state, state,
`${window.location.origin}/auth/space/callback`,
workspaceUuid, workspaceUuid,
launchAssertion, launchAssertion,
); );
@@ -195,7 +203,11 @@ function SpaceOAuthCallbackContent() {
async (authCode: string, state: string) => { async (authCode: string, state: string) => {
setIsProcessing(true); setIsProcessing(true);
try { 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) { if (!isMountedRef.current) {
return; return;
} }
+6 -9
View File
@@ -1365,6 +1365,7 @@ export class BackendClient extends BaseHttpClient {
public async bindSpaceAccount( public async bindSpaceAccount(
code: string, code: string,
state: string, state: string,
redirectUri: string,
): Promise<{ ): Promise<{
token: string; token: string;
user: string; user: string;
@@ -1372,7 +1373,7 @@ export class BackendClient extends BaseHttpClient {
}> { }> {
const response = await this.instance.post( const response = await this.instance.post(
'/api/v1/user/bind-space', '/api/v1/user/bind-space',
{ code, state }, { code, state, redirect_uri: redirectUri },
{ skipWorkspace: true } as RequestConfig, { skipWorkspace: true } as RequestConfig,
); );
if (response.data.code !== 0) { if (response.data.code !== 0) {
@@ -1385,18 +1386,12 @@ export class BackendClient extends BaseHttpClient {
} }
// ============ Space OAuth API (Redirect Flow) ============ // ============ Space OAuth API (Redirect Flow) ============
public getSpaceAuthorizeUrl( public getSpaceAuthorizeUrl(redirectUri: string): Promise<{
redirectUri: string,
options?: { cloudEntry?: boolean },
): Promise<{
authorize_url: string; authorize_url: string;
}> { }> {
return this.get( return this.get(
'/api/v1/user/space/authorize-url', '/api/v1/user/space/authorize-url',
{ { redirect_uri: redirectUri },
redirect_uri: redirectUri,
...(options?.cloudEntry ? { cloud_entry: '1' } : {}),
},
{ skipWorkspace: true }, { skipWorkspace: true },
); );
} }
@@ -1414,6 +1409,7 @@ export class BackendClient extends BaseHttpClient {
public async exchangeSpaceOAuthCode( public async exchangeSpaceOAuthCode(
code: string, code: string,
state: string, state: string,
redirectUri: string,
workspaceUuid?: string, workspaceUuid?: string,
launchAssertion?: string, launchAssertion?: string,
): Promise<{ ): Promise<{
@@ -1428,6 +1424,7 @@ export class BackendClient extends BaseHttpClient {
{ {
code, code,
state, state,
redirect_uri: redirectUri,
workspace_uuid: workspaceUuid, workspace_uuid: workspaceUuid,
launch_assertion: launchAssertion, launch_assertion: launchAssertion,
}, },
+1 -7
View File
@@ -202,13 +202,7 @@ export default function Login() {
try { try {
const currentOrigin = window.location.origin; const currentOrigin = window.location.origin;
const redirectUri = `${currentOrigin}/auth/space/callback`; const redirectUri = `${currentOrigin}/auth/space/callback`;
const response = await httpClient.getSpaceAuthorizeUrl(redirectUri, { 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(),
});
window.location.href = response.authorize_url; window.location.href = response.authorize_url;
} catch { } catch {
toast.error(t('common.spaceLoginFailed')); toast.error(t('common.spaceLoginFailed'));
@@ -7,11 +7,13 @@ const source = fs.readFileSync(
'utf8', 'utf8',
); );
test('normal Cloud login enters through the Space Workspace launcher', () => { test('normal Cloud login uses the standard Space OAuth callback path', () => {
assert.match(source, /cloudEntry:\s*!getPendingInvitationToken\(\)/); assert.doesNotMatch(source, /cloudEntry/);
assert.match(source, /getSpaceAuthorizeUrl\(redirectUri,\s*\{/); assert.match(source, /getSpaceAuthorizeUrl\(redirectUri\)/);
}); });
test('invitation login remains on the OAuth callback path', () => { test('invitation login uses the same OAuth callback before accepting the invitation', () => {
assert.match(source, /cloudEntry:\s*!getPendingInvitationToken\(\)/); assert.doesNotMatch(source, /cloudEntry/);
assert.match(source, /const invitationToken = getPendingInvitationToken\(\)/);
assert.match(source, /acceptWorkspaceInvitation\(invitationToken\)/);
}); });