mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-09-16 14:57:15 +00:00
fix(plugin): preserve explicit nested installation scope (#2528)
* fix(plugin): preserve explicit nested installation scope * fix(deps): pin released RAG runtime SDK 0.5.8 --------- Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
This commit is contained in:
Hyu
@@ -0,0 +1,193 @@
|
||||
"""Exercise nested installation routing through real Core/SDK wire envelopes."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from langbot_plugin.entities.io.actions.enums import CommonAction, LangBotToRuntimeAction, PluginToRuntimeAction
|
||||
from langbot_plugin.entities.io.req import ActionRequest
|
||||
from langbot_plugin.entities.io.resp import ActionResponse
|
||||
from langbot_plugin.runtime.io import handler as sdk_handler
|
||||
|
||||
from langbot.pkg.plugin.connector import PluginRuntimeConnector
|
||||
from tests.unit_tests.plugin.test_handler_tenancy import RecordingConnection, make_handler, workspace_context
|
||||
|
||||
|
||||
class ReplyingConnection(RecordingConnection):
|
||||
"""Replace only the transport, retaining serialization and response routing."""
|
||||
|
||||
async def send(self, message: str) -> None:
|
||||
await super().send(message)
|
||||
request = json.loads(message)
|
||||
if 'action' in request:
|
||||
response = ActionResponse.success({'elements': []})
|
||||
response.seq_id = request['seq_id']
|
||||
await self.handler._route_response(response.seq_id, response.model_dump())
|
||||
|
||||
@property
|
||||
def requests(self):
|
||||
return [request for message in self.sent if 'action' in (request := json.loads(message))]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bridge(monkeypatch):
|
||||
runtime_handler, app, binding_a = make_handler()
|
||||
connection = ReplyingConnection()
|
||||
connection.handler = runtime_handler
|
||||
runtime_handler.conn = connection
|
||||
monkeypatch.setattr(sdk_handler, 'FILE_CHUNK_LENGTH', 4)
|
||||
binding_b = binding_a.model_copy(
|
||||
update={
|
||||
'installation_uuid': '00000000-0000-4000-8000-000000000002',
|
||||
'runtime_revision': 2,
|
||||
'artifact_digest': 'b' * 64,
|
||||
}
|
||||
)
|
||||
return runtime_handler, app, connection, binding_a, binding_b
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('mode', ['managed', 'legacy'])
|
||||
async def test_nested_invoke_parser_uses_target_for_every_chunk_and_parse(bridge, mode):
|
||||
runtime_handler, app, connection, binding_a, binding_b = bridge
|
||||
app.instance_config = SimpleNamespace(data={'plugin': {'enable': True}})
|
||||
app.deployment.mode = 'cloud' if mode == 'managed' else 'oss'
|
||||
connector = PluginRuntimeConnector(app, AsyncMock())
|
||||
connector.handler = runtime_handler
|
||||
app.plugin_connector = connector
|
||||
execution_context = runtime_handler._execution_context(binding_a)
|
||||
setting_b = SimpleNamespace(
|
||||
installation_uuid=binding_b.installation_uuid,
|
||||
runtime_revision=binding_b.runtime_revision,
|
||||
artifact_digest=binding_b.artifact_digest,
|
||||
install_info={'_artifact_storage': 'tenant_binary_storage_v1'} if mode == 'managed' else {},
|
||||
)
|
||||
connector._setting_for_plugin = AsyncMock(return_value=(execution_context, setting_b))
|
||||
connector.require_workspace_context = AsyncMock(return_value=execution_context)
|
||||
file_bytes = b'parser document'
|
||||
app.rag_runtime_service = SimpleNamespace(get_file_stream=AsyncMock(return_value=file_bytes))
|
||||
inbound_context = binding_a
|
||||
if mode == 'legacy':
|
||||
inbound_context = workspace_context().for_installation(binding_a.installation_uuid)
|
||||
setting_a = SimpleNamespace(
|
||||
plugin_author='author-a',
|
||||
plugin_name='plugin-a',
|
||||
installation_uuid=binding_a.installation_uuid,
|
||||
runtime_revision=binding_a.runtime_revision,
|
||||
artifact_digest=binding_a.artifact_digest,
|
||||
)
|
||||
app.persistence_mgr.execute_async.return_value = SimpleNamespace(first=lambda: setting_a)
|
||||
expected = binding_b if mode == 'managed' else connector._legacy_oss_bridge_binding(execution_context)
|
||||
request = ActionRequest.make_request(
|
||||
101,
|
||||
PluginToRuntimeAction.INVOKE_PARSER.value,
|
||||
{'plugin_author': 'author-b', 'plugin_name': 'parser-b', 'storage_path': 'file-a'},
|
||||
inbound_context,
|
||||
)
|
||||
|
||||
await runtime_handler._handle_action(request.model_dump())
|
||||
|
||||
response = json.loads(connection.sent[-1])
|
||||
assert response['code'] == 0, response
|
||||
chunks = connection.requests[:-1]
|
||||
parse = connection.requests[-1]
|
||||
assert len(chunks) == 4
|
||||
assert all(chunk['action'] == CommonAction.FILE_CHUNK.value for chunk in chunks)
|
||||
assert parse['action'] == LangBotToRuntimeAction.PARSE_DOCUMENT.value
|
||||
assert all(request['context'] == expected.model_dump() for request in connection.requests)
|
||||
assert b''.join(base64.b64decode(chunk['data']['chunk_base64']) for chunk in chunks) == file_bytes
|
||||
assert {chunk['data']['file_key'] for chunk in chunks} == {parse['data']['context']['file_key']}
|
||||
connector._setting_for_plugin.assert_awaited_once_with('author-b', 'parser-b', require_enabled=True)
|
||||
assert runtime_handler.current_action_context is None
|
||||
assert runtime_handler.resolve_outbound_action_context(None) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_explicit_argument_overrides_scope_and_inbound_falls_back(bridge):
|
||||
runtime_handler, _, connection, binding_a, binding_b = bridge
|
||||
token = runtime_handler._current_action_context.set(binding_a)
|
||||
try:
|
||||
with runtime_handler.installation_scope(binding_b):
|
||||
await runtime_handler.call_action(
|
||||
LangBotToRuntimeAction.LIST_PARSERS, {}, action_context=binding_a.model_dump()
|
||||
)
|
||||
await runtime_handler.list_parsers()
|
||||
finally:
|
||||
runtime_handler._current_action_context.reset(token)
|
||||
assert [request['context'] for request in connection.requests] == [binding_a.model_dump()] * 2
|
||||
assert runtime_handler.resolve_outbound_action_context(None) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_explicit_none_scope_clears_inbound_and_restores_outer_scope(bridge):
|
||||
runtime_handler, _, connection, binding_a, binding_b = bridge
|
||||
token = runtime_handler._current_action_context.set(binding_a)
|
||||
try:
|
||||
with runtime_handler.installation_scope(binding_b):
|
||||
await runtime_handler.ping()
|
||||
await runtime_handler.list_parsers()
|
||||
await runtime_handler.list_parsers()
|
||||
finally:
|
||||
runtime_handler._current_action_context.reset(token)
|
||||
assert [request.get('context') for request in connection.requests] == [
|
||||
None,
|
||||
binding_b.model_dump(),
|
||||
binding_a.model_dump(),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('failure', [RuntimeError, asyncio.CancelledError])
|
||||
async def test_scope_restores_after_exception_or_cancellation(bridge, failure):
|
||||
runtime_handler, _, connection, binding_a, binding_b = bridge
|
||||
with runtime_handler.installation_scope(binding_a):
|
||||
with pytest.raises(failure):
|
||||
with runtime_handler.installation_scope(binding_b):
|
||||
await runtime_handler.list_parsers()
|
||||
raise failure()
|
||||
await runtime_handler.list_parsers()
|
||||
await runtime_handler.list_parsers()
|
||||
assert [request.get('context') for request in connection.requests] == [
|
||||
binding_b.model_dump(),
|
||||
binding_a.model_dump(),
|
||||
None,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_nested_scopes_do_not_leak_on_task_cancellation(bridge):
|
||||
runtime_handler, _, connection, binding_a, binding_b = bridge
|
||||
entered = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def cancelled_invocation():
|
||||
with runtime_handler.installation_scope(binding_b):
|
||||
await runtime_handler.list_parsers()
|
||||
entered.set()
|
||||
await release.wait()
|
||||
|
||||
token = runtime_handler._current_action_context.set(binding_a)
|
||||
task = asyncio.create_task(cancelled_invocation())
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), timeout=2)
|
||||
with runtime_handler.installation_scope(None):
|
||||
await runtime_handler.list_parsers()
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
await runtime_handler.list_parsers()
|
||||
finally:
|
||||
runtime_handler._current_action_context.reset(token)
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
assert [request.get('context') for request in connection.requests] == [
|
||||
binding_b.model_dump(),
|
||||
None,
|
||||
binding_a.model_dump(),
|
||||
]
|
||||
assert runtime_handler.resolve_outbound_action_context(None) is None
|
||||
Reference in New Issue
Block a user