mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-22 12:26:08 +00:00
huanghuoguoguo
1184 lines
41 KiB
Python
1184 lines
41 KiB
Python
"""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 asyncio
|
|
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)
|
|
pkg_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(mcp_path))))
|
|
sys.modules['langbot.pkg'].__path__ = [pkg_root]
|
|
sys.modules['langbot.pkg.provider.tools.loaders'].__path__ = [os.path.dirname(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)
|
|
sys.modules.pop('langbot.pkg.provider.tools.loaders.mcp_stdio', None)
|
|
sys.modules.pop('langbot.pkg.box.workspace', 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 == 300
|
|
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']
|
|
|
|
|
|
# ── Python Workspace Preparation ────────────────────────────────────
|
|
|
|
|
|
class TestPythonWorkspacePreparation:
|
|
def test_requirements_workspace_uses_venv_bootstrap(self, mcp_module, tmp_path):
|
|
host_path = tmp_path / 'mcp-source'
|
|
host_path.mkdir()
|
|
(host_path / 'requirements.txt').write_text('mcp==1.26.0\n', encoding='utf-8')
|
|
|
|
command = mcp_module.BoxStdioSessionRuntime.detect_install_command(
|
|
str(host_path),
|
|
'/workspace/.mcp/u1/workspace',
|
|
)
|
|
|
|
assert command is not None
|
|
assert '_LB_SYSTEM_PYTHON="$(command -v python3 || command -v python || true)"' in command
|
|
assert '"$_LB_SYSTEM_PYTHON" -m venv "$_LB_VENV_DIR"' in command
|
|
assert 'python -m pip install -r "/workspace/.mcp/u1/workspace/requirements.txt"' in command
|
|
assert 'pip install --no-cache-dir -r' not in command
|
|
|
|
def test_staging_refresh_removes_stale_source_files_but_preserves_runtime_dirs(self, mcp_module, tmp_path):
|
|
source = tmp_path / 'source'
|
|
source.mkdir()
|
|
(source / 'server.py').write_text('print("new")\n', encoding='utf-8')
|
|
(source / 'requirements.txt').write_text('mcp==1.26.0\n', encoding='utf-8')
|
|
(source / '.env').write_text('TOKEN=new\n', encoding='utf-8')
|
|
|
|
process_root = tmp_path / 'shared' / '.mcp' / 'u1'
|
|
workspace = process_root / 'workspace'
|
|
(workspace / '.venv' / 'bin').mkdir(parents=True)
|
|
(workspace / '.venv' / 'bin' / 'python').write_text('', encoding='utf-8')
|
|
(workspace / '.langbot').mkdir()
|
|
(workspace / '.langbot' / 'python-env.lock').mkdir()
|
|
(workspace / '.env').write_text('TOKEN=old\n', encoding='utf-8')
|
|
(workspace / 'server.py').write_text('print("old")\n', encoding='utf-8')
|
|
(workspace / 'removed.py').write_text('stale\n', encoding='utf-8')
|
|
(workspace / 'removed_dir').mkdir()
|
|
(workspace / 'removed_dir' / 'old.txt').write_text('stale\n', encoding='utf-8')
|
|
|
|
mcp_module.BoxStdioSessionRuntime._copy_workspace_tree(str(source), str(process_root), str(workspace))
|
|
|
|
assert (workspace / 'server.py').read_text(encoding='utf-8') == 'print("new")\n'
|
|
assert (workspace / 'requirements.txt').read_text(encoding='utf-8') == 'mcp==1.26.0\n'
|
|
assert (workspace / '.env').read_text(encoding='utf-8') == 'TOKEN=new\n'
|
|
assert not (workspace / 'removed.py').exists()
|
|
assert not (workspace / 'removed_dir').exists()
|
|
assert (workspace / '.venv' / 'bin' / 'python').exists()
|
|
assert (workspace / '.langbot' / 'python-env.lock').is_dir()
|
|
|
|
def test_staging_refresh_ignores_unlink_race(self, mcp_module, tmp_path, monkeypatch):
|
|
mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio']
|
|
|
|
source = tmp_path / 'source'
|
|
source.mkdir()
|
|
(source / 'server.py').write_text('print("new")\n', encoding='utf-8')
|
|
|
|
process_root = tmp_path / 'shared' / '.mcp' / 'u1'
|
|
workspace = process_root / 'workspace'
|
|
workspace.mkdir(parents=True)
|
|
stale_file = workspace / 'removed.py'
|
|
stale_file.write_text('stale\n', encoding='utf-8')
|
|
|
|
real_unlink = os.unlink
|
|
|
|
def unlink_with_race(path):
|
|
if os.fspath(path) == str(stale_file):
|
|
real_unlink(path)
|
|
raise FileNotFoundError(path)
|
|
real_unlink(path)
|
|
|
|
monkeypatch.setattr(mcp_stdio_module.os, 'unlink', unlink_with_race)
|
|
|
|
mcp_module.BoxStdioSessionRuntime._copy_workspace_tree(str(source), str(process_root), str(workspace))
|
|
|
|
assert not stale_file.exists()
|
|
assert (workspace / 'server.py').read_text(encoding='utf-8') == 'print("new")\n'
|
|
|
|
|
|
# ── 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_runtime_tools_include_parameters(self, mcp_module):
|
|
s = _make_session(
|
|
mcp_module,
|
|
{
|
|
'name': 'test',
|
|
'uuid': 'test-uuid',
|
|
'mode': 'sse',
|
|
'command': 'python',
|
|
'args': [],
|
|
},
|
|
)
|
|
s.functions = [
|
|
SimpleNamespace(
|
|
name='create-service',
|
|
description='Create a service',
|
|
parameters={
|
|
'type': 'object',
|
|
'properties': {
|
|
'project_id': {'type': 'string'},
|
|
},
|
|
'required': ['project_id'],
|
|
},
|
|
)
|
|
]
|
|
|
|
info = s.get_runtime_info_dict()
|
|
|
|
assert info['tools'][0]['parameters']['properties']['project_id']['type'] == 'string'
|
|
assert info['tools'][0]['parameters']['required'] == ['project_id']
|
|
|
|
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-shared'
|
|
assert info['box_enabled'] is True
|
|
|
|
def test_transient_test_shares_session_but_isolated_by_process(self, mcp_module):
|
|
"""A transient config-page "test" now shares the same 'mcp-shared' Box
|
|
session as live servers (so a test reuses the running container / live
|
|
process instead of a cold per-test session bootstrap). Isolation is at
|
|
the PROCESS level: the test runs under its own process_id and only ever
|
|
stops that process_id, so it cannot disturb another server's live
|
|
process or the shared session itself."""
|
|
ap = _make_ap()
|
|
ap.box_service.available = True
|
|
transient = _make_session(
|
|
mcp_module,
|
|
{
|
|
'name': 'test',
|
|
'uuid': 'gen-uuid-123',
|
|
'mode': 'stdio',
|
|
'command': 'uvx',
|
|
'args': ['mcp-server-time'],
|
|
'_transient': True,
|
|
},
|
|
ap=ap,
|
|
)
|
|
live = _make_session(
|
|
mcp_module,
|
|
{
|
|
'name': 'time',
|
|
'uuid': 'real-uuid',
|
|
'mode': 'stdio',
|
|
'command': 'uvx',
|
|
'args': ['mcp-server-time'],
|
|
},
|
|
ap=ap,
|
|
)
|
|
assert transient.is_transient is True
|
|
assert live.is_transient is False
|
|
# Both share ONE Box session ...
|
|
assert transient._build_box_session_id() == 'mcp-shared'
|
|
assert live._build_box_session_id() == 'mcp-shared'
|
|
assert transient._build_box_session_id() == live._build_box_session_id()
|
|
# ... but are isolated by distinct process_ids within that session.
|
|
assert transient._box_stdio_runtime.process_id != live._box_stdio_runtime.process_id
|
|
|
|
def test_stdio_session_keeps_box_transport_while_runtime_reconnects(self, mcp_module):
|
|
ap = _make_ap()
|
|
ap.box_service.available = False
|
|
ap.box_service.enabled = 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-shared'
|
|
assert info['box_enabled'] is True
|
|
|
|
def test_stdio_session_refuses_when_box_is_disabled(self, mcp_module):
|
|
ap = _make_ap()
|
|
ap.box_service.available = False
|
|
ap.box_service.enabled = False
|
|
s = _make_session(
|
|
mcp_module,
|
|
{
|
|
'name': 'test',
|
|
'uuid': 'test-uuid',
|
|
'mode': 'stdio',
|
|
'command': 'python',
|
|
'args': [],
|
|
},
|
|
ap=ap,
|
|
)
|
|
info = s.get_runtime_info_dict()
|
|
assert 'box_session_id' not in info
|
|
assert 'box_enabled' not in info
|
|
|
|
def test_stdio_session_without_box_service_uses_local_stdio(self, mcp_module):
|
|
ap = _make_ap()
|
|
del ap.box_service
|
|
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_lifecycle_cleanup_timeout_does_not_block_box_stdio_retry(mcp_module, monkeypatch):
|
|
class HangingExitStack:
|
|
async def aclose(self):
|
|
await asyncio.Event().wait()
|
|
|
|
session = _make_session(
|
|
mcp_module,
|
|
{
|
|
'name': 'cleanup-timeout',
|
|
'uuid': 'cleanup-timeout-uuid',
|
|
'mode': 'stdio',
|
|
'command': 'python',
|
|
'args': [],
|
|
},
|
|
)
|
|
session.exit_stack = HangingExitStack()
|
|
session.functions.append(Mock())
|
|
session.resources.append({'uri': 'test://stale'})
|
|
session.session = Mock()
|
|
session._cleanup_box_stdio_session = AsyncMock()
|
|
monkeypatch.setattr(mcp_module, 'MCP_TRANSPORT_CLEANUP_TIMEOUT_SECONDS', 0.01)
|
|
|
|
await asyncio.wait_for(session._cleanup_lifecycle_attempt(), timeout=1)
|
|
|
|
assert isinstance(session.exit_stack, mcp_module.AsyncExitStack)
|
|
assert session.functions == []
|
|
assert session.resources == []
|
|
assert session.session is None
|
|
session._cleanup_box_stdio_session.assert_awaited_once()
|
|
session.ap.logger.warning.assert_called_once_with(
|
|
'Timed out cleaning up MCP transport for cleanup-timeout; continuing lifecycle recovery'
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lifecycle_transport_cleanup_error_is_recovery_warning(mcp_module):
|
|
class FailingExitStack:
|
|
async def aclose(self):
|
|
raise RuntimeError('transport already closed')
|
|
|
|
session = _make_session(
|
|
mcp_module,
|
|
{
|
|
'name': 'cleanup-error',
|
|
'uuid': 'cleanup-error-uuid',
|
|
'mode': 'stdio',
|
|
'command': 'python',
|
|
'args': [],
|
|
},
|
|
)
|
|
session.exit_stack = FailingExitStack()
|
|
session._cleanup_box_stdio_session = AsyncMock()
|
|
|
|
await session._cleanup_lifecycle_attempt()
|
|
|
|
session.ap.logger.warning.assert_called_once_with(
|
|
'Error cleaning up MCP transport for cleanup-error; '
|
|
'continuing lifecycle recovery: RuntimeError: transport already closed'
|
|
)
|
|
session.ap.logger.error.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unexpected_box_stdio_transport_cancellation_becomes_retryable_error(mcp_module, monkeypatch):
|
|
session = _make_session(
|
|
mcp_module,
|
|
{
|
|
'name': 'cancelled-transport',
|
|
'uuid': 'cancelled-transport-uuid',
|
|
'mode': 'stdio',
|
|
'command': 'python',
|
|
'args': [],
|
|
},
|
|
)
|
|
session._uses_box_stdio = Mock(return_value=True)
|
|
lifecycle_calls = 0
|
|
|
|
async def cancelled_then_shutdown():
|
|
nonlocal lifecycle_calls
|
|
lifecycle_calls += 1
|
|
if lifecycle_calls == 1:
|
|
raise asyncio.CancelledError()
|
|
session._shutdown_event.set()
|
|
|
|
session._lifecycle_loop = AsyncMock(side_effect=cancelled_then_shutdown)
|
|
session._cleanup_box_stdio_session = AsyncMock()
|
|
session.functions.append(Mock())
|
|
session.resources.append({'uri': 'test://stale'})
|
|
session.resource_templates.append({'uri_template': 'test://{id}'})
|
|
session.resource_capabilities = {'subscribe': True}
|
|
session._resource_cache[('test://stale', 1, None, False)] = {'content': 'stale'}
|
|
session.session = Mock()
|
|
monkeypatch.setattr(session, '_RETRY_DELAYS', [0, 0, 0])
|
|
|
|
await session._lifecycle_loop_with_retry()
|
|
|
|
assert session._lifecycle_loop.await_count == 2
|
|
assert session.functions == []
|
|
assert session.resources == []
|
|
assert session.resource_templates == []
|
|
assert session.resource_capabilities == {}
|
|
assert session._resource_cache == {}
|
|
assert session.session is None
|
|
session.ap.logger.error.assert_called_once_with(
|
|
'Error in MCP session lifecycle cancelled-transport: '
|
|
'Box MCP transport task was cancelled unexpectedly'
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_lifecycle_closes_transport_in_its_own_task(mcp_module):
|
|
class TrackingExitStack:
|
|
def __init__(self):
|
|
self.close_task = None
|
|
|
|
async def aclose(self):
|
|
self.close_task = asyncio.current_task()
|
|
|
|
session = _make_session(
|
|
mcp_module,
|
|
{
|
|
'name': 'same-task-cleanup',
|
|
'uuid': 'same-task-cleanup-uuid',
|
|
'mode': 'stdio',
|
|
'command': 'python',
|
|
'args': [],
|
|
},
|
|
)
|
|
transport_stack = TrackingExitStack()
|
|
session.exit_stack = transport_stack
|
|
session._init_stdio_python_server = AsyncMock(side_effect=asyncio.CancelledError())
|
|
session._cleanup_box_stdio_session = AsyncMock()
|
|
|
|
lifecycle_task = asyncio.create_task(session._lifecycle_loop())
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await lifecycle_task
|
|
|
|
assert transport_stack.close_task is lifecycle_task
|
|
session._cleanup_box_stdio_session.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_connected_lifecycle_failures_receive_fresh_retry_budgets(mcp_module, monkeypatch):
|
|
session = _make_session(
|
|
mcp_module,
|
|
{
|
|
'name': 'repeated-runtime-recovery',
|
|
'uuid': 'repeated-runtime-recovery-uuid',
|
|
'mode': 'stdio',
|
|
'command': 'python',
|
|
'args': [],
|
|
},
|
|
)
|
|
session._uses_box_stdio = Mock(return_value=True)
|
|
session._cleanup_box_stdio_session = AsyncMock()
|
|
lifecycle_calls = 0
|
|
|
|
async def connected_then_failed_repeatedly():
|
|
nonlocal lifecycle_calls
|
|
lifecycle_calls += 1
|
|
if lifecycle_calls <= 4:
|
|
session._connection_generation += 1
|
|
raise RuntimeError(f'runtime failure {lifecycle_calls}')
|
|
session._shutdown_event.set()
|
|
|
|
session._lifecycle_loop = AsyncMock(side_effect=connected_then_failed_repeatedly)
|
|
monkeypatch.setattr(session, '_MAX_RETRIES', 1)
|
|
monkeypatch.setattr(session, '_RETRY_DELAYS', [0])
|
|
|
|
await session._lifecycle_loop_with_retry()
|
|
|
|
assert session._lifecycle_loop.await_count == 5
|
|
assert session.retry_count == 4
|
|
assert session.status != mcp_module.MCPSessionStatus.ERROR
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unexpected_box_stdio_lifecycle_return_is_retried(mcp_module, monkeypatch):
|
|
session = _make_session(
|
|
mcp_module,
|
|
{
|
|
'name': 'ended-transport',
|
|
'uuid': 'ended-transport-uuid',
|
|
'mode': 'stdio',
|
|
'command': 'python',
|
|
'args': [],
|
|
},
|
|
)
|
|
session._uses_box_stdio = Mock(return_value=True)
|
|
session._cleanup_box_stdio_session = AsyncMock()
|
|
lifecycle_calls = 0
|
|
|
|
async def ended_then_shutdown():
|
|
nonlocal lifecycle_calls
|
|
lifecycle_calls += 1
|
|
if lifecycle_calls == 2:
|
|
session._shutdown_event.set()
|
|
|
|
session._lifecycle_loop = AsyncMock(side_effect=ended_then_shutdown)
|
|
monkeypatch.setattr(session, '_RETRY_DELAYS', [0, 0, 0])
|
|
|
|
await session._lifecycle_loop_with_retry()
|
|
|
|
assert session._lifecycle_loop.await_count == 2
|
|
session.ap.logger.error.assert_called_once_with(
|
|
'Error in MCP session lifecycle ended-transport: Box MCP lifecycle ended unexpectedly'
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_box_stdio_server_stages_host_path_in_shared_workspace(mcp_module, tmp_path):
|
|
mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio']
|
|
|
|
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_stdio_module.ClientSession = FakeClientSession
|
|
mcp_stdio_module.websocket_client = fake_websocket_client
|
|
|
|
ap = _make_ap()
|
|
ap.box_service.available = True
|
|
ap.box_service.default_workspace = str(tmp_path / 'shared-box-workspace')
|
|
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')
|
|
|
|
host_path = tmp_path / 'mcp-source'
|
|
host_path.mkdir()
|
|
server_file = host_path / 'server.py'
|
|
server_file.write_text('print("hello")\n', encoding='utf-8')
|
|
|
|
session = _make_session(
|
|
mcp_module,
|
|
{
|
|
'name': 'test',
|
|
'uuid': 'u1',
|
|
'mode': 'stdio',
|
|
'command': str(host_path / '.venv' / 'bin' / 'python'),
|
|
'args': [str(server_file)],
|
|
'box': {'host_path': str(host_path)},
|
|
},
|
|
ap=ap,
|
|
)
|
|
|
|
await session._init_box_stdio_server()
|
|
await session.exit_stack.aclose()
|
|
|
|
assert ap.box_service.create_session.await_count == 1
|
|
session_payload = ap.box_service.create_session.await_args.args[0]
|
|
assert session_payload['session_id'] == 'mcp-shared'
|
|
assert 'host_path' not in session_payload
|
|
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
|
|
assert ap.box_service.build_spec.call_args.args[0]['host_path'] == str(host_path)
|
|
|
|
staged_file = tmp_path / 'shared-box-workspace' / '.mcp' / 'u1' / 'workspace' / 'server.py'
|
|
assert staged_file.read_text(encoding='utf-8') == 'print("hello")\n'
|
|
|
|
process_payload = ap.box_service.start_managed_process.await_args.args[1]
|
|
assert process_payload['process_id'] == 'u1'
|
|
assert process_payload['command'] == 'python'
|
|
assert process_payload['args'] == ['/workspace/.mcp/u1/workspace/server.py']
|
|
assert process_payload['cwd'] == '/workspace/.mcp/u1/workspace'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stdio_handshake_raises_coldstart_retry_while_process_alive(mcp_module, tmp_path, monkeypatch):
|
|
"""During a slow (npx) cold start the handshake fails while the managed
|
|
process is still alive. initialize() must raise _ColdStartRetry (so the
|
|
outer lifecycle loop reuses the live process and retries without stopping it
|
|
or consuming the fatal budget), NOT a fatal error."""
|
|
from contextlib import asynccontextmanager
|
|
|
|
mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio']
|
|
|
|
class ColdClientSession:
|
|
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):
|
|
# Process still cold-starting: handshake fails.
|
|
raise Exception('Connection closed')
|
|
|
|
@asynccontextmanager
|
|
async def fake_websocket_client(_url: str):
|
|
yield ('read-stream', 'write-stream')
|
|
|
|
monkeypatch.setattr(mcp_stdio_module, 'ClientSession', ColdClientSession)
|
|
monkeypatch.setattr(mcp_stdio_module, 'websocket_client', fake_websocket_client)
|
|
monkeypatch.setattr(mcp_stdio_module, '_HANDSHAKE_ATTEMPT_TIMEOUT_SEC', 1.0, raising=False)
|
|
|
|
ap = _make_ap()
|
|
ap.box_service.available = True
|
|
ap.box_service.create_session = AsyncMock(return_value={})
|
|
ap.box_service.start_managed_process = AsyncMock(return_value={})
|
|
ap.box_service.get_managed_process_websocket_url = Mock(return_value='ws://box/p')
|
|
|
|
session = _make_session(
|
|
mcp_module,
|
|
{
|
|
'name': 'slow',
|
|
'uuid': 'slow-uuid',
|
|
'mode': 'stdio',
|
|
'command': 'npx',
|
|
'args': ['-y', 'some-mcp'],
|
|
},
|
|
ap=ap,
|
|
)
|
|
|
|
# Process is NOT exited (still cold-starting) and not yet running for reuse.
|
|
async def _not_exited():
|
|
return False
|
|
|
|
session._box_stdio_runtime._managed_process_has_exited = _not_exited
|
|
|
|
async def _not_running():
|
|
return False
|
|
|
|
session._box_stdio_runtime._managed_process_is_running = _not_running
|
|
|
|
with pytest.raises(mcp_stdio_module._ColdStartRetry):
|
|
await session._init_box_stdio_server()
|
|
|
|
# Process was started exactly once (the retry will reuse it, not rebuild).
|
|
assert ap.box_service.start_managed_process.await_count == 1
|
|
await session.exit_stack.aclose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stdio_handshake_raises_fatal_when_process_exited(mcp_module, tmp_path, monkeypatch):
|
|
"""If the handshake fails AND the process has definitively exited, that is a
|
|
real failure — initialize() must NOT swallow it as a cold-start retry."""
|
|
from contextlib import asynccontextmanager
|
|
|
|
mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio']
|
|
|
|
class DeadClientSession:
|
|
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):
|
|
raise Exception('Connection closed')
|
|
|
|
@asynccontextmanager
|
|
async def fake_websocket_client(_url: str):
|
|
yield ('read-stream', 'write-stream')
|
|
|
|
monkeypatch.setattr(mcp_stdio_module, 'ClientSession', DeadClientSession)
|
|
monkeypatch.setattr(mcp_stdio_module, 'websocket_client', fake_websocket_client)
|
|
monkeypatch.setattr(mcp_stdio_module, '_HANDSHAKE_ATTEMPT_TIMEOUT_SEC', 1.0, raising=False)
|
|
|
|
ap = _make_ap()
|
|
ap.box_service.available = True
|
|
ap.box_service.create_session = AsyncMock(return_value={})
|
|
ap.box_service.start_managed_process = AsyncMock(return_value={})
|
|
ap.box_service.get_managed_process_websocket_url = Mock(return_value='ws://box/p')
|
|
|
|
session = _make_session(
|
|
mcp_module,
|
|
{'name': 'dead', 'uuid': 'dead-uuid', 'mode': 'stdio', 'command': 'npx', 'args': ['-y', 'x']},
|
|
ap=ap,
|
|
)
|
|
|
|
async def _exited():
|
|
return True
|
|
|
|
session._box_stdio_runtime._managed_process_has_exited = _exited
|
|
|
|
async def _not_running():
|
|
return False
|
|
|
|
session._box_stdio_runtime._managed_process_is_running = _not_running
|
|
|
|
with pytest.raises(Exception) as ei:
|
|
await session._init_box_stdio_server()
|
|
assert not isinstance(ei.value, mcp_stdio_module._ColdStartRetry)
|
|
await session.exit_stack.aclose()
|