From 1765c432625c1b2925a10b5611d74478bf18d8de Mon Sep 17 00:00:00 2001 From: Junyan Qin Date: Tue, 21 Jul 2026 18:43:01 +0800 Subject: [PATCH] Revert "fix(runtime): make plugin and box connectors resilient" This reverts commit 0b461e5830dc3851830df8a65ad45be84918ae46. --- docker/kubernetes.yaml | 16 +- src/langbot/pkg/box/connector.py | 139 +----- src/langbot/pkg/box/service.py | 37 +- src/langbot/pkg/core/app.py | 36 +- src/langbot/pkg/core/boot.py | 31 +- src/langbot/pkg/core/stages/build_app.py | 13 +- src/langbot/pkg/plugin/connector.py | 417 +++++++----------- src/langbot/pkg/provider/tools/loaders/mcp.py | 46 +- src/langbot/pkg/utils/managed_runtime.py | 23 - src/langbot/templates/config.yaml | 5 - tests/unit_tests/box/test_box_connector.py | 20 +- tests/unit_tests/box/test_box_service.py | 20 +- tests/unit_tests/core/test_boot.py | 54 +-- .../plugin/test_connector_methods.py | 31 -- .../provider/test_mcp_box_integration.py | 47 +- 15 files changed, 236 insertions(+), 699 deletions(-) diff --git a/docker/kubernetes.yaml b/docker/kubernetes.yaml index d92ef8ce9..6adc50510 100644 --- a/docker/kubernetes.yaml +++ b/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 diff --git a/src/langbot/pkg/box/connector.py b/src/langbot/pkg/box/connector.py index 2e37b84f7..2257910d1 100644 --- a/src/langbot/pkg/box/connector.py +++ b/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 diff --git a/src/langbot/pkg/box/service.py b/src/langbot/pkg/box/service.py index 2b900df89..ca30eb930 100644 --- a/src/langbot/pkg/box/service.py +++ b/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: diff --git a/src/langbot/pkg/core/app.py b/src/langbot/pkg/core/app.py index e5bc1abf0..b0adb5594 100644 --- a/src/langbot/pkg/core/app.py +++ b/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: diff --git a/src/langbot/pkg/core/boot.py b/src/langbot/pkg/core/boot.py index fb6919630..f2b335ac8 100644 --- a/src/langbot/pkg/core/boot.py +++ b/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() diff --git a/src/langbot/pkg/core/stages/build_app.py b/src/langbot/pkg/core/stages/build_app.py index 2af7fa579..a8d53b7b3 100644 --- a/src/langbot/pkg/core/stages/build_app.py +++ b/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) diff --git a/src/langbot/pkg/plugin/connector.py b/src/langbot/pkg/plugin/connector.py index 1b962bd0a..6075d4b68 100644 --- a/src/langbot/pkg/plugin/connector.py +++ b/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) diff --git a/src/langbot/pkg/provider/tools/loaders/mcp.py b/src/langbot/pkg/provider/tools/loaders/mcp.py index edb25b03b..11a943545 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp.py +++ b/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) diff --git a/src/langbot/pkg/utils/managed_runtime.py b/src/langbot/pkg/utils/managed_runtime.py index 374529874..77f59be4c 100644 --- a/src/langbot/pkg/utils/managed_runtime.py +++ b/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 diff --git a/src/langbot/templates/config.yaml b/src/langbot/templates/config.yaml index b5870b8eb..f4ae79bf8 100644 --- a/src/langbot/templates/config.yaml +++ b/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. diff --git a/tests/unit_tests/box/test_box_connector.py b/tests/unit_tests/box/test_box_connector.py index 0f2f8c340..ddd4899b0 100644 --- a/tests/unit_tests/box/test_box_connector.py +++ b/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() diff --git a/tests/unit_tests/box/test_box_service.py b/tests/unit_tests/box/test_box_service.py index 5c8c1da85..4d66ec8f8 100644 --- a/tests/unit_tests/box/test_box_service.py +++ b/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() diff --git a/tests/unit_tests/core/test_boot.py b/tests/unit_tests/core/test_boot.py index b2b717006..461458c0e 100644 --- a/tests/unit_tests/core/test_boot.py +++ b/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 diff --git a/tests/unit_tests/plugin/test_connector_methods.py b/tests/unit_tests/plugin/test_connector_methods.py index 88de1608d..34cab5271 100644 --- a/tests/unit_tests/plugin/test_connector_methods.py +++ b/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.""" diff --git a/tests/unit_tests/provider/test_mcp_box_integration.py b/tests/unit_tests/provider/test_mcp_box_integration.py index 4a790908b..6b0c141ed 100644 --- a/tests/unit_tests/provider/test_mcp_box_integration.py +++ b/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.