mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-21 20:06:06 +00:00
feat(tenancy): harden shared cloud runtime boundaries
This commit is contained in:
Junyan Qin
@@ -16,19 +16,30 @@ from __future__ import annotations
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
import asyncio
|
||||
import contextlib
|
||||
import datetime
|
||||
import hashlib
|
||||
import typing
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from quart import Quart
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
|
||||
from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
|
||||
from sqlalchemy import text
|
||||
|
||||
from langbot.pkg.entity.persistence.base import Base
|
||||
from langbot.pkg.entity.persistence.user import User
|
||||
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
|
||||
from langbot.pkg.persistence.tenant_uow import TENANT_POLICY_NAME, TENANT_TABLE_COLUMNS, TenantUnitOfWork
|
||||
from langbot.pkg.persistence.tenant_uow import (
|
||||
TENANT_POLICY_NAME,
|
||||
TENANT_TABLE_COLUMNS,
|
||||
TenantUnitOfWork,
|
||||
TransactionRollbackOnlyError,
|
||||
)
|
||||
from langbot.pkg.persistence.alembic_runner import (
|
||||
run_alembic_upgrade,
|
||||
run_alembic_stamp,
|
||||
@@ -39,6 +50,32 @@ from alembic.config import Config
|
||||
from alembic.script import ScriptDirectory
|
||||
from langbot.pkg.utils import constants
|
||||
from langbot.pkg.workspace.collaboration import normalize_email
|
||||
from langbot.pkg.workspace.collaboration import WorkspaceCollaborationService
|
||||
from langbot.pkg.workspace.policy import CloudWorkspacePolicy
|
||||
from langbot.pkg.workspace.service import WorkspaceService
|
||||
from langbot.pkg.api.http.authz import Permission
|
||||
from langbot.pkg.api.http.controller import group as http_group
|
||||
from langbot.pkg.api.http.controller.groups.system import SystemRouterGroup
|
||||
from langbot.pkg.api.http.controller.groups.webhooks import WebhookRouterGroup
|
||||
from langbot.pkg.api.http.context import ExecutionContext, RequestContext
|
||||
from langbot.pkg.api.http.service.apikey import ApiKeyService
|
||||
from langbot.pkg.api.http.service.monitoring import MonitoringService
|
||||
from langbot.pkg.api.http.service.user import UserService
|
||||
from langbot.pkg.api.mcp.context import get_request_context as get_mcp_request_context
|
||||
from langbot.pkg.api.mcp.mount import MCPMount
|
||||
from langbot.pkg.entity.persistence.apikey import ApiKey
|
||||
from langbot.pkg.entity.persistence.metadata import WorkspaceMetadata
|
||||
from langbot.pkg.entity.persistence.monitoring import MonitoringFeedback
|
||||
from langbot.pkg.entity.persistence.workspace import (
|
||||
Workspace,
|
||||
WorkspaceExecutionState,
|
||||
WorkspaceMembership,
|
||||
)
|
||||
from langbot.pkg.platform.botmgr import PlatformManager
|
||||
from langbot.pkg.pipeline.pipelinemgr import PipelineManager
|
||||
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
|
||||
from langbot.pkg.provider.tools.loaders.mcp import MCPLoader
|
||||
from langbot.pkg.rag.knowledge.kbmgr import RAGManager
|
||||
|
||||
from .resource_migration_support import TENANT_TABLES, create_legacy_resource_schema
|
||||
|
||||
@@ -54,6 +91,42 @@ def _get_script_head() -> str:
|
||||
return ScriptDirectory.from_config(cfg).get_current_head()
|
||||
|
||||
|
||||
async def _grant_runtime_role_business_objects(
|
||||
conn: AsyncConnection,
|
||||
role_name: str,
|
||||
quote: typing.Callable[[str], str],
|
||||
) -> None:
|
||||
"""Mirror the release job's object ACLs without overgranting Alembic."""
|
||||
|
||||
business_tables = tuple(sorted({table.name for table in Base.metadata.tables.values()} | {'langbot_vectors'}))
|
||||
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(role_name)}'))
|
||||
await conn.execute(text(f'GRANT SELECT ON TABLE public.alembic_version TO {quote(role_name)}'))
|
||||
sequence_query = text(
|
||||
"""
|
||||
SELECT DISTINCT sequence.relname
|
||||
FROM pg_class sequence
|
||||
JOIN pg_namespace sequence_namespace ON sequence_namespace.oid = sequence.relnamespace
|
||||
JOIN pg_depend dependency
|
||||
ON dependency.classid = 'pg_class'::regclass
|
||||
AND dependency.objid = sequence.oid
|
||||
AND dependency.refclassid = 'pg_class'::regclass
|
||||
AND dependency.deptype IN ('a', 'i')
|
||||
JOIN pg_class business_table ON business_table.oid = dependency.refobjid
|
||||
JOIN pg_namespace table_namespace ON table_namespace.oid = business_table.relnamespace
|
||||
WHERE sequence.relkind = 'S'
|
||||
AND sequence_namespace.nspname = 'public'
|
||||
AND table_namespace.nspname = 'public'
|
||||
AND business_table.relname IN :table_names
|
||||
ORDER BY sequence.relname
|
||||
"""
|
||||
).bindparams(sa.bindparam('table_names', expanding=True))
|
||||
sequence_names = tuple((await conn.execute(sequence_query, {'table_names': business_tables})).scalars().all())
|
||||
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(role_name)}'))
|
||||
|
||||
|
||||
def _application_for_postgres_url(postgres_url: str, logger_name: str) -> SimpleNamespace:
|
||||
url = sa.engine.make_url(postgres_url)
|
||||
return SimpleNamespace(
|
||||
@@ -115,15 +188,20 @@ async def postgres_engine(postgres_url):
|
||||
@pytest.fixture
|
||||
async def clean_tables(postgres_engine):
|
||||
"""Drop all tables before and after each test for isolation."""
|
||||
# Drop all tables before test
|
||||
async with postgres_engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.drop_all)
|
||||
|
||||
async def drop_all_tables() -> None:
|
||||
# Alembic can create tables (notably langbot_vectors) outside the ORM
|
||||
# metadata, and legacy migration tests intentionally alter constraints.
|
||||
# Reflect the dedicated test schema instead of relying on stale ORM DDL.
|
||||
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 drop_all_tables()
|
||||
yield
|
||||
|
||||
# Drop all tables after test
|
||||
async with postgres_engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.drop_all)
|
||||
await drop_all_tables()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -442,10 +520,16 @@ class TestPostgreSQLResourceTenancyMigration:
|
||||
await conn.execute(
|
||||
text(
|
||||
'INSERT INTO plugin_settings '
|
||||
'(workspace_uuid, plugin_author, plugin_name, enabled) '
|
||||
"VALUES (:workspace_uuid, 'author', 'plugin', true)"
|
||||
'(workspace_uuid, plugin_author, plugin_name, enabled, '
|
||||
'installation_uuid, artifact_digest, runtime_revision) '
|
||||
"VALUES (:workspace_uuid, 'author', 'plugin', true, "
|
||||
':installation_uuid, :artifact_digest, 1)'
|
||||
),
|
||||
{'workspace_uuid': second_workspace_uuid},
|
||||
{
|
||||
'workspace_uuid': second_workspace_uuid,
|
||||
'installation_uuid': str(uuid.uuid4()),
|
||||
'artifact_digest': hashlib.sha256(b'test-plugin-artifact').hexdigest(),
|
||||
},
|
||||
)
|
||||
|
||||
with pytest.raises(IntegrityError):
|
||||
@@ -556,10 +640,7 @@ class TestPostgreSQLTenantRuntime:
|
||||
)
|
||||
)
|
||||
await conn.execute(text(f'GRANT USAGE ON SCHEMA public TO {quote(role_name)}'))
|
||||
await conn.execute(
|
||||
text(f'GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO {quote(role_name)}')
|
||||
)
|
||||
await conn.execute(text(f'GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO {quote(role_name)}'))
|
||||
await _grant_runtime_role_business_objects(conn, role_name, quote)
|
||||
created_roles.append(role_name)
|
||||
|
||||
try:
|
||||
@@ -703,6 +784,31 @@ class TestPostgreSQLTenantRuntime:
|
||||
cloud_manager.create_tables.assert_not_awaited()
|
||||
cloud_manager._run_alembic_migrations.assert_not_awaited()
|
||||
|
||||
with pytest.raises(TransactionRollbackOnlyError, match='after-commit work was cancelled'):
|
||||
async with cloud_manager.tenant_uow(workspace_a):
|
||||
duplicate_statement = sa.insert(WorkspaceMetadata).values(
|
||||
workspace_uuid=workspace_a,
|
||||
key='rollback-only-unique',
|
||||
value='must-not-commit',
|
||||
)
|
||||
await cloud_manager.execute_async(duplicate_statement)
|
||||
after_commit_gate = cloud_manager.create_after_commit_gate()
|
||||
assert after_commit_gate is not None
|
||||
try:
|
||||
await cloud_manager.execute_async(duplicate_statement)
|
||||
except IntegrityError:
|
||||
pass
|
||||
|
||||
assert after_commit_gate.cancelled()
|
||||
async with cloud_manager.tenant_uow(workspace_a):
|
||||
assert (
|
||||
await cloud_manager.execute_async(
|
||||
sa.select(sa.func.count())
|
||||
.select_from(WorkspaceMetadata)
|
||||
.where(WorkspaceMetadata.key == 'rollback-only-unique')
|
||||
)
|
||||
).scalar_one() == 0
|
||||
|
||||
superuser_manager = PersistenceManager(
|
||||
_application_for_postgres_url(postgres_url, 'postgres-superuser-runtime-test'),
|
||||
mode=PersistenceMode.CLOUD_RUNTIME,
|
||||
@@ -743,3 +849,644 @@ class TestPostgreSQLTenantRuntime:
|
||||
for role_name in reversed(created_roles):
|
||||
await conn.execute(text(f'DROP OWNED BY {quote(role_name)}'))
|
||||
await conn.execute(text(f'DROP ROLE {quote(role_name)}'))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cloud_discovery_http_and_transaction_contract(
|
||||
self,
|
||||
postgres_url,
|
||||
postgres_engine,
|
||||
clean_tables,
|
||||
clean_alembic_version,
|
||||
monkeypatch,
|
||||
):
|
||||
"""Exercise the complete request/discovery path as a non-owner role."""
|
||||
|
||||
instance_uuid = 'cloud-request-rls-test'
|
||||
other_instance_uuid = 'other-cloud-instance'
|
||||
workspace_a = '31000000-0000-0000-0000-000000000001'
|
||||
workspace_b = '32000000-0000-0000-0000-000000000002'
|
||||
workspace_other = '33000000-0000-0000-0000-000000000003'
|
||||
workspace_fenced = '34000000-0000-0000-0000-000000000004'
|
||||
account_a_uuid = '41000000-0000-0000-0000-000000000001'
|
||||
shared_account_uuid = '42000000-0000-0000-0000-000000000002'
|
||||
active_secret = 'lbk_active-request-key'
|
||||
revoked_secret = 'lbk_revoked-request-key'
|
||||
expired_secret = 'lbk_expired-request-key'
|
||||
role_name = f'lb_request_{uuid.uuid4().hex[:12]}'
|
||||
role_password = f'Lb{uuid.uuid4().hex}'
|
||||
quote = postgres_engine.dialect.identifier_preparer.quote
|
||||
managers: list[PersistenceManager] = []
|
||||
role_created = False
|
||||
|
||||
_restore_postgres_manager_registry(monkeypatch)
|
||||
monkeypatch.setattr(constants, 'instance_id', instance_uuid)
|
||||
release_manager = PersistenceManager(
|
||||
_application_for_postgres_url(postgres_url, 'postgres-request-release-test'),
|
||||
mode=PersistenceMode.RELEASE_MIGRATION,
|
||||
)
|
||||
managers.append(release_manager)
|
||||
|
||||
def role_url() -> str:
|
||||
return (
|
||||
sa.engine.make_url(postgres_url)
|
||||
.set(username=role_name, password=role_password)
|
||||
.render_as_string(hide_password=False)
|
||||
)
|
||||
|
||||
async def seed_workspace(
|
||||
workspace_uuid: str,
|
||||
*,
|
||||
target_instance: str,
|
||||
state: str = 'active',
|
||||
write_fenced: bool = False,
|
||||
) -> None:
|
||||
async with release_manager.tenant_uow(workspace_uuid) as uow:
|
||||
uow.session.add(
|
||||
Workspace(
|
||||
uuid=workspace_uuid,
|
||||
instance_uuid=target_instance,
|
||||
name=workspace_uuid[-4:],
|
||||
slug=f'workspace-{workspace_uuid[-4:]}',
|
||||
type='team',
|
||||
status='active',
|
||||
source='cloud_projection',
|
||||
projection_revision=1,
|
||||
)
|
||||
)
|
||||
await uow.session.flush()
|
||||
uow.session.add(
|
||||
WorkspaceExecutionState(
|
||||
workspace_uuid=workspace_uuid,
|
||||
instance_uuid=target_instance,
|
||||
active_generation=1,
|
||||
state=state,
|
||||
write_fenced=write_fenced,
|
||||
source='cloud',
|
||||
desired_state_revision=1,
|
||||
)
|
||||
)
|
||||
uow.session.add(
|
||||
WorkspaceMetadata(
|
||||
workspace_uuid=workspace_uuid,
|
||||
key='tenant-marker',
|
||||
value=workspace_uuid,
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
await release_manager.initialize()
|
||||
async with release_manager.get_db_engine().begin() as conn:
|
||||
await conn.execute(
|
||||
sa.insert(User),
|
||||
[
|
||||
{
|
||||
'uuid': account_a_uuid,
|
||||
'user': 'account-a@example.com',
|
||||
'normalized_email': 'account-a@example.com',
|
||||
'password': 'closed-directory',
|
||||
'account_type': 'local',
|
||||
'status': 'active',
|
||||
'source': 'cloud_projection',
|
||||
'projection_revision': 1,
|
||||
},
|
||||
{
|
||||
'uuid': shared_account_uuid,
|
||||
'user': 'shared@example.com',
|
||||
'normalized_email': 'shared@example.com',
|
||||
'password': 'closed-directory',
|
||||
'account_type': 'local',
|
||||
'status': 'active',
|
||||
'source': 'cloud_projection',
|
||||
'projection_revision': 1,
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
await seed_workspace(workspace_a, target_instance=instance_uuid)
|
||||
await seed_workspace(workspace_b, target_instance=instance_uuid)
|
||||
await seed_workspace(workspace_other, target_instance=other_instance_uuid)
|
||||
await seed_workspace(
|
||||
workspace_fenced,
|
||||
target_instance=instance_uuid,
|
||||
write_fenced=True,
|
||||
)
|
||||
|
||||
async with release_manager.tenant_uow(workspace_a) as uow:
|
||||
uow.session.add_all(
|
||||
[
|
||||
WorkspaceMembership(
|
||||
uuid=str(uuid.uuid4()),
|
||||
workspace_uuid=workspace_a,
|
||||
account_uuid=account_a_uuid,
|
||||
role='owner',
|
||||
status='active',
|
||||
projection_revision=1,
|
||||
),
|
||||
WorkspaceMembership(
|
||||
uuid=str(uuid.uuid4()),
|
||||
workspace_uuid=workspace_a,
|
||||
account_uuid=shared_account_uuid,
|
||||
role='viewer',
|
||||
status='active',
|
||||
projection_revision=1,
|
||||
),
|
||||
ApiKey(
|
||||
uuid=str(uuid.uuid4()),
|
||||
workspace_uuid=workspace_a,
|
||||
name='active-key',
|
||||
key_hash=hashlib.sha256(active_secret.encode()).hexdigest(),
|
||||
scopes=[Permission.WORKSPACE_VIEW.value],
|
||||
status='active',
|
||||
),
|
||||
]
|
||||
)
|
||||
async with release_manager.tenant_uow(workspace_b) as uow:
|
||||
uow.session.add_all(
|
||||
[
|
||||
WorkspaceMembership(
|
||||
uuid=str(uuid.uuid4()),
|
||||
workspace_uuid=workspace_b,
|
||||
account_uuid=shared_account_uuid,
|
||||
role='viewer',
|
||||
status='active',
|
||||
projection_revision=1,
|
||||
),
|
||||
ApiKey(
|
||||
uuid=str(uuid.uuid4()),
|
||||
workspace_uuid=workspace_b,
|
||||
name='revoked-key',
|
||||
key_hash=hashlib.sha256(revoked_secret.encode()).hexdigest(),
|
||||
scopes=[Permission.WORKSPACE_VIEW.value],
|
||||
status='revoked',
|
||||
),
|
||||
ApiKey(
|
||||
uuid=str(uuid.uuid4()),
|
||||
workspace_uuid=workspace_b,
|
||||
name='expired-key',
|
||||
key_hash=hashlib.sha256(expired_secret.encode()).hexdigest(),
|
||||
scopes=[Permission.WORKSPACE_VIEW.value],
|
||||
status='active',
|
||||
expires_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
|
||||
- datetime.timedelta(minutes=1),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
async with postgres_engine.connect() as conn:
|
||||
await conn.execute(text(f"CREATE ROLE {quote(role_name)} LOGIN PASSWORD '{role_password}'"))
|
||||
role_created = True
|
||||
await conn.execute(
|
||||
text(
|
||||
f'GRANT CONNECT ON DATABASE {quote(sa.engine.make_url(postgres_url).database)} '
|
||||
f'TO {quote(role_name)}'
|
||||
)
|
||||
)
|
||||
await conn.execute(text(f'GRANT USAGE ON SCHEMA public TO {quote(role_name)}'))
|
||||
await _grant_runtime_role_business_objects(conn, role_name, quote)
|
||||
|
||||
runtime_application = _application_for_postgres_url(role_url(), 'postgres-request-runtime-test')
|
||||
runtime_application.instance_config.data.update(
|
||||
{
|
||||
'system': {
|
||||
'jwt': {'secret': 'postgres-request-jwt', 'expire': 3600},
|
||||
},
|
||||
'api': {'global_api_key': ''},
|
||||
}
|
||||
)
|
||||
runtime_application.logger = logging.getLogger('postgres-request-runtime-test')
|
||||
cloud_manager = PersistenceManager(runtime_application, mode=PersistenceMode.CLOUD_RUNTIME)
|
||||
managers.append(cloud_manager)
|
||||
await cloud_manager.initialize()
|
||||
|
||||
runtime_application.persistence_mgr = cloud_manager
|
||||
cloud_manager.ap = runtime_application
|
||||
runtime_application.workspace_service = WorkspaceService(
|
||||
runtime_application,
|
||||
policy=CloudWorkspacePolicy(),
|
||||
instance_uuid=instance_uuid,
|
||||
)
|
||||
runtime_application.workspace_collaboration_service = WorkspaceCollaborationService(
|
||||
runtime_application,
|
||||
runtime_application.workspace_service,
|
||||
policy=CloudWorkspacePolicy(),
|
||||
)
|
||||
runtime_application.user_service = UserService(runtime_application)
|
||||
runtime_application.apikey_service = ApiKeyService(runtime_application)
|
||||
runtime_application.monitoring_service = MonitoringService(runtime_application)
|
||||
|
||||
bindings = await runtime_application.workspace_service.list_active_execution_bindings()
|
||||
assert {binding.workspace_uuid for binding in bindings} == {workspace_a, workspace_b}
|
||||
|
||||
# No scope is an application error before SQL reaches PostgreSQL.
|
||||
with pytest.raises(RuntimeError, match='explicit Workspace or discovery'):
|
||||
await cloud_manager.execute_async(sa.select(WorkspaceMetadata))
|
||||
|
||||
# Discovery exposes only its index rows and cannot write.
|
||||
with pytest.raises(TransactionRollbackOnlyError, match='transaction was rolled back'):
|
||||
async with cloud_manager.account_discovery_uow(shared_account_uuid) as discovery:
|
||||
assert set(
|
||||
(
|
||||
await discovery.session.scalars(
|
||||
sa.select(WorkspaceMembership.workspace_uuid).order_by(
|
||||
WorkspaceMembership.workspace_uuid
|
||||
)
|
||||
)
|
||||
).all()
|
||||
) == {workspace_a, workspace_b}
|
||||
with pytest.raises(sa.exc.DBAPIError):
|
||||
await discovery.session.execute(
|
||||
sa.insert(WorkspaceMembership).values(
|
||||
uuid=str(uuid.uuid4()),
|
||||
workspace_uuid=workspace_a,
|
||||
account_uuid=shared_account_uuid,
|
||||
role='viewer',
|
||||
status='active',
|
||||
projection_revision=1,
|
||||
)
|
||||
)
|
||||
|
||||
async with cloud_manager.api_key_discovery_uow(
|
||||
hashlib.sha256(active_secret.encode()).hexdigest()
|
||||
) as discovery:
|
||||
assert await discovery.session.scalar(sa.select(ApiKey.workspace_uuid)) == workspace_a
|
||||
update_result = await discovery.session.execute(
|
||||
sa.update(ApiKey)
|
||||
.where(ApiKey.key_hash == hashlib.sha256(active_secret.encode()).hexdigest())
|
||||
.values(name='discovery-must-not-write')
|
||||
)
|
||||
assert update_result.rowcount == 0
|
||||
async with cloud_manager.api_key_discovery_uow(
|
||||
hashlib.sha256(revoked_secret.encode()).hexdigest()
|
||||
) as discovery:
|
||||
assert await discovery.session.scalar(sa.select(ApiKey.workspace_uuid)) is None
|
||||
async with cloud_manager.api_key_discovery_uow(
|
||||
hashlib.sha256(expired_secret.encode()).hexdigest()
|
||||
) as discovery:
|
||||
assert await discovery.session.scalar(sa.select(ApiKey.workspace_uuid)) is None
|
||||
|
||||
async with cloud_manager.instance_discovery_uow(instance_uuid) as discovery:
|
||||
assert set(
|
||||
(await discovery.session.scalars(sa.select(WorkspaceExecutionState.workspace_uuid))).all()
|
||||
) == {workspace_a, workspace_b}
|
||||
update_result = await discovery.session.execute(
|
||||
sa.update(WorkspaceExecutionState)
|
||||
.where(WorkspaceExecutionState.workspace_uuid == workspace_a)
|
||||
.values(write_fenced=True)
|
||||
)
|
||||
assert update_result.rowcount == 0
|
||||
# Instance discovery deliberately cannot see business rows.
|
||||
assert (await discovery.session.execute(sa.select(WorkspaceMetadata))).all() == []
|
||||
|
||||
platform_manager = PlatformManager(runtime_application)
|
||||
platform_manager._load_workspace_bots = AsyncMock()
|
||||
await platform_manager.load_bots_from_db()
|
||||
assert {call.args[0] for call in platform_manager._load_workspace_bots.await_args_list} == {
|
||||
workspace_a,
|
||||
workspace_b,
|
||||
}
|
||||
|
||||
# Every startup cache loader traverses the instance index first,
|
||||
# then reads business resources in one tenant transaction at a time.
|
||||
await ModelManager(runtime_application).load_models_from_db()
|
||||
await PipelineManager(runtime_application).load_pipelines_from_db()
|
||||
await MCPLoader(runtime_application).load_mcp_servers_from_db()
|
||||
await RAGManager(runtime_application).load_knowledge_bases_from_db()
|
||||
|
||||
# An omitted Workspace predicate remains isolated by RLS.
|
||||
async with cloud_manager.tenant_uow(workspace_a):
|
||||
values = (await cloud_manager.execute_async(sa.select(WorkspaceMetadata.value))).scalars().all()
|
||||
assert values == [workspace_a]
|
||||
|
||||
async def read_tenant_repeatedly(workspace_uuid: str) -> list[str]:
|
||||
observed: list[str] = []
|
||||
for _ in range(5):
|
||||
async with cloud_manager.tenant_uow(workspace_uuid):
|
||||
observed.extend(
|
||||
(await cloud_manager.execute_async(sa.select(WorkspaceMetadata.value))).scalars().all()
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
return observed
|
||||
|
||||
observed_a, observed_b = await asyncio.gather(
|
||||
read_tenant_repeatedly(workspace_a),
|
||||
read_tenant_repeatedly(workspace_b),
|
||||
)
|
||||
assert observed_a == [workspace_a] * 5
|
||||
assert observed_b == [workspace_b] * 5
|
||||
|
||||
accesses = await runtime_application.workspace_collaboration_service.list_account_workspaces(
|
||||
shared_account_uuid
|
||||
)
|
||||
assert {access.workspace.uuid for access in accesses} == {workspace_a, workspace_b}
|
||||
assert await runtime_application.apikey_service.authenticate_api_key(revoked_secret) is None
|
||||
assert await runtime_application.apikey_service.authenticate_api_key(expired_secret) is None
|
||||
active_identity = await runtime_application.apikey_service.authenticate_api_key(active_secret)
|
||||
assert active_identity is not None
|
||||
assert active_identity.workspace_uuid == workspace_a
|
||||
|
||||
async def record_feedback(feedback_type: int) -> str | None:
|
||||
async with cloud_manager.tenant_scope(workspace_a):
|
||||
return await runtime_application.monitoring_service.record_feedback(
|
||||
ExecutionContext(
|
||||
instance_uuid=instance_uuid,
|
||||
workspace_uuid=workspace_a,
|
||||
placement_generation=1,
|
||||
),
|
||||
feedback_id='concurrent-feedback',
|
||||
feedback_type=feedback_type,
|
||||
)
|
||||
|
||||
feedback_ids = await asyncio.gather(record_feedback(1), record_feedback(2))
|
||||
assert feedback_ids[0] == feedback_ids[1]
|
||||
async with cloud_manager.tenant_uow(workspace_a) as uow:
|
||||
feedback_rows = (
|
||||
(
|
||||
await uow.execute(
|
||||
sa.select(MonitoringFeedback).where(MonitoringFeedback.feedback_id == 'concurrent-feedback')
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
assert len(feedback_rows) == 1
|
||||
assert feedback_rows[0].feedback_type in {1, 2}
|
||||
await uow.execute(
|
||||
sa.delete(MonitoringFeedback).where(MonitoringFeedback.feedback_id == 'concurrent-feedback')
|
||||
)
|
||||
|
||||
class TenantRuntimeRouter(http_group.RouterGroup):
|
||||
name = 'postgres-tenant-runtime'
|
||||
path = '/tenant-runtime'
|
||||
|
||||
async def initialize(self) -> None:
|
||||
@self.route('/account', permission=Permission.WORKSPACE_VIEW)
|
||||
async def account_route(request_context: RequestContext):
|
||||
assert self.ap.persistence_mgr.current_session() is None
|
||||
if self.quart_app.config.get('FORCE_HANDLER_FAILURE'):
|
||||
async with self.ap.persistence_mgr.tenant_uow(request_context.workspace_uuid):
|
||||
await self.ap.persistence_mgr.execute_async(
|
||||
sa.insert(WorkspaceMetadata).values(
|
||||
workspace_uuid=request_context.workspace_uuid,
|
||||
key='rolled-back-handler',
|
||||
value='must-not-commit',
|
||||
)
|
||||
)
|
||||
raise RuntimeError('forced handler failure')
|
||||
values = (
|
||||
(await self.ap.persistence_mgr.execute_async(sa.select(WorkspaceMetadata.value)))
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
assert self.ap.persistence_mgr.current_session() is None
|
||||
return self.success(data={'workspace_uuid': request_context.workspace_uuid, 'values': values})
|
||||
|
||||
@self.route('/key', auth_type=http_group.AuthType.API_KEY)
|
||||
async def key_route(request_context: RequestContext):
|
||||
values = (
|
||||
(await self.ap.persistence_mgr.execute_async(sa.select(WorkspaceMetadata.value)))
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
return self.success(data={'workspace_uuid': request_context.workspace_uuid, 'values': values})
|
||||
|
||||
@self.route('/bootstrap', auth_type=http_group.AuthType.ACCOUNT_TOKEN)
|
||||
async def bootstrap_route(user_email: str):
|
||||
account = await self.ap.user_service.get_user_by_email(user_email)
|
||||
accesses = await self.ap.workspace_collaboration_service.list_account_workspaces(account.uuid)
|
||||
return self.success(data=sorted(access.workspace.uuid for access in accesses))
|
||||
|
||||
quart_app = Quart(__name__)
|
||||
await TenantRuntimeRouter(runtime_application, quart_app).initialize()
|
||||
|
||||
webhook_bot_uuid = str(uuid.uuid4())
|
||||
|
||||
class TenantAwareWebhookAdapter:
|
||||
async def handle_unified_webhook(self, **_kwargs):
|
||||
assert cloud_manager.current_session() is None
|
||||
await asyncio.sleep(0)
|
||||
assert cloud_manager.current_session() is None
|
||||
values = (
|
||||
(
|
||||
await cloud_manager.execute_async(
|
||||
sa.select(WorkspaceMetadata.value).where(WorkspaceMetadata.key == 'tenant-marker')
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
assert cloud_manager.current_session() is None
|
||||
return {'values': values}
|
||||
|
||||
runtime_application.platform_mgr = SimpleNamespace(
|
||||
resolve_public_bot=AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
workspace_uuid=workspace_a,
|
||||
placement_generation=1,
|
||||
enable=True,
|
||||
adapter=TenantAwareWebhookAdapter(),
|
||||
)
|
||||
)
|
||||
)
|
||||
await WebhookRouterGroup(runtime_application, quart_app).initialize()
|
||||
await SystemRouterGroup(runtime_application, quart_app).initialize()
|
||||
client = quart_app.test_client()
|
||||
account_a = await runtime_application.user_service.get_user_by_uuid(account_a_uuid)
|
||||
shared_account = await runtime_application.user_service.get_user_by_uuid(shared_account_uuid)
|
||||
assert account_a is not None and shared_account is not None
|
||||
account_token = await runtime_application.user_service.generate_jwt_token(account_a)
|
||||
shared_token = await runtime_application.user_service.generate_jwt_token(shared_account)
|
||||
|
||||
async with cloud_manager.tenant_uow(workspace_a):
|
||||
await cloud_manager.execute_async(
|
||||
sa.insert(WorkspaceMetadata).values(
|
||||
workspace_uuid=workspace_a,
|
||||
key='wizard_status',
|
||||
value='completed',
|
||||
)
|
||||
)
|
||||
response = await client.get(
|
||||
'/api/v1/system/info',
|
||||
headers={
|
||||
'Authorization': f'Bearer {account_token}',
|
||||
'X-Workspace-Id': workspace_a,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert (await response.get_json())['data']['wizard_status'] == 'completed'
|
||||
async with cloud_manager.tenant_uow(workspace_a):
|
||||
await cloud_manager.execute_async(
|
||||
sa.delete(WorkspaceMetadata).where(
|
||||
WorkspaceMetadata.workspace_uuid == workspace_a,
|
||||
WorkspaceMetadata.key == 'wizard_status',
|
||||
)
|
||||
)
|
||||
|
||||
response = await client.post(f'/bots/{webhook_bot_uuid}')
|
||||
assert response.status_code == 200
|
||||
assert await response.get_json() == {'values': [workspace_a]}
|
||||
|
||||
response = await client.get(
|
||||
'/tenant-runtime/account',
|
||||
headers={
|
||||
'Authorization': f'Bearer {account_token}',
|
||||
'X-Workspace-Id': workspace_a,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert (await response.get_json())['data'] == {
|
||||
'workspace_uuid': workspace_a,
|
||||
'values': [workspace_a],
|
||||
}
|
||||
|
||||
response = await client.get(
|
||||
'/tenant-runtime/account',
|
||||
headers={
|
||||
'Authorization': f'Bearer {account_token}',
|
||||
'X-Workspace-Id': workspace_b,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 404
|
||||
response = await client.get(
|
||||
'/tenant-runtime/account',
|
||||
headers={'Authorization': f'Bearer {shared_token}'},
|
||||
)
|
||||
assert response.status_code == 404
|
||||
|
||||
response = await client.get(
|
||||
'/tenant-runtime/bootstrap',
|
||||
headers={'Authorization': f'Bearer {shared_token}'},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert (await response.get_json())['data'] == [workspace_a, workspace_b]
|
||||
|
||||
response = await client.get(
|
||||
'/tenant-runtime/key',
|
||||
headers={'X-API-Key': active_secret, 'X-Workspace-Id': workspace_b},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert (await response.get_json())['data'] == {
|
||||
'workspace_uuid': workspace_a,
|
||||
'values': [workspace_a],
|
||||
}
|
||||
|
||||
# The parallel MCP ASGI entrypoint authenticates the same key and
|
||||
# retains only a trusted Workspace scope for the full tool request.
|
||||
# Each DB call gets its own short RLS transaction.
|
||||
mcp_observation: dict[str, typing.Any] = {}
|
||||
|
||||
async def fake_mcp_asgi(scope, receive, send):
|
||||
del scope, receive
|
||||
context = get_mcp_request_context()
|
||||
assert cloud_manager.current_session() is None
|
||||
values = (await cloud_manager.execute_async(sa.select(WorkspaceMetadata.value))).scalars().all()
|
||||
assert cloud_manager.current_session() is None
|
||||
await asyncio.sleep(0)
|
||||
assert cloud_manager.current_session() is None
|
||||
repeated_values = (
|
||||
(await cloud_manager.execute_async(sa.select(WorkspaceMetadata.value))).scalars().all()
|
||||
)
|
||||
assert repeated_values == values
|
||||
assert cloud_manager.current_session() is None
|
||||
mcp_observation.update(
|
||||
workspace_uuid=context.workspace_uuid,
|
||||
values=values,
|
||||
)
|
||||
await send({'type': 'http.response.start', 'status': 200, 'headers': []})
|
||||
await send({'type': 'http.response.body', 'body': b'{}'})
|
||||
|
||||
async def unused_quart_asgi(scope, receive, send): # pragma: no cover - routing assertion
|
||||
del scope, receive, send
|
||||
raise AssertionError('MCP request was routed to Quart')
|
||||
|
||||
mount = MCPMount.__new__(MCPMount)
|
||||
mount.ap = runtime_application
|
||||
mount._mcp_asgi = fake_mcp_asgi
|
||||
sent_messages: list[dict[str, typing.Any]] = []
|
||||
|
||||
async def receive():
|
||||
return {'type': 'http.request', 'body': b'', 'more_body': False}
|
||||
|
||||
async def send(message):
|
||||
sent_messages.append(message)
|
||||
|
||||
await mount.wrap(unused_quart_asgi)(
|
||||
{
|
||||
'type': 'http',
|
||||
'path': '/mcp',
|
||||
'headers': [
|
||||
(b'x-api-key', active_secret.encode()),
|
||||
(b'x-workspace-id', workspace_b.encode()),
|
||||
],
|
||||
},
|
||||
receive,
|
||||
send,
|
||||
)
|
||||
assert sent_messages[0]['status'] == 200
|
||||
assert mcp_observation == {'workspace_uuid': workspace_a, 'values': [workspace_a]}
|
||||
|
||||
original_api_key_discovery = cloud_manager.api_key_discovery_uow
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def revoke_after_discovery(key_hash: str):
|
||||
async with original_api_key_discovery(key_hash) as discovery:
|
||||
yield discovery
|
||||
async with cloud_manager.tenant_uow(workspace_a):
|
||||
await cloud_manager.execute_async(
|
||||
sa.update(ApiKey).where(ApiKey.key_hash == key_hash).values(status='revoked')
|
||||
)
|
||||
|
||||
monkeypatch.setattr(cloud_manager, 'api_key_discovery_uow', revoke_after_discovery)
|
||||
assert await runtime_application.apikey_service.authenticate_api_key(active_secret) is None
|
||||
monkeypatch.setattr(cloud_manager, 'api_key_discovery_uow', original_api_key_discovery)
|
||||
|
||||
quart_app.config['FORCE_HANDLER_FAILURE'] = True
|
||||
response = await client.get(
|
||||
'/tenant-runtime/account',
|
||||
headers={
|
||||
'Authorization': f'Bearer {account_token}',
|
||||
'X-Workspace-Id': workspace_a,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 500
|
||||
quart_app.config['FORCE_HANDLER_FAILURE'] = False
|
||||
async with cloud_manager.tenant_uow(workspace_a):
|
||||
assert (
|
||||
await cloud_manager.execute_async(
|
||||
sa.select(sa.func.count())
|
||||
.select_from(WorkspaceMetadata)
|
||||
.where(WorkspaceMetadata.key == 'rolled-back-handler')
|
||||
)
|
||||
).scalar_one() == 0
|
||||
|
||||
# Transaction-local settings are gone when pooled connections are reused.
|
||||
async with cloud_manager.get_db_engine().connect() as conn:
|
||||
assert await conn.scalar(text("SELECT current_setting('langbot.workspace_uuid', true)")) in (
|
||||
None,
|
||||
'',
|
||||
)
|
||||
assert await conn.scalar(text('SELECT COUNT(*) FROM workspace_metadata')) == 0
|
||||
|
||||
# Runtime validation rejects both extra permissive policies and a
|
||||
# modified expression even if the expected policy name remains.
|
||||
async with postgres_engine.connect() as conn:
|
||||
await conn.execute(text('CREATE POLICY injected_policy ON bots FOR SELECT USING (true)'))
|
||||
with pytest.raises(RuntimeError, match='policy set does not match'):
|
||||
await cloud_manager._validate_postgres_tenant_schema(validate_runtime_role=True)
|
||||
async with postgres_engine.connect() as conn:
|
||||
await conn.execute(text('DROP POLICY injected_policy ON bots'))
|
||||
await conn.execute(text('DROP POLICY langbot_workspace_isolation ON workspace_metadata'))
|
||||
await conn.execute(
|
||||
text(
|
||||
'CREATE POLICY langbot_workspace_isolation ON workspace_metadata '
|
||||
'FOR ALL TO PUBLIC USING (true) WITH CHECK (true)'
|
||||
)
|
||||
)
|
||||
with pytest.raises(RuntimeError, match='policy definitions are invalid'):
|
||||
await cloud_manager._validate_postgres_tenant_schema(validate_runtime_role=True)
|
||||
finally:
|
||||
for manager in reversed(managers):
|
||||
await _dispose_manager(manager)
|
||||
if role_created:
|
||||
async with postgres_engine.connect() as conn:
|
||||
await conn.execute(text(f'DROP OWNED BY {quote(role_name)}'))
|
||||
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(role_name)}'))
|
||||
|
||||
Reference in New Issue
Block a user