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/pyproject.toml b/pyproject.toml index 1f1c3e438..53425aa01 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,7 +70,7 @@ dependencies = [ "chromadb>=1.0.0,<2.0.0", "qdrant-client (>=1.15.1,<2.0.0)", "pyseekdb==1.1.0.post3", - "langbot-plugin==0.4.13", + "langbot-plugin==0.4.17", "asyncpg>=0.30.0", "line-bot-sdk>=3.19.0", "matrix-nio>=0.25.2", diff --git a/src/langbot/pkg/box/connector.py b/src/langbot/pkg/box/connector.py index 2257910d1..f24287477 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,67 @@ 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') + self._generation += 1 + await self._stop_transport() + 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') + self._generation += 1 + await self._stop_transport() + 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._heartbeat_task is None or self._heartbeat_task.done(): + 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.""" - 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: @@ -196,9 +238,11 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): # mirroring `rt -s`. args=['-m', 'langbot_plugin.cli.__init__', 'box', '-s', '--ws-control-port', str(self._relay_port)], env=env, + capture_stderr=False, ) + 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 +326,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,9 +344,29 @@ 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) + connection_ready = False + disconnect_notified = False + + async def notify_disconnect() -> None: + nonlocal disconnect_notified + if ( + connection_ready + and not disconnect_notified + and generation == self._generation + and not self._closing + and self.runtime_disconnect_callback is not None + ): + disconnect_notified = True + self.ap.logger.error('Disconnected from Box runtime, trying to reconnect...') + await self.runtime_disconnect_callback(self) + self._handler = handler self.client.set_handler(handler) self._handler_task = asyncio.create_task(handler.run()) @@ -311,33 +376,74 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): await handler.call_action(LangBotToBoxAction.INIT, self._filtered_box_config) self.ap.logger.debug('Sent box configuration to Box runtime via INIT.') self.ap.logger.info(f'Connected to Box runtime via {transport_name}.') + connection_ready = True connected.set() await self._handler_task + except asyncio.CancelledError: + raise except Exception as exc: if not connected.is_set(): connect_error.append(exc) connected.set() return - - # 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.' - ) + finally: + if getattr(self, '_handler', None) is handler: + self._handler = None + self.client.set_handler(None) + await notify_disconnect() 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 + self._generation += 1 + 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..809a94130 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,29 +123,34 @@ 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._ensure_default_workspace() + await self._purge_attachment_dirs() self._available = True self._connector_error = '' + skill_mgr = getattr(self.ap, 'skill_mgr', None) + reload_skills = getattr(skill_mgr, 'reload_skills', None) + if callable(reload_skills): + await reload_skills() self.ap.logger.info('Box runtime reconnected, sandbox features restored.') return except Exception as exc: @@ -149,6 +159,7 @@ class BoxService: delay = min(delay * 2, max_delay) finally: self._reconnecting = False + self._reconnect_task = None @property def available(self) -> bool: @@ -838,14 +849,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..8d2c8597c 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,235 @@ 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 + + self._generation += 1 + generation = self._generation + await self._stop_transport() + 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 + connection_ready = False + disconnect_notified = False + + async def notify_disconnect() -> None: + nonlocal disconnect_notified + if ( + connection_ready + and not disconnect_notified + and generation == self._generation + and not self._closing + ): + disconnect_notified = True + self._connected.clear() + await self.runtime_disconnect_callback(self) + + async def disconnect_callback( + rchandler: handler.RuntimeConnectionHandler, + ) -> bool: + await notify_disconnect() 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: + connection_ready = True + 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 notify_disconnect() - 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(), + capture_stderr=False, + ) + 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: + self._connected.clear() + runtime_handler = getattr(self, 'handler', None) + if runtime_handler is not None: + with contextlib.suppress(Exception): + await runtime_handler.close() + if getattr(self, 'handler', None) is runtime_handler: + 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 + self._generation += 1 + 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 +587,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 +610,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 +649,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 +680,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 +698,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 +713,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 +724,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 +799,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 +820,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 +835,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 +852,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 +871,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 +901,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 +922,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 +945,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 +1001,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 +1033,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..1a594260b 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp.py @@ -3,14 +3,17 @@ from __future__ import annotations import base64 import enum import json +import math import re import time import typing from contextlib import AsyncExitStack, asynccontextmanager +from datetime import timedelta 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 +29,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. @@ -44,6 +53,7 @@ MCP_RESOURCE_CONTEXT_MAX_BYTES = 96 * 1024 MCP_RESOURCE_TRACE_QUERY_KEY = '_mcp_resource_reads' MCP_RESOURCE_LINKS_QUERY_KEY = '_mcp_resource_links' MCP_RESOURCE_CONTEXT_QUERY_KEY = '_mcp_resource_context' +MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS = 300.0 TEXT_LIKE_MIME_TYPES = { 'application/json', @@ -206,6 +216,10 @@ class _CallerReconnect(Exception): """ +class MCPToolCallTimeoutError(TimeoutError): + """An MCP tool call exceeded its configured per-server deadline.""" + + class RuntimeMCPSession: """运行时 MCP 会话""" @@ -255,6 +269,9 @@ class RuntimeMCPSession: self.ap = ap self.enable = enable self.session = None + self.tool_call_timeout_sec = self._parse_tool_call_timeout( + server_config.get('tool_call_timeout_sec', MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS) + ) # Transient test sessions (created from the config page "test" button, # which carry no persisted server UUID) must NOT share the live @@ -295,21 +312,39 @@ class RuntimeMCPSession: self._box_stdio_runtime = BoxStdioSessionRuntime(self) self.box_config = self._box_stdio_runtime.config + def _parse_tool_call_timeout(self, value: typing.Any) -> float: + """Return a safe tool-call timeout; zero explicitly disables it.""" + try: + timeout = -1 if isinstance(value, bool) else float(value) + if timeout > 0: + # Validate the exact conversion used for each call here, so a + # finite-but-enormous manual config cannot fail at invocation. + timedelta(seconds=timeout) + except (TypeError, ValueError, OverflowError): + timeout = -1 + + if not math.isfinite(timeout) or timeout < 0: + self.ap.logger.warning( + f'Invalid MCP tool call timeout {value!r} for {self.server_name}; ' + f'using {MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS:g} seconds' + ) + return MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS + return timeout + async def _init_stdio_python_server(self): if self._uses_box_stdio(): await self._box_stdio_runtime.initialize() return - # Box is configured (ap.box_service exists) but currently unavailable - # (disabled by config or connection failed). Refuse stdio MCP rather - # than silently falling through to host-stdio — the operator asked - # for the sandbox and the failure mode should be visible. + # Box is configured but explicitly disabled. Refuse stdio MCP rather + # than silently falling through to host-stdio — the operator asked for + # the sandbox and the failure mode should be visible. An enabled Box + # that is reconnecting is handled above and waits for availability. # # Set ``error_phase = BOX_UNAVAILABLE`` BEFORE raising so the retry - # wrapper can short-circuit (retrying is pointless when Box is - # deliberately off) and the frontend can render a localized, - # actionable message instead of this raw RuntimeError. Keep the - # message itself short — the frontend ignores it for this phase. + # wrapper can distinguish a deliberately disabled Box from an enabled + # runtime that is still reconnecting. Keep the message itself short — + # the frontend ignores it for this phase. box_service = getattr(self.ap, 'box_service', None) if box_service is not None and not getattr(box_service, 'available', False): self.error_phase = MCPSessionErrorPhase.BOX_UNAVAILABLE @@ -613,17 +648,29 @@ class RuntimeMCPSession: await asyncio.sleep(2) continue except Exception as e: - self.retry_count = attempt + 1 if self._shutdown_event.is_set(): return # Shutdown requested, don't retry - # BOX_UNAVAILABLE is a deliberate refusal, not a transient - # failure — retrying produces log spam and a misleading - # "Failed after N attempts" message. Surface it immediately. if self.error_phase == MCPSessionErrorPhase.BOX_UNAVAILABLE: + box_service = getattr(self.ap, 'box_service', None) + if box_service is not None and getattr(box_service, 'enabled', True): + # Box is configured and may recover independently of + # this MCP session. Keep retrying without consuming the + # fatal budget; _wait_for_box_runtime() rate-limits the + # loop to one warning per startup timeout. + self.status = MCPSessionStatus.CONNECTING + self.error_message = None + self.error_phase = None + await asyncio.sleep(1) + continue + # Explicitly disabled Box is a deliberate refusal, not a + # transient failure. Surface it immediately without log + # spam or a misleading "Failed after N attempts" message. + self.retry_count = attempt + 1 self.status = MCPSessionStatus.ERROR self.error_message = str(e) self._ready_event.set() return + self.retry_count = attempt + 1 if attempt >= self._MAX_RETRIES: self.status = MCPSessionStatus.ERROR self.error_message = f'Failed after {self._MAX_RETRIES + 1} attempts: {self._describe_exception(e)}' @@ -950,8 +997,22 @@ class RuntimeMCPSession: raise Exception('MCP session is not connected') try: - result = await self.session.call_tool(tool_name, arguments) + read_timeout = timedelta(seconds=self.tool_call_timeout_sec) if self.tool_call_timeout_sec > 0 else None + result = await self.session.call_tool( + tool_name, + arguments, + read_timeout_seconds=read_timeout, + ) except Exception as e: + if self._is_tool_call_timeout(e): + self.ap.logger.warning( + f'MCP tool {tool_name} on {self.server_name} timed out after ' + f'{self.tool_call_timeout_sec:g} seconds' + ) + raise MCPToolCallTimeoutError( + f"MCP tool '{tool_name}' on server '{self.server_name}' timed out after " + f'{self.tool_call_timeout_sec:g} seconds' + ) from e if attempt == 0 and self._is_session_terminated(e): self.ap.logger.warning( f'MCP tool {tool_name} on {self.server_name} got session terminated, triggering reconnect...' @@ -974,6 +1035,16 @@ class RuntimeMCPSession: raise Exception('MCP session is not connected') + @staticmethod + def _is_tool_call_timeout(exc: BaseException) -> bool: + """Recognize the MCP SDK's per-request timeout without retrying it.""" + return any( + isinstance(leaf, McpError) + and leaf.error.code == httpx.codes.REQUEST_TIMEOUT + and leaf.error.message.startswith('Timed out while waiting for response') + for leaf in RuntimeMCPSession._iter_exception_leaves(exc) + ) + def get_tools(self) -> list[resource_tool.LLMTool]: return self.functions @@ -1206,15 +1277,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) @@ -1797,11 +1887,23 @@ class MCPLoader(loader.ToolLoader): async def shutdown(self): """关闭所有工具""" self.ap.logger.info('Shutting down all MCP sessions...') - for server_name, session in list(self.sessions.items()): + + hosted_tasks = [task for task in self._hosted_mcp_tasks if not task.done()] + for task in hosted_tasks: + task.cancel() + if hosted_tasks: + await asyncio.gather(*hosted_tasks, return_exceptions=True) + self._hosted_mcp_tasks.clear() + + async def shutdown_session(server_name: str, session: RuntimeMCPSession) -> None: try: await session.shutdown() self.ap.logger.debug(f'Shutdown MCP session: {server_name}') except Exception as e: self.ap.logger.error(f'Error shutting down MCP session {server_name}: {e}\n{traceback.format_exc()}') + + await asyncio.gather( + *(shutdown_session(server_name, session) for server_name, session in list(self.sessions.items())) + ) self.sessions.clear() self.ap.logger.info('All MCP sessions shutdown complete') diff --git a/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py b/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py index 134c60a88..8c956fda3 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py @@ -185,11 +185,10 @@ class BoxStdioSessionRuntime: box_service = getattr(self.ap, 'box_service', None) if box_service is None: return False - # When Box is configured but currently unavailable (disabled or - # connection failed), do NOT silently fall through to host-stdio — - # that would bypass the sandbox the operator asked for. The caller - # is expected to refuse the stdio MCP server with a clear error. - return bool(getattr(box_service, 'available', False)) + # An enabled Box service remains the required transport while it is + # reconnecting. initialize() waits for availability instead of + # permanently failing the MCP server or falling through to host stdio. + return bool(getattr(box_service, 'enabled', True)) async def initialize(self) -> None: await self._wait_for_box_runtime() @@ -488,7 +487,7 @@ class BoxStdioSessionRuntime: ) warned = True if asyncio.get_running_loop().time() >= deadline: - self.owner.error_phase = MCPSessionErrorPhase.SESSION_CREATE + self.owner.error_phase = MCPSessionErrorPhase.BOX_UNAVAILABLE raise Exception(f'Box runtime is not available after {int(timeout_sec)} seconds') await asyncio.sleep(1) diff --git a/src/langbot/pkg/provider/tools/loaders/native.py b/src/langbot/pkg/provider/tools/loaders/native.py index 7f5ee4226..a1a8c027b 100644 --- a/src/langbot/pkg/provider/tools/loaders/native.py +++ b/src/langbot/pkg/provider/tools/loaders/native.py @@ -57,7 +57,7 @@ class NativeToolLoader(loader.ToolLoader): return await is_box_backend_available(self.ap) async def get_tools(self, bound_plugins: list[str] | None = None) -> list[resource_tool.LLMTool]: - if not self._is_sandbox_available(): + if not await self._is_sandbox_available(): return [] if self._tools is None: self._tools = [ @@ -71,7 +71,7 @@ class NativeToolLoader(loader.ToolLoader): return list(self._tools) async def has_tool(self, name: str) -> bool: - return name in _ALL_TOOL_NAMES and self._is_sandbox_available() + return name in _ALL_TOOL_NAMES and await self._is_sandbox_available() async def invoke_tool(self, name: str, parameters: dict, query: pipeline_query.Query): if name == EXEC_TOOL_NAME: @@ -652,12 +652,9 @@ else: if callable(refresh_skill): refresh_skill(selected_skill.get('name', '')) - def _is_sandbox_available(self) -> bool: - """Check if sandbox backend is available. - - This checks the cached backend availability from initialization, - not just whether the box_service process is running. - """ + async def _is_sandbox_available(self) -> bool: + """Refresh backend availability so Box reconnects restore tool exposure.""" + self._backend_available = await self._check_backend_available() return bool(self._backend_available) def _build_exec_tool(self) -> resource_tool.LLMTool: diff --git a/src/langbot/pkg/provider/tools/loaders/skill_authoring.py b/src/langbot/pkg/provider/tools/loaders/skill_authoring.py index d53721785..fe0f65620 100644 --- a/src/langbot/pkg/provider/tools/loaders/skill_authoring.py +++ b/src/langbot/pkg/provider/tools/loaders/skill_authoring.py @@ -49,19 +49,27 @@ class SkillToolLoader(loader.ToolLoader): return await is_box_backend_available(self.ap) async def get_tools(self, bound_plugins: list[str] | None = None) -> list[resource_tool.LLMTool]: - if not self._is_available(): + if not await self._is_available(): return [] + if not self._tools: + self._tools = [ + self._build_activate_skill_tool(), + self._build_register_skill_tool(), + ] return list(self._tools) async def has_tool(self, name: str) -> bool: - return self._is_available() and name in SKILL_TOOL_NAMES + return await self._is_available() and name in SKILL_TOOL_NAMES - def _is_available(self) -> bool: + async def _is_available(self) -> bool: """Check if skill tools should be available. Skill tools require both a skill manager and a sandbox backend. """ - return self._has_skill_manager() and self._sandbox_available + if not self._has_skill_manager(): + return False + self._sandbox_available = await self._check_sandbox_available() + return self._sandbox_available async def invoke_tool(self, name: str, parameters: dict, query) -> typing.Any: if name == ACTIVATE_SKILL_TOOL_NAME: 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..91016291d 100644 --- a/tests/unit_tests/box/test_box_connector.py +++ b/tests/unit_tests/box/test_box_connector.py @@ -1,11 +1,13 @@ from __future__ import annotations +import asyncio from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import AsyncMock, Mock import pytest from langbot_plugin.box.client import ActionRPCBoxClient +from langbot.pkg.box import connector as connector_module from langbot.pkg.box.connector import BoxRuntimeConnector @@ -104,3 +106,136 @@ 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() + + +@pytest.mark.asyncio +async def test_box_runtime_connector_starts_heartbeat_after_reconnect( + 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'), None]) + connector._stop_transport = AsyncMock() + connector._close_managed_subprocess = AsyncMock() + + with pytest.raises(RuntimeError, match='bind failed'): + await connector.initialize() + + assert connector._heartbeat_task is None + + await connector.reconnect() + + assert connector._heartbeat_task is not None + assert not connector._heartbeat_task.done() + await connector.aclose() + + +@pytest.mark.asyncio +async def test_box_stdio_connection_does_not_capture_unconsumed_stderr( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux') + monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False) + created = {} + + class FakeHandler: + def __init__(self, connection): + self.release = asyncio.Event() + + async def call_action(self, action, data): + return None + + async def run(self): + await self.release.wait() + + async def close(self): + self.release.set() + + class FakeController: + def __init__(self, **kwargs): + created.update(kwargs) + self.process = SimpleNamespace(returncode=0) + + async def run(self, callback): + await callback(object()) + + async def close(self): + return None + + monkeypatch.setattr(connector_module, 'Handler', FakeHandler) + monkeypatch.setattr( + 'langbot_plugin.runtime.io.controllers.stdio.client.StdioClientController', + FakeController, + ) + connector = BoxRuntimeConnector(make_app(Mock())) + + await connector.initialize() + + assert created['capture_stderr'] is False + assert connector._handler is not None + await connector.aclose() + + +@pytest.mark.asyncio +async def test_box_disconnect_notifies_once_and_clears_handler( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux') + monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False) + disconnect = AsyncMock() + + class FakeHandler: + def __init__(self, connection): + pass + + async def call_action(self, action, data): + return None + + async def run(self): + return None + + async def close(self): + return None + + class FakeController: + def __init__(self, **kwargs): + self.process = SimpleNamespace(returncode=0) + + async def run(self, callback): + await callback(object()) + + async def close(self): + return None + + monkeypatch.setattr(connector_module, 'Handler', FakeHandler) + monkeypatch.setattr( + 'langbot_plugin.runtime.io.controllers.stdio.client.StdioClientController', + FakeController, + ) + connector = BoxRuntimeConnector(make_app(Mock()), runtime_disconnect_callback=disconnect) + + await connector.initialize() + await asyncio.sleep(0) + + disconnect.assert_awaited_once_with(connector) + assert connector._handler is None + await connector.aclose() diff --git a/tests/unit_tests/box/test_box_service.py b/tests/unit_tests/box/test_box_service.py index 4d66ec8f8..4737dab15 100644 --- a/tests/unit_tests/box/test_box_service.py +++ b/tests/unit_tests/box/test_box_service.py @@ -325,10 +325,50 @@ 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_service_reconnect_restores_workspace_and_runs_cleanup( + monkeypatch: pytest.MonkeyPatch, +): + app = make_app(Mock()) + app.skill_mgr = SimpleNamespace(reload_skills=AsyncMock()) + service = BoxService(app, client=Mock(spec=BoxRuntimeClient)) + connector = Mock() + connector.reconnect = AsyncMock() + service._ensure_default_workspace = Mock() + service._purge_attachment_dirs = AsyncMock() + monkeypatch.setattr('langbot.pkg.box.service.asyncio.sleep', AsyncMock()) + + await service._reconnect_loop(connector) + + connector.reconnect.assert_awaited_once() + service._ensure_default_workspace.assert_called_once() + service._purge_attachment_dirs.assert_awaited_once() + app.skill_mgr.reload_skills.assert_awaited_once() + assert service.available is True + + @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/plugin/test_connector_ping.py b/tests/unit_tests/plugin/test_connector_ping.py index 766e51f8d..0a8e74207 100644 --- a/tests/unit_tests/plugin/test_connector_ping.py +++ b/tests/unit_tests/plugin/test_connector_ping.py @@ -1,15 +1,20 @@ from __future__ import annotations +import asyncio from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock import pytest +from langbot.pkg.plugin import connector as connector_module from langbot.pkg.plugin.connector import PluginRuntimeConnector, PluginRuntimeNotConnectedError def make_connector() -> PluginRuntimeConnector: - app = SimpleNamespace(instance_config=SimpleNamespace(data={'plugin': {'enable': True}})) + app = SimpleNamespace( + logger=Mock(), + instance_config=SimpleNamespace(data={'plugin': {'enable': True}, 'space': {'url': ''}}), + ) return PluginRuntimeConnector(app, AsyncMock()) @@ -30,3 +35,133 @@ async def test_ping_plugin_runtime_delegates_to_connected_handler(): assert result == 'pong' connector.handler.ping.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_stop_transport_tolerates_handler_callback_removing_attribute(): + connector = make_connector() + + class Handler: + async def close(self): + del connector.handler + + connector.handler = Handler() + + await connector._stop_transport() + + assert not hasattr(connector, 'handler') + + +@pytest.mark.asyncio +async def test_stdio_runtime_connection_does_not_capture_unconsumed_stderr( + monkeypatch: pytest.MonkeyPatch, +): + connector = make_connector() + created = {} + + class FakeRuntimeHandler: + def __init__(self, connection, disconnect_callback, ap): + self.release = asyncio.Event() + + async def ping(self): + return None + + async def set_runtime_config(self, **kwargs): + return None + + async def run(self): + await self.release.wait() + + async def close(self): + self.release.set() + + class FakeController: + def __init__(self, **kwargs): + created.update(kwargs) + + async def run(self, callback): + await callback(object()) + + async def close(self): + return None + + monkeypatch.setattr(connector_module.platform, 'get_platform', lambda: 'linux') + monkeypatch.setattr( + connector_module.platform, + 'use_websocket_to_connect_plugin_runtime', + lambda: False, + ) + monkeypatch.setattr( + connector_module.handler, + 'RuntimeConnectionHandler', + FakeRuntimeHandler, + ) + monkeypatch.setattr( + connector_module.stdio_client_controller, + 'StdioClientController', + FakeController, + ) + + await connector.initialize() + + assert created['capture_stderr'] is False + assert connector._connected.is_set() + await connector.aclose() + + +@pytest.mark.asyncio +async def test_runtime_disconnect_notifies_once_and_clears_handler( + monkeypatch: pytest.MonkeyPatch, +): + disconnect = AsyncMock() + connector = PluginRuntimeConnector(make_connector().ap, disconnect) + + class FakeRuntimeHandler: + def __init__(self, connection, disconnect_callback, ap): + self.disconnect_callback = disconnect_callback + + async def ping(self): + return None + + async def set_runtime_config(self, **kwargs): + return None + + async def run(self): + await self.disconnect_callback(self) + + async def close(self): + return None + + class FakeController: + def __init__(self, **kwargs): + pass + + async def run(self, callback): + await callback(object()) + + async def close(self): + return None + + monkeypatch.setattr(connector_module.platform, 'get_platform', lambda: 'linux') + monkeypatch.setattr( + connector_module.platform, + 'use_websocket_to_connect_plugin_runtime', + lambda: False, + ) + monkeypatch.setattr( + connector_module.handler, + 'RuntimeConnectionHandler', + FakeRuntimeHandler, + ) + monkeypatch.setattr( + connector_module.stdio_client_controller, + 'StdioClientController', + FakeController, + ) + + await connector.initialize() + await asyncio.sleep(0) + + disconnect.assert_awaited_once_with(connector) + assert not hasattr(connector, 'handler') + await connector.aclose() diff --git a/tests/unit_tests/provider/test_mcp_box_integration.py b/tests/unit_tests/provider/test_mcp_box_integration.py index 6b0c141ed..caf06bf6d 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,14 +680,55 @@ 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_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. - ``_init_stdio_python_server`` will raise a clear refusal at start - time; until then, the runtime info simply omits box_session_id so the - UI can render the disabled state cleanly.""" + 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_waits_for_enabled_box_when_temporarily_unavailable(self, mcp_module): ap = _make_ap() ap.box_service.available = False + ap.box_service.enabled = True s = _make_session( mcp_module, { @@ -700,9 +741,125 @@ class TestGetRuntimeInfoDict: ap=ap, ) info = s.get_runtime_info_dict() + assert info['box_session_id'] == 'mcp-shared' + assert info['box_enabled'] is True + + def test_stdio_session_refuses_when_box_is_disabled(self, mcp_module): + ap = _make_ap() + ap.box_service.available = False + ap.box_service.enabled = False + s = _make_session( + mcp_module, + { + 'name': 'test', + 'uuid': 'test-uuid', + 'mode': 'stdio', + 'command': 'python', + 'args': [], + }, + ap=ap, + ) + + info = s.get_runtime_info_dict() + assert 'box_session_id' not in info assert 'box_enabled' not in info + @pytest.mark.asyncio + async def test_stdio_session_waits_until_box_reconnects(self, mcp_module, monkeypatch): + mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio'] + ap = _make_ap() + ap.box_service.available = False + ap.box_service.enabled = True + session = _make_session( + mcp_module, + { + 'name': 'test', + 'uuid': 'test-uuid', + 'mode': 'stdio', + 'command': 'python', + 'args': [], + 'box': {'startup_timeout_sec': 1}, + }, + ap=ap, + ) + + async def reconnect_box(_delay): + ap.box_service.available = True + + monkeypatch.setattr(mcp_stdio_module.asyncio, 'sleep', reconnect_box) + + await session._box_stdio_runtime._wait_for_box_runtime() + + assert ap.box_service.available is True + + @pytest.mark.asyncio + async def test_enabled_box_timeout_does_not_exhaust_mcp_retry_budget(self, mcp_module, monkeypatch): + ap = _make_ap() + ap.box_service.available = False + ap.box_service.enabled = True + session = _make_session( + mcp_module, + { + 'name': 'test', + 'uuid': 'test-uuid', + 'mode': 'stdio', + 'command': 'python', + 'args': [], + }, + ap=ap, + ) + attempts = 0 + + async def lifecycle(): + nonlocal attempts + attempts += 1 + if attempts == 1: + session.error_phase = mcp_module.MCPSessionErrorPhase.BOX_UNAVAILABLE + raise RuntimeError('Box runtime is not available after 1 seconds') + + session._lifecycle_loop = lifecycle + sleep = AsyncMock() + monkeypatch.setattr(mcp_module.asyncio, 'sleep', sleep) + + await session._lifecycle_loop_with_retry() + + assert attempts == 2 + assert session.retry_count == 0 + sleep.assert_awaited_once_with(1) + + @pytest.mark.asyncio + async def test_disabled_box_still_stops_mcp_retry_loop(self, mcp_module): + ap = _make_ap() + ap.box_service.available = False + ap.box_service.enabled = False + session = _make_session( + mcp_module, + { + 'name': 'test', + 'uuid': 'test-uuid', + 'mode': 'stdio', + 'command': 'python', + 'args': [], + }, + ap=ap, + ) + attempts = 0 + + async def lifecycle(): + nonlocal attempts + attempts += 1 + session.error_phase = mcp_module.MCPSessionErrorPhase.BOX_UNAVAILABLE + raise RuntimeError('box_disabled_in_config') + + session._lifecycle_loop = lifecycle + + await session._lifecycle_loop_with_retry() + + assert attempts == 1 + assert session.status == mcp_module.MCPSessionStatus.ERROR + assert session.retry_count == 1 + def test_stdio_session_without_box_service_uses_local_stdio(self, mcp_module): ap = _make_ap() del ap.box_service diff --git a/tests/unit_tests/provider/test_mcp_remote_transport.py b/tests/unit_tests/provider/test_mcp_remote_transport.py index b2f1d2e14..b7761ec8a 100644 --- a/tests/unit_tests/provider/test_mcp_remote_transport.py +++ b/tests/unit_tests/provider/test_mcp_remote_transport.py @@ -12,7 +12,7 @@ import pytest from aiohttp import web from mcp import types as mcp_types -from langbot.pkg.provider.tools.loaders.mcp import RuntimeMCPSession +from langbot.pkg.provider.tools.loaders.mcp import MCPToolCallTimeoutError, RuntimeMCPSession class _TransportProbe: @@ -65,6 +65,25 @@ class _TransportProbe: }, } ) + if method == 'tools/call': + tool_name = message.get('params', {}).get('name') + if tool_name == 'hang': + return web.Response(status=202) + return web.json_response( + { + 'jsonrpc': '2.0', + 'id': message['id'], + 'result': { + 'content': [ + { + 'type': 'text', + 'text': 'healthy', + } + ], + 'isError': False, + }, + } + ) return web.Response(status=202) return web.Response(status=self.streamable_status) @@ -126,11 +145,22 @@ async def _transport_server(streamable_status: int | None): await runner.cleanup() -def _session(url: str, *, timeout: float = 2) -> RuntimeMCPSession: +def _session( + url: str, + *, + timeout: float = 2, + tool_call_timeout_sec: float = 300, +) -> RuntimeMCPSession: app = cast(Any, SimpleNamespace(logger=Mock())) return RuntimeMCPSession( 'remote-transport-test', - {'uuid': 'srv-1', 'mode': 'remote', 'url': url, 'timeout': timeout}, + { + 'uuid': 'srv-1', + 'mode': 'remote', + 'url': url, + 'timeout': timeout, + 'tool_call_timeout_sec': tool_call_timeout_sec, + }, True, app, ) @@ -164,6 +194,24 @@ async def test_remote_transport_real_streamable_http_success_keeps_session_usabl await _close_session(session) +@pytest.mark.asyncio +async def test_remote_transport_tool_timeout_does_not_poison_session(): + async with _transport_server(200) as (probe, url): + session = _session(url, tool_call_timeout_sec=0.05) + try: + await session._init_remote_server() + + with pytest.raises(MCPToolCallTimeoutError, match='timed out after 0.05 seconds'): + await session.invoke_mcp_tool('hang', {}) + + result = await session.invoke_mcp_tool('health_check', {}) + + assert result[0].text == 'healthy' + assert probe.streamable_messages.count('tools/call') == 2 + finally: + await _close_session(session) + + @pytest.mark.asyncio @pytest.mark.parametrize('status_code', [400, 404, 405]) async def test_remote_transport_real_streamable_http_error_falls_back_to_legacy_sse(status_code: int): diff --git a/tests/unit_tests/provider/test_mcp_resources.py b/tests/unit_tests/provider/test_mcp_resources.py index 773efe71c..cee241b75 100644 --- a/tests/unit_tests/provider/test_mcp_resources.py +++ b/tests/unit_tests/provider/test_mcp_resources.py @@ -1,20 +1,25 @@ from __future__ import annotations +import asyncio import base64 +from datetime import timedelta from types import SimpleNamespace from unittest.mock import AsyncMock, Mock import httpx import pytest from mcp import types as mcp_types +from mcp.shared.exceptions import McpError from langbot.pkg.provider.tools.loaders.mcp import ( MCP_RESOURCE_CONTEXT_QUERY_KEY, MCP_RESOURCE_TRACE_QUERY_KEY, + MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS, MCP_TOOL_LIST_RESOURCES, MCP_TOOL_READ_RESOURCE, MCPLoader, MCPSessionStatus, + MCPToolCallTimeoutError, RuntimeMCPSession, ) from langbot.pkg.telemetry import features as telemetry_features @@ -61,6 +66,111 @@ def _http_status_error(status_code: int) -> httpx.HTTPStatusError: return httpx.HTTPStatusError(f'HTTP {status_code}', request=request, response=response) +def _tool_result(text: str = 'ok') -> mcp_types.CallToolResult: + return mcp_types.CallToolResult( + content=[mcp_types.TextContent(type='text', text=text)], + isError=False, + ) + + +@pytest.mark.asyncio +async def test_invoke_mcp_tool_uses_configurable_request_timeout(): + session = RuntimeMCPSession( + 'slow-tools', + { + 'uuid': 'srv-1', + 'mode': 'remote', + 'tool_call_timeout_sec': 900, + }, + True, + _app(), + ) + session.session = SimpleNamespace(call_tool=AsyncMock(return_value=_tool_result())) + + result = await session.invoke_mcp_tool('render_video', {'quality': 'high'}) + + assert result[0].text == 'ok' + session.session.call_tool.assert_awaited_once_with( + 'render_video', + {'quality': 'high'}, + read_timeout_seconds=timedelta(seconds=900), + ) + + +@pytest.mark.asyncio +async def test_invoke_mcp_tool_zero_timeout_disables_request_deadline(): + session = RuntimeMCPSession( + 'unbounded-tools', + { + 'uuid': 'srv-1', + 'mode': 'remote', + 'tool_call_timeout_sec': 0, + }, + True, + _app(), + ) + session.session = SimpleNamespace(call_tool=AsyncMock(return_value=_tool_result())) + + await session.invoke_mcp_tool('long_job', {}) + + session.session.call_tool.assert_awaited_once_with( + 'long_job', + {}, + read_timeout_seconds=None, + ) + + +@pytest.mark.asyncio +async def test_invoke_mcp_tool_timeout_is_not_retried_and_session_remains_usable(): + session = RuntimeMCPSession( + 'recoverable-tools', + { + 'uuid': 'srv-1', + 'mode': 'remote', + 'tool_call_timeout_sec': 5, + }, + True, + _app(), + ) + timeout = McpError( + mcp_types.ErrorData( + code=httpx.codes.REQUEST_TIMEOUT, + message='Timed out while waiting for response to ClientRequest. Waited 5 seconds.', + ) + ) + call_tool = AsyncMock(side_effect=[timeout, _tool_result('recovered')]) + session.session = SimpleNamespace(call_tool=call_tool) + + with pytest.raises( + MCPToolCallTimeoutError, + match="MCP tool 'long_job' on server 'recoverable-tools' timed out after 5 seconds", + ): + await session.invoke_mcp_tool('long_job', {}) + + assert call_tool.await_count == 1 + second_result = await session.invoke_mcp_tool('health_check', {}) + assert second_result[0].text == 'recovered' + assert call_tool.await_count == 2 + + +@pytest.mark.parametrize('invalid_timeout', [-1, float('inf'), 1e300, True, 'not-a-number']) +def test_invalid_tool_call_timeout_falls_back_to_default(invalid_timeout): + ap = _app() + session = RuntimeMCPSession( + 'invalid-timeout', + { + 'uuid': 'srv-1', + 'mode': 'remote', + 'tool_call_timeout_sec': invalid_timeout, + }, + True, + ap, + ) + + assert session.tool_call_timeout_sec == MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS + ap.logger.warning.assert_called_once() + + @pytest.mark.asyncio async def test_remote_transport_falls_back_to_sse_for_compatible_http_status_in_exception_group(): session = RuntimeMCPSession( @@ -321,3 +431,42 @@ async def test_build_resource_context_for_query_uses_only_bound_attached_text_re assert query.variables[MCP_RESOURCE_CONTEXT_QUERY_KEY]['resource_count'] == 1 docs.session.read_resource.assert_awaited_once() other.session.read_resource.assert_not_called() + + +@pytest.mark.asyncio +async def test_mcp_loader_shutdown_cancels_startup_tasks_and_closes_sessions_concurrently(): + loader = MCPLoader(_app()) + hosted_cancelled = asyncio.Event() + + async def pending_host(): + try: + await asyncio.Event().wait() + finally: + hosted_cancelled.set() + + hosted_task = asyncio.create_task(pending_host()) + await asyncio.sleep(0) + loader._hosted_mcp_tasks = [hosted_task] + + started: set[str] = set() + all_started = asyncio.Event() + + class Session: + def __init__(self, name: str): + self.name = name + + async def shutdown(self): + started.add(self.name) + if len(started) == 2: + all_started.set() + await all_started.wait() + + loader.sessions = {'one': Session('one'), 'two': Session('two')} + + await asyncio.wait_for(loader.shutdown(), timeout=1) + + assert hosted_cancelled.is_set() + assert hosted_task.cancelled() + assert started == {'one', 'two'} + assert loader._hosted_mcp_tasks == [] + assert loader.sessions == {} diff --git a/tests/unit_tests/provider/test_skill_tools.py b/tests/unit_tests/provider/test_skill_tools.py index 9db7b945e..405f11d17 100644 --- a/tests/unit_tests/provider/test_skill_tools.py +++ b/tests/unit_tests/provider/test_skill_tools.py @@ -458,6 +458,25 @@ class TestSkillToolLoader: assert await loader.has_tool('activate') is True assert await loader.has_tool('register_skill') is True + @pytest.mark.asyncio + async def test_tools_reappear_after_box_backend_recovers(self): + from langbot.pkg.provider.tools.loaders.skill_authoring import SkillToolLoader + + ap = _make_ap() + ap.skill_mgr = SimpleNamespace(skills={'demo': _make_skill_data(name='demo')}) + ap.box_service = SimpleNamespace( + available=False, + get_status=AsyncMock(return_value={'backend': {'available': True}}), + ) + + loader = SkillToolLoader(ap) + await loader.initialize() + assert await loader.get_tools() == [] + + ap.box_service.available = True + + assert sorted(tool.name for tool in await loader.get_tools()) == ['activate', 'register_skill'] + class TestNativeToolLoaderSkillPaths: @pytest.mark.asyncio diff --git a/tests/unit_tests/provider/test_tool_manager_native.py b/tests/unit_tests/provider/test_tool_manager_native.py index 2085a9e8c..d73c66e4f 100644 --- a/tests/unit_tests/provider/test_tool_manager_native.py +++ b/tests/unit_tests/provider/test_tool_manager_native.py @@ -151,6 +151,21 @@ async def test_native_tool_loader_exposes_all_tools_when_box_available(): assert await loader.has_tool(tool_name) is True +@pytest.mark.asyncio +async def test_native_tool_loader_refreshes_after_box_recovers(): + box_service = SimpleNamespace( + available=False, + get_status=AsyncMock(return_value={'backend': {'available': True}}), + ) + loader = NativeToolLoader(SimpleNamespace(box_service=box_service, logger=Mock())) + await loader.initialize() + assert await loader.get_tools() == [] + + box_service.available = True + + assert [tool.name for tool in await loader.get_tools()] == ['exec', 'read', 'write', 'edit', 'glob', 'grep'] + + # ── read/write/edit file tool tests ───────────────────────────── diff --git a/uv.lock b/uv.lock index 6ddc22e3e..cc71d0888 100644 --- a/uv.lock +++ b/uv.lock @@ -2124,7 +2124,7 @@ requires-dist = [ { name = "ebooklib", specifier = ">=0.18" }, { name = "gewechat-client", specifier = ">=0.1.5" }, { name = "html2text", specifier = ">=2024.2.26" }, - { name = "langbot-plugin", specifier = "==0.4.13" }, + { name = "langbot-plugin", specifier = "==0.4.17" }, { name = "langchain", specifier = ">=1.3.9" }, { name = "langchain-core", specifier = ">=1.3.3" }, { name = "langchain-text-splitters", specifier = ">=1.1.2" }, @@ -2189,7 +2189,7 @@ dev = [ [[package]] name = "langbot-plugin" -version = "0.4.13" +version = "0.4.17" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiofiles" }, @@ -2210,9 +2210,9 @@ dependencies = [ { name = "watchdog" }, { name = "websockets" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/40/a6/1eaf77c3b81e9de3390c504c5f627dc41f43bff6df9aff0e1e31d796b6f0/langbot_plugin-0.4.13.tar.gz", hash = "sha256:f936340e67679c21f1e7e7f1447339f31a0a2c965db060ecfbd9d0c51bb0d6fe", size = 334887, upload-time = "2026-07-04T05:38:59.942Z" } +sdist = { url = "https://files.pythonhosted.org/packages/3a/96/336d7ac97ff5a9c413e7aaf0c329a7047ec91dec5fa4f9e0b1d0abb57d0e/langbot_plugin-0.4.17.tar.gz", hash = "sha256:b1539444f16568c0b3244f68c498ac9498c56e2c20f12e220e6ee735b8d6b5c2", size = 347080, upload-time = "2026-07-23T09:23:55.941Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e9/bf/fc9671a7afbd933440c38403c84d918c1022fdeed16e22a6ab3b2aec83ff/langbot_plugin-0.4.13-py3-none-any.whl", hash = "sha256:9d45ebc7a7ee0413d6db9baa009fcbf0ad07e2e1753a6f0a27f37b8b665cd1ee", size = 221884, upload-time = "2026-07-04T05:38:58.525Z" }, + { url = "https://files.pythonhosted.org/packages/e1/98/ac372a238e59894ac54189cb220a3fb75da48f319881eeb49d583712549a/langbot_plugin-0.4.17-py3-none-any.whl", hash = "sha256:6fd6c9d7e0583f04659f023a233c4201cb4c450d1ef696e79fa1e338933ee156", size = 229881, upload-time = "2026-07-23T09:23:54.582Z" }, ] [[package]] diff --git a/web/src/app/home/mcp/components/mcp-form/MCPForm.tsx b/web/src/app/home/mcp/components/mcp-form/MCPForm.tsx index 753e7a73c..058946395 100644 --- a/web/src/app/home/mcp/components/mcp-form/MCPForm.tsx +++ b/web/src/app/home/mcp/components/mcp-form/MCPForm.tsx @@ -435,6 +435,10 @@ const getFormSchema = (t: TFunction) => .number({ invalid_type_error: t('mcp.timeoutMustBeNumber') }) .positive({ message: t('mcp.timeoutMustBePositive') }) .default(30), + tool_call_timeout_sec: z + .number({ invalid_type_error: t('mcp.timeoutMustBeNumber') }) + .nonnegative({ message: t('mcp.timeoutNonNegative') }) + .default(300), ssereadtimeout: z .number({ invalid_type_error: t('mcp.sseTimeoutMustBeNumber') }) .positive({ message: t('mcp.timeoutMustBePositive') }) @@ -474,6 +478,7 @@ const getFormSchema = (t: TFunction) => type FormValues = z.infer> & { timeout: number; + tool_call_timeout_sec: number; ssereadtimeout: number; }; @@ -535,6 +540,7 @@ const MCPForm = forwardRef(function MCPForm( command: '', args: [], timeout: 30, + tool_call_timeout_sec: 300, ssereadtimeout: 300, extra_args: [], ...initialDraftRef.current, @@ -609,6 +615,7 @@ const MCPForm = forwardRef(function MCPForm( command: '', args: [], timeout: 30, + tool_call_timeout_sec: 300, ssereadtimeout: 300, extra_args: [], ...initialDraftRef.current, @@ -687,10 +694,16 @@ const MCPForm = forwardRef(function MCPForm( command: '', args: [], timeout: 30, + tool_call_timeout_sec: 300, ssereadtimeout: 300, extra_args: [], }; + if (typeof server.extra_args.tool_call_timeout_sec === 'number') { + formValues.tool_call_timeout_sec = + server.extra_args.tool_call_timeout_sec; + } + let newExtraArgs: { key: string; type: 'string' | 'number' | 'boolean'; @@ -770,6 +783,7 @@ const MCPForm = forwardRef(function MCPForm( url: value.url!, headers, timeout: value.timeout, + tool_call_timeout_sec: value.tool_call_timeout_sec, }, }; } else { @@ -786,6 +800,7 @@ const MCPForm = forwardRef(function MCPForm( command: value.command!, args: value.args?.map((arg) => arg.value) || [], env, + tool_call_timeout_sec: value.tool_call_timeout_sec, }, }; } @@ -835,6 +850,7 @@ const MCPForm = forwardRef(function MCPForm( extraArgsData = { url: form.getValues('url')!, timeout: form.getValues('timeout'), + tool_call_timeout_sec: form.getValues('tool_call_timeout_sec'), headers: Object.fromEntries( formExtraArgs.map((arg) => [arg.key, arg.value]), ), @@ -846,6 +862,7 @@ const MCPForm = forwardRef(function MCPForm( env: Object.fromEntries( formExtraArgs.map((arg) => [arg.key, arg.value]), ), + tool_call_timeout_sec: form.getValues('tool_call_timeout_sec'), }; } @@ -1047,6 +1064,30 @@ const MCPForm = forwardRef(function MCPForm( )} /> + ( + + {t('mcp.toolCallTimeout')} + + field.onChange(Number(e.target.value))} + /> + + + {t('mcp.toolCallTimeoutDescription')} + + + + )} + /> + {watchMode === 'remote' && ( <> ; timeout: number; ssereadtimeout: number; + tool_call_timeout_sec?: number; } export interface MCPServerExtraArgsStdio { command: string; args: string[]; env: Record; + tool_call_timeout_sec?: number; } export interface MCPServerExtraArgsHttp { url: string; headers: Record; timeout: number; + tool_call_timeout_sec?: number; } // "remote" mode: the user only supplies a URL; the backend auto-detects the @@ -540,6 +543,7 @@ export interface MCPServerExtraArgsRemote { url: string; headers?: Record; timeout?: number; + tool_call_timeout_sec?: number; } export enum MCPSessionStatus { diff --git a/web/src/i18n/locales/en-US.ts b/web/src/i18n/locales/en-US.ts index 2f9a71077..dbd1aea14 100644 --- a/web/src/i18n/locales/en-US.ts +++ b/web/src/i18n/locales/en-US.ts @@ -797,6 +797,9 @@ const enUS = { url: 'URL', headers: 'Headers', timeout: 'Timeout', + toolCallTimeout: 'Tool call timeout (seconds)', + toolCallTimeoutDescription: + 'Maximum wait for one tool call. Set to 0 for no timeout. Defaults to 300 seconds.', addArgument: 'Add Argument', addEnvVar: 'Add Environment Variable', addHeader: 'Add Header', diff --git a/web/src/i18n/locales/es-ES.ts b/web/src/i18n/locales/es-ES.ts index 6d25d0112..74f1985f1 100644 --- a/web/src/i18n/locales/es-ES.ts +++ b/web/src/i18n/locales/es-ES.ts @@ -811,6 +811,9 @@ const esES = { url: 'URL', headers: 'Encabezados', timeout: 'Tiempo de espera', + toolCallTimeout: 'Tiempo de espera de herramienta (segundos)', + toolCallTimeoutDescription: + 'Espera máxima para una llamada de herramienta. Use 0 para no limitar. El valor predeterminado es 300 segundos.', addArgument: 'Añadir argumento', addEnvVar: 'Añadir variable de entorno', addHeader: 'Añadir encabezado', diff --git a/web/src/i18n/locales/ja-JP.ts b/web/src/i18n/locales/ja-JP.ts index e3d59c55f..614891d98 100644 --- a/web/src/i18n/locales/ja-JP.ts +++ b/web/src/i18n/locales/ja-JP.ts @@ -803,6 +803,9 @@ const jaJP = { url: 'URL', headers: 'ヘッダー', timeout: 'タイムアウト', + toolCallTimeout: 'ツール呼び出しタイムアウト(秒)', + toolCallTimeoutDescription: + '1 回のツール呼び出しの最大待機時間です。0 で無制限、既定値は 300 秒です。', addArgument: '引数を追加', addEnvVar: '環境変数を追加', addHeader: 'ヘッダーを追加', diff --git a/web/src/i18n/locales/ru-RU.ts b/web/src/i18n/locales/ru-RU.ts index 0f6072c49..eef09cdee 100644 --- a/web/src/i18n/locales/ru-RU.ts +++ b/web/src/i18n/locales/ru-RU.ts @@ -808,6 +808,9 @@ const ruRU = { url: 'URL', headers: 'Заголовки', timeout: 'Таймаут', + toolCallTimeout: 'Таймаут вызова инструмента (секунды)', + toolCallTimeoutDescription: + 'Максимальное ожидание одного вызова. 0 отключает ограничение. По умолчанию 300 секунд.', addArgument: 'Добавить аргумент', addEnvVar: 'Добавить переменную окружения', addHeader: 'Добавить заголовок', diff --git a/web/src/i18n/locales/th-TH.ts b/web/src/i18n/locales/th-TH.ts index 704daf9cd..4ced7ace4 100644 --- a/web/src/i18n/locales/th-TH.ts +++ b/web/src/i18n/locales/th-TH.ts @@ -786,6 +786,9 @@ const thTH = { url: 'URL', headers: 'ส่วนหัว', timeout: 'หมดเวลา', + toolCallTimeout: 'หมดเวลาการเรียกเครื่องมือ (วินาที)', + toolCallTimeoutDescription: + 'เวลารอสูงสุดต่อการเรียกเครื่องมือหนึ่งครั้ง ตั้งเป็น 0 เพื่อไม่จำกัด ค่าเริ่มต้นคือ 300 วินาที', addArgument: 'เพิ่มอาร์กิวเมนต์', addEnvVar: 'เพิ่มตัวแปรสภาพแวดล้อม', addHeader: 'เพิ่มส่วนหัว', diff --git a/web/src/i18n/locales/vi-VN.ts b/web/src/i18n/locales/vi-VN.ts index b9770314e..26c415688 100644 --- a/web/src/i18n/locales/vi-VN.ts +++ b/web/src/i18n/locales/vi-VN.ts @@ -801,6 +801,9 @@ const viVN = { url: 'URL', headers: 'Tiêu đề', timeout: 'Thời gian chờ', + toolCallTimeout: 'Thời gian chờ gọi công cụ (giây)', + toolCallTimeoutDescription: + 'Thời gian chờ tối đa cho một lần gọi công cụ. Đặt 0 để không giới hạn. Mặc định là 300 giây.', addArgument: 'Thêm tham số', addEnvVar: 'Thêm biến môi trường', addHeader: 'Thêm tiêu đề', diff --git a/web/src/i18n/locales/zh-Hans.ts b/web/src/i18n/locales/zh-Hans.ts index a4a04e0de..96e24759d 100644 --- a/web/src/i18n/locales/zh-Hans.ts +++ b/web/src/i18n/locales/zh-Hans.ts @@ -763,6 +763,9 @@ const zhHans = { url: 'URL地址', headers: '请求头', timeout: '超时时间', + toolCallTimeout: '工具调用超时(秒)', + toolCallTimeoutDescription: + '单次工具调用的最长等待时间。设为 0 表示不限制,默认 300 秒。', addArgument: '添加参数', addEnvVar: '添加环境变量', addHeader: '添加请求头', diff --git a/web/src/i18n/locales/zh-Hant.ts b/web/src/i18n/locales/zh-Hant.ts index 1ae98f758..57f0bc58e 100644 --- a/web/src/i18n/locales/zh-Hant.ts +++ b/web/src/i18n/locales/zh-Hant.ts @@ -762,6 +762,9 @@ const zhHant = { url: 'URL位址', headers: '請求標頭', timeout: '逾時時間', + toolCallTimeout: '工具呼叫逾時(秒)', + toolCallTimeoutDescription: + '單次工具呼叫的最長等待時間。設為 0 表示不限制,預設 300 秒。', addArgument: '新增參數', addEnvVar: '新增環境變數', addHeader: '新增請求標頭',