chore(version): bump to 1.7.0 (#22830)

Signed-off-by: -LAN- <laipz8200@outlook.com>
Co-authored-by: -LAN- <laipz8200@outlook.com>

.devcontainer/post_create_command.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-npm add -g pnpm@10.11.1
+npm add -g pnpm@10.13.1
 cd web && pnpm install
 pipx install uv
 
@@ -12,3 +12,4 @@ echo 'alias start-containers="cd /workspaces/dify/docker && docker-compose -f do
 echo 'alias stop-containers="cd /workspaces/dify/docker && docker-compose -f docker-compose.middleware.yaml -p dify --env-file middleware.env down"' >> ~/.bashrc
 
 source /home/vscode/.bashrc
+

.github/workflows/style.yml

@@ -28,7 +28,7 @@ jobs:
 
       - name: Check changed files
         id: changed-files
-        uses: tj-actions/changed-files@v45
+        uses: tj-actions/changed-files@v46
         with:
           files: |
             api/**
@@ -75,7 +75,7 @@ jobs:
 
       - name: Check changed files
         id: changed-files
-        uses: tj-actions/changed-files@v45
+        uses: tj-actions/changed-files@v46
         with:
           files: web/**
 
@@ -113,7 +113,7 @@ jobs:
 
       - name: Check changed files
         id: changed-files
-        uses: tj-actions/changed-files@v45
+        uses: tj-actions/changed-files@v46
         with:
           files: |
             docker/generate_docker_compose
@@ -144,7 +144,7 @@ jobs:
 
       - name: Check changed files
         id: changed-files
-        uses: tj-actions/changed-files@v45
+        uses: tj-actions/changed-files@v46
         with:
           files: |
             **.sh
@@ -152,13 +152,15 @@ jobs:
             **.yml
             **Dockerfile
             dev/**
+            .editorconfig
 
       - name: Super-linter
-        uses: super-linter/super-linter/slim@v7
+        uses: super-linter/super-linter/slim@v8
         if: steps.changed-files.outputs.any_changed == 'true'
         env:
           BASH_SEVERITY: warning
-          DEFAULT_BRANCH: main
+          DEFAULT_BRANCH: origin/main
+          EDITORCONFIG_FILE_NAME: editorconfig-checker.json
           FILTER_REGEX_INCLUDE: pnpm-lock.yaml
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           IGNORE_GENERATED_FILES: true
@@ -168,16 +170,6 @@ jobs:
           # FIXME: temporarily disabled until api-docker.yaml's run script is fixed for shellcheck
           # VALIDATE_GITHUB_ACTIONS: true
           VALIDATE_DOCKERFILE_HADOLINT: true
+          VALIDATE_EDITORCONFIG: true
           VALIDATE_XML: true
           VALIDATE_YAML: true
-
-      - name: EditorConfig checks
-        uses: super-linter/super-linter/slim@v7
-        env:
-          DEFAULT_BRANCH: main
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          IGNORE_GENERATED_FILES: true
-          IGNORE_GITIGNORED_FILES: true
-          # EditorConfig validation
-          VALIDATE_EDITORCONFIG: true
-          EDITORCONFIG_FILE_NAME: editorconfig-checker.json

.github/workflows/web-tests.yml

@@ -27,7 +27,7 @@ jobs:
 
       - name: Check changed files
         id: changed-files
-        uses: tj-actions/changed-files@v45
+        uses: tj-actions/changed-files@v46
         with:
           files: web/**
 

api/.env.example

@@ -54,7 +54,7 @@ REDIS_CLUSTERS_PASSWORD=
 
 # celery configuration
 CELERY_BROKER_URL=redis://:difyai123456@localhost:${REDIS_PORT}/1
-
+CELERY_BACKEND=redis
 # PostgreSQL database configuration
 DB_USERNAME=postgres
 DB_PASSWORD=difyai123456
@@ -142,8 +142,10 @@ WEB_API_CORS_ALLOW_ORIGINS=http://localhost:3000,*
 CONSOLE_CORS_ALLOW_ORIGINS=http://localhost:3000,*
 
 # Vector database configuration
-# support: weaviate, qdrant, milvus, myscale, relyt, pgvecto_rs, pgvector, pgvector, chroma, opensearch, tidb_vector, couchbase, vikingdb, upstash, lindorm, oceanbase, opengauss, tablestore, matrixone
+# Supported values are `weaviate`, `qdrant`, `milvus`, `myscale`, `relyt`, `pgvector`, `pgvecto-rs`, `chroma`, `opensearch`, `oracle`, `tencent`, `elasticsearch`, `elasticsearch-ja`, `analyticdb`, `couchbase`, `vikingdb`, `oceanbase`, `opengauss`, `tablestore`,`vastbase`,`tidb`,`tidb_on_qdrant`,`baidu`,`lindorm`,`huawei_cloud`,`upstash`, `matrixone`.
 VECTOR_STORE=weaviate
+# Prefix used to create collection name in vector database
+VECTOR_INDEX_NAME_PREFIX=Vector_index
 
 # Weaviate configuration
 WEAVIATE_ENDPOINT=http://localhost:8080
@@ -469,6 +471,16 @@ APP_MAX_ACTIVE_REQUESTS=0
 # Celery beat configuration
 CELERY_BEAT_SCHEDULER_TIME=1
 
+# Celery schedule tasks configuration
+ENABLE_CLEAN_EMBEDDING_CACHE_TASK=false
+ENABLE_CLEAN_UNUSED_DATASETS_TASK=false
+ENABLE_CREATE_TIDB_SERVERLESS_TASK=false
+ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK=false
+ENABLE_CLEAN_MESSAGES=false
+ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK=false
+ENABLE_DATASETS_QUEUE_MONITOR=false
+ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK=true
+
 # Position configuration
 POSITION_TOOL_PINS=
 POSITION_TOOL_INCLUDES=

api/Dockerfile

@@ -47,6 +47,8 @@ RUN \
         curl nodejs libgmp-dev libmpfr-dev libmpc-dev \
         # For Security
         expat libldap-2.5-0 perl libsqlite3-0 zlib1g \
+        # install fonts to support the use of tools like pypdfium2
+        fonts-noto-cjk \
         # install a package to improve the accuracy of guessing mime type and file extension
         media-types \
         # install libmagic to support the use of python-magic guess MIMETYPE

api/README.md

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

api/configs/feature/__init__.py

@@ -832,6 +832,41 @@ class CeleryBeatConfig(BaseSettings):
     )
 
 
+class CeleryScheduleTasksConfig(BaseSettings):
+    ENABLE_CLEAN_EMBEDDING_CACHE_TASK: bool = Field(
+        description="Enable clean embedding cache task",
+        default=False,
+    )
+    ENABLE_CLEAN_UNUSED_DATASETS_TASK: bool = Field(
+        description="Enable clean unused datasets task",
+        default=False,
+    )
+    ENABLE_CREATE_TIDB_SERVERLESS_TASK: bool = Field(
+        description="Enable create tidb service job task",
+        default=False,
+    )
+    ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK: bool = Field(
+        description="Enable update tidb service job status task",
+        default=False,
+    )
+    ENABLE_CLEAN_MESSAGES: bool = Field(
+        description="Enable clean messages task",
+        default=False,
+    )
+    ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK: bool = Field(
+        description="Enable mail clean document notify task",
+        default=False,
+    )
+    ENABLE_DATASETS_QUEUE_MONITOR: bool = Field(
+        description="Enable queue monitor task",
+        default=False,
+    )
+    ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK: bool = Field(
+        description="Enable check upgradable plugin task",
+        default=True,
+    )
+
+
 class PositionConfig(BaseSettings):
     POSITION_PROVIDER_PINS: str = Field(
         description="Comma-separated list of pinned model providers",
@@ -961,5 +996,6 @@ class FeatureConfig(
     # hosted services config
     HostedServiceConfig,
     CeleryBeatConfig,
+    CeleryScheduleTasksConfig,
 ):
     pass

api/configs/middleware/__init__.py

@@ -85,6 +85,11 @@ class VectorStoreConfig(BaseSettings):
         default=False,
     )
 
+    VECTOR_INDEX_NAME_PREFIX: Optional[str] = Field(
+        description="Prefix used to create collection name in vector database",
+        default="Vector_index",
+    )
+
 
 class KeywordStoreConfig(BaseSettings):
     KEYWORD_STORE: str = Field(
@@ -211,7 +216,7 @@ class DatabaseConfig(BaseSettings):
 class CeleryConfig(DatabaseConfig):
     CELERY_BACKEND: str = Field(
         description="Backend for Celery task results. Options: 'database', 'redis'.",
-        default="database",
+        default="redis",
     )
 
     CELERY_BROKER_URL: Optional[str] = Field(

api/controllers/console/app/conversation.py

@@ -1,4 +1,4 @@
-from datetime import UTC, datetime
+from datetime import datetime
 
 import pytz  # pip install pytz
 from flask_login import current_user
@@ -19,6 +19,7 @@ from fields.conversation_fields import (
     conversation_pagination_fields,
     conversation_with_summary_pagination_fields,
 )
+from libs.datetime_utils import naive_utc_now
 from libs.helper import DatetimeString
 from libs.login import login_required
 from models import Conversation, EndUser, Message, MessageAnnotation
@@ -315,7 +316,7 @@ def _get_conversation(app_model, conversation_id):
         raise NotFound("Conversation Not Exists.")
 
     if not conversation.read_at:
-        conversation.read_at = datetime.now(UTC).replace(tzinfo=None)
+        conversation.read_at = naive_utc_now()
         conversation.read_account_id = current_user.id
         db.session.commit()
 

api/controllers/console/app/message.py

@@ -5,6 +5,7 @@ from flask_restful import Resource, fields, marshal_with, reqparse
 from flask_restful.inputs import int_range
 from werkzeug.exceptions import Forbidden, InternalServerError, NotFound
 
+import services
 from controllers.console import api
 from controllers.console.app.error import (
     CompletionRequestError,
@@ -27,7 +28,7 @@ from fields.conversation_fields import annotation_fields, message_detail_fields
 from libs.helper import uuid_value
 from libs.infinite_scroll_pagination import InfiniteScrollPagination
 from libs.login import login_required
-from models.model import AppMode, Conversation, Message, MessageAnnotation, MessageFeedback
+from models.model import AppMode, Conversation, Message, MessageAnnotation
 from services.annotation_service import AppAnnotationService
 from services.errors.conversation import ConversationNotExistsError
 from services.errors.message import MessageNotExistsError, SuggestedQuestionsAfterAnswerDisabledError
@@ -124,33 +125,16 @@ class MessageFeedbackApi(Resource):
         parser.add_argument("rating", type=str, choices=["like", "dislike", None], location="json")
         args = parser.parse_args()
 
-        message_id = str(args["message_id"])
-
-        message = db.session.query(Message).filter(Message.id == message_id, Message.app_id == app_model.id).first()
-
-        if not message:
-            raise NotFound("Message Not Exists.")
-
-        feedback = message.admin_feedback
-
-        if not args["rating"] and feedback:
-            db.session.delete(feedback)
-        elif args["rating"] and feedback:
-            feedback.rating = args["rating"]
-        elif not args["rating"] and not feedback:
-            raise ValueError("rating cannot be None when feedback not exists")
-        else:
-            feedback = MessageFeedback(
-                app_id=app_model.id,
-                conversation_id=message.conversation_id,
-                message_id=message.id,
-                rating=args["rating"],
-                from_source="admin",
-                from_account_id=current_user.id,
+        try:
+            MessageService.create_feedback(
+                app_model=app_model,
+                message_id=str(args["message_id"]),
+                user=current_user,
+                rating=args.get("rating"),
+                content=None,
             )
-            db.session.add(feedback)
-
-        db.session.commit()
+        except services.errors.message.MessageNotExistsError:
+            raise NotFound("Message Not Exists.")
 
         return {"result": "success"}
 

api/controllers/console/app/site.py

@@ -1,5 +1,3 @@
-from datetime import UTC, datetime
-
 from flask_login import current_user
 from flask_restful import Resource, marshal_with, reqparse
 from werkzeug.exceptions import Forbidden, NotFound
@@ -10,6 +8,7 @@ from controllers.console.app.wraps import get_app_model
 from controllers.console.wraps import account_initialization_required, setup_required
 from extensions.ext_database import db
 from fields.app_fields import app_site_fields
+from libs.datetime_utils import naive_utc_now
 from libs.login import login_required
 from models import Site
 
@@ -77,7 +76,7 @@ class AppSite(Resource):
                 setattr(site, attr_name, value)
 
         site.updated_by = current_user.id
-        site.updated_at = datetime.now(UTC).replace(tzinfo=None)
+        site.updated_at = naive_utc_now()
         db.session.commit()
 
         return site
@@ -101,7 +100,7 @@ class AppSiteAccessTokenReset(Resource):
 
         site.code = Site.generate_code(16)
         site.updated_by = current_user.id
-        site.updated_at = datetime.now(UTC).replace(tzinfo=None)
+        site.updated_at = naive_utc_now()
         db.session.commit()
 
         return site

api/controllers/console/auth/activate.py

@@ -1,5 +1,3 @@
-import datetime
-
 from flask import request
 from flask_restful import Resource, reqparse
 
@@ -7,6 +5,7 @@ from constants.languages import supported_language
 from controllers.console import api
 from controllers.console.error import AlreadyActivateError
 from extensions.ext_database import db
+from libs.datetime_utils import naive_utc_now
 from libs.helper import StrLen, email, extract_remote_ip, timezone
 from models.account import AccountStatus
 from services.account_service import AccountService, RegisterService
@@ -65,7 +64,7 @@ class ActivateApi(Resource):
         account.timezone = args["timezone"]
         account.interface_theme = "light"
         account.status = AccountStatus.ACTIVE.value
-        account.initialized_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+        account.initialized_at = naive_utc_now()
         db.session.commit()
 
         token_pair = AccountService.login(account, ip_address=extract_remote_ip(request))

api/controllers/console/auth/oauth.py

@@ -1,5 +1,4 @@
 import logging
-from datetime import UTC, datetime
 from typing import Optional
 
 import requests
@@ -13,6 +12,7 @@ from configs import dify_config
 from constants.languages import languages
 from events.tenant_event import tenant_was_created
 from extensions.ext_database import db
+from libs.datetime_utils import naive_utc_now
 from libs.helper import extract_remote_ip
 from libs.oauth import GitHubOAuth, GoogleOAuth, OAuthUserInfo
 from models import Account
@@ -110,7 +110,7 @@ class OAuthCallback(Resource):
 
         if account.status == AccountStatus.PENDING.value:
             account.status = AccountStatus.ACTIVE.value
-            account.initialized_at = datetime.now(UTC).replace(tzinfo=None)
+            account.initialized_at = naive_utc_now()
             db.session.commit()
 
         try:

api/controllers/console/datasets/data_source.py

@@ -1,4 +1,3 @@
-import datetime
 import json
 
 from flask import request
@@ -15,6 +14,7 @@ from core.rag.extractor.entity.extract_setting import ExtractSetting
 from core.rag.extractor.notion_extractor import NotionExtractor
 from extensions.ext_database import db
 from fields.data_source_fields import integrate_list_fields, integrate_notion_info_list_fields
+from libs.datetime_utils import naive_utc_now
 from libs.login import login_required
 from models import DataSourceOauthBinding, Document
 from services.dataset_service import DatasetService, DocumentService
@@ -88,7 +88,7 @@ class DataSourceApi(Resource):
         if action == "enable":
             if data_source_binding.disabled:
                 data_source_binding.disabled = False
-                data_source_binding.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+                data_source_binding.updated_at = naive_utc_now()
                 db.session.add(data_source_binding)
                 db.session.commit()
             else:
@@ -97,7 +97,7 @@ class DataSourceApi(Resource):
         if action == "disable":
             if not data_source_binding.disabled:
                 data_source_binding.disabled = True
-                data_source_binding.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+                data_source_binding.updated_at = naive_utc_now()
                 db.session.add(data_source_binding)
                 db.session.commit()
             else:

api/controllers/console/datasets/datasets.py

@@ -211,10 +211,6 @@ class DatasetApi(Resource):
         else:
             data["embedding_available"] = True
 
-        if data.get("permission") == "partial_members":
-            part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)
-            data.update({"partial_member_list": part_users_list})
-
         return data, 200
 
     @setup_required

api/controllers/console/datasets/datasets_document.py

@@ -1,6 +1,5 @@
 import logging
 from argparse import ArgumentTypeError
-from datetime import UTC, datetime
 from typing import cast
 
 from flask import request
@@ -49,6 +48,7 @@ from fields.document_fields import (
     document_status_fields,
     document_with_segments_fields,
 )
+from libs.datetime_utils import naive_utc_now
 from libs.login import login_required
 from models import Dataset, DatasetProcessRule, Document, DocumentSegment, UploadFile
 from services.dataset_service import DatasetService, DocumentService
@@ -750,7 +750,7 @@ class DocumentProcessingApi(DocumentResource):
                 raise InvalidActionError("Document not in indexing state.")
 
             document.paused_by = current_user.id
-            document.paused_at = datetime.now(UTC).replace(tzinfo=None)
+            document.paused_at = naive_utc_now()
             document.is_paused = True
             db.session.commit()
 
@@ -830,7 +830,7 @@ class DocumentMetadataApi(DocumentResource):
                     document.doc_metadata[key] = value
 
         document.doc_type = doc_type
-        document.updated_at = datetime.now(UTC).replace(tzinfo=None)
+        document.updated_at = naive_utc_now()
         db.session.commit()
 
         return {"result": "success", "message": "Document metadata updated."}, 200

api/controllers/console/datasets/website.py

@@ -4,7 +4,7 @@ from controllers.console import api
 from controllers.console.datasets.error import WebsiteCrawlError
 from controllers.console.wraps import account_initialization_required, setup_required
 from libs.login import login_required
-from services.website_service import WebsiteService
+from services.website_service import WebsiteCrawlApiRequest, WebsiteCrawlStatusApiRequest, WebsiteService
 
 
 class WebsiteCrawlApi(Resource):
@@ -24,10 +24,16 @@ class WebsiteCrawlApi(Resource):
         parser.add_argument("url", type=str, required=True, nullable=True, location="json")
         parser.add_argument("options", type=dict, required=True, nullable=True, location="json")
         args = parser.parse_args()
-        WebsiteService.document_create_args_validate(args)
-        # crawl url
+
+        # Create typed request and validate
         try:
-            result = WebsiteService.crawl_url(args)
+            api_request = WebsiteCrawlApiRequest.from_args(args)
+        except ValueError as e:
+            raise WebsiteCrawlError(str(e))
+
+        # Crawl URL using typed request
+        try:
+            result = WebsiteService.crawl_url(api_request)
         except Exception as e:
             raise WebsiteCrawlError(str(e))
         return result, 200
@@ -43,9 +49,16 @@ class WebsiteCrawlStatusApi(Resource):
             "provider", type=str, choices=["firecrawl", "watercrawl", "jinareader"], required=True, location="args"
         )
         args = parser.parse_args()
-        # get crawl status
+
+        # Create typed request and validate
         try:
-            result = WebsiteService.get_crawl_status(job_id, args["provider"])
+            api_request = WebsiteCrawlStatusApiRequest.from_args(args, job_id)
+        except ValueError as e:
+            raise WebsiteCrawlError(str(e))
+
+        # Get crawl status using typed request
+        try:
+            result = WebsiteService.get_crawl_status_typed(api_request)
         except Exception as e:
             raise WebsiteCrawlError(str(e))
         return result, 200

api/controllers/console/explore/completion.py

@@ -1,5 +1,4 @@
 import logging
-from datetime import UTC, datetime
 
 from flask_login import current_user
 from flask_restful import reqparse
@@ -27,6 +26,7 @@ from core.errors.error import (
 from core.model_runtime.errors.invoke import InvokeError
 from extensions.ext_database import db
 from libs import helper
+from libs.datetime_utils import naive_utc_now
 from libs.helper import uuid_value
 from models.model import AppMode
 from services.app_generate_service import AppGenerateService
@@ -51,7 +51,7 @@ class CompletionApi(InstalledAppResource):
         streaming = args["response_mode"] == "streaming"
         args["auto_generate_name"] = False
 
-        installed_app.last_used_at = datetime.now(UTC).replace(tzinfo=None)
+        installed_app.last_used_at = naive_utc_now()
         db.session.commit()
 
         try:
@@ -111,7 +111,7 @@ class ChatApi(InstalledAppResource):
 
         args["auto_generate_name"] = False
 
-        installed_app.last_used_at = datetime.now(UTC).replace(tzinfo=None)
+        installed_app.last_used_at = naive_utc_now()
         db.session.commit()
 
         try:

api/controllers/console/explore/installed_app.py

@@ -1,5 +1,4 @@
 import logging
-from datetime import UTC, datetime
 from typing import Any
 
 from flask import request
@@ -13,6 +12,7 @@ from controllers.console.explore.wraps import InstalledAppResource
 from controllers.console.wraps import account_initialization_required, cloud_edition_billing_resource_check
 from extensions.ext_database import db
 from fields.installed_app_fields import installed_app_list_fields
+from libs.datetime_utils import naive_utc_now
 from libs.login import login_required
 from models import App, InstalledApp, RecommendedApp
 from services.account_service import TenantService
@@ -122,7 +122,7 @@ class InstalledAppsListApi(Resource):
                 tenant_id=current_tenant_id,
                 app_owner_tenant_id=app.tenant_id,
                 is_pinned=False,
-                last_used_at=datetime.now(UTC).replace(tzinfo=None),
+                last_used_at=naive_utc_now(),
             )
             db.session.add(new_installed_app)
             db.session.commit()

api/controllers/console/workspace/account.py

@@ -1,5 +1,3 @@
-import datetime
-
 import pytz
 from flask import request
 from flask_login import current_user
@@ -35,6 +33,7 @@ from controllers.console.wraps import (
 )
 from extensions.ext_database import db
 from fields.member_fields import account_fields
+from libs.datetime_utils import naive_utc_now
 from libs.helper import TimestampField, email, extract_remote_ip, timezone
 from libs.login import login_required
 from models import AccountIntegrate, InvitationCode
@@ -80,7 +79,7 @@ class AccountInitApi(Resource):
                 raise InvalidInvitationCodeError()
 
             invitation_code.status = "used"
-            invitation_code.used_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+            invitation_code.used_at = naive_utc_now()
             invitation_code.used_by_tenant_id = account.current_tenant_id
             invitation_code.used_by_account_id = account.id
 
@@ -88,7 +87,7 @@ class AccountInitApi(Resource):
         account.timezone = args["timezone"]
         account.interface_theme = "light"
         account.status = "active"
-        account.initialized_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+        account.initialized_at = naive_utc_now()
         db.session.commit()
 
         return {"result": "success"}

api/controllers/console/workspace/plugin.py

@@ -12,7 +12,8 @@ from controllers.console.wraps import account_initialization_required, setup_req
 from core.model_runtime.utils.encoders import jsonable_encoder
 from core.plugin.impl.exc import PluginDaemonClientSideError
 from libs.login import login_required
-from models.account import TenantPluginPermission
+from models.account import TenantPluginAutoUpgradeStrategy, TenantPluginPermission
+from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService
 from services.plugin.plugin_parameter_service import PluginParameterService
 from services.plugin.plugin_permission_service import PluginPermissionService
 from services.plugin.plugin_service import PluginService
@@ -534,6 +535,114 @@ class PluginFetchDynamicSelectOptionsApi(Resource):
         return jsonable_encoder({"options": options})
 
 
+class PluginChangePreferencesApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self):
+        user = current_user
+        if not user.is_admin_or_owner:
+            raise Forbidden()
+
+        req = reqparse.RequestParser()
+        req.add_argument("permission", type=dict, required=True, location="json")
+        req.add_argument("auto_upgrade", type=dict, required=True, location="json")
+        args = req.parse_args()
+
+        tenant_id = user.current_tenant_id
+
+        permission = args["permission"]
+
+        install_permission = TenantPluginPermission.InstallPermission(permission.get("install_permission", "everyone"))
+        debug_permission = TenantPluginPermission.DebugPermission(permission.get("debug_permission", "everyone"))
+
+        auto_upgrade = args["auto_upgrade"]
+
+        strategy_setting = TenantPluginAutoUpgradeStrategy.StrategySetting(
+            auto_upgrade.get("strategy_setting", "fix_only")
+        )
+        upgrade_time_of_day = auto_upgrade.get("upgrade_time_of_day", 0)
+        upgrade_mode = TenantPluginAutoUpgradeStrategy.UpgradeMode(auto_upgrade.get("upgrade_mode", "exclude"))
+        exclude_plugins = auto_upgrade.get("exclude_plugins", [])
+        include_plugins = auto_upgrade.get("include_plugins", [])
+
+        # set permission
+        set_permission_result = PluginPermissionService.change_permission(
+            tenant_id,
+            install_permission,
+            debug_permission,
+        )
+        if not set_permission_result:
+            return jsonable_encoder({"success": False, "message": "Failed to set permission"})
+
+        # set auto upgrade strategy
+        set_auto_upgrade_strategy_result = PluginAutoUpgradeService.change_strategy(
+            tenant_id,
+            strategy_setting,
+            upgrade_time_of_day,
+            upgrade_mode,
+            exclude_plugins,
+            include_plugins,
+        )
+        if not set_auto_upgrade_strategy_result:
+            return jsonable_encoder({"success": False, "message": "Failed to set auto upgrade strategy"})
+
+        return jsonable_encoder({"success": True})
+
+
+class PluginFetchPreferencesApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self):
+        tenant_id = current_user.current_tenant_id
+
+        permission = PluginPermissionService.get_permission(tenant_id)
+        permission_dict = {
+            "install_permission": TenantPluginPermission.InstallPermission.EVERYONE,
+            "debug_permission": TenantPluginPermission.DebugPermission.EVERYONE,
+        }
+
+        if permission:
+            permission_dict["install_permission"] = permission.install_permission
+            permission_dict["debug_permission"] = permission.debug_permission
+
+        auto_upgrade = PluginAutoUpgradeService.get_strategy(tenant_id)
+        auto_upgrade_dict = {
+            "strategy_setting": TenantPluginAutoUpgradeStrategy.StrategySetting.DISABLED,
+            "upgrade_time_of_day": 0,
+            "upgrade_mode": TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE,
+            "exclude_plugins": [],
+            "include_plugins": [],
+        }
+
+        if auto_upgrade:
+            auto_upgrade_dict = {
+                "strategy_setting": auto_upgrade.strategy_setting,
+                "upgrade_time_of_day": auto_upgrade.upgrade_time_of_day,
+                "upgrade_mode": auto_upgrade.upgrade_mode,
+                "exclude_plugins": auto_upgrade.exclude_plugins,
+                "include_plugins": auto_upgrade.include_plugins,
+            }
+
+        return jsonable_encoder({"permission": permission_dict, "auto_upgrade": auto_upgrade_dict})
+
+
+class PluginAutoUpgradeExcludePluginApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self):
+        # exclude one single plugin
+        tenant_id = current_user.current_tenant_id
+
+        req = reqparse.RequestParser()
+        req.add_argument("plugin_id", type=str, required=True, location="json")
+        args = req.parse_args()
+
+        return jsonable_encoder({"success": PluginAutoUpgradeService.exclude_plugin(tenant_id, args["plugin_id"])})
+
+
 api.add_resource(PluginDebuggingKeyApi, "/workspaces/current/plugin/debugging-key")
 api.add_resource(PluginListApi, "/workspaces/current/plugin/list")
 api.add_resource(PluginListLatestVersionsApi, "/workspaces/current/plugin/list/latest-versions")
@@ -560,3 +669,7 @@ api.add_resource(PluginChangePermissionApi, "/workspaces/current/plugin/permissi
 api.add_resource(PluginFetchPermissionApi, "/workspaces/current/plugin/permission/fetch")
 
 api.add_resource(PluginFetchDynamicSelectOptionsApi, "/workspaces/current/plugin/parameters/dynamic-options")
+
+api.add_resource(PluginFetchPreferencesApi, "/workspaces/current/plugin/preferences/fetch")
+api.add_resource(PluginChangePreferencesApi, "/workspaces/current/plugin/preferences/change")
+api.add_resource(PluginAutoUpgradeExcludePluginApi, "/workspaces/current/plugin/preferences/autoupgrade/exclude")

api/controllers/console/workspace/tool_providers.py

@@ -29,7 +29,7 @@ from libs.login import login_required
 from services.plugin.oauth_service import OAuthProxyService
 from services.tools.api_tools_manage_service import ApiToolManageService
 from services.tools.builtin_tools_manage_service import BuiltinToolManageService
-from services.tools.mcp_tools_mange_service import MCPToolManageService
+from services.tools.mcp_tools_manage_service import MCPToolManageService
 from services.tools.tool_labels_service import ToolLabelsService
 from services.tools.tools_manage_service import ToolCommonService
 from services.tools.tools_transform_service import ToolTransformService
@@ -739,7 +739,7 @@ class ToolOAuthCallback(Resource):
             raise Forbidden("no oauth available client config found for this tool provider")
 
         redirect_uri = f"{dify_config.CONSOLE_API_URL}/console/api/oauth/plugin/{provider}/tool/callback"
-        credentials = oauth_handler.get_credentials(
+        credentials_response = oauth_handler.get_credentials(
             tenant_id=tenant_id,
             user_id=user_id,
             plugin_id=plugin_id,
@@ -747,7 +747,10 @@ class ToolOAuthCallback(Resource):
             redirect_uri=redirect_uri,
             system_credentials=oauth_client_params,
             request=request,
-        ).credentials
+        )
+
+        credentials = credentials_response.credentials
+        expires_at = credentials_response.expires_at
 
         if not credentials:
             raise Exception("the plugin credentials failed")
@@ -758,6 +761,7 @@ class ToolOAuthCallback(Resource):
             tenant_id=tenant_id,
             provider=provider,
             credentials=dict(credentials),
+            expires_at=expires_at,
             api_type=CredentialType.OAUTH2,
         )
         return redirect(f"{dify_config.CONSOLE_WEB_URL}/oauth-callback")

api/controllers/service_api/app/completion.py

@@ -1,5 +1,6 @@
 import logging
 
+from flask import request
 from flask_restful import Resource, reqparse
 from werkzeug.exceptions import InternalServerError, NotFound
 
@@ -23,6 +24,7 @@ from core.errors.error import (
     ProviderTokenNotInitError,
     QuotaExceededError,
 )
+from core.helper.trace_id_helper import get_external_trace_id
 from core.model_runtime.errors.invoke import InvokeError
 from libs import helper
 from libs.helper import uuid_value
@@ -111,6 +113,10 @@ class ChatApi(Resource):
 
         args = parser.parse_args()
 
+        external_trace_id = get_external_trace_id(request)
+        if external_trace_id:
+            args["external_trace_id"] = external_trace_id
+
         streaming = args["response_mode"] == "streaming"
 
         try:

api/controllers/service_api/app/workflow.py

@@ -1,6 +1,7 @@
 import logging
 
 from dateutil.parser import isoparse
+from flask import request
 from flask_restful import Resource, fields, marshal_with, reqparse
 from flask_restful.inputs import int_range
 from sqlalchemy.orm import Session, sessionmaker
@@ -23,6 +24,7 @@ from core.errors.error import (
     ProviderTokenNotInitError,
     QuotaExceededError,
 )
+from core.helper.trace_id_helper import get_external_trace_id
 from core.model_runtime.errors.invoke import InvokeError
 from core.workflow.entities.workflow_execution import WorkflowExecutionStatus
 from extensions.ext_database import db
@@ -90,7 +92,9 @@ class WorkflowRunApi(Resource):
         parser.add_argument("files", type=list, required=False, location="json")
         parser.add_argument("response_mode", type=str, choices=["blocking", "streaming"], location="json")
         args = parser.parse_args()
-
+        external_trace_id = get_external_trace_id(request)
+        if external_trace_id:
+            args["external_trace_id"] = external_trace_id
         streaming = args.get("response_mode") == "streaming"
 
         try:

api/controllers/service_api/wraps.py

@@ -1,6 +1,6 @@
 import time
 from collections.abc import Callable
-from datetime import UTC, datetime, timedelta
+from datetime import timedelta
 from enum import Enum
 from functools import wraps
 from typing import Optional
@@ -15,6 +15,7 @@ from werkzeug.exceptions import Forbidden, NotFound, Unauthorized
 
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
+from libs.datetime_utils import naive_utc_now
 from libs.login import _get_user
 from models.account import Account, Tenant, TenantAccountJoin, TenantStatus
 from models.dataset import Dataset, RateLimitLog
@@ -256,7 +257,7 @@ def validate_and_get_api_token(scope: str | None = None):
     if auth_scheme != "bearer":
         raise Unauthorized("Authorization scheme must be 'Bearer'")
 
-    current_time = datetime.now(UTC).replace(tzinfo=None)
+    current_time = naive_utc_now()
     cutoff_time = current_time - timedelta(minutes=1)
     with Session(db.engine, expire_on_commit=False) as session:
         update_stmt = (

api/core/app/apps/README.md

@@ -1,48 +0,0 @@
-## Guidelines for Database Connection Management in App Runner and Task Pipeline
-
-Due to the presence of tasks in App Runner that require long execution times, such as LLM generation and external requests, Flask-Sqlalchemy's strategy for database connection pooling is to allocate one connection (transaction) per request. This approach keeps a connection occupied even during non-DB tasks, leading to the inability to acquire new connections during high concurrency requests due to multiple long-running tasks.
-
-Therefore, the database operations in App Runner and Task Pipeline must ensure connections are closed immediately after use, and it's better to pass IDs rather than Model objects to avoid detach errors.
-
-Examples:
-
-1. Creating a new record:
-
-   ```python
-   app = App(id=1)
-   db.session.add(app)
-   db.session.commit()
-   db.session.refresh(app)  # Retrieve table default values, like created_at, cached in the app object, won't affect after close
-   
-   # Handle non-long-running tasks or store the content of the App instance in memory (via variable assignment).
-   
-   db.session.close()
-   
-   return app.id
-   ```
-
-2. Fetching a record from the table:
-
-   ```python
-   app = db.session.query(App).filter(App.id == app_id).first()
-    
-   created_at = app.created_at
-    
-   db.session.close()
-   
-   # Handle tasks (include long-running).
-   
-   ```
-
-3. Updating a table field:
-
-   ```python
-   app = db.session.query(App).filter(App.id == app_id).first()
-
-   app.updated_at = time.utcnow()
-   db.session.commit()
-   db.session.close()
-
-   return app_id
-   ```
-   

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

@@ -7,7 +7,8 @@ from typing import Any, Literal, Optional, Union, overload
 
 from flask import Flask, current_app
 from pydantic import ValidationError
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy import select
+from sqlalchemy.orm import Session, sessionmaker
 
 import contexts
 from configs import dify_config
@@ -17,11 +18,13 @@ from core.app.apps.advanced_chat.app_config_manager import AdvancedChatAppConfig
 from core.app.apps.advanced_chat.app_runner import AdvancedChatAppRunner
 from core.app.apps.advanced_chat.generate_response_converter import AdvancedChatAppGenerateResponseConverter
 from core.app.apps.advanced_chat.generate_task_pipeline import AdvancedChatAppGenerateTaskPipeline
-from core.app.apps.base_app_queue_manager import AppQueueManager, GenerateTaskStoppedError, PublishFrom
+from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
+from core.app.apps.exc import GenerateTaskStoppedError
 from core.app.apps.message_based_app_generator import MessageBasedAppGenerator
 from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager
 from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom
 from core.app.entities.task_entities import ChatbotAppBlockingResponse, ChatbotAppStreamResponse
+from core.helper.trace_id_helper import extract_external_trace_id_from_args
 from core.model_runtime.errors.invoke import InvokeAuthorizationError
 from core.ops.ops_trace_manager import TraceQueueManager
 from core.prompt.utils.get_thread_messages_length import get_thread_messages_length
@@ -111,7 +114,10 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
         query = query.replace("\x00", "")
         inputs = args["inputs"]
 
-        extras = {"auto_generate_conversation_name": args.get("auto_generate_name", False)}
+        extras = {
+            "auto_generate_conversation_name": args.get("auto_generate_name", False),
+            **extract_external_trace_id_from_args(args),
+        }
 
         # get conversation
         conversation = None
@@ -481,21 +487,52 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
         """
 
         with preserve_flask_contexts(flask_app, context_vars=context):
-            try:
-                # get conversation and message
-                conversation = self._get_conversation(conversation_id)
-                message = self._get_message(message_id)
+            # get conversation and message
+            conversation = self._get_conversation(conversation_id)
+            message = self._get_message(message_id)
 
-                # chatbot app
-                runner = AdvancedChatAppRunner(
-                    application_generate_entity=application_generate_entity,
-                    queue_manager=queue_manager,
-                    conversation=conversation,
-                    message=message,
-                    dialogue_count=self._dialogue_count,
-                    variable_loader=variable_loader,
+            with Session(db.engine, expire_on_commit=False) as session:
+                workflow = session.scalar(
+                    select(Workflow).where(
+                        Workflow.tenant_id == application_generate_entity.app_config.tenant_id,
+                        Workflow.app_id == application_generate_entity.app_config.app_id,
+                        Workflow.id == application_generate_entity.app_config.workflow_id,
+                    )
                 )
+                if workflow is None:
+                    raise ValueError("Workflow not found")
 
+                # Determine system_user_id based on invocation source
+                is_external_api_call = application_generate_entity.invoke_from in {
+                    InvokeFrom.WEB_APP,
+                    InvokeFrom.SERVICE_API,
+                }
+
+                if is_external_api_call:
+                    # For external API calls, use end user's session ID
+                    end_user = session.scalar(select(EndUser).where(EndUser.id == application_generate_entity.user_id))
+                    system_user_id = end_user.session_id if end_user else ""
+                else:
+                    # For internal calls, use the original user ID
+                    system_user_id = application_generate_entity.user_id
+
+                app = session.scalar(select(App).where(App.id == application_generate_entity.app_config.app_id))
+                if app is None:
+                    raise ValueError("App not found")
+
+            runner = AdvancedChatAppRunner(
+                application_generate_entity=application_generate_entity,
+                queue_manager=queue_manager,
+                conversation=conversation,
+                message=message,
+                dialogue_count=self._dialogue_count,
+                variable_loader=variable_loader,
+                workflow=workflow,
+                system_user_id=system_user_id,
+                app=app,
+            )
+
+            try:
                 runner.run()
             except GenerateTaskStoppedError:
                 pass

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

@@ -1,6 +1,6 @@
 import logging
 from collections.abc import Mapping
-from typing import Any, cast
+from typing import Any, Optional, cast
 
 from sqlalchemy import select
 from sqlalchemy.orm import Session
@@ -9,13 +9,19 @@ from configs import dify_config
 from core.app.apps.advanced_chat.app_config_manager import AdvancedChatAppConfig
 from core.app.apps.base_app_queue_manager import AppQueueManager
 from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner
-from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom
+from core.app.entities.app_invoke_entities import (
+    AdvancedChatAppGenerateEntity,
+    AppGenerateEntity,
+    InvokeFrom,
+)
 from core.app.entities.queue_entities import (
     QueueAnnotationReplyEvent,
     QueueStopEvent,
     QueueTextChunkEvent,
 )
+from core.app.features.annotation_reply.annotation_reply import AnnotationReplyFeature
 from core.moderation.base import ModerationError
+from core.moderation.input_moderation import InputModeration
 from core.variables.variables import VariableUnion
 from core.workflow.callbacks import WorkflowCallback, WorkflowLoggingCallback
 from core.workflow.entities.variable_pool import VariablePool
@@ -23,8 +29,9 @@ from core.workflow.system_variable import SystemVariable
 from core.workflow.variable_loader import VariableLoader
 from core.workflow.workflow_entry import WorkflowEntry
 from extensions.ext_database import db
+from models import Workflow
 from models.enums import UserFrom
-from models.model import App, Conversation, EndUser, Message
+from models.model import App, Conversation, Message, MessageAnnotation
 from models.workflow import ConversationVariable, WorkflowType
 
 logger = logging.getLogger(__name__)
@@ -37,21 +44,29 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
 
     def __init__(
         self,
+        *,
         application_generate_entity: AdvancedChatAppGenerateEntity,
         queue_manager: AppQueueManager,
         conversation: Conversation,
         message: Message,
         dialogue_count: int,
         variable_loader: VariableLoader,
+        workflow: Workflow,
+        system_user_id: str,
+        app: App,
     ) -> None:
-        super().__init__(queue_manager, variable_loader)
+        super().__init__(
+            queue_manager=queue_manager,
+            variable_loader=variable_loader,
+            app_id=application_generate_entity.app_config.app_id,
+        )
         self.application_generate_entity = application_generate_entity
         self.conversation = conversation
         self.message = message
         self._dialogue_count = dialogue_count
-
-    def _get_app_id(self) -> str:
-        return self.application_generate_entity.app_config.app_id
+        self._workflow = workflow
+        self.system_user_id = system_user_id
+        self._app = app
 
     def run(self) -> None:
         app_config = self.application_generate_entity.app_config
@@ -61,18 +76,6 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
         if not app_record:
             raise ValueError("App not found")
 
-        workflow = self.get_workflow(app_model=app_record, workflow_id=app_config.workflow_id)
-        if not workflow:
-            raise ValueError("Workflow not initialized")
-
-        user_id: str | None = None
-        if self.application_generate_entity.invoke_from in {InvokeFrom.WEB_APP, InvokeFrom.SERVICE_API}:
-            end_user = db.session.query(EndUser).filter(EndUser.id == self.application_generate_entity.user_id).first()
-            if end_user:
-                user_id = end_user.session_id
-        else:
-            user_id = self.application_generate_entity.user_id
-
         workflow_callbacks: list[WorkflowCallback] = []
         if dify_config.DEBUG:
             workflow_callbacks.append(WorkflowLoggingCallback())
@@ -80,14 +83,14 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
         if self.application_generate_entity.single_iteration_run:
             # if only single iteration run is requested
             graph, variable_pool = self._get_graph_and_variable_pool_of_single_iteration(
-                workflow=workflow,
+                workflow=self._workflow,
                 node_id=self.application_generate_entity.single_iteration_run.node_id,
                 user_inputs=dict(self.application_generate_entity.single_iteration_run.inputs),
             )
         elif self.application_generate_entity.single_loop_run:
             # if only single loop run is requested
             graph, variable_pool = self._get_graph_and_variable_pool_of_single_loop(
-                workflow=workflow,
+                workflow=self._workflow,
                 node_id=self.application_generate_entity.single_loop_run.node_id,
                 user_inputs=dict(self.application_generate_entity.single_loop_run.inputs),
             )
@@ -98,7 +101,7 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
 
             # moderation
             if self.handle_input_moderation(
-                app_record=app_record,
+                app_record=self._app,
                 app_generate_entity=self.application_generate_entity,
                 inputs=inputs,
                 query=query,
@@ -108,7 +111,7 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
 
             # annotation reply
             if self.handle_annotation_reply(
-                app_record=app_record,
+                app_record=self._app,
                 message=self.message,
                 query=query,
                 app_generate_entity=self.application_generate_entity,
@@ -128,7 +131,7 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
                         ConversationVariable.from_variable(
                             app_id=self.conversation.app_id, conversation_id=self.conversation.id, variable=variable
                         )
-                        for variable in workflow.conversation_variables
+                        for variable in self._workflow.conversation_variables
                     ]
                     session.add_all(db_conversation_variables)
                 # Convert database entities to variables.
@@ -141,7 +144,7 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
                 query=query,
                 files=files,
                 conversation_id=self.conversation.id,
-                user_id=user_id,
+                user_id=self.system_user_id,
                 dialogue_count=self._dialogue_count,
                 app_id=app_config.app_id,
                 workflow_id=app_config.workflow_id,
@@ -152,25 +155,25 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
             variable_pool = VariablePool(
                 system_variables=system_inputs,
                 user_inputs=inputs,
-                environment_variables=workflow.environment_variables,
+                environment_variables=self._workflow.environment_variables,
                 # Based on the definition of `VariableUnion`,
                 # `list[Variable]` can be safely used as `list[VariableUnion]` since they are compatible.
                 conversation_variables=cast(list[VariableUnion], conversation_variables),
             )
 
             # init graph
-            graph = self._init_graph(graph_config=workflow.graph_dict)
+            graph = self._init_graph(graph_config=self._workflow.graph_dict)
 
         db.session.close()
 
         # RUN WORKFLOW
         workflow_entry = WorkflowEntry(
-            tenant_id=workflow.tenant_id,
-            app_id=workflow.app_id,
-            workflow_id=workflow.id,
-            workflow_type=WorkflowType.value_of(workflow.type),
+            tenant_id=self._workflow.tenant_id,
+            app_id=self._workflow.app_id,
+            workflow_id=self._workflow.id,
+            workflow_type=WorkflowType.value_of(self._workflow.type),
             graph=graph,
-            graph_config=workflow.graph_dict,
+            graph_config=self._workflow.graph_dict,
             user_id=self.application_generate_entity.user_id,
             user_from=(
                 UserFrom.ACCOUNT
@@ -241,3 +244,51 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
         self._publish_event(QueueTextChunkEvent(text=text))
 
         self._publish_event(QueueStopEvent(stopped_by=stopped_by))
+
+    def query_app_annotations_to_reply(
+        self, app_record: App, message: Message, query: str, user_id: str, invoke_from: InvokeFrom
+    ) -> Optional[MessageAnnotation]:
+        """
+        Query app annotations to reply
+        :param app_record: app record
+        :param message: message
+        :param query: query
+        :param user_id: user id
+        :param invoke_from: invoke from
+        :return:
+        """
+        annotation_reply_feature = AnnotationReplyFeature()
+        return annotation_reply_feature.query(
+            app_record=app_record, message=message, query=query, user_id=user_id, invoke_from=invoke_from
+        )
+
+    def moderation_for_inputs(
+        self,
+        *,
+        app_id: str,
+        tenant_id: str,
+        app_generate_entity: AppGenerateEntity,
+        inputs: Mapping[str, Any],
+        query: str | None = None,
+        message_id: str,
+    ) -> tuple[bool, Mapping[str, Any], str]:
+        """
+        Process sensitive_word_avoidance.
+        :param app_id: app id
+        :param tenant_id: tenant id
+        :param app_generate_entity: app generate entity
+        :param inputs: inputs
+        :param query: query
+        :param message_id: message id
+        :return:
+        """
+        moderation_feature = InputModeration()
+        return moderation_feature.check(
+            app_id=app_id,
+            tenant_id=tenant_id,
+            app_config=app_generate_entity.app_config,
+            inputs=dict(inputs),
+            query=query or "",
+            message_id=message_id,
+            trace_manager=app_generate_entity.trace_manager,
+        )

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

@@ -15,7 +15,8 @@ from core.app.app_config.features.file_upload.manager import FileUploadConfigMan
 from core.app.apps.agent_chat.app_config_manager import AgentChatAppConfigManager
 from core.app.apps.agent_chat.app_runner import AgentChatAppRunner
 from core.app.apps.agent_chat.generate_response_converter import AgentChatAppGenerateResponseConverter
-from core.app.apps.base_app_queue_manager import AppQueueManager, GenerateTaskStoppedError, PublishFrom
+from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
+from core.app.apps.exc import GenerateTaskStoppedError
 from core.app.apps.message_based_app_generator import MessageBasedAppGenerator
 from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager
 from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, InvokeFrom

api/core/app/apps/base_app_queue_manager.py

@@ -169,7 +169,3 @@ class AppQueueManager:
                 raise TypeError(
                     "Critical Error: Passing SQLAlchemy Model instances that cause thread safety issues is not allowed."
                 )
-
-
-class GenerateTaskStoppedError(Exception):
-    pass

api/core/app/apps/base_app_runner.py

@@ -118,7 +118,7 @@ class AppRunner:
         else:
             memory_config = MemoryConfig(window=MemoryConfig.WindowConfig(enabled=False))
 
-            model_mode = ModelMode.value_of(model_config.mode)
+            model_mode = ModelMode(model_config.mode)
             prompt_template: Union[CompletionModelPromptTemplate, list[ChatModelMessage]]
             if model_mode == ModelMode.COMPLETION:
                 advanced_completion_prompt_template = prompt_template_entity.advanced_completion_prompt_template

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

@@ -11,10 +11,11 @@ from configs import dify_config
 from constants import UUID_NIL
 from core.app.app_config.easy_ui_based_app.model_config.converter import ModelConfigConverter
 from core.app.app_config.features.file_upload.manager import FileUploadConfigManager
-from core.app.apps.base_app_queue_manager import AppQueueManager, GenerateTaskStoppedError, PublishFrom
+from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
 from core.app.apps.chat.app_config_manager import ChatAppConfigManager
 from core.app.apps.chat.app_runner import ChatAppRunner
 from core.app.apps.chat.generate_response_converter import ChatAppGenerateResponseConverter
+from core.app.apps.exc import GenerateTaskStoppedError
 from core.app.apps.message_based_app_generator import MessageBasedAppGenerator
 from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager
 from core.app.entities.app_invoke_entities import ChatAppGenerateEntity, InvokeFrom

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

@@ -10,10 +10,11 @@ from pydantic import ValidationError
 from configs import dify_config
 from core.app.app_config.easy_ui_based_app.model_config.converter import ModelConfigConverter
 from core.app.app_config.features.file_upload.manager import FileUploadConfigManager
-from core.app.apps.base_app_queue_manager import AppQueueManager, GenerateTaskStoppedError, PublishFrom
+from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
 from core.app.apps.completion.app_config_manager import CompletionAppConfigManager
 from core.app.apps.completion.app_runner import CompletionAppRunner
 from core.app.apps.completion.generate_response_converter import CompletionAppGenerateResponseConverter
+from core.app.apps.exc import GenerateTaskStoppedError
 from core.app.apps.message_based_app_generator import MessageBasedAppGenerator
 from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager
 from core.app.entities.app_invoke_entities import CompletionAppGenerateEntity, InvokeFrom

api/core/app/apps/exc.py

@@ -0,0 +1,2 @@
+class GenerateTaskStoppedError(Exception):
+    pass

api/core/app/apps/message_based_app_generator.py

@@ -1,12 +1,12 @@
 import json
 import logging
 from collections.abc import Generator
-from datetime import UTC, datetime
 from typing import Optional, Union, cast
 
 from core.app.app_config.entities import EasyUIBasedAppConfig, EasyUIBasedAppModelConfigFrom
 from core.app.apps.base_app_generator import BaseAppGenerator
-from core.app.apps.base_app_queue_manager import AppQueueManager, GenerateTaskStoppedError
+from core.app.apps.base_app_queue_manager import AppQueueManager
+from core.app.apps.exc import GenerateTaskStoppedError
 from core.app.entities.app_invoke_entities import (
     AdvancedChatAppGenerateEntity,
     AgentChatAppGenerateEntity,
@@ -24,6 +24,7 @@ from core.app.entities.task_entities import (
 from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBasedGenerateTaskPipeline
 from core.prompt.utils.prompt_template_parser import PromptTemplateParser
 from extensions.ext_database import db
+from libs.datetime_utils import naive_utc_now
 from models import Account
 from models.enums import CreatorUserRole
 from models.model import App, AppMode, AppModelConfig, Conversation, EndUser, Message, MessageFile
@@ -183,7 +184,7 @@ class MessageBasedAppGenerator(BaseAppGenerator):
             db.session.commit()
             db.session.refresh(conversation)
         else:
-            conversation.updated_at = datetime.now(UTC).replace(tzinfo=None)
+            conversation.updated_at = naive_utc_now()
             db.session.commit()
 
         message = Message(

api/core/app/apps/message_based_app_queue_manager.py

@@ -1,4 +1,5 @@
-from core.app.apps.base_app_queue_manager import AppQueueManager, GenerateTaskStoppedError, PublishFrom
+from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
+from core.app.apps.exc import GenerateTaskStoppedError
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.app.entities.queue_entities import (
     AppQueueEvent,

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

@@ -7,13 +7,15 @@ from typing import Any, Literal, Optional, Union, overload
 
 from flask import Flask, current_app
 from pydantic import ValidationError
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy import select
+from sqlalchemy.orm import Session, sessionmaker
 
 import contexts
 from configs import dify_config
 from core.app.app_config.features.file_upload.manager import FileUploadConfigManager
 from core.app.apps.base_app_generator import BaseAppGenerator
-from core.app.apps.base_app_queue_manager import AppQueueManager, GenerateTaskStoppedError, PublishFrom
+from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
+from core.app.apps.exc import GenerateTaskStoppedError
 from core.app.apps.workflow.app_config_manager import WorkflowAppConfigManager
 from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager
 from core.app.apps.workflow.app_runner import WorkflowAppRunner
@@ -21,6 +23,7 @@ from core.app.apps.workflow.generate_response_converter import WorkflowAppGenera
 from core.app.apps.workflow.generate_task_pipeline import WorkflowAppGenerateTaskPipeline
 from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity
 from core.app.entities.task_entities import WorkflowAppBlockingResponse, WorkflowAppStreamResponse
+from core.helper.trace_id_helper import extract_external_trace_id_from_args
 from core.model_runtime.errors.invoke import InvokeAuthorizationError
 from core.ops.ops_trace_manager import TraceQueueManager
 from core.repositories import DifyCoreRepositoryFactory
@@ -122,6 +125,10 @@ class WorkflowAppGenerator(BaseAppGenerator):
         )
 
         inputs: Mapping[str, Any] = args["inputs"]
+
+        extras = {
+            **extract_external_trace_id_from_args(args),
+        }
         workflow_run_id = str(uuid.uuid4())
         # init application generate entity
         application_generate_entity = WorkflowAppGenerateEntity(
@@ -141,6 +148,7 @@ class WorkflowAppGenerator(BaseAppGenerator):
             call_depth=call_depth,
             trace_manager=trace_manager,
             workflow_execution_id=workflow_run_id,
+            extras=extras,
         )
 
         contexts.plugin_tool_providers.set({})
@@ -438,17 +446,44 @@ class WorkflowAppGenerator(BaseAppGenerator):
         """
 
         with preserve_flask_contexts(flask_app, context_vars=context):
-            try:
-                # workflow app
-                runner = WorkflowAppRunner(
-                    application_generate_entity=application_generate_entity,
-                    queue_manager=queue_manager,
-                    workflow_thread_pool_id=workflow_thread_pool_id,
-                    variable_loader=variable_loader,
+            with Session(db.engine, expire_on_commit=False) as session:
+                workflow = session.scalar(
+                    select(Workflow).where(
+                        Workflow.tenant_id == application_generate_entity.app_config.tenant_id,
+                        Workflow.app_id == application_generate_entity.app_config.app_id,
+                        Workflow.id == application_generate_entity.app_config.workflow_id,
+                    )
                 )
+                if workflow is None:
+                    raise ValueError("Workflow not found")
 
+                # Determine system_user_id based on invocation source
+                is_external_api_call = application_generate_entity.invoke_from in {
+                    InvokeFrom.WEB_APP,
+                    InvokeFrom.SERVICE_API,
+                }
+
+                if is_external_api_call:
+                    # For external API calls, use end user's session ID
+                    end_user = session.scalar(select(EndUser).where(EndUser.id == application_generate_entity.user_id))
+                    system_user_id = end_user.session_id if end_user else ""
+                else:
+                    # For internal calls, use the original user ID
+                    system_user_id = application_generate_entity.user_id
+
+            runner = WorkflowAppRunner(
+                application_generate_entity=application_generate_entity,
+                queue_manager=queue_manager,
+                workflow_thread_pool_id=workflow_thread_pool_id,
+                variable_loader=variable_loader,
+                workflow=workflow,
+                system_user_id=system_user_id,
+            )
+
+            try:
                 runner.run()
-            except GenerateTaskStoppedError:
+            except GenerateTaskStoppedError as e:
+                logger.warning(f"Task stopped: {str(e)}")
                 pass
             except InvokeAuthorizationError:
                 queue_manager.publish_error(
@@ -464,8 +499,6 @@ class WorkflowAppGenerator(BaseAppGenerator):
             except Exception as e:
                 logger.exception("Unknown Error when generating")
                 queue_manager.publish_error(e, PublishFrom.APPLICATION_MANAGER)
-            finally:
-                db.session.close()
 
     def _handle_response(
         self,

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

@@ -1,4 +1,5 @@
-from core.app.apps.base_app_queue_manager import AppQueueManager, GenerateTaskStoppedError, PublishFrom
+from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
+from core.app.apps.exc import GenerateTaskStoppedError
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.app.entities.queue_entities import (
     AppQueueEvent,

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

@@ -14,10 +14,8 @@ from core.workflow.entities.variable_pool import VariablePool
 from core.workflow.system_variable import SystemVariable
 from core.workflow.variable_loader import VariableLoader
 from core.workflow.workflow_entry import WorkflowEntry
-from extensions.ext_database import db
 from models.enums import UserFrom
-from models.model import App, EndUser
-from models.workflow import WorkflowType
+from models.workflow import Workflow, WorkflowType
 
 logger = logging.getLogger(__name__)
 
@@ -29,22 +27,23 @@ class WorkflowAppRunner(WorkflowBasedAppRunner):
 
     def __init__(
         self,
+        *,
         application_generate_entity: WorkflowAppGenerateEntity,
         queue_manager: AppQueueManager,
         variable_loader: VariableLoader,
         workflow_thread_pool_id: Optional[str] = None,
+        workflow: Workflow,
+        system_user_id: str,
     ) -> None:
-        """
-        :param application_generate_entity: application generate entity
-        :param queue_manager: application queue manager
-        :param workflow_thread_pool_id: workflow thread pool id
-        """
-        super().__init__(queue_manager, variable_loader)
+        super().__init__(
+            queue_manager=queue_manager,
+            variable_loader=variable_loader,
+            app_id=application_generate_entity.app_config.app_id,
+        )
         self.application_generate_entity = application_generate_entity
         self.workflow_thread_pool_id = workflow_thread_pool_id
-
-    def _get_app_id(self) -> str:
-        return self.application_generate_entity.app_config.app_id
+        self._workflow = workflow
+        self._sys_user_id = system_user_id
 
     def run(self) -> None:
         """
@@ -53,24 +52,6 @@ class WorkflowAppRunner(WorkflowBasedAppRunner):
         app_config = self.application_generate_entity.app_config
         app_config = cast(WorkflowAppConfig, app_config)
 
-        user_id = None
-        if self.application_generate_entity.invoke_from in {InvokeFrom.WEB_APP, InvokeFrom.SERVICE_API}:
-            end_user = db.session.query(EndUser).filter(EndUser.id == self.application_generate_entity.user_id).first()
-            if end_user:
-                user_id = end_user.session_id
-        else:
-            user_id = self.application_generate_entity.user_id
-
-        app_record = db.session.query(App).filter(App.id == app_config.app_id).first()
-        if not app_record:
-            raise ValueError("App not found")
-
-        workflow = self.get_workflow(app_model=app_record, workflow_id=app_config.workflow_id)
-        if not workflow:
-            raise ValueError("Workflow not initialized")
-
-        db.session.close()
-
         workflow_callbacks: list[WorkflowCallback] = []
         if dify_config.DEBUG:
             workflow_callbacks.append(WorkflowLoggingCallback())
@@ -79,14 +60,14 @@ class WorkflowAppRunner(WorkflowBasedAppRunner):
         if self.application_generate_entity.single_iteration_run:
             # if only single iteration run is requested
             graph, variable_pool = self._get_graph_and_variable_pool_of_single_iteration(
-                workflow=workflow,
+                workflow=self._workflow,
                 node_id=self.application_generate_entity.single_iteration_run.node_id,
                 user_inputs=self.application_generate_entity.single_iteration_run.inputs,
             )
         elif self.application_generate_entity.single_loop_run:
             # if only single loop run is requested
             graph, variable_pool = self._get_graph_and_variable_pool_of_single_loop(
-                workflow=workflow,
+                workflow=self._workflow,
                 node_id=self.application_generate_entity.single_loop_run.node_id,
                 user_inputs=self.application_generate_entity.single_loop_run.inputs,
             )
@@ -98,7 +79,7 @@ class WorkflowAppRunner(WorkflowBasedAppRunner):
 
             system_inputs = SystemVariable(
                 files=files,
-                user_id=user_id,
+                user_id=self._sys_user_id,
                 app_id=app_config.app_id,
                 workflow_id=app_config.workflow_id,
                 workflow_execution_id=self.application_generate_entity.workflow_execution_id,
@@ -107,21 +88,21 @@ class WorkflowAppRunner(WorkflowBasedAppRunner):
             variable_pool = VariablePool(
                 system_variables=system_inputs,
                 user_inputs=inputs,
-                environment_variables=workflow.environment_variables,
+                environment_variables=self._workflow.environment_variables,
                 conversation_variables=[],
             )
 
             # init graph
-            graph = self._init_graph(graph_config=workflow.graph_dict)
+            graph = self._init_graph(graph_config=self._workflow.graph_dict)
 
         # RUN WORKFLOW
         workflow_entry = WorkflowEntry(
-            tenant_id=workflow.tenant_id,
-            app_id=workflow.app_id,
-            workflow_id=workflow.id,
-            workflow_type=WorkflowType.value_of(workflow.type),
+            tenant_id=self._workflow.tenant_id,
+            app_id=self._workflow.app_id,
+            workflow_id=self._workflow.id,
+            workflow_type=WorkflowType.value_of(self._workflow.type),
             graph=graph,
-            graph_config=workflow.graph_dict,
+            graph_config=self._workflow.graph_dict,
             user_id=self.application_generate_entity.user_id,
             user_from=(
                 UserFrom.ACCOUNT

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

@@ -1,7 +1,8 @@
 import logging
 import time
-from collections.abc import Generator
-from typing import Optional, Union
+from collections.abc import Callable, Generator
+from contextlib import contextmanager
+from typing import Any, Optional, Union
 
 from sqlalchemy.orm import Session
 
@@ -13,6 +14,7 @@ from core.app.entities.app_invoke_entities import (
     WorkflowAppGenerateEntity,
 )
 from core.app.entities.queue_entities import (
+    MessageQueueMessage,
     QueueAgentLogEvent,
     QueueErrorEvent,
     QueueIterationCompletedEvent,
@@ -38,11 +40,13 @@ from core.app.entities.queue_entities import (
     QueueWorkflowPartialSuccessEvent,
     QueueWorkflowStartedEvent,
     QueueWorkflowSucceededEvent,
+    WorkflowQueueMessage,
 )
 from core.app.entities.task_entities import (
     ErrorStreamResponse,
     MessageAudioEndStreamResponse,
     MessageAudioStreamResponse,
+    PingStreamResponse,
     StreamResponse,
     TextChunkStreamResponse,
     WorkflowAppBlockingResponse,
@@ -54,6 +58,7 @@ from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTas
 from core.base.tts import AppGeneratorTTSPublisher, AudioTrunk
 from core.ops.ops_trace_manager import TraceQueueManager
 from core.workflow.entities.workflow_execution import WorkflowExecution, WorkflowExecutionStatus, WorkflowType
+from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState
 from core.workflow.repositories.draft_variable_repository import DraftVariableSaverFactory
 from core.workflow.repositories.workflow_execution_repository import WorkflowExecutionRepository
 from core.workflow.repositories.workflow_node_execution_repository import WorkflowNodeExecutionRepository
@@ -246,315 +251,495 @@ class WorkflowAppGenerateTaskPipeline:
         if tts_publisher:
             yield MessageAudioEndStreamResponse(audio="", task_id=task_id)
 
+    @contextmanager
+    def _database_session(self):
+        """Context manager for database sessions."""
+        with Session(db.engine, expire_on_commit=False) as session:
+            try:
+                yield session
+                session.commit()
+            except Exception:
+                session.rollback()
+                raise
+
+    def _ensure_workflow_initialized(self) -> None:
+        """Fluent validation for workflow state."""
+        if not self._workflow_run_id:
+            raise ValueError("workflow run not initialized.")
+
+    def _ensure_graph_runtime_initialized(self, graph_runtime_state: Optional[GraphRuntimeState]) -> GraphRuntimeState:
+        """Fluent validation for graph runtime state."""
+        if not graph_runtime_state:
+            raise ValueError("graph runtime state not initialized.")
+        return graph_runtime_state
+
+    def _handle_ping_event(self, event: QueuePingEvent, **kwargs) -> Generator[PingStreamResponse, None, None]:
+        """Handle ping events."""
+        yield self._base_task_pipeline._ping_stream_response()
+
+    def _handle_error_event(self, event: QueueErrorEvent, **kwargs) -> Generator[ErrorStreamResponse, None, None]:
+        """Handle error events."""
+        err = self._base_task_pipeline._handle_error(event=event)
+        yield self._base_task_pipeline._error_to_stream_response(err)
+
+    def _handle_workflow_started_event(
+        self, event: QueueWorkflowStartedEvent, **kwargs
+    ) -> Generator[StreamResponse, None, None]:
+        """Handle workflow started events."""
+        # init workflow run
+        workflow_execution = self._workflow_cycle_manager.handle_workflow_run_start()
+        self._workflow_run_id = workflow_execution.id_
+        start_resp = self._workflow_response_converter.workflow_start_to_stream_response(
+            task_id=self._application_generate_entity.task_id,
+            workflow_execution=workflow_execution,
+        )
+        yield start_resp
+
+    def _handle_node_retry_event(self, event: QueueNodeRetryEvent, **kwargs) -> Generator[StreamResponse, None, None]:
+        """Handle node retry events."""
+        self._ensure_workflow_initialized()
+
+        with self._database_session() as session:
+            workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_retried(
+                workflow_execution_id=self._workflow_run_id,
+                event=event,
+            )
+            response = self._workflow_response_converter.workflow_node_retry_to_stream_response(
+                event=event,
+                task_id=self._application_generate_entity.task_id,
+                workflow_node_execution=workflow_node_execution,
+            )
+
+        if response:
+            yield response
+
+    def _handle_node_started_event(
+        self, event: QueueNodeStartedEvent, **kwargs
+    ) -> Generator[StreamResponse, None, None]:
+        """Handle node started events."""
+        self._ensure_workflow_initialized()
+
+        workflow_node_execution = self._workflow_cycle_manager.handle_node_execution_start(
+            workflow_execution_id=self._workflow_run_id, event=event
+        )
+        node_start_response = self._workflow_response_converter.workflow_node_start_to_stream_response(
+            event=event,
+            task_id=self._application_generate_entity.task_id,
+            workflow_node_execution=workflow_node_execution,
+        )
+
+        if node_start_response:
+            yield node_start_response
+
+    def _handle_node_succeeded_event(
+        self, event: QueueNodeSucceededEvent, **kwargs
+    ) -> Generator[StreamResponse, None, None]:
+        """Handle node succeeded events."""
+        workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_success(event=event)
+        node_success_response = self._workflow_response_converter.workflow_node_finish_to_stream_response(
+            event=event,
+            task_id=self._application_generate_entity.task_id,
+            workflow_node_execution=workflow_node_execution,
+        )
+
+        self._save_output_for_event(event, workflow_node_execution.id)
+
+        if node_success_response:
+            yield node_success_response
+
+    def _handle_node_failed_events(
+        self,
+        event: Union[
+            QueueNodeFailedEvent, QueueNodeInIterationFailedEvent, QueueNodeInLoopFailedEvent, QueueNodeExceptionEvent
+        ],
+        **kwargs,
+    ) -> Generator[StreamResponse, None, None]:
+        """Handle various node failure events."""
+        workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_failed(
+            event=event,
+        )
+        node_failed_response = self._workflow_response_converter.workflow_node_finish_to_stream_response(
+            event=event,
+            task_id=self._application_generate_entity.task_id,
+            workflow_node_execution=workflow_node_execution,
+        )
+
+        if isinstance(event, QueueNodeExceptionEvent):
+            self._save_output_for_event(event, workflow_node_execution.id)
+
+        if node_failed_response:
+            yield node_failed_response
+
+    def _handle_parallel_branch_started_event(
+        self, event: QueueParallelBranchRunStartedEvent, **kwargs
+    ) -> Generator[StreamResponse, None, None]:
+        """Handle parallel branch started events."""
+        self._ensure_workflow_initialized()
+
+        parallel_start_resp = self._workflow_response_converter.workflow_parallel_branch_start_to_stream_response(
+            task_id=self._application_generate_entity.task_id,
+            workflow_execution_id=self._workflow_run_id,
+            event=event,
+        )
+        yield parallel_start_resp
+
+    def _handle_parallel_branch_finished_events(
+        self, event: Union[QueueParallelBranchRunSucceededEvent, QueueParallelBranchRunFailedEvent], **kwargs
+    ) -> Generator[StreamResponse, None, None]:
+        """Handle parallel branch finished events."""
+        self._ensure_workflow_initialized()
+
+        parallel_finish_resp = self._workflow_response_converter.workflow_parallel_branch_finished_to_stream_response(
+            task_id=self._application_generate_entity.task_id,
+            workflow_execution_id=self._workflow_run_id,
+            event=event,
+        )
+        yield parallel_finish_resp
+
+    def _handle_iteration_start_event(
+        self, event: QueueIterationStartEvent, **kwargs
+    ) -> Generator[StreamResponse, None, None]:
+        """Handle iteration start events."""
+        self._ensure_workflow_initialized()
+
+        iter_start_resp = self._workflow_response_converter.workflow_iteration_start_to_stream_response(
+            task_id=self._application_generate_entity.task_id,
+            workflow_execution_id=self._workflow_run_id,
+            event=event,
+        )
+        yield iter_start_resp
+
+    def _handle_iteration_next_event(
+        self, event: QueueIterationNextEvent, **kwargs
+    ) -> Generator[StreamResponse, None, None]:
+        """Handle iteration next events."""
+        self._ensure_workflow_initialized()
+
+        iter_next_resp = self._workflow_response_converter.workflow_iteration_next_to_stream_response(
+            task_id=self._application_generate_entity.task_id,
+            workflow_execution_id=self._workflow_run_id,
+            event=event,
+        )
+        yield iter_next_resp
+
+    def _handle_iteration_completed_event(
+        self, event: QueueIterationCompletedEvent, **kwargs
+    ) -> Generator[StreamResponse, None, None]:
+        """Handle iteration completed events."""
+        self._ensure_workflow_initialized()
+
+        iter_finish_resp = self._workflow_response_converter.workflow_iteration_completed_to_stream_response(
+            task_id=self._application_generate_entity.task_id,
+            workflow_execution_id=self._workflow_run_id,
+            event=event,
+        )
+        yield iter_finish_resp
+
+    def _handle_loop_start_event(self, event: QueueLoopStartEvent, **kwargs) -> Generator[StreamResponse, None, None]:
+        """Handle loop start events."""
+        self._ensure_workflow_initialized()
+
+        loop_start_resp = self._workflow_response_converter.workflow_loop_start_to_stream_response(
+            task_id=self._application_generate_entity.task_id,
+            workflow_execution_id=self._workflow_run_id,
+            event=event,
+        )
+        yield loop_start_resp
+
+    def _handle_loop_next_event(self, event: QueueLoopNextEvent, **kwargs) -> Generator[StreamResponse, None, None]:
+        """Handle loop next events."""
+        self._ensure_workflow_initialized()
+
+        loop_next_resp = self._workflow_response_converter.workflow_loop_next_to_stream_response(
+            task_id=self._application_generate_entity.task_id,
+            workflow_execution_id=self._workflow_run_id,
+            event=event,
+        )
+        yield loop_next_resp
+
+    def _handle_loop_completed_event(
+        self, event: QueueLoopCompletedEvent, **kwargs
+    ) -> Generator[StreamResponse, None, None]:
+        """Handle loop completed events."""
+        self._ensure_workflow_initialized()
+
+        loop_finish_resp = self._workflow_response_converter.workflow_loop_completed_to_stream_response(
+            task_id=self._application_generate_entity.task_id,
+            workflow_execution_id=self._workflow_run_id,
+            event=event,
+        )
+        yield loop_finish_resp
+
+    def _handle_workflow_succeeded_event(
+        self,
+        event: QueueWorkflowSucceededEvent,
+        *,
+        graph_runtime_state: Optional[GraphRuntimeState] = None,
+        trace_manager: Optional[TraceQueueManager] = None,
+        **kwargs,
+    ) -> Generator[StreamResponse, None, None]:
+        """Handle workflow succeeded events."""
+        self._ensure_workflow_initialized()
+        validated_state = self._ensure_graph_runtime_initialized(graph_runtime_state)
+
+        with self._database_session() as session:
+            workflow_execution = self._workflow_cycle_manager.handle_workflow_run_success(
+                workflow_run_id=self._workflow_run_id,
+                total_tokens=validated_state.total_tokens,
+                total_steps=validated_state.node_run_steps,
+                outputs=event.outputs,
+                conversation_id=None,
+                trace_manager=trace_manager,
+                external_trace_id=self._application_generate_entity.extras.get("external_trace_id"),
+            )
+
+            # save workflow app log
+            self._save_workflow_app_log(session=session, workflow_execution=workflow_execution)
+
+            workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
+                session=session,
+                task_id=self._application_generate_entity.task_id,
+                workflow_execution=workflow_execution,
+            )
+
+        yield workflow_finish_resp
+
+    def _handle_workflow_partial_success_event(
+        self,
+        event: QueueWorkflowPartialSuccessEvent,
+        *,
+        graph_runtime_state: Optional[GraphRuntimeState] = None,
+        trace_manager: Optional[TraceQueueManager] = None,
+        **kwargs,
+    ) -> Generator[StreamResponse, None, None]:
+        """Handle workflow partial success events."""
+        self._ensure_workflow_initialized()
+        validated_state = self._ensure_graph_runtime_initialized(graph_runtime_state)
+
+        with self._database_session() as session:
+            workflow_execution = self._workflow_cycle_manager.handle_workflow_run_partial_success(
+                workflow_run_id=self._workflow_run_id,
+                total_tokens=validated_state.total_tokens,
+                total_steps=validated_state.node_run_steps,
+                outputs=event.outputs,
+                exceptions_count=event.exceptions_count,
+                conversation_id=None,
+                trace_manager=trace_manager,
+                external_trace_id=self._application_generate_entity.extras.get("external_trace_id"),
+            )
+
+            # save workflow app log
+            self._save_workflow_app_log(session=session, workflow_execution=workflow_execution)
+
+            workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
+                session=session,
+                task_id=self._application_generate_entity.task_id,
+                workflow_execution=workflow_execution,
+            )
+
+        yield workflow_finish_resp
+
+    def _handle_workflow_failed_and_stop_events(
+        self,
+        event: Union[QueueWorkflowFailedEvent, QueueStopEvent],
+        *,
+        graph_runtime_state: Optional[GraphRuntimeState] = None,
+        trace_manager: Optional[TraceQueueManager] = None,
+        **kwargs,
+    ) -> Generator[StreamResponse, None, None]:
+        """Handle workflow failed and stop events."""
+        self._ensure_workflow_initialized()
+        validated_state = self._ensure_graph_runtime_initialized(graph_runtime_state)
+
+        with self._database_session() as session:
+            workflow_execution = self._workflow_cycle_manager.handle_workflow_run_failed(
+                workflow_run_id=self._workflow_run_id,
+                total_tokens=validated_state.total_tokens,
+                total_steps=validated_state.node_run_steps,
+                status=WorkflowExecutionStatus.FAILED
+                if isinstance(event, QueueWorkflowFailedEvent)
+                else WorkflowExecutionStatus.STOPPED,
+                error_message=event.error if isinstance(event, QueueWorkflowFailedEvent) else event.get_stop_reason(),
+                conversation_id=None,
+                trace_manager=trace_manager,
+                exceptions_count=event.exceptions_count if isinstance(event, QueueWorkflowFailedEvent) else 0,
+                external_trace_id=self._application_generate_entity.extras.get("external_trace_id"),
+            )
+
+            # save workflow app log
+            self._save_workflow_app_log(session=session, workflow_execution=workflow_execution)
+
+            workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
+                session=session,
+                task_id=self._application_generate_entity.task_id,
+                workflow_execution=workflow_execution,
+            )
+
+        yield workflow_finish_resp
+
+    def _handle_text_chunk_event(
+        self,
+        event: QueueTextChunkEvent,
+        *,
+        tts_publisher: Optional[AppGeneratorTTSPublisher] = None,
+        queue_message: Optional[Union[WorkflowQueueMessage, MessageQueueMessage]] = None,
+        **kwargs,
+    ) -> Generator[StreamResponse, None, None]:
+        """Handle text chunk events."""
+        delta_text = event.text
+        if delta_text is None:
+            return
+
+        # only publish tts message at text chunk streaming
+        if tts_publisher and queue_message:
+            tts_publisher.publish(queue_message)
+
+        yield self._text_chunk_to_stream_response(delta_text, from_variable_selector=event.from_variable_selector)
+
+    def _handle_agent_log_event(self, event: QueueAgentLogEvent, **kwargs) -> Generator[StreamResponse, None, None]:
+        """Handle agent log events."""
+        yield self._workflow_response_converter.handle_agent_log(
+            task_id=self._application_generate_entity.task_id, event=event
+        )
+
+    def _get_event_handlers(self) -> dict[type, Callable]:
+        """Get mapping of event types to their handlers using fluent pattern."""
+        return {
+            # Basic events
+            QueuePingEvent: self._handle_ping_event,
+            QueueErrorEvent: self._handle_error_event,
+            QueueTextChunkEvent: self._handle_text_chunk_event,
+            # Workflow events
+            QueueWorkflowStartedEvent: self._handle_workflow_started_event,
+            QueueWorkflowSucceededEvent: self._handle_workflow_succeeded_event,
+            QueueWorkflowPartialSuccessEvent: self._handle_workflow_partial_success_event,
+            # Node events
+            QueueNodeRetryEvent: self._handle_node_retry_event,
+            QueueNodeStartedEvent: self._handle_node_started_event,
+            QueueNodeSucceededEvent: self._handle_node_succeeded_event,
+            # Parallel branch events
+            QueueParallelBranchRunStartedEvent: self._handle_parallel_branch_started_event,
+            # Iteration events
+            QueueIterationStartEvent: self._handle_iteration_start_event,
+            QueueIterationNextEvent: self._handle_iteration_next_event,
+            QueueIterationCompletedEvent: self._handle_iteration_completed_event,
+            # Loop events
+            QueueLoopStartEvent: self._handle_loop_start_event,
+            QueueLoopNextEvent: self._handle_loop_next_event,
+            QueueLoopCompletedEvent: self._handle_loop_completed_event,
+            # Agent events
+            QueueAgentLogEvent: self._handle_agent_log_event,
+        }
+
+    def _dispatch_event(
+        self,
+        event: Any,
+        *,
+        graph_runtime_state: Optional[GraphRuntimeState] = None,
+        tts_publisher: Optional[AppGeneratorTTSPublisher] = None,
+        trace_manager: Optional[TraceQueueManager] = None,
+        queue_message: Optional[Union[WorkflowQueueMessage, MessageQueueMessage]] = None,
+    ) -> Generator[StreamResponse, None, None]:
+        """Dispatch events using elegant pattern matching."""
+        handlers = self._get_event_handlers()
+        event_type = type(event)
+
+        # Direct handler lookup
+        if handler := handlers.get(event_type):
+            yield from handler(
+                event,
+                graph_runtime_state=graph_runtime_state,
+                tts_publisher=tts_publisher,
+                trace_manager=trace_manager,
+                queue_message=queue_message,
+            )
+            return
+
+        # Handle node failure events with isinstance check
+        if isinstance(
+            event,
+            (
+                QueueNodeFailedEvent,
+                QueueNodeInIterationFailedEvent,
+                QueueNodeInLoopFailedEvent,
+                QueueNodeExceptionEvent,
+            ),
+        ):
+            yield from self._handle_node_failed_events(
+                event,
+                graph_runtime_state=graph_runtime_state,
+                tts_publisher=tts_publisher,
+                trace_manager=trace_manager,
+                queue_message=queue_message,
+            )
+            return
+
+        # Handle parallel branch finished events with isinstance check
+        if isinstance(event, (QueueParallelBranchRunSucceededEvent, QueueParallelBranchRunFailedEvent)):
+            yield from self._handle_parallel_branch_finished_events(
+                event,
+                graph_runtime_state=graph_runtime_state,
+                tts_publisher=tts_publisher,
+                trace_manager=trace_manager,
+                queue_message=queue_message,
+            )
+            return
+
+        # Handle workflow failed and stop events with isinstance check
+        if isinstance(event, (QueueWorkflowFailedEvent, QueueStopEvent)):
+            yield from self._handle_workflow_failed_and_stop_events(
+                event,
+                graph_runtime_state=graph_runtime_state,
+                tts_publisher=tts_publisher,
+                trace_manager=trace_manager,
+                queue_message=queue_message,
+            )
+            return
+
+        # For unhandled events, we continue (original behavior)
+        return
+
     def _process_stream_response(
         self,
         tts_publisher: Optional[AppGeneratorTTSPublisher] = None,
         trace_manager: Optional[TraceQueueManager] = None,
     ) -> Generator[StreamResponse, None, None]:
         """
-        Process stream response.
-        :return:
+        Process stream response using elegant Fluent Python patterns.
+        Maintains exact same functionality as original 44-if-statement version.
         """
+        # Initialize graph runtime state
         graph_runtime_state = None
 
         for queue_message in self._base_task_pipeline._queue_manager.listen():
             event = queue_message.event
 
-            if isinstance(event, QueuePingEvent):
-                yield self._base_task_pipeline._ping_stream_response()
-            elif isinstance(event, QueueErrorEvent):
-                err = self._base_task_pipeline._handle_error(event=event)
-                yield self._base_task_pipeline._error_to_stream_response(err)
-                break
-            elif isinstance(event, QueueWorkflowStartedEvent):
-                # override graph runtime state
-                graph_runtime_state = event.graph_runtime_state
+            match event:
+                case QueueWorkflowStartedEvent():
+                    graph_runtime_state = event.graph_runtime_state
+                    yield from self._handle_workflow_started_event(event)
 
-                # init workflow run
-                workflow_execution = self._workflow_cycle_manager.handle_workflow_run_start()
-                self._workflow_run_id = workflow_execution.id_
-                start_resp = self._workflow_response_converter.workflow_start_to_stream_response(
-                    task_id=self._application_generate_entity.task_id,
-                    workflow_execution=workflow_execution,
-                )
-
-                yield start_resp
-            elif isinstance(
-                event,
-                QueueNodeRetryEvent,
-            ):
-                if not self._workflow_run_id:
-                    raise ValueError("workflow run not initialized.")
-                with Session(db.engine, expire_on_commit=False) as session:
-                    workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_retried(
-                        workflow_execution_id=self._workflow_run_id,
-                        event=event,
-                    )
-                    response = self._workflow_response_converter.workflow_node_retry_to_stream_response(
-                        event=event,
-                        task_id=self._application_generate_entity.task_id,
-                        workflow_node_execution=workflow_node_execution,
-                    )
-                    session.commit()
-
-                if response:
-                    yield response
-            elif isinstance(event, QueueNodeStartedEvent):
-                if not self._workflow_run_id:
-                    raise ValueError("workflow run not initialized.")
-
-                workflow_node_execution = self._workflow_cycle_manager.handle_node_execution_start(
-                    workflow_execution_id=self._workflow_run_id, event=event
-                )
-                node_start_response = self._workflow_response_converter.workflow_node_start_to_stream_response(
-                    event=event,
-                    task_id=self._application_generate_entity.task_id,
-                    workflow_node_execution=workflow_node_execution,
-                )
-
-                if node_start_response:
-                    yield node_start_response
-            elif isinstance(event, QueueNodeSucceededEvent):
-                workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_success(
-                    event=event
-                )
-                node_success_response = self._workflow_response_converter.workflow_node_finish_to_stream_response(
-                    event=event,
-                    task_id=self._application_generate_entity.task_id,
-                    workflow_node_execution=workflow_node_execution,
-                )
-
-                self._save_output_for_event(event, workflow_node_execution.id)
-
-                if node_success_response:
-                    yield node_success_response
-            elif isinstance(
-                event,
-                QueueNodeFailedEvent
-                | QueueNodeInIterationFailedEvent
-                | QueueNodeInLoopFailedEvent
-                | QueueNodeExceptionEvent,
-            ):
-                workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_failed(
-                    event=event,
-                )
-                node_failed_response = self._workflow_response_converter.workflow_node_finish_to_stream_response(
-                    event=event,
-                    task_id=self._application_generate_entity.task_id,
-                    workflow_node_execution=workflow_node_execution,
-                )
-                if isinstance(event, QueueNodeExceptionEvent):
-                    self._save_output_for_event(event, workflow_node_execution.id)
-
-                if node_failed_response:
-                    yield node_failed_response
-
-            elif isinstance(event, QueueParallelBranchRunStartedEvent):
-                if not self._workflow_run_id:
-                    raise ValueError("workflow run not initialized.")
-
-                parallel_start_resp = (
-                    self._workflow_response_converter.workflow_parallel_branch_start_to_stream_response(
-                        task_id=self._application_generate_entity.task_id,
-                        workflow_execution_id=self._workflow_run_id,
-                        event=event,
-                    )
-                )
-
-                yield parallel_start_resp
-
-            elif isinstance(event, QueueParallelBranchRunSucceededEvent | QueueParallelBranchRunFailedEvent):
-                if not self._workflow_run_id:
-                    raise ValueError("workflow run not initialized.")
-
-                parallel_finish_resp = (
-                    self._workflow_response_converter.workflow_parallel_branch_finished_to_stream_response(
-                        task_id=self._application_generate_entity.task_id,
-                        workflow_execution_id=self._workflow_run_id,
-                        event=event,
-                    )
-                )
-
-                yield parallel_finish_resp
-
-            elif isinstance(event, QueueIterationStartEvent):
-                if not self._workflow_run_id:
-                    raise ValueError("workflow run not initialized.")
-
-                iter_start_resp = self._workflow_response_converter.workflow_iteration_start_to_stream_response(
-                    task_id=self._application_generate_entity.task_id,
-                    workflow_execution_id=self._workflow_run_id,
-                    event=event,
-                )
-
-                yield iter_start_resp
-
-            elif isinstance(event, QueueIterationNextEvent):
-                if not self._workflow_run_id:
-                    raise ValueError("workflow run not initialized.")
-
-                iter_next_resp = self._workflow_response_converter.workflow_iteration_next_to_stream_response(
-                    task_id=self._application_generate_entity.task_id,
-                    workflow_execution_id=self._workflow_run_id,
-                    event=event,
-                )
-
-                yield iter_next_resp
-
-            elif isinstance(event, QueueIterationCompletedEvent):
-                if not self._workflow_run_id:
-                    raise ValueError("workflow run not initialized.")
-
-                iter_finish_resp = self._workflow_response_converter.workflow_iteration_completed_to_stream_response(
-                    task_id=self._application_generate_entity.task_id,
-                    workflow_execution_id=self._workflow_run_id,
-                    event=event,
-                )
-
-                yield iter_finish_resp
-
-            elif isinstance(event, QueueLoopStartEvent):
-                if not self._workflow_run_id:
-                    raise ValueError("workflow run not initialized.")
-
-                loop_start_resp = self._workflow_response_converter.workflow_loop_start_to_stream_response(
-                    task_id=self._application_generate_entity.task_id,
-                    workflow_execution_id=self._workflow_run_id,
-                    event=event,
-                )
-
-                yield loop_start_resp
-
-            elif isinstance(event, QueueLoopNextEvent):
-                if not self._workflow_run_id:
-                    raise ValueError("workflow run not initialized.")
-
-                loop_next_resp = self._workflow_response_converter.workflow_loop_next_to_stream_response(
-                    task_id=self._application_generate_entity.task_id,
-                    workflow_execution_id=self._workflow_run_id,
-                    event=event,
-                )
-
-                yield loop_next_resp
-
-            elif isinstance(event, QueueLoopCompletedEvent):
-                if not self._workflow_run_id:
-                    raise ValueError("workflow run not initialized.")
-
-                loop_finish_resp = self._workflow_response_converter.workflow_loop_completed_to_stream_response(
-                    task_id=self._application_generate_entity.task_id,
-                    workflow_execution_id=self._workflow_run_id,
-                    event=event,
-                )
-
-                yield loop_finish_resp
-
-            elif isinstance(event, QueueWorkflowSucceededEvent):
-                if not self._workflow_run_id:
-                    raise ValueError("workflow run not initialized.")
-                if not graph_runtime_state:
-                    raise ValueError("graph runtime state not initialized.")
-
-                with Session(db.engine, expire_on_commit=False) as session:
-                    workflow_execution = self._workflow_cycle_manager.handle_workflow_run_success(
-                        workflow_run_id=self._workflow_run_id,
-                        total_tokens=graph_runtime_state.total_tokens,
-                        total_steps=graph_runtime_state.node_run_steps,
-                        outputs=event.outputs,
-                        conversation_id=None,
-                        trace_manager=trace_manager,
+                case QueueTextChunkEvent():
+                    yield from self._handle_text_chunk_event(
+                        event, tts_publisher=tts_publisher, queue_message=queue_message
                     )
 
-                    # save workflow app log
-                    self._save_workflow_app_log(session=session, workflow_execution=workflow_execution)
+                case QueueErrorEvent():
+                    yield from self._handle_error_event(event)
+                    break
 
-                    workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
-                        session=session,
-                        task_id=self._application_generate_entity.task_id,
-                        workflow_execution=workflow_execution,
-                    )
-                    session.commit()
-
-                yield workflow_finish_resp
-            elif isinstance(event, QueueWorkflowPartialSuccessEvent):
-                if not self._workflow_run_id:
-                    raise ValueError("workflow run not initialized.")
-                if not graph_runtime_state:
-                    raise ValueError("graph runtime state not initialized.")
-
-                with Session(db.engine, expire_on_commit=False) as session:
-                    workflow_execution = self._workflow_cycle_manager.handle_workflow_run_partial_success(
-                        workflow_run_id=self._workflow_run_id,
-                        total_tokens=graph_runtime_state.total_tokens,
-                        total_steps=graph_runtime_state.node_run_steps,
-                        outputs=event.outputs,
-                        exceptions_count=event.exceptions_count,
-                        conversation_id=None,
-                        trace_manager=trace_manager,
-                    )
-
-                    # save workflow app log
-                    self._save_workflow_app_log(session=session, workflow_execution=workflow_execution)
-
-                    workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
-                        session=session,
-                        task_id=self._application_generate_entity.task_id,
-                        workflow_execution=workflow_execution,
-                    )
-                    session.commit()
-
-                yield workflow_finish_resp
-            elif isinstance(event, QueueWorkflowFailedEvent | QueueStopEvent):
-                if not self._workflow_run_id:
-                    raise ValueError("workflow run not initialized.")
-                if not graph_runtime_state:
-                    raise ValueError("graph runtime state not initialized.")
-
-                with Session(db.engine, expire_on_commit=False) as session:
-                    workflow_execution = self._workflow_cycle_manager.handle_workflow_run_failed(
-                        workflow_run_id=self._workflow_run_id,
-                        total_tokens=graph_runtime_state.total_tokens,
-                        total_steps=graph_runtime_state.node_run_steps,
-                        status=WorkflowExecutionStatus.FAILED
-                        if isinstance(event, QueueWorkflowFailedEvent)
-                        else WorkflowExecutionStatus.STOPPED,
-                        error_message=event.error
-                        if isinstance(event, QueueWorkflowFailedEvent)
-                        else event.get_stop_reason(),
-                        conversation_id=None,
-                        trace_manager=trace_manager,
-                        exceptions_count=event.exceptions_count if isinstance(event, QueueWorkflowFailedEvent) else 0,
-                    )
-
-                    # save workflow app log
-                    self._save_workflow_app_log(session=session, workflow_execution=workflow_execution)
-
-                    workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
-                        session=session,
-                        task_id=self._application_generate_entity.task_id,
-                        workflow_execution=workflow_execution,
-                    )
-                    session.commit()
-
-                yield workflow_finish_resp
-            elif isinstance(event, QueueTextChunkEvent):
-                delta_text = event.text
-                if delta_text is None:
-                    continue
-
-                # only publish tts message at text chunk streaming
-                if tts_publisher:
-                    tts_publisher.publish(queue_message)
-
-                yield self._text_chunk_to_stream_response(
-                    delta_text, from_variable_selector=event.from_variable_selector
-                )
-            elif isinstance(event, QueueAgentLogEvent):
-                yield self._workflow_response_converter.handle_agent_log(
-                    task_id=self._application_generate_entity.task_id, event=event
-                )
-            else:
-                continue
+                # Handle all other events through elegant dispatch
+                case _:
+                    if responses := list(
+                        self._dispatch_event(
+                            event,
+                            graph_runtime_state=graph_runtime_state,
+                            tts_publisher=tts_publisher,
+                            trace_manager=trace_manager,
+                            queue_message=queue_message,
+                        )
+                    ):
+                        yield from responses
 
         if tts_publisher:
             tts_publisher.publish(None)

api/core/app/apps/workflow_app_runner.py

@@ -1,8 +1,7 @@
 from collections.abc import Mapping
-from typing import Any, Optional, cast
+from typing import Any, cast
 
 from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
-from core.app.apps.base_app_runner import AppRunner
 from core.app.entities.queue_entities import (
     AppQueueEvent,
     QueueAgentLogEvent,
@@ -65,18 +64,20 @@ from core.workflow.nodes.node_mapping import NODE_TYPE_CLASSES_MAPPING
 from core.workflow.system_variable import SystemVariable
 from core.workflow.variable_loader import DUMMY_VARIABLE_LOADER, VariableLoader, load_into_variable_pool
 from core.workflow.workflow_entry import WorkflowEntry
-from extensions.ext_database import db
-from models.model import App
 from models.workflow import Workflow
 
 
-class WorkflowBasedAppRunner(AppRunner):
-    def __init__(self, queue_manager: AppQueueManager, variable_loader: VariableLoader = DUMMY_VARIABLE_LOADER) -> None:
-        self.queue_manager = queue_manager
+class WorkflowBasedAppRunner:
+    def __init__(
+        self,
+        *,
+        queue_manager: AppQueueManager,
+        variable_loader: VariableLoader = DUMMY_VARIABLE_LOADER,
+        app_id: str,
+    ) -> None:
+        self._queue_manager = queue_manager
         self._variable_loader = variable_loader
-
-    def _get_app_id(self) -> str:
-        raise NotImplementedError("not implemented")
+        self._app_id = app_id
 
     def _init_graph(self, graph_config: Mapping[str, Any]) -> Graph:
         """
@@ -693,21 +694,5 @@ class WorkflowBasedAppRunner(AppRunner):
                 )
             )
 
-    def get_workflow(self, app_model: App, workflow_id: str) -> Optional[Workflow]:
-        """
-        Get workflow
-        """
-        # fetch workflow by workflow_id
-        workflow = (
-            db.session.query(Workflow)
-            .filter(
-                Workflow.tenant_id == app_model.tenant_id, Workflow.app_id == app_model.id, Workflow.id == workflow_id
-            )
-            .first()
-        )
-
-        # return workflow
-        return workflow
-
     def _publish_event(self, event: AppQueueEvent) -> None:
-        self.queue_manager.publish(event, PublishFrom.APPLICATION_MANAGER)
+        self._queue_manager.publish(event, PublishFrom.APPLICATION_MANAGER)

api/core/file/file_manager.py

@@ -7,6 +7,7 @@ from core.model_runtime.entities import (
     AudioPromptMessageContent,
     DocumentPromptMessageContent,
     ImagePromptMessageContent,
+    TextPromptMessageContent,
     VideoPromptMessageContent,
 )
 from core.model_runtime.entities.message_entities import PromptMessageContentUnionTypes
@@ -44,11 +45,44 @@ def to_prompt_message_content(
     *,
     image_detail_config: ImagePromptMessageContent.DETAIL | None = None,
 ) -> PromptMessageContentUnionTypes:
+    """
+    Convert a file to prompt message content.
+
+    This function converts files to their appropriate prompt message content types.
+    For supported file types (IMAGE, AUDIO, VIDEO, DOCUMENT), it creates the
+    corresponding message content with proper encoding/URL.
+
+    For unsupported file types, instead of raising an error, it returns a
+    TextPromptMessageContent with a descriptive message about the file.
+
+    Args:
+        f: The file to convert
+        image_detail_config: Optional detail configuration for image files
+
+    Returns:
+        PromptMessageContentUnionTypes: The appropriate message content type
+
+    Raises:
+        ValueError: If file extension or mime_type is missing
+    """
     if f.extension is None:
         raise ValueError("Missing file extension")
     if f.mime_type is None:
         raise ValueError("Missing file mime_type")
 
+    prompt_class_map: Mapping[FileType, type[PromptMessageContentUnionTypes]] = {
+        FileType.IMAGE: ImagePromptMessageContent,
+        FileType.AUDIO: AudioPromptMessageContent,
+        FileType.VIDEO: VideoPromptMessageContent,
+        FileType.DOCUMENT: DocumentPromptMessageContent,
+    }
+
+    # Check if file type is supported
+    if f.type not in prompt_class_map:
+        # For unsupported file types, return a text description
+        return TextPromptMessageContent(data=f"[Unsupported file type: {f.filename} ({f.type.value})]")
+
+    # Process supported file types
     params = {
         "base64_data": _get_encoded_string(f) if dify_config.MULTIMODAL_SEND_FORMAT == "base64" else "",
         "url": _to_url(f) if dify_config.MULTIMODAL_SEND_FORMAT == "url" else "",
@@ -58,17 +92,7 @@ def to_prompt_message_content(
     if f.type == FileType.IMAGE:
         params["detail"] = image_detail_config or ImagePromptMessageContent.DETAIL.LOW
 
-    prompt_class_map: Mapping[FileType, type[PromptMessageContentUnionTypes]] = {
-        FileType.IMAGE: ImagePromptMessageContent,
-        FileType.AUDIO: AudioPromptMessageContent,
-        FileType.VIDEO: VideoPromptMessageContent,
-        FileType.DOCUMENT: DocumentPromptMessageContent,
-    }
-
-    try:
-        return prompt_class_map[f.type].model_validate(params)
-    except KeyError:
-        raise ValueError(f"file type {f.type} is not supported")
+    return prompt_class_map[f.type].model_validate(params)
 
 
 def download(f: File, /):

api/core/helper/encrypter.py

@@ -21,7 +21,7 @@ def encrypt_token(tenant_id: str, token: str):
     return base64.b64encode(encrypted_token).decode()
 
 
-def decrypt_token(tenant_id: str, token: str):
+def decrypt_token(tenant_id: str, token: str) -> str:
     return rsa.decrypt(base64.b64decode(token), tenant_id)
 
 

api/core/helper/marketplace.py

@@ -25,9 +25,29 @@ def batch_fetch_plugin_manifests(plugin_ids: list[str]) -> Sequence[MarketplaceP
     url = str(marketplace_api_url / "api/v1/plugins/batch")
     response = requests.post(url, json={"plugin_ids": plugin_ids})
     response.raise_for_status()
+
     return [MarketplacePluginDeclaration(**plugin) for plugin in response.json()["data"]["plugins"]]
 
 
+def batch_fetch_plugin_manifests_ignore_deserialization_error(
+    plugin_ids: list[str],
+) -> Sequence[MarketplacePluginDeclaration]:
+    if len(plugin_ids) == 0:
+        return []
+
+    url = str(marketplace_api_url / "api/v1/plugins/batch")
+    response = requests.post(url, json={"plugin_ids": plugin_ids})
+    response.raise_for_status()
+    result: list[MarketplacePluginDeclaration] = []
+    for plugin in response.json()["data"]["plugins"]:
+        try:
+            result.append(MarketplacePluginDeclaration(**plugin))
+        except Exception as e:
+            pass
+
+    return result
+
+
 def record_install_plugin_event(plugin_unique_identifier: str):
     url = str(marketplace_api_url / "api/v1/stats/plugins/install_count")
     response = requests.post(url, json={"unique_identifier": plugin_unique_identifier})

api/core/helper/trace_id_helper.py

@@ -0,0 +1,42 @@
+import re
+from collections.abc import Mapping
+from typing import Any, Optional
+
+
+def is_valid_trace_id(trace_id: str) -> bool:
+    """
+    Check if the trace_id is valid.
+
+    Requirements: 1-128 characters, only letters, numbers, '-', and '_'.
+    """
+    return bool(re.match(r"^[a-zA-Z0-9\-_]{1,128}$", trace_id))
+
+
+def get_external_trace_id(request: Any) -> Optional[str]:
+    """
+    Retrieve the trace_id from the request.
+
+    Priority: header ('X-Trace-Id'), then parameters, then JSON body. Returns None if not provided or invalid.
+    """
+    trace_id = request.headers.get("X-Trace-Id")
+    if not trace_id:
+        trace_id = request.args.get("trace_id")
+    if not trace_id and getattr(request, "is_json", False):
+        json_data = getattr(request, "json", None)
+        if json_data:
+            trace_id = json_data.get("trace_id")
+    if isinstance(trace_id, str) and is_valid_trace_id(trace_id):
+        return trace_id
+    return None
+
+
+def extract_external_trace_id_from_args(args: Mapping[str, Any]) -> dict:
+    """
+    Extract 'external_trace_id' from args.
+
+    Returns a dict suitable for use in extras. Returns an empty dict if not found.
+    """
+    trace_id = args.get("external_trace_id")
+    if trace_id:
+        return {"external_trace_id": trace_id}
+    return {}

api/core/indexing_runner.py

@@ -672,8 +672,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)
+        db.session.query(DatasetDocument).filter_by(id=document_id).update(update_params)  # type: ignore
         db.session.commit()
 
     @staticmethod

api/core/llm_generator/llm_generator.py

@@ -114,7 +114,8 @@ class LLMGenerator:
                 ),
             )
 
-            questions = output_parser.parse(cast(str, response.message.content))
+            text_content = response.message.get_text_content()
+            questions = output_parser.parse(text_content) if text_content else []
         except InvokeError:
             questions = []
         except Exception:

api/core/llm_generator/output_parser/suggested_questions_after_answer.py

@@ -15,5 +15,4 @@ class SuggestedQuestionsAfterAnswerOutputParser:
             json_obj = json.loads(action_match.group(0).strip())
         else:
             json_obj = []
-
         return json_obj

api/core/mcp/auth/auth_provider.py

@@ -8,7 +8,7 @@ from core.mcp.types import (
     OAuthTokens,
 )
 from models.tools import MCPToolProvider
-from services.tools.mcp_tools_mange_service import MCPToolManageService
+from services.tools.mcp_tools_manage_service import MCPToolManageService
 
 LATEST_PROTOCOL_VERSION = "1.0"
 

api/core/mcp/mcp_client.py

@@ -68,15 +68,17 @@ class MCPClient:
         }
 
         parsed_url = urlparse(self.server_url)
-        path = parsed_url.path
+        path = parsed_url.path or ""
         method_name = path.rstrip("/").split("/")[-1] if path else ""
-        try:
+        if method_name in connection_methods:
             client_factory = connection_methods[method_name]
             self.connect_server(client_factory, method_name)
-        except KeyError:
+        else:
             try:
+                logger.debug(f"Not supported method {method_name} found in URL path, trying default 'mcp' method.")
                 self.connect_server(sse_client, "sse")
             except MCPConnectionError:
+                logger.debug("MCP connection failed with 'sse', falling back to 'mcp' method.")
                 self.connect_server(streamablehttp_client, "mcp")
 
     def connect_server(
@@ -91,7 +93,7 @@ class MCPClient:
                 else {}
             )
             self._streams_context = client_factory(url=self.server_url, headers=headers)
-            if self._streams_context is None:
+            if not self._streams_context:
                 raise MCPConnectionError("Failed to create connection context")
 
             # Use exit_stack to manage context managers properly
@@ -141,10 +143,11 @@ class MCPClient:
         try:
             # ExitStack will handle proper cleanup of all managed context managers
             self.exit_stack.close()
+        except Exception as e:
+            logging.exception("Error during cleanup")
+            raise ValueError(f"Error during cleanup: {e}")
+        finally:
             self._session = None
             self._session_context = None
             self._streams_context = None
             self._initialized = False
-        except Exception as e:
-            logging.exception("Error during cleanup")
-            raise ValueError(f"Error during cleanup: {e}")

api/core/model_runtime/entities/message_entities.py

@@ -156,6 +156,23 @@ class PromptMessage(ABC, BaseModel):
         """
         return not self.content
 
+    def get_text_content(self) -> str:
+        """
+        Get text content from prompt message.
+
+        :return: Text content as string, empty string if no text content
+        """
+        if isinstance(self.content, str):
+            return self.content
+        elif isinstance(self.content, list):
+            text_parts = []
+            for item in self.content:
+                if isinstance(item, TextPromptMessageContent):
+                    text_parts.append(item.data)
+            return "".join(text_parts)
+        else:
+            return ""
+
     @field_validator("content", mode="before")
     @classmethod
     def validate_content(cls, v):

api/core/ops/aliyun_trace/aliyun_trace.py

@@ -101,7 +101,8 @@ class AliyunDataTrace(BaseTraceInstance):
             raise ValueError(f"Aliyun get run url failed: {str(e)}")
 
     def workflow_trace(self, trace_info: WorkflowTraceInfo):
-        trace_id = convert_to_trace_id(trace_info.workflow_run_id)
+        external_trace_id = trace_info.metadata.get("external_trace_id")
+        trace_id = external_trace_id or convert_to_trace_id(trace_info.workflow_run_id)
         workflow_span_id = convert_to_span_id(trace_info.workflow_run_id, "workflow")
         self.add_workflow_span(trace_id, workflow_span_id, trace_info)
 

api/core/ops/arize_phoenix_trace/arize_phoenix_trace.py

@@ -3,7 +3,7 @@ import json
 import logging
 import os
 from datetime import datetime, timedelta
-from typing import Optional, Union, cast
+from typing import Any, Optional, Union, cast
 
 from openinference.semconv.trace import OpenInferenceSpanKindValues, SpanAttributes
 from opentelemetry import trace
@@ -142,11 +142,8 @@ class ArizePhoenixDataTrace(BaseTraceInstance):
             raise
 
     def workflow_trace(self, trace_info: WorkflowTraceInfo):
-        if trace_info.message_data is None:
-            return
-
         workflow_metadata = {
-            "workflow_id": trace_info.workflow_run_id or "",
+            "workflow_run_id": trace_info.workflow_run_id or "",
             "message_id": trace_info.message_id or "",
             "workflow_app_log_id": trace_info.workflow_app_log_id or "",
             "status": trace_info.workflow_run_status or "",
@@ -156,7 +153,8 @@ class ArizePhoenixDataTrace(BaseTraceInstance):
         }
         workflow_metadata.update(trace_info.metadata)
 
-        trace_id = uuid_to_trace_id(trace_info.message_id)
+        external_trace_id = trace_info.metadata.get("external_trace_id")
+        trace_id = external_trace_id or uuid_to_trace_id(trace_info.workflow_run_id)
         span_id = RandomIdGenerator().generate_span_id()
         context = SpanContext(
             trace_id=trace_id,
@@ -213,7 +211,7 @@ class ArizePhoenixDataTrace(BaseTraceInstance):
                     if model:
                         node_metadata["ls_model_name"] = model
 
-                    outputs = json.loads(node_execution.outputs).get("usage", {})
+                    outputs = json.loads(node_execution.outputs).get("usage", {}) if "outputs" in node_execution else {}
                     usage_data = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {})
                     if usage_data:
                         node_metadata["total_tokens"] = usage_data.get("total_tokens", 0)
@@ -236,31 +234,34 @@ class ArizePhoenixDataTrace(BaseTraceInstance):
                         SpanAttributes.SESSION_ID: trace_info.conversation_id or "",
                     },
                     start_time=datetime_to_nanos(created_at),
+                    context=trace.set_span_in_context(trace.NonRecordingSpan(context)),
                 )
 
                 try:
                     if node_execution.node_type == "llm":
+                        llm_attributes: dict[str, Any] = {
+                            SpanAttributes.INPUT_VALUE: json.dumps(process_data.get("prompts", []), ensure_ascii=False),
+                        }
                         provider = process_data.get("model_provider")
                         model = process_data.get("model_name")
                         if provider:
-                            node_span.set_attribute(SpanAttributes.LLM_PROVIDER, provider)
+                            llm_attributes[SpanAttributes.LLM_PROVIDER] = provider
                         if model:
-                            node_span.set_attribute(SpanAttributes.LLM_MODEL_NAME, model)
-
-                        outputs = json.loads(node_execution.outputs).get("usage", {})
+                            llm_attributes[SpanAttributes.LLM_MODEL_NAME] = model
+                        outputs = (
+                            json.loads(node_execution.outputs).get("usage", {}) if "outputs" in node_execution else {}
+                        )
                         usage_data = (
                             process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {})
                         )
                         if usage_data:
-                            node_span.set_attribute(
-                                SpanAttributes.LLM_TOKEN_COUNT_TOTAL, usage_data.get("total_tokens", 0)
-                            )
-                            node_span.set_attribute(
-                                SpanAttributes.LLM_TOKEN_COUNT_PROMPT, usage_data.get("prompt_tokens", 0)
-                            )
-                            node_span.set_attribute(
-                                SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, usage_data.get("completion_tokens", 0)
+                            llm_attributes[SpanAttributes.LLM_TOKEN_COUNT_TOTAL] = usage_data.get("total_tokens", 0)
+                            llm_attributes[SpanAttributes.LLM_TOKEN_COUNT_PROMPT] = usage_data.get("prompt_tokens", 0)
+                            llm_attributes[SpanAttributes.LLM_TOKEN_COUNT_COMPLETION] = usage_data.get(
+                                "completion_tokens", 0
                             )
+                        llm_attributes.update(self._construct_llm_attributes(process_data.get("prompts", [])))
+                        node_span.set_attributes(llm_attributes)
                 finally:
                     node_span.end(end_time=datetime_to_nanos(finished_at))
         finally:
@@ -352,25 +353,7 @@ class ArizePhoenixDataTrace(BaseTraceInstance):
                 SpanAttributes.METADATA: json.dumps(message_metadata, ensure_ascii=False),
                 SpanAttributes.SESSION_ID: trace_info.message_data.conversation_id,
             }
-
-            if isinstance(trace_info.inputs, list):
-                for i, msg in enumerate(trace_info.inputs):
-                    if isinstance(msg, dict):
-                        llm_attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.{i}.message.content"] = msg.get("text", "")
-                        llm_attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.{i}.message.role"] = msg.get(
-                            "role", "user"
-                        )
-                        # todo: handle assistant and tool role messages, as they don't always
-                        # have a text field, but may have a tool_calls field instead
-                        # e.g. 'tool_calls': [{'id': '98af3a29-b066-45a5-b4b1-46c74ddafc58',
-                        # 'type': 'function', 'function': {'name': 'current_time', 'arguments': '{}'}}]}
-            elif isinstance(trace_info.inputs, dict):
-                llm_attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.message.content"] = json.dumps(trace_info.inputs)
-                llm_attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.message.role"] = "user"
-            elif isinstance(trace_info.inputs, str):
-                llm_attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.message.content"] = trace_info.inputs
-                llm_attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.message.role"] = "user"
-
+            llm_attributes.update(self._construct_llm_attributes(trace_info.inputs))
             if trace_info.total_tokens is not None and trace_info.total_tokens > 0:
                 llm_attributes[SpanAttributes.LLM_TOKEN_COUNT_TOTAL] = trace_info.total_tokens
             if trace_info.message_tokens is not None and trace_info.message_tokens > 0:
@@ -724,3 +707,24 @@ class ArizePhoenixDataTrace(BaseTraceInstance):
             .all()
         )
         return workflow_nodes
+
+    def _construct_llm_attributes(self, prompts: dict | list | str | None) -> dict[str, str]:
+        """Helper method to construct LLM attributes with passed prompts."""
+        attributes = {}
+        if isinstance(prompts, list):
+            for i, msg in enumerate(prompts):
+                if isinstance(msg, dict):
+                    attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.{i}.message.content"] = msg.get("text", "")
+                    attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.{i}.message.role"] = msg.get("role", "user")
+                    # todo: handle assistant and tool role messages, as they don't always
+                    # have a text field, but may have a tool_calls field instead
+                    # e.g. 'tool_calls': [{'id': '98af3a29-b066-45a5-b4b1-46c74ddafc58',
+                    # 'type': 'function', 'function': {'name': 'current_time', 'arguments': '{}'}}]}
+        elif isinstance(prompts, dict):
+            attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.message.content"] = json.dumps(prompts)
+            attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.message.role"] = "user"
+        elif isinstance(prompts, str):
+            attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.message.content"] = prompts
+            attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.message.role"] = "user"
+
+        return attributes

api/core/ops/langfuse_trace/langfuse_trace.py

@@ -67,13 +67,14 @@ class LangFuseDataTrace(BaseTraceInstance):
             self.generate_name_trace(trace_info)
 
     def workflow_trace(self, trace_info: WorkflowTraceInfo):
-        trace_id = trace_info.workflow_run_id
+        external_trace_id = trace_info.metadata.get("external_trace_id")
+        trace_id = external_trace_id or trace_info.workflow_run_id
         user_id = trace_info.metadata.get("user_id")
         metadata = trace_info.metadata
         metadata["workflow_app_log_id"] = trace_info.workflow_app_log_id
 
         if trace_info.message_id:
-            trace_id = trace_info.message_id
+            trace_id = external_trace_id or trace_info.message_id
             name = TraceTaskName.MESSAGE_TRACE.value
             trace_data = LangfuseTrace(
                 id=trace_id,

api/core/ops/langsmith_trace/langsmith_trace.py

@@ -65,7 +65,8 @@ class LangSmithDataTrace(BaseTraceInstance):
             self.generate_name_trace(trace_info)
 
     def workflow_trace(self, trace_info: WorkflowTraceInfo):
-        trace_id = trace_info.message_id or trace_info.workflow_run_id
+        external_trace_id = trace_info.metadata.get("external_trace_id")
+        trace_id = external_trace_id or trace_info.message_id or trace_info.workflow_run_id
         if trace_info.start_time is None:
             trace_info.start_time = datetime.now()
         message_dotted_order = (

api/core/ops/opik_trace/opik_trace.py

@@ -96,7 +96,8 @@ class OpikDataTrace(BaseTraceInstance):
             self.generate_name_trace(trace_info)
 
     def workflow_trace(self, trace_info: WorkflowTraceInfo):
-        dify_trace_id = trace_info.workflow_run_id
+        external_trace_id = trace_info.metadata.get("external_trace_id")
+        dify_trace_id = external_trace_id or trace_info.workflow_run_id
         opik_trace_id = prepare_opik_uuid(trace_info.start_time, dify_trace_id)
         workflow_metadata = wrap_metadata(
             trace_info.metadata, message_id=trace_info.message_id, workflow_app_log_id=trace_info.workflow_app_log_id
@@ -104,7 +105,7 @@ class OpikDataTrace(BaseTraceInstance):
         root_span_id = None
 
         if trace_info.message_id:
-            dify_trace_id = trace_info.message_id
+            dify_trace_id = external_trace_id or trace_info.message_id
             opik_trace_id = prepare_opik_uuid(trace_info.start_time, dify_trace_id)
 
             trace_data = {

api/core/ops/ops_trace_manager.py

@@ -520,6 +520,10 @@ class TraceTask:
                 "app_id": workflow_run.app_id,
             }
 
+            external_trace_id = self.kwargs.get("external_trace_id")
+            if external_trace_id:
+                metadata["external_trace_id"] = external_trace_id
+
             workflow_trace_info = WorkflowTraceInfo(
                 workflow_data=workflow_run.to_dict(),
                 conversation_id=conversation_id,

api/core/ops/weave_trace/weave_trace.py

@@ -87,7 +87,8 @@ class WeaveDataTrace(BaseTraceInstance):
             self.generate_name_trace(trace_info)
 
     def workflow_trace(self, trace_info: WorkflowTraceInfo):
-        trace_id = trace_info.message_id or trace_info.workflow_run_id
+        external_trace_id = trace_info.metadata.get("external_trace_id")
+        trace_id = external_trace_id or trace_info.message_id or trace_info.workflow_run_id
         if trace_info.start_time is None:
             trace_info.start_time = datetime.now()
 

api/core/plugin/entities/marketplace.py

@@ -32,6 +32,13 @@ class MarketplacePluginDeclaration(BaseModel):
     latest_package_identifier: str = Field(
         ..., description="Unique identifier for the latest package release of the plugin"
     )
+    status: str = Field(..., description="Indicate the status of marketplace plugin, enum from `active` `deleted`")
+    deprecated_reason: str = Field(
+        ..., description="Not empty when status='deleted', indicates the reason why this plugin is deleted(deprecated)"
+    )
+    alternative_plugin_id: str = Field(
+        ..., description="Optional, indicates the alternative plugin for user to switch to"
+    )
 
     @model_validator(mode="before")
     @classmethod

api/core/plugin/entities/plugin_daemon.py

@@ -182,6 +182,10 @@ class PluginOAuthAuthorizationUrlResponse(BaseModel):
 
 
 class PluginOAuthCredentialsResponse(BaseModel):
+    metadata: Mapping[str, Any] = Field(
+        default_factory=dict, description="The metadata of the OAuth, like avatar url, name, etc."
+    )
+    expires_at: int = Field(default=-1, description="The expires at time of the credentials. UTC timestamp.")
     credentials: Mapping[str, Any] = Field(description="The credentials of the OAuth.")
 
 

api/core/plugin/impl/oauth.py

@@ -84,6 +84,41 @@ class OAuthHandler(BasePluginClient):
         except Exception as e:
             raise ValueError(f"Error getting credentials: {e}")
 
+    def refresh_credentials(
+        self,
+        tenant_id: str,
+        user_id: str,
+        plugin_id: str,
+        provider: str,
+        redirect_uri: str,
+        system_credentials: Mapping[str, Any],
+        credentials: Mapping[str, Any],
+    ) -> PluginOAuthCredentialsResponse:
+        try:
+            response = self._request_with_plugin_daemon_response_stream(
+                "POST",
+                f"plugin/{tenant_id}/dispatch/oauth/refresh_credentials",
+                PluginOAuthCredentialsResponse,
+                data={
+                    "user_id": user_id,
+                    "data": {
+                        "provider": provider,
+                        "redirect_uri": redirect_uri,
+                        "system_credentials": system_credentials,
+                        "credentials": credentials,
+                    },
+                },
+                headers={
+                    "X-Plugin-ID": plugin_id,
+                    "Content-Type": "application/json",
+                },
+            )
+            for resp in response:
+                return resp
+            raise ValueError("No response received from plugin daemon for refresh credentials request.")
+        except Exception as e:
+            raise ValueError(f"Error refreshing credentials: {e}")
+
     def _convert_request_to_raw_data(self, request: Request) -> bytes:
         """
         Convert a Request object to raw HTTP data.

api/core/prompt/simple_prompt_transform.py

@@ -29,19 +29,6 @@ class ModelMode(enum.StrEnum):
     COMPLETION = "completion"
     CHAT = "chat"
 
-    @classmethod
-    def value_of(cls, value: str) -> "ModelMode":
-        """
-        Get value of given mode.
-
-        :param value: mode value
-        :return: mode
-        """
-        for mode in cls:
-            if mode.value == value:
-                return mode
-        raise ValueError(f"invalid mode value {value}")
-
 
 prompt_file_contents: dict[str, Any] = {}
 
@@ -65,7 +52,7 @@ class SimplePromptTransform(PromptTransform):
     ) -> tuple[list[PromptMessage], Optional[list[str]]]:
         inputs = {key: str(value) for key, value in inputs.items()}
 
-        model_mode = ModelMode.value_of(model_config.mode)
+        model_mode = ModelMode(model_config.mode)
         if model_mode == ModelMode.CHAT:
             prompt_messages, stops = self._get_chat_model_prompt_messages(
                 app_mode=app_mode,

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

@@ -233,6 +233,12 @@ class AnalyticdbVectorOpenAPI:
     def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]:
         from alibabacloud_gpdb20160503 import models as gpdb_20160503_models
 
+        document_ids_filter = kwargs.get("document_ids_filter")
+        where_clause = ""
+        if document_ids_filter:
+            document_ids = ", ".join(f"'{id}'" for id in document_ids_filter)
+            where_clause += f"metadata_->>'document_id' IN ({document_ids})"
+
         score_threshold = kwargs.get("score_threshold") or 0.0
         request = gpdb_20160503_models.QueryCollectionDataRequest(
             dbinstance_id=self.config.instance_id,
@@ -245,7 +251,7 @@ class AnalyticdbVectorOpenAPI:
             vector=query_vector,
             content=None,
             top_k=kwargs.get("top_k", 4),
-            filter=None,
+            filter=where_clause,
         )
         response = self._client.query_collection_data(request)
         documents = []
@@ -265,6 +271,11 @@ class AnalyticdbVectorOpenAPI:
     def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
         from alibabacloud_gpdb20160503 import models as gpdb_20160503_models
 
+        document_ids_filter = kwargs.get("document_ids_filter")
+        where_clause = ""
+        if document_ids_filter:
+            document_ids = ", ".join(f"'{id}'" for id in document_ids_filter)
+            where_clause += f"metadata_->>'document_id' IN ({document_ids})"
         score_threshold = float(kwargs.get("score_threshold") or 0.0)
         request = gpdb_20160503_models.QueryCollectionDataRequest(
             dbinstance_id=self.config.instance_id,
@@ -277,7 +288,7 @@ class AnalyticdbVectorOpenAPI:
             vector=None,
             content=query,
             top_k=kwargs.get("top_k", 4),
-            filter=None,
+            filter=where_clause,
         )
         response = self._client.query_collection_data(request)
         documents = []

api/core/rag/datasource/vdb/elasticsearch/elasticsearch_vector.py

@@ -147,10 +147,17 @@ class ElasticSearchVector(BaseVector):
         return docs
 
     def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
-        query_str = {"match": {Field.CONTENT_KEY.value: query}}
+        query_str: dict[str, Any] = {"match": {Field.CONTENT_KEY.value: query}}
         document_ids_filter = kwargs.get("document_ids_filter")
+
         if document_ids_filter:
-            query_str["filter"] = {"terms": {"metadata.document_id": document_ids_filter}}  # type: ignore
+            query_str = {
+                "bool": {
+                    "must": {"match": {Field.CONTENT_KEY.value: query}},
+                    "filter": {"terms": {"metadata.document_id": document_ids_filter}},
+                }
+            }
+
         results = self._client.search(index=self._collection_name, query=query_str, size=kwargs.get("top_k", 4))
         docs = []
         for hit in results["hits"]["hits"]:

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

@@ -6,7 +6,7 @@ from uuid import UUID, uuid4
 from numpy import ndarray
 from pgvecto_rs.sqlalchemy import VECTOR  # type: ignore
 from pydantic import BaseModel, model_validator
-from sqlalchemy import Float, String, create_engine, insert, select, text
+from sqlalchemy import Float, create_engine, insert, select, text
 from sqlalchemy import text as sql_text
 from sqlalchemy.dialects import postgresql
 from sqlalchemy.orm import Mapped, Session, mapped_column
@@ -67,7 +67,7 @@ class PGVectoRS(BaseVector):
                 postgresql.UUID(as_uuid=True),
                 primary_key=True,
             )
-            text: Mapped[str] = mapped_column(String)
+            text: Mapped[str]
             meta: Mapped[dict] = mapped_column(postgresql.JSONB)
             vector: Mapped[ndarray] = mapped_column(VECTOR(dim))
 

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

@@ -118,10 +118,21 @@ class TableStoreVector(BaseVector):
 
     def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]:
         top_k = kwargs.get("top_k", 4)
-        return self._search_by_vector(query_vector, top_k)
+        document_ids_filter = kwargs.get("document_ids_filter")
+        filtered_list = None
+        if document_ids_filter:
+            filtered_list = ["document_id=" + item for item in document_ids_filter]
+        score_threshold = float(kwargs.get("score_threshold") or 0.0)
+        return self._search_by_vector(query_vector, filtered_list, top_k, score_threshold)
 
     def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
-        return self._search_by_full_text(query)
+        top_k = kwargs.get("top_k", 4)
+        document_ids_filter = kwargs.get("document_ids_filter")
+        filtered_list = None
+        if document_ids_filter:
+            filtered_list = ["document_id=" + item for item in document_ids_filter]
+
+        return self._search_by_full_text(query, filtered_list, top_k)
 
     def delete(self) -> None:
         self._delete_table_if_exist()
@@ -230,32 +241,51 @@ class TableStoreVector(BaseVector):
         primary_key = [("id", id)]
         row = tablestore.Row(primary_key)
         self._tablestore_client.delete_row(self._table_name, row, None)
-        logging.info("Tablestore delete row successfully. id:%s", id)
 
     def _search_by_metadata(self, key: str, value: str) -> list[str]:
         query = tablestore.SearchQuery(
             tablestore.TermQuery(self._tags_field, str(key) + "=" + str(value)),
-            limit=100,
+            limit=1000,
             get_total_count=False,
         )
+        rows: list[str] = []
+        next_token = None
+        while True:
+            if next_token is not None:
+                query.next_token = next_token
 
-        search_response = self._tablestore_client.search(
-            table_name=self._table_name,
-            index_name=self._index_name,
-            search_query=query,
-            columns_to_get=tablestore.ColumnsToGet(return_type=tablestore.ColumnReturnType.ALL_FROM_INDEX),
-        )
+            search_response = self._tablestore_client.search(
+                table_name=self._table_name,
+                index_name=self._index_name,
+                search_query=query,
+                columns_to_get=tablestore.ColumnsToGet(
+                    column_names=[Field.PRIMARY_KEY.value], return_type=tablestore.ColumnReturnType.SPECIFIED
+                ),
+            )
 
-        return [row[0][0][1] for row in search_response.rows]
+            if search_response is not None:
+                rows.extend([row[0][0][1] for row in search_response.rows])
 
-    def _search_by_vector(self, query_vector: list[float], top_k: int) -> list[Document]:
-        ots_query = tablestore.KnnVectorQuery(
+            if search_response is None or search_response.next_token == b"":
+                break
+            else:
+                next_token = search_response.next_token
+
+        return rows
+
+    def _search_by_vector(
+        self, query_vector: list[float], document_ids_filter: list[str] | None, top_k: int, score_threshold: float
+    ) -> list[Document]:
+        knn_vector_query = tablestore.KnnVectorQuery(
             field_name=Field.VECTOR.value,
             top_k=top_k,
             float32_query_vector=query_vector,
         )
+        if document_ids_filter:
+            knn_vector_query.filter = tablestore.TermsQuery(self._tags_field, document_ids_filter)
+
         sort = tablestore.Sort(sorters=[tablestore.ScoreSort(sort_order=tablestore.SortOrder.DESC)])
-        search_query = tablestore.SearchQuery(ots_query, limit=top_k, get_total_count=False, sort=sort)
+        search_query = tablestore.SearchQuery(knn_vector_query, limit=top_k, get_total_count=False, sort=sort)
 
         search_response = self._tablestore_client.search(
             table_name=self._table_name,
@@ -263,30 +293,32 @@ class TableStoreVector(BaseVector):
             search_query=search_query,
             columns_to_get=tablestore.ColumnsToGet(return_type=tablestore.ColumnReturnType.ALL_FROM_INDEX),
         )
-        logging.info(
-            "Tablestore search successfully. request_id:%s",
-            search_response.request_id,
-        )
-        return self._to_query_result(search_response)
-
-    def _to_query_result(self, search_response: tablestore.SearchResponse) -> list[Document]:
         documents = []
-        for row in search_response.rows:
-            documents.append(
-                Document(
-                    page_content=row[1][2][1],
-                    vector=json.loads(row[1][3][1]),
-                    metadata=json.loads(row[1][0][1]),
+        for search_hit in search_response.search_hits:
+            if search_hit.score > score_threshold:
+                metadata = json.loads(search_hit.row[1][0][1])
+                metadata["score"] = search_hit.score
+                documents.append(
+                    Document(
+                        page_content=search_hit.row[1][2][1],
+                        vector=json.loads(search_hit.row[1][3][1]),
+                        metadata=metadata,
+                    )
                 )
-            )
-
+        documents = sorted(documents, key=lambda x: x.metadata["score"] if x.metadata else 0, reverse=True)
         return documents
 
-    def _search_by_full_text(self, query: str) -> list[Document]:
+    def _search_by_full_text(self, query: str, document_ids_filter: list[str] | None, top_k: int) -> list[Document]:
+        bool_query = tablestore.BoolQuery()
+        bool_query.must_queries.append(tablestore.MatchQuery(text=query, field_name=Field.CONTENT_KEY.value))
+
+        if document_ids_filter:
+            bool_query.filter_queries.append(tablestore.TermsQuery(self._tags_field, document_ids_filter))
+
         search_query = tablestore.SearchQuery(
-            query=tablestore.MatchQuery(text=query, field_name=Field.CONTENT_KEY.value),
+            query=bool_query,
             sort=tablestore.Sort(sorters=[tablestore.ScoreSort(sort_order=tablestore.SortOrder.DESC)]),
-            limit=100,
+            limit=top_k,
         )
         search_response = self._tablestore_client.search(
             table_name=self._table_name,
@@ -295,7 +327,16 @@ class TableStoreVector(BaseVector):
             columns_to_get=tablestore.ColumnsToGet(return_type=tablestore.ColumnReturnType.ALL_FROM_INDEX),
         )
 
-        return self._to_query_result(search_response)
+        documents = []
+        for search_hit in search_response.search_hits:
+            documents.append(
+                Document(
+                    page_content=search_hit.row[1][2][1],
+                    vector=json.loads(search_hit.row[1][3][1]),
+                    metadata=json.loads(search_hit.row[1][0][1]),
+                )
+            )
+        return documents
 
 
 class TableStoreVectorFactory(AbstractVectorFactory):

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

@@ -206,9 +206,19 @@ class TencentVector(BaseVector):
     def delete_by_ids(self, ids: list[str]) -> None:
         if not ids:
             return
-        self._client.delete(
-            database_name=self._client_config.database, collection_name=self.collection_name, document_ids=ids
-        )
+
+        total_count = len(ids)
+        batch_size = self._client_config.max_upsert_batch_size
+        batch = math.ceil(total_count / batch_size)
+
+        for j in range(batch):
+            start_idx = j * batch_size
+            end_idx = min(total_count, (j + 1) * batch_size)
+            batch_ids = ids[start_idx:end_idx]
+
+            self._client.delete(
+                database_name=self._client_config.database, collection_name=self.collection_name, document_ids=batch_ids
+            )
 
     def delete_by_metadata_field(self, key: str, value: str) -> None:
         self._client.delete(
@@ -274,7 +284,8 @@ class TencentVector(BaseVector):
                 # Compatible with version 1.1.3 and below.
                 meta = json.loads(meta)
                 score = 1 - result.get("score", 0.0)
-            score = result.get("score", 0.0)
+            else:
+                score = result.get("score", 0.0)
             if score > score_threshold:
                 meta["score"] = score
                 doc = Document(page_content=result.get(self.field_text), metadata=meta)

api/core/rag/extractor/notion_extractor.py

@@ -331,9 +331,10 @@ class NotionExtractor(BaseExtractor):
         last_edited_time = self.get_notion_last_edited_time()
         data_source_info = document_model.data_source_info_dict
         data_source_info["last_edited_time"] = last_edited_time
-        update_params = {DocumentModel.data_source_info: json.dumps(data_source_info)}
 
-        db.session.query(DocumentModel).filter_by(id=document_model.id).update(update_params)
+        db.session.query(DocumentModel).filter_by(id=document_model.id).update(
+            {DocumentModel.data_source_info: json.dumps(data_source_info)}
+        )  # type: ignore
         db.session.commit()
 
     def get_notion_last_edited_time(self) -> str:

api/core/rag/extractor/word_extractor.py

@@ -238,9 +238,11 @@ class WordExtractor(BaseExtractor):
             paragraph_content = []
             for run in paragraph.runs:
                 if hasattr(run.element, "tag") and isinstance(run.element.tag, str) and run.element.tag.endswith("r"):
+                    # Process drawing type images
                     drawing_elements = run.element.findall(
                         ".//{http://schemas.openxmlformats.org/wordprocessingml/2006/main}drawing"
                     )
+                    has_drawing = False
                     for drawing in drawing_elements:
                         blip_elements = drawing.findall(
                             ".//{http://schemas.openxmlformats.org/drawingml/2006/main}blip"
@@ -252,6 +254,34 @@ class WordExtractor(BaseExtractor):
                             if embed_id:
                                 image_part = doc.part.related_parts.get(embed_id)
                                 if image_part in image_map:
+                                    has_drawing = True
+                                    paragraph_content.append(image_map[image_part])
+                    # Process pict type images
+                    shape_elements = run.element.findall(
+                        ".//{http://schemas.openxmlformats.org/wordprocessingml/2006/main}pict"
+                    )
+                    for shape in shape_elements:
+                        # Find image data in VML
+                        shape_image = shape.find(
+                            ".//{http://schemas.openxmlformats.org/wordprocessingml/2006/main}binData"
+                        )
+                        if shape_image is not None and shape_image.text:
+                            image_id = shape_image.get(
+                                "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}id"
+                            )
+                            if image_id and image_id in doc.part.rels:
+                                image_part = doc.part.rels[image_id].target_part
+                                if image_part in image_map and not has_drawing:
+                                    paragraph_content.append(image_map[image_part])
+                        # Find imagedata element in VML
+                        image_data = shape.find(".//{urn:schemas-microsoft-com:vml}imagedata")
+                        if image_data is not None:
+                            image_id = image_data.get("id") or image_data.get(
+                                "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}id"
+                            )
+                            if image_id and image_id in doc.part.rels:
+                                image_part = doc.part.rels[image_id].target_part
+                                if image_part in image_map and not has_drawing:
                                     paragraph_content.append(image_map[image_part])
                 if run.text.strip():
                     paragraph_content.append(run.text.strip())

api/core/rag/retrieval/dataset_retrieval.py

@@ -1137,7 +1137,7 @@ class DatasetRetrieval:
     def _get_prompt_template(
         self, model_config: ModelConfigWithCredentialsEntity, mode: str, metadata_fields: list, query: str
     ):
-        model_mode = ModelMode.value_of(mode)
+        model_mode = ModelMode(mode)
         input_text = query
 
         prompt_template: Union[CompletionModelPromptTemplate, list[ChatModelMessage]]

api/core/rag/splitter/fixed_text_splitter.py

@@ -102,6 +102,7 @@ class FixedRecursiveCharacterTextSplitter(EnhanceRecursiveCharacterTextSplitter)
                 splits = text.split()
             else:
                 splits = text.split(separator)
+                splits = [item + separator if i < len(splits) else item for i, item in enumerate(splits)]
         else:
             splits = list(text)
         splits = [s for s in splits if (s not in {"", "\n"})]

api/core/repositories/sqlalchemy_workflow_execution_repository.py

@@ -6,7 +6,6 @@ import json
 import logging
 from typing import Optional, Union
 
-from sqlalchemy import select
 from sqlalchemy.engine import Engine
 from sqlalchemy.orm import sessionmaker
 
@@ -206,44 +205,3 @@ class SQLAlchemyWorkflowExecutionRepository(WorkflowExecutionRepository):
             # Update the in-memory cache for faster subsequent lookups
             logger.debug(f"Updating cache for execution_id: {db_model.id}")
             self._execution_cache[db_model.id] = db_model
-
-    def get(self, execution_id: str) -> Optional[WorkflowExecution]:
-        """
-        Retrieve a WorkflowExecution by its ID.
-
-        First checks the in-memory cache, and if not found, queries the database.
-        If found in the database, adds it to the cache for future lookups.
-
-        Args:
-            execution_id: The workflow execution ID
-
-        Returns:
-            The WorkflowExecution instance if found, None otherwise
-        """
-        # First check the cache
-        if execution_id in self._execution_cache:
-            logger.debug(f"Cache hit for execution_id: {execution_id}")
-            # Convert cached DB model to domain model
-            cached_db_model = self._execution_cache[execution_id]
-            return self._to_domain_model(cached_db_model)
-
-        # If not in cache, query the database
-        logger.debug(f"Cache miss for execution_id: {execution_id}, querying database")
-        with self._session_factory() as session:
-            stmt = select(WorkflowRun).where(
-                WorkflowRun.id == execution_id,
-                WorkflowRun.tenant_id == self._tenant_id,
-            )
-
-            if self._app_id:
-                stmt = stmt.where(WorkflowRun.app_id == self._app_id)
-
-            db_model = session.scalar(stmt)
-            if db_model:
-                # Add DB model to cache
-                self._execution_cache[execution_id] = db_model
-
-                # Convert to domain model and return
-                return self._to_domain_model(db_model)
-
-            return None

api/core/repositories/sqlalchemy_workflow_node_execution_repository.py

@@ -7,7 +7,7 @@ import logging
 from collections.abc import Sequence
 from typing import Optional, Union
 
-from sqlalchemy import UnaryExpression, asc, delete, desc, select
+from sqlalchemy import UnaryExpression, asc, desc, select
 from sqlalchemy.engine import Engine
 from sqlalchemy.orm import sessionmaker
 
@@ -218,47 +218,6 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository)
                 logger.debug(f"Updating cache for node_execution_id: {db_model.node_execution_id}")
                 self._node_execution_cache[db_model.node_execution_id] = db_model
 
-    def get_by_node_execution_id(self, node_execution_id: str) -> Optional[WorkflowNodeExecution]:
-        """
-        Retrieve a NodeExecution by its node_execution_id.
-
-        First checks the in-memory cache, and if not found, queries the database.
-        If found in the database, adds it to the cache for future lookups.
-
-        Args:
-            node_execution_id: The node execution ID
-
-        Returns:
-            The NodeExecution instance if found, None otherwise
-        """
-        # First check the cache
-        if node_execution_id in self._node_execution_cache:
-            logger.debug(f"Cache hit for node_execution_id: {node_execution_id}")
-            # Convert cached DB model to domain model
-            cached_db_model = self._node_execution_cache[node_execution_id]
-            return self._to_domain_model(cached_db_model)
-
-        # If not in cache, query the database
-        logger.debug(f"Cache miss for node_execution_id: {node_execution_id}, querying database")
-        with self._session_factory() as session:
-            stmt = select(WorkflowNodeExecutionModel).where(
-                WorkflowNodeExecutionModel.node_execution_id == node_execution_id,
-                WorkflowNodeExecutionModel.tenant_id == self._tenant_id,
-            )
-
-            if self._app_id:
-                stmt = stmt.where(WorkflowNodeExecutionModel.app_id == self._app_id)
-
-            db_model = session.scalar(stmt)
-            if db_model:
-                # Add DB model to cache
-                self._node_execution_cache[node_execution_id] = db_model
-
-                # Convert to domain model and return
-                return self._to_domain_model(db_model)
-
-            return None
-
     def get_db_models_by_workflow_run(
         self,
         workflow_run_id: str,
@@ -344,68 +303,3 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository)
             domain_models.append(domain_model)
 
         return domain_models
-
-    def get_running_executions(self, workflow_run_id: str) -> Sequence[WorkflowNodeExecution]:
-        """
-        Retrieve all running NodeExecution instances for a specific workflow run.
-
-        This method queries the database directly and updates the cache with any
-        retrieved executions that have a node_execution_id.
-
-        Args:
-            workflow_run_id: The workflow run ID
-
-        Returns:
-            A list of running NodeExecution instances
-        """
-        with self._session_factory() as session:
-            stmt = select(WorkflowNodeExecutionModel).where(
-                WorkflowNodeExecutionModel.workflow_run_id == workflow_run_id,
-                WorkflowNodeExecutionModel.tenant_id == self._tenant_id,
-                WorkflowNodeExecutionModel.status == WorkflowNodeExecutionStatus.RUNNING,
-                WorkflowNodeExecutionModel.triggered_from == WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
-            )
-
-            if self._app_id:
-                stmt = stmt.where(WorkflowNodeExecutionModel.app_id == self._app_id)
-
-            db_models = session.scalars(stmt).all()
-            domain_models = []
-
-            for model in db_models:
-                # Update cache if node_execution_id is present
-                if model.node_execution_id:
-                    self._node_execution_cache[model.node_execution_id] = model
-
-                # Convert to domain model
-                domain_model = self._to_domain_model(model)
-                domain_models.append(domain_model)
-
-            return domain_models
-
-    def clear(self) -> None:
-        """
-        Clear all WorkflowNodeExecution records for the current tenant_id and app_id.
-
-        This method deletes all WorkflowNodeExecution records that match the tenant_id
-        and app_id (if provided) associated with this repository instance.
-        It also clears the in-memory cache.
-        """
-        with self._session_factory() as session:
-            stmt = delete(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.tenant_id == self._tenant_id)
-
-            if self._app_id:
-                stmt = stmt.where(WorkflowNodeExecutionModel.app_id == self._app_id)
-
-            result = session.execute(stmt)
-            session.commit()
-
-            deleted_count = result.rowcount
-            logger.info(
-                f"Cleared {deleted_count} workflow node execution records for tenant {self._tenant_id}"
-                + (f" and app {self._app_id}" if self._app_id else "")
-            )
-
-            # Clear the in-memory cache
-            self._node_execution_cache.clear()
-            logger.info("Cleared in-memory node execution cache")

api/core/tools/tool_manager.py

@@ -1,16 +1,19 @@
 import json
 import logging
 import mimetypes
-from collections.abc import Generator
+import time
+from collections.abc import Generator, Mapping
 from os import listdir, path
 from threading import Lock
 from typing import TYPE_CHECKING, Any, Literal, Optional, Union, cast
 
+from pydantic import TypeAdapter
 from yarl import URL
 
 import contexts
 from core.helper.provider_cache import ToolProviderCredentialsCache
 from core.plugin.entities.plugin import ToolProviderID
+from core.plugin.impl.oauth import OAuthHandler
 from core.plugin.impl.tool import PluginToolManager
 from core.tools.__base.tool_provider import ToolProviderController
 from core.tools.__base.tool_runtime import ToolRuntime
@@ -21,7 +24,7 @@ from core.tools.plugin_tool.tool import PluginTool
 from core.tools.utils.uuid_utils import is_valid_uuid
 from core.tools.workflow_as_tool.provider import WorkflowToolProviderController
 from core.workflow.entities.variable_pool import VariablePool
-from services.tools.mcp_tools_mange_service import MCPToolManageService
+from services.tools.mcp_tools_manage_service import MCPToolManageService
 
 if TYPE_CHECKING:
     from core.workflow.nodes.tool.entities import ToolEntity
@@ -244,12 +247,47 @@ class ToolManager:
                     tenant_id=tenant_id, provider=provider_id, credential_id=builtin_provider.id
                 ),
             )
+
+            # decrypt the credentials
+            decrypted_credentials: Mapping[str, Any] = encrypter.decrypt(builtin_provider.credentials)
+
+            # check if the credentials is expired
+            if builtin_provider.expires_at != -1 and (builtin_provider.expires_at - 60) < int(time.time()):
+                # TODO: circular import
+                from services.tools.builtin_tools_manage_service import BuiltinToolManageService
+
+                # refresh the credentials
+                tool_provider = ToolProviderID(provider_id)
+                provider_name = tool_provider.provider_name
+                redirect_uri = f"{dify_config.CONSOLE_API_URL}/console/api/oauth/plugin/{provider_id}/tool/callback"
+                system_credentials = BuiltinToolManageService.get_oauth_client(tenant_id, provider_id)
+                oauth_handler = OAuthHandler()
+                # refresh the credentials
+                refreshed_credentials = oauth_handler.refresh_credentials(
+                    tenant_id=tenant_id,
+                    user_id=builtin_provider.user_id,
+                    plugin_id=tool_provider.plugin_id,
+                    provider=provider_name,
+                    redirect_uri=redirect_uri,
+                    system_credentials=system_credentials or {},
+                    credentials=decrypted_credentials,
+                )
+                # update the credentials
+                builtin_provider.encrypted_credentials = (
+                    TypeAdapter(dict[str, Any])
+                    .dump_json(encrypter.encrypt(dict(refreshed_credentials.credentials)))
+                    .decode("utf-8")
+                )
+                builtin_provider.expires_at = refreshed_credentials.expires_at
+                db.session.commit()
+                decrypted_credentials = refreshed_credentials.credentials
+
             return cast(
                 BuiltinTool,
                 builtin_tool.fork_tool_runtime(
                     runtime=ToolRuntime(
                         tenant_id=tenant_id,
-                        credentials=encrypter.decrypt(builtin_provider.credentials),
+                        credentials=dict(decrypted_credentials),
                         credential_type=CredentialType.of(builtin_provider.credential_type),
                         runtime_parameters={},
                         invoke_from=invoke_from,

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

@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import Any, Optional
+from typing import Optional
 
 from msal_extensions.persistence import ABC  # type: ignore
 from pydantic import BaseModel, ConfigDict
@@ -21,11 +21,7 @@ class DatasetRetrieverBaseTool(BaseModel, ABC):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     @abstractmethod
-    def _run(
-        self,
-        *args: Any,
-        **kwargs: Any,
-    ) -> Any:
+    def _run(self, query: str) -> str:
         """Use the tool.
 
         Add run_manager: Optional[CallbackManagerForToolRun] = None

api/core/workflow/errors.py

@@ -2,7 +2,7 @@ from core.workflow.nodes.base import BaseNode
 
 
 class WorkflowNodeRunFailedError(Exception):
-    def __init__(self, node_instance: BaseNode, error: str):
-        self.node_instance = node_instance
-        self.error = error
-        super().__init__(f"Node {node_instance.node_data.title} run failed: {error}")
+    def __init__(self, node: BaseNode, err_msg: str):
+        self._node = node
+        self._error = err_msg
+        super().__init__(f"Node {node.title} run failed: {err_msg}")

api/core/workflow/graph_engine/__init__.py

@@ -1,3 +1,4 @@
 from .entities import Graph, GraphInitParams, GraphRuntimeState, RuntimeRouteState
+from .graph_engine import GraphEngine
 
-__all__ = ["Graph", "GraphInitParams", "GraphRuntimeState", "RuntimeRouteState"]
+__all__ = ["Graph", "GraphEngine", "GraphInitParams", "GraphRuntimeState", "RuntimeRouteState"]

api/core/workflow/graph_engine/graph_engine.py

@@ -12,7 +12,7 @@ from typing import Any, Optional, cast
 from flask import Flask, current_app
 
 from configs import dify_config
-from core.app.apps.base_app_queue_manager import GenerateTaskStoppedError
+from core.app.apps.exc import GenerateTaskStoppedError
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.workflow.entities.node_entities import AgentNodeStrategyInit, NodeRunResult
 from core.workflow.entities.variable_pool import VariablePool, VariableValue
@@ -48,11 +48,9 @@ from core.workflow.nodes.agent.entities import AgentNodeData
 from core.workflow.nodes.answer.answer_stream_processor import AnswerStreamProcessor
 from core.workflow.nodes.answer.base_stream_processor import StreamProcessor
 from core.workflow.nodes.base import BaseNode
-from core.workflow.nodes.base.entities import BaseNodeData
 from core.workflow.nodes.end.end_stream_processor import EndStreamProcessor
 from core.workflow.nodes.enums import ErrorStrategy, FailBranchSourceHandle
 from core.workflow.nodes.event import RunCompletedEvent, RunRetrieverResourceEvent, RunStreamChunkEvent
-from core.workflow.nodes.node_mapping import NODE_TYPE_CLASSES_MAPPING
 from core.workflow.utils import variable_utils
 from libs.flask_utils import preserve_flask_contexts
 from models.enums import UserFrom
@@ -260,12 +258,16 @@ class GraphEngine:
             # convert to specific node
             node_type = NodeType(node_config.get("data", {}).get("type"))
             node_version = node_config.get("data", {}).get("version", "1")
+
+            # Import here to avoid circular import
+            from core.workflow.nodes.node_mapping import NODE_TYPE_CLASSES_MAPPING
+
             node_cls = NODE_TYPE_CLASSES_MAPPING[node_type][node_version]
 
             previous_node_id = previous_route_node_state.node_id if previous_route_node_state else None
 
             # init workflow run state
-            node_instance = node_cls(  # type: ignore
+            node = node_cls(
                 id=route_node_state.id,
                 config=node_config,
                 graph_init_params=self.init_params,
@@ -274,11 +276,11 @@ class GraphEngine:
                 previous_node_id=previous_node_id,
                 thread_pool_id=self.thread_pool_id,
             )
-            node_instance = cast(BaseNode[BaseNodeData], node_instance)
+            node.init_node_data(node_config.get("data", {}))
             try:
                 # run node
                 generator = self._run_node(
-                    node_instance=node_instance,
+                    node=node,
                     route_node_state=route_node_state,
                     parallel_id=in_parallel_id,
                     parallel_start_node_id=parallel_start_node_id,
@@ -306,16 +308,16 @@ class GraphEngine:
                 route_node_state.failed_reason = str(e)
                 yield NodeRunFailedEvent(
                     error=str(e),
-                    id=node_instance.id,
+                    id=node.id,
                     node_id=next_node_id,
                     node_type=node_type,
-                    node_data=node_instance.node_data,
+                    node_data=node.get_base_node_data(),
                     route_node_state=route_node_state,
                     parallel_id=in_parallel_id,
                     parallel_start_node_id=parallel_start_node_id,
                     parent_parallel_id=parent_parallel_id,
                     parent_parallel_start_node_id=parent_parallel_start_node_id,
-                    node_version=node_instance.version(),
+                    node_version=node.version(),
                 )
                 raise e
 
@@ -337,7 +339,7 @@ class GraphEngine:
                 edge = edge_mappings[0]
                 if (
                     previous_route_node_state.status == RouteNodeState.Status.EXCEPTION
-                    and node_instance.node_data.error_strategy == ErrorStrategy.FAIL_BRANCH
+                    and node.error_strategy == ErrorStrategy.FAIL_BRANCH
                     and edge.run_condition is None
                 ):
                     break
@@ -413,8 +415,8 @@ class GraphEngine:
 
                     next_node_id = final_node_id
                 elif (
-                    node_instance.node_data.error_strategy == ErrorStrategy.FAIL_BRANCH
-                    and node_instance.should_continue_on_error
+                    node.continue_on_error
+                    and node.error_strategy == ErrorStrategy.FAIL_BRANCH
                     and previous_route_node_state.status == RouteNodeState.Status.EXCEPTION
                 ):
                     break
@@ -597,7 +599,7 @@ class GraphEngine:
 
     def _run_node(
         self,
-        node_instance: BaseNode[BaseNodeData],
+        node: BaseNode,
         route_node_state: RouteNodeState,
         parallel_id: Optional[str] = None,
         parallel_start_node_id: Optional[str] = None,
@@ -611,29 +613,29 @@ class GraphEngine:
         # trigger node run start event
         agent_strategy = (
             AgentNodeStrategyInit(
-                name=cast(AgentNodeData, node_instance.node_data).agent_strategy_name,
-                icon=cast(AgentNode, node_instance).agent_strategy_icon,
+                name=cast(AgentNodeData, node.get_base_node_data()).agent_strategy_name,
+                icon=cast(AgentNode, node).agent_strategy_icon,
             )
-            if node_instance.node_type == NodeType.AGENT
+            if node.type_ == NodeType.AGENT
             else None
         )
         yield NodeRunStartedEvent(
-            id=node_instance.id,
-            node_id=node_instance.node_id,
-            node_type=node_instance.node_type,
-            node_data=node_instance.node_data,
+            id=node.id,
+            node_id=node.node_id,
+            node_type=node.type_,
+            node_data=node.get_base_node_data(),
             route_node_state=route_node_state,
-            predecessor_node_id=node_instance.previous_node_id,
+            predecessor_node_id=node.previous_node_id,
             parallel_id=parallel_id,
             parallel_start_node_id=parallel_start_node_id,
             parent_parallel_id=parent_parallel_id,
             parent_parallel_start_node_id=parent_parallel_start_node_id,
             agent_strategy=agent_strategy,
-            node_version=node_instance.version(),
+            node_version=node.version(),
         )
 
-        max_retries = node_instance.node_data.retry_config.max_retries
-        retry_interval = node_instance.node_data.retry_config.retry_interval_seconds
+        max_retries = node.retry_config.max_retries
+        retry_interval = node.retry_config.retry_interval_seconds
         retries = 0
         should_continue_retry = True
         while should_continue_retry and retries <= max_retries:
@@ -642,7 +644,7 @@ class GraphEngine:
                 retry_start_at = datetime.now(UTC).replace(tzinfo=None)
                 # yield control to other threads
                 time.sleep(0.001)
-                event_stream = node_instance.run()
+                event_stream = node.run()
                 for event in event_stream:
                     if isinstance(event, GraphEngineEvent):
                         # add parallel info to iteration event
@@ -658,21 +660,21 @@ class GraphEngine:
                             if run_result.status == WorkflowNodeExecutionStatus.FAILED:
                                 if (
                                     retries == max_retries
-                                    and node_instance.node_type == NodeType.HTTP_REQUEST
+                                    and node.type_ == NodeType.HTTP_REQUEST
                                     and run_result.outputs
-                                    and not node_instance.should_continue_on_error
+                                    and not node.continue_on_error
                                 ):
                                     run_result.status = WorkflowNodeExecutionStatus.SUCCEEDED
-                                if node_instance.should_retry and retries < max_retries:
+                                if node.retry and retries < max_retries:
                                     retries += 1
                                     route_node_state.node_run_result = run_result
                                     yield NodeRunRetryEvent(
                                         id=str(uuid.uuid4()),
-                                        node_id=node_instance.node_id,
-                                        node_type=node_instance.node_type,
-                                        node_data=node_instance.node_data,
+                                        node_id=node.node_id,
+                                        node_type=node.type_,
+                                        node_data=node.get_base_node_data(),
                                         route_node_state=route_node_state,
-                                        predecessor_node_id=node_instance.previous_node_id,
+                                        predecessor_node_id=node.previous_node_id,
                                         parallel_id=parallel_id,
                                         parallel_start_node_id=parallel_start_node_id,
                                         parent_parallel_id=parent_parallel_id,
@@ -680,17 +682,17 @@ class GraphEngine:
                                         error=run_result.error or "Unknown error",
                                         retry_index=retries,
                                         start_at=retry_start_at,
-                                        node_version=node_instance.version(),
+                                        node_version=node.version(),
                                     )
                                     time.sleep(retry_interval)
                                     break
                             route_node_state.set_finished(run_result=run_result)
 
                             if run_result.status == WorkflowNodeExecutionStatus.FAILED:
-                                if node_instance.should_continue_on_error:
+                                if node.continue_on_error:
                                     # if run failed, handle error
                                     run_result = self._handle_continue_on_error(
-                                        node_instance,
+                                        node,
                                         event.run_result,
                                         self.graph_runtime_state.variable_pool,
                                         handle_exceptions=handle_exceptions,
@@ -701,44 +703,44 @@ class GraphEngine:
                                         for variable_key, variable_value in run_result.outputs.items():
                                             # append variables to variable pool recursively
                                             self._append_variables_recursively(
-                                                node_id=node_instance.node_id,
+                                                node_id=node.node_id,
                                                 variable_key_list=[variable_key],
                                                 variable_value=variable_value,
                                             )
                                     yield NodeRunExceptionEvent(
                                         error=run_result.error or "System Error",
-                                        id=node_instance.id,
-                                        node_id=node_instance.node_id,
-                                        node_type=node_instance.node_type,
-                                        node_data=node_instance.node_data,
+                                        id=node.id,
+                                        node_id=node.node_id,
+                                        node_type=node.type_,
+                                        node_data=node.get_base_node_data(),
                                         route_node_state=route_node_state,
                                         parallel_id=parallel_id,
                                         parallel_start_node_id=parallel_start_node_id,
                                         parent_parallel_id=parent_parallel_id,
                                         parent_parallel_start_node_id=parent_parallel_start_node_id,
-                                        node_version=node_instance.version(),
+                                        node_version=node.version(),
                                     )
                                     should_continue_retry = False
                                 else:
                                     yield NodeRunFailedEvent(
                                         error=route_node_state.failed_reason or "Unknown error.",
-                                        id=node_instance.id,
-                                        node_id=node_instance.node_id,
-                                        node_type=node_instance.node_type,
-                                        node_data=node_instance.node_data,
+                                        id=node.id,
+                                        node_id=node.node_id,
+                                        node_type=node.type_,
+                                        node_data=node.get_base_node_data(),
                                         route_node_state=route_node_state,
                                         parallel_id=parallel_id,
                                         parallel_start_node_id=parallel_start_node_id,
                                         parent_parallel_id=parent_parallel_id,
                                         parent_parallel_start_node_id=parent_parallel_start_node_id,
-                                        node_version=node_instance.version(),
+                                        node_version=node.version(),
                                     )
                                 should_continue_retry = False
                             elif run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED:
                                 if (
-                                    node_instance.should_continue_on_error
-                                    and self.graph.edge_mapping.get(node_instance.node_id)
-                                    and node_instance.node_data.error_strategy is ErrorStrategy.FAIL_BRANCH
+                                    node.continue_on_error
+                                    and self.graph.edge_mapping.get(node.node_id)
+                                    and node.error_strategy is ErrorStrategy.FAIL_BRANCH
                                 ):
                                     run_result.edge_source_handle = FailBranchSourceHandle.SUCCESS
                                 if run_result.metadata and run_result.metadata.get(
@@ -758,7 +760,7 @@ class GraphEngine:
                                     for variable_key, variable_value in run_result.outputs.items():
                                         # append variables to variable pool recursively
                                         self._append_variables_recursively(
-                                            node_id=node_instance.node_id,
+                                            node_id=node.node_id,
                                             variable_key_list=[variable_key],
                                             variable_value=variable_value,
                                         )
@@ -783,26 +785,26 @@ class GraphEngine:
                                     run_result.metadata = metadata_dict
 
                                 yield NodeRunSucceededEvent(
-                                    id=node_instance.id,
-                                    node_id=node_instance.node_id,
-                                    node_type=node_instance.node_type,
-                                    node_data=node_instance.node_data,
+                                    id=node.id,
+                                    node_id=node.node_id,
+                                    node_type=node.type_,
+                                    node_data=node.get_base_node_data(),
                                     route_node_state=route_node_state,
                                     parallel_id=parallel_id,
                                     parallel_start_node_id=parallel_start_node_id,
                                     parent_parallel_id=parent_parallel_id,
                                     parent_parallel_start_node_id=parent_parallel_start_node_id,
-                                    node_version=node_instance.version(),
+                                    node_version=node.version(),
                                 )
                                 should_continue_retry = False
 
                             break
                         elif isinstance(event, RunStreamChunkEvent):
                             yield NodeRunStreamChunkEvent(
-                                id=node_instance.id,
-                                node_id=node_instance.node_id,
-                                node_type=node_instance.node_type,
-                                node_data=node_instance.node_data,
+                                id=node.id,
+                                node_id=node.node_id,
+                                node_type=node.type_,
+                                node_data=node.get_base_node_data(),
                                 chunk_content=event.chunk_content,
                                 from_variable_selector=event.from_variable_selector,
                                 route_node_state=route_node_state,
@@ -810,14 +812,14 @@ class GraphEngine:
                                 parallel_start_node_id=parallel_start_node_id,
                                 parent_parallel_id=parent_parallel_id,
                                 parent_parallel_start_node_id=parent_parallel_start_node_id,
-                                node_version=node_instance.version(),
+                                node_version=node.version(),
                             )
                         elif isinstance(event, RunRetrieverResourceEvent):
                             yield NodeRunRetrieverResourceEvent(
-                                id=node_instance.id,
-                                node_id=node_instance.node_id,
-                                node_type=node_instance.node_type,
-                                node_data=node_instance.node_data,
+                                id=node.id,
+                                node_id=node.node_id,
+                                node_type=node.type_,
+                                node_data=node.get_base_node_data(),
                                 retriever_resources=event.retriever_resources,
                                 context=event.context,
                                 route_node_state=route_node_state,
@@ -825,7 +827,7 @@ class GraphEngine:
                                 parallel_start_node_id=parallel_start_node_id,
                                 parent_parallel_id=parent_parallel_id,
                                 parent_parallel_start_node_id=parent_parallel_start_node_id,
-                                node_version=node_instance.version(),
+                                node_version=node.version(),
                             )
             except GenerateTaskStoppedError:
                 # trigger node run failed event
@@ -833,20 +835,20 @@ class GraphEngine:
                 route_node_state.failed_reason = "Workflow stopped."
                 yield NodeRunFailedEvent(
                     error="Workflow stopped.",
-                    id=node_instance.id,
-                    node_id=node_instance.node_id,
-                    node_type=node_instance.node_type,
-                    node_data=node_instance.node_data,
+                    id=node.id,
+                    node_id=node.node_id,
+                    node_type=node.type_,
+                    node_data=node.get_base_node_data(),
                     route_node_state=route_node_state,
                     parallel_id=parallel_id,
                     parallel_start_node_id=parallel_start_node_id,
                     parent_parallel_id=parent_parallel_id,
                     parent_parallel_start_node_id=parent_parallel_start_node_id,
-                    node_version=node_instance.version(),
+                    node_version=node.version(),
                 )
                 return
             except Exception as e:
-                logger.exception(f"Node {node_instance.node_data.title} run failed")
+                logger.exception(f"Node {node.title} run failed")
                 raise e
 
     def _append_variables_recursively(self, node_id: str, variable_key_list: list[str], variable_value: VariableValue):
@@ -886,22 +888,14 @@ class GraphEngine:
 
     def _handle_continue_on_error(
         self,
-        node_instance: BaseNode[BaseNodeData],
+        node: BaseNode,
         error_result: NodeRunResult,
         variable_pool: VariablePool,
         handle_exceptions: list[str] = [],
     ) -> NodeRunResult:
-        """
-        handle continue on error when self._should_continue_on_error is True
-
-
-        :param    error_result (NodeRunResult): error run result
-        :param    variable_pool (VariablePool): variable pool
-        :return:  excption run result
-        """
         # add error message and error type to variable pool
-        variable_pool.add([node_instance.node_id, "error_message"], error_result.error)
-        variable_pool.add([node_instance.node_id, "error_type"], error_result.error_type)
+        variable_pool.add([node.node_id, "error_message"], error_result.error)
+        variable_pool.add([node.node_id, "error_type"], error_result.error_type)
         # add error message to handle_exceptions
         handle_exceptions.append(error_result.error or "")
         node_error_args: dict[str, Any] = {
@@ -909,21 +903,21 @@ class GraphEngine:
             "error": error_result.error,
             "inputs": error_result.inputs,
             "metadata": {
-                WorkflowNodeExecutionMetadataKey.ERROR_STRATEGY: node_instance.node_data.error_strategy,
+                WorkflowNodeExecutionMetadataKey.ERROR_STRATEGY: node.error_strategy,
             },
         }
 
-        if node_instance.node_data.error_strategy is ErrorStrategy.DEFAULT_VALUE:
+        if node.error_strategy is ErrorStrategy.DEFAULT_VALUE:
             return NodeRunResult(
                 **node_error_args,
                 outputs={
-                    **node_instance.node_data.default_value_dict,
+                    **node.default_value_dict,
                     "error_message": error_result.error,
                     "error_type": error_result.error_type,
                 },
             )
-        elif node_instance.node_data.error_strategy is ErrorStrategy.FAIL_BRANCH:
-            if self.graph.edge_mapping.get(node_instance.node_id):
+        elif node.error_strategy is ErrorStrategy.FAIL_BRANCH:
+            if self.graph.edge_mapping.get(node.node_id):
                 node_error_args["edge_source_handle"] = FailBranchSourceHandle.FAILED
             return NodeRunResult(
                 **node_error_args,

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

@@ -1,5 +1,4 @@
 import json
-import uuid
 from collections.abc import Generator, Mapping, Sequence
 from typing import Any, Optional, cast
 
@@ -11,8 +10,10 @@ from sqlalchemy.orm import Session
 from core.agent.entities import AgentToolEntity
 from core.agent.plugin_entities import AgentStrategyParameter
 from core.agent.strategy.plugin import PluginAgentStrategy
+from core.file import File, FileTransferMethod
 from core.memory.token_buffer_memory import TokenBufferMemory
 from core.model_manager import ModelInstance, ModelManager
+from core.model_runtime.entities.llm_entities import LLMUsage
 from core.model_runtime.entities.model_entities import AIModelEntity, ModelType
 from core.plugin.entities.request import InvokeCredentials
 from core.plugin.impl.exc import PluginDaemonClientSideError
@@ -25,45 +26,75 @@ from core.tools.entities.tool_entities import (
     ToolProviderType,
 )
 from core.tools.tool_manager import ToolManager
-from core.variables.segments import StringSegment
+from core.tools.utils.message_transformer import ToolFileMessageTransformer
+from core.variables.segments import ArrayFileSegment, StringSegment
 from core.workflow.entities.node_entities import NodeRunResult
 from core.workflow.entities.variable_pool import VariablePool
-from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
+from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus
 from core.workflow.enums import SystemVariableKey
+from core.workflow.graph_engine.entities.event import AgentLogEvent
 from core.workflow.nodes.agent.entities import AgentNodeData, AgentOldVersionModelFeatures, ParamsAutoGenerated
-from core.workflow.nodes.base.entities import BaseNodeData
-from core.workflow.nodes.enums import NodeType
-from core.workflow.nodes.event.event import RunCompletedEvent
-from core.workflow.nodes.tool.tool_node import ToolNode
+from core.workflow.nodes.base import BaseNode
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
+from core.workflow.nodes.event import RunCompletedEvent, RunStreamChunkEvent
 from core.workflow.utils.variable_template_parser import VariableTemplateParser
 from extensions.ext_database import db
+from factories import file_factory
 from factories.agent_factory import get_plugin_agent_strategy
+from models import ToolFile
 from models.model import Conversation
+from services.tools.builtin_tools_manage_service import BuiltinToolManageService
+
+from .exc import (
+    AgentInputTypeError,
+    AgentInvocationError,
+    AgentMessageTransformError,
+    AgentVariableNotFoundError,
+    AgentVariableTypeError,
+    ToolFileNotFoundError,
+)
 
 
-class AgentNode(ToolNode):
+class AgentNode(BaseNode):
     """
     Agent Node
     """
 
-    _node_data_cls = AgentNodeData  # type: ignore
     _node_type = NodeType.AGENT
+    _node_data: AgentNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = AgentNodeData.model_validate(data)
+
+    def _get_error_strategy(self) -> Optional[ErrorStrategy]:
+        return self._node_data.error_strategy
+
+    def _get_retry_config(self) -> RetryConfig:
+        return self._node_data.retry_config
+
+    def _get_title(self) -> str:
+        return self._node_data.title
+
+    def _get_description(self) -> Optional[str]:
+        return self._node_data.desc
+
+    def _get_default_value_dict(self) -> dict[str, Any]:
+        return self._node_data.default_value_dict
+
+    def get_base_node_data(self) -> BaseNodeData:
+        return self._node_data
 
     @classmethod
     def version(cls) -> str:
         return "1"
 
     def _run(self) -> Generator:
-        """
-        Run the agent node
-        """
-        node_data = cast(AgentNodeData, self.node_data)
-
         try:
             strategy = get_plugin_agent_strategy(
                 tenant_id=self.tenant_id,
-                agent_strategy_provider_name=node_data.agent_strategy_provider_name,
-                agent_strategy_name=node_data.agent_strategy_name,
+                agent_strategy_provider_name=self._node_data.agent_strategy_provider_name,
+                agent_strategy_name=self._node_data.agent_strategy_name,
             )
         except Exception as e:
             yield RunCompletedEvent(
@@ -81,13 +112,13 @@ class AgentNode(ToolNode):
         parameters = self._generate_agent_parameters(
             agent_parameters=agent_parameters,
             variable_pool=self.graph_runtime_state.variable_pool,
-            node_data=node_data,
+            node_data=self._node_data,
             strategy=strategy,
         )
         parameters_for_log = self._generate_agent_parameters(
             agent_parameters=agent_parameters,
             variable_pool=self.graph_runtime_state.variable_pool,
-            node_data=node_data,
+            node_data=self._node_data,
             for_log=True,
             strategy=strategy,
         )
@@ -105,59 +136,39 @@ class AgentNode(ToolNode):
                 credentials=credentials,
             )
         except Exception as e:
+            error = AgentInvocationError(f"Failed to invoke agent: {str(e)}", original_error=e)
             yield RunCompletedEvent(
                 run_result=NodeRunResult(
                     status=WorkflowNodeExecutionStatus.FAILED,
                     inputs=parameters_for_log,
-                    error=f"Failed to invoke agent: {str(e)}",
+                    error=str(error),
                 )
             )
             return
 
         try:
-            # convert tool messages
-            agent_thoughts: list = []
-
-            thought_log_message = ToolInvokeMessage(
-                type=ToolInvokeMessage.MessageType.LOG,
-                message=ToolInvokeMessage.LogMessage(
-                    id=str(uuid.uuid4()),
-                    label=f"Agent Strategy: {cast(AgentNodeData, self.node_data).agent_strategy_name}",
-                    parent_id=None,
-                    error=None,
-                    status=ToolInvokeMessage.LogMessage.LogStatus.START,
-                    data={
-                        "strategy": cast(AgentNodeData, self.node_data).agent_strategy_name,
-                        "parameters": parameters_for_log,
-                        "thought_process": "Agent strategy execution started",
-                    },
-                    metadata={
-                        "icon": self.agent_strategy_icon,
-                        "agent_strategy": cast(AgentNodeData, self.node_data).agent_strategy_name,
-                    },
-                ),
-            )
-
-            def enhanced_message_stream():
-                yield thought_log_message
-
-                yield from message_stream
-
             yield from self._transform_message(
-                message_stream,
-                {
+                messages=message_stream,
+                tool_info={
                     "icon": self.agent_strategy_icon,
-                    "agent_strategy": cast(AgentNodeData, self.node_data).agent_strategy_name,
+                    "agent_strategy": cast(AgentNodeData, self._node_data).agent_strategy_name,
                 },
-                parameters_for_log,
-                agent_thoughts,
+                parameters_for_log=parameters_for_log,
+                user_id=self.user_id,
+                tenant_id=self.tenant_id,
+                node_type=self.type_,
+                node_id=self.node_id,
+                node_execution_id=self.id,
             )
         except PluginDaemonClientSideError as e:
+            transform_error = AgentMessageTransformError(
+                f"Failed to transform agent message: {str(e)}", original_error=e
+            )
             yield RunCompletedEvent(
                 run_result=NodeRunResult(
                     status=WorkflowNodeExecutionStatus.FAILED,
                     inputs=parameters_for_log,
-                    error=f"Failed to transform agent message: {str(e)}",
+                    error=str(transform_error),
                 )
             )
 
@@ -194,7 +205,7 @@ class AgentNode(ToolNode):
             if agent_input.type == "variable":
                 variable = variable_pool.get(agent_input.value)  # type: ignore
                 if variable is None:
-                    raise ValueError(f"Variable {agent_input.value} does not exist")
+                    raise AgentVariableNotFoundError(str(agent_input.value))
                 parameter_value = variable.value
             elif agent_input.type in {"mixed", "constant"}:
                 # variable_pool.convert_template expects a string template,
@@ -216,7 +227,7 @@ class AgentNode(ToolNode):
                 except json.JSONDecodeError:
                     parameter_value = parameter_value
             else:
-                raise ValueError(f"Unknown agent input type '{agent_input.type}'")
+                raise AgentInputTypeError(agent_input.type)
             value = parameter_value
             if parameter.type == "array[tools]":
                 value = cast(list[dict[str, Any]], value)
@@ -259,7 +270,14 @@ class AgentNode(ToolNode):
                         )
 
                         extra = tool.get("extra", {})
-                        runtime_variable_pool = variable_pool if self.node_data.version != "1" else None
+
+                        # This is an issue that caused problems before.
+                        # Logically, we shouldn't use the node_data.version field for judgment
+                        # But for backward compatibility with historical data
+                        # this version field judgment is still preserved here.
+                        runtime_variable_pool: VariablePool | None = None
+                        if node_data.version != "1" or node_data.tool_node_version != "1":
+                            runtime_variable_pool = variable_pool
                         tool_runtime = ToolManager.get_agent_tool_runtime(
                             self.tenant_id, self.app_id, entity, self.invoke_from, runtime_variable_pool
                         )
@@ -343,19 +361,14 @@ class AgentNode(ToolNode):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: BaseNodeData,
+        node_data: Mapping[str, Any],
     ) -> Mapping[str, Sequence[str]]:
-        """
-        Extract variable selector to variable mapping
-        :param graph_config: graph config
-        :param node_id: node id
-        :param node_data: node data
-        :return:
-        """
-        node_data = cast(AgentNodeData, node_data)
+        # Create typed NodeData from dict
+        typed_node_data = AgentNodeData.model_validate(node_data)
+
         result: dict[str, Any] = {}
-        for parameter_name in node_data.agent_parameters:
-            input = node_data.agent_parameters[parameter_name]
+        for parameter_name in typed_node_data.agent_parameters:
+            input = typed_node_data.agent_parameters[parameter_name]
             if input.type in ["mixed", "constant"]:
                 selectors = VariableTemplateParser(str(input.value)).extract_variable_selectors()
                 for selector in selectors:
@@ -380,7 +393,7 @@ class AgentNode(ToolNode):
                 plugin
                 for plugin in plugins
                 if f"{plugin.plugin_id}/{plugin.name}"
-                == cast(AgentNodeData, self.node_data).agent_strategy_provider_name
+                == cast(AgentNodeData, self._node_data).agent_strategy_provider_name
             )
             icon = current_plugin.declaration.icon
         except StopIteration:
@@ -448,3 +461,236 @@ class AgentNode(ToolNode):
             return tools
         else:
             return [tool for tool in tools if tool.get("type") != ToolProviderType.MCP.value]
+
+    def _transform_message(
+        self,
+        messages: Generator[ToolInvokeMessage, None, None],
+        tool_info: Mapping[str, Any],
+        parameters_for_log: dict[str, Any],
+        user_id: str,
+        tenant_id: str,
+        node_type: NodeType,
+        node_id: str,
+        node_execution_id: str,
+    ) -> Generator:
+        """
+        Convert ToolInvokeMessages into tuple[plain_text, files]
+        """
+        # transform message and handle file storage
+        message_stream = ToolFileMessageTransformer.transform_tool_invoke_messages(
+            messages=messages,
+            user_id=user_id,
+            tenant_id=tenant_id,
+            conversation_id=None,
+        )
+
+        text = ""
+        files: list[File] = []
+        json_list: list[dict] = []
+
+        agent_logs: list[AgentLogEvent] = []
+        agent_execution_metadata: Mapping[WorkflowNodeExecutionMetadataKey, Any] = {}
+        llm_usage: LLMUsage | None = None
+        variables: dict[str, Any] = {}
+
+        for message in message_stream:
+            if message.type in {
+                ToolInvokeMessage.MessageType.IMAGE_LINK,
+                ToolInvokeMessage.MessageType.BINARY_LINK,
+                ToolInvokeMessage.MessageType.IMAGE,
+            }:
+                assert isinstance(message.message, ToolInvokeMessage.TextMessage)
+
+                url = message.message.text
+                if message.meta:
+                    transfer_method = message.meta.get("transfer_method", FileTransferMethod.TOOL_FILE)
+                else:
+                    transfer_method = FileTransferMethod.TOOL_FILE
+
+                tool_file_id = str(url).split("/")[-1].split(".")[0]
+
+                with Session(db.engine) as session:
+                    stmt = select(ToolFile).where(ToolFile.id == tool_file_id)
+                    tool_file = session.scalar(stmt)
+                    if tool_file is None:
+                        raise ToolFileNotFoundError(tool_file_id)
+
+                mapping = {
+                    "tool_file_id": tool_file_id,
+                    "type": file_factory.get_file_type_by_mime_type(tool_file.mimetype),
+                    "transfer_method": transfer_method,
+                    "url": url,
+                }
+                file = file_factory.build_from_mapping(
+                    mapping=mapping,
+                    tenant_id=tenant_id,
+                )
+                files.append(file)
+            elif message.type == ToolInvokeMessage.MessageType.BLOB:
+                # get tool file id
+                assert isinstance(message.message, ToolInvokeMessage.TextMessage)
+                assert message.meta
+
+                tool_file_id = message.message.text.split("/")[-1].split(".")[0]
+                with Session(db.engine) as session:
+                    stmt = select(ToolFile).where(ToolFile.id == tool_file_id)
+                    tool_file = session.scalar(stmt)
+                    if tool_file is None:
+                        raise ToolFileNotFoundError(tool_file_id)
+
+                mapping = {
+                    "tool_file_id": tool_file_id,
+                    "transfer_method": FileTransferMethod.TOOL_FILE,
+                }
+
+                files.append(
+                    file_factory.build_from_mapping(
+                        mapping=mapping,
+                        tenant_id=tenant_id,
+                    )
+                )
+            elif message.type == ToolInvokeMessage.MessageType.TEXT:
+                assert isinstance(message.message, ToolInvokeMessage.TextMessage)
+                text += message.message.text
+                yield RunStreamChunkEvent(chunk_content=message.message.text, from_variable_selector=[node_id, "text"])
+            elif message.type == ToolInvokeMessage.MessageType.JSON:
+                assert isinstance(message.message, ToolInvokeMessage.JsonMessage)
+                if node_type == NodeType.AGENT:
+                    msg_metadata: dict[str, Any] = message.message.json_object.pop("execution_metadata", {})
+                    llm_usage = LLMUsage.from_metadata(msg_metadata)
+                    agent_execution_metadata = {
+                        WorkflowNodeExecutionMetadataKey(key): value
+                        for key, value in msg_metadata.items()
+                        if key in WorkflowNodeExecutionMetadataKey.__members__.values()
+                    }
+                if message.message.json_object is not None:
+                    json_list.append(message.message.json_object)
+            elif message.type == ToolInvokeMessage.MessageType.LINK:
+                assert isinstance(message.message, ToolInvokeMessage.TextMessage)
+                stream_text = f"Link: {message.message.text}\n"
+                text += stream_text
+                yield RunStreamChunkEvent(chunk_content=stream_text, from_variable_selector=[node_id, "text"])
+            elif message.type == ToolInvokeMessage.MessageType.VARIABLE:
+                assert isinstance(message.message, ToolInvokeMessage.VariableMessage)
+                variable_name = message.message.variable_name
+                variable_value = message.message.variable_value
+                if message.message.stream:
+                    if not isinstance(variable_value, str):
+                        raise AgentVariableTypeError(
+                            "When 'stream' is True, 'variable_value' must be a string.",
+                            variable_name=variable_name,
+                            expected_type="str",
+                            actual_type=type(variable_value).__name__,
+                        )
+                    if variable_name not in variables:
+                        variables[variable_name] = ""
+                    variables[variable_name] += variable_value
+
+                    yield RunStreamChunkEvent(
+                        chunk_content=variable_value, from_variable_selector=[node_id, variable_name]
+                    )
+                else:
+                    variables[variable_name] = variable_value
+            elif message.type == ToolInvokeMessage.MessageType.FILE:
+                assert message.meta is not None
+                assert isinstance(message.meta, File)
+                files.append(message.meta["file"])
+            elif message.type == ToolInvokeMessage.MessageType.LOG:
+                assert isinstance(message.message, ToolInvokeMessage.LogMessage)
+                if message.message.metadata:
+                    icon = tool_info.get("icon", "")
+                    dict_metadata = dict(message.message.metadata)
+                    if dict_metadata.get("provider"):
+                        manager = PluginInstaller()
+                        plugins = manager.list_plugins(tenant_id)
+                        try:
+                            current_plugin = next(
+                                plugin
+                                for plugin in plugins
+                                if f"{plugin.plugin_id}/{plugin.name}" == dict_metadata["provider"]
+                            )
+                            icon = current_plugin.declaration.icon
+                        except StopIteration:
+                            pass
+                        icon_dark = None
+                        try:
+                            builtin_tool = next(
+                                provider
+                                for provider in BuiltinToolManageService.list_builtin_tools(
+                                    user_id,
+                                    tenant_id,
+                                )
+                                if provider.name == dict_metadata["provider"]
+                            )
+                            icon = builtin_tool.icon
+                            icon_dark = builtin_tool.icon_dark
+                        except StopIteration:
+                            pass
+
+                        dict_metadata["icon"] = icon
+                        dict_metadata["icon_dark"] = icon_dark
+                        message.message.metadata = dict_metadata
+                agent_log = AgentLogEvent(
+                    id=message.message.id,
+                    node_execution_id=node_execution_id,
+                    parent_id=message.message.parent_id,
+                    error=message.message.error,
+                    status=message.message.status.value,
+                    data=message.message.data,
+                    label=message.message.label,
+                    metadata=message.message.metadata,
+                    node_id=node_id,
+                )
+
+                # check if the agent log is already in the list
+                for log in agent_logs:
+                    if log.id == agent_log.id:
+                        # update the log
+                        log.data = agent_log.data
+                        log.status = agent_log.status
+                        log.error = agent_log.error
+                        log.label = agent_log.label
+                        log.metadata = agent_log.metadata
+                        break
+                else:
+                    agent_logs.append(agent_log)
+
+                yield agent_log
+
+        # Add agent_logs to outputs['json'] to ensure frontend can access thinking process
+        json_output: list[dict[str, Any]] = []
+
+        # Step 1: append each agent log as its own dict.
+        if agent_logs:
+            for log in agent_logs:
+                json_output.append(
+                    {
+                        "id": log.id,
+                        "parent_id": log.parent_id,
+                        "error": log.error,
+                        "status": log.status,
+                        "data": log.data,
+                        "label": log.label,
+                        "metadata": log.metadata,
+                        "node_id": log.node_id,
+                    }
+                )
+        # Step 2: normalize JSON into {"data": [...]}.change json to list[dict]
+        if json_list:
+            json_output.extend(json_list)
+        else:
+            json_output.append({"data": []})
+
+        yield RunCompletedEvent(
+            run_result=NodeRunResult(
+                status=WorkflowNodeExecutionStatus.SUCCEEDED,
+                outputs={"text": text, "files": ArrayFileSegment(value=files), "json": json_output, **variables},
+                metadata={
+                    **agent_execution_metadata,
+                    WorkflowNodeExecutionMetadataKey.TOOL_INFO: tool_info,
+                    WorkflowNodeExecutionMetadataKey.AGENT_LOG: agent_logs,
+                },
+                inputs=parameters_for_log,
+                llm_usage=llm_usage,
+            )
+        )

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

@@ -13,6 +13,10 @@ class AgentNodeData(BaseNodeData):
     agent_strategy_name: str
     agent_strategy_label: str  # redundancy
     memory: MemoryConfig | None = None
+    # The version of the tool parameter.
+    # If this value is None, it indicates this is a previous version
+    # and requires using the legacy parameter parsing rules.
+    tool_node_version: str | None = None
 
     class AgentInput(BaseModel):
         value: Union[list[str], list[ToolSelector], Any]

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

@@ -0,0 +1,124 @@
+from typing import Optional
+
+
+class AgentNodeError(Exception):
+    """Base exception for all agent node errors."""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(self.message)
+
+
+class AgentStrategyError(AgentNodeError):
+    """Exception raised when there's an error with the agent strategy."""
+
+    def __init__(self, message: str, strategy_name: Optional[str] = None, provider_name: Optional[str] = None):
+        self.strategy_name = strategy_name
+        self.provider_name = provider_name
+        super().__init__(message)
+
+
+class AgentStrategyNotFoundError(AgentStrategyError):
+    """Exception raised when the specified agent strategy is not found."""
+
+    def __init__(self, strategy_name: str, provider_name: Optional[str] = None):
+        super().__init__(
+            f"Agent strategy '{strategy_name}' not found"
+            + (f" for provider '{provider_name}'" if provider_name else ""),
+            strategy_name,
+            provider_name,
+        )
+
+
+class AgentInvocationError(AgentNodeError):
+    """Exception raised when there's an error invoking the agent."""
+
+    def __init__(self, message: str, original_error: Optional[Exception] = None):
+        self.original_error = original_error
+        super().__init__(message)
+
+
+class AgentParameterError(AgentNodeError):
+    """Exception raised when there's an error with agent parameters."""
+
+    def __init__(self, message: str, parameter_name: Optional[str] = None):
+        self.parameter_name = parameter_name
+        super().__init__(message)
+
+
+class AgentVariableError(AgentNodeError):
+    """Exception raised when there's an error with variables in the agent node."""
+
+    def __init__(self, message: str, variable_name: Optional[str] = None):
+        self.variable_name = variable_name
+        super().__init__(message)
+
+
+class AgentVariableNotFoundError(AgentVariableError):
+    """Exception raised when a variable is not found in the variable pool."""
+
+    def __init__(self, variable_name: str):
+        super().__init__(f"Variable '{variable_name}' does not exist", variable_name)
+
+
+class AgentInputTypeError(AgentNodeError):
+    """Exception raised when an unknown agent input type is encountered."""
+
+    def __init__(self, input_type: str):
+        super().__init__(f"Unknown agent input type '{input_type}'")
+
+
+class ToolFileError(AgentNodeError):
+    """Exception raised when there's an error with a tool file."""
+
+    def __init__(self, message: str, file_id: Optional[str] = None):
+        self.file_id = file_id
+        super().__init__(message)
+
+
+class ToolFileNotFoundError(ToolFileError):
+    """Exception raised when a tool file is not found."""
+
+    def __init__(self, file_id: str):
+        super().__init__(f"Tool file '{file_id}' does not exist", file_id)
+
+
+class AgentMessageTransformError(AgentNodeError):
+    """Exception raised when there's an error transforming agent messages."""
+
+    def __init__(self, message: str, original_error: Optional[Exception] = None):
+        self.original_error = original_error
+        super().__init__(message)
+
+
+class AgentModelError(AgentNodeError):
+    """Exception raised when there's an error with the model used by the agent."""
+
+    def __init__(self, message: str, model_name: Optional[str] = None, provider: Optional[str] = None):
+        self.model_name = model_name
+        self.provider = provider
+        super().__init__(message)
+
+
+class AgentMemoryError(AgentNodeError):
+    """Exception raised when there's an error with the agent's memory."""
+
+    def __init__(self, message: str, conversation_id: Optional[str] = None):
+        self.conversation_id = conversation_id
+        super().__init__(message)
+
+
+class AgentVariableTypeError(AgentNodeError):
+    """Exception raised when a variable has an unexpected type."""
+
+    def __init__(
+        self,
+        message: str,
+        variable_name: Optional[str] = None,
+        expected_type: Optional[str] = None,
+        actual_type: Optional[str] = None,
+    ):
+        self.variable_name = variable_name
+        self.expected_type = expected_type
+        self.actual_type = actual_type
+        super().__init__(message)

api/core/workflow/nodes/answer/answer_node.py

@@ -1,5 +1,5 @@
 from collections.abc import Mapping, Sequence
-from typing import Any, cast
+from typing import Any, Optional, cast
 
 from core.variables import ArrayFileSegment, FileSegment
 from core.workflow.entities.node_entities import NodeRunResult
@@ -12,14 +12,37 @@ from core.workflow.nodes.answer.entities import (
     VarGenerateRouteChunk,
 )
 from core.workflow.nodes.base import BaseNode
-from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
 from core.workflow.utils.variable_template_parser import VariableTemplateParser
 
 
-class AnswerNode(BaseNode[AnswerNodeData]):
-    _node_data_cls = AnswerNodeData
+class AnswerNode(BaseNode):
     _node_type = NodeType.ANSWER
 
+    _node_data: AnswerNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = AnswerNodeData.model_validate(data)
+
+    def _get_error_strategy(self) -> Optional[ErrorStrategy]:
+        return self._node_data.error_strategy
+
+    def _get_retry_config(self) -> RetryConfig:
+        return self._node_data.retry_config
+
+    def _get_title(self) -> str:
+        return self._node_data.title
+
+    def _get_description(self) -> Optional[str]:
+        return self._node_data.desc
+
+    def _get_default_value_dict(self) -> dict[str, Any]:
+        return self._node_data.default_value_dict
+
+    def get_base_node_data(self) -> BaseNodeData:
+        return self._node_data
+
     @classmethod
     def version(cls) -> str:
         return "1"
@@ -30,7 +53,7 @@ class AnswerNode(BaseNode[AnswerNodeData]):
         :return:
         """
         # generate routes
-        generate_routes = AnswerStreamGeneratorRouter.extract_generate_route_from_node_data(self.node_data)
+        generate_routes = AnswerStreamGeneratorRouter.extract_generate_route_from_node_data(self._node_data)
 
         answer = ""
         files = []
@@ -60,16 +83,12 @@ class AnswerNode(BaseNode[AnswerNodeData]):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: AnswerNodeData,
+        node_data: Mapping[str, Any],
     ) -> Mapping[str, Sequence[str]]:
-        """
-        Extract variable selector to variable mapping
-        :param graph_config: graph config
-        :param node_id: node id
-        :param node_data: node data
-        :return:
-        """
-        variable_template_parser = VariableTemplateParser(template=node_data.answer)
+        # Create typed NodeData from dict
+        typed_node_data = AnswerNodeData.model_validate(node_data)
+
+        variable_template_parser = VariableTemplateParser(template=typed_node_data.answer)
         variable_selectors = variable_template_parser.extract_variable_selectors()
 
         variable_mapping = {}

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

@@ -122,13 +122,13 @@ class RetryConfig(BaseModel):
 class BaseNodeData(ABC, BaseModel):
     title: str
     desc: Optional[str] = None
+    version: str = "1"
     error_strategy: Optional[ErrorStrategy] = None
     default_value: Optional[list[DefaultValue]] = None
-    version: str = "1"
     retry_config: RetryConfig = RetryConfig()
 
     @property
-    def default_value_dict(self):
+    def default_value_dict(self) -> dict[str, Any]:
         if self.default_value:
             return {item.key: item.value for item in self.default_value}
         return {}

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

@@ -1,28 +1,22 @@
 import logging
 from abc import abstractmethod
 from collections.abc import Generator, Mapping, Sequence
-from typing import TYPE_CHECKING, Any, ClassVar, Generic, Optional, TypeVar, Union, cast
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union
 
 from core.workflow.entities.node_entities import NodeRunResult
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
-from core.workflow.nodes.enums import CONTINUE_ON_ERROR_NODE_TYPE, RETRY_ON_ERROR_NODE_TYPE, NodeType
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
 from core.workflow.nodes.event import NodeEvent, RunCompletedEvent
 
-from .entities import BaseNodeData
-
 if TYPE_CHECKING:
+    from core.workflow.graph_engine import Graph, GraphInitParams, GraphRuntimeState
     from core.workflow.graph_engine.entities.event import InNodeEvent
-    from core.workflow.graph_engine.entities.graph import Graph
-    from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams
-    from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState
 
 logger = logging.getLogger(__name__)
 
-GenericNodeData = TypeVar("GenericNodeData", bound=BaseNodeData)
 
-
-class BaseNode(Generic[GenericNodeData]):
-    _node_data_cls: type[GenericNodeData]
+class BaseNode:
     _node_type: ClassVar[NodeType]
 
     def __init__(
@@ -56,8 +50,8 @@ class BaseNode(Generic[GenericNodeData]):
 
         self.node_id = node_id
 
-        node_data = self._node_data_cls.model_validate(config.get("data", {}))
-        self.node_data = node_data
+    @abstractmethod
+    def init_node_data(self, data: Mapping[str, Any]) -> None: ...
 
     @abstractmethod
     def _run(self) -> NodeRunResult | Generator[Union[NodeEvent, "InNodeEvent"], None, None]:
@@ -130,9 +124,9 @@ class BaseNode(Generic[GenericNodeData]):
         if not node_id:
             raise ValueError("Node ID is required when extracting variable selector to variable mapping.")
 
-        node_data = cls._node_data_cls(**config.get("data", {}))
+        # Pass raw dict data instead of creating NodeData instance
         data = cls._extract_variable_selector_to_variable_mapping(
-            graph_config=graph_config, node_id=node_id, node_data=cast(GenericNodeData, node_data)
+            graph_config=graph_config, node_id=node_id, node_data=config.get("data", {})
         )
         return data
 
@@ -142,32 +136,16 @@ class BaseNode(Generic[GenericNodeData]):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: GenericNodeData,
+        node_data: Mapping[str, Any],
     ) -> Mapping[str, Sequence[str]]:
-        """
-        Extract variable selector to variable mapping
-        :param graph_config: graph config
-        :param node_id: node id
-        :param node_data: node data
-        :return:
-        """
         return {}
 
     @classmethod
     def get_default_config(cls, filters: Optional[dict] = None) -> dict:
-        """
-        Get default config of node.
-        :param filters: filter by node config parameters.
-        :return:
-        """
         return {}
 
     @property
-    def node_type(self) -> NodeType:
-        """
-        Get node type
-        :return:
-        """
+    def type_(self) -> NodeType:
         return self._node_type
 
     @classmethod
@@ -181,19 +159,68 @@ class BaseNode(Generic[GenericNodeData]):
         raise NotImplementedError("subclasses of BaseNode must implement `version` method.")
 
     @property
-    def should_continue_on_error(self) -> bool:
-        """judge if should continue on error
-
-        Returns:
-            bool: if should continue on error
-        """
-        return self.node_data.error_strategy is not None and self.node_type in CONTINUE_ON_ERROR_NODE_TYPE
+    def continue_on_error(self) -> bool:
+        return False
 
     @property
-    def should_retry(self) -> bool:
-        """judge if should retry
+    def retry(self) -> bool:
+        return False
 
-        Returns:
-            bool: if should retry
-        """
-        return self.node_data.retry_config.retry_enabled and self.node_type in RETRY_ON_ERROR_NODE_TYPE
+    # Abstract methods that subclasses must implement to provide access
+    # to BaseNodeData properties in a type-safe way
+
+    @abstractmethod
+    def _get_error_strategy(self) -> Optional[ErrorStrategy]:
+        """Get the error strategy for this node."""
+        ...
+
+    @abstractmethod
+    def _get_retry_config(self) -> RetryConfig:
+        """Get the retry configuration for this node."""
+        ...
+
+    @abstractmethod
+    def _get_title(self) -> str:
+        """Get the node title."""
+        ...
+
+    @abstractmethod
+    def _get_description(self) -> Optional[str]:
+        """Get the node description."""
+        ...
+
+    @abstractmethod
+    def _get_default_value_dict(self) -> dict[str, Any]:
+        """Get the default values dictionary for this node."""
+        ...
+
+    @abstractmethod
+    def get_base_node_data(self) -> BaseNodeData:
+        """Get the BaseNodeData object for this node."""
+        ...
+
+    # Public interface properties that delegate to abstract methods
+    @property
+    def error_strategy(self) -> Optional[ErrorStrategy]:
+        """Get the error strategy for this node."""
+        return self._get_error_strategy()
+
+    @property
+    def retry_config(self) -> RetryConfig:
+        """Get the retry configuration for this node."""
+        return self._get_retry_config()
+
+    @property
+    def title(self) -> str:
+        """Get the node title."""
+        return self._get_title()
+
+    @property
+    def description(self) -> Optional[str]:
+        """Get the node description."""
+        return self._get_description()
+
+    @property
+    def default_value_dict(self) -> dict[str, Any]:
+        """Get the default values dictionary for this node."""
+        return self._get_default_value_dict()

api/core/workflow/nodes/code/code_node.py

@@ -11,8 +11,9 @@ from core.variables.segments import ArrayFileSegment
 from core.workflow.entities.node_entities import NodeRunResult
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
 from core.workflow.nodes.base import BaseNode
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
 from core.workflow.nodes.code.entities import CodeNodeData
-from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
 
 from .exc import (
     CodeNodeError,
@@ -21,10 +22,32 @@ from .exc import (
 )
 
 
-class CodeNode(BaseNode[CodeNodeData]):
-    _node_data_cls = CodeNodeData
+class CodeNode(BaseNode):
     _node_type = NodeType.CODE
 
+    _node_data: CodeNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = CodeNodeData.model_validate(data)
+
+    def _get_error_strategy(self) -> Optional[ErrorStrategy]:
+        return self._node_data.error_strategy
+
+    def _get_retry_config(self) -> RetryConfig:
+        return self._node_data.retry_config
+
+    def _get_title(self) -> str:
+        return self._node_data.title
+
+    def _get_description(self) -> Optional[str]:
+        return self._node_data.desc
+
+    def _get_default_value_dict(self) -> dict[str, Any]:
+        return self._node_data.default_value_dict
+
+    def get_base_node_data(self) -> BaseNodeData:
+        return self._node_data
+
     @classmethod
     def get_default_config(cls, filters: Optional[dict] = None) -> dict:
         """
@@ -47,12 +70,12 @@ class CodeNode(BaseNode[CodeNodeData]):
 
     def _run(self) -> NodeRunResult:
         # Get code language
-        code_language = self.node_data.code_language
-        code = self.node_data.code
+        code_language = self._node_data.code_language
+        code = self._node_data.code
 
         # Get variables
         variables = {}
-        for variable_selector in self.node_data.variables:
+        for variable_selector in self._node_data.variables:
             variable_name = variable_selector.variable
             variable = self.graph_runtime_state.variable_pool.get(variable_selector.value_selector)
             if isinstance(variable, ArrayFileSegment):
@@ -68,7 +91,7 @@ class CodeNode(BaseNode[CodeNodeData]):
             )
 
             # Transform result
-            result = self._transform_result(result=result, output_schema=self.node_data.outputs)
+            result = self._transform_result(result=result, output_schema=self._node_data.outputs)
         except (CodeExecutionError, CodeNodeError) as e:
             return NodeRunResult(
                 status=WorkflowNodeExecutionStatus.FAILED, inputs=variables, error=str(e), error_type=type(e).__name__
@@ -334,16 +357,20 @@ class CodeNode(BaseNode[CodeNodeData]):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: CodeNodeData,
+        node_data: Mapping[str, Any],
     ) -> Mapping[str, Sequence[str]]:
-        """
-        Extract variable selector to variable mapping
-        :param graph_config: graph config
-        :param node_id: node id
-        :param node_data: node data
-        :return:
-        """
+        # Create typed NodeData from dict
+        typed_node_data = CodeNodeData.model_validate(node_data)
+
         return {
             node_id + "." + variable_selector.variable: variable_selector.value_selector
-            for variable_selector in node_data.variables
+            for variable_selector in typed_node_data.variables
         }
+
+    @property
+    def continue_on_error(self) -> bool:
+        return self._node_data.error_strategy is not None
+
+    @property
+    def retry(self) -> bool:
+        return self._node_data.retry_config.retry_enabled

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

@@ -5,7 +5,7 @@ import logging
 import os
 import tempfile
 from collections.abc import Mapping, Sequence
-from typing import Any, cast
+from typing import Any, Optional, cast
 
 import chardet
 import docx
@@ -28,7 +28,8 @@ from core.variables.segments import ArrayStringSegment, FileSegment
 from core.workflow.entities.node_entities import NodeRunResult
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
 from core.workflow.nodes.base import BaseNode
-from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
 
 from .entities import DocumentExtractorNodeData
 from .exc import DocumentExtractorError, FileDownloadError, TextExtractionError, UnsupportedFileTypeError
@@ -36,21 +37,43 @@ from .exc import DocumentExtractorError, FileDownloadError, TextExtractionError,
 logger = logging.getLogger(__name__)
 
 
-class DocumentExtractorNode(BaseNode[DocumentExtractorNodeData]):
+class DocumentExtractorNode(BaseNode):
     """
     Extracts text content from various file types.
     Supports plain text, PDF, and DOC/DOCX files.
     """
 
-    _node_data_cls = DocumentExtractorNodeData
     _node_type = NodeType.DOCUMENT_EXTRACTOR
 
+    _node_data: DocumentExtractorNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = DocumentExtractorNodeData.model_validate(data)
+
+    def _get_error_strategy(self) -> Optional[ErrorStrategy]:
+        return self._node_data.error_strategy
+
+    def _get_retry_config(self) -> RetryConfig:
+        return self._node_data.retry_config
+
+    def _get_title(self) -> str:
+        return self._node_data.title
+
+    def _get_description(self) -> Optional[str]:
+        return self._node_data.desc
+
+    def _get_default_value_dict(self) -> dict[str, Any]:
+        return self._node_data.default_value_dict
+
+    def get_base_node_data(self) -> BaseNodeData:
+        return self._node_data
+
     @classmethod
     def version(cls) -> str:
         return "1"
 
     def _run(self):
-        variable_selector = self.node_data.variable_selector
+        variable_selector = self._node_data.variable_selector
         variable = self.graph_runtime_state.variable_pool.get(variable_selector)
 
         if variable is None:
@@ -97,16 +120,12 @@ class DocumentExtractorNode(BaseNode[DocumentExtractorNodeData]):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: DocumentExtractorNodeData,
+        node_data: Mapping[str, Any],
     ) -> Mapping[str, Sequence[str]]:
-        """
-        Extract variable selector to variable mapping
-        :param graph_config: graph config
-        :param node_id: node id
-        :param node_data: node data
-        :return:
-        """
-        return {node_id + ".files": node_data.variable_selector}
+        # Create typed NodeData from dict
+        typed_node_data = DocumentExtractorNodeData.model_validate(node_data)
+
+        return {node_id + ".files": typed_node_data.variable_selector}
 
 
 def _extract_text_by_mime_type(*, file_content: bytes, mime_type: str) -> str:

api/core/workflow/nodes/end/end_node.py

@@ -1,14 +1,40 @@
+from collections.abc import Mapping
+from typing import Any, Optional
+
 from core.workflow.entities.node_entities import NodeRunResult
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
 from core.workflow.nodes.base import BaseNode
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
 from core.workflow.nodes.end.entities import EndNodeData
-from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
 
 
-class EndNode(BaseNode[EndNodeData]):
-    _node_data_cls = EndNodeData
+class EndNode(BaseNode):
     _node_type = NodeType.END
 
+    _node_data: EndNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = EndNodeData(**data)
+
+    def _get_error_strategy(self) -> Optional[ErrorStrategy]:
+        return self._node_data.error_strategy
+
+    def _get_retry_config(self) -> RetryConfig:
+        return self._node_data.retry_config
+
+    def _get_title(self) -> str:
+        return self._node_data.title
+
+    def _get_description(self) -> Optional[str]:
+        return self._node_data.desc
+
+    def _get_default_value_dict(self) -> dict[str, Any]:
+        return self._node_data.default_value_dict
+
+    def get_base_node_data(self) -> BaseNodeData:
+        return self._node_data
+
     @classmethod
     def version(cls) -> str:
         return "1"
@@ -18,7 +44,7 @@ class EndNode(BaseNode[EndNodeData]):
         Run node
         :return:
         """
-        output_variables = self.node_data.outputs
+        output_variables = self._node_data.outputs
 
         outputs = {}
         for variable_selector in output_variables:

api/core/workflow/nodes/enums.py

@@ -35,7 +35,3 @@ class ErrorStrategy(StrEnum):
 class FailBranchSourceHandle(StrEnum):
     FAILED = "fail-branch"
     SUCCESS = "success-branch"
-
-
-CONTINUE_ON_ERROR_NODE_TYPE = [NodeType.LLM, NodeType.CODE, NodeType.TOOL, NodeType.HTTP_REQUEST]
-RETRY_ON_ERROR_NODE_TYPE = CONTINUE_ON_ERROR_NODE_TYPE

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

@@ -11,7 +11,8 @@ from core.workflow.entities.node_entities import NodeRunResult
 from core.workflow.entities.variable_entities import VariableSelector
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
 from core.workflow.nodes.base import BaseNode
-from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
 from core.workflow.nodes.http_request.executor import Executor
 from core.workflow.utils import variable_template_parser
 from factories import file_factory
@@ -32,10 +33,32 @@ HTTP_REQUEST_DEFAULT_TIMEOUT = HttpRequestNodeTimeout(
 logger = logging.getLogger(__name__)
 
 
-class HttpRequestNode(BaseNode[HttpRequestNodeData]):
-    _node_data_cls = HttpRequestNodeData
+class HttpRequestNode(BaseNode):
     _node_type = NodeType.HTTP_REQUEST
 
+    _node_data: HttpRequestNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = HttpRequestNodeData.model_validate(data)
+
+    def _get_error_strategy(self) -> Optional[ErrorStrategy]:
+        return self._node_data.error_strategy
+
+    def _get_retry_config(self) -> RetryConfig:
+        return self._node_data.retry_config
+
+    def _get_title(self) -> str:
+        return self._node_data.title
+
+    def _get_description(self) -> Optional[str]:
+        return self._node_data.desc
+
+    def _get_default_value_dict(self) -> dict[str, Any]:
+        return self._node_data.default_value_dict
+
+    def get_base_node_data(self) -> BaseNodeData:
+        return self._node_data
+
     @classmethod
     def get_default_config(cls, filters: Optional[dict[str, Any]] = None) -> dict:
         return {
@@ -69,8 +92,8 @@ class HttpRequestNode(BaseNode[HttpRequestNodeData]):
         process_data = {}
         try:
             http_executor = Executor(
-                node_data=self.node_data,
-                timeout=self._get_request_timeout(self.node_data),
+                node_data=self._node_data,
+                timeout=self._get_request_timeout(self._node_data),
                 variable_pool=self.graph_runtime_state.variable_pool,
                 max_retries=0,
             )
@@ -78,7 +101,7 @@ class HttpRequestNode(BaseNode[HttpRequestNodeData]):
 
             response = http_executor.invoke()
             files = self.extract_files(url=http_executor.url, response=response)
-            if not response.response.is_success and (self.should_continue_on_error or self.should_retry):
+            if not response.response.is_success and (self.continue_on_error or self.retry):
                 return NodeRunResult(
                     status=WorkflowNodeExecutionStatus.FAILED,
                     outputs={
@@ -131,15 +154,18 @@ class HttpRequestNode(BaseNode[HttpRequestNodeData]):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: HttpRequestNodeData,
+        node_data: Mapping[str, Any],
     ) -> Mapping[str, Sequence[str]]:
+        # Create typed NodeData from dict
+        typed_node_data = HttpRequestNodeData.model_validate(node_data)
+
         selectors: list[VariableSelector] = []
-        selectors += variable_template_parser.extract_selectors_from_template(node_data.url)
-        selectors += variable_template_parser.extract_selectors_from_template(node_data.headers)
-        selectors += variable_template_parser.extract_selectors_from_template(node_data.params)
-        if node_data.body:
-            body_type = node_data.body.type
-            data = node_data.body.data
+        selectors += variable_template_parser.extract_selectors_from_template(typed_node_data.url)
+        selectors += variable_template_parser.extract_selectors_from_template(typed_node_data.headers)
+        selectors += variable_template_parser.extract_selectors_from_template(typed_node_data.params)
+        if typed_node_data.body:
+            body_type = typed_node_data.body.type
+            data = typed_node_data.body.data
             match body_type:
                 case "binary":
                     if len(data) != 1:
@@ -217,3 +243,11 @@ class HttpRequestNode(BaseNode[HttpRequestNodeData]):
         files.append(file)
 
         return ArrayFileSegment(value=files)
+
+    @property
+    def continue_on_error(self) -> bool:
+        return self._node_data.error_strategy is not None
+
+    @property
+    def retry(self) -> bool:
+        return self._node_data.retry_config.retry_enabled

api/core/workflow/nodes/if_else/if_else_node.py

@@ -1,5 +1,5 @@
 from collections.abc import Mapping, Sequence
-from typing import Any, Literal
+from typing import Any, Literal, Optional
 
 from typing_extensions import deprecated
 
@@ -7,16 +7,39 @@ from core.workflow.entities.node_entities import NodeRunResult
 from core.workflow.entities.variable_pool import VariablePool
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
 from core.workflow.nodes.base import BaseNode
-from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
 from core.workflow.nodes.if_else.entities import IfElseNodeData
 from core.workflow.utils.condition.entities import Condition
 from core.workflow.utils.condition.processor import ConditionProcessor
 
 
-class IfElseNode(BaseNode[IfElseNodeData]):
-    _node_data_cls = IfElseNodeData
+class IfElseNode(BaseNode):
     _node_type = NodeType.IF_ELSE
 
+    _node_data: IfElseNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = IfElseNodeData.model_validate(data)
+
+    def _get_error_strategy(self) -> Optional[ErrorStrategy]:
+        return self._node_data.error_strategy
+
+    def _get_retry_config(self) -> RetryConfig:
+        return self._node_data.retry_config
+
+    def _get_title(self) -> str:
+        return self._node_data.title
+
+    def _get_description(self) -> Optional[str]:
+        return self._node_data.desc
+
+    def _get_default_value_dict(self) -> dict[str, Any]:
+        return self._node_data.default_value_dict
+
+    def get_base_node_data(self) -> BaseNodeData:
+        return self._node_data
+
     @classmethod
     def version(cls) -> str:
         return "1"
@@ -36,8 +59,8 @@ class IfElseNode(BaseNode[IfElseNodeData]):
         condition_processor = ConditionProcessor()
         try:
             # Check if the new cases structure is used
-            if self.node_data.cases:
-                for case in self.node_data.cases:
+            if self._node_data.cases:
+                for case in self._node_data.cases:
                     input_conditions, group_result, final_result = condition_processor.process_conditions(
                         variable_pool=self.graph_runtime_state.variable_pool,
                         conditions=case.conditions,
@@ -63,8 +86,8 @@ class IfElseNode(BaseNode[IfElseNodeData]):
                 input_conditions, group_result, final_result = _should_not_use_old_function(
                     condition_processor=condition_processor,
                     variable_pool=self.graph_runtime_state.variable_pool,
-                    conditions=self.node_data.conditions or [],
-                    operator=self.node_data.logical_operator or "and",
+                    conditions=self._node_data.conditions or [],
+                    operator=self._node_data.logical_operator or "and",
                 )
 
                 selected_case_id = "true" if final_result else "false"
@@ -98,10 +121,13 @@ class IfElseNode(BaseNode[IfElseNodeData]):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: IfElseNodeData,
+        node_data: Mapping[str, Any],
     ) -> Mapping[str, Sequence[str]]:
+        # Create typed NodeData from dict
+        typed_node_data = IfElseNodeData.model_validate(node_data)
+
         var_mapping: dict[str, list[str]] = {}
-        for case in node_data.cases or []:
+        for case in typed_node_data.cases or []:
             for condition in case.conditions:
                 key = "{}.#{}#".format(node_id, ".".join(condition.variable_selector))
                 var_mapping[key] = condition.variable_selector

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

@@ -36,7 +36,8 @@ from core.workflow.graph_engine.entities.event import (
 )
 from core.workflow.graph_engine.entities.graph import Graph
 from core.workflow.nodes.base import BaseNode
-from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
 from core.workflow.nodes.event import NodeEvent, RunCompletedEvent
 from core.workflow.nodes.iteration.entities import ErrorHandleMode, IterationNodeData
 from factories.variable_factory import build_segment
@@ -56,14 +57,36 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-class IterationNode(BaseNode[IterationNodeData]):
+class IterationNode(BaseNode):
     """
     Iteration Node.
     """
 
-    _node_data_cls = IterationNodeData
     _node_type = NodeType.ITERATION
 
+    _node_data: IterationNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = IterationNodeData.model_validate(data)
+
+    def _get_error_strategy(self) -> Optional[ErrorStrategy]:
+        return self._node_data.error_strategy
+
+    def _get_retry_config(self) -> RetryConfig:
+        return self._node_data.retry_config
+
+    def _get_title(self) -> str:
+        return self._node_data.title
+
+    def _get_description(self) -> Optional[str]:
+        return self._node_data.desc
+
+    def _get_default_value_dict(self) -> dict[str, Any]:
+        return self._node_data.default_value_dict
+
+    def get_base_node_data(self) -> BaseNodeData:
+        return self._node_data
+
     @classmethod
     def get_default_config(cls, filters: Optional[dict] = None) -> dict:
         return {
@@ -83,10 +106,10 @@ class IterationNode(BaseNode[IterationNodeData]):
         """
         Run the node.
         """
-        variable = self.graph_runtime_state.variable_pool.get(self.node_data.iterator_selector)
+        variable = self.graph_runtime_state.variable_pool.get(self._node_data.iterator_selector)
 
         if not variable:
-            raise IteratorVariableNotFoundError(f"iterator variable {self.node_data.iterator_selector} not found")
+            raise IteratorVariableNotFoundError(f"iterator variable {self._node_data.iterator_selector} not found")
 
         if not isinstance(variable, ArrayVariable) and not isinstance(variable, NoneVariable):
             raise InvalidIteratorValueError(f"invalid iterator value: {variable}, please provide a list.")
@@ -116,10 +139,10 @@ class IterationNode(BaseNode[IterationNodeData]):
 
         graph_config = self.graph_config
 
-        if not self.node_data.start_node_id:
+        if not self._node_data.start_node_id:
             raise StartNodeIdNotFoundError(f"field start_node_id in iteration {self.node_id} not found")
 
-        root_node_id = self.node_data.start_node_id
+        root_node_id = self._node_data.start_node_id
 
         # init graph
         iteration_graph = Graph.init(graph_config=graph_config, root_node_id=root_node_id)
@@ -161,8 +184,8 @@ class IterationNode(BaseNode[IterationNodeData]):
         yield IterationRunStartedEvent(
             iteration_id=self.id,
             iteration_node_id=self.node_id,
-            iteration_node_type=self.node_type,
-            iteration_node_data=self.node_data,
+            iteration_node_type=self.type_,
+            iteration_node_data=self._node_data,
             start_at=start_at,
             inputs=inputs,
             metadata={"iterator_length": len(iterator_list_value)},
@@ -172,8 +195,8 @@ class IterationNode(BaseNode[IterationNodeData]):
         yield IterationRunNextEvent(
             iteration_id=self.id,
             iteration_node_id=self.node_id,
-            iteration_node_type=self.node_type,
-            iteration_node_data=self.node_data,
+            iteration_node_type=self.type_,
+            iteration_node_data=self._node_data,
             index=0,
             pre_iteration_output=None,
             duration=None,
@@ -181,11 +204,11 @@ class IterationNode(BaseNode[IterationNodeData]):
         iter_run_map: dict[str, float] = {}
         outputs: list[Any] = [None] * len(iterator_list_value)
         try:
-            if self.node_data.is_parallel:
+            if self._node_data.is_parallel:
                 futures: list[Future] = []
                 q: Queue = Queue()
                 thread_pool = GraphEngineThreadPool(
-                    max_workers=self.node_data.parallel_nums, max_submit_count=dify_config.MAX_SUBMIT_COUNT
+                    max_workers=self._node_data.parallel_nums, max_submit_count=dify_config.MAX_SUBMIT_COUNT
                 )
                 for index, item in enumerate(iterator_list_value):
                     future: Future = thread_pool.submit(
@@ -242,7 +265,7 @@ class IterationNode(BaseNode[IterationNodeData]):
                         iteration_graph=iteration_graph,
                         iter_run_map=iter_run_map,
                     )
-            if self.node_data.error_handle_mode == ErrorHandleMode.REMOVE_ABNORMAL_OUTPUT:
+            if self._node_data.error_handle_mode == ErrorHandleMode.REMOVE_ABNORMAL_OUTPUT:
                 outputs = [output for output in outputs if output is not None]
 
             # Flatten the list of lists
@@ -253,8 +276,8 @@ class IterationNode(BaseNode[IterationNodeData]):
             yield IterationRunSucceededEvent(
                 iteration_id=self.id,
                 iteration_node_id=self.node_id,
-                iteration_node_type=self.node_type,
-                iteration_node_data=self.node_data,
+                iteration_node_type=self.type_,
+                iteration_node_data=self._node_data,
                 start_at=start_at,
                 inputs=inputs,
                 outputs={"output": outputs},
@@ -278,8 +301,8 @@ class IterationNode(BaseNode[IterationNodeData]):
             yield IterationRunFailedEvent(
                 iteration_id=self.id,
                 iteration_node_id=self.node_id,
-                iteration_node_type=self.node_type,
-                iteration_node_data=self.node_data,
+                iteration_node_type=self.type_,
+                iteration_node_data=self._node_data,
                 start_at=start_at,
                 inputs=inputs,
                 outputs={"output": outputs},
@@ -305,21 +328,17 @@ class IterationNode(BaseNode[IterationNodeData]):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: IterationNodeData,
+        node_data: Mapping[str, Any],
     ) -> Mapping[str, Sequence[str]]:
-        """
-        Extract variable selector to variable mapping
-        :param graph_config: graph config
-        :param node_id: node id
-        :param node_data: node data
-        :return:
-        """
+        # Create typed NodeData from dict
+        typed_node_data = IterationNodeData.model_validate(node_data)
+
         variable_mapping: dict[str, Sequence[str]] = {
-            f"{node_id}.input_selector": node_data.iterator_selector,
+            f"{node_id}.input_selector": typed_node_data.iterator_selector,
         }
 
         # init graph
-        iteration_graph = Graph.init(graph_config=graph_config, root_node_id=node_data.start_node_id)
+        iteration_graph = Graph.init(graph_config=graph_config, root_node_id=typed_node_data.start_node_id)
 
         if not iteration_graph:
             raise IterationGraphNotFoundError("iteration graph not found")
@@ -375,7 +394,7 @@ class IterationNode(BaseNode[IterationNodeData]):
         """
         if not isinstance(event, BaseNodeEvent):
             return event
-        if self.node_data.is_parallel and isinstance(event, NodeRunStartedEvent):
+        if self._node_data.is_parallel and isinstance(event, NodeRunStartedEvent):
             event.parallel_mode_run_id = parallel_mode_run_id
 
         iter_metadata = {
@@ -438,12 +457,12 @@ class IterationNode(BaseNode[IterationNodeData]):
                 elif isinstance(event, BaseGraphEvent):
                     if isinstance(event, GraphRunFailedEvent):
                         # iteration run failed
-                        if self.node_data.is_parallel:
+                        if self._node_data.is_parallel:
                             yield IterationRunFailedEvent(
                                 iteration_id=self.id,
                                 iteration_node_id=self.node_id,
-                                iteration_node_type=self.node_type,
-                                iteration_node_data=self.node_data,
+                                iteration_node_type=self.type_,
+                                iteration_node_data=self._node_data,
                                 parallel_mode_run_id=parallel_mode_run_id,
                                 start_at=start_at,
                                 inputs=inputs,
@@ -456,8 +475,8 @@ class IterationNode(BaseNode[IterationNodeData]):
                             yield IterationRunFailedEvent(
                                 iteration_id=self.id,
                                 iteration_node_id=self.node_id,
-                                iteration_node_type=self.node_type,
-                                iteration_node_data=self.node_data,
+                                iteration_node_type=self.type_,
+                                iteration_node_data=self._node_data,
                                 start_at=start_at,
                                 inputs=inputs,
                                 outputs={"output": outputs},
@@ -478,7 +497,7 @@ class IterationNode(BaseNode[IterationNodeData]):
                         event=event, iter_run_index=current_index, parallel_mode_run_id=parallel_mode_run_id
                     )
                     if isinstance(event, NodeRunFailedEvent):
-                        if self.node_data.error_handle_mode == ErrorHandleMode.CONTINUE_ON_ERROR:
+                        if self._node_data.error_handle_mode == ErrorHandleMode.CONTINUE_ON_ERROR:
                             yield NodeInIterationFailedEvent(
                                 **metadata_event.model_dump(),
                             )
@@ -491,15 +510,15 @@ class IterationNode(BaseNode[IterationNodeData]):
                             yield IterationRunNextEvent(
                                 iteration_id=self.id,
                                 iteration_node_id=self.node_id,
-                                iteration_node_type=self.node_type,
-                                iteration_node_data=self.node_data,
+                                iteration_node_type=self.type_,
+                                iteration_node_data=self._node_data,
                                 index=next_index,
                                 parallel_mode_run_id=parallel_mode_run_id,
                                 pre_iteration_output=None,
                                 duration=duration,
                             )
                             return
-                        elif self.node_data.error_handle_mode == ErrorHandleMode.REMOVE_ABNORMAL_OUTPUT:
+                        elif self._node_data.error_handle_mode == ErrorHandleMode.REMOVE_ABNORMAL_OUTPUT:
                             yield NodeInIterationFailedEvent(
                                 **metadata_event.model_dump(),
                             )
@@ -512,15 +531,15 @@ class IterationNode(BaseNode[IterationNodeData]):
                             yield IterationRunNextEvent(
                                 iteration_id=self.id,
                                 iteration_node_id=self.node_id,
-                                iteration_node_type=self.node_type,
-                                iteration_node_data=self.node_data,
+                                iteration_node_type=self.type_,
+                                iteration_node_data=self._node_data,
                                 index=next_index,
                                 parallel_mode_run_id=parallel_mode_run_id,
                                 pre_iteration_output=None,
                                 duration=duration,
                             )
                             return
-                        elif self.node_data.error_handle_mode == ErrorHandleMode.TERMINATED:
+                        elif self._node_data.error_handle_mode == ErrorHandleMode.TERMINATED:
                             yield NodeInIterationFailedEvent(
                                 **metadata_event.model_dump(),
                             )
@@ -531,12 +550,12 @@ class IterationNode(BaseNode[IterationNodeData]):
                                 variable_pool.remove([node_id])
 
                             # iteration run failed
-                            if self.node_data.is_parallel:
+                            if self._node_data.is_parallel:
                                 yield IterationRunFailedEvent(
                                     iteration_id=self.id,
                                     iteration_node_id=self.node_id,
-                                    iteration_node_type=self.node_type,
-                                    iteration_node_data=self.node_data,
+                                    iteration_node_type=self.type_,
+                                    iteration_node_data=self._node_data,
                                     parallel_mode_run_id=parallel_mode_run_id,
                                     start_at=start_at,
                                     inputs=inputs,
@@ -549,8 +568,8 @@ class IterationNode(BaseNode[IterationNodeData]):
                                 yield IterationRunFailedEvent(
                                     iteration_id=self.id,
                                     iteration_node_id=self.node_id,
-                                    iteration_node_type=self.node_type,
-                                    iteration_node_data=self.node_data,
+                                    iteration_node_type=self.type_,
+                                    iteration_node_data=self._node_data,
                                     start_at=start_at,
                                     inputs=inputs,
                                     outputs={"output": outputs},
@@ -569,7 +588,7 @@ class IterationNode(BaseNode[IterationNodeData]):
                             return
                     yield metadata_event
 
-            current_output_segment = variable_pool.get(self.node_data.output_selector)
+            current_output_segment = variable_pool.get(self._node_data.output_selector)
             if current_output_segment is None:
                 raise IterationNodeError("iteration output selector not found")
             current_iteration_output = current_output_segment.value
@@ -588,8 +607,8 @@ class IterationNode(BaseNode[IterationNodeData]):
             yield IterationRunNextEvent(
                 iteration_id=self.id,
                 iteration_node_id=self.node_id,
-                iteration_node_type=self.node_type,
-                iteration_node_data=self.node_data,
+                iteration_node_type=self.type_,
+                iteration_node_data=self._node_data,
                 index=next_index,
                 parallel_mode_run_id=parallel_mode_run_id,
                 pre_iteration_output=current_iteration_output or None,
@@ -601,8 +620,8 @@ class IterationNode(BaseNode[IterationNodeData]):
             yield IterationRunFailedEvent(
                 iteration_id=self.id,
                 iteration_node_id=self.node_id,
-                iteration_node_type=self.node_type,
-                iteration_node_data=self.node_data,
+                iteration_node_type=self.type_,
+                iteration_node_data=self._node_data,
                 start_at=start_at,
                 inputs=inputs,
                 outputs={"output": None},

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

@@ -1,18 +1,44 @@
+from collections.abc import Mapping
+from typing import Any, Optional
+
 from core.workflow.entities.node_entities import NodeRunResult
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
 from core.workflow.nodes.base import BaseNode
-from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
 from core.workflow.nodes.iteration.entities import IterationStartNodeData
 
 
-class IterationStartNode(BaseNode[IterationStartNodeData]):
+class IterationStartNode(BaseNode):
     """
     Iteration Start Node.
     """
 
-    _node_data_cls = IterationStartNodeData
     _node_type = NodeType.ITERATION_START
 
+    _node_data: IterationStartNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = IterationStartNodeData(**data)
+
+    def _get_error_strategy(self) -> Optional[ErrorStrategy]:
+        return self._node_data.error_strategy
+
+    def _get_retry_config(self) -> RetryConfig:
+        return self._node_data.retry_config
+
+    def _get_title(self) -> str:
+        return self._node_data.title
+
+    def _get_description(self) -> Optional[str]:
+        return self._node_data.desc
+
+    def _get_default_value_dict(self) -> dict[str, Any]:
+        return self._node_data.default_value_dict
+
+    def get_base_node_data(self) -> BaseNodeData:
+        return self._node_data
+
     @classmethod
     def version(cls) -> str:
         return "1"

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

@@ -1,10 +1,10 @@
 from collections.abc import Sequence
-from typing import Any, Literal, Optional
+from typing import Literal, Optional
 
 from pydantic import BaseModel, Field
 
 from core.workflow.nodes.base import BaseNodeData
-from core.workflow.nodes.llm.entities import VisionConfig
+from core.workflow.nodes.llm.entities import ModelConfig, VisionConfig
 
 
 class RerankingModelConfig(BaseModel):
@@ -56,17 +56,6 @@ class MultipleRetrievalConfig(BaseModel):
     weights: Optional[WeightedScoreConfig] = None
 
 
-class ModelConfig(BaseModel):
-    """
-    Model Config.
-    """
-
-    provider: str
-    name: str
-    mode: str
-    completion_params: dict[str, Any] = {}
-
-
 class SingleRetrievalConfig(BaseModel):
     """
     Single Retrieval Config.

api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py

@@ -4,7 +4,7 @@ import re
 import time
 from collections import defaultdict
 from collections.abc import Mapping, Sequence
-from typing import Any, Optional, cast
+from typing import TYPE_CHECKING, Any, Optional, cast
 
 from sqlalchemy import Float, and_, func, or_, text
 from sqlalchemy import cast as sqlalchemy_cast
@@ -15,20 +15,31 @@ from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEnti
 from core.entities.agent_entities import PlanningStrategy
 from core.entities.model_entities import ModelStatus
 from core.model_manager import ModelInstance, ModelManager
-from core.model_runtime.entities.message_entities import PromptMessageRole
-from core.model_runtime.entities.model_entities import ModelFeature, ModelType
+from core.model_runtime.entities.message_entities import (
+    PromptMessageRole,
+)
+from core.model_runtime.entities.model_entities import (
+    ModelFeature,
+    ModelType,
+)
 from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
 from core.prompt.simple_prompt_transform import ModelMode
 from core.rag.datasource.retrieval_service import RetrievalService
 from core.rag.entities.metadata_entities import Condition, MetadataCondition
 from core.rag.retrieval.dataset_retrieval import DatasetRetrieval
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
-from core.variables import StringSegment
+from core.variables import (
+    StringSegment,
+)
 from core.variables.segments import ArrayObjectSegment
 from core.workflow.entities.node_entities import NodeRunResult
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
-from core.workflow.nodes.enums import NodeType
-from core.workflow.nodes.event.event import ModelInvokeCompletedEvent
+from core.workflow.nodes.base import BaseNode
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
+from core.workflow.nodes.event import (
+    ModelInvokeCompletedEvent,
+)
 from core.workflow.nodes.knowledge_retrieval.template_prompts import (
     METADATA_FILTER_ASSISTANT_PROMPT_1,
     METADATA_FILTER_ASSISTANT_PROMPT_2,
@@ -38,7 +49,8 @@ from core.workflow.nodes.knowledge_retrieval.template_prompts import (
     METADATA_FILTER_USER_PROMPT_2,
     METADATA_FILTER_USER_PROMPT_3,
 )
-from core.workflow.nodes.llm.entities import LLMNodeChatModelMessage, LLMNodeCompletionModelPromptTemplate
+from core.workflow.nodes.llm.entities import LLMNodeChatModelMessage, LLMNodeCompletionModelPromptTemplate, ModelConfig
+from core.workflow.nodes.llm.file_saver import FileSaverImpl, LLMFileSaver
 from core.workflow.nodes.llm.node import LLMNode
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
@@ -46,7 +58,7 @@ from libs.json_in_md_parser import parse_and_check_json_markdown
 from models.dataset import Dataset, DatasetMetadata, Document, RateLimitLog
 from services.feature_service import FeatureService
 
-from .entities import KnowledgeRetrievalNodeData, ModelConfig
+from .entities import KnowledgeRetrievalNodeData
 from .exc import (
     InvalidModelTypeError,
     KnowledgeRetrievalNodeError,
@@ -56,6 +68,10 @@ from .exc import (
     ModelQuotaExceededError,
 )
 
+if TYPE_CHECKING:
+    from core.file.models import File
+    from core.workflow.graph_engine import Graph, GraphInitParams, GraphRuntimeState
+
 logger = logging.getLogger(__name__)
 
 default_retrieval_model = {
@@ -67,18 +83,76 @@ default_retrieval_model = {
 }
 
 
-class KnowledgeRetrievalNode(LLMNode):
-    _node_data_cls = KnowledgeRetrievalNodeData  # type: ignore
+class KnowledgeRetrievalNode(BaseNode):
     _node_type = NodeType.KNOWLEDGE_RETRIEVAL
 
+    _node_data: KnowledgeRetrievalNodeData
+
+    # Instance attributes specific to LLMNode.
+    # Output variable for file
+    _file_outputs: list["File"]
+
+    _llm_file_saver: LLMFileSaver
+
+    def __init__(
+        self,
+        id: str,
+        config: Mapping[str, Any],
+        graph_init_params: "GraphInitParams",
+        graph: "Graph",
+        graph_runtime_state: "GraphRuntimeState",
+        previous_node_id: Optional[str] = None,
+        thread_pool_id: Optional[str] = None,
+        *,
+        llm_file_saver: LLMFileSaver | None = None,
+    ) -> None:
+        super().__init__(
+            id=id,
+            config=config,
+            graph_init_params=graph_init_params,
+            graph=graph,
+            graph_runtime_state=graph_runtime_state,
+            previous_node_id=previous_node_id,
+            thread_pool_id=thread_pool_id,
+        )
+        # LLM file outputs, used for MultiModal outputs.
+        self._file_outputs: list[File] = []
+
+        if llm_file_saver is None:
+            llm_file_saver = FileSaverImpl(
+                user_id=graph_init_params.user_id,
+                tenant_id=graph_init_params.tenant_id,
+            )
+        self._llm_file_saver = llm_file_saver
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = KnowledgeRetrievalNodeData.model_validate(data)
+
+    def _get_error_strategy(self) -> Optional[ErrorStrategy]:
+        return self._node_data.error_strategy
+
+    def _get_retry_config(self) -> RetryConfig:
+        return self._node_data.retry_config
+
+    def _get_title(self) -> str:
+        return self._node_data.title
+
+    def _get_description(self) -> Optional[str]:
+        return self._node_data.desc
+
+    def _get_default_value_dict(self) -> dict[str, Any]:
+        return self._node_data.default_value_dict
+
+    def get_base_node_data(self) -> BaseNodeData:
+        return self._node_data
+
     @classmethod
     def version(cls):
         return "1"
 
     def _run(self) -> NodeRunResult:  # type: ignore
-        node_data = cast(KnowledgeRetrievalNodeData, self.node_data)
         # extract variables
-        variable = self.graph_runtime_state.variable_pool.get(node_data.query_variable_selector)
+        variable = self.graph_runtime_state.variable_pool.get(self._node_data.query_variable_selector)
         if not isinstance(variable, StringSegment):
             return NodeRunResult(
                 status=WorkflowNodeExecutionStatus.FAILED,
@@ -119,7 +193,7 @@ class KnowledgeRetrievalNode(LLMNode):
 
         # retrieve knowledge
         try:
-            results = self._fetch_dataset_retriever(node_data=node_data, query=query)
+            results = self._fetch_dataset_retriever(node_data=self._node_data, query=query)
             outputs = {"result": ArrayObjectSegment(value=results)}
             return NodeRunResult(
                 status=WorkflowNodeExecutionStatus.SUCCEEDED,
@@ -388,7 +462,7 @@ class KnowledgeRetrievalNode(LLMNode):
                                 expected_value = self.graph_runtime_state.variable_pool.convert_template(
                                     expected_value
                                 ).value[0]
-                                if expected_value.value_type == "number":  # type: ignore
+                                if expected_value.value_type in {"number", "integer", "float"}:  # type: ignore
                                     expected_value = expected_value.value  # type: ignore
                                 elif expected_value.value_type == "string":  # type: ignore
                                     expected_value = re.sub(r"[\r\n\t]+", " ", expected_value.text).strip()  # type: ignore
@@ -435,20 +509,17 @@ class KnowledgeRetrievalNode(LLMNode):
         # get all metadata field
         metadata_fields = db.session.query(DatasetMetadata).filter(DatasetMetadata.dataset_id.in_(dataset_ids)).all()
         all_metadata_fields = [metadata_field.name for metadata_field in metadata_fields]
-        # get metadata model config
-        metadata_model_config = node_data.metadata_model_config
-        if metadata_model_config is None:
+        if node_data.metadata_model_config is None:
             raise ValueError("metadata_model_config is required")
-        # get metadata model instance
-        # fetch model config
-        model_instance, model_config = self.get_model_config(metadata_model_config)
+        # get metadata model instance and fetch model config
+        model_instance, model_config = self.get_model_config(node_data.metadata_model_config)
         # fetch prompt messages
         prompt_template = self._get_prompt_template(
             node_data=node_data,
             metadata_fields=all_metadata_fields,
             query=query or "",
         )
-        prompt_messages, stop = self._fetch_prompt_messages(
+        prompt_messages, stop = LLMNode.fetch_prompt_messages(
             prompt_template=prompt_template,
             sys_query=query,
             memory=None,
@@ -458,16 +529,23 @@ class KnowledgeRetrievalNode(LLMNode):
             vision_detail=node_data.vision.configs.detail,
             variable_pool=self.graph_runtime_state.variable_pool,
             jinja2_variables=[],
+            tenant_id=self.tenant_id,
         )
 
         result_text = ""
         try:
             # handle invoke result
-            generator = self._invoke_llm(
-                node_data_model=node_data.metadata_model_config,  # type: ignore
+            generator = LLMNode.invoke_llm(
+                node_data_model=node_data.metadata_model_config,
                 model_instance=model_instance,
                 prompt_messages=prompt_messages,
                 stop=stop,
+                user_id=self.user_id,
+                structured_output_enabled=self._node_data.structured_output_enabled,
+                structured_output=None,
+                file_saver=self._llm_file_saver,
+                file_outputs=self._file_outputs,
+                node_id=self.node_id,
             )
 
             for event in generator:
@@ -557,17 +635,13 @@ class KnowledgeRetrievalNode(LLMNode):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: KnowledgeRetrievalNodeData,  # type: ignore
+        node_data: Mapping[str, Any],
     ) -> Mapping[str, Sequence[str]]:
-        """
-        Extract variable selector to variable mapping
-        :param graph_config: graph config
-        :param node_id: node id
-        :param node_data: node data
-        :return:
-        """
+        # Create typed NodeData from dict
+        typed_node_data = KnowledgeRetrievalNodeData.model_validate(node_data)
+
         variable_mapping = {}
-        variable_mapping[node_id + ".query"] = node_data.query_variable_selector
+        variable_mapping[node_id + ".query"] = typed_node_data.query_variable_selector
         return variable_mapping
 
     def get_model_config(self, model: ModelConfig) -> tuple[ModelInstance, ModelConfigWithCredentialsEntity]:
@@ -629,7 +703,7 @@ class KnowledgeRetrievalNode(LLMNode):
         )
 
     def _get_prompt_template(self, node_data: KnowledgeRetrievalNodeData, metadata_fields: list, query: str):
-        model_mode = ModelMode.value_of(node_data.metadata_model_config.mode)  # type: ignore
+        model_mode = ModelMode(node_data.metadata_model_config.mode)  # type: ignore
         input_text = query
 
         prompt_messages: list[LLMNodeChatModelMessage] = []

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

@@ -1,5 +1,5 @@
-from collections.abc import Callable, Sequence
-from typing import Any, Literal, Union
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any, Literal, Optional, Union
 
 from core.file import File
 from core.variables import ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment
@@ -7,16 +7,39 @@ from core.variables.segments import ArrayAnySegment, ArraySegment
 from core.workflow.entities.node_entities import NodeRunResult
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
 from core.workflow.nodes.base import BaseNode
-from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
 
 from .entities import ListOperatorNodeData
 from .exc import InvalidConditionError, InvalidFilterValueError, InvalidKeyError, ListOperatorError
 
 
-class ListOperatorNode(BaseNode[ListOperatorNodeData]):
-    _node_data_cls = ListOperatorNodeData
+class ListOperatorNode(BaseNode):
     _node_type = NodeType.LIST_OPERATOR
 
+    _node_data: ListOperatorNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = ListOperatorNodeData(**data)
+
+    def _get_error_strategy(self) -> Optional[ErrorStrategy]:
+        return self._node_data.error_strategy
+
+    def _get_retry_config(self) -> RetryConfig:
+        return self._node_data.retry_config
+
+    def _get_title(self) -> str:
+        return self._node_data.title
+
+    def _get_description(self) -> Optional[str]:
+        return self._node_data.desc
+
+    def _get_default_value_dict(self) -> dict[str, Any]:
+        return self._node_data.default_value_dict
+
+    def get_base_node_data(self) -> BaseNodeData:
+        return self._node_data
+
     @classmethod
     def version(cls) -> str:
         return "1"
@@ -26,9 +49,9 @@ class ListOperatorNode(BaseNode[ListOperatorNodeData]):
         process_data: dict[str, list] = {}
         outputs: dict[str, Any] = {}
 
-        variable = self.graph_runtime_state.variable_pool.get(self.node_data.variable)
+        variable = self.graph_runtime_state.variable_pool.get(self._node_data.variable)
         if variable is None:
-            error_message = f"Variable not found for selector: {self.node_data.variable}"
+            error_message = f"Variable not found for selector: {self._node_data.variable}"
             return NodeRunResult(
                 status=WorkflowNodeExecutionStatus.FAILED, error=error_message, inputs=inputs, outputs=outputs
             )
@@ -48,7 +71,7 @@ class ListOperatorNode(BaseNode[ListOperatorNodeData]):
             )
         if not isinstance(variable, ArrayFileSegment | ArrayNumberSegment | ArrayStringSegment):
             error_message = (
-                f"Variable {self.node_data.variable} is not an ArrayFileSegment, ArrayNumberSegment "
+                f"Variable {self._node_data.variable} is not an ArrayFileSegment, ArrayNumberSegment "
                 "or ArrayStringSegment"
             )
             return NodeRunResult(
@@ -64,19 +87,19 @@ class ListOperatorNode(BaseNode[ListOperatorNodeData]):
 
         try:
             # Filter
-            if self.node_data.filter_by.enabled:
+            if self._node_data.filter_by.enabled:
                 variable = self._apply_filter(variable)
 
             # Extract
-            if self.node_data.extract_by.enabled:
+            if self._node_data.extract_by.enabled:
                 variable = self._extract_slice(variable)
 
             # Order
-            if self.node_data.order_by.enabled:
+            if self._node_data.order_by.enabled:
                 variable = self._apply_order(variable)
 
             # Slice
-            if self.node_data.limit.enabled:
+            if self._node_data.limit.enabled:
                 variable = self._apply_slice(variable)
 
             outputs = {
@@ -104,7 +127,7 @@ class ListOperatorNode(BaseNode[ListOperatorNodeData]):
     ) -> Union[ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment]:
         filter_func: Callable[[Any], bool]
         result: list[Any] = []
-        for condition in self.node_data.filter_by.conditions:
+        for condition in self._node_data.filter_by.conditions:
             if isinstance(variable, ArrayStringSegment):
                 if not isinstance(condition.value, str):
                     raise InvalidFilterValueError(f"Invalid filter value: {condition.value}")
@@ -137,14 +160,14 @@ class ListOperatorNode(BaseNode[ListOperatorNodeData]):
         self, variable: Union[ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment]
     ) -> Union[ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment]:
         if isinstance(variable, ArrayStringSegment):
-            result = _order_string(order=self.node_data.order_by.value, array=variable.value)
+            result = _order_string(order=self._node_data.order_by.value, array=variable.value)
             variable = variable.model_copy(update={"value": result})
         elif isinstance(variable, ArrayNumberSegment):
-            result = _order_number(order=self.node_data.order_by.value, array=variable.value)
+            result = _order_number(order=self._node_data.order_by.value, array=variable.value)
             variable = variable.model_copy(update={"value": result})
         elif isinstance(variable, ArrayFileSegment):
             result = _order_file(
-                order=self.node_data.order_by.value, order_by=self.node_data.order_by.key, array=variable.value
+                order=self._node_data.order_by.value, order_by=self._node_data.order_by.key, array=variable.value
             )
             variable = variable.model_copy(update={"value": result})
         return variable
@@ -152,13 +175,13 @@ class ListOperatorNode(BaseNode[ListOperatorNodeData]):
     def _apply_slice(
         self, variable: Union[ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment]
     ) -> Union[ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment]:
-        result = variable.value[: self.node_data.limit.size]
+        result = variable.value[: self._node_data.limit.size]
         return variable.model_copy(update={"value": result})
 
     def _extract_slice(
         self, variable: Union[ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment]
     ) -> Union[ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment]:
-        value = int(self.graph_runtime_state.variable_pool.convert_template(self.node_data.extract_by.serial).text)
+        value = int(self.graph_runtime_state.variable_pool.convert_template(self._node_data.extract_by.serial).text)
         if value < 1:
             raise ValueError(f"Invalid serial index: must be >= 1, got {value}")
         value -= 1

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

@@ -1,4 +1,4 @@
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from typing import Any, Optional
 
 from pydantic import BaseModel, Field, field_validator
@@ -65,7 +65,7 @@ class LLMNodeData(BaseNodeData):
     memory: Optional[MemoryConfig] = None
     context: ContextConfig
     vision: VisionConfig = Field(default_factory=VisionConfig)
-    structured_output: dict | None = None
+    structured_output: Mapping[str, Any] | None = None
     # We used 'structured_output_enabled' in the past, but it's not a good name.
     structured_output_switch_on: bool = Field(False, alias="structured_output_enabled")
 

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

@@ -59,7 +59,8 @@ from core.workflow.entities.workflow_node_execution import WorkflowNodeExecution
 from core.workflow.enums import SystemVariableKey
 from core.workflow.graph_engine.entities.event import InNodeEvent
 from core.workflow.nodes.base import BaseNode
-from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
 from core.workflow.nodes.event import (
     ModelInvokeCompletedEvent,
     NodeEvent,
@@ -90,17 +91,16 @@ from .file_saver import FileSaverImpl, LLMFileSaver
 
 if TYPE_CHECKING:
     from core.file.models import File
-    from core.workflow.graph_engine.entities.graph import Graph
-    from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams
-    from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState
+    from core.workflow.graph_engine import Graph, GraphInitParams, GraphRuntimeState
 
 logger = logging.getLogger(__name__)
 
 
-class LLMNode(BaseNode[LLMNodeData]):
-    _node_data_cls = LLMNodeData
+class LLMNode(BaseNode):
     _node_type = NodeType.LLM
 
+    _node_data: LLMNodeData
+
     # Instance attributes specific to LLMNode.
     # Output variable for file
     _file_outputs: list["File"]
@@ -138,6 +138,27 @@ class LLMNode(BaseNode[LLMNodeData]):
             )
         self._llm_file_saver = llm_file_saver
 
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = LLMNodeData.model_validate(data)
+
+    def _get_error_strategy(self) -> Optional[ErrorStrategy]:
+        return self._node_data.error_strategy
+
+    def _get_retry_config(self) -> RetryConfig:
+        return self._node_data.retry_config
+
+    def _get_title(self) -> str:
+        return self._node_data.title
+
+    def _get_description(self) -> Optional[str]:
+        return self._node_data.desc
+
+    def _get_default_value_dict(self) -> dict[str, Any]:
+        return self._node_data.default_value_dict
+
+    def get_base_node_data(self) -> BaseNodeData:
+        return self._node_data
+
     @classmethod
     def version(cls) -> str:
         return "1"
@@ -152,13 +173,13 @@ class LLMNode(BaseNode[LLMNodeData]):
 
         try:
             # init messages template
-            self.node_data.prompt_template = self._transform_chat_messages(self.node_data.prompt_template)
+            self._node_data.prompt_template = self._transform_chat_messages(self._node_data.prompt_template)
 
             # fetch variables and fetch values from variable pool
-            inputs = self._fetch_inputs(node_data=self.node_data)
+            inputs = self._fetch_inputs(node_data=self._node_data)
 
             # fetch jinja2 inputs
-            jinja_inputs = self._fetch_jinja_inputs(node_data=self.node_data)
+            jinja_inputs = self._fetch_jinja_inputs(node_data=self._node_data)
 
             # merge inputs
             inputs.update(jinja_inputs)
@@ -169,9 +190,9 @@ class LLMNode(BaseNode[LLMNodeData]):
             files = (
                 llm_utils.fetch_files(
                     variable_pool=variable_pool,
-                    selector=self.node_data.vision.configs.variable_selector,
+                    selector=self._node_data.vision.configs.variable_selector,
                 )
-                if self.node_data.vision.enabled
+                if self._node_data.vision.enabled
                 else []
             )
 
@@ -179,7 +200,7 @@ class LLMNode(BaseNode[LLMNodeData]):
                 node_inputs["#files#"] = [file.to_dict() for file in files]
 
             # fetch context value
-            generator = self._fetch_context(node_data=self.node_data)
+            generator = self._fetch_context(node_data=self._node_data)
             context = None
             for event in generator:
                 if isinstance(event, RunRetrieverResourceEvent):
@@ -189,44 +210,54 @@ class LLMNode(BaseNode[LLMNodeData]):
                 node_inputs["#context#"] = context
 
             # fetch model config
-            model_instance, model_config = self._fetch_model_config(self.node_data.model)
+            model_instance, model_config = LLMNode._fetch_model_config(
+                node_data_model=self._node_data.model,
+                tenant_id=self.tenant_id,
+            )
 
             # fetch memory
             memory = llm_utils.fetch_memory(
                 variable_pool=variable_pool,
                 app_id=self.app_id,
-                node_data_memory=self.node_data.memory,
+                node_data_memory=self._node_data.memory,
                 model_instance=model_instance,
             )
 
             query = None
-            if self.node_data.memory:
-                query = self.node_data.memory.query_prompt_template
+            if self._node_data.memory:
+                query = self._node_data.memory.query_prompt_template
                 if not query and (
                     query_variable := variable_pool.get((SYSTEM_VARIABLE_NODE_ID, SystemVariableKey.QUERY))
                 ):
                     query = query_variable.text
 
-            prompt_messages, stop = self._fetch_prompt_messages(
+            prompt_messages, stop = LLMNode.fetch_prompt_messages(
                 sys_query=query,
                 sys_files=files,
                 context=context,
                 memory=memory,
                 model_config=model_config,
-                prompt_template=self.node_data.prompt_template,
-                memory_config=self.node_data.memory,
-                vision_enabled=self.node_data.vision.enabled,
-                vision_detail=self.node_data.vision.configs.detail,
+                prompt_template=self._node_data.prompt_template,
+                memory_config=self._node_data.memory,
+                vision_enabled=self._node_data.vision.enabled,
+                vision_detail=self._node_data.vision.configs.detail,
                 variable_pool=variable_pool,
-                jinja2_variables=self.node_data.prompt_config.jinja2_variables,
+                jinja2_variables=self._node_data.prompt_config.jinja2_variables,
+                tenant_id=self.tenant_id,
             )
 
             # handle invoke result
-            generator = self._invoke_llm(
-                node_data_model=self.node_data.model,
+            generator = LLMNode.invoke_llm(
+                node_data_model=self._node_data.model,
                 model_instance=model_instance,
                 prompt_messages=prompt_messages,
                 stop=stop,
+                user_id=self.user_id,
+                structured_output_enabled=self._node_data.structured_output_enabled,
+                structured_output=self._node_data.structured_output,
+                file_saver=self._llm_file_saver,
+                file_outputs=self._file_outputs,
+                node_id=self.node_id,
             )
 
             structured_output: LLMStructuredOutput | None = None
@@ -296,12 +327,19 @@ class LLMNode(BaseNode[LLMNodeData]):
                 )
             )
 
-    def _invoke_llm(
-        self,
+    @staticmethod
+    def invoke_llm(
+        *,
         node_data_model: ModelConfig,
         model_instance: ModelInstance,
         prompt_messages: Sequence[PromptMessage],
         stop: Optional[Sequence[str]] = None,
+        user_id: str,
+        structured_output_enabled: bool,
+        structured_output: Optional[Mapping[str, Any]] = None,
+        file_saver: LLMFileSaver,
+        file_outputs: list["File"],
+        node_id: str,
     ) -> Generator[NodeEvent | LLMStructuredOutput, None, None]:
         model_schema = model_instance.model_type_instance.get_model_schema(
             node_data_model.name, model_instance.credentials
@@ -309,8 +347,10 @@ class LLMNode(BaseNode[LLMNodeData]):
         if not model_schema:
             raise ValueError(f"Model schema not found for {node_data_model.name}")
 
-        if self.node_data.structured_output_enabled:
-            output_schema = self._fetch_structured_output_schema()
+        if structured_output_enabled:
+            output_schema = LLMNode.fetch_structured_output_schema(
+                structured_output=structured_output or {},
+            )
             invoke_result = invoke_llm_with_structured_output(
                 provider=model_instance.provider,
                 model_schema=model_schema,
@@ -320,7 +360,7 @@ class LLMNode(BaseNode[LLMNodeData]):
                 model_parameters=node_data_model.completion_params,
                 stop=list(stop or []),
                 stream=True,
-                user=self.user_id,
+                user=user_id,
             )
         else:
             invoke_result = model_instance.invoke_llm(
@@ -328,17 +368,31 @@ class LLMNode(BaseNode[LLMNodeData]):
                 model_parameters=node_data_model.completion_params,
                 stop=list(stop or []),
                 stream=True,
-                user=self.user_id,
+                user=user_id,
             )
 
-        return self._handle_invoke_result(invoke_result=invoke_result)
+        return LLMNode.handle_invoke_result(
+            invoke_result=invoke_result,
+            file_saver=file_saver,
+            file_outputs=file_outputs,
+            node_id=node_id,
+        )
 
-    def _handle_invoke_result(
-        self, invoke_result: LLMResult | Generator[LLMResultChunk | LLMStructuredOutput, None, None]
+    @staticmethod
+    def handle_invoke_result(
+        *,
+        invoke_result: LLMResult | Generator[LLMResultChunk | LLMStructuredOutput, None, None],
+        file_saver: LLMFileSaver,
+        file_outputs: list["File"],
+        node_id: str,
     ) -> Generator[NodeEvent | LLMStructuredOutput, None, None]:
         # For blocking mode
         if isinstance(invoke_result, LLMResult):
-            event = self._handle_blocking_result(invoke_result=invoke_result)
+            event = LLMNode.handle_blocking_result(
+                invoke_result=invoke_result,
+                saver=file_saver,
+                file_outputs=file_outputs,
+            )
             yield event
             return
 
@@ -356,11 +410,13 @@ class LLMNode(BaseNode[LLMNodeData]):
                     yield result
                 if isinstance(result, LLMResultChunk):
                     contents = result.delta.message.content
-                    for text_part in self._save_multimodal_output_and_convert_result_to_markdown(contents):
+                    for text_part in LLMNode._save_multimodal_output_and_convert_result_to_markdown(
+                        contents=contents,
+                        file_saver=file_saver,
+                        file_outputs=file_outputs,
+                    ):
                         full_text_buffer.write(text_part)
-                        yield RunStreamChunkEvent(
-                            chunk_content=text_part, from_variable_selector=[self.node_id, "text"]
-                        )
+                        yield RunStreamChunkEvent(chunk_content=text_part, from_variable_selector=[node_id, "text"])
 
                     # Update the whole metadata
                     if not model and result.model:
@@ -378,7 +434,8 @@ class LLMNode(BaseNode[LLMNodeData]):
 
         yield ModelInvokeCompletedEvent(text=full_text_buffer.getvalue(), usage=usage, finish_reason=finish_reason)
 
-    def _image_file_to_markdown(self, file: "File", /):
+    @staticmethod
+    def _image_file_to_markdown(file: "File", /):
         text_chunk = f"![]({file.generate_url()})"
         return text_chunk
 
@@ -508,7 +565,7 @@ class LLMNode(BaseNode[LLMNodeData]):
                     retriever_resources=original_retriever_resource, context=context_str.strip()
                 )
 
-    def _convert_to_original_retriever_resource(self, context_dict: dict):
+    def _convert_to_original_retriever_resource(self, context_dict: dict) -> RetrievalSourceMetadata | None:
         if (
             "metadata" in context_dict
             and "_source" in context_dict["metadata"]
@@ -539,11 +596,14 @@ class LLMNode(BaseNode[LLMNodeData]):
 
         return None
 
+    @staticmethod
     def _fetch_model_config(
-        self, node_data_model: ModelConfig
+        *,
+        node_data_model: ModelConfig,
+        tenant_id: str,
     ) -> tuple[ModelInstance, ModelConfigWithCredentialsEntity]:
         model, model_config_with_cred = llm_utils.fetch_model_config(
-            tenant_id=self.tenant_id, node_data_model=node_data_model
+            tenant_id=tenant_id, node_data_model=node_data_model
         )
         completion_params = model_config_with_cred.parameters
 
@@ -556,8 +616,8 @@ class LLMNode(BaseNode[LLMNodeData]):
         node_data_model.completion_params = completion_params
         return model, model_config_with_cred
 
-    def _fetch_prompt_messages(
-        self,
+    @staticmethod
+    def fetch_prompt_messages(
         *,
         sys_query: str | None = None,
         sys_files: Sequence["File"],
@@ -570,13 +630,14 @@ class LLMNode(BaseNode[LLMNodeData]):
         vision_detail: ImagePromptMessageContent.DETAIL,
         variable_pool: VariablePool,
         jinja2_variables: Sequence[VariableSelector],
+        tenant_id: str,
     ) -> tuple[Sequence[PromptMessage], Optional[Sequence[str]]]:
         prompt_messages: list[PromptMessage] = []
 
         if isinstance(prompt_template, list):
             # For chat model
             prompt_messages.extend(
-                self._handle_list_messages(
+                LLMNode.handle_list_messages(
                     messages=prompt_template,
                     context=context,
                     jinja2_variables=jinja2_variables,
@@ -602,7 +663,7 @@ class LLMNode(BaseNode[LLMNodeData]):
                     edition_type="basic",
                 )
                 prompt_messages.extend(
-                    self._handle_list_messages(
+                    LLMNode.handle_list_messages(
                         messages=[message],
                         context="",
                         jinja2_variables=[],
@@ -731,7 +792,7 @@ class LLMNode(BaseNode[LLMNodeData]):
             )
 
         model = ModelManager().get_model_instance(
-            tenant_id=self.tenant_id,
+            tenant_id=tenant_id,
             model_type=ModelType.LLM,
             provider=model_config.provider,
             model=model_config.model,
@@ -750,10 +811,12 @@ class LLMNode(BaseNode[LLMNodeData]):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: LLMNodeData,
+        node_data: Mapping[str, Any],
     ) -> Mapping[str, Sequence[str]]:
-        prompt_template = node_data.prompt_template
+        # Create typed NodeData from dict
+        typed_node_data = LLMNodeData.model_validate(node_data)
 
+        prompt_template = typed_node_data.prompt_template
         variable_selectors = []
         if isinstance(prompt_template, list) and all(
             isinstance(prompt, LLMNodeChatModelMessage) for prompt in prompt_template
@@ -773,7 +836,7 @@ class LLMNode(BaseNode[LLMNodeData]):
         for variable_selector in variable_selectors:
             variable_mapping[variable_selector.variable] = variable_selector.value_selector
 
-        memory = node_data.memory
+        memory = typed_node_data.memory
         if memory and memory.query_prompt_template:
             query_variable_selectors = VariableTemplateParser(
                 template=memory.query_prompt_template
@@ -781,16 +844,16 @@ class LLMNode(BaseNode[LLMNodeData]):
             for variable_selector in query_variable_selectors:
                 variable_mapping[variable_selector.variable] = variable_selector.value_selector
 
-        if node_data.context.enabled:
-            variable_mapping["#context#"] = node_data.context.variable_selector
+        if typed_node_data.context.enabled:
+            variable_mapping["#context#"] = typed_node_data.context.variable_selector
 
-        if node_data.vision.enabled:
-            variable_mapping["#files#"] = node_data.vision.configs.variable_selector
+        if typed_node_data.vision.enabled:
+            variable_mapping["#files#"] = typed_node_data.vision.configs.variable_selector
 
-        if node_data.memory:
+        if typed_node_data.memory:
             variable_mapping["#sys.query#"] = ["sys", SystemVariableKey.QUERY.value]
 
-        if node_data.prompt_config:
+        if typed_node_data.prompt_config:
             enable_jinja = False
 
             if isinstance(prompt_template, list):
@@ -803,7 +866,7 @@ class LLMNode(BaseNode[LLMNodeData]):
                     enable_jinja = True
 
             if enable_jinja:
-                for variable_selector in node_data.prompt_config.jinja2_variables or []:
+                for variable_selector in typed_node_data.prompt_config.jinja2_variables or []:
                     variable_mapping[variable_selector.variable] = variable_selector.value_selector
 
         variable_mapping = {node_id + "." + key: value for key, value in variable_mapping.items()}
@@ -835,8 +898,8 @@ class LLMNode(BaseNode[LLMNodeData]):
             },
         }
 
-    def _handle_list_messages(
-        self,
+    @staticmethod
+    def handle_list_messages(
         *,
         messages: Sequence[LLMNodeChatModelMessage],
         context: Optional[str],
@@ -849,7 +912,7 @@ class LLMNode(BaseNode[LLMNodeData]):
             if message.edition_type == "jinja2":
                 result_text = _render_jinja2_message(
                     template=message.jinja2_text or "",
-                    jinjia2_variables=jinja2_variables,
+                    jinja2_variables=jinja2_variables,
                     variable_pool=variable_pool,
                 )
                 prompt_message = _combine_message_content_with_role(
@@ -897,9 +960,19 @@ class LLMNode(BaseNode[LLMNodeData]):
 
         return prompt_messages
 
-    def _handle_blocking_result(self, *, invoke_result: LLMResult) -> ModelInvokeCompletedEvent:
+    @staticmethod
+    def handle_blocking_result(
+        *,
+        invoke_result: LLMResult,
+        saver: LLMFileSaver,
+        file_outputs: list["File"],
+    ) -> ModelInvokeCompletedEvent:
         buffer = io.StringIO()
-        for text_part in self._save_multimodal_output_and_convert_result_to_markdown(invoke_result.message.content):
+        for text_part in LLMNode._save_multimodal_output_and_convert_result_to_markdown(
+            contents=invoke_result.message.content,
+            file_saver=saver,
+            file_outputs=file_outputs,
+        ):
             buffer.write(text_part)
 
         return ModelInvokeCompletedEvent(
@@ -908,7 +981,12 @@ class LLMNode(BaseNode[LLMNodeData]):
             finish_reason=None,
         )
 
-    def _save_multimodal_image_output(self, content: ImagePromptMessageContent) -> "File":
+    @staticmethod
+    def save_multimodal_image_output(
+        *,
+        content: ImagePromptMessageContent,
+        file_saver: LLMFileSaver,
+    ) -> "File":
         """_save_multimodal_output saves multi-modal contents generated by LLM plugins.
 
         There are two kinds of multimodal outputs:
@@ -918,26 +996,21 @@ class LLMNode(BaseNode[LLMNodeData]):
 
         Currently, only image files are supported.
         """
-        # Inject the saver somehow...
-        _saver = self._llm_file_saver
-
-        # If this
         if content.url != "":
-            saved_file = _saver.save_remote_url(content.url, FileType.IMAGE)
+            saved_file = file_saver.save_remote_url(content.url, FileType.IMAGE)
         else:
-            saved_file = _saver.save_binary_string(
+            saved_file = file_saver.save_binary_string(
                 data=base64.b64decode(content.base64_data),
                 mime_type=content.mime_type,
                 file_type=FileType.IMAGE,
             )
-        self._file_outputs.append(saved_file)
         return saved_file
 
     def _fetch_model_schema(self, provider: str) -> AIModelEntity | None:
         """
         Fetch model schema
         """
-        model_name = self.node_data.model.name
+        model_name = self._node_data.model.name
         model_manager = ModelManager()
         model_instance = model_manager.get_model_instance(
             tenant_id=self.tenant_id, model_type=ModelType.LLM, provider=provider, model=model_name
@@ -948,16 +1021,20 @@ class LLMNode(BaseNode[LLMNodeData]):
         model_schema = model_type_instance.get_model_schema(model_name, model_credentials)
         return model_schema
 
-    def _fetch_structured_output_schema(self) -> dict[str, Any]:
+    @staticmethod
+    def fetch_structured_output_schema(
+        *,
+        structured_output: Mapping[str, Any],
+    ) -> dict[str, Any]:
         """
         Fetch the structured output schema from the node data.
 
         Returns:
             dict[str, Any]: The structured output schema
         """
-        if not self.node_data.structured_output:
+        if not structured_output:
             raise LLMNodeError("Please provide a valid structured output schema")
-        structured_output_schema = json.dumps(self.node_data.structured_output.get("schema", {}), ensure_ascii=False)
+        structured_output_schema = json.dumps(structured_output.get("schema", {}), ensure_ascii=False)
         if not structured_output_schema:
             raise LLMNodeError("Please provide a valid structured output schema")
 
@@ -969,9 +1046,12 @@ class LLMNode(BaseNode[LLMNodeData]):
         except json.JSONDecodeError:
             raise LLMNodeError("structured_output_schema is not valid JSON format")
 
+    @staticmethod
     def _save_multimodal_output_and_convert_result_to_markdown(
-        self,
+        *,
         contents: str | list[PromptMessageContentUnionTypes] | None,
+        file_saver: LLMFileSaver,
+        file_outputs: list["File"],
     ) -> Generator[str, None, None]:
         """Convert intermediate prompt messages into strings and yield them to the caller.
 
@@ -994,9 +1074,12 @@ class LLMNode(BaseNode[LLMNodeData]):
                 if isinstance(item, TextPromptMessageContent):
                     yield item.data
                 elif isinstance(item, ImagePromptMessageContent):
-                    file = self._save_multimodal_image_output(item)
-                    self._file_outputs.append(file)
-                    yield self._image_file_to_markdown(file)
+                    file = LLMNode.save_multimodal_image_output(
+                        content=item,
+                        file_saver=file_saver,
+                    )
+                    file_outputs.append(file)
+                    yield LLMNode._image_file_to_markdown(file)
                 else:
                     logger.warning("unknown item type encountered, type=%s", type(item))
                     yield str(item)
@@ -1004,6 +1087,14 @@ class LLMNode(BaseNode[LLMNodeData]):
             logger.warning("unknown contents type encountered, type=%s", type(contents))
             yield str(contents)
 
+    @property
+    def continue_on_error(self) -> bool:
+        return self._node_data.error_strategy is not None
+
+    @property
+    def retry(self) -> bool:
+        return self._node_data.retry_config.retry_enabled
+
 
 def _combine_message_content_with_role(
     *, contents: Optional[str | list[PromptMessageContentUnionTypes]] = None, role: PromptMessageRole
@@ -1021,20 +1112,20 @@ def _combine_message_content_with_role(
 def _render_jinja2_message(
     *,
     template: str,
-    jinjia2_variables: Sequence[VariableSelector],
+    jinja2_variables: Sequence[VariableSelector],
     variable_pool: VariablePool,
 ):
     if not template:
         return ""
 
-    jinjia2_inputs = {}
-    for jinja2_variable in jinjia2_variables:
+    jinja2_inputs = {}
+    for jinja2_variable in jinja2_variables:
         variable = variable_pool.get(jinja2_variable.value_selector)
-        jinjia2_inputs[jinja2_variable.variable] = variable.to_object() if variable else ""
+        jinja2_inputs[jinja2_variable.variable] = variable.to_object() if variable else ""
     code_execute_resp = CodeExecutor.execute_workflow_code_template(
         language=CodeLanguage.JINJA2,
         code=template,
-        inputs=jinjia2_inputs,
+        inputs=jinja2_inputs,
     )
     result_text = code_execute_resp["result"]
     return result_text
@@ -1130,7 +1221,7 @@ def _handle_completion_template(
     if template.edition_type == "jinja2":
         result_text = _render_jinja2_message(
             template=template.jinja2_text or "",
-            jinjia2_variables=jinja2_variables,
+            jinja2_variables=jinja2_variables,
             variable_pool=variable_pool,
         )
     else:

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

@@ -1,18 +1,44 @@
+from collections.abc import Mapping
+from typing import Any, Optional
+
 from core.workflow.entities.node_entities import NodeRunResult
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
 from core.workflow.nodes.base import BaseNode
-from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
 from core.workflow.nodes.loop.entities import LoopEndNodeData
 
 
-class LoopEndNode(BaseNode[LoopEndNodeData]):
+class LoopEndNode(BaseNode):
     """
     Loop End Node.
     """
 
-    _node_data_cls = LoopEndNodeData
     _node_type = NodeType.LOOP_END
 
+    _node_data: LoopEndNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = LoopEndNodeData(**data)
+
+    def _get_error_strategy(self) -> Optional[ErrorStrategy]:
+        return self._node_data.error_strategy
+
+    def _get_retry_config(self) -> RetryConfig:
+        return self._node_data.retry_config
+
+    def _get_title(self) -> str:
+        return self._node_data.title
+
+    def _get_description(self) -> Optional[str]:
+        return self._node_data.desc
+
+    def _get_default_value_dict(self) -> dict[str, Any]:
+        return self._node_data.default_value_dict
+
+    def get_base_node_data(self) -> BaseNodeData:
+        return self._node_data
+
     @classmethod
     def version(cls) -> str:
         return "1"

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

@@ -3,7 +3,7 @@ import logging
 import time
 from collections.abc import Generator, Mapping, Sequence
 from datetime import UTC, datetime
-from typing import TYPE_CHECKING, Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal, Optional, cast
 
 from configs import dify_config
 from core.variables import (
@@ -30,7 +30,8 @@ from core.workflow.graph_engine.entities.event import (
 )
 from core.workflow.graph_engine.entities.graph import Graph
 from core.workflow.nodes.base import BaseNode
-from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
 from core.workflow.nodes.event import NodeEvent, RunCompletedEvent
 from core.workflow.nodes.loop.entities import LoopNodeData
 from core.workflow.utils.condition.processor import ConditionProcessor
@@ -43,14 +44,36 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-class LoopNode(BaseNode[LoopNodeData]):
+class LoopNode(BaseNode):
     """
     Loop Node.
     """
 
-    _node_data_cls = LoopNodeData
     _node_type = NodeType.LOOP
 
+    _node_data: LoopNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = LoopNodeData.model_validate(data)
+
+    def _get_error_strategy(self) -> Optional[ErrorStrategy]:
+        return self._node_data.error_strategy
+
+    def _get_retry_config(self) -> RetryConfig:
+        return self._node_data.retry_config
+
+    def _get_title(self) -> str:
+        return self._node_data.title
+
+    def _get_description(self) -> Optional[str]:
+        return self._node_data.desc
+
+    def _get_default_value_dict(self) -> dict[str, Any]:
+        return self._node_data.default_value_dict
+
+    def get_base_node_data(self) -> BaseNodeData:
+        return self._node_data
+
     @classmethod
     def version(cls) -> str:
         return "1"
@@ -58,17 +81,17 @@ class LoopNode(BaseNode[LoopNodeData]):
     def _run(self) -> Generator[NodeEvent | InNodeEvent, None, None]:
         """Run the node."""
         # Get inputs
-        loop_count = self.node_data.loop_count
-        break_conditions = self.node_data.break_conditions
-        logical_operator = self.node_data.logical_operator
+        loop_count = self._node_data.loop_count
+        break_conditions = self._node_data.break_conditions
+        logical_operator = self._node_data.logical_operator
 
         inputs = {"loop_count": loop_count}
 
-        if not self.node_data.start_node_id:
+        if not self._node_data.start_node_id:
             raise ValueError(f"field start_node_id in loop {self.node_id} not found")
 
         # Initialize graph
-        loop_graph = Graph.init(graph_config=self.graph_config, root_node_id=self.node_data.start_node_id)
+        loop_graph = Graph.init(graph_config=self.graph_config, root_node_id=self._node_data.start_node_id)
         if not loop_graph:
             raise ValueError("loop graph not found")
 
@@ -78,8 +101,8 @@ class LoopNode(BaseNode[LoopNodeData]):
 
         # Initialize loop variables
         loop_variable_selectors = {}
-        if self.node_data.loop_variables:
-            for loop_variable in self.node_data.loop_variables:
+        if self._node_data.loop_variables:
+            for loop_variable in self._node_data.loop_variables:
                 value_processor = {
                     "constant": lambda var=loop_variable: self._get_segment_for_constant(var.var_type, var.value),
                     "variable": lambda var=loop_variable: variable_pool.get(var.value),
@@ -127,8 +150,8 @@ class LoopNode(BaseNode[LoopNodeData]):
         yield LoopRunStartedEvent(
             loop_id=self.id,
             loop_node_id=self.node_id,
-            loop_node_type=self.node_type,
-            loop_node_data=self.node_data,
+            loop_node_type=self.type_,
+            loop_node_data=self._node_data,
             start_at=start_at,
             inputs=inputs,
             metadata={"loop_length": loop_count},
@@ -184,11 +207,11 @@ class LoopNode(BaseNode[LoopNodeData]):
             yield LoopRunSucceededEvent(
                 loop_id=self.id,
                 loop_node_id=self.node_id,
-                loop_node_type=self.node_type,
-                loop_node_data=self.node_data,
+                loop_node_type=self.type_,
+                loop_node_data=self._node_data,
                 start_at=start_at,
                 inputs=inputs,
-                outputs=self.node_data.outputs,
+                outputs=self._node_data.outputs,
                 steps=loop_count,
                 metadata={
                     WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: graph_engine.graph_runtime_state.total_tokens,
@@ -206,7 +229,7 @@ class LoopNode(BaseNode[LoopNodeData]):
                         WorkflowNodeExecutionMetadataKey.LOOP_DURATION_MAP: loop_duration_map,
                         WorkflowNodeExecutionMetadataKey.LOOP_VARIABLE_MAP: single_loop_variable_map,
                     },
-                    outputs=self.node_data.outputs,
+                    outputs=self._node_data.outputs,
                     inputs=inputs,
                 )
             )
@@ -217,8 +240,8 @@ class LoopNode(BaseNode[LoopNodeData]):
             yield LoopRunFailedEvent(
                 loop_id=self.id,
                 loop_node_id=self.node_id,
-                loop_node_type=self.node_type,
-                loop_node_data=self.node_data,
+                loop_node_type=self.type_,
+                loop_node_data=self._node_data,
                 start_at=start_at,
                 inputs=inputs,
                 steps=loop_count,
@@ -320,8 +343,8 @@ class LoopNode(BaseNode[LoopNodeData]):
                     yield LoopRunFailedEvent(
                         loop_id=self.id,
                         loop_node_id=self.node_id,
-                        loop_node_type=self.node_type,
-                        loop_node_data=self.node_data,
+                        loop_node_type=self.type_,
+                        loop_node_data=self._node_data,
                         start_at=start_at,
                         inputs=inputs,
                         steps=current_index,
@@ -351,8 +374,8 @@ class LoopNode(BaseNode[LoopNodeData]):
                 yield LoopRunFailedEvent(
                     loop_id=self.id,
                     loop_node_id=self.node_id,
-                    loop_node_type=self.node_type,
-                    loop_node_data=self.node_data,
+                    loop_node_type=self.type_,
+                    loop_node_data=self._node_data,
                     start_at=start_at,
                     inputs=inputs,
                     steps=current_index,
@@ -388,7 +411,7 @@ class LoopNode(BaseNode[LoopNodeData]):
                 _outputs[loop_variable_key] = None
 
         _outputs["loop_round"] = current_index + 1
-        self.node_data.outputs = _outputs
+        self._node_data.outputs = _outputs
 
         if check_break_result:
             return {"check_break_result": True}
@@ -400,10 +423,10 @@ class LoopNode(BaseNode[LoopNodeData]):
         yield LoopRunNextEvent(
             loop_id=self.id,
             loop_node_id=self.node_id,
-            loop_node_type=self.node_type,
-            loop_node_data=self.node_data,
+            loop_node_type=self.type_,
+            loop_node_data=self._node_data,
             index=next_index,
-            pre_loop_output=self.node_data.outputs,
+            pre_loop_output=self._node_data.outputs,
         )
 
         return {"check_break_result": False}
@@ -438,19 +461,15 @@ class LoopNode(BaseNode[LoopNodeData]):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: LoopNodeData,
+        node_data: Mapping[str, Any],
     ) -> Mapping[str, Sequence[str]]:
-        """
-        Extract variable selector to variable mapping
-        :param graph_config: graph config
-        :param node_id: node id
-        :param node_data: node data
-        :return:
-        """
+        # Create typed NodeData from dict
+        typed_node_data = LoopNodeData.model_validate(node_data)
+
         variable_mapping = {}
 
         # init graph
-        loop_graph = Graph.init(graph_config=graph_config, root_node_id=node_data.start_node_id)
+        loop_graph = Graph.init(graph_config=graph_config, root_node_id=typed_node_data.start_node_id)
 
         if not loop_graph:
             raise ValueError("loop graph not found")
@@ -486,7 +505,7 @@ class LoopNode(BaseNode[LoopNodeData]):
 
             variable_mapping.update(sub_node_variable_mapping)
 
-        for loop_variable in node_data.loop_variables or []:
+        for loop_variable in typed_node_data.loop_variables or []:
             if loop_variable.value_type == "variable":
                 assert loop_variable.value is not None, "Loop variable value must be provided for variable type"
                 # add loop variable to variable mapping