mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-14 08:26:07 +00:00
Hyu
ed9343c686
A node/npx stdio MCP server (e.g. firecrawl-mcp via npx -y) failed on first connect with Connection closed / Failed after 4 attempts, even though the process was fine. An npx cold start downloads+installs the package before the server can answer the MCP handshake (measured ~27s for a simple official server; longer for heavier ones). The old code attached the WS and called session.initialize() the instant the process was started, so the handshake ran before the process could answer and failed; the outer lifecycle retry then rebuilt the process, churning it in a loop. Verified decisively: attaching + initialize() against a mid-cold-start process times out on attempt 1 (process still installing) but SUCCEEDS at t+0.6s on attempt 2 once the process is ready. So the fix is to retry the handshake in place, not to rebuild the process. Changes (mcp_stdio.initialize): - Start the managed process ONCE, then loop attach WS -> ClientSession -> initialize() within the startup_timeout budget, tearing down each failed attempt cleanly, until the handshake succeeds or the budget elapses. A successful transport/session is transferred into the owner exit stack via a small _TransferredStack adapter. - Bound each attempt with asyncio.wait_for(initialize, _HANDSHAKE_ATTEMPT_TIMEOUT_SEC=10s) so a cold-starting process fails fast and retries instead of hanging until the transport drops. - Stop retrying ONLY when the process has DEFINITIVELY exited: new _managed_process_has_exited() (checks EXITED status) replaces the previous not-_managed_process_is_running() test, which false-negatived on a just-spawned process that had not yet reported RUNNING and made the loop bail to the outer rebuild path (relay then rejected the early re-attach with HTTP 400). Adds a unit test that fails the first two handshakes with the process alive and asserts the loop retries to success while starting the process exactly once. Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
923 lines
32 KiB
Python
923 lines
32 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 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_refuses_when_box_unavailable(self, mcp_module):
|
|
"""Policy: when Box is configured but unavailable (disabled in config
|
|
OR connection failed), stdio MCP servers are NOT treated as box-stdio.
|
|
``_init_stdio_python_server`` will raise a clear refusal at start
|
|
time; until then, the runtime info simply omits box_session_id so the
|
|
UI can render the disabled state cleanly."""
|
|
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
|
|
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_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_retries_during_cold_start(mcp_module, tmp_path, monkeypatch):
|
|
"""A slow (npx-style) cold start: the first handshake attempts fail while
|
|
the process is still starting, but the managed process is NOT exited, so
|
|
the handshake retry loop keeps trying and succeeds once the process is
|
|
ready — rather than raising and letting the outer loop churn the process."""
|
|
import asyncio as _asyncio
|
|
from contextlib import asynccontextmanager
|
|
|
|
mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio']
|
|
|
|
calls = {'init': 0}
|
|
|
|
class FlakyClientSession:
|
|
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):
|
|
calls['init'] += 1
|
|
# Fail the first two attempts (process still cold-starting),
|
|
# succeed on the third.
|
|
if calls['init'] < 3:
|
|
raise Exception('Connection closed')
|
|
return None
|
|
|
|
@asynccontextmanager
|
|
async def fake_websocket_client(_url: str):
|
|
yield ('read-stream', 'write-stream')
|
|
|
|
monkeypatch.setattr(mcp_stdio_module, 'ClientSession', FlakyClientSession)
|
|
monkeypatch.setattr(mcp_stdio_module, 'websocket_client', fake_websocket_client)
|
|
# Make the retry loop fast.
|
|
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,
|
|
)
|
|
|
|
# The process never reports 'exited' during the cold start.
|
|
async def _not_exited():
|
|
return False
|
|
|
|
session._box_stdio_runtime._managed_process_has_exited = _not_exited
|
|
# First check (is it already running for reuse?) -> False so we start it.
|
|
running_flags = iter([False])
|
|
|
|
async def _is_running():
|
|
try:
|
|
return next(running_flags)
|
|
except StopIteration:
|
|
return True
|
|
|
|
session._box_stdio_runtime._managed_process_is_running = _is_running
|
|
|
|
# Speed up the inter-attempt sleep.
|
|
orig_sleep = _asyncio.sleep
|
|
|
|
async def _fast_sleep(_s):
|
|
await orig_sleep(0)
|
|
|
|
monkeypatch.setattr(mcp_stdio_module.asyncio, 'sleep', _fast_sleep)
|
|
|
|
await session._init_box_stdio_server()
|
|
await session.exit_stack.aclose()
|
|
|
|
# Handshake was retried until success (3 initialize calls) and the process
|
|
# was started exactly once (never rebuilt).
|
|
assert calls['init'] == 3
|
|
assert ap.box_service.start_managed_process.await_count == 1
|