fix(cloud): show owner model balance and enforce single owner (#2384)

* fix(cloud): show owner model balance and enforce single owner

* fix(migrations): create owner index idempotently

---------

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

tests/integration/api/test_support_admin_launch.py

@@ -333,7 +333,6 @@ async def test_support_admin_request_context_has_actor_owner_and_no_membership(s
     assert Permission.RESOURCE_MANAGE.value in permissions
     assert not permissions.intersection(
         {
-            Permission.OWNER_TRANSFER.value,
             Permission.MEMBER_VIEW.value,
             Permission.MEMBER_INVITE.value,
             Permission.MEMBER_UPDATE_ROLE.value,

tests/integration/api/test_user_space_oauth.py

@@ -289,6 +289,9 @@ async def test_cloud_workspace_owner_is_always_space_bound_after_login(space_oau
     application.deployment.mode = 'cloud'
     application.user_service.get_workspace_owner = AsyncMock(return_value=None)
     application.space_service.get_credits = AsyncMock()
+    application.cloud_model_catalog_service = SimpleNamespace(
+        get_workspace_credits=lambda workspace_uuid: 25000 if workspace_uuid == WORKSPACE_UUID else None
+    )
 
     response = await client.get(
         '/api/v1/user/space-credits',
@@ -298,7 +301,7 @@ async def test_cloud_workspace_owner_is_always_space_bound_after_login(space_oau
 
     assert response.status_code == 200
     assert payload['data'] == {
-        'credits': None,
+        'credits': 25000,
         'owner_space_bound': True,
         'is_workspace_owner': True,
     }

tests/integration/api/test_workspaces.py

@@ -188,6 +188,7 @@ async def test_owner_invites_second_account_and_secret_is_not_persisted(workspac
     workspace_uuid = current['workspace']['uuid']
     assert current['membership']['role'] == 'owner'
     assert 'member.invite' in current['permissions']
+    assert 'owner.transfer' not in current['permissions']
 
     invite_response = await client.post(
         f'/api/v1/workspaces/{workspace_uuid}/invitations',
@@ -263,6 +264,14 @@ async def test_owner_invites_second_account_and_secret_is_not_persisted(workspac
     assert member_current['membership']['role'] == 'viewer'
     assert 'member.invite' not in member_current['permissions']
 
+    transfer_response = await client.patch(
+        f'/api/v1/workspaces/{workspace_uuid}/members/{member_current["membership"]["account_uuid"]}',
+        headers=_auth(owner_token, workspace_uuid),
+        json={'role': 'owner'},
+    )
+    assert transfer_response.status_code == 403
+    assert (await transfer_response.get_json())['code'] == 'permission_denied'
+
     forbidden_invite = await client.post(
         f'/api/v1/workspaces/{workspace_uuid}/invitations',
         headers=_auth(member_token, workspace_uuid),

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() == '0018_merge_launch_replay'
+        assert _get_script_head() == '0019_single_workspace_owner'
 
     @pytest.mark.asyncio
     async def test_upgrade_from_baseline_to_head(self, sqlite_engine):

tests/integration/persistence/test_single_owner_migration.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import pytest
+import sqlalchemy as sa
+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, WorkspaceMembership
+from langbot.pkg.persistence.alembic_runner import run_alembic_stamp, run_alembic_upgrade
+
+
+@pytest.mark.asyncio
+async def test_single_owner_migration_demotes_historical_extra_owner_and_installs_unique_index(tmp_path):
+    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "single-owner.db"}')
+    try:
+        async with engine.begin() as connection:
+            await connection.run_sync(Base.metadata.create_all)
+            await connection.execute(sa.text('DROP INDEX uq_workspace_memberships_one_active_owner'))
+
+        session_factory = async_sessionmaker(engine, expire_on_commit=False)
+        workspace_uuid = '00000000-0000-4000-8000-000000000001'
+        creator_uuid = '00000000-0000-4000-8000-000000000010'
+        promoted_uuid = '00000000-0000-4000-8000-000000000020'
+        async with session_factory() as session:
+            session.add_all(
+                [
+                    User(
+                        uuid=creator_uuid,
+                        user='creator@example.test',
+                        normalized_email='creator@example.test',
+                        password='hash',
+                        account_type='local',
+                    ),
+                    User(
+                        uuid=promoted_uuid,
+                        user='promoted@example.test',
+                        normalized_email='promoted@example.test',
+                        password='hash',
+                        account_type='local',
+                    ),
+                    Workspace(
+                        uuid=workspace_uuid,
+                        instance_uuid='instance-test',
+                        name='Workspace',
+                        slug='workspace',
+                        type='team',
+                        status='active',
+                        source='local',
+                        created_by_account_uuid=creator_uuid,
+                    ),
+                    WorkspaceMembership(
+                        uuid='00000000-0000-4000-8000-000000000100',
+                        workspace_uuid=workspace_uuid,
+                        account_uuid=creator_uuid,
+                        role='owner',
+                        status='active',
+                    ),
+                    WorkspaceMembership(
+                        uuid='00000000-0000-4000-8000-000000000200',
+                        workspace_uuid=workspace_uuid,
+                        account_uuid=promoted_uuid,
+                        role='owner',
+                        status='active',
+                    ),
+                ]
+            )
+            await session.commit()
+
+        await run_alembic_stamp(engine, '0018_merge_launch_replay')
+        await run_alembic_upgrade(engine, 'head')
+
+        async with engine.connect() as connection:
+            roles = dict(
+                (
+                    await connection.execute(
+                        sa.text(
+                            'SELECT account_uuid, role FROM workspace_memberships '
+                            'WHERE workspace_uuid = :workspace_uuid ORDER BY account_uuid'
+                        ),
+                        {'workspace_uuid': workspace_uuid},
+                    )
+                ).all()
+            )
+            assert roles == {creator_uuid: 'owner', promoted_uuid: 'admin'}
+            indexes = await connection.run_sync(
+                lambda sync_connection: {
+                    index['name'] for index in sa.inspect(sync_connection).get_indexes('workspace_memberships')
+                }
+            )
+            assert 'uq_workspace_memberships_one_active_owner' in indexes
+
+        with pytest.raises(sa.exc.IntegrityError):
+            async with engine.begin() as connection:
+                await connection.execute(
+                    sa.text("UPDATE workspace_memberships SET role = 'owner' WHERE account_uuid = :account_uuid"),
+                    {'account_uuid': promoted_uuid},
+                )
+    finally:
+        await engine.dispose()