diff --git a/api/core/mcp/client_manager.py b/api/core/mcp/client_manager.py new file mode 100644 index 00000000000000..d33c1771c1fad8 --- /dev/null +++ b/api/core/mcp/client_manager.py @@ -0,0 +1,310 @@ +"""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 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 +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__) + + +class _PooledClient: + def __init__(self, client: MCPClientWithAuthRetry, key: str): + self.client = client + # 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() + + +class MCPClientManager: + """A process-wide pool of authenticated MCP clients.""" + + def __init__( + self, + idle_ttl_seconds: float | None = None, + max_size: int | None = None, + ): + self._clients: dict[str, _PooledClient] = {} + self._lock = threading.Lock() + # 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, + provider_id: str | None, + timeout: float | None, + sse_read_timeout: float | None, + ) -> str: + """Identify a pooled connection by stable scope, never by credentials. + + 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, + "provider_id": provider_id, + "timeout": timeout, + "sse_read_timeout": sse_read_timeout, + }, + sort_keys=True, + ) + 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: + """Close `entry`; must be called without holding the registry lock. + + 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] + victims.append(entry) + return victims + + 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 [] + victims = [] + for key, entry in sorted(self._clients.items(), key=lambda item: item[1].last_used)[:overflow]: + del self._clients[key] + victims.append(entry) + return victims + + def _acquire( + self, + key: str, + 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.""" + with self._lock: + expired = self._evict_expired() + entry = self._clients.get(key) + if entry is not None: + entry.last_used = time.monotonic() + for victim in expired: + 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.cleanup() + except Exception: + 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 + 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. + self._close_entry(candidate) + return existing + 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, + *, + 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, + 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 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 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. + """ + 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) + except MCPAuthError: + raise + except MCPConnectionError: + logger.info("Pooled MCP connection to %s died; reconnecting once", server_url) + self._evict(entry) + entry = self._acquire( + 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 (registered as an atexit hook).""" + 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 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() + + +# 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 84235403d61cfe..c20fe81b803125 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,29 @@ 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 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: - with MCPClientWithAuthRetry( + 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, 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..8c3c9f0535e04f --- /dev/null +++ b/api/tests/unit_tests/core/mcp/test_client_manager.py @@ -0,0 +1,302 @@ +"""Unit tests for the pooled MCP client manager.""" + +import threading +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, **kwargs): + params = { + "tenant_id": "tenant-1", + "user_id": "user-1", + "server_url": "http://test.example.com/mcp", + "provider_id": "provider-1", + "headers": {"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_scope_reuses_connection(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_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, 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: + 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_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: + 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, user_id="user-a") + _invoke(manager, user_id="user-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, user_id="user-a") + _invoke(manager, user_id="user-b") + + manager.close_all() + + 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] 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"):