mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-02 03:55:55 +00:00
chore(sandbox): move MCP loader changes to follow-up branch
This commit is contained in:
youhuanghe
@@ -1,7 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import enum
|
||||
import os
|
||||
import typing
|
||||
from contextlib import AsyncExitStack
|
||||
import traceback
|
||||
@@ -10,13 +9,11 @@ import sqlalchemy
|
||||
import asyncio
|
||||
import httpx
|
||||
|
||||
import pydantic
|
||||
import uuid as uuid_module
|
||||
from mcp import ClientSession, StdioServerParameters
|
||||
from mcp.client.stdio import stdio_client
|
||||
from mcp.client.sse import sse_client
|
||||
from mcp.client.streamable_http import streamable_http_client
|
||||
from mcp.client.websocket import websocket_client
|
||||
|
||||
from .. import loader
|
||||
from ....core import app
|
||||
@@ -31,39 +28,6 @@ class MCPSessionStatus(enum.Enum):
|
||||
ERROR = 'error'
|
||||
|
||||
|
||||
class MCPSessionErrorPhase(enum.Enum):
|
||||
"""Which phase of the MCP lifecycle failed."""
|
||||
|
||||
SESSION_CREATE = 'session_create'
|
||||
DEP_INSTALL = 'dep_install'
|
||||
PROCESS_START = 'process_start'
|
||||
RELAY_CONNECT = 'relay_connect'
|
||||
MCP_INIT = 'mcp_init'
|
||||
RUNTIME = 'runtime'
|
||||
TOOL_CALL = 'tool_call'
|
||||
|
||||
|
||||
_VENV_DIRS = frozenset({'.venv', 'venv', 'env', '.env'})
|
||||
_VENV_BIN_DIRS = frozenset({'bin', 'Scripts'})
|
||||
|
||||
|
||||
class MCPServerBoxConfig(pydantic.BaseModel):
|
||||
"""Structured configuration for running an MCP server inside a Box container."""
|
||||
|
||||
image: str | None = None
|
||||
network: str = 'on' # MCP servers need network for dependency installation
|
||||
host_path: str | None = None
|
||||
host_path_mode: str = 'ro' # MCP servers default to read-only mount
|
||||
env: dict[str, str] = pydantic.Field(default_factory=dict)
|
||||
startup_timeout_sec: int = 120 # Longer default to allow pip install
|
||||
cpus: float | None = None
|
||||
memory_mb: int | None = None
|
||||
pids_limit: int | None = None
|
||||
read_only_rootfs: bool | None = None
|
||||
|
||||
model_config = pydantic.ConfigDict(extra='ignore')
|
||||
|
||||
|
||||
class RuntimeMCPSession:
|
||||
"""运行时 MCP 会话"""
|
||||
|
||||
@@ -94,10 +58,6 @@ class RuntimeMCPSession:
|
||||
|
||||
error_message: str | None = None
|
||||
|
||||
error_phase: MCPSessionErrorPhase | None = None
|
||||
|
||||
retry_count: int = 0
|
||||
|
||||
def __init__(self, server_name: str, server_config: dict, enable: bool, ap: app.Application):
|
||||
self.server_name = server_name
|
||||
self.server_uuid = server_config.get('uuid', '')
|
||||
@@ -115,14 +75,7 @@ class RuntimeMCPSession:
|
||||
self._shutdown_event = asyncio.Event()
|
||||
self._ready_event = asyncio.Event()
|
||||
|
||||
# Parse box config once
|
||||
self.box_config = MCPServerBoxConfig.model_validate(server_config.get('box', {}))
|
||||
|
||||
async def _init_stdio_python_server(self):
|
||||
if self._uses_box_stdio():
|
||||
await self._init_box_stdio_server()
|
||||
return
|
||||
|
||||
server_params = StdioServerParameters(
|
||||
command=self.server_config['command'],
|
||||
args=self.server_config['args'],
|
||||
@@ -137,68 +90,6 @@ class RuntimeMCPSession:
|
||||
|
||||
await self.session.initialize()
|
||||
|
||||
async def _init_box_stdio_server(self):
|
||||
box_service = self.ap.box_service
|
||||
session_id = self._build_box_session_id()
|
||||
host_path = self._resolve_host_path()
|
||||
session_payload = self._build_box_session_payload(session_id, host_path)
|
||||
|
||||
# Phase: session creation
|
||||
try:
|
||||
await box_service.create_session(
|
||||
session_payload,
|
||||
)
|
||||
except Exception:
|
||||
self.error_phase = MCPSessionErrorPhase.SESSION_CREATE
|
||||
raise
|
||||
|
||||
# Phase: dependency installation
|
||||
if host_path:
|
||||
install_cmd = self._detect_install_command(host_path)
|
||||
if install_cmd:
|
||||
self.ap.logger.info(
|
||||
f'MCP server {self.server_name}: installing dependencies in Box with: {install_cmd}'
|
||||
)
|
||||
exec_payload = dict(session_payload)
|
||||
exec_payload['cmd'] = install_cmd
|
||||
exec_payload['timeout_sec'] = self.box_config.startup_timeout_sec or 120
|
||||
try:
|
||||
result = await box_service.client.execute(box_service.build_spec(exec_payload))
|
||||
except Exception:
|
||||
self.error_phase = MCPSessionErrorPhase.DEP_INSTALL
|
||||
raise
|
||||
if not result.ok:
|
||||
self.error_phase = MCPSessionErrorPhase.DEP_INSTALL
|
||||
stderr_preview = (result.stderr or '')[:500]
|
||||
raise Exception(f'Dependency install failed (exit code {result.exit_code}): {stderr_preview}')
|
||||
|
||||
# Phase: managed process start
|
||||
try:
|
||||
await box_service.start_managed_process(
|
||||
session_id,
|
||||
self._build_box_process_payload(host_path),
|
||||
)
|
||||
except Exception:
|
||||
self.error_phase = MCPSessionErrorPhase.PROCESS_START
|
||||
raise
|
||||
|
||||
# Phase: WebSocket relay connection
|
||||
try:
|
||||
websocket_url = box_service.get_managed_process_websocket_url(session_id)
|
||||
transport = await self.exit_stack.enter_async_context(websocket_client(websocket_url))
|
||||
read_stream, write_stream = transport
|
||||
self.session = await self.exit_stack.enter_async_context(ClientSession(read_stream, write_stream))
|
||||
except Exception:
|
||||
self.error_phase = MCPSessionErrorPhase.RELAY_CONNECT
|
||||
raise
|
||||
|
||||
# Phase: MCP protocol initialization
|
||||
try:
|
||||
await self.session.initialize()
|
||||
except Exception:
|
||||
self.error_phase = MCPSessionErrorPhase.MCP_INIT
|
||||
raise
|
||||
|
||||
async def _init_sse_server(self):
|
||||
sse_transport = await self.exit_stack.enter_async_context(
|
||||
sse_client(
|
||||
@@ -233,11 +124,8 @@ class RuntimeMCPSession:
|
||||
|
||||
await self.session.initialize()
|
||||
|
||||
_MAX_RETRIES = 3
|
||||
_RETRY_DELAYS = [2, 4, 8]
|
||||
|
||||
async def _lifecycle_loop(self):
|
||||
"""Manage the full MCP session lifecycle in a background task."""
|
||||
"""在后台任务中管理整个MCP会话的生命周期"""
|
||||
try:
|
||||
if self.server_config['mode'] == 'stdio':
|
||||
await self._init_stdio_python_server()
|
||||
@@ -246,125 +134,49 @@ class RuntimeMCPSession:
|
||||
elif self.server_config['mode'] == 'http':
|
||||
await self._init_streamable_http_server()
|
||||
else:
|
||||
raise ValueError(f'Unknown MCP server mode: {self.server_name}: {self.server_config}')
|
||||
raise ValueError(f'无法识别 MCP 服务器类型: {self.server_name}: {self.server_config}')
|
||||
|
||||
await self.refresh()
|
||||
|
||||
self.status = MCPSessionStatus.CONNECTED
|
||||
|
||||
# Notify start() that connection is established
|
||||
# 通知start()方法连接已建立
|
||||
self._ready_event.set()
|
||||
|
||||
# Wait for shutdown signal, with optional health monitoring for Box stdio
|
||||
if self._uses_box_stdio():
|
||||
monitor_task = asyncio.create_task(self._monitor_box_process_health())
|
||||
shutdown_task = asyncio.create_task(self._shutdown_event.wait())
|
||||
done, pending = await asyncio.wait(
|
||||
[shutdown_task, monitor_task],
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
for task in done:
|
||||
if task is monitor_task and not self._shutdown_event.is_set():
|
||||
self.error_phase = MCPSessionErrorPhase.RUNTIME
|
||||
raise Exception('Box managed process exited unexpectedly')
|
||||
else:
|
||||
await self._shutdown_event.wait()
|
||||
# 等待shutdown信号
|
||||
await self._shutdown_event.wait()
|
||||
|
||||
except Exception as e:
|
||||
self.status = MCPSessionStatus.ERROR
|
||||
self.error_message = str(e)
|
||||
self.ap.logger.error(f'Error in MCP session lifecycle {self.server_name}: {e}\n{traceback.format_exc()}')
|
||||
# Do NOT set _ready_event here — let _lifecycle_loop_with_retry
|
||||
# handle retries first. It will set the event when all retries
|
||||
# are exhausted or on success.
|
||||
raise # Re-raise so _lifecycle_loop_with_retry can catch it
|
||||
# 即使出错也要设置ready事件，让start()方法知道初始化已完成
|
||||
self._ready_event.set()
|
||||
finally:
|
||||
# Clean up all resources in the same task
|
||||
# 在同一个任务中清理所有资源
|
||||
try:
|
||||
if self.exit_stack:
|
||||
await self.exit_stack.aclose()
|
||||
self.exit_stack = AsyncExitStack()
|
||||
self.functions.clear()
|
||||
self.session = None
|
||||
except Exception as e:
|
||||
self.ap.logger.error(f'Error cleaning up MCP session {self.server_name}: {e}\n{traceback.format_exc()}')
|
||||
finally:
|
||||
await self._cleanup_box_stdio_session()
|
||||
|
||||
async def _lifecycle_loop_with_retry(self):
|
||||
"""Wrap _lifecycle_loop with retry and exponential backoff."""
|
||||
for attempt in range(self._MAX_RETRIES + 1):
|
||||
try:
|
||||
await self._lifecycle_loop()
|
||||
return # Normal shutdown, don't retry
|
||||
except Exception as e:
|
||||
self.retry_count = attempt + 1
|
||||
if self._shutdown_event.is_set():
|
||||
return # Shutdown requested, don't retry
|
||||
if attempt >= self._MAX_RETRIES:
|
||||
self.status = MCPSessionStatus.ERROR
|
||||
self.error_message = f'Failed after {self._MAX_RETRIES + 1} attempts: {e}'
|
||||
self._ready_event.set()
|
||||
return
|
||||
delay = self._RETRY_DELAYS[attempt]
|
||||
self.ap.logger.warning(
|
||||
f'MCP session {self.server_name} failed (attempt {attempt + 1}), retrying in {delay}s: {e}'
|
||||
)
|
||||
await self._cleanup_box_stdio_session()
|
||||
# Reset status for retry
|
||||
self.status = MCPSessionStatus.CONNECTING
|
||||
self.error_message = None
|
||||
self.error_phase = None
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
_MONITOR_POLL_INTERVAL = 5
|
||||
_MONITOR_MAX_CONSECUTIVE_ERRORS = 3
|
||||
|
||||
async def _monitor_box_process_health(self):
|
||||
"""Poll managed process status; return when process exits."""
|
||||
from langbot_plugin.box.models import BoxManagedProcessStatus
|
||||
|
||||
session_id = self._build_box_session_id()
|
||||
consecutive_errors = 0
|
||||
while not self._shutdown_event.is_set():
|
||||
try:
|
||||
info = await self.ap.box_service.client.get_managed_process(session_id)
|
||||
if isinstance(info, dict):
|
||||
status = info.get('status', '')
|
||||
else:
|
||||
status = getattr(info, 'status', '')
|
||||
if status == BoxManagedProcessStatus.EXITED.value or status == BoxManagedProcessStatus.EXITED:
|
||||
return
|
||||
consecutive_errors = 0
|
||||
except Exception as exc:
|
||||
consecutive_errors += 1
|
||||
self.ap.logger.warning(
|
||||
f'MCP monitor for {self.server_name}: get_managed_process failed '
|
||||
f'({consecutive_errors}/{self._MONITOR_MAX_CONSECUTIVE_ERRORS}): '
|
||||
f'{type(exc).__name__}: {exc}'
|
||||
)
|
||||
if consecutive_errors >= self._MONITOR_MAX_CONSECUTIVE_ERRORS:
|
||||
return
|
||||
await asyncio.sleep(self._MONITOR_POLL_INTERVAL)
|
||||
|
||||
async def start(self):
|
||||
if not self.enable:
|
||||
return
|
||||
|
||||
# Create background task for lifecycle management with retry
|
||||
self._lifecycle_task = asyncio.create_task(self._lifecycle_loop_with_retry())
|
||||
# 创建后台任务来管理生命周期
|
||||
self._lifecycle_task = asyncio.create_task(self._lifecycle_loop())
|
||||
|
||||
# Wait for connection or failure (with timeout)
|
||||
startup_timeout = self.box_config.startup_timeout_sec if self._uses_box_stdio() else 30.0
|
||||
# 等待连接建立或失败（带超时）
|
||||
try:
|
||||
await asyncio.wait_for(self._ready_event.wait(), timeout=startup_timeout)
|
||||
await asyncio.wait_for(self._ready_event.wait(), timeout=30.0)
|
||||
except asyncio.TimeoutError:
|
||||
self.status = MCPSessionStatus.ERROR
|
||||
raise Exception(f'Connection timeout after {startup_timeout} seconds')
|
||||
raise Exception('Connection timeout after 30 seconds')
|
||||
|
||||
# Check for errors
|
||||
# 检查是否有错误
|
||||
if self.status == MCPSessionStatus.ERROR:
|
||||
raise Exception('Connection failed, please check URL')
|
||||
|
||||
@@ -420,11 +232,9 @@ class RuntimeMCPSession:
|
||||
return self.functions
|
||||
|
||||
def get_runtime_info_dict(self) -> dict:
|
||||
info = {
|
||||
return {
|
||||
'status': self.status.value,
|
||||
'error_message': self.error_message,
|
||||
'error_phase': self.error_phase.value if self.error_phase else None,
|
||||
'retry_count': self.retry_count,
|
||||
'tool_count': len(self.get_tools()),
|
||||
'tools': [
|
||||
{
|
||||
@@ -434,10 +244,6 @@ class RuntimeMCPSession:
|
||||
for tool in self.get_tools()
|
||||
],
|
||||
}
|
||||
if self._uses_box_stdio():
|
||||
info['box_session_id'] = self._build_box_session_id()
|
||||
info['box_enabled'] = True
|
||||
return info
|
||||
|
||||
async def shutdown(self):
|
||||
"""关闭会话并清理资源"""
|
||||
@@ -461,182 +267,6 @@ class RuntimeMCPSession:
|
||||
except Exception as e:
|
||||
self.ap.logger.error(f'Error shutting down MCP session {self.server_name}: {e}\n{traceback.format_exc()}')
|
||||
|
||||
def _uses_box_stdio(self) -> bool:
|
||||
"""Check whether this stdio MCP server should run inside a Box container.
|
||||
|
||||
Returns True when mode is stdio AND the Box runtime is available.
|
||||
An explicit ``box`` key in server_config is NOT required — if the
|
||||
runtime is reachable, stdio servers default to Box isolation.
|
||||
"""
|
||||
if self.server_config.get('mode') != 'stdio':
|
||||
return False
|
||||
try:
|
||||
return getattr(self.ap.box_service, 'available', False)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def _build_box_session_id(self) -> str:
|
||||
return f'mcp-{self.server_uuid}'
|
||||
|
||||
def _rewrite_path(self, path: str, host_path: str | None) -> str:
|
||||
"""Rewrite host path prefix to container /workspace prefix."""
|
||||
if not host_path or not path:
|
||||
return path
|
||||
normalized_host = os.path.realpath(host_path)
|
||||
if path.startswith(normalized_host + '/'):
|
||||
return '/workspace' + path[len(normalized_host) :]
|
||||
if path == normalized_host:
|
||||
return '/workspace'
|
||||
return path
|
||||
|
||||
def _infer_host_path(self) -> str | None:
|
||||
"""Try to infer host_path from command and args absolute paths.
|
||||
|
||||
Detects virtualenv patterns (e.g. .venv/bin/python) and walks up
|
||||
to the project root rather than using the bin directory.
|
||||
"""
|
||||
candidates = []
|
||||
parts = [self.server_config.get('command', '')] + self.server_config.get('args', [])
|
||||
for part in parts:
|
||||
if not os.path.isabs(part):
|
||||
continue
|
||||
# Use the raw path for venv detection (before resolving symlinks)
|
||||
# because .venv/bin/python is often a symlink to the system python.
|
||||
if os.path.exists(part):
|
||||
directory = os.path.dirname(part)
|
||||
directory = self._unwrap_venv_path(directory)
|
||||
candidates.append(os.path.realpath(directory))
|
||||
if not candidates:
|
||||
return None
|
||||
common = os.path.commonpath(candidates)
|
||||
return common if common != '/' else None
|
||||
|
||||
@staticmethod
|
||||
def _unwrap_venv_path(directory: str) -> str:
|
||||
"""If directory looks like a virtualenv bin dir, return the project root.
|
||||
|
||||
Recognized patterns:
|
||||
/project/.venv/bin -> /project
|
||||
/project/venv/bin -> /project
|
||||
/project/.venv/Scripts -> /project (Windows)
|
||||
/project/env/bin -> /project
|
||||
"""
|
||||
parts = directory.replace('\\', '/').split('/')
|
||||
# Look for patterns like .../(.venv|venv|env)/(bin|Scripts)
|
||||
for i in range(len(parts) - 1, 0, -1):
|
||||
if parts[i] in _VENV_BIN_DIRS and i >= 1:
|
||||
venv_dir = parts[i - 1]
|
||||
if venv_dir in _VENV_DIRS:
|
||||
# Return everything before the venv directory
|
||||
project_root = '/'.join(parts[: i - 1])
|
||||
return project_root if project_root else '/'
|
||||
return directory
|
||||
|
||||
def _resolve_host_path(self) -> str | None:
|
||||
"""Resolve the effective host_path: explicit config > inference."""
|
||||
return self.box_config.host_path or self._infer_host_path()
|
||||
|
||||
@staticmethod
|
||||
def _detect_install_command(host_path: str) -> str | None:
|
||||
"""Detect how to install dependencies from the mounted project.
|
||||
|
||||
Copies the project to a writable temp directory before installing,
|
||||
because /workspace may be mounted read-only and pip needs to write
|
||||
build artifacts in the source tree.
|
||||
"""
|
||||
# Use /opt instead of /tmp — /tmp is often a small tmpfs (64 MB)
|
||||
# and cannot hold the copied source tree plus pip build artifacts.
|
||||
_COPY_AND_INSTALL = (
|
||||
'mkdir -p /opt/_mcp_src'
|
||||
' && tar -C /workspace'
|
||||
' --exclude=.venv --exclude=.git --exclude=__pycache__'
|
||||
' --exclude=node_modules --exclude=.tox --exclude=.nox'
|
||||
' --exclude="*.egg-info" --exclude=.uv-cache'
|
||||
' -cf - .'
|
||||
' | tar -C /opt/_mcp_src -xf -'
|
||||
' && pip install --no-cache-dir /opt/_mcp_src'
|
||||
' && rm -rf /opt/_mcp_src'
|
||||
)
|
||||
_INSTALL_REQUIREMENTS = 'pip install --no-cache-dir -r /workspace/requirements.txt'
|
||||
|
||||
if os.path.isfile(os.path.join(host_path, 'pyproject.toml')):
|
||||
return _COPY_AND_INSTALL
|
||||
if os.path.isfile(os.path.join(host_path, 'setup.py')):
|
||||
return _COPY_AND_INSTALL
|
||||
if os.path.isfile(os.path.join(host_path, 'requirements.txt')):
|
||||
return _INSTALL_REQUIREMENTS
|
||||
return None
|
||||
|
||||
def _build_box_session_payload(self, session_id: str, host_path: str | None = None) -> dict:
|
||||
bc = self.box_config
|
||||
if host_path is None:
|
||||
host_path = self._resolve_host_path()
|
||||
|
||||
payload: dict[str, typing.Any] = {
|
||||
'session_id': session_id,
|
||||
'workdir': '/workspace',
|
||||
'env': bc.env,
|
||||
# MCP sessions need network for dependency install and writable rootfs
|
||||
'network': bc.network,
|
||||
'read_only_rootfs': bc.read_only_rootfs if bc.read_only_rootfs is not None else False,
|
||||
}
|
||||
if host_path:
|
||||
payload['host_path'] = host_path
|
||||
payload['host_path_mode'] = bc.host_path_mode
|
||||
for key in ('image', 'cpus', 'memory_mb', 'pids_limit'):
|
||||
val = getattr(bc, key)
|
||||
if val is not None:
|
||||
payload[key] = val if not isinstance(val, enum.Enum) else val.value
|
||||
return payload
|
||||
|
||||
def _build_box_process_payload(self, host_path: str | None = None) -> dict:
|
||||
if host_path is None:
|
||||
host_path = self._resolve_host_path()
|
||||
|
||||
command = self.server_config['command']
|
||||
args = self.server_config.get('args', [])
|
||||
cwd = '/workspace'
|
||||
|
||||
if host_path:
|
||||
# When host_path is resolved, we install deps in-container rather
|
||||
# than relying on the host venv. Rewrite paths so the container
|
||||
# sees /workspace/... but replace venv python with plain "python".
|
||||
command = self._rewrite_venv_command(command, host_path)
|
||||
args = [self._rewrite_path(a, host_path) for a in args]
|
||||
cwd = self._rewrite_path(cwd, host_path)
|
||||
|
||||
return {
|
||||
'command': command,
|
||||
'args': args,
|
||||
'env': self.server_config.get('env', {}),
|
||||
'cwd': cwd,
|
||||
}
|
||||
|
||||
def _rewrite_venv_command(self, command: str, host_path: str) -> str:
|
||||
"""Rewrite command: if it points to a venv python, use plain 'python'."""
|
||||
if not host_path or not command:
|
||||
return command
|
||||
normalized_host = os.path.realpath(host_path)
|
||||
if not command.startswith(normalized_host + '/'):
|
||||
return command
|
||||
# Check if command is a venv python interpreter
|
||||
rel = command[len(normalized_host) + 1 :] # e.g. ".venv/bin/python"
|
||||
parts = rel.replace('\\', '/').split('/')
|
||||
# Match patterns like .venv/bin/python*, venv/bin/python*, etc.
|
||||
if len(parts) >= 3 and parts[0] in _VENV_DIRS and parts[1] in _VENV_BIN_DIRS and parts[2].startswith('python'):
|
||||
return 'python'
|
||||
# Not a venv python — do normal path rewrite
|
||||
return self._rewrite_path(command, host_path)
|
||||
|
||||
async def _cleanup_box_stdio_session(self) -> None:
|
||||
if not self._uses_box_stdio():
|
||||
return
|
||||
|
||||
try:
|
||||
await self.ap.box_service.client.delete_session(self._build_box_session_id())
|
||||
except Exception as e:
|
||||
self.ap.logger.warning(f'Failed to cleanup Box session for MCP server {self.server_name}: {e}')
|
||||
|
||||
|
||||
# @loader.loader_class('mcp')
|
||||
class MCPLoader(loader.ToolLoader):
|
||||
@@ -702,7 +332,7 @@ class MCPLoader(loader.ToolLoader):
|
||||
Args:
|
||||
server_config: 服务器配置字典，必须包含:
|
||||
- name: 服务器名称
|
||||
- mode: 连接模式 (stdio/sse/http)
|
||||
- mode: 连接模式 (stdio/sse)
|
||||
- enable: 是否启用
|
||||
- extra_args: 额外的配置参数 (可选)
|
||||
"""
|
||||
@@ -801,13 +431,12 @@ class MCPLoader(loader.ToolLoader):
|
||||
"""获取所有服务器的信息"""
|
||||
info = {}
|
||||
for server_name, session in self.sessions.items():
|
||||
tools = session.get_tools()
|
||||
info[server_name] = {
|
||||
'name': server_name,
|
||||
'mode': session.server_config.get('mode'),
|
||||
'enable': session.enable,
|
||||
'tools_count': len(tools),
|
||||
'tool_names': [f.name for f in tools],
|
||||
'tools_count': len(session.get_tools()),
|
||||
'tool_names': [f.name for f in session.get_tools()],
|
||||
}
|
||||
return info
|
||||
|
||||
|
||||
@@ -1,361 +0,0 @@
|
||||
"""Integration tests for Box MCP-related features.
|
||||
|
||||
These tests verify managed process lifecycle, WebSocket stdio attach,
|
||||
session cleanup, and the single-session query API using a real container
|
||||
runtime.
|
||||
|
||||
CI only runs ``tests/unit_tests/``, so these tests never execute in the
|
||||
CI pipeline. Run them locally with::
|
||||
|
||||
pytest tests/integration_tests/box/test_box_mcp_integration.py -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import shutil
|
||||
import socket
|
||||
import subprocess
|
||||
|
||||
import aiohttp
|
||||
import pytest
|
||||
from aiohttp.test_utils import TestServer
|
||||
|
||||
from langbot_plugin.box.client import ActionRPCBoxClient
|
||||
from langbot_plugin.box.errors import BoxSessionNotFoundError
|
||||
from langbot_plugin.box.models import BoxManagedProcessSpec, BoxManagedProcessStatus, BoxSpec
|
||||
from langbot_plugin.box.runtime import BoxRuntime
|
||||
from langbot_plugin.box.server import BoxServerHandler, create_ws_relay_app
|
||||
|
||||
_logger = logging.getLogger('test.box.mcp_integration')
|
||||
|
||||
_TEST_IMAGE = 'alpine:latest'
|
||||
|
||||
|
||||
# ── Skip helpers ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _has_container_runtime() -> bool:
|
||||
for cmd in ('podman', 'docker'):
|
||||
if shutil.which(cmd) is None:
|
||||
continue
|
||||
try:
|
||||
result = subprocess.run([cmd, 'info'], capture_output=True, timeout=10)
|
||||
if result.returncode == 0:
|
||||
return True
|
||||
except Exception:
|
||||
continue
|
||||
return False
|
||||
|
||||
|
||||
def _can_open_test_socket() -> bool:
|
||||
try:
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
except OSError:
|
||||
return False
|
||||
sock.close()
|
||||
return True
|
||||
|
||||
|
||||
requires_container = pytest.mark.skipif(
|
||||
not _has_container_runtime(),
|
||||
reason='no container runtime (podman/docker) available',
|
||||
)
|
||||
|
||||
requires_socket = pytest.mark.skipif(
|
||||
not _can_open_test_socket(),
|
||||
reason='local test environment does not permit opening TCP sockets',
|
||||
)
|
||||
|
||||
|
||||
# ── Helpers ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class _QueueConnection:
|
||||
"""In-process Connection backed by asyncio Queues — no real IO."""
|
||||
|
||||
def __init__(self, rx: asyncio.Queue[str], tx: asyncio.Queue[str]):
|
||||
self._rx = rx
|
||||
self._tx = tx
|
||||
|
||||
async def send(self, message: str) -> None:
|
||||
await self._tx.put(message)
|
||||
|
||||
async def receive(self) -> str:
|
||||
return await self._rx.get()
|
||||
|
||||
async def close(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
async def _make_rpc_pair(runtime: BoxRuntime):
|
||||
"""Create an in-process RPC pair connected via queues."""
|
||||
from langbot_plugin.runtime.io.handler import Handler
|
||||
|
||||
c2s: asyncio.Queue[str] = asyncio.Queue()
|
||||
s2c: asyncio.Queue[str] = asyncio.Queue()
|
||||
client_conn = _QueueConnection(rx=s2c, tx=c2s)
|
||||
server_conn = _QueueConnection(rx=c2s, tx=s2c)
|
||||
|
||||
server_handler = BoxServerHandler(server_conn, runtime)
|
||||
server_task = asyncio.create_task(server_handler.run())
|
||||
|
||||
client_handler = Handler.__new__(Handler)
|
||||
Handler.__init__(client_handler, client_conn)
|
||||
client_task = asyncio.create_task(client_handler.run())
|
||||
|
||||
client = ActionRPCBoxClient(logger=_logger)
|
||||
client.set_handler(client_handler)
|
||||
|
||||
return client, server_task, client_task
|
||||
|
||||
|
||||
# ── Fixtures ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def box_server():
|
||||
"""Yield a (ws_relay_url, ActionRPCBoxClient) backed by a real BoxRuntime."""
|
||||
runtime = BoxRuntime(logger=_logger)
|
||||
await runtime.initialize()
|
||||
|
||||
# Start ws relay for managed process attach
|
||||
ws_app = create_ws_relay_app(runtime)
|
||||
ws_server = TestServer(ws_app)
|
||||
await ws_server.start_server()
|
||||
|
||||
client, server_task, client_task = await _make_rpc_pair(runtime)
|
||||
|
||||
ws_relay_url = str(ws_server.make_url(''))
|
||||
yield ws_relay_url, client
|
||||
|
||||
server_task.cancel()
|
||||
client_task.cancel()
|
||||
await runtime.shutdown()
|
||||
await ws_server.close()
|
||||
|
||||
|
||||
# ── 1. Managed process lifecycle ─────────────────────────────────────
|
||||
|
||||
|
||||
@requires_container
|
||||
@requires_socket
|
||||
@pytest.mark.asyncio
|
||||
async def test_managed_process_start_and_query(box_server):
|
||||
"""Start a managed process and query its status."""
|
||||
ws_relay_url, client = box_server
|
||||
|
||||
# Create session
|
||||
spec = BoxSpec(
|
||||
cmd='',
|
||||
session_id='mcp-int-lifecycle',
|
||||
workdir='/tmp',
|
||||
image=_TEST_IMAGE,
|
||||
)
|
||||
await client.create_session(spec)
|
||||
|
||||
# Start a managed process that stays alive
|
||||
proc_spec = BoxManagedProcessSpec(
|
||||
command='sh',
|
||||
args=['-c', 'while true; do sleep 1; done'],
|
||||
cwd='/tmp',
|
||||
)
|
||||
info = await client.start_managed_process('mcp-int-lifecycle', proc_spec)
|
||||
assert info.status == BoxManagedProcessStatus.RUNNING
|
||||
|
||||
# Query it
|
||||
info2 = await client.get_managed_process('mcp-int-lifecycle')
|
||||
assert info2.status == BoxManagedProcessStatus.RUNNING
|
||||
assert info2.command == 'sh'
|
||||
|
||||
# Cleanup
|
||||
await client.delete_session('mcp-int-lifecycle')
|
||||
|
||||
|
||||
# ── 2. WebSocket stdio attach ────────────────────────────────────────
|
||||
|
||||
|
||||
@requires_container
|
||||
@requires_socket
|
||||
@pytest.mark.asyncio
|
||||
async def test_ws_stdio_attach_echo(box_server):
|
||||
"""Attach to a managed process via WebSocket and verify bidirectional IO."""
|
||||
ws_relay_url, client = box_server
|
||||
|
||||
spec = BoxSpec(
|
||||
cmd='',
|
||||
session_id='mcp-int-ws',
|
||||
workdir='/tmp',
|
||||
image=_TEST_IMAGE,
|
||||
)
|
||||
await client.create_session(spec)
|
||||
|
||||
# Start a cat process (echoes stdin to stdout)
|
||||
proc_spec = BoxManagedProcessSpec(
|
||||
command='cat',
|
||||
args=[],
|
||||
cwd='/tmp',
|
||||
)
|
||||
await client.start_managed_process('mcp-int-ws', proc_spec)
|
||||
|
||||
# Connect via WebSocket (ws relay)
|
||||
ws_url = client.get_managed_process_websocket_url('mcp-int-ws', ws_relay_url)
|
||||
session = aiohttp.ClientSession()
|
||||
try:
|
||||
async with session.ws_connect(ws_url) as ws:
|
||||
# Send a line
|
||||
await ws.send_str('hello from test')
|
||||
|
||||
# Expect to receive it back (cat echoes)
|
||||
msg = await asyncio.wait_for(ws.receive(), timeout=5)
|
||||
assert msg.type == aiohttp.WSMsgType.TEXT
|
||||
assert 'hello from test' in msg.data
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
await client.delete_session('mcp-int-ws')
|
||||
|
||||
|
||||
# ── 3. Session cleanup removes container ─────────────────────────────
|
||||
|
||||
|
||||
@requires_container
|
||||
@requires_socket
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_session_cleans_up(box_server):
|
||||
"""After deleting a session, it should no longer exist."""
|
||||
ws_relay_url, client = box_server
|
||||
|
||||
spec = BoxSpec(
|
||||
cmd='',
|
||||
session_id='mcp-int-cleanup',
|
||||
workdir='/tmp',
|
||||
image=_TEST_IMAGE,
|
||||
)
|
||||
await client.create_session(spec)
|
||||
|
||||
# Start a process
|
||||
proc_spec = BoxManagedProcessSpec(
|
||||
command='sleep',
|
||||
args=['3600'],
|
||||
cwd='/tmp',
|
||||
)
|
||||
await client.start_managed_process('mcp-int-cleanup', proc_spec)
|
||||
|
||||
# Delete
|
||||
await client.delete_session('mcp-int-cleanup')
|
||||
|
||||
# Session should be gone
|
||||
with pytest.raises(BoxSessionNotFoundError):
|
||||
await client.get_session('mcp-int-cleanup')
|
||||
|
||||
|
||||
# ── 4. GET session details ────────────────────────────────────────
|
||||
|
||||
|
||||
@requires_container
|
||||
@requires_socket
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_session_returns_details(box_server):
|
||||
"""Get single session returns session details and managed process info."""
|
||||
ws_relay_url, client = box_server
|
||||
|
||||
spec = BoxSpec(
|
||||
cmd='',
|
||||
session_id='mcp-int-get',
|
||||
workdir='/tmp',
|
||||
image=_TEST_IMAGE,
|
||||
)
|
||||
await client.create_session(spec)
|
||||
|
||||
# Query without managed process
|
||||
info = await client.get_session('mcp-int-get')
|
||||
assert info['session_id'] == 'mcp-int-get'
|
||||
assert info['image'] == _TEST_IMAGE
|
||||
assert 'managed_process' not in info
|
||||
|
||||
# Start a process and query again
|
||||
proc_spec = BoxManagedProcessSpec(
|
||||
command='sleep',
|
||||
args=['3600'],
|
||||
cwd='/tmp',
|
||||
)
|
||||
await client.start_managed_process('mcp-int-get', proc_spec)
|
||||
|
||||
info2 = await client.get_session('mcp-int-get')
|
||||
assert info2['session_id'] == 'mcp-int-get'
|
||||
assert 'managed_process' in info2
|
||||
assert info2['managed_process']['status'] == BoxManagedProcessStatus.RUNNING.value
|
||||
|
||||
await client.delete_session('mcp-int-get')
|
||||
|
||||
|
||||
# ── 5. Process exit detected ────────────────────────────────────────
|
||||
|
||||
|
||||
@requires_container
|
||||
@requires_socket
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_exit_detected(box_server):
|
||||
"""When a managed process exits, its status should reflect EXITED."""
|
||||
ws_relay_url, client = box_server
|
||||
|
||||
spec = BoxSpec(
|
||||
cmd='',
|
||||
session_id='mcp-int-exit',
|
||||
workdir='/tmp',
|
||||
image=_TEST_IMAGE,
|
||||
)
|
||||
await client.create_session(spec)
|
||||
|
||||
# Start a process that exits immediately
|
||||
proc_spec = BoxManagedProcessSpec(
|
||||
command='sh',
|
||||
args=['-c', 'echo done && exit 0'],
|
||||
cwd='/tmp',
|
||||
)
|
||||
await client.start_managed_process('mcp-int-exit', proc_spec)
|
||||
|
||||
# Wait a bit for process to exit
|
||||
await asyncio.sleep(2)
|
||||
|
||||
info = await client.get_managed_process('mcp-int-exit')
|
||||
assert info.status == BoxManagedProcessStatus.EXITED
|
||||
assert info.exit_code == 0
|
||||
|
||||
await client.delete_session('mcp-int-exit')
|
||||
|
||||
|
||||
# ── 6. Instance ID orphan cleanup ───────────────────────────────────
|
||||
|
||||
|
||||
@requires_container
|
||||
@requires_socket
|
||||
@pytest.mark.asyncio
|
||||
async def test_orphan_cleanup_preserves_own_containers(box_server):
|
||||
"""Orphan cleanup should not remove containers belonging to the current instance."""
|
||||
ws_relay_url, client = box_server
|
||||
|
||||
# Create a session (container gets current instance ID label)
|
||||
spec = BoxSpec(
|
||||
cmd='',
|
||||
session_id='mcp-int-orphan',
|
||||
workdir='/tmp',
|
||||
image=_TEST_IMAGE,
|
||||
)
|
||||
await client.create_session(spec)
|
||||
|
||||
# Verify session exists
|
||||
sessions = await client.get_sessions()
|
||||
assert any(s['session_id'] == 'mcp-int-orphan' for s in sessions)
|
||||
|
||||
# Trigger status check (which doesn't clean up own containers)
|
||||
status = await client.get_status()
|
||||
assert status['active_sessions'] >= 1
|
||||
|
||||
# Our session should still exist
|
||||
sessions = await client.get_sessions()
|
||||
assert any(s['session_id'] == 'mcp-int-orphan' for s in sessions)
|
||||
|
||||
await client.delete_session('mcp-int-orphan')
|
||||
@@ -1,635 +0,0 @@
|
||||
"""Tests for MCP Box integration: path rewriting, host_path inference, config model, payloads.
|
||||
|
||||
Uses importlib.util.spec_from_file_location to load mcp.py directly without
|
||||
triggering the circular import chain through the app module.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import importlib.util
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Load mcp.py directly from file path, with stub dependencies
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _stub_module(fqn: str, attrs: dict | None = None, is_package: bool = False):
|
||||
"""Create or return a stub module and register it in sys.modules."""
|
||||
if fqn in sys.modules:
|
||||
mod = sys.modules[fqn]
|
||||
else:
|
||||
mod = types.ModuleType(fqn)
|
||||
mod.__spec__ = importlib.machinery.ModuleSpec(fqn, None, is_package=is_package)
|
||||
if is_package:
|
||||
mod.__path__ = []
|
||||
sys.modules[fqn] = mod
|
||||
parts = fqn.rsplit('.', 1)
|
||||
if len(parts) == 2 and parts[0] in sys.modules:
|
||||
setattr(sys.modules[parts[0]], parts[1], mod)
|
||||
if attrs:
|
||||
for k, v in attrs.items():
|
||||
setattr(mod, k, v)
|
||||
return mod
|
||||
|
||||
|
||||
@pytest.fixture(scope='module', autouse=True)
|
||||
def mcp_module():
|
||||
"""Load mcp.py with minimal stubs to avoid circular imports."""
|
||||
saved = {}
|
||||
|
||||
def _save_and_stub(name, attrs=None, is_package=False):
|
||||
saved[name] = sys.modules.get(name)
|
||||
# Don't overwrite modules that already exist (from other test modules)
|
||||
if name in sys.modules:
|
||||
return
|
||||
_stub_module(name, attrs, is_package)
|
||||
|
||||
# Stub entire dependency chains as packages / modules
|
||||
_save_and_stub('langbot_plugin', is_package=True)
|
||||
_save_and_stub('langbot_plugin.api', is_package=True)
|
||||
_save_and_stub('langbot_plugin.api.entities', is_package=True)
|
||||
_save_and_stub('langbot_plugin.api.entities.events', is_package=True)
|
||||
_save_and_stub('langbot_plugin.api.entities.events.pipeline_query', {})
|
||||
_save_and_stub('langbot_plugin.api.entities.builtin', is_package=True)
|
||||
_save_and_stub('langbot_plugin.api.entities.builtin.resource', is_package=True)
|
||||
_save_and_stub(
|
||||
'langbot_plugin.api.entities.builtin.resource.tool',
|
||||
{
|
||||
'LLMTool': type('LLMTool', (), {}),
|
||||
},
|
||||
)
|
||||
_save_and_stub('langbot_plugin.api.entities.builtin.provider', is_package=True)
|
||||
_save_and_stub('langbot_plugin.api.entities.builtin.provider.message', {})
|
||||
_save_and_stub('sqlalchemy', {'select': Mock()})
|
||||
_save_and_stub('httpx', {'AsyncClient': Mock()})
|
||||
_save_and_stub('mcp', {'ClientSession': Mock, 'StdioServerParameters': Mock}, is_package=True)
|
||||
_save_and_stub('mcp.client', is_package=True)
|
||||
_save_and_stub('mcp.client.stdio', {'stdio_client': Mock()})
|
||||
_save_and_stub('mcp.client.sse', {'sse_client': Mock()})
|
||||
_save_and_stub('mcp.client.streamable_http', {'streamable_http_client': Mock()})
|
||||
_save_and_stub('mcp.client.websocket', {'websocket_client': Mock()})
|
||||
|
||||
# Stub the provider.tools.loader (source of circular import)
|
||||
_save_and_stub('langbot', is_package=True)
|
||||
_save_and_stub('langbot.pkg', is_package=True)
|
||||
_save_and_stub('langbot.pkg.provider', is_package=True)
|
||||
_save_and_stub('langbot.pkg.provider.tools', is_package=True)
|
||||
_save_and_stub(
|
||||
'langbot.pkg.provider.tools.loader',
|
||||
{
|
||||
'ToolLoader': type('ToolLoader', (), {'__init__': lambda self, ap: None}),
|
||||
},
|
||||
)
|
||||
_save_and_stub('langbot.pkg.provider.tools.loaders', is_package=True)
|
||||
_save_and_stub('langbot.pkg.core', is_package=True)
|
||||
_save_and_stub('langbot.pkg.core.app', {'Application': type('Application', (), {})})
|
||||
_save_and_stub('langbot.pkg.entity', is_package=True)
|
||||
_save_and_stub('langbot.pkg.entity.persistence', is_package=True)
|
||||
_save_and_stub('langbot.pkg.entity.persistence.mcp', {})
|
||||
|
||||
# box models
|
||||
import enum as _enum
|
||||
|
||||
class _BPS(str, _enum.Enum):
|
||||
RUNNING = 'running'
|
||||
EXITED = 'exited'
|
||||
|
||||
_save_and_stub('langbot_plugin.box', is_package=True)
|
||||
_save_and_stub('langbot_plugin.box.models', {'BoxManagedProcessStatus': _BPS})
|
||||
|
||||
# Now load mcp.py via spec_from_file_location
|
||||
mod_fqn = 'langbot.pkg.provider.tools.loaders.mcp'
|
||||
sys.modules.pop(mod_fqn, None)
|
||||
mcp_path = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
'..',
|
||||
'..',
|
||||
'..',
|
||||
'src',
|
||||
'langbot',
|
||||
'pkg',
|
||||
'provider',
|
||||
'tools',
|
||||
'loaders',
|
||||
'mcp.py',
|
||||
)
|
||||
mcp_path = os.path.normpath(mcp_path)
|
||||
spec = importlib.util.spec_from_file_location(mod_fqn, mcp_path)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
sys.modules[mod_fqn] = mod
|
||||
spec.loader.exec_module(mod)
|
||||
|
||||
yield mod
|
||||
|
||||
# Cleanup
|
||||
sys.modules.pop(mod_fqn, None)
|
||||
for name in reversed(list(saved)):
|
||||
if saved[name] is None:
|
||||
sys.modules.pop(name, None)
|
||||
else:
|
||||
sys.modules[name] = saved[name]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_ap():
|
||||
ap = Mock()
|
||||
ap.logger = Mock()
|
||||
ap.box_service = Mock()
|
||||
return ap
|
||||
|
||||
|
||||
def _make_session(mcp_module, server_config: dict, ap=None):
|
||||
if ap is None:
|
||||
ap = _make_ap()
|
||||
return mcp_module.RuntimeMCPSession(
|
||||
server_name=server_config.get('name', 'test-server'),
|
||||
server_config=server_config,
|
||||
enable=True,
|
||||
ap=ap,
|
||||
)
|
||||
|
||||
|
||||
# ── MCPServerBoxConfig ──────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestMCPServerBoxConfig:
|
||||
def test_default_values(self, mcp_module):
|
||||
cfg = mcp_module.MCPServerBoxConfig.model_validate({})
|
||||
assert cfg.image is None
|
||||
assert cfg.network == 'on'
|
||||
assert cfg.host_path is None
|
||||
assert cfg.host_path_mode == 'ro'
|
||||
assert cfg.env == {}
|
||||
assert cfg.startup_timeout_sec == 120
|
||||
assert cfg.cpus is None
|
||||
assert cfg.memory_mb is None
|
||||
assert cfg.pids_limit is None
|
||||
assert cfg.read_only_rootfs is None
|
||||
|
||||
def test_custom_values(self, mcp_module):
|
||||
cfg = mcp_module.MCPServerBoxConfig.model_validate(
|
||||
{
|
||||
'image': 'node:20',
|
||||
'network': 'on',
|
||||
'host_path': '/home/user/mcp',
|
||||
'host_path_mode': 'rw',
|
||||
'env': {'FOO': 'bar'},
|
||||
'startup_timeout_sec': 60,
|
||||
'cpus': 2.0,
|
||||
'memory_mb': 1024,
|
||||
'pids_limit': 256,
|
||||
'read_only_rootfs': False,
|
||||
}
|
||||
)
|
||||
assert cfg.image == 'node:20'
|
||||
assert cfg.network == 'on'
|
||||
assert cfg.cpus == 2.0
|
||||
assert cfg.memory_mb == 1024
|
||||
|
||||
def test_extra_fields_ignored(self, mcp_module):
|
||||
cfg = mcp_module.MCPServerBoxConfig.model_validate(
|
||||
{
|
||||
'image': 'node:20',
|
||||
'unknown_field': 'whatever',
|
||||
}
|
||||
)
|
||||
assert cfg.image == 'node:20'
|
||||
assert not hasattr(cfg, 'unknown_field')
|
||||
|
||||
|
||||
# ── Path Rewriting ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestRewritePath:
|
||||
def test_no_host_path_returns_unchanged(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
assert s._rewrite_path('/some/path', None) == '/some/path'
|
||||
|
||||
def test_empty_path_returns_empty(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
assert s._rewrite_path('', '/home/user/mcp') == ''
|
||||
|
||||
def test_prefix_match_rewrites(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
result = s._rewrite_path('/home/user/mcp/server.py', '/home/user/mcp')
|
||||
assert result == '/workspace/server.py'
|
||||
|
||||
def test_exact_match_rewrites_to_workspace(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
result = s._rewrite_path('/home/user/mcp', '/home/user/mcp')
|
||||
assert result == '/workspace'
|
||||
|
||||
def test_non_matching_path_unchanged(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
result = s._rewrite_path('/opt/other/server.py', '/home/user/mcp')
|
||||
assert result == '/opt/other/server.py'
|
||||
|
||||
def test_similar_prefix_not_rewritten(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
result = s._rewrite_path('/home/user/mcp-other/file.py', '/home/user/mcp')
|
||||
assert result == '/home/user/mcp-other/file.py'
|
||||
|
||||
def test_nested_subpath_rewrites(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
result = s._rewrite_path('/home/user/mcp/src/lib/main.py', '/home/user/mcp')
|
||||
assert result == '/workspace/src/lib/main.py'
|
||||
|
||||
|
||||
# ── host_path Inference ─────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestInferHostPath:
|
||||
def test_no_absolute_paths_returns_none(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': ['server.py'],
|
||||
},
|
||||
)
|
||||
assert s._infer_host_path() is None
|
||||
|
||||
def test_nonexistent_path_returns_none(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': '/nonexistent/path/to/python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
assert s._infer_host_path() is None
|
||||
|
||||
def test_existing_absolute_path_infers_directory(self, mcp_module):
|
||||
with tempfile.NamedTemporaryFile(suffix='.py') as f:
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [f.name],
|
||||
},
|
||||
)
|
||||
result = s._infer_host_path()
|
||||
assert result is not None
|
||||
assert result == os.path.dirname(os.path.realpath(f.name))
|
||||
|
||||
|
||||
# ── Build Box Session Payload ───────────────────────────────────────
|
||||
|
||||
|
||||
class TestBuildBoxSessionPayload:
|
||||
def test_minimal_config(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
payload = s._build_box_session_payload('session-123')
|
||||
assert payload['session_id'] == 'session-123'
|
||||
assert payload['workdir'] == '/workspace'
|
||||
assert payload['env'] == {}
|
||||
assert 'host_path' not in payload
|
||||
|
||||
def test_with_host_path(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
'box': {'host_path': '/home/user/mcp', 'host_path_mode': 'ro'},
|
||||
},
|
||||
)
|
||||
payload = s._build_box_session_payload('session-123')
|
||||
assert payload['host_path'] == '/home/user/mcp'
|
||||
assert payload['host_path_mode'] == 'ro'
|
||||
|
||||
def test_optional_fields_included_when_set(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
'box': {'image': 'node:20', 'cpus': 2.0, 'memory_mb': 1024, 'pids_limit': 256},
|
||||
},
|
||||
)
|
||||
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['pids_limit'] == 256
|
||||
|
||||
def test_none_fields_excluded(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
payload = s._build_box_session_payload('session-123')
|
||||
assert 'image' not in payload
|
||||
assert 'cpus' not in payload
|
||||
|
||||
|
||||
# ── Build Box Process Payload ───────────────────────────────────────
|
||||
|
||||
|
||||
class TestBuildBoxProcessPayload:
|
||||
def test_basic_payload(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': ['server.py'],
|
||||
'env': {'KEY': 'val'},
|
||||
},
|
||||
)
|
||||
payload = s._build_box_process_payload()
|
||||
assert payload['command'] == 'python'
|
||||
assert payload['args'] == ['server.py']
|
||||
assert payload['env'] == {'KEY': 'val'}
|
||||
assert payload['cwd'] == '/workspace'
|
||||
|
||||
def test_path_rewriting_applied(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': '/home/user/mcp/venv/bin/python',
|
||||
'args': ['/home/user/mcp/server.py', '--config', '/home/user/mcp/config.json'],
|
||||
'env': {},
|
||||
'box': {'host_path': '/home/user/mcp'},
|
||||
},
|
||||
)
|
||||
payload = s._build_box_process_payload()
|
||||
# venv python is replaced with plain 'python' (deps installed in-container)
|
||||
assert payload['command'] == 'python'
|
||||
assert payload['args'] == ['/workspace/server.py', '--config', '/workspace/config.json']
|
||||
|
||||
def test_non_matching_args_not_rewritten(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': ['/opt/other/server.py', '--flag'],
|
||||
'env': {},
|
||||
'box': {'host_path': '/home/user/mcp'},
|
||||
},
|
||||
)
|
||||
payload = s._build_box_process_payload()
|
||||
assert payload['command'] == 'python'
|
||||
assert payload['args'] == ['/opt/other/server.py', '--flag']
|
||||
|
||||
|
||||
# ── get_runtime_info_dict ───────────────────────────────────────────
|
||||
|
||||
|
||||
class TestGetRuntimeInfoDict:
|
||||
def test_non_stdio_session(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'test-uuid',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
info = s.get_runtime_info_dict()
|
||||
assert info['status'] == 'connecting'
|
||||
assert 'box_session_id' not in info
|
||||
|
||||
def test_stdio_session_includes_box_info(self, mcp_module):
|
||||
ap = _make_ap()
|
||||
ap.box_service.available = True
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'test-uuid',
|
||||
'mode': 'stdio',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
ap=ap,
|
||||
)
|
||||
info = s.get_runtime_info_dict()
|
||||
assert info['box_session_id'] == 'mcp-test-uuid'
|
||||
assert info['box_enabled'] is True
|
||||
|
||||
def test_stdio_session_without_box_runtime(self, mcp_module):
|
||||
ap = _make_ap()
|
||||
ap.box_service.available = 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
|
||||
|
||||
|
||||
# ── Box config parsing ──────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestBoxConfigParsing:
|
||||
def test_box_config_parsed_from_server_config(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
'box': {'image': 'node:20', 'host_path': '/home/user/mcp'},
|
||||
},
|
||||
)
|
||||
assert isinstance(s.box_config, mcp_module.MCPServerBoxConfig)
|
||||
assert s.box_config.image == 'node:20'
|
||||
assert s.box_config.host_path == '/home/user/mcp'
|
||||
|
||||
def test_missing_box_key_uses_defaults(self, mcp_module):
|
||||
s = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'sse',
|
||||
'command': 'python',
|
||||
'args': [],
|
||||
},
|
||||
)
|
||||
assert isinstance(s.box_config, mcp_module.MCPServerBoxConfig)
|
||||
assert s.box_config.image is None
|
||||
assert s.box_config.host_path_mode == 'ro'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_box_stdio_server_keeps_host_mount_validation_enabled(mcp_module):
|
||||
class FakeClientSession:
|
||||
def __init__(self, *_args):
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def initialize(self):
|
||||
return None
|
||||
|
||||
@asynccontextmanager
|
||||
async def fake_websocket_client(_url: str):
|
||||
yield ('read-stream', 'write-stream')
|
||||
|
||||
mcp_module.ClientSession = FakeClientSession
|
||||
mcp_module.websocket_client = fake_websocket_client
|
||||
|
||||
ap = _make_ap()
|
||||
ap.box_service.available = True
|
||||
ap.box_service.create_session = AsyncMock(return_value={})
|
||||
ap.box_service.build_spec = Mock(return_value='validated-spec')
|
||||
ap.box_service.client = SimpleNamespace(
|
||||
execute=AsyncMock(return_value=SimpleNamespace(ok=True, stderr='', exit_code=0))
|
||||
)
|
||||
ap.box_service.start_managed_process = AsyncMock(return_value={})
|
||||
ap.box_service.get_managed_process_websocket_url = Mock(return_value='ws://box.example/process')
|
||||
|
||||
session = _make_session(
|
||||
mcp_module,
|
||||
{
|
||||
'name': 'test',
|
||||
'uuid': 'u1',
|
||||
'mode': 'stdio',
|
||||
'command': '/home/user/mcp/.venv/bin/python',
|
||||
'args': ['/home/user/mcp/server.py'],
|
||||
'box': {'host_path': '/home/user/mcp'},
|
||||
},
|
||||
ap=ap,
|
||||
)
|
||||
session._detect_install_command = Mock(return_value='pip install --no-cache-dir -r /workspace/requirements.txt')
|
||||
|
||||
await session._init_box_stdio_server()
|
||||
await session.exit_stack.aclose()
|
||||
|
||||
assert ap.box_service.create_session.await_count == 1
|
||||
assert ap.box_service.create_session.await_args.kwargs.get('skip_host_mount_validation', False) is False
|
||||
assert ap.box_service.build_spec.call_count == 1
|
||||
assert ap.box_service.build_spec.call_args.kwargs.get('skip_host_mount_validation', False) is False
|
||||
Reference in New Issue
Block a user