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