mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-22 20:36:08 +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]
|
||||
@@ -0,0 +1,230 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime as dt
|
||||
import hashlib
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot_plugin.box.models import SandboxAdmissionPolicy
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.box.admission import (
|
||||
BoxAdmissionError,
|
||||
SandboxAdmissionController,
|
||||
require_cloud_admission_policy,
|
||||
)
|
||||
from langbot.pkg.box.service import BoxService
|
||||
from langbot.pkg.cloud.entitlements import (
|
||||
EntitlementResolver,
|
||||
EntitlementSnapshot,
|
||||
EntitlementUnavailableError,
|
||||
)
|
||||
|
||||
|
||||
_UTC = dt.timezone.utc
|
||||
_CONTEXT = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=3,
|
||||
entitlement_revision=7,
|
||||
)
|
||||
|
||||
|
||||
def _snapshot(
|
||||
*,
|
||||
revision: int = 7,
|
||||
managed: bool = True,
|
||||
sessions: int = 1,
|
||||
expires_at: int = 2_000,
|
||||
) -> EntitlementSnapshot:
|
||||
return EntitlementSnapshot(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
entitlement_revision=revision,
|
||||
status='active',
|
||||
not_before=1,
|
||||
expires_at=expires_at,
|
||||
features={'managed_sandbox': managed},
|
||||
limits={'managed_sandbox_sessions': sessions},
|
||||
)
|
||||
|
||||
|
||||
def _controller(snapshot: EntitlementSnapshot, *, now: float = 1_000.25):
|
||||
provider = SimpleNamespace(get_workspace_entitlement=AsyncMock(return_value=snapshot))
|
||||
resolver = EntitlementResolver('instance-a', provider)
|
||||
client = SimpleNamespace(
|
||||
upsert_sandbox_admission_grant=AsyncMock(
|
||||
side_effect=lambda grant: {
|
||||
'installed': True,
|
||||
'workspace_uuid': grant.workspace_uuid,
|
||||
'execution_generation': grant.execution_generation,
|
||||
'entitlement_revision': grant.entitlement_revision,
|
||||
'max_sessions': grant.max_sessions,
|
||||
'max_managed_processes': grant.max_managed_processes,
|
||||
}
|
||||
),
|
||||
revoke_sandbox_admission_grant=AsyncMock(
|
||||
side_effect=lambda revocation: {
|
||||
'revoked': True,
|
||||
'workspace_uuid': revocation.workspace_uuid,
|
||||
'entitlement_revision': revocation.entitlement_revision,
|
||||
}
|
||||
),
|
||||
)
|
||||
app = SimpleNamespace(entitlement_resolver=resolver, logger=Mock())
|
||||
controller = SandboxAdmissionController(
|
||||
app,
|
||||
client,
|
||||
policy=SandboxAdmissionPolicy(required=True, max_grant_ttl_sec=300),
|
||||
wall_time=lambda: now,
|
||||
)
|
||||
return controller, client, provider
|
||||
|
||||
|
||||
def test_cloud_admission_policy_requires_positive_workspace_quota():
|
||||
with pytest.raises(BoxAdmissionError, match='workspace quota must be a positive integer'):
|
||||
require_cloud_admission_policy(
|
||||
{
|
||||
'required': True,
|
||||
'workspace_quota_mb': 0,
|
||||
}
|
||||
)
|
||||
|
||||
policy = require_cloud_admission_policy(
|
||||
{
|
||||
'required': True,
|
||||
'workspace_quota_mb': 32,
|
||||
}
|
||||
)
|
||||
assert policy.workspace_quota_mb == 32
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_active_generic_entitlement_installs_short_lived_numeric_grant():
|
||||
controller, client, provider = _controller(_snapshot())
|
||||
|
||||
grant = await controller.require(_CONTEXT)
|
||||
|
||||
assert grant.instance_uuid == _CONTEXT.instance_uuid
|
||||
assert grant.workspace_uuid == _CONTEXT.workspace_uuid
|
||||
assert grant.execution_generation == _CONTEXT.placement_generation
|
||||
assert grant.entitlement_revision == 7
|
||||
assert grant.max_sessions == 1
|
||||
assert grant.max_managed_processes == 0
|
||||
assert grant.expires_at == dt.datetime.fromtimestamp(1_300, tz=_UTC)
|
||||
assert (grant.expires_at - dt.datetime.fromtimestamp(1_000.25, tz=_UTC)).total_seconds() < 300
|
||||
provider.get_workspace_entitlement.assert_awaited_once_with('workspace-a')
|
||||
client.upsert_sandbox_admission_grant.assert_awaited_once_with(grant)
|
||||
client.revoke_sandbox_admission_grant.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
'snapshot',
|
||||
[
|
||||
_snapshot(managed=False),
|
||||
_snapshot(sessions=0),
|
||||
_snapshot(sessions=2),
|
||||
],
|
||||
)
|
||||
async def test_non_eligible_entitlement_revokes_and_fails_closed(snapshot):
|
||||
controller, client, _provider = _controller(snapshot)
|
||||
|
||||
with pytest.raises(EntitlementUnavailableError):
|
||||
await controller.require(_CONTEXT)
|
||||
|
||||
client.upsert_sandbox_admission_grant.assert_not_awaited()
|
||||
revocation = client.revoke_sandbox_admission_grant.await_args.args[0]
|
||||
assert revocation.entitlement_revision == snapshot.entitlement_revision
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_transient_entitlement_failure_does_not_tombstone_valid_revision():
|
||||
controller, client, provider = _controller(_snapshot())
|
||||
await controller.require(_CONTEXT)
|
||||
provider.get_workspace_entitlement.side_effect = RuntimeError('control plane unavailable')
|
||||
|
||||
with pytest.raises(RuntimeError, match='control plane unavailable'):
|
||||
await controller.require(_CONTEXT)
|
||||
|
||||
client.revoke_sandbox_admission_grant.assert_not_awaited()
|
||||
|
||||
provider.get_workspace_entitlement.side_effect = None
|
||||
provider.get_workspace_entitlement.return_value = _snapshot()
|
||||
recovered = await controller.require(_CONTEXT)
|
||||
assert recovered.entitlement_revision == 7
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_receipt_mismatch_is_revoked_and_never_admitted():
|
||||
controller, client, _provider = _controller(_snapshot())
|
||||
client.upsert_sandbox_admission_grant.return_value = {'installed': True, 'workspace_uuid': 'other'}
|
||||
client.upsert_sandbox_admission_grant.side_effect = None
|
||||
|
||||
with pytest.raises(Exception, match='invalid sandbox admission receipt'):
|
||||
await controller.require(_CONTEXT)
|
||||
|
||||
client.revoke_sandbox_admission_grant.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authoritative_cancelled_revision_is_revoked():
|
||||
cancelled = _snapshot(revision=8).model_copy(update={'status': 'cancelled'})
|
||||
controller, client, _provider = _controller(cancelled)
|
||||
|
||||
with pytest.raises(EntitlementUnavailableError, match='not active'):
|
||||
await controller.require(_CONTEXT)
|
||||
|
||||
revocation = client.revoke_sandbox_admission_grant.await_args.args[0]
|
||||
assert revocation.entitlement_revision == 8
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cloud_box_readiness_failure_aborts_service_initialization(tmp_path):
|
||||
workspace_root = tmp_path / 'box' / 'workspaces'
|
||||
workspace_root.mkdir(parents=True)
|
||||
box_config = {
|
||||
'enabled': True,
|
||||
'backend': 'nsjail',
|
||||
'runtime': {'endpoint': 'ws://box:5410'},
|
||||
'local': {
|
||||
'host_root': str(tmp_path / 'box'),
|
||||
'default_workspace': str(workspace_root),
|
||||
'allowed_mount_roots': [str(tmp_path / 'box')],
|
||||
},
|
||||
'admission': {
|
||||
'required': True,
|
||||
'logical_session_id': 'global',
|
||||
'required_backend': 'nsjail',
|
||||
'max_sessions': 1,
|
||||
'max_managed_processes': 0,
|
||||
'max_grant_ttl_sec': 300,
|
||||
'workspace_quota_mb': 32,
|
||||
},
|
||||
}
|
||||
client = SimpleNamespace(
|
||||
initialize=AsyncMock(),
|
||||
verify_shared_workspace=AsyncMock(
|
||||
side_effect=lambda marker_name: {
|
||||
'marker_name': marker_name,
|
||||
'size': (workspace_root / marker_name).stat().st_size,
|
||||
'sha256': hashlib.sha256((workspace_root / marker_name).read_bytes()).hexdigest(),
|
||||
}
|
||||
),
|
||||
get_backend_info=AsyncMock(return_value={'name': 'docker', 'available': True}),
|
||||
)
|
||||
app = SimpleNamespace(
|
||||
logger=Mock(),
|
||||
deployment=SimpleNamespace(multi_workspace_enabled=True),
|
||||
entitlement_resolver=Mock(),
|
||||
workspace_service=SimpleNamespace(instance_uuid='instance-a'),
|
||||
instance_config=SimpleNamespace(data={'box': box_config}),
|
||||
)
|
||||
service = BoxService(app, client=client)
|
||||
|
||||
with pytest.raises(Exception, match='nsjail isolation readiness failed'):
|
||||
await service.initialize()
|
||||
|
||||
assert service.available is False
|
||||
@@ -117,6 +117,9 @@ class _InProcessBoxRuntimeClient(BoxRuntimeClient):
|
||||
async def init(self, config: dict) -> None:
|
||||
self._runtime.init(config)
|
||||
|
||||
async def verify_shared_workspace(self, marker_name: str) -> dict:
|
||||
return self._runtime.verify_shared_workspace(marker_name)
|
||||
|
||||
|
||||
class FakeBackend(BaseSandboxBackend):
|
||||
def __init__(self, logger: Mock, available: bool = True):
|
||||
@@ -322,6 +325,43 @@ def test_separated_box_runtime_does_not_create_default_workspace_in_langbot(tmp_
|
||||
assert not (host_root / 'default').exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cloud_initialize_fails_when_core_and_runtime_volumes_are_separated(tmp_path):
|
||||
logger = Mock()
|
||||
core_root = tmp_path / 'core-box'
|
||||
runtime_root = tmp_path / 'runtime-box'
|
||||
(core_root / 'default').mkdir(parents=True)
|
||||
runtime = BoxRuntime(logger=logger, backends=[FakeBackend(logger)], session_ttl_sec=300)
|
||||
runtime.init(
|
||||
{
|
||||
'local': {
|
||||
'host_root': str(runtime_root),
|
||||
'default_workspace': 'default',
|
||||
'allowed_mount_roots': [str(runtime_root)],
|
||||
}
|
||||
}
|
||||
)
|
||||
app = make_app(logger, host_root=str(core_root))
|
||||
app.deployment = SimpleNamespace(multi_workspace_enabled=True)
|
||||
app.instance_config.data['box'].update(
|
||||
{
|
||||
'backend': 'nsjail',
|
||||
'admission': {'required': True, 'workspace_quota_mb': 32},
|
||||
}
|
||||
)
|
||||
service = BoxService(
|
||||
app,
|
||||
client=_InProcessBoxRuntimeClient(logger, runtime),
|
||||
)
|
||||
|
||||
with pytest.raises(BoxValidationError, match='shared durable Workspace volume'):
|
||||
await service.initialize()
|
||||
|
||||
assert service.available is False
|
||||
assert list((core_root / 'default').glob('.langbot-box-volume-probe-*')) == []
|
||||
await runtime.shutdown()
|
||||
|
||||
|
||||
def test_separated_box_runtime_allows_box_owned_missing_host_path(tmp_path):
|
||||
logger = Mock()
|
||||
runtime = BoxRuntime(logger=logger, backends=[FakeBackend(logger)], session_ttl_sec=300)
|
||||
@@ -1888,7 +1928,7 @@ class TestInboundOutboundRoundTrip:
|
||||
assert '/workspace/inbox/' in parameters['command']
|
||||
return {
|
||||
'ok': True,
|
||||
'stdout': '["/workspace/inbox/42/image_1.png"]',
|
||||
'stdout': '["/workspace/inbox/query-42/image_1.png"]',
|
||||
'stderr': '',
|
||||
}
|
||||
|
||||
@@ -1898,7 +1938,7 @@ class TestInboundOutboundRoundTrip:
|
||||
assert len(descriptors) == 1
|
||||
d = descriptors[0]
|
||||
assert d['type'] == 'Image'
|
||||
assert d['path'] == '/workspace/inbox/42/image_1.png'
|
||||
assert d['path'] == '/workspace/inbox/query-42/image_1.png'
|
||||
assert d['size'] == len(img_bytes)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -2011,7 +2051,7 @@ class TestAttachmentHostPath:
|
||||
assert d['type'] == 'Image'
|
||||
assert d['size'] == len(big)
|
||||
# File actually landed on the host workspace.
|
||||
host_file = os.path.join(ws, 'inbox', str(query.query_id), d['name'])
|
||||
host_file = os.path.join(ws, 'inbox', str(query.query_uuid), d['name'])
|
||||
assert os.path.isfile(host_file)
|
||||
assert open(host_file, 'rb').read() == big
|
||||
|
||||
@@ -2023,7 +2063,7 @@ class TestAttachmentHostPath:
|
||||
|
||||
service, ws = self._service_with_workspace(tmp_path)
|
||||
# Seed a stale file under the same query_id (simulates webchat id reuse).
|
||||
stale_dir = os.path.join(ws, 'inbox', '42')
|
||||
stale_dir = os.path.join(ws, 'inbox', 'query-42')
|
||||
os.makedirs(stale_dir, exist_ok=True)
|
||||
open(os.path.join(stale_dir, 'image_1.png'), 'wb').write(b'STALE-OLD-IMAGE')
|
||||
|
||||
@@ -2039,11 +2079,40 @@ class TestAttachmentHostPath:
|
||||
# No leftover content from the stale image.
|
||||
assert b'STALE-OLD-IMAGE' not in open(host_file, 'rb').read()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inbound_host_replaces_query_symlink_without_touching_other_workspace(self, tmp_path):
|
||||
import base64
|
||||
|
||||
import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
|
||||
service, ws = self._service_with_workspace(tmp_path)
|
||||
other_workspace = tmp_path / 'other-workspace'
|
||||
other_workspace.mkdir()
|
||||
protected = other_workspace / 'protected.txt'
|
||||
protected.write_bytes(b'workspace-b-secret')
|
||||
inbox = os.path.join(ws, 'inbox')
|
||||
os.makedirs(inbox, exist_ok=True)
|
||||
os.symlink(other_workspace, os.path.join(inbox, 'query-42'))
|
||||
|
||||
query = make_query()
|
||||
payload = b'workspace-a-input'
|
||||
query.message_chain = platform_message.MessageChain(
|
||||
[platform_message.File(name='input.bin', base64=base64.b64encode(payload).decode())]
|
||||
)
|
||||
service.execute_tool = AsyncMock(side_effect=AssertionError('exec must not be used on host path'))
|
||||
|
||||
descriptors = await service.materialize_inbound_attachments(query)
|
||||
|
||||
assert descriptors[0]['path'] == '/workspace/inbox/query-42/input.bin'
|
||||
assert protected.read_bytes() == b'workspace-b-secret'
|
||||
assert not os.path.islink(os.path.join(inbox, 'query-42'))
|
||||
assert open(os.path.join(inbox, 'query-42', 'input.bin'), 'rb').read() == payload
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_outbound_reads_host_and_clears(self, tmp_path):
|
||||
service, ws = self._service_with_workspace(tmp_path)
|
||||
query = make_query()
|
||||
outbox = os.path.join(ws, 'outbox', str(query.query_id))
|
||||
outbox = os.path.join(ws, 'outbox', str(query.query_uuid))
|
||||
os.makedirs(outbox, exist_ok=True)
|
||||
# A large file that would be truncated on the exec/stdout path:
|
||||
big_png = b'\x89PNG\r\n\x1a\n' + b'y' * (400 * 1024)
|
||||
@@ -2063,13 +2132,69 @@ class TestAttachmentHostPath:
|
||||
# Outbox cleared after collection.
|
||||
assert os.listdir(outbox) == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_outbound_host_never_follows_query_or_file_symlinks(self, tmp_path):
|
||||
service, ws = self._service_with_workspace(tmp_path)
|
||||
query = make_query()
|
||||
other_workspace = tmp_path / 'other-workspace'
|
||||
other_workspace.mkdir()
|
||||
secret = other_workspace / 'secret.txt'
|
||||
secret.write_bytes(b'workspace-b-secret')
|
||||
outbox_root = os.path.join(ws, 'outbox')
|
||||
os.makedirs(outbox_root, exist_ok=True)
|
||||
|
||||
# A hostile query-directory replacement is rejected rather than read.
|
||||
query_dir = os.path.join(outbox_root, str(query.query_uuid))
|
||||
os.symlink(other_workspace, query_dir)
|
||||
with pytest.raises(BoxValidationError, match='symbolic link'):
|
||||
await service.collect_outbound_attachments(query)
|
||||
assert secret.read_bytes() == b'workspace-b-secret'
|
||||
|
||||
os.unlink(query_dir)
|
||||
os.makedirs(query_dir)
|
||||
os.symlink(secret, os.path.join(query_dir, 'leak.txt'))
|
||||
service.execute_tool = AsyncMock(side_effect=AssertionError('exec must not be used on host path'))
|
||||
assert await service.collect_outbound_attachments(query) == []
|
||||
assert secret.read_bytes() == b'workspace-b-secret'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_outbound_host_fails_closed_on_inode_bomb(self, tmp_path):
|
||||
service, ws = self._service_with_workspace(tmp_path)
|
||||
query = make_query()
|
||||
outbox = os.path.join(ws, 'outbox', str(query.query_uuid))
|
||||
os.makedirs(outbox, exist_ok=True)
|
||||
harmless_target = tmp_path / 'harmless-target'
|
||||
harmless_target.write_bytes(b'x')
|
||||
# Symlinks do not count toward the 20 returned files, so this proves
|
||||
# traversal itself has a bounded entry budget.
|
||||
for index in range(513):
|
||||
os.symlink(harmless_target, os.path.join(outbox, f'entry-{index}'))
|
||||
|
||||
with pytest.raises(BoxValidationError, match='symbolic link'):
|
||||
await service.collect_outbound_attachments(query)
|
||||
assert harmless_target.read_bytes() == b'x'
|
||||
|
||||
def test_host_attachment_directories_use_query_uuid_not_process_local_id(self, tmp_path):
|
||||
service, _ws = self._service_with_workspace(tmp_path)
|
||||
first = make_query(query_id=7)
|
||||
second = make_query(query_id=7)
|
||||
object.__setattr__(first, 'query_uuid', 'replica-a-query')
|
||||
object.__setattr__(second, 'query_uuid', 'replica-b-query')
|
||||
|
||||
first_path = service._host_query_dir(service.OUTBOX_SUBDIR, first)
|
||||
second_path = service._host_query_dir(service.OUTBOX_SUBDIR, second)
|
||||
|
||||
assert first_path is not None and first_path.endswith('/outbox/replica-a-query')
|
||||
assert second_path is not None and second_path.endswith('/outbox/replica-b-query')
|
||||
assert first_path != second_path
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_outbound_empty_clears_stale_host_dir(self, tmp_path):
|
||||
# Reusing a query_id (counter resets on restart) must not re-send files
|
||||
# a previous run left in the outbox: an empty collection still clears it.
|
||||
service, ws = self._service_with_workspace(tmp_path)
|
||||
query = make_query()
|
||||
outbox = os.path.join(ws, 'outbox', str(query.query_id))
|
||||
outbox = os.path.join(ws, 'outbox', str(query.query_uuid))
|
||||
os.makedirs(outbox, exist_ok=True)
|
||||
# Stale file from a prior turn; the agent produced nothing this turn —
|
||||
# but _read_outbox_host would still pick it up, so collection must drop
|
||||
@@ -2118,16 +2243,16 @@ class TestAttachmentHostPath:
|
||||
os.makedirs(os.path.join(outbox, '0'), exist_ok=True)
|
||||
|
||||
# Simulate a host delete that cannot remove the root-owned outbox.
|
||||
import shutil as _shutil
|
||||
from langbot.pkg.box import secure_fs
|
||||
|
||||
real_rmtree = _shutil.rmtree
|
||||
real_purge = secure_fs.purge_subdirectory
|
||||
|
||||
def fake_rmtree(path, *a, **k):
|
||||
if os.path.abspath(path) == os.path.abspath(outbox):
|
||||
return # "permission denied" — silently leaves the dir
|
||||
return real_rmtree(path, *a, **k)
|
||||
def fake_purge(root, subdir):
|
||||
if os.path.abspath(os.path.join(root, subdir)) == os.path.abspath(outbox):
|
||||
raise PermissionError('root-owned')
|
||||
return real_purge(root, subdir)
|
||||
|
||||
monkeypatch.setattr(_shutil, 'rmtree', fake_rmtree)
|
||||
monkeypatch.setattr(secure_fs, 'purge_subdirectory', fake_purge)
|
||||
|
||||
service.build_spec = Mock()
|
||||
service.client.execute = AsyncMock()
|
||||
|
||||
@@ -1,11 +1,14 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.cloud.bootstrap import (
|
||||
CloudBootstrapError,
|
||||
CloudRuntimeUnavailableError,
|
||||
DeploymentAdmissionGuard,
|
||||
OpenSourceDeployment,
|
||||
VerifiedCloudDeployment,
|
||||
resolve_deployment,
|
||||
@@ -62,8 +65,34 @@ class _EntryPoints(list):
|
||||
def _cloud_config() -> dict:
|
||||
return {
|
||||
'database': {'use': 'postgresql'},
|
||||
'vdb': {'use': 'pgvector'},
|
||||
'vdb': {
|
||||
'use': 'pgvector',
|
||||
'pgvector': {
|
||||
'use_business_database': True,
|
||||
'allowed_dimensions': [384, 768, 1536],
|
||||
},
|
||||
},
|
||||
'mcp': {'stdio': {'enabled': False}},
|
||||
'plugin': {'worker': {'require_hard_limits': True}},
|
||||
'box': {
|
||||
'enabled': True,
|
||||
'backend': 'nsjail',
|
||||
'runtime': {'endpoint': 'ws://langbot-box:5410'},
|
||||
'admission': {
|
||||
'required': True,
|
||||
'logical_session_id': 'global',
|
||||
'required_backend': 'nsjail',
|
||||
'max_sessions': 1,
|
||||
'max_managed_processes': 0,
|
||||
'max_grant_ttl_sec': 300,
|
||||
'workspace_quota_mb': 32,
|
||||
},
|
||||
'local': {
|
||||
'host_root': '/var/lib/langbot/box',
|
||||
'default_workspace': '/var/lib/langbot/box/workspaces',
|
||||
'allowed_mount_roots': ['/var/lib/langbot/box'],
|
||||
},
|
||||
},
|
||||
# Proves mutable product metadata does not participate in selection.
|
||||
'system': {'edition': 'community'},
|
||||
}
|
||||
@@ -99,6 +128,7 @@ async def test_verified_closed_entry_point_activates_cloud_policy():
|
||||
('database', {'use': 'sqlite'}, 'database.use=postgresql'),
|
||||
('vdb', {'use': 'chroma'}, 'vdb.use=pgvector'),
|
||||
('mcp', {'stdio': {'enabled': True}}, 'mcp.stdio.enabled=false'),
|
||||
('plugin', {'worker': {'require_hard_limits': False}}, 'plugin.worker.require_hard_limits=true'),
|
||||
],
|
||||
)
|
||||
async def test_cloud_runtime_config_is_fail_closed(field, value, message):
|
||||
@@ -114,6 +144,79 @@ async def test_cloud_runtime_config_is_fail_closed(field, value, message):
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('pgvector_config', 'message'),
|
||||
[
|
||||
({'use_business_database': False, 'allowed_dimensions': [1536]}, 'use_business_database=true'),
|
||||
({'use_business_database': True, 'allowed_dimensions': []}, 'allowed_dimensions'),
|
||||
({'use_business_database': True, 'allowed_dimensions': [3072]}, 'allowed_dimensions'),
|
||||
({'use_business_database': True, 'allowed_dimensions': [True]}, 'allowed_dimensions'),
|
||||
],
|
||||
)
|
||||
async def test_cloud_pgvector_contract_is_fail_closed(pgvector_config, message):
|
||||
config = _cloud_config()
|
||||
config['vdb']['pgvector'] = pgvector_config
|
||||
|
||||
with pytest.raises(CloudBootstrapError, match=message):
|
||||
await resolve_deployment(
|
||||
instance_uuid='instance-a',
|
||||
instance_config=config,
|
||||
entry_points=lambda: _EntryPoints([_EntryPoint(_Provider())]),
|
||||
now=1_000,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('mutate', 'message'),
|
||||
[
|
||||
(lambda config: config['box'].update(enabled=False), 'box.enabled=true'),
|
||||
(lambda config: config['box'].update(backend='docker'), 'box.backend=nsjail'),
|
||||
(lambda config: config['box']['runtime'].update(endpoint=''), 'box.runtime.endpoint'),
|
||||
(
|
||||
lambda config: config['box']['admission'].update(max_sessions=2),
|
||||
'grant-enforced Box admission',
|
||||
),
|
||||
(
|
||||
lambda config: config['box']['admission'].update(max_managed_processes=1),
|
||||
'zero managed processes',
|
||||
),
|
||||
(
|
||||
lambda config: config['box']['admission'].update(max_grant_ttl_sec=301),
|
||||
'max_grant_ttl_sec',
|
||||
),
|
||||
(
|
||||
lambda config: config['box']['admission'].update(workspace_quota_mb=0),
|
||||
'workspace_quota_mb must be a positive integer',
|
||||
),
|
||||
(
|
||||
lambda config: config['box']['admission'].update(workspace_quota_mb=True),
|
||||
'workspace_quota_mb must be a positive integer',
|
||||
),
|
||||
(
|
||||
lambda config: config['box']['local'].update(default_workspace='relative/workspaces'),
|
||||
'default_workspace must be an absolute',
|
||||
),
|
||||
(
|
||||
lambda config: config['box']['local'].update(
|
||||
default_workspace='/other/workspaces',
|
||||
),
|
||||
'under allowed_mount_roots',
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_cloud_box_contract_is_fail_closed(mutate, message):
|
||||
config = _cloud_config()
|
||||
mutate(config)
|
||||
|
||||
with pytest.raises(CloudBootstrapError, match=message):
|
||||
await resolve_deployment(
|
||||
instance_uuid='instance-a',
|
||||
instance_config=config,
|
||||
entry_points=lambda: _EntryPoints([_EntryPoint(_Provider())]),
|
||||
now=1_000,
|
||||
)
|
||||
|
||||
|
||||
async def test_invalid_provider_never_falls_back_to_oss():
|
||||
provider = SimpleNamespace(bootstrap=lambda **_: object())
|
||||
|
||||
@@ -134,3 +237,55 @@ async def test_duplicate_closed_providers_fail_closed():
|
||||
entry_points=lambda: _EntryPoints([_EntryPoint(_Provider()), _EntryPoint(_Provider())]),
|
||||
now=1_000,
|
||||
)
|
||||
|
||||
|
||||
async def test_deployment_admission_expires_even_after_wall_clock_rollback():
|
||||
wall = [1_000.0]
|
||||
monotonic = [50.0]
|
||||
deployment = dataclasses.replace(
|
||||
_Provider().bootstrap(instance_uuid='instance-a', instance_config={}),
|
||||
expires_at=1_010,
|
||||
)
|
||||
guard = DeploymentAdmissionGuard(
|
||||
'instance-a',
|
||||
deployment,
|
||||
wall_time=lambda: wall[0],
|
||||
monotonic_time=lambda: monotonic[0],
|
||||
)
|
||||
|
||||
assert guard.require_active() is deployment
|
||||
wall[0] = 900.0
|
||||
monotonic[0] = 60.0
|
||||
with pytest.raises(CloudRuntimeUnavailableError, match='expired'):
|
||||
guard.require_active()
|
||||
|
||||
|
||||
async def test_deployment_admission_accepts_only_monotonic_non_conflicting_renewal():
|
||||
wall = [1_000.0]
|
||||
monotonic = [50.0]
|
||||
current = dataclasses.replace(
|
||||
_Provider().bootstrap(instance_uuid='instance-a', instance_config={}),
|
||||
expires_at=1_010,
|
||||
)
|
||||
guard = DeploymentAdmissionGuard(
|
||||
'instance-a',
|
||||
current,
|
||||
wall_time=lambda: wall[0],
|
||||
monotonic_time=lambda: monotonic[0],
|
||||
)
|
||||
renewed = dataclasses.replace(
|
||||
current,
|
||||
manifest_jti='manifest-b',
|
||||
manifest_generation=4,
|
||||
expires_at=2_000,
|
||||
)
|
||||
guard.replace(renewed)
|
||||
assert guard.require_active() is renewed
|
||||
|
||||
rollback = dataclasses.replace(current, manifest_generation=2)
|
||||
with pytest.raises(CloudRuntimeUnavailableError, match='rolled back'):
|
||||
guard.replace(rollback)
|
||||
|
||||
conflicting = dataclasses.replace(renewed, manifest_jti='different')
|
||||
with pytest.raises(CloudRuntimeUnavailableError, match='conflicting'):
|
||||
guard.replace(conflicting)
|
||||
|
||||
@@ -84,3 +84,26 @@ async def test_resolver_rejects_same_revision_with_different_contents():
|
||||
await resolver.resolve('workspace-a', now=150)
|
||||
with pytest.raises(EntitlementUnavailableError, match='conflicting contents'):
|
||||
await resolver.resolve('workspace-a', now=150)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolver_checks_deployment_admission_before_and_after_provider_call():
|
||||
checks = 0
|
||||
|
||||
def require_admission() -> None:
|
||||
nonlocal checks
|
||||
checks += 1
|
||||
if checks == 2:
|
||||
raise RuntimeError('manifest expired during provider call')
|
||||
|
||||
provider = AsyncMock()
|
||||
provider.get_workspace_entitlement = AsyncMock(return_value=_snapshot())
|
||||
resolver = EntitlementResolver(
|
||||
'instance-a',
|
||||
provider,
|
||||
deployment_admission=require_admission,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match='expired during provider call'):
|
||||
await resolver.resolve('workspace-a', now=150)
|
||||
assert checks == 2
|
||||
|
||||
@@ -35,6 +35,22 @@ class TestApplyEnvOverridesToConfig:
|
||||
|
||||
assert result['system']['name'] == 'custom_name'
|
||||
|
||||
def test_override_log_never_prints_secret_value(self, capsys):
|
||||
"""Environment-backed credentials must not be copied into logs."""
|
||||
load_config = get_load_config_module()
|
||||
|
||||
secret = 'database-password-that-must-not-leak'
|
||||
cfg = {'database': {'postgresql': {'password': ''}}}
|
||||
env = {'DATABASE__POSTGRESQL__PASSWORD': secret}
|
||||
|
||||
with patch.dict(os.environ, env, clear=True):
|
||||
result = load_config._apply_env_overrides_to_config(cfg)
|
||||
|
||||
captured = capsys.readouterr().out
|
||||
assert result['database']['postgresql']['password'] == secret
|
||||
assert 'DATABASE__POSTGRESQL__PASSWORD' in captured
|
||||
assert secret not in captured
|
||||
|
||||
def test_override_int_value(self):
|
||||
"""Test overriding an int value with proper conversion."""
|
||||
load_config = get_load_config_module()
|
||||
@@ -196,6 +212,19 @@ class TestApplyEnvOverridesToConfig:
|
||||
|
||||
assert result['system']['name'] == 'default'
|
||||
|
||||
def test_skip_env_vars_with_empty_path_segments(self, capsys):
|
||||
"""Platform variables such as __CF_USER_TEXT_ENCODING are not config."""
|
||||
load_config = get_load_config_module()
|
||||
|
||||
cfg = {'system': {'name': 'default'}}
|
||||
env = {'__CF_USER_TEXT_ENCODING': '0x1F5:0x0:0x64'}
|
||||
|
||||
with patch.dict(os.environ, env, clear=True):
|
||||
result = load_config._apply_env_overrides_to_config(cfg)
|
||||
|
||||
assert result == cfg
|
||||
assert capsys.readouterr().out == ''
|
||||
|
||||
def test_nested_config_path(self):
|
||||
"""Test overriding deeply nested config."""
|
||||
load_config = get_load_config_module()
|
||||
|
||||
@@ -12,6 +12,7 @@ from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
import contextvars
|
||||
import sys
|
||||
from unittest.mock import Mock, MagicMock
|
||||
from contextlib import contextmanager
|
||||
@@ -360,6 +361,53 @@ class TestAsyncTaskManager:
|
||||
|
||||
wrapper.cancel()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_task_does_not_inherit_request_context(self):
|
||||
"""Long-lived tasks must receive identity through explicit arguments."""
|
||||
|
||||
_, _, AsyncTaskManager = get_taskmgr_classes()
|
||||
mock_app = create_mock_app()
|
||||
manager = AsyncTaskManager(mock_app)
|
||||
request_value = contextvars.ContextVar('request_value', default=None)
|
||||
token = request_value.set('request-scoped-transaction')
|
||||
observed = []
|
||||
|
||||
async def detached_task(captured_workspace: str) -> None:
|
||||
observed.append((request_value.get(), captured_workspace))
|
||||
|
||||
try:
|
||||
wrapper = manager.create_task(detached_task('workspace-a'))
|
||||
await wrapper.task
|
||||
finally:
|
||||
request_value.reset(token)
|
||||
|
||||
assert observed == [(None, 'workspace-a')]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_task_waits_for_registered_transaction_commit(self):
|
||||
_, _, AsyncTaskManager = get_taskmgr_classes()
|
||||
mock_app = create_mock_app()
|
||||
gate = asyncio.get_running_loop().create_future()
|
||||
|
||||
class PersistenceManagerStub:
|
||||
def create_after_commit_gate(self):
|
||||
return gate
|
||||
|
||||
mock_app.persistence_mgr = PersistenceManagerStub()
|
||||
manager = AsyncTaskManager(mock_app)
|
||||
observed = []
|
||||
|
||||
async def background_work() -> None:
|
||||
observed.append('started')
|
||||
|
||||
wrapper = manager.create_task(background_work())
|
||||
await asyncio.sleep(0)
|
||||
assert observed == []
|
||||
|
||||
gate.set_result(None)
|
||||
await wrapper.task
|
||||
assert observed == ['started']
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_stats_counts_correctly(self):
|
||||
"""Test get_stats returns correct counts."""
|
||||
|
||||
@@ -11,9 +11,19 @@ Note: Uses import isolation to break circular import chains.
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from unittest.mock import Mock, MagicMock
|
||||
from contextlib import contextmanager
|
||||
from typing import Generator
|
||||
from unittest.mock import MagicMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolate_database_manager_registry(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Keep decorator tests from mutating the process-wide manager registry."""
|
||||
from langbot.pkg.persistence import database
|
||||
|
||||
monkeypatch.setattr(database, 'preregistered_managers', list(database.preregistered_managers))
|
||||
|
||||
|
||||
@contextmanager
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
# Persistence manager performs the package's database-manager registration;
|
||||
# importing a concrete manager first would enter the historical app/mgr cycle.
|
||||
from langbot.pkg.persistence import mgr as _persistence_mgr # noqa: F401
|
||||
from langbot.pkg.persistence.databases import postgresql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_postgresql_manager_parses_explicit_url_without_string_reassembly(monkeypatch) -> None:
|
||||
captured = None
|
||||
sentinel_engine = object()
|
||||
|
||||
def create_engine(url):
|
||||
nonlocal captured
|
||||
captured = url
|
||||
return sentinel_engine
|
||||
|
||||
monkeypatch.setattr(postgresql.sqlalchemy_asyncio, 'create_async_engine', create_engine)
|
||||
ap = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(
|
||||
data={
|
||||
'database': {
|
||||
'postgresql': {
|
||||
'url': 'postgresql://runtime:p%40ss@db.internal:5432/langbot?sslmode=require',
|
||||
}
|
||||
}
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
manager = postgresql.PostgreSQLDatabaseManager(ap)
|
||||
await manager.initialize()
|
||||
|
||||
assert captured.drivername == 'postgresql+asyncpg'
|
||||
assert captured.password == 'p@ss'
|
||||
assert captured.query['ssl'] == 'require'
|
||||
assert 'sslmode' not in captured.query
|
||||
assert manager.engine is sentinel_engine
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_postgresql_manager_builds_structured_url_with_special_password(monkeypatch) -> None:
|
||||
captured = None
|
||||
|
||||
def create_engine(url):
|
||||
nonlocal captured
|
||||
captured = url
|
||||
return object()
|
||||
|
||||
monkeypatch.setattr(postgresql.sqlalchemy_asyncio, 'create_async_engine', create_engine)
|
||||
ap = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(
|
||||
data={
|
||||
'database': {
|
||||
'postgresql': {
|
||||
'host': 'db.internal',
|
||||
'port': 5432,
|
||||
'user': 'runtime',
|
||||
'password': 'p@ss:/?#word',
|
||||
'database': 'langbot',
|
||||
}
|
||||
}
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
await postgresql.PostgreSQLDatabaseManager(ap).initialize()
|
||||
|
||||
assert captured.password == 'p@ss:/?#word'
|
||||
assert captured.host == 'db.internal'
|
||||
assert captured.database == 'langbot'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_postgresql_manager_rejects_non_postgresql_url_without_echoing_secret() -> None:
|
||||
ap = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(
|
||||
data={'database': {'postgresql': {'url': 'sqlite:///operator-super-secret.db'}}}
|
||||
)
|
||||
)
|
||||
|
||||
manager = postgresql.PostgreSQLDatabaseManager(ap)
|
||||
with pytest.raises(ValueError, match='valid PostgreSQL') as exc_info:
|
||||
await manager.initialize()
|
||||
assert 'operator-super-secret' not in str(exc_info.value)
|
||||
@@ -0,0 +1,153 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import sqlalchemy
|
||||
|
||||
from langbot.__main__ import _build_parser
|
||||
from langbot.pkg.persistence import release_migration
|
||||
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
|
||||
|
||||
|
||||
def _cloud_config(*, database_use: str = 'postgresql', runtime_user: str = 'langbot_runtime') -> dict:
|
||||
return {
|
||||
'database': {
|
||||
'use': database_use,
|
||||
'postgresql': {
|
||||
'host': 'runtime-db',
|
||||
'port': 5432,
|
||||
'user': runtime_user,
|
||||
'password': 'runtime-secret',
|
||||
'database': 'langbot',
|
||||
},
|
||||
'cloud_migration': {
|
||||
'operator_dsn_env': 'TEST_LANGBOT_OPERATOR_DSN',
|
||||
},
|
||||
},
|
||||
'vdb': {
|
||||
'use': 'pgvector',
|
||||
'pgvector': {
|
||||
'use_business_database': True,
|
||||
'allowed_dimensions': [384, 1536],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _operator_environ(
|
||||
*,
|
||||
user: str = 'langbot_migrator',
|
||||
database: str = 'langbot',
|
||||
host: str = 'runtime-db',
|
||||
port: int = 5432,
|
||||
) -> dict[str, str]:
|
||||
return {
|
||||
'TEST_LANGBOT_OPERATOR_DSN': (f'postgresql://{user}:operator%40secret@{host}:{port}/{database}?sslmode=require')
|
||||
}
|
||||
|
||||
|
||||
def test_cloud_migration_cli_is_explicit() -> None:
|
||||
args = _build_parser().parse_args(['migrate', '--cloud'])
|
||||
assert args.command == 'migrate'
|
||||
assert args.cloud is True
|
||||
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
_build_parser().parse_args(['migrate'])
|
||||
assert exc_info.value.code == 2
|
||||
|
||||
|
||||
def test_operator_url_is_separate_and_preserves_escaped_secret() -> None:
|
||||
url = release_migration._operator_database_url(
|
||||
_cloud_config(),
|
||||
environ=_operator_environ(),
|
||||
)
|
||||
|
||||
assert url.drivername == 'postgresql+asyncpg'
|
||||
assert url.username == 'langbot_migrator'
|
||||
assert url.password == 'operator@secret'
|
||||
assert url.host == 'runtime-db'
|
||||
assert url.port == 5432
|
||||
assert url.database == 'langbot'
|
||||
assert url.query['ssl'] == 'require'
|
||||
assert 'sslmode' not in url.query
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('config', 'environ', 'message'),
|
||||
[
|
||||
(_cloud_config(database_use='sqlite'), _operator_environ(), 'SQLite fallback is forbidden'),
|
||||
(_cloud_config(), {}, 'requires the operator DSN'),
|
||||
(_cloud_config(), {'TEST_LANGBOT_OPERATOR_DSN': 'not a secret://operator-password'}, 'DSN is invalid'),
|
||||
(
|
||||
_cloud_config(),
|
||||
{'TEST_LANGBOT_OPERATOR_DSN': 'postgresql://operator:secret@runtime-db:not-a-port/langbot'},
|
||||
'DSN is invalid',
|
||||
),
|
||||
(_cloud_config(), _operator_environ(user='langbot_runtime'), 'distinct operator role'),
|
||||
(_cloud_config(), _operator_environ(database='another_database'), 'configured runtime database'),
|
||||
(_cloud_config(), _operator_environ(host='other-cluster'), 'runtime PostgreSQL endpoint'),
|
||||
(_cloud_config(), _operator_environ(port=6432), 'runtime PostgreSQL endpoint'),
|
||||
],
|
||||
)
|
||||
def test_operator_url_rejects_unsafe_configuration(config: dict, environ: dict[str, str], message: str) -> None:
|
||||
with pytest.raises(release_migration.CloudReleaseMigrationConfigurationError, match=message) as exc_info:
|
||||
release_migration._operator_database_url(config, environ=environ)
|
||||
assert 'operator-password' not in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_migration_disposes_operator_engine_on_failure(monkeypatch) -> None:
|
||||
engine = SimpleNamespace(dispose=AsyncMock())
|
||||
manager = SimpleNamespace(
|
||||
db=SimpleNamespace(engine=engine),
|
||||
initialize=AsyncMock(side_effect=RuntimeError('migration failed')),
|
||||
)
|
||||
|
||||
def manager_factory(*args, **kwargs):
|
||||
del args, kwargs
|
||||
return manager
|
||||
|
||||
monkeypatch.setattr(release_migration, 'PersistenceManager', manager_factory)
|
||||
ap = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(data=_cloud_config()),
|
||||
logger=logging.getLogger('release-migration-disposal-test'),
|
||||
persistence_mgr=None,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match='migration failed'):
|
||||
await release_migration.run_cloud_release_migration(ap, environ=_operator_environ())
|
||||
|
||||
assert ap.persistence_mgr is manager
|
||||
engine.dispose.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_mode_rejects_sqlite_before_schema_changes(tmp_path, monkeypatch) -> None:
|
||||
from langbot.pkg.persistence import mgr as persistence_mgr_module
|
||||
from langbot.pkg.persistence.databases.sqlite import SQLiteDatabaseManager
|
||||
|
||||
monkeypatch.setattr(persistence_mgr_module.database, 'preregistered_managers', [SQLiteDatabaseManager])
|
||||
ap = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(
|
||||
data={
|
||||
'database': {
|
||||
'use': 'sqlite',
|
||||
'sqlite': {'path': str(tmp_path / 'must-not-migrate.db')},
|
||||
}
|
||||
}
|
||||
),
|
||||
logger=logging.getLogger('release-migration-sqlite-rejection-test'),
|
||||
)
|
||||
manager = PersistenceManager(ap, mode=PersistenceMode.RELEASE_MIGRATION)
|
||||
with pytest.raises(RuntimeError, match='requires PostgreSQL'):
|
||||
await manager.initialize()
|
||||
await manager.get_db_engine().dispose()
|
||||
|
||||
engine = sqlalchemy.create_engine(f'sqlite:///{tmp_path / "must-not-migrate.db"}')
|
||||
try:
|
||||
assert sqlalchemy.inspect(engine).get_table_names() == []
|
||||
finally:
|
||||
engine.dispose()
|
||||
@@ -1,11 +1,23 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextvars
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from langbot.pkg.core.task_boundary import create_detached_task
|
||||
from langbot.pkg.entity.persistence.user import User
|
||||
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
|
||||
from langbot.pkg.persistence.tenant_uow import TenantUnitOfWork
|
||||
from langbot.pkg.persistence.tenant_uow import (
|
||||
CrossScopeTransactionError,
|
||||
PersistenceScopeKind,
|
||||
TenantScopeRequiredError,
|
||||
TenantUnitOfWork,
|
||||
TransactionRollbackOnlyError,
|
||||
)
|
||||
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
@@ -61,3 +73,340 @@ def test_persistence_mode_must_be_a_trusted_enum() -> None:
|
||||
|
||||
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||
assert manager.mode is PersistenceMode.CLOUD_RUNTIME
|
||||
|
||||
|
||||
async def test_manager_reuses_one_session_and_rejects_cross_workspace(tmp_path) -> None:
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "manager-uow.db"}')
|
||||
table = sa.Table(
|
||||
'manager_rows',
|
||||
sa.MetaData(),
|
||||
sa.Column('id', sa.Integer, primary_key=True),
|
||||
sa.Column('workspace_uuid', sa.String(36), nullable=False),
|
||||
)
|
||||
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||
manager.db = SimpleNamespace(get_engine=lambda: engine)
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(table.metadata.create_all)
|
||||
|
||||
with pytest.raises(TenantScopeRequiredError, match='explicit Workspace or discovery'):
|
||||
await manager.execute_async(sa.select(table))
|
||||
|
||||
async with manager.tenant_uow('workspace-a') as outer:
|
||||
await manager.execute_async(sa.insert(table).values(id=1, workspace_uuid='workspace-a'))
|
||||
async with manager.tenant_uow('workspace-a') as inner:
|
||||
assert inner.session is outer.session
|
||||
assert manager.current_session() is outer.session
|
||||
assert (await manager.execute_async(sa.select(table.c.id))).scalar_one() == 1
|
||||
with pytest.raises(CrossScopeTransactionError, match='while workspace scope is active'):
|
||||
async with manager.tenant_uow('workspace-b'):
|
||||
pass
|
||||
|
||||
async with engine.connect() as conn:
|
||||
assert (await conn.execute(sa.select(table.c.id))).scalars().all() == [1]
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_manager_scoped_execute_preserves_core_row_and_scalar_contract(tmp_path) -> None:
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "manager-result-contract.db"}')
|
||||
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||
manager.db = SimpleNamespace(get_engine=lambda: engine)
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(User.__table__.create)
|
||||
await conn.execute(
|
||||
sa.insert(User).values(
|
||||
uuid='account-a',
|
||||
user='owner@example.com',
|
||||
normalized_email='owner@example.com',
|
||||
password='hashed-password',
|
||||
)
|
||||
)
|
||||
|
||||
async with manager.tenant_uow('workspace-a') as uow:
|
||||
result = await manager.execute_async(sa.select(User))
|
||||
row = result.first()
|
||||
assert row is not None
|
||||
assert row.uuid == 'account-a'
|
||||
assert row.user == 'owner@example.com'
|
||||
|
||||
scalar_result = await manager.execute_async(sa.select(User.uuid))
|
||||
assert scalar_result.scalar_one() == 'account-a'
|
||||
|
||||
list_result = await manager.execute_async(sa.select(User.user))
|
||||
assert list_result.scalars().all() == ['owner@example.com']
|
||||
|
||||
# Direct UoW execution remains an ORM API for code that opts into it.
|
||||
orm_result = await uow.execute(sa.select(User))
|
||||
assert orm_result.scalars().one().uuid == 'account-a'
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_transaction_free_tenant_scope_opens_one_short_uow_per_database_call(tmp_path) -> None:
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "short-scope.db"}')
|
||||
table = sa.Table(
|
||||
'short_scope_rows',
|
||||
sa.MetaData(),
|
||||
sa.Column('id', sa.Integer, primary_key=True),
|
||||
sa.Column('workspace_uuid', sa.String(36), nullable=False),
|
||||
)
|
||||
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||
manager.db = SimpleNamespace(get_engine=lambda: engine)
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(table.metadata.create_all)
|
||||
|
||||
async with manager.tenant_scope('workspace-a'):
|
||||
assert manager.current_scope() is not None
|
||||
assert manager.current_scope().kind is PersistenceScopeKind.WORKSPACE
|
||||
assert manager.current_scope().settings == (('langbot.workspace_uuid', 'workspace-a'),)
|
||||
assert manager.current_session() is None
|
||||
|
||||
await manager.execute_async(sa.insert(table).values(id=1, workspace_uuid='workspace-a'))
|
||||
assert manager.current_session() is None
|
||||
|
||||
# The first statement has already committed. Long external waits
|
||||
# inside this boundary retain only identity, never a DB session.
|
||||
await asyncio.sleep(0)
|
||||
assert manager.current_session() is None
|
||||
|
||||
assert (await manager.execute_async(sa.select(table.c.id))).scalar_one() == 1
|
||||
assert manager.current_session() is None
|
||||
|
||||
async with manager.tenant_scope('workspace-a'):
|
||||
assert manager.current_session() is None
|
||||
with pytest.raises(CrossScopeTransactionError, match='while workspace scope is active'):
|
||||
async with manager.tenant_scope('workspace-b'):
|
||||
pass
|
||||
with pytest.raises(CrossScopeTransactionError, match='while workspace scope is active'):
|
||||
async with manager.tenant_uow('workspace-b'):
|
||||
pass
|
||||
|
||||
assert manager.current_scope() is None
|
||||
with pytest.raises(TenantScopeRequiredError, match='explicit Workspace'):
|
||||
await manager.execute_async(sa.select(table))
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_transaction_free_scope_requires_explicit_child_task_scope(tmp_path) -> None:
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "short-scope-child.db"}')
|
||||
table = sa.Table('child_scope_rows', sa.MetaData(), sa.Column('id', sa.Integer, primary_key=True))
|
||||
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||
manager.db = SimpleNamespace(get_engine=lambda: engine)
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(table.metadata.create_all)
|
||||
|
||||
async with manager.tenant_scope('workspace-a'):
|
||||
|
||||
async def inherited_access() -> None:
|
||||
await manager.execute_async(sa.select(table))
|
||||
|
||||
with pytest.raises(CrossScopeTransactionError, match='cannot be inherited by child tasks'):
|
||||
await asyncio.create_task(inherited_access())
|
||||
|
||||
async def explicitly_scoped_access() -> None:
|
||||
async with manager.tenant_scope('workspace-a'):
|
||||
await manager.execute_async(sa.insert(table).values(id=1))
|
||||
assert manager.current_session() is None
|
||||
|
||||
await asyncio.create_task(explicitly_scoped_access())
|
||||
assert manager.current_session() is None
|
||||
|
||||
async with engine.connect() as conn:
|
||||
assert (await conn.execute(sa.select(table.c.id))).scalars().all() == [1]
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_caught_nested_failure_marks_outer_transaction_rollback_only(tmp_path) -> None:
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "rollback-only.db"}')
|
||||
table = sa.Table('rollback_rows', sa.MetaData(), sa.Column('id', sa.Integer, primary_key=True))
|
||||
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||
manager.db = SimpleNamespace(get_engine=lambda: engine)
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(table.metadata.create_all)
|
||||
|
||||
with pytest.raises(TransactionRollbackOnlyError, match='transaction was rolled back'):
|
||||
async with manager.tenant_uow('workspace-a'):
|
||||
await manager.execute_async(sa.insert(table).values(id=1))
|
||||
try:
|
||||
async with manager.tenant_uow('workspace-a'):
|
||||
raise ValueError('caught nested failure')
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
async with engine.connect() as conn:
|
||||
assert (await conn.execute(sa.select(table))).all() == []
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest.mark.parametrize('executor_kind', ['manager', 'uow', 'session'])
|
||||
async def test_caught_database_error_rolls_back_and_cancels_after_commit_gate(tmp_path, executor_kind: str) -> None:
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / f"db-error-{executor_kind}.db"}')
|
||||
table = sa.Table('unique_rows', sa.MetaData(), sa.Column('id', sa.Integer, primary_key=True))
|
||||
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||
manager.db = SimpleNamespace(get_engine=lambda: engine)
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(table.metadata.create_all)
|
||||
|
||||
with pytest.raises(TransactionRollbackOnlyError, match='after-commit work was cancelled'):
|
||||
async with manager.tenant_uow('workspace-a') as uow:
|
||||
statement = sa.insert(table).values(id=1)
|
||||
if executor_kind == 'manager':
|
||||
await manager.execute_async(statement)
|
||||
elif executor_kind == 'uow':
|
||||
await uow.execute(statement)
|
||||
else:
|
||||
await uow.session.execute(statement)
|
||||
|
||||
gate = manager.create_after_commit_gate()
|
||||
assert gate is not None
|
||||
try:
|
||||
if executor_kind == 'manager':
|
||||
await manager.execute_async(statement)
|
||||
elif executor_kind == 'uow':
|
||||
await uow.execute(statement)
|
||||
else:
|
||||
await uow.session.execute(statement)
|
||||
except sa.exc.IntegrityError:
|
||||
pass
|
||||
|
||||
assert gate.cancelled()
|
||||
async with engine.connect() as conn:
|
||||
assert (await conn.execute(sa.select(table))).all() == []
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_child_task_must_open_its_own_explicit_uow(tmp_path) -> None:
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "child-task.db"}')
|
||||
table = sa.Table(
|
||||
'child_rows',
|
||||
sa.MetaData(),
|
||||
sa.Column('id', sa.Integer, primary_key=True),
|
||||
sa.Column('workspace_uuid', sa.String(36), nullable=False),
|
||||
)
|
||||
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||
manager.db = SimpleNamespace(get_engine=lambda: engine)
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(table.metadata.create_all)
|
||||
|
||||
async with manager.tenant_uow('workspace-a'):
|
||||
|
||||
async def inherited_access() -> None:
|
||||
await manager.execute_async(sa.select(table))
|
||||
|
||||
with pytest.raises(CrossScopeTransactionError, match='cannot be inherited by child tasks'):
|
||||
await asyncio.create_task(inherited_access())
|
||||
|
||||
async def explicitly_scoped_access() -> None:
|
||||
async with manager.tenant_uow('workspace-a'):
|
||||
await manager.execute_async(sa.insert(table).values(id=2, workspace_uuid='workspace-a'))
|
||||
|
||||
await asyncio.create_task(explicitly_scoped_access())
|
||||
|
||||
async with engine.connect() as conn:
|
||||
assert (await conn.execute(sa.select(table.c.id))).scalars().all() == [2]
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_detached_task_starts_without_parent_scope_and_rolls_back_its_uow(tmp_path) -> None:
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "detached-task.db"}')
|
||||
table = sa.Table(
|
||||
'detached_rows',
|
||||
sa.MetaData(),
|
||||
sa.Column('id', sa.Integer, primary_key=True),
|
||||
sa.Column('workspace_uuid', sa.String(36), nullable=False),
|
||||
)
|
||||
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||
manager.db = SimpleNamespace(get_engine=lambda: engine)
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(table.metadata.create_all)
|
||||
|
||||
async with manager.tenant_uow('workspace-a'):
|
||||
|
||||
async def detached_write() -> None:
|
||||
# Before the detached boundary this access raises
|
||||
# CrossScopeTransactionError because asyncio copies the
|
||||
# parent's ActiveScopedTransaction into this child task.
|
||||
assert manager.current_scope() is None
|
||||
async with manager.tenant_uow('workspace-a'):
|
||||
await manager.execute_async(sa.insert(table).values(id=2, workspace_uuid='workspace-a'))
|
||||
raise RuntimeError('roll back detached write')
|
||||
|
||||
task = create_detached_task(detached_write())
|
||||
with pytest.raises(RuntimeError, match='roll back detached write'):
|
||||
await task
|
||||
|
||||
await manager.execute_async(sa.insert(table).values(id=1, workspace_uuid='workspace-a'))
|
||||
|
||||
async with engine.connect() as conn:
|
||||
assert (await conn.execute(sa.select(table.c.id))).scalars().all() == [1]
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_after_commit_task_waits_for_commit_and_starts_with_empty_context(tmp_path) -> None:
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "after-commit.db"}')
|
||||
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||
manager.db = SimpleNamespace(get_engine=lambda: engine)
|
||||
request_value = contextvars.ContextVar('after_commit_request_value', default=None)
|
||||
observed = []
|
||||
try:
|
||||
async with manager.tenant_uow('workspace-a'):
|
||||
token = request_value.set('request-scope')
|
||||
|
||||
async def after_commit_work() -> None:
|
||||
observed.append((request_value.get(), manager.current_scope()))
|
||||
|
||||
try:
|
||||
task = create_detached_task(
|
||||
after_commit_work(),
|
||||
after_commit_manager=manager,
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
assert observed == []
|
||||
finally:
|
||||
request_value.reset(token)
|
||||
|
||||
await task
|
||||
assert observed == [(None, None)]
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_after_commit_task_is_cancelled_and_coroutine_closed_on_rollback(tmp_path) -> None:
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "after-rollback.db"}')
|
||||
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||
manager.db = SimpleNamespace(get_engine=lambda: engine)
|
||||
started = False
|
||||
|
||||
async def should_not_start() -> None:
|
||||
nonlocal started
|
||||
started = True
|
||||
|
||||
coro = should_not_start()
|
||||
try:
|
||||
with pytest.raises(RuntimeError, match='rollback request'):
|
||||
async with manager.tenant_uow('workspace-a'):
|
||||
task = create_detached_task(coro, after_commit_manager=manager)
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
raise RuntimeError('rollback request')
|
||||
|
||||
await asyncio.sleep(0)
|
||||
assert task.cancelled()
|
||||
assert not started
|
||||
assert coro.cr_frame is None
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
@@ -13,8 +13,11 @@ from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
import contextvars
|
||||
from contextlib import asynccontextmanager
|
||||
from unittest.mock import Mock, AsyncMock
|
||||
from importlib import import_module
|
||||
from types import SimpleNamespace
|
||||
|
||||
from tests.factories import (
|
||||
FakeApp,
|
||||
@@ -894,17 +897,48 @@ class TestMessageAggregatorWorkspaceIsolation:
|
||||
app = make_aggregator_app()
|
||||
enable_aggregation(app)
|
||||
agg = get_aggregator_module().MessageAggregator(app)
|
||||
delayed_flush = AsyncMock()
|
||||
request_value = contextvars.ContextVar('aggregator_request_value', default=None)
|
||||
token = request_value.set('request-scope')
|
||||
observed = []
|
||||
|
||||
async def delayed_flush(*args):
|
||||
observed.append((request_value.get(), args[2]))
|
||||
|
||||
monkeypatch.setattr(agg, '_delayed_flush', delayed_flush)
|
||||
context = execution_context(pipeline_uuid='test-pipeline')
|
||||
|
||||
await agg.add_message(**scoped_message_kwargs(context))
|
||||
await asyncio.sleep(0)
|
||||
try:
|
||||
await agg.add_message(**scoped_message_kwargs(context))
|
||||
await asyncio.sleep(0)
|
||||
finally:
|
||||
request_value.reset(token)
|
||||
|
||||
delayed_flush.assert_awaited_once()
|
||||
assert delayed_flush.await_args.args[2] is context
|
||||
assert observed == [(None, context)]
|
||||
await agg.flush_all()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delayed_flush_opens_explicit_workspace_uow(self, monkeypatch):
|
||||
app = make_aggregator_app()
|
||||
app.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime')
|
||||
scopes = []
|
||||
|
||||
@asynccontextmanager
|
||||
async def tenant_uow(workspace_uuid):
|
||||
scopes.append(workspace_uuid)
|
||||
yield
|
||||
|
||||
app.persistence_mgr.tenant_uow = tenant_uow
|
||||
agg = get_aggregator_module().MessageAggregator(app)
|
||||
flush = AsyncMock()
|
||||
monkeypatch.setattr(agg, '_flush_buffer', flush)
|
||||
context = execution_context('workspace-a', pipeline_uuid='test-pipeline')
|
||||
key = aggregation_key(context, pipeline_uuid='test-pipeline')
|
||||
|
||||
await agg._delayed_flush(key, 0, context)
|
||||
|
||||
assert scopes == ['workspace-a']
|
||||
flush.assert_awaited_once_with(key, context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_rejects_context_from_another_workspace(self):
|
||||
app = make_aggregator_app()
|
||||
|
||||
@@ -1,10 +1,15 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
|
||||
from langbot.pkg.persistence.tenant_uow import PersistenceScopeKind
|
||||
from langbot.pkg.pipeline.controller import Controller
|
||||
from langbot.pkg.workspace.errors import WorkspaceGenerationMismatchError
|
||||
|
||||
@@ -44,6 +49,79 @@ async def test_controller_drops_stale_query_before_pipeline_lookup(
|
||||
query_pool.condition.notify_all.assert_called_once_with()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cloud_controller_releases_database_connection_during_pipeline_wait(
|
||||
tmp_path,
|
||||
mock_app,
|
||||
sample_query,
|
||||
):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "pipeline-short-scope.db"}')
|
||||
table = sa.Table('pipeline_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)
|
||||
|
||||
_prepare_scheduler(mock_app)
|
||||
mock_app.persistence_mgr = manager
|
||||
pipeline_waiting = asyncio.Event()
|
||||
release_pipeline = asyncio.Event()
|
||||
|
||||
async def get_binding(*_args, **_kwargs):
|
||||
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
|
||||
return SimpleNamespace(
|
||||
instance_uuid='test-instance',
|
||||
workspace_uuid='test-workspace',
|
||||
placement_generation=1,
|
||||
)
|
||||
|
||||
async def run_pipeline(_query):
|
||||
await manager.execute_async(sa.select(table.c.id))
|
||||
assert manager.current_session() is None
|
||||
pipeline_waiting.set()
|
||||
await release_pipeline.wait()
|
||||
assert manager.current_scope().kind is PersistenceScopeKind.WORKSPACE
|
||||
assert manager.current_session() is None
|
||||
|
||||
runtime_pipeline = SimpleNamespace(run=AsyncMock(side_effect=run_pipeline))
|
||||
|
||||
async def get_pipeline(*_args, **_kwargs):
|
||||
await manager.execute_async(sa.select(table.c.id))
|
||||
assert manager.current_session() is None
|
||||
return runtime_pipeline
|
||||
|
||||
mock_app.workspace_service.get_execution_binding = AsyncMock(side_effect=get_binding)
|
||||
mock_app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(side_effect=get_pipeline)
|
||||
controller = Controller(mock_app)
|
||||
|
||||
task = asyncio.create_task(controller._process_query(sample_query))
|
||||
await asyncio.wait_for(pipeline_waiting.wait(), timeout=2)
|
||||
assert checked_out == 0
|
||||
assert not task.done()
|
||||
release_pipeline.set()
|
||||
await asyncio.wait_for(task, timeout=2)
|
||||
assert checked_out == 0
|
||||
runtime_pipeline.run.assert_awaited_once_with(sample_query)
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_controller_revalidates_generation_before_running_pipeline(
|
||||
mock_app,
|
||||
|
||||
@@ -1,12 +1,15 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.api.http.authz import WorkspaceRequiredError
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.platform.botmgr import PlatformManager, RuntimeBot
|
||||
import langbot_plugin.api.entities.builtin.platform.events as platform_events
|
||||
|
||||
|
||||
WORKSPACE_A = '00000000-0000-0000-0000-00000000000a'
|
||||
@@ -112,3 +115,71 @@ def test_runtime_bot_rejects_workspace_mismatch():
|
||||
logger=SimpleNamespace(),
|
||||
execution_context=_context(WORKSPACE_B, BOT_A),
|
||||
)
|
||||
|
||||
|
||||
class _ScopeOnlyPersistenceManager:
|
||||
mode = SimpleNamespace(value='cloud_runtime')
|
||||
|
||||
def __init__(self):
|
||||
self.active_workspace = None
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def tenant_scope(self, workspace_uuid: str):
|
||||
assert self.active_workspace is None
|
||||
self.active_workspace = workspace_uuid
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self.active_workspace = None
|
||||
|
||||
def current_session(self):
|
||||
return None
|
||||
|
||||
|
||||
class _ListenerAdapter:
|
||||
def __init__(self):
|
||||
self.listeners = {}
|
||||
|
||||
def register_listener(self, event_type, listener):
|
||||
self.listeners[event_type] = listener
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_platform_callback_carries_scope_without_holding_database_session():
|
||||
persistence_mgr = _ScopeOnlyPersistenceManager()
|
||||
adapter = _ListenerAdapter()
|
||||
|
||||
async def push_person_message(*_args, **_kwargs):
|
||||
assert persistence_mgr.active_workspace == WORKSPACE_A
|
||||
assert persistence_mgr.current_session() is None
|
||||
return True
|
||||
|
||||
application = SimpleNamespace(
|
||||
persistence_mgr=persistence_mgr,
|
||||
workspace_service=_WorkspaceService(),
|
||||
webhook_pusher=SimpleNamespace(push_person_message=push_person_message),
|
||||
)
|
||||
entity = SimpleNamespace(
|
||||
uuid=BOT_A,
|
||||
workspace_uuid=WORKSPACE_A,
|
||||
name='Bot',
|
||||
enable=True,
|
||||
pipeline_routing_rules=[],
|
||||
use_pipeline_uuid=None,
|
||||
)
|
||||
logger = SimpleNamespace(info=AsyncMock(), error=AsyncMock())
|
||||
runtime = RuntimeBot(
|
||||
ap=application,
|
||||
bot_entity=entity,
|
||||
adapter=adapter,
|
||||
logger=logger,
|
||||
execution_context=_context(WORKSPACE_A, BOT_A),
|
||||
)
|
||||
await runtime.initialize()
|
||||
|
||||
listener = adapter.listeners[platform_events.FriendMessage]
|
||||
event = SimpleNamespace(message_chain=[], sender=SimpleNamespace(id='user'))
|
||||
await listener(event, adapter)
|
||||
|
||||
assert persistence_mgr.active_workspace is None
|
||||
logger.info.assert_awaited()
|
||||
|
||||
@@ -9,19 +9,32 @@ Tests cover:
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import nullcontext
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock, AsyncMock
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
from importlib import import_module
|
||||
|
||||
from tests.factories import text_query
|
||||
from langbot_plugin.entities.io.context import ActionContext
|
||||
from langbot_plugin.entities.io.context import InstallationBinding
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
|
||||
|
||||
TEST_ACTION_CONTEXT = ActionContext(
|
||||
TEST_EXECUTION_CONTEXT = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
)
|
||||
TEST_INSTALLATION_BINDING = InstallationBinding(
|
||||
instance_uuid=TEST_EXECUTION_CONTEXT.instance_uuid,
|
||||
workspace_uuid=TEST_EXECUTION_CONTEXT.workspace_uuid,
|
||||
placement_generation=TEST_EXECUTION_CONTEXT.placement_generation,
|
||||
installation_uuid='00000000-0000-4000-8000-000000000001',
|
||||
runtime_revision=1,
|
||||
artifact_digest='a' * 64,
|
||||
)
|
||||
|
||||
|
||||
def get_connector_module():
|
||||
@@ -37,6 +50,7 @@ def create_mock_app():
|
||||
mock_app.instance_config.data = {'plugin': {'enable': True}}
|
||||
mock_app.persistence_mgr = AsyncMock()
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock()
|
||||
mock_app.persistence_mgr.tenant_uow = None
|
||||
return mock_app
|
||||
|
||||
|
||||
@@ -47,7 +61,19 @@ def create_mock_connector():
|
||||
async def mock_disconnect_callback(conn):
|
||||
pass
|
||||
|
||||
return connector.PluginRuntimeConnector(create_mock_app(), mock_disconnect_callback)
|
||||
instance = connector.PluginRuntimeConnector(create_mock_app(), mock_disconnect_callback)
|
||||
instance._execution_context.set(TEST_EXECUTION_CONTEXT)
|
||||
instance._operation_bindings = AsyncMock(return_value=[TEST_INSTALLATION_BINDING])
|
||||
instance._target_binding = AsyncMock(return_value=TEST_INSTALLATION_BINDING)
|
||||
instance._load_workspace_settings = AsyncMock(return_value=[])
|
||||
instance.require_workspace_context = AsyncMock(side_effect=lambda context: context)
|
||||
return instance
|
||||
|
||||
|
||||
def configure_handler(connector, runtime_handler):
|
||||
runtime_handler.installation_scope = Mock(side_effect=lambda _binding: nullcontext())
|
||||
connector.handler = runtime_handler
|
||||
return runtime_handler
|
||||
|
||||
|
||||
class TestListPlugins:
|
||||
@@ -76,8 +102,7 @@ class TestListPlugins:
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
connector.handler.require_bound_action_context = Mock(return_value=TEST_ACTION_CONTEXT)
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.list_plugins = AsyncMock(
|
||||
return_value=[{'manifest': {'manifest': {'metadata': {'author': 'test', 'name': 'plugin'}}}}]
|
||||
)
|
||||
@@ -86,8 +111,7 @@ class TestListPlugins:
|
||||
|
||||
connector.handler.list_plugins.assert_called_once()
|
||||
assert result == [{'manifest': {'manifest': {'metadata': {'author': 'test', 'name': 'plugin'}}}}]
|
||||
statement = connector.ap.persistence_mgr.execute_async.await_args.args[0]
|
||||
assert 'workspace-a' in statement.compile().params.values()
|
||||
connector._load_workspace_settings.assert_awaited_once_with(TEST_EXECUTION_CONTEXT)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filters_by_component_kinds(self):
|
||||
@@ -95,8 +119,7 @@ class TestListPlugins:
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
connector.handler.require_bound_action_context = Mock(return_value=TEST_ACTION_CONTEXT)
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.list_plugins = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
@@ -123,8 +146,7 @@ class TestListPlugins:
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
connector.handler.require_bound_action_context = Mock(return_value=TEST_ACTION_CONTEXT)
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.list_plugins = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
@@ -171,7 +193,7 @@ class TestPluginDiagnostics:
|
||||
'response_sources': response_sources,
|
||||
}
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.emit_event = AsyncMock(side_effect=emit_event_response)
|
||||
|
||||
fake_event_ctx = Mock()
|
||||
@@ -215,7 +237,7 @@ class TestPluginDiagnostics:
|
||||
],
|
||||
}
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.emit_event = AsyncMock(side_effect=emit_event_response)
|
||||
|
||||
fake_event_ctx = Mock()
|
||||
@@ -238,7 +260,7 @@ class TestPluginDiagnostics:
|
||||
connector_module.context.EventContext.from_event = original_from_event
|
||||
connector_module.context.EventContext.model_validate = original_model_validate
|
||||
|
||||
assert '_response_sources' not in vars(event_ctx)
|
||||
assert event_ctx._response_sources == []
|
||||
assert event_ctx._emitted_plugins == [
|
||||
{'manifest': {'metadata': {'author': 'tester', 'name': 'demo'}}},
|
||||
]
|
||||
@@ -253,7 +275,7 @@ class TestPluginDiagnostics:
|
||||
mock_app = create_mock_app()
|
||||
mock_app.instance_config.data = {'plugin': {'enable': False}}
|
||||
connector = connector_module.PluginRuntimeConnector(mock_app, mock_disconnect)
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
|
||||
await connector.notify_plugin_diagnostic({'code': 'response_delivery_failed'})
|
||||
|
||||
@@ -262,7 +284,7 @@ class TestPluginDiagnostics:
|
||||
@pytest.mark.asyncio
|
||||
async def test_notify_plugin_diagnostic_is_best_effort(self):
|
||||
connector = create_mock_connector()
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.notify_plugin_diagnostic = AsyncMock(side_effect=RuntimeError('action not found'))
|
||||
|
||||
await connector.notify_plugin_diagnostic({'code': 'response_delivery_failed'})
|
||||
@@ -297,7 +319,7 @@ class TestListKnowledgeEngines:
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.list_knowledge_engines = AsyncMock(
|
||||
return_value=[{'plugin_id': 'author/engine', 'name': 'Engine'}]
|
||||
)
|
||||
@@ -334,7 +356,7 @@ class TestListParsers:
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.list_parsers = AsyncMock(
|
||||
return_value=[{'plugin_id': 'author/parser', 'supported_mime_types': ['text/plain']}]
|
||||
)
|
||||
@@ -354,7 +376,7 @@ class TestCallParser:
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.parse_document = AsyncMock(return_value={'content': 'parsed'})
|
||||
|
||||
result = await connector.call_parser(
|
||||
@@ -381,7 +403,7 @@ class TestRAGMethods:
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.rag_ingest_document = AsyncMock(return_value={'status': 'success'})
|
||||
|
||||
result = await connector.call_rag_ingest('author/engine', {'file': 'test.pdf'})
|
||||
@@ -395,7 +417,7 @@ class TestRAGMethods:
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.retrieve_knowledge = AsyncMock(
|
||||
return_value={
|
||||
'results': [
|
||||
@@ -424,7 +446,7 @@ class TestRAGMethods:
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.get_rag_creation_schema = AsyncMock(return_value={'properties': {'name': {'type': 'string'}}})
|
||||
|
||||
result = await connector.get_rag_creation_schema('author/engine')
|
||||
@@ -438,7 +460,7 @@ class TestRAGMethods:
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.get_rag_retrieval_schema = AsyncMock(
|
||||
return_value={'properties': {'top_k': {'type': 'integer'}}}
|
||||
)
|
||||
@@ -454,7 +476,7 @@ class TestRAGMethods:
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.rag_on_kb_create = AsyncMock(return_value={'status': 'ok'})
|
||||
|
||||
await connector.rag_on_kb_create('author/engine', 'kb-uuid', {'model': 'test'})
|
||||
@@ -467,7 +489,7 @@ class TestRAGMethods:
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.rag_on_kb_delete = AsyncMock(return_value={'status': 'ok'})
|
||||
|
||||
await connector.rag_on_kb_delete('author/engine', 'kb-uuid')
|
||||
@@ -480,7 +502,7 @@ class TestRAGMethods:
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.rag_delete_document = AsyncMock(return_value=True)
|
||||
|
||||
result = await connector.call_rag_delete_document('author/engine', 'doc-uuid', 'kb-uuid')
|
||||
@@ -574,7 +596,7 @@ class TestGetPluginInfo:
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.get_plugin_info = AsyncMock(return_value={'manifest': {'metadata': {'name': 'plugin'}}})
|
||||
|
||||
result = await connector.get_plugin_info('author', 'plugin')
|
||||
@@ -587,17 +609,39 @@ class TestSetPluginConfig:
|
||||
"""Tests for set_plugin_config method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_calls_handler_set_plugin_config(self):
|
||||
"""Test that handler.set_plugin_config is called."""
|
||||
async def test_updates_revision_then_applies_desired_state(self):
|
||||
"""Config changes are fenced by a new runtime revision."""
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
connector.handler.set_plugin_config = AsyncMock(return_value={'status': 'ok'})
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.register_installation_binding = Mock()
|
||||
connector.handler.apply_plugin_installation = AsyncMock(return_value={'state': 'running'})
|
||||
setting = SimpleNamespace(
|
||||
installation_uuid=TEST_INSTALLATION_BINDING.installation_uuid,
|
||||
runtime_revision=1,
|
||||
artifact_digest=TEST_INSTALLATION_BINDING.artifact_digest,
|
||||
enabled=True,
|
||||
install_info={'_artifact_storage': 'tenant_binary_storage_v1'},
|
||||
)
|
||||
connector._setting_for_plugin = AsyncMock(return_value=(TEST_EXECUTION_CONTEXT, setting))
|
||||
connector.ap.persistence_mgr.execute_async = AsyncMock(return_value=SimpleNamespace(rowcount=1))
|
||||
|
||||
await connector.set_plugin_config('author', 'plugin', {'setting': 'value'})
|
||||
|
||||
connector.handler.set_plugin_config.assert_called_once_with('author', 'plugin', {'setting': 'value'})
|
||||
applied_binding = connector.handler.apply_plugin_installation.await_args.args[0]
|
||||
assert applied_binding.runtime_revision == 2
|
||||
assert applied_binding.installation_uuid == TEST_INSTALLATION_BINDING.installation_uuid
|
||||
connector.handler.register_installation_binding.assert_called_once_with(
|
||||
applied_binding,
|
||||
plugin_author='author',
|
||||
plugin_name='plugin',
|
||||
)
|
||||
connector.handler.apply_plugin_installation.assert_awaited_once_with(
|
||||
applied_binding,
|
||||
artifact_package=None,
|
||||
enabled=True,
|
||||
)
|
||||
|
||||
|
||||
class TestPingPluginRuntime:
|
||||
@@ -621,7 +665,7 @@ class TestPingPluginRuntime:
|
||||
get_connector_module()
|
||||
connector = create_mock_connector()
|
||||
|
||||
connector.handler = AsyncMock()
|
||||
configure_handler(connector, AsyncMock())
|
||||
connector.handler.ping = AsyncMock(return_value={'status': 'ok'})
|
||||
|
||||
await connector.ping_plugin_runtime()
|
||||
|
||||
@@ -7,7 +7,6 @@ import pytest
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.plugin.connector import PluginRuntimeConnector, PluginRuntimeNotConnectedError
|
||||
from langbot_plugin.entities.io.context import ActionContext
|
||||
from langbot_plugin.runtime.security import (
|
||||
PLUGIN_RUNTIME_CONTROL_TOKEN_ENV,
|
||||
PLUGIN_RUNTIME_CONTROL_TOKEN_HEADER,
|
||||
@@ -91,7 +90,7 @@ async def test_enabled_connector_reports_not_connected_after_workspace_validatio
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connector_resolves_oss_singleton_binding_from_workspace_service():
|
||||
async def test_oss_connector_resolves_singleton_only_for_legacy_callers():
|
||||
connector = make_connector()
|
||||
connector.ap.workspace_service = SimpleNamespace(
|
||||
get_local_execution_binding=AsyncMock(
|
||||
@@ -103,31 +102,29 @@ async def test_connector_resolves_oss_singleton_binding_from_workspace_service()
|
||||
)
|
||||
)
|
||||
|
||||
assert await connector._resolve_action_context() == ActionContext(
|
||||
assert await connector._current_execution_context() == ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=3,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edition_metadata_cannot_enable_multi_workspace_runtime_binding():
|
||||
def test_edition_metadata_cannot_enable_shared_runtime_profile():
|
||||
connector = make_connector()
|
||||
connector.ap.instance_config.data['system'] = {'edition': 'cloud'}
|
||||
connector.ap.workspace_service = SimpleNamespace(
|
||||
policy=SimpleNamespace(multi_workspace_enabled=False),
|
||||
get_local_execution_binding=AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
)
|
||||
),
|
||||
|
||||
assert connector.runtime_profile == 'oss_dev'
|
||||
|
||||
|
||||
def test_closed_deployment_selects_instance_scoped_shared_profile():
|
||||
app = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
|
||||
deployment=SimpleNamespace(mode='cloud'),
|
||||
)
|
||||
|
||||
context = await connector._resolve_action_context()
|
||||
connector = PluginRuntimeConnector(app, AsyncMock())
|
||||
|
||||
assert context.workspace_uuid == 'workspace-a'
|
||||
assert connector.runtime_profile == 'shared'
|
||||
|
||||
|
||||
def test_external_runtime_control_headers_require_strong_secret(monkeypatch):
|
||||
@@ -148,48 +145,61 @@ def test_local_runtime_control_headers_generate_ephemeral_secret(monkeypatch):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connector_fails_closed_without_trusted_workspace_service():
|
||||
async def test_oss_legacy_fallback_fails_without_workspace_service():
|
||||
connector = make_connector()
|
||||
|
||||
with pytest.raises(
|
||||
RuntimeError,
|
||||
match='Plugin Runtime Workspace binding is unavailable',
|
||||
):
|
||||
await connector._resolve_action_context()
|
||||
with pytest.raises(AttributeError):
|
||||
await connector._current_execution_context()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cloud_connector_never_falls_back_to_ghost_local_workspace():
|
||||
connector = make_connector()
|
||||
connector.ap.instance_config.data['system'] = {'edition': 'cloud'}
|
||||
app = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
|
||||
deployment=SimpleNamespace(mode='cloud'),
|
||||
)
|
||||
get_local_binding = AsyncMock()
|
||||
connector.ap.workspace_service = SimpleNamespace(
|
||||
policy=SimpleNamespace(multi_workspace_enabled=True),
|
||||
app.workspace_service = SimpleNamespace(
|
||||
get_local_execution_binding=get_local_binding,
|
||||
)
|
||||
connector = PluginRuntimeConnector(app, AsyncMock())
|
||||
|
||||
with pytest.raises(
|
||||
RuntimeError,
|
||||
match='require an explicit projected Workspace binding',
|
||||
):
|
||||
await connector._resolve_action_context()
|
||||
with pytest.raises(Exception, match='Plugin resource not found'):
|
||||
await connector._current_execution_context()
|
||||
|
||||
get_local_binding.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cloud_connector_initialize_fails_before_opening_unbound_runtime():
|
||||
connector = make_connector()
|
||||
connector.ap.instance_config.data['system'] = {'edition': 'cloud'}
|
||||
connector.ap.workspace_service = SimpleNamespace(
|
||||
policy=SimpleNamespace(multi_workspace_enabled=True),
|
||||
def test_worker_policy_is_loaded_only_from_instance_configuration():
|
||||
app = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(
|
||||
data={
|
||||
'plugin': {
|
||||
'enable': True,
|
||||
'worker': {
|
||||
'max_cpus': 1.5,
|
||||
'max_memory_mb': 768,
|
||||
'max_pids': 64,
|
||||
'max_open_files': 128,
|
||||
'max_file_size_mb': 32,
|
||||
'require_hard_limits': True,
|
||||
},
|
||||
# A plugin-controlled value at any other path is ignored.
|
||||
'manifest': {'max_memory_mb': 99999},
|
||||
}
|
||||
}
|
||||
)
|
||||
)
|
||||
connector = PluginRuntimeConnector(app, AsyncMock())
|
||||
|
||||
with pytest.raises(
|
||||
RuntimeError,
|
||||
match='require an explicit projected Workspace binding',
|
||||
):
|
||||
await connector.initialize()
|
||||
policy = connector._load_worker_policy()
|
||||
|
||||
assert policy.max_cpus == 1.5
|
||||
assert policy.max_memory_mb == 768
|
||||
assert policy.max_pids == 64
|
||||
assert policy.max_open_files == 128
|
||||
assert policy.max_file_size_mb == 32
|
||||
assert policy.require_hard_limits is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -198,17 +208,11 @@ async def test_explicit_cloud_binding_is_revalidated_against_projection():
|
||||
instance_config=SimpleNamespace(
|
||||
data={
|
||||
'plugin': {'enable': True},
|
||||
'system': {'edition': 'cloud'},
|
||||
}
|
||||
)
|
||||
)
|
||||
configured = ActionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-cloud-a',
|
||||
placement_generation=7,
|
||||
),
|
||||
deployment=SimpleNamespace(mode='cloud'),
|
||||
)
|
||||
app.workspace_service = SimpleNamespace(
|
||||
policy=SimpleNamespace(multi_workspace_enabled=True),
|
||||
get_execution_binding=AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
instance_uuid='instance-a',
|
||||
@@ -216,13 +220,19 @@ async def test_explicit_cloud_binding_is_revalidated_against_projection():
|
||||
placement_generation=7,
|
||||
)
|
||||
),
|
||||
get_local_execution_binding=AsyncMock(),
|
||||
)
|
||||
connector = PluginRuntimeConnector(app, AsyncMock(), action_context=configured)
|
||||
connector = PluginRuntimeConnector(app, AsyncMock())
|
||||
connector.handler = SimpleNamespace()
|
||||
connector._synchronize_workspace = AsyncMock()
|
||||
configured = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-cloud-a',
|
||||
placement_generation=7,
|
||||
)
|
||||
|
||||
assert await connector._resolve_action_context() == configured
|
||||
assert await connector.require_workspace_context(configured) == configured
|
||||
app.workspace_service.get_execution_binding.assert_awaited_once_with(
|
||||
'workspace-cloud-a',
|
||||
expected_generation=7,
|
||||
)
|
||||
app.workspace_service.get_local_execution_binding.assert_not_awaited()
|
||||
connector._synchronize_workspace.assert_awaited_once_with(configured)
|
||||
|
||||
@@ -0,0 +1,482 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
import hashlib
|
||||
from contextlib import nullcontext
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
from langbot_plugin.entities.io.context import InstallationBinding
|
||||
from langbot_plugin.runtime.plugin.mgr import PluginInstallSource
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.plugin.connector import (
|
||||
PluginInstallationFailedError,
|
||||
PluginRuntimeConnector,
|
||||
)
|
||||
|
||||
|
||||
def connection_result_connector(execute_async: AsyncMock) -> PluginRuntimeConnector:
|
||||
app = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
|
||||
deployment=SimpleNamespace(mode='cloud'),
|
||||
persistence_mgr=SimpleNamespace(
|
||||
tenant_uow=None,
|
||||
execute_async=execute_async,
|
||||
),
|
||||
logger=Mock(),
|
||||
)
|
||||
return PluginRuntimeConnector(app, AsyncMock())
|
||||
|
||||
|
||||
def execution_binding(workspace_uuid: str, generation: int = 1) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid=workspace_uuid,
|
||||
placement_generation=generation,
|
||||
)
|
||||
|
||||
|
||||
def plugin_setting(
|
||||
workspace_suffix: str,
|
||||
artifact_digest: str,
|
||||
*,
|
||||
durable: bool = True,
|
||||
) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
plugin_author='author',
|
||||
plugin_name=f'plugin-{workspace_suffix}',
|
||||
installation_uuid=f'00000000-0000-4000-8000-0000000000{workspace_suffix}',
|
||||
runtime_revision=1,
|
||||
artifact_digest=artifact_digest,
|
||||
enabled=True,
|
||||
priority=0,
|
||||
created_at=datetime.datetime(2026, 1, 1),
|
||||
install_source='local',
|
||||
install_info={'_artifact_storage': 'tenant_binary_storage_v1'} if durable else {},
|
||||
)
|
||||
|
||||
|
||||
def runtime_handler(
|
||||
*,
|
||||
missing_artifacts: list[str] | None = None,
|
||||
failed_installations: list[dict[str, str]] | None = None,
|
||||
) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
register_installation_binding=Mock(),
|
||||
unregister_installation_binding=Mock(),
|
||||
reconcile_plugin_installations=AsyncMock(
|
||||
return_value={
|
||||
'applied': [],
|
||||
'removed': [],
|
||||
'missing_artifacts': missing_artifacts or [],
|
||||
'failed_installations': failed_installations or [],
|
||||
}
|
||||
),
|
||||
apply_plugin_installation=AsyncMock(return_value={'state': 'starting'}),
|
||||
installation_scope=Mock(side_effect=lambda _binding: nullcontext()),
|
||||
list_plugins=AsyncMock(return_value=[]),
|
||||
)
|
||||
|
||||
|
||||
def shared_connector(
|
||||
projected_bindings: list[list[SimpleNamespace]],
|
||||
settings: dict[str, list[SimpleNamespace]],
|
||||
) -> PluginRuntimeConnector:
|
||||
async def get_execution_binding(workspace_uuid: str, *, expected_generation: int | None = None):
|
||||
for binding_set in projected_bindings:
|
||||
for binding in binding_set:
|
||||
if binding.workspace_uuid == workspace_uuid:
|
||||
assert expected_generation in (None, binding.placement_generation)
|
||||
return binding
|
||||
raise AssertionError(f'unexpected Workspace {workspace_uuid}')
|
||||
|
||||
app = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
|
||||
deployment=SimpleNamespace(mode='cloud'),
|
||||
workspace_service=SimpleNamespace(
|
||||
list_active_execution_bindings=AsyncMock(side_effect=projected_bindings),
|
||||
get_execution_binding=AsyncMock(side_effect=get_execution_binding),
|
||||
),
|
||||
persistence_mgr=SimpleNamespace(tenant_uow=None),
|
||||
logger=Mock(),
|
||||
)
|
||||
connector = PluginRuntimeConnector(app, AsyncMock())
|
||||
connector._load_workspace_settings = AsyncMock(side_effect=lambda context: settings[context.workspace_uuid])
|
||||
return connector
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shared_reconnect_replays_two_workspaces_and_removes_missing_projection():
|
||||
binding_a = execution_binding('workspace-a')
|
||||
binding_b = execution_binding('workspace-b')
|
||||
setting_a = plugin_setting('01', 'a' * 64)
|
||||
setting_b = plugin_setting('02', 'b' * 64)
|
||||
connector = shared_connector(
|
||||
[[binding_a, binding_b], [binding_a]],
|
||||
{'workspace-a': [setting_a], 'workspace-b': [setting_b]},
|
||||
)
|
||||
|
||||
first_handler = runtime_handler()
|
||||
connector.handler = first_handler
|
||||
await connector._prepare_connected_runtime()
|
||||
|
||||
first_desired = first_handler.reconcile_plugin_installations.await_args.args[0]
|
||||
assert {state.binding.workspace_uuid for state in first_desired} == {'workspace-a', 'workspace-b'}
|
||||
|
||||
second_handler = runtime_handler()
|
||||
connector.handler = second_handler
|
||||
await connector._prepare_connected_runtime()
|
||||
|
||||
second_desired = second_handler.reconcile_plugin_installations.await_args.args[0]
|
||||
assert [state.binding.workspace_uuid for state in second_desired] == ['workspace-a']
|
||||
second_handler.unregister_installation_binding.assert_called_once_with(first_desired[1].binding)
|
||||
assert set(connector._workspace_installations) == {'workspace-a'}
|
||||
assert set(connector._known_desired_states) == {setting_a.installation_uuid}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fresh_shared_runtime_cache_replays_persisted_local_package():
|
||||
package = b'local-lbpkg-bytes'
|
||||
digest = hashlib.sha256(package).hexdigest()
|
||||
binding = execution_binding('workspace-a')
|
||||
setting = plugin_setting('01', digest)
|
||||
connector = shared_connector([[binding]], {'workspace-a': [setting]})
|
||||
connector.handler = runtime_handler(missing_artifacts=[setting.installation_uuid])
|
||||
connector._load_artifact_package = AsyncMock(return_value=package)
|
||||
|
||||
await connector._prepare_connected_runtime()
|
||||
|
||||
desired = connector.handler.reconcile_plugin_installations.await_args.args[0][0]
|
||||
connector._load_artifact_package.assert_awaited_once()
|
||||
connector.handler.apply_plugin_installation.assert_awaited_once_with(
|
||||
desired.binding,
|
||||
artifact_package=package,
|
||||
enabled=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oss_upgrade_keeps_legacy_data_plugins_when_no_lbpkg_was_backfilled():
|
||||
binding = execution_binding('workspace-a')
|
||||
setting = plugin_setting('01', hashlib.sha256(b'legacy-installation').hexdigest(), durable=False)
|
||||
legacy_plugin = {
|
||||
'debug': False,
|
||||
'manifest': {'manifest': {'metadata': {'author': 'author', 'name': 'plugin-01'}}},
|
||||
'components': [],
|
||||
}
|
||||
app = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
|
||||
deployment=SimpleNamespace(mode='oss'),
|
||||
workspace_service=SimpleNamespace(get_local_execution_binding=AsyncMock(return_value=binding)),
|
||||
persistence_mgr=SimpleNamespace(tenant_uow=None),
|
||||
logger=Mock(),
|
||||
)
|
||||
connector = PluginRuntimeConnector(app, AsyncMock())
|
||||
connector.handler = runtime_handler(missing_artifacts=[setting.installation_uuid])
|
||||
connector.handler.list_plugins = AsyncMock(return_value=[legacy_plugin])
|
||||
connector.handler.apply_plugin_installation = AsyncMock(return_value={'state': 'artifact_missing'})
|
||||
connector._load_workspace_settings = AsyncMock(return_value=[setting])
|
||||
connector._load_artifact_package = AsyncMock(return_value=None)
|
||||
|
||||
await connector._prepare_connected_runtime()
|
||||
plugins = await connector.list_plugins()
|
||||
|
||||
assert plugins == [legacy_plugin]
|
||||
assert connector.handler.list_plugins.await_count >= 2
|
||||
connector._load_artifact_package.assert_awaited_once()
|
||||
assert not hasattr(connector.handler, 'delete_plugin')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_install_persists_verified_package_before_runtime_apply():
|
||||
package = b'local-lbpkg-bytes'
|
||||
digest = hashlib.sha256(package).hexdigest()
|
||||
execution_context = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
)
|
||||
binding = InstallationBinding(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
installation_uuid='00000000-0000-4000-8000-000000000001',
|
||||
runtime_revision=1,
|
||||
artifact_digest=digest,
|
||||
)
|
||||
app = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
|
||||
deployment=SimpleNamespace(mode='cloud'),
|
||||
logger=Mock(),
|
||||
)
|
||||
connector = PluginRuntimeConnector(app, AsyncMock())
|
||||
connector.handler = runtime_handler()
|
||||
connector._current_execution_context = AsyncMock(return_value=execution_context)
|
||||
connector._inspect_plugin_package = Mock(return_value=('author', 'plugin'))
|
||||
connector._store_artifact_package = AsyncMock()
|
||||
connector._persist_installation_package = AsyncMock(return_value=(binding, None, False))
|
||||
connector._wait_for_installed_plugin_ready = AsyncMock()
|
||||
|
||||
await connector.install_plugin(
|
||||
PluginInstallSource.LOCAL,
|
||||
{'plugin_file': package},
|
||||
)
|
||||
|
||||
connector._store_artifact_package.assert_awaited_once_with(execution_context, digest, package)
|
||||
connector.handler.apply_plugin_installation.assert_awaited_once_with(
|
||||
binding,
|
||||
artifact_package=package,
|
||||
enabled=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(('remaining_references', 'statement_count'), [(1, 1), (0, 2)])
|
||||
async def test_artifact_cleanup_is_reference_counted_within_workspace(
|
||||
remaining_references: int,
|
||||
statement_count: int,
|
||||
):
|
||||
app = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
|
||||
deployment=SimpleNamespace(mode='cloud'),
|
||||
)
|
||||
connector = PluginRuntimeConnector(app, AsyncMock())
|
||||
execution_context = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
)
|
||||
statements = []
|
||||
|
||||
async def execute(statement):
|
||||
statements.append(statement)
|
||||
return SimpleNamespace(scalar_one=lambda: remaining_references)
|
||||
|
||||
await connector._delete_artifact_if_unreferenced(
|
||||
execution_context,
|
||||
'a' * 64,
|
||||
execute=execute,
|
||||
)
|
||||
|
||||
assert len(statements) == statement_count
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_artifact_store_and_load_accept_connection_scalar_results():
|
||||
package = b'persisted-lbpkg'
|
||||
digest = hashlib.sha256(package).hexdigest()
|
||||
execution_context = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
)
|
||||
execute_async = AsyncMock(return_value=SimpleNamespace(scalar_one_or_none=lambda: package))
|
||||
connector = connection_result_connector(execute_async)
|
||||
|
||||
await connector._store_artifact_package(execution_context, digest, package)
|
||||
loaded = await connector._load_artifact_package(execution_context, digest)
|
||||
|
||||
assert loaded == package
|
||||
assert execute_async.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_workspace_settings_accept_connection_mapping_results():
|
||||
created_at = datetime.datetime(2026, 1, 1)
|
||||
row = {
|
||||
'workspace_uuid': 'workspace-a',
|
||||
'plugin_author': 'author',
|
||||
'plugin_name': 'plugin',
|
||||
'installation_uuid': '00000000-0000-4000-8000-000000000001',
|
||||
'artifact_digest': 'a' * 64,
|
||||
'runtime_revision': 2,
|
||||
'enabled': True,
|
||||
'priority': 3,
|
||||
'config': {'key': 'value'},
|
||||
'install_source': 'local',
|
||||
'install_info': {'_artifact_storage': 'tenant_binary_storage_v1'},
|
||||
'created_at': created_at,
|
||||
'updated_at': created_at,
|
||||
}
|
||||
mapped_result = SimpleNamespace(mappings=lambda: SimpleNamespace(all=lambda: [row]))
|
||||
connector = connection_result_connector(AsyncMock(return_value=mapped_result))
|
||||
|
||||
settings = await connector._load_workspace_settings(
|
||||
ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
)
|
||||
)
|
||||
|
||||
assert len(settings) == 1
|
||||
assert settings[0].installation_uuid == row['installation_uuid']
|
||||
assert settings[0].runtime_revision == 2
|
||||
assert settings[0].created_at == created_at
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_installation_update_accepts_connection_row_result():
|
||||
old_digest = 'a' * 64
|
||||
new_digest = 'b' * 64
|
||||
existing = SimpleNamespace(
|
||||
installation_uuid='00000000-0000-4000-8000-000000000001',
|
||||
runtime_revision=4,
|
||||
artifact_digest=old_digest,
|
||||
install_info={'_artifact_storage': 'tenant_binary_storage_v1'},
|
||||
)
|
||||
execute_async = AsyncMock(
|
||||
side_effect=[
|
||||
SimpleNamespace(first=lambda: existing),
|
||||
SimpleNamespace(rowcount=1),
|
||||
]
|
||||
)
|
||||
connector = connection_result_connector(execute_async)
|
||||
execution_context = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=7,
|
||||
)
|
||||
|
||||
binding, previous_digest, previous_was_durable = await connector._persist_installation_package(
|
||||
execution_context,
|
||||
plugin_author='author',
|
||||
plugin_name='plugin',
|
||||
install_source=PluginInstallSource.LOCAL,
|
||||
install_info={},
|
||||
artifact_digest=new_digest,
|
||||
)
|
||||
|
||||
assert binding.installation_uuid == existing.installation_uuid
|
||||
assert binding.runtime_revision == 5
|
||||
assert binding.artifact_digest == new_digest
|
||||
assert previous_digest == old_digest
|
||||
assert previous_was_durable is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_dependency_failure_raises_stable_observable_error():
|
||||
binding = execution_binding('workspace-a')
|
||||
setting = plugin_setting('01', 'a' * 64)
|
||||
connector = shared_connector([[binding]], {'workspace-a': [setting]})
|
||||
connector.handler = runtime_handler()
|
||||
desired = (
|
||||
await connector._load_workspace_desired_states(
|
||||
ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
)
|
||||
)
|
||||
)[0]
|
||||
connector.handler.apply_plugin_installation.return_value = {
|
||||
'installation_uuid': setting.installation_uuid,
|
||||
'state': 'failed',
|
||||
'error_code': 'dependency_prepare_failed',
|
||||
'message': 'Plugin dependency installer exited with code 1',
|
||||
}
|
||||
|
||||
with pytest.raises(PluginInstallationFailedError) as exc_info:
|
||||
await connector._apply_desired_state(desired)
|
||||
|
||||
error = exc_info.value
|
||||
assert error.installation_uuid == setting.installation_uuid
|
||||
assert error.error_code == 'dependency_prepare_failed'
|
||||
assert '[dependency_prepare_failed]' in str(error)
|
||||
assert connector._installation_failures[setting.installation_uuid] == {
|
||||
'installation_uuid': setting.installation_uuid,
|
||||
'error_code': 'dependency_prepare_failed',
|
||||
'message': 'Plugin dependency installer exited with code 1',
|
||||
}
|
||||
connector.ap.logger.error.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shared_reconcile_records_one_failure_without_blocking_other_state():
|
||||
binding_a = execution_binding('workspace-a')
|
||||
binding_b = execution_binding('workspace-b')
|
||||
setting_a = plugin_setting('01', 'a' * 64)
|
||||
setting_b = plugin_setting('02', 'b' * 64)
|
||||
failure = {
|
||||
'installation_uuid': setting_a.installation_uuid,
|
||||
'error_code': 'dependency_prepare_failed',
|
||||
'message': 'Plugin dependency installer exited with code 1',
|
||||
}
|
||||
connector = shared_connector(
|
||||
[[binding_a, binding_b]],
|
||||
{'workspace-a': [setting_a], 'workspace-b': [setting_b]},
|
||||
)
|
||||
connector.handler = runtime_handler(failed_installations=[failure])
|
||||
|
||||
await connector._prepare_connected_runtime()
|
||||
|
||||
desired = connector.handler.reconcile_plugin_installations.await_args.args[0]
|
||||
assert {item.binding.installation_uuid for item in desired} == {
|
||||
setting_a.installation_uuid,
|
||||
setting_b.installation_uuid,
|
||||
}
|
||||
assert connector._installation_failures == {
|
||||
setting_a.installation_uuid: failure,
|
||||
}
|
||||
assert set(connector._known_desired_states) == {
|
||||
setting_a.installation_uuid,
|
||||
setting_b.installation_uuid,
|
||||
}
|
||||
connector.ap.logger.error.assert_called_once_with(
|
||||
'Plugin installation %s failed during reconcile [%s]: %s',
|
||||
setting_a.installation_uuid,
|
||||
'dependency_prepare_failed',
|
||||
'Plugin dependency installer exited with code 1',
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_artifact_repair_adds_dependency_failure_and_continues():
|
||||
package_a = b'package-a'
|
||||
package_b = b'package-b'
|
||||
binding = execution_binding('workspace-a')
|
||||
setting_a = plugin_setting('01', hashlib.sha256(package_a).hexdigest())
|
||||
setting_b = plugin_setting('02', hashlib.sha256(package_b).hexdigest())
|
||||
connector = shared_connector(
|
||||
[[binding]],
|
||||
{'workspace-a': [setting_a, setting_b]},
|
||||
)
|
||||
connector.handler = runtime_handler(
|
||||
missing_artifacts=[
|
||||
setting_a.installation_uuid,
|
||||
setting_b.installation_uuid,
|
||||
]
|
||||
)
|
||||
connector._load_artifact_package = AsyncMock(side_effect=[package_a, package_b])
|
||||
connector.handler.apply_plugin_installation.side_effect = [
|
||||
{
|
||||
'installation_uuid': setting_a.installation_uuid,
|
||||
'state': 'failed',
|
||||
'error_code': 'dependency_prepare_failed',
|
||||
'message': 'Plugin dependency installer exited with code 1',
|
||||
},
|
||||
{'installation_uuid': setting_b.installation_uuid, 'state': 'starting'},
|
||||
]
|
||||
|
||||
result = await connector.reconcile_projected_workspaces(
|
||||
[
|
||||
ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert connector.handler.apply_plugin_installation.await_count == 2
|
||||
assert result['failed_installations'] == [
|
||||
{
|
||||
'installation_uuid': setting_a.installation_uuid,
|
||||
'error_code': 'dependency_prepare_failed',
|
||||
'message': 'Plugin dependency installer exited with code 1',
|
||||
}
|
||||
]
|
||||
assert setting_a.installation_uuid in connector._installation_failures
|
||||
assert setting_b.installation_uuid not in connector._installation_failures
|
||||
@@ -0,0 +1,174 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.plugin import connector as connector_module
|
||||
|
||||
from .test_connector_methods import create_mock_connector
|
||||
|
||||
|
||||
TRUSTED_LEGACY_REQUEST = {
|
||||
'owner': 'langbot-app',
|
||||
'repo': 'demo-plugin',
|
||||
'release_tag': 'v1.0.0',
|
||||
'asset_url': 'https://github.com/langbot-app/demo-plugin/releases/download/v1.0.0/demo.lbpkg',
|
||||
}
|
||||
|
||||
|
||||
def _patch_client(monkeypatch, handler):
|
||||
real_async_client = httpx.AsyncClient
|
||||
observed: list[dict] = []
|
||||
|
||||
def client_factory(*args, **kwargs):
|
||||
observed.append(dict(kwargs))
|
||||
return real_async_client(
|
||||
transport=httpx.MockTransport(handler),
|
||||
follow_redirects=kwargs.get('follow_redirects', False),
|
||||
trust_env=kwargs.get('trust_env', True),
|
||||
timeout=kwargs.get('timeout'),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(connector_module.httpx, 'AsyncClient', client_factory)
|
||||
return observed
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'asset_url',
|
||||
[
|
||||
'http://127.0.0.1/internal.lbpkg',
|
||||
'https://169.254.169.254/latest/meta-data',
|
||||
'https://github.com@127.0.0.1/internal.lbpkg',
|
||||
'https://evil.example/langbot-app/demo-plugin/releases/download/v1.0.0/demo.lbpkg',
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_github_install_rejects_internal_or_untrusted_asset_url_before_network(
|
||||
monkeypatch,
|
||||
asset_url,
|
||||
):
|
||||
connector = create_mock_connector()
|
||||
monkeypatch.setattr(
|
||||
connector_module.httpx,
|
||||
'AsyncClient',
|
||||
lambda *_args, **_kwargs: pytest.fail('network client must not be created'),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match='GitHub release asset URL'):
|
||||
await connector._download_github_package(
|
||||
{**TRUSTED_LEGACY_REQUEST, 'asset_url': asset_url},
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_github_asset_id_is_resolved_server_side_and_redirect_escape_is_rejected(
|
||||
monkeypatch,
|
||||
):
|
||||
connector = create_mock_connector()
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
if request.url.path.endswith('/releases/42'):
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
'id': 42,
|
||||
'tag_name': 'v1.0.0',
|
||||
'assets': [{'id': 99, 'size': 128, 'state': 'uploaded'}],
|
||||
},
|
||||
)
|
||||
if request.url.path.endswith('/releases/assets/99'):
|
||||
return httpx.Response(
|
||||
302,
|
||||
headers={'location': 'http://169.254.169.254/latest/meta-data'},
|
||||
)
|
||||
raise AssertionError(f'unexpected request: {request.url}')
|
||||
|
||||
observed = _patch_client(monkeypatch, handler)
|
||||
with pytest.raises(ValueError, match='untrusted host'):
|
||||
await connector._download_github_package(
|
||||
{
|
||||
'owner': 'langbot-app',
|
||||
'repo': 'demo-plugin',
|
||||
'release_tag': 'v1.0.0',
|
||||
'release_id': 42,
|
||||
'asset_id': 99,
|
||||
'asset_url': 'https://attacker.invalid/ignored',
|
||||
},
|
||||
None,
|
||||
)
|
||||
|
||||
assert len(observed) == 1
|
||||
assert observed[0]['trust_env'] is False
|
||||
assert observed[0]['follow_redirects'] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_github_download_rejects_oversized_content_length(monkeypatch):
|
||||
connector = create_mock_connector()
|
||||
monkeypatch.setattr(connector_module, '_GITHUB_PLUGIN_DOWNLOAD_MAX_BYTES', 8)
|
||||
|
||||
def handler(_request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, headers={'content-length': '9'}, content=b'')
|
||||
|
||||
_patch_client(monkeypatch, handler)
|
||||
with pytest.raises(ValueError, match='10 MiB download limit'):
|
||||
await connector._download_github_package(TRUSTED_LEGACY_REQUEST, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_github_download_counts_chunked_stream_bytes(monkeypatch):
|
||||
connector = create_mock_connector()
|
||||
monkeypatch.setattr(connector_module, '_GITHUB_PLUGIN_DOWNLOAD_MAX_BYTES', 8)
|
||||
|
||||
class ChunkedBody(httpx.AsyncByteStream):
|
||||
async def __aiter__(self):
|
||||
yield b'1234'
|
||||
yield b'56789'
|
||||
|
||||
def handler(_request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, stream=ChunkedBody())
|
||||
|
||||
_patch_client(monkeypatch, handler)
|
||||
with pytest.raises(ValueError, match='10 MiB download limit'):
|
||||
await connector._download_github_package(TRUSTED_LEGACY_REQUEST, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_github_asset_id_download_allows_github_object_redirect(monkeypatch):
|
||||
connector = create_mock_connector()
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
if request.url.path.endswith('/releases/42'):
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
'id': 42,
|
||||
'tag_name': 'v1.0.0',
|
||||
'assets': [{'id': 99, 'size': 7, 'state': 'uploaded'}],
|
||||
},
|
||||
)
|
||||
if request.url.path.endswith('/releases/assets/99'):
|
||||
return httpx.Response(
|
||||
302,
|
||||
headers={
|
||||
'location': 'https://release-assets.githubusercontent.com/github-production-release-asset/demo'
|
||||
},
|
||||
)
|
||||
if request.url.host == 'release-assets.githubusercontent.com':
|
||||
return httpx.Response(200, content=b'package')
|
||||
raise AssertionError(f'unexpected request: {request.url}')
|
||||
|
||||
_patch_client(monkeypatch, handler)
|
||||
package = await connector._download_github_package(
|
||||
{
|
||||
'owner': 'langbot-app',
|
||||
'repo': 'demo-plugin',
|
||||
'release_tag': 'v1.0.0',
|
||||
'release_id': 42,
|
||||
'asset_id': 99,
|
||||
},
|
||||
None,
|
||||
)
|
||||
|
||||
assert package == b'package'
|
||||
@@ -10,7 +10,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock
|
||||
import pytest
|
||||
|
||||
from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction
|
||||
from langbot_plugin.entities.io.context import ActionContext
|
||||
from langbot_plugin.entities.io.context import ActionContext, InstallationBinding
|
||||
|
||||
|
||||
def make_handler(app):
|
||||
@@ -35,14 +35,19 @@ def make_handler(app):
|
||||
Mock(),
|
||||
AsyncMock(return_value=True),
|
||||
app,
|
||||
workspace_context,
|
||||
)
|
||||
installation_uuid = runtime_handler._remember_installation(
|
||||
workspace_context,
|
||||
'test-author',
|
||||
'test-plugin',
|
||||
installation_binding = InstallationBinding(
|
||||
**workspace_context.model_dump(exclude_none=True),
|
||||
installation_uuid='00000000-0000-4000-8000-000000000001',
|
||||
runtime_revision=1,
|
||||
artifact_digest='a' * 64,
|
||||
)
|
||||
runtime_handler.bind_action_context(workspace_context.for_installation(installation_uuid))
|
||||
runtime_handler.register_installation_binding(
|
||||
installation_binding,
|
||||
plugin_author='test-author',
|
||||
plugin_name='test-plugin',
|
||||
)
|
||||
runtime_handler._current_action_context.set(installation_binding)
|
||||
query_pool = getattr(app, 'query_pool', None)
|
||||
if query_pool is not None and hasattr(query_pool, 'cached_queries'):
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@ from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction, RuntimeToLangBotAction
|
||||
from langbot_plugin.entities.io.context import ActionContext
|
||||
from langbot_plugin.entities.io.context import ActionContext, InstallationBinding
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.storage.mgr import StorageMgr
|
||||
@@ -52,14 +52,19 @@ def make_handler(app, workspace_context: ActionContext | None = None):
|
||||
Mock(),
|
||||
AsyncMock(return_value=True),
|
||||
app,
|
||||
workspace_context,
|
||||
)
|
||||
installation_uuid = runtime_handler._remember_installation(
|
||||
workspace_context,
|
||||
'test-author',
|
||||
'test-plugin',
|
||||
installation_binding = InstallationBinding(
|
||||
**workspace_context.model_dump(exclude_none=True),
|
||||
installation_uuid='00000000-0000-4000-8000-000000000001',
|
||||
runtime_revision=1,
|
||||
artifact_digest='a' * 64,
|
||||
)
|
||||
runtime_handler.bind_action_context(workspace_context.for_installation(installation_uuid))
|
||||
runtime_handler.register_installation_binding(
|
||||
installation_binding,
|
||||
plugin_author='test-author',
|
||||
plugin_name='test-plugin',
|
||||
)
|
||||
runtime_handler._current_action_context.set(installation_binding)
|
||||
query_pool = getattr(app, 'query_pool', None)
|
||||
if query_pool is not None and hasattr(query_pool, 'cached_queries'):
|
||||
|
||||
@@ -169,13 +174,10 @@ class TestInitializePluginSettings:
|
||||
return mock_app
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_creates_new_setting_when_not_exists(self, app):
|
||||
"""New plugin settings use default enabled, priority and config values."""
|
||||
async def test_rejects_desired_installation_when_setting_not_exists(self, app):
|
||||
"""A desired-state worker cannot create an unowned Core setting row."""
|
||||
runtime_handler = make_handler(app)
|
||||
app.persistence_mgr.execute_async.side_effect = [
|
||||
make_result(),
|
||||
Mock(),
|
||||
]
|
||||
app.persistence_mgr.execute_async.return_value = make_result()
|
||||
|
||||
response = await runtime_handler.actions[RuntimeToLangBotAction.INITIALIZE_PLUGIN_SETTINGS.value](
|
||||
{
|
||||
@@ -186,34 +188,23 @@ class TestInitializePluginSettings:
|
||||
}
|
||||
)
|
||||
|
||||
assert response.code == 0
|
||||
assert app.persistence_mgr.execute_async.await_count == 2
|
||||
insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[1].args[0])
|
||||
assert insert_params == {
|
||||
'workspace_uuid': 'workspace-a',
|
||||
'plugin_author': 'test-author',
|
||||
'plugin_name': 'test-plugin',
|
||||
'install_source': 'local',
|
||||
'install_info': {'path': '/test'},
|
||||
'enabled': True,
|
||||
'priority': 0,
|
||||
'config': {},
|
||||
}
|
||||
assert response.code != 0
|
||||
assert 'Plugin installation setting was not found' in response.message
|
||||
app.persistence_mgr.execute_async.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inherits_values_from_existing_setting(self, app):
|
||||
"""Existing settings are replaced while preserving user-controlled values."""
|
||||
async def test_existing_desired_setting_remains_core_owned(self, app):
|
||||
"""Runtime initialization validates identity without rewriting Core state."""
|
||||
runtime_handler = make_handler(app)
|
||||
existing_setting = SimpleNamespace(
|
||||
enabled=False,
|
||||
priority=5,
|
||||
config={'key': 'value'},
|
||||
installation_uuid='00000000-0000-4000-8000-000000000001',
|
||||
runtime_revision=1,
|
||||
artifact_digest='a' * 64,
|
||||
)
|
||||
app.persistence_mgr.execute_async.side_effect = [
|
||||
make_result(existing_setting),
|
||||
Mock(),
|
||||
Mock(),
|
||||
]
|
||||
app.persistence_mgr.execute_async.return_value = make_result(existing_setting)
|
||||
|
||||
response = await runtime_handler.actions[RuntimeToLangBotAction.INITIALIZE_PLUGIN_SETTINGS.value](
|
||||
{
|
||||
@@ -225,13 +216,7 @@ class TestInitializePluginSettings:
|
||||
)
|
||||
|
||||
assert response.code == 0
|
||||
assert app.persistence_mgr.execute_async.await_count == 3
|
||||
insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[2].args[0])
|
||||
assert insert_params['enabled'] is False
|
||||
assert insert_params['priority'] == 5
|
||||
assert insert_params['config'] == {'key': 'value'}
|
||||
assert insert_params['install_source'] == 'github'
|
||||
assert insert_params['install_info'] == {'repo': 'author/name'}
|
||||
app.persistence_mgr.execute_async.assert_awaited_once()
|
||||
|
||||
|
||||
class TestSetBinaryStorage:
|
||||
@@ -367,31 +352,18 @@ class TestGetPluginSettings:
|
||||
return mock_app
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_defaults_when_setting_not_found(self, app):
|
||||
"""Default plugin settings are returned when no persisted row exists."""
|
||||
async def test_rejects_desired_installation_when_setting_not_found(self, app):
|
||||
"""A desired-state worker cannot synthesize settings for a missing row."""
|
||||
runtime_handler = make_handler(app)
|
||||
app.persistence_mgr.execute_async.return_value = make_result()
|
||||
|
||||
response = await runtime_handler.actions[RuntimeToLangBotAction.GET_PLUGIN_SETTINGS.value](
|
||||
{
|
||||
'plugin_author': 'test-author',
|
||||
'plugin_name': 'test-plugin',
|
||||
}
|
||||
)
|
||||
|
||||
assert response.code == 0
|
||||
assert response.data == {
|
||||
'enabled': True,
|
||||
'priority': 0,
|
||||
'plugin_config': {},
|
||||
'install_source': 'local',
|
||||
'install_info': {},
|
||||
'installation_uuid': runtime_handler.derive_installation_uuid(
|
||||
runtime_handler.require_bound_action_context(),
|
||||
'test-author',
|
||||
'test-plugin',
|
||||
),
|
||||
}
|
||||
with pytest.raises(ValueError, match='Plugin installation setting was not found'):
|
||||
await runtime_handler.actions[RuntimeToLangBotAction.GET_PLUGIN_SETTINGS.value](
|
||||
{
|
||||
'plugin_author': 'test-author',
|
||||
'plugin_name': 'test-plugin',
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_actual_values_when_setting_exists(self, app):
|
||||
@@ -403,6 +375,9 @@ class TestGetPluginSettings:
|
||||
config={'custom': 'config'},
|
||||
install_source='github',
|
||||
install_info={'repo': 'test/repo'},
|
||||
installation_uuid='00000000-0000-4000-8000-000000000001',
|
||||
runtime_revision=1,
|
||||
artifact_digest='a' * 64,
|
||||
)
|
||||
app.persistence_mgr.execute_async.return_value = make_result(setting)
|
||||
|
||||
@@ -420,11 +395,9 @@ class TestGetPluginSettings:
|
||||
'plugin_config': {'custom': 'config'},
|
||||
'install_source': 'github',
|
||||
'install_info': {'repo': 'test/repo'},
|
||||
'installation_uuid': runtime_handler.derive_installation_uuid(
|
||||
runtime_handler.require_bound_action_context(),
|
||||
'test-author',
|
||||
'test-plugin',
|
||||
),
|
||||
'installation_uuid': '00000000-0000-4000-8000-000000000001',
|
||||
'runtime_revision': 1,
|
||||
'artifact_digest': 'a' * 64,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -1,19 +1,22 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction
|
||||
from langbot_plugin.entities.io.context import ActionContext
|
||||
from langbot_plugin.entities.io.context import ActionContext, InstallationBinding, PluginWorkerPolicy, RuntimeIdentity
|
||||
from langbot_plugin.entities.io.resp import ActionResponse
|
||||
from langbot_plugin.runtime.io.connection import Connection
|
||||
|
||||
from langbot.pkg.plugin.handler import RuntimeConnectionHandler
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
|
||||
|
||||
|
||||
class EmptyResult:
|
||||
@@ -49,6 +52,7 @@ def workspace_context(workspace_uuid: str = 'workspace-a') -> ActionContext:
|
||||
def make_handler(workspace_uuid: str = 'workspace-a'):
|
||||
context = workspace_context(workspace_uuid)
|
||||
app = SimpleNamespace(
|
||||
deployment=SimpleNamespace(mode='cloud'),
|
||||
persistence_mgr=SimpleNamespace(execute_async=AsyncMock(return_value=EmptyResult())),
|
||||
logger=Mock(),
|
||||
workspace_service=SimpleNamespace(
|
||||
@@ -65,14 +69,19 @@ def make_handler(workspace_uuid: str = 'workspace-a'):
|
||||
Mock(),
|
||||
AsyncMock(return_value=True),
|
||||
app,
|
||||
context,
|
||||
)
|
||||
installation_uuid = runtime_handler._remember_installation(
|
||||
context,
|
||||
'author-a',
|
||||
'plugin-a',
|
||||
installation_context = InstallationBinding(
|
||||
**context.model_dump(exclude_none=True),
|
||||
installation_uuid='00000000-0000-4000-8000-000000000001',
|
||||
runtime_revision=1,
|
||||
artifact_digest='a' * 64,
|
||||
)
|
||||
return runtime_handler, app, context.for_installation(installation_uuid)
|
||||
runtime_handler.register_installation_binding(
|
||||
installation_context,
|
||||
plugin_author='author-a',
|
||||
plugin_name='plugin-a',
|
||||
)
|
||||
return runtime_handler, app, installation_context
|
||||
|
||||
|
||||
async def invoke_with_context(
|
||||
@@ -92,10 +101,98 @@ async def invoke_with_context(
|
||||
async def test_plugin_action_requires_installation_capability():
|
||||
runtime_handler, _app, _installation_context = make_handler()
|
||||
|
||||
with pytest.raises(ValueError, match='missing installation_uuid'):
|
||||
with pytest.raises(ValueError, match='trusted Workspace context'):
|
||||
await runtime_handler.actions[PluginToRuntimeAction.GET_LANGBOT_VERSION.value]({})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_action_enters_trusted_workspace_scope():
|
||||
runtime_handler, app, installation_context = make_handler()
|
||||
scope_events: list[tuple[str, str]] = []
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def tenant_scope(workspace_uuid: str):
|
||||
scope_events.append(('enter', workspace_uuid))
|
||||
try:
|
||||
yield SimpleNamespace()
|
||||
finally:
|
||||
scope_events.append(('exit', workspace_uuid))
|
||||
|
||||
class RecordingPersistenceManager:
|
||||
def __init__(self, execute_async):
|
||||
self.execute_async = execute_async
|
||||
|
||||
def tenant_scope(self, workspace_uuid: str):
|
||||
return tenant_scope(workspace_uuid)
|
||||
|
||||
app.persistence_mgr = RecordingPersistenceManager(app.persistence_mgr.execute_async)
|
||||
|
||||
response = await invoke_with_context(
|
||||
runtime_handler,
|
||||
installation_context,
|
||||
PluginToRuntimeAction.GET_LANGBOT_VERSION,
|
||||
{},
|
||||
)
|
||||
|
||||
assert response.code == 0
|
||||
assert scope_events == [('enter', 'workspace-a'), ('exit', 'workspace-a')]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_blocked_llm_provider_does_not_hold_tenant_database_session():
|
||||
entered = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
observations: list[bool] = []
|
||||
engine = create_async_engine('sqlite+aiosqlite:///:memory:')
|
||||
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||
manager.db = SimpleNamespace(get_engine=lambda: engine)
|
||||
runtime_handler, app, installation_context = make_handler()
|
||||
app.persistence_mgr = manager
|
||||
|
||||
async def invoke_llm(**_kwargs):
|
||||
observations.append(manager.current_session() is None)
|
||||
entered.set()
|
||||
await release.wait()
|
||||
observations.append(manager.current_session() is None)
|
||||
return SimpleNamespace(model_dump=lambda: {'role': 'assistant', 'content': 'ok'})
|
||||
|
||||
app.model_mgr = SimpleNamespace(
|
||||
get_model_by_uuid=AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
model_entity=SimpleNamespace(workspace_uuid='workspace-a'),
|
||||
provider=SimpleNamespace(invoke_llm=invoke_llm),
|
||||
)
|
||||
)
|
||||
)
|
||||
runtime_handler._require_plugin_action_context = AsyncMock(return_value=(installation_context, SimpleNamespace()))
|
||||
runtime_handler._require_active_action_context = AsyncMock()
|
||||
runtime_handler._resource_exists = AsyncMock(return_value=True)
|
||||
|
||||
action = asyncio.create_task(
|
||||
invoke_with_context(
|
||||
runtime_handler,
|
||||
installation_context,
|
||||
PluginToRuntimeAction.INVOKE_LLM,
|
||||
{
|
||||
'llm_model_uuid': 'model-a',
|
||||
'messages': [],
|
||||
},
|
||||
)
|
||||
)
|
||||
try:
|
||||
await entered.wait()
|
||||
assert observations == [True]
|
||||
release.set()
|
||||
response = await action
|
||||
assert response.code == 0
|
||||
assert observations == [True, True]
|
||||
finally:
|
||||
release.set()
|
||||
if not action.done():
|
||||
await action
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plugin_action_rejects_forged_installation_capability():
|
||||
runtime_handler, _app, installation_context = make_handler()
|
||||
@@ -203,14 +300,12 @@ async def test_legacy_query_id_fallback_is_workspace_scoped():
|
||||
)
|
||||
|
||||
|
||||
def test_runtime_connection_rejects_workspace_rebinding():
|
||||
def test_runtime_connection_is_instance_scoped_and_unbound():
|
||||
runtime_handler, _app, _installation_context = make_handler()
|
||||
|
||||
with pytest.raises(ValueError, match='another Workspace'):
|
||||
runtime_handler.bind_action_context(workspace_context('workspace-b'))
|
||||
assert runtime_handler.bound_action_context is None
|
||||
|
||||
|
||||
def test_inbound_installation_envelope_is_preserved_only_inside_bound_workspace():
|
||||
def test_inbound_tenant_action_requires_complete_installation_envelope():
|
||||
runtime_handler, _app, installation_context = make_handler()
|
||||
|
||||
assert (
|
||||
@@ -220,35 +315,69 @@ def test_inbound_installation_envelope_is_preserved_only_inside_bound_workspace(
|
||||
)
|
||||
== installation_context
|
||||
)
|
||||
with pytest.raises(ValueError, match='does not match connection Workspace'):
|
||||
with pytest.raises(ValueError, match='complete InstallationBinding'):
|
||||
runtime_handler.validate_inbound_action_context(
|
||||
PluginToRuntimeAction.GET_BOTS.value,
|
||||
workspace_context('workspace-b').for_installation('installation-b'),
|
||||
)
|
||||
|
||||
|
||||
def test_installation_uuid_is_stable_and_workspace_specific():
|
||||
context_a = workspace_context('workspace-a')
|
||||
context_b = workspace_context('workspace-b')
|
||||
@pytest.mark.asyncio
|
||||
async def test_legacy_oss_worker_capability_remains_usable_after_identity_migration():
|
||||
runtime_handler, app, installation_context = make_handler()
|
||||
app.deployment = SimpleNamespace(mode='oss')
|
||||
setting = SimpleNamespace(
|
||||
plugin_author='author-a',
|
||||
plugin_name='plugin-a',
|
||||
installation_uuid=installation_context.installation_uuid,
|
||||
runtime_revision=installation_context.runtime_revision,
|
||||
artifact_digest=installation_context.artifact_digest,
|
||||
)
|
||||
result = Mock()
|
||||
result.first.return_value = setting
|
||||
app.persistence_mgr.execute_async.return_value = result
|
||||
legacy_context = workspace_context().for_installation(installation_context.installation_uuid)
|
||||
|
||||
first = RuntimeConnectionHandler.derive_installation_uuid(
|
||||
context_a,
|
||||
'author-a',
|
||||
'plugin-a',
|
||||
assert (
|
||||
runtime_handler.validate_inbound_action_context(
|
||||
PluginToRuntimeAction.GET_LANGBOT_VERSION.value,
|
||||
legacy_context,
|
||||
)
|
||||
== legacy_context
|
||||
)
|
||||
repeated = RuntimeConnectionHandler.derive_installation_uuid(
|
||||
context_a,
|
||||
'author-a',
|
||||
'plugin-a',
|
||||
)
|
||||
other_workspace = RuntimeConnectionHandler.derive_installation_uuid(
|
||||
context_b,
|
||||
'author-a',
|
||||
'plugin-a',
|
||||
response = await invoke_with_context(
|
||||
runtime_handler,
|
||||
legacy_context,
|
||||
PluginToRuntimeAction.GET_LANGBOT_VERSION,
|
||||
{},
|
||||
)
|
||||
|
||||
assert first == repeated
|
||||
assert first != other_workspace
|
||||
assert response.code == 0
|
||||
|
||||
|
||||
def test_installation_uuid_cannot_move_between_workspaces():
|
||||
runtime_handler, _app, binding = make_handler()
|
||||
moved = binding.model_copy(update={'workspace_uuid': 'workspace-b', 'runtime_revision': 2})
|
||||
|
||||
with pytest.raises(ValueError, match='cannot move between Workspaces'):
|
||||
runtime_handler.register_installation_binding(
|
||||
moved,
|
||||
plugin_author='author-a',
|
||||
plugin_name='plugin-a',
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_newer_revision_fences_old_binding():
|
||||
runtime_handler, _app, binding = make_handler()
|
||||
newer = binding.model_copy(update={'runtime_revision': 2, 'artifact_digest': 'b' * 64})
|
||||
runtime_handler.register_installation_binding(
|
||||
newer,
|
||||
plugin_author='author-a',
|
||||
plugin_name='plugin-a',
|
||||
)
|
||||
with pytest.raises(ValueError, match='revision or artifact is stale'):
|
||||
await runtime_handler._resolve_installation_identity(binding)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -281,16 +410,27 @@ async def test_plugin_vector_action_forwards_trusted_context_and_logical_collect
|
||||
@pytest.mark.asyncio
|
||||
async def test_host_to_runtime_action_carries_trusted_connector_context():
|
||||
app = SimpleNamespace(logger=Mock())
|
||||
context = workspace_context()
|
||||
connection = RecordingConnection()
|
||||
runtime_handler = RuntimeConnectionHandler(
|
||||
connection,
|
||||
AsyncMock(return_value=True),
|
||||
app,
|
||||
context,
|
||||
)
|
||||
|
||||
task = asyncio.create_task(runtime_handler.set_runtime_config(None))
|
||||
task = asyncio.create_task(
|
||||
runtime_handler.set_runtime_config(
|
||||
runtime_identity=RuntimeIdentity(instance_uuid='instance-a', runtime_id='runtime-a'),
|
||||
worker_policy=PluginWorkerPolicy(
|
||||
max_cpus=1.0,
|
||||
max_memory_mb=256,
|
||||
max_pids=32,
|
||||
max_open_files=64,
|
||||
max_file_size_mb=128,
|
||||
),
|
||||
runtime_profile='oss_dev',
|
||||
cloud_service_url=None,
|
||||
)
|
||||
)
|
||||
for _ in range(10):
|
||||
if connection.sent:
|
||||
break
|
||||
@@ -299,5 +439,8 @@ async def test_host_to_runtime_action_carries_trusted_connector_context():
|
||||
runtime_handler.resp_waiters[request['seq_id']].set_result(ActionResponse.success({}))
|
||||
await task
|
||||
|
||||
assert request['data'] == {}
|
||||
assert request['context'] == context.model_dump(exclude_none=True)
|
||||
assert request['data']['runtime_identity'] == {
|
||||
'instance_uuid': 'instance-a',
|
||||
'runtime_id': 'runtime-a',
|
||||
}
|
||||
assert request.get('context') is None
|
||||
|
||||
@@ -1,7 +1,34 @@
|
||||
"""Test plugin list filtering by component kinds."""
|
||||
|
||||
from contextlib import nullcontext
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from langbot_plugin.entities.io.context import InstallationBinding
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
|
||||
|
||||
TEST_EXECUTION_CONTEXT = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
)
|
||||
TEST_INSTALLATION_BINDING = InstallationBinding(
|
||||
instance_uuid=TEST_EXECUTION_CONTEXT.instance_uuid,
|
||||
workspace_uuid=TEST_EXECUTION_CONTEXT.workspace_uuid,
|
||||
placement_generation=TEST_EXECUTION_CONTEXT.placement_generation,
|
||||
installation_uuid='00000000-0000-4000-8000-000000000001',
|
||||
runtime_revision=1,
|
||||
artifact_digest='a' * 64,
|
||||
)
|
||||
|
||||
|
||||
def configure_connector(connector) -> None:
|
||||
connector._execution_context.set(TEST_EXECUTION_CONTEXT)
|
||||
connector._operation_bindings = AsyncMock(return_value=[TEST_INSTALLATION_BINDING])
|
||||
connector._load_workspace_settings = AsyncMock(return_value=[])
|
||||
connector.handler.installation_scope = MagicMock(side_effect=lambda _binding: nullcontext())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -17,6 +44,7 @@ async def test_plugin_list_filter_by_component_kinds():
|
||||
# Create connector
|
||||
connector = PluginRuntimeConnector(mock_app, AsyncMock())
|
||||
connector.handler = MagicMock()
|
||||
configure_connector(connector)
|
||||
|
||||
# Mock plugin data with different component kinds
|
||||
mock_plugins = [
|
||||
@@ -90,7 +118,7 @@ async def test_plugin_list_filter_by_component_kinds():
|
||||
# Mock database query
|
||||
async def mock_execute_async(query):
|
||||
mock_result = MagicMock()
|
||||
mock_result.__iter__ = lambda self: iter([])
|
||||
mock_result.scalars.return_value.all.return_value = []
|
||||
return mock_result
|
||||
|
||||
mock_app.persistence_mgr.execute_async = mock_execute_async
|
||||
@@ -123,6 +151,7 @@ async def test_plugin_list_filter_no_filter():
|
||||
# Create connector
|
||||
connector = PluginRuntimeConnector(mock_app, AsyncMock())
|
||||
connector.handler = MagicMock()
|
||||
configure_connector(connector)
|
||||
|
||||
# Mock plugin data with different component kinds
|
||||
mock_plugins = [
|
||||
@@ -157,7 +186,7 @@ async def test_plugin_list_filter_no_filter():
|
||||
# Mock database query
|
||||
async def mock_execute_async(query):
|
||||
mock_result = MagicMock()
|
||||
mock_result.__iter__ = lambda self: iter([])
|
||||
mock_result.scalars.return_value.all.return_value = []
|
||||
return mock_result
|
||||
|
||||
mock_app.persistence_mgr.execute_async = mock_execute_async
|
||||
@@ -184,6 +213,7 @@ async def test_plugin_list_filter_empty_result():
|
||||
# Create connector
|
||||
connector = PluginRuntimeConnector(mock_app, AsyncMock())
|
||||
connector.handler = MagicMock()
|
||||
configure_connector(connector)
|
||||
|
||||
# Mock plugin data - only KnowledgeEngine plugins
|
||||
mock_plugins = [
|
||||
@@ -206,7 +236,7 @@ async def test_plugin_list_filter_empty_result():
|
||||
# Mock database query
|
||||
async def mock_execute_async(query):
|
||||
mock_result = MagicMock()
|
||||
mock_result.__iter__ = lambda self: iter([])
|
||||
mock_result.scalars.return_value.all.return_value = []
|
||||
return mock_result
|
||||
|
||||
mock_app.persistence_mgr.execute_async = mock_execute_async
|
||||
@@ -230,6 +260,7 @@ async def test_plugin_list_filter_plugin_without_components():
|
||||
# Create connector
|
||||
connector = PluginRuntimeConnector(mock_app, AsyncMock())
|
||||
connector.handler = MagicMock()
|
||||
configure_connector(connector)
|
||||
|
||||
# Mock plugin data - one with components, one without
|
||||
mock_plugins = [
|
||||
@@ -264,7 +295,7 @@ async def test_plugin_list_filter_plugin_without_components():
|
||||
# Mock database query
|
||||
async def mock_execute_async(query):
|
||||
mock_result = MagicMock()
|
||||
mock_result.__iter__ = lambda self: iter([])
|
||||
mock_result.scalars.return_value.all.return_value = []
|
||||
return mock_result
|
||||
|
||||
mock_app.persistence_mgr.execute_async = mock_execute_async
|
||||
|
||||
@@ -1,8 +1,34 @@
|
||||
"""Test plugin list sorting functionality."""
|
||||
|
||||
from contextlib import nullcontext
|
||||
from datetime import datetime, timedelta
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from langbot_plugin.entities.io.context import InstallationBinding
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
|
||||
|
||||
TEST_EXECUTION_CONTEXT = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
)
|
||||
TEST_INSTALLATION_BINDING = InstallationBinding(
|
||||
instance_uuid=TEST_EXECUTION_CONTEXT.instance_uuid,
|
||||
workspace_uuid=TEST_EXECUTION_CONTEXT.workspace_uuid,
|
||||
placement_generation=TEST_EXECUTION_CONTEXT.placement_generation,
|
||||
installation_uuid='00000000-0000-4000-8000-000000000001',
|
||||
runtime_revision=1,
|
||||
artifact_digest='a' * 64,
|
||||
)
|
||||
|
||||
|
||||
def configure_connector(connector) -> None:
|
||||
connector._execution_context.set(TEST_EXECUTION_CONTEXT)
|
||||
connector._operation_bindings = AsyncMock(return_value=[TEST_INSTALLATION_BINDING])
|
||||
connector.handler.installation_scope = MagicMock(side_effect=lambda _binding: nullcontext())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -18,6 +44,7 @@ async def test_plugin_list_sorting_debug_first():
|
||||
# Create connector
|
||||
connector = PluginRuntimeConnector(mock_app, AsyncMock())
|
||||
connector.handler = MagicMock()
|
||||
configure_connector(connector)
|
||||
|
||||
# Mock plugin data with different debug states and timestamps
|
||||
now = datetime.now()
|
||||
@@ -60,9 +87,7 @@ async def test_plugin_list_sorting_debug_first():
|
||||
connector.handler.list_plugins = AsyncMock(return_value=mock_plugins)
|
||||
|
||||
# Mock database query to return all timestamps in a single batch
|
||||
async def mock_execute_async(query):
|
||||
mock_result = MagicMock()
|
||||
|
||||
async def mock_load_workspace_settings(_execution_context):
|
||||
# Create mock rows for all plugins with timestamps
|
||||
mock_rows = []
|
||||
|
||||
@@ -85,12 +110,9 @@ async def test_plugin_list_sorting_debug_first():
|
||||
mock_row3.created_at = now
|
||||
mock_rows.append(mock_row3)
|
||||
|
||||
# Make the result iterable
|
||||
mock_result.__iter__ = lambda self: iter(mock_rows)
|
||||
return mock_rows
|
||||
|
||||
return mock_result
|
||||
|
||||
mock_app.persistence_mgr.execute_async = mock_execute_async
|
||||
connector._load_workspace_settings = AsyncMock(side_effect=mock_load_workspace_settings)
|
||||
|
||||
# Call list_plugins
|
||||
result = await connector.list_plugins()
|
||||
@@ -120,6 +142,7 @@ async def test_plugin_list_sorting_by_installation_time():
|
||||
# Create connector
|
||||
connector = PluginRuntimeConnector(mock_app, AsyncMock())
|
||||
connector.handler = MagicMock()
|
||||
configure_connector(connector)
|
||||
|
||||
# Mock plugin data - all non-debug with different installation times
|
||||
now = datetime.now()
|
||||
@@ -162,9 +185,7 @@ async def test_plugin_list_sorting_by_installation_time():
|
||||
connector.handler.list_plugins = AsyncMock(return_value=mock_plugins)
|
||||
|
||||
# Mock database query to return all timestamps in a single batch
|
||||
async def mock_execute_async(query):
|
||||
mock_result = MagicMock()
|
||||
|
||||
async def mock_load_workspace_settings(_execution_context):
|
||||
# Create mock rows for all plugins with timestamps
|
||||
mock_rows = []
|
||||
|
||||
@@ -187,12 +208,9 @@ async def test_plugin_list_sorting_by_installation_time():
|
||||
mock_row3.created_at = now
|
||||
mock_rows.append(mock_row3)
|
||||
|
||||
# Make the result iterable
|
||||
mock_result.__iter__ = lambda self: iter(mock_rows)
|
||||
return mock_rows
|
||||
|
||||
return mock_result
|
||||
|
||||
mock_app.persistence_mgr.execute_async = mock_execute_async
|
||||
connector._load_workspace_settings = AsyncMock(side_effect=mock_load_workspace_settings)
|
||||
|
||||
# Call list_plugins
|
||||
result = await connector.list_plugins()
|
||||
@@ -217,6 +235,7 @@ async def test_plugin_list_empty():
|
||||
# Create connector
|
||||
connector = PluginRuntimeConnector(mock_app, AsyncMock())
|
||||
connector.handler = MagicMock()
|
||||
configure_connector(connector)
|
||||
|
||||
# Mock empty plugin list
|
||||
connector.handler.list_plugins = AsyncMock(return_value=[])
|
||||
|
||||
@@ -304,6 +304,22 @@ class TestSkillPathHelpers:
|
||||
assert 'export VIRTUAL_ENV="$_LB_VENV_DIR"' in command
|
||||
assert command.rstrip().endswith('python scripts/run.py')
|
||||
|
||||
def test_wrap_skill_python_env_keeps_state_outside_read_only_source(self):
|
||||
from langbot.pkg.provider.tools.loaders.skill import wrap_skill_command_with_python_env
|
||||
|
||||
command = wrap_skill_command_with_python_env(
|
||||
'python scripts/run.py',
|
||||
mount_path='/workspace/.skills/demo',
|
||||
state_path='/workspace/.skill-envs/demo',
|
||||
)
|
||||
|
||||
assert '_LB_VENV_DIR="/workspace/.skill-envs/demo/.venv"' in command
|
||||
assert '_LB_META_DIR="/workspace/.skill-envs/demo/.langbot"' in command
|
||||
assert '_LB_TMP_DIR="/workspace/.skill-envs/demo/.tmp"' in command
|
||||
assert '_LB_PIP_CACHE_DIR="/workspace/.skill-envs/demo/.cache/pip"' in command
|
||||
assert 'root = "/workspace/.skills/demo"' in command
|
||||
assert 'pip install "/workspace/.skills/demo"' in command
|
||||
|
||||
|
||||
class TestSkillToolLoader:
|
||||
"""The skill tool surface is now just ``activate`` + ``register_skill``.
|
||||
@@ -497,7 +513,11 @@ class TestNativeToolLoaderSkillPaths:
|
||||
f.write('demo instructions')
|
||||
|
||||
ap = _make_ap()
|
||||
ap.box_service = SimpleNamespace(available=True, default_workspace=tmpdir)
|
||||
ap.box_service = SimpleNamespace(
|
||||
available=True,
|
||||
default_workspace=tmpdir,
|
||||
shares_filesystem_with_box=True,
|
||||
)
|
||||
ap.skill_mgr = _make_skill_manager({'demo': _make_skill_data(name='demo', package_root=tmpdir)})
|
||||
loader = NativeToolLoader(ap)
|
||||
|
||||
@@ -511,6 +531,70 @@ class TestNativeToolLoaderSkillPaths:
|
||||
assert result['content'] == 'demo instructions'
|
||||
assert result['truncated'] is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_external_runtime_read_never_interprets_package_root_on_core_host(self):
|
||||
from langbot.pkg.provider.tools.loaders.native import NativeToolLoader
|
||||
from langbot.pkg.provider.tools.loaders.skill import PIPELINE_BOUND_SKILLS_KEY
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
with open(os.path.join(tmpdir, 'SKILL.md'), 'w', encoding='utf-8') as file_obj:
|
||||
file_obj.write('core-host-secret')
|
||||
|
||||
ap = _make_ap()
|
||||
ap.box_service = SimpleNamespace(
|
||||
available=True,
|
||||
shares_filesystem_with_box=False,
|
||||
read_skill_file=AsyncMock(return_value={'content': 'runtime-owned-content'}),
|
||||
)
|
||||
ap.skill_mgr = _make_skill_manager({'demo': _make_skill_data(name='demo', package_root=tmpdir)})
|
||||
loader = NativeToolLoader(ap)
|
||||
query = _make_query(
|
||||
query_id='q-external-read',
|
||||
variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']},
|
||||
)
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'read',
|
||||
{'path': '/workspace/.skills/demo/SKILL.md'},
|
||||
query,
|
||||
)
|
||||
|
||||
assert result['ok'] is True
|
||||
assert result['content'] == 'runtime-owned-content'
|
||||
assert 'core-host-secret' not in repr(result)
|
||||
ap.box_service.read_skill_file.assert_awaited_once_with(_CONTEXT, 'demo', 'SKILL.md')
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_external_runtime_rejects_skill_host_fallback_without_protocol_capability(self):
|
||||
from langbot.pkg.provider.tools.loaders.native import NativeToolLoader
|
||||
from langbot.pkg.provider.tools.loaders.skill import PIPELINE_BOUND_SKILLS_KEY
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
with open(os.path.join(tmpdir, 'secret.txt'), 'w', encoding='utf-8') as file_obj:
|
||||
file_obj.write('core-host-secret')
|
||||
|
||||
ap = _make_ap()
|
||||
ap.box_service = SimpleNamespace(
|
||||
available=True,
|
||||
shares_filesystem_with_box=False,
|
||||
)
|
||||
ap.skill_mgr = _make_skill_manager({'demo': _make_skill_data(name='demo', package_root=tmpdir)})
|
||||
loader = NativeToolLoader(ap)
|
||||
query = _make_query(
|
||||
query_id='q-external-no-protocol',
|
||||
variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']},
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match='owned by the Box Runtime'):
|
||||
await loader.invoke_tool(
|
||||
'grep',
|
||||
{
|
||||
'path': '/workspace/.skills/demo',
|
||||
'pattern': 'core-host-secret',
|
||||
},
|
||||
query,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exec_in_activated_skill_mount_rewrites_command_and_refreshes(self):
|
||||
from langbot.pkg.provider.tools.loaders.native import NativeToolLoader
|
||||
@@ -542,8 +626,49 @@ class TestNativeToolLoaderSkillPaths:
|
||||
tool_parameters = ap.box_service.execute_tool.await_args.args[0]
|
||||
assert tool_parameters['command'] == 'python /workspace/.skills/demo/scripts/run.py'
|
||||
assert tool_parameters['workdir'] == '/workspace/.skills/demo'
|
||||
assert ap.box_service.execute_tool.await_args.kwargs['skill_name'] == 'demo'
|
||||
ap.skill_mgr.refresh_skill_from_disk.assert_called_once_with(_CONTEXT, 'demo')
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_external_runtime_python_skill_uses_trusted_metadata_and_writable_env(self):
|
||||
from langbot.pkg.provider.tools.loaders.native import NativeToolLoader
|
||||
from langbot.pkg.provider.tools.loaders.skill import register_activated_skill
|
||||
|
||||
ap = _make_ap()
|
||||
ap.box_service = SimpleNamespace(
|
||||
available=True,
|
||||
shares_filesystem_with_box=False,
|
||||
execute_tool=AsyncMock(return_value={'ok': True}),
|
||||
)
|
||||
ap.skill_mgr = SimpleNamespace(refresh_skill_from_disk=Mock())
|
||||
loader = NativeToolLoader(ap)
|
||||
query = _make_query(query_id='q-external', launcher_type='person', launcher_id='123')
|
||||
register_activated_skill(
|
||||
query,
|
||||
_make_skill_data(
|
||||
name='demo',
|
||||
package_root='/box-runtime/skills/tenants/workspace/demo',
|
||||
python_project=True,
|
||||
),
|
||||
)
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'exec',
|
||||
{
|
||||
'command': 'python /workspace/.skills/demo/scripts/run.py',
|
||||
'workdir': '/workspace/.skills/demo',
|
||||
},
|
||||
query,
|
||||
)
|
||||
|
||||
assert result['ok'] is True
|
||||
tool_parameters = ap.box_service.execute_tool.await_args.args[0]
|
||||
wrapped = tool_parameters['command']
|
||||
assert '_LB_VENV_DIR="/workspace/.skill-envs/demo/.venv"' in wrapped
|
||||
assert 'root = "/workspace/.skills/demo"' in wrapped
|
||||
assert '/box-runtime/skills/tenants/workspace/demo' not in wrapped
|
||||
assert ap.box_service.execute_tool.await_args.kwargs['skill_name'] == 'demo'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_requires_skill_activation(self):
|
||||
from langbot.pkg.provider.tools.loaders.native import NativeToolLoader
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import contextlib
|
||||
import os
|
||||
import tempfile
|
||||
from types import SimpleNamespace
|
||||
@@ -11,6 +12,7 @@ import pytest
|
||||
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
|
||||
|
||||
from langbot.pkg.provider.tools.loaders.native import NativeToolLoader
|
||||
from langbot.pkg.provider.tools.loaders import native as native_loader
|
||||
from langbot.pkg.provider.tools.toolmgr import ToolManager
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
|
||||
@@ -52,7 +54,8 @@ class StubLoader:
|
||||
for tool in self._tools
|
||||
]
|
||||
|
||||
async def has_tool(self, name: str) -> bool:
|
||||
async def has_tool(self, *args) -> bool:
|
||||
name = args[-1]
|
||||
return any(tool.name == name for tool in self._tools)
|
||||
|
||||
async def invoke_tool(self, name: str, parameters: dict, query):
|
||||
@@ -129,11 +132,55 @@ async def test_tool_manager_routes_native_tool_calls():
|
||||
manager.plugin_tool_loader = StubLoader([make_tool('plugin_tool')])
|
||||
manager.mcp_tool_loader = StubLoader([make_tool('mcp_tool')])
|
||||
|
||||
result = await manager.execute_func_call('exec', {'command': 'pwd'}, query=Mock())
|
||||
query = SimpleNamespace(
|
||||
_execution_context=_CONTEXT,
|
||||
bot_uuid=None,
|
||||
pipeline_uuid=None,
|
||||
query_uuid=None,
|
||||
)
|
||||
result = await manager.execute_func_call('exec', {'command': 'pwd'}, query=query)
|
||||
|
||||
assert result == {'backend': 'fake'}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_manager_hides_sandbox_and_skill_tools_without_workspace_entitlement():
|
||||
box_service = SimpleNamespace(is_workspace_sandbox_available=AsyncMock(return_value=False))
|
||||
manager = ToolManager(SimpleNamespace(box_service=box_service))
|
||||
manager.native_tool_loader = StubLoader([make_tool('exec')])
|
||||
manager.skill_tool_loader = StubLoader([make_tool('activate')])
|
||||
manager.plugin_tool_loader = StubLoader([make_tool('plugin_tool')])
|
||||
manager.mcp_tool_loader = StubLoader([make_tool('mcp_tool')])
|
||||
|
||||
tools = await manager.get_all_tools(_CONTEXT, include_skill_authoring=True)
|
||||
catalog = await manager.get_tool_catalog(_CONTEXT, include_skill_authoring=True)
|
||||
|
||||
assert [tool.name for tool in tools] == ['plugin_tool', 'mcp_tool']
|
||||
assert [item['name'] for item in catalog] == ['plugin_tool', 'mcp_tool']
|
||||
assert box_service.is_workspace_sandbox_available.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_manager_rechecks_workspace_entitlement_before_native_invocation():
|
||||
box_service = SimpleNamespace(is_workspace_sandbox_available=AsyncMock(return_value=False))
|
||||
manager = ToolManager(SimpleNamespace(box_service=box_service))
|
||||
manager.native_tool_loader = StubLoader([make_tool('exec')], invoke_result={'unexpected': True})
|
||||
manager.skill_tool_loader = StubLoader([])
|
||||
manager.plugin_tool_loader = StubLoader([])
|
||||
manager.mcp_tool_loader = StubLoader([])
|
||||
query = SimpleNamespace(
|
||||
_execution_context=_CONTEXT,
|
||||
bot_uuid=None,
|
||||
pipeline_uuid=None,
|
||||
query_uuid=None,
|
||||
)
|
||||
|
||||
with pytest.raises(Exception, match='exec'):
|
||||
await manager.execute_func_call('exec', {'command': 'pwd'}, query=query)
|
||||
|
||||
box_service.is_workspace_sandbox_available.assert_awaited_once_with(_CONTEXT)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_tool_loader_hides_tools_when_box_unavailable():
|
||||
loader = NativeToolLoader(SimpleNamespace(box_service=SimpleNamespace(available=False)))
|
||||
@@ -159,6 +206,26 @@ async def test_native_tool_loader_exposes_all_tools_when_box_available():
|
||||
assert await loader.has_tool(tool_name) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_tool_loader_rechecks_admission_at_the_final_invoke_boundary():
|
||||
box_service = SimpleNamespace(
|
||||
available=True,
|
||||
require_workspace_sandbox=AsyncMock(side_effect=RuntimeError('entitlement expired')),
|
||||
)
|
||||
loader = NativeToolLoader(SimpleNamespace(box_service=box_service, logger=Mock()))
|
||||
query = SimpleNamespace(
|
||||
_execution_context=_CONTEXT,
|
||||
bot_uuid=None,
|
||||
pipeline_uuid=None,
|
||||
query_uuid=None,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match='entitlement expired'):
|
||||
await loader.invoke_tool('read', {'path': '/workspace/private.txt'}, query)
|
||||
|
||||
box_service.require_workspace_sandbox.assert_awaited_once_with(_CONTEXT)
|
||||
|
||||
|
||||
# ── read/write/edit file tool tests ─────────────────────────────
|
||||
|
||||
|
||||
@@ -386,6 +453,103 @@ async def test_path_escape_blocked():
|
||||
await loader.invoke_tool('read', {'path': '/workspace/../../etc/passwd'}, _make_query())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('tool_name', 'parameters'),
|
||||
[
|
||||
('read', {'path': '/workspace/shared/tenant-b-only.txt'}),
|
||||
(
|
||||
'write',
|
||||
{'path': '/workspace/shared/tenant-b-only.txt', 'content': 'overwritten by tenant a'},
|
||||
),
|
||||
(
|
||||
'edit',
|
||||
{
|
||||
'path': '/workspace/shared/tenant-b-only.txt',
|
||||
'old_string': 'tenant-b-secret',
|
||||
'new_string': 'overwritten by tenant a',
|
||||
},
|
||||
),
|
||||
('glob', {'path': '/workspace/shared', 'pattern': '*'}),
|
||||
('grep', {'path': '/workspace/shared', 'pattern': 'tenant-b-secret'}),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_host_workspace_operations_do_not_follow_a_swapped_ancestor(
|
||||
monkeypatch,
|
||||
tool_name: str,
|
||||
parameters: dict,
|
||||
):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tenant_a = os.path.join(tmpdir, 'tenant-a')
|
||||
tenant_b = os.path.join(tmpdir, 'tenant-b')
|
||||
os.makedirs(os.path.join(tenant_a, 'shared'))
|
||||
os.makedirs(os.path.join(tenant_b, 'shared'))
|
||||
tenant_b_file = os.path.join(tenant_b, 'shared', 'tenant-b-only.txt')
|
||||
with open(os.path.join(tenant_a, 'shared', 'tenant-a-only.txt'), 'w', encoding='utf-8') as file_obj:
|
||||
file_obj.write('tenant-a-content')
|
||||
with open(tenant_b_file, 'w', encoding='utf-8') as file_obj:
|
||||
file_obj.write('tenant-b-secret')
|
||||
|
||||
box_service = SimpleNamespace(
|
||||
available=True,
|
||||
default_workspace=tmpdir,
|
||||
_tenant_workspace=Mock(return_value=tenant_a),
|
||||
)
|
||||
loader = NativeToolLoader(SimpleNamespace(box_service=box_service, logger=Mock()))
|
||||
original_open_host_root = native_loader._open_host_root
|
||||
|
||||
@contextlib.contextmanager
|
||||
def open_host_root_after_swap(location, *, create):
|
||||
with original_open_host_root(location, create=create) as root_fd:
|
||||
original_ancestor = os.path.join(tenant_a, 'shared-original')
|
||||
os.rename(os.path.join(tenant_a, 'shared'), original_ancestor)
|
||||
os.symlink(os.path.join(tenant_b, 'shared'), os.path.join(tenant_a, 'shared'))
|
||||
yield root_fd
|
||||
|
||||
monkeypatch.setattr(native_loader, '_open_host_root', open_host_root_after_swap)
|
||||
|
||||
try:
|
||||
result = await loader.invoke_tool(tool_name, parameters, _make_query())
|
||||
except ValueError as exc:
|
||||
result = {'ok': False, 'error': str(exc)}
|
||||
|
||||
assert result.get('ok') is False
|
||||
assert 'tenant-b-secret' not in repr(result)
|
||||
assert 'tenant-b-only.txt' not in repr(result)
|
||||
with open(tenant_b_file, encoding='utf-8') as file_obj:
|
||||
assert file_obj.read() == 'tenant-b-secret'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_host_file_api_falls_back_to_tenant_box_when_openat_is_unavailable(monkeypatch):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
box_service = SimpleNamespace(
|
||||
available=True,
|
||||
default_workspace=tmpdir,
|
||||
_tenant_workspace=Mock(return_value=tmpdir),
|
||||
execute_tool=AsyncMock(
|
||||
return_value={
|
||||
'ok': True,
|
||||
'stdout': '{"ok": true, "content": "box-owned", "truncated": false}',
|
||||
'stderr': '',
|
||||
}
|
||||
),
|
||||
)
|
||||
loader = NativeToolLoader(SimpleNamespace(box_service=box_service, logger=Mock()))
|
||||
monkeypatch.setattr(native_loader, '_SECURE_HOST_FILE_OPS_AVAILABLE', False)
|
||||
|
||||
result = await loader.invoke_tool(
|
||||
'read',
|
||||
{'path': '/workspace/file.txt'},
|
||||
_make_query(),
|
||||
)
|
||||
|
||||
assert result['ok'] is True
|
||||
assert result['content'] == 'box-owned'
|
||||
command = box_service.execute_tool.await_args.args[0]['command']
|
||||
assert 'path = "/workspace/file.txt"' in command
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_box_availability_helper_handles_unavailable_and_errors():
|
||||
from langbot.pkg.provider.tools.loaders.availability import is_box_backend_available
|
||||
|
||||
@@ -2,8 +2,10 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextvars
|
||||
import io
|
||||
import zipfile
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
@@ -54,7 +56,7 @@ def _make_app() -> Mock:
|
||||
storage_mgr.storage_provider.delete = AsyncMock()
|
||||
app.storage_mgr = storage_mgr
|
||||
app.persistence_mgr = Mock()
|
||||
app.persistence_mgr.execute_async = AsyncMock()
|
||||
app.persistence_mgr.execute_async = AsyncMock(return_value=SimpleNamespace(rowcount=1))
|
||||
app.plugin_connector = Mock()
|
||||
app.plugin_connector.require_workspace_context = AsyncMock(side_effect=lambda context: context)
|
||||
app.workspace_service = SimpleNamespace(
|
||||
@@ -202,6 +204,36 @@ class TestStoreZipFile:
|
||||
|
||||
|
||||
class TestStoreFileTask:
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_file_task_opens_uow_before_first_database_helper(self):
|
||||
kb = _make_kb()
|
||||
active_workspace = contextvars.ContextVar('rag_task_workspace', default=None)
|
||||
observed = []
|
||||
|
||||
@asynccontextmanager
|
||||
async def tenant_uow(workspace_uuid):
|
||||
token = active_workspace.set(workspace_uuid)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
active_workspace.reset(token)
|
||||
|
||||
kb.ap.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime')
|
||||
kb.ap.persistence_mgr.tenant_uow = tenant_uow
|
||||
|
||||
async def assert_execution_context(_context):
|
||||
observed.append(active_workspace.get())
|
||||
|
||||
kb._assert_execution_context = AsyncMock(side_effect=assert_execution_context)
|
||||
kb._set_file_status = AsyncMock(side_effect=[True, True])
|
||||
kb._ingest_document = AsyncMock(return_value={'status': 'completed'})
|
||||
object_key = _upload_key('scoped.pdf')
|
||||
file_obj = SimpleNamespace(uuid='file-uuid', file_name=object_key, extension='pdf')
|
||||
|
||||
await kb._store_file_task(CONTEXT, file_obj, Mock())
|
||||
|
||||
assert observed[0] == WORKSPACE_A
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_file_task_marks_completed_and_cleans_storage(self):
|
||||
kb = _make_kb()
|
||||
|
||||
@@ -1,12 +1,13 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
import sqlalchemy
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
|
||||
|
||||
from langbot.pkg.api.http.authz import WorkspaceRequiredError
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
@@ -17,6 +18,7 @@ from langbot.pkg.entity.persistence.workspace import Workspace, WorkspaceExecuti
|
||||
from langbot.pkg.rag.knowledge.kbmgr import RAGManager
|
||||
from langbot.pkg.rag.service.runtime import RAGRuntimeService
|
||||
from langbot.pkg.vector.mgr import VectorDBManager
|
||||
from langbot.pkg.vector.vdbs.pgvector_db import PgVectorDatabase
|
||||
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
|
||||
from langbot.pkg.workspace.policy import CloudWorkspacePolicy, SingleWorkspacePolicy
|
||||
from langbot.pkg.workspace.service import WorkspaceService
|
||||
@@ -42,6 +44,13 @@ class _PersistenceManager:
|
||||
await connection.commit()
|
||||
return result
|
||||
|
||||
@asynccontextmanager
|
||||
async def tenant_uow(self, _workspace_uuid):
|
||||
# This lightweight fixture does not emulate PostgreSQL RLS; production
|
||||
# persistence tests cover the transaction-bound unit of work itself.
|
||||
async with AsyncSession(self.engine, expire_on_commit=False) as session, session.begin():
|
||||
yield SimpleNamespace(session=session)
|
||||
|
||||
@staticmethod
|
||||
def serialize_model(model, row, masked_columns=()):
|
||||
return {
|
||||
@@ -348,3 +357,34 @@ async def test_stale_generation_is_rejected_before_vector_access(tenant_rag):
|
||||
with pytest.raises(Exception, match='generation'):
|
||||
await manager.upsert(stale, 'kb-a', [[0.1]], ['a'])
|
||||
assert database.collections == []
|
||||
|
||||
|
||||
async def test_pgvector_first_write_binds_dimension_and_later_mismatch_fails(tenant_rag):
|
||||
app, engine = tenant_rag
|
||||
manager = VectorDBManager(app)
|
||||
pgvector = object.__new__(PgVectorDatabase)
|
||||
pgvector.allowed_dimensions = frozenset({1, 2})
|
||||
pgvector.add_embeddings = AsyncMock()
|
||||
pgvector.search = AsyncMock(return_value={'ids': [[]], 'distances': [[]], 'metadatas': [[]]})
|
||||
manager.vector_db = pgvector
|
||||
context = _context(WORKSPACE_A)
|
||||
|
||||
await manager.upsert(context, 'kb-a', [[0.1]], ['chunk-a'])
|
||||
scope = pgvector.add_embeddings.await_args.kwargs['scope']
|
||||
assert scope.workspace_uuid == WORKSPACE_A
|
||||
assert scope.knowledge_base_uuid == 'kb-a'
|
||||
assert scope.embedding_dimension == 1
|
||||
|
||||
async with engine.connect() as connection:
|
||||
selected_dimension = await connection.scalar(
|
||||
sqlalchemy.select(KnowledgeBase.embedding_dimension).where(
|
||||
KnowledgeBase.workspace_uuid == WORKSPACE_A,
|
||||
KnowledgeBase.uuid == 'kb-a',
|
||||
)
|
||||
)
|
||||
assert selected_dimension == 1
|
||||
|
||||
with pytest.raises(ValueError, match='dimension is 1, not 2'):
|
||||
await manager.upsert(context, 'kb-a', [[0.1, 0.2]], ['chunk-b'])
|
||||
with pytest.raises(ValueError, match='not enabled'):
|
||||
await manager.search(context, 'kb-a', [0.1, 0.2, 0.3], 3)
|
||||
|
||||
@@ -32,6 +32,12 @@ def create_mock_app():
|
||||
return mock_app
|
||||
|
||||
|
||||
def scalar_result(value=None):
|
||||
"""Return the scalar-only shape used by Connection column selects."""
|
||||
|
||||
return Mock(scalar_one_or_none=Mock(return_value=value))
|
||||
|
||||
|
||||
class TestSurveyManagerInit:
|
||||
"""Tests for SurveyManager initialization."""
|
||||
|
||||
@@ -67,7 +73,7 @@ class TestSurveyManagerInit:
|
||||
"""Test that initialize loads space URL from config."""
|
||||
survey_module = get_survey_module()
|
||||
mock_app = create_mock_app()
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=Mock(first=Mock(return_value=None)))
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=scalar_result())
|
||||
|
||||
manager = survey_module.SurveyManager(mock_app)
|
||||
await manager.initialize()
|
||||
@@ -80,7 +86,7 @@ class TestSurveyManagerInit:
|
||||
survey_module = get_survey_module()
|
||||
mock_app = create_mock_app()
|
||||
mock_app.instance_config.data = {'space': {'url': 'https://space.example.com/'}}
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=Mock(first=Mock(return_value=None)))
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=scalar_result())
|
||||
|
||||
manager = survey_module.SurveyManager(mock_app)
|
||||
await manager.initialize()
|
||||
@@ -93,7 +99,7 @@ class TestSurveyManagerInit:
|
||||
survey_module = get_survey_module()
|
||||
mock_app = create_mock_app()
|
||||
mock_app.instance_config.data = {}
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=Mock(first=Mock(return_value=None)))
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=scalar_result())
|
||||
|
||||
manager = survey_module.SurveyManager(mock_app)
|
||||
await manager.initialize()
|
||||
@@ -110,12 +116,7 @@ class TestLoadTriggeredEvents:
|
||||
survey_module = get_survey_module()
|
||||
mock_app = create_mock_app()
|
||||
|
||||
# Mock existing metadata row
|
||||
mock_row = Mock()
|
||||
mock_row.value = json.dumps(['event1', 'event2'])
|
||||
mock_result = Mock()
|
||||
mock_result.first = Mock(return_value=(mock_row,))
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=scalar_result(json.dumps(['event1', 'event2'])))
|
||||
|
||||
manager = survey_module.SurveyManager(mock_app)
|
||||
await manager._load_triggered_events()
|
||||
@@ -128,7 +129,7 @@ class TestLoadTriggeredEvents:
|
||||
"""Test that empty set is used when no events stored."""
|
||||
survey_module = get_survey_module()
|
||||
mock_app = create_mock_app()
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=Mock(first=Mock(return_value=None)))
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=scalar_result())
|
||||
|
||||
manager = survey_module.SurveyManager(mock_app)
|
||||
await manager._load_triggered_events()
|
||||
@@ -218,7 +219,7 @@ class TestTriggerEvent:
|
||||
"""Test that new event is added and saved."""
|
||||
survey_module = get_survey_module()
|
||||
mock_app = create_mock_app()
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=Mock(first=Mock(return_value=None)))
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=scalar_result())
|
||||
|
||||
manager = survey_module.SurveyManager(mock_app)
|
||||
manager._space_url = 'https://space.example.com'
|
||||
@@ -235,7 +236,7 @@ class TestRecordBotResponseSuccess:
|
||||
manager = survey_module.SurveyManager(mock_app)
|
||||
manager._space_url = 'https://space.example.com'
|
||||
# No existing metadata rows: select returns no row
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=Mock(first=Mock(return_value=None)))
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=scalar_result())
|
||||
return manager
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -304,19 +305,13 @@ class TestRecordBotResponseSuccess:
|
||||
survey_module = get_survey_module()
|
||||
mock_app = create_mock_app()
|
||||
|
||||
count_row = Mock()
|
||||
count_row.value = '42'
|
||||
|
||||
def execute_side_effect(stmt):
|
||||
result = Mock()
|
||||
# Both _load_triggered_events and _load_bot_response_count select
|
||||
# from Metadata; return the count row only for the count key.
|
||||
# Metadata.value; return the count only for the count key.
|
||||
stmt_str = str(stmt.compile(compile_kwargs={'literal_binds': True}))
|
||||
if survey_module.BOT_RESPONSE_COUNT_KEY in stmt_str:
|
||||
result.first.return_value = (count_row,)
|
||||
else:
|
||||
result.first.return_value = None
|
||||
return result
|
||||
return scalar_result('42')
|
||||
return scalar_result()
|
||||
|
||||
mock_app.persistence_mgr.execute_async = AsyncMock(side_effect=execute_side_effect)
|
||||
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
import io
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
import zipfile
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
@@ -115,3 +118,52 @@ class TestRequireBoxForWrite:
|
||||
service = SkillService(self._ap_with_disabled_box())
|
||||
with pytest.raises(ValueError, match='Reading a skill file'):
|
||||
await service.read_skill_file(_CONTEXT, 'x', 'a.txt')
|
||||
|
||||
|
||||
class TestGithubSkillArchiveLimits:
|
||||
@staticmethod
|
||||
def _service() -> SkillService:
|
||||
return SkillService(SimpleNamespace())
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_rejects_declared_oversized_archive(self, monkeypatch):
|
||||
import langbot.pkg.api.http.service.skill as skill_module
|
||||
|
||||
real_client = httpx.AsyncClient
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={'content-length': str(10 * 1024 * 1024 + 1)},
|
||||
content=b'',
|
||||
request=request,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
skill_module.httpx,
|
||||
'AsyncClient',
|
||||
lambda **_kwargs: real_client(transport=httpx.MockTransport(handler)),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match='compressed size limit'):
|
||||
await self._service()._download_github_asset('https://codeload.github.com/o/r/zip/main')
|
||||
|
||||
def test_copy_rejects_high_compression_ratio_before_extracting(self):
|
||||
source_buffer = io.BytesIO()
|
||||
with zipfile.ZipFile(source_buffer, 'w', zipfile.ZIP_DEFLATED) as archive:
|
||||
archive.writestr('repo-main/skill/SKILL.md', b'---\nname: safe\n---\n')
|
||||
archive.writestr('repo-main/skill/bomb.bin', b'0' * (1024 * 1024))
|
||||
|
||||
source_buffer.seek(0)
|
||||
target_buffer = io.BytesIO()
|
||||
with (
|
||||
zipfile.ZipFile(source_buffer, 'r') as source_zip,
|
||||
zipfile.ZipFile(target_buffer, 'w', zipfile.ZIP_DEFLATED) as target_zip,
|
||||
):
|
||||
with pytest.raises(ValueError, match='compression-ratio limit'):
|
||||
self._service()._copy_github_skill_directory_to_zip(
|
||||
source_zip,
|
||||
target_zip,
|
||||
'repo-main/skill',
|
||||
'safe',
|
||||
)
|
||||
|
||||
@@ -138,6 +138,7 @@ class TestVectorDBManagerInitialization:
|
||||
mgr = VectorDBManager(mock_app)
|
||||
|
||||
import asyncio
|
||||
|
||||
asyncio.get_event_loop().run_until_complete(mgr.initialize())
|
||||
|
||||
mock_valkey_class.assert_called_once_with(mock_app)
|
||||
@@ -207,7 +208,10 @@ class TestVectorDBManagerInitialization:
|
||||
asyncio.get_event_loop().run_until_complete(mgr.initialize())
|
||||
|
||||
mock_pgvector_class.assert_called_once_with(
|
||||
mock_app, connection_string='postgresql://user:pass@host:5432/langbot'
|
||||
mock_app,
|
||||
connection_string='postgresql://user:pass@host:5432/langbot',
|
||||
use_business_database=False,
|
||||
allowed_dimensions=[384, 512, 768, 1024, 1536],
|
||||
)
|
||||
|
||||
def test_initialize_pgvector_with_individual_params(self):
|
||||
@@ -238,7 +242,14 @@ class TestVectorDBManagerInitialization:
|
||||
asyncio.get_event_loop().run_until_complete(mgr.initialize())
|
||||
|
||||
mock_pgvector_class.assert_called_once_with(
|
||||
mock_app, host='db.example.com', port=5433, database='vectordb', user='admin', password='secret'
|
||||
mock_app,
|
||||
host='db.example.com',
|
||||
port=5433,
|
||||
database='vectordb',
|
||||
user='admin',
|
||||
password='secret',
|
||||
use_business_database=False,
|
||||
allowed_dimensions=[384, 512, 768, 1024, 1536],
|
||||
)
|
||||
|
||||
def test_initialize_pgvector_defaults(self):
|
||||
@@ -260,7 +271,42 @@ class TestVectorDBManagerInitialization:
|
||||
asyncio.get_event_loop().run_until_complete(mgr.initialize())
|
||||
|
||||
mock_pgvector_class.assert_called_once_with(
|
||||
mock_app, host='localhost', port=5432, database='langbot', user='postgres', password='postgres'
|
||||
mock_app,
|
||||
host='localhost',
|
||||
port=5432,
|
||||
database='langbot',
|
||||
user='postgres',
|
||||
password='postgres',
|
||||
use_business_database=False,
|
||||
allowed_dimensions=[384, 512, 768, 1024, 1536],
|
||||
)
|
||||
|
||||
def test_initialize_pgvector_with_shared_business_database(self):
|
||||
vdb_config = {
|
||||
'use': 'pgvector',
|
||||
'pgvector': {
|
||||
'use_business_database': True,
|
||||
'allowed_dimensions': [768, 1536],
|
||||
},
|
||||
}
|
||||
mock_app = self._create_mock_app(vdb_config)
|
||||
mocks = self._make_vector_import_mocks()
|
||||
mock_pgvector_class = MagicMock()
|
||||
mocks['langbot.pkg.vector.vdbs.pgvector_db'].PgVectorDatabase = mock_pgvector_class
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.vector.mgr import VectorDBManager
|
||||
|
||||
mgr = VectorDBManager(mock_app)
|
||||
|
||||
import asyncio
|
||||
|
||||
asyncio.get_event_loop().run_until_complete(mgr.initialize())
|
||||
|
||||
mock_pgvector_class.assert_called_once_with(
|
||||
mock_app,
|
||||
use_business_database=True,
|
||||
allowed_dimensions=[768, 1536],
|
||||
)
|
||||
|
||||
def test_initialize_unknown_backend_defaults_to_chroma(self):
|
||||
|
||||
Reference in New Issue
Block a user