510 lines
20 KiB
Python
510 lines
20 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import sys
|
||
import tempfile
|
||
import textwrap
|
||
import unittest
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
|
||
from core.main_terminal import MainTerminal
|
||
from modules.mcp_client_manager import MCPClientManager, _StdioMCPClient
|
||
from modules.mcp_server_registry import MCPServerRegistry
|
||
|
||
|
||
class MCPIntegrationTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.temp_dir = tempfile.TemporaryDirectory()
|
||
root = Path(self.temp_dir.name)
|
||
self.project_dir = root / "project"
|
||
self.project_dir.mkdir(parents=True, exist_ok=True)
|
||
self.data_dir = root / "data"
|
||
self.data_dir.mkdir(parents=True, exist_ok=True)
|
||
self.registry_path = self.data_dir / "mcp_servers.json"
|
||
|
||
self.registry = MCPServerRegistry(config_path=str(self.registry_path), enabled=True)
|
||
self.manager = MCPClientManager(self.registry)
|
||
|
||
demo_script = Path(__file__).resolve().parents[1] / "scripts" / "mcp_demo_server.py"
|
||
self.assertTrue(demo_script.exists(), f"missing demo script: {demo_script}")
|
||
self.demo_script = demo_script
|
||
|
||
self.registry.upsert_server(
|
||
{
|
||
"id": "demo_stdio",
|
||
"name": "Demo MCP",
|
||
"transport": "stdio",
|
||
"enabled": True,
|
||
"command": sys.executable,
|
||
"args": [str(self.demo_script)],
|
||
"timeout_seconds": 10,
|
||
}
|
||
)
|
||
|
||
def _write_temp_server(self, filename: str, content: str) -> Path:
|
||
path = self.project_dir / filename
|
||
path.write_text(textwrap.dedent(content).strip() + "\n", encoding="utf-8")
|
||
return path
|
||
|
||
def tearDown(self):
|
||
try:
|
||
self.manager.close_all_clients()
|
||
except Exception:
|
||
pass
|
||
self.temp_dir.cleanup()
|
||
|
||
def test_manager_discover_and_call(self):
|
||
sync_result = self.manager.sync_servers(server_id="demo_stdio")
|
||
self.assertTrue(sync_result.get("success"), sync_result)
|
||
self.assertGreaterEqual(sync_result.get("tools_count", 0), 2, sync_result)
|
||
|
||
tools, alias_map = self.manager.build_openai_tools(ensure_discovery=False)
|
||
self.assertTrue(tools)
|
||
|
||
echo_alias = None
|
||
for alias, binding in alias_map.items():
|
||
if binding.remote_name == "echo_upper":
|
||
echo_alias = alias
|
||
break
|
||
self.assertTrue(echo_alias, f"missing echo_upper alias: {list(alias_map.keys())}")
|
||
|
||
call_result = self.manager.call_tool_by_alias(
|
||
echo_alias,
|
||
{"text": "hello mcp"},
|
||
alias_map=alias_map,
|
||
)
|
||
self.assertTrue(call_result.get("success"), call_result)
|
||
payload = json.dumps(call_result, ensure_ascii=False)
|
||
self.assertIn("HELLO MCP", payload)
|
||
|
||
def test_main_terminal_tool_execution(self):
|
||
sync_result = self.manager.sync_servers(server_id="demo_stdio")
|
||
self.assertTrue(sync_result.get("success"), sync_result)
|
||
|
||
terminal = MainTerminal(
|
||
project_path=str(self.project_dir),
|
||
thinking_mode=False,
|
||
data_dir=str(self.data_dir),
|
||
)
|
||
terminal.mcp_tools_enabled = True
|
||
terminal.mcp_server_registry = self.registry
|
||
terminal.mcp_client_manager = self.manager
|
||
|
||
tool_defs = terminal.define_tools()
|
||
alias = None
|
||
for item in tool_defs:
|
||
name = ((item or {}).get("function") or {}).get("name")
|
||
if isinstance(name, str) and name.startswith("mcp__") and "echo_upper" in name:
|
||
alias = name
|
||
break
|
||
if not alias:
|
||
for item in tool_defs:
|
||
name = ((item or {}).get("function") or {}).get("name")
|
||
if isinstance(name, str) and name.startswith("mcp__"):
|
||
alias = name
|
||
break
|
||
self.assertTrue(alias, "no mcp alias in define_tools")
|
||
|
||
raw = asyncio.run(terminal.handle_tool_call(alias, {"text": "abc"}))
|
||
data = json.loads(raw)
|
||
self.assertTrue(data.get("success"), data)
|
||
self.assertIn("ABC", json.dumps(data, ensure_ascii=False))
|
||
|
||
# 原生 MCP 列表工具
|
||
list_raw = asyncio.run(terminal.handle_tool_call("list_mcp_servers", {"refresh": False}))
|
||
list_data = json.loads(list_raw)
|
||
self.assertTrue(list_data.get("success"), list_data)
|
||
self.assertGreaterEqual(int(list_data.get("count", 0)), 1, list_data)
|
||
|
||
def test_main_terminal_mcp_blocked_in_docker_mode(self):
|
||
terminal = MainTerminal(
|
||
project_path=str(self.project_dir),
|
||
thinking_mode=False,
|
||
data_dir=str(self.data_dir),
|
||
)
|
||
terminal.mcp_tools_enabled = True
|
||
terminal.mcp_server_registry = self.registry
|
||
terminal.mcp_client_manager = self.manager
|
||
terminal.container_session = SimpleNamespace(mode="docker")
|
||
|
||
message = "当前为docker模式，MCP仅支持宿主机模式"
|
||
|
||
list_raw = asyncio.run(terminal.handle_tool_call("list_mcp_servers", {"refresh": True}))
|
||
list_data = json.loads(list_raw)
|
||
self.assertFalse(list_data.get("success"), list_data)
|
||
self.assertEqual(list_data.get("error"), message, list_data)
|
||
|
||
call_raw = asyncio.run(terminal.handle_tool_call("mcp__demo_stdio__echo_upper", {"text": "abc"}))
|
||
call_data = json.loads(call_raw)
|
||
self.assertFalse(call_data.get("success"), call_data)
|
||
self.assertEqual(call_data.get("error"), message, call_data)
|
||
|
||
def test_calculator_mcp_server(self):
|
||
calculator_script = Path(__file__).resolve().parents[1] / "scripts" / "mcp_calculator_server.py"
|
||
self.assertTrue(calculator_script.exists(), f"missing calculator script: {calculator_script}")
|
||
self.registry.upsert_server(
|
||
{
|
||
"id": "calc_stdio",
|
||
"name": "Calculator MCP",
|
||
"transport": "stdio",
|
||
"enabled": True,
|
||
"command": sys.executable,
|
||
"args": [str(calculator_script)],
|
||
"timeout_seconds": 10,
|
||
}
|
||
)
|
||
|
||
sync_result = self.manager.sync_servers(server_id="calc_stdio")
|
||
self.assertTrue(sync_result.get("success"), sync_result)
|
||
self.assertEqual(sync_result.get("tools_count"), 1, sync_result)
|
||
|
||
_, alias_map = self.manager.build_openai_tools(ensure_discovery=False)
|
||
calc_alias = None
|
||
for alias, binding in alias_map.items():
|
||
if binding.server_id == "calc_stdio" and binding.remote_name == "calculator":
|
||
calc_alias = alias
|
||
break
|
||
self.assertTrue(calc_alias, f"missing calculator alias: {list(alias_map.keys())}")
|
||
|
||
call_result = self.manager.call_tool_by_alias(
|
||
calc_alias,
|
||
{"operation": "mul", "a": 6, "b": 7},
|
||
alias_map=alias_map,
|
||
)
|
||
self.assertTrue(call_result.get("success"), call_result)
|
||
self.assertIn("42", json.dumps(call_result, ensure_ascii=False))
|
||
|
||
def test_sync_servers_reload_external_file_changes(self):
|
||
# 模拟“外部工具直接写 mcp_servers.json”，绕过 registry 内存缓存
|
||
payload = {
|
||
"servers": [
|
||
{
|
||
"id": "demo_stdio",
|
||
"name": "Demo MCP",
|
||
"transport": "stdio",
|
||
"enabled": True,
|
||
"command": sys.executable,
|
||
"args": [str(self.demo_script)],
|
||
"timeout_seconds": 10,
|
||
},
|
||
{
|
||
"id": "extra_stdio",
|
||
"name": "Extra MCP",
|
||
"transport": "stdio",
|
||
"enabled": True,
|
||
"command": sys.executable,
|
||
"args": [str(self.demo_script)],
|
||
"timeout_seconds": 10,
|
||
},
|
||
]
|
||
}
|
||
self.registry_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
|
||
sync_result = self.manager.sync_servers()
|
||
self.assertTrue(sync_result.get("success"), sync_result)
|
||
self.assertEqual(sync_result.get("servers_total"), 2, sync_result)
|
||
|
||
by_id = {item.get("server_id"): item for item in sync_result.get("results") or []}
|
||
self.assertIn("extra_stdio", by_id, sync_result)
|
||
self.assertTrue(by_id["extra_stdio"].get("success"), sync_result)
|
||
|
||
def test_stdio_container_path_rewrite_from_workspace(self):
|
||
fake_session = SimpleNamespace(
|
||
mode="docker",
|
||
container_name="fake-container",
|
||
mount_path="/workspace",
|
||
workspace_path=str(self.project_dir),
|
||
sandbox_bin="docker",
|
||
)
|
||
server = {
|
||
"command": "python3",
|
||
"args": [
|
||
str(self.project_dir / "scripts" / "demo.py"),
|
||
f"--root={self.project_dir / 'subdir'}",
|
||
],
|
||
"cwd": str(self.project_dir),
|
||
"env": {"HELLO": "1"},
|
||
}
|
||
client = _StdioMCPClient(server, timeout_seconds=10, protocol_version="2025-06-18", container_session=fake_session)
|
||
launch_cmd, cwd, _ = client._prepare_launch()
|
||
|
||
self.assertTrue(str(launch_cmd[0]).endswith("docker"), launch_cmd)
|
||
self.assertEqual(launch_cmd[1:3], ["exec", "-i"], launch_cmd)
|
||
self.assertIn("-w", launch_cmd)
|
||
self.assertEqual(cwd, None)
|
||
joined = " ".join(launch_cmd)
|
||
self.assertIn("/workspace/scripts/demo.py", joined)
|
||
self.assertIn("--root=/workspace/subdir", joined)
|
||
self.assertIn("-e HELLO=1", joined)
|
||
|
||
def test_stdio_container_command_fallback_for_host_abs_path(self):
|
||
fake_session = SimpleNamespace(
|
||
mode="docker",
|
||
container_name="fake-container",
|
||
mount_path="/workspace",
|
||
workspace_path=str(self.project_dir),
|
||
sandbox_bin="docker",
|
||
)
|
||
server = {
|
||
"command": "/opt/homebrew/bin/npx",
|
||
"args": ["-y", "@modelcontextprotocol/server-filesystem", "/Users/jojo/Desktop"],
|
||
"cwd": "",
|
||
"env": {},
|
||
}
|
||
client = _StdioMCPClient(server, timeout_seconds=10, protocol_version="2025-06-18", container_session=fake_session)
|
||
launch_cmd, cwd, _ = client._prepare_launch()
|
||
|
||
self.assertEqual(cwd, None)
|
||
self.assertIn("fake-container", launch_cmd)
|
||
self.assertIn("npx", launch_cmd)
|
||
self.assertNotIn("/opt/homebrew/bin/npx", launch_cmd)
|
||
self.assertIn("/Users/jojo/Desktop", launch_cmd)
|
||
|
||
def test_stateful_stdio_server_reuses_persistent_session(self):
|
||
stateful_script = self._write_temp_server(
|
||
"stateful_mcp_server.py",
|
||
"""
|
||
import json
|
||
import sys
|
||
|
||
counter = 0
|
||
|
||
def send(payload):
|
||
sys.stdout.write(json.dumps(payload, ensure_ascii=False) + "\\n")
|
||
sys.stdout.flush()
|
||
|
||
while True:
|
||
line = sys.stdin.readline()
|
||
if not line:
|
||
break
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
req = json.loads(line)
|
||
method = req.get("method")
|
||
msg_id = req.get("id")
|
||
params = req.get("params") or {}
|
||
if msg_id is None:
|
||
continue
|
||
if method == "initialize":
|
||
send({
|
||
"jsonrpc": "2.0",
|
||
"id": msg_id,
|
||
"result": {
|
||
"protocolVersion": "2025-06-18",
|
||
"capabilities": {"tools": {}},
|
||
"serverInfo": {"name": "stateful", "version": "1.0.0"},
|
||
},
|
||
})
|
||
continue
|
||
if method == "tools/list":
|
||
send({
|
||
"jsonrpc": "2.0",
|
||
"id": msg_id,
|
||
"result": {
|
||
"tools": [
|
||
{
|
||
"name": "next_counter",
|
||
"description": "返回递增计数器",
|
||
"inputSchema": {"type": "object", "properties": {}},
|
||
}
|
||
]
|
||
},
|
||
})
|
||
continue
|
||
if method == "tools/call":
|
||
name = str(params.get("name") or "")
|
||
if name == "next_counter":
|
||
counter += 1
|
||
send({
|
||
"jsonrpc": "2.0",
|
||
"id": msg_id,
|
||
"result": {
|
||
"isError": False,
|
||
"content": [{"type": "text", "text": str(counter)}],
|
||
"structuredContent": {"counter": counter},
|
||
},
|
||
})
|
||
continue
|
||
send({
|
||
"jsonrpc": "2.0",
|
||
"id": msg_id,
|
||
"result": {
|
||
"isError": True,
|
||
"content": [{"type": "text", "text": "unknown tool"}],
|
||
},
|
||
})
|
||
continue
|
||
send({
|
||
"jsonrpc": "2.0",
|
||
"id": msg_id,
|
||
"error": {"code": -32601, "message": f"Method not found: {method}"},
|
||
})
|
||
""",
|
||
)
|
||
self.registry.upsert_server(
|
||
{
|
||
"id": "stateful_stdio",
|
||
"name": "Stateful MCP",
|
||
"transport": "stdio",
|
||
"enabled": True,
|
||
"command": sys.executable,
|
||
"args": [str(stateful_script)],
|
||
"timeout_seconds": 10,
|
||
}
|
||
)
|
||
|
||
sync_result = self.manager.sync_servers(server_id="stateful_stdio")
|
||
self.assertTrue(sync_result.get("success"), sync_result)
|
||
|
||
_, alias_map = self.manager.build_openai_tools(ensure_discovery=False)
|
||
alias = None
|
||
for key, binding in alias_map.items():
|
||
if binding.server_id == "stateful_stdio" and binding.remote_name == "next_counter":
|
||
alias = key
|
||
break
|
||
self.assertTrue(alias, f"missing alias for stateful server: {list(alias_map.keys())}")
|
||
|
||
first = self.manager.call_tool_by_alias(alias, {}, alias_map=alias_map)
|
||
second = self.manager.call_tool_by_alias(alias, {}, alias_map=alias_map)
|
||
|
||
self.assertTrue(first.get("success"), first)
|
||
self.assertTrue(second.get("success"), second)
|
||
self.assertEqual((first.get("structured_content") or {}).get("counter"), 1, first)
|
||
self.assertEqual((second.get("structured_content") or {}).get("counter"), 2, second)
|
||
|
||
def test_stdio_server_request_ping_is_handled(self):
|
||
ping_script = self._write_temp_server(
|
||
"ping_request_mcp_server.py",
|
||
"""
|
||
import json
|
||
import select
|
||
import sys
|
||
import time
|
||
|
||
def send(payload):
|
||
sys.stdout.write(json.dumps(payload, ensure_ascii=False) + "\\n")
|
||
sys.stdout.flush()
|
||
|
||
def read_json_line(timeout_seconds):
|
||
deadline = time.monotonic() + timeout_seconds
|
||
while time.monotonic() < deadline:
|
||
remain = max(0.05, deadline - time.monotonic())
|
||
readable, _, _ = select.select([sys.stdin], [], [], min(0.2, remain))
|
||
if not readable:
|
||
continue
|
||
line = sys.stdin.readline()
|
||
if not line:
|
||
return None
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
return json.loads(line)
|
||
except Exception:
|
||
continue
|
||
return None
|
||
|
||
while True:
|
||
line = sys.stdin.readline()
|
||
if not line:
|
||
break
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
req = json.loads(line)
|
||
method = req.get("method")
|
||
msg_id = req.get("id")
|
||
params = req.get("params") or {}
|
||
if msg_id is None:
|
||
continue
|
||
if method == "initialize":
|
||
send({
|
||
"jsonrpc": "2.0",
|
||
"id": msg_id,
|
||
"result": {
|
||
"protocolVersion": "2025-06-18",
|
||
"capabilities": {"tools": {}},
|
||
"serverInfo": {"name": "ping-stdio", "version": "1.0.0"},
|
||
},
|
||
})
|
||
continue
|
||
if method == "tools/list":
|
||
send({
|
||
"jsonrpc": "2.0",
|
||
"id": msg_id,
|
||
"result": {
|
||
"tools": [
|
||
{
|
||
"name": "trigger_ping",
|
||
"description": "触发服务端 ping 请求",
|
||
"inputSchema": {"type": "object", "properties": {}},
|
||
}
|
||
]
|
||
},
|
||
})
|
||
continue
|
||
if method == "tools/call":
|
||
ping_id = 9001
|
||
send({
|
||
"jsonrpc": "2.0",
|
||
"id": ping_id,
|
||
"method": "ping",
|
||
"params": {},
|
||
})
|
||
response = read_json_line(2.5) or {}
|
||
ok = (
|
||
isinstance(response, dict)
|
||
and response.get("id") == ping_id
|
||
and isinstance(response.get("result"), dict)
|
||
)
|
||
send({
|
||
"jsonrpc": "2.0",
|
||
"id": msg_id,
|
||
"result": {
|
||
"isError": not ok,
|
||
"content": [{"type": "text", "text": "ping_ok" if ok else "ping_failed"}],
|
||
"structuredContent": {"ping_ok": ok},
|
||
},
|
||
})
|
||
continue
|
||
send({
|
||
"jsonrpc": "2.0",
|
||
"id": msg_id,
|
||
"error": {"code": -32601, "message": f"Method not found: {method}"},
|
||
})
|
||
""",
|
||
)
|
||
self.registry.upsert_server(
|
||
{
|
||
"id": "ping_stdio",
|
||
"name": "Ping MCP",
|
||
"transport": "stdio",
|
||
"enabled": True,
|
||
"command": sys.executable,
|
||
"args": [str(ping_script)],
|
||
"timeout_seconds": 10,
|
||
}
|
||
)
|
||
|
||
sync_result = self.manager.sync_servers(server_id="ping_stdio")
|
||
self.assertTrue(sync_result.get("success"), sync_result)
|
||
|
||
_, alias_map = self.manager.build_openai_tools(ensure_discovery=False)
|
||
alias = None
|
||
for key, binding in alias_map.items():
|
||
if binding.server_id == "ping_stdio" and binding.remote_name == "trigger_ping":
|
||
alias = key
|
||
break
|
||
self.assertTrue(alias, f"missing alias for ping server: {list(alias_map.keys())}")
|
||
|
||
result = self.manager.call_tool_by_alias(alias, {}, alias_map=alias_map)
|
||
self.assertTrue(result.get("success"), result)
|
||
self.assertTrue((result.get("structured_content") or {}).get("ping_ok"), result)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|