fix: improve MCP compatibility with persistent sessions and server callbacks

This commit is contained in:
JOJO 2026-04-27 11:34:58 +08:00
parent e3e2d22dc4
commit dd334549c7
3 changed files with 629 additions and 42 deletions

View File

@ -15,6 +15,7 @@
2. 系统通过 `tools/list` 拉取远端工具 2. 系统通过 `tools/list` 拉取远端工具
3. 本地自动生成工具别名:`mcp__<server_id>__<tool_name>` 3. 本地自动生成工具别名:`mcp__<server_id>__<tool_name>`
4. 模型调用该别名时,桥接层转发到对应服务的 `tools/call` 4. 模型调用该别名时,桥接层转发到对应服务的 `tools/call`
5. 对于同一服务,桥接层会复用长连接会话(避免每次调用都重启 MCP 进程),并处理常见 Server→Client 请求(如 `ping`、`roots/list`
内置了一个原生工具: 内置了一个原生工具:

View File

@ -8,10 +8,11 @@ import os
import select import select
import shutil import shutil
import subprocess import subprocess
import threading
import time import time
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple, TYPE_CHECKING from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, TYPE_CHECKING
import httpx import httpx
@ -36,6 +37,18 @@ class MCPToolBinding:
timeout_seconds: int timeout_seconds: int
@dataclass
class MCPClientPoolEntry:
"""长连接 MCP 客户端池条目。"""
server_id: str
signature: str
client: Any
timeout_seconds: int
created_at: float
last_used_at: float
class _StdioMCPClient: class _StdioMCPClient:
"""最小可用 stdio MCP 客户端JSON-RPC over newline-delimited JSON""" """最小可用 stdio MCP 客户端JSON-RPC over newline-delimited JSON"""
@ -75,6 +88,9 @@ class _StdioMCPClient:
self.container_session = container_session self.container_session = container_session
self.process: Optional[subprocess.Popen] = None self.process: Optional[subprocess.Popen] = None
self._seq = 0 self._seq = 0
self._initialized = False
self._initialize_result: Dict[str, Any] = {}
self._negotiated_protocol_version: Optional[str] = None
def _next_id(self) -> int: def _next_id(self) -> int:
self._seq += 1 self._seq += 1
@ -248,6 +264,8 @@ class _StdioMCPClient:
return cmd, None, dict(os.environ) return cmd, None, dict(os.environ)
def start(self) -> None: def start(self) -> None:
if self.process and self.process.poll() is None:
return
launch_cmd, cwd, env = self._prepare_launch() launch_cmd, cwd, env = self._prepare_launch()
try: try:
self.process = subprocess.Popen( self.process = subprocess.Popen(
@ -262,6 +280,9 @@ class _StdioMCPClient:
) )
except Exception as exc: except Exception as exc:
raise MCPClientError(f"启动 stdio MCP 服务失败: {exc}") from exc raise MCPClientError(f"启动 stdio MCP 服务失败: {exc}") from exc
self._initialized = False
self._initialize_result = {}
self._negotiated_protocol_version = None
def _send(self, payload: Dict[str, Any]) -> None: def _send(self, payload: Dict[str, Any]) -> None:
if not self.process or not self.process.stdin: if not self.process or not self.process.stdin:
@ -316,41 +337,101 @@ class _StdioMCPClient:
last_text_line = line last_text_line = line
continue continue
@staticmethod
def _build_server_request_result(method: str, params: Dict[str, Any]) -> Tuple[bool, Dict[str, Any]]:
"""处理服务端发起的请求Server -> Client Request"""
name = str(method or "").strip()
if name == "ping":
return True, {}
if name == "roots/list":
return True, {"roots": []}
if name == "completion/complete":
return False, {"code": -32601, "message": "completion/complete not supported by client"}
if name == "sampling/createMessage":
return False, {"code": -32601, "message": "sampling not supported by client"}
if name == "elicitation/create":
return False, {"code": -32601, "message": "elicitation not supported by client"}
return False, {"code": -32601, "message": f"Client method not implemented: {name}"}
@staticmethod
def _handle_server_notification(method: str, params: Dict[str, Any]) -> None:
"""处理服务端通知Server -> Client Notification。当前策略静默接收。"""
_ = method, params
def _respond_server_request(self, msg: Dict[str, Any]) -> None:
msg_id = msg.get("id")
method = str(msg.get("method") or "").strip()
params = msg.get("params")
if not isinstance(params, dict):
params = {}
try:
ok, payload = self._build_server_request_result(method, params)
if ok:
self._send({"jsonrpc": "2.0", "id": msg_id, "result": payload})
else:
self._send({"jsonrpc": "2.0", "id": msg_id, "error": payload})
except Exception as exc:
self._send(
{
"jsonrpc": "2.0",
"id": msg_id,
"error": {"code": -32603, "message": f"Client internal error: {exc}"},
}
)
def request(self, method: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: def request(self, method: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
if not self.process or self.process.poll() is not None:
self.start()
req_id = self._next_id() req_id = self._next_id()
payload = {"jsonrpc": "2.0", "id": req_id, "method": method, "params": params or {}} payload = {"jsonrpc": "2.0", "id": req_id, "method": method, "params": params or {}}
self._send(payload) self._send(payload)
while True: while True:
msg = self._read_message() msg = self._read_message()
if msg.get("id") == req_id:
# 收到目标响应
if msg.get("id") == req_id and "method" not in msg:
if "error" in msg: if "error" in msg:
error = msg.get("error") or {} error = msg.get("error") or {}
raise MCPClientError(f"{method} 调用失败: {error}") raise MCPClientError(f"{method} 调用失败: {error}")
return msg.get("result") or {} return msg.get("result") or {}
# 服务器主动向客户端发起请求时,返回 method not found避免阻塞
# 服务端请求:回包后继续等待目标响应
if "id" in msg and "method" in msg: if "id" in msg and "method" in msg:
self._send( self._respond_server_request(msg)
{ continue
"jsonrpc": "2.0",
"id": msg.get("id"), # 服务端通知:消费后继续
"error": {"code": -32601, "message": "Client method not implemented"}, if "method" in msg:
} self._handle_server_notification(str(msg.get("method") or ""), msg.get("params") or {})
) continue
def notify(self, method: str, params: Optional[Dict[str, Any]] = None) -> None: def notify(self, method: str, params: Optional[Dict[str, Any]] = None) -> None:
if not self.process or self.process.poll() is not None:
self.start()
self._send({"jsonrpc": "2.0", "method": method, "params": params or {}}) self._send({"jsonrpc": "2.0", "method": method, "params": params or {}})
def initialize(self) -> Dict[str, Any]: def initialize(self) -> Dict[str, Any]:
if self._initialized:
return dict(self._initialize_result)
if not self.process or self.process.poll() is not None:
self.start()
result = self.request( result = self.request(
"initialize", "initialize",
{ {
"protocolVersion": self.protocol_version, "protocolVersion": self.protocol_version,
"capabilities": {}, "capabilities": {
"roots": {"listChanged": False},
},
"clientInfo": {"name": "agents-mcp-bridge", "version": "1.0.0"}, "clientInfo": {"name": "agents-mcp-bridge", "version": "1.0.0"},
}, },
) )
self.notify("notifications/initialized") self.notify("notifications/initialized")
return result self._initialize_result = dict(result or {})
self._initialized = True
negotiated = result.get("protocolVersion") if isinstance(result, dict) else None
if negotiated:
self._negotiated_protocol_version = str(negotiated)
return dict(self._initialize_result)
def close(self) -> None: def close(self) -> None:
if not self.process: if not self.process:
@ -381,6 +462,9 @@ class _StdioMCPClient:
except Exception: except Exception:
pass pass
self.process = None self.process = None
self._initialized = False
self._initialize_result = {}
self._negotiated_protocol_version = None
class _StreamableHTTPMCPClient: class _StreamableHTTPMCPClient:
@ -394,6 +478,9 @@ class _StreamableHTTPMCPClient:
self._session_id: Optional[str] = None self._session_id: Optional[str] = None
self._seq = 0 self._seq = 0
self._client = httpx.Client(timeout=self.timeout_seconds) self._client = httpx.Client(timeout=self.timeout_seconds)
self._initialized = False
self._initialize_result: Dict[str, Any] = {}
self._negotiated_protocol_version: Optional[str] = None
def _next_id(self) -> int: def _next_id(self) -> int:
self._seq += 1 self._seq += 1
@ -406,14 +493,35 @@ class _StreamableHTTPMCPClient:
return hv return hv
return None return None
def _build_headers(self) -> Dict[str, str]: @staticmethod
def _infer_mcp_name(method: str, params: Optional[Dict[str, Any]]) -> Optional[str]:
if not isinstance(params, dict):
params = {}
method = str(method or "").strip()
if method == "tools/call":
value = params.get("name")
return str(value).strip() if value is not None else None
if method == "resources/read":
value = params.get("uri")
return str(value).strip() if value is not None else None
if method == "prompts/get":
value = params.get("name")
return str(value).strip() if value is not None else None
return None
def _build_headers(self, *, method: Optional[str] = None, params: Optional[Dict[str, Any]] = None) -> Dict[str, str]:
headers = { headers = {
"Accept": "application/json, text/event-stream", "Accept": "application/json, text/event-stream",
"Content-Type": "application/json", "Content-Type": "application/json",
"MCP-Protocol-Version": self.protocol_version, "MCP-Protocol-Version": self._negotiated_protocol_version or self.protocol_version,
} }
if self._session_id: if self._session_id:
headers["Mcp-Session-Id"] = self._session_id headers["Mcp-Session-Id"] = self._session_id
if method:
headers["Mcp-Method"] = str(method)
name = self._infer_mcp_name(str(method), params)
if name:
headers["Mcp-Name"] = name
headers.update(self.server.get("headers") or {}) headers.update(self.server.get("headers") or {})
return headers return headers
@ -439,7 +547,62 @@ class _StreamableHTTPMCPClient:
events.append(parsed) events.append(parsed)
return events return events
def _request(self, method: str, params: Optional[Dict[str, Any]] = None, notification: bool = False) -> Dict[str, Any]: @staticmethod
def _build_server_request_result(method: str, params: Dict[str, Any]) -> Tuple[bool, Dict[str, Any]]:
name = str(method or "").strip()
if name == "ping":
return True, {}
if name == "roots/list":
return True, {"roots": []}
if name == "completion/complete":
return False, {"code": -32601, "message": "completion/complete not supported by client"}
if name == "sampling/createMessage":
return False, {"code": -32601, "message": "sampling not supported by client"}
if name == "elicitation/create":
return False, {"code": -32601, "message": "elicitation not supported by client"}
return False, {"code": -32601, "message": f"Client method not implemented: {name}"}
def _post_jsonrpc_message(self, payload: Dict[str, Any]) -> None:
try:
resp = self._client.post(
self.url,
headers=self._build_headers(),
json=payload,
)
except Exception as exc:
raise MCPClientError(f"HTTP MCP 回写失败: {exc}") from exc
if resp.status_code >= 400:
raise MCPClientError(f"HTTP MCP 回写状态异常: {resp.status_code} {resp.text[:400]}")
def _handle_server_request_message(self, message: Dict[str, Any]) -> None:
msg_id = message.get("id")
method = str(message.get("method") or "").strip()
params = message.get("params")
if not isinstance(params, dict):
params = {}
try:
ok, payload = self._build_server_request_result(method, params)
if ok:
self._post_jsonrpc_message({"jsonrpc": "2.0", "id": msg_id, "result": payload})
else:
self._post_jsonrpc_message({"jsonrpc": "2.0", "id": msg_id, "error": payload})
except Exception as exc:
self._post_jsonrpc_message(
{
"jsonrpc": "2.0",
"id": msg_id,
"error": {"code": -32603, "message": f"Client internal error: {exc}"},
}
)
def _request(
self,
method: str,
params: Optional[Dict[str, Any]] = None,
notification: bool = False,
*,
_allow_session_retry: bool = True,
) -> Dict[str, Any]:
if not self.url: if not self.url:
raise MCPClientError("streamable_http 服务缺少 url") raise MCPClientError("streamable_http 服务缺少 url")
req_id = None if notification else self._next_id() req_id = None if notification else self._next_id()
@ -447,7 +610,11 @@ class _StreamableHTTPMCPClient:
if req_id is not None: if req_id is not None:
payload["id"] = req_id payload["id"] = req_id
try: try:
resp = self._client.post(self.url, headers=self._build_headers(), json=payload) resp = self._client.post(
self.url,
headers=self._build_headers(method=method, params=params or {}),
json=payload,
)
except Exception as exc: except Exception as exc:
raise MCPClientError(f"HTTP MCP 请求失败: {exc}") from exc raise MCPClientError(f"HTTP MCP 请求失败: {exc}") from exc
@ -455,6 +622,25 @@ class _StreamableHTTPMCPClient:
if session_id: if session_id:
self._session_id = session_id self._session_id = session_id
# 会话失效(服务端返回 404按规范重新建立 session 并重试一次
if (
resp.status_code == 404
and self._session_id
and method != "initialize"
and _allow_session_retry
):
self._session_id = None
self._initialized = False
self._initialize_result = {}
self._negotiated_protocol_version = None
self.initialize()
return self._request(
method,
params,
notification=notification,
_allow_session_retry=False,
)
if resp.status_code >= 400: if resp.status_code >= 400:
raise MCPClientError(f"HTTP MCP 状态异常: {resp.status_code} {resp.text[:400]}") raise MCPClientError(f"HTTP MCP 状态异常: {resp.status_code} {resp.text[:400]}")
@ -469,6 +655,13 @@ class _StreamableHTTPMCPClient:
message = parsed message = parsed
elif "text/event-stream" in content_type: elif "text/event-stream" in content_type:
for event in self._parse_sse_json_events(resp.text): for event in self._parse_sse_json_events(resp.text):
if not isinstance(event, dict):
continue
if "method" in event:
if "id" in event:
self._handle_server_request_message(event)
# 通知直接消费
continue
if event.get("id") == req_id: if event.get("id") == req_id:
message = event message = event
break break
@ -487,16 +680,25 @@ class _StreamableHTTPMCPClient:
return message.get("result") or {} return message.get("result") or {}
def initialize(self) -> Dict[str, Any]: def initialize(self) -> Dict[str, Any]:
if self._initialized:
return dict(self._initialize_result)
result = self._request( result = self._request(
"initialize", "initialize",
{ {
"protocolVersion": self.protocol_version, "protocolVersion": self.protocol_version,
"capabilities": {}, "capabilities": {
"roots": {"listChanged": False},
},
"clientInfo": {"name": "agents-mcp-bridge", "version": "1.0.0"}, "clientInfo": {"name": "agents-mcp-bridge", "version": "1.0.0"},
}, },
) )
self._request("notifications/initialized", {}, notification=True) self._request("notifications/initialized", {}, notification=True)
return result self._initialize_result = dict(result or {})
self._initialized = True
negotiated = result.get("protocolVersion") if isinstance(result, dict) else None
if negotiated:
self._negotiated_protocol_version = str(negotiated)
return dict(self._initialize_result)
def request(self, method: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: def request(self, method: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
return self._request(method, params, notification=False) return self._request(method, params, notification=False)
@ -511,6 +713,10 @@ class _StreamableHTTPMCPClient:
self._client.close() self._client.close()
except Exception: except Exception:
pass pass
self._session_id = None
self._initialized = False
self._initialize_result = {}
self._negotiated_protocol_version = None
class MCPClientManager: class MCPClientManager:
@ -526,10 +732,141 @@ class MCPClientManager:
self.protocol_version = str(protocol_version or MCP_PROTOCOL_VERSION) self.protocol_version = str(protocol_version or MCP_PROTOCOL_VERSION)
self.container_session = container_session self.container_session = container_session
self._latest_alias_map: Dict[str, MCPToolBinding] = {} self._latest_alias_map: Dict[str, MCPToolBinding] = {}
self._client_pool: Dict[str, MCPClientPoolEntry] = {}
self._pool_lock = threading.RLock()
def __del__(self):
try:
self.close_all_clients()
except Exception:
pass
def set_container_session(self, session: Optional["ContainerHandle"]) -> None: def set_container_session(self, session: Optional["ContainerHandle"]) -> None:
if session is self.container_session:
return
self.close_all_clients()
self.container_session = session self.container_session = session
@staticmethod
def _server_signature(server: Dict[str, Any]) -> str:
payload = {
"id": server.get("id"),
"transport": server.get("transport"),
"command": server.get("command"),
"args": server.get("args") or [],
"cwd": server.get("cwd"),
"env": server.get("env") or {},
"url": server.get("url"),
"headers": server.get("headers") or {},
"timeout_seconds": server.get("timeout_seconds"),
}
raw = json.dumps(payload, ensure_ascii=False, sort_keys=True)
return hashlib.sha1(raw.encode("utf-8")).hexdigest()
@staticmethod
def _is_client_alive(client: Any) -> bool:
process = getattr(client, "process", None)
if process is not None:
return bool(process.poll() is None)
# HTTP 客户端无法直接探活,按可用处理,失败时走重连兜底
return True
@staticmethod
def _safe_close_client(client: Any) -> None:
try:
if client:
client.close()
except Exception:
pass
def _close_pool_entry_locked(self, server_id: str) -> None:
entry = self._client_pool.pop(server_id, None)
if entry:
self._safe_close_client(entry.client)
def _ensure_client_locked(self, server: Dict[str, Any], timeout_seconds: int, *, force_recreate: bool = False):
sid = str(server.get("id") or "").strip()
if not sid:
raise MCPClientError("MCP 服务缺少 id")
signature = self._server_signature(server)
entry = self._client_pool.get(sid)
needs_recreate = force_recreate or (entry is None)
if entry and not needs_recreate:
if entry.signature != signature:
needs_recreate = True
elif not self._is_client_alive(entry.client):
needs_recreate = True
if needs_recreate:
if entry:
self._safe_close_client(entry.client)
client = self._make_client(server, timeout_seconds)
if hasattr(client, "start"):
client.start()
init_result = client.initialize()
now = time.time()
entry = MCPClientPoolEntry(
server_id=sid,
signature=signature,
client=client,
timeout_seconds=timeout_seconds,
created_at=now,
last_used_at=now,
)
self._client_pool[sid] = entry
return client, init_result
# 复用连接:动态刷新超时,并确保处于 initialized 状态
entry.timeout_seconds = timeout_seconds
try:
if hasattr(entry.client, "timeout_seconds"):
entry.client.timeout_seconds = max(3, int(timeout_seconds))
if hasattr(entry.client, "_client"):
entry.client._client.timeout = httpx.Timeout(max(3, int(timeout_seconds)))
except Exception:
pass
init_result = {}
if hasattr(entry.client, "initialize"):
init_result = entry.client.initialize() or {}
entry.last_used_at = time.time()
return entry.client, init_result
def _run_with_client(
self,
server: Dict[str, Any],
timeout_seconds: int,
operation: Callable[[Any], Any],
) -> Tuple[Any, Dict[str, Any]]:
last_exc: Optional[Exception] = None
for attempt in range(2):
with self._pool_lock:
client, init_result = self._ensure_client_locked(
server,
timeout_seconds,
force_recreate=attempt > 0,
)
sid = str(server.get("id") or "").strip()
try:
output = operation(client)
entry = self._client_pool.get(sid)
if entry:
entry.last_used_at = time.time()
return output, init_result
except Exception as exc: # pragma: no cover - 仅异常链路
last_exc = exc
self._close_pool_entry_locked(sid)
continue
if last_exc:
raise last_exc
raise MCPClientError("MCP 客户端执行失败")
def close_all_clients(self) -> None:
with self._pool_lock:
entries = list(self._client_pool.values())
self._client_pool.clear()
for entry in entries:
self._safe_close_client(entry.client)
@staticmethod @staticmethod
def _sanitize_input_schema(schema: Any) -> Dict[str, Any]: def _sanitize_input_schema(schema: Any) -> Dict[str, Any]:
if not isinstance(schema, dict): if not isinstance(schema, dict):
@ -625,12 +962,12 @@ class MCPClientManager:
return {"success": False, "error": f"未找到 MCP 服务: {server_id}"} return {"success": False, "error": f"未找到 MCP 服务: {server_id}"}
timeout = int(timeout_seconds or server.get("timeout_seconds") or MCP_DEFAULT_TIMEOUT_SECONDS) timeout = int(timeout_seconds or server.get("timeout_seconds") or MCP_DEFAULT_TIMEOUT_SECONDS)
started = time.time() started = time.time()
client = None
try: try:
client = self._make_client(server, timeout) tools, init_result = self._run_with_client(
client.start() if hasattr(client, "start") else None server,
init_result = client.initialize() timeout,
tools = self._paged_list_tools(client) lambda client: self._paged_list_tools(client),
)
self.registry.update_tools_cache(server_id, tools=tools, error="", touched=True) self.registry.update_tools_cache(server_id, tools=tools, error="", touched=True)
elapsed_ms = int((time.time() - started) * 1000) elapsed_ms = int((time.time() - started) * 1000)
return { return {
@ -644,17 +981,20 @@ class MCPClientManager:
except Exception as exc: except Exception as exc:
self.registry.update_tools_cache(server_id, error=str(exc), touched=True) self.registry.update_tools_cache(server_id, error=str(exc), touched=True)
return {"success": False, "server_id": server_id, "error": str(exc)} return {"success": False, "server_id": server_id, "error": str(exc)}
finally:
try:
if client:
client.close()
except Exception:
pass
def sync_servers(self, server_id: Optional[str] = None) -> Dict[str, Any]: def sync_servers(self, server_id: Optional[str] = None) -> Dict[str, Any]:
# 先从磁盘重载一次,避免「文件已改但进程内缓存未刷新」导致同步遗漏, # 先从磁盘重载一次,避免「文件已改但进程内缓存未刷新」导致同步遗漏,
# 以及后续 update_tools_cache 把旧缓存反写覆盖新配置。 # 以及后续 update_tools_cache 把旧缓存反写覆盖新配置。
self.registry.reload() self.registry.reload()
with self._pool_lock:
enabled_ids = {
str(item.get("id") or "").strip()
for item in self.registry.list_enabled_servers()
if str(item.get("id") or "").strip()
}
stale_ids = [sid for sid in list(self._client_pool.keys()) if sid not in enabled_ids]
for sid in stale_ids:
self._close_pool_entry_locked(sid)
if server_id: if server_id:
return self.discover_tools_for_server(server_id) return self.discover_tools_for_server(server_id)
results = [] results = []
@ -775,15 +1115,15 @@ class MCPClientManager:
return {"success": False, "error": f"MCP 服务已禁用: {binding.server_id}"} return {"success": False, "error": f"MCP 服务已禁用: {binding.server_id}"}
timeout = int(server.get("timeout_seconds") or binding.timeout_seconds or MCP_DEFAULT_TIMEOUT_SECONDS) timeout = int(server.get("timeout_seconds") or binding.timeout_seconds or MCP_DEFAULT_TIMEOUT_SECONDS)
client = None
started = time.time() started = time.time()
try: try:
client = self._make_client(server, timeout) call_result, _ = self._run_with_client(
client.start() if hasattr(client, "start") else None server,
client.initialize() timeout,
call_result = client.request( lambda client: client.request(
"tools/call", "tools/call",
{"name": binding.remote_name, "arguments": arguments or {}}, {"name": binding.remote_name, "arguments": arguments or {}},
),
) )
is_error = bool((call_result or {}).get("isError")) is_error = bool((call_result or {}).get("isError"))
content_items = (call_result or {}).get("content") or [] content_items = (call_result or {}).get("content") or []
@ -816,9 +1156,3 @@ class MCPClientManager:
"transport": binding.transport, "transport": binding.transport,
"error": str(exc), "error": str(exc),
} }
finally:
try:
if client:
client.close()
except Exception:
pass

View File

@ -4,6 +4,7 @@ import asyncio
import json import json
import sys import sys
import tempfile import tempfile
import textwrap
import unittest import unittest
from pathlib import Path from pathlib import Path
from types import SimpleNamespace from types import SimpleNamespace
@ -42,7 +43,16 @@ class MCPIntegrationTest(unittest.TestCase):
} }
) )
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): def tearDown(self):
try:
self.manager.close_all_clients()
except Exception:
pass
self.temp_dir.cleanup() self.temp_dir.cleanup()
def test_manager_discover_and_call(self): def test_manager_discover_and_call(self):
@ -252,6 +262,248 @@ class MCPIntegrationTest(unittest.TestCase):
self.assertNotIn("/opt/homebrew/bin/npx", launch_cmd) self.assertNotIn("/opt/homebrew/bin/npx", launch_cmd)
self.assertIn("/Users/jojo/Desktop", 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__": if __name__ == "__main__":
unittest.main() unittest.main()