From d497defbf77bb98e57e80852502ed6d0067e5b5a Mon Sep 17 00:00:00 2001 From: Hyu Date: Tue, 15 Sep 2026 13:51:27 +0800 Subject: [PATCH] fix(runner): keep installation context within each stream resume --- src/langbot/pkg/agent/runner/invoker.py | 17 +- src/langbot/pkg/plugin/connector.py | 42 +++-- src/langbot/pkg/plugin/handler.py | 10 +- .../plugin/test_runner_stream_context.py | 158 ++++++++++++++++++ 4 files changed, 203 insertions(+), 24 deletions(-) create mode 100644 tests/unit_tests/plugin/test_runner_stream_context.py diff --git a/src/langbot/pkg/agent/runner/invoker.py b/src/langbot/pkg/agent/runner/invoker.py index 00f521ff9..c4e54307b 100644 --- a/src/langbot/pkg/agent/runner/invoker.py +++ b/src/langbot/pkg/agent/runner/invoker.py @@ -48,12 +48,15 @@ class RunnerInvoker: context=context, ) - while True: - try: - result_dict = await self._next_with_deadline(gen, descriptor, context) - except StopAsyncIteration: - break - yield result_dict + try: + while True: + try: + result_dict = await self._next_with_deadline(gen, descriptor, context) + except StopAsyncIteration: + break + yield result_dict + finally: + await self._close_generator(gen, descriptor) except asyncio.TimeoutError as e: raise RunnerExecutionError( @@ -128,4 +131,4 @@ class RunnerInvoker: try: await gen.aclose() except Exception as e: - self.ap.logger.warning(f'Failed to close timed-out runner {descriptor.id}: {e}') + self.ap.logger.warning(f'Failed to close runner {descriptor.id}: {e}') diff --git a/src/langbot/pkg/plugin/connector.py b/src/langbot/pkg/plugin/connector.py index 2b1cb4447..a2eebbb52 100644 --- a/src/langbot/pkg/plugin/connector.py +++ b/src/langbot/pkg/plugin/connector.py @@ -2237,12 +2237,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): include_plugins=bound_plugins, ) runtime_handler = self._runtime_handler() - with runtime_handler.installation_scope(binding): - gen = runtime_handler.execute_command( - command_ctx.model_dump(serialize_as_any=True), - include_plugins=bound_plugins, - ) - async for ret in gen: + gen = runtime_handler.execute_command( + command_ctx.model_dump(serialize_as_any=True), + include_plugins=bound_plugins, + ) + async with contextlib.aclosing(self._installation_scoped_stream(runtime_handler, binding, gen)) as scoped: + async for ret in scoped: yield command_context.CommandReturn.model_validate(ret) # Runner methods @@ -2308,15 +2308,31 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): require_enabled=True, ) runtime_handler = self._runtime_handler() - with runtime_handler.installation_scope(binding): - async for ret in runtime_handler.run_runner( - plugin_author, - plugin_name, - runner_name, - context, - ): + gen = runtime_handler.run_runner(plugin_author, plugin_name, runner_name, context) + async with contextlib.aclosing(self._installation_scoped_stream(runtime_handler, binding, gen)) as scoped: + async for ret in scoped: yield ret + @staticmethod + async def _installation_scoped_stream(runtime_handler, binding, gen): + """Keep ContextVar tokens inside a single resume, never across yields. + + Consumers may use a different Task for each anext (e.g. wait_for). + Reset the installation before exposing a result to the consumer, and + re-enter the same immutable binding for transport cleanup. + """ + try: + while True: + with runtime_handler.installation_scope(binding): + try: + result = await anext(gen) + except StopAsyncIteration: + return + yield result + finally: + with runtime_handler.installation_scope(binding): + await gen.aclose() + async def retrieve_knowledge( self, plugin_author: str, diff --git a/src/langbot/pkg/plugin/handler.py b/src/langbot/pkg/plugin/handler.py index 9d9c55cc7..ddc784acb 100644 --- a/src/langbot/pkg/plugin/handler.py +++ b/src/langbot/pkg/plugin/handler.py @@ -2933,8 +2933,9 @@ class RuntimeConnectionHandler(handler.Handler): timeout=timeout, ) - async for ret in gen: - yield ret + async with contextlib.aclosing(gen): + async for ret in gen: + yield ret def _get_runner_action_timeout(self, context: dict[str, Any]) -> float: """Use the run deadline as the transport idle timeout when available.""" @@ -3143,8 +3144,9 @@ class RuntimeConnectionHandler(handler.Handler): timeout=180, ) - async for ret in gen: - yield ret + async with contextlib.aclosing(gen): + async for ret in gen: + yield ret async def retrieve_knowledge( self, diff --git a/tests/unit_tests/plugin/test_runner_stream_context.py b/tests/unit_tests/plugin/test_runner_stream_context.py new file mode 100644 index 000000000..76df1bad8 --- /dev/null +++ b/tests/unit_tests/plugin/test_runner_stream_context.py @@ -0,0 +1,158 @@ +"""Regression coverage for installation scopes across async-generator resumes.""" + +import asyncio +import contextvars +import time +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from langbot.pkg.agent.runner.errors import RunnerExecutionError +from langbot.pkg.agent.runner.invoker import RunnerInvoker +from langbot.pkg.plugin.handler import RuntimeConnectionHandler +from langbot.pkg.plugin.connector import PluginRuntimeConnector +from tests.unit_tests.plugin.test_connector_methods import ( + TEST_EXECUTION_CONTEXT, + TEST_INSTALLATION_BINDING, + create_mock_connector, +) + + +def make_stream(*, deadline=True, failure=None, blocked=False): + handler = object.__new__(RuntimeConnectionHandler) + handler._outbound_installation_context = contextvars.ContextVar('test_installation', default=None) + entered = asyncio.Event() + closed = [] + observed = [] + + async def wire_stream(*args, **kwargs): + try: + for index in range(3): + observed.append(handler._outbound_installation_context.get()) + if index == 1: + entered.set() + if blocked: + await asyncio.Event().wait() + if failure: + raise failure + yield {'type': 'message.delta', 'sequence': index} + finally: + closed.append(handler._outbound_installation_context.get()) + + handler.call_action_generator = wire_stream + connector = create_mock_connector() + connector.handler = handler + invoker = RunnerInvoker(SimpleNamespace(plugin_connector=connector, logger=Mock())) + descriptor = SimpleNamespace( + id='plugin:qa/runner/default', plugin_author='qa', plugin_name='runner', runner_name='default' + ) + context = { + 'conversation': {'workspace_id': TEST_EXECUTION_CONTEXT.workspace_uuid}, + 'runtime': {'deadline_at': time.time() + 10 if deadline else None}, + } + return invoker.invoke(descriptor, context), handler, context, entered, closed, observed + + +@pytest.mark.asyncio +@pytest.mark.parametrize('deadline', [True, False]) +async def test_scope_is_reset_before_yield_and_stream_finishes(deadline): + stream, handler, _, _, closed, observed = make_stream(deadline=deadline) + frames = [] + async for frame in stream: + frames.append(frame) + assert handler._outbound_installation_context.get() is None + await asyncio.sleep(0) + assert [frame['sequence'] for frame in frames] == [0, 1, 2] + assert observed == [TEST_INSTALLATION_BINDING] * 3 + assert closed == [TEST_INSTALLATION_BINDING] + + +@pytest.mark.asyncio +@pytest.mark.parametrize('deadline', [True, False]) +async def test_early_close_releases_wire_stream_in_scope(deadline): + stream, handler, _, _, closed, _ = make_stream(deadline=deadline) + await anext(stream) + await stream.aclose() + assert closed == [TEST_INSTALLATION_BINDING] + assert handler._outbound_installation_context.get() is None + + +@pytest.mark.asyncio +async def test_deadline_expired_between_frames_closes_in_scope(): + stream, handler, context, _, closed, _ = make_stream() + await anext(stream) + context['runtime']['deadline_at'] = time.time() - 1 + with pytest.raises(RunnerExecutionError) as exc: + await anext(stream) + assert exc.value.error_code == 'runner.timeout' + assert closed == [TEST_INSTALLATION_BINDING] + assert handler._outbound_installation_context.get() is None + + +@pytest.mark.asyncio +async def test_timeout_during_next_frame_keeps_timeout_error(): + stream, handler, context, _, closed, _ = make_stream(blocked=True) + await anext(stream) + context['runtime']['deadline_at'] = time.time() + 0.05 + with pytest.raises(RunnerExecutionError) as exc: + await anext(stream) + assert exc.value.error_code == 'runner.timeout' + assert closed == [TEST_INSTALLATION_BINDING] + assert handler._outbound_installation_context.get() is None + + +@pytest.mark.asyncio +async def test_cancellation_preserves_cancelled_error_and_cleans_up(): + stream, handler, _, entered, closed, _ = make_stream(blocked=True) + await anext(stream) + task = asyncio.create_task(anext(stream)) + await asyncio.wait_for(entered.wait(), 1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert closed == [TEST_INSTALLATION_BINDING] + assert handler._outbound_installation_context.get() is None + + +@pytest.mark.asyncio +async def test_transport_error_is_not_masked_by_context_reset(): + stream, handler, _, _, closed, _ = make_stream(failure=RuntimeError('wire failed')) + await anext(stream) + with pytest.raises(RunnerExecutionError, match='wire failed'): + await anext(stream) + assert closed == [TEST_INSTALLATION_BINDING] + assert handler._outbound_installation_context.get() is None + + +@pytest.mark.asyncio +async def test_interleaved_installations_share_handler_without_scope_leakage(): + handler = object.__new__(RuntimeConnectionHandler) + handler._outbound_installation_context = contextvars.ContextVar('shared_installation', default=None) + other_binding = TEST_INSTALLATION_BINDING.model_copy( + update={'installation_uuid': '00000000-0000-4000-8000-000000000002'} + ) + closed = [] + + async def wire(binding): + try: + for index in range(3): + assert handler._outbound_installation_context.get() == binding + await asyncio.sleep(0) + assert handler._outbound_installation_context.get() == binding + yield index + finally: + assert handler._outbound_installation_context.get() == binding + closed.append(binding) + + async def consume(binding): + stream = PluginRuntimeConnector._installation_scoped_stream(handler, binding, wire(binding)) + for index in range(3): + assert await asyncio.wait_for(anext(stream), 1) == index + assert handler._outbound_installation_context.get() is None + # Closing in yet another task must also restore that task's context. + await asyncio.wait_for(stream.aclose(), 1) + + await asyncio.gather(consume(TEST_INSTALLATION_BINDING), consume(other_binding)) + assert set(closed) == {TEST_INSTALLATION_BINDING, other_binding} + assert handler._outbound_installation_context.get() is None