xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-08-09 04:40:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(tenancy): add Workspace multi-tenant foundation (#2353)

Browse Source 
* 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
2026-07-30 21:43:35 +08:00
committed by GitHub
parent 463b120923
commit e1ac5e0fc8
468 changed files with 78320 additions and 13137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/integration/api/test_bots.py
+52 -1
View File
@@ -10,6 +10,7 @@ from __future__ import annotations
import pytest
from unittest.mock import MagicMock, AsyncMock, Mock
from types import SimpleNamespace
from tests.factories import FakeApp
@@ -65,12 +66,43 @@ def fake_bot_app():
)
# Auth services
account = SimpleNamespace(uuid='account-test', user='test@example.com')
app.user_service = Mock()
app.user_service.is_initialized = AsyncMock(return_value=True)
app.user_service.verify_jwt_token = AsyncMock(return_value='test@example.com')
app.user_service.get_user_by_email = AsyncMock(return_value=Mock(email='test@example.com'))
app.user_service.get_user_by_email = AsyncMock(return_value=account)
app.user_service.get_authenticated_account = AsyncMock(return_value=account)
app.workspace_collaboration_service = SimpleNamespace(
resolve_account_workspace=AsyncMock(
return_value=SimpleNamespace(
workspace=SimpleNamespace(uuid='workspace-test'),
membership=SimpleNamespace(
uuid='membership-test',
role='owner',
projection_revision=0,
),
execution=SimpleNamespace(instance_uuid='instance-test', placement_generation=1),
)
)
)
app.apikey_service = Mock()
app.apikey_service.verify_api_key = AsyncMock(return_value=True)
app.apikey_service.authenticate_api_key = AsyncMock(
return_value=SimpleNamespace(
instance_uuid='instance-test',
placement_generation=1,
api_key_uuid='api-key-test',
workspace_uuid='workspace-test',
permissions=frozenset(
{
'resource.view',
'resource.manage',
'runtime.operate',
'provider_secret.manage',
}
),
)
)
# Bot service
app.bot_service = Mock()
@@ -198,6 +230,25 @@ class TestBotLogsEndpoint:
assert 'logs' in data['data']
assert 'total_count' in data['data']
@pytest.mark.asyncio
async def test_viewer_can_read_ordinary_bot_logs(self, quart_test_client, fake_bot_app):
access = fake_bot_app.workspace_collaboration_service.resolve_account_workspace.return_value
original_role = access.membership.role
access.membership.role = 'viewer'
fake_bot_app.bot_service.list_event_logs.reset_mock()
try:
response = await quart_test_client.post(
'/api/v1/platform/bots/test-bot-uuid/logs',
headers={'Authorization': 'Bearer test_token'},
json={'from_index': -1, 'max_count': 10},
)
finally:
access.membership.role = original_role
assert response.status_code == 200
assert (await response.get_json())['code'] == 0
fake_bot_app.bot_service.list_event_logs.assert_awaited_once()
@pytest.mark.usefixtures('mock_circular_import_chain')
class TestBotSendMessageEndpoint:
tests/integration/api/test_box_security.py
+119
View File
@@ -0,0 +1,119 @@
"""Authorization tests for sensitive Box runtime observability."""
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
import quart
from langbot.pkg.api.http.controller.groups.box import BoxRouterGroup
from langbot.pkg.cloud.entitlements import EntitlementUnavailableError
pytestmark = pytest.mark.integration
WORKSPACE_UUID = '11111111-1111-4111-8111-111111111111'
def _access(account_uuid: str):
return SimpleNamespace(
workspace=SimpleNamespace(uuid=WORKSPACE_UUID),
membership=SimpleNamespace(
uuid=f'member-{account_uuid}',
role='viewer' if account_uuid == 'viewer-account' else 'owner',
projection_revision=1,
),
execution=SimpleNamespace(instance_uuid='instance-a', placement_generation=1),
)
@pytest.fixture
async def box_security_api():
accounts = {
'viewer-token': SimpleNamespace(uuid='viewer-account', user='viewer@example.com'),
'owner-token': SimpleNamespace(uuid='owner-account', user='owner@example.com'),
}
application = Mock()
application.deployment = SimpleNamespace(multi_workspace_enabled=False)
application.persistence_mgr = SimpleNamespace(tenant_uow=None)
application.user_service.get_authenticated_account = AsyncMock(side_effect=lambda token: accounts[token])
application.workspace_collaboration_service.resolve_account_workspace = AsyncMock(
side_effect=lambda account_uuid, _workspace_uuid: _access(account_uuid)
)
application.box_service.get_status = AsyncMock(return_value={'enabled': True})
application.box_service.get_backend_status = AsyncMock(
return_value={'available': True, 'enabled': True, 'backend': {'name': 'nsjail'}}
)
application.box_service.get_sessions = AsyncMock(return_value=[{'session_id': 'private-session'}])
application.box_service.get_recent_errors = Mock(return_value=[{'error': 'private error'}])
application.box_service.managed_admission_required = False
quart_app = quart.Quart(__name__)
router = BoxRouterGroup(application, quart_app)
await router.initialize()
return application, quart_app.test_client()
def _headers(token: str) -> dict[str, str]:
return {
'Authorization': f'Bearer {token}',
'X-Workspace-Id': WORKSPACE_UUID,
}
@pytest.mark.asyncio
async def test_viewer_can_read_status_but_not_sessions_or_errors(box_security_api):
application, client = box_security_api
status = await client.get('/api/v1/box/status', headers=_headers('viewer-token'))
sessions = await client.get('/api/v1/box/sessions', headers=_headers('viewer-token'))
errors = await client.get('/api/v1/box/errors', headers=_headers('viewer-token'))
assert status.status_code == 200
assert sessions.status_code == 403
assert errors.status_code == 403
application.box_service.get_sessions.assert_not_awaited()
application.box_service.get_recent_errors.assert_not_called()
@pytest.mark.asyncio
async def test_owner_can_audit_box_sessions_and_errors(box_security_api):
application, client = box_security_api
sessions = await client.get('/api/v1/box/sessions', headers=_headers('owner-token'))
errors = await client.get('/api/v1/box/errors', headers=_headers('owner-token'))
assert sessions.status_code == 200
assert errors.status_code == 200
application.box_service.get_sessions.assert_awaited_once()
application.box_service.get_recent_errors.assert_called_once()
@pytest.mark.asyncio
async def test_box_status_returns_explicit_403_when_workspace_has_no_managed_sandbox(box_security_api):
application, client = box_security_api
application.box_service.get_status.side_effect = EntitlementUnavailableError(
'Workspace entitlement does not grant managed_sandbox'
)
response = await client.get('/api/v1/box/status', headers=_headers('viewer-token'))
assert response.status_code == 403
payload = await response.get_json()
assert payload['code'] == 'managed_sandbox_unavailable'
@pytest.mark.asyncio
async def test_runtime_status_reports_connector_health_without_consuming_workspace_entitlement(box_security_api):
application, client = box_security_api
application.box_service.get_status.side_effect = EntitlementUnavailableError(
'Workspace entitlement does not grant managed_sandbox'
)
response = await client.get('/api/v1/box/runtime-status', headers=_headers('viewer-token'))
assert response.status_code == 200
assert (await response.get_json())['data']['available'] is True
application.box_service.get_backend_status.assert_awaited_once()
application.box_service.get_status.assert_not_awaited()
tests/integration/api/test_embed.py
+109 -7
View File
@@ -8,8 +8,11 @@ Run: uv run pytest tests/integration/api/test_embed.py -q
from __future__ import annotations
import json
import pytest
from unittest.mock import MagicMock, AsyncMock, Mock
from types import SimpleNamespace
from tests.factories import FakeApp
@@ -80,10 +83,18 @@ def fake_embed_app():
mock_runtime_bot = Mock()
mock_runtime_bot.bot_entity = mock_bot_entity
mock_runtime_bot.execution_context = SimpleNamespace(
instance_uuid='instance-test',
workspace_uuid='workspace-test',
placement_generation=1,
)
# Platform manager with bots
app.platform_mgr = Mock()
app.platform_mgr.bots = [mock_runtime_bot]
app.platform_mgr.resolve_public_bot = AsyncMock(
side_effect=lambda route_key: mock_runtime_bot if route_key == mock_bot_entity.uuid else None
)
# WebSocket proxy bot with adapter
mock_websocket_adapter = Mock()
@@ -94,6 +105,16 @@ def fake_embed_app():
mock_ws_proxy_bot = Mock()
mock_ws_proxy_bot.adapter = mock_websocket_adapter
app.platform_mgr.websocket_proxy_bot = mock_ws_proxy_bot
app.platform_mgr.get_websocket_proxy_bot = AsyncMock(return_value=mock_ws_proxy_bot)
app.workspace_service = SimpleNamespace(
get_execution_binding=AsyncMock(
return_value=SimpleNamespace(
instance_uuid='instance-test',
workspace_uuid='workspace-test',
placement_generation=1,
)
)
)
# Monitoring service for feedback
app.monitoring_service = Mock()
@@ -117,12 +138,13 @@ class TestEmbedWidgetEndpoint:
"""Tests for widget.js endpoint."""
@pytest.mark.asyncio
async def test_get_widget_js_success(self, quart_test_client):
async def test_get_widget_js_success(self, quart_test_client, fake_embed_app):
"""GET /api/v1/embed/{bot_uuid}/widget.js returns JS."""
response = await quart_test_client.get('/api/v1/embed/a1b2c3d4-5678-90ab-cdef-123456789abc/widget.js')
assert response.status_code == 200
assert 'javascript' in response.content_type
fake_embed_app.platform_mgr.resolve_public_bot.assert_any_await('a1b2c3d4-5678-90ab-cdef-123456789abc')
@pytest.mark.asyncio
async def test_get_widget_js_invalid_uuid(self, quart_test_client):
@@ -203,9 +225,8 @@ class TestEmbedMessagesEndpoint:
data = await response.get_json()
assert data['code'] == 0
assert 'messages' in data['data']
fake_embed_app.platform_mgr.websocket_proxy_bot.adapter.get_websocket_messages.assert_called_with(
'test-pipeline-uuid', 'person', SESSION_ID
)
proxy_bot = fake_embed_app.platform_mgr.get_websocket_proxy_bot.return_value
proxy_bot.adapter.get_websocket_messages.assert_called_with('test-pipeline-uuid', 'person', SESSION_ID)
@pytest.mark.asyncio
async def test_get_messages_group_success(self, quart_test_client):
@@ -253,9 +274,8 @@ class TestEmbedResetEndpoint:
assert response.status_code == 200
data = await response.get_json()
assert data['code'] == 0
fake_embed_app.platform_mgr.websocket_proxy_bot.adapter.reset_session.assert_called_with(
'test-pipeline-uuid', 'person', SESSION_ID
)
proxy_bot = fake_embed_app.platform_mgr.get_websocket_proxy_bot.return_value
proxy_bot.adapter.reset_session.assert_called_with('test-pipeline-uuid', 'person', SESSION_ID)
@pytest.mark.asyncio
async def test_reset_session_requires_session_id(self, quart_test_client):
@@ -316,3 +336,85 @@ class TestEmbedFeedbackEndpoint:
)
assert response.status_code == 400
@pytest.mark.usefixtures('mock_circular_import_chain')
class TestEmbedWebSocketEndpoint:
"""The public socket authenticates before resolving shared runtime state."""
@pytest.mark.asyncio
async def test_authenticates_before_connecting(self, quart_test_client, fake_embed_app):
async with quart_test_client.websocket(
f'/api/v1/embed/a1b2c3d4-5678-90ab-cdef-123456789abc/ws/connect'
f'?session_type=person&session_id={SESSION_ID}',
headers={'Origin': 'http://localhost'},
) as websocket:
await websocket.send(json.dumps({'type': 'authenticate', 'token': ''}))
connected = json.loads(await websocket.receive())
assert connected['type'] == 'connected'
assert connected['bot_uuid'] == 'a1b2c3d4-5678-90ab-cdef-123456789abc'
await websocket.send(json.dumps({'type': 'disconnect'}))
fake_embed_app.workspace_service.get_execution_binding.assert_awaited_with(
'workspace-test',
expected_generation=1,
)
@pytest.mark.asyncio
async def test_rejects_non_auth_first_frame_before_runtime_lookup(self, quart_test_client, fake_embed_app):
fake_embed_app.platform_mgr.get_websocket_proxy_bot.reset_mock()
async with quart_test_client.websocket(
f'/api/v1/embed/a1b2c3d4-5678-90ab-cdef-123456789abc/ws/connect'
f'?session_type=person&session_id={SESSION_ID}',
headers={'Origin': 'http://localhost'},
) as websocket:
await websocket.send(json.dumps({'type': 'message', 'message': []}))
response = json.loads(await websocket.receive())
assert response == {'type': 'error', 'message': 'Unauthorized'}
fake_embed_app.platform_mgr.get_websocket_proxy_bot.assert_not_awaited()
@pytest.mark.asyncio
async def test_rejects_invalid_turnstile_session_before_runtime_lookup(self, quart_test_client, fake_embed_app):
fake_embed_app.platform_mgr.get_websocket_proxy_bot.reset_mock()
config = fake_embed_app.platform_mgr.resolve_public_bot.side_effect(
'a1b2c3d4-5678-90ab-cdef-123456789abc'
).bot_entity.adapter_config
config['turnstile_secret_key'] = 'test-secret'
try:
async with quart_test_client.websocket(
f'/api/v1/embed/a1b2c3d4-5678-90ab-cdef-123456789abc/ws/connect'
f'?session_type=person&session_id={SESSION_ID}',
headers={'Origin': 'http://localhost'},
) as websocket:
await websocket.send(json.dumps({'type': 'authenticate', 'token': 'invalid'}))
response = json.loads(await websocket.receive())
assert response == {'type': 'error', 'message': 'Unauthorized'}
finally:
config['turnstile_secret_key'] = ''
fake_embed_app.platform_mgr.get_websocket_proxy_bot.assert_not_awaited()
@pytest.mark.asyncio
async def test_rejects_message_when_bot_is_disabled_after_connect(self, quart_test_client, fake_embed_app):
runtime_bot = fake_embed_app.platform_mgr.resolve_public_bot.side_effect('a1b2c3d4-5678-90ab-cdef-123456789abc')
adapter = fake_embed_app.platform_mgr.get_websocket_proxy_bot.return_value.adapter
adapter.handle_websocket_message.reset_mock()
async with quart_test_client.websocket(
f'/api/v1/embed/a1b2c3d4-5678-90ab-cdef-123456789abc/ws/connect'
f'?session_type=person&session_id={SESSION_ID}',
headers={'Origin': 'http://localhost'},
) as websocket:
await websocket.send(json.dumps({'type': 'authenticate', 'token': ''}))
assert json.loads(await websocket.receive())['type'] == 'connected'
runtime_bot.bot_entity.enable = False
try:
await websocket.send(json.dumps({'type': 'message', 'message': [{'type': 'text', 'text': 'hi'}]}))
response = json.loads(await websocket.receive())
assert response == {'type': 'error', 'message': 'Bot is unavailable'}
finally:
runtime_bot.bot_entity.enable = True
adapter.handle_websocket_message.assert_not_awaited()
tests/integration/api/test_fresh_oss_workspace_journey.py
+214
View File
@@ -0,0 +1,214 @@
from __future__ import annotations
import json
import logging
from types import SimpleNamespace
import pytest
import sqlalchemy as sa
from quart import Quart
from langbot.pkg.api.http.controller.groups.system import SystemRouterGroup
from langbot.pkg.api.http.controller.groups.user import UserRouterGroup
from langbot.pkg.api.http.controller.groups.workspaces import WorkspacesRouterGroup
from langbot.pkg.api.http.service.user import UserService
from langbot.pkg.entity.persistence.metadata import WorkspaceMetadata
from langbot.pkg.persistence.mgr import PersistenceManager
from langbot.pkg.utils import constants
from langbot.pkg.workspace.collaboration import WorkspaceCollaborationService
from langbot.pkg.workspace.service import WorkspaceService
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
def _authorization(token: str, workspace_uuid: str | None = None) -> dict[str, str]:
headers = {'Authorization': f'Bearer {token}'}
if workspace_uuid is not None:
headers['X-Workspace-Id'] = workspace_uuid
return headers
async def test_fresh_oss_workspace_http_journey_uses_real_sqlite_persistence(
tmp_path,
monkeypatch,
):
"""Exercise the first-run Workspace journey through real HTTP handlers."""
instance_uuid = 'fresh-oss-workspace-journey'
monkeypatch.setattr(constants, 'instance_id', instance_uuid)
application = SimpleNamespace(
logger=logging.getLogger('fresh-oss-workspace-journey-test'),
instance_config=SimpleNamespace(
data={
'database': {
'use': 'sqlite',
'sqlite': {'path': str(tmp_path / 'langbot.db')},
},
'system': {
'jwt': {'secret': 'fresh-oss-workspace-secret', 'expire': 3600},
'allow_modify_login_info': True,
'limitation': {},
'outbound_ips': [],
},
'api': {'global_api_key': ''},
'plugin': {'enable_marketplace': True},
'space': {
'url': 'https://space.langbot.app',
'models_gateway_api_url': 'https://api.langbot.cloud/v1',
'disable_models_service': False,
},
'mcp': {'stdio': {'enabled': True}},
}
),
)
persistence = PersistenceManager(application)
application.persistence_mgr = persistence
await persistence.initialize()
try:
application.workspace_service = WorkspaceService(
application,
instance_uuid=instance_uuid,
)
application.workspace_collaboration_service = WorkspaceCollaborationService(
application,
application.workspace_service,
)
application.user_service = UserService(application)
quart_app = Quart(__name__)
await UserRouterGroup(application, quart_app).initialize()
await WorkspacesRouterGroup(application, quart_app).initialize()
await SystemRouterGroup(application, quart_app).initialize()
client = quart_app.test_client()
initialization = await client.get('/api/v1/user/init')
assert initialization.status_code == 200
assert (await initialization.get_json())['data'] == {'initialized': False}
initialized = await client.post(
'/api/v1/user/init',
json={'user': 'owner@example.com', 'password': 'owner-password'},
)
assert initialized.status_code == 200
authenticated = await client.post(
'/api/v1/user/auth',
json={'user': 'owner@example.com', 'password': 'owner-password'},
)
assert authenticated.status_code == 200
token = (await authenticated.get_json())['data']['token']
bootstrap = await client.get(
'/api/v1/workspaces/bootstrap',
headers=_authorization(token),
)
assert bootstrap.status_code == 200
bootstrap_workspaces = (await bootstrap.get_json())['data']['workspaces']
assert len(bootstrap_workspaces) == 1
bootstrap_access = bootstrap_workspaces[0]
workspace_uuid = bootstrap_access['workspace']['uuid']
assert bootstrap_access['workspace'] == {
'uuid': workspace_uuid,
'instance_uuid': instance_uuid,
'name': 'Default Workspace',
'slug': 'default',
'type': 'team',
'status': 'active',
'source': 'local',
}
assert bootstrap_access['membership']['email'] == 'owner@example.com'
assert bootstrap_access['membership']['role'] == 'owner'
assert 'workspace.update' in bootstrap_access['permissions']
current = await client.get(
'/api/v1/workspaces/current',
headers=_authorization(token, workspace_uuid),
)
assert current.status_code == 200
current_data = (await current.get_json())['data']
assert current_data['workspace']['uuid'] == workspace_uuid
assert current_data['membership']['account_uuid'] == bootstrap_access['membership']['account_uuid']
assert current_data['membership']['role'] == 'owner'
user_info = await client.get(
'/api/v1/user/info',
headers=_authorization(token, workspace_uuid),
)
assert user_info.status_code == 200
assert (await user_info.get_json())['data'] == {
'account_uuid': bootstrap_access['membership']['account_uuid'],
'user': 'owner@example.com',
'account_type': 'local',
'has_password': True,
}
initial_system_info = await client.get(
'/api/v1/system/info',
headers=_authorization(token, workspace_uuid),
)
assert initial_system_info.status_code == 200
assert (await initial_system_info.get_json())['data']['wizard_status'] == 'none'
assert (await initial_system_info.get_json())['data']['wizard_progress'] is None
progress = {'step': 2, 'selected_adapter': 'telegram', 'bot_saved': False}
updated_progress = await client.put(
'/api/v1/system/wizard/progress',
headers=_authorization(token, workspace_uuid),
json=progress,
)
assert updated_progress.status_code == 200
persisted_progress = await persistence.execute_async(
sa.select(WorkspaceMetadata.value).where(
WorkspaceMetadata.workspace_uuid == workspace_uuid,
WorkspaceMetadata.key == 'wizard_progress',
)
)
assert json.loads(persisted_progress.scalar_one()) == progress
system_info_with_progress = await client.get(
'/api/v1/system/info',
headers=_authorization(token, workspace_uuid),
)
assert system_info_with_progress.status_code == 200
progress_data = (await system_info_with_progress.get_json())['data']
assert progress_data['wizard_status'] == 'none'
assert progress_data['wizard_progress'] == progress
completed = await client.post(
'/api/v1/system/wizard/completed',
headers=_authorization(token, workspace_uuid),
json={'status': 'completed'},
)
assert completed.status_code == 200
completed_system_info = await client.get(
'/api/v1/system/info',
headers=_authorization(token, workspace_uuid),
)
assert completed_system_info.status_code == 200
completed_data = (await completed_system_info.get_json())['data']
assert completed_data['wizard_status'] == 'completed'
assert completed_data['wizard_progress'] is None
rejected_workspace = await client.post(
'/api/v1/workspaces',
headers=_authorization(token, workspace_uuid),
json={'name': 'Second Workspace'},
)
assert rejected_workspace.status_code == 403
rejected_data = await rejected_workspace.get_json()
assert rejected_data['code'] == 'edition_limit'
persisted_wizard_status = await persistence.execute_async(
sa.select(WorkspaceMetadata.value).where(
WorkspaceMetadata.workspace_uuid == workspace_uuid,
WorkspaceMetadata.key == 'wizard_status',
)
)
assert persisted_wizard_status.scalar_one() == 'completed'
finally:
await persistence.get_db_engine().dispose()
tests/integration/api/test_knowledge.py
+23 -1
View File
@@ -10,6 +10,7 @@ from __future__ import annotations
import pytest
from unittest.mock import MagicMock, AsyncMock, Mock
from types import SimpleNamespace
from tests.factories import FakeApp
@@ -69,7 +70,28 @@ def fake_knowledge_app():
app.user_service = Mock()
app.user_service.is_initialized = AsyncMock(return_value=True)
app.user_service.verify_jwt_token = AsyncMock(return_value='test@example.com')
app.user_service.get_user_by_email = AsyncMock(return_value=Mock(email='test@example.com'))
account = SimpleNamespace(
uuid='00000000-0000-0000-0000-000000000001',
user='test@example.com',
)
app.user_service.get_authenticated_account = AsyncMock(return_value=account)
app.user_service.get_user_by_email = AsyncMock(return_value=account)
app.workspace_collaboration_service = SimpleNamespace(
resolve_account_workspace=AsyncMock(
return_value=SimpleNamespace(
execution=SimpleNamespace(
instance_uuid='instance-knowledge-api',
placement_generation=1,
),
workspace=SimpleNamespace(uuid='00000000-0000-0000-0000-00000000000a'),
membership=SimpleNamespace(
uuid='00000000-0000-0000-0000-000000000010',
role='owner',
projection_revision=0,
),
)
)
)
app.apikey_service = Mock()
app.apikey_service.verify_api_key = AsyncMock(return_value=True)
tests/integration/api/test_monitoring.py
+49 -1
View File
@@ -10,6 +10,7 @@ from __future__ import annotations
import pytest
from unittest.mock import MagicMock, AsyncMock, Mock
from types import SimpleNamespace
from tests.factories import FakeApp
@@ -66,7 +67,26 @@ def fake_monitoring_app():
app.user_service = Mock()
app.user_service.is_initialized = AsyncMock(return_value=True)
app.user_service.verify_jwt_token = AsyncMock(return_value='test@example.com')
app.user_service.get_user_by_email = AsyncMock(return_value=Mock(email='test@example.com'))
app.user_service.get_user_by_email = AsyncMock(
return_value=SimpleNamespace(
uuid='account-uuid',
user='test@example.com',
email='test@example.com',
)
)
app.workspace_collaboration_service = SimpleNamespace(
resolve_account_workspace=AsyncMock(
return_value=SimpleNamespace(
execution=SimpleNamespace(instance_uuid='instance', placement_generation=1),
workspace=SimpleNamespace(uuid='00000000-0000-0000-0000-00000000000a'),
membership=SimpleNamespace(
uuid='membership-uuid',
role='owner',
projection_revision=1,
),
)
)
)
# Monitoring service
app.monitoring_service = Mock()
@@ -135,6 +155,34 @@ class TestMonitoringOverviewEndpoint:
data = await response.get_json()
assert data['code'] == 0
@pytest.mark.asyncio
async def test_viewer_can_read_monitoring_but_cannot_export(
self,
quart_test_client,
fake_monitoring_app,
):
"""Ordinary monitoring is resource.view; export remains data.export."""
membership = (
fake_monitoring_app.workspace_collaboration_service.resolve_account_workspace.return_value.membership
)
original_role = membership.role
membership.role = 'viewer'
try:
response = await quart_test_client.get(
'/api/v1/monitoring/overview',
headers={'Authorization': 'Bearer test_token'},
)
assert response.status_code == 200
export_response = await quart_test_client.get(
'/api/v1/monitoring/export?type=messages',
headers={'Authorization': 'Bearer test_token'},
)
assert export_response.status_code == 403
assert (await export_response.get_json())['code'] == 'permission_denied'
finally:
membership.role = original_role
@pytest.mark.usefixtures('mock_circular_import_chain')
class TestMonitoringMessagesEndpoint:
tests/integration/api/test_pipelines.py
+173 -1
View File
@@ -9,10 +9,13 @@ Run: uv run pytest tests/integration/api/test_pipelines.py -q
from __future__ import annotations
import json
import pytest
from unittest.mock import MagicMock, AsyncMock, Mock
from types import SimpleNamespace
from tests.factories import FakeApp
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
pytestmark = pytest.mark.integration
@@ -54,6 +57,7 @@ def mock_circular_import_chain():
):
# Import groups after mocking to populate preregistered_groups
import langbot.pkg.api.http.controller.groups.pipelines.pipelines as _pipelines # noqa: E402, F401
import langbot.pkg.api.http.controller.groups.pipelines.websocket_chat as _websocket_chat # noqa: E402, F401
yield
@@ -75,10 +79,25 @@ def fake_pipeline_app():
)
# Auth services
account = SimpleNamespace(uuid='account-test', user='test@example.com')
app.user_service = Mock()
app.user_service.is_initialized = AsyncMock(return_value=True)
app.user_service.verify_jwt_token = AsyncMock(return_value='test@example.com')
app.user_service.get_user_by_email = AsyncMock(return_value=Mock(email='test@example.com'))
app.user_service.get_user_by_email = AsyncMock(return_value=account)
app.user_service.get_authenticated_account = AsyncMock(return_value=account)
app.workspace_collaboration_service = SimpleNamespace(
resolve_account_workspace=AsyncMock(
return_value=SimpleNamespace(
workspace=SimpleNamespace(uuid='workspace-test'),
membership=SimpleNamespace(
uuid='membership-test',
role='owner',
projection_revision=0,
),
execution=SimpleNamespace(instance_uuid='instance-test', placement_generation=1),
)
)
)
app.apikey_service = Mock()
app.apikey_service.verify_api_key = AsyncMock(return_value=True)
@@ -119,6 +138,15 @@ def fake_pipeline_app():
app.bot_service.get_bots = AsyncMock(return_value=[])
app.bot_service.create_bot = AsyncMock(return_value={'uuid': 'new-bot-uuid'})
# Workspace-scoped dashboard WebSocket proxy
websocket_adapter = Mock()
websocket_adapter.get_websocket_messages = Mock(return_value=[])
websocket_adapter.reset_session = Mock()
websocket_adapter.handle_websocket_message = AsyncMock()
websocket_proxy_bot = Mock(adapter=websocket_adapter)
app.platform_mgr = Mock()
app.platform_mgr.get_websocket_proxy_bot = AsyncMock(return_value=websocket_proxy_bot)
# MCP service (for extensions endpoint)
app.mcp_service = Mock()
app.mcp_service.get_mcp_servers = AsyncMock(return_value=[])
@@ -278,3 +306,147 @@ class TestPipelineExtensionsEndpoint:
assert response.status_code == 200
data = await response.get_json()
assert data['code'] == 0
@pytest.mark.asyncio
async def test_get_extensions_redacts_available_plugin_secrets(
self,
quart_test_client,
fake_pipeline_app,
):
connector = fake_pipeline_app.plugin_connector
raw_plugin = {
'plugin_config': {'apiKey': 'plugin-secret'},
'debug': {'plugin_debug_key': 'debug-secret'},
}
connector.list_plugins.return_value = [raw_plugin]
try:
response = await quart_test_client.get(
'/api/v1/pipelines/test-pipeline-uuid/extensions',
headers={'Authorization': 'Bearer test_token'},
)
finally:
connector.list_plugins.return_value = []
assert response.status_code == 200
plugin = (await response.get_json())['data']['available_plugins'][0]
assert plugin['plugin_config']['apiKey'] == '***'
assert plugin['debug']['plugin_debug_key'] == '***'
assert raw_plugin['plugin_config']['apiKey'] == 'plugin-secret'
@pytest.mark.asyncio
async def test_get_extensions_hides_connector_bound_to_another_workspace(
self,
quart_test_client,
fake_pipeline_app,
):
connector = fake_pipeline_app.plugin_connector
original_enabled = connector.is_enable_plugin
connector.is_enable_plugin = True
connector.require_workspace_context.reset_mock()
connector.list_plugins.reset_mock()
connector.require_workspace_context.side_effect = WorkspaceNotFoundError('Plugin resource not found')
try:
response = await quart_test_client.get(
'/api/v1/pipelines/test-pipeline-uuid/extensions',
headers={'Authorization': 'Bearer test_token'},
)
finally:
connector.require_workspace_context.side_effect = None
connector.is_enable_plugin = original_enabled
assert response.status_code == 404
connector.list_plugins.assert_not_awaited()
@pytest.mark.usefixtures('mock_circular_import_chain')
class TestPipelineDashboardWebSocket:
@pytest.mark.asyncio
async def test_websocket_authenticates_before_registering(self, quart_test_client, fake_pipeline_app):
async with quart_test_client.websocket(
'/api/v1/pipelines/test-pipeline-uuid/ws/connect?session_type=person',
headers={'Origin': 'http://localhost'},
) as websocket:
await websocket.send(
json.dumps(
{
'type': 'authenticate',
'token': 'test_token',
'workspace_uuid': 'workspace-test',
}
)
)
connected = json.loads(await websocket.receive())
assert connected['type'] == 'connected'
assert connected['pipeline_uuid'] == 'test-pipeline-uuid'
await websocket.send(json.dumps({'type': 'disconnect'}))
fake_pipeline_app.workspace_collaboration_service.resolve_account_workspace.assert_awaited_with(
'account-test',
'workspace-test',
)
fake_pipeline_app.platform_mgr.get_websocket_proxy_bot.assert_awaited()
@pytest.mark.asyncio
async def test_websocket_rejects_non_auth_first_frame(self, quart_test_client):
async with quart_test_client.websocket(
'/api/v1/pipelines/test-pipeline-uuid/ws/connect?session_type=person',
headers={'Origin': 'http://localhost'},
) as websocket:
await websocket.send(json.dumps({'type': 'message', 'message': []}))
response = json.loads(await websocket.receive())
assert response == {'type': 'error', 'message': 'Unauthorized'}
@pytest.mark.asyncio
async def test_websocket_rechecks_revocable_membership_before_each_message(
self,
quart_test_client,
fake_pipeline_app,
):
access = fake_pipeline_app.workspace_collaboration_service.resolve_account_workspace.return_value
adapter = fake_pipeline_app.platform_mgr.get_websocket_proxy_bot.return_value.adapter
original_role = access.membership.role
adapter.handle_websocket_message.reset_mock()
try:
async with quart_test_client.websocket(
'/api/v1/pipelines/test-pipeline-uuid/ws/connect?session_type=person',
headers={'Origin': 'http://localhost'},
) as websocket:
await websocket.send(
json.dumps(
{
'type': 'authenticate',
'token': 'test_token',
'workspace_uuid': 'workspace-test',
}
)
)
assert json.loads(await websocket.receive())['type'] == 'connected'
access.membership.role = 'viewer'
await websocket.send(json.dumps({'type': 'message', 'message': []}))
assert json.loads(await websocket.receive()) == {
'type': 'error',
'message': 'Unauthorized',
}
finally:
access.membership.role = original_role
adapter.handle_websocket_message.assert_not_awaited()
@pytest.mark.asyncio
async def test_dashboard_history_requires_runtime_permission(self, quart_test_client, fake_pipeline_app):
access = fake_pipeline_app.workspace_collaboration_service.resolve_account_workspace.return_value
original_role = access.membership.role
access.membership.role = 'viewer'
try:
response = await quart_test_client.get(
'/api/v1/pipelines/test-pipeline-uuid/ws/messages/person',
headers={
'Authorization': 'Bearer test_token',
'X-Workspace-Id': 'workspace-test',
},
)
finally:
access.membership.role = original_role
assert response.status_code == 403
tests/integration/api/test_plugins_security.py
+272
View File
@@ -0,0 +1,272 @@
"""Security regression tests for plugin configuration HTTP responses."""
from __future__ import annotations
import copy
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock, call
import pytest
import quart
pytestmark = pytest.mark.integration
WORKSPACE_UUID = '11111111-1111-4111-8111-111111111111'
RAW_CONFIG = {
'apiKey': 'api-secret',
'nested': {
'headers': {'Authorization': 'Bearer nested-secret', 'Accept': 'application/json'},
'refresh_token': 'refresh-secret',
'public_key': 'public-material',
'tokenizer': 'not-a-secret',
},
'credentials': {'username': 'service-user', 'password': 'service-password'},
'secret_list': ['first-secret', {'value': 'second-secret'}],
'empty_secret': '',
'enabled': True,
}
def _access(account_uuid: str):
roles = {
'viewer-account': 'viewer',
'operator-account': 'operator',
'manager-account': 'developer',
}
return SimpleNamespace(
workspace=SimpleNamespace(uuid=WORKSPACE_UUID),
membership=SimpleNamespace(
uuid=f'membership-{account_uuid}',
role=roles[account_uuid],
projection_revision=1,
),
execution=SimpleNamespace(instance_uuid='instance-test', placement_generation=1),
)
@pytest.fixture(scope='module')
def plugin_module():
"""Import the plugin router without following the core HTTP cycle."""
from tests.utils.import_isolation import MockLifecycleControlScope, isolated_sys_modules
class FakeMinimalApplication:
pass
mock_app = Mock(Application=FakeMinimalApplication)
mock_entities = Mock(LifecycleControlScope=MockLifecycleControlScope)
clear = [
'langbot.pkg.core.taskmgr',
'langbot.pkg.api.http.controller.group',
'langbot.pkg.api.http.controller.groups',
'langbot.pkg.api.http.controller.groups.plugins',
'langbot.pkg.api.http.controller.main',
]
with isolated_sys_modules(
mocks={
'langbot.pkg.core.app': mock_app,
'langbot.pkg.core.entities': mock_entities,
},
clear=clear,
):
import langbot.pkg.api.http.controller.groups.plugins as plugins
yield plugins
@pytest.fixture
async def plugin_security_api(plugin_module):
viewer = SimpleNamespace(uuid='viewer-account', user='viewer@example.com')
operator = SimpleNamespace(uuid='operator-account', user='operator@example.com')
manager = SimpleNamespace(uuid='manager-account', user='manager@example.com')
accounts = {
'viewer-token': viewer,
'operator-token': operator,
'manager-token': manager,
}
application = Mock()
application.deployment = SimpleNamespace(multi_workspace_enabled=False)
application.user_service.get_authenticated_account = AsyncMock(side_effect=lambda token: accounts[token])
application.workspace_collaboration_service.resolve_account_workspace = AsyncMock(
side_effect=lambda account_uuid, _workspace_uuid: _access(account_uuid)
)
application.apikey_service.verify_api_key = AsyncMock(return_value=False)
application.instance_config.data = {
'plugin': {'display_plugin_debug_url': 'http://localhost:5401'},
'system': {'limitation': {}},
}
raw_plugin = {
'author': 'example',
'name': 'secure-plugin',
'plugin_config': RAW_CONFIG,
'debug': {'plugin_debug_key': 'list-debug-secret'},
}
application.plugin_connector.require_workspace_context = AsyncMock()
application.plugin_connector.list_plugins = AsyncMock(return_value=[raw_plugin])
application.plugin_connector.get_plugin_info = AsyncMock(return_value=raw_plugin)
application.plugin_connector.get_debug_info = AsyncMock(return_value={'plugin_debug_key': 'runtime-debug-secret'})
application.plugin_connector.get_plugin_logs = AsyncMock(return_value=['private runtime line'])
application.plugin_connector.set_plugin_config = AsyncMock()
persistence_result = Mock()
persistence_result.scalar_one_or_none.return_value = RAW_CONFIG
application.persistence_mgr.execute_async = AsyncMock(return_value=persistence_result)
application.persistence_mgr.tenant_uow = None
quart_app = quart.Quart(__name__)
router = plugin_module.PluginsRouterGroup(application, quart_app)
await router.initialize()
return application, quart_app.test_client(), raw_plugin
def _headers(token: str) -> dict[str, str]:
return {
'Authorization': f'Bearer {token}',
'X-Workspace-Id': WORKSPACE_UUID,
}
def test_recursive_redaction_preserves_structure_without_mutating_input(plugin_module):
redacted = plugin_module.redact_plugin_secrets(RAW_CONFIG)
assert redacted['apiKey'] == '***'
assert redacted['nested']['headers'] == {
'Authorization': '***',
'Accept': 'application/json',
}
assert redacted['nested']['refresh_token'] == '***'
assert redacted['nested']['public_key'] == 'public-material'
assert redacted['nested']['tokenizer'] == 'not-a-secret'
assert redacted['credentials'] == {'username': '***', 'password': '***'}
assert redacted['secret_list'] == ['***', {'value': '***'}]
assert redacted['empty_secret'] == ''
assert redacted['enabled'] is True
assert RAW_CONFIG['apiKey'] == 'api-secret'
assert RAW_CONFIG['nested']['headers']['Authorization'] == 'Bearer nested-secret'
with pytest.raises(ValueError, match='no existing value'):
plugin_module.restore_plugin_secret_placeholders({'api_key': '***'}, {})
@pytest.mark.asyncio
async def test_viewer_plugin_reads_are_recursively_redacted(plugin_security_api):
application, client, raw_plugin = plugin_security_api
list_response = await client.get('/api/v1/plugins', headers=_headers('viewer-token'))
detail_response = await client.get(
'/api/v1/plugins/example/secure-plugin',
headers=_headers('viewer-token'),
)
config_response = await client.get(
'/api/v1/plugins/example/secure-plugin/config',
headers=_headers('viewer-token'),
)
assert list_response.status_code == 200
assert detail_response.status_code == 200
assert config_response.status_code == 200
listed_plugin = (await list_response.get_json())['data']['plugins'][0]
detailed_plugin = (await detail_response.get_json())['data']['plugin']
config = (await config_response.get_json())['data']['config']
for plugin in (listed_plugin, detailed_plugin):
assert plugin['plugin_config']['apiKey'] == '***'
assert plugin['debug']['plugin_debug_key'] == '***'
assert config['apiKey'] == '***'
assert config['nested']['headers']['Authorization'] == '***'
assert raw_plugin['plugin_config']['apiKey'] == 'api-secret'
application.plugin_connector.require_workspace_context.assert_awaited()
@pytest.mark.asyncio
async def test_manager_read_is_redacted_but_write_preserves_or_replaces_secrets(
plugin_security_api,
plugin_module,
):
application, client, _ = plugin_security_api
read_response = await client.get(
'/api/v1/plugins/example/secure-plugin/config',
headers=_headers('manager-token'),
)
masked_update = plugin_module.redact_plugin_secrets(RAW_CONFIG)
masked_update['enabled'] = False
preserved_write = await client.put(
'/api/v1/plugins/example/secure-plugin/config',
headers=_headers('manager-token'),
json=masked_update,
)
replacement = copy.deepcopy(RAW_CONFIG)
replacement['apiKey'] = 'replacement-secret'
replaced_write = await client.put(
'/api/v1/plugins/example/secure-plugin/config',
headers=_headers('manager-token'),
json=replacement,
)
preserved = copy.deepcopy(RAW_CONFIG)
preserved['enabled'] = False
assert read_response.status_code == 200
assert (await read_response.get_json())['data']['config']['apiKey'] == '***'
assert preserved_write.status_code == 200
assert replaced_write.status_code == 200
assert application.plugin_connector.set_plugin_config.await_args_list == [
call('example', 'secure-plugin', preserved),
call('example', 'secure-plugin', replacement),
]
@pytest.mark.asyncio
async def test_debug_key_requires_resource_manage_permission(plugin_security_api):
application, client, _ = plugin_security_api
viewer_denied = await client.get('/api/v1/plugins/debug-info', headers=_headers('viewer-token'))
operator_denied = await client.get('/api/v1/plugins/debug-info', headers=_headers('operator-token'))
application.plugin_connector.get_debug_info.assert_not_awaited()
allowed = await client.get('/api/v1/plugins/debug-info', headers=_headers('manager-token'))
assert viewer_denied.status_code == 403
assert operator_denied.status_code == 403
assert allowed.status_code == 200
assert (await allowed.get_json())['data'] == {
'debug_url': 'http://localhost:5401',
'plugin_debug_key': 'runtime-debug-secret',
}
application.plugin_connector.get_debug_info.assert_awaited_once_with()
@pytest.mark.asyncio
async def test_viewer_cannot_read_plugin_runtime_logs(plugin_security_api):
application, client, _ = plugin_security_api
response = await client.get(
'/api/v1/plugins/example/secure-plugin/logs',
headers=_headers('viewer-token'),
)
assert response.status_code == 403
assert (await response.get_json())['code'] == 'permission_denied'
application.plugin_connector.get_plugin_logs.assert_not_awaited()
@pytest.mark.asyncio
async def test_github_install_rejects_internal_asset_url_before_task_creation(
plugin_security_api,
):
application, client, _ = plugin_security_api
response = await client.post(
'/api/v1/plugins/install/github',
headers=_headers('manager-token'),
json={
'asset_url': 'http://169.254.169.254/latest/meta-data',
'owner': 'langbot-app',
'repo': 'demo-plugin',
'release_tag': 'v1.0.0',
},
)
assert response.status_code == 400
assert 'HTTPS GitHub release asset URL' in (await response.get_json())['msg']
application.task_mgr.create_user_task.assert_not_called()
tests/integration/api/test_providers.py
+17 -1
View File
@@ -10,6 +10,7 @@ from __future__ import annotations
import pytest
from unittest.mock import MagicMock, AsyncMock, Mock
from types import SimpleNamespace
from tests.factories import FakeApp
@@ -66,10 +67,25 @@ def fake_provider_app():
)
# Auth services
account = SimpleNamespace(uuid='account-test', user='test@example.com')
app.user_service = Mock()
app.user_service.is_initialized = AsyncMock(return_value=True)
app.user_service.verify_jwt_token = AsyncMock(return_value='test@example.com')
app.user_service.get_user_by_email = AsyncMock(return_value=Mock(email='test@example.com'))
app.user_service.get_user_by_email = AsyncMock(return_value=account)
app.user_service.get_authenticated_account = AsyncMock(return_value=account)
app.workspace_collaboration_service = SimpleNamespace(
resolve_account_workspace=AsyncMock(
return_value=SimpleNamespace(
workspace=SimpleNamespace(uuid='workspace-test'),
membership=SimpleNamespace(
uuid='membership-test',
role='owner',
projection_revision=0,
),
execution=SimpleNamespace(instance_uuid='instance-test', placement_generation=1),
)
)
)
app.apikey_service = Mock()
app.apikey_service.verify_api_key = AsyncMock(return_value=True)
tests/integration/api/test_skills_entitlements.py
+79
View File
@@ -0,0 +1,79 @@
"""Skills API behavior when a workspace plan has no managed sandbox."""
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
import quart
from langbot.pkg.api.http.controller.groups.skills import SkillsRouterGroup
from langbot.pkg.cloud.entitlements import (
EntitlementFeatureUnavailableError,
EntitlementUnavailableError,
)
pytestmark = pytest.mark.integration
WORKSPACE_UUID = '11111111-1111-4111-8111-111111111111'
@pytest.fixture
async def skills_api():
account = SimpleNamespace(uuid='owner-account', user='owner@example.com')
access = SimpleNamespace(
workspace=SimpleNamespace(uuid=WORKSPACE_UUID),
membership=SimpleNamespace(uuid='member-owner', role='owner', projection_revision=1),
execution=SimpleNamespace(instance_uuid='instance-a', placement_generation=1),
)
application = Mock()
application.deployment = SimpleNamespace(multi_workspace_enabled=False)
application.persistence_mgr = SimpleNamespace(tenant_uow=None)
application.user_service.get_authenticated_account = AsyncMock(return_value=account)
application.workspace_collaboration_service.resolve_account_workspace = AsyncMock(return_value=access)
application.skill_service.list_skills = AsyncMock(
side_effect=EntitlementFeatureUnavailableError(
'managed_sandbox',
entitlement_revision=1,
)
)
quart_app = quart.Quart(__name__)
router = SkillsRouterGroup(application, quart_app)
await router.initialize()
return application, quart_app.test_client()
@pytest.mark.asyncio
async def test_list_skills_is_empty_when_plan_has_no_managed_sandbox(skills_api):
application, client = skills_api
response = await client.get(
'/api/v1/skills',
headers={
'Authorization': 'Bearer owner-token',
'X-Workspace-Id': WORKSPACE_UUID,
},
)
assert response.status_code == 200
payload = await response.get_json()
assert payload['data'] == {'skills': []}
application.skill_service.list_skills.assert_awaited_once()
@pytest.mark.asyncio
async def test_list_skills_does_not_hide_other_entitlement_failures(skills_api):
application, client = skills_api
application.skill_service.list_skills.side_effect = EntitlementUnavailableError(
'Workspace entitlement revision rolled back'
)
response = await client.get(
'/api/v1/skills',
headers={
'Authorization': 'Bearer owner-token',
'X-Workspace-Id': WORKSPACE_UUID,
},
)
assert response.status_code == 500
tests/integration/api/test_smoke.py
+43 -2
View File
@@ -86,7 +86,7 @@ def fake_api_app():
'api': {'port': 5300},
'plugin': {'enable_marketplace': True},
'space': {'url': 'https://space.langbot.app'},
'system': {'allow_modify_login_info': True, 'limitation': {}},
'system': {'allow_modify_login_info': True, 'recovery_key': 'recovery-secret', 'limitation': {}},
}
)
@@ -160,7 +160,7 @@ class TestHealthEndpoint:
assert response.status_code == 200
data = await response.get_json()
assert data == {'code': 0, 'msg': 'ok'}
assert data == {'code': 0, 'msg': 'ok', 'resources': {}}
@pytest.mark.asyncio
async def test_healthz_no_auth_required(self, quart_test_client):
@@ -288,6 +288,47 @@ class TestUserInitEndpoint:
assert data['msg'] == 'ok'
assert data['data']['initialized'] is False
@pytest.mark.asyncio
async def test_account_info_exposes_instance_capabilities_not_first_account(self, quart_test_client, fake_api_app):
fake_api_app.user_service.is_initialized.return_value = True
fake_api_app.user_service.get_login_capabilities = AsyncMock(
return_value={'password_login_enabled': True, 'space_login_enabled': False}
)
fake_api_app.user_service.get_first_user = AsyncMock(
side_effect=AssertionError('public login bootstrap must not inspect an account')
)
response = await quart_test_client.get('/api/v1/user/account-info')
assert response.status_code == 200
data = await response.get_json()
assert data['data'] == {
'initialized': True,
'password_login_enabled': True,
'space_login_enabled': False,
}
fake_api_app.user_service.get_login_capabilities.assert_awaited_once_with()
fake_api_app.user_service.get_first_user.assert_not_awaited()
@pytest.mark.asyncio
async def test_recovery_key_resets_any_existing_account(self, quart_test_client, fake_api_app, monkeypatch):
fake_api_app.user_service.is_initialized.return_value = True
fake_api_app.user_service.get_user_by_email.return_value = Mock(user='member@example.com')
fake_api_app.user_service.reset_password = AsyncMock()
monkeypatch.setattr('langbot.pkg.api.http.controller.groups.user.asyncio.sleep', AsyncMock())
response = await quart_test_client.post(
'/api/v1/user/reset-password',
json={
'user': 'member@example.com',
'recovery_key': 'recovery-secret',
'new_password': 'new-member-password',
},
)
assert response.status_code == 200
fake_api_app.user_service.reset_password.assert_awaited_once_with('member@example.com', 'new-member-password')
@pytest.mark.usefixtures('mock_circular_import_chain')
class TestRealImports:
tests/integration/api/test_user_space_oauth.py
+338
View File
@@ -0,0 +1,338 @@
"""Security tests for the LangBot-to-Space OAuth redirect boundary."""
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
from urllib.parse import parse_qs, urlsplit
import pytest
import quart
from langbot.pkg.api.http.controller.groups.user import UserRouterGroup
pytestmark = pytest.mark.integration
WORKSPACE_UUID = '11111111-1111-4111-8111-111111111111'
@pytest.fixture
async def space_oauth_api():
account = SimpleNamespace(uuid='account-a', user='owner@example.com')
access = SimpleNamespace(
workspace=SimpleNamespace(uuid=WORKSPACE_UUID),
membership=SimpleNamespace(uuid='member-a', role='owner', projection_revision=1),
execution=SimpleNamespace(instance_uuid='instance-a', placement_generation=1),
)
application = Mock()
application.deployment = SimpleNamespace(multi_workspace_enabled=False)
application.persistence_mgr = None
application.user_service.get_authenticated_account = AsyncMock(return_value=account)
application.user_service.issue_space_oauth_state = AsyncMock(
side_effect=lambda purpose, **_: f'opaque-{purpose}-state'
)
local_account = SimpleNamespace(
uuid='account-a',
user='owner@example.com',
account_type='local',
)
bound_account = SimpleNamespace(
uuid='account-a',
user='owner@example.com',
account_type='space',
)
application.user_service.consume_space_oauth_state = AsyncMock(
side_effect=lambda state, purpose: (
local_account if (state, purpose) == ('opaque-bind-state', 'bind') else None
)
)
application.user_service.consume_space_oauth_state_details = AsyncMock(
return_value=SimpleNamespace(launch_workspace_uuid=None)
)
application.user_service.bind_space_account = AsyncMock(return_value=bound_account)
application.user_service.generate_jwt_token = AsyncMock(return_value='rotated-account-token')
application.user_service.get_user_by_uuid = AsyncMock(return_value=bound_account)
application.user_service.authenticate_space_user = AsyncMock(return_value=('space-login-token', bound_account))
application.user_service.verify_jwt_token = AsyncMock()
application.space_launch_service.consume_assertion = AsyncMock(
return_value={'account_uuid': 'account-a', 'workspace_uuid': WORKSPACE_UUID}
)
application.workspace_collaboration_service.resolve_account_workspace = AsyncMock(return_value=access)
application.space_service.get_oauth_authorize_url = Mock(
side_effect=lambda redirect_uri, state: f'https://space.example/authorize?state={state}'
)
application.space_service.exchange_oauth_code = AsyncMock(
return_value={
'access_token': 'space-access-token',
'refresh_token': 'space-refresh-token',
'expires_in': 3600,
}
)
application.instance_config.data = {
'api': {'webui_url': 'http://localhost'},
'system': {'allow_modify_login_info': True},
}
quart_app = quart.Quart(__name__)
router = UserRouterGroup(application, quart_app)
await router.initialize()
return application, quart_app.test_client()
@pytest.mark.asyncio
async def test_public_login_state_is_server_issued(space_oauth_api):
application, client = space_oauth_api
response = await client.get(
'/api/v1/user/space/authorize-url',
query_string={'redirect_uri': 'http://localhost/auth/space/callback'},
headers={'Origin': 'http://localhost'},
)
assert response.status_code == 200
authorize_url = (await response.get_json())['data']['authorize_url']
assert parse_qs(urlsplit(authorize_url).query)['state'] == ['opaque-login-state']
application.user_service.issue_space_oauth_state.assert_awaited_once_with('login')
@pytest.mark.asyncio
async def test_cloud_launch_state_is_server_issued_and_workspace_bound(space_oauth_api):
application, client = space_oauth_api
application.deployment.multi_workspace_enabled = True
response = await client.get(
'/api/v1/user/space/authorize-url',
query_string={
'redirect_uri': 'http://localhost/auth/space/callback',
'launch_workspace_uuid': WORKSPACE_UUID,
},
headers={'Origin': 'http://localhost'},
)
assert response.status_code == 200
authorize_url = (await response.get_json())['data']['authorize_url']
assert parse_qs(urlsplit(authorize_url).query)['state'] == ['opaque-login-state']
application.user_service.issue_space_oauth_state.assert_awaited_once_with(
'login',
launch_workspace_uuid=WORKSPACE_UUID,
)
@pytest.mark.asyncio
async def test_public_login_rejects_caller_supplied_state(space_oauth_api):
application, client = space_oauth_api
response = await client.get(
'/api/v1/user/space/authorize-url',
query_string={
'redirect_uri': 'http://localhost/auth/space/callback',
'state': 'jwt.must-not-be-used',
},
headers={'Origin': 'http://localhost'},
)
assert response.status_code == 200
assert (await response.get_json())['code'] == 1
application.space_service.get_oauth_authorize_url.assert_not_called()
@pytest.mark.asyncio
async def test_bind_state_is_account_bound_and_requires_authentication(space_oauth_api):
application, client = space_oauth_api
path = '/api/v1/user/space/bind-authorize-url'
query = {'redirect_uri': 'http://localhost/auth/space/callback?mode=bind'}
unauthorized = await client.get(path, query_string=query, headers={'Origin': 'http://localhost'})
response = await client.get(
path,
query_string=query,
headers={
'Origin': 'http://localhost',
'Authorization': 'Bearer user-token',
'X-Workspace-Id': WORKSPACE_UUID,
},
)
assert unauthorized.status_code == 401
assert response.status_code == 200
application.user_service.issue_space_oauth_state.assert_awaited_once_with('bind', account_uuid='account-a')
@pytest.mark.asyncio
async def test_redirect_origin_and_callback_path_are_restricted(space_oauth_api):
_, client = space_oauth_api
wrong_origin = await client.get(
'/api/v1/user/space/authorize-url',
query_string={'redirect_uri': 'https://evil.example/auth/space/callback'},
headers={'Origin': 'http://localhost'},
)
wrong_path = await client.get(
'/api/v1/user/space/authorize-url',
query_string={'redirect_uri': 'http://localhost/arbitrary'},
headers={'Origin': 'http://localhost'},
)
forged_origin = await client.get(
'/api/v1/user/space/authorize-url',
query_string={'redirect_uri': 'https://evil.example/auth/space/callback'},
headers={'Origin': 'https://evil.example'},
)
forged_host = await client.get(
'/api/v1/user/space/authorize-url',
query_string={'redirect_uri': 'https://evil.example/auth/space/callback'},
headers={'Host': 'evil.example'},
)
assert (await wrong_origin.get_json())['code'] == 1
assert (await wrong_path.get_json())['code'] == 1
assert (await forged_origin.get_json())['code'] == 1
assert (await forged_host.get_json())['code'] == 1
@pytest.mark.asyncio
async def test_explicit_server_side_webui_origin_supports_split_dev_server(space_oauth_api):
application, client = space_oauth_api
application.instance_config.data['api'] = {'webui_url': 'http://localhost:5173'}
response = await client.get(
'/api/v1/user/space/authorize-url',
query_string={'redirect_uri': 'http://localhost:5173/auth/space/callback'},
headers={'Origin': 'https://irrelevant.example'},
)
assert response.status_code == 200
assert (await response.get_json())['code'] == 0
@pytest.mark.asyncio
async def test_server_side_webhook_origin_supports_bundled_ui(space_oauth_api):
application, client = space_oauth_api
application.instance_config.data['api'] = {
'webui_url': '',
'webhook_prefix': 'https://langbot.example/base/path',
}
response = await client.get(
'/api/v1/user/space/authorize-url',
query_string={'redirect_uri': 'https://langbot.example/auth/space/callback'},
headers={'Host': 'attacker.example'},
)
assert response.status_code == 200
assert (await response.get_json())['code'] == 0
@pytest.mark.asyncio
async def test_login_callback_requires_and_consumes_server_state(space_oauth_api):
application, client = space_oauth_api
missing = await client.post('/api/v1/user/space/callback', json={'code': 'oauth-code'})
response = await client.post(
'/api/v1/user/space/callback',
json={'code': 'oauth-code', 'state': 'opaque-login-state'},
)
assert (await missing.get_json())['code'] == 1
assert response.status_code == 200
assert (await response.get_json())['data']['token'] == 'space-login-token'
application.user_service.consume_space_oauth_state_details.assert_awaited_once_with('opaque-login-state', 'login')
application.space_service.exchange_oauth_code.assert_awaited_once_with('oauth-code')
@pytest.mark.asyncio
async def test_login_callback_launch_state_selects_asserted_workspace(space_oauth_api):
application, client = space_oauth_api
application.user_service.consume_space_oauth_state_details.reset_mock()
application.user_service.consume_space_oauth_state_details.return_value = SimpleNamespace(
launch_workspace_uuid=WORKSPACE_UUID
)
response = await client.post(
'/api/v1/user/space/callback',
json={'code': 'oauth-code', 'state': 'opaque-login-state'},
)
assert response.status_code == 200
data = (await response.get_json())['data']
assert data['token'] == 'space-login-token'
assert data['workspace_uuid'] == WORKSPACE_UUID
application.workspace_collaboration_service.resolve_account_workspace.assert_awaited_with(
'account-a',
WORKSPACE_UUID,
)
@pytest.mark.asyncio
async def test_space_credits_are_resolved_from_workspace_owner(space_oauth_api):
application, client = space_oauth_api
application.user_service.get_workspace_owner = AsyncMock(
return_value=SimpleNamespace(user='owner@example.com', space_account_uuid='space-owner')
)
application.space_service.get_credits = AsyncMock(return_value=25000)
response = await client.get(
'/api/v1/user/space-credits',
headers={'Authorization': 'Bearer account-token', 'X-Workspace-UUID': WORKSPACE_UUID},
)
assert response.status_code == 200
assert (await response.get_json())['data'] == {
'credits': 25000,
'owner_space_bound': True,
'is_workspace_owner': True,
}
application.space_service.get_credits.assert_awaited_once_with('owner@example.com')
@pytest.mark.asyncio
async def test_bind_callback_uses_opaque_state_and_never_treats_it_as_jwt(space_oauth_api):
application, client = space_oauth_api
application.user_service.consume_space_oauth_state.reset_mock()
application.user_service.consume_space_oauth_state.side_effect = [
ValueError('invalid state'),
SimpleNamespace(
uuid='account-a',
user='owner@example.com',
account_type='local',
),
]
rejected = await client.post(
'/api/v1/user/bind-space',
json={'code': 'attacker-code', 'state': 'jwt.must-not-be-used'},
)
response = await client.post(
'/api/v1/user/bind-space',
json={'code': 'oauth-code', 'state': 'opaque-bind-state'},
)
assert rejected.status_code == 401
assert response.status_code == 200
assert (await response.get_json())['data']['token'] == 'rotated-account-token'
application.user_service.verify_jwt_token.assert_not_awaited()
application.user_service.bind_space_account.assert_awaited_once_with('owner@example.com', 'oauth-code')
@pytest.mark.asyncio
async def test_direct_launch_assertion_does_not_consume_normal_oauth_state(space_oauth_api):
application, client = space_oauth_api
application.user_service.consume_space_oauth_state.reset_mock()
application.space_service.exchange_oauth_code.reset_mock()
response = await client.post(
'/api/v1/user/space/callback',
json={
'state': 'space-generated-state-is-not-oauth-state',
'workspace_uuid': WORKSPACE_UUID,
'launch_assertion': 'signed-launch-token',
},
)
assert response.status_code == 200
data = (await response.get_json())['data']
assert data['token'] == 'rotated-account-token'
assert data['workspace_uuid'] == WORKSPACE_UUID
application.space_launch_service.consume_assertion.assert_awaited_once_with(
'signed-launch-token',
expected_workspace_uuid=WORKSPACE_UUID,
)
application.user_service.consume_space_oauth_state.assert_not_awaited()
application.space_service.exchange_oauth_code.assert_not_awaited()
tests/integration/api/test_workspaces.py
+616
View File
@@ -0,0 +1,616 @@
from __future__ import annotations
import json
import logging
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import sqlalchemy
import jwt
from quart import Quart
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.api.http.controller.groups.workspaces import (
InvitationsRouterGroup,
WorkspacesRouterGroup,
)
from langbot.pkg.api.http.controller.groups.system import SystemRouterGroup
from langbot.pkg.api.http.controller.groups.apikeys import ApiKeysRouterGroup
from langbot.pkg.api.http.controller.groups.user import UserRouterGroup
from langbot.pkg.api.http.service.apikey import ApiKeyService
from langbot.pkg.api.http.service.user import ControlPlaneDirectoryRequiredError, UserService
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.metadata import WorkspaceMetadata
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.entity.persistence.workspace import (
Workspace,
WorkspaceExecutionState,
WorkspaceInvitation,
WorkspaceMembership,
)
from langbot.pkg.persistence.mgr import PersistenceManager
from langbot.pkg.workspace.collaboration import WorkspaceCollaborationService
from langbot.pkg.workspace.service import WorkspaceService
from langbot.pkg.workspace.policy import CloudWorkspacePolicy
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
@pytest.fixture
async def workspace_api(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-api.db"}')
async with engine.begin() as connection:
await connection.run_sync(Base.metadata.create_all)
application = SimpleNamespace()
application.persistence_mgr = PersistenceManager(application)
application.persistence_mgr.db = SimpleNamespace(get_engine=lambda: engine)
application.instance_config = SimpleNamespace(
data={
'system': {
'jwt': {'secret': 'workspace-api-secret', 'expire': 3600},
'allow_modify_login_info': True,
},
'api': {'global_api_key': '', 'webui_url': 'https://langbot.example'},
}
)
application.logger = logging.getLogger('workspace-api-test')
application.workspace_service = WorkspaceService(
application,
instance_uuid='instance-workspace-api',
)
await application.workspace_service.ensure_singleton_workspace()
application.workspace_collaboration_service = WorkspaceCollaborationService(
application,
application.workspace_service,
)
application.user_service = UserService(application)
application.apikey_service = ApiKeyService(application)
quart_app = Quart(__name__)
await WorkspacesRouterGroup(application, quart_app).initialize()
await InvitationsRouterGroup(application, quart_app).initialize()
await ApiKeysRouterGroup(application, quart_app).initialize()
await UserRouterGroup(application, quart_app).initialize()
await SystemRouterGroup(application, quart_app).initialize()
client = quart_app.test_client()
init_response = await client.post(
'/api/v1/user/init',
json={'user': 'owner@example.com', 'password': 'owner-password'},
)
assert init_response.status_code == 200
auth_response = await client.post(
'/api/v1/user/auth',
json={'user': 'owner@example.com', 'password': 'owner-password'},
)
assert auth_response.status_code == 200
owner_token = (await auth_response.get_json())['data']['token']
yield application, client, engine, owner_token
await engine.dispose()
def _auth(token: str, workspace_uuid: str | None = None) -> dict[str, str]:
headers = {'Authorization': f'Bearer {token}'}
if workspace_uuid is not None:
headers['X-Workspace-Id'] = workspace_uuid
return headers
async def test_account_bootstrap_uses_the_account_resolved_from_the_token(workspace_api):
application, client, _, owner_token = workspace_api
account = await application.user_service.get_authenticated_account(owner_token)
application.user_service.get_authenticated_account = AsyncMock(return_value=account)
application.user_service.get_user_by_email = AsyncMock(return_value=None)
response = await client.get('/api/v1/workspaces/bootstrap', headers=_auth(owner_token))
assert response.status_code == 200
application.user_service.get_user_by_email.assert_not_awaited()
async def test_fresh_sqlite_login_returns_current_workspace_and_user_info(workspace_api):
_, client, _, owner_token = workspace_api
current_response = await client.get('/api/v1/workspaces/current', headers=_auth(owner_token))
assert current_response.status_code == 200
current = (await current_response.get_json())['data']
assert current['workspace']['uuid']
assert current['membership']['email'] == 'owner@example.com'
assert current['membership']['role'] == 'owner'
info_response = await client.get('/api/v1/user/info', headers=_auth(owner_token))
assert info_response.status_code == 200
info = (await info_response.get_json())['data']
assert info['account_uuid'] == current['membership']['account_uuid']
assert info['user'] == 'owner@example.com'
async def test_user_info_uses_the_account_resolved_from_the_token(workspace_api):
application, client, _, owner_token = workspace_api
account = await application.user_service.get_authenticated_account(owner_token)
application.user_service.get_authenticated_account = AsyncMock(return_value=account)
application.user_service.get_user_by_email = AsyncMock(return_value=None)
response = await client.get('/api/v1/user/info', headers=_auth(owner_token))
assert response.status_code == 200
info = (await response.get_json())['data']
assert info['account_uuid'] == account.uuid
assert info['user'] == account.user
application.user_service.get_user_by_email.assert_not_awaited()
async def test_authenticated_system_info_reads_workspace_wizard_metadata(workspace_api):
application, client, _, owner_token = workspace_api
current_response = await client.get('/api/v1/workspaces/current', headers=_auth(owner_token))
workspace_uuid = (await current_response.get_json())['data']['workspace']['uuid']
progress = {'step': 3, 'selected_adapter': 'telegram'}
await application.persistence_mgr.execute_async(
sqlalchemy.insert(WorkspaceMetadata),
[
{
'workspace_uuid': workspace_uuid,
'key': 'wizard_status',
'value': 'completed',
},
{
'workspace_uuid': workspace_uuid,
'key': 'wizard_progress',
'value': json.dumps(progress),
},
],
)
response = await client.get(
'/api/v1/system/info',
headers=_auth(owner_token, workspace_uuid),
)
assert response.status_code == 200
data = (await response.get_json())['data']
assert data['wizard_status'] == 'completed'
assert data['wizard_progress'] == progress
async def test_owner_invites_second_account_and_secret_is_not_persisted(workspace_api):
application, client, engine, owner_token = workspace_api
current_response = await client.get('/api/v1/workspaces/current', headers=_auth(owner_token))
assert current_response.status_code == 200
current = (await current_response.get_json())['data']
workspace_uuid = current['workspace']['uuid']
assert current['membership']['role'] == 'owner'
assert 'member.invite' in current['permissions']
invite_response = await client.post(
f'/api/v1/workspaces/{workspace_uuid}/invitations',
headers=_auth(owner_token, workspace_uuid),
json={'email': 'member@example.com', 'role': 'viewer'},
)
assert invite_response.status_code == 200
invite_data = (await invite_response.get_json())['data']
invitation_token = invite_data['token']
assert invitation_token.startswith('lbi_')
assert invite_data['link'] == f'https://langbot.example/invitations/accept#token={invitation_token}'
assert invite_data['delivery'] == {'status': 'link_only', 'provider': None}
assert 'token_hash' not in invite_data['invitation']
async with engine.connect() as connection:
persisted_token_hash = await connection.scalar(
sqlalchemy.select(WorkspaceInvitation.token_hash).where(
WorkspaceInvitation.uuid == invite_data['invitation']['uuid']
)
)
assert persisted_token_hash is not None
assert persisted_token_hash != invitation_token
inspect_response = await client.post(
'/api/v1/invitations/inspect',
json={'token': invitation_token},
)
assert inspect_response.status_code == 200
inspected = (await inspect_response.get_json())['data']
assert inspected['workspace']['uuid'] == workspace_uuid
assert inspected['invitation']['normalized_email'] == 'member@example.com'
accept_response = await client.post(
'/api/v1/invitations/accept',
json={
'token': invitation_token,
'registration': {
'email': 'member@example.com',
'password': 'member-password',
},
},
)
assert accept_response.status_code == 200
member_registration = (await accept_response.get_json())['data']
assert member_registration == {'workspace_uuid': workspace_uuid, 'login_required': True}
member_login_response = await client.post(
'/api/v1/user/auth',
json={'user': 'member@example.com', 'password': 'member-password'},
)
assert member_login_response.status_code == 200
member_token = (await member_login_response.get_json())['data']['token']
reused_response = await client.post(
'/api/v1/invitations/accept',
json={
'token': invitation_token,
'registration': {
'email': 'member@example.com',
'password': 'member-password',
},
},
)
assert reused_response.status_code == 400
assert (await reused_response.get_json())['code'] == 'invitation_used'
member_current_response = await client.get(
'/api/v1/workspaces/current',
headers=_auth(member_token, workspace_uuid),
)
assert member_current_response.status_code == 200
member_current = (await member_current_response.get_json())['data']
assert member_current['membership']['role'] == 'viewer'
assert 'member.invite' not in member_current['permissions']
forbidden_invite = await client.post(
f'/api/v1/workspaces/{workspace_uuid}/invitations',
headers=_auth(member_token, workspace_uuid),
json={'email': 'third@example.com', 'role': 'viewer'},
)
assert forbidden_invite.status_code == 403
assert (await forbidden_invite.get_json())['code'] == 'permission_denied'
async def test_oss_invitation_accept_requires_logout_before_registration(workspace_api):
_, client, _, owner_token = workspace_api
response = await client.post(
'/api/v1/invitations/accept',
headers={'Authorization': f'Bearer {owner_token}'},
json={'token': 'lbi_pending-invitation'},
)
assert response.status_code == 409
assert (await response.get_json())['code'] == 'invitation_logout_required'
async def test_invalid_bearer_on_cloud_invitation_is_authentication_failure(workspace_api):
application, client, _, _ = workspace_api
application.deployment = SimpleNamespace(mode='cloud')
response = await client.post(
'/api/v1/invitations/accept',
headers={'Authorization': 'Bearer definitely-not-a-jwt'},
json={'token': 'lbi_not-a-real-invitation'},
)
assert response.status_code == 401
assert await response.get_json() == {
'code': 'invalid_authentication',
'msg': 'Invalid authentication credentials',
}
async def test_workspace_selector_and_path_cannot_escape_membership(workspace_api):
_, client, _, owner_token = workspace_api
unknown_uuid = '00000000-0000-0000-0000-000000000099'
selector_response = await client.get(
'/api/v1/workspaces/current',
headers=_auth(owner_token, unknown_uuid),
)
assert selector_response.status_code == 404
assert (await selector_response.get_json())['code'] == 'resource_not_found'
path_response = await client.get(
f'/api/v1/workspaces/{unknown_uuid}',
headers=_auth(owner_token),
)
assert path_response.status_code == 404
assert (await path_response.get_json())['code'] == 'resource_not_found'
async def test_oss_rejects_second_workspace(workspace_api):
_, client, _, owner_token = workspace_api
response = await client.post('/api/v1/workspaces', headers=_auth(owner_token), json={'name': 'Second'})
assert response.status_code == 403
assert (await response.get_json())['code'] == 'edition_limit'
async def test_jwt_uses_account_uuid_and_disabled_account_is_rejected(workspace_api):
_, client, engine, owner_token = workspace_api
payload = jwt.decode(
owner_token,
'workspace-api-secret',
algorithms=['HS256'],
audience='langbot-instance:instance-workspace-api',
issuer='langbot-core',
)
assert payload['sub']
assert payload['sub'] != payload['user']
async with engine.begin() as connection:
await connection.execute(sqlalchemy.update(User).where(User.uuid == payload['sub']).values(status='disabled'))
response = await client.get('/api/v1/workspaces/current', headers=_auth(owner_token))
assert response.status_code == 401
assert (await response.get_json())['code'] == 'invalid_authentication'
async def test_api_key_secret_is_one_time_and_viewer_cannot_manage_keys(workspace_api):
application, client, _engine, owner_token = workspace_api
current_response = await client.get('/api/v1/workspaces/current', headers=_auth(owner_token))
workspace_uuid = (await current_response.get_json())['data']['workspace']['uuid']
create_response = await client.post(
'/api/v1/apikeys',
headers=_auth(owner_token, workspace_uuid),
json={'name': 'E2E automation', 'scopes': ['resource.view']},
)
assert create_response.status_code == 200
created = (await create_response.get_json())['data']['key']
assert created['key'].startswith('lbk_')
assert created['secret_available'] is True
assert 'key_hash' not in created
list_response = await client.get('/api/v1/apikeys', headers=_auth(owner_token, workspace_uuid))
listed = (await list_response.get_json())['data']['keys']
assert len(listed) == 1
assert 'key' not in listed[0]
assert 'key_hash' not in listed[0]
assert listed[0]['secret_available'] is False
identity = await application.apikey_service.authenticate_api_key(created['key'])
assert identity is not None
assert identity.workspace_uuid == workspace_uuid
assert identity.permissions == frozenset({'resource.view'})
invite_response = await client.post(
f'/api/v1/workspaces/{workspace_uuid}/invitations',
headers=_auth(owner_token, workspace_uuid),
json={'email': 'viewer@example.com', 'role': 'viewer'},
)
invitation_token = (await invite_response.get_json())['data']['token']
accept_response = await client.post(
'/api/v1/invitations/accept',
json={
'token': invitation_token,
'registration': {'email': 'viewer@example.com', 'password': 'viewer-password'},
},
)
assert accept_response.status_code == 200
assert (await accept_response.get_json())['data']['login_required'] is True
login_response = await client.post(
'/api/v1/user/auth',
json={'user': 'viewer@example.com', 'password': 'viewer-password'},
)
assert login_response.status_code == 200
viewer_token = (await login_response.get_json())['data']['token']
forbidden = await client.post(
'/api/v1/apikeys',
headers=_auth(viewer_token, workspace_uuid),
json={'name': 'forbidden'},
)
assert forbidden.status_code == 403
assert (await forbidden.get_json())['code'] == 'permission_denied'
async def test_cloud_projection_is_selected_explicitly_and_collaboration_runs_in_core(
workspace_api,
):
application, client, engine, owner_token = workspace_api
application.deployment = SimpleNamespace(mode='cloud')
owner_uuid = jwt.decode(
owner_token,
'workspace-api-secret',
algorithms=['HS256'],
audience='langbot-instance:instance-workspace-api',
issuer='langbot-core',
)['sub']
cloud_workspace_uuid = '00000000-0000-0000-0000-000000000777'
async with engine.begin() as connection:
await connection.execute(
sqlalchemy.insert(Workspace).values(
uuid=cloud_workspace_uuid,
instance_uuid='instance-workspace-api',
name='Cloud Team',
slug='cloud-team',
type='team',
status='active',
source='cloud_projection',
projection_revision=12,
)
)
await connection.execute(
sqlalchemy.insert(WorkspaceExecutionState).values(
workspace_uuid=cloud_workspace_uuid,
instance_uuid='instance-workspace-api',
active_generation=12,
state='active',
write_fenced=False,
source='cloud',
desired_state_revision=12,
)
)
await connection.execute(
sqlalchemy.insert(WorkspaceMembership).values(
uuid='00000000-0000-0000-0000-000000000778',
workspace_uuid=cloud_workspace_uuid,
account_uuid=owner_uuid,
role='owner',
status='active',
projection_revision=12,
)
)
policy = CloudWorkspacePolicy()
application.workspace_service.policy = policy
application.workspace_collaboration_service.policy = policy
with pytest.raises(ControlPlaneDirectoryRequiredError):
await application.user_service.create_initial_account(
'forbidden-cloud-local@example.com',
'password',
)
omitted = await client.get('/api/v1/workspaces/current', headers=_auth(owner_token))
assert omitted.status_code == 404
refreshed_token = await client.get('/api/v1/user/check-token', headers=_auth(owner_token))
assert refreshed_token.status_code == 200
assert (await refreshed_token.get_json())['data']['token']
bootstrap_response = await client.get(
'/api/v1/workspaces/bootstrap',
headers=_auth(owner_token),
)
assert bootstrap_response.status_code == 200
bootstrap = (await bootstrap_response.get_json())['data']
singleton_uuid = (await application.workspace_service.get_singleton_workspace()).uuid
workspace_uuids = [item['workspace']['uuid'] for item in bootstrap['workspaces']]
assert set(workspace_uuids) == {singleton_uuid, cloud_workspace_uuid}
repeated = await client.get('/api/v1/workspaces/bootstrap', headers=_auth(owner_token))
assert [item['workspace']['uuid'] for item in (await repeated.get_json())['data']['workspaces']] == workspace_uuids
by_uuid = {item['workspace']['uuid']: item for item in bootstrap['workspaces']}
assert by_uuid[singleton_uuid]['membership']['account_uuid'] == owner_uuid
assert by_uuid[singleton_uuid]['membership']['email'] == 'owner@example.com'
assert by_uuid[singleton_uuid]['permissions']
assert by_uuid[cloud_workspace_uuid]['placement_generation'] == 12
list_response = await client.get(
'/api/v1/workspaces',
headers=_auth(owner_token, singleton_uuid),
)
assert list_response.status_code == 200
assert {workspace['uuid'] for workspace in (await list_response.get_json())['data']['workspaces']} == {
singleton_uuid,
cloud_workspace_uuid,
}
class PlanResolver:
async def resolve(self, workspace_uuid: str, *, minimum_revision: int = 0):
from langbot.pkg.cloud.entitlements import EntitlementSnapshot
assert workspace_uuid == cloud_workspace_uuid
return EntitlementSnapshot(
instance_uuid='instance-workspace-api',
workspace_uuid=workspace_uuid,
entitlement_revision=max(12, minimum_revision),
status='active',
not_before=1,
expires_at=4102444800,
plan_name='free',
)
application.entitlement_resolver = PlanResolver()
bootstrap_with_plans = await client.get('/api/v1/workspaces/bootstrap', headers=_auth(owner_token))
bootstrap_by_uuid = {
item['workspace']['uuid']: item for item in (await bootstrap_with_plans.get_json())['data']['workspaces']
}
assert bootstrap_by_uuid[cloud_workspace_uuid]['plan_name'] == 'free'
assert bootstrap_by_uuid[singleton_uuid]['plan_name'] is None
current_response = await client.get(
'/api/v1/workspaces/current',
headers=_auth(owner_token, cloud_workspace_uuid),
)
assert current_response.status_code == 200
current = (await current_response.get_json())['data']
assert current['workspace']['uuid'] == cloud_workspace_uuid
assert current['workspace']['source'] == 'cloud_projection'
assert current['placement_generation'] == 12
assert current['plan_name'] == 'free'
create_workspace = await client.post(
'/api/v1/workspaces',
headers=_auth(owner_token, cloud_workspace_uuid),
json={'name': 'Not in Core'},
)
assert create_workspace.status_code == 409
assert (await create_workspace.get_json())['code'] == 'control_plane_required'
create_invitation = await client.post(
f'/api/v1/workspaces/{cloud_workspace_uuid}/invitations',
headers=_auth(owner_token, cloud_workspace_uuid),
json={'email': 'member@example.com', 'role': 'viewer'},
)
assert create_invitation.status_code == 200
created_invitation = (await create_invitation.get_json())['data']
assert created_invitation['invitation']['workspace_uuid'] == cloud_workspace_uuid
assert created_invitation['link'].startswith('https://langbot.example/invitations/accept#token=lbi_')
assert created_invitation['delivery'] == {'status': 'link_only', 'provider': None}
accept_response = await client.post(
'/api/v1/invitations/accept',
headers=_auth(owner_token, cloud_workspace_uuid),
json={'token': created_invitation['token']},
)
assert accept_response.status_code == 400
assert (await accept_response.get_json())['code'] == 'invitation_email_mismatch'
registration_response = await client.post(
'/api/v1/invitations/accept',
json={
'token': created_invitation['token'],
'registration': {'email': 'member@example.com', 'password': 'member-password'},
},
)
assert registration_response.status_code == 401
assert (await registration_response.get_json())['code'] == 'account_exists_login_required'
async def test_account_bootstrap_does_not_disclose_non_member_workspaces(workspace_api):
application, client, engine, owner_token = workspace_api
foreign_workspace_uuid = '00000000-0000-0000-0000-000000000880'
async with engine.begin() as connection:
await connection.execute(
sqlalchemy.insert(Workspace).values(
uuid=foreign_workspace_uuid,
instance_uuid='instance-workspace-api',
name='Foreign Team',
slug='foreign-team',
type='team',
status='active',
source='cloud_projection',
projection_revision=1,
)
)
await connection.execute(
sqlalchemy.insert(WorkspaceExecutionState).values(
workspace_uuid=foreign_workspace_uuid,
instance_uuid='instance-workspace-api',
active_generation=1,
state='active',
write_fenced=False,
source='cloud',
desired_state_revision=1,
)
)
policy = CloudWorkspacePolicy()
application.workspace_service.policy = policy
application.workspace_collaboration_service.policy = policy
response = await client.get('/api/v1/workspaces/bootstrap', headers=_auth(owner_token))
assert response.status_code == 200
workspace_uuids = {item['workspace']['uuid'] for item in (await response.get_json())['data']['workspaces']}
assert foreign_workspace_uuid not in workspace_uuids
current = await client.get(
'/api/v1/workspaces/current',
headers=_auth(owner_token, foreign_workspace_uuid),
)
assert current.status_code == 404
assert (await current.get_json())['code'] == 'resource_not_found'
tests/integration/persistence/resource_migration_support.py
+253
View File
@@ -0,0 +1,253 @@
from __future__ import annotations
import datetime
import sqlalchemy as sa
TENANT_TABLES = (
'api_keys',
'bots',
'bot_admins',
'binary_storages',
'mcp_servers',
'model_providers',
'llm_models',
'embedding_models',
'rerank_models',
'legacy_pipelines',
'pipeline_run_records',
'plugin_settings',
'knowledge_bases',
'knowledge_base_files',
'knowledge_base_chunks',
'webhooks',
'monitoring_messages',
'monitoring_llm_calls',
'monitoring_tool_calls',
'monitoring_sessions',
'monitoring_errors',
'monitoring_embedding_calls',
'monitoring_feedback',
)
def _uuid_table(metadata: sa.MetaData, name: str, *columns: sa.Column) -> sa.Table:
return sa.Table(name, metadata, sa.Column('uuid', sa.String(255), primary_key=True), *columns)
async def create_legacy_resource_schema(engine, *, instance_uuid: str) -> None:
"""Create the smallest representative pre-0010 schema with one row/table."""
metadata = sa.MetaData()
system_metadata = sa.Table(
'metadata',
metadata,
sa.Column('key', sa.String(255), primary_key=True),
sa.Column('value', sa.String(255)),
)
users = sa.Table(
'users',
metadata,
sa.Column('id', sa.Integer, primary_key=True),
sa.Column('user', sa.String(255), nullable=False),
sa.Column('password', sa.String(255), nullable=False),
)
api_keys = sa.Table(
'api_keys',
metadata,
sa.Column('id', sa.Integer, primary_key=True, autoincrement=True),
sa.Column('name', sa.String(255), nullable=False),
sa.Column('key', sa.String(255), nullable=False, unique=True),
)
bots = _uuid_table(
metadata,
'bots',
sa.Column('name', sa.String(255), nullable=False),
sa.Column('updated_at', sa.DateTime, nullable=False),
)
bot_admins = sa.Table(
'bot_admins',
metadata,
sa.Column('id', sa.Integer, primary_key=True, autoincrement=True),
sa.Column('bot_uuid', sa.String(255), nullable=False),
sa.Column('launcher_type', sa.String(64), nullable=False),
sa.Column('launcher_id', sa.String(255), nullable=False),
sa.UniqueConstraint('bot_uuid', 'launcher_type', 'launcher_id', name='uq_bot_admin'),
)
binary_storages = sa.Table(
'binary_storages',
metadata,
sa.Column('unique_key', sa.String(255), primary_key=True),
sa.Column('key', sa.String(255), nullable=False),
sa.Column('owner_type', sa.String(255), nullable=False),
sa.Column('owner', sa.String(255), nullable=False),
)
mcp_servers = _uuid_table(
metadata,
'mcp_servers',
sa.Column('name', sa.String(255), nullable=False),
sa.Column('enable', sa.Boolean, nullable=False),
sa.Column('updated_at', sa.DateTime, nullable=False),
)
model_providers = _uuid_table(
metadata,
'model_providers',
sa.Column('name', sa.String(255), nullable=False),
sa.Column('requester', sa.String(255), nullable=False),
)
llm_models = _uuid_table(
metadata,
'llm_models',
sa.Column('name', sa.String(255), nullable=False),
sa.Column('provider_uuid', sa.String(255), nullable=False),
)
embedding_models = _uuid_table(
metadata,
'embedding_models',
sa.Column('name', sa.String(255), nullable=False),
sa.Column('provider_uuid', sa.String(255), nullable=False),
)
rerank_models = _uuid_table(
metadata,
'rerank_models',
sa.Column('name', sa.String(255), nullable=False),
sa.Column('provider_uuid', sa.String(255), nullable=False),
)
legacy_pipelines = _uuid_table(
metadata,
'legacy_pipelines',
sa.Column('name', sa.String(255), nullable=False),
sa.Column('is_default', sa.Boolean, nullable=False),
sa.Column('updated_at', sa.DateTime, nullable=False),
)
pipeline_run_records = _uuid_table(
metadata,
'pipeline_run_records',
sa.Column('pipeline_uuid', sa.String(255), nullable=False),
sa.Column('created_at', sa.DateTime, nullable=False),
)
plugin_settings = sa.Table(
'plugin_settings',
metadata,
sa.Column('plugin_author', sa.String(255), primary_key=True),
sa.Column('plugin_name', sa.String(255), primary_key=True),
sa.Column('enabled', sa.Boolean, nullable=False),
)
knowledge_bases = _uuid_table(
metadata,
'knowledge_bases',
sa.Column('name', sa.String(255), nullable=False),
sa.Column('collection_id', sa.String(255), nullable=True),
)
knowledge_base_files = _uuid_table(
metadata,
'knowledge_base_files',
sa.Column('kb_id', sa.String(255), nullable=True),
)
knowledge_base_chunks = _uuid_table(
metadata,
'knowledge_base_chunks',
sa.Column('file_id', sa.String(255), nullable=True),
)
webhooks = sa.Table(
'webhooks',
metadata,
sa.Column('id', sa.Integer, primary_key=True, autoincrement=True),
sa.Column('name', sa.String(255), nullable=False),
sa.Column('enabled', sa.Boolean, nullable=False),
sa.Column('created_at', sa.DateTime, nullable=False),
)
monitoring_tables: dict[str, sa.Table] = {}
for table_name in (
'monitoring_messages',
'monitoring_llm_calls',
'monitoring_tool_calls',
'monitoring_errors',
'monitoring_embedding_calls',
):
monitoring_tables[table_name] = sa.Table(
table_name,
metadata,
sa.Column('id', sa.String(255), primary_key=True),
sa.Column('timestamp', sa.DateTime, nullable=False),
sa.Column('session_id', sa.String(255), nullable=True),
sa.Column('message_id', sa.String(255), nullable=True),
)
monitoring_tables['monitoring_sessions'] = sa.Table(
'monitoring_sessions',
metadata,
sa.Column('session_id', sa.String(255), primary_key=True),
sa.Column('bot_id', sa.String(255), nullable=False),
sa.Column('last_activity', sa.DateTime, nullable=False),
sa.Column('is_active', sa.Boolean, nullable=False),
)
monitoring_tables['monitoring_feedback'] = sa.Table(
'monitoring_feedback',
metadata,
sa.Column('id', sa.String(255), primary_key=True),
sa.Column('feedback_id', sa.String(255), nullable=False, unique=True),
sa.Column('timestamp', sa.DateTime, nullable=False),
sa.Column('session_id', sa.String(255), nullable=True),
sa.Column('message_id', sa.String(255), nullable=True),
)
now = datetime.datetime(2026, 1, 1)
async with engine.begin() as conn:
await conn.run_sync(metadata.create_all)
await conn.execute(
system_metadata.insert(),
[
{'key': 'database_version', 'value': '25'},
{'key': 'instance_uuid', 'value': instance_uuid},
{'key': 'wizard_status', 'value': 'completed'},
{'key': 'wizard_progress', 'value': '3'},
{'key': 'rag_plugin_migration_needed', 'value': 'true'},
],
)
await conn.execute(users.insert().values(user='Owner@Example.COM', password='hash'))
await conn.execute(api_keys.insert().values(name='legacy', key='lbk_legacy-secret'))
await conn.execute(bots.insert().values(uuid='bot-1', name='bot', updated_at=now))
await conn.execute(bot_admins.insert().values(bot_uuid='bot-1', launcher_type='person', launcher_id='owner'))
await conn.execute(
binary_storages.insert().values(unique_key='plugin:demo:key', key='key', owner_type='plugin', owner='demo')
)
await conn.execute(mcp_servers.insert().values(uuid='mcp-1', name='shared-name', enable=True, updated_at=now))
await conn.execute(model_providers.insert().values(uuid='provider-1', name='provider', requester='openai'))
for table in (llm_models, embedding_models, rerank_models):
await conn.execute(table.insert().values(uuid=f'{table.name}-1', name='model', provider_uuid='provider-1'))
await conn.execute(
legacy_pipelines.insert().values(uuid='pipeline-1', name='pipeline', is_default=True, updated_at=now)
)
await conn.execute(
pipeline_run_records.insert().values(uuid='run-1', pipeline_uuid='pipeline-1', created_at=now)
)
await conn.execute(plugin_settings.insert().values(plugin_author='author', plugin_name='plugin', enabled=True))
await conn.execute(knowledge_bases.insert().values(uuid='kb-1', name='knowledge', collection_id='collection-1'))
await conn.execute(knowledge_base_files.insert().values(uuid='file-1', kb_id='kb-1'))
await conn.execute(knowledge_base_chunks.insert().values(uuid='chunk-1', file_id='file-1'))
await conn.execute(webhooks.insert().values(name='hook', enabled=True, created_at=now))
for table_name, table in monitoring_tables.items():
if table_name == 'monitoring_sessions':
values = {
'session_id': 'session-1',
'bot_id': 'bot-1',
'last_activity': now,
'is_active': True,
}
elif table_name == 'monitoring_feedback':
values = {
'id': 'feedback-row-1',
'feedback_id': 'feedback-1',
'timestamp': now,
'session_id': 'session-1',
'message_id': 'message-1',
}
else:
values = {
'id': f'{table_name}-1',
'timestamp': now,
'session_id': 'session-1',
'message_id': 'message-1',
}
await conn.execute(table.insert().values(**values))
tests/integration/persistence/test_migrations.py
+41
View File
@@ -10,10 +10,12 @@ Run: uv run pytest tests/integration/persistence/test_migrations.py -q
from __future__ import annotations
import pytest
import sqlalchemy
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,
@@ -149,6 +151,45 @@ class TestSQLiteMigrationUpgrade:
rev2 = await get_alembic_current(sqlite_engine)
assert rev2 == rev1, f'Expected {rev1}, got {rev2}'
@pytest.mark.asyncio
async def test_upgrade_from_0012_adds_knowledge_base_embedding_dimension(self, sqlite_engine):
"""The PostgreSQL pgvector revision also evolves the OSS ORM schema."""
async with sqlite_engine.begin() as conn:
await conn.exec_driver_sql(
'CREATE TABLE knowledge_bases ('
'uuid VARCHAR(255) PRIMARY KEY, workspace_uuid VARCHAR(36) NOT NULL, name VARCHAR(255) NOT NULL)'
)
await run_alembic_stamp(sqlite_engine, '0012_plugin_identity')
await run_alembic_upgrade(sqlite_engine, 'head')
async with sqlite_engine.connect() as conn:
columns = await conn.run_sync(
lambda sync_conn: {
item['name'] for item in sqlalchemy.inspect(sync_conn).get_columns('knowledge_bases')
}
)
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
+1962 -8
View File
File diff suppressed because it is too large Load Diff
tests/integration/persistence/test_pgvector_postgres.py
+553
View File
@@ -0,0 +1,553 @@
"""Real PostgreSQL/pgvector tenant isolation and CRUD verification."""
from __future__ import annotations
import logging
import os
import uuid
from types import SimpleNamespace
import pytest
import sqlalchemy as sa
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.rag import KnowledgeBase
from langbot.pkg.entity.persistence.workspace import Workspace
from langbot.pkg.persistence.alembic_runner import get_alembic_current, run_alembic_stamp, run_alembic_upgrade
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
from langbot.pkg.utils import constants
from langbot.pkg.vector.vdbs.pgvector_db import PgVectorDatabase, PgVectorEntry, PgVectorScope
pytestmark = [pytest.mark.integration, pytest.mark.slow, pytest.mark.asyncio]
def _application(postgres_url: str, logger_name: str) -> SimpleNamespace:
url = sa.engine.make_url(postgres_url)
return SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'use': 'postgresql',
'postgresql': {
'host': url.host,
'port': url.port,
'user': url.username,
'password': url.password,
'database': url.database,
},
}
}
),
logger=logging.getLogger(logger_name),
)
def _restore_postgres_registry(monkeypatch) -> None:
from langbot.pkg.persistence import mgr as persistence_mgr_module
from langbot.pkg.persistence.databases.postgresql import PostgreSQLDatabaseManager
monkeypatch.setattr(
persistence_mgr_module.database,
'preregistered_managers',
[PostgreSQLDatabaseManager],
)
@pytest.fixture
def postgres_url() -> str:
url = os.environ.get('TEST_POSTGRES_URL')
if not url:
pytest.skip('TEST_POSTGRES_URL not set')
return url
@pytest.fixture
async def postgres_engine(postgres_url: str):
engine = create_async_engine(postgres_url, isolation_level='AUTOCOMMIT')
yield engine
await engine.dispose()
@pytest.fixture
async def clean_database(postgres_engine: AsyncEngine):
async def clean() -> None:
async with postgres_engine.begin() as conn:
await conn.execute(text('DROP TABLE IF EXISTS langbot_vectors_legacy_0013 CASCADE'))
await conn.execute(text('DROP TABLE IF EXISTS langbot_vectors CASCADE'))
await conn.run_sync(Base.metadata.drop_all)
await conn.execute(text('DROP TABLE IF EXISTS alembic_version'))
await clean()
yield
await clean()
async def test_legacy_upgrade_temporarily_suspends_and_restores_source_rls_for_unprivileged_owner(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
) -> None:
workspace_uuid = '30000000-0000-0000-0000-000000000303'
knowledge_base_uuid = 'legacy-knowledge-base'
migrator_role = f'lb_vector_migrator_{uuid.uuid4().hex[:12]}'
migrator_password = f'Lb{uuid.uuid4().hex}'
quote = postgres_engine.dialect.identifier_preparer.quote
database_name = sa.engine.make_url(postgres_url).database
migrator_url = (
sa.engine.make_url(postgres_url)
.set(username=migrator_role, password=migrator_password)
.render_as_string(hide_password=False)
)
migrator_engine: AsyncEngine | None = None
role_created = False
admin_role: str | None = None
source_tables = ('knowledge_bases', 'knowledge_base_files', 'knowledge_base_chunks')
async def read_rls_states(conn, table_names: tuple[str, ...]) -> dict[str, tuple[bool, bool]]:
rows = (
(
await conn.execute(
text(
"""
SELECT c.relname, c.relrowsecurity, c.relforcerowsecurity
FROM pg_class AS c
JOIN pg_namespace AS n ON n.oid = c.relnamespace
WHERE n.nspname = current_schema()
AND c.relname IN :table_names
"""
).bindparams(sa.bindparam('table_names', expanding=True)),
{'table_names': table_names},
)
)
.mappings()
.all()
)
return {str(row['relname']): (bool(row['relrowsecurity']), bool(row['relforcerowsecurity'])) for row in rows}
try:
async with postgres_engine.begin() as conn:
admin_role = await conn.scalar(text('SELECT current_user'))
await conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
await conn.run_sync(Base.metadata.create_all)
await conn.execute(
text(
'ALTER TABLE knowledge_bases DROP CONSTRAINT IF EXISTS '
'ck_knowledge_bases_embedding_dimension_positive'
)
)
await conn.execute(text('ALTER TABLE knowledge_bases DROP COLUMN IF EXISTS embedding_dimension'))
await run_alembic_stamp(postgres_engine, '0010_scope_resources')
await run_alembic_upgrade(postgres_engine, '0012_plugin_identity')
assert await get_alembic_current(postgres_engine) == '0012_plugin_identity'
embedding_a = '[' + ','.join(['0.125'] * 384) + ']'
embedding_b = '[' + ','.join(['0.25'] * 384) + ']'
async with postgres_engine.begin() as conn:
await conn.execute(
text(
"""
INSERT INTO workspaces
(uuid, instance_uuid, name, slug, type, status, source, projection_revision)
VALUES
(:uuid, 'legacy-vector-instance', 'legacy', 'legacy-vector',
'team', 'active', 'cloud_projection', 0)
"""
),
{'uuid': workspace_uuid},
)
await conn.execute(
text(
"""
INSERT INTO knowledge_bases
(uuid, workspace_uuid, name, collection_id, legacy_vector_collection)
VALUES
(:uuid, :workspace_uuid, 'legacy', 'legacy-collection', true)
"""
),
{'uuid': knowledge_base_uuid, 'workspace_uuid': workspace_uuid},
)
await conn.execute(
text(
"""
INSERT INTO knowledge_base_files
(uuid, workspace_uuid, kb_id, file_name, extension, status)
VALUES
('legacy-file', :workspace_uuid, :kb_uuid, 'legacy.txt', 'txt', 'completed')
"""
),
{'workspace_uuid': workspace_uuid, 'kb_uuid': knowledge_base_uuid},
)
await conn.execute(
text(
"""
INSERT INTO knowledge_base_chunks (uuid, workspace_uuid, file_id, text)
VALUES ('legacy-chunk', :workspace_uuid, 'legacy-file', 'chunk text')
"""
),
{'workspace_uuid': workspace_uuid},
)
await conn.execute(
text(
"""
CREATE TABLE langbot_vectors (
id VARCHAR(255) PRIMARY KEY,
collection VARCHAR(255),
embedding vector NOT NULL,
text TEXT,
file_id VARCHAR(255),
chunk_uuid VARCHAR(255)
)
"""
)
)
await conn.execute(
text(
"""
INSERT INTO langbot_vectors
(id, collection, embedding, text, file_id, chunk_uuid)
VALUES
('legacy-by-collection', 'legacy-collection', CAST(:embedding_a AS vector),
'collection row', NULL, NULL),
('legacy-by-chunk', 'unmatched-collection', CAST(:embedding_b AS vector),
'chunk row', NULL, 'legacy-chunk')
"""
),
{'embedding_a': embedding_a, 'embedding_b': embedding_b},
)
# Exercise exact restoration rather than assuming all source tables
# arrived with identical flags.
await conn.execute(text('ALTER TABLE knowledge_base_files NO FORCE ROW LEVEL SECURITY'))
await conn.execute(text('ALTER TABLE knowledge_base_chunks NO FORCE ROW LEVEL SECURITY'))
await conn.execute(text('ALTER TABLE knowledge_base_chunks DISABLE ROW LEVEL SECURITY'))
expected_source_rls = await read_rls_states(conn, source_tables)
assert expected_source_rls == {
'knowledge_bases': (True, True),
'knowledge_base_files': (True, False),
'knowledge_base_chunks': (False, False),
}
await conn.execute(
text(f"CREATE ROLE {quote(migrator_role)} LOGIN PASSWORD '{migrator_password}' NOSUPERUSER NOBYPASSRLS")
)
role_created = True
await conn.execute(text(f'GRANT CONNECT ON DATABASE {quote(database_name)} TO {quote(migrator_role)}'))
await conn.execute(text(f'GRANT USAGE, CREATE ON SCHEMA public TO {quote(migrator_role)}'))
await conn.execute(text(f'GRANT SELECT, UPDATE ON alembic_version TO {quote(migrator_role)}'))
for table_name in (*source_tables, 'langbot_vectors'):
await conn.execute(text(f'ALTER TABLE {quote(table_name)} OWNER TO {quote(migrator_role)}'))
role = (
await conn.execute(
text('SELECT rolsuper, rolbypassrls FROM pg_roles WHERE rolname = :role'),
{'role': migrator_role},
)
).one()
assert role == (False, False)
owned_tables = set(
(
await conn.execute(
text(
"""
SELECT c.relname
FROM pg_class AS c
JOIN pg_namespace AS n ON n.oid = c.relnamespace
WHERE n.nspname = current_schema()
AND c.relname IN :table_names
AND pg_get_userbyid(c.relowner) = :role
"""
).bindparams(sa.bindparam('table_names', expanding=True)),
{'table_names': (*source_tables, 'langbot_vectors'), 'role': migrator_role},
)
).scalars()
)
assert owned_tables == {*source_tables, 'langbot_vectors'}
migrator_engine = create_async_engine(migrator_url)
await run_alembic_upgrade(migrator_engine, '0013_tenant_pgvector')
assert await get_alembic_current(migrator_engine) == '0013_tenant_pgvector'
async with postgres_engine.connect() as conn:
migrated_rows = (
(
await conn.execute(
text(
"""
SELECT workspace_uuid, knowledge_base_uuid, vector_id,
embedding_dimension, text, file_id, chunk_uuid
FROM langbot_vectors
ORDER BY vector_id
"""
)
)
)
.mappings()
.all()
)
assert [dict(row) for row in migrated_rows] == [
{
'workspace_uuid': workspace_uuid,
'knowledge_base_uuid': knowledge_base_uuid,
'vector_id': 'legacy-by-chunk',
'embedding_dimension': 384,
'text': 'chunk row',
'file_id': None,
'chunk_uuid': 'legacy-chunk',
},
{
'workspace_uuid': workspace_uuid,
'knowledge_base_uuid': knowledge_base_uuid,
'vector_id': 'legacy-by-collection',
'embedding_dimension': 384,
'text': 'collection row',
'file_id': None,
'chunk_uuid': None,
},
]
assert (
await conn.scalar(
text('SELECT embedding_dimension FROM knowledge_bases WHERE uuid = :uuid'),
{'uuid': knowledge_base_uuid},
)
== 384
)
assert await conn.scalar(text("SELECT to_regclass('langbot_vectors_legacy_0013') IS NULL")) is True
assert await read_rls_states(conn, source_tables) == expected_source_rls
assert await read_rls_states(conn, ('langbot_vectors',)) == {'langbot_vectors': (True, True)}
assert (
await conn.scalar(
text(
"""
SELECT COUNT(*)
FROM pg_policy AS p
JOIN pg_class AS c ON c.oid = p.polrelid
JOIN pg_namespace AS n ON n.oid = c.relnamespace
WHERE n.nspname = current_schema()
AND c.relname = 'langbot_vectors'
AND p.polname = 'langbot_workspace_isolation'
"""
)
)
== 1
)
async with migrator_engine.connect() as conn:
assert await conn.scalar(text('SELECT COUNT(*) FROM langbot_vectors')) == 0
async with migrator_engine.begin() as conn:
await conn.execute(
text("SELECT set_config('langbot.workspace_uuid', :workspace_uuid, true)"),
{'workspace_uuid': workspace_uuid},
)
assert await conn.scalar(text('SELECT COUNT(*) FROM langbot_vectors')) == 2
finally:
if migrator_engine is not None:
await migrator_engine.dispose()
if role_created:
async with postgres_engine.connect() as conn:
if admin_role is None: # pragma: no cover - setup cannot create the role without an admin
admin_role = await conn.scalar(text('SELECT current_user'))
await conn.execute(text(f'REASSIGN OWNED BY {quote(migrator_role)} TO {quote(admin_role)}'))
await conn.execute(text(f'DROP OWNED BY {quote(migrator_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(migrator_role)}'))
async def test_pgvector_shared_database_is_scoped_indexed_and_ddl_free_at_runtime(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
monkeypatch,
) -> None:
instance_uuid = 'pgvector-tenant-integration'
workspace_a = '10000000-0000-0000-0000-000000000101'
workspace_b = '20000000-0000-0000-0000-000000000202'
kb_a = 'knowledge-base-a'
kb_b = 'knowledge-base-b'
role_suffix = uuid.uuid4().hex[:12]
runtime_role = f'lb_vector_runtime_{role_suffix}'
role_password = f'Lb{uuid.uuid4().hex}'
release_manager: PersistenceManager | None = None
runtime_manager: PersistenceManager | None = None
role_created = False
_restore_postgres_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', instance_uuid)
quote = postgres_engine.dialect.identifier_preparer.quote
database_name = sa.engine.make_url(postgres_url).database
runtime_url = (
sa.engine.make_url(postgres_url)
.set(username=runtime_role, password=role_password)
.render_as_string(hide_password=False)
)
try:
release_app = _application(postgres_url, 'pgvector-release-migration-test')
release_manager = PersistenceManager(release_app, mode=PersistenceMode.RELEASE_MIGRATION)
release_app.persistence_mgr = release_manager
await release_manager.initialize()
for workspace_uuid, kb_uuid, slug in (
(workspace_a, kb_a, 'vector-a'),
(workspace_b, kb_b, 'vector-b'),
):
async with release_manager.tenant_uow(workspace_uuid) as uow:
await uow.execute(
sa.insert(Workspace).values(
uuid=workspace_uuid,
instance_uuid=instance_uuid,
name=slug,
slug=slug,
type='team',
status='active',
source='cloud_projection',
projection_revision=0,
)
)
await uow.execute(
sa.insert(KnowledgeBase).values(
uuid=kb_uuid,
workspace_uuid=workspace_uuid,
name=slug,
embedding_dimension=384,
)
)
async with postgres_engine.connect() as conn:
await conn.execute(text(f"CREATE ROLE {quote(runtime_role)} LOGIN PASSWORD '{role_password}'"))
await conn.execute(text(f'GRANT CONNECT ON DATABASE {quote(database_name)} TO {quote(runtime_role)}'))
await conn.execute(text(f'GRANT USAGE ON SCHEMA public TO {quote(runtime_role)}'))
business_tables = release_manager._runtime_business_table_names()
quoted_tables = ', '.join(f'public.{quote(table_name)}' for table_name in business_tables)
await conn.execute(
text(f'GRANT SELECT, INSERT, UPDATE, DELETE ON TABLE {quoted_tables} TO {quote(runtime_role)}')
)
await conn.execute(text(f'GRANT SELECT ON TABLE public.alembic_version TO {quote(runtime_role)}'))
sequence_names = await release_manager._runtime_business_sequence_names(conn, business_tables)
if sequence_names:
quoted_sequences = ', '.join(f'public.{quote(sequence_name)}' for sequence_name in sequence_names)
await conn.execute(text(f'GRANT USAGE, SELECT ON SEQUENCE {quoted_sequences} TO {quote(runtime_role)}'))
role_created = True
runtime_app = _application(runtime_url, 'pgvector-runtime-test')
runtime_manager = PersistenceManager(runtime_app, mode=PersistenceMode.CLOUD_RUNTIME)
runtime_app.persistence_mgr = runtime_manager
await runtime_manager.initialize()
adapter = PgVectorDatabase(
runtime_app,
use_business_database=True,
allowed_dimensions=[384],
)
scope_a = PgVectorScope(workspace_a, kb_a, 384)
scope_b = PgVectorScope(workspace_b, kb_b, 384)
# The same vector ID is valid in two Workspaces because the relational
# primary key includes Workspace and knowledge base.
await adapter.add_embeddings(
'opaque-a',
['same-vector'],
[[0.1] * 384],
[{'text': 'workspace-a', 'file_id': 'file-a', 'uuid': 'chunk-a'}],
scope=scope_a,
)
await adapter.add_embeddings(
'opaque-b',
['same-vector'],
[[0.2] * 384],
[{'text': 'workspace-b', 'file_id': 'file-b', 'uuid': 'chunk-b'}],
scope=scope_b,
)
result_a = await adapter.search('opaque-a', [0.1] * 384, scope=scope_a)
result_b = await adapter.search('opaque-b', [0.2] * 384, scope=scope_b)
assert result_a['metadatas'][0][0]['text'] == 'workspace-a'
assert result_b['metadatas'][0][0]['text'] == 'workspace-b'
# Guessing another knowledge-base UUID while retaining A's Workspace
# cannot escape either the explicit conditions or PostgreSQL RLS.
guessed = await adapter.search(
'attacker-controlled-name',
[0.2] * 384,
scope=PgVectorScope(workspace_a, kb_b, 384),
)
assert guessed['ids'] == [[]]
with pytest.raises(ValueError, match='trusted PgVectorScope'):
await adapter.search('opaque-a', [0.1] * 384)
with pytest.raises(ValueError, match='selected dimension'):
await adapter.add_embeddings(
'opaque-a',
['bad-dimension'],
[[0.1] * 383],
[{}],
scope=scope_a,
)
# Deliberately omit the application Workspace predicate. FORCE RLS is
# still the second isolation boundary and returns only A.
async with runtime_manager.tenant_uow(workspace_a) as uow:
rows = (
await uow.execute(
sa.select(PgVectorEntry.workspace_uuid, PgVectorEntry.vector_id).where(
PgVectorEntry.vector_id == 'same-vector'
)
)
).all()
assert rows == [(workspace_a, 'same-vector')]
# Index-plan inspection is deployment diagnostics, not a public tenant
# Session capability. Establish the same transaction-local RLS scope on
# a test-only connection and keep raw EXPLAIN outside TenantUnitOfWork.
runtime_engine = runtime_manager.get_db_engine()
async with runtime_engine.begin() as conn:
await conn.execute(
text('SELECT set_config(:setting_name, :setting_value, true)'),
{'setting_name': 'langbot.workspace_uuid', 'setting_value': workspace_a},
)
await conn.execute(text('SET LOCAL enable_seqscan = off'))
plan = '\n'.join(
(
await conn.execute(
text(
"""
EXPLAIN SELECT vector_id
FROM langbot_vectors
WHERE workspace_uuid = :workspace_uuid
AND knowledge_base_uuid = :knowledge_base_uuid
AND embedding_dimension = 384
ORDER BY (embedding::vector(384)) <=> CAST(:query AS vector(384))
LIMIT 5
"""
),
{
'workspace_uuid': workspace_a,
'knowledge_base_uuid': kb_a,
'query': '[' + ','.join(['0.1'] * 384) + ']',
},
)
).scalars()
)
assert 'ix_langbot_vectors_hnsw_cosine_384' in plan
async with runtime_engine.connect() as conn:
assert await conn.scalar(text('SELECT COUNT(*) FROM langbot_vectors')) == 0
assert await conn.scalar(text("SELECT current_setting('langbot.workspace_uuid', true)")) in (None, '')
items_a, total_a = await adapter.list_by_filter('opaque-a', scope=scope_a)
assert total_a == 1
assert items_a[0]['metadata']['file_id'] == 'file-a'
await adapter.delete_by_file_id('opaque-a', 'file-a', scope=scope_a)
assert (await adapter.list_by_filter('opaque-a', scope=scope_a))[1] == 0
assert (await adapter.list_by_filter('opaque-b', scope=scope_b))[1] == 1
finally:
if runtime_manager is not None and getattr(runtime_manager, 'db', None) is not None:
await runtime_manager.get_db_engine().dispose()
if release_manager is not None and getattr(release_manager, 'db', None) is not None:
await release_manager.get_db_engine().dispose()
if role_created:
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}'))
tests/integration/persistence/test_plugin_identity_migration.py
+94
View File
@@ -0,0 +1,94 @@
from __future__ import annotations
import hashlib
import uuid
import pytest
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.persistence.alembic_runner import (
get_alembic_current,
run_alembic_stamp,
run_alembic_upgrade,
)
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
async def test_legacy_plugin_settings_receive_stable_random_installation_identities(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "plugin-identity.db"}')
metadata = sa.MetaData()
plugin_settings = sa.Table(
'plugin_settings',
metadata,
sa.Column('workspace_uuid', sa.String(36), primary_key=True),
sa.Column('plugin_author', sa.String(255), primary_key=True),
sa.Column('plugin_name', sa.String(255), primary_key=True),
sa.Column('enabled', sa.Boolean, nullable=False, server_default=sa.true()),
sa.Column('priority', sa.Integer, nullable=False, server_default='0'),
sa.Column('config', sa.JSON, nullable=False, server_default='{}'),
sa.Column('install_source', sa.String(255), nullable=False, server_default='local'),
sa.Column('install_info', sa.JSON, nullable=False, server_default='{}'),
)
try:
async with engine.begin() as connection:
await connection.run_sync(metadata.create_all)
await connection.execute(
plugin_settings.insert(),
[
{
'workspace_uuid': '11111111-1111-4111-8111-111111111111',
'plugin_author': 'author',
'plugin_name': 'one',
},
{
'workspace_uuid': '22222222-2222-4222-8222-222222222222',
'plugin_author': 'author',
'plugin_name': 'two',
},
],
)
await run_alembic_stamp(engine, '0011_postgres_tenant_rls')
await run_alembic_upgrade(engine, '0012_plugin_identity')
async with engine.connect() as connection:
rows = (
(
await connection.execute(
sa.text(
'SELECT installation_uuid, artifact_digest, runtime_revision '
'FROM plugin_settings ORDER BY workspace_uuid'
)
)
)
.mappings()
.all()
)
columns = await connection.run_sync(
lambda sync_connection: {
column['name']: column for column in sa.inspect(sync_connection).get_columns('plugin_settings')
}
)
indexes = await connection.run_sync(
lambda sync_connection: {
index['name']: index for index in sa.inspect(sync_connection).get_indexes('plugin_settings')
}
)
assert await get_alembic_current(engine) == '0012_plugin_identity'
assert columns['installation_uuid']['nullable'] is False
assert columns['artifact_digest']['nullable'] is False
assert columns['runtime_revision']['nullable'] is False
assert indexes['ix_plugin_settings_workspace_installation']['unique'] == 1
assert len({row['installation_uuid'] for row in rows}) == 2
for row in rows:
uuid.UUID(row['installation_uuid'])
assert row['runtime_revision'] == 1
assert (
row['artifact_digest']
== hashlib.sha256(f'legacy-installation:{row["installation_uuid"]}'.encode()).hexdigest()
)
finally:
await engine.dispose()
tests/integration/persistence/test_release_migration_postgres.py
+734
View File
@@ -0,0 +1,734 @@
"""Real PostgreSQL coverage for the one-shot Cloud release migration job."""
from __future__ import annotations
import logging
import os
import uuid
from types import SimpleNamespace
import pytest
import sqlalchemy as sa
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from langbot.pkg.persistence import release_migration
from langbot.pkg.persistence.alembic_runner import (
get_alembic_current,
get_alembic_head,
run_alembic_stamp,
run_alembic_upgrade,
)
from langbot.pkg.persistence.mgr import (
PersistenceManager,
PersistenceMode,
_RELEASE_MIGRATION_ADVISORY_LOCK_ID,
)
from langbot.pkg.utils import constants
pytestmark = [pytest.mark.integration, pytest.mark.slow, pytest.mark.asyncio]
_RUNTIME_PASSWORD = 'runtime-secret-not-used-by-migration'
@pytest.fixture
def postgres_url() -> str:
url = os.environ.get('TEST_POSTGRES_URL')
if not url:
pytest.skip('TEST_POSTGRES_URL not set')
return url
@pytest.fixture
async def postgres_engine(postgres_url: str):
engine = create_async_engine(postgres_url, isolation_level='AUTOCOMMIT')
yield engine
await engine.dispose()
@pytest.fixture
async def clean_database(postgres_engine: AsyncEngine):
async def clean() -> None:
async with postgres_engine.begin() as conn:
table_names = await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names())
quote = postgres_engine.dialect.identifier_preparer.quote
for table_name in table_names:
await conn.execute(text(f'DROP TABLE {quote(table_name)} CASCADE'))
await clean()
yield
await clean()
def _restore_postgres_registry(monkeypatch) -> None:
from langbot.pkg.persistence import mgr as persistence_mgr_module
from langbot.pkg.persistence.databases.postgresql import PostgreSQLDatabaseManager
monkeypatch.setattr(
persistence_mgr_module.database,
'preregistered_managers',
[PostgreSQLDatabaseManager],
)
def _application(postgres_url: str, *, runtime_role: str = 'langbot_runtime_not_used_by_migration') -> SimpleNamespace:
url = sa.engine.make_url(postgres_url)
return SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'use': 'postgresql',
'postgresql': {
'host': url.host,
'port': url.port,
# This is deliberately not the operator role in the DSN.
'user': runtime_role,
'password': _RUNTIME_PASSWORD,
'database': url.database,
},
'cloud_migration': {'operator_dsn_env': 'TEST_RELEASE_OPERATOR_DSN'},
},
'vdb': {
'use': 'pgvector',
'pgvector': {
'use_business_database': True,
'allowed_dimensions': [384, 512, 768, 1024, 1536],
},
},
}
),
logger=logging.getLogger('cloud-release-migration-entrypoint-test'),
persistence_mgr=None,
)
async def test_release_entrypoint_holds_lock_migrates_validates_and_disposes(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
monkeypatch,
) -> None:
_restore_postgres_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', 'release-migration-entrypoint-test')
original_validate = PersistenceManager._validate_release_schema
validation_observed_lock = False
runtime_role = f'lb_release_runtime_{uuid.uuid4().hex[:12]}'
quote = postgres_engine.dialect.identifier_preparer.quote
async def validate_while_asserting_lock(self: PersistenceManager) -> None:
nonlocal validation_observed_lock
async with postgres_engine.connect() as conn:
acquired = await conn.scalar(
text('SELECT pg_try_advisory_lock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
if acquired:
await conn.scalar(
text('SELECT pg_advisory_unlock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
assert acquired is False
validation_observed_lock = True
await original_validate(self)
monkeypatch.setattr(PersistenceManager, '_validate_release_schema', validate_while_asserting_lock)
async with postgres_engine.connect() as conn:
await conn.execute(text(f"CREATE ROLE {quote(runtime_role)} LOGIN PASSWORD '{_RUNTIME_PASSWORD}'"))
# The release job must make this otherwise bare LOGIN usable without
# relying on pre-provisioned object ACLs.
assert (
await conn.scalar(
text(
"""
SELECT NOT EXISTS (
SELECT 1
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
CROSS JOIN LATERAL aclexplode(c.relacl) acl
JOIN pg_roles grantee ON grantee.oid = acl.grantee
WHERE n.nspname = current_schema()
AND grantee.rolname = :runtime_role
)
"""
),
{'runtime_role': runtime_role},
)
is True
)
ap = _application(postgres_url, runtime_role=runtime_role)
try:
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
assert validation_observed_lock is True
assert await get_alembic_current(postgres_engine) == get_alembic_head()
manager = ap.persistence_mgr
assert isinstance(manager, PersistenceManager)
business_tables = set(manager._runtime_business_table_names())
async with postgres_engine.connect() as conn:
assert await conn.scalar(text("SELECT to_regclass('langbot_vectors') IS NOT NULL")) is True
assert (
await conn.scalar(text("SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'vector')")) is True
)
runtime_role_state = (
(
await conn.execute(
text(
"""
SELECT
rolcanlogin,
rolsuper,
rolbypassrls,
rolcreatedb,
rolcreaterole,
rolreplication
FROM pg_roles
WHERE rolname = :runtime_role
"""
),
{'runtime_role': runtime_role},
)
)
.mappings()
.one()
)
assert dict(runtime_role_state) == {
'rolcanlogin': True,
'rolsuper': False,
'rolbypassrls': False,
'rolcreatedb': False,
'rolcreaterole': False,
'rolreplication': False,
}
database_privileges = set(
(
await conn.execute(
text(
"""
SELECT acl.privilege_type
FROM pg_database database
CROSS JOIN LATERAL aclexplode(database.datacl) acl
JOIN pg_roles grantee ON grantee.oid = acl.grantee
WHERE database.datname = current_database()
AND grantee.rolname = :runtime_role
AND acl.is_grantable IS FALSE
"""
),
{'runtime_role': runtime_role},
)
)
.scalars()
.all()
)
schema_privileges = set(
(
await conn.execute(
text(
"""
SELECT acl.privilege_type
FROM pg_namespace namespace
CROSS JOIN LATERAL aclexplode(namespace.nspacl) acl
JOIN pg_roles grantee ON grantee.oid = acl.grantee
WHERE namespace.nspname = current_schema()
AND grantee.rolname = :runtime_role
AND acl.is_grantable IS FALSE
"""
),
{'runtime_role': runtime_role},
)
)
.scalars()
.all()
)
assert database_privileges == {'CONNECT'}
assert schema_privileges == {'USAGE'}
assert (
await conn.scalar(
text("SELECT has_database_privilege(:runtime_role, current_database(), 'CREATE')"),
{'runtime_role': runtime_role},
)
is False
)
assert (
await conn.scalar(
text("SELECT has_schema_privilege(:runtime_role, current_schema(), 'CREATE')"),
{'runtime_role': runtime_role},
)
is False
)
# PostgreSQL grants TEMP to PUBLIC by default. The first Cloud
# release deliberately tolerates that inherited compatibility
# privilege while granting no direct TEMP ACL to the runtime role.
assert (
await conn.scalar(
text("SELECT has_database_privilege(:runtime_role, current_database(), 'TEMP')"),
{'runtime_role': runtime_role},
)
is True
)
object_grants = (
(
await conn.execute(
text(
"""
SELECT
c.relname,
c.relkind::text AS relkind,
acl.privilege_type,
acl.is_grantable
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
CROSS JOIN LATERAL aclexplode(c.relacl) acl
JOIN pg_roles grantee ON grantee.oid = acl.grantee
WHERE n.nspname = current_schema()
AND c.relkind IN ('r', 'p', 'S')
AND grantee.rolname = :runtime_role
ORDER BY c.relname, acl.privilege_type
"""
),
{'runtime_role': runtime_role},
)
)
.mappings()
.all()
)
direct_table_grants: dict[str, set[str]] = {}
direct_sequence_grants: dict[str, set[str]] = {}
for grant in object_grants:
assert grant['is_grantable'] is False
target = direct_sequence_grants if grant['relkind'] == 'S' else direct_table_grants
target.setdefault(grant['relname'], set()).add(grant['privilege_type'])
assert set(direct_table_grants) == business_tables | {'alembic_version'}
assert all(
privileges == {'SELECT', 'INSERT', 'UPDATE', 'DELETE'}
for table_name, privileges in direct_table_grants.items()
if table_name != 'alembic_version'
)
assert direct_table_grants['alembic_version'] == {'SELECT'}
assert direct_sequence_grants
assert all(privileges == {'USAGE', 'SELECT'} for privileges in direct_sequence_grants.values())
# The session-level lock must be released before the one-shot job exits.
assert (
await conn.scalar(
text('SELECT pg_try_advisory_lock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
is True
)
assert (
await conn.scalar(
text('SELECT pg_advisory_unlock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
is True
)
runtime_url = (
sa.engine.make_url(postgres_url)
.set(username=runtime_role, password=_RUNTIME_PASSWORD)
.render_as_string(hide_password=False)
)
runtime_engine = create_async_engine(runtime_url)
try:
async with runtime_engine.begin() as conn:
assert await conn.scalar(text('SELECT current_user')) == runtime_role
await conn.execute(text("INSERT INTO metadata (key, value) VALUES ('runtime-grant-smoke', 'created')"))
await conn.execute(text("UPDATE metadata SET value = 'updated' WHERE key = 'runtime-grant-smoke'"))
assert (
await conn.scalar(text("SELECT value FROM metadata WHERE key = 'runtime-grant-smoke'")) == 'updated'
)
await conn.execute(text("DELETE FROM metadata WHERE key = 'runtime-grant-smoke'"))
await conn.scalar(
text('SELECT nextval(CAST(:sequence_name AS regclass))'),
{'sequence_name': f'public.{sorted(direct_sequence_grants)[0]}'},
)
finally:
await runtime_engine.dispose()
runtime_application = _application(postgres_url, runtime_role=runtime_role)
runtime_manager = PersistenceManager(runtime_application, mode=PersistenceMode.CLOUD_RUNTIME)
runtime_application.persistence_mgr = runtime_manager
try:
# This reads alembic_version as the actual runtime role and then
# reruns the complete grant/catalog validator before startup.
await runtime_manager.initialize()
finally:
await runtime_manager.get_db_engine().dispose()
# The catalog validator must fail closed if the role is no longer
# deployable, even though the remaining grants still look plausible.
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE DELETE ON TABLE public.metadata FROM {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match="table 'metadata' grants are incomplete"):
await manager._validate_configured_runtime_postgres_role(require_grants=True)
finally:
manager = getattr(ap, 'persistence_mgr', None)
if isinstance(manager, PersistenceManager):
# The one-shot entrypoint disposes its pool before returning. This
# test deliberately reuses the manager for catalog mutation checks,
# which can open a fresh pool and therefore owns a second shutdown.
await manager.shutdown()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}'))
async def test_release_entrypoint_rejects_privileged_or_table_owning_runtime_role(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
monkeypatch,
) -> None:
_restore_postgres_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', 'release-runtime-role-validation-test')
runtime_role = f'lb_release_runtime_{uuid.uuid4().hex[:12]}'
quote = postgres_engine.dialect.identifier_preparer.quote
async with postgres_engine.connect() as conn:
operator_role = await conn.scalar(text('SELECT current_user'))
await conn.execute(text(f'CREATE ROLE {quote(runtime_role)} LOGIN'))
ap = _application(postgres_url, runtime_role=runtime_role)
try:
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} SUPERUSER'))
with pytest.raises(RuntimeError, match='must not be superuser or BYPASSRLS'):
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} NOSUPERUSER BYPASSRLS'))
with pytest.raises(RuntimeError, match='must not be superuser or BYPASSRLS'):
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} NOBYPASSRLS'))
await conn.execute(text(f'ALTER TABLE bots OWNER TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='owns tenant tables'):
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
finally:
async with postgres_engine.connect() as conn:
table_exists = await conn.scalar(text("SELECT to_regclass('bots') IS NOT NULL"))
if table_exists:
await conn.execute(text(f'ALTER TABLE bots OWNER TO {quote(operator_role)}'))
await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}'))
async def test_runtime_role_catalog_validator_rejects_delegation_and_escape_hatches(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
monkeypatch,
) -> None:
_restore_postgres_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', 'release-runtime-role-catalog-test')
suffix = uuid.uuid4().hex[:12]
runtime_role = f'lb_release_runtime_{suffix}'
delegated_role = f'lb_release_delegate_{suffix}'
extra_schema = f'lb_release_schema_{suffix}'
extra_view = f'lb_release_view_{suffix}'
owned_routine = f'lb_release_owned_routine_{suffix}'
security_definer = f'lb_release_definer_{suffix}'
foreign_wrapper = f'lb_release_fdw_{suffix}'
foreign_server = f'lb_release_server_{suffix}'
persistent_setting_secret = f'lb_release_secret_{suffix}'
database_name = sa.engine.make_url(postgres_url).database
assert database_name
quote = postgres_engine.dialect.identifier_preparer.quote
async with postgres_engine.connect() as conn:
await conn.execute(text(f"CREATE ROLE {quote(runtime_role)} LOGIN PASSWORD '{_RUNTIME_PASSWORD}'"))
await conn.execute(text(f'CREATE ROLE {quote(delegated_role)}'))
ap = _application(postgres_url, runtime_role=runtime_role)
try:
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
manager = ap.persistence_mgr
assert isinstance(manager, PersistenceManager)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'GRANT pg_read_all_data TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='must not participate in role memberships'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE pg_read_all_data FROM {quote(runtime_role)}'))
# Reject delegation in the other direction too: no role may be
# allowed to SET ROLE to the runtime identity or administer it.
await conn.execute(text(f'GRANT {quote(runtime_role)} TO {quote(delegated_role)} WITH ADMIN OPTION'))
with pytest.raises(RuntimeError, match='must not participate in role memberships'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE {quote(runtime_role)} FROM {quote(delegated_role)}'))
await conn.execute(
text(f'GRANT SELECT ON TABLE public.metadata TO {quote(runtime_role)} WITH GRANT OPTION')
)
with pytest.raises(RuntimeError, match='GRANT OPTION'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(
text(f'REVOKE GRANT OPTION FOR SELECT ON TABLE public.metadata FROM {quote(runtime_role)}')
)
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} SET search_path TO public'))
with pytest.raises(RuntimeError, match='persistent session overrides'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET search_path'))
await conn.execute(text(f'ALTER DATABASE {quote(database_name)} SET search_path TO public'))
with pytest.raises(RuntimeError, match='persistent session overrides'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER DATABASE {quote(database_name)} RESET search_path'))
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} SET session_replication_role TO replica'))
with pytest.raises(RuntimeError, match='persistent session overrides'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET session_replication_role'))
await conn.execute(
text(f"ALTER ROLE {quote(runtime_role)} SET application_name TO '{persistent_setting_secret}'")
)
with pytest.raises(RuntimeError, match='persistent session overrides') as error:
await manager._validate_configured_runtime_postgres_role()
assert persistent_setting_secret not in str(error.value)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET application_name'))
await conn.execute(text('CREATE EXTENSION dblink'))
with pytest.raises(RuntimeError, match='extensions must include vector and be limited'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text('DROP EXTENSION dblink'))
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'GRANT CREATE ON DATABASE {quote(database_name)} TO {quote(runtime_role)}'))
await conn.execute(text(f'GRANT CREATE ON SCHEMA public TO {quote(runtime_role)}'))
runtime_url = (
sa.engine.make_url(postgres_url)
.set(username=runtime_role, password=_RUNTIME_PASSWORD)
.render_as_string(hide_password=False)
)
runtime_engine = create_async_engine(runtime_url)
try:
async with runtime_engine.begin() as conn:
# hstore is a trusted extension in the production PG16 image,
# so this creates a real runtime-owned extension catalog row.
await conn.execute(text('CREATE EXTENSION hstore'))
finally:
await runtime_engine.dispose()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE CREATE ON SCHEMA public FROM {quote(runtime_role)}'))
await conn.execute(text(f'REVOKE CREATE ON DATABASE {quote(database_name)} FROM {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='must not own extensions'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text('DROP EXTENSION hstore CASCADE'))
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'CREATE FOREIGN DATA WRAPPER {quote(foreign_wrapper)}'))
await conn.execute(
text(f'CREATE SERVER {quote(foreign_server)} FOREIGN DATA WRAPPER {quote(foreign_wrapper)}')
)
await conn.execute(text(f'CREATE USER MAPPING FOR {quote(runtime_role)} SERVER {quote(foreign_server)}'))
with pytest.raises(RuntimeError, match='foreign data wrappers, servers, or user mappings'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP FOREIGN DATA WRAPPER {quote(foreign_wrapper)} CASCADE'))
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'CREATE SCHEMA {quote(extra_schema)} AUTHORIZATION {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='non-business schemas'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP SCHEMA {quote(extra_schema)} CASCADE'))
await conn.execute(text(f'CREATE VIEW public.{quote(extra_view)} AS SELECT key FROM public.metadata'))
await conn.execute(text(f'GRANT SELECT ON public.{quote(extra_view)} TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='non-business objects|table privileges are unsafe'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP VIEW public.{quote(extra_view)}'))
await conn.execute(text(f'GRANT SELECT (key) ON public.metadata TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='column-level ACLs'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE SELECT (key) ON public.metadata FROM {quote(runtime_role)}'))
# System file access functions are not SECURITY DEFINER, so the
# validator must reject their explicit EXECUTE ACL independently.
await conn.execute(
text(f'GRANT EXECUTE ON FUNCTION pg_catalog.pg_read_file(text) TO {quote(runtime_role)}')
)
runtime_application = _application(postgres_url, runtime_role=runtime_role)
runtime_manager = PersistenceManager(runtime_application, mode=PersistenceMode.CLOUD_RUNTIME)
runtime_application.persistence_mgr = runtime_manager
try:
with pytest.raises(RuntimeError, match='explicit EXECUTE privileges on routines'):
await runtime_manager.initialize()
finally:
await runtime_manager.get_db_engine().dispose()
async with postgres_engine.connect() as conn:
await conn.execute(
text(f'REVOKE EXECUTE ON FUNCTION pg_catalog.pg_read_file(text) FROM {quote(runtime_role)}')
)
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
# replica disables ordinary triggers/rules and foreign-key
# enforcement; it must never reach the runtime identity.
await conn.execute(text(f'GRANT SET ON PARAMETER session_replication_role TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='explicit SET or ALTER SYSTEM parameter privileges'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE SET ON PARAMETER session_replication_role FROM {quote(runtime_role)}'))
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(
text(f"CREATE FUNCTION public.{quote(owned_routine)}() RETURNS integer LANGUAGE sql AS 'SELECT 1'")
)
await conn.execute(text(f'ALTER FUNCTION public.{quote(owned_routine)}() OWNER TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='must not own routines'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP FUNCTION public.{quote(owned_routine)}()'))
async with postgres_engine.connect() as conn:
await conn.execute(
text(
f'CREATE FUNCTION public.{quote(security_definer)}() RETURNS integer '
"LANGUAGE sql SECURITY DEFINER AS 'SELECT 1'"
)
)
# Extension membership must not exempt an executable definer from
# the runtime audit, even for an allowlisted extension.
await conn.execute(text(f'ALTER EXTENSION vector ADD FUNCTION public.{quote(security_definer)}()'))
with pytest.raises(RuntimeError, match='SECURITY DEFINER'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER EXTENSION vector DROP FUNCTION public.{quote(security_definer)}()'))
finally:
manager = getattr(ap, 'persistence_mgr', None)
if isinstance(manager, PersistenceManager):
# The entrypoint disposed the release pool before returning; all
# validator calls above happened afterward and can reopen it.
await manager.shutdown()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER DATABASE {quote(database_name)} RESET search_path'))
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET search_path'))
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET session_replication_role'))
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET application_name'))
await conn.execute(text(f'REVOKE pg_read_all_data FROM {quote(runtime_role)}'))
await conn.execute(text(f'REVOKE {quote(runtime_role)} FROM {quote(delegated_role)}'))
await conn.execute(text(f'REVOKE CREATE ON SCHEMA public FROM {quote(runtime_role)}'))
await conn.execute(text(f'REVOKE CREATE ON DATABASE {quote(database_name)} FROM {quote(runtime_role)}'))
await conn.execute(
text(f'REVOKE EXECUTE ON FUNCTION pg_catalog.pg_read_file(text) FROM {quote(runtime_role)}')
)
await conn.execute(text(f'REVOKE SET ON PARAMETER session_replication_role FROM {quote(runtime_role)}'))
await conn.execute(text('DROP EXTENSION IF EXISTS dblink CASCADE'))
await conn.execute(text('DROP EXTENSION IF EXISTS hstore CASCADE'))
await conn.execute(text(f'DROP FOREIGN DATA WRAPPER IF EXISTS {quote(foreign_wrapper)} CASCADE'))
await conn.execute(text(f'DROP FUNCTION IF EXISTS public.{quote(owned_routine)}()'))
security_definer_is_extension_member = await conn.scalar(
text(
"""
SELECT EXISTS (
SELECT 1
FROM pg_depend dependency
JOIN pg_extension extension ON extension.oid = dependency.refobjid
WHERE dependency.classid = 'pg_proc'::regclass
AND dependency.objid = to_regprocedure(:routine)
AND dependency.refclassid = 'pg_extension'::regclass
AND dependency.deptype = 'e'
AND extension.extname = 'vector'
)
"""
),
{'routine': f'public.{security_definer}()'},
)
if security_definer_is_extension_member:
await conn.execute(text(f'ALTER EXTENSION vector DROP FUNCTION public.{quote(security_definer)}()'))
await conn.execute(text(f'DROP FUNCTION IF EXISTS public.{quote(security_definer)}()'))
await conn.execute(text(f'DROP VIEW IF EXISTS public.{quote(extra_view)}'))
await conn.execute(text(f'DROP SCHEMA IF EXISTS {quote(extra_schema)} CASCADE'))
await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(delegated_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}'))
async def test_direct_postgres_head_stamp_fails_without_business_schema(
postgres_engine: AsyncEngine,
clean_database,
) -> None:
await run_alembic_stamp(postgres_engine, '0012_plugin_identity')
with pytest.raises(RuntimeError, match='requires the knowledge_bases table'):
await run_alembic_upgrade(postgres_engine)
assert await get_alembic_current(postgres_engine) == '0012_plugin_identity'
async def test_release_entrypoint_fails_immediately_when_another_job_holds_lock(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
monkeypatch,
) -> None:
_restore_postgres_registry(monkeypatch)
ap = _application(postgres_url)
async with postgres_engine.connect() as lock_connection:
assert (
await lock_connection.scalar(
text('SELECT pg_try_advisory_lock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
is True
)
try:
with pytest.raises(RuntimeError, match='already holds the advisory lock'):
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
finally:
assert (
await lock_connection.scalar(
text('SELECT pg_advisory_unlock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
is True
)
async with postgres_engine.connect() as conn:
assert await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names()) == []
tests/integration/persistence/test_resource_tenancy_migration.py
+292
View File
@@ -0,0 +1,292 @@
from __future__ import annotations
import hashlib
import json
import uuid
import pytest
import sqlalchemy as sa
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.entity import persistence
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.persistence.alembic_runner import run_alembic_stamp, run_alembic_upgrade
from langbot.pkg.utils import importutil
from .resource_migration_support import TENANT_TABLES, create_legacy_resource_schema
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
async def _inspect(engine, callback):
async with engine.connect() as conn:
return await conn.run_sync(callback)
async def test_legacy_sqlite_resources_are_backfilled_and_contracted(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "legacy-resources.db"}')
try:
await create_legacy_resource_schema(engine, instance_uuid='resource-migration-test')
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
await run_alembic_upgrade(engine, 'head')
async with engine.connect() as conn:
workspace_uuid = await conn.scalar(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'"))
assert workspace_uuid is not None
for table_name in TENANT_TABLES:
count, distinct_workspaces = (
await conn.execute(sa.text(f'SELECT COUNT(*), COUNT(DISTINCT workspace_uuid) FROM {table_name}'))
).one()
assert count == 1, table_name
assert distinct_workspaces == 1, table_name
assert (
await conn.scalar(
sa.text(f'SELECT COUNT(*) FROM {table_name} WHERE workspace_uuid != :workspace_uuid'),
{'workspace_uuid': workspace_uuid},
)
== 0
)
api_key = (
(await conn.execute(sa.text('SELECT key_hash, scopes, status, created_by_account_uuid FROM api_keys')))
.mappings()
.one()
)
assert api_key['key_hash'] == hashlib.sha256(b'lbk_legacy-secret').hexdigest()
stored_scopes = api_key['scopes']
if isinstance(stored_scopes, str):
stored_scopes = json.loads(stored_scopes)
assert stored_scopes == ['*']
assert api_key['status'] == 'active'
assert api_key['created_by_account_uuid'] is not None
assert await conn.scalar(sa.text('SELECT normalized_email FROM users')) == 'owner@example.com'
legacy_kb = (
(
await conn.execute(
sa.text(
'SELECT collection_id, legacy_vector_collection FROM knowledge_bases WHERE uuid = :uuid'
),
{'uuid': 'kb-1'},
)
)
.mappings()
.one()
)
assert legacy_kb['collection_id'] == 'collection-1'
assert legacy_kb['legacy_vector_collection'] == 1
assert (
await conn.scalar(
sa.text(
'SELECT COUNT(*) FROM metadata '
"WHERE key IN ('wizard_status', 'wizard_progress', 'rag_plugin_migration_needed')"
)
)
== 0
)
assert (
await conn.scalar(
sa.text('SELECT COUNT(*) FROM workspace_metadata WHERE workspace_uuid = :workspace_uuid'),
{'workspace_uuid': workspace_uuid},
)
== 3
)
api_columns = await _inspect(
engine,
lambda conn: {column['name'] for column in sa.inspect(conn).get_columns('api_keys')},
)
assert 'key' not in api_columns
assert {'uuid', 'key_hash', 'scopes', 'status', 'expires_at', 'last_used_at'} <= api_columns
for table_name in TENANT_TABLES:
columns = await _inspect(
engine,
lambda conn, name=table_name: {column['name']: column for column in sa.inspect(conn).get_columns(name)},
)
assert columns['workspace_uuid']['nullable'] is False, table_name
if table_name == 'knowledge_bases':
assert columns['legacy_vector_collection']['nullable'] is False
pk_columns = {
table_name: tuple(
(
await _inspect(
engine,
lambda conn, name=table_name: sa.inspect(conn).get_pk_constraint(name),
)
)['constrained_columns']
)
for table_name in ('binary_storages', 'plugin_settings', 'monitoring_sessions')
}
assert pk_columns == {
'binary_storages': ('workspace_uuid', 'unique_key'),
'plugin_settings': ('workspace_uuid', 'plugin_author', 'plugin_name'),
'monitoring_sessions': ('workspace_uuid', 'session_id'),
}
pipeline_run_foreign_keys = await _inspect(
engine,
lambda conn: sa.inspect(conn).get_foreign_keys('pipeline_run_records'),
)
assert any(
tuple(foreign_key['constrained_columns']) == ('workspace_uuid', 'pipeline_uuid')
and foreign_key['referred_table'] == 'legacy_pipelines'
and tuple(foreign_key['referred_columns']) == ('workspace_uuid', 'uuid')
for foreign_key in pipeline_run_foreign_keys
)
finally:
await engine.dispose()
async def test_legacy_vector_marker_backfill_resumes_from_nullable_expand_step(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "legacy-vector-retry.db"}')
try:
await create_legacy_resource_schema(engine, instance_uuid='legacy-vector-retry')
async with engine.begin() as conn:
await conn.execute(sa.text('ALTER TABLE knowledge_bases ADD COLUMN legacy_vector_collection BOOLEAN NULL'))
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
await run_alembic_upgrade(engine, 'head')
async with engine.connect() as conn:
assert (
await conn.scalar(
sa.text('SELECT legacy_vector_collection FROM knowledge_bases WHERE uuid = :uuid'),
{'uuid': 'kb-1'},
)
== 1
)
columns = await _inspect(
engine,
lambda conn: {column['name']: column for column in sa.inspect(conn).get_columns('knowledge_bases')},
)
assert columns['legacy_vector_collection']['nullable'] is False
finally:
await engine.dispose()
async def test_sqlite_scoped_keys_allow_cross_workspace_but_reject_same_workspace(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "scoped-keys.db"}')
try:
await create_legacy_resource_schema(engine, instance_uuid='scoped-key-test')
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
await run_alembic_upgrade(engine, 'head')
second_workspace_uuid = str(uuid.uuid4())
async with engine.begin() as conn:
await conn.execute(sa.text('PRAGMA foreign_keys=ON'))
first_workspace_uuid = await conn.scalar(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'"))
await conn.execute(
sa.text(
'INSERT INTO workspaces '
'(uuid, instance_uuid, name, slug, type, status, source, projection_revision) '
"VALUES (:uuid, 'scoped-key-test', 'Second', 'second', 'team', 'active', "
"'cloud_projection', 0)"
),
{'uuid': second_workspace_uuid},
)
await conn.execute(
sa.text(
'INSERT INTO mcp_servers (uuid, workspace_uuid, name, enable, updated_at) '
"VALUES ('mcp-2', :workspace_uuid, 'shared-name', 1, CURRENT_TIMESTAMP)"
),
{'workspace_uuid': second_workspace_uuid},
)
await conn.execute(
sa.text(
'INSERT INTO plugin_settings '
'(workspace_uuid, plugin_author, plugin_name, enabled, '
'installation_uuid, artifact_digest, runtime_revision) '
"VALUES (:workspace_uuid, 'author', 'plugin', 1, "
':installation_uuid, :artifact_digest, 1)'
),
{
'workspace_uuid': second_workspace_uuid,
'installation_uuid': str(uuid.uuid4()),
'artifact_digest': hashlib.sha256(b'test-plugin-artifact').hexdigest(),
},
)
await conn.execute(
sa.text(
'INSERT INTO binary_storages '
'(workspace_uuid, unique_key, key, owner_type, owner) '
"VALUES (:workspace_uuid, 'plugin:demo:key', 'key', 'plugin', 'demo')"
),
{'workspace_uuid': second_workspace_uuid},
)
await conn.execute(
sa.text(
'INSERT INTO monitoring_sessions '
'(workspace_uuid, session_id, bot_id, last_activity, is_active) '
"VALUES (:workspace_uuid, 'session-1', 'bot-2', CURRENT_TIMESTAMP, 1)"
),
{'workspace_uuid': second_workspace_uuid},
)
with pytest.raises(IntegrityError):
async with engine.begin() as conn:
await conn.execute(sa.text('PRAGMA foreign_keys=ON'))
await conn.execute(
sa.text(
'INSERT INTO mcp_servers (uuid, workspace_uuid, name, enable, updated_at) '
"VALUES ('mcp-duplicate', :workspace_uuid, 'shared-name', 1, CURRENT_TIMESTAMP)"
),
{'workspace_uuid': first_workspace_uuid},
)
with pytest.raises(IntegrityError):
async with engine.begin() as conn:
await conn.execute(sa.text('PRAGMA foreign_keys=ON'))
await conn.execute(
sa.text(
'INSERT INTO llm_models (uuid, workspace_uuid, name, provider_uuid) '
"VALUES ('cross-workspace-model', :workspace_uuid, 'model', 'provider-1')"
),
{'workspace_uuid': second_workspace_uuid},
)
with pytest.raises(IntegrityError):
async with engine.begin() as conn:
await conn.execute(sa.text('PRAGMA foreign_keys=ON'))
await conn.execute(
sa.text(
'INSERT INTO pipeline_run_records '
'(uuid, workspace_uuid, pipeline_uuid, created_at) '
"VALUES ('cross-workspace-run', :workspace_uuid, 'pipeline-1', CURRENT_TIMESTAMP)"
),
{'workspace_uuid': second_workspace_uuid},
)
with pytest.raises(IntegrityError):
async with engine.begin() as conn:
await conn.execute(
sa.text(
'INSERT INTO mcp_servers (uuid, name, enable, updated_at) '
"VALUES ('unscoped-mcp', 'unscoped', 1, CURRENT_TIMESTAMP)"
)
)
finally:
await engine.dispose()
async def test_fresh_sqlite_schema_matches_resource_tenancy_contract(tmp_path):
importutil.import_modules_in_pkg(persistence)
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "fresh-resources.db"}')
try:
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
await run_alembic_stamp(engine, '0001_baseline')
await run_alembic_upgrade(engine, 'head')
tables = await _inspect(engine, lambda conn: set(sa.inspect(conn).get_table_names()))
assert set(TENANT_TABLES) | {'workspace_metadata'} <= tables
for table_name in TENANT_TABLES:
columns = await _inspect(
engine,
lambda conn, name=table_name: {column['name']: column for column in sa.inspect(conn).get_columns(name)},
)
assert columns['workspace_uuid']['nullable'] is False, table_name
if table_name == 'knowledge_bases':
assert columns['legacy_vector_collection']['nullable'] is False
finally:
await engine.dispose()
tests/integration/persistence/test_sqlite_migration_backup.py
+107
View File
@@ -0,0 +1,107 @@
from __future__ import annotations
import json
import logging
import pathlib
import sqlite3
import pytest
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.persistence import alembic_runner
from langbot.pkg.persistence.mgr import PersistenceManager
from .resource_migration_support import create_legacy_resource_schema
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
def _manager(engine) -> PersistenceManager:
database = type('Database', (), {'get_engine': lambda self: engine})()
application = type('Application', (), {})()
application.logger = logging.getLogger('sqlite-migration-backup-test')
manager = PersistenceManager(application)
manager.db = database
return manager
def _manifest_payloads(backup_directory) -> list[dict]:
return [json.loads(path.read_text(encoding='utf-8')) for path in sorted(backup_directory.glob('*.json'))]
def _assert_verified_backup(payload: dict) -> None:
backup_path = pathlib.Path(payload['backup_path'])
with sqlite3.connect(f'{backup_path.as_uri()}?mode=ro', uri=True) as connection:
assert connection.execute('PRAGMA quick_check').fetchall() == [('ok',)]
assert connection.execute('SELECT version_num FROM alembic_version').fetchone()[0] == payload['source_revision']
async def test_tenancy_migrations_retain_verified_boundary_backups(tmp_path):
database_path = tmp_path / 'legacy-with-backups.db'
engine = create_async_engine(f'sqlite+aiosqlite:///{database_path}')
try:
await create_legacy_resource_schema(engine, instance_uuid='backup-success')
await alembic_runner.run_alembic_stamp(engine, '0008_mcp_resource_prefs')
await _manager(engine)._run_alembic_migrations()
assert await alembic_runner.get_alembic_current(engine) == alembic_runner.get_alembic_head()
payloads = _manifest_payloads(tmp_path / 'migration-backups')
assert len(payloads) == 2
assert {
(payload['source_revision'], payload['target_revision'], payload['status']) for payload in payloads
} == {
('0008_mcp_resource_prefs', '0009_workspace_tenancy', 'migration_succeeded'),
('0009_workspace_tenancy', '0010_scope_resources', 'migration_succeeded'),
}
for payload in payloads:
_assert_verified_backup(payload)
finally:
await engine.dispose()
async def test_failed_tenancy_migration_restores_backup_and_revision(
tmp_path,
monkeypatch,
):
database_path = tmp_path / 'legacy-fault-injection.db'
engine = create_async_engine(f'sqlite+aiosqlite:///{database_path}')
real_upgrade = alembic_runner.run_alembic_upgrade
async def injected_upgrade(async_engine, revision='head'):
if revision != '0010_scope_resources':
return await real_upgrade(async_engine, revision)
async with async_engine.begin() as connection:
await connection.execute(sa.text('CREATE TABLE injected_partial_migration (value TEXT NOT NULL)'))
await alembic_runner.run_alembic_stamp(async_engine, '0010_scope_resources')
raise RuntimeError('injected migration failure after a fake revision stamp')
try:
await create_legacy_resource_schema(engine, instance_uuid='backup-failure')
await alembic_runner.run_alembic_stamp(engine, '0008_mcp_resource_prefs')
monkeypatch.setattr(alembic_runner, 'run_alembic_upgrade', injected_upgrade)
with pytest.raises(RuntimeError, match='injected migration failure'):
await _manager(engine)._run_alembic_migrations()
assert await alembic_runner.get_alembic_current(engine) == '0009_workspace_tenancy'
async with engine.connect() as connection:
tables = set(
await connection.run_sync(lambda sync_connection: sa.inspect(sync_connection).get_table_names())
)
assert 'injected_partial_migration' not in tables
payloads = _manifest_payloads(tmp_path / 'migration-backups')
restored = [payload for payload in payloads if payload['target_revision'] == '0010_scope_resources']
assert len(restored) == 1
assert restored[0]['status'] == 'restored_after_failure'
assert restored[0]['source_revision'] == '0009_workspace_tenancy'
_assert_verified_backup(restored[0])
monkeypatch.setattr(alembic_runner, 'run_alembic_upgrade', real_upgrade)
await _manager(engine)._run_alembic_migrations()
assert await alembic_runner.get_alembic_current(engine) == alembic_runner.get_alembic_head()
finally:
await engine.dispose()
tests/integration/persistence/test_workspace_migration.py
+380
View File
@@ -0,0 +1,380 @@
from __future__ import annotations
import logging
import uuid
import pytest
import sqlalchemy as sa
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.entity import persistence
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.persistence.mgr import PersistenceManager
from langbot.pkg.persistence.alembic_runner import (
get_alembic_head,
get_alembic_current,
run_alembic_downgrade,
run_alembic_stamp,
run_alembic_upgrade,
)
from langbot.pkg.utils import constants
from langbot.pkg.utils import importutil
from langbot.pkg.workspace.collaboration import normalize_email
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
async def _create_legacy_schema(
engine,
*,
include_instance_uuid: bool = True,
include_users: bool = True,
) -> None:
legacy_metadata = sa.MetaData()
metadata_table = sa.Table(
'metadata',
legacy_metadata,
sa.Column('key', sa.String(255), primary_key=True),
sa.Column('value', sa.String(255)),
)
users = sa.Table(
'users',
legacy_metadata,
sa.Column('id', sa.Integer, primary_key=True),
sa.Column('user', sa.String(255), nullable=False),
sa.Column('password', sa.String(255), nullable=False),
sa.Column('account_type', sa.String(32), nullable=False, server_default='local'),
sa.Column('space_account_uuid', sa.String(255), nullable=True),
sa.Column('space_access_token', sa.Text, nullable=True),
sa.Column('space_refresh_token', sa.Text, nullable=True),
sa.Column('space_access_token_expires_at', sa.DateTime, nullable=True),
sa.Column('space_api_key', sa.String(255), nullable=True),
sa.Column('created_at', sa.DateTime, nullable=False, server_default=sa.func.now()),
sa.Column('updated_at', sa.DateTime, nullable=False, server_default=sa.func.now()),
)
async with engine.begin() as conn:
await conn.run_sync(legacy_metadata.create_all)
await conn.execute(metadata_table.insert().values(key='database_version', value='25'))
if include_instance_uuid:
await conn.execute(metadata_table.insert().values(key='instance_uuid', value='instance_migration_test'))
if include_users:
await conn.execute(
users.insert(),
[
{'user': 'owner@example.com', 'password': 'owner-hash'},
{'user': 'member@example.com', 'password': 'member-hash'},
],
)
@pytest.fixture
async def legacy_engine(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "legacy-workspace.db"}')
await _create_legacy_schema(engine)
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
yield engine
await engine.dispose()
async def test_legacy_instance_gets_stable_accounts_and_default_workspace(legacy_engine):
await run_alembic_upgrade(legacy_engine, 'head')
async with legacy_engine.connect() as conn:
tables = set(await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names()))
assert {
'workspaces',
'workspace_memberships',
'workspace_invitations',
'workspace_execution_states',
}.issubset(tables)
accounts = (
(await conn.execute(sa.text('SELECT id, uuid, status, source, projection_revision FROM users ORDER BY id')))
.mappings()
.all()
)
assert len(accounts) == 2
assert len({account['uuid'] for account in accounts}) == 2
for account in accounts:
uuid.UUID(account['uuid'])
assert account['status'] == 'active'
assert account['source'] == 'local'
assert account['projection_revision'] == 0
workspace = (
(await conn.execute(sa.text('SELECT * FROM workspaces WHERE source = :source'), {'source': 'local'}))
.mappings()
.one()
)
assert workspace['instance_uuid'] == 'instance_migration_test'
assert workspace['slug'] == 'default'
assert workspace['status'] == 'active'
assert workspace['created_by_account_uuid'] == accounts[0]['uuid']
membership = (await conn.execute(sa.text('SELECT * FROM workspace_memberships'))).mappings().one()
assert membership['workspace_uuid'] == workspace['uuid']
assert membership['account_uuid'] == accounts[0]['uuid']
assert membership['role'] == 'owner'
assert membership['status'] == 'active'
execution_state = (await conn.execute(sa.text('SELECT * FROM workspace_execution_states'))).mappings().one()
assert execution_state['workspace_uuid'] == workspace['uuid']
assert execution_state['instance_uuid'] == 'instance_migration_test'
assert execution_state['active_generation'] == 1
assert execution_state['state'] == 'active'
assert execution_state['write_fenced'] in (False, 0)
assert await get_alembic_current(legacy_engine) == get_alembic_head()
async def test_workspace_upgrade_is_idempotent_and_preserves_identifiers(legacy_engine):
await run_alembic_upgrade(legacy_engine, 'head')
async with legacy_engine.connect() as conn:
account_uuids_before = (await conn.execute(sa.text('SELECT uuid FROM users ORDER BY id'))).scalars().all()
workspace_uuid_before = (
await conn.execute(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'"))
).scalar_one()
await run_alembic_upgrade(legacy_engine, 'head')
async with legacy_engine.connect() as conn:
account_uuids_after = (await conn.execute(sa.text('SELECT uuid FROM users ORDER BY id'))).scalars().all()
workspace_uuid_after = (
await conn.execute(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'"))
).scalar_one()
assert account_uuids_after == account_uuids_before
assert workspace_uuid_after == workspace_uuid_before
async def test_workspace_kernel_upgrade_downgrade_upgrade_round_trip(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-round-trip.db"}')
try:
await _create_legacy_schema(engine)
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
await run_alembic_upgrade(engine, '0009_workspace_tenancy')
assert await get_alembic_current(engine) == '0009_workspace_tenancy'
await run_alembic_downgrade(engine, '0008_mcp_resource_prefs')
assert await get_alembic_current(engine) == '0008_mcp_resource_prefs'
async with engine.connect() as conn:
tables = set(await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names()))
user_columns = {
column['name']
for column in await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_columns('users'))
}
accounts = (await conn.execute(sa.text('SELECT user, password FROM users ORDER BY id'))).all()
assert (
not {
'workspaces',
'workspace_memberships',
'workspace_invitations',
'workspace_execution_states',
}
& tables
)
assert not {'uuid', 'status', 'source', 'projection_revision'} & user_columns
assert accounts == [
('owner@example.com', 'owner-hash'),
('member@example.com', 'member-hash'),
]
await run_alembic_upgrade(engine, '0009_workspace_tenancy')
assert await get_alembic_current(engine) == '0009_workspace_tenancy'
async with engine.connect() as conn:
assert await conn.scalar(sa.text('SELECT COUNT(*) FROM workspaces')) == 1
assert await conn.scalar(sa.text('SELECT COUNT(*) FROM workspace_memberships')) == 1
finally:
await engine.dispose()
@pytest.mark.parametrize(
('raw_email', 'expected_email'),
[
('Straße@Example.COM', 'strasse@example.com'),
('Ꭰ@Example.COM', 'Ꭰ@example.com'),
],
)
async def test_workspace_upgrade_uses_runtime_unicode_email_normalization(
tmp_path,
raw_email,
expected_email,
):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "unicode-email.db"}')
try:
await _create_legacy_schema(engine, include_users=False)
async with engine.begin() as conn:
await conn.execute(
sa.text('INSERT INTO users (user, password, account_type) VALUES (:email, :password, :type)'),
{'email': raw_email, 'password': 'owner-hash', 'type': 'local'},
)
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
await run_alembic_upgrade(engine, 'head')
async with engine.connect() as conn:
assert await conn.scalar(sa.text('SELECT normalized_email FROM users')) == expected_email
finally:
await engine.dispose()
async def test_fresh_sqlite_schema_accepts_application_casefold_identity(tmp_path):
importutil.import_modules_in_pkg(persistence)
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "fresh-unicode-email.db"}')
canonical_email = normalize_email('Ꭰ@Example.COM')
try:
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
await conn.execute(
sa.insert(User).values(
uuid='00000000-0000-0000-0000-000000000099',
user=canonical_email,
normalized_email=canonical_email,
password='hash',
)
)
async with engine.connect() as conn:
assert await conn.scalar(sa.select(User.normalized_email)) == 'Ꭰ@example.com'
finally:
await engine.dispose()
async def test_workspace_upgrade_rejects_unicode_casefold_duplicate_accounts(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "unicode-email-duplicate.db"}')
try:
await _create_legacy_schema(engine, include_users=False)
async with engine.begin() as conn:
await conn.execute(
sa.text(
'INSERT INTO users (user, password, account_type) VALUES '
"('Straße@Example.COM', 'first-hash', 'local'), "
"('STRASSE@example.com', 'second-hash', 'local')"
)
)
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
with pytest.raises(RuntimeError, match='both normalize'):
await run_alembic_upgrade(engine, 'head')
finally:
await engine.dispose()
async def test_uninitialized_instance_gets_ownerless_default_workspace(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "uninitialized-instance.db"}')
try:
await _create_legacy_schema(engine, include_users=False)
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
await run_alembic_upgrade(engine, 'head')
async with engine.connect() as conn:
workspace = (await conn.execute(sa.text('SELECT * FROM workspaces'))).mappings().one()
membership_count = await conn.scalar(sa.text('SELECT COUNT(*) FROM workspace_memberships'))
execution_state = (await conn.execute(sa.text('SELECT * FROM workspace_execution_states'))).mappings().one()
assert workspace['created_by_account_uuid'] is None
assert membership_count == 0
assert execution_state['workspace_uuid'] == workspace['uuid']
assert execution_state['active_generation'] == 1
finally:
await engine.dispose()
async def test_local_workspace_unique_index_allows_cloud_projections(legacy_engine):
await run_alembic_upgrade(legacy_engine, 'head')
async with legacy_engine.begin() as conn:
await conn.execute(
sa.text(
'INSERT INTO workspaces '
'(uuid, instance_uuid, name, slug, type, status, source, projection_revision) '
'VALUES (:uuid, :instance_uuid, :name, :slug, :type, :status, :source, 0)'
),
{
'uuid': str(uuid.uuid4()),
'instance_uuid': 'instance_migration_test',
'name': 'Cloud Projection',
'slug': 'cloud-projection',
'type': 'team',
'status': 'active',
'source': 'cloud_projection',
},
)
with pytest.raises(IntegrityError):
async with legacy_engine.begin() as conn:
await conn.execute(
sa.text(
'INSERT INTO workspaces '
'(uuid, instance_uuid, name, slug, type, status, source, projection_revision) '
'VALUES (:uuid, :instance_uuid, :name, :slug, :type, :status, :source, 0)'
),
{
'uuid': str(uuid.uuid4()),
'instance_uuid': 'instance_migration_test',
'name': 'Second Local',
'slug': 'second-local',
'type': 'team',
'status': 'active',
'source': 'local',
},
)
async def test_legacy_instance_without_bound_instance_uuid_fails_closed(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "missing-instance.db"}')
try:
await _create_legacy_schema(engine, include_instance_uuid=False)
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
with pytest.raises(RuntimeError, match='instance_uuid'):
await run_alembic_upgrade(engine, 'head')
finally:
await engine.dispose()
async def test_persistence_startup_defers_workspace_tables_until_account_upgrade(tmp_path, monkeypatch):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "startup-order.db"}')
try:
await _create_legacy_schema(engine)
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
monkeypatch.setattr(constants, 'instance_id', 'instance_migration_test')
database = type('Database', (), {'get_engine': lambda self: engine})()
application = type('Application', (), {})()
application.logger = logging.getLogger('workspace-startup-test')
manager = PersistenceManager(application)
manager.db = database
await manager.create_tables()
async with engine.connect() as conn:
tables_before_migration = set(
await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names())
)
assert 'workspaces' not in tables_before_migration
await manager._run_alembic_migrations()
async with engine.connect() as conn:
workspace = (
(await conn.execute(sa.text("SELECT * FROM workspaces WHERE source = 'local'"))).mappings().one()
)
assert workspace['instance_uuid'] == 'instance_migration_test'
finally:
await engine.dispose()
async def test_persistence_startup_rejects_instance_uuid_drift(tmp_path, monkeypatch):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "instance-drift.db"}')
try:
await _create_legacy_schema(engine)
monkeypatch.setattr(constants, 'instance_id', 'different_instance')
database = type('Database', (), {'get_engine': lambda self: engine})()
application = type('Application', (), {})()
application.logger = logging.getLogger('workspace-instance-drift-test')
manager = PersistenceManager(application)
manager.db = database
with pytest.raises(RuntimeError, match='does not match'):
await manager.create_tables()
finally:
await engine.dispose()
tests/integration/pipeline/test_full_flow.py
+9 -1
View File
@@ -210,7 +210,15 @@ def pipeline_app():
mock_conversation.update_time = None
mock_conversation.create_time = None
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
async def get_scoped_session(query):
context = query._execution_context
mock_session.instance_uuid = context.instance_uuid
mock_session.workspace_uuid = context.workspace_uuid
mock_session.placement_generation = context.placement_generation
mock_session.bot_uuid = query.bot_uuid
return mock_session
app.sess_mgr.get_session = AsyncMock(side_effect=get_scoped_session)
app.sess_mgr.get_conversation = AsyncMock(return_value=mock_conversation)
# Model mock for PreProcessor