feat(tenancy): implement workspace isolation

tests/integration/persistence/test_workspace_migration.py

@@ -0,0 +1,379 @@
+from __future__ import annotations
+
+import logging
+import uuid
+
+import pytest
+import sqlalchemy as sa
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.ext.asyncio import create_async_engine
+
+from langbot.pkg.entity import persistence
+from langbot.pkg.entity.persistence.base import Base
+from langbot.pkg.entity.persistence.user import User
+from langbot.pkg.persistence.mgr import PersistenceManager
+from langbot.pkg.persistence.alembic_runner import (
+    get_alembic_current,
+    run_alembic_downgrade,
+    run_alembic_stamp,
+    run_alembic_upgrade,
+)
+from langbot.pkg.utils import constants
+from langbot.pkg.utils import importutil
+from langbot.pkg.workspace.collaboration import normalize_email
+
+
+pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
+
+
+async def _create_legacy_schema(
+    engine,
+    *,
+    include_instance_uuid: bool = True,
+    include_users: bool = True,
+) -> None:
+    legacy_metadata = sa.MetaData()
+    metadata_table = sa.Table(
+        'metadata',
+        legacy_metadata,
+        sa.Column('key', sa.String(255), primary_key=True),
+        sa.Column('value', sa.String(255)),
+    )
+    users = sa.Table(
+        'users',
+        legacy_metadata,
+        sa.Column('id', sa.Integer, primary_key=True),
+        sa.Column('user', sa.String(255), nullable=False),
+        sa.Column('password', sa.String(255), nullable=False),
+        sa.Column('account_type', sa.String(32), nullable=False, server_default='local'),
+        sa.Column('space_account_uuid', sa.String(255), nullable=True),
+        sa.Column('space_access_token', sa.Text, nullable=True),
+        sa.Column('space_refresh_token', sa.Text, nullable=True),
+        sa.Column('space_access_token_expires_at', sa.DateTime, nullable=True),
+        sa.Column('space_api_key', sa.String(255), nullable=True),
+        sa.Column('created_at', sa.DateTime, nullable=False, server_default=sa.func.now()),
+        sa.Column('updated_at', sa.DateTime, nullable=False, server_default=sa.func.now()),
+    )
+    async with engine.begin() as conn:
+        await conn.run_sync(legacy_metadata.create_all)
+        await conn.execute(metadata_table.insert().values(key='database_version', value='25'))
+        if include_instance_uuid:
+            await conn.execute(metadata_table.insert().values(key='instance_uuid', value='instance_migration_test'))
+        if include_users:
+            await conn.execute(
+                users.insert(),
+                [
+                    {'user': 'owner@example.com', 'password': 'owner-hash'},
+                    {'user': 'member@example.com', 'password': 'member-hash'},
+                ],
+            )
+
+
+@pytest.fixture
+async def legacy_engine(tmp_path):
+    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "legacy-workspace.db"}')
+    await _create_legacy_schema(engine)
+    await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
+    yield engine
+    await engine.dispose()
+
+
+async def test_legacy_instance_gets_stable_accounts_and_default_workspace(legacy_engine):
+    await run_alembic_upgrade(legacy_engine, 'head')
+
+    async with legacy_engine.connect() as conn:
+        tables = set(await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names()))
+        assert {
+            'workspaces',
+            'workspace_memberships',
+            'workspace_invitations',
+            'workspace_execution_states',
+        }.issubset(tables)
+
+        accounts = (
+            (await conn.execute(sa.text('SELECT id, uuid, status, source, projection_revision FROM users ORDER BY id')))
+            .mappings()
+            .all()
+        )
+        assert len(accounts) == 2
+        assert len({account['uuid'] for account in accounts}) == 2
+        for account in accounts:
+            uuid.UUID(account['uuid'])
+            assert account['status'] == 'active'
+            assert account['source'] == 'local'
+            assert account['projection_revision'] == 0
+
+        workspace = (
+            (await conn.execute(sa.text('SELECT * FROM workspaces WHERE source = :source'), {'source': 'local'}))
+            .mappings()
+            .one()
+        )
+        assert workspace['instance_uuid'] == 'instance_migration_test'
+        assert workspace['slug'] == 'default'
+        assert workspace['status'] == 'active'
+        assert workspace['created_by_account_uuid'] == accounts[0]['uuid']
+
+        membership = (await conn.execute(sa.text('SELECT * FROM workspace_memberships'))).mappings().one()
+        assert membership['workspace_uuid'] == workspace['uuid']
+        assert membership['account_uuid'] == accounts[0]['uuid']
+        assert membership['role'] == 'owner'
+        assert membership['status'] == 'active'
+
+        execution_state = (await conn.execute(sa.text('SELECT * FROM workspace_execution_states'))).mappings().one()
+        assert execution_state['workspace_uuid'] == workspace['uuid']
+        assert execution_state['instance_uuid'] == 'instance_migration_test'
+        assert execution_state['active_generation'] == 1
+        assert execution_state['state'] == 'active'
+        assert execution_state['write_fenced'] in (False, 0)
+
+    assert await get_alembic_current(legacy_engine) == '0010_scope_resources'
+
+
+async def test_workspace_upgrade_is_idempotent_and_preserves_identifiers(legacy_engine):
+    await run_alembic_upgrade(legacy_engine, 'head')
+    async with legacy_engine.connect() as conn:
+        account_uuids_before = (await conn.execute(sa.text('SELECT uuid FROM users ORDER BY id'))).scalars().all()
+        workspace_uuid_before = (
+            await conn.execute(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'"))
+        ).scalar_one()
+
+    await run_alembic_upgrade(legacy_engine, 'head')
+
+    async with legacy_engine.connect() as conn:
+        account_uuids_after = (await conn.execute(sa.text('SELECT uuid FROM users ORDER BY id'))).scalars().all()
+        workspace_uuid_after = (
+            await conn.execute(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'"))
+        ).scalar_one()
+    assert account_uuids_after == account_uuids_before
+    assert workspace_uuid_after == workspace_uuid_before
+
+
+async def test_workspace_kernel_upgrade_downgrade_upgrade_round_trip(tmp_path):
+    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-round-trip.db"}')
+    try:
+        await _create_legacy_schema(engine)
+        await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
+        await run_alembic_upgrade(engine, '0009_workspace_tenancy')
+        assert await get_alembic_current(engine) == '0009_workspace_tenancy'
+
+        await run_alembic_downgrade(engine, '0008_mcp_resource_prefs')
+        assert await get_alembic_current(engine) == '0008_mcp_resource_prefs'
+        async with engine.connect() as conn:
+            tables = set(await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names()))
+            user_columns = {
+                column['name']
+                for column in await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_columns('users'))
+            }
+            accounts = (await conn.execute(sa.text('SELECT user, password FROM users ORDER BY id'))).all()
+        assert (
+            not {
+                'workspaces',
+                'workspace_memberships',
+                'workspace_invitations',
+                'workspace_execution_states',
+            }
+            & tables
+        )
+        assert not {'uuid', 'status', 'source', 'projection_revision'} & user_columns
+        assert accounts == [
+            ('owner@example.com', 'owner-hash'),
+            ('member@example.com', 'member-hash'),
+        ]
+
+        await run_alembic_upgrade(engine, '0009_workspace_tenancy')
+        assert await get_alembic_current(engine) == '0009_workspace_tenancy'
+        async with engine.connect() as conn:
+            assert await conn.scalar(sa.text('SELECT COUNT(*) FROM workspaces')) == 1
+            assert await conn.scalar(sa.text('SELECT COUNT(*) FROM workspace_memberships')) == 1
+    finally:
+        await engine.dispose()
+
+
+@pytest.mark.parametrize(
+    ('raw_email', 'expected_email'),
+    [
+        ('Straße@Example.COM', 'strasse@example.com'),
+        ('Ꭰ@Example.COM', 'Ꭰ@example.com'),
+    ],
+)
+async def test_workspace_upgrade_uses_runtime_unicode_email_normalization(
+    tmp_path,
+    raw_email,
+    expected_email,
+):
+    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "unicode-email.db"}')
+    try:
+        await _create_legacy_schema(engine, include_users=False)
+        async with engine.begin() as conn:
+            await conn.execute(
+                sa.text('INSERT INTO users (user, password, account_type) VALUES (:email, :password, :type)'),
+                {'email': raw_email, 'password': 'owner-hash', 'type': 'local'},
+            )
+        await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
+        await run_alembic_upgrade(engine, 'head')
+
+        async with engine.connect() as conn:
+            assert await conn.scalar(sa.text('SELECT normalized_email FROM users')) == expected_email
+    finally:
+        await engine.dispose()
+
+
+async def test_fresh_sqlite_schema_accepts_application_casefold_identity(tmp_path):
+    importutil.import_modules_in_pkg(persistence)
+    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "fresh-unicode-email.db"}')
+    canonical_email = normalize_email('Ꭰ@Example.COM')
+    try:
+        async with engine.begin() as conn:
+            await conn.run_sync(Base.metadata.create_all)
+            await conn.execute(
+                sa.insert(User).values(
+                    uuid='00000000-0000-0000-0000-000000000099',
+                    user=canonical_email,
+                    normalized_email=canonical_email,
+                    password='hash',
+                )
+            )
+        async with engine.connect() as conn:
+            assert await conn.scalar(sa.select(User.normalized_email)) == 'Ꭰ@example.com'
+    finally:
+        await engine.dispose()
+
+
+async def test_workspace_upgrade_rejects_unicode_casefold_duplicate_accounts(tmp_path):
+    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "unicode-email-duplicate.db"}')
+    try:
+        await _create_legacy_schema(engine, include_users=False)
+        async with engine.begin() as conn:
+            await conn.execute(
+                sa.text(
+                    'INSERT INTO users (user, password, account_type) VALUES '
+                    "('Straße@Example.COM', 'first-hash', 'local'), "
+                    "('STRASSE@example.com', 'second-hash', 'local')"
+                )
+            )
+        await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
+        with pytest.raises(RuntimeError, match='both normalize'):
+            await run_alembic_upgrade(engine, 'head')
+    finally:
+        await engine.dispose()
+
+
+async def test_uninitialized_instance_gets_ownerless_default_workspace(tmp_path):
+    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "uninitialized-instance.db"}')
+    try:
+        await _create_legacy_schema(engine, include_users=False)
+        await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
+        await run_alembic_upgrade(engine, 'head')
+
+        async with engine.connect() as conn:
+            workspace = (await conn.execute(sa.text('SELECT * FROM workspaces'))).mappings().one()
+            membership_count = await conn.scalar(sa.text('SELECT COUNT(*) FROM workspace_memberships'))
+            execution_state = (await conn.execute(sa.text('SELECT * FROM workspace_execution_states'))).mappings().one()
+
+        assert workspace['created_by_account_uuid'] is None
+        assert membership_count == 0
+        assert execution_state['workspace_uuid'] == workspace['uuid']
+        assert execution_state['active_generation'] == 1
+    finally:
+        await engine.dispose()
+
+
+async def test_local_workspace_unique_index_allows_cloud_projections(legacy_engine):
+    await run_alembic_upgrade(legacy_engine, 'head')
+
+    async with legacy_engine.begin() as conn:
+        await conn.execute(
+            sa.text(
+                'INSERT INTO workspaces '
+                '(uuid, instance_uuid, name, slug, type, status, source, projection_revision) '
+                'VALUES (:uuid, :instance_uuid, :name, :slug, :type, :status, :source, 0)'
+            ),
+            {
+                'uuid': str(uuid.uuid4()),
+                'instance_uuid': 'instance_migration_test',
+                'name': 'Cloud Projection',
+                'slug': 'cloud-projection',
+                'type': 'team',
+                'status': 'active',
+                'source': 'cloud_projection',
+            },
+        )
+
+    with pytest.raises(IntegrityError):
+        async with legacy_engine.begin() as conn:
+            await conn.execute(
+                sa.text(
+                    'INSERT INTO workspaces '
+                    '(uuid, instance_uuid, name, slug, type, status, source, projection_revision) '
+                    'VALUES (:uuid, :instance_uuid, :name, :slug, :type, :status, :source, 0)'
+                ),
+                {
+                    'uuid': str(uuid.uuid4()),
+                    'instance_uuid': 'instance_migration_test',
+                    'name': 'Second Local',
+                    'slug': 'second-local',
+                    'type': 'team',
+                    'status': 'active',
+                    'source': 'local',
+                },
+            )
+
+
+async def test_legacy_instance_without_bound_instance_uuid_fails_closed(tmp_path):
+    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "missing-instance.db"}')
+    try:
+        await _create_legacy_schema(engine, include_instance_uuid=False)
+        await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
+        with pytest.raises(RuntimeError, match='instance_uuid'):
+            await run_alembic_upgrade(engine, 'head')
+    finally:
+        await engine.dispose()
+
+
+async def test_persistence_startup_defers_workspace_tables_until_account_upgrade(tmp_path, monkeypatch):
+    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "startup-order.db"}')
+    try:
+        await _create_legacy_schema(engine)
+        await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
+        monkeypatch.setattr(constants, 'instance_id', 'instance_migration_test')
+
+        database = type('Database', (), {'get_engine': lambda self: engine})()
+        application = type('Application', (), {})()
+        application.logger = logging.getLogger('workspace-startup-test')
+        manager = PersistenceManager(application)
+        manager.db = database
+
+        await manager.create_tables()
+        async with engine.connect() as conn:
+            tables_before_migration = set(
+                await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names())
+            )
+        assert 'workspaces' not in tables_before_migration
+
+        await manager._run_alembic_migrations()
+
+        async with engine.connect() as conn:
+            workspace = (
+                (await conn.execute(sa.text("SELECT * FROM workspaces WHERE source = 'local'"))).mappings().one()
+            )
+        assert workspace['instance_uuid'] == 'instance_migration_test'
+    finally:
+        await engine.dispose()
+
+
+async def test_persistence_startup_rejects_instance_uuid_drift(tmp_path, monkeypatch):
+    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "instance-drift.db"}')
+    try:
+        await _create_legacy_schema(engine)
+        monkeypatch.setattr(constants, 'instance_id', 'different_instance')
+
+        database = type('Database', (), {'get_engine': lambda self: engine})()
+        application = type('Application', (), {})()
+        application.logger = logging.getLogger('workspace-instance-drift-test')
+        manager = PersistenceManager(application)
+        manager.db = database
+
+        with pytest.raises(RuntimeError, match='does not match'):
+            await manager.create_tables()
+    finally:
+        await engine.dispose()