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:
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
@@ -101,6 +102,60 @@ async def test_public_webhook_error_uses_same_generic_error_contract():
|
||||
assert 'do-not-return' not in (await response.get_data(as_text=True))
|
||||
|
||||
|
||||
async def test_public_webhook_carries_scope_without_holding_database_session():
|
||||
class ScopeOnlyPersistenceManager:
|
||||
mode = SimpleNamespace(value='cloud_runtime')
|
||||
|
||||
def __init__(self):
|
||||
self.active_workspace = None
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def tenant_scope(self, workspace_uuid):
|
||||
self.active_workspace = workspace_uuid
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self.active_workspace = None
|
||||
|
||||
def current_session(self):
|
||||
return None
|
||||
|
||||
persistence_mgr = ScopeOnlyPersistenceManager()
|
||||
workspace_uuid = '00000000-0000-0000-0000-00000000000a'
|
||||
bot_uuid = '11111111-1111-4111-8111-111111111111'
|
||||
|
||||
class Adapter:
|
||||
async def handle_unified_webhook(self, **_kwargs):
|
||||
assert persistence_mgr.active_workspace == workspace_uuid
|
||||
assert persistence_mgr.current_session() is None
|
||||
return {'ok': True}
|
||||
|
||||
async def get_execution_binding(resolved_workspace_uuid, expected_generation=None):
|
||||
assert resolved_workspace_uuid == workspace_uuid
|
||||
assert expected_generation == 4
|
||||
|
||||
runtime_bot = SimpleNamespace(
|
||||
workspace_uuid=workspace_uuid,
|
||||
placement_generation=4,
|
||||
enable=True,
|
||||
adapter=Adapter(),
|
||||
)
|
||||
application = SimpleNamespace(
|
||||
logger=Mock(),
|
||||
persistence_mgr=persistence_mgr,
|
||||
platform_mgr=SimpleNamespace(resolve_public_bot=AsyncMock(return_value=runtime_bot)),
|
||||
workspace_service=SimpleNamespace(get_execution_binding=get_execution_binding),
|
||||
)
|
||||
quart_app = quart.Quart(__name__)
|
||||
await WebhookRouterGroup(application, quart_app).initialize()
|
||||
|
||||
response = await quart_app.test_client().post(f'/bots/{bot_uuid}')
|
||||
|
||||
assert response.status_code == 200
|
||||
assert await response.get_json() == {'ok': True}
|
||||
assert persistence_mgr.active_workspace is None
|
||||
|
||||
|
||||
async def test_authentication_failure_does_not_return_internal_exception_text():
|
||||
logger = Mock()
|
||||
application = SimpleNamespace(
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
import quart
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from langbot.pkg.api.http.controller import group
|
||||
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
|
||||
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
async def test_authenticated_route_does_not_hold_database_session_during_external_wait():
|
||||
entered = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
observations: list[bool] = []
|
||||
engine = create_async_engine('sqlite+aiosqlite:///:memory:')
|
||||
persistence = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||
persistence.db = SimpleNamespace(get_engine=lambda: engine)
|
||||
|
||||
class BlockingRouter(group.RouterGroup):
|
||||
name = 'blocking-route-test'
|
||||
path = '/blocking-route-test'
|
||||
|
||||
async def initialize(self) -> None:
|
||||
@self.route('', methods=['GET'], auth_type=group.AuthType.USER_TOKEN)
|
||||
async def _():
|
||||
observations.append(persistence.current_session() is None)
|
||||
entered.set()
|
||||
await release.wait()
|
||||
observations.append(persistence.current_session() is None)
|
||||
return self.success(data={})
|
||||
|
||||
account = SimpleNamespace(uuid='account-a', user='owner@example.com')
|
||||
access = SimpleNamespace(
|
||||
execution=SimpleNamespace(instance_uuid='instance-a', placement_generation=1),
|
||||
workspace=SimpleNamespace(uuid='workspace-a'),
|
||||
membership=SimpleNamespace(uuid='membership-a', role='owner', projection_revision=1),
|
||||
)
|
||||
application = SimpleNamespace(
|
||||
persistence_mgr=persistence,
|
||||
deployment=SimpleNamespace(multi_workspace_enabled=False),
|
||||
user_service=SimpleNamespace(get_authenticated_account=AsyncMock(return_value=account)),
|
||||
workspace_collaboration_service=SimpleNamespace(resolve_account_workspace=AsyncMock(return_value=access)),
|
||||
logger=Mock(),
|
||||
)
|
||||
quart_app = quart.Quart(__name__)
|
||||
await BlockingRouter(application, quart_app).initialize()
|
||||
client = quart_app.test_client()
|
||||
|
||||
request = asyncio.create_task(
|
||||
client.get(
|
||||
'/blocking-route-test',
|
||||
headers={'Authorization': 'Bearer token', 'X-Workspace-Id': 'workspace-a'},
|
||||
)
|
||||
)
|
||||
try:
|
||||
await entered.wait()
|
||||
assert observations == [True]
|
||||
release.set()
|
||||
response = await request
|
||||
assert response.status_code == 200
|
||||
assert observations == [True, True]
|
||||
finally:
|
||||
release.set()
|
||||
if not request.done():
|
||||
await request
|
||||
await engine.dispose()
|
||||
@@ -166,6 +166,29 @@ async def test_revoked_expired_and_unknown_keys_fail_closed(api_key_context):
|
||||
assert await service.authenticate_api_key(expired_secret) is None
|
||||
|
||||
|
||||
async def test_revoke_winning_last_used_update_race_fails_authentication(api_key_context):
|
||||
application, service, context, _engine = api_key_context
|
||||
created = await service.create_api_key(context, 'Racing revoke')
|
||||
original_execute = application.persistence_mgr.execute_async
|
||||
injected_revoke = False
|
||||
|
||||
async def execute_with_revoke(statement, *args, **kwargs):
|
||||
nonlocal injected_revoke
|
||||
if (
|
||||
not injected_revoke
|
||||
and isinstance(statement, sqlalchemy.sql.dml.Update)
|
||||
and statement.table.name == ApiKey.__tablename__
|
||||
):
|
||||
injected_revoke = True
|
||||
await original_execute(sqlalchemy.update(ApiKey).where(ApiKey.id == created['id']).values(status='revoked'))
|
||||
return await original_execute(statement, *args, **kwargs)
|
||||
|
||||
application.persistence_mgr.execute_async = execute_with_revoke
|
||||
|
||||
assert await service.authenticate_api_key(created['key']) is None
|
||||
assert injected_revoke is True
|
||||
|
||||
|
||||
async def test_cross_workspace_crud_and_secret_guessing_are_isolated(api_key_context):
|
||||
application, service, first_context, engine = api_key_context
|
||||
second_workspace_uuid = str(uuid.uuid4())
|
||||
|
||||
@@ -13,6 +13,7 @@ Source: src/langbot/pkg/api/http/service/maintenance.py
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, Mock, patch, MagicMock
|
||||
from types import SimpleNamespace
|
||||
@@ -196,6 +197,51 @@ class TestMaintenanceServiceCleanupExpiredFiles:
|
||||
assert ap.logger.warning.called
|
||||
assert 'uploaded_files' in result
|
||||
|
||||
async def test_cloud_cleanup_carries_scope_without_holding_database_session(self):
|
||||
class ScopeOnlyPersistenceManager:
|
||||
mode = SimpleNamespace(value='cloud_runtime')
|
||||
|
||||
def __init__(self):
|
||||
self.active_workspace = None
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def tenant_scope(self, workspace_uuid):
|
||||
self.active_workspace = workspace_uuid
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self.active_workspace = None
|
||||
|
||||
def current_session(self):
|
||||
return None
|
||||
|
||||
persistence_mgr = ScopeOnlyPersistenceManager()
|
||||
application = SimpleNamespace(
|
||||
persistence_mgr=persistence_mgr,
|
||||
instance_config=SimpleNamespace(data={}),
|
||||
logger=SimpleNamespace(warning=Mock()),
|
||||
)
|
||||
service = MaintenanceService(application)
|
||||
|
||||
async def cleanup_uploads(_context, _retention_days):
|
||||
assert persistence_mgr.active_workspace == TEST_CONTEXT.workspace_uuid
|
||||
assert persistence_mgr.current_session() is None
|
||||
return 2
|
||||
|
||||
def cleanup_logs(_retention_days):
|
||||
assert persistence_mgr.active_workspace == TEST_CONTEXT.workspace_uuid
|
||||
assert persistence_mgr.current_session() is None
|
||||
return 1
|
||||
|
||||
service._cleanup_expired_uploaded_files = cleanup_uploads
|
||||
service._cleanup_expired_log_files = cleanup_logs
|
||||
|
||||
assert await service.cleanup_expired_files(TEST_CONTEXT) == {
|
||||
'uploaded_files': 2,
|
||||
'log_files': 1,
|
||||
}
|
||||
assert persistence_mgr.active_workspace is None
|
||||
|
||||
|
||||
class TestMaintenanceServiceGetStorageAnalysis:
|
||||
"""Tests for get_storage_analysis method."""
|
||||
|
||||
@@ -13,6 +13,7 @@ Source: src/langbot/pkg/api/http/service/mcp.py
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, Mock, MagicMock
|
||||
@@ -485,6 +486,54 @@ class TestMCPServiceCreateMCPServer:
|
||||
# Verify - host_mcp_server was called
|
||||
ap.tool_mgr.mcp_tool_loader.host_mcp_server.assert_called_once()
|
||||
|
||||
async def test_create_mcp_server_does_not_start_host_until_transaction_commits(self):
|
||||
"""The Runtime must not observe a server row that can still roll back."""
|
||||
|
||||
gate = asyncio.get_running_loop().create_future()
|
||||
|
||||
class PersistenceManagerStub:
|
||||
def create_after_commit_gate(self):
|
||||
return gate
|
||||
|
||||
ap = SimpleNamespace()
|
||||
ap.persistence_mgr = PersistenceManagerStub()
|
||||
ap.instance_config = SimpleNamespace(data={'system': {'limitation': {'max_extensions': -1}}})
|
||||
observed = []
|
||||
|
||||
async def host_mcp_server(context, config):
|
||||
observed.append((context, config))
|
||||
|
||||
ap.tool_mgr = SimpleNamespace(
|
||||
mcp_tool_loader=SimpleNamespace(
|
||||
host_mcp_server=host_mcp_server,
|
||||
_hosted_mcp_tasks=[],
|
||||
)
|
||||
)
|
||||
server_entity = _create_mock_mcp_server(server_uuid='new-uuid', enable=True)
|
||||
results = [
|
||||
_create_mock_result([]),
|
||||
Mock(),
|
||||
_create_mock_result(first_item=server_entity),
|
||||
]
|
||||
ap.persistence_mgr.execute_async = AsyncMock(side_effect=results)
|
||||
ap.persistence_mgr.serialize_model = Mock(
|
||||
return_value={'uuid': 'new-uuid', 'name': 'New Server', 'enable': True}
|
||||
)
|
||||
service = _service(ap)
|
||||
|
||||
await service.create_mcp_server(_CONTEXT, {'name': 'New Server', 'enable': True})
|
||||
await asyncio.sleep(0)
|
||||
assert observed == []
|
||||
|
||||
gate.set_result(None)
|
||||
await ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks[0]
|
||||
assert observed == [
|
||||
(
|
||||
_CONTEXT,
|
||||
{'uuid': 'new-uuid', 'name': 'New Server', 'enable': True},
|
||||
)
|
||||
]
|
||||
|
||||
async def test_create_mcp_server_disabled_no_load(self):
|
||||
"""Does not load server when disabled."""
|
||||
# Setup
|
||||
|
||||
@@ -11,7 +11,9 @@ from langbot.pkg.api.http.authz import WorkspaceRequiredError
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.api.http.service.monitoring import MonitoringService
|
||||
from langbot.pkg.entity.persistence.base import Base
|
||||
from langbot.pkg.entity.persistence.monitoring import MonitoringMessage
|
||||
from langbot.pkg.entity.persistence.workspace import Workspace
|
||||
from langbot.pkg.persistence.mgr import PersistenceManager
|
||||
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
@@ -136,6 +138,26 @@ async def test_same_session_and_resource_ids_do_not_collide(service):
|
||||
assert (await service.get_message_details(context_a, message_b))['found'] is False
|
||||
|
||||
|
||||
async def test_tool_call_inherits_context_from_connection_message_row(service):
|
||||
context = _context(WORKSPACE_A)
|
||||
message_id = await _record_message(service, context, 'tool context')
|
||||
|
||||
await service.record_tool_call(
|
||||
context,
|
||||
tool_name='search',
|
||||
tool_source='native',
|
||||
duration=12,
|
||||
message_id=message_id,
|
||||
)
|
||||
|
||||
tool_calls, total = await service.get_tool_calls(context)
|
||||
assert total == 1
|
||||
assert tool_calls[0]['bot_id'] == 'same-bot'
|
||||
assert tool_calls[0]['pipeline_id'] == 'same-pipeline'
|
||||
assert tool_calls[0]['session_id'] == 'same-session'
|
||||
assert tool_calls[0]['message_id'] == message_id
|
||||
|
||||
|
||||
async def test_feedback_upsert_and_cancel_are_workspace_scoped(service):
|
||||
context_a = _context(WORKSPACE_A)
|
||||
context_b = _context(WORKSPACE_B)
|
||||
@@ -152,3 +174,58 @@ async def test_feedback_upsert_and_cancel_are_workspace_scoped(service):
|
||||
await service.record_feedback(context_a, feedback_id='same-feedback', feedback_type=3)
|
||||
assert (await service.get_feedback_stats(context_a))['total_feedback'] == 0
|
||||
assert (await service.get_feedback_stats(context_b))['total_feedback'] == 1
|
||||
|
||||
|
||||
async def test_cleanup_commits_sqlite_delete_before_vacuum(tmp_path):
|
||||
engine = create_async_engine(
|
||||
f'sqlite+aiosqlite:///{tmp_path / "monitoring-cleanup.db"}',
|
||||
connect_args={'timeout': 0.1},
|
||||
)
|
||||
application = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(data={'database': {'use': 'sqlite'}}),
|
||||
)
|
||||
manager = PersistenceManager(application)
|
||||
manager.db = SimpleNamespace(get_engine=lambda: engine)
|
||||
application.persistence_mgr = manager
|
||||
try:
|
||||
async with engine.begin() as connection:
|
||||
await connection.run_sync(Base.metadata.create_all)
|
||||
await connection.execute(
|
||||
sqlalchemy.insert(Workspace).values(
|
||||
uuid=WORKSPACE_A,
|
||||
instance_uuid='instance',
|
||||
name='A',
|
||||
slug='a',
|
||||
source='cloud_projection',
|
||||
)
|
||||
)
|
||||
await connection.execute(
|
||||
sqlalchemy.insert(MonitoringMessage).values(
|
||||
id='expired-message',
|
||||
workspace_uuid=WORKSPACE_A,
|
||||
timestamp=datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
|
||||
- datetime.timedelta(days=30),
|
||||
bot_id='bot',
|
||||
bot_name='Bot',
|
||||
pipeline_id='pipeline',
|
||||
pipeline_name='Pipeline',
|
||||
message_content='expired',
|
||||
session_id='session',
|
||||
status='success',
|
||||
level='info',
|
||||
)
|
||||
)
|
||||
|
||||
deleted = await MonitoringService(application).cleanup_expired_records(
|
||||
_context(WORKSPACE_A),
|
||||
retention_days=1,
|
||||
)
|
||||
|
||||
assert deleted['monitoring_messages'] == 1
|
||||
async with engine.connect() as connection:
|
||||
remaining = await connection.scalar(
|
||||
sqlalchemy.select(sqlalchemy.func.count()).select_from(MonitoringMessage)
|
||||
)
|
||||
assert remaining == 0
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
@@ -45,9 +45,16 @@ async def test_verify_api_key_keeps_db_validation_for_lbk_keys(key_exists):
|
||||
if key_exists
|
||||
else None
|
||||
)
|
||||
query_result = Mock()
|
||||
query_result.first.return_value = key
|
||||
persistence_mgr = SimpleNamespace(execute_async=AsyncMock(side_effect=[query_result, Mock(rowcount=1)]))
|
||||
discovery_result = Mock()
|
||||
discovery_result.first.return_value = key
|
||||
query_results = [discovery_result]
|
||||
if key_exists:
|
||||
scoped_result = Mock()
|
||||
scoped_result.first.return_value = key
|
||||
update_result = Mock()
|
||||
update_result.scalar_one_or_none.return_value = key.id
|
||||
query_results.extend([scoped_result, update_result])
|
||||
persistence_mgr = SimpleNamespace(execute_async=AsyncMock(side_effect=query_results))
|
||||
instance_config = SimpleNamespace(data={'api': {'global_api_key': ''}})
|
||||
workspace_service = SimpleNamespace(
|
||||
get_execution_binding=AsyncMock(
|
||||
@@ -69,7 +76,7 @@ async def test_verify_api_key_keeps_db_validation_for_lbk_keys(key_exists):
|
||||
result = await service.verify_api_key('lbk_valid_format')
|
||||
|
||||
assert result is key_exists
|
||||
assert persistence_mgr.execute_async.await_count == (2 if key_exists else 1)
|
||||
assert persistence_mgr.execute_async.await_count == (3 if key_exists else 1)
|
||||
if key_exists:
|
||||
workspace_service.get_execution_binding.assert_awaited_once_with('workspace-a')
|
||||
else:
|
||||
|
||||
@@ -0,0 +1,121 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from langbot.pkg.api.http.service.apikey import ApiKeyIdentity
|
||||
from langbot.pkg.api.mcp.context import get_request_context
|
||||
from langbot.pkg.api.mcp.mount import MCPMount
|
||||
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
|
||||
from langbot.pkg.persistence.tenant_uow import PersistenceScopeKind
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_mount_keeps_request_context_but_no_session_during_stream_wait(tmp_path) -> None:
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "mcp-short-scope.db"}')
|
||||
table = sa.Table('mcp_scope_probe', sa.MetaData(), sa.Column('id', sa.Integer, primary_key=True))
|
||||
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||
manager.db = SimpleNamespace(get_engine=lambda: engine)
|
||||
checked_out = 0
|
||||
|
||||
def on_checkout(*_args):
|
||||
nonlocal checked_out
|
||||
checked_out += 1
|
||||
|
||||
def on_checkin(*_args):
|
||||
nonlocal checked_out
|
||||
checked_out -= 1
|
||||
|
||||
sa.event.listen(engine.sync_engine, 'checkout', on_checkout)
|
||||
sa.event.listen(engine.sync_engine, 'checkin', on_checkin)
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(table.metadata.create_all)
|
||||
|
||||
identity = ApiKeyIdentity(
|
||||
instance_uuid='instance-1',
|
||||
workspace_uuid='workspace-1',
|
||||
placement_generation=7,
|
||||
api_key_uuid='key-1',
|
||||
permissions=frozenset({'pipelines:read'}),
|
||||
)
|
||||
app = SimpleNamespace(
|
||||
apikey_service=SimpleNamespace(authenticate_api_key=AsyncMock(return_value=identity)),
|
||||
persistence_mgr=manager,
|
||||
deployment_admission=None,
|
||||
deployment=None,
|
||||
)
|
||||
stream_waiting = asyncio.Event()
|
||||
release_stream = asyncio.Event()
|
||||
observations: list[tuple[str, str, bool]] = []
|
||||
|
||||
async def fake_mcp_asgi(scope, receive, send):
|
||||
del scope, receive
|
||||
context = get_request_context()
|
||||
assert manager.current_scope().kind is PersistenceScopeKind.WORKSPACE
|
||||
assert manager.current_session() is None
|
||||
await manager.execute_async(sa.select(table.c.id))
|
||||
assert manager.current_session() is None
|
||||
observations.append((context.request_id, context.workspace_uuid, manager.current_session() is None))
|
||||
stream_waiting.set()
|
||||
await release_stream.wait()
|
||||
preserved_context = get_request_context()
|
||||
observations.append(
|
||||
(
|
||||
preserved_context.request_id,
|
||||
preserved_context.workspace_uuid,
|
||||
manager.current_session() is None,
|
||||
)
|
||||
)
|
||||
await manager.execute_async(sa.select(table.c.id))
|
||||
assert manager.current_session() is None
|
||||
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):
|
||||
del scope, receive, send
|
||||
raise AssertionError('MCP request was routed to Quart')
|
||||
|
||||
mount = MCPMount.__new__(MCPMount)
|
||||
mount.ap = app
|
||||
mount._mcp_asgi = fake_mcp_asgi
|
||||
sent_messages: list[dict] = []
|
||||
|
||||
async def receive():
|
||||
return {'type': 'http.request', 'body': b'', 'more_body': False}
|
||||
|
||||
async def send(message):
|
||||
sent_messages.append(message)
|
||||
|
||||
async def release_after_observation() -> None:
|
||||
await asyncio.wait_for(stream_waiting.wait(), timeout=2)
|
||||
assert checked_out == 0
|
||||
release_stream.set()
|
||||
|
||||
release_task = asyncio.create_task(release_after_observation())
|
||||
await mount.wrap(unused_quart_asgi)(
|
||||
{
|
||||
'type': 'http',
|
||||
'path': '/mcp',
|
||||
'headers': [(b'x-api-key', b'secret')],
|
||||
},
|
||||
receive,
|
||||
send,
|
||||
)
|
||||
await release_task
|
||||
|
||||
assert sent_messages[0]['status'] == 200
|
||||
assert len(observations) == 2
|
||||
assert observations[0] == observations[1]
|
||||
assert observations[0][1:] == ('workspace-1', True)
|
||||
assert checked_out == 0
|
||||
assert manager.current_scope() is None
|
||||
with pytest.raises(RuntimeError, match='context is unavailable'):
|
||||
get_request_context()
|
||||
finally:
|
||||
await engine.dispose()
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
@@ -106,3 +107,33 @@ async def test_background_plugin_operation_refences_captured_generation(plugin_r
|
||||
await router._run_fenced_plugin_operation(CONTEXT, operation)
|
||||
|
||||
operation.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_background_plugin_operation_revalidates_inside_short_tenant_uow(plugin_router_cls):
|
||||
scopes = []
|
||||
|
||||
@asynccontextmanager
|
||||
async def tenant_uow(workspace_uuid):
|
||||
scopes.append(workspace_uuid)
|
||||
yield
|
||||
|
||||
connector = SimpleNamespace(
|
||||
require_workspace_context=AsyncMock(side_effect=lambda context: context),
|
||||
)
|
||||
operation = AsyncMock(return_value='done')
|
||||
router = object.__new__(plugin_router_cls)
|
||||
router.ap = SimpleNamespace(
|
||||
plugin_connector=connector,
|
||||
persistence_mgr=SimpleNamespace(
|
||||
mode=SimpleNamespace(value='cloud_runtime'),
|
||||
tenant_uow=tenant_uow,
|
||||
),
|
||||
)
|
||||
|
||||
result = await router._run_fenced_plugin_operation(CONTEXT, operation)
|
||||
|
||||
assert result == 'done'
|
||||
assert scopes == [CONTEXT.workspace_uuid]
|
||||
connector.require_workspace_context.assert_awaited_once_with(CONTEXT)
|
||||
operation.assert_awaited_once()
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.api.http.context import (
|
||||
PrincipalContext,
|
||||
PrincipalType,
|
||||
RequestContext,
|
||||
WorkspaceContext,
|
||||
)
|
||||
from langbot.pkg.api.http.controller.groups.pipelines.websocket_chat import WebSocketChatRouterGroup
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_pipeline_lookup_opens_workspace_uow_after_auth_scope_closed() -> None:
|
||||
workspace_uuid = 'workspace-a'
|
||||
scopes: list[str] = []
|
||||
in_scope = False
|
||||
|
||||
@asynccontextmanager
|
||||
async def tenant_uow(selected_workspace_uuid: str):
|
||||
nonlocal in_scope
|
||||
assert not in_scope
|
||||
in_scope = True
|
||||
scopes.append(selected_workspace_uuid)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
in_scope = False
|
||||
|
||||
async def get_pipeline(_context, _pipeline_uuid):
|
||||
assert in_scope
|
||||
return {'uuid': 'pipeline-a'}
|
||||
|
||||
adapter = Mock()
|
||||
router = object.__new__(WebSocketChatRouterGroup)
|
||||
router.ap = SimpleNamespace(
|
||||
persistence_mgr=SimpleNamespace(
|
||||
mode=SimpleNamespace(value='cloud_runtime'),
|
||||
tenant_uow=tenant_uow,
|
||||
),
|
||||
pipeline_service=SimpleNamespace(get_pipeline=AsyncMock(side_effect=get_pipeline)),
|
||||
platform_mgr=SimpleNamespace(get_websocket_proxy_bot=AsyncMock(return_value=SimpleNamespace(adapter=adapter))),
|
||||
)
|
||||
request_context = RequestContext(
|
||||
instance_uuid='instance-a',
|
||||
placement_generation=1,
|
||||
request_id='request-a',
|
||||
auth_type='user_token',
|
||||
principal=PrincipalContext(
|
||||
principal_type=PrincipalType.ACCOUNT,
|
||||
account_uuid='account-a',
|
||||
),
|
||||
workspace=WorkspaceContext(
|
||||
workspace_uuid=workspace_uuid,
|
||||
membership_uuid='membership-a',
|
||||
role='owner',
|
||||
permissions=frozenset(),
|
||||
),
|
||||
)
|
||||
|
||||
result = await router._get_scoped_adapter(request_context, 'pipeline-a')
|
||||
|
||||
assert result is adapter
|
||||
assert scopes == [workspace_uuid]
|
||||
Reference in New Issue
Block a user