mirror of
https://github.com/zhayujie/chatgpt-on-wechat.git
synced 2026-07-19 12:47:25 +08:00
fix(mcp): enable concurrent calls for SSE and streamable-http transports
_stdio_send (single pipe) must remain serialized under _call_lock, but SSE and streamable-http use independent HTTP requests and can safely execute concurrently across sessions. - Scope _call_lock to stdio transport only - Add _http_lock with double-checked pattern to protect _http_session_id initialization during concurrent streamable-http requests
This commit is contained in:
@@ -59,7 +59,13 @@ class McpClient:
|
|||||||
# Shared state
|
# Shared state
|
||||||
self._next_id = 1
|
self._next_id = 1
|
||||||
self._id_lock = threading.Lock()
|
self._id_lock = threading.Lock()
|
||||||
|
# _call_lock serializes all requests on the single stdio pipe.
|
||||||
|
# SSE and streamable-http use independent HTTP requests, so they
|
||||||
|
# do not acquire this lock (see _send_request).
|
||||||
self._call_lock = threading.Lock()
|
self._call_lock = threading.Lock()
|
||||||
|
# _http_lock protects _http_session_id initialization across
|
||||||
|
# concurrent streamable-http requests.
|
||||||
|
self._http_lock = threading.Lock()
|
||||||
self._initialized = False
|
self._initialized = False
|
||||||
|
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
@@ -338,8 +344,12 @@ class McpClient:
|
|||||||
"Content-Type": "application/json",
|
"Content-Type": "application/json",
|
||||||
"Accept": "application/json, text/event-stream",
|
"Accept": "application/json, text/event-stream",
|
||||||
}
|
}
|
||||||
if self._http_session_id:
|
# Read session id under lock to avoid racing with the
|
||||||
headers["Mcp-Session-Id"] = self._http_session_id
|
# initialization write below during concurrent requests.
|
||||||
|
with self._http_lock:
|
||||||
|
sid = self._http_session_id
|
||||||
|
if sid:
|
||||||
|
headers["Mcp-Session-Id"] = sid
|
||||||
headers.update(self._http_headers)
|
headers.update(self._http_headers)
|
||||||
|
|
||||||
req = urllib.request.Request(
|
req = urllib.request.Request(
|
||||||
@@ -365,8 +375,13 @@ class McpClient:
|
|||||||
with resp:
|
with resp:
|
||||||
# Capture session id assigned by the server (if any)
|
# Capture session id assigned by the server (if any)
|
||||||
session_id = resp.headers.get("Mcp-Session-Id")
|
session_id = resp.headers.get("Mcp-Session-Id")
|
||||||
|
# Double-checked lock: only the first response sets the
|
||||||
|
# session id, preventing concurrent initializers from
|
||||||
|
# overwriting each other.
|
||||||
if session_id and not self._http_session_id:
|
if session_id and not self._http_session_id:
|
||||||
self._http_session_id = session_id
|
with self._http_lock:
|
||||||
|
if not self._http_session_id:
|
||||||
|
self._http_session_id = session_id
|
||||||
|
|
||||||
status = resp.status if hasattr(resp, "status") else resp.getcode()
|
status = resp.status if hasattr(resp, "status") else resp.getcode()
|
||||||
|
|
||||||
@@ -445,15 +460,18 @@ class McpClient:
|
|||||||
|
|
||||||
message = self._build_request(method, params)
|
message = self._build_request(method, params)
|
||||||
|
|
||||||
with self._call_lock:
|
# stdio transport uses a single pipe and must be serialized.
|
||||||
if self.transport == "stdio":
|
# SSE and streamable-http use independent HTTP requests and
|
||||||
|
# can safely run concurrently across sessions.
|
||||||
|
if self.transport == "stdio":
|
||||||
|
with self._call_lock:
|
||||||
return self._stdio_send(message)
|
return self._stdio_send(message)
|
||||||
elif self.transport == "sse":
|
elif self.transport == "sse":
|
||||||
return self._sse_send(message)
|
return self._sse_send(message)
|
||||||
elif self.transport == "streamable-http":
|
elif self.transport == "streamable-http":
|
||||||
return self._streamable_http_send(message)
|
return self._streamable_http_send(message)
|
||||||
else:
|
else:
|
||||||
raise ValueError(f"[MCP:{self.name}] Unsupported transport: {self.transport}")
|
raise ValueError(f"[MCP:{self.name}] Unsupported transport: {self.transport}")
|
||||||
|
|
||||||
def _send_notification(self, method: str, params: dict):
|
def _send_notification(self, method: str, params: dict):
|
||||||
"""Fire-and-forget notification (no response expected)."""
|
"""Fire-and-forget notification (no response expected)."""
|
||||||
|
|||||||
Reference in New Issue
Block a user