fix: improve MCP compatibility with persistent sessions and server callbacks
This commit is contained in:
parent
e3e2d22dc4
commit
dd334549c7
@ -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`)
|
||||||
|
|
||||||
内置了一个原生工具:
|
内置了一个原生工具:
|
||||||
|
|
||||||
|
|||||||
@ -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
|
|
||||||
|
|||||||
@ -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()
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user