mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 04:40:57 +00:00
chore(merge): sync master into dev/4.11.x
This commit is contained in:
huanghuoguoguo
@@ -0,0 +1,454 @@
|
||||
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, 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:
|
||||
def first(self):
|
||||
return None
|
||||
|
||||
def all(self):
|
||||
return []
|
||||
|
||||
|
||||
class RecordingConnection(Connection):
|
||||
def __init__(self):
|
||||
self.sent: list[str] = []
|
||||
|
||||
async def send(self, message: str) -> None:
|
||||
self.sent.append(message)
|
||||
|
||||
async def receive(self) -> str:
|
||||
raise NotImplementedError
|
||||
|
||||
async def close(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def workspace_context(workspace_uuid: str = 'workspace-a') -> ActionContext:
|
||||
return ActionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid=workspace_uuid,
|
||||
placement_generation=7,
|
||||
)
|
||||
|
||||
|
||||
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(
|
||||
get_execution_binding=AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
instance_uuid=context.instance_uuid,
|
||||
workspace_uuid=context.workspace_uuid,
|
||||
placement_generation=context.placement_generation,
|
||||
)
|
||||
)
|
||||
),
|
||||
)
|
||||
runtime_handler = RuntimeConnectionHandler(
|
||||
Mock(),
|
||||
AsyncMock(return_value=True),
|
||||
app,
|
||||
)
|
||||
installation_context = InstallationBinding(
|
||||
**context.model_dump(exclude_none=True),
|
||||
installation_uuid='00000000-0000-4000-8000-000000000001',
|
||||
runtime_revision=1,
|
||||
artifact_digest='a' * 64,
|
||||
)
|
||||
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(
|
||||
runtime_handler: RuntimeConnectionHandler,
|
||||
action_context: ActionContext,
|
||||
action: PluginToRuntimeAction,
|
||||
data: dict,
|
||||
):
|
||||
token = runtime_handler._current_action_context.set(action_context)
|
||||
try:
|
||||
return await runtime_handler.actions[action.value](data)
|
||||
finally:
|
||||
runtime_handler._current_action_context.reset(token)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plugin_action_requires_installation_capability():
|
||||
runtime_handler, _app, _installation_context = make_handler()
|
||||
|
||||
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',
|
||||
abilities=[],
|
||||
extra_args={},
|
||||
),
|
||||
provider=SimpleNamespace(invoke_llm=invoke_llm),
|
||||
)
|
||||
)
|
||||
)
|
||||
runtime_handler._require_plugin_action_context = AsyncMock(
|
||||
return_value=(
|
||||
installation_context,
|
||||
SimpleNamespace(plugin_author='author-a', plugin_name='plugin-a'),
|
||||
)
|
||||
)
|
||||
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 asyncio.wait_for(entered.wait(), timeout=2)
|
||||
assert observations == [True]
|
||||
release.set()
|
||||
response = await action
|
||||
assert response.code == 0
|
||||
assert observations == [True, True]
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.wait_for(action, timeout=2)
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plugin_action_rejects_forged_installation_capability():
|
||||
runtime_handler, _app, installation_context = make_handler()
|
||||
forged = installation_context.model_copy(update={'installation_uuid': 'forged-installation'})
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match='installation is not registered in this Workspace',
|
||||
):
|
||||
await invoke_with_context(
|
||||
runtime_handler,
|
||||
forged,
|
||||
PluginToRuntimeAction.GET_LANGBOT_VERSION,
|
||||
{},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plugin_action_rejects_stale_workspace_generation():
|
||||
runtime_handler, app, installation_context = make_handler()
|
||||
app.workspace_service.get_execution_binding.side_effect = ValueError('generation is fenced')
|
||||
|
||||
with pytest.raises(ValueError, match='generation is fenced'):
|
||||
await invoke_with_context(
|
||||
runtime_handler,
|
||||
installation_context,
|
||||
PluginToRuntimeAction.GET_LANGBOT_VERSION,
|
||||
{},
|
||||
)
|
||||
|
||||
app.workspace_service.get_execution_binding.assert_awaited_once_with(
|
||||
'workspace-a',
|
||||
expected_generation=7,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_uuid_and_forged_payload_workspace_cannot_cross_tenants():
|
||||
runtime_handler, app, installation_context = make_handler()
|
||||
query_a = SimpleNamespace(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=7,
|
||||
bot_uuid='bot-a',
|
||||
)
|
||||
query_b = SimpleNamespace(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-b',
|
||||
placement_generation=7,
|
||||
bot_uuid='bot-b',
|
||||
)
|
||||
|
||||
async def get_query(workspace_uuid, query_uuid):
|
||||
return {
|
||||
('workspace-a', 'query-a'): query_a,
|
||||
('workspace-b', 'query-b'): query_b,
|
||||
}.get((workspace_uuid, query_uuid))
|
||||
|
||||
app.query_pool = SimpleNamespace(
|
||||
get_query=AsyncMock(side_effect=get_query),
|
||||
get_query_by_legacy_id=AsyncMock(),
|
||||
)
|
||||
|
||||
response = await invoke_with_context(
|
||||
runtime_handler,
|
||||
installation_context,
|
||||
PluginToRuntimeAction.GET_BOT_UUID,
|
||||
{
|
||||
'query_id': 2,
|
||||
'query_uuid': 'query-b',
|
||||
'workspace_uuid': 'workspace-b',
|
||||
},
|
||||
)
|
||||
|
||||
assert response.code != 0
|
||||
app.query_pool.get_query.assert_awaited_once_with('workspace-a', 'query-b')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_legacy_query_id_fallback_is_workspace_scoped():
|
||||
runtime_handler, app, installation_context = make_handler()
|
||||
query = SimpleNamespace(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=7,
|
||||
bot_uuid='bot-a',
|
||||
)
|
||||
app.query_pool = SimpleNamespace(
|
||||
get_query=AsyncMock(),
|
||||
get_query_by_legacy_id=AsyncMock(return_value=query),
|
||||
)
|
||||
|
||||
response = await invoke_with_context(
|
||||
runtime_handler,
|
||||
installation_context,
|
||||
PluginToRuntimeAction.GET_BOT_UUID,
|
||||
{'query_id': 19},
|
||||
)
|
||||
|
||||
assert response.code == 0
|
||||
assert response.data == {'bot_uuid': 'bot-a'}
|
||||
app.query_pool.get_query_by_legacy_id.assert_awaited_once_with(
|
||||
'workspace-a',
|
||||
19,
|
||||
)
|
||||
|
||||
|
||||
def test_runtime_connection_is_instance_scoped_and_unbound():
|
||||
runtime_handler, _app, _installation_context = make_handler()
|
||||
assert runtime_handler.bound_action_context is None
|
||||
|
||||
|
||||
def test_inbound_tenant_action_requires_complete_installation_envelope():
|
||||
runtime_handler, _app, installation_context = make_handler()
|
||||
|
||||
assert (
|
||||
runtime_handler.validate_inbound_action_context(
|
||||
PluginToRuntimeAction.GET_BOTS.value,
|
||||
installation_context,
|
||||
)
|
||||
== installation_context
|
||||
)
|
||||
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'),
|
||||
)
|
||||
|
||||
|
||||
@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)
|
||||
|
||||
assert (
|
||||
runtime_handler.validate_inbound_action_context(
|
||||
PluginToRuntimeAction.GET_LANGBOT_VERSION.value,
|
||||
legacy_context,
|
||||
)
|
||||
== legacy_context
|
||||
)
|
||||
response = await invoke_with_context(
|
||||
runtime_handler,
|
||||
legacy_context,
|
||||
PluginToRuntimeAction.GET_LANGBOT_VERSION,
|
||||
{},
|
||||
)
|
||||
|
||||
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
|
||||
async def test_plugin_vector_action_forwards_trusted_context_and_logical_collection():
|
||||
runtime_handler, app, installation_context = make_handler()
|
||||
app.rag_runtime_service = SimpleNamespace(vector_upsert=AsyncMock())
|
||||
|
||||
response = await invoke_with_context(
|
||||
runtime_handler,
|
||||
installation_context,
|
||||
PluginToRuntimeAction.VECTOR_UPSERT,
|
||||
{
|
||||
'workspace_uuid': 'workspace-forged',
|
||||
'collection_id': 'plugin-supplied-name',
|
||||
'vectors': [[0.1]],
|
||||
'ids': ['point-a'],
|
||||
},
|
||||
)
|
||||
|
||||
assert response.code == 0
|
||||
execution_context = app.rag_runtime_service.vector_upsert.await_args.args[0]
|
||||
assert execution_context == ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=7,
|
||||
)
|
||||
assert app.rag_runtime_service.vector_upsert.await_args.args[1] == 'plugin-supplied-name'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_host_to_runtime_action_carries_trusted_connector_context():
|
||||
app = SimpleNamespace(logger=Mock())
|
||||
connection = RecordingConnection()
|
||||
runtime_handler = RuntimeConnectionHandler(
|
||||
connection,
|
||||
AsyncMock(return_value=True),
|
||||
app,
|
||||
)
|
||||
|
||||
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
|
||||
await asyncio.sleep(0)
|
||||
request = json.loads(connection.sent[0])
|
||||
runtime_handler.resp_waiters[request['seq_id']].set_result(ActionResponse.success({}))
|
||||
await task
|
||||
|
||||
assert request['data']['runtime_identity'] == {
|
||||
'instance_uuid': 'instance-a',
|
||||
'runtime_id': 'runtime-a',
|
||||
}
|
||||
assert request.get('context') is None
|
||||
Reference in New Issue
Block a user