xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-09-04 00:27:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(cloud): provision login workspace just in time (#2505)

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files
src/langbot/pkg/api/http/controller/groups/user.py
+39 -52
View File
@@ -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={
src/langbot/pkg/api/http/service/space.py
+4 -6
View File
@@ -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.
src/langbot/pkg/api/http/service/user.py
+3 -2
View File
@@ -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)
src/langbot/pkg/cloud/directory_projection.py
+84 -1
View File
@@ -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)
tests/integration/api/test_user_space_oauth.py
+159 -74
View File
@@ -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')
tests/unit_tests/api/service/test_space_service.py
+6 -5
View File
@@ -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}
tests/unit_tests/cloud/test_directory_projection.py
+83 -1
View File
@@ -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()
web/src/app/auth/space/callback/page.tsx
+15 -3
View File
@@ -43,17 +43,24 @@ const pendingSpaceOAuthLogins = new Map<
function getOrCreateSpaceOAuthLoginPromise(
authCode: string,
state: string,
redirectUri: string,
workspaceUuid?: string,
launchAssertion?: string,
): Promise<SpaceOAuthLoginResult> {
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;
}
web/src/app/infra/http/BackendClient.ts
+6 -9
View File
@@ -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,
},
web/src/app/login/page.tsx
+1 -7
View File
@@ -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'));
web/tests/unit/cloud-new-account-entry.test.mjs
+7 -5
View File
@@ -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\)/);
});