feat(cloud): enforce workspace resource quotas

src/langbot/pkg/api/http/controller/group.py

@@ -14,6 +14,7 @@ from ....utils import bounded_executor
 from ....workspace.collaboration import MembershipPermissionError, WorkspaceCollaborationError
 from ....workspace.errors import WorkspaceNotFoundError
 from ....cloud.entitlements import EntitlementUnavailableError
+from ....cloud.quotas import WorkspaceQuotaExceededError
 from ....core.errors import TaskCapacityError
 from ..authz import AuthorizationError, Permission, permissions_for_role, require_permission
 from ..context import PrincipalContext, PrincipalType, RequestContext, WorkspaceContext
@@ -219,6 +220,8 @@ class RouterGroup(abc.ABC):
                         return self.http_status(403, e.code, str(e))
                     if isinstance(e, WorkspaceCollaborationError):
                         return self.http_status(400, e.code, str(e))
+                    if isinstance(e, WorkspaceQuotaExceededError):
+                        return self.http_status(409, e.error_code, str(e))
                     if isinstance(e, TaskCapacityError):
                         return self.http_status(429, 'task_capacity_exceeded', str(e))
                     if isinstance(

src/langbot/pkg/api/http/service/bot.py

@@ -4,6 +4,7 @@ import uuid
 import sqlalchemy
 
 from ....core import app
+from ....cloud.quotas import require_resource_capacity, resolve_workspace_quota
 from ....entity.persistence import bot as persistence_bot
 from ....entity.persistence import pipeline as persistence_pipeline
 from ....workspace.errors import WorkspaceNotFoundError
@@ -101,20 +102,21 @@ class BotService:
     async def create_bot(self, context: TenantContext, bot_data: dict) -> str:
         """Create bot"""
         workspace_uuid = require_workspace_uuid(context)
-        # Check limitation
         limitation = self.ap.instance_config.data.get('system', {}).get('limitation', {})
-        max_bots = limitation.get('max_bots', -1)
+        quota = await resolve_workspace_quota(
-        if max_bots >= 0:
+            self.ap,
-            existing_bots = await self.get_bots(context)
+            workspace_uuid,
-            if len(existing_bots) >= max_bots:
+            'bots.max',
-                raise ValueError(f'Maximum number of bots ({max_bots}) reached')
+            fallback=limitation.get('max_bots', -1),
+        )
 
         # TODO: 检查配置信息格式
         bot_data = bot_data.copy()
         bot_data['uuid'] = str(uuid.uuid4())
         bot_data['workspace_uuid'] = workspace_uuid
 
-        # bind the most recently updated pipeline if any exist
+        # Preserve the legacy flat-row result shape for this optional lookup;
+        # quota admission and insertion below still share one transaction.
         result = await self.ap.persistence_mgr.execute_async(
             scope_statement(
                 sqlalchemy.select(persistence_pipeline.LegacyPipeline),
@@ -129,7 +131,25 @@ class BotService:
             bot_data['use_pipeline_uuid'] = pipeline.uuid
             bot_data['use_pipeline_name'] = pipeline.name
 
-        await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_bot.Bot).values(bot_data))
+        async def persist(execute) -> None:
+            await require_resource_capacity(
+                execute,
+                workspace_uuid=workspace_uuid,
+                model=persistence_bot.Bot,
+                quota=quota,
+                resource_name='bots',
+            )
+
+            await execute(sqlalchemy.insert(persistence_bot.Bot).values(bot_data))
+
+        tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
+        if quota.requires_transaction_lock:
+            if not callable(tenant_uow):
+                raise RuntimeError('Cloud bot quota enforcement requires transactional persistence')
+            async with tenant_uow(workspace_uuid) as uow:
+                await persist(uow.execute)
+        else:
+            await persist(self.ap.persistence_mgr.execute_async)
 
         bot = await self.get_bot(context, bot_data['uuid'], include_secret=True)
 

src/langbot/pkg/cloud/quotas.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Awaitable, Callable
+
+import sqlalchemy
+
+from ..entity.persistence import workspace as persistence_workspace
+from .entitlements import EntitlementResolver
+
+
+Execute = Callable[[Any], Awaitable[Any]]
+
+
+class WorkspaceQuotaExceededError(ValueError):
+    """A stable business error raised when a workspace has no free slots."""
+
+    error_code = 'workspace_quota_exceeded'
+
+    def __init__(self, resource_name: str, limit: int) -> None:
+        self.resource_name = resource_name
+        self.limit = limit
+        super().__init__(f'Maximum number of {resource_name} ({limit}) reached')
+
+
+@dataclass(frozen=True, slots=True)
+class WorkspaceQuota:
+    limit: int
+    requires_transaction_lock: bool
+
+
+async def resolve_workspace_quota(
+    ap: Any,
+    workspace_uuid: str,
+    limit_name: str,
+    *,
+    fallback: int = -1,
+) -> WorkspaceQuota:
+    """Resolve a plan-agnostic Cloud limit while preserving OSS configuration."""
+
+    resolver = getattr(ap, 'entitlement_resolver', None)
+    if isinstance(resolver, EntitlementResolver):
+        snapshot = await resolver.resolve(workspace_uuid)
+        return WorkspaceQuota(
+            limit=snapshot.limit(limit_name),
+            requires_transaction_lock=True,
+        )
+    return WorkspaceQuota(limit=fallback, requires_transaction_lock=False)
+
+
+async def lock_workspace_for_quota(execute: Execute, workspace_uuid: str) -> None:
+    """Serialize quota checks on the durable Workspace row within one transaction."""
+
+    result = await execute(
+        sqlalchemy.select(persistence_workspace.Workspace.uuid)
+        .where(persistence_workspace.Workspace.uuid == workspace_uuid)
+        .with_for_update()
+    )
+    if result.first() is None:
+        raise ValueError('Workspace does not exist')
+
+
+async def require_resource_capacity(
+    execute: Execute,
+    *,
+    workspace_uuid: str,
+    model: type,
+    quota: WorkspaceQuota,
+    resource_name: str,
+    workspace_locked: bool = False,
+) -> None:
+    if quota.limit < 0:
+        return
+    if quota.requires_transaction_lock and not workspace_locked:
+        await lock_workspace_for_quota(execute, workspace_uuid)
+    result = await execute(
+        sqlalchemy.select(sqlalchemy.func.count())
+        .select_from(model)
+        .where(model.workspace_uuid == workspace_uuid)
+    )
+    if int(result.scalar_one()) >= quota.limit:
+        raise WorkspaceQuotaExceededError(resource_name, quota.limit)

src/langbot/pkg/plugin/connector.py

@@ -19,6 +19,11 @@ from urllib.parse import urljoin, urlparse
 from langbot_plugin.api.entities.builtin.pipeline.query import provider_session
 
 from ..core import app
+from ..cloud.quotas import (
+    lock_workspace_for_quota,
+    require_resource_capacity,
+    resolve_workspace_quota,
+)
 from . import handler
 from .archive import inspect_plugin_archive_metadata
 from .github import (
@@ -1295,6 +1300,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
         install_info: dict[str, Any],
         artifact_digest: str,
     ) -> tuple[InstallationBinding, str | None, bool]:
+        quota = await resolve_workspace_quota(
+            self.ap,
+            execution_context.workspace_uuid,
+            'plugins.max',
+        )
         safe_install_info = {
             key: value
             for key, value in install_info.items()
@@ -1316,9 +1326,19 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
         )
 
         async def persist(execute):
+            if quota.requires_transaction_lock:
+                await lock_workspace_for_quota(execute, execution_context.workspace_uuid)
             result = await execute(statement)
             setting = result.first()
             if setting is None:
+                await require_resource_capacity(
+                    execute,
+                    workspace_uuid=execution_context.workspace_uuid,
+                    model=persistence_plugin.PluginSetting,
+                    quota=quota,
+                    resource_name='plugins',
+                    workspace_locked=quota.requires_transaction_lock,
+                )
                 installation_uuid = str(uuid.uuid4())
                 runtime_revision = 1
                 previous_digest = None
@@ -1373,6 +1393,8 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
             )
 
         tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
+        if quota.requires_transaction_lock and not callable(tenant_uow):
+            raise RuntimeError('Cloud plugin quota enforcement requires transactional persistence')
         if callable(tenant_uow):
             async with tenant_uow(execution_context.workspace_uuid) as uow:
                 return await persist(uow.execute)

tests/integration/persistence/test_workspace_quota_postgres.py

@@ -0,0 +1,137 @@
+"""PostgreSQL integration coverage for durable workspace quota locking.
+
+Run with TEST_POSTGRES_URL=postgresql+asyncpg://... pytest ...
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import uuid
+
+import pytest
+import sqlalchemy as sa
+from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+
+from langbot.pkg.cloud import quotas as quota_module
+from langbot.pkg.cloud.quotas import WorkspaceQuota, WorkspaceQuotaExceededError, require_resource_capacity
+
+
+pytestmark = [pytest.mark.integration, pytest.mark.slow, pytest.mark.asyncio]
+
+
+class _Base(DeclarativeBase):
+    pass
+
+
+class _Workspace(_Base):
+    __tablename__ = 'quota_integration_workspaces'
+
+    uuid: Mapped[str] = mapped_column(sa.String(36), primary_key=True)
+
+
+class _Resource(_Base):
+    __tablename__ = 'quota_integration_resources'
+
+    uuid: Mapped[str] = mapped_column(sa.String(36), primary_key=True)
+    workspace_uuid: Mapped[str] = mapped_column(
+        sa.String(36),
+        sa.ForeignKey('quota_integration_workspaces.uuid', ondelete='CASCADE'),
+        nullable=False,
+        index=True,
+    )
+
+
+@pytest.fixture
+async def quota_postgres(monkeypatch):
+    url = os.environ.get('TEST_POSTGRES_URL')
+    if not url:
+        pytest.skip('TEST_POSTGRES_URL not set')
+    if url.startswith('postgresql://'):
+        url = url.replace('postgresql://', 'postgresql+asyncpg://', 1)
+
+    engine = create_async_engine(url, pool_size=5, max_overflow=0)
+    monkeypatch.setattr(quota_module.persistence_workspace, 'Workspace', _Workspace)
+    async with engine.begin() as connection:
+        await connection.run_sync(_Base.metadata.drop_all)
+        await connection.run_sync(_Base.metadata.create_all)
+    try:
+        yield url, engine
+    finally:
+        async with engine.begin() as connection:
+            await connection.run_sync(_Base.metadata.drop_all)
+        await engine.dispose()
+
+
+async def test_workspace_row_lock_is_atomic_isolated_and_survives_pool_restart(quota_postgres) -> None:
+    url, engine = quota_postgres
+    workspace_a = str(uuid.uuid4())
+    workspace_b = str(uuid.uuid4())
+    quota = WorkspaceQuota(limit=1, requires_transaction_lock=True)
+    sessions = async_sessionmaker(engine, expire_on_commit=False)
+
+    async with engine.begin() as connection:
+        await connection.execute(sa.insert(_Workspace), [{'uuid': workspace_a}, {'uuid': workspace_b}])
+
+    lock_acquired = asyncio.Event()
+    release_first = asyncio.Event()
+
+    async def admit(workspace_uuid: str, *, hold: bool = False) -> None:
+        async with sessions() as session:
+            async with session.begin():
+                await require_resource_capacity(
+                    session.execute,
+                    workspace_uuid=workspace_uuid,
+                    model=_Resource,
+                    quota=quota,
+                    resource_name='resources',
+                )
+                if hold:
+                    lock_acquired.set()
+                    await release_first.wait()
+                await session.execute(
+                    sa.insert(_Resource).values(uuid=str(uuid.uuid4()), workspace_uuid=workspace_uuid)
+                )
+
+    first = asyncio.create_task(admit(workspace_a, hold=True))
+    await asyncio.wait_for(lock_acquired.wait(), timeout=2)
+    same_workspace = asyncio.create_task(admit(workspace_a))
+    other_workspace = asyncio.create_task(admit(workspace_b))
+
+    await asyncio.wait_for(other_workspace, timeout=2)
+    assert not same_workspace.done(), 'same-workspace transaction bypassed SELECT FOR UPDATE'
+
+    release_first.set()
+    await first
+    with pytest.raises(WorkspaceQuotaExceededError, match=r'Maximum number of resources \(1\) reached'):
+        await same_workspace
+
+    async with sessions() as session:
+        counts = dict(
+            (
+                await session.execute(
+                    sa.select(_Resource.workspace_uuid, sa.func.count())
+                    .group_by(_Resource.workspace_uuid)
+                    .order_by(_Resource.workspace_uuid)
+                )
+            ).all()
+        )
+    assert counts == {workspace_a: 1, workspace_b: 1}
+
+    await engine.dispose()
+    restarted_engine = create_async_engine(url, pool_size=2, max_overflow=0)
+    restarted_sessions = async_sessionmaker(restarted_engine, expire_on_commit=False)
+    try:
+        async with restarted_sessions() as session:
+            async with session.begin():
+                with pytest.raises(WorkspaceQuotaExceededError):
+                    await require_resource_capacity(
+                        session.execute,
+                        workspace_uuid=workspace_a,
+                        model=_Resource,
+                        quota=quota,
+                        resource_name='resources',
+                    )
+    finally:
+        await restarted_engine.dispose()

tests/unit_tests/api/http/test_internal_error_responses.py

@@ -9,6 +9,7 @@ import quart
 
 from langbot.pkg.api.http.controller import group
 from langbot.pkg.api.http.controller.groups.webhooks import WebhookRouterGroup
+from langbot.pkg.cloud.quotas import WorkspaceQuotaExceededError
 from langbot.pkg.utils.bounded_executor import (
     BlockingWorkCapacityError,
     current_blocking_work_scope,
@@ -48,6 +49,16 @@ class _BlockingCapacityRouterGroup(group.RouterGroup):
             raise BlockingWorkCapacityError('Workspace blocking executor capacity reached')
 
 
+class _QuotaRouterGroup(group.RouterGroup):
+    name = 'quota-test'
+    path = '/quota-test'
+
+    async def initialize(self) -> None:
+        @self.route('', methods=['POST'], auth_type=group.AuthType.NONE)
+        async def _():
+            raise WorkspaceQuotaExceededError('bots', 2)
+
+
 class _InvalidAccountRouterGroup(group.RouterGroup):
     name = 'invalid-account-test'
     path = '/invalid-account-test'
@@ -130,6 +141,20 @@ async def test_blocking_work_capacity_maps_to_retryable_http_response():
     }
 
 
+async def test_workspace_quota_maps_to_stable_conflict_response():
+    application = SimpleNamespace(logger=Mock())
+    quart_app = quart.Quart(__name__)
+    await _QuotaRouterGroup(application, quart_app).initialize()
+
+    response = await quart_app.test_client().post('/quota-test')
+
+    assert response.status_code == 409
+    assert await response.get_json() == {
+        'code': 'workspace_quota_exceeded',
+        'msg': 'Maximum number of bots (2) reached',
+    }
+
+
 async def test_public_webhook_carries_scope_without_holding_database_session():
     class ScopeOnlyPersistenceManager:
         mode = SimpleNamespace(value='cloud_runtime')

tests/unit_tests/api/service/test_bot_service.py

@@ -311,10 +311,9 @@ class TestBotServiceCreateBot:
         ap.platform_mgr = SimpleNamespace()
         ap.platform_mgr.load_bot = AsyncMock()
 
-        # Mock get_bots to return 2 bots already
+        # Mock the atomic count query to report 2 existing bots.
-        bot1 = _create_mock_bot(bot_uuid='uuid-1')
+        mock_result = _create_mock_result()
-        bot2 = _create_mock_bot(bot_uuid='uuid-2')
+        mock_result.scalar_one = Mock(return_value=2)
-        mock_result = _create_mock_result([bot1, bot2])
         ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
         ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'uuid-1', 'name': 'Bot 1'})
 

tests/unit_tests/api/service/test_cloud_bot_quota_concurrency.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import asyncio
+from collections import defaultdict
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+import sqlalchemy
+
+from langbot.pkg.api.http.service.bot import BotService
+from langbot.pkg.cloud.entitlements import EntitlementResolver, EntitlementSnapshot
+
+
+INSTANCE_UUID = 'cloud-instance'
+WORKSPACE_A = '11111111-1111-1111-1111-111111111111'
+WORKSPACE_B = '22222222-2222-2222-2222-222222222222'
+
+
+class _Provider:
+    async def get_workspace_entitlement(self, workspace_uuid: str) -> EntitlementSnapshot:
+        return EntitlementSnapshot(
+            instance_uuid=INSTANCE_UUID,
+            workspace_uuid=workspace_uuid,
+            entitlement_revision=1,
+            status='active',
+            not_before=0,
+            expires_at=4_102_444_800,
+            features={},
+            limits={'bots.max': 2},
+        )
+
+
+class _Result:
+    def __init__(self, *, first=None, scalar=None) -> None:
+        self._first = first
+        self._scalar = scalar
+
+    def first(self):
+        return self._first
+
+    def scalar_one(self):
+        return self._scalar
+
+
+class _TenantUow:
+    def __init__(self, manager: '_Persistence', workspace_uuid: str) -> None:
+        self.manager = manager
+        self.workspace_uuid = workspace_uuid
+        self.lock = manager.locks[workspace_uuid]
+
+    async def __aenter__(self):
+        await self.lock.acquire()
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        self.lock.release()
+
+    async def execute(self, statement):
+        sql = str(statement)
+        if isinstance(statement, sqlalchemy.sql.dml.Insert):
+            assert statement.table.name == 'bots'
+            self.manager.bots[self.workspace_uuid].append(statement.compile().params)
+            return _Result()
+        if 'FROM workspaces' in sql:
+            assert statement._for_update_arg is not None
+            self.manager.workspace_locks_seen += 1
+            return _Result(first=(self.workspace_uuid,))
+        if 'count(' in sql.lower() and 'FROM bots' in sql:
+            return _Result(scalar=len(self.manager.bots[self.workspace_uuid]))
+        if 'FROM legacy_pipelines' in sql:
+            return _Result(first=None)
+        raise AssertionError(f'unexpected statement: {sql}')
+
+
+class _Persistence:
+    def __init__(self) -> None:
+        self.locks = defaultdict(asyncio.Lock)
+        self.bots = defaultdict(list)
+        self.workspace_locks_seen = 0
+
+    def tenant_uow(self, workspace_uuid: str) -> _TenantUow:
+        return _TenantUow(self, workspace_uuid)
+
+    async def execute_async(self, statement):
+        assert 'FROM legacy_pipelines' in str(statement)
+        return _Result(first=None)
+
+
+async def _service(manager: _Persistence) -> BotService:
+    resolver = EntitlementResolver(INSTANCE_UUID, _Provider())
+    await resolver.reconcile_active_workspaces({WORKSPACE_A, WORKSPACE_B})
+    ap = SimpleNamespace(
+        entitlement_resolver=resolver,
+        persistence_mgr=manager,
+        instance_config=SimpleNamespace(data={'system': {'limitation': {'max_bots': 99}}}),
+        platform_mgr=SimpleNamespace(load_bot=AsyncMock()),
+    )
+    service = BotService(ap)
+    service.get_bot = AsyncMock(return_value={'uuid': 'created'})
+    return service
+
+
+@pytest.mark.asyncio
+async def test_cloud_bot_quota_is_atomic_isolated_and_persists_across_service_restart() -> None:
+    manager = _Persistence()
+    service = await _service(manager)
+
+    async def create(workspace_uuid: str, index: int):
+        return await service.create_bot(workspace_uuid, {'name': f'bot-{index}'})
+
+    results = await asyncio.gather(
+        *(create(WORKSPACE_A, index) for index in range(8)),
+        *(create(WORKSPACE_B, index) for index in range(8)),
+        return_exceptions=True,
+    )
+
+    successes = [result for result in results if isinstance(result, str)]
+    failures = [result for result in results if isinstance(result, ValueError)]
+    assert len(successes) == 4
+    assert len(failures) == 12
+    assert len(manager.bots[WORKSPACE_A]) == 2
+    assert len(manager.bots[WORKSPACE_B]) == 2
+    assert manager.workspace_locks_seen == 16
+
+    restarted_service = await _service(manager)
+    with pytest.raises(ValueError, match=r'Maximum number of bots \(2\) reached'):
+        await restarted_service.create_bot(WORKSPACE_A, {'name': 'after-restart'})
+    assert len(manager.bots[WORKSPACE_A]) == 2

tests/unit_tests/cloud/test_workspace_quotas.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+import sqlalchemy
+
+from langbot.pkg.cloud.entitlements import EntitlementResolver, EntitlementSnapshot
+from langbot.pkg.cloud.quotas import (
+    WorkspaceQuota,
+    require_resource_capacity,
+    resolve_workspace_quota,
+)
+from langbot.pkg.entity.persistence.bot import Bot
+
+
+WORKSPACE_UUID = '11111111-1111-1111-1111-111111111111'
+INSTANCE_UUID = 'cloud-instance'
+
+
+class _Provider:
+    def __init__(self, limits: dict[str, int]) -> None:
+        self.limits = limits
+
+    async def get_workspace_entitlement(self, workspace_uuid: str) -> EntitlementSnapshot:
+        return EntitlementSnapshot(
+            instance_uuid=INSTANCE_UUID,
+            workspace_uuid=workspace_uuid,
+            entitlement_revision=1,
+            status='active',
+            not_before=0,
+            expires_at=4_102_444_800,
+            features={},
+            limits=self.limits,
+            plan_name='test',
+        )
+
+
+async def _resolver(limits: dict[str, int]) -> EntitlementResolver:
+    resolver = EntitlementResolver(INSTANCE_UUID, _Provider(limits))
+    await resolver.reconcile_active_workspaces({WORKSPACE_UUID})
+    return resolver
+
+
+@pytest.mark.asyncio
+async def test_resolve_workspace_quota_uses_signed_cloud_limit() -> None:
+    ap = SimpleNamespace(entitlement_resolver=await _resolver({'bots.max': 2}))
+
+    quota = await resolve_workspace_quota(ap, WORKSPACE_UUID, 'bots.max', fallback=99)
+
+    assert quota == WorkspaceQuota(limit=2, requires_transaction_lock=True)
+
+
+@pytest.mark.asyncio
+async def test_resolve_workspace_quota_preserves_oss_fallback() -> None:
+    quota = await resolve_workspace_quota(SimpleNamespace(), WORKSPACE_UUID, 'bots.max', fallback=7)
+
+    assert quota == WorkspaceQuota(limit=7, requires_transaction_lock=False)
+
+
+@pytest.mark.asyncio
+async def test_require_resource_capacity_locks_workspace_before_counting() -> None:
+    statements: list[object] = []
+    lock_result = Mock()
+    lock_result.first.return_value = (WORKSPACE_UUID,)
+    count_result = Mock()
+    count_result.scalar_one.return_value = 1
+    execute = AsyncMock(side_effect=[lock_result, count_result])
+
+    await require_resource_capacity(
+        execute,
+        workspace_uuid=WORKSPACE_UUID,
+        model=Bot,
+        quota=WorkspaceQuota(limit=2, requires_transaction_lock=True),
+        resource_name='bots',
+    )
+
+    statements.extend(call.args[0] for call in execute.await_args_list)
+    assert len(statements) == 2
+    assert isinstance(statements[0], sqlalchemy.sql.Select)
+    assert statements[0]._for_update_arg is not None
+    assert 'workspaces' in str(statements[0])
+    assert 'count' in str(statements[1]).lower()
+
+
+@pytest.mark.asyncio
+async def test_require_resource_capacity_rejects_at_boundary() -> None:
+    lock_result = Mock()
+    lock_result.first.return_value = (WORKSPACE_UUID,)
+    count_result = Mock()
+    count_result.scalar_one.return_value = 2
+    execute = AsyncMock(side_effect=[lock_result, count_result])
+
+    with pytest.raises(ValueError, match=r'Maximum number of bots \(2\) reached'):
+        await require_resource_capacity(
+            execute,
+            workspace_uuid=WORKSPACE_UUID,
+            model=Bot,
+            quota=WorkspaceQuota(limit=2, requires_transaction_lock=True),
+            resource_name='bots',
+        )

tests/unit_tests/plugin/test_cloud_plugin_quota_concurrency.py

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+import asyncio
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+import sqlalchemy
+
+from langbot.pkg.api.http.context import ExecutionContext
+from langbot.pkg.cloud.entitlements import EntitlementResolver, EntitlementSnapshot
+from langbot.pkg.cloud.quotas import WorkspaceQuotaExceededError
+from langbot.pkg.plugin.connector import PluginRuntimeConnector
+from langbot_plugin.runtime.plugin.mgr import PluginInstallSource
+
+
+INSTANCE_UUID = 'cloud-instance'
+WORKSPACE_A = '11111111-1111-1111-1111-111111111111'
+WORKSPACE_B = '22222222-2222-2222-2222-222222222222'
+
+
+class _Provider:
+    async def get_workspace_entitlement(self, workspace_uuid: str) -> EntitlementSnapshot:
+        return EntitlementSnapshot(
+            instance_uuid=INSTANCE_UUID,
+            workspace_uuid=workspace_uuid,
+            entitlement_revision=1,
+            status='active',
+            not_before=0,
+            expires_at=4_102_444_800,
+            features={},
+            limits={'plugins.max': 3},
+        )
+
+
+class _Result:
+    def __init__(self, *, first=None, scalar=None) -> None:
+        self._first = first
+        self._scalar = scalar
+
+    def first(self):
+        return self._first
+
+    def scalar_one(self):
+        return self._scalar
+
+
+class _TenantUow:
+    def __init__(self, manager: '_Persistence', workspace_uuid: str) -> None:
+        self.manager = manager
+        self.workspace_uuid = workspace_uuid
+        self.lock = manager.locks[workspace_uuid]
+
+    async def __aenter__(self):
+        await self.lock.acquire()
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        self.lock.release()
+
+    async def execute(self, statement):
+        sql = str(statement)
+        params = statement.compile().params
+        if isinstance(statement, sqlalchemy.sql.dml.Insert):
+            assert statement.table.name == 'plugin_settings'
+            key = (params['plugin_author'], params['plugin_name'])
+            self.manager.plugins[self.workspace_uuid][key] = dict(params)
+            return _Result()
+        if isinstance(statement, sqlalchemy.sql.dml.Update):
+            return _Result()
+        if 'FROM workspaces' in sql:
+            assert statement._for_update_arg is not None
+            self.manager.workspace_locks_seen += 1
+            return _Result(first=(self.workspace_uuid,))
+        if 'count(' in sql.lower() and 'FROM plugin_settings' in sql:
+            return _Result(scalar=len(self.manager.plugins[self.workspace_uuid]))
+        if 'FROM plugin_settings' in sql:
+            author = next(value for name, value in params.items() if 'plugin_author' in name)
+            name = next(value for param, value in params.items() if 'plugin_name' in param)
+            row = self.manager.plugins[self.workspace_uuid].get((author, name))
+            if row is None:
+                return _Result(first=None)
+            return _Result(
+                first=SimpleNamespace(
+                    installation_uuid=row['installation_uuid'],
+                    runtime_revision=row['runtime_revision'],
+                    artifact_digest=row['artifact_digest'],
+                    install_info=row['install_info'],
+                )
+            )
+        raise AssertionError(f'unexpected statement: {sql}')
+
+
+class _Persistence:
+    def __init__(self) -> None:
+        self.locks = defaultdict(asyncio.Lock)
+        self.plugins = defaultdict(dict)
+        self.workspace_locks_seen = 0
+
+    def tenant_uow(self, workspace_uuid: str) -> _TenantUow:
+        return _TenantUow(self, workspace_uuid)
+
+
+async def _connector(manager: _Persistence) -> PluginRuntimeConnector:
+    resolver = EntitlementResolver(INSTANCE_UUID, _Provider())
+    await resolver.reconcile_active_workspaces({WORKSPACE_A, WORKSPACE_B})
+    connector = object.__new__(PluginRuntimeConnector)
+    connector.ap = SimpleNamespace(entitlement_resolver=resolver, persistence_mgr=manager)
+    return connector
+
+
+def _context(workspace_uuid: str) -> ExecutionContext:
+    return ExecutionContext(
+        instance_uuid=INSTANCE_UUID,
+        workspace_uuid=workspace_uuid,
+        placement_generation=1,
+        entitlement_revision=1,
+    )
+
+
+@pytest.mark.asyncio
+async def test_cloud_plugin_quota_is_atomic_isolated_and_persists_across_connector_restart() -> None:
+    manager = _Persistence()
+    connector = await _connector(manager)
+
+    async def install(workspace_uuid: str, index: int):
+        return await connector._persist_installation_package(
+            _context(workspace_uuid),
+            plugin_author='test-author',
+            plugin_name=f'plugin-{index}',
+            install_source=PluginInstallSource.MARKETPLACE,
+            install_info={'author': 'test-author', 'name': f'plugin-{index}'},
+            artifact_digest=f'{index:064x}',
+        )
+
+    results = await asyncio.gather(
+        *(install(WORKSPACE_A, index) for index in range(10)),
+        *(install(WORKSPACE_B, index) for index in range(10)),
+        return_exceptions=True,
+    )
+
+    successes = [result for result in results if isinstance(result, tuple)]
+    failures = [result for result in results if isinstance(result, WorkspaceQuotaExceededError)]
+    assert len(successes) == 6
+    assert len(failures) == 14
+    assert len(manager.plugins[WORKSPACE_A]) == 3
+    assert len(manager.plugins[WORKSPACE_B]) == 3
+    assert manager.workspace_locks_seen == 20
+
+    restarted_connector = await _connector(manager)
+    with pytest.raises(WorkspaceQuotaExceededError, match=r'Maximum number of plugins \(3\) reached'):
+        await restarted_connector._persist_installation_package(
+            _context(WORKSPACE_A),
+            plugin_author='test-author',
+            plugin_name='after-restart',
+            install_source=PluginInstallSource.MARKETPLACE,
+            install_info={},
+            artifact_digest='f' * 64,
+        )
+    assert len(manager.plugins[WORKSPACE_A]) == 3
+
+    installed_name = next(iter(manager.plugins[WORKSPACE_A]))[1]
+
+    async def reinstall():
+        return await restarted_connector._persist_installation_package(
+            _context(WORKSPACE_A),
+            plugin_author='test-author',
+            plugin_name=installed_name,
+            install_source=PluginInstallSource.MARKETPLACE,
+            install_info={'author': 'test-author', 'name': installed_name, 'revision': 2},
+            artifact_digest='e' * 64,
+        )
+
+    reinstall_results = await asyncio.gather(reinstall(), reinstall())
+    assert all(result[2] is True for result in reinstall_results)
+
+    mixed_results = await asyncio.gather(
+        reinstall(),
+        restarted_connector._persist_installation_package(
+            _context(WORKSPACE_A),
+            plugin_author='test-author',
+            plugin_name='new-at-capacity',
+            install_source=PluginInstallSource.MARKETPLACE,
+            install_info={},
+            artifact_digest='d' * 64,
+        ),
+        return_exceptions=True,
+    )
+    assert isinstance(mixed_results[0], tuple)
+    assert isinstance(mixed_results[1], WorkspaceQuotaExceededError)
+    assert len(manager.plugins[WORKSPACE_A]) == 3