mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 12:40:59 +00:00
chore(merge): sync master into dev/4.11.x
This commit is contained in:
@@ -9,7 +9,8 @@ Run: uv run pytest tests/integration/api/test_bots.py -q
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, AsyncMock, Mock
|
||||
from unittest.mock import ANY, 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()
|
||||
@@ -252,6 +284,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 TestBotEventRouteDryRunEndpoint:
|
||||
@@ -286,6 +337,7 @@ class TestBotEventRouteDryRunEndpoint:
|
||||
assert data['data']['target']['target_name'] == 'Member Agent'
|
||||
assert data['data']['diagnostic_steps'][0].startswith('Route 1')
|
||||
fake_bot_app.bot_service.dry_run_event_route.assert_awaited_with(
|
||||
ANY,
|
||||
bot_uuid='test-bot-uuid',
|
||||
event_type='platform.member.joined',
|
||||
event_data={'room': {'id': 'room-1'}},
|
||||
@@ -318,7 +370,7 @@ class TestBotEventRouteStatusEndpoint:
|
||||
assert data['code'] == 0
|
||||
assert data['data']['routes'][0]['binding_id'] == 'binding-1'
|
||||
assert data['data']['routes'][0]['last_status'] == 'delivered'
|
||||
fake_bot_app.bot_service.list_event_route_statuses.assert_awaited_with('test-bot-uuid')
|
||||
fake_bot_app.bot_service.list_event_route_statuses.assert_awaited_with(ANY, 'test-bot-uuid')
|
||||
|
||||
|
||||
@pytest.mark.usefixtures('mock_circular_import_chain')
|
||||
@@ -343,6 +395,7 @@ class TestBotEventRouteTestEndpoint:
|
||||
assert data['data']['dispatched'] is True
|
||||
assert data['data']['event_type'] == 'message.received'
|
||||
fake_bot_app.bot_service.dispatch_test_event_route.assert_awaited_with(
|
||||
ANY,
|
||||
bot_uuid='test-bot-uuid',
|
||||
event_type='message.received',
|
||||
payload={'message_text': 'hello'},
|
||||
|
||||
@@ -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()
|
||||
@@ -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,19 @@ def fake_embed_app():
|
||||
|
||||
mock_runtime_bot = Mock()
|
||||
mock_runtime_bot.bot_entity = mock_bot_entity
|
||||
mock_runtime_bot.get_pipeline_target_for_event_type = Mock(return_value='test-pipeline-uuid')
|
||||
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 +106,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 +139,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 +226,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 +275,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 +337,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()
|
||||
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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=[])
|
||||
@@ -280,15 +308,145 @@ class TestPipelineExtensionsEndpoint:
|
||||
assert data['code'] == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_extensions_tolerates_skill_list_failure(self, quart_test_client, fake_pipeline_app):
|
||||
"""Pipeline plugin/MCP binding UI should not fail when Box skill listing is unavailable."""
|
||||
fake_pipeline_app.skill_service.list_skills = AsyncMock(side_effect=TimeoutError('box_list_skills timeout'))
|
||||
|
||||
response = await quart_test_client.get(
|
||||
'/api/v1/pipelines/test-pipeline-uuid/extensions', headers={'Authorization': 'Bearer test_token'}
|
||||
)
|
||||
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
|
||||
data = await response.get_json()
|
||||
assert data['code'] == 0
|
||||
assert data['data']['available_skills'] == []
|
||||
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
|
||||
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
@@ -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'
|
||||
@@ -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))
|
||||
@@ -9,62 +9,36 @@ Run: uv run pytest tests/integration/persistence/test_migrations.py -q
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import hashlib
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from alembic.config import Config
|
||||
from alembic.script import ScriptDirectory
|
||||
import sqlalchemy
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from langbot.pkg.entity.persistence import (
|
||||
agent as agent_models,
|
||||
agent_interaction as agent_interaction_models,
|
||||
agent_run as agent_run_models,
|
||||
agent_runner_state as agent_runner_state_models,
|
||||
bot as bot_models,
|
||||
metadata as metadata_models,
|
||||
monitoring as monitoring_models,
|
||||
)
|
||||
from langbot.pkg.entity.persistence.base import Base
|
||||
from langbot.pkg.persistence.alembic_runner import (
|
||||
_ALEMBIC_DIR,
|
||||
get_alembic_current,
|
||||
run_alembic_stamp,
|
||||
run_alembic_downgrade,
|
||||
run_alembic_upgrade,
|
||||
run_alembic_stamp,
|
||||
get_alembic_current,
|
||||
_ALEMBIC_DIR,
|
||||
)
|
||||
|
||||
|
||||
def _get_script_directory() -> ScriptDirectory:
|
||||
"""Load the repository's Alembic revision graph."""
|
||||
cfg = Config()
|
||||
cfg.set_main_option('script_location', _ALEMBIC_DIR)
|
||||
return ScriptDirectory.from_config(cfg)
|
||||
from alembic.config import Config
|
||||
from alembic.script import ScriptDirectory
|
||||
|
||||
|
||||
def _get_script_head() -> str:
|
||||
"""Resolve the only Alembic head without hardcoding a revision."""
|
||||
return _get_script_directory().get_current_head()
|
||||
"""Resolve the current Alembic head revision from the script directory.
|
||||
|
||||
Avoids hardcoding a revision number in assertions so adding a new
|
||||
migration doesn't require editing the migration tests.
|
||||
"""
|
||||
cfg = Config()
|
||||
cfg.set_main_option('script_location', _ALEMBIC_DIR)
|
||||
return ScriptDirectory.from_config(cfg).get_current_head()
|
||||
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
|
||||
class TestAlembicRevisionGraph:
|
||||
"""Static release gates for the Alembic graph."""
|
||||
|
||||
def test_revision_ids_fit_alembic_version_column_and_graph_has_one_head(self):
|
||||
script = _get_script_directory()
|
||||
revisions = list(script.walk_revisions())
|
||||
|
||||
assert script.get_bases() == ['0001_baseline']
|
||||
assert script.get_heads() == ['0015_official_runner_ids']
|
||||
assert all(len(item.revision) <= 32 for item in revisions), {
|
||||
item.revision: len(item.revision) for item in revisions if len(item.revision) > 32
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sqlite_db_url(tmp_path):
|
||||
"""Create SQLite URL with temporary database file."""
|
||||
@@ -178,333 +152,43 @@ class TestSQLiteMigrationUpgrade:
|
||||
assert rev2 == rev1, f'Expected {rev1}, got {rev2}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upgrade_from_mcp_resource_branch_creates_agent_and_monitoring_schema(self, sqlite_engine):
|
||||
"""The MCP branch can converge into the Agent branch and the current head."""
|
||||
await run_alembic_stamp(sqlite_engine, '0008_mcp_resource_prefs')
|
||||
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')
|
||||
|
||||
def inspect_schema(sync_conn):
|
||||
inspector = sa.inspect(sync_conn)
|
||||
tables = set(inspector.get_table_names())
|
||||
monitoring_indexes = {
|
||||
index['name'] for index in inspector.get_indexes(monitoring_models.MonitoringToolCall.__tablename__)
|
||||
}
|
||||
return tables, monitoring_indexes
|
||||
|
||||
async with sqlite_engine.connect() as conn:
|
||||
tables, monitoring_indexes = await conn.run_sync(inspect_schema)
|
||||
|
||||
expected_agent_tables = {
|
||||
agent_models.Agent.__tablename__,
|
||||
agent_interaction_models.AgentInteraction.__tablename__,
|
||||
agent_run_models.AgentRun.__tablename__,
|
||||
agent_run_models.AgentRunEvent.__tablename__,
|
||||
agent_run_models.AgentRuntime.__tablename__,
|
||||
agent_runner_state_models.AgentRunnerState.__tablename__,
|
||||
}
|
||||
expected_monitoring_indexes = {index.name for index in monitoring_models.MonitoringToolCall.__table__.indexes}
|
||||
|
||||
assert expected_agent_tables <= tables
|
||||
assert monitoring_models.MonitoringToolCall.__tablename__ in tables
|
||||
assert expected_monitoring_indexes <= monitoring_indexes
|
||||
assert await get_alembic_current(sqlite_engine) == _get_script_head()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bot_admin_data_migrates_when_create_all_already_created_table(self, sqlite_engine):
|
||||
"""0007 must migrate config admins even when the ORM created its table first."""
|
||||
config = {
|
||||
'admins': ['group_admin-1', 'person_user_with_underscore', 'malformed'],
|
||||
'preserved': True,
|
||||
}
|
||||
|
||||
async with sqlite_engine.begin() as conn:
|
||||
await conn.run_sync(bot_models.Bot.__table__.create)
|
||||
await conn.run_sync(bot_models.BotAdmin.__table__.create)
|
||||
await conn.run_sync(metadata_models.Metadata.__table__.create)
|
||||
await conn.execute(
|
||||
sa.insert(bot_models.Bot).values(
|
||||
uuid='bot-1',
|
||||
name='Bot',
|
||||
description='',
|
||||
adapter='test',
|
||||
adapter_config={},
|
||||
enable=True,
|
||||
)
|
||||
)
|
||||
await conn.execute(
|
||||
sa.insert(bot_models.BotAdmin).values(
|
||||
bot_uuid='bot-1',
|
||||
launcher_type='group',
|
||||
launcher_id='admin-1',
|
||||
)
|
||||
)
|
||||
await conn.execute(
|
||||
sa.insert(metadata_models.Metadata).values(
|
||||
key='instance_config',
|
||||
value=json.dumps(config),
|
||||
)
|
||||
)
|
||||
|
||||
await run_alembic_stamp(sqlite_engine, '0006_normalize_mcp_remote_mode')
|
||||
await run_alembic_upgrade(sqlite_engine, '0007_add_bot_admins')
|
||||
|
||||
async with sqlite_engine.connect() as conn:
|
||||
admin_rows = (
|
||||
await conn.execute(
|
||||
sa.select(
|
||||
bot_models.BotAdmin.launcher_type,
|
||||
bot_models.BotAdmin.launcher_id,
|
||||
).order_by(bot_models.BotAdmin.launcher_type, bot_models.BotAdmin.launcher_id)
|
||||
)
|
||||
).all()
|
||||
stored_config = (
|
||||
await conn.execute(
|
||||
sa.select(metadata_models.Metadata.value).where(metadata_models.Metadata.key == 'instance_config')
|
||||
)
|
||||
).scalar_one()
|
||||
|
||||
assert admin_rows == [('group', 'admin-1'), ('person', 'user_with_underscore')]
|
||||
assert json.loads(stored_config) == {'preserved': True}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pipeline_routing_rules_preserve_message_filters(self, sqlite_engine):
|
||||
"""Every legacy routing rule keeps its matching semantics in event bindings."""
|
||||
routing_rules = [
|
||||
{
|
||||
'type': 'launcher_type',
|
||||
'operator': 'eq',
|
||||
'value': 'group',
|
||||
'pipeline_uuid': 'pipeline-group',
|
||||
},
|
||||
{
|
||||
'type': 'launcher_id',
|
||||
'operator': 'regex',
|
||||
'value': '^room-',
|
||||
'pipeline_uuid': 'pipeline-room',
|
||||
},
|
||||
{
|
||||
'type': 'message_content',
|
||||
'operator': 'contains',
|
||||
'value': 'urgent',
|
||||
'pipeline_uuid': 'pipeline-content',
|
||||
},
|
||||
{
|
||||
'type': 'message_has_element',
|
||||
'operator': 'eq',
|
||||
'value': 'Image',
|
||||
'pipeline_uuid': 'pipeline-image',
|
||||
},
|
||||
{
|
||||
'type': 'message_has_element',
|
||||
'operator': 'neq',
|
||||
'value': 'Voice',
|
||||
'pipeline_uuid': 'pipeline-no-voice',
|
||||
},
|
||||
]
|
||||
|
||||
async with sqlite_engine.begin() as conn:
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'CREATE TABLE bots ('
|
||||
'uuid VARCHAR(255) PRIMARY KEY, '
|
||||
'use_pipeline_uuid VARCHAR(255), '
|
||||
'pipeline_routing_rules JSON NOT NULL, '
|
||||
'event_bindings JSON NOT NULL'
|
||||
')'
|
||||
)
|
||||
)
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'INSERT INTO bots '
|
||||
'(uuid, use_pipeline_uuid, pipeline_routing_rules, event_bindings) '
|
||||
'VALUES (:uuid, :default_pipeline, :rules, :bindings)'
|
||||
),
|
||||
{
|
||||
'uuid': 'bot-routing',
|
||||
'default_pipeline': 'pipeline-default',
|
||||
'rules': json.dumps(routing_rules),
|
||||
'bindings': '[]',
|
||||
},
|
||||
)
|
||||
|
||||
await run_alembic_stamp(sqlite_engine, '0008_agent_product_surface')
|
||||
await run_alembic_upgrade(sqlite_engine, '0009_migrate_event_bindings')
|
||||
|
||||
async with sqlite_engine.connect() as conn:
|
||||
raw_bindings = (
|
||||
await conn.execute(sa.text("SELECT event_bindings FROM bots WHERE uuid = 'bot-routing'"))
|
||||
).scalar_one()
|
||||
|
||||
bindings = json.loads(raw_bindings)
|
||||
filters_by_pipeline = {binding['target_uuid']: binding['filters'] for binding in bindings}
|
||||
assert filters_by_pipeline == {
|
||||
'pipeline-group': [{'field': 'chat_type', 'operator': 'eq', 'value': 'group'}],
|
||||
'pipeline-room': [{'field': 'chat_id', 'operator': 'regex', 'value': '^room-'}],
|
||||
'pipeline-content': [{'field': 'message_text', 'operator': 'contains', 'value': 'urgent'}],
|
||||
'pipeline-image': [{'field': 'message_element_types', 'operator': 'contains', 'value': 'Image'}],
|
||||
'pipeline-no-voice': [{'field': 'message_element_types', 'operator': 'not_contains', 'value': 'Voice'}],
|
||||
'pipeline-default': [],
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_official_runner_ids_migrate_without_losing_config_or_state(self, sqlite_engine):
|
||||
"""0015 rewrites persisted official Runner identities and state scope keys."""
|
||||
old_dify = 'plugin:langbot/dify-agent/default'
|
||||
new_dify = 'plugin:langbot-team/DifyAgent/default'
|
||||
old_codex = 'plugin:langbot/codex-agent/default'
|
||||
new_codex = 'plugin:langbot-team/CodexAgent/default'
|
||||
binding_id = f'pipeline_pipeline-1_{old_dify}'
|
||||
scope_payload = {
|
||||
'version': 2,
|
||||
'scope': 'conversation',
|
||||
'runner_id': old_dify,
|
||||
'binding_identity': binding_id,
|
||||
'bot_id': 'bot-1',
|
||||
'workspace_id': None,
|
||||
'conversation_id': 'conversation-1',
|
||||
'thread_id': None,
|
||||
}
|
||||
old_scope_key = (
|
||||
'conversation:v2:'
|
||||
+ hashlib.sha256(json.dumps(scope_payload, sort_keys=True, separators=(',', ':')).encode()).hexdigest()
|
||||
)
|
||||
|
||||
async with sqlite_engine.begin() as conn:
|
||||
await conn.execute(sa.text('CREATE TABLE legacy_pipelines (uuid TEXT PRIMARY KEY, config JSON NOT NULL)'))
|
||||
await conn.execute(
|
||||
sa.text('CREATE TABLE agents (uuid TEXT PRIMARY KEY, component_ref TEXT, config JSON NOT NULL)')
|
||||
)
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'CREATE TABLE workflows ('
|
||||
'uuid TEXT PRIMARY KEY, definition JSON NOT NULL, global_config JSON NOT NULL'
|
||||
')'
|
||||
)
|
||||
)
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'CREATE TABLE agent_runner_state ('
|
||||
'id INTEGER PRIMARY KEY, runner_id TEXT NOT NULL, binding_identity TEXT NOT NULL, '
|
||||
'scope TEXT NOT NULL, scope_key TEXT NOT NULL, state_key TEXT NOT NULL, '
|
||||
'bot_id TEXT, workspace_id TEXT, conversation_id TEXT, thread_id TEXT, '
|
||||
'actor_type TEXT, actor_id TEXT, subject_type TEXT, subject_id TEXT, '
|
||||
'UNIQUE(scope_key, state_key)'
|
||||
')'
|
||||
)
|
||||
)
|
||||
for table_name, columns in {
|
||||
'agent_run': 'id INTEGER PRIMARY KEY, runner_id TEXT NOT NULL, binding_id TEXT',
|
||||
'agent_interaction': 'id INTEGER PRIMARY KEY, runner_id TEXT NOT NULL, binding_id TEXT',
|
||||
'event_log': 'id INTEGER PRIMARY KEY, runner_id TEXT',
|
||||
'transcript': 'id INTEGER PRIMARY KEY, runner_id TEXT',
|
||||
}.items():
|
||||
await conn.execute(sa.text(f'CREATE TABLE {table_name} ({columns})'))
|
||||
|
||||
pipeline_config = {
|
||||
'ai': {
|
||||
'runner': {'id': old_dify},
|
||||
'runner_config': {
|
||||
old_dify: {'api-key': 'preserved-secret'},
|
||||
old_codex: {'workspace': 'K:/workspace'},
|
||||
},
|
||||
columns = await conn.run_sync(
|
||||
lambda sync_conn: {
|
||||
item['name'] for item in sqlalchemy.inspect(sync_conn).get_columns('knowledge_bases')
|
||||
}
|
||||
}
|
||||
await conn.execute(
|
||||
sa.text('INSERT INTO legacy_pipelines (uuid, config) VALUES (:uuid, :config)'),
|
||||
{'uuid': 'pipeline-1', 'config': json.dumps(pipeline_config)},
|
||||
)
|
||||
await conn.execute(
|
||||
sa.text('INSERT INTO agents (uuid, component_ref, config) VALUES (:uuid, :ref, :config)'),
|
||||
{'uuid': 'agent-1', 'ref': old_codex, 'config': json.dumps({'runner': {'id': old_codex}})},
|
||||
)
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'INSERT INTO workflows (uuid, definition, global_config) '
|
||||
'VALUES (:uuid, :definition, :global_config)'
|
||||
),
|
||||
{
|
||||
'uuid': 'workflow-1',
|
||||
'definition': json.dumps({'runner_id': old_dify}),
|
||||
'global_config': json.dumps({'runner_id': old_codex}),
|
||||
},
|
||||
)
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'INSERT INTO agent_runner_state ('
|
||||
'id, runner_id, binding_identity, scope, scope_key, state_key, bot_id, '
|
||||
'workspace_id, conversation_id, thread_id, actor_type, actor_id, subject_type, subject_id'
|
||||
') VALUES ('
|
||||
'1, :runner_id, :binding_identity, :scope, :scope_key, :state_key, :bot_id, '
|
||||
'NULL, :conversation_id, NULL, NULL, NULL, NULL, NULL'
|
||||
')'
|
||||
),
|
||||
{
|
||||
'runner_id': old_dify,
|
||||
'binding_identity': binding_id,
|
||||
'scope': 'conversation',
|
||||
'scope_key': old_scope_key,
|
||||
'state_key': 'external.conversation_id',
|
||||
'bot_id': 'bot-1',
|
||||
'conversation_id': 'conversation-1',
|
||||
},
|
||||
)
|
||||
for table_name in ('agent_run', 'agent_interaction'):
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
f'INSERT INTO {table_name} (id, runner_id, binding_id) VALUES (1, :runner_id, :binding_id)'
|
||||
),
|
||||
{'runner_id': old_dify, 'binding_id': binding_id},
|
||||
)
|
||||
for table_name in ('event_log', 'transcript'):
|
||||
await conn.execute(
|
||||
sa.text(f'INSERT INTO {table_name} (id, runner_id) VALUES (1, :runner_id)'),
|
||||
{'runner_id': old_codex},
|
||||
)
|
||||
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_stamp(sqlite_engine, '0014_interaction_delivery')
|
||||
await run_alembic_upgrade(sqlite_engine, 'head')
|
||||
|
||||
async with sqlite_engine.connect() as conn:
|
||||
pipeline_raw = (
|
||||
await conn.execute(sa.text("SELECT config FROM legacy_pipelines WHERE uuid = 'pipeline-1'"))
|
||||
).scalar_one()
|
||||
agent_row = (
|
||||
await conn.execute(sa.text("SELECT component_ref, config FROM agents WHERE uuid = 'agent-1'"))
|
||||
).one()
|
||||
workflow_row = (
|
||||
await conn.execute(sa.text("SELECT definition, global_config FROM workflows WHERE uuid = 'workflow-1'"))
|
||||
).one()
|
||||
state_row = (
|
||||
await conn.execute(
|
||||
sa.text('SELECT runner_id, binding_identity, scope_key FROM agent_runner_state WHERE id = 1')
|
||||
)
|
||||
).one()
|
||||
direct_runner_ids = {
|
||||
table_name: (
|
||||
await conn.execute(sa.text(f'SELECT runner_id FROM {table_name} WHERE id = 1'))
|
||||
).scalar_one()
|
||||
for table_name in ('agent_run', 'agent_interaction', 'event_log', 'transcript')
|
||||
}
|
||||
tables = await conn.run_sync(lambda sync_conn: set(sqlalchemy.inspect(sync_conn).get_table_names()))
|
||||
assert {'directory_projection_states', 'directory_projection_inbox'} <= tables
|
||||
|
||||
pipeline = json.loads(pipeline_raw)
|
||||
assert pipeline['ai']['runner']['id'] == new_dify
|
||||
assert pipeline['ai']['runner_config'][new_dify] == {'api-key': 'preserved-secret'}
|
||||
assert pipeline['ai']['runner_config'][new_codex] == {'workspace': 'K:/workspace'}
|
||||
assert old_dify not in pipeline['ai']['runner_config']
|
||||
assert agent_row.component_ref == new_codex
|
||||
assert json.loads(agent_row.config)['runner']['id'] == new_codex
|
||||
assert json.loads(workflow_row.definition)['runner_id'] == new_dify
|
||||
assert json.loads(workflow_row.global_config)['runner_id'] == new_codex
|
||||
assert state_row.runner_id == new_dify
|
||||
assert state_row.binding_identity == f'pipeline_pipeline-1_{new_dify}'
|
||||
assert state_row.scope_key != old_scope_key
|
||||
assert direct_runner_ids == {
|
||||
'agent_run': new_dify,
|
||||
'agent_interaction': new_dify,
|
||||
'event_log': new_codex,
|
||||
'transcript': new_codex,
|
||||
}
|
||||
assert await get_alembic_current(sqlite_engine) == '0015_official_runner_ids'
|
||||
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:
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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)}'))
|
||||
@@ -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()
|
||||
@@ -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()) == []
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -202,7 +202,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
|
||||
|
||||
Reference in New Issue
Block a user