feat(tenancy): implement workspace isolation

2026-07-21 20:06:06 +00:00 · 2026-07-19 09:58:59 +08:00
parent 9eb292992d
commit c6f826fe2d
271 changed files with 31162 additions and 6106 deletions
@@ -0,0 +1,310 @@
+from __future__ import annotations
+
+import asyncio
+import uuid
+from types import SimpleNamespace
+
+import pytest
+import sqlalchemy
+from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
+
+from langbot.pkg.entity.persistence.base import Base
+from langbot.pkg.entity.persistence.user import User
+from langbot.pkg.entity.persistence.workspace import (
+    Workspace,
+    WorkspaceExecutionState,
+    WorkspaceInvitation,
+    WorkspaceMembership,
+)
+from langbot.pkg.workspace.collaboration import (
+    InvitationEmailMismatchError,
+    InvitationRoleError,
+    InvitationUsedError,
+    LastOwnerError,
+    MembershipPermissionError,
+    WorkspaceCollaborationService,
+)
+from langbot.pkg.workspace.errors import WorkspaceNotFoundError
+from langbot.pkg.workspace.service import WorkspaceService
+from langbot.pkg.workspace.policy import CloudWorkspacePolicy
+
+
+pytestmark = pytest.mark.asyncio
+
+
+@pytest.fixture
+async def collaboration_context(tmp_path):
+    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-collaboration.db"}')
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+
+    persistence_mgr = SimpleNamespace(get_db_engine=lambda: engine)
+    application = SimpleNamespace(persistence_mgr=persistence_mgr)
+    workspace_service = WorkspaceService(application, instance_uuid='instance-collaboration-test')
+    service = WorkspaceCollaborationService(application, workspace_service)
+    session_factory = async_sessionmaker(engine, expire_on_commit=False)
+
+    async with session_factory() as session:
+        async with session.begin():
+            owner = User(
+                uuid=str(uuid.uuid4()),
+                user='owner@example.com',
+                normalized_email='owner@example.com',
+                password='owner-hash',
+                account_type='local',
+            )
+            session.add(owner)
+            await session.flush()
+            workspace, owner_membership = await workspace_service.bootstrap_local_account(
+                owner.uuid,
+                session=session,
+            )
+
+    yield service, workspace_service, session_factory, owner, workspace, owner_membership
+    await engine.dispose()
+
+
+async def _add_account(session_factory, email: str) -> User:
+    async with session_factory() as session:
+        async with session.begin():
+            account = User(
+                uuid=str(uuid.uuid4()),
+                user=email,
+                normalized_email=email.strip().casefold(),
+                password='member-hash',
+                account_type='local',
+            )
+            session.add(account)
+            await session.flush()
+            return account
+
+
+async def test_invitation_secret_is_hashed_and_acceptance_is_one_time(collaboration_context):
+    service, _, session_factory, _, workspace, owner_membership = collaboration_context
+    created = await service.create_invitation(
+        workspace.uuid,
+        owner_membership,
+        'member@example.com',
+        'developer',
+    )
+
+    assert created.token.startswith('lbi_')
+    assert created.invitation.token_hash != created.token
+    assert created.token not in created.invitation.token_hash
+
+    account = await _add_account(session_factory, 'MEMBER@example.com')
+    membership = await service.accept_invitation(created.token, account.uuid)
+    assert membership.workspace_uuid == workspace.uuid
+    assert membership.role == 'developer'
+
+    with pytest.raises(InvitationUsedError):
+        await service.accept_invitation(created.token, account.uuid)
+
+    async with session_factory() as session:
+        persisted = await session.get(WorkspaceInvitation, created.invitation.uuid)
+        assert persisted is not None
+        assert persisted.status == 'accepted'
+        assert not hasattr(persisted, 'token')
+
+
+async def test_concurrent_invitation_acceptance_creates_one_membership(collaboration_context):
+    service, _, session_factory, _, workspace, owner_membership = collaboration_context
+    created = await service.create_invitation(
+        workspace.uuid,
+        owner_membership,
+        'race@example.com',
+        'viewer',
+    )
+    account = await _add_account(session_factory, 'race@example.com')
+
+    results = await asyncio.gather(
+        service.accept_invitation(created.token, account.uuid),
+        service.accept_invitation(created.token, account.uuid),
+        return_exceptions=True,
+    )
+
+    assert sum(isinstance(result, WorkspaceMembership) for result in results) == 1
+    assert sum(isinstance(result, InvitationUsedError) for result in results) == 1
+    async with session_factory() as session:
+        count = await session.scalar(
+            sqlalchemy.select(sqlalchemy.func.count())
+            .select_from(WorkspaceMembership)
+            .where(
+                WorkspaceMembership.workspace_uuid == workspace.uuid,
+                WorkspaceMembership.account_uuid == account.uuid,
+            )
+        )
+        assert count == 1
+
+
+async def test_invitation_rejects_owner_role_and_email_mismatch(collaboration_context):
+    service, _, session_factory, _, workspace, owner_membership = collaboration_context
+    with pytest.raises(InvitationRoleError):
+        await service.create_invitation(
+            workspace.uuid,
+            owner_membership,
+            'member@example.com',
+            'owner',
+        )
+
+    created = await service.create_invitation(
+        workspace.uuid,
+        owner_membership,
+        'expected@example.com',
+        'viewer',
+    )
+    wrong_account = await _add_account(session_factory, 'wrong@example.com')
+    with pytest.raises(InvitationEmailMismatchError):
+        await service.accept_invitation(created.token, wrong_account.uuid)
+
+
+async def test_last_owner_cannot_be_demoted(collaboration_context):
+    service, _, session_factory, _, workspace, owner_membership = collaboration_context
+    created = await service.create_invitation(
+        workspace.uuid,
+        owner_membership,
+        'second@example.com',
+        'admin',
+    )
+    second = await _add_account(session_factory, 'second@example.com')
+    second_membership = await service.accept_invitation(created.token, second.uuid)
+
+    with pytest.raises(LastOwnerError):
+        await service.update_member_role(
+            workspace.uuid,
+            owner_membership.account_uuid,
+            'admin',
+            owner_membership,
+        )
+
+    with pytest.raises(MembershipPermissionError):
+        await service.update_member_role(
+            workspace.uuid,
+            owner_membership.account_uuid,
+            'viewer',
+            second_membership,
+        )
+
+    promoted = await service.update_member_role(
+        workspace.uuid,
+        second.uuid,
+        'owner',
+        owner_membership,
+    )
+    assert promoted.role == 'owner'
+    demoted = await service.update_member_role(
+        workspace.uuid,
+        owner_membership.account_uuid,
+        'admin',
+        owner_membership,
+    )
+    assert demoted.role == 'admin'
+
+
+async def test_workspace_selector_requires_membership(collaboration_context):
+    service, _, session_factory, _, workspace, _ = collaboration_context
+    outsider = await _add_account(session_factory, 'outsider@example.com')
+
+    with pytest.raises(WorkspaceNotFoundError):
+        await service.resolve_account_workspace(outsider.uuid, workspace.uuid)
+
+
+async def test_cloud_directory_requires_explicit_selector_and_rejects_local_mutation(
+    collaboration_context,
+):
+    _service, workspace_service, session_factory, owner, _workspace, _membership = collaboration_context
+    cloud_workspace_uuid = str(uuid.uuid4())
+    cloud_membership = WorkspaceMembership(
+        uuid=str(uuid.uuid4()),
+        workspace_uuid=cloud_workspace_uuid,
+        account_uuid=owner.uuid,
+        role='owner',
+        status='active',
+        projection_revision=4,
+    )
+    async with session_factory.begin() as session:
+        session.add(
+            Workspace(
+                uuid=cloud_workspace_uuid,
+                instance_uuid='instance-collaboration-test',
+                name='Projected Workspace',
+                slug='projected-workspace',
+                source='cloud_projection',
+                projection_revision=4,
+            )
+        )
+        session.add(
+            WorkspaceExecutionState(
+                workspace_uuid=cloud_workspace_uuid,
+                instance_uuid='instance-collaboration-test',
+                active_generation=8,
+                state='active',
+                write_fenced=False,
+                source='cloud',
+                desired_state_revision=8,
+            )
+        )
+        session.add(cloud_membership)
+
+    cloud_service = WorkspaceCollaborationService(
+        workspace_service.ap,
+        WorkspaceService(
+            workspace_service.ap,
+            policy=CloudWorkspacePolicy(),
+            instance_uuid='instance-collaboration-test',
+        ),
+    )
+    with pytest.raises(WorkspaceNotFoundError):
+        await cloud_service.resolve_account_workspace(owner.uuid, None)
+
+    access = await cloud_service.resolve_account_workspace(owner.uuid, cloud_workspace_uuid)
+    assert access.workspace.uuid == cloud_workspace_uuid
+    assert access.execution.placement_generation == 8
+
+    with pytest.raises(MembershipPermissionError, match='control plane'):
+        await cloud_service.create_invitation(
+            cloud_workspace_uuid,
+            cloud_membership,
+            'member@example.com',
+            'viewer',
+        )
+
+
+async def test_invitation_lock_registry_does_not_retain_sequential_tokens(collaboration_context):
+    service, *_ = collaboration_context
+
+    for index in range(100):
+        async with service._invitation_lock(f'token-{index}'):
+            pass
+
+    assert service._invitation_locks == {}
+
+
+async def test_invitation_lock_registry_keeps_waiters_serialized(collaboration_context):
+    service, *_ = collaboration_context
+    first_entered = asyncio.Event()
+    release_first = asyncio.Event()
+    order: list[str] = []
+
+    async def first_worker() -> None:
+        async with service._invitation_lock('same-token'):
+            order.append('first')
+            first_entered.set()
+            await release_first.wait()
+
+    async def second_worker() -> None:
+        await first_entered.wait()
+        async with service._invitation_lock('same-token'):
+            order.append('second')
+
+    first_task = asyncio.create_task(first_worker())
+    second_task = asyncio.create_task(second_worker())
+    await first_entered.wait()
+    await asyncio.sleep(0)
+
+    assert order == ['first']
+    release_first.set()
+    await asyncio.gather(first_task, second_task)
+
+    assert order == ['first', 'second']
+    assert service._invitation_locks == {}
@@ -0,0 +1,19 @@
+from langbot.pkg.workspace.policy import (
+    CloudWorkspacePolicy,
+    SingleWorkspacePolicy,
+    open_core_workspace_policy,
+)
+
+
+def test_open_source_policy_is_single_workspace() -> None:
+    policy = open_core_workspace_policy()
+
+    assert policy.workspace_limit == 1
+    assert policy.multi_workspace_enabled is False
+
+
+def test_cloud_policy_is_not_exported_as_the_default_policy() -> None:
+    # CloudWorkspacePolicy remains a contract for the future verified, closed
+    # bootstrap.  Constructing it explicitly must not change the OSS default.
+    assert CloudWorkspacePolicy().multi_workspace_enabled is True
+    assert isinstance(open_core_workspace_policy(), SingleWorkspacePolicy)
@@ -0,0 +1,237 @@
+from __future__ import annotations
+
+import uuid
+from types import SimpleNamespace
+
+import pytest
+import sqlalchemy
+from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
+
+from langbot.pkg.entity.persistence.base import Base
+from langbot.pkg.entity.persistence.user import User
+from langbot.pkg.entity.persistence.workspace import (
+    Workspace,
+    WorkspaceExecutionSource,
+    WorkspaceExecutionState,
+    WorkspaceMembership,
+    WorkspaceSource,
+)
+from langbot.pkg.workspace import (
+    WorkspaceExecutionUnavailableError,
+    WorkspaceGenerationMismatchError,
+    WorkspaceInvariantError,
+    WorkspaceLimitExceededError,
+    WorkspaceOwnerAlreadyExistsError,
+    WorkspaceService,
+)
+from langbot.pkg.workspace.policy import CloudWorkspacePolicy
+
+
+pytestmark = pytest.mark.asyncio
+
+
+@pytest.fixture
+async def workspace_test_context(tmp_path):
+    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-service.db"}')
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+
+    persistence_mgr = SimpleNamespace(get_db_engine=lambda: engine)
+    application = SimpleNamespace(persistence_mgr=persistence_mgr)
+    session_factory = async_sessionmaker(engine, expire_on_commit=False)
+    service = WorkspaceService(application, instance_uuid='instance_service_test')
+    yield service, session_factory
+    await engine.dispose()
+
+
+async def _insert_account(session, email: str) -> str:
+    account_uuid = str(uuid.uuid4())
+    session.add(
+        User(
+            uuid=account_uuid,
+            user=email,
+            normalized_email=email.strip().casefold(),
+            password='hashed-password',
+            account_type='local',
+        )
+    )
+    await session.flush()
+    return account_uuid
+
+
+async def test_account_uuid_default_supports_existing_core_insert_path(workspace_test_context):
+    _service, session_factory = workspace_test_context
+
+    async with session_factory.begin() as session:
+        await session.execute(
+            sqlalchemy.insert(User).values(
+                user='core-insert@example.com',
+                normalized_email='core-insert@example.com',
+                password='hashed-password',
+                account_type='local',
+            )
+        )
+
+    async with session_factory() as session:
+        account = await session.scalar(sqlalchemy.select(User))
+        assert account is not None
+        uuid.UUID(account.uuid)
+        assert account.status == 'active'
+        assert account.source == 'local'
+
+
+async def test_bootstrap_local_account_uses_callers_transaction(workspace_test_context):
+    service, session_factory = workspace_test_context
+
+    async with session_factory() as session:
+        async with session.begin():
+            account_uuid = await _insert_account(session, 'owner@example.com')
+            workspace, membership = await service.bootstrap_local_account(account_uuid, session=session)
+            assert session.in_transaction()
+            assert workspace.created_by_account_uuid == account_uuid
+            assert membership.account_uuid == account_uuid
+            assert membership.role == 'owner'
+
+    async with session_factory() as session:
+        persisted_workspace = await session.scalar(sqlalchemy.select(Workspace))
+        persisted_membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
+        execution_state = await session.scalar(sqlalchemy.select(WorkspaceExecutionState))
+        assert persisted_workspace is not None
+        assert persisted_membership is not None
+        assert execution_state is not None
+        assert execution_state.workspace_uuid == persisted_workspace.uuid
+        assert execution_state.instance_uuid == 'instance_service_test'
+        assert execution_state.active_generation == 1
+        assert execution_state.write_fenced is False
+
+
+async def test_ensure_singleton_workspace_is_idempotent(workspace_test_context):
+    service, session_factory = workspace_test_context
+
+    first = await service.ensure_singleton_workspace()
+    second = await service.ensure_singleton_workspace()
+
+    assert second.uuid == first.uuid
+    async with session_factory() as session:
+        assert await session.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(Workspace)) == 1
+        assert (
+            await session.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(WorkspaceExecutionState)) == 1
+        )
+
+
+async def test_create_local_workspace_rejects_second_workspace(workspace_test_context):
+    service, session_factory = workspace_test_context
+
+    await service.create_local_workspace(name='First', slug='first')
+    with pytest.raises(WorkspaceLimitExceededError) as exc_info:
+        await service.create_local_workspace(name='Second', slug='second')
+
+    assert exc_info.value.code == 'edition_limit'
+    async with session_factory() as session:
+        assert await session.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(Workspace)) == 1
+
+
+async def test_initial_owner_cannot_be_claimed_by_another_account(workspace_test_context):
+    service, session_factory = workspace_test_context
+
+    async with session_factory() as session:
+        async with session.begin():
+            first_account_uuid = await _insert_account(session, 'first@example.com')
+            second_account_uuid = await _insert_account(session, 'second@example.com')
+            await service.bootstrap_local_account(first_account_uuid, session=session)
+
+    with pytest.raises(WorkspaceOwnerAlreadyExistsError):
+        await service.claim_initial_owner(second_account_uuid)
+
+    async with session_factory() as session:
+        owners = (
+            await session.scalars(sqlalchemy.select(WorkspaceMembership).where(WorkspaceMembership.role == 'owner'))
+        ).all()
+        assert len(owners) == 1
+        assert owners[0].account_uuid == first_account_uuid
+
+
+async def test_execution_binding_returns_persisted_generation(workspace_test_context):
+    service, session_factory = workspace_test_context
+    workspace = await service.ensure_singleton_workspace()
+
+    async with session_factory.begin() as session:
+        execution_state = await session.get(WorkspaceExecutionState, workspace.uuid)
+        execution_state.active_generation = 7
+
+    binding = await service.get_local_execution_binding(
+        workspace.uuid,
+        expected_generation=7,
+    )
+
+    assert binding.instance_uuid == 'instance_service_test'
+    assert binding.workspace_uuid == workspace.uuid
+    assert binding.placement_generation == 7
+    assert binding.state == 'active'
+    assert binding.write_fenced is False
+
+    with pytest.raises(WorkspaceGenerationMismatchError):
+        await service.get_local_execution_binding(workspace.uuid, expected_generation=6)
+
+
+async def test_execution_binding_fails_closed_when_fenced(workspace_test_context):
+    service, session_factory = workspace_test_context
+    workspace = await service.ensure_singleton_workspace()
+
+    async with session_factory.begin() as session:
+        execution_state = await session.get(WorkspaceExecutionState, workspace.uuid)
+        execution_state.write_fenced = True
+
+    with pytest.raises(WorkspaceExecutionUnavailableError):
+        await service.get_local_execution_context(workspace.uuid)
+
+
+async def test_cloud_projection_requires_explicit_general_binding(workspace_test_context):
+    service, session_factory = workspace_test_context
+    await service.ensure_singleton_workspace()
+    cloud_workspace_uuid = str(uuid.uuid4())
+
+    async with session_factory.begin() as session:
+        session.add(
+            Workspace(
+                uuid=cloud_workspace_uuid,
+                instance_uuid='instance_service_test',
+                name='Cloud Workspace',
+                slug='cloud-workspace',
+                source=WorkspaceSource.CLOUD_PROJECTION.value,
+            )
+        )
+        session.add(
+            WorkspaceExecutionState(
+                workspace_uuid=cloud_workspace_uuid,
+                instance_uuid='instance_service_test',
+                active_generation=9,
+                state='active',
+                write_fenced=False,
+                source=WorkspaceExecutionSource.CLOUD.value,
+            )
+        )
+
+    binding = await service.get_execution_binding(
+        cloud_workspace_uuid,
+        expected_generation=9,
+    )
+    assert binding.workspace_uuid == cloud_workspace_uuid
+    assert binding.placement_generation == 9
+
+    with pytest.raises(WorkspaceInvariantError):
+        await service.get_local_execution_binding(cloud_workspace_uuid)
+
+
+async def test_cloud_policy_never_creates_or_guesses_a_workspace(workspace_test_context):
+    service, _session_factory = workspace_test_context
+    cloud_service = WorkspaceService(
+        service.ap,
+        policy=CloudWorkspacePolicy(),
+        instance_uuid='instance_service_test',
+    )
+
+    with pytest.raises(WorkspaceLimitExceededError):
+        await cloud_service.create_local_workspace(name='Forbidden', slug='forbidden')
+
+    assert cloud_service.policy.multi_workspace_enabled is True