fix(cloud): track workspace membership provenance (#2406)

* fix(cloud): converge legacy revision-zero members

* fix(cloud): track workspace membership provenance

* test(persistence): track current migration head

---------

Co-authored-by: Chan <dadachann@users.noreply.github.com>

src/langbot/pkg/cloud/directory_projection.py

@@ -15,6 +15,7 @@ from ..entity.persistence.cloud_directory import DirectoryProjectionInbox, Direc
 from ..entity.persistence.user import AccountSource, AccountStatus, User
 from ..entity.persistence.workspace import (
     MembershipRole,
+    MembershipSource,
     MembershipStatus,
     Workspace,
     WorkspaceExecutionSource,
@@ -886,15 +887,15 @@ class DirectoryProjectionService:
                         account_uuid=member.account_uuid,
                         role=role,
                         status=status,
+                        source=MembershipSource.CLOUD_PROJECTION.value,
                         joined_at=joined_at,
                         projection_revision=member.projection_revision,
                     )
                 )
                 continue
-            if membership.projection_revision == 0:
-                # Revision zero is Core-owned collaboration state. Directory
-                # projection seeds memberships, but must not overwrite later
-                # invitation, role, or removal decisions made by Core.
+            if membership.source != MembershipSource.CLOUD_PROJECTION.value:
+                # Core-owned collaboration state is never adopted based on
+                # account provenance, revision, or matching account identity.
                 continue
             if membership.uuid != member.membership_uuid:
                 raise DirectoryProjectionUnavailableError('Directory membership UUID changed for one account')
@@ -906,11 +907,12 @@ class DirectoryProjectionService:
                 raise DirectoryProjectionUnavailableError('Directory membership revision has conflicting contents')
             membership.role = role
             membership.status = status
+            membership.source = MembershipSource.CLOUD_PROJECTION.value
             membership.joined_at = joined_at
             membership.projection_revision = member.projection_revision
 
         for account_uuid, membership in existing.items():
-            if account_uuid not in included_accounts and membership.projection_revision != 0:
+            if account_uuid not in included_accounts and membership.source == MembershipSource.CLOUD_PROJECTION.value:
                 membership.status = MembershipStatus.REMOVED.value
                 membership.projection_revision = max(
                     int(membership.projection_revision),

src/langbot/pkg/entity/persistence/workspace.py

@@ -40,6 +40,11 @@ class MembershipStatus(enum.StrEnum):
     REMOVED = 'removed'
 
 
+class MembershipSource(enum.StrEnum):
+    LOCAL = 'local'
+    CLOUD_PROJECTION = 'cloud_projection'
+
+
 class InvitationStatus(enum.StrEnum):
     PENDING = 'pending'
     ACCEPTED = 'accepted'
@@ -151,6 +156,11 @@ class WorkspaceMembership(Base):
         nullable=True,
     )
     joined_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True)
+    source = sqlalchemy.Column(
+        sqlalchemy.String(32),
+        nullable=False,
+        server_default=MembershipSource.LOCAL.value,
+    )
     projection_revision = sqlalchemy.Column(sqlalchemy.BigInteger, nullable=False, server_default='0')
     created_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False, server_default=sqlalchemy.func.now())
     updated_at = sqlalchemy.Column(
@@ -178,6 +188,10 @@ class WorkspaceMembership(Base):
             "status IN ('active', 'disabled', 'removed')",
             name='ck_workspace_memberships_status',
         ),
+        sqlalchemy.CheckConstraint(
+            "source IN ('local', 'cloud_projection')",
+            name='ck_workspace_memberships_source',
+        ),
     )
 
 

src/langbot/pkg/persistence/alembic/versions/0020_workspace_membership_source.py

@@ -0,0 +1,49 @@
+"""add explicit Workspace membership source
+
+Revision ID: 0020_membership_source
+Revises: 001a_pgvector_dimension_3072
+Create Date: 2026-08-06
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+revision = '0020_membership_source'
+down_revision = '001a_pgvector_dimension_3072'
+branch_labels = None
+depends_on = None
+
+_CONSTRAINT_NAME = 'ck_workspace_memberships_source'
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    inspector = sa.inspect(conn)
+    if 'workspace_memberships' not in inspector.get_table_names():
+        return
+    if 'source' in {column['name'] for column in inspector.get_columns('workspace_memberships')}:
+        return
+
+    # No durable historical field distinguishes Directory-created revision-zero
+    # rows from Core invitations. Protect every existing row; production can
+    # reclassify separately after UUIDs have been verified against Space.
+    with op.batch_alter_table('workspace_memberships') as batch_op:
+        batch_op.add_column(sa.Column('source', sa.String(length=32), nullable=False, server_default='local'))
+        batch_op.create_check_constraint(
+            _CONSTRAINT_NAME,
+            "source IN ('local', 'cloud_projection')",
+        )
+
+
+def downgrade() -> None:
+    conn = op.get_bind()
+    inspector = sa.inspect(conn)
+    if 'workspace_memberships' not in inspector.get_table_names():
+        return
+    if 'source' not in {column['name'] for column in inspector.get_columns('workspace_memberships')}:
+        return
+    with op.batch_alter_table('workspace_memberships') as batch_op:
+        batch_op.drop_constraint(_CONSTRAINT_NAME, type_='check')
+        batch_op.drop_column('source')

src/langbot/pkg/workspace/collaboration.py

@@ -17,6 +17,7 @@ from ..entity.persistence.user import AccountStatus, User
 from ..entity.persistence.workspace import (
     InvitationStatus,
     MembershipRole,
+    MembershipSource,
     MembershipStatus,
     Workspace,
     WorkspaceInvitation,
@@ -483,6 +484,7 @@ class WorkspaceCollaborationService:
                     account_uuid=account_uuid,
                     role=invitation.role,
                     status=MembershipStatus.ACTIVE.value,
+                    source=MembershipSource.LOCAL.value,
                     invited_by_account_uuid=invitation.created_by_account_uuid,
                     joined_at=now,
                     projection_revision=0,
@@ -491,6 +493,7 @@ class WorkspaceCollaborationService:
             elif membership.status != MembershipStatus.ACTIVE.value:
                 membership.role = invitation.role
                 membership.status = MembershipStatus.ACTIVE.value
+                membership.source = MembershipSource.LOCAL.value
                 membership.invited_by_account_uuid = invitation.created_by_account_uuid
                 membership.joined_at = now
 

src/langbot/pkg/workspace/service.py

@@ -11,6 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
 
 from ..entity.persistence.workspace import (
     MembershipRole,
+    MembershipSource,
     MembershipStatus,
     Workspace,
     WorkspaceExecutionSource,
@@ -451,6 +452,7 @@ class WorkspaceService:
                 account_uuid=account_uuid,
                 role=MembershipRole.OWNER.value,
                 status=MembershipStatus.ACTIVE.value,
+                source=MembershipSource.LOCAL.value,
                 joined_at=joined_at,
                 projection_revision=0,
             )
@@ -458,6 +460,7 @@ class WorkspaceService:
         else:
             membership.role = MembershipRole.OWNER.value
             membership.status = MembershipStatus.ACTIVE.value
+            membership.source = MembershipSource.LOCAL.value
             membership.joined_at = membership.joined_at or joined_at
 
         if workspace.created_by_account_uuid is None:

tests/integration/persistence/test_membership_source_migration.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import pytest
+import sqlalchemy as sa
+from sqlalchemy.ext.asyncio import create_async_engine
+
+from langbot.pkg.persistence.alembic_runner import run_alembic_stamp, run_alembic_upgrade
+
+
+@pytest.mark.asyncio
+async def test_membership_source_migration_backfills_existing_rows_as_local_and_enforces_constraint(tmp_path):
+    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "membership-source.db"}')
+    try:
+        async with engine.begin() as connection:
+            await connection.execute(
+                sa.text(
+                    """
+                    CREATE TABLE workspace_memberships (
+                        uuid VARCHAR(36) PRIMARY KEY,
+                        workspace_uuid VARCHAR(36) NOT NULL,
+                        account_uuid VARCHAR(36) NOT NULL,
+                        role VARCHAR(32) NOT NULL,
+                        status VARCHAR(32) NOT NULL,
+                        projection_revision BIGINT NOT NULL DEFAULT 0,
+                        created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+                        updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
+                    )
+                    """
+                )
+            )
+            await connection.execute(
+                sa.text(
+                    """
+                    INSERT INTO workspace_memberships
+                        (uuid, workspace_uuid, account_uuid, role, status, projection_revision)
+                    VALUES
+                        ('00000000-0000-4000-8000-000000000001', 'workspace', 'local-account',
+                         'viewer', 'active', 0),
+                        ('00000000-0000-4000-8000-000000000002', 'workspace', 'cloud-account',
+                         'viewer', 'active', 0)
+                    """
+                )
+            )
+
+        await run_alembic_stamp(engine, '0019_single_workspace_owner')
+        await run_alembic_upgrade(engine, 'head')
+
+        async with engine.connect() as connection:
+            rows = (
+                await connection.execute(sa.text('SELECT uuid, source FROM workspace_memberships ORDER BY uuid'))
+            ).all()
+            columns = await connection.run_sync(
+                lambda sync_connection: {
+                    column['name']: column
+                    for column in sa.inspect(sync_connection).get_columns('workspace_memberships')
+                }
+            )
+        assert rows == [
+            ('00000000-0000-4000-8000-000000000001', 'local'),
+            ('00000000-0000-4000-8000-000000000002', 'local'),
+        ]
+        assert columns['source']['nullable'] is False
+
+        with pytest.raises(sa.exc.IntegrityError):
+            async with engine.begin() as connection:
+                await connection.execute(
+                    sa.text("UPDATE workspace_memberships SET source = 'guessed-from-user-source'")
+                )
+    finally:
+        await engine.dispose()

tests/integration/persistence/test_migrations.py

@@ -105,7 +105,7 @@ class TestSQLiteMigrationUpgrade:
         await run_alembic_upgrade(sqlite_engine, 'head')
 
         assert await get_alembic_current(sqlite_engine) == _get_script_head()
-        assert _get_script_head() == '001a_pgvector_dimension_3072'
+        assert _get_script_head() == '0020_membership_source'
 
     @pytest.mark.asyncio
     async def test_upgrade_from_baseline_to_head(self, sqlite_engine):

tests/unit_tests/cloud/test_directory_projection.py

@@ -1056,7 +1056,7 @@ async def test_snapshot_for_another_instance_is_rejected(projection_context):
         await service.initialize()
 
 
-async def test_core_owned_membership_survives_directory_updates_and_omission(projection_context):
+async def test_directory_revision_zero_membership_is_adopted(projection_context):
     application, session_factory = projection_context
     service = DirectoryProjectionService(application, _Provider([_snapshot(1)]), INSTANCE_UUID)
     await service.initialize()
@@ -1067,29 +1067,125 @@ async def test_core_owned_membership_survives_directory_updates_and_omission(pro
             membership.role = 'viewer'
             membership.status = 'active'
             membership.projection_revision = 0
-            session.add(
-                WorkspaceMembership(
-                    uuid=SECOND_MEMBERSHIP_UUID,
-                    workspace_uuid=WORKSPACE_UUID,
-                    account_uuid='20000000-0000-0000-0000-000000000099',
-                    role='viewer',
-                    status='active',
-                    joined_at=membership.joined_at,
-                    projection_revision=0,
-                )
-            )
 
     projected_member = _member(revision=2).model_copy(update={'role': 'owner', 'membership_status': 'removed'})
     projected_workspace = _workspace(revision=2).model_copy(update={'members': (projected_member,)})
     await service.apply_snapshot(_snapshot(2, workspaces=[projected_workspace]))
 
     async with session_factory() as session:
-        memberships = {
-            membership.uuid: membership
-            for membership in (await session.scalars(sqlalchemy.select(WorkspaceMembership))).all()
-        }
-    assert memberships[MEMBERSHIP_UUID].role == 'viewer'
-    assert memberships[MEMBERSHIP_UUID].status == 'active'
-    assert memberships[MEMBERSHIP_UUID].projection_revision == 0
-    assert memberships[SECOND_MEMBERSHIP_UUID].status == 'active'
-    assert memberships[SECOND_MEMBERSHIP_UUID].projection_revision == 0
+        membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
+    assert membership.source == 'cloud_projection'
+    assert membership.role == 'owner'
+    assert membership.status == 'removed'
+    assert membership.projection_revision == 2
+
+
+async def test_directory_revision_zero_membership_omitted_from_snapshot_is_removed(projection_context):
+    application, session_factory = projection_context
+    service = DirectoryProjectionService(application, _Provider([_snapshot(1)]), INSTANCE_UUID)
+    await service.initialize()
+
+    historical_account_uuid = '20000000-0000-0000-0000-000000000099'
+    async with session_factory() as session:
+        async with session.begin():
+            membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
+            session.add(
+                User(
+                    uuid=historical_account_uuid,
+                    user='Historical Space Member',
+                    normalized_email='historical@example.com',
+                    password='',
+                    status='active',
+                    source='cloud_projection',
+                    projection_revision=1,
+                    account_type='space',
+                    space_account_uuid=historical_account_uuid,
+                )
+            )
+            session.add(
+                WorkspaceMembership(
+                    uuid=SECOND_MEMBERSHIP_UUID,
+                    workspace_uuid=WORKSPACE_UUID,
+                    account_uuid=historical_account_uuid,
+                    role='viewer',
+                    status='active',
+                    source='cloud_projection',
+                    joined_at=membership.joined_at,
+                    projection_revision=0,
+                )
+            )
+
+    await service.apply_snapshot(_snapshot(2))
+
+    async with session_factory() as session:
+        historical = await session.get(WorkspaceMembership, SECOND_MEMBERSHIP_UUID)
+    assert historical.status == 'removed'
+    assert historical.projection_revision == 2
+
+
+async def test_cloud_account_core_invitation_membership_survives_directory_omission(projection_context):
+    application, session_factory = projection_context
+    service = DirectoryProjectionService(application, _Provider([_snapshot(1)]), INSTANCE_UUID)
+    await service.initialize()
+
+    invited_account_uuid = '20000000-0000-0000-0000-000000000098'
+    async with session_factory() as session:
+        async with session.begin():
+            projected_membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
+            session.add(
+                User(
+                    uuid=invited_account_uuid,
+                    user='Invited Cloud Account',
+                    normalized_email='invited-cloud@example.com',
+                    password='',
+                    status='active',
+                    source='cloud_projection',
+                    projection_revision=1,
+                    account_type='space',
+                    space_account_uuid=invited_account_uuid,
+                )
+            )
+            session.add(
+                WorkspaceMembership(
+                    uuid=SECOND_MEMBERSHIP_UUID,
+                    workspace_uuid=WORKSPACE_UUID,
+                    account_uuid=invited_account_uuid,
+                    role='viewer',
+                    status='active',
+                    source='local',
+                    joined_at=projected_membership.joined_at,
+                    projection_revision=0,
+                )
+            )
+
+    await service.apply_snapshot(_snapshot(2))
+
+    async with session_factory() as session:
+        membership = await session.get(WorkspaceMembership, SECOND_MEMBERSHIP_UUID)
+    assert membership.source == 'local'
+    assert membership.status == 'active'
+    assert membership.projection_revision == 0
+
+
+async def test_directory_does_not_adopt_local_membership_with_different_uuid_for_same_cloud_account(projection_context):
+    application, session_factory = projection_context
+    service = DirectoryProjectionService(application, _Provider([_snapshot(1)]), INSTANCE_UUID)
+    await service.initialize()
+
+    async with session_factory() as session:
+        async with session.begin():
+            membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
+            membership.uuid = SECOND_MEMBERSHIP_UUID
+            membership.source = 'local'
+            membership.projection_revision = 0
+
+    projected_member = _member(revision=2).model_copy(update={'role': 'owner', 'membership_status': 'removed'})
+    projected_workspace = _workspace(revision=2).model_copy(update={'members': (projected_member,)})
+    await service.apply_snapshot(_snapshot(2, workspaces=[projected_workspace]))
+
+    async with session_factory() as session:
+        membership = await session.get(WorkspaceMembership, SECOND_MEMBERSHIP_UUID)
+    assert membership.source == 'local'
+    assert membership.role == 'developer'
+    assert membership.status == 'active'
+    assert membership.projection_revision == 0

tests/unit_tests/workspace/test_workspace_collaboration.py

@@ -99,6 +99,7 @@ async def test_invitation_secret_is_hashed_and_acceptance_is_one_time(collaborat
     membership = await service.accept_invitation(created.token, account.uuid)
     assert membership.workspace_uuid == workspace.uuid
     assert membership.role == 'developer'
+    assert membership.source == 'local'
 
     with pytest.raises(InvitationUsedError):
         await service.accept_invitation(created.token, account.uuid)

tests/unit_tests/workspace/test_workspace_service.py

@@ -153,6 +153,33 @@ async def test_initial_owner_cannot_be_claimed_by_another_account(workspace_test
         ).all()
         assert len(owners) == 1
         assert owners[0].account_uuid == first_account_uuid
+        assert owners[0].source == 'local'
+
+
+async def test_claim_initial_owner_reclassifies_existing_membership_as_local(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, 'reclaimed@example.com')
+            workspace = await service.ensure_singleton_workspace(session=session)
+            session.add(
+                WorkspaceMembership(
+                    uuid='44444444-4444-4444-8444-444444444444',
+                    workspace_uuid=workspace.uuid,
+                    account_uuid=account_uuid,
+                    role='viewer',
+                    status='removed',
+                    source='cloud_projection',
+                    projection_revision=4,
+                )
+            )
+
+    membership = await service.claim_initial_owner(account_uuid)
+
+    assert membership.role == 'owner'
+    assert membership.status == 'active'
+    assert membership.source == 'local'
 
 
 async def test_execution_binding_returns_persisted_generation(workspace_test_context):