fix(runtime): make plugin and box connectors resilient

src/langbot/pkg/box/connector.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import json
 import os
 import sys
@@ -28,6 +29,7 @@ _DOCKER_BOX_HOST = 'langbot_box'
 _DEFAULT_PORT = 5410
 
 _HEARTBEAT_INTERVAL_SEC = 20
+_HEARTBEAT_FAILURE_THRESHOLD = 3
 
 # Top-level keys under ``box`` that are LangBot-internal and should not be
 # forwarded to the Box runtime.
@@ -113,6 +115,8 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
         self._handler_task: asyncio.Task | None = None
         self._ctrl_task: asyncio.Task | None = None
         self._heartbeat_task: asyncio.Task | None = None
+        self._ctrl = None
+        self._generation = 0
 
         # Parse the relay URL once for reuse.
         parsed = urlparse(self.ws_relay_base_url)
@@ -145,29 +149,64 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
         return self.uses_websocket()
 
     async def initialize(self) -> None:
-        if self._uses_websocket():
-            if platform.get_platform() == 'win32' and not self.configured_runtime_endpoint:
-                await self._start_subprocess_then_ws()
-            else:
-                await self._connect_remote_ws()
-        else:
-            await self._start_local_stdio()
+        async with self._lifecycle_lock:
+            if self._closing:
+                raise BoxRuntimeUnavailableError('box runtime connector is shutting down')
+            await self._stop_transport()
+            self._generation += 1
+            try:
+                if self._uses_websocket():
+                    if platform.get_platform() == 'win32' and not self.configured_runtime_endpoint:
+                        await self._start_subprocess_then_ws()
+                    else:
+                        await self._connect_remote_ws()
+                else:
+                    await self._start_local_stdio()
+            except BaseException:
+                await self._stop_transport()
+                await self._close_managed_subprocess()
+                raise
 
-        # Start heartbeat after successful connection
-        if self._heartbeat_task is None:
-            self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
+            if self._heartbeat_task is None or self._heartbeat_task.done():
+                self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
+
+    async def reconnect(self) -> None:
+        async with self._lifecycle_lock:
+            if self._closing:
+                raise BoxRuntimeUnavailableError('box runtime connector is shutting down')
+            await self._stop_transport()
+            self._generation += 1
+            try:
+                if self._uses_websocket():
+                    if platform.get_platform() == 'win32' and not self.configured_runtime_endpoint:
+                        await self._start_subprocess_then_ws()
+                    else:
+                        await self._connect_remote_ws()
+                else:
+                    await self._start_local_stdio()
+            except BaseException:
+                await self._stop_transport()
+                await self._close_managed_subprocess()
+                raise
 
     # -- heartbeat -----------------------------------------------------------
 
     async def _heartbeat_loop(self) -> None:
         """Periodically ping the Box runtime to detect silent disconnections."""
-        while True:
+        failures = 0
+        while not self._closing:
             await asyncio.sleep(_HEARTBEAT_INTERVAL_SEC)
             try:
                 await self.ping()
+                failures = 0
                 self.ap.logger.debug('Heartbeat to Box runtime success.')
             except Exception as e:
-                self.ap.logger.debug(f'Failed to heartbeat to Box runtime: {e}')
+                failures += 1
+                self.ap.logger.warning(f'Box runtime heartbeat failed ({failures}/{_HEARTBEAT_FAILURE_THRESHOLD}): {e}')
+                if failures >= _HEARTBEAT_FAILURE_THRESHOLD:
+                    failures = 0
+                    if self.runtime_disconnect_callback is not None:
+                        await self.runtime_disconnect_callback(self)
 
     async def ping(self) -> None:
         if self._handler is None:
@@ -197,8 +236,9 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
             args=['-m', 'langbot_plugin.cli.__init__', 'box', '-s', '--ws-control-port', str(self._relay_port)],
             env=env,
         )
+        self._ctrl = ctrl
         self._ctrl_task = asyncio.create_task(
-            ctrl.run(self._make_connection_callback('stdio', connected, connect_error))
+            ctrl.run(self._make_connection_callback('stdio', connected, connect_error, self._generation))
         )
 
         try:
@@ -282,8 +322,9 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
                 await self.runtime_disconnect_callback(self)
 
         ctrl = WebSocketClientController(ws_url=ws_url, make_connection_failed_callback=on_connect_failed)
+        self._ctrl = ctrl
         self._ctrl_task = asyncio.create_task(
-            ctrl.run(self._make_connection_callback(transport_name, connected, connect_error))
+            ctrl.run(self._make_connection_callback(transport_name, connected, connect_error, self._generation))
         )
 
         try:
@@ -299,8 +340,12 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
         transport_name: str,
         connected: asyncio.Event,
         connect_error: list[Exception],
+        generation: int,
     ):
         async def new_connection_callback(connection: Connection) -> None:
+            if generation != self._generation or self._closing:
+                await connection.close()
+                return
             handler = Handler(connection)
             self._handler = handler
             self.client.set_handler(handler)
@@ -313,6 +358,8 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
                 self.ap.logger.info(f'Connected to Box runtime via {transport_name}.')
                 connected.set()
                 await self._handler_task
+            except asyncio.CancelledError:
+                raise
             except Exception as exc:
                 if not connected.is_set():
                     connect_error.append(exc)
@@ -322,22 +369,66 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
             # 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():
-                if self._uses_websocket():
-                    self.ap.logger.error('Disconnected from Box runtime, trying to reconnect...')
-                    if self.runtime_disconnect_callback is not None:
-                        await self.runtime_disconnect_callback(self)
-                else:
-                    self.ap.logger.error(
-                        'Disconnected from Box runtime via stdio. '
-                        'Cannot automatically reconnect — please restart LangBot.'
-                    )
+            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)
 
         return new_connection_callback
 
     # -- lifecycle -----------------------------------------------------------
 
+    async def _stop_transport(self) -> None:
+        if self._handler is not None:
+            with contextlib.suppress(Exception):
+                await self._handler.close()
+        self.client.set_handler(None)
+        tasks = [
+            task
+            for task in (self._handler_task, self._ctrl_task)
+            if task is not None and task is not asyncio.current_task()
+        ]
+        for task in tasks:
+            task.cancel()
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+        close_ctrl = getattr(self._ctrl, 'close', None)
+        if close_ctrl is not None:
+            with contextlib.suppress(Exception):
+                await close_ctrl()
+        self._handler = None
+        self._handler_task = None
+        self._ctrl_task = None
+        self._ctrl = None
+
+    async def aclose(self) -> None:
+        self._closing = True
+        if self._heartbeat_task is not None:
+            self._heartbeat_task.cancel()
+            await asyncio.gather(self._heartbeat_task, return_exceptions=True)
+            self._heartbeat_task = None
+        await self._stop_transport()
+
+        process = getattr(self, '_subprocess', None)
+        if process is not None and process.returncode is None:
+            with contextlib.suppress(ProcessLookupError):
+                process.terminate()
+            try:
+                await asyncio.wait_for(process.wait(), timeout=3)
+            except asyncio.TimeoutError:
+                with contextlib.suppress(ProcessLookupError):
+                    process.kill()
+                await process.wait()
+        self._subprocess = None
+        await self._close_managed_subprocess()
+
     def dispose(self) -> None:
+        """Best-effort synchronous compatibility wrapper; prefer ``aclose``."""
+        self._closing = True
         if self._heartbeat_task is not None:
             self._heartbeat_task.cancel()
             self._heartbeat_task = None