opensource
/
dify
mirror of https://github.com/langgenius/dify.git
2
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

refactor(mcp): clean the client code

This commit is contained in:
Novice 2025-09-12 15:23:53 +08:00
parent aa44c38b58
commit f16151ea29
5 changed files with 160 additions and 131 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
api/controllers/console/workspace/tool_providers.py
View File

@ -947,15 +947,19 @@ class ToolMCPAuthApi(Resource):
provider = MCPToolManageService.get_mcp_provider_by_provider_id(provider_id, tenant_id)
if not provider:
raise ValueError("provider not found")
# headers1: if headers is provided, use it and don't need to get token
headers = provider.decrypted_headers or {}
# headers2: Add OAuth token if authed and no headers provided
if not provider.decrypted_headers and provider.authed:
token = OAuthClientProvider(provider_id, tenant_id, for_list=True).tokens()
if token:
headers["Authorization"] = f"{token.token_type.capitalize()} {token.access_token}"
try:
# try to connect to MCP server with headers
with MCPClient(
provider.decrypted_server_url,
provider_id,
tenant_id,
authed=False,
authorization_code=args["authorization_code"],
for_list=True,
headers=provider.decrypted_headers,
headers=headers,
timeout=provider.timeout,
sse_read_timeout=provider.sse_read_timeout,
):
@ -966,8 +970,11 @@ class ToolMCPAuthApi(Resource):
)
return {"result": "success"}
except MCPAuthError:
except MCPAuthError as e:
try:
if provider.decrypted_headers:
raise ValueError(f"Failed to authenticate, please check your headers: {e}") from e
# if auth failed, try to auth with OAuth or exchange token
auth_provider = OAuthClientProvider(provider_id, tenant_id, for_list=True)
return auth(auth_provider, provider.decrypted_server_url, args["authorization_code"])
except Exception as e:

1
api/core/mcp/auth/auth_flow.py
View File

@ -299,7 +299,6 @@ def auth(
server_url: str,
authorization_code: Optional[str] = None,
state_param: Optional[str] = None,
for_list: bool = False,
) -> dict[str, str]:
"""Orchestrates the full auth flow with a server using secure Redis state storage."""
metadata = discover_oauth_metadata(server_url)

113
api/core/mcp/mcp_client.py
View File

@ -2,12 +2,12 @@ import logging
from collections.abc import Callable
from contextlib import AbstractContextManager, ExitStack
from types import TracebackType
from typing import Any, Optional
from typing import Any
from urllib.parse import urlparse
from core.mcp.client.sse_client import sse_client
from core.mcp.client.streamable_client import streamablehttp_client
from core.mcp.error import MCPAuthError, MCPConnectionError
from core.mcp.error import MCPConnectionError
from core.mcp.session.client_session import ClientSession
from core.mcp.types import CallToolResult, Tool
@ -18,40 +18,18 @@ class MCPClient:
def __init__(
self,
server_url: str,
provider_id: str,
tenant_id: str,
authed: bool = True,
authorization_code: Optional[str] = None,
for_list: bool = False,
headers: Optional[dict[str, str]] = None,
timeout: Optional[float] = None,
sse_read_timeout: Optional[float] = None,
headers: dict[str, str] | None = None,
timeout: float | None = None,
sse_read_timeout: float | None = None,
):
# Initialize info
self.provider_id = provider_id
self.tenant_id = tenant_id
self.client_type = "streamable"
self.server_url = server_url
self.headers = headers or {}
self.timeout = timeout
self.sse_read_timeout = sse_read_timeout
# Authentication info
self.authed = authed
self.authorization_code = authorization_code
if authed:
from core.mcp.auth.auth_provider import OAuthClientProvider
self.provider = OAuthClientProvider(self.provider_id, self.tenant_id, for_list=for_list)
self.token = self.provider.tokens()
# Initialize session and client objects
self._session: Optional[ClientSession] = None
self._streams_context: Optional[AbstractContextManager[Any]] = None
self._session_context: Optional[ClientSession] = None
self._session: ClientSession | None = None
self._exit_stack = ExitStack()
# Whether the client has been initialized
self._initialized = False
def __enter__(self):
@ -59,9 +37,7 @@ class MCPClient:
self._initialized = True
return self
def __exit__(
self, exc_type: Optional[type], exc_value: Optional[BaseException], traceback: Optional[TracebackType]
):
def __exit__(self, exc_type: type | None, exc_value: BaseException | None, traceback: TracebackType | None):
self.cleanup()
def _initialize(
@ -87,61 +63,42 @@ class MCPClient:
logger.debug("MCP connection failed with 'sse', falling back to 'mcp' method.")
self.connect_server(streamablehttp_client, "mcp")
def connect_server(
self, client_factory: Callable[..., AbstractContextManager[Any]], method_name: str, first_try: bool = True
):
from core.mcp.auth.auth_flow import auth
def connect_server(self, client_factory: Callable[..., AbstractContextManager[Any]], method_name: str) -> None:
"""
Connect to the MCP server using streamable http or sse.
Default to streamable http.
Args:
client_factory: The client factory to use(streamablehttp_client or sse_client).
method_name: The method name to use(mcp or sse).
"""
streams_context = client_factory(
url=self.server_url,
headers=self.headers,
timeout=self.timeout,
sse_read_timeout=self.sse_read_timeout,
)
try:
headers = (
{"Authorization": f"{self.token.token_type.capitalize()} {self.token.access_token}"}
if self.authed and self.token
else self.headers
)
self._streams_context = client_factory(
url=self.server_url,
headers=headers,
timeout=self.timeout,
sse_read_timeout=self.sse_read_timeout,
)
if not self._streams_context:
raise MCPConnectionError("Failed to create connection context")
# Use exit_stack to manage context managers properly
if method_name == "mcp":
read_stream, write_stream, _ = self._exit_stack.enter_context(streams_context)
streams = (read_stream, write_stream)
else: # sse_client
streams = self._exit_stack.enter_context(streams_context)
# Use exit_stack to manage context managers properly
if method_name == "mcp":
read_stream, write_stream, _ = self._exit_stack.enter_context(self._streams_context)
streams = (read_stream, write_stream)
else: # sse_client
streams = self._exit_stack.enter_context(self._streams_context)
self._session_context = ClientSession(*streams)
self._session = self._exit_stack.enter_context(self._session_context)
self._session.initialize()
return
except MCPAuthError:
if not self.authed:
raise
try:
auth(self.provider, self.server_url, self.authorization_code)
except Exception as e:
raise ValueError(f"Failed to authenticate: {e}")
self.token = self.provider.tokens()
if first_try:
return self.connect_server(client_factory, method_name, first_try=False)
session_context = ClientSession(*streams)
self._session = self._exit_stack.enter_context(session_context)
self._session.initialize()
def list_tools(self) -> list[Tool]:
"""Connect to an MCP server running with SSE transport"""
# List available tools to verify connection
if not self._initialized or not self._session:
"""List available tools from the MCP server"""
if not self._session:
raise ValueError("Session not initialized.")
response = self._session.list_tools()
tools = response.tools
return tools
return response.tools
def invoke_tool(self, tool_name: str, tool_args: dict[str, Any]) -> CallToolResult:
"""Call a tool"""
if not self._initialized or not self._session:
if not self._session:
raise ValueError("Session not initialized.")
return self._session.call_tool(tool_name, tool_args)
@ -155,6 +112,4 @@ class MCPClient:
raise ValueError(f"Error during cleanup: {e}")
finally:
self._session = None
self._session_context = None
self._streams_context = None
self._initialized = False

89
api/core/tools/mcp_tool/tool.py
View File

@ -3,12 +3,14 @@ import json
from collections.abc import Generator
from typing import Any, Optional
from core.mcp.auth.auth_flow import auth
from core.mcp.error import MCPAuthError, MCPConnectionError
from core.mcp.mcp_client import MCPClient
from core.mcp.types import ImageContent, TextContent
from core.mcp.types import CallToolResult, ImageContent, TextContent
from core.tools.__base.tool import Tool
from core.tools.__base.tool_runtime import ToolRuntime
from core.tools.entities.tool_entities import ToolEntity, ToolInvokeMessage, ToolProviderType
from core.tools.errors import ToolInvokeError
class MCPTool(Tool):
@ -44,26 +46,7 @@ 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,
headers=self.headers,
timeout=self.timeout,
sse_read_timeout=self.sse_read_timeout,
) 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 ToolInvokeError("Please auth the tool first") from e
except MCPConnectionError as 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
result = self.invoke_remote_mcp_tool(tool_parameters)
# handle dify tool output
for content in result.content:
if isinstance(content, TextContent):
@ -95,7 +78,7 @@ class MCPTool(Tool):
def _process_json_list(self, json_list: list) -> Generator[ToolInvokeMessage, None, None]:
"""Process a list of JSON items."""
if any(not isinstance(item, dict[str, Any]) for item in json_list):
if any(not isinstance(item, dict) for item in json_list):
# If the list contains any non-dict item, treat the entire list as a text message.
yield self.create_text_message(str(json_list))
return
@ -130,3 +113,65 @@ class MCPTool(Tool):
for key, value in parameter.items()
if value is not None and not (isinstance(value, str) and value.strip() == "")
}
def invoke_remote_mcp_tool(self, tool_parameters: dict[str, Any]) -> CallToolResult:
headers = self.headers.copy() if self.headers else {}
tool_parameters = self._handle_none_parameter(tool_parameters)
# Initialize auth provider
from core.mcp.auth.auth_provider import OAuthClientProvider
provider = None
try:
provider = OAuthClientProvider(self.provider_id, self.tenant_id, for_list=False)
except Exception as e:
# If provider initialization fails, continue without auth
pass
# Try to get existing token and add to headers
if provider:
try:
token = provider.tokens()
if token:
headers["Authorization"] = f"{token.token_type.capitalize()} {token.access_token}"
except Exception:
# If token retrieval fails, continue without auth header
pass
# Define a helper function to invoke the tool
def _invoke_with_client(client_headers: dict[str, str]) -> CallToolResult:
with MCPClient(
self.server_url,
headers=client_headers,
timeout=self.timeout,
sse_read_timeout=self.sse_read_timeout,
) as mcp_client:
return mcp_client.invoke_tool(tool_name=self.entity.identity.name, tool_args=tool_parameters)
try:
# First attempt with current headers
return _invoke_with_client(headers)
except MCPAuthError as e:
# Authentication required - try to authenticate
if not provider:
raise ToolInvokeError("Authentication required but no auth provider available") from e
try:
# Perform authentication flow
auth(provider, self.server_url, None, None, False)
token = provider.tokens()
if not token:
raise ToolInvokeError("Authentication failed - no token received")
# Update headers with new token while preserving other headers
headers["Authorization"] = f"{token.token_type.capitalize()} {token.access_token}"
# Retry with authenticated headers
return _invoke_with_client(headers)
except MCPAuthError as auth_error:
raise ToolInvokeError("Authentication failed") from auth_error
except MCPConnectionError as 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

67
api/services/tools/mcp_tools_manage_service.py
View File

@ -1,7 +1,7 @@
import hashlib
import json
from datetime import datetime
from typing import Any, cast
from typing import Any
from sqlalchemy import or_
from sqlalchemy.exc import IntegrityError
@ -10,10 +10,10 @@ from core.helper import encrypter
from core.helper.provider_cache import NoOpProviderCredentialCache
from core.mcp.error import MCPAuthError, MCPError
from core.mcp.mcp_client import MCPClient
from core.mcp.types import OAuthTokens
from core.tools.entities.api_entities import ToolProviderApiEntity
from core.tools.entities.common_entities import I18nObject
from core.tools.entities.tool_entities import ToolProviderType
from core.tools.mcp_tool.provider import MCPToolProviderController
from core.tools.utils.encryption import ProviderConfigEncrypter
from extensions.ext_database import db
from models.tools import MCPToolProvider
@ -55,7 +55,25 @@ class MCPToolManageService:
cache=NoOpProviderCredentialCache(),
)
return cast(dict[str, str], encrypter_instance.encrypt(headers))
return encrypter_instance.encrypt(headers)
@staticmethod
def _retrieve_remote_mcp_tools(server_url: str, headers: dict[str, str], timeout: float, sse_read_timeout: float):
with MCPClient(
server_url,
headers=headers,
timeout=timeout,
sse_read_timeout=sse_read_timeout,
) as mcp_client:
tools = mcp_client.list_tools()
return tools
@staticmethod
def _process_headers(headers: dict[str, str], tokens: OAuthTokens | None = None):
headers = headers or {}
if tokens:
headers["Authorization"] = f"{tokens.token_type.capitalize()} {tokens.access_token}"
return headers
@staticmethod
def get_mcp_provider_by_provider_id(provider_id: str, tenant_id: str) -> MCPToolProvider:
@ -153,6 +171,9 @@ class MCPToolManageService:
@classmethod
def list_mcp_tool_from_remote_server(cls, tenant_id: str, provider_id: str) -> ToolProviderApiEntity:
from core.mcp.auth.auth_flow import auth
from core.mcp.auth.auth_provider import OAuthClientProvider
mcp_provider = cls.get_mcp_provider_by_provider_id(provider_id, tenant_id)
server_url = mcp_provider.decrypted_server_url
authed = mcp_provider.authed
@ -160,20 +181,24 @@ class MCPToolManageService:
timeout = mcp_provider.timeout
sse_read_timeout = mcp_provider.sse_read_timeout
try:
with MCPClient(
server_url,
provider_id,
tenant_id,
authed=authed,
for_list=True,
headers=headers,
timeout=timeout,
sse_read_timeout=sse_read_timeout,
) as mcp_client:
tools = mcp_client.list_tools()
except MCPAuthError:
# Handle authentication headers if authed
if not authed:
raise ValueError("Please auth the tool first")
provider = OAuthClientProvider(provider_id, tenant_id, for_list=True)
tokens = provider.tokens()
headers = cls._process_headers(headers, tokens)
try:
tools = cls._retrieve_remote_mcp_tools(server_url, headers, timeout, sse_read_timeout)
except MCPAuthError:
try:
auth(provider, server_url, None, None, False)
tokens = provider.tokens()
re_authed_headers = cls._process_headers(headers, tokens)
tools = cls._retrieve_remote_mcp_tools(server_url, re_authed_headers, timeout, sse_read_timeout)
except Exception:
raise ValueError("Please auth the tool first")
except MCPError as e:
raise ValueError(f"Failed to connect to MCP server: {e}")
@ -284,10 +309,12 @@ class MCPToolManageService:
def update_mcp_provider_credentials(
cls, mcp_provider: MCPToolProvider, credentials: dict[str, Any], authed: bool = False
):
from core.tools.mcp_tool.provider import MCPToolProviderController
provider_controller = MCPToolProviderController.from_db(mcp_provider)
tool_configuration = ProviderConfigEncrypter(
tenant_id=mcp_provider.tenant_id,
config=list(provider_controller.get_credentials_schema()), # ty: ignore [invalid-argument-type]
config=list(provider_controller.get_credentials_schema()),
provider_config_cache=NoOpProviderCredentialCache(),
)
credentials = tool_configuration.encrypt(credentials)
@ -310,7 +337,7 @@ class MCPToolManageService:
db.session.commit()
@classmethod
def _re_connect_mcp_provider(cls, server_url: str, provider_id: str, tenant_id: str):
def _re_connect_mcp_provider(cls, server_url: str, provider_id: str, tenant_id: str) -> dict[str, Any]:
# Get the existing provider to access headers and timeout settings
mcp_provider = cls.get_mcp_provider_by_provider_id(provider_id, tenant_id)
headers = mcp_provider.decrypted_headers
@ -320,10 +347,6 @@ class MCPToolManageService:
try:
with MCPClient(
server_url,
provider_id,
tenant_id,
authed=False,
for_list=True,
headers=headers,
timeout=timeout,
sse_read_timeout=sse_read_timeout,