From 4bd899e77b5d925bb65fdc78303784486825c8db Mon Sep 17 00:00:00 2001 From: Hyu Date: Fri, 7 Aug 2026 11:26:06 +0800 Subject: [PATCH] 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 --- src/langbot/pkg/cloud/directory_projection.py | 12 +- .../pkg/entity/persistence/workspace.py | 14 ++ .../0020_workspace_membership_source.py | 49 +++++++ src/langbot/pkg/workspace/collaboration.py | 3 + src/langbot/pkg/workspace/service.py | 3 + .../test_membership_source_migration.py | 70 +++++++++ .../persistence/test_migrations.py | 2 +- .../cloud/test_directory_projection.py | 138 +++++++++++++++--- .../workspace/test_workspace_collaboration.py | 1 + .../workspace/test_workspace_service.py | 27 ++++ 10 files changed, 292 insertions(+), 27 deletions(-) create mode 100644 src/langbot/pkg/persistence/alembic/versions/0020_workspace_membership_source.py create mode 100644 tests/integration/persistence/test_membership_source_migration.py diff --git a/src/langbot/pkg/cloud/directory_projection.py b/src/langbot/pkg/cloud/directory_projection.py index 46d06f153..47fad19fd 100644 --- a/src/langbot/pkg/cloud/directory_projection.py +++ b/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), diff --git a/src/langbot/pkg/entity/persistence/workspace.py b/src/langbot/pkg/entity/persistence/workspace.py index 822c8e9d4..d088ae7d6 100644 --- a/src/langbot/pkg/entity/persistence/workspace.py +++ b/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', + ), ) diff --git a/src/langbot/pkg/persistence/alembic/versions/0020_workspace_membership_source.py b/src/langbot/pkg/persistence/alembic/versions/0020_workspace_membership_source.py new file mode 100644 index 000000000..03db6a72a --- /dev/null +++ b/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') diff --git a/src/langbot/pkg/workspace/collaboration.py b/src/langbot/pkg/workspace/collaboration.py index 5d88935bf..c66ac3866 100644 --- a/src/langbot/pkg/workspace/collaboration.py +++ b/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 diff --git a/src/langbot/pkg/workspace/service.py b/src/langbot/pkg/workspace/service.py index 96439b2af..c20fdc228 100644 --- a/src/langbot/pkg/workspace/service.py +++ b/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: diff --git a/tests/integration/persistence/test_membership_source_migration.py b/tests/integration/persistence/test_membership_source_migration.py new file mode 100644 index 000000000..f72c7a1b4 --- /dev/null +++ b/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() diff --git a/tests/integration/persistence/test_migrations.py b/tests/integration/persistence/test_migrations.py index 6e2b723ff..e7106f939 100644 --- a/tests/integration/persistence/test_migrations.py +++ b/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): diff --git a/tests/unit_tests/cloud/test_directory_projection.py b/tests/unit_tests/cloud/test_directory_projection.py index 0256ae595..69619a109 100644 --- a/tests/unit_tests/cloud/test_directory_projection.py +++ b/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 diff --git a/tests/unit_tests/workspace/test_workspace_collaboration.py b/tests/unit_tests/workspace/test_workspace_collaboration.py index 77974f4ff..487095639 100644 --- a/tests/unit_tests/workspace/test_workspace_collaboration.py +++ b/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) diff --git a/tests/unit_tests/workspace/test_workspace_service.py b/tests/unit_tests/workspace/test_workspace_service.py index f24f2ca44..5d6633d32 100644 --- a/tests/unit_tests/workspace/test_workspace_service.py +++ b/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):