From 87e8e0f69a8f049f9c4edc570f18ae280915fada Mon Sep 17 00:00:00 2001 From: Novice Date: Wed, 9 Jul 2025 14:07:02 +0800 Subject: [PATCH] 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)