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

This reverts commit 0b461e5830.

docker/kubernetes.yaml

@@ -139,8 +139,7 @@ spec:
             cpu: "1000m"
         # Liveness probe to restart container if it becomes unresponsive
         livenessProbe:
-          httpGet:
-            path: /healthz
+          tcpSocket:
             port: 5400
           initialDelaySeconds: 30
           periodSeconds: 10
@@ -148,8 +147,7 @@ spec:
           failureThreshold: 3
         # Readiness probe to know when container is ready to accept traffic
         readinessProbe:
-          httpGet:
-            path: /healthz
+          tcpSocket:
             port: 5400
           initialDelaySeconds: 10
           periodSeconds: 5
@@ -267,16 +265,14 @@ spec:
             memory: "1Gi"
             cpu: "1000m"
         livenessProbe:
-          httpGet:
-            path: /healthz
+          tcpSocket:
             port: 5410
           initialDelaySeconds: 20
           periodSeconds: 10
           timeoutSeconds: 5
           failureThreshold: 3
         readinessProbe:
-          httpGet:
-            path: /healthz
+          tcpSocket:
             port: 5410
           initialDelaySeconds: 10
           periodSeconds: 5
@@ -323,10 +319,6 @@ metadata:
     app: langbot
 spec:
   replicas: 1
-  # Plugin Runtime has a single active LangBot control owner. Recreate avoids
-  # two LangBot pods fighting over that connection during a rolling update.
-  strategy:
-    type: Recreate
   selector:
     matchLabels:
       app: langbot

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:

src/langbot/pkg/core/app.py

@@ -4,7 +4,6 @@ import logging
 import asyncio
 import traceback
 import os
-import contextlib
 
 from ..platform import botmgr as im_mgr
 from ..platform.webhook_pusher import WebhookPusher
@@ -167,8 +166,7 @@ class Application:
     maintenance_service: maintenance_service.MaintenanceService = None
 
     def __init__(self):
-        self._shutdown_lock = asyncio.Lock()
-        self._shutdown_complete = False
+        pass
 
     async def initialize(self):
         pass
@@ -320,39 +318,7 @@ class Application:
             return default
         return parsed
 
-    async def shutdown(self):
-        """Stop application work and deterministically release runtime resources."""
-        async with self._shutdown_lock:
-            if self._shutdown_complete:
-                return
-
-            if self.task_mgr is not None:
-                self.task_mgr.cancel_by_scope(core_entities.LifecycleControlScope.APPLICATION)
-            if self.platform_mgr is not None:
-                with contextlib.suppress(Exception):
-                    await self.platform_mgr.shutdown()
-            if self.tool_mgr is not None:
-                with contextlib.suppress(Exception):
-                    await self.tool_mgr.shutdown()
-            if self.box_service is not None:
-                with contextlib.suppress(Exception):
-                    await self.box_service.shutdown()
-            if self.plugin_connector is not None:
-                with contextlib.suppress(Exception):
-                    await self.plugin_connector.aclose()
-
-            if self.task_mgr is not None:
-                tasks = [wrapper.task for wrapper in self.task_mgr.tasks if not wrapper.task.done()]
-                if tasks:
-                    await asyncio.gather(*tasks, return_exceptions=True)
-            self._shutdown_complete = True
-
     def dispose(self):
-        """Compatibility wrapper for callers that cannot await shutdown."""
-        loop = self.event_loop
-        if loop is not None and not loop.is_closed():
-            loop.create_task(self.shutdown())
-            return
         if self.plugin_connector is not None:
             self.plugin_connector.dispose()
         if self.box_service is not None:

src/langbot/pkg/core/boot.py

@@ -46,38 +46,21 @@ async def make_app(loop: asyncio.AbstractEventLoop) -> app.Application:
 
 async def main(loop: asyncio.AbstractEventLoop):
     app_inst: app.Application | None = None
-    runtime_loop = asyncio.get_running_loop()
-    shutdown_requested = asyncio.Event()
-    run_task: asyncio.Task | None = None
     try:
+        # Hang system signal processing
         import signal
 
         def signal_handler(sig, frame):
+            if app_inst is not None:
+                app_inst.dispose()
             print('[Signal] Program exit.')
-            runtime_loop.call_soon_threadsafe(shutdown_requested.set)
+            os._exit(0)
 
         signal.signal(signal.SIGINT, signal_handler)
-        if hasattr(signal, 'SIGTERM'):
-            signal.signal(signal.SIGTERM, signal_handler)
 
         app_inst = await make_app(loop)
-        if app_inst is None:
-            return
-        run_task = asyncio.create_task(app_inst.run())
-        shutdown_task = asyncio.create_task(shutdown_requested.wait())
-        done, pending = await asyncio.wait((run_task, shutdown_task), return_when=asyncio.FIRST_COMPLETED)
-        if shutdown_task in done:
-            await app_inst.shutdown()
-            if not run_task.done():
-                run_task.cancel()
-        for task in pending:
-            task.cancel()
-        results = await asyncio.gather(run_task, shutdown_task, return_exceptions=True)
-        run_result = results[0]
-        if isinstance(run_result, BaseException) and not isinstance(run_result, asyncio.CancelledError):
-            raise run_result
+        await app_inst.run()
     except Exception:
-        traceback.print_exc()
-    finally:
         if app_inst is not None:
-            await app_inst.shutdown()
+            app_inst.dispose()
+        traceback.print_exc()

src/langbot/pkg/core/stages/build_app.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import asyncio
+
 from .. import stage, app
 from ...utils import version, proxy
 from ...pipeline import pool, controller, pipelinemgr
@@ -185,16 +187,11 @@ class BuildAppStage(stage.BootingStage):
         ap.maintenance_service = maintenance_service_inst
 
         async def runtime_disconnect_callback(connector: plugin_connector.PluginRuntimeConnector) -> None:
-            connector.schedule_reconnect()
+            await asyncio.sleep(3)
+            await plugin_connector_inst.initialize()
 
         plugin_connector_inst = plugin_connector.PluginRuntimeConnector(ap, runtime_disconnect_callback)
-        try:
-            await plugin_connector_inst.initialize()
-        except Exception as exc:
-            # Keep the API/UI available while an external or managed runtime is
-            # starting, then recover in the background with bounded backoff.
-            ap.logger.warning(f'Plugin runtime unavailable during startup; reconnecting in background: {exc}')
-            plugin_connector_inst.schedule_reconnect()
+        await plugin_connector_inst.initialize()
         ap.plugin_connector = plugin_connector_inst
 
         ctrl = controller.Controller(ap)

src/langbot/pkg/plugin/connector.py

@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 import asyncio
-import contextlib
 import io
 import time
 import zipfile
@@ -37,12 +36,6 @@ from ..core import taskmgr
 from ..entity.persistence import plugin as persistence_plugin
 
 
-_CONNECT_TIMEOUT_SEC = 30.0
-_HEARTBEAT_INTERVAL_SEC = 20.0
-_HEARTBEAT_FAILURE_THRESHOLD = 3
-_RECONNECT_MAX_DELAY_SEC = 60.0
-
-
 class PluginRuntimeNotConnectedError(RuntimeError):
     """Raised when plugin runtime operations are requested before connection."""
 
@@ -77,218 +70,128 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
         super().__init__(ap)
         self.runtime_disconnect_callback = runtime_disconnect_callback
         self.is_enable_plugin = self.ap.instance_config.data.get('plugin', {}).get('enable', True)
-        self._transport_task: asyncio.Task | None = None
-        self._reconnect_task: asyncio.Task | None = None
-        self._generation = 0
-        self._connected = asyncio.Event()
-
-    def _runtime_handler(self) -> handler.RuntimeConnectionHandler:
-        runtime_handler = getattr(self, 'handler', None)
-        if runtime_handler is None:
-            raise PluginRuntimeNotConnectedError('Plugin runtime is not connected')
-        return runtime_handler
-
-    def _runtime_available(self) -> bool:
-        runtime_handler = getattr(self, 'handler', None)
-        if runtime_handler is None:
-            return False
-        # Unit-level and explicitly injected handlers don't own a transport.
-        # A managed transport must also have completed its handshake.
-        return self._transport_task is None or self._connected.is_set()
 
     async def heartbeat_loop(self):
-        failures = 0
-        while not self._closing:
-            await asyncio.sleep(_HEARTBEAT_INTERVAL_SEC)
+        while True:
+            await asyncio.sleep(20)
             try:
                 await self.ping_plugin_runtime()
-                failures = 0
                 self.ap.logger.debug('Heartbeat to plugin runtime success.')
             except Exception as e:
-                failures += 1
-                self.ap.logger.warning(
-                    f'Plugin runtime heartbeat failed ({failures}/{_HEARTBEAT_FAILURE_THRESHOLD}): {e}'
-                )
-                if failures >= _HEARTBEAT_FAILURE_THRESHOLD:
-                    self._connected.clear()
-                    self.schedule_reconnect()
-                    failures = 0
+                self.ap.logger.debug(f'Failed to heartbeat to plugin runtime: {e}')
 
     async def initialize(self):
         if not self.is_enable_plugin:
             self.ap.logger.info('Plugin system is disabled.')
             return
 
-        async with self._lifecycle_lock:
-            if self._closing:
-                raise PluginRuntimeNotConnectedError('Plugin runtime connector is shutting down')
-            if self._connected.is_set() and hasattr(self, 'handler'):
-                return
-
-            await self._stop_transport()
-            self._generation += 1
-            generation = self._generation
-            self._connected = asyncio.Event()
-            connect_errors: list[Exception] = []
-
-            async def new_connection_callback(
-                connection: base_connection.Connection,
-            ):
-                if generation != self._generation or self._closing:
-                    await connection.close()
-                    return
-
-                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)
+        async def new_connection_callback(connection: base_connection.Connection):
+            async def disconnect_callback(
+                rchandler: handler.RuntimeConnectionHandler,
+            ) -> bool:
+                if platform.get_platform() == 'docker' or platform.use_websocket_to_connect_plugin_runtime():
+                    self.ap.logger.error('Disconnected from plugin runtime, trying to reconnect...')
+                    await self.runtime_disconnect_callback(self)
+                    return False
+                else:
+                    self.ap.logger.error(
+                        'Disconnected from plugin runtime, cannot automatically reconnect while LangBot connects to plugin runtime via stdio, please restart LangBot.'
+                    )
                     return False
 
-                runtime_handler = handler.RuntimeConnectionHandler(connection, disconnect_callback, self.ap)
-                self.handler = runtime_handler
-                self.handler_task = asyncio.create_task(runtime_handler.run())
+            self.handler = handler.RuntimeConnectionHandler(connection, disconnect_callback, self.ap)
+
+            self.handler_task = asyncio.create_task(self.handler.run())
+            _ = await self.handler.ping()
+            # Push the configured marketplace (Space) URL to the runtime so it
+            # downloads plugins from the same Space LangBot is bound to, rather
+            # than relying on the runtime's own env/default.
+            space_url = self.ap.instance_config.data.get('space', {}).get('url', '').rstrip('/')
+            if space_url:
                 try:
-                    await runtime_handler.ping()
-                    space_url = self.ap.instance_config.data.get('space', {}).get('url', '').rstrip('/')
-                    if space_url:
-                        await runtime_handler.set_runtime_config(cloud_service_url=space_url)
-                    if generation == self._generation and not self._closing:
-                        self._connected.set()
-                        self.ap.logger.info('Connected to plugin runtime.')
-                    await self.handler_task
-                except asyncio.CancelledError:
-                    raise
-                except Exception as exc:
-                    if not self._connected.is_set():
-                        connect_errors.append(exc)
-                        self._connected.set()
-                finally:
-                    if generation == self._generation and not self._closing:
-                        self._connected.clear()
-                        if getattr(self, 'handler', None) is runtime_handler:
-                            del self.handler
-                        await self.runtime_disconnect_callback(self)
+                    await self.handler.set_runtime_config(cloud_service_url=space_url)
+                    self.ap.logger.info(f'Pushed marketplace URL to plugin runtime: {space_url}')
+                except Exception as e:
+                    self.ap.logger.warning(f'Failed to push runtime config: {e}')
+            self.ap.logger.info('Connected to plugin runtime.')
+            await self.handler_task
 
-            task_coro: typing.Coroutine
-            if platform.get_platform() == 'docker' or platform.use_websocket_to_connect_plugin_runtime():
-                ws_url = self.ap.instance_config.data.get('plugin', {}).get(
-                    'runtime_ws_url',
-                    'ws://langbot_plugin_runtime:5400/control/ws',
-                )
+        task: asyncio.Task | None = None
 
-                async def connection_failed(ctrl, exc=None):
-                    error = exc or RuntimeError('WebSocket connection failed')
-                    connect_errors.append(error)
-                    self._connected.set()
-
-                self.ctrl = ws_client_controller.WebSocketClientController(
-                    ws_url=ws_url,
-                    make_connection_failed_callback=connection_failed,
-                )
-                task_coro = self.ctrl.run(new_connection_callback)
-            elif platform.get_platform() == 'win32':
-                await self._start_runtime_subprocess('-m', 'langbot_plugin.cli.__init__', 'rt')
-                ws_url = 'ws://localhost:5400/control/ws'
-
-                async def connection_failed(ctrl, exc=None):
-                    error = exc or RuntimeError('WebSocket connection failed')
-                    connect_errors.append(error)
-                    self._connected.set()
-
-                self.ctrl = ws_client_controller.WebSocketClientController(
-                    ws_url=ws_url,
-                    make_connection_failed_callback=connection_failed,
-                )
-                task_coro = self.ctrl.run(new_connection_callback)
-            else:
-                self.ctrl = stdio_client_controller.StdioClientController(
-                    command=sys.executable,
-                    args=['-m', 'langbot_plugin.cli.__init__', 'rt', '-s'],
-                    env=os.environ.copy(),
-                )
-                task_coro = self.ctrl.run(new_connection_callback)
-
-            self._transport_task = asyncio.create_task(task_coro)
-            try:
-                await asyncio.wait_for(self._connected.wait(), timeout=_CONNECT_TIMEOUT_SEC)
-            except asyncio.TimeoutError as exc:
-                await self._stop_transport()
-                raise PluginRuntimeNotConnectedError('Plugin runtime did not become ready within 30 seconds') from exc
-            if connect_errors:
-                await self._stop_transport()
-                raise PluginRuntimeNotConnectedError(f'Plugin runtime connection failed: {connect_errors[-1]}')
-
-            if self.heartbeat_task is None or self.heartbeat_task.done():
-                self.heartbeat_task = asyncio.create_task(self.heartbeat_loop())
-
-    def schedule_reconnect(self) -> None:
-        if self._closing or not self.is_enable_plugin:
-            return
-        if self._reconnect_task is not None and not self._reconnect_task.done():
-            return
-        self._reconnect_task = asyncio.create_task(self._reconnect_loop())
-
-    async def _reconnect_loop(self) -> None:
-        delay = 1.0
-        try:
-            while not self._closing:
-                try:
-                    await self.initialize()
-                    return
-                except Exception as exc:
-                    self.ap.logger.warning(f'Plugin runtime reconnection failed: {exc}; retrying in {delay:.0f}s')
-                    await asyncio.sleep(delay)
-                    delay = min(delay * 2, _RECONNECT_MAX_DELAY_SEC)
-        finally:
-            self._reconnect_task = None
-
-    async def _stop_transport(self) -> None:
-        runtime_handler = getattr(self, 'handler', None)
-        if runtime_handler is not None:
-            with contextlib.suppress(Exception):
-                await runtime_handler.close()
-            del self.handler
-        tasks = [
-            task
-            for task in (
-                getattr(self, 'handler_task', None),
-                self._transport_task,
+        if platform.get_platform() == 'docker' or platform.use_websocket_to_connect_plugin_runtime():  # use websocket
+            self.ap.logger.info('use websocket to connect to plugin runtime')
+            ws_url = self.ap.instance_config.data.get('plugin', {}).get(
+                'runtime_ws_url', 'ws://langbot_plugin_runtime:5400/control/ws'
             )
-            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)
-        self._transport_task = None
-        if hasattr(self, 'handler_task'):
-            del self.handler_task
-        close_ctrl = getattr(getattr(self, 'ctrl', None), 'close', None)
-        if close_ctrl is not None:
-            with contextlib.suppress(Exception):
-                await close_ctrl()
 
-    async def aclose(self) -> None:
-        self._closing = True
-        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)
-        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()
-        await self._close_managed_subprocess()
+            async def make_connection_failed_callback(
+                ctrl: ws_client_controller.WebSocketClientController,
+                exc: Exception = None,
+            ) -> None:
+                if exc is not None:
+                    self.ap.logger.error(f'Failed to connect to plugin runtime({ws_url}): {exc}')
+                else:
+                    self.ap.logger.error(f'Failed to connect to plugin runtime({ws_url}), trying to reconnect...')
+                await self.runtime_disconnect_callback(self)
+
+            self.ctrl = ws_client_controller.WebSocketClientController(
+                ws_url=ws_url,
+                make_connection_failed_callback=make_connection_failed_callback,
+            )
+            task = self.ctrl.run(new_connection_callback)
+        elif platform.get_platform() == 'win32':
+            # Due to Windows's lack of supports for both stdio and subprocess:
+            # See also: https://docs.python.org/zh-cn/3.13/library/asyncio-platforms.html
+            # We have to launch runtime via cmd but communicate via ws.
+            self.ap.logger.info('(windows) use cmd to launch plugin runtime and communicate via ws')
+
+            await self._start_runtime_subprocess('-m', 'langbot_plugin.cli.__init__', 'rt')
+
+            ws_url = 'ws://localhost:5400/control/ws'
+
+            async def make_connection_failed_callback(
+                ctrl: ws_client_controller.WebSocketClientController,
+                exc: Exception = None,
+            ) -> None:
+                if exc is not None:
+                    self.ap.logger.error(f'(windows) Failed to connect to plugin runtime({ws_url}): {exc}')
+                else:
+                    self.ap.logger.error(
+                        f'(windows) Failed to connect to plugin runtime({ws_url}), trying to reconnect...'
+                    )
+                await self.runtime_disconnect_callback(self)
+
+            self.ctrl = ws_client_controller.WebSocketClientController(
+                ws_url=ws_url,
+                make_connection_failed_callback=make_connection_failed_callback,
+            )
+            task = self.ctrl.run(new_connection_callback)
+
+        else:  # stdio
+            self.ap.logger.info('use stdio to connect to plugin runtime')
+            # cmd: lbp rt -s
+            python_path = sys.executable
+            env = os.environ.copy()
+            self.ctrl = stdio_client_controller.StdioClientController(
+                command=python_path,
+                args=['-m', 'langbot_plugin.cli.__init__', 'rt', '-s'],
+                env=env,
+            )
+            task = self.ctrl.run(new_connection_callback)
+
+        if self.heartbeat_task is None:
+            self.heartbeat_task = asyncio.create_task(self.heartbeat_loop())
+
+        asyncio.create_task(task)
 
     async def initialize_plugins(self):
         pass
 
     async def ping_plugin_runtime(self):
-        return await self._runtime_handler().ping()
+        if not hasattr(self, 'handler'):
+            raise PluginRuntimeNotConnectedError('Plugin runtime is not connected')
+
+        return await self.handler.ping()
 
     def _inspect_plugin_package(
         self,
@@ -570,7 +473,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
 
                                     file_bytes = download_resp.content
                                     self._inspect_plugin_package(file_bytes, task_context)
-                                    file_key = await self._runtime_handler().send_file(file_bytes, 'lbpkg')
+                                    file_key = await self.handler.send_file(file_bytes, 'lbpkg')
                                     install_info['plugin_file_key'] = file_key
                                     self.ap.logger.info(f'Transfered file {file_key} to plugin runtime')
                                     # Continue to install via runtime
@@ -593,7 +496,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
             plugin_author, plugin_name = self._inspect_plugin_package(file_bytes, task_context)
             if task_context is not None and plugin_author and plugin_name:
                 task_context.metadata['plugin_name'] = f'{plugin_author}/{plugin_name}'
-            file_key = await self._runtime_handler().send_file(file_bytes, 'lbpkg')
+            file_key = await self.handler.send_file(file_bytes, 'lbpkg')
             install_info['plugin_file_key'] = file_key
             del install_info['plugin_file']
             self.ap.logger.info(f'Transfered file {file_key} to plugin runtime')
@@ -632,14 +535,14 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
                     plugin_author, plugin_name = self._inspect_plugin_package(file_bytes, task_context)
                     if task_context is not None and plugin_author and plugin_name:
                         task_context.metadata['plugin_name'] = f'{plugin_author}/{plugin_name}'
-                    file_key = await self._runtime_handler().send_file(file_bytes, 'lbpkg')
+                    file_key = await self.handler.send_file(file_bytes, 'lbpkg')
                     install_info['plugin_file_key'] = file_key
                     self.ap.logger.info(f'Transfered file {file_key} to plugin runtime')
             except Exception as e:
                 self.ap.logger.error(f'Failed to download file from GitHub: {e}')
                 raise Exception(f'Failed to download file from GitHub: {e}')
 
-        async for ret in self._runtime_handler().install_plugin(install_source.value, install_info):
+        async for ret in self.handler.install_plugin(install_source.value, install_info):
             current_action = ret.get('current_action', None)
             if current_action is not None:
                 if task_context is not None:
@@ -663,7 +566,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
         plugin_name: str,
         task_context: taskmgr.TaskContext | None = None,
     ) -> dict[str, Any]:
-        async for ret in self._runtime_handler().upgrade_plugin(plugin_author, plugin_name):
+        async for ret in self.handler.upgrade_plugin(plugin_author, plugin_name):
             current_action = ret.get('current_action', None)
             if current_action is not None:
                 if task_context is not None:
@@ -681,7 +584,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
         delete_data: bool = False,
         task_context: taskmgr.TaskContext | None = None,
     ) -> dict[str, Any]:
-        async for ret in self._runtime_handler().delete_plugin(plugin_author, plugin_name):
+        async for ret in self.handler.delete_plugin(plugin_author, plugin_name):
             current_action = ret.get('current_action', None)
             if current_action is not None:
                 if task_context is not None:
@@ -696,7 +599,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
         if delete_data:
             if task_context is not None:
                 task_context.trace('Cleaning up plugin configuration and storage...')
-            await self._runtime_handler().cleanup_plugin_data(plugin_author, plugin_name)
+            await self.handler.cleanup_plugin_data(plugin_author, plugin_name)
 
     async def list_plugins(self, component_kinds: list[str] | None = None) -> list[dict[str, Any]]:
         """List plugins, optionally filtered by component kinds.
@@ -707,10 +610,10 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
                            component of the specified kinds will be returned.
                            E.g., ['Command', 'EventListener', 'Tool'] for pipeline-related plugins.
         """
-        if not self.is_enable_plugin or not self._runtime_available():
+        if not self.is_enable_plugin:
             return []
 
-        plugins = await self._runtime_handler().list_plugins()
+        plugins = await self.handler.list_plugins()
 
         # Filter plugins by component kinds if specified
         if component_kinds is not None:
@@ -782,18 +685,18 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
         return plugins
 
     async def get_plugin_info(self, author: str, plugin_name: str) -> dict[str, Any]:
-        return await self._runtime_handler().get_plugin_info(author, plugin_name)
+        return await self.handler.get_plugin_info(author, plugin_name)
 
     async def set_plugin_config(self, plugin_author: str, plugin_name: str, config: dict[str, Any]) -> dict[str, Any]:
-        return await self._runtime_handler().set_plugin_config(plugin_author, plugin_name, config)
+        return await self.handler.set_plugin_config(plugin_author, plugin_name, config)
 
     @alru_cache(ttl=5 * 60)  # 5 minutes
     async def get_plugin_icon(self, plugin_author: str, plugin_name: str) -> dict[str, Any]:
-        return await self._runtime_handler().get_plugin_icon(plugin_author, plugin_name)
+        return await self.handler.get_plugin_icon(plugin_author, plugin_name)
 
     @alru_cache(ttl=5 * 60)  # 5 minutes
     async def get_plugin_readme(self, plugin_author: str, plugin_name: str, language: str = 'en') -> str:
-        return await self._runtime_handler().get_plugin_readme(plugin_author, plugin_name, language)
+        return await self.handler.get_plugin_readme(plugin_author, plugin_name, language)
 
     async def get_plugin_logs(
         self,
@@ -803,11 +706,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
         level: str | None = None,
     ) -> list[dict[str, Any]]:
         # Not cached: logs are live and change constantly.
-        return await self._runtime_handler().get_plugin_logs(plugin_author, plugin_name, limit, level)
+        return await self.handler.get_plugin_logs(plugin_author, plugin_name, limit, level)
 
     @alru_cache(ttl=5 * 60)
     async def get_plugin_assets(self, plugin_author: str, plugin_name: str, filepath: str) -> dict[str, Any]:
-        return await self._runtime_handler().get_plugin_assets(plugin_author, plugin_name, filepath)
+        return await self.handler.get_plugin_assets(plugin_author, plugin_name, filepath)
 
     async def handle_page_api(
         self,
@@ -818,15 +721,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
         method: str,
         body: Any = None,
     ) -> dict[str, Any]:
-        return await self._runtime_handler().handle_page_api(
-            plugin_author, plugin_name, page_id, endpoint, method, body
-        )
+        return await self.handler.handle_page_api(plugin_author, plugin_name, page_id, endpoint, method, body)
 
     async def get_debug_info(self) -> dict[str, Any]:
         """Get debug information including debug key and WS URL"""
-        if not self.is_enable_plugin or not self._runtime_available():
+        if not self.is_enable_plugin:
             return {}
-        return await self._runtime_handler().get_debug_info()
+        return await self.handler.get_debug_info()
 
     async def emit_event(
         self,
@@ -835,13 +736,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
     ) -> context.EventContext:
         event_ctx = context.EventContext.from_event(event)
 
-        if not self.is_enable_plugin or not self._runtime_available():
+        if not self.is_enable_plugin:
             event_ctx._emitted_plugins = []
             event_ctx._response_sources = []
             return event_ctx
 
         # Pass include_plugins to runtime for filtering
-        event_ctx_result = await self._runtime_handler().emit_event(
+        event_ctx_result = await self.handler.emit_event(
             event_ctx.model_dump(serialize_as_any=False), include_plugins=bound_plugins
         )
 
@@ -854,19 +755,19 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
 
     async def notify_plugin_diagnostic(self, diagnostic: dict[str, Any]) -> None:
         """Best-effort diagnostic forwarding to the plugin runtime."""
-        if not self.is_enable_plugin or not self._runtime_available():
+        if not self.is_enable_plugin:
             return
         try:
-            await self._runtime_handler().notify_plugin_diagnostic(diagnostic)
+            await self.handler.notify_plugin_diagnostic(diagnostic)
         except Exception as e:
             self.ap.logger.debug(f'Plugin diagnostic forwarding skipped: {e}')
 
     async def list_tools(self, bound_plugins: list[str] | None = None) -> list[ComponentManifest]:
-        if not self.is_enable_plugin or not self._runtime_available():
+        if not self.is_enable_plugin:
             return []
 
         # Pass include_plugins to runtime for filtering
-        list_tools_data = await self._runtime_handler().list_tools(include_plugins=bound_plugins)
+        list_tools_data = await self.handler.list_tools(include_plugins=bound_plugins)
 
         tools = [ComponentManifest.model_validate(tool) for tool in list_tools_data]
 
@@ -884,19 +785,16 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
             return {'error': 'Tool not found: plugin system is disabled'}
 
         # Pass include_plugins to runtime for validation
-        if not self._runtime_available():
-            return {'error': 'Plugin runtime is temporarily unavailable'}
-
-        return await self._runtime_handler().call_tool(
+        return await self.handler.call_tool(
             tool_name, parameters, session.model_dump(serialize_as_any=True), query_id, include_plugins=bound_plugins
         )
 
     async def list_commands(self, bound_plugins: list[str] | None = None) -> list[ComponentManifest]:
-        if not self.is_enable_plugin or not self._runtime_available():
+        if not self.is_enable_plugin:
             return []
 
         # Pass include_plugins to runtime for filtering
-        list_commands_data = await self._runtime_handler().list_commands(include_plugins=bound_plugins)
+        list_commands_data = await self.handler.list_commands(include_plugins=bound_plugins)
 
         commands = [ComponentManifest.model_validate(command) for command in list_commands_data]
 
@@ -905,15 +803,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
     async def execute_command(
         self, command_ctx: command_context.ExecuteContext, bound_plugins: list[str] | None = None
     ) -> typing.AsyncGenerator[command_context.CommandReturn, None]:
-        if not self.is_enable_plugin or not self._runtime_available():
+        if not self.is_enable_plugin:
             yield command_context.CommandReturn(error=command_errors.CommandNotFoundError(command_ctx.command))
             return
 
         # Pass include_plugins to runtime for validation
-        gen = self._runtime_handler().execute_command(
-            command_ctx.model_dump(serialize_as_any=True),
-            include_plugins=bound_plugins,
-        )
+        gen = self.handler.execute_command(command_ctx.model_dump(serialize_as_any=True), include_plugins=bound_plugins)
 
         async for ret in gen:
             cmd_ret = command_context.CommandReturn.model_validate(ret)
@@ -928,33 +823,27 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
         retrieval_context: dict[str, Any],
     ) -> dict[str, Any]:
         """Retrieve knowledge using a KnowledgeEngine instance."""
-        if not self.is_enable_plugin or not self._runtime_available():
+        if not self.is_enable_plugin:
             return {'results': []}
 
-        return await self._runtime_handler().retrieve_knowledge(
-            plugin_author, plugin_name, retriever_name, retrieval_context
-        )
+        return await self.handler.retrieve_knowledge(plugin_author, plugin_name, retriever_name, retrieval_context)
 
     def dispose(self):
-        """Best-effort synchronous compatibility wrapper; prefer ``aclose``."""
-        self._closing = True
+        # On non-Windows stdio mode, terminate via the controller's process handle.
+        # On Windows, the managed subprocess is cleaned up by the base class.
+        if (
+            self.is_enable_plugin
+            and hasattr(self, 'ctrl')
+            and isinstance(self.ctrl, stdio_client_controller.StdioClientController)
+        ):
+            self.ap.logger.info('Terminating plugin runtime process...')
+            self.ctrl.process.terminate()
+
+        self._dispose_subprocess()
+
         if self.heartbeat_task is not None:
             self.heartbeat_task.cancel()
             self.heartbeat_task = None
-        if self._reconnect_task is not None:
-            self._reconnect_task.cancel()
-            self._reconnect_task = None
-        for task in (
-            getattr(self, 'handler_task', None),
-            self._transport_task,
-        ):
-            if task is not None:
-                task.cancel()
-        ctrl = getattr(self, 'ctrl', None)
-        process = getattr(ctrl, 'process', None)
-        if process is not None and process.returncode is None:
-            process.terminate()
-        self._dispose_subprocess()
 
     @staticmethod
     def _parse_plugin_id(plugin_id: str) -> tuple[str, str]:
@@ -984,29 +873,29 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
             context_data: IngestionContext data.
         """
         plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
-        return await self._runtime_handler().rag_ingest_document(plugin_author, plugin_name, context_data)
+        return await self.handler.rag_ingest_document(plugin_author, plugin_name, context_data)
 
     async def call_rag_delete_document(self, plugin_id: str, document_id: str, kb_id: str) -> bool:
         plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
-        return await self._runtime_handler().rag_delete_document(plugin_author, plugin_name, document_id, kb_id)
+        return await self.handler.rag_delete_document(plugin_author, plugin_name, document_id, kb_id)
 
     async def get_rag_creation_schema(self, plugin_id: str) -> dict[str, Any]:
         plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
-        return await self._runtime_handler().get_rag_creation_schema(plugin_author, plugin_name)
+        return await self.handler.get_rag_creation_schema(plugin_author, plugin_name)
 
     async def get_rag_retrieval_schema(self, plugin_id: str) -> dict[str, Any]:
         plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
-        return await self._runtime_handler().get_rag_retrieval_schema(plugin_author, plugin_name)
+        return await self.handler.get_rag_retrieval_schema(plugin_author, plugin_name)
 
     async def rag_on_kb_create(self, plugin_id: str, kb_id: str, config: dict[str, Any]) -> dict[str, Any]:
         """Notify plugin about KB creation."""
         plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
-        return await self._runtime_handler().rag_on_kb_create(plugin_author, plugin_name, kb_id, config)
+        return await self.handler.rag_on_kb_create(plugin_author, plugin_name, kb_id, config)
 
     async def rag_on_kb_delete(self, plugin_id: str, kb_id: str) -> dict[str, Any]:
         """Notify plugin about KB deletion."""
         plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
-        return await self._runtime_handler().rag_on_kb_delete(plugin_author, plugin_name, kb_id)
+        return await self.handler.rag_on_kb_delete(plugin_author, plugin_name, kb_id)
 
     async def call_rag_retrieve(self, plugin_id: str, retrieval_context: dict[str, Any]) -> dict[str, Any]:
         """Call plugin to retrieve knowledge.
@@ -1016,25 +905,25 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
             retrieval_context: RetrievalContext data.
         """
         plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
-        return await self._runtime_handler().retrieve_knowledge(plugin_author, plugin_name, '', retrieval_context)
+        return await self.handler.retrieve_knowledge(plugin_author, plugin_name, '', retrieval_context)
 
     async def list_knowledge_engines(self) -> list[dict[str, Any]]:
         """List all available Knowledge Engines from plugins.
 
         Returns a list of Knowledge Engines with their capabilities and configuration schemas.
         """
-        if not self.is_enable_plugin or not self._runtime_available():
+        if not self.is_enable_plugin:
             return []
 
-        return await self._runtime_handler().list_knowledge_engines()
+        return await self.handler.list_knowledge_engines()
 
     async def list_parsers(self) -> list[dict[str, Any]]:
         """List all available parsers from plugins."""
-        if not self.is_enable_plugin or not self._runtime_available():
+        if not self.is_enable_plugin:
             return []
-        return await self._runtime_handler().list_parsers()
+        return await self.handler.list_parsers()
 
     async def call_parser(self, plugin_id: str, context_data: dict[str, Any], file_bytes: bytes) -> dict[str, Any]:
         """Call plugin to parse a document."""
         plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
-        return await self._runtime_handler().parse_document(plugin_author, plugin_name, context_data, file_bytes)
+        return await self.handler.parse_document(plugin_author, plugin_name, context_data, file_bytes)

src/langbot/pkg/provider/tools/loaders/mcp.py

@@ -11,7 +11,6 @@ import traceback
 from langbot_plugin.api.entities.events import pipeline_query
 import sqlalchemy
 import asyncio
-import hashlib
 import httpx
 
 import uuid as uuid_module
@@ -27,13 +26,7 @@ from ....core import app
 import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
 import langbot_plugin.api.entities.builtin.provider.message as provider_message
 from ....entity.persistence import mcp as persistence_mcp
-from .mcp_stdio import (
-    BoxStdioSessionRuntime,
-    MCPServerBoxConfig as MCPServerBoxConfig,  # noqa: F401 - public re-export
-    MCPSessionErrorPhase,
-    _ColdStartRetry,
-    _get_default_memory_mb,
-)  # noqa: F401
+from .mcp_stdio import BoxStdioSessionRuntime, MCPServerBoxConfig, MCPSessionErrorPhase, _ColdStartRetry  # noqa: F401
 
 # Synthesized LLM tools for MCP resources (not from server tools/list).
 # Dispatched in MCPLoader.invoke_tool; placeholder func on LLMTool is never used.
@@ -1213,34 +1206,15 @@ class RuntimeMCPSession:
         return self._box_stdio_runtime.uses_box_stdio()
 
     def _build_box_session_id(self) -> str:
-        # Compatible MCP servers share a session and remain isolated by
-        # process_id. A server with a different immutable resource profile gets
-        # another session; Docker/E2B cannot change memory/image/etc. after a
-        # session has been created.
-        config = self._box_stdio_runtime.config
-        default_memory = _get_default_memory_mb(self.ap)
-        profile = {
-            'image': config.image,
-            'network': config.network,
-            'host_path_mode': config.host_path_mode,
-            'cpus': config.cpus,
-            'memory_mb': config.memory_mb or default_memory,
-            'pids_limit': config.pids_limit,
-            'read_only_rootfs': (config.read_only_rootfs if config.read_only_rootfs is not None else False),
-        }
-        default_profile = {
-            'image': None,
-            'network': 'on',
-            'host_path_mode': 'ro',
-            'cpus': None,
-            'memory_mb': default_memory,
-            'pids_limit': None,
-            'read_only_rootfs': False,
-        }
-        if profile == default_profile:
-            return 'mcp-shared'
-        digest = hashlib.sha256(json.dumps(profile, sort_keys=True).encode('utf-8')).hexdigest()[:12]
-        return f'mcp-shared-{digest}'
+        # Both live servers and transient config-page tests share ONE Box
+        # session ('mcp-shared'). A test therefore reuses the already-running
+        # container (and, for an existing server, its live managed process)
+        # instead of paying a full per-test session cold-start + dependency
+        # bootstrap. Isolation between a test and the live servers is provided
+        # at the *process* level: each server/test has its own process_id and a
+        # test only ever stops its own process_id (see cleanup_session), so it
+        # never disturbs another server's process or the shared session itself.
+        return 'mcp-shared'
 
     def _rewrite_path(self, path: str, host_path: str | None) -> str:
         return self._box_stdio_runtime.rewrite_path(path, host_path)

src/langbot/pkg/utils/managed_runtime.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import asyncio
-import contextlib
 import os
 import sys
 from typing import TYPE_CHECKING, Awaitable, Callable
@@ -28,8 +27,6 @@ class ManagedRuntimeConnector:
         self.ap = ap
         self.runtime_subprocess = None
         self.runtime_subprocess_task = None
-        self._lifecycle_lock = asyncio.Lock()
-        self._closing = False
 
     async def _start_runtime_subprocess(self, *args: str) -> None:
         """Launch a local runtime as a subprocess of the current Python interpreter.
@@ -89,23 +86,3 @@ class ManagedRuntimeConnector:
         if self.runtime_subprocess_task is not None:
             self.runtime_subprocess_task.cancel()
             self.runtime_subprocess_task = None
-
-    async def _close_managed_subprocess(self) -> None:
-        """Terminate, escalate, and reap the optional owned subprocess."""
-        process = self.runtime_subprocess
-        wait_task = self.runtime_subprocess_task
-        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()
-        if wait_task is not None and wait_task is not asyncio.current_task():
-            if not wait_task.done():
-                wait_task.cancel()
-            await asyncio.gather(wait_task, return_exceptions=True)
-        self.runtime_subprocess = None
-        self.runtime_subprocess_task = None

src/langbot/templates/config.yaml

@@ -143,11 +143,6 @@ box:
     backend: 'local'  # 'local' (Docker/nsjail), 'docker', 'nsjail', or 'e2b'. Can be written via BOX__BACKEND.
     runtime:
         endpoint: ''  # External Box Runtime base URL, e.g. 'ws://127.0.0.1:5410'. Leave empty for local auto-managed runtime.
-    limits:
-        max_sessions: 64  # Includes persistent sessions. New sessions fail explicitly when this cap is reached.
-        max_managed_processes: 64  # Maximum concurrently running stdio MCP / managed processes.
-        max_completed_processes: 256  # Global cap for retained exited-process diagnostics.
-        completed_process_retention_sec: 300  # Keep exited-process diagnostics before releasing memory.
     local:
         profile: 'default'
         image: ''  # Custom local sandbox image. Leave empty to use the profile default.

tests/unit_tests/box/test_box_connector.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from types import SimpleNamespace
-from unittest.mock import AsyncMock, Mock
+from unittest.mock import Mock
 
 import pytest
 
@@ -104,21 +104,3 @@ def test_box_runtime_connector_dispose_terminates_subprocess(monkeypatch: pytest
     ctrl_task.cancel.assert_called_once()
     assert connector._handler_task is None
     assert connector._ctrl_task is None
-
-
-@pytest.mark.asyncio
-async def test_box_runtime_connector_cleans_partial_transport_on_connect_failure(
-    monkeypatch: pytest.MonkeyPatch,
-):
-    monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
-    monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
-    connector = BoxRuntimeConnector(make_app(Mock()))
-    connector._start_local_stdio = AsyncMock(side_effect=RuntimeError('bind failed'))
-    connector._stop_transport = AsyncMock()
-    connector._close_managed_subprocess = AsyncMock()
-
-    with pytest.raises(RuntimeError, match='bind failed'):
-        await connector.initialize()
-
-    assert connector._stop_transport.await_count == 2
-    connector._close_managed_subprocess.assert_awaited_once()

tests/unit_tests/box/test_box_service.py

@@ -325,28 +325,10 @@ async def test_box_service_dispose_schedules_shutdown_on_event_loop(monkeypatch:
     service.dispose()
     await asyncio.sleep(0)
 
-    connector.dispose.assert_not_called()
+    connector.dispose.assert_called_once()
     service.shutdown.assert_awaited_once()
 
 
-@pytest.mark.asyncio
-async def test_box_service_shutdown_reaps_connector_when_runtime_rpc_is_offline(
-    monkeypatch: pytest.MonkeyPatch,
-):
-    connector = Mock()
-    connector.client = Mock()
-    connector.client.shutdown = AsyncMock(side_effect=RuntimeError('offline'))
-    connector.aclose = AsyncMock()
-
-    monkeypatch.setattr('langbot.pkg.box.service.BoxRuntimeConnector', Mock(return_value=connector))
-
-    service = BoxService(make_app(Mock()))
-    await service.shutdown()
-
-    connector.client.shutdown.assert_awaited_once()
-    connector.aclose.assert_awaited_once()
-
-
 @pytest.mark.asyncio
 async def test_box_runtime_reuses_request_session():
     logger = Mock()

tests/unit_tests/core/test_boot.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import signal
 from types import SimpleNamespace
-from unittest.mock import Mock
 
 import pytest
 
@@ -19,60 +18,47 @@ async def test_main_signal_handler_handles_sigint_before_app_created(monkeypatch
     async def fake_make_app(loop):
         captured_handler[signal.SIGINT](signal.SIGINT, None)
 
+    def fake_exit(code):
+        raise SystemExit(code)
+
     monkeypatch.setattr(signal, 'signal', fake_signal)
     monkeypatch.setattr(boot, 'make_app', fake_make_app)
+    monkeypatch.setattr(boot.os, '_exit', fake_exit)
 
-    await boot.main(SimpleNamespace())
+    with pytest.raises(SystemExit) as exc_info:
+        await boot.main(SimpleNamespace())
+
+    assert exc_info.value.code == 0
 
 
 @pytest.mark.asyncio
 async def test_main_signal_handler_disposes_created_app(monkeypatch):
     captured_handler = {}
-    app_inst = SimpleNamespace(shutdown_called=False)
+    app_inst = SimpleNamespace(disposed=False)
 
     def fake_signal(sig, handler):
         captured_handler[sig] = handler
 
-    async def shutdown():
-        app_inst.shutdown_called = True
+    def dispose():
+        app_inst.disposed = True
 
     async def run():
         captured_handler[signal.SIGINT](signal.SIGINT, None)
 
     async def fake_make_app(loop):
-        app_inst.shutdown = shutdown
+        app_inst.dispose = dispose
         app_inst.run = run
         return app_inst
 
+    def fake_exit(code):
+        raise SystemExit(code)
+
     monkeypatch.setattr(signal, 'signal', fake_signal)
     monkeypatch.setattr(boot, 'make_app', fake_make_app)
+    monkeypatch.setattr(boot.os, '_exit', fake_exit)
 
-    await boot.main(SimpleNamespace())
+    with pytest.raises(SystemExit) as exc_info:
+        await boot.main(SimpleNamespace())
 
-    assert app_inst.shutdown_called is True
-
-
-@pytest.mark.asyncio
-async def test_main_reports_app_run_failure_and_still_shuts_down(monkeypatch):
-    app_inst = SimpleNamespace(shutdown_called=False)
-
-    async def shutdown():
-        app_inst.shutdown_called = True
-
-    async def run():
-        raise RuntimeError('run failed')
-
-    async def fake_make_app(loop):
-        app_inst.shutdown = shutdown
-        app_inst.run = run
-        return app_inst
-
-    print_exc = Mock()
-    monkeypatch.setattr(signal, 'signal', lambda *_args: None)
-    monkeypatch.setattr(boot, 'make_app', fake_make_app)
-    monkeypatch.setattr(boot.traceback, 'print_exc', print_exc)
-
-    await boot.main(SimpleNamespace())
-
-    print_exc.assert_called_once()
-    assert app_inst.shutdown_called is True
+    assert exc_info.value.code == 0
+    assert app_inst.disposed is True

tests/unit_tests/plugin/test_connector_methods.py

@@ -62,25 +62,6 @@ class TestListPlugins:
 
         assert result == []
 
-    @pytest.mark.asyncio
-    async def test_returns_empty_while_runtime_is_disconnected(self):
-        connector = create_mock_connector()
-
-        result = await connector.list_plugins()
-
-        assert result == []
-
-    @pytest.mark.asyncio
-    async def test_returns_empty_after_managed_transport_disconnects(self):
-        connector = create_mock_connector()
-        connector.handler = AsyncMock()
-        connector._transport_task = Mock()
-
-        result = await connector.list_plugins()
-
-        assert result == []
-        connector.handler.list_plugins.assert_not_called()
-
     @pytest.mark.asyncio
     async def test_calls_handler_list_plugins(self):
         """Test that handler.list_plugins is called."""
@@ -313,12 +294,6 @@ class TestListKnowledgeEngines:
         connector.handler.list_knowledge_engines.assert_called_once()
         assert result == [{'plugin_id': 'author/engine', 'name': 'Engine'}]
 
-    @pytest.mark.asyncio
-    async def test_returns_empty_while_runtime_is_disconnected(self):
-        connector = create_mock_connector()
-
-        assert await connector.list_knowledge_engines() == []
-
 
 class TestListParsers:
     """Tests for list_parsers method."""
@@ -356,12 +331,6 @@ class TestListParsers:
         connector.handler.list_parsers.assert_called_once()
         assert result == [{'plugin_id': 'author/parser', 'supported_mime_types': ['text/plain']}]
 
-    @pytest.mark.asyncio
-    async def test_returns_empty_while_runtime_is_disconnected(self):
-        connector = create_mock_connector()
-
-        assert await connector.list_parsers() == []
-
 
 class TestCallParser:
     """Tests for call_parser method."""

tests/unit_tests/provider/test_mcp_box_integration.py

@@ -417,7 +417,7 @@ class TestBuildBoxSessionPayload:
         payload = s._build_box_session_payload('session-123')
         assert payload['image'] == 'node:20'
         assert payload['cpus'] == 2.0
-        assert payload['memory_mb'] == 1024
+        assert payload["memory_mb"] == 1024
         assert payload['pids_limit'] == 256
 
     def test_none_fields_excluded(self, mcp_module):
@@ -680,51 +680,6 @@ class TestGetRuntimeInfoDict:
         # ... but are isolated by distinct process_ids within that session.
         assert transient._box_stdio_runtime.process_id != live._box_stdio_runtime.process_id
 
-    def test_different_resource_profiles_use_different_box_sessions(self, mcp_module):
-        ap = _make_ap()
-        ap.box_service.available = True
-        default = _make_session(
-            mcp_module,
-            {
-                'name': 'default',
-                'uuid': 'default-uuid',
-                'mode': 'stdio',
-                'command': 'uvx',
-                'args': ['mcp-server-time'],
-            },
-            ap=ap,
-        )
-        constrained = _make_session(
-            mcp_module,
-            {
-                'name': 'constrained',
-                'uuid': 'constrained-uuid',
-                'mode': 'stdio',
-                'command': 'uvx',
-                'args': ['mcp-server-time'],
-                'box': {'memory_mb': 2048},
-            },
-            ap=ap,
-        )
-
-        assert default._build_box_session_id() == 'mcp-shared'
-        assert constrained._build_box_session_id().startswith('mcp-shared-')
-        assert constrained._build_box_session_id() != default._build_box_session_id()
-
-        writable = _make_session(
-            mcp_module,
-            {
-                'name': 'writable',
-                'uuid': 'writable-uuid',
-                'mode': 'stdio',
-                'command': 'uvx',
-                'args': ['mcp-server-time'],
-                'box': {'host_path_mode': 'rw'},
-            },
-            ap=ap,
-        )
-        assert writable._build_box_session_id() != default._build_box_session_id()
-
     def test_stdio_session_refuses_when_box_unavailable(self, mcp_module):
         """Policy: when Box is configured but unavailable (disabled in config
         OR connection failed), stdio MCP servers are NOT treated as box-stdio.