diff --git a/src/langbot/pkg/box/connector.py b/src/langbot/pkg/box/connector.py index 2e37b84f7..d56270105 100644 --- a/src/langbot/pkg/box/connector.py +++ b/src/langbot/pkg/box/connector.py @@ -152,8 +152,8 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): async with self._lifecycle_lock: if self._closing: raise BoxRuntimeUnavailableError('box runtime connector is shutting down') - await self._stop_transport() self._generation += 1 + await self._stop_transport() try: if self._uses_websocket(): if platform.get_platform() == 'win32' and not self.configured_runtime_endpoint: @@ -174,8 +174,8 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): async with self._lifecycle_lock: if self._closing: raise BoxRuntimeUnavailableError('box runtime connector is shutting down') - await self._stop_transport() self._generation += 1 + await self._stop_transport() try: if self._uses_websocket(): if platform.get_platform() == 'win32' and not self.configured_runtime_endpoint: @@ -235,6 +235,7 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): # mirroring `rt -s`. args=['-m', 'langbot_plugin.cli.__init__', 'box', '-s', '--ws-control-port', str(self._relay_port)], env=env, + capture_stderr=False, ) self._ctrl = ctrl self._ctrl_task = asyncio.create_task( @@ -347,6 +348,22 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): await connection.close() return handler = Handler(connection) + connection_ready = False + disconnect_notified = False + + async def notify_disconnect() -> None: + nonlocal disconnect_notified + if ( + connection_ready + and not disconnect_notified + and generation == self._generation + and not self._closing + and self.runtime_disconnect_callback is not None + ): + disconnect_notified = True + self.ap.logger.error('Disconnected from Box runtime, trying to reconnect...') + await self.runtime_disconnect_callback(self) + self._handler = handler self.client.set_handler(handler) self._handler_task = asyncio.create_task(handler.run()) @@ -356,6 +373,7 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): await handler.call_action(LangBotToBoxAction.INIT, self._filtered_box_config) self.ap.logger.debug('Sent box configuration to Box runtime via INIT.') self.ap.logger.info(f'Connected to Box runtime via {transport_name}.') + connection_ready = True connected.set() await self._handler_task except asyncio.CancelledError: @@ -365,18 +383,11 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): connect_error.append(exc) connected.set() return - - # If we reach here, handler.run() returned normally (connection - # closed) or raised after the initial handshake succeeded. - # Either way, treat it as a disconnect. - if ( - connected.is_set() - and generation == self._generation - and not self._closing - and self.runtime_disconnect_callback is not None - ): - self.ap.logger.error('Disconnected from Box runtime, trying to reconnect...') - await self.runtime_disconnect_callback(self) + finally: + if getattr(self, '_handler', None) is handler: + self._handler = None + self.client.set_handler(None) + await notify_disconnect() return new_connection_callback @@ -407,6 +418,7 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): async def aclose(self) -> None: self._closing = True + self._generation += 1 if self._heartbeat_task is not None: self._heartbeat_task.cancel() await asyncio.gather(self._heartbeat_task, return_exceptions=True) diff --git a/src/langbot/pkg/plugin/connector.py b/src/langbot/pkg/plugin/connector.py index 1b962bd0a..8d2c8597c 100644 --- a/src/langbot/pkg/plugin/connector.py +++ b/src/langbot/pkg/plugin/connector.py @@ -125,9 +125,9 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): if self._connected.is_set() and hasattr(self, 'handler'): return - await self._stop_transport() self._generation += 1 generation = self._generation + await self._stop_transport() self._connected = asyncio.Event() connect_errors: list[Exception] = [] @@ -137,13 +137,25 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): if generation != self._generation or self._closing: await connection.close() return + connection_ready = False + disconnect_notified = False + + async def notify_disconnect() -> None: + nonlocal disconnect_notified + if ( + connection_ready + and not disconnect_notified + and generation == self._generation + and not self._closing + ): + disconnect_notified = True + self._connected.clear() + await self.runtime_disconnect_callback(self) async def disconnect_callback( rchandler: handler.RuntimeConnectionHandler, ) -> bool: - if generation == self._generation and not self._closing: - self._connected.clear() - await self.runtime_disconnect_callback(self) + await notify_disconnect() return False runtime_handler = handler.RuntimeConnectionHandler(connection, disconnect_callback, self.ap) @@ -155,6 +167,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): if space_url: await runtime_handler.set_runtime_config(cloud_service_url=space_url) if generation == self._generation and not self._closing: + connection_ready = True self._connected.set() self.ap.logger.info('Connected to plugin runtime.') await self.handler_task @@ -169,7 +182,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): self._connected.clear() if getattr(self, 'handler', None) is runtime_handler: del self.handler - await self.runtime_disconnect_callback(self) + await notify_disconnect() task_coro: typing.Coroutine if platform.get_platform() == 'docker' or platform.use_websocket_to_connect_plugin_runtime(): @@ -207,6 +220,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): command=sys.executable, args=['-m', 'langbot_plugin.cli.__init__', 'rt', '-s'], env=os.environ.copy(), + capture_stderr=False, ) task_coro = self.ctrl.run(new_connection_callback) @@ -245,11 +259,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): self._reconnect_task = None async def _stop_transport(self) -> None: + self._connected.clear() runtime_handler = getattr(self, 'handler', None) if runtime_handler is not None: with contextlib.suppress(Exception): await runtime_handler.close() - del self.handler + if getattr(self, 'handler', None) is runtime_handler: + del self.handler tasks = [ task for task in ( @@ -272,6 +288,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): async def aclose(self) -> None: self._closing = True + self._generation += 1 reconnect_task = self._reconnect_task self._reconnect_task = None if reconnect_task is not None and reconnect_task is not asyncio.current_task(): diff --git a/tests/unit_tests/box/test_box_connector.py b/tests/unit_tests/box/test_box_connector.py index 0f2f8c340..9cb7d6249 100644 --- a/tests/unit_tests/box/test_box_connector.py +++ b/tests/unit_tests/box/test_box_connector.py @@ -1,11 +1,13 @@ from __future__ import annotations +import asyncio from types import SimpleNamespace from unittest.mock import AsyncMock, Mock import pytest from langbot_plugin.box.client import ActionRPCBoxClient +from langbot.pkg.box import connector as connector_module from langbot.pkg.box.connector import BoxRuntimeConnector @@ -122,3 +124,95 @@ async def test_box_runtime_connector_cleans_partial_transport_on_connect_failure assert connector._stop_transport.await_count == 2 connector._close_managed_subprocess.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_box_stdio_connection_does_not_capture_unconsumed_stderr( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux') + monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False) + created = {} + + class FakeHandler: + def __init__(self, connection): + self.release = asyncio.Event() + + async def call_action(self, action, data): + return None + + async def run(self): + await self.release.wait() + + async def close(self): + self.release.set() + + class FakeController: + def __init__(self, **kwargs): + created.update(kwargs) + self.process = SimpleNamespace(returncode=0) + + async def run(self, callback): + await callback(object()) + + async def close(self): + return None + + monkeypatch.setattr(connector_module, 'Handler', FakeHandler) + monkeypatch.setattr( + 'langbot_plugin.runtime.io.controllers.stdio.client.StdioClientController', + FakeController, + ) + connector = BoxRuntimeConnector(make_app(Mock())) + + await connector.initialize() + + assert created['capture_stderr'] is False + assert connector._handler is not None + await connector.aclose() + + +@pytest.mark.asyncio +async def test_box_disconnect_notifies_once_and_clears_handler( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux') + monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False) + disconnect = AsyncMock() + + class FakeHandler: + def __init__(self, connection): + pass + + async def call_action(self, action, data): + return None + + async def run(self): + return None + + async def close(self): + return None + + class FakeController: + def __init__(self, **kwargs): + self.process = SimpleNamespace(returncode=0) + + async def run(self, callback): + await callback(object()) + + async def close(self): + return None + + monkeypatch.setattr(connector_module, 'Handler', FakeHandler) + monkeypatch.setattr( + 'langbot_plugin.runtime.io.controllers.stdio.client.StdioClientController', + FakeController, + ) + connector = BoxRuntimeConnector(make_app(Mock()), runtime_disconnect_callback=disconnect) + + await connector.initialize() + await asyncio.sleep(0) + + disconnect.assert_awaited_once_with(connector) + assert connector._handler is None + await connector.aclose() diff --git a/tests/unit_tests/plugin/test_connector_ping.py b/tests/unit_tests/plugin/test_connector_ping.py index 766e51f8d..0a8e74207 100644 --- a/tests/unit_tests/plugin/test_connector_ping.py +++ b/tests/unit_tests/plugin/test_connector_ping.py @@ -1,15 +1,20 @@ from __future__ import annotations +import asyncio from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock import pytest +from langbot.pkg.plugin import connector as connector_module from langbot.pkg.plugin.connector import PluginRuntimeConnector, PluginRuntimeNotConnectedError def make_connector() -> PluginRuntimeConnector: - app = SimpleNamespace(instance_config=SimpleNamespace(data={'plugin': {'enable': True}})) + app = SimpleNamespace( + logger=Mock(), + instance_config=SimpleNamespace(data={'plugin': {'enable': True}, 'space': {'url': ''}}), + ) return PluginRuntimeConnector(app, AsyncMock()) @@ -30,3 +35,133 @@ async def test_ping_plugin_runtime_delegates_to_connected_handler(): assert result == 'pong' connector.handler.ping.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_stop_transport_tolerates_handler_callback_removing_attribute(): + connector = make_connector() + + class Handler: + async def close(self): + del connector.handler + + connector.handler = Handler() + + await connector._stop_transport() + + assert not hasattr(connector, 'handler') + + +@pytest.mark.asyncio +async def test_stdio_runtime_connection_does_not_capture_unconsumed_stderr( + monkeypatch: pytest.MonkeyPatch, +): + connector = make_connector() + created = {} + + class FakeRuntimeHandler: + def __init__(self, connection, disconnect_callback, ap): + self.release = asyncio.Event() + + async def ping(self): + return None + + async def set_runtime_config(self, **kwargs): + return None + + async def run(self): + await self.release.wait() + + async def close(self): + self.release.set() + + class FakeController: + def __init__(self, **kwargs): + created.update(kwargs) + + async def run(self, callback): + await callback(object()) + + async def close(self): + return None + + monkeypatch.setattr(connector_module.platform, 'get_platform', lambda: 'linux') + monkeypatch.setattr( + connector_module.platform, + 'use_websocket_to_connect_plugin_runtime', + lambda: False, + ) + monkeypatch.setattr( + connector_module.handler, + 'RuntimeConnectionHandler', + FakeRuntimeHandler, + ) + monkeypatch.setattr( + connector_module.stdio_client_controller, + 'StdioClientController', + FakeController, + ) + + await connector.initialize() + + assert created['capture_stderr'] is False + assert connector._connected.is_set() + await connector.aclose() + + +@pytest.mark.asyncio +async def test_runtime_disconnect_notifies_once_and_clears_handler( + monkeypatch: pytest.MonkeyPatch, +): + disconnect = AsyncMock() + connector = PluginRuntimeConnector(make_connector().ap, disconnect) + + class FakeRuntimeHandler: + def __init__(self, connection, disconnect_callback, ap): + self.disconnect_callback = disconnect_callback + + async def ping(self): + return None + + async def set_runtime_config(self, **kwargs): + return None + + async def run(self): + await self.disconnect_callback(self) + + async def close(self): + return None + + class FakeController: + def __init__(self, **kwargs): + pass + + async def run(self, callback): + await callback(object()) + + async def close(self): + return None + + monkeypatch.setattr(connector_module.platform, 'get_platform', lambda: 'linux') + monkeypatch.setattr( + connector_module.platform, + 'use_websocket_to_connect_plugin_runtime', + lambda: False, + ) + monkeypatch.setattr( + connector_module.handler, + 'RuntimeConnectionHandler', + FakeRuntimeHandler, + ) + monkeypatch.setattr( + connector_module.stdio_client_controller, + 'StdioClientController', + FakeController, + ) + + await connector.initialize() + await asyncio.sleep(0) + + disconnect.assert_awaited_once_with(connector) + assert not hasattr(connector, 'handler') + await connector.aclose()