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>

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')

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}

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()