From 7865ad51f7bf3561059d818b658d761fb3285d29 Mon Sep 17 00:00:00 2001 From: huanghuoguoguo <1051233107@qq.com> Date: Tue, 27 Oct 2026 11:36:47 +0800 Subject: [PATCH] perf(plugin): negotiate ordered batch event dispatch (#2778) --- src/langbot/pkg/plugin/connector.py | 12 ++- src/langbot/pkg/plugin/handler.py | 25 +++--- .../plugin/test_connector_methods.py | 77 +++++++++++++++++++ .../unit_tests/plugin/test_handler_tenancy.py | 46 ++++++++++- 4 files changed, 145 insertions(+), 15 deletions(-) diff --git a/src/langbot/pkg/plugin/connector.py b/src/langbot/pkg/plugin/connector.py index 706d9e1b8..b880e86da 100644 --- a/src/langbot/pkg/plugin/connector.py +++ b/src/langbot/pkg/plugin/connector.py @@ -2368,7 +2368,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): _execution_context, settings = await self._prepare_workspace_context(get_query_execution_context(query)) event_ctx = context.EventContext.from_event(event) - if not self.is_enable_plugin or not self._runtime_available(): + if (bound_plugins == [] and query is not None) or not self.is_enable_plugin or not self._runtime_available(): event_ctx._emitted_plugins = [] event_ctx._response_sources = [] return event_ctx @@ -2377,16 +2377,22 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): emitted_plugins: list[Any] = [] response_sources: list[dict[str, Any]] = [] # Reuse the settings read under the reconciliation lock for this dispatch. - for binding in await self._operation_bindings(include_plugins=bound_plugins, settings=settings): + bindings = await self._operation_bindings(include_plugins=bound_plugins, settings=settings) + batch = ( + bindings if len(bindings) > 1 and getattr(runtime_handler, 'supports_event_batch', False) is True else None + ) + batch_args = {'installation_bindings': batch} if batch is not None else {} + for binding in bindings: with runtime_handler.installation_scope(binding): result = await runtime_handler.emit_event( event_ctx.model_dump(serialize_as_any=False), include_plugins=bound_plugins, + **batch_args, ) event_ctx = context.EventContext.model_validate(result['event_context']) emitted_plugins.extend(result.get('emitted_plugins', [])) response_sources.extend(result.get('response_sources', [])) - if event_ctx.is_prevented_postorder(): + if batch is not None or event_ctx.is_prevented_postorder(): break if query is not None: event_ctx.event.query = query diff --git a/src/langbot/pkg/plugin/handler.py b/src/langbot/pkg/plugin/handler.py index f0253c4e3..0b87218d0 100644 --- a/src/langbot/pkg/plugin/handler.py +++ b/src/langbot/pkg/plugin/handler.py @@ -894,6 +894,7 @@ class RuntimeConnectionHandler(handler.Handler): ): super().__init__(connection, disconnect_callback) self.ap = ap + self.supports_event_batch = False self._outbound_installation_context: contextvars.ContextVar[InstallationBinding | None | object] = ( contextvars.ContextVar( f'{self.__class__.__name__}_{id(self)}_outbound_installation', @@ -2801,12 +2802,15 @@ class RuntimeConnectionHandler(handler.Handler): runtime_profile=runtime_profile, cloud_service_url=cloud_service_url, ) + self.supports_event_batch = False with self.installation_scope(None): - return await self.call_action( + result = await self.call_action( LangBotToRuntimeAction.SET_RUNTIME_CONFIG, runtime_config.model_dump(exclude_none=True), timeout=10, ) + self.supports_event_batch = result.get('event_batch') is True + return result async def get_storage_analysis(self) -> dict[str, Any]: """Return instance-scoped filesystem usage measured by Plugin Runtime.""" @@ -2960,19 +2964,20 @@ class RuntimeConnectionHandler(handler.Handler): self, event_context: dict[str, Any], include_plugins: list[str] | None = None, + *, + installation_bindings: list[InstallationBinding] | None = None, ) -> dict[str, Any]: - """Emit event""" - result = await self.call_action( + """Emit one event, optionally across an ordered negotiated installation batch.""" + data = {'event_context': event_context, 'include_plugins': include_plugins} + if installation_bindings is not None: + data['installation_bindings'] = [binding.model_dump() for binding in installation_bindings] + return await self.call_action( LangBotToRuntimeAction.EMIT_EVENT, - { - 'event_context': event_context, - 'include_plugins': include_plugins, - }, - timeout=180, + data, + timeout=180 * (len(installation_bindings) if installation_bindings else 1), + cancel_peer_on_cancel=installation_bindings is not None, ) - return result - async def notify_plugin_diagnostic(self, diagnostic: dict[str, Any]) -> dict[str, Any]: """Notify the plugin runtime about a best-effort plugin diagnostic. diff --git a/tests/unit_tests/plugin/test_connector_methods.py b/tests/unit_tests/plugin/test_connector_methods.py index 82450a033..6b5da083c 100644 --- a/tests/unit_tests/plugin/test_connector_methods.py +++ b/tests/unit_tests/plugin/test_connector_methods.py @@ -822,3 +822,80 @@ class TestPingPluginRuntime: await connector.ping_plugin_runtime() connector.handler.ping.assert_called_once() + + +@pytest.mark.asyncio +async def test_empty_event_selection_keeps_workspace_preparation_and_query(): + from langbot_plugin.api.entities.events import PersonMessageReceived + + connector = create_mock_connector() + runtime_handler = configure_handler(connector, Mock()) + runtime_handler.emit_event = AsyncMock() + query = text_query('hello') + event = PersonMessageReceived( + query=query, + launcher_type=query.launcher_type.value, + launcher_id=query.launcher_id, + sender_id=query.sender_id, + message_event=query.message_event, + message_chain=query.message_chain, + ) + + returned = await connector.emit_event(event, bound_plugins=[]) + + connector._prepare_workspace_context.assert_awaited_once() + connector._operation_bindings.assert_not_awaited() + runtime_handler.emit_event.assert_not_awaited() + assert returned.event.query is query + assert returned._emitted_plugins == [] + assert returned._response_sources == [] + assert not returned.is_prevented_default() + assert not returned.is_prevented_postorder() + + connector._prepare_workspace_context.side_effect = WorkspaceNotFoundError('Plugin resource not found') + with pytest.raises(WorkspaceNotFoundError): + await connector.emit_event(event, bound_plugins=[]) + runtime_handler.emit_event.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('batch_support', [True, False, None]) +async def test_event_dispatch_negotiates_batch_and_retains_query(batch_support): + from langbot_plugin.api.entities.events import PersonMessageReceived + + connector = create_mock_connector() + runtime_handler = configure_handler(connector, Mock()) + runtime_handler.supports_event_batch = batch_support + bindings = [ + TEST_INSTALLATION_BINDING, + TEST_INSTALLATION_BINDING.model_copy(update={'installation_uuid': '00000000-0000-4000-8000-000000000002'}), + ] + connector._operation_bindings.return_value = bindings + query = text_query('hello') + event = PersonMessageReceived( + query=query, + launcher_type=query.launcher_type.value, + launcher_id=query.launcher_id, + sender_id=query.sender_id, + message_event=query.message_event, + message_chain=query.message_chain, + ) + + async def emit(event_context, include_plugins=None, **kwargs): + if batch_support is True: + assert kwargs == {'installation_bindings': bindings} + else: + assert kwargs == {} + return { + 'event_context': event_context, + 'emitted_plugins': ['listener'], + 'response_sources': [{'kind': 'reply'}], + } + + runtime_handler.emit_event = AsyncMock(side_effect=emit) + returned = await connector.emit_event(event) + count = 1 if batch_support is True else 2 + assert runtime_handler.emit_event.await_count == count + assert returned.event.query is query + assert returned._emitted_plugins == ['listener'] * count + assert returned._response_sources == [{'kind': 'reply'}] * count diff --git a/tests/unit_tests/plugin/test_handler_tenancy.py b/tests/unit_tests/plugin/test_handler_tenancy.py index 263e3ac46..370ab4cb7 100644 --- a/tests/unit_tests/plugin/test_handler_tenancy.py +++ b/tests/unit_tests/plugin/test_handler_tenancy.py @@ -467,7 +467,8 @@ 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(): +@pytest.mark.parametrize('capability', [None, True, False, 1]) +async def test_host_to_runtime_action_carries_trusted_connector_context(capability): app = SimpleNamespace(logger=Mock()) connection = RecordingConnection() runtime_handler = RuntimeConnectionHandler( @@ -495,8 +496,11 @@ async def test_host_to_runtime_action_carries_trusted_connector_context(): break await asyncio.sleep(0) request = json.loads(connection.sent[0]) - runtime_handler.resp_waiters[request['seq_id']].set_result(ActionResponse.success({})) + runtime_handler.resp_waiters[request['seq_id']].set_result( + ActionResponse.success({} if capability is None else {'event_batch': capability}) + ) await task + assert runtime_handler.supports_event_batch is (capability is True) assert request['data']['runtime_identity'] == { 'instance_uuid': 'instance-a', @@ -523,3 +527,41 @@ async def test_get_debug_info_converts_execution_context_to_sdk_action_context() workspace_uuid='workspace-a', placement_generation=7, ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize('failure', ['cancel', 'timeout']) +async def test_event_batch_cancels_remote_dispatch_and_cleans_waiters(failure): + from langbot_plugin.entities.io.errors import ActionCallTimeoutError + + runtime_handler, _, binding = make_handler() + connection = RecordingConnection() + runtime_handler.conn = connection + second = binding.model_copy(update={'installation_uuid': '00000000-0000-4000-8000-000000000002'}) + call_action = runtime_handler.call_action + + async def short_call(action, data, **kwargs): + if action.value == 'emit_event' and failure == 'timeout': + kwargs['timeout'] = 0.01 + return await call_action(action, data, **kwargs) + + runtime_handler.call_action = short_call + with runtime_handler.installation_scope(binding): + task = asyncio.create_task(runtime_handler.emit_event({}, installation_bindings=[binding, second])) + while not connection.sent: + await asyncio.sleep(0) + original = json.loads(connection.sent[0]) + if failure == 'cancel': + task.cancel() + async with asyncio.timeout(2): + while len(connection.sent) < 2: + await asyncio.sleep(0) + cancellation = json.loads(connection.sent[1]) + assert cancellation['action'] == '__cancel_action' + assert cancellation['data'] == {'seq_id': original['seq_id']} + assert cancellation['context'] == original['context'] == binding.model_dump() + runtime_handler.resp_waiters[cancellation['seq_id']].set_result(ActionResponse.success({'accepted': True})) + with pytest.raises(asyncio.CancelledError if failure == 'cancel' else ActionCallTimeoutError): + await task + assert runtime_handler.resp_waiters == {} + assert runtime_handler.resp_queues == {}