mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 04:40:57 +00:00
feat(tenancy): add Workspace multi-tenant foundation (#2353)
* Document multi-tenant workspace architecture * Add OSS and commercial workspace boundaries * docs: redesign multi-tenant workspace architecture * feat(tenancy): implement workspace isolation * docs(tenancy): record verification evidence * docs(tenancy): revise single-instance SaaS topology * docs(tenancy): refine architecture options * docs: finalize cloud v2 multi-tenant decisions * feat(tenancy): establish cloud isolation foundations * feat(tenancy): harden shared cloud runtime boundaries * docs(tenancy): record final isolation verification * fix(tenancy): close isolation and permission gaps * docs(tenancy): record final isolation verification * feat(tenancy): connect cloud workspace control plane * fix(build): install git for pinned SDK * docs(cloud): update control plane verification * chore: update multi-tenant SDK pin * fix(cloud): skip legacy model sync during startup * test(cloud): preserve minimal model manager fixtures * fix(cloud): preserve authenticated account context * fix(cloud): reuse authenticated account for user info * feat(cloud): complete Workspace settings navigation * test(web): cover Workspace dropdown menu * feat(web): place workspace controls in sidebar * refactor(web): streamline workspace controls * style(web): format workspace layout test * fix(cloud): surface runtime and workspace plan status * fix(plugin): keep runtime identity stable across restarts * fix(ui): widen and center workspace switcher * fix(ui): hide roles from workspace switcher * fix(ui): align workspace switcher with sidebar entries * feat(workspace): add in-product collaboration and direct Cloud launch * style: format collaboration changes * fix(workspace): bind collaboration APIs to tenant UoW * fix(cloud): preserve Core-owned collaboration state * test(cloud): require Space identity for invite registration * feat(cloud): complete secure invitation experience * style(web): format invitation flows * fix(cloud): recover box runtime without unscoped skill reload * feat(oss): enforce invitation account and owner billing flows * style: format OSS account service * test(oss): cover invitation logout handoff * fix(oss): resolve workspace owner in scoped session * feat(cloud): harden multi-tenant runtime resources * fix(cloud): bound runtime restart storms * fix(cloud): eliminate periodic runtime CPU spikes * fix(cloud): enforce instance capacity ceilings * fix(cloud): scope public login capability discovery * fix(cloud): bound tenant maintenance and monitoring work * fix(runtime): bound tenant resource amplification * fix(deps): pin green multi-tenant plugin SDK * fix(cloud): handle unavailable skill capability * fix(security): require authentication for image file endpoint (H-2) - Changed /api/v1/files/image from AuthType.NONE to USER_TOKEN_OR_API_KEY - Added Permission.RESOURCE_VIEW requirement - Prevents unauthenticated cross-tenant file access via leaked keys - Fixes HIGH severity finding from multi-tenant security review docs: add comprehensive database migration guide - Complete migration steps for OSS → multi-tenant - Backup, execution, verification procedures - Rollback scenarios and recovery plans - Performance tuning recommendations * test: add comprehensive cross-tenant isolation tests Added 7 critical test scenarios for multi-tenant boundaries: - Cross-tenant bot access prevention - Viewer role read-only enforcement - Removed member immediate access revocation - Model provider credential isolation - WebSocket message isolation - Invitation token workspace scoping - Multi-workspace context validation These tests address P0-2 coverage gaps for: - workspaces.py (membership & invitation flows) - user.py (authentication & authorization) - websocket_chat.py (real-time isolation) - plugins.py (resource access control) docs: finalize database migration guide * fix(security): resolve M-1, M-2, M-3 security findings M-1: WebSocket authorization TOCTOU race (FIXED) - Changed _revalidate_websocket_authorization to return RequestContext - Ensures validated context is used immediately without race window - Prevents removed members from sending messages during revalidation gap M-2: Model Manager cache workspace isolation (VERIFIED) - Confirmed _CacheKey already uses 4-tuple: (instance, workspace, generation, resource) - Cache is properly scoped per workspace, no cross-tenant leakage possible - No code change needed, documented as working correctly M-3: Invitation lock workspace scoping (FIXED) - Changed lock key from token_digest to workspace_uuid:token_digest - Prevents DoS where attacker locks token in Workspace A to block Workspace B - Locks now isolated per workspace All MEDIUM severity findings from security review now resolved. * fix(cloud): unblock tenant CI and enforce knowledge quotas * fix(tenancy): scope rerank model sync --------- Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
This commit is contained in:
RockChinQ
@@ -0,0 +1,105 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.workspace.invitation_delivery import (
|
||||
InvitationDeliveryResult,
|
||||
InvitationDeliveryService,
|
||||
)
|
||||
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
def _app(config: dict) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
instance_config=SimpleNamespace(data=config),
|
||||
logger=SimpleNamespace(warning=lambda *args, **kwargs: None),
|
||||
)
|
||||
|
||||
|
||||
async def test_link_only_delivery_builds_public_web_link_without_secret_capability():
|
||||
service = InvitationDeliveryService(
|
||||
_app(
|
||||
{
|
||||
'api': {
|
||||
'webui_url': 'https://public.langbot.example/',
|
||||
'webhook_prefix': 'http://internal:5300',
|
||||
}
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
assert service.build_invitation_link('lbi_secret') == (
|
||||
'https://public.langbot.example/invitations/accept#token=lbi_secret'
|
||||
)
|
||||
assert service.capability() == {'enabled': False, 'provider': None}
|
||||
result = await service.deliver_invitation(
|
||||
recipient_email='member@example.com',
|
||||
workspace_name='Workspace',
|
||||
invitation_link='https://public.langbot.example/invitations/accept#token=lbi_secret',
|
||||
)
|
||||
assert result == InvitationDeliveryResult(status='link_only', provider=None)
|
||||
|
||||
|
||||
async def test_configured_provider_failure_returns_failed_without_raising():
|
||||
service = InvitationDeliveryService(
|
||||
_app(
|
||||
{
|
||||
'workspace': {
|
||||
'invitations': {
|
||||
'email': {
|
||||
'provider': 'resend',
|
||||
'from': 'LangBot <noreply@example.com>',
|
||||
'resend': {'api_key': 'resend-secret'},
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
)
|
||||
)
|
||||
service._send_resend = AsyncMock(return_value=False)
|
||||
|
||||
assert service.capability() == {'enabled': True, 'provider': 'resend'}
|
||||
result = await service.deliver_invitation(
|
||||
recipient_email='member@example.com',
|
||||
workspace_name='Workspace',
|
||||
invitation_link='https://public.langbot.example/invitations/accept#token=lbi_secret',
|
||||
)
|
||||
|
||||
assert result == InvitationDeliveryResult(status='failed', provider='resend')
|
||||
service._send_resend.assert_awaited_once()
|
||||
|
||||
|
||||
async def test_environment_mapping_enables_provider_without_leaking_secret(monkeypatch):
|
||||
monkeypatch.setenv('WORKSPACE__INVITATIONS__PUBLIC_WEB_URL', 'https://env.langbot.example/')
|
||||
monkeypatch.setenv('WORKSPACE__INVITATIONS__EMAIL__PROVIDER', 'smtp')
|
||||
monkeypatch.setenv('WORKSPACE__INVITATIONS__EMAIL__FROM', 'LangBot <noreply@example.com>')
|
||||
monkeypatch.setenv('WORKSPACE__INVITATIONS__EMAIL__SMTP__HOST', 'smtp.example.com')
|
||||
monkeypatch.setenv('WORKSPACE__INVITATIONS__EMAIL__SMTP__PASSWORD', 'smtp-secret')
|
||||
service = InvitationDeliveryService(_app({'api': {'webui_url': 'https://config.example'}}))
|
||||
|
||||
assert service.build_invitation_link('lbi_secret') == (
|
||||
'https://env.langbot.example/invitations/accept#token=lbi_secret'
|
||||
)
|
||||
assert service.capability() == {'enabled': True, 'provider': 'smtp'}
|
||||
|
||||
|
||||
async def test_cloud_invitation_email_has_branded_html_plain_fallback_and_expiry_copy():
|
||||
service = InvitationDeliveryService(_app({}))
|
||||
link = 'https://cloud.langbot.app/invitations/accept#token=lbi_secret&next=<unsafe>'
|
||||
|
||||
text = service._plain_text('Research & Development', link)
|
||||
html = service._html('Research & Development', link)
|
||||
|
||||
assert 'LangBot Cloud' in text
|
||||
assert 'Research & Development' in text
|
||||
assert '7 days' in text
|
||||
assert link in text
|
||||
assert 'Accept invitation' in html
|
||||
assert 'Research & Development' in html
|
||||
assert 'expires in 7 days' in html
|
||||
assert 'lbi_secret&next=<unsafe>' in html
|
||||
@@ -0,0 +1,375 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import datetime
|
||||
import uuid
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import sqlalchemy
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||||
|
||||
from langbot.pkg.entity.persistence.base import Base
|
||||
from langbot.pkg.entity.persistence.user import User
|
||||
from langbot.pkg.entity.persistence.workspace import (
|
||||
Workspace,
|
||||
WorkspaceExecutionState,
|
||||
WorkspaceInvitation,
|
||||
WorkspaceMembership,
|
||||
)
|
||||
from langbot.pkg.workspace.collaboration import (
|
||||
InvitationEmailMismatchError,
|
||||
InvitationExpiredError,
|
||||
InvitationRoleError,
|
||||
InvitationUsedError,
|
||||
LastOwnerError,
|
||||
MembershipPermissionError,
|
||||
WorkspaceCollaborationService,
|
||||
)
|
||||
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
|
||||
from langbot.pkg.workspace.service import WorkspaceService
|
||||
from langbot.pkg.workspace.policy import CloudWorkspacePolicy
|
||||
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def collaboration_context(tmp_path):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-collaboration.db"}')
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
|
||||
persistence_mgr = SimpleNamespace(get_db_engine=lambda: engine)
|
||||
application = SimpleNamespace(persistence_mgr=persistence_mgr)
|
||||
workspace_service = WorkspaceService(application, instance_uuid='instance-collaboration-test')
|
||||
service = WorkspaceCollaborationService(application, workspace_service)
|
||||
session_factory = async_sessionmaker(engine, expire_on_commit=False)
|
||||
|
||||
async with session_factory() as session:
|
||||
async with session.begin():
|
||||
owner = User(
|
||||
uuid=str(uuid.uuid4()),
|
||||
user='owner@example.com',
|
||||
normalized_email='owner@example.com',
|
||||
password='owner-hash',
|
||||
account_type='local',
|
||||
)
|
||||
session.add(owner)
|
||||
await session.flush()
|
||||
workspace, owner_membership = await workspace_service.bootstrap_local_account(
|
||||
owner.uuid,
|
||||
session=session,
|
||||
)
|
||||
|
||||
yield service, workspace_service, session_factory, owner, workspace, owner_membership
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def _add_account(session_factory, email: str) -> User:
|
||||
async with session_factory() as session:
|
||||
async with session.begin():
|
||||
account = User(
|
||||
uuid=str(uuid.uuid4()),
|
||||
user=email,
|
||||
normalized_email=email.strip().casefold(),
|
||||
password='member-hash',
|
||||
account_type='local',
|
||||
)
|
||||
session.add(account)
|
||||
await session.flush()
|
||||
return account
|
||||
|
||||
|
||||
async def test_invitation_secret_is_hashed_and_acceptance_is_one_time(collaboration_context):
|
||||
service, _, session_factory, _, workspace, owner_membership = collaboration_context
|
||||
created = await service.create_invitation(
|
||||
workspace.uuid,
|
||||
owner_membership,
|
||||
'member@example.com',
|
||||
'developer',
|
||||
)
|
||||
|
||||
assert created.token.startswith('lbi_')
|
||||
assert created.invitation.token_hash != created.token
|
||||
assert created.token not in created.invitation.token_hash
|
||||
|
||||
account = await _add_account(session_factory, 'MEMBER@example.com')
|
||||
membership = await service.accept_invitation(created.token, account.uuid)
|
||||
assert membership.workspace_uuid == workspace.uuid
|
||||
assert membership.role == 'developer'
|
||||
|
||||
with pytest.raises(InvitationUsedError):
|
||||
await service.accept_invitation(created.token, account.uuid)
|
||||
|
||||
async with session_factory() as session:
|
||||
persisted = await session.get(WorkspaceInvitation, created.invitation.uuid)
|
||||
assert persisted is not None
|
||||
assert persisted.status == 'accepted'
|
||||
assert not hasattr(persisted, 'token')
|
||||
|
||||
|
||||
async def test_concurrent_invitation_acceptance_creates_one_membership(collaboration_context):
|
||||
service, _, session_factory, _, workspace, owner_membership = collaboration_context
|
||||
created = await service.create_invitation(
|
||||
workspace.uuid,
|
||||
owner_membership,
|
||||
'race@example.com',
|
||||
'viewer',
|
||||
)
|
||||
account = await _add_account(session_factory, 'race@example.com')
|
||||
|
||||
results = await asyncio.gather(
|
||||
service.accept_invitation(created.token, account.uuid),
|
||||
service.accept_invitation(created.token, account.uuid),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
assert sum(isinstance(result, WorkspaceMembership) for result in results) == 1
|
||||
assert sum(isinstance(result, InvitationUsedError) for result in results) == 1
|
||||
async with session_factory() as session:
|
||||
count = await session.scalar(
|
||||
sqlalchemy.select(sqlalchemy.func.count())
|
||||
.select_from(WorkspaceMembership)
|
||||
.where(
|
||||
WorkspaceMembership.workspace_uuid == workspace.uuid,
|
||||
WorkspaceMembership.account_uuid == account.uuid,
|
||||
)
|
||||
)
|
||||
assert count == 1
|
||||
|
||||
|
||||
async def test_invitation_rejects_owner_role_and_email_mismatch(collaboration_context):
|
||||
service, _, session_factory, _, workspace, owner_membership = collaboration_context
|
||||
with pytest.raises(InvitationRoleError):
|
||||
await service.create_invitation(
|
||||
workspace.uuid,
|
||||
owner_membership,
|
||||
'member@example.com',
|
||||
'owner',
|
||||
)
|
||||
|
||||
created = await service.create_invitation(
|
||||
workspace.uuid,
|
||||
owner_membership,
|
||||
'expected@example.com',
|
||||
'viewer',
|
||||
)
|
||||
wrong_account = await _add_account(session_factory, 'wrong@example.com')
|
||||
with pytest.raises(InvitationEmailMismatchError):
|
||||
await service.accept_invitation(created.token, wrong_account.uuid)
|
||||
|
||||
|
||||
async def test_expired_invitation_is_inspectable_before_periodic_cleanup_deletes_it(collaboration_context):
|
||||
service, _, session_factory, _, workspace, owner_membership = collaboration_context
|
||||
created = await service.create_invitation(workspace.uuid, owner_membership, 'expired@example.com', 'viewer')
|
||||
async with session_factory.begin() as session:
|
||||
invitation = await session.get(WorkspaceInvitation, created.invitation.uuid)
|
||||
invitation.expires_at = service._utcnow() - datetime.timedelta(minutes=1)
|
||||
|
||||
with pytest.raises(InvitationExpiredError):
|
||||
await service.inspect_invitation(created.token)
|
||||
|
||||
assert await service.cleanup_expired_invitations(retention=datetime.timedelta(0)) == 1
|
||||
async with session_factory() as session:
|
||||
assert await session.get(WorkspaceInvitation, created.invitation.uuid) is None
|
||||
|
||||
|
||||
async def test_last_owner_cannot_be_demoted(collaboration_context):
|
||||
service, _, session_factory, _, workspace, owner_membership = collaboration_context
|
||||
created = await service.create_invitation(
|
||||
workspace.uuid,
|
||||
owner_membership,
|
||||
'second@example.com',
|
||||
'admin',
|
||||
)
|
||||
second = await _add_account(session_factory, 'second@example.com')
|
||||
second_membership = await service.accept_invitation(created.token, second.uuid)
|
||||
|
||||
with pytest.raises(LastOwnerError):
|
||||
await service.update_member_role(
|
||||
workspace.uuid,
|
||||
owner_membership.account_uuid,
|
||||
'admin',
|
||||
owner_membership,
|
||||
)
|
||||
|
||||
with pytest.raises(MembershipPermissionError):
|
||||
await service.update_member_role(
|
||||
workspace.uuid,
|
||||
owner_membership.account_uuid,
|
||||
'viewer',
|
||||
second_membership,
|
||||
)
|
||||
|
||||
promoted = await service.update_member_role(
|
||||
workspace.uuid,
|
||||
second.uuid,
|
||||
'owner',
|
||||
owner_membership,
|
||||
)
|
||||
assert promoted.role == 'owner'
|
||||
demoted = await service.update_member_role(
|
||||
workspace.uuid,
|
||||
owner_membership.account_uuid,
|
||||
'admin',
|
||||
owner_membership,
|
||||
)
|
||||
assert demoted.role == 'admin'
|
||||
|
||||
|
||||
async def test_workspace_selector_requires_membership(collaboration_context):
|
||||
service, _, session_factory, _, workspace, _ = collaboration_context
|
||||
outsider = await _add_account(session_factory, 'outsider@example.com')
|
||||
|
||||
with pytest.raises(WorkspaceNotFoundError):
|
||||
await service.resolve_account_workspace(outsider.uuid, workspace.uuid)
|
||||
|
||||
|
||||
async def test_cloud_member_listing_opens_workspace_uow_when_request_scope_has_closed(
|
||||
collaboration_context,
|
||||
):
|
||||
service, _, session_factory, _, workspace, owner_membership = collaboration_context
|
||||
entered_workspaces: list[str] = []
|
||||
|
||||
@asynccontextmanager
|
||||
async def tenant_uow(workspace_uuid: str):
|
||||
entered_workspaces.append(workspace_uuid)
|
||||
async with session_factory.begin() as session:
|
||||
yield SimpleNamespace(session=session)
|
||||
|
||||
service.ap.persistence_mgr = SimpleNamespace(
|
||||
mode=SimpleNamespace(value='cloud_runtime'),
|
||||
current_session=lambda: None,
|
||||
tenant_uow=tenant_uow,
|
||||
require_current_session=lambda: (_ for _ in ()).throw(AssertionError('list_members must open a Workspace UoW')),
|
||||
get_db_engine=lambda: session_factory.kw['bind'],
|
||||
)
|
||||
|
||||
members = await service.list_members(workspace.uuid, owner_membership)
|
||||
|
||||
assert entered_workspaces == [workspace.uuid]
|
||||
assert [(member.email, member.membership.role) for member in members] == [('owner@example.com', 'owner')]
|
||||
|
||||
|
||||
async def test_cloud_directory_requires_explicit_selector_and_allows_core_membership_management(
|
||||
collaboration_context,
|
||||
):
|
||||
_service, workspace_service, session_factory, owner, _workspace, _membership = collaboration_context
|
||||
cloud_workspace_uuid = str(uuid.uuid4())
|
||||
cloud_membership = WorkspaceMembership(
|
||||
uuid=str(uuid.uuid4()),
|
||||
workspace_uuid=cloud_workspace_uuid,
|
||||
account_uuid=owner.uuid,
|
||||
role='owner',
|
||||
status='active',
|
||||
projection_revision=4,
|
||||
)
|
||||
async with session_factory.begin() as session:
|
||||
session.add(
|
||||
Workspace(
|
||||
uuid=cloud_workspace_uuid,
|
||||
instance_uuid='instance-collaboration-test',
|
||||
name='Projected Workspace',
|
||||
slug='projected-workspace',
|
||||
source='cloud_projection',
|
||||
projection_revision=4,
|
||||
)
|
||||
)
|
||||
session.add(
|
||||
WorkspaceExecutionState(
|
||||
workspace_uuid=cloud_workspace_uuid,
|
||||
instance_uuid='instance-collaboration-test',
|
||||
active_generation=8,
|
||||
state='active',
|
||||
write_fenced=False,
|
||||
source='cloud',
|
||||
desired_state_revision=8,
|
||||
)
|
||||
)
|
||||
session.add(cloud_membership)
|
||||
|
||||
cloud_service = WorkspaceCollaborationService(
|
||||
workspace_service.ap,
|
||||
WorkspaceService(
|
||||
workspace_service.ap,
|
||||
policy=CloudWorkspacePolicy(),
|
||||
instance_uuid='instance-collaboration-test',
|
||||
),
|
||||
)
|
||||
with pytest.raises(WorkspaceNotFoundError):
|
||||
await cloud_service.resolve_account_workspace(owner.uuid, None)
|
||||
|
||||
access = await cloud_service.resolve_account_workspace(owner.uuid, cloud_workspace_uuid)
|
||||
assert access.workspace.uuid == cloud_workspace_uuid
|
||||
assert access.execution.placement_generation == 8
|
||||
|
||||
created = await cloud_service.create_invitation(
|
||||
cloud_workspace_uuid,
|
||||
cloud_membership,
|
||||
'member@example.com',
|
||||
'viewer',
|
||||
)
|
||||
assert created.invitation.workspace_uuid == cloud_workspace_uuid
|
||||
assert created.token.startswith('lbi_')
|
||||
|
||||
member = await _add_account(session_factory, 'member@example.com')
|
||||
membership = await cloud_service.accept_invitation(created.token, member.uuid)
|
||||
assert membership.workspace_uuid == cloud_workspace_uuid
|
||||
assert membership.role == 'viewer'
|
||||
|
||||
updated = await cloud_service.update_member_role(
|
||||
cloud_workspace_uuid,
|
||||
member.uuid,
|
||||
'developer',
|
||||
cloud_membership,
|
||||
)
|
||||
assert updated.role == 'developer'
|
||||
|
||||
removed = await cloud_service.remove_member(
|
||||
cloud_workspace_uuid,
|
||||
member.uuid,
|
||||
cloud_membership,
|
||||
)
|
||||
assert removed.status == 'removed'
|
||||
|
||||
|
||||
async def test_invitation_lock_registry_does_not_retain_sequential_tokens(collaboration_context):
|
||||
service, *_ = collaboration_context
|
||||
|
||||
for index in range(100):
|
||||
async with service._invitation_lock(f'token-{index}'):
|
||||
pass
|
||||
|
||||
assert service._invitation_locks == {}
|
||||
|
||||
|
||||
async def test_invitation_lock_registry_keeps_waiters_serialized(collaboration_context):
|
||||
service, *_ = collaboration_context
|
||||
first_entered = asyncio.Event()
|
||||
release_first = asyncio.Event()
|
||||
order: list[str] = []
|
||||
|
||||
async def first_worker() -> None:
|
||||
async with service._invitation_lock('same-token'):
|
||||
order.append('first')
|
||||
first_entered.set()
|
||||
await release_first.wait()
|
||||
|
||||
async def second_worker() -> None:
|
||||
await first_entered.wait()
|
||||
async with service._invitation_lock('same-token'):
|
||||
order.append('second')
|
||||
|
||||
first_task = asyncio.create_task(first_worker())
|
||||
second_task = asyncio.create_task(second_worker())
|
||||
await first_entered.wait()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert order == ['first']
|
||||
release_first.set()
|
||||
await asyncio.gather(first_task, second_task)
|
||||
|
||||
assert order == ['first', 'second']
|
||||
assert service._invitation_locks == {}
|
||||
@@ -0,0 +1,19 @@
|
||||
from langbot.pkg.workspace.policy import (
|
||||
CloudWorkspacePolicy,
|
||||
SingleWorkspacePolicy,
|
||||
open_core_workspace_policy,
|
||||
)
|
||||
|
||||
|
||||
def test_open_source_policy_is_single_workspace() -> None:
|
||||
policy = open_core_workspace_policy()
|
||||
|
||||
assert policy.workspace_limit == 1
|
||||
assert policy.multi_workspace_enabled is False
|
||||
|
||||
|
||||
def test_cloud_policy_is_not_exported_as_the_default_policy() -> None:
|
||||
# CloudWorkspacePolicy remains a contract for the future verified, closed
|
||||
# bootstrap. Constructing it explicitly must not change the OSS default.
|
||||
assert CloudWorkspacePolicy().multi_workspace_enabled is True
|
||||
assert isinstance(open_core_workspace_policy(), SingleWorkspacePolicy)
|
||||
@@ -0,0 +1,264 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import sqlalchemy
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||||
|
||||
from langbot.pkg.entity.persistence.base import Base
|
||||
from langbot.pkg.entity.persistence.user import User
|
||||
from langbot.pkg.entity.persistence.workspace import (
|
||||
Workspace,
|
||||
WorkspaceExecutionSource,
|
||||
WorkspaceExecutionState,
|
||||
WorkspaceMembership,
|
||||
WorkspaceSource,
|
||||
)
|
||||
from langbot.pkg.workspace import (
|
||||
WorkspaceExecutionUnavailableError,
|
||||
WorkspaceExecutionBinding,
|
||||
WorkspaceGenerationMismatchError,
|
||||
WorkspaceInvariantError,
|
||||
WorkspaceLimitExceededError,
|
||||
WorkspaceOwnerAlreadyExistsError,
|
||||
WorkspaceService,
|
||||
)
|
||||
from langbot.pkg.workspace.policy import CloudWorkspacePolicy
|
||||
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def workspace_test_context(tmp_path):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-service.db"}')
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
|
||||
persistence_mgr = SimpleNamespace(get_db_engine=lambda: engine)
|
||||
application = SimpleNamespace(persistence_mgr=persistence_mgr)
|
||||
session_factory = async_sessionmaker(engine, expire_on_commit=False)
|
||||
service = WorkspaceService(application, instance_uuid='instance_service_test')
|
||||
yield service, session_factory
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def _insert_account(session, email: str) -> str:
|
||||
account_uuid = str(uuid.uuid4())
|
||||
session.add(
|
||||
User(
|
||||
uuid=account_uuid,
|
||||
user=email,
|
||||
normalized_email=email.strip().casefold(),
|
||||
password='hashed-password',
|
||||
account_type='local',
|
||||
)
|
||||
)
|
||||
await session.flush()
|
||||
return account_uuid
|
||||
|
||||
|
||||
async def test_account_uuid_default_supports_existing_core_insert_path(workspace_test_context):
|
||||
_service, session_factory = workspace_test_context
|
||||
|
||||
async with session_factory.begin() as session:
|
||||
await session.execute(
|
||||
sqlalchemy.insert(User).values(
|
||||
user='core-insert@example.com',
|
||||
normalized_email='core-insert@example.com',
|
||||
password='hashed-password',
|
||||
account_type='local',
|
||||
)
|
||||
)
|
||||
|
||||
async with session_factory() as session:
|
||||
account = await session.scalar(sqlalchemy.select(User))
|
||||
assert account is not None
|
||||
uuid.UUID(account.uuid)
|
||||
assert account.status == 'active'
|
||||
assert account.source == 'local'
|
||||
|
||||
|
||||
async def test_bootstrap_local_account_uses_callers_transaction(workspace_test_context):
|
||||
service, session_factory = workspace_test_context
|
||||
|
||||
async with session_factory() as session:
|
||||
async with session.begin():
|
||||
account_uuid = await _insert_account(session, 'owner@example.com')
|
||||
workspace, membership = await service.bootstrap_local_account(account_uuid, session=session)
|
||||
assert session.in_transaction()
|
||||
assert workspace.created_by_account_uuid == account_uuid
|
||||
assert membership.account_uuid == account_uuid
|
||||
assert membership.role == 'owner'
|
||||
|
||||
async with session_factory() as session:
|
||||
persisted_workspace = await session.scalar(sqlalchemy.select(Workspace))
|
||||
persisted_membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
|
||||
execution_state = await session.scalar(sqlalchemy.select(WorkspaceExecutionState))
|
||||
assert persisted_workspace is not None
|
||||
assert persisted_membership is not None
|
||||
assert execution_state is not None
|
||||
assert execution_state.workspace_uuid == persisted_workspace.uuid
|
||||
assert execution_state.instance_uuid == 'instance_service_test'
|
||||
assert execution_state.active_generation == 1
|
||||
assert execution_state.write_fenced is False
|
||||
|
||||
|
||||
async def test_ensure_singleton_workspace_is_idempotent(workspace_test_context):
|
||||
service, session_factory = workspace_test_context
|
||||
|
||||
first = await service.ensure_singleton_workspace()
|
||||
second = await service.ensure_singleton_workspace()
|
||||
|
||||
assert second.uuid == first.uuid
|
||||
async with session_factory() as session:
|
||||
assert await session.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(Workspace)) == 1
|
||||
assert (
|
||||
await session.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(WorkspaceExecutionState)) == 1
|
||||
)
|
||||
|
||||
|
||||
async def test_create_local_workspace_rejects_second_workspace(workspace_test_context):
|
||||
service, session_factory = workspace_test_context
|
||||
|
||||
await service.create_local_workspace(name='First', slug='first')
|
||||
with pytest.raises(WorkspaceLimitExceededError) as exc_info:
|
||||
await service.create_local_workspace(name='Second', slug='second')
|
||||
|
||||
assert exc_info.value.code == 'edition_limit'
|
||||
async with session_factory() as session:
|
||||
assert await session.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(Workspace)) == 1
|
||||
|
||||
|
||||
async def test_initial_owner_cannot_be_claimed_by_another_account(workspace_test_context):
|
||||
service, session_factory = workspace_test_context
|
||||
|
||||
async with session_factory() as session:
|
||||
async with session.begin():
|
||||
first_account_uuid = await _insert_account(session, 'first@example.com')
|
||||
second_account_uuid = await _insert_account(session, 'second@example.com')
|
||||
await service.bootstrap_local_account(first_account_uuid, session=session)
|
||||
|
||||
with pytest.raises(WorkspaceOwnerAlreadyExistsError):
|
||||
await service.claim_initial_owner(second_account_uuid)
|
||||
|
||||
async with session_factory() as session:
|
||||
owners = (
|
||||
await session.scalars(sqlalchemy.select(WorkspaceMembership).where(WorkspaceMembership.role == 'owner'))
|
||||
).all()
|
||||
assert len(owners) == 1
|
||||
assert owners[0].account_uuid == first_account_uuid
|
||||
|
||||
|
||||
async def test_execution_binding_returns_persisted_generation(workspace_test_context):
|
||||
service, session_factory = workspace_test_context
|
||||
workspace = await service.ensure_singleton_workspace()
|
||||
|
||||
async with session_factory.begin() as session:
|
||||
execution_state = await session.get(WorkspaceExecutionState, workspace.uuid)
|
||||
execution_state.active_generation = 7
|
||||
|
||||
binding = await service.get_local_execution_binding(
|
||||
workspace.uuid,
|
||||
expected_generation=7,
|
||||
)
|
||||
|
||||
assert binding.instance_uuid == 'instance_service_test'
|
||||
assert binding.workspace_uuid == workspace.uuid
|
||||
assert binding.placement_generation == 7
|
||||
assert binding.state == 'active'
|
||||
assert binding.write_fenced is False
|
||||
|
||||
with pytest.raises(WorkspaceGenerationMismatchError):
|
||||
await service.get_local_execution_binding(workspace.uuid, expected_generation=6)
|
||||
|
||||
|
||||
async def test_execution_binding_fails_closed_when_fenced(workspace_test_context):
|
||||
service, session_factory = workspace_test_context
|
||||
workspace = await service.ensure_singleton_workspace()
|
||||
|
||||
async with session_factory.begin() as session:
|
||||
execution_state = await session.get(WorkspaceExecutionState, workspace.uuid)
|
||||
execution_state.write_fenced = True
|
||||
|
||||
with pytest.raises(WorkspaceExecutionUnavailableError):
|
||||
await service.get_local_execution_context(workspace.uuid)
|
||||
|
||||
|
||||
async def test_cloud_projection_requires_explicit_general_binding(workspace_test_context):
|
||||
service, session_factory = workspace_test_context
|
||||
await service.ensure_singleton_workspace()
|
||||
cloud_workspace_uuid = str(uuid.uuid4())
|
||||
|
||||
async with session_factory.begin() as session:
|
||||
session.add(
|
||||
Workspace(
|
||||
uuid=cloud_workspace_uuid,
|
||||
instance_uuid='instance_service_test',
|
||||
name='Cloud Workspace',
|
||||
slug='cloud-workspace',
|
||||
source=WorkspaceSource.CLOUD_PROJECTION.value,
|
||||
)
|
||||
)
|
||||
session.add(
|
||||
WorkspaceExecutionState(
|
||||
workspace_uuid=cloud_workspace_uuid,
|
||||
instance_uuid='instance_service_test',
|
||||
active_generation=9,
|
||||
state='active',
|
||||
write_fenced=False,
|
||||
source=WorkspaceExecutionSource.CLOUD.value,
|
||||
)
|
||||
)
|
||||
|
||||
binding = await service.get_execution_binding(
|
||||
cloud_workspace_uuid,
|
||||
expected_generation=9,
|
||||
)
|
||||
assert binding.workspace_uuid == cloud_workspace_uuid
|
||||
assert binding.placement_generation == 9
|
||||
|
||||
with pytest.raises(WorkspaceInvariantError):
|
||||
await service.get_local_execution_binding(cloud_workspace_uuid)
|
||||
|
||||
|
||||
async def test_cloud_policy_never_creates_or_guesses_a_workspace(workspace_test_context):
|
||||
service, _session_factory = workspace_test_context
|
||||
cloud_service = WorkspaceService(
|
||||
service.ap,
|
||||
policy=CloudWorkspacePolicy(),
|
||||
instance_uuid='instance_service_test',
|
||||
)
|
||||
|
||||
with pytest.raises(WorkspaceLimitExceededError):
|
||||
await cloud_service.create_local_workspace(name='Forbidden', slug='forbidden')
|
||||
|
||||
assert cloud_service.policy.multi_workspace_enabled is True
|
||||
|
||||
|
||||
async def test_startup_binding_snapshot_avoids_repeated_discovery():
|
||||
service = WorkspaceService(
|
||||
SimpleNamespace(persistence_mgr=SimpleNamespace()),
|
||||
policy=CloudWorkspacePolicy(),
|
||||
instance_uuid='instance-service-test',
|
||||
)
|
||||
binding = WorkspaceExecutionBinding(
|
||||
instance_uuid='instance-service-test',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
write_fenced=False,
|
||||
state='active',
|
||||
)
|
||||
service._discover_active_execution_bindings = AsyncMock(return_value=[binding])
|
||||
|
||||
assert await service.prime_startup_execution_bindings() == [binding]
|
||||
assert await service.list_active_execution_bindings() == [binding]
|
||||
assert await service.list_active_execution_bindings() == [binding]
|
||||
service._discover_active_execution_bindings.assert_awaited_once()
|
||||
|
||||
service.release_startup_execution_bindings()
|
||||
assert await service.list_active_execution_bindings() == [binding]
|
||||
assert service._discover_active_execution_bindings.await_count == 2
|
||||
Reference in New Issue
Block a user