agent-Specialization/test/test_mcp_integration.py

784 lines
31 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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_call_drops_intent_argument(self):
strict_script = self._write_temp_server(
"strict_args_mcp_server.py",
"""
import json
import sys
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": "strict-args", "version": "1.0.0"},
},
})
continue
if method == "tools/list":
send({
"jsonrpc": "2.0",
"id": msg_id,
"result": {
"tools": [
{
"name": "echo_strict",
"description": "仅接受 text 参数",
"inputSchema": {
"type": "object",
"properties": {"text": {"type": "string"}},
"required": ["text"],
},
}
]
},
})
continue
if method == "tools/call":
name = str(params.get("name") or "")
arguments = params.get("arguments") if isinstance(params.get("arguments"), dict) else {}
keys = sorted(arguments.keys())
if name != "echo_strict":
send({
"jsonrpc": "2.0",
"id": msg_id,
"result": {
"isError": True,
"content": [{"type": "text", "text": "unknown tool"}],
},
})
continue
if keys != ["text"]:
send({
"jsonrpc": "2.0",
"id": msg_id,
"result": {
"isError": True,
"content": [{"type": "text", "text": "unexpected keys:" + ",".join(keys)}],
"structuredContent": {"received_keys": keys},
},
})
continue
text = str(arguments.get("text") or "")
send({
"jsonrpc": "2.0",
"id": msg_id,
"result": {
"isError": False,
"content": [{"type": "text", "text": text.upper()}],
"structuredContent": {"received_keys": keys, "text": text.upper()},
},
})
continue
send({
"jsonrpc": "2.0",
"id": msg_id,
"error": {"code": -32601, "message": f"Method not found: {method}"},
})
""",
)
self.registry.upsert_server(
{
"id": "strict_stdio",
"name": "Strict MCP",
"transport": "stdio",
"enabled": True,
"command": sys.executable,
"args": [str(strict_script)],
"timeout_seconds": 10,
}
)
sync_result = self.manager.sync_servers(server_id="strict_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
alias = None
for item in terminal.define_tools():
name = ((item or {}).get("function") or {}).get("name")
if isinstance(name, str) and name.startswith("mcp__strict_stdio__"):
alias = name
break
self.assertTrue(alias, "missing strict mcp alias")
raw = asyncio.run(terminal.handle_tool_call(alias, {"intent": "测试调用", "text": "abc"}))
data = json.loads(raw)
self.assertTrue(data.get("success"), data)
self.assertIn("ABC", json.dumps(data, ensure_ascii=False))
self.assertEqual((data.get("structured_content") or {}).get("received_keys"), ["text"], data)
def test_define_tools_skip_intent_in_mcp_alias(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
terminal.tool_intent_enabled = True
tool_defs = terminal.define_tools()
mcp_tool = None
native_tool = None
for item in tool_defs:
fn = (item or {}).get("function") or {}
name = str(fn.get("name") or "")
if not native_tool and name == "run_command":
native_tool = item
if not mcp_tool and name.startswith("mcp__demo_stdio__"):
mcp_tool = item
if mcp_tool and native_tool:
break
self.assertIsNotNone(native_tool, "missing native run_command tool")
self.assertIsNotNone(mcp_tool, "missing mcp alias tool")
native_params = ((native_tool or {}).get("function") or {}).get("parameters") or {}
native_props = native_params.get("properties") if isinstance(native_params, dict) else {}
self.assertIn("intent", native_props or {})
mcp_params = ((mcp_tool or {}).get("function") or {}).get("parameters") or {}
mcp_props = mcp_params.get("properties") if isinstance(mcp_params, dict) else {}
self.assertNotIn("intent", mcp_props or {})
self.assertNotIn("intent", mcp_params.get("required") or [])
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_build_openai_tools_reload_external_file_changes(self):
# 模拟外部直接改配置文件registry._cache 仍是旧内容)
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")
tools, alias_map = self.manager.build_openai_tools(
ensure_discovery=True,
discovery_timeout_seconds=10,
)
self.assertTrue(tools)
self.assertTrue(alias_map)
self.assertTrue(
any(getattr(binding, "server_id", "") == "extra_stdio" for binding in alias_map.values()),
alias_map,
)
# 关键断言:外部新增服务不应被旧缓存反写覆盖
after = json.loads(self.registry_path.read_text(encoding="utf-8"))
ids = {str((item or {}).get("id") or "") for item in (after.get("servers") or [])}
self.assertIn("demo_stdio", ids)
self.assertIn("extra_stdio", ids)
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_stdio_host_mode_uses_workspace_as_default_cwd(self):
workspace = (self.project_dir / "workspace_a").resolve()
workspace.mkdir(parents=True, exist_ok=True)
fake_session = SimpleNamespace(
mode="host",
workspace_path=str(workspace),
mount_path=str(workspace),
container_name=None,
sandbox_bin="docker",
)
server = {
"command": sys.executable,
"args": ["-V"],
"cwd": "",
"env": {},
}
client = _StdioMCPClient(server, timeout_seconds=10, protocol_version="2025-06-18", container_session=fake_session)
_, cwd, _ = client._prepare_launch()
self.assertEqual(cwd, str(workspace))
def test_set_container_session_detects_inplace_workspace_change(self):
manager = MCPClientManager(self.registry)
close_calls = []
original_close = manager.close_all_clients
manager.close_all_clients = lambda: close_calls.append("closed") # type: ignore[assignment]
try:
session = SimpleNamespace(
mode="host",
workspace_path=str((self.project_dir / "ws_a").resolve()),
mount_path="/workspace",
container_name="",
sandbox_bin="docker",
)
manager.set_container_session(session)
self.assertEqual(len(close_calls), 1)
session.workspace_path = str((self.project_dir / "ws_b").resolve())
manager.set_container_session(session)
self.assertEqual(len(close_calls), 2)
manager.set_container_session(session)
self.assertEqual(len(close_calls), 2)
finally:
manager.close_all_clients = original_close # type: ignore[assignment]
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)
def test_build_call_message_empty_content_returns_empty_string(self):
message = MCPClientManager._build_call_message(
{
"content": [
{
"type": "text",
"text": "",
}
]
}
)
self.assertEqual(message, "")
if __name__ == "__main__":
unittest.main()