rm useless router.replace (#27386)

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

api/.env.example

@@ -434,6 +434,9 @@ CODE_EXECUTION_SSL_VERIFY=True
 CODE_EXECUTION_POOL_MAX_CONNECTIONS=100
 CODE_EXECUTION_POOL_MAX_KEEPALIVE_CONNECTIONS=20
 CODE_EXECUTION_POOL_KEEPALIVE_EXPIRY=5.0
+CODE_EXECUTION_CONNECT_TIMEOUT=10
+CODE_EXECUTION_READ_TIMEOUT=60
+CODE_EXECUTION_WRITE_TIMEOUT=10
 CODE_MAX_NUMBER=9223372036854775807
 CODE_MIN_NUMBER=-9223372036854775808
 CODE_MAX_STRING_LENGTH=400000

api/configs/middleware/__init__.py

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

api/constants/__init__.py

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

api/constants/languages.py

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

api/controllers/common/helpers.py

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

api/controllers/console/auth/login.py

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

api/controllers/console/explore/workflow.py

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

api/controllers/console/setup.py

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

api/controllers/mcp/mcp.py

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

api/controllers/web/login.py

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

api/controllers/web/passport.py

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

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

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

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

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

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

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

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

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

api/core/app/task_pipeline/based_generate_task_pipeline.py

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

api/core/entities/provider_configuration.py

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

api/core/helper/module_import_helper.py

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

api/core/indexing_runner.py

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

api/core/ops/langfuse_trace/langfuse_trace.py

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

api/core/plugin/entities/parameters.py

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

api/core/plugin/impl/base.py

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

api/core/plugin/impl/exc.py

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

api/core/rag/retrieval/dataset_retrieval.py

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

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

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

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

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

api/core/repositories/celery_workflow_execution_repository.py

@@ -74,7 +74,7 @@ class CeleryWorkflowExecutionRepository(WorkflowExecutionRepository):
         tenant_id = extract_tenant_id(user)
         if not tenant_id:
             raise ValueError("User must have a tenant_id or current_tenant_id")
-        self._tenant_id = tenant_id  # type: ignore[assignment]  # We've already checked tenant_id is not None
+        self._tenant_id = tenant_id
 
         # Store app context
         self._app_id = app_id

api/core/repositories/celery_workflow_node_execution_repository.py

@@ -81,7 +81,7 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository):
         tenant_id = extract_tenant_id(user)
         if not tenant_id:
             raise ValueError("User must have a tenant_id or current_tenant_id")
-        self._tenant_id = tenant_id  # type: ignore[assignment]  # We've already checked tenant_id is not None
+        self._tenant_id = tenant_id
 
         # Store app context
         self._app_id = app_id

api/core/repositories/factory.py

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

api/core/tools/builtin_tool/provider.py

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

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

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

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

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

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

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

api/core/tools/mcp_tool/provider.py

@@ -105,7 +105,7 @@ class MCPToolProviderController(ToolProviderController):
         """
         pass
 
-    def get_tool(self, tool_name: str) -> MCPTool:  # type: ignore
+    def get_tool(self, tool_name: str) -> MCPTool:
         """
         return tool with given name
         """
@@ -128,7 +128,7 @@ class MCPToolProviderController(ToolProviderController):
             sse_read_timeout=self.sse_read_timeout,
         )
 
-    def get_tools(self) -> list[MCPTool]:  # type: ignore
+    def get_tools(self) -> list[MCPTool]:
         """
         get all tools
         """

api/core/tools/tool_label_manager.py

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

api/core/tools/tool_manager.py

@@ -326,7 +326,8 @@ class ToolManager:
             workflow_provider_stmt = select(WorkflowToolProvider).where(
                 WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id
             )
-            workflow_provider = db.session.scalar(workflow_provider_stmt)
+            with Session(db.engine, expire_on_commit=False) as session, session.begin():
+                workflow_provider = session.scalar(workflow_provider_stmt)
 
             if workflow_provider is None:
                 raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found")

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

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

api/core/tools/utils/parser.py

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

api/core/tools/utils/web_reader_tool.py

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

api/core/tools/utils/yaml_utils.py

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

api/core/tools/workflow_as_tool/provider.py

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

api/core/tools/workflow_as_tool/tool.py

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

api/core/variables/segment_group.py

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

api/core/variables/segments.py

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

api/core/workflow/entities/__init__.py

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

api/core/workflow/graph/graph.py

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

api/core/workflow/graph/validation.py

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

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

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

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

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

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

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

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

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

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

@@ -10,10 +10,10 @@ from typing import Any
 import chardet
 import docx
 import pandas as pd
-import pypandoc  # type: ignore
-import pypdfium2  # type: ignore
-import webvtt  # type: ignore
-import yaml  # type: ignore
+import pypandoc
+import pypdfium2
+import webvtt
+import yaml
 from docx.document import Document
 from docx.oxml.table import CT_Tbl
 from docx.oxml.text.paragraph import CT_P

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

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any, NewType, cast
 from flask import Flask, current_app
 from typing_extensions import TypeIs
 
+from core.model_runtime.entities.llm_entities import LLMUsage
 from core.variables import IntegerVariable, NoneSegment
 from core.variables.segments import ArrayAnySegment, ArraySegment
 from core.variables.variables import VariableUnion
@@ -34,6 +35,7 @@ from core.workflow.node_events import (
     NodeRunResult,
     StreamCompletedEvent,
 )
+from core.workflow.nodes.base import LLMUsageTrackingMixin
 from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.iteration.entities import ErrorHandleMode, IterationNodeData
@@ -58,7 +60,7 @@ logger = logging.getLogger(__name__)
 EmptyArraySegment = NewType("EmptyArraySegment", ArraySegment)
 
 
-class IterationNode(Node):
+class IterationNode(LLMUsageTrackingMixin, Node):
     """
     Iteration Node.
     """
@@ -118,6 +120,7 @@ class IterationNode(Node):
         started_at = naive_utc_now()
         iter_run_map: dict[str, float] = {}
         outputs: list[object] = []
+        usage_accumulator = [LLMUsage.empty_usage()]
 
         yield IterationStartedEvent(
             start_at=started_at,
@@ -130,22 +133,27 @@ class IterationNode(Node):
                 iterator_list_value=iterator_list_value,
                 outputs=outputs,
                 iter_run_map=iter_run_map,
+                usage_accumulator=usage_accumulator,
             )
 
+            self._accumulate_usage(usage_accumulator[0])
             yield from self._handle_iteration_success(
                 started_at=started_at,
                 inputs=inputs,
                 outputs=outputs,
                 iterator_list_value=iterator_list_value,
                 iter_run_map=iter_run_map,
+                usage=usage_accumulator[0],
             )
         except IterationNodeError as e:
+            self._accumulate_usage(usage_accumulator[0])
             yield from self._handle_iteration_failure(
                 started_at=started_at,
                 inputs=inputs,
                 outputs=outputs,
                 iterator_list_value=iterator_list_value,
                 iter_run_map=iter_run_map,
+                usage=usage_accumulator[0],
                 error=e,
             )
 
@@ -196,6 +204,7 @@ class IterationNode(Node):
         iterator_list_value: Sequence[object],
         outputs: list[object],
         iter_run_map: dict[str, float],
+        usage_accumulator: list[LLMUsage],
     ) -> Generator[GraphNodeEventBase | NodeEventBase, None, None]:
         if self._node_data.is_parallel:
             # Parallel mode execution
@@ -203,6 +212,7 @@ class IterationNode(Node):
                 iterator_list_value=iterator_list_value,
                 outputs=outputs,
                 iter_run_map=iter_run_map,
+                usage_accumulator=usage_accumulator,
             )
         else:
             # Sequential mode execution
@@ -228,6 +238,9 @@ class IterationNode(Node):
 
                 # Update the total tokens from this iteration
                 self.graph_runtime_state.total_tokens += graph_engine.graph_runtime_state.total_tokens
+                usage_accumulator[0] = self._merge_usage(
+                    usage_accumulator[0], graph_engine.graph_runtime_state.llm_usage
+                )
                 iter_run_map[str(index)] = (datetime.now(UTC).replace(tzinfo=None) - iter_start_at).total_seconds()
 
     def _execute_parallel_iterations(
@@ -235,6 +248,7 @@ class IterationNode(Node):
         iterator_list_value: Sequence[object],
         outputs: list[object],
         iter_run_map: dict[str, float],
+        usage_accumulator: list[LLMUsage],
     ) -> Generator[GraphNodeEventBase | NodeEventBase, None, None]:
         # Initialize outputs list with None values to maintain order
         outputs.extend([None] * len(iterator_list_value))
@@ -245,7 +259,16 @@ class IterationNode(Node):
         with ThreadPoolExecutor(max_workers=max_workers) as executor:
             # Submit all iteration tasks
             future_to_index: dict[
-                Future[tuple[datetime, list[GraphNodeEventBase], object | None, int, dict[str, VariableUnion]]],
+                Future[
+                    tuple[
+                        datetime,
+                        list[GraphNodeEventBase],
+                        object | None,
+                        int,
+                        dict[str, VariableUnion],
+                        LLMUsage,
+                    ]
+                ],
                 int,
             ] = {}
             for index, item in enumerate(iterator_list_value):
@@ -264,7 +287,14 @@ class IterationNode(Node):
                 index = future_to_index[future]
                 try:
                     result = future.result()
-                    iter_start_at, events, output_value, tokens_used, conversation_snapshot = result
+                    (
+                        iter_start_at,
+                        events,
+                        output_value,
+                        tokens_used,
+                        conversation_snapshot,
+                        iteration_usage,
+                    ) = result
 
                     # Update outputs at the correct index
                     outputs[index] = output_value
@@ -276,6 +306,8 @@ class IterationNode(Node):
                     self.graph_runtime_state.total_tokens += tokens_used
                     iter_run_map[str(index)] = (datetime.now(UTC).replace(tzinfo=None) - iter_start_at).total_seconds()
 
+                    usage_accumulator[0] = self._merge_usage(usage_accumulator[0], iteration_usage)
+
                     # Sync conversation variables after iteration completion
                     self._sync_conversation_variables_from_snapshot(conversation_snapshot)
 
@@ -303,7 +335,7 @@ class IterationNode(Node):
         item: object,
         flask_app: Flask,
         context_vars: contextvars.Context,
-    ) -> tuple[datetime, list[GraphNodeEventBase], object | None, int, dict[str, VariableUnion]]:
+    ) -> tuple[datetime, list[GraphNodeEventBase], object | None, int, dict[str, VariableUnion], LLMUsage]:
         """Execute a single iteration in parallel mode and return results."""
         with preserve_flask_contexts(flask_app=flask_app, context_vars=context_vars):
             iter_start_at = datetime.now(UTC).replace(tzinfo=None)
@@ -332,6 +364,7 @@ class IterationNode(Node):
                 output_value,
                 graph_engine.graph_runtime_state.total_tokens,
                 conversation_snapshot,
+                graph_engine.graph_runtime_state.llm_usage,
             )
 
     def _handle_iteration_success(
@@ -341,6 +374,8 @@ class IterationNode(Node):
         outputs: list[object],
         iterator_list_value: Sequence[object],
         iter_run_map: dict[str, float],
+        *,
+        usage: LLMUsage,
     ) -> Generator[NodeEventBase, None, None]:
         # Flatten the list of lists if all outputs are lists
         flattened_outputs = self._flatten_outputs_if_needed(outputs)
@@ -351,7 +386,9 @@ class IterationNode(Node):
             outputs={"output": flattened_outputs},
             steps=len(iterator_list_value),
             metadata={
-                WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: self.graph_runtime_state.total_tokens,
+                WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: usage.total_tokens,
+                WorkflowNodeExecutionMetadataKey.TOTAL_PRICE: usage.total_price,
+                WorkflowNodeExecutionMetadataKey.CURRENCY: usage.currency,
                 WorkflowNodeExecutionMetadataKey.ITERATION_DURATION_MAP: iter_run_map,
             },
         )
@@ -362,8 +399,11 @@ class IterationNode(Node):
                 status=WorkflowNodeExecutionStatus.SUCCEEDED,
                 outputs={"output": flattened_outputs},
                 metadata={
-                    WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: self.graph_runtime_state.total_tokens,
+                    WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: usage.total_tokens,
+                    WorkflowNodeExecutionMetadataKey.TOTAL_PRICE: usage.total_price,
+                    WorkflowNodeExecutionMetadataKey.CURRENCY: usage.currency,
                 },
+                llm_usage=usage,
             )
         )
 
@@ -400,6 +440,8 @@ class IterationNode(Node):
         outputs: list[object],
         iterator_list_value: Sequence[object],
         iter_run_map: dict[str, float],
+        *,
+        usage: LLMUsage,
         error: IterationNodeError,
     ) -> Generator[NodeEventBase, None, None]:
         # Flatten the list of lists if all outputs are lists (even in failure case)
@@ -411,7 +453,9 @@ class IterationNode(Node):
             outputs={"output": flattened_outputs},
             steps=len(iterator_list_value),
             metadata={
-                WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: self.graph_runtime_state.total_tokens,
+                WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: usage.total_tokens,
+                WorkflowNodeExecutionMetadataKey.TOTAL_PRICE: usage.total_price,
+                WorkflowNodeExecutionMetadataKey.CURRENCY: usage.currency,
                 WorkflowNodeExecutionMetadataKey.ITERATION_DURATION_MAP: iter_run_map,
             },
             error=str(error),
@@ -420,6 +464,12 @@ class IterationNode(Node):
             node_run_result=NodeRunResult(
                 status=WorkflowNodeExecutionStatus.FAILED,
                 error=str(error),
+                metadata={
+                    WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: usage.total_tokens,
+                    WorkflowNodeExecutionMetadataKey.TOTAL_PRICE: usage.total_price,
+                    WorkflowNodeExecutionMetadataKey.CURRENCY: usage.currency,
+                },
+                llm_usage=usage,
             )
         )
 

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

@@ -15,14 +15,11 @@ 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.llm_entities import LLMUsage
+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.model_runtime.utils.encoders import jsonable_encoder
 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
@@ -33,8 +30,14 @@ from core.variables import (
 )
 from core.variables.segments import ArrayObjectSegment
 from core.workflow.entities import GraphInitParams
-from core.workflow.enums import ErrorStrategy, NodeType, WorkflowNodeExecutionStatus
+from core.workflow.enums import (
+    ErrorStrategy,
+    NodeType,
+    WorkflowNodeExecutionMetadataKey,
+    WorkflowNodeExecutionStatus,
+)
 from core.workflow.node_events import ModelInvokeCompletedEvent, NodeRunResult
+from core.workflow.nodes.base import LLMUsageTrackingMixin
 from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.knowledge_retrieval.template_prompts import (
@@ -80,7 +83,7 @@ default_retrieval_model = {
 }
 
 
-class KnowledgeRetrievalNode(Node):
+class KnowledgeRetrievalNode(LLMUsageTrackingMixin, Node):
     node_type = NodeType.KNOWLEDGE_RETRIEVAL
 
     _node_data: KnowledgeRetrievalNodeData
@@ -141,7 +144,7 @@ class KnowledgeRetrievalNode(Node):
     def version(cls):
         return "1"
 
-    def _run(self) -> NodeRunResult:  # type: ignore
+    def _run(self) -> NodeRunResult:
         # extract variables
         variable = self.graph_runtime_state.variable_pool.get(self._node_data.query_variable_selector)
         if not isinstance(variable, StringSegment):
@@ -182,14 +185,21 @@ class KnowledgeRetrievalNode(Node):
                 )
 
         # retrieve knowledge
+        usage = LLMUsage.empty_usage()
         try:
-            results = self._fetch_dataset_retriever(node_data=self._node_data, query=query)
+            results, usage = self._fetch_dataset_retriever(node_data=self._node_data, query=query)
             outputs = {"result": ArrayObjectSegment(value=results)}
             return NodeRunResult(
                 status=WorkflowNodeExecutionStatus.SUCCEEDED,
                 inputs=variables,
-                process_data={},
+                process_data={"usage": jsonable_encoder(usage)},
                 outputs=outputs,  # type: ignore
+                metadata={
+                    WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: usage.total_tokens,
+                    WorkflowNodeExecutionMetadataKey.TOTAL_PRICE: usage.total_price,
+                    WorkflowNodeExecutionMetadataKey.CURRENCY: usage.currency,
+                },
+                llm_usage=usage,
             )
 
         except KnowledgeRetrievalNodeError as e:
@@ -199,6 +209,7 @@ class KnowledgeRetrievalNode(Node):
                 inputs=variables,
                 error=str(e),
                 error_type=type(e).__name__,
+                llm_usage=usage,
             )
         # Temporary handle all exceptions from DatasetRetrieval class here.
         except Exception as e:
@@ -207,11 +218,15 @@ class KnowledgeRetrievalNode(Node):
                 inputs=variables,
                 error=str(e),
                 error_type=type(e).__name__,
+                llm_usage=usage,
             )
         finally:
             db.session.close()
 
-    def _fetch_dataset_retriever(self, node_data: KnowledgeRetrievalNodeData, query: str) -> list[dict[str, Any]]:
+    def _fetch_dataset_retriever(
+        self, node_data: KnowledgeRetrievalNodeData, query: str
+    ) -> tuple[list[dict[str, Any]], LLMUsage]:
+        usage = LLMUsage.empty_usage()
         available_datasets = []
         dataset_ids = node_data.dataset_ids
 
@@ -245,9 +260,10 @@ class KnowledgeRetrievalNode(Node):
             if not dataset:
                 continue
             available_datasets.append(dataset)
-        metadata_filter_document_ids, metadata_condition = self._get_metadata_filter_condition(
+        metadata_filter_document_ids, metadata_condition, metadata_usage = self._get_metadata_filter_condition(
             [dataset.id for dataset in available_datasets], query, node_data
         )
+        usage = self._merge_usage(usage, metadata_usage)
         all_documents = []
         dataset_retrieval = DatasetRetrieval()
         if node_data.retrieval_mode == DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE:
@@ -330,6 +346,8 @@ class KnowledgeRetrievalNode(Node):
                 metadata_filter_document_ids=metadata_filter_document_ids,
                 metadata_condition=metadata_condition,
             )
+        usage = self._merge_usage(usage, dataset_retrieval.llm_usage)
+
         dify_documents = [item for item in all_documents if item.provider == "dify"]
         external_documents = [item for item in all_documents if item.provider == "external"]
         retrieval_resource_list = []
@@ -406,11 +424,12 @@ class KnowledgeRetrievalNode(Node):
             )
             for position, item in enumerate(retrieval_resource_list, start=1):
                 item["metadata"]["position"] = position
-        return retrieval_resource_list
+        return retrieval_resource_list, usage
 
     def _get_metadata_filter_condition(
         self, dataset_ids: list, query: str, node_data: KnowledgeRetrievalNodeData
-    ) -> tuple[dict[str, list[str]] | None, MetadataCondition | None]:
+    ) -> tuple[dict[str, list[str]] | None, MetadataCondition | None, LLMUsage]:
+        usage = LLMUsage.empty_usage()
         document_query = db.session.query(Document).where(
             Document.dataset_id.in_(dataset_ids),
             Document.indexing_status == "completed",
@@ -420,9 +439,12 @@ class KnowledgeRetrievalNode(Node):
         filters: list[Any] = []
         metadata_condition = None
         if node_data.metadata_filtering_mode == "disabled":
-            return None, None
+            return None, None, usage
         elif node_data.metadata_filtering_mode == "automatic":
-            automatic_metadata_filters = self._automatic_metadata_filter_func(dataset_ids, query, node_data)
+            automatic_metadata_filters, automatic_usage = self._automatic_metadata_filter_func(
+                dataset_ids, query, node_data
+            )
+            usage = self._merge_usage(usage, automatic_usage)
             if automatic_metadata_filters:
                 conditions = []
                 for sequence, filter in enumerate(automatic_metadata_filters):
@@ -443,7 +465,7 @@ class KnowledgeRetrievalNode(Node):
                 metadata_condition = MetadataCondition(
                     logical_operator=node_data.metadata_filtering_conditions.logical_operator
                     if node_data.metadata_filtering_conditions
-                    else "or",  # type: ignore
+                    else "or",
                     conditions=conditions,
                 )
         elif node_data.metadata_filtering_mode == "manual":
@@ -457,10 +479,10 @@ class KnowledgeRetrievalNode(Node):
                             expected_value = self.graph_runtime_state.variable_pool.convert_template(
                                 expected_value
                             ).value[0]
-                            if expected_value.value_type in {"number", "integer", "float"}:  # type: ignore
-                                expected_value = expected_value.value  # type: ignore
-                            elif expected_value.value_type == "string":  # type: ignore
-                                expected_value = re.sub(r"[\r\n\t]+", " ", expected_value.text).strip()  # type: ignore
+                            if expected_value.value_type in {"number", "integer", "float"}:
+                                expected_value = expected_value.value
+                            elif expected_value.value_type == "string":
+                                expected_value = re.sub(r"[\r\n\t]+", " ", expected_value.text).strip()
                             else:
                                 raise ValueError("Invalid expected metadata value type")
                     conditions.append(
@@ -487,7 +509,7 @@ class KnowledgeRetrievalNode(Node):
             if (
                 node_data.metadata_filtering_conditions
                 and node_data.metadata_filtering_conditions.logical_operator == "and"
-            ):  # type: ignore
+            ):
                 document_query = document_query.where(and_(*filters))
             else:
                 document_query = document_query.where(or_(*filters))
@@ -496,11 +518,12 @@ class KnowledgeRetrievalNode(Node):
         metadata_filter_document_ids = defaultdict(list) if documents else None  # type: ignore
         for document in documents:
             metadata_filter_document_ids[document.dataset_id].append(document.id)  # type: ignore
-        return metadata_filter_document_ids, metadata_condition
+        return metadata_filter_document_ids, metadata_condition, usage
 
     def _automatic_metadata_filter_func(
         self, dataset_ids: list, query: str, node_data: KnowledgeRetrievalNodeData
-    ) -> list[dict[str, Any]]:
+    ) -> tuple[list[dict[str, Any]], LLMUsage]:
+        usage = LLMUsage.empty_usage()
         # get all metadata field
         stmt = select(DatasetMetadata).where(DatasetMetadata.dataset_id.in_(dataset_ids))
         metadata_fields = db.session.scalars(stmt).all()
@@ -548,6 +571,7 @@ class KnowledgeRetrievalNode(Node):
             for event in generator:
                 if isinstance(event, ModelInvokeCompletedEvent):
                     result_text = event.text
+                    usage = self._merge_usage(usage, event.usage)
                     break
 
             result_text_json = parse_and_check_json_markdown(result_text, [])
@@ -564,8 +588,8 @@ class KnowledgeRetrievalNode(Node):
                             }
                         )
         except Exception:
-            return []
-        return automatic_metadata_filters
+            return [], usage
+        return automatic_metadata_filters, usage
 
     def _process_metadata_filter_func(
         self, sequence: int, condition: str, metadata_name: str, value: Any, filters: list[Any]

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

@@ -5,6 +5,7 @@ from collections.abc import Callable, Generator, Mapping, Sequence
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Literal, cast
 
+from core.model_runtime.entities.llm_entities import LLMUsage
 from core.variables import Segment, SegmentType
 from core.workflow.enums import (
     ErrorStrategy,
@@ -27,6 +28,7 @@ from core.workflow.node_events import (
     NodeRunResult,
     StreamCompletedEvent,
 )
+from core.workflow.nodes.base import LLMUsageTrackingMixin
 from core.workflow.nodes.base.entities import BaseNodeData, RetryConfig
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.loop.entities import LoopNodeData, LoopVariableData
@@ -40,7 +42,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-class LoopNode(Node):
+class LoopNode(LLMUsageTrackingMixin, Node):
     """
     Loop Node.
     """
@@ -108,7 +110,7 @@ class LoopNode(Node):
                     raise ValueError(f"Invalid value for loop variable {loop_variable.label}")
                 variable_selector = [self._node_id, loop_variable.label]
                 variable = segment_to_variable(segment=processed_segment, selector=variable_selector)
-                self.graph_runtime_state.variable_pool.add(variable_selector, variable)
+                self.graph_runtime_state.variable_pool.add(variable_selector, variable.value)
                 loop_variable_selectors[loop_variable.label] = variable_selector
                 inputs[loop_variable.label] = processed_segment.value
 
@@ -117,6 +119,7 @@ class LoopNode(Node):
 
         loop_duration_map: dict[str, float] = {}
         single_loop_variable_map: dict[str, dict[str, Any]] = {}  # single loop variable output
+        loop_usage = LLMUsage.empty_usage()
 
         # Start Loop event
         yield LoopStartedEvent(
@@ -163,6 +166,9 @@ class LoopNode(Node):
                 # Update the total tokens from this iteration
                 cost_tokens += graph_engine.graph_runtime_state.total_tokens
 
+                # Accumulate usage from the sub-graph execution
+                loop_usage = self._merge_usage(loop_usage, graph_engine.graph_runtime_state.llm_usage)
+
                 # Collect loop variable values after iteration
                 single_loop_variable = {}
                 for key, selector in loop_variable_selectors.items():
@@ -189,6 +195,7 @@ class LoopNode(Node):
                 )
 
             self.graph_runtime_state.total_tokens += cost_tokens
+            self._accumulate_usage(loop_usage)
             # Loop completed successfully
             yield LoopSucceededEvent(
                 start_at=start_at,
@@ -196,7 +203,9 @@ class LoopNode(Node):
                 outputs=self._node_data.outputs,
                 steps=loop_count,
                 metadata={
-                    WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: cost_tokens,
+                    WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: loop_usage.total_tokens,
+                    WorkflowNodeExecutionMetadataKey.TOTAL_PRICE: loop_usage.total_price,
+                    WorkflowNodeExecutionMetadataKey.CURRENCY: loop_usage.currency,
                     "completed_reason": "loop_break" if reach_break_condition else "loop_completed",
                     WorkflowNodeExecutionMetadataKey.LOOP_DURATION_MAP: loop_duration_map,
                     WorkflowNodeExecutionMetadataKey.LOOP_VARIABLE_MAP: single_loop_variable_map,
@@ -207,22 +216,28 @@ class LoopNode(Node):
                 node_run_result=NodeRunResult(
                     status=WorkflowNodeExecutionStatus.SUCCEEDED,
                     metadata={
-                        WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: self.graph_runtime_state.total_tokens,
+                        WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: loop_usage.total_tokens,
+                        WorkflowNodeExecutionMetadataKey.TOTAL_PRICE: loop_usage.total_price,
+                        WorkflowNodeExecutionMetadataKey.CURRENCY: loop_usage.currency,
                         WorkflowNodeExecutionMetadataKey.LOOP_DURATION_MAP: loop_duration_map,
                         WorkflowNodeExecutionMetadataKey.LOOP_VARIABLE_MAP: single_loop_variable_map,
                     },
                     outputs=self._node_data.outputs,
                     inputs=inputs,
+                    llm_usage=loop_usage,
                 )
             )
 
         except Exception as e:
+            self._accumulate_usage(loop_usage)
             yield LoopFailedEvent(
                 start_at=start_at,
                 inputs=inputs,
                 steps=loop_count,
                 metadata={
-                    WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: self.graph_runtime_state.total_tokens,
+                    WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: loop_usage.total_tokens,
+                    WorkflowNodeExecutionMetadataKey.TOTAL_PRICE: loop_usage.total_price,
+                    WorkflowNodeExecutionMetadataKey.CURRENCY: loop_usage.currency,
                     "completed_reason": "error",
                     WorkflowNodeExecutionMetadataKey.LOOP_DURATION_MAP: loop_duration_map,
                     WorkflowNodeExecutionMetadataKey.LOOP_VARIABLE_MAP: single_loop_variable_map,
@@ -235,10 +250,13 @@ class LoopNode(Node):
                     status=WorkflowNodeExecutionStatus.FAILED,
                     error=str(e),
                     metadata={
-                        WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: self.graph_runtime_state.total_tokens,
+                        WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: loop_usage.total_tokens,
+                        WorkflowNodeExecutionMetadataKey.TOTAL_PRICE: loop_usage.total_price,
+                        WorkflowNodeExecutionMetadataKey.CURRENCY: loop_usage.currency,
                         WorkflowNodeExecutionMetadataKey.LOOP_DURATION_MAP: loop_duration_map,
                         WorkflowNodeExecutionMetadataKey.LOOP_VARIABLE_MAP: single_loop_variable_map,
                     },
+                    llm_usage=loop_usage,
                 )
             )
 

api/core/workflow/nodes/node_factory.py

@@ -2,7 +2,7 @@ from typing import TYPE_CHECKING, final
 
 from typing_extensions import override
 
-from core.workflow.enums import ErrorStrategy, NodeExecutionType, NodeType
+from core.workflow.enums import NodeType
 from core.workflow.graph import NodeFactory
 from core.workflow.nodes.base.node import Node
 from libs.typing import is_str, is_str_dict
@@ -82,8 +82,4 @@ class DifyNodeFactory(NodeFactory):
             raise ValueError(f"Node {node_id} missing data information")
         node_instance.init_node_data(node_data)
 
-        # If node has fail branch, change execution type to branch
-        if node_instance.error_strategy == ErrorStrategy.FAIL_BRANCH:
-            node_instance.execution_type = NodeExecutionType.BRANCH
-
         return node_instance

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

@@ -193,15 +193,19 @@ class QuestionClassifierNode(Node):
                     finish_reason = event.finish_reason
                     break
 
-            category_name = node_data.classes[0].name
-            category_id = node_data.classes[0].id
+            rendered_classes = [
+                c.model_copy(update={"name": variable_pool.convert_template(c.name).text}) for c in node_data.classes
+            ]
+
+            category_name = rendered_classes[0].name
+            category_id = rendered_classes[0].id
             if "<think>" in result_text:
                 result_text = re.sub(r"<think[^>]*>[\s\S]*?</think>", "", result_text, flags=re.IGNORECASE)
             result_text_json = parse_and_check_json_markdown(result_text, [])
             # result_text_json = json.loads(result_text.strip('```JSON\n'))
             if "category_name" in result_text_json and "category_id" in result_text_json:
                 category_id_result = result_text_json["category_id"]
-                classes = node_data.classes
+                classes = rendered_classes
                 classes_map = {class_.id: class_.name for class_ in classes}
                 category_ids = [_class.id for _class in classes]
                 if category_id_result in category_ids:

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

@@ -6,10 +6,13 @@ 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.tools.__base.tool import Tool
 from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter
 from core.tools.errors import ToolInvokeError
 from core.tools.tool_engine import ToolEngine
 from core.tools.utils.message_transformer import ToolFileMessageTransformer
+from core.tools.workflow_as_tool.tool import WorkflowTool
 from core.variables.segments import ArrayAnySegment, ArrayFileSegment
 from core.variables.variables import ArrayAnyVariable
 from core.workflow.enums import (
@@ -136,13 +139,14 @@ class ToolNode(Node):
 
         try:
             # convert tool messages
-            yield from self._transform_message(
+            _ = 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,
+                tool_runtime=tool_runtime,
             )
         except ToolInvokeError as e:
             yield StreamCompletedEvent(
@@ -236,7 +240,8 @@ class ToolNode(Node):
         user_id: str,
         tenant_id: str,
         node_id: str,
-    ) -> Generator:
+        tool_runtime: Tool,
+    ) -> Generator[NodeEventBase, None, LLMUsage]:
         """
         Convert ToolInvokeMessages into tuple[plain_text, files]
         """
@@ -424,17 +429,34 @@ class ToolNode(Node):
                 is_final=True,
             )
 
+        usage = self._extract_tool_usage(tool_runtime)
+
+        metadata: dict[WorkflowNodeExecutionMetadataKey, Any] = {
+            WorkflowNodeExecutionMetadataKey.TOOL_INFO: tool_info,
+        }
+        if usage.total_tokens > 0:
+            metadata[WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS] = usage.total_tokens
+            metadata[WorkflowNodeExecutionMetadataKey.TOTAL_PRICE] = usage.total_price
+            metadata[WorkflowNodeExecutionMetadataKey.CURRENCY] = usage.currency
+
         yield StreamCompletedEvent(
             node_run_result=NodeRunResult(
                 status=WorkflowNodeExecutionStatus.SUCCEEDED,
                 outputs={"text": text, "files": ArrayFileSegment(value=files), "json": json_output, **variables},
-                metadata={
-                    WorkflowNodeExecutionMetadataKey.TOOL_INFO: tool_info,
-                },
+                metadata=metadata,
                 inputs=parameters_for_log,
+                llm_usage=usage,
             )
         )
 
+        return usage
+
+    @staticmethod
+    def _extract_tool_usage(tool_runtime: Tool) -> LLMUsage:
+        if isinstance(tool_runtime, WorkflowTool):
+            return tool_runtime.latest_usage
+        return LLMUsage.empty_usage()
+
     @classmethod
     def _extract_variable_selector_to_variable_mapping(
         cls,

api/core/workflow/runtime/graph_runtime_state.py

@@ -5,6 +5,7 @@ import json
 from collections.abc import Mapping, Sequence
 from collections.abc import Mapping as TypingMapping
 from copy import deepcopy
+from dataclasses import dataclass
 from typing import Any, Protocol
 
 from pydantic.json import pydantic_encoder
@@ -106,6 +107,23 @@ class GraphProtocol(Protocol):
     def get_outgoing_edges(self, node_id: str) -> Sequence[object]: ...
 
 
+@dataclass(slots=True)
+class _GraphRuntimeStateSnapshot:
+    """Immutable view of a serialized runtime state snapshot."""
+
+    start_at: float
+    total_tokens: int
+    node_run_steps: int
+    llm_usage: LLMUsage
+    outputs: dict[str, Any]
+    variable_pool: VariablePool
+    has_variable_pool: bool
+    ready_queue_dump: str | None
+    graph_execution_dump: str | None
+    response_coordinator_dump: str | None
+    paused_nodes: tuple[str, ...]
+
+
 class GraphRuntimeState:
     """Mutable runtime state shared across graph execution components."""
 
@@ -293,69 +311,28 @@ class GraphRuntimeState:
 
         return json.dumps(snapshot, default=pydantic_encoder)
 
-    def loads(self, data: str | Mapping[str, Any]) -> None:
+    @classmethod
+    def from_snapshot(cls, data: str | Mapping[str, Any]) -> GraphRuntimeState:
         """Restore runtime state from a serialized snapshot."""
 
-        payload: dict[str, Any]
-        if isinstance(data, str):
-            payload = json.loads(data)
-        else:
-            payload = dict(data)
+        snapshot = cls._parse_snapshot_payload(data)
 
-        version = payload.get("version")
-        if version != "1.0":
-            raise ValueError(f"Unsupported GraphRuntimeState snapshot version: {version}")
+        state = cls(
+            variable_pool=snapshot.variable_pool,
+            start_at=snapshot.start_at,
+            total_tokens=snapshot.total_tokens,
+            llm_usage=snapshot.llm_usage,
+            outputs=snapshot.outputs,
+            node_run_steps=snapshot.node_run_steps,
+        )
+        state._apply_snapshot(snapshot)
+        return state
 
-        self._start_at = float(payload.get("start_at", 0.0))
-        total_tokens = int(payload.get("total_tokens", 0))
-        if total_tokens < 0:
-            raise ValueError("total_tokens must be non-negative")
-        self._total_tokens = total_tokens
+    def loads(self, data: str | Mapping[str, Any]) -> None:
+        """Restore runtime state from a serialized snapshot (legacy API)."""
 
-        node_run_steps = int(payload.get("node_run_steps", 0))
-        if node_run_steps < 0:
-            raise ValueError("node_run_steps must be non-negative")
-        self._node_run_steps = node_run_steps
-
-        llm_usage_payload = payload.get("llm_usage", {})
-        self._llm_usage = LLMUsage.model_validate(llm_usage_payload)
-
-        self._outputs = deepcopy(payload.get("outputs", {}))
-
-        variable_pool_payload = payload.get("variable_pool")
-        if variable_pool_payload is not None:
-            self._variable_pool = VariablePool.model_validate(variable_pool_payload)
-
-        ready_queue_payload = payload.get("ready_queue")
-        if ready_queue_payload is not None:
-            self._ready_queue = self._build_ready_queue()
-            self._ready_queue.loads(ready_queue_payload)
-        else:
-            self._ready_queue = None
-
-        graph_execution_payload = payload.get("graph_execution")
-        self._graph_execution = None
-        self._pending_graph_execution_workflow_id = None
-        if graph_execution_payload is not None:
-            try:
-                execution_payload = json.loads(graph_execution_payload)
-                self._pending_graph_execution_workflow_id = execution_payload.get("workflow_id")
-            except (json.JSONDecodeError, TypeError, AttributeError):
-                self._pending_graph_execution_workflow_id = None
-            self.graph_execution.loads(graph_execution_payload)
-
-        response_payload = payload.get("response_coordinator")
-        if response_payload is not None:
-            if self._graph is not None:
-                self.response_coordinator.loads(response_payload)
-            else:
-                self._pending_response_coordinator_dump = response_payload
-        else:
-            self._pending_response_coordinator_dump = None
-            self._response_coordinator = None
-
-        paused_nodes_payload = payload.get("paused_nodes", [])
-        self._paused_nodes = set(map(str, paused_nodes_payload))
+        snapshot = self._parse_snapshot_payload(data)
+        self._apply_snapshot(snapshot)
 
     def register_paused_node(self, node_id: str) -> None:
         """Record a node that should resume when execution is continued."""
@@ -391,3 +368,106 @@ class GraphRuntimeState:
         module = importlib.import_module("core.workflow.graph_engine.response_coordinator")
         coordinator_cls = module.ResponseStreamCoordinator
         return coordinator_cls(variable_pool=self.variable_pool, graph=graph)
+
+    # ------------------------------------------------------------------
+    # Snapshot helpers
+    # ------------------------------------------------------------------
+    @classmethod
+    def _parse_snapshot_payload(cls, data: str | Mapping[str, Any]) -> _GraphRuntimeStateSnapshot:
+        payload: dict[str, Any]
+        if isinstance(data, str):
+            payload = json.loads(data)
+        else:
+            payload = dict(data)
+
+        version = payload.get("version")
+        if version != "1.0":
+            raise ValueError(f"Unsupported GraphRuntimeState snapshot version: {version}")
+
+        start_at = float(payload.get("start_at", 0.0))
+
+        total_tokens = int(payload.get("total_tokens", 0))
+        if total_tokens < 0:
+            raise ValueError("total_tokens must be non-negative")
+
+        node_run_steps = int(payload.get("node_run_steps", 0))
+        if node_run_steps < 0:
+            raise ValueError("node_run_steps must be non-negative")
+
+        llm_usage_payload = payload.get("llm_usage", {})
+        llm_usage = LLMUsage.model_validate(llm_usage_payload)
+
+        outputs_payload = deepcopy(payload.get("outputs", {}))
+
+        variable_pool_payload = payload.get("variable_pool")
+        has_variable_pool = variable_pool_payload is not None
+        variable_pool = VariablePool.model_validate(variable_pool_payload) if has_variable_pool else VariablePool()
+
+        ready_queue_payload = payload.get("ready_queue")
+        graph_execution_payload = payload.get("graph_execution")
+        response_payload = payload.get("response_coordinator")
+        paused_nodes_payload = payload.get("paused_nodes", [])
+
+        return _GraphRuntimeStateSnapshot(
+            start_at=start_at,
+            total_tokens=total_tokens,
+            node_run_steps=node_run_steps,
+            llm_usage=llm_usage,
+            outputs=outputs_payload,
+            variable_pool=variable_pool,
+            has_variable_pool=has_variable_pool,
+            ready_queue_dump=ready_queue_payload,
+            graph_execution_dump=graph_execution_payload,
+            response_coordinator_dump=response_payload,
+            paused_nodes=tuple(map(str, paused_nodes_payload)),
+        )
+
+    def _apply_snapshot(self, snapshot: _GraphRuntimeStateSnapshot) -> None:
+        self._start_at = snapshot.start_at
+        self._total_tokens = snapshot.total_tokens
+        self._node_run_steps = snapshot.node_run_steps
+        self._llm_usage = snapshot.llm_usage.model_copy()
+        self._outputs = deepcopy(snapshot.outputs)
+        if snapshot.has_variable_pool or self._variable_pool is None:
+            self._variable_pool = snapshot.variable_pool
+
+        self._restore_ready_queue(snapshot.ready_queue_dump)
+        self._restore_graph_execution(snapshot.graph_execution_dump)
+        self._restore_response_coordinator(snapshot.response_coordinator_dump)
+        self._paused_nodes = set(snapshot.paused_nodes)
+
+    def _restore_ready_queue(self, payload: str | None) -> None:
+        if payload is not None:
+            self._ready_queue = self._build_ready_queue()
+            self._ready_queue.loads(payload)
+        else:
+            self._ready_queue = None
+
+    def _restore_graph_execution(self, payload: str | None) -> None:
+        self._graph_execution = None
+        self._pending_graph_execution_workflow_id = None
+
+        if payload is None:
+            return
+
+        try:
+            execution_payload = json.loads(payload)
+            self._pending_graph_execution_workflow_id = execution_payload.get("workflow_id")
+        except (json.JSONDecodeError, TypeError, AttributeError):
+            self._pending_graph_execution_workflow_id = None
+
+        self.graph_execution.loads(payload)
+
+    def _restore_response_coordinator(self, payload: str | None) -> None:
+        if payload is None:
+            self._pending_response_coordinator_dump = None
+            self._response_coordinator = None
+            return
+
+        if self._graph is not None:
+            self.response_coordinator.loads(payload)
+            self._pending_response_coordinator_dump = None
+            return
+
+        self._pending_response_coordinator_dump = payload
+        self._response_coordinator = None

api/core/workflow/runtime/variable_pool.py

@@ -260,7 +260,7 @@ class VariablePool(BaseModel):
             # This ensures that we can keep the id of the system variables intact.
             if self._has(selector):
                 continue
-            self.add(selector, value)  # type: ignore
+            self.add(selector, value)
 
     @classmethod
     def empty(cls) -> "VariablePool":

api/docker/entrypoint.sh

@@ -32,7 +32,8 @@ if [[ "${MODE}" == "worker" ]]; then
 
   exec celery -A celery_entrypoint.celery worker -P ${CELERY_WORKER_CLASS:-gevent} $CONCURRENCY_OPTION \
     --max-tasks-per-child ${MAX_TASKS_PER_CHILD:-50} --loglevel ${LOG_LEVEL:-INFO} \
-    -Q ${CELERY_QUEUES:-dataset,pipeline,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation}
+    -Q ${CELERY_QUEUES:-dataset,pipeline,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation} \
+    --prefetch-multiplier=1
 
 elif [[ "${MODE}" == "beat" ]]; then
   exec celery -A app.celery beat --loglevel ${LOG_LEVEL:-INFO}

api/extensions/ext_compress.py

@@ -7,7 +7,7 @@ def is_enabled() -> bool:
 
 
 def init_app(app: DifyApp):
-    from flask_compress import Compress  # type: ignore
+    from flask_compress import Compress
 
     compress = Compress()
     compress.init_app(app)

api/extensions/ext_login.py

@@ -1,15 +1,16 @@
 import json
 
-import flask_login  # type: ignore
+import flask_login
 from flask import Response, request
 from flask_login import user_loaded_from_request, user_logged_in
 from werkzeug.exceptions import NotFound, Unauthorized
 
 from configs import dify_config
+from constants import HEADER_NAME_APP_CODE
 from dify_app import DifyApp
 from extensions.ext_database import db
 from libs.passport import PassportService
-from libs.token import extract_access_token
+from libs.token import extract_access_token, extract_webapp_passport
 from models import Account, Tenant, TenantAccountJoin
 from models.model import AppMCPServer, EndUser
 from services.account_service import AccountService
@@ -61,14 +62,30 @@ def load_user_from_request(request_from_flask_login):
         logged_in_account = AccountService.load_logged_in_account(account_id=user_id)
         return logged_in_account
     elif request.blueprint == "web":
-        decoded = PassportService().verify(auth_token)
-        end_user_id = decoded.get("end_user_id")
-        if not end_user_id:
-            raise Unauthorized("Invalid Authorization token.")
-        end_user = db.session.query(EndUser).where(EndUser.id == decoded["end_user_id"]).first()
-        if not end_user:
-            raise NotFound("End user not found.")
-        return end_user
+        app_code = request.headers.get(HEADER_NAME_APP_CODE)
+        webapp_token = extract_webapp_passport(app_code, request) if app_code else None
+
+        if webapp_token:
+            decoded = PassportService().verify(webapp_token)
+            end_user_id = decoded.get("end_user_id")
+            if not end_user_id:
+                raise Unauthorized("Invalid Authorization token.")
+            end_user = db.session.query(EndUser).where(EndUser.id == end_user_id).first()
+            if not end_user:
+                raise NotFound("End user not found.")
+            return end_user
+        else:
+            if not auth_token:
+                raise Unauthorized("Invalid Authorization token.")
+            decoded = PassportService().verify(auth_token)
+            end_user_id = decoded.get("end_user_id")
+            if end_user_id:
+                end_user = db.session.query(EndUser).where(EndUser.id == end_user_id).first()
+                if not end_user:
+                    raise NotFound("End user not found.")
+                return end_user
+            else:
+                raise Unauthorized("Invalid Authorization token for web API.")
     elif request.blueprint == "mcp":
         server_code = request.view_args.get("server_code") if request.view_args else None
         if not server_code:

api/extensions/ext_migrate.py

@@ -2,7 +2,7 @@ from dify_app import DifyApp
 
 
 def init_app(app: DifyApp):
-    import flask_migrate  # type: ignore
+    import flask_migrate
 
     from extensions.ext_database import db
 

api/extensions/ext_otel.py

@@ -103,7 +103,7 @@ def init_app(app: DifyApp):
     def shutdown_tracer():
         provider = trace.get_tracer_provider()
         if hasattr(provider, "force_flush"):
-            provider.force_flush()  # ty: ignore [call-non-callable]
+            provider.force_flush()
 
     class ExceptionLoggingHandler(logging.Handler):
         """Custom logging handler that creates spans for logging.exception() calls"""

api/extensions/ext_proxy_fix.py

@@ -6,4 +6,4 @@ def init_app(app: DifyApp):
     if dify_config.RESPECT_XFORWARD_HEADERS_ENABLED:
         from werkzeug.middleware.proxy_fix import ProxyFix
 
-        app.wsgi_app = ProxyFix(app.wsgi_app, x_port=1)  # type: ignore
+        app.wsgi_app = ProxyFix(app.wsgi_app, x_port=1)  # type: ignore[method-assign]

api/extensions/ext_sentry.py

@@ -5,7 +5,7 @@ from dify_app import DifyApp
 def init_app(app: DifyApp):
     if dify_config.SENTRY_DSN:
         import sentry_sdk
-        from langfuse import parse_error  # type: ignore
+        from langfuse import parse_error
         from sentry_sdk.integrations.celery import CeleryIntegration
         from sentry_sdk.integrations.flask import FlaskIntegration
         from werkzeug.exceptions import HTTPException

api/extensions/storage/aliyun_oss_storage.py

@@ -1,7 +1,7 @@
 import posixpath
 from collections.abc import Generator
 
-import oss2 as aliyun_s3  # type: ignore
+import oss2 as aliyun_s3
 
 from configs import dify_config
 from extensions.storage.base_storage import BaseStorage

api/extensions/storage/baidu_obs_storage.py

@@ -2,9 +2,9 @@ import base64
 import hashlib
 from collections.abc import Generator
 
-from baidubce.auth.bce_credentials import BceCredentials  # type: ignore
-from baidubce.bce_client_configuration import BceClientConfiguration  # type: ignore
-from baidubce.services.bos.bos_client import BosClient  # type: ignore
+from baidubce.auth.bce_credentials import BceCredentials
+from baidubce.bce_client_configuration import BceClientConfiguration
+from baidubce.services.bos.bos_client import BosClient
 
 from configs import dify_config
 from extensions.storage.base_storage import BaseStorage

api/extensions/storage/clickzetta_volume/clickzetta_volume_storage.py

@@ -11,7 +11,7 @@ from collections.abc import Generator
 from io import BytesIO
 from pathlib import Path
 
-import clickzetta  # type: ignore[import]
+import clickzetta
 from pydantic import BaseModel, model_validator
 
 from extensions.storage.base_storage import BaseStorage

api/extensions/storage/clickzetta_volume/volume_permissions.py

@@ -34,7 +34,7 @@ class VolumePermissionManager:
         # Support two initialization methods: connection object or configuration dictionary
         if isinstance(connection_or_config, dict):
             # Create connection from configuration dictionary
-            import clickzetta  # type: ignore[import-untyped]
+            import clickzetta
 
             config = connection_or_config
             self._connection = clickzetta.connect(

api/extensions/storage/google_cloud_storage.py

@@ -3,7 +3,7 @@ import io
 import json
 from collections.abc import Generator
 
-from google.cloud import storage as google_cloud_storage  # type: ignore
+from google.cloud import storage as google_cloud_storage
 
 from configs import dify_config
 from extensions.storage.base_storage import BaseStorage

api/extensions/storage/huawei_obs_storage.py

@@ -1,6 +1,6 @@
 from collections.abc import Generator
 
-from obs import ObsClient  # type: ignore
+from obs import ObsClient
 
 from configs import dify_config
 from extensions.storage.base_storage import BaseStorage

api/extensions/storage/oracle_oci_storage.py

@@ -1,7 +1,7 @@
 from collections.abc import Generator
 
-import boto3  # type: ignore
-from botocore.exceptions import ClientError  # type: ignore
+import boto3
+from botocore.exceptions import ClientError
 
 from configs import dify_config
 from extensions.storage.base_storage import BaseStorage

api/extensions/storage/tencent_cos_storage.py

@@ -1,6 +1,6 @@
 from collections.abc import Generator
 
-from qcloud_cos import CosConfig, CosS3Client  # type: ignore
+from qcloud_cos import CosConfig, CosS3Client
 
 from configs import dify_config
 from extensions.storage.base_storage import BaseStorage

api/extensions/storage/volcengine_tos_storage.py

@@ -1,6 +1,6 @@
 from collections.abc import Generator
 
-import tos  # type: ignore
+import tos
 
 from configs import dify_config
 from extensions.storage.base_storage import BaseStorage

api/libs/external_api.py

@@ -146,6 +146,6 @@ class ExternalApi(Api):
         kwargs["doc"] = dify_config.SWAGGER_UI_PATH if dify_config.SWAGGER_UI_ENABLED else False
 
         # manual separate call on construction and init_app to ensure configs in kwargs effective
-        super().__init__(app=None, *args, **kwargs)  # type: ignore
+        super().__init__(app=None, *args, **kwargs)
         self.init_app(app, **kwargs)
         register_external_error_handlers(self)

api/libs/gmpy2_pkcs10aep_cipher.py

@@ -23,7 +23,7 @@ from hashlib import sha1
 
 import Crypto.Hash.SHA1
 import Crypto.Util.number
-import gmpy2  # type: ignore
+import gmpy2
 from Crypto import Random
 from Crypto.Signature.pss import MGF1
 from Crypto.Util.number import bytes_to_long, ceil_div, long_to_bytes
@@ -136,7 +136,7 @@ class PKCS1OAepCipher:
         # Step 3a (OS2IP)
         em_int = bytes_to_long(em)
         # Step 3b (RSAEP)
-        m_int = gmpy2.powmod(em_int, self._key.e, self._key.n)  # ty: ignore [unresolved-attribute]
+        m_int = gmpy2.powmod(em_int, self._key.e, self._key.n)
         # Step 3c (I2OSP)
         c = long_to_bytes(m_int, k)
         return c
@@ -169,7 +169,7 @@ class PKCS1OAepCipher:
         ct_int = bytes_to_long(ciphertext)
         # Step 2b (RSADP)
         # m_int = self._key._decrypt(ct_int)
-        m_int = gmpy2.powmod(ct_int, self._key.d, self._key.n)  # ty: ignore [unresolved-attribute]
+        m_int = gmpy2.powmod(ct_int, self._key.d, self._key.n)
         # Complete step 2c (I2OSP)
         em = long_to_bytes(m_int, k)
         # Step 3a
@@ -191,12 +191,12 @@ class PKCS1OAepCipher:
         # Step 3g
         one_pos = hLen + db[hLen:].find(b"\x01")
         lHash1 = db[:hLen]
-        invalid = bord(y) | int(one_pos < hLen)  # type: ignore
+        invalid = bord(y) | int(one_pos < hLen)  # type: ignore[arg-type]
         hash_compare = strxor(lHash1, lHash)
         for x in hash_compare:
-            invalid |= bord(x)  # type: ignore
+            invalid |= bord(x)  # type: ignore[arg-type]
         for x in db[hLen:one_pos]:
-            invalid |= bord(x)  # type: ignore
+            invalid |= bord(x)  # type: ignore[arg-type]
         if invalid != 0:
             raise ValueError("Incorrect decryption.")
         # Step 4

api/libs/helper.py

@@ -81,6 +81,8 @@ class AvatarUrlField(fields.Raw):
         from models import Account
 
         if isinstance(obj, Account) and obj.avatar is not None:
+            if obj.avatar.startswith(("http://", "https://")):
+                return obj.avatar
             return file_helpers.get_signed_file_url(obj.avatar)
         return None
 

api/libs/login.py

@@ -3,7 +3,7 @@ from functools import wraps
 from typing import Any
 
 from flask import current_app, g, has_request_context, request
-from flask_login.config import EXEMPT_METHODS  # type: ignore
+from flask_login.config import EXEMPT_METHODS
 from werkzeug.local import LocalProxy
 
 from configs import dify_config
@@ -87,7 +87,7 @@ def _get_user() -> EndUser | Account | None:
         if "_login_user" not in g:
             current_app.login_manager._load_user()  # type: ignore
 
-        return g._login_user  # type: ignore
+        return g._login_user
 
     return None
 

api/libs/sendgrid.py

@@ -1,8 +1,8 @@
 import logging
 
-import sendgrid  # type: ignore
+import sendgrid
 from python_http_client.exceptions import ForbiddenError, UnauthorizedError
-from sendgrid.helpers.mail import Content, Email, Mail, To  # type: ignore
+from sendgrid.helpers.mail import Content, Email, Mail, To
 
 logger = logging.getLogger(__name__)
 

api/libs/token.py

@@ -12,6 +12,7 @@ from constants import (
     COOKIE_NAME_CSRF_TOKEN,
     COOKIE_NAME_PASSPORT,
     COOKIE_NAME_REFRESH_TOKEN,
+    COOKIE_NAME_WEBAPP_ACCESS_TOKEN,
     HEADER_NAME_CSRF_TOKEN,
     HEADER_NAME_PASSPORT,
 )
@@ -37,9 +38,6 @@ def _real_cookie_name(cookie_name: str) -> str:
 
 
 def _try_extract_from_header(request: Request) -> str | None:
-    """
-    Try to extract access token from header
-    """
     auth_header = request.headers.get("Authorization")
     if auth_header:
         if " " not in auth_header:
@@ -54,40 +52,30 @@ def _try_extract_from_header(request: Request) -> str | None:
     return None
 
 
+def extract_refresh_token(request: Request) -> str | None:
+    return request.cookies.get(_real_cookie_name(COOKIE_NAME_REFRESH_TOKEN))
+
+
 def extract_csrf_token(request: Request) -> str | None:
-    """
-    Try to extract CSRF token from header or cookie.
-    """
     return request.headers.get(HEADER_NAME_CSRF_TOKEN)
 
 
 def extract_csrf_token_from_cookie(request: Request) -> str | None:
-    """
-    Try to extract CSRF token from cookie.
-    """
     return request.cookies.get(_real_cookie_name(COOKIE_NAME_CSRF_TOKEN))
 
 
 def extract_access_token(request: Request) -> str | None:
-    """
-    Try to extract access token from cookie, header or params.
-
-    Access token is either for console session or webapp passport exchange.
-    """
-
     def _try_extract_from_cookie(request: Request) -> str | None:
         return request.cookies.get(_real_cookie_name(COOKIE_NAME_ACCESS_TOKEN))
 
     return _try_extract_from_cookie(request) or _try_extract_from_header(request)
 
 
+def extract_webapp_access_token(request: Request) -> str | None:
+    return request.cookies.get(_real_cookie_name(COOKIE_NAME_WEBAPP_ACCESS_TOKEN)) or _try_extract_from_header(request)
+
+
 def extract_webapp_passport(app_code: str, request: Request) -> str | None:
-    """
-    Try to extract app token from header or params.
-
-    Webapp access token (part of passport) is only used for webapp session.
-    """
-
     def _try_extract_passport_token_from_cookie(request: Request) -> str | None:
         return request.cookies.get(_real_cookie_name(COOKIE_NAME_PASSPORT + "-" + app_code))
 
@@ -155,6 +143,10 @@ def clear_access_token_from_cookie(response: Response, samesite: str = "Lax"):
     _clear_cookie(response, COOKIE_NAME_ACCESS_TOKEN, samesite)
 
 
+def clear_webapp_access_token_from_cookie(response: Response, samesite: str = "Lax"):
+    _clear_cookie(response, COOKIE_NAME_WEBAPP_ACCESS_TOKEN, samesite)
+
+
 def clear_refresh_token_from_cookie(response: Response):
     _clear_cookie(response, COOKIE_NAME_REFRESH_TOKEN)
 

api/migrations/versions/2025_10_21_1430-ae662b25d9bc_remove_builtin_template_user.py

@@ -0,0 +1,36 @@
+"""remove-builtin-template-user
+
+Revision ID: ae662b25d9bc
+Revises: d98acf217d43
+Create Date: 2025-10-21 14:30:28.566192
+
+"""
+from alembic import op
+import models as models
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = 'ae662b25d9bc'
+down_revision = 'd98acf217d43'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+
+    with op.batch_alter_table('pipeline_built_in_templates', schema=None) as batch_op:
+        batch_op.drop_column('updated_by')
+        batch_op.drop_column('created_by')
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('pipeline_built_in_templates', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('created_by', sa.UUID(), autoincrement=False, nullable=False))
+        batch_op.add_column(sa.Column('updated_by', sa.UUID(), autoincrement=False, nullable=True))
+
+    # ### end Alembic commands ###

api/models/account.py

@@ -5,7 +5,7 @@ from datetime import datetime
 from typing import Any, Optional
 
 import sqlalchemy as sa
-from flask_login import UserMixin  # type: ignore[import-untyped]
+from flask_login import UserMixin
 from sqlalchemy import DateTime, String, func, select
 from sqlalchemy.orm import Mapped, Session, mapped_column
 from typing_extensions import deprecated

api/models/dataset.py

@@ -1239,15 +1239,6 @@ class PipelineBuiltInTemplate(Base):  # type: ignore[name-defined]
     language = mapped_column(db.String(255), nullable=False)
     created_at = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
     updated_at = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
-    created_by = mapped_column(StringUUID, nullable=False)
-    updated_by = mapped_column(StringUUID, nullable=True)
-
-    @property
-    def created_user_name(self):
-        account = db.session.query(Account).where(Account.id == self.created_by).first()
-        if account:
-            return account.name
-        return ""
 
 
 class PipelineCustomizedTemplate(Base):  # type: ignore[name-defined]

api/models/model.py

@@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Any, Literal, Optional, cast
 
 import sqlalchemy as sa
 from flask import request
-from flask_login import UserMixin  # type: ignore[import-untyped]
+from flask_login import UserMixin
 from sqlalchemy import Float, Index, PrimaryKeyConstraint, String, exists, func, select, text
 from sqlalchemy.orm import Mapped, Session, mapped_column
 

api/models/tools.py

@@ -222,7 +222,7 @@ class WorkflowToolProvider(TypeBase):
         sa.UniqueConstraint("tenant_id", "app_id", name="unique_workflow_tool_provider_app_id"),
     )
 
-    id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"), init=False)
+    id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
     # name of the workflow provider
     name: Mapped[str] = mapped_column(String(255), nullable=False)
     # label of the workflow provider

api/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dify-api"
-version = "1.9.1"
+version = "1.9.2"
 requires-python = ">=3.11,<3.13"
 
 dependencies = [

api/pyrightconfig.json

@@ -16,7 +16,25 @@
     "opentelemetry.instrumentation.requests",
     "opentelemetry.instrumentation.sqlalchemy",
     "opentelemetry.instrumentation.redis",
-    "opentelemetry.instrumentation.httpx"
+    "langfuse",
+    "cloudscraper",
+    "readabilipy",
+    "pypandoc",
+    "pypdfium2",
+    "webvtt",
+    "flask_compress",
+    "oss2",
+    "baidubce.auth.bce_credentials",
+    "baidubce.bce_client_configuration",
+    "baidubce.services.bos.bos_client",
+    "clickzetta",
+    "google.cloud",
+    "obs",
+    "qcloud_cos",
+    "tos",
+    "gmpy2",
+    "sendgrid",
+    "sendgrid.helpers.mail"
   ],
   "reportUnknownMemberType": "hint",
   "reportUnknownParameterType": "hint",
@@ -28,7 +46,7 @@
   "reportUnnecessaryComparison": "hint",
   "reportUnnecessaryIsInstance": "hint",
   "reportUntypedFunctionDecorator": "hint",
-
+  "reportUnnecessaryTypeIgnoreComment": "hint",
   "reportAttributeAccessIssue": "hint",
   "pythonVersion": "3.11",
   "pythonPlatform": "All"

api/repositories/factory.py

@@ -48,7 +48,7 @@ class DifyAPIRepositoryFactory(DifyCoreRepositoryFactory):
 
         try:
             repository_class = import_string(class_path)
-            return repository_class(session_maker=session_maker)  # type: ignore[no-any-return]
+            return repository_class(session_maker=session_maker)
         except (ImportError, Exception) as e:
             raise RepositoryImportError(
                 f"Failed to create DifyAPIWorkflowNodeExecutionRepository from '{class_path}': {e}"
@@ -77,6 +77,6 @@ class DifyAPIRepositoryFactory(DifyCoreRepositoryFactory):
 
         try:
             repository_class = import_string(class_path)
-            return repository_class(session_maker=session_maker)  # type: ignore[no-any-return]
+            return repository_class(session_maker=session_maker)
         except (ImportError, Exception) as e:
             raise RepositoryImportError(f"Failed to create APIWorkflowRunRepository from '{class_path}': {e}") from e

api/services/account_service.py

@@ -13,7 +13,7 @@ from sqlalchemy.orm import Session
 from werkzeug.exceptions import Unauthorized
 
 from configs import dify_config
-from constants.languages import language_timezone_mapping, languages
+from constants.languages import get_valid_language, language_timezone_mapping
 from events.tenant_event import tenant_was_created
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client, redis_fallback
@@ -1259,7 +1259,7 @@ class RegisterService:
         return f"member_invite:token:{token}"
 
     @classmethod
-    def setup(cls, email: str, name: str, password: str, ip_address: str):
+    def setup(cls, email: str, name: str, password: str, ip_address: str, language: str):
         """
         Setup dify
 
@@ -1269,11 +1269,10 @@ class RegisterService:
         :param ip_address: ip address
         """
         try:
-            # Register
             account = AccountService.create_account(
                 email=email,
                 name=name,
-                interface_language=languages[0],
+                interface_language=get_valid_language(language),
                 password=password,
                 is_setup=True,
             )
@@ -1315,7 +1314,7 @@ class RegisterService:
             account = AccountService.create_account(
                 email=email,
                 name=name,
-                interface_language=language or languages[0],
+                interface_language=get_valid_language(language),
                 password=password,
                 is_setup=is_setup,
             )

api/services/app_dsl_service.py

@@ -7,7 +7,7 @@ from enum import StrEnum
 from urllib.parse import urlparse
 from uuid import uuid4
 
-import yaml  # type: ignore
+import yaml
 from Crypto.Cipher import AES
 from Crypto.Util.Padding import pad, unpad
 from packaging import version
@@ -563,7 +563,7 @@ class AppDslService:
         else:
             cls._append_model_config_export_data(export_data, app_model)
 
-        return yaml.dump(export_data, allow_unicode=True)  # type: ignore
+        return yaml.dump(export_data, allow_unicode=True)
 
     @classmethod
     def _append_workflow_export_data(

api/services/audio_service.py

@@ -82,54 +82,51 @@ class AudioService:
         message_id: str | None = None,
         is_draft: bool = False,
     ):
-        from app import app
-
         def invoke_tts(text_content: str, app_model: App, voice: str | None = None, is_draft: bool = False):
-            with app.app_context():
-                if voice is None:
-                    if app_model.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}:
-                        if is_draft:
-                            workflow = WorkflowService().get_draft_workflow(app_model=app_model)
-                        else:
-                            workflow = app_model.workflow
-                        if (
-                            workflow is None
-                            or "text_to_speech" not in workflow.features_dict
-                            or not workflow.features_dict["text_to_speech"].get("enabled")
-                        ):
+            if voice is None:
+                if app_model.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}:
+                    if is_draft:
+                        workflow = WorkflowService().get_draft_workflow(app_model=app_model)
+                    else:
+                        workflow = app_model.workflow
+                    if (
+                        workflow is None
+                        or "text_to_speech" not in workflow.features_dict
+                        or not workflow.features_dict["text_to_speech"].get("enabled")
+                    ):
+                        raise ValueError("TTS is not enabled")
+
+                    voice = workflow.features_dict["text_to_speech"].get("voice")
+                else:
+                    if not is_draft:
+                        if app_model.app_model_config is None:
+                            raise ValueError("AppModelConfig not found")
+                        text_to_speech_dict = app_model.app_model_config.text_to_speech_dict
+
+                        if not text_to_speech_dict.get("enabled"):
                             raise ValueError("TTS is not enabled")
 
-                        voice = workflow.features_dict["text_to_speech"].get("voice")
-                    else:
-                        if not is_draft:
-                            if app_model.app_model_config is None:
-                                raise ValueError("AppModelConfig not found")
-                            text_to_speech_dict = app_model.app_model_config.text_to_speech_dict
+                        voice = text_to_speech_dict.get("voice")
 
-                            if not text_to_speech_dict.get("enabled"):
-                                raise ValueError("TTS is not enabled")
-
-                            voice = text_to_speech_dict.get("voice")
-
-                model_manager = ModelManager()
-                model_instance = model_manager.get_default_model_instance(
-                    tenant_id=app_model.tenant_id, model_type=ModelType.TTS
-                )
-                try:
-                    if not voice:
-                        voices = model_instance.get_tts_voices()
-                        if voices:
-                            voice = voices[0].get("value")
-                            if not voice:
-                                raise ValueError("Sorry, no voice available.")
-                        else:
+            model_manager = ModelManager()
+            model_instance = model_manager.get_default_model_instance(
+                tenant_id=app_model.tenant_id, model_type=ModelType.TTS
+            )
+            try:
+                if not voice:
+                    voices = model_instance.get_tts_voices()
+                    if voices:
+                        voice = voices[0].get("value")
+                        if not voice:
                             raise ValueError("Sorry, no voice available.")
+                    else:
+                        raise ValueError("Sorry, no voice available.")
 
-                    return model_instance.invoke_tts(
-                        content_text=text_content.strip(), user=end_user, tenant_id=app_model.tenant_id, voice=voice
-                    )
-                except Exception as e:
-                    raise e
+                return model_instance.invoke_tts(
+                    content_text=text_content.strip(), user=end_user, tenant_id=app_model.tenant_id, voice=voice
+                )
+            except Exception as e:
+                raise e
 
         if message_id:
             try:

api/services/dataset_service.py

@@ -241,9 +241,9 @@ class DatasetService:
         dataset.created_by = account.id
         dataset.updated_by = account.id
         dataset.tenant_id = tenant_id
-        dataset.embedding_model_provider = embedding_model.provider if embedding_model else None  # type: ignore
-        dataset.embedding_model = embedding_model.model if embedding_model else None  # type: ignore
-        dataset.retrieval_model = retrieval_model.model_dump() if retrieval_model else None  # type: ignore
+        dataset.embedding_model_provider = embedding_model.provider if embedding_model else None
+        dataset.embedding_model = embedding_model.model if embedding_model else None
+        dataset.retrieval_model = retrieval_model.model_dump() if retrieval_model else None
         dataset.permission = permission or DatasetPermissionEnum.ONLY_ME
         dataset.provider = provider
         db.session.add(dataset)
@@ -1416,6 +1416,8 @@ class DocumentService:
         # check document limit
         assert isinstance(current_user, Account)
         assert current_user.current_tenant_id is not None
+        assert knowledge_config.data_source
+        assert knowledge_config.data_source.info_list
 
         features = FeatureService.get_features(current_user.current_tenant_id)
 
@@ -1424,15 +1426,18 @@ class DocumentService:
                 count = 0
                 if knowledge_config.data_source:
                     if knowledge_config.data_source.info_list.data_source_type == "upload_file":
-                        upload_file_list = knowledge_config.data_source.info_list.file_info_list.file_ids  # type: ignore
+                        if not knowledge_config.data_source.info_list.file_info_list:
+                            raise ValueError("File source info is required")
+                        upload_file_list = knowledge_config.data_source.info_list.file_info_list.file_ids
                         count = len(upload_file_list)
                     elif knowledge_config.data_source.info_list.data_source_type == "notion_import":
-                        notion_info_list = knowledge_config.data_source.info_list.notion_info_list
-                        for notion_info in notion_info_list:  # type: ignore
+                        notion_info_list = knowledge_config.data_source.info_list.notion_info_list or []
+                        for notion_info in notion_info_list:
                             count = count + len(notion_info.pages)
                     elif knowledge_config.data_source.info_list.data_source_type == "website_crawl":
                         website_info = knowledge_config.data_source.info_list.website_info_list
-                        count = len(website_info.urls)  # type: ignore
+                        assert website_info
+                        count = len(website_info.urls)
                     batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT)
 
                     if features.billing.subscription.plan == "sandbox" and count > 1:
@@ -1444,7 +1449,7 @@ class DocumentService:
 
         # if dataset is empty, update dataset data_source_type
         if not dataset.data_source_type:
-            dataset.data_source_type = knowledge_config.data_source.info_list.data_source_type  # type: ignore
+            dataset.data_source_type = knowledge_config.data_source.info_list.data_source_type
 
         if not dataset.indexing_technique:
             if knowledge_config.indexing_technique not in Dataset.INDEXING_TECHNIQUE_LIST:
@@ -1481,7 +1486,7 @@ class DocumentService:
                         knowledge_config.retrieval_model.model_dump()
                         if knowledge_config.retrieval_model
                         else default_retrieval_model
-                    )  # type: ignore
+                    )
 
         documents = []
         if knowledge_config.original_document_id:
@@ -1523,11 +1528,14 @@ class DocumentService:
                     db.session.flush()
             lock_name = f"add_document_lock_dataset_id_{dataset.id}"
             with redis_client.lock(lock_name, timeout=600):
+                assert dataset_process_rule
                 position = DocumentService.get_documents_position(dataset.id)
                 document_ids = []
                 duplicate_document_ids = []
-                if knowledge_config.data_source.info_list.data_source_type == "upload_file":  # type: ignore
-                    upload_file_list = knowledge_config.data_source.info_list.file_info_list.file_ids  # type: ignore
+                if knowledge_config.data_source.info_list.data_source_type == "upload_file":
+                    if not knowledge_config.data_source.info_list.file_info_list:
+                        raise ValueError("File source info is required")
+                    upload_file_list = knowledge_config.data_source.info_list.file_info_list.file_ids
                     for file_id in upload_file_list:
                         file = (
                             db.session.query(UploadFile)
@@ -1540,7 +1548,7 @@ class DocumentService:
                             raise FileNotExistsError()
 
                         file_name = file.name
-                        data_source_info = {
+                        data_source_info: dict[str, str | bool] = {
                             "upload_file_id": file_id,
                         }
                         # check duplicate
@@ -1557,7 +1565,7 @@ class DocumentService:
                                 .first()
                             )
                             if document:
-                                document.dataset_process_rule_id = dataset_process_rule.id  # type: ignore
+                                document.dataset_process_rule_id = dataset_process_rule.id
                                 document.updated_at = naive_utc_now()
                                 document.created_from = created_from
                                 document.doc_form = knowledge_config.doc_form
@@ -1571,8 +1579,8 @@ class DocumentService:
                                 continue
                         document = DocumentService.build_document(
                             dataset,
-                            dataset_process_rule.id,  # type: ignore
-                            knowledge_config.data_source.info_list.data_source_type,  # type: ignore
+                            dataset_process_rule.id,
+                            knowledge_config.data_source.info_list.data_source_type,
                             knowledge_config.doc_form,
                             knowledge_config.doc_language,
                             data_source_info,
@@ -1587,7 +1595,7 @@ class DocumentService:
                         document_ids.append(document.id)
                         documents.append(document)
                         position += 1
-                elif knowledge_config.data_source.info_list.data_source_type == "notion_import":  # type: ignore
+                elif knowledge_config.data_source.info_list.data_source_type == "notion_import":
                     notion_info_list = knowledge_config.data_source.info_list.notion_info_list  # type: ignore
                     if not notion_info_list:
                         raise ValueError("No notion info list found.")
@@ -1616,15 +1624,15 @@ class DocumentService:
                                     "credential_id": notion_info.credential_id,
                                     "notion_workspace_id": workspace_id,
                                     "notion_page_id": page.page_id,
-                                    "notion_page_icon": page.page_icon.model_dump() if page.page_icon else None,
+                                    "notion_page_icon": page.page_icon.model_dump() if page.page_icon else None,  # type: ignore
                                     "type": page.type,
                                 }
                                 # Truncate page name to 255 characters to prevent DB field length errors
                                 truncated_page_name = page.page_name[:255] if page.page_name else "nopagename"
                                 document = DocumentService.build_document(
                                     dataset,
-                                    dataset_process_rule.id,  # type: ignore
-                                    knowledge_config.data_source.info_list.data_source_type,  # type: ignore
+                                    dataset_process_rule.id,
+                                    knowledge_config.data_source.info_list.data_source_type,
                                     knowledge_config.doc_form,
                                     knowledge_config.doc_language,
                                     data_source_info,
@@ -1644,8 +1652,8 @@ class DocumentService:
                     # delete not selected documents
                     if len(exist_document) > 0:
                         clean_notion_document_task.delay(list(exist_document.values()), dataset.id)
-                elif knowledge_config.data_source.info_list.data_source_type == "website_crawl":  # type: ignore
-                    website_info = knowledge_config.data_source.info_list.website_info_list  # type: ignore
+                elif knowledge_config.data_source.info_list.data_source_type == "website_crawl":
+                    website_info = knowledge_config.data_source.info_list.website_info_list
                     if not website_info:
                         raise ValueError("No website info list found.")
                     urls = website_info.urls
@@ -1663,8 +1671,8 @@ class DocumentService:
                             document_name = url
                         document = DocumentService.build_document(
                             dataset,
-                            dataset_process_rule.id,  # type: ignore
-                            knowledge_config.data_source.info_list.data_source_type,  # type: ignore
+                            dataset_process_rule.id,
+                            knowledge_config.data_source.info_list.data_source_type,
                             knowledge_config.doc_form,
                             knowledge_config.doc_language,
                             data_source_info,
@@ -2071,7 +2079,7 @@ class DocumentService:
         # update document data source
         if document_data.data_source:
             file_name = ""
-            data_source_info = {}
+            data_source_info: dict[str, str | bool] = {}
             if document_data.data_source.info_list.data_source_type == "upload_file":
                 if not document_data.data_source.info_list.file_info_list:
                     raise ValueError("No file info list found.")
@@ -2128,7 +2136,7 @@ class DocumentService:
                             "url": url,
                             "provider": website_info.provider,
                             "job_id": website_info.job_id,
-                            "only_main_content": website_info.only_main_content,  # type: ignore
+                            "only_main_content": website_info.only_main_content,
                             "mode": "crawl",
                         }
             document.data_source_type = document_data.data_source.info_list.data_source_type
@@ -2154,7 +2162,7 @@ class DocumentService:
 
         db.session.query(DocumentSegment).filter_by(document_id=document.id).update(
             {DocumentSegment.status: "re_segment"}
-        )  # type: ignore
+        )
         db.session.commit()
         # trigger async task
         document_indexing_update_task.delay(document.dataset_id, document.id)
@@ -2164,25 +2172,26 @@ class DocumentService:
     def save_document_without_dataset_id(tenant_id: str, knowledge_config: KnowledgeConfig, account: Account):
         assert isinstance(current_user, Account)
         assert current_user.current_tenant_id is not None
+        assert knowledge_config.data_source
 
         features = FeatureService.get_features(current_user.current_tenant_id)
 
         if features.billing.enabled:
             count = 0
-            if knowledge_config.data_source.info_list.data_source_type == "upload_file":  # type: ignore
+            if knowledge_config.data_source.info_list.data_source_type == "upload_file":
                 upload_file_list = (
-                    knowledge_config.data_source.info_list.file_info_list.file_ids  # type: ignore
-                    if knowledge_config.data_source.info_list.file_info_list  # type: ignore
+                    knowledge_config.data_source.info_list.file_info_list.file_ids
+                    if knowledge_config.data_source.info_list.file_info_list
                     else []
                 )
                 count = len(upload_file_list)
-            elif knowledge_config.data_source.info_list.data_source_type == "notion_import":  # type: ignore
-                notion_info_list = knowledge_config.data_source.info_list.notion_info_list  # type: ignore
+            elif knowledge_config.data_source.info_list.data_source_type == "notion_import":
+                notion_info_list = knowledge_config.data_source.info_list.notion_info_list
                 if notion_info_list:
                     for notion_info in notion_info_list:
                         count = count + len(notion_info.pages)
-            elif knowledge_config.data_source.info_list.data_source_type == "website_crawl":  # type: ignore
-                website_info = knowledge_config.data_source.info_list.website_info_list  # type: ignore
+            elif knowledge_config.data_source.info_list.data_source_type == "website_crawl":
+                website_info = knowledge_config.data_source.info_list.website_info_list
                 if website_info:
                     count = len(website_info.urls)
             if features.billing.subscription.plan == "sandbox" and count > 1:
@@ -2196,9 +2205,11 @@ class DocumentService:
         dataset_collection_binding_id = None
         retrieval_model = None
         if knowledge_config.indexing_technique == "high_quality":
+            assert knowledge_config.embedding_model_provider
+            assert knowledge_config.embedding_model
             dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
-                knowledge_config.embedding_model_provider,  # type: ignore
-                knowledge_config.embedding_model,  # type: ignore
+                knowledge_config.embedding_model_provider,
+                knowledge_config.embedding_model,
             )
             dataset_collection_binding_id = dataset_collection_binding.id
         if knowledge_config.retrieval_model:
@@ -2215,7 +2226,7 @@ class DocumentService:
         dataset = Dataset(
             tenant_id=tenant_id,
             name="",
-            data_source_type=knowledge_config.data_source.info_list.data_source_type,  # type: ignore
+            data_source_type=knowledge_config.data_source.info_list.data_source_type,
             indexing_technique=knowledge_config.indexing_technique,
             created_by=account.id,
             embedding_model=knowledge_config.embedding_model,
@@ -2224,7 +2235,7 @@ class DocumentService:
             retrieval_model=retrieval_model.model_dump() if retrieval_model else None,
         )
 
-        db.session.add(dataset)  # type: ignore
+        db.session.add(dataset)
         db.session.flush()
 
         documents, batch = DocumentService.save_document_with_dataset_id(dataset, knowledge_config, account)

api/services/feature_service.py

@@ -174,6 +174,7 @@ class FeatureService:
 
         if dify_config.ENTERPRISE_ENABLED:
             features.webapp_copyright_enabled = True
+            features.knowledge_pipeline.publish_enabled = True
             cls._fulfill_params_from_workspace_info(features, tenant_id)
 
         return features

api/services/hit_testing_service.py

@@ -88,7 +88,7 @@ class HitTestingService:
         db.session.add(dataset_query)
         db.session.commit()
 
-        return cls.compact_retrieve_response(query, all_documents)  # type: ignore
+        return cls.compact_retrieve_response(query, all_documents)
 
     @classmethod
     def external_retrieve(

api/services/knowledge_service.py

@@ -1,4 +1,4 @@
-import boto3  # type: ignore
+import boto3
 
 from configs import dify_config
 

api/services/metadata_service.py

@@ -89,7 +89,7 @@ class MetadataService:
                     document.doc_metadata = doc_metadata
                     db.session.add(document)
             db.session.commit()
-            return metadata  # type: ignore
+            return metadata
         except Exception:
             logger.exception("Update metadata name failed")
         finally:

api/services/model_provider_service.py

@@ -137,7 +137,7 @@ class ModelProviderService:
         :return:
         """
         provider_configuration = self._get_provider_configuration(tenant_id, provider)
-        return provider_configuration.get_provider_credential(credential_id=credential_id)  # type: ignore
+        return provider_configuration.get_provider_credential(credential_id=credential_id)
 
     def validate_provider_credentials(self, tenant_id: str, provider: str, credentials: dict):
         """
@@ -225,7 +225,7 @@ class ModelProviderService:
         :return:
         """
         provider_configuration = self._get_provider_configuration(tenant_id, provider)
-        return provider_configuration.get_custom_model_credential(  # type: ignore
+        return provider_configuration.get_custom_model_credential(
             model_type=ModelType.value_of(model_type), model=model, credential_id=credential_id
         )
 

api/services/plugin/plugin_migration.py

@@ -146,7 +146,7 @@ class PluginMigration:
                     futures.append(
                         thread_pool.submit(
                             process_tenant,
-                            current_app._get_current_object(),  # type: ignore[attr-defined]
+                            current_app._get_current_object(),  # type: ignore
                             tenant_id,
                         )
                     )