xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-07-21 11:56:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(tenancy): establish cloud isolation foundations

Browse Source
This commit is contained in:
Junyan Qin
2026-07-19 22:39:58 +08:00
parent 59b1570ead
commit 41772920ef
41 changed files with 1984 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/integration/persistence/test_migrations_postgres.py
+321 -3
View File
@@ -15,15 +15,20 @@ from __future__ import annotations
import logging
import os
import uuid
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import sqlalchemy as sa
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from sqlalchemy import text
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.persistence.mgr import PersistenceManager
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
from langbot.pkg.persistence.tenant_uow import TENANT_POLICY_NAME, TENANT_TABLE_COLUMNS, TenantUnitOfWork
from langbot.pkg.persistence.alembic_runner import (
run_alembic_upgrade,
run_alembic_stamp,
@@ -49,6 +54,44 @@ def _get_script_head() -> str:
return ScriptDirectory.from_config(cfg).get_current_head()
def _application_for_postgres_url(postgres_url: str, logger_name: str) -> SimpleNamespace:
url = sa.engine.make_url(postgres_url)
return SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'use': 'postgresql',
'postgresql': {
'host': url.host,
'port': url.port,
'user': url.username,
'password': url.password,
'database': url.database,
},
}
}
),
logger=logging.getLogger(logger_name),
)
async def _dispose_manager(manager: PersistenceManager | None) -> None:
if manager is not None and getattr(manager, 'db', None) is not None:
await manager.get_db_engine().dispose()
def _restore_postgres_manager_registry(monkeypatch) -> None:
"""Undo the registry isolation used by test_database_decorator.py."""
from langbot.pkg.persistence import mgr as persistence_mgr_module
from langbot.pkg.persistence.databases.postgresql import PostgreSQLDatabaseManager
monkeypatch.setattr(
persistence_mgr_module.database,
'preregistered_managers',
[PostgreSQLDatabaseManager],
)
pytestmark = [pytest.mark.integration, pytest.mark.slow]
@@ -309,7 +352,7 @@ class TestPostgreSQLWorkspaceMigration:
)
assert 'workspaces' not in tables_before_migration
await manager._run_alembic_migrations()
await manager._initialize_managed_schema()
async with postgres_engine.connect() as conn:
account = (await conn.execute(text('SELECT uuid, status, source FROM users'))).mappings().one()
@@ -425,3 +468,278 @@ class TestPostgreSQLResourceTenancyMigration:
),
{'workspace_uuid': second_workspace_uuid},
)
class TestPostgreSQLTenantRuntime:
"""Release bootstrap, RLS enforcement, and runtime-role safety."""
@pytest.mark.asyncio
async def test_oss_postgres_defaults_to_the_singleton_workspace(
self,
postgres_url,
clean_tables,
clean_alembic_version,
monkeypatch,
):
instance_uuid = 'oss-postgres-rls-compatibility-test'
_restore_postgres_manager_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', instance_uuid)
manager = PersistenceManager(
_application_for_postgres_url(postgres_url, 'postgres-oss-rls-compatibility-test'),
mode=PersistenceMode.OSS_COMPAT,
)
try:
await manager.initialize()
async with manager.get_db_engine().connect() as conn:
workspace_uuid = await conn.scalar(text("SELECT uuid FROM workspaces WHERE source = 'local'"))
tenant_setting = await conn.scalar(text("SELECT current_setting('langbot.workspace_uuid', true)"))
visible_workspaces = await conn.scalar(text('SELECT COUNT(*) FROM workspaces'))
assert workspace_uuid == tenant_setting
assert visible_workspaces == 1
finally:
await _dispose_manager(manager)
@pytest.mark.asyncio
async def test_release_bootstrap_and_runtime_isolation(
self,
postgres_url,
postgres_engine,
clean_tables,
clean_alembic_version,
monkeypatch,
):
instance_uuid = 'cloud-runtime-persistence-test'
workspace_a = '10000000-0000-0000-0000-000000000001'
workspace_b = '20000000-0000-0000-0000-000000000002'
role_suffix = uuid.uuid4().hex[:12]
runtime_role = f'lb_runtime_{role_suffix}'
bypass_role = f'lb_bypass_{role_suffix}'
owner_role = f'lb_owner_{role_suffix}'
role_password = f'Lb{uuid.uuid4().hex}'
created_roles: list[str] = []
managers: list[PersistenceManager] = []
runtime_engine: AsyncEngine | None = None
owner_changed = False
_restore_postgres_manager_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', instance_uuid)
release_manager = PersistenceManager(
_application_for_postgres_url(postgres_url, 'postgres-release-bootstrap-test'),
mode=PersistenceMode.RELEASE_MIGRATION,
)
managers.append(release_manager)
async with postgres_engine.connect() as conn:
admin_user = await conn.scalar(text('SELECT current_user'))
quote = postgres_engine.dialect.identifier_preparer.quote
def role_url(role_name: str) -> str:
return (
sa.engine.make_url(postgres_url)
.set(
username=role_name,
password=role_password,
)
.render_as_string(hide_password=False)
)
async def create_role(role_name: str, *, bypass_rls: bool = False) -> None:
bypass_clause = ' BYPASSRLS' if bypass_rls else ''
async with postgres_engine.connect() as conn:
await conn.execute(
text(f"CREATE ROLE {quote(role_name)} LOGIN{bypass_clause} PASSWORD '{role_password}'")
)
await conn.execute(
text(
f'GRANT CONNECT ON DATABASE {quote(sa.engine.make_url(postgres_url).database)} '
f'TO {quote(role_name)}'
)
)
await conn.execute(text(f'GRANT USAGE ON SCHEMA public TO {quote(role_name)}'))
await conn.execute(
text(f'GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO {quote(role_name)}')
)
await conn.execute(text(f'GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO {quote(role_name)}'))
created_roles.append(role_name)
try:
await release_manager.initialize()
release_engine = release_manager.get_db_engine()
assert await get_alembic_current(release_engine) == _get_script_head()
async with release_engine.connect() as conn:
assert await conn.scalar(text('SELECT COUNT(*) FROM workspaces')) == 0
rls_rows = (
(
await conn.execute(
text(
"""
SELECT c.relname, c.relrowsecurity, c.relforcerowsecurity,
EXISTS (
SELECT 1 FROM pg_policy p
WHERE p.polrelid = c.oid AND p.polname = :policy_name
) AS has_policy
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE n.nspname = current_schema() AND c.relname IN :table_names
"""
).bindparams(sa.bindparam('table_names', expanding=True)),
{
'policy_name': TENANT_POLICY_NAME,
'table_names': tuple(TENANT_TABLE_COLUMNS),
},
)
)
.mappings()
.all()
)
assert {row['relname'] for row in rls_rows} == set(TENANT_TABLE_COLUMNS)
assert all(row['relrowsecurity'] and row['relforcerowsecurity'] and row['has_policy'] for row in rls_rows)
for workspace_uuid, slug in ((workspace_a, 'workspace-a'), (workspace_b, 'workspace-b')):
async with TenantUnitOfWork(release_engine, workspace_uuid) as uow:
await uow.execute(
text(
"""
INSERT INTO workspaces
(uuid, instance_uuid, name, slug, type, status, source, projection_revision)
VALUES
(:uuid, :instance_uuid, :name, :slug, 'team', 'active', 'cloud_projection', 0)
"""
),
{
'uuid': workspace_uuid,
'instance_uuid': instance_uuid,
'name': slug,
'slug': slug,
},
)
await uow.execute(
text(
'INSERT INTO workspace_metadata (workspace_uuid, key, value) '
"VALUES (:workspace_uuid, 'seed', :value)"
),
{'workspace_uuid': workspace_uuid, 'value': slug},
)
await create_role(runtime_role)
runtime_engine = create_async_engine(role_url(runtime_role), pool_size=1, max_overflow=0)
async with runtime_engine.connect() as conn:
assert (await conn.execute(text('SELECT uuid FROM workspaces'))).all() == []
assert (await conn.execute(text('SELECT * FROM workspace_metadata'))).all() == []
with pytest.raises(sa.exc.DBAPIError):
async with runtime_engine.begin() as conn:
await conn.execute(
text(
'INSERT INTO workspace_metadata (workspace_uuid, key, value) '
"VALUES (:workspace_uuid, 'no-scope', 'rejected')"
),
{'workspace_uuid': workspace_a},
)
async with TenantUnitOfWork(runtime_engine, workspace_a) as uow:
assert (await uow.execute(text('SELECT uuid FROM workspaces'))).scalars().all() == [workspace_a]
assert (
await uow.execute(
text('SELECT uuid FROM workspaces WHERE uuid = :workspace_uuid'),
{'workspace_uuid': workspace_b},
)
).all() == []
with pytest.raises(sa.exc.DBAPIError):
async with TenantUnitOfWork(runtime_engine, workspace_a) as uow:
await uow.execute(
text(
'INSERT INTO workspace_metadata (workspace_uuid, key, value) '
"VALUES (:workspace_uuid, 'cross-scope', 'rejected')"
),
{'workspace_uuid': workspace_b},
)
async with TenantUnitOfWork(runtime_engine, workspace_a) as uow:
assert (await uow.execute(text('SELECT value FROM workspace_metadata'))).scalars().all() == [
'workspace-a'
]
async with TenantUnitOfWork(runtime_engine, workspace_b) as uow:
assert (await uow.execute(text('SELECT value FROM workspace_metadata'))).scalars().all() == [
'workspace-b'
]
async with runtime_engine.connect() as conn:
setting = await conn.scalar(text("SELECT current_setting('langbot.workspace_uuid', true)"))
assert setting in (None, '')
assert await conn.scalar(text('SELECT COUNT(*) FROM workspace_metadata')) == 0
with pytest.raises(RuntimeError, match='force rollback'):
async with TenantUnitOfWork(runtime_engine, workspace_a) as uow:
await uow.execute(
text(
'INSERT INTO workspace_metadata (workspace_uuid, key, value) '
"VALUES (:workspace_uuid, 'rolled-back', 'no')"
),
{'workspace_uuid': workspace_a},
)
raise RuntimeError('force rollback')
async with TenantUnitOfWork(runtime_engine, workspace_a) as uow:
assert (
await uow.session.scalar(text("SELECT COUNT(*) FROM workspace_metadata WHERE key = 'rolled-back'"))
== 0
)
async with runtime_engine.connect() as conn:
setting = await conn.scalar(text("SELECT current_setting('langbot.workspace_uuid', true)"))
assert setting in (None, '')
cloud_manager = PersistenceManager(
_application_for_postgres_url(role_url(runtime_role), 'postgres-cloud-runtime-test'),
mode=PersistenceMode.CLOUD_RUNTIME,
)
cloud_manager.create_tables = AsyncMock(side_effect=AssertionError('Cloud runtime attempted create_all'))
cloud_manager._run_alembic_migrations = AsyncMock(
side_effect=AssertionError('Cloud runtime attempted an Alembic upgrade')
)
managers.append(cloud_manager)
await cloud_manager.initialize()
cloud_manager.create_tables.assert_not_awaited()
cloud_manager._run_alembic_migrations.assert_not_awaited()
superuser_manager = PersistenceManager(
_application_for_postgres_url(postgres_url, 'postgres-superuser-runtime-test'),
mode=PersistenceMode.CLOUD_RUNTIME,
)
managers.append(superuser_manager)
with pytest.raises(RuntimeError, match='must not be a superuser'):
await superuser_manager.initialize()
await create_role(bypass_role, bypass_rls=True)
bypass_manager = PersistenceManager(
_application_for_postgres_url(role_url(bypass_role), 'postgres-bypass-runtime-test'),
mode=PersistenceMode.CLOUD_RUNTIME,
)
managers.append(bypass_manager)
with pytest.raises(RuntimeError, match='must not have BYPASSRLS'):
await bypass_manager.initialize()
await create_role(owner_role)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER TABLE workspace_metadata OWNER TO {quote(owner_role)}'))
owner_changed = True
owner_manager = PersistenceManager(
_application_for_postgres_url(role_url(owner_role), 'postgres-owner-runtime-test'),
mode=PersistenceMode.CLOUD_RUNTIME,
)
managers.append(owner_manager)
with pytest.raises(RuntimeError, match='must not own tenant tables'):
await owner_manager.initialize()
finally:
if runtime_engine is not None:
await runtime_engine.dispose()
for manager in reversed(managers):
await _dispose_manager(manager)
async with postgres_engine.connect() as conn:
if owner_changed:
await conn.execute(text(f'ALTER TABLE workspace_metadata OWNER TO {quote(admin_user)}'))
for role_name in reversed(created_roles):
await conn.execute(text(f'DROP OWNED BY {quote(role_name)}'))
await conn.execute(text(f'DROP ROLE {quote(role_name)}'))
tests/integration/persistence/test_sqlite_migration_backup.py
+2 -2
View File
@@ -47,7 +47,7 @@ async def test_tenancy_migrations_retain_verified_boundary_backups(tmp_path):
await _manager(engine)._run_alembic_migrations()
assert await alembic_runner.get_alembic_current(engine) == '0010_scope_resources'
assert await alembic_runner.get_alembic_current(engine) == alembic_runner.get_alembic_head()
payloads = _manifest_payloads(tmp_path / 'migration-backups')
assert len(payloads) == 2
assert {
@@ -102,6 +102,6 @@ async def test_failed_tenancy_migration_restores_backup_and_revision(
monkeypatch.setattr(alembic_runner, 'run_alembic_upgrade', real_upgrade)
await _manager(engine)._run_alembic_migrations()
assert await alembic_runner.get_alembic_current(engine) == '0010_scope_resources'
assert await alembic_runner.get_alembic_current(engine) == alembic_runner.get_alembic_head()
finally:
await engine.dispose()
tests/integration/persistence/test_workspace_migration.py
+2 -1
View File
@@ -13,6 +13,7 @@ from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.persistence.mgr import PersistenceManager
from langbot.pkg.persistence.alembic_runner import (
get_alembic_head,
get_alembic_current,
run_alembic_downgrade,
run_alembic_stamp,
@@ -126,7 +127,7 @@ async def test_legacy_instance_gets_stable_accounts_and_default_workspace(legacy
assert execution_state['state'] == 'active'
assert execution_state['write_fenced'] in (False, 0)
assert await get_alembic_current(legacy_engine) == '0010_scope_resources'
assert await get_alembic_current(legacy_engine) == get_alembic_head()
async def test_workspace_upgrade_is_idempotent_and_preserves_identifiers(legacy_engine):
tests/unit_tests/api/http/test_entitlement_context.py
+78
View File
@@ -0,0 +1,78 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import quart
from langbot.pkg.api.http.context import (
ExecutionContext,
PrincipalContext,
PrincipalType,
RequestContext,
WorkspaceContext,
)
from langbot.pkg.api.http.controller.group import RouterGroup
from langbot.pkg.cloud.entitlements import EntitlementSnapshot, EntitlementUnavailableError
from langbot.pkg.cloud.entitlements import EntitlementResolver
class _Group(RouterGroup):
async def initialize(self) -> None:
return None
def _router(deployment) -> _Group:
provider = getattr(deployment, 'entitlement_provider', None)
resolver = EntitlementResolver('instance-a', provider) if provider is not None else None
ap = SimpleNamespace(deployment=deployment, entitlement_resolver=resolver)
return _Group(ap, quart.Quart(__name__))
@pytest.mark.asyncio
async def test_cloud_request_resolves_verified_entitlement_revision():
snapshot = EntitlementSnapshot(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
entitlement_revision=9,
status='active',
not_before=1,
expires_at=4_000_000_000,
features={},
limits={},
)
provider = SimpleNamespace(get_workspace_entitlement=AsyncMock(return_value=snapshot))
router = _router(SimpleNamespace(multi_workspace_enabled=True, entitlement_provider=provider))
revision = await router._resolve_entitlement_revision('instance-a', 'workspace-a')
assert revision == 9
provider.get_workspace_entitlement.assert_awaited_once_with('workspace-a')
@pytest.mark.asyncio
async def test_cloud_request_fails_closed_without_entitlement_provider():
router = _router(SimpleNamespace(multi_workspace_enabled=True, entitlement_provider=None))
with pytest.raises(EntitlementUnavailableError):
await router._resolve_entitlement_revision('instance-a', 'workspace-a')
def test_execution_context_preserves_entitlement_revision():
request = RequestContext(
instance_uuid='instance-a',
placement_generation=1,
request_id='request-a',
auth_type='user-token',
principal=PrincipalContext(PrincipalType.ACCOUNT, account_uuid='account-a'),
workspace=WorkspaceContext(
workspace_uuid='workspace-a',
membership_uuid='membership-a',
role='owner',
permissions=frozenset(),
),
entitlement_revision=11,
)
assert ExecutionContext.from_request(request).entitlement_revision == 11
tests/unit_tests/api/service/test_mcp_service.py
+40
View File
@@ -29,6 +29,7 @@ from langbot.pkg.api.http.context import (
)
from langbot.pkg.api.http.service.mcp import MCPService, redact_mcp_secrets, restore_mcp_secret_placeholders
from langbot.pkg.entity.persistence.mcp import MCPServer
from langbot.pkg.provider.tools.loaders.mcp_policy import MCPStdioDisabledError
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
@@ -351,6 +352,27 @@ class TestMCPServiceGetMCPServers:
class TestMCPServiceCreateMCPServer:
"""Tests for create_mcp_server method."""
async def test_create_stdio_rejected_by_independent_instance_gate(self):
ap = SimpleNamespace(
instance_config=SimpleNamespace(
data={
'mcp': {'stdio': {'enabled': False}},
'system': {'limitation': {'max_extensions': -1}},
}
),
persistence_mgr=SimpleNamespace(execute_async=AsyncMock()),
tool_mgr=None,
)
service = _service(ap)
with pytest.raises(MCPStdioDisabledError, match='disabled by instance policy'):
await service.create_mcp_server(
_CONTEXT,
{'name': 'local', 'mode': 'stdio', 'enable': True, 'extra_args': {}},
)
ap.persistence_mgr.execute_async.assert_not_awaited()
async def test_create_mcp_server_max_extensions_reached_raises(self):
"""Raises ValueError when max_extensions limit reached."""
# Setup
@@ -887,6 +909,24 @@ class TestMCPServiceDeleteMCPServer:
class TestMCPServiceTestMCPServer:
"""Tests for test_mcp_server method."""
async def test_transient_stdio_test_rejected_by_instance_gate(self):
ap = SimpleNamespace(
instance_config=SimpleNamespace(data={'mcp': {'stdio': {'enabled': False}}}),
tool_mgr=SimpleNamespace(mcp_tool_loader=SimpleNamespace(load_mcp_server=AsyncMock())),
task_mgr=SimpleNamespace(create_user_task=Mock()),
)
service = _service(ap)
with pytest.raises(MCPStdioDisabledError, match='disabled by instance policy'):
await service.test_mcp_server(
_CONTEXT,
'_',
{'name': 'local', 'mode': 'stdio', 'enable': True, 'extra_args': {}},
)
ap.tool_mgr.mcp_tool_loader.load_mcp_server.assert_not_awaited()
ap.task_mgr.create_user_task.assert_not_called()
async def test_test_mcp_server_existing_server(self):
"""Tests existing MCP server connection."""
# Setup
tests/unit_tests/cloud/test_bootstrap.py
+136
View File
@@ -0,0 +1,136 @@
from __future__ import annotations
from types import SimpleNamespace
import pytest
from langbot.pkg.cloud.bootstrap import (
CloudBootstrapError,
OpenSourceDeployment,
VerifiedCloudDeployment,
resolve_deployment,
)
from langbot.pkg.cloud.entitlements import EntitlementSnapshot
pytestmark = pytest.mark.asyncio
class _Entitlements:
async def get_workspace_entitlement(self, workspace_uuid: str) -> EntitlementSnapshot:
return EntitlementSnapshot(
instance_uuid='instance-a',
workspace_uuid=workspace_uuid,
entitlement_revision=1,
status='active',
not_before=1,
expires_at=4_000_000_000,
features={'managed_sandbox': True},
limits={'managed_sandbox_sessions': 1},
)
class _Provider:
def bootstrap(self, *, instance_uuid: str, instance_config: dict):
del instance_config
return VerifiedCloudDeployment(
instance_uuid=instance_uuid,
manifest_jti='manifest-a',
manifest_generation=3,
expires_at=4_000_000_000,
release='cloud-v2',
capabilities=frozenset({'multi_workspace_v2'}),
tenant_isolation_version=2,
entitlement_provider=_Entitlements(),
verification_key_id='root-2026',
)
class _EntryPoint:
def __init__(self, value):
self.value = value
def load(self):
return self.value
class _EntryPoints(list):
def select(self, *, group: str):
return self if group == 'langbot.cloud_bootstrap' else []
def _cloud_config() -> dict:
return {
'database': {'use': 'postgresql'},
'vdb': {'use': 'pgvector'},
'mcp': {'stdio': {'enabled': False}},
# Proves mutable product metadata does not participate in selection.
'system': {'edition': 'community'},
}
async def test_no_closed_entry_point_selects_oss_singleton_even_if_edition_says_cloud():
deployment = await resolve_deployment(
instance_uuid='instance-a',
instance_config={'system': {'edition': 'cloud'}},
entry_points=lambda: _EntryPoints(),
)
assert isinstance(deployment, OpenSourceDeployment)
assert deployment.multi_workspace_enabled is False
async def test_verified_closed_entry_point_activates_cloud_policy():
deployment = await resolve_deployment(
instance_uuid='instance-a',
instance_config=_cloud_config(),
entry_points=lambda: _EntryPoints([_EntryPoint(_Provider)]),
now=1_000,
)
assert isinstance(deployment, VerifiedCloudDeployment)
assert deployment.multi_workspace_enabled is True
assert deployment.persistence_mode == 'cloud_runtime'
@pytest.mark.parametrize(
('field', 'value', 'message'),
[
('database', {'use': 'sqlite'}, 'database.use=postgresql'),
('vdb', {'use': 'chroma'}, 'vdb.use=pgvector'),
('mcp', {'stdio': {'enabled': True}}, 'mcp.stdio.enabled=false'),
],
)
async def test_cloud_runtime_config_is_fail_closed(field, value, message):
config = _cloud_config()
config[field] = value
with pytest.raises(CloudBootstrapError, match=message):
await resolve_deployment(
instance_uuid='instance-a',
instance_config=config,
entry_points=lambda: _EntryPoints([_EntryPoint(_Provider())]),
now=1_000,
)
async def test_invalid_provider_never_falls_back_to_oss():
provider = SimpleNamespace(bootstrap=lambda **_: object())
with pytest.raises(CloudBootstrapError, match='must return VerifiedCloudDeployment'):
await resolve_deployment(
instance_uuid='instance-a',
instance_config=_cloud_config(),
entry_points=lambda: _EntryPoints([_EntryPoint(provider)]),
now=1_000,
)
async def test_duplicate_closed_providers_fail_closed():
with pytest.raises(CloudBootstrapError, match='Exactly one'):
await resolve_deployment(
instance_uuid='instance-a',
instance_config=_cloud_config(),
entry_points=lambda: _EntryPoints([_EntryPoint(_Provider()), _EntryPoint(_Provider())]),
now=1_000,
)
tests/unit_tests/cloud/test_entitlements.py
+86
View File
@@ -0,0 +1,86 @@
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock
from langbot.pkg.cloud.entitlements import EntitlementResolver, EntitlementSnapshot, EntitlementUnavailableError
def _snapshot(**overrides) -> EntitlementSnapshot:
values = {
'instance_uuid': 'instance-a',
'workspace_uuid': 'workspace-a',
'entitlement_revision': 7,
'status': 'active',
'not_before': 100,
'expires_at': 200,
'features': {'managed_sandbox': True, 'mcp_stdio': False},
'limits': {'managed_sandbox_sessions': 1},
}
values.update(overrides)
return EntitlementSnapshot(**values)
def test_active_snapshot_exposes_only_generic_features_and_limits():
snapshot = _snapshot().require_active(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
now=150,
)
snapshot.require_feature('managed_sandbox')
assert snapshot.limit('managed_sandbox_sessions') == 1
assert 'plan' not in snapshot.model_fields
@pytest.mark.parametrize(
'snapshot,now',
[
(_snapshot(status='suspended'), 150),
(_snapshot(), 99),
(_snapshot(), 200),
],
)
def test_inactive_or_expired_snapshot_fails_closed(snapshot, now):
with pytest.raises(EntitlementUnavailableError):
snapshot.require_active(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
now=now,
)
def test_scope_mismatch_fails_closed():
with pytest.raises(EntitlementUnavailableError, match='scope'):
_snapshot().require_active(
instance_uuid='instance-a',
workspace_uuid='workspace-b',
now=150,
)
@pytest.mark.asyncio
async def test_resolver_rejects_revision_rollback():
provider = AsyncMock()
provider.get_workspace_entitlement = AsyncMock(side_effect=[_snapshot(), _snapshot(entitlement_revision=6)])
resolver = EntitlementResolver('instance-a', provider)
await resolver.resolve('workspace-a', now=150)
with pytest.raises(EntitlementUnavailableError, match='rolled back'):
await resolver.resolve('workspace-a', now=150)
@pytest.mark.asyncio
async def test_resolver_rejects_same_revision_with_different_contents():
provider = AsyncMock()
provider.get_workspace_entitlement = AsyncMock(
side_effect=[
_snapshot(),
_snapshot(features={'managed_sandbox': False}),
]
)
resolver = EntitlementResolver('instance-a', provider)
await resolver.resolve('workspace-a', now=150)
with pytest.raises(EntitlementUnavailableError, match='conflicting contents'):
await resolver.resolve('workspace-a', now=150)
tests/unit_tests/core/test_load_config.py
+54
View File
@@ -259,6 +259,60 @@ class TestApplyEnvOverridesToConfig:
assert result['system']['enable'] is False
assert result['concurrency']['pipeline'] == 10
def test_plugin_worker_and_stdio_policy_native_env_overrides(self):
load_config = get_load_config_module()
cfg = {
'plugin': {
'worker': {
'max_cpus': 1.0,
'max_memory_mb': 512,
'max_pids': 128,
'max_open_files': 256,
'max_file_size_mb': 512,
}
},
'mcp': {'stdio': {'enabled': True}},
}
env = {
'PLUGIN__WORKER__MAX_CPUS': '2.5',
'PLUGIN__WORKER__MAX_MEMORY_MB': '1024',
'PLUGIN__WORKER__MAX_PIDS': '64',
'PLUGIN__WORKER__MAX_OPEN_FILES': '128',
'PLUGIN__WORKER__MAX_FILE_SIZE_MB': '256',
'MCP__STDIO__ENABLED': 'false',
}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['plugin']['worker'] == {
'max_cpus': 2.5,
'max_memory_mb': 1024,
'max_pids': 64,
'max_open_files': 128,
'max_file_size_mb': 256,
}
assert result['mcp']['stdio']['enabled'] is False
def test_runtime_policy_defaults_preserve_env_types_for_upgraded_config(self):
load_config = get_load_config_module()
cfg = {'plugin': {'enable': True}}
completed = load_config._complete_runtime_policy_defaults(cfg)
with patch.dict(
os.environ,
{
'PLUGIN__WORKER__MAX_MEMORY_MB': '768',
'MCP__STDIO__ENABLED': 'false',
},
clear=True,
):
result = load_config._apply_env_overrides_to_config(completed)
assert result['plugin']['worker']['max_memory_mb'] == 768
assert isinstance(result['plugin']['worker']['max_memory_mb'], int)
assert result['mcp']['stdio']['enabled'] is False
def test_webhook_prefix_override(self):
"""Test overriding webhook_prefix via environment variable."""
load_config = get_load_config_module()
tests/unit_tests/persistence/test_tenant_uow.py
+63
View File
@@ -0,0 +1,63 @@
from __future__ import annotations
import pytest
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
from langbot.pkg.persistence.tenant_uow import TenantUnitOfWork
pytestmark = pytest.mark.asyncio
async def test_sqlite_tenant_uow_commits_and_rolls_back() -> None:
engine = create_async_engine('sqlite+aiosqlite:///:memory:')
table = sa.Table(
'uow_rows',
sa.MetaData(),
sa.Column('id', sa.Integer, primary_key=True),
sa.Column('workspace_uuid', sa.String(36), nullable=False),
)
try:
async with engine.begin() as conn:
await conn.run_sync(table.metadata.create_all)
async with TenantUnitOfWork(engine, 'workspace-a') as uow:
await uow.execute(sa.insert(table).values(id=1, workspace_uuid='workspace-a'))
with pytest.raises(RuntimeError, match='roll back'):
async with TenantUnitOfWork(engine, 'workspace-a') as uow:
await uow.execute(sa.insert(table).values(id=2, workspace_uuid='workspace-a'))
raise RuntimeError('roll back this transaction')
async with engine.connect() as conn:
rows = (await conn.execute(sa.select(table.c.id).order_by(table.c.id))).scalars().all()
assert rows == [1]
finally:
await engine.dispose()
async def test_tenant_uow_is_single_use_and_requires_an_active_scope() -> None:
engine = create_async_engine('sqlite+aiosqlite:///:memory:')
uow = TenantUnitOfWork(engine, 'workspace-a')
try:
with pytest.raises(RuntimeError, match='not active'):
_ = uow.session
async with uow:
assert uow.session.in_transaction()
with pytest.raises(RuntimeError, match='cannot be reused'):
async with uow:
pass
finally:
await engine.dispose()
def test_persistence_mode_must_be_a_trusted_enum() -> None:
with pytest.raises(TypeError, match='trusted PersistenceMode'):
PersistenceManager(object(), mode='cloud_runtime') # type: ignore[arg-type]
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
assert manager.mode is PersistenceMode.CLOUD_RUNTIME
tests/unit_tests/provider/test_mcp_box_integration.py
+26
View File
@@ -154,6 +154,7 @@ def mcp_module():
def _make_ap():
ap = Mock()
ap.logger = Mock()
ap.instance_config = SimpleNamespace(data={'mcp': {'stdio': {'enabled': True}}})
ap.workspace_service = Mock()
ap.workspace_service.get_execution_binding = AsyncMock(
return_value=SimpleNamespace(
@@ -791,6 +792,31 @@ class TestBoxConfigParsing:
assert s.box_config.host_path_mode == 'ro'
@pytest.mark.asyncio
async def test_stdio_instance_gate_runs_before_box_transport(mcp_module):
ap = _make_ap()
ap.instance_config.data['mcp']['stdio']['enabled'] = False
ap.box_service.available = True
session = _make_session(
mcp_module,
{
'name': 'blocked',
'uuid': 'blocked-uuid',
'mode': 'stdio',
'command': 'python',
'args': [],
'env': {},
},
ap=ap,
)
session._box_stdio_runtime.initialize = AsyncMock()
with pytest.raises(RuntimeError, match='disabled by instance policy'):
await session._init_stdio_python_server()
session._box_stdio_runtime.initialize.assert_not_awaited()
@pytest.mark.asyncio
async def test_init_box_stdio_server_stages_host_path_in_shared_workspace(mcp_module, tmp_path):
mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio']
tests/unit_tests/provider/test_mcp_stdio_policy.py
+70
View File
@@ -0,0 +1,70 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.provider.tools.loaders.mcp_policy import (
MCPStdioDisabledError,
require_stdio_mcp_enabled,
stdio_mcp_enabled,
)
from langbot.pkg.provider.tools.loaders.mcp import MCPLoader
def _app(config: dict) -> SimpleNamespace:
return SimpleNamespace(instance_config=SimpleNamespace(data=config))
def test_oss_default_remains_enabled_when_key_is_absent():
assert stdio_mcp_enabled(_app({})) is True
@pytest.mark.parametrize(
'value',
[False, 'false', 0, None, {}, []],
)
def test_disabled_or_invalid_values_fail_closed(value):
ap = _app({'mcp': {'stdio': {'enabled': value}}})
assert stdio_mcp_enabled(ap) is False
with pytest.raises(MCPStdioDisabledError, match='disabled by instance policy'):
require_stdio_mcp_enabled(ap, {'mode': 'stdio'})
def test_remote_transport_is_independent_of_stdio_gate():
ap = _app({'mcp': {'stdio': {'enabled': False}}})
require_stdio_mcp_enabled(ap, {'mode': 'remote'})
@pytest.mark.asyncio
async def test_bootstrap_retains_but_does_not_launch_disabled_stdio_rows():
server = SimpleNamespace(uuid='server-a', workspace_uuid='workspace-a')
result = Mock()
result.all.return_value = [server]
ap = _app({'mcp': {'stdio': {'enabled': False}}})
ap.logger = Mock()
ap.persistence_mgr = SimpleNamespace(
execute_async=AsyncMock(return_value=result),
serialize_model=Mock(
return_value={
'uuid': 'server-a',
'workspace_uuid': 'workspace-a',
'name': 'local',
'mode': 'stdio',
'enable': True,
'extra_args': {},
}
),
)
ap.workspace_service = SimpleNamespace(get_execution_binding=AsyncMock())
loader = MCPLoader(ap)
loader.host_mcp_server = AsyncMock()
await loader.load_mcp_servers_from_db()
loader.host_mcp_server.assert_not_awaited()
ap.workspace_service.get_execution_binding.assert_not_awaited()
assert loader.sessions == {}