From 87e8e0f69a8f049f9c4edc570f18ae280915fada Mon Sep 17 00:00:00 2001 From: Novice Date: Wed, 9 Jul 2025 14:07:02 +0800 Subject: [PATCH 1/2] fix: change the mcp tool node error type --- api/core/mcp/client/sse_client.py | 2 +- api/core/mcp/client/streamable_client.py | 2 -- api/core/mcp/session/base_session.py | 5 +---- api/core/mcp/utils.py | 5 +---- api/core/tools/mcp_tool/tool.py | 9 +++++++-- 5 files changed, 10 insertions(+), 13 deletions(-) diff --git a/api/core/mcp/client/sse_client.py b/api/core/mcp/client/sse_client.py index a142a3ce48..91debcc8f9 100644 --- a/api/core/mcp/client/sse_client.py +++ b/api/core/mcp/client/sse_client.py @@ -284,7 +284,7 @@ def sse_client( try: with create_ssrf_proxy_mcp_http_client(headers=transport.headers) as client: with ssrf_proxy_sse_connect( - url, 2, timeout=httpx.Timeout(timeout, read=sse_read_timeout), client=client + url, timeout=httpx.Timeout(timeout, read=sse_read_timeout), client=client ) as event_source: event_source.response.raise_for_status() diff --git a/api/core/mcp/client/streamable_client.py b/api/core/mcp/client/streamable_client.py index ea8fe15831..fbd8d05f9e 100644 --- a/api/core/mcp/client/streamable_client.py +++ b/api/core/mcp/client/streamable_client.py @@ -185,7 +185,6 @@ class StreamableHTTPTransport: with ssrf_proxy_sse_connect( self.url, - 2, headers=headers, timeout=httpx.Timeout(self.timeout.seconds, read=self.sse_read_timeout.seconds), client=client, @@ -215,7 +214,6 @@ class StreamableHTTPTransport: with ssrf_proxy_sse_connect( self.url, - 2, headers=headers, timeout=httpx.Timeout(self.timeout.seconds, read=ctx.sse_read_timeout.seconds), client=ctx.client, diff --git a/api/core/mcp/session/base_session.py b/api/core/mcp/session/base_session.py index ac344ec395..1c0f582501 100644 --- a/api/core/mcp/session/base_session.py +++ b/api/core/mcp/session/base_session.py @@ -179,10 +179,7 @@ class BaseSession( def check_receiver_status(self) -> None: if self._receiver_future.done(): - try: - self._receiver_future.result() - except Exception as e: - raise e + self._receiver_future.result() def __exit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None diff --git a/api/core/mcp/utils.py b/api/core/mcp/utils.py index b177f34f9b..a54badcd4c 100644 --- a/api/core/mcp/utils.py +++ b/api/core/mcp/utils.py @@ -6,8 +6,6 @@ from configs import dify_config from core.mcp.types import ErrorData, JSONRPCError from core.model_runtime.utils.encoders import jsonable_encoder -SSRF_DEFAULT_MAX_RETRIES = dify_config.SSRF_DEFAULT_MAX_RETRIES - HTTP_REQUEST_NODE_SSL_VERIFY = dify_config.HTTP_REQUEST_NODE_SSL_VERIFY STATUS_FORCELIST = [429, 500, 502, 503, 504] @@ -57,7 +55,7 @@ def create_ssrf_proxy_mcp_http_client( ) -def ssrf_proxy_sse_connect(url, max_retries=SSRF_DEFAULT_MAX_RETRIES, **kwargs): +def ssrf_proxy_sse_connect(url, **kwargs): """Connect to SSE endpoint with SSRF proxy protection. This function creates an SSE connection using the configured proxy settings @@ -65,7 +63,6 @@ def ssrf_proxy_sse_connect(url, max_retries=SSRF_DEFAULT_MAX_RETRIES, **kwargs): Args: url: The SSE endpoint URL - max_retries: Maximum number of retry attempts **kwargs: Additional arguments passed to the SSE connection Returns: diff --git a/api/core/tools/mcp_tool/tool.py b/api/core/tools/mcp_tool/tool.py index 175e0133fb..da66da673b 100644 --- a/api/core/tools/mcp_tool/tool.py +++ b/api/core/tools/mcp_tool/tool.py @@ -39,14 +39,19 @@ class MCPTool(Tool): app_id: Optional[str] = None, message_id: Optional[str] = None, ) -> Generator[ToolInvokeMessage, None, None]: + from core.tools.errors import ToolInvokeError + try: with MCPClient(self.server_url, self.provider_id, self.tenant_id, authed=True) as mcp_client: tool_parameters = self._handle_none_parameter(tool_parameters) result = mcp_client.invoke_tool(tool_name=self.entity.identity.name, tool_args=tool_parameters) except MCPAuthError as e: - raise ValueError("Please auth the tool first") + raise ToolInvokeError("Please auth the tool first") from e except MCPConnectionError as e: - raise ValueError(f"Failed to connect to MCP server: {e}") + raise ToolInvokeError(f"Failed to connect to MCP server: {e}") from e + except Exception as e: + raise ToolInvokeError(f"Failed to invoke tool: {e}") from e + for content in result.content: if isinstance(content, TextContent): yield self.create_text_message(content.text) From 7b3bb1d9606bcaf4800c5e72d274729e931c36fa Mon Sep 17 00:00:00 2001 From: Novice Date: Wed, 9 Jul 2025 14:44:51 +0800 Subject: [PATCH 2/2] fix: auth api error handle --- .../console/workspace/tool_providers.py | 10 ++++++++-- api/services/tools/mcp_tools_mange_service.py | 20 ++++++++++++------- 2 files changed, 21 insertions(+), 9 deletions(-) diff --git a/api/controllers/console/workspace/tool_providers.py b/api/controllers/console/workspace/tool_providers.py index 6ac3c4b20b..df50871a38 100644 --- a/api/controllers/console/workspace/tool_providers.py +++ b/api/controllers/console/workspace/tool_providers.py @@ -12,7 +12,7 @@ from controllers.console import api from controllers.console.wraps import account_initialization_required, enterprise_license_required, setup_required from core.mcp.auth.auth_flow import auth, handle_callback from core.mcp.auth.auth_provider import OAuthClientProvider -from core.mcp.error import MCPAuthError +from core.mcp.error import MCPAuthError, MCPError from core.mcp.mcp_client import MCPClient from core.model_runtime.utils.encoders import jsonable_encoder from extensions.ext_database import db @@ -733,8 +733,14 @@ class ToolMCPAuthApi(Resource): except MCPAuthError: auth_provider = OAuthClientProvider(provider_id, tenant_id, for_list=True) - return auth(auth_provider, provider.decrypted_server_url, args["authorization_code"]) + except MCPError as e: + MCPToolManageService.update_mcp_provider_credentials( + mcp_provider=provider, + credentials={}, + authed=False, + ) + raise ValueError(f"Failed to connect to MCP server: {e}") from e class ToolMCPDetailApi(Resource): diff --git a/api/services/tools/mcp_tools_mange_service.py b/api/services/tools/mcp_tools_mange_service.py index b2a88738f6..3b1592230a 100644 --- a/api/services/tools/mcp_tools_mange_service.py +++ b/api/services/tools/mcp_tools_mange_service.py @@ -1,12 +1,13 @@ import hashlib import json from datetime import datetime +from typing import Any from sqlalchemy import or_ from sqlalchemy.exc import IntegrityError from core.helper import encrypter -from core.mcp.error import MCPAuthError, MCPConnectionError +from core.mcp.error import MCPAuthError, MCPError from core.mcp.mcp_client import MCPClient from core.tools.entities.api_entities import ToolProviderApiEntity from core.tools.entities.common_entities import I18nObject @@ -119,7 +120,7 @@ class MCPToolManageService: tools = mcp_client.list_tools() except MCPAuthError as e: raise ValueError("Please auth the tool first") - except MCPConnectionError as e: + except MCPError as e: raise ValueError(f"Failed to connect to MCP server: {e}") mcp_provider.tools = json.dumps([tool.model_dump() for tool in tools]) mcp_provider.authed = True @@ -173,7 +174,7 @@ class MCPToolManageService: server_url_hash = hashlib.sha256(server_url.encode()).hexdigest() if server_url_hash != mcp_provider.server_url_hash: - cls._re_auth_mcp_provider(mcp_provider, provider_id, tenant_id) + cls._re_connect_mcp_provider(mcp_provider, provider_id, tenant_id) mcp_provider.server_url_hash = server_url_hash try: db.session.commit() @@ -190,7 +191,9 @@ class MCPToolManageService: raise @classmethod - def update_mcp_provider_credentials(cls, mcp_provider: MCPToolProvider, credentials: dict, authed: bool = False): + def update_mcp_provider_credentials( + cls, mcp_provider: MCPToolProvider, credentials: dict[str, Any], authed: bool = False + ): provider_controller = MCPToolProviderController._from_db(mcp_provider) tool_configuration = ProviderConfigEncrypter( tenant_id=mcp_provider.tenant_id, @@ -202,11 +205,13 @@ class MCPToolManageService: mcp_provider.updated_at = datetime.now() mcp_provider.encrypted_credentials = json.dumps({**mcp_provider.credentials, **credentials}) mcp_provider.authed = authed + if not authed: + mcp_provider.tools = "[]" db.session.commit() @classmethod - def _re_auth_mcp_provider(cls, mcp_provider: MCPToolProvider, provider_id: str, tenant_id: str): - """re-auth mcp provider""" + def _re_connect_mcp_provider(cls, mcp_provider: MCPToolProvider, provider_id: str, tenant_id: str): + """re-connect mcp provider""" try: with MCPClient( mcp_provider.decrypted_server_url, @@ -221,6 +226,7 @@ class MCPToolManageService: except MCPAuthError: mcp_provider.authed = False mcp_provider.tools = "[]" - + except MCPError as e: + raise ValueError(f"Failed to re-connect MCP server: {e}") from e # reset credentials mcp_provider.encrypted_credentials = "{}"