xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-07-24 13:26:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(tenancy): connect cloud workspace control plane

Browse Source
This commit is contained in:
Junyan Qin
2026-07-24 19:11:33 +08:00
parent d7cdd206c2
commit 98f45aa88e
40 changed files with 4159 additions and 452 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/integration/persistence/test_migrations.py
+19
View File
@@ -15,6 +15,7 @@ from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.persistence.alembic_runner import (
run_alembic_downgrade,
run_alembic_upgrade,
run_alembic_stamp,
get_alembic_current,
@@ -171,6 +172,24 @@ class TestSQLiteMigrationUpgrade:
)
assert 'embedding_dimension' in columns
@pytest.mark.asyncio
async def test_directory_projection_upgrade_downgrade_round_trip(self, sqlite_engine):
await run_alembic_stamp(sqlite_engine, '0013_tenant_pgvector')
await run_alembic_upgrade(sqlite_engine, 'head')
async with sqlite_engine.connect() as conn:
tables = await conn.run_sync(lambda sync_conn: set(sqlalchemy.inspect(sync_conn).get_table_names()))
assert {'directory_projection_states', 'directory_projection_inbox'} <= tables
await run_alembic_downgrade(sqlite_engine, '0013_tenant_pgvector')
async with sqlite_engine.connect() as conn:
tables = await conn.run_sync(lambda sync_conn: set(sqlalchemy.inspect(sync_conn).get_table_names()))
assert 'directory_projection_states' not in tables
assert 'directory_projection_inbox' not in tables
await run_alembic_upgrade(sqlite_engine, 'head')
assert await get_alembic_current(sqlite_engine) == _get_script_head()
class TestSQLiteMigrationFreshDatabase:
"""Tests for fresh database workflow."""
tests/integration/persistence/test_migrations_postgres.py
+208 -3
View File
@@ -31,7 +31,13 @@ from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
from sqlalchemy import text
from langbot.pkg.cloud.directory import DirectoryEvent, DirectoryEventBatch
from langbot.pkg.cloud.directory_projection import DirectoryProjectionService
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.cloud_directory import (
DirectoryProjectionInbox,
DirectoryProjectionState,
)
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
from langbot.pkg.persistence.tenant_uow import (
@@ -82,6 +88,17 @@ from langbot.pkg.rag.knowledge.kbmgr import RAGManager
from .resource_migration_support import TENANT_TABLES, create_legacy_resource_schema
class _NoopDirectoryProjectionProvider:
async def fetch_snapshot(self, instance_uuid: str):
raise AssertionError(f'Unexpected snapshot fetch for {instance_uuid}')
async def fetch_events(self, instance_uuid: str, after_cursor: int, limit: int):
raise AssertionError(f'Unexpected event fetch for {instance_uuid} after {after_cursor} with limit {limit}')
async def fetch_workspaces(self, instance_uuid: str, workspace_uuids: tuple[str, ...]):
raise AssertionError(f'Unexpected delta fetch for {instance_uuid}: {workspace_uuids!r}')
def _get_script_head() -> str:
"""Resolve the current Alembic head revision from the script directory.
@@ -648,6 +665,10 @@ class TestPostgreSQLTenantRuntime:
instance_uuid = 'cloud-runtime-persistence-test'
workspace_a = '10000000-0000-0000-0000-000000000001'
workspace_b = '20000000-0000-0000-0000-000000000002'
workspace_other = '30000000-0000-0000-0000-000000000003'
workspace_local = '40000000-0000-0000-0000-000000000004'
directory_account = '50000000-0000-0000-0000-000000000005'
workspace_projected = '60000000-0000-0000-0000-000000000006'
role_suffix = uuid.uuid4().hex[:12]
runtime_role = f'lb_runtime_{role_suffix}'
bypass_role = f'lb_bypass_{role_suffix}'
@@ -730,17 +751,22 @@ class TestPostgreSQLTenantRuntime:
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')):
for workspace_uuid, slug, target_instance, source in (
(workspace_a, 'workspace-a', instance_uuid, 'cloud_projection'),
(workspace_b, 'workspace-b', instance_uuid, 'cloud_projection'),
(workspace_other, 'workspace-other', 'other-instance', 'cloud_projection'),
(workspace_local, 'workspace-local', instance_uuid, 'local'),
):
async with TenantUnitOfWork(release_engine, workspace_uuid) as uow:
await uow.execute(
sa.insert(Workspace).values(
uuid=workspace_uuid,
instance_uuid=instance_uuid,
instance_uuid=target_instance,
name=slug,
slug=slug,
type='team',
status='active',
source='cloud_projection',
source=source,
projection_revision=0,
)
)
@@ -751,6 +777,30 @@ class TestPostgreSQLTenantRuntime:
value=slug,
)
)
async with release_engine.begin() as conn:
await conn.execute(
sa.insert(User).values(
uuid=directory_account,
user='directory@example.com',
normalized_email='directory@example.com',
password='closed-directory',
account_type='space',
status='active',
source='cloud_projection',
projection_revision=1,
)
)
async with TenantUnitOfWork(release_engine, workspace_local) as uow:
uow.session.add(
WorkspaceMembership(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_local,
account_uuid=directory_account,
role='owner',
status='active',
projection_revision=1,
)
)
await create_role(runtime_role)
runtime_engine = create_async_engine(role_url(runtime_role), pool_size=1, max_overflow=0)
@@ -844,6 +894,161 @@ class TestPostgreSQLTenantRuntime:
cloud_manager.create_tables.assert_not_awaited()
cloud_manager._run_alembic_migrations.assert_not_awaited()
directory_fingerprint = hashlib.sha256(b'directory-snapshot').hexdigest()
async with cloud_manager.directory_projection_uow(instance_uuid) as directory:
assert set((await directory.session.scalars(sa.select(Workspace.uuid))).all()) == {
workspace_a,
workspace_b,
}
directory.session.add(
Workspace(
uuid=workspace_projected,
instance_uuid=instance_uuid,
name='workspace-projected',
slug='workspace-projected',
type='team',
status='active',
source='cloud_projection',
projection_revision=1,
)
)
await directory.session.flush()
directory.session.add(
WorkspaceMembership(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_projected,
account_uuid=directory_account,
role='owner',
status='active',
projection_revision=1,
)
)
directory.session.add(
WorkspaceExecutionState(
workspace_uuid=workspace_projected,
instance_uuid=instance_uuid,
active_generation=1,
state='active',
write_fenced=False,
source='cloud',
desired_state_revision=1,
)
)
directory.session.add(
DirectoryProjectionState(
instance_uuid=instance_uuid,
cursor=1,
snapshot_fingerprint=directory_fingerprint,
last_applied_at=datetime.datetime.now(datetime.UTC),
)
)
directory.session.add(
DirectoryProjectionInbox(
instance_uuid=instance_uuid,
event_uuid=str(uuid.uuid4()),
cursor=1,
event_type='workspace.changed',
revision=1,
fingerprint=directory_fingerprint,
)
)
await directory.session.flush()
assert await directory.session.scalar(sa.select(sa.func.count()).select_from(WorkspaceMembership)) == 1
assert (await directory.session.execute(sa.select(WorkspaceMetadata))).all() == []
async with cloud_manager.tenant_uow(workspace_projected) as tenant:
workspace_update = await tenant.session.execute(
sa.update(Workspace)
.where(Workspace.uuid == workspace_projected)
.values(name='tenant-must-not-update-cloud')
)
membership_update = await tenant.session.execute(
sa.update(WorkspaceMembership)
.where(WorkspaceMembership.workspace_uuid == workspace_projected)
.values(role='admin')
)
execution_update = await tenant.session.execute(
sa.update(WorkspaceExecutionState)
.where(WorkspaceExecutionState.workspace_uuid == workspace_projected)
.values(write_fenced=True)
)
assert workspace_update.rowcount == 0
assert membership_update.rowcount == 0
assert execution_update.rowcount == 0
async with cloud_manager.tenant_uow(workspace_local) as tenant:
local_update = await tenant.session.execute(
sa.update(Workspace).where(Workspace.uuid == workspace_local).values(name='tenant-can-update-local')
)
assert local_update.rowcount == 1
async with cloud_manager.directory_projection_uow(instance_uuid) as directory:
local_update = await directory.session.execute(
sa.update(Workspace)
.where(Workspace.uuid == workspace_local)
.values(name='directory-must-not-update-local')
)
assert local_update.rowcount == 0
projection_application = SimpleNamespace(
persistence_mgr=cloud_manager,
logger=logging.getLogger('postgres-directory-projection-concurrency'),
)
first_projection = DirectoryProjectionService(
projection_application,
_NoopDirectoryProjectionProvider(),
instance_uuid,
)
second_projection = DirectoryProjectionService(
projection_application,
_NoopDirectoryProjectionProvider(),
instance_uuid,
)
concurrent_event = DirectoryEvent(
cursor=2,
uuid='70000000-0000-0000-0000-000000000007',
aggregate_uuid=workspace_projected,
event_type='entitlement.changed',
revision=2,
payload={
'workspace_uuid': workspace_projected,
'entitlement_revision': 2,
},
created_at=datetime.datetime.now(datetime.UTC),
)
concurrent_batch = DirectoryEventBatch(
instance_uuid=instance_uuid,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[concurrent_event],
)
await asyncio.gather(
first_projection.apply_event_batch(concurrent_batch),
second_projection.apply_event_batch(concurrent_batch),
)
async with cloud_manager.directory_projection_uow(instance_uuid) as directory:
projection_state = await directory.session.get(DirectoryProjectionState, instance_uuid)
concurrent_receipts = (
await directory.session.scalars(
sa.select(DirectoryProjectionInbox).where(
DirectoryProjectionInbox.event_uuid == concurrent_event.uuid
)
)
).all()
assert projection_state.cursor == 2
assert len(concurrent_receipts) == 1
assert concurrent_receipts[0].applied_at is not None
async with cloud_manager.tenant_uow(workspace_a) as tenant:
assert (await tenant.session.execute(sa.select(DirectoryProjectionState))).all() == []
assert (await tenant.session.execute(sa.select(DirectoryProjectionInbox))).all() == []
async with cloud_manager.instance_discovery_uow(instance_uuid) as discovery:
assert (await discovery.session.execute(sa.select(DirectoryProjectionState))).all() == []
update_result = await discovery.session.execute(sa.update(DirectoryProjectionState).values(cursor=2))
assert update_result.rowcount == 0
with pytest.raises(TransactionRollbackOnlyError, match='after-commit work was cancelled'):
async with cloud_manager.tenant_uow(workspace_a):
duplicate_statement = sa.insert(WorkspaceMetadata).values(
tests/unit_tests/api/service/test_user_service.py
+61 -2
View File
@@ -19,8 +19,11 @@ import datetime
from unittest.mock import AsyncMock, Mock
from types import SimpleNamespace
from langbot.pkg.api.http.service.user import UserService
from langbot.pkg.entity.persistence.user import AccountStatus, User
from langbot.pkg.api.http.service.user import (
ControlPlaneDirectoryRequiredError,
UserService,
)
from langbot.pkg.entity.persistence.user import AccountSource, AccountStatus, User
from langbot.pkg.entity.errors.account import AccountEmailMismatchError
@@ -563,6 +566,62 @@ class TestUserServiceCreateOrUpdateSpaceUser:
ap.persistence_mgr.execute_async.assert_called()
assert updated_user.space_account_uuid == 'existing-space-uuid'
async def test_cloud_login_updates_only_the_projected_space_account(self):
projected = SimpleNamespace(
uuid='projected-space-uuid',
user='Cloud Owner',
normalized_email='owner@example.com',
password='',
account_type='space',
status=AccountStatus.ACTIVE.value,
source=AccountSource.CLOUD_PROJECTION.value,
projection_revision=7,
space_account_uuid='projected-space-uuid',
)
persistence = SimpleNamespace(execute_async=AsyncMock())
ap = SimpleNamespace(
persistence_mgr=persistence,
workspace_service=SimpleNamespace(policy=SimpleNamespace(multi_workspace_enabled=True)),
)
service = UserService(ap)
service.get_user_by_space_account_uuid = AsyncMock(side_effect=[projected, projected])
result = await service.create_or_update_space_user(
space_account_uuid='projected-space-uuid',
email='OWNER@example.com',
access_token='access-token',
refresh_token='refresh-token',
api_key='api-key',
expires_in=3600,
)
assert result is projected
persistence.execute_async.assert_awaited_once()
async def test_cloud_login_never_creates_an_unprojected_account(self):
persistence = SimpleNamespace(execute_async=AsyncMock())
ap = SimpleNamespace(
persistence_mgr=persistence,
workspace_service=SimpleNamespace(policy=SimpleNamespace(multi_workspace_enabled=True)),
)
service = UserService(ap)
service.get_user_by_space_account_uuid = AsyncMock(return_value=None)
with pytest.raises(
ControlPlaneDirectoryRequiredError,
match='verified Cloud directory',
):
await service.create_or_update_space_user(
space_account_uuid='unknown-space-uuid',
email='unknown@example.com',
access_token='access-token',
refresh_token='refresh-token',
api_key='api-key',
expires_in=3600,
)
persistence.execute_async.assert_not_awaited()
async def test_create_or_update_new_space_user_first_init(self):
"""Creates new Space user on first initialization."""
# Setup
tests/unit_tests/box/test_box_connector.py
+136 -2
View File
@@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.box import connector as connector_module
from langbot_plugin.box.client import ActionRPCBoxClient
from langbot_plugin.box.errors import BoxRuntimeUnavailableError
from langbot_plugin.box.security import (
@@ -121,6 +122,139 @@ def test_box_runtime_connector_dispose_terminates_subprocess(monkeypatch: pytest
assert connector._ctrl_task is None
@pytest.mark.asyncio
async def test_box_runtime_connector_cleans_partial_transport_on_connect_failure(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
connector = BoxRuntimeConnector(make_app(Mock()))
connector._start_local_stdio = AsyncMock(side_effect=RuntimeError('bind failed'))
connector._stop_transport = AsyncMock()
connector._close_managed_subprocess = AsyncMock()
with pytest.raises(RuntimeError, match='bind failed'):
await connector.initialize()
assert connector._stop_transport.await_count == 2
connector._close_managed_subprocess.assert_awaited_once()
@pytest.mark.asyncio
async def test_box_runtime_connector_starts_heartbeat_after_reconnect(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
connector = BoxRuntimeConnector(make_app(Mock()))
connector._start_local_stdio = AsyncMock(side_effect=[RuntimeError('bind failed'), None])
connector._stop_transport = AsyncMock()
connector._close_managed_subprocess = AsyncMock()
with pytest.raises(RuntimeError, match='bind failed'):
await connector.initialize()
assert connector._heartbeat_task is None
await connector.reconnect()
assert connector._heartbeat_task is not None
assert not connector._heartbeat_task.done()
await connector.aclose()
@pytest.mark.asyncio
async def test_box_stdio_connection_does_not_capture_unconsumed_stderr(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
created = {}
class FakeHandler:
def __init__(self, connection):
self.release = asyncio.Event()
async def call_action(self, action, data):
return None
async def run(self):
await self.release.wait()
async def close(self):
self.release.set()
class FakeController:
def __init__(self, **kwargs):
created.update(kwargs)
self.process = SimpleNamespace(returncode=0)
async def run(self, callback):
await callback(object())
async def close(self):
return None
monkeypatch.setattr(connector_module, 'Handler', FakeHandler)
monkeypatch.setattr(
'langbot_plugin.runtime.io.controllers.stdio.client.StdioClientController',
FakeController,
)
connector = BoxRuntimeConnector(make_app(Mock()))
await connector.initialize()
assert created['capture_stderr'] is False
assert connector._handler is not None
await connector.aclose()
@pytest.mark.asyncio
async def test_box_disconnect_notifies_once_and_clears_handler(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
disconnect = AsyncMock()
class FakeHandler:
def __init__(self, connection):
pass
async def call_action(self, action, data):
return None
async def run(self):
return None
async def close(self):
return None
class FakeController:
def __init__(self, **kwargs):
self.process = SimpleNamespace(returncode=0)
async def run(self, callback):
await callback(object())
async def close(self):
return None
monkeypatch.setattr(connector_module, 'Handler', FakeHandler)
monkeypatch.setattr(
'langbot_plugin.runtime.io.controllers.stdio.client.StdioClientController',
FakeController,
)
connector = BoxRuntimeConnector(make_app(Mock()), runtime_disconnect_callback=disconnect)
await connector.initialize()
await asyncio.sleep(0)
disconnect.assert_awaited_once_with(connector)
assert connector._handler is None
await connector.aclose()
def test_box_runtime_connector_builds_host_control_headers(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv(BOX_CONTROL_TOKEN_ENV, _CONTROL_TOKEN)
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
@@ -200,7 +334,7 @@ async def test_local_stdio_injects_generated_token_and_trusted_instance(
)
connector = BoxRuntimeConnector(make_app(Mock()))
def fake_callback(_transport_name, connected, _connect_error):
def fake_callback(_transport_name, connected, _connect_error, _generation):
async def callback(_connection):
connected.set()
@@ -231,7 +365,7 @@ async def test_websocket_controller_receives_control_headers(monkeypatch: pytest
)
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
def fake_callback(_transport_name, connected, _connect_error):
def fake_callback(_transport_name, connected, _connect_error, _generation):
async def callback(_connection):
connected.set()
tests/unit_tests/box/test_box_service.py
+91
View File
@@ -253,6 +253,46 @@ async def test_box_service_without_explicit_client_initializes_internal_connecto
connector.initialize.assert_awaited_once()
@pytest.mark.asyncio
async def test_cloud_initialize_validation_failure_closes_connector_and_cancels_reconnect(
monkeypatch: pytest.MonkeyPatch,
):
logger = Mock()
app = make_app(logger)
app.deployment = SimpleNamespace(multi_workspace_enabled=True)
app.instance_config.data['box'].update(
{
'backend': 'nsjail',
'admission': {'required': True, 'workspace_quota_mb': 32},
}
)
connector = Mock()
connector.client = Mock(spec=BoxRuntimeClient)
connector.initialize = AsyncMock()
connector.aclose = AsyncMock()
connector.runtime_disconnect_callback = Mock()
monkeypatch.setattr('langbot.pkg.box.service.BoxRuntimeConnector', Mock(return_value=connector))
service = BoxService(app)
service._ensure_default_workspace = Mock()
readiness_error = BoxValidationError('Cloud Box nsjail isolation readiness failed')
service._verify_cloud_runtime = AsyncMock(side_effect=readiness_error)
reconnect_task = asyncio.create_task(asyncio.Event().wait())
service._reconnect_task = reconnect_task
service._reconnecting = True
with pytest.raises(BoxValidationError) as exc_info:
await service.initialize()
assert exc_info.value is readiness_error
assert service.available is False
assert service._closing is True
assert service._reconnecting is False
assert service._reconnect_task is None
assert reconnect_task.cancelled()
assert connector.runtime_disconnect_callback is None
connector.aclose.assert_awaited_once()
class TestSharesFilesystemWithBox:
"""``shares_filesystem_with_box`` must reflect the real LangBot<->Box
filesystem topology, which is derived from the connector transport:
@@ -1788,6 +1828,57 @@ class TestBoxDisabledByConfig:
assert service._reconnecting is False
@pytest.mark.asyncio
async def test_disconnect_callback_does_not_schedule_on_closing_event_loop(
monkeypatch: pytest.MonkeyPatch,
):
service = BoxService(make_app(Mock()), client=Mock(spec=BoxRuntimeClient))
closed_loop = Mock()
closed_loop.is_closed.return_value = True
monkeypatch.setattr('langbot.pkg.box.service.asyncio.get_running_loop', Mock(return_value=closed_loop))
await service._on_runtime_disconnect(connector=Mock())
closed_loop.create_task.assert_not_called()
assert service._reconnect_task is None
assert service._reconnecting is False
@pytest.mark.asyncio
async def test_disconnect_callback_closes_reconnect_coroutine_when_task_creation_races_with_loop_close(
monkeypatch: pytest.MonkeyPatch,
):
service = BoxService(make_app(Mock()), client=Mock(spec=BoxRuntimeClient))
loop = Mock()
loop.is_closed.return_value = False
loop.create_task.side_effect = RuntimeError('event loop is closed')
monkeypatch.setattr('langbot.pkg.box.service.asyncio.get_running_loop', Mock(return_value=loop))
async def reconnect():
await asyncio.Event().wait()
reconnect_coroutine = reconnect()
service._reconnect_loop = Mock(return_value=reconnect_coroutine)
await service._on_runtime_disconnect(connector=Mock())
loop.create_task.assert_called_once_with(reconnect_coroutine)
assert reconnect_coroutine.cr_frame is None
assert service._reconnect_task is None
assert service._reconnecting is False
def test_disconnect_callback_does_not_schedule_without_running_event_loop():
service = BoxService(make_app(Mock()), client=Mock(spec=BoxRuntimeClient))
callback = service._on_runtime_disconnect(connector=Mock())
with pytest.raises(StopIteration):
callback.send(None)
assert service._reconnect_task is None
assert service._reconnecting is False
class TestBuildSkillExtraMounts:
"""Robustness of skill mount construction against a stale skill cache.
tests/unit_tests/cloud/test_bootstrap.py
+61
View File
@@ -7,6 +7,7 @@ import pytest
from langbot.pkg.cloud.bootstrap import (
CloudBootstrapError,
CloudManifestRefreshService,
CloudRuntimeUnavailableError,
DeploymentAdmissionGuard,
OpenSourceDeployment,
@@ -33,7 +34,38 @@ class _Entitlements:
)
class _Directory:
async def fetch_snapshot(self, instance_uuid: str):
del instance_uuid
raise AssertionError('not used by bootstrap contract tests')
async def fetch_events(self, instance_uuid: str, after_cursor: int, limit: int):
del instance_uuid, after_cursor, limit
raise AssertionError('not used by bootstrap contract tests')
async def fetch_workspaces(self, instance_uuid: str, workspace_uuids: tuple[str, ...]):
del instance_uuid, workspace_uuids
raise AssertionError('not used by bootstrap contract tests')
class _Manifest:
def __init__(self):
self.candidate = None
self.closed = False
async def refresh_manifest(self):
if self.candidate is None:
raise AssertionError('no refreshed Manifest was configured')
return self.candidate
async def aclose(self) -> None:
self.closed = True
class _Provider:
def __init__(self):
self.manifest_provider = _Manifest()
def bootstrap(self, *, instance_uuid: str, instance_config: dict):
del instance_config
return VerifiedCloudDeployment(
@@ -45,6 +77,8 @@ class _Provider:
capabilities=frozenset({'multi_workspace_v2'}),
tenant_isolation_version=2,
entitlement_provider=_Entitlements(),
directory_provider=_Directory(),
manifest_provider=self.manifest_provider,
verification_key_id='root-2026',
)
@@ -289,3 +323,30 @@ async def test_deployment_admission_accepts_only_monotonic_non_conflicting_renew
conflicting = dataclasses.replace(renewed, manifest_jti='different')
with pytest.raises(CloudRuntimeUnavailableError, match='conflicting'):
guard.replace(conflicting)
async def test_manifest_refresh_replaces_receipt_before_short_ttl_expires():
wall = [1_000.0]
provider = _Provider()
current = dataclasses.replace(
provider.bootstrap(instance_uuid='instance-a', instance_config={}),
expires_at=1_300,
)
guard = DeploymentAdmissionGuard('instance-a', current, wall_time=lambda: wall[0])
renewed = dataclasses.replace(
current,
manifest_jti='manifest-renewed',
manifest_generation=current.manifest_generation + 1,
expires_at=2_000,
)
provider.manifest_provider.candidate = renewed
service = CloudManifestRefreshService(
guard,
provider.manifest_provider,
SimpleNamespace(exception=lambda *_: None),
wall_time=lambda: wall[0],
)
assert service.next_refresh_delay() == 120
assert await service.refresh_once() is renewed
assert guard.deployment is renewed
tests/unit_tests/cloud/test_directory_projection.py
+817
View File
@@ -0,0 +1,817 @@
from __future__ import annotations
import datetime
import logging
from types import SimpleNamespace
import pytest
import sqlalchemy
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from langbot.pkg.cloud.directory import (
DirectoryDelta,
DirectoryEvent,
DirectoryEventBatch,
DirectoryMember,
DirectoryProjectionUnavailableError,
DirectorySnapshot,
DirectoryWorkspace,
)
from langbot.pkg.cloud.directory_projection import DirectoryProjectionService
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.cloud_directory import DirectoryProjectionInbox, DirectoryProjectionState
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.entity.persistence.workspace import (
Workspace,
WorkspaceExecutionState,
WorkspaceMembership,
)
from langbot.pkg.persistence.tenant_uow import PersistenceScope, TenantUnitOfWork
pytestmark = pytest.mark.asyncio
INSTANCE_UUID = 'instance-directory-test'
WORKSPACE_UUID = '10000000-0000-0000-0000-000000000001'
SECOND_WORKSPACE_UUID = '10000000-0000-0000-0000-000000000002'
ACCOUNT_UUID = '20000000-0000-0000-0000-000000000001'
MEMBERSHIP_UUID = '30000000-0000-0000-0000-000000000001'
SECOND_MEMBERSHIP_UUID = '30000000-0000-0000-0000-000000000002'
class _Persistence:
def __init__(self, engine):
self.engine = engine
def directory_projection_uow(self, instance_uuid: str) -> TenantUnitOfWork:
return TenantUnitOfWork(
self.engine,
scope=PersistenceScope.directory(instance_uuid),
)
def get_db_engine(self):
return self.engine
class _Provider:
def __init__(
self,
snapshots: list[DirectorySnapshot],
batches: list[DirectoryEventBatch] | None = None,
deltas: list[DirectoryDelta] | None = None,
) -> None:
self.snapshots = snapshots
self.batches = list(batches or [])
self.deltas = list(deltas or [])
self.snapshot_calls = 0
self.delta_calls = 0
self.after_cursors: list[int] = []
async def fetch_snapshot(self, instance_uuid: str) -> DirectorySnapshot:
assert instance_uuid == INSTANCE_UUID
snapshot = self.snapshots[min(self.snapshot_calls, len(self.snapshots) - 1)]
self.snapshot_calls += 1
return snapshot
async def fetch_events(
self,
instance_uuid: str,
after_cursor: int,
limit: int,
) -> DirectoryEventBatch:
assert instance_uuid == INSTANCE_UUID
assert limit == 100
self.after_cursors.append(after_cursor)
if self.batches:
return self.batches.pop(0)
return DirectoryEventBatch(
instance_uuid=instance_uuid,
after_cursor=after_cursor,
cursor=after_cursor,
high_water_cursor=after_cursor,
events=[],
)
async def fetch_workspaces(
self,
instance_uuid: str,
workspace_uuids: tuple[str, ...],
) -> DirectoryDelta:
assert instance_uuid == INSTANCE_UUID
assert workspace_uuids == (WORKSPACE_UUID,)
delta = self.deltas[min(self.delta_calls, len(self.deltas) - 1)]
self.delta_calls += 1
return delta
@pytest.fixture
async def projection_context(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "directory-projection.db"}')
async with engine.begin() as connection:
await connection.run_sync(Base.metadata.create_all)
application = SimpleNamespace(
persistence_mgr=_Persistence(engine),
logger=logging.getLogger('test-directory-projection'),
)
session_factory = async_sessionmaker(engine, expire_on_commit=False)
yield application, session_factory
await engine.dispose()
def _member(*, revision: int = 1, role: str = 'member') -> DirectoryMember:
return DirectoryMember(
membership_uuid=MEMBERSHIP_UUID,
account_uuid=ACCOUNT_UUID,
normalized_email='owner@example.com',
display_name='Workspace Owner',
account_status='active',
role=role,
membership_status='active',
projection_revision=revision,
joined_at=datetime.datetime(2026, 7, 24, tzinfo=datetime.UTC),
)
def _workspace(
*,
revision: int = 1,
status: str = 'active',
name: str = 'Personal Workspace',
role: str = 'member',
execution_generation: int = 1,
) -> DirectoryWorkspace:
return DirectoryWorkspace(
uuid=WORKSPACE_UUID,
name=name,
slug='personal-workspace',
type='personal',
status=status,
created_by_account_uuid=ACCOUNT_UUID,
projection_revision=revision,
execution_generation=execution_generation,
members=[_member(revision=revision, role=role)],
)
def _snapshot(
cursor: int,
*,
workspaces: list[DirectoryWorkspace] | None = None,
) -> DirectorySnapshot:
return DirectorySnapshot(
instance_uuid=INSTANCE_UUID,
cursor=cursor,
generated_at=datetime.datetime(2026, 7, 24, 12, cursor % 60, tzinfo=datetime.UTC),
workspaces=[_workspace(revision=max(cursor, 1))] if workspaces is None else workspaces,
)
def _delta(
*,
workspaces: list[DirectoryWorkspace] | None = None,
requested_workspace_uuids: tuple[str, ...] = (WORKSPACE_UUID,),
) -> DirectoryDelta:
return DirectoryDelta(
instance_uuid=INSTANCE_UUID,
requested_workspace_uuids=requested_workspace_uuids,
generated_at=datetime.datetime(2026, 7, 24, 12, 30, tzinfo=datetime.UTC),
workspaces=[_workspace(revision=2)] if workspaces is None else workspaces,
)
async def test_initial_snapshot_projects_core_owned_rows(projection_context):
application, session_factory = projection_context
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
)
await service.initialize()
service.require_ready()
async with session_factory() as session:
account = await session.scalar(sqlalchemy.select(User))
workspace = await session.get(Workspace, WORKSPACE_UUID)
membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
execution = await session.get(WorkspaceExecutionState, WORKSPACE_UUID)
state = await session.get(DirectoryProjectionState, INSTANCE_UUID)
assert account is not None
assert account.uuid == ACCOUNT_UUID
assert account.source == 'cloud_projection'
assert account.account_type == 'space'
assert account.password == ''
assert workspace is not None
assert workspace.source == 'cloud_projection'
assert membership is not None
assert membership.role == 'developer'
assert execution is not None
assert execution.source == 'cloud'
assert execution.write_fenced is False
assert state is not None
assert state.cursor == 1
assert state.snapshot_coverage_cursor == 1
async def test_same_cursor_equivocation_and_rollback_fail_closed(projection_context):
application, _session_factory = projection_context
service = DirectoryProjectionService(
application,
_Provider([_snapshot(2)]),
INSTANCE_UUID,
)
await service.initialize()
conflicting = _snapshot(
2,
workspaces=[_workspace(revision=2, name='Conflicting Name')],
)
with pytest.raises(DirectoryProjectionUnavailableError, match='conflicting contents'):
await service.apply_snapshot(conflicting)
with pytest.raises(DirectoryProjectionUnavailableError, match='rolled back'):
await service.apply_snapshot(_snapshot(1))
async def test_execution_generation_cannot_change_without_directory_revision(projection_context):
application, _session_factory = projection_context
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
)
await service.initialize()
conflicting = _snapshot(
2,
workspaces=[_workspace(revision=1, execution_generation=2)],
)
with pytest.raises(DirectoryProjectionUnavailableError, match='execution state has conflicting contents'):
await service.apply_snapshot(conflicting)
async def test_event_limit_matches_single_delta_request_limit(projection_context):
application, _session_factory = projection_context
with pytest.raises(ValueError, match='between 1 and 100'):
DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
event_limit=101,
)
async def test_archived_and_absent_workspaces_are_execution_fenced(projection_context):
application, session_factory = projection_context
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
)
await service.initialize()
await service.apply_snapshot(
_snapshot(
2,
workspaces=[_workspace(revision=2, status='archived')],
)
)
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
execution = await session.get(WorkspaceExecutionState, WORKSPACE_UUID)
membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
assert workspace.status == 'archived'
assert execution.state == 'inactive'
assert execution.write_fenced is True
assert membership.status == 'removed'
await service.apply_snapshot(_snapshot(3, workspaces=[]))
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
execution = await session.get(WorkspaceExecutionState, WORKSPACE_UUID)
assert workspace.status == 'archived'
assert execution.state == 'inactive'
assert execution.write_fenced is True
async def test_event_poll_fetches_workspace_delta_and_records_receipt(projection_context):
application, session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000001',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[event],
)
provider = _Provider(
[_snapshot(1)],
[batch],
[_delta(workspaces=[_workspace(revision=2, name='Updated Workspace')])],
)
service = DirectoryProjectionService(application, provider, INSTANCE_UUID)
await service.initialize()
await service.sync_once()
async with session_factory() as session:
account = await session.scalar(sqlalchemy.select(User).where(User.uuid == ACCOUNT_UUID))
workspace = await session.get(Workspace, WORKSPACE_UUID)
state = await session.get(DirectoryProjectionState, INSTANCE_UUID)
inbox = await session.scalar(sqlalchemy.select(DirectoryProjectionInbox))
assert account.projection_revision == 1
assert workspace.name == 'Updated Workspace'
assert state.cursor == 2
assert inbox.event_uuid == event.uuid
assert inbox.applied_at is not None
assert provider.snapshot_calls == 1
assert provider.delta_calls == 1
async def test_directory_delta_does_not_skip_unfetched_event_cursors(projection_context):
application, session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000003',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=3,
events=[event],
)
provider = _Provider(
[_snapshot(1)],
[batch],
[_delta(workspaces=[_workspace(revision=3, name='Latest Workspace')])],
)
service = DirectoryProjectionService(application, provider, INSTANCE_UUID)
await service.initialize()
await service.sync_once()
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
state = await session.get(DirectoryProjectionState, INSTANCE_UUID)
assert workspace.name == 'Latest Workspace'
assert state.cursor == 2
async def test_entitlement_event_advances_cursor_without_refetching_directory(projection_context):
application, session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000002',
aggregate_uuid=WORKSPACE_UUID,
event_type='entitlement.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'entitlement_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[event],
)
provider = _Provider([_snapshot(1)], [batch])
service = DirectoryProjectionService(application, provider, INSTANCE_UUID)
await service.initialize()
await service.sync_once()
await service.apply_snapshot(_snapshot(2, workspaces=[_workspace(revision=1)]))
async with session_factory() as session:
state = await session.get(DirectoryProjectionState, INSTANCE_UUID)
inbox = await session.scalar(sqlalchemy.select(DirectoryProjectionInbox))
assert provider.snapshot_calls == 1
assert state.cursor == 2
assert inbox.event_uuid == event.uuid
assert inbox.applied_at is not None
async def test_missing_workspace_in_delta_fences_only_requested_workspace(projection_context):
application, session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000004',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[event],
)
provider = _Provider([_snapshot(1)], [batch], [_delta(workspaces=[])])
service = DirectoryProjectionService(application, provider, INSTANCE_UUID)
await service.initialize()
await service.sync_once()
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
execution = await session.get(WorkspaceExecutionState, WORKSPACE_UUID)
membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
assert workspace.status == 'archived'
assert workspace.projection_revision == 2
assert execution.state == 'inactive'
assert execution.write_fenced is True
assert execution.desired_state_revision == 2
assert membership.status == 'removed'
assert membership.projection_revision == 2
stale_same_cursor_snapshot = _snapshot(
2,
workspaces=[_workspace(revision=2, status='active')],
)
with pytest.raises(DirectoryProjectionUnavailableError, match='conflicting contents'):
await service.apply_snapshot(stale_same_cursor_snapshot)
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
execution = await session.get(WorkspaceExecutionState, WORKSPACE_UUID)
assert workspace.status == 'archived'
assert execution.write_fenced is True
async def test_each_replica_consumes_events_with_its_own_cursor(projection_context):
application, session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000005',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[event],
)
updated = _delta(workspaces=[_workspace(revision=2, name='Replica-safe Workspace')])
first_provider = _Provider([_snapshot(1)], [batch], [updated])
second_provider = _Provider([_snapshot(1)], [batch], [updated])
first = DirectoryProjectionService(application, first_provider, INSTANCE_UUID)
second = DirectoryProjectionService(application, second_provider, INSTANCE_UUID)
await first.initialize()
await second.initialize()
await first.sync_once()
await second.sync_once()
await second.sync_once()
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
state = await session.get(DirectoryProjectionState, INSTANCE_UUID)
inbox_rows = (await session.scalars(sqlalchemy.select(DirectoryProjectionInbox))).all()
assert workspace.name == 'Replica-safe Workspace'
assert state.cursor == 2
assert len(inbox_rows) == 1
assert first_provider.after_cursors == [1]
assert second_provider.after_cursors == [1, 2]
async def test_snapshot_coverage_allows_lagging_replica_to_replay_receipts(projection_context):
application, session_factory = projection_context
event_two = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000008',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
event_three = DirectoryEvent(
cursor=3,
uuid='40000000-0000-0000-0000-000000000009',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=3,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 3},
created_at=datetime.datetime(2026, 7, 24, 12, 3, tzinfo=datetime.UTC),
)
batch_two = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=3,
events=[event_two],
)
batch_three = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=2,
cursor=3,
high_water_cursor=3,
events=[event_three],
)
lagging_provider = _Provider(
[_snapshot(1)],
[batch_two, batch_three],
[
_delta(workspaces=[_workspace(revision=2, name='Intermediate Workspace')]),
_delta(workspaces=[_workspace(revision=3, name='Current Workspace')]),
],
)
lagging = DirectoryProjectionService(application, lagging_provider, INSTANCE_UUID)
leading = DirectoryProjectionService(
application,
_Provider([_snapshot(3, workspaces=[_workspace(revision=3, name='Current Workspace')])]),
INSTANCE_UUID,
)
await lagging.initialize()
await leading.initialize()
await lagging.sync_once()
await lagging.sync_once()
lagging.require_ready()
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
state = await session.get(DirectoryProjectionState, INSTANCE_UUID)
inbox_rows = (
await session.scalars(sqlalchemy.select(DirectoryProjectionInbox).order_by(DirectoryProjectionInbox.cursor))
).all()
assert workspace.name == 'Current Workspace'
assert state.cursor == 3
assert state.snapshot_coverage_cursor == 3
assert [row.cursor for row in inbox_rows] == [2, 3]
assert lagging_provider.after_cursors == [1, 2]
async def test_initialize_retries_snapshot_superseded_by_another_replica(projection_context):
application, _session_factory = projection_context
leading = DirectoryProjectionService(
application,
_Provider([_snapshot(2)]),
INSTANCE_UUID,
)
await leading.initialize()
racing_provider = _Provider([_snapshot(1), _snapshot(2)])
racing = DirectoryProjectionService(application, racing_provider, INSTANCE_UUID)
await racing.initialize()
await racing.sync_once()
racing.require_ready()
assert racing_provider.snapshot_calls == 2
assert racing_provider.after_cursors == [2]
async def test_page_below_signed_high_water_does_not_renew_readiness(projection_context):
application, _session_factory = projection_context
monotonic = [10.0]
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000010',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=3,
events=[event],
)
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)], [batch], [_delta()]),
INSTANCE_UUID,
sync_interval_seconds=5,
max_staleness_seconds=60,
monotonic_time=lambda: monotonic[0],
)
await service.initialize()
await service.sync_once()
monotonic[0] = 70.0
with pytest.raises(DirectoryProjectionUnavailableError, match='stale'):
service.require_ready()
async def test_empty_page_cannot_hide_shared_projection_progress(projection_context):
application, _session_factory = projection_context
lagging = DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
)
leading = DirectoryProjectionService(
application,
_Provider([_snapshot(2)]),
INSTANCE_UUID,
)
await lagging.initialize()
await leading.initialize()
with pytest.raises(DirectoryProjectionUnavailableError, match='has not consumed'):
await lagging.sync_once()
async def test_event_payload_revision_mismatch_fails_closed(projection_context):
application, _session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000006',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 3},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[event],
)
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)], [batch], [_delta()]),
INSTANCE_UUID,
)
await service.initialize()
with pytest.raises(DirectoryProjectionUnavailableError, match='conflicting revision'):
await service.sync_once()
async def test_workspace_delta_older_than_event_fails_closed(projection_context):
application, _session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000007',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=3,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 3},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[event],
)
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)], [batch], [_delta(workspaces=[_workspace(revision=2)])]),
INSTANCE_UUID,
)
await service.initialize()
with pytest.raises(DirectoryProjectionUnavailableError, match='older'):
await service.sync_once()
async def test_stale_workspace_tombstone_revision_fails_closed(projection_context):
application, session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000011',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[event],
)
service = DirectoryProjectionService(
application,
_Provider(
[_snapshot(1, workspaces=[_workspace(revision=3)])],
[batch],
[_delta(workspaces=[])],
),
INSTANCE_UUID,
)
await service.initialize()
with pytest.raises(DirectoryProjectionUnavailableError, match='tombstone revision rolled back'):
await service.sync_once()
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
execution = await session.get(WorkspaceExecutionState, WORKSPACE_UUID)
assert workspace.status == 'active'
assert membership.status == 'active'
assert execution.write_fenced is False
async def test_conflicting_account_projection_across_workspaces_fails_closed(projection_context):
application, _session_factory = projection_context
conflicting_member = DirectoryMember(
membership_uuid=SECOND_MEMBERSHIP_UUID,
account_uuid=ACCOUNT_UUID,
normalized_email='owner@example.com',
display_name='Workspace Owner',
account_status='blocked',
role='member',
membership_status='active',
projection_revision=99,
joined_at=datetime.datetime(2026, 7, 24, tzinfo=datetime.UTC),
)
second_workspace = DirectoryWorkspace(
uuid=SECOND_WORKSPACE_UUID,
name='Second Workspace',
slug='second-workspace',
type='team',
status='active',
created_by_account_uuid=ACCOUNT_UUID,
projection_revision=99,
execution_generation=1,
members=[conflicting_member],
)
snapshot = _snapshot(
1,
workspaces=[_workspace(revision=1), second_workspace],
)
service = DirectoryProjectionService(application, _Provider([snapshot]), INSTANCE_UUID)
with pytest.raises(DirectoryProjectionUnavailableError, match='conflicting account projections'):
await service.initialize()
async def test_account_projection_revision_advances_when_account_fields_change(projection_context):
application, session_factory = projection_context
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
)
await service.initialize()
blocked_member = _member(revision=2).model_copy(update={'account_status': 'blocked'})
blocked_workspace = _workspace(revision=2).model_copy(update={'members': (blocked_member,)})
await service.apply_snapshot(_snapshot(2, workspaces=[blocked_workspace]))
async with session_factory() as session:
account = await session.scalar(sqlalchemy.select(User).where(User.uuid == ACCOUNT_UUID))
assert account.status == 'disabled'
assert account.projection_revision == 2
async def test_readiness_expires_on_monotonic_staleness(projection_context):
application, _session_factory = projection_context
monotonic = [10.0]
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
sync_interval_seconds=5,
max_staleness_seconds=60,
monotonic_time=lambda: monotonic[0],
)
await service.initialize()
service.require_ready()
monotonic[0] = 70.0
with pytest.raises(DirectoryProjectionUnavailableError, match='stale'):
service.require_ready()
async def test_snapshot_for_another_instance_is_rejected(projection_context):
application, _session_factory = projection_context
other = _snapshot(1).model_copy(update={'instance_uuid': 'other-instance'})
service = DirectoryProjectionService(application, _Provider([other]), INSTANCE_UUID)
with pytest.raises(DirectoryProjectionUnavailableError, match='another LangBot instance'):
await service.initialize()
tests/unit_tests/persistence/test_tenant_uow.py
+53
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import asyncio
import contextvars
import datetime
from types import SimpleNamespace
import pytest
@@ -15,6 +16,10 @@ from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, registry, rel
from sqlalchemy.sql import quoted_name
from langbot.pkg.core.task_boundary import create_detached_task
from langbot.pkg.entity.persistence.cloud_directory import (
DirectoryProjectionInbox,
DirectoryProjectionState,
)
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
from langbot.pkg.persistence.tenant_uow import (
@@ -159,6 +164,54 @@ async def test_manager_reuses_one_session_and_rejects_cross_workspace(tmp_path)
await engine.dispose()
async def test_directory_projection_uow_is_instance_scoped_and_not_workspace_nestable(tmp_path) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "directory-projection-uow.db"}')
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
fingerprint = 'a' * 64
try:
async with engine.begin() as conn:
await conn.run_sync(DirectoryProjectionState.__table__.create)
await conn.run_sync(DirectoryProjectionInbox.__table__.create)
async with manager.directory_projection_uow('instance-a') as uow:
assert manager.current_scope() is not None
assert manager.current_scope().kind is PersistenceScopeKind.DIRECTORY_PROJECTION
assert manager.current_scope().settings == (('langbot.directory_instance_uuid', 'instance-a'),)
manager.require_current_session(PersistenceScopeKind.DIRECTORY_PROJECTION)
uow.session.add(
DirectoryProjectionState(
instance_uuid='instance-a',
cursor=1,
snapshot_fingerprint=fingerprint,
last_applied_at=datetime.datetime.now(datetime.UTC),
)
)
uow.session.add(
DirectoryProjectionInbox(
instance_uuid='instance-a',
event_uuid='00000000-0000-0000-0000-000000000001',
cursor=1,
event_type='workspace.changed',
revision=1,
fingerprint=fingerprint,
)
)
with pytest.raises(CrossScopeTransactionError, match='while directory_projection scope is active'):
async with manager.tenant_uow('workspace-a'):
pass
async with manager.tenant_scope('workspace-a'):
with pytest.raises(CrossScopeTransactionError, match='while workspace scope is active'):
async with manager.directory_projection_uow('instance-a'):
pass
with pytest.raises(ValueError, match='must not be empty'):
manager.directory_projection_uow(' ')
finally:
await engine.dispose()
async def test_manager_scoped_execute_preserves_core_row_and_scalar_contract(tmp_path) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "manager-result-contract.db"}')
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
tests/unit_tests/plugin/test_connector_ping.py
+152 -1
View File
@@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.plugin import connector as connector_module
from langbot.pkg.plugin.connector import PluginRuntimeConnector, PluginRuntimeNotConnectedError
from langbot_plugin.runtime.security import (
PLUGIN_RUNTIME_CONTROL_TOKEN_ENV,
@@ -17,7 +18,22 @@ from langbot_plugin.runtime.security import (
def make_connector() -> PluginRuntimeConnector:
app = SimpleNamespace(
logger=Mock(),
instance_config=SimpleNamespace(data={'plugin': {'enable': True}, 'space': {'url': ''}}),
instance_config=SimpleNamespace(
data={
'plugin': {
'enable': True,
'worker': {
'max_cpus': 1.0,
'max_memory_mb': 512,
'max_pids': 128,
'max_open_files': 256,
'max_file_size_mb': 512,
'require_hard_limits': False,
},
},
'space': {'url': ''},
}
),
)
return PluginRuntimeConnector(app, AsyncMock())
@@ -41,6 +57,141 @@ async def test_ping_plugin_runtime_delegates_to_connected_handler():
connector.handler.ping.assert_awaited_once()
@pytest.mark.asyncio
async def test_stop_transport_tolerates_handler_callback_removing_attribute():
connector = make_connector()
class Handler:
async def close(self):
del connector.handler
connector.handler = Handler()
await connector._stop_transport()
assert not hasattr(connector, 'handler')
@pytest.mark.asyncio
async def test_stdio_runtime_connection_does_not_capture_unconsumed_stderr(
monkeypatch: pytest.MonkeyPatch,
):
connector = make_connector()
connector._prepare_connected_runtime = AsyncMock()
created = {}
monkeypatch.setattr(connector_module.constants, 'instance_id', 'instance-a')
class FakeRuntimeHandler:
def __init__(self, connection, disconnect_callback, ap):
self.release = asyncio.Event()
async def ping(self):
return None
async def set_runtime_config(self, **kwargs):
return None
async def run(self):
await self.release.wait()
async def close(self):
self.release.set()
class FakeController:
def __init__(self, **kwargs):
created.update(kwargs)
async def run(self, callback):
await callback(object())
async def close(self):
return None
monkeypatch.setattr(connector_module.platform, 'get_platform', lambda: 'linux')
monkeypatch.setattr(
connector_module.platform,
'use_websocket_to_connect_plugin_runtime',
lambda: False,
)
monkeypatch.setattr(
connector_module.handler,
'RuntimeConnectionHandler',
FakeRuntimeHandler,
)
monkeypatch.setattr(
connector_module.stdio_client_controller,
'StdioClientController',
FakeController,
)
await connector.initialize()
assert created['capture_stderr'] is False
assert connector._connected.is_set()
connector._prepare_connected_runtime.assert_awaited_once()
await connector.aclose()
@pytest.mark.asyncio
async def test_runtime_disconnect_notifies_once_and_clears_handler(
monkeypatch: pytest.MonkeyPatch,
):
disconnect = AsyncMock()
connector = PluginRuntimeConnector(make_connector().ap, disconnect)
connector._prepare_connected_runtime = AsyncMock()
monkeypatch.setattr(connector_module.constants, 'instance_id', 'instance-a')
class FakeRuntimeHandler:
def __init__(self, connection, disconnect_callback, ap):
self.disconnect_callback = disconnect_callback
async def ping(self):
return None
async def set_runtime_config(self, **kwargs):
return None
async def run(self):
await self.disconnect_callback(self)
async def close(self):
return None
class FakeController:
def __init__(self, **kwargs):
pass
async def run(self, callback):
await callback(object())
async def close(self):
return None
monkeypatch.setattr(connector_module.platform, 'get_platform', lambda: 'linux')
monkeypatch.setattr(
connector_module.platform,
'use_websocket_to_connect_plugin_runtime',
lambda: False,
)
monkeypatch.setattr(
connector_module.handler,
'RuntimeConnectionHandler',
FakeRuntimeHandler,
)
monkeypatch.setattr(
connector_module.stdio_client_controller,
'StdioClientController',
FakeController,
)
await connector.initialize()
await asyncio.sleep(0)
disconnect.assert_awaited_once_with(connector)
assert not hasattr(connector, 'handler')
await connector.aclose()
@pytest.mark.asyncio
async def test_disabled_connector_validates_workspace_without_runtime_handler():
app = SimpleNamespace(
tests/unit_tests/provider/test_localagent_no_duplicate.py
+13 -1
View File
@@ -10,6 +10,7 @@ import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.provider.message as provider_message
import langbot_plugin.api.entities.builtin.provider.session as provider_session
from langbot.pkg.api.http.context import ExecutionContext, PrincipalContext, PrincipalType
from langbot.pkg.provider.runners.localagent import LocalAgentRunner
@@ -97,7 +98,7 @@ def make_query() -> pipeline_query.Query:
adapter = AsyncMock()
adapter.is_stream_output_supported = AsyncMock(return_value=False)
return pipeline_query.Query.model_construct(
query = pipeline_query.Query.model_construct(
query_id='no-dup-query',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -124,6 +125,17 @@ def make_query() -> pipeline_query.Query:
use_llm_model_uuid='test-model-uuid',
variables={},
)
object.__setattr__(
query,
'_execution_context',
ExecutionContext(
instance_uuid='instance-test',
workspace_uuid='workspace-test',
placement_generation=1,
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
),
)
return query
def _make_app(provider) -> SimpleNamespace:
tests/unit_tests/provider/test_mcp_box_integration.py
+3 -4
View File
@@ -829,7 +829,7 @@ class TestGetRuntimeInfoDict:
assert ap.box_service.available is True
@pytest.mark.asyncio
async def test_enabled_box_timeout_does_not_exhaust_mcp_retry_budget(self, mcp_module, monkeypatch):
async def test_enabled_box_timeout_does_not_exhaust_mcp_retry_budget(self, mcp_module):
ap = _make_ap()
ap.box_service.available = False
ap.box_service.enabled = True
@@ -854,14 +854,13 @@ class TestGetRuntimeInfoDict:
raise RuntimeError('Box runtime is not available after 1 seconds')
session._lifecycle_loop = lifecycle
sleep = AsyncMock()
monkeypatch.setattr(mcp_module.asyncio, 'sleep', sleep)
session._sleep_with_execution_fence = AsyncMock()
await session._lifecycle_loop_with_retry()
assert attempts == 2
assert session.retry_count == 0
sleep.assert_awaited_once_with(1)
session._sleep_with_execution_fence.assert_awaited_once_with(1)
@pytest.mark.asyncio
async def test_disabled_box_still_stops_mcp_retry_loop(self, mcp_module):
tests/unit_tests/provider/test_mcp_remote_transport.py
+16 -2
View File
@@ -13,7 +13,7 @@ from aiohttp import web
from mcp import types as mcp_types
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.provider.tools.loaders.mcp import RuntimeMCPSession
from langbot.pkg.provider.tools.loaders.mcp import MCPToolCallTimeoutError, RuntimeMCPSession
TEST_EXECUTION_CONTEXT = ExecutionContext(
@@ -159,7 +159,21 @@ def _session(
timeout: float = 2,
tool_call_timeout_sec: float = 300,
) -> RuntimeMCPSession:
app = cast(Any, SimpleNamespace(logger=Mock()))
app = cast(
Any,
SimpleNamespace(
logger=Mock(),
workspace_service=SimpleNamespace(
get_execution_binding=AsyncMock(
return_value=SimpleNamespace(
instance_uuid=TEST_EXECUTION_CONTEXT.instance_uuid,
workspace_uuid=TEST_EXECUTION_CONTEXT.workspace_uuid,
placement_generation=TEST_EXECUTION_CONTEXT.placement_generation,
)
)
),
),
)
return RuntimeMCPSession(
'remote-transport-test',
{
tests/unit_tests/provider/test_mcp_resources.py
+49 -1
View File
@@ -124,6 +124,7 @@ async def test_invoke_mcp_tool_uses_configurable_request_timeout():
},
True,
_app(),
TEST_EXECUTION_CONTEXT,
)
session.session = SimpleNamespace(call_tool=AsyncMock(return_value=_tool_result()))
@@ -148,6 +149,7 @@ async def test_invoke_mcp_tool_zero_timeout_disables_request_deadline():
},
True,
_app(),
TEST_EXECUTION_CONTEXT,
)
session.session = SimpleNamespace(call_tool=AsyncMock(return_value=_tool_result()))
@@ -171,6 +173,7 @@ async def test_invoke_mcp_tool_timeout_is_not_retried_and_session_remains_usable
},
True,
_app(),
TEST_EXECUTION_CONTEXT,
)
timeout = McpError(
mcp_types.ErrorData(
@@ -205,6 +208,7 @@ def test_invalid_tool_call_timeout_falls_back_to_default(invalid_timeout):
},
True,
ap,
TEST_EXECUTION_CONTEXT,
)
assert session.tool_call_timeout_sec == MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS
@@ -525,7 +529,11 @@ async def test_mcp_tool_result_is_discarded_when_generation_changes_during_call(
with pytest.raises(WorkspaceGenerationMismatchError):
await session.invoke_mcp_tool('side_effecting_tool', {})
session.session.call_tool.assert_awaited_once_with('side_effecting_tool', {})
session.session.call_tool.assert_awaited_once_with(
'side_effecting_tool',
{},
read_timeout_seconds=timedelta(seconds=MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS),
)
@pytest.mark.asyncio
@@ -567,3 +575,43 @@ async def test_mcp_idle_lifecycle_stops_without_retry_after_generation_bump():
assert session.status == MCPSessionStatus.ERROR
assert session.error_message == 'Workspace execution binding is stale'
assert session._shutdown_event.is_set()
@pytest.mark.asyncio
async def test_mcp_loader_shutdown_cancels_startup_tasks_and_closes_sessions_concurrently():
loader = MCPLoader(_app())
hosted_cancelled = asyncio.Event()
async def pending_host():
try:
await asyncio.Event().wait()
finally:
hosted_cancelled.set()
hosted_task = asyncio.create_task(pending_host())
await asyncio.sleep(0)
loader._hosted_mcp_tasks = [hosted_task]
started: set[str] = set()
all_started = asyncio.Event()
class Session:
def __init__(self, name: str):
self.name = name
self.server_name = name
async def shutdown(self):
started.add(self.name)
if len(started) == 2:
all_started.set()
await all_started.wait()
loader.sessions = {'one': Session('one'), 'two': Session('two')}
await asyncio.wait_for(loader.shutdown(), timeout=1)
assert hosted_cancelled.is_set()
assert hosted_task.cancelled()
assert started == {'one', 'two'}
assert loader._hosted_mcp_tasks == []
assert loader.sessions == {}
tests/unit_tests/provider/test_skill_tools.py
+1 -1
View File
@@ -508,7 +508,7 @@ class TestSkillToolLoader:
ap.skill_mgr = SimpleNamespace(skills={'demo': _make_skill_data(name='demo')})
ap.box_service = SimpleNamespace(
available=False,
get_status=AsyncMock(return_value={'backend': {'available': True}}),
get_backend_status=AsyncMock(return_value={'backend': {'available': True}}),
)
loader = SkillToolLoader(ap)
tests/unit_tests/provider/test_tool_manager_native.py
+15
View File
@@ -206,6 +206,21 @@ async def test_native_tool_loader_exposes_all_tools_when_box_available():
assert await loader.has_tool(tool_name) is True
@pytest.mark.asyncio
async def test_native_tool_loader_refreshes_after_box_recovers():
box_service = SimpleNamespace(
available=False,
get_backend_status=AsyncMock(return_value={'backend': {'available': True}}),
)
loader = NativeToolLoader(SimpleNamespace(box_service=box_service, logger=Mock()))
await loader.initialize()
assert await loader.get_tools() == []
box_service.available = True
assert [tool.name for tool in await loader.get_tools()] == ['exec', 'read', 'write', 'edit', 'glob', 'grep']
@pytest.mark.asyncio
async def test_native_tool_loader_rechecks_admission_at_the_final_invoke_boundary():
box_service = SimpleNamespace(