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` 拉取远端工具
|
||||
3. 本地自动生成工具别名:`mcp__<server_id>__<tool_name>`
|
||||
4. 模型调用该别名时,桥接层转发到对应服务的 `tools/call`
|
||||
5. 对于同一服务,桥接层会复用长连接会话(避免每次调用都重启 MCP 进程),并处理常见 Server→Client 请求(如 `ping`、`roots/list`)
|
||||
|
||||
内置了一个原生工具:
|
||||
|
||||
|
||||
@ -8,10 +8,11 @@ import os
|
||||
import select
|
||||
import shutil
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
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
|
||||
|
||||
@ -36,6 +37,18 @@ class MCPToolBinding:
|
||||
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:
|
||||
"""最小可用 stdio MCP 客户端(JSON-RPC over newline-delimited JSON)。"""
|
||||
|
||||
@ -75,6 +88,9 @@ class _StdioMCPClient:
|
||||
self.container_session = container_session
|
||||
self.process: Optional[subprocess.Popen] = None
|
||||
self._seq = 0
|
||||
self._initialized = False
|
||||
self._initialize_result: Dict[str, Any] = {}
|
||||
self._negotiated_protocol_version: Optional[str] = None
|
||||
|
||||
def _next_id(self) -> int:
|
||||
self._seq += 1
|
||||
@ -248,6 +264,8 @@ class _StdioMCPClient:
|
||||
return cmd, None, dict(os.environ)
|
||||
|
||||
def start(self) -> None:
|
||||
if self.process and self.process.poll() is None:
|
||||
return
|
||||
launch_cmd, cwd, env = self._prepare_launch()
|
||||
try:
|
||||
self.process = subprocess.Popen(
|
||||
@ -262,6 +280,9 @@ class _StdioMCPClient:
|
||||
)
|
||||
except Exception as 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:
|
||||
if not self.process or not self.process.stdin:
|
||||
@ -316,41 +337,101 @@ class _StdioMCPClient:
|
||||
last_text_line = line
|
||||
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]:
|
||||
if not self.process or self.process.poll() is not None:
|
||||
self.start()
|
||||
req_id = self._next_id()
|
||||
payload = {"jsonrpc": "2.0", "id": req_id, "method": method, "params": params or {}}
|
||||
self._send(payload)
|
||||
while True:
|
||||
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:
|
||||
error = msg.get("error") or {}
|
||||
raise MCPClientError(f"{method} 调用失败: {error}")
|
||||
return msg.get("result") or {}
|
||||
# 服务器主动向客户端发起请求时,返回 method not found,避免阻塞
|
||||
|
||||
# 服务端请求:回包后继续等待目标响应
|
||||
if "id" in msg and "method" in msg:
|
||||
self._send(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": msg.get("id"),
|
||||
"error": {"code": -32601, "message": "Client method not implemented"},
|
||||
}
|
||||
)
|
||||
self._respond_server_request(msg)
|
||||
continue
|
||||
|
||||
# 服务端通知:消费后继续
|
||||
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:
|
||||
if not self.process or self.process.poll() is not None:
|
||||
self.start()
|
||||
self._send({"jsonrpc": "2.0", "method": method, "params": params or {}})
|
||||
|
||||
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(
|
||||
"initialize",
|
||||
{
|
||||
"protocolVersion": self.protocol_version,
|
||||
"capabilities": {},
|
||||
"capabilities": {
|
||||
"roots": {"listChanged": False},
|
||||
},
|
||||
"clientInfo": {"name": "agents-mcp-bridge", "version": "1.0.0"},
|
||||
},
|
||||
)
|
||||
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:
|
||||
if not self.process:
|
||||
@ -381,6 +462,9 @@ class _StdioMCPClient:
|
||||
except Exception:
|
||||
pass
|
||||
self.process = None
|
||||
self._initialized = False
|
||||
self._initialize_result = {}
|
||||
self._negotiated_protocol_version = None
|
||||
|
||||
|
||||
class _StreamableHTTPMCPClient:
|
||||
@ -394,6 +478,9 @@ class _StreamableHTTPMCPClient:
|
||||
self._session_id: Optional[str] = None
|
||||
self._seq = 0
|
||||
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:
|
||||
self._seq += 1
|
||||
@ -406,14 +493,35 @@ class _StreamableHTTPMCPClient:
|
||||
return hv
|
||||
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 = {
|
||||
"Accept": "application/json, text/event-stream",
|
||||
"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:
|
||||
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 {})
|
||||
return headers
|
||||
|
||||
@ -439,7 +547,62 @@ class _StreamableHTTPMCPClient:
|
||||
events.append(parsed)
|
||||
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:
|
||||
raise MCPClientError("streamable_http 服务缺少 url")
|
||||
req_id = None if notification else self._next_id()
|
||||
@ -447,7 +610,11 @@ class _StreamableHTTPMCPClient:
|
||||
if req_id is not None:
|
||||
payload["id"] = req_id
|
||||
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:
|
||||
raise MCPClientError(f"HTTP MCP 请求失败: {exc}") from exc
|
||||
|
||||
@ -455,6 +622,25 @@ class _StreamableHTTPMCPClient:
|
||||
if 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:
|
||||
raise MCPClientError(f"HTTP MCP 状态异常: {resp.status_code} {resp.text[:400]}")
|
||||
|
||||
@ -469,6 +655,13 @@ class _StreamableHTTPMCPClient:
|
||||
message = parsed
|
||||
elif "text/event-stream" in content_type:
|
||||
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:
|
||||
message = event
|
||||
break
|
||||
@ -487,16 +680,25 @@ class _StreamableHTTPMCPClient:
|
||||
return message.get("result") or {}
|
||||
|
||||
def initialize(self) -> Dict[str, Any]:
|
||||
if self._initialized:
|
||||
return dict(self._initialize_result)
|
||||
result = self._request(
|
||||
"initialize",
|
||||
{
|
||||
"protocolVersion": self.protocol_version,
|
||||
"capabilities": {},
|
||||
"capabilities": {
|
||||
"roots": {"listChanged": False},
|
||||
},
|
||||
"clientInfo": {"name": "agents-mcp-bridge", "version": "1.0.0"},
|
||||
},
|
||||
)
|
||||
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]:
|
||||
return self._request(method, params, notification=False)
|
||||
@ -511,6 +713,10 @@ class _StreamableHTTPMCPClient:
|
||||
self._client.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._session_id = None
|
||||
self._initialized = False
|
||||
self._initialize_result = {}
|
||||
self._negotiated_protocol_version = None
|
||||
|
||||
|
||||
class MCPClientManager:
|
||||
@ -526,10 +732,141 @@ class MCPClientManager:
|
||||
self.protocol_version = str(protocol_version or MCP_PROTOCOL_VERSION)
|
||||
self.container_session = container_session
|
||||
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:
|
||||
if session is self.container_session:
|
||||
return
|
||||
self.close_all_clients()
|
||||
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
|
||||
def _sanitize_input_schema(schema: Any) -> Dict[str, Any]:
|
||||
if not isinstance(schema, dict):
|
||||
@ -625,12 +962,12 @@ class MCPClientManager:
|
||||
return {"success": False, "error": f"未找到 MCP 服务: {server_id}"}
|
||||
timeout = int(timeout_seconds or server.get("timeout_seconds") or MCP_DEFAULT_TIMEOUT_SECONDS)
|
||||
started = time.time()
|
||||
client = None
|
||||
try:
|
||||
client = self._make_client(server, timeout)
|
||||
client.start() if hasattr(client, "start") else None
|
||||
init_result = client.initialize()
|
||||
tools = self._paged_list_tools(client)
|
||||
tools, init_result = self._run_with_client(
|
||||
server,
|
||||
timeout,
|
||||
lambda client: self._paged_list_tools(client),
|
||||
)
|
||||
self.registry.update_tools_cache(server_id, tools=tools, error="", touched=True)
|
||||
elapsed_ms = int((time.time() - started) * 1000)
|
||||
return {
|
||||
@ -644,17 +981,20 @@ class MCPClientManager:
|
||||
except Exception as exc:
|
||||
self.registry.update_tools_cache(server_id, error=str(exc), touched=True)
|
||||
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]:
|
||||
# 先从磁盘重载一次,避免「文件已改但进程内缓存未刷新」导致同步遗漏,
|
||||
# 以及后续 update_tools_cache 把旧缓存反写覆盖新配置。
|
||||
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:
|
||||
return self.discover_tools_for_server(server_id)
|
||||
results = []
|
||||
@ -775,15 +1115,15 @@ class MCPClientManager:
|
||||
return {"success": False, "error": f"MCP 服务已禁用: {binding.server_id}"}
|
||||
|
||||
timeout = int(server.get("timeout_seconds") or binding.timeout_seconds or MCP_DEFAULT_TIMEOUT_SECONDS)
|
||||
client = None
|
||||
started = time.time()
|
||||
try:
|
||||
client = self._make_client(server, timeout)
|
||||
client.start() if hasattr(client, "start") else None
|
||||
client.initialize()
|
||||
call_result = client.request(
|
||||
"tools/call",
|
||||
{"name": binding.remote_name, "arguments": arguments or {}},
|
||||
call_result, _ = self._run_with_client(
|
||||
server,
|
||||
timeout,
|
||||
lambda client: client.request(
|
||||
"tools/call",
|
||||
{"name": binding.remote_name, "arguments": arguments or {}},
|
||||
),
|
||||
)
|
||||
is_error = bool((call_result or {}).get("isError"))
|
||||
content_items = (call_result or {}).get("content") or []
|
||||
@ -816,9 +1156,3 @@ class MCPClientManager:
|
||||
"transport": binding.transport,
|
||||
"error": str(exc),
|
||||
}
|
||||
finally:
|
||||
try:
|
||||
if client:
|
||||
client.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@ -4,6 +4,7 @@ import asyncio
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
import textwrap
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
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):
|
||||
try:
|
||||
self.manager.close_all_clients()
|
||||
except Exception:
|
||||
pass
|
||||
self.temp_dir.cleanup()
|
||||
|
||||
def test_manager_discover_and_call(self):
|
||||
@ -252,6 +262,248 @@ class MCPIntegrationTest(unittest.TestCase):
|
||||
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()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user