if session unauthorized, rejoin

The collaboration feature increased workflow operator z-index from z-10 to z-[60].
This caused the AppInfo ContentDialog (z-30) to appear below the operator buttons.
Increased ContentDialog z-index to z-[70] to ensure proper layer hierarchy.

.devcontainer/devcontainer.json

@@ -11,7 +11,7 @@
 			"nodeGypDependencies": true,
 			"version": "lts"
 		},
-		"ghcr.io/devcontainers-extra/features/npm-package:1": {
+		"ghcr.io/devcontainers-contrib/features/npm-package:1": {
 			"package": "typescript",
 			"version": "latest"
 		},

.github/workflows/style.yml

@@ -103,11 +103,6 @@ jobs:
         run: |
           pnpm run lint
 
-      - name: Web type check
-        if: steps.changed-files.outputs.any_changed == 'true'
-        working-directory: ./web
-        run: pnpm run type-check
-
   docker-compose-template:
     name: Docker Compose Template
     runs-on: ubuntu-latest

.gitignore

@@ -97,7 +97,6 @@ __pypackages__/
 
 # Celery stuff
 celerybeat-schedule
-celerybeat-schedule.db
 celerybeat.pid
 
 # SageMath parsed files

.vscode/launch.json.template

@@ -40,7 +40,7 @@
                 "-c",
                 "1",
                 "-Q",
-                "dataset,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation,priority_pipeline,pipeline",
+                "dataset,generation,mail,ops_trace",
                 "--loglevel",
                 "INFO"
             ],

README.md

@@ -63,7 +63,7 @@ Dify is an open-source platform for developing LLM applications. Its intuitive i
 > - CPU >= 2 Core
 > - RAM >= 4 GiB
 
-<br/>
+</br>
 
 The easiest way to start the Dify server is through [Docker Compose](docker/docker-compose.yaml). Before running Dify with the following commands, make sure that [Docker](https://docs.docker.com/get-docker/) and [Docker Compose](https://docs.docker.com/compose/install/) are installed on your machine:
 
@@ -109,15 +109,15 @@ All of Dify's offerings come with corresponding APIs, so you could effortlessly
 
 ## Using Dify
 
-- **Cloud <br/>**
+- **Cloud </br>**
   We host a [Dify Cloud](https://dify.ai) service for anyone to try with zero setup. It provides all the capabilities of the self-deployed version, and includes 200 free GPT-4 calls in the sandbox plan.
 
-- **Self-hosting Dify Community Edition<br/>**
+- **Self-hosting Dify Community Edition</br>**
   Quickly get Dify running in your environment with this [starter guide](#quick-start).
   Use our [documentation](https://docs.dify.ai) for further references and more in-depth instructions.
 
-- **Dify for enterprise / organizations<br/>**
-  We provide additional enterprise-centric features. [Log your questions for us through this chatbot](https://udify.app/chat/22L1zSxg6yW1cWQg) or [send us an email](mailto:business@dify.ai?subject=%5BGitHub%5DBusiness%20License%20Inquiry) to discuss enterprise needs. <br/>
+- **Dify for enterprise / organizations</br>**
+  We provide additional enterprise-centric features. [Log your questions for us through this chatbot](https://udify.app/chat/22L1zSxg6yW1cWQg) or [send us an email](mailto:business@dify.ai?subject=%5BGitHub%5DBusiness%20License%20Inquiry) to discuss enterprise needs. </br>
 
   > For startups and small businesses using AWS, check out [Dify Premium on AWS Marketplace](https://aws.amazon.com/marketplace/pp/prodview-t22mebxzwjhu6) and deploy it to your own AWS VPC with one click. It's an affordable AMI offering with the option to create apps with custom logo and branding.
 

api/.env.example

@@ -30,6 +30,9 @@ INTERNAL_FILES_URL=http://127.0.0.1:5001
 # The time in seconds after the signature is rejected
 FILES_ACCESS_TIMEOUT=300
 
+# Collaboration mode toggle
+ENABLE_COLLABORATION_MODE=false
+
 # Access token expiration time in minutes
 ACCESS_TOKEN_EXPIRE_MINUTES=60
 
@@ -156,9 +159,6 @@ SUPABASE_URL=your-server-url
 # CORS configuration
 WEB_API_CORS_ALLOW_ORIGINS=http://localhost:3000,*
 CONSOLE_CORS_ALLOW_ORIGINS=http://localhost:3000,*
-# Set COOKIE_DOMAIN when the console frontend and API are on different subdomains.
-# Provide the registrable domain (e.g. example.com); leading dots are optional.
-COOKIE_DOMAIN=
 
 # Vector database configuration
 # Supported values are `weaviate`, `qdrant`, `milvus`, `myscale`, `relyt`, `pgvector`, `pgvecto-rs`, `chroma`, `opensearch`, `oracle`, `tencent`, `elasticsearch`, `elasticsearch-ja`, `analyticdb`, `couchbase`, `vikingdb`, `oceanbase`, `opengauss`, `tablestore`,`vastbase`,`tidb`,`tidb_on_qdrant`,`baidu`,`lindorm`,`huawei_cloud`,`upstash`, `matrixone`.
@@ -437,9 +437,6 @@ CODE_EXECUTION_SSL_VERIFY=True
 CODE_EXECUTION_POOL_MAX_CONNECTIONS=100
 CODE_EXECUTION_POOL_MAX_KEEPALIVE_CONNECTIONS=20
 CODE_EXECUTION_POOL_KEEPALIVE_EXPIRY=5.0
-CODE_EXECUTION_CONNECT_TIMEOUT=10
-CODE_EXECUTION_READ_TIMEOUT=60
-CODE_EXECUTION_WRITE_TIMEOUT=10
 CODE_MAX_NUMBER=9223372036854775807
 CODE_MIN_NUMBER=-9223372036854775808
 CODE_MAX_STRING_LENGTH=400000

api/.vscode/launch.json.example

@@ -54,7 +54,7 @@
                 "--loglevel",
                 "DEBUG",
                 "-Q",
-                "dataset,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation,priority_pipeline,pipeline"
+                "dataset,generation,mail,ops_trace,app_deletion"
             ]
         }
     ]

api/README.md

@@ -80,7 +80,7 @@
 1. If you need to handle and debug the async tasks (e.g. dataset importing and documents indexing), please start the worker service.
 
 ```bash
-uv run celery -A app.celery worker -P gevent -c 2 --loglevel INFO -Q dataset,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation,priority_pipeline,pipeline
+uv run celery -A app.celery worker -P gevent -c 2 --loglevel INFO -Q dataset,generation,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation
 ```
 
 Additionally, if you want to debug the celery scheduled tasks, you can run the following command in another terminal to start the beat service:

api/app.py

@@ -1,3 +1,4 @@
+import os
 import sys
 
 
@@ -8,10 +9,16 @@ def is_db_command():
 
 
 # create app
+celery = None
+flask_app = None
+socketio_app = None
+
 if is_db_command():
     from app_factory import create_migrations_app
 
     app = create_migrations_app()
+    socketio_app = app
+    flask_app = app
 else:
     # It seems that JetBrains Python debugger does not work well with gevent,
     # so we need to disable gevent in debug mode.
@@ -33,8 +40,15 @@ else:
 
     from app_factory import create_app
 
-    app = create_app()
-    celery = app.extensions["celery"]
+    socketio_app, flask_app = create_app()
+    app = flask_app
+    celery = flask_app.extensions["celery"]
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=5001)
+    from gevent import pywsgi
+    from geventwebsocket.handler import WebSocketHandler
+
+    host = os.environ.get("HOST", "0.0.0.0")
+    port = int(os.environ.get("PORT", 5001))
+    server = pywsgi.WSGIServer((host, port), socketio_app, handler_class=WebSocketHandler)
+    server.serve_forever()

api/app_factory.py

@@ -31,14 +31,22 @@ def create_flask_app_with_configs() -> DifyApp:
     return dify_app
 
 
-def create_app() -> DifyApp:
+def create_app() -> tuple[any, DifyApp]:
     start_time = time.perf_counter()
     app = create_flask_app_with_configs()
     initialize_extensions(app)
+
+    import socketio
+
+    from extensions.ext_socketio import sio
+
+    sio.app = app
+    socketio_app = socketio.WSGIApp(sio, app)
+
     end_time = time.perf_counter()
     if dify_config.DEBUG:
         logger.info("Finished create_app (%s ms)", round((end_time - start_time) * 1000, 2))
-    return app
+    return socketio_app, app
 
 
 def initialize_extensions(app: DifyApp):

api/commands.py

@@ -321,8 +321,6 @@ def migrate_knowledge_vector_database():
             )
 
             datasets = db.paginate(select=stmt, page=page, per_page=50, max_per_page=50, error_out=False)
-            if not datasets.items:
-                break
         except SQLAlchemyError:
             raise
 

api/configs/feature/__init__.py

@@ -337,11 +337,6 @@ class HttpConfig(BaseSettings):
     HTTP-related configurations for the application
     """
 
-    COOKIE_DOMAIN: str = Field(
-        description="Explicit cookie domain for console/service cookies when sharing across subdomains",
-        default="",
-    )
-
     API_COMPRESSION_ENABLED: bool = Field(
         description="Enable or disable gzip compression for HTTP responses",
         default=False,
@@ -1053,6 +1048,13 @@ class PositionConfig(BaseSettings):
         return {item.strip() for item in self.POSITION_TOOL_EXCLUDES.split(",") if item.strip() != ""}
 
 
+class CollaborationConfig(BaseSettings):
+    ENABLE_COLLABORATION_MODE: bool = Field(
+        description="Whether to enable collaboration mode features across the workspace",
+        default=False,
+    )
+
+
 class LoginConfig(BaseSettings):
     ENABLE_EMAIL_CODE_LOGIN: bool = Field(
         description="whether to enable email code login",
@@ -1141,6 +1143,7 @@ class FeatureConfig(
     WorkflowConfig,
     WorkflowNodeExecutionConfig,
     WorkspaceConfig,
+    CollaborationConfig,
     LoginConfig,
     AccountConfig,
     SwaggerUIConfig,

api/configs/middleware/__init__.py

@@ -145,7 +145,7 @@ class DatabaseConfig(BaseSettings):
         default="postgresql",
     )
 
-    @computed_field  # type: ignore[prop-decorator]
+    @computed_field  # type: ignore[misc]
     @property
     def SQLALCHEMY_DATABASE_URI(self) -> str:
         db_extras = (
@@ -198,7 +198,7 @@ class DatabaseConfig(BaseSettings):
         default=os.cpu_count() or 1,
     )
 
-    @computed_field  # type: ignore[prop-decorator]
+    @computed_field  # type: ignore[misc]
     @property
     def SQLALCHEMY_ENGINE_OPTIONS(self) -> dict[str, Any]:
         # Parse DB_EXTRAS for 'options'

api/constants/__init__.py

@@ -56,14 +56,10 @@ else:
     }
 DOCUMENT_EXTENSIONS: set[str] = convert_to_lower_and_upper_set(_doc_extensions)
 
-# console
 COOKIE_NAME_ACCESS_TOKEN = "access_token"
 COOKIE_NAME_REFRESH_TOKEN = "refresh_token"
-COOKIE_NAME_CSRF_TOKEN = "csrf_token"
-
-# webapp
-COOKIE_NAME_WEBAPP_ACCESS_TOKEN = "webapp_access_token"
 COOKIE_NAME_PASSPORT = "passport"
+COOKIE_NAME_CSRF_TOKEN = "csrf_token"
 
 HEADER_NAME_CSRF_TOKEN = "X-CSRF-Token"
 HEADER_NAME_APP_CODE = "X-App-Code"

api/constants/languages.py

@@ -31,9 +31,3 @@ def supported_language(lang):
 
     error = f"{lang} is not a valid language."
     raise ValueError(error)
-
-
-def get_valid_language(lang: str | None) -> str:
-    if lang and lang in languages:
-        return lang
-    return languages[0]

api/controllers/common/helpers.py

@@ -24,7 +24,7 @@ except ImportError:
         )
     else:
         warnings.warn("To use python-magic guess MIMETYPE, you need to install `libmagic`", stacklevel=2)
-    magic = None  # type: ignore[assignment]
+    magic = None  # type: ignore
 
 from pydantic import BaseModel
 

api/controllers/console/__init__.py

@@ -58,11 +58,13 @@ from .app import (
     mcp_server,
     message,
     model_config,
+    online_user,
     ops_trace,
     site,
     statistic,
     workflow,
     workflow_app_log,
+    workflow_comment,
     workflow_draft_variable,
     workflow_run,
     workflow_statistic,

api/controllers/console/app/online_user.py

@@ -0,0 +1,339 @@
+import json
+import time
+
+from werkzeug.wrappers import Request as WerkzeugRequest
+
+from extensions.ext_redis import redis_client
+from extensions.ext_socketio import sio
+from libs.passport import PassportService
+from libs.token import extract_access_token
+from services.account_service import AccountService
+
+SESSION_STATE_TTL_SECONDS = 3600
+WORKFLOW_ONLINE_USERS_PREFIX = "workflow_online_users:"
+WORKFLOW_LEADER_PREFIX = "workflow_leader:"
+WS_SID_MAP_PREFIX = "ws_sid_map:"
+
+
+def _workflow_key(workflow_id: str) -> str:
+    return f"{WORKFLOW_ONLINE_USERS_PREFIX}{workflow_id}"
+
+
+def _leader_key(workflow_id: str) -> str:
+    return f"{WORKFLOW_LEADER_PREFIX}{workflow_id}"
+
+
+def _sid_key(sid: str) -> str:
+    return f"{WS_SID_MAP_PREFIX}{sid}"
+
+
+def _refresh_session_state(workflow_id: str, sid: str) -> None:
+    """
+    Refresh TTLs for workflow + session keys so healthy sessions do not linger forever after crashes.
+    """
+    workflow_key = _workflow_key(workflow_id)
+    sid_key = _sid_key(sid)
+    if redis_client.exists(workflow_key):
+        redis_client.expire(workflow_key, SESSION_STATE_TTL_SECONDS)
+    if redis_client.exists(sid_key):
+        redis_client.expire(sid_key, SESSION_STATE_TTL_SECONDS)
+
+
+@sio.on("connect")
+def socket_connect(sid, environ, auth):
+    """
+    WebSocket connect event, do authentication here.
+    """
+    token = None
+    if auth and isinstance(auth, dict):
+        token = auth.get("token")
+
+    if not token:
+        try:
+            request_environ = WerkzeugRequest(environ)
+            token = extract_access_token(request_environ)
+        except Exception:
+            token = None
+
+    if not token:
+        return False
+
+    try:
+        decoded = PassportService().verify(token)
+        user_id = decoded.get("user_id")
+        if not user_id:
+            return False
+
+        with sio.app.app_context():
+            user = AccountService.load_logged_in_account(account_id=user_id)
+            if not user:
+                return False
+
+            sio.save_session(sid, {"user_id": user.id, "username": user.name, "avatar": user.avatar})
+
+            return True
+
+    except Exception:
+        return False
+
+
+@sio.on("user_connect")
+def handle_user_connect(sid, data):
+    """
+    Handle user connect event. Each session (tab) is treated as an independent collaborator.
+    """
+
+    workflow_id = data.get("workflow_id")
+    if not workflow_id:
+        return {"msg": "workflow_id is required"}, 400
+
+    session = sio.get_session(sid)
+    user_id = session.get("user_id")
+
+    if not user_id:
+        return {"msg": "unauthorized"}, 401
+
+    # Each session is stored independently with sid as key
+    session_info = {
+        "user_id": user_id,
+        "username": session.get("username", "Unknown"),
+        "avatar": session.get("avatar", None),
+        "sid": sid,
+        "connected_at": int(time.time()),  # Add timestamp to differentiate tabs
+    }
+
+    workflow_key = _workflow_key(workflow_id)
+    # Store session info with sid as key
+    redis_client.hset(workflow_key, sid, json.dumps(session_info))
+    redis_client.set(
+        _sid_key(sid),
+        json.dumps({"workflow_id": workflow_id, "user_id": user_id}),
+        ex=SESSION_STATE_TTL_SECONDS,
+    )
+    _refresh_session_state(workflow_id, sid)
+
+    # Leader election: first session becomes the leader
+    leader_sid = get_or_set_leader(workflow_id, sid)
+    is_leader = leader_sid == sid
+
+    sio.enter_room(sid, workflow_id)
+    broadcast_online_users(workflow_id)
+
+    # Notify this session of their leader status
+    sio.emit("status", {"isLeader": is_leader}, room=sid)
+
+    return {"msg": "connected", "user_id": user_id, "sid": sid, "isLeader": is_leader}
+
+
+@sio.on("disconnect")
+def handle_disconnect(sid):
+    """
+    Handle session disconnect event. Remove the specific session from online users.
+    """
+    mapping = redis_client.get(_sid_key(sid))
+    if mapping:
+        data = json.loads(mapping)
+        workflow_id = data["workflow_id"]
+
+        # Remove this specific session
+        redis_client.hdel(_workflow_key(workflow_id), sid)
+        redis_client.delete(_sid_key(sid))
+
+        # Handle leader re-election if the leader session disconnected
+        handle_leader_disconnect(workflow_id, sid)
+
+        broadcast_online_users(workflow_id)
+
+
+def _clear_session_state(workflow_id: str, sid: str) -> None:
+    redis_client.hdel(_workflow_key(workflow_id), sid)
+    redis_client.delete(_sid_key(sid))
+
+
+def _is_session_active(workflow_id: str, sid: str) -> bool:
+    if not sid:
+        return False
+
+    try:
+        if not sio.manager.is_connected(sid, "/"):
+            return False
+    except AttributeError:
+        return False
+
+    if not redis_client.hexists(_workflow_key(workflow_id), sid):
+        return False
+
+    if not redis_client.exists(_sid_key(sid)):
+        return False
+
+    return True
+
+
+def get_or_set_leader(workflow_id: str, sid: str) -> str:
+    """
+    Get current leader session or set this session as leader if no valid leader exists.
+    Returns the leader session id (sid).
+    """
+    raw_leader = redis_client.get(_leader_key(workflow_id))
+    current_leader = raw_leader.decode("utf-8") if isinstance(raw_leader, bytes) else raw_leader
+    leader_replaced = False
+
+    if current_leader and not _is_session_active(workflow_id, current_leader):
+        _clear_session_state(workflow_id, current_leader)
+        redis_client.delete(_leader_key(workflow_id))
+        current_leader = None
+        leader_replaced = True
+
+    if not current_leader:
+        redis_client.set(_leader_key(workflow_id), sid, ex=SESSION_STATE_TTL_SECONDS)  # Expire in 1 hour
+        if leader_replaced:
+            broadcast_leader_change(workflow_id, sid)
+        return sid
+
+    return current_leader
+
+
+def handle_leader_disconnect(workflow_id, disconnected_sid):
+    """
+    Handle leader re-election when a session disconnects.
+    If the disconnected session was the leader, elect a new leader from remaining sessions.
+    """
+    current_leader = redis_client.get(_leader_key(workflow_id))
+
+    if current_leader:
+        current_leader = current_leader.decode("utf-8") if isinstance(current_leader, bytes) else current_leader
+
+        if current_leader == disconnected_sid:
+            # Leader session disconnected, elect a new leader
+            sessions_json = redis_client.hgetall(_workflow_key(workflow_id))
+
+            if sessions_json:
+                # Get the first remaining session as new leader
+                new_leader_sid = list(sessions_json.keys())[0]
+                if isinstance(new_leader_sid, bytes):
+                    new_leader_sid = new_leader_sid.decode("utf-8")
+
+                redis_client.set(_leader_key(workflow_id), new_leader_sid, ex=SESSION_STATE_TTL_SECONDS)
+
+                # Notify all sessions about the new leader
+                broadcast_leader_change(workflow_id, new_leader_sid)
+            else:
+                # No sessions left, remove leader
+                redis_client.delete(_leader_key(workflow_id))
+
+
+def broadcast_leader_change(workflow_id, new_leader_sid):
+    """
+    Broadcast leader change to all sessions in the workflow.
+    """
+    sessions_json = redis_client.hgetall(_workflow_key(workflow_id))
+
+    for sid, session_info_json in sessions_json.items():
+        try:
+            sid_str = sid.decode("utf-8") if isinstance(sid, bytes) else sid
+            is_leader = sid_str == new_leader_sid
+            # Emit to each session whether they are the new leader
+            sio.emit("status", {"isLeader": is_leader}, room=sid_str)
+        except Exception:
+            continue
+
+
+def get_current_leader(workflow_id):
+    """
+    Get the current leader for a workflow.
+    """
+    leader = redis_client.get(_leader_key(workflow_id))
+    return leader.decode("utf-8") if leader and isinstance(leader, bytes) else leader
+
+
+def broadcast_online_users(workflow_id):
+    """
+    Broadcast online users to the workflow room.
+    Each session is shown as a separate user (even if same person has multiple tabs).
+    """
+    sessions_json = redis_client.hgetall(_workflow_key(workflow_id))
+    users = []
+
+    for sid, session_info_json in sessions_json.items():
+        try:
+            session_info = json.loads(session_info_json)
+            # Each session appears as a separate "user" in the UI
+            users.append(
+                {
+                    "user_id": session_info["user_id"],
+                    "username": session_info["username"],
+                    "avatar": session_info.get("avatar"),
+                    "sid": session_info["sid"],
+                    "connected_at": session_info.get("connected_at"),
+                }
+            )
+        except Exception:
+            continue
+
+    # Sort by connection time to maintain consistent order
+    users.sort(key=lambda x: x.get("connected_at") or 0)
+
+    # Get current leader session
+    leader_sid = get_current_leader(workflow_id)
+
+    sio.emit("online_users", {"workflow_id": workflow_id, "users": users, "leader": leader_sid}, room=workflow_id)
+
+
+@sio.on("collaboration_event")
+def handle_collaboration_event(sid, data):
+    """
+    Handle general collaboration events, include:
+    1. mouse_move
+    2. vars_and_features_update
+    3. sync_request (ask leader to update graph)
+    4. app_state_update
+    5. mcp_server_update
+    6. workflow_update
+    7. comments_update
+    8. node_panel_presence
+
+    """
+    mapping = redis_client.get(_sid_key(sid))
+
+    if not mapping:
+        return {"msg": "unauthorized"}, 401
+
+    mapping_data = json.loads(mapping)
+    workflow_id = mapping_data["workflow_id"]
+    user_id = mapping_data["user_id"]
+    _refresh_session_state(workflow_id, sid)
+
+    event_type = data.get("type")
+    event_data = data.get("data")
+    timestamp = data.get("timestamp", int(time.time()))
+
+    if not event_type:
+        return {"msg": "invalid event type"}, 400
+
+    sio.emit(
+        "collaboration_update",
+        {"type": event_type, "userId": user_id, "data": event_data, "timestamp": timestamp},
+        room=workflow_id,
+        skip_sid=sid,
+    )
+
+    return {"msg": "event_broadcasted"}
+
+
+@sio.on("graph_event")
+def handle_graph_event(sid, data):
+    """
+    Handle graph events - simple broadcast relay.
+    """
+    mapping = redis_client.get(_sid_key(sid))
+
+    if not mapping:
+        return {"msg": "unauthorized"}, 401
+
+    mapping_data = json.loads(mapping)
+    workflow_id = mapping_data["workflow_id"]
+    _refresh_session_state(workflow_id, sid)
+
+    sio.emit("graph_update", data, room=workflow_id, skip_sid=sid)
+
+    return {"msg": "graph_update_broadcasted"}

api/controllers/console/app/workflow.py

@@ -9,6 +9,7 @@ from sqlalchemy.orm import Session
 from werkzeug.exceptions import Forbidden, InternalServerError, NotFound
 
 import services
+from configs import dify_config
 from controllers.console import api, console_ns
 from controllers.console.app.error import ConversationCompletedError, DraftWorkflowNotExist, DraftWorkflowNotSync
 from controllers.console.app.wraps import get_app_model
@@ -21,7 +22,9 @@ from core.file.models import File
 from core.helper.trace_id_helper import get_external_trace_id
 from core.workflow.graph_engine.manager import GraphEngineManager
 from extensions.ext_database import db
+from extensions.ext_redis import redis_client
 from factories import file_factory, variable_factory
+from fields.online_user_fields import online_user_list_fields
 from fields.workflow_fields import workflow_fields, workflow_pagination_fields
 from fields.workflow_run_fields import workflow_run_node_execution_fields
 from libs import helper
@@ -122,6 +125,7 @@ class DraftWorkflowApi(Resource):
                 .add_argument("hash", type=str, required=False, location="json")
                 .add_argument("environment_variables", type=list, required=True, location="json")
                 .add_argument("conversation_variables", type=list, required=False, location="json")
+                .add_argument("force_upload", type=bool, required=False, default=False, location="json")
             )
             args = parser.parse_args()
         elif "text/plain" in content_type:
@@ -139,6 +143,7 @@ class DraftWorkflowApi(Resource):
                     "hash": data.get("hash"),
                     "environment_variables": data.get("environment_variables"),
                     "conversation_variables": data.get("conversation_variables"),
+                    "force_upload": data.get("force_upload", False),
                 }
             except json.JSONDecodeError:
                 return {"message": "Invalid JSON data"}, 400
@@ -163,6 +168,7 @@ class DraftWorkflowApi(Resource):
                 account=current_user,
                 environment_variables=environment_variables,
                 conversation_variables=conversation_variables,
+                force_upload=args.get("force_upload", False),
             )
         except WorkflowHashNotEqualError:
             raise DraftWorkflowNotSync()
@@ -734,6 +740,46 @@ class ConvertToWorkflowApi(Resource):
         }
 
 
+@console_ns.route("/apps/<uuid:app_id>/workflows/draft/config")
+class WorkflowConfigApi(Resource):
+    """Resource for workflow configuration."""
+
+    @api.doc("get_workflow_config")
+    @api.doc(description="Get workflow configuration")
+    @api.doc(params={"app_id": "Application ID"})
+    @api.response(200, "Workflow configuration retrieved successfully")
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW])
+    def get(self, app_model: App):
+        return {
+            "parallel_depth_limit": dify_config.WORKFLOW_PARALLEL_DEPTH_LIMIT,
+        }
+
+
+@console_ns.route("/apps/<uuid:app_id>/workflows/draft/features")
+class WorkflowFeaturesApi(Resource):
+    """Update draft workflow features."""
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW])
+    def post(self, app_model: App):
+        current_user, _ = current_account_with_tenant()
+
+        parser = reqparse.RequestParser().add_argument("features", type=dict, required=True, location="json")
+        args = parser.parse_args()
+
+        features = args.get("features")
+
+        workflow_service = WorkflowService()
+        workflow_service.update_draft_workflow_features(app_model=app_model, features=features, account=current_user)
+
+        return {"result": "success"}
+
+
 @console_ns.route("/apps/<uuid:app_id>/workflows")
 class PublishedAllWorkflowApi(Resource):
     @api.doc("get_all_published_workflows")
@@ -915,3 +961,30 @@ class DraftWorkflowNodeLastRunApi(Resource):
         if node_exec is None:
             raise NotFound("last run not found")
         return node_exec
+
+
+@console_ns.route("/apps/workflows/online-users")
+class WorkflowOnlineUsersApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(online_user_list_fields)
+    def get(self):
+        parser = reqparse.RequestParser().add_argument("workflow_ids", type=str, required=True, location="args")
+        args = parser.parse_args()
+
+        workflow_ids = [workflow_id.strip() for workflow_id in args["workflow_ids"].split(",")]
+
+        results = []
+        for workflow_id in workflow_ids:
+            users_json = redis_client.hgetall(f"workflow_online_users:{workflow_id}")
+
+            users = []
+            for _, user_info_json in users_json.items():
+                try:
+                    users.append(json.loads(user_info_json))
+                except Exception:
+                    continue
+            results.append({"workflow_id": workflow_id, "users": users})
+
+        return {"data": results}

api/controllers/console/app/workflow_comment.py

@@ -0,0 +1,240 @@
+import logging
+
+from flask_restx import Resource, fields, marshal_with, reqparse
+
+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 fields.member_fields import account_with_role_fields
+from fields.workflow_comment_fields import (
+    workflow_comment_basic_fields,
+    workflow_comment_create_fields,
+    workflow_comment_detail_fields,
+    workflow_comment_reply_create_fields,
+    workflow_comment_reply_update_fields,
+    workflow_comment_resolve_fields,
+    workflow_comment_update_fields,
+)
+from libs.login import current_user, login_required
+from models import App
+from services.account_service import TenantService
+from services.workflow_comment_service import WorkflowCommentService
+
+logger = logging.getLogger(__name__)
+
+
+class WorkflowCommentListApi(Resource):
+    """API for listing and creating workflow comments."""
+
+    @login_required
+    @setup_required
+    @account_initialization_required
+    @get_app_model
+    @marshal_with(workflow_comment_basic_fields, envelope="data")
+    def get(self, app_model: App):
+        """Get all comments for a workflow."""
+        comments = WorkflowCommentService.get_comments(tenant_id=current_user.current_tenant_id, app_id=app_model.id)
+
+        return comments
+
+    @login_required
+    @setup_required
+    @account_initialization_required
+    @get_app_model
+    @marshal_with(workflow_comment_create_fields)
+    def post(self, app_model: App):
+        """Create a new workflow comment."""
+        parser = reqparse.RequestParser()
+        parser.add_argument("position_x", type=float, required=True, location="json")
+        parser.add_argument("position_y", type=float, required=True, location="json")
+        parser.add_argument("content", type=str, required=True, location="json")
+        parser.add_argument("mentioned_user_ids", type=list, location="json", default=[])
+        args = parser.parse_args()
+
+        result = WorkflowCommentService.create_comment(
+            tenant_id=current_user.current_tenant_id,
+            app_id=app_model.id,
+            created_by=current_user.id,
+            content=args.content,
+            position_x=args.position_x,
+            position_y=args.position_y,
+            mentioned_user_ids=args.mentioned_user_ids,
+        )
+
+        return result, 201
+
+
+class WorkflowCommentDetailApi(Resource):
+    """API for managing individual workflow comments."""
+
+    @login_required
+    @setup_required
+    @account_initialization_required
+    @get_app_model
+    @marshal_with(workflow_comment_detail_fields)
+    def get(self, app_model: App, comment_id: str):
+        """Get a specific workflow comment."""
+        comment = WorkflowCommentService.get_comment(
+            tenant_id=current_user.current_tenant_id, app_id=app_model.id, comment_id=comment_id
+        )
+
+        return comment
+
+    @login_required
+    @setup_required
+    @account_initialization_required
+    @get_app_model
+    @marshal_with(workflow_comment_update_fields)
+    def put(self, app_model: App, comment_id: str):
+        """Update a workflow comment."""
+        parser = reqparse.RequestParser()
+        parser.add_argument("content", type=str, required=True, location="json")
+        parser.add_argument("position_x", type=float, required=False, location="json")
+        parser.add_argument("position_y", type=float, required=False, location="json")
+        parser.add_argument("mentioned_user_ids", type=list, location="json", default=[])
+        args = parser.parse_args()
+
+        result = WorkflowCommentService.update_comment(
+            tenant_id=current_user.current_tenant_id,
+            app_id=app_model.id,
+            comment_id=comment_id,
+            user_id=current_user.id,
+            content=args.content,
+            position_x=args.position_x,
+            position_y=args.position_y,
+            mentioned_user_ids=args.mentioned_user_ids,
+        )
+
+        return result
+
+    @login_required
+    @setup_required
+    @account_initialization_required
+    @get_app_model
+    def delete(self, app_model: App, comment_id: str):
+        """Delete a workflow comment."""
+        WorkflowCommentService.delete_comment(
+            tenant_id=current_user.current_tenant_id,
+            app_id=app_model.id,
+            comment_id=comment_id,
+            user_id=current_user.id,
+        )
+
+        return {"result": "success"}, 204
+
+
+class WorkflowCommentResolveApi(Resource):
+    """API for resolving and reopening workflow comments."""
+
+    @login_required
+    @setup_required
+    @account_initialization_required
+    @get_app_model
+    @marshal_with(workflow_comment_resolve_fields)
+    def post(self, app_model: App, comment_id: str):
+        """Resolve a workflow comment."""
+        comment = WorkflowCommentService.resolve_comment(
+            tenant_id=current_user.current_tenant_id,
+            app_id=app_model.id,
+            comment_id=comment_id,
+            user_id=current_user.id,
+        )
+
+        return comment
+
+
+class WorkflowCommentReplyApi(Resource):
+    """API for managing comment replies."""
+
+    @login_required
+    @setup_required
+    @account_initialization_required
+    @get_app_model
+    @marshal_with(workflow_comment_reply_create_fields)
+    def post(self, app_model: App, comment_id: str):
+        """Add a reply to a workflow comment."""
+        # Validate comment access first
+        WorkflowCommentService.validate_comment_access(
+            comment_id=comment_id, tenant_id=current_user.current_tenant_id, app_id=app_model.id
+        )
+
+        parser = reqparse.RequestParser()
+        parser.add_argument("content", type=str, required=True, location="json")
+        parser.add_argument("mentioned_user_ids", type=list, location="json", default=[])
+        args = parser.parse_args()
+
+        result = WorkflowCommentService.create_reply(
+            comment_id=comment_id,
+            content=args.content,
+            created_by=current_user.id,
+            mentioned_user_ids=args.mentioned_user_ids,
+        )
+
+        return result, 201
+
+
+class WorkflowCommentReplyDetailApi(Resource):
+    """API for managing individual comment replies."""
+
+    @login_required
+    @setup_required
+    @account_initialization_required
+    @get_app_model
+    @marshal_with(workflow_comment_reply_update_fields)
+    def put(self, app_model: App, comment_id: str, reply_id: str):
+        """Update a comment reply."""
+        # Validate comment access first
+        WorkflowCommentService.validate_comment_access(
+            comment_id=comment_id, tenant_id=current_user.current_tenant_id, app_id=app_model.id
+        )
+
+        parser = reqparse.RequestParser()
+        parser.add_argument("content", type=str, required=True, location="json")
+        parser.add_argument("mentioned_user_ids", type=list, location="json", default=[])
+        args = parser.parse_args()
+
+        reply = WorkflowCommentService.update_reply(
+            reply_id=reply_id, user_id=current_user.id, content=args.content, mentioned_user_ids=args.mentioned_user_ids
+        )
+
+        return reply
+
+    @login_required
+    @setup_required
+    @account_initialization_required
+    @get_app_model
+    def delete(self, app_model: App, comment_id: str, reply_id: str):
+        """Delete a comment reply."""
+        # Validate comment access first
+        WorkflowCommentService.validate_comment_access(
+            comment_id=comment_id, tenant_id=current_user.current_tenant_id, app_id=app_model.id
+        )
+
+        WorkflowCommentService.delete_reply(reply_id=reply_id, user_id=current_user.id)
+
+        return {"result": "success"}, 204
+
+
+class WorkflowCommentMentionUsersApi(Resource):
+    """API for getting mentionable users for workflow comments."""
+
+    @login_required
+    @setup_required
+    @account_initialization_required
+    @get_app_model
+    @marshal_with({"users": fields.List(fields.Nested(account_with_role_fields))})
+    def get(self, app_model: App):
+        """Get all users in current tenant for mentions."""
+        members = TenantService.get_tenant_members(current_user.current_tenant)
+        return {"users": members}
+
+
+# Register API routes
+api.add_resource(WorkflowCommentListApi, "/apps/<uuid:app_id>/workflow/comments")
+api.add_resource(WorkflowCommentDetailApi, "/apps/<uuid:app_id>/workflow/comments/<string:comment_id>")
+api.add_resource(WorkflowCommentResolveApi, "/apps/<uuid:app_id>/workflow/comments/<string:comment_id>/resolve")
+api.add_resource(WorkflowCommentReplyApi, "/apps/<uuid:app_id>/workflow/comments/<string:comment_id>/replies")
+api.add_resource(
+    WorkflowCommentReplyDetailApi, "/apps/<uuid:app_id>/workflow/comments/<string:comment_id>/replies/<string:reply_id>"
+)
+api.add_resource(WorkflowCommentMentionUsersApi, "/apps/<uuid:app_id>/workflow/comments/mention-users")

api/controllers/console/app/workflow_draft_variable.py

@@ -19,8 +19,8 @@ from core.variables.segments import ArrayFileSegment, FileSegment, Segment
 from core.variables.types import SegmentType
 from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID, SYSTEM_VARIABLE_NODE_ID
 from extensions.ext_database import db
+from factories import variable_factory
 from factories.file_factory import build_from_mapping, build_from_mappings
-from factories.variable_factory import build_segment_with_type
 from libs.login import current_user, login_required
 from models import Account, App, AppMode
 from models.workflow import WorkflowDraftVariable
@@ -355,7 +355,7 @@ class VariableApi(Resource):
                 if len(raw_value) > 0 and not isinstance(raw_value[0], dict):
                     raise InvalidArgumentError(description=f"expected dict for files[0], got {type(raw_value)}")
                 raw_value = build_from_mappings(mappings=raw_value, tenant_id=app_model.tenant_id)
-            new_value = build_segment_with_type(variable.value_type, raw_value)
+            new_value = variable_factory.build_segment_with_type(variable.value_type, raw_value)
         draft_var_srv.update_variable(variable, name=new_name, value=new_value)
         db.session.commit()
         return variable
@@ -448,8 +448,35 @@ class ConversationVariableCollectionApi(Resource):
         db.session.commit()
         return _get_variable_list(app_model, CONVERSATION_VARIABLE_NODE_ID)
 
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @get_app_model(mode=AppMode.ADVANCED_CHAT)
+    def post(self, app_model: App):
+        # The role of the current user in the ta table must be admin, owner, or editor
+        if not current_user.is_editor:
+            raise Forbidden()
+
+        parser = reqparse.RequestParser()
+        parser.add_argument("conversation_variables", type=list, required=True, location="json")
+        args = parser.parse_args()
+
+        workflow_service = WorkflowService()
+
+        conversation_variables_list = args.get("conversation_variables") or []
+        conversation_variables = [
+            variable_factory.build_conversation_variable_from_mapping(obj) for obj in conversation_variables_list
+        ]
+
+        workflow_service.update_draft_workflow_conversation_variables(
+            app_model=app_model,
+            account=current_user,
+            conversation_variables=conversation_variables,
+        )
+
+        return {"result": "success"}
+
 
-@console_ns.route("/apps/<uuid:app_id>/workflows/draft/system-variables")
 class SystemVariableCollectionApi(Resource):
     @api.doc("get_system_variables")
     @api.doc(description="Get system variables for workflow")
@@ -499,3 +526,44 @@ class EnvironmentVariableCollectionApi(Resource):
             )
 
         return {"items": env_vars_list}
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW])
+    def post(self, app_model: App):
+        # The role of the current user in the ta table must be admin, owner, or editor
+        if not current_user.is_editor:
+            raise Forbidden()
+
+        parser = reqparse.RequestParser()
+        parser.add_argument("environment_variables", type=list, required=True, location="json")
+        args = parser.parse_args()
+
+        workflow_service = WorkflowService()
+
+        environment_variables_list = args.get("environment_variables") or []
+        environment_variables = [
+            variable_factory.build_environment_variable_from_mapping(obj) for obj in environment_variables_list
+        ]
+
+        workflow_service.update_draft_workflow_environment_variables(
+            app_model=app_model,
+            account=current_user,
+            environment_variables=environment_variables,
+        )
+
+        return {"result": "success"}
+
+
+api.add_resource(
+    WorkflowVariableCollectionApi,
+    "/apps/<uuid:app_id>/workflows/draft/variables",
+)
+api.add_resource(NodeVariableCollectionApi, "/apps/<uuid:app_id>/workflows/draft/nodes/<string:node_id>/variables")
+api.add_resource(VariableApi, "/apps/<uuid:app_id>/workflows/draft/variables/<uuid:variable_id>")
+api.add_resource(VariableResetApi, "/apps/<uuid:app_id>/workflows/draft/variables/<uuid:variable_id>/reset")
+
+api.add_resource(ConversationVariableCollectionApi, "/apps/<uuid:app_id>/workflows/draft/conversation-variables")
+api.add_resource(SystemVariableCollectionApi, "/apps/<uuid:app_id>/workflows/draft/system-variables")
+api.add_resource(EnvironmentVariableCollectionApi, "/apps/<uuid:app_id>/workflows/draft/environment-variables")

api/controllers/console/app/workflow_statistic.py

@@ -1,9 +1,10 @@
 from datetime import datetime
+from decimal import Decimal
 
 import pytz
+import sqlalchemy as sa
 from flask import jsonify
 from flask_restx import Resource, reqparse
-from sqlalchemy.orm import sessionmaker
 
 from controllers.console import api, console_ns
 from controllers.console.app.wraps import get_app_model
@@ -13,16 +14,10 @@ from libs.helper import DatetimeString
 from libs.login import current_account_with_tenant, login_required
 from models.enums import WorkflowRunTriggeredFrom
 from models.model import AppMode
-from repositories.factory import DifyAPIRepositoryFactory
 
 
 @console_ns.route("/apps/<uuid:app_id>/workflow/statistics/daily-conversations")
 class WorkflowDailyRunsStatistic(Resource):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        session_maker = sessionmaker(bind=db.engine, expire_on_commit=False)
-        self._workflow_run_repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker)
-
     @api.doc("get_workflow_daily_runs_statistic")
     @api.doc(description="Get workflow daily runs statistics")
     @api.doc(params={"app_id": "Application ID"})
@@ -42,44 +37,57 @@ class WorkflowDailyRunsStatistic(Resource):
         )
         args = parser.parse_args()
 
+        sql_query = """SELECT
+    DATE(DATE_TRUNC('day', created_at AT TIME ZONE 'UTC' AT TIME ZONE :tz )) AS date,
+    COUNT(id) AS runs
+FROM
+    workflow_runs
+WHERE
+    app_id = :app_id
+    AND triggered_from = :triggered_from"""
+        arg_dict = {
+            "tz": account.timezone,
+            "app_id": app_model.id,
+            "triggered_from": WorkflowRunTriggeredFrom.APP_RUN,
+        }
         assert account.timezone is not None
         timezone = pytz.timezone(account.timezone)
         utc_timezone = pytz.utc
 
-        start_date = None
-        end_date = None
-
         if args["start"]:
             start_datetime = datetime.strptime(args["start"], "%Y-%m-%d %H:%M")
             start_datetime = start_datetime.replace(second=0)
+
             start_datetime_timezone = timezone.localize(start_datetime)
-            start_date = start_datetime_timezone.astimezone(utc_timezone)
+            start_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
+
+            sql_query += " AND created_at >= :start"
+            arg_dict["start"] = start_datetime_utc
 
         if args["end"]:
             end_datetime = datetime.strptime(args["end"], "%Y-%m-%d %H:%M")
             end_datetime = end_datetime.replace(second=0)
-            end_datetime_timezone = timezone.localize(end_datetime)
-            end_date = end_datetime_timezone.astimezone(utc_timezone)
 
-        response_data = self._workflow_run_repo.get_daily_runs_statistics(
-            tenant_id=app_model.tenant_id,
-            app_id=app_model.id,
-            triggered_from=WorkflowRunTriggeredFrom.APP_RUN,
-            start_date=start_date,
-            end_date=end_date,
-            timezone=account.timezone,
-        )
+            end_datetime_timezone = timezone.localize(end_datetime)
+            end_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
+
+            sql_query += " AND created_at < :end"
+            arg_dict["end"] = end_datetime_utc
+
+        sql_query += " GROUP BY date ORDER BY date"
+
+        response_data = []
+
+        with db.engine.begin() as conn:
+            rs = conn.execute(sa.text(sql_query), arg_dict)
+            for i in rs:
+                response_data.append({"date": str(i.date), "runs": i.runs})
 
         return jsonify({"data": response_data})
 
 
 @console_ns.route("/apps/<uuid:app_id>/workflow/statistics/daily-terminals")
 class WorkflowDailyTerminalsStatistic(Resource):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        session_maker = sessionmaker(bind=db.engine, expire_on_commit=False)
-        self._workflow_run_repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker)
-
     @api.doc("get_workflow_daily_terminals_statistic")
     @api.doc(description="Get workflow daily terminals statistics")
     @api.doc(params={"app_id": "Application ID"})
@@ -99,44 +107,57 @@ class WorkflowDailyTerminalsStatistic(Resource):
         )
         args = parser.parse_args()
 
+        sql_query = """SELECT
+    DATE(DATE_TRUNC('day', created_at AT TIME ZONE 'UTC' AT TIME ZONE :tz )) AS date,
+    COUNT(DISTINCT workflow_runs.created_by) AS terminal_count
+FROM
+    workflow_runs
+WHERE
+    app_id = :app_id
+    AND triggered_from = :triggered_from"""
+        arg_dict = {
+            "tz": account.timezone,
+            "app_id": app_model.id,
+            "triggered_from": WorkflowRunTriggeredFrom.APP_RUN,
+        }
         assert account.timezone is not None
         timezone = pytz.timezone(account.timezone)
         utc_timezone = pytz.utc
 
-        start_date = None
-        end_date = None
-
         if args["start"]:
             start_datetime = datetime.strptime(args["start"], "%Y-%m-%d %H:%M")
             start_datetime = start_datetime.replace(second=0)
+
             start_datetime_timezone = timezone.localize(start_datetime)
-            start_date = start_datetime_timezone.astimezone(utc_timezone)
+            start_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
+
+            sql_query += " AND created_at >= :start"
+            arg_dict["start"] = start_datetime_utc
 
         if args["end"]:
             end_datetime = datetime.strptime(args["end"], "%Y-%m-%d %H:%M")
             end_datetime = end_datetime.replace(second=0)
-            end_datetime_timezone = timezone.localize(end_datetime)
-            end_date = end_datetime_timezone.astimezone(utc_timezone)
 
-        response_data = self._workflow_run_repo.get_daily_terminals_statistics(
-            tenant_id=app_model.tenant_id,
-            app_id=app_model.id,
-            triggered_from=WorkflowRunTriggeredFrom.APP_RUN,
-            start_date=start_date,
-            end_date=end_date,
-            timezone=account.timezone,
-        )
+            end_datetime_timezone = timezone.localize(end_datetime)
+            end_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
+
+            sql_query += " AND created_at < :end"
+            arg_dict["end"] = end_datetime_utc
+
+        sql_query += " GROUP BY date ORDER BY date"
+
+        response_data = []
+
+        with db.engine.begin() as conn:
+            rs = conn.execute(sa.text(sql_query), arg_dict)
+            for i in rs:
+                response_data.append({"date": str(i.date), "terminal_count": i.terminal_count})
 
         return jsonify({"data": response_data})
 
 
 @console_ns.route("/apps/<uuid:app_id>/workflow/statistics/token-costs")
 class WorkflowDailyTokenCostStatistic(Resource):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        session_maker = sessionmaker(bind=db.engine, expire_on_commit=False)
-        self._workflow_run_repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker)
-
     @api.doc("get_workflow_daily_token_cost_statistic")
     @api.doc(description="Get workflow daily token cost statistics")
     @api.doc(params={"app_id": "Application ID"})
@@ -156,44 +177,62 @@ class WorkflowDailyTokenCostStatistic(Resource):
         )
         args = parser.parse_args()
 
+        sql_query = """SELECT
+    DATE(DATE_TRUNC('day', created_at AT TIME ZONE 'UTC' AT TIME ZONE :tz )) AS date,
+    SUM(workflow_runs.total_tokens) AS token_count
+FROM
+    workflow_runs
+WHERE
+    app_id = :app_id
+    AND triggered_from = :triggered_from"""
+        arg_dict = {
+            "tz": account.timezone,
+            "app_id": app_model.id,
+            "triggered_from": WorkflowRunTriggeredFrom.APP_RUN,
+        }
         assert account.timezone is not None
         timezone = pytz.timezone(account.timezone)
         utc_timezone = pytz.utc
 
-        start_date = None
-        end_date = None
-
         if args["start"]:
             start_datetime = datetime.strptime(args["start"], "%Y-%m-%d %H:%M")
             start_datetime = start_datetime.replace(second=0)
+
             start_datetime_timezone = timezone.localize(start_datetime)
-            start_date = start_datetime_timezone.astimezone(utc_timezone)
+            start_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
+
+            sql_query += " AND created_at >= :start"
+            arg_dict["start"] = start_datetime_utc
 
         if args["end"]:
             end_datetime = datetime.strptime(args["end"], "%Y-%m-%d %H:%M")
             end_datetime = end_datetime.replace(second=0)
-            end_datetime_timezone = timezone.localize(end_datetime)
-            end_date = end_datetime_timezone.astimezone(utc_timezone)
 
-        response_data = self._workflow_run_repo.get_daily_token_cost_statistics(
-            tenant_id=app_model.tenant_id,
-            app_id=app_model.id,
-            triggered_from=WorkflowRunTriggeredFrom.APP_RUN,
-            start_date=start_date,
-            end_date=end_date,
-            timezone=account.timezone,
-        )
+            end_datetime_timezone = timezone.localize(end_datetime)
+            end_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
+
+            sql_query += " AND created_at < :end"
+            arg_dict["end"] = end_datetime_utc
+
+        sql_query += " GROUP BY date ORDER BY date"
+
+        response_data = []
+
+        with db.engine.begin() as conn:
+            rs = conn.execute(sa.text(sql_query), arg_dict)
+            for i in rs:
+                response_data.append(
+                    {
+                        "date": str(i.date),
+                        "token_count": i.token_count,
+                    }
+                )
 
         return jsonify({"data": response_data})
 
 
 @console_ns.route("/apps/<uuid:app_id>/workflow/statistics/average-app-interactions")
 class WorkflowAverageAppInteractionStatistic(Resource):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        session_maker = sessionmaker(bind=db.engine, expire_on_commit=False)
-        self._workflow_run_repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker)
-
     @api.doc("get_workflow_average_app_interaction_statistic")
     @api.doc(description="Get workflow average app interaction statistics")
     @api.doc(params={"app_id": "Application ID"})
@@ -213,32 +252,67 @@ class WorkflowAverageAppInteractionStatistic(Resource):
         )
         args = parser.parse_args()
 
+        sql_query = """SELECT
+    AVG(sub.interactions) AS interactions,
+    sub.date
+FROM
+    (
+        SELECT
+            DATE(DATE_TRUNC('day', c.created_at AT TIME ZONE 'UTC' AT TIME ZONE :tz )) AS date,
+            c.created_by,
+            COUNT(c.id) AS interactions
+        FROM
+            workflow_runs c
+        WHERE
+            c.app_id = :app_id
+            AND c.triggered_from = :triggered_from
+            {{start}}
+            {{end}}
+        GROUP BY
+            date, c.created_by
+    ) sub
+GROUP BY
+    sub.date"""
+        arg_dict = {
+            "tz": account.timezone,
+            "app_id": app_model.id,
+            "triggered_from": WorkflowRunTriggeredFrom.APP_RUN,
+        }
         assert account.timezone is not None
         timezone = pytz.timezone(account.timezone)
         utc_timezone = pytz.utc
 
-        start_date = None
-        end_date = None
-
         if args["start"]:
             start_datetime = datetime.strptime(args["start"], "%Y-%m-%d %H:%M")
             start_datetime = start_datetime.replace(second=0)
+
             start_datetime_timezone = timezone.localize(start_datetime)
-            start_date = start_datetime_timezone.astimezone(utc_timezone)
+            start_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
+
+            sql_query = sql_query.replace("{{start}}", " AND c.created_at >= :start")
+            arg_dict["start"] = start_datetime_utc
+        else:
+            sql_query = sql_query.replace("{{start}}", "")
 
         if args["end"]:
             end_datetime = datetime.strptime(args["end"], "%Y-%m-%d %H:%M")
             end_datetime = end_datetime.replace(second=0)
-            end_datetime_timezone = timezone.localize(end_datetime)
-            end_date = end_datetime_timezone.astimezone(utc_timezone)
 
-        response_data = self._workflow_run_repo.get_average_app_interaction_statistics(
-            tenant_id=app_model.tenant_id,
-            app_id=app_model.id,
-            triggered_from=WorkflowRunTriggeredFrom.APP_RUN,
-            start_date=start_date,
-            end_date=end_date,
-            timezone=account.timezone,
-        )
+            end_datetime_timezone = timezone.localize(end_datetime)
+            end_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
+
+            sql_query = sql_query.replace("{{end}}", " AND c.created_at < :end")
+            arg_dict["end"] = end_datetime_utc
+        else:
+            sql_query = sql_query.replace("{{end}}", "")
+
+        response_data = []
+
+        with db.engine.begin() as conn:
+            rs = conn.execute(sa.text(sql_query), arg_dict)
+            for i in rs:
+                response_data.append(
+                    {"date": str(i.date), "interactions": float(i.interactions.quantize(Decimal("0.01")))}
+                )
 
         return jsonify({"data": response_data})

api/controllers/console/auth/login.py

@@ -4,7 +4,7 @@ from flask_restx import Resource, reqparse
 
 import services
 from configs import dify_config
-from constants.languages import get_valid_language
+from constants.languages import languages
 from controllers.console import console_ns
 from controllers.console.auth.error import (
     AuthenticationFailedError,
@@ -29,7 +29,8 @@ from libs.token import (
     clear_access_token_from_cookie,
     clear_csrf_token_from_cookie,
     clear_refresh_token_from_cookie,
-    extract_refresh_token,
+    extract_access_token,
+    extract_csrf_token,
     set_access_token_to_cookie,
     set_csrf_token_to_cookie,
     set_refresh_token_to_cookie,
@@ -205,12 +206,10 @@ class EmailCodeLoginApi(Resource):
             .add_argument("email", type=str, required=True, location="json")
             .add_argument("code", type=str, required=True, location="json")
             .add_argument("token", type=str, required=True, location="json")
-            .add_argument("language", type=str, required=False, location="json")
         )
         args = parser.parse_args()
 
         user_email = args["email"]
-        language = args["language"]
 
         token_data = AccountService.get_email_code_login_data(args["token"])
         if token_data is None:
@@ -244,9 +243,7 @@ class EmailCodeLoginApi(Resource):
         if account is None:
             try:
                 account = AccountService.create_account_and_tenant(
-                    email=user_email,
-                    name=user_email,
-                    interface_language=get_valid_language(language),
+                    email=user_email, name=user_email, interface_language=languages[0]
                 )
             except WorkSpaceNotAllowedCreateError:
                 raise NotAllowedCreateWorkspace()
@@ -271,7 +268,7 @@ class EmailCodeLoginApi(Resource):
 class RefreshTokenApi(Resource):
     def post(self):
         # Get refresh token from cookie instead of request body
-        refresh_token = extract_refresh_token(request)
+        refresh_token = request.cookies.get("refresh_token")
 
         if not refresh_token:
             return {"result": "fail", "message": "No refresh token provided"}, 401
@@ -289,3 +286,13 @@ class RefreshTokenApi(Resource):
             return response
         except Exception as e:
             return {"result": "fail", "message": str(e)}, 401
+
+
+# this api helps frontend to check whether user is authenticated
+# TODO: remove in the future. frontend should redirect to login page by catching 401 status
+@console_ns.route("/login/status")
+class LoginStatus(Resource):
+    def get(self):
+        token = extract_access_token(request)
+        csrf_token = extract_csrf_token(request)
+        return {"logged_in": bool(token) and bool(csrf_token)}

api/controllers/console/datasets/rag_pipeline/datasource_content_preview.py

@@ -4,7 +4,7 @@ from flask_restx import (  # type: ignore
 )
 from werkzeug.exceptions import Forbidden
 
-from controllers.console import api, console_ns
+from controllers.console import console_ns
 from controllers.console.datasets.wraps import get_rag_pipeline
 from controllers.console.wraps import account_initialization_required, setup_required
 from libs.login import current_user, login_required
@@ -12,17 +12,9 @@ from models import Account
 from models.dataset import Pipeline
 from services.rag_pipeline.rag_pipeline import RagPipelineService
 
-parser = (
-    reqparse.RequestParser()
-    .add_argument("inputs", type=dict, required=True, nullable=False, location="json")
-    .add_argument("datasource_type", type=str, required=True, location="json")
-    .add_argument("credential_id", type=str, required=False, location="json")
-)
-
 
 @console_ns.route("/rag/pipelines/<uuid:pipeline_id>/workflows/published/datasource/nodes/<string:node_id>/preview")
 class DataSourceContentPreviewApi(Resource):
-    @api.expect(parser)
     @setup_required
     @login_required
     @account_initialization_required
@@ -34,6 +26,12 @@ class DataSourceContentPreviewApi(Resource):
         if not isinstance(current_user, Account):
             raise Forbidden()
 
+        parser = (
+            reqparse.RequestParser()
+            .add_argument("inputs", type=dict, required=True, nullable=False, location="json")
+            .add_argument("datasource_type", type=str, required=True, location="json")
+            .add_argument("credential_id", type=str, required=False, location="json")
+        )
         args = parser.parse_args()
 
         inputs = args.get("inputs")

api/controllers/console/explore/workflow.py

@@ -22,7 +22,7 @@ from core.errors.error import (
 from core.model_runtime.errors.invoke import InvokeError
 from core.workflow.graph_engine.manager import GraphEngineManager
 from libs import helper
-from libs.login import current_account_with_tenant
+from libs.login import current_user as current_user_
 from models.model import AppMode, InstalledApp
 from services.app_generate_service import AppGenerateService
 from services.errors.llm import InvokeRateLimitError
@@ -31,6 +31,8 @@ from .. import console_ns
 
 logger = logging.getLogger(__name__)
 
+current_user = current_user_._get_current_object()  # type: ignore
+
 
 @console_ns.route("/installed-apps/<uuid:installed_app_id>/workflows/run")
 class InstalledAppWorkflowRunApi(InstalledAppResource):
@@ -38,7 +40,6 @@ class InstalledAppWorkflowRunApi(InstalledAppResource):
         """
         Run workflow
         """
-        current_user, _ = current_account_with_tenant()
         app_model = installed_app.app
         if not app_model:
             raise NotWorkflowAppError()
@@ -52,6 +53,7 @@ class InstalledAppWorkflowRunApi(InstalledAppResource):
             .add_argument("files", type=list, required=False, location="json")
         )
         args = parser.parse_args()
+        assert current_user is not None
         try:
             response = AppGenerateService.generate(
                 app_model=app_model, user=current_user, args=args, invoke_from=InvokeFrom.EXPLORE, streaming=True
@@ -87,6 +89,7 @@ class InstalledAppWorkflowTaskStopApi(InstalledAppResource):
         app_mode = AppMode.value_of(app_model.mode)
         if app_mode != AppMode.WORKFLOW:
             raise NotWorkflowAppError()
+        assert current_user is not None
 
         # Stop using both mechanisms for backward compatibility
         # Legacy stop flag mechanism (without user check)

api/controllers/console/extension.py

@@ -66,7 +66,13 @@ class APIBasedExtensionAPI(Resource):
     @account_initialization_required
     @marshal_with(api_based_extension_fields)
     def post(self):
-        args = api.payload
+        parser = (
+            reqparse.RequestParser()
+            .add_argument("name", type=str, required=True, location="json")
+            .add_argument("api_endpoint", type=str, required=True, location="json")
+            .add_argument("api_key", type=str, required=True, location="json")
+        )
+        args = parser.parse_args()
         _, current_tenant_id = current_account_with_tenant()
 
         extension_data = APIBasedExtension(
@@ -119,7 +125,13 @@ class APIBasedExtensionDetailAPI(Resource):
 
         extension_data_from_db = APIBasedExtensionService.get_with_tenant_id(current_tenant_id, api_based_extension_id)
 
-        args = api.payload
+        parser = (
+            reqparse.RequestParser()
+            .add_argument("name", type=str, required=True, location="json")
+            .add_argument("api_endpoint", type=str, required=True, location="json")
+            .add_argument("api_key", type=str, required=True, location="json")
+        )
+        args = parser.parse_args()
 
         extension_data_from_db.name = args["name"]
         extension_data_from_db.api_endpoint = args["api_endpoint"]

api/controllers/console/files.py

@@ -39,7 +39,6 @@ class FileApi(Resource):
         return {
             "file_size_limit": dify_config.UPLOAD_FILE_SIZE_LIMIT,
             "batch_count_limit": dify_config.UPLOAD_FILE_BATCH_LIMIT,
-            "file_upload_limit": dify_config.BATCH_UPLOAD_LIMIT,
             "image_file_size_limit": dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT,
             "video_file_size_limit": dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT,
             "audio_file_size_limit": dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT,

api/controllers/console/setup.py

@@ -74,17 +74,12 @@ class SetupApi(Resource):
             .add_argument("email", type=email, required=True, location="json")
             .add_argument("name", type=StrLen(30), required=True, location="json")
             .add_argument("password", type=valid_password, required=True, location="json")
-            .add_argument("language", type=str, required=False, location="json")
         )
         args = parser.parse_args()
 
         # setup
         RegisterService.setup(
-            email=args["email"],
-            name=args["name"],
-            password=args["password"],
-            ip_address=extract_remote_ip(request),
-            language=args["language"],
+            email=args["email"], name=args["name"], password=args["password"], ip_address=extract_remote_ip(request)
         )
 
         return {"result": "success"}, 201

api/controllers/console/workspace/account.py

@@ -32,6 +32,7 @@ from controllers.console.wraps import (
     only_edition_cloud,
     setup_required,
 )
+from core.file import helpers as file_helpers
 from extensions.ext_database import db
 from fields.member_fields import account_fields
 from libs.datetime_utils import naive_utc_now
@@ -128,6 +129,17 @@ class AccountNameApi(Resource):
 
 @console_ns.route("/account/avatar")
 class AccountAvatarApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument("avatar", type=str, required=True, location="args")
+        args = parser.parse_args()
+
+        avatar_url = file_helpers.get_signed_file_url(args["avatar"])
+        return {"avatar_url": avatar_url}
+    
     @setup_required
     @login_required
     @account_initialization_required

api/controllers/console/workspace/tool_providers.py

@@ -6,7 +6,6 @@ from flask_restx import (
     Resource,
     reqparse,
 )
-from sqlalchemy.orm import Session
 from werkzeug.exceptions import Forbidden
 
 from configs import dify_config
@@ -16,21 +15,20 @@ from controllers.console.wraps import (
     enterprise_license_required,
     setup_required,
 )
-from core.entities.mcp_provider import MCPAuthentication, MCPConfiguration
 from core.mcp.auth.auth_flow import auth, handle_callback
-from core.mcp.error import MCPAuthError, MCPError, MCPRefreshTokenError
+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.impl.oauth import OAuthHandler
 from core.tools.entities.tool_entities import CredentialType
-from extensions.ext_database import db
 from libs.helper import StrLen, alphanumeric, uuid_value
 from libs.login import current_account_with_tenant, login_required
 from models.provider_ids import ToolProviderID
 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_manage_service import MCPToolManageService, OAuthDataType
+from services.tools.mcp_tools_manage_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
@@ -44,9 +42,7 @@ def is_valid_url(url: str) -> bool:
     try:
         parsed = urlparse(url)
         return all([parsed.scheme, parsed.netloc]) and parsed.scheme in ["http", "https"]
-    except (ValueError, TypeError):
-        # ValueError: Invalid URL format
-        # TypeError: url is not a string
+    except Exception:
         return False
 
 
@@ -890,34 +886,29 @@ class ToolProviderMCPApi(Resource):
             .add_argument("icon_type", type=str, required=True, nullable=False, location="json")
             .add_argument("icon_background", type=str, required=False, nullable=True, location="json", default="")
             .add_argument("server_identifier", type=str, required=True, nullable=False, location="json")
-            .add_argument("configuration", type=dict, required=False, nullable=True, location="json", default={})
+            .add_argument("timeout", type=float, required=False, nullable=False, location="json", default=30)
+            .add_argument("sse_read_timeout", type=float, required=False, nullable=False, location="json", default=300)
             .add_argument("headers", type=dict, required=False, nullable=True, location="json", default={})
-            .add_argument("authentication", type=dict, required=False, nullable=True, location="json", default={})
         )
         args = parser.parse_args()
         user, tenant_id = current_account_with_tenant()
-
-        # Parse and validate models
-        configuration = MCPConfiguration.model_validate(args["configuration"])
-        authentication = MCPAuthentication.model_validate(args["authentication"]) if args["authentication"] else None
-
-        # Create provider
-        with Session(db.engine) as session, session.begin():
-            service = MCPToolManageService(session=session)
-            result = service.create_provider(
+        if not is_valid_url(args["server_url"]):
+            raise ValueError("Server URL is not valid.")
+        return jsonable_encoder(
+            MCPToolManageService.create_mcp_provider(
                 tenant_id=tenant_id,
-                user_id=user.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"],
+                timeout=args["timeout"],
+                sse_read_timeout=args["sse_read_timeout"],
                 headers=args["headers"],
-                configuration=configuration,
-                authentication=authentication,
             )
-            return jsonable_encoder(result)
+        )
 
     @setup_required
     @login_required
@@ -932,43 +923,31 @@ class ToolProviderMCPApi(Resource):
             .add_argument("icon_background", type=str, required=False, nullable=True, location="json")
             .add_argument("provider_id", type=str, required=True, nullable=False, location="json")
             .add_argument("server_identifier", type=str, required=True, nullable=False, location="json")
-            .add_argument("configuration", type=dict, required=False, nullable=True, location="json", default={})
-            .add_argument("headers", type=dict, required=False, nullable=True, location="json", default={})
-            .add_argument("authentication", type=dict, required=False, nullable=True, location="json", default={})
+            .add_argument("timeout", type=float, required=False, nullable=True, location="json")
+            .add_argument("sse_read_timeout", type=float, required=False, nullable=True, location="json")
+            .add_argument("headers", type=dict, required=False, nullable=True, location="json")
         )
         args = parser.parse_args()
-        configuration = MCPConfiguration.model_validate(args["configuration"])
-        authentication = MCPAuthentication.model_validate(args["authentication"]) if args["authentication"] else None
+        if not is_valid_url(args["server_url"]):
+            if "[__HIDDEN__]" in args["server_url"]:
+                pass
+            else:
+                raise ValueError("Server URL is not valid.")
         _, current_tenant_id = current_account_with_tenant()
-
-        # Step 1: Validate server URL change if needed (includes URL format validation and network operation)
-        validation_result = None
-        with Session(db.engine) as session:
-            service = MCPToolManageService(session=session)
-            validation_result = service.validate_server_url_change(
-                tenant_id=current_tenant_id, provider_id=args["provider_id"], new_server_url=args["server_url"]
-            )
-
-            # No need to check for errors here, exceptions will be raised directly
-
-        # Step 2: Perform database update in a transaction
-        with Session(db.engine) as session, session.begin():
-            service = MCPToolManageService(session=session)
-            service.update_provider(
-                tenant_id=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"],
-                headers=args["headers"],
-                configuration=configuration,
-                authentication=authentication,
-                validation_result=validation_result,
-            )
-            return {"result": "success"}
+        MCPToolManageService.update_mcp_provider(
+            tenant_id=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"],
+            timeout=args.get("timeout"),
+            sse_read_timeout=args.get("sse_read_timeout"),
+            headers=args.get("headers"),
+        )
+        return {"result": "success"}
 
     @setup_required
     @login_required
@@ -979,11 +958,8 @@ class ToolProviderMCPApi(Resource):
         )
         args = parser.parse_args()
         _, current_tenant_id = current_account_with_tenant()
-
-        with Session(db.engine) as session, session.begin():
-            service = MCPToolManageService(session=session)
-            service.delete_provider(tenant_id=current_tenant_id, provider_id=args["provider_id"])
-            return {"result": "success"}
+        MCPToolManageService.delete_mcp_tool(tenant_id=current_tenant_id, provider_id=args["provider_id"])
+        return {"result": "success"}
 
 
 @console_ns.route("/workspaces/current/tool-provider/mcp/auth")
@@ -1000,53 +976,37 @@ class ToolMCPAuthApi(Resource):
         args = parser.parse_args()
         provider_id = args["provider_id"]
         _, tenant_id = current_account_with_tenant()
-
-        with Session(db.engine) as session, session.begin():
-            service = MCPToolManageService(session=session)
-            db_provider = service.get_provider(provider_id=provider_id, tenant_id=tenant_id)
-            if not db_provider:
-                raise ValueError("provider not found")
-
-            # Convert to entity
-            provider_entity = db_provider.to_entity()
-            server_url = provider_entity.decrypt_server_url()
-            headers = provider_entity.decrypt_authentication()
-
-        # Try to connect without active transaction
+        provider = MCPToolManageService.get_mcp_provider_by_provider_id(provider_id, tenant_id)
+        if not provider:
+            raise ValueError("provider not found")
         try:
-            # Use MCPClientWithAuthRetry to handle authentication automatically
             with MCPClient(
-                server_url=server_url,
-                headers=headers,
-                timeout=provider_entity.timeout,
-                sse_read_timeout=provider_entity.sse_read_timeout,
+                provider.decrypted_server_url,
+                provider_id,
+                tenant_id,
+                authed=False,
+                authorization_code=args["authorization_code"],
+                for_list=True,
+                headers=provider.decrypted_headers,
+                timeout=provider.timeout,
+                sse_read_timeout=provider.sse_read_timeout,
             ):
-                # Update credentials in new transaction
-                with Session(db.engine) as session, session.begin():
-                    service = MCPToolManageService(session=session)
-                    service.update_provider_credentials(
-                        provider_id=provider_id,
-                        tenant_id=tenant_id,
-                        credentials=provider_entity.credentials,
-                        authed=True,
-                    )
+                MCPToolManageService.update_mcp_provider_credentials(
+                    mcp_provider=provider,
+                    credentials=provider.decrypted_credentials,
+                    authed=True,
+                )
                 return {"result": "success"}
-        except MCPAuthError as e:
-            try:
-                auth_result = auth(provider_entity, args.get("authorization_code"))
-                with Session(db.engine) as session, session.begin():
-                    service = MCPToolManageService(session=session)
-                    response = service.execute_auth_actions(auth_result)
-                    return response
-            except MCPRefreshTokenError as e:
-                with Session(db.engine) as session, session.begin():
-                    service = MCPToolManageService(session=session)
-                    service.clear_provider_credentials(provider_id=provider_id, tenant_id=tenant_id)
-                raise ValueError(f"Failed to refresh token, please try to authorize again: {e}") from e
+
+        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:
-            with Session(db.engine) as session, session.begin():
-                service = MCPToolManageService(session=session)
-                service.clear_provider_credentials(provider_id=provider_id, tenant_id=tenant_id)
+            MCPToolManageService.update_mcp_provider_credentials(
+                mcp_provider=provider,
+                credentials={},
+                authed=False,
+            )
             raise ValueError(f"Failed to connect to MCP server: {e}") from e
 
 
@@ -1057,10 +1017,8 @@ class ToolMCPDetailApi(Resource):
     @account_initialization_required
     def get(self, provider_id):
         _, tenant_id = current_account_with_tenant()
-        with Session(db.engine) as session, session.begin():
-            service = MCPToolManageService(session=session)
-            provider = service.get_provider(provider_id=provider_id, tenant_id=tenant_id)
-            return jsonable_encoder(ToolTransformService.mcp_provider_to_user_provider(provider, for_list=True))
+        provider = MCPToolManageService.get_mcp_provider_by_provider_id(provider_id, tenant_id)
+        return jsonable_encoder(ToolTransformService.mcp_provider_to_user_provider(provider, for_list=True))
 
 
 @console_ns.route("/workspaces/current/tools/mcp")
@@ -1071,12 +1029,9 @@ class ToolMCPListAllApi(Resource):
     def get(self):
         _, tenant_id = current_account_with_tenant()
 
-        with Session(db.engine) as session, session.begin():
-            service = MCPToolManageService(session=session)
-            # Skip sensitive data decryption for list view to improve performance
-            tools = service.list_providers(tenant_id=tenant_id, include_sensitive=False)
+        tools = MCPToolManageService.retrieve_mcp_tools(tenant_id=tenant_id)
 
-            return [tool.to_dict() for tool in tools]
+        return [tool.to_dict() for tool in tools]
 
 
 @console_ns.route("/workspaces/current/tool-provider/mcp/update/<path:provider_id>")
@@ -1086,13 +1041,11 @@ class ToolMCPUpdateApi(Resource):
     @account_initialization_required
     def get(self, provider_id):
         _, tenant_id = current_account_with_tenant()
-        with Session(db.engine) as session, session.begin():
-            service = MCPToolManageService(session=session)
-            tools = service.list_provider_tools(
-                tenant_id=tenant_id,
-                provider_id=provider_id,
-            )
-            return jsonable_encoder(tools)
+        tools = MCPToolManageService.list_mcp_tool_from_remote_server(
+            tenant_id=tenant_id,
+            provider_id=provider_id,
+        )
+        return jsonable_encoder(tools)
 
 
 @console_ns.route("/mcp/oauth/callback")
@@ -1106,15 +1059,5 @@ class ToolMCPCallbackApi(Resource):
         args = parser.parse_args()
         state_key = args["state"]
         authorization_code = args["code"]
-
-        # Create service instance for handle_callback
-        with Session(db.engine) as session, session.begin():
-            mcp_service = MCPToolManageService(session=session)
-            # handle_callback now returns state data and tokens
-            state_data, tokens = handle_callback(state_key, authorization_code)
-            # Save tokens using the service layer
-            mcp_service.save_oauth_data(
-                state_data.provider_id, state_data.tenant_id, tokens.model_dump(), OAuthDataType.TOKENS
-            )
-
+        handle_callback(state_key, authorization_code)
         return redirect(f"{dify_config.CONSOLE_WEB_URL}/oauth-callback")

api/controllers/files/image_preview.py

@@ -14,25 +14,10 @@ from services.file_service import FileService
 
 @files_ns.route("/<uuid:file_id>/image-preview")
 class ImagePreviewApi(Resource):
-    """Deprecated endpoint for retrieving image previews."""
+    """
+    Deprecated
+    """
 
-    @files_ns.doc("get_image_preview")
-    @files_ns.doc(description="Retrieve a signed image preview for a file")
-    @files_ns.doc(
-        params={
-            "file_id": "ID of the file to preview",
-            "timestamp": "Unix timestamp used in the signature",
-            "nonce": "Random string used in the signature",
-            "sign": "HMAC signature verifying the request",
-        }
-    )
-    @files_ns.doc(
-        responses={
-            200: "Image preview returned successfully",
-            400: "Missing or invalid signature parameters",
-            415: "Unsupported file type",
-        }
-    )
     def get(self, file_id):
         file_id = str(file_id)
 
@@ -58,25 +43,6 @@ class ImagePreviewApi(Resource):
 
 @files_ns.route("/<uuid:file_id>/file-preview")
 class FilePreviewApi(Resource):
-    @files_ns.doc("get_file_preview")
-    @files_ns.doc(description="Download a file preview or attachment using signed parameters")
-    @files_ns.doc(
-        params={
-            "file_id": "ID of the file to preview",
-            "timestamp": "Unix timestamp used in the signature",
-            "nonce": "Random string used in the signature",
-            "sign": "HMAC signature verifying the request",
-            "as_attachment": "Whether to download the file as an attachment",
-        }
-    )
-    @files_ns.doc(
-        responses={
-            200: "File stream returned successfully",
-            400: "Missing or invalid signature parameters",
-            404: "File not found",
-            415: "Unsupported file type",
-        }
-    )
     def get(self, file_id):
         file_id = str(file_id)
 
@@ -135,20 +101,6 @@ class FilePreviewApi(Resource):
 
 @files_ns.route("/workspaces/<uuid:workspace_id>/webapp-logo")
 class WorkspaceWebappLogoApi(Resource):
-    @files_ns.doc("get_workspace_webapp_logo")
-    @files_ns.doc(description="Fetch the custom webapp logo for a workspace")
-    @files_ns.doc(
-        params={
-            "workspace_id": "Workspace identifier",
-        }
-    )
-    @files_ns.doc(
-        responses={
-            200: "Logo returned successfully",
-            404: "Webapp logo not configured",
-            415: "Unsupported file type",
-        }
-    )
     def get(self, workspace_id):
         workspace_id = str(workspace_id)
 

api/controllers/files/tool_files.py

@@ -13,26 +13,6 @@ from extensions.ext_database import db as global_db
 
 @files_ns.route("/tools/<uuid:file_id>.<string:extension>")
 class ToolFileApi(Resource):
-    @files_ns.doc("get_tool_file")
-    @files_ns.doc(description="Download a tool file by ID using signed parameters")
-    @files_ns.doc(
-        params={
-            "file_id": "Tool file identifier",
-            "extension": "Expected file extension",
-            "timestamp": "Unix timestamp used in the signature",
-            "nonce": "Random string used in the signature",
-            "sign": "HMAC signature verifying the request",
-            "as_attachment": "Whether to download the file as an attachment",
-        }
-    )
-    @files_ns.doc(
-        responses={
-            200: "Tool file stream returned successfully",
-            403: "Forbidden - invalid signature",
-            404: "File not found",
-            415: "Unsupported file type",
-        }
-    )
     def get(self, file_id, extension):
         file_id = str(file_id)
 

api/controllers/mcp/mcp.py

@@ -193,16 +193,15 @@ class MCPAppApi(Resource):
             except ValidationError as e:
                 raise MCPRequestError(mcp_types.INVALID_PARAMS, f"Invalid MCP request: {str(e)}")
 
-    def _retrieve_end_user(self, tenant_id: str, mcp_server_id: str) -> EndUser | None:
-        """Get end user - manages its own database session"""
-        with Session(db.engine, expire_on_commit=False) as session, session.begin():
-            return (
-                session.query(EndUser)
-                .where(EndUser.tenant_id == tenant_id)
-                .where(EndUser.session_id == mcp_server_id)
-                .where(EndUser.type == "mcp")
-                .first()
-            )
+    def _retrieve_end_user(self, tenant_id: str, mcp_server_id: str, session: Session) -> EndUser | None:
+        """Get end user from existing session - optimized query"""
+        return (
+            session.query(EndUser)
+            .where(EndUser.tenant_id == tenant_id)
+            .where(EndUser.session_id == mcp_server_id)
+            .where(EndUser.type == "mcp")
+            .first()
+        )
 
     def _create_end_user(
         self, client_name: str, tenant_id: str, app_id: str, mcp_server_id: str, session: Session
@@ -230,7 +229,7 @@ class MCPAppApi(Resource):
         request_id: Union[int, str],
     ) -> mcp_types.JSONRPCResponse | mcp_types.JSONRPCError | None:
         """Handle MCP request and return response"""
-        end_user = self._retrieve_end_user(mcp_server.tenant_id, mcp_server.id)
+        end_user = self._retrieve_end_user(mcp_server.tenant_id, mcp_server.id, session)
 
         if not end_user and isinstance(mcp_request.root, mcp_types.InitializeRequest):
             client_info = mcp_request.root.params.clientInfo

api/controllers/web/login.py

@@ -17,8 +17,8 @@ from libs.helper import email
 from libs.passport import PassportService
 from libs.password import valid_password
 from libs.token import (
-    clear_webapp_access_token_from_cookie,
-    extract_webapp_access_token,
+    clear_access_token_from_cookie,
+    extract_access_token,
 )
 from services.account_service import AccountService
 from services.app_service import AppService
@@ -81,7 +81,7 @@ class LoginStatusApi(Resource):
     )
     def get(self):
         app_code = request.args.get("app_code")
-        token = extract_webapp_access_token(request)
+        token = extract_access_token(request)
         if not app_code:
             return {
                 "logged_in": bool(token),
@@ -128,7 +128,7 @@ class LogoutApi(Resource):
         response = make_response({"result": "success"})
         # enterprise SSO sets same site to None in https deployment
         # so we need to logout by calling api
-        clear_webapp_access_token_from_cookie(response, samesite="None")
+        clear_access_token_from_cookie(response, samesite="None")
         return response
 
 

api/controllers/web/passport.py

@@ -12,8 +12,10 @@ from controllers.web import web_ns
 from controllers.web.error import WebAppAuthRequiredError
 from extensions.ext_database import db
 from libs.passport import PassportService
-from libs.token import extract_webapp_access_token
+from libs.token import extract_access_token
 from models.model import App, EndUser, Site
+from services.app_service import AppService
+from services.enterprise.enterprise_service import EnterpriseService
 from services.feature_service import FeatureService
 from services.webapp_auth_service import WebAppAuthService, WebAppAuthType
 
@@ -35,18 +37,23 @@ class PassportResource(Resource):
         system_features = FeatureService.get_system_features()
         app_code = request.headers.get(HEADER_NAME_APP_CODE)
         user_id = request.args.get("user_id")
-        access_token = extract_webapp_access_token(request)
+        access_token = extract_access_token(request)
+
         if app_code is None:
             raise Unauthorized("X-App-Code header is missing.")
+        app_id = AppService.get_app_id_by_code(app_code)
+        # exchange token for enterprise logined web user
+        enterprise_user_decoded = decode_enterprise_webapp_user_id(access_token)
+        if enterprise_user_decoded:
+            # a web user has already logged in, exchange a token for this app without redirecting to the login page
+            return exchange_token_for_existing_web_user(
+                app_code=app_code, enterprise_user_decoded=enterprise_user_decoded
+            )
+
         if system_features.webapp_auth.enabled:
-            enterprise_user_decoded = decode_enterprise_webapp_user_id(access_token)
-            app_auth_type = WebAppAuthService.get_app_auth_type(app_code=app_code)
-            if app_auth_type != WebAppAuthType.PUBLIC:
-                if not enterprise_user_decoded:
-                    raise WebAppAuthRequiredError()
-                return exchange_token_for_existing_web_user(
-                    app_code=app_code, enterprise_user_decoded=enterprise_user_decoded, auth_type=app_auth_type
-                )
+            app_settings = EnterpriseService.WebAppAuth.get_app_access_mode_by_id(app_id=app_id)
+            if not app_settings or not app_settings.access_mode == "public":
+                raise WebAppAuthRequiredError()
 
         # get site from db and check if it is normal
         site = db.session.scalar(select(Site).where(Site.code == app_code, Site.status == "normal"))
@@ -117,7 +124,7 @@ def decode_enterprise_webapp_user_id(jwt_token: str | None):
     return decoded
 
 
-def exchange_token_for_existing_web_user(app_code: str, enterprise_user_decoded: dict, auth_type: WebAppAuthType):
+def exchange_token_for_existing_web_user(app_code: str, enterprise_user_decoded: dict):
     """
     Exchange a token for an existing web user session.
     """
@@ -138,11 +145,13 @@ def exchange_token_for_existing_web_user(app_code: str, enterprise_user_decoded:
     if not app_model or app_model.status != "normal" or not app_model.enable_site:
         raise NotFound()
 
-    if auth_type == WebAppAuthType.PUBLIC:
+    app_auth_type = WebAppAuthService.get_app_auth_type(app_code=app_code)
+
+    if app_auth_type == WebAppAuthType.PUBLIC:
         return _exchange_for_public_app_token(app_model, site, enterprise_user_decoded)
-    elif auth_type == WebAppAuthType.EXTERNAL and user_auth_type != "external":
+    elif app_auth_type == WebAppAuthType.EXTERNAL and user_auth_type != "external":
         raise WebAppAuthRequiredError("Please login as external user.")
-    elif auth_type == WebAppAuthType.INTERNAL and user_auth_type != "internal":
+    elif app_auth_type == WebAppAuthType.INTERNAL and user_auth_type != "internal":
         raise WebAppAuthRequiredError("Please login as internal user.")
 
     end_user = None

api/core/app/apps/agent_chat/app_generator.py

@@ -211,7 +211,8 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
             user=user,
             stream=streaming,
         )
-        return AgentChatAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from)
+        # FIXME: Type hinting issue here, ignore it for now, will fix it later
+        return AgentChatAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from)  # type: ignore
 
     def _generate_worker(
         self,

api/core/app/apps/common/workflow_response_converter.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 from datetime import datetime
 from typing import Any, NewType, Union
 
-from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom, WorkflowAppGenerateEntity
+from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity
 from core.app.entities.queue_entities import (
     QueueAgentLogEvent,
     QueueIterationCompletedEvent,
@@ -51,7 +51,7 @@ from core.workflow.workflow_entry import WorkflowEntry
 from core.workflow.workflow_type_encoder import WorkflowRuntimeTypeConverter
 from libs.datetime_utils import naive_utc_now
 from models import Account, EndUser
-from services.variable_truncator import BaseTruncator, DummyVariableTruncator, VariableTruncator
+from services.variable_truncator import VariableTruncator
 
 NodeExecutionId = NewType("NodeExecutionId", str)
 
@@ -70,8 +70,6 @@ class _NodeSnapshot:
 
 
 class WorkflowResponseConverter:
-    _truncator: BaseTruncator
-
     def __init__(
         self,
         *,
@@ -83,13 +81,7 @@ class WorkflowResponseConverter:
         self._user = user
         self._system_variables = system_variables
         self._workflow_inputs = self._prepare_workflow_inputs()
-
-        # Disable truncation for SERVICE_API calls to keep backward compatibility.
-        if application_generate_entity.invoke_from == InvokeFrom.SERVICE_API:
-            self._truncator = DummyVariableTruncator()
-        else:
-            self._truncator = VariableTruncator.default()
-
+        self._truncator = VariableTruncator.default()
         self._node_snapshots: dict[NodeExecutionId, _NodeSnapshot] = {}
         self._workflow_execution_id: str | None = None
         self._workflow_started_at: datetime | None = None

api/core/app/apps/pipeline/pipeline_generator.py

@@ -255,7 +255,7 @@ class PipelineGenerator(BaseAppGenerator):
             json_text = json.dumps(text)
             upload_file = FileService(db.engine).upload_text(json_text, name, user.id, dataset.tenant_id)
             features = FeatureService.get_features(dataset.tenant_id)
-            if features.billing.enabled and features.billing.subscription.plan == "sandbox":
+            if features.billing.subscription.plan == "sandbox":
                 tenant_pipeline_task_key = f"tenant_pipeline_task:{dataset.tenant_id}"
                 tenant_self_pipeline_task_queue = f"tenant_self_pipeline_task_queue:{dataset.tenant_id}"
 

api/core/app/apps/workflow/generate_response_converter.py

@@ -89,7 +89,7 @@ class WorkflowAppGenerateResponseConverter(AppGenerateResponseConverter):
                 data = cls._error_to_stream_response(sub_stream_response.err)
                 response_chunk.update(data)
             elif isinstance(sub_stream_response, NodeStartStreamResponse | NodeFinishStreamResponse):
-                response_chunk.update(sub_stream_response.to_ignore_detail_dict())
+                response_chunk.update(sub_stream_response.to_ignore_detail_dict())  # ty: ignore [unresolved-attribute]
             else:
                 response_chunk.update(sub_stream_response.model_dump(mode="json"))
             yield response_chunk

api/core/app/features/rate_limiting/rate_limit.py

@@ -98,7 +98,7 @@ class RateLimit:
         else:
             return RateLimitGenerator(
                 rate_limit=self,
-                generator=generator,
+                generator=generator,  # ty: ignore [invalid-argument-type]
                 request_id=request_id,
             )
 

api/core/app/task_pipeline/based_generate_task_pipeline.py

@@ -49,7 +49,7 @@ class BasedGenerateTaskPipeline:
         if isinstance(e, InvokeAuthorizationError):
             err = InvokeAuthorizationError("Incorrect API key provided")
         elif isinstance(e, InvokeError | ValueError):
-            err = e
+            err = e  # ty: ignore [invalid-assignment]
         else:
             description = getattr(e, "description", None)
             err = Exception(description if description is not None else str(e))

api/core/entities/mcp_provider.py

@@ -1,329 +0,0 @@
-import json
-from datetime import datetime
-from enum import StrEnum
-from typing import TYPE_CHECKING, Any
-from urllib.parse import urlparse
-
-from pydantic import BaseModel
-
-from configs import dify_config
-from core.entities.provider_entities import BasicProviderConfig
-from core.file import helpers as file_helpers
-from core.helper import encrypter
-from core.helper.provider_cache import NoOpProviderCredentialCache
-from core.mcp.types import OAuthClientInformation, OAuthClientMetadata, OAuthTokens
-from core.tools.entities.common_entities import I18nObject
-from core.tools.entities.tool_entities import ToolProviderType
-
-if TYPE_CHECKING:
-    from models.tools import MCPToolProvider
-
-# Constants
-CLIENT_NAME = "Dify"
-CLIENT_URI = "https://github.com/langgenius/dify"
-DEFAULT_TOKEN_TYPE = "Bearer"
-DEFAULT_EXPIRES_IN = 3600
-MASK_CHAR = "*"
-MIN_UNMASK_LENGTH = 6
-
-
-class MCPSupportGrantType(StrEnum):
-    """The supported grant types for MCP"""
-
-    AUTHORIZATION_CODE = "authorization_code"
-    CLIENT_CREDENTIALS = "client_credentials"
-    REFRESH_TOKEN = "refresh_token"
-
-
-class MCPAuthentication(BaseModel):
-    client_id: str
-    client_secret: str | None = None
-
-
-class MCPConfiguration(BaseModel):
-    timeout: float = 30
-    sse_read_timeout: float = 300
-
-
-class MCPProviderEntity(BaseModel):
-    """MCP Provider domain entity for business logic operations"""
-
-    # Basic identification
-    id: str
-    provider_id: str  # server_identifier
-    name: str
-    tenant_id: str
-    user_id: str
-
-    # Server connection info
-    server_url: str  # encrypted URL
-    headers: dict[str, str]  # encrypted headers
-    timeout: float
-    sse_read_timeout: float
-
-    # Authentication related
-    authed: bool
-    credentials: dict[str, Any]  # encrypted credentials
-    code_verifier: str | None = None  # for OAuth
-
-    # Tools and display info
-    tools: list[dict[str, Any]]  # parsed tools list
-    icon: str | dict[str, str]  # parsed icon
-
-    # Timestamps
-    created_at: datetime
-    updated_at: datetime
-
-    @classmethod
-    def from_db_model(cls, db_provider: "MCPToolProvider") -> "MCPProviderEntity":
-        """Create entity from database model with decryption"""
-
-        return cls(
-            id=db_provider.id,
-            provider_id=db_provider.server_identifier,
-            name=db_provider.name,
-            tenant_id=db_provider.tenant_id,
-            user_id=db_provider.user_id,
-            server_url=db_provider.server_url,
-            headers=db_provider.headers,
-            timeout=db_provider.timeout,
-            sse_read_timeout=db_provider.sse_read_timeout,
-            authed=db_provider.authed,
-            credentials=db_provider.credentials,
-            tools=db_provider.tool_dict,
-            icon=db_provider.icon or "",
-            created_at=db_provider.created_at,
-            updated_at=db_provider.updated_at,
-        )
-
-    @property
-    def redirect_url(self) -> str:
-        """OAuth redirect URL"""
-        return dify_config.CONSOLE_API_URL + "/console/api/mcp/oauth/callback"
-
-    @property
-    def client_metadata(self) -> OAuthClientMetadata:
-        """Metadata about this OAuth client."""
-        # Get grant type from credentials
-        credentials = self.decrypt_credentials()
-
-        # Try to get grant_type from different locations
-        grant_type = credentials.get("grant_type", MCPSupportGrantType.AUTHORIZATION_CODE)
-
-        # For nested structure, check if client_information has grant_types
-        if "client_information" in credentials and isinstance(credentials["client_information"], dict):
-            client_info = credentials["client_information"]
-            # If grant_types is specified in client_information, use it to determine grant_type
-            if "grant_types" in client_info and isinstance(client_info["grant_types"], list):
-                if "client_credentials" in client_info["grant_types"]:
-                    grant_type = MCPSupportGrantType.CLIENT_CREDENTIALS
-                elif "authorization_code" in client_info["grant_types"]:
-                    grant_type = MCPSupportGrantType.AUTHORIZATION_CODE
-
-        # Configure based on grant type
-        is_client_credentials = grant_type == MCPSupportGrantType.CLIENT_CREDENTIALS
-
-        grant_types = ["refresh_token"]
-        grant_types.append("client_credentials" if is_client_credentials else "authorization_code")
-
-        response_types = [] if is_client_credentials else ["code"]
-        redirect_uris = [] if is_client_credentials else [self.redirect_url]
-
-        return OAuthClientMetadata(
-            redirect_uris=redirect_uris,
-            token_endpoint_auth_method="none",
-            grant_types=grant_types,
-            response_types=response_types,
-            client_name=CLIENT_NAME,
-            client_uri=CLIENT_URI,
-        )
-
-    @property
-    def provider_icon(self) -> dict[str, str] | str:
-        """Get provider icon, handling both dict and string formats"""
-        if isinstance(self.icon, dict):
-            return self.icon
-        try:
-            return json.loads(self.icon)
-        except (json.JSONDecodeError, TypeError):
-            # If not JSON, assume it's a file path
-            return file_helpers.get_signed_file_url(self.icon)
-
-    def to_api_response(self, user_name: str | None = None, include_sensitive: bool = True) -> dict[str, Any]:
-        """Convert to API response format
-
-        Args:
-            user_name: User name to display
-            include_sensitive: If False, skip expensive decryption operations (for list view optimization)
-        """
-        response = {
-            "id": self.id,
-            "author": user_name or "Anonymous",
-            "name": self.name,
-            "icon": self.provider_icon,
-            "type": ToolProviderType.MCP.value,
-            "is_team_authorization": self.authed,
-            "server_url": self.masked_server_url(),
-            "server_identifier": self.provider_id,
-            "updated_at": int(self.updated_at.timestamp()),
-            "label": I18nObject(en_US=self.name, zh_Hans=self.name).to_dict(),
-            "description": I18nObject(en_US="", zh_Hans="").to_dict(),
-        }
-
-        # Add configuration
-        response["configuration"] = {
-            "timeout": str(self.timeout),
-            "sse_read_timeout": str(self.sse_read_timeout),
-        }
-
-        # Skip expensive operations when sensitive data is not needed (e.g., list view)
-        if not include_sensitive:
-            response["masked_headers"] = {}
-            response["is_dynamic_registration"] = True
-        else:
-            # Add masked headers
-            response["masked_headers"] = self.masked_headers()
-
-            # Add authentication info if available
-            masked_creds = self.masked_credentials()
-            if masked_creds:
-                response["authentication"] = masked_creds
-            response["is_dynamic_registration"] = self.credentials.get("client_information", {}).get(
-                "is_dynamic_registration", True
-            )
-
-        return response
-
-    def retrieve_client_information(self) -> OAuthClientInformation | None:
-        """OAuth client information if available"""
-        credentials = self.decrypt_credentials()
-        if not credentials:
-            return None
-
-        # Check if we have nested client_information structure
-        if "client_information" not in credentials:
-            return None
-        client_info_data = credentials["client_information"]
-        if isinstance(client_info_data, dict):
-            if "encrypted_client_secret" in client_info_data:
-                client_info_data["client_secret"] = encrypter.decrypt_token(
-                    self.tenant_id, client_info_data["encrypted_client_secret"]
-                )
-            return OAuthClientInformation.model_validate(client_info_data)
-        return None
-
-    def retrieve_tokens(self) -> OAuthTokens | None:
-        """OAuth tokens if available"""
-        if not self.credentials:
-            return None
-        credentials = self.decrypt_credentials()
-        return OAuthTokens(
-            access_token=credentials.get("access_token", ""),
-            token_type=credentials.get("token_type", DEFAULT_TOKEN_TYPE),
-            expires_in=int(credentials.get("expires_in", str(DEFAULT_EXPIRES_IN)) or DEFAULT_EXPIRES_IN),
-            refresh_token=credentials.get("refresh_token", ""),
-        )
-
-    def masked_server_url(self) -> str:
-        """Masked server URL for display"""
-        parsed = urlparse(self.decrypt_server_url())
-        if parsed.path and parsed.path != "/":
-            masked = parsed._replace(path="/******")
-            return masked.geturl()
-        return parsed.geturl()
-
-    def _mask_value(self, value: str) -> str:
-        """Mask a sensitive value for display"""
-        if len(value) > MIN_UNMASK_LENGTH:
-            return value[:2] + MASK_CHAR * (len(value) - 4) + value[-2:]
-        else:
-            return MASK_CHAR * len(value)
-
-    def masked_headers(self) -> dict[str, str]:
-        """Masked headers for display"""
-        return {key: self._mask_value(value) for key, value in self.decrypt_headers().items()}
-
-    def masked_credentials(self) -> dict[str, str]:
-        """Masked credentials for display"""
-        credentials = self.decrypt_credentials()
-        if not credentials:
-            return {}
-
-        masked = {}
-
-        if "client_information" not in credentials or not isinstance(credentials["client_information"], dict):
-            return {}
-        client_info = credentials["client_information"]
-        # Mask sensitive fields from nested structure
-        if client_info.get("client_id"):
-            masked["client_id"] = self._mask_value(client_info["client_id"])
-        if client_info.get("encrypted_client_secret"):
-            masked["client_secret"] = self._mask_value(
-                encrypter.decrypt_token(self.tenant_id, client_info["encrypted_client_secret"])
-            )
-        if client_info.get("client_secret"):
-            masked["client_secret"] = self._mask_value(client_info["client_secret"])
-        return masked
-
-    def decrypt_server_url(self) -> str:
-        """Decrypt server URL"""
-        return encrypter.decrypt_token(self.tenant_id, self.server_url)
-
-    def _decrypt_dict(self, data: dict[str, Any]) -> dict[str, Any]:
-        """Generic method to decrypt dictionary fields"""
-        from core.tools.utils.encryption import create_provider_encrypter
-
-        if not data:
-            return {}
-
-        # Only decrypt fields that are actually encrypted
-        # For nested structures, client_information is not encrypted as a whole
-        encrypted_fields = []
-        for key, value in data.items():
-            # Skip nested objects - they are not encrypted
-            if isinstance(value, dict):
-                continue
-            # Only process string values that might be encrypted
-            if isinstance(value, str) and value:
-                encrypted_fields.append(key)
-
-        if not encrypted_fields:
-            return data
-
-        # Create dynamic config only for encrypted fields
-        config = [BasicProviderConfig(type=BasicProviderConfig.Type.SECRET_INPUT, name=key) for key in encrypted_fields]
-
-        encrypter_instance, _ = create_provider_encrypter(
-            tenant_id=self.tenant_id,
-            config=config,
-            cache=NoOpProviderCredentialCache(),
-        )
-
-        # Decrypt only the encrypted fields
-        decrypted_data = encrypter_instance.decrypt({k: data[k] for k in encrypted_fields})
-
-        # Merge decrypted data with original data (preserving non-encrypted fields)
-        result = data.copy()
-        result.update(decrypted_data)
-
-        return result
-
-    def decrypt_headers(self) -> dict[str, Any]:
-        """Decrypt headers"""
-        return self._decrypt_dict(self.headers)
-
-    def decrypt_credentials(self) -> dict[str, Any]:
-        """Decrypt credentials"""
-        return self._decrypt_dict(self.credentials)
-
-    def decrypt_authentication(self) -> dict[str, Any]:
-        """Decrypt authentication"""
-        # Option 1: if headers is provided, use it and don't need to get token
-        headers = self.decrypt_headers()
-
-        # Option 2: Add OAuth token if authed and no headers provided
-        if not self.headers and self.authed:
-            token = self.retrieve_tokens()
-            if token:
-                headers["Authorization"] = f"{token.token_type.capitalize()} {token.access_token}"
-        return headers

api/core/entities/provider_configuration.py

@@ -1868,7 +1868,7 @@ class ProviderConfigurations(BaseModel):
         if "/" not in key:
             key = str(ModelProviderID(key))
 
-        return self.configurations.get(key, default)
+        return self.configurations.get(key, default)  # type: ignore
 
 
 class ProviderModelBundle(BaseModel):

api/core/helper/module_import_helper.py

@@ -20,7 +20,7 @@ def import_module_from_source(*, module_name: str, py_file_path: AnyStr, use_laz
         else:
             # Refer to: https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly
             # FIXME: mypy does not support the type of spec.loader
-            spec = importlib.util.spec_from_file_location(module_name, py_file_path)  # type: ignore[assignment]
+            spec = importlib.util.spec_from_file_location(module_name, py_file_path)  # type: ignore
             if not spec or not spec.loader:
                 raise Exception(f"Failed to load module {module_name} from {py_file_path!r}")
             if use_lazy_loader:

api/core/indexing_runner.py

@@ -415,6 +415,7 @@ class IndexingRunner:
             document_id=dataset_document.id,
             after_indexing_status="splitting",
             extra_update_params={
+                DatasetDocument.word_count: sum(len(text_doc.page_content) for text_doc in text_docs),
                 DatasetDocument.parsing_completed_at: naive_utc_now(),
             },
         )
@@ -754,7 +755,6 @@ class IndexingRunner:
             extra_update_params={
                 DatasetDocument.cleaning_completed_at: cur_time,
                 DatasetDocument.splitting_completed_at: cur_time,
-                DatasetDocument.word_count: sum(len(doc.page_content) for doc in documents),
             },
         )
 

api/core/mcp/auth/auth_flow.py

@@ -6,15 +6,11 @@ import secrets
 import urllib.parse
 from urllib.parse import urljoin, urlparse
 
-from httpx import ConnectError, HTTPStatusError, RequestError
-from pydantic import ValidationError
+import httpx
+from pydantic import BaseModel, ValidationError
 
-from core.entities.mcp_provider import MCPProviderEntity, MCPSupportGrantType
-from core.helper import ssrf_proxy
-from core.mcp.entities import AuthAction, AuthActionType, AuthResult, OAuthCallbackState
-from core.mcp.error import MCPRefreshTokenError
+from core.mcp.auth.auth_provider import OAuthClientProvider
 from core.mcp.types import (
-    LATEST_PROTOCOL_VERSION,
     OAuthClientInformation,
     OAuthClientInformationFull,
     OAuthClientMetadata,
@@ -23,10 +19,21 @@ from core.mcp.types import (
 )
 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")
@@ -73,13 +80,8 @@ def _retrieve_redis_state(state_key: str) -> OAuthCallbackState:
         raise ValueError(f"Invalid state parameter: {str(e)}")
 
 
-def handle_callback(state_key: str, authorization_code: str) -> tuple[OAuthCallbackState, OAuthTokens]:
-    """
-    Handle the callback from the OAuth provider.
-
-    Returns:
-        A tuple of (callback_state, tokens) that can be used by the caller to save data.
-    """
+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)
 
@@ -91,32 +93,30 @@ def handle_callback(state_key: str, authorization_code: str) -> tuple[OAuthCallb
         full_state_data.code_verifier,
         full_state_data.redirect_uri,
     )
-
-    return full_state_data, tokens
+    provider = OAuthClientProvider(full_state_data.provider_id, full_state_data.tenant_id, for_list=True)
+    provider.save_tokens(tokens)
+    return full_state_data
 
 
 def check_support_resource_discovery(server_url: str) -> tuple[bool, str]:
     """Check if the server supports OAuth 2.0 Resource Discovery."""
-    b_scheme, b_netloc, _, _, b_query, b_fragment = urlparse(server_url, "", True)
-    url_for_resource_discovery = f"{b_scheme}://{b_netloc}/.well-known/oauth-protected-resource"
+    b_scheme, b_netloc, b_path, _, b_query, b_fragment = urlparse(server_url, "", True)
+    url_for_resource_discovery = f"{b_scheme}://{b_netloc}/.well-known/oauth-protected-resource{b_path}"
     if b_query:
         url_for_resource_discovery += f"?{b_query}"
     if b_fragment:
         url_for_resource_discovery += f"#{b_fragment}"
     try:
         headers = {"MCP-Protocol-Version": LATEST_PROTOCOL_VERSION, "User-Agent": "Dify"}
-        response = ssrf_proxy.get(url_for_resource_discovery, headers=headers)
+        response = httpx.get(url_for_resource_discovery, headers=headers)
         if 200 <= response.status_code < 300:
             body = response.json()
-            # Support both singular and plural forms
-            if body.get("authorization_servers"):
-                return True, body["authorization_servers"][0]
-            elif body.get("authorization_server_url"):
+            if "authorization_server_url" in body:
                 return True, body["authorization_server_url"][0]
             else:
                 return False, ""
         return False, ""
-    except RequestError:
+    except httpx.RequestError:
         # Not support resource discovery, fall back to well-known OAuth metadata
         return False, ""
 
@@ -126,37 +126,27 @@ def discover_oauth_metadata(server_url: str, protocol_version: str | None = None
     # First check if the server supports OAuth 2.0 Resource Discovery
     support_resource_discovery, oauth_discovery_url = check_support_resource_discovery(server_url)
     if support_resource_discovery:
-        # The oauth_discovery_url is the authorization server base URL
-        # Try OpenID Connect discovery first (more common), then OAuth 2.0
-        urls_to_try = [
-            urljoin(oauth_discovery_url + "/", ".well-known/oauth-authorization-server"),
-            urljoin(oauth_discovery_url + "/", ".well-known/openid-configuration"),
-        ]
+        url = oauth_discovery_url
     else:
-        urls_to_try = [urljoin(server_url, "/.well-known/oauth-authorization-server")]
+        url = urljoin(server_url, "/.well-known/oauth-authorization-server")
 
-    headers = {"MCP-Protocol-Version": protocol_version or LATEST_PROTOCOL_VERSION}
-
-    for url in urls_to_try:
-        try:
-            response = ssrf_proxy.get(url, headers=headers)
+    try:
+        headers = {"MCP-Protocol-Version": protocol_version or LATEST_PROTOCOL_VERSION}
+        response = httpx.get(url, headers=headers)
+        if response.status_code == 404:
+            return None
+        if not response.is_success:
+            raise ValueError(f"HTTP {response.status_code} trying to load well-known OAuth metadata")
+        return OAuthMetadata.model_validate(response.json())
+    except httpx.RequestError as e:
+        if isinstance(e, httpx.ConnectError):
+            response = httpx.get(url)
             if response.status_code == 404:
-                continue
+                return None
             if not response.is_success:
-                response.raise_for_status()
+                raise ValueError(f"HTTP {response.status_code} trying to load well-known OAuth metadata")
             return OAuthMetadata.model_validate(response.json())
-        except (RequestError, HTTPStatusError) as e:
-            if isinstance(e, ConnectError):
-                response = ssrf_proxy.get(url)
-                if response.status_code == 404:
-                    continue  # Try next URL
-                if not response.is_success:
-                    raise ValueError(f"HTTP {response.status_code} trying to load well-known OAuth metadata")
-                return OAuthMetadata.model_validate(response.json())
-            # For other errors, try next URL
-            continue
-
-    return None  # No metadata found
+        raise
 
 
 def start_authorization(
@@ -223,7 +213,7 @@ def exchange_authorization(
     redirect_uri: str,
 ) -> OAuthTokens:
     """Exchanges an authorization code for an access token."""
-    grant_type = MCPSupportGrantType.AUTHORIZATION_CODE.value
+    grant_type = "authorization_code"
 
     if metadata:
         token_url = metadata.token_endpoint
@@ -243,7 +233,7 @@ def exchange_authorization(
     if client_information.client_secret:
         params["client_secret"] = client_information.client_secret
 
-    response = ssrf_proxy.post(token_url, data=params)
+    response = httpx.post(token_url, data=params)
     if not response.is_success:
         raise ValueError(f"Token exchange failed: HTTP {response.status_code}")
     return OAuthTokens.model_validate(response.json())
@@ -256,7 +246,7 @@ def refresh_authorization(
     refresh_token: str,
 ) -> OAuthTokens:
     """Exchange a refresh token for an updated access token."""
-    grant_type = MCPSupportGrantType.REFRESH_TOKEN.value
+    grant_type = "refresh_token"
 
     if metadata:
         token_url = metadata.token_endpoint
@@ -273,55 +263,10 @@ def refresh_authorization(
 
     if client_information.client_secret:
         params["client_secret"] = client_information.client_secret
-    try:
-        response = ssrf_proxy.post(token_url, data=params)
-    except ssrf_proxy.MaxRetriesExceededError as e:
-        raise MCPRefreshTokenError(e) from e
+
+    response = httpx.post(token_url, data=params)
     if not response.is_success:
-        raise MCPRefreshTokenError(response.text)
-    return OAuthTokens.model_validate(response.json())
-
-
-def client_credentials_flow(
-    server_url: str,
-    metadata: OAuthMetadata | None,
-    client_information: OAuthClientInformation,
-    scope: str | None = None,
-) -> OAuthTokens:
-    """Execute Client Credentials Flow to get access token."""
-    grant_type = MCPSupportGrantType.CLIENT_CREDENTIALS.value
-
-    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")
-
-    # Support both Basic Auth and body parameters for client authentication
-    headers = {"Content-Type": "application/x-www-form-urlencoded"}
-    data = {"grant_type": grant_type}
-
-    if scope:
-        data["scope"] = scope
-
-    # If client_secret is provided, use Basic Auth (preferred method)
-    if client_information.client_secret:
-        credentials = f"{client_information.client_id}:{client_information.client_secret}"
-        encoded_credentials = base64.b64encode(credentials.encode()).decode()
-        headers["Authorization"] = f"Basic {encoded_credentials}"
-    else:
-        # Fall back to including credentials in the body
-        data["client_id"] = client_information.client_id
-        if client_information.client_secret:
-            data["client_secret"] = client_information.client_secret
-
-    response = ssrf_proxy.post(token_url, headers=headers, data=data)
-    if not response.is_success:
-        raise ValueError(
-            f"Client credentials token request failed: HTTP {response.status_code}, Response: {response.text}"
-        )
-
+        raise ValueError(f"Token refresh failed: HTTP {response.status_code}")
     return OAuthTokens.model_validate(response.json())
 
 
@@ -338,7 +283,7 @@ def register_client(
     else:
         registration_url = urljoin(server_url, "/register")
 
-    response = ssrf_proxy.post(
+    response = httpx.post(
         registration_url,
         json=client_metadata.model_dump(),
         headers={"Content-Type": "application/json"},
@@ -349,111 +294,28 @@ def register_client(
 
 
 def auth(
-    provider: MCPProviderEntity,
+    provider: OAuthClientProvider,
+    server_url: str,
     authorization_code: str | None = None,
     state_param: str | None = None,
-) -> AuthResult:
-    """
-    Orchestrates the full auth flow with a server using secure Redis state storage.
-
-    This function performs only network operations and returns actions that need
-    to be performed by the caller (such as saving data to database).
-
-    Args:
-        provider: The MCP provider entity
-        authorization_code: Optional authorization code from OAuth callback
-        state_param: Optional state parameter from OAuth callback
-
-    Returns:
-        AuthResult containing actions to be performed and response data
-    """
-    actions: list[AuthAction] = []
-    server_url = provider.decrypt_server_url()
-    server_metadata = discover_oauth_metadata(server_url)
-    client_metadata = provider.client_metadata
-    provider_id = provider.id
-    tenant_id = provider.tenant_id
-    client_information = provider.retrieve_client_information()
-    redirect_url = provider.redirect_url
-
-    # Determine grant type based on server metadata
-    if not server_metadata:
-        raise ValueError("Failed to discover OAuth metadata from server")
-
-    supported_grant_types = server_metadata.grant_types_supported or []
-
-    # Convert to lowercase for comparison
-    supported_grant_types_lower = [gt.lower() for gt in supported_grant_types]
-
-    # Determine which grant type to use
-    effective_grant_type = None
-    if MCPSupportGrantType.AUTHORIZATION_CODE.value in supported_grant_types_lower:
-        effective_grant_type = MCPSupportGrantType.AUTHORIZATION_CODE.value
-    else:
-        effective_grant_type = MCPSupportGrantType.CLIENT_CREDENTIALS.value
-
-    # Get stored credentials
-    credentials = provider.decrypt_credentials()
+    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")
-
-        # For client credentials flow, we don't need to register client dynamically
-        if effective_grant_type == MCPSupportGrantType.CLIENT_CREDENTIALS.value:
-            # Client should provide client_id and client_secret directly
-            raise ValueError("Client credentials flow requires client_id and client_secret to be provided")
-
         try:
-            full_information = register_client(server_url, server_metadata, client_metadata)
-        except RequestError as e:
+            full_information = register_client(server_url, metadata, provider.client_metadata)
+        except httpx.RequestError as e:
             raise ValueError(f"Could not register OAuth client: {e}")
-
-        # Return action to save client information
-        actions.append(
-            AuthAction(
-                action_type=AuthActionType.SAVE_CLIENT_INFO,
-                data={"client_information": full_information.model_dump()},
-                provider_id=provider_id,
-                tenant_id=tenant_id,
-            )
-        )
-
+        provider.save_client_information(full_information)
         client_information = full_information
 
-    # Handle client credentials flow
-    if effective_grant_type == MCPSupportGrantType.CLIENT_CREDENTIALS.value:
-        # Direct token request without user interaction
-        try:
-            scope = credentials.get("scope")
-            tokens = client_credentials_flow(
-                server_url,
-                server_metadata,
-                client_information,
-                scope,
-            )
-
-            # Return action to save tokens and grant type
-            token_data = tokens.model_dump()
-            token_data["grant_type"] = MCPSupportGrantType.CLIENT_CREDENTIALS.value
-
-            actions.append(
-                AuthAction(
-                    action_type=AuthActionType.SAVE_TOKENS,
-                    data=token_data,
-                    provider_id=provider_id,
-                    tenant_id=tenant_id,
-                )
-            )
-
-            return AuthResult(actions=actions, response={"result": "success"})
-        except (RequestError, ValueError, KeyError) as e:
-            # RequestError: HTTP request failed
-            # ValueError: Invalid response data
-            # KeyError: Missing required fields in response
-            raise ValueError(f"Client credentials flow failed: {e}")
-
-    # Exchange authorization code for tokens (Authorization Code flow)
+    # 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")
@@ -473,69 +335,35 @@ def auth(
 
         tokens = exchange_authorization(
             server_url,
-            server_metadata,
+            metadata,
             client_information,
             authorization_code,
             code_verifier,
             redirect_uri,
         )
+        provider.save_tokens(tokens)
+        return {"result": "success"}
 
-        # Return action to save tokens
-        actions.append(
-            AuthAction(
-                action_type=AuthActionType.SAVE_TOKENS,
-                data=tokens.model_dump(),
-                provider_id=provider_id,
-                tenant_id=tenant_id,
-            )
-        )
-
-        return AuthResult(actions=actions, response={"result": "success"})
-
-    provider_tokens = provider.retrieve_tokens()
+    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, server_metadata, client_information, provider_tokens.refresh_token
-            )
-
-            # Return action to save new tokens
-            actions.append(
-                AuthAction(
-                    action_type=AuthActionType.SAVE_TOKENS,
-                    data=new_tokens.model_dump(),
-                    provider_id=provider_id,
-                    tenant_id=tenant_id,
-                )
-            )
-
-            return AuthResult(actions=actions, response={"result": "success"})
-        except (RequestError, ValueError, KeyError) as e:
-            # RequestError: HTTP request failed
-            # ValueError: Invalid response data
-            # KeyError: Missing required fields in response
+            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 (only for authorization code flow)
+    # Start new authorization flow
     authorization_url, code_verifier = start_authorization(
         server_url,
-        server_metadata,
+        metadata,
         client_information,
-        redirect_url,
-        provider_id,
-        tenant_id,
+        provider.redirect_url,
+        provider.mcp_provider.id,
+        provider.mcp_provider.tenant_id,
     )
 
-    # Return action to save code verifier
-    actions.append(
-        AuthAction(
-            action_type=AuthActionType.SAVE_CODE_VERIFIER,
-            data={"code_verifier": code_verifier},
-            provider_id=provider_id,
-            tenant_id=tenant_id,
-        )
-    )
-
-    return AuthResult(actions=actions, response={"authorization_url": authorization_url})
+    provider.save_code_verifier(code_verifier)
+    return {"authorization_url": authorization_url}

api/core/mcp/auth/auth_provider.py

@@ -0,0 +1,77 @@
+from configs import dify_config
+from core.mcp.types import (
+    OAuthClientInformation,
+    OAuthClientInformationFull,
+    OAuthClientMetadata,
+    OAuthTokens,
+)
+from models.tools import MCPToolProvider
+from services.tools.mcp_tools_manage_service import MCPToolManageService
+
+
+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) -> OAuthClientInformation | None:
+        """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):
+        """Saves client information after dynamic registration."""
+        MCPToolManageService.update_mcp_provider_credentials(
+            self.mcp_provider,
+            {"client_information": client_information.model_dump()},
+        )
+
+    def tokens(self) -> OAuthTokens | None:
+        """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):
+        """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):
+        """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/auth_client.py

@@ -1,191 +0,0 @@
-"""
-MCP Client with Authentication Retry Support
-
-This module provides an enhanced MCPClient that automatically handles
-authentication failures and retries operations after refreshing tokens.
-"""
-
-import logging
-from collections.abc import Callable
-from typing import Any
-
-from sqlalchemy.orm import Session
-
-from core.entities.mcp_provider import MCPProviderEntity
-from core.mcp.error import MCPAuthError
-from core.mcp.mcp_client import MCPClient
-from core.mcp.types import CallToolResult, Tool
-from extensions.ext_database import db
-
-logger = logging.getLogger(__name__)
-
-
-class MCPClientWithAuthRetry(MCPClient):
-    """
-    An enhanced MCPClient that provides automatic authentication retry.
-
-    This class extends MCPClient and intercepts MCPAuthError exceptions
-    to refresh authentication before retrying failed operations.
-
-    Note: This class uses lazy session creation - database sessions are only
-    created when authentication retry is actually needed, not on every request.
-    """
-
-    def __init__(
-        self,
-        server_url: str,
-        headers: dict[str, str] | None = None,
-        timeout: float | None = None,
-        sse_read_timeout: float | None = None,
-        provider_entity: MCPProviderEntity | None = None,
-        authorization_code: str | None = None,
-        by_server_id: bool = False,
-    ):
-        """
-        Initialize the MCP client with auth retry capability.
-
-        Args:
-            server_url: The MCP server URL
-            headers: Optional headers for requests
-            timeout: Request timeout
-            sse_read_timeout: SSE read timeout
-            provider_entity: Provider entity for authentication
-            authorization_code: Optional authorization code for initial auth
-            by_server_id: Whether to look up provider by server ID
-        """
-        super().__init__(server_url, headers, timeout, sse_read_timeout)
-
-        self.provider_entity = provider_entity
-        self.authorization_code = authorization_code
-        self.by_server_id = by_server_id
-        self._has_retried = False
-
-    def _handle_auth_error(self, error: MCPAuthError) -> None:
-        """
-        Handle authentication error by refreshing tokens.
-
-        This method creates a short-lived database session only when authentication
-        retry is needed, minimizing database connection hold time.
-
-        Args:
-            error: The authentication error
-
-        Raises:
-            MCPAuthError: If authentication fails or max retries reached
-        """
-        if not self.provider_entity:
-            raise error
-        if self._has_retried:
-            raise error
-
-        self._has_retried = True
-
-        try:
-            # Create a temporary session only for auth retry
-            # This session is short-lived and only exists during the auth operation
-
-            from services.tools.mcp_tools_manage_service import MCPToolManageService
-
-            with Session(db.engine) as session, session.begin():
-                mcp_service = MCPToolManageService(session=session)
-
-                # Perform authentication using the service's auth method
-                mcp_service.auth_with_actions(self.provider_entity, self.authorization_code)
-
-                # Retrieve new tokens
-                self.provider_entity = mcp_service.get_provider_entity(
-                    self.provider_entity.id, self.provider_entity.tenant_id, by_server_id=self.by_server_id
-                )
-
-            # Session is closed here, before we update headers
-            token = self.provider_entity.retrieve_tokens()
-            if not token:
-                raise MCPAuthError("Authentication failed - no token received")
-
-            # Update headers with new token
-            self.headers["Authorization"] = f"{token.token_type.capitalize()} {token.access_token}"
-
-            # Clear authorization code after first use
-            self.authorization_code = None
-
-        except MCPAuthError:
-            # Re-raise MCPAuthError as is
-            raise
-        except Exception as e:
-            # Catch all exceptions during auth retry
-            logger.exception("Authentication retry failed")
-            raise MCPAuthError(f"Authentication retry failed: {e}") from e
-
-    def _execute_with_retry(self, func: Callable[..., Any], *args, **kwargs) -> Any:
-        """
-        Execute a function with authentication retry logic.
-
-        Args:
-            func: The function to execute
-            *args: Positional arguments for the function
-            **kwargs: Keyword arguments for the function
-
-        Returns:
-            The result of the function call
-
-        Raises:
-            MCPAuthError: If authentication fails after retries
-            Any other exceptions from the function
-        """
-        try:
-            return func(*args, **kwargs)
-        except MCPAuthError as e:
-            self._handle_auth_error(e)
-
-            # Re-initialize the connection with new headers
-            if self._initialized:
-                # Clean up existing connection
-                self._exit_stack.close()
-                self._session = None
-                self._initialized = False
-
-                # Re-initialize with new headers
-                self._initialize()
-                self._initialized = True
-
-            return func(*args, **kwargs)
-        finally:
-            # Reset retry flag after operation completes
-            self._has_retried = False
-
-    def __enter__(self):
-        """Enter the context manager with retry support."""
-
-        def initialize_with_retry():
-            super(MCPClientWithAuthRetry, self).__enter__()
-            return self
-
-        return self._execute_with_retry(initialize_with_retry)
-
-    def list_tools(self) -> list[Tool]:
-        """
-        List available tools from the MCP server with auth retry.
-
-        Returns:
-            List of available tools
-
-        Raises:
-            MCPAuthError: If authentication fails after retries
-        """
-        return self._execute_with_retry(super().list_tools)
-
-    def invoke_tool(self, tool_name: str, tool_args: dict[str, Any]) -> CallToolResult:
-        """
-        Invoke a tool on the MCP server with auth retry.
-
-        Args:
-            tool_name: Name of the tool to invoke
-            tool_args: Arguments for the tool
-
-        Returns:
-            Result of the tool invocation
-
-        Raises:
-            MCPAuthError: If authentication fails after retries
-        """
-        return self._execute_with_retry(super().invoke_tool, tool_name, tool_args)

api/core/mcp/client/sse_client.py

@@ -46,7 +46,7 @@ class SSETransport:
         url: str,
         headers: dict[str, Any] | None = None,
         timeout: float = 5.0,
-        sse_read_timeout: float = 1 * 60,
+        sse_read_timeout: float = 5 * 60,
     ):
         """Initialize the SSE transport.
 
@@ -255,7 +255,7 @@ def sse_client(
     url: str,
     headers: dict[str, Any] | None = None,
     timeout: float = 5.0,
-    sse_read_timeout: float = 1 * 60,
+    sse_read_timeout: float = 5 * 60,
 ) -> Generator[tuple[ReadQueue, WriteQueue], None, None]:
     """
     Client transport for SSE.
@@ -276,34 +276,31 @@ def sse_client(
     read_queue: ReadQueue | None = None
     write_queue: WriteQueue | None = None
 
-    executor = ThreadPoolExecutor()
-    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()
+    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)
+                    read_queue, write_queue = transport.connect(executor, client, event_source)
 
-                yield read_queue, write_queue
+                    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)
-
-        # Shutdown executor without waiting to prevent hanging
-        executor.shutdown(wait=False)
+        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):

api/core/mcp/client/streamable_client.py

@@ -434,48 +434,45 @@ def streamablehttp_client(
     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
 
-    executor = ThreadPoolExecutor(max_workers=2)
-    try:
-        with create_ssrf_proxy_mcp_http_client(
-            headers=transport.request_headers,
-            timeout=httpx.Timeout(transport.timeout, read=transport.sse_read_timeout),
-        ) as client:
-            # Define callbacks that need access to thread pool
-            def start_get_stream():
-                """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
+    with ThreadPoolExecutor(max_workers=2) as executor:
         try:
-            while not client_to_server_queue.empty():
-                client_to_server_queue.get_nowait()
-        except queue.Empty:
-            pass
+            with create_ssrf_proxy_mcp_http_client(
+                headers=transport.request_headers,
+                timeout=httpx.Timeout(transport.timeout, read=transport.sse_read_timeout),
+            ) as client:
+                # Define callbacks that need access to thread pool
+                def start_get_stream():
+                    """Start a worker thread to handle server-initiated messages."""
+                    executor.submit(transport.handle_get_stream, client, server_to_client_queue)
 
-        client_to_server_queue.put(None)
-        server_to_client_queue.put(None)
+                # 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,
+                )
 
-        # Shutdown executor without waiting to prevent hanging
-        executor.shutdown(wait=False)
+                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

@@ -1,13 +1,10 @@
 from dataclasses import dataclass
-from enum import StrEnum
 from typing import Any, Generic, TypeVar
 
-from pydantic import BaseModel
-
 from core.mcp.session.base_session import BaseSession
-from core.mcp.types import LATEST_PROTOCOL_VERSION, OAuthClientInformation, OAuthMetadata, RequestId, RequestParams
+from core.mcp.types import LATEST_PROTOCOL_VERSION, RequestId, RequestParams
 
-SUPPORTED_PROTOCOL_VERSIONS: list[str] = ["2024-11-05", "2025-03-26", LATEST_PROTOCOL_VERSION]
+SUPPORTED_PROTOCOL_VERSIONS: list[str] = ["2024-11-05", LATEST_PROTOCOL_VERSION]
 
 
 SessionT = TypeVar("SessionT", bound=BaseSession[Any, Any, Any, Any, Any])
@@ -20,41 +17,3 @@ class RequestContext(Generic[SessionT, LifespanContextT]):
     meta: RequestParams.Meta | None
     session: SessionT
     lifespan_context: LifespanContextT
-
-
-class AuthActionType(StrEnum):
-    """Types of actions that can be performed during auth flow."""
-
-    SAVE_CLIENT_INFO = "save_client_info"
-    SAVE_TOKENS = "save_tokens"
-    SAVE_CODE_VERIFIER = "save_code_verifier"
-    START_AUTHORIZATION = "start_authorization"
-    SUCCESS = "success"
-
-
-class AuthAction(BaseModel):
-    """Represents an action that needs to be performed as a result of auth flow."""
-
-    action_type: AuthActionType
-    data: dict[str, Any]
-    provider_id: str | None = None
-    tenant_id: str | None = None
-
-
-class AuthResult(BaseModel):
-    """Result of auth function containing actions to be performed and response data."""
-
-    actions: list[AuthAction]
-    response: dict[str, str]
-
-
-class OAuthCallbackState(BaseModel):
-    """State data stored in Redis during OAuth callback flow."""
-
-    provider_id: str
-    tenant_id: str
-    server_url: str
-    metadata: OAuthMetadata | None = None
-    client_information: OAuthClientInformation
-    code_verifier: str
-    redirect_uri: str

api/core/mcp/error.py

@@ -8,7 +8,3 @@ class MCPConnectionError(MCPError):
 
 class MCPAuthError(MCPConnectionError):
     pass
-
-
-class MCPRefreshTokenError(MCPError):
-    pass

api/core/mcp/mcp_client.py

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

api/core/mcp/session/base_session.py

@@ -201,14 +201,11 @@ class BaseSession(
                 self._receiver_future.result(timeout=5.0)  # Wait up to 5 seconds
             except TimeoutError:
                 # If the receiver loop is still running after timeout, we'll force shutdown
-                # Cancel the future to interrupt the receiver loop
-                self._receiver_future.cancel()
+                pass
 
         # Shutdown the executor
         if self._executor:
-            # Use non-blocking shutdown to prevent hanging
-            # The receiver thread should have already exited due to the None message in the queue
-            self._executor.shutdown(wait=False)
+            self._executor.shutdown(wait=True)
 
     def send_request(
         self,

api/core/mcp/session/client_session.py

@@ -109,16 +109,12 @@ class ClientSession(
         self._message_handler = message_handler or _default_message_handler
 
     def initialize(self) -> types.InitializeResult:
-        # Only set capabilities if non-default callbacks are provided
-        # This prevents servers from attempting callbacks when we don't actually support them
-        sampling = types.SamplingCapability() if self._sampling_callback is not _default_sampling_callback else None
-        roots = (
-            types.RootsCapability(
-                # Only enable listChanged if we have a custom callback
-                listChanged=True,
-            )
-            if self._list_roots_callback is not _default_list_roots_callback
-            else None
+        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(
@@ -288,7 +284,7 @@ class ClientSession(
 
     def complete(
         self,
-        ref: types.ResourceTemplateReference | types.PromptReference,
+        ref: types.ResourceReference | types.PromptReference,
         argument: dict[str, str],
     ) -> types.CompleteResult:
         """Send a completion/complete request."""

api/core/mcp/types.py

@@ -1,6 +1,13 @@
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Annotated, Any, Generic, Literal, TypeAlias, TypeVar
+from typing import (
+    Annotated,
+    Any,
+    Generic,
+    Literal,
+    TypeAlias,
+    TypeVar,
+)
 
 from pydantic import BaseModel, ConfigDict, Field, FileUrl, RootModel
 from pydantic.networks import AnyUrl, UrlConstraints
@@ -26,7 +33,6 @@ for reference.
 LATEST_PROTOCOL_VERSION = "2025-03-26"
 # Server support 2024-11-05 to allow claude to use.
 SERVER_LATEST_PROTOCOL_VERSION = "2024-11-05"
-DEFAULT_NEGOTIATED_VERSION = "2025-03-26"
 ProgressToken = str | int
 Cursor = str
 Role = Literal["user", "assistant"]
@@ -49,22 +55,14 @@ class RequestParams(BaseModel):
     meta: Meta | None = Field(alias="_meta", default=None)
 
 
-class PaginatedRequestParams(RequestParams):
-    cursor: Cursor | None = None
-    """
-    An opaque token representing the current pagination position.
-    If provided, the server should return results starting after this cursor.
-    """
-
-
 class NotificationParams(BaseModel):
     class Meta(BaseModel):
         model_config = ConfigDict(extra="allow")
 
     meta: Meta | None = Field(alias="_meta", default=None)
     """
-    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
-    for notes on _meta usage.
+    This parameter name is reserved by MCP to allow clients and servers to attach
+    additional metadata to their notifications.
     """
 
 
@@ -81,11 +79,12 @@ class Request(BaseModel, Generic[RequestParamsT, MethodT]):
     model_config = ConfigDict(extra="allow")
 
 
-class PaginatedRequest(Request[PaginatedRequestParams | None, MethodT], Generic[MethodT]):
-    """Base class for paginated requests,
-    matching the schema's PaginatedRequest interface."""
-
-    params: PaginatedRequestParams | None = None
+class PaginatedRequest(Request[RequestParamsT, MethodT]):
+    cursor: Cursor | None = None
+    """
+    An opaque token representing the current pagination position.
+    If provided, the server should return results starting after this cursor.
+    """
 
 
 class Notification(BaseModel, Generic[NotificationParamsT, MethodT]):
@@ -99,12 +98,13 @@ class Notification(BaseModel, Generic[NotificationParamsT, MethodT]):
 class Result(BaseModel):
     """Base class for JSON-RPC results."""
 
+    model_config = ConfigDict(extra="allow")
+
     meta: dict[str, Any] | None = Field(alias="_meta", default=None)
     """
-    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
-    for notes on _meta usage.
+    This result property is reserved by the protocol to allow clients and servers to
+    attach additional metadata to their responses.
     """
-    model_config = ConfigDict(extra="allow")
 
 
 class PaginatedResult(Result):
@@ -186,26 +186,10 @@ class EmptyResult(Result):
     """A response that indicates success but carries no data."""
 
 
-class BaseMetadata(BaseModel):
-    """Base class for entities with name and optional title fields."""
-
-    name: str
-    """The programmatic name of the entity."""
-
-    title: str | None = None
-    """
-    Intended for UI and end-user contexts — optimized to be human-readable and easily understood,
-    even by those unfamiliar with domain-specific terminology.
-
-    If not provided, the name should be used for display (except for Tool,
-    where `annotations.title` should be given precedence over using `name`,
-    if present).
-    """
-
-
-class Implementation(BaseMetadata):
+class Implementation(BaseModel):
     """Describes the name and version of an MCP implementation."""
 
+    name: str
     version: str
     model_config = ConfigDict(extra="allow")
 
@@ -219,7 +203,7 @@ class RootsCapability(BaseModel):
 
 
 class SamplingCapability(BaseModel):
-    """Capability for sampling operations."""
+    """Capability for logging operations."""
 
     model_config = ConfigDict(extra="allow")
 
@@ -268,12 +252,6 @@ class LoggingCapability(BaseModel):
     model_config = ConfigDict(extra="allow")
 
 
-class CompletionsCapability(BaseModel):
-    """Capability for completions operations."""
-
-    model_config = ConfigDict(extra="allow")
-
-
 class ServerCapabilities(BaseModel):
     """Capabilities that a server may support."""
 
@@ -287,8 +265,6 @@ class ServerCapabilities(BaseModel):
     """Present if the server offers any resources to read."""
     tools: ToolsCapability | None = None
     """Present if the server offers any tools to call."""
-    completions: CompletionsCapability | None = None
-    """Present if the server offers autocompletion suggestions for prompts and resources."""
     model_config = ConfigDict(extra="allow")
 
 
@@ -308,7 +284,7 @@ class InitializeRequest(Request[InitializeRequestParams, Literal["initialize"]])
     to begin initialization.
     """
 
-    method: Literal["initialize"] = "initialize"
+    method: Literal["initialize"]
     params: InitializeRequestParams
 
 
@@ -329,7 +305,7 @@ class InitializedNotification(Notification[NotificationParams | None, Literal["n
     finished.
     """
 
-    method: Literal["notifications/initialized"] = "notifications/initialized"
+    method: Literal["notifications/initialized"]
     params: NotificationParams | None = None
 
 
@@ -339,7 +315,7 @@ class PingRequest(Request[RequestParams | None, Literal["ping"]]):
     still alive.
     """
 
-    method: Literal["ping"] = "ping"
+    method: Literal["ping"]
     params: RequestParams | None = None
 
 
@@ -358,11 +334,6 @@ class ProgressNotificationParams(NotificationParams):
     """
     total: float | None = None
     """Total number of items to process (or total progress required), if known."""
-    message: str | None = None
-    """
-    Message related to progress. This should provide relevant human readable
-    progress information.
-    """
     model_config = ConfigDict(extra="allow")
 
 
@@ -372,14 +343,15 @@ class ProgressNotification(Notification[ProgressNotificationParams, Literal["not
     long-running request.
     """
 
-    method: Literal["notifications/progress"] = "notifications/progress"
+    method: Literal["notifications/progress"]
     params: ProgressNotificationParams
 
 
-class ListResourcesRequest(PaginatedRequest[Literal["resources/list"]]):
+class ListResourcesRequest(PaginatedRequest[RequestParams | None, Literal["resources/list"]]):
     """Sent from the client to request a list of resources the server has."""
 
-    method: Literal["resources/list"] = "resources/list"
+    method: Literal["resources/list"]
+    params: RequestParams | None = None
 
 
 class Annotations(BaseModel):
@@ -388,11 +360,13 @@ class Annotations(BaseModel):
     model_config = ConfigDict(extra="allow")
 
 
-class Resource(BaseMetadata):
+class Resource(BaseModel):
     """A known resource that the server is capable of reading."""
 
     uri: Annotated[AnyUrl, UrlConstraints(host_required=False)]
     """The URI of this resource."""
+    name: str
+    """A human-readable name for this resource."""
     description: str | None = None
     """A description of what this resource represents."""
     mimeType: str | None = None
@@ -405,15 +379,10 @@ class Resource(BaseMetadata):
     This can be used by Hosts to display file sizes and estimate context window usage.
     """
     annotations: Annotations | None = None
-    meta: dict[str, Any] | None = Field(alias="_meta", default=None)
-    """
-    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
-    for notes on _meta usage.
-    """
     model_config = ConfigDict(extra="allow")
 
 
-class ResourceTemplate(BaseMetadata):
+class ResourceTemplate(BaseModel):
     """A template description for resources available on the server."""
 
     uriTemplate: str
@@ -421,6 +390,8 @@ class ResourceTemplate(BaseMetadata):
     A URI template (according to RFC 6570) that can be used to construct resource
     URIs.
     """
+    name: str
+    """A human-readable name for the type of resource this template refers to."""
     description: str | None = None
     """A human-readable description of what this template is for."""
     mimeType: str | None = None
@@ -429,11 +400,6 @@ class ResourceTemplate(BaseMetadata):
     included if all resources matching this template have the same type.
     """
     annotations: Annotations | None = None
-    meta: dict[str, Any] | None = Field(alias="_meta", default=None)
-    """
-    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
-    for notes on _meta usage.
-    """
     model_config = ConfigDict(extra="allow")
 
 
@@ -443,10 +409,11 @@ class ListResourcesResult(PaginatedResult):
     resources: list[Resource]
 
 
-class ListResourceTemplatesRequest(PaginatedRequest[Literal["resources/templates/list"]]):
+class ListResourceTemplatesRequest(PaginatedRequest[RequestParams | None, Literal["resources/templates/list"]]):
     """Sent from the client to request a list of resource templates the server has."""
 
-    method: Literal["resources/templates/list"] = "resources/templates/list"
+    method: Literal["resources/templates/list"]
+    params: RequestParams | None = None
 
 
 class ListResourceTemplatesResult(PaginatedResult):
@@ -469,7 +436,7 @@ class ReadResourceRequestParams(RequestParams):
 class ReadResourceRequest(Request[ReadResourceRequestParams, Literal["resources/read"]]):
     """Sent from the client to the server, to read a specific resource URI."""
 
-    method: Literal["resources/read"] = "resources/read"
+    method: Literal["resources/read"]
     params: ReadResourceRequestParams
 
 
@@ -480,11 +447,6 @@ class ResourceContents(BaseModel):
     """The URI of this resource."""
     mimeType: str | None = None
     """The MIME type of this resource, if known."""
-    meta: dict[str, Any] | None = Field(alias="_meta", default=None)
-    """
-    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
-    for notes on _meta usage.
-    """
     model_config = ConfigDict(extra="allow")
 
 
@@ -519,7 +481,7 @@ class ResourceListChangedNotification(
     of resources it can read from has changed.
     """
 
-    method: Literal["notifications/resources/list_changed"] = "notifications/resources/list_changed"
+    method: Literal["notifications/resources/list_changed"]
     params: NotificationParams | None = None
 
 
@@ -540,7 +502,7 @@ class SubscribeRequest(Request[SubscribeRequestParams, Literal["resources/subscr
     whenever a particular resource changes.
     """
 
-    method: Literal["resources/subscribe"] = "resources/subscribe"
+    method: Literal["resources/subscribe"]
     params: SubscribeRequestParams
 
 
@@ -558,7 +520,7 @@ class UnsubscribeRequest(Request[UnsubscribeRequestParams, Literal["resources/un
     the server.
     """
 
-    method: Literal["resources/unsubscribe"] = "resources/unsubscribe"
+    method: Literal["resources/unsubscribe"]
     params: UnsubscribeRequestParams
 
 
@@ -581,14 +543,15 @@ class ResourceUpdatedNotification(
     changed and may need to be read again.
     """
 
-    method: Literal["notifications/resources/updated"] = "notifications/resources/updated"
+    method: Literal["notifications/resources/updated"]
     params: ResourceUpdatedNotificationParams
 
 
-class ListPromptsRequest(PaginatedRequest[Literal["prompts/list"]]):
+class ListPromptsRequest(PaginatedRequest[RequestParams | None, Literal["prompts/list"]]):
     """Sent from the client to request a list of prompts and prompt templates."""
 
-    method: Literal["prompts/list"] = "prompts/list"
+    method: Literal["prompts/list"]
+    params: RequestParams | None = None
 
 
 class PromptArgument(BaseModel):
@@ -603,18 +566,15 @@ class PromptArgument(BaseModel):
     model_config = ConfigDict(extra="allow")
 
 
-class Prompt(BaseMetadata):
+class Prompt(BaseModel):
     """A prompt or prompt template that the server offers."""
 
+    name: str
+    """The name of the prompt or prompt template."""
     description: str | None = None
     """An optional description of what this prompt provides."""
     arguments: list[PromptArgument] | None = None
     """A list of arguments to use for templating the prompt."""
-    meta: dict[str, Any] | None = Field(alias="_meta", default=None)
-    """
-    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
-    for notes on _meta usage.
-    """
     model_config = ConfigDict(extra="allow")
 
 
@@ -637,7 +597,7 @@ class GetPromptRequestParams(RequestParams):
 class GetPromptRequest(Request[GetPromptRequestParams, Literal["prompts/get"]]):
     """Used by the client to get a prompt provided by the server."""
 
-    method: Literal["prompts/get"] = "prompts/get"
+    method: Literal["prompts/get"]
     params: GetPromptRequestParams
 
 
@@ -648,11 +608,6 @@ class TextContent(BaseModel):
     text: str
     """The text content of the message."""
     annotations: Annotations | None = None
-    meta: dict[str, Any] | None = Field(alias="_meta", default=None)
-    """
-    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
-    for notes on _meta usage.
-    """
     model_config = ConfigDict(extra="allow")
 
 
@@ -668,31 +623,6 @@ class ImageContent(BaseModel):
     image types.
     """
     annotations: Annotations | None = None
-    meta: dict[str, Any] | None = Field(alias="_meta", default=None)
-    """
-    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
-    for notes on _meta usage.
-    """
-    model_config = ConfigDict(extra="allow")
-
-
-class AudioContent(BaseModel):
-    """Audio content for a message."""
-
-    type: Literal["audio"]
-    data: str
-    """The base64-encoded audio data."""
-    mimeType: str
-    """
-    The MIME type of the audio. Different providers may support different
-    audio types.
-    """
-    annotations: Annotations | None = None
-    meta: dict[str, Any] | None = Field(alias="_meta", default=None)
-    """
-    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
-    for notes on _meta usage.
-    """
     model_config = ConfigDict(extra="allow")
 
 
@@ -700,7 +630,7 @@ class SamplingMessage(BaseModel):
     """Describes a message issued to or received from an LLM API."""
 
     role: Role
-    content: TextContent | ImageContent | AudioContent
+    content: TextContent | ImageContent
     model_config = ConfigDict(extra="allow")
 
 
@@ -715,36 +645,14 @@ class EmbeddedResource(BaseModel):
     type: Literal["resource"]
     resource: TextResourceContents | BlobResourceContents
     annotations: Annotations | None = None
-    meta: dict[str, Any] | None = Field(alias="_meta", default=None)
-    """
-    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
-    for notes on _meta usage.
-    """
     model_config = ConfigDict(extra="allow")
 
 
-class ResourceLink(Resource):
-    """
-    A resource that the server is capable of reading, included in a prompt or tool call result.
-
-    Note: resource links returned by tools are not guaranteed to appear in the results of `resources/list` requests.
-    """
-
-    type: Literal["resource_link"]
-
-
-ContentBlock = TextContent | ImageContent | AudioContent | ResourceLink | EmbeddedResource
-"""A content block that can be used in prompts and tool results."""
-
-Content: TypeAlias = ContentBlock
-# """DEPRECATED: Content is deprecated, you should use ContentBlock directly."""
-
-
 class PromptMessage(BaseModel):
     """Describes a message returned as part of a prompt."""
 
     role: Role
-    content: ContentBlock
+    content: TextContent | ImageContent | EmbeddedResource
     model_config = ConfigDict(extra="allow")
 
 
@@ -764,14 +672,15 @@ class PromptListChangedNotification(
     of prompts it offers has changed.
     """
 
-    method: Literal["notifications/prompts/list_changed"] = "notifications/prompts/list_changed"
+    method: Literal["notifications/prompts/list_changed"]
     params: NotificationParams | None = None
 
 
-class ListToolsRequest(PaginatedRequest[Literal["tools/list"]]):
+class ListToolsRequest(PaginatedRequest[RequestParams | None, Literal["tools/list"]]):
     """Sent from the client to request a list of tools the server has."""
 
-    method: Literal["tools/list"] = "tools/list"
+    method: Literal["tools/list"]
+    params: RequestParams | None = None
 
 
 class ToolAnnotations(BaseModel):
@@ -822,25 +731,17 @@ class ToolAnnotations(BaseModel):
     model_config = ConfigDict(extra="allow")
 
 
-class Tool(BaseMetadata):
+class Tool(BaseModel):
     """Definition for a tool the client can call."""
 
+    name: str
+    """The name of the tool."""
     description: str | None = None
     """A human-readable description of the tool."""
     inputSchema: dict[str, Any]
     """A JSON Schema object defining the expected parameters for the tool."""
-    outputSchema: dict[str, Any] | None = None
-    """
-    An optional JSON Schema object defining the structure of the tool's output
-    returned in the structuredContent field of a CallToolResult.
-    """
     annotations: ToolAnnotations | None = None
     """Optional additional tool information."""
-    meta: dict[str, Any] | None = Field(alias="_meta", default=None)
-    """
-    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
-    for notes on _meta usage.
-    """
     model_config = ConfigDict(extra="allow")
 
 
@@ -861,16 +762,14 @@ class CallToolRequestParams(RequestParams):
 class CallToolRequest(Request[CallToolRequestParams, Literal["tools/call"]]):
     """Used by the client to invoke a tool provided by the server."""
 
-    method: Literal["tools/call"] = "tools/call"
+    method: Literal["tools/call"]
     params: CallToolRequestParams
 
 
 class CallToolResult(Result):
     """The server's response to a tool call."""
 
-    content: list[ContentBlock]
-    structuredContent: dict[str, Any] | None = None
-    """An optional JSON object that represents the structured result of the tool call."""
+    content: list[TextContent | ImageContent | EmbeddedResource]
     isError: bool = False
 
 
@@ -880,7 +779,7 @@ class ToolListChangedNotification(Notification[NotificationParams | None, Litera
     of tools it offers has changed.
     """
 
-    method: Literal["notifications/tools/list_changed"] = "notifications/tools/list_changed"
+    method: Literal["notifications/tools/list_changed"]
     params: NotificationParams | None = None
 
 
@@ -898,7 +797,7 @@ class SetLevelRequestParams(RequestParams):
 class SetLevelRequest(Request[SetLevelRequestParams, Literal["logging/setLevel"]]):
     """A request from the client to the server, to enable or adjust logging."""
 
-    method: Literal["logging/setLevel"] = "logging/setLevel"
+    method: Literal["logging/setLevel"]
     params: SetLevelRequestParams
 
 
@@ -909,7 +808,7 @@ class LoggingMessageNotificationParams(NotificationParams):
     """The severity of this log message."""
     logger: str | None = None
     """An optional name of the logger issuing this message."""
-    data: Any
+    data: Any = None
     """
     The data to be logged, such as a string message or an object. Any JSON serializable
     type is allowed here.
@@ -920,7 +819,7 @@ class LoggingMessageNotificationParams(NotificationParams):
 class LoggingMessageNotification(Notification[LoggingMessageNotificationParams, Literal["notifications/message"]]):
     """Notification of a log message passed from server to client."""
 
-    method: Literal["notifications/message"] = "notifications/message"
+    method: Literal["notifications/message"]
     params: LoggingMessageNotificationParams
 
 
@@ -1015,7 +914,7 @@ class CreateMessageRequestParams(RequestParams):
 class CreateMessageRequest(Request[CreateMessageRequestParams, Literal["sampling/createMessage"]]):
     """A request from the server to sample an LLM via the client."""
 
-    method: Literal["sampling/createMessage"] = "sampling/createMessage"
+    method: Literal["sampling/createMessage"]
     params: CreateMessageRequestParams
 
 
@@ -1026,14 +925,14 @@ class CreateMessageResult(Result):
     """The client's response to a sampling/create_message request from the server."""
 
     role: Role
-    content: TextContent | ImageContent | AudioContent
+    content: TextContent | ImageContent
     model: str
     """The name of the model that generated the message."""
     stopReason: StopReason | None = None
     """The reason why sampling stopped, if known."""
 
 
-class ResourceTemplateReference(BaseModel):
+class ResourceReference(BaseModel):
     """A reference to a resource or resource template definition."""
 
     type: Literal["ref/resource"]
@@ -1061,28 +960,18 @@ class CompletionArgument(BaseModel):
     model_config = ConfigDict(extra="allow")
 
 
-class CompletionContext(BaseModel):
-    """Additional, optional context for completions."""
-
-    arguments: dict[str, str] | None = None
-    """Previously-resolved variables in a URI template or prompt."""
-    model_config = ConfigDict(extra="allow")
-
-
 class CompleteRequestParams(RequestParams):
     """Parameters for completion requests."""
 
-    ref: ResourceTemplateReference | PromptReference
+    ref: ResourceReference | PromptReference
     argument: CompletionArgument
-    context: CompletionContext | None = None
-    """Additional, optional context for completions"""
     model_config = ConfigDict(extra="allow")
 
 
 class CompleteRequest(Request[CompleteRequestParams, Literal["completion/complete"]]):
     """A request from the client to the server, to ask for completion options."""
 
-    method: Literal["completion/complete"] = "completion/complete"
+    method: Literal["completion/complete"]
     params: CompleteRequestParams
 
 
@@ -1121,7 +1010,7 @@ class ListRootsRequest(Request[RequestParams | None, Literal["roots/list"]]):
     structure or access specific locations that the client has permission to read from.
     """
 
-    method: Literal["roots/list"] = "roots/list"
+    method: Literal["roots/list"]
     params: RequestParams | None = None
 
 
@@ -1140,11 +1029,6 @@ class Root(BaseModel):
     identifier for the root, which may be useful for display purposes or for
     referencing the root in other parts of the application.
     """
-    meta: dict[str, Any] | None = Field(alias="_meta", default=None)
-    """
-    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
-    for notes on _meta usage.
-    """
     model_config = ConfigDict(extra="allow")
 
 
@@ -1170,7 +1054,7 @@ class RootsListChangedNotification(
     using the ListRootsRequest.
     """
 
-    method: Literal["notifications/roots/list_changed"] = "notifications/roots/list_changed"
+    method: Literal["notifications/roots/list_changed"]
     params: NotificationParams | None = None
 
 
@@ -1190,7 +1074,7 @@ class CancelledNotification(Notification[CancelledNotificationParams, Literal["n
     previously-issued request.
     """
 
-    method: Literal["notifications/cancelled"] = "notifications/cancelled"
+    method: Literal["notifications/cancelled"]
     params: CancelledNotificationParams
 
 

api/core/memory/token_buffer_memory.py

@@ -1,7 +1,6 @@
 from collections.abc import Sequence
 
 from sqlalchemy import select
-from sqlalchemy.orm import sessionmaker
 
 from core.app.app_config.features.file_upload.manager import FileUploadConfigManager
 from core.file import file_manager
@@ -19,9 +18,7 @@ from core.prompt.utils.extract_thread_messages import extract_thread_messages
 from extensions.ext_database import db
 from factories import file_factory
 from models.model import AppMode, Conversation, Message, MessageFile
-from models.workflow import Workflow
-from repositories.api_workflow_run_repository import APIWorkflowRunRepository
-from repositories.factory import DifyAPIRepositoryFactory
+from models.workflow import Workflow, WorkflowRun
 
 
 class TokenBufferMemory:
@@ -32,14 +29,6 @@ class TokenBufferMemory:
     ):
         self.conversation = conversation
         self.model_instance = model_instance
-        self._workflow_run_repo: APIWorkflowRunRepository | None = None
-
-    @property
-    def workflow_run_repo(self) -> APIWorkflowRunRepository:
-        if self._workflow_run_repo is None:
-            session_maker = sessionmaker(bind=db.engine, expire_on_commit=False)
-            self._workflow_run_repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker)
-        return self._workflow_run_repo
 
     def _build_prompt_message_with_files(
         self,
@@ -61,16 +50,7 @@ class TokenBufferMemory:
         if self.conversation.mode in {AppMode.AGENT_CHAT, AppMode.COMPLETION, AppMode.CHAT}:
             file_extra_config = FileUploadConfigManager.convert(self.conversation.model_config)
         elif self.conversation.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}:
-            app = self.conversation.app
-            if not app:
-                raise ValueError("App not found for conversation")
-
-            if not message.workflow_run_id:
-                raise ValueError("Workflow run ID not found")
-
-            workflow_run = self.workflow_run_repo.get_workflow_run_by_id(
-                tenant_id=app.tenant_id, app_id=app.id, run_id=message.workflow_run_id
-            )
+            workflow_run = db.session.scalar(select(WorkflowRun).where(WorkflowRun.id == message.workflow_run_id))
             if not workflow_run:
                 raise ValueError(f"Workflow run not found: {message.workflow_run_id}")
             workflow = db.session.scalar(select(Workflow).where(Workflow.id == workflow_run.workflow_id))

api/core/ops/langfuse_trace/langfuse_trace.py

@@ -2,7 +2,7 @@ import logging
 import os
 from datetime import datetime, timedelta
 
-from langfuse import Langfuse
+from langfuse import Langfuse  # type: ignore
 from sqlalchemy.orm import sessionmaker
 
 from core.ops.base_trace_instance import BaseTraceInstance

api/core/ops/ops_trace_manager.py

@@ -12,7 +12,7 @@ from uuid import UUID, uuid4
 from cachetools import LRUCache
 from flask import current_app
 from sqlalchemy import select
-from sqlalchemy.orm import Session, sessionmaker
+from sqlalchemy.orm import Session
 
 from core.helper.encrypter import decrypt_token, encrypt_token, obfuscated_token
 from core.ops.entities.config_entity import (
@@ -34,8 +34,7 @@ from core.ops.utils import get_message_data
 from extensions.ext_database import db
 from extensions.ext_storage import storage
 from models.model import App, AppModelConfig, Conversation, Message, MessageFile, TraceAppConfig
-from models.workflow import WorkflowAppLog
-from repositories.factory import DifyAPIRepositoryFactory
+from models.workflow import WorkflowAppLog, WorkflowRun
 from tasks.ops_trace_task import process_trace_tasks
 
 if TYPE_CHECKING:
@@ -420,18 +419,6 @@ class OpsTraceManager:
 
 
 class TraceTask:
-    _workflow_run_repo = None
-    _repo_lock = threading.Lock()
-
-    @classmethod
-    def _get_workflow_run_repo(cls):
-        if cls._workflow_run_repo is None:
-            with cls._repo_lock:
-                if cls._workflow_run_repo is None:
-                    session_maker = sessionmaker(bind=db.engine, expire_on_commit=False)
-                    cls._workflow_run_repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker)
-        return cls._workflow_run_repo
-
     def __init__(
         self,
         trace_type: Any,
@@ -499,27 +486,27 @@ class TraceTask:
         if not workflow_run_id:
             return {}
 
-        workflow_run_repo = self._get_workflow_run_repo()
-        workflow_run = workflow_run_repo.get_workflow_run_by_id_without_tenant(run_id=workflow_run_id)
-        if not workflow_run:
-            raise ValueError("Workflow run not found")
-
-        workflow_id = workflow_run.workflow_id
-        tenant_id = workflow_run.tenant_id
-        workflow_run_id = workflow_run.id
-        workflow_run_elapsed_time = workflow_run.elapsed_time
-        workflow_run_status = workflow_run.status
-        workflow_run_inputs = workflow_run.inputs_dict
-        workflow_run_outputs = workflow_run.outputs_dict
-        workflow_run_version = workflow_run.version
-        error = workflow_run.error or ""
-
-        total_tokens = workflow_run.total_tokens
-
-        file_list = workflow_run_inputs.get("sys.file") or []
-        query = workflow_run_inputs.get("query") or workflow_run_inputs.get("sys.query") or ""
-
         with Session(db.engine) as session:
+            workflow_run_stmt = select(WorkflowRun).where(WorkflowRun.id == workflow_run_id)
+            workflow_run = session.scalars(workflow_run_stmt).first()
+            if not workflow_run:
+                raise ValueError("Workflow run not found")
+
+            workflow_id = workflow_run.workflow_id
+            tenant_id = workflow_run.tenant_id
+            workflow_run_id = workflow_run.id
+            workflow_run_elapsed_time = workflow_run.elapsed_time
+            workflow_run_status = workflow_run.status
+            workflow_run_inputs = workflow_run.inputs_dict
+            workflow_run_outputs = workflow_run.outputs_dict
+            workflow_run_version = workflow_run.version
+            error = workflow_run.error or ""
+
+            total_tokens = workflow_run.total_tokens
+
+            file_list = workflow_run_inputs.get("sys.file") or []
+            query = workflow_run_inputs.get("query") or workflow_run_inputs.get("sys.query") or ""
+
             # get workflow_app_log_id
             workflow_app_log_data_stmt = select(WorkflowAppLog.id).where(
                 WorkflowAppLog.tenant_id == tenant_id,
@@ -536,43 +523,43 @@ class TraceTask:
                 )
                 message_id = session.scalar(message_data_stmt)
 
-        metadata = {
-            "workflow_id": workflow_id,
-            "conversation_id": conversation_id,
-            "workflow_run_id": workflow_run_id,
-            "tenant_id": tenant_id,
-            "elapsed_time": workflow_run_elapsed_time,
-            "status": workflow_run_status,
-            "version": workflow_run_version,
-            "total_tokens": total_tokens,
-            "file_list": file_list,
-            "triggered_from": workflow_run.triggered_from,
-            "user_id": user_id,
-            "app_id": workflow_run.app_id,
-        }
+            metadata = {
+                "workflow_id": workflow_id,
+                "conversation_id": conversation_id,
+                "workflow_run_id": workflow_run_id,
+                "tenant_id": tenant_id,
+                "elapsed_time": workflow_run_elapsed_time,
+                "status": workflow_run_status,
+                "version": workflow_run_version,
+                "total_tokens": total_tokens,
+                "file_list": file_list,
+                "triggered_from": workflow_run.triggered_from,
+                "user_id": user_id,
+                "app_id": workflow_run.app_id,
+            }
 
-        workflow_trace_info = WorkflowTraceInfo(
-            trace_id=self.trace_id,
-            workflow_data=workflow_run.to_dict(),
-            conversation_id=conversation_id,
-            workflow_id=workflow_id,
-            tenant_id=tenant_id,
-            workflow_run_id=workflow_run_id,
-            workflow_run_elapsed_time=workflow_run_elapsed_time,
-            workflow_run_status=workflow_run_status,
-            workflow_run_inputs=workflow_run_inputs,
-            workflow_run_outputs=workflow_run_outputs,
-            workflow_run_version=workflow_run_version,
-            error=error,
-            total_tokens=total_tokens,
-            file_list=file_list,
-            query=query,
-            metadata=metadata,
-            workflow_app_log_id=workflow_app_log_id,
-            message_id=message_id,
-            start_time=workflow_run.created_at,
-            end_time=workflow_run.finished_at,
-        )
+            workflow_trace_info = WorkflowTraceInfo(
+                trace_id=self.trace_id,
+                workflow_data=workflow_run.to_dict(),
+                conversation_id=conversation_id,
+                workflow_id=workflow_id,
+                tenant_id=tenant_id,
+                workflow_run_id=workflow_run_id,
+                workflow_run_elapsed_time=workflow_run_elapsed_time,
+                workflow_run_status=workflow_run_status,
+                workflow_run_inputs=workflow_run_inputs,
+                workflow_run_outputs=workflow_run_outputs,
+                workflow_run_version=workflow_run_version,
+                error=error,
+                total_tokens=total_tokens,
+                file_list=file_list,
+                query=query,
+                metadata=metadata,
+                workflow_app_log_id=workflow_app_log_id,
+                message_id=message_id,
+                start_time=workflow_run.created_at,
+                end_time=workflow_run.finished_at,
+            )
         return workflow_trace_info
 
     def message_trace(self, message_id: str | None):

api/core/plugin/entities/parameters.py

@@ -76,7 +76,7 @@ class PluginParameter(BaseModel):
     auto_generate: PluginParameterAutoGenerate | None = None
     template: PluginParameterTemplate | None = None
     required: bool = False
-    default: Union[float, int, str, bool] | None = None
+    default: Union[float, int, str] | None = None
     min: Union[float, int] | None = None
     max: Union[float, int] | None = None
     precision: int | None = None

api/core/plugin/impl/base.py

@@ -180,7 +180,7 @@ class BasePluginClient:
         Make a request to the plugin daemon inner API and return the response as a model.
         """
         response = self._request(method, path, headers, data, params, files)
-        return type_(**response.json())  # type: ignore[return-value]
+        return type_(**response.json())  # type: ignore
 
     def _request_with_plugin_daemon_response(
         self,

api/core/plugin/impl/exc.py

@@ -40,7 +40,7 @@ class PluginDaemonBadRequestError(PluginDaemonClientSideError):
     description: str = "Bad Request"
 
 
-class PluginInvokeError(PluginDaemonClientSideError, ValueError):
+class PluginInvokeError(PluginDaemonClientSideError):
     description: str = "Invoke Error"
 
     def _get_error_object(self) -> Mapping:

api/core/rag/datasource/vdb/opensearch/opensearch_vector.py

@@ -161,7 +161,7 @@ class OpenSearchVector(BaseVector):
                         logger.exception("Error deleting document: %s", error)
 
     def delete(self):
-        self._client.indices.delete(index=self._collection_name.lower(), ignore_unavailable=True)
+        self._client.indices.delete(index=self._collection_name.lower())
 
     def text_exists(self, id: str) -> bool:
         try:

api/core/rag/datasource/vdb/oracle/oraclevector.py

@@ -289,7 +289,8 @@ class OracleVector(BaseVector):
                 words = pseg.cut(query)
                 current_entity = ""
                 for word, pos in words:
-                    if pos in {"nr", "Ng", "eng", "nz", "n", "ORG", "v"}:  # nr: 人名，ns: 地名，nt: 机构名
+                    # nr: person name, ns: place name, nt: organization name
+                    if pos in {"nr", "Ng", "eng", "nz", "n", "ORG", "v"}:
                         current_entity += word
                     else:
                         if current_entity:

api/core/rag/datasource/vdb/pyvastbase/vastbase_vector.py

@@ -213,7 +213,7 @@ class VastbaseVector(BaseVector):
 
             with self._get_cursor() as cur:
                 cur.execute(SQL_CREATE_TABLE.format(table_name=self.table_name, dimension=dimension))
-                # Vastbase 支持的向量维度取值范围为 [1,16000]
+                # Vastbase supports vector dimensions in range [1, 16000]
                 if dimension <= 16000:
                     cur.execute(SQL_CREATE_INDEX.format(table_name=self.table_name))
             redis_client.set(collection_exist_cache_key, 1, ex=3600)

api/core/rag/retrieval/dataset_retrieval.py

@@ -72,19 +72,6 @@ default_retrieval_model: dict[str, Any] = {
 class DatasetRetrieval:
     def __init__(self, application_generate_entity=None):
         self.application_generate_entity = application_generate_entity
-        self._llm_usage = LLMUsage.empty_usage()
-
-    @property
-    def llm_usage(self) -> LLMUsage:
-        return self._llm_usage.model_copy()
-
-    def _record_usage(self, usage: LLMUsage | None) -> None:
-        if usage is None or usage.total_tokens <= 0:
-            return
-        if self._llm_usage.total_tokens == 0:
-            self._llm_usage = usage
-        else:
-            self._llm_usage = self._llm_usage.plus(usage)
 
     def retrieve(
         self,
@@ -325,18 +312,15 @@ class DatasetRetrieval:
             )
             tools.append(message_tool)
         dataset_id = None
-        router_usage = LLMUsage.empty_usage()
         if planning_strategy == PlanningStrategy.REACT_ROUTER:
             react_multi_dataset_router = ReactMultiDatasetRouter()
-            dataset_id, router_usage = react_multi_dataset_router.invoke(
+            dataset_id = react_multi_dataset_router.invoke(
                 query, tools, model_config, model_instance, user_id, tenant_id
             )
 
         elif planning_strategy == PlanningStrategy.ROUTER:
             function_call_router = FunctionCallMultiDatasetRouter()
-            dataset_id, router_usage = function_call_router.invoke(query, tools, model_config, model_instance)
-
-        self._record_usage(router_usage)
+            dataset_id = function_call_router.invoke(query, tools, model_config, model_instance)
 
         if dataset_id:
             # get retrieval model config
@@ -999,8 +983,7 @@ class DatasetRetrieval:
             )
 
             # handle invoke result
-            result_text, usage = self._handle_invoke_result(invoke_result=invoke_result)
-            self._record_usage(usage)
+            result_text, _ = self._handle_invoke_result(invoke_result=invoke_result)
 
             result_text_json = parse_and_check_json_markdown(result_text, [])
             automatic_metadata_filters = []

api/core/rag/retrieval/router/multi_dataset_function_call_router.py

@@ -2,7 +2,7 @@ from typing import Union
 
 from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
 from core.model_manager import ModelInstance
-from core.model_runtime.entities.llm_entities import LLMResult, LLMUsage
+from core.model_runtime.entities.llm_entities import LLMResult
 from core.model_runtime.entities.message_entities import PromptMessageTool, SystemPromptMessage, UserPromptMessage
 
 
@@ -13,15 +13,15 @@ class FunctionCallMultiDatasetRouter:
         dataset_tools: list[PromptMessageTool],
         model_config: ModelConfigWithCredentialsEntity,
         model_instance: ModelInstance,
-    ) -> tuple[Union[str, None], LLMUsage]:
+    ) -> Union[str, None]:
         """Given input, decided what to do.
         Returns:
             Action specifying what tool to use.
         """
         if len(dataset_tools) == 0:
-            return None, LLMUsage.empty_usage()
+            return None
         elif len(dataset_tools) == 1:
-            return dataset_tools[0].name, LLMUsage.empty_usage()
+            return dataset_tools[0].name
 
         try:
             prompt_messages = [
@@ -34,10 +34,9 @@ class FunctionCallMultiDatasetRouter:
                 stream=False,
                 model_parameters={"temperature": 0.2, "top_p": 0.3, "max_tokens": 1500},
             )
-            usage = result.usage or LLMUsage.empty_usage()
             if result.message.tool_calls:
                 # get retrieval model config
-                return result.message.tool_calls[0].function.name, usage
-            return None, usage
+                return result.message.tool_calls[0].function.name
+            return None
         except Exception:
-            return None, LLMUsage.empty_usage()
+            return None

api/core/rag/retrieval/router/multi_dataset_react_route.py

@@ -58,15 +58,15 @@ class ReactMultiDatasetRouter:
         model_instance: ModelInstance,
         user_id: str,
         tenant_id: str,
-    ) -> tuple[Union[str, None], LLMUsage]:
+    ) -> Union[str, None]:
         """Given input, decided what to do.
         Returns:
             Action specifying what tool to use.
         """
         if len(dataset_tools) == 0:
-            return None, LLMUsage.empty_usage()
+            return None
         elif len(dataset_tools) == 1:
-            return dataset_tools[0].name, LLMUsage.empty_usage()
+            return dataset_tools[0].name
 
         try:
             return self._react_invoke(
@@ -78,7 +78,7 @@ class ReactMultiDatasetRouter:
                 tenant_id=tenant_id,
             )
         except Exception:
-            return None, LLMUsage.empty_usage()
+            return None
 
     def _react_invoke(
         self,
@@ -91,7 +91,7 @@ class ReactMultiDatasetRouter:
         prefix: str = PREFIX,
         suffix: str = SUFFIX,
         format_instructions: str = FORMAT_INSTRUCTIONS,
-    ) -> tuple[Union[str, None], LLMUsage]:
+    ) -> Union[str, None]:
         prompt: Union[list[ChatModelMessage], CompletionModelPromptTemplate]
         if model_config.mode == "chat":
             prompt = self.create_chat_prompt(
@@ -120,7 +120,7 @@ class ReactMultiDatasetRouter:
             memory=None,
             model_config=model_config,
         )
-        result_text, usage = self._invoke_llm(
+        result_text, _ = self._invoke_llm(
             completion_param=model_config.parameters,
             model_instance=model_instance,
             prompt_messages=prompt_messages,
@@ -131,8 +131,8 @@ class ReactMultiDatasetRouter:
         output_parser = StructuredChatOutputParser()
         react_decision = output_parser.parse(result_text)
         if isinstance(react_decision, ReactAction):
-            return react_decision.tool, usage
-        return None, usage
+            return react_decision.tool
+        return None
 
     def _invoke_llm(
         self,

api/core/repositories/celery_workflow_execution_repository.py

@@ -74,7 +74,7 @@ class CeleryWorkflowExecutionRepository(WorkflowExecutionRepository):
         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
+        self._tenant_id = tenant_id  # type: ignore[assignment]  # We've already checked tenant_id is not None
 
         # Store app context
         self._app_id = app_id

api/core/repositories/celery_workflow_node_execution_repository.py

@@ -81,7 +81,7 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository):
         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
+        self._tenant_id = tenant_id  # type: ignore[assignment]  # We've already checked tenant_id is not None
 
         # Store app context
         self._app_id = app_id

api/core/repositories/factory.py

@@ -60,7 +60,7 @@ class DifyCoreRepositoryFactory:
 
         try:
             repository_class = import_string(class_path)
-            return repository_class(
+            return repository_class(  # type: ignore[no-any-return]
                 session_factory=session_factory,
                 user=user,
                 app_id=app_id,
@@ -96,7 +96,7 @@ class DifyCoreRepositoryFactory:
 
         try:
             repository_class = import_string(class_path)
-            return repository_class(
+            return repository_class(  # type: ignore[no-any-return]
                 session_factory=session_factory,
                 user=user,
                 app_id=app_id,

api/core/tools/__base/tool.py

@@ -217,16 +217,3 @@ class Tool(ABC):
         return ToolInvokeMessage(
             type=ToolInvokeMessage.MessageType.JSON, message=ToolInvokeMessage.JsonMessage(json_object=object)
         )
-
-    def create_variable_message(
-        self, variable_name: str, variable_value: Any, stream: bool = False
-    ) -> ToolInvokeMessage:
-        """
-        create a variable message
-        """
-        return ToolInvokeMessage(
-            type=ToolInvokeMessage.MessageType.VARIABLE,
-            message=ToolInvokeMessage.VariableMessage(
-                variable_name=variable_name, variable_value=variable_value, stream=stream
-            ),
-        )

api/core/tools/builtin_tool/provider.py

@@ -157,7 +157,7 @@ class BuiltinToolProviderController(ToolProviderController):
         """
         returns the tool that the provider can provide
         """
-        return next(filter(lambda x: x.entity.identity.name == tool_name, self.get_tools()), None)
+        return next(filter(lambda x: x.entity.identity.name == tool_name, self.get_tools()), None)  # type: ignore
 
     @property
     def need_credentials(self) -> bool:

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

@@ -43,7 +43,7 @@ class TTSTool(BuiltinTool):
             content_text=tool_parameters.get("text"),  # type: ignore
             user=user_id,
             tenant_id=self.runtime.tenant_id,
-            voice=voice,
+            voice=voice,  # type: ignore
         )
         buffer = io.BytesIO()
         for chunk in tts:

api/core/tools/builtin_tool/providers/time/tools/localtime_to_timestamp.py

@@ -34,7 +34,6 @@ class LocaltimeToTimestampTool(BuiltinTool):
 
         yield self.create_text_message(f"{timestamp}")
 
-    # TODO: this method's type is messy
     @staticmethod
     def localtime_to_timestamp(localtime: str, time_format: str, local_tz=None) -> int | None:
         try:

api/core/tools/builtin_tool/providers/time/tools/timezone_conversion.py

@@ -48,6 +48,6 @@ class TimezoneConversionTool(BuiltinTool):
             datetime_with_tz = input_timezone.localize(local_time)
             # timezone convert
             converted_datetime = datetime_with_tz.astimezone(output_timezone)
-            return converted_datetime.strftime(time_format)
+            return converted_datetime.strftime(format=time_format)  # type: ignore
         except Exception as e:
             raise ToolInvokeError(str(e))

api/core/tools/entities/api_entities.py

@@ -4,7 +4,6 @@ from typing import Any, Literal
 
 from pydantic import BaseModel, Field, field_validator
 
-from core.entities.mcp_provider import MCPAuthentication, MCPConfiguration
 from core.model_runtime.utils.encoders import jsonable_encoder
 from core.tools.__base.tool import ToolParameter
 from core.tools.entities.common_entities import I18nObject
@@ -45,14 +44,10 @@ class ToolProviderApiEntity(BaseModel):
     server_url: str | None = Field(default="", description="The server url of the tool")
     updated_at: int = Field(default_factory=lambda: int(datetime.now().timestamp()))
     server_identifier: str | None = Field(default="", description="The server identifier of the MCP tool")
-
+    timeout: float | None = Field(default=30.0, description="The timeout of the MCP tool")
+    sse_read_timeout: float | None = Field(default=300.0, description="The SSE read timeout of the MCP tool")
     masked_headers: dict[str, str] | None = Field(default=None, description="The masked headers of the MCP tool")
     original_headers: dict[str, str] | None = Field(default=None, description="The original headers of the MCP tool")
-    authentication: MCPAuthentication | None = Field(default=None, description="The OAuth config of the MCP tool")
-    is_dynamic_registration: bool = Field(default=True, description="Whether the MCP tool is dynamically registered")
-    configuration: MCPConfiguration | None = Field(
-        default=None, description="The timeout and sse_read_timeout of the MCP tool"
-    )
 
     @field_validator("tools", mode="before")
     @classmethod
@@ -75,15 +70,8 @@ class ToolProviderApiEntity(BaseModel):
         if self.type == ToolProviderType.MCP:
             optional_fields.update(self.optional_field("updated_at", self.updated_at))
             optional_fields.update(self.optional_field("server_identifier", self.server_identifier))
-            optional_fields.update(
-                self.optional_field(
-                    "configuration", self.configuration.model_dump() if self.configuration else MCPConfiguration()
-                )
-            )
-            optional_fields.update(
-                self.optional_field("authentication", self.authentication.model_dump() if self.authentication else None)
-            )
-            optional_fields.update(self.optional_field("is_dynamic_registration", self.is_dynamic_registration))
+            optional_fields.update(self.optional_field("timeout", self.timeout))
+            optional_fields.update(self.optional_field("sse_read_timeout", self.sse_read_timeout))
             optional_fields.update(self.optional_field("masked_headers", self.masked_headers))
             optional_fields.update(self.optional_field("original_headers", self.original_headers))
         return {

api/core/tools/mcp_tool/provider.py

@@ -1,6 +1,6 @@
+import json
 from typing import Any, Self
 
-from core.entities.mcp_provider import MCPProviderEntity
 from core.mcp.types import Tool as RemoteMCPTool
 from core.tools.__base.tool_provider import ToolProviderController
 from core.tools.__base.tool_runtime import ToolRuntime
@@ -52,25 +52,18 @@ class MCPToolProviderController(ToolProviderController):
         """
         from db provider
         """
-        # Convert to entity first
-        provider_entity = db_provider.to_entity()
-        return cls.from_entity(provider_entity)
-
-    @classmethod
-    def from_entity(cls, entity: MCPProviderEntity) -> Self:
-        """
-        create a MCPToolProviderController from a MCPProviderEntity
-        """
-        remote_mcp_tools = [RemoteMCPTool(**tool) for tool in entity.tools]
-
+        tools = []
+        tools_data = json.loads(db_provider.tools)
+        remote_mcp_tools = [RemoteMCPTool.model_validate(tool) for tool in tools_data]
+        user = db_provider.load_user()
         tools = [
             ToolEntity(
                 identity=ToolIdentity(
-                    author="Anonymous",  # Tool level author is not stored
+                    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=entity.provider_id,
-                    icon=entity.icon if isinstance(entity.icon, str) else "",
+                    provider=db_provider.server_identifier,
+                    icon=db_provider.icon,
                 ),
                 parameters=ToolTransformService.convert_mcp_schema_to_parameter(remote_mcp_tool.inputSchema),
                 description=ToolDescription(
@@ -79,32 +72,31 @@ class MCPToolProviderController(ToolProviderController):
                     ),
                     llm=remote_mcp_tool.description or "",
                 ),
-                output_schema=remote_mcp_tool.outputSchema or {},
                 has_runtime_parameters=len(remote_mcp_tool.inputSchema) > 0,
             )
             for remote_mcp_tool in remote_mcp_tools
         ]
-        if not entity.icon:
+        if not db_provider.icon:
             raise ValueError("Database provider icon is required")
         return cls(
             entity=ToolProviderEntityWithPlugin(
                 identity=ToolProviderIdentity(
-                    author="Anonymous",  # Provider level author is not stored in entity
-                    name=entity.name,
-                    label=I18nObject(en_US=entity.name, zh_Hans=entity.name),
+                    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=entity.icon if isinstance(entity.icon, str) else "",
+                    icon=db_provider.icon,
                 ),
                 plugin_id=None,
                 credentials_schema=[],
                 tools=tools,
             ),
-            provider_id=entity.provider_id,
-            tenant_id=entity.tenant_id,
-            server_url=entity.server_url,
-            headers=entity.headers,
-            timeout=entity.timeout,
-            sse_read_timeout=entity.sse_read_timeout,
+            provider_id=db_provider.server_identifier or "",
+            tenant_id=db_provider.tenant_id or "",
+            server_url=db_provider.decrypted_server_url,
+            headers=db_provider.decrypted_headers or {},
+            timeout=db_provider.timeout,
+            sse_read_timeout=db_provider.sse_read_timeout,
         )
 
     def _validate_credentials(self, user_id: str, credentials: dict[str, Any]):
@@ -113,7 +105,7 @@ class MCPToolProviderController(ToolProviderController):
         """
         pass
 
-    def get_tool(self, tool_name: str) -> MCPTool:
+    def get_tool(self, tool_name: str) -> MCPTool:  # type: ignore
         """
         return tool with given name
         """
@@ -136,7 +128,7 @@ class MCPToolProviderController(ToolProviderController):
             sse_read_timeout=self.sse_read_timeout,
         )
 
-    def get_tools(self) -> list[MCPTool]:
+    def get_tools(self) -> list[MCPTool]:  # type: ignore
         """
         get all tools
         """

api/core/tools/mcp_tool/tool.py

@@ -3,13 +3,12 @@ import json
 from collections.abc import Generator
 from typing import Any
 
-from core.mcp.auth_client import MCPClientWithAuthRetry
-from core.mcp.error import MCPConnectionError
-from core.mcp.types import CallToolResult, ImageContent, TextContent
+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, ToolProviderType
-from core.tools.errors import ToolInvokeError
 
 
 class MCPTool(Tool):
@@ -45,32 +44,40 @@ class MCPTool(Tool):
         app_id: str | None = None,
         message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
-        result = self.invoke_remote_mcp_tool(tool_parameters)
-        # handle dify tool output
+        from core.tools.errors import ToolInvokeError
+
+        try:
+            with MCPClient(
+                self.server_url,
+                self.provider_id,
+                self.tenant_id,
+                authed=True,
+                headers=self.headers,
+                timeout=self.timeout,
+                sse_read_timeout=self.sse_read_timeout,
+            ) as mcp_client:
+                tool_parameters = self._handle_none_parameter(tool_parameters)
+                result = mcp_client.invoke_tool(tool_name=self.entity.identity.name, tool_args=tool_parameters)
+        except MCPAuthError as e:
+            raise ToolInvokeError("Please auth the tool first") from e
+        except MCPConnectionError as e:
+            raise ToolInvokeError(f"Failed to connect to MCP server: {e}") from e
+        except Exception as e:
+            raise ToolInvokeError(f"Failed to invoke tool: {e}") from e
+
         for content in result.content:
             if isinstance(content, TextContent):
                 yield from self._process_text_content(content)
             elif isinstance(content, ImageContent):
                 yield self._process_image_content(content)
-        # handle MCP structured output
-        if self.entity.output_schema and result.structuredContent:
-            for k, v in result.structuredContent.items():
-                yield self.create_variable_message(k, v)
 
     def _process_text_content(self, content: TextContent) -> Generator[ToolInvokeMessage, None, None]:
         """Process text content and yield appropriate messages."""
-        # Check if content looks like JSON before attempting to parse
-        text = content.text.strip()
-        if text and text[0] in ("{", "[") and text[-1] in ("}", "]"):
-            try:
-                content_json = json.loads(text)
-                yield from self._process_json_content(content_json)
-                return
-            except json.JSONDecodeError:
-                pass
-
-        # If not JSON or parsing failed, treat as plain text
-        yield self.create_text_message(content.text)
+        try:
+            content_json = json.loads(content.text)
+            yield from self._process_json_content(content_json)
+        except json.JSONDecodeError:
+            yield self.create_text_message(content.text)
 
     def _process_json_content(self, content_json: Any) -> Generator[ToolInvokeMessage, None, None]:
         """Process JSON content based on its type."""
@@ -119,44 +126,3 @@ class MCPTool(Tool):
             for key, value in parameter.items()
             if value is not None and not (isinstance(value, str) and value.strip() == "")
         }
-
-    def invoke_remote_mcp_tool(self, tool_parameters: dict[str, Any]) -> CallToolResult:
-        headers = self.headers.copy() if self.headers else {}
-        tool_parameters = self._handle_none_parameter(tool_parameters)
-
-        from sqlalchemy.orm import Session
-
-        from extensions.ext_database import db
-        from services.tools.mcp_tools_manage_service import MCPToolManageService
-
-        # Step 1: Load provider entity and credentials in a short-lived session
-        # This minimizes database connection hold time
-        with Session(db.engine, expire_on_commit=False) as session:
-            mcp_service = MCPToolManageService(session=session)
-            provider_entity = mcp_service.get_provider_entity(self.provider_id, self.tenant_id, by_server_id=True)
-
-            # Decrypt and prepare all credentials before closing session
-            server_url = provider_entity.decrypt_server_url()
-            headers = provider_entity.decrypt_headers()
-
-            # Try to get existing token and add to headers
-            if not headers:
-                tokens = provider_entity.retrieve_tokens()
-                if tokens and tokens.access_token:
-                    headers["Authorization"] = f"{tokens.token_type.capitalize()} {tokens.access_token}"
-
-        # Step 2: Session is now closed, perform network operations without holding database connection
-        # MCPClientWithAuthRetry will create a new session lazily only if auth retry is needed
-        try:
-            with MCPClientWithAuthRetry(
-                server_url=server_url,
-                headers=headers,
-                timeout=self.timeout,
-                sse_read_timeout=self.sse_read_timeout,
-                provider_entity=provider_entity,
-            ) as mcp_client:
-                return mcp_client.invoke_tool(tool_name=self.entity.identity.name, tool_args=tool_parameters)
-        except MCPConnectionError as e:
-            raise ToolInvokeError(f"Failed to connect to MCP server: {e}") from e
-        except Exception as e:
-            raise ToolInvokeError(f"Failed to invoke tool: {e}") from e

api/core/tools/tool_label_manager.py

@@ -26,7 +26,7 @@ class ToolLabelManager:
         labels = cls.filter_tool_labels(labels)
 
         if isinstance(controller, ApiToolProviderController | WorkflowToolProviderController):
-            provider_id = controller.provider_id
+            provider_id = controller.provider_id  # ty: ignore [unresolved-attribute]
         else:
             raise ValueError("Unsupported tool type")
 
@@ -51,7 +51,7 @@ class ToolLabelManager:
         Get tool labels
         """
         if isinstance(controller, ApiToolProviderController | WorkflowToolProviderController):
-            provider_id = controller.provider_id
+            provider_id = controller.provider_id  # ty: ignore [unresolved-attribute]
         elif isinstance(controller, BuiltinToolProviderController):
             return controller.tool_labels
         else:
@@ -85,7 +85,7 @@ class ToolLabelManager:
         provider_ids = []
         for controller in tool_providers:
             assert isinstance(controller, ApiToolProviderController | WorkflowToolProviderController)
-            provider_ids.append(controller.provider_id)
+            provider_ids.append(controller.provider_id)  # ty: ignore [unresolved-attribute]
 
         labels = db.session.scalars(select(ToolLabelBinding).where(ToolLabelBinding.tool_id.in_(provider_ids))).all()
 

api/core/tools/tool_manager.py

@@ -14,32 +14,17 @@ from sqlalchemy.orm import Session
 from yarl import URL
 
 import contexts
-from core.helper.provider_cache import ToolProviderCredentialsCache
-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.utils.uuid_utils import is_valid_uuid
-from core.tools.workflow_as_tool.provider import WorkflowToolProviderController
-from core.workflow.runtime.variable_pool import VariablePool
-from extensions.ext_database import db
-from models.provider_ids import ToolProviderID
-from services.enterprise.plugin_manager_service import PluginCredentialType
-from services.tools.mcp_tools_manage_service import MCPToolManageService
-
-if TYPE_CHECKING:
-    from core.workflow.nodes.tool.entities import ToolEntity
-
 from configs import dify_config
 from core.agent.entities import AgentToolEntity
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.helper.module_import_helper import load_single_subclass_from_source
 from core.helper.position_helper import is_filtered
+from core.helper.provider_cache import ToolProviderCredentialsCache
 from core.model_runtime.utils.encoders import jsonable_encoder
+from core.plugin.impl.tool import PluginToolManager
 from core.tools.__base.tool import Tool
+from core.tools.__base.tool_provider import ToolProviderController
+from core.tools.__base.tool_runtime import ToolRuntime
 from core.tools.builtin_tool.provider import BuiltinToolProviderController
 from core.tools.builtin_tool.providers._positions import BuiltinToolProviderSort
 from core.tools.builtin_tool.tool import BuiltinTool
@@ -55,11 +40,21 @@ from core.tools.entities.tool_entities import (
     ToolProviderType,
 )
 from core.tools.errors import ToolProviderNotFoundError
+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.tool_label_manager import ToolLabelManager
 from core.tools.utils.configuration import ToolParameterConfigurationManager
 from core.tools.utils.encryption import create_provider_encrypter, create_tool_provider_encrypter
+from core.tools.utils.uuid_utils import is_valid_uuid
+from core.tools.workflow_as_tool.provider import WorkflowToolProviderController
 from core.tools.workflow_as_tool.tool import WorkflowTool
-from models.tools import ApiToolProvider, BuiltinToolProvider, WorkflowToolProvider
+from extensions.ext_database import db
+from models.provider_ids import ToolProviderID
+from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider, WorkflowToolProvider
+from services.enterprise.plugin_manager_service import PluginCredentialType
+from services.tools.mcp_tools_manage_service import MCPToolManageService
 from services.tools.tools_transform_service import ToolTransformService
 
 if TYPE_CHECKING:
@@ -331,8 +326,7 @@ class ToolManager:
             workflow_provider_stmt = select(WorkflowToolProvider).where(
                 WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id
             )
-            with Session(db.engine, expire_on_commit=False) as session, session.begin():
-                workflow_provider = session.scalar(workflow_provider_stmt)
+            workflow_provider = db.session.scalar(workflow_provider_stmt)
 
             if workflow_provider is None:
                 raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found")
@@ -724,9 +718,7 @@ class ToolManager:
                     )
                     result_providers[f"workflow_provider.{user_provider.name}"] = user_provider
             if "mcp" in filters:
-                with Session(db.engine) as session:
-                    mcp_service = MCPToolManageService(session=session)
-                    mcp_providers = mcp_service.list_providers(tenant_id=tenant_id, for_list=True)
+                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
 
@@ -781,12 +773,17 @@ class ToolManager:
 
         :return: the provider controller, the credentials
         """
-        with Session(db.engine) as session:
-            mcp_service = MCPToolManageService(session=session)
-            try:
-                provider = mcp_service.get_provider(server_identifier=provider_id, tenant_id=tenant_id)
-            except ValueError:
-                raise ToolProviderNotFoundError(f"mcp provider {provider_id} not found")
+        provider: MCPToolProvider | None = (
+            db.session.query(MCPToolProvider)
+            .where(
+                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)
 
@@ -924,15 +921,16 @@ class ToolManager:
     @classmethod
     def generate_mcp_tool_icon_url(cls, tenant_id: str, provider_id: str) -> Mapping[str, str] | str:
         try:
-            with Session(db.engine) as session:
-                mcp_service = MCPToolManageService(session=session)
-                try:
-                    mcp_provider = mcp_service.get_provider_entity(
-                        provider_id=provider_id, tenant_id=tenant_id, by_server_id=True
-                    )
-                    return mcp_provider.provider_icon
-                except ValueError:
-                    raise ToolProviderNotFoundError(f"mcp provider {provider_id} not found")
+            mcp_provider: MCPToolProvider | None = (
+                db.session.query(MCPToolProvider)
+                .where(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"}
 

api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py

@@ -193,18 +193,18 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
                                 DatasetDocument.enabled == True,
                                 DatasetDocument.archived == False,
                             )
-                            document = db.session.scalar(dataset_document_stmt)
+                            document = db.session.scalar(dataset_document_stmt)  # type: ignore
                             if dataset and document:
                                 source = RetrievalSourceMetadata(
                                     dataset_id=dataset.id,
                                     dataset_name=dataset.name,
-                                    document_id=document.id,
-                                    document_name=document.name,
-                                    data_source_type=document.data_source_type,
+                                    document_id=document.id,  # type: ignore
+                                    document_name=document.name,  # type: ignore
+                                    data_source_type=document.data_source_type,  # type: ignore
                                     segment_id=segment.id,
                                     retriever_from=self.retriever_from,
                                     score=record.score or 0.0,
-                                    doc_metadata=document.doc_metadata,
+                                    doc_metadata=document.doc_metadata,  # type: ignore
                                 )
 
                                 if self.retriever_from == "dev":

api/core/tools/utils/parser.py

@@ -62,11 +62,6 @@ class ApiBasedToolSchemaParser:
                             root = root[ref]
                         interface["operation"]["parameters"][i] = root
                 for parameter in interface["operation"]["parameters"]:
-                    # Handle complex type defaults that are not supported by PluginParameter
-                    default_value = None
-                    if "schema" in parameter and "default" in parameter["schema"]:
-                        default_value = ApiBasedToolSchemaParser._sanitize_default_value(parameter["schema"]["default"])
-
                     tool_parameter = ToolParameter(
                         name=parameter["name"],
                         label=I18nObject(en_US=parameter["name"], zh_Hans=parameter["name"]),
@@ -77,7 +72,9 @@ class ApiBasedToolSchemaParser:
                         required=parameter.get("required", False),
                         form=ToolParameter.ToolParameterForm.LLM,
                         llm_description=parameter.get("description"),
-                        default=default_value,
+                        default=parameter["schema"]["default"]
+                        if "schema" in parameter and "default" in parameter["schema"]
+                        else None,
                         placeholder=I18nObject(
                             en_US=parameter.get("description", ""), zh_Hans=parameter.get("description", "")
                         ),
@@ -137,11 +134,6 @@ class ApiBasedToolSchemaParser:
                             required = body_schema.get("required", [])
                             properties = body_schema.get("properties", {})
                             for name, property in properties.items():
-                                # Handle complex type defaults that are not supported by PluginParameter
-                                default_value = ApiBasedToolSchemaParser._sanitize_default_value(
-                                    property.get("default", None)
-                                )
-
                                 tool = ToolParameter(
                                     name=name,
                                     label=I18nObject(en_US=name, zh_Hans=name),
@@ -152,11 +144,12 @@ class ApiBasedToolSchemaParser:
                                     required=name in required,
                                     form=ToolParameter.ToolParameterForm.LLM,
                                     llm_description=property.get("description", ""),
-                                    default=default_value,
+                                    default=property.get("default", None),
                                     placeholder=I18nObject(
                                         en_US=property.get("description", ""), zh_Hans=property.get("description", "")
                                     ),
                                 )
+
                                 # check if there is a type
                                 typ = ApiBasedToolSchemaParser._get_tool_parameter_type(property)
                                 if typ:
@@ -204,22 +197,6 @@ class ApiBasedToolSchemaParser:
 
         return bundles
 
-    @staticmethod
-    def _sanitize_default_value(value):
-        """
-        Sanitize default values for PluginParameter compatibility.
-        Complex types (list, dict) are converted to None to avoid validation errors.
-
-        Args:
-            value: The default value from OpenAPI schema
-
-        Returns:
-            None for complex types (list, dict), otherwise the original value
-        """
-        if isinstance(value, (list, dict)):
-            return None
-        return value
-
     @staticmethod
     def _get_tool_parameter_type(parameter: dict) -> ToolParameter.ToolParameterType | None:
         parameter = parameter or {}
@@ -240,11 +217,7 @@ class ApiBasedToolSchemaParser:
             return ToolParameter.ToolParameterType.STRING
         elif typ == "array":
             items = parameter.get("items") or parameter.get("schema", {}).get("items")
-            if items and items.get("format") == "binary":
-                return ToolParameter.ToolParameterType.FILES
-            else:
-                # For regular arrays, return ARRAY type instead of None
-                return ToolParameter.ToolParameterType.ARRAY
+            return ToolParameter.ToolParameterType.FILES if items and items.get("format") == "binary" else None
         else:
             return None
 

api/core/tools/utils/web_reader_tool.py

@@ -6,8 +6,8 @@ from typing import Any, cast
 from urllib.parse import unquote
 
 import chardet
-import cloudscraper
-from readabilipy import simple_json_from_html_string
+import cloudscraper  # type: ignore
+from readabilipy import simple_json_from_html_string  # type: ignore
 
 from core.helper import ssrf_proxy
 from core.rag.extractor import extract_processor
@@ -63,8 +63,8 @@ def get_url(url: str, user_agent: str | None = None) -> str:
         response = ssrf_proxy.get(url, headers=headers, follow_redirects=True, timeout=(120, 300))
     elif response.status_code == 403:
         scraper = cloudscraper.create_scraper()
-        scraper.perform_request = ssrf_proxy.make_request
-        response = scraper.get(url, headers=headers, timeout=(120, 300))
+        scraper.perform_request = ssrf_proxy.make_request  # type: ignore
+        response = scraper.get(url, headers=headers, follow_redirects=True, timeout=(120, 300))  # type: ignore
 
     if response.status_code != 200:
         return f"URL returned status code {response.status_code}."

api/core/tools/utils/yaml_utils.py

@@ -3,7 +3,7 @@ from functools import lru_cache
 from pathlib import Path
 from typing import Any
 
-import yaml
+import yaml  # type: ignore
 from yaml import YAMLError
 
 logger = logging.getLogger(__name__)

api/core/tools/workflow_as_tool/provider.py

@@ -1,7 +1,6 @@
 from collections.abc import Mapping
 
 from pydantic import Field
-from sqlalchemy.orm import Session
 
 from core.app.app_config.entities import VariableEntity, VariableEntityType
 from core.app.apps.workflow.app_config_manager import WorkflowAppConfigManager
@@ -21,7 +20,6 @@ from core.tools.entities.tool_entities import (
 from core.tools.utils.workflow_configuration_sync import WorkflowToolConfigurationUtils
 from core.tools.workflow_as_tool.tool import WorkflowTool
 from extensions.ext_database import db
-from models.account import Account
 from models.model import App, AppMode
 from models.tools import WorkflowToolProvider
 from models.workflow import Workflow
@@ -31,7 +29,6 @@ VARIABLE_TO_PARAMETER_TYPE_MAPPING = {
     VariableEntityType.PARAGRAPH: ToolParameter.ToolParameterType.STRING,
     VariableEntityType.SELECT: ToolParameter.ToolParameterType.SELECT,
     VariableEntityType.NUMBER: ToolParameter.ToolParameterType.NUMBER,
-    VariableEntityType.CHECKBOX: ToolParameter.ToolParameterType.BOOLEAN,
     VariableEntityType.FILE: ToolParameter.ToolParameterType.FILE,
     VariableEntityType.FILE_LIST: ToolParameter.ToolParameterType.FILES,
 }
@@ -47,34 +44,29 @@ class WorkflowToolProviderController(ToolProviderController):
 
     @classmethod
     def from_db(cls, db_provider: WorkflowToolProvider) -> "WorkflowToolProviderController":
-        with Session(db.engine, expire_on_commit=False) as session, session.begin():
-            provider = session.get(WorkflowToolProvider, db_provider.id) if db_provider.id else None
-            if not provider:
-                raise ValueError("workflow provider not found")
-            app = session.get(App, provider.app_id)
-            if not app:
-                raise ValueError("app not found")
+        app = db_provider.app
 
-            user = session.get(Account, provider.user_id) if provider.user_id else None
+        if not app:
+            raise ValueError("app not found")
 
-            controller = WorkflowToolProviderController(
-                entity=ToolProviderEntity(
-                    identity=ToolProviderIdentity(
-                        author=user.name if user else "",
-                        name=provider.label,
-                        label=I18nObject(en_US=provider.label, zh_Hans=provider.label),
-                        description=I18nObject(en_US=provider.description, zh_Hans=provider.description),
-                        icon=provider.icon,
-                    ),
-                    credentials_schema=[],
-                    plugin_id=None,
+        controller = WorkflowToolProviderController(
+            entity=ToolProviderEntity(
+                identity=ToolProviderIdentity(
+                    author=db_provider.user.name if db_provider.user_id and db_provider.user else "",
+                    name=db_provider.label,
+                    label=I18nObject(en_US=db_provider.label, zh_Hans=db_provider.label),
+                    description=I18nObject(en_US=db_provider.description, zh_Hans=db_provider.description),
+                    icon=db_provider.icon,
                 ),
-                provider_id=provider.id or "",
-            )
+                credentials_schema=[],
+                plugin_id=None,
+            ),
+            provider_id=db_provider.id or "",
+        )
 
-            controller.tools = [
-                controller._get_db_provider_tool(provider, app, session=session, user=user),
-            ]
+        # init tools
+
+        controller.tools = [controller._get_db_provider_tool(db_provider, app)]
 
         return controller
 
@@ -82,14 +74,7 @@ class WorkflowToolProviderController(ToolProviderController):
     def provider_type(self) -> ToolProviderType:
         return ToolProviderType.WORKFLOW
 
-    def _get_db_provider_tool(
-        self,
-        db_provider: WorkflowToolProvider,
-        app: App,
-        *,
-        session: Session,
-        user: Account | None = None,
-    ) -> WorkflowTool:
+    def _get_db_provider_tool(self, db_provider: WorkflowToolProvider, app: App) -> WorkflowTool:
         """
         get db provider tool
         :param db_provider: the db provider
@@ -97,7 +82,7 @@ class WorkflowToolProviderController(ToolProviderController):
         :return: the tool
         """
         workflow: Workflow | None = (
-            session.query(Workflow)
+            db.session.query(Workflow)
             .where(Workflow.app_id == db_provider.app_id, Workflow.version == db_provider.version)
             .first()
         )
@@ -114,7 +99,9 @@ class WorkflowToolProviderController(ToolProviderController):
         variables = WorkflowToolConfigurationUtils.get_workflow_graph_variables(graph)
 
         def fetch_workflow_variable(variable_name: str) -> VariableEntity | None:
-            return next(filter(lambda x: x.variable == variable_name, variables), None)
+            return next(filter(lambda x: x.variable == variable_name, variables), None)  # type: ignore
+
+        user = db_provider.user
 
         workflow_tool_parameters = []
         for parameter in parameters:
@@ -200,25 +187,22 @@ class WorkflowToolProviderController(ToolProviderController):
         if self.tools is not None:
             return self.tools
 
-        with Session(db.engine, expire_on_commit=False) as session, session.begin():
-            db_provider: WorkflowToolProvider | None = (
-                session.query(WorkflowToolProvider)
-                .where(
-                    WorkflowToolProvider.tenant_id == tenant_id,
-                    WorkflowToolProvider.app_id == self.provider_id,
-                )
-                .first()
+        db_providers: WorkflowToolProvider | None = (
+            db.session.query(WorkflowToolProvider)
+            .where(
+                WorkflowToolProvider.tenant_id == tenant_id,
+                WorkflowToolProvider.app_id == self.provider_id,
             )
+            .first()
+        )
 
-            if not db_provider:
-                return []
+        if not db_providers:
+            return []
+        if not db_providers.app:
+            raise ValueError("app not found")
 
-            app = session.get(App, db_provider.app_id)
-            if not app:
-                raise ValueError("app not found")
-
-            user = session.get(Account, db_provider.user_id) if db_provider.user_id else None
-            self.tools = [self._get_db_provider_tool(db_provider, app, session=session, user=user)]
+        app = db_providers.app
+        self.tools = [self._get_db_provider_tool(db_providers, app)]
 
         return self.tools
 

api/core/tools/workflow_as_tool/tool.py

@@ -1,14 +1,12 @@
 import json
 import logging
-from collections.abc import Generator, Mapping, Sequence
-from typing import Any, cast
+from collections.abc import Generator
+from typing import Any
 
 from flask import has_request_context
 from sqlalchemy import select
-from sqlalchemy.orm import Session
 
 from core.file import FILE_MODEL_IDENTITY, File, FileTransferMethod
-from core.model_runtime.entities.llm_entities import LLMUsage, LLMUsageMetadata
 from core.tools.__base.tool import Tool
 from core.tools.__base.tool_runtime import ToolRuntime
 from core.tools.entities.tool_entities import (
@@ -50,7 +48,6 @@ class WorkflowTool(Tool):
         self.workflow_entities = workflow_entities
         self.workflow_call_depth = workflow_call_depth
         self.label = label
-        self._latest_usage = LLMUsage.empty_usage()
 
         super().__init__(entity=entity, runtime=runtime)
 
@@ -86,11 +83,10 @@ class WorkflowTool(Tool):
         assert self.runtime.invoke_from is not None
 
         user = self._resolve_user(user_id=user_id)
+
         if user is None:
             raise ToolInvokeError("User not found")
 
-        self._latest_usage = LLMUsage.empty_usage()
-
         result = generator.generate(
             app_model=app,
             workflow=workflow,
@@ -114,68 +110,9 @@ class WorkflowTool(Tool):
             for file in files:
                 yield self.create_file_message(file)  # type: ignore
 
-        self._latest_usage = self._derive_usage_from_result(data)
-
         yield self.create_text_message(json.dumps(outputs, ensure_ascii=False))
         yield self.create_json_message(outputs)
 
-    @property
-    def latest_usage(self) -> LLMUsage:
-        return self._latest_usage
-
-    @classmethod
-    def _derive_usage_from_result(cls, data: Mapping[str, Any]) -> LLMUsage:
-        usage_dict = cls._extract_usage_dict(data)
-        if usage_dict is not None:
-            return LLMUsage.from_metadata(cast(LLMUsageMetadata, dict(usage_dict)))
-
-        total_tokens = data.get("total_tokens")
-        total_price = data.get("total_price")
-        if total_tokens is None and total_price is None:
-            return LLMUsage.empty_usage()
-
-        usage_metadata: dict[str, Any] = {}
-        if total_tokens is not None:
-            try:
-                usage_metadata["total_tokens"] = int(str(total_tokens))
-            except (TypeError, ValueError):
-                pass
-        if total_price is not None:
-            usage_metadata["total_price"] = str(total_price)
-        currency = data.get("currency")
-        if currency is not None:
-            usage_metadata["currency"] = currency
-
-        if not usage_metadata:
-            return LLMUsage.empty_usage()
-
-        return LLMUsage.from_metadata(cast(LLMUsageMetadata, usage_metadata))
-
-    @classmethod
-    def _extract_usage_dict(cls, payload: Mapping[str, Any]) -> Mapping[str, Any] | None:
-        usage_candidate = payload.get("usage")
-        if isinstance(usage_candidate, Mapping):
-            return usage_candidate
-
-        metadata_candidate = payload.get("metadata")
-        if isinstance(metadata_candidate, Mapping):
-            usage_candidate = metadata_candidate.get("usage")
-            if isinstance(usage_candidate, Mapping):
-                return usage_candidate
-
-        for value in payload.values():
-            if isinstance(value, Mapping):
-                found = cls._extract_usage_dict(value)
-                if found is not None:
-                    return found
-            elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
-                for item in value:
-                    if isinstance(item, Mapping):
-                        found = cls._extract_usage_dict(item)
-                        if found is not None:
-                            return found
-        return None
-
     def fork_tool_runtime(self, runtime: ToolRuntime) -> "WorkflowTool":
         """
         fork a new tool with metadata
@@ -242,17 +179,16 @@ class WorkflowTool(Tool):
         """
         get the workflow by app id and version
         """
-        with Session(db.engine, expire_on_commit=False) as session, session.begin():
-            if not version:
-                stmt = (
-                    select(Workflow)
-                    .where(Workflow.app_id == app_id, Workflow.version != Workflow.VERSION_DRAFT)
-                    .order_by(Workflow.created_at.desc())
-                )
-                workflow = session.scalars(stmt).first()
-            else:
-                stmt = select(Workflow).where(Workflow.app_id == app_id, Workflow.version == version)
-                workflow = session.scalar(stmt)
+        if not version:
+            workflow = (
+                db.session.query(Workflow)
+                .where(Workflow.app_id == app_id, Workflow.version != Workflow.VERSION_DRAFT)
+                .order_by(Workflow.created_at.desc())
+                .first()
+            )
+        else:
+            stmt = select(Workflow).where(Workflow.app_id == app_id, Workflow.version == version)
+            workflow = db.session.scalar(stmt)
 
         if not workflow:
             raise ValueError("workflow not found or not published")
@@ -264,8 +200,7 @@ class WorkflowTool(Tool):
         get the app by app id
         """
         stmt = select(App).where(App.id == app_id)
-        with Session(db.engine, expire_on_commit=False) as session, session.begin():
-            app = session.scalar(stmt)
+        app = db.session.scalar(stmt)
         if not app:
             raise ValueError("app not found")
 

api/core/variables/segment_group.py

@@ -4,7 +4,7 @@ from .types import SegmentType
 
 class SegmentGroup(Segment):
     value_type: SegmentType = SegmentType.GROUP
-    value: list[Segment]
+    value: list[Segment] = None  # type: ignore
 
     @property
     def text(self):

api/core/variables/segments.py

@@ -19,7 +19,7 @@ class Segment(BaseModel):
     model_config = ConfigDict(frozen=True)
 
     value_type: SegmentType
-    value: Any
+    value: Any = None
 
     @field_validator("value_type")
     @classmethod
@@ -74,12 +74,12 @@ class NoneSegment(Segment):
 
 class StringSegment(Segment):
     value_type: SegmentType = SegmentType.STRING
-    value: str
+    value: str = None  # type: ignore
 
 
 class FloatSegment(Segment):
     value_type: SegmentType = SegmentType.FLOAT
-    value: float
+    value: float = None  # type: ignore
     # NOTE(QuantumGhost): seems that the equality for FloatSegment with `NaN` value has some problems.
     # The following tests cannot pass.
     #
@@ -98,12 +98,12 @@ class FloatSegment(Segment):
 
 class IntegerSegment(Segment):
     value_type: SegmentType = SegmentType.INTEGER
-    value: int
+    value: int = None  # type: ignore
 
 
 class ObjectSegment(Segment):
     value_type: SegmentType = SegmentType.OBJECT
-    value: Mapping[str, Any]
+    value: Mapping[str, Any] = None  # type: ignore
 
     @property
     def text(self) -> str:
@@ -136,7 +136,7 @@ class ArraySegment(Segment):
 
 class FileSegment(Segment):
     value_type: SegmentType = SegmentType.FILE
-    value: File
+    value: File = None  # type: ignore
 
     @property
     def markdown(self) -> str:
@@ -153,17 +153,17 @@ class FileSegment(Segment):
 
 class BooleanSegment(Segment):
     value_type: SegmentType = SegmentType.BOOLEAN
-    value: bool
+    value: bool = None  # type: ignore
 
 
 class ArrayAnySegment(ArraySegment):
     value_type: SegmentType = SegmentType.ARRAY_ANY
-    value: Sequence[Any]
+    value: Sequence[Any] = None  # type: ignore
 
 
 class ArrayStringSegment(ArraySegment):
     value_type: SegmentType = SegmentType.ARRAY_STRING
-    value: Sequence[str]
+    value: Sequence[str] = None  # type: ignore
 
     @property
     def text(self) -> str:
@@ -175,17 +175,17 @@ class ArrayStringSegment(ArraySegment):
 
 class ArrayNumberSegment(ArraySegment):
     value_type: SegmentType = SegmentType.ARRAY_NUMBER
-    value: Sequence[float | int]
+    value: Sequence[float | int] = None  # type: ignore
 
 
 class ArrayObjectSegment(ArraySegment):
     value_type: SegmentType = SegmentType.ARRAY_OBJECT
-    value: Sequence[Mapping[str, Any]]
+    value: Sequence[Mapping[str, Any]] = None  # type: ignore
 
 
 class ArrayFileSegment(ArraySegment):
     value_type: SegmentType = SegmentType.ARRAY_FILE
-    value: Sequence[File]
+    value: Sequence[File] = None  # type: ignore
 
     @property
     def markdown(self) -> str:
@@ -205,7 +205,7 @@ class ArrayFileSegment(ArraySegment):
 
 class ArrayBooleanSegment(ArraySegment):
     value_type: SegmentType = SegmentType.ARRAY_BOOLEAN
-    value: Sequence[bool]
+    value: Sequence[bool] = None  # type: ignore
 
 
 def get_segment_discriminator(v: Any) -> SegmentType | None:

api/core/workflow/entities/__init__.py

@@ -1,5 +1,3 @@
-from ..runtime.graph_runtime_state import GraphRuntimeState
-from ..runtime.variable_pool import VariablePool
 from .agent import AgentNodeStrategyInit
 from .graph_init_params import GraphInitParams
 from .workflow_execution import WorkflowExecution
@@ -8,8 +6,6 @@ from .workflow_node_execution import WorkflowNodeExecution
 __all__ = [
     "AgentNodeStrategyInit",
     "GraphInitParams",
-    "GraphRuntimeState",
-    "VariablePool",
     "WorkflowExecution",
     "WorkflowNodeExecution",
 ]

api/core/workflow/graph/graph.py

@@ -3,12 +3,11 @@ from collections import defaultdict
 from collections.abc import Mapping, Sequence
 from typing import Protocol, cast, final
 
-from core.workflow.enums import ErrorStrategy, NodeExecutionType, NodeState, NodeType
+from core.workflow.enums import NodeExecutionType, NodeState, NodeType
 from core.workflow.nodes.base.node import Node
 from libs.typing import is_str, is_str_dict
 
 from .edge import Edge
-from .validation import get_graph_validator
 
 logger = logging.getLogger(__name__)
 
@@ -202,17 +201,6 @@ class Graph:
 
         return GraphBuilder(graph_cls=cls)
 
-    @classmethod
-    def _promote_fail_branch_nodes(cls, nodes: dict[str, Node]) -> None:
-        """
-        Promote nodes configured with FAIL_BRANCH error strategy to branch execution type.
-
-        :param nodes: mapping of node ID to node instance
-        """
-        for node in nodes.values():
-            if node.error_strategy == ErrorStrategy.FAIL_BRANCH:
-                node.execution_type = NodeExecutionType.BRANCH
-
     @classmethod
     def _mark_inactive_root_branches(
         cls,
@@ -319,9 +307,6 @@ class Graph:
         # Create node instances
         nodes = cls._create_node_instances(node_configs_map, node_factory)
 
-        # Promote fail-branch nodes to branch execution type at graph level
-        cls._promote_fail_branch_nodes(nodes)
-
         # Get root node instance
         root_node = nodes[root_node_id]
 
@@ -329,7 +314,7 @@ class Graph:
         cls._mark_inactive_root_branches(nodes, edges, in_edges, out_edges, root_node_id)
 
         # Create and return the graph
-        graph = cls(
+        return cls(
             nodes=nodes,
             edges=edges,
             in_edges=in_edges,
@@ -337,11 +322,6 @@ class Graph:
             root_node=root_node,
         )
 
-        # Validate the graph structure using built-in validators
-        get_graph_validator().validate(graph)
-
-        return graph
-
     @property
     def node_ids(self) -> list[str]:
         """

api/core/workflow/graph/validation.py

@@ -1,125 +0,0 @@
-from __future__ import annotations
-
-from collections.abc import Sequence
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Protocol
-
-from core.workflow.enums import NodeExecutionType, NodeType
-
-if TYPE_CHECKING:
-    from .graph import Graph
-
-
-@dataclass(frozen=True, slots=True)
-class GraphValidationIssue:
-    """Immutable value object describing a single validation issue."""
-
-    code: str
-    message: str
-    node_id: str | None = None
-
-
-class GraphValidationError(ValueError):
-    """Raised when graph validation fails."""
-
-    def __init__(self, issues: Sequence[GraphValidationIssue]) -> None:
-        if not issues:
-            raise ValueError("GraphValidationError requires at least one issue.")
-        self.issues: tuple[GraphValidationIssue, ...] = tuple(issues)
-        message = "; ".join(f"[{issue.code}] {issue.message}" for issue in self.issues)
-        super().__init__(message)
-
-
-class GraphValidationRule(Protocol):
-    """Protocol that individual validation rules must satisfy."""
-
-    def validate(self, graph: Graph) -> Sequence[GraphValidationIssue]:
-        """Validate the provided graph and return any discovered issues."""
-        ...
-
-
-@dataclass(frozen=True, slots=True)
-class _EdgeEndpointValidator:
-    """Ensures all edges reference existing nodes."""
-
-    missing_node_code: str = "MISSING_NODE"
-
-    def validate(self, graph: Graph) -> Sequence[GraphValidationIssue]:
-        issues: list[GraphValidationIssue] = []
-        for edge in graph.edges.values():
-            if edge.tail not in graph.nodes:
-                issues.append(
-                    GraphValidationIssue(
-                        code=self.missing_node_code,
-                        message=f"Edge {edge.id} references unknown source node '{edge.tail}'.",
-                        node_id=edge.tail,
-                    )
-                )
-            if edge.head not in graph.nodes:
-                issues.append(
-                    GraphValidationIssue(
-                        code=self.missing_node_code,
-                        message=f"Edge {edge.id} references unknown target node '{edge.head}'.",
-                        node_id=edge.head,
-                    )
-                )
-        return issues
-
-
-@dataclass(frozen=True, slots=True)
-class _RootNodeValidator:
-    """Validates root node invariants."""
-
-    invalid_root_code: str = "INVALID_ROOT"
-    container_entry_types: tuple[NodeType, ...] = (NodeType.ITERATION_START, NodeType.LOOP_START)
-
-    def validate(self, graph: Graph) -> Sequence[GraphValidationIssue]:
-        root_node = graph.root_node
-        issues: list[GraphValidationIssue] = []
-        if root_node.id not in graph.nodes:
-            issues.append(
-                GraphValidationIssue(
-                    code=self.invalid_root_code,
-                    message=f"Root node '{root_node.id}' is missing from the node registry.",
-                    node_id=root_node.id,
-                )
-            )
-            return issues
-
-        node_type = getattr(root_node, "node_type", None)
-        if root_node.execution_type != NodeExecutionType.ROOT and node_type not in self.container_entry_types:
-            issues.append(
-                GraphValidationIssue(
-                    code=self.invalid_root_code,
-                    message=f"Root node '{root_node.id}' must declare execution type 'root'.",
-                    node_id=root_node.id,
-                )
-            )
-        return issues
-
-
-@dataclass(frozen=True, slots=True)
-class GraphValidator:
-    """Coordinates execution of graph validation rules."""
-
-    rules: tuple[GraphValidationRule, ...]
-
-    def validate(self, graph: Graph) -> None:
-        """Validate the graph against all configured rules."""
-        issues: list[GraphValidationIssue] = []
-        for rule in self.rules:
-            issues.extend(rule.validate(graph))
-
-        if issues:
-            raise GraphValidationError(issues)
-
-
-_DEFAULT_RULES: tuple[GraphValidationRule, ...] = (
-    _EdgeEndpointValidator(),
-    _RootNodeValidator(),
-)
-
-
-def get_graph_validator() -> GraphValidator:
-    """Construct the validator composed of default rules."""
-    return GraphValidator(_DEFAULT_RULES)

api/core/workflow/graph_engine/event_management/event_handlers.py

@@ -24,7 +24,6 @@ from core.workflow.graph_events import (
     NodeRunLoopStartedEvent,
     NodeRunLoopSucceededEvent,
     NodeRunPauseRequestedEvent,
-    NodeRunRetrieverResourceEvent,
     NodeRunRetryEvent,
     NodeRunStartedEvent,
     NodeRunStreamChunkEvent,
@@ -113,7 +112,6 @@ class EventHandler:
     @_dispatch.register(NodeRunLoopSucceededEvent)
     @_dispatch.register(NodeRunLoopFailedEvent)
     @_dispatch.register(NodeRunAgentLogEvent)
-    @_dispatch.register(NodeRunRetrieverResourceEvent)
     def _(self, event: GraphNodeEventBase) -> None:
         self._event_collector.collect(event)
 

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

@@ -26,8 +26,8 @@ class AgentNodeData(BaseNodeData):
 
 
 class ParamsAutoGenerated(IntEnum):
-    CLOSE = 0
-    OPEN = 1
+    CLOSE = auto()
+    OPEN = auto()
 
 
 class AgentOldVersionModelFeatures(StrEnum):

api/core/workflow/nodes/base/__init__.py

@@ -1,5 +1,4 @@
 from .entities import BaseIterationNodeData, BaseIterationState, BaseLoopNodeData, BaseLoopState, BaseNodeData
-from .usage_tracking_mixin import LLMUsageTrackingMixin
 
 __all__ = [
     "BaseIterationNodeData",
@@ -7,5 +6,4 @@ __all__ = [
     "BaseLoopNodeData",
     "BaseLoopState",
     "BaseNodeData",
-    "LLMUsageTrackingMixin",
 ]

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

@@ -1,6 +1,5 @@
 import json
 from abc import ABC
-from builtins import type as type_
 from collections.abc import Sequence
 from enum import StrEnum
 from typing import Any, Union
@@ -59,9 +58,10 @@ class DefaultValue(BaseModel):
             raise DefaultValueTypeError(f"Invalid JSON format for value: {value}")
 
     @staticmethod
-    def _validate_array(value: Any, element_type: type_ | tuple[type_, ...]) -> bool:
+    def _validate_array(value: Any, element_type: DefaultValueType) -> bool:
         """Unified array type validation"""
-        return isinstance(value, list) and all(isinstance(x, element_type) for x in value)
+        # FIXME, type ignore here for do not find the reason mypy complain, if find the root cause, please fix it
+        return isinstance(value, list) and all(isinstance(x, element_type) for x in value)  # type: ignore
 
     @staticmethod
     def _convert_number(value: str) -> float:

api/core/workflow/nodes/base/usage_tracking_mixin.py

@@ -1,28 +0,0 @@
-from core.model_runtime.entities.llm_entities import LLMUsage
-from core.workflow.runtime import GraphRuntimeState
-
-
-class LLMUsageTrackingMixin:
-    """Provides shared helpers for merging and recording LLM usage within workflow nodes."""
-
-    graph_runtime_state: GraphRuntimeState
-
-    @staticmethod
-    def _merge_usage(current: LLMUsage, new_usage: LLMUsage | None) -> LLMUsage:
-        """Return a combined usage snapshot, preserving zero-value inputs."""
-        if new_usage is None or new_usage.total_tokens <= 0:
-            return current
-        if current.total_tokens == 0:
-            return new_usage
-        return current.plus(new_usage)
-
-    def _accumulate_usage(self, usage: LLMUsage) -> None:
-        """Push usage into the graph runtime accumulator for downstream reporting."""
-        if usage.total_tokens <= 0:
-            return
-
-        current_usage = self.graph_runtime_state.llm_usage
-        if current_usage.total_tokens == 0:
-            self.graph_runtime_state.llm_usage = usage.model_copy()
-        else:
-            self.graph_runtime_state.llm_usage = current_usage.plus(usage)

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

@@ -10,10 +10,10 @@ from typing import Any
 import chardet
 import docx
 import pandas as pd
-import pypandoc
-import pypdfium2
-import webvtt
-import yaml
+import pypandoc  # type: ignore
+import pypdfium2  # type: ignore
+import webvtt  # type: ignore
+import yaml  # type: ignore
 from docx.document import Document
 from docx.oxml.table import CT_Tbl
 from docx.oxml.text.paragraph import CT_P
@@ -171,7 +171,6 @@ def _extract_text_by_file_extension(*, file_content: bytes, file_extension: str)
             ".txt"
             | ".markdown"
             | ".md"
-            | ".mdx"
             | ".html"
             | ".htm"
             | ".xml"

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

@@ -23,7 +23,6 @@ class IterationNodeData(BaseIterationNodeData):
     is_parallel: bool = False  # open the parallel mode or not
     parallel_nums: int = 10  # the numbers of parallel
     error_handle_mode: ErrorHandleMode = ErrorHandleMode.TERMINATED  # how to handle the error
-    flatten_output: bool = True  # whether to flatten the output array if all elements are lists
 
 
 class IterationStartNodeData(BaseNodeData):