fix: change default value of expires_at field in tool_builtin_providers to -1

LGTM

.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

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/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/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/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/advanced_chat/app_generator.py

@@ -17,7 +17,8 @@ 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

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

@@ -13,7 +13,8 @@ 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

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/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,492 @@ 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,
+            )
+
+            # 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,
+            )
+
+            # 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,
+            )
+
+            # 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/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/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/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,7 @@ class ArizePhoenixDataTrace(BaseTraceInstance):
         }
         workflow_metadata.update(trace_info.metadata)
 
-        trace_id = uuid_to_trace_id(trace_info.message_id)
+        trace_id = uuid_to_trace_id(trace_info.workflow_run_id)
         span_id = RandomIdGenerator().generate_span_id()
         context = SpanContext(
             trace_id=trace_id,
@@ -213,7 +210,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 +233,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 +352,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 +706,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/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/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(

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/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,
@@ -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
 
@@ -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

api/core/workflow/nodes/loop/loop_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.loop.entities import LoopStartNodeData
 
 
-class LoopStartNode(BaseNode[LoopStartNodeData]):
+class LoopStartNode(BaseNode):
     """
     Loop Start Node.
     """
 
-    _node_data_cls = LoopStartNodeData
     _node_type = NodeType.LOOP_START
 
+    _node_data: LoopStartNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = LoopStartNodeData(**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/node_mapping.py

@@ -73,6 +73,9 @@ NODE_TYPE_CLASSES_MAPPING: Mapping[NodeType, Mapping[str, type[BaseNode]]] = {
     },
     NodeType.TOOL: {
         LATEST_VERSION: ToolNode,
+        # This is an issue that caused problems before.
+        # Logically, we shouldn't use two different versions to point to the same class here,
+        # but in order to maintain compatibility with historical data, this approach has been retained.
         "2": ToolNode,
         "1": ToolNode,
     },
@@ -123,6 +126,9 @@ NODE_TYPE_CLASSES_MAPPING: Mapping[NodeType, Mapping[str, type[BaseNode]]] = {
     },
     NodeType.AGENT: {
         LATEST_VERSION: AgentNode,
+        # This is an issue that caused problems before.
+        # Logically, we shouldn't use two different versions to point to the same class here,
+        # but in order to maintain compatibility with historical data, this approach has been retained.
         "2": AgentNode,
         "1": AgentNode,
     },

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

@@ -29,8 +29,9 @@ from core.variables.types import SegmentType
 from core.workflow.entities.node_entities import NodeRunResult
 from core.workflow.entities.variable_pool import VariablePool
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
 from core.workflow.nodes.base.node import BaseNode
-from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
 from core.workflow.nodes.llm import ModelConfig, llm_utils
 from core.workflow.utils import variable_template_parser
 from factories.variable_factory import build_segment_with_type
@@ -91,10 +92,31 @@ class ParameterExtractorNode(BaseNode):
     Parameter Extractor Node.
     """
 
-    # FIXME: figure out why here is different from super class
-    _node_data_cls = ParameterExtractorNodeData  # type: ignore
     _node_type = NodeType.PARAMETER_EXTRACTOR
 
+    _node_data: ParameterExtractorNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = ParameterExtractorNodeData.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
+
     _model_instance: Optional[ModelInstance] = None
     _model_config: Optional[ModelConfigWithCredentialsEntity] = None
 
@@ -119,7 +141,7 @@ class ParameterExtractorNode(BaseNode):
         """
         Run the node.
         """
-        node_data = cast(ParameterExtractorNodeData, self.node_data)
+        node_data = cast(ParameterExtractorNodeData, self._node_data)
         variable = self.graph_runtime_state.variable_pool.get(node_data.query)
         query = variable.text if variable else ""
 
@@ -398,7 +420,7 @@ class ParameterExtractorNode(BaseNode):
         """
         Generate prompt engineering prompt.
         """
-        model_mode = ModelMode.value_of(data.model.mode)
+        model_mode = ModelMode(data.model.mode)
 
         if model_mode == ModelMode.COMPLETION:
             return self._generate_prompt_engineering_completion_prompt(
@@ -694,7 +716,7 @@ class ParameterExtractorNode(BaseNode):
         memory: Optional[TokenBufferMemory],
         max_token_limit: int = 2000,
     ) -> list[ChatModelMessage]:
-        model_mode = ModelMode.value_of(node_data.model.mode)
+        model_mode = ModelMode(node_data.model.mode)
         input_text = query
         memory_str = ""
         instruction = variable_pool.convert_template(node_data.instruction or "").text
@@ -721,7 +743,7 @@ class ParameterExtractorNode(BaseNode):
         memory: Optional[TokenBufferMemory],
         max_token_limit: int = 2000,
     ):
-        model_mode = ModelMode.value_of(node_data.model.mode)
+        model_mode = ModelMode(node_data.model.mode)
         input_text = query
         memory_str = ""
         instruction = variable_pool.convert_template(node_data.instruction or "").text
@@ -827,19 +849,15 @@ class ParameterExtractorNode(BaseNode):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: ParameterExtractorNodeData,  # 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:
-        """
-        variable_mapping: dict[str, Sequence[str]] = {"query": node_data.query}
+        # Create typed NodeData from dict
+        typed_node_data = ParameterExtractorNodeData.model_validate(node_data)
 
-        if node_data.instruction:
-            selectors = variable_template_parser.extract_selectors_from_template(node_data.instruction)
+        variable_mapping: dict[str, Sequence[str]] = {"query": typed_node_data.query}
+
+        if typed_node_data.instruction:
+            selectors = variable_template_parser.extract_selectors_from_template(typed_node_data.instruction)
             for selector in selectors:
                 variable_mapping[selector.variable] = selector.value_selector
 

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

@@ -1,6 +1,6 @@
 import json
 from collections.abc import Mapping, Sequence
-from typing import Any, Optional, cast
+from typing import TYPE_CHECKING, Any, Optional, cast
 
 from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
 from core.memory.token_buffer_memory import TokenBufferMemory
@@ -11,8 +11,11 @@ from core.prompt.advanced_prompt_transform import AdvancedPromptTransform
 from core.prompt.simple_prompt_transform import ModelMode
 from core.prompt.utils.prompt_message_util import PromptMessageUtil
 from core.workflow.entities.node_entities import NodeRunResult
+from core.workflow.entities.variable_entities import VariableSelector
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus
-from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
+from core.workflow.nodes.base.node import BaseNode
+from core.workflow.nodes.enums import ErrorStrategy, NodeType
 from core.workflow.nodes.event import ModelInvokeCompletedEvent
 from core.workflow.nodes.llm import (
     LLMNode,
@@ -20,6 +23,7 @@ from core.workflow.nodes.llm import (
     LLMNodeCompletionModelPromptTemplate,
     llm_utils,
 )
+from core.workflow.nodes.llm.file_saver import FileSaverImpl, LLMFileSaver
 from core.workflow.utils.variable_template_parser import VariableTemplateParser
 from libs.json_in_md_parser import parse_and_check_json_markdown
 
@@ -35,17 +39,77 @@ from .template_prompts import (
     QUESTION_CLASSIFIER_USER_PROMPT_3,
 )
 
+if TYPE_CHECKING:
+    from core.file.models import File
+    from core.workflow.graph_engine import Graph, GraphInitParams, GraphRuntimeState
 
-class QuestionClassifierNode(LLMNode):
-    _node_data_cls = QuestionClassifierNodeData  # type: ignore
+
+class QuestionClassifierNode(BaseNode):
     _node_type = NodeType.QUESTION_CLASSIFIER
 
+    _node_data: QuestionClassifierNodeData
+
+    _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 = QuestionClassifierNodeData.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):
-        node_data = cast(QuestionClassifierNodeData, self.node_data)
+        node_data = cast(QuestionClassifierNodeData, self._node_data)
         variable_pool = self.graph_runtime_state.variable_pool
 
         # extract variables
@@ -53,7 +117,10 @@ class QuestionClassifierNode(LLMNode):
         query = variable.value if variable else None
         variables = {"query": query}
         # fetch model config
-        model_instance, model_config = self._fetch_model_config(node_data.model)
+        model_instance, model_config = LLMNode._fetch_model_config(
+            node_data_model=node_data.model,
+            tenant_id=self.tenant_id,
+        )
         # fetch memory
         memory = llm_utils.fetch_memory(
             variable_pool=variable_pool,
@@ -91,7 +158,7 @@ class QuestionClassifierNode(LLMNode):
         # If both self._get_prompt_template and self._fetch_prompt_messages append a user prompt,
         # two consecutive user prompts will be generated, causing model's error.
         # To avoid this, set sys_query to an empty string so that only one user prompt is appended at the end.
-        prompt_messages, stop = self._fetch_prompt_messages(
+        prompt_messages, stop = LLMNode.fetch_prompt_messages(
             prompt_template=prompt_template,
             sys_query="",
             memory=memory,
@@ -101,6 +168,7 @@ class QuestionClassifierNode(LLMNode):
             vision_detail=node_data.vision.configs.detail,
             variable_pool=variable_pool,
             jinja2_variables=[],
+            tenant_id=self.tenant_id,
         )
 
         result_text = ""
@@ -109,11 +177,17 @@ class QuestionClassifierNode(LLMNode):
 
         try:
             # handle invoke result
-            generator = self._invoke_llm(
+            generator = LLMNode.invoke_llm(
                 node_data_model=node_data.model,
                 model_instance=model_instance,
                 prompt_messages=prompt_messages,
                 stop=stop,
+                user_id=self.user_id,
+                structured_output_enabled=False,
+                structured_output=None,
+                file_saver=self._llm_file_saver,
+                file_outputs=self._file_outputs,
+                node_id=self.node_id,
             )
 
             for event in generator:
@@ -183,23 +257,18 @@ class QuestionClassifierNode(LLMNode):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: Any,
+        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(QuestionClassifierNodeData, node_data)
-        variable_mapping = {"query": node_data.query_variable_selector}
-        variable_selectors = []
-        if node_data.instruction:
-            variable_template_parser = VariableTemplateParser(template=node_data.instruction)
+        # Create typed NodeData from dict
+        typed_node_data = QuestionClassifierNodeData.model_validate(node_data)
+
+        variable_mapping = {"query": typed_node_data.query_variable_selector}
+        variable_selectors: list[VariableSelector] = []
+        if typed_node_data.instruction:
+            variable_template_parser = VariableTemplateParser(template=typed_node_data.instruction)
             variable_selectors.extend(variable_template_parser.extract_variable_selectors())
         for variable_selector in variable_selectors:
-            variable_mapping[variable_selector.variable] = variable_selector.value_selector
+            variable_mapping[variable_selector.variable] = list(variable_selector.value_selector)
 
         variable_mapping = {node_id + "." + key: value for key, value in variable_mapping.items()}
 
@@ -265,7 +334,7 @@ class QuestionClassifierNode(LLMNode):
         memory: Optional[TokenBufferMemory],
         max_token_limit: int = 2000,
     ):
-        model_mode = ModelMode.value_of(node_data.model.mode)
+        model_mode = ModelMode(node_data.model.mode)
         classes = node_data.classes
         categories = []
         for class_ in classes:

api/core/workflow/nodes/start/start_node.py

@@ -1,15 +1,41 @@
+from collections.abc import Mapping
+from typing import Any, Optional
+
 from core.workflow.constants import SYSTEM_VARIABLE_NODE_ID
 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.start.entities import StartNodeData
 
 
-class StartNode(BaseNode[StartNodeData]):
-    _node_data_cls = StartNodeData
+class StartNode(BaseNode):
     _node_type = NodeType.START
 
+    _node_data: StartNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = StartNodeData(**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/template_transform/template_transform_node.py

@@ -6,16 +6,39 @@ from core.helper.code_executor.code_executor import CodeExecutionError, CodeExec
 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.template_transform.entities import TemplateTransformNodeData
 
 MAX_TEMPLATE_TRANSFORM_OUTPUT_LENGTH = int(os.environ.get("TEMPLATE_TRANSFORM_MAX_LENGTH", "80000"))
 
 
-class TemplateTransformNode(BaseNode[TemplateTransformNodeData]):
-    _node_data_cls = TemplateTransformNodeData
+class TemplateTransformNode(BaseNode):
     _node_type = NodeType.TEMPLATE_TRANSFORM
 
+    _node_data: TemplateTransformNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = TemplateTransformNodeData.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:
         """
@@ -35,14 +58,14 @@ class TemplateTransformNode(BaseNode[TemplateTransformNodeData]):
     def _run(self) -> NodeRunResult:
         # Get variables
         variables = {}
-        for variable_selector in self.node_data.variables:
+        for variable_selector in self._node_data.variables:
             variable_name = variable_selector.variable
             value = self.graph_runtime_state.variable_pool.get(variable_selector.value_selector)
             variables[variable_name] = value.to_object() if value else None
         # Run code
         try:
             result = CodeExecutor.execute_workflow_code_template(
-                language=CodeLanguage.JINJA2, code=self.node_data.template, inputs=variables
+                language=CodeLanguage.JINJA2, code=self._node_data.template, inputs=variables
             )
         except CodeExecutionError as e:
             return NodeRunResult(inputs=variables, status=WorkflowNodeExecutionStatus.FAILED, error=str(e))
@@ -60,16 +83,12 @@ class TemplateTransformNode(BaseNode[TemplateTransformNodeData]):
 
     @classmethod
     def _extract_variable_selector_to_variable_mapping(
-        cls, *, graph_config: Mapping[str, Any], node_id: str, node_data: TemplateTransformNodeData
+        cls, *, graph_config: Mapping[str, Any], node_id: str, 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 = TemplateTransformNodeData.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
         }

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

@@ -59,6 +59,10 @@ class ToolNodeData(BaseNodeData, ToolEntity):
             return typ
 
     tool_parameters: dict[str, ToolInput]
+    # 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
 
     @field_validator("tool_parameters", mode="before")
     @classmethod

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

@@ -6,7 +6,6 @@ from sqlalchemy.orm import Session
 
 from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler
 from core.file import File, FileTransferMethod
-from core.model_runtime.entities.llm_entities import LLMUsage
 from core.plugin.impl.exc import PluginDaemonClientSideError
 from core.plugin.impl.plugin import PluginInstaller
 from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter
@@ -19,10 +18,10 @@ from core.workflow.entities.node_entities import NodeRunResult
 from core.workflow.entities.variable_pool import VariablePool
 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.base import BaseNode
-from core.workflow.nodes.enums import NodeType
-from core.workflow.nodes.event import RunCompletedEvent, RunRetrieverResourceEvent, RunStreamChunkEvent
+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
@@ -37,14 +36,18 @@ from .exc import (
 )
 
 
-class ToolNode(BaseNode[ToolNodeData]):
+class ToolNode(BaseNode):
     """
     Tool Node
     """
 
-    _node_data_cls = ToolNodeData
     _node_type = NodeType.TOOL
 
+    _node_data: ToolNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = ToolNodeData.model_validate(data)
+
     @classmethod
     def version(cls) -> str:
         return "1"
@@ -54,7 +57,7 @@ class ToolNode(BaseNode[ToolNodeData]):
         Run the tool node
         """
 
-        node_data = cast(ToolNodeData, self.node_data)
+        node_data = cast(ToolNodeData, self._node_data)
 
         # fetch tool icon
         tool_info = {
@@ -67,9 +70,15 @@ class ToolNode(BaseNode[ToolNodeData]):
         try:
             from core.tools.tool_manager import ToolManager
 
-            variable_pool = self.graph_runtime_state.variable_pool if self.node_data.version != "1" else None
+            # 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.
+            variable_pool: VariablePool | None = None
+            if node_data.version != "1" or node_data.tool_node_version != "1":
+                variable_pool = self.graph_runtime_state.variable_pool
             tool_runtime = ToolManager.get_workflow_tool_runtime(
-                self.tenant_id, self.app_id, self.node_id, self.node_data, self.invoke_from, variable_pool
+                self.tenant_id, self.app_id, self.node_id, self._node_data, self.invoke_from, variable_pool
             )
         except ToolNodeError as e:
             yield RunCompletedEvent(
@@ -88,12 +97,12 @@ class ToolNode(BaseNode[ToolNodeData]):
         parameters = self._generate_parameters(
             tool_parameters=tool_parameters,
             variable_pool=self.graph_runtime_state.variable_pool,
-            node_data=self.node_data,
+            node_data=self._node_data,
         )
         parameters_for_log = self._generate_parameters(
             tool_parameters=tool_parameters,
             variable_pool=self.graph_runtime_state.variable_pool,
-            node_data=self.node_data,
+            node_data=self._node_data,
             for_log=True,
         )
         # get conversation id
@@ -124,7 +133,14 @@ class ToolNode(BaseNode[ToolNodeData]):
 
         try:
             # convert tool messages
-            yield from self._transform_message(message_stream, tool_info, parameters_for_log)
+            yield from self._transform_message(
+                messages=message_stream,
+                tool_info=tool_info,
+                parameters_for_log=parameters_for_log,
+                user_id=self.user_id,
+                tenant_id=self.tenant_id,
+                node_id=self.node_id,
+            )
         except (PluginDaemonClientSideError, ToolInvokeError) as e:
             yield RunCompletedEvent(
                 run_result=NodeRunResult(
@@ -191,7 +207,9 @@ class ToolNode(BaseNode[ToolNodeData]):
         messages: Generator[ToolInvokeMessage, None, None],
         tool_info: Mapping[str, Any],
         parameters_for_log: dict[str, Any],
-        agent_thoughts: Optional[list] = None,
+        user_id: str,
+        tenant_id: str,
+        node_id: str,
     ) -> Generator:
         """
         Convert ToolInvokeMessages into tuple[plain_text, files]
@@ -199,8 +217,8 @@ class ToolNode(BaseNode[ToolNodeData]):
         # transform message and handle file storage
         message_stream = ToolFileMessageTransformer.transform_tool_invoke_messages(
             messages=messages,
-            user_id=self.user_id,
-            tenant_id=self.tenant_id,
+            user_id=user_id,
+            tenant_id=tenant_id,
             conversation_id=None,
         )
 
@@ -208,9 +226,6 @@ class ToolNode(BaseNode[ToolNodeData]):
         files: list[File] = []
         json: 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:
@@ -243,7 +258,7 @@ class ToolNode(BaseNode[ToolNodeData]):
                 }
                 file = file_factory.build_from_mapping(
                     mapping=mapping,
-                    tenant_id=self.tenant_id,
+                    tenant_id=tenant_id,
                 )
                 files.append(file)
             elif message.type == ToolInvokeMessage.MessageType.BLOB:
@@ -266,45 +281,36 @@ class ToolNode(BaseNode[ToolNodeData]):
                 files.append(
                     file_factory.build_from_mapping(
                         mapping=mapping,
-                        tenant_id=self.tenant_id,
+                        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=[self.node_id, "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 self.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()
-                    }
+                # JSON message handling for tool node
                 if message.message.json_object is not None:
                     json.append(message.message.json_object)
             elif message.type == ToolInvokeMessage.MessageType.LINK:
                 assert isinstance(message.message, ToolInvokeMessage.TextMessage)
                 stream_text = f"Link: {message.message.text}\n"
                 text += stream_text
-                yield RunStreamChunkEvent(chunk_content=stream_text, from_variable_selector=[self.node_id, "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 ValueError("When 'stream' is True, 'variable_value' must be a string.")
+                        raise ToolNodeError("When 'stream' is True, 'variable_value' must be a string.")
                     if variable_name not in variables:
                         variables[variable_name] = ""
                     variables[variable_name] += variable_value
 
                     yield RunStreamChunkEvent(
-                        chunk_content=variable_value, from_variable_selector=[self.node_id, variable_name]
+                        chunk_content=variable_value, from_variable_selector=[node_id, variable_name]
                     )
                 else:
                     variables[variable_name] = variable_value
@@ -319,7 +325,7 @@ class ToolNode(BaseNode[ToolNodeData]):
                     dict_metadata = dict(message.message.metadata)
                     if dict_metadata.get("provider"):
                         manager = PluginInstaller()
-                        plugins = manager.list_plugins(self.tenant_id)
+                        plugins = manager.list_plugins(tenant_id)
                         try:
                             current_plugin = next(
                                 plugin
@@ -334,8 +340,8 @@ class ToolNode(BaseNode[ToolNodeData]):
                             builtin_tool = next(
                                 provider
                                 for provider in BuiltinToolManageService.list_builtin_tools(
-                                    self.user_id,
-                                    self.tenant_id,
+                                    user_id,
+                                    tenant_id,
                                 )
                                 if provider.name == dict_metadata["provider"]
                             )
@@ -347,57 +353,10 @@ class ToolNode(BaseNode[ToolNodeData]):
                         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=self.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=self.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
-            elif message.type == ToolInvokeMessage.MessageType.RETRIEVER_RESOURCES:
-                assert isinstance(message.message, ToolInvokeMessage.RetrieverResourceMessage)
-                yield RunRetrieverResourceEvent(
-                    retriever_resources=message.message.retriever_resources,
-                    context=message.message.context,
-                )
 
         # Add agent_logs to outputs['json'] to ensure frontend can access thinking process
         json_output: list[dict[str, Any]] = []
 
-        # Step 1: append each agent log as its own dict.
-        if agent_logs:
-            for log in agent_logs:
-                json_output.append(
-                    {
-                        "id": log.id,
-                        "parent_id": log.parent_id,
-                        "error": log.error,
-                        "status": log.status,
-                        "data": log.data,
-                        "label": log.label,
-                        "metadata": log.metadata,
-                        "node_id": log.node_id,
-                    }
-                )
         # Step 2: normalize JSON into {"data": [...]}.change json to list[dict]
         if json:
             json_output.extend(json)
@@ -409,12 +368,9 @@ class ToolNode(BaseNode[ToolNodeData]):
                 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,
             )
         )
 
@@ -424,7 +380,7 @@ class ToolNode(BaseNode[ToolNodeData]):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: ToolNodeData,
+        node_data: Mapping[str, Any],
     ) -> Mapping[str, Sequence[str]]:
         """
         Extract variable selector to variable mapping
@@ -433,9 +389,12 @@ class ToolNode(BaseNode[ToolNodeData]):
         :param node_data: node data
         :return:
         """
+        # Create typed NodeData from dict
+        typed_node_data = ToolNodeData.model_validate(node_data)
+
         result = {}
-        for parameter_name in node_data.tool_parameters:
-            input = node_data.tool_parameters[parameter_name]
+        for parameter_name in typed_node_data.tool_parameters:
+            input = typed_node_data.tool_parameters[parameter_name]
             if input.type == "mixed":
                 assert isinstance(input.value, str)
                 selectors = VariableTemplateParser(input.value).extract_variable_selectors()
@@ -449,3 +408,29 @@ class ToolNode(BaseNode[ToolNodeData]):
         result = {node_id + "." + key: value for key, value in result.items()}
 
         return result
+
+    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
+
+    @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/variable_aggregator/variable_aggregator_node.py

@@ -1,17 +1,41 @@
 from collections.abc import Mapping
+from typing import Any, Optional
 
 from core.variables.segments import Segment
 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.variable_aggregator.entities import VariableAssignerNodeData
 
 
-class VariableAggregatorNode(BaseNode[VariableAssignerNodeData]):
-    _node_data_cls = VariableAssignerNodeData
+class VariableAggregatorNode(BaseNode):
     _node_type = NodeType.VARIABLE_AGGREGATOR
 
+    _node_data: VariableAssignerNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = VariableAssignerNodeData(**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"
@@ -21,8 +45,8 @@ class VariableAggregatorNode(BaseNode[VariableAssignerNodeData]):
         outputs: dict[str, Segment | Mapping[str, Segment]] = {}
         inputs = {}
 
-        if not self.node_data.advanced_settings or not self.node_data.advanced_settings.group_enabled:
-            for selector in self.node_data.variables:
+        if not self._node_data.advanced_settings or not self._node_data.advanced_settings.group_enabled:
+            for selector in self._node_data.variables:
                 variable = self.graph_runtime_state.variable_pool.get(selector)
                 if variable is not None:
                     outputs = {"output": variable}
@@ -30,7 +54,7 @@ class VariableAggregatorNode(BaseNode[VariableAssignerNodeData]):
                     inputs = {".".join(selector[1:]): variable.to_object()}
                     break
         else:
-            for group in self.node_data.advanced_settings.groups:
+            for group in self._node_data.advanced_settings.groups:
                 for selector in group.variables:
                     variable = self.graph_runtime_state.variable_pool.get(selector)
 

api/core/workflow/nodes/variable_assigner/v1/node.py

@@ -7,7 +7,8 @@ from core.workflow.conversation_variable_updater import ConversationVariableUpda
 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.variable_assigner.common import helpers as common_helpers
 from core.workflow.nodes.variable_assigner.common.exc import VariableOperatorNodeError
 from factories import variable_factory
@@ -22,11 +23,33 @@ if TYPE_CHECKING:
 _CONV_VAR_UPDATER_FACTORY: TypeAlias = Callable[[], ConversationVariableUpdater]
 
 
-class VariableAssignerNode(BaseNode[VariableAssignerData]):
-    _node_data_cls = VariableAssignerData
+class VariableAssignerNode(BaseNode):
     _node_type = NodeType.VARIABLE_ASSIGNER
     _conv_var_updater_factory: _CONV_VAR_UPDATER_FACTORY
 
+    _node_data: VariableAssignerData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = VariableAssignerData.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
+
     def __init__(
         self,
         id: str,
@@ -59,36 +82,39 @@ class VariableAssignerNode(BaseNode[VariableAssignerData]):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: VariableAssignerData,
+        node_data: Mapping[str, Any],
     ) -> Mapping[str, Sequence[str]]:
-        mapping = {}
-        assigned_variable_node_id = node_data.assigned_variable_selector[0]
-        if assigned_variable_node_id == CONVERSATION_VARIABLE_NODE_ID:
-            selector_key = ".".join(node_data.assigned_variable_selector)
-            key = f"{node_id}.#{selector_key}#"
-            mapping[key] = node_data.assigned_variable_selector
+        # Create typed NodeData from dict
+        typed_node_data = VariableAssignerData.model_validate(node_data)
 
-        selector_key = ".".join(node_data.input_variable_selector)
+        mapping = {}
+        assigned_variable_node_id = typed_node_data.assigned_variable_selector[0]
+        if assigned_variable_node_id == CONVERSATION_VARIABLE_NODE_ID:
+            selector_key = ".".join(typed_node_data.assigned_variable_selector)
+            key = f"{node_id}.#{selector_key}#"
+            mapping[key] = typed_node_data.assigned_variable_selector
+
+        selector_key = ".".join(typed_node_data.input_variable_selector)
         key = f"{node_id}.#{selector_key}#"
-        mapping[key] = node_data.input_variable_selector
+        mapping[key] = typed_node_data.input_variable_selector
         return mapping
 
     def _run(self) -> NodeRunResult:
-        assigned_variable_selector = self.node_data.assigned_variable_selector
+        assigned_variable_selector = self._node_data.assigned_variable_selector
         # Should be String, Number, Object, ArrayString, ArrayNumber, ArrayObject
         original_variable = self.graph_runtime_state.variable_pool.get(assigned_variable_selector)
         if not isinstance(original_variable, Variable):
             raise VariableOperatorNodeError("assigned variable not found")
 
-        match self.node_data.write_mode:
+        match self._node_data.write_mode:
             case WriteMode.OVER_WRITE:
-                income_value = self.graph_runtime_state.variable_pool.get(self.node_data.input_variable_selector)
+                income_value = self.graph_runtime_state.variable_pool.get(self._node_data.input_variable_selector)
                 if not income_value:
                     raise VariableOperatorNodeError("input value not found")
                 updated_variable = original_variable.model_copy(update={"value": income_value.value})
 
             case WriteMode.APPEND:
-                income_value = self.graph_runtime_state.variable_pool.get(self.node_data.input_variable_selector)
+                income_value = self.graph_runtime_state.variable_pool.get(self._node_data.input_variable_selector)
                 if not income_value:
                     raise VariableOperatorNodeError("input value not found")
                 updated_value = original_variable.value + [income_value.value]
@@ -101,7 +127,7 @@ class VariableAssignerNode(BaseNode[VariableAssignerData]):
                 updated_variable = original_variable.model_copy(update={"value": income_value.to_object()})
 
             case _:
-                raise VariableOperatorNodeError(f"unsupported write mode: {self.node_data.write_mode}")
+                raise VariableOperatorNodeError(f"unsupported write mode: {self._node_data.write_mode}")
 
         # Over write the variable.
         self.graph_runtime_state.variable_pool.add(assigned_variable_selector, updated_variable)

api/core/workflow/nodes/variable_assigner/v2/node.py

@@ -1,6 +1,6 @@
 import json
-from collections.abc import Callable, Mapping, MutableMapping, Sequence
-from typing import Any, TypeAlias, cast
+from collections.abc import Mapping, MutableMapping, Sequence
+from typing import Any, Optional, cast
 
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.variables import SegmentType, Variable
@@ -10,7 +10,8 @@ from core.workflow.conversation_variable_updater import ConversationVariableUpda
 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.variable_assigner.common import helpers as common_helpers
 from core.workflow.nodes.variable_assigner.common.exc import VariableOperatorNodeError
 from core.workflow.nodes.variable_assigner.common.impl import conversation_variable_updater_factory
@@ -28,8 +29,6 @@ from .exc import (
     VariableNotFoundError,
 )
 
-_CONV_VAR_UPDATER_FACTORY: TypeAlias = Callable[[], ConversationVariableUpdater]
-
 
 def _target_mapping_from_item(mapping: MutableMapping[str, Sequence[str]], node_id: str, item: VariableOperationItem):
     selector_node_id = item.variable_selector[0]
@@ -54,10 +53,32 @@ def _source_mapping_from_item(mapping: MutableMapping[str, Sequence[str]], node_
     mapping[key] = selector
 
 
-class VariableAssignerNode(BaseNode[VariableAssignerNodeData]):
-    _node_data_cls = VariableAssignerNodeData
+class VariableAssignerNode(BaseNode):
     _node_type = NodeType.VARIABLE_ASSIGNER
 
+    _node_data: VariableAssignerNodeData
+
+    def init_node_data(self, data: Mapping[str, Any]) -> None:
+        self._node_data = VariableAssignerNodeData.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
+
     def _conv_var_updater_factory(self) -> ConversationVariableUpdater:
         return conversation_variable_updater_factory()
 
@@ -71,22 +92,25 @@ class VariableAssignerNode(BaseNode[VariableAssignerNodeData]):
         *,
         graph_config: Mapping[str, Any],
         node_id: str,
-        node_data: VariableAssignerNodeData,
+        node_data: Mapping[str, Any],
     ) -> Mapping[str, Sequence[str]]:
+        # Create typed NodeData from dict
+        typed_node_data = VariableAssignerNodeData.model_validate(node_data)
+
         var_mapping: dict[str, Sequence[str]] = {}
-        for item in node_data.items:
+        for item in typed_node_data.items:
             _target_mapping_from_item(var_mapping, node_id, item)
             _source_mapping_from_item(var_mapping, node_id, item)
         return var_mapping
 
     def _run(self) -> NodeRunResult:
-        inputs = self.node_data.model_dump()
+        inputs = self._node_data.model_dump()
         process_data: dict[str, Any] = {}
         # NOTE: This node has no outputs
         updated_variable_selectors: list[Sequence[str]] = []
 
         try:
-            for item in self.node_data.items:
+            for item in self._node_data.items:
                 variable = self.graph_runtime_state.variable_pool.get(item.variable_selector)
 
                 # ==================== Validation Part

api/core/workflow/repositories/workflow_execution_repository.py

@@ -1,4 +1,4 @@
-from typing import Optional, Protocol
+from typing import Protocol
 
 from core.workflow.entities.workflow_execution import WorkflowExecution
 
@@ -28,15 +28,3 @@ class WorkflowExecutionRepository(Protocol):
             execution: The WorkflowExecution instance to save or update
         """
         ...
-
-    def get(self, execution_id: str) -> Optional[WorkflowExecution]:
-        """
-        Retrieve a WorkflowExecution by its ID.
-
-        Args:
-            execution_id: The workflow execution ID
-
-        Returns:
-            The WorkflowExecution instance if found, None otherwise
-        """
-        ...

api/core/workflow/repositories/workflow_node_execution_repository.py

@@ -39,18 +39,6 @@ class WorkflowNodeExecutionRepository(Protocol):
         """
         ...
 
-    def get_by_node_execution_id(self, node_execution_id: str) -> Optional[WorkflowNodeExecution]:
-        """
-        Retrieve a NodeExecution by its node_execution_id.
-
-        Args:
-            node_execution_id: The node execution ID
-
-        Returns:
-            The NodeExecution instance if found, None otherwise
-        """
-        ...
-
     def get_by_workflow_run(
         self,
         workflow_run_id: str,
@@ -69,24 +57,3 @@ class WorkflowNodeExecutionRepository(Protocol):
             A list of NodeExecution instances
         """
         ...
-
-    def get_running_executions(self, workflow_run_id: str) -> Sequence[WorkflowNodeExecution]:
-        """
-        Retrieve all running NodeExecution instances for a specific workflow run.
-
-        Args:
-            workflow_run_id: The workflow run ID
-
-        Returns:
-            A list of running NodeExecution instances
-        """
-        ...
-
-    def clear(self) -> None:
-        """
-        Clear all NodeExecution records based on implementation-specific criteria.
-
-        This method is intended to be used for bulk deletion operations, such as removing
-        all records associated with a specific app_id and tenant_id in multi-tenant implementations.
-        """
-        ...

api/core/workflow/workflow_cycle_manager.py

@@ -1,6 +1,6 @@
 from collections.abc import Mapping
 from dataclasses import dataclass
-from datetime import UTC, datetime
+from datetime import datetime
 from typing import Any, Optional, Union
 from uuid import uuid4
 
@@ -55,24 +55,15 @@ class WorkflowCycleManager:
         self._workflow_execution_repository = workflow_execution_repository
         self._workflow_node_execution_repository = workflow_node_execution_repository
 
+        # Initialize caches for workflow execution cycle
+        # These caches avoid redundant repository calls during a single workflow execution
+        self._workflow_execution_cache: dict[str, WorkflowExecution] = {}
+        self._node_execution_cache: dict[str, WorkflowNodeExecution] = {}
+
     def handle_workflow_run_start(self) -> WorkflowExecution:
-        inputs = {**self._application_generate_entity.inputs}
+        inputs = self._prepare_workflow_inputs()
+        execution_id = self._get_or_generate_execution_id()
 
-        # Iterate over SystemVariable fields using Pydantic's model_fields
-        if self._workflow_system_variables:
-            for field_name, value in self._workflow_system_variables.to_dict().items():
-                if field_name == SystemVariableKey.CONVERSATION_ID:
-                    continue
-                inputs[f"sys.{field_name}"] = value
-
-        # handle special values
-        inputs = dict(WorkflowEntry.handle_special_values(inputs) or {})
-
-        # init workflow run
-        # TODO: This workflow_run_id should always not be None, maybe we can use a more elegant way to handle this
-        execution_id = str(
-            self._workflow_system_variables.workflow_execution_id if self._workflow_system_variables else None
-        ) or str(uuid4())
         execution = WorkflowExecution.new(
             id_=execution_id,
             workflow_id=self._workflow_info.workflow_id,
@@ -80,12 +71,10 @@ class WorkflowCycleManager:
             workflow_version=self._workflow_info.version,
             graph=self._workflow_info.graph_data,
             inputs=inputs,
-            started_at=datetime.now(UTC).replace(tzinfo=None),
+            started_at=naive_utc_now(),
         )
 
-        self._workflow_execution_repository.save(execution)
-
-        return execution
+        return self._save_and_cache_workflow_execution(execution)
 
     def handle_workflow_run_success(
         self,
@@ -99,23 +88,15 @@ class WorkflowCycleManager:
     ) -> WorkflowExecution:
         workflow_execution = self._get_workflow_execution_or_raise_error(workflow_run_id)
 
-        # outputs = WorkflowEntry.handle_special_values(outputs)
+        self._update_workflow_execution_completion(
+            workflow_execution,
+            status=WorkflowExecutionStatus.SUCCEEDED,
+            outputs=outputs,
+            total_tokens=total_tokens,
+            total_steps=total_steps,
+        )
 
-        workflow_execution.status = WorkflowExecutionStatus.SUCCEEDED
-        workflow_execution.outputs = outputs or {}
-        workflow_execution.total_tokens = total_tokens
-        workflow_execution.total_steps = total_steps
-        workflow_execution.finished_at = datetime.now(UTC).replace(tzinfo=None)
-
-        if trace_manager:
-            trace_manager.add_trace_task(
-                TraceTask(
-                    TraceTaskName.WORKFLOW_TRACE,
-                    workflow_execution=workflow_execution,
-                    conversation_id=conversation_id,
-                    user_id=trace_manager.user_id,
-                )
-            )
+        self._add_trace_task_if_needed(trace_manager, workflow_execution, conversation_id)
 
         self._workflow_execution_repository.save(workflow_execution)
         return workflow_execution
@@ -132,24 +113,17 @@ class WorkflowCycleManager:
         trace_manager: Optional[TraceQueueManager] = None,
     ) -> WorkflowExecution:
         execution = self._get_workflow_execution_or_raise_error(workflow_run_id)
-        # outputs = WorkflowEntry.handle_special_values(dict(outputs) if outputs else None)
 
-        execution.status = WorkflowExecutionStatus.PARTIAL_SUCCEEDED
-        execution.outputs = outputs or {}
-        execution.total_tokens = total_tokens
-        execution.total_steps = total_steps
-        execution.finished_at = datetime.now(UTC).replace(tzinfo=None)
-        execution.exceptions_count = exceptions_count
+        self._update_workflow_execution_completion(
+            execution,
+            status=WorkflowExecutionStatus.PARTIAL_SUCCEEDED,
+            outputs=outputs,
+            total_tokens=total_tokens,
+            total_steps=total_steps,
+            exceptions_count=exceptions_count,
+        )
 
-        if trace_manager:
-            trace_manager.add_trace_task(
-                TraceTask(
-                    TraceTaskName.WORKFLOW_TRACE,
-                    workflow_execution=execution,
-                    conversation_id=conversation_id,
-                    user_id=trace_manager.user_id,
-                )
-            )
+        self._add_trace_task_if_needed(trace_manager, execution, conversation_id)
 
         self._workflow_execution_repository.save(execution)
         return execution
@@ -169,39 +143,18 @@ class WorkflowCycleManager:
         workflow_execution = self._get_workflow_execution_or_raise_error(workflow_run_id)
         now = naive_utc_now()
 
-        workflow_execution.status = WorkflowExecutionStatus(status.value)
-        workflow_execution.error_message = error_message
-        workflow_execution.total_tokens = total_tokens
-        workflow_execution.total_steps = total_steps
-        workflow_execution.finished_at = now
-        workflow_execution.exceptions_count = exceptions_count
-
-        # Use the instance repository to find running executions for a workflow run
-        running_node_executions = self._workflow_node_execution_repository.get_running_executions(
-            workflow_run_id=workflow_execution.id_
+        self._update_workflow_execution_completion(
+            workflow_execution,
+            status=status,
+            total_tokens=total_tokens,
+            total_steps=total_steps,
+            error_message=error_message,
+            exceptions_count=exceptions_count,
+            finished_at=now,
         )
 
-        # Update the domain models
-        for node_execution in running_node_executions:
-            if node_execution.node_execution_id:
-                # Update the domain model
-                node_execution.status = WorkflowNodeExecutionStatus.FAILED
-                node_execution.error = error_message
-                node_execution.finished_at = now
-                node_execution.elapsed_time = (now - node_execution.created_at).total_seconds()
-
-                # Update the repository with the domain model
-                self._workflow_node_execution_repository.save(node_execution)
-
-        if trace_manager:
-            trace_manager.add_trace_task(
-                TraceTask(
-                    TraceTaskName.WORKFLOW_TRACE,
-                    workflow_execution=workflow_execution,
-                    conversation_id=conversation_id,
-                    user_id=trace_manager.user_id,
-                )
-            )
+        self._fail_running_node_executions(workflow_execution.id_, error_message, now)
+        self._add_trace_task_if_needed(trace_manager, workflow_execution, conversation_id)
 
         self._workflow_execution_repository.save(workflow_execution)
         return workflow_execution
@@ -214,8 +167,198 @@ class WorkflowCycleManager:
     ) -> WorkflowNodeExecution:
         workflow_execution = self._get_workflow_execution_or_raise_error(workflow_execution_id)
 
-        # Create a domain model
-        created_at = datetime.now(UTC).replace(tzinfo=None)
+        domain_execution = self._create_node_execution_from_event(
+            workflow_execution=workflow_execution,
+            event=event,
+            status=WorkflowNodeExecutionStatus.RUNNING,
+        )
+
+        return self._save_and_cache_node_execution(domain_execution)
+
+    def handle_workflow_node_execution_success(self, *, event: QueueNodeSucceededEvent) -> WorkflowNodeExecution:
+        domain_execution = self._get_node_execution_from_cache(event.node_execution_id)
+
+        self._update_node_execution_completion(
+            domain_execution,
+            event=event,
+            status=WorkflowNodeExecutionStatus.SUCCEEDED,
+        )
+
+        self._workflow_node_execution_repository.save(domain_execution)
+        return domain_execution
+
+    def handle_workflow_node_execution_failed(
+        self,
+        *,
+        event: QueueNodeFailedEvent
+        | QueueNodeInIterationFailedEvent
+        | QueueNodeInLoopFailedEvent
+        | QueueNodeExceptionEvent,
+    ) -> WorkflowNodeExecution:
+        """
+        Workflow node execution failed
+        :param event: queue node failed event
+        :return:
+        """
+        domain_execution = self._get_node_execution_from_cache(event.node_execution_id)
+
+        status = (
+            WorkflowNodeExecutionStatus.EXCEPTION
+            if isinstance(event, QueueNodeExceptionEvent)
+            else WorkflowNodeExecutionStatus.FAILED
+        )
+
+        self._update_node_execution_completion(
+            domain_execution,
+            event=event,
+            status=status,
+            error=event.error,
+            handle_special_values=True,
+        )
+
+        self._workflow_node_execution_repository.save(domain_execution)
+        return domain_execution
+
+    def handle_workflow_node_execution_retried(
+        self, *, workflow_execution_id: str, event: QueueNodeRetryEvent
+    ) -> WorkflowNodeExecution:
+        workflow_execution = self._get_workflow_execution_or_raise_error(workflow_execution_id)
+
+        domain_execution = self._create_node_execution_from_event(
+            workflow_execution=workflow_execution,
+            event=event,
+            status=WorkflowNodeExecutionStatus.RETRY,
+            error=event.error,
+            created_at=event.start_at,
+        )
+
+        # Handle inputs and outputs
+        inputs = WorkflowEntry.handle_special_values(event.inputs)
+        outputs = event.outputs
+        metadata = self._merge_event_metadata(event)
+
+        domain_execution.update_from_mapping(inputs=inputs, outputs=outputs, metadata=metadata)
+
+        return self._save_and_cache_node_execution(domain_execution)
+
+    def _get_workflow_execution_or_raise_error(self, id: str, /) -> WorkflowExecution:
+        # Check cache first
+        if id in self._workflow_execution_cache:
+            return self._workflow_execution_cache[id]
+
+        raise WorkflowRunNotFoundError(id)
+
+    def _prepare_workflow_inputs(self) -> dict[str, Any]:
+        """Prepare workflow inputs by merging application inputs with system variables."""
+        inputs = {**self._application_generate_entity.inputs}
+
+        if self._workflow_system_variables:
+            for field_name, value in self._workflow_system_variables.to_dict().items():
+                if field_name != SystemVariableKey.CONVERSATION_ID:
+                    inputs[f"sys.{field_name}"] = value
+
+        return dict(WorkflowEntry.handle_special_values(inputs) or {})
+
+    def _get_or_generate_execution_id(self) -> str:
+        """Get execution ID from system variables or generate a new one."""
+        if self._workflow_system_variables and self._workflow_system_variables.workflow_execution_id:
+            return str(self._workflow_system_variables.workflow_execution_id)
+        return str(uuid4())
+
+    def _save_and_cache_workflow_execution(self, execution: WorkflowExecution) -> WorkflowExecution:
+        """Save workflow execution to repository and cache it."""
+        self._workflow_execution_repository.save(execution)
+        self._workflow_execution_cache[execution.id_] = execution
+        return execution
+
+    def _save_and_cache_node_execution(self, execution: WorkflowNodeExecution) -> WorkflowNodeExecution:
+        """Save node execution to repository and cache it if it has an ID."""
+        self._workflow_node_execution_repository.save(execution)
+        if execution.node_execution_id:
+            self._node_execution_cache[execution.node_execution_id] = execution
+        return execution
+
+    def _get_node_execution_from_cache(self, node_execution_id: str) -> WorkflowNodeExecution:
+        """Get node execution from cache or raise error if not found."""
+        domain_execution = self._node_execution_cache.get(node_execution_id)
+        if not domain_execution:
+            raise ValueError(f"Domain node execution not found: {node_execution_id}")
+        return domain_execution
+
+    def _update_workflow_execution_completion(
+        self,
+        execution: WorkflowExecution,
+        *,
+        status: WorkflowExecutionStatus,
+        total_tokens: int,
+        total_steps: int,
+        outputs: Mapping[str, Any] | None = None,
+        error_message: Optional[str] = None,
+        exceptions_count: int = 0,
+        finished_at: Optional[datetime] = None,
+    ) -> None:
+        """Update workflow execution with completion data."""
+        execution.status = status
+        execution.outputs = outputs or {}
+        execution.total_tokens = total_tokens
+        execution.total_steps = total_steps
+        execution.finished_at = finished_at or naive_utc_now()
+        execution.exceptions_count = exceptions_count
+        if error_message:
+            execution.error_message = error_message
+
+    def _add_trace_task_if_needed(
+        self,
+        trace_manager: Optional[TraceQueueManager],
+        workflow_execution: WorkflowExecution,
+        conversation_id: Optional[str],
+    ) -> None:
+        """Add trace task if trace manager is provided."""
+        if trace_manager:
+            trace_manager.add_trace_task(
+                TraceTask(
+                    TraceTaskName.WORKFLOW_TRACE,
+                    workflow_execution=workflow_execution,
+                    conversation_id=conversation_id,
+                    user_id=trace_manager.user_id,
+                )
+            )
+
+    def _fail_running_node_executions(
+        self,
+        workflow_execution_id: str,
+        error_message: str,
+        now: datetime,
+    ) -> None:
+        """Fail all running node executions for a workflow."""
+        running_node_executions = [
+            node_exec
+            for node_exec in self._node_execution_cache.values()
+            if node_exec.workflow_execution_id == workflow_execution_id
+            and node_exec.status == WorkflowNodeExecutionStatus.RUNNING
+        ]
+
+        for node_execution in running_node_executions:
+            if node_execution.node_execution_id:
+                node_execution.status = WorkflowNodeExecutionStatus.FAILED
+                node_execution.error = error_message
+                node_execution.finished_at = now
+                node_execution.elapsed_time = (now - node_execution.created_at).total_seconds()
+                self._workflow_node_execution_repository.save(node_execution)
+
+    def _create_node_execution_from_event(
+        self,
+        *,
+        workflow_execution: WorkflowExecution,
+        event: Union[QueueNodeStartedEvent, QueueNodeRetryEvent],
+        status: WorkflowNodeExecutionStatus,
+        error: Optional[str] = None,
+        created_at: Optional[datetime] = None,
+    ) -> WorkflowNodeExecution:
+        """Create a node execution from an event."""
+        now = naive_utc_now()
+        created_at = created_at or now
+
         metadata = {
             WorkflowNodeExecutionMetadataKey.PARALLEL_MODE_RUN_ID: event.parallel_mode_run_id,
             WorkflowNodeExecutionMetadataKey.ITERATION_ID: event.in_iteration_id,
@@ -232,152 +375,76 @@ class WorkflowCycleManager:
             node_id=event.node_id,
             node_type=event.node_type,
             title=event.node_data.title,
-            status=WorkflowNodeExecutionStatus.RUNNING,
+            status=status,
             metadata=metadata,
             created_at=created_at,
+            error=error,
         )
 
-        # Use the instance repository to save the domain model
-        self._workflow_node_execution_repository.save(domain_execution)
+        if status == WorkflowNodeExecutionStatus.RETRY:
+            domain_execution.finished_at = now
+            domain_execution.elapsed_time = (now - created_at).total_seconds()
 
         return domain_execution
 
-    def handle_workflow_node_execution_success(self, *, event: QueueNodeSucceededEvent) -> WorkflowNodeExecution:
-        # Get the domain model from repository
-        domain_execution = self._workflow_node_execution_repository.get_by_node_execution_id(event.node_execution_id)
-        if not domain_execution:
-            raise ValueError(f"Domain node execution not found: {event.node_execution_id}")
-
-        # Process data
-        inputs = event.inputs
-        process_data = event.process_data
-        outputs = event.outputs
-
-        # Convert metadata keys to strings
-        execution_metadata_dict = {}
-        if event.execution_metadata:
-            for key, value in event.execution_metadata.items():
-                execution_metadata_dict[key] = value
-
-        finished_at = datetime.now(UTC).replace(tzinfo=None)
-        elapsed_time = (finished_at - event.start_at).total_seconds()
-
-        # Update domain model
-        domain_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED
-        domain_execution.update_from_mapping(
-            inputs=inputs, process_data=process_data, outputs=outputs, metadata=execution_metadata_dict
-        )
-        domain_execution.finished_at = finished_at
-        domain_execution.elapsed_time = elapsed_time
-
-        # Update the repository with the domain model
-        self._workflow_node_execution_repository.save(domain_execution)
-
-        return domain_execution
-
-    def handle_workflow_node_execution_failed(
+    def _update_node_execution_completion(
         self,
+        domain_execution: WorkflowNodeExecution,
         *,
-        event: QueueNodeFailedEvent
-        | QueueNodeInIterationFailedEvent
-        | QueueNodeInLoopFailedEvent
-        | QueueNodeExceptionEvent,
-    ) -> WorkflowNodeExecution:
-        """
-        Workflow node execution failed
-        :param event: queue node failed event
-        :return:
-        """
-        # Get the domain model from repository
-        domain_execution = self._workflow_node_execution_repository.get_by_node_execution_id(event.node_execution_id)
-        if not domain_execution:
-            raise ValueError(f"Domain node execution not found: {event.node_execution_id}")
-
-        # Process data
-        inputs = WorkflowEntry.handle_special_values(event.inputs)
-        process_data = WorkflowEntry.handle_special_values(event.process_data)
-        outputs = event.outputs
-
-        # Convert metadata keys to strings
-        execution_metadata_dict = {}
-        if event.execution_metadata:
-            for key, value in event.execution_metadata.items():
-                execution_metadata_dict[key] = value
-
-        finished_at = datetime.now(UTC).replace(tzinfo=None)
+        event: Union[
+            QueueNodeSucceededEvent,
+            QueueNodeFailedEvent,
+            QueueNodeInIterationFailedEvent,
+            QueueNodeInLoopFailedEvent,
+            QueueNodeExceptionEvent,
+        ],
+        status: WorkflowNodeExecutionStatus,
+        error: Optional[str] = None,
+        handle_special_values: bool = False,
+    ) -> None:
+        """Update node execution with completion data."""
+        finished_at = naive_utc_now()
         elapsed_time = (finished_at - event.start_at).total_seconds()
 
+        # Process data
+        if handle_special_values:
+            inputs = WorkflowEntry.handle_special_values(event.inputs)
+            process_data = WorkflowEntry.handle_special_values(event.process_data)
+        else:
+            inputs = event.inputs
+            process_data = event.process_data
+
+        outputs = event.outputs
+
+        # Convert metadata
+        execution_metadata_dict: dict[WorkflowNodeExecutionMetadataKey, Any] = {}
+        if event.execution_metadata:
+            execution_metadata_dict.update(event.execution_metadata)
+
         # Update domain model
-        domain_execution.status = (
-            WorkflowNodeExecutionStatus.FAILED
-            if not isinstance(event, QueueNodeExceptionEvent)
-            else WorkflowNodeExecutionStatus.EXCEPTION
-        )
-        domain_execution.error = event.error
+        domain_execution.status = status
         domain_execution.update_from_mapping(
-            inputs=inputs, process_data=process_data, outputs=outputs, metadata=execution_metadata_dict
+            inputs=inputs,
+            process_data=process_data,
+            outputs=outputs,
+            metadata=execution_metadata_dict,
         )
         domain_execution.finished_at = finished_at
         domain_execution.elapsed_time = elapsed_time
 
-        # Update the repository with the domain model
-        self._workflow_node_execution_repository.save(domain_execution)
+        if error:
+            domain_execution.error = error
 
-        return domain_execution
-
-    def handle_workflow_node_execution_retried(
-        self, *, workflow_execution_id: str, event: QueueNodeRetryEvent
-    ) -> WorkflowNodeExecution:
-        workflow_execution = self._get_workflow_execution_or_raise_error(workflow_execution_id)
-        created_at = event.start_at
-        finished_at = datetime.now(UTC).replace(tzinfo=None)
-        elapsed_time = (finished_at - created_at).total_seconds()
-        inputs = WorkflowEntry.handle_special_values(event.inputs)
-        outputs = event.outputs
-
-        # Convert metadata keys to strings
+    def _merge_event_metadata(self, event: QueueNodeRetryEvent) -> dict[WorkflowNodeExecutionMetadataKey, str | None]:
+        """Merge event metadata with origin metadata."""
         origin_metadata = {
             WorkflowNodeExecutionMetadataKey.ITERATION_ID: event.in_iteration_id,
             WorkflowNodeExecutionMetadataKey.PARALLEL_MODE_RUN_ID: event.parallel_mode_run_id,
             WorkflowNodeExecutionMetadataKey.LOOP_ID: event.in_loop_id,
         }
 
-        # Convert execution metadata keys to strings
         execution_metadata_dict: dict[WorkflowNodeExecutionMetadataKey, str | None] = {}
         if event.execution_metadata:
-            for key, value in event.execution_metadata.items():
-                execution_metadata_dict[key] = value
+            execution_metadata_dict.update(event.execution_metadata)
 
-        merged_metadata = {**execution_metadata_dict, **origin_metadata} if execution_metadata_dict else origin_metadata
-
-        # Create a domain model
-        domain_execution = WorkflowNodeExecution(
-            id=str(uuid4()),
-            workflow_id=workflow_execution.workflow_id,
-            workflow_execution_id=workflow_execution.id_,
-            predecessor_node_id=event.predecessor_node_id,
-            node_execution_id=event.node_execution_id,
-            node_id=event.node_id,
-            node_type=event.node_type,
-            title=event.node_data.title,
-            status=WorkflowNodeExecutionStatus.RETRY,
-            created_at=created_at,
-            finished_at=finished_at,
-            elapsed_time=elapsed_time,
-            error=event.error,
-            index=event.node_run_index,
-        )
-
-        # Update with mappings
-        domain_execution.update_from_mapping(inputs=inputs, outputs=outputs, metadata=merged_metadata)
-
-        # Use the instance repository to save the domain model
-        self._workflow_node_execution_repository.save(domain_execution)
-
-        return domain_execution
-
-    def _get_workflow_execution_or_raise_error(self, id: str, /) -> WorkflowExecution:
-        execution = self._workflow_execution_repository.get(id)
-        if not execution:
-            raise WorkflowRunNotFoundError(id)
-        return execution
+        return {**execution_metadata_dict, **origin_metadata} if execution_metadata_dict else origin_metadata

api/core/workflow/workflow_entry.py

@@ -5,7 +5,7 @@ from collections.abc import Generator, Mapping, Sequence
 from typing import Any, Optional, cast
 
 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.file.models import File
 from core.workflow.callbacks import WorkflowCallback
@@ -146,7 +146,7 @@ class WorkflowEntry:
         graph = Graph.init(graph_config=workflow.graph_dict)
 
         # init workflow run state
-        node_instance = node_cls(
+        node = node_cls(
             id=str(uuid.uuid4()),
             config=node_config,
             graph_init_params=GraphInitParams(
@@ -163,6 +163,7 @@ class WorkflowEntry:
             graph=graph,
             graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()),
         )
+        node.init_node_data(node_config_data)
 
         try:
             # variable selector to variable mapping
@@ -190,17 +191,11 @@ class WorkflowEntry:
 
         try:
             # run node
-            generator = node_instance.run()
+            generator = node.run()
         except Exception as e:
-            logger.exception(
-                "error while running node_instance, workflow_id=%s, node_id=%s, type=%s, version=%s",
-                workflow.id,
-                node_instance.id,
-                node_instance.node_type,
-                node_instance.version(),
-            )
-            raise WorkflowNodeRunFailedError(node_instance=node_instance, error=str(e))
-        return node_instance, generator
+            logger.exception(f"error while running node, {workflow.id=}, {node.id=}, {node.type_=}, {node.version()=}")
+            raise WorkflowNodeRunFailedError(node=node, err_msg=str(e))
+        return node, generator
 
     @classmethod
     def run_free_node(
@@ -262,7 +257,7 @@ class WorkflowEntry:
 
         node_cls = cast(type[BaseNode], node_cls)
         # init workflow run state
-        node_instance: BaseNode = node_cls(
+        node: BaseNode = node_cls(
             id=str(uuid.uuid4()),
             config=node_config,
             graph_init_params=GraphInitParams(
@@ -279,6 +274,7 @@ class WorkflowEntry:
             graph=graph,
             graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()),
         )
+        node.init_node_data(node_data)
 
         try:
             # variable selector to variable mapping
@@ -297,17 +293,12 @@ class WorkflowEntry:
             )
 
             # run node
-            generator = node_instance.run()
+            generator = node.run()
 
-            return node_instance, generator
+            return node, generator
         except Exception as e:
-            logger.exception(
-                "error while running node_instance, node_id=%s, type=%s, version=%s",
-                node_instance.id,
-                node_instance.node_type,
-                node_instance.version(),
-            )
-            raise WorkflowNodeRunFailedError(node_instance=node_instance, error=str(e))
+            logger.exception(f"error while running node, {node.id=}, {node.type_=}, {node.version()=}")
+            raise WorkflowNodeRunFailedError(node=node, err_msg=str(e))
 
     @staticmethod
     def handle_special_values(value: Optional[Mapping[str, Any]]) -> Mapping[str, Any] | None:

api/events/event_handlers/create_document_index.py

@@ -1,4 +1,3 @@
-import datetime
 import logging
 import time
 
@@ -8,6 +7,7 @@ from werkzeug.exceptions import NotFound
 from core.indexing_runner import DocumentIsPausedError, IndexingRunner
 from events.event_handlers.document_index_event import document_index_created
 from extensions.ext_database import db
+from libs.datetime_utils import naive_utc_now
 from models.dataset import Document
 
 
@@ -33,7 +33,7 @@ def handle(sender, **kwargs):
             raise NotFound("Document not found")
 
         document.indexing_status = "parsing"
-        document.processing_started_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+        document.processing_started_at = naive_utc_now()
         documents.append(document)
         db.session.add(document)
     db.session.commit()

api/extensions/storage/azure_blob_storage.py

@@ -1,5 +1,5 @@
 from collections.abc import Generator
-from datetime import UTC, datetime, timedelta
+from datetime import timedelta
 from typing import Optional
 
 from azure.identity import ChainedTokenCredential, DefaultAzureCredential
@@ -8,6 +8,7 @@ from azure.storage.blob import AccountSasPermissions, BlobServiceClient, Resourc
 from configs import dify_config
 from extensions.ext_redis import redis_client
 from extensions.storage.base_storage import BaseStorage
+from libs.datetime_utils import naive_utc_now
 
 
 class AzureBlobStorage(BaseStorage):
@@ -78,7 +79,7 @@ class AzureBlobStorage(BaseStorage):
                 account_key=self.account_key or "",
                 resource_types=ResourceTypes(service=True, container=True, object=True),
                 permission=AccountSasPermissions(read=True, write=True, delete=True, list=True, add=True, create=True),
-                expiry=datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=1),
+                expiry=naive_utc_now() + timedelta(hours=1),
             )
             redis_client.set(cache_key, sas_token, ex=3000)
         return BlobServiceClient(account_url=self.account_url or "", credential=sas_token)

api/factories/file_factory.py

@@ -148,9 +148,7 @@ def _build_from_local_file(
     if strict_type_validation and detected_file_type.value != specified_type:
         raise ValueError("Detected file type does not match the specified type. Please verify the file.")
 
-    file_type = (
-        FileType(specified_type) if specified_type and specified_type != FileType.CUSTOM.value else detected_file_type
-    )
+    file_type = FileType(specified_type) if specified_type and specified_type != FileType.CUSTOM else detected_file_type
 
     return File(
         id=mapping.get("id"),
@@ -199,9 +197,7 @@ def _build_from_remote_url(
             raise ValueError("Detected file type does not match the specified type. Please verify the file.")
 
         file_type = (
-            FileType(specified_type)
-            if specified_type and specified_type != FileType.CUSTOM.value
-            else detected_file_type
+            FileType(specified_type) if specified_type and specified_type != FileType.CUSTOM else detected_file_type
         )
 
         return File(
@@ -286,9 +282,7 @@ def _build_from_tool_file(
     if strict_type_validation and specified_type and detected_file_type.value != specified_type:
         raise ValueError("Detected file type does not match the specified type. Please verify the file.")
 
-    file_type = (
-        FileType(specified_type) if specified_type and specified_type != FileType.CUSTOM.value else detected_file_type
-    )
+    file_type = FileType(specified_type) if specified_type and specified_type != FileType.CUSTOM else detected_file_type
 
     return File(
         id=mapping.get("id"),

api/libs/oauth_data_source.py

@@ -1,4 +1,3 @@
-import datetime
 import urllib.parse
 from typing import Any
 
@@ -6,6 +5,7 @@ import requests
 from flask_login import current_user
 
 from extensions.ext_database import db
+from libs.datetime_utils import naive_utc_now
 from models.source import DataSourceOauthBinding
 
 
@@ -75,7 +75,7 @@ class NotionOAuth(OAuthDataSource):
         if data_source_binding:
             data_source_binding.source_info = source_info
             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.commit()
         else:
             new_data_source_binding = DataSourceOauthBinding(
@@ -115,7 +115,7 @@ class NotionOAuth(OAuthDataSource):
         if data_source_binding:
             data_source_binding.source_info = source_info
             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.commit()
         else:
             new_data_source_binding = DataSourceOauthBinding(
@@ -154,7 +154,7 @@ class NotionOAuth(OAuthDataSource):
             }
             data_source_binding.source_info = new_source_info
             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.commit()
         else:
             raise ValueError("Data source binding not found")

api/libs/rsa.py

@@ -1,4 +1,5 @@
 import hashlib
+from typing import Union
 
 from Crypto.Cipher import AES
 from Crypto.PublicKey import RSA
@@ -9,7 +10,7 @@ from extensions.ext_storage import storage
 from libs import gmpy2_pkcs10aep_cipher
 
 
-def generate_key_pair(tenant_id):
+def generate_key_pair(tenant_id: str) -> str:
     private_key = RSA.generate(2048)
     public_key = private_key.publickey()
 
@@ -26,7 +27,7 @@ def generate_key_pair(tenant_id):
 prefix_hybrid = b"HYBRID:"
 
 
-def encrypt(text, public_key):
+def encrypt(text: str, public_key: Union[str, bytes]) -> bytes:
     if isinstance(public_key, str):
         public_key = public_key.encode()
 
@@ -38,14 +39,14 @@ def encrypt(text, public_key):
     rsa_key = RSA.import_key(public_key)
     cipher_rsa = gmpy2_pkcs10aep_cipher.new(rsa_key)
 
-    enc_aes_key = cipher_rsa.encrypt(aes_key)
+    enc_aes_key: bytes = cipher_rsa.encrypt(aes_key)
 
     encrypted_data = enc_aes_key + cipher_aes.nonce + tag + ciphertext
 
     return prefix_hybrid + encrypted_data
 
 
-def get_decrypt_decoding(tenant_id):
+def get_decrypt_decoding(tenant_id: str) -> tuple[RSA.RsaKey, object]:
     filepath = "privkeys/{tenant_id}".format(tenant_id=tenant_id) + "/private.pem"
 
     cache_key = "tenant_privkey:{hash}".format(hash=hashlib.sha3_256(filepath.encode()).hexdigest())
@@ -64,7 +65,7 @@ def get_decrypt_decoding(tenant_id):
     return rsa_key, cipher_rsa
 
 
-def decrypt_token_with_decoding(encrypted_text, rsa_key, cipher_rsa):
+def decrypt_token_with_decoding(encrypted_text: bytes, rsa_key: RSA.RsaKey, cipher_rsa) -> str:
     if encrypted_text.startswith(prefix_hybrid):
         encrypted_text = encrypted_text[len(prefix_hybrid) :]
 
@@ -83,10 +84,10 @@ def decrypt_token_with_decoding(encrypted_text, rsa_key, cipher_rsa):
     return decrypted_text.decode()
 
 
-def decrypt(encrypted_text, tenant_id):
+def decrypt(encrypted_text: bytes, tenant_id: str) -> str:
     rsa_key, cipher_rsa = get_decrypt_decoding(tenant_id)
 
-    return decrypt_token_with_decoding(encrypted_text, rsa_key, cipher_rsa)
+    return decrypt_token_with_decoding(encrypted_text=encrypted_text, rsa_key=rsa_key, cipher_rsa=cipher_rsa)
 
 
 class PrivkeyNotFoundError(Exception):

api/migrations/versions/2025_05_15_1635-16081485540c_.py

@@ -1,41 +0,0 @@
-"""empty message
-
-Revision ID: 16081485540c
-Revises: d28f2004b072
-Create Date: 2025-05-15 16:35:39.113777
-
-"""
-from alembic import op
-import models as models
-import sqlalchemy as sa
-
-
-# revision identifiers, used by Alembic.
-revision = '16081485540c'
-down_revision = '2adcbe1f5dfb'
-branch_labels = None
-depends_on = None
-
-
-def upgrade():
-    # ### commands auto generated by Alembic - please adjust! ###
-    op.create_table('tenant_plugin_auto_upgrade_strategies',
-    sa.Column('id', models.types.StringUUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False),
-    sa.Column('tenant_id', models.types.StringUUID(), nullable=False),
-    sa.Column('strategy_setting', sa.String(length=16), server_default='fix_only', nullable=False),
-    sa.Column('upgrade_time_of_day', sa.Integer(), nullable=False),
-    sa.Column('upgrade_mode', sa.String(length=16), server_default='exclude', nullable=False),
-    sa.Column('exclude_plugins', sa.ARRAY(sa.String(length=255)), nullable=False),
-    sa.Column('include_plugins', sa.ARRAY(sa.String(length=255)), nullable=False),
-    sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False),
-    sa.Column('updated_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False),
-    sa.PrimaryKeyConstraint('id', name='tenant_plugin_auto_upgrade_strategy_pkey'),
-    sa.UniqueConstraint('tenant_id', name='unique_tenant_plugin_auto_upgrade_strategy')
-    )
-    # ### end Alembic commands ###
-
-
-def downgrade():
-    # ### commands auto generated by Alembic - please adjust! ###
-    op.drop_table('tenant_plugin_auto_upgrade_strategies')
-    # ### end Alembic commands ###

api/migrations/versions/2025_06_06_1424-4474872b0ee6_workflow_draft_varaibles_add_node_execution_id.py

@@ -12,7 +12,7 @@ import sqlalchemy as sa
 
 # revision identifiers, used by Alembic.
 revision = '4474872b0ee6'
-down_revision = '16081485540c'
+down_revision = '2adcbe1f5dfb'
 branch_labels = None
 depends_on = None
 

api/migrations/versions/2025_07_21_0935-1a83934ad6d1_update_models.py

@@ -0,0 +1,51 @@
+"""update models
+
+Revision ID: 1a83934ad6d1
+Revises: 71f5020c6470
+Create Date: 2025-07-21 09:35:48.774794
+
+"""
+from alembic import op
+import models as models
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '1a83934ad6d1'
+down_revision = '71f5020c6470'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('tool_mcp_providers', schema=None) as batch_op:
+        batch_op.alter_column('server_identifier',
+               existing_type=sa.VARCHAR(length=24),
+               type_=sa.String(length=64),
+               existing_nullable=False)
+
+    with op.batch_alter_table('tool_model_invokes', schema=None) as batch_op:
+        batch_op.alter_column('tool_name',
+               existing_type=sa.VARCHAR(length=40),
+               type_=sa.String(length=128),
+               existing_nullable=False)
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('tool_model_invokes', schema=None) as batch_op:
+        batch_op.alter_column('tool_name',
+               existing_type=sa.String(length=128),
+               type_=sa.VARCHAR(length=40),
+               existing_nullable=False)
+
+    with op.batch_alter_table('tool_mcp_providers', schema=None) as batch_op:
+        batch_op.alter_column('server_identifier',
+               existing_type=sa.String(length=64),
+               type_=sa.VARCHAR(length=24),
+               existing_nullable=False)
+
+    # ### end Alembic commands ###

api/migrations/versions/2025_07_22_0019-375fe79ead14_oauth_refresh_token.py

@@ -0,0 +1,34 @@
+"""oauth_refresh_token
+
+Revision ID: 375fe79ead14
+Revises: 1a83934ad6d1
+Create Date: 2025-07-22 00:19:45.599636
+
+"""
+from alembic import op
+import models as models
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = '375fe79ead14'
+down_revision = '1a83934ad6d1'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('tool_builtin_providers', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('expires_at', sa.BigInteger(), server_default=sa.text('-1'), nullable=False))
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+
+    with op.batch_alter_table('tool_builtin_providers', schema=None) as batch_op:
+        batch_op.drop_column('expires_at')
+
+    # ### end Alembic commands ###

api/models/account.py

@@ -196,7 +196,7 @@ class Tenant(Base):
     __tablename__ = "tenants"
     __table_args__ = (db.PrimaryKeyConstraint("id", name="tenant_pkey"),)
 
-    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))
+    id: Mapped[str] = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))
     name = db.Column(db.String(255), nullable=False)
     encrypt_public_key = db.Column(db.Text)
     plan = db.Column(db.String(255), nullable=False, server_default=db.text("'basic'::character varying"))