mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 04:40:57 +00:00
feat(cloud): enforce workspace resource quotas
This commit is contained in:
dadachann
@@ -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')
|
||||
|
||||
@@ -311,10 +311,9 @@ class TestBotServiceCreateBot:
|
||||
ap.platform_mgr = SimpleNamespace()
|
||||
ap.platform_mgr.load_bot = AsyncMock()
|
||||
|
||||
# Mock get_bots to return 2 bots already
|
||||
bot1 = _create_mock_bot(bot_uuid='uuid-1')
|
||||
bot2 = _create_mock_bot(bot_uuid='uuid-2')
|
||||
mock_result = _create_mock_result([bot1, bot2])
|
||||
# Mock the atomic count query to report 2 existing bots.
|
||||
mock_result = _create_mock_result()
|
||||
mock_result.scalar_one = Mock(return_value=2)
|
||||
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
|
||||
ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'uuid-1', 'name': 'Bot 1'})
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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',
|
||||
)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user