xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-07-21 11:56:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(tenancy): harden shared cloud runtime boundaries

Browse Source
This commit is contained in:
Junyan Qin
2026-07-20 01:47:42 +08:00
parent 41772920ef
commit a47bfe8167
121 changed files with 18152 additions and 5588 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit_tests/plugin/test_handler_tenancy.py
+180 -37
View File
@@ -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