diff --git a/docker/kubernetes.yaml b/docker/kubernetes.yaml index 6adc50510..d92ef8ce9 100644 --- a/docker/kubernetes.yaml +++ b/docker/kubernetes.yaml @@ -139,7 +139,8 @@ spec: cpu: "1000m" # Liveness probe to restart container if it becomes unresponsive livenessProbe: - tcpSocket: + httpGet: + path: /healthz port: 5400 initialDelaySeconds: 30 periodSeconds: 10 @@ -147,7 +148,8 @@ spec: failureThreshold: 3 # Readiness probe to know when container is ready to accept traffic readinessProbe: - tcpSocket: + httpGet: + path: /healthz port: 5400 initialDelaySeconds: 10 periodSeconds: 5 @@ -265,14 +267,16 @@ spec: memory: "1Gi" cpu: "1000m" livenessProbe: - tcpSocket: + httpGet: + path: /healthz port: 5410 initialDelaySeconds: 20 periodSeconds: 10 timeoutSeconds: 5 failureThreshold: 3 readinessProbe: - tcpSocket: + httpGet: + path: /healthz port: 5410 initialDelaySeconds: 10 periodSeconds: 5 @@ -319,6 +323,10 @@ 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 2257910d1..2e37b84f7 100644 --- a/src/langbot/pkg/box/connector.py +++ b/src/langbot/pkg/box/connector.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import contextlib import json import os import sys @@ -28,6 +29,7 @@ _DOCKER_BOX_HOST = 'langbot_box' _DEFAULT_PORT = 5410 _HEARTBEAT_INTERVAL_SEC = 20 +_HEARTBEAT_FAILURE_THRESHOLD = 3 # Top-level keys under ``box`` that are LangBot-internal and should not be # forwarded to the Box runtime. @@ -113,6 +115,8 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): self._handler_task: asyncio.Task | None = None self._ctrl_task: asyncio.Task | None = None self._heartbeat_task: asyncio.Task | None = None + self._ctrl = None + self._generation = 0 # Parse the relay URL once for reuse. parsed = urlparse(self.ws_relay_base_url) @@ -145,29 +149,64 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): return self.uses_websocket() async def initialize(self) -> None: - if self._uses_websocket(): - if platform.get_platform() == 'win32' and not self.configured_runtime_endpoint: - await self._start_subprocess_then_ws() - else: - await self._connect_remote_ws() - else: - await self._start_local_stdio() + async with self._lifecycle_lock: + if self._closing: + raise BoxRuntimeUnavailableError('box runtime connector is shutting down') + await self._stop_transport() + self._generation += 1 + try: + if self._uses_websocket(): + if platform.get_platform() == 'win32' and not self.configured_runtime_endpoint: + await self._start_subprocess_then_ws() + else: + await self._connect_remote_ws() + else: + await self._start_local_stdio() + except BaseException: + await self._stop_transport() + await self._close_managed_subprocess() + raise - # Start heartbeat after successful connection - if self._heartbeat_task is None: - self._heartbeat_task = asyncio.create_task(self._heartbeat_loop()) + if self._heartbeat_task is None or self._heartbeat_task.done(): + self._heartbeat_task = asyncio.create_task(self._heartbeat_loop()) + + async def reconnect(self) -> None: + async with self._lifecycle_lock: + if self._closing: + raise BoxRuntimeUnavailableError('box runtime connector is shutting down') + await self._stop_transport() + self._generation += 1 + try: + if self._uses_websocket(): + if platform.get_platform() == 'win32' and not self.configured_runtime_endpoint: + await self._start_subprocess_then_ws() + else: + await self._connect_remote_ws() + else: + await self._start_local_stdio() + except BaseException: + await self._stop_transport() + await self._close_managed_subprocess() + raise # -- heartbeat ----------------------------------------------------------- async def _heartbeat_loop(self) -> None: """Periodically ping the Box runtime to detect silent disconnections.""" - while True: + failures = 0 + while not self._closing: await asyncio.sleep(_HEARTBEAT_INTERVAL_SEC) try: await self.ping() + failures = 0 self.ap.logger.debug('Heartbeat to Box runtime success.') except Exception as e: - self.ap.logger.debug(f'Failed to heartbeat to Box runtime: {e}') + failures += 1 + self.ap.logger.warning(f'Box runtime heartbeat failed ({failures}/{_HEARTBEAT_FAILURE_THRESHOLD}): {e}') + if failures >= _HEARTBEAT_FAILURE_THRESHOLD: + failures = 0 + if self.runtime_disconnect_callback is not None: + await self.runtime_disconnect_callback(self) async def ping(self) -> None: if self._handler is None: @@ -197,8 +236,9 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): args=['-m', 'langbot_plugin.cli.__init__', 'box', '-s', '--ws-control-port', str(self._relay_port)], env=env, ) + self._ctrl = ctrl self._ctrl_task = asyncio.create_task( - ctrl.run(self._make_connection_callback('stdio', connected, connect_error)) + ctrl.run(self._make_connection_callback('stdio', connected, connect_error, self._generation)) ) try: @@ -282,8 +322,9 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): await self.runtime_disconnect_callback(self) ctrl = WebSocketClientController(ws_url=ws_url, make_connection_failed_callback=on_connect_failed) + self._ctrl = ctrl self._ctrl_task = asyncio.create_task( - ctrl.run(self._make_connection_callback(transport_name, connected, connect_error)) + ctrl.run(self._make_connection_callback(transport_name, connected, connect_error, self._generation)) ) try: @@ -299,8 +340,12 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): transport_name: str, connected: asyncio.Event, connect_error: list[Exception], + generation: int, ): async def new_connection_callback(connection: Connection) -> None: + if generation != self._generation or self._closing: + await connection.close() + return handler = Handler(connection) self._handler = handler self.client.set_handler(handler) @@ -313,6 +358,8 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): self.ap.logger.info(f'Connected to Box runtime via {transport_name}.') connected.set() await self._handler_task + except asyncio.CancelledError: + raise except Exception as exc: if not connected.is_set(): connect_error.append(exc) @@ -322,22 +369,66 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): # If we reach here, handler.run() returned normally (connection # closed) or raised after the initial handshake succeeded. # Either way, treat it as a disconnect. - if connected.is_set(): - if self._uses_websocket(): - self.ap.logger.error('Disconnected from Box runtime, trying to reconnect...') - if self.runtime_disconnect_callback is not None: - await self.runtime_disconnect_callback(self) - else: - self.ap.logger.error( - 'Disconnected from Box runtime via stdio. ' - 'Cannot automatically reconnect — please restart LangBot.' - ) + if ( + connected.is_set() + and generation == self._generation + and not self._closing + and self.runtime_disconnect_callback is not None + ): + self.ap.logger.error('Disconnected from Box runtime, trying to reconnect...') + await self.runtime_disconnect_callback(self) return new_connection_callback # -- lifecycle ----------------------------------------------------------- + async def _stop_transport(self) -> None: + if self._handler is not None: + with contextlib.suppress(Exception): + await self._handler.close() + self.client.set_handler(None) + tasks = [ + task + for task in (self._handler_task, self._ctrl_task) + if task is not None and task is not asyncio.current_task() + ] + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + close_ctrl = getattr(self._ctrl, 'close', None) + if close_ctrl is not None: + with contextlib.suppress(Exception): + await close_ctrl() + self._handler = None + self._handler_task = None + self._ctrl_task = None + self._ctrl = None + + async def aclose(self) -> None: + self._closing = True + if self._heartbeat_task is not None: + self._heartbeat_task.cancel() + await asyncio.gather(self._heartbeat_task, return_exceptions=True) + self._heartbeat_task = None + await self._stop_transport() + + process = getattr(self, '_subprocess', None) + if process is not None and process.returncode is None: + with contextlib.suppress(ProcessLookupError): + process.terminate() + try: + await asyncio.wait_for(process.wait(), timeout=3) + except asyncio.TimeoutError: + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + self._subprocess = None + await self._close_managed_subprocess() + def dispose(self) -> None: + """Best-effort synchronous compatibility wrapper; prefer ``aclose``.""" + self._closing = True if self._heartbeat_task is not None: self._heartbeat_task.cancel() self._heartbeat_task = None diff --git a/src/langbot/pkg/box/service.py b/src/langbot/pkg/box/service.py index ca30eb930..2b900df89 100644 --- a/src/langbot/pkg/box/service.py +++ b/src/langbot/pkg/box/service.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import collections +import contextlib import datetime as _dt import enum import json @@ -65,6 +66,8 @@ class BoxService: self.workspace_quota_mb = self._load_workspace_quota_mb() self._recent_errors: collections.deque[dict] = collections.deque(maxlen=_MAX_RECENT_ERRORS) self._shutdown_task = None + self._reconnect_task: asyncio.Task | None = None + self._closing = False self._available = False self._connector_error: str = '' self._reconnecting = False @@ -110,6 +113,8 @@ class BoxService: self.ap.logger.warning(f'LangBot Box runtime unavailable, sandbox features disabled: {exc}') self._available = False self._connector_error = str(exc) + if self._runtime_connector is not None: + await self._on_runtime_disconnect(self._runtime_connector) async def _on_runtime_disconnect(self, connector: BoxRuntimeConnector) -> None: """Called by the connector when the Box runtime connection drops. @@ -118,27 +123,26 @@ class BoxService: Skipped entirely when Box is disabled by config — that path should never have connected in the first place. """ - if not self._enabled: + if not self._enabled or self._closing: return - if self._reconnecting: + if self._reconnect_task is not None and not self._reconnect_task.done(): return # Another reconnect loop is already running self._reconnecting = True self._available = False self._connector_error = 'Disconnected from Box runtime' self.ap.logger.warning('Box runtime disconnected, sandbox features temporarily disabled.') - asyncio.create_task(self._reconnect_loop(connector)) + self._reconnect_task = asyncio.create_task(self._reconnect_loop(connector)) async def _reconnect_loop(self, connector: BoxRuntimeConnector) -> None: """Retry reconnection with exponential backoff (3s → 60s max).""" delay = 3 max_delay = 60 try: - while True: + while not self._closing: self.ap.logger.info(f'Attempting to reconnect to Box runtime in {delay}s...') await asyncio.sleep(delay) try: - connector.dispose() - await connector.initialize() + await connector.reconnect() self._available = True self._connector_error = '' self.ap.logger.info('Box runtime reconnected, sandbox features restored.') @@ -149,6 +153,7 @@ class BoxService: delay = min(delay * 2, max_delay) finally: self._reconnecting = False + self._reconnect_task = None @property def available(self) -> bool: @@ -838,14 +843,28 @@ class BoxService: return attachments async def shutdown(self): - await self.client.shutdown() + if self._closing: + return + self._closing = True + self._available = False + reconnect_task = self._reconnect_task + self._reconnect_task = None + if reconnect_task is not None and reconnect_task is not asyncio.current_task(): + reconnect_task.cancel() + await asyncio.gather(reconnect_task, return_exceptions=True) + # The runtime may already be offline. A failed best-effort SHUTDOWN RPC + # must not prevent us from cancelling transports and reaping children. + with contextlib.suppress(Exception): + await self.client.shutdown() + if self._runtime_connector is not None: + await self._runtime_connector.aclose() def dispose(self): - if self._runtime_connector is not None: - self._runtime_connector.dispose() loop = getattr(self.ap, 'event_loop', None) if loop is not None and not loop.is_closed() and (self._shutdown_task is None or self._shutdown_task.done()): self._shutdown_task = loop.create_task(self.shutdown()) + elif self._runtime_connector is not None: + self._runtime_connector.dispose() async def get_sessions(self) -> list[dict]: if not self._available: diff --git a/src/langbot/pkg/core/app.py b/src/langbot/pkg/core/app.py index b0adb5594..e5bc1abf0 100644 --- a/src/langbot/pkg/core/app.py +++ b/src/langbot/pkg/core/app.py @@ -4,6 +4,7 @@ import logging import asyncio import traceback import os +import contextlib from ..platform import botmgr as im_mgr from ..platform.webhook_pusher import WebhookPusher @@ -166,7 +167,8 @@ class Application: maintenance_service: maintenance_service.MaintenanceService = None def __init__(self): - pass + self._shutdown_lock = asyncio.Lock() + self._shutdown_complete = False async def initialize(self): pass @@ -318,7 +320,39 @@ 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 f2b335ac8..fb6919630 100644 --- a/src/langbot/pkg/core/boot.py +++ b/src/langbot/pkg/core/boot.py @@ -46,21 +46,38 @@ 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.') - os._exit(0) + runtime_loop.call_soon_threadsafe(shutdown_requested.set) signal.signal(signal.SIGINT, signal_handler) + if hasattr(signal, 'SIGTERM'): + signal.signal(signal.SIGTERM, signal_handler) app_inst = await make_app(loop) - await app_inst.run() + 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 except Exception: - if app_inst is not None: - app_inst.dispose() traceback.print_exc() + finally: + if app_inst is not None: + await app_inst.shutdown() diff --git a/src/langbot/pkg/core/stages/build_app.py b/src/langbot/pkg/core/stages/build_app.py index a8d53b7b3..2af7fa579 100644 --- a/src/langbot/pkg/core/stages/build_app.py +++ b/src/langbot/pkg/core/stages/build_app.py @@ -1,7 +1,5 @@ from __future__ import annotations -import asyncio - from .. import stage, app from ...utils import version, proxy from ...pipeline import pool, controller, pipelinemgr @@ -187,11 +185,16 @@ class BuildAppStage(stage.BootingStage): ap.maintenance_service = maintenance_service_inst async def runtime_disconnect_callback(connector: plugin_connector.PluginRuntimeConnector) -> None: - await asyncio.sleep(3) - await plugin_connector_inst.initialize() + connector.schedule_reconnect() plugin_connector_inst = plugin_connector.PluginRuntimeConnector(ap, runtime_disconnect_callback) - await plugin_connector_inst.initialize() + 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() 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 6075d4b68..1b962bd0a 100644 --- a/src/langbot/pkg/plugin/connector.py +++ b/src/langbot/pkg/plugin/connector.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio +import contextlib import io import time import zipfile @@ -36,6 +37,12 @@ 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.""" @@ -70,128 +77,218 @@ 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): - while True: - await asyncio.sleep(20) + failures = 0 + while not self._closing: + await asyncio.sleep(_HEARTBEAT_INTERVAL_SEC) try: await self.ping_plugin_runtime() + failures = 0 self.ap.logger.debug('Heartbeat to plugin runtime success.') except Exception as e: - self.ap.logger.debug(f'Failed to heartbeat to plugin runtime: {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 async def initialize(self): if not self.is_enable_plugin: self.ap.logger.info('Plugin system is disabled.') return - 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.' - ) + 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) return False - 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: + runtime_handler = handler.RuntimeConnectionHandler(connection, disconnect_callback, self.ap) + self.handler = runtime_handler + self.handler_task = asyncio.create_task(runtime_handler.run()) try: - 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 + 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) - task: asyncio.Task | None = None + 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', + ) - 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' + 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 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 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 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 initialize_plugins(self): pass async def ping_plugin_runtime(self): - if not hasattr(self, 'handler'): - raise PluginRuntimeNotConnectedError('Plugin runtime is not connected') - - return await self.handler.ping() + return await self._runtime_handler().ping() def _inspect_plugin_package( self, @@ -473,7 +570,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): file_bytes = download_resp.content self._inspect_plugin_package(file_bytes, task_context) - file_key = await self.handler.send_file(file_bytes, 'lbpkg') + file_key = await self._runtime_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 @@ -496,7 +593,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.handler.send_file(file_bytes, 'lbpkg') + file_key = await self._runtime_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') @@ -535,14 +632,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.handler.send_file(file_bytes, 'lbpkg') + file_key = await self._runtime_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.handler.install_plugin(install_source.value, install_info): + async for ret in self._runtime_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: @@ -566,7 +663,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): plugin_name: str, task_context: taskmgr.TaskContext | None = None, ) -> dict[str, Any]: - async for ret in self.handler.upgrade_plugin(plugin_author, plugin_name): + async for ret in self._runtime_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: @@ -584,7 +681,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): delete_data: bool = False, task_context: taskmgr.TaskContext | None = None, ) -> dict[str, Any]: - async for ret in self.handler.delete_plugin(plugin_author, plugin_name): + async for ret in self._runtime_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: @@ -599,7 +696,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): if delete_data: if task_context is not None: task_context.trace('Cleaning up plugin configuration and storage...') - await self.handler.cleanup_plugin_data(plugin_author, plugin_name) + await self._runtime_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. @@ -610,10 +707,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: + if not self.is_enable_plugin or not self._runtime_available(): return [] - plugins = await self.handler.list_plugins() + plugins = await self._runtime_handler().list_plugins() # Filter plugins by component kinds if specified if component_kinds is not None: @@ -685,18 +782,18 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): return plugins async def get_plugin_info(self, author: str, plugin_name: str) -> dict[str, Any]: - return await self.handler.get_plugin_info(author, plugin_name) + return await self._runtime_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.handler.set_plugin_config(plugin_author, plugin_name, config) + return await self._runtime_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.handler.get_plugin_icon(plugin_author, plugin_name) + return await self._runtime_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.handler.get_plugin_readme(plugin_author, plugin_name, language) + return await self._runtime_handler().get_plugin_readme(plugin_author, plugin_name, language) async def get_plugin_logs( self, @@ -706,11 +803,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): level: str | None = None, ) -> list[dict[str, Any]]: # Not cached: logs are live and change constantly. - return await self.handler.get_plugin_logs(plugin_author, plugin_name, limit, level) + return await self._runtime_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.handler.get_plugin_assets(plugin_author, plugin_name, filepath) + return await self._runtime_handler().get_plugin_assets(plugin_author, plugin_name, filepath) async def handle_page_api( self, @@ -721,13 +818,15 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): method: str, body: Any = None, ) -> dict[str, Any]: - return await self.handler.handle_page_api(plugin_author, plugin_name, page_id, endpoint, method, body) + return await self._runtime_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: + if not self.is_enable_plugin or not self._runtime_available(): return {} - return await self.handler.get_debug_info() + return await self._runtime_handler().get_debug_info() async def emit_event( self, @@ -736,13 +835,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): ) -> context.EventContext: event_ctx = context.EventContext.from_event(event) - if not self.is_enable_plugin: + if not self.is_enable_plugin or not self._runtime_available(): event_ctx._emitted_plugins = [] event_ctx._response_sources = [] return event_ctx # Pass include_plugins to runtime for filtering - event_ctx_result = await self.handler.emit_event( + event_ctx_result = await self._runtime_handler().emit_event( event_ctx.model_dump(serialize_as_any=False), include_plugins=bound_plugins ) @@ -755,19 +854,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: + if not self.is_enable_plugin or not self._runtime_available(): return try: - await self.handler.notify_plugin_diagnostic(diagnostic) + await self._runtime_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: + if not self.is_enable_plugin or not self._runtime_available(): return [] # Pass include_plugins to runtime for filtering - list_tools_data = await self.handler.list_tools(include_plugins=bound_plugins) + list_tools_data = await self._runtime_handler().list_tools(include_plugins=bound_plugins) tools = [ComponentManifest.model_validate(tool) for tool in list_tools_data] @@ -785,16 +884,19 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): return {'error': 'Tool not found: plugin system is disabled'} # Pass include_plugins to runtime for validation - return await self.handler.call_tool( + if not self._runtime_available(): + return {'error': 'Plugin runtime is temporarily unavailable'} + + return await self._runtime_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: + if not self.is_enable_plugin or not self._runtime_available(): return [] # Pass include_plugins to runtime for filtering - list_commands_data = await self.handler.list_commands(include_plugins=bound_plugins) + list_commands_data = await self._runtime_handler().list_commands(include_plugins=bound_plugins) commands = [ComponentManifest.model_validate(command) for command in list_commands_data] @@ -803,12 +905,15 @@ 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: + if not self.is_enable_plugin or not self._runtime_available(): yield command_context.CommandReturn(error=command_errors.CommandNotFoundError(command_ctx.command)) return # Pass include_plugins to runtime for validation - gen = self.handler.execute_command(command_ctx.model_dump(serialize_as_any=True), include_plugins=bound_plugins) + gen = self._runtime_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) @@ -823,27 +928,33 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): retrieval_context: dict[str, Any], ) -> dict[str, Any]: """Retrieve knowledge using a KnowledgeEngine instance.""" - if not self.is_enable_plugin: + if not self.is_enable_plugin or not self._runtime_available(): return {'results': []} - return await self.handler.retrieve_knowledge(plugin_author, plugin_name, retriever_name, retrieval_context) + return await self._runtime_handler().retrieve_knowledge( + plugin_author, plugin_name, retriever_name, retrieval_context + ) def dispose(self): - # 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() - + """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 + 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]: @@ -873,29 +984,29 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): context_data: IngestionContext data. """ plugin_author, plugin_name = self._parse_plugin_id(plugin_id) - return await self.handler.rag_ingest_document(plugin_author, plugin_name, context_data) + return await self._runtime_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.handler.rag_delete_document(plugin_author, plugin_name, document_id, kb_id) + return await self._runtime_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.handler.get_rag_creation_schema(plugin_author, plugin_name) + return await self._runtime_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.handler.get_rag_retrieval_schema(plugin_author, plugin_name) + return await self._runtime_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.handler.rag_on_kb_create(plugin_author, plugin_name, kb_id, config) + return await self._runtime_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.handler.rag_on_kb_delete(plugin_author, plugin_name, kb_id) + return await self._runtime_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. @@ -905,25 +1016,25 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): retrieval_context: RetrievalContext data. """ plugin_author, plugin_name = self._parse_plugin_id(plugin_id) - return await self.handler.retrieve_knowledge(plugin_author, plugin_name, '', retrieval_context) + return await self._runtime_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: + if not self.is_enable_plugin or not self._runtime_available(): return [] - return await self.handler.list_knowledge_engines() + return await self._runtime_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: + if not self.is_enable_plugin or not self._runtime_available(): return [] - return await self.handler.list_parsers() + return await self._runtime_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.handler.parse_document(plugin_author, plugin_name, context_data, file_bytes) + return await self._runtime_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 11a943545..edb25b03b 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp.py @@ -11,6 +11,7 @@ import traceback from langbot_plugin.api.entities.events import pipeline_query import sqlalchemy import asyncio +import hashlib import httpx import uuid as uuid_module @@ -26,7 +27,13 @@ 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, MCPSessionErrorPhase, _ColdStartRetry # noqa: F401 +from .mcp_stdio import ( + BoxStdioSessionRuntime, + MCPServerBoxConfig as MCPServerBoxConfig, # noqa: F401 - public re-export + MCPSessionErrorPhase, + _ColdStartRetry, + _get_default_memory_mb, +) # 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. @@ -1206,15 +1213,34 @@ class RuntimeMCPSession: return self._box_stdio_runtime.uses_box_stdio() def _build_box_session_id(self) -> str: - # 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' + # 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}' 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 77f59be4c..374529874 100644 --- a/src/langbot/pkg/utils/managed_runtime.py +++ b/src/langbot/pkg/utils/managed_runtime.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import contextlib import os import sys from typing import TYPE_CHECKING, Awaitable, Callable @@ -27,6 +28,8 @@ 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. @@ -86,3 +89,23 @@ 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 f4ae79bf8..b5870b8eb 100644 --- a/src/langbot/templates/config.yaml +++ b/src/langbot/templates/config.yaml @@ -143,6 +143,11 @@ 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 ddd4899b0..0f2f8c340 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 Mock +from unittest.mock import AsyncMock, Mock import pytest @@ -104,3 +104,21 @@ 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 4d66ec8f8..5c8c1da85 100644 --- a/tests/unit_tests/box/test_box_service.py +++ b/tests/unit_tests/box/test_box_service.py @@ -325,10 +325,28 @@ async def test_box_service_dispose_schedules_shutdown_on_event_loop(monkeypatch: service.dispose() await asyncio.sleep(0) - connector.dispose.assert_called_once() + connector.dispose.assert_not_called() 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 461458c0e..b2b717006 100644 --- a/tests/unit_tests/core/test_boot.py +++ b/tests/unit_tests/core/test_boot.py @@ -2,6 +2,7 @@ from __future__ import annotations import signal from types import SimpleNamespace +from unittest.mock import Mock import pytest @@ -18,47 +19,60 @@ 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) - with pytest.raises(SystemExit) as exc_info: - await boot.main(SimpleNamespace()) - - assert exc_info.value.code == 0 + await boot.main(SimpleNamespace()) @pytest.mark.asyncio async def test_main_signal_handler_disposes_created_app(monkeypatch): captured_handler = {} - app_inst = SimpleNamespace(disposed=False) + app_inst = SimpleNamespace(shutdown_called=False) def fake_signal(sig, handler): captured_handler[sig] = handler - def dispose(): - app_inst.disposed = True + async def shutdown(): + app_inst.shutdown_called = True async def run(): captured_handler[signal.SIGINT](signal.SIGINT, None) async def fake_make_app(loop): - app_inst.dispose = dispose + app_inst.shutdown = shutdown 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) - with pytest.raises(SystemExit) as exc_info: - await boot.main(SimpleNamespace()) + await boot.main(SimpleNamespace()) - assert exc_info.value.code == 0 - assert app_inst.disposed is True + 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 diff --git a/tests/unit_tests/plugin/test_connector_methods.py b/tests/unit_tests/plugin/test_connector_methods.py index 34cab5271..88de1608d 100644 --- a/tests/unit_tests/plugin/test_connector_methods.py +++ b/tests/unit_tests/plugin/test_connector_methods.py @@ -62,6 +62,25 @@ 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.""" @@ -294,6 +313,12 @@ 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.""" @@ -331,6 +356,12 @@ 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 6b0c141ed..4a790908b 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,6 +680,51 @@ 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.