Revert "fix(runtime): make plugin and box connectors resilient"

This reverts commit 0b461e5830.

src/langbot/pkg/box/connector.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import asyncio
-import contextlib
 import json
 import os
 import sys
@@ -29,7 +28,6 @@ _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.
@@ -115,8 +113,6 @@ 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)
@@ -149,64 +145,29 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
         return self.uses_websocket()
 
     async def initialize(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
+        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()
 
-            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
+        # Start heartbeat after successful connection
+        if self._heartbeat_task is None:
+            self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
 
     # -- heartbeat -----------------------------------------------------------
 
     async def _heartbeat_loop(self) -> None:
         """Periodically ping the Box runtime to detect silent disconnections."""
-        failures = 0
-        while not self._closing:
+        while True:
             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:
-                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)
+                self.ap.logger.debug(f'Failed to heartbeat to Box runtime: {e}')
 
     async def ping(self) -> None:
         if self._handler is None:
@@ -236,9 +197,8 @@ 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, self._generation))
+            ctrl.run(self._make_connection_callback('stdio', connected, connect_error))
         )
 
         try:
@@ -322,9 +282,8 @@ 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, self._generation))
+            ctrl.run(self._make_connection_callback(transport_name, connected, connect_error))
         )
 
         try:
@@ -340,12 +299,8 @@ 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)
@@ -358,8 +313,6 @@ 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)
@@ -369,66 +322,22 @@ 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()
-                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)
+            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.'
+                    )
 
         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,7 +2,6 @@ from __future__ import annotations
 
 import asyncio
 import collections
-import contextlib
 import datetime as _dt
 import enum
 import json
@@ -66,8 +65,6 @@ 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
@@ -113,8 +110,6 @@ 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.
@@ -123,26 +118,27 @@ class BoxService:
         Skipped entirely when Box is disabled by config — that path should
         never have connected in the first place.
         """
-        if not self._enabled or self._closing:
+        if not self._enabled:
             return
-        if self._reconnect_task is not None and not self._reconnect_task.done():
+        if self._reconnecting:
             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.')
-        self._reconnect_task = asyncio.create_task(self._reconnect_loop(connector))
+        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 not self._closing:
+            while True:
                 self.ap.logger.info(f'Attempting to reconnect to Box runtime in {delay}s...')
                 await asyncio.sleep(delay)
                 try:
-                    await connector.reconnect()
+                    connector.dispose()
+                    await connector.initialize()
                     self._available = True
                     self._connector_error = ''
                     self.ap.logger.info('Box runtime reconnected, sandbox features restored.')
@@ -153,7 +149,6 @@ class BoxService:
                     delay = min(delay * 2, max_delay)
         finally:
             self._reconnecting = False
-            self._reconnect_task = None
 
     @property
     def available(self) -> bool:
@@ -843,28 +838,14 @@ class BoxService:
         return attachments
 
     async def shutdown(self):
-        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()
+        await self.client.shutdown()
 
     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: