258 lines
10 KiB
Python
258 lines
10 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import sys
|
||
import tempfile
|
||
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 tearDown(self):
|
||
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)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|