Merge branch 'main' into jzh

.agents/skills/frontend-query-mutation/references/runtime-rules.md

@@ -64,7 +64,7 @@ export const useUpdateAccessMode = () => {
 
 // Component only adds UI behavior.
 updateAccessMode({ appId, mode }, {
-  onSuccess: () => Toast.notify({ type: 'success', message: '...' }),
+  onSuccess: () => toast.success('...'),
 })
 
 // Avoid putting invalidation knowledge in the component.
@@ -114,10 +114,7 @@ try {
   router.push(`/orders/${order.id}`)
 }
 catch (error) {
-  Toast.notify({
-    type: 'error',
-    message: error instanceof Error ? error.message : 'Unknown error',
-  })
+  toast.error(error instanceof Error ? error.message : 'Unknown error')
 }
 ```
 

.gitignore

@@ -212,7 +212,8 @@ api/.vscode
 
 # pnpm
 /.pnpm-store
-/node_modules
+node_modules
+.vite-hooks/_
 
 # plugin migrate
 plugins.jsonl

.npmrc

@@ -0,0 +1 @@
+save-exact=true

.vite-hooks/pre-commit

@@ -77,7 +77,7 @@ if $web_modified; then
     fi
 
     cd ./web || exit 1
-    lint-staged
+    vp staged
 
     if $web_ts_modified; then
         echo "Running TypeScript type-check:tsgo"
@@ -89,6 +89,12 @@ if $web_modified; then
         echo "No staged TypeScript changes detected, skipping type-check:tsgo"
     fi
 
+    echo "Running knip"
+    if ! pnpm run knip; then
+        echo "Knip check failed. Please run 'pnpm run knip' to fix the errors."
+        exit 1
+    fi
+
     echo "Running unit tests check"
     modified_files=$(git diff --cached --name-only -- utils | grep -v '\.spec\.ts$' || true)
 

api/.ruff.toml

@@ -115,12 +115,6 @@ ignore = [
 "controllers/console/human_input_form.py" = ["TID251"]
 "controllers/web/human_input_form.py" = ["TID251"]
 
-[lint.pyflakes]
-allowed-unused-imports = [
-    "tests.integration_tests",
-    "tests.unit_tests",
-]
-
 [lint.flake8-tidy-imports]
 
 [lint.flake8-tidy-imports.banned-api."flask_restx.reqparse"]

api/context/execution_context.py

@@ -10,7 +10,7 @@ import threading
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Generator
 from contextlib import AbstractContextManager, contextmanager
-from typing import Any, Protocol, TypeVar, final, runtime_checkable
+from typing import Any, Protocol, final, runtime_checkable
 
 from pydantic import BaseModel
 
@@ -188,8 +188,6 @@ class ExecutionContextBuilder:
 _capturer: Callable[[], IExecutionContext] | None = None
 _tenant_context_providers: dict[tuple[str, str], Callable[[], BaseModel]] = {}
 
-T = TypeVar("T", bound=BaseModel)
-
 
 class ContextProviderNotFoundError(KeyError):
     """Raised when a tenant-scoped context provider is missing."""

api/contexts/wrapper.py

@@ -1,7 +1,4 @@
 from contextvars import ContextVar
-from typing import Generic, TypeVar
-
-T = TypeVar("T")
 
 
 class HiddenValue:
@@ -11,7 +8,7 @@ class HiddenValue:
 _default = HiddenValue()
 
 
-class RecyclableContextVar(Generic[T]):
+class RecyclableContextVar[T]:
     """
     RecyclableContextVar is a wrapper around ContextVar
     It's safe to use in gunicorn with thread recycling, but features like `reset` are not available for now

api/controllers/common/fields.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
-from typing import Any, TypeAlias
+from typing import Any
 
 from graphon.file import helpers as file_helpers
 from pydantic import BaseModel, ConfigDict, computed_field
 
 from models.model import IconType
 
-JSONValue: TypeAlias = str | int | float | bool | None | dict[str, Any] | list[Any]
-JSONObject: TypeAlias = dict[str, Any]
+type JSONValue = str | int | float | bool | None | dict[str, Any] | list[Any]
+type JSONObject = dict[str, Any]
 
 
 class SystemParameters(BaseModel):

api/controllers/console/admin.py

@@ -2,7 +2,6 @@ import csv
 import io
 from collections.abc import Callable
 from functools import wraps
-from typing import ParamSpec, TypeVar
 
 from flask import request
 from flask_restx import Resource
@@ -20,9 +19,6 @@ from libs.token import extract_access_token
 from models.model import App, ExporleBanner, InstalledApp, RecommendedApp, TrialApp
 from services.billing_service import BillingService
 
-P = ParamSpec("P")
-R = TypeVar("R")
-
 DEFAULT_REF_TEMPLATE_SWAGGER_2_0 = "#/definitions/{model}"
 
 
@@ -72,9 +68,9 @@ console_ns.schema_model(
 )
 
 
-def admin_required(view: Callable[P, R]):
+def admin_required[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
-    def decorated(*args: P.args, **kwargs: P.kwargs):
+    def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
         if not dify_config.ADMIN_API_KEY:
             raise Unauthorized("API key is invalid.")
 

api/controllers/console/app/app.py

@@ -1,7 +1,7 @@
 import logging
 import uuid
 from datetime import datetime
-from typing import Any, Literal, TypeAlias
+from typing import Any, Literal
 
 from flask import request
 from flask_restx import Resource
@@ -152,7 +152,7 @@ class AppTracePayload(BaseModel):
         return value
 
 
-JSONValue: TypeAlias = Any
+type JSONValue = Any
 
 
 class ResponseModel(BaseModel):

api/controllers/console/app/app_import.py

@@ -1,6 +1,6 @@
 from flask_restx import Resource, fields, marshal_with
 from pydantic import BaseModel, Field
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import Session, sessionmaker
 
 from controllers.console.app.wraps import get_app_model
 from controllers.console.wraps import (
@@ -71,7 +71,7 @@ class AppImportApi(Resource):
         args = AppImportPayload.model_validate(console_ns.payload)
 
         # Create service with session
-        with sessionmaker(db.engine).begin() as session:
+        with Session(db.engine) as session:
             import_service = AppDslService(session)
             # Import app
             account = current_user

api/controllers/console/app/workflow.py

@@ -9,7 +9,7 @@ from graphon.enums import NodeType
 from graphon.file import File
 from graphon.graph_engine.manager import GraphEngineManager
 from graphon.model_runtime.utils.encoders import jsonable_encoder
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, Field, ValidationError, field_validator
 from sqlalchemy.orm import sessionmaker
 from werkzeug.exceptions import BadRequest, Forbidden, InternalServerError, NotFound
 
@@ -268,22 +268,18 @@ class DraftWorkflowApi(Resource):
 
         content_type = request.headers.get("Content-Type", "")
 
-        payload_data: dict[str, Any] | None = None
         if "application/json" in content_type:
             payload_data = request.get_json(silent=True)
             if not isinstance(payload_data, dict):
                 return {"message": "Invalid JSON data"}, 400
+            args_model = SyncDraftWorkflowPayload.model_validate(payload_data)
         elif "text/plain" in content_type:
             try:
-                payload_data = json.loads(request.data.decode("utf-8"))
-            except json.JSONDecodeError:
-                return {"message": "Invalid JSON data"}, 400
-            if not isinstance(payload_data, dict):
+                args_model = SyncDraftWorkflowPayload.model_validate_json(request.data)
+            except (ValueError, ValidationError):
                 return {"message": "Invalid JSON data"}, 400
         else:
             abort(415)
-
-        args_model = SyncDraftWorkflowPayload.model_validate(payload_data)
         args = args_model.model_dump()
         workflow_service = WorkflowService()
 

api/controllers/console/app/workflow_draft_variable.py

@@ -1,7 +1,7 @@
 import logging
 from collections.abc import Callable
 from functools import wraps
-from typing import Any, NoReturn, ParamSpec, TypeVar
+from typing import Any
 
 from flask import Response, request
 from flask_restx import Resource, fields, marshal, marshal_with
@@ -192,11 +192,8 @@ workflow_draft_variable_list_model = console_ns.model(
     "WorkflowDraftVariableList", workflow_draft_variable_list_fields_copy
 )
 
-P = ParamSpec("P")
-R = TypeVar("R")
 
-
-def _api_prerequisite(f: Callable[P, R]):
+def _api_prerequisite[**P, R](f: Callable[P, R]) -> Callable[P, R | Response]:
     """Common prerequisites for all draft workflow variable APIs.
 
     It ensures the following conditions are satisfied:
@@ -213,7 +210,7 @@ def _api_prerequisite(f: Callable[P, R]):
     @edit_permission_required
     @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW])
     @wraps(f)
-    def wrapper(*args: P.args, **kwargs: P.kwargs):
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R | Response:
         return f(*args, **kwargs)
 
     return wrapper
@@ -270,7 +267,7 @@ class WorkflowVariableCollectionApi(Resource):
         return Response("", 204)
 
 
-def validate_node_id(node_id: str) -> NoReturn | None:
+def validate_node_id(node_id: str) -> None:
     if node_id in [
         CONVERSATION_VARIABLE_NODE_ID,
         SYSTEM_VARIABLE_NODE_ID,
@@ -285,7 +282,6 @@ def validate_node_id(node_id: str) -> NoReturn | None:
         raise InvalidArgumentError(
             f"invalid node_id, please use correspond api for conversation and system variables, node_id={node_id}",
         )
-    return None
 
 
 @console_ns.route("/apps/<uuid:app_id>/workflows/draft/nodes/<string:node_id>/variables")

api/controllers/console/app/wraps.py

@@ -1,6 +1,6 @@
 from collections.abc import Callable
 from functools import wraps
-from typing import ParamSpec, TypeVar, Union
+from typing import overload
 
 from sqlalchemy import select
 
@@ -9,11 +9,6 @@ from extensions.ext_database import db
 from libs.login import current_account_with_tenant
 from models import App, AppMode
 
-P = ParamSpec("P")
-R = TypeVar("R")
-P1 = ParamSpec("P1")
-R1 = TypeVar("R1")
-
 
 def _load_app_model(app_id: str) -> App | None:
     _, current_tenant_id = current_account_with_tenant()
@@ -28,10 +23,30 @@ def _load_app_model_with_trial(app_id: str) -> App | None:
     return app_model
 
 
-def get_app_model(view: Callable[P, R] | None = None, *, mode: Union[AppMode, list[AppMode], None] = None):
-    def decorator(view_func: Callable[P1, R1]):
+@overload
+def get_app_model[**P, R](
+    view: Callable[P, R],
+    *,
+    mode: AppMode | list[AppMode] | None = None,
+) -> Callable[P, R]: ...
+
+
+@overload
+def get_app_model[**P, R](
+    view: None = None,
+    *,
+    mode: AppMode | list[AppMode] | None = None,
+) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
+
+
+def get_app_model[**P, R](
+    view: Callable[P, R] | None = None,
+    *,
+    mode: AppMode | list[AppMode] | None = None,
+) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
+    def decorator(view_func: Callable[P, R]) -> Callable[P, R]:
         @wraps(view_func)
-        def decorated_view(*args: P1.args, **kwargs: P1.kwargs):
+        def decorated_view(*args: P.args, **kwargs: P.kwargs) -> R:
             if not kwargs.get("app_id"):
                 raise ValueError("missing app_id in path parameters")
 
@@ -69,10 +84,30 @@ def get_app_model(view: Callable[P, R] | None = None, *, mode: Union[AppMode, li
         return decorator(view)
 
 
-def get_app_model_with_trial(view: Callable[P, R] | None = None, *, mode: Union[AppMode, list[AppMode], None] = None):
-    def decorator(view_func: Callable[P, R]):
+@overload
+def get_app_model_with_trial[**P, R](
+    view: Callable[P, R],
+    *,
+    mode: AppMode | list[AppMode] | None = None,
+) -> Callable[P, R]: ...
+
+
+@overload
+def get_app_model_with_trial[**P, R](
+    view: None = None,
+    *,
+    mode: AppMode | list[AppMode] | None = None,
+) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
+
+
+def get_app_model_with_trial[**P, R](
+    view: Callable[P, R] | None = None,
+    *,
+    mode: AppMode | list[AppMode] | None = None,
+) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
+    def decorator(view_func: Callable[P, R]) -> Callable[P, R]:
         @wraps(view_func)
-        def decorated_view(*args: P.args, **kwargs: P.kwargs):
+        def decorated_view(*args: P.args, **kwargs: P.kwargs) -> R:
             if not kwargs.get("app_id"):
                 raise ValueError("missing app_id in path parameters")
 

api/controllers/console/auth/oauth_server.py

@@ -1,8 +1,9 @@
 from collections.abc import Callable
 from functools import wraps
-from typing import Concatenate, ParamSpec, TypeVar
+from typing import Concatenate
 
 from flask import jsonify, request
+from flask.typing import ResponseReturnValue
 from flask_restx import Resource
 from graphon.model_runtime.utils.encoders import jsonable_encoder
 from pydantic import BaseModel
@@ -16,10 +17,6 @@ from services.oauth_server import OAUTH_ACCESS_TOKEN_EXPIRES_IN, OAuthGrantType,
 
 from .. import console_ns
 
-P = ParamSpec("P")
-R = TypeVar("R")
-T = TypeVar("T")
-
 
 class OAuthClientPayload(BaseModel):
     client_id: str
@@ -39,9 +36,11 @@ class OAuthTokenRequest(BaseModel):
     refresh_token: str | None = None
 
 
-def oauth_server_client_id_required(view: Callable[Concatenate[T, OAuthProviderApp, P], R]):
+def oauth_server_client_id_required[T, **P, R](
+    view: Callable[Concatenate[T, OAuthProviderApp, P], R],
+) -> Callable[Concatenate[T, P], R]:
     @wraps(view)
-    def decorated(self: T, *args: P.args, **kwargs: P.kwargs):
+    def decorated(self: T, *args: P.args, **kwargs: P.kwargs) -> R:
         json_data = request.get_json()
         if json_data is None:
             raise BadRequest("client_id is required")
@@ -58,9 +57,13 @@ def oauth_server_client_id_required(view: Callable[Concatenate[T, OAuthProviderA
     return decorated
 
 
-def oauth_server_access_token_required(view: Callable[Concatenate[T, OAuthProviderApp, Account, P], R]):
+def oauth_server_access_token_required[T, **P, R](
+    view: Callable[Concatenate[T, OAuthProviderApp, Account, P], R],
+) -> Callable[Concatenate[T, OAuthProviderApp, P], R | ResponseReturnValue]:
     @wraps(view)
-    def decorated(self: T, oauth_provider_app: OAuthProviderApp, *args: P.args, **kwargs: P.kwargs):
+    def decorated(
+        self: T, oauth_provider_app: OAuthProviderApp, *args: P.args, **kwargs: P.kwargs
+    ) -> R | ResponseReturnValue:
         if not isinstance(oauth_provider_app, OAuthProviderApp):
             raise BadRequest("Invalid oauth_provider_app")
 

api/controllers/console/datasets/data_source.py

@@ -158,10 +158,11 @@ class DataSourceApi(Resource):
     @login_required
     @account_initialization_required
     def patch(self, binding_id, action: Literal["enable", "disable"]):
+        _, current_tenant_id = current_account_with_tenant()
         binding_id = str(binding_id)
         with sessionmaker(db.engine, expire_on_commit=False).begin() as session:
             data_source_binding = session.execute(
-                select(DataSourceOauthBinding).filter_by(id=binding_id)
+                select(DataSourceOauthBinding).filter_by(id=binding_id, tenant_id=current_tenant_id)
             ).scalar_one_or_none()
         if data_source_binding is None:
             raise NotFound("Data source binding not found.")

api/controllers/console/datasets/external.py

@@ -173,8 +173,11 @@ class ExternalApiTemplateApi(Resource):
     @login_required
     @account_initialization_required
     def get(self, external_knowledge_api_id):
+        _, current_tenant_id = current_account_with_tenant()
         external_knowledge_api_id = str(external_knowledge_api_id)
-        external_knowledge_api = ExternalDatasetService.get_external_knowledge_api(external_knowledge_api_id)
+        external_knowledge_api = ExternalDatasetService.get_external_knowledge_api(
+            external_knowledge_api_id, current_tenant_id
+        )
         if external_knowledge_api is None:
             raise NotFound("API template not found.")
 

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

@@ -1,4 +1,5 @@
 import logging
+from collections.abc import Callable
 from typing import Any, NoReturn
 
 from flask import Response, request
@@ -55,7 +56,7 @@ class WorkflowDraftVariablePatchPayload(BaseModel):
 register_schema_models(console_ns, WorkflowDraftVariablePatchPayload)
 
 
-def _api_prerequisite(f):
+def _api_prerequisite[**P, R](f: Callable[P, R]) -> Callable[P, R | Response]:
     """Common prerequisites for all draft workflow variable APIs.
 
     It ensures the following conditions are satisfied:
@@ -70,7 +71,7 @@ def _api_prerequisite(f):
     @login_required
     @account_initialization_required
     @get_rag_pipeline
-    def wrapper(*args, **kwargs):
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R | Response:
         if not isinstance(current_user, Account) or not current_user.has_edit_permission:
             raise Forbidden()
         return f(*args, **kwargs)

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

@@ -5,7 +5,7 @@ from typing import Any, Literal, cast
 from flask import abort, request
 from flask_restx import Resource, marshal_with  # type: ignore
 from graphon.model_runtime.utils.encoders import jsonable_encoder
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ValidationError
 from sqlalchemy.orm import sessionmaker
 from werkzeug.exceptions import BadRequest, Forbidden, InternalServerError, NotFound
 
@@ -186,29 +186,14 @@ class DraftRagPipelineApi(Resource):
 
         if "application/json" in content_type:
             payload_dict = console_ns.payload or {}
+            payload = DraftWorkflowSyncPayload.model_validate(payload_dict)
         elif "text/plain" in content_type:
             try:
-                data = json.loads(request.data.decode("utf-8"))
-                if "graph" not in data or "features" not in data:
-                    raise ValueError("graph or features not found in data")
-
-                if not isinstance(data.get("graph"), dict):
-                    raise ValueError("graph is not a dict")
-
-                payload_dict = {
-                    "graph": data.get("graph"),
-                    "features": data.get("features"),
-                    "hash": data.get("hash"),
-                    "environment_variables": data.get("environment_variables"),
-                    "conversation_variables": data.get("conversation_variables"),
-                    "rag_pipeline_variables": data.get("rag_pipeline_variables"),
-                }
-            except json.JSONDecodeError:
+                payload = DraftWorkflowSyncPayload.model_validate_json(request.data)
+            except (ValueError, ValidationError):
                 return {"message": "Invalid JSON data"}, 400
         else:
             abort(415)
-
-        payload = DraftWorkflowSyncPayload.model_validate(payload_dict)
         rag_pipeline_service = RagPipelineService()
 
         try:
@@ -608,17 +593,15 @@ class PublishedRagPipelineApi(Resource):
         # The role of the current user in the ta table must be admin, owner, or editor
         current_user, _ = current_account_with_tenant()
         rag_pipeline_service = RagPipelineService()
-        with sessionmaker(db.engine).begin() as session:
-            pipeline = session.merge(pipeline)
-            workflow = rag_pipeline_service.publish_workflow(
-                session=session,
-                pipeline=pipeline,
-                account=current_user,
-            )
-            pipeline.is_published = True
-            pipeline.workflow_id = workflow.id
-            session.add(pipeline)
-            workflow_created_at = TimestampField().format(workflow.created_at)
+        workflow = rag_pipeline_service.publish_workflow(
+            session=db.session,  # type: ignore[reportArgumentType,arg-type]
+            pipeline=pipeline,
+            account=current_user,
+        )
+        pipeline.is_published = True
+        pipeline.workflow_id = workflow.id
+        db.session.commit()
+        workflow_created_at = TimestampField().format(workflow.created_at)
 
         return {
             "result": "success",

api/controllers/console/datasets/wraps.py

@@ -1,6 +1,5 @@
 from collections.abc import Callable
 from functools import wraps
-from typing import ParamSpec, TypeVar
 
 from sqlalchemy import select
 
@@ -9,13 +8,10 @@ from extensions.ext_database import db
 from libs.login import current_account_with_tenant
 from models.dataset import Pipeline
 
-P = ParamSpec("P")
-R = TypeVar("R")
 
-
-def get_rag_pipeline(view_func: Callable[P, R]):
+def get_rag_pipeline[**P, R](view_func: Callable[P, R]) -> Callable[P, R]:
     @wraps(view_func)
-    def decorated_view(*args: P.args, **kwargs: P.kwargs):
+    def decorated_view(*args: P.args, **kwargs: P.kwargs) -> R:
         if not kwargs.get("pipeline_id"):
             raise ValueError("missing pipeline_id in path parameters")
 

api/controllers/console/explore/wraps.py

@@ -1,6 +1,6 @@
 from collections.abc import Callable
 from functools import wraps
-from typing import Concatenate, ParamSpec, TypeVar
+from typing import Concatenate
 
 from flask import abort
 from flask_restx import Resource
@@ -15,12 +15,8 @@ from models import AccountTrialAppRecord, App, InstalledApp, TrialApp
 from services.enterprise.enterprise_service import EnterpriseService
 from services.feature_service import FeatureService
 
-P = ParamSpec("P")
-R = TypeVar("R")
-T = TypeVar("T")
 
-
-def installed_app_required(view: Callable[Concatenate[InstalledApp, P], R] | None = None):
+def installed_app_required[**P, R](view: Callable[Concatenate[InstalledApp, P], R] | None = None):
     def decorator(view: Callable[Concatenate[InstalledApp, P], R]):
         @wraps(view)
         def decorated(installed_app_id: str, *args: P.args, **kwargs: P.kwargs):
@@ -49,7 +45,7 @@ def installed_app_required(view: Callable[Concatenate[InstalledApp, P], R] | Non
     return decorator
 
 
-def user_allowed_to_access_app(view: Callable[Concatenate[InstalledApp, P], R] | None = None):
+def user_allowed_to_access_app[**P, R](view: Callable[Concatenate[InstalledApp, P], R] | None = None):
     def decorator(view: Callable[Concatenate[InstalledApp, P], R]):
         @wraps(view)
         def decorated(installed_app: InstalledApp, *args: P.args, **kwargs: P.kwargs):
@@ -73,7 +69,7 @@ def user_allowed_to_access_app(view: Callable[Concatenate[InstalledApp, P], R] |
     return decorator
 
 
-def trial_app_required(view: Callable[Concatenate[App, P], R] | None = None):
+def trial_app_required[**P, R](view: Callable[Concatenate[App, P], R] | None = None):
     def decorator(view: Callable[Concatenate[App, P], R]):
         @wraps(view)
         def decorated(app_id: str, *args: P.args, **kwargs: P.kwargs):
@@ -106,7 +102,7 @@ def trial_app_required(view: Callable[Concatenate[App, P], R] | None = None):
     return decorator
 
 
-def trial_feature_enable(view: Callable[P, R]):
+def trial_feature_enable[**P, R](view: Callable[P, R]):
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs):
         features = FeatureService.get_system_features()
@@ -117,7 +113,7 @@ def trial_feature_enable(view: Callable[P, R]):
     return decorated
 
 
-def explore_banner_enabled(view: Callable[P, R]):
+def explore_banner_enabled[**P, R](view: Callable[P, R]):
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs):
         features = FeatureService.get_system_features()

api/controllers/console/workspace/__init__.py

@@ -1,6 +1,5 @@
 from collections.abc import Callable
 from functools import wraps
-from typing import ParamSpec, TypeVar
 
 from sqlalchemy.orm import sessionmaker
 from werkzeug.exceptions import Forbidden
@@ -9,17 +8,14 @@ from extensions.ext_database import db
 from libs.login import current_account_with_tenant
 from models.account import TenantPluginPermission
 
-P = ParamSpec("P")
-R = TypeVar("R")
-
 
 def plugin_permission_required(
     install_required: bool = False,
     debug_required: bool = False,
 ):
-    def interceptor(view: Callable[P, R]):
+    def interceptor[**P, R](view: Callable[P, R]) -> Callable[P, R]:
         @wraps(view)
-        def decorated(*args: P.args, **kwargs: P.kwargs):
+        def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
             current_user, current_tenant_id = current_account_with_tenant()
             user = current_user
             tenant_id = current_tenant_id

api/controllers/console/wraps.py

@@ -4,7 +4,6 @@ import os
 import time
 from collections.abc import Callable
 from functools import wraps
-from typing import ParamSpec, TypeVar
 
 from flask import abort, request
 from sqlalchemy import select
@@ -25,9 +24,6 @@ from services.operation_service import OperationService
 
 from .error import NotInitValidateError, NotSetupError, UnauthorizedAndForceLogout
 
-P = ParamSpec("P")
-R = TypeVar("R")
-
 # Field names for decryption
 FIELD_NAME_PASSWORD = "password"
 FIELD_NAME_CODE = "code"
@@ -37,7 +33,7 @@ ERROR_MSG_INVALID_ENCRYPTED_DATA = "Invalid encrypted data"
 ERROR_MSG_INVALID_ENCRYPTED_CODE = "Invalid encrypted code"
 
 
-def account_initialization_required(view: Callable[P, R]) -> Callable[P, R]:
+def account_initialization_required[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
         # check account initialization
@@ -50,7 +46,7 @@ def account_initialization_required(view: Callable[P, R]) -> Callable[P, R]:
     return decorated
 
 
-def only_edition_cloud(view: Callable[P, R]):
+def only_edition_cloud[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs):
         if dify_config.EDITION != "CLOUD":
@@ -61,7 +57,7 @@ def only_edition_cloud(view: Callable[P, R]):
     return decorated
 
 
-def only_edition_enterprise(view: Callable[P, R]):
+def only_edition_enterprise[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs):
         if not dify_config.ENTERPRISE_ENABLED:
@@ -72,7 +68,7 @@ def only_edition_enterprise(view: Callable[P, R]):
     return decorated
 
 
-def only_edition_self_hosted(view: Callable[P, R]):
+def only_edition_self_hosted[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs):
         if dify_config.EDITION != "SELF_HOSTED":
@@ -83,7 +79,7 @@ def only_edition_self_hosted(view: Callable[P, R]):
     return decorated
 
 
-def cloud_edition_billing_enabled(view: Callable[P, R]):
+def cloud_edition_billing_enabled[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs):
         _, current_tenant_id = current_account_with_tenant()
@@ -95,7 +91,7 @@ def cloud_edition_billing_enabled(view: Callable[P, R]):
     return decorated
 
 
-def cloud_edition_billing_resource_check(resource: str):
+def cloud_edition_billing_resource_check[**P, R](resource: str) -> Callable[[Callable[P, R]], Callable[P, R]]:
     def interceptor(view: Callable[P, R]):
         @wraps(view)
         def decorated(*args: P.args, **kwargs: P.kwargs):
@@ -137,7 +133,9 @@ def cloud_edition_billing_resource_check(resource: str):
     return interceptor
 
 
-def cloud_edition_billing_knowledge_limit_check(resource: str):
+def cloud_edition_billing_knowledge_limit_check[**P, R](
+    resource: str,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
     def interceptor(view: Callable[P, R]):
         @wraps(view)
         def decorated(*args: P.args, **kwargs: P.kwargs):
@@ -160,7 +158,7 @@ def cloud_edition_billing_knowledge_limit_check(resource: str):
     return interceptor
 
 
-def cloud_edition_billing_rate_limit_check(resource: str):
+def cloud_edition_billing_rate_limit_check[**P, R](resource: str) -> Callable[[Callable[P, R]], Callable[P, R]]:
     def interceptor(view: Callable[P, R]):
         @wraps(view)
         def decorated(*args: P.args, **kwargs: P.kwargs):
@@ -196,7 +194,7 @@ def cloud_edition_billing_rate_limit_check(resource: str):
     return interceptor
 
 
-def cloud_utm_record(view: Callable[P, R]):
+def cloud_utm_record[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs):
         with contextlib.suppress(Exception):
@@ -215,7 +213,7 @@ def cloud_utm_record(view: Callable[P, R]):
     return decorated
 
 
-def setup_required(view: Callable[P, R]) -> Callable[P, R]:
+def setup_required[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
         # check setup
@@ -229,7 +227,7 @@ def setup_required(view: Callable[P, R]) -> Callable[P, R]:
     return decorated
 
 
-def enterprise_license_required(view: Callable[P, R]):
+def enterprise_license_required[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs):
         settings = FeatureService.get_system_features()
@@ -241,7 +239,7 @@ def enterprise_license_required(view: Callable[P, R]):
     return decorated
 
 
-def email_password_login_enabled(view: Callable[P, R]):
+def email_password_login_enabled[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs):
         features = FeatureService.get_system_features()
@@ -254,7 +252,7 @@ def email_password_login_enabled(view: Callable[P, R]):
     return decorated
 
 
-def email_register_enabled(view: Callable[P, R]):
+def email_register_enabled[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs):
         features = FeatureService.get_system_features()
@@ -267,7 +265,7 @@ def email_register_enabled(view: Callable[P, R]):
     return decorated
 
 
-def enable_change_email(view: Callable[P, R]):
+def enable_change_email[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs):
         features = FeatureService.get_system_features()
@@ -280,7 +278,7 @@ def enable_change_email(view: Callable[P, R]):
     return decorated
 
 
-def is_allow_transfer_owner(view: Callable[P, R]):
+def is_allow_transfer_owner[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs):
         from libs.workspace_permission import check_workspace_owner_transfer_permission
@@ -293,7 +291,7 @@ def is_allow_transfer_owner(view: Callable[P, R]):
     return decorated
 
 
-def knowledge_pipeline_publish_enabled(view: Callable[P, R]):
+def knowledge_pipeline_publish_enabled[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs):
         _, current_tenant_id = current_account_with_tenant()
@@ -305,7 +303,7 @@ def knowledge_pipeline_publish_enabled(view: Callable[P, R]):
     return decorated
 
 
-def edit_permission_required(f: Callable[P, R]):
+def edit_permission_required[**P, R](f: Callable[P, R]) -> Callable[P, R]:
     @wraps(f)
     def decorated_function(*args: P.args, **kwargs: P.kwargs):
         from werkzeug.exceptions import Forbidden
@@ -323,7 +321,7 @@ def edit_permission_required(f: Callable[P, R]):
     return decorated_function
 
 
-def is_admin_or_owner_required(f: Callable[P, R]):
+def is_admin_or_owner_required[**P, R](f: Callable[P, R]) -> Callable[P, R]:
     @wraps(f)
     def decorated_function(*args: P.args, **kwargs: P.kwargs):
         from werkzeug.exceptions import Forbidden
@@ -339,7 +337,7 @@ def is_admin_or_owner_required(f: Callable[P, R]):
     return decorated_function
 
 
-def annotation_import_rate_limit(view: Callable[P, R]):
+def annotation_import_rate_limit[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     """
     Rate limiting decorator for annotation import operations.
 
@@ -388,7 +386,7 @@ def annotation_import_rate_limit(view: Callable[P, R]):
     return decorated
 
 
-def annotation_import_concurrency_limit(view: Callable[P, R]):
+def annotation_import_concurrency_limit[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     """
     Concurrency control decorator for annotation import operations.
 
@@ -455,7 +453,7 @@ def _decrypt_field(field_name: str, error_class: type[Exception], error_message:
     payload[field_name] = decoded_value
 
 
-def decrypt_password_field(view: Callable[P, R]):
+def decrypt_password_field[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     """
     Decorator to decrypt password field in request payload.
 
@@ -477,7 +475,7 @@ def decrypt_password_field(view: Callable[P, R]):
     return decorated
 
 
-def decrypt_code_field(view: Callable[P, R]):
+def decrypt_code_field[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     """
     Decorator to decrypt verification code field in request payload.
 

api/controllers/inner_api/plugin/wraps.py

@@ -1,21 +1,17 @@
 from collections.abc import Callable
 from functools import wraps
-from typing import ParamSpec, TypeVar
 
 from flask import current_app, request
 from flask_login import user_logged_in
 from pydantic import BaseModel
 from sqlalchemy import select
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import sessionmaker
 
 from extensions.ext_database import db
 from libs.login import current_user
 from models.account import Tenant
 from models.model import DefaultEndUserSessionID, EndUser
 
-P = ParamSpec("P")
-R = TypeVar("R")
-
 
 class TenantUserPayload(BaseModel):
     tenant_id: str
@@ -33,7 +29,7 @@ def get_user(tenant_id: str, user_id: str | None) -> EndUser:
         user_id = DefaultEndUserSessionID.DEFAULT_SESSION_ID
     is_anonymous = user_id == DefaultEndUserSessionID.DEFAULT_SESSION_ID
     try:
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine, expire_on_commit=False).begin() as session:
             user_model = None
 
             if is_anonymous:
@@ -56,7 +52,7 @@ def get_user(tenant_id: str, user_id: str | None) -> EndUser:
                     session_id=user_id,
                 )
                 session.add(user_model)
-                session.commit()
+                session.flush()
                 session.refresh(user_model)
 
     except Exception:
@@ -65,9 +61,9 @@ def get_user(tenant_id: str, user_id: str | None) -> EndUser:
     return user_model
 
 
-def get_user_tenant(view_func: Callable[P, R]):
+def get_user_tenant[**P, R](view_func: Callable[P, R]) -> Callable[P, R]:
     @wraps(view_func)
-    def decorated_view(*args: P.args, **kwargs: P.kwargs):
+    def decorated_view(*args: P.args, **kwargs: P.kwargs) -> R:
         payload = TenantUserPayload.model_validate(request.get_json(silent=True) or {})
 
         user_id = payload.user_id
@@ -97,10 +93,14 @@ def get_user_tenant(view_func: Callable[P, R]):
     return decorated_view
 
 
-def plugin_data(view: Callable[P, R] | None = None, *, payload_type: type[BaseModel]):
-    def decorator(view_func: Callable[P, R]):
+def plugin_data[**P, R](
+    view: Callable[P, R] | None = None,
+    *,
+    payload_type: type[BaseModel],
+) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
+    def decorator(view_func: Callable[P, R]) -> Callable[P, R]:
         @wraps(view_func)
-        def decorated_view(*args: P.args, **kwargs: P.kwargs):
+        def decorated_view(*args: P.args, **kwargs: P.kwargs) -> R:
             try:
                 data = request.get_json()
             except Exception:

api/controllers/inner_api/wraps.py

@@ -3,10 +3,7 @@ from collections.abc import Callable
 from functools import wraps
 from hashlib import sha1
 from hmac import new as hmac_new
-from typing import ParamSpec, TypeVar
 
-P = ParamSpec("P")
-R = TypeVar("R")
 from flask import abort, request
 
 from configs import dify_config
@@ -14,9 +11,9 @@ from extensions.ext_database import db
 from models.model import EndUser
 
 
-def billing_inner_api_only(view: Callable[P, R]):
+def billing_inner_api_only[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
-    def decorated(*args: P.args, **kwargs: P.kwargs):
+    def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
         if not dify_config.INNER_API:
             abort(404)
 
@@ -30,9 +27,9 @@ def billing_inner_api_only(view: Callable[P, R]):
     return decorated
 
 
-def enterprise_inner_api_only(view: Callable[P, R]):
+def enterprise_inner_api_only[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
-    def decorated(*args: P.args, **kwargs: P.kwargs):
+    def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
         if not dify_config.INNER_API:
             abort(404)
 
@@ -46,9 +43,9 @@ def enterprise_inner_api_only(view: Callable[P, R]):
     return decorated
 
 
-def enterprise_inner_api_user_auth(view: Callable[P, R]):
+def enterprise_inner_api_user_auth[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
-    def decorated(*args: P.args, **kwargs: P.kwargs):
+    def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
         if not dify_config.INNER_API:
             return view(*args, **kwargs)
 
@@ -82,9 +79,9 @@ def enterprise_inner_api_user_auth(view: Callable[P, R]):
     return decorated
 
 
-def plugin_inner_api_only(view: Callable[P, R]):
+def plugin_inner_api_only[**P, R](view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
-    def decorated(*args: P.args, **kwargs: P.kwargs):
+    def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
         if not dify_config.PLUGIN_DAEMON_KEY:
             abort(404)
 

api/controllers/service_api/app/conversation.py

@@ -3,7 +3,7 @@ from typing import Any, Literal
 from flask import request
 from flask_restx import Resource
 from pydantic import BaseModel, Field, TypeAdapter, field_validator, model_validator
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import sessionmaker
 from werkzeug.exceptions import BadRequest, NotFound
 
 import services
@@ -116,7 +116,7 @@ class ConversationApi(Resource):
         last_id = str(query_args.last_id) if query_args.last_id else None
 
         try:
-            with Session(db.engine) as session:
+            with sessionmaker(db.engine).begin() as session:
                 pagination = ConversationService.pagination_by_last_id(
                     session=session,
                     app_model=app_model,

api/controllers/service_api/app/workflow.py

@@ -8,7 +8,7 @@ from graphon.enums import WorkflowExecutionStatus
 from graphon.graph_engine.manager import GraphEngineManager
 from graphon.model_runtime.errors.invoke import InvokeError
 from pydantic import BaseModel, Field
-from sqlalchemy.orm import Session, sessionmaker
+from sqlalchemy.orm import sessionmaker
 from werkzeug.exceptions import BadRequest, InternalServerError, NotFound
 
 from controllers.common.schema import register_schema_models
@@ -314,7 +314,7 @@ class WorkflowAppLogApi(Resource):
 
         # get paginate workflow app logs
         workflow_app_service = WorkflowAppService()
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine).begin() as session:
             workflow_app_log_pagination = workflow_app_service.get_paginate_workflow_app_logs(
                 session=session,
                 app_model=app_model,

api/controllers/service_api/wraps.py

@@ -1,9 +1,10 @@
+import inspect
 import logging
 import time
 from collections.abc import Callable
 from enum import StrEnum, auto
 from functools import wraps
-from typing import Concatenate, ParamSpec, TypeVar, cast, overload
+from typing import cast, overload
 
 from flask import current_app, request
 from flask_login import user_logged_in
@@ -23,10 +24,6 @@ from services.api_token_service import ApiTokenCache, fetch_token_with_single_fl
 from services.end_user_service import EndUserService
 from services.feature_service import FeatureService
 
-P = ParamSpec("P")
-R = TypeVar("R")
-T = TypeVar("T")
-
 logger = logging.getLogger(__name__)
 
 
@@ -46,16 +43,16 @@ class FetchUserArg(BaseModel):
 
 
 @overload
-def validate_app_token(view: Callable[P, R]) -> Callable[P, R]: ...
+def validate_app_token[**P, R](view: Callable[P, R]) -> Callable[P, R]: ...
 
 
 @overload
-def validate_app_token(
+def validate_app_token[**P, R](
     view: None = None, *, fetch_user_arg: FetchUserArg | None = None
 ) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
 
 
-def validate_app_token(
+def validate_app_token[**P, R](
     view: Callable[P, R] | None = None, *, fetch_user_arg: FetchUserArg | None = None
 ) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
     def decorator(view_func: Callable[P, R]) -> Callable[P, R]:
@@ -136,7 +133,10 @@ def validate_app_token(
         return decorator(view)
 
 
-def cloud_edition_billing_resource_check(resource: str, api_token_type: str):
+def cloud_edition_billing_resource_check[**P, R](
+    resource: str,
+    api_token_type: str,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
     def interceptor(view: Callable[P, R]):
         def decorated(*args: P.args, **kwargs: P.kwargs):
             api_token = validate_and_get_api_token(api_token_type)
@@ -166,7 +166,10 @@ def cloud_edition_billing_resource_check(resource: str, api_token_type: str):
     return interceptor
 
 
-def cloud_edition_billing_knowledge_limit_check(resource: str, api_token_type: str):
+def cloud_edition_billing_knowledge_limit_check[**P, R](
+    resource: str,
+    api_token_type: str,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
     def interceptor(view: Callable[P, R]):
         @wraps(view)
         def decorated(*args: P.args, **kwargs: P.kwargs):
@@ -188,7 +191,10 @@ def cloud_edition_billing_knowledge_limit_check(resource: str, api_token_type: s
     return interceptor
 
 
-def cloud_edition_billing_rate_limit_check(resource: str, api_token_type: str):
+def cloud_edition_billing_rate_limit_check[**P, R](
+    resource: str,
+    api_token_type: str,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
     def interceptor(view: Callable[P, R]):
         @wraps(view)
         def decorated(*args: P.args, **kwargs: P.kwargs):
@@ -225,99 +231,73 @@ def cloud_edition_billing_rate_limit_check(resource: str, api_token_type: str):
     return interceptor
 
 
-@overload
-def validate_dataset_token(view: Callable[Concatenate[T, P], R]) -> Callable[P, R]: ...
+def validate_dataset_token[R](view: Callable[..., R]) -> Callable[..., R]:
+    positional_parameters = [
+        parameter
+        for parameter in inspect.signature(view).parameters.values()
+        if parameter.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+    ]
+    expects_bound_instance = bool(positional_parameters and positional_parameters[0].name in {"self", "cls"})
 
+    @wraps(view)
+    def decorated(*args: object, **kwargs: object) -> R:
+        api_token = validate_and_get_api_token("dataset")
 
-@overload
-def validate_dataset_token(view: None = None) -> Callable[[Callable[Concatenate[T, P], R]], Callable[P, R]]: ...
+        # Flask may pass URL path parameters positionally, so inspect both kwargs and args.
+        dataset_id = kwargs.get("dataset_id")
 
+        if not dataset_id and args:
+            potential_id = args[0]
+            try:
+                str_id = str(potential_id)
+                if len(str_id) == 36 and str_id.count("-") == 4:
+                    dataset_id = str_id
+            except Exception:
+                logger.exception("Failed to parse dataset_id from positional args")
 
-def validate_dataset_token(
-    view: Callable[Concatenate[T, P], R] | None = None,
-) -> Callable[P, R] | Callable[[Callable[Concatenate[T, P], R]], Callable[P, R]]:
-    def decorator(view_func: Callable[Concatenate[T, P], R]) -> Callable[P, R]:
-        @wraps(view_func)
-        def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
-            api_token = validate_and_get_api_token("dataset")
-
-            # get url path dataset_id from positional args or kwargs
-            # Flask passes URL path parameters as positional arguments
-            dataset_id = None
-
-            # First try to get from kwargs (explicit parameter)
-            dataset_id = kwargs.get("dataset_id")
-
-            # If not in kwargs, try to extract from positional args
-            if not dataset_id and args:
-                # For class methods: args[0] is self, args[1] is dataset_id (if exists)
-                # Check if first arg is likely a class instance (has __dict__ or __class__)
-                if len(args) > 1 and hasattr(args[0], "__dict__"):
-                    # This is a class method, dataset_id should be in args[1]
-                    potential_id = args[1]
-                    # Validate it's a string-like UUID, not another object
-                    try:
-                        # Try to convert to string and check if it's a valid UUID format
-                        str_id = str(potential_id)
-                        # Basic check: UUIDs are 36 chars with hyphens
-                        if len(str_id) == 36 and str_id.count("-") == 4:
-                            dataset_id = str_id
-                    except Exception:
-                        logger.exception("Failed to parse dataset_id from class method args")
-                elif len(args) > 0:
-                    # Not a class method, check if args[0] looks like a UUID
-                    potential_id = args[0]
-                    try:
-                        str_id = str(potential_id)
-                        if len(str_id) == 36 and str_id.count("-") == 4:
-                            dataset_id = str_id
-                    except Exception:
-                        logger.exception("Failed to parse dataset_id from positional args")
-
-            # Validate dataset if dataset_id is provided
-            if dataset_id:
-                dataset_id = str(dataset_id)
-                dataset = db.session.scalar(
-                    select(Dataset)
-                    .where(
-                        Dataset.id == dataset_id,
-                        Dataset.tenant_id == api_token.tenant_id,
-                    )
-                    .limit(1)
+        if dataset_id:
+            dataset_id = str(dataset_id)
+            dataset = db.session.scalar(
+                select(Dataset)
+                .where(
+                    Dataset.id == dataset_id,
+                    Dataset.tenant_id == api_token.tenant_id,
                 )
-                if not dataset:
-                    raise NotFound("Dataset not found.")
-                if not dataset.enable_api:
-                    raise Forbidden("Dataset api access is not enabled.")
-            tenant_account_join = db.session.execute(
-                select(Tenant, TenantAccountJoin)
-                .where(Tenant.id == api_token.tenant_id)
-                .where(TenantAccountJoin.tenant_id == Tenant.id)
-                .where(TenantAccountJoin.role.in_(["owner"]))
-                .where(Tenant.status == TenantStatus.NORMAL)
-            ).one_or_none()  # TODO: only owner information is required, so only one is returned.
-            if tenant_account_join:
-                tenant, ta = tenant_account_join
-                account = db.session.get(Account, ta.account_id)
-                # Login admin
-                if account:
-                    account.current_tenant = tenant
-                    current_app.login_manager._update_request_context_with_user(account)  # type: ignore
-                    user_logged_in.send(current_app._get_current_object(), user=current_user)  # type: ignore
-                else:
-                    raise Unauthorized("Tenant owner account does not exist.")
+                .limit(1)
+            )
+            if not dataset:
+                raise NotFound("Dataset not found.")
+            if not dataset.enable_api:
+                raise Forbidden("Dataset api access is not enabled.")
+
+        tenant_account_join = db.session.execute(
+            select(Tenant, TenantAccountJoin)
+            .where(Tenant.id == api_token.tenant_id)
+            .where(TenantAccountJoin.tenant_id == Tenant.id)
+            .where(TenantAccountJoin.role.in_(["owner"]))
+            .where(Tenant.status == TenantStatus.NORMAL)
+        ).one_or_none()  # TODO: only owner information is required, so only one is returned.
+        if tenant_account_join:
+            tenant, ta = tenant_account_join
+            account = db.session.get(Account, ta.account_id)
+            # Login admin
+            if account:
+                account.current_tenant = tenant
+                current_app.login_manager._update_request_context_with_user(account)  # type: ignore
+                user_logged_in.send(current_app._get_current_object(), user=current_user)  # type: ignore
             else:
-                raise Unauthorized("Tenant does not exist.")
-            return view_func(api_token.tenant_id, *args, **kwargs)  # type: ignore[arg-type]
+                raise Unauthorized("Tenant owner account does not exist.")
+        else:
+            raise Unauthorized("Tenant does not exist.")
 
-        return decorated
+        if expects_bound_instance:
+            if not args:
+                raise TypeError("validate_dataset_token expected a bound resource instance.")
+            return view(args[0], api_token.tenant_id, *args[1:], **kwargs)
 
-    if view:
-        return decorator(view)
+        return view(api_token.tenant_id, *args, **kwargs)
 
-    # if view is None, it means that the decorator is used without parentheses
-    # use the decorator as a function for method_decorators
-    return decorator
+    return decorated
 
 
 def validate_and_get_api_token(scope: str | None = None):

api/controllers/trigger/webhook.py

@@ -7,7 +7,7 @@ from werkzeug.exceptions import NotFound, RequestEntityTooLarge
 from controllers.trigger import bp
 from core.trigger.debug.event_bus import TriggerDebugEventBus
 from core.trigger.debug.events import WebhookDebugEvent, build_webhook_pool_key
-from services.trigger.webhook_service import WebhookService
+from services.trigger.webhook_service import RawWebhookDataDict, WebhookService
 
 logger = logging.getLogger(__name__)
 
@@ -23,6 +23,7 @@ def _prepare_webhook_execution(webhook_id: str, is_debug: bool = False):
         webhook_id, is_debug=is_debug
     )
 
+    webhook_data: RawWebhookDataDict
     try:
         # Use new unified extraction and validation
         webhook_data = WebhookService.extract_and_validate_webhook_data(webhook_trigger, node_config)

api/controllers/web/wraps.py

@@ -1,7 +1,7 @@
 from collections.abc import Callable
 from datetime import UTC, datetime
 from functools import wraps
-from typing import Concatenate, ParamSpec, TypeVar
+from typing import Concatenate
 
 from flask import request
 from flask_restx import Resource
@@ -20,14 +20,13 @@ from services.enterprise.enterprise_service import EnterpriseService, WebAppSett
 from services.feature_service import FeatureService
 from services.webapp_auth_service import WebAppAuthService
 
-P = ParamSpec("P")
-R = TypeVar("R")
 
-
-def validate_jwt_token(view: Callable[Concatenate[App, EndUser, P], R] | None = None):
-    def decorator(view: Callable[Concatenate[App, EndUser, P], R]):
+def validate_jwt_token[**P, R](
+    view: Callable[Concatenate[App, EndUser, P], R] | None = None,
+) -> Callable[P, R] | Callable[[Callable[Concatenate[App, EndUser, P], R]], Callable[P, R]]:
+    def decorator(view: Callable[Concatenate[App, EndUser, P], R]) -> Callable[P, R]:
         @wraps(view)
-        def decorated(*args: P.args, **kwargs: P.kwargs):
+        def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
             app_model, end_user = decode_jwt_token()
             return view(app_model, end_user, *args, **kwargs)
 
@@ -38,7 +37,7 @@ def validate_jwt_token(view: Callable[Concatenate[App, EndUser, P], R] | None =
     return decorator
 
 
-def decode_jwt_token(app_code: str | None = None, user_id: str | None = None):
+def decode_jwt_token(app_code: str | None = None, user_id: str | None = None) -> tuple[App, EndUser]:
     system_features = FeatureService.get_system_features()
     if not app_code:
         app_code = str(request.headers.get(HEADER_NAME_APP_CODE))

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

@@ -5,7 +5,7 @@ import logging
 import threading
 import uuid
 from collections.abc import Generator, Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Literal, Union, overload
+from typing import TYPE_CHECKING, Any, Literal, overload
 
 from flask import Flask, current_app
 from pydantic import ValidationError
@@ -68,7 +68,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
         self,
         app_model: App,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         workflow_run_id: str,
@@ -81,7 +81,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
         self,
         app_model: App,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         workflow_run_id: str,
@@ -94,7 +94,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
         self,
         app_model: App,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         workflow_run_id: str,
@@ -106,7 +106,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
         self,
         app_model: App,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         workflow_run_id: str,
@@ -239,7 +239,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
         *,
         app_model: App,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         conversation: Conversation,
         message: Message,
         application_generate_entity: AdvancedChatAppGenerateEntity,
@@ -271,9 +271,9 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
         workflow: Workflow,
         node_id: str,
         user: Account | EndUser,
-        args: Mapping,
+        args: Mapping[str, Any],
         streaming: bool = True,
-    ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], Any, None]:
+    ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
         """
         Generate App response.
 
@@ -359,7 +359,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
         user: Account | EndUser,
         args: LoopNodeRunPayload,
         streaming: bool = True,
-    ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], Any, None]:
+    ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
         """
         Generate App response.
 
@@ -439,7 +439,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
         self,
         *,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         invoke_from: InvokeFrom,
         application_generate_entity: AdvancedChatAppGenerateEntity,
         workflow_execution_repository: WorkflowExecutionRepository,
@@ -451,7 +451,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
         pause_state_config: PauseStateLayerConfig | None = None,
         graph_runtime_state: GraphRuntimeState | None = None,
         graph_engine_layers: Sequence[GraphEngineLayer] = (),
-    ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], Any, None]:
+    ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
         """
         Generate App response.
 
@@ -653,10 +653,10 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
         queue_manager: AppQueueManager,
         conversation: ConversationSnapshot,
         message: MessageSnapshot,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         draft_var_saver_factory: DraftVariableSaverFactory,
         stream: bool = False,
-    ) -> Union[ChatbotAppBlockingResponse, Generator[ChatbotAppStreamResponse, None, None]]:
+    ) -> ChatbotAppBlockingResponse | Generator[ChatbotAppStreamResponse, None, None]:
         """
         Handle response.
         :param application_generate_entity: application generate entity

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

@@ -3,7 +3,7 @@ import logging
 import threading
 import uuid
 from collections.abc import Generator, Mapping
-from typing import Any, Literal, Union, overload
+from typing import Any, Literal, overload
 
 from flask import Flask, current_app
 from graphon.model_runtime.errors.invoke import InvokeAuthorizationError
@@ -37,7 +37,7 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
         self,
         *,
         app_model: App,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: Literal[False],
@@ -48,7 +48,7 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
         self,
         *,
         app_model: App,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: Literal[True],
@@ -59,21 +59,21 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
         self,
         *,
         app_model: App,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: bool,
-    ) -> Union[Mapping, Generator[Mapping | str, None, None]]: ...
+    ) -> Mapping | Generator[Mapping | str, None, None]: ...
 
     def generate(
         self,
         *,
         app_model: App,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: bool = True,
-    ) -> Union[Mapping, Generator[Mapping | str, None, None]]:
+    ) -> Mapping | Generator[Mapping | str, None, None]:
         """
         Generate App response.
 

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

@@ -3,7 +3,7 @@ import logging
 import threading
 import uuid
 from collections.abc import Generator, Mapping
-from typing import Any, Literal, Union, overload
+from typing import Any, Literal, overload
 
 from flask import Flask, copy_current_request_context, current_app
 from graphon.model_runtime.errors.invoke import InvokeAuthorizationError
@@ -36,7 +36,7 @@ class ChatAppGenerator(MessageBasedAppGenerator):
     def generate(
         self,
         app_model: App,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: Literal[True],
@@ -46,7 +46,7 @@ class ChatAppGenerator(MessageBasedAppGenerator):
     def generate(
         self,
         app_model: App,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: Literal[False],
@@ -56,20 +56,20 @@ class ChatAppGenerator(MessageBasedAppGenerator):
     def generate(
         self,
         app_model: App,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: bool,
-    ) -> Union[Mapping[str, Any], Generator[Mapping[str, Any] | str, None, None]]: ...
+    ) -> Mapping[str, Any] | Generator[Mapping[str, Any] | str, None, None]: ...
 
     def generate(
         self,
         app_model: App,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: bool = True,
-    ) -> Union[Mapping[str, Any], Generator[Mapping[str, Any] | str, None, None]]:
+    ) -> Mapping[str, Any] | Generator[Mapping[str, Any] | str, None, None]:
         """
         Generate App response.
 

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

@@ -3,7 +3,7 @@ import logging
 import threading
 import uuid
 from collections.abc import Generator, Mapping
-from typing import Any, Literal, Union, overload
+from typing import Any, Literal, overload
 
 from flask import Flask, copy_current_request_context, current_app
 from graphon.model_runtime.errors.invoke import InvokeAuthorizationError
@@ -36,7 +36,7 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
     def generate(
         self,
         app_model: App,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: Literal[True],
@@ -46,7 +46,7 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
     def generate(
         self,
         app_model: App,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: Literal[False],
@@ -56,20 +56,20 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
     def generate(
         self,
         app_model: App,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: bool = False,
-    ) -> Union[Mapping[str, Any], Generator[str | Mapping[str, Any], None, None]]: ...
+    ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]: ...
 
     def generate(
         self,
         app_model: App,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: bool = True,
-    ) -> Union[Mapping[str, Any], Generator[str | Mapping[str, Any], None, None]]:
+    ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
         """
         Generate App response.
 
@@ -244,10 +244,10 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
         self,
         app_model: App,
         message_id: str,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         invoke_from: InvokeFrom,
         stream: bool = True,
-    ) -> Union[Mapping, Generator[Mapping | str, None, None]]:
+    ) -> Mapping | Generator[Mapping | str, None, None]:
         """
         Generate App response.
 

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

@@ -7,7 +7,7 @@ import threading
 import time
 import uuid
 from collections.abc import Generator, Mapping
-from typing import Any, Literal, Union, cast, overload
+from typing import Any, Literal, cast, overload
 
 from flask import Flask, current_app
 from graphon.model_runtime.errors.invoke import InvokeAuthorizationError
@@ -62,7 +62,7 @@ class PipelineGenerator(BaseAppGenerator):
         *,
         pipeline: Pipeline,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: Literal[True],
@@ -77,7 +77,7 @@ class PipelineGenerator(BaseAppGenerator):
         *,
         pipeline: Pipeline,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: Literal[False],
@@ -92,28 +92,28 @@ class PipelineGenerator(BaseAppGenerator):
         *,
         pipeline: Pipeline,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: bool,
         call_depth: int,
         workflow_thread_pool_id: str | None,
         is_retry: bool = False,
-    ) -> Union[Mapping[str, Any], Generator[Mapping | str, None, None]]: ...
+    ) -> Mapping[str, Any] | Generator[Mapping | str, None, None]: ...
 
     def generate(
         self,
         *,
         pipeline: Pipeline,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: bool = True,
         call_depth: int = 0,
         workflow_thread_pool_id: str | None = None,
         is_retry: bool = False,
-    ) -> Union[Mapping[str, Any], Generator[Mapping | str, None, None], None]:
+    ) -> Mapping[str, Any] | Generator[Mapping | str, None, None] | None:
         # Add null check for dataset
 
         with Session(db.engine, expire_on_commit=False) as session:
@@ -278,7 +278,7 @@ class PipelineGenerator(BaseAppGenerator):
         context: contextvars.Context,
         pipeline: Pipeline,
         workflow_id: str,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         application_generate_entity: RagPipelineGenerateEntity,
         invoke_from: InvokeFrom,
         workflow_execution_repository: WorkflowExecutionRepository,
@@ -286,7 +286,7 @@ class PipelineGenerator(BaseAppGenerator):
         streaming: bool = True,
         variable_loader: VariableLoader = DUMMY_VARIABLE_LOADER,
         workflow_thread_pool_id: str | None = None,
-    ) -> Union[Mapping[str, Any], Generator[str | Mapping[str, Any], None, None]]:
+    ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
         """
         Generate App response.
 
@@ -302,7 +302,7 @@ class PipelineGenerator(BaseAppGenerator):
         """
         with preserve_flask_contexts(flask_app, context_vars=context):
             # init queue manager
-            workflow = db.session.query(Workflow).where(Workflow.id == workflow_id).first()
+            workflow = db.session.get(Workflow, workflow_id)
             if not workflow:
                 raise ValueError(f"Workflow not found: {workflow_id}")
             queue_manager = PipelineQueueManager(
@@ -624,10 +624,10 @@ class PipelineGenerator(BaseAppGenerator):
         application_generate_entity: RagPipelineGenerateEntity,
         workflow: Workflow,
         queue_manager: AppQueueManager,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         draft_var_saver_factory: DraftVariableSaverFactory,
         stream: bool = False,
-    ) -> Union[WorkflowAppBlockingResponse, Generator[WorkflowAppStreamResponse, None, None]]:
+    ) -> WorkflowAppBlockingResponse | Generator[WorkflowAppStreamResponse, None, None]:
         """
         Handle response.
         :param application_generate_entity: application generate entity
@@ -668,7 +668,7 @@ class PipelineGenerator(BaseAppGenerator):
         datasource_info: Mapping[str, Any],
         created_from: str,
         position: int,
-        account: Union[Account, EndUser],
+        account: Account | EndUser,
         batch: str,
         document_form: str,
     ):
@@ -715,7 +715,7 @@ class PipelineGenerator(BaseAppGenerator):
         pipeline: Pipeline,
         workflow: Workflow,
         start_node_id: str,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
     ) -> list[Mapping[str, Any]]:
         """
         Format datasource info list.

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

@@ -9,6 +9,7 @@ from graphon.graph_events import GraphEngineEvent, GraphRunFailedEvent
 from graphon.runtime import GraphRuntimeState, VariablePool
 from graphon.variable_loader import VariableLoader
 from graphon.variables.variables import RAGPipelineVariable, RAGPipelineVariableInput
+from sqlalchemy import select
 
 from core.app.apps.base_app_queue_manager import AppQueueManager
 from core.app.apps.pipeline.pipeline_config_manager import PipelineConfig
@@ -84,13 +85,13 @@ class PipelineRunner(WorkflowBasedAppRunner):
 
         user_id = None
         if invoke_from in {InvokeFrom.WEB_APP, InvokeFrom.SERVICE_API}:
-            end_user = db.session.query(EndUser).where(EndUser.id == self.application_generate_entity.user_id).first()
+            end_user = db.session.get(EndUser, self.application_generate_entity.user_id)
             if end_user:
                 user_id = end_user.session_id
         else:
             user_id = self.application_generate_entity.user_id
 
-        pipeline = db.session.query(Pipeline).where(Pipeline.id == app_config.app_id).first()
+        pipeline = db.session.get(Pipeline, app_config.app_id)
         if not pipeline:
             raise ValueError("Pipeline not found")
 
@@ -213,10 +214,10 @@ class PipelineRunner(WorkflowBasedAppRunner):
         Get workflow
         """
         # fetch workflow by workflow_id
-        workflow = (
-            db.session.query(Workflow)
+        workflow = db.session.scalar(
+            select(Workflow)
             .where(Workflow.tenant_id == pipeline.tenant_id, Workflow.app_id == pipeline.id, Workflow.id == workflow_id)
-            .first()
+            .limit(1)
         )
 
         # return workflow
@@ -297,10 +298,8 @@ class PipelineRunner(WorkflowBasedAppRunner):
         """
         if isinstance(event, GraphRunFailedEvent):
             if document_id and dataset_id:
-                document = (
-                    db.session.query(Document)
-                    .where(Document.id == document_id, Document.dataset_id == dataset_id)
-                    .first()
+                document = db.session.scalar(
+                    select(Document).where(Document.id == document_id, Document.dataset_id == dataset_id).limit(1)
                 )
                 if document:
                     document.indexing_status = "error"

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

@@ -5,7 +5,7 @@ import logging
 import threading
 import uuid
 from collections.abc import Generator, Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Literal, Union, overload
+from typing import TYPE_CHECKING, Any, Literal, overload
 
 from flask import Flask, current_app
 from graphon.graph_engine.layers import GraphEngineLayer
@@ -64,7 +64,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
         *,
         app_model: App,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: Literal[True],
@@ -82,7 +82,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
         *,
         app_model: App,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: Literal[False],
@@ -100,7 +100,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
         *,
         app_model: App,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: bool,
@@ -110,14 +110,14 @@ class WorkflowAppGenerator(BaseAppGenerator):
         root_node_id: str | None = None,
         graph_engine_layers: Sequence[GraphEngineLayer] = (),
         pause_state_config: PauseStateLayerConfig | None = None,
-    ) -> Union[Mapping[str, Any], Generator[Mapping[str, Any] | str, None, None]]: ...
+    ) -> Mapping[str, Any] | Generator[Mapping[str, Any] | str, None, None]: ...
 
     def generate(
         self,
         *,
         app_model: App,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         args: Mapping[str, Any],
         invoke_from: InvokeFrom,
         streaming: bool = True,
@@ -127,7 +127,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
         root_node_id: str | None = None,
         graph_engine_layers: Sequence[GraphEngineLayer] = (),
         pause_state_config: PauseStateLayerConfig | None = None,
-    ) -> Union[Mapping[str, Any], Generator[Mapping[str, Any] | str, None, None]]:
+    ) -> Mapping[str, Any] | Generator[Mapping[str, Any] | str, None, None]:
         with self._bind_file_access_scope(tenant_id=app_model.tenant_id, user=user, invoke_from=invoke_from):
             files: Sequence[Mapping[str, Any]] = args.get("files") or []
 
@@ -237,7 +237,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
         *,
         app_model: App,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         application_generate_entity: WorkflowAppGenerateEntity,
         graph_runtime_state: GraphRuntimeState,
         workflow_execution_repository: WorkflowExecutionRepository,
@@ -245,7 +245,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
         graph_engine_layers: Sequence[GraphEngineLayer] = (),
         pause_state_config: PauseStateLayerConfig | None = None,
         variable_loader: VariableLoader = DUMMY_VARIABLE_LOADER,
-    ) -> Union[Mapping[str, Any], Generator[str | Mapping[str, Any], None, None]]:
+    ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
         """
         Resume a paused workflow execution using the persisted runtime state.
         """
@@ -269,7 +269,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
         *,
         app_model: App,
         workflow: Workflow,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         application_generate_entity: WorkflowAppGenerateEntity,
         invoke_from: InvokeFrom,
         workflow_execution_repository: WorkflowExecutionRepository,
@@ -280,7 +280,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
         graph_engine_layers: Sequence[GraphEngineLayer] = (),
         graph_runtime_state: GraphRuntimeState | None = None,
         pause_state_config: PauseStateLayerConfig | None = None,
-    ) -> Union[Mapping[str, Any], Generator[str | Mapping[str, Any], None, None]]:
+    ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
         """
         Generate App response.
 
@@ -609,10 +609,10 @@ class WorkflowAppGenerator(BaseAppGenerator):
         application_generate_entity: WorkflowAppGenerateEntity,
         workflow: Workflow,
         queue_manager: AppQueueManager,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         draft_var_saver_factory: DraftVariableSaverFactory,
         stream: bool = False,
-    ) -> Union[WorkflowAppBlockingResponse, Generator[WorkflowAppStreamResponse, None, None]]:
+    ) -> WorkflowAppBlockingResponse | Generator[WorkflowAppStreamResponse, None, None]:
         """
         Handle response.
         :param application_generate_entity: application generate entity

api/core/app/layers/pause_state_persist_layer.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Annotated, Literal, Self, TypeAlias
+from typing import Annotated, Literal, Self
 
 from graphon.graph_engine.layers import GraphEngineLayer
 from graphon.graph_events import GraphEngineEvent, GraphRunPausedEvent
@@ -27,7 +27,7 @@ class _AdvancedChatAppGenerateEntityWrapper(BaseModel):
     entity: AdvancedChatAppGenerateEntity
 
 
-_GenerateEntityUnion: TypeAlias = Annotated[
+type _GenerateEntityUnion = Annotated[
     _WorkflowGenerateEntityWrapper | _AdvancedChatAppGenerateEntityWrapper,
     Field(discriminator="type"),
 ]

api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py

@@ -2,7 +2,7 @@ import logging
 import time
 from collections.abc import Generator
 from threading import Thread
-from typing import Any, Union, cast
+from typing import Any, cast
 
 from graphon.file import FileTransferMethod
 from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
@@ -72,14 +72,12 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline):
     """
 
     _task_state: EasyUITaskState
-    _application_generate_entity: Union[ChatAppGenerateEntity, CompletionAppGenerateEntity, AgentChatAppGenerateEntity]
+    _application_generate_entity: ChatAppGenerateEntity | CompletionAppGenerateEntity | AgentChatAppGenerateEntity
     _precomputed_event_type: StreamEvent | None = None
 
     def __init__(
         self,
-        application_generate_entity: Union[
-            ChatAppGenerateEntity, CompletionAppGenerateEntity, AgentChatAppGenerateEntity
-        ],
+        application_generate_entity: ChatAppGenerateEntity | CompletionAppGenerateEntity | AgentChatAppGenerateEntity,
         queue_manager: AppQueueManager,
         conversation: Conversation,
         message: Message,
@@ -117,11 +115,11 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline):
 
     def process(
         self,
-    ) -> Union[
-        ChatbotAppBlockingResponse,
-        CompletionAppBlockingResponse,
-        Generator[Union[ChatbotAppStreamResponse, CompletionAppStreamResponse], None, None],
-    ]:
+    ) -> (
+        ChatbotAppBlockingResponse
+        | CompletionAppBlockingResponse
+        | Generator[ChatbotAppStreamResponse | CompletionAppStreamResponse, None, None]
+    ):
         if self._application_generate_entity.app_config.app_mode != AppMode.COMPLETION:
             # start generate conversation name thread
             self._conversation_name_generate_thread = self._message_cycle_manager.generate_conversation_name(
@@ -136,7 +134,7 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline):
 
     def _to_blocking_response(
         self, generator: Generator[StreamResponse, None, None]
-    ) -> Union[ChatbotAppBlockingResponse, CompletionAppBlockingResponse]:
+    ) -> ChatbotAppBlockingResponse | CompletionAppBlockingResponse:
         """
         Process blocking response.
         :return:
@@ -148,7 +146,7 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline):
                 extras = {"usage": self._task_state.llm_result.usage.model_dump()}
                 if self._task_state.metadata:
                     extras["metadata"] = self._task_state.metadata.model_dump()
-                response: Union[ChatbotAppBlockingResponse, CompletionAppBlockingResponse]
+                response: ChatbotAppBlockingResponse | CompletionAppBlockingResponse
                 if self._conversation_mode == AppMode.COMPLETION:
                     response = CompletionAppBlockingResponse(
                         task_id=self._application_generate_entity.task_id,
@@ -183,7 +181,7 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline):
 
     def _to_stream_response(
         self, generator: Generator[StreamResponse, None, None]
-    ) -> Generator[Union[ChatbotAppStreamResponse, CompletionAppStreamResponse], None, None]:
+    ) -> Generator[ChatbotAppStreamResponse | CompletionAppStreamResponse, None, None]:
         """
         To stream response.
         :return:

api/core/app/workflow/layers/llm_quota.py

@@ -5,14 +5,13 @@ This layer centralizes model-quota deduction outside node implementations.
 """
 
 import logging
-from typing import TYPE_CHECKING, cast, final
+from typing import TYPE_CHECKING, cast, final, override
 
 from graphon.enums import BuiltinNodeTypes
 from graphon.graph_engine.entities.commands import AbortCommand, CommandType
 from graphon.graph_engine.layers import GraphEngineLayer
 from graphon.graph_events import GraphEngineEvent, GraphNodeEventBase, NodeRunSucceededEvent
 from graphon.nodes.base.node import Node
-from typing_extensions import override
 
 from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, DifyRunContext
 from core.app.llm import deduct_llm_quota, ensure_llm_quota_available

api/core/app/workflow/layers/observability.py

@@ -10,7 +10,7 @@ associates with the node span.
 import logging
 from contextvars import Token
 from dataclasses import dataclass
-from typing import cast, final
+from typing import cast, final, override
 
 from graphon.enums import BuiltinNodeTypes, NodeType
 from graphon.graph_engine.layers import GraphEngineLayer
@@ -18,7 +18,6 @@ from graphon.graph_events import GraphNodeEventBase
 from graphon.nodes.base.node import Node
 from opentelemetry import context as context_api
 from opentelemetry.trace import Span, SpanKind, Tracer, get_tracer, set_span_in_context
-from typing_extensions import override
 
 from configs import dify_config
 from extensions.otel.parser import (

api/core/datasource/datasource_file_manager.py

@@ -153,7 +153,7 @@ class DatasourceFileManager:
 
         :return: the binary of the file, mime type
         """
-        upload_file: UploadFile | None = db.session.query(UploadFile).where(UploadFile.id == id).first()
+        upload_file: UploadFile | None = db.session.get(UploadFile, id)
 
         if not upload_file:
             return None
@@ -171,7 +171,7 @@ class DatasourceFileManager:
 
         :return: the binary of the file, mime type
         """
-        message_file: MessageFile | None = db.session.query(MessageFile).where(MessageFile.id == id).first()
+        message_file: MessageFile | None = db.session.get(MessageFile, id)
 
         # Check if message_file is not None
         if message_file is not None:
@@ -185,7 +185,7 @@ class DatasourceFileManager:
         else:
             tool_file_id = None
 
-        tool_file: ToolFile | None = db.session.query(ToolFile).where(ToolFile.id == tool_file_id).first()
+        tool_file: ToolFile | None = db.session.get(ToolFile, tool_file_id)
 
         if not tool_file:
             return None
@@ -203,7 +203,7 @@ class DatasourceFileManager:
 
         :return: the binary of the file, mime type
         """
-        upload_file: UploadFile | None = db.session.query(UploadFile).where(UploadFile.id == upload_file_id).first()
+        upload_file: UploadFile | None = db.session.get(UploadFile, upload_file_id)
 
         if not upload_file:
             return None, None

api/core/entities/execution_extra_content.py

@@ -44,7 +44,8 @@ class HumanInputContent(BaseModel):
     type: ExecutionContentType = Field(default=ExecutionContentType.HUMAN_INPUT)
 
 
-ExecutionExtraContentDomainModel: TypeAlias = HumanInputContent
+# Keep a runtime alias here: callers and tests expect identity with HumanInputContent.
+ExecutionExtraContentDomainModel: TypeAlias = HumanInputContent  # noqa: UP040
 
 __all__ = [
     "ExecutionExtraContentDomainModel",

api/core/helper/module_import_helper.py

@@ -2,12 +2,13 @@ import importlib.util
 import logging
 import sys
 from types import ModuleType
-from typing import AnyStr
 
 logger = logging.getLogger(__name__)
 
 
-def import_module_from_source(*, module_name: str, py_file_path: AnyStr, use_lazy_loader: bool = False) -> ModuleType:
+def import_module_from_source[T: (str, bytes)](
+    *, module_name: str, py_file_path: T, use_lazy_loader: bool = False
+) -> ModuleType:
     """
     Importing a module from the source file directly
     """

api/core/helper/position_helper.py

@@ -2,7 +2,6 @@ import os
 from collections import OrderedDict
 from collections.abc import Callable
 from functools import lru_cache
-from typing import TypeVar
 
 from configs import dify_config
 from core.tools.utils.yaml_utils import load_yaml_file_cached
@@ -65,10 +64,7 @@ def pin_position_map(original_position_map: dict[str, int], pin_list: list[str])
     return position_map
 
 
-T = TypeVar("T")
-
-
-def is_filtered(
+def is_filtered[T](
     include_set: set[str],
     exclude_set: set[str],
     data: T,
@@ -97,11 +93,11 @@ def is_filtered(
     return False
 
 
-def sort_by_position_map(
+def sort_by_position_map[T](
     position_map: dict[str, int],
     data: list[T],
     name_func: Callable[[T], str],
-):
+) -> list[T]:
     """
     Sort the objects by the position map.
     If the name of the object is not in the position map, it will be put at the end.
@@ -116,11 +112,11 @@ def sort_by_position_map(
     return sorted(data, key=lambda x: position_map.get(name_func(x), float("inf")))
 
 
-def sort_to_dict_by_position_map(
+def sort_to_dict_by_position_map[T](
     position_map: dict[str, int],
     data: list[T],
     name_func: Callable[[T], str],
-):
+) -> OrderedDict[str, T]:
     """
     Sort the objects into a ordered dict by the position map.
     If the name of the object is not in the position map, it will be put at the end.

api/core/helper/ssrf_proxy.py

@@ -4,7 +4,7 @@ Proxy requests to avoid SSRF
 
 import logging
 import time
-from typing import Any, TypeAlias
+from typing import Any
 
 import httpx
 from pydantic import TypeAdapter, ValidationError
@@ -20,8 +20,8 @@ SSRF_DEFAULT_MAX_RETRIES = dify_config.SSRF_DEFAULT_MAX_RETRIES
 BACKOFF_FACTOR = 0.5
 STATUS_FORCELIST = [429, 500, 502, 503, 504]
 
-Headers: TypeAlias = dict[str, str]
-_HEADERS_ADAPTER = TypeAdapter(Headers)
+type Headers = dict[str, str]
+_HEADERS_ADAPTER: TypeAdapter[Headers] = TypeAdapter(Headers)
 
 _SSL_VERIFIED_POOL_KEY = "ssrf:verified"
 _SSL_UNVERIFIED_POOL_KEY = "ssrf:unverified"

api/core/indexing_runner.py

@@ -10,7 +10,7 @@ from typing import Any
 
 from flask import Flask, current_app
 from graphon.model_runtime.entities.model_entities import ModelType
-from sqlalchemy import select
+from sqlalchemy import delete, func, select, update
 from sqlalchemy.orm.exc import ObjectDeletedError
 
 from configs import dify_config
@@ -78,7 +78,7 @@ class IndexingRunner:
                     continue
 
                 # get dataset
-                dataset = db.session.query(Dataset).filter_by(id=requeried_document.dataset_id).first()
+                dataset = db.session.get(Dataset, requeried_document.dataset_id)
 
                 if not dataset:
                     raise ValueError("no dataset found")
@@ -95,7 +95,7 @@ class IndexingRunner:
                 text_docs = self._extract(index_processor, requeried_document, processing_rule.to_dict())
 
                 # transform
-                current_user = db.session.query(Account).filter_by(id=requeried_document.created_by).first()
+                current_user = db.session.get(Account, requeried_document.created_by)
                 if not current_user:
                     raise ValueError("no current user found")
                 current_user.set_tenant_id(dataset.tenant_id)
@@ -137,23 +137,24 @@ class IndexingRunner:
                 return
 
             # get dataset
-            dataset = db.session.query(Dataset).filter_by(id=requeried_document.dataset_id).first()
+            dataset = db.session.get(Dataset, requeried_document.dataset_id)
 
             if not dataset:
                 raise ValueError("no dataset found")
 
             # get exist document_segment list and delete
-            document_segments = (
-                db.session.query(DocumentSegment)
-                .filter_by(dataset_id=dataset.id, document_id=requeried_document.id)
-                .all()
-            )
+            document_segments = db.session.scalars(
+                select(DocumentSegment).where(
+                    DocumentSegment.dataset_id == dataset.id,
+                    DocumentSegment.document_id == requeried_document.id,
+                )
+            ).all()
 
             for document_segment in document_segments:
                 db.session.delete(document_segment)
                 if requeried_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
                     # delete child chunks
-                    db.session.query(ChildChunk).where(ChildChunk.segment_id == document_segment.id).delete()
+                    db.session.execute(delete(ChildChunk).where(ChildChunk.segment_id == document_segment.id))
             db.session.commit()
             # get the process rule
             stmt = select(DatasetProcessRule).where(DatasetProcessRule.id == requeried_document.dataset_process_rule_id)
@@ -167,7 +168,7 @@ class IndexingRunner:
             text_docs = self._extract(index_processor, requeried_document, processing_rule.to_dict())
 
             # transform
-            current_user = db.session.query(Account).filter_by(id=requeried_document.created_by).first()
+            current_user = db.session.get(Account, requeried_document.created_by)
             if not current_user:
                 raise ValueError("no current user found")
             current_user.set_tenant_id(dataset.tenant_id)
@@ -207,17 +208,18 @@ class IndexingRunner:
                 return
 
             # get dataset
-            dataset = db.session.query(Dataset).filter_by(id=requeried_document.dataset_id).first()
+            dataset = db.session.get(Dataset, requeried_document.dataset_id)
 
             if not dataset:
                 raise ValueError("no dataset found")
 
             # get exist document_segment list and delete
-            document_segments = (
-                db.session.query(DocumentSegment)
-                .filter_by(dataset_id=dataset.id, document_id=requeried_document.id)
-                .all()
-            )
+            document_segments = db.session.scalars(
+                select(DocumentSegment).where(
+                    DocumentSegment.dataset_id == dataset.id,
+                    DocumentSegment.document_id == requeried_document.id,
+                )
+            ).all()
 
             documents = []
             if document_segments:
@@ -289,7 +291,7 @@ class IndexingRunner:
 
         embedding_model_instance = None
         if dataset_id:
-            dataset = db.session.query(Dataset).filter_by(id=dataset_id).first()
+            dataset = db.session.get(Dataset, dataset_id)
             if not dataset:
                 raise ValueError("Dataset not found.")
             if IndexTechniqueType.HIGH_QUALITY in {dataset.indexing_technique, indexing_technique}:
@@ -652,24 +654,26 @@ class IndexingRunner:
     @staticmethod
     def _process_keyword_index(flask_app, dataset_id, document_id, documents):
         with flask_app.app_context():
-            dataset = db.session.query(Dataset).filter_by(id=dataset_id).first()
+            dataset = db.session.get(Dataset, dataset_id)
             if not dataset:
                 raise ValueError("no dataset found")
             keyword = Keyword(dataset)
             keyword.create(documents)
             if dataset.indexing_technique != IndexTechniqueType.HIGH_QUALITY:
                 document_ids = [document.metadata["doc_id"] for document in documents]
-                db.session.query(DocumentSegment).where(
-                    DocumentSegment.document_id == document_id,
-                    DocumentSegment.dataset_id == dataset_id,
-                    DocumentSegment.index_node_id.in_(document_ids),
-                    DocumentSegment.status == SegmentStatus.INDEXING,
-                ).update(
-                    {
-                        DocumentSegment.status: SegmentStatus.COMPLETED,
-                        DocumentSegment.enabled: True,
-                        DocumentSegment.completed_at: naive_utc_now(),
-                    }
+                db.session.execute(
+                    update(DocumentSegment)
+                    .where(
+                        DocumentSegment.document_id == document_id,
+                        DocumentSegment.dataset_id == dataset_id,
+                        DocumentSegment.index_node_id.in_(document_ids),
+                        DocumentSegment.status == SegmentStatus.INDEXING,
+                    )
+                    .values(
+                        status=SegmentStatus.COMPLETED,
+                        enabled=True,
+                        completed_at=naive_utc_now(),
+                    )
                 )
 
                 db.session.commit()
@@ -703,17 +707,19 @@ class IndexingRunner:
             )
 
             document_ids = [document.metadata["doc_id"] for document in chunk_documents]
-            db.session.query(DocumentSegment).where(
-                DocumentSegment.document_id == dataset_document.id,
-                DocumentSegment.dataset_id == dataset.id,
-                DocumentSegment.index_node_id.in_(document_ids),
-                DocumentSegment.status == SegmentStatus.INDEXING,
-            ).update(
-                {
-                    DocumentSegment.status: SegmentStatus.COMPLETED,
-                    DocumentSegment.enabled: True,
-                    DocumentSegment.completed_at: naive_utc_now(),
-                }
+            db.session.execute(
+                update(DocumentSegment)
+                .where(
+                    DocumentSegment.document_id == dataset_document.id,
+                    DocumentSegment.dataset_id == dataset.id,
+                    DocumentSegment.index_node_id.in_(document_ids),
+                    DocumentSegment.status == SegmentStatus.INDEXING,
+                )
+                .values(
+                    status=SegmentStatus.COMPLETED,
+                    enabled=True,
+                    completed_at=naive_utc_now(),
+                )
             )
 
             db.session.commit()
@@ -734,10 +740,17 @@ class IndexingRunner:
         """
         Update the document indexing status.
         """
-        count = db.session.query(DatasetDocument).filter_by(id=document_id, is_paused=True).count()
+        count = (
+            db.session.scalar(
+                select(func.count())
+                .select_from(DatasetDocument)
+                .where(DatasetDocument.id == document_id, DatasetDocument.is_paused == True)
+            )
+            or 0
+        )
         if count > 0:
             raise DocumentIsPausedError()
-        document = db.session.query(DatasetDocument).filter_by(id=document_id).first()
+        document = db.session.get(DatasetDocument, document_id)
         if not document:
             raise DocumentIsDeletedPausedError()
 
@@ -745,7 +758,7 @@ class IndexingRunner:
 
         if extra_update_params:
             update_params.update(extra_update_params)
-        db.session.query(DatasetDocument).filter_by(id=document_id).update(update_params)  # type: ignore
+        db.session.execute(update(DatasetDocument).where(DatasetDocument.id == document_id).values(update_params))  # type: ignore
         db.session.commit()
 
     @staticmethod
@@ -753,7 +766,9 @@ class IndexingRunner:
         """
         Update the document segment by document id.
         """
-        db.session.query(DocumentSegment).filter_by(document_id=dataset_document_id).update(update_params)
+        db.session.execute(
+            update(DocumentSegment).where(DocumentSegment.document_id == dataset_document_id).values(update_params)
+        )
         db.session.commit()
 
     def _transform(

api/core/logging/structured_formatter.py

@@ -3,13 +3,19 @@
 import logging
 import traceback
 from datetime import UTC, datetime
-from typing import Any
+from typing import Any, TypedDict
 
 import orjson
 
 from configs import dify_config
 
 
+class IdentityDict(TypedDict, total=False):
+    tenant_id: str
+    user_id: str
+    user_type: str
+
+
 class StructuredJSONFormatter(logging.Formatter):
     """
     JSON log formatter following the specified schema:
@@ -84,7 +90,7 @@ class StructuredJSONFormatter(logging.Formatter):
 
         return log_dict
 
-    def _extract_identity(self, record: logging.LogRecord) -> dict[str, str] | None:
+    def _extract_identity(self, record: logging.LogRecord) -> IdentityDict | None:
         tenant_id = getattr(record, "tenant_id", None)
         user_id = getattr(record, "user_id", None)
         user_type = getattr(record, "user_type", None)
@@ -92,7 +98,7 @@ class StructuredJSONFormatter(logging.Formatter):
         if not any([tenant_id, user_id, user_type]):
             return None
 
-        identity: dict[str, str] = {}
+        identity: IdentityDict = {}
         if tenant_id:
             identity["tenant_id"] = tenant_id
         if user_id:

api/core/mcp/client/sse_client.py

@@ -3,7 +3,7 @@ import queue
 from collections.abc import Generator
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
-from typing import Any, TypeAlias, final
+from typing import Any, final
 from urllib.parse import urljoin, urlparse
 
 import httpx
@@ -33,9 +33,9 @@ class _StatusError:
 
 
 # Type aliases for better readability
-ReadQueue: TypeAlias = queue.Queue[SessionMessage | Exception | None]
-WriteQueue: TypeAlias = queue.Queue[SessionMessage | Exception | None]
-StatusQueue: TypeAlias = queue.Queue[_StatusReady | _StatusError]
+type ReadQueue = queue.Queue[SessionMessage | Exception | None]
+type WriteQueue = queue.Queue[SessionMessage | Exception | None]
+type StatusQueue = queue.Queue[_StatusReady | _StatusError]
 
 
 class SSETransport:

api/core/mcp/entities.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from enum import StrEnum
-from typing import Any, Generic, TypeVar
+from typing import Any, TypeVar
 
 from pydantic import BaseModel
 
@@ -9,13 +9,12 @@ from core.mcp.types import LATEST_PROTOCOL_VERSION, OAuthClientInformation, OAut
 
 SUPPORTED_PROTOCOL_VERSIONS: list[str] = ["2024-11-05", "2025-03-26", LATEST_PROTOCOL_VERSION]
 
-
 SessionT = TypeVar("SessionT", bound=BaseSession[Any, Any, Any, Any, Any])
 LifespanContextT = TypeVar("LifespanContextT")
 
 
 @dataclass
-class RequestContext(Generic[SessionT, LifespanContextT]):
+class RequestContext[SessionT: BaseSession[Any, Any, Any, Any, Any], LifespanContextT]:
     request_id: RequestId
     meta: RequestParams.Meta | None
     session: SessionT

api/core/mcp/session/base_session.py

@@ -4,7 +4,7 @@ from collections.abc import Callable
 from concurrent.futures import Future, ThreadPoolExecutor, TimeoutError
 from datetime import timedelta
 from types import TracebackType
-from typing import Any, Generic, Self, TypeVar
+from typing import Any, Self
 
 from httpx import HTTPStatusError
 from pydantic import BaseModel
@@ -34,16 +34,10 @@ from core.mcp.types import (
 logger = logging.getLogger(__name__)
 
 
-SendRequestT = TypeVar("SendRequestT", ClientRequest, ServerRequest)
-SendResultT = TypeVar("SendResultT", ClientResult, ServerResult)
-SendNotificationT = TypeVar("SendNotificationT", ClientNotification, ServerNotification)
-ReceiveRequestT = TypeVar("ReceiveRequestT", ClientRequest, ServerRequest)
-ReceiveResultT = TypeVar("ReceiveResultT", bound=BaseModel)
-ReceiveNotificationT = TypeVar("ReceiveNotificationT", ClientNotification, ServerNotification)
 DEFAULT_RESPONSE_READ_TIMEOUT = 1.0
 
 
-class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
+class RequestResponder[ReceiveRequestT: ClientRequest | ServerRequest, SendResultT: ClientResult | ServerResult]:
     """Handles responding to MCP requests and manages request lifecycle.
 
     This class MUST be used as a context manager to ensure proper cleanup and
@@ -60,7 +54,7 @@ class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
     """
 
     request: ReceiveRequestT
-    _session: Any
+    _session: "BaseSession[Any, Any, SendResultT, ReceiveRequestT, Any]"
     _on_complete: Callable[["RequestResponder[ReceiveRequestT, SendResultT]"], Any]
 
     def __init__(
@@ -68,7 +62,7 @@ class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
         request_id: RequestId,
         request_meta: RequestParams.Meta | None,
         request: ReceiveRequestT,
-        session: """BaseSession[SendRequestT, SendNotificationT, SendResultT, ReceiveRequestT, ReceiveNotificationT]""",
+        session: "BaseSession[Any, Any, SendResultT, ReceiveRequestT, Any]",
         on_complete: Callable[["RequestResponder[ReceiveRequestT, SendResultT]"], Any],
     ):
         self.request_id = request_id
@@ -111,7 +105,7 @@ class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
 
         self.completed = True
 
-        self._session._send_response(request_id=self.request_id, response=response)
+        self._session.send_response(request_id=self.request_id, response=response)
 
     def cancel(self):
         """Cancel this request and mark it as completed."""
@@ -120,21 +114,19 @@ class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
 
         self.completed = True  # Mark as completed so it's removed from in_flight
         # Send an error response to indicate cancellation
-        self._session._send_response(
+        self._session.send_response(
             request_id=self.request_id,
             response=ErrorData(code=0, message="Request cancelled", data=None),
         )
 
 
-class BaseSession(
-    Generic[
-        SendRequestT,
-        SendNotificationT,
-        SendResultT,
-        ReceiveRequestT,
-        ReceiveNotificationT,
-    ],
-):
+class BaseSession[
+    SendRequestT: ClientRequest | ServerRequest,
+    SendNotificationT: ClientNotification | ServerNotification,
+    SendResultT: ClientResult | ServerResult,
+    ReceiveRequestT: ClientRequest | ServerRequest,
+    ReceiveNotificationT: ClientNotification | ServerNotification,
+]:
     """
     Implements an MCP "session" on top of read/write streams, including features
     like request/response linking, notifications, and progress.
@@ -204,13 +196,13 @@ class BaseSession(
             # The receiver thread should have already exited due to the None message in the queue
             self._executor.shutdown(wait=False)
 
-    def send_request(
+    def send_request[T: BaseModel](
         self,
         request: SendRequestT,
-        result_type: type[ReceiveResultT],
+        result_type: type[T],
         request_read_timeout_seconds: timedelta | None = None,
         metadata: MessageMetadata | None = None,
-    ) -> ReceiveResultT:
+    ) -> T:
         """
         Sends a request and wait for a response. Raises an McpError if the
         response contains an error. If a request read timeout is provided, it
@@ -299,7 +291,7 @@ class BaseSession(
         )
         self._write_stream.put(session_message)
 
-    def _send_response(self, request_id: RequestId, response: SendResultT | ErrorData):
+    def send_response(self, request_id: RequestId, response: SendResultT | ErrorData):
         if isinstance(response, ErrorData):
             jsonrpc_error = JSONRPCError(jsonrpc="2.0", id=request_id, error=response)
             session_message = SessionMessage(message=JSONRPCMessage(jsonrpc_error))
@@ -350,7 +342,7 @@ class BaseSession(
                     responder = RequestResponder[ReceiveRequestT, SendResultT](
                         request_id=message.message.root.id,
                         request_meta=validated_request.root.params.meta if validated_request.root.params else None,
-                        request=validated_request,
+                        request=validated_request,  # type: ignore[arg-type]  # mypy can't narrow constrained TypeVar from model_validate
                         session=self,
                         on_complete=lambda r: self._in_flight.pop(r.request_id, None),
                     )
@@ -372,8 +364,8 @@ class BaseSession(
                             if cancelled_id in self._in_flight:
                                 self._in_flight[cancelled_id].cancel()
                         else:
-                            self._received_notification(notification)
-                            self._handle_incoming(notification)
+                            self._received_notification(notification)  # type: ignore[arg-type]
+                            self._handle_incoming(notification)  # type: ignore[arg-type]
                     except Exception as e:
                         # For other validation errors, log and continue
                         logger.warning("Failed to validate notification: %s. Message was: %s", e, message.message.root)

api/core/mcp/types.py

@@ -1,6 +1,6 @@
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Annotated, Any, Generic, Literal, TypeAlias, TypeVar
+from typing import Annotated, Any, Literal
 
 from pydantic import BaseModel, ConfigDict, Field, FileUrl, RootModel
 from pydantic.networks import AnyUrl, UrlConstraints
@@ -31,7 +31,7 @@ ProgressToken = str | int
 Cursor = str
 Role = Literal["user", "assistant"]
 RequestId = Annotated[int | str, Field(union_mode="left_to_right")]
-AnyFunction: TypeAlias = Callable[..., Any]
+type AnyFunction = Callable[..., Any]
 
 
 class RequestParams(BaseModel):
@@ -68,12 +68,7 @@ class NotificationParams(BaseModel):
     """
 
 
-RequestParamsT = TypeVar("RequestParamsT", bound=RequestParams | dict[str, Any] | None)
-NotificationParamsT = TypeVar("NotificationParamsT", bound=NotificationParams | dict[str, Any] | None)
-MethodT = TypeVar("MethodT", bound=str)
-
-
-class Request(BaseModel, Generic[RequestParamsT, MethodT]):
+class Request[RequestParamsT: RequestParams | dict[str, Any] | None, MethodT: str](BaseModel):
     """Base class for JSON-RPC requests."""
 
     method: MethodT
@@ -81,14 +76,14 @@ class Request(BaseModel, Generic[RequestParamsT, MethodT]):
     model_config = ConfigDict(extra="allow")
 
 
-class PaginatedRequest(Request[PaginatedRequestParams | None, MethodT], Generic[MethodT]):
+class PaginatedRequest[T: str](Request[PaginatedRequestParams | None, T]):
     """Base class for paginated requests,
     matching the schema's PaginatedRequest interface."""
 
     params: PaginatedRequestParams | None = None
 
 
-class Notification(BaseModel, Generic[NotificationParamsT, MethodT]):
+class Notification[NotificationParamsT: NotificationParams | dict[str, Any] | None, MethodT: str](BaseModel):
     """Base class for JSON-RPC notifications."""
 
     method: MethodT
@@ -736,7 +731,7 @@ class ResourceLink(Resource):
 ContentBlock = TextContent | ImageContent | AudioContent | ResourceLink | EmbeddedResource
 """A content block that can be used in prompts and tool results."""
 
-Content: TypeAlias = ContentBlock
+type Content = ContentBlock
 # """DEPRECATED: Content is deprecated, you should use ContentBlock directly."""
 
 

api/core/ops/arize_phoenix_trace/arize_phoenix_trace.py

@@ -38,6 +38,7 @@ from core.ops.entities.trace_entity import (
     TraceTaskName,
     WorkflowTraceInfo,
 )
+from core.ops.utils import JSON_DICT_ADAPTER
 from core.repositories import DifyCoreRepositoryFactory
 from extensions.ext_database import db
 from models.model import EndUser, MessageFile
@@ -469,7 +470,7 @@ class ArizePhoenixDataTrace(BaseTraceInstance):
                 llm_attributes[SpanAttributes.LLM_PROVIDER] = trace_info.message_data.model_provider
 
             if trace_info.message_data and trace_info.message_data.message_metadata:
-                metadata_dict = json.loads(trace_info.message_data.message_metadata)
+                metadata_dict = JSON_DICT_ADAPTER.validate_json(trace_info.message_data.message_metadata)
                 if model_params := metadata_dict.get("model_parameters"):
                     llm_attributes[SpanAttributes.LLM_INVOCATION_PARAMETERS] = json.dumps(model_params)
 

api/core/ops/mlflow_trace/mlflow_trace.py

@@ -1,4 +1,3 @@
-import json
 import logging
 import os
 from datetime import datetime, timedelta
@@ -25,6 +24,7 @@ from core.ops.entities.trace_entity import (
     TraceTaskName,
     WorkflowTraceInfo,
 )
+from core.ops.utils import JSON_DICT_ADAPTER
 from extensions.ext_database import db
 from models import EndUser
 from models.workflow import WorkflowNodeExecutionModel
@@ -153,7 +153,7 @@ class MLflowDataTrace(BaseTraceInstance):
                     inputs = node.process_data  # contains request URL
 
                 if not inputs:
-                    inputs = json.loads(node.inputs) if node.inputs else {}
+                    inputs = JSON_DICT_ADAPTER.validate_json(node.inputs) if node.inputs else {}
 
                 node_span = start_span_no_context(
                     name=node.title,
@@ -180,7 +180,7 @@ class MLflowDataTrace(BaseTraceInstance):
 
                 # End node span
                 finished_at = node.created_at + timedelta(seconds=node.elapsed_time)
-                outputs = json.loads(node.outputs) if node.outputs else {}
+                outputs = JSON_DICT_ADAPTER.validate_json(node.outputs) if node.outputs else {}
                 if node.node_type == BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL:
                     outputs = self._parse_knowledge_retrieval_outputs(outputs)
                 elif node.node_type == BuiltinNodeTypes.LLM:
@@ -216,8 +216,8 @@ class MLflowDataTrace(BaseTraceInstance):
             return {}, {}
 
         try:
-            data = json.loads(node.process_data)
-        except (json.JSONDecodeError, TypeError):
+            data = JSON_DICT_ADAPTER.validate_json(node.process_data)
+        except (ValueError, TypeError):
             return {}, {}
 
         inputs = self._parse_prompts(data.get("prompts"))

api/core/ops/ops_trace_manager.py

@@ -6,17 +6,19 @@ import queue
 import threading
 import time
 from datetime import timedelta
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any, TypedDict
 from uuid import UUID, uuid4
 
 from cachetools import LRUCache
 from flask import current_app
+from pydantic import TypeAdapter
 from sqlalchemy import select
 from sqlalchemy.orm import Session, sessionmaker
 
 from core.helper.encrypter import batch_decrypt_token, encrypt_token, obfuscated_token
 from core.ops.entities.config_entity import (
     OPS_FILE_PATH,
+    BaseTracingConfig,
     TracingProviderEnum,
 )
 from core.ops.entities.trace_entity import (
@@ -33,7 +35,7 @@ from core.ops.entities.trace_entity import (
     WorkflowNodeTraceInfo,
     WorkflowTraceInfo,
 )
-from core.ops.utils import get_message_data
+from core.ops.utils import JSON_DICT_ADAPTER, get_message_data
 from extensions.ext_database import db
 from extensions.ext_storage import storage
 from models.account import Tenant
@@ -50,6 +52,14 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+class _AppTracingConfig(TypedDict, total=False):
+    enabled: bool
+    tracing_provider: str | None
+
+
+_app_tracing_config_adapter: TypeAdapter[_AppTracingConfig] = TypeAdapter(_AppTracingConfig)
+
+
 def _lookup_app_and_workspace_names(app_id: str | None, tenant_id: str | None) -> tuple[str, str]:
     """Return (app_name, workspace_name) for the given IDs. Falls back to empty strings."""
     app_name = ""
@@ -185,8 +195,15 @@ def _lookup_llm_credential_info(
         return None, ""
 
 
-class OpsTraceProviderConfigMap(collections.UserDict[str, dict[str, Any]]):
-    def __getitem__(self, provider: str) -> dict[str, Any]:
+class TracingProviderConfigEntry(TypedDict):
+    config_class: type[BaseTracingConfig]
+    secret_keys: list[str]
+    other_keys: list[str]
+    trace_instance: type[Any]
+
+
+class OpsTraceProviderConfigMap(collections.UserDict[str, TracingProviderConfigEntry]):
+    def __getitem__(self, provider: str) -> TracingProviderConfigEntry:
         match provider:
             case TracingProviderEnum.LANGFUSE:
                 from core.ops.entities.config_entity import LangfuseConfig
@@ -446,7 +463,7 @@ class OpsTraceManager:
     @classmethod
     def get_ops_trace_instance(
         cls,
-        app_id: Union[UUID, str] | None = None,
+        app_id: UUID | str | None = None,
     ):
         """
         Get ops trace through model config
@@ -468,7 +485,7 @@ class OpsTraceManager:
         if app is None:
             return None
 
-        app_ops_trace_config = json.loads(app.tracing) if app.tracing else None
+        app_ops_trace_config = _app_tracing_config_adapter.validate_json(app.tracing) if app.tracing else None
         if app_ops_trace_config is None:
             return None
         if not app_ops_trace_config.get("enabled"):
@@ -560,7 +577,7 @@ class OpsTraceManager:
             raise ValueError("App not found")
         if not app.tracing:
             return {"enabled": False, "tracing_provider": None}
-        app_trace_config = json.loads(app.tracing)
+        app_trace_config = _app_tracing_config_adapter.validate_json(app.tracing)
         return app_trace_config
 
     @staticmethod
@@ -575,8 +592,8 @@ class OpsTraceManager:
             provider_config_map[tracing_provider]["config_class"],
             provider_config_map[tracing_provider]["trace_instance"],
         )
-        tracing_config = config_type(**tracing_config)
-        return trace_instance(tracing_config).api_check()
+        config = config_type(**tracing_config)
+        return trace_instance(config).api_check()
 
     @staticmethod
     def get_trace_config_project_key(tracing_config: dict, tracing_provider: str):
@@ -590,8 +607,8 @@ class OpsTraceManager:
             provider_config_map[tracing_provider]["config_class"],
             provider_config_map[tracing_provider]["trace_instance"],
         )
-        tracing_config = config_type(**tracing_config)
-        return trace_instance(tracing_config).get_project_key()
+        config = config_type(**tracing_config)
+        return trace_instance(config).get_project_key()
 
     @staticmethod
     def get_trace_config_project_url(tracing_config: dict, tracing_provider: str):
@@ -605,8 +622,8 @@ class OpsTraceManager:
             provider_config_map[tracing_provider]["config_class"],
             provider_config_map[tracing_provider]["trace_instance"],
         )
-        tracing_config = config_type(**tracing_config)
-        return trace_instance(tracing_config).get_project_url()
+        config = config_type(**tracing_config)
+        return trace_instance(config).get_project_url()
 
 
 class TraceTask:
@@ -636,7 +653,6 @@ class TraceTask:
         carries ``total_tokens``.  Projects only the ``outputs`` column to avoid loading
         large JSON blobs unnecessarily.
         """
-        import json
 
         from models.workflow import WorkflowNodeExecutionModel
 
@@ -658,7 +674,7 @@ class TraceTask:
             if not raw:
                 continue
             try:
-                outputs = json.loads(raw) if isinstance(raw, str) else raw
+                outputs = JSON_DICT_ADAPTER.validate_json(raw) if isinstance(raw, str) else raw
             except (ValueError, TypeError):
                 continue
             if not isinstance(outputs, dict):
@@ -700,7 +716,7 @@ class TraceTask:
         self,
         trace_type: Any,
         message_id: str | None = None,
-        workflow_execution: Optional["WorkflowExecution"] = None,
+        workflow_execution: "WorkflowExecution | None" = None,
         conversation_id: str | None = None,
         user_id: str | None = None,
         timer: Any | None = None,
@@ -1420,7 +1436,7 @@ class TraceTask:
             return {}
 
         try:
-            metadata = json.loads(message_data.message_metadata)
+            metadata = JSON_DICT_ADAPTER.validate_json(message_data.message_metadata)
             usage = metadata.get("usage", {})
             time_to_first_token = usage.get("time_to_first_token")
             time_to_generate = usage.get("time_to_generate")
@@ -1430,7 +1446,7 @@ class TraceTask:
                 "llm_streaming_time_to_generate": time_to_generate,
                 "is_streaming_request": time_to_first_token is not None,
             }
-        except (json.JSONDecodeError, AttributeError):
+        except (ValueError, AttributeError):
             return {}
 
 

api/core/ops/utils.py

@@ -3,11 +3,14 @@ from datetime import datetime
 from typing import Any, Union
 from urllib.parse import urlparse
 
+from pydantic import TypeAdapter
 from sqlalchemy import select
 
 from models.engine import db
 from models.model import Message
 
+JSON_DICT_ADAPTER: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any])
+
 
 def filter_none_values(data: dict[str, Any]) -> dict[str, Any]:
     new_data = {}

api/core/plugin/backwards_invocation/base.py

@@ -1,5 +1,4 @@
 from collections.abc import Generator, Mapping
-from typing import Generic, TypeVar
 
 from pydantic import BaseModel
 
@@ -19,9 +18,6 @@ class BaseBackwardsInvocation:
             yield BaseBackwardsInvocationResponse(data=response).model_dump_json().encode()
 
 
-T = TypeVar("T", bound=dict | Mapping | str | bool | int | BaseModel)
-
-
-class BaseBackwardsInvocationResponse(BaseModel, Generic[T]):
+class BaseBackwardsInvocationResponse[T: dict | Mapping | str | bool | int | BaseModel](BaseModel):
     data: T | None = None
     error: str = ""

api/core/plugin/entities/plugin_daemon.py

@@ -4,7 +4,7 @@ import enum
 from collections.abc import Mapping, Sequence
 from datetime import datetime
 from enum import StrEnum
-from typing import Any, Generic, TypeVar
+from typing import Any
 
 from graphon.model_runtime.entities.model_entities import AIModelEntity
 from graphon.model_runtime.entities.provider_entities import ProviderEntity
@@ -19,10 +19,8 @@ from core.tools.entities.common_entities import I18nObject
 from core.tools.entities.tool_entities import ToolProviderEntityWithPlugin
 from core.trigger.entities.entities import TriggerProviderEntity
 
-T = TypeVar("T", bound=(BaseModel | dict | list | bool | str))
 
-
-class PluginDaemonBasicResponse(BaseModel, Generic[T]):
+class PluginDaemonBasicResponse[T: BaseModel | dict | list | bool | str](BaseModel):
     """
     Basic response from plugin daemon.
     """

api/core/plugin/impl/base.py

@@ -2,7 +2,7 @@ import inspect
 import json
 import logging
 from collections.abc import Callable, Generator
-from typing import Any, TypeVar, cast
+from typing import Any, cast
 
 import httpx
 from graphon.model_runtime.errors.invoke import (
@@ -17,6 +17,7 @@ from pydantic import BaseModel
 from yarl import URL
 
 from configs import dify_config
+from core.helper.http_client_pooling import get_pooled_http_client
 from core.plugin.endpoint.exc import EndpointSetupFailedError
 from core.plugin.entities.plugin_daemon import PluginDaemonBasicResponse, PluginDaemonError, PluginDaemonInnerError
 from core.plugin.impl.exc import (
@@ -50,10 +51,13 @@ elif isinstance(_plugin_daemon_timeout_config, httpx.Timeout):
 else:
     plugin_daemon_request_timeout = httpx.Timeout(_plugin_daemon_timeout_config)
 
-T = TypeVar("T", bound=(BaseModel | dict[str, Any] | list[Any] | bool | str))
-
 logger = logging.getLogger(__name__)
 
+_httpx_client: httpx.Client = get_pooled_http_client(
+    "plugin_daemon",
+    lambda: httpx.Client(limits=httpx.Limits(max_keepalive_connections=50, max_connections=100), trust_env=False),
+)
+
 
 class BasePluginClient:
     def _request(
@@ -84,7 +88,7 @@ class BasePluginClient:
             request_kwargs["content"] = prepared_data
 
         try:
-            response = httpx.request(**request_kwargs)
+            response = _httpx_client.request(**request_kwargs)
         except httpx.RequestError:
             logger.exception("Request to Plugin Daemon Service failed")
             raise PluginDaemonInnerError(code=-500, message="Request to Plugin Daemon Service failed")
@@ -171,7 +175,7 @@ class BasePluginClient:
             stream_kwargs["content"] = prepared_data
 
         try:
-            with httpx.stream(**stream_kwargs) as response:
+            with _httpx_client.stream(**stream_kwargs) as response:
                 for raw_line in response.iter_lines():
                     if not raw_line:
                         continue
@@ -185,7 +189,7 @@ class BasePluginClient:
             logger.exception("Stream request to Plugin Daemon Service failed")
             raise PluginDaemonInnerError(code=-500, message="Request to Plugin Daemon Service failed")
 
-    def _stream_request_with_model(
+    def _stream_request_with_model[T: BaseModel | dict[str, Any] | list[Any] | bool | str](
         self,
         method: str,
         path: str,
@@ -201,7 +205,7 @@ class BasePluginClient:
         for line in self._stream_request(method, path, params, headers, data, files):
             yield type_(**json.loads(line))  # type: ignore
 
-    def _request_with_model(
+    def _request_with_model[T: BaseModel | dict[str, Any] | list[Any] | bool | str](
         self,
         method: str,
         path: str,
@@ -217,7 +221,7 @@ class BasePluginClient:
         response = self._request(method, path, headers, data, params, files)
         return type_(**response.json())  # type: ignore[return-value]
 
-    def _request_with_plugin_daemon_response(
+    def _request_with_plugin_daemon_response[T: BaseModel | dict[str, Any] | list[Any] | bool | str](
         self,
         method: str,
         path: str,
@@ -272,7 +276,7 @@ class BasePluginClient:
 
         return rep.data
 
-    def _request_with_plugin_daemon_response_stream(
+    def _request_with_plugin_daemon_response_stream[T: BaseModel | dict[str, Any] | list[Any] | bool | str](
         self,
         method: str,
         path: str,

api/core/plugin/utils/chunk_merger.py

@@ -1,12 +1,9 @@
 from collections.abc import Generator
 from dataclasses import dataclass, field
-from typing import TypeVar, Union
 
 from core.agent.entities import AgentInvokeMessage
 from core.tools.entities.tool_entities import ToolInvokeMessage
 
-MessageType = TypeVar("MessageType", bound=Union[ToolInvokeMessage, AgentInvokeMessage])
-
 
 @dataclass
 class FileChunk:
@@ -22,11 +19,11 @@ class FileChunk:
         self.data = bytearray(self.total_length)
 
 
-def merge_blob_chunks(
-    response: Generator[MessageType, None, None],
+def merge_blob_chunks[T: ToolInvokeMessage | AgentInvokeMessage](
+    response: Generator[T, None, None],
     max_file_size: int = 30 * 1024 * 1024,
     max_chunk_size: int = 8192,
-) -> Generator[MessageType, None, None]:
+) -> Generator[T, None, None]:
     """
     Merge streaming blob chunks into complete blob messages.
 

api/core/rag/data_post_processor/data_post_processor.py

@@ -1,6 +1,7 @@
+from typing import TypedDict
+
 from graphon.model_runtime.entities.model_entities import ModelType
 from graphon.model_runtime.errors.invoke import InvokeAuthorizationError
-from typing_extensions import TypedDict
 
 from core.model_manager import ModelInstance, ModelManager
 from core.rag.data_post_processor.reorder import ReorderRunner

api/core/rag/datasource/keyword/jieba/jieba.py

@@ -1,10 +1,9 @@
 from collections import defaultdict
-from typing import Any
+from typing import Any, TypedDict
 
 import orjson
 from pydantic import BaseModel
 from sqlalchemy import select
-from typing_extensions import TypedDict
 
 from configs import dify_config
 from core.rag.datasource.keyword.jieba.jieba_keyword_table_handler import JiebaKeywordTableHandler

api/core/rag/datasource/retrieval_service.py

@@ -1,13 +1,12 @@
 import concurrent.futures
 import logging
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, NotRequired
+from typing import Any, NotRequired, TypedDict
 
 from flask import Flask, current_app
 from graphon.model_runtime.entities.model_entities import ModelType
 from sqlalchemy import select
 from sqlalchemy.orm import Session, load_only
-from typing_extensions import TypedDict
 
 from configs import dify_config
 from core.db.session_factory import session_factory

api/core/rag/datasource/vdb/alibabacloud_mysql/alibabacloud_mysql_vector.py

@@ -10,6 +10,7 @@ from mysql.connector import Error as MySQLError
 from pydantic import BaseModel, model_validator
 
 from configs import dify_config
+from core.rag.datasource.vdb.field import parse_metadata_json
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
@@ -178,9 +179,7 @@ class AlibabaCloudMySQLVector(BaseVector):
             cur.execute(f"SELECT meta, text FROM {self.table_name} WHERE id IN ({placeholders})", ids)
             docs = []
             for record in cur:
-                metadata = record["meta"]
-                if isinstance(metadata, str):
-                    metadata = json.loads(metadata)
+                metadata = parse_metadata_json(record["meta"])
                 docs.append(Document(page_content=record["text"], metadata=metadata))
         return docs
 
@@ -263,15 +262,13 @@ class AlibabaCloudMySQLVector(BaseVector):
                         # similarity = 1 / (1 + distance)
                         similarity = 1.0 / (1.0 + distance)
 
-                    metadata = record["meta"]
-                    if isinstance(metadata, str):
-                        metadata = json.loads(metadata)
+                    metadata = parse_metadata_json(record["meta"])
                     metadata["score"] = similarity
                     metadata["distance"] = distance
 
                     if similarity >= score_threshold:
                         docs.append(Document(page_content=record["text"], metadata=metadata))
-                except (ValueError, json.JSONDecodeError) as e:
+                except (ValueError, TypeError) as e:
                     logger.warning("Error processing search result: %s", e)
                     continue
 
@@ -306,9 +303,7 @@ class AlibabaCloudMySQLVector(BaseVector):
             )
             docs = []
             for record in cur:
-                metadata = record["meta"]
-                if isinstance(metadata, str):
-                    metadata = json.loads(metadata)
+                metadata = parse_metadata_json(record["meta"])
                 metadata["score"] = float(record["score"])
                 docs.append(Document(page_content=record["text"], metadata=metadata))
         return docs

api/core/rag/datasource/vdb/analyticdb/analyticdb_vector_openapi.py

@@ -8,6 +8,7 @@ _import_err_msg = (
     "please run `pip install alibabacloud_gpdb20160503 alibabacloud_tea_openapi`"
 )
 
+from core.rag.datasource.vdb.field import parse_metadata_json
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 
@@ -257,7 +258,7 @@ class AnalyticdbVectorOpenAPI:
         documents = []
         for match in response.body.matches.match:
             if match.score >= score_threshold:
-                metadata = json.loads(match.metadata.get("metadata_"))
+                metadata = parse_metadata_json(match.metadata.get("metadata_"))
                 metadata["score"] = match.score
                 doc = Document(
                     page_content=match.metadata.get("page_content"),
@@ -294,7 +295,7 @@ class AnalyticdbVectorOpenAPI:
         documents = []
         for match in response.body.matches.match:
             if match.score >= score_threshold:
-                metadata = json.loads(match.metadata.get("metadata_"))
+                metadata = parse_metadata_json(match.metadata.get("metadata_"))
                 metadata["score"] = match.score
                 doc = Document(
                     page_content=match.metadata.get("page_content"),

api/core/rag/datasource/vdb/baidu/baidu_vector.py

@@ -29,6 +29,7 @@ from pymochow.model.table import AnnSearch, BM25SearchRequest, HNSWSearchParams,
 
 from configs import dify_config
 from core.rag.datasource.vdb.field import Field as VDBField
+from core.rag.datasource.vdb.field import parse_metadata_json
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
@@ -173,15 +174,9 @@ class BaiduVector(BaseVector):
             score = row.get("score", 0.0)
             meta = row_data.get(VDBField.METADATA_KEY, {})
 
-            # Handle both JSON string and dict formats for backward compatibility
-            if isinstance(meta, str):
-                try:
-                    import json
-
-                    meta = json.loads(meta)
-                except (json.JSONDecodeError, TypeError):
-                    meta = {}
-            elif not isinstance(meta, dict):
+            try:
+                meta = parse_metadata_json(meta)
+            except (ValueError, TypeError):
                 meta = {}
 
             if score >= score_threshold:
@@ -200,7 +195,11 @@ class BaiduVector(BaseVector):
                 raise
 
     def _init_client(self, config) -> MochowClient:
-        config = Configuration(credentials=BceCredentials(config.account, config.api_key), endpoint=config.endpoint)
+        config = Configuration(
+            credentials=BceCredentials(config.account, config.api_key),
+            endpoint=config.endpoint,
+            connection_timeout_in_mills=config.connection_timeout_in_mills,
+        )
         client = MochowClient(config)
         return client
 

api/core/rag/datasource/vdb/clickzetta/clickzetta_vector.py

@@ -17,7 +17,7 @@ if TYPE_CHECKING:
     from clickzetta.connector.v0.connection import Connection  # type: ignore
 
 from configs import dify_config
-from core.rag.datasource.vdb.field import Field
+from core.rag.datasource.vdb.field import Field, parse_metadata_json
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.embedding.embedding_base import Embeddings
@@ -357,18 +357,19 @@ class ClickzettaVector(BaseVector):
         """
         try:
             if raw_metadata:
-                metadata = json.loads(raw_metadata)
+                # First parse may yield a string (double-encoded JSON) so use json.loads
+                first_pass = json.loads(raw_metadata)
 
                 # Handle double-encoded JSON
-                if isinstance(metadata, str):
-                    metadata = json.loads(metadata)
-
-                # Ensure we have a dict
-                if not isinstance(metadata, dict):
+                if isinstance(first_pass, str):
+                    metadata = parse_metadata_json(first_pass)
+                elif isinstance(first_pass, dict):
+                    metadata = first_pass
+                else:
                     metadata = {}
             else:
                 metadata = {}
-        except (json.JSONDecodeError, TypeError):
+        except (json.JSONDecodeError, ValueError, TypeError):
             logger.exception("JSON parsing failed for metadata")
             # Fallback: extract document_id with regex
             doc_id_match = re.search(r'"document_id":\s*"([^"]+)"', raw_metadata or "")
@@ -930,17 +931,18 @@ class ClickzettaVector(BaseVector):
                         # Parse metadata from JSON string (may be double-encoded)
                         try:
                             if row[2]:
-                                metadata = json.loads(row[2])
+                                # First parse may yield a string (double-encoded JSON)
+                                first_pass = json.loads(row[2])
 
-                                # If result is a string, it's double-encoded JSON - parse again
-                                if isinstance(metadata, str):
-                                    metadata = json.loads(metadata)
-
-                                if not isinstance(metadata, dict):
+                                if isinstance(first_pass, str):
+                                    metadata = parse_metadata_json(first_pass)
+                                elif isinstance(first_pass, dict):
+                                    metadata = first_pass
+                                else:
                                     metadata = {}
                             else:
                                 metadata = {}
-                        except (json.JSONDecodeError, TypeError):
+                        except (json.JSONDecodeError, ValueError, TypeError):
                             logger.exception("JSON parsing failed")
                             # Fallback: extract document_id with regex
 

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

@@ -1,4 +1,24 @@
 from enum import StrEnum, auto
+from typing import Any
+
+from pydantic import TypeAdapter
+
+_metadata_adapter: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any])
+
+
+def parse_metadata_json(raw: Any) -> dict[str, Any]:
+    """Parse metadata from a JSON string or pass through an existing dict.
+
+    Many VDB drivers return metadata as either a JSON string or an already-
+    decoded dict depending on the column type and driver version.
+    """
+    if raw is None or raw in ("", b""):
+        return {}
+    if isinstance(raw, dict):
+        return raw
+    if not isinstance(raw, (str, bytes, bytearray)):
+        return {}
+    return _metadata_adapter.validate_json(raw)
 
 
 class Field(StrEnum):

api/core/rag/datasource/vdb/hologres/hologres_vector.py

@@ -9,6 +9,7 @@ from psycopg import sql as psql
 from pydantic import BaseModel, model_validator
 
 from configs import dify_config
+from core.rag.datasource.vdb.field import parse_metadata_json
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
@@ -217,8 +218,7 @@ class HologresVector(BaseVector):
             text = row[2]
             meta = row[3]
 
-            if isinstance(meta, str):
-                meta = json.loads(meta)
+            meta = parse_metadata_json(meta)
 
             # Convert distance to similarity score (consistent with pgvector)
             score = 1 - distance
@@ -265,8 +265,7 @@ class HologresVector(BaseVector):
             meta = row[2]
             score = row[-1]  # score is the last column from return_score
 
-            if isinstance(meta, str):
-                meta = json.loads(meta)
+            meta = parse_metadata_json(meta)
 
             meta["score"] = score
             docs.append(Document(page_content=text, metadata=meta))

api/core/rag/datasource/vdb/iris/iris_vector.py

@@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any
 
 from configs import dify_config
 from configs.middleware.vdb.iris_config import IrisVectorConfig
+from core.rag.datasource.vdb.field import parse_metadata_json
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
@@ -269,7 +270,7 @@ class IrisVector(BaseVector):
                 if len(row) >= 4:
                     text, meta_str, score = row[1], row[2], float(row[3])
                     if score >= score_threshold:
-                        metadata = json.loads(meta_str) if meta_str else {}
+                        metadata = parse_metadata_json(meta_str)
                         metadata["score"] = score
                         docs.append(Document(page_content=text, metadata=metadata))
             return docs
@@ -384,7 +385,7 @@ class IrisVector(BaseVector):
                     meta_str = row[2]
                     score_value = row[3]
 
-                    metadata = json.loads(meta_str) if meta_str else {}
+                    metadata = parse_metadata_json(meta_str)
                     # Add score to metadata for hybrid search compatibility
                     score = float(score_value) if score_value is not None else 0.0
                     metadata["score"] = score

api/core/rag/datasource/vdb/matrixone/matrixone_vector.py

@@ -3,12 +3,13 @@ import logging
 import uuid
 from collections.abc import Callable
 from functools import wraps
-from typing import Any, Concatenate, ParamSpec, TypeVar
+from typing import Any, Concatenate
 
 from mo_vector.client import MoVectorClient  # type: ignore
 from pydantic import BaseModel, model_validator
 
 from configs import dify_config
+from core.rag.datasource.vdb.field import parse_metadata_json
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
@@ -19,15 +20,12 @@ from models.dataset import Dataset
 
 logger = logging.getLogger(__name__)
 
-P = ParamSpec("P")
-R = TypeVar("R")
 
-T = TypeVar("T", bound="MatrixoneVector")
-
-
-def ensure_client(func: Callable[Concatenate[T, P], R]):
+def ensure_client[T: MatrixoneVector, **P, R](
+    func: Callable[Concatenate[T, P], R],
+) -> Callable[Concatenate[T, P], R]:
     @wraps(func)
-    def wrapper(self: T, *args: P.args, **kwargs: P.kwargs):
+    def wrapper(self: T, *args: P.args, **kwargs: P.kwargs) -> R:
         if self.client is None:
             self.client = self._get_client(None, False)
         return func(self, *args, **kwargs)
@@ -196,11 +194,7 @@ class MatrixoneVector(BaseVector):
 
         docs = []
         for result in results:
-            metadata = result.metadata
-            if isinstance(metadata, str):
-                import json
-
-                metadata = json.loads(metadata)
+            metadata = parse_metadata_json(result.metadata)
             score = 1 - result.distance
             if score >= score_threshold:
                 metadata["score"] = score

api/core/rag/datasource/vdb/oceanbase/oceanbase_vector.py

@@ -10,6 +10,7 @@ from sqlalchemy.dialects.mysql import LONGTEXT
 from sqlalchemy.exc import SQLAlchemyError
 
 from configs import dify_config
+from core.rag.datasource.vdb.field import parse_metadata_json
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
@@ -366,8 +367,8 @@ class OceanBaseVector(BaseVector):
 
             # Parse metadata JSON
             try:
-                metadata = json.loads(metadata_str) if isinstance(metadata_str, str) else metadata_str
-            except json.JSONDecodeError:
+                metadata = parse_metadata_json(metadata_str)
+            except (ValueError, TypeError):
                 logger.warning("Invalid JSON metadata: %s", metadata_str)
                 metadata = {}
 

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

@@ -3,7 +3,7 @@ import os
 import uuid
 from collections.abc import Generator, Iterable, Sequence
 from itertools import islice
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any
 
 import qdrant_client
 from flask import current_app
@@ -36,8 +36,8 @@ if TYPE_CHECKING:
     from qdrant_client.conversions import common_types
     from qdrant_client.http import models as rest
 
-    DictFilter = dict[str, Union[str, int, bool, dict, list]]
-    MetadataFilter = Union[DictFilter, common_types.Filter]
+    type DictFilter = dict[str, str | int | bool | dict | list]
+    type MetadataFilter = DictFilter | common_types.Filter
 
 
 class PathQdrantParams(BaseModel):

api/core/rag/datasource/vdb/tablestore/tablestore_vector.py

@@ -9,7 +9,7 @@ from pydantic import BaseModel, model_validator
 from tablestore import BatchGetRowRequest, TableInBatchGetRowItem
 
 from configs import dify_config
-from core.rag.datasource.vdb.field import Field
+from core.rag.datasource.vdb.field import Field, parse_metadata_json
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
@@ -73,7 +73,8 @@ class TableStoreVector(BaseVector):
         for item in table_result:
             if item.is_ok and item.row:
                 kv = {k: v for k, v, _ in item.row.attribute_columns}
-                docs.append(Document(page_content=kv[Field.CONTENT_KEY], metadata=json.loads(kv[Field.METADATA_KEY])))
+                metadata = parse_metadata_json(kv[Field.METADATA_KEY])
+                docs.append(Document(page_content=kv[Field.CONTENT_KEY], metadata=metadata))
         return docs
 
     def get_type(self) -> str:
@@ -311,7 +312,7 @@ class TableStoreVector(BaseVector):
                 metadata_str = ots_column_map.get(Field.METADATA_KEY)
 
                 vector = json.loads(vector_str) if vector_str else None
-                metadata = json.loads(metadata_str) if metadata_str else {}
+                metadata = parse_metadata_json(metadata_str)
 
                 metadata["score"] = search_hit.score
 
@@ -371,7 +372,7 @@ class TableStoreVector(BaseVector):
                 ots_column_map[col[0]] = col[1]
 
             metadata_str = ots_column_map.get(Field.METADATA_KEY)
-            metadata = json.loads(metadata_str) if metadata_str else {}
+            metadata = parse_metadata_json(metadata_str)
 
             vector_str = ots_column_map.get(Field.VECTOR)
             vector = json.loads(vector_str) if vector_str else None

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

@@ -11,6 +11,7 @@ from tcvectordb.model import index as vdb_index  # type: ignore
 from tcvectordb.model.document import AnnSearch, Filter, KeywordSearch, WeightedRerank  # type: ignore
 
 from configs import dify_config
+from core.rag.datasource.vdb.field import parse_metadata_json
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
@@ -286,13 +287,10 @@ class TencentVector(BaseVector):
             return docs
 
         for result in res[0]:
-            meta = result.get(self.field_metadata)
-            if isinstance(meta, str):
-                # Compatible with version 1.1.3 and below.
-                meta = json.loads(meta)
-                score = 1 - result.get("score", 0.0)
-            else:
-                score = result.get("score", 0.0)
+            raw_meta = result.get(self.field_metadata)
+            # Compatible with version 1.1.3 and below: str means old driver.
+            score = (1 - result.get("score", 0.0)) if isinstance(raw_meta, str) else result.get("score", 0.0)
+            meta = parse_metadata_json(raw_meta)
             if score >= score_threshold:
                 meta["score"] = score
                 doc = Document(page_content=result.get(self.field_text), metadata=meta)

api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_on_qdrant_vector.py

@@ -3,7 +3,7 @@ import os
 import uuid
 from collections.abc import Generator, Iterable, Sequence
 from itertools import islice
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any
 
 import httpx
 import qdrant_client
@@ -40,8 +40,8 @@ if TYPE_CHECKING:
     from qdrant_client.conversions import common_types
     from qdrant_client.http import models as rest
 
-    DictFilter = dict[str, Union[str, int, bool, dict, list]]
-    MetadataFilter = Union[DictFilter, common_types.Filter]
+    type DictFilter = dict[str, str | int | bool | dict | list]
+    type MetadataFilter = DictFilter | common_types.Filter
 
 
 class TidbOnQdrantConfig(BaseModel):

api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_service.py

@@ -6,11 +6,18 @@ import httpx
 from httpx import DigestAuth
 
 from configs import dify_config
+from core.helper.http_client_pooling import get_pooled_http_client
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from models.dataset import TidbAuthBinding
 from models.enums import TidbAuthBindingStatus
 
+# Reuse a pooled HTTP client for all TiDB Cloud requests to minimize connection churn
+_tidb_http_client: httpx.Client = get_pooled_http_client(
+    "tidb:cloud",
+    lambda: httpx.Client(limits=httpx.Limits(max_keepalive_connections=50, max_connections=100)),
+)
+
 
 class TidbService:
     @staticmethod
@@ -50,7 +57,9 @@ class TidbService:
             "rootPassword": password,
         }
 
-        response = httpx.post(f"{api_url}/clusters", json=cluster_data, auth=DigestAuth(public_key, private_key))
+        response = _tidb_http_client.post(
+            f"{api_url}/clusters", json=cluster_data, auth=DigestAuth(public_key, private_key)
+        )
 
         if response.status_code == 200:
             response_data = response.json()
@@ -84,7 +93,9 @@ class TidbService:
         :return: The response from the API.
         """
 
-        response = httpx.delete(f"{api_url}/clusters/{cluster_id}", auth=DigestAuth(public_key, private_key))
+        response = _tidb_http_client.delete(
+            f"{api_url}/clusters/{cluster_id}", auth=DigestAuth(public_key, private_key)
+        )
 
         if response.status_code == 200:
             return response.json()
@@ -103,7 +114,7 @@ class TidbService:
         :return: The response from the API.
         """
 
-        response = httpx.get(f"{api_url}/clusters/{cluster_id}", auth=DigestAuth(public_key, private_key))
+        response = _tidb_http_client.get(f"{api_url}/clusters/{cluster_id}", auth=DigestAuth(public_key, private_key))
 
         if response.status_code == 200:
             return response.json()
@@ -128,7 +139,7 @@ class TidbService:
 
         body = {"password": new_password, "builtinRole": "role_admin", "customRoles": []}
 
-        response = httpx.patch(
+        response = _tidb_http_client.patch(
             f"{api_url}/clusters/{cluster_id}/sqlUsers/{account}",
             json=body,
             auth=DigestAuth(public_key, private_key),
@@ -162,7 +173,9 @@ class TidbService:
         tidb_serverless_list_map = {item.cluster_id: item for item in tidb_serverless_list}
         cluster_ids = [item.cluster_id for item in tidb_serverless_list]
         params = {"clusterIds": cluster_ids, "view": "BASIC"}
-        response = httpx.get(f"{api_url}/clusters:batchGet", params=params, auth=DigestAuth(public_key, private_key))
+        response = _tidb_http_client.get(
+            f"{api_url}/clusters:batchGet", params=params, auth=DigestAuth(public_key, private_key)
+        )
 
         if response.status_code == 200:
             response_data = response.json()
@@ -223,7 +236,7 @@ class TidbService:
             clusters.append(cluster_data)
 
         request_body = {"requests": clusters}
-        response = httpx.post(
+        response = _tidb_http_client.post(
             f"{api_url}/clusters:batchCreate", json=request_body, auth=DigestAuth(public_key, private_key)
         )
 

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

@@ -9,7 +9,7 @@ from sqlalchemy import text as sql_text
 from sqlalchemy.orm import Session, declarative_base
 
 from configs import dify_config
-from core.rag.datasource.vdb.field import Field
+from core.rag.datasource.vdb.field import Field, parse_metadata_json
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
@@ -228,7 +228,7 @@ class TiDBVector(BaseVector):
             )
             results = [(row[0], row[1], row[2]) for row in res]
             for meta, text, distance in results:
-                metadata = json.loads(meta)
+                metadata = parse_metadata_json(meta)
                 metadata["score"] = 1 - distance
                 docs.append(Document(page_content=text, metadata=metadata))
         return docs

api/core/rag/datasource/vdb/vikingdb/vikingdb_vector.py

@@ -15,6 +15,7 @@ from volcengine.viking_db import (  # type: ignore
 
 from configs import dify_config
 from core.rag.datasource.vdb.field import Field as vdb_Field
+from core.rag.datasource.vdb.field import parse_metadata_json
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
@@ -163,7 +164,7 @@ class VikingDBVector(BaseVector):
         for result in results:
             metadata = result.fields.get(vdb_Field.METADATA_KEY)
             if metadata is not None:
-                metadata = json.loads(metadata)
+                metadata = parse_metadata_json(metadata)
                 if metadata.get(key) == value:
                     ids.append(result.id)
         return ids
@@ -189,9 +190,7 @@ class VikingDBVector(BaseVector):
 
         docs = []
         for result in results:
-            metadata = result.fields.get(vdb_Field.METADATA_KEY)
-            if metadata is not None:
-                metadata = json.loads(metadata)
+            metadata = parse_metadata_json(result.fields.get(vdb_Field.METADATA_KEY))
             if result.score >= score_threshold:
                 metadata["score"] = result.score
                 doc = Document(page_content=result.fields.get(vdb_Field.CONTENT_KEY), metadata=metadata)

api/core/rag/embedding/retrieval.py

@@ -1,5 +1,6 @@
+from typing import TypedDict
+
 from pydantic import BaseModel
-from typing_extensions import TypedDict
 
 from models.dataset import DocumentSegment
 

api/core/rag/extractor/blob/blob.py

@@ -12,11 +12,11 @@ import mimetypes
 from collections.abc import Generator, Mapping
 from io import BufferedReader, BytesIO
 from pathlib import Path, PurePath
-from typing import Any, Union
+from typing import Any
 
 from pydantic import BaseModel, ConfigDict, model_validator
 
-PathLike = Union[str, PurePath]
+type PathLike = str | PurePath
 
 
 class Blob(BaseModel):
@@ -29,7 +29,7 @@ class Blob(BaseModel):
     Inspired by: https://developer.mozilla.org/en-US/docs/Web/API/Blob
     """
 
-    data: Union[bytes, str, None] = None  # Raw data
+    data: bytes | str | None = None  # Raw data
     mimetype: str | None = None  # Not to be confused with a file extension
     encoding: str = "utf-8"  # Use utf-8 as default encoding, if decoding to string
     # Location where the original content was found
@@ -75,7 +75,7 @@ class Blob(BaseModel):
             raise ValueError(f"Unable to get bytes for blob {self}")
 
     @contextlib.contextmanager
-    def as_bytes_io(self) -> Generator[Union[BytesIO, BufferedReader], None, None]:
+    def as_bytes_io(self) -> Generator[BytesIO | BufferedReader, None, None]:
         """Read data as a byte stream."""
         if isinstance(self.data, bytes):
             yield BytesIO(self.data)
@@ -117,7 +117,7 @@ class Blob(BaseModel):
     @classmethod
     def from_data(
         cls,
-        data: Union[str, bytes],
+        data: str | bytes,
         *,
         encoding: str = "utf-8",
         mime_type: str | None = None,

api/core/rag/extractor/firecrawl/firecrawl_app.py

@@ -1,9 +1,8 @@
 import json
 import time
-from typing import Any, NotRequired, cast
+from typing import Any, NotRequired, TypedDict, cast
 
 import httpx
-from typing_extensions import TypedDict
 
 from extensions.ext_storage import storage
 

api/core/rag/extractor/watercrawl/client.py

@@ -1,11 +1,10 @@
 import json
 from collections.abc import Generator
-from typing import Any, Union
+from typing import Any, TypedDict
 from urllib.parse import urljoin
 
 import httpx
 from httpx import Response
-from typing_extensions import TypedDict
 
 from core.rag.extractor.watercrawl.exceptions import (
     WaterCrawlAuthenticationError,
@@ -142,7 +141,7 @@ class WaterCrawlAPIClient(BaseAPIClient):
 
     def create_crawl_request(
         self,
-        url: Union[list, str] | None = None,
+        url: list | str | None = None,
         spider_options: SpiderOptions | None = None,
         page_options: PageOptions | None = None,
         plugin_options: dict[str, Any] | None = None,

api/core/rag/extractor/watercrawl/provider.py

@@ -1,8 +1,6 @@
 from collections.abc import Generator
 from datetime import datetime
-from typing import Any
-
-from typing_extensions import TypedDict
+from typing import Any, TypedDict
 
 from core.rag.extractor.watercrawl.client import PageOptions, SpiderOptions, WaterCrawlAPIClient
 

api/core/rag/index_processor/index_processor.py

@@ -35,7 +35,10 @@ class IndexProcessor:
         if "parent_mode" in preview:
             data.parent_mode = preview["parent_mode"]
 
-        for item in preview["preview"]:
+        # Different index processors return different preview shapes:
+        # - paragraph/parent-child processors: {"preview": [...]}
+        # - QA processor: {"qa_preview": [...]} (no "preview" key)
+        for item in preview.get("preview", []):
             if "content" in item and "child_chunks" in item:
                 data.preview.append(
                     PreviewItem(content=item["content"], child_chunks=item["child_chunks"], summary=None)
@@ -44,6 +47,10 @@ class IndexProcessor:
                 data.qa_preview.append(QaPreview(question=item["question"], answer=item["answer"]))
             elif "content" in item:
                 data.preview.append(PreviewItem(content=item["content"], child_chunks=None, summary=None))
+
+        for item in preview.get("qa_preview", []):
+            if "question" in item and "answer" in item:
+                data.qa_preview.append(QaPreview(question=item["question"], answer=item["answer"]))
         return data
 
     def index_and_clean(

api/core/rag/index_processor/index_processor_base.py

@@ -7,12 +7,11 @@ import os
 import re
 from abc import ABC, abstractmethod
 from collections.abc import Mapping
-from typing import TYPE_CHECKING, Any, NotRequired, Optional
+from typing import TYPE_CHECKING, Any, NotRequired, TypedDict
 from urllib.parse import unquote, urlparse
 
 import httpx
 from sqlalchemy import select
-from typing_extensions import TypedDict
 
 from configs import dify_config
 from core.entities.knowledge_entities import PreviewDetail
@@ -118,11 +117,12 @@ class BaseIndexProcessor(ABC):
         max_tokens: int,
         chunk_overlap: int,
         separator: str,
-        embedding_model_instance: Optional["ModelInstance"],
+        embedding_model_instance: "ModelInstance | None",
     ) -> TextSplitter:
         """
         Get the NodeParser object according to the processing rule.
         """
+        character_splitter: TextSplitter
         if processing_rule_mode in ["custom", "hierarchical"]:
             # The user-defined segmentation rule
             max_segmentation_tokens_length = dify_config.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH
@@ -148,7 +148,7 @@ class BaseIndexProcessor(ABC):
                 embedding_model_instance=embedding_model_instance,
             )
 
-        return character_splitter  # type: ignore
+        return character_splitter
 
     def _get_content_files(self, document: Document, current_user: Account | None = None) -> list[AttachmentDocument]:
         """

api/core/rag/splitter/fixed_text_splitter.py

@@ -4,19 +4,13 @@ from __future__ import annotations
 
 import codecs
 import re
-from typing import Any
+from collections.abc import Collection
+from typing import Any, Literal
 
 from graphon.model_runtime.model_providers.__base.tokenizers.gpt2_tokenizer import GPT2Tokenizer
 
 from core.model_manager import ModelInstance
-from core.rag.splitter.text_splitter import (
-    TS,
-    Collection,
-    Literal,
-    RecursiveCharacterTextSplitter,
-    Set,
-    Union,
-)
+from core.rag.splitter.text_splitter import RecursiveCharacterTextSplitter
 
 
 class EnhanceRecursiveCharacterTextSplitter(RecursiveCharacterTextSplitter):
@@ -25,13 +19,13 @@ class EnhanceRecursiveCharacterTextSplitter(RecursiveCharacterTextSplitter):
     """
 
     @classmethod
-    def from_encoder(
-        cls: type[TS],
+    def from_encoder[T: EnhanceRecursiveCharacterTextSplitter](
+        cls: type[T],
         embedding_model_instance: ModelInstance | None,
-        allowed_special: Union[Literal["all"], Set[str]] = set(),  # noqa: UP037
-        disallowed_special: Union[Literal["all"], Collection[str]] = "all",  # noqa: UP037
+        allowed_special: Literal["all"] | set[str] = set(),
+        disallowed_special: Literal["all"] | Collection[str] = "all",
         **kwargs: Any,
-    ):
+    ) -> T:
         def _token_encoder(texts: list[str]) -> list[int]:
             if not texts:
                 return []

api/core/rag/splitter/text_splitter.py

@@ -6,19 +6,12 @@ import re
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Collection, Iterable, Sequence, Set
 from dataclasses import dataclass
-from typing import (
-    Any,
-    Literal,
-    TypeVar,
-    Union,
-)
+from typing import Any, Literal
 
 from core.rag.models.document import BaseDocumentTransformer, Document
 
 logger = logging.getLogger(__name__)
 
-TS = TypeVar("TS", bound="TextSplitter")
-
 
 def _split_text_with_regex(text: str, separator: str, keep_separator: bool) -> list[str]:
     # Now that we have the separator, split the text
@@ -194,8 +187,8 @@ class TokenTextSplitter(TextSplitter):
         self,
         encoding_name: str = "gpt2",
         model_name: str | None = None,
-        allowed_special: Union[Literal["all"], Set[str]] = set(),
-        disallowed_special: Union[Literal["all"], Collection[str]] = "all",
+        allowed_special: Literal["all"] | Set[str] = set(),
+        disallowed_special: Literal["all"] | Collection[str] = "all",
         **kwargs: Any,
     ):
         """Create a new TextSplitter."""

api/core/repositories/celery_workflow_execution_repository.py

@@ -6,7 +6,6 @@ providing improved performance by offloading database operations to background w
 """
 
 import logging
-from typing import Union
 
 from graphon.entities import WorkflowExecution
 from sqlalchemy.engine import Engine
@@ -47,7 +46,7 @@ class CeleryWorkflowExecutionRepository(WorkflowExecutionRepository):
     def __init__(
         self,
         session_factory: sessionmaker | Engine,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         app_id: str | None,
         triggered_from: WorkflowRunTriggeredFrom | None,
     ):

api/core/repositories/celery_workflow_node_execution_repository.py

@@ -7,7 +7,6 @@ providing improved performance by offloading database operations to background w
 
 import logging
 from collections.abc import Sequence
-from typing import Union
 
 from graphon.entities import WorkflowNodeExecution
 from sqlalchemy.engine import Engine
@@ -54,7 +53,7 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository):
     def __init__(
         self,
         session_factory: sessionmaker | Engine,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         app_id: str | None,
         triggered_from: WorkflowNodeExecutionTriggeredFrom | None,
     ):

api/core/repositories/factory.py

@@ -7,7 +7,7 @@ allowing users to configure different repository backends through string paths.
 
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import Literal, Protocol, Union
+from typing import Literal, Protocol
 
 from graphon.entities import WorkflowExecution, WorkflowNodeExecution
 from sqlalchemy.engine import Engine
@@ -61,8 +61,8 @@ class DifyCoreRepositoryFactory:
     @classmethod
     def create_workflow_execution_repository(
         cls,
-        session_factory: Union[sessionmaker, Engine],
-        user: Union[Account, EndUser],
+        session_factory: sessionmaker | Engine,
+        user: Account | EndUser,
         app_id: str,
         triggered_from: WorkflowRunTriggeredFrom,
     ) -> WorkflowExecutionRepository:
@@ -97,8 +97,8 @@ class DifyCoreRepositoryFactory:
     @classmethod
     def create_workflow_node_execution_repository(
         cls,
-        session_factory: Union[sessionmaker, Engine],
-        user: Union[Account, EndUser],
+        session_factory: sessionmaker | Engine,
+        user: Account | EndUser,
         app_id: str,
         triggered_from: WorkflowNodeExecutionTriggeredFrom,
     ) -> WorkflowNodeExecutionRepository:

api/core/repositories/sqlalchemy_workflow_execution_repository.py

@@ -4,7 +4,6 @@ SQLAlchemy implementation of the WorkflowExecutionRepository.
 
 import json
 import logging
-from typing import Union
 
 from graphon.entities import WorkflowExecution
 from graphon.enums import WorkflowExecutionStatus, WorkflowType
@@ -40,7 +39,7 @@ class SQLAlchemyWorkflowExecutionRepository(WorkflowExecutionRepository):
     def __init__(
         self,
         session_factory: sessionmaker | Engine,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         app_id: str | None,
         triggered_from: WorkflowRunTriggeredFrom | None,
     ):

api/core/repositories/sqlalchemy_workflow_node_execution_repository.py

@@ -7,7 +7,7 @@ import json
 import logging
 from collections.abc import Callable, Mapping, Sequence
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, TypeVar, Union
+from typing import Any
 
 import psycopg2.errors
 from graphon.entities import WorkflowNodeExecution
@@ -63,7 +63,7 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository)
     def __init__(
         self,
         session_factory: sessionmaker | Engine,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         app_id: str | None,
         triggered_from: WorkflowNodeExecutionTriggeredFrom | None,
     ):
@@ -551,10 +551,7 @@ def _deterministic_json_dump(value: Mapping[str, Any]) -> str:
     return json.dumps(value, sort_keys=True)
 
 
-_T = TypeVar("_T")
-
-
-def _find_first(seq: Sequence[_T], pred: Callable[[_T], bool]) -> _T | None:
+def _find_first[T](seq: Sequence[T], pred: Callable[[T], bool]) -> T | None:
     filtered = [i for i in seq if pred(i)]
     if filtered:
         return filtered[0]

api/core/schemas/resolver.py

@@ -3,15 +3,15 @@ import re
 import threading
 from collections import deque
 from dataclasses import dataclass
-from typing import Any, Union
+from typing import Any
 
 from core.schemas.registry import SchemaRegistry
 
 logger = logging.getLogger(__name__)
 
 # Type aliases for better clarity
-SchemaType = Union[dict[str, Any], list[Any], str, int, float, bool, None]
-SchemaDict = dict[str, Any]
+type SchemaType = dict[str, Any] | list[Any] | str | int | float | bool | None
+type SchemaDict = dict[str, Any]
 
 # Pre-compiled pattern for better performance
 _DIFY_SCHEMA_PATTERN = re.compile(r"^https://dify\.ai/schemas/(v\d+)/(.+)\.json$")
@@ -54,7 +54,7 @@ class QueueItem:
 
     current: Any
     parent: Any | None
-    key: Union[str, int] | None
+    key: str | int | None
     depth: int
     ref_path: set[str]
 

api/core/trigger/debug/event_bus.py

@@ -1,6 +1,5 @@
 import hashlib
 import logging
-from typing import TypeVar
 
 from redis import RedisError
 
@@ -11,8 +10,6 @@ logger = logging.getLogger(__name__)
 
 TRIGGER_DEBUG_EVENT_TTL = 300
 
-TTriggerDebugEvent = TypeVar("TTriggerDebugEvent", bound="BaseDebugEvent")
-
 
 class TriggerDebugEventBus:
     """
@@ -81,15 +78,15 @@ class TriggerDebugEventBus:
             return 0
 
     @classmethod
-    def poll(
+    def poll[T: BaseDebugEvent](
         cls,
-        event_type: type[TTriggerDebugEvent],
+        event_type: type[T],
         pool_key: str,
         tenant_id: str,
         user_id: str,
         app_id: str,
         node_id: str,
-    ) -> TTriggerDebugEvent | None:
+    ) -> T | None:
         """
         Poll for an event or register to the waiting pool.
 

api/core/workflow/node_factory.py

@@ -2,7 +2,7 @@ import importlib
 import pkgutil
 from collections.abc import Callable, Iterator, Mapping, MutableMapping
 from functools import lru_cache
-from typing import TYPE_CHECKING, Any, TypeAlias, cast, final
+from typing import TYPE_CHECKING, Any, cast, final, override
 
 from graphon.entities.base_node_data import BaseNodeData
 from graphon.entities.graph_config import NodeConfigDict, NodeConfigDictAdapter
@@ -22,7 +22,6 @@ from graphon.nodes.parameter_extractor.entities import ParameterExtractorNodeDat
 from graphon.nodes.question_classifier.entities import QuestionClassifierNodeData
 from sqlalchemy import select
 from sqlalchemy.orm import Session
-from typing_extensions import override
 
 from configs import dify_config
 from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, DifyRunContext
@@ -192,7 +191,7 @@ class _LazyNodeTypeClassesMapping(MutableMapping[NodeType, Mapping[str, type[Nod
 NODE_TYPE_CLASSES_MAPPING: MutableMapping[NodeType, Mapping[str, type[Node]]] = _LazyNodeTypeClassesMapping()
 
 
-LLMCompatibleNodeData: TypeAlias = LLMNodeData | QuestionClassifierNodeData | ParameterExtractorNodeData
+type LLMCompatibleNodeData = LLMNodeData | QuestionClassifierNodeData | ParameterExtractorNodeData
 
 
 def fetch_memory(

api/dify_app.py

@@ -1,5 +1,14 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from flask import Flask
 
+if TYPE_CHECKING:
+    from extensions.ext_login import DifyLoginManager
+
 
 class DifyApp(Flask):
-    pass
+    """Flask application type with Dify-specific extension attributes."""
+
+    login_manager: DifyLoginManager

api/extensions/ext_login.py

@@ -1,7 +1,8 @@
 import json
+from typing import cast
 
 import flask_login
-from flask import Response, request
+from flask import Request, Response, request
 from flask_login import user_loaded_from_request, user_logged_in
 from sqlalchemy import select
 from werkzeug.exceptions import NotFound, Unauthorized
@@ -16,13 +17,35 @@ from models import Account, Tenant, TenantAccountJoin
 from models.model import AppMCPServer, EndUser
 from services.account_service import AccountService
 
-login_manager = flask_login.LoginManager()
+type LoginUser = Account | EndUser
+
+
+class DifyLoginManager(flask_login.LoginManager):
+    """Project-specific Flask-Login manager with a stable unauthorized contract.
+
+    Dify registers `unauthorized_handler` below to always return a JSON `Response`.
+    Overriding this method lets callers rely on that narrower return type instead of
+    Flask-Login's broader callback contract.
+    """
+
+    def unauthorized(self) -> Response:
+        """Return the registered unauthorized handler result as a Flask `Response`."""
+        return cast(Response, super().unauthorized())
+
+    def load_user_from_request_context(self) -> None:
+        """Populate Flask-Login's request-local user cache for the current request."""
+        self._load_user()
+
+
+login_manager = DifyLoginManager()
 
 
 # Flask-Login configuration
 @login_manager.request_loader
-def load_user_from_request(request_from_flask_login):
+def load_user_from_request(request_from_flask_login: Request) -> LoginUser | None:
     """Load user based on the request."""
+    del request_from_flask_login
+
     # Skip authentication for documentation endpoints
     if dify_config.SWAGGER_UI_ENABLED and request.path.endswith((dify_config.SWAGGER_UI_PATH, "/swagger.json")):
         return None
@@ -100,10 +123,12 @@ def load_user_from_request(request_from_flask_login):
             raise NotFound("End user not found.")
         return end_user
 
+    return None
+
 
 @user_logged_in.connect
 @user_loaded_from_request.connect
-def on_user_logged_in(_sender, user):
+def on_user_logged_in(_sender: object, user: LoginUser) -> None:
     """Called when a user logged in.
 
     Note: AccountService.load_logged_in_account will populate user.current_tenant_id
@@ -114,8 +139,10 @@ def on_user_logged_in(_sender, user):
 
 
 @login_manager.unauthorized_handler
-def unauthorized_handler():
+def unauthorized_handler() -> Response:
     """Handle unauthorized requests."""
+    # Keep this as a concrete `Response`; `DifyLoginManager.unauthorized()` narrows
+    # Flask-Login's callback contract based on this override.
     return Response(
         json.dumps({"code": "unauthorized", "message": "Unauthorized."}),
         status=401,
@@ -123,5 +150,5 @@ def unauthorized_handler():
     )
 
 
-def init_app(app: DifyApp):
+def init_app(app: DifyApp) -> None:
     login_manager.init_app(app)

api/extensions/ext_redis.py

@@ -3,7 +3,7 @@ import logging
 import ssl
 from collections.abc import Callable
 from datetime import timedelta
-from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Union
 
 import redis
 from redis import RedisError
@@ -297,12 +297,7 @@ def get_pubsub_broadcast_channel() -> BroadcastChannelProtocol:
     return RedisBroadcastChannel(_pubsub_redis_client)
 
 
-P = ParamSpec("P")
-R = TypeVar("R")
-T = TypeVar("T")
-
-
-def redis_fallback(default_return: T | None = None):  # type: ignore
+def redis_fallback[T](default_return: T | None = None):  # type: ignore
     """
     decorator to handle Redis operation exceptions and return a default value when Redis is unavailable.
 
@@ -310,9 +305,9 @@ def redis_fallback(default_return: T | None = None):  # type: ignore
         default_return: The value to return when a Redis operation fails. Defaults to None.
     """
 
-    def decorator(func: Callable[P, R]):
+    def decorator[**P, R](func: Callable[P, R]) -> Callable[P, R | T | None]:
         @functools.wraps(func)
-        def wrapper(*args: P.args, **kwargs: P.kwargs):
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R | T | None:
             try:
                 return func(*args, **kwargs)
             except RedisError as e:

api/extensions/logstore/repositories/logstore_workflow_execution_repository.py

@@ -2,7 +2,6 @@ import json
 import logging
 import os
 import time
-from typing import Union
 
 from graphon.entities import WorkflowExecution
 from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter
@@ -27,7 +26,7 @@ class LogstoreWorkflowExecutionRepository(WorkflowExecutionRepository):
     def __init__(
         self,
         session_factory: sessionmaker | Engine,
-        user: Union[Account, EndUser],
+        user: Account | EndUser,
         app_id: str | None,
         triggered_from: WorkflowRunTriggeredFrom | None,
     ):