From 37b969ac9a957969345fe845e67ef65f82a898f7 Mon Sep 17 00:00:00 2001 From: pengzeyu Date: Mon, 21 Sep 2026 17:32:36 +0800 Subject: [PATCH 1/4] fix(api): reuse MCP client connections across tool invocations --- api/core/mcp/client_manager.py | 237 ++++++++++++++++++ api/core/tools/mcp_tool/tool.py | 17 +- .../core/mcp/test_client_manager.py | 130 ++++++++++ .../unit_tests/core/tools/test_mcp_tool.py | 8 +- 4 files changed, 383 insertions(+), 9 deletions(-) create mode 100644 api/core/mcp/client_manager.py create mode 100644 api/tests/unit_tests/core/mcp/test_client_manager.py diff --git a/api/core/mcp/client_manager.py b/api/core/mcp/client_manager.py new file mode 100644 index 00000000000000..3df5c18459b294 --- /dev/null +++ b/api/core/mcp/client_manager.py @@ -0,0 +1,237 @@ +"""Process-wide reuse of MCP client connections. + +Every `MCPTool` invocation used to open an MCP connection and close it as soon +as the call returned, which discards any per-session state a stateful MCP +server keeps (for example the browser context of the Playwright MCP server: +`browser_navigate` succeeds, the next call lands on a fresh `about:blank` +page). This module keeps successfully initialized clients in a process-wide +pool keyed by server URL and effective request headers, so consecutive calls +within the idle TTL reuse the same session. + +The pool is deliberately conservative: each pooled client is guarded by a +lock (the sync `ClientSession` processes messages on a single worker), +entries are evicted after an idle TTL or when the pool exceeds its size cap, +and a pooled connection that died is dropped and retried once on a fresh +connection, keeping the per-call robustness of connect-per-invocation. +""" + +import hashlib +import json +import logging +import os +import threading +import time +from typing import Any + +from core.entities.mcp_provider import MCPProviderEntity +from core.mcp.auth_client import MCPClientWithAuthRetry +from core.mcp.error import MCPAuthError, MCPConnectionError +from core.mcp.types import CallToolResult + +logger = logging.getLogger(__name__) + +# Seconds after which an idle pooled client is closed on its next acquire, and +# the maximum number of clients kept open per process. Overridable for tests +# and for deployments that rely on long-lived stateful MCP servers. +DEFAULT_IDLE_TTL_SECONDS = float(os.getenv("MCP_CLIENT_POOL_IDLE_TTL", "300")) +DEFAULT_MAX_SIZE = int(os.getenv("MCP_CLIENT_POOL_MAX_SIZE", "100")) + + +class _PooledClient: + def __init__(self, client: MCPClientWithAuthRetry, key: str): + self.client = client + # The pool key this entry was registered under. Kept explicitly because + # the client may refresh its own Authorization header after an auth + # retry, which would change a recomputed key. + self.key = key + self.last_used = time.monotonic() + # ClientSession drives a single-worker executor; serialize calls that + # share one pooled connection. + self.lock = threading.Lock() + + +class MCPClientManager: + """A process-wide pool of authenticated MCP clients.""" + + def __init__( + self, + idle_ttl_seconds: float = DEFAULT_IDLE_TTL_SECONDS, + max_size: int = DEFAULT_MAX_SIZE, + ): + self._clients: dict[str, _PooledClient] = {} + self._lock = threading.Lock() + self._idle_ttl_seconds = idle_ttl_seconds + self._max_size = max_size + + @staticmethod + def _make_key( + server_url: str, + headers: dict[str, str] | None, + timeout: float | None, + sse_read_timeout: float | None, + ) -> str: + """Identify a connection by everything that goes into building it. + + Headers carry the (possibly per-user) credentials, so distinct tenants, + users or tokens naturally map to distinct pooled connections. + """ + material = json.dumps( + { + "server_url": server_url, + "headers": dict(sorted((headers or {}).items())), + "timeout": timeout, + "sse_read_timeout": sse_read_timeout, + }, + sort_keys=True, + ) + return hashlib.sha256(material.encode("utf-8")).hexdigest() + + def _close_entry(self, entry: _PooledClient) -> None: + with entry.lock: + try: + entry.client.cleanup() + except Exception: + logger.warning("Failed to clean up a pooled MCP client", exc_info=True) + + def _evict_expired(self) -> None: + now = time.monotonic() + for key, entry in list(self._clients.items()): + if now - entry.last_used >= self._idle_ttl_seconds: + del self._clients[key] + self._close_entry(entry) + + def _evict_to_capacity(self) -> None: + overflow = len(self._clients) - self._max_size + if overflow <= 0: + return + for key, entry in sorted(self._clients.items(), key=lambda item: item[1].last_used)[:overflow]: + del self._clients[key] + self._close_entry(entry) + + def _acquire( + self, + server_url: str, + headers: dict[str, str] | None, + timeout: float | None, + sse_read_timeout: float | None, + provider_entity: MCPProviderEntity | None, + forward_identity_active: bool, + ) -> _PooledClient: + """Return a live pooled entry, connecting (with auth retry) if needed.""" + key = self._make_key(server_url, headers, timeout, sse_read_timeout) + while True: + with self._lock: + self._evict_expired() + entry = self._clients.get(key) + if entry is not None: + entry.last_used = time.monotonic() + return entry + + # Connect outside the registry lock: initialization performs + # network I/O (and possibly an OAuth refresh) that must not block + # unrelated MCP calls. + client = MCPClientWithAuthRetry( + server_url=server_url, + headers=headers, + timeout=timeout, + sse_read_timeout=sse_read_timeout, + provider_entity=provider_entity, + forward_identity_active=forward_identity_active, + ) + try: + client.__enter__() + except Exception: + # Never pool a client that failed to initialize; keep the + # connect-per-invocation behavior for failing servers. + try: + client.cleanup() + except Exception: + logger.warning("Failed to clean up an MCP client whose initialization failed", exc_info=True) + raise + candidate = _PooledClient(client, key) + with self._lock: + existing = self._clients.get(key) + if existing is not None: + # Lost a race creating the same connection: use the winner. + self._close_entry(candidate) + existing.last_used = time.monotonic() + return existing + self._clients[key] = candidate + self._evict_to_capacity() + return candidate + + def _evict(self, entry: _PooledClient) -> None: + with self._lock: + pooled = self._clients.get(entry.key) + if pooled is entry: + del self._clients[entry.key] + self._close_entry(entry) + + def invoke_tool( + self, + *, + server_url: str, + headers: dict[str, str] | None, + timeout: float | None, + sse_read_timeout: float | None, + provider_entity: MCPProviderEntity | None, + forward_identity_active: bool, + tool_name: str, + tool_args: dict[str, Any], + ) -> CallToolResult: + """Invoke a tool on a pooled connection. + + A pooled connection that turns out to be dead (`MCPConnectionError`) + is closed and the call is retried once on a fresh connection, so a + stale pool entry fails no more calls than connect-per-invocation did. + Auth failures must not retry here (`MCPAuthError` also extends + `MCPConnectionError`): the pooled client refreshes tokens internally, + so an error that escapes it means re-auth is not possible, and + re-invoking could execute the tool twice. They propagate untouched. + """ + entry = self._acquire(server_url, headers, timeout, sse_read_timeout, provider_entity, forward_identity_active) + try: + with entry.lock: + return entry.client.invoke_tool(tool_name=tool_name, tool_args=tool_args) + except MCPAuthError: + raise + except MCPConnectionError: + logger.info("Pooled MCP connection to %s died; reconnecting once", server_url) + self._evict(entry) + entry = self._acquire( + server_url, headers, timeout, sse_read_timeout, provider_entity, forward_identity_active + ) + with entry.lock: + return entry.client.invoke_tool(tool_name=tool_name, tool_args=tool_args) + finally: + entry.last_used = time.monotonic() + + def close_all(self) -> None: + """Close every pooled client (used on shutdown and in tests).""" + with self._lock: + entries = list(self._clients.values()) + self._clients.clear() + for entry in entries: + self._close_entry(entry) + + +_manager: MCPClientManager | None = None +_manager_lock = threading.Lock() + + +def get_mcp_client_manager() -> MCPClientManager: + global _manager + with _manager_lock: + if _manager is None: + _manager = MCPClientManager() + return _manager + + +def reset_mcp_client_manager() -> None: + """Drop the process-wide manager (tests only).""" + global _manager + with _manager_lock: + manager = _manager + _manager = None + if manager is not None: + manager.close_all() diff --git a/api/core/tools/mcp_tool/tool.py b/api/core/tools/mcp_tool/tool.py index 84235403d61cfe..fc8a54bda9fd19 100644 --- a/api/core/tools/mcp_tool/tool.py +++ b/api/core/tools/mcp_tool/tool.py @@ -10,7 +10,6 @@ from configs import dify_config from core.entities.mcp_provider import IdentityMode -from core.mcp.auth_client import MCPClientWithAuthRetry from core.mcp.error import MCPConnectionError from core.mcp.types import ( AudioContent, @@ -322,18 +321,24 @@ def invoke_remote_mcp_tool( self._inject_forwarded_identity(headers, user_id=user_id, app_id=app_id, audience=server_url) forward_identity_active = True - # Step 2: Session is now closed, perform network operations without holding database connection - # MCPClientWithAuthRetry will create a new session lazily only if auth retry is needed + # Step 2: Session is now closed, perform network operations without holding database connection. + # The pooled manager reuses one connection across calls so stateful MCP + # servers (e.g. Playwright) keep their session state between calls; it + # reconnects once when a pooled connection has died, and auth retries + # are handled inside the pooled client. + from core.mcp.client_manager import get_mcp_client_manager + try: - with MCPClientWithAuthRetry( + return get_mcp_client_manager().invoke_tool( server_url=server_url, headers=headers, timeout=self.timeout, sse_read_timeout=self.sse_read_timeout, provider_entity=provider_entity, forward_identity_active=forward_identity_active, - ) as mcp_client: - return mcp_client.invoke_tool(tool_name=self.entity.identity.name, tool_args=tool_parameters) + tool_name=self.entity.identity.name, + tool_args=tool_parameters, + ) except MCPConnectionError as e: raise ToolInvokeError(f"Failed to connect to MCP server: {e}") from e except Exception as e: diff --git a/api/tests/unit_tests/core/mcp/test_client_manager.py b/api/tests/unit_tests/core/mcp/test_client_manager.py new file mode 100644 index 00000000000000..1e2fdb85f87897 --- /dev/null +++ b/api/tests/unit_tests/core/mcp/test_client_manager.py @@ -0,0 +1,130 @@ +"""Unit tests for the pooled MCP client manager.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from core.mcp.client_manager import MCPClientManager +from core.mcp.error import MCPAuthError, MCPConnectionError + + +def _invoke(manager: MCPClientManager, headers: dict[str, str] | None = None, **kwargs): + params = { + "server_url": "http://test.example.com/mcp", + "headers": headers if headers is not None else {"Authorization": "Bearer token"}, + "timeout": 30.0, + "sse_read_timeout": 60.0, + "provider_entity": None, + "forward_identity_active": False, + "tool_name": "browser_navigate", + "tool_args": {"url": "https://example.com"}, + } + params.update(kwargs) + return manager.invoke_tool(**params) + + +class TestMCPClientManager: + def test_same_connection_reused_across_calls(self): + manager = MCPClientManager() + with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: + client = MagicMock() + factory.return_value = client + + first = _invoke(manager) + second = _invoke(manager) + + assert first is second + assert factory.call_count == 1 + assert client.invoke_tool.call_count == 2 + client.cleanup.assert_not_called() + + def test_distinct_headers_get_distinct_connections(self): + manager = MCPClientManager() + with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: + factory.side_effect = [MagicMock(), MagicMock()] + + _invoke(manager, headers={"Authorization": "Bearer user-a"}) + _invoke(manager, headers={"Authorization": "Bearer user-b"}) + + assert factory.call_count == 2 + + def test_dead_connection_is_evicted_and_retried(self): + manager = MCPClientManager() + with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: + dead, alive = MagicMock(), MagicMock() + dead.invoke_tool.side_effect = MCPConnectionError("connection reset") + factory.side_effect = [dead, alive] + + result = _invoke(manager) + + assert result is alive.invoke_tool.return_value + assert factory.call_count == 2 + dead.cleanup.assert_called_once() + + def test_auth_error_propagates_but_keeps_connection(self): + manager = MCPClientManager() + with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: + client = MagicMock() + client.invoke_tool.side_effect = MCPAuthError("401") + factory.return_value = client + + with pytest.raises(MCPAuthError): + _invoke(manager) + + # The client refreshes tokens internally; it stays pooled. + assert factory.call_count == 1 + + client.invoke_tool.side_effect = None + _invoke(manager) + assert factory.call_count == 1 + + def test_idle_ttl_closes_stale_connection(self): + manager = MCPClientManager(idle_ttl_seconds=0) + with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: + stale, fresh = MagicMock(), MagicMock() + factory.side_effect = [stale, fresh] + + _invoke(manager) + _invoke(manager) + + assert factory.call_count == 2 + stale.cleanup.assert_called_once() + + def test_max_size_evicts_least_recently_used(self): + manager = MCPClientManager(max_size=1) + with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: + oldest, newest = MagicMock(), MagicMock() + factory.side_effect = [oldest, newest] + + _invoke(manager, headers={"Authorization": "Bearer a"}) + _invoke(manager, headers={"Authorization": "Bearer b"}) + + assert factory.call_count == 2 + oldest.cleanup.assert_called_once() + + def test_failed_initialize_is_not_pooled(self): + manager = MCPClientManager() + with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: + failing = MagicMock() + failing.__enter__.side_effect = MCPConnectionError("server down") + ok = MagicMock() + factory.side_effect = [failing, ok] + + with pytest.raises(MCPConnectionError): + _invoke(manager) + + result = _invoke(manager) + assert result is ok.invoke_tool.return_value + + def test_close_all_closes_every_connection(self): + manager = MCPClientManager() + with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: + factory.side_effect = [MagicMock(), MagicMock()] + _invoke(manager, headers={"Authorization": "Bearer a"}) + _invoke(manager, headers={"Authorization": "Bearer b"}) + + manager.close_all() + + for client in factory.side_effect: + client.cleanup.assert_called_once() + assert manager._clients == {} diff --git a/api/tests/unit_tests/core/tools/test_mcp_tool.py b/api/tests/unit_tests/core/tools/test_mcp_tool.py index 2944c378a43503..997ed4eaf11be1 100644 --- a/api/tests/unit_tests/core/tools/test_mcp_tool.py +++ b/api/tests/unit_tests/core/tools/test_mcp_tool.py @@ -278,11 +278,13 @@ def test_invoke_skips_forwarding_outside_enterprise_edition(config_overrides): # The fail-closed branch must NOT fire (no enterprise → no forwarding). # The function will still try the legacy DB-load path; we patch that # to keep the test unit-scoped. - with patch("core.tools.mcp_tool.tool.MCPClientWithAuthRetry") as client_cls: - client_cls.return_value.__enter__.return_value.invoke_tool.return_value = CallToolResult( + with patch( + "core.mcp.client_manager.MCPClientManager.invoke_tool", + return_value=CallToolResult( content=[], _meta=None, - ) + ), + ): with patch.object(tool, "_inject_forwarded_identity") as inject: with patch("services.tools.mcp_tools_manage_service.MCPToolManageService"): with patch("core.entities.mcp_provider.MCPProviderEntity.decrypt_server_url", return_value="u"): From 155d75ee3891086e37cf686502dc240dd02b9fe3 Mon Sep 17 00:00:00 2001 From: pengzeyu Date: Mon, 21 Sep 2026 19:25:43 +0800 Subject: [PATCH 2/4] fix(api): key MCP client pool by stable scope and close entries off the registry lock --- api/core/mcp/client_manager.py | 246 ++++++++++++------ api/core/tools/mcp_tool/tool.py | 13 +- .../core/mcp/test_client_manager.py | 143 +++++++++- 3 files changed, 306 insertions(+), 96 deletions(-) diff --git a/api/core/mcp/client_manager.py b/api/core/mcp/client_manager.py index 3df5c18459b294..285641ea222842 100644 --- a/api/core/mcp/client_manager.py +++ b/api/core/mcp/client_manager.py @@ -5,16 +5,28 @@ server keeps (for example the browser context of the Playwright MCP server: `browser_navigate` succeeds, the next call lands on a fresh `about:blank` page). This module keeps successfully initialized clients in a process-wide -pool keyed by server URL and effective request headers, so consecutive calls -within the idle TTL reuse the same session. - -The pool is deliberately conservative: each pooled client is guarded by a -lock (the sync `ClientSession` processes messages on a single worker), -entries are evicted after an idle TTL or when the pool exceeds its size cap, -and a pooled connection that died is dropped and retried once on a fresh -connection, keeping the per-call robustness of connect-per-invocation. +pool so consecutive calls reuse the same session. + +The pool key is the execution scope — tenant, end user, provider config — +plus the server URL and timeouts: stable identities, never the per-call +credentials. Credentials change between calls (a freshly minted +forwarded-identity token, or an OAuth token after a refresh); keying on them +would prevent all reuse, while the session they authenticated at +`initialize()` keeps working and the pooled client refreshes its own token on +401. Distinct tenants, end users and provider configs therefore map to +distinct connections, so a stateful server never shares session state across +them. + +Locking is two-level and strictly ordered: the registry lock only guards the +`_clients` dict, and per-entry locks serialize calls on one connection. +Entries are never closed while holding the registry lock — an in-flight call +holds the entry lock for as long as the server takes to answer (the session +has no read timeout), so eviction removes the entry under the registry lock +and closes it afterwards; a busy entry is marked doomed and closed by its +in-flight call on the way out. """ +import atexit import hashlib import json import logging @@ -30,21 +42,18 @@ logger = logging.getLogger(__name__) -# Seconds after which an idle pooled client is closed on its next acquire, and -# the maximum number of clients kept open per process. Overridable for tests -# and for deployments that rely on long-lived stateful MCP servers. -DEFAULT_IDLE_TTL_SECONDS = float(os.getenv("MCP_CLIENT_POOL_IDLE_TTL", "300")) -DEFAULT_MAX_SIZE = int(os.getenv("MCP_CLIENT_POOL_MAX_SIZE", "100")) - class _PooledClient: def __init__(self, client: MCPClientWithAuthRetry, key: str): self.client = client - # The pool key this entry was registered under. Kept explicitly because - # the client may refresh its own Authorization header after an auth - # retry, which would change a recomputed key. + # The pool key this entry was registered under, kept explicitly so + # eviction never has to recompute it from mutable client state. self.key = key self.last_used = time.monotonic() + # Set when the entry has been removed from the registry but its lock + # is held by an in-flight call; that call closes it on the way out. + self.doomed = False + self.closed = False # ClientSession drives a single-worker executor; serialize calls that # share one pooled connection. self.lock = threading.Lock() @@ -55,30 +64,41 @@ class MCPClientManager: def __init__( self, - idle_ttl_seconds: float = DEFAULT_IDLE_TTL_SECONDS, - max_size: int = DEFAULT_MAX_SIZE, + idle_ttl_seconds: float | None = None, + max_size: int | None = None, ): self._clients: dict[str, _PooledClient] = {} self._lock = threading.Lock() - self._idle_ttl_seconds = idle_ttl_seconds - self._max_size = max_size + # Read when the manager is constructed, so the knobs respond to + # environment changes without re-importing the module. + self._idle_ttl_seconds = ( + float(os.getenv("MCP_CLIENT_POOL_IDLE_TTL", "300")) if idle_ttl_seconds is None else idle_ttl_seconds + ) + self._max_size = int(os.getenv("MCP_CLIENT_POOL_MAX_SIZE", "100")) if max_size is None else max_size @staticmethod def _make_key( + tenant_id: str, + user_id: str | None, server_url: str, - headers: dict[str, str] | None, + provider_id: str | None, timeout: float | None, sse_read_timeout: float | None, ) -> str: - """Identify a connection by everything that goes into building it. + """Identify a pooled connection by stable scope, never by credentials. - Headers carry the (possibly per-user) credentials, so distinct tenants, - users or tokens naturally map to distinct pooled connections. + tenant/user/provider separate tenants, end users (including forwarded + identities) and provider configs; server URL and timeouts are the + connection fingerprint. Per-call credentials are deliberately + excluded: forwarded-identity tokens are minted per call and OAuth + tokens rotate, so keying on either would prevent all reuse. """ material = json.dumps( { + "tenant_id": tenant_id, + "user_id": user_id, "server_url": server_url, - "headers": dict(sorted((headers or {}).items())), + "provider_id": provider_id, "timeout": timeout, "sse_read_timeout": sse_read_timeout, }, @@ -86,30 +106,56 @@ def _make_key( ) return hashlib.sha256(material.encode("utf-8")).hexdigest() + def _shutdown_entry(self, entry: _PooledClient) -> None: + """Idempotently close the underlying client.""" + if entry.closed: + return + entry.closed = True + try: + entry.client.cleanup() + except Exception: + logger.warning("Failed to clean up a pooled MCP client", exc_info=True) + def _close_entry(self, entry: _PooledClient) -> None: - with entry.lock: - try: - entry.client.cleanup() - except Exception: - logger.warning("Failed to clean up a pooled MCP client", exc_info=True) + """Close `entry`; must be called without holding the registry lock. - def _evict_expired(self) -> None: + An in-flight call holds the entry lock for as long as the server takes + to answer, so the registry lock must never wait on it. When the entry + is busy, mark it doomed and let the in-flight call close it in its + finally block instead. + """ + if not entry.lock.acquire(blocking=False): + entry.doomed = True + return + try: + self._shutdown_entry(entry) + finally: + entry.lock.release() + + def _evict_expired(self) -> list[_PooledClient]: + """Remove idle-expired entries; caller closes them outside the lock.""" now = time.monotonic() + victims = [] for key, entry in list(self._clients.items()): if now - entry.last_used >= self._idle_ttl_seconds: del self._clients[key] - self._close_entry(entry) + victims.append(entry) + return victims - def _evict_to_capacity(self) -> None: + def _evict_to_capacity(self) -> list[_PooledClient]: + """Drop least-recently-used entries above the size cap; caller closes them.""" overflow = len(self._clients) - self._max_size if overflow <= 0: - return + return [] + victims = [] for key, entry in sorted(self._clients.items(), key=lambda item: item[1].last_used)[:overflow]: del self._clients[key] - self._close_entry(entry) + victims.append(entry) + return victims def _acquire( self, + key: str, server_url: str, headers: dict[str, str] | None, timeout: float | None, @@ -118,47 +164,52 @@ def _acquire( forward_identity_active: bool, ) -> _PooledClient: """Return a live pooled entry, connecting (with auth retry) if needed.""" - key = self._make_key(server_url, headers, timeout, sse_read_timeout) - while True: - with self._lock: - self._evict_expired() - entry = self._clients.get(key) - if entry is not None: - entry.last_used = time.monotonic() - return entry - - # Connect outside the registry lock: initialization performs - # network I/O (and possibly an OAuth refresh) that must not block - # unrelated MCP calls. - client = MCPClientWithAuthRetry( - server_url=server_url, - headers=headers, - timeout=timeout, - sse_read_timeout=sse_read_timeout, - provider_entity=provider_entity, - forward_identity_active=forward_identity_active, - ) + with self._lock: + victims = self._evict_expired() + entry = self._clients.get(key) + if entry is not None: + entry.last_used = time.monotonic() + for victim in victims: + self._close_entry(victim) + if entry is not None: + return entry + + # Connect outside the registry lock: initialization performs network + # I/O (and possibly an OAuth refresh) that must not block unrelated + # MCP calls. Credentials may differ between calls of the same scope; + # the first caller's headers authenticate the session and the pooled + # client refreshes its own token on 401. + client = MCPClientWithAuthRetry( + server_url=server_url, + headers=headers, + timeout=timeout, + sse_read_timeout=sse_read_timeout, + provider_entity=provider_entity, + forward_identity_active=forward_identity_active, + ) + try: + client.__enter__() + except Exception: + # Never pool a client that failed to initialize; keep the + # connect-per-invocation behavior for failing servers. try: - client.__enter__() + client.cleanup() except Exception: - # Never pool a client that failed to initialize; keep the - # connect-per-invocation behavior for failing servers. - try: - client.cleanup() - except Exception: - logger.warning("Failed to clean up an MCP client whose initialization failed", exc_info=True) - raise - candidate = _PooledClient(client, key) - with self._lock: - existing = self._clients.get(key) - if existing is not None: - # Lost a race creating the same connection: use the winner. - self._close_entry(candidate) - existing.last_used = time.monotonic() - return existing + logger.warning("Failed to clean up an MCP client whose initialization failed", exc_info=True) + raise + candidate = _PooledClient(client, key) + with self._lock: + existing = self._clients.get(key) + if existing is None: self._clients[key] = candidate - self._evict_to_capacity() - return candidate + victims = self._evict_to_capacity() + for victim in victims: + self._close_entry(victim) + if existing is not None: + # Lost a race creating the same connection: use the winner. + self._close_entry(candidate) + return existing + return candidate def _evict(self, entry: _PooledClient) -> None: with self._lock: @@ -170,7 +221,10 @@ def _evict(self, entry: _PooledClient) -> None: def invoke_tool( self, *, + tenant_id: str, + user_id: str | None, server_url: str, + provider_id: str | None, headers: dict[str, str] | None, timeout: float | None, sse_read_timeout: float | None, @@ -182,14 +236,20 @@ def invoke_tool( """Invoke a tool on a pooled connection. A pooled connection that turns out to be dead (`MCPConnectionError`) - is closed and the call is retried once on a fresh connection, so a + is evicted and the call retried once on a fresh connection, so a stale pool entry fails no more calls than connect-per-invocation did. - Auth failures must not retry here (`MCPAuthError` also extends - `MCPConnectionError`): the pooled client refreshes tokens internally, - so an error that escapes it means re-auth is not possible, and - re-invoking could execute the tool twice. They propagate untouched. + Auth failures are deliberately not retried (`MCPAuthError` also + extends `MCPConnectionError`): token refresh already happened inside + the client, so an auth error that escapes it means re-auth is not + possible, and re-invoking could execute the tool twice. Any other + error escaping a session (e.g. a `ValueError` from transport-level + parsing) evicts the entry without a retry — the next call reconnects + instead of failing on the same broken session for a full TTL. """ - entry = self._acquire(server_url, headers, timeout, sse_read_timeout, provider_entity, forward_identity_active) + key = self._make_key(tenant_id, user_id, server_url, provider_id, timeout, sse_read_timeout) + entry = self._acquire( + key, server_url, headers, timeout, sse_read_timeout, provider_entity, forward_identity_active + ) try: with entry.lock: return entry.client.invoke_tool(tool_name=tool_name, tool_args=tool_args) @@ -199,15 +259,23 @@ def invoke_tool( logger.info("Pooled MCP connection to %s died; reconnecting once", server_url) self._evict(entry) entry = self._acquire( - server_url, headers, timeout, sse_read_timeout, provider_entity, forward_identity_active + key, server_url, headers, timeout, sse_read_timeout, provider_entity, forward_identity_active ) with entry.lock: return entry.client.invoke_tool(tool_name=tool_name, tool_args=tool_args) + except ValueError: + # The session may be corrupted (e.g. an unexpected content type in + # the streamable client kills its receive loop): drop it so the + # next call reconnects. No retry — the tool may already have run. + self._evict(entry) + raise finally: entry.last_used = time.monotonic() + if entry.doomed: + self._shutdown_entry(entry) def close_all(self) -> None: - """Close every pooled client (used on shutdown and in tests).""" + """Close every pooled client (registered as an atexit hook).""" with self._lock: entries = list(self._clients.values()) self._clients.clear() @@ -227,6 +295,14 @@ def get_mcp_client_manager() -> MCPClientManager: return _manager +def close_all_mcp_clients() -> None: + """Close every pooled client, so stateful servers do not outlive the process.""" + with _manager_lock: + manager = _manager + if manager is not None: + manager.close_all() + + def reset_mcp_client_manager() -> None: """Drop the process-wide manager (tests only).""" global _manager @@ -235,3 +311,9 @@ def reset_mcp_client_manager() -> None: _manager = None if manager is not None: manager.close_all() + + +# Mirror core/helper/http_client_pooling.py: pooled resources get an atexit +# hook so long-lived sessions (e.g. a Playwright browser) are closed when the +# process exits instead of idling until the TTL. +atexit.register(close_all_mcp_clients) diff --git a/api/core/tools/mcp_tool/tool.py b/api/core/tools/mcp_tool/tool.py index fc8a54bda9fd19..c20fe81b803125 100644 --- a/api/core/tools/mcp_tool/tool.py +++ b/api/core/tools/mcp_tool/tool.py @@ -322,15 +322,20 @@ def invoke_remote_mcp_tool( forward_identity_active = True # Step 2: Session is now closed, perform network operations without holding database connection. - # The pooled manager reuses one connection across calls so stateful MCP - # servers (e.g. Playwright) keep their session state between calls; it - # reconnects once when a pooled connection has died, and auth retries - # are handled inside the pooled client. + # The pooled manager reuses one connection per stable scope (tenant, + # end user, provider config, server) so stateful MCP servers (e.g. + # Playwright) keep their session state between calls without sharing + # it across users; per-call credentials are used to connect but not + # to key the pool. It reconnects once when a pooled connection has + # died, and auth retries are handled inside the pooled client. from core.mcp.client_manager import get_mcp_client_manager try: return get_mcp_client_manager().invoke_tool( + tenant_id=self.tenant_id, + user_id=user_id, server_url=server_url, + provider_id=provider_entity.id, headers=headers, timeout=self.timeout, sse_read_timeout=self.sse_read_timeout, diff --git a/api/tests/unit_tests/core/mcp/test_client_manager.py b/api/tests/unit_tests/core/mcp/test_client_manager.py index 1e2fdb85f87897..1098ac8bf562c2 100644 --- a/api/tests/unit_tests/core/mcp/test_client_manager.py +++ b/api/tests/unit_tests/core/mcp/test_client_manager.py @@ -1,5 +1,6 @@ """Unit tests for the pooled MCP client manager.""" +import threading from unittest.mock import MagicMock, patch import pytest @@ -8,10 +9,13 @@ from core.mcp.error import MCPAuthError, MCPConnectionError -def _invoke(manager: MCPClientManager, headers: dict[str, str] | None = None, **kwargs): +def _invoke(manager: MCPClientManager, **kwargs): params = { + "tenant_id": "tenant-1", + "user_id": "user-1", "server_url": "http://test.example.com/mcp", - "headers": headers if headers is not None else {"Authorization": "Bearer token"}, + "provider_id": "provider-1", + "headers": {"Authorization": "Bearer token"}, "timeout": 30.0, "sse_read_timeout": 60.0, "provider_entity": None, @@ -24,7 +28,7 @@ def _invoke(manager: MCPClientManager, headers: dict[str, str] | None = None, ** class TestMCPClientManager: - def test_same_connection_reused_across_calls(self): + def test_same_scope_reuses_connection(self): manager = MCPClientManager() with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: client = MagicMock() @@ -38,16 +42,40 @@ def test_same_connection_reused_across_calls(self): assert client.invoke_tool.call_count == 2 client.cleanup.assert_not_called() - def test_distinct_headers_get_distinct_connections(self): + def test_changing_credentials_still_reuse_connection(self): + """Per-call credentials (minted JWTs, refreshed OAuth tokens) must not split the pool.""" + manager = MCPClientManager() + with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: + client = MagicMock() + factory.return_value = client + + _invoke(manager, headers={"Authorization": "Bearer minted-1"}) + _invoke(manager, headers={"Authorization": "Bearer minted-2"}) + + assert factory.call_count == 1 + + def test_distinct_users_get_distinct_connections(self): + """A stateful server must not share one session across end users.""" manager = MCPClientManager() with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: factory.side_effect = [MagicMock(), MagicMock()] - _invoke(manager, headers={"Authorization": "Bearer user-a"}) - _invoke(manager, headers={"Authorization": "Bearer user-b"}) + _invoke(manager, user_id="user-1") + _invoke(manager, user_id="user-2") assert factory.call_count == 2 + def test_distinct_tenants_and_providers_get_distinct_connections(self): + manager = MCPClientManager() + with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: + factory.side_effect = [MagicMock(), MagicMock(), MagicMock()] + + _invoke(manager, tenant_id="tenant-1") + _invoke(manager, tenant_id="tenant-2") + _invoke(manager, provider_id="provider-2") + + assert factory.call_count == 3 + def test_dead_connection_is_evicted_and_retried(self): manager = MCPClientManager() with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: @@ -78,6 +106,101 @@ def test_auth_error_propagates_but_keeps_connection(self): _invoke(manager) assert factory.call_count == 1 + def test_value_error_evicts_without_retry(self): + """Transport-level parsing failures corrupt the session: drop it, do not re-invoke.""" + manager = MCPClientManager() + with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: + broken = MagicMock() + broken.invoke_tool.side_effect = ValueError("Unexpected content type") + factory.side_effect = [broken, MagicMock()] + + with pytest.raises(ValueError): + _invoke(manager) + + # No retry happened, but the entry was dropped for the next call. + assert factory.call_count == 1 + _invoke(manager) + assert factory.call_count == 2 + broken.cleanup.assert_called_once() + + def test_blocked_call_does_not_block_other_scopes(self): + """A hung call on one connection must not wedge calls on other connections.""" + manager = MCPClientManager() + release = threading.Event() + hung_entered = threading.Event() + + def hung_invoke(**_kwargs): + hung_entered.set() + release.wait(10) + return "hung-ok" + + with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: + + def make_client(**kwargs): + client = MagicMock() + if kwargs.get("headers", {}).get("X-User") == "hung": + client.invoke_tool.side_effect = hung_invoke + return client + + factory.side_effect = make_client + + hung_thread = threading.Thread( + target=lambda: _invoke(manager, user_id="user-hung", headers={"X-User": "hung"}) + ) + hung_thread.start() + assert hung_entered.wait(2) + + other_results = [] + other_thread = threading.Thread(target=lambda: other_results.append(_invoke(manager, user_id="user-2"))) + other_thread.start() + other_thread.join(timeout=3) + + assert not other_thread.is_alive(), "a hung call on one connection blocked an unrelated scope" + assert other_results[0] is not None + + release.set() + hung_thread.join(timeout=3) + + def test_evicting_in_flight_entry_does_not_wedge_new_acquires(self): + """Idle eviction of an entry whose call is still running must not block the same key.""" + manager = MCPClientManager(idle_ttl_seconds=0) + release = threading.Event() + hung_entered = threading.Event() + + def hung_invoke(**_kwargs): + hung_entered.set() + release.wait(10) + return "hung-ok" + + with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: + + def make_client(**kwargs): + client = MagicMock() + if kwargs.get("headers", {}).get("X-User") == "hung": + client.invoke_tool.side_effect = hung_invoke + return client + + factory.side_effect = make_client + + hung_thread = threading.Thread( + target=lambda: _invoke(manager, user_id="user-a", headers={"X-User": "hung"}) + ) + hung_thread.start() + assert hung_entered.wait(2) + + same_key_results = [] + same_key_thread = threading.Thread( + target=lambda: same_key_results.append(_invoke(manager, user_id="user-a")) + ) + same_key_thread.start() + same_key_thread.join(timeout=3) + + assert not same_key_thread.is_alive(), "eviction of an in-flight entry blocked its own key" + assert same_key_results[0] is not None + + release.set() + hung_thread.join(timeout=3) + def test_idle_ttl_closes_stale_connection(self): manager = MCPClientManager(idle_ttl_seconds=0) with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: @@ -96,8 +219,8 @@ def test_max_size_evicts_least_recently_used(self): oldest, newest = MagicMock(), MagicMock() factory.side_effect = [oldest, newest] - _invoke(manager, headers={"Authorization": "Bearer a"}) - _invoke(manager, headers={"Authorization": "Bearer b"}) + _invoke(manager, user_id="user-a") + _invoke(manager, user_id="user-b") assert factory.call_count == 2 oldest.cleanup.assert_called_once() @@ -120,8 +243,8 @@ def test_close_all_closes_every_connection(self): manager = MCPClientManager() with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: factory.side_effect = [MagicMock(), MagicMock()] - _invoke(manager, headers={"Authorization": "Bearer a"}) - _invoke(manager, headers={"Authorization": "Bearer b"}) + _invoke(manager, user_id="user-a") + _invoke(manager, user_id="user-b") manager.close_all() From a496b6c82f84ec4116303d3c765c165e83375af9 Mon Sep 17 00:00:00 2001 From: pengzeyu Date: Mon, 21 Sep 2026 19:40:08 +0800 Subject: [PATCH 3/4] refactor(api): separate eviction victim lists and drop the unused manager reset --- api/core/mcp/client_manager.py | 19 +++++-------------- 1 file changed, 5 insertions(+), 14 deletions(-) diff --git a/api/core/mcp/client_manager.py b/api/core/mcp/client_manager.py index 285641ea222842..d33c1771c1fad8 100644 --- a/api/core/mcp/client_manager.py +++ b/api/core/mcp/client_manager.py @@ -165,11 +165,11 @@ def _acquire( ) -> _PooledClient: """Return a live pooled entry, connecting (with auth retry) if needed.""" with self._lock: - victims = self._evict_expired() + expired = self._evict_expired() entry = self._clients.get(key) if entry is not None: entry.last_used = time.monotonic() - for victim in victims: + for victim in expired: self._close_entry(victim) if entry is not None: return entry @@ -198,12 +198,13 @@ def _acquire( logger.warning("Failed to clean up an MCP client whose initialization failed", exc_info=True) raise candidate = _PooledClient(client, key) + overflow: list[_PooledClient] = [] with self._lock: existing = self._clients.get(key) if existing is None: self._clients[key] = candidate - victims = self._evict_to_capacity() - for victim in victims: + overflow = self._evict_to_capacity() + for victim in overflow: self._close_entry(victim) if existing is not None: # Lost a race creating the same connection: use the winner. @@ -303,16 +304,6 @@ def close_all_mcp_clients() -> None: manager.close_all() -def reset_mcp_client_manager() -> None: - """Drop the process-wide manager (tests only).""" - global _manager - with _manager_lock: - manager = _manager - _manager = None - if manager is not None: - manager.close_all() - - # Mirror core/helper/http_client_pooling.py: pooled resources get an atexit # hook so long-lived sessions (e.g. a Playwright browser) are closed when the # process exits instead of idling until the TTL. From ed09bba32091aac768ab8f0b96a0af03fae5c71c Mon Sep 17 00:00:00 2001 From: pengzeyu Date: Mon, 21 Sep 2026 20:02:32 +0800 Subject: [PATCH 4/4] test(api): cover the same-key registration race in the MCP client pool --- .../core/mcp/test_client_manager.py | 49 +++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/api/tests/unit_tests/core/mcp/test_client_manager.py b/api/tests/unit_tests/core/mcp/test_client_manager.py index 1098ac8bf562c2..8c3c9f0535e04f 100644 --- a/api/tests/unit_tests/core/mcp/test_client_manager.py +++ b/api/tests/unit_tests/core/mcp/test_client_manager.py @@ -251,3 +251,52 @@ def test_close_all_closes_every_connection(self): for client in factory.side_effect: client.cleanup.assert_called_once() assert manager._clients == {} + + def test_same_key_race_neither_call_raises_and_loser_is_closed(self): + """Two callers racing on one key must both succeed on one pooled connection. + + A Barrier inside the mocked ``__enter__`` forces both callers to + connect concurrently before either registers, mirroring the race the + loser branch exists for. + """ + manager = MCPClientManager() + both_connecting = threading.Barrier(2, timeout=5) + created: list = [] + created_lock = threading.Lock() + + with patch("core.mcp.client_manager.MCPClientWithAuthRetry") as factory: + + def make_client(**_kwargs): + client = MagicMock() + + def slow_enter(): + both_connecting.wait() + return client + + client.__enter__.side_effect = slow_enter + with created_lock: + created.append(client) + return client + + factory.side_effect = make_client + + results: list = [] + errors: list = [] + + def call(): + try: + results.append(_invoke(manager)) + except Exception as exc: + errors.append(exc) + + threads = [threading.Thread(target=call) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=5) + + assert not any(thread.is_alive() for thread in threads) + assert errors == [] + assert len(manager._clients) == 1 + assert sum(client.cleanup.call_count for client in created) == 1 + assert results[0] is results[1]