xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-07-21 11:56:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(runtime): make plugin and box connectors resilient

Browse Source
This commit is contained in:
Junyan Qin
2026-07-21 18:41:22 +08:00
parent 76c5003c21
commit 0b461e5830
15 changed files with 699 additions and 236 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docker/kubernetes.yaml
+12 -4
View File
@@ -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
src/langbot/pkg/box/connector.py
+115 -24
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import asyncio
import contextlib
import json
import os
import sys
@@ -28,6 +29,7 @@ _DOCKER_BOX_HOST = 'langbot_box'
_DEFAULT_PORT = 5410
_HEARTBEAT_INTERVAL_SEC = 20
_HEARTBEAT_FAILURE_THRESHOLD = 3
# Top-level keys under ``box`` that are LangBot-internal and should not be
# forwarded to the Box runtime.
@@ -113,6 +115,8 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
self._handler_task: asyncio.Task | None = None
self._ctrl_task: asyncio.Task | None = None
self._heartbeat_task: asyncio.Task | None = None
self._ctrl = None
self._generation = 0
# Parse the relay URL once for reuse.
parsed = urlparse(self.ws_relay_base_url)
@@ -145,29 +149,64 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
return self.uses_websocket()
async def initialize(self) -> None:
if self._uses_websocket():
if platform.get_platform() == 'win32' and not self.configured_runtime_endpoint:
await self._start_subprocess_then_ws()
else:
await self._connect_remote_ws()
else:
await self._start_local_stdio()
async with self._lifecycle_lock:
if self._closing:
raise BoxRuntimeUnavailableError('box runtime connector is shutting down')
await self._stop_transport()
self._generation += 1
try:
if self._uses_websocket():
if platform.get_platform() == 'win32' and not self.configured_runtime_endpoint:
await self._start_subprocess_then_ws()
else:
await self._connect_remote_ws()
else:
await self._start_local_stdio()
except BaseException:
await self._stop_transport()
await self._close_managed_subprocess()
raise
# Start heartbeat after successful connection
if self._heartbeat_task is None:
self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
if self._heartbeat_task is None or self._heartbeat_task.done():
self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
async def reconnect(self) -> None:
async with self._lifecycle_lock:
if self._closing:
raise BoxRuntimeUnavailableError('box runtime connector is shutting down')
await self._stop_transport()
self._generation += 1
try:
if self._uses_websocket():
if platform.get_platform() == 'win32' and not self.configured_runtime_endpoint:
await self._start_subprocess_then_ws()
else:
await self._connect_remote_ws()
else:
await self._start_local_stdio()
except BaseException:
await self._stop_transport()
await self._close_managed_subprocess()
raise
# -- heartbeat -----------------------------------------------------------
async def _heartbeat_loop(self) -> None:
"""Periodically ping the Box runtime to detect silent disconnections."""
while True:
failures = 0
while not self._closing:
await asyncio.sleep(_HEARTBEAT_INTERVAL_SEC)
try:
await self.ping()
failures = 0
self.ap.logger.debug('Heartbeat to Box runtime success.')
except Exception as e:
self.ap.logger.debug(f'Failed to heartbeat to Box runtime: {e}')
failures += 1
self.ap.logger.warning(f'Box runtime heartbeat failed ({failures}/{_HEARTBEAT_FAILURE_THRESHOLD}): {e}')
if failures >= _HEARTBEAT_FAILURE_THRESHOLD:
failures = 0
if self.runtime_disconnect_callback is not None:
await self.runtime_disconnect_callback(self)
async def ping(self) -> None:
if self._handler is None:
@@ -197,8 +236,9 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
args=['-m', 'langbot_plugin.cli.__init__', 'box', '-s', '--ws-control-port', str(self._relay_port)],
env=env,
)
self._ctrl = ctrl
self._ctrl_task = asyncio.create_task(
ctrl.run(self._make_connection_callback('stdio', connected, connect_error))
ctrl.run(self._make_connection_callback('stdio', connected, connect_error, self._generation))
)
try:
@@ -282,8 +322,9 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
await self.runtime_disconnect_callback(self)
ctrl = WebSocketClientController(ws_url=ws_url, make_connection_failed_callback=on_connect_failed)
self._ctrl = ctrl
self._ctrl_task = asyncio.create_task(
ctrl.run(self._make_connection_callback(transport_name, connected, connect_error))
ctrl.run(self._make_connection_callback(transport_name, connected, connect_error, self._generation))
)
try:
@@ -299,8 +340,12 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
transport_name: str,
connected: asyncio.Event,
connect_error: list[Exception],
generation: int,
):
async def new_connection_callback(connection: Connection) -> None:
if generation != self._generation or self._closing:
await connection.close()
return
handler = Handler(connection)
self._handler = handler
self.client.set_handler(handler)
@@ -313,6 +358,8 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
self.ap.logger.info(f'Connected to Box runtime via {transport_name}.')
connected.set()
await self._handler_task
except asyncio.CancelledError:
raise
except Exception as exc:
if not connected.is_set():
connect_error.append(exc)
@@ -322,22 +369,66 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
# If we reach here, handler.run() returned normally (connection
# closed) or raised after the initial handshake succeeded.
# Either way, treat it as a disconnect.
if connected.is_set():
if self._uses_websocket():
self.ap.logger.error('Disconnected from Box runtime, trying to reconnect...')
if self.runtime_disconnect_callback is not None:
await self.runtime_disconnect_callback(self)
else:
self.ap.logger.error(
'Disconnected from Box runtime via stdio. '
'Cannot automatically reconnect — please restart LangBot.'
)
if (
connected.is_set()
and generation == self._generation
and not self._closing
and self.runtime_disconnect_callback is not None
):
self.ap.logger.error('Disconnected from Box runtime, trying to reconnect...')
await self.runtime_disconnect_callback(self)
return new_connection_callback
# -- lifecycle -----------------------------------------------------------
async def _stop_transport(self) -> None:
if self._handler is not None:
with contextlib.suppress(Exception):
await self._handler.close()
self.client.set_handler(None)
tasks = [
task
for task in (self._handler_task, self._ctrl_task)
if task is not None and task is not asyncio.current_task()
]
for task in tasks:
task.cancel()
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
close_ctrl = getattr(self._ctrl, 'close', None)
if close_ctrl is not None:
with contextlib.suppress(Exception):
await close_ctrl()
self._handler = None
self._handler_task = None
self._ctrl_task = None
self._ctrl = None
async def aclose(self) -> None:
self._closing = True
if self._heartbeat_task is not None:
self._heartbeat_task.cancel()
await asyncio.gather(self._heartbeat_task, return_exceptions=True)
self._heartbeat_task = None
await self._stop_transport()
process = getattr(self, '_subprocess', None)
if process is not None and process.returncode is None:
with contextlib.suppress(ProcessLookupError):
process.terminate()
try:
await asyncio.wait_for(process.wait(), timeout=3)
except asyncio.TimeoutError:
with contextlib.suppress(ProcessLookupError):
process.kill()
await process.wait()
self._subprocess = None
await self._close_managed_subprocess()
def dispose(self) -> None:
"""Best-effort synchronous compatibility wrapper; prefer ``aclose``."""
self._closing = True
if self._heartbeat_task is not None:
self._heartbeat_task.cancel()
self._heartbeat_task = None
src/langbot/pkg/box/service.py
+28 -9
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import asyncio
import collections
import contextlib
import datetime as _dt
import enum
import json
@@ -65,6 +66,8 @@ class BoxService:
self.workspace_quota_mb = self._load_workspace_quota_mb()
self._recent_errors: collections.deque[dict] = collections.deque(maxlen=_MAX_RECENT_ERRORS)
self._shutdown_task = None
self._reconnect_task: asyncio.Task | None = None
self._closing = False
self._available = False
self._connector_error: str = ''
self._reconnecting = False
@@ -110,6 +113,8 @@ class BoxService:
self.ap.logger.warning(f'LangBot Box runtime unavailable, sandbox features disabled: {exc}')
self._available = False
self._connector_error = str(exc)
if self._runtime_connector is not None:
await self._on_runtime_disconnect(self._runtime_connector)
async def _on_runtime_disconnect(self, connector: BoxRuntimeConnector) -> None:
"""Called by the connector when the Box runtime connection drops.
@@ -118,27 +123,26 @@ class BoxService:
Skipped entirely when Box is disabled by config — that path should
never have connected in the first place.
"""
if not self._enabled:
if not self._enabled or self._closing:
return
if self._reconnecting:
if self._reconnect_task is not None and not self._reconnect_task.done():
return # Another reconnect loop is already running
self._reconnecting = True
self._available = False
self._connector_error = 'Disconnected from Box runtime'
self.ap.logger.warning('Box runtime disconnected, sandbox features temporarily disabled.')
asyncio.create_task(self._reconnect_loop(connector))
self._reconnect_task = asyncio.create_task(self._reconnect_loop(connector))
async def _reconnect_loop(self, connector: BoxRuntimeConnector) -> None:
"""Retry reconnection with exponential backoff (3s → 60s max)."""
delay = 3
max_delay = 60
try:
while True:
while not self._closing:
self.ap.logger.info(f'Attempting to reconnect to Box runtime in {delay}s...')
await asyncio.sleep(delay)
try:
connector.dispose()
await connector.initialize()
await connector.reconnect()
self._available = True
self._connector_error = ''
self.ap.logger.info('Box runtime reconnected, sandbox features restored.')
@@ -149,6 +153,7 @@ class BoxService:
delay = min(delay * 2, max_delay)
finally:
self._reconnecting = False
self._reconnect_task = None
@property
def available(self) -> bool:
@@ -838,14 +843,28 @@ class BoxService:
return attachments
async def shutdown(self):
await self.client.shutdown()
if self._closing:
return
self._closing = True
self._available = False
reconnect_task = self._reconnect_task
self._reconnect_task = None
if reconnect_task is not None and reconnect_task is not asyncio.current_task():
reconnect_task.cancel()
await asyncio.gather(reconnect_task, return_exceptions=True)
# The runtime may already be offline. A failed best-effort SHUTDOWN RPC
# must not prevent us from cancelling transports and reaping children.
with contextlib.suppress(Exception):
await self.client.shutdown()
if self._runtime_connector is not None:
await self._runtime_connector.aclose()
def dispose(self):
if self._runtime_connector is not None:
self._runtime_connector.dispose()
loop = getattr(self.ap, 'event_loop', None)
if loop is not None and not loop.is_closed() and (self._shutdown_task is None or self._shutdown_task.done()):
self._shutdown_task = loop.create_task(self.shutdown())
elif self._runtime_connector is not None:
self._runtime_connector.dispose()
async def get_sessions(self) -> list[dict]:
if not self._available:
src/langbot/pkg/core/app.py
+35 -1
View File
@@ -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:
src/langbot/pkg/core/boot.py
+24 -7
View File
@@ -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()
src/langbot/pkg/core/stages/build_app.py
+8 -5
View File
@@ -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)
src/langbot/pkg/plugin/connector.py
+264 -153
View File
@@ -2,6 +2,7 @@
from __future__ import annotations
import asyncio
import contextlib
import io
import time
import zipfile
@@ -36,6 +37,12 @@ from ..core import taskmgr
from ..entity.persistence import plugin as persistence_plugin
_CONNECT_TIMEOUT_SEC = 30.0
_HEARTBEAT_INTERVAL_SEC = 20.0
_HEARTBEAT_FAILURE_THRESHOLD = 3
_RECONNECT_MAX_DELAY_SEC = 60.0
class PluginRuntimeNotConnectedError(RuntimeError):
"""Raised when plugin runtime operations are requested before connection."""
@@ -70,128 +77,218 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
super().__init__(ap)
self.runtime_disconnect_callback = runtime_disconnect_callback
self.is_enable_plugin = self.ap.instance_config.data.get('plugin', {}).get('enable', True)
self._transport_task: asyncio.Task | None = None
self._reconnect_task: asyncio.Task | None = None
self._generation = 0
self._connected = asyncio.Event()
def _runtime_handler(self) -> handler.RuntimeConnectionHandler:
runtime_handler = getattr(self, 'handler', None)
if runtime_handler is None:
raise PluginRuntimeNotConnectedError('Plugin runtime is not connected')
return runtime_handler
def _runtime_available(self) -> bool:
runtime_handler = getattr(self, 'handler', None)
if runtime_handler is None:
return False
# Unit-level and explicitly injected handlers don't own a transport.
# A managed transport must also have completed its handshake.
return self._transport_task is None or self._connected.is_set()
async def heartbeat_loop(self):
while True:
await asyncio.sleep(20)
failures = 0
while not self._closing:
await asyncio.sleep(_HEARTBEAT_INTERVAL_SEC)
try:
await self.ping_plugin_runtime()
failures = 0
self.ap.logger.debug('Heartbeat to plugin runtime success.')
except Exception as e:
self.ap.logger.debug(f'Failed to heartbeat to plugin runtime: {e}')
failures += 1
self.ap.logger.warning(
f'Plugin runtime heartbeat failed ({failures}/{_HEARTBEAT_FAILURE_THRESHOLD}): {e}'
)
if failures >= _HEARTBEAT_FAILURE_THRESHOLD:
self._connected.clear()
self.schedule_reconnect()
failures = 0
async def initialize(self):
if not self.is_enable_plugin:
self.ap.logger.info('Plugin system is disabled.')
return
async def new_connection_callback(connection: base_connection.Connection):
async def disconnect_callback(
rchandler: handler.RuntimeConnectionHandler,
) -> bool:
if platform.get_platform() == 'docker' or platform.use_websocket_to_connect_plugin_runtime():
self.ap.logger.error('Disconnected from plugin runtime, trying to reconnect...')
await self.runtime_disconnect_callback(self)
return False
else:
self.ap.logger.error(
'Disconnected from plugin runtime, cannot automatically reconnect while LangBot connects to plugin runtime via stdio, please restart LangBot.'
)
async with self._lifecycle_lock:
if self._closing:
raise PluginRuntimeNotConnectedError('Plugin runtime connector is shutting down')
if self._connected.is_set() and hasattr(self, 'handler'):
return
await self._stop_transport()
self._generation += 1
generation = self._generation
self._connected = asyncio.Event()
connect_errors: list[Exception] = []
async def new_connection_callback(
connection: base_connection.Connection,
):
if generation != self._generation or self._closing:
await connection.close()
return
async def disconnect_callback(
rchandler: handler.RuntimeConnectionHandler,
) -> bool:
if generation == self._generation and not self._closing:
self._connected.clear()
await self.runtime_disconnect_callback(self)
return False
self.handler = handler.RuntimeConnectionHandler(connection, disconnect_callback, self.ap)
self.handler_task = asyncio.create_task(self.handler.run())
_ = await self.handler.ping()
# Push the configured marketplace (Space) URL to the runtime so it
# downloads plugins from the same Space LangBot is bound to, rather
# than relying on the runtime's own env/default.
space_url = self.ap.instance_config.data.get('space', {}).get('url', '').rstrip('/')
if space_url:
runtime_handler = handler.RuntimeConnectionHandler(connection, disconnect_callback, self.ap)
self.handler = runtime_handler
self.handler_task = asyncio.create_task(runtime_handler.run())
try:
await self.handler.set_runtime_config(cloud_service_url=space_url)
self.ap.logger.info(f'Pushed marketplace URL to plugin runtime: {space_url}')
except Exception as e:
self.ap.logger.warning(f'Failed to push runtime config: {e}')
self.ap.logger.info('Connected to plugin runtime.')
await self.handler_task
await runtime_handler.ping()
space_url = self.ap.instance_config.data.get('space', {}).get('url', '').rstrip('/')
if space_url:
await runtime_handler.set_runtime_config(cloud_service_url=space_url)
if generation == self._generation and not self._closing:
self._connected.set()
self.ap.logger.info('Connected to plugin runtime.')
await self.handler_task
except asyncio.CancelledError:
raise
except Exception as exc:
if not self._connected.is_set():
connect_errors.append(exc)
self._connected.set()
finally:
if generation == self._generation and not self._closing:
self._connected.clear()
if getattr(self, 'handler', None) is runtime_handler:
del self.handler
await self.runtime_disconnect_callback(self)
task: asyncio.Task | None = None
task_coro: typing.Coroutine
if platform.get_platform() == 'docker' or platform.use_websocket_to_connect_plugin_runtime():
ws_url = self.ap.instance_config.data.get('plugin', {}).get(
'runtime_ws_url',
'ws://langbot_plugin_runtime:5400/control/ws',
)
if platform.get_platform() == 'docker' or platform.use_websocket_to_connect_plugin_runtime(): # use websocket
self.ap.logger.info('use websocket to connect to plugin runtime')
ws_url = self.ap.instance_config.data.get('plugin', {}).get(
'runtime_ws_url', 'ws://langbot_plugin_runtime:5400/control/ws'
async def connection_failed(ctrl, exc=None):
error = exc or RuntimeError('WebSocket connection failed')
connect_errors.append(error)
self._connected.set()
self.ctrl = ws_client_controller.WebSocketClientController(
ws_url=ws_url,
make_connection_failed_callback=connection_failed,
)
task_coro = self.ctrl.run(new_connection_callback)
elif platform.get_platform() == 'win32':
await self._start_runtime_subprocess('-m', 'langbot_plugin.cli.__init__', 'rt')
ws_url = 'ws://localhost:5400/control/ws'
async def connection_failed(ctrl, exc=None):
error = exc or RuntimeError('WebSocket connection failed')
connect_errors.append(error)
self._connected.set()
self.ctrl = ws_client_controller.WebSocketClientController(
ws_url=ws_url,
make_connection_failed_callback=connection_failed,
)
task_coro = self.ctrl.run(new_connection_callback)
else:
self.ctrl = stdio_client_controller.StdioClientController(
command=sys.executable,
args=['-m', 'langbot_plugin.cli.__init__', 'rt', '-s'],
env=os.environ.copy(),
)
task_coro = self.ctrl.run(new_connection_callback)
self._transport_task = asyncio.create_task(task_coro)
try:
await asyncio.wait_for(self._connected.wait(), timeout=_CONNECT_TIMEOUT_SEC)
except asyncio.TimeoutError as exc:
await self._stop_transport()
raise PluginRuntimeNotConnectedError('Plugin runtime did not become ready within 30 seconds') from exc
if connect_errors:
await self._stop_transport()
raise PluginRuntimeNotConnectedError(f'Plugin runtime connection failed: {connect_errors[-1]}')
if self.heartbeat_task is None or self.heartbeat_task.done():
self.heartbeat_task = asyncio.create_task(self.heartbeat_loop())
def schedule_reconnect(self) -> None:
if self._closing or not self.is_enable_plugin:
return
if self._reconnect_task is not None and not self._reconnect_task.done():
return
self._reconnect_task = asyncio.create_task(self._reconnect_loop())
async def _reconnect_loop(self) -> None:
delay = 1.0
try:
while not self._closing:
try:
await self.initialize()
return
except Exception as exc:
self.ap.logger.warning(f'Plugin runtime reconnection failed: {exc}; retrying in {delay:.0f}s')
await asyncio.sleep(delay)
delay = min(delay * 2, _RECONNECT_MAX_DELAY_SEC)
finally:
self._reconnect_task = None
async def _stop_transport(self) -> None:
runtime_handler = getattr(self, 'handler', None)
if runtime_handler is not None:
with contextlib.suppress(Exception):
await runtime_handler.close()
del self.handler
tasks = [
task
for task in (
getattr(self, 'handler_task', None),
self._transport_task,
)
if task is not None and task is not asyncio.current_task()
]
for task in tasks:
task.cancel()
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
self._transport_task = None
if hasattr(self, 'handler_task'):
del self.handler_task
close_ctrl = getattr(getattr(self, 'ctrl', None), 'close', None)
if close_ctrl is not None:
with contextlib.suppress(Exception):
await close_ctrl()
async def make_connection_failed_callback(
ctrl: ws_client_controller.WebSocketClientController,
exc: Exception = None,
) -> None:
if exc is not None:
self.ap.logger.error(f'Failed to connect to plugin runtime({ws_url}): {exc}')
else:
self.ap.logger.error(f'Failed to connect to plugin runtime({ws_url}), trying to reconnect...')
await self.runtime_disconnect_callback(self)
self.ctrl = ws_client_controller.WebSocketClientController(
ws_url=ws_url,
make_connection_failed_callback=make_connection_failed_callback,
)
task = self.ctrl.run(new_connection_callback)
elif platform.get_platform() == 'win32':
# Due to Windows's lack of supports for both stdio and subprocess:
# See also: https://docs.python.org/zh-cn/3.13/library/asyncio-platforms.html
# We have to launch runtime via cmd but communicate via ws.
self.ap.logger.info('(windows) use cmd to launch plugin runtime and communicate via ws')
await self._start_runtime_subprocess('-m', 'langbot_plugin.cli.__init__', 'rt')
ws_url = 'ws://localhost:5400/control/ws'
async def make_connection_failed_callback(
ctrl: ws_client_controller.WebSocketClientController,
exc: Exception = None,
) -> None:
if exc is not None:
self.ap.logger.error(f'(windows) Failed to connect to plugin runtime({ws_url}): {exc}')
else:
self.ap.logger.error(
f'(windows) Failed to connect to plugin runtime({ws_url}), trying to reconnect...'
)
await self.runtime_disconnect_callback(self)
self.ctrl = ws_client_controller.WebSocketClientController(
ws_url=ws_url,
make_connection_failed_callback=make_connection_failed_callback,
)
task = self.ctrl.run(new_connection_callback)
else: # stdio
self.ap.logger.info('use stdio to connect to plugin runtime')
# cmd: lbp rt -s
python_path = sys.executable
env = os.environ.copy()
self.ctrl = stdio_client_controller.StdioClientController(
command=python_path,
args=['-m', 'langbot_plugin.cli.__init__', 'rt', '-s'],
env=env,
)
task = self.ctrl.run(new_connection_callback)
if self.heartbeat_task is None:
self.heartbeat_task = asyncio.create_task(self.heartbeat_loop())
asyncio.create_task(task)
async def aclose(self) -> None:
self._closing = True
reconnect_task = self._reconnect_task
self._reconnect_task = None
if reconnect_task is not None and reconnect_task is not asyncio.current_task():
reconnect_task.cancel()
await asyncio.gather(reconnect_task, return_exceptions=True)
if self.heartbeat_task is not None:
self.heartbeat_task.cancel()
await asyncio.gather(self.heartbeat_task, return_exceptions=True)
self.heartbeat_task = None
await self._stop_transport()
await self._close_managed_subprocess()
async def initialize_plugins(self):
pass
async def ping_plugin_runtime(self):
if not hasattr(self, 'handler'):
raise PluginRuntimeNotConnectedError('Plugin runtime is not connected')
return await self.handler.ping()
return await self._runtime_handler().ping()
def _inspect_plugin_package(
self,
@@ -473,7 +570,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
file_bytes = download_resp.content
self._inspect_plugin_package(file_bytes, task_context)
file_key = await self.handler.send_file(file_bytes, 'lbpkg')
file_key = await self._runtime_handler().send_file(file_bytes, 'lbpkg')
install_info['plugin_file_key'] = file_key
self.ap.logger.info(f'Transfered file {file_key} to plugin runtime')
# Continue to install via runtime
@@ -496,7 +593,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
plugin_author, plugin_name = self._inspect_plugin_package(file_bytes, task_context)
if task_context is not None and plugin_author and plugin_name:
task_context.metadata['plugin_name'] = f'{plugin_author}/{plugin_name}'
file_key = await self.handler.send_file(file_bytes, 'lbpkg')
file_key = await self._runtime_handler().send_file(file_bytes, 'lbpkg')
install_info['plugin_file_key'] = file_key
del install_info['plugin_file']
self.ap.logger.info(f'Transfered file {file_key} to plugin runtime')
@@ -535,14 +632,14 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
plugin_author, plugin_name = self._inspect_plugin_package(file_bytes, task_context)
if task_context is not None and plugin_author and plugin_name:
task_context.metadata['plugin_name'] = f'{plugin_author}/{plugin_name}'
file_key = await self.handler.send_file(file_bytes, 'lbpkg')
file_key = await self._runtime_handler().send_file(file_bytes, 'lbpkg')
install_info['plugin_file_key'] = file_key
self.ap.logger.info(f'Transfered file {file_key} to plugin runtime')
except Exception as e:
self.ap.logger.error(f'Failed to download file from GitHub: {e}')
raise Exception(f'Failed to download file from GitHub: {e}')
async for ret in self.handler.install_plugin(install_source.value, install_info):
async for ret in self._runtime_handler().install_plugin(install_source.value, install_info):
current_action = ret.get('current_action', None)
if current_action is not None:
if task_context is not None:
@@ -566,7 +663,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
plugin_name: str,
task_context: taskmgr.TaskContext | None = None,
) -> dict[str, Any]:
async for ret in self.handler.upgrade_plugin(plugin_author, plugin_name):
async for ret in self._runtime_handler().upgrade_plugin(plugin_author, plugin_name):
current_action = ret.get('current_action', None)
if current_action is not None:
if task_context is not None:
@@ -584,7 +681,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
delete_data: bool = False,
task_context: taskmgr.TaskContext | None = None,
) -> dict[str, Any]:
async for ret in self.handler.delete_plugin(plugin_author, plugin_name):
async for ret in self._runtime_handler().delete_plugin(plugin_author, plugin_name):
current_action = ret.get('current_action', None)
if current_action is not None:
if task_context is not None:
@@ -599,7 +696,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if delete_data:
if task_context is not None:
task_context.trace('Cleaning up plugin configuration and storage...')
await self.handler.cleanup_plugin_data(plugin_author, plugin_name)
await self._runtime_handler().cleanup_plugin_data(plugin_author, plugin_name)
async def list_plugins(self, component_kinds: list[str] | None = None) -> list[dict[str, Any]]:
"""List plugins, optionally filtered by component kinds.
@@ -610,10 +707,10 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
component of the specified kinds will be returned.
E.g., ['Command', 'EventListener', 'Tool'] for pipeline-related plugins.
"""
if not self.is_enable_plugin:
if not self.is_enable_plugin or not self._runtime_available():
return []
plugins = await self.handler.list_plugins()
plugins = await self._runtime_handler().list_plugins()
# Filter plugins by component kinds if specified
if component_kinds is not None:
@@ -685,18 +782,18 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
return plugins
async def get_plugin_info(self, author: str, plugin_name: str) -> dict[str, Any]:
return await self.handler.get_plugin_info(author, plugin_name)
return await self._runtime_handler().get_plugin_info(author, plugin_name)
async def set_plugin_config(self, plugin_author: str, plugin_name: str, config: dict[str, Any]) -> dict[str, Any]:
return await self.handler.set_plugin_config(plugin_author, plugin_name, config)
return await self._runtime_handler().set_plugin_config(plugin_author, plugin_name, config)
@alru_cache(ttl=5 * 60) # 5 minutes
async def get_plugin_icon(self, plugin_author: str, plugin_name: str) -> dict[str, Any]:
return await self.handler.get_plugin_icon(plugin_author, plugin_name)
return await self._runtime_handler().get_plugin_icon(plugin_author, plugin_name)
@alru_cache(ttl=5 * 60) # 5 minutes
async def get_plugin_readme(self, plugin_author: str, plugin_name: str, language: str = 'en') -> str:
return await self.handler.get_plugin_readme(plugin_author, plugin_name, language)
return await self._runtime_handler().get_plugin_readme(plugin_author, plugin_name, language)
async def get_plugin_logs(
self,
@@ -706,11 +803,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
level: str | None = None,
) -> list[dict[str, Any]]:
# Not cached: logs are live and change constantly.
return await self.handler.get_plugin_logs(plugin_author, plugin_name, limit, level)
return await self._runtime_handler().get_plugin_logs(plugin_author, plugin_name, limit, level)
@alru_cache(ttl=5 * 60)
async def get_plugin_assets(self, plugin_author: str, plugin_name: str, filepath: str) -> dict[str, Any]:
return await self.handler.get_plugin_assets(plugin_author, plugin_name, filepath)
return await self._runtime_handler().get_plugin_assets(plugin_author, plugin_name, filepath)
async def handle_page_api(
self,
@@ -721,13 +818,15 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
method: str,
body: Any = None,
) -> dict[str, Any]:
return await self.handler.handle_page_api(plugin_author, plugin_name, page_id, endpoint, method, body)
return await self._runtime_handler().handle_page_api(
plugin_author, plugin_name, page_id, endpoint, method, body
)
async def get_debug_info(self) -> dict[str, Any]:
"""Get debug information including debug key and WS URL"""
if not self.is_enable_plugin:
if not self.is_enable_plugin or not self._runtime_available():
return {}
return await self.handler.get_debug_info()
return await self._runtime_handler().get_debug_info()
async def emit_event(
self,
@@ -736,13 +835,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
) -> context.EventContext:
event_ctx = context.EventContext.from_event(event)
if not self.is_enable_plugin:
if not self.is_enable_plugin or not self._runtime_available():
event_ctx._emitted_plugins = []
event_ctx._response_sources = []
return event_ctx
# Pass include_plugins to runtime for filtering
event_ctx_result = await self.handler.emit_event(
event_ctx_result = await self._runtime_handler().emit_event(
event_ctx.model_dump(serialize_as_any=False), include_plugins=bound_plugins
)
@@ -755,19 +854,19 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
async def notify_plugin_diagnostic(self, diagnostic: dict[str, Any]) -> None:
"""Best-effort diagnostic forwarding to the plugin runtime."""
if not self.is_enable_plugin:
if not self.is_enable_plugin or not self._runtime_available():
return
try:
await self.handler.notify_plugin_diagnostic(diagnostic)
await self._runtime_handler().notify_plugin_diagnostic(diagnostic)
except Exception as e:
self.ap.logger.debug(f'Plugin diagnostic forwarding skipped: {e}')
async def list_tools(self, bound_plugins: list[str] | None = None) -> list[ComponentManifest]:
if not self.is_enable_plugin:
if not self.is_enable_plugin or not self._runtime_available():
return []
# Pass include_plugins to runtime for filtering
list_tools_data = await self.handler.list_tools(include_plugins=bound_plugins)
list_tools_data = await self._runtime_handler().list_tools(include_plugins=bound_plugins)
tools = [ComponentManifest.model_validate(tool) for tool in list_tools_data]
@@ -785,16 +884,19 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
return {'error': 'Tool not found: plugin system is disabled'}
# Pass include_plugins to runtime for validation
return await self.handler.call_tool(
if not self._runtime_available():
return {'error': 'Plugin runtime is temporarily unavailable'}
return await self._runtime_handler().call_tool(
tool_name, parameters, session.model_dump(serialize_as_any=True), query_id, include_plugins=bound_plugins
)
async def list_commands(self, bound_plugins: list[str] | None = None) -> list[ComponentManifest]:
if not self.is_enable_plugin:
if not self.is_enable_plugin or not self._runtime_available():
return []
# Pass include_plugins to runtime for filtering
list_commands_data = await self.handler.list_commands(include_plugins=bound_plugins)
list_commands_data = await self._runtime_handler().list_commands(include_plugins=bound_plugins)
commands = [ComponentManifest.model_validate(command) for command in list_commands_data]
@@ -803,12 +905,15 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
async def execute_command(
self, command_ctx: command_context.ExecuteContext, bound_plugins: list[str] | None = None
) -> typing.AsyncGenerator[command_context.CommandReturn, None]:
if not self.is_enable_plugin:
if not self.is_enable_plugin or not self._runtime_available():
yield command_context.CommandReturn(error=command_errors.CommandNotFoundError(command_ctx.command))
return
# Pass include_plugins to runtime for validation
gen = self.handler.execute_command(command_ctx.model_dump(serialize_as_any=True), include_plugins=bound_plugins)
gen = self._runtime_handler().execute_command(
command_ctx.model_dump(serialize_as_any=True),
include_plugins=bound_plugins,
)
async for ret in gen:
cmd_ret = command_context.CommandReturn.model_validate(ret)
@@ -823,27 +928,33 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
retrieval_context: dict[str, Any],
) -> dict[str, Any]:
"""Retrieve knowledge using a KnowledgeEngine instance."""
if not self.is_enable_plugin:
if not self.is_enable_plugin or not self._runtime_available():
return {'results': []}
return await self.handler.retrieve_knowledge(plugin_author, plugin_name, retriever_name, retrieval_context)
return await self._runtime_handler().retrieve_knowledge(
plugin_author, plugin_name, retriever_name, retrieval_context
)
def dispose(self):
# On non-Windows stdio mode, terminate via the controller's process handle.
# On Windows, the managed subprocess is cleaned up by the base class.
if (
self.is_enable_plugin
and hasattr(self, 'ctrl')
and isinstance(self.ctrl, stdio_client_controller.StdioClientController)
):
self.ap.logger.info('Terminating plugin runtime process...')
self.ctrl.process.terminate()
self._dispose_subprocess()
"""Best-effort synchronous compatibility wrapper; prefer ``aclose``."""
self._closing = True
if self.heartbeat_task is not None:
self.heartbeat_task.cancel()
self.heartbeat_task = None
if self._reconnect_task is not None:
self._reconnect_task.cancel()
self._reconnect_task = None
for task in (
getattr(self, 'handler_task', None),
self._transport_task,
):
if task is not None:
task.cancel()
ctrl = getattr(self, 'ctrl', None)
process = getattr(ctrl, 'process', None)
if process is not None and process.returncode is None:
process.terminate()
self._dispose_subprocess()
@staticmethod
def _parse_plugin_id(plugin_id: str) -> tuple[str, str]:
@@ -873,29 +984,29 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
context_data: IngestionContext data.
"""
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
return await self.handler.rag_ingest_document(plugin_author, plugin_name, context_data)
return await self._runtime_handler().rag_ingest_document(plugin_author, plugin_name, context_data)
async def call_rag_delete_document(self, plugin_id: str, document_id: str, kb_id: str) -> bool:
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
return await self.handler.rag_delete_document(plugin_author, plugin_name, document_id, kb_id)
return await self._runtime_handler().rag_delete_document(plugin_author, plugin_name, document_id, kb_id)
async def get_rag_creation_schema(self, plugin_id: str) -> dict[str, Any]:
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
return await self.handler.get_rag_creation_schema(plugin_author, plugin_name)
return await self._runtime_handler().get_rag_creation_schema(plugin_author, plugin_name)
async def get_rag_retrieval_schema(self, plugin_id: str) -> dict[str, Any]:
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
return await self.handler.get_rag_retrieval_schema(plugin_author, plugin_name)
return await self._runtime_handler().get_rag_retrieval_schema(plugin_author, plugin_name)
async def rag_on_kb_create(self, plugin_id: str, kb_id: str, config: dict[str, Any]) -> dict[str, Any]:
"""Notify plugin about KB creation."""
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
return await self.handler.rag_on_kb_create(plugin_author, plugin_name, kb_id, config)
return await self._runtime_handler().rag_on_kb_create(plugin_author, plugin_name, kb_id, config)
async def rag_on_kb_delete(self, plugin_id: str, kb_id: str) -> dict[str, Any]:
"""Notify plugin about KB deletion."""
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
return await self.handler.rag_on_kb_delete(plugin_author, plugin_name, kb_id)
return await self._runtime_handler().rag_on_kb_delete(plugin_author, plugin_name, kb_id)
async def call_rag_retrieve(self, plugin_id: str, retrieval_context: dict[str, Any]) -> dict[str, Any]:
"""Call plugin to retrieve knowledge.
@@ -905,25 +1016,25 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
retrieval_context: RetrievalContext data.
"""
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
return await self.handler.retrieve_knowledge(plugin_author, plugin_name, '', retrieval_context)
return await self._runtime_handler().retrieve_knowledge(plugin_author, plugin_name, '', retrieval_context)
async def list_knowledge_engines(self) -> list[dict[str, Any]]:
"""List all available Knowledge Engines from plugins.
Returns a list of Knowledge Engines with their capabilities and configuration schemas.
"""
if not self.is_enable_plugin:
if not self.is_enable_plugin or not self._runtime_available():
return []
return await self.handler.list_knowledge_engines()
return await self._runtime_handler().list_knowledge_engines()
async def list_parsers(self) -> list[dict[str, Any]]:
"""List all available parsers from plugins."""
if not self.is_enable_plugin:
if not self.is_enable_plugin or not self._runtime_available():
return []
return await self.handler.list_parsers()
return await self._runtime_handler().list_parsers()
async def call_parser(self, plugin_id: str, context_data: dict[str, Any], file_bytes: bytes) -> dict[str, Any]:
"""Call plugin to parse a document."""
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
return await self.handler.parse_document(plugin_author, plugin_name, context_data, file_bytes)
return await self._runtime_handler().parse_document(plugin_author, plugin_name, context_data, file_bytes)
src/langbot/pkg/provider/tools/loaders/mcp.py
+36 -10
View File
@@ -11,6 +11,7 @@ import traceback
from langbot_plugin.api.entities.events import pipeline_query
import sqlalchemy
import asyncio
import hashlib
import httpx
import uuid as uuid_module
@@ -26,7 +27,13 @@ from ....core import app
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
import langbot_plugin.api.entities.builtin.provider.message as provider_message
from ....entity.persistence import mcp as persistence_mcp
from .mcp_stdio import BoxStdioSessionRuntime, MCPServerBoxConfig, MCPSessionErrorPhase, _ColdStartRetry # noqa: F401
from .mcp_stdio import (
BoxStdioSessionRuntime,
MCPServerBoxConfig as MCPServerBoxConfig, # noqa: F401 - public re-export
MCPSessionErrorPhase,
_ColdStartRetry,
_get_default_memory_mb,
) # noqa: F401
# Synthesized LLM tools for MCP resources (not from server tools/list).
# Dispatched in MCPLoader.invoke_tool; placeholder func on LLMTool is never used.
@@ -1206,15 +1213,34 @@ class RuntimeMCPSession:
return self._box_stdio_runtime.uses_box_stdio()
def _build_box_session_id(self) -> str:
# Both live servers and transient config-page tests share ONE Box
# session ('mcp-shared'). A test therefore reuses the already-running
# container (and, for an existing server, its live managed process)
# instead of paying a full per-test session cold-start + dependency
# bootstrap. Isolation between a test and the live servers is provided
# at the *process* level: each server/test has its own process_id and a
# test only ever stops its own process_id (see cleanup_session), so it
# never disturbs another server's process or the shared session itself.
return 'mcp-shared'
# Compatible MCP servers share a session and remain isolated by
# process_id. A server with a different immutable resource profile gets
# another session; Docker/E2B cannot change memory/image/etc. after a
# session has been created.
config = self._box_stdio_runtime.config
default_memory = _get_default_memory_mb(self.ap)
profile = {
'image': config.image,
'network': config.network,
'host_path_mode': config.host_path_mode,
'cpus': config.cpus,
'memory_mb': config.memory_mb or default_memory,
'pids_limit': config.pids_limit,
'read_only_rootfs': (config.read_only_rootfs if config.read_only_rootfs is not None else False),
}
default_profile = {
'image': None,
'network': 'on',
'host_path_mode': 'ro',
'cpus': None,
'memory_mb': default_memory,
'pids_limit': None,
'read_only_rootfs': False,
}
if profile == default_profile:
return 'mcp-shared'
digest = hashlib.sha256(json.dumps(profile, sort_keys=True).encode('utf-8')).hexdigest()[:12]
return f'mcp-shared-{digest}'
def _rewrite_path(self, path: str, host_path: str | None) -> str:
return self._box_stdio_runtime.rewrite_path(path, host_path)
src/langbot/pkg/utils/managed_runtime.py
+23
View File
@@ -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
src/langbot/templates/config.yaml
+5
View File
@@ -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.
tests/unit_tests/box/test_box_connector.py
+19 -1
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import Mock
from unittest.mock import AsyncMock, Mock
import pytest
@@ -104,3 +104,21 @@ def test_box_runtime_connector_dispose_terminates_subprocess(monkeypatch: pytest
ctrl_task.cancel.assert_called_once()
assert connector._handler_task is None
assert connector._ctrl_task is None
@pytest.mark.asyncio
async def test_box_runtime_connector_cleans_partial_transport_on_connect_failure(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
connector = BoxRuntimeConnector(make_app(Mock()))
connector._start_local_stdio = AsyncMock(side_effect=RuntimeError('bind failed'))
connector._stop_transport = AsyncMock()
connector._close_managed_subprocess = AsyncMock()
with pytest.raises(RuntimeError, match='bind failed'):
await connector.initialize()
assert connector._stop_transport.await_count == 2
connector._close_managed_subprocess.assert_awaited_once()
tests/unit_tests/box/test_box_service.py
+19 -1
View File
@@ -325,10 +325,28 @@ async def test_box_service_dispose_schedules_shutdown_on_event_loop(monkeypatch:
service.dispose()
await asyncio.sleep(0)
connector.dispose.assert_called_once()
connector.dispose.assert_not_called()
service.shutdown.assert_awaited_once()
@pytest.mark.asyncio
async def test_box_service_shutdown_reaps_connector_when_runtime_rpc_is_offline(
monkeypatch: pytest.MonkeyPatch,
):
connector = Mock()
connector.client = Mock()
connector.client.shutdown = AsyncMock(side_effect=RuntimeError('offline'))
connector.aclose = AsyncMock()
monkeypatch.setattr('langbot.pkg.box.service.BoxRuntimeConnector', Mock(return_value=connector))
service = BoxService(make_app(Mock()))
await service.shutdown()
connector.client.shutdown.assert_awaited_once()
connector.aclose.assert_awaited_once()
@pytest.mark.asyncio
async def test_box_runtime_reuses_request_session():
logger = Mock()
tests/unit_tests/core/test_boot.py
+34 -20
View File
@@ -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
tests/unit_tests/plugin/test_connector_methods.py
+31
View File
@@ -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."""
tests/unit_tests/provider/test_mcp_box_integration.py
+46 -1
View File
@@ -417,7 +417,7 @@ class TestBuildBoxSessionPayload:
payload = s._build_box_session_payload('session-123')
assert payload['image'] == 'node:20'
assert payload['cpus'] == 2.0
assert payload["memory_mb"] == 1024
assert payload['memory_mb'] == 1024
assert payload['pids_limit'] == 256
def test_none_fields_excluded(self, mcp_module):
@@ -680,6 +680,51 @@ class TestGetRuntimeInfoDict:
# ... but are isolated by distinct process_ids within that session.
assert transient._box_stdio_runtime.process_id != live._box_stdio_runtime.process_id
def test_different_resource_profiles_use_different_box_sessions(self, mcp_module):
ap = _make_ap()
ap.box_service.available = True
default = _make_session(
mcp_module,
{
'name': 'default',
'uuid': 'default-uuid',
'mode': 'stdio',
'command': 'uvx',
'args': ['mcp-server-time'],
},
ap=ap,
)
constrained = _make_session(
mcp_module,
{
'name': 'constrained',
'uuid': 'constrained-uuid',
'mode': 'stdio',
'command': 'uvx',
'args': ['mcp-server-time'],
'box': {'memory_mb': 2048},
},
ap=ap,
)
assert default._build_box_session_id() == 'mcp-shared'
assert constrained._build_box_session_id().startswith('mcp-shared-')
assert constrained._build_box_session_id() != default._build_box_session_id()
writable = _make_session(
mcp_module,
{
'name': 'writable',
'uuid': 'writable-uuid',
'mode': 'stdio',
'command': 'uvx',
'args': ['mcp-server-time'],
'box': {'host_path_mode': 'rw'},
},
ap=ap,
)
assert writable._build_box_session_id() != default._build_box_session_id()
def test_stdio_session_refuses_when_box_unavailable(self, mcp_module):
"""Policy: when Box is configured but unavailable (disabled in config
OR connection failed), stdio MCP servers are NOT treated as box-stdio.