Merge remote-tracking branch 'origin/main' into feat/tool-plugin-oauth

# Conflicts:
#	api/controllers/console/workspace/tool_providers.py
#	api/core/tools/entities/api_entities.py
#	api/core/tools/tool_manager.py
#	api/core/tools/utils/configuration.py
#	api/services/tools/tools_transform_service.py

api/.env.example

@@ -17,6 +17,11 @@ APP_WEB_URL=http://127.0.0.1:3000
 # Files URL
 FILES_URL=http://127.0.0.1:5001
 
+# INTERNAL_FILES_URL is used for plugin daemon communication within Docker network.
+# Set this to the internal Docker service URL for proper plugin file access.
+# Example: INTERNAL_FILES_URL=http://api:5001
+INTERNAL_FILES_URL=http://127.0.0.1:5001
+
 # The time in seconds after the signature is rejected
 FILES_ACCESS_TIMEOUT=300
 

api/configs/feature/__init__.py

@@ -237,6 +237,13 @@ class FileAccessConfig(BaseSettings):
         default="",
     )
 
+    INTERNAL_FILES_URL: str = Field(
+        description="Internal base URL for file access within Docker network,"
+        " used for plugin daemon and internal service communication."
+        " Falls back to FILES_URL if not specified.",
+        default="",
+    )
+
     FILES_ACCESS_TIMEOUT: int = Field(
         description="Expiration time in seconds for file access URLs",
         default=300,

api/configs/packaging/__init__.py

@@ -8,11 +8,6 @@ class PackagingInfo(PyProjectTomlConfig):
     Packaging build information
     """
 
-    CURRENT_VERSION: str = Field(
-        description="Dify version",
-        default="1.5.0",
-    )
-
     COMMIT_SHA: str = Field(
         description="SHA-1 checksum of the git commit used to build the app",
         default="",

api/controllers/console/__init__.py

@@ -56,6 +56,7 @@ from .app import (
     conversation,
     conversation_variables,
     generator,
+    mcp_server,
     message,
     model_config,
     ops_trace,

api/controllers/console/app/audio.py

@@ -90,23 +90,11 @@ class ChatMessageTextApi(Resource):
 
             message_id = args.get("message_id", None)
             text = args.get("text", None)
-            if (
-                app_model.mode in {AppMode.ADVANCED_CHAT.value, AppMode.WORKFLOW.value}
-                and app_model.workflow
-                and app_model.workflow.features_dict
-            ):
-                text_to_speech = app_model.workflow.features_dict.get("text_to_speech")
-                if text_to_speech is None:
-                    raise ValueError("TTS is not enabled")
-                voice = args.get("voice") or text_to_speech.get("voice")
-            else:
-                try:
-                    if app_model.app_model_config is None:
-                        raise ValueError("AppModelConfig not found")
-                    voice = args.get("voice") or app_model.app_model_config.text_to_speech_dict.get("voice")
-                except Exception:
-                    voice = None
-            response = AudioService.transcript_tts(app_model=app_model, text=text, message_id=message_id, voice=voice)
+            voice = args.get("voice", None)
+
+            response = AudioService.transcript_tts(
+                app_model=app_model, text=text, voice=voice, message_id=message_id, is_draft=True
+            )
             return response
         except services.errors.app_model_config.AppModelConfigBrokenError:
             logging.exception("App model config broken.")

api/controllers/console/app/mcp_server.py

@@ -0,0 +1,106 @@
+import json
+from enum import StrEnum
+
+from flask_login import current_user
+from flask_restful import Resource, marshal_with, reqparse
+from werkzeug.exceptions import NotFound
+
+from controllers.console import api
+from controllers.console.app.wraps import get_app_model
+from controllers.console.wraps import account_initialization_required, setup_required
+from extensions.ext_database import db
+from fields.app_fields import app_server_fields
+from libs.login import login_required
+from models.model import AppMCPServer
+
+
+class AppMCPServerStatus(StrEnum):
+    ACTIVE = "active"
+    INACTIVE = "inactive"
+
+
+class AppMCPServerController(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @get_app_model
+    @marshal_with(app_server_fields)
+    def get(self, app_model):
+        server = db.session.query(AppMCPServer).filter(AppMCPServer.app_id == app_model.id).first()
+        return server
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @get_app_model
+    @marshal_with(app_server_fields)
+    def post(self, app_model):
+        # The role of the current user in the ta table must be editor, admin, or owner
+        if not current_user.is_editor:
+            raise NotFound()
+        parser = reqparse.RequestParser()
+        parser.add_argument("description", type=str, required=True, location="json")
+        parser.add_argument("parameters", type=dict, required=True, location="json")
+        args = parser.parse_args()
+        server = AppMCPServer(
+            name=app_model.name,
+            description=args["description"],
+            parameters=json.dumps(args["parameters"], ensure_ascii=False),
+            status=AppMCPServerStatus.ACTIVE,
+            app_id=app_model.id,
+            tenant_id=current_user.current_tenant_id,
+            server_code=AppMCPServer.generate_server_code(16),
+        )
+        db.session.add(server)
+        db.session.commit()
+        return server
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @get_app_model
+    @marshal_with(app_server_fields)
+    def put(self, app_model):
+        if not current_user.is_editor:
+            raise NotFound()
+        parser = reqparse.RequestParser()
+        parser.add_argument("id", type=str, required=True, location="json")
+        parser.add_argument("description", type=str, required=True, location="json")
+        parser.add_argument("parameters", type=dict, required=True, location="json")
+        parser.add_argument("status", type=str, required=False, location="json")
+        args = parser.parse_args()
+        server = db.session.query(AppMCPServer).filter(AppMCPServer.id == args["id"]).first()
+        if not server:
+            raise NotFound()
+        server.description = args["description"]
+        server.parameters = json.dumps(args["parameters"], ensure_ascii=False)
+        if args["status"]:
+            if args["status"] not in [status.value for status in AppMCPServerStatus]:
+                raise ValueError("Invalid status")
+            server.status = args["status"]
+        db.session.commit()
+        return server
+
+
+class AppMCPServerRefreshController(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(app_server_fields)
+    def get(self, server_id):
+        if not current_user.is_editor:
+            raise NotFound()
+        server = (
+            db.session.query(AppMCPServer)
+            .filter(AppMCPServer.id == server_id and AppMCPServer.tenant_id == current_user.current_tenant_id)
+            .first()
+        )
+        if not server:
+            raise NotFound()
+        server.server_code = AppMCPServer.generate_server_code(16)
+        db.session.commit()
+        return server
+
+
+api.add_resource(AppMCPServerController, "/apps/<uuid:app_id>/server")
+api.add_resource(AppMCPServerRefreshController, "/apps/<uuid:server_id>/server/refresh")

api/controllers/console/explore/audio.py

@@ -18,7 +18,6 @@ from controllers.console.app.error import (
 from controllers.console.explore.wraps import InstalledAppResource
 from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
 from core.model_runtime.errors.invoke import InvokeError
-from models.model import AppMode
 from services.audio_service import AudioService
 from services.errors.audio import (
     AudioTooLargeServiceError,
@@ -79,19 +78,9 @@ class ChatTextApi(InstalledAppResource):
 
             message_id = args.get("message_id", None)
             text = args.get("text", None)
-            if (
-                app_model.mode in {AppMode.ADVANCED_CHAT.value, AppMode.WORKFLOW.value}
-                and app_model.workflow
-                and app_model.workflow.features_dict
-            ):
-                text_to_speech = app_model.workflow.features_dict.get("text_to_speech")
-                voice = args.get("voice") or text_to_speech.get("voice")
-            else:
-                try:
-                    voice = args.get("voice") or app_model.app_model_config.text_to_speech_dict.get("voice")
-                except Exception:
-                    voice = None
-            response = AudioService.transcript_tts(app_model=app_model, message_id=message_id, voice=voice, text=text)
+            voice = args.get("voice", None)
+
+            response = AudioService.transcript_tts(app_model=app_model, text=text, voice=voice, message_id=message_id)
             return response
         except services.errors.app_model_config.AppModelConfigBrokenError:
             logging.exception("App model config broken.")

api/controllers/console/workspace/tool_providers.py

@@ -1,5 +1,7 @@
 import io
+from urllib.parse import urlparse
 
+from flask import redirect, send_file
 from flask import make_response, redirect, request, send_file
 from flask_login import current_user
 from flask_restful import (
@@ -16,6 +18,11 @@ from controllers.console.wraps import (
     enterprise_license_required,
     setup_required,
 )
+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, MCPError
+from core.mcp.mcp_client import MCPClient
 from core.model_runtime.utils.encoders import jsonable_encoder
 from core.plugin.entities.plugin import ToolProviderID
 from core.plugin.impl.oauth import OAuthHandler
@@ -26,11 +33,24 @@ from libs.login import login_required
 from services.plugin.oauth_service import OAuthProxyService
 from services.tools.api_tools_manage_service import ApiToolManageService
 from services.tools.builtin_tools_manage_service import BuiltinToolManageService
+from services.tools.mcp_tools_mange_service import MCPToolManageService
 from services.tools.tool_labels_service import ToolLabelsService
 from services.tools.tools_manage_service import ToolCommonService
+from services.tools.tools_transform_service import ToolTransformService
 from services.tools.workflow_tools_manage_service import WorkflowToolManageService
 
 
+def is_valid_url(url: str) -> bool:
+    if not url:
+        return False
+
+    try:
+        parsed = urlparse(url)
+        return all([parsed.scheme, parsed.netloc]) and parsed.scheme in ["http", "https"]
+    except Exception:
+        return False
+
+
 class ToolProviderListApi(Resource):
     @setup_required
     @login_required
@@ -45,7 +65,7 @@ class ToolProviderListApi(Resource):
         req.add_argument(
             "type",
             type=str,
-            choices=["builtin", "model", "api", "workflow"],
+            choices=["builtin", "model", "api", "workflow", "mcp"],
             required=False,
             nullable=True,
             location="args",
@@ -831,6 +851,166 @@ api.add_resource(ToolOAuthCallback, "/oauth/plugin/<path:provider>/tool/callback
 
 api.add_resource(ToolOAuthCustomClient, "/workspaces/current/tool-provider/builtin/<path:provider>/oauth/custom-client")
 
+class ToolProviderMCPApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument("server_url", type=str, required=True, nullable=False, location="json")
+        parser.add_argument("name", type=str, required=True, nullable=False, location="json")
+        parser.add_argument("icon", type=str, required=True, nullable=False, location="json")
+        parser.add_argument("icon_type", type=str, required=True, nullable=False, location="json")
+        parser.add_argument("icon_background", type=str, required=False, nullable=True, location="json", default="")
+        parser.add_argument("server_identifier", type=str, required=True, nullable=False, location="json")
+        args = parser.parse_args()
+        user = current_user
+        if not is_valid_url(args["server_url"]):
+            raise ValueError("Server URL is not valid.")
+        return jsonable_encoder(
+            MCPToolManageService.create_mcp_provider(
+                tenant_id=user.current_tenant_id,
+                server_url=args["server_url"],
+                name=args["name"],
+                icon=args["icon"],
+                icon_type=args["icon_type"],
+                icon_background=args["icon_background"],
+                user_id=user.id,
+                server_identifier=args["server_identifier"],
+            )
+        )
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def put(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument("server_url", type=str, required=True, nullable=False, location="json")
+        parser.add_argument("name", type=str, required=True, nullable=False, location="json")
+        parser.add_argument("icon", type=str, required=True, nullable=False, location="json")
+        parser.add_argument("icon_type", type=str, required=True, nullable=False, location="json")
+        parser.add_argument("icon_background", type=str, required=False, nullable=True, location="json")
+        parser.add_argument("provider_id", type=str, required=True, nullable=False, location="json")
+        parser.add_argument("server_identifier", type=str, required=True, nullable=False, location="json")
+        args = parser.parse_args()
+        if not is_valid_url(args["server_url"]):
+            if "[__HIDDEN__]" in args["server_url"]:
+                pass
+            else:
+                raise ValueError("Server URL is not valid.")
+        MCPToolManageService.update_mcp_provider(
+            tenant_id=current_user.current_tenant_id,
+            provider_id=args["provider_id"],
+            server_url=args["server_url"],
+            name=args["name"],
+            icon=args["icon"],
+            icon_type=args["icon_type"],
+            icon_background=args["icon_background"],
+            server_identifier=args["server_identifier"],
+        )
+        return {"result": "success"}
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def delete(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument("provider_id", type=str, required=True, nullable=False, location="json")
+        args = parser.parse_args()
+        MCPToolManageService.delete_mcp_tool(tenant_id=current_user.current_tenant_id, provider_id=args["provider_id"])
+        return {"result": "success"}
+
+
+class ToolMCPAuthApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument("provider_id", type=str, required=True, nullable=False, location="json")
+        parser.add_argument("authorization_code", type=str, required=False, nullable=True, location="json")
+        args = parser.parse_args()
+        provider_id = args["provider_id"]
+        tenant_id = current_user.current_tenant_id
+        provider = MCPToolManageService.get_mcp_provider_by_provider_id(provider_id, tenant_id)
+        if not provider:
+            raise ValueError("provider not found")
+        try:
+            with MCPClient(
+                provider.decrypted_server_url,
+                provider_id,
+                tenant_id,
+                authed=False,
+                authorization_code=args["authorization_code"],
+                for_list=True,
+            ):
+                MCPToolManageService.update_mcp_provider_credentials(
+                    mcp_provider=provider,
+                    credentials=provider.decrypted_credentials,
+                    authed=True,
+                )
+                return {"result": "success"}
+
+        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):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, provider_id):
+        user = current_user
+        provider = MCPToolManageService.get_mcp_provider_by_provider_id(provider_id, user.current_tenant_id)
+        return jsonable_encoder(ToolTransformService.mcp_provider_to_user_provider(provider, for_list=True))
+
+
+class ToolMCPListAllApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self):
+        user = current_user
+        tenant_id = user.current_tenant_id
+
+        tools = MCPToolManageService.retrieve_mcp_tools(tenant_id=tenant_id)
+
+        return [tool.to_dict() for tool in tools]
+
+
+class ToolMCPUpdateApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, provider_id):
+        tenant_id = current_user.current_tenant_id
+        tools = MCPToolManageService.list_mcp_tool_from_remote_server(
+            tenant_id=tenant_id,
+            provider_id=provider_id,
+        )
+        return jsonable_encoder(tools)
+
+
+class ToolMCPCallbackApi(Resource):
+    def get(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument("code", type=str, required=True, nullable=False, location="args")
+        parser.add_argument("state", type=str, required=True, nullable=False, location="args")
+        args = parser.parse_args()
+        state_key = args["state"]
+        authorization_code = args["code"]
+        handle_callback(state_key, authorization_code)
+        return redirect(f"{dify_config.CONSOLE_WEB_URL}/oauth-callback")
+
+
 # tool provider
 api.add_resource(ToolProviderListApi, "/workspaces/current/tool-providers")
 
@@ -876,8 +1056,15 @@ api.add_resource(ToolWorkflowProviderDeleteApi, "/workspaces/current/tool-provid
 api.add_resource(ToolWorkflowProviderGetApi, "/workspaces/current/tool-provider/workflow/get")
 api.add_resource(ToolWorkflowProviderListToolApi, "/workspaces/current/tool-provider/workflow/tools")
 
+# mcp tool provider
+api.add_resource(ToolMCPDetailApi, "/workspaces/current/tool-provider/mcp/tools/<path:provider_id>")
+api.add_resource(ToolProviderMCPApi, "/workspaces/current/tool-provider/mcp")
+api.add_resource(ToolMCPUpdateApi, "/workspaces/current/tool-provider/mcp/update/<path:provider_id>")
+api.add_resource(ToolMCPAuthApi, "/workspaces/current/tool-provider/mcp/auth")
+api.add_resource(ToolMCPCallbackApi, "/mcp/oauth/callback")
+
 api.add_resource(ToolBuiltinListApi, "/workspaces/current/tools/builtin")
 api.add_resource(ToolApiListApi, "/workspaces/current/tools/api")
+api.add_resource(ToolMCPListAllApi, "/workspaces/current/tools/mcp")
 api.add_resource(ToolWorkflowListApi, "/workspaces/current/tools/workflow")
-
 api.add_resource(ToolLabelsApi, "/workspaces/current/tool-labels")

api/controllers/files/upload.py

@@ -87,7 +87,5 @@ class PluginUploadFileApi(Resource):
         except services.errors.file.UnsupportedFileTypeError:
             raise UnsupportedFileTypeError()
 
-        return tool_file, 201
-
 
 api.add_resource(PluginUploadFileApi, "/files/upload/for-plugin")

api/controllers/mcp/__init__.py

@@ -0,0 +1,8 @@
+from flask import Blueprint
+
+from libs.external_api import ExternalApi
+
+bp = Blueprint("mcp", __name__, url_prefix="/mcp")
+api = ExternalApi(bp)
+
+from . import mcp

api/controllers/mcp/mcp.py

@@ -0,0 +1,104 @@
+from flask_restful import Resource, reqparse
+from pydantic import ValidationError
+
+from controllers.console.app.mcp_server import AppMCPServerStatus
+from controllers.mcp import api
+from core.app.app_config.entities import VariableEntity
+from core.mcp import types
+from core.mcp.server.streamable_http import MCPServerStreamableHTTPRequestHandler
+from core.mcp.types import ClientNotification, ClientRequest
+from core.mcp.utils import create_mcp_error_response
+from extensions.ext_database import db
+from libs import helper
+from models.model import App, AppMCPServer, AppMode
+
+
+class MCPAppApi(Resource):
+    def post(self, server_code):
+        def int_or_str(value):
+            if isinstance(value, (int, str)):
+                return value
+            else:
+                return None
+
+        parser = reqparse.RequestParser()
+        parser.add_argument("jsonrpc", type=str, required=True, location="json")
+        parser.add_argument("method", type=str, required=True, location="json")
+        parser.add_argument("params", type=dict, required=False, location="json")
+        parser.add_argument("id", type=int_or_str, required=False, location="json")
+        args = parser.parse_args()
+
+        request_id = args.get("id")
+
+        server = db.session.query(AppMCPServer).filter(AppMCPServer.server_code == server_code).first()
+        if not server:
+            return helper.compact_generate_response(
+                create_mcp_error_response(request_id, types.INVALID_REQUEST, "Server Not Found")
+            )
+
+        if server.status != AppMCPServerStatus.ACTIVE:
+            return helper.compact_generate_response(
+                create_mcp_error_response(request_id, types.INVALID_REQUEST, "Server is not active")
+            )
+
+        app = db.session.query(App).filter(App.id == server.app_id).first()
+        if not app:
+            return helper.compact_generate_response(
+                create_mcp_error_response(request_id, types.INVALID_REQUEST, "App Not Found")
+            )
+
+        if app.mode in {AppMode.ADVANCED_CHAT.value, AppMode.WORKFLOW.value}:
+            workflow = app.workflow
+            if workflow is None:
+                return helper.compact_generate_response(
+                    create_mcp_error_response(request_id, types.INVALID_REQUEST, "App is unavailable")
+                )
+
+            user_input_form = workflow.user_input_form(to_old_structure=True)
+        else:
+            app_model_config = app.app_model_config
+            if app_model_config is None:
+                return helper.compact_generate_response(
+                    create_mcp_error_response(request_id, types.INVALID_REQUEST, "App is unavailable")
+                )
+
+            features_dict = app_model_config.to_dict()
+            user_input_form = features_dict.get("user_input_form", [])
+        converted_user_input_form: list[VariableEntity] = []
+        try:
+            for item in user_input_form:
+                variable_type = item.get("type", "") or list(item.keys())[0]
+                variable = item[variable_type]
+                converted_user_input_form.append(
+                    VariableEntity(
+                        type=variable_type,
+                        variable=variable.get("variable"),
+                        description=variable.get("description") or "",
+                        label=variable.get("label"),
+                        required=variable.get("required", False),
+                        max_length=variable.get("max_length"),
+                        options=variable.get("options") or [],
+                    )
+                )
+        except ValidationError as e:
+            return helper.compact_generate_response(
+                create_mcp_error_response(request_id, types.INVALID_PARAMS, f"Invalid user_input_form: {str(e)}")
+            )
+
+        try:
+            request: ClientRequest | ClientNotification = ClientRequest.model_validate(args)
+        except ValidationError as e:
+            try:
+                notification = ClientNotification.model_validate(args)
+                request = notification
+            except ValidationError as e:
+                return helper.compact_generate_response(
+                    create_mcp_error_response(request_id, types.INVALID_PARAMS, f"Invalid MCP request: {str(e)}")
+                )
+
+        mcp_server_handler = MCPServerStreamableHTTPRequestHandler(app, request, converted_user_input_form)
+        response = mcp_server_handler.handle()
+        return helper.compact_generate_response(response)
+
+
+api.add_resource(MCPAppApi, "/server/<string:server_code>/mcp")

api/controllers/service_api/app/audio.py

@@ -20,7 +20,7 @@ from controllers.service_api.app.error import (
 from controllers.service_api.wraps import FetchUserArg, WhereisUserArg, validate_app_token
 from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
 from core.model_runtime.errors.invoke import InvokeError
-from models.model import App, AppMode, EndUser
+from models.model import App, EndUser
 from services.audio_service import AudioService
 from services.errors.audio import (
     AudioTooLargeServiceError,
@@ -78,20 +78,9 @@ class TextApi(Resource):
 
             message_id = args.get("message_id", None)
             text = args.get("text", None)
-            if (
-                app_model.mode in {AppMode.ADVANCED_CHAT.value, AppMode.WORKFLOW.value}
-                and app_model.workflow
-                and app_model.workflow.features_dict
-            ):
-                text_to_speech = app_model.workflow.features_dict.get("text_to_speech", {})
-                voice = args.get("voice") or text_to_speech.get("voice")
-            else:
-                try:
-                    voice = args.get("voice") or app_model.app_model_config.text_to_speech_dict.get("voice")
-                except Exception:
-                    voice = None
+            voice = args.get("voice", None)
             response = AudioService.transcript_tts(
-                app_model=app_model, message_id=message_id, end_user=end_user.external_user_id, voice=voice, text=text
+                app_model=app_model, text=text, voice=voice, end_user=end_user.external_user_id, message_id=message_id
             )
 
             return response

api/controllers/service_api/dataset/document.py

@@ -211,6 +211,9 @@ class DocumentAddByFileApi(DatasetApiResource):
         if not dataset:
             raise ValueError("Dataset does not exist.")
 
+        if dataset.provider == "external":
+            raise ValueError("External datasets are not supported.")
+
         indexing_technique = args.get("indexing_technique") or dataset.indexing_technique
         if not indexing_technique:
             raise ValueError("indexing_technique is required.")
@@ -301,6 +304,9 @@ class DocumentUpdateByFileApi(DatasetApiResource):
         if not dataset:
             raise ValueError("Dataset does not exist.")
 
+        if dataset.provider == "external":
+            raise ValueError("External datasets are not supported.")
+
         # indexing_technique is already set in dataset since this is an update
         args["indexing_technique"] = dataset.indexing_technique
 

api/controllers/web/audio.py

@@ -19,7 +19,7 @@ from controllers.web.error import (
 from controllers.web.wraps import WebApiResource
 from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
 from core.model_runtime.errors.invoke import InvokeError
-from models.model import App, AppMode
+from models.model import App
 from services.audio_service import AudioService
 from services.errors.audio import (
     AudioTooLargeServiceError,
@@ -77,21 +77,9 @@ class TextApi(WebApiResource):
 
             message_id = args.get("message_id", None)
             text = args.get("text", None)
-            if (
-                app_model.mode in {AppMode.ADVANCED_CHAT.value, AppMode.WORKFLOW.value}
-                and app_model.workflow
-                and app_model.workflow.features_dict
-            ):
-                text_to_speech = app_model.workflow.features_dict.get("text_to_speech", {})
-                voice = args.get("voice") or text_to_speech.get("voice")
-            else:
-                try:
-                    voice = args.get("voice") or app_model.app_model_config.text_to_speech_dict.get("voice")
-                except Exception:
-                    voice = None
-
+            voice = args.get("voice", None)
             response = AudioService.transcript_tts(
-                app_model=app_model, message_id=message_id, end_user=end_user.external_user_id, voice=voice, text=text
+                app_model=app_model, text=text, voice=voice, end_user=end_user.external_user_id, message_id=message_id
             )
 
             return response

api/core/agent/base_agent_runner.py

@@ -161,10 +161,14 @@ class BaseAgentRunner(AppRunner):
             if parameter.type == ToolParameter.ToolParameterType.SELECT:
                 enum = [option.value for option in parameter.options] if parameter.options else []
 
-            message_tool.parameters["properties"][parameter.name] = {
-                "type": parameter_type,
-                "description": parameter.llm_description or "",
-            }
+            message_tool.parameters["properties"][parameter.name] = (
+                {
+                    "type": parameter_type,
+                    "description": parameter.llm_description or "",
+                }
+                if parameter.input_schema is None
+                else parameter.input_schema
+            )
 
             if len(enum) > 0:
                 message_tool.parameters["properties"][parameter.name]["enum"] = enum
@@ -254,10 +258,14 @@ class BaseAgentRunner(AppRunner):
             if parameter.type == ToolParameter.ToolParameterType.SELECT:
                 enum = [option.value for option in parameter.options] if parameter.options else []
 
-            prompt_tool.parameters["properties"][parameter.name] = {
-                "type": parameter_type,
-                "description": parameter.llm_description or "",
-            }
+            prompt_tool.parameters["properties"][parameter.name] = (
+                {
+                    "type": parameter_type,
+                    "description": parameter.llm_description or "",
+                }
+                if parameter.input_schema is None
+                else parameter.input_schema
+            )
 
             if len(enum) > 0:
                 prompt_tool.parameters["properties"][parameter.name]["enum"] = enum

api/core/agent/plugin_entities.py

@@ -85,7 +85,7 @@ class AgentStrategyEntity(BaseModel):
     description: I18nObject = Field(..., description="The description of the agent strategy")
     output_schema: Optional[dict] = None
     features: Optional[list[AgentFeature]] = None
-
+    meta_version: Optional[str] = None
     # pydantic configs
     model_config = ConfigDict(protected_namespaces=())
 

api/core/agent/strategy/plugin.py

@@ -15,10 +15,12 @@ class PluginAgentStrategy(BaseAgentStrategy):
 
     tenant_id: str
     declaration: AgentStrategyEntity
+    meta_version: str | None = None
 
-    def __init__(self, tenant_id: str, declaration: AgentStrategyEntity):
+    def __init__(self, tenant_id: str, declaration: AgentStrategyEntity, meta_version: str | None):
         self.tenant_id = tenant_id
         self.declaration = declaration
+        self.meta_version = meta_version
 
     def get_parameters(self) -> Sequence[AgentStrategyParameter]:
         return self.declaration.parameters

api/core/app/task_pipeline/based_generate_task_pipeline.py

@@ -19,6 +19,7 @@ from core.app.entities.task_entities import (
 from core.errors.error import QuotaExceededError
 from core.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError
 from core.moderation.output_moderation import ModerationRule, OutputModeration
+from models.enums import MessageStatus
 from models.model import Message
 
 logger = logging.getLogger(__name__)
@@ -62,7 +63,7 @@ class BasedGenerateTaskPipeline:
             return err
 
         err_desc = self._error_to_desc(err)
-        message.status = "error"
+        message.status = MessageStatus.ERROR
         message.error = err_desc
         return err
 

api/core/entities/parameter_entities.py

@@ -21,6 +21,9 @@ class CommonParameterType(StrEnum):
     DYNAMIC_SELECT = "dynamic-select"
 
     # TOOL_SELECTOR = "tool-selector"
+    # MCP object and array type parameters
+    ARRAY = "array"
+    OBJECT = "object"
 
 
 class AppSelectorScope(StrEnum):

api/core/file/helpers.py

@@ -21,7 +21,9 @@ def get_signed_file_url(upload_file_id: str) -> str:
 
 
 def get_signed_file_url_for_plugin(filename: str, mimetype: str, tenant_id: str, user_id: str) -> str:
-    url = f"{dify_config.FILES_URL}/files/upload/for-plugin"
+    # Plugin access should use internal URL for Docker network communication
+    base_url = dify_config.INTERNAL_FILES_URL or dify_config.FILES_URL
+    url = f"{base_url}/files/upload/for-plugin"
 
     if user_id is None:
         user_id = "DEFAULT-USER"

api/core/file/models.py

@@ -51,7 +51,7 @@ class File(BaseModel):
     # It should be set to `ToolFile.id` when `transfer_method` is `tool_file`.
     related_id: Optional[str] = None
     filename: Optional[str] = None
-    extension: Optional[str] = Field(default=None, description="File extension, should contains dot")
+    extension: Optional[str] = Field(default=None, description="File extension, should contain dot")
     mime_type: Optional[str] = None
     size: int = -1
 

api/core/file/upload_file_parser.py

@@ -1,67 +0,0 @@
-import base64
-import logging
-import time
-from typing import Optional
-
-from configs import dify_config
-from constants import IMAGE_EXTENSIONS
-from core.helper.url_signer import UrlSigner
-from extensions.ext_storage import storage
-
-
-class UploadFileParser:
-    @classmethod
-    def get_image_data(cls, upload_file, force_url: bool = False) -> Optional[str]:
-        if not upload_file:
-            return None
-
-        if upload_file.extension not in IMAGE_EXTENSIONS:
-            return None
-
-        if dify_config.MULTIMODAL_SEND_FORMAT == "url" or force_url:
-            return cls.get_signed_temp_image_url(upload_file.id)
-        else:
-            # get image file base64
-            try:
-                data = storage.load(upload_file.key)
-            except FileNotFoundError:
-                logging.exception(f"File not found: {upload_file.key}")
-                return None
-
-            encoded_string = base64.b64encode(data).decode("utf-8")
-            return f"data:{upload_file.mime_type};base64,{encoded_string}"
-
-    @classmethod
-    def get_signed_temp_image_url(cls, upload_file_id) -> str:
-        """
-        get signed url from upload file
-
-        :param upload_file_id: the id of UploadFile object
-        :return:
-        """
-        base_url = dify_config.FILES_URL
-        image_preview_url = f"{base_url}/files/{upload_file_id}/image-preview"
-
-        return UrlSigner.get_signed_url(url=image_preview_url, sign_key=upload_file_id, prefix="image-preview")
-
-    @classmethod
-    def verify_image_file_signature(cls, upload_file_id: str, timestamp: str, nonce: str, sign: str) -> bool:
-        """
-        verify signature
-
-        :param upload_file_id: file id
-        :param timestamp: timestamp
-        :param nonce: nonce
-        :param sign: signature
-        :return:
-        """
-        result = UrlSigner.verify(
-            sign_key=upload_file_id, timestamp=timestamp, nonce=nonce, sign=sign, prefix="image-preview"
-        )
-
-        # verify signature
-        if not result:
-            return False
-
-        current_time = int(time.time())
-        return current_time - int(timestamp) <= dify_config.FILES_ACCESS_TIMEOUT

api/core/helper/code_executor/template_transformer.py

@@ -28,7 +28,7 @@ class TemplateTransformer(ABC):
     def extract_result_str_from_response(cls, response: str):
         result = re.search(rf"{cls._result_tag}(.*){cls._result_tag}", response, re.DOTALL)
         if not result:
-            raise ValueError("Failed to parse result")
+            raise ValueError(f"Failed to parse result: no result tag found in response. Response: {response[:200]}...")
         return result.group(1)
 
     @classmethod
@@ -38,16 +38,53 @@ class TemplateTransformer(ABC):
         :param response: response
         :return:
         """
+
         try:
-            result = json.loads(cls.extract_result_str_from_response(response))
-        except json.JSONDecodeError:
-            raise ValueError("failed to parse response")
+            result_str = cls.extract_result_str_from_response(response)
+            result = json.loads(result_str)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Failed to parse JSON response: {str(e)}. Response content: {result_str[:200]}...")
+        except ValueError as e:
+            # Re-raise ValueError from extract_result_str_from_response
+            raise e
+        except Exception as e:
+            raise ValueError(f"Unexpected error during response transformation: {str(e)}")
+
+        # Check if the result contains an error
+        if isinstance(result, dict) and "error" in result:
+            raise ValueError(f"JavaScript execution error: {result['error']}")
+
         if not isinstance(result, dict):
-            raise ValueError("result must be a dict")
+            raise ValueError(f"Result must be a dict, got {type(result).__name__}")
         if not all(isinstance(k, str) for k in result):
-            raise ValueError("result keys must be strings")
+            raise ValueError("Result keys must be strings")
+
+        # Post-process the result to convert scientific notation strings back to numbers
+        result = cls._post_process_result(result)
         return result
 
+    @classmethod
+    def _post_process_result(cls, result: dict[Any, Any]) -> dict[Any, Any]:
+        """
+        Post-process the result to convert scientific notation strings back to numbers
+        """
+
+        def convert_scientific_notation(value):
+            if isinstance(value, str):
+                # Check if the string looks like scientific notation
+                if re.match(r"^-?\d+\.?\d*e[+-]\d+$", value, re.IGNORECASE):
+                    try:
+                        return float(value)
+                    except ValueError:
+                        pass
+            elif isinstance(value, dict):
+                return {k: convert_scientific_notation(v) for k, v in value.items()}
+            elif isinstance(value, list):
+                return [convert_scientific_notation(v) for v in value]
+            return value
+
+        return convert_scientific_notation(result)  # type: ignore[no-any-return]
+
     @classmethod
     @abstractmethod
     def get_runner_script(cls) -> str:

api/core/helper/lru_cache.py

@@ -1,22 +0,0 @@
-from collections import OrderedDict
-from typing import Any
-
-
-class LRUCache:
-    def __init__(self, capacity: int):
-        self.cache: OrderedDict[Any, Any] = OrderedDict()
-        self.capacity = capacity
-
-    def get(self, key: Any) -> Any:
-        if key not in self.cache:
-            return None
-        else:
-            self.cache.move_to_end(key)  # move the key to the end of the OrderedDict
-            return self.cache[key]
-
-    def put(self, key: Any, value: Any) -> None:
-        if key in self.cache:
-            self.cache.move_to_end(key)
-        self.cache[key] = value
-        if len(self.cache) > self.capacity:
-            self.cache.popitem(last=False)  # pop the first item

api/core/indexing_runner.py

@@ -317,9 +317,10 @@ class IndexingRunner:
                 image_upload_file_ids = get_image_upload_file_ids(document.page_content)
                 for upload_file_id in image_upload_file_ids:
                     image_file = db.session.query(UploadFile).filter(UploadFile.id == upload_file_id).first()
+                    if image_file is None:
+                        continue
                     try:
-                        if image_file:
-                            storage.delete(image_file.key)
+                        storage.delete(image_file.key)
                     except Exception:
                         logging.exception(
                             "Delete image_files failed while indexing_estimate, \

api/core/llm_generator/output_parser/structured_output.py

@@ -23,6 +23,7 @@ from core.model_runtime.entities.message_entities import (
     PromptMessage,
     PromptMessageTool,
     SystemPromptMessage,
+    TextPromptMessageContent,
 )
 from core.model_runtime.entities.model_entities import AIModelEntity, ParameterRule
 
@@ -170,10 +171,15 @@ def invoke_llm_with_structured_output(
             system_fingerprint: Optional[str] = None
             for event in llm_result:
                 if isinstance(event, LLMResultChunk):
+                    prompt_messages = event.prompt_messages
+                    system_fingerprint = event.system_fingerprint
+
                     if isinstance(event.delta.message.content, str):
                         result_text += event.delta.message.content
-                        prompt_messages = event.prompt_messages
-                        system_fingerprint = event.system_fingerprint
+                    elif isinstance(event.delta.message.content, list):
+                        for item in event.delta.message.content:
+                            if isinstance(item, TextPromptMessageContent):
+                                result_text += item.data
 
                 yield LLMResultChunkWithStructuredOutput(
                     model=model_schema.model,

api/core/mcp/auth/auth_flow.py

@@ -0,0 +1,342 @@
+import base64
+import hashlib
+import json
+import os
+import secrets
+import urllib.parse
+from typing import Optional
+from urllib.parse import urljoin
+
+import requests
+from pydantic import BaseModel, ValidationError
+
+from core.mcp.auth.auth_provider import OAuthClientProvider
+from core.mcp.types import (
+    OAuthClientInformation,
+    OAuthClientInformationFull,
+    OAuthClientMetadata,
+    OAuthMetadata,
+    OAuthTokens,
+)
+from extensions.ext_redis import redis_client
+
+LATEST_PROTOCOL_VERSION = "1.0"
+OAUTH_STATE_EXPIRY_SECONDS = 5 * 60  # 5 minutes expiry
+OAUTH_STATE_REDIS_KEY_PREFIX = "oauth_state:"
+
+
+class OAuthCallbackState(BaseModel):
+    provider_id: str
+    tenant_id: str
+    server_url: str
+    metadata: OAuthMetadata | None = None
+    client_information: OAuthClientInformation
+    code_verifier: str
+    redirect_uri: str
+
+
+def generate_pkce_challenge() -> tuple[str, str]:
+    """Generate PKCE challenge and verifier."""
+    code_verifier = base64.urlsafe_b64encode(os.urandom(40)).decode("utf-8")
+    code_verifier = code_verifier.replace("=", "").replace("+", "-").replace("/", "_")
+
+    code_challenge_hash = hashlib.sha256(code_verifier.encode("utf-8")).digest()
+    code_challenge = base64.urlsafe_b64encode(code_challenge_hash).decode("utf-8")
+    code_challenge = code_challenge.replace("=", "").replace("+", "-").replace("/", "_")
+
+    return code_verifier, code_challenge
+
+
+def _create_secure_redis_state(state_data: OAuthCallbackState) -> str:
+    """Create a secure state parameter by storing state data in Redis and returning a random state key."""
+    # Generate a secure random state key
+    state_key = secrets.token_urlsafe(32)
+
+    # Store the state data in Redis with expiration
+    redis_key = f"{OAUTH_STATE_REDIS_KEY_PREFIX}{state_key}"
+    redis_client.setex(redis_key, OAUTH_STATE_EXPIRY_SECONDS, state_data.model_dump_json())
+
+    return state_key
+
+
+def _retrieve_redis_state(state_key: str) -> OAuthCallbackState:
+    """Retrieve and decode OAuth state data from Redis using the state key, then delete it."""
+    redis_key = f"{OAUTH_STATE_REDIS_KEY_PREFIX}{state_key}"
+
+    # Get state data from Redis
+    state_data = redis_client.get(redis_key)
+
+    if not state_data:
+        raise ValueError("State parameter has expired or does not exist")
+
+    # Delete the state data from Redis immediately after retrieval to prevent reuse
+    redis_client.delete(redis_key)
+
+    try:
+        # Parse and validate the state data
+        oauth_state = OAuthCallbackState.model_validate_json(state_data)
+
+        return oauth_state
+    except ValidationError as e:
+        raise ValueError(f"Invalid state parameter: {str(e)}")
+
+
+def handle_callback(state_key: str, authorization_code: str) -> OAuthCallbackState:
+    """Handle the callback from the OAuth provider."""
+    # Retrieve state data from Redis (state is automatically deleted after retrieval)
+    full_state_data = _retrieve_redis_state(state_key)
+
+    tokens = exchange_authorization(
+        full_state_data.server_url,
+        full_state_data.metadata,
+        full_state_data.client_information,
+        authorization_code,
+        full_state_data.code_verifier,
+        full_state_data.redirect_uri,
+    )
+    provider = OAuthClientProvider(full_state_data.provider_id, full_state_data.tenant_id, for_list=True)
+    provider.save_tokens(tokens)
+    return full_state_data
+
+
+def discover_oauth_metadata(server_url: str, protocol_version: Optional[str] = None) -> Optional[OAuthMetadata]:
+    """Looks up RFC 8414 OAuth 2.0 Authorization Server Metadata."""
+    url = urljoin(server_url, "/.well-known/oauth-authorization-server")
+
+    try:
+        headers = {"MCP-Protocol-Version": protocol_version or LATEST_PROTOCOL_VERSION}
+        response = requests.get(url, headers=headers)
+        if response.status_code == 404:
+            return None
+        if not response.ok:
+            raise ValueError(f"HTTP {response.status_code} trying to load well-known OAuth metadata")
+        return OAuthMetadata.model_validate(response.json())
+    except requests.RequestException as e:
+        if isinstance(e, requests.ConnectionError):
+            response = requests.get(url)
+            if response.status_code == 404:
+                return None
+            if not response.ok:
+                raise ValueError(f"HTTP {response.status_code} trying to load well-known OAuth metadata")
+            return OAuthMetadata.model_validate(response.json())
+        raise
+
+
+def start_authorization(
+    server_url: str,
+    metadata: Optional[OAuthMetadata],
+    client_information: OAuthClientInformation,
+    redirect_url: str,
+    provider_id: str,
+    tenant_id: str,
+) -> tuple[str, str]:
+    """Begins the authorization flow with secure Redis state storage."""
+    response_type = "code"
+    code_challenge_method = "S256"
+
+    if metadata:
+        authorization_url = metadata.authorization_endpoint
+        if response_type not in metadata.response_types_supported:
+            raise ValueError(f"Incompatible auth server: does not support response type {response_type}")
+        if (
+            not metadata.code_challenge_methods_supported
+            or code_challenge_method not in metadata.code_challenge_methods_supported
+        ):
+            raise ValueError(
+                f"Incompatible auth server: does not support code challenge method {code_challenge_method}"
+            )
+    else:
+        authorization_url = urljoin(server_url, "/authorize")
+
+    code_verifier, code_challenge = generate_pkce_challenge()
+
+    # Prepare state data with all necessary information
+    state_data = OAuthCallbackState(
+        provider_id=provider_id,
+        tenant_id=tenant_id,
+        server_url=server_url,
+        metadata=metadata,
+        client_information=client_information,
+        code_verifier=code_verifier,
+        redirect_uri=redirect_url,
+    )
+
+    # Store state data in Redis and generate secure state key
+    state_key = _create_secure_redis_state(state_data)
+
+    params = {
+        "response_type": response_type,
+        "client_id": client_information.client_id,
+        "code_challenge": code_challenge,
+        "code_challenge_method": code_challenge_method,
+        "redirect_uri": redirect_url,
+        "state": state_key,
+    }
+
+    authorization_url = f"{authorization_url}?{urllib.parse.urlencode(params)}"
+    return authorization_url, code_verifier
+
+
+def exchange_authorization(
+    server_url: str,
+    metadata: Optional[OAuthMetadata],
+    client_information: OAuthClientInformation,
+    authorization_code: str,
+    code_verifier: str,
+    redirect_uri: str,
+) -> OAuthTokens:
+    """Exchanges an authorization code for an access token."""
+    grant_type = "authorization_code"
+
+    if metadata:
+        token_url = metadata.token_endpoint
+        if metadata.grant_types_supported and grant_type not in metadata.grant_types_supported:
+            raise ValueError(f"Incompatible auth server: does not support grant type {grant_type}")
+    else:
+        token_url = urljoin(server_url, "/token")
+
+    params = {
+        "grant_type": grant_type,
+        "client_id": client_information.client_id,
+        "code": authorization_code,
+        "code_verifier": code_verifier,
+        "redirect_uri": redirect_uri,
+    }
+
+    if client_information.client_secret:
+        params["client_secret"] = client_information.client_secret
+
+    response = requests.post(token_url, data=params)
+    if not response.ok:
+        raise ValueError(f"Token exchange failed: HTTP {response.status_code}")
+    return OAuthTokens.model_validate(response.json())
+
+
+def refresh_authorization(
+    server_url: str,
+    metadata: Optional[OAuthMetadata],
+    client_information: OAuthClientInformation,
+    refresh_token: str,
+) -> OAuthTokens:
+    """Exchange a refresh token for an updated access token."""
+    grant_type = "refresh_token"
+
+    if metadata:
+        token_url = metadata.token_endpoint
+        if metadata.grant_types_supported and grant_type not in metadata.grant_types_supported:
+            raise ValueError(f"Incompatible auth server: does not support grant type {grant_type}")
+    else:
+        token_url = urljoin(server_url, "/token")
+
+    params = {
+        "grant_type": grant_type,
+        "client_id": client_information.client_id,
+        "refresh_token": refresh_token,
+    }
+
+    if client_information.client_secret:
+        params["client_secret"] = client_information.client_secret
+
+    response = requests.post(token_url, data=params)
+    if not response.ok:
+        raise ValueError(f"Token refresh failed: HTTP {response.status_code}")
+    return OAuthTokens.parse_obj(response.json())
+
+
+def register_client(
+    server_url: str,
+    metadata: Optional[OAuthMetadata],
+    client_metadata: OAuthClientMetadata,
+) -> OAuthClientInformationFull:
+    """Performs OAuth 2.0 Dynamic Client Registration."""
+    if metadata:
+        if not metadata.registration_endpoint:
+            raise ValueError("Incompatible auth server: does not support dynamic client registration")
+        registration_url = metadata.registration_endpoint
+    else:
+        registration_url = urljoin(server_url, "/register")
+
+    response = requests.post(
+        registration_url,
+        json=client_metadata.model_dump(),
+        headers={"Content-Type": "application/json"},
+    )
+    if not response.ok:
+        response.raise_for_status()
+    return OAuthClientInformationFull.model_validate(response.json())
+
+
+def auth(
+    provider: OAuthClientProvider,
+    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)
+
+    # Handle client registration if needed
+    client_information = provider.client_information()
+    if not client_information:
+        if authorization_code is not None:
+            raise ValueError("Existing OAuth client information is required when exchanging an authorization code")
+        try:
+            full_information = register_client(server_url, metadata, provider.client_metadata)
+        except requests.RequestException as e:
+            raise ValueError(f"Could not register OAuth client: {e}")
+        provider.save_client_information(full_information)
+        client_information = full_information
+
+    # Exchange authorization code for tokens
+    if authorization_code is not None:
+        if not state_param:
+            raise ValueError("State parameter is required when exchanging authorization code")
+
+        try:
+            # Retrieve state data from Redis using state key
+            full_state_data = _retrieve_redis_state(state_param)
+
+            code_verifier = full_state_data.code_verifier
+            redirect_uri = full_state_data.redirect_uri
+
+            if not code_verifier or not redirect_uri:
+                raise ValueError("Missing code_verifier or redirect_uri in state data")
+
+        except (json.JSONDecodeError, ValueError) as e:
+            raise ValueError(f"Invalid state parameter: {e}")
+
+        tokens = exchange_authorization(
+            server_url,
+            metadata,
+            client_information,
+            authorization_code,
+            code_verifier,
+            redirect_uri,
+        )
+        provider.save_tokens(tokens)
+        return {"result": "success"}
+
+    provider_tokens = provider.tokens()
+
+    # Handle token refresh or new authorization
+    if provider_tokens and provider_tokens.refresh_token:
+        try:
+            new_tokens = refresh_authorization(server_url, metadata, client_information, provider_tokens.refresh_token)
+            provider.save_tokens(new_tokens)
+            return {"result": "success"}
+        except Exception as e:
+            raise ValueError(f"Could not refresh OAuth tokens: {e}")
+
+    # Start new authorization flow
+    authorization_url, code_verifier = start_authorization(
+        server_url,
+        metadata,
+        client_information,
+        provider.redirect_url,
+        provider.mcp_provider.id,
+        provider.mcp_provider.tenant_id,
+    )
+
+    provider.save_code_verifier(code_verifier)
+    return {"authorization_url": authorization_url}

api/core/mcp/auth/auth_provider.py

@@ -0,0 +1,81 @@
+from typing import Optional
+
+from configs import dify_config
+from core.mcp.types import (
+    OAuthClientInformation,
+    OAuthClientInformationFull,
+    OAuthClientMetadata,
+    OAuthTokens,
+)
+from models.tools import MCPToolProvider
+from services.tools.mcp_tools_mange_service import MCPToolManageService
+
+LATEST_PROTOCOL_VERSION = "1.0"
+
+
+class OAuthClientProvider:
+    mcp_provider: MCPToolProvider
+
+    def __init__(self, provider_id: str, tenant_id: str, for_list: bool = False):
+        if for_list:
+            self.mcp_provider = MCPToolManageService.get_mcp_provider_by_provider_id(provider_id, tenant_id)
+        else:
+            self.mcp_provider = MCPToolManageService.get_mcp_provider_by_server_identifier(provider_id, tenant_id)
+
+    @property
+    def redirect_url(self) -> str:
+        """The URL to redirect the user agent to after authorization."""
+        return dify_config.CONSOLE_API_URL + "/console/api/mcp/oauth/callback"
+
+    @property
+    def client_metadata(self) -> OAuthClientMetadata:
+        """Metadata about this OAuth client."""
+        return OAuthClientMetadata(
+            redirect_uris=[self.redirect_url],
+            token_endpoint_auth_method="none",
+            grant_types=["authorization_code", "refresh_token"],
+            response_types=["code"],
+            client_name="Dify",
+            client_uri="https://github.com/langgenius/dify",
+        )
+
+    def client_information(self) -> Optional[OAuthClientInformation]:
+        """Loads information about this OAuth client."""
+        client_information = self.mcp_provider.decrypted_credentials.get("client_information", {})
+        if not client_information:
+            return None
+        return OAuthClientInformation.model_validate(client_information)
+
+    def save_client_information(self, client_information: OAuthClientInformationFull) -> None:
+        """Saves client information after dynamic registration."""
+        MCPToolManageService.update_mcp_provider_credentials(
+            self.mcp_provider,
+            {"client_information": client_information.model_dump()},
+        )
+
+    def tokens(self) -> Optional[OAuthTokens]:
+        """Loads any existing OAuth tokens for the current session."""
+        credentials = self.mcp_provider.decrypted_credentials
+        if not credentials:
+            return None
+        return OAuthTokens(
+            access_token=credentials.get("access_token", ""),
+            token_type=credentials.get("token_type", "Bearer"),
+            expires_in=int(credentials.get("expires_in", "3600") or 3600),
+            refresh_token=credentials.get("refresh_token", ""),
+        )
+
+    def save_tokens(self, tokens: OAuthTokens) -> None:
+        """Stores new OAuth tokens for the current session."""
+        # update mcp provider credentials
+        token_dict = tokens.model_dump()
+        MCPToolManageService.update_mcp_provider_credentials(self.mcp_provider, token_dict, authed=True)
+
+    def save_code_verifier(self, code_verifier: str) -> None:
+        """Saves a PKCE code verifier for the current session."""
+        MCPToolManageService.update_mcp_provider_credentials(self.mcp_provider, {"code_verifier": code_verifier})
+
+    def code_verifier(self) -> str:
+        """Loads the PKCE code verifier for the current session."""
+        # get code verifier from mcp provider credentials
+        return str(self.mcp_provider.decrypted_credentials.get("code_verifier", ""))

api/core/mcp/client/sse_client.py

@@ -0,0 +1,361 @@
+import logging
+import queue
+from collections.abc import Generator
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
+from typing import Any, TypeAlias, final
+from urllib.parse import urljoin, urlparse
+
+import httpx
+from sseclient import SSEClient
+
+from core.mcp import types
+from core.mcp.error import MCPAuthError, MCPConnectionError
+from core.mcp.types import SessionMessage
+from core.mcp.utils import create_ssrf_proxy_mcp_http_client, ssrf_proxy_sse_connect
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_QUEUE_READ_TIMEOUT = 3
+
+
+@final
+class _StatusReady:
+    def __init__(self, endpoint_url: str):
+        self._endpoint_url = endpoint_url
+
+
+@final
+class _StatusError:
+    def __init__(self, exc: Exception):
+        self._exc = exc
+
+
+# Type aliases for better readability
+ReadQueue: TypeAlias = queue.Queue[SessionMessage | Exception | None]
+WriteQueue: TypeAlias = queue.Queue[SessionMessage | Exception | None]
+StatusQueue: TypeAlias = queue.Queue[_StatusReady | _StatusError]
+
+
+def remove_request_params(url: str) -> str:
+    """Remove request parameters from URL, keeping only the path."""
+    return urljoin(url, urlparse(url).path)
+
+
+class SSETransport:
+    """SSE client transport implementation."""
+
+    def __init__(
+        self,
+        url: str,
+        headers: dict[str, Any] | None = None,
+        timeout: float = 5.0,
+        sse_read_timeout: float = 5 * 60,
+    ) -> None:
+        """Initialize the SSE transport.
+
+        Args:
+            url: The SSE endpoint URL.
+            headers: Optional headers to include in requests.
+            timeout: HTTP timeout for regular operations.
+            sse_read_timeout: Timeout for SSE read operations.
+        """
+        self.url = url
+        self.headers = headers or {}
+        self.timeout = timeout
+        self.sse_read_timeout = sse_read_timeout
+        self.endpoint_url: str | None = None
+
+    def _validate_endpoint_url(self, endpoint_url: str) -> bool:
+        """Validate that the endpoint URL matches the connection origin.
+
+        Args:
+            endpoint_url: The endpoint URL to validate.
+
+        Returns:
+            True if valid, False otherwise.
+        """
+        url_parsed = urlparse(self.url)
+        endpoint_parsed = urlparse(endpoint_url)
+
+        return url_parsed.netloc == endpoint_parsed.netloc and url_parsed.scheme == endpoint_parsed.scheme
+
+    def _handle_endpoint_event(self, sse_data: str, status_queue: StatusQueue) -> None:
+        """Handle an 'endpoint' SSE event.
+
+        Args:
+            sse_data: The SSE event data.
+            status_queue: Queue to put status updates.
+        """
+        endpoint_url = urljoin(self.url, sse_data)
+        logger.info(f"Received endpoint URL: {endpoint_url}")
+
+        if not self._validate_endpoint_url(endpoint_url):
+            error_msg = f"Endpoint origin does not match connection origin: {endpoint_url}"
+            logger.error(error_msg)
+            status_queue.put(_StatusError(ValueError(error_msg)))
+            return
+
+        status_queue.put(_StatusReady(endpoint_url))
+
+    def _handle_message_event(self, sse_data: str, read_queue: ReadQueue) -> None:
+        """Handle a 'message' SSE event.
+
+        Args:
+            sse_data: The SSE event data.
+            read_queue: Queue to put parsed messages.
+        """
+        try:
+            message = types.JSONRPCMessage.model_validate_json(sse_data)
+            logger.debug(f"Received server message: {message}")
+            session_message = SessionMessage(message)
+            read_queue.put(session_message)
+        except Exception as exc:
+            logger.exception("Error parsing server message")
+            read_queue.put(exc)
+
+    def _handle_sse_event(self, sse, read_queue: ReadQueue, status_queue: StatusQueue) -> None:
+        """Handle a single SSE event.
+
+        Args:
+            sse: The SSE event object.
+            read_queue: Queue for message events.
+            status_queue: Queue for status events.
+        """
+        match sse.event:
+            case "endpoint":
+                self._handle_endpoint_event(sse.data, status_queue)
+            case "message":
+                self._handle_message_event(sse.data, read_queue)
+            case _:
+                logger.warning(f"Unknown SSE event: {sse.event}")
+
+    def sse_reader(self, event_source, read_queue: ReadQueue, status_queue: StatusQueue) -> None:
+        """Read and process SSE events.
+
+        Args:
+            event_source: The SSE event source.
+            read_queue: Queue to put received messages.
+            status_queue: Queue to put status updates.
+        """
+        try:
+            for sse in event_source.iter_sse():
+                self._handle_sse_event(sse, read_queue, status_queue)
+        except httpx.ReadError as exc:
+            logger.debug(f"SSE reader shutting down normally: {exc}")
+        except Exception as exc:
+            read_queue.put(exc)
+        finally:
+            read_queue.put(None)
+
+    def _send_message(self, client: httpx.Client, endpoint_url: str, message: SessionMessage) -> None:
+        """Send a single message to the server.
+
+        Args:
+            client: HTTP client to use.
+            endpoint_url: The endpoint URL to send to.
+            message: The message to send.
+        """
+        response = client.post(
+            endpoint_url,
+            json=message.message.model_dump(
+                by_alias=True,
+                mode="json",
+                exclude_none=True,
+            ),
+        )
+        response.raise_for_status()
+        logger.debug(f"Client message sent successfully: {response.status_code}")
+
+    def post_writer(self, client: httpx.Client, endpoint_url: str, write_queue: WriteQueue) -> None:
+        """Handle writing messages to the server.
+
+        Args:
+            client: HTTP client to use.
+            endpoint_url: The endpoint URL to send messages to.
+            write_queue: Queue to read messages from.
+        """
+        try:
+            while True:
+                try:
+                    message = write_queue.get(timeout=DEFAULT_QUEUE_READ_TIMEOUT)
+                    if message is None:
+                        break
+                    if isinstance(message, Exception):
+                        write_queue.put(message)
+                        continue
+
+                    self._send_message(client, endpoint_url, message)
+
+                except queue.Empty:
+                    continue
+        except httpx.ReadError as exc:
+            logger.debug(f"Post writer shutting down normally: {exc}")
+        except Exception as exc:
+            logger.exception("Error writing messages")
+            write_queue.put(exc)
+        finally:
+            write_queue.put(None)
+
+    def _wait_for_endpoint(self, status_queue: StatusQueue) -> str:
+        """Wait for the endpoint URL from the status queue.
+
+        Args:
+            status_queue: Queue to read status from.
+
+        Returns:
+            The endpoint URL.
+
+        Raises:
+            ValueError: If endpoint URL is not received or there's an error.
+        """
+        try:
+            status = status_queue.get(timeout=1)
+        except queue.Empty:
+            raise ValueError("failed to get endpoint URL")
+
+        if isinstance(status, _StatusReady):
+            return status._endpoint_url
+        elif isinstance(status, _StatusError):
+            raise status._exc
+        else:
+            raise ValueError("failed to get endpoint URL")
+
+    def connect(
+        self,
+        executor: ThreadPoolExecutor,
+        client: httpx.Client,
+        event_source,
+    ) -> tuple[ReadQueue, WriteQueue]:
+        """Establish connection and start worker threads.
+
+        Args:
+            executor: Thread pool executor.
+            client: HTTP client.
+            event_source: SSE event source.
+
+        Returns:
+            Tuple of (read_queue, write_queue).
+        """
+        read_queue: ReadQueue = queue.Queue()
+        write_queue: WriteQueue = queue.Queue()
+        status_queue: StatusQueue = queue.Queue()
+
+        # Start SSE reader thread
+        executor.submit(self.sse_reader, event_source, read_queue, status_queue)
+
+        # Wait for endpoint URL
+        endpoint_url = self._wait_for_endpoint(status_queue)
+        self.endpoint_url = endpoint_url
+
+        # Start post writer thread
+        executor.submit(self.post_writer, client, endpoint_url, write_queue)
+
+        return read_queue, write_queue
+
+
+@contextmanager
+def sse_client(
+    url: str,
+    headers: dict[str, Any] | None = None,
+    timeout: float = 5.0,
+    sse_read_timeout: float = 5 * 60,
+) -> Generator[tuple[ReadQueue, WriteQueue], None, None]:
+    """
+    Client transport for SSE.
+    `sse_read_timeout` determines how long (in seconds) the client will wait for a new
+    event before disconnecting. All other HTTP operations are controlled by `timeout`.
+
+    Args:
+        url: The SSE endpoint URL.
+        headers: Optional headers to include in requests.
+        timeout: HTTP timeout for regular operations.
+        sse_read_timeout: Timeout for SSE read operations.
+
+    Yields:
+        Tuple of (read_queue, write_queue) for message communication.
+    """
+    transport = SSETransport(url, headers, timeout, sse_read_timeout)
+
+    read_queue: ReadQueue | None = None
+    write_queue: WriteQueue | None = None
+
+    with ThreadPoolExecutor() as executor:
+        try:
+            with create_ssrf_proxy_mcp_http_client(headers=transport.headers) as client:
+                with ssrf_proxy_sse_connect(
+                    url, timeout=httpx.Timeout(timeout, read=sse_read_timeout), client=client
+                ) as event_source:
+                    event_source.response.raise_for_status()
+
+                    read_queue, write_queue = transport.connect(executor, client, event_source)
+
+                    yield read_queue, write_queue
+
+        except httpx.HTTPStatusError as exc:
+            if exc.response.status_code == 401:
+                raise MCPAuthError()
+            raise MCPConnectionError()
+        except Exception:
+            logger.exception("Error connecting to SSE endpoint")
+            raise
+        finally:
+            # Clean up queues
+            if read_queue:
+                read_queue.put(None)
+            if write_queue:
+                write_queue.put(None)
+
+
+def send_message(http_client: httpx.Client, endpoint_url: str, session_message: SessionMessage) -> None:
+    """
+    Send a message to the server using the provided HTTP client.
+
+    Args:
+        http_client: The HTTP client to use for sending
+        endpoint_url: The endpoint URL to send the message to
+        session_message: The message to send
+    """
+    try:
+        response = http_client.post(
+            endpoint_url,
+            json=session_message.message.model_dump(
+                by_alias=True,
+                mode="json",
+                exclude_none=True,
+            ),
+        )
+        response.raise_for_status()
+        logger.debug(f"Client message sent successfully: {response.status_code}")
+    except Exception as exc:
+        logger.exception("Error sending message")
+        raise
+
+
+def read_messages(
+    sse_client: SSEClient,
+) -> Generator[SessionMessage | Exception, None, None]:
+    """
+    Read messages from the SSE client.
+
+    Args:
+        sse_client: The SSE client to read from
+
+    Yields:
+        SessionMessage or Exception for each event received
+    """
+    try:
+        for sse in sse_client.events():
+            if sse.event == "message":
+                try:
+                    message = types.JSONRPCMessage.model_validate_json(sse.data)
+                    logger.debug(f"Received server message: {message}")
+                    yield SessionMessage(message)
+                except Exception as exc:
+                    logger.exception("Error parsing server message")
+                    yield exc
+            else:
+                logger.warning(f"Unknown SSE event: {sse.event}")
+    except Exception as exc:
+        logger.exception("Error reading SSE messages")
+        yield exc

api/core/mcp/client/streamable_client.py

@@ -0,0 +1,476 @@
+"""
+StreamableHTTP Client Transport Module
+
+This module implements the StreamableHTTP transport for MCP clients,
+providing support for HTTP POST requests with optional SSE streaming responses
+and session management.
+"""
+
+import logging
+import queue
+from collections.abc import Callable, Generator
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Any, cast
+
+import httpx
+from httpx_sse import EventSource, ServerSentEvent
+
+from core.mcp.types import (
+    ClientMessageMetadata,
+    ErrorData,
+    JSONRPCError,
+    JSONRPCMessage,
+    JSONRPCNotification,
+    JSONRPCRequest,
+    JSONRPCResponse,
+    RequestId,
+    SessionMessage,
+)
+from core.mcp.utils import create_ssrf_proxy_mcp_http_client, ssrf_proxy_sse_connect
+
+logger = logging.getLogger(__name__)
+
+
+SessionMessageOrError = SessionMessage | Exception | None
+# Queue types with clearer names for their roles
+ServerToClientQueue = queue.Queue[SessionMessageOrError]  # Server to client messages
+ClientToServerQueue = queue.Queue[SessionMessage | None]  # Client to server messages
+GetSessionIdCallback = Callable[[], str | None]
+
+MCP_SESSION_ID = "mcp-session-id"
+LAST_EVENT_ID = "last-event-id"
+CONTENT_TYPE = "content-type"
+ACCEPT = "Accept"
+
+
+JSON = "application/json"
+SSE = "text/event-stream"
+
+DEFAULT_QUEUE_READ_TIMEOUT = 3
+
+
+class StreamableHTTPError(Exception):
+    """Base exception for StreamableHTTP transport errors."""
+
+    pass
+
+
+class ResumptionError(StreamableHTTPError):
+    """Raised when resumption request is invalid."""
+
+    pass
+
+
+@dataclass
+class RequestContext:
+    """Context for a request operation."""
+
+    client: httpx.Client
+    headers: dict[str, str]
+    session_id: str | None
+    session_message: SessionMessage
+    metadata: ClientMessageMetadata | None
+    server_to_client_queue: ServerToClientQueue  # Renamed for clarity
+    sse_read_timeout: timedelta
+
+
+class StreamableHTTPTransport:
+    """StreamableHTTP client transport implementation."""
+
+    def __init__(
+        self,
+        url: str,
+        headers: dict[str, Any] | None = None,
+        timeout: timedelta = timedelta(seconds=30),
+        sse_read_timeout: timedelta = timedelta(seconds=60 * 5),
+    ) -> None:
+        """Initialize the StreamableHTTP transport.
+
+        Args:
+            url: The endpoint URL.
+            headers: Optional headers to include in requests.
+            timeout: HTTP timeout for regular operations.
+            sse_read_timeout: Timeout for SSE read operations.
+        """
+        self.url = url
+        self.headers = headers or {}
+        self.timeout = timeout
+        self.sse_read_timeout = sse_read_timeout
+        self.session_id: str | None = None
+        self.request_headers = {
+            ACCEPT: f"{JSON}, {SSE}",
+            CONTENT_TYPE: JSON,
+            **self.headers,
+        }
+
+    def _update_headers_with_session(self, base_headers: dict[str, str]) -> dict[str, str]:
+        """Update headers with session ID if available."""
+        headers = base_headers.copy()
+        if self.session_id:
+            headers[MCP_SESSION_ID] = self.session_id
+        return headers
+
+    def _is_initialization_request(self, message: JSONRPCMessage) -> bool:
+        """Check if the message is an initialization request."""
+        return isinstance(message.root, JSONRPCRequest) and message.root.method == "initialize"
+
+    def _is_initialized_notification(self, message: JSONRPCMessage) -> bool:
+        """Check if the message is an initialized notification."""
+        return isinstance(message.root, JSONRPCNotification) and message.root.method == "notifications/initialized"
+
+    def _maybe_extract_session_id_from_response(
+        self,
+        response: httpx.Response,
+    ) -> None:
+        """Extract and store session ID from response headers."""
+        new_session_id = response.headers.get(MCP_SESSION_ID)
+        if new_session_id:
+            self.session_id = new_session_id
+            logger.info(f"Received session ID: {self.session_id}")
+
+    def _handle_sse_event(
+        self,
+        sse: ServerSentEvent,
+        server_to_client_queue: ServerToClientQueue,
+        original_request_id: RequestId | None = None,
+        resumption_callback: Callable[[str], None] | None = None,
+    ) -> bool:
+        """Handle an SSE event, returning True if the response is complete."""
+        if sse.event == "message":
+            try:
+                message = JSONRPCMessage.model_validate_json(sse.data)
+                logger.debug(f"SSE message: {message}")
+
+                # If this is a response and we have original_request_id, replace it
+                if original_request_id is not None and isinstance(message.root, JSONRPCResponse | JSONRPCError):
+                    message.root.id = original_request_id
+
+                session_message = SessionMessage(message)
+                # Put message in queue that goes to client
+                server_to_client_queue.put(session_message)
+
+                # Call resumption token callback if we have an ID
+                if sse.id and resumption_callback:
+                    resumption_callback(sse.id)
+
+                # If this is a response or error return True indicating completion
+                # Otherwise, return False to continue listening
+                return isinstance(message.root, JSONRPCResponse | JSONRPCError)
+
+            except Exception as exc:
+                # Put exception in queue that goes to client
+                server_to_client_queue.put(exc)
+                return False
+        elif sse.event == "ping":
+            logger.debug("Received ping event")
+            return False
+        else:
+            logger.warning(f"Unknown SSE event: {sse.event}")
+            return False
+
+    def handle_get_stream(
+        self,
+        client: httpx.Client,
+        server_to_client_queue: ServerToClientQueue,
+    ) -> None:
+        """Handle GET stream for server-initiated messages."""
+        try:
+            if not self.session_id:
+                return
+
+            headers = self._update_headers_with_session(self.request_headers)
+
+            with ssrf_proxy_sse_connect(
+                self.url,
+                headers=headers,
+                timeout=httpx.Timeout(self.timeout.seconds, read=self.sse_read_timeout.seconds),
+                client=client,
+                method="GET",
+            ) as event_source:
+                event_source.response.raise_for_status()
+                logger.debug("GET SSE connection established")
+
+                for sse in event_source.iter_sse():
+                    self._handle_sse_event(sse, server_to_client_queue)
+
+        except Exception as exc:
+            logger.debug(f"GET stream error (non-fatal): {exc}")
+
+    def _handle_resumption_request(self, ctx: RequestContext) -> None:
+        """Handle a resumption request using GET with SSE."""
+        headers = self._update_headers_with_session(ctx.headers)
+        if ctx.metadata and ctx.metadata.resumption_token:
+            headers[LAST_EVENT_ID] = ctx.metadata.resumption_token
+        else:
+            raise ResumptionError("Resumption request requires a resumption token")
+
+        # Extract original request ID to map responses
+        original_request_id = None
+        if isinstance(ctx.session_message.message.root, JSONRPCRequest):
+            original_request_id = ctx.session_message.message.root.id
+
+        with ssrf_proxy_sse_connect(
+            self.url,
+            headers=headers,
+            timeout=httpx.Timeout(self.timeout.seconds, read=ctx.sse_read_timeout.seconds),
+            client=ctx.client,
+            method="GET",
+        ) as event_source:
+            event_source.response.raise_for_status()
+            logger.debug("Resumption GET SSE connection established")
+
+            for sse in event_source.iter_sse():
+                is_complete = self._handle_sse_event(
+                    sse,
+                    ctx.server_to_client_queue,
+                    original_request_id,
+                    ctx.metadata.on_resumption_token_update if ctx.metadata else None,
+                )
+                if is_complete:
+                    break
+
+    def _handle_post_request(self, ctx: RequestContext) -> None:
+        """Handle a POST request with response processing."""
+        headers = self._update_headers_with_session(ctx.headers)
+        message = ctx.session_message.message
+        is_initialization = self._is_initialization_request(message)
+
+        with ctx.client.stream(
+            "POST",
+            self.url,
+            json=message.model_dump(by_alias=True, mode="json", exclude_none=True),
+            headers=headers,
+        ) as response:
+            if response.status_code == 202:
+                logger.debug("Received 202 Accepted")
+                return
+
+            if response.status_code == 404:
+                if isinstance(message.root, JSONRPCRequest):
+                    self._send_session_terminated_error(
+                        ctx.server_to_client_queue,
+                        message.root.id,
+                    )
+                return
+
+            response.raise_for_status()
+            if is_initialization:
+                self._maybe_extract_session_id_from_response(response)
+
+            content_type = cast(str, response.headers.get(CONTENT_TYPE, "").lower())
+
+            if content_type.startswith(JSON):
+                self._handle_json_response(response, ctx.server_to_client_queue)
+            elif content_type.startswith(SSE):
+                self._handle_sse_response(response, ctx)
+            else:
+                self._handle_unexpected_content_type(
+                    content_type,
+                    ctx.server_to_client_queue,
+                )
+
+    def _handle_json_response(
+        self,
+        response: httpx.Response,
+        server_to_client_queue: ServerToClientQueue,
+    ) -> None:
+        """Handle JSON response from the server."""
+        try:
+            content = response.read()
+            message = JSONRPCMessage.model_validate_json(content)
+            session_message = SessionMessage(message)
+            server_to_client_queue.put(session_message)
+        except Exception as exc:
+            server_to_client_queue.put(exc)
+
+    def _handle_sse_response(self, response: httpx.Response, ctx: RequestContext) -> None:
+        """Handle SSE response from the server."""
+        try:
+            event_source = EventSource(response)
+            for sse in event_source.iter_sse():
+                is_complete = self._handle_sse_event(
+                    sse,
+                    ctx.server_to_client_queue,
+                    resumption_callback=(ctx.metadata.on_resumption_token_update if ctx.metadata else None),
+                )
+                if is_complete:
+                    break
+        except Exception as e:
+            ctx.server_to_client_queue.put(e)
+
+    def _handle_unexpected_content_type(
+        self,
+        content_type: str,
+        server_to_client_queue: ServerToClientQueue,
+    ) -> None:
+        """Handle unexpected content type in response."""
+        error_msg = f"Unexpected content type: {content_type}"
+        logger.error(error_msg)
+        server_to_client_queue.put(ValueError(error_msg))
+
+    def _send_session_terminated_error(
+        self,
+        server_to_client_queue: ServerToClientQueue,
+        request_id: RequestId,
+    ) -> None:
+        """Send a session terminated error response."""
+        jsonrpc_error = JSONRPCError(
+            jsonrpc="2.0",
+            id=request_id,
+            error=ErrorData(code=32600, message="Session terminated by server"),
+        )
+        session_message = SessionMessage(JSONRPCMessage(jsonrpc_error))
+        server_to_client_queue.put(session_message)
+
+    def post_writer(
+        self,
+        client: httpx.Client,
+        client_to_server_queue: ClientToServerQueue,
+        server_to_client_queue: ServerToClientQueue,
+        start_get_stream: Callable[[], None],
+    ) -> None:
+        """Handle writing requests to the server.
+
+        This method processes messages from the client_to_server_queue and sends them to the server.
+        Responses are written to the server_to_client_queue.
+        """
+        while True:
+            try:
+                # Read message from client queue with timeout to check stop_event periodically
+                session_message = client_to_server_queue.get(timeout=DEFAULT_QUEUE_READ_TIMEOUT)
+                if session_message is None:
+                    break
+
+                message = session_message.message
+                metadata = (
+                    session_message.metadata if isinstance(session_message.metadata, ClientMessageMetadata) else None
+                )
+
+                # Check if this is a resumption request
+                is_resumption = bool(metadata and metadata.resumption_token)
+
+                logger.debug(f"Sending client message: {message}")
+
+                # Handle initialized notification
+                if self._is_initialized_notification(message):
+                    start_get_stream()
+
+                ctx = RequestContext(
+                    client=client,
+                    headers=self.request_headers,
+                    session_id=self.session_id,
+                    session_message=session_message,
+                    metadata=metadata,
+                    server_to_client_queue=server_to_client_queue,  # Queue to write responses to client
+                    sse_read_timeout=self.sse_read_timeout,
+                )
+
+                if is_resumption:
+                    self._handle_resumption_request(ctx)
+                else:
+                    self._handle_post_request(ctx)
+            except queue.Empty:
+                continue
+            except Exception as exc:
+                server_to_client_queue.put(exc)
+
+    def terminate_session(self, client: httpx.Client) -> None:
+        """Terminate the session by sending a DELETE request."""
+        if not self.session_id:
+            return
+
+        try:
+            headers = self._update_headers_with_session(self.request_headers)
+            response = client.delete(self.url, headers=headers)
+
+            if response.status_code == 405:
+                logger.debug("Server does not allow session termination")
+            elif response.status_code != 200:
+                logger.warning(f"Session termination failed: {response.status_code}")
+        except Exception as exc:
+            logger.warning(f"Session termination failed: {exc}")
+
+    def get_session_id(self) -> str | None:
+        """Get the current session ID."""
+        return self.session_id
+
+
+@contextmanager
+def streamablehttp_client(
+    url: str,
+    headers: dict[str, Any] | None = None,
+    timeout: timedelta = timedelta(seconds=30),
+    sse_read_timeout: timedelta = timedelta(seconds=60 * 5),
+    terminate_on_close: bool = True,
+) -> Generator[
+    tuple[
+        ServerToClientQueue,  # Queue for receiving messages FROM server
+        ClientToServerQueue,  # Queue for sending messages TO server
+        GetSessionIdCallback,
+    ],
+    None,
+    None,
+]:
+    """
+    Client transport for StreamableHTTP.
+
+    `sse_read_timeout` determines how long (in seconds) the client will wait for a new
+    event before disconnecting. All other HTTP operations are controlled by `timeout`.
+
+    Yields:
+        Tuple containing:
+            - server_to_client_queue: Queue for reading messages FROM the server
+            - client_to_server_queue: Queue for sending messages TO the server
+            - get_session_id_callback: Function to retrieve the current session ID
+    """
+    transport = StreamableHTTPTransport(url, headers, timeout, sse_read_timeout)
+
+    # Create queues with clear directional meaning
+    server_to_client_queue: ServerToClientQueue = queue.Queue()  # For messages FROM server TO client
+    client_to_server_queue: ClientToServerQueue = queue.Queue()  # For messages FROM client TO server
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        try:
+            with create_ssrf_proxy_mcp_http_client(
+                headers=transport.request_headers,
+                timeout=httpx.Timeout(transport.timeout.seconds, read=transport.sse_read_timeout.seconds),
+            ) as client:
+                # Define callbacks that need access to thread pool
+                def start_get_stream() -> None:
+                    """Start a worker thread to handle server-initiated messages."""
+                    executor.submit(transport.handle_get_stream, client, server_to_client_queue)
+
+                # Start the post_writer worker thread
+                executor.submit(
+                    transport.post_writer,
+                    client,
+                    client_to_server_queue,  # Queue for messages FROM client TO server
+                    server_to_client_queue,  # Queue for messages FROM server TO client
+                    start_get_stream,
+                )
+
+                try:
+                    yield (
+                        server_to_client_queue,  # Queue for receiving messages FROM server
+                        client_to_server_queue,  # Queue for sending messages TO server
+                        transport.get_session_id,
+                    )
+                finally:
+                    if transport.session_id and terminate_on_close:
+                        transport.terminate_session(client)
+
+                    # Signal threads to stop
+                    client_to_server_queue.put(None)
+        finally:
+            # Clear any remaining items and add None sentinel to unblock any waiting threads
+            try:
+                while not client_to_server_queue.empty():
+                    client_to_server_queue.get_nowait()
+            except queue.Empty:
+                pass
+
+            client_to_server_queue.put(None)
+            server_to_client_queue.put(None)

api/core/mcp/entities.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+from typing import Any, Generic, TypeVar
+
+from core.mcp.session.base_session import BaseSession
+from core.mcp.types import LATEST_PROTOCOL_VERSION, RequestId, RequestParams
+
+SUPPORTED_PROTOCOL_VERSIONS: list[str] = ["2024-11-05", LATEST_PROTOCOL_VERSION]
+
+
+SessionT = TypeVar("SessionT", bound=BaseSession[Any, Any, Any, Any, Any])
+LifespanContextT = TypeVar("LifespanContextT")
+
+
+@dataclass
+class RequestContext(Generic[SessionT, LifespanContextT]):
+    request_id: RequestId
+    meta: RequestParams.Meta | None
+    session: SessionT
+    lifespan_context: LifespanContextT

api/core/mcp/error.py

@@ -0,0 +1,10 @@
+class MCPError(Exception):
+    pass
+
+
+class MCPConnectionError(MCPError):
+    pass
+
+
+class MCPAuthError(MCPConnectionError):
+    pass

api/core/mcp/mcp_client.py

@@ -0,0 +1,150 @@
+import logging
+from collections.abc import Callable
+from contextlib import AbstractContextManager, ExitStack
+from types import TracebackType
+from typing import Any, Optional, cast
+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.session.client_session import ClientSession
+from core.mcp.types import Tool
+
+logger = logging.getLogger(__name__)
+
+
+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,
+    ):
+        # Initialize info
+        self.provider_id = provider_id
+        self.tenant_id = tenant_id
+        self.client_type = "streamable"
+        self.server_url = server_url
+
+        # 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.exit_stack = ExitStack()
+
+        # Whether the client has been initialized
+        self._initialized = False
+
+    def __enter__(self):
+        self._initialize()
+        self._initialized = True
+        return self
+
+    def __exit__(
+        self, exc_type: Optional[type], exc_value: Optional[BaseException], traceback: Optional[TracebackType]
+    ):
+        self.cleanup()
+
+    def _initialize(
+        self,
+    ):
+        """Initialize the client with fallback to SSE if streamable connection fails"""
+        connection_methods: dict[str, Callable[..., AbstractContextManager[Any]]] = {
+            "mcp": streamablehttp_client,
+            "sse": sse_client,
+        }
+
+        parsed_url = urlparse(self.server_url)
+        path = parsed_url.path
+        method_name = path.rstrip("/").split("/")[-1] if path else ""
+        try:
+            client_factory = connection_methods[method_name]
+            self.connect_server(client_factory, method_name)
+        except KeyError:
+            try:
+                self.connect_server(sse_client, "sse")
+            except MCPConnectionError:
+                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
+
+        try:
+            headers = (
+                {"Authorization": f"{self.token.token_type.capitalize()} {self.token.access_token}"}
+                if self.authed and self.token
+                else {}
+            )
+            self._streams_context = client_factory(url=self.server_url, headers=headers)
+            if self._streams_context is None:
+                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(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)
+            session = cast(ClientSession, self._session)
+            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)
+
+        except MCPConnectionError:
+            raise
+
+    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:
+            raise ValueError("Session not initialized.")
+        response = self._session.list_tools()
+        tools = response.tools
+        return tools
+
+    def invoke_tool(self, tool_name: str, tool_args: dict):
+        """Call a tool"""
+        if not self._initialized or not self._session:
+            raise ValueError("Session not initialized.")
+        return self._session.call_tool(tool_name, tool_args)
+
+    def cleanup(self):
+        """Clean up resources"""
+        try:
+            # ExitStack will handle proper cleanup of all managed context managers
+            self.exit_stack.close()
+            self._session = None
+            self._session_context = None
+            self._streams_context = None
+            self._initialized = False
+        except Exception as e:
+            logging.exception("Error during cleanup")
+            raise ValueError(f"Error during cleanup: {e}")

api/core/mcp/server/streamable_http.py

@@ -0,0 +1,228 @@
+import json
+import logging
+from collections.abc import Mapping
+from typing import Any, cast
+
+from configs import dify_config
+from controllers.web.passport import generate_session_id
+from core.app.app_config.entities import VariableEntity, VariableEntityType
+from core.app.entities.app_invoke_entities import InvokeFrom
+from core.app.features.rate_limiting.rate_limit import RateLimitGenerator
+from core.mcp import types
+from core.mcp.types import INTERNAL_ERROR, INVALID_PARAMS, METHOD_NOT_FOUND
+from core.mcp.utils import create_mcp_error_response
+from core.model_runtime.utils.encoders import jsonable_encoder
+from extensions.ext_database import db
+from models.model import App, AppMCPServer, AppMode, EndUser
+from services.app_generate_service import AppGenerateService
+
+"""
+Apply to MCP HTTP streamable server with stateless http
+"""
+logger = logging.getLogger(__name__)
+
+
+class MCPServerStreamableHTTPRequestHandler:
+    def __init__(
+        self, app: App, request: types.ClientRequest | types.ClientNotification, user_input_form: list[VariableEntity]
+    ):
+        self.app = app
+        self.request = request
+        mcp_server = db.session.query(AppMCPServer).filter(AppMCPServer.app_id == self.app.id).first()
+        if not mcp_server:
+            raise ValueError("MCP server not found")
+        self.mcp_server: AppMCPServer = mcp_server
+        self.end_user = self.retrieve_end_user()
+        self.user_input_form = user_input_form
+
+    @property
+    def request_type(self):
+        return type(self.request.root)
+
+    @property
+    def parameter_schema(self):
+        parameters, required = self._convert_input_form_to_parameters(self.user_input_form)
+        if self.app.mode in {AppMode.COMPLETION.value, AppMode.WORKFLOW.value}:
+            return {
+                "type": "object",
+                "properties": parameters,
+                "required": required,
+            }
+        return {
+            "type": "object",
+            "properties": {
+                "query": {"type": "string", "description": "User Input/Question content"},
+                **parameters,
+            },
+            "required": ["query", *required],
+        }
+
+    @property
+    def capabilities(self):
+        return types.ServerCapabilities(
+            tools=types.ToolsCapability(listChanged=False),
+        )
+
+    def response(self, response: types.Result | str):
+        if isinstance(response, str):
+            sse_content = f"event: ping\ndata: {response}\n\n".encode()
+            yield sse_content
+            return
+        json_response = types.JSONRPCResponse(
+            jsonrpc="2.0",
+            id=(self.request.root.model_extra or {}).get("id", 1),
+            result=response.model_dump(by_alias=True, mode="json", exclude_none=True),
+        )
+        json_data = json.dumps(jsonable_encoder(json_response))
+
+        sse_content = f"event: message\ndata: {json_data}\n\n".encode()
+
+        yield sse_content
+
+    def error_response(self, code: int, message: str, data=None):
+        request_id = (self.request.root.model_extra or {}).get("id", 1) or 1
+        return create_mcp_error_response(request_id, code, message, data)
+
+    def handle(self):
+        handle_map = {
+            types.InitializeRequest: self.initialize,
+            types.ListToolsRequest: self.list_tools,
+            types.CallToolRequest: self.invoke_tool,
+            types.InitializedNotification: self.handle_notification,
+            types.PingRequest: self.handle_ping,
+        }
+        try:
+            if self.request_type in handle_map:
+                return self.response(handle_map[self.request_type]())
+            else:
+                return self.error_response(METHOD_NOT_FOUND, f"Method not found: {self.request_type}")
+        except ValueError as e:
+            logger.exception("Invalid params")
+            return self.error_response(INVALID_PARAMS, str(e))
+        except Exception as e:
+            logger.exception("Internal server error")
+            return self.error_response(INTERNAL_ERROR, f"Internal server error: {str(e)}")
+
+    def handle_notification(self):
+        return "ping"
+
+    def handle_ping(self):
+        return types.EmptyResult()
+
+    def initialize(self):
+        request = cast(types.InitializeRequest, self.request.root)
+        client_info = request.params.clientInfo
+        client_name = f"{client_info.name}@{client_info.version}"
+        if not self.end_user:
+            end_user = EndUser(
+                tenant_id=self.app.tenant_id,
+                app_id=self.app.id,
+                type="mcp",
+                name=client_name,
+                session_id=generate_session_id(),
+                external_user_id=self.mcp_server.id,
+            )
+            db.session.add(end_user)
+            db.session.commit()
+        return types.InitializeResult(
+            protocolVersion=types.SERVER_LATEST_PROTOCOL_VERSION,
+            capabilities=self.capabilities,
+            serverInfo=types.Implementation(name="Dify", version=dify_config.project.version),
+            instructions=self.mcp_server.description,
+        )
+
+    def list_tools(self):
+        if not self.end_user:
+            raise ValueError("User not found")
+        return types.ListToolsResult(
+            tools=[
+                types.Tool(
+                    name=self.app.name,
+                    description=self.mcp_server.description,
+                    inputSchema=self.parameter_schema,
+                )
+            ],
+        )
+
+    def invoke_tool(self):
+        if not self.end_user:
+            raise ValueError("User not found")
+        request = cast(types.CallToolRequest, self.request.root)
+        args = request.params.arguments
+        if not args:
+            raise ValueError("No arguments provided")
+        if self.app.mode in {AppMode.WORKFLOW.value}:
+            args = {"inputs": args}
+        elif self.app.mode in {AppMode.COMPLETION.value}:
+            args = {"query": "", "inputs": args}
+        else:
+            args = {"query": args["query"], "inputs": {k: v for k, v in args.items() if k != "query"}}
+        response = AppGenerateService.generate(
+            self.app,
+            self.end_user,
+            args,
+            InvokeFrom.SERVICE_API,
+            streaming=self.app.mode == AppMode.AGENT_CHAT.value,
+        )
+        answer = ""
+        if isinstance(response, RateLimitGenerator):
+            for item in response.generator:
+                data = item
+                if isinstance(data, str) and data.startswith("data: "):
+                    try:
+                        json_str = data[6:].strip()
+                        parsed_data = json.loads(json_str)
+                        if parsed_data.get("event") == "agent_thought":
+                            answer += parsed_data.get("thought", "")
+                    except json.JSONDecodeError:
+                        continue
+        if isinstance(response, Mapping):
+            if self.app.mode in {
+                AppMode.ADVANCED_CHAT.value,
+                AppMode.COMPLETION.value,
+                AppMode.CHAT.value,
+                AppMode.AGENT_CHAT.value,
+            }:
+                answer = response["answer"]
+            elif self.app.mode in {AppMode.WORKFLOW.value}:
+                answer = json.dumps(response["data"]["outputs"], ensure_ascii=False)
+            else:
+                raise ValueError("Invalid app mode")
+            # Not support image yet
+        return types.CallToolResult(content=[types.TextContent(text=answer, type="text")])
+
+    def retrieve_end_user(self):
+        return (
+            db.session.query(EndUser)
+            .filter(EndUser.external_user_id == self.mcp_server.id, EndUser.type == "mcp")
+            .first()
+        )
+
+    def _convert_input_form_to_parameters(self, user_input_form: list[VariableEntity]):
+        parameters: dict[str, dict[str, Any]] = {}
+        required = []
+        for item in user_input_form:
+            parameters[item.variable] = {}
+            if item.type in (
+                VariableEntityType.FILE,
+                VariableEntityType.FILE_LIST,
+                VariableEntityType.EXTERNAL_DATA_TOOL,
+            ):
+                continue
+            if item.required:
+                required.append(item.variable)
+            # if the workflow republished, the parameters not changed
+            # we should not raise error here
+            try:
+                description = self.mcp_server.parameters_dict[item.variable]
+            except KeyError:
+                description = ""
+            parameters[item.variable]["description"] = description
+            if item.type in (VariableEntityType.TEXT_INPUT, VariableEntityType.PARAGRAPH):
+                parameters[item.variable]["type"] = "string"
+            elif item.type == VariableEntityType.SELECT:
+                parameters[item.variable]["type"] = "string"
+                parameters[item.variable]["enum"] = item.options
+            elif item.type == VariableEntityType.NUMBER:
+                parameters[item.variable]["type"] = "float"
+        return parameters, required

api/core/mcp/session/base_session.py

@@ -0,0 +1,397 @@
+import logging
+import queue
+from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import ExitStack
+from datetime import timedelta
+from types import TracebackType
+from typing import Any, Generic, Self, TypeVar
+
+from httpx import HTTPStatusError
+from pydantic import BaseModel
+
+from core.mcp.error import MCPAuthError, MCPConnectionError
+from core.mcp.types import (
+    CancelledNotification,
+    ClientNotification,
+    ClientRequest,
+    ClientResult,
+    ErrorData,
+    JSONRPCError,
+    JSONRPCMessage,
+    JSONRPCNotification,
+    JSONRPCRequest,
+    JSONRPCResponse,
+    MessageMetadata,
+    RequestId,
+    RequestParams,
+    ServerMessageMetadata,
+    ServerNotification,
+    ServerRequest,
+    ServerResult,
+    SessionMessage,
+)
+
+SendRequestT = TypeVar("SendRequestT", ClientRequest, ServerRequest)
+SendResultT = TypeVar("SendResultT", ClientResult, ServerResult)
+SendNotificationT = TypeVar("SendNotificationT", ClientNotification, ServerNotification)
+ReceiveRequestT = TypeVar("ReceiveRequestT", ClientRequest, ServerRequest)
+ReceiveResultT = TypeVar("ReceiveResultT", bound=BaseModel)
+ReceiveNotificationT = TypeVar("ReceiveNotificationT", ClientNotification, ServerNotification)
+DEFAULT_RESPONSE_READ_TIMEOUT = 1.0
+
+
+class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
+    """Handles responding to MCP requests and manages request lifecycle.
+
+    This class MUST be used as a context manager to ensure proper cleanup and
+    cancellation handling:
+
+    Example:
+        with request_responder as resp:
+            resp.respond(result)
+
+    The context manager ensures:
+    1. Proper cancellation scope setup and cleanup
+    2. Request completion tracking
+    3. Cleanup of in-flight requests
+    """
+
+    request: ReceiveRequestT
+    _session: Any
+    _on_complete: Callable[["RequestResponder[ReceiveRequestT, SendResultT]"], Any]
+
+    def __init__(
+        self,
+        request_id: RequestId,
+        request_meta: RequestParams.Meta | None,
+        request: ReceiveRequestT,
+        session: """BaseSession[
+            SendRequestT,
+            SendNotificationT,
+            SendResultT,
+            ReceiveRequestT,
+            ReceiveNotificationT
+        ]""",
+        on_complete: Callable[["RequestResponder[ReceiveRequestT, SendResultT]"], Any],
+    ) -> None:
+        self.request_id = request_id
+        self.request_meta = request_meta
+        self.request = request
+        self._session = session
+        self._completed = False
+        self._on_complete = on_complete
+        self._entered = False  # Track if we're in a context manager
+
+    def __enter__(self) -> "RequestResponder[ReceiveRequestT, SendResultT]":
+        """Enter the context manager, enabling request cancellation tracking."""
+        self._entered = True
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exit the context manager, performing cleanup and notifying completion."""
+        try:
+            if self._completed:
+                self._on_complete(self)
+        finally:
+            self._entered = False
+
+    def respond(self, response: SendResultT | ErrorData) -> None:
+        """Send a response for this request.
+
+        Must be called within a context manager block.
+        Raises:
+            RuntimeError: If not used within a context manager
+            AssertionError: If request was already responded to
+        """
+        if not self._entered:
+            raise RuntimeError("RequestResponder must be used as a context manager")
+        assert not self._completed, "Request already responded to"
+
+        self._completed = True
+
+        self._session._send_response(request_id=self.request_id, response=response)
+
+    def cancel(self) -> None:
+        """Cancel this request and mark it as completed."""
+        if not self._entered:
+            raise RuntimeError("RequestResponder must be used as a context manager")
+
+        self._completed = True  # Mark as completed so it's removed from in_flight
+        # Send an error response to indicate cancellation
+        self._session._send_response(
+            request_id=self.request_id,
+            response=ErrorData(code=0, message="Request cancelled", data=None),
+        )
+
+
+class BaseSession(
+    Generic[
+        SendRequestT,
+        SendNotificationT,
+        SendResultT,
+        ReceiveRequestT,
+        ReceiveNotificationT,
+    ],
+):
+    """
+    Implements an MCP "session" on top of read/write streams, including features
+    like request/response linking, notifications, and progress.
+
+    This class is a context manager that automatically starts processing
+    messages when entered.
+    """
+
+    _response_streams: dict[RequestId, queue.Queue[JSONRPCResponse | JSONRPCError]]
+    _request_id: int
+    _in_flight: dict[RequestId, RequestResponder[ReceiveRequestT, SendResultT]]
+    _receive_request_type: type[ReceiveRequestT]
+    _receive_notification_type: type[ReceiveNotificationT]
+
+    def __init__(
+        self,
+        read_stream: queue.Queue,
+        write_stream: queue.Queue,
+        receive_request_type: type[ReceiveRequestT],
+        receive_notification_type: type[ReceiveNotificationT],
+        # If none, reading will never time out
+        read_timeout_seconds: timedelta | None = None,
+    ) -> None:
+        self._read_stream = read_stream
+        self._write_stream = write_stream
+        self._response_streams = {}
+        self._request_id = 0
+        self._receive_request_type = receive_request_type
+        self._receive_notification_type = receive_notification_type
+        self._session_read_timeout_seconds = read_timeout_seconds
+        self._in_flight = {}
+        self._exit_stack = ExitStack()
+
+    def __enter__(self) -> Self:
+        self._executor = ThreadPoolExecutor()
+        self._receiver_future = self._executor.submit(self._receive_loop)
+        return self
+
+    def check_receiver_status(self) -> None:
+        if self._receiver_future.done():
+            self._receiver_future.result()
+
+    def __exit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
+    ) -> None:
+        self._exit_stack.close()
+        self._read_stream.put(None)
+        self._write_stream.put(None)
+
+    def send_request(
+        self,
+        request: SendRequestT,
+        result_type: type[ReceiveResultT],
+        request_read_timeout_seconds: timedelta | None = None,
+        metadata: MessageMetadata = None,
+    ) -> ReceiveResultT:
+        """
+        Sends a request and wait for a response. Raises an McpError if the
+        response contains an error. If a request read timeout is provided, it
+        will take precedence over the session read timeout.
+
+        Do not use this method to emit notifications! Use send_notification()
+        instead.
+        """
+        self.check_receiver_status()
+
+        request_id = self._request_id
+        self._request_id = request_id + 1
+
+        response_queue: queue.Queue[JSONRPCResponse | JSONRPCError] = queue.Queue()
+        self._response_streams[request_id] = response_queue
+
+        try:
+            jsonrpc_request = JSONRPCRequest(
+                jsonrpc="2.0",
+                id=request_id,
+                **request.model_dump(by_alias=True, mode="json", exclude_none=True),
+            )
+
+            self._write_stream.put(SessionMessage(message=JSONRPCMessage(jsonrpc_request), metadata=metadata))
+            timeout = DEFAULT_RESPONSE_READ_TIMEOUT
+            if request_read_timeout_seconds is not None:
+                timeout = float(request_read_timeout_seconds.total_seconds())
+            elif self._session_read_timeout_seconds is not None:
+                timeout = float(self._session_read_timeout_seconds.total_seconds())
+            while True:
+                try:
+                    response_or_error = response_queue.get(timeout=timeout)
+                    break
+                except queue.Empty:
+                    self.check_receiver_status()
+                    continue
+
+            if response_or_error is None:
+                raise MCPConnectionError(
+                    ErrorData(
+                        code=500,
+                        message="No response received",
+                    )
+                )
+            elif isinstance(response_or_error, JSONRPCError):
+                if response_or_error.error.code == 401:
+                    raise MCPAuthError(
+                        ErrorData(code=response_or_error.error.code, message=response_or_error.error.message)
+                    )
+                else:
+                    raise MCPConnectionError(
+                        ErrorData(code=response_or_error.error.code, message=response_or_error.error.message)
+                    )
+            else:
+                return result_type.model_validate(response_or_error.result)
+
+        finally:
+            self._response_streams.pop(request_id, None)
+
+    def send_notification(
+        self,
+        notification: SendNotificationT,
+        related_request_id: RequestId | None = None,
+    ) -> None:
+        """
+        Emits a notification, which is a one-way message that does not expect
+        a response.
+        """
+        self.check_receiver_status()
+
+        # Some transport implementations may need to set the related_request_id
+        # to attribute to the notifications to the request that triggered them.
+        jsonrpc_notification = JSONRPCNotification(
+            jsonrpc="2.0",
+            **notification.model_dump(by_alias=True, mode="json", exclude_none=True),
+        )
+        session_message = SessionMessage(
+            message=JSONRPCMessage(jsonrpc_notification),
+            metadata=ServerMessageMetadata(related_request_id=related_request_id) if related_request_id else None,
+        )
+        self._write_stream.put(session_message)
+
+    def _send_response(self, request_id: RequestId, response: SendResultT | ErrorData) -> None:
+        if isinstance(response, ErrorData):
+            jsonrpc_error = JSONRPCError(jsonrpc="2.0", id=request_id, error=response)
+            session_message = SessionMessage(message=JSONRPCMessage(jsonrpc_error))
+            self._write_stream.put(session_message)
+        else:
+            jsonrpc_response = JSONRPCResponse(
+                jsonrpc="2.0",
+                id=request_id,
+                result=response.model_dump(by_alias=True, mode="json", exclude_none=True),
+            )
+            session_message = SessionMessage(message=JSONRPCMessage(jsonrpc_response))
+            self._write_stream.put(session_message)
+
+    def _receive_loop(self) -> None:
+        """
+        Main message processing loop.
+        In a real synchronous implementation, this would likely run in a separate thread.
+        """
+        while True:
+            try:
+                # Attempt to receive a message (this would be blocking in a synchronous context)
+                message = self._read_stream.get(timeout=DEFAULT_RESPONSE_READ_TIMEOUT)
+                if message is None:
+                    break
+                if isinstance(message, HTTPStatusError):
+                    response_queue = self._response_streams.get(self._request_id - 1)
+                    if response_queue is not None:
+                        response_queue.put(
+                            JSONRPCError(
+                                jsonrpc="2.0",
+                                id=self._request_id - 1,
+                                error=ErrorData(code=message.response.status_code, message=message.args[0]),
+                            )
+                        )
+                    else:
+                        self._handle_incoming(RuntimeError(f"Received response with an unknown request ID: {message}"))
+                elif isinstance(message, Exception):
+                    self._handle_incoming(message)
+                elif isinstance(message.message.root, JSONRPCRequest):
+                    validated_request = self._receive_request_type.model_validate(
+                        message.message.root.model_dump(by_alias=True, mode="json", exclude_none=True)
+                    )
+
+                    responder = RequestResponder(
+                        request_id=message.message.root.id,
+                        request_meta=validated_request.root.params.meta if validated_request.root.params else None,
+                        request=validated_request,
+                        session=self,
+                        on_complete=lambda r: self._in_flight.pop(r.request_id, None),
+                    )
+
+                    self._in_flight[responder.request_id] = responder
+                    self._received_request(responder)
+
+                    if not responder._completed:
+                        self._handle_incoming(responder)
+
+                elif isinstance(message.message.root, JSONRPCNotification):
+                    try:
+                        notification = self._receive_notification_type.model_validate(
+                            message.message.root.model_dump(by_alias=True, mode="json", exclude_none=True)
+                        )
+                        # Handle cancellation notifications
+                        if isinstance(notification.root, CancelledNotification):
+                            cancelled_id = notification.root.params.requestId
+                            if cancelled_id in self._in_flight:
+                                self._in_flight[cancelled_id].cancel()
+                        else:
+                            self._received_notification(notification)
+                            self._handle_incoming(notification)
+                    except Exception as e:
+                        # For other validation errors, log and continue
+                        logging.warning(f"Failed to validate notification: {e}. Message was: {message.message.root}")
+                else:  # Response or error
+                    response_queue = self._response_streams.get(message.message.root.id)
+                    if response_queue is not None:
+                        response_queue.put(message.message.root)
+                    else:
+                        self._handle_incoming(RuntimeError(f"Server Error: {message}"))
+            except queue.Empty:
+                continue
+            except Exception as e:
+                logging.exception("Error in message processing loop")
+                raise
+
+    def _received_request(self, responder: RequestResponder[ReceiveRequestT, SendResultT]) -> None:
+        """
+        Can be overridden by subclasses to handle a request without needing to
+        listen on the message stream.
+
+        If the request is responded to within this method, it will not be
+        forwarded on to the message stream.
+        """
+        pass
+
+    def _received_notification(self, notification: ReceiveNotificationT) -> None:
+        """
+        Can be overridden by subclasses to handle a notification without needing
+        to listen on the message stream.
+        """
+        pass
+
+    def send_progress_notification(
+        self, progress_token: str | int, progress: float, total: float | None = None
+    ) -> None:
+        """
+        Sends a progress notification for a request that is currently being
+        processed.
+        """
+        pass
+
+    def _handle_incoming(
+        self,
+        req: RequestResponder[ReceiveRequestT, SendResultT] | ReceiveNotificationT | Exception,
+    ) -> None:
+        """A generic handler for incoming messages. Overwritten by subclasses."""
+        pass

api/core/mcp/session/client_session.py

@@ -0,0 +1,365 @@
+from datetime import timedelta
+from typing import Any, Protocol
+
+from pydantic import AnyUrl, TypeAdapter
+
+from configs import dify_config
+from core.mcp import types
+from core.mcp.entities import SUPPORTED_PROTOCOL_VERSIONS, RequestContext
+from core.mcp.session.base_session import BaseSession, RequestResponder
+
+DEFAULT_CLIENT_INFO = types.Implementation(name="Dify", version=dify_config.project.version)
+
+
+class SamplingFnT(Protocol):
+    def __call__(
+        self,
+        context: RequestContext["ClientSession", Any],
+        params: types.CreateMessageRequestParams,
+    ) -> types.CreateMessageResult | types.ErrorData: ...
+
+
+class ListRootsFnT(Protocol):
+    def __call__(self, context: RequestContext["ClientSession", Any]) -> types.ListRootsResult | types.ErrorData: ...
+
+
+class LoggingFnT(Protocol):
+    def __call__(
+        self,
+        params: types.LoggingMessageNotificationParams,
+    ) -> None: ...
+
+
+class MessageHandlerFnT(Protocol):
+    def __call__(
+        self,
+        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
+    ) -> None: ...
+
+
+def _default_message_handler(
+    message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
+) -> None:
+    if isinstance(message, Exception):
+        raise ValueError(str(message))
+    elif isinstance(message, (types.ServerNotification | RequestResponder)):
+        pass
+
+
+def _default_sampling_callback(
+    context: RequestContext["ClientSession", Any],
+    params: types.CreateMessageRequestParams,
+) -> types.CreateMessageResult | types.ErrorData:
+    return types.ErrorData(
+        code=types.INVALID_REQUEST,
+        message="Sampling not supported",
+    )
+
+
+def _default_list_roots_callback(
+    context: RequestContext["ClientSession", Any],
+) -> types.ListRootsResult | types.ErrorData:
+    return types.ErrorData(
+        code=types.INVALID_REQUEST,
+        message="List roots not supported",
+    )
+
+
+def _default_logging_callback(
+    params: types.LoggingMessageNotificationParams,
+) -> None:
+    pass
+
+
+ClientResponse: TypeAdapter[types.ClientResult | types.ErrorData] = TypeAdapter(types.ClientResult | types.ErrorData)
+
+
+class ClientSession(
+    BaseSession[
+        types.ClientRequest,
+        types.ClientNotification,
+        types.ClientResult,
+        types.ServerRequest,
+        types.ServerNotification,
+    ]
+):
+    def __init__(
+        self,
+        read_stream,
+        write_stream,
+        read_timeout_seconds: timedelta | None = None,
+        sampling_callback: SamplingFnT | None = None,
+        list_roots_callback: ListRootsFnT | None = None,
+        logging_callback: LoggingFnT | None = None,
+        message_handler: MessageHandlerFnT | None = None,
+        client_info: types.Implementation | None = None,
+    ) -> None:
+        super().__init__(
+            read_stream,
+            write_stream,
+            types.ServerRequest,
+            types.ServerNotification,
+            read_timeout_seconds=read_timeout_seconds,
+        )
+        self._client_info = client_info or DEFAULT_CLIENT_INFO
+        self._sampling_callback = sampling_callback or _default_sampling_callback
+        self._list_roots_callback = list_roots_callback or _default_list_roots_callback
+        self._logging_callback = logging_callback or _default_logging_callback
+        self._message_handler = message_handler or _default_message_handler
+
+    def initialize(self) -> types.InitializeResult:
+        sampling = types.SamplingCapability()
+        roots = types.RootsCapability(
+            # TODO: Should this be based on whether we
+            # _will_ send notifications, or only whether
+            # they're supported?
+            listChanged=True,
+        )
+
+        result = self.send_request(
+            types.ClientRequest(
+                types.InitializeRequest(
+                    method="initialize",
+                    params=types.InitializeRequestParams(
+                        protocolVersion=types.LATEST_PROTOCOL_VERSION,
+                        capabilities=types.ClientCapabilities(
+                            sampling=sampling,
+                            experimental=None,
+                            roots=roots,
+                        ),
+                        clientInfo=self._client_info,
+                    ),
+                )
+            ),
+            types.InitializeResult,
+        )
+
+        if result.protocolVersion not in SUPPORTED_PROTOCOL_VERSIONS:
+            raise RuntimeError(f"Unsupported protocol version from the server: {result.protocolVersion}")
+
+        self.send_notification(
+            types.ClientNotification(types.InitializedNotification(method="notifications/initialized"))
+        )
+
+        return result
+
+    def send_ping(self) -> types.EmptyResult:
+        """Send a ping request."""
+        return self.send_request(
+            types.ClientRequest(
+                types.PingRequest(
+                    method="ping",
+                )
+            ),
+            types.EmptyResult,
+        )
+
+    def send_progress_notification(
+        self, progress_token: str | int, progress: float, total: float | None = None
+    ) -> None:
+        """Send a progress notification."""
+        self.send_notification(
+            types.ClientNotification(
+                types.ProgressNotification(
+                    method="notifications/progress",
+                    params=types.ProgressNotificationParams(
+                        progressToken=progress_token,
+                        progress=progress,
+                        total=total,
+                    ),
+                ),
+            )
+        )
+
+    def set_logging_level(self, level: types.LoggingLevel) -> types.EmptyResult:
+        """Send a logging/setLevel request."""
+        return self.send_request(
+            types.ClientRequest(
+                types.SetLevelRequest(
+                    method="logging/setLevel",
+                    params=types.SetLevelRequestParams(level=level),
+                )
+            ),
+            types.EmptyResult,
+        )
+
+    def list_resources(self) -> types.ListResourcesResult:
+        """Send a resources/list request."""
+        return self.send_request(
+            types.ClientRequest(
+                types.ListResourcesRequest(
+                    method="resources/list",
+                )
+            ),
+            types.ListResourcesResult,
+        )
+
+    def list_resource_templates(self) -> types.ListResourceTemplatesResult:
+        """Send a resources/templates/list request."""
+        return self.send_request(
+            types.ClientRequest(
+                types.ListResourceTemplatesRequest(
+                    method="resources/templates/list",
+                )
+            ),
+            types.ListResourceTemplatesResult,
+        )
+
+    def read_resource(self, uri: AnyUrl) -> types.ReadResourceResult:
+        """Send a resources/read request."""
+        return self.send_request(
+            types.ClientRequest(
+                types.ReadResourceRequest(
+                    method="resources/read",
+                    params=types.ReadResourceRequestParams(uri=uri),
+                )
+            ),
+            types.ReadResourceResult,
+        )
+
+    def subscribe_resource(self, uri: AnyUrl) -> types.EmptyResult:
+        """Send a resources/subscribe request."""
+        return self.send_request(
+            types.ClientRequest(
+                types.SubscribeRequest(
+                    method="resources/subscribe",
+                    params=types.SubscribeRequestParams(uri=uri),
+                )
+            ),
+            types.EmptyResult,
+        )
+
+    def unsubscribe_resource(self, uri: AnyUrl) -> types.EmptyResult:
+        """Send a resources/unsubscribe request."""
+        return self.send_request(
+            types.ClientRequest(
+                types.UnsubscribeRequest(
+                    method="resources/unsubscribe",
+                    params=types.UnsubscribeRequestParams(uri=uri),
+                )
+            ),
+            types.EmptyResult,
+        )
+
+    def call_tool(
+        self,
+        name: str,
+        arguments: dict[str, Any] | None = None,
+        read_timeout_seconds: timedelta | None = None,
+    ) -> types.CallToolResult:
+        """Send a tools/call request."""
+
+        return self.send_request(
+            types.ClientRequest(
+                types.CallToolRequest(
+                    method="tools/call",
+                    params=types.CallToolRequestParams(name=name, arguments=arguments),
+                )
+            ),
+            types.CallToolResult,
+            request_read_timeout_seconds=read_timeout_seconds,
+        )
+
+    def list_prompts(self) -> types.ListPromptsResult:
+        """Send a prompts/list request."""
+        return self.send_request(
+            types.ClientRequest(
+                types.ListPromptsRequest(
+                    method="prompts/list",
+                )
+            ),
+            types.ListPromptsResult,
+        )
+
+    def get_prompt(self, name: str, arguments: dict[str, str] | None = None) -> types.GetPromptResult:
+        """Send a prompts/get request."""
+        return self.send_request(
+            types.ClientRequest(
+                types.GetPromptRequest(
+                    method="prompts/get",
+                    params=types.GetPromptRequestParams(name=name, arguments=arguments),
+                )
+            ),
+            types.GetPromptResult,
+        )
+
+    def complete(
+        self,
+        ref: types.ResourceReference | types.PromptReference,
+        argument: dict[str, str],
+    ) -> types.CompleteResult:
+        """Send a completion/complete request."""
+        return self.send_request(
+            types.ClientRequest(
+                types.CompleteRequest(
+                    method="completion/complete",
+                    params=types.CompleteRequestParams(
+                        ref=ref,
+                        argument=types.CompletionArgument(**argument),
+                    ),
+                )
+            ),
+            types.CompleteResult,
+        )
+
+    def list_tools(self) -> types.ListToolsResult:
+        """Send a tools/list request."""
+        return self.send_request(
+            types.ClientRequest(
+                types.ListToolsRequest(
+                    method="tools/list",
+                )
+            ),
+            types.ListToolsResult,
+        )
+
+    def send_roots_list_changed(self) -> None:
+        """Send a roots/list_changed notification."""
+        self.send_notification(
+            types.ClientNotification(
+                types.RootsListChangedNotification(
+                    method="notifications/roots/list_changed",
+                )
+            )
+        )
+
+    def _received_request(self, responder: RequestResponder[types.ServerRequest, types.ClientResult]) -> None:
+        ctx = RequestContext[ClientSession, Any](
+            request_id=responder.request_id,
+            meta=responder.request_meta,
+            session=self,
+            lifespan_context=None,
+        )
+
+        match responder.request.root:
+            case types.CreateMessageRequest(params=params):
+                with responder:
+                    response = self._sampling_callback(ctx, params)
+                    client_response = ClientResponse.validate_python(response)
+                    responder.respond(client_response)
+
+            case types.ListRootsRequest():
+                with responder:
+                    list_roots_response = self._list_roots_callback(ctx)
+                    client_response = ClientResponse.validate_python(list_roots_response)
+                    responder.respond(client_response)
+
+            case types.PingRequest():
+                with responder:
+                    return responder.respond(types.ClientResult(root=types.EmptyResult()))
+
+    def _handle_incoming(
+        self,
+        req: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
+    ) -> None:
+        """Handle incoming messages by forwarding to the message handler."""
+        self._message_handler(req)
+
+    def _received_notification(self, notification: types.ServerNotification) -> None:
+        """Handle notifications from the server."""
+        # Process specific notification types
+        match notification.root:
+            case types.LoggingMessageNotification(params=params):
+                self._logging_callback(params)
+            case _:
+                pass

api/core/mcp/utils.py

@@ -0,0 +1,114 @@
+import json
+
+import httpx
+
+from configs import dify_config
+from core.mcp.types import ErrorData, JSONRPCError
+from core.model_runtime.utils.encoders import jsonable_encoder
+
+HTTP_REQUEST_NODE_SSL_VERIFY = dify_config.HTTP_REQUEST_NODE_SSL_VERIFY
+
+STATUS_FORCELIST = [429, 500, 502, 503, 504]
+
+
+def create_ssrf_proxy_mcp_http_client(
+    headers: dict[str, str] | None = None,
+    timeout: httpx.Timeout | None = None,
+) -> httpx.Client:
+    """Create an HTTPX client with SSRF proxy configuration for MCP connections.
+
+    Args:
+        headers: Optional headers to include in the client
+        timeout: Optional timeout configuration
+
+    Returns:
+        Configured httpx.Client with proxy settings
+    """
+    if dify_config.SSRF_PROXY_ALL_URL:
+        return httpx.Client(
+            verify=HTTP_REQUEST_NODE_SSL_VERIFY,
+            headers=headers or {},
+            timeout=timeout,
+            follow_redirects=True,
+            proxy=dify_config.SSRF_PROXY_ALL_URL,
+        )
+    elif dify_config.SSRF_PROXY_HTTP_URL and dify_config.SSRF_PROXY_HTTPS_URL:
+        proxy_mounts = {
+            "http://": httpx.HTTPTransport(proxy=dify_config.SSRF_PROXY_HTTP_URL, verify=HTTP_REQUEST_NODE_SSL_VERIFY),
+            "https://": httpx.HTTPTransport(
+                proxy=dify_config.SSRF_PROXY_HTTPS_URL, verify=HTTP_REQUEST_NODE_SSL_VERIFY
+            ),
+        }
+        return httpx.Client(
+            verify=HTTP_REQUEST_NODE_SSL_VERIFY,
+            headers=headers or {},
+            timeout=timeout,
+            follow_redirects=True,
+            mounts=proxy_mounts,
+        )
+    else:
+        return httpx.Client(
+            verify=HTTP_REQUEST_NODE_SSL_VERIFY,
+            headers=headers or {},
+            timeout=timeout,
+            follow_redirects=True,
+        )
+
+
+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
+    to prevent SSRF attacks when connecting to external endpoints.
+
+    Args:
+        url: The SSE endpoint URL
+        **kwargs: Additional arguments passed to the SSE connection
+
+    Returns:
+        EventSource object for SSE streaming
+    """
+    from httpx_sse import connect_sse
+
+    # Extract client if provided, otherwise create one
+    client = kwargs.pop("client", None)
+    if client is None:
+        # Create client with SSRF proxy configuration
+        timeout = kwargs.pop(
+            "timeout",
+            httpx.Timeout(
+                timeout=dify_config.SSRF_DEFAULT_TIME_OUT,
+                connect=dify_config.SSRF_DEFAULT_CONNECT_TIME_OUT,
+                read=dify_config.SSRF_DEFAULT_READ_TIME_OUT,
+                write=dify_config.SSRF_DEFAULT_WRITE_TIME_OUT,
+            ),
+        )
+        headers = kwargs.pop("headers", {})
+        client = create_ssrf_proxy_mcp_http_client(headers=headers, timeout=timeout)
+        client_provided = False
+    else:
+        client_provided = True
+
+    # Extract method if provided, default to GET
+    method = kwargs.pop("method", "GET")
+
+    try:
+        return connect_sse(client, method, url, **kwargs)
+    except Exception:
+        # If we created the client, we need to clean it up on error
+        if not client_provided:
+            client.close()
+        raise
+
+
+def create_mcp_error_response(request_id: int | str | None, code: int, message: str, data=None):
+    """Create MCP error response"""
+    error_data = ErrorData(code=code, message=message, data=data)
+    json_response = JSONRPCError(
+        jsonrpc="2.0",
+        id=request_id or 1,
+        error=error_data,
+    )
+    json_data = json.dumps(jsonable_encoder(json_response))
+    sse_content = f"event: message\ndata: {json_data}\n\n".encode()
+    yield sse_content

api/core/model_runtime/entities/llm_entities.py

@@ -53,6 +53,37 @@ class LLMUsage(ModelUsage):
             latency=0.0,
         )
 
+    @classmethod
+    def from_metadata(cls, metadata: dict) -> "LLMUsage":
+        """
+        Create LLMUsage instance from metadata dictionary with default values.
+
+        Args:
+            metadata: Dictionary containing usage metadata
+
+        Returns:
+            LLMUsage instance with values from metadata or defaults
+        """
+        total_tokens = metadata.get("total_tokens", 0)
+        completion_tokens = metadata.get("completion_tokens", 0)
+        if total_tokens > 0 and completion_tokens == 0:
+            completion_tokens = total_tokens
+
+        return cls(
+            prompt_tokens=metadata.get("prompt_tokens", 0),
+            completion_tokens=completion_tokens,
+            total_tokens=total_tokens,
+            prompt_unit_price=Decimal(str(metadata.get("prompt_unit_price", 0))),
+            completion_unit_price=Decimal(str(metadata.get("completion_unit_price", 0))),
+            total_price=Decimal(str(metadata.get("total_price", 0))),
+            currency=metadata.get("currency", "USD"),
+            prompt_price_unit=Decimal(str(metadata.get("prompt_price_unit", 0))),
+            completion_price_unit=Decimal(str(metadata.get("completion_price_unit", 0))),
+            prompt_price=Decimal(str(metadata.get("prompt_price", 0))),
+            completion_price=Decimal(str(metadata.get("completion_price", 0))),
+            latency=metadata.get("latency", 0.0),
+        )
+
     def plus(self, other: "LLMUsage") -> "LLMUsage":
         """
         Add two LLMUsage instances together.

api/core/model_runtime/entities/provider_entities.py

@@ -123,6 +123,8 @@ class ProviderEntity(BaseModel):
     description: Optional[I18nObject] = None
     icon_small: Optional[I18nObject] = None
     icon_large: Optional[I18nObject] = None
+    icon_small_dark: Optional[I18nObject] = None
+    icon_large_dark: Optional[I18nObject] = None
     background: Optional[str] = None
     help: Optional[ProviderHelpEntity] = None
     supported_model_types: Sequence[ModelType]

api/core/ops/aliyun_trace/aliyun_trace.py

@@ -0,0 +1,487 @@
+import json
+import logging
+from collections.abc import Sequence
+from typing import Optional
+from urllib.parse import urljoin
+
+from opentelemetry.trace import Status, StatusCode
+from sqlalchemy.orm import Session, sessionmaker
+
+from core.ops.aliyun_trace.data_exporter.traceclient import (
+    TraceClient,
+    convert_datetime_to_nanoseconds,
+    convert_to_span_id,
+    convert_to_trace_id,
+    generate_span_id,
+)
+from core.ops.aliyun_trace.entities.aliyun_trace_entity import SpanData
+from core.ops.aliyun_trace.entities.semconv import (
+    GEN_AI_COMPLETION,
+    GEN_AI_FRAMEWORK,
+    GEN_AI_MODEL_NAME,
+    GEN_AI_PROMPT,
+    GEN_AI_PROMPT_TEMPLATE_TEMPLATE,
+    GEN_AI_PROMPT_TEMPLATE_VARIABLE,
+    GEN_AI_RESPONSE_FINISH_REASON,
+    GEN_AI_SESSION_ID,
+    GEN_AI_SPAN_KIND,
+    GEN_AI_SYSTEM,
+    GEN_AI_USAGE_INPUT_TOKENS,
+    GEN_AI_USAGE_OUTPUT_TOKENS,
+    GEN_AI_USAGE_TOTAL_TOKENS,
+    GEN_AI_USER_ID,
+    INPUT_VALUE,
+    OUTPUT_VALUE,
+    RETRIEVAL_DOCUMENT,
+    RETRIEVAL_QUERY,
+    TOOL_DESCRIPTION,
+    TOOL_NAME,
+    TOOL_PARAMETERS,
+    GenAISpanKind,
+)
+from core.ops.base_trace_instance import BaseTraceInstance
+from core.ops.entities.config_entity import AliyunConfig
+from core.ops.entities.trace_entity import (
+    BaseTraceInfo,
+    DatasetRetrievalTraceInfo,
+    GenerateNameTraceInfo,
+    MessageTraceInfo,
+    ModerationTraceInfo,
+    SuggestedQuestionTraceInfo,
+    ToolTraceInfo,
+    WorkflowTraceInfo,
+)
+from core.rag.models.document import Document
+from core.repositories import SQLAlchemyWorkflowNodeExecutionRepository
+from core.workflow.entities.workflow_node_execution import (
+    WorkflowNodeExecution,
+    WorkflowNodeExecutionMetadataKey,
+    WorkflowNodeExecutionStatus,
+)
+from core.workflow.nodes import NodeType
+from models import Account, App, EndUser, TenantAccountJoin, WorkflowNodeExecutionTriggeredFrom, db
+
+logger = logging.getLogger(__name__)
+
+
+class AliyunDataTrace(BaseTraceInstance):
+    def __init__(
+        self,
+        aliyun_config: AliyunConfig,
+    ):
+        super().__init__(aliyun_config)
+        base_url = aliyun_config.endpoint.rstrip("/")
+        endpoint = urljoin(base_url, f"adapt_{aliyun_config.license_key}/api/otlp/traces")
+        self.trace_client = TraceClient(service_name=aliyun_config.app_name, endpoint=endpoint)
+
+    def trace(self, trace_info: BaseTraceInfo):
+        if isinstance(trace_info, WorkflowTraceInfo):
+            self.workflow_trace(trace_info)
+        if isinstance(trace_info, MessageTraceInfo):
+            self.message_trace(trace_info)
+        if isinstance(trace_info, ModerationTraceInfo):
+            pass
+        if isinstance(trace_info, SuggestedQuestionTraceInfo):
+            self.suggested_question_trace(trace_info)
+        if isinstance(trace_info, DatasetRetrievalTraceInfo):
+            self.dataset_retrieval_trace(trace_info)
+        if isinstance(trace_info, ToolTraceInfo):
+            self.tool_trace(trace_info)
+        if isinstance(trace_info, GenerateNameTraceInfo):
+            pass
+
+    def api_check(self):
+        return self.trace_client.api_check()
+
+    def get_project_url(self):
+        try:
+            return self.trace_client.get_project_url()
+        except Exception as e:
+            logger.info(f"Aliyun get run url failed: {str(e)}", exc_info=True)
+            raise ValueError(f"Aliyun get run url failed: {str(e)}")
+
+    def workflow_trace(self, trace_info: WorkflowTraceInfo):
+        trace_id = convert_to_trace_id(trace_info.workflow_run_id)
+        workflow_span_id = convert_to_span_id(trace_info.workflow_run_id, "workflow")
+        self.add_workflow_span(trace_id, workflow_span_id, trace_info)
+
+        workflow_node_executions = self.get_workflow_node_executions(trace_info)
+        for node_execution in workflow_node_executions:
+            node_span = self.build_workflow_node_span(node_execution, trace_id, trace_info, workflow_span_id)
+            self.trace_client.add_span(node_span)
+
+    def message_trace(self, trace_info: MessageTraceInfo):
+        message_data = trace_info.message_data
+        if message_data is None:
+            return
+        message_id = trace_info.message_id
+
+        user_id = message_data.from_account_id
+        if message_data.from_end_user_id:
+            end_user_data: Optional[EndUser] = (
+                db.session.query(EndUser).filter(EndUser.id == message_data.from_end_user_id).first()
+            )
+            if end_user_data is not None:
+                user_id = end_user_data.session_id
+
+        status: Status = Status(StatusCode.OK)
+        if trace_info.error:
+            status = Status(StatusCode.ERROR, trace_info.error)
+
+        trace_id = convert_to_trace_id(message_id)
+        message_span_id = convert_to_span_id(message_id, "message")
+        message_span = SpanData(
+            trace_id=trace_id,
+            parent_span_id=None,
+            span_id=message_span_id,
+            name="message",
+            start_time=convert_datetime_to_nanoseconds(trace_info.start_time),
+            end_time=convert_datetime_to_nanoseconds(trace_info.end_time),
+            attributes={
+                GEN_AI_SESSION_ID: trace_info.metadata.get("conversation_id", ""),
+                GEN_AI_USER_ID: str(user_id),
+                GEN_AI_SPAN_KIND: GenAISpanKind.CHAIN.value,
+                GEN_AI_FRAMEWORK: "dify",
+                INPUT_VALUE: json.dumps(trace_info.inputs, ensure_ascii=False),
+                OUTPUT_VALUE: str(trace_info.outputs),
+            },
+            status=status,
+        )
+        self.trace_client.add_span(message_span)
+
+        app_model_config = getattr(trace_info.message_data, "app_model_config", {})
+        pre_prompt = getattr(app_model_config, "pre_prompt", "")
+        inputs_data = getattr(trace_info.message_data, "inputs", {})
+        llm_span = SpanData(
+            trace_id=trace_id,
+            parent_span_id=message_span_id,
+            span_id=convert_to_span_id(message_id, "llm"),
+            name="llm",
+            start_time=convert_datetime_to_nanoseconds(trace_info.start_time),
+            end_time=convert_datetime_to_nanoseconds(trace_info.end_time),
+            attributes={
+                GEN_AI_SESSION_ID: trace_info.metadata.get("conversation_id", ""),
+                GEN_AI_USER_ID: str(user_id),
+                GEN_AI_SPAN_KIND: GenAISpanKind.LLM.value,
+                GEN_AI_FRAMEWORK: "dify",
+                GEN_AI_MODEL_NAME: trace_info.metadata.get("ls_model_name", ""),
+                GEN_AI_SYSTEM: trace_info.metadata.get("ls_provider", ""),
+                GEN_AI_USAGE_INPUT_TOKENS: str(trace_info.message_tokens),
+                GEN_AI_USAGE_OUTPUT_TOKENS: str(trace_info.answer_tokens),
+                GEN_AI_USAGE_TOTAL_TOKENS: str(trace_info.total_tokens),
+                GEN_AI_PROMPT_TEMPLATE_VARIABLE: json.dumps(inputs_data, ensure_ascii=False),
+                GEN_AI_PROMPT_TEMPLATE_TEMPLATE: pre_prompt,
+                GEN_AI_PROMPT: json.dumps(trace_info.inputs, ensure_ascii=False),
+                GEN_AI_COMPLETION: str(trace_info.outputs),
+                INPUT_VALUE: json.dumps(trace_info.inputs, ensure_ascii=False),
+                OUTPUT_VALUE: str(trace_info.outputs),
+            },
+            status=status,
+        )
+        self.trace_client.add_span(llm_span)
+
+    def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo):
+        if trace_info.message_data is None:
+            return
+        message_id = trace_info.message_id
+
+        documents_data = extract_retrieval_documents(trace_info.documents)
+        dataset_retrieval_span = SpanData(
+            trace_id=convert_to_trace_id(message_id),
+            parent_span_id=convert_to_span_id(message_id, "message"),
+            span_id=generate_span_id(),
+            name="dataset_retrieval",
+            start_time=convert_datetime_to_nanoseconds(trace_info.start_time),
+            end_time=convert_datetime_to_nanoseconds(trace_info.end_time),
+            attributes={
+                GEN_AI_SPAN_KIND: GenAISpanKind.RETRIEVER.value,
+                GEN_AI_FRAMEWORK: "dify",
+                RETRIEVAL_QUERY: str(trace_info.inputs),
+                RETRIEVAL_DOCUMENT: json.dumps(documents_data, ensure_ascii=False),
+                INPUT_VALUE: str(trace_info.inputs),
+                OUTPUT_VALUE: json.dumps(documents_data, ensure_ascii=False),
+            },
+        )
+        self.trace_client.add_span(dataset_retrieval_span)
+
+    def tool_trace(self, trace_info: ToolTraceInfo):
+        if trace_info.message_data is None:
+            return
+        message_id = trace_info.message_id
+
+        status: Status = Status(StatusCode.OK)
+        if trace_info.error:
+            status = Status(StatusCode.ERROR, trace_info.error)
+
+        tool_span = SpanData(
+            trace_id=convert_to_trace_id(message_id),
+            parent_span_id=convert_to_span_id(message_id, "message"),
+            span_id=generate_span_id(),
+            name=trace_info.tool_name,
+            start_time=convert_datetime_to_nanoseconds(trace_info.start_time),
+            end_time=convert_datetime_to_nanoseconds(trace_info.end_time),
+            attributes={
+                GEN_AI_SPAN_KIND: GenAISpanKind.TOOL.value,
+                GEN_AI_FRAMEWORK: "dify",
+                TOOL_NAME: trace_info.tool_name,
+                TOOL_DESCRIPTION: json.dumps(trace_info.tool_config, ensure_ascii=False),
+                TOOL_PARAMETERS: json.dumps(trace_info.tool_inputs, ensure_ascii=False),
+                INPUT_VALUE: json.dumps(trace_info.inputs, ensure_ascii=False),
+                OUTPUT_VALUE: str(trace_info.tool_outputs),
+            },
+            status=status,
+        )
+        self.trace_client.add_span(tool_span)
+
+    def get_workflow_node_executions(self, trace_info: WorkflowTraceInfo) -> Sequence[WorkflowNodeExecution]:
+        # through workflow_run_id get all_nodes_execution using repository
+        session_factory = sessionmaker(bind=db.engine)
+        # Find the app's creator account
+        with Session(db.engine, expire_on_commit=False) as session:
+            # Get the app to find its creator
+            app_id = trace_info.metadata.get("app_id")
+            if not app_id:
+                raise ValueError("No app_id found in trace_info metadata")
+
+            app = session.query(App).filter(App.id == app_id).first()
+            if not app:
+                raise ValueError(f"App with id {app_id} not found")
+
+            if not app.created_by:
+                raise ValueError(f"App with id {app_id} has no creator (created_by is None)")
+
+            service_account = session.query(Account).filter(Account.id == app.created_by).first()
+            if not service_account:
+                raise ValueError(f"Creator account with id {app.created_by} not found for app {app_id}")
+            current_tenant = (
+                session.query(TenantAccountJoin).filter_by(account_id=service_account.id, current=True).first()
+            )
+            if not current_tenant:
+                raise ValueError(f"Current tenant not found for account {service_account.id}")
+            service_account.set_tenant_id(current_tenant.tenant_id)
+        workflow_node_execution_repository = SQLAlchemyWorkflowNodeExecutionRepository(
+            session_factory=session_factory,
+            user=service_account,
+            app_id=trace_info.metadata.get("app_id"),
+            triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
+        )
+        # Get all executions for this workflow run
+        workflow_node_executions = workflow_node_execution_repository.get_by_workflow_run(
+            workflow_run_id=trace_info.workflow_run_id
+        )
+        return workflow_node_executions
+
+    def build_workflow_node_span(
+        self, node_execution: WorkflowNodeExecution, trace_id: int, trace_info: WorkflowTraceInfo, workflow_span_id: int
+    ):
+        try:
+            if node_execution.node_type == NodeType.LLM:
+                node_span = self.build_workflow_llm_span(trace_id, workflow_span_id, trace_info, node_execution)
+            elif node_execution.node_type == NodeType.KNOWLEDGE_RETRIEVAL:
+                node_span = self.build_workflow_retrieval_span(trace_id, workflow_span_id, trace_info, node_execution)
+            elif node_execution.node_type == NodeType.TOOL:
+                node_span = self.build_workflow_tool_span(trace_id, workflow_span_id, trace_info, node_execution)
+            else:
+                node_span = self.build_workflow_task_span(trace_id, workflow_span_id, trace_info, node_execution)
+            return node_span
+        except Exception:
+            return None
+
+    def get_workflow_node_status(self, node_execution: WorkflowNodeExecution) -> Status:
+        span_status: Status = Status(StatusCode.UNSET)
+        if node_execution.status == WorkflowNodeExecutionStatus.SUCCEEDED:
+            span_status = Status(StatusCode.OK)
+        elif node_execution.status in [WorkflowNodeExecutionStatus.FAILED, WorkflowNodeExecutionStatus.EXCEPTION]:
+            span_status = Status(StatusCode.ERROR, str(node_execution.error))
+        return span_status
+
+    def build_workflow_task_span(
+        self, trace_id: int, workflow_span_id: int, trace_info: WorkflowTraceInfo, node_execution: WorkflowNodeExecution
+    ) -> SpanData:
+        return SpanData(
+            trace_id=trace_id,
+            parent_span_id=workflow_span_id,
+            span_id=convert_to_span_id(node_execution.id, "node"),
+            name=node_execution.title,
+            start_time=convert_datetime_to_nanoseconds(node_execution.created_at),
+            end_time=convert_datetime_to_nanoseconds(node_execution.finished_at),
+            attributes={
+                GEN_AI_SESSION_ID: trace_info.metadata.get("conversation_id", ""),
+                GEN_AI_SPAN_KIND: GenAISpanKind.TASK.value,
+                GEN_AI_FRAMEWORK: "dify",
+                INPUT_VALUE: json.dumps(node_execution.inputs, ensure_ascii=False),
+                OUTPUT_VALUE: json.dumps(node_execution.outputs, ensure_ascii=False),
+            },
+            status=self.get_workflow_node_status(node_execution),
+        )
+
+    def build_workflow_tool_span(
+        self, trace_id: int, workflow_span_id: int, trace_info: WorkflowTraceInfo, node_execution: WorkflowNodeExecution
+    ) -> SpanData:
+        tool_des = {}
+        if node_execution.metadata:
+            tool_des = node_execution.metadata.get(WorkflowNodeExecutionMetadataKey.TOOL_INFO, {})
+        return SpanData(
+            trace_id=trace_id,
+            parent_span_id=workflow_span_id,
+            span_id=convert_to_span_id(node_execution.id, "node"),
+            name=node_execution.title,
+            start_time=convert_datetime_to_nanoseconds(node_execution.created_at),
+            end_time=convert_datetime_to_nanoseconds(node_execution.finished_at),
+            attributes={
+                GEN_AI_SPAN_KIND: GenAISpanKind.TOOL.value,
+                GEN_AI_FRAMEWORK: "dify",
+                TOOL_NAME: node_execution.title,
+                TOOL_DESCRIPTION: json.dumps(tool_des, ensure_ascii=False),
+                TOOL_PARAMETERS: json.dumps(node_execution.inputs if node_execution.inputs else {}, ensure_ascii=False),
+                INPUT_VALUE: json.dumps(node_execution.inputs if node_execution.inputs else {}, ensure_ascii=False),
+                OUTPUT_VALUE: json.dumps(node_execution.outputs, ensure_ascii=False),
+            },
+            status=self.get_workflow_node_status(node_execution),
+        )
+
+    def build_workflow_retrieval_span(
+        self, trace_id: int, workflow_span_id: int, trace_info: WorkflowTraceInfo, node_execution: WorkflowNodeExecution
+    ) -> SpanData:
+        input_value = ""
+        if node_execution.inputs:
+            input_value = str(node_execution.inputs.get("query", ""))
+        output_value = ""
+        if node_execution.outputs:
+            output_value = json.dumps(node_execution.outputs.get("result", []), ensure_ascii=False)
+        return SpanData(
+            trace_id=trace_id,
+            parent_span_id=workflow_span_id,
+            span_id=convert_to_span_id(node_execution.id, "node"),
+            name=node_execution.title,
+            start_time=convert_datetime_to_nanoseconds(node_execution.created_at),
+            end_time=convert_datetime_to_nanoseconds(node_execution.finished_at),
+            attributes={
+                GEN_AI_SPAN_KIND: GenAISpanKind.RETRIEVER.value,
+                GEN_AI_FRAMEWORK: "dify",
+                RETRIEVAL_QUERY: input_value,
+                RETRIEVAL_DOCUMENT: output_value,
+                INPUT_VALUE: input_value,
+                OUTPUT_VALUE: output_value,
+            },
+            status=self.get_workflow_node_status(node_execution),
+        )
+
+    def build_workflow_llm_span(
+        self, trace_id: int, workflow_span_id: int, trace_info: WorkflowTraceInfo, node_execution: WorkflowNodeExecution
+    ) -> SpanData:
+        process_data = node_execution.process_data or {}
+        outputs = node_execution.outputs or {}
+        usage_data = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {})
+        return SpanData(
+            trace_id=trace_id,
+            parent_span_id=workflow_span_id,
+            span_id=convert_to_span_id(node_execution.id, "node"),
+            name=node_execution.title,
+            start_time=convert_datetime_to_nanoseconds(node_execution.created_at),
+            end_time=convert_datetime_to_nanoseconds(node_execution.finished_at),
+            attributes={
+                GEN_AI_SESSION_ID: trace_info.metadata.get("conversation_id", ""),
+                GEN_AI_SPAN_KIND: GenAISpanKind.LLM.value,
+                GEN_AI_FRAMEWORK: "dify",
+                GEN_AI_MODEL_NAME: process_data.get("model_name", ""),
+                GEN_AI_SYSTEM: process_data.get("model_provider", ""),
+                GEN_AI_USAGE_INPUT_TOKENS: str(usage_data.get("prompt_tokens", 0)),
+                GEN_AI_USAGE_OUTPUT_TOKENS: str(usage_data.get("completion_tokens", 0)),
+                GEN_AI_USAGE_TOTAL_TOKENS: str(usage_data.get("total_tokens", 0)),
+                GEN_AI_PROMPT: json.dumps(process_data.get("prompts", []), ensure_ascii=False),
+                GEN_AI_COMPLETION: str(outputs.get("text", "")),
+                GEN_AI_RESPONSE_FINISH_REASON: outputs.get("finish_reason", ""),
+                INPUT_VALUE: json.dumps(process_data.get("prompts", []), ensure_ascii=False),
+                OUTPUT_VALUE: str(outputs.get("text", "")),
+            },
+            status=self.get_workflow_node_status(node_execution),
+        )
+
+    def add_workflow_span(self, trace_id: int, workflow_span_id: int, trace_info: WorkflowTraceInfo):
+        message_span_id = None
+        if trace_info.message_id:
+            message_span_id = convert_to_span_id(trace_info.message_id, "message")
+        user_id = trace_info.metadata.get("user_id")
+        status: Status = Status(StatusCode.OK)
+        if trace_info.error:
+            status = Status(StatusCode.ERROR, trace_info.error)
+        if message_span_id:  # chatflow
+            message_span = SpanData(
+                trace_id=trace_id,
+                parent_span_id=None,
+                span_id=message_span_id,
+                name="message",
+                start_time=convert_datetime_to_nanoseconds(trace_info.start_time),
+                end_time=convert_datetime_to_nanoseconds(trace_info.end_time),
+                attributes={
+                    GEN_AI_SESSION_ID: trace_info.metadata.get("conversation_id", ""),
+                    GEN_AI_USER_ID: str(user_id),
+                    GEN_AI_SPAN_KIND: GenAISpanKind.CHAIN.value,
+                    GEN_AI_FRAMEWORK: "dify",
+                    INPUT_VALUE: trace_info.workflow_run_inputs.get("sys.query", ""),
+                    OUTPUT_VALUE: json.dumps(trace_info.workflow_run_outputs, ensure_ascii=False),
+                },
+                status=status,
+            )
+            self.trace_client.add_span(message_span)
+
+        workflow_span = SpanData(
+            trace_id=trace_id,
+            parent_span_id=message_span_id,
+            span_id=workflow_span_id,
+            name="workflow",
+            start_time=convert_datetime_to_nanoseconds(trace_info.start_time),
+            end_time=convert_datetime_to_nanoseconds(trace_info.end_time),
+            attributes={
+                GEN_AI_USER_ID: str(user_id),
+                GEN_AI_SPAN_KIND: GenAISpanKind.CHAIN.value,
+                GEN_AI_FRAMEWORK: "dify",
+                INPUT_VALUE: json.dumps(trace_info.workflow_run_inputs, ensure_ascii=False),
+                OUTPUT_VALUE: json.dumps(trace_info.workflow_run_outputs, ensure_ascii=False),
+            },
+            status=status,
+        )
+        self.trace_client.add_span(workflow_span)
+
+    def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo):
+        message_id = trace_info.message_id
+        status: Status = Status(StatusCode.OK)
+        if trace_info.error:
+            status = Status(StatusCode.ERROR, trace_info.error)
+        suggested_question_span = SpanData(
+            trace_id=convert_to_trace_id(message_id),
+            parent_span_id=convert_to_span_id(message_id, "message"),
+            span_id=convert_to_span_id(message_id, "suggested_question"),
+            name="suggested_question",
+            start_time=convert_datetime_to_nanoseconds(trace_info.start_time),
+            end_time=convert_datetime_to_nanoseconds(trace_info.end_time),
+            attributes={
+                GEN_AI_SPAN_KIND: GenAISpanKind.LLM.value,
+                GEN_AI_FRAMEWORK: "dify",
+                GEN_AI_MODEL_NAME: trace_info.metadata.get("ls_model_name", ""),
+                GEN_AI_SYSTEM: trace_info.metadata.get("ls_provider", ""),
+                GEN_AI_PROMPT: json.dumps(trace_info.inputs, ensure_ascii=False),
+                GEN_AI_COMPLETION: json.dumps(trace_info.suggested_question, ensure_ascii=False),
+                INPUT_VALUE: json.dumps(trace_info.inputs, ensure_ascii=False),
+                OUTPUT_VALUE: json.dumps(trace_info.suggested_question, ensure_ascii=False),
+            },
+            status=status,
+        )
+        self.trace_client.add_span(suggested_question_span)
+
+
+def extract_retrieval_documents(documents: list[Document]):
+    documents_data = []
+    for document in documents:
+        document_data = {
+            "content": document.page_content,
+            "metadata": {
+                "dataset_id": document.metadata.get("dataset_id"),
+                "doc_id": document.metadata.get("doc_id"),
+                "document_id": document.metadata.get("document_id"),
+            },
+            "score": document.metadata.get("score"),
+        }
+        documents_data.append(document_data)
+    return documents_data

api/core/ops/aliyun_trace/data_exporter/traceclient.py

@@ -0,0 +1,200 @@
+import hashlib
+import logging
+import random
+import socket
+import threading
+import uuid
+from collections import deque
+from collections.abc import Sequence
+from datetime import datetime
+from typing import Optional
+
+import requests
+from opentelemetry import trace as trace_api
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.util.instrumentation import InstrumentationScope
+from opentelemetry.semconv.resource import ResourceAttributes
+
+from configs import dify_config
+from core.ops.aliyun_trace.entities.aliyun_trace_entity import SpanData
+
+INVALID_SPAN_ID = 0x0000000000000000
+INVALID_TRACE_ID = 0x00000000000000000000000000000000
+
+logger = logging.getLogger(__name__)
+
+
+class TraceClient:
+    def __init__(
+        self,
+        service_name: str,
+        endpoint: str,
+        max_queue_size: int = 1000,
+        schedule_delay_sec: int = 5,
+        max_export_batch_size: int = 50,
+    ):
+        self.endpoint = endpoint
+        self.resource = Resource(
+            attributes={
+                ResourceAttributes.SERVICE_NAME: service_name,
+                ResourceAttributes.SERVICE_VERSION: f"dify-{dify_config.project.version}-{dify_config.COMMIT_SHA}",
+                ResourceAttributes.DEPLOYMENT_ENVIRONMENT: f"{dify_config.DEPLOY_ENV}-{dify_config.EDITION}",
+                ResourceAttributes.HOST_NAME: socket.gethostname(),
+            }
+        )
+        self.span_builder = SpanBuilder(self.resource)
+        self.exporter = OTLPSpanExporter(endpoint=endpoint)
+
+        self.max_queue_size = max_queue_size
+        self.schedule_delay_sec = schedule_delay_sec
+        self.max_export_batch_size = max_export_batch_size
+
+        self.queue: deque = deque(maxlen=max_queue_size)
+        self.condition = threading.Condition(threading.Lock())
+        self.done = False
+
+        self.worker_thread = threading.Thread(target=self._worker, daemon=True)
+        self.worker_thread.start()
+
+        self._spans_dropped = False
+
+    def export(self, spans: Sequence[ReadableSpan]):
+        self.exporter.export(spans)
+
+    def api_check(self):
+        try:
+            response = requests.head(self.endpoint, timeout=5)
+            if response.status_code == 405:
+                return True
+            else:
+                logger.debug(f"AliyunTrace API check failed: Unexpected status code: {response.status_code}")
+                return False
+        except requests.exceptions.RequestException as e:
+            logger.debug(f"AliyunTrace API check failed: {str(e)}")
+            raise ValueError(f"AliyunTrace API check failed: {str(e)}")
+
+    def get_project_url(self):
+        return "https://arms.console.aliyun.com/#/llm"
+
+    def add_span(self, span_data: SpanData):
+        if span_data is None:
+            return
+        span: ReadableSpan = self.span_builder.build_span(span_data)
+        with self.condition:
+            if len(self.queue) == self.max_queue_size:
+                if not self._spans_dropped:
+                    logger.warning("Queue is full, likely spans will be dropped.")
+                    self._spans_dropped = True
+
+            self.queue.appendleft(span)
+            if len(self.queue) >= self.max_export_batch_size:
+                self.condition.notify()
+
+    def _worker(self):
+        while not self.done:
+            with self.condition:
+                if len(self.queue) < self.max_export_batch_size and not self.done:
+                    self.condition.wait(timeout=self.schedule_delay_sec)
+            self._export_batch()
+
+    def _export_batch(self):
+        spans_to_export: list[ReadableSpan] = []
+        with self.condition:
+            while len(spans_to_export) < self.max_export_batch_size and self.queue:
+                spans_to_export.append(self.queue.pop())
+
+        if spans_to_export:
+            try:
+                self.exporter.export(spans_to_export)
+            except Exception as e:
+                logger.debug(f"Error exporting spans: {e}")
+
+    def shutdown(self):
+        with self.condition:
+            self.done = True
+            self.condition.notify_all()
+        self.worker_thread.join()
+        self._export_batch()
+        self.exporter.shutdown()
+
+
+class SpanBuilder:
+    def __init__(self, resource):
+        self.resource = resource
+        self.instrumentation_scope = InstrumentationScope(
+            __name__,
+            "",
+            None,
+            None,
+        )
+
+    def build_span(self, span_data: SpanData) -> ReadableSpan:
+        span_context = trace_api.SpanContext(
+            trace_id=span_data.trace_id,
+            span_id=span_data.span_id,
+            is_remote=False,
+            trace_flags=trace_api.TraceFlags(trace_api.TraceFlags.SAMPLED),
+            trace_state=None,
+        )
+
+        parent_span_context = None
+        if span_data.parent_span_id is not None:
+            parent_span_context = trace_api.SpanContext(
+                trace_id=span_data.trace_id,
+                span_id=span_data.parent_span_id,
+                is_remote=False,
+                trace_flags=trace_api.TraceFlags(trace_api.TraceFlags.SAMPLED),
+                trace_state=None,
+            )
+
+        span = ReadableSpan(
+            name=span_data.name,
+            context=span_context,
+            parent=parent_span_context,
+            resource=self.resource,
+            attributes=span_data.attributes,
+            events=span_data.events,
+            links=span_data.links,
+            kind=trace_api.SpanKind.INTERNAL,
+            status=span_data.status,
+            start_time=span_data.start_time,
+            end_time=span_data.end_time,
+            instrumentation_scope=self.instrumentation_scope,
+        )
+        return span
+
+
+def generate_span_id() -> int:
+    span_id = random.getrandbits(64)
+    while span_id == INVALID_SPAN_ID:
+        span_id = random.getrandbits(64)
+    return span_id
+
+
+def convert_to_trace_id(uuid_v4: Optional[str]) -> int:
+    try:
+        uuid_obj = uuid.UUID(uuid_v4)
+        return uuid_obj.int
+    except Exception as e:
+        raise ValueError(f"Invalid UUID input: {e}")
+
+
+def convert_to_span_id(uuid_v4: Optional[str], span_type: str) -> int:
+    try:
+        uuid_obj = uuid.UUID(uuid_v4)
+    except Exception as e:
+        raise ValueError(f"Invalid UUID input: {e}")
+    combined_key = f"{uuid_obj.hex}-{span_type}"
+    hash_bytes = hashlib.sha256(combined_key.encode("utf-8")).digest()
+    span_id = int.from_bytes(hash_bytes[:8], byteorder="big", signed=False)
+    return span_id
+
+
+def convert_datetime_to_nanoseconds(start_time_a: Optional[datetime]) -> Optional[int]:
+    if start_time_a is None:
+        return None
+    timestamp_in_seconds = start_time_a.timestamp()
+    timestamp_in_nanoseconds = int(timestamp_in_seconds * 1e9)
+    return timestamp_in_nanoseconds

api/core/ops/aliyun_trace/entities/aliyun_trace_entity.py

@@ -0,0 +1,21 @@
+from collections.abc import Sequence
+from typing import Optional
+
+from opentelemetry import trace as trace_api
+from opentelemetry.sdk.trace import Event, Status, StatusCode
+from pydantic import BaseModel, Field
+
+
+class SpanData(BaseModel):
+    model_config = {"arbitrary_types_allowed": True}
+
+    trace_id: int = Field(..., description="The unique identifier for the trace.")
+    parent_span_id: Optional[int] = Field(None, description="The ID of the parent span, if any.")
+    span_id: int = Field(..., description="The unique identifier for this span.")
+    name: str = Field(..., description="The name of the span.")
+    attributes: dict[str, str] = Field(default_factory=dict, description="Attributes associated with the span.")
+    events: Sequence[Event] = Field(default_factory=list, description="Events recorded in the span.")
+    links: Sequence[trace_api.Link] = Field(default_factory=list, description="Links to other spans.")
+    status: Status = Field(default=Status(StatusCode.UNSET), description="The status of the span.")
+    start_time: Optional[int] = Field(..., description="The start time of the span in nanoseconds.")
+    end_time: Optional[int] = Field(..., description="The end time of the span in nanoseconds.")

api/core/ops/aliyun_trace/entities/semconv.py

@@ -0,0 +1,64 @@
+from enum import Enum
+
+# public
+GEN_AI_SESSION_ID = "gen_ai.session.id"
+
+GEN_AI_USER_ID = "gen_ai.user.id"
+
+GEN_AI_USER_NAME = "gen_ai.user.name"
+
+GEN_AI_SPAN_KIND = "gen_ai.span.kind"
+
+GEN_AI_FRAMEWORK = "gen_ai.framework"
+
+
+# Chain
+INPUT_VALUE = "input.value"
+
+OUTPUT_VALUE = "output.value"
+
+
+# Retriever
+RETRIEVAL_QUERY = "retrieval.query"
+
+RETRIEVAL_DOCUMENT = "retrieval.document"
+
+
+# LLM
+GEN_AI_MODEL_NAME = "gen_ai.model_name"
+
+GEN_AI_SYSTEM = "gen_ai.system"
+
+GEN_AI_USAGE_INPUT_TOKENS = "gen_ai.usage.input_tokens"
+
+GEN_AI_USAGE_OUTPUT_TOKENS = "gen_ai.usage.output_tokens"
+
+GEN_AI_USAGE_TOTAL_TOKENS = "gen_ai.usage.total_tokens"
+
+GEN_AI_PROMPT_TEMPLATE_TEMPLATE = "gen_ai.prompt_template.template"
+
+GEN_AI_PROMPT_TEMPLATE_VARIABLE = "gen_ai.prompt_template.variable"
+
+GEN_AI_PROMPT = "gen_ai.prompt"
+
+GEN_AI_COMPLETION = "gen_ai.completion"
+
+GEN_AI_RESPONSE_FINISH_REASON = "gen_ai.response.finish_reason"
+
+# Tool
+TOOL_NAME = "tool.name"
+
+TOOL_DESCRIPTION = "tool.description"
+
+TOOL_PARAMETERS = "tool.parameters"
+
+
+class GenAISpanKind(Enum):
+    CHAIN = "CHAIN"
+    RETRIEVER = "RETRIEVER"
+    RERANKER = "RERANKER"
+    LLM = "LLM"
+    EMBEDDING = "EMBEDDING"
+    TOOL = "TOOL"
+    AGENT = "AGENT"
+    TASK = "TASK"

api/core/ops/arize_phoenix_trace/arize_phoenix_trace.py

@@ -0,0 +1,726 @@
+import hashlib
+import json
+import logging
+import os
+from datetime import datetime, timedelta
+from typing import Optional, Union, cast
+
+from openinference.semconv.trace import OpenInferenceSpanKindValues, SpanAttributes
+from opentelemetry import trace
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter as GrpcOTLPSpanExporter
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter as HttpOTLPSpanExporter
+from opentelemetry.sdk import trace as trace_sdk
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
+from opentelemetry.trace import SpanContext, TraceFlags, TraceState
+
+from core.ops.base_trace_instance import BaseTraceInstance
+from core.ops.entities.config_entity import ArizeConfig, PhoenixConfig
+from core.ops.entities.trace_entity import (
+    BaseTraceInfo,
+    DatasetRetrievalTraceInfo,
+    GenerateNameTraceInfo,
+    MessageTraceInfo,
+    ModerationTraceInfo,
+    SuggestedQuestionTraceInfo,
+    ToolTraceInfo,
+    TraceTaskName,
+    WorkflowTraceInfo,
+)
+from extensions.ext_database import db
+from models.model import EndUser, MessageFile
+from models.workflow import WorkflowNodeExecutionModel
+
+logger = logging.getLogger(__name__)
+
+
+def setup_tracer(arize_phoenix_config: ArizeConfig | PhoenixConfig) -> tuple[trace_sdk.Tracer, SimpleSpanProcessor]:
+    """Configure OpenTelemetry tracer with OTLP exporter for Arize/Phoenix."""
+    try:
+        # Choose the appropriate exporter based on config type
+        exporter: Union[GrpcOTLPSpanExporter, HttpOTLPSpanExporter]
+        if isinstance(arize_phoenix_config, ArizeConfig):
+            arize_endpoint = f"{arize_phoenix_config.endpoint}/v1"
+            arize_headers = {
+                "api_key": arize_phoenix_config.api_key or "",
+                "space_id": arize_phoenix_config.space_id or "",
+                "authorization": f"Bearer {arize_phoenix_config.api_key or ''}",
+            }
+            exporter = GrpcOTLPSpanExporter(
+                endpoint=arize_endpoint,
+                headers=arize_headers,
+                timeout=30,
+            )
+        else:
+            phoenix_endpoint = f"{arize_phoenix_config.endpoint}/v1/traces"
+            phoenix_headers = {
+                "api_key": arize_phoenix_config.api_key or "",
+                "authorization": f"Bearer {arize_phoenix_config.api_key or ''}",
+            }
+            exporter = HttpOTLPSpanExporter(
+                endpoint=phoenix_endpoint,
+                headers=phoenix_headers,
+                timeout=30,
+            )
+
+        attributes = {
+            "openinference.project.name": arize_phoenix_config.project or "",
+            "model_id": arize_phoenix_config.project or "",
+        }
+        resource = Resource(attributes=attributes)
+        provider = trace_sdk.TracerProvider(resource=resource)
+        processor = SimpleSpanProcessor(
+            exporter,
+        )
+        provider.add_span_processor(processor)
+
+        # Create a named tracer instead of setting the global provider
+        tracer_name = f"arize_phoenix_tracer_{arize_phoenix_config.project}"
+        logger.info(f"[Arize/Phoenix] Created tracer with name: {tracer_name}")
+        return cast(trace_sdk.Tracer, provider.get_tracer(tracer_name)), processor
+    except Exception as e:
+        logger.error(f"[Arize/Phoenix] Failed to setup the tracer: {str(e)}", exc_info=True)
+        raise
+
+
+def datetime_to_nanos(dt: Optional[datetime]) -> int:
+    """Convert datetime to nanoseconds since epoch. If None, use current time."""
+    if dt is None:
+        dt = datetime.now()
+    return int(dt.timestamp() * 1_000_000_000)
+
+
+def uuid_to_trace_id(string: Optional[str]) -> int:
+    """Convert UUID string to a valid trace ID (16-byte integer)."""
+    if string is None:
+        string = ""
+    hash_object = hashlib.sha256(string.encode())
+
+    # Take the first 16 bytes (128 bits) of the hash
+    digest = hash_object.digest()[:16]
+
+    # Convert to integer (128 bits)
+    return int.from_bytes(digest, byteorder="big")
+
+
+class ArizePhoenixDataTrace(BaseTraceInstance):
+    def __init__(
+        self,
+        arize_phoenix_config: ArizeConfig | PhoenixConfig,
+    ):
+        super().__init__(arize_phoenix_config)
+        import logging
+
+        logging.basicConfig()
+        logging.getLogger().setLevel(logging.DEBUG)
+        self.arize_phoenix_config = arize_phoenix_config
+        self.tracer, self.processor = setup_tracer(arize_phoenix_config)
+        self.project = arize_phoenix_config.project
+        self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001")
+
+    def trace(self, trace_info: BaseTraceInfo):
+        logger.info(f"[Arize/Phoenix] Trace: {trace_info}")
+        try:
+            if isinstance(trace_info, WorkflowTraceInfo):
+                self.workflow_trace(trace_info)
+            if isinstance(trace_info, MessageTraceInfo):
+                self.message_trace(trace_info)
+            if isinstance(trace_info, ModerationTraceInfo):
+                self.moderation_trace(trace_info)
+            if isinstance(trace_info, SuggestedQuestionTraceInfo):
+                self.suggested_question_trace(trace_info)
+            if isinstance(trace_info, DatasetRetrievalTraceInfo):
+                self.dataset_retrieval_trace(trace_info)
+            if isinstance(trace_info, ToolTraceInfo):
+                self.tool_trace(trace_info)
+            if isinstance(trace_info, GenerateNameTraceInfo):
+                self.generate_name_trace(trace_info)
+
+        except Exception as e:
+            logger.error(f"[Arize/Phoenix] Error in the trace: {str(e)}", exc_info=True)
+            raise
+
+    def workflow_trace(self, trace_info: WorkflowTraceInfo):
+        if trace_info.message_data is None:
+            return
+
+        workflow_metadata = {
+            "workflow_id": trace_info.workflow_run_id or "",
+            "message_id": trace_info.message_id or "",
+            "workflow_app_log_id": trace_info.workflow_app_log_id or "",
+            "status": trace_info.workflow_run_status or "",
+            "status_message": trace_info.error or "",
+            "level": "ERROR" if trace_info.error else "DEFAULT",
+            "total_tokens": trace_info.total_tokens or 0,
+        }
+        workflow_metadata.update(trace_info.metadata)
+
+        trace_id = uuid_to_trace_id(trace_info.message_id)
+        span_id = RandomIdGenerator().generate_span_id()
+        context = SpanContext(
+            trace_id=trace_id,
+            span_id=span_id,
+            is_remote=False,
+            trace_flags=TraceFlags(TraceFlags.SAMPLED),
+            trace_state=TraceState(),
+        )
+
+        workflow_span = self.tracer.start_span(
+            name=TraceTaskName.WORKFLOW_TRACE.value,
+            attributes={
+                SpanAttributes.INPUT_VALUE: json.dumps(trace_info.workflow_run_inputs, ensure_ascii=False),
+                SpanAttributes.OUTPUT_VALUE: json.dumps(trace_info.workflow_run_outputs, ensure_ascii=False),
+                SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.CHAIN.value,
+                SpanAttributes.METADATA: json.dumps(workflow_metadata, ensure_ascii=False),
+                SpanAttributes.SESSION_ID: trace_info.conversation_id or "",
+            },
+            start_time=datetime_to_nanos(trace_info.start_time),
+            context=trace.set_span_in_context(trace.NonRecordingSpan(context)),
+        )
+
+        try:
+            # Process workflow nodes
+            for node_execution in self._get_workflow_nodes(trace_info.workflow_run_id):
+                created_at = node_execution.created_at or datetime.now()
+                elapsed_time = node_execution.elapsed_time
+                finished_at = created_at + timedelta(seconds=elapsed_time)
+
+                process_data = json.loads(node_execution.process_data) if node_execution.process_data else {}
+
+                node_metadata = {
+                    "node_id": node_execution.id,
+                    "node_type": node_execution.node_type,
+                    "node_status": node_execution.status,
+                    "tenant_id": node_execution.tenant_id,
+                    "app_id": node_execution.app_id,
+                    "app_name": node_execution.title,
+                    "status": node_execution.status,
+                    "level": "ERROR" if node_execution.status != "succeeded" else "DEFAULT",
+                }
+
+                if node_execution.execution_metadata:
+                    node_metadata.update(json.loads(node_execution.execution_metadata))
+
+                # Determine the correct span kind based on node type
+                span_kind = OpenInferenceSpanKindValues.CHAIN.value
+                if node_execution.node_type == "llm":
+                    span_kind = OpenInferenceSpanKindValues.LLM.value
+                    provider = process_data.get("model_provider")
+                    model = process_data.get("model_name")
+                    if provider:
+                        node_metadata["ls_provider"] = provider
+                    if model:
+                        node_metadata["ls_model_name"] = model
+
+                    outputs = json.loads(node_execution.outputs).get("usage", {})
+                    usage_data = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {})
+                    if usage_data:
+                        node_metadata["total_tokens"] = usage_data.get("total_tokens", 0)
+                        node_metadata["prompt_tokens"] = usage_data.get("prompt_tokens", 0)
+                        node_metadata["completion_tokens"] = usage_data.get("completion_tokens", 0)
+                elif node_execution.node_type == "dataset_retrieval":
+                    span_kind = OpenInferenceSpanKindValues.RETRIEVER.value
+                elif node_execution.node_type == "tool":
+                    span_kind = OpenInferenceSpanKindValues.TOOL.value
+                else:
+                    span_kind = OpenInferenceSpanKindValues.CHAIN.value
+
+                node_span = self.tracer.start_span(
+                    name=node_execution.node_type,
+                    attributes={
+                        SpanAttributes.INPUT_VALUE: node_execution.inputs or "{}",
+                        SpanAttributes.OUTPUT_VALUE: node_execution.outputs or "{}",
+                        SpanAttributes.OPENINFERENCE_SPAN_KIND: span_kind,
+                        SpanAttributes.METADATA: json.dumps(node_metadata, ensure_ascii=False),
+                        SpanAttributes.SESSION_ID: trace_info.conversation_id or "",
+                    },
+                    start_time=datetime_to_nanos(created_at),
+                )
+
+                try:
+                    if node_execution.node_type == "llm":
+                        provider = process_data.get("model_provider")
+                        model = process_data.get("model_name")
+                        if provider:
+                            node_span.set_attribute(SpanAttributes.LLM_PROVIDER, provider)
+                        if model:
+                            node_span.set_attribute(SpanAttributes.LLM_MODEL_NAME, model)
+
+                        outputs = json.loads(node_execution.outputs).get("usage", {})
+                        usage_data = (
+                            process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {})
+                        )
+                        if usage_data:
+                            node_span.set_attribute(
+                                SpanAttributes.LLM_TOKEN_COUNT_TOTAL, usage_data.get("total_tokens", 0)
+                            )
+                            node_span.set_attribute(
+                                SpanAttributes.LLM_TOKEN_COUNT_PROMPT, usage_data.get("prompt_tokens", 0)
+                            )
+                            node_span.set_attribute(
+                                SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, usage_data.get("completion_tokens", 0)
+                            )
+                finally:
+                    node_span.end(end_time=datetime_to_nanos(finished_at))
+        finally:
+            workflow_span.end(end_time=datetime_to_nanos(trace_info.end_time))
+
+    def message_trace(self, trace_info: MessageTraceInfo):
+        if trace_info.message_data is None:
+            return
+
+        file_list = cast(list[str], trace_info.file_list) or []
+        message_file_data: Optional[MessageFile] = trace_info.message_file_data
+
+        if message_file_data is not None:
+            file_url = f"{self.file_base_url}/{message_file_data.url}" if message_file_data else ""
+            file_list.append(file_url)
+
+        message_metadata = {
+            "message_id": trace_info.message_id or "",
+            "conversation_mode": str(trace_info.conversation_mode or ""),
+            "user_id": trace_info.message_data.from_account_id or "",
+            "file_list": json.dumps(file_list),
+            "status": trace_info.message_data.status or "",
+            "status_message": trace_info.error or "",
+            "level": "ERROR" if trace_info.error else "DEFAULT",
+            "total_tokens": trace_info.total_tokens or 0,
+            "prompt_tokens": trace_info.message_tokens or 0,
+            "completion_tokens": trace_info.answer_tokens or 0,
+            "ls_provider": trace_info.message_data.model_provider or "",
+            "ls_model_name": trace_info.message_data.model_id or "",
+        }
+        message_metadata.update(trace_info.metadata)
+
+        # Add end user data if available
+        if trace_info.message_data.from_end_user_id:
+            end_user_data: Optional[EndUser] = (
+                db.session.query(EndUser).filter(EndUser.id == trace_info.message_data.from_end_user_id).first()
+            )
+            if end_user_data is not None:
+                message_metadata["end_user_id"] = end_user_data.session_id
+
+        attributes = {
+            SpanAttributes.INPUT_VALUE: trace_info.message_data.query,
+            SpanAttributes.OUTPUT_VALUE: trace_info.message_data.answer,
+            SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.CHAIN.value,
+            SpanAttributes.METADATA: json.dumps(message_metadata, ensure_ascii=False),
+            SpanAttributes.SESSION_ID: trace_info.message_data.conversation_id,
+        }
+
+        trace_id = uuid_to_trace_id(trace_info.message_id)
+        message_span_id = RandomIdGenerator().generate_span_id()
+        span_context = SpanContext(
+            trace_id=trace_id,
+            span_id=message_span_id,
+            is_remote=False,
+            trace_flags=TraceFlags(TraceFlags.SAMPLED),
+            trace_state=TraceState(),
+        )
+
+        message_span = self.tracer.start_span(
+            name=TraceTaskName.MESSAGE_TRACE.value,
+            attributes=attributes,
+            start_time=datetime_to_nanos(trace_info.start_time),
+            context=trace.set_span_in_context(trace.NonRecordingSpan(span_context)),
+        )
+
+        try:
+            if trace_info.error:
+                message_span.add_event(
+                    "exception",
+                    attributes={
+                        "exception.message": trace_info.error,
+                        "exception.type": "Error",
+                        "exception.stacktrace": trace_info.error,
+                    },
+                )
+
+            # Convert outputs to string based on type
+            if isinstance(trace_info.outputs, dict | list):
+                outputs_str = json.dumps(trace_info.outputs, ensure_ascii=False)
+            elif isinstance(trace_info.outputs, str):
+                outputs_str = trace_info.outputs
+            else:
+                outputs_str = str(trace_info.outputs)
+
+            llm_attributes = {
+                SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.LLM.value,
+                SpanAttributes.INPUT_VALUE: json.dumps(trace_info.inputs, ensure_ascii=False),
+                SpanAttributes.OUTPUT_VALUE: outputs_str,
+                SpanAttributes.METADATA: json.dumps(message_metadata, ensure_ascii=False),
+                SpanAttributes.SESSION_ID: trace_info.message_data.conversation_id,
+            }
+
+            if isinstance(trace_info.inputs, list):
+                for i, msg in enumerate(trace_info.inputs):
+                    if isinstance(msg, dict):
+                        llm_attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.{i}.message.content"] = msg.get("text", "")
+                        llm_attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.{i}.message.role"] = msg.get(
+                            "role", "user"
+                        )
+                        # todo: handle assistant and tool role messages, as they don't always
+                        # have a text field, but may have a tool_calls field instead
+                        # e.g. 'tool_calls': [{'id': '98af3a29-b066-45a5-b4b1-46c74ddafc58',
+                        # 'type': 'function', 'function': {'name': 'current_time', 'arguments': '{}'}}]}
+            elif isinstance(trace_info.inputs, dict):
+                llm_attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.message.content"] = json.dumps(trace_info.inputs)
+                llm_attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.message.role"] = "user"
+            elif isinstance(trace_info.inputs, str):
+                llm_attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.message.content"] = trace_info.inputs
+                llm_attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.message.role"] = "user"
+
+            if trace_info.total_tokens is not None and trace_info.total_tokens > 0:
+                llm_attributes[SpanAttributes.LLM_TOKEN_COUNT_TOTAL] = trace_info.total_tokens
+            if trace_info.message_tokens is not None and trace_info.message_tokens > 0:
+                llm_attributes[SpanAttributes.LLM_TOKEN_COUNT_PROMPT] = trace_info.message_tokens
+            if trace_info.answer_tokens is not None and trace_info.answer_tokens > 0:
+                llm_attributes[SpanAttributes.LLM_TOKEN_COUNT_COMPLETION] = trace_info.answer_tokens
+
+            if trace_info.message_data.model_id is not None:
+                llm_attributes[SpanAttributes.LLM_MODEL_NAME] = trace_info.message_data.model_id
+            if trace_info.message_data.model_provider is not None:
+                llm_attributes[SpanAttributes.LLM_PROVIDER] = trace_info.message_data.model_provider
+
+            if trace_info.message_data and trace_info.message_data.message_metadata:
+                metadata_dict = json.loads(trace_info.message_data.message_metadata)
+                if model_params := metadata_dict.get("model_parameters"):
+                    llm_attributes[SpanAttributes.LLM_INVOCATION_PARAMETERS] = json.dumps(model_params)
+
+            llm_span = self.tracer.start_span(
+                name="llm",
+                attributes=llm_attributes,
+                start_time=datetime_to_nanos(trace_info.start_time),
+                context=trace.set_span_in_context(trace.NonRecordingSpan(span_context)),
+            )
+
+            try:
+                if trace_info.error:
+                    llm_span.add_event(
+                        "exception",
+                        attributes={
+                            "exception.message": trace_info.error,
+                            "exception.type": "Error",
+                            "exception.stacktrace": trace_info.error,
+                        },
+                    )
+            finally:
+                llm_span.end(end_time=datetime_to_nanos(trace_info.end_time))
+        finally:
+            message_span.end(end_time=datetime_to_nanos(trace_info.end_time))
+
+    def moderation_trace(self, trace_info: ModerationTraceInfo):
+        if trace_info.message_data is None:
+            return
+
+        metadata = {
+            "message_id": trace_info.message_id,
+            "tool_name": "moderation",
+            "status": trace_info.message_data.status,
+            "status_message": trace_info.message_data.error or "",
+            "level": "ERROR" if trace_info.message_data.error else "DEFAULT",
+        }
+        metadata.update(trace_info.metadata)
+
+        trace_id = uuid_to_trace_id(trace_info.message_id)
+        span_id = RandomIdGenerator().generate_span_id()
+        context = SpanContext(
+            trace_id=trace_id,
+            span_id=span_id,
+            is_remote=False,
+            trace_flags=TraceFlags(TraceFlags.SAMPLED),
+            trace_state=TraceState(),
+        )
+
+        span = self.tracer.start_span(
+            name=TraceTaskName.MODERATION_TRACE.value,
+            attributes={
+                SpanAttributes.INPUT_VALUE: json.dumps(trace_info.inputs, ensure_ascii=False),
+                SpanAttributes.OUTPUT_VALUE: json.dumps(
+                    {
+                        "action": trace_info.action,
+                        "flagged": trace_info.flagged,
+                        "preset_response": trace_info.preset_response,
+                        "inputs": trace_info.inputs,
+                    },
+                    ensure_ascii=False,
+                ),
+                SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.CHAIN.value,
+                SpanAttributes.METADATA: json.dumps(metadata, ensure_ascii=False),
+            },
+            start_time=datetime_to_nanos(trace_info.start_time),
+            context=trace.set_span_in_context(trace.NonRecordingSpan(context)),
+        )
+
+        try:
+            if trace_info.message_data.error:
+                span.add_event(
+                    "exception",
+                    attributes={
+                        "exception.message": trace_info.message_data.error,
+                        "exception.type": "Error",
+                        "exception.stacktrace": trace_info.message_data.error,
+                    },
+                )
+        finally:
+            span.end(end_time=datetime_to_nanos(trace_info.end_time))
+
+    def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo):
+        if trace_info.message_data is None:
+            return
+
+        start_time = trace_info.start_time or trace_info.message_data.created_at
+        end_time = trace_info.end_time or trace_info.message_data.updated_at
+
+        metadata = {
+            "message_id": trace_info.message_id,
+            "tool_name": "suggested_question",
+            "status": trace_info.status,
+            "status_message": trace_info.error or "",
+            "level": "ERROR" if trace_info.error else "DEFAULT",
+            "total_tokens": trace_info.total_tokens,
+            "ls_provider": trace_info.model_provider or "",
+            "ls_model_name": trace_info.model_id or "",
+        }
+        metadata.update(trace_info.metadata)
+
+        trace_id = uuid_to_trace_id(trace_info.message_id)
+        span_id = RandomIdGenerator().generate_span_id()
+        context = SpanContext(
+            trace_id=trace_id,
+            span_id=span_id,
+            is_remote=False,
+            trace_flags=TraceFlags(TraceFlags.SAMPLED),
+            trace_state=TraceState(),
+        )
+
+        span = self.tracer.start_span(
+            name=TraceTaskName.SUGGESTED_QUESTION_TRACE.value,
+            attributes={
+                SpanAttributes.INPUT_VALUE: json.dumps(trace_info.inputs, ensure_ascii=False),
+                SpanAttributes.OUTPUT_VALUE: json.dumps(trace_info.suggested_question, ensure_ascii=False),
+                SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.CHAIN.value,
+                SpanAttributes.METADATA: json.dumps(metadata, ensure_ascii=False),
+            },
+            start_time=datetime_to_nanos(start_time),
+            context=trace.set_span_in_context(trace.NonRecordingSpan(context)),
+        )
+
+        try:
+            if trace_info.error:
+                span.add_event(
+                    "exception",
+                    attributes={
+                        "exception.message": trace_info.error,
+                        "exception.type": "Error",
+                        "exception.stacktrace": trace_info.error,
+                    },
+                )
+        finally:
+            span.end(end_time=datetime_to_nanos(end_time))
+
+    def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo):
+        if trace_info.message_data is None:
+            return
+
+        start_time = trace_info.start_time or trace_info.message_data.created_at
+        end_time = trace_info.end_time or trace_info.message_data.updated_at
+
+        metadata = {
+            "message_id": trace_info.message_id,
+            "tool_name": "dataset_retrieval",
+            "status": trace_info.message_data.status,
+            "status_message": trace_info.message_data.error or "",
+            "level": "ERROR" if trace_info.message_data.error else "DEFAULT",
+            "ls_provider": trace_info.message_data.model_provider or "",
+            "ls_model_name": trace_info.message_data.model_id or "",
+        }
+        metadata.update(trace_info.metadata)
+
+        trace_id = uuid_to_trace_id(trace_info.message_id)
+        span_id = RandomIdGenerator().generate_span_id()
+        context = SpanContext(
+            trace_id=trace_id,
+            span_id=span_id,
+            is_remote=False,
+            trace_flags=TraceFlags(TraceFlags.SAMPLED),
+            trace_state=TraceState(),
+        )
+
+        span = self.tracer.start_span(
+            name=TraceTaskName.DATASET_RETRIEVAL_TRACE.value,
+            attributes={
+                SpanAttributes.INPUT_VALUE: json.dumps(trace_info.inputs, ensure_ascii=False),
+                SpanAttributes.OUTPUT_VALUE: json.dumps({"documents": trace_info.documents}, ensure_ascii=False),
+                SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.RETRIEVER.value,
+                SpanAttributes.METADATA: json.dumps(metadata, ensure_ascii=False),
+                "start_time": start_time.isoformat() if start_time else "",
+                "end_time": end_time.isoformat() if end_time else "",
+            },
+            start_time=datetime_to_nanos(start_time),
+            context=trace.set_span_in_context(trace.NonRecordingSpan(context)),
+        )
+
+        try:
+            if trace_info.message_data.error:
+                span.add_event(
+                    "exception",
+                    attributes={
+                        "exception.message": trace_info.message_data.error,
+                        "exception.type": "Error",
+                        "exception.stacktrace": trace_info.message_data.error,
+                    },
+                )
+        finally:
+            span.end(end_time=datetime_to_nanos(end_time))
+
+    def tool_trace(self, trace_info: ToolTraceInfo):
+        if trace_info.message_data is None:
+            logger.warning("[Arize/Phoenix] Message data is None, skipping tool trace.")
+            return
+
+        metadata = {
+            "message_id": trace_info.message_id,
+            "tool_config": json.dumps(trace_info.tool_config, ensure_ascii=False),
+        }
+
+        trace_id = uuid_to_trace_id(trace_info.message_id)
+        tool_span_id = RandomIdGenerator().generate_span_id()
+        logger.info(f"[Arize/Phoenix] Creating tool trace with trace_id: {trace_id}, span_id: {tool_span_id}")
+
+        # Create span context with the same trace_id as the parent
+        # todo: Create with the appropriate parent span context, so that the tool span is
+        # a child of the appropriate span (e.g. message span)
+        span_context = SpanContext(
+            trace_id=trace_id,
+            span_id=tool_span_id,
+            is_remote=False,
+            trace_flags=TraceFlags(TraceFlags.SAMPLED),
+            trace_state=TraceState(),
+        )
+
+        tool_params_str = (
+            json.dumps(trace_info.tool_parameters, ensure_ascii=False)
+            if isinstance(trace_info.tool_parameters, dict)
+            else str(trace_info.tool_parameters)
+        )
+
+        span = self.tracer.start_span(
+            name=trace_info.tool_name,
+            attributes={
+                SpanAttributes.INPUT_VALUE: json.dumps(trace_info.tool_inputs, ensure_ascii=False),
+                SpanAttributes.OUTPUT_VALUE: trace_info.tool_outputs,
+                SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.TOOL.value,
+                SpanAttributes.METADATA: json.dumps(metadata, ensure_ascii=False),
+                SpanAttributes.TOOL_NAME: trace_info.tool_name,
+                SpanAttributes.TOOL_PARAMETERS: tool_params_str,
+            },
+            start_time=datetime_to_nanos(trace_info.start_time),
+            context=trace.set_span_in_context(trace.NonRecordingSpan(span_context)),
+        )
+
+        try:
+            if trace_info.error:
+                span.add_event(
+                    "exception",
+                    attributes={
+                        "exception.message": trace_info.error,
+                        "exception.type": "Error",
+                        "exception.stacktrace": trace_info.error,
+                    },
+                )
+        finally:
+            span.end(end_time=datetime_to_nanos(trace_info.end_time))
+
+    def generate_name_trace(self, trace_info: GenerateNameTraceInfo):
+        if trace_info.message_data is None:
+            return
+
+        metadata = {
+            "project_name": self.project,
+            "message_id": trace_info.message_id,
+            "status": trace_info.message_data.status,
+            "status_message": trace_info.message_data.error or "",
+            "level": "ERROR" if trace_info.message_data.error else "DEFAULT",
+        }
+        metadata.update(trace_info.metadata)
+
+        trace_id = uuid_to_trace_id(trace_info.message_id)
+        span_id = RandomIdGenerator().generate_span_id()
+        context = SpanContext(
+            trace_id=trace_id,
+            span_id=span_id,
+            is_remote=False,
+            trace_flags=TraceFlags(TraceFlags.SAMPLED),
+            trace_state=TraceState(),
+        )
+
+        span = self.tracer.start_span(
+            name=TraceTaskName.GENERATE_NAME_TRACE.value,
+            attributes={
+                SpanAttributes.INPUT_VALUE: json.dumps(trace_info.inputs, ensure_ascii=False),
+                SpanAttributes.OUTPUT_VALUE: json.dumps(trace_info.outputs, ensure_ascii=False),
+                SpanAttributes.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.CHAIN.value,
+                SpanAttributes.METADATA: json.dumps(metadata, ensure_ascii=False),
+                SpanAttributes.SESSION_ID: trace_info.message_data.conversation_id,
+                "start_time": trace_info.start_time.isoformat() if trace_info.start_time else "",
+                "end_time": trace_info.end_time.isoformat() if trace_info.end_time else "",
+            },
+            start_time=datetime_to_nanos(trace_info.start_time),
+            context=trace.set_span_in_context(trace.NonRecordingSpan(context)),
+        )
+
+        try:
+            if trace_info.message_data.error:
+                span.add_event(
+                    "exception",
+                    attributes={
+                        "exception.message": trace_info.message_data.error,
+                        "exception.type": "Error",
+                        "exception.stacktrace": trace_info.message_data.error,
+                    },
+                )
+        finally:
+            span.end(end_time=datetime_to_nanos(trace_info.end_time))
+
+    def api_check(self):
+        try:
+            with self.tracer.start_span("api_check") as span:
+                span.set_attribute("test", "true")
+            return True
+        except Exception as e:
+            logger.info(f"[Arize/Phoenix] API check failed: {str(e)}", exc_info=True)
+            raise ValueError(f"[Arize/Phoenix] API check failed: {str(e)}")
+
+    def get_project_url(self):
+        try:
+            if self.arize_phoenix_config.endpoint == "https://otlp.arize.com":
+                return "https://app.arize.com/"
+            else:
+                return f"{self.arize_phoenix_config.endpoint}/projects/"
+        except Exception as e:
+            logger.info(f"[Arize/Phoenix] Get run url failed: {str(e)}", exc_info=True)
+            raise ValueError(f"[Arize/Phoenix] Get run url failed: {str(e)}")
+
+    def _get_workflow_nodes(self, workflow_run_id: str):
+        """Helper method to get workflow nodes"""
+        workflow_nodes = (
+            db.session.query(
+                WorkflowNodeExecutionModel.id,
+                WorkflowNodeExecutionModel.tenant_id,
+                WorkflowNodeExecutionModel.app_id,
+                WorkflowNodeExecutionModel.title,
+                WorkflowNodeExecutionModel.node_type,
+                WorkflowNodeExecutionModel.status,
+                WorkflowNodeExecutionModel.inputs,
+                WorkflowNodeExecutionModel.outputs,
+                WorkflowNodeExecutionModel.created_at,
+                WorkflowNodeExecutionModel.elapsed_time,
+                WorkflowNodeExecutionModel.process_data,
+                WorkflowNodeExecutionModel.execution_metadata,
+            )
+            .filter(WorkflowNodeExecutionModel.workflow_run_id == workflow_run_id)
+            .all()
+        )
+        return workflow_nodes

api/core/ops/entities/config_entity.py

@@ -2,20 +2,92 @@ from enum import StrEnum
 
 from pydantic import BaseModel, ValidationInfo, field_validator
 
+from core.ops.utils import validate_project_name, validate_url, validate_url_with_path
+
 
 class TracingProviderEnum(StrEnum):
+    ARIZE = "arize"
+    PHOENIX = "phoenix"
     LANGFUSE = "langfuse"
     LANGSMITH = "langsmith"
     OPIK = "opik"
     WEAVE = "weave"
+    ALIYUN = "aliyun"
 
 
 class BaseTracingConfig(BaseModel):
     """
-    Base model class for tracing
+    Base model class for tracing configurations
     """
 
-    ...
+    @classmethod
+    def validate_endpoint_url(cls, v: str, default_url: str) -> str:
+        """
+        Common endpoint URL validation logic
+
+        Args:
+            v: URL value to validate
+            default_url: Default URL to use if input is None or empty
+
+        Returns:
+            Validated and normalized URL
+        """
+        return validate_url(v, default_url)
+
+    @classmethod
+    def validate_project_field(cls, v: str, default_name: str) -> str:
+        """
+        Common project name validation logic
+
+        Args:
+            v: Project name to validate
+            default_name: Default name to use if input is None or empty
+
+        Returns:
+            Validated project name
+        """
+        return validate_project_name(v, default_name)
+
+
+class ArizeConfig(BaseTracingConfig):
+    """
+    Model class for Arize tracing config.
+    """
+
+    api_key: str | None = None
+    space_id: str | None = None
+    project: str | None = None
+    endpoint: str = "https://otlp.arize.com"
+
+    @field_validator("project")
+    @classmethod
+    def project_validator(cls, v, info: ValidationInfo):
+        return cls.validate_project_field(v, "default")
+
+    @field_validator("endpoint")
+    @classmethod
+    def endpoint_validator(cls, v, info: ValidationInfo):
+        return cls.validate_endpoint_url(v, "https://otlp.arize.com")
+
+
+class PhoenixConfig(BaseTracingConfig):
+    """
+    Model class for Phoenix tracing config.
+    """
+
+    api_key: str | None = None
+    project: str | None = None
+    endpoint: str = "https://app.phoenix.arize.com"
+
+    @field_validator("project")
+    @classmethod
+    def project_validator(cls, v, info: ValidationInfo):
+        return cls.validate_project_field(v, "default")
+
+    @field_validator("endpoint")
+    @classmethod
+    def endpoint_validator(cls, v, info: ValidationInfo):
+        return cls.validate_endpoint_url(v, "https://app.phoenix.arize.com")
 
 
 class LangfuseConfig(BaseTracingConfig):
@@ -29,13 +101,8 @@ class LangfuseConfig(BaseTracingConfig):
 
     @field_validator("host")
     @classmethod
-    def set_value(cls, v, info: ValidationInfo):
-        if v is None or v == "":
-            v = "https://api.langfuse.com"
-        if not v.startswith("https://") and not v.startswith("http://"):
-            raise ValueError("host must start with https:// or http://")
-
-        return v
+    def host_validator(cls, v, info: ValidationInfo):
+        return cls.validate_endpoint_url(v, "https://api.langfuse.com")
 
 
 class LangSmithConfig(BaseTracingConfig):
@@ -49,13 +116,9 @@ class LangSmithConfig(BaseTracingConfig):
 
     @field_validator("endpoint")
     @classmethod
-    def set_value(cls, v, info: ValidationInfo):
-        if v is None or v == "":
-            v = "https://api.smith.langchain.com"
-        if not v.startswith("https://"):
-            raise ValueError("endpoint must start with https://")
-
-        return v
+    def endpoint_validator(cls, v, info: ValidationInfo):
+        # LangSmith only allows HTTPS
+        return validate_url(v, "https://api.smith.langchain.com", allowed_schemes=("https",))
 
 
 class OpikConfig(BaseTracingConfig):
@@ -71,22 +134,12 @@ class OpikConfig(BaseTracingConfig):
     @field_validator("project")
     @classmethod
     def project_validator(cls, v, info: ValidationInfo):
-        if v is None or v == "":
-            v = "Default Project"
-
-        return v
+        return cls.validate_project_field(v, "Default Project")
 
     @field_validator("url")
     @classmethod
     def url_validator(cls, v, info: ValidationInfo):
-        if v is None or v == "":
-            v = "https://www.comet.com/opik/api/"
-        if not v.startswith(("https://", "http://")):
-            raise ValueError("url must start with https:// or http://")
-        if not v.endswith("/api/"):
-            raise ValueError("url should ends with /api/")
-
-        return v
+        return validate_url_with_path(v, "https://www.comet.com/opik/api/", required_suffix="/api/")
 
 
 class WeaveConfig(BaseTracingConfig):
@@ -102,22 +155,44 @@ class WeaveConfig(BaseTracingConfig):
 
     @field_validator("endpoint")
     @classmethod
-    def set_value(cls, v, info: ValidationInfo):
-        if v is None or v == "":
-            v = "https://trace.wandb.ai"
-        if not v.startswith("https://"):
-            raise ValueError("endpoint must start with https://")
-
-        return v
+    def endpoint_validator(cls, v, info: ValidationInfo):
+        # Weave only allows HTTPS for endpoint
+        return validate_url(v, "https://trace.wandb.ai", allowed_schemes=("https",))
 
     @field_validator("host")
     @classmethod
-    def validate_host(cls, v, info: ValidationInfo):
-        if v is not None and v != "":
-            if not v.startswith(("https://", "http://")):
-                raise ValueError("host must start with https:// or http://")
+    def host_validator(cls, v, info: ValidationInfo):
+        if v is not None and v.strip() != "":
+            return validate_url(v, v, allowed_schemes=("https", "http"))
         return v
 
 
+class AliyunConfig(BaseTracingConfig):
+    """
+    Model class for Aliyun tracing config.
+    """
+
+    app_name: str = "dify_app"
+    license_key: str
+    endpoint: str
+
+    @field_validator("app_name")
+    @classmethod
+    def app_name_validator(cls, v, info: ValidationInfo):
+        return cls.validate_project_field(v, "dify_app")
+
+    @field_validator("license_key")
+    @classmethod
+    def license_key_validator(cls, v, info: ValidationInfo):
+        if not v or v.strip() == "":
+            raise ValueError("License key cannot be empty")
+        return v
+
+    @field_validator("endpoint")
+    @classmethod
+    def endpoint_validator(cls, v, info: ValidationInfo):
+        return cls.validate_endpoint_url(v, "https://tracing-analysis-dc-hz.aliyuncs.com")
+
+
 OPS_FILE_PATH = "ops_trace/"
 OPS_TRACE_FAILED_KEY = "FAILED_OPS_TRACE"

api/core/ops/langfuse_trace/langfuse_trace.py

@@ -32,6 +32,7 @@ from core.repositories import SQLAlchemyWorkflowNodeExecutionRepository
 from core.workflow.nodes.enums import NodeType
 from extensions.ext_database import db
 from models import EndUser, WorkflowNodeExecutionTriggeredFrom
+from models.enums import MessageStatus
 
 logger = logging.getLogger(__name__)
 
@@ -180,12 +181,9 @@ class LangFuseDataTrace(BaseTraceInstance):
                 prompt_tokens = 0
                 completion_tokens = 0
                 try:
-                    if outputs.get("usage"):
-                        prompt_tokens = outputs.get("usage", {}).get("prompt_tokens", 0)
-                        completion_tokens = outputs.get("usage", {}).get("completion_tokens", 0)
-                    else:
-                        prompt_tokens = process_data.get("usage", {}).get("prompt_tokens", 0)
-                        completion_tokens = process_data.get("usage", {}).get("completion_tokens", 0)
+                    usage_data = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {})
+                    prompt_tokens = usage_data.get("prompt_tokens", 0)
+                    completion_tokens = usage_data.get("completion_tokens", 0)
                 except Exception:
                     logger.error("Failed to extract usage", exc_info=True)
 
@@ -293,7 +291,7 @@ class LangFuseDataTrace(BaseTraceInstance):
             input=trace_info.inputs,
             output=message_data.answer,
             metadata=metadata,
-            level=(LevelEnum.DEFAULT if message_data.status != "error" else LevelEnum.ERROR),
+            level=(LevelEnum.DEFAULT if message_data.status != MessageStatus.ERROR else LevelEnum.ERROR),
             status_message=message_data.error or "",
             usage=generation_usage,
         )
@@ -339,7 +337,7 @@ class LangFuseDataTrace(BaseTraceInstance):
             start_time=trace_info.start_time,
             end_time=trace_info.end_time,
             metadata=trace_info.metadata,
-            level=(LevelEnum.DEFAULT if message_data.status != "error" else LevelEnum.ERROR),
+            level=(LevelEnum.DEFAULT if message_data.status != MessageStatus.ERROR else LevelEnum.ERROR),
             status_message=message_data.error or "",
             usage=generation_usage,
         )

api/core/ops/langsmith_trace/langsmith_trace.py

@@ -206,12 +206,9 @@ class LangSmithDataTrace(BaseTraceInstance):
             prompt_tokens = 0
             completion_tokens = 0
             try:
-                if outputs.get("usage"):
-                    prompt_tokens = outputs.get("usage", {}).get("prompt_tokens", 0)
-                    completion_tokens = outputs.get("usage", {}).get("completion_tokens", 0)
-                else:
-                    prompt_tokens = process_data.get("usage", {}).get("prompt_tokens", 0)
-                    completion_tokens = process_data.get("usage", {}).get("completion_tokens", 0)
+                usage_data = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {})
+                prompt_tokens = usage_data.get("prompt_tokens", 0)
+                completion_tokens = usage_data.get("completion_tokens", 0)
             except Exception:
                 logger.error("Failed to extract usage", exc_info=True)
 

api/core/ops/opik_trace/opik_trace.py

@@ -222,10 +222,10 @@ class OpikDataTrace(BaseTraceInstance):
                 )
 
                 try:
-                    if outputs.get("usage"):
-                        total_tokens = outputs["usage"].get("total_tokens", 0)
-                        prompt_tokens = outputs["usage"].get("prompt_tokens", 0)
-                        completion_tokens = outputs["usage"].get("completion_tokens", 0)
+                    usage_data = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {})
+                    total_tokens = usage_data.get("total_tokens", 0)
+                    prompt_tokens = usage_data.get("prompt_tokens", 0)
+                    completion_tokens = usage_data.get("completion_tokens", 0)
                 except Exception:
                     logger.error("Failed to extract usage", exc_info=True)
 

api/core/ops/ops_trace_manager.py

@@ -84,6 +84,36 @@ class OpsTraceProviderConfigMap(dict[str, dict[str, Any]]):
                     "other_keys": ["project", "entity", "endpoint", "host"],
                     "trace_instance": WeaveDataTrace,
                 }
+            case TracingProviderEnum.ARIZE:
+                from core.ops.arize_phoenix_trace.arize_phoenix_trace import ArizePhoenixDataTrace
+                from core.ops.entities.config_entity import ArizeConfig
+
+                return {
+                    "config_class": ArizeConfig,
+                    "secret_keys": ["api_key", "space_id"],
+                    "other_keys": ["project", "endpoint"],
+                    "trace_instance": ArizePhoenixDataTrace,
+                }
+            case TracingProviderEnum.PHOENIX:
+                from core.ops.arize_phoenix_trace.arize_phoenix_trace import ArizePhoenixDataTrace
+                from core.ops.entities.config_entity import PhoenixConfig
+
+                return {
+                    "config_class": PhoenixConfig,
+                    "secret_keys": ["api_key"],
+                    "other_keys": ["project", "endpoint"],
+                    "trace_instance": ArizePhoenixDataTrace,
+                }
+            case TracingProviderEnum.ALIYUN:
+                from core.ops.aliyun_trace.aliyun_trace import AliyunDataTrace
+                from core.ops.entities.config_entity import AliyunConfig
+
+                return {
+                    "config_class": AliyunConfig,
+                    "secret_keys": ["license_key"],
+                    "other_keys": ["endpoint", "app_name"],
+                    "trace_instance": AliyunDataTrace,
+                }
 
             case _:
                 raise KeyError(f"Unsupported tracing provider: {provider}")

api/core/ops/utils.py

@@ -1,6 +1,7 @@
 from contextlib import contextmanager
 from datetime import datetime
 from typing import Optional, Union
+from urllib.parse import urlparse
 
 from extensions.ext_database import db
 from models.model import Message
@@ -60,3 +61,83 @@ def generate_dotted_order(
         return current_segment
 
     return f"{parent_dotted_order}.{current_segment}"
+
+
+def validate_url(url: str, default_url: str, allowed_schemes: tuple = ("https", "http")) -> str:
+    """
+    Validate and normalize URL with proper error handling
+
+    Args:
+        url: The URL to validate
+        default_url: Default URL to use if input is None or empty
+        allowed_schemes: Tuple of allowed URL schemes (default: https, http)
+
+    Returns:
+        Normalized URL string
+
+    Raises:
+        ValueError: If URL format is invalid or scheme not allowed
+    """
+    if not url or url.strip() == "":
+        return default_url
+
+    # Parse URL to validate format
+    parsed = urlparse(url)
+
+    # Check if scheme is allowed
+    if parsed.scheme not in allowed_schemes:
+        raise ValueError(f"URL scheme must be one of: {', '.join(allowed_schemes)}")
+
+    # Reconstruct URL with only scheme, netloc (removing path, query, fragment)
+    normalized_url = f"{parsed.scheme}://{parsed.netloc}"
+
+    return normalized_url
+
+
+def validate_url_with_path(url: str, default_url: str, required_suffix: str | None = None) -> str:
+    """
+    Validate URL that may include path components
+
+    Args:
+        url: The URL to validate
+        default_url: Default URL to use if input is None or empty
+        required_suffix: Optional suffix that URL must end with
+
+    Returns:
+        Validated URL string
+
+    Raises:
+        ValueError: If URL format is invalid or doesn't match required suffix
+    """
+    if not url or url.strip() == "":
+        return default_url
+
+    # Parse URL to validate format
+    parsed = urlparse(url)
+
+    # Check if scheme is allowed
+    if parsed.scheme not in ("https", "http"):
+        raise ValueError("URL must start with https:// or http://")
+
+    # Check required suffix if specified
+    if required_suffix and not url.endswith(required_suffix):
+        raise ValueError(f"URL should end with {required_suffix}")
+
+    return url
+
+
+def validate_project_name(project: str, default_name: str) -> str:
+    """
+    Validate and normalize project name
+
+    Args:
+        project: Project name to validate
+        default_name: Default name to use if input is None or empty
+
+    Returns:
+        Normalized project name
+    """
+    if not project or project.strip() == "":
+        return default_name
+
+    return project.strip()

api/core/plugin/entities/parameters.py

@@ -43,6 +43,19 @@ class PluginParameterType(enum.StrEnum):
     # deprecated, should not use.
     SYSTEM_FILES = CommonParameterType.SYSTEM_FILES.value
 
+    # MCP object and array type parameters
+    ARRAY = CommonParameterType.ARRAY.value
+    OBJECT = CommonParameterType.OBJECT.value
+
+
+class MCPServerParameterType(enum.StrEnum):
+    """
+    MCP server got complex parameter types
+    """
+
+    ARRAY = "array"
+    OBJECT = "object"
+
 
 class PluginParameterAutoGenerate(BaseModel):
     class Type(enum.StrEnum):
@@ -138,6 +151,34 @@ def cast_parameter_value(typ: enum.StrEnum, value: Any, /):
                 if value and not isinstance(value, list):
                     raise ValueError("The tools selector must be a list.")
                 return value
+            case PluginParameterType.ARRAY:
+                if not isinstance(value, list):
+                    # Try to parse JSON string for arrays
+                    if isinstance(value, str):
+                        try:
+                            import json
+
+                            parsed_value = json.loads(value)
+                            if isinstance(parsed_value, list):
+                                return parsed_value
+                        except (json.JSONDecodeError, ValueError):
+                            pass
+                    return [value]
+                return value
+            case PluginParameterType.OBJECT:
+                if not isinstance(value, dict):
+                    # Try to parse JSON string for objects
+                    if isinstance(value, str):
+                        try:
+                            import json
+
+                            parsed_value = json.loads(value)
+                            if isinstance(parsed_value, dict):
+                                return parsed_value
+                        except (json.JSONDecodeError, ValueError):
+                            pass
+                    return {}
+                return value
             case _:
                 return str(value)
     except ValueError:

api/core/plugin/entities/plugin.py

@@ -72,12 +72,14 @@ class PluginDeclaration(BaseModel):
 
     class Meta(BaseModel):
         minimum_dify_version: Optional[str] = Field(default=None, pattern=r"^\d{1,4}(\.\d{1,4}){1,3}(-\w{1,16})?$")
+        version: Optional[str] = Field(default=None)
 
     version: str = Field(..., pattern=r"^\d{1,4}(\.\d{1,4}){1,3}(-\w{1,16})?$")
     author: Optional[str] = Field(..., pattern=r"^[a-zA-Z0-9_-]{1,64}$")
     name: str = Field(..., pattern=r"^[a-z0-9_-]{1,128}$")
     description: I18nObject
     icon: str
+    icon_dark: Optional[str] = Field(default=None)
     label: I18nObject
     category: PluginCategory
     created_at: datetime.datetime

api/core/plugin/entities/plugin_daemon.py

@@ -53,6 +53,7 @@ class PluginAgentProviderEntity(BaseModel):
     plugin_unique_identifier: str
     plugin_id: str
     declaration: AgentProviderEntityWithPlugin
+    meta: PluginDeclaration.Meta
 
 
 class PluginBasicBooleanResponse(BaseModel):

api/core/plugin/entities/request.py

@@ -32,7 +32,7 @@ class RequestInvokeTool(BaseModel):
     Request to invoke a tool
     """
 
-    tool_type: Literal["builtin", "workflow", "api"]
+    tool_type: Literal["builtin", "workflow", "api", "mcp"]
     provider: str
     tool: str
     tool_parameters: dict

api/core/rag/datasource/vdb/qdrant/qdrant_vector.py

@@ -47,6 +47,7 @@ class QdrantConfig(BaseModel):
     grpc_port: int = 6334
     prefer_grpc: bool = False
     replication_factor: int = 1
+    write_consistency_factor: int = 1
 
     def to_qdrant_params(self):
         if self.endpoint and self.endpoint.startswith("path:"):
@@ -127,6 +128,7 @@ class QdrantVector(BaseVector):
                     hnsw_config=hnsw_config,
                     timeout=int(self._client_config.timeout),
                     replication_factor=self._client_config.replication_factor,
+                    write_consistency_factor=self._client_config.write_consistency_factor,
                 )
 
                 # create group_id payload index

api/core/rag/datasource/vdb/tencent/tencent_vector.py

@@ -122,7 +122,6 @@ class TencentVector(BaseVector):
                 metric_type,
                 params,
             )
-            index_text = vdb_index.FilterIndex(self.field_text, enum.FieldType.String, enum.IndexType.FILTER)
             index_metadate = vdb_index.FilterIndex(self.field_metadata, enum.FieldType.Json, enum.IndexType.FILTER)
             index_sparse_vector = vdb_index.SparseIndex(
                 name="sparse_vector",
@@ -130,7 +129,7 @@ class TencentVector(BaseVector):
                 index_type=enum.IndexType.SPARSE_INVERTED,
                 metric_type=enum.MetricType.IP,
             )
-            indexes = [index_id, index_vector, index_text, index_metadate]
+            indexes = [index_id, index_vector, index_metadate]
             if self._enable_hybrid_search:
                 indexes.append(index_sparse_vector)
             try:
@@ -149,7 +148,7 @@ class TencentVector(BaseVector):
                 index_metadate = vdb_index.FilterIndex(
                     self.field_metadata, enum.FieldType.String, enum.IndexType.FILTER
                 )
-                indexes = [index_id, index_vector, index_text, index_metadate]
+                indexes = [index_id, index_vector, index_metadate]
                 if self._enable_hybrid_search:
                     indexes.append(index_sparse_vector)
                 self._client.create_collection(

api/core/rag/datasource/vdb/vector_factory.py

@@ -1,3 +1,5 @@
+import logging
+import time
 from abc import ABC, abstractmethod
 from typing import Any, Optional
 
@@ -13,6 +15,8 @@ from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset, Whitelist
 
+logger = logging.getLogger(__name__)
+
 
 class AbstractVectorFactory(ABC):
     @abstractmethod
@@ -173,8 +177,20 @@ class Vector:
 
     def create(self, texts: Optional[list] = None, **kwargs):
         if texts:
-            embeddings = self._embeddings.embed_documents([document.page_content for document in texts])
-            self._vector_processor.create(texts=texts, embeddings=embeddings, **kwargs)
+            start = time.time()
+            logger.info(f"start embedding {len(texts)} texts {start}")
+            batch_size = 1000
+            total_batches = len(texts) + batch_size - 1
+            for i in range(0, len(texts), batch_size):
+                batch = texts[i : i + batch_size]
+                batch_start = time.time()
+                logger.info(f"Processing batch {i // batch_size + 1}/{total_batches} ({len(batch)} texts)")
+                batch_embeddings = self._embeddings.embed_documents([document.page_content for document in batch])
+                logger.info(
+                    f"Embedding batch {i // batch_size + 1}/{total_batches} took {time.time() - batch_start:.3f}s"
+                )
+                self._vector_processor.create(texts=batch, embeddings=batch_embeddings, **kwargs)
+            logger.info(f"Embedding {len(texts)} texts took {time.time() - start:.3f}s")
 
     def add_texts(self, documents: list[Document], **kwargs):
         if kwargs.get("duplicate_check", False):

api/core/rag/extractor/helpers.py

@@ -1,7 +1,6 @@
 """Document loader helpers."""
 
 import concurrent.futures
-from pathlib import Path
 from typing import NamedTuple, Optional, cast
 
 
@@ -16,7 +15,7 @@ class FileEncoding(NamedTuple):
     """The language of the file."""
 
 
-def detect_file_encodings(file_path: str, timeout: int = 5) -> list[FileEncoding]:
+def detect_file_encodings(file_path: str, timeout: int = 5, sample_size: int = 1024 * 1024) -> list[FileEncoding]:
     """Try to detect the file encoding.
 
     Returns a list of `FileEncoding` tuples with the detected encodings ordered
@@ -25,11 +24,16 @@ def detect_file_encodings(file_path: str, timeout: int = 5) -> list[FileEncoding
     Args:
         file_path: The path to the file to detect the encoding for.
         timeout: The timeout in seconds for the encoding detection.
+        sample_size: The number of bytes to read for encoding detection. Default is 1MB.
+                    For large files, reading only a sample is sufficient and prevents timeout.
     """
     import chardet
 
     def read_and_detect(file_path: str) -> list[dict]:
-        rawdata = Path(file_path).read_bytes()
+        with open(file_path, "rb") as f:
+            # Read only a sample of the file for encoding detection
+            # This prevents timeout on large files while still providing accurate encoding detection
+            rawdata = f.read(sample_size)
         return cast(list[dict], chardet.detect_all(rawdata))
 
     with concurrent.futures.ThreadPoolExecutor() as executor:

api/core/rag/extractor/text_extractor.py

@@ -36,8 +36,12 @@ class TextExtractor(BaseExtractor):
                         break
                     except UnicodeDecodeError:
                         continue
+                else:
+                    raise RuntimeError(
+                        f"Decode failed: {self._file_path}, all detected encodings failed. Original error: {e}"
+                    )
             else:
-                raise RuntimeError(f"Error loading {self._file_path}") from e
+                raise RuntimeError(f"Decode failed: {self._file_path}, specified encoding failed. Original error: {e}")
         except Exception as e:
             raise RuntimeError(f"Error loading {self._file_path}") from e
 

api/core/repositories/sqlalchemy_workflow_execution_repository.py

@@ -17,6 +17,7 @@ from core.workflow.entities.workflow_execution import (
 )
 from core.workflow.repositories.workflow_execution_repository import WorkflowExecutionRepository
 from core.workflow.workflow_type_encoder import WorkflowRuntimeTypeConverter
+from libs.helper import extract_tenant_id
 from models import (
     Account,
     CreatorUserRole,
@@ -67,7 +68,7 @@ class SQLAlchemyWorkflowExecutionRepository(WorkflowExecutionRepository):
             )
 
         # Extract tenant_id from user
-        tenant_id: str | None = user.tenant_id if isinstance(user, EndUser) else user.current_tenant_id
+        tenant_id = extract_tenant_id(user)
         if not tenant_id:
             raise ValueError("User must have a tenant_id or current_tenant_id")
         self._tenant_id = tenant_id

api/core/repositories/sqlalchemy_workflow_node_execution_repository.py

@@ -20,6 +20,7 @@ from core.workflow.entities.workflow_node_execution import (
 from core.workflow.nodes.enums import NodeType
 from core.workflow.repositories.workflow_node_execution_repository import OrderConfig, WorkflowNodeExecutionRepository
 from core.workflow.workflow_type_encoder import WorkflowRuntimeTypeConverter
+from libs.helper import extract_tenant_id
 from models import (
     Account,
     CreatorUserRole,
@@ -70,7 +71,7 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository)
             )
 
         # Extract tenant_id from user
-        tenant_id: str | None = user.tenant_id if isinstance(user, EndUser) else user.current_tenant_id
+        tenant_id = extract_tenant_id(user)
         if not tenant_id:
             raise ValueError("User must have a tenant_id or current_tenant_id")
         self._tenant_id = tenant_id

api/core/tools/builtin_tool/providers/audio/tools/tts.py

@@ -31,6 +31,14 @@ class TTSTool(BuiltinTool):
             model_type=ModelType.TTS,
             model=model,
         )
+        if not voice:
+            voices = model_instance.get_tts_voices()
+            if voices:
+                voice = voices[0].get("value")
+                if not voice:
+                    raise ValueError("Sorry, no voice available.")
+            else:
+                raise ValueError("Sorry, no voice available.")
         tts = model_instance.invoke_tts(
             content_text=tool_parameters.get("text"),  # type: ignore
             user=user_id,

api/core/tools/custom_tool/provider.py

@@ -39,19 +39,22 @@ class ApiToolProviderController(ToolProviderController):
                 type=ProviderConfig.Type.SELECT,
                 options=[
                     ProviderConfig.Option(value="none", label=I18nObject(en_US="None", zh_Hans="无")),
-                    ProviderConfig.Option(value="api_key", label=I18nObject(en_US="api_key", zh_Hans="api_key")),
+                    ProviderConfig.Option(value="api_key_header", label=I18nObject(en_US="Header", zh_Hans="请求头")),
+                    ProviderConfig.Option(
+                        value="api_key_query", label=I18nObject(en_US="Query Param", zh_Hans="查询参数")
+                    ),
                 ],
                 default="none",
                 help=I18nObject(en_US="The auth type of the api provider", zh_Hans="api provider 的认证类型"),
             )
         ]
-        if auth_type == ApiProviderAuthType.API_KEY:
+        if auth_type == ApiProviderAuthType.API_KEY_HEADER:
             credentials_schema = [
                 *credentials_schema,
                 ProviderConfig(
                     name="api_key_header",
                     required=False,
-                    default="api_key",
+                    default="Authorization",
                     type=ProviderConfig.Type.TEXT_INPUT,
                     help=I18nObject(en_US="The header name of the api key", zh_Hans="携带 api key 的 header 名称"),
                 ),
@@ -74,6 +77,25 @@ class ApiToolProviderController(ToolProviderController):
                     ],
                 ),
             ]
+        elif auth_type == ApiProviderAuthType.API_KEY_QUERY:
+            credentials_schema = [
+                *credentials_schema,
+                ProviderConfig(
+                    name="api_key_query_param",
+                    required=False,
+                    default="key",
+                    type=ProviderConfig.Type.TEXT_INPUT,
+                    help=I18nObject(
+                        en_US="The query parameter name of the api key", zh_Hans="携带 api key 的查询参数名称"
+                    ),
+                ),
+                ProviderConfig(
+                    name="api_key_value",
+                    required=True,
+                    type=ProviderConfig.Type.SECRET_INPUT,
+                    help=I18nObject(en_US="The api key", zh_Hans="api key 的值"),
+                ),
+            ]
         elif auth_type == ApiProviderAuthType.NONE:
             pass
 

api/core/tools/custom_tool/tool.py

@@ -78,8 +78,8 @@ class ApiTool(Tool):
         if "auth_type" not in credentials:
             raise ToolProviderCredentialValidationError("Missing auth_type")
 
-        if credentials["auth_type"] == "api_key":
-            api_key_header = "api_key"
+        if credentials["auth_type"] in ("api_key_header", "api_key"):  # backward compatibility:
+            api_key_header = "Authorization"
 
             if "api_key_header" in credentials:
                 api_key_header = credentials["api_key_header"]
@@ -100,6 +100,11 @@ class ApiTool(Tool):
 
             headers[api_key_header] = credentials["api_key_value"]
 
+        elif credentials["auth_type"] == "api_key_query":
+            # For query parameter authentication, we don't add anything to headers
+            # The query parameter will be added in do_http_request method
+            pass
+
         needed_parameters = [parameter for parameter in (self.api_bundle.parameters or []) if parameter.required]
         for parameter in needed_parameters:
             if parameter.required and parameter.name not in parameters:
@@ -154,6 +159,15 @@ class ApiTool(Tool):
         cookies = {}
         files = []
 
+        # Add API key to query parameters if auth_type is api_key_query
+        if self.runtime and self.runtime.credentials:
+            credentials = self.runtime.credentials
+            if credentials.get("auth_type") == "api_key_query":
+                api_key_query_param = credentials.get("api_key_query_param", "key")
+                api_key_value = credentials.get("api_key_value")
+                if api_key_value:
+                    params[api_key_query_param] = api_key_value
+
         # check parameters
         for parameter in self.api_bundle.openapi.get("parameters", []):
             value = self.get_parameter_value(parameter, parameters)
@@ -213,7 +227,8 @@ class ApiTool(Tool):
                         elif "default" in property:
                             body[name] = property["default"]
                         else:
-                            body[name] = None
+                            # omit optional parameters that weren't provided, instead of setting them to None
+                            pass
                     break
 
         # replace path parameters

api/core/tools/entities/api_entities.py

@@ -1,4 +1,5 @@
-from typing import Literal, Optional
+from datetime import datetime
+from typing import Any, Literal, Optional
 
 from pydantic import BaseModel, Field, field_validator
 
@@ -18,7 +19,7 @@ class ToolApiEntity(BaseModel):
     output_schema: Optional[dict] = None
 
 
-ToolProviderTypeApiLiteral = Optional[Literal["builtin", "api", "workflow"]]
+ToolProviderTypeApiLiteral = Optional[Literal["builtin", "api", "workflow", "mcp"]]
 
 
 class ToolProviderApiEntity(BaseModel):
@@ -27,6 +28,7 @@ class ToolProviderApiEntity(BaseModel):
     name: str  # identifier
     description: I18nObject
     icon: str | dict
+    icon_dark: Optional[str | dict] = Field(default=None, description="The dark icon of the tool")
     label: I18nObject  # label
     type: ToolProviderType
     masked_credentials: Optional[dict] = None
@@ -37,6 +39,10 @@ class ToolProviderApiEntity(BaseModel):
     plugin_unique_identifier: Optional[str] = Field(default="", description="The unique identifier of the tool")
     tools: list[ToolApiEntity] = Field(default_factory=list)
     labels: list[str] = Field(default_factory=list)
+    # MCP
+    server_url: Optional[str] = Field(default="", description="The server url of the tool")
+    updated_at: int = Field(default_factory=lambda: int(datetime.now().timestamp()))
+    server_identifier: Optional[str] = Field(default="", description="The server identifier of the MCP tool")
 
     @field_validator("tools", mode="before")
     @classmethod
@@ -52,8 +58,13 @@ class ToolProviderApiEntity(BaseModel):
                 for parameter in tool.get("parameters"):
                     if parameter.get("type") == ToolParameter.ToolParameterType.SYSTEM_FILES.value:
                         parameter["type"] = "files"
+                    if parameter.get("input_schema") is None:
+                        parameter.pop("input_schema", None)
         # -------------
-
+        optional_fields = self.optional_field("server_url", self.server_url)
+        if self.type == ToolProviderType.MCP.value:
+            optional_fields.update(self.optional_field("updated_at", self.updated_at))
+            optional_fields.update(self.optional_field("server_identifier", self.server_identifier))
         return {
             "id": self.id,
             "author": self.author,
@@ -62,6 +73,7 @@ class ToolProviderApiEntity(BaseModel):
             "plugin_unique_identifier": self.plugin_unique_identifier,
             "description": self.description.to_dict(),
             "icon": self.icon,
+            "icon_dark": self.icon_dark,
             "label": self.label.to_dict(),
             "type": self.type.value,
             "team_credentials": self.masked_credentials,
@@ -69,8 +81,13 @@ class ToolProviderApiEntity(BaseModel):
             "allow_delete": self.allow_delete,
             "tools": tools,
             "labels": self.labels,
+            **optional_fields,
         }
 
+    def optional_field(self, key: str, value: Any) -> dict:
+        """Return dict with key-value if value is truthy, empty dict otherwise."""
+        return {key: value} if value else {}
+
 
 class ToolProviderCredentialApiEntity(BaseModel):
     id: str = Field(description="The unique id of the credential")

api/core/tools/entities/tool_entities.py

@@ -8,6 +8,7 @@ from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_seriali
 
 from core.entities.provider_entities import ProviderConfig
 from core.plugin.entities.parameters import (
+    MCPServerParameterType,
     PluginParameter,
     PluginParameterOption,
     PluginParameterType,
@@ -49,6 +50,7 @@ class ToolProviderType(enum.StrEnum):
     API = "api"
     APP = "app"
     DATASET_RETRIEVAL = "dataset-retrieval"
+    MCP = "mcp"
 
     @classmethod
     def value_of(cls, value: str) -> "ToolProviderType":
@@ -94,7 +96,8 @@ class ApiProviderAuthType(Enum):
     """
 
     NONE = "none"
-    API_KEY = "api_key"
+    API_KEY_HEADER = "api_key_header"
+    API_KEY_QUERY = "api_key_query"
 
     @classmethod
     def value_of(cls, value: str) -> "ApiProviderAuthType":
@@ -242,6 +245,10 @@ class ToolParameter(PluginParameter):
         MODEL_SELECTOR = PluginParameterType.MODEL_SELECTOR.value
         DYNAMIC_SELECT = PluginParameterType.DYNAMIC_SELECT.value
 
+        # MCP object and array type parameters
+        ARRAY = MCPServerParameterType.ARRAY.value
+        OBJECT = MCPServerParameterType.OBJECT.value
+
         # deprecated, should not use.
         SYSTEM_FILES = PluginParameterType.SYSTEM_FILES.value
 
@@ -260,6 +267,8 @@ class ToolParameter(PluginParameter):
     human_description: Optional[I18nObject] = Field(default=None, description="The description presented to the user")
     form: ToolParameterForm = Field(..., description="The form of the parameter, schema/form/llm")
     llm_description: Optional[str] = None
+    # MCP object and array type parameters use this field to store the schema
+    input_schema: Optional[dict] = None
 
     @classmethod
     def get_simple_instance(
@@ -309,6 +318,7 @@ class ToolProviderIdentity(BaseModel):
     name: str = Field(..., description="The name of the tool")
     description: I18nObject = Field(..., description="The description of the tool")
     icon: str = Field(..., description="The icon of the tool")
+    icon_dark: Optional[str] = Field(default=None, description="The dark icon of the tool")
     label: I18nObject = Field(..., description="The label of the tool")
     tags: Optional[list[ToolLabelEnum]] = Field(
         default=[],

api/core/tools/mcp_tool/provider.py

@@ -0,0 +1,130 @@
+import json
+from typing import Any
+
+from core.mcp.types import Tool as RemoteMCPTool
+from core.tools.__base.tool_provider import ToolProviderController
+from core.tools.__base.tool_runtime import ToolRuntime
+from core.tools.entities.common_entities import I18nObject
+from core.tools.entities.tool_entities import (
+    ToolDescription,
+    ToolEntity,
+    ToolIdentity,
+    ToolProviderEntityWithPlugin,
+    ToolProviderIdentity,
+    ToolProviderType,
+)
+from core.tools.mcp_tool.tool import MCPTool
+from models.tools import MCPToolProvider
+from services.tools.tools_transform_service import ToolTransformService
+
+
+class MCPToolProviderController(ToolProviderController):
+    provider_id: str
+    entity: ToolProviderEntityWithPlugin
+
+    def __init__(self, entity: ToolProviderEntityWithPlugin, provider_id: str, tenant_id: str, server_url: str) -> None:
+        super().__init__(entity)
+        self.entity = entity
+        self.tenant_id = tenant_id
+        self.provider_id = provider_id
+        self.server_url = server_url
+
+    @property
+    def provider_type(self) -> ToolProviderType:
+        """
+        returns the type of the provider
+
+        :return: type of the provider
+        """
+        return ToolProviderType.MCP
+
+    @classmethod
+    def _from_db(cls, db_provider: MCPToolProvider) -> "MCPToolProviderController":
+        """
+        from db provider
+        """
+        tools = []
+        tools_data = json.loads(db_provider.tools)
+        remote_mcp_tools = [RemoteMCPTool(**tool) for tool in tools_data]
+        user = db_provider.load_user()
+        tools = [
+            ToolEntity(
+                identity=ToolIdentity(
+                    author=user.name if user else "Anonymous",
+                    name=remote_mcp_tool.name,
+                    label=I18nObject(en_US=remote_mcp_tool.name, zh_Hans=remote_mcp_tool.name),
+                    provider=db_provider.server_identifier,
+                    icon=db_provider.icon,
+                ),
+                parameters=ToolTransformService.convert_mcp_schema_to_parameter(remote_mcp_tool.inputSchema),
+                description=ToolDescription(
+                    human=I18nObject(
+                        en_US=remote_mcp_tool.description or "", zh_Hans=remote_mcp_tool.description or ""
+                    ),
+                    llm=remote_mcp_tool.description or "",
+                ),
+                output_schema=None,
+                has_runtime_parameters=len(remote_mcp_tool.inputSchema) > 0,
+            )
+            for remote_mcp_tool in remote_mcp_tools
+        ]
+
+        return cls(
+            entity=ToolProviderEntityWithPlugin(
+                identity=ToolProviderIdentity(
+                    author=user.name if user else "Anonymous",
+                    name=db_provider.name,
+                    label=I18nObject(en_US=db_provider.name, zh_Hans=db_provider.name),
+                    description=I18nObject(en_US="", zh_Hans=""),
+                    icon=db_provider.icon,
+                ),
+                plugin_id=None,
+                credentials_schema=[],
+                tools=tools,
+            ),
+            provider_id=db_provider.server_identifier or "",
+            tenant_id=db_provider.tenant_id or "",
+            server_url=db_provider.decrypted_server_url,
+        )
+
+    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]) -> None:
+        """
+        validate the credentials of the provider
+        """
+        pass
+
+    def get_tool(self, tool_name: str) -> MCPTool:  # type: ignore
+        """
+        return tool with given name
+        """
+        tool_entity = next(
+            (tool_entity for tool_entity in self.entity.tools if tool_entity.identity.name == tool_name), None
+        )
+
+        if not tool_entity:
+            raise ValueError(f"Tool with name {tool_name} not found")
+
+        return MCPTool(
+            entity=tool_entity,
+            runtime=ToolRuntime(tenant_id=self.tenant_id),
+            tenant_id=self.tenant_id,
+            icon=self.entity.identity.icon,
+            server_url=self.server_url,
+            provider_id=self.provider_id,
+        )
+
+    def get_tools(self) -> list[MCPTool]:  # type: ignore
+        """
+        get all tools
+        """
+        return [
+            MCPTool(
+                entity=tool_entity,
+                runtime=ToolRuntime(tenant_id=self.tenant_id),
+                tenant_id=self.tenant_id,
+                icon=self.entity.identity.icon,
+                server_url=self.server_url,
+                provider_id=self.provider_id,
+            )
+            for tool_entity in self.entity.tools
+        ]

api/core/tools/mcp_tool/tool.py

@@ -0,0 +1,92 @@
+import base64
+import json
+from collections.abc import Generator
+from typing import Any, Optional
+
+from core.mcp.error import MCPAuthError, MCPConnectionError
+from core.mcp.mcp_client import MCPClient
+from core.mcp.types import 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, ToolParameter, ToolProviderType
+
+
+class MCPTool(Tool):
+    tenant_id: str
+    icon: str
+    runtime_parameters: Optional[list[ToolParameter]]
+    server_url: str
+    provider_id: str
+
+    def __init__(
+        self, entity: ToolEntity, runtime: ToolRuntime, tenant_id: str, icon: str, server_url: str, provider_id: str
+    ) -> None:
+        super().__init__(entity, runtime)
+        self.tenant_id = tenant_id
+        self.icon = icon
+        self.runtime_parameters = None
+        self.server_url = server_url
+        self.provider_id = provider_id
+
+    def tool_provider_type(self) -> ToolProviderType:
+        return ToolProviderType.MCP
+
+    def _invoke(
+        self,
+        user_id: str,
+        tool_parameters: dict[str, Any],
+        conversation_id: Optional[str] = None,
+        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 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
+
+        for content in result.content:
+            if isinstance(content, TextContent):
+                try:
+                    content_json = json.loads(content.text)
+                    if isinstance(content_json, dict):
+                        yield self.create_json_message(content_json)
+                    elif isinstance(content_json, list):
+                        for item in content_json:
+                            yield self.create_json_message(item)
+                    else:
+                        yield self.create_text_message(content.text)
+                except json.JSONDecodeError:
+                    yield self.create_text_message(content.text)
+
+            elif isinstance(content, ImageContent):
+                yield self.create_blob_message(
+                    blob=base64.b64decode(content.data), meta={"mime_type": content.mimeType}
+                )
+
+    def fork_tool_runtime(self, runtime: ToolRuntime) -> "MCPTool":
+        return MCPTool(
+            entity=self.entity,
+            runtime=runtime,
+            tenant_id=self.tenant_id,
+            icon=self.icon,
+            server_url=self.server_url,
+            provider_id=self.provider_id,
+        )
+
+    def _handle_none_parameter(self, parameter: dict[str, Any]) -> dict[str, Any]:
+        """
+        in mcp tool invoke, if the parameter is empty, it will be set to None
+        """
+        return {
+            key: value
+            for key, value in parameter.items()
+            if value is not None and not (isinstance(value, str) and value.strip() == "")
+        }

api/core/tools/signature.py

@@ -9,9 +9,10 @@ from configs import dify_config
 
 def sign_tool_file(tool_file_id: str, extension: str) -> str:
     """
-    sign file to get a temporary url
+    sign file to get a temporary url for plugin access
     """
-    base_url = dify_config.FILES_URL
+    # Use internal URL for plugin/tool file access in Docker environments
+    base_url = dify_config.INTERNAL_FILES_URL or dify_config.FILES_URL
     file_preview_url = f"{base_url}/files/tools/{tool_file_id}{extension}"
 
     timestamp = str(int(time.time()))

api/core/tools/tool_file_manager.py

@@ -35,9 +35,10 @@ class ToolFileManager:
     @staticmethod
     def sign_file(tool_file_id: str, extension: str) -> str:
         """
-        sign file to get a temporary url
+        sign file to get a temporary url for plugin access
         """
-        base_url = dify_config.FILES_URL
+        # Use internal URL for plugin/tool file access in Docker environments
+        base_url = dify_config.INTERNAL_FILES_URL or dify_config.FILES_URL
         file_preview_url = f"{base_url}/files/tools/{tool_file_id}{extension}"
 
         timestamp = str(int(time.time()))

api/core/tools/tool_manager.py

@@ -4,7 +4,7 @@ import mimetypes
 from collections.abc import Generator
 from os import listdir, path
 from threading import Lock
-from typing import TYPE_CHECKING, Any, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union, cast
 
 from yarl import URL
 
@@ -14,9 +14,13 @@ from core.plugin.entities.plugin import ToolProviderID
 from core.plugin.impl.tool import PluginToolManager
 from core.tools.__base.tool_provider import ToolProviderController
 from core.tools.__base.tool_runtime import ToolRuntime
+from core.tools.mcp_tool.provider import MCPToolProviderController
+from core.tools.mcp_tool.tool import MCPTool
 from core.tools.plugin_tool.provider import PluginToolProviderController
 from core.tools.plugin_tool.tool import PluginTool
 from core.tools.workflow_as_tool.provider import WorkflowToolProviderController
+from core.workflow.entities.variable_pool import VariablePool
+from services.tools.mcp_tools_mange_service import MCPToolManageService
 
 if TYPE_CHECKING:
     from core.workflow.nodes.tool.entities import ToolEntity
@@ -42,7 +46,7 @@ from core.tools.entities.tool_entities import (
     ToolParameter,
     ToolProviderType,
 )
-from core.tools.errors import ToolProviderNotFoundError
+from core.tools.errors import ToolNotFoundError, ToolProviderNotFoundError
 from core.tools.tool_label_manager import ToolLabelManager
 from core.tools.utils.configuration import (
     ToolParameterConfigurationManager,
@@ -50,7 +54,7 @@ from core.tools.utils.configuration import (
 from core.tools.utils.encryption import create_provider_encrypter, create_tool_provider_encrypter
 from core.tools.workflow_as_tool.tool import WorkflowTool
 from extensions.ext_database import db
-from models.tools import ApiToolProvider, BuiltinToolProvider, WorkflowToolProvider
+from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider, WorkflowToolProvider
 from services.tools.tools_transform_service import ToolTransformService
 
 logger = logging.getLogger(__name__)
@@ -148,7 +152,7 @@ class ToolManager:
         invoke_from: InvokeFrom = InvokeFrom.DEBUGGER,
         tool_invoke_from: ToolInvokeFrom = ToolInvokeFrom.AGENT,
         credential_id: Optional[str] = None,
-    ) -> Union[BuiltinTool, PluginTool, ApiTool, WorkflowTool]:
+    ) -> Union[BuiltinTool, PluginTool, ApiTool, WorkflowTool, MCPTool]:
         """
         get the tool runtime
 
@@ -293,6 +297,8 @@ class ToolManager:
             raise NotImplementedError("app provider not implemented")
         elif provider_type == ToolProviderType.PLUGIN:
             return cls.get_plugin_provider(provider_id, tenant_id).get_tool(tool_name)
+        elif provider_type == ToolProviderType.MCP:
+            return cls.get_mcp_provider_controller(tenant_id, provider_id).get_tool(tool_name)
         else:
             raise ToolProviderNotFoundError(f"provider type {provider_type.value} not found")
 
@@ -303,6 +309,7 @@ class ToolManager:
         app_id: str,
         agent_tool: AgentToolEntity,
         invoke_from: InvokeFrom = InvokeFrom.DEBUGGER,
+        variable_pool: Optional[VariablePool] = None,
     ) -> Tool:
         """
         get the agent tool runtime
@@ -317,24 +324,9 @@ class ToolManager:
         )
         runtime_parameters = {}
         parameters = tool_entity.get_merged_runtime_parameters()
-        for parameter in parameters:
-            # check file types
-            if (
-                parameter.type
-                in {
-                    ToolParameter.ToolParameterType.SYSTEM_FILES,
-                    ToolParameter.ToolParameterType.FILE,
-                    ToolParameter.ToolParameterType.FILES,
-                }
-                and parameter.required
-            ):
-                raise ValueError(f"file type parameter {parameter.name} not supported in agent")
-
-            if parameter.form == ToolParameter.ToolParameterForm.FORM:
-                # save tool parameter to tool entity memory
-                value = parameter.init_frontend_parameter(agent_tool.tool_parameters.get(parameter.name))
-                runtime_parameters[parameter.name] = value
-
+        runtime_parameters = cls._convert_tool_parameters_type(
+            parameters, variable_pool, agent_tool.tool_parameters, typ="agent"
+        )
         # decrypt runtime parameters
         encryption_manager = ToolParameterConfigurationManager(
             tenant_id=tenant_id,
@@ -358,10 +350,12 @@ class ToolManager:
         node_id: str,
         workflow_tool: "ToolEntity",
         invoke_from: InvokeFrom = InvokeFrom.DEBUGGER,
+        variable_pool: Optional[VariablePool] = None,
     ) -> Tool:
         """
         get the workflow tool runtime
         """
+
         tool_runtime = cls.get_tool_runtime(
             provider_type=workflow_tool.provider_type,
             provider_id=workflow_tool.provider_id,
@@ -371,15 +365,11 @@ class ToolManager:
             tool_invoke_from=ToolInvokeFrom.WORKFLOW,
             credential_id=workflow_tool.credential_id,
         )
-        runtime_parameters = {}
+
         parameters = tool_runtime.get_merged_runtime_parameters()
-
-        for parameter in parameters:
-            # save tool parameter to tool entity memory
-            if parameter.form == ToolParameter.ToolParameterForm.FORM:
-                value = parameter.init_frontend_parameter(workflow_tool.tool_configurations.get(parameter.name))
-                runtime_parameters[parameter.name] = value
-
+        runtime_parameters = cls._convert_tool_parameters_type(
+            parameters, variable_pool, workflow_tool.tool_configurations, typ="workflow"
+        )
         # decrypt runtime parameters
         encryption_manager = ToolParameterConfigurationManager(
             tenant_id=tenant_id,
@@ -587,7 +577,7 @@ class ToolManager:
 
         filters = []
         if not typ:
-            filters.extend(["builtin", "api", "workflow"])
+            filters.extend(["builtin", "api", "workflow", "mcp"])
         else:
             filters.append(typ)
 
@@ -671,6 +661,10 @@ class ToolManager:
                         labels=labels.get(provider_controller.provider_id, []),
                     )
                     result_providers[f"workflow_provider.{user_provider.name}"] = user_provider
+            if "mcp" in filters:
+                mcp_providers = MCPToolManageService.retrieve_mcp_tools(tenant_id, for_list=True)
+                for mcp_provider in mcp_providers:
+                    result_providers[f"mcp_provider.{mcp_provider.name}"] = mcp_provider
 
         return BuiltinToolProviderSort.sort(list(result_providers.values()))
 
@@ -698,14 +692,47 @@ class ToolManager:
         if provider is None:
             raise ToolProviderNotFoundError(f"api provider {provider_id} not found")
 
+        auth_type = ApiProviderAuthType.NONE
+        provider_auth_type = provider.credentials.get("auth_type")
+        if provider_auth_type in ("api_key_header", "api_key"):  # backward compatibility
+            auth_type = ApiProviderAuthType.API_KEY_HEADER
+        elif provider_auth_type == "api_key_query":
+            auth_type = ApiProviderAuthType.API_KEY_QUERY
+
         controller = ApiToolProviderController.from_db(
             provider,
-            ApiProviderAuthType.API_KEY if provider.credentials["auth_type"] == "api_key" else ApiProviderAuthType.NONE,
+            auth_type,
         )
         controller.load_bundled_tools(provider.tools)
 
         return controller, provider.credentials
 
+    @classmethod
+    def get_mcp_provider_controller(cls, tenant_id: str, provider_id: str) -> MCPToolProviderController:
+        """
+        get the api provider
+
+        :param tenant_id: the id of the tenant
+        :param provider_id: the id of the provider
+
+        :return: the provider controller, the credentials
+        """
+        provider: MCPToolProvider | None = (
+            db.session.query(MCPToolProvider)
+            .filter(
+                MCPToolProvider.server_identifier == provider_id,
+                MCPToolProvider.tenant_id == tenant_id,
+            )
+            .first()
+        )
+
+        if provider is None:
+            raise ToolProviderNotFoundError(f"mcp provider {provider_id} not found")
+
+        controller = MCPToolProviderController._from_db(provider)
+
+        return controller
+
     @classmethod
     def user_get_api_provider(cls, provider: str, tenant_id: str) -> dict:
         """
@@ -733,9 +760,16 @@ class ToolManager:
             credentials = {}
 
         # package tool provider controller
+        auth_type = ApiProviderAuthType.NONE
+        credentials_auth_type = credentials.get("auth_type")
+        if credentials_auth_type in ("api_key_header", "api_key"):  # backward compatibility
+            auth_type = ApiProviderAuthType.API_KEY_HEADER
+        elif credentials_auth_type == "api_key_query":
+            auth_type = ApiProviderAuthType.API_KEY_QUERY
+
         controller = ApiToolProviderController.from_db(
             provider_obj,
-            ApiProviderAuthType.API_KEY if credentials["auth_type"] == "api_key" else ApiProviderAuthType.NONE,
+            auth_type,
         )
         # init tool configuration
         encrypter, _ = create_tool_provider_encrypter(
@@ -831,6 +865,22 @@ class ToolManager:
         except Exception:
             return {"background": "#252525", "content": "\ud83d\ude01"}
 
+    @classmethod
+    def generate_mcp_tool_icon_url(cls, tenant_id: str, provider_id: str) -> dict[str, str] | str:
+        try:
+            mcp_provider: MCPToolProvider | None = (
+                db.session.query(MCPToolProvider)
+                .filter(MCPToolProvider.tenant_id == tenant_id, MCPToolProvider.server_identifier == provider_id)
+                .first()
+            )
+
+            if mcp_provider is None:
+                raise ToolProviderNotFoundError(f"mcp provider {provider_id} not found")
+
+            return mcp_provider.provider_icon
+        except Exception:
+            return {"background": "#252525", "content": "\ud83d\ude01"}
+
     @classmethod
     def get_tool_icon(
         cls,
@@ -868,8 +918,61 @@ class ToolManager:
                 except Exception:
                     return {"background": "#252525", "content": "\ud83d\ude01"}
             raise ValueError(f"plugin provider {provider_id} not found")
+        elif provider_type == ToolProviderType.MCP:
+            return cls.generate_mcp_tool_icon_url(tenant_id, provider_id)
         else:
             raise ValueError(f"provider type {provider_type} not found")
 
+    @classmethod
+    def _convert_tool_parameters_type(
+        cls,
+        parameters: list[ToolParameter],
+        variable_pool: Optional[VariablePool],
+        tool_configurations: dict[str, Any],
+        typ: Literal["agent", "workflow", "tool"] = "workflow",
+    ) -> dict[str, Any]:
+        """
+        Convert tool parameters type
+        """
+        from core.workflow.nodes.tool.entities import ToolNodeData
+        from core.workflow.nodes.tool.exc import ToolParameterError
+
+        runtime_parameters = {}
+        for parameter in parameters:
+            if (
+                parameter.type
+                in {
+                    ToolParameter.ToolParameterType.SYSTEM_FILES,
+                    ToolParameter.ToolParameterType.FILE,
+                    ToolParameter.ToolParameterType.FILES,
+                }
+                and parameter.required
+                and typ == "agent"
+            ):
+                raise ValueError(f"file type parameter {parameter.name} not supported in agent")
+            # save tool parameter to tool entity memory
+            if parameter.form == ToolParameter.ToolParameterForm.FORM:
+                if variable_pool:
+                    config = tool_configurations.get(parameter.name, {})
+                    if not (config and isinstance(config, dict) and config.get("value") is not None):
+                        continue
+                    tool_input = ToolNodeData.ToolInput(**tool_configurations.get(parameter.name, {}))
+                    if tool_input.type == "variable":
+                        variable = variable_pool.get(tool_input.value)
+                        if variable is None:
+                            raise ToolParameterError(f"Variable {tool_input.value} does not exist")
+                        parameter_value = variable.value
+                    elif tool_input.type in {"mixed", "constant"}:
+                        segment_group = variable_pool.convert_template(str(tool_input.value))
+                        parameter_value = segment_group.text
+                    else:
+                        raise ToolParameterError(f"Unknown tool input type '{tool_input.type}'")
+                    runtime_parameters[parameter.name] = parameter_value
+
+                else:
+                    value = parameter.init_frontend_parameter(tool_configurations.get(parameter.name))
+                    runtime_parameters[parameter.name] = value
+        return runtime_parameters
+
 
 ToolManager.load_hardcoded_providers_cache()

api/core/tools/workflow_as_tool/tool.py

@@ -8,7 +8,12 @@ from flask_login import current_user
 from core.file import FILE_MODEL_IDENTITY, File, FileTransferMethod
 from core.tools.__base.tool import Tool
 from core.tools.__base.tool_runtime import ToolRuntime
-from core.tools.entities.tool_entities import ToolEntity, ToolInvokeMessage, ToolParameter, ToolProviderType
+from core.tools.entities.tool_entities import (
+    ToolEntity,
+    ToolInvokeMessage,
+    ToolParameter,
+    ToolProviderType,
+)
 from core.tools.errors import ToolInvokeError
 from extensions.ext_database import db
 from factories.file_factory import build_from_mapping

api/core/workflow/callbacks/workflow_logging_callback.py

@@ -232,14 +232,14 @@ class WorkflowLoggingCallback(WorkflowCallback):
         Publish loop started
         """
         self.print_text("\n[LoopRunStartedEvent]", color="blue")
-        self.print_text(f"Loop Node ID: {event.loop_id}", color="blue")
+        self.print_text(f"Loop Node ID: {event.loop_node_id}", color="blue")
 
     def on_workflow_loop_next(self, event: LoopRunNextEvent) -> None:
         """
         Publish loop next
         """
         self.print_text("\n[LoopRunNextEvent]", color="blue")
-        self.print_text(f"Loop Node ID: {event.loop_id}", color="blue")
+        self.print_text(f"Loop Node ID: {event.loop_node_id}", color="blue")
         self.print_text(f"Loop Index: {event.index}", color="blue")
 
     def on_workflow_loop_completed(self, event: LoopRunSucceededEvent | LoopRunFailedEvent) -> None:
@@ -250,7 +250,7 @@ class WorkflowLoggingCallback(WorkflowCallback):
             "\n[LoopRunSucceededEvent]" if isinstance(event, LoopRunSucceededEvent) else "\n[LoopRunFailedEvent]",
             color="blue",
         )
-        self.print_text(f"Node ID: {event.loop_id}", color="blue")
+        self.print_text(f"Loop Node ID: {event.loop_node_id}", color="blue")
 
     def print_text(self, text: str, color: Optional[str] = None, end: str = "\n") -> None:
         """Print text with highlighting and no end characters."""

api/core/workflow/graph_engine/entities/graph.py

@@ -334,7 +334,7 @@ class Graph(BaseModel):
 
                     parallel = GraphParallel(
                         start_from_node_id=start_node_id,
-                        parent_parallel_id=parent_parallel.id if parent_parallel else None,
+                        parent_parallel_id=parent_parallel_id,
                         parent_parallel_start_node_id=parent_parallel.start_from_node_id if parent_parallel else None,
                     )
                     parallel_mapping[parallel.id] = parallel

api/core/workflow/graph_engine/graph_engine.py

@@ -103,7 +103,7 @@ class GraphEngine:
         call_depth: int,
         graph: Graph,
         graph_config: Mapping[str, Any],
-        variable_pool: VariablePool,
+        graph_runtime_state: GraphRuntimeState,
         max_execution_steps: int,
         max_execution_time: int,
         thread_pool_id: Optional[str] = None,
@@ -140,7 +140,7 @@ class GraphEngine:
             call_depth=call_depth,
         )
 
-        self.graph_runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter())
+        self.graph_runtime_state = graph_runtime_state
 
         self.max_execution_steps = max_execution_steps
         self.max_execution_time = max_execution_time

api/core/workflow/nodes/agent/agent_node.py

@@ -1,19 +1,22 @@
 import json
+import uuid
 from collections.abc import Generator, Mapping, Sequence
 from typing import Any, Optional, cast
 
+from packaging.version import Version
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from core.agent.entities import AgentToolEntity
 from core.agent.plugin_entities import AgentStrategyParameter
+from core.agent.strategy.plugin import PluginAgentStrategy
 from core.memory.token_buffer_memory import TokenBufferMemory
 from core.model_manager import ModelInstance, ModelManager
 from core.model_runtime.entities.model_entities import AIModelEntity, ModelType
 from core.plugin.impl.exc import PluginDaemonClientSideError
 from core.plugin.impl.plugin import PluginInstaller
 from core.provider_manager import ProviderManager
-from core.tools.entities.tool_entities import ToolParameter, ToolProviderType
+from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter, ToolProviderType
 from core.tools.tool_manager import ToolManager
 from core.variables.segments import StringSegment
 from core.workflow.entities.node_entities import NodeRunResult
@@ -72,12 +75,14 @@ class AgentNode(ToolNode):
             agent_parameters=agent_parameters,
             variable_pool=self.graph_runtime_state.variable_pool,
             node_data=node_data,
+            strategy=strategy,
         )
         parameters_for_log = self._generate_agent_parameters(
             agent_parameters=agent_parameters,
             variable_pool=self.graph_runtime_state.variable_pool,
             node_data=node_data,
             for_log=True,
+            strategy=strategy,
         )
 
         # get conversation id
@@ -102,6 +107,32 @@ class AgentNode(ToolNode):
 
         try:
             # convert tool messages
+            agent_thoughts: list = []
+
+            thought_log_message = ToolInvokeMessage(
+                type=ToolInvokeMessage.MessageType.LOG,
+                message=ToolInvokeMessage.LogMessage(
+                    id=str(uuid.uuid4()),
+                    label=f"Agent Strategy: {cast(AgentNodeData, self.node_data).agent_strategy_name}",
+                    parent_id=None,
+                    error=None,
+                    status=ToolInvokeMessage.LogMessage.LogStatus.START,
+                    data={
+                        "strategy": cast(AgentNodeData, self.node_data).agent_strategy_name,
+                        "parameters": parameters_for_log,
+                        "thought_process": "Agent strategy execution started",
+                    },
+                    metadata={
+                        "icon": self.agent_strategy_icon,
+                        "agent_strategy": cast(AgentNodeData, self.node_data).agent_strategy_name,
+                    },
+                ),
+            )
+
+            def enhanced_message_stream():
+                yield thought_log_message
+
+                yield from message_stream
 
             yield from self._transform_message(
                 message_stream,
@@ -110,6 +141,7 @@ class AgentNode(ToolNode):
                     "agent_strategy": cast(AgentNodeData, self.node_data).agent_strategy_name,
                 },
                 parameters_for_log,
+                agent_thoughts,
             )
         except PluginDaemonClientSideError as e:
             yield RunCompletedEvent(
@@ -127,6 +159,7 @@ class AgentNode(ToolNode):
         variable_pool: VariablePool,
         node_data: AgentNodeData,
         for_log: bool = False,
+        strategy: PluginAgentStrategy,
     ) -> dict[str, Any]:
         """
         Generate parameters based on the given tool parameters, variable pool, and node data.
@@ -179,7 +212,7 @@ class AgentNode(ToolNode):
             if parameter.type == "array[tools]":
                 value = cast(list[dict[str, Any]], value)
                 value = [tool for tool in value if tool.get("enabled", False)]
-
+                value = self._filter_mcp_type_tool(strategy, value)
                 for tool in value:
                     if "schemas" in tool:
                         tool.pop("schemas")
@@ -216,9 +249,9 @@ class AgentNode(ToolNode):
                         )
 
                         extra = tool.get("extra", {})
-
+                        runtime_variable_pool = variable_pool if self.node_data.version != "1" else None
                         tool_runtime = ToolManager.get_agent_tool_runtime(
-                            self.tenant_id, self.app_id, entity, self.invoke_from
+                            self.tenant_id, self.app_id, entity, self.invoke_from, runtime_variable_pool
                         )
                         if tool_runtime.entity.description:
                             tool_runtime.entity.description.llm = (
@@ -370,3 +403,16 @@ class AgentNode(ToolNode):
                 except ValueError:
                     model_schema.features.remove(feature)
         return model_schema
+
+    def _filter_mcp_type_tool(self, strategy: PluginAgentStrategy, tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        """
+        Filter MCP type tool
+        :param strategy: plugin agent strategy
+        :param tool: tool
+        :return: filtered tool dict
+        """
+        meta_version = strategy.meta_version
+        if meta_version and Version(meta_version) > Version("0.0.1"):
+            return tools
+        else:
+            return [tool for tool in tools if tool.get("type") != ToolProviderType.MCP.value]

api/core/workflow/nodes/http_request/executor.py

@@ -8,6 +8,7 @@ from typing import Any, Literal
 from urllib.parse import urlencode, urlparse
 
 import httpx
+from json_repair import repair_json
 
 from configs import dify_config
 from core.file import file_manager
@@ -178,7 +179,8 @@ class Executor:
                         raise RequestBodyError("json body type should have exactly one item")
                     json_string = self.variable_pool.convert_template(data[0].value).text
                     try:
-                        json_object = json.loads(json_string, strict=False)
+                        repaired = repair_json(json_string)
+                        json_object = json.loads(repaired, strict=False)
                     except json.JSONDecodeError as e:
                         raise RequestBodyError(f"Failed to parse JSON: {json_string}") from e
                     self.json = json_object

api/core/workflow/nodes/iteration/iteration_node.py

@@ -1,5 +1,6 @@
 import contextvars
 import logging
+import time
 import uuid
 from collections.abc import Generator, Mapping, Sequence
 from concurrent.futures import Future, wait
@@ -133,8 +134,11 @@ class IterationNode(BaseNode[IterationNodeData]):
         variable_pool.add([self.node_id, "item"], iterator_list_value[0])
 
         # init graph engine
+        from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState
         from core.workflow.graph_engine.graph_engine import GraphEngine, GraphEngineThreadPool
 
+        graph_runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter())
+
         graph_engine = GraphEngine(
             tenant_id=self.tenant_id,
             app_id=self.app_id,
@@ -146,7 +150,7 @@ class IterationNode(BaseNode[IterationNodeData]):
             call_depth=self.workflow_call_depth,
             graph=iteration_graph,
             graph_config=graph_config,
-            variable_pool=variable_pool,
+            graph_runtime_state=graph_runtime_state,
             max_execution_steps=dify_config.WORKFLOW_MAX_EXECUTION_STEPS,
             max_execution_time=dify_config.WORKFLOW_MAX_EXECUTION_TIME,
             thread_pool_id=self.thread_pool_id,

api/core/workflow/nodes/llm/node.py

@@ -221,15 +221,6 @@ class LLMNode(BaseNode[LLMNodeData]):
                 jinja2_variables=self.node_data.prompt_config.jinja2_variables,
             )
 
-            process_data = {
-                "model_mode": model_config.mode,
-                "prompts": PromptMessageUtil.prompt_messages_to_prompt_for_saving(
-                    model_mode=model_config.mode, prompt_messages=prompt_messages
-                ),
-                "model_provider": model_config.provider,
-                "model_name": model_config.model,
-            }
-
             # handle invoke result
             generator = self._invoke_llm(
                 node_data_model=self.node_data.model,
@@ -253,6 +244,17 @@ class LLMNode(BaseNode[LLMNodeData]):
                 elif isinstance(event, LLMStructuredOutput):
                     structured_output = event
 
+            process_data = {
+                "model_mode": model_config.mode,
+                "prompts": PromptMessageUtil.prompt_messages_to_prompt_for_saving(
+                    model_mode=model_config.mode, prompt_messages=prompt_messages
+                ),
+                "usage": jsonable_encoder(usage),
+                "finish_reason": finish_reason,
+                "model_provider": model_config.provider,
+                "model_name": model_config.model,
+            }
+
             outputs = {"text": result_text, "usage": jsonable_encoder(usage), "finish_reason": finish_reason}
             if structured_output:
                 outputs["structured_output"] = structured_output.structured_output

api/core/workflow/nodes/loop/loop_node.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import time
 from collections.abc import Generator, Mapping, Sequence
 from datetime import UTC, datetime
 from typing import TYPE_CHECKING, Any, Literal, cast
@@ -101,8 +102,11 @@ class LoopNode(BaseNode[LoopNodeData]):
                 loop_variable_selectors[loop_variable.label] = variable_selector
                 inputs[loop_variable.label] = processed_segment.value
 
+        from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState
         from core.workflow.graph_engine.graph_engine import GraphEngine
 
+        graph_runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter())
+
         graph_engine = GraphEngine(
             tenant_id=self.tenant_id,
             app_id=self.app_id,
@@ -114,7 +118,7 @@ class LoopNode(BaseNode[LoopNodeData]):
             call_depth=self.workflow_call_depth,
             graph=loop_graph,
             graph_config=self.graph_config,
-            variable_pool=variable_pool,
+            graph_runtime_state=graph_runtime_state,
             max_execution_steps=dify_config.WORKFLOW_MAX_EXECUTION_STEPS,
             max_execution_time=dify_config.WORKFLOW_MAX_EXECUTION_TIME,
             thread_pool_id=self.thread_pool_id,

api/core/workflow/nodes/node_mapping.py

@@ -73,6 +73,7 @@ NODE_TYPE_CLASSES_MAPPING: Mapping[NodeType, Mapping[str, type[BaseNode]]] = {
     },
     NodeType.TOOL: {
         LATEST_VERSION: ToolNode,
+        "2": ToolNode,
         "1": ToolNode,
     },
     NodeType.VARIABLE_AGGREGATOR: {
@@ -122,6 +123,7 @@ NODE_TYPE_CLASSES_MAPPING: Mapping[NodeType, Mapping[str, type[BaseNode]]] = {
     },
     NodeType.AGENT: {
         LATEST_VERSION: AgentNode,
+        "2": AgentNode,
         "1": AgentNode,
     },
 }

api/core/workflow/nodes/parameter_extractor/parameter_extractor_node.py

@@ -253,7 +253,12 @@ class ParameterExtractorNode(BaseNode):
             status=WorkflowNodeExecutionStatus.SUCCEEDED,
             inputs=inputs,
             process_data=process_data,
-            outputs={"__is_success": 1 if not error else 0, "__reason": error, **result},
+            outputs={
+                "__is_success": 1 if not error else 0,
+                "__reason": error,
+                "__usage": jsonable_encoder(usage),
+                **result,
+            },
             metadata={
                 WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: usage.total_tokens,
                 WorkflowNodeExecutionMetadataKey.TOTAL_PRICE: usage.total_price,

api/core/workflow/nodes/question_classifier/question_classifier_node.py

@@ -145,7 +145,11 @@ class QuestionClassifierNode(LLMNode):
                 "model_provider": model_config.provider,
                 "model_name": model_config.model,
             }
-            outputs = {"class_name": category_name, "class_id": category_id}
+            outputs = {
+                "class_name": category_name,
+                "class_id": category_id,
+                "usage": jsonable_encoder(usage),
+            }
 
             return NodeRunResult(
                 status=WorkflowNodeExecutionStatus.SUCCEEDED,

api/core/workflow/nodes/tool/entities.py

@@ -42,6 +42,10 @@ class ToolNodeData(BaseNodeData, ToolEntity):
         def check_type(cls, value, validation_info: ValidationInfo):
             typ = value
             value = validation_info.data.get("value")
+
+            if value is None:
+                return typ
+
             if typ == "mixed" and not isinstance(value, str):
                 raise ValueError("value must be a string")
             elif typ == "variable":
@@ -55,3 +59,22 @@ class ToolNodeData(BaseNodeData, ToolEntity):
             return typ
 
     tool_parameters: dict[str, ToolInput]
+
+    @field_validator("tool_parameters", mode="before")
+    @classmethod
+    def filter_none_tool_inputs(cls, value):
+        if not isinstance(value, dict):
+            return value
+
+        return {
+            key: tool_input
+            for key, tool_input in value.items()
+            if tool_input is not None and cls._has_valid_value(tool_input)
+        }
+
+    @staticmethod
+    def _has_valid_value(tool_input):
+        """Check if the value is valid"""
+        if isinstance(tool_input, dict):
+            return tool_input.get("value") is not None
+        return getattr(tool_input, "value", None) is not None

api/core/workflow/nodes/tool/tool_node.py

@@ -1,11 +1,12 @@
 from collections.abc import Generator, Mapping, Sequence
-from typing import Any, cast
+from typing import Any, Optional, cast
 
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler
 from core.file import File, FileTransferMethod
+from core.model_runtime.entities.llm_entities import LLMUsage
 from core.plugin.impl.exc import PluginDaemonClientSideError
 from core.plugin.impl.plugin import PluginInstaller
 from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter
@@ -66,8 +67,9 @@ class ToolNode(BaseNode[ToolNodeData]):
         try:
             from core.tools.tool_manager import ToolManager
 
+            variable_pool = self.graph_runtime_state.variable_pool if self.node_data.version != "1" else None
             tool_runtime = ToolManager.get_workflow_tool_runtime(
-                self.tenant_id, self.app_id, self.node_id, self.node_data, self.invoke_from
+                self.tenant_id, self.app_id, self.node_id, self.node_data, self.invoke_from, variable_pool
             )
         except ToolNodeError as e:
             yield RunCompletedEvent(
@@ -94,7 +96,6 @@ class ToolNode(BaseNode[ToolNodeData]):
             node_data=self.node_data,
             for_log=True,
         )
-
         # get conversation id
         conversation_id = self.graph_runtime_state.variable_pool.get(["sys", SystemVariableKey.CONVERSATION_ID])
 
@@ -190,6 +191,7 @@ class ToolNode(BaseNode[ToolNodeData]):
         messages: Generator[ToolInvokeMessage, None, None],
         tool_info: Mapping[str, Any],
         parameters_for_log: dict[str, Any],
+        agent_thoughts: Optional[list] = None,
     ) -> Generator:
         """
         Convert ToolInvokeMessages into tuple[plain_text, files]
@@ -208,7 +210,7 @@ class ToolNode(BaseNode[ToolNodeData]):
 
         agent_logs: list[AgentLogEvent] = []
         agent_execution_metadata: Mapping[WorkflowNodeExecutionMetadataKey, Any] = {}
-
+        llm_usage: LLMUsage | None = None
         variables: dict[str, Any] = {}
 
         for message in message_stream:
@@ -276,13 +278,15 @@ class ToolNode(BaseNode[ToolNodeData]):
             elif message.type == ToolInvokeMessage.MessageType.JSON:
                 assert isinstance(message.message, ToolInvokeMessage.JsonMessage)
                 if self.node_type == NodeType.AGENT:
-                    msg_metadata = message.message.json_object.pop("execution_metadata", {})
+                    msg_metadata: dict[str, Any] = message.message.json_object.pop("execution_metadata", {})
+                    llm_usage = LLMUsage.from_metadata(msg_metadata)
                     agent_execution_metadata = {
-                        key: value
+                        WorkflowNodeExecutionMetadataKey(key): value
                         for key, value in msg_metadata.items()
                         if key in WorkflowNodeExecutionMetadataKey.__members__.values()
                     }
-                json.append(message.message.json_object)
+                if message.message.json_object is not None:
+                    json.append(message.message.json_object)
             elif message.type == ToolInvokeMessage.MessageType.LINK:
                 assert isinstance(message.message, ToolInvokeMessage.TextMessage)
                 stream_text = f"Link: {message.message.text}\n"
@@ -325,6 +329,7 @@ class ToolNode(BaseNode[ToolNodeData]):
                             icon = current_plugin.declaration.icon
                         except StopIteration:
                             pass
+                        icon_dark = None
                         try:
                             builtin_tool = next(
                                 provider
@@ -335,10 +340,12 @@ class ToolNode(BaseNode[ToolNodeData]):
                                 if provider.name == dict_metadata["provider"]
                             )
                             icon = builtin_tool.icon
+                            icon_dark = builtin_tool.icon_dark
                         except StopIteration:
                             pass
 
                         dict_metadata["icon"] = icon
+                        dict_metadata["icon_dark"] = icon_dark
                         message.message.metadata = dict_metadata
                 agent_log = AgentLogEvent(
                     id=message.message.id,
@@ -367,16 +374,41 @@ class ToolNode(BaseNode[ToolNodeData]):
 
                 yield agent_log
 
+        # Add agent_logs to outputs['json'] to ensure frontend can access thinking process
+        json_output: list[dict[str, Any]] = []
+
+        # Step 1: append each agent log as its own dict.
+        if agent_logs:
+            for log in agent_logs:
+                json_output.append(
+                    {
+                        "id": log.id,
+                        "parent_id": log.parent_id,
+                        "error": log.error,
+                        "status": log.status,
+                        "data": log.data,
+                        "label": log.label,
+                        "metadata": log.metadata,
+                        "node_id": log.node_id,
+                    }
+                )
+        # Step 2: normalize JSON into {"data": [...]}.change json to list[dict]
+        if json:
+            json_output.extend(json)
+        else:
+            json_output.append({"data": []})
+
         yield RunCompletedEvent(
             run_result=NodeRunResult(
                 status=WorkflowNodeExecutionStatus.SUCCEEDED,
-                outputs={"text": text, "files": ArrayFileSegment(value=files), "json": json, **variables},
+                outputs={"text": text, "files": ArrayFileSegment(value=files), "json": json_output, **variables},
                 metadata={
                     **agent_execution_metadata,
                     WorkflowNodeExecutionMetadataKey.TOOL_INFO: tool_info,
                     WorkflowNodeExecutionMetadataKey.AGENT_LOG: agent_logs,
                 },
                 inputs=parameters_for_log,
+                llm_usage=llm_usage,
             )
         )
 

api/core/workflow/workflow_entry.py

@@ -69,6 +69,7 @@ class WorkflowEntry:
             raise ValueError("Max workflow call depth {} reached.".format(workflow_call_max_depth))
 
         # init workflow run state
+        graph_runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter())
         self.graph_engine = GraphEngine(
             tenant_id=tenant_id,
             app_id=app_id,
@@ -80,7 +81,7 @@ class WorkflowEntry:
             call_depth=call_depth,
             graph=graph,
             graph_config=graph_config,
-            variable_pool=variable_pool,
+            graph_runtime_state=graph_runtime_state,
             max_execution_steps=dify_config.WORKFLOW_MAX_EXECUTION_STEPS,
             max_execution_time=dify_config.WORKFLOW_MAX_EXECUTION_TIME,
             thread_pool_id=thread_pool_id,

api/extensions/ext_blueprints.py

@@ -10,6 +10,7 @@ def init_app(app: DifyApp):
     from controllers.console import bp as console_app_bp
     from controllers.files import bp as files_bp
     from controllers.inner_api import bp as inner_api_bp
+    from controllers.mcp import bp as mcp_bp
     from controllers.service_api import bp as service_api_bp
     from controllers.web import bp as web_bp
 
@@ -46,3 +47,4 @@ def init_app(app: DifyApp):
     app.register_blueprint(files_bp)
 
     app.register_blueprint(inner_api_bp)
+    app.register_blueprint(mcp_bp)

api/extensions/ext_login.py

@@ -10,7 +10,7 @@ from dify_app import DifyApp
 from extensions.ext_database import db
 from libs.passport import PassportService
 from models.account import Account, Tenant, TenantAccountJoin
-from models.model import EndUser
+from models.model import AppMCPServer, EndUser
 from services.account_service import AccountService
 
 login_manager = flask_login.LoginManager()
@@ -74,6 +74,21 @@ def load_user_from_request(request_from_flask_login):
         if not end_user:
             raise NotFound("End user not found.")
         return end_user
+    elif request.blueprint == "mcp":
+        server_code = request.view_args.get("server_code") if request.view_args else None
+        if not server_code:
+            raise Unauthorized("Invalid Authorization token.")
+        app_mcp_server = db.session.query(AppMCPServer).filter(AppMCPServer.server_code == server_code).first()
+        if not app_mcp_server:
+            raise NotFound("App MCP server not found.")
+        end_user = (
+            db.session.query(EndUser)
+            .filter(EndUser.external_user_id == app_mcp_server.id, EndUser.type == "mcp")
+            .first()
+        )
+        if not end_user:
+            raise NotFound("End user not found.")
+        return end_user
 
 
 @user_logged_in.connect

api/extensions/ext_otel.py

@@ -12,6 +12,7 @@ from flask_login import user_loaded_from_request, user_logged_in  # type: ignore
 
 from configs import dify_config
 from dify_app import DifyApp
+from libs.helper import extract_tenant_id
 from models import Account, EndUser
 
 
@@ -24,11 +25,8 @@ def on_user_loaded(_sender, user: Union["Account", "EndUser"]):
         if user:
             try:
                 current_span = get_current_span()
-                if isinstance(user, Account) and user.current_tenant_id:
-                    tenant_id = user.current_tenant_id
-                elif isinstance(user, EndUser):
-                    tenant_id = user.tenant_id
-                else:
+                tenant_id = extract_tenant_id(user)
+                if not tenant_id:
                     return
                 if current_span:
                     current_span.set_attribute("service.tenant.id", tenant_id)

api/extensions/ext_redis.py

@@ -1,6 +1,10 @@
+import functools
+import logging
+from collections.abc import Callable
 from typing import Any, Union
 
 import redis
+from redis import RedisError
 from redis.cache import CacheConfig
 from redis.cluster import ClusterNode, RedisCluster
 from redis.connection import Connection, SSLConnection
@@ -9,6 +13,8 @@ from redis.sentinel import Sentinel
 from configs import dify_config
 from dify_app import DifyApp
 
+logger = logging.getLogger(__name__)
+
 
 class RedisClientWrapper:
     """
@@ -115,3 +121,25 @@ def init_app(app: DifyApp):
         redis_client.initialize(redis.Redis(connection_pool=pool))
 
     app.extensions["redis"] = redis_client
+
+
+def redis_fallback(default_return: Any = None):
+    """
+    decorator to handle Redis operation exceptions and return a default value when Redis is unavailable.
+
+    Args:
+        default_return: The value to return when a Redis operation fails. Defaults to None.
+    """
+
+    def decorator(func: Callable):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except RedisError as e:
+                logger.warning(f"Redis operation failed in {func.__name__}: {str(e)}", exc_info=True)
+                return default_return
+
+        return wrapper
+
+    return decorator

api/factories/agent_factory.py

@@ -10,6 +10,6 @@ def get_plugin_agent_strategy(
     agent_provider = manager.fetch_agent_strategy_provider(tenant_id, agent_strategy_provider_name)
     for agent_strategy in agent_provider.declaration.strategies:
         if agent_strategy.identity.name == agent_strategy_name:
-            return PluginAgentStrategy(tenant_id, agent_strategy)
+            return PluginAgentStrategy(tenant_id, agent_strategy, agent_provider.meta.version)
 
     raise ValueError(f"Agent strategy {agent_strategy_name} not found")

api/fields/app_fields.py

@@ -1,8 +1,21 @@
+import json
+
 from flask_restful import fields
 
 from fields.workflow_fields import workflow_partial_fields
 from libs.helper import AppIconUrlField, TimestampField
 
+
+class JsonStringField(fields.Raw):
+    def format(self, value):
+        if isinstance(value, str):
+            try:
+                return json.loads(value)
+            except (json.JSONDecodeError, TypeError):
+                return value
+        return value
+
+
 app_detail_kernel_fields = {
     "id": fields.String,
     "name": fields.String,
@@ -218,3 +231,14 @@ app_import_fields = {
 app_import_check_dependencies_fields = {
     "leaked_dependencies": fields.List(fields.Nested(leaked_dependency_fields)),
 }
+
+app_server_fields = {
+    "id": fields.String,
+    "name": fields.String,
+    "server_code": fields.String,
+    "description": fields.String,
+    "status": fields.String,
+    "parameters": JsonStringField,
+    "created_at": TimestampField,
+    "updated_at": TimestampField,
+}

api/fields/workflow_fields.py

@@ -17,6 +17,7 @@ class EnvironmentVariableField(fields.Raw):
                 "name": value.name,
                 "value": encrypter.obfuscated_token(value.value),
                 "value_type": value.value_type.value,
+                "description": value.description,
             }
         if isinstance(value, Variable):
             return {
@@ -24,6 +25,7 @@ class EnvironmentVariableField(fields.Raw):
                 "name": value.name,
                 "value": value.value,
                 "value_type": value.value_type.value,
+                "description": value.description,
             }
         if isinstance(value, dict):
             value_type = value.get("value_type")

api/libs/helper.py

@@ -25,6 +25,31 @@ from extensions.ext_redis import redis_client
 
 if TYPE_CHECKING:
     from models.account import Account
+    from models.model import EndUser
+
+
+def extract_tenant_id(user: Union["Account", "EndUser"]) -> str | None:
+    """
+    Extract tenant_id from Account or EndUser object.
+
+    Args:
+        user: Account or EndUser object
+
+    Returns:
+        tenant_id string if available, None otherwise
+
+    Raises:
+        ValueError: If user is neither Account nor EndUser
+    """
+    from models.account import Account
+    from models.model import EndUser
+
+    if isinstance(user, Account):
+        return user.current_tenant_id
+    elif isinstance(user, EndUser):
+        return user.tenant_id
+    else:
+        raise ValueError(f"Invalid user type: {type(user)}. Expected Account or EndUser.")
 
 
 def run(script):