fix(runtime): harden connector lifecycle

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
-
+            finally:
-            # If we reach here, handler.run() returned normally (connection
+                if getattr(self, '_handler', None) is handler:
-            # closed) or raised after the initial handshake succeeded.
+                    self._handler = None
-            # Either way, treat it as a disconnect.
+                    self.client.set_handler(None)
-            if (
+                await notify_disconnect()
-                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)
 
         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)

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:
+                    await notify_disconnect()
-                        self._connected.clear()
-                        await self.runtime_disconnect_callback(self)
                     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():

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()

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()