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

src/langbot/pkg/box/service.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import collections
+import contextlib
 import datetime as _dt
 import enum
 import json
@@ -65,6 +66,8 @@ class BoxService:
         self.workspace_quota_mb = self._load_workspace_quota_mb()
         self._recent_errors: collections.deque[dict] = collections.deque(maxlen=_MAX_RECENT_ERRORS)
         self._shutdown_task = None
+        self._reconnect_task: asyncio.Task | None = None
+        self._closing = False
         self._available = False
         self._connector_error: str = ''
         self._reconnecting = False
@@ -110,6 +113,8 @@ class BoxService:
             self.ap.logger.warning(f'LangBot Box runtime unavailable, sandbox features disabled: {exc}')
             self._available = False
             self._connector_error = str(exc)
+            if self._runtime_connector is not None:
+                await self._on_runtime_disconnect(self._runtime_connector)
 
     async def _on_runtime_disconnect(self, connector: BoxRuntimeConnector) -> None:
         """Called by the connector when the Box runtime connection drops.
@@ -118,27 +123,26 @@ class BoxService:
         Skipped entirely when Box is disabled by config — that path should
         never have connected in the first place.
         """
-        if not self._enabled:
+        if not self._enabled or self._closing:
             return
-        if self._reconnecting:
+        if self._reconnect_task is not None and not self._reconnect_task.done():
             return  # Another reconnect loop is already running
         self._reconnecting = True
         self._available = False
         self._connector_error = 'Disconnected from Box runtime'
         self.ap.logger.warning('Box runtime disconnected, sandbox features temporarily disabled.')
-        asyncio.create_task(self._reconnect_loop(connector))
+        self._reconnect_task = asyncio.create_task(self._reconnect_loop(connector))
 
     async def _reconnect_loop(self, connector: BoxRuntimeConnector) -> None:
         """Retry reconnection with exponential backoff (3s → 60s max)."""
         delay = 3
         max_delay = 60
         try:
-            while True:
+            while not self._closing:
                 self.ap.logger.info(f'Attempting to reconnect to Box runtime in {delay}s...')
                 await asyncio.sleep(delay)
                 try:
-                    connector.dispose()
-                    await connector.initialize()
+                    await connector.reconnect()
                     self._available = True
                     self._connector_error = ''
                     self.ap.logger.info('Box runtime reconnected, sandbox features restored.')
@@ -149,6 +153,7 @@ class BoxService:
                     delay = min(delay * 2, max_delay)
         finally:
             self._reconnecting = False
+            self._reconnect_task = None
 
     @property
     def available(self) -> bool:
@@ -838,14 +843,28 @@ class BoxService:
         return attachments
 
     async def shutdown(self):
-        await self.client.shutdown()
+        if self._closing:
+            return
+        self._closing = True
+        self._available = False
+        reconnect_task = self._reconnect_task
+        self._reconnect_task = None
+        if reconnect_task is not None and reconnect_task is not asyncio.current_task():
+            reconnect_task.cancel()
+            await asyncio.gather(reconnect_task, return_exceptions=True)
+        # The runtime may already be offline. A failed best-effort SHUTDOWN RPC
+        # must not prevent us from cancelling transports and reaping children.
+        with contextlib.suppress(Exception):
+            await self.client.shutdown()
+        if self._runtime_connector is not None:
+            await self._runtime_connector.aclose()
 
     def dispose(self):
-        if self._runtime_connector is not None:
-            self._runtime_connector.dispose()
         loop = getattr(self.ap, 'event_loop', None)
         if loop is not None and not loop.is_closed() and (self._shutdown_task is None or self._shutdown_task.done()):
             self._shutdown_task = loop.create_task(self.shutdown())
+        elif self._runtime_connector is not None:
+            self._runtime_connector.dispose()
 
     async def get_sessions(self) -> list[dict]:
         if not self._available: