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()