fix typo

fix: rm useless script

.devcontainer/devcontainer.json

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

.github/workflows/style.yml

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

.gitignore

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

.vscode/launch.json.template

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

README.md

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

api/.env.example

@@ -156,9 +156,6 @@ SUPABASE_URL=your-server-url
 # CORS configuration
 WEB_API_CORS_ALLOW_ORIGINS=http://localhost:3000,*
 CONSOLE_CORS_ALLOW_ORIGINS=http://localhost:3000,*
-# Set COOKIE_DOMAIN when the console frontend and API are on different subdomains.
-# Provide the registrable domain (e.g. example.com); leading dots are optional.
-COOKIE_DOMAIN=
 
 # Vector database configuration
 # Supported values are `weaviate`, `qdrant`, `milvus`, `myscale`, `relyt`, `pgvector`, `pgvecto-rs`, `chroma`, `opensearch`, `oracle`, `tencent`, `elasticsearch`, `elasticsearch-ja`, `analyticdb`, `couchbase`, `vikingdb`, `oceanbase`, `opengauss`, `tablestore`,`vastbase`,`tidb`,`tidb_on_qdrant`,`baidu`,`lindorm`,`huawei_cloud`,`upstash`, `matrixone`.

api/.vscode/launch.json.example

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

api/README.md

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

api/commands.py

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

api/configs/feature/__init__.py

@@ -337,11 +337,6 @@ class HttpConfig(BaseSettings):
     HTTP-related configurations for the application
     """
 
-    COOKIE_DOMAIN: str = Field(
-        description="Explicit cookie domain for console/service cookies when sharing across subdomains",
-        default="",
-    )
-
     API_COMPRESSION_ENABLED: bool = Field(
         description="Enable or disable gzip compression for HTTP responses",
         default=False,
@@ -1112,6 +1107,13 @@ class SwaggerUIConfig(BaseSettings):
     )
 
 
+class TenantSelfTaskQueueConfig(BaseSettings):
+    TENANT_SELF_TASK_QUEUE_PULL_SIZE: int = Field(
+        description="Default batch size for tenant self task queue pull operations",
+        default=1,
+    )
+
+
 class FeatureConfig(
     # place the configs in alphabet order
     AppExecutionConfig,
@@ -1136,6 +1138,7 @@ class FeatureConfig(
     RagEtlConfig,
     RepositoryConfig,
     SecurityConfig,
+    TenantSelfTaskQueueConfig,
     ToolConfig,
     UpdateConfig,
     WorkflowConfig,

api/controllers/console/app/workflow_statistic.py

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

api/controllers/console/auth/login.py

@@ -29,7 +29,6 @@ from libs.token import (
     clear_access_token_from_cookie,
     clear_csrf_token_from_cookie,
     clear_refresh_token_from_cookie,
-    extract_refresh_token,
     set_access_token_to_cookie,
     set_csrf_token_to_cookie,
     set_refresh_token_to_cookie,
@@ -271,7 +270,7 @@ class EmailCodeLoginApi(Resource):
 class RefreshTokenApi(Resource):
     def post(self):
         # Get refresh token from cookie instead of request body
-        refresh_token = extract_refresh_token(request)
+        refresh_token = request.cookies.get("refresh_token")
 
         if not refresh_token:
             return {"result": "fail", "message": "No refresh token provided"}, 401

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

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

api/controllers/console/extension.py

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

api/controllers/console/files.py

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

api/controllers/console/workspace/tool_providers.py

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

api/controllers/files/image_preview.py

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

api/controllers/files/tool_files.py

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

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

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

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

@@ -41,18 +41,14 @@ from core.workflow.repositories.workflow_execution_repository import WorkflowExe
 from core.workflow.repositories.workflow_node_execution_repository import WorkflowNodeExecutionRepository
 from core.workflow.variable_loader import DUMMY_VARIABLE_LOADER, VariableLoader
 from extensions.ext_database import db
-from extensions.ext_redis import redis_client
 from libs.flask_utils import preserve_flask_contexts
 from models import Account, EndUser, Workflow, WorkflowNodeExecutionTriggeredFrom
 from models.dataset import Document, DocumentPipelineExecutionLog, Pipeline
 from models.enums import WorkflowRunTriggeredFrom
 from models.model import AppMode
 from services.datasource_provider_service import DatasourceProviderService
-from services.feature_service import FeatureService
-from services.file_service import FileService
+from services.rag_pipeline.rag_pipeline_task_proxy import RagPipelineTaskProxy
 from services.workflow_draft_variable_service import DraftVarLoader, WorkflowDraftVariableService
-from tasks.rag_pipeline.priority_rag_pipeline_run_task import priority_rag_pipeline_run_task
-from tasks.rag_pipeline.rag_pipeline_run_task import rag_pipeline_run_task
 
 logger = logging.getLogger(__name__)
 
@@ -248,34 +244,7 @@ class PipelineGenerator(BaseAppGenerator):
                 )
 
         if rag_pipeline_invoke_entities:
-            # store the rag_pipeline_invoke_entities to object storage
-            text = [item.model_dump() for item in rag_pipeline_invoke_entities]
-            name = "rag_pipeline_invoke_entities.json"
-            # Convert list to proper JSON string
-            json_text = json.dumps(text)
-            upload_file = FileService(db.engine).upload_text(json_text, name, user.id, dataset.tenant_id)
-            features = FeatureService.get_features(dataset.tenant_id)
-            if features.billing.enabled and features.billing.subscription.plan == "sandbox":
-                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}"
-
-                if redis_client.get(tenant_pipeline_task_key):
-                    # Add to waiting queue using List operations (lpush)
-                    redis_client.lpush(tenant_self_pipeline_task_queue, upload_file.id)
-                else:
-                    # Set flag and execute task
-                    redis_client.set(tenant_pipeline_task_key, 1, ex=60 * 60)
-                    rag_pipeline_run_task.delay(  # type: ignore
-                        rag_pipeline_invoke_entities_file_id=upload_file.id,
-                        tenant_id=dataset.tenant_id,
-                    )
-
-            else:
-                priority_rag_pipeline_run_task.delay(  # type: ignore
-                    rag_pipeline_invoke_entities_file_id=upload_file.id,
-                    tenant_id=dataset.tenant_id,
-                )
-
+            RagPipelineTaskProxy(dataset.tenant_id, user.id, rag_pipeline_invoke_entities).delay()
         # return batch, dataset, documents
         return {
             "batch": batch,

api/core/entities/document_task.py

@@ -0,0 +1,13 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class DocumentTask:
+    """Document task entity for document indexing operations.
+    
+    This class represents a document indexing task that can be queued
+    and processed by the document indexing system.
+    """
+    tenant_id: str
+    dataset_id: str
+    document_ids: list[str]

api/core/entities/mcp_provider.py

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

api/core/mcp/auth/auth_flow.py

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

api/core/mcp/auth/auth_provider.py

@@ -0,0 +1,77 @@
+from configs import dify_config
+from core.mcp.types import (
+    OAuthClientInformation,
+    OAuthClientInformationFull,
+    OAuthClientMetadata,
+    OAuthTokens,
+)
+from models.tools import MCPToolProvider
+from services.tools.mcp_tools_manage_service import MCPToolManageService
+
+
+class OAuthClientProvider:
+    mcp_provider: MCPToolProvider
+
+    def __init__(self, provider_id: str, tenant_id: str, for_list: bool = False):
+        if for_list:
+            self.mcp_provider = MCPToolManageService.get_mcp_provider_by_provider_id(provider_id, tenant_id)
+        else:
+            self.mcp_provider = MCPToolManageService.get_mcp_provider_by_server_identifier(provider_id, tenant_id)
+
+    @property
+    def redirect_url(self) -> str:
+        """The URL to redirect the user agent to after authorization."""
+        return dify_config.CONSOLE_API_URL + "/console/api/mcp/oauth/callback"
+
+    @property
+    def client_metadata(self) -> OAuthClientMetadata:
+        """Metadata about this OAuth client."""
+        return OAuthClientMetadata(
+            redirect_uris=[self.redirect_url],
+            token_endpoint_auth_method="none",
+            grant_types=["authorization_code", "refresh_token"],
+            response_types=["code"],
+            client_name="Dify",
+            client_uri="https://github.com/langgenius/dify",
+        )
+
+    def client_information(self) -> OAuthClientInformation | None:
+        """Loads information about this OAuth client."""
+        client_information = self.mcp_provider.decrypted_credentials.get("client_information", {})
+        if not client_information:
+            return None
+        return OAuthClientInformation.model_validate(client_information)
+
+    def save_client_information(self, client_information: OAuthClientInformationFull):
+        """Saves client information after dynamic registration."""
+        MCPToolManageService.update_mcp_provider_credentials(
+            self.mcp_provider,
+            {"client_information": client_information.model_dump()},
+        )
+
+    def tokens(self) -> OAuthTokens | None:
+        """Loads any existing OAuth tokens for the current session."""
+        credentials = self.mcp_provider.decrypted_credentials
+        if not credentials:
+            return None
+        return OAuthTokens(
+            access_token=credentials.get("access_token", ""),
+            token_type=credentials.get("token_type", "Bearer"),
+            expires_in=int(credentials.get("expires_in", "3600") or 3600),
+            refresh_token=credentials.get("refresh_token", ""),
+        )
+
+    def save_tokens(self, tokens: OAuthTokens):
+        """Stores new OAuth tokens for the current session."""
+        # update mcp provider credentials
+        token_dict = tokens.model_dump()
+        MCPToolManageService.update_mcp_provider_credentials(self.mcp_provider, token_dict, authed=True)
+
+    def save_code_verifier(self, code_verifier: str):
+        """Saves a PKCE code verifier for the current session."""
+        MCPToolManageService.update_mcp_provider_credentials(self.mcp_provider, {"code_verifier": code_verifier})
+
+    def code_verifier(self) -> str:
+        """Loads the PKCE code verifier for the current session."""
+        # get code verifier from mcp provider credentials
+        return str(self.mcp_provider.decrypted_credentials.get("code_verifier", ""))

api/core/mcp/auth_client.py

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

api/core/mcp/client/sse_client.py

@@ -46,7 +46,7 @@ class SSETransport:
         url: str,
         headers: dict[str, Any] | None = None,
         timeout: float = 5.0,
-        sse_read_timeout: float = 1 * 60,
+        sse_read_timeout: float = 5 * 60,
     ):
         """Initialize the SSE transport.
 
@@ -255,7 +255,7 @@ def sse_client(
     url: str,
     headers: dict[str, Any] | None = None,
     timeout: float = 5.0,
-    sse_read_timeout: float = 1 * 60,
+    sse_read_timeout: float = 5 * 60,
 ) -> Generator[tuple[ReadQueue, WriteQueue], None, None]:
     """
     Client transport for SSE.
@@ -276,34 +276,31 @@ def sse_client(
     read_queue: ReadQueue | None = None
     write_queue: WriteQueue | None = None
 
-    executor = ThreadPoolExecutor()
-    try:
-        with create_ssrf_proxy_mcp_http_client(headers=transport.headers) as client:
-            with ssrf_proxy_sse_connect(
-                url, timeout=httpx.Timeout(timeout, read=sse_read_timeout), client=client
-            ) as event_source:
-                event_source.response.raise_for_status()
+    with ThreadPoolExecutor() as executor:
+        try:
+            with create_ssrf_proxy_mcp_http_client(headers=transport.headers) as client:
+                with ssrf_proxy_sse_connect(
+                    url, timeout=httpx.Timeout(timeout, read=sse_read_timeout), client=client
+                ) as event_source:
+                    event_source.response.raise_for_status()
 
-                read_queue, write_queue = transport.connect(executor, client, event_source)
+                    read_queue, write_queue = transport.connect(executor, client, event_source)
 
-                yield read_queue, write_queue
+                    yield read_queue, write_queue
 
-    except httpx.HTTPStatusError as exc:
-        if exc.response.status_code == 401:
-            raise MCPAuthError()
-        raise MCPConnectionError()
-    except Exception:
-        logger.exception("Error connecting to SSE endpoint")
-        raise
-    finally:
-        # Clean up queues
-        if read_queue:
-            read_queue.put(None)
-        if write_queue:
-            write_queue.put(None)
-
-        # Shutdown executor without waiting to prevent hanging
-        executor.shutdown(wait=False)
+        except httpx.HTTPStatusError as exc:
+            if exc.response.status_code == 401:
+                raise MCPAuthError()
+            raise MCPConnectionError()
+        except Exception:
+            logger.exception("Error connecting to SSE endpoint")
+            raise
+        finally:
+            # Clean up queues
+            if read_queue:
+                read_queue.put(None)
+            if write_queue:
+                write_queue.put(None)
 
 
 def send_message(http_client: httpx.Client, endpoint_url: str, session_message: SessionMessage):

api/core/mcp/client/streamable_client.py

@@ -434,48 +434,45 @@ def streamablehttp_client(
     server_to_client_queue: ServerToClientQueue = queue.Queue()  # For messages FROM server TO client
     client_to_server_queue: ClientToServerQueue = queue.Queue()  # For messages FROM client TO server
 
-    executor = ThreadPoolExecutor(max_workers=2)
-    try:
-        with create_ssrf_proxy_mcp_http_client(
-            headers=transport.request_headers,
-            timeout=httpx.Timeout(transport.timeout, read=transport.sse_read_timeout),
-        ) as client:
-            # Define callbacks that need access to thread pool
-            def start_get_stream():
-                """Start a worker thread to handle server-initiated messages."""
-                executor.submit(transport.handle_get_stream, client, server_to_client_queue)
-
-            # Start the post_writer worker thread
-            executor.submit(
-                transport.post_writer,
-                client,
-                client_to_server_queue,  # Queue for messages FROM client TO server
-                server_to_client_queue,  # Queue for messages FROM server TO client
-                start_get_stream,
-            )
-
-            try:
-                yield (
-                    server_to_client_queue,  # Queue for receiving messages FROM server
-                    client_to_server_queue,  # Queue for sending messages TO server
-                    transport.get_session_id,
-                )
-            finally:
-                if transport.session_id and terminate_on_close:
-                    transport.terminate_session(client)
-
-                # Signal threads to stop
-                client_to_server_queue.put(None)
-    finally:
-        # Clear any remaining items and add None sentinel to unblock any waiting threads
+    with ThreadPoolExecutor(max_workers=2) as executor:
         try:
-            while not client_to_server_queue.empty():
-                client_to_server_queue.get_nowait()
-        except queue.Empty:
-            pass
+            with create_ssrf_proxy_mcp_http_client(
+                headers=transport.request_headers,
+                timeout=httpx.Timeout(transport.timeout, read=transport.sse_read_timeout),
+            ) as client:
+                # Define callbacks that need access to thread pool
+                def start_get_stream():
+                    """Start a worker thread to handle server-initiated messages."""
+                    executor.submit(transport.handle_get_stream, client, server_to_client_queue)
 
-        client_to_server_queue.put(None)
-        server_to_client_queue.put(None)
+                # Start the post_writer worker thread
+                executor.submit(
+                    transport.post_writer,
+                    client,
+                    client_to_server_queue,  # Queue for messages FROM client TO server
+                    server_to_client_queue,  # Queue for messages FROM server TO client
+                    start_get_stream,
+                )
 
-        # Shutdown executor without waiting to prevent hanging
-        executor.shutdown(wait=False)
+                try:
+                    yield (
+                        server_to_client_queue,  # Queue for receiving messages FROM server
+                        client_to_server_queue,  # Queue for sending messages TO server
+                        transport.get_session_id,
+                    )
+                finally:
+                    if transport.session_id and terminate_on_close:
+                        transport.terminate_session(client)
+
+                    # Signal threads to stop
+                    client_to_server_queue.put(None)
+        finally:
+            # Clear any remaining items and add None sentinel to unblock any waiting threads
+            try:
+                while not client_to_server_queue.empty():
+                    client_to_server_queue.get_nowait()
+            except queue.Empty:
+                pass
+
+            client_to_server_queue.put(None)
+            server_to_client_queue.put(None)

api/core/mcp/entities.py

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

api/core/mcp/error.py

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

api/core/mcp/mcp_client.py

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

api/core/mcp/session/base_session.py

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

api/core/mcp/session/client_session.py

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

api/core/mcp/types.py

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

api/core/memory/token_buffer_memory.py

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

api/core/ops/ops_trace_manager.py

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

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

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

api/core/rag/pipeline/queue.py

@@ -0,0 +1,89 @@
+import json
+from dataclasses import dataclass
+from typing import Any
+
+from extensions.ext_redis import redis_client
+
+TASK_WRAPPER_PREFIX = "__WRAPPER__:"
+
+
+@dataclass
+class TaskWrapper:
+    data: Any
+    
+    def serialize(self) -> str:
+        return json.dumps(self.data, ensure_ascii=False)
+    
+    @classmethod
+    def deserialize(cls, serialized_data: str) -> 'TaskWrapper':
+        data = json.loads(serialized_data)
+        return cls(data)
+
+
+class TenantSelfTaskQueue:
+    """
+    Simple queue for tenant self tasks, used for tenant self task isolation.
+    It uses Redis list to store tasks, and Redis key to store task waiting flag.
+    Support tasks that can be serialized by json.
+    """
+    DEFAULT_TASK_TTL = 60 * 60
+
+    def __init__(self, tenant_id: str, unique_key: str):
+        self.tenant_id = tenant_id
+        self.unique_key = unique_key
+        self.queue = f"tenant_self_{unique_key}_task_queue:{tenant_id}"
+        self.task_key = f"tenant_{unique_key}_task:{tenant_id}"
+
+    def get_task_key(self):
+        return redis_client.get(self.task_key)
+
+    def set_task_waiting_time(self, ttl: int | None = None):
+        ttl = ttl or self.DEFAULT_TASK_TTL
+        redis_client.setex(self.task_key, ttl, 1)
+
+    def delete_task_key(self):
+        redis_client.delete(self.task_key)
+
+    def push_tasks(self, tasks: list):
+        serialized_tasks = []
+        for task in tasks:
+            # Store str list directly, maintaining full compatibility for pipeline scenarios
+            if isinstance(task, str):
+                serialized_tasks.append(task)
+            else:
+                # Use TaskWrapper to do JSON serialization, add prefix for identification
+                wrapper = TaskWrapper(task)
+                serialized_data = wrapper.serialize()
+                serialized_tasks.append(f"{TASK_WRAPPER_PREFIX}{serialized_data}")
+        
+        redis_client.lpush(self.queue, *serialized_tasks)
+    
+    def pull_tasks(self, count: int = 1) -> list:
+        if count <= 0:
+            return []
+        
+        tasks = []
+        for _ in range(count):
+            serialized_task = redis_client.rpop(self.queue)
+            if not serialized_task:
+                break
+            
+            if isinstance(serialized_task, bytes):
+                serialized_task = serialized_task.decode('utf-8')
+            
+            # Check if use TaskWrapper or not
+            if serialized_task.startswith(TASK_WRAPPER_PREFIX):
+                try:
+                    wrapper_data = serialized_task[len(TASK_WRAPPER_PREFIX):]
+                    wrapper = TaskWrapper.deserialize(wrapper_data)
+                    tasks.append(wrapper.data)
+                except (json.JSONDecodeError, TypeError, ValueError):
+                    tasks.append(serialized_task)
+            else:
+                tasks.append(serialized_task)
+        
+        return tasks
+
+    def get_next_task(self):
+        tasks = self.pull_tasks(1)
+        return tasks[0] if tasks else None

api/core/tools/__base/tool.py

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

api/core/tools/entities/api_entities.py

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

api/core/tools/mcp_tool/provider.py

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

api/core/tools/mcp_tool/tool.py

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

api/core/tools/tool_manager.py

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

api/core/tools/workflow_as_tool/provider.py

@@ -31,7 +31,6 @@ VARIABLE_TO_PARAMETER_TYPE_MAPPING = {
     VariableEntityType.PARAGRAPH: ToolParameter.ToolParameterType.STRING,
     VariableEntityType.SELECT: ToolParameter.ToolParameterType.SELECT,
     VariableEntityType.NUMBER: ToolParameter.ToolParameterType.NUMBER,
-    VariableEntityType.CHECKBOX: ToolParameter.ToolParameterType.BOOLEAN,
     VariableEntityType.FILE: ToolParameter.ToolParameterType.FILE,
     VariableEntityType.FILE_LIST: ToolParameter.ToolParameterType.FILES,
 }

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

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

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

@@ -171,7 +171,6 @@ def _extract_text_by_file_extension(*, file_content: bytes, file_extension: str)
             ".txt"
             | ".markdown"
             | ".md"
-            | ".mdx"
             | ".html"
             | ".htm"
             | ".xml"

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

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

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

@@ -98,7 +98,6 @@ class IterationNode(LLMUsageTrackingMixin, Node):
                 "is_parallel": False,
                 "parallel_nums": 10,
                 "error_handle_mode": ErrorHandleMode.TERMINATED,
-                "flatten_output": True,
             },
         }
 
@@ -412,14 +411,7 @@ class IterationNode(LLMUsageTrackingMixin, Node):
         """
         Flatten the outputs list if all elements are lists.
         This maintains backward compatibility with version 1.8.1 behavior.
-
-        If flatten_output is False, returns outputs as-is (nested structure).
-        If flatten_output is True (default), flattens the list if all elements are lists.
         """
-        # If flatten_output is disabled, return outputs as-is
-        if not self._node_data.flatten_output:
-            return outputs
-
         if not outputs:
             return outputs
 

api/docker/entrypoint.sh

@@ -32,7 +32,7 @@ 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,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation,priority_pipeline,pipeline} \
+    -Q ${CELERY_QUEUES:-dataset,pipeline,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation} \
     --prefetch-multiplier=1
 
 elif [[ "${MODE}" == "beat" ]]; then

api/extensions/ext_login.py

@@ -6,11 +6,10 @@ 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, extract_webapp_passport
+from libs.token import extract_access_token
 from models import Account, Tenant, TenantAccountJoin
 from models.model import AppMCPServer, EndUser
 from services.account_service import AccountService
@@ -62,30 +61,14 @@ 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":
-        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.")
+        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
     elif request.blueprint == "mcp":
         server_code = request.view_args.get("server_code") if request.view_args else None
         if not server_code:

api/libs/external_api.py

@@ -9,8 +9,9 @@ from werkzeug.exceptions import HTTPException
 from werkzeug.http import HTTP_STATUS_CODES
 
 from configs import dify_config
+from constants import COOKIE_NAME_ACCESS_TOKEN, COOKIE_NAME_CSRF_TOKEN, COOKIE_NAME_REFRESH_TOKEN
 from core.errors.error import AppInvokeQuotaExceededError
-from libs.token import build_force_logout_cookie_headers
+from libs.token import is_secure
 
 
 def http_status_message(code):
@@ -72,7 +73,15 @@ def register_external_error_handlers(api: Api):
                 error_code = getattr(e, "error_code", None)
                 if error_code == "unauthorized_and_force_logout":
                     # Add Set-Cookie headers to clear auth cookies
-                    headers["Set-Cookie"] = build_force_logout_cookie_headers()
+
+                    secure = is_secure()
+                    # response is not accessible, so we need to do it ugly
+                    common_part = "Path=/; Expires=Thu, 01 Jan 1970 00:00:00 GMT; HttpOnly"
+                    headers["Set-Cookie"] = [
+                        f'{COOKIE_NAME_ACCESS_TOKEN}=""; {common_part}{"; Secure" if secure else ""}; SameSite=Lax',
+                        f'{COOKIE_NAME_CSRF_TOKEN}=""; {common_part}{"; Secure" if secure else ""}; SameSite=Lax',
+                        f'{COOKIE_NAME_REFRESH_TOKEN}=""; {common_part}{"; Secure" if secure else ""}; SameSite=Lax',
+                    ]
             return data, status_code, headers
 
     _ = handle_http_exception

api/libs/token.py

@@ -30,28 +30,17 @@ def is_secure() -> bool:
     return dify_config.CONSOLE_WEB_URL.startswith("https") and dify_config.CONSOLE_API_URL.startswith("https")
 
 
-def _cookie_domain() -> str | None:
-    """
-    Returns the normalized cookie domain.
-
-    Leading dots are stripped from the configured domain. Historically, a leading dot
-    indicated that a cookie should be sent to all subdomains, but modern browsers treat
-    'example.com' and '.example.com' identically. This normalization ensures consistent
-    behavior and avoids confusion.
-    """
-    domain = dify_config.COOKIE_DOMAIN.strip()
-    domain = domain.removeprefix(".")
-    return domain or None
-
-
 def _real_cookie_name(cookie_name: str) -> str:
-    if is_secure() and _cookie_domain() is None:
+    if is_secure():
         return "__Host-" + cookie_name
     else:
         return cookie_name
 
 
 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:
@@ -66,19 +55,27 @@ 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))
 
@@ -86,10 +83,20 @@ def extract_access_token(request: Request) -> str | None:
 
 
 def extract_webapp_access_token(request: Request) -> str | None:
+    """
+    Try to extract webapp access token from cookie, then header.
+    """
+
     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))
 
@@ -105,7 +112,6 @@ def set_access_token_to_cookie(request: Request, response: Response, token: str,
         _real_cookie_name(COOKIE_NAME_ACCESS_TOKEN),
         value=token,
         httponly=True,
-        domain=_cookie_domain(),
         secure=is_secure(),
         samesite=samesite,
         max_age=int(dify_config.ACCESS_TOKEN_EXPIRE_MINUTES * 60),
@@ -118,7 +124,6 @@ def set_refresh_token_to_cookie(request: Request, response: Response, token: str
         _real_cookie_name(COOKIE_NAME_REFRESH_TOKEN),
         value=token,
         httponly=True,
-        domain=_cookie_domain(),
         secure=is_secure(),
         samesite="Lax",
         max_age=int(60 * 60 * 24 * dify_config.REFRESH_TOKEN_EXPIRE_DAYS),
@@ -131,7 +136,6 @@ def set_csrf_token_to_cookie(request: Request, response: Response, token: str):
         _real_cookie_name(COOKIE_NAME_CSRF_TOKEN),
         value=token,
         httponly=False,
-        domain=_cookie_domain(),
         secure=is_secure(),
         samesite="Lax",
         max_age=int(60 * dify_config.ACCESS_TOKEN_EXPIRE_MINUTES),
@@ -150,7 +154,6 @@ def _clear_cookie(
         "",
         expires=0,
         path="/",
-        domain=_cookie_domain(),
         secure=is_secure(),
         httponly=http_only,
         samesite=samesite,
@@ -173,19 +176,6 @@ def clear_csrf_token_from_cookie(response: Response):
     _clear_cookie(response, COOKIE_NAME_CSRF_TOKEN, http_only=False)
 
 
-def build_force_logout_cookie_headers() -> list[str]:
-    """
-    Generate Set-Cookie header values that clear all auth-related cookies.
-    This mirrors the behavior of the standard cookie clearing helpers while
-    allowing callers that do not have a Response instance to reuse the logic.
-    """
-    response = Response()
-    clear_access_token_from_cookie(response)
-    clear_csrf_token_from_cookie(response)
-    clear_refresh_token_from_cookie(response)
-    return response.headers.getlist("Set-Cookie")
-
-
 def check_csrf_token(request: Request, user_id: str):
     # some apis are sent by beacon, so we need to bypass csrf token check
     # since these APIs are post, they are already protected by SameSite: Lax, so csrf is not required.

api/models/tools.py

@@ -1,13 +1,16 @@
 import json
+from collections.abc import Mapping
 from datetime import datetime
 from decimal import Decimal
 from typing import TYPE_CHECKING, Any, cast
+from urllib.parse import urlparse
 
 import sqlalchemy as sa
 from deprecated import deprecated
 from sqlalchemy import ForeignKey, String, func
 from sqlalchemy.orm import Mapped, mapped_column
 
+from core.helper import encrypter
 from core.tools.entities.common_entities import I18nObject
 from core.tools.entities.tool_bundle import ApiToolBundle
 from core.tools.entities.tool_entities import ApiProviderSchemaType, WorkflowToolParameterConfiguration
@@ -18,7 +21,7 @@ from .model import Account, App, Tenant
 from .types import StringUUID
 
 if TYPE_CHECKING:
-    from core.entities.mcp_provider import MCPProviderEntity
+    from core.mcp.types import Tool as MCPTool
     from core.tools.entities.common_entities import I18nObject
     from core.tools.entities.tool_bundle import ApiToolBundle
     from core.tools.entities.tool_entities import ApiProviderSchemaType, WorkflowToolParameterConfiguration
@@ -328,36 +331,126 @@ class MCPToolProvider(TypeBase):
     def load_user(self) -> Account | None:
         return db.session.query(Account).where(Account.id == self.user_id).first()
 
+    @property
+    def tenant(self) -> Tenant | None:
+        return db.session.query(Tenant).where(Tenant.id == self.tenant_id).first()
+
     @property
     def credentials(self) -> dict[str, Any]:
         if not self.encrypted_credentials:
             return {}
         try:
-            return json.loads(self.encrypted_credentials)
+            return cast(dict[str, Any], json.loads(self.encrypted_credentials)) or {}
+        except json.JSONDecodeError:
+            return {}
+
+    @property
+    def mcp_tools(self) -> list["MCPTool"]:
+        from core.mcp.types import Tool as MCPTool
+
+        return [MCPTool.model_validate(tool) for tool in json.loads(self.tools)]
+
+    @property
+    def provider_icon(self) -> Mapping[str, str] | str:
+        from core.file import helpers as file_helpers
+
+        assert self.icon
+        try:
+            return json.loads(self.icon)
+        except json.JSONDecodeError:
+            return file_helpers.get_signed_file_url(self.icon)
+
+    @property
+    def decrypted_server_url(self) -> str:
+        return encrypter.decrypt_token(self.tenant_id, self.server_url)
+
+    @property
+    def decrypted_headers(self) -> dict[str, Any]:
+        """Get decrypted headers for MCP server requests."""
+        from core.entities.provider_entities import BasicProviderConfig
+        from core.helper.provider_cache import NoOpProviderCredentialCache
+        from core.tools.utils.encryption import create_provider_encrypter
+
+        try:
+            if not self.encrypted_headers:
+                return {}
+
+            headers_data = json.loads(self.encrypted_headers)
+
+            # Create dynamic config for all headers as SECRET_INPUT
+            config = [BasicProviderConfig(type=BasicProviderConfig.Type.SECRET_INPUT, name=key) for key in headers_data]
+
+            encrypter_instance, _ = create_provider_encrypter(
+                tenant_id=self.tenant_id,
+                config=config,
+                cache=NoOpProviderCredentialCache(),
+            )
+
+            result = encrypter_instance.decrypt(headers_data)
+            return result
         except Exception:
             return {}
 
     @property
-    def headers(self) -> dict[str, Any]:
-        if self.encrypted_headers is None:
-            return {}
+    def masked_headers(self) -> dict[str, Any]:
+        """Get masked headers for frontend display."""
+        from core.entities.provider_entities import BasicProviderConfig
+        from core.helper.provider_cache import NoOpProviderCredentialCache
+        from core.tools.utils.encryption import create_provider_encrypter
+
         try:
-            return json.loads(self.encrypted_headers)
+            if not self.encrypted_headers:
+                return {}
+
+            headers_data = json.loads(self.encrypted_headers)
+
+            # Create dynamic config for all headers as SECRET_INPUT
+            config = [BasicProviderConfig(type=BasicProviderConfig.Type.SECRET_INPUT, name=key) for key in headers_data]
+
+            encrypter_instance, _ = create_provider_encrypter(
+                tenant_id=self.tenant_id,
+                config=config,
+                cache=NoOpProviderCredentialCache(),
+            )
+
+            # First decrypt, then mask
+            decrypted_headers = encrypter_instance.decrypt(headers_data)
+            result = encrypter_instance.mask_tool_credentials(decrypted_headers)
+            return result
         except Exception:
             return {}
 
     @property
-    def tool_dict(self) -> list[dict[str, Any]]:
-        try:
-            return json.loads(self.tools) if self.tools else []
-        except (json.JSONDecodeError, TypeError):
-            return []
+    def masked_server_url(self) -> str:
+        def mask_url(url: str, mask_char: str = "*") -> str:
+            """
+            mask the url to a simple string
+            """
+            parsed = urlparse(url)
+            base_url = f"{parsed.scheme}://{parsed.netloc}"
 
-    def to_entity(self) -> "MCPProviderEntity":
-        """Convert to domain entity"""
-        from core.entities.mcp_provider import MCPProviderEntity
+            if parsed.path and parsed.path != "/":
+                return f"{base_url}/{mask_char * 6}"
+            else:
+                return base_url
 
-        return MCPProviderEntity.from_db_model(self)
+        return mask_url(self.decrypted_server_url)
+
+    @property
+    def decrypted_credentials(self) -> dict[str, Any]:
+        from core.helper.provider_cache import NoOpProviderCredentialCache
+        from core.tools.mcp_tool.provider import MCPToolProviderController
+        from core.tools.utils.encryption import create_provider_encrypter
+
+        provider_controller = MCPToolProviderController.from_db(self)
+
+        encrypter, _ = create_provider_encrypter(
+            tenant_id=self.tenant_id,
+            config=[x.to_basic_provider_config() for x in provider_controller.get_credentials_schema()],
+            cache=NoOpProviderCredentialCache(),
+        )
+
+        return encrypter.decrypt(self.credentials)
 
 
 class ToolModelInvoke(TypeBase):

api/pyproject.toml

@@ -117,7 +117,7 @@ dev = [
     "pytest-cov~=4.1.0",
     "pytest-env~=1.1.3",
     "pytest-mock~=3.14.0",
-    "testcontainers~=4.13.2",
+    "testcontainers~=4.10.0",
     "types-aiofiles~=24.1.0",
     "types-beautifulsoup4~=4.12.0",
     "types-cachetools~=5.5.0",

api/repositories/api_workflow_run_repository.py

@@ -28,7 +28,7 @@ Example:
     runs = repo.get_paginated_workflow_runs(
         tenant_id="tenant-123",
         app_id="app-456",
-        triggered_from=WorkflowRunTriggeredFrom.DEBUGGING,
+        triggered_from="debugging",
         limit=20
     )
     ```
@@ -40,14 +40,7 @@ from typing import Protocol
 
 from core.workflow.repositories.workflow_execution_repository import WorkflowExecutionRepository
 from libs.infinite_scroll_pagination import InfiniteScrollPagination
-from models.enums import WorkflowRunTriggeredFrom
 from models.workflow import WorkflowRun
-from repositories.types import (
-    AverageInteractionStats,
-    DailyRunsStats,
-    DailyTerminalsStats,
-    DailyTokenCostStats,
-)
 
 
 class APIWorkflowRunRepository(WorkflowExecutionRepository, Protocol):
@@ -63,7 +56,7 @@ class APIWorkflowRunRepository(WorkflowExecutionRepository, Protocol):
         self,
         tenant_id: str,
         app_id: str,
-        triggered_from: WorkflowRunTriggeredFrom | Sequence[WorkflowRunTriggeredFrom],
+        triggered_from: str,
         limit: int = 20,
         last_id: str | None = None,
         status: str | None = None,
@@ -78,7 +71,7 @@ class APIWorkflowRunRepository(WorkflowExecutionRepository, Protocol):
         Args:
             tenant_id: Tenant identifier for multi-tenant isolation
             app_id: Application identifier
-            triggered_from: Filter by trigger source(s) (e.g., "debugging", "app-run", or list of values)
+            triggered_from: Filter by trigger source (e.g., "debugging", "app-run")
             limit: Maximum number of records to return (default: 20)
             last_id: Cursor for pagination - ID of the last record from previous page
             status: Optional filter by status (e.g., "running", "succeeded", "failed")
@@ -116,31 +109,6 @@ class APIWorkflowRunRepository(WorkflowExecutionRepository, Protocol):
         """
         ...
 
-    def get_workflow_run_by_id_without_tenant(
-        self,
-        run_id: str,
-    ) -> WorkflowRun | None:
-        """
-        Get a specific workflow run by ID without tenant/app context.
-
-        Retrieves a single workflow run using only the run ID, without
-        requiring tenant_id or app_id. This method is intended for internal
-        system operations like tracing and monitoring where the tenant context
-        is not available upfront.
-
-        Args:
-            run_id: Workflow run identifier
-
-        Returns:
-            WorkflowRun object if found, None otherwise
-
-        Note:
-            This method bypasses tenant isolation checks and should only be used
-            in trusted system contexts like ops trace collection. For user-facing
-            operations, use get_workflow_run_by_id() with proper tenant isolation.
-        """
-        ...
-
     def get_workflow_runs_count(
         self,
         tenant_id: str,
@@ -250,119 +218,3 @@ class APIWorkflowRunRepository(WorkflowExecutionRepository, Protocol):
             and ensure proper data retention policies are followed.
         """
         ...
-
-    def get_daily_runs_statistics(
-        self,
-        tenant_id: str,
-        app_id: str,
-        triggered_from: str,
-        start_date: datetime | None = None,
-        end_date: datetime | None = None,
-        timezone: str = "UTC",
-    ) -> list[DailyRunsStats]:
-        """
-        Get daily runs statistics.
-
-        Retrieves daily workflow runs count grouped by date for a specific app
-        and trigger source. Used for workflow statistics dashboard.
-
-        Args:
-            tenant_id: Tenant identifier for multi-tenant isolation
-            app_id: Application identifier
-            triggered_from: Filter by trigger source (e.g., "app-run")
-            start_date: Optional start date filter
-            end_date: Optional end date filter
-            timezone: Timezone for date grouping (default: "UTC")
-
-        Returns:
-            List of dictionaries containing date and runs count:
-            [{"date": "2024-01-01", "runs": 10}, ...]
-        """
-        ...
-
-    def get_daily_terminals_statistics(
-        self,
-        tenant_id: str,
-        app_id: str,
-        triggered_from: str,
-        start_date: datetime | None = None,
-        end_date: datetime | None = None,
-        timezone: str = "UTC",
-    ) -> list[DailyTerminalsStats]:
-        """
-        Get daily terminals statistics.
-
-        Retrieves daily unique terminal count grouped by date for a specific app
-        and trigger source. Used for workflow statistics dashboard.
-
-        Args:
-            tenant_id: Tenant identifier for multi-tenant isolation
-            app_id: Application identifier
-            triggered_from: Filter by trigger source (e.g., "app-run")
-            start_date: Optional start date filter
-            end_date: Optional end date filter
-            timezone: Timezone for date grouping (default: "UTC")
-
-        Returns:
-            List of dictionaries containing date and terminal count:
-            [{"date": "2024-01-01", "terminal_count": 5}, ...]
-        """
-        ...
-
-    def get_daily_token_cost_statistics(
-        self,
-        tenant_id: str,
-        app_id: str,
-        triggered_from: str,
-        start_date: datetime | None = None,
-        end_date: datetime | None = None,
-        timezone: str = "UTC",
-    ) -> list[DailyTokenCostStats]:
-        """
-        Get daily token cost statistics.
-
-        Retrieves daily total token count grouped by date for a specific app
-        and trigger source. Used for workflow statistics dashboard.
-
-        Args:
-            tenant_id: Tenant identifier for multi-tenant isolation
-            app_id: Application identifier
-            triggered_from: Filter by trigger source (e.g., "app-run")
-            start_date: Optional start date filter
-            end_date: Optional end date filter
-            timezone: Timezone for date grouping (default: "UTC")
-
-        Returns:
-            List of dictionaries containing date and token count:
-            [{"date": "2024-01-01", "token_count": 1000}, ...]
-        """
-        ...
-
-    def get_average_app_interaction_statistics(
-        self,
-        tenant_id: str,
-        app_id: str,
-        triggered_from: str,
-        start_date: datetime | None = None,
-        end_date: datetime | None = None,
-        timezone: str = "UTC",
-    ) -> list[AverageInteractionStats]:
-        """
-        Get average app interaction statistics.
-
-        Retrieves daily average interactions per user grouped by date for a specific app
-        and trigger source. Used for workflow statistics dashboard.
-
-        Args:
-            tenant_id: Tenant identifier for multi-tenant isolation
-            app_id: Application identifier
-            triggered_from: Filter by trigger source (e.g., "app-run")
-            start_date: Optional start date filter
-            end_date: Optional end date filter
-            timezone: Timezone for date grouping (default: "UTC")
-
-        Returns:
-            List of dictionaries containing date and average interactions:
-            [{"date": "2024-01-01", "interactions": 2.5}, ...]
-        """
-        ...

api/repositories/sqlalchemy_api_workflow_run_repository.py

@@ -22,25 +22,16 @@ Implementation Notes:
 import logging
 from collections.abc import Sequence
 from datetime import datetime
-from decimal import Decimal
-from typing import Any, cast
+from typing import cast
 
-import sqlalchemy as sa
 from sqlalchemy import delete, func, select
 from sqlalchemy.engine import CursorResult
 from sqlalchemy.orm import Session, sessionmaker
 
 from libs.infinite_scroll_pagination import InfiniteScrollPagination
 from libs.time_parser import get_time_threshold
-from models.enums import WorkflowRunTriggeredFrom
 from models.workflow import WorkflowRun
 from repositories.api_workflow_run_repository import APIWorkflowRunRepository
-from repositories.types import (
-    AverageInteractionStats,
-    DailyRunsStats,
-    DailyTerminalsStats,
-    DailyTokenCostStats,
-)
 
 logger = logging.getLogger(__name__)
 
@@ -70,7 +61,7 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository):
         self,
         tenant_id: str,
         app_id: str,
-        triggered_from: WorkflowRunTriggeredFrom | Sequence[WorkflowRunTriggeredFrom],
+        triggered_from: str,
         limit: int = 20,
         last_id: str | None = None,
         status: str | None = None,
@@ -87,14 +78,9 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository):
             base_stmt = select(WorkflowRun).where(
                 WorkflowRun.tenant_id == tenant_id,
                 WorkflowRun.app_id == app_id,
+                WorkflowRun.triggered_from == triggered_from,
             )
 
-            # Handle triggered_from values
-            if isinstance(triggered_from, WorkflowRunTriggeredFrom):
-                triggered_from = [triggered_from]
-            if triggered_from:
-                base_stmt = base_stmt.where(WorkflowRun.triggered_from.in_(triggered_from))
-
             # Add optional status filter
             if status:
                 base_stmt = base_stmt.where(WorkflowRun.status == status)
@@ -140,17 +126,6 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository):
             )
             return session.scalar(stmt)
 
-    def get_workflow_run_by_id_without_tenant(
-        self,
-        run_id: str,
-    ) -> WorkflowRun | None:
-        """
-        Get a specific workflow run by ID without tenant/app context.
-        """
-        with self._session_maker() as session:
-            stmt = select(WorkflowRun).where(WorkflowRun.id == run_id)
-            return session.scalar(stmt)
-
     def get_workflow_runs_count(
         self,
         tenant_id: str,
@@ -300,213 +275,3 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository):
 
         logger.info("Total deleted %s workflow runs for app %s", total_deleted, app_id)
         return total_deleted
-
-    def get_daily_runs_statistics(
-        self,
-        tenant_id: str,
-        app_id: str,
-        triggered_from: str,
-        start_date: datetime | None = None,
-        end_date: datetime | None = None,
-        timezone: str = "UTC",
-    ) -> list[DailyRunsStats]:
-        """
-        Get daily runs statistics using raw SQL for optimal performance.
-        """
-        sql_query = """SELECT
-    DATE(DATE_TRUNC('day', created_at AT TIME ZONE 'UTC' AT TIME ZONE :tz )) AS date,
-    COUNT(id) AS runs
-FROM
-    workflow_runs
-WHERE
-    tenant_id = :tenant_id
-    AND app_id = :app_id
-    AND triggered_from = :triggered_from"""
-
-        arg_dict: dict[str, Any] = {
-            "tz": timezone,
-            "tenant_id": tenant_id,
-            "app_id": app_id,
-            "triggered_from": triggered_from,
-        }
-
-        if start_date:
-            sql_query += " AND created_at >= :start_date"
-            arg_dict["start_date"] = start_date
-
-        if end_date:
-            sql_query += " AND created_at < :end_date"
-            arg_dict["end_date"] = end_date
-
-        sql_query += " GROUP BY date ORDER BY date"
-
-        response_data = []
-        with self._session_maker() as session:
-            rs = session.execute(sa.text(sql_query), arg_dict)
-            for row in rs:
-                response_data.append({"date": str(row.date), "runs": row.runs})
-
-        return cast(list[DailyRunsStats], response_data)
-
-    def get_daily_terminals_statistics(
-        self,
-        tenant_id: str,
-        app_id: str,
-        triggered_from: str,
-        start_date: datetime | None = None,
-        end_date: datetime | None = None,
-        timezone: str = "UTC",
-    ) -> list[DailyTerminalsStats]:
-        """
-        Get daily terminals statistics using raw SQL for optimal performance.
-        """
-        sql_query = """SELECT
-    DATE(DATE_TRUNC('day', created_at AT TIME ZONE 'UTC' AT TIME ZONE :tz )) AS date,
-    COUNT(DISTINCT created_by) AS terminal_count
-FROM
-    workflow_runs
-WHERE
-    tenant_id = :tenant_id
-    AND app_id = :app_id
-    AND triggered_from = :triggered_from"""
-
-        arg_dict: dict[str, Any] = {
-            "tz": timezone,
-            "tenant_id": tenant_id,
-            "app_id": app_id,
-            "triggered_from": triggered_from,
-        }
-
-        if start_date:
-            sql_query += " AND created_at >= :start_date"
-            arg_dict["start_date"] = start_date
-
-        if end_date:
-            sql_query += " AND created_at < :end_date"
-            arg_dict["end_date"] = end_date
-
-        sql_query += " GROUP BY date ORDER BY date"
-
-        response_data = []
-        with self._session_maker() as session:
-            rs = session.execute(sa.text(sql_query), arg_dict)
-            for row in rs:
-                response_data.append({"date": str(row.date), "terminal_count": row.terminal_count})
-
-        return cast(list[DailyTerminalsStats], response_data)
-
-    def get_daily_token_cost_statistics(
-        self,
-        tenant_id: str,
-        app_id: str,
-        triggered_from: str,
-        start_date: datetime | None = None,
-        end_date: datetime | None = None,
-        timezone: str = "UTC",
-    ) -> list[DailyTokenCostStats]:
-        """
-        Get daily token cost statistics using raw SQL for optimal performance.
-        """
-        sql_query = """SELECT
-    DATE(DATE_TRUNC('day', created_at AT TIME ZONE 'UTC' AT TIME ZONE :tz )) AS date,
-    SUM(total_tokens) AS token_count
-FROM
-    workflow_runs
-WHERE
-    tenant_id = :tenant_id
-    AND app_id = :app_id
-    AND triggered_from = :triggered_from"""
-
-        arg_dict: dict[str, Any] = {
-            "tz": timezone,
-            "tenant_id": tenant_id,
-            "app_id": app_id,
-            "triggered_from": triggered_from,
-        }
-
-        if start_date:
-            sql_query += " AND created_at >= :start_date"
-            arg_dict["start_date"] = start_date
-
-        if end_date:
-            sql_query += " AND created_at < :end_date"
-            arg_dict["end_date"] = end_date
-
-        sql_query += " GROUP BY date ORDER BY date"
-
-        response_data = []
-        with self._session_maker() as session:
-            rs = session.execute(sa.text(sql_query), arg_dict)
-            for row in rs:
-                response_data.append(
-                    {
-                        "date": str(row.date),
-                        "token_count": row.token_count,
-                    }
-                )
-
-        return cast(list[DailyTokenCostStats], response_data)
-
-    def get_average_app_interaction_statistics(
-        self,
-        tenant_id: str,
-        app_id: str,
-        triggered_from: str,
-        start_date: datetime | None = None,
-        end_date: datetime | None = None,
-        timezone: str = "UTC",
-    ) -> list[AverageInteractionStats]:
-        """
-        Get average app interaction statistics using raw SQL for optimal performance.
-        """
-        sql_query = """SELECT
-    AVG(sub.interactions) AS interactions,
-    sub.date
-FROM
-    (
-        SELECT
-            DATE(DATE_TRUNC('day', c.created_at AT TIME ZONE 'UTC' AT TIME ZONE :tz )) AS date,
-            c.created_by,
-            COUNT(c.id) AS interactions
-        FROM
-            workflow_runs c
-        WHERE
-            c.tenant_id = :tenant_id
-            AND c.app_id = :app_id
-            AND c.triggered_from = :triggered_from
-            {{start}}
-            {{end}}
-        GROUP BY
-            date, c.created_by
-    ) sub
-GROUP BY
-    sub.date"""
-
-        arg_dict: dict[str, Any] = {
-            "tz": timezone,
-            "tenant_id": tenant_id,
-            "app_id": app_id,
-            "triggered_from": triggered_from,
-        }
-
-        if start_date:
-            sql_query = sql_query.replace("{{start}}", " AND c.created_at >= :start_date")
-            arg_dict["start_date"] = start_date
-        else:
-            sql_query = sql_query.replace("{{start}}", "")
-
-        if end_date:
-            sql_query = sql_query.replace("{{end}}", " AND c.created_at < :end_date")
-            arg_dict["end_date"] = end_date
-        else:
-            sql_query = sql_query.replace("{{end}}", "")
-
-        response_data = []
-        with self._session_maker() as session:
-            rs = session.execute(sa.text(sql_query), arg_dict)
-            for row in rs:
-                response_data.append(
-                    {"date": str(row.date), "interactions": float(row.interactions.quantize(Decimal("0.01")))}
-                )
-
-        return cast(list[AverageInteractionStats], response_data)

api/repositories/types.py

@@ -1,21 +0,0 @@
-from typing import TypedDict
-
-
-class DailyRunsStats(TypedDict):
-    date: str
-    runs: int
-
-
-class DailyTerminalsStats(TypedDict):
-    date: str
-    terminal_count: int
-
-
-class DailyTokenCostStats(TypedDict):
-    date: str
-    token_count: int
-
-
-class AverageInteractionStats(TypedDict):
-    date: str
-    interactions: float

api/services/audio_service.py

@@ -82,51 +82,54 @@ 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):
-            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"):
+            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")
+                        ):
                             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.")
+                        voice = workflow.features_dict["text_to_speech"].get("voice")
                     else:
-                        raise ValueError("Sorry, no voice available.")
+                        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
 
-                return model_instance.invoke_tts(
-                    content_text=text_content.strip(), user=end_user, tenant_id=app_model.tenant_id, voice=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
                 )
-            except Exception as e:
-                raise e
+                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
 
         if message_id:
             try:

api/services/dataset_service.py

@@ -49,6 +49,7 @@ from models.model import UploadFile
 from models.provider_ids import ModelProviderID
 from models.source import DataSourceOauthBinding
 from models.workflow import Workflow
+from services.document_indexing_task_proxy import DocumentIndexingTaskProxy
 from services.entities.knowledge_entities.knowledge_entities import (
     ChildChunkUpdateArgs,
     KnowledgeConfig,
@@ -78,7 +79,6 @@ from tasks.deal_dataset_vector_index_task import deal_dataset_vector_index_task
 from tasks.delete_segment_from_index_task import delete_segment_from_index_task
 from tasks.disable_segment_from_index_task import disable_segment_from_index_task
 from tasks.disable_segments_from_index_task import disable_segments_from_index_task
-from tasks.document_indexing_task import document_indexing_task
 from tasks.document_indexing_update_task import document_indexing_update_task
 from tasks.duplicate_document_indexing_task import duplicate_document_indexing_task
 from tasks.enable_segments_to_index_task import enable_segments_to_index_task
@@ -1417,7 +1417,7 @@ class DocumentService:
         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
+        assert knowledge_config.data_source.info_list.file_info_list
 
         features = FeatureService.get_features(current_user.current_tenant_id)
 
@@ -1426,8 +1426,6 @@ class DocumentService:
                 count = 0
                 if knowledge_config.data_source:
                     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
                         count = len(upload_file_list)
                     elif knowledge_config.data_source.info_list.data_source_type == "notion_import":
@@ -1533,8 +1531,6 @@ class DocumentService:
                 document_ids = []
                 duplicate_document_ids = []
                 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 = (
@@ -1691,7 +1687,7 @@ class DocumentService:
 
                 # trigger async task
                 if document_ids:
-                    document_indexing_task.delay(dataset.id, document_ids)
+                    DocumentIndexingTaskProxy(dataset.tenant_id, dataset.id, document_ids).delay()
                 if duplicate_document_ids:
                     duplicate_document_indexing_task.delay(dataset.id, duplicate_document_ids)
 

api/services/document_indexing_task_proxy.py

@@ -0,0 +1,82 @@
+import logging
+from collections.abc import Callable
+from dataclasses import asdict
+from functools import cached_property
+
+from core.entities.document_task import DocumentTask
+from core.rag.pipeline.queue import TenantSelfTaskQueue
+from services.feature_service import FeatureService
+from tasks.document_indexing_task import normal_document_indexing_task, priority_document_indexing_task
+
+logger = logging.getLogger(__name__)
+
+
+class DocumentIndexingTaskProxy:
+    def __init__(self, tenant_id: str, dataset_id: str, document_ids: list[str]):
+        self.tenant_id = tenant_id
+        self.dataset_id = dataset_id
+        self.document_ids = document_ids
+        self.tenant_self_task_queue = TenantSelfTaskQueue(tenant_id, "document_indexing")
+
+    @cached_property
+    def features(self):
+        return FeatureService.get_features(self.tenant_id)
+
+    def _send_to_direct_queue(self, task_func: Callable):
+        logger.info("send dataset %s to direct queue", self.dataset_id)
+        task_func.delay(  # type: ignore
+            tenant_id=self.tenant_id,
+            dataset_id=self.dataset_id,
+            document_ids=self.document_ids
+        )
+
+    def _send_to_tenant_queue(self, task_func: Callable):
+        logger.info("send dataset %s to tenant queue", self.dataset_id)
+        if self.tenant_self_task_queue.get_task_key():
+            # Add to waiting queue using List operations (lpush)
+            self.tenant_self_task_queue.push_tasks([
+                asdict(
+                    DocumentTask(tenant_id=self.tenant_id, dataset_id=self.dataset_id, document_ids=self.document_ids)
+                )
+            ])
+            logger.info("push tasks: %s - %s", self.dataset_id, self.document_ids)
+        else:
+            # Set flag and execute task
+            self.tenant_self_task_queue.set_task_waiting_time()
+            task_func.delay(  # type: ignore
+                tenant_id=self.tenant_id,
+                dataset_id=self.dataset_id,
+                document_ids=self.document_ids
+            )
+            logger.info("init tasks: %s - %s", self.dataset_id, self.document_ids)
+
+    def _send_to_default_tenant_queue(self):
+        self._send_to_tenant_queue(normal_document_indexing_task)
+    
+    def _send_to_priority_tenant_queue(self):
+        self._send_to_tenant_queue(priority_document_indexing_task)
+
+    def _send_to_priority_direct_queue(self):
+        self._send_to_direct_queue(priority_document_indexing_task)
+
+    def _dispatch(self):
+        logger.info(
+            "dispatch args: %s - %s - %s",
+            self.tenant_id,
+            self.features.billing.enabled,
+            self.features.billing.subscription.plan
+        )
+        # dispatch to different indexing queue with tenant isolation when billing enabled
+        if self.features.billing.enabled:
+            if self.features.billing.subscription.plan == "sandbox":
+                # dispatch to normal pipeline queue with tenant self sub queue for sandbox plan
+                self._send_to_default_tenant_queue()
+            else:
+                # dispatch to priority pipeline queue with tenant self sub queue for other plans
+                self._send_to_priority_tenant_queue()
+        else:
+            # dispatch to priority queue without tenant isolation for others, e.g.: self-hosted or enterprise
+            self._send_to_priority_direct_queue()
+
+    def delay(self):
+        self._dispatch()

api/services/rag_pipeline/rag_pipeline.py

@@ -9,7 +9,7 @@ from typing import Any, Union, cast
 from uuid import uuid4
 
 from flask_login import current_user
-from sqlalchemy import func, select
+from sqlalchemy import func, or_, select
 from sqlalchemy.orm import Session, sessionmaker
 
 import contexts
@@ -94,7 +94,6 @@ class RagPipelineService:
         self._node_execution_service_repo = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(
             session_maker
         )
-        self._workflow_run_repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker)
 
     @classmethod
     def get_pipeline_templates(cls, type: str = "built-in", language: str = "en-US") -> dict:
@@ -1016,21 +1015,48 @@ class RagPipelineService:
         :param args: request args
         """
         limit = int(args.get("limit", 20))
-        last_id = args.get("last_id")
 
-        triggered_from_values = [
-            WorkflowRunTriggeredFrom.RAG_PIPELINE_RUN,
-            WorkflowRunTriggeredFrom.RAG_PIPELINE_DEBUGGING,
-        ]
-
-        return self._workflow_run_repo.get_paginated_workflow_runs(
-            tenant_id=pipeline.tenant_id,
-            app_id=pipeline.id,
-            triggered_from=triggered_from_values,
-            limit=limit,
-            last_id=last_id,
+        base_query = db.session.query(WorkflowRun).where(
+            WorkflowRun.tenant_id == pipeline.tenant_id,
+            WorkflowRun.app_id == pipeline.id,
+            or_(
+                WorkflowRun.triggered_from == WorkflowRunTriggeredFrom.RAG_PIPELINE_RUN.value,
+                WorkflowRun.triggered_from == WorkflowRunTriggeredFrom.RAG_PIPELINE_DEBUGGING.value,
+            ),
         )
 
+        if args.get("last_id"):
+            last_workflow_run = base_query.where(
+                WorkflowRun.id == args.get("last_id"),
+            ).first()
+
+            if not last_workflow_run:
+                raise ValueError("Last workflow run not exists")
+
+            workflow_runs = (
+                base_query.where(
+                    WorkflowRun.created_at < last_workflow_run.created_at, WorkflowRun.id != last_workflow_run.id
+                )
+                .order_by(WorkflowRun.created_at.desc())
+                .limit(limit)
+                .all()
+            )
+        else:
+            workflow_runs = base_query.order_by(WorkflowRun.created_at.desc()).limit(limit).all()
+
+        has_more = False
+        if len(workflow_runs) == limit:
+            current_page_first_workflow_run = workflow_runs[-1]
+            rest_count = base_query.where(
+                WorkflowRun.created_at < current_page_first_workflow_run.created_at,
+                WorkflowRun.id != current_page_first_workflow_run.id,
+            ).count()
+
+            if rest_count > 0:
+                has_more = True
+
+        return InfiniteScrollPagination(data=workflow_runs, limit=limit, has_more=has_more)
+
     def get_rag_pipeline_workflow_run(self, pipeline: Pipeline, run_id: str) -> WorkflowRun | None:
         """
         Get workflow run detail
@@ -1038,12 +1064,18 @@ class RagPipelineService:
         :param app_model: app model
         :param run_id: workflow run id
         """
-        return self._workflow_run_repo.get_workflow_run_by_id(
-            tenant_id=pipeline.tenant_id,
-            app_id=pipeline.id,
-            run_id=run_id,
+        workflow_run = (
+            db.session.query(WorkflowRun)
+            .where(
+                WorkflowRun.tenant_id == pipeline.tenant_id,
+                WorkflowRun.app_id == pipeline.id,
+                WorkflowRun.id == run_id,
+            )
+            .first()
         )
 
+        return workflow_run
+
     def get_rag_pipeline_workflow_run_node_executions(
         self,
         pipeline: Pipeline,

api/services/rag_pipeline/rag_pipeline_task_proxy.py

@@ -0,0 +1,101 @@
+import json
+import logging
+from collections.abc import Callable
+from functools import cached_property
+
+from core.app.entities.rag_pipeline_invoke_entities import RagPipelineInvokeEntity
+from core.rag.pipeline.queue import TenantSelfTaskQueue
+from extensions.ext_database import db
+from services.feature_service import FeatureService
+from services.file_service import FileService
+from tasks.rag_pipeline.priority_rag_pipeline_run_task import priority_rag_pipeline_run_task
+from tasks.rag_pipeline.rag_pipeline_run_task import rag_pipeline_run_task
+
+logger = logging.getLogger(__name__)
+
+
+class RagPipelineTaskProxy:
+    def __init__(
+        self, dataset_tenant_id: str, user_id: str, rag_pipeline_invoke_entities: list[RagPipelineInvokeEntity]
+    ):
+        self.dataset_tenant_id = dataset_tenant_id
+        self.user_id = user_id
+        self.rag_pipeline_invoke_entities = rag_pipeline_invoke_entities
+        self.tenant_self_pipeline_task_queue = TenantSelfTaskQueue(dataset_tenant_id, "pipeline")
+
+    @cached_property
+    def features(self):
+        return FeatureService.get_features(self.dataset_tenant_id)
+    
+    def _upload_invoke_entities(self) -> str:
+        text = [item.model_dump() for item in self.rag_pipeline_invoke_entities]
+        name = "rag_pipeline_invoke_entities.json"
+        # Convert list to proper JSON string
+        json_text = json.dumps(text)
+        upload_file = FileService(db.engine).upload_text(json_text, name, self.user_id, self.dataset_tenant_id)
+        return upload_file.id
+    
+    def _send_to_direct_queue(self, upload_file_id: str, task_func: Callable):
+        logger.info("send file %s to direct queue", upload_file_id)
+        task_func.delay(  # type: ignore
+            rag_pipeline_invoke_entities_file_id=upload_file_id,
+            tenant_id=self.dataset_tenant_id,
+        )
+
+    def _send_to_tenant_queue(self, upload_file_id: str, task_func: Callable):
+        logger.info("send file %s to tenant queue", upload_file_id)
+        if self.tenant_self_pipeline_task_queue.get_task_key():
+            # Add to waiting queue using List operations (lpush)
+            self.tenant_self_pipeline_task_queue.push_tasks([upload_file_id])
+            logger.info("push tasks: %s", upload_file_id)
+        else:
+            # Set flag and execute task
+            self.tenant_self_pipeline_task_queue.set_task_waiting_time()
+            task_func.delay(  # type: ignore
+                rag_pipeline_invoke_entities_file_id=upload_file_id,
+                tenant_id=self.dataset_tenant_id,
+            )
+            logger.info("init tasks: %s", upload_file_id)
+
+    def _send_to_default_tenant_queue(self, upload_file_id: str):
+        self._send_to_tenant_queue(upload_file_id, rag_pipeline_run_task)
+
+    def _send_to_priority_tenant_queue(self, upload_file_id: str):
+        self._send_to_tenant_queue(upload_file_id, priority_rag_pipeline_run_task)
+
+    def _send_to_priority_direct_queue(self, upload_file_id: str):
+        self._send_to_direct_queue(upload_file_id, priority_rag_pipeline_run_task)
+
+    def _dispatch(self):
+        upload_file_id = self._upload_invoke_entities()
+        if not upload_file_id:
+            raise ValueError("upload_file_id is empty")
+
+        logger.info(
+            "dispatch args: %s - %s - %s",
+            self.dataset_tenant_id,
+            self.features.billing.enabled,
+            self.features.billing.subscription.plan
+        )
+
+        # dispatch to different pipeline queue with tenant isolation when billing enabled
+        if self.features.billing.enabled:
+            if self.features.billing.subscription.plan == "sandbox":
+                # dispatch to normal pipeline queue with tenant isolation for sandbox plan
+                self._send_to_default_tenant_queue(upload_file_id)
+            else:
+                # dispatch to priority pipeline queue with tenant isolation for other plans
+                self._send_to_priority_tenant_queue(upload_file_id)
+        else:
+            # dispatch to priority pipeline queue without tenant isolation for others, e.g.: self-hosted or enterprise
+            self._send_to_priority_direct_queue(upload_file_id)
+
+    def delay(self):
+        if not self.rag_pipeline_invoke_entities:
+            logger.warning(
+                "Received empty rag pipeline invoke entities, no tasks delivered: %s %s",
+                self.dataset_tenant_id,
+                self.user_id
+            )
+            return
+        self._dispatch()

api/services/tools/mcp_tools_manage_service.py

@@ -1,118 +1,86 @@
 import hashlib
 import json
-import logging
 from datetime import datetime
-from enum import StrEnum
 from typing import Any
-from urllib.parse import urlparse
 
-from pydantic import BaseModel, Field
-from sqlalchemy import or_, select
+from sqlalchemy import or_
 from sqlalchemy.exc import IntegrityError
-from sqlalchemy.orm import Session
 
-from core.entities.mcp_provider import MCPAuthentication, MCPConfiguration, MCPProviderEntity
 from core.helper import encrypter
 from core.helper.provider_cache import NoOpProviderCredentialCache
-from core.mcp.auth.auth_flow import auth
-from core.mcp.auth_client import MCPClientWithAuthRetry
 from core.mcp.error import MCPAuthError, MCPError
+from core.mcp.mcp_client import MCPClient
 from core.tools.entities.api_entities import ToolProviderApiEntity
+from core.tools.entities.common_entities import I18nObject
+from core.tools.entities.tool_entities import ToolProviderType
+from core.tools.mcp_tool.provider import MCPToolProviderController
 from core.tools.utils.encryption import ProviderConfigEncrypter
+from extensions.ext_database import db
 from models.tools import MCPToolProvider
 from services.tools.tools_transform_service import ToolTransformService
 
-logger = logging.getLogger(__name__)
-
-# Constants
 UNCHANGED_SERVER_URL_PLACEHOLDER = "[__HIDDEN__]"
-CLIENT_NAME = "Dify"
-EMPTY_TOOLS_JSON = "[]"
-EMPTY_CREDENTIALS_JSON = "{}"
-
-
-class OAuthDataType(StrEnum):
-    """Types of OAuth data that can be saved."""
-
-    TOKENS = "tokens"
-    CLIENT_INFO = "client_info"
-    CODE_VERIFIER = "code_verifier"
-    MIXED = "mixed"
-
-
-class ReconnectResult(BaseModel):
-    """Result of reconnecting to an MCP provider"""
-
-    authed: bool = Field(description="Whether the provider is authenticated")
-    tools: str = Field(description="JSON string of tool list")
-    encrypted_credentials: str = Field(description="JSON string of encrypted credentials")
-
-
-class ServerUrlValidationResult(BaseModel):
-    """Result of server URL validation check"""
-
-    needs_validation: bool
-    validation_passed: bool = False
-    reconnect_result: ReconnectResult | None = None
-    encrypted_server_url: str | None = None
-    server_url_hash: str | None = None
-
-    @property
-    def should_update_server_url(self) -> bool:
-        """Check if server URL should be updated based on validation result"""
-        return self.needs_validation and self.validation_passed and self.reconnect_result is not None
 
 
 class MCPToolManageService:
-    """Service class for managing MCP tools and providers."""
+    """
+    Service class for managing mcp tools.
+    """
 
-    def __init__(self, session: Session):
-        self._session = session
-
-    # ========== Provider CRUD Operations ==========
-
-    def get_provider(
-        self, *, provider_id: str | None = None, server_identifier: str | None = None, tenant_id: str
-    ) -> MCPToolProvider:
+    @staticmethod
+    def _encrypt_headers(headers: dict[str, str], tenant_id: str) -> dict[str, str]:
         """
-        Get MCP provider by ID or server identifier.
+        Encrypt headers using ProviderConfigEncrypter with all headers as SECRET_INPUT.
 
         Args:
-            provider_id: Provider ID (UUID)
-            server_identifier: Server identifier
-            tenant_id: Tenant ID
+            headers: Dictionary of headers to encrypt
+            tenant_id: Tenant ID for encryption
 
         Returns:
-            MCPToolProvider instance
-
-        Raises:
-            ValueError: If provider not found
+            Dictionary with all headers encrypted
         """
-        if server_identifier:
-            stmt = select(MCPToolProvider).where(
-                MCPToolProvider.tenant_id == tenant_id, MCPToolProvider.server_identifier == server_identifier
-            )
-        else:
-            stmt = select(MCPToolProvider).where(
-                MCPToolProvider.tenant_id == tenant_id, MCPToolProvider.id == provider_id
-            )
+        if not headers:
+            return {}
 
-        provider = self._session.scalar(stmt)
-        if not provider:
+        from core.entities.provider_entities import BasicProviderConfig
+        from core.helper.provider_cache import NoOpProviderCredentialCache
+        from core.tools.utils.encryption import create_provider_encrypter
+
+        # Create dynamic config for all headers as SECRET_INPUT
+        config = [BasicProviderConfig(type=BasicProviderConfig.Type.SECRET_INPUT, name=key) for key in headers]
+
+        encrypter_instance, _ = create_provider_encrypter(
+            tenant_id=tenant_id,
+            config=config,
+            cache=NoOpProviderCredentialCache(),
+        )
+
+        return encrypter_instance.encrypt(headers)
+
+    @staticmethod
+    def get_mcp_provider_by_provider_id(provider_id: str, tenant_id: str) -> MCPToolProvider:
+        res = (
+            db.session.query(MCPToolProvider)
+            .where(MCPToolProvider.tenant_id == tenant_id, MCPToolProvider.id == provider_id)
+            .first()
+        )
+        if not res:
             raise ValueError("MCP tool not found")
-        return provider
+        return res
 
-    def get_provider_entity(self, provider_id: str, tenant_id: str, by_server_id: bool = False) -> MCPProviderEntity:
-        """Get provider entity by ID or server identifier."""
-        if by_server_id:
-            db_provider = self.get_provider(server_identifier=provider_id, tenant_id=tenant_id)
-        else:
-            db_provider = self.get_provider(provider_id=provider_id, tenant_id=tenant_id)
-        return db_provider.to_entity()
+    @staticmethod
+    def get_mcp_provider_by_server_identifier(server_identifier: str, tenant_id: str) -> MCPToolProvider:
+        res = (
+            db.session.query(MCPToolProvider)
+            .where(MCPToolProvider.tenant_id == tenant_id, MCPToolProvider.server_identifier == server_identifier)
+            .first()
+        )
+        if not res:
+            raise ValueError("MCP tool not found")
+        return res
 
-    def create_provider(
-        self,
-        *,
+    @staticmethod
+    def create_mcp_provider(
         tenant_id: str,
         name: str,
         server_url: str,
@@ -121,30 +89,37 @@ class MCPToolManageService:
         icon_type: str,
         icon_background: str,
         server_identifier: str,
-        configuration: MCPConfiguration,
-        authentication: MCPAuthentication | None = None,
+        timeout: float,
+        sse_read_timeout: float,
         headers: dict[str, str] | None = None,
     ) -> ToolProviderApiEntity:
-        """Create a new MCP provider."""
-        # Validate URL format
-        if not self._is_valid_url(server_url):
-            raise ValueError("Server URL is not valid.")
-
         server_url_hash = hashlib.sha256(server_url.encode()).hexdigest()
-
-        # Check for existing provider
-        self._check_provider_exists(tenant_id, name, server_url_hash, server_identifier)
-
-        # Encrypt sensitive data
-        encrypted_server_url = encrypter.encrypt_token(tenant_id, server_url)
-        encrypted_headers = self._prepare_encrypted_dict(headers, tenant_id) if headers else None
-        encrypted_credentials = None
-        if authentication is not None and authentication.client_id:
-            encrypted_credentials = self._build_and_encrypt_credentials(
-                authentication.client_id, authentication.client_secret, tenant_id
+        existing_provider = (
+            db.session.query(MCPToolProvider)
+            .where(
+                MCPToolProvider.tenant_id == tenant_id,
+                or_(
+                    MCPToolProvider.name == name,
+                    MCPToolProvider.server_url_hash == server_url_hash,
+                    MCPToolProvider.server_identifier == server_identifier,
+                ),
             )
+            .first()
+        )
+        if existing_provider:
+            if existing_provider.name == name:
+                raise ValueError(f"MCP tool {name} already exists")
+            if existing_provider.server_url_hash == server_url_hash:
+                raise ValueError(f"MCP tool {server_url} already exists")
+            if existing_provider.server_identifier == server_identifier:
+                raise ValueError(f"MCP tool {server_identifier} already exists")
+        encrypted_server_url = encrypter.encrypt_token(tenant_id, server_url)
+        # Encrypt headers
+        encrypted_headers = None
+        if headers:
+            encrypted_headers_dict = MCPToolManageService._encrypt_headers(headers, tenant_id)
+            encrypted_headers = json.dumps(encrypted_headers_dict)
 
-        # Create provider
         mcp_tool = MCPToolProvider(
             tenant_id=tenant_id,
             name=name,
@@ -152,23 +127,93 @@ class MCPToolManageService:
             server_url_hash=server_url_hash,
             user_id=user_id,
             authed=False,
-            tools=EMPTY_TOOLS_JSON,
-            icon=self._prepare_icon(icon, icon_type, icon_background),
+            tools="[]",
+            icon=json.dumps({"content": icon, "background": icon_background}) if icon_type == "emoji" else icon,
             server_identifier=server_identifier,
-            timeout=configuration.timeout,
-            sse_read_timeout=configuration.sse_read_timeout,
+            timeout=timeout,
+            sse_read_timeout=sse_read_timeout,
             encrypted_headers=encrypted_headers,
-            encrypted_credentials=encrypted_credentials,
+        )
+        db.session.add(mcp_tool)
+        db.session.commit()
+        return ToolTransformService.mcp_provider_to_user_provider(mcp_tool, for_list=True)
+
+    @staticmethod
+    def retrieve_mcp_tools(tenant_id: str, for_list: bool = False) -> list[ToolProviderApiEntity]:
+        mcp_providers = (
+            db.session.query(MCPToolProvider)
+            .where(MCPToolProvider.tenant_id == tenant_id)
+            .order_by(MCPToolProvider.name)
+            .all()
+        )
+        return [
+            ToolTransformService.mcp_provider_to_user_provider(mcp_provider, for_list=for_list)
+            for mcp_provider in mcp_providers
+        ]
+
+    @classmethod
+    def list_mcp_tool_from_remote_server(cls, tenant_id: str, provider_id: str) -> ToolProviderApiEntity:
+        mcp_provider = cls.get_mcp_provider_by_provider_id(provider_id, tenant_id)
+        server_url = mcp_provider.decrypted_server_url
+        authed = mcp_provider.authed
+        headers = mcp_provider.decrypted_headers
+        timeout = mcp_provider.timeout
+        sse_read_timeout = mcp_provider.sse_read_timeout
+
+        try:
+            with MCPClient(
+                server_url,
+                provider_id,
+                tenant_id,
+                authed=authed,
+                for_list=True,
+                headers=headers,
+                timeout=timeout,
+                sse_read_timeout=sse_read_timeout,
+            ) as mcp_client:
+                tools = mcp_client.list_tools()
+        except MCPAuthError:
+            raise ValueError("Please auth the tool first")
+        except MCPError as e:
+            raise ValueError(f"Failed to connect to MCP server: {e}")
+
+        try:
+            mcp_provider = cls.get_mcp_provider_by_provider_id(provider_id, tenant_id)
+            mcp_provider.tools = json.dumps([tool.model_dump() for tool in tools])
+            mcp_provider.authed = True
+            mcp_provider.updated_at = datetime.now()
+            db.session.commit()
+        except Exception:
+            db.session.rollback()
+            raise
+
+        user = mcp_provider.load_user()
+        if not mcp_provider.icon:
+            raise ValueError("MCP provider icon is required")
+        return ToolProviderApiEntity(
+            id=mcp_provider.id,
+            name=mcp_provider.name,
+            tools=ToolTransformService.mcp_tool_to_user_tool(mcp_provider, tools),
+            type=ToolProviderType.MCP,
+            icon=mcp_provider.icon,
+            author=user.name if user else "Anonymous",
+            server_url=mcp_provider.masked_server_url,
+            updated_at=int(mcp_provider.updated_at.timestamp()),
+            description=I18nObject(en_US="", zh_Hans=""),
+            label=I18nObject(en_US=mcp_provider.name, zh_Hans=mcp_provider.name),
+            plugin_unique_identifier=mcp_provider.server_identifier,
         )
 
-        self._session.add(mcp_tool)
-        self._session.flush()
-        mcp_providers = ToolTransformService.mcp_provider_to_user_provider(mcp_tool, for_list=True)
-        return mcp_providers
+    @classmethod
+    def delete_mcp_tool(cls, tenant_id: str, provider_id: str):
+        mcp_tool = cls.get_mcp_provider_by_provider_id(provider_id, tenant_id)
 
-    def update_provider(
-        self,
-        *,
+        db.session.delete(mcp_tool)
+        db.session.commit()
+
+    @classmethod
+    def update_mcp_provider(
+        cls,
         tenant_id: str,
         provider_id: str,
         name: str,
@@ -177,546 +222,129 @@ class MCPToolManageService:
         icon_type: str,
         icon_background: str,
         server_identifier: str,
+        timeout: float | None = None,
+        sse_read_timeout: float | None = None,
         headers: dict[str, str] | None = None,
-        configuration: MCPConfiguration,
-        authentication: MCPAuthentication | None = None,
-        validation_result: ServerUrlValidationResult | None = None,
-    ) -> None:
-        """
-        Update an MCP provider.
+    ):
+        mcp_provider = cls.get_mcp_provider_by_provider_id(provider_id, tenant_id)
 
-        Args:
-            validation_result: Pre-validation result from validate_server_url_change.
-                              If provided and contains reconnect_result, it will be used
-                              instead of performing network operations.
-        """
-        mcp_provider = self.get_provider(provider_id=provider_id, tenant_id=tenant_id)
-
-        # Check for duplicate name (excluding current provider)
-        if name != mcp_provider.name:
-            stmt = select(MCPToolProvider).where(
-                MCPToolProvider.tenant_id == tenant_id,
-                MCPToolProvider.name == name,
-                MCPToolProvider.id != provider_id,
-            )
-            existing_provider = self._session.scalar(stmt)
-            if existing_provider:
-                raise ValueError(f"MCP tool {name} already exists")
-
-        # Get URL update data from validation result
+        reconnect_result = None
         encrypted_server_url = None
         server_url_hash = None
-        reconnect_result = None
 
-        if validation_result and validation_result.encrypted_server_url:
-            # Use all data from validation result
-            encrypted_server_url = validation_result.encrypted_server_url
-            server_url_hash = validation_result.server_url_hash
-            reconnect_result = validation_result.reconnect_result
+        if UNCHANGED_SERVER_URL_PLACEHOLDER not in server_url:
+            encrypted_server_url = encrypter.encrypt_token(tenant_id, server_url)
+            server_url_hash = hashlib.sha256(server_url.encode()).hexdigest()
+
+            if server_url_hash != mcp_provider.server_url_hash:
+                reconnect_result = cls._re_connect_mcp_provider(server_url, provider_id, tenant_id)
 
         try:
-            # Update basic fields
             mcp_provider.updated_at = datetime.now()
             mcp_provider.name = name
-            mcp_provider.icon = self._prepare_icon(icon, icon_type, icon_background)
+            mcp_provider.icon = (
+                json.dumps({"content": icon, "background": icon_background}) if icon_type == "emoji" else icon
+            )
             mcp_provider.server_identifier = server_identifier
 
-            # Update server URL if changed
-            if encrypted_server_url and server_url_hash:
+            if encrypted_server_url is not None and server_url_hash is not None:
                 mcp_provider.server_url = encrypted_server_url
                 mcp_provider.server_url_hash = server_url_hash
 
                 if reconnect_result:
-                    mcp_provider.authed = reconnect_result.authed
-                    mcp_provider.tools = reconnect_result.tools
-                    mcp_provider.encrypted_credentials = reconnect_result.encrypted_credentials
+                    mcp_provider.authed = reconnect_result["authed"]
+                    mcp_provider.tools = reconnect_result["tools"]
+                    mcp_provider.encrypted_credentials = reconnect_result["encrypted_credentials"]
 
-            # Update optional configuration fields
-            self._update_optional_fields(mcp_provider, configuration)
-
-            # Update headers if provided
+            if timeout is not None:
+                mcp_provider.timeout = timeout
+            if sse_read_timeout is not None:
+                mcp_provider.sse_read_timeout = sse_read_timeout
             if headers is not None:
-                mcp_provider.encrypted_headers = self._process_headers(headers, mcp_provider, tenant_id)
+                # Merge masked headers from frontend with existing real values
+                if headers:
+                    # existing decrypted and masked headers
+                    existing_decrypted = mcp_provider.decrypted_headers
+                    existing_masked = mcp_provider.masked_headers
 
-            # Update credentials if provided
-            if authentication and authentication.client_id:
-                mcp_provider.encrypted_credentials = self._process_credentials(authentication, mcp_provider, tenant_id)
+                    # Build final headers: if value equals masked existing, keep original decrypted value
+                    final_headers: dict[str, str] = {}
+                    for key, incoming_value in headers.items():
+                        if (
+                            key in existing_masked
+                            and key in existing_decrypted
+                            and isinstance(incoming_value, str)
+                            and incoming_value == existing_masked.get(key)
+                        ):
+                            # unchanged, use original decrypted value
+                            final_headers[key] = str(existing_decrypted[key])
+                        else:
+                            final_headers[key] = incoming_value
 
-            # Flush changes to database
-            self._session.flush()
+                    encrypted_headers_dict = MCPToolManageService._encrypt_headers(final_headers, tenant_id)
+                    mcp_provider.encrypted_headers = json.dumps(encrypted_headers_dict)
+                else:
+                    # Explicitly clear headers if empty dict passed
+                    mcp_provider.encrypted_headers = None
+            db.session.commit()
         except IntegrityError as e:
-            self._handle_integrity_error(e, name, server_url, server_identifier)
+            db.session.rollback()
+            error_msg = str(e.orig)
+            if "unique_mcp_provider_name" in error_msg:
+                raise ValueError(f"MCP tool {name} already exists")
+            if "unique_mcp_provider_server_url" in error_msg:
+                raise ValueError(f"MCP tool {server_url} already exists")
+            if "unique_mcp_provider_server_identifier" in error_msg:
+                raise ValueError(f"MCP tool {server_identifier} already exists")
+            raise
+        except Exception:
+            db.session.rollback()
+            raise
 
-    def delete_provider(self, *, tenant_id: str, provider_id: str) -> None:
-        """Delete an MCP provider."""
-        mcp_tool = self.get_provider(provider_id=provider_id, tenant_id=tenant_id)
-        self._session.delete(mcp_tool)
-
-    def list_providers(
-        self, *, tenant_id: str, for_list: bool = False, include_sensitive: bool = True
-    ) -> list[ToolProviderApiEntity]:
-        """List all MCP providers for a tenant.
-
-        Args:
-            tenant_id: Tenant ID
-            for_list: If True, return provider ID; if False, return server identifier
-            include_sensitive: If False, skip expensive decryption operations (default: True for backward compatibility)
-        """
-        from models.account import Account
-
-        stmt = select(MCPToolProvider).where(MCPToolProvider.tenant_id == tenant_id).order_by(MCPToolProvider.name)
-        mcp_providers = self._session.scalars(stmt).all()
-
-        if not mcp_providers:
-            return []
-
-        # Batch query all users to avoid N+1 problem
-        user_ids = {provider.user_id for provider in mcp_providers}
-        users = self._session.query(Account).where(Account.id.in_(user_ids)).all()
-        user_name_map = {user.id: user.name for user in users}
-
-        return [
-            ToolTransformService.mcp_provider_to_user_provider(
-                provider,
-                for_list=for_list,
-                user_name=user_name_map.get(provider.user_id),
-                include_sensitive=include_sensitive,
-            )
-            for provider in mcp_providers
-        ]
-
-    # ========== Tool Operations ==========
-
-    def list_provider_tools(self, *, tenant_id: str, provider_id: str) -> ToolProviderApiEntity:
-        """List tools from remote MCP server."""
-        # Load provider and convert to entity
-        db_provider = self.get_provider(provider_id=provider_id, tenant_id=tenant_id)
-        provider_entity = db_provider.to_entity()
-
-        # Verify authentication
-        if not provider_entity.authed:
-            raise ValueError("Please auth the tool first")
-
-        # Prepare headers with auth token
-        headers = self._prepare_auth_headers(provider_entity)
-
-        # Retrieve tools from remote server
-        server_url = provider_entity.decrypt_server_url()
-        try:
-            tools = self._retrieve_remote_mcp_tools(server_url, headers, provider_entity)
-        except MCPError as e:
-            raise ValueError(f"Failed to connect to MCP server: {e}")
-
-        # Update database with retrieved tools
-        db_provider.tools = json.dumps([tool.model_dump() for tool in tools])
-        db_provider.authed = True
-        db_provider.updated_at = datetime.now()
-        self._session.flush()
-
-        # Build API response
-        return self._build_tool_provider_response(db_provider, provider_entity, tools)
-
-    # ========== OAuth and Credentials Operations ==========
-
-    def update_provider_credentials(
-        self, *, provider_id: str, tenant_id: str, credentials: dict[str, Any], authed: bool | None = None
-    ) -> None:
-        """
-        Update provider credentials with encryption.
-
-        Args:
-            provider_id: Provider ID
-            tenant_id: Tenant ID
-            credentials: Credentials to save
-            authed: Whether provider is authenticated (None means keep current state)
-        """
-        from core.tools.mcp_tool.provider import MCPToolProviderController
-
-        # Get provider from current session
-        provider = self.get_provider(provider_id=provider_id, tenant_id=tenant_id)
-
-        # Encrypt new credentials
-        provider_controller = MCPToolProviderController.from_db(provider)
+    @classmethod
+    def update_mcp_provider_credentials(
+        cls, mcp_provider: MCPToolProvider, credentials: dict[str, Any], authed: bool = False
+    ):
+        provider_controller = MCPToolProviderController.from_db(mcp_provider)
         tool_configuration = ProviderConfigEncrypter(
-            tenant_id=provider.tenant_id,
+            tenant_id=mcp_provider.tenant_id,
             config=list(provider_controller.get_credentials_schema()),
             provider_config_cache=NoOpProviderCredentialCache(),
         )
-        encrypted_credentials = tool_configuration.encrypt(credentials)
+        credentials = tool_configuration.encrypt(credentials)
+        mcp_provider.updated_at = datetime.now()
+        mcp_provider.encrypted_credentials = json.dumps({**mcp_provider.credentials, **credentials})
+        mcp_provider.authed = authed
+        if not authed:
+            mcp_provider.tools = "[]"
+        db.session.commit()
 
-        # Update provider
-        provider.updated_at = datetime.now()
-        provider.encrypted_credentials = json.dumps({**provider.credentials, **encrypted_credentials})
-
-        if authed is not None:
-            provider.authed = authed
-            if not authed:
-                provider.tools = EMPTY_TOOLS_JSON
-
-        # Flush changes to database
-        self._session.flush()
-
-    def save_oauth_data(
-        self, provider_id: str, tenant_id: str, data: dict[str, Any], data_type: OAuthDataType = OAuthDataType.MIXED
-    ) -> None:
-        """
-        Save OAuth-related data (tokens, client info, code verifier).
-
-        Args:
-            provider_id: Provider ID
-            tenant_id: Tenant ID
-            data: Data to save (tokens, client info, or code verifier)
-            data_type: Type of OAuth data to save
-        """
-        # Determine if this makes the provider authenticated
-        authed = (
-            data_type == OAuthDataType.TOKENS or (data_type == OAuthDataType.MIXED and "access_token" in data) or None
-        )
-
-        # update_provider_credentials will validate provider existence
-        self.update_provider_credentials(provider_id=provider_id, tenant_id=tenant_id, credentials=data, authed=authed)
-
-    def clear_provider_credentials(self, *, provider_id: str, tenant_id: str) -> None:
-        """
-        Clear all credentials for a provider.
-
-        Args:
-            provider_id: Provider ID
-            tenant_id: Tenant ID
-        """
-        # Get provider from current session
-        provider = self.get_provider(provider_id=provider_id, tenant_id=tenant_id)
-
-        provider.tools = EMPTY_TOOLS_JSON
-        provider.encrypted_credentials = EMPTY_CREDENTIALS_JSON
-        provider.updated_at = datetime.now()
-        provider.authed = False
-
-    # ========== Private Helper Methods ==========
-
-    def _check_provider_exists(self, tenant_id: str, name: str, server_url_hash: str, server_identifier: str) -> None:
-        """Check if provider with same attributes already exists."""
-        stmt = select(MCPToolProvider).where(
-            MCPToolProvider.tenant_id == tenant_id,
-            or_(
-                MCPToolProvider.name == name,
-                MCPToolProvider.server_url_hash == server_url_hash,
-                MCPToolProvider.server_identifier == server_identifier,
-            ),
-        )
-        existing_provider = self._session.scalar(stmt)
-
-        if existing_provider:
-            if existing_provider.name == name:
-                raise ValueError(f"MCP tool {name} already exists")
-            if existing_provider.server_url_hash == server_url_hash:
-                raise ValueError("MCP tool with this server URL already exists")
-            if existing_provider.server_identifier == server_identifier:
-                raise ValueError(f"MCP tool {server_identifier} already exists")
-
-    def _prepare_icon(self, icon: str, icon_type: str, icon_background: str) -> str:
-        """Prepare icon data for storage."""
-        if icon_type == "emoji":
-            return json.dumps({"content": icon, "background": icon_background})
-        return icon
-
-    def _encrypt_dict_fields(self, data: dict[str, Any], secret_fields: list[str], tenant_id: str) -> dict[str, str]:
-        """Encrypt specified fields in a dictionary.
-
-        Args:
-            data: Dictionary containing data to encrypt
-            secret_fields: List of field names to encrypt
-            tenant_id: Tenant ID for encryption
-
-        Returns:
-            JSON string of encrypted data
-        """
-        from core.entities.provider_entities import BasicProviderConfig
-        from core.tools.utils.encryption import create_provider_encrypter
-
-        # Create config for secret fields
-        config = [
-            BasicProviderConfig(type=BasicProviderConfig.Type.SECRET_INPUT, name=field) for field in secret_fields
-        ]
-
-        encrypter_instance, _ = create_provider_encrypter(
-            tenant_id=tenant_id,
-            config=config,
-            cache=NoOpProviderCredentialCache(),
-        )
-
-        encrypted_data = encrypter_instance.encrypt(data)
-        return encrypted_data
-
-    def _prepare_encrypted_dict(self, headers: dict[str, str], tenant_id: str) -> str:
-        """Encrypt headers and prepare for storage."""
-        # All headers are treated as secret
-        return json.dumps(self._encrypt_dict_fields(headers, list(headers.keys()), tenant_id))
-
-    def _prepare_auth_headers(self, provider_entity: MCPProviderEntity) -> dict[str, str]:
-        """Prepare headers with OAuth token if available."""
-        headers = provider_entity.decrypt_headers()
-        tokens = provider_entity.retrieve_tokens()
-        if tokens:
-            headers["Authorization"] = f"{tokens.token_type.capitalize()} {tokens.access_token}"
-        return headers
-
-    def _retrieve_remote_mcp_tools(
-        self,
-        server_url: str,
-        headers: dict[str, str],
-        provider_entity: MCPProviderEntity,
-    ):
-        """Retrieve tools from remote MCP server."""
-        with MCPClientWithAuthRetry(
-            server_url=server_url,
-            headers=headers,
-            timeout=provider_entity.timeout,
-            sse_read_timeout=provider_entity.sse_read_timeout,
-            provider_entity=provider_entity,
-        ) as mcp_client:
-            return mcp_client.list_tools()
-
-    def execute_auth_actions(self, auth_result: Any) -> dict[str, str]:
-        """
-        Execute the actions returned by the auth function.
-
-        This method processes the AuthResult and performs the necessary database operations.
-
-        Args:
-            auth_result: The result from the auth function
-
-        Returns:
-            The response from the auth result
-        """
-        from core.mcp.entities import AuthAction, AuthActionType
-
-        action: AuthAction
-        for action in auth_result.actions:
-            if action.provider_id is None or action.tenant_id is None:
-                continue
-
-            if action.action_type == AuthActionType.SAVE_CLIENT_INFO:
-                self.save_oauth_data(action.provider_id, action.tenant_id, action.data, OAuthDataType.CLIENT_INFO)
-            elif action.action_type == AuthActionType.SAVE_TOKENS:
-                self.save_oauth_data(action.provider_id, action.tenant_id, action.data, OAuthDataType.TOKENS)
-            elif action.action_type == AuthActionType.SAVE_CODE_VERIFIER:
-                self.save_oauth_data(action.provider_id, action.tenant_id, action.data, OAuthDataType.CODE_VERIFIER)
-
-        return auth_result.response
-
-    def auth_with_actions(
-        self, provider_entity: MCPProviderEntity, authorization_code: str | None = None
-    ) -> dict[str, str]:
-        """
-        Perform authentication and execute all resulting actions.
-
-        This method is used by MCPClientWithAuthRetry for automatic re-authentication.
-
-        Args:
-            provider_entity: The MCP provider entity
-            authorization_code: Optional authorization code
-
-        Returns:
-            Response dictionary from auth result
-        """
-        auth_result = auth(provider_entity, authorization_code)
-        return self.execute_auth_actions(auth_result)
-
-    def _reconnect_provider(self, *, server_url: str, provider: MCPToolProvider) -> ReconnectResult:
-        """Attempt to reconnect to MCP provider with new server URL."""
-        provider_entity = provider.to_entity()
-        headers = provider_entity.headers
+    @classmethod
+    def _re_connect_mcp_provider(cls, server_url: str, provider_id: str, tenant_id: str):
+        # Get the existing provider to access headers and timeout settings
+        mcp_provider = cls.get_mcp_provider_by_provider_id(provider_id, tenant_id)
+        headers = mcp_provider.decrypted_headers
+        timeout = mcp_provider.timeout
+        sse_read_timeout = mcp_provider.sse_read_timeout
 
         try:
-            tools = self._retrieve_remote_mcp_tools(server_url, headers, provider_entity)
-            return ReconnectResult(
-                authed=True,
-                tools=json.dumps([tool.model_dump() for tool in tools]),
-                encrypted_credentials=EMPTY_CREDENTIALS_JSON,
-            )
+            with MCPClient(
+                server_url,
+                provider_id,
+                tenant_id,
+                authed=False,
+                for_list=True,
+                headers=headers,
+                timeout=timeout,
+                sse_read_timeout=sse_read_timeout,
+            ) as mcp_client:
+                tools = mcp_client.list_tools()
+                return {
+                    "authed": True,
+                    "tools": json.dumps([tool.model_dump() for tool in tools]),
+                    "encrypted_credentials": "{}",
+                }
         except MCPAuthError:
-            return ReconnectResult(authed=False, tools=EMPTY_TOOLS_JSON, encrypted_credentials=EMPTY_CREDENTIALS_JSON)
+            return {"authed": False, "tools": "[]", "encrypted_credentials": "{}"}
         except MCPError as e:
             raise ValueError(f"Failed to re-connect MCP server: {e}") from e
-
-    def validate_server_url_change(
-        self, *, tenant_id: str, provider_id: str, new_server_url: str
-    ) -> ServerUrlValidationResult:
-        """
-        Validate server URL change by attempting to connect to the new server.
-        This method should be called BEFORE update_provider to perform network operations
-        outside of the database transaction.
-
-        Returns:
-            ServerUrlValidationResult: Validation result with connection status and tools if successful
-        """
-        # Handle hidden/unchanged URL
-        if UNCHANGED_SERVER_URL_PLACEHOLDER in new_server_url:
-            return ServerUrlValidationResult(needs_validation=False)
-
-        # Validate URL format
-        if not self._is_valid_url(new_server_url):
-            raise ValueError("Server URL is not valid.")
-
-        # Always encrypt and hash the URL
-        encrypted_server_url = encrypter.encrypt_token(tenant_id, new_server_url)
-        new_server_url_hash = hashlib.sha256(new_server_url.encode()).hexdigest()
-
-        # Get current provider
-        provider = self.get_provider(provider_id=provider_id, tenant_id=tenant_id)
-
-        # Check if URL is actually different
-        if new_server_url_hash == provider.server_url_hash:
-            # URL hasn't changed, but still return the encrypted data
-            return ServerUrlValidationResult(
-                needs_validation=False, encrypted_server_url=encrypted_server_url, server_url_hash=new_server_url_hash
-            )
-
-        # Perform validation by attempting to connect
-        reconnect_result = self._reconnect_provider(server_url=new_server_url, provider=provider)
-        return ServerUrlValidationResult(
-            needs_validation=True,
-            validation_passed=True,
-            reconnect_result=reconnect_result,
-            encrypted_server_url=encrypted_server_url,
-            server_url_hash=new_server_url_hash,
-        )
-
-    def _build_tool_provider_response(
-        self, db_provider: MCPToolProvider, provider_entity: MCPProviderEntity, tools: list
-    ) -> ToolProviderApiEntity:
-        """Build API response for tool provider."""
-        user = db_provider.load_user()
-        response = provider_entity.to_api_response(
-            user_name=user.name if user else None,
-        )
-        response["tools"] = ToolTransformService.mcp_tool_to_user_tool(db_provider, tools)
-        response["plugin_unique_identifier"] = provider_entity.provider_id
-        return ToolProviderApiEntity(**response)
-
-    def _handle_integrity_error(
-        self, error: IntegrityError, name: str, server_url: str, server_identifier: str
-    ) -> None:
-        """Handle database integrity errors with user-friendly messages."""
-        error_msg = str(error.orig)
-        if "unique_mcp_provider_name" in error_msg:
-            raise ValueError(f"MCP tool {name} already exists")
-        if "unique_mcp_provider_server_url" in error_msg:
-            raise ValueError(f"MCP tool {server_url} already exists")
-        if "unique_mcp_provider_server_identifier" in error_msg:
-            raise ValueError(f"MCP tool {server_identifier} already exists")
-        raise
-
-    def _is_valid_url(self, url: str) -> bool:
-        """Validate URL format."""
-        if not url:
-            return False
-        try:
-            parsed = urlparse(url)
-            return all([parsed.scheme, parsed.netloc]) and parsed.scheme in ["http", "https"]
-        except (ValueError, TypeError):
-            return False
-
-    def _update_optional_fields(self, mcp_provider: MCPToolProvider, configuration: MCPConfiguration) -> None:
-        """Update optional configuration fields using setattr for cleaner code."""
-        field_mapping = {"timeout": configuration.timeout, "sse_read_timeout": configuration.sse_read_timeout}
-
-        for field, value in field_mapping.items():
-            if value is not None:
-                setattr(mcp_provider, field, value)
-
-    def _process_headers(self, headers: dict[str, str], mcp_provider: MCPToolProvider, tenant_id: str) -> str | None:
-        """Process headers update, handling empty dict to clear headers."""
-        if not headers:
-            return None
-
-        # Merge with existing headers to preserve masked values
-        final_headers = self._merge_headers_with_masked(incoming_headers=headers, mcp_provider=mcp_provider)
-        return self._prepare_encrypted_dict(final_headers, tenant_id)
-
-    def _process_credentials(
-        self, authentication: MCPAuthentication, mcp_provider: MCPToolProvider, tenant_id: str
-    ) -> str:
-        """Process credentials update, handling masked values."""
-        # Merge with existing credentials
-        final_client_id, final_client_secret = self._merge_credentials_with_masked(
-            authentication.client_id, authentication.client_secret, mcp_provider
-        )
-
-        # Build and encrypt
-        return self._build_and_encrypt_credentials(final_client_id, final_client_secret, tenant_id)
-
-    def _merge_headers_with_masked(
-        self, incoming_headers: dict[str, str], mcp_provider: MCPToolProvider
-    ) -> dict[str, str]:
-        """Merge incoming headers with existing ones, preserving unchanged masked values.
-
-        Args:
-            incoming_headers: Headers from frontend (may contain masked values)
-            mcp_provider: The MCP provider instance
-
-        Returns:
-            Final headers dict with proper values (original for unchanged masked, new for changed)
-        """
-        mcp_provider_entity = mcp_provider.to_entity()
-        existing_decrypted = mcp_provider_entity.decrypt_headers()
-        existing_masked = mcp_provider_entity.masked_headers()
-
-        return {
-            key: (str(existing_decrypted[key]) if key in existing_masked and value == existing_masked[key] else value)
-            for key, value in incoming_headers.items()
-            if key in existing_decrypted or value != existing_masked.get(key)
-        }
-
-    def _merge_credentials_with_masked(
-        self,
-        client_id: str,
-        client_secret: str | None,
-        mcp_provider: MCPToolProvider,
-    ) -> tuple[
-        str,
-        str | None,
-    ]:
-        """Merge incoming credentials with existing ones, preserving unchanged masked values.
-
-        Args:
-            client_id: Client ID from frontend (may be masked)
-            client_secret: Client secret from frontend (may be masked)
-            mcp_provider: The MCP provider instance
-
-        Returns:
-            Tuple of (final_client_id, final_client_secret)
-        """
-        mcp_provider_entity = mcp_provider.to_entity()
-        existing_decrypted = mcp_provider_entity.decrypt_credentials()
-        existing_masked = mcp_provider_entity.masked_credentials()
-
-        # Check if client_id is masked and unchanged
-        final_client_id = client_id
-        if existing_masked.get("client_id") and client_id == existing_masked["client_id"]:
-            # Use existing decrypted value
-            final_client_id = existing_decrypted.get("client_id", client_id)
-
-        # Check if client_secret is masked and unchanged
-        final_client_secret = client_secret
-        if existing_masked.get("client_secret") and client_secret == existing_masked["client_secret"]:
-            # Use existing decrypted value
-            final_client_secret = existing_decrypted.get("client_secret", client_secret)
-
-        return final_client_id, final_client_secret
-
-    def _build_and_encrypt_credentials(self, client_id: str, client_secret: str | None, tenant_id: str) -> str:
-        """Build credentials and encrypt sensitive fields."""
-        # Create a flat structure with all credential data
-        credentials_data = {
-            "client_id": client_id,
-            "client_name": CLIENT_NAME,
-            "is_dynamic_registration": False,
-        }
-        secret_fields = []
-        if client_secret is not None:
-            credentials_data["encrypted_client_secret"] = client_secret
-            secret_fields = ["encrypted_client_secret"]
-        client_info = self._encrypt_dict_fields(credentials_data, secret_fields, tenant_id)
-        return json.dumps({"client_information": client_info})

api/services/tools/tools_transform_service.py

@@ -3,7 +3,6 @@ import logging
 from collections.abc import Mapping
 from typing import Any, Union
 
-from pydantic import ValidationError
 from yarl import URL
 
 from configs import dify_config
@@ -233,59 +232,40 @@ class ToolTransformService:
         )
 
     @staticmethod
-    def mcp_provider_to_user_provider(
-        db_provider: MCPToolProvider,
-        for_list: bool = False,
-        user_name: str | None = None,
-        include_sensitive: bool = True,
-    ) -> ToolProviderApiEntity:
-        from core.entities.mcp_provider import MCPConfiguration
-
-        # Use provided user_name to avoid N+1 query, fallback to load_user() if not provided
-        if user_name is None:
-            user = db_provider.load_user()
-            user_name = user.name if user else None
-
-        # Convert to entity and use its API response method
-        provider_entity = db_provider.to_entity()
-
-        response = provider_entity.to_api_response(user_name=user_name, include_sensitive=include_sensitive)
-        try:
-            mcp_tools = [MCPTool(**tool) for tool in json.loads(db_provider.tools)]
-        except (ValidationError, json.JSONDecodeError):
-            mcp_tools = []
-        # Add additional fields specific to the transform
-        response["id"] = db_provider.server_identifier if not for_list else db_provider.id
-        response["tools"] = ToolTransformService.mcp_tool_to_user_tool(db_provider, mcp_tools, user_name=user_name)
-        response["server_identifier"] = db_provider.server_identifier
-
-        # Convert configuration dict to MCPConfiguration object
-        if "configuration" in response and isinstance(response["configuration"], dict):
-            response["configuration"] = MCPConfiguration(
-                timeout=float(response["configuration"]["timeout"]),
-                sse_read_timeout=float(response["configuration"]["sse_read_timeout"]),
-            )
-
-        return ToolProviderApiEntity(**response)
+    def mcp_provider_to_user_provider(db_provider: MCPToolProvider, for_list: bool = False) -> ToolProviderApiEntity:
+        user = db_provider.load_user()
+        return ToolProviderApiEntity(
+            id=db_provider.server_identifier if not for_list else db_provider.id,
+            author=user.name if user else "Anonymous",
+            name=db_provider.name,
+            icon=db_provider.provider_icon,
+            type=ToolProviderType.MCP,
+            is_team_authorization=db_provider.authed,
+            server_url=db_provider.masked_server_url,
+            tools=ToolTransformService.mcp_tool_to_user_tool(
+                db_provider, [MCPTool.model_validate(tool) for tool in json.loads(db_provider.tools)]
+            ),
+            updated_at=int(db_provider.updated_at.timestamp()),
+            label=I18nObject(en_US=db_provider.name, zh_Hans=db_provider.name),
+            description=I18nObject(en_US="", zh_Hans=""),
+            server_identifier=db_provider.server_identifier,
+            timeout=db_provider.timeout,
+            sse_read_timeout=db_provider.sse_read_timeout,
+            masked_headers=db_provider.masked_headers,
+            original_headers=db_provider.decrypted_headers,
+        )
 
     @staticmethod
-    def mcp_tool_to_user_tool(
-        mcp_provider: MCPToolProvider, tools: list[MCPTool], user_name: str | None = None
-    ) -> list[ToolApiEntity]:
-        # Use provided user_name to avoid N+1 query, fallback to load_user() if not provided
-        if user_name is None:
-            user = mcp_provider.load_user()
-            user_name = user.name if user else "Anonymous"
-
+    def mcp_tool_to_user_tool(mcp_provider: MCPToolProvider, tools: list[MCPTool]) -> list[ToolApiEntity]:
+        user = mcp_provider.load_user()
         return [
             ToolApiEntity(
-                author=user_name or "Anonymous",
+                author=user.name if user else "Anonymous",
                 name=tool.name,
                 label=I18nObject(en_US=tool.name, zh_Hans=tool.name),
                 description=I18nObject(en_US=tool.description or "", zh_Hans=tool.description or ""),
                 parameters=ToolTransformService.convert_mcp_schema_to_parameter(tool.inputSchema),
                 labels=[],
-                output_schema=tool.outputSchema or {},
             )
             for tool in tools
         ]
@@ -432,7 +412,7 @@ class ToolTransformService:
         )
 
     @staticmethod
-    def convert_mcp_schema_to_parameter(schema: dict[str, Any]) -> list["ToolParameter"]:
+    def convert_mcp_schema_to_parameter(schema: dict) -> list["ToolParameter"]:
         """
         Convert MCP JSON schema to tool parameters
 
@@ -441,7 +421,7 @@ class ToolTransformService:
         """
 
         def create_parameter(
-            name: str, description: str, param_type: str, required: bool, input_schema: dict[str, Any] | None = None
+            name: str, description: str, param_type: str, required: bool, input_schema: dict | None = None
         ) -> ToolParameter:
             """Create a ToolParameter instance with given attributes"""
             input_schema_dict: dict[str, Any] = {"input_schema": input_schema} if input_schema else {}
@@ -456,9 +436,7 @@ class ToolTransformService:
                 **input_schema_dict,
             )
 
-        def process_properties(
-            props: dict[str, dict[str, Any]], required: list[str], prefix: str = ""
-        ) -> list[ToolParameter]:
+        def process_properties(props: dict, required: list, prefix: str = "") -> list[ToolParameter]:
             """Process properties recursively"""
             TYPE_MAPPING = {"integer": "number", "float": "number"}
             COMPLEX_TYPES = ["array", "object"]

api/services/variable_truncator.py

@@ -1,5 +1,4 @@
 import dataclasses
-from abc import ABC, abstractmethod
 from collections.abc import Mapping
 from typing import Any, Generic, TypeAlias, TypeVar, overload
 
@@ -67,17 +66,7 @@ class TruncationResult:
     truncated: bool
 
 
-class BaseTruncator(ABC):
-    @abstractmethod
-    def truncate(self, segment: Segment) -> TruncationResult:
-        pass
-
-    @abstractmethod
-    def truncate_variable_mapping(self, v: Mapping[str, Any]) -> tuple[Mapping[str, Any], bool]:
-        pass
-
-
-class VariableTruncator(BaseTruncator):
+class VariableTruncator:
     """
     Handles variable truncation with structure-preserving strategies.
 
@@ -429,38 +418,3 @@ class VariableTruncator(BaseTruncator):
             return _PartResult(val, self.calculate_json_size(val), False)
         else:
             raise AssertionError("this statement should be unreachable.")
-
-
-class DummyVariableTruncator(BaseTruncator):
-    """
-    A no-op variable truncator that doesn't truncate any data.
-
-    This is used for Service API calls where truncation should be disabled
-    to maintain backward compatibility and provide complete data.
-    """
-
-    def truncate_variable_mapping(self, v: Mapping[str, Any]) -> tuple[Mapping[str, Any], bool]:
-        """
-        Return original mapping without truncation.
-
-        Args:
-            v: The variable mapping to process
-
-        Returns:
-            Tuple of (original_mapping, False) where False indicates no truncation occurred
-        """
-        return v, False
-
-    def truncate(self, segment: Segment) -> TruncationResult:
-        """
-        Return original segment without truncation.
-
-        Args:
-            segment: The segment to process
-
-        Returns:
-            The original segment unchanged
-        """
-        # For Service API, we want to preserve the original segment
-        # without any truncation, so just return it as-is
-        return TruncationResult(result=segment, truncated=False)

api/tasks/document_indexing_task.py

@@ -1,11 +1,14 @@
 import logging
 import time
+from collections.abc import Callable
 
 import click
 from celery import shared_task
 
 from configs import dify_config
+from core.entities.document_task import DocumentTask
 from core.indexing_runner import DocumentIsPausedError, IndexingRunner
+from core.rag.pipeline.queue import TenantSelfTaskQueue
 from extensions.ext_database import db
 from libs.datetime_utils import naive_utc_now
 from models.dataset import Dataset, Document
@@ -21,8 +24,24 @@ def document_indexing_task(dataset_id: str, document_ids: list):
     :param dataset_id:
     :param document_ids:
 
+    .. warning:: TO BE DEPRECATED
+        This function will be deprecated and removed in a future version.
+        Use normal_document_indexing_task or priority_document_indexing_task instead.
+
     Usage: document_indexing_task.delay(dataset_id, document_ids)
     """
+    logger.warning("document indexing legacy mode received: %s - %s", dataset_id, document_ids)
+    _document_indexing(dataset_id, document_ids)
+
+
+def _document_indexing(dataset_id: str, document_ids: list):
+    """
+    Process document for tasks
+    :param dataset_id:
+    :param document_ids:
+
+    Usage: _document_indexing(dataset_id, document_ids)
+    """
     documents = []
     start_at = time.perf_counter()
 
@@ -86,3 +105,61 @@ def document_indexing_task(dataset_id: str, document_ids: list):
         logger.exception("Document indexing task failed, dataset_id: %s", dataset_id)
     finally:
         db.session.close()
+
+
+def _document_indexing_with_tenant_queue(tenant_id: str, dataset_id: str, document_ids: list, task_func: Callable):
+    try:
+        _document_indexing(dataset_id, document_ids)
+    except Exception:
+        logger.exception("Error processing document indexing %s for tenant %s: %s", dataset_id, tenant_id)
+    finally:
+        tenant_self_task_queue = TenantSelfTaskQueue(tenant_id, "document_indexing")
+
+        # Check if there are waiting tasks in the queue
+        # Use rpop to get the next task from the queue (FIFO order)
+        next_tasks = tenant_self_task_queue.pull_tasks(count=dify_config.TENANT_SELF_TASK_QUEUE_PULL_SIZE)
+
+        logger.info("document indexing tenant isolation queue next tasks: %s", next_tasks)
+
+        if next_tasks:
+            for next_task in next_tasks:
+                document_task = DocumentTask(**next_task)
+                # Process the next waiting task
+                # Keep the flag set to indicate a task is running
+                tenant_self_task_queue.set_task_waiting_time()
+                task_func.delay(  # type: ignore
+                    tenant_id=document_task.tenant_id,
+                    dataset_id=document_task.dataset_id,
+                    document_ids=document_task.document_ids,
+                )
+        else:
+            # No more waiting tasks, clear the flag
+            tenant_self_task_queue.delete_task_key()
+
+
+@shared_task(queue="dataset")
+def normal_document_indexing_task(tenant_id: str, dataset_id: str, document_ids: list):
+    """
+    Async process document
+    :param tenant_id:
+    :param dataset_id:
+    :param document_ids:
+
+    Usage: normal_document_indexing_task.delay(tenant_id, dataset_id, document_ids)
+    """
+    logger.info("normal document indexing task received: %s - %s - %s", tenant_id, dataset_id, document_ids)
+    _document_indexing_with_tenant_queue(tenant_id, dataset_id, document_ids, normal_document_indexing_task)
+
+
+@shared_task(queue="priority_dataset")
+def priority_document_indexing_task(tenant_id: str, dataset_id: str, document_ids: list):
+    """
+    Priority async process document
+    :param tenant_id:
+    :param dataset_id:
+    :param document_ids:
+
+    Usage: priority_document_indexing_task.delay(tenant_id, dataset_id, document_ids)
+    """
+    logger.info("priority document indexing task received: %s - %s - %s", tenant_id, dataset_id, document_ids)
+    _document_indexing_with_tenant_queue(tenant_id, dataset_id, document_ids, priority_document_indexing_task)

api/tasks/rag_pipeline/priority_rag_pipeline_run_task.py

@@ -12,8 +12,10 @@ from celery import shared_task  # type: ignore
 from flask import current_app, g
 from sqlalchemy.orm import Session, sessionmaker
 
+from configs import dify_config
 from core.app.entities.app_invoke_entities import InvokeFrom, RagPipelineGenerateEntity
 from core.app.entities.rag_pipeline_invoke_entities import RagPipelineInvokeEntity
+from core.rag.pipeline.queue import TenantSelfTaskQueue
 from core.repositories.factory import DifyCoreRepositoryFactory
 from extensions.ext_database import db
 from models import Account, Tenant
@@ -22,6 +24,8 @@ from models.enums import WorkflowRunTriggeredFrom
 from models.workflow import Workflow, WorkflowNodeExecutionTriggeredFrom
 from services.file_service import FileService
 
+logger = logging.getLogger(__name__)
+
 
 @shared_task(queue="priority_pipeline")
 def priority_rag_pipeline_run_task(
@@ -69,6 +73,27 @@ def priority_rag_pipeline_run_task(
         logging.exception(click.style(f"Error running rag pipeline, tenant_id: {tenant_id}", fg="red"))
         raise
     finally:
+        tenant_self_pipeline_task_queue = TenantSelfTaskQueue(tenant_id, "pipeline")
+
+        # Check if there are waiting tasks in the queue
+        # Use rpop to get the next task from the queue (FIFO order)
+        next_file_ids = tenant_self_pipeline_task_queue.pull_tasks(count=dify_config.TENANT_SELF_TASK_QUEUE_PULL_SIZE)
+        logger.info("priority rag pipeline tenant isolation queue next files: %s", next_file_ids)
+
+        if next_file_ids:
+            for next_file_id in next_file_ids:
+                # Process the next waiting task
+                # Keep the flag set to indicate a task is running
+                tenant_self_pipeline_task_queue.set_task_waiting_time()
+                priority_rag_pipeline_run_task.delay(  # type: ignore
+                    rag_pipeline_invoke_entities_file_id=next_file_id.decode("utf-8")
+                    if isinstance(next_file_id, bytes)
+                    else next_file_id,
+                    tenant_id=tenant_id,
+                )
+        else:
+            # No more waiting tasks, clear the flag
+            tenant_self_pipeline_task_queue.delete_task_key()
         file_service = FileService(db.engine)
         file_service.delete_file(rag_pipeline_invoke_entities_file_id)
         db.session.close()

api/tasks/rag_pipeline/rag_pipeline_run_task.py

@@ -12,17 +12,20 @@ from celery import shared_task  # type: ignore
 from flask import current_app, g
 from sqlalchemy.orm import Session, sessionmaker
 
+from configs import dify_config
 from core.app.entities.app_invoke_entities import InvokeFrom, RagPipelineGenerateEntity
 from core.app.entities.rag_pipeline_invoke_entities import RagPipelineInvokeEntity
+from core.rag.pipeline.queue import TenantSelfTaskQueue
 from core.repositories.factory import DifyCoreRepositoryFactory
 from extensions.ext_database import db
-from extensions.ext_redis import redis_client
 from models import Account, Tenant
 from models.dataset import Pipeline
 from models.enums import WorkflowRunTriggeredFrom
 from models.workflow import Workflow, WorkflowNodeExecutionTriggeredFrom
 from services.file_service import FileService
 
+logger = logging.getLogger(__name__)
+
 
 @shared_task(queue="pipeline")
 def rag_pipeline_run_task(
@@ -70,26 +73,27 @@ def rag_pipeline_run_task(
         logging.exception(click.style(f"Error running rag pipeline, tenant_id: {tenant_id}", fg="red"))
         raise
     finally:
-        tenant_self_pipeline_task_queue = f"tenant_self_pipeline_task_queue:{tenant_id}"
-        tenant_pipeline_task_key = f"tenant_pipeline_task:{tenant_id}"
+        tenant_self_pipeline_task_queue = TenantSelfTaskQueue(tenant_id, "pipeline")
 
         # Check if there are waiting tasks in the queue
         # Use rpop to get the next task from the queue (FIFO order)
-        next_file_id = redis_client.rpop(tenant_self_pipeline_task_queue)
+        next_file_ids = tenant_self_pipeline_task_queue.pull_tasks(count=dify_config.TENANT_SELF_TASK_QUEUE_PULL_SIZE)
+        logger.info("rag pipeline tenant isolation queue next files: %s", next_file_ids)
 
-        if next_file_id:
-            # Process the next waiting task
-            # Keep the flag set to indicate a task is running
-            redis_client.setex(tenant_pipeline_task_key, 60 * 60, 1)
-            rag_pipeline_run_task.delay(  # type: ignore
-                rag_pipeline_invoke_entities_file_id=next_file_id.decode("utf-8")
-                if isinstance(next_file_id, bytes)
-                else next_file_id,
-                tenant_id=tenant_id,
-            )
+        if next_file_ids:
+            for next_file_id in next_file_ids:
+                # Process the next waiting task
+                # Keep the flag set to indicate a task is running
+                tenant_self_pipeline_task_queue.set_task_waiting_time()
+                rag_pipeline_run_task.delay(  # type: ignore
+                    rag_pipeline_invoke_entities_file_id=next_file_id.decode("utf-8")
+                    if isinstance(next_file_id, bytes)
+                    else next_file_id,
+                    tenant_id=tenant_id,
+                )
         else:
             # No more waiting tasks, clear the flag
-            redis_client.delete(tenant_pipeline_task_key)
+            tenant_self_pipeline_task_queue.delete_task_key()
         file_service = FileService(db.engine)
         file_service.delete_file(rag_pipeline_invoke_entities_file_id)
         db.session.close()

api/tests/fixtures/workflow/iteration_flatten_output_disabled_workflow.yml

@@ -1,258 +0,0 @@
-app:
-  description: 'This workflow tests the iteration node with flatten_output=False.
-
-
-    It processes [1, 2, 3], outputs [item, item*2] for each iteration.
-
-
-    With flatten_output=False, it should output nested arrays:
-
-
-    ```
-
-    {"output": [[1, 2], [2, 4], [3, 6]]}
-
-    ```'
-  icon: 🤖
-  icon_background: '#FFEAD5'
-  mode: workflow
-  name: test_iteration_flatten_disabled
-  use_icon_as_answer_icon: false
-dependencies: []
-kind: app
-version: 0.3.1
-workflow:
-  conversation_variables: []
-  environment_variables: []
-  features:
-    file_upload:
-      enabled: false
-    opening_statement: ''
-    retriever_resource:
-      enabled: true
-    sensitive_word_avoidance:
-      enabled: false
-    speech_to_text:
-      enabled: false
-    suggested_questions: []
-    suggested_questions_after_answer:
-      enabled: false
-    text_to_speech:
-      enabled: false
-  graph:
-    edges:
-    - data:
-        isInIteration: false
-        isInLoop: false
-        sourceType: start
-        targetType: code
-      id: start-source-code-target
-      source: start_node
-      sourceHandle: source
-      target: code_node
-      targetHandle: target
-      type: custom
-      zIndex: 0
-    - data:
-        isInIteration: false
-        isInLoop: false
-        sourceType: code
-        targetType: iteration
-      id: code-source-iteration-target
-      source: code_node
-      sourceHandle: source
-      target: iteration_node
-      targetHandle: target
-      type: custom
-      zIndex: 0
-    - data:
-        isInIteration: true
-        isInLoop: false
-        iteration_id: iteration_node
-        sourceType: iteration-start
-        targetType: code
-      id: iteration-start-source-code-inner-target
-      source: iteration_nodestart
-      sourceHandle: source
-      target: code_inner_node
-      targetHandle: target
-      type: custom
-      zIndex: 1002
-    - data:
-        isInIteration: false
-        isInLoop: false
-        sourceType: iteration
-        targetType: end
-      id: iteration-source-end-target
-      source: iteration_node
-      sourceHandle: source
-      target: end_node
-      targetHandle: target
-      type: custom
-      zIndex: 0
-    nodes:
-    - data:
-        desc: ''
-        selected: false
-        title: Start
-        type: start
-        variables: []
-      height: 54
-      id: start_node
-      position:
-        x: 80
-        y: 282
-      positionAbsolute:
-        x: 80
-        y: 282
-      sourcePosition: right
-      targetPosition: left
-      type: custom
-      width: 244
-    - data:
-        code: "\ndef main() -> dict:\n    return {\n        \"result\": [1, 2, 3],\n\
-          \    }\n"
-        code_language: python3
-        desc: ''
-        outputs:
-          result:
-            children: null
-            type: array[number]
-        selected: false
-        title: Generate Array
-        type: code
-        variables: []
-      height: 54
-      id: code_node
-      position:
-        x: 384
-        y: 282
-      positionAbsolute:
-        x: 384
-        y: 282
-      selected: false
-      sourcePosition: right
-      targetPosition: left
-      type: custom
-      width: 244
-    - data:
-        desc: ''
-        error_handle_mode: terminated
-        flatten_output: false
-        height: 178
-        is_parallel: false
-        iterator_input_type: array[number]
-        iterator_selector:
-        - code_node
-        - result
-        output_selector:
-        - code_inner_node
-        - result
-        output_type: array[array[number]]
-        parallel_nums: 10
-        selected: false
-        start_node_id: iteration_nodestart
-        title: Iteration with Flatten Disabled
-        type: iteration
-        width: 388
-      height: 178
-      id: iteration_node
-      position:
-        x: 684
-        y: 282
-      positionAbsolute:
-        x: 684
-        y: 282
-      selected: false
-      sourcePosition: right
-      targetPosition: left
-      type: custom
-      width: 388
-      zIndex: 1
-    - data:
-        desc: ''
-        isInIteration: true
-        selected: false
-        title: ''
-        type: iteration-start
-      draggable: false
-      height: 48
-      id: iteration_nodestart
-      parentId: iteration_node
-      position:
-        x: 24
-        y: 68
-      positionAbsolute:
-        x: 708
-        y: 350
-      selectable: false
-      sourcePosition: right
-      targetPosition: left
-      type: custom-iteration-start
-      width: 44
-      zIndex: 1002
-    - data:
-        code: "\ndef main(arg1: int) -> dict:\n    return {\n        \"result\": [arg1,\
-          \ arg1 * 2],\n    }\n"
-        code_language: python3
-        desc: ''
-        isInIteration: true
-        isInLoop: false
-        iteration_id: iteration_node
-        outputs:
-          result:
-            children: null
-            type: array[number]
-        selected: false
-        title: Generate Pair
-        type: code
-        variables:
-        - value_selector:
-          - iteration_node
-          - item
-          value_type: number
-          variable: arg1
-      height: 54
-      id: code_inner_node
-      parentId: iteration_node
-      position:
-        x: 128
-        y: 68
-      positionAbsolute:
-        x: 812
-        y: 350
-      selected: false
-      sourcePosition: right
-      targetPosition: left
-      type: custom
-      width: 244
-      zIndex: 1002
-    - data:
-        desc: ''
-        outputs:
-        - value_selector:
-          - iteration_node
-          - output
-          value_type: array[array[number]]
-          variable: output
-        selected: false
-        title: End
-        type: end
-      height: 90
-      id: end_node
-      position:
-        x: 1132
-        y: 282
-      positionAbsolute:
-        x: 1132
-        y: 282
-      selected: true
-      sourcePosition: right
-      targetPosition: left
-      type: custom
-      width: 244
-    viewport:
-      x: -476
-      y: 3
-      zoom: 1
-

api/tests/fixtures/workflow/iteration_flatten_output_enabled_workflow.yml

@@ -1,258 +0,0 @@
-app:
-  description: 'This workflow tests the iteration node with flatten_output=True.
-
-
-    It processes [1, 2, 3], outputs [item, item*2] for each iteration.
-
-
-    With flatten_output=True (default), it should output:
-
-
-    ```
-
-    {"output": [1, 2, 2, 4, 3, 6]}
-
-    ```'
-  icon: 🤖
-  icon_background: '#FFEAD5'
-  mode: workflow
-  name: test_iteration_flatten_enabled
-  use_icon_as_answer_icon: false
-dependencies: []
-kind: app
-version: 0.3.1
-workflow:
-  conversation_variables: []
-  environment_variables: []
-  features:
-    file_upload:
-      enabled: false
-    opening_statement: ''
-    retriever_resource:
-      enabled: true
-    sensitive_word_avoidance:
-      enabled: false
-    speech_to_text:
-      enabled: false
-    suggested_questions: []
-    suggested_questions_after_answer:
-      enabled: false
-    text_to_speech:
-      enabled: false
-  graph:
-    edges:
-    - data:
-        isInIteration: false
-        isInLoop: false
-        sourceType: start
-        targetType: code
-      id: start-source-code-target
-      source: start_node
-      sourceHandle: source
-      target: code_node
-      targetHandle: target
-      type: custom
-      zIndex: 0
-    - data:
-        isInIteration: false
-        isInLoop: false
-        sourceType: code
-        targetType: iteration
-      id: code-source-iteration-target
-      source: code_node
-      sourceHandle: source
-      target: iteration_node
-      targetHandle: target
-      type: custom
-      zIndex: 0
-    - data:
-        isInIteration: true
-        isInLoop: false
-        iteration_id: iteration_node
-        sourceType: iteration-start
-        targetType: code
-      id: iteration-start-source-code-inner-target
-      source: iteration_nodestart
-      sourceHandle: source
-      target: code_inner_node
-      targetHandle: target
-      type: custom
-      zIndex: 1002
-    - data:
-        isInIteration: false
-        isInLoop: false
-        sourceType: iteration
-        targetType: end
-      id: iteration-source-end-target
-      source: iteration_node
-      sourceHandle: source
-      target: end_node
-      targetHandle: target
-      type: custom
-      zIndex: 0
-    nodes:
-    - data:
-        desc: ''
-        selected: false
-        title: Start
-        type: start
-        variables: []
-      height: 54
-      id: start_node
-      position:
-        x: 80
-        y: 282
-      positionAbsolute:
-        x: 80
-        y: 282
-      sourcePosition: right
-      targetPosition: left
-      type: custom
-      width: 244
-    - data:
-        code: "\ndef main() -> dict:\n    return {\n        \"result\": [1, 2, 3],\n\
-          \    }\n"
-        code_language: python3
-        desc: ''
-        outputs:
-          result:
-            children: null
-            type: array[number]
-        selected: false
-        title: Generate Array
-        type: code
-        variables: []
-      height: 54
-      id: code_node
-      position:
-        x: 384
-        y: 282
-      positionAbsolute:
-        x: 384
-        y: 282
-      selected: false
-      sourcePosition: right
-      targetPosition: left
-      type: custom
-      width: 244
-    - data:
-        desc: ''
-        error_handle_mode: terminated
-        flatten_output: true
-        height: 178
-        is_parallel: false
-        iterator_input_type: array[number]
-        iterator_selector:
-        - code_node
-        - result
-        output_selector:
-        - code_inner_node
-        - result
-        output_type: array[array[number]]
-        parallel_nums: 10
-        selected: false
-        start_node_id: iteration_nodestart
-        title: Iteration with Flatten Enabled
-        type: iteration
-        width: 388
-      height: 178
-      id: iteration_node
-      position:
-        x: 684
-        y: 282
-      positionAbsolute:
-        x: 684
-        y: 282
-      selected: false
-      sourcePosition: right
-      targetPosition: left
-      type: custom
-      width: 388
-      zIndex: 1
-    - data:
-        desc: ''
-        isInIteration: true
-        selected: false
-        title: ''
-        type: iteration-start
-      draggable: false
-      height: 48
-      id: iteration_nodestart
-      parentId: iteration_node
-      position:
-        x: 24
-        y: 68
-      positionAbsolute:
-        x: 708
-        y: 350
-      selectable: false
-      sourcePosition: right
-      targetPosition: left
-      type: custom-iteration-start
-      width: 44
-      zIndex: 1002
-    - data:
-        code: "\ndef main(arg1: int) -> dict:\n    return {\n        \"result\": [arg1,\
-          \ arg1 * 2],\n    }\n"
-        code_language: python3
-        desc: ''
-        isInIteration: true
-        isInLoop: false
-        iteration_id: iteration_node
-        outputs:
-          result:
-            children: null
-            type: array[number]
-        selected: false
-        title: Generate Pair
-        type: code
-        variables:
-        - value_selector:
-          - iteration_node
-          - item
-          value_type: number
-          variable: arg1
-      height: 54
-      id: code_inner_node
-      parentId: iteration_node
-      position:
-        x: 128
-        y: 68
-      positionAbsolute:
-        x: 812
-        y: 350
-      selected: false
-      sourcePosition: right
-      targetPosition: left
-      type: custom
-      width: 244
-      zIndex: 1002
-    - data:
-        desc: ''
-        outputs:
-        - value_selector:
-          - iteration_node
-          - output
-          value_type: array[number]
-          variable: output
-        selected: false
-        title: End
-        type: end
-      height: 90
-      id: end_node
-      position:
-        x: 1132
-        y: 282
-      positionAbsolute:
-        x: 1132
-        y: 282
-      selected: true
-      sourcePosition: right
-      targetPosition: left
-      type: custom
-      width: 244
-    viewport:
-      x: -476
-      y: 3
-      zoom: 1
-

api/tests/integration_tests/vdb/opensearch/test_opensearch.py

@@ -182,28 +182,6 @@ class TestOpenSearchVector:
         assert len(ids) == 1
         assert ids[0] == "mock_id"
 
-    def test_delete_nonexistent_index(self):
-        """Test deleting a non-existent index."""
-        # Create a vector instance with a non-existent collection name
-        self.vector._client.indices.exists.return_value = False
-
-        # Should not raise an exception
-        self.vector.delete()
-
-        # Verify that exists was called but delete was not
-        self.vector._client.indices.exists.assert_called_once_with(index=self.collection_name.lower())
-        self.vector._client.indices.delete.assert_not_called()
-
-    def test_delete_existing_index(self):
-        """Test deleting an existing index."""
-        self.vector._client.indices.exists.return_value = True
-
-        self.vector.delete()
-
-        # Verify both exists and delete were called
-        self.vector._client.indices.exists.assert_called_once_with(index=self.collection_name.lower())
-        self.vector._client.indices.delete.assert_called_once_with(index=self.collection_name.lower())
-
 
 @pytest.mark.usefixtures("setup_mock_redis")
 class TestOpenSearchVectorWithRedis:

api/tests/test_containers_integration_tests/core/__init__.py

@@ -0,0 +1 @@
+# Test containers integration tests for core RAG pipeline components

api/tests/test_containers_integration_tests/core/rag/__init__.py

@@ -0,0 +1 @@
+# Test containers integration tests for RAG pipeline

api/tests/test_containers_integration_tests/core/rag/pipeline/__init__.py

@@ -0,0 +1 @@
+# Test containers integration tests for RAG pipeline queue

api/tests/test_containers_integration_tests/core/rag/pipeline/test_queue_integration.py

@@ -0,0 +1,663 @@
+"""
+Integration tests for TenantSelfTaskQueue using testcontainers.
+
+These tests verify the Redis-based task queue functionality with real Redis instances,
+testing tenant isolation, task serialization, and queue operations in a realistic environment.
+Includes compatibility tests for migrating from legacy string-only queues.
+
+All tests use generic naming to avoid coupling to specific business implementations.
+"""
+
+import time
+from dataclasses import dataclass
+from typing import Any
+from uuid import uuid4
+
+import pytest
+from faker import Faker
+
+from core.rag.pipeline.queue import TASK_WRAPPER_PREFIX, TaskWrapper, TenantSelfTaskQueue
+from extensions.ext_redis import redis_client
+from models import Account, Tenant, TenantAccountJoin, TenantAccountRole
+
+
+@dataclass
+class TestTask:
+    """Test task data structure for testing complex object serialization."""
+    task_id: str
+    tenant_id: str
+    data: dict[str, Any]
+    metadata: dict[str, Any]
+
+
+class TestTenantSelfTaskQueueIntegration:
+    """Integration tests for TenantSelfTaskQueue using testcontainers."""
+
+    @pytest.fixture
+    def fake(self):
+        """Faker instance for generating test data."""
+        return Faker()
+
+    @pytest.fixture
+    def test_tenant_and_account(self, db_session_with_containers, fake):
+        """Create test tenant and account for testing."""
+        # Create account
+        account = Account(
+            email=fake.email(),
+            name=fake.name(),
+            interface_language="en-US",
+            status="active",
+        )
+        db_session_with_containers.add(account)
+        db_session_with_containers.commit()
+
+        # Create tenant
+        tenant = Tenant(
+            name=fake.company(),
+            status="normal",
+        )
+        db_session_with_containers.add(tenant)
+        db_session_with_containers.commit()
+
+        # Create tenant-account join
+        join = TenantAccountJoin(
+            tenant_id=tenant.id,
+            account_id=account.id,
+            role=TenantAccountRole.OWNER,
+            current=True,
+        )
+        db_session_with_containers.add(join)
+        db_session_with_containers.commit()
+
+        return tenant, account
+
+    @pytest.fixture
+    def test_queue(self, test_tenant_and_account):
+        """Create a generic test queue for testing."""
+        tenant, _ = test_tenant_and_account
+        return TenantSelfTaskQueue(tenant.id, "test_queue")
+
+    @pytest.fixture
+    def secondary_queue(self, test_tenant_and_account):
+        """Create a secondary test queue for testing isolation."""
+        tenant, _ = test_tenant_and_account
+        return TenantSelfTaskQueue(tenant.id, "secondary_queue")
+
+    def test_queue_initialization(self, test_tenant_and_account):
+        """Test queue initialization with correct key generation."""
+        tenant, _ = test_tenant_and_account
+        queue = TenantSelfTaskQueue(tenant.id, "test-key")
+        
+        assert queue.tenant_id == tenant.id
+        assert queue.unique_key == "test-key"
+        assert queue.queue == f"tenant_self_test-key_task_queue:{tenant.id}"
+        assert queue.task_key == f"tenant_test-key_task:{tenant.id}"
+        assert queue.DEFAULT_TASK_TTL == 60 * 60
+
+    def test_tenant_isolation(self, test_tenant_and_account, db_session_with_containers, fake):
+        """Test that different tenants have isolated queues."""
+        tenant1, _ = test_tenant_and_account
+        
+        # Create second tenant
+        tenant2 = Tenant(
+            name=fake.company(),
+            status="normal",
+        )
+        db_session_with_containers.add(tenant2)
+        db_session_with_containers.commit()
+        
+        queue1 = TenantSelfTaskQueue(tenant1.id, "same-key")
+        queue2 = TenantSelfTaskQueue(tenant2.id, "same-key")
+        
+        assert queue1.queue != queue2.queue
+        assert queue1.task_key != queue2.task_key
+        assert queue1.queue == f"tenant_self_same-key_task_queue:{tenant1.id}"
+        assert queue2.queue == f"tenant_self_same-key_task_queue:{tenant2.id}"
+
+    def test_key_isolation(self, test_tenant_and_account):
+        """Test that different keys have isolated queues."""
+        tenant, _ = test_tenant_and_account
+        queue1 = TenantSelfTaskQueue(tenant.id, "key1")
+        queue2 = TenantSelfTaskQueue(tenant.id, "key2")
+        
+        assert queue1.queue != queue2.queue
+        assert queue1.task_key != queue2.task_key
+        assert queue1.queue == f"tenant_self_key1_task_queue:{tenant.id}"
+        assert queue2.queue == f"tenant_self_key2_task_queue:{tenant.id}"
+
+    def test_task_key_operations(self, test_queue):
+        """Test task key operations (get, set, delete)."""
+        # Initially no task key should exist
+        assert test_queue.get_task_key() is None
+        
+        # Set task waiting time with default TTL
+        test_queue.set_task_waiting_time()
+        task_key = test_queue.get_task_key()
+        # Redis returns bytes, convert to string for comparison
+        assert task_key in (b"1", "1")
+        
+        # Set task waiting time with custom TTL
+        custom_ttl = 30
+        test_queue.set_task_waiting_time(custom_ttl)
+        task_key = test_queue.get_task_key()
+        assert task_key in (b"1", "1")
+        
+        # Delete task key
+        test_queue.delete_task_key()
+        assert test_queue.get_task_key() is None
+
+    def test_push_and_pull_string_tasks(self, test_queue):
+        """Test pushing and pulling string tasks."""
+        tasks = ["task1", "task2", "task3"]
+        
+        # Push tasks
+        test_queue.push_tasks(tasks)
+        
+        # Pull tasks one by one (FIFO order)
+        pulled_tasks = []
+        for _ in range(3):
+            task = test_queue.get_next_task()
+            if task:
+                pulled_tasks.append(task)
+        
+        # Should get tasks in FIFO order (lpush + rpop = FIFO)
+        assert pulled_tasks == ["task1", "task2", "task3"]
+
+    def test_push_and_pull_multiple_tasks(self, test_queue):
+        """Test pushing and pulling multiple tasks at once."""
+        tasks = ["task1", "task2", "task3", "task4", "task5"]
+        
+        # Push tasks
+        test_queue.push_tasks(tasks)
+        
+        # Pull multiple tasks
+        pulled_tasks = test_queue.pull_tasks(3)
+        assert len(pulled_tasks) == 3
+        assert pulled_tasks == ["task1", "task2", "task3"]
+        
+        # Pull remaining tasks
+        remaining_tasks = test_queue.pull_tasks(5)
+        assert len(remaining_tasks) == 2
+        assert remaining_tasks == ["task4", "task5"]
+
+    def test_push_and_pull_complex_objects(self, test_queue, fake):
+        """Test pushing and pulling complex object tasks."""
+        # Create complex task objects as dictionaries (not dataclass instances)
+        tasks = [
+            {
+                "task_id": str(uuid4()),
+                "tenant_id": test_queue.tenant_id,
+                "data": {
+                    "file_id": str(uuid4()),
+                    "content": fake.text(),
+                    "metadata": {"size": fake.random_int(1000, 10000)}
+                },
+                "metadata": {
+                    "created_at": fake.iso8601(),
+                    "tags": fake.words(3)
+                }
+            },
+            {
+                "task_id": str(uuid4()),
+                "tenant_id": test_queue.tenant_id,
+                "data": {
+                    "file_id": str(uuid4()),
+                    "content": "测试中文内容",
+                    "metadata": {"size": fake.random_int(1000, 10000)}
+                },
+                "metadata": {
+                    "created_at": fake.iso8601(),
+                    "tags": ["中文", "测试", "emoji🚀"]
+                }
+            }
+        ]
+        
+        # Push complex tasks
+        test_queue.push_tasks(tasks)
+        
+        # Pull tasks
+        pulled_tasks = test_queue.pull_tasks(2)
+        assert len(pulled_tasks) == 2
+        
+        # Verify deserialized tasks match original (FIFO order)
+        for i, pulled_task in enumerate(pulled_tasks):
+            original_task = tasks[i]  # FIFO order
+            assert isinstance(pulled_task, dict)
+            assert pulled_task["task_id"] == original_task["task_id"]
+            assert pulled_task["tenant_id"] == original_task["tenant_id"]
+            assert pulled_task["data"] == original_task["data"]
+            assert pulled_task["metadata"] == original_task["metadata"]
+
+    def test_mixed_task_types(self, test_queue, fake):
+        """Test pushing and pulling mixed string and object tasks."""
+        string_task = "simple_string_task"
+        object_task = {
+            "task_id": str(uuid4()),
+            "dataset_id": str(uuid4()),
+            "document_ids": [str(uuid4()) for _ in range(3)]
+        }
+        
+        tasks = [string_task, object_task, "another_string"]
+        
+        # Push mixed tasks
+        test_queue.push_tasks(tasks)
+        
+        # Pull all tasks
+        pulled_tasks = test_queue.pull_tasks(3)
+        assert len(pulled_tasks) == 3
+        
+        # Verify types and content
+        assert pulled_tasks[0] == string_task
+        assert isinstance(pulled_tasks[1], dict)
+        assert pulled_tasks[1] == object_task
+        assert pulled_tasks[2] == "another_string"
+
+    def test_empty_queue_operations(self, test_queue):
+        """Test operations on empty queue."""
+        # Pull from empty queue
+        tasks = test_queue.pull_tasks(5)
+        assert tasks == []
+        
+        # Get next task from empty queue
+        task = test_queue.get_next_task()
+        assert task is None
+        
+        # Pull zero or negative count
+        assert test_queue.pull_tasks(0) == []
+        assert test_queue.pull_tasks(-1) == []
+
+    def test_task_ttl_expiration(self, test_queue):
+        """Test task key TTL expiration."""
+        # Set task with short TTL
+        short_ttl = 2
+        test_queue.set_task_waiting_time(short_ttl)
+        
+        # Verify task key exists
+        assert test_queue.get_task_key() == b"1" or test_queue.get_task_key() == "1"
+        
+        # Wait for TTL to expire
+        time.sleep(short_ttl + 1)
+        
+        # Verify task key has expired
+        assert test_queue.get_task_key() is None
+
+    def test_large_task_batch(self, test_queue, fake):
+        """Test handling large batches of tasks."""
+        # Create large batch of tasks
+        large_batch = []
+        for i in range(100):
+            task = {
+                "task_id": str(uuid4()),
+                "index": i,
+                "data": fake.text(max_nb_chars=100),
+                "metadata": {"batch_id": str(uuid4())}
+            }
+            large_batch.append(task)
+        
+        # Push large batch
+        test_queue.push_tasks(large_batch)
+        
+        # Pull all tasks
+        pulled_tasks = test_queue.pull_tasks(100)
+        assert len(pulled_tasks) == 100
+        
+        # Verify all tasks were retrieved correctly (FIFO order)
+        for i, task in enumerate(pulled_tasks):
+            assert isinstance(task, dict)
+            assert task["index"] == i  # FIFO order
+
+    def test_queue_operations_isolation(self, test_tenant_and_account, fake):
+        """Test concurrent operations on different queues."""
+        tenant, _ = test_tenant_and_account
+        
+        # Create multiple queues for the same tenant
+        queue1 = TenantSelfTaskQueue(tenant.id, "queue1")
+        queue2 = TenantSelfTaskQueue(tenant.id, "queue2")
+        
+        # Push tasks to different queues
+        queue1.push_tasks(["task1_queue1", "task2_queue1"])
+        queue2.push_tasks(["task1_queue2", "task2_queue2"])
+        
+        # Verify queues are isolated
+        tasks1 = queue1.pull_tasks(2)
+        tasks2 = queue2.pull_tasks(2)
+        
+        assert tasks1 == ["task1_queue1", "task2_queue1"]
+        assert tasks2 == ["task1_queue2", "task2_queue2"]
+        assert tasks1 != tasks2
+
+    def test_task_wrapper_serialization_roundtrip(self, test_queue, fake):
+        """Test TaskWrapper serialization and deserialization roundtrip."""
+        # Create complex nested data
+        complex_data = {
+            "id": str(uuid4()),
+            "nested": {
+                "deep": {
+                    "value": "test",
+                    "numbers": [1, 2, 3, 4, 5],
+                    "unicode": "测试中文",
+                    "emoji": "🚀"
+                }
+            },
+            "metadata": {
+                "created_at": fake.iso8601(),
+                "tags": ["tag1", "tag2", "tag3"]
+            }
+        }
+        
+        # Create wrapper and serialize
+        wrapper = TaskWrapper(complex_data)
+        serialized = wrapper.serialize()
+        
+        # Verify serialization
+        assert isinstance(serialized, str)
+        assert "测试中文" in serialized
+        assert "🚀" in serialized
+        
+        # Deserialize and verify
+        deserialized_wrapper = TaskWrapper.deserialize(serialized)
+        assert deserialized_wrapper.data == complex_data
+
+    def test_error_handling_invalid_json(self, test_queue):
+        """Test error handling for invalid JSON in wrapped tasks."""
+        # Manually create invalid wrapped task
+        invalid_wrapped_task = f"{TASK_WRAPPER_PREFIX}invalid json data"
+        
+        # Push invalid task directly to Redis
+        redis_client.lpush(test_queue.queue, invalid_wrapped_task)
+        
+        # Pull task - should fall back to string
+        task = test_queue.get_next_task()
+        assert task == invalid_wrapped_task
+
+    def test_real_world_processing_scenario(self, test_queue, fake):
+        """Test realistic task processing scenario."""
+        # Simulate various task types
+        tasks = []
+        for i in range(5):
+            task = {
+                "tenant_id": test_queue.tenant_id,
+                "resource_id": str(uuid4()),
+                "resource_ids": [str(uuid4()) for _ in range(fake.random_int(1, 5))]
+            }
+            tasks.append(task)
+        
+        # Push all tasks
+        test_queue.push_tasks(tasks)
+        
+        # Simulate processing tasks one by one
+        processed_tasks = []
+        while True:
+            task = test_queue.get_next_task()
+            if task is None:
+                break
+            
+            processed_tasks.append(task)
+            
+            # Simulate task processing time
+            time.sleep(0.01)
+        
+        # Verify all tasks were processed
+        assert len(processed_tasks) == 5
+        
+        # Verify task content
+        for task in processed_tasks:
+            assert isinstance(task, dict)
+            assert "tenant_id" in task
+            assert "resource_id" in task
+            assert "resource_ids" in task
+            assert task["tenant_id"] == test_queue.tenant_id
+
+    def test_real_world_batch_processing_scenario(self, test_queue, fake):
+        """Test realistic batch processing scenario."""
+        # Simulate batch processing tasks
+        batch_tasks = []
+        for i in range(3):
+            task = {
+                "file_id": str(uuid4()),
+                "tenant_id": test_queue.tenant_id,
+                "user_id": str(uuid4()),
+                "processing_config": {
+                    "model": fake.random_element(["model_a", "model_b", "model_c"]),
+                    "temperature": fake.random.uniform(0.1, 1.0),
+                    "max_tokens": fake.random_int(1000, 4000)
+                },
+                "metadata": {
+                    "source": fake.random_element(["upload", "api", "webhook"]),
+                    "priority": fake.random_element(["low", "normal", "high"])
+                }
+            }
+            batch_tasks.append(task)
+        
+        # Push tasks
+        test_queue.push_tasks(batch_tasks)
+        
+        # Process tasks in batches
+        batch_size = 2
+        processed_tasks = []
+        
+        while True:
+            batch = test_queue.pull_tasks(batch_size)
+            if not batch:
+                break
+            
+            processed_tasks.extend(batch)
+        
+        # Verify all tasks were processed
+        assert len(processed_tasks) == 3
+        
+        # Verify task structure
+        for task in processed_tasks:
+            assert isinstance(task, dict)
+            assert "file_id" in task
+            assert "tenant_id" in task
+            assert "processing_config" in task
+            assert "metadata" in task
+            assert task["tenant_id"] == test_queue.tenant_id
+
+
+class TestTenantSelfTaskQueueCompatibility:
+    """Compatibility tests for migrating from legacy string-only queues."""
+
+    @pytest.fixture
+    def fake(self):
+        """Faker instance for generating test data."""
+        return Faker()
+
+    @pytest.fixture
+    def test_tenant_and_account(self, db_session_with_containers, fake):
+        """Create test tenant and account for testing."""
+        # Create account
+        account = Account(
+            email=fake.email(),
+            name=fake.name(),
+            interface_language="en-US",
+            status="active",
+        )
+        db_session_with_containers.add(account)
+        db_session_with_containers.commit()
+
+        # Create tenant
+        tenant = Tenant(
+            name=fake.company(),
+            status="normal",
+        )
+        db_session_with_containers.add(tenant)
+        db_session_with_containers.commit()
+
+        # Create tenant-account join
+        join = TenantAccountJoin(
+            tenant_id=tenant.id,
+            account_id=account.id,
+            role=TenantAccountRole.OWNER,
+            current=True,
+        )
+        db_session_with_containers.add(join)
+        db_session_with_containers.commit()
+
+        return tenant, account
+
+    def test_legacy_string_queue_compatibility(self, test_tenant_and_account, fake):
+        """
+        Test compatibility with legacy queues containing only string data.
+        
+        This simulates the scenario where Redis queues already contain string data
+        from the old architecture, and we need to ensure the new code can read them.
+        """
+        tenant, _ = test_tenant_and_account
+        queue = TenantSelfTaskQueue(tenant.id, "legacy_queue")
+        
+        # Simulate legacy string data in Redis queue (using old format)
+        legacy_strings = [
+            "legacy_task_1",
+            "legacy_task_2", 
+            "legacy_task_3",
+            "legacy_task_4",
+            "legacy_task_5"
+        ]
+        
+        # Manually push legacy strings directly to Redis (simulating old system)
+        for legacy_string in legacy_strings:
+            redis_client.lpush(queue.queue, legacy_string)
+        
+        # Verify new code can read legacy string data
+        pulled_tasks = queue.pull_tasks(5)
+        assert len(pulled_tasks) == 5
+        
+        # Verify all tasks are strings (not wrapped)
+        for task in pulled_tasks:
+            assert isinstance(task, str)
+            assert task.startswith("legacy_task_")
+        
+        # Verify order (FIFO from Redis list)
+        expected_order = ["legacy_task_1", "legacy_task_2", "legacy_task_3", "legacy_task_4", "legacy_task_5"]
+        assert pulled_tasks == expected_order
+
+    def test_legacy_queue_migration_scenario(self, test_tenant_and_account, fake):
+        """
+        Test complete migration scenario from legacy to new system.
+        
+        This simulates the real-world scenario where:
+        1. Legacy system has string data in Redis
+        2. New system starts processing the same queue
+        3. Both legacy and new tasks coexist during migration
+        4. New system can handle both formats seamlessly
+        """
+        tenant, _ = test_tenant_and_account
+        queue = TenantSelfTaskQueue(tenant.id, "migration_queue")
+        
+        # Phase 1: Legacy system has data
+        legacy_tasks = [f"legacy_resource_{i}" for i in range(1, 6)]
+        redis_client.lpush(queue.queue, *legacy_tasks)
+        
+        # Phase 2: New system starts processing legacy data
+        processed_legacy = []
+        while True:
+            task = queue.get_next_task()
+            if task is None:
+                break
+            processed_legacy.append(task)
+        
+        # Verify legacy data was processed correctly
+        assert len(processed_legacy) == 5
+        for task in processed_legacy:
+            assert isinstance(task, str)
+            assert task.startswith("legacy_resource_")
+        
+        # Phase 3: New system adds new tasks (mixed types)
+        new_string_tasks = ["new_resource_1", "new_resource_2"]
+        new_object_tasks = [
+            {
+                "resource_id": str(uuid4()),
+                "tenant_id": tenant.id,
+                "processing_type": "new_system",
+                "metadata": {"version": "2.0", "features": ["ai", "ml"]}
+            },
+            {
+                "resource_id": str(uuid4()),
+                "tenant_id": tenant.id,
+                "processing_type": "new_system",
+                "metadata": {"version": "2.0", "features": ["ai", "ml"]}
+            }
+        ]
+        
+        # Push new tasks using new system
+        queue.push_tasks(new_string_tasks)
+        queue.push_tasks(new_object_tasks)
+        
+        # Phase 4: Process all new tasks
+        processed_new = []
+        while True:
+            task = queue.get_next_task()
+            if task is None:
+                break
+            processed_new.append(task)
+        
+        # Verify new tasks were processed correctly
+        assert len(processed_new) == 4
+        
+        string_tasks = [task for task in processed_new if isinstance(task, str)]
+        object_tasks = [task for task in processed_new if isinstance(task, dict)]
+        
+        assert len(string_tasks) == 2
+        assert len(object_tasks) == 2
+        
+        # Verify string tasks
+        for task in string_tasks:
+            assert task.startswith("new_resource_")
+        
+        # Verify object tasks
+        for task in object_tasks:
+            assert isinstance(task, dict)
+            assert "resource_id" in task
+            assert "tenant_id" in task
+            assert task["tenant_id"] == tenant.id
+            assert task["processing_type"] == "new_system"
+
+    def test_legacy_queue_error_recovery(self, test_tenant_and_account, fake):
+        """
+        Test error recovery when legacy queue contains malformed data.
+        
+        This ensures the new system can gracefully handle corrupted or
+        malformed legacy data without crashing.
+        """
+        tenant, _ = test_tenant_and_account
+        queue = TenantSelfTaskQueue(tenant.id, "error_recovery_queue")
+        
+        # Create mix of valid and malformed legacy data
+        mixed_legacy_data = [
+            "valid_legacy_task_1",
+            "valid_legacy_task_2",
+            "malformed_data_without_prefix",  # This should be treated as string
+            "valid_legacy_task_3",
+            f"{TASK_WRAPPER_PREFIX}invalid_json_data",  # This should fall back to string
+            "valid_legacy_task_4"
+        ]
+        
+        # Manually push mixed data directly to Redis
+        redis_client.lpush(queue.queue, *mixed_legacy_data)
+        
+        # Process all tasks
+        processed_tasks = []
+        while True:
+            task = queue.get_next_task()
+            if task is None:
+                break
+            processed_tasks.append(task)
+        
+        # Verify all tasks were processed (no crashes)
+        assert len(processed_tasks) == 6
+        
+        # Verify all tasks are strings (malformed data falls back to string)
+        for task in processed_tasks:
+            assert isinstance(task, str)
+        
+        # Verify valid tasks are preserved
+        valid_tasks = [task for task in processed_tasks if task.startswith("valid_legacy_task_")]
+        assert len(valid_tasks) == 4
+        
+        # Verify malformed data is handled gracefully
+        malformed_tasks = [task for task in processed_tasks if not task.startswith("valid_legacy_task_")]
+        assert len(malformed_tasks) == 2
+        assert "malformed_data_without_prefix" in malformed_tasks
+        assert f"{TASK_WRAPPER_PREFIX}invalid_json_data" in malformed_tasks

api/tests/test_containers_integration_tests/services/tools/test_mcp_tools_manage_service.py

@@ -20,21 +20,12 @@ class TestMCPToolManageService:
             patch("services.tools.mcp_tools_manage_service.ToolTransformService") as mock_tool_transform_service,
         ):
             # Setup default mock returns
-            from core.tools.entities.api_entities import ToolProviderApiEntity
-            from core.tools.entities.common_entities import I18nObject
-
             mock_encrypter.encrypt_token.return_value = "encrypted_server_url"
-            mock_tool_transform_service.mcp_provider_to_user_provider.return_value = ToolProviderApiEntity(
-                id="test_id",
-                author="test_author",
-                name="test_name",
-                type=ToolProviderType.MCP,
-                description=I18nObject(en_US="Test Description", zh_Hans="测试描述"),
-                icon={"type": "emoji", "content": "🤖"},
-                label=I18nObject(en_US="Test Label", zh_Hans="测试标签"),
-                labels=[],
-                tools=[],
-            )
+            mock_tool_transform_service.mcp_provider_to_user_provider.return_value = {
+                "id": "test_id",
+                "name": "test_name",
+                "type": ToolProviderType.MCP,
+            }
 
             yield {
                 "encrypter": mock_encrypter,
@@ -113,9 +104,9 @@ class TestMCPToolManageService:
         mcp_provider = MCPToolProvider(
             tenant_id=tenant_id,
             name=fake.company(),
-            server_identifier=str(fake.uuid4()),
+            server_identifier=fake.uuid4(),
             server_url="encrypted_server_url",
-            server_url_hash=str(fake.sha256()),
+            server_url_hash=fake.sha256(),
             user_id=user_id,
             authed=False,
             tools="[]",
@@ -153,10 +144,7 @@ class TestMCPToolManageService:
         )
 
         # Act: Execute the method under test
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
-        result = service.get_provider(provider_id=mcp_provider.id, tenant_id=tenant.id)
+        result = MCPToolManageService.get_mcp_provider_by_provider_id(mcp_provider.id, tenant.id)
 
         # Assert: Verify the expected outcomes
         assert result is not None
@@ -166,6 +154,8 @@ class TestMCPToolManageService:
         assert result.user_id == account.id
 
         # Verify database state
+        from extensions.ext_database import db
+
         db.session.refresh(result)
         assert result.id is not None
         assert result.server_identifier == mcp_provider.server_identifier
@@ -187,14 +177,11 @@ class TestMCPToolManageService:
             db_session_with_containers, mock_external_service_dependencies
         )
 
-        non_existent_id = str(fake.uuid4())
+        non_existent_id = fake.uuid4()
 
         # Act & Assert: Verify proper error handling
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
         with pytest.raises(ValueError, match="MCP tool not found"):
-            service.get_provider(provider_id=non_existent_id, tenant_id=tenant.id)
+            MCPToolManageService.get_mcp_provider_by_provider_id(non_existent_id, tenant.id)
 
     def test_get_mcp_provider_by_provider_id_tenant_isolation(
         self, db_session_with_containers, mock_external_service_dependencies
@@ -223,11 +210,8 @@ class TestMCPToolManageService:
         )
 
         # Act & Assert: Verify tenant isolation
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
         with pytest.raises(ValueError, match="MCP tool not found"):
-            service.get_provider(provider_id=mcp_provider1.id, tenant_id=tenant2.id)
+            MCPToolManageService.get_mcp_provider_by_provider_id(mcp_provider1.id, tenant2.id)
 
     def test_get_mcp_provider_by_server_identifier_success(
         self, db_session_with_containers, mock_external_service_dependencies
@@ -251,10 +235,7 @@ class TestMCPToolManageService:
         )
 
         # Act: Execute the method under test
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
-        result = service.get_provider(server_identifier=mcp_provider.server_identifier, tenant_id=tenant.id)
+        result = MCPToolManageService.get_mcp_provider_by_server_identifier(mcp_provider.server_identifier, tenant.id)
 
         # Assert: Verify the expected outcomes
         assert result is not None
@@ -264,6 +245,8 @@ class TestMCPToolManageService:
         assert result.user_id == account.id
 
         # Verify database state
+        from extensions.ext_database import db
+
         db.session.refresh(result)
         assert result.id is not None
         assert result.name == mcp_provider.name
@@ -285,14 +268,11 @@ class TestMCPToolManageService:
             db_session_with_containers, mock_external_service_dependencies
         )
 
-        non_existent_identifier = str(fake.uuid4())
+        non_existent_identifier = fake.uuid4()
 
         # Act & Assert: Verify proper error handling
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
         with pytest.raises(ValueError, match="MCP tool not found"):
-            service.get_provider(server_identifier=non_existent_identifier, tenant_id=tenant.id)
+            MCPToolManageService.get_mcp_provider_by_server_identifier(non_existent_identifier, tenant.id)
 
     def test_get_mcp_provider_by_server_identifier_tenant_isolation(
         self, db_session_with_containers, mock_external_service_dependencies
@@ -321,11 +301,8 @@ class TestMCPToolManageService:
         )
 
         # Act & Assert: Verify tenant isolation
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
         with pytest.raises(ValueError, match="MCP tool not found"):
-            service.get_provider(server_identifier=mcp_provider1.server_identifier, tenant_id=tenant2.id)
+            MCPToolManageService.get_mcp_provider_by_server_identifier(mcp_provider1.server_identifier, tenant2.id)
 
     def test_create_mcp_provider_success(self, db_session_with_containers, mock_external_service_dependencies):
         """
@@ -345,30 +322,15 @@ class TestMCPToolManageService:
         )
 
         # Setup mocks for provider creation
-        from core.tools.entities.api_entities import ToolProviderApiEntity
-        from core.tools.entities.common_entities import I18nObject
-
         mock_external_service_dependencies["encrypter"].encrypt_token.return_value = "encrypted_server_url"
-        mock_external_service_dependencies[
-            "tool_transform_service"
-        ].mcp_provider_to_user_provider.return_value = ToolProviderApiEntity(
-            id="new_provider_id",
-            author=account.name,
-            name="Test MCP Provider",
-            type=ToolProviderType.MCP,
-            description=I18nObject(en_US="Test MCP Provider Description", zh_Hans="测试MCP提供者描述"),
-            icon={"type": "emoji", "content": "🤖"},
-            label=I18nObject(en_US="Test MCP Provider", zh_Hans="测试MCP提供者"),
-            labels=[],
-            tools=[],
-        )
+        mock_external_service_dependencies["tool_transform_service"].mcp_provider_to_user_provider.return_value = {
+            "id": "new_provider_id",
+            "name": "Test MCP Provider",
+            "type": ToolProviderType.MCP,
+        }
 
         # Act: Execute the method under test
-        from core.entities.mcp_provider import MCPConfiguration
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
-        result = service.create_provider(
+        result = MCPToolManageService.create_mcp_provider(
             tenant_id=tenant.id,
             name="Test MCP Provider",
             server_url="https://example.com/mcp",
@@ -377,16 +339,14 @@ class TestMCPToolManageService:
             icon_type="emoji",
             icon_background="#FF6B6B",
             server_identifier="test_identifier_123",
-            configuration=MCPConfiguration(
-                timeout=30.0,
-                sse_read_timeout=300.0,
-            ),
+            timeout=30.0,
+            sse_read_timeout=300.0,
         )
 
         # Assert: Verify the expected outcomes
         assert result is not None
-        assert result.name == "Test MCP Provider"
-        assert result.type == ToolProviderType.MCP
+        assert result["name"] == "Test MCP Provider"
+        assert result["type"] == ToolProviderType.MCP
 
         # Verify database state
         from extensions.ext_database import db
@@ -426,11 +386,7 @@ class TestMCPToolManageService:
         )
 
         # Create first provider
-        from core.entities.mcp_provider import MCPConfiguration
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
-        service.create_provider(
+        MCPToolManageService.create_mcp_provider(
             tenant_id=tenant.id,
             name="Test MCP Provider",
             server_url="https://example1.com/mcp",
@@ -439,15 +395,13 @@ class TestMCPToolManageService:
             icon_type="emoji",
             icon_background="#FF6B6B",
             server_identifier="test_identifier_1",
-            configuration=MCPConfiguration(
-                timeout=30.0,
-                sse_read_timeout=300.0,
-            ),
+            timeout=30.0,
+            sse_read_timeout=300.0,
         )
 
         # Act & Assert: Verify proper error handling for duplicate name
         with pytest.raises(ValueError, match="MCP tool Test MCP Provider already exists"):
-            service.create_provider(
+            MCPToolManageService.create_mcp_provider(
                 tenant_id=tenant.id,
                 name="Test MCP Provider",  # Duplicate name
                 server_url="https://example2.com/mcp",
@@ -456,10 +410,8 @@ class TestMCPToolManageService:
                 icon_type="emoji",
                 icon_background="#4ECDC4",
                 server_identifier="test_identifier_2",
-                configuration=MCPConfiguration(
-                    timeout=45.0,
-                    sse_read_timeout=400.0,
-                ),
+                timeout=45.0,
+                sse_read_timeout=400.0,
             )
 
     def test_create_mcp_provider_duplicate_server_url(
@@ -480,11 +432,7 @@ class TestMCPToolManageService:
         )
 
         # Create first provider
-        from core.entities.mcp_provider import MCPConfiguration
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
-        service.create_provider(
+        MCPToolManageService.create_mcp_provider(
             tenant_id=tenant.id,
             name="Test MCP Provider 1",
             server_url="https://example.com/mcp",
@@ -493,15 +441,13 @@ class TestMCPToolManageService:
             icon_type="emoji",
             icon_background="#FF6B6B",
             server_identifier="test_identifier_1",
-            configuration=MCPConfiguration(
-                timeout=30.0,
-                sse_read_timeout=300.0,
-            ),
+            timeout=30.0,
+            sse_read_timeout=300.0,
         )
 
         # Act & Assert: Verify proper error handling for duplicate server URL
-        with pytest.raises(ValueError, match="MCP tool with this server URL already exists"):
-            service.create_provider(
+        with pytest.raises(ValueError, match="MCP tool https://example.com/mcp already exists"):
+            MCPToolManageService.create_mcp_provider(
                 tenant_id=tenant.id,
                 name="Test MCP Provider 2",
                 server_url="https://example.com/mcp",  # Duplicate URL
@@ -510,10 +456,8 @@ class TestMCPToolManageService:
                 icon_type="emoji",
                 icon_background="#4ECDC4",
                 server_identifier="test_identifier_2",
-                configuration=MCPConfiguration(
-                    timeout=45.0,
-                    sse_read_timeout=400.0,
-                ),
+                timeout=45.0,
+                sse_read_timeout=400.0,
             )
 
     def test_create_mcp_provider_duplicate_server_identifier(
@@ -534,11 +478,7 @@ class TestMCPToolManageService:
         )
 
         # Create first provider
-        from core.entities.mcp_provider import MCPConfiguration
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
-        service.create_provider(
+        MCPToolManageService.create_mcp_provider(
             tenant_id=tenant.id,
             name="Test MCP Provider 1",
             server_url="https://example1.com/mcp",
@@ -547,15 +487,13 @@ class TestMCPToolManageService:
             icon_type="emoji",
             icon_background="#FF6B6B",
             server_identifier="test_identifier_123",
-            configuration=MCPConfiguration(
-                timeout=30.0,
-                sse_read_timeout=300.0,
-            ),
+            timeout=30.0,
+            sse_read_timeout=300.0,
         )
 
         # Act & Assert: Verify proper error handling for duplicate server identifier
         with pytest.raises(ValueError, match="MCP tool test_identifier_123 already exists"):
-            service.create_provider(
+            MCPToolManageService.create_mcp_provider(
                 tenant_id=tenant.id,
                 name="Test MCP Provider 2",
                 server_url="https://example2.com/mcp",
@@ -564,10 +502,8 @@ class TestMCPToolManageService:
                 icon_type="emoji",
                 icon_background="#4ECDC4",
                 server_identifier="test_identifier_123",  # Duplicate identifier
-                configuration=MCPConfiguration(
-                    timeout=45.0,
-                    sse_read_timeout=400.0,
-                ),
+                timeout=45.0,
+                sse_read_timeout=400.0,
             )
 
     def test_retrieve_mcp_tools_success(self, db_session_with_containers, mock_external_service_dependencies):
@@ -607,59 +543,23 @@ class TestMCPToolManageService:
         db.session.commit()
 
         # Setup mock for transformation service
-        from core.tools.entities.api_entities import ToolProviderApiEntity
-        from core.tools.entities.common_entities import I18nObject
-
         mock_external_service_dependencies["tool_transform_service"].mcp_provider_to_user_provider.side_effect = [
-            ToolProviderApiEntity(
-                id=provider1.id,
-                author=account.name,
-                name=provider1.name,
-                type=ToolProviderType.MCP,
-                description=I18nObject(en_US="Alpha Provider Description", zh_Hans="Alpha提供者描述"),
-                icon={"type": "emoji", "content": "🅰️"},
-                label=I18nObject(en_US=provider1.name, zh_Hans=provider1.name),
-                labels=[],
-                tools=[],
-            ),
-            ToolProviderApiEntity(
-                id=provider2.id,
-                author=account.name,
-                name=provider2.name,
-                type=ToolProviderType.MCP,
-                description=I18nObject(en_US="Beta Provider Description", zh_Hans="Beta提供者描述"),
-                icon={"type": "emoji", "content": "🅱️"},
-                label=I18nObject(en_US=provider2.name, zh_Hans=provider2.name),
-                labels=[],
-                tools=[],
-            ),
-            ToolProviderApiEntity(
-                id=provider3.id,
-                author=account.name,
-                name=provider3.name,
-                type=ToolProviderType.MCP,
-                description=I18nObject(en_US="Gamma Provider Description", zh_Hans="Gamma提供者描述"),
-                icon={"type": "emoji", "content": "Γ"},
-                label=I18nObject(en_US=provider3.name, zh_Hans=provider3.name),
-                labels=[],
-                tools=[],
-            ),
+            {"id": provider1.id, "name": provider1.name, "type": ToolProviderType.MCP},
+            {"id": provider2.id, "name": provider2.name, "type": ToolProviderType.MCP},
+            {"id": provider3.id, "name": provider3.name, "type": ToolProviderType.MCP},
         ]
 
         # Act: Execute the method under test
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
-        result = service.list_providers(tenant_id=tenant.id, for_list=True)
+        result = MCPToolManageService.retrieve_mcp_tools(tenant.id, for_list=True)
 
         # Assert: Verify the expected outcomes
         assert result is not None
         assert len(result) == 3
 
         # Verify correct ordering by name
-        assert result[0].name == "Alpha Provider"
-        assert result[1].name == "Beta Provider"
-        assert result[2].name == "Gamma Provider"
+        assert result[0]["name"] == "Alpha Provider"
+        assert result[1]["name"] == "Beta Provider"
+        assert result[2]["name"] == "Gamma Provider"
 
         # Verify mock interactions
         assert (
@@ -684,10 +584,7 @@ class TestMCPToolManageService:
         # No MCP providers created for this tenant
 
         # Act: Execute the method under test
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
-        result = service.list_providers(tenant_id=tenant.id, for_list=False)
+        result = MCPToolManageService.retrieve_mcp_tools(tenant.id, for_list=False)
 
         # Assert: Verify the expected outcomes
         assert result is not None
@@ -727,46 +624,20 @@ class TestMCPToolManageService:
         )
 
         # Setup mock for transformation service
-        from core.tools.entities.api_entities import ToolProviderApiEntity
-        from core.tools.entities.common_entities import I18nObject
-
         mock_external_service_dependencies["tool_transform_service"].mcp_provider_to_user_provider.side_effect = [
-            ToolProviderApiEntity(
-                id=provider1.id,
-                author=account1.name,
-                name=provider1.name,
-                type=ToolProviderType.MCP,
-                description=I18nObject(en_US="Provider 1 Description", zh_Hans="提供者1描述"),
-                icon={"type": "emoji", "content": "1️⃣"},
-                label=I18nObject(en_US=provider1.name, zh_Hans=provider1.name),
-                labels=[],
-                tools=[],
-            ),
-            ToolProviderApiEntity(
-                id=provider2.id,
-                author=account2.name,
-                name=provider2.name,
-                type=ToolProviderType.MCP,
-                description=I18nObject(en_US="Provider 2 Description", zh_Hans="提供者2描述"),
-                icon={"type": "emoji", "content": "2️⃣"},
-                label=I18nObject(en_US=provider2.name, zh_Hans=provider2.name),
-                labels=[],
-                tools=[],
-            ),
+            {"id": provider1.id, "name": provider1.name, "type": ToolProviderType.MCP},
+            {"id": provider2.id, "name": provider2.name, "type": ToolProviderType.MCP},
         ]
 
         # Act: Execute the method under test for both tenants
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
-        result1 = service.list_providers(tenant_id=tenant1.id, for_list=True)
-        result2 = service.list_providers(tenant_id=tenant2.id, for_list=True)
+        result1 = MCPToolManageService.retrieve_mcp_tools(tenant1.id, for_list=True)
+        result2 = MCPToolManageService.retrieve_mcp_tools(tenant2.id, for_list=True)
 
         # Assert: Verify tenant isolation
         assert len(result1) == 1
         assert len(result2) == 1
-        assert result1[0].id == provider1.id
-        assert result2[0].id == provider2.id
+        assert result1[0]["id"] == provider1.id
+        assert result2[0]["id"] == provider2.id
 
     def test_list_mcp_tool_from_remote_server_success(
         self, db_session_with_containers, mock_external_service_dependencies
@@ -790,20 +661,17 @@ class TestMCPToolManageService:
         mcp_provider = self._create_test_mcp_provider(
             db_session_with_containers, mock_external_service_dependencies, tenant.id, account.id
         )
-        # Use a valid base64 encoded string to avoid decryption errors
-        import base64
-
-        mcp_provider.server_url = base64.b64encode(b"encrypted_server_url").decode()
-        mcp_provider.authed = True  # Provider must be authenticated to list tools
+        mcp_provider.server_url = "encrypted_server_url"
+        mcp_provider.authed = False
         mcp_provider.tools = "[]"
 
         from extensions.ext_database import db
 
         db.session.commit()
 
-        # Mock the decryption process at the rsa level to avoid key file issues
-        with patch("libs.rsa.decrypt") as mock_decrypt:
-            mock_decrypt.return_value = "https://example.com/mcp"
+        # Mock the decrypted_server_url property to avoid encryption issues
+        with patch("models.tools.encrypter") as mock_encrypter:
+            mock_encrypter.decrypt_token.return_value = "https://example.com/mcp"
 
             # Mock MCPClient and its context manager
             mock_tools = [
@@ -815,16 +683,13 @@ class TestMCPToolManageService:
                 )(),
             ]
 
-            with patch("services.tools.mcp_tools_manage_service.MCPClientWithAuthRetry") as mock_mcp_client:
+            with patch("services.tools.mcp_tools_manage_service.MCPClient") as mock_mcp_client:
                 # Setup mock client
                 mock_client_instance = mock_mcp_client.return_value.__enter__.return_value
                 mock_client_instance.list_tools.return_value = mock_tools
 
                 # Act: Execute the method under test
-                from extensions.ext_database import db
-
-                service = MCPToolManageService(db.session())
-                result = service.list_provider_tools(tenant_id=tenant.id, provider_id=mcp_provider.id)
+                result = MCPToolManageService.list_mcp_tool_from_remote_server(tenant.id, mcp_provider.id)
 
         # Assert: Verify the expected outcomes
         assert result is not None
@@ -840,8 +705,16 @@ class TestMCPToolManageService:
         assert mcp_provider.updated_at is not None
 
         # Verify mock interactions
-        # MCPClientWithAuthRetry is called with different parameters
-        mock_mcp_client.assert_called_once()
+        mock_mcp_client.assert_called_once_with(
+            "https://example.com/mcp",
+            mcp_provider.id,
+            tenant.id,
+            authed=False,
+            for_list=True,
+            headers={},
+            timeout=30.0,
+            sse_read_timeout=300.0,
+        )
 
     def test_list_mcp_tool_from_remote_server_auth_error(
         self, db_session_with_containers, mock_external_service_dependencies
@@ -864,10 +737,7 @@ class TestMCPToolManageService:
         mcp_provider = self._create_test_mcp_provider(
             db_session_with_containers, mock_external_service_dependencies, tenant.id, account.id
         )
-        # Use a valid base64 encoded string to avoid decryption errors
-        import base64
-
-        mcp_provider.server_url = base64.b64encode(b"encrypted_server_url").decode()
+        mcp_provider.server_url = "encrypted_server_url"
         mcp_provider.authed = False
         mcp_provider.tools = "[]"
 
@@ -875,23 +745,20 @@ class TestMCPToolManageService:
 
         db.session.commit()
 
-        # Mock the decryption process at the rsa level to avoid key file issues
-        with patch("libs.rsa.decrypt") as mock_decrypt:
-            mock_decrypt.return_value = "https://example.com/mcp"
+        # Mock the decrypted_server_url property to avoid encryption issues
+        with patch("models.tools.encrypter") as mock_encrypter:
+            mock_encrypter.decrypt_token.return_value = "https://example.com/mcp"
 
             # Mock MCPClient to raise authentication error
-            with patch("services.tools.mcp_tools_manage_service.MCPClientWithAuthRetry") as mock_mcp_client:
+            with patch("services.tools.mcp_tools_manage_service.MCPClient") as mock_mcp_client:
                 from core.mcp.error import MCPAuthError
 
                 mock_client_instance = mock_mcp_client.return_value.__enter__.return_value
                 mock_client_instance.list_tools.side_effect = MCPAuthError("Authentication required")
 
                 # Act & Assert: Verify proper error handling
-                from extensions.ext_database import db
-
-                service = MCPToolManageService(db.session())
                 with pytest.raises(ValueError, match="Please auth the tool first"):
-                    service.list_provider_tools(tenant_id=tenant.id, provider_id=mcp_provider.id)
+                    MCPToolManageService.list_mcp_tool_from_remote_server(tenant.id, mcp_provider.id)
 
         # Verify database state was not changed
         db.session.refresh(mcp_provider)
@@ -919,38 +786,32 @@ class TestMCPToolManageService:
         mcp_provider = self._create_test_mcp_provider(
             db_session_with_containers, mock_external_service_dependencies, tenant.id, account.id
         )
-        # Use a valid base64 encoded string to avoid decryption errors
-        import base64
-
-        mcp_provider.server_url = base64.b64encode(b"encrypted_server_url").decode()
-        mcp_provider.authed = True  # Provider must be authenticated to test connection errors
+        mcp_provider.server_url = "encrypted_server_url"
+        mcp_provider.authed = False
         mcp_provider.tools = "[]"
 
         from extensions.ext_database import db
 
         db.session.commit()
 
-        # Mock the decryption process at the rsa level to avoid key file issues
-        with patch("libs.rsa.decrypt") as mock_decrypt:
-            mock_decrypt.return_value = "https://example.com/mcp"
+        # Mock the decrypted_server_url property to avoid encryption issues
+        with patch("models.tools.encrypter") as mock_encrypter:
+            mock_encrypter.decrypt_token.return_value = "https://example.com/mcp"
 
             # Mock MCPClient to raise connection error
-            with patch("services.tools.mcp_tools_manage_service.MCPClientWithAuthRetry") as mock_mcp_client:
+            with patch("services.tools.mcp_tools_manage_service.MCPClient") as mock_mcp_client:
                 from core.mcp.error import MCPError
 
                 mock_client_instance = mock_mcp_client.return_value.__enter__.return_value
                 mock_client_instance.list_tools.side_effect = MCPError("Connection failed")
 
                 # Act & Assert: Verify proper error handling
-                from extensions.ext_database import db
-
-                service = MCPToolManageService(db.session())
                 with pytest.raises(ValueError, match="Failed to connect to MCP server: Connection failed"):
-                    service.list_provider_tools(tenant_id=tenant.id, provider_id=mcp_provider.id)
+                    MCPToolManageService.list_mcp_tool_from_remote_server(tenant.id, mcp_provider.id)
 
         # Verify database state was not changed
         db.session.refresh(mcp_provider)
-        assert mcp_provider.authed is True  # Provider remains authenticated
+        assert mcp_provider.authed is False
         assert mcp_provider.tools == "[]"
 
     def test_delete_mcp_tool_success(self, db_session_with_containers, mock_external_service_dependencies):
@@ -979,8 +840,7 @@ class TestMCPToolManageService:
         assert db.session.query(MCPToolProvider).filter_by(id=mcp_provider.id).first() is not None
 
         # Act: Execute the method under test
-        service = MCPToolManageService(db.session())
-        service.delete_provider(tenant_id=tenant.id, provider_id=mcp_provider.id)
+        MCPToolManageService.delete_mcp_tool(tenant.id, mcp_provider.id)
 
         # Assert: Verify the expected outcomes
         # Provider should be deleted from database
@@ -1002,14 +862,11 @@ class TestMCPToolManageService:
             db_session_with_containers, mock_external_service_dependencies
         )
 
-        non_existent_id = str(fake.uuid4())
+        non_existent_id = fake.uuid4()
 
         # Act & Assert: Verify proper error handling
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
         with pytest.raises(ValueError, match="MCP tool not found"):
-            service.delete_provider(tenant_id=tenant.id, provider_id=non_existent_id)
+            MCPToolManageService.delete_mcp_tool(tenant.id, non_existent_id)
 
     def test_delete_mcp_tool_tenant_isolation(self, db_session_with_containers, mock_external_service_dependencies):
         """
@@ -1036,11 +893,8 @@ class TestMCPToolManageService:
         )
 
         # Act & Assert: Verify tenant isolation
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
         with pytest.raises(ValueError, match="MCP tool not found"):
-            service.delete_provider(tenant_id=tenant2.id, provider_id=mcp_provider1.id)
+            MCPToolManageService.delete_mcp_tool(tenant2.id, mcp_provider1.id)
 
         # Verify provider still exists in tenant1
         from extensions.ext_database import db
@@ -1075,10 +929,7 @@ class TestMCPToolManageService:
         db.session.commit()
 
         # Act: Execute the method under test
-        from core.entities.mcp_provider import MCPConfiguration
-
-        service = MCPToolManageService(db.session())
-        service.update_provider(
+        MCPToolManageService.update_mcp_provider(
             tenant_id=tenant.id,
             provider_id=mcp_provider.id,
             name="Updated MCP Provider",
@@ -1087,10 +938,8 @@ class TestMCPToolManageService:
             icon_type="emoji",
             icon_background="#4ECDC4",
             server_identifier="updated_identifier_123",
-            configuration=MCPConfiguration(
-                timeout=45.0,
-                sse_read_timeout=400.0,
-            ),
+            timeout=45.0,
+            sse_read_timeout=400.0,
         )
 
         # Assert: Verify the expected outcomes
@@ -1104,10 +953,70 @@ class TestMCPToolManageService:
         # Verify icon was updated
         import json
 
-        icon_data = json.loads(mcp_provider.icon or "{}")
+        icon_data = json.loads(mcp_provider.icon)
         assert icon_data["content"] == "🚀"
         assert icon_data["background"] == "#4ECDC4"
 
+    def test_update_mcp_provider_with_server_url_change(
+        self, db_session_with_containers, mock_external_service_dependencies
+    ):
+        """
+        Test successful update of MCP provider with server URL change.
+
+        This test verifies:
+        - Proper handling of server URL changes
+        - Correct reconnection logic
+        - Database state updates
+        - External service integration
+        """
+        # Arrange: Create test data
+        fake = Faker()
+        account, tenant = self._create_test_account_and_tenant(
+            db_session_with_containers, mock_external_service_dependencies
+        )
+
+        # Create MCP provider
+        mcp_provider = self._create_test_mcp_provider(
+            db_session_with_containers, mock_external_service_dependencies, tenant.id, account.id
+        )
+
+        from extensions.ext_database import db
+
+        db.session.commit()
+
+        # Mock the reconnection method
+        with patch.object(MCPToolManageService, "_re_connect_mcp_provider") as mock_reconnect:
+            mock_reconnect.return_value = {
+                "authed": True,
+                "tools": '[{"name": "test_tool"}]',
+                "encrypted_credentials": "{}",
+            }
+
+            # Act: Execute the method under test
+            MCPToolManageService.update_mcp_provider(
+                tenant_id=tenant.id,
+                provider_id=mcp_provider.id,
+                name="Updated MCP Provider",
+                server_url="https://new-example.com/mcp",
+                icon="🚀",
+                icon_type="emoji",
+                icon_background="#4ECDC4",
+                server_identifier="updated_identifier_123",
+                timeout=45.0,
+                sse_read_timeout=400.0,
+            )
+
+        # Assert: Verify the expected outcomes
+        db.session.refresh(mcp_provider)
+        assert mcp_provider.name == "Updated MCP Provider"
+        assert mcp_provider.server_identifier == "updated_identifier_123"
+        assert mcp_provider.timeout == 45.0
+        assert mcp_provider.sse_read_timeout == 400.0
+        assert mcp_provider.updated_at is not None
+
+        # Verify reconnection was called
+        mock_reconnect.assert_called_once_with("https://new-example.com/mcp", mcp_provider.id, tenant.id)
+
     def test_update_mcp_provider_duplicate_name(self, db_session_with_containers, mock_external_service_dependencies):
         """
         Test error handling when updating MCP provider with duplicate name.
@@ -1139,12 +1048,8 @@ class TestMCPToolManageService:
         db.session.commit()
 
         # Act & Assert: Verify proper error handling for duplicate name
-        from core.entities.mcp_provider import MCPConfiguration
-        from extensions.ext_database import db
-
-        service = MCPToolManageService(db.session())
         with pytest.raises(ValueError, match="MCP tool First Provider already exists"):
-            service.update_provider(
+            MCPToolManageService.update_mcp_provider(
                 tenant_id=tenant.id,
                 provider_id=provider2.id,
                 name="First Provider",  # Duplicate name
@@ -1153,10 +1058,8 @@ class TestMCPToolManageService:
                 icon_type="emoji",
                 icon_background="#4ECDC4",
                 server_identifier="unique_identifier",
-                configuration=MCPConfiguration(
-                    timeout=45.0,
-                    sse_read_timeout=400.0,
-                ),
+                timeout=45.0,
+                sse_read_timeout=400.0,
             )
 
     def test_update_mcp_provider_credentials_success(
@@ -1191,25 +1094,19 @@ class TestMCPToolManageService:
 
         # Mock the provider controller and encryption
         with (
-            patch("core.tools.mcp_tool.provider.MCPToolProviderController") as mock_controller,
-            patch("core.tools.utils.encryption.ProviderConfigEncrypter") as mock_encrypter,
+            patch("services.tools.mcp_tools_manage_service.MCPToolProviderController") as mock_controller,
+            patch("services.tools.mcp_tools_manage_service.ProviderConfigEncrypter") as mock_encrypter,
         ):
             # Setup mocks
-            mock_controller_instance = mock_controller.from_db.return_value
+            mock_controller_instance = mock_controller._from_db.return_value
             mock_controller_instance.get_credentials_schema.return_value = []
 
             mock_encrypter_instance = mock_encrypter.return_value
             mock_encrypter_instance.encrypt.return_value = {"new_key": "encrypted_value"}
 
             # Act: Execute the method under test
-            from extensions.ext_database import db
-
-            service = MCPToolManageService(db.session())
-            service.update_provider_credentials(
-                provider_id=mcp_provider.id,
-                tenant_id=tenant.id,
-                credentials={"new_key": "new_value"},
-                authed=True,
+            MCPToolManageService.update_mcp_provider_credentials(
+                mcp_provider=mcp_provider, credentials={"new_key": "new_value"}, authed=True
             )
 
         # Assert: Verify the expected outcomes
@@ -1220,7 +1117,7 @@ class TestMCPToolManageService:
         # Verify credentials were encrypted and merged
         import json
 
-        credentials = json.loads(mcp_provider.encrypted_credentials or "{}")
+        credentials = json.loads(mcp_provider.encrypted_credentials)
         assert "existing_key" in credentials
         assert "new_key" in credentials
 
@@ -1255,25 +1152,19 @@ class TestMCPToolManageService:
 
         # Mock the provider controller and encryption
         with (
-            patch("core.tools.mcp_tool.provider.MCPToolProviderController") as mock_controller,
-            patch("core.tools.utils.encryption.ProviderConfigEncrypter") as mock_encrypter,
+            patch("services.tools.mcp_tools_manage_service.MCPToolProviderController") as mock_controller,
+            patch("services.tools.mcp_tools_manage_service.ProviderConfigEncrypter") as mock_encrypter,
         ):
             # Setup mocks
-            mock_controller_instance = mock_controller.from_db.return_value
+            mock_controller_instance = mock_controller._from_db.return_value
             mock_controller_instance.get_credentials_schema.return_value = []
 
             mock_encrypter_instance = mock_encrypter.return_value
             mock_encrypter_instance.encrypt.return_value = {"new_key": "encrypted_value"}
 
             # Act: Execute the method under test
-            from extensions.ext_database import db
-
-            service = MCPToolManageService(db.session())
-            service.update_provider_credentials(
-                provider_id=mcp_provider.id,
-                tenant_id=tenant.id,
-                credentials={"new_key": "new_value"},
-                authed=False,
+            MCPToolManageService.update_mcp_provider_credentials(
+                mcp_provider=mcp_provider, credentials={"new_key": "new_value"}, authed=False
             )
 
         # Assert: Verify the expected outcomes
@@ -1308,37 +1199,41 @@ class TestMCPToolManageService:
             type("MockTool", (), {"model_dump": lambda self: {"name": "test_tool_2", "description": "Test tool 2"}})(),
         ]
 
-        with patch("services.tools.mcp_tools_manage_service.MCPClientWithAuthRetry") as mock_mcp_client:
+        with patch("services.tools.mcp_tools_manage_service.MCPClient") as mock_mcp_client:
             # Setup mock client
             mock_client_instance = mock_mcp_client.return_value.__enter__.return_value
             mock_client_instance.list_tools.return_value = mock_tools
 
             # Act: Execute the method under test
-            from extensions.ext_database import db
-
-            service = MCPToolManageService(db.session())
-            result = service._reconnect_provider(
-                server_url="https://example.com/mcp",
-                provider=mcp_provider,
+            result = MCPToolManageService._re_connect_mcp_provider(
+                "https://example.com/mcp", mcp_provider.id, tenant.id
             )
 
         # Assert: Verify the expected outcomes
         assert result is not None
-        assert result.authed is True
-        assert result.tools is not None
-        assert result.encrypted_credentials == "{}"
+        assert result["authed"] is True
+        assert result["tools"] is not None
+        assert result["encrypted_credentials"] == "{}"
 
         # Verify tools were properly serialized
         import json
 
-        tools_data = json.loads(result.tools)
+        tools_data = json.loads(result["tools"])
         assert len(tools_data) == 2
         assert tools_data[0]["name"] == "test_tool_1"
         assert tools_data[1]["name"] == "test_tool_2"
 
         # Verify mock interactions
-        provider_entity = mcp_provider.to_entity()
-        mock_mcp_client.assert_called_once()
+        mock_mcp_client.assert_called_once_with(
+            "https://example.com/mcp",
+            mcp_provider.id,
+            tenant.id,
+            authed=False,
+            for_list=True,
+            headers={},
+            timeout=30.0,
+            sse_read_timeout=300.0,
+        )
 
     def test_re_connect_mcp_provider_auth_error(self, db_session_with_containers, mock_external_service_dependencies):
         """
@@ -1361,26 +1256,22 @@ class TestMCPToolManageService:
         )
 
         # Mock MCPClient to raise authentication error
-        with patch("services.tools.mcp_tools_manage_service.MCPClientWithAuthRetry") as mock_mcp_client:
+        with patch("services.tools.mcp_tools_manage_service.MCPClient") as mock_mcp_client:
             from core.mcp.error import MCPAuthError
 
             mock_client_instance = mock_mcp_client.return_value.__enter__.return_value
             mock_client_instance.list_tools.side_effect = MCPAuthError("Authentication required")
 
             # Act: Execute the method under test
-            from extensions.ext_database import db
-
-            service = MCPToolManageService(db.session())
-            result = service._reconnect_provider(
-                server_url="https://example.com/mcp",
-                provider=mcp_provider,
+            result = MCPToolManageService._re_connect_mcp_provider(
+                "https://example.com/mcp", mcp_provider.id, tenant.id
             )
 
         # Assert: Verify the expected outcomes
         assert result is not None
-        assert result.authed is False
-        assert result.tools == "[]"
-        assert result.encrypted_credentials == "{}"
+        assert result["authed"] is False
+        assert result["tools"] == "[]"
+        assert result["encrypted_credentials"] == "{}"
 
     def test_re_connect_mcp_provider_connection_error(
         self, db_session_with_containers, mock_external_service_dependencies
@@ -1404,18 +1295,12 @@ class TestMCPToolManageService:
         )
 
         # Mock MCPClient to raise connection error
-        with patch("services.tools.mcp_tools_manage_service.MCPClientWithAuthRetry") as mock_mcp_client:
+        with patch("services.tools.mcp_tools_manage_service.MCPClient") as mock_mcp_client:
             from core.mcp.error import MCPError
 
             mock_client_instance = mock_mcp_client.return_value.__enter__.return_value
             mock_client_instance.list_tools.side_effect = MCPError("Connection failed")
 
             # Act & Assert: Verify proper error handling
-            from extensions.ext_database import db
-
-            service = MCPToolManageService(db.session())
             with pytest.raises(ValueError, match="Failed to re-connect MCP server: Connection failed"):
-                service._reconnect_provider(
-                    server_url="https://example.com/mcp",
-                    provider=mcp_provider,
-                )
+                MCPToolManageService._re_connect_mcp_provider("https://example.com/mcp", mcp_provider.id, tenant.id)

api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py

@@ -1,16 +1,32 @@
+from dataclasses import asdict
 from unittest.mock import MagicMock, patch
 
 import pytest
 from faker import Faker
 
+from core.entities.document_task import DocumentTask
 from extensions.ext_database import db
 from models import Account, Tenant, TenantAccountJoin, TenantAccountRole
 from models.dataset import Dataset, Document
-from tasks.document_indexing_task import document_indexing_task
+from tasks.document_indexing_task import (
+    _document_indexing,  # Core function
+    _document_indexing_with_tenant_queue,  # Tenant queue wrapper function
+    document_indexing_task,  # Deprecated old interface
+    normal_document_indexing_task,  # New normal task
+    priority_document_indexing_task,  # New priority task
+)
 
 
-class TestDocumentIndexingTask:
-    """Integration tests for document_indexing_task using testcontainers."""
+class TestDocumentIndexingTasks:
+    """Integration tests for document indexing tasks using testcontainers.
+    
+    This test class covers:
+    - Core _document_indexing function
+    - Deprecated document_indexing_task function
+    - New normal_document_indexing_task function
+    - New priority_document_indexing_task function
+    - Tenant queue wrapper _document_indexing_with_tenant_queue function
+    """
 
     @pytest.fixture
     def mock_external_service_dependencies(self):
@@ -223,7 +239,7 @@ class TestDocumentIndexingTask:
         document_ids = [doc.id for doc in documents]
 
         # Act: Execute the task
-        document_indexing_task(dataset.id, document_ids)
+        _document_indexing(dataset.id, document_ids)
 
         # Assert: Verify the expected outcomes
         # Verify indexing runner was called correctly
@@ -231,10 +247,11 @@ class TestDocumentIndexingTask:
         mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
 
         # Verify documents were updated to parsing status
-        for document in documents:
-            db.session.refresh(document)
-            assert document.indexing_status == "parsing"
-            assert document.processing_started_at is not None
+        # Re-query documents from database since _document_indexing uses a different session
+        for doc_id in document_ids:
+            updated_document = db.session.query(Document).filter(Document.id == doc_id).first()
+            assert updated_document.indexing_status == "parsing"
+            assert updated_document.processing_started_at is not None
 
         # Verify the run method was called with correct documents
         call_args = mock_external_service_dependencies["indexing_runner_instance"].run.call_args
@@ -260,7 +277,7 @@ class TestDocumentIndexingTask:
         document_ids = [fake.uuid4() for _ in range(3)]
 
         # Act: Execute the task with non-existent dataset
-        document_indexing_task(non_existent_dataset_id, document_ids)
+        _document_indexing(non_existent_dataset_id, document_ids)
 
         # Assert: Verify no processing occurred
         mock_external_service_dependencies["indexing_runner"].assert_not_called()
@@ -290,17 +307,18 @@ class TestDocumentIndexingTask:
         all_document_ids = existing_document_ids + non_existent_document_ids
 
         # Act: Execute the task with mixed document IDs
-        document_indexing_task(dataset.id, all_document_ids)
+        _document_indexing(dataset.id, all_document_ids)
 
         # Assert: Verify only existing documents were processed
         mock_external_service_dependencies["indexing_runner"].assert_called_once()
         mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
 
         # Verify only existing documents were updated
-        for document in documents:
-            db.session.refresh(document)
-            assert document.indexing_status == "parsing"
-            assert document.processing_started_at is not None
+        # Re-query documents from database since _document_indexing uses a different session
+        for doc_id in existing_document_ids:
+            updated_document = db.session.query(Document).filter(Document.id == doc_id).first()
+            assert updated_document.indexing_status == "parsing"
+            assert updated_document.processing_started_at is not None
 
         # Verify the run method was called with only existing documents
         call_args = mock_external_service_dependencies["indexing_runner_instance"].run.call_args
@@ -332,7 +350,7 @@ class TestDocumentIndexingTask:
         )
 
         # Act: Execute the task
-        document_indexing_task(dataset.id, document_ids)
+        _document_indexing(dataset.id, document_ids)
 
         # Assert: Verify exception was handled gracefully
         # The task should complete without raising exceptions
@@ -340,10 +358,11 @@ class TestDocumentIndexingTask:
         mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
 
         # Verify documents were still updated to parsing status before the exception
-        for document in documents:
-            db.session.refresh(document)
-            assert document.indexing_status == "parsing"
-            assert document.processing_started_at is not None
+        # Re-query documents from database since _document_indexing close the session
+        for doc_id in document_ids:
+            updated_document = db.session.query(Document).filter(Document.id == doc_id).first()
+            assert updated_document.indexing_status == "parsing"
+            assert updated_document.processing_started_at is not None
 
     def test_document_indexing_task_mixed_document_states(
         self, db_session_with_containers, mock_external_service_dependencies
@@ -406,17 +425,18 @@ class TestDocumentIndexingTask:
         document_ids = [doc.id for doc in all_documents]
 
         # Act: Execute the task with mixed document states
-        document_indexing_task(dataset.id, document_ids)
+        _document_indexing(dataset.id, document_ids)
 
         # Assert: Verify processing
         mock_external_service_dependencies["indexing_runner"].assert_called_once()
         mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
 
         # Verify all documents were updated to parsing status
-        for document in all_documents:
-            db.session.refresh(document)
-            assert document.indexing_status == "parsing"
-            assert document.processing_started_at is not None
+        # Re-query documents from database since _document_indexing uses a different session
+        for doc_id in document_ids:
+            updated_document = db.session.query(Document).filter(Document.id == doc_id).first()
+            assert updated_document.indexing_status == "parsing"
+            assert updated_document.processing_started_at is not None
 
         # Verify the run method was called with all documents
         call_args = mock_external_service_dependencies["indexing_runner_instance"].run.call_args
@@ -469,15 +489,16 @@ class TestDocumentIndexingTask:
         document_ids = [doc.id for doc in all_documents]
 
         # Act: Execute the task with too many documents for sandbox plan
-        document_indexing_task(dataset.id, document_ids)
+        _document_indexing(dataset.id, document_ids)
 
         # Assert: Verify error handling
-        for document in all_documents:
-            db.session.refresh(document)
-            assert document.indexing_status == "error"
-            assert document.error is not None
-            assert "batch upload" in document.error
-            assert document.stopped_at is not None
+        # Re-query documents from database since _document_indexing uses a different session
+        for doc_id in document_ids:
+            updated_document = db.session.query(Document).filter(Document.id == doc_id).first()
+            assert updated_document.indexing_status == "error"
+            assert updated_document.error is not None
+            assert "batch upload" in updated_document.error
+            assert updated_document.stopped_at is not None
 
         # Verify no indexing runner was called
         mock_external_service_dependencies["indexing_runner"].assert_not_called()
@@ -502,17 +523,18 @@ class TestDocumentIndexingTask:
         document_ids = [doc.id for doc in documents]
 
         # Act: Execute the task with billing disabled
-        document_indexing_task(dataset.id, document_ids)
+        _document_indexing(dataset.id, document_ids)
 
         # Assert: Verify successful processing
         mock_external_service_dependencies["indexing_runner"].assert_called_once()
         mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
 
         # Verify documents were updated to parsing status
-        for document in documents:
-            db.session.refresh(document)
-            assert document.indexing_status == "parsing"
-            assert document.processing_started_at is not None
+        # Re-query documents from database since _document_indexing uses a different session
+        for doc_id in document_ids:
+            updated_document = db.session.query(Document).filter(Document.id == doc_id).first()
+            assert updated_document.indexing_status == "parsing"
+            assert updated_document.processing_started_at is not None
 
     def test_document_indexing_task_document_is_paused_error(
         self, db_session_with_containers, mock_external_service_dependencies
@@ -540,7 +562,7 @@ class TestDocumentIndexingTask:
         )
 
         # Act: Execute the task
-        document_indexing_task(dataset.id, document_ids)
+        _document_indexing(dataset.id, document_ids)
 
         # Assert: Verify exception was handled gracefully
         # The task should complete without raising exceptions
@@ -548,7 +570,314 @@ class TestDocumentIndexingTask:
         mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
 
         # Verify documents were still updated to parsing status before the exception
-        for document in documents:
-            db.session.refresh(document)
-            assert document.indexing_status == "parsing"
-            assert document.processing_started_at is not None
+        # Re-query documents from database since _document_indexing uses a different session
+        for doc_id in document_ids:
+            updated_document = db.session.query(Document).filter(Document.id == doc_id).first()
+            assert updated_document.indexing_status == "parsing"
+            assert updated_document.processing_started_at is not None
+
+    # ==================== NEW TESTS FOR REFACTORED FUNCTIONS ====================
+    def test_old_document_indexing_task_success(self, db_session_with_containers, mock_external_service_dependencies):
+        """
+        Test document_indexing_task basic functionality.
+
+        This test verifies:
+        - Task function calls the wrapper correctly
+        - Basic parameter passing works
+        """
+        # Arrange: Create test data
+        dataset, documents = self._create_test_dataset_and_documents(
+            db_session_with_containers, mock_external_service_dependencies, document_count=1
+        )
+        document_ids = [doc.id for doc in documents]
+
+        # Act: Execute the deprecated task (it only takes 2 parameters)
+        document_indexing_task(dataset.id, document_ids)
+
+        # Assert: Verify processing occurred (core logic is tested in _document_indexing tests)
+        mock_external_service_dependencies["indexing_runner"].assert_called_once()
+        mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
+    
+    def test_normal_document_indexing_task_success(
+        self, db_session_with_containers, mock_external_service_dependencies
+    ):
+        """
+        Test normal_document_indexing_task basic functionality.
+
+        This test verifies:
+        - Task function calls the wrapper correctly
+        - Basic parameter passing works
+        """
+        # Arrange: Create test data
+        dataset, documents = self._create_test_dataset_and_documents(
+            db_session_with_containers, mock_external_service_dependencies, document_count=1
+        )
+        document_ids = [doc.id for doc in documents]
+        tenant_id = dataset.tenant_id
+
+        # Act: Execute the new normal task
+        normal_document_indexing_task(tenant_id, dataset.id, document_ids)
+
+        # Assert: Verify processing occurred (core logic is tested in _document_indexing tests)
+        mock_external_service_dependencies["indexing_runner"].assert_called_once()
+        mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
+
+    def test_priority_document_indexing_task_success(
+        self, db_session_with_containers, mock_external_service_dependencies
+    ):
+        """
+        Test priority_document_indexing_task basic functionality.
+
+        This test verifies:
+        - Task function calls the wrapper correctly
+        - Basic parameter passing works
+        """
+        # Arrange: Create test data
+        dataset, documents = self._create_test_dataset_and_documents(
+            db_session_with_containers, mock_external_service_dependencies, document_count=1
+        )
+        document_ids = [doc.id for doc in documents]
+        tenant_id = dataset.tenant_id
+
+        # Act: Execute the new priority task
+        priority_document_indexing_task(tenant_id, dataset.id, document_ids)
+
+        # Assert: Verify processing occurred (core logic is tested in _document_indexing tests)
+        mock_external_service_dependencies["indexing_runner"].assert_called_once()
+        mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
+
+    def test_document_indexing_with_tenant_queue_success(
+        self, db_session_with_containers, mock_external_service_dependencies
+    ):
+        """
+        Test _document_indexing_with_tenant_queue function with no waiting tasks.
+
+        This test verifies:
+        - Core indexing logic execution (same as _document_indexing)
+        - Tenant queue cleanup when no waiting tasks
+        - Task function parameter passing
+        - Queue management after processing
+        """
+        # Arrange: Create test data
+        dataset, documents = self._create_test_dataset_and_documents(
+            db_session_with_containers, mock_external_service_dependencies, document_count=2
+        )
+        document_ids = [doc.id for doc in documents]
+        tenant_id = dataset.tenant_id
+
+        # Mock the task function
+        from unittest.mock import MagicMock
+        mock_task_func = MagicMock()
+
+        # Act: Execute the wrapper function
+        _document_indexing_with_tenant_queue(tenant_id, dataset.id, document_ids, mock_task_func)
+
+        # Assert: Verify core processing occurred (same as _document_indexing)
+        mock_external_service_dependencies["indexing_runner"].assert_called_once()
+        mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
+
+        # Verify documents were updated (same as _document_indexing)
+        # Re-query documents from database since _document_indexing uses a different session
+        for doc_id in document_ids:
+            updated_document = db.session.query(Document).filter(Document.id == doc_id).first()
+            assert updated_document.indexing_status == "parsing"
+            assert updated_document.processing_started_at is not None
+
+        # Verify the run method was called with correct documents
+        call_args = mock_external_service_dependencies["indexing_runner_instance"].run.call_args
+        assert call_args is not None
+        processed_documents = call_args[0][0]
+        assert len(processed_documents) == 2
+
+        # Verify task function was not called (no waiting tasks)
+        mock_task_func.delay.assert_not_called()
+
+    def test_document_indexing_with_tenant_queue_with_waiting_tasks(
+        self, db_session_with_containers, mock_external_service_dependencies
+    ):
+        """
+        Test _document_indexing_with_tenant_queue function with waiting tasks in queue using real Redis.
+
+        This test verifies:
+        - Core indexing logic execution
+        - Real Redis-based tenant queue processing of waiting tasks
+        - Task function calls for waiting tasks
+        - Queue management with multiple tasks using actual Redis operations
+        """
+        # Arrange: Create test data
+        dataset, documents = self._create_test_dataset_and_documents(
+            db_session_with_containers, mock_external_service_dependencies, document_count=1
+        )
+        document_ids = [doc.id for doc in documents]
+        tenant_id = dataset.tenant_id
+        dataset_id = dataset.id
+
+        # Mock the task function
+        from unittest.mock import MagicMock
+        mock_task_func = MagicMock()
+
+        # Use real Redis for TenantSelfTaskQueue
+        from core.rag.pipeline.queue import TenantSelfTaskQueue
+        
+        # Create real queue instance
+        queue = TenantSelfTaskQueue(tenant_id, "document_indexing")
+        
+        # Add waiting tasks to the real Redis queue
+        waiting_tasks = [
+            DocumentTask(tenant_id=tenant_id, dataset_id=dataset.id, document_ids=["waiting-doc-1"]),
+            DocumentTask(tenant_id=tenant_id, dataset_id=dataset.id, document_ids=["waiting-doc-2"])
+        ]
+        # Convert DocumentTask objects to dictionaries for serialization
+        waiting_task_dicts = [asdict(task) for task in waiting_tasks]
+        queue.push_tasks(waiting_task_dicts)
+
+        # Act: Execute the wrapper function
+        _document_indexing_with_tenant_queue(tenant_id, dataset.id, document_ids, mock_task_func)
+
+        # Assert: Verify core processing occurred
+        mock_external_service_dependencies["indexing_runner"].assert_called_once()
+        mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
+
+        # Verify task function was called for each waiting task
+        assert mock_task_func.delay.call_count == 1
+        
+        # Verify correct parameters for each call
+        calls = mock_task_func.delay.call_args_list
+        assert calls[0][1] == {"tenant_id": tenant_id, "dataset_id": dataset_id, "document_ids": ["waiting-doc-1"]}
+
+        # Verify queue is empty after processing (tasks were pulled)
+        remaining_tasks = queue.pull_tasks(count=10)  # Pull more than we added
+        assert len(remaining_tasks) == 1
+
+    def test_document_indexing_with_tenant_queue_error_handling(
+        self, db_session_with_containers, mock_external_service_dependencies
+    ):
+        """
+        Test error handling in _document_indexing_with_tenant_queue using real Redis.
+
+        This test verifies:
+        - Exception handling during core processing
+        - Tenant queue cleanup even on errors using real Redis
+        - Proper error logging
+        - Function completes without raising exceptions
+        - Queue management continues despite core processing errors
+        """
+        # Arrange: Create test data
+        dataset, documents = self._create_test_dataset_and_documents(
+            db_session_with_containers, mock_external_service_dependencies, document_count=1
+        )
+        document_ids = [doc.id for doc in documents]
+        tenant_id = dataset.tenant_id
+        dataset_id = dataset.id
+
+        # Mock IndexingRunner to raise an exception
+        mock_external_service_dependencies["indexing_runner_instance"].run.side_effect = Exception("Test error")
+
+        # Mock the task function
+        from unittest.mock import MagicMock
+        mock_task_func = MagicMock()
+
+        # Use real Redis for TenantSelfTaskQueue
+        from core.rag.pipeline.queue import TenantSelfTaskQueue
+
+        # Create real queue instance
+        queue = TenantSelfTaskQueue(tenant_id, "document_indexing")
+        
+        # Add waiting task to the real Redis queue
+        waiting_task = DocumentTask(tenant_id=tenant_id, dataset_id=dataset.id, document_ids=["waiting-doc-1"])
+        queue.push_tasks([asdict(waiting_task)])
+
+        # Act: Execute the wrapper function
+        _document_indexing_with_tenant_queue(tenant_id, dataset.id, document_ids, mock_task_func)
+
+        # Assert: Verify error was handled gracefully
+        # The function should not raise exceptions
+        mock_external_service_dependencies["indexing_runner"].assert_called_once()
+        mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
+
+        # Verify documents were still updated to parsing status before the exception
+        # Re-query documents from database since _document_indexing uses a different session
+        for doc_id in document_ids:
+            updated_document = db.session.query(Document).filter(Document.id == doc_id).first()
+            assert updated_document.indexing_status == "parsing"
+            assert updated_document.processing_started_at is not None
+
+        # Verify waiting task was still processed despite core processing error
+        mock_task_func.delay.assert_called_once()
+        
+        # Verify correct parameters for the call
+        call = mock_task_func.delay.call_args
+        assert call[1] == {"tenant_id": tenant_id, "dataset_id": dataset_id, "document_ids": ["waiting-doc-1"]}
+
+        # Verify queue is empty after processing (task was pulled)
+        remaining_tasks = queue.pull_tasks(count=10)
+        assert len(remaining_tasks) == 0
+
+    def test_document_indexing_with_tenant_queue_tenant_isolation(
+        self, db_session_with_containers, mock_external_service_dependencies
+    ):
+        """
+        Test tenant isolation in _document_indexing_with_tenant_queue using real Redis.
+
+        This test verifies:
+        - Different tenants have isolated queues
+        - Tasks from one tenant don't affect another tenant's queue
+        - Queue operations are properly scoped to tenant
+        """
+        # Arrange: Create test data for two different tenants
+        dataset1, documents1 = self._create_test_dataset_and_documents(
+            db_session_with_containers, mock_external_service_dependencies, document_count=1
+        )
+        dataset2, documents2 = self._create_test_dataset_and_documents(
+            db_session_with_containers, mock_external_service_dependencies, document_count=1
+        )
+        
+        tenant1_id = dataset1.tenant_id
+        tenant2_id = dataset2.tenant_id
+        dataset1_id = dataset1.id
+        dataset2_id = dataset2.id
+        document_ids1 = [doc.id for doc in documents1]
+        document_ids2 = [doc.id for doc in documents2]
+
+        # Mock the task function
+        from unittest.mock import MagicMock
+        mock_task_func = MagicMock()
+
+        # Use real Redis for TenantSelfTaskQueue
+        from core.rag.pipeline.queue import TenantSelfTaskQueue
+        
+        # Create queue instances for both tenants
+        queue1 = TenantSelfTaskQueue(tenant1_id, "document_indexing")
+        queue2 = TenantSelfTaskQueue(tenant2_id, "document_indexing")
+        
+        # Add waiting tasks to both queues
+        waiting_task1 = DocumentTask(tenant_id=tenant1_id, dataset_id=dataset1.id, document_ids=["tenant1-doc-1"])
+        waiting_task2 = DocumentTask(tenant_id=tenant2_id, dataset_id=dataset2.id, document_ids=["tenant2-doc-1"])
+        
+        queue1.push_tasks([asdict(waiting_task1)])
+        queue2.push_tasks([asdict(waiting_task2)])
+
+        # Act: Execute the wrapper function for tenant1 only
+        _document_indexing_with_tenant_queue(tenant1_id, dataset1.id, document_ids1, mock_task_func)
+
+        # Assert: Verify core processing occurred for tenant1
+        mock_external_service_dependencies["indexing_runner"].assert_called_once()
+        mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
+
+        # Verify only tenant1's waiting task was processed
+        mock_task_func.delay.assert_called_once()
+        call = mock_task_func.delay.call_args
+        assert call[1] == {"tenant_id": tenant1_id, "dataset_id": dataset1_id, "document_ids": ["tenant1-doc-1"]}
+
+        # Verify tenant1's queue is empty
+        remaining_tasks1 = queue1.pull_tasks(count=10)
+        assert len(remaining_tasks1) == 0
+
+        # Verify tenant2's queue still has its task (isolation)
+        remaining_tasks2 = queue2.pull_tasks(count=10)
+        assert len(remaining_tasks2) == 1
+
+        # Verify queue keys are different
+        assert queue1.queue != queue2.queue
+        assert queue1.task_key != queue2.task_key
+

api/tests/test_containers_integration_tests/tasks/test_rag_pipeline_run_tasks.py

@@ -0,0 +1,939 @@
+import json
+import uuid
+from unittest.mock import patch
+
+import pytest
+from faker import Faker
+
+from core.app.entities.app_invoke_entities import InvokeFrom, RagPipelineGenerateEntity
+from core.app.entities.rag_pipeline_invoke_entities import RagPipelineInvokeEntity
+from core.rag.pipeline.queue import TenantSelfTaskQueue
+from extensions.ext_database import db
+from models import Account, Tenant, TenantAccountJoin, TenantAccountRole
+from models.dataset import Pipeline
+from models.workflow import Workflow
+from tasks.rag_pipeline.priority_rag_pipeline_run_task import (
+    priority_rag_pipeline_run_task,
+    run_single_rag_pipeline_task,
+)
+from tasks.rag_pipeline.rag_pipeline_run_task import rag_pipeline_run_task
+
+
+class TestRagPipelineRunTasks:
+    """Integration tests for RAG pipeline run tasks using testcontainers.
+    
+    This test class covers:
+    - priority_rag_pipeline_run_task function
+    - rag_pipeline_run_task function
+    - run_single_rag_pipeline_task function
+    - Real Redis-based TenantSelfTaskQueue operations
+    - PipelineGenerator._generate method mocking and parameter validation
+    - File operations and cleanup
+    - Error handling and queue management
+    """
+
+    @pytest.fixture
+    def mock_pipeline_generator(self):
+        """Mock PipelineGenerator._generate method."""
+        with patch("core.app.apps.pipeline.pipeline_generator.PipelineGenerator._generate") as mock_generate:
+            # Mock the _generate method to return a simple response
+            mock_generate.return_value = {
+                "answer": "Test response",
+                "metadata": {"test": "data"}
+            }
+            yield mock_generate
+
+    @pytest.fixture
+    def mock_file_service(self):
+        """Mock FileService for file operations."""
+        with (
+            patch("services.file_service.FileService.get_file_content") as mock_get_content,
+            patch("services.file_service.FileService.delete_file") as mock_delete_file,
+        ):
+            yield {
+                "get_content": mock_get_content,
+                "delete_file": mock_delete_file,
+            }
+
+    def _create_test_pipeline_and_workflow(self, db_session_with_containers):
+        """
+        Helper method to create test pipeline and workflow for testing.
+
+        Args:
+            db_session_with_containers: Database session from testcontainers infrastructure
+
+        Returns:
+            tuple: (account, tenant, pipeline, workflow) - Created entities
+        """
+        fake = Faker()
+
+        # Create account and tenant
+        account = Account(
+            email=fake.email(),
+            name=fake.name(),
+            interface_language="en-US",
+            status="active",
+        )
+        db.session.add(account)
+        db.session.commit()
+
+        tenant = Tenant(
+            name=fake.company(),
+            status="normal",
+        )
+        db.session.add(tenant)
+        db.session.commit()
+
+        # Create tenant-account join
+        join = TenantAccountJoin(
+            tenant_id=tenant.id,
+            account_id=account.id,
+            role=TenantAccountRole.OWNER,
+            current=True,
+        )
+        db.session.add(join)
+        db.session.commit()
+
+        # Create workflow
+        workflow = Workflow(
+            id=str(uuid.uuid4()),
+            tenant_id=tenant.id,
+            app_id=str(uuid.uuid4()),
+            type="workflow",
+            version="draft",
+            graph="{}",
+            features="{}",
+            marked_name=fake.company(),
+            marked_comment=fake.text(max_nb_chars=100),
+            created_by=account.id,
+            environment_variables=[],
+            conversation_variables=[],
+            rag_pipeline_variables=[],
+        )
+        db.session.add(workflow)
+        db.session.commit()
+
+        # Create pipeline
+        pipeline = Pipeline(
+            id=str(uuid.uuid4()),
+            tenant_id=tenant.id,
+            workflow_id=workflow.id,
+            name=fake.company(),
+            description=fake.text(max_nb_chars=100),
+            created_by=account.id,
+        )
+        db.session.add(pipeline)
+        db.session.commit()
+
+        # Refresh entities to ensure they're properly loaded
+        db.session.refresh(account)
+        db.session.refresh(tenant)
+        db.session.refresh(workflow)
+        db.session.refresh(pipeline)
+
+        return account, tenant, pipeline, workflow
+
+    def _create_rag_pipeline_invoke_entities(self, account, tenant, pipeline, workflow, count=2):
+        """
+        Helper method to create RAG pipeline invoke entities for testing.
+
+        Args:
+            account: Account instance
+            tenant: Tenant instance
+            pipeline: Pipeline instance
+            workflow: Workflow instance
+            count: Number of entities to create
+
+        Returns:
+            list: List of RagPipelineInvokeEntity instances
+        """
+        fake = Faker()
+        entities = []
+
+        for i in range(count):
+            # Create application generate entity
+            app_config = {
+                "app_id": str(uuid.uuid4()),
+                "app_name": fake.company(),
+                "mode": "workflow",
+                "workflow_id": workflow.id,
+                "tenant_id": tenant.id,
+                "app_mode": "workflow",
+            }
+
+            application_generate_entity = {
+                "task_id": str(uuid.uuid4()),
+                "app_config": app_config,
+                "inputs": {"query": f"Test query {i}"},
+                "files": [],
+                "user_id": account.id,
+                "stream": False,
+                "invoke_from": "published",
+                "workflow_execution_id": str(uuid.uuid4()),
+                "pipeline_config": {
+                    "app_id": str(uuid.uuid4()),
+                    "app_name": fake.company(),
+                    "mode": "workflow",
+                    "workflow_id": workflow.id,
+                    "tenant_id": tenant.id,
+                    "app_mode": "workflow",
+                },
+                "datasource_type": "upload_file",
+                "datasource_info": {},
+                "dataset_id": str(uuid.uuid4()),
+                "batch": "test_batch",
+            }
+
+            entity = RagPipelineInvokeEntity(
+                pipeline_id=pipeline.id,
+                application_generate_entity=application_generate_entity,
+                user_id=account.id,
+                tenant_id=tenant.id,
+                workflow_id=workflow.id,
+                streaming=False,
+                workflow_execution_id=str(uuid.uuid4()),
+                workflow_thread_pool_id=str(uuid.uuid4()),
+            )
+            entities.append(entity)
+
+        return entities
+
+    def _create_file_content_for_entities(self, entities):
+        """
+        Helper method to create file content for RAG pipeline invoke entities.
+
+        Args:
+            entities: List of RagPipelineInvokeEntity instances
+
+        Returns:
+            str: JSON string containing serialized entities
+        """
+        entities_data = [entity.model_dump() for entity in entities]
+        return json.dumps(entities_data)
+
+    def test_priority_rag_pipeline_run_task_success(
+        self, db_session_with_containers, mock_pipeline_generator, mock_file_service
+    ):
+        """
+        Test successful priority RAG pipeline run task execution.
+
+        This test verifies:
+        - Task execution with multiple RAG pipeline invoke entities
+        - File content retrieval and parsing
+        - PipelineGenerator._generate method calls with correct parameters
+        - Thread pool execution
+        - File cleanup after execution
+        - Queue management with no waiting tasks
+        """
+        # Arrange: Create test data
+        account, tenant, pipeline, workflow = self._create_test_pipeline_and_workflow(db_session_with_containers)
+        entities = self._create_rag_pipeline_invoke_entities(account, tenant, pipeline, workflow, count=2)
+        file_content = self._create_file_content_for_entities(entities)
+        
+        # Mock file service
+        file_id = str(uuid.uuid4())
+        mock_file_service["get_content"].return_value = file_content
+
+        # Act: Execute the priority task
+        priority_rag_pipeline_run_task(file_id, tenant.id)
+
+        # Assert: Verify expected outcomes
+        # Verify file operations
+        mock_file_service["get_content"].assert_called_once_with(file_id)
+        mock_file_service["delete_file"].assert_called_once_with(file_id)
+
+        # Verify PipelineGenerator._generate was called for each entity
+        assert mock_pipeline_generator.call_count == 2
+
+        # Verify call parameters for each entity
+        calls = mock_pipeline_generator.call_args_list
+        for call in calls:
+            call_kwargs = call[1]  # Get keyword arguments
+            assert call_kwargs["pipeline"].id == pipeline.id
+            assert call_kwargs["workflow_id"] == workflow.id
+            assert call_kwargs["user"].id == account.id
+            assert call_kwargs["invoke_from"] == InvokeFrom.PUBLISHED
+            assert call_kwargs["streaming"] == False
+            assert isinstance(call_kwargs["application_generate_entity"], RagPipelineGenerateEntity)
+
+    def test_rag_pipeline_run_task_success(
+        self, db_session_with_containers, mock_pipeline_generator, mock_file_service
+    ):
+        """
+        Test successful regular RAG pipeline run task execution.
+
+        This test verifies:
+        - Task execution with multiple RAG pipeline invoke entities
+        - File content retrieval and parsing
+        - PipelineGenerator._generate method calls with correct parameters
+        - Thread pool execution
+        - File cleanup after execution
+        - Queue management with no waiting tasks
+        """
+        # Arrange: Create test data
+        account, tenant, pipeline, workflow = self._create_test_pipeline_and_workflow(db_session_with_containers)
+        entities = self._create_rag_pipeline_invoke_entities(account, tenant, pipeline, workflow, count=3)
+        file_content = self._create_file_content_for_entities(entities)
+        
+        # Mock file service
+        file_id = str(uuid.uuid4())
+        mock_file_service["get_content"].return_value = file_content
+
+        # Act: Execute the regular task
+        rag_pipeline_run_task(file_id, tenant.id)
+
+        # Assert: Verify expected outcomes
+        # Verify file operations
+        mock_file_service["get_content"].assert_called_once_with(file_id)
+        mock_file_service["delete_file"].assert_called_once_with(file_id)
+
+        # Verify PipelineGenerator._generate was called for each entity
+        assert mock_pipeline_generator.call_count == 3
+
+        # Verify call parameters for each entity
+        calls = mock_pipeline_generator.call_args_list
+        for call in calls:
+            call_kwargs = call[1]  # Get keyword arguments
+            assert call_kwargs["pipeline"].id == pipeline.id
+            assert call_kwargs["workflow_id"] == workflow.id
+            assert call_kwargs["user"].id == account.id
+            assert call_kwargs["invoke_from"] == InvokeFrom.PUBLISHED
+            assert call_kwargs["streaming"] == False
+            assert isinstance(call_kwargs["application_generate_entity"], RagPipelineGenerateEntity)
+
+    def test_priority_rag_pipeline_run_task_with_waiting_tasks(
+        self, db_session_with_containers, mock_pipeline_generator, mock_file_service
+    ):
+        """
+        Test priority RAG pipeline run task with waiting tasks in queue using real Redis.
+
+        This test verifies:
+        - Core task execution
+        - Real Redis-based tenant queue processing of waiting tasks
+        - Task function calls for waiting tasks
+        - Queue management with multiple tasks using actual Redis operations
+        """
+        # Arrange: Create test data
+        account, tenant, pipeline, workflow = self._create_test_pipeline_and_workflow(db_session_with_containers)
+        entities = self._create_rag_pipeline_invoke_entities(account, tenant, pipeline, workflow, count=1)
+        file_content = self._create_file_content_for_entities(entities)
+        
+        # Mock file service
+        file_id = str(uuid.uuid4())
+        mock_file_service["get_content"].return_value = file_content
+
+        # Use real Redis for TenantSelfTaskQueue
+        queue = TenantSelfTaskQueue(tenant.id, "pipeline")
+        
+        # Add waiting tasks to the real Redis queue
+        waiting_file_ids = [str(uuid.uuid4()) for _ in range(2)]
+        queue.push_tasks(waiting_file_ids)
+
+        # Mock the task function calls
+        with patch(
+            "tasks.rag_pipeline.priority_rag_pipeline_run_task.priority_rag_pipeline_run_task.delay"
+        ) as mock_delay:
+            # Act: Execute the priority task
+            priority_rag_pipeline_run_task(file_id, tenant.id)
+
+            # Assert: Verify core processing occurred
+            mock_file_service["get_content"].assert_called_once_with(file_id)
+            mock_file_service["delete_file"].assert_called_once_with(file_id)
+            assert mock_pipeline_generator.call_count == 1
+
+            # Verify waiting tasks were processed, pull 1 task a time by default
+            assert mock_delay.call_count == 1
+            
+            # Verify correct parameters for the call
+            call_kwargs = mock_delay.call_args[1] if mock_delay.call_args else {}
+            assert call_kwargs.get('rag_pipeline_invoke_entities_file_id') == waiting_file_ids[0]
+            assert call_kwargs.get('tenant_id') == tenant.id
+
+            # Verify queue still has remaining tasks (only 1 was pulled)
+            remaining_tasks = queue.pull_tasks(count=10)
+            assert len(remaining_tasks) == 1  # 2 original - 1 pulled = 1 remaining
+
+    def test_rag_pipeline_run_task_legacy_compatibility(
+        self, db_session_with_containers, mock_pipeline_generator, mock_file_service
+    ):
+        """
+        Test regular RAG pipeline run task with legacy Redis queue format for backward compatibility.
+
+        This test simulates the scenario where:
+        - Old code writes file IDs directly to Redis list using lpush
+        - New worker processes these legacy queue entries
+        - Ensures backward compatibility during deployment transition
+
+        Legacy format: redis_client.lpush(tenant_self_pipeline_task_queue, upload_file.id)
+        New format: TenantSelfTaskQueue.push_tasks([file_id])
+        """
+        # Arrange: Create test data
+        account, tenant, pipeline, workflow = self._create_test_pipeline_and_workflow(db_session_with_containers)
+        entities = self._create_rag_pipeline_invoke_entities(account, tenant, pipeline, workflow, count=1)
+        file_content = self._create_file_content_for_entities(entities)
+        
+        # Mock file service
+        file_id = str(uuid.uuid4())
+        mock_file_service["get_content"].return_value = file_content
+
+        # Simulate legacy Redis queue format - direct file IDs in Redis list
+        from extensions.ext_redis import redis_client
+        
+        # Legacy queue key format (old code)
+        legacy_queue_key = f"tenant_self_pipeline_task_queue:{tenant.id}"
+        legacy_task_key = f"tenant_pipeline_task:{tenant.id}"
+        
+        # Add legacy format data to Redis (simulating old code behavior)
+        legacy_file_ids = [str(uuid.uuid4()) for _ in range(3)]
+        for file_id_legacy in legacy_file_ids:
+            redis_client.lpush(legacy_queue_key, file_id_legacy)
+        
+        # Set the task key to indicate there are waiting tasks (legacy behavior)
+        redis_client.set(legacy_task_key, 1, ex=60 * 60)
+
+        # Mock the task function calls
+        with patch("tasks.rag_pipeline.rag_pipeline_run_task.rag_pipeline_run_task.delay") as mock_delay:
+            # Act: Execute the priority task with new code but legacy queue data
+            rag_pipeline_run_task(file_id, tenant.id)
+
+            # Assert: Verify core processing occurred
+            mock_file_service["get_content"].assert_called_once_with(file_id)
+            mock_file_service["delete_file"].assert_called_once_with(file_id)
+            assert mock_pipeline_generator.call_count == 1
+
+            # Verify waiting tasks were processed, pull 1 task a time by default
+            assert mock_delay.call_count == 1
+
+            # Verify correct parameters for the call
+            call_kwargs = mock_delay.call_args[1] if mock_delay.call_args else {}
+            assert call_kwargs.get('rag_pipeline_invoke_entities_file_id') == legacy_file_ids[0]
+            assert call_kwargs.get('tenant_id') == tenant.id
+
+            # Verify that new code can process legacy queue entries
+            # The new TenantSelfTaskQueue should be able to read from the legacy format
+            queue = TenantSelfTaskQueue(tenant.id, "pipeline")
+
+            # Verify queue still has remaining tasks (only 1 was pulled)
+            remaining_tasks = queue.pull_tasks(count=10)
+            assert len(remaining_tasks) == 2  # 3 original - 1 pulled = 2 remaining
+
+        # Cleanup: Remove legacy test data
+        redis_client.delete(legacy_queue_key)
+        redis_client.delete(legacy_task_key)
+
+    def test_rag_pipeline_run_task_with_waiting_tasks(
+        self, db_session_with_containers, mock_pipeline_generator, mock_file_service
+    ):
+        """
+        Test regular RAG pipeline run task with waiting tasks in queue using real Redis.
+
+        This test verifies:
+        - Core task execution
+        - Real Redis-based tenant queue processing of waiting tasks
+        - Task function calls for waiting tasks
+        - Queue management with multiple tasks using actual Redis operations
+        """
+        # Arrange: Create test data
+        account, tenant, pipeline, workflow = self._create_test_pipeline_and_workflow(db_session_with_containers)
+        entities = self._create_rag_pipeline_invoke_entities(account, tenant, pipeline, workflow, count=1)
+        file_content = self._create_file_content_for_entities(entities)
+        
+        # Mock file service
+        file_id = str(uuid.uuid4())
+        mock_file_service["get_content"].return_value = file_content
+
+        # Use real Redis for TenantSelfTaskQueue
+        queue = TenantSelfTaskQueue(tenant.id, "pipeline")
+        
+        # Add waiting tasks to the real Redis queue
+        waiting_file_ids = [str(uuid.uuid4()) for _ in range(3)]
+        queue.push_tasks(waiting_file_ids)
+
+        # Mock the task function calls
+        with patch("tasks.rag_pipeline.rag_pipeline_run_task.rag_pipeline_run_task.delay") as mock_delay:
+            # Act: Execute the regular task
+            rag_pipeline_run_task(file_id, tenant.id)
+
+            # Assert: Verify core processing occurred
+            mock_file_service["get_content"].assert_called_once_with(file_id)
+            mock_file_service["delete_file"].assert_called_once_with(file_id)
+            assert mock_pipeline_generator.call_count == 1
+
+            # Verify waiting tasks were processed, pull 1 task a time by default
+            assert mock_delay.call_count == 1
+            
+            # Verify correct parameters for the call
+            call_kwargs = mock_delay.call_args[1] if mock_delay.call_args else {}
+            assert call_kwargs.get('rag_pipeline_invoke_entities_file_id') == waiting_file_ids[0]
+            assert call_kwargs.get('tenant_id') == tenant.id
+
+            # Verify queue still has remaining tasks (only 1 was pulled)
+            remaining_tasks = queue.pull_tasks(count=10)
+            assert len(remaining_tasks) == 2  # 3 original - 1 pulled = 2 remaining
+
+    def test_priority_rag_pipeline_run_task_error_handling(
+        self, db_session_with_containers, mock_pipeline_generator, mock_file_service
+    ):
+        """
+        Test error handling in priority RAG pipeline run task using real Redis.
+
+        This test verifies:
+        - Exception handling during core processing
+        - Tenant queue cleanup even on errors using real Redis
+        - Proper error logging
+        - Function completes without raising exceptions
+        - Queue management continues despite core processing errors
+        """
+        # Arrange: Create test data
+        account, tenant, pipeline, workflow = self._create_test_pipeline_and_workflow(db_session_with_containers)
+        entities = self._create_rag_pipeline_invoke_entities(account, tenant, pipeline, workflow, count=1)
+        file_content = self._create_file_content_for_entities(entities)
+        
+        # Mock file service
+        file_id = str(uuid.uuid4())
+        mock_file_service["get_content"].return_value = file_content
+
+        # Mock PipelineGenerator to raise an exception
+        mock_pipeline_generator.side_effect = Exception("Pipeline generation failed")
+
+        # Use real Redis for TenantSelfTaskQueue
+        queue = TenantSelfTaskQueue(tenant.id, "pipeline")
+        
+        # Add waiting task to the real Redis queue
+        waiting_file_id = str(uuid.uuid4())
+        queue.push_tasks([waiting_file_id])
+
+        # Mock the task function calls
+        with patch(
+            "tasks.rag_pipeline.priority_rag_pipeline_run_task.priority_rag_pipeline_run_task.delay"
+        ) as mock_delay:
+            # Act: Execute the priority task (should not raise exception)
+            priority_rag_pipeline_run_task(file_id, tenant.id)
+
+            # Assert: Verify error was handled gracefully
+            # The function should not raise exceptions
+            mock_file_service["get_content"].assert_called_once_with(file_id)
+            mock_file_service["delete_file"].assert_called_once_with(file_id)
+            assert mock_pipeline_generator.call_count == 1
+
+            # Verify waiting task was still processed despite core processing error
+            mock_delay.assert_called_once()
+            
+            # Verify correct parameters for the call
+            call_kwargs = mock_delay.call_args[1] if mock_delay.call_args else {}
+            assert call_kwargs.get('rag_pipeline_invoke_entities_file_id') == waiting_file_id
+            assert call_kwargs.get('tenant_id') == tenant.id
+
+            # Verify queue is empty after processing (task was pulled)
+            remaining_tasks = queue.pull_tasks(count=10)
+            assert len(remaining_tasks) == 0
+
+    def test_rag_pipeline_run_task_error_handling(
+        self, db_session_with_containers, mock_pipeline_generator, mock_file_service
+    ):
+        """
+        Test error handling in regular RAG pipeline run task using real Redis.
+
+        This test verifies:
+        - Exception handling during core processing
+        - Tenant queue cleanup even on errors using real Redis
+        - Proper error logging
+        - Function completes without raising exceptions
+        - Queue management continues despite core processing errors
+        """
+        # Arrange: Create test data
+        account, tenant, pipeline, workflow = self._create_test_pipeline_and_workflow(db_session_with_containers)
+        entities = self._create_rag_pipeline_invoke_entities(account, tenant, pipeline, workflow, count=1)
+        file_content = self._create_file_content_for_entities(entities)
+        
+        # Mock file service
+        file_id = str(uuid.uuid4())
+        mock_file_service["get_content"].return_value = file_content
+
+        # Mock PipelineGenerator to raise an exception
+        mock_pipeline_generator.side_effect = Exception("Pipeline generation failed")
+
+        # Use real Redis for TenantSelfTaskQueue
+        queue = TenantSelfTaskQueue(tenant.id, "pipeline")
+        
+        # Add waiting task to the real Redis queue
+        waiting_file_id = str(uuid.uuid4())
+        queue.push_tasks([waiting_file_id])
+
+        # Mock the task function calls
+        with patch("tasks.rag_pipeline.rag_pipeline_run_task.rag_pipeline_run_task.delay") as mock_delay:
+            # Act: Execute the regular task (should not raise exception)
+            rag_pipeline_run_task(file_id, tenant.id)
+
+            # Assert: Verify error was handled gracefully
+            # The function should not raise exceptions
+            mock_file_service["get_content"].assert_called_once_with(file_id)
+            mock_file_service["delete_file"].assert_called_once_with(file_id)
+            assert mock_pipeline_generator.call_count == 1
+
+            # Verify waiting task was still processed despite core processing error
+            mock_delay.assert_called_once()
+            
+            # Verify correct parameters for the call
+            call_kwargs = mock_delay.call_args[1] if mock_delay.call_args else {}
+            assert call_kwargs.get('rag_pipeline_invoke_entities_file_id') == waiting_file_id
+            assert call_kwargs.get('tenant_id') == tenant.id
+
+            # Verify queue is empty after processing (task was pulled)
+            remaining_tasks = queue.pull_tasks(count=10)
+            assert len(remaining_tasks) == 0
+
+    def test_priority_rag_pipeline_run_task_tenant_isolation(
+        self, db_session_with_containers, mock_pipeline_generator, mock_file_service
+    ):
+        """
+        Test tenant isolation in priority RAG pipeline run task using real Redis.
+
+        This test verifies:
+        - Different tenants have isolated queues
+        - Tasks from one tenant don't affect another tenant's queue
+        - Queue operations are properly scoped to tenant
+        """
+        # Arrange: Create test data for two different tenants
+        account1, tenant1, pipeline1, workflow1 = self._create_test_pipeline_and_workflow(db_session_with_containers)
+        account2, tenant2, pipeline2, workflow2 = self._create_test_pipeline_and_workflow(db_session_with_containers)
+        
+        entities1 = self._create_rag_pipeline_invoke_entities(account1, tenant1, pipeline1, workflow1, count=1)
+        entities2 = self._create_rag_pipeline_invoke_entities(account2, tenant2, pipeline2, workflow2, count=1)
+        
+        file_content1 = self._create_file_content_for_entities(entities1)
+        file_content2 = self._create_file_content_for_entities(entities2)
+        
+        # Mock file service
+        file_id1 = str(uuid.uuid4())
+        file_id2 = str(uuid.uuid4())
+        mock_file_service["get_content"].side_effect = [file_content1, file_content2]
+
+        # Use real Redis for TenantSelfTaskQueue
+        queue1 = TenantSelfTaskQueue(tenant1.id, "pipeline")
+        queue2 = TenantSelfTaskQueue(tenant2.id, "pipeline")
+        
+        # Add waiting tasks to both queues
+        waiting_file_id1 = str(uuid.uuid4())
+        waiting_file_id2 = str(uuid.uuid4())
+        
+        queue1.push_tasks([waiting_file_id1])
+        queue2.push_tasks([waiting_file_id2])
+
+        # Mock the task function calls
+        with patch(
+            "tasks.rag_pipeline.priority_rag_pipeline_run_task.priority_rag_pipeline_run_task.delay"
+        ) as mock_delay:
+            # Act: Execute the priority task for tenant1 only
+            priority_rag_pipeline_run_task(file_id1, tenant1.id)
+
+            # Assert: Verify core processing occurred for tenant1
+            assert mock_file_service["get_content"].call_count == 1
+            assert mock_file_service["delete_file"].call_count == 1
+            assert mock_pipeline_generator.call_count == 1
+
+            # Verify only tenant1's waiting task was processed
+            mock_delay.assert_called_once()
+            call_kwargs = mock_delay.call_args[1] if mock_delay.call_args else {}
+            assert call_kwargs.get('rag_pipeline_invoke_entities_file_id') == waiting_file_id1
+            assert call_kwargs.get('tenant_id') == tenant1.id
+
+            # Verify tenant1's queue is empty
+            remaining_tasks1 = queue1.pull_tasks(count=10)
+            assert len(remaining_tasks1) == 0
+
+            # Verify tenant2's queue still has its task (isolation)
+            remaining_tasks2 = queue2.pull_tasks(count=10)
+            assert len(remaining_tasks2) == 1
+
+            # Verify queue keys are different
+            assert queue1.queue != queue2.queue
+            assert queue1.task_key != queue2.task_key
+
+    def test_rag_pipeline_run_task_tenant_isolation(
+        self, db_session_with_containers, mock_pipeline_generator, mock_file_service
+    ):
+        """
+        Test tenant isolation in regular RAG pipeline run task using real Redis.
+
+        This test verifies:
+        - Different tenants have isolated queues
+        - Tasks from one tenant don't affect another tenant's queue
+        - Queue operations are properly scoped to tenant
+        """
+        # Arrange: Create test data for two different tenants
+        account1, tenant1, pipeline1, workflow1 = self._create_test_pipeline_and_workflow(db_session_with_containers)
+        account2, tenant2, pipeline2, workflow2 = self._create_test_pipeline_and_workflow(db_session_with_containers)
+        
+        entities1 = self._create_rag_pipeline_invoke_entities(account1, tenant1, pipeline1, workflow1, count=1)
+        entities2 = self._create_rag_pipeline_invoke_entities(account2, tenant2, pipeline2, workflow2, count=1)
+        
+        file_content1 = self._create_file_content_for_entities(entities1)
+        file_content2 = self._create_file_content_for_entities(entities2)
+        
+        # Mock file service
+        file_id1 = str(uuid.uuid4())
+        file_id2 = str(uuid.uuid4())
+        mock_file_service["get_content"].side_effect = [file_content1, file_content2]
+
+        # Use real Redis for TenantSelfTaskQueue
+        queue1 = TenantSelfTaskQueue(tenant1.id, "pipeline")
+        queue2 = TenantSelfTaskQueue(tenant2.id, "pipeline")
+        
+        # Add waiting tasks to both queues
+        waiting_file_id1 = str(uuid.uuid4())
+        waiting_file_id2 = str(uuid.uuid4())
+        
+        queue1.push_tasks([waiting_file_id1])
+        queue2.push_tasks([waiting_file_id2])
+
+        # Mock the task function calls
+        with patch("tasks.rag_pipeline.rag_pipeline_run_task.rag_pipeline_run_task.delay") as mock_delay:
+            # Act: Execute the regular task for tenant1 only
+            rag_pipeline_run_task(file_id1, tenant1.id)
+
+            # Assert: Verify core processing occurred for tenant1
+            assert mock_file_service["get_content"].call_count == 1
+            assert mock_file_service["delete_file"].call_count == 1
+            assert mock_pipeline_generator.call_count == 1
+
+            # Verify only tenant1's waiting task was processed
+            mock_delay.assert_called_once()
+            call_kwargs = mock_delay.call_args[1] if mock_delay.call_args else {}
+            assert call_kwargs.get('rag_pipeline_invoke_entities_file_id') == waiting_file_id1
+            assert call_kwargs.get('tenant_id') == tenant1.id
+
+            # Verify tenant1's queue is empty
+            remaining_tasks1 = queue1.pull_tasks(count=10)
+            assert len(remaining_tasks1) == 0
+
+            # Verify tenant2's queue still has its task (isolation)
+            remaining_tasks2 = queue2.pull_tasks(count=10)
+            assert len(remaining_tasks2) == 1
+
+            # Verify queue keys are different
+            assert queue1.queue != queue2.queue
+            assert queue1.task_key != queue2.task_key
+
+    def test_run_single_rag_pipeline_task_success(
+        self, db_session_with_containers, mock_pipeline_generator, flask_app_with_containers
+    ):
+        """
+        Test successful run_single_rag_pipeline_task execution.
+
+        This test verifies:
+        - Single RAG pipeline task execution within Flask app context
+        - Entity validation and database queries
+        - PipelineGenerator._generate method call with correct parameters
+        - Proper Flask context handling
+        """
+        # Arrange: Create test data
+        account, tenant, pipeline, workflow = self._create_test_pipeline_and_workflow(db_session_with_containers)
+        entities = self._create_rag_pipeline_invoke_entities(account, tenant, pipeline, workflow, count=1)
+        entity_data = entities[0].model_dump()
+
+        # Act: Execute the single task
+        with flask_app_with_containers.app_context():
+            run_single_rag_pipeline_task(entity_data, flask_app_with_containers)
+
+        # Assert: Verify expected outcomes
+        # Verify PipelineGenerator._generate was called
+        assert mock_pipeline_generator.call_count == 1
+
+        # Verify call parameters
+        call = mock_pipeline_generator.call_args
+        call_kwargs = call[1]  # Get keyword arguments
+        assert call_kwargs["pipeline"].id == pipeline.id
+        assert call_kwargs["workflow_id"] == workflow.id
+        assert call_kwargs["user"].id == account.id
+        assert call_kwargs["invoke_from"] == InvokeFrom.PUBLISHED
+        assert call_kwargs["streaming"] == False
+        assert isinstance(call_kwargs["application_generate_entity"], RagPipelineGenerateEntity)
+
+    def test_run_single_rag_pipeline_task_entity_validation_error(
+        self, db_session_with_containers, mock_pipeline_generator, flask_app_with_containers
+    ):
+        """
+        Test run_single_rag_pipeline_task with invalid entity data.
+
+        This test verifies:
+        - Proper error handling for invalid entity data
+        - Exception logging
+        - Function raises ValueError for missing entities
+        """
+        # Arrange: Create entity data with valid UUIDs but non-existent entities
+        fake = Faker()
+        invalid_entity_data = {
+            "pipeline_id": str(uuid.uuid4()),
+            "application_generate_entity": {
+                "app_config": {
+                    "app_id": str(uuid.uuid4()),
+                    "app_name": "Test App",
+                    "mode": "workflow",
+                    "workflow_id": str(uuid.uuid4()),
+                },
+                "inputs": {"query": "Test query"},
+                "query": "Test query",
+                "response_mode": "blocking",
+                "user": str(uuid.uuid4()),
+                "files": [],
+                "conversation_id": str(uuid.uuid4()),
+            },
+            "user_id": str(uuid.uuid4()),
+            "tenant_id": str(uuid.uuid4()),
+            "workflow_id": str(uuid.uuid4()),
+            "streaming": False,
+            "workflow_execution_id": str(uuid.uuid4()),
+            "workflow_thread_pool_id": str(uuid.uuid4()),
+        }
+
+        # Act & Assert: Execute the single task with non-existent entities (should raise ValueError)
+        with flask_app_with_containers.app_context():
+            with pytest.raises(ValueError, match="Account .* not found"):
+                run_single_rag_pipeline_task(invalid_entity_data, flask_app_with_containers)
+
+        # Assert: Pipeline generator should not be called
+        mock_pipeline_generator.assert_not_called()
+
+    def test_run_single_rag_pipeline_task_database_entity_not_found(
+        self, db_session_with_containers, mock_pipeline_generator, flask_app_with_containers
+    ):
+        """
+        Test run_single_rag_pipeline_task with non-existent database entities.
+
+        This test verifies:
+        - Proper error handling for missing database entities
+        - Exception logging
+        - Function raises ValueError for missing entities
+        """
+        # Arrange: Create test data with non-existent IDs
+        fake = Faker()
+        entity_data = {
+            "pipeline_id": str(uuid.uuid4()),
+            "application_generate_entity": {
+                "app_config": {
+                    "app_id": str(uuid.uuid4()),
+                    "app_name": "Test App",
+                    "mode": "workflow",
+                    "workflow_id": str(uuid.uuid4()),
+                },
+                "inputs": {"query": "Test query"},
+                "query": "Test query",
+                "response_mode": "blocking",
+                "user": str(uuid.uuid4()),
+                "files": [],
+                "conversation_id": str(uuid.uuid4()),
+            },
+            "user_id": str(uuid.uuid4()),
+            "tenant_id": str(uuid.uuid4()),
+            "workflow_id": str(uuid.uuid4()),
+            "streaming": False,
+            "workflow_execution_id": str(uuid.uuid4()),
+            "workflow_thread_pool_id": str(uuid.uuid4()),
+        }
+
+        # Act & Assert: Execute the single task with non-existent entities (should raise ValueError)
+        with flask_app_with_containers.app_context():
+            with pytest.raises(ValueError, match="Account .* not found"):
+                run_single_rag_pipeline_task(entity_data, flask_app_with_containers)
+
+        # Assert: Pipeline generator should not be called
+        mock_pipeline_generator.assert_not_called()
+
+    def test_priority_rag_pipeline_run_task_file_not_found(
+        self, db_session_with_containers, mock_pipeline_generator, mock_file_service
+    ):
+        """
+        Test priority RAG pipeline run task with non-existent file.
+
+        This test verifies:
+        - Proper error handling for missing files
+        - Exception logging
+        - Function raises Exception for file errors
+        - Queue management continues despite file errors
+        """
+        # Arrange: Create test data
+        account, tenant, pipeline, workflow = self._create_test_pipeline_and_workflow(db_session_with_containers)
+        
+        # Mock file service to raise exception
+        file_id = str(uuid.uuid4())
+        mock_file_service["get_content"].side_effect = Exception("File not found")
+
+        # Use real Redis for TenantSelfTaskQueue
+        queue = TenantSelfTaskQueue(tenant.id, "pipeline")
+        
+        # Add waiting task to the real Redis queue
+        waiting_file_id = str(uuid.uuid4())
+        queue.push_tasks([waiting_file_id])
+
+        # Mock the task function calls
+        with patch(
+            "tasks.rag_pipeline.priority_rag_pipeline_run_task.priority_rag_pipeline_run_task.delay"
+        ) as mock_delay:
+            # Act & Assert: Execute the priority task (should raise Exception)
+            with pytest.raises(Exception, match="File not found"):
+                priority_rag_pipeline_run_task(file_id, tenant.id)
+
+            # Assert: Verify error was handled gracefully
+            mock_file_service["get_content"].assert_called_once_with(file_id)
+            mock_pipeline_generator.assert_not_called()
+
+            # Verify waiting task was still processed despite file error
+            mock_delay.assert_called_once()
+            
+            # Verify correct parameters for the call
+            call_kwargs = mock_delay.call_args[1] if mock_delay.call_args else {}
+            assert call_kwargs.get('rag_pipeline_invoke_entities_file_id') == waiting_file_id
+            assert call_kwargs.get('tenant_id') == tenant.id
+
+            # Verify queue is empty after processing (task was pulled)
+            remaining_tasks = queue.pull_tasks(count=10)
+            assert len(remaining_tasks) == 0
+
+    def test_rag_pipeline_run_task_file_not_found(
+        self, db_session_with_containers, mock_pipeline_generator, mock_file_service
+    ):
+        """
+        Test regular RAG pipeline run task with non-existent file.
+
+        This test verifies:
+        - Proper error handling for missing files
+        - Exception logging
+        - Function raises Exception for file errors
+        - Queue management continues despite file errors
+        """
+        # Arrange: Create test data
+        account, tenant, pipeline, workflow = self._create_test_pipeline_and_workflow(db_session_with_containers)
+        
+        # Mock file service to raise exception
+        file_id = str(uuid.uuid4())
+        mock_file_service["get_content"].side_effect = Exception("File not found")
+
+        # Use real Redis for TenantSelfTaskQueue
+        queue = TenantSelfTaskQueue(tenant.id, "pipeline")
+        
+        # Add waiting task to the real Redis queue
+        waiting_file_id = str(uuid.uuid4())
+        queue.push_tasks([waiting_file_id])
+
+        # Mock the task function calls
+        with patch("tasks.rag_pipeline.rag_pipeline_run_task.rag_pipeline_run_task.delay") as mock_delay:
+            # Act & Assert: Execute the regular task (should raise Exception)
+            with pytest.raises(Exception, match="File not found"):
+                rag_pipeline_run_task(file_id, tenant.id)
+
+            # Assert: Verify error was handled gracefully
+            mock_file_service["get_content"].assert_called_once_with(file_id)
+            mock_pipeline_generator.assert_not_called()
+
+            # Verify waiting task was still processed despite file error
+            mock_delay.assert_called_once()
+            
+            # Verify correct parameters for the call
+            call_kwargs = mock_delay.call_args[1] if mock_delay.call_args else {}
+            assert call_kwargs.get('rag_pipeline_invoke_entities_file_id') == waiting_file_id
+            assert call_kwargs.get('tenant_id') == tenant.id
+
+            # Verify queue is empty after processing (task was pulled)
+            remaining_tasks = queue.pull_tasks(count=10)
+            assert len(remaining_tasks) == 0

api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter_process_data.py

@@ -0,0 +1,324 @@
+"""
+Unit tests for WorkflowResponseConverter focusing on process_data truncation functionality.
+"""
+
+import uuid
+from collections.abc import Mapping
+from typing import Any
+from unittest.mock import Mock
+
+import pytest
+
+from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter
+from core.app.entities.app_invoke_entities import WorkflowAppGenerateEntity
+from core.app.entities.queue_entities import (
+    QueueNodeRetryEvent,
+    QueueNodeStartedEvent,
+    QueueNodeSucceededEvent,
+)
+from core.workflow.enums import NodeType
+from core.workflow.system_variable import SystemVariable
+from libs.datetime_utils import naive_utc_now
+from models import Account
+
+
+class TestWorkflowResponseConverterCenarios:
+    """Test process_data truncation in WorkflowResponseConverter."""
+
+    def create_mock_generate_entity(self) -> WorkflowAppGenerateEntity:
+        """Create a mock WorkflowAppGenerateEntity."""
+        mock_entity = Mock(spec=WorkflowAppGenerateEntity)
+        mock_app_config = Mock()
+        mock_app_config.tenant_id = "test-tenant-id"
+        mock_entity.app_config = mock_app_config
+        mock_entity.inputs = {}
+        return mock_entity
+
+    def create_workflow_response_converter(self) -> WorkflowResponseConverter:
+        """Create a WorkflowResponseConverter for testing."""
+
+        mock_entity = self.create_mock_generate_entity()
+        mock_user = Mock(spec=Account)
+        mock_user.id = "test-user-id"
+        mock_user.name = "Test User"
+        mock_user.email = "test@example.com"
+
+        system_variables = SystemVariable(workflow_id="wf-id", workflow_execution_id="initial-run-id")
+        return WorkflowResponseConverter(
+            application_generate_entity=mock_entity,
+            user=mock_user,
+            system_variables=system_variables,
+        )
+
+    def create_node_started_event(self, *, node_execution_id: str | None = None) -> QueueNodeStartedEvent:
+        """Create a QueueNodeStartedEvent for testing."""
+        return QueueNodeStartedEvent(
+            node_execution_id=node_execution_id or str(uuid.uuid4()),
+            node_id="test-node-id",
+            node_title="Test Node",
+            node_type=NodeType.CODE,
+            start_at=naive_utc_now(),
+            predecessor_node_id=None,
+            in_iteration_id=None,
+            in_loop_id=None,
+            provider_type="built-in",
+            provider_id="code",
+        )
+
+    def create_node_succeeded_event(
+        self,
+        *,
+        node_execution_id: str,
+        process_data: Mapping[str, Any] | None = None,
+    ) -> QueueNodeSucceededEvent:
+        """Create a QueueNodeSucceededEvent for testing."""
+        return QueueNodeSucceededEvent(
+            node_id="test-node-id",
+            node_type=NodeType.CODE,
+            node_execution_id=node_execution_id,
+            start_at=naive_utc_now(),
+            in_iteration_id=None,
+            in_loop_id=None,
+            inputs={},
+            process_data=process_data or {},
+            outputs={},
+            execution_metadata={},
+        )
+
+    def create_node_retry_event(
+        self,
+        *,
+        node_execution_id: str,
+        process_data: Mapping[str, Any] | None = None,
+    ) -> QueueNodeRetryEvent:
+        """Create a QueueNodeRetryEvent for testing."""
+        return QueueNodeRetryEvent(
+            inputs={"data": "inputs"},
+            outputs={"data": "outputs"},
+            process_data=process_data or {},
+            error="oops",
+            retry_index=1,
+            node_id="test-node-id",
+            node_type=NodeType.CODE,
+            node_title="test code",
+            provider_type="built-in",
+            provider_id="code",
+            node_execution_id=node_execution_id,
+            start_at=naive_utc_now(),
+            in_iteration_id=None,
+            in_loop_id=None,
+        )
+
+    def test_workflow_node_finish_response_uses_truncated_process_data(self):
+        """Test that node finish response uses get_response_process_data()."""
+        converter = self.create_workflow_response_converter()
+
+        original_data = {"large_field": "x" * 10000, "metadata": "info"}
+        truncated_data = {"large_field": "[TRUNCATED]", "metadata": "info"}
+
+        converter.workflow_start_to_stream_response(task_id="bootstrap", workflow_run_id="run-id", workflow_id="wf-id")
+        start_event = self.create_node_started_event()
+        converter.workflow_node_start_to_stream_response(
+            event=start_event,
+            task_id="test-task-id",
+        )
+
+        event = self.create_node_succeeded_event(
+            node_execution_id=start_event.node_execution_id,
+            process_data=original_data,
+        )
+
+        def fake_truncate(mapping):
+            if mapping == dict(original_data):
+                return truncated_data, True
+            return mapping, False
+
+        converter._truncator.truncate_variable_mapping = fake_truncate  # type: ignore[assignment]
+
+        response = converter.workflow_node_finish_to_stream_response(
+            event=event,
+            task_id="test-task-id",
+        )
+
+        # Response should use truncated data, not original
+        assert response is not None
+        assert response.data.process_data == truncated_data
+        assert response.data.process_data != original_data
+        assert response.data.process_data_truncated is True
+
+    def test_workflow_node_finish_response_without_truncation(self):
+        """Test node finish response when no truncation is applied."""
+        converter = self.create_workflow_response_converter()
+
+        original_data = {"small": "data"}
+
+        converter.workflow_start_to_stream_response(task_id="bootstrap", workflow_run_id="run-id", workflow_id="wf-id")
+        start_event = self.create_node_started_event()
+        converter.workflow_node_start_to_stream_response(
+            event=start_event,
+            task_id="test-task-id",
+        )
+
+        event = self.create_node_succeeded_event(
+            node_execution_id=start_event.node_execution_id,
+            process_data=original_data,
+        )
+
+        def fake_truncate(mapping):
+            return mapping, False
+
+        converter._truncator.truncate_variable_mapping = fake_truncate  # type: ignore[assignment]
+
+        response = converter.workflow_node_finish_to_stream_response(
+            event=event,
+            task_id="test-task-id",
+        )
+
+        # Response should use original data
+        assert response is not None
+        assert response.data.process_data == original_data
+        assert response.data.process_data_truncated is False
+
+    def test_workflow_node_finish_response_with_none_process_data(self):
+        """Test node finish response when process_data is None."""
+        converter = self.create_workflow_response_converter()
+
+        converter.workflow_start_to_stream_response(task_id="bootstrap", workflow_run_id="run-id", workflow_id="wf-id")
+        start_event = self.create_node_started_event()
+        converter.workflow_node_start_to_stream_response(
+            event=start_event,
+            task_id="test-task-id",
+        )
+
+        event = self.create_node_succeeded_event(
+            node_execution_id=start_event.node_execution_id,
+            process_data=None,
+        )
+
+        def fake_truncate(mapping):
+            return mapping, False
+
+        converter._truncator.truncate_variable_mapping = fake_truncate  # type: ignore[assignment]
+
+        response = converter.workflow_node_finish_to_stream_response(
+            event=event,
+            task_id="test-task-id",
+        )
+
+        # Response should normalize missing process_data to an empty mapping
+        assert response is not None
+        assert response.data.process_data == {}
+        assert response.data.process_data_truncated is False
+
+    def test_workflow_node_retry_response_uses_truncated_process_data(self):
+        """Test that node retry response uses get_response_process_data()."""
+        converter = self.create_workflow_response_converter()
+
+        original_data = {"large_field": "x" * 10000, "metadata": "info"}
+        truncated_data = {"large_field": "[TRUNCATED]", "metadata": "info"}
+
+        converter.workflow_start_to_stream_response(task_id="bootstrap", workflow_run_id="run-id", workflow_id="wf-id")
+        start_event = self.create_node_started_event()
+        converter.workflow_node_start_to_stream_response(
+            event=start_event,
+            task_id="test-task-id",
+        )
+
+        event = self.create_node_retry_event(
+            node_execution_id=start_event.node_execution_id,
+            process_data=original_data,
+        )
+
+        def fake_truncate(mapping):
+            if mapping == dict(original_data):
+                return truncated_data, True
+            return mapping, False
+
+        converter._truncator.truncate_variable_mapping = fake_truncate  # type: ignore[assignment]
+
+        response = converter.workflow_node_retry_to_stream_response(
+            event=event,
+            task_id="test-task-id",
+        )
+
+        # Response should use truncated data, not original
+        assert response is not None
+        assert response.data.process_data == truncated_data
+        assert response.data.process_data != original_data
+        assert response.data.process_data_truncated is True
+
+    def test_workflow_node_retry_response_without_truncation(self):
+        """Test node retry response when no truncation is applied."""
+        converter = self.create_workflow_response_converter()
+
+        original_data = {"small": "data"}
+
+        converter.workflow_start_to_stream_response(task_id="bootstrap", workflow_run_id="run-id", workflow_id="wf-id")
+        start_event = self.create_node_started_event()
+        converter.workflow_node_start_to_stream_response(
+            event=start_event,
+            task_id="test-task-id",
+        )
+
+        event = self.create_node_retry_event(
+            node_execution_id=start_event.node_execution_id,
+            process_data=original_data,
+        )
+
+        def fake_truncate(mapping):
+            return mapping, False
+
+        converter._truncator.truncate_variable_mapping = fake_truncate  # type: ignore[assignment]
+
+        response = converter.workflow_node_retry_to_stream_response(
+            event=event,
+            task_id="test-task-id",
+        )
+
+        assert response is not None
+        assert response.data.process_data == original_data
+        assert response.data.process_data_truncated is False
+
+    def test_iteration_and_loop_nodes_return_none(self):
+        """Test that iteration and loop nodes return None (no streaming events)."""
+        converter = self.create_workflow_response_converter()
+
+        iteration_event = QueueNodeSucceededEvent(
+            node_id="iteration-node",
+            node_type=NodeType.ITERATION,
+            node_execution_id=str(uuid.uuid4()),
+            start_at=naive_utc_now(),
+            in_iteration_id=None,
+            in_loop_id=None,
+            inputs={},
+            process_data={},
+            outputs={},
+            execution_metadata={},
+        )
+
+        response = converter.workflow_node_finish_to_stream_response(
+            event=iteration_event,
+            task_id="test-task-id",
+        )
+        assert response is None
+
+        loop_event = iteration_event.model_copy(update={"node_type": NodeType.LOOP})
+        response = converter.workflow_node_finish_to_stream_response(
+            event=loop_event,
+            task_id="test-task-id",
+        )
+        assert response is None
+
+    def test_finish_without_start_raises(self):
+        """Ensure finish responses require a prior workflow start."""
+        converter = self.create_workflow_response_converter()
+        event = self.create_node_succeeded_event(
+            node_execution_id=str(uuid.uuid4()),
+            process_data={},
+        )
+
+        with pytest.raises(ValueError):
+            converter.workflow_node_finish_to_stream_response(
+                event=event,
+                task_id="test-task-id",
+            )

api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter_truncation.py

@@ -1,810 +0,0 @@
-"""
-Unit tests for WorkflowResponseConverter focusing on process_data truncation functionality.
-"""
-
-import uuid
-from collections.abc import Mapping
-from dataclasses import dataclass
-from typing import Any
-from unittest.mock import Mock
-
-import pytest
-
-from core.app.app_config.entities import WorkflowUIBasedAppConfig
-from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter
-from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity
-from core.app.entities.queue_entities import (
-    QueueEvent,
-    QueueIterationStartEvent,
-    QueueLoopStartEvent,
-    QueueNodeExceptionEvent,
-    QueueNodeFailedEvent,
-    QueueNodeRetryEvent,
-    QueueNodeStartedEvent,
-    QueueNodeSucceededEvent,
-)
-from core.workflow.enums import NodeType
-from core.workflow.system_variable import SystemVariable
-from libs.datetime_utils import naive_utc_now
-from models import Account
-from models.model import AppMode
-
-
-class TestWorkflowResponseConverter:
-    """Test truncation in WorkflowResponseConverter."""
-
-    def create_mock_generate_entity(self) -> WorkflowAppGenerateEntity:
-        """Create a mock WorkflowAppGenerateEntity."""
-        mock_entity = Mock(spec=WorkflowAppGenerateEntity)
-        mock_app_config = Mock()
-        mock_app_config.tenant_id = "test-tenant-id"
-        mock_entity.invoke_from = InvokeFrom.WEB_APP
-        mock_entity.app_config = mock_app_config
-        mock_entity.inputs = {}
-        return mock_entity
-
-    def create_workflow_response_converter(self) -> WorkflowResponseConverter:
-        """Create a WorkflowResponseConverter for testing."""
-
-        mock_entity = self.create_mock_generate_entity()
-        mock_user = Mock(spec=Account)
-        mock_user.id = "test-user-id"
-        mock_user.name = "Test User"
-        mock_user.email = "test@example.com"
-
-        system_variables = SystemVariable(workflow_id="wf-id", workflow_execution_id="initial-run-id")
-        return WorkflowResponseConverter(
-            application_generate_entity=mock_entity,
-            user=mock_user,
-            system_variables=system_variables,
-        )
-
-    def create_node_started_event(self, *, node_execution_id: str | None = None) -> QueueNodeStartedEvent:
-        """Create a QueueNodeStartedEvent for testing."""
-        return QueueNodeStartedEvent(
-            node_execution_id=node_execution_id or str(uuid.uuid4()),
-            node_id="test-node-id",
-            node_title="Test Node",
-            node_type=NodeType.CODE,
-            start_at=naive_utc_now(),
-            in_iteration_id=None,
-            in_loop_id=None,
-            provider_type="built-in",
-            provider_id="code",
-        )
-
-    def create_node_succeeded_event(
-        self,
-        *,
-        node_execution_id: str,
-        process_data: Mapping[str, Any] | None = None,
-    ) -> QueueNodeSucceededEvent:
-        """Create a QueueNodeSucceededEvent for testing."""
-        return QueueNodeSucceededEvent(
-            node_id="test-node-id",
-            node_type=NodeType.CODE,
-            node_execution_id=node_execution_id,
-            start_at=naive_utc_now(),
-            in_iteration_id=None,
-            in_loop_id=None,
-            inputs={},
-            process_data=process_data or {},
-            outputs={},
-            execution_metadata={},
-        )
-
-    def create_node_retry_event(
-        self,
-        *,
-        node_execution_id: str,
-        process_data: Mapping[str, Any] | None = None,
-    ) -> QueueNodeRetryEvent:
-        """Create a QueueNodeRetryEvent for testing."""
-        return QueueNodeRetryEvent(
-            inputs={"data": "inputs"},
-            outputs={"data": "outputs"},
-            process_data=process_data or {},
-            error="oops",
-            retry_index=1,
-            node_id="test-node-id",
-            node_type=NodeType.CODE,
-            node_title="test code",
-            provider_type="built-in",
-            provider_id="code",
-            node_execution_id=node_execution_id,
-            start_at=naive_utc_now(),
-            in_iteration_id=None,
-            in_loop_id=None,
-        )
-
-    def test_workflow_node_finish_response_uses_truncated_process_data(self):
-        """Test that node finish response uses get_response_process_data()."""
-        converter = self.create_workflow_response_converter()
-
-        original_data = {"large_field": "x" * 10000, "metadata": "info"}
-        truncated_data = {"large_field": "[TRUNCATED]", "metadata": "info"}
-
-        converter.workflow_start_to_stream_response(task_id="bootstrap", workflow_run_id="run-id", workflow_id="wf-id")
-        start_event = self.create_node_started_event()
-        converter.workflow_node_start_to_stream_response(
-            event=start_event,
-            task_id="test-task-id",
-        )
-
-        event = self.create_node_succeeded_event(
-            node_execution_id=start_event.node_execution_id,
-            process_data=original_data,
-        )
-
-        def fake_truncate(mapping):
-            if mapping == dict(original_data):
-                return truncated_data, True
-            return mapping, False
-
-        converter._truncator.truncate_variable_mapping = fake_truncate  # type: ignore[assignment]
-
-        response = converter.workflow_node_finish_to_stream_response(
-            event=event,
-            task_id="test-task-id",
-        )
-
-        # Response should use truncated data, not original
-        assert response is not None
-        assert response.data.process_data == truncated_data
-        assert response.data.process_data != original_data
-        assert response.data.process_data_truncated is True
-
-    def test_workflow_node_finish_response_without_truncation(self):
-        """Test node finish response when no truncation is applied."""
-        converter = self.create_workflow_response_converter()
-
-        original_data = {"small": "data"}
-
-        converter.workflow_start_to_stream_response(task_id="bootstrap", workflow_run_id="run-id", workflow_id="wf-id")
-        start_event = self.create_node_started_event()
-        converter.workflow_node_start_to_stream_response(
-            event=start_event,
-            task_id="test-task-id",
-        )
-
-        event = self.create_node_succeeded_event(
-            node_execution_id=start_event.node_execution_id,
-            process_data=original_data,
-        )
-
-        def fake_truncate(mapping):
-            return mapping, False
-
-        converter._truncator.truncate_variable_mapping = fake_truncate  # type: ignore[assignment]
-
-        response = converter.workflow_node_finish_to_stream_response(
-            event=event,
-            task_id="test-task-id",
-        )
-
-        # Response should use original data
-        assert response is not None
-        assert response.data.process_data == original_data
-        assert response.data.process_data_truncated is False
-
-    def test_workflow_node_finish_response_with_none_process_data(self):
-        """Test node finish response when process_data is None."""
-        converter = self.create_workflow_response_converter()
-
-        converter.workflow_start_to_stream_response(task_id="bootstrap", workflow_run_id="run-id", workflow_id="wf-id")
-        start_event = self.create_node_started_event()
-        converter.workflow_node_start_to_stream_response(
-            event=start_event,
-            task_id="test-task-id",
-        )
-
-        event = self.create_node_succeeded_event(
-            node_execution_id=start_event.node_execution_id,
-            process_data=None,
-        )
-
-        def fake_truncate(mapping):
-            return mapping, False
-
-        converter._truncator.truncate_variable_mapping = fake_truncate  # type: ignore[assignment]
-
-        response = converter.workflow_node_finish_to_stream_response(
-            event=event,
-            task_id="test-task-id",
-        )
-
-        # Response should normalize missing process_data to an empty mapping
-        assert response is not None
-        assert response.data.process_data == {}
-        assert response.data.process_data_truncated is False
-
-    def test_workflow_node_retry_response_uses_truncated_process_data(self):
-        """Test that node retry response uses get_response_process_data()."""
-        converter = self.create_workflow_response_converter()
-
-        original_data = {"large_field": "x" * 10000, "metadata": "info"}
-        truncated_data = {"large_field": "[TRUNCATED]", "metadata": "info"}
-
-        converter.workflow_start_to_stream_response(task_id="bootstrap", workflow_run_id="run-id", workflow_id="wf-id")
-        start_event = self.create_node_started_event()
-        converter.workflow_node_start_to_stream_response(
-            event=start_event,
-            task_id="test-task-id",
-        )
-
-        event = self.create_node_retry_event(
-            node_execution_id=start_event.node_execution_id,
-            process_data=original_data,
-        )
-
-        def fake_truncate(mapping):
-            if mapping == dict(original_data):
-                return truncated_data, True
-            return mapping, False
-
-        converter._truncator.truncate_variable_mapping = fake_truncate  # type: ignore[assignment]
-
-        response = converter.workflow_node_retry_to_stream_response(
-            event=event,
-            task_id="test-task-id",
-        )
-
-        # Response should use truncated data, not original
-        assert response is not None
-        assert response.data.process_data == truncated_data
-        assert response.data.process_data != original_data
-        assert response.data.process_data_truncated is True
-
-    def test_workflow_node_retry_response_without_truncation(self):
-        """Test node retry response when no truncation is applied."""
-        converter = self.create_workflow_response_converter()
-
-        original_data = {"small": "data"}
-
-        converter.workflow_start_to_stream_response(task_id="bootstrap", workflow_run_id="run-id", workflow_id="wf-id")
-        start_event = self.create_node_started_event()
-        converter.workflow_node_start_to_stream_response(
-            event=start_event,
-            task_id="test-task-id",
-        )
-
-        event = self.create_node_retry_event(
-            node_execution_id=start_event.node_execution_id,
-            process_data=original_data,
-        )
-
-        def fake_truncate(mapping):
-            return mapping, False
-
-        converter._truncator.truncate_variable_mapping = fake_truncate  # type: ignore[assignment]
-
-        response = converter.workflow_node_retry_to_stream_response(
-            event=event,
-            task_id="test-task-id",
-        )
-
-        assert response is not None
-        assert response.data.process_data == original_data
-        assert response.data.process_data_truncated is False
-
-    def test_iteration_and_loop_nodes_return_none(self):
-        """Test that iteration and loop nodes return None (no streaming events)."""
-        converter = self.create_workflow_response_converter()
-
-        iteration_event = QueueNodeSucceededEvent(
-            node_id="iteration-node",
-            node_type=NodeType.ITERATION,
-            node_execution_id=str(uuid.uuid4()),
-            start_at=naive_utc_now(),
-            in_iteration_id=None,
-            in_loop_id=None,
-            inputs={},
-            process_data={},
-            outputs={},
-            execution_metadata={},
-        )
-
-        response = converter.workflow_node_finish_to_stream_response(
-            event=iteration_event,
-            task_id="test-task-id",
-        )
-        assert response is None
-
-        loop_event = iteration_event.model_copy(update={"node_type": NodeType.LOOP})
-        response = converter.workflow_node_finish_to_stream_response(
-            event=loop_event,
-            task_id="test-task-id",
-        )
-        assert response is None
-
-    def test_finish_without_start_raises(self):
-        """Ensure finish responses require a prior workflow start."""
-        converter = self.create_workflow_response_converter()
-        event = self.create_node_succeeded_event(
-            node_execution_id=str(uuid.uuid4()),
-            process_data={},
-        )
-
-        with pytest.raises(ValueError):
-            converter.workflow_node_finish_to_stream_response(
-                event=event,
-                task_id="test-task-id",
-            )
-
-
-@dataclass
-class TestCase:
-    """Test case data for table-driven tests."""
-
-    name: str
-    invoke_from: InvokeFrom
-    expected_truncation_enabled: bool
-    description: str
-
-
-class TestWorkflowResponseConverterServiceApiTruncation:
-    """Test class for Service API truncation functionality in WorkflowResponseConverter."""
-
-    def create_test_app_generate_entity(self, invoke_from: InvokeFrom) -> WorkflowAppGenerateEntity:
-        """Create a test WorkflowAppGenerateEntity with specified invoke_from."""
-        # Create a minimal WorkflowUIBasedAppConfig for testing
-        app_config = WorkflowUIBasedAppConfig(
-            tenant_id="test_tenant",
-            app_id="test_app",
-            app_mode=AppMode.WORKFLOW,
-            workflow_id="test_workflow_id",
-        )
-
-        entity = WorkflowAppGenerateEntity(
-            task_id="test_task_id",
-            app_id="test_app_id",
-            app_config=app_config,
-            tenant_id="test_tenant",
-            app_mode=AppMode.WORKFLOW,
-            invoke_from=invoke_from,
-            inputs={"test_input": "test_value"},
-            user_id="test_user_id",
-            stream=True,
-            files=[],
-            workflow_execution_id="test_workflow_exec_id",
-        )
-        return entity
-
-    def create_test_user(self) -> Account:
-        """Create a test user account."""
-        account = Account(
-            name="Test User",
-            email="test@example.com",
-        )
-        # Manually set the ID for testing purposes
-        account.id = "test_user_id"
-        return account
-
-    def create_test_system_variables(self) -> SystemVariable:
-        """Create test system variables."""
-        return SystemVariable()
-
-    def create_test_converter(self, invoke_from: InvokeFrom) -> WorkflowResponseConverter:
-        """Create WorkflowResponseConverter with specified invoke_from."""
-        entity = self.create_test_app_generate_entity(invoke_from)
-        user = self.create_test_user()
-        system_variables = self.create_test_system_variables()
-
-        converter = WorkflowResponseConverter(
-            application_generate_entity=entity,
-            user=user,
-            system_variables=system_variables,
-        )
-        # ensure `workflow_run_id` is set.
-        converter.workflow_start_to_stream_response(
-            task_id="test-task-id",
-            workflow_run_id="test-workflow-run-id",
-            workflow_id="test-workflow-id",
-        )
-        return converter
-
-    @pytest.mark.parametrize(
-        "test_case",
-        [
-            TestCase(
-                name="service_api_truncation_disabled",
-                invoke_from=InvokeFrom.SERVICE_API,
-                expected_truncation_enabled=False,
-                description="Service API calls should have truncation disabled",
-            ),
-            TestCase(
-                name="web_app_truncation_enabled",
-                invoke_from=InvokeFrom.WEB_APP,
-                expected_truncation_enabled=True,
-                description="Web app calls should have truncation enabled",
-            ),
-            TestCase(
-                name="debugger_truncation_enabled",
-                invoke_from=InvokeFrom.DEBUGGER,
-                expected_truncation_enabled=True,
-                description="Debugger calls should have truncation enabled",
-            ),
-            TestCase(
-                name="explore_truncation_enabled",
-                invoke_from=InvokeFrom.EXPLORE,
-                expected_truncation_enabled=True,
-                description="Explore calls should have truncation enabled",
-            ),
-            TestCase(
-                name="published_truncation_enabled",
-                invoke_from=InvokeFrom.PUBLISHED,
-                expected_truncation_enabled=True,
-                description="Published app calls should have truncation enabled",
-            ),
-        ],
-        ids=lambda x: x.name,
-    )
-    def test_truncator_selection_based_on_invoke_from(self, test_case: TestCase):
-        """Test that the correct truncator is selected based on invoke_from."""
-        converter = self.create_test_converter(test_case.invoke_from)
-
-        # Test truncation behavior instead of checking private attribute
-
-        # Create a test event with large data
-        large_value = {"key": ["x"] * 2000}  # Large data that would be truncated
-
-        event = QueueNodeSucceededEvent(
-            node_execution_id="test_node_exec_id",
-            node_id="test_node",
-            node_type=NodeType.LLM,
-            start_at=naive_utc_now(),
-            inputs=large_value,
-            process_data=large_value,
-            outputs=large_value,
-            error=None,
-            execution_metadata=None,
-            in_iteration_id=None,
-            in_loop_id=None,
-        )
-
-        response = converter.workflow_node_finish_to_stream_response(
-            event=event,
-            task_id="test_task",
-        )
-
-        # Verify response is not None
-        assert response is not None
-
-        # Verify truncation behavior matches expectations
-        if test_case.expected_truncation_enabled:
-            # Truncation should be enabled for non-service-api calls
-            assert response.data.inputs_truncated
-            assert response.data.process_data_truncated
-            assert response.data.outputs_truncated
-        else:
-            # SERVICE_API should not truncate
-            assert not response.data.inputs_truncated
-            assert not response.data.process_data_truncated
-            assert not response.data.outputs_truncated
-
-    def test_service_api_truncator_no_op_mapping(self):
-        """Test that Service API truncator doesn't truncate variable mappings."""
-        converter = self.create_test_converter(InvokeFrom.SERVICE_API)
-
-        # Create a test event with large data
-        large_value: dict[str, Any] = {
-            "large_string": "x" * 10000,  # Large string
-            "large_list": list(range(2000)),  # Large array
-            "nested_data": {"deep_nested": {"very_deep": {"value": "x" * 5000}}},
-        }
-
-        event = QueueNodeSucceededEvent(
-            node_execution_id="test_node_exec_id",
-            node_id="test_node",
-            node_type=NodeType.LLM,
-            start_at=naive_utc_now(),
-            inputs=large_value,
-            process_data=large_value,
-            outputs=large_value,
-            error=None,
-            execution_metadata=None,
-            in_iteration_id=None,
-            in_loop_id=None,
-        )
-
-        response = converter.workflow_node_finish_to_stream_response(
-            event=event,
-            task_id="test_task",
-        )
-
-        # Verify response is not None
-        data = response.data
-        assert data.inputs == large_value
-        assert data.process_data == large_value
-        assert data.outputs == large_value
-        # Service API should not truncate
-        assert data.inputs_truncated is False
-        assert data.process_data_truncated is False
-        assert data.outputs_truncated is False
-
-    def test_web_app_truncator_works_normally(self):
-        """Test that web app truncator still works normally."""
-        converter = self.create_test_converter(InvokeFrom.WEB_APP)
-
-        # Create a test event with large data
-        large_value = {
-            "large_string": "x" * 10000,  # Large string
-            "large_list": list(range(2000)),  # Large array
-        }
-
-        event = QueueNodeSucceededEvent(
-            node_execution_id="test_node_exec_id",
-            node_id="test_node",
-            node_type=NodeType.LLM,
-            start_at=naive_utc_now(),
-            inputs=large_value,
-            process_data=large_value,
-            outputs=large_value,
-            error=None,
-            execution_metadata=None,
-            in_iteration_id=None,
-            in_loop_id=None,
-        )
-
-        response = converter.workflow_node_finish_to_stream_response(
-            event=event,
-            task_id="test_task",
-        )
-
-        # Verify response is not None
-        assert response is not None
-
-        # Web app should truncate
-        data = response.data
-        assert data.inputs != large_value
-        assert data.process_data != large_value
-        assert data.outputs != large_value
-        # The exact behavior depends on VariableTruncator implementation
-        # Just verify that truncation flags are present
-        assert data.inputs_truncated is True
-        assert data.process_data_truncated is True
-        assert data.outputs_truncated is True
-
-    @staticmethod
-    def _create_event_by_type(
-        type_: QueueEvent, inputs: Mapping[str, Any], process_data: Mapping[str, Any], outputs: Mapping[str, Any]
-    ) -> QueueNodeSucceededEvent | QueueNodeFailedEvent | QueueNodeExceptionEvent:
-        if type_ == QueueEvent.NODE_SUCCEEDED:
-            return QueueNodeSucceededEvent(
-                node_execution_id="test_node_exec_id",
-                node_id="test_node",
-                node_type=NodeType.LLM,
-                start_at=naive_utc_now(),
-                inputs=inputs,
-                process_data=process_data,
-                outputs=outputs,
-                error=None,
-                execution_metadata=None,
-                in_iteration_id=None,
-                in_loop_id=None,
-            )
-        elif type_ == QueueEvent.NODE_FAILED:
-            return QueueNodeFailedEvent(
-                node_execution_id="test_node_exec_id",
-                node_id="test_node",
-                node_type=NodeType.LLM,
-                start_at=naive_utc_now(),
-                inputs=inputs,
-                process_data=process_data,
-                outputs=outputs,
-                error="oops",
-                execution_metadata=None,
-                in_iteration_id=None,
-                in_loop_id=None,
-            )
-        elif type_ == QueueEvent.NODE_EXCEPTION:
-            return QueueNodeExceptionEvent(
-                node_execution_id="test_node_exec_id",
-                node_id="test_node",
-                node_type=NodeType.LLM,
-                start_at=naive_utc_now(),
-                inputs=inputs,
-                process_data=process_data,
-                outputs=outputs,
-                error="oops",
-                execution_metadata=None,
-                in_iteration_id=None,
-                in_loop_id=None,
-            )
-        else:
-            raise Exception("unknown type.")
-
-    @pytest.mark.parametrize(
-        "event_type",
-        [
-            QueueEvent.NODE_SUCCEEDED,
-            QueueEvent.NODE_FAILED,
-            QueueEvent.NODE_EXCEPTION,
-        ],
-    )
-    def test_service_api_node_finish_event_no_truncation(self, event_type: QueueEvent):
-        """Test that Service API doesn't truncate node finish events."""
-        converter = self.create_test_converter(InvokeFrom.SERVICE_API)
-        # Create test event with large data
-        large_inputs = {"input1": "x" * 5000, "input2": list(range(2000))}
-        large_process_data = {"process1": "y" * 5000, "process2": {"nested": ["z"] * 2000}}
-        large_outputs = {"output1": "result" * 1000, "output2": list(range(2000))}
-
-        event = TestWorkflowResponseConverterServiceApiTruncation._create_event_by_type(
-            event_type, large_inputs, large_process_data, large_outputs
-        )
-
-        response = converter.workflow_node_finish_to_stream_response(
-            event=event,
-            task_id="test_task",
-        )
-
-        # Verify response is not None
-        assert response is not None
-
-        # Verify response contains full data (not truncated)
-        assert response.data.inputs == large_inputs
-        assert response.data.process_data == large_process_data
-        assert response.data.outputs == large_outputs
-        assert not response.data.inputs_truncated
-        assert not response.data.process_data_truncated
-        assert not response.data.outputs_truncated
-
-    def test_service_api_node_retry_event_no_truncation(self):
-        """Test that Service API doesn't truncate node retry events."""
-        converter = self.create_test_converter(InvokeFrom.SERVICE_API)
-
-        # Create test event with large data
-        large_inputs = {"retry_input": "x" * 5000}
-        large_process_data = {"retry_process": "y" * 5000}
-        large_outputs = {"retry_output": "z" * 5000}
-
-        # First, we need to store a snapshot by simulating a start event
-        start_event = QueueNodeStartedEvent(
-            node_execution_id="test_node_exec_id",
-            node_id="test_node",
-            node_type=NodeType.LLM,
-            node_title="Test Node",
-            node_run_index=1,
-            start_at=naive_utc_now(),
-            in_iteration_id=None,
-            in_loop_id=None,
-            agent_strategy=None,
-            provider_type="plugin",
-            provider_id="test/test_plugin",
-        )
-        converter.workflow_node_start_to_stream_response(event=start_event, task_id="test_task")
-
-        # Now create retry event
-        event = QueueNodeRetryEvent(
-            node_execution_id="test_node_exec_id",
-            node_id="test_node",
-            node_type=NodeType.LLM,
-            node_title="Test Node",
-            node_run_index=1,
-            start_at=naive_utc_now(),
-            inputs=large_inputs,
-            process_data=large_process_data,
-            outputs=large_outputs,
-            error="Retry error",
-            execution_metadata=None,
-            in_iteration_id=None,
-            in_loop_id=None,
-            retry_index=1,
-            provider_type="plugin",
-            provider_id="test/test_plugin",
-        )
-
-        response = converter.workflow_node_retry_to_stream_response(
-            event=event,
-            task_id="test_task",
-        )
-
-        # Verify response is not None
-        assert response is not None
-
-        # Verify response contains full data (not truncated)
-        assert response.data.inputs == large_inputs
-        assert response.data.process_data == large_process_data
-        assert response.data.outputs == large_outputs
-        assert not response.data.inputs_truncated
-        assert not response.data.process_data_truncated
-        assert not response.data.outputs_truncated
-
-    def test_service_api_iteration_events_no_truncation(self):
-        """Test that Service API doesn't truncate iteration events."""
-        converter = self.create_test_converter(InvokeFrom.SERVICE_API)
-
-        # Test iteration start event
-        large_value = {"iteration_input": ["x"] * 2000}
-
-        start_event = QueueIterationStartEvent(
-            node_execution_id="test_iter_exec_id",
-            node_id="test_iteration",
-            node_type=NodeType.ITERATION,
-            node_title="Test Iteration",
-            node_run_index=0,
-            start_at=naive_utc_now(),
-            inputs=large_value,
-            metadata={},
-        )
-
-        response = converter.workflow_iteration_start_to_stream_response(
-            task_id="test_task",
-            workflow_execution_id="test_workflow_exec_id",
-            event=start_event,
-        )
-
-        assert response is not None
-        assert response.data.inputs == large_value
-        assert not response.data.inputs_truncated
-
-    def test_service_api_loop_events_no_truncation(self):
-        """Test that Service API doesn't truncate loop events."""
-        converter = self.create_test_converter(InvokeFrom.SERVICE_API)
-
-        # Test loop start event
-        large_inputs = {"loop_input": ["x"] * 2000}
-
-        start_event = QueueLoopStartEvent(
-            node_execution_id="test_loop_exec_id",
-            node_id="test_loop",
-            node_type=NodeType.LOOP,
-            node_title="Test Loop",
-            start_at=naive_utc_now(),
-            inputs=large_inputs,
-            metadata={},
-            node_run_index=0,
-        )
-
-        response = converter.workflow_loop_start_to_stream_response(
-            task_id="test_task",
-            workflow_execution_id="test_workflow_exec_id",
-            event=start_event,
-        )
-
-        assert response is not None
-        assert response.data.inputs == large_inputs
-        assert not response.data.inputs_truncated
-
-    def test_web_app_node_finish_event_truncation_works(self):
-        """Test that web app still truncates node finish events."""
-        converter = self.create_test_converter(InvokeFrom.WEB_APP)
-
-        # Create test event with large data that should be truncated
-        large_inputs = {"input1": ["x"] * 2000}
-        large_process_data = {"process1": ["y"] * 2000}
-        large_outputs = {"output1": ["z"] * 2000}
-
-        event = QueueNodeSucceededEvent(
-            node_execution_id="test_node_exec_id",
-            node_id="test_node",
-            node_type=NodeType.LLM,
-            start_at=naive_utc_now(),
-            inputs=large_inputs,
-            process_data=large_process_data,
-            outputs=large_outputs,
-            error=None,
-            execution_metadata=None,
-            in_iteration_id=None,
-            in_loop_id=None,
-        )
-
-        response = converter.workflow_node_finish_to_stream_response(
-            event=event,
-            task_id="test_task",
-        )
-
-        # Verify response is not None
-        assert response is not None
-
-        # Verify response contains truncated data
-        # The exact behavior depends on VariableTruncator implementation
-        # Just verify truncation flags are set correctly (may or may not be truncated depending on size)
-        # At minimum, the truncation mechanism should work
-        assert isinstance(response.data.inputs, dict)
-        assert response.data.inputs_truncated
-        assert isinstance(response.data.process_data, dict)
-        assert response.data.process_data_truncated
-        assert isinstance(response.data.outputs, dict)
-        assert response.data.outputs_truncated

api/tests/unit_tests/core/mcp/auth/test_auth_flow.py

@@ -1,740 +0,0 @@
-"""Unit tests for MCP OAuth authentication flow."""
-
-from unittest.mock import Mock, patch
-
-import pytest
-
-from core.entities.mcp_provider import MCPProviderEntity
-from core.mcp.auth.auth_flow import (
-    OAUTH_STATE_EXPIRY_SECONDS,
-    OAUTH_STATE_REDIS_KEY_PREFIX,
-    OAuthCallbackState,
-    _create_secure_redis_state,
-    _retrieve_redis_state,
-    auth,
-    check_support_resource_discovery,
-    discover_oauth_metadata,
-    exchange_authorization,
-    generate_pkce_challenge,
-    handle_callback,
-    refresh_authorization,
-    register_client,
-    start_authorization,
-)
-from core.mcp.entities import AuthActionType, AuthResult
-from core.mcp.types import (
-    OAuthClientInformation,
-    OAuthClientInformationFull,
-    OAuthClientMetadata,
-    OAuthMetadata,
-    OAuthTokens,
-)
-
-
-class TestPKCEGeneration:
-    """Test PKCE challenge generation."""
-
-    def test_generate_pkce_challenge(self):
-        """Test PKCE challenge and verifier generation."""
-        code_verifier, code_challenge = generate_pkce_challenge()
-
-        # Verify format - should be URL-safe base64 without padding
-        assert "=" not in code_verifier
-        assert "+" not in code_verifier
-        assert "/" not in code_verifier
-        assert "=" not in code_challenge
-        assert "+" not in code_challenge
-        assert "/" not in code_challenge
-
-        # Verify length
-        assert len(code_verifier) > 40  # Should be around 54 characters
-        assert len(code_challenge) > 40  # Should be around 43 characters
-
-    def test_generate_pkce_challenge_uniqueness(self):
-        """Test that PKCE generation produces unique values."""
-        results = set()
-        for _ in range(10):
-            code_verifier, code_challenge = generate_pkce_challenge()
-            results.add((code_verifier, code_challenge))
-
-        # All should be unique
-        assert len(results) == 10
-
-
-class TestRedisStateManagement:
-    """Test Redis state management functions."""
-
-    @patch("core.mcp.auth.auth_flow.redis_client")
-    def test_create_secure_redis_state(self, mock_redis):
-        """Test creating secure Redis state."""
-        state_data = OAuthCallbackState(
-            provider_id="test-provider",
-            tenant_id="test-tenant",
-            server_url="https://example.com",
-            metadata=None,
-            client_information=OAuthClientInformation(client_id="test-client"),
-            code_verifier="test-verifier",
-            redirect_uri="https://redirect.example.com",
-        )
-
-        state_key = _create_secure_redis_state(state_data)
-
-        # Verify state key format
-        assert len(state_key) > 20  # Should be a secure random token
-
-        # Verify Redis call
-        mock_redis.setex.assert_called_once()
-        call_args = mock_redis.setex.call_args
-        assert call_args[0][0].startswith(OAUTH_STATE_REDIS_KEY_PREFIX)
-        assert call_args[0][1] == OAUTH_STATE_EXPIRY_SECONDS
-        assert state_data.model_dump_json() in call_args[0][2]
-
-    @patch("core.mcp.auth.auth_flow.redis_client")
-    def test_retrieve_redis_state_success(self, mock_redis):
-        """Test retrieving state from Redis."""
-        state_data = OAuthCallbackState(
-            provider_id="test-provider",
-            tenant_id="test-tenant",
-            server_url="https://example.com",
-            metadata=None,
-            client_information=OAuthClientInformation(client_id="test-client"),
-            code_verifier="test-verifier",
-            redirect_uri="https://redirect.example.com",
-        )
-        mock_redis.get.return_value = state_data.model_dump_json()
-
-        result = _retrieve_redis_state("test-state-key")
-
-        # Verify result
-        assert result.provider_id == "test-provider"
-        assert result.tenant_id == "test-tenant"
-        assert result.server_url == "https://example.com"
-
-        # Verify Redis calls
-        mock_redis.get.assert_called_once_with(f"{OAUTH_STATE_REDIS_KEY_PREFIX}test-state-key")
-        mock_redis.delete.assert_called_once_with(f"{OAUTH_STATE_REDIS_KEY_PREFIX}test-state-key")
-
-    @patch("core.mcp.auth.auth_flow.redis_client")
-    def test_retrieve_redis_state_not_found(self, mock_redis):
-        """Test retrieving non-existent state from Redis."""
-        mock_redis.get.return_value = None
-
-        with pytest.raises(ValueError) as exc_info:
-            _retrieve_redis_state("nonexistent-key")
-
-        assert "State parameter has expired or does not exist" in str(exc_info.value)
-
-    @patch("core.mcp.auth.auth_flow.redis_client")
-    def test_retrieve_redis_state_invalid_json(self, mock_redis):
-        """Test retrieving invalid JSON state from Redis."""
-        mock_redis.get.return_value = '{"invalid": json}'
-
-        with pytest.raises(ValueError) as exc_info:
-            _retrieve_redis_state("test-key")
-
-        assert "Invalid state parameter" in str(exc_info.value)
-        # State should still be deleted
-        mock_redis.delete.assert_called_once()
-
-
-class TestOAuthDiscovery:
-    """Test OAuth discovery functions."""
-
-    @patch("core.helper.ssrf_proxy.get")
-    def test_check_support_resource_discovery_success(self, mock_get):
-        """Test successful resource discovery check."""
-        mock_response = Mock()
-        mock_response.status_code = 200
-        mock_response.json.return_value = {"authorization_server_url": ["https://auth.example.com"]}
-        mock_get.return_value = mock_response
-
-        supported, auth_url = check_support_resource_discovery("https://api.example.com/endpoint")
-
-        assert supported is True
-        assert auth_url == "https://auth.example.com"
-        mock_get.assert_called_once_with(
-            "https://api.example.com/.well-known/oauth-protected-resource",
-            headers={"MCP-Protocol-Version": "2025-03-26", "User-Agent": "Dify"},
-        )
-
-    @patch("core.helper.ssrf_proxy.get")
-    def test_check_support_resource_discovery_not_supported(self, mock_get):
-        """Test resource discovery not supported."""
-        mock_response = Mock()
-        mock_response.status_code = 404
-        mock_get.return_value = mock_response
-
-        supported, auth_url = check_support_resource_discovery("https://api.example.com")
-
-        assert supported is False
-        assert auth_url == ""
-
-    @patch("core.helper.ssrf_proxy.get")
-    def test_check_support_resource_discovery_with_query_fragment(self, mock_get):
-        """Test resource discovery with query and fragment."""
-        mock_response = Mock()
-        mock_response.status_code = 200
-        mock_response.json.return_value = {"authorization_server_url": ["https://auth.example.com"]}
-        mock_get.return_value = mock_response
-
-        supported, auth_url = check_support_resource_discovery("https://api.example.com/path?query=1#fragment")
-
-        assert supported is True
-        assert auth_url == "https://auth.example.com"
-        mock_get.assert_called_once_with(
-            "https://api.example.com/.well-known/oauth-protected-resource?query=1#fragment",
-            headers={"MCP-Protocol-Version": "2025-03-26", "User-Agent": "Dify"},
-        )
-
-    @patch("core.helper.ssrf_proxy.get")
-    def test_discover_oauth_metadata_with_resource_discovery(self, mock_get):
-        """Test OAuth metadata discovery with resource discovery support."""
-        with patch("core.mcp.auth.auth_flow.check_support_resource_discovery") as mock_check:
-            mock_check.return_value = (True, "https://auth.example.com")
-
-            mock_response = Mock()
-            mock_response.status_code = 200
-            mock_response.is_success = True
-            mock_response.json.return_value = {
-                "authorization_endpoint": "https://auth.example.com/authorize",
-                "token_endpoint": "https://auth.example.com/token",
-                "response_types_supported": ["code"],
-            }
-            mock_get.return_value = mock_response
-
-            metadata = discover_oauth_metadata("https://api.example.com")
-
-            assert metadata is not None
-            assert metadata.authorization_endpoint == "https://auth.example.com/authorize"
-            assert metadata.token_endpoint == "https://auth.example.com/token"
-            mock_get.assert_called_once_with(
-                "https://auth.example.com/.well-known/oauth-authorization-server",
-                headers={"MCP-Protocol-Version": "2025-03-26"},
-            )
-
-    @patch("core.helper.ssrf_proxy.get")
-    def test_discover_oauth_metadata_without_resource_discovery(self, mock_get):
-        """Test OAuth metadata discovery without resource discovery."""
-        with patch("core.mcp.auth.auth_flow.check_support_resource_discovery") as mock_check:
-            mock_check.return_value = (False, "")
-
-            mock_response = Mock()
-            mock_response.status_code = 200
-            mock_response.is_success = True
-            mock_response.json.return_value = {
-                "authorization_endpoint": "https://api.example.com/oauth/authorize",
-                "token_endpoint": "https://api.example.com/oauth/token",
-                "response_types_supported": ["code"],
-            }
-            mock_get.return_value = mock_response
-
-            metadata = discover_oauth_metadata("https://api.example.com")
-
-            assert metadata is not None
-            assert metadata.authorization_endpoint == "https://api.example.com/oauth/authorize"
-            mock_get.assert_called_once_with(
-                "https://api.example.com/.well-known/oauth-authorization-server",
-                headers={"MCP-Protocol-Version": "2025-03-26"},
-            )
-
-    @patch("core.helper.ssrf_proxy.get")
-    def test_discover_oauth_metadata_not_found(self, mock_get):
-        """Test OAuth metadata discovery when not found."""
-        with patch("core.mcp.auth.auth_flow.check_support_resource_discovery") as mock_check:
-            mock_check.return_value = (False, "")
-
-            mock_response = Mock()
-            mock_response.status_code = 404
-            mock_get.return_value = mock_response
-
-            metadata = discover_oauth_metadata("https://api.example.com")
-
-            assert metadata is None
-
-
-class TestAuthorizationFlow:
-    """Test authorization flow functions."""
-
-    @patch("core.mcp.auth.auth_flow._create_secure_redis_state")
-    def test_start_authorization_with_metadata(self, mock_create_state):
-        """Test starting authorization with metadata."""
-        mock_create_state.return_value = "secure-state-key"
-
-        metadata = OAuthMetadata(
-            authorization_endpoint="https://auth.example.com/authorize",
-            token_endpoint="https://auth.example.com/token",
-            response_types_supported=["code"],
-            code_challenge_methods_supported=["S256"],
-        )
-        client_info = OAuthClientInformation(client_id="test-client-id")
-
-        auth_url, code_verifier = start_authorization(
-            "https://api.example.com",
-            metadata,
-            client_info,
-            "https://redirect.example.com",
-            "provider-id",
-            "tenant-id",
-        )
-
-        # Verify URL format
-        assert auth_url.startswith("https://auth.example.com/authorize?")
-        assert "response_type=code" in auth_url
-        assert "client_id=test-client-id" in auth_url
-        assert "code_challenge=" in auth_url
-        assert "code_challenge_method=S256" in auth_url
-        assert "redirect_uri=https%3A%2F%2Fredirect.example.com" in auth_url
-        assert "state=secure-state-key" in auth_url
-
-        # Verify code verifier
-        assert len(code_verifier) > 40
-
-        # Verify state was stored
-        mock_create_state.assert_called_once()
-        state_data = mock_create_state.call_args[0][0]
-        assert state_data.provider_id == "provider-id"
-        assert state_data.tenant_id == "tenant-id"
-        assert state_data.code_verifier == code_verifier
-
-    def test_start_authorization_without_metadata(self):
-        """Test starting authorization without metadata."""
-        with patch("core.mcp.auth.auth_flow._create_secure_redis_state") as mock_create_state:
-            mock_create_state.return_value = "secure-state-key"
-
-            client_info = OAuthClientInformation(client_id="test-client-id")
-
-            auth_url, code_verifier = start_authorization(
-                "https://api.example.com",
-                None,
-                client_info,
-                "https://redirect.example.com",
-                "provider-id",
-                "tenant-id",
-            )
-
-            # Should use default authorization endpoint
-            assert auth_url.startswith("https://api.example.com/authorize?")
-
-    def test_start_authorization_invalid_metadata(self):
-        """Test starting authorization with invalid metadata."""
-        metadata = OAuthMetadata(
-            authorization_endpoint="https://auth.example.com/authorize",
-            token_endpoint="https://auth.example.com/token",
-            response_types_supported=["token"],  # No "code" support
-            code_challenge_methods_supported=["plain"],  # No "S256" support
-        )
-        client_info = OAuthClientInformation(client_id="test-client-id")
-
-        with pytest.raises(ValueError) as exc_info:
-            start_authorization(
-                "https://api.example.com",
-                metadata,
-                client_info,
-                "https://redirect.example.com",
-                "provider-id",
-                "tenant-id",
-            )
-
-        assert "does not support response type code" in str(exc_info.value)
-
-    @patch("core.helper.ssrf_proxy.post")
-    def test_exchange_authorization_success(self, mock_post):
-        """Test successful authorization code exchange."""
-        mock_response = Mock()
-        mock_response.is_success = True
-        mock_response.json.return_value = {
-            "access_token": "new-access-token",
-            "token_type": "Bearer",
-            "expires_in": 3600,
-            "refresh_token": "new-refresh-token",
-        }
-        mock_post.return_value = mock_response
-
-        metadata = OAuthMetadata(
-            authorization_endpoint="https://auth.example.com/authorize",
-            token_endpoint="https://auth.example.com/token",
-            response_types_supported=["code"],
-            grant_types_supported=["authorization_code"],
-        )
-        client_info = OAuthClientInformation(client_id="test-client-id", client_secret="test-secret")
-
-        tokens = exchange_authorization(
-            "https://api.example.com",
-            metadata,
-            client_info,
-            "auth-code-123",
-            "code-verifier-xyz",
-            "https://redirect.example.com",
-        )
-
-        assert tokens.access_token == "new-access-token"
-        assert tokens.token_type == "Bearer"
-        assert tokens.expires_in == 3600
-        assert tokens.refresh_token == "new-refresh-token"
-
-        # Verify request
-        mock_post.assert_called_once_with(
-            "https://auth.example.com/token",
-            data={
-                "grant_type": "authorization_code",
-                "client_id": "test-client-id",
-                "client_secret": "test-secret",
-                "code": "auth-code-123",
-                "code_verifier": "code-verifier-xyz",
-                "redirect_uri": "https://redirect.example.com",
-            },
-        )
-
-    @patch("core.helper.ssrf_proxy.post")
-    def test_exchange_authorization_failure(self, mock_post):
-        """Test failed authorization code exchange."""
-        mock_response = Mock()
-        mock_response.is_success = False
-        mock_response.status_code = 400
-        mock_post.return_value = mock_response
-
-        client_info = OAuthClientInformation(client_id="test-client-id")
-
-        with pytest.raises(ValueError) as exc_info:
-            exchange_authorization(
-                "https://api.example.com",
-                None,
-                client_info,
-                "invalid-code",
-                "code-verifier",
-                "https://redirect.example.com",
-            )
-
-        assert "Token exchange failed: HTTP 400" in str(exc_info.value)
-
-    @patch("core.helper.ssrf_proxy.post")
-    def test_refresh_authorization_success(self, mock_post):
-        """Test successful token refresh."""
-        mock_response = Mock()
-        mock_response.is_success = True
-        mock_response.json.return_value = {
-            "access_token": "refreshed-access-token",
-            "token_type": "Bearer",
-            "expires_in": 3600,
-            "refresh_token": "new-refresh-token",
-        }
-        mock_post.return_value = mock_response
-
-        metadata = OAuthMetadata(
-            authorization_endpoint="https://auth.example.com/authorize",
-            token_endpoint="https://auth.example.com/token",
-            response_types_supported=["code"],
-            grant_types_supported=["refresh_token"],
-        )
-        client_info = OAuthClientInformation(client_id="test-client-id")
-
-        tokens = refresh_authorization("https://api.example.com", metadata, client_info, "old-refresh-token")
-
-        assert tokens.access_token == "refreshed-access-token"
-        assert tokens.refresh_token == "new-refresh-token"
-
-        # Verify request
-        mock_post.assert_called_once_with(
-            "https://auth.example.com/token",
-            data={
-                "grant_type": "refresh_token",
-                "client_id": "test-client-id",
-                "refresh_token": "old-refresh-token",
-            },
-        )
-
-    @patch("core.helper.ssrf_proxy.post")
-    def test_register_client_success(self, mock_post):
-        """Test successful client registration."""
-        mock_response = Mock()
-        mock_response.is_success = True
-        mock_response.json.return_value = {
-            "client_id": "new-client-id",
-            "client_secret": "new-client-secret",
-            "client_name": "Dify",
-            "redirect_uris": ["https://redirect.example.com"],
-        }
-        mock_post.return_value = mock_response
-
-        metadata = OAuthMetadata(
-            authorization_endpoint="https://auth.example.com/authorize",
-            token_endpoint="https://auth.example.com/token",
-            registration_endpoint="https://auth.example.com/register",
-            response_types_supported=["code"],
-        )
-        client_metadata = OAuthClientMetadata(
-            client_name="Dify",
-            redirect_uris=["https://redirect.example.com"],
-            grant_types=["authorization_code"],
-            response_types=["code"],
-        )
-
-        client_info = register_client("https://api.example.com", metadata, client_metadata)
-
-        assert isinstance(client_info, OAuthClientInformationFull)
-        assert client_info.client_id == "new-client-id"
-        assert client_info.client_secret == "new-client-secret"
-
-        # Verify request
-        mock_post.assert_called_once_with(
-            "https://auth.example.com/register",
-            json=client_metadata.model_dump(),
-            headers={"Content-Type": "application/json"},
-        )
-
-    def test_register_client_no_endpoint(self):
-        """Test client registration when no endpoint available."""
-        metadata = OAuthMetadata(
-            authorization_endpoint="https://auth.example.com/authorize",
-            token_endpoint="https://auth.example.com/token",
-            registration_endpoint=None,
-            response_types_supported=["code"],
-        )
-        client_metadata = OAuthClientMetadata(client_name="Dify", redirect_uris=["https://redirect.example.com"])
-
-        with pytest.raises(ValueError) as exc_info:
-            register_client("https://api.example.com", metadata, client_metadata)
-
-        assert "does not support dynamic client registration" in str(exc_info.value)
-
-
-class TestCallbackHandling:
-    """Test OAuth callback handling."""
-
-    @patch("core.mcp.auth.auth_flow._retrieve_redis_state")
-    @patch("core.mcp.auth.auth_flow.exchange_authorization")
-    def test_handle_callback_success(self, mock_exchange, mock_retrieve_state):
-        """Test successful callback handling."""
-        # Setup state
-        state_data = OAuthCallbackState(
-            provider_id="test-provider",
-            tenant_id="test-tenant",
-            server_url="https://api.example.com",
-            metadata=None,
-            client_information=OAuthClientInformation(client_id="test-client"),
-            code_verifier="test-verifier",
-            redirect_uri="https://redirect.example.com",
-        )
-        mock_retrieve_state.return_value = state_data
-
-        # Setup token exchange
-        tokens = OAuthTokens(
-            access_token="new-token",
-            token_type="Bearer",
-            expires_in=3600,
-        )
-        mock_exchange.return_value = tokens
-
-        # Setup service
-        mock_service = Mock()
-
-        state_result, tokens_result = handle_callback("state-key", "auth-code")
-
-        assert state_result == state_data
-        assert tokens_result == tokens
-
-        # Verify calls
-        mock_retrieve_state.assert_called_once_with("state-key")
-        mock_exchange.assert_called_once_with(
-            "https://api.example.com",
-            None,
-            state_data.client_information,
-            "auth-code",
-            "test-verifier",
-            "https://redirect.example.com",
-        )
-        # Note: handle_callback no longer saves tokens directly, it just returns them
-        # The caller (e.g., controller) is responsible for saving via execute_auth_actions
-
-
-class TestAuthOrchestration:
-    """Test the main auth orchestration function."""
-
-    @pytest.fixture
-    def mock_provider(self):
-        """Create a mock provider entity."""
-        provider = Mock(spec=MCPProviderEntity)
-        provider.id = "provider-id"
-        provider.tenant_id = "tenant-id"
-        provider.decrypt_server_url.return_value = "https://api.example.com"
-        provider.client_metadata = OAuthClientMetadata(
-            client_name="Dify",
-            redirect_uris=["https://redirect.example.com"],
-        )
-        provider.redirect_url = "https://redirect.example.com"
-        provider.retrieve_client_information.return_value = None
-        provider.retrieve_tokens.return_value = None
-        return provider
-
-    @pytest.fixture
-    def mock_service(self):
-        """Create a mock MCP service."""
-        return Mock()
-
-    @patch("core.mcp.auth.auth_flow.discover_oauth_metadata")
-    @patch("core.mcp.auth.auth_flow.register_client")
-    @patch("core.mcp.auth.auth_flow.start_authorization")
-    def test_auth_new_registration(self, mock_start_auth, mock_register, mock_discover, mock_provider, mock_service):
-        """Test auth flow for new client registration."""
-        # Setup
-        mock_discover.return_value = OAuthMetadata(
-            authorization_endpoint="https://auth.example.com/authorize",
-            token_endpoint="https://auth.example.com/token",
-            response_types_supported=["code"],
-            grant_types_supported=["authorization_code"],
-        )
-        mock_register.return_value = OAuthClientInformationFull(
-            client_id="new-client-id",
-            client_name="Dify",
-            redirect_uris=["https://redirect.example.com"],
-        )
-        mock_start_auth.return_value = ("https://auth.example.com/authorize?...", "code-verifier")
-
-        result = auth(mock_provider)
-
-        # auth() now returns AuthResult
-        assert isinstance(result, AuthResult)
-        assert result.response == {"authorization_url": "https://auth.example.com/authorize?..."}
-
-        # Verify that the result contains the correct actions
-        assert len(result.actions) == 2
-        # Check for SAVE_CLIENT_INFO action
-        client_info_action = next(a for a in result.actions if a.action_type == AuthActionType.SAVE_CLIENT_INFO)
-        assert client_info_action.data == {"client_information": mock_register.return_value.model_dump()}
-        assert client_info_action.provider_id == "provider-id"
-        assert client_info_action.tenant_id == "tenant-id"
-
-        # Check for SAVE_CODE_VERIFIER action
-        verifier_action = next(a for a in result.actions if a.action_type == AuthActionType.SAVE_CODE_VERIFIER)
-        assert verifier_action.data == {"code_verifier": "code-verifier"}
-        assert verifier_action.provider_id == "provider-id"
-        assert verifier_action.tenant_id == "tenant-id"
-
-        # Verify calls
-        mock_register.assert_called_once()
-
-    @patch("core.mcp.auth.auth_flow.discover_oauth_metadata")
-    @patch("core.mcp.auth.auth_flow._retrieve_redis_state")
-    @patch("core.mcp.auth.auth_flow.exchange_authorization")
-    def test_auth_exchange_code(self, mock_exchange, mock_retrieve_state, mock_discover, mock_provider, mock_service):
-        """Test auth flow for exchanging authorization code."""
-        # Setup metadata discovery
-        mock_discover.return_value = OAuthMetadata(
-            authorization_endpoint="https://auth.example.com/authorize",
-            token_endpoint="https://auth.example.com/token",
-            response_types_supported=["code"],
-            grant_types_supported=["authorization_code"],
-        )
-
-        # Setup existing client
-        mock_provider.retrieve_client_information.return_value = OAuthClientInformation(client_id="existing-client")
-
-        # Setup state retrieval
-        state_data = OAuthCallbackState(
-            provider_id="provider-id",
-            tenant_id="tenant-id",
-            server_url="https://api.example.com",
-            metadata=None,
-            client_information=OAuthClientInformation(client_id="existing-client"),
-            code_verifier="test-verifier",
-            redirect_uri="https://redirect.example.com",
-        )
-        mock_retrieve_state.return_value = state_data
-
-        # Setup token exchange
-        tokens = OAuthTokens(access_token="new-token", token_type="Bearer", expires_in=3600)
-        mock_exchange.return_value = tokens
-
-        result = auth(mock_provider, authorization_code="auth-code", state_param="state-key")
-
-        # auth() now returns AuthResult, not a dict
-        assert isinstance(result, AuthResult)
-        assert result.response == {"result": "success"}
-
-        # Verify that the result contains the correct action
-        assert len(result.actions) == 1
-        assert result.actions[0].action_type == AuthActionType.SAVE_TOKENS
-        assert result.actions[0].data == tokens.model_dump()
-        assert result.actions[0].provider_id == "provider-id"
-        assert result.actions[0].tenant_id == "tenant-id"
-
-    @patch("core.mcp.auth.auth_flow.discover_oauth_metadata")
-    def test_auth_exchange_code_without_state(self, mock_discover, mock_provider, mock_service):
-        """Test auth flow fails when exchanging code without state."""
-        # Setup metadata discovery
-        mock_discover.return_value = OAuthMetadata(
-            authorization_endpoint="https://auth.example.com/authorize",
-            token_endpoint="https://auth.example.com/token",
-            response_types_supported=["code"],
-            grant_types_supported=["authorization_code"],
-        )
-
-        mock_provider.retrieve_client_information.return_value = OAuthClientInformation(client_id="existing-client")
-
-        with pytest.raises(ValueError) as exc_info:
-            auth(mock_provider, authorization_code="auth-code")
-
-        assert "State parameter is required" in str(exc_info.value)
-
-    @patch("core.mcp.auth.auth_flow.refresh_authorization")
-    def test_auth_refresh_token(self, mock_refresh, mock_provider, mock_service):
-        """Test auth flow for refreshing tokens."""
-        # Setup existing client and tokens
-        mock_provider.retrieve_client_information.return_value = OAuthClientInformation(client_id="existing-client")
-        mock_provider.retrieve_tokens.return_value = OAuthTokens(
-            access_token="old-token",
-            token_type="Bearer",
-            expires_in=0,
-            refresh_token="refresh-token",
-        )
-
-        # Setup refresh
-        new_tokens = OAuthTokens(
-            access_token="refreshed-token",
-            token_type="Bearer",
-            expires_in=3600,
-            refresh_token="new-refresh-token",
-        )
-        mock_refresh.return_value = new_tokens
-
-        with patch("core.mcp.auth.auth_flow.discover_oauth_metadata") as mock_discover:
-            mock_discover.return_value = OAuthMetadata(
-                authorization_endpoint="https://auth.example.com/authorize",
-                token_endpoint="https://auth.example.com/token",
-                response_types_supported=["code"],
-                grant_types_supported=["authorization_code"],
-            )
-
-            result = auth(mock_provider)
-
-            # auth() now returns AuthResult
-            assert isinstance(result, AuthResult)
-            assert result.response == {"result": "success"}
-
-            # Verify that the result contains the correct action
-            assert len(result.actions) == 1
-            assert result.actions[0].action_type == AuthActionType.SAVE_TOKENS
-            assert result.actions[0].data == new_tokens.model_dump()
-            assert result.actions[0].provider_id == "provider-id"
-            assert result.actions[0].tenant_id == "tenant-id"
-
-            # Verify refresh was called
-            mock_refresh.assert_called_once()
-
-    @patch("core.mcp.auth.auth_flow.discover_oauth_metadata")
-    def test_auth_registration_fails_with_code(self, mock_discover, mock_provider, mock_service):
-        """Test auth fails when no client info exists but code is provided."""
-        # Setup metadata discovery
-        mock_discover.return_value = OAuthMetadata(
-            authorization_endpoint="https://auth.example.com/authorize",
-            token_endpoint="https://auth.example.com/token",
-            response_types_supported=["code"],
-            grant_types_supported=["authorization_code"],
-        )
-
-        mock_provider.retrieve_client_information.return_value = None
-
-        with pytest.raises(ValueError) as exc_info:
-            auth(mock_provider, authorization_code="auth-code")
-
-        assert "Existing OAuth client information is required" in str(exc_info.value)

api/tests/unit_tests/core/mcp/client/test_session.py

@@ -395,6 +395,9 @@ def test_client_capabilities_default():
 
     # Assert default capabilities
     assert received_capabilities is not None
+    assert received_capabilities.sampling is not None
+    assert received_capabilities.roots is not None
+    assert received_capabilities.roots.listChanged is True
 
 
 def test_client_capabilities_with_custom_callbacks():

api/tests/unit_tests/core/mcp/test_entities.py

@@ -1,239 +0,0 @@
-"""Unit tests for MCP entities module."""
-
-from unittest.mock import Mock
-
-from core.mcp.entities import (
-    SUPPORTED_PROTOCOL_VERSIONS,
-    LifespanContextT,
-    RequestContext,
-    SessionT,
-)
-from core.mcp.session.base_session import BaseSession
-from core.mcp.types import LATEST_PROTOCOL_VERSION, RequestParams
-
-
-class TestProtocolVersions:
-    """Test protocol version constants."""
-
-    def test_supported_protocol_versions(self):
-        """Test supported protocol versions list."""
-        assert isinstance(SUPPORTED_PROTOCOL_VERSIONS, list)
-        assert len(SUPPORTED_PROTOCOL_VERSIONS) >= 3
-        assert "2024-11-05" in SUPPORTED_PROTOCOL_VERSIONS
-        assert "2025-03-26" in SUPPORTED_PROTOCOL_VERSIONS
-        assert LATEST_PROTOCOL_VERSION in SUPPORTED_PROTOCOL_VERSIONS
-
-    def test_latest_protocol_version_is_supported(self):
-        """Test that latest protocol version is in supported versions."""
-        assert LATEST_PROTOCOL_VERSION in SUPPORTED_PROTOCOL_VERSIONS
-
-
-class TestRequestContext:
-    """Test RequestContext dataclass."""
-
-    def test_request_context_creation(self):
-        """Test creating a RequestContext instance."""
-        mock_session = Mock(spec=BaseSession)
-        mock_lifespan = {"key": "value"}
-        mock_meta = RequestParams.Meta(progressToken="test-token")
-
-        context = RequestContext(
-            request_id="test-request-123",
-            meta=mock_meta,
-            session=mock_session,
-            lifespan_context=mock_lifespan,
-        )
-
-        assert context.request_id == "test-request-123"
-        assert context.meta == mock_meta
-        assert context.session == mock_session
-        assert context.lifespan_context == mock_lifespan
-
-    def test_request_context_with_none_meta(self):
-        """Test creating RequestContext with None meta."""
-        mock_session = Mock(spec=BaseSession)
-
-        context = RequestContext(
-            request_id=42,  # Can be int or string
-            meta=None,
-            session=mock_session,
-            lifespan_context=None,
-        )
-
-        assert context.request_id == 42
-        assert context.meta is None
-        assert context.session == mock_session
-        assert context.lifespan_context is None
-
-    def test_request_context_attributes(self):
-        """Test RequestContext attributes are accessible."""
-        mock_session = Mock(spec=BaseSession)
-
-        context = RequestContext(
-            request_id="test-123",
-            meta=None,
-            session=mock_session,
-            lifespan_context=None,
-        )
-
-        # Verify attributes are accessible
-        assert hasattr(context, "request_id")
-        assert hasattr(context, "meta")
-        assert hasattr(context, "session")
-        assert hasattr(context, "lifespan_context")
-
-        # Verify values
-        assert context.request_id == "test-123"
-        assert context.meta is None
-        assert context.session == mock_session
-        assert context.lifespan_context is None
-
-    def test_request_context_generic_typing(self):
-        """Test RequestContext with different generic types."""
-        # Create a mock session with specific type
-        mock_session = Mock(spec=BaseSession)
-
-        # Create context with string lifespan context
-        context_str = RequestContext[BaseSession, str](
-            request_id="test-1",
-            meta=None,
-            session=mock_session,
-            lifespan_context="string-context",
-        )
-        assert isinstance(context_str.lifespan_context, str)
-
-        # Create context with dict lifespan context
-        context_dict = RequestContext[BaseSession, dict](
-            request_id="test-2",
-            meta=None,
-            session=mock_session,
-            lifespan_context={"key": "value"},
-        )
-        assert isinstance(context_dict.lifespan_context, dict)
-
-        # Create context with custom object lifespan context
-        class CustomLifespan:
-            def __init__(self, data):
-                self.data = data
-
-        custom_lifespan = CustomLifespan("test-data")
-        context_custom = RequestContext[BaseSession, CustomLifespan](
-            request_id="test-3",
-            meta=None,
-            session=mock_session,
-            lifespan_context=custom_lifespan,
-        )
-        assert isinstance(context_custom.lifespan_context, CustomLifespan)
-        assert context_custom.lifespan_context.data == "test-data"
-
-    def test_request_context_with_progress_meta(self):
-        """Test RequestContext with progress metadata."""
-        mock_session = Mock(spec=BaseSession)
-        progress_meta = RequestParams.Meta(progressToken="progress-123")
-
-        context = RequestContext(
-            request_id="req-456",
-            meta=progress_meta,
-            session=mock_session,
-            lifespan_context=None,
-        )
-
-        assert context.meta is not None
-        assert context.meta.progressToken == "progress-123"
-
-    def test_request_context_equality(self):
-        """Test RequestContext equality comparison."""
-        mock_session1 = Mock(spec=BaseSession)
-        mock_session2 = Mock(spec=BaseSession)
-
-        context1 = RequestContext(
-            request_id="test-123",
-            meta=None,
-            session=mock_session1,
-            lifespan_context="context",
-        )
-
-        context2 = RequestContext(
-            request_id="test-123",
-            meta=None,
-            session=mock_session1,
-            lifespan_context="context",
-        )
-
-        context3 = RequestContext(
-            request_id="test-456",
-            meta=None,
-            session=mock_session1,
-            lifespan_context="context",
-        )
-
-        # Same values should be equal
-        assert context1 == context2
-
-        # Different request_id should not be equal
-        assert context1 != context3
-
-        # Different session should not be equal
-        context4 = RequestContext(
-            request_id="test-123",
-            meta=None,
-            session=mock_session2,
-            lifespan_context="context",
-        )
-        assert context1 != context4
-
-    def test_request_context_repr(self):
-        """Test RequestContext string representation."""
-        mock_session = Mock(spec=BaseSession)
-        mock_session.__repr__ = Mock(return_value="<MockSession>")
-
-        context = RequestContext(
-            request_id="test-123",
-            meta=None,
-            session=mock_session,
-            lifespan_context={"data": "test"},
-        )
-
-        repr_str = repr(context)
-        assert "RequestContext" in repr_str
-        assert "test-123" in repr_str
-        assert "MockSession" in repr_str
-
-
-class TestTypeVariables:
-    """Test type variables defined in the module."""
-
-    def test_session_type_var(self):
-        """Test SessionT type variable."""
-
-        # Create a custom session class
-        class CustomSession(BaseSession):
-            pass
-
-        # Use in generic context
-        def process_session(session: SessionT) -> SessionT:
-            return session
-
-        mock_session = Mock(spec=CustomSession)
-        result = process_session(mock_session)
-        assert result == mock_session
-
-    def test_lifespan_context_type_var(self):
-        """Test LifespanContextT type variable."""
-
-        # Use in generic context
-        def process_lifespan(context: LifespanContextT) -> LifespanContextT:
-            return context
-
-        # Test with different types
-        str_context = "string-context"
-        assert process_lifespan(str_context) == str_context
-
-        dict_context = {"key": "value"}
-        assert process_lifespan(dict_context) == dict_context
-
-        class CustomContext:
-            pass
-
-        custom_context = CustomContext()
-        assert process_lifespan(custom_context) == custom_context

api/tests/unit_tests/core/mcp/test_error.py

@@ -1,205 +0,0 @@
-"""Unit tests for MCP error classes."""
-
-import pytest
-
-from core.mcp.error import MCPAuthError, MCPConnectionError, MCPError
-
-
-class TestMCPError:
-    """Test MCPError base exception class."""
-
-    def test_mcp_error_creation(self):
-        """Test creating MCPError instance."""
-        error = MCPError("Test error message")
-        assert str(error) == "Test error message"
-        assert isinstance(error, Exception)
-
-    def test_mcp_error_inheritance(self):
-        """Test MCPError inherits from Exception."""
-        error = MCPError()
-        assert isinstance(error, Exception)
-        assert type(error).__name__ == "MCPError"
-
-    def test_mcp_error_with_empty_message(self):
-        """Test MCPError with empty message."""
-        error = MCPError()
-        assert str(error) == ""
-
-    def test_mcp_error_raise(self):
-        """Test raising MCPError."""
-        with pytest.raises(MCPError) as exc_info:
-            raise MCPError("Something went wrong")
-
-        assert str(exc_info.value) == "Something went wrong"
-
-
-class TestMCPConnectionError:
-    """Test MCPConnectionError exception class."""
-
-    def test_mcp_connection_error_creation(self):
-        """Test creating MCPConnectionError instance."""
-        error = MCPConnectionError("Connection failed")
-        assert str(error) == "Connection failed"
-        assert isinstance(error, MCPError)
-        assert isinstance(error, Exception)
-
-    def test_mcp_connection_error_inheritance(self):
-        """Test MCPConnectionError inheritance chain."""
-        error = MCPConnectionError()
-        assert isinstance(error, MCPConnectionError)
-        assert isinstance(error, MCPError)
-        assert isinstance(error, Exception)
-
-    def test_mcp_connection_error_raise(self):
-        """Test raising MCPConnectionError."""
-        with pytest.raises(MCPConnectionError) as exc_info:
-            raise MCPConnectionError("Unable to connect to server")
-
-        assert str(exc_info.value) == "Unable to connect to server"
-
-    def test_mcp_connection_error_catch_as_mcp_error(self):
-        """Test catching MCPConnectionError as MCPError."""
-        with pytest.raises(MCPError) as exc_info:
-            raise MCPConnectionError("Connection issue")
-
-        assert isinstance(exc_info.value, MCPConnectionError)
-        assert str(exc_info.value) == "Connection issue"
-
-
-class TestMCPAuthError:
-    """Test MCPAuthError exception class."""
-
-    def test_mcp_auth_error_creation(self):
-        """Test creating MCPAuthError instance."""
-        error = MCPAuthError("Authentication failed")
-        assert str(error) == "Authentication failed"
-        assert isinstance(error, MCPConnectionError)
-        assert isinstance(error, MCPError)
-        assert isinstance(error, Exception)
-
-    def test_mcp_auth_error_inheritance(self):
-        """Test MCPAuthError inheritance chain."""
-        error = MCPAuthError()
-        assert isinstance(error, MCPAuthError)
-        assert isinstance(error, MCPConnectionError)
-        assert isinstance(error, MCPError)
-        assert isinstance(error, Exception)
-
-    def test_mcp_auth_error_raise(self):
-        """Test raising MCPAuthError."""
-        with pytest.raises(MCPAuthError) as exc_info:
-            raise MCPAuthError("Invalid credentials")
-
-        assert str(exc_info.value) == "Invalid credentials"
-
-    def test_mcp_auth_error_catch_hierarchy(self):
-        """Test catching MCPAuthError at different levels."""
-        # Catch as MCPAuthError
-        with pytest.raises(MCPAuthError) as exc_info:
-            raise MCPAuthError("Auth specific error")
-        assert str(exc_info.value) == "Auth specific error"
-
-        # Catch as MCPConnectionError
-        with pytest.raises(MCPConnectionError) as exc_info:
-            raise MCPAuthError("Auth connection error")
-        assert isinstance(exc_info.value, MCPAuthError)
-        assert str(exc_info.value) == "Auth connection error"
-
-        # Catch as MCPError
-        with pytest.raises(MCPError) as exc_info:
-            raise MCPAuthError("Auth base error")
-        assert isinstance(exc_info.value, MCPAuthError)
-        assert str(exc_info.value) == "Auth base error"
-
-
-class TestErrorHierarchy:
-    """Test the complete error hierarchy."""
-
-    def test_exception_hierarchy(self):
-        """Test the complete exception hierarchy."""
-        # Create instances
-        base_error = MCPError("base")
-        connection_error = MCPConnectionError("connection")
-        auth_error = MCPAuthError("auth")
-
-        # Test type relationships
-        assert not isinstance(base_error, MCPConnectionError)
-        assert not isinstance(base_error, MCPAuthError)
-
-        assert isinstance(connection_error, MCPError)
-        assert not isinstance(connection_error, MCPAuthError)
-
-        assert isinstance(auth_error, MCPError)
-        assert isinstance(auth_error, MCPConnectionError)
-
-    def test_error_handling_patterns(self):
-        """Test common error handling patterns."""
-
-        def raise_auth_error():
-            raise MCPAuthError("401 Unauthorized")
-
-        def raise_connection_error():
-            raise MCPConnectionError("Connection timeout")
-
-        def raise_base_error():
-            raise MCPError("Generic error")
-
-        # Pattern 1: Catch specific errors first
-        errors_caught = []
-
-        for error_func in [raise_auth_error, raise_connection_error, raise_base_error]:
-            try:
-                error_func()
-            except MCPAuthError:
-                errors_caught.append("auth")
-            except MCPConnectionError:
-                errors_caught.append("connection")
-            except MCPError:
-                errors_caught.append("base")
-
-        assert errors_caught == ["auth", "connection", "base"]
-
-        # Pattern 2: Catch all as base error
-        for error_func in [raise_auth_error, raise_connection_error, raise_base_error]:
-            with pytest.raises(MCPError) as exc_info:
-                error_func()
-            assert isinstance(exc_info.value, MCPError)
-
-    def test_error_with_cause(self):
-        """Test errors with cause (chained exceptions)."""
-        original_error = ValueError("Original error")
-
-        def raise_chained_error():
-            try:
-                raise original_error
-            except ValueError as e:
-                raise MCPConnectionError("Connection failed") from e
-
-        with pytest.raises(MCPConnectionError) as exc_info:
-            raise_chained_error()
-
-        assert str(exc_info.value) == "Connection failed"
-        assert exc_info.value.__cause__ == original_error
-
-    def test_error_comparison(self):
-        """Test error instance comparison."""
-        error1 = MCPError("Test message")
-        error2 = MCPError("Test message")
-        error3 = MCPError("Different message")
-
-        # Errors are not equal even with same message (different instances)
-        assert error1 != error2
-        assert error1 != error3
-
-        # But they have the same type
-        assert type(error1) == type(error2) == type(error3)
-
-    def test_error_representation(self):
-        """Test error string representation."""
-        base_error = MCPError("Base error message")
-        connection_error = MCPConnectionError("Connection error message")
-        auth_error = MCPAuthError("Auth error message")
-
-        assert repr(base_error) == "MCPError('Base error message')"
-        assert repr(connection_error) == "MCPConnectionError('Connection error message')"
-        assert repr(auth_error) == "MCPAuthError('Auth error message')"

api/tests/unit_tests/core/mcp/test_mcp_client.py

@@ -1,382 +0,0 @@
-"""Unit tests for MCP client."""
-
-from contextlib import ExitStack
-from types import TracebackType
-from unittest.mock import Mock, patch
-
-import pytest
-
-from core.mcp.error import MCPConnectionError
-from core.mcp.mcp_client import MCPClient
-from core.mcp.types import CallToolResult, ListToolsResult, TextContent, Tool, ToolAnnotations
-
-
-class TestMCPClient:
-    """Test suite for MCPClient."""
-
-    def test_init(self):
-        """Test client initialization."""
-        client = MCPClient(
-            server_url="http://test.example.com/mcp",
-            headers={"Authorization": "Bearer test"},
-            timeout=30.0,
-            sse_read_timeout=60.0,
-        )
-
-        assert client.server_url == "http://test.example.com/mcp"
-        assert client.headers == {"Authorization": "Bearer test"}
-        assert client.timeout == 30.0
-        assert client.sse_read_timeout == 60.0
-        assert client._session is None
-        assert isinstance(client._exit_stack, ExitStack)
-        assert client._initialized is False
-
-    def test_init_defaults(self):
-        """Test client initialization with defaults."""
-        client = MCPClient(server_url="http://test.example.com")
-
-        assert client.server_url == "http://test.example.com"
-        assert client.headers == {}
-        assert client.timeout is None
-        assert client.sse_read_timeout is None
-
-    @patch("core.mcp.mcp_client.streamablehttp_client")
-    @patch("core.mcp.mcp_client.ClientSession")
-    def test_initialize_with_mcp_url(self, mock_client_session, mock_streamable_client):
-        """Test initialization with MCP URL."""
-        # Setup mocks
-        mock_read_stream = Mock()
-        mock_write_stream = Mock()
-        mock_client_context = Mock()
-        mock_streamable_client.return_value.__enter__.return_value = (
-            mock_read_stream,
-            mock_write_stream,
-            mock_client_context,
-        )
-
-        mock_session = Mock()
-        mock_client_session.return_value.__enter__.return_value = mock_session
-
-        client = MCPClient(server_url="http://test.example.com/mcp")
-        client._initialize()
-
-        # Verify streamable client was called
-        mock_streamable_client.assert_called_once_with(
-            url="http://test.example.com/mcp",
-            headers={},
-            timeout=None,
-            sse_read_timeout=None,
-        )
-
-        # Verify session was created
-        mock_client_session.assert_called_once_with(mock_read_stream, mock_write_stream)
-        mock_session.initialize.assert_called_once()
-        assert client._session == mock_session
-
-    @patch("core.mcp.mcp_client.sse_client")
-    @patch("core.mcp.mcp_client.ClientSession")
-    def test_initialize_with_sse_url(self, mock_client_session, mock_sse_client):
-        """Test initialization with SSE URL."""
-        # Setup mocks
-        mock_read_stream = Mock()
-        mock_write_stream = Mock()
-        mock_sse_client.return_value.__enter__.return_value = (mock_read_stream, mock_write_stream)
-
-        mock_session = Mock()
-        mock_client_session.return_value.__enter__.return_value = mock_session
-
-        client = MCPClient(server_url="http://test.example.com/sse")
-        client._initialize()
-
-        # Verify SSE client was called
-        mock_sse_client.assert_called_once_with(
-            url="http://test.example.com/sse",
-            headers={},
-            timeout=None,
-            sse_read_timeout=None,
-        )
-
-        # Verify session was created
-        mock_client_session.assert_called_once_with(mock_read_stream, mock_write_stream)
-        mock_session.initialize.assert_called_once()
-        assert client._session == mock_session
-
-    @patch("core.mcp.mcp_client.sse_client")
-    @patch("core.mcp.mcp_client.streamablehttp_client")
-    @patch("core.mcp.mcp_client.ClientSession")
-    def test_initialize_with_unknown_method_fallback_to_sse(
-        self, mock_client_session, mock_streamable_client, mock_sse_client
-    ):
-        """Test initialization with unknown method falls back to SSE."""
-        # Setup mocks
-        mock_read_stream = Mock()
-        mock_write_stream = Mock()
-        mock_sse_client.return_value.__enter__.return_value = (mock_read_stream, mock_write_stream)
-
-        mock_session = Mock()
-        mock_client_session.return_value.__enter__.return_value = mock_session
-
-        client = MCPClient(server_url="http://test.example.com/unknown")
-        client._initialize()
-
-        # Verify SSE client was tried
-        mock_sse_client.assert_called_once()
-        mock_streamable_client.assert_not_called()
-
-        # Verify session was created
-        assert client._session == mock_session
-
-    @patch("core.mcp.mcp_client.sse_client")
-    @patch("core.mcp.mcp_client.streamablehttp_client")
-    @patch("core.mcp.mcp_client.ClientSession")
-    def test_initialize_fallback_from_sse_to_mcp(self, mock_client_session, mock_streamable_client, mock_sse_client):
-        """Test initialization falls back from SSE to MCP on connection error."""
-        # Setup SSE to fail
-        mock_sse_client.side_effect = MCPConnectionError("SSE connection failed")
-
-        # Setup MCP to succeed
-        mock_read_stream = Mock()
-        mock_write_stream = Mock()
-        mock_client_context = Mock()
-        mock_streamable_client.return_value.__enter__.return_value = (
-            mock_read_stream,
-            mock_write_stream,
-            mock_client_context,
-        )
-
-        mock_session = Mock()
-        mock_client_session.return_value.__enter__.return_value = mock_session
-
-        client = MCPClient(server_url="http://test.example.com/unknown")
-        client._initialize()
-
-        # Verify both were tried
-        mock_sse_client.assert_called_once()
-        mock_streamable_client.assert_called_once()
-
-        # Verify session was created with MCP
-        assert client._session == mock_session
-
-    @patch("core.mcp.mcp_client.streamablehttp_client")
-    @patch("core.mcp.mcp_client.ClientSession")
-    def test_connect_server_mcp(self, mock_client_session, mock_streamable_client):
-        """Test connect_server with MCP method."""
-        # Setup mocks
-        mock_read_stream = Mock()
-        mock_write_stream = Mock()
-        mock_client_context = Mock()
-        mock_streamable_client.return_value.__enter__.return_value = (
-            mock_read_stream,
-            mock_write_stream,
-            mock_client_context,
-        )
-
-        mock_session = Mock()
-        mock_client_session.return_value.__enter__.return_value = mock_session
-
-        client = MCPClient(server_url="http://test.example.com")
-        client.connect_server(mock_streamable_client, "mcp")
-
-        # Verify correct streams were passed
-        mock_client_session.assert_called_once_with(mock_read_stream, mock_write_stream)
-        mock_session.initialize.assert_called_once()
-
-    @patch("core.mcp.mcp_client.sse_client")
-    @patch("core.mcp.mcp_client.ClientSession")
-    def test_connect_server_sse(self, mock_client_session, mock_sse_client):
-        """Test connect_server with SSE method."""
-        # Setup mocks
-        mock_read_stream = Mock()
-        mock_write_stream = Mock()
-        mock_sse_client.return_value.__enter__.return_value = (mock_read_stream, mock_write_stream)
-
-        mock_session = Mock()
-        mock_client_session.return_value.__enter__.return_value = mock_session
-
-        client = MCPClient(server_url="http://test.example.com")
-        client.connect_server(mock_sse_client, "sse")
-
-        # Verify correct streams were passed
-        mock_client_session.assert_called_once_with(mock_read_stream, mock_write_stream)
-        mock_session.initialize.assert_called_once()
-
-    def test_context_manager_enter(self):
-        """Test context manager enter."""
-        client = MCPClient(server_url="http://test.example.com")
-
-        with patch.object(client, "_initialize") as mock_initialize:
-            result = client.__enter__()
-
-            assert result == client
-            assert client._initialized is True
-            mock_initialize.assert_called_once()
-
-    def test_context_manager_exit(self):
-        """Test context manager exit."""
-        client = MCPClient(server_url="http://test.example.com")
-
-        with patch.object(client, "cleanup") as mock_cleanup:
-            exc_type: type[BaseException] | None = None
-            exc_val: BaseException | None = None
-            exc_tb: TracebackType | None = None
-            client.__exit__(exc_type, exc_val, exc_tb)
-
-            mock_cleanup.assert_called_once()
-
-    def test_list_tools_not_initialized(self):
-        """Test list_tools when session not initialized."""
-        client = MCPClient(server_url="http://test.example.com")
-
-        with pytest.raises(ValueError) as exc_info:
-            client.list_tools()
-
-        assert "Session not initialized" in str(exc_info.value)
-
-    def test_list_tools_success(self):
-        """Test successful list_tools call."""
-        client = MCPClient(server_url="http://test.example.com")
-
-        # Setup mock session
-        mock_session = Mock()
-        expected_tools = [
-            Tool(
-                name="test-tool",
-                description="A test tool",
-                inputSchema={"type": "object", "properties": {}},
-                annotations=ToolAnnotations(title="Test Tool"),
-            )
-        ]
-        mock_session.list_tools.return_value = ListToolsResult(tools=expected_tools)
-        client._session = mock_session
-
-        result = client.list_tools()
-
-        assert result == expected_tools
-        mock_session.list_tools.assert_called_once()
-
-    def test_invoke_tool_not_initialized(self):
-        """Test invoke_tool when session not initialized."""
-        client = MCPClient(server_url="http://test.example.com")
-
-        with pytest.raises(ValueError) as exc_info:
-            client.invoke_tool("test-tool", {"arg": "value"})
-
-        assert "Session not initialized" in str(exc_info.value)
-
-    def test_invoke_tool_success(self):
-        """Test successful invoke_tool call."""
-        client = MCPClient(server_url="http://test.example.com")
-
-        # Setup mock session
-        mock_session = Mock()
-        expected_result = CallToolResult(
-            content=[TextContent(type="text", text="Tool executed successfully")],
-            isError=False,
-        )
-        mock_session.call_tool.return_value = expected_result
-        client._session = mock_session
-
-        result = client.invoke_tool("test-tool", {"arg": "value"})
-
-        assert result == expected_result
-        mock_session.call_tool.assert_called_once_with("test-tool", {"arg": "value"})
-
-    def test_cleanup(self):
-        """Test cleanup method."""
-        client = MCPClient(server_url="http://test.example.com")
-        mock_exit_stack = Mock(spec=ExitStack)
-        client._exit_stack = mock_exit_stack
-        client._session = Mock()
-        client._initialized = True
-
-        client.cleanup()
-
-        mock_exit_stack.close.assert_called_once()
-        assert client._session is None
-        assert client._initialized is False
-
-    def test_cleanup_with_error(self):
-        """Test cleanup method with error."""
-        client = MCPClient(server_url="http://test.example.com")
-        mock_exit_stack = Mock(spec=ExitStack)
-        mock_exit_stack.close.side_effect = Exception("Cleanup error")
-        client._exit_stack = mock_exit_stack
-        client._session = Mock()
-        client._initialized = True
-
-        with pytest.raises(ValueError) as exc_info:
-            client.cleanup()
-
-        assert "Error during cleanup: Cleanup error" in str(exc_info.value)
-        assert client._session is None
-        assert client._initialized is False
-
-    @patch("core.mcp.mcp_client.streamablehttp_client")
-    @patch("core.mcp.mcp_client.ClientSession")
-    def test_full_context_manager_flow(self, mock_client_session, mock_streamable_client):
-        """Test full context manager flow."""
-        # Setup mocks
-        mock_read_stream = Mock()
-        mock_write_stream = Mock()
-        mock_client_context = Mock()
-        mock_streamable_client.return_value.__enter__.return_value = (
-            mock_read_stream,
-            mock_write_stream,
-            mock_client_context,
-        )
-
-        mock_session = Mock()
-        mock_client_session.return_value.__enter__.return_value = mock_session
-
-        expected_tools = [Tool(name="test-tool", description="Test", inputSchema={})]
-        mock_session.list_tools.return_value = ListToolsResult(tools=expected_tools)
-
-        with MCPClient(server_url="http://test.example.com/mcp") as client:
-            assert client._initialized is True
-            assert client._session == mock_session
-
-            # Test tool operations
-            tools = client.list_tools()
-            assert tools == expected_tools
-
-        # After exit, should be cleaned up
-        assert client._initialized is False
-        assert client._session is None
-
-    def test_headers_passed_to_clients(self):
-        """Test that headers are properly passed to underlying clients."""
-        custom_headers = {
-            "Authorization": "Bearer test-token",
-            "X-Custom-Header": "test-value",
-        }
-
-        with patch("core.mcp.mcp_client.streamablehttp_client") as mock_streamable_client:
-            with patch("core.mcp.mcp_client.ClientSession") as mock_client_session:
-                # Setup mocks
-                mock_read_stream = Mock()
-                mock_write_stream = Mock()
-                mock_client_context = Mock()
-                mock_streamable_client.return_value.__enter__.return_value = (
-                    mock_read_stream,
-                    mock_write_stream,
-                    mock_client_context,
-                )
-
-                mock_session = Mock()
-                mock_client_session.return_value.__enter__.return_value = mock_session
-
-                client = MCPClient(
-                    server_url="http://test.example.com/mcp",
-                    headers=custom_headers,
-                    timeout=30.0,
-                    sse_read_timeout=60.0,
-                )
-                client._initialize()
-
-                # Verify headers were passed
-                mock_streamable_client.assert_called_once_with(
-                    url="http://test.example.com/mcp",
-                    headers=custom_headers,
-                    timeout=30.0,
-                    sse_read_timeout=60.0,
-                )

api/tests/unit_tests/core/mcp/test_types.py

@@ -1,492 +0,0 @@
-"""Unit tests for MCP types module."""
-
-import pytest
-from pydantic import ValidationError
-
-from core.mcp.types import (
-    INTERNAL_ERROR,
-    INVALID_PARAMS,
-    INVALID_REQUEST,
-    LATEST_PROTOCOL_VERSION,
-    METHOD_NOT_FOUND,
-    PARSE_ERROR,
-    SERVER_LATEST_PROTOCOL_VERSION,
-    Annotations,
-    CallToolRequest,
-    CallToolRequestParams,
-    CallToolResult,
-    ClientCapabilities,
-    CompleteRequest,
-    CompleteRequestParams,
-    CompleteResult,
-    Completion,
-    CompletionArgument,
-    CompletionContext,
-    ErrorData,
-    ImageContent,
-    Implementation,
-    InitializeRequest,
-    InitializeRequestParams,
-    InitializeResult,
-    JSONRPCError,
-    JSONRPCMessage,
-    JSONRPCNotification,
-    JSONRPCRequest,
-    JSONRPCResponse,
-    ListToolsRequest,
-    ListToolsResult,
-    OAuthClientInformation,
-    OAuthClientMetadata,
-    OAuthMetadata,
-    OAuthTokens,
-    PingRequest,
-    ProgressNotification,
-    ProgressNotificationParams,
-    PromptReference,
-    RequestParams,
-    ResourceTemplateReference,
-    Result,
-    ServerCapabilities,
-    TextContent,
-    Tool,
-    ToolAnnotations,
-)
-
-
-class TestConstants:
-    """Test module constants."""
-
-    def test_protocol_versions(self):
-        """Test protocol version constants."""
-        assert LATEST_PROTOCOL_VERSION == "2025-03-26"
-        assert SERVER_LATEST_PROTOCOL_VERSION == "2024-11-05"
-
-    def test_error_codes(self):
-        """Test JSON-RPC error code constants."""
-        assert PARSE_ERROR == -32700
-        assert INVALID_REQUEST == -32600
-        assert METHOD_NOT_FOUND == -32601
-        assert INVALID_PARAMS == -32602
-        assert INTERNAL_ERROR == -32603
-
-
-class TestRequestParams:
-    """Test RequestParams and related classes."""
-
-    def test_request_params_basic(self):
-        """Test basic RequestParams creation."""
-        params = RequestParams()
-        assert params.meta is None
-
-    def test_request_params_with_meta(self):
-        """Test RequestParams with meta."""
-        meta = RequestParams.Meta(progressToken="test-token")
-        params = RequestParams(_meta=meta)
-        assert params.meta is not None
-        assert params.meta.progressToken == "test-token"
-
-    def test_request_params_meta_extra_fields(self):
-        """Test RequestParams.Meta allows extra fields."""
-        meta = RequestParams.Meta(progressToken="token", customField="value")
-        assert meta.progressToken == "token"
-        assert meta.customField == "value"  # type: ignore
-
-    def test_request_params_serialization(self):
-        """Test RequestParams serialization with _meta alias."""
-        meta = RequestParams.Meta(progressToken="test")
-        params = RequestParams(_meta=meta)
-
-        # Model dump should use the alias
-        dumped = params.model_dump(by_alias=True)
-        assert "_meta" in dumped
-        assert dumped["_meta"] is not None
-        assert dumped["_meta"]["progressToken"] == "test"
-
-
-class TestJSONRPCMessages:
-    """Test JSON-RPC message types."""
-
-    def test_jsonrpc_request(self):
-        """Test JSONRPCRequest creation and validation."""
-        request = JSONRPCRequest(jsonrpc="2.0", id="test-123", method="test_method", params={"key": "value"})
-
-        assert request.jsonrpc == "2.0"
-        assert request.id == "test-123"
-        assert request.method == "test_method"
-        assert request.params == {"key": "value"}
-
-    def test_jsonrpc_request_numeric_id(self):
-        """Test JSONRPCRequest with numeric ID."""
-        request = JSONRPCRequest(jsonrpc="2.0", id=123, method="test", params=None)
-        assert request.id == 123
-
-    def test_jsonrpc_notification(self):
-        """Test JSONRPCNotification creation."""
-        notification = JSONRPCNotification(jsonrpc="2.0", method="notification_method", params={"data": "test"})
-
-        assert notification.jsonrpc == "2.0"
-        assert notification.method == "notification_method"
-        assert not hasattr(notification, "id")  # Notifications don't have ID
-
-    def test_jsonrpc_response(self):
-        """Test JSONRPCResponse creation."""
-        response = JSONRPCResponse(jsonrpc="2.0", id="req-123", result={"success": True})
-
-        assert response.jsonrpc == "2.0"
-        assert response.id == "req-123"
-        assert response.result == {"success": True}
-
-    def test_jsonrpc_error(self):
-        """Test JSONRPCError creation."""
-        error_data = ErrorData(code=INVALID_PARAMS, message="Invalid parameters", data={"field": "missing"})
-
-        error = JSONRPCError(jsonrpc="2.0", id="req-123", error=error_data)
-
-        assert error.jsonrpc == "2.0"
-        assert error.id == "req-123"
-        assert error.error.code == INVALID_PARAMS
-        assert error.error.message == "Invalid parameters"
-        assert error.error.data == {"field": "missing"}
-
-    def test_jsonrpc_message_parsing(self):
-        """Test JSONRPCMessage parsing different message types."""
-        # Parse request
-        request_json = '{"jsonrpc": "2.0", "id": 1, "method": "test", "params": null}'
-        msg = JSONRPCMessage.model_validate_json(request_json)
-        assert isinstance(msg.root, JSONRPCRequest)
-
-        # Parse response
-        response_json = '{"jsonrpc": "2.0", "id": 1, "result": {"data": "test"}}'
-        msg = JSONRPCMessage.model_validate_json(response_json)
-        assert isinstance(msg.root, JSONRPCResponse)
-
-        # Parse error
-        error_json = '{"jsonrpc": "2.0", "id": 1, "error": {"code": -32600, "message": "Invalid Request"}}'
-        msg = JSONRPCMessage.model_validate_json(error_json)
-        assert isinstance(msg.root, JSONRPCError)
-
-
-class TestCapabilities:
-    """Test capability classes."""
-
-    def test_client_capabilities(self):
-        """Test ClientCapabilities creation."""
-        caps = ClientCapabilities(
-            experimental={"feature": {"enabled": True}},
-            sampling={"model_config": {"extra": "allow"}},
-            roots={"listChanged": True},
-        )
-
-        assert caps.experimental == {"feature": {"enabled": True}}
-        assert caps.sampling is not None
-        assert caps.roots.listChanged is True  # type: ignore
-
-    def test_server_capabilities(self):
-        """Test ServerCapabilities creation."""
-        caps = ServerCapabilities(
-            tools={"listChanged": True},
-            resources={"subscribe": True, "listChanged": False},
-            prompts={"listChanged": True},
-            logging={},
-            completions={},
-        )
-
-        assert caps.tools.listChanged is True  # type: ignore
-        assert caps.resources.subscribe is True  # type: ignore
-        assert caps.resources.listChanged is False  # type: ignore
-
-
-class TestInitialization:
-    """Test initialization request/response types."""
-
-    def test_initialize_request(self):
-        """Test InitializeRequest creation."""
-        client_info = Implementation(name="test-client", version="1.0.0")
-        capabilities = ClientCapabilities()
-
-        params = InitializeRequestParams(
-            protocolVersion=LATEST_PROTOCOL_VERSION, capabilities=capabilities, clientInfo=client_info
-        )
-
-        request = InitializeRequest(params=params)
-
-        assert request.method == "initialize"
-        assert request.params.protocolVersion == LATEST_PROTOCOL_VERSION
-        assert request.params.clientInfo.name == "test-client"
-
-    def test_initialize_result(self):
-        """Test InitializeResult creation."""
-        server_info = Implementation(name="test-server", version="1.0.0")
-        capabilities = ServerCapabilities()
-
-        result = InitializeResult(
-            protocolVersion=LATEST_PROTOCOL_VERSION,
-            capabilities=capabilities,
-            serverInfo=server_info,
-            instructions="Welcome to test server",
-        )
-
-        assert result.protocolVersion == LATEST_PROTOCOL_VERSION
-        assert result.serverInfo.name == "test-server"
-        assert result.instructions == "Welcome to test server"
-
-
-class TestTools:
-    """Test tool-related types."""
-
-    def test_tool_creation(self):
-        """Test Tool creation with all fields."""
-        tool = Tool(
-            name="test_tool",
-            title="Test Tool",
-            description="A tool for testing",
-            inputSchema={"type": "object", "properties": {"input": {"type": "string"}}, "required": ["input"]},
-            outputSchema={"type": "object", "properties": {"result": {"type": "string"}}},
-            annotations=ToolAnnotations(
-                title="Test Tool", readOnlyHint=False, destructiveHint=False, idempotentHint=True
-            ),
-        )
-
-        assert tool.name == "test_tool"
-        assert tool.title == "Test Tool"
-        assert tool.description == "A tool for testing"
-        assert tool.inputSchema["properties"]["input"]["type"] == "string"
-        assert tool.annotations.idempotentHint is True
-
-    def test_call_tool_request(self):
-        """Test CallToolRequest creation."""
-        params = CallToolRequestParams(name="test_tool", arguments={"input": "test value"})
-
-        request = CallToolRequest(params=params)
-
-        assert request.method == "tools/call"
-        assert request.params.name == "test_tool"
-        assert request.params.arguments == {"input": "test value"}
-
-    def test_call_tool_result(self):
-        """Test CallToolResult creation."""
-        result = CallToolResult(
-            content=[TextContent(type="text", text="Tool executed successfully")],
-            structuredContent={"status": "success", "data": "test"},
-            isError=False,
-        )
-
-        assert len(result.content) == 1
-        assert result.content[0].text == "Tool executed successfully"  # type: ignore
-        assert result.structuredContent == {"status": "success", "data": "test"}
-        assert result.isError is False
-
-    def test_list_tools_request(self):
-        """Test ListToolsRequest creation."""
-        request = ListToolsRequest()
-        assert request.method == "tools/list"
-
-    def test_list_tools_result(self):
-        """Test ListToolsResult creation."""
-        tool1 = Tool(name="tool1", inputSchema={})
-        tool2 = Tool(name="tool2", inputSchema={})
-
-        result = ListToolsResult(tools=[tool1, tool2])
-
-        assert len(result.tools) == 2
-        assert result.tools[0].name == "tool1"
-        assert result.tools[1].name == "tool2"
-
-
-class TestContent:
-    """Test content types."""
-
-    def test_text_content(self):
-        """Test TextContent creation."""
-        annotations = Annotations(audience=["user"], priority=0.8)
-        content = TextContent(type="text", text="Hello, world!", annotations=annotations)
-
-        assert content.type == "text"
-        assert content.text == "Hello, world!"
-        assert content.annotations is not None
-        assert content.annotations.priority == 0.8
-
-    def test_image_content(self):
-        """Test ImageContent creation."""
-        content = ImageContent(type="image", data="base64encodeddata", mimeType="image/png")
-
-        assert content.type == "image"
-        assert content.data == "base64encodeddata"
-        assert content.mimeType == "image/png"
-
-
-class TestOAuth:
-    """Test OAuth-related types."""
-
-    def test_oauth_client_metadata(self):
-        """Test OAuthClientMetadata creation."""
-        metadata = OAuthClientMetadata(
-            client_name="Test Client",
-            redirect_uris=["https://example.com/callback"],
-            grant_types=["authorization_code", "refresh_token"],
-            response_types=["code"],
-            token_endpoint_auth_method="none",
-            client_uri="https://example.com",
-            scope="read write",
-        )
-
-        assert metadata.client_name == "Test Client"
-        assert len(metadata.redirect_uris) == 1
-        assert "authorization_code" in metadata.grant_types
-
-    def test_oauth_client_information(self):
-        """Test OAuthClientInformation creation."""
-        info = OAuthClientInformation(client_id="test-client-id", client_secret="test-secret")
-
-        assert info.client_id == "test-client-id"
-        assert info.client_secret == "test-secret"
-
-    def test_oauth_client_information_without_secret(self):
-        """Test OAuthClientInformation without secret."""
-        info = OAuthClientInformation(client_id="public-client")
-
-        assert info.client_id == "public-client"
-        assert info.client_secret is None
-
-    def test_oauth_tokens(self):
-        """Test OAuthTokens creation."""
-        tokens = OAuthTokens(
-            access_token="access-token-123",
-            token_type="Bearer",
-            expires_in=3600,
-            refresh_token="refresh-token-456",
-            scope="read write",
-        )
-
-        assert tokens.access_token == "access-token-123"
-        assert tokens.token_type == "Bearer"
-        assert tokens.expires_in == 3600
-        assert tokens.refresh_token == "refresh-token-456"
-        assert tokens.scope == "read write"
-
-    def test_oauth_metadata(self):
-        """Test OAuthMetadata creation."""
-        metadata = OAuthMetadata(
-            authorization_endpoint="https://auth.example.com/authorize",
-            token_endpoint="https://auth.example.com/token",
-            registration_endpoint="https://auth.example.com/register",
-            response_types_supported=["code", "token"],
-            grant_types_supported=["authorization_code", "refresh_token"],
-            code_challenge_methods_supported=["plain", "S256"],
-        )
-
-        assert metadata.authorization_endpoint == "https://auth.example.com/authorize"
-        assert "code" in metadata.response_types_supported
-        assert "S256" in metadata.code_challenge_methods_supported
-
-
-class TestNotifications:
-    """Test notification types."""
-
-    def test_progress_notification(self):
-        """Test ProgressNotification creation."""
-        params = ProgressNotificationParams(
-            progressToken="progress-123", progress=50.0, total=100.0, message="Processing... 50%"
-        )
-
-        notification = ProgressNotification(params=params)
-
-        assert notification.method == "notifications/progress"
-        assert notification.params.progressToken == "progress-123"
-        assert notification.params.progress == 50.0
-        assert notification.params.total == 100.0
-        assert notification.params.message == "Processing... 50%"
-
-    def test_ping_request(self):
-        """Test PingRequest creation."""
-        request = PingRequest()
-        assert request.method == "ping"
-        assert request.params is None
-
-
-class TestCompletion:
-    """Test completion-related types."""
-
-    def test_completion_context(self):
-        """Test CompletionContext creation."""
-        context = CompletionContext(arguments={"template_var": "value"})
-        assert context.arguments == {"template_var": "value"}
-
-    def test_resource_template_reference(self):
-        """Test ResourceTemplateReference creation."""
-        ref = ResourceTemplateReference(type="ref/resource", uri="file:///path/to/{filename}")
-        assert ref.type == "ref/resource"
-        assert ref.uri == "file:///path/to/{filename}"
-
-    def test_prompt_reference(self):
-        """Test PromptReference creation."""
-        ref = PromptReference(type="ref/prompt", name="test_prompt")
-        assert ref.type == "ref/prompt"
-        assert ref.name == "test_prompt"
-
-    def test_complete_request(self):
-        """Test CompleteRequest creation."""
-        ref = PromptReference(type="ref/prompt", name="test_prompt")
-        arg = CompletionArgument(name="arg1", value="val")
-
-        params = CompleteRequestParams(ref=ref, argument=arg, context=CompletionContext(arguments={"key": "value"}))
-
-        request = CompleteRequest(params=params)
-
-        assert request.method == "completion/complete"
-        assert request.params.ref.name == "test_prompt"  # type: ignore
-        assert request.params.argument.name == "arg1"
-
-    def test_complete_result(self):
-        """Test CompleteResult creation."""
-        completion = Completion(values=["option1", "option2", "option3"], total=10, hasMore=True)
-
-        result = CompleteResult(completion=completion)
-
-        assert len(result.completion.values) == 3
-        assert result.completion.total == 10
-        assert result.completion.hasMore is True
-
-
-class TestValidation:
-    """Test validation of various types."""
-
-    def test_invalid_jsonrpc_version(self):
-        """Test invalid JSON-RPC version validation."""
-        with pytest.raises(ValidationError):
-            JSONRPCRequest(
-                jsonrpc="1.0",  # Invalid version
-                id=1,
-                method="test",
-            )
-
-    def test_tool_annotations_validation(self):
-        """Test ToolAnnotations with invalid values."""
-        # Valid annotations
-        annotations = ToolAnnotations(
-            title="Test", readOnlyHint=True, destructiveHint=False, idempotentHint=True, openWorldHint=False
-        )
-        assert annotations.title == "Test"
-
-    def test_extra_fields_allowed(self):
-        """Test that extra fields are allowed in models."""
-        # Most models should allow extra fields
-        tool = Tool(
-            name="test",
-            inputSchema={},
-            customField="allowed",  # type: ignore
-        )
-        assert tool.customField == "allowed"  # type: ignore
-
-    def test_result_meta_alias(self):
-        """Test Result model with _meta alias."""
-        # Create with the field name (not alias)
-        result = Result(_meta={"key": "value"})
-
-        # Verify the field is set correctly
-        assert result.meta == {"key": "value"}
-
-        # Dump with alias
-        dumped = result.model_dump(by_alias=True)
-        assert "_meta" in dumped
-        assert dumped["_meta"] == {"key": "value"}

api/tests/unit_tests/core/mcp/test_utils.py

@@ -1,355 +0,0 @@
-"""Unit tests for MCP utils module."""
-
-import json
-from collections.abc import Generator
-from unittest.mock import MagicMock, Mock, patch
-
-import httpx
-import httpx_sse
-import pytest
-
-from core.mcp.utils import (
-    STATUS_FORCELIST,
-    create_mcp_error_response,
-    create_ssrf_proxy_mcp_http_client,
-    ssrf_proxy_sse_connect,
-)
-
-
-class TestConstants:
-    """Test module constants."""
-
-    def test_status_forcelist(self):
-        """Test STATUS_FORCELIST contains expected HTTP status codes."""
-        assert STATUS_FORCELIST == [429, 500, 502, 503, 504]
-        assert 429 in STATUS_FORCELIST  # Too Many Requests
-        assert 500 in STATUS_FORCELIST  # Internal Server Error
-        assert 502 in STATUS_FORCELIST  # Bad Gateway
-        assert 503 in STATUS_FORCELIST  # Service Unavailable
-        assert 504 in STATUS_FORCELIST  # Gateway Timeout
-
-
-class TestCreateSSRFProxyMCPHTTPClient:
-    """Test create_ssrf_proxy_mcp_http_client function."""
-
-    @patch("core.mcp.utils.dify_config")
-    def test_create_client_with_all_url_proxy(self, mock_config):
-        """Test client creation with SSRF_PROXY_ALL_URL configured."""
-        mock_config.SSRF_PROXY_ALL_URL = "http://proxy.example.com:8080"
-        mock_config.HTTP_REQUEST_NODE_SSL_VERIFY = True
-
-        client = create_ssrf_proxy_mcp_http_client(
-            headers={"Authorization": "Bearer token"}, timeout=httpx.Timeout(30.0)
-        )
-
-        assert isinstance(client, httpx.Client)
-        assert client.headers["Authorization"] == "Bearer token"
-        assert client.timeout.connect == 30.0
-        assert client.follow_redirects is True
-
-        # Clean up
-        client.close()
-
-    @patch("core.mcp.utils.dify_config")
-    def test_create_client_with_http_https_proxies(self, mock_config):
-        """Test client creation with separate HTTP/HTTPS proxies."""
-        mock_config.SSRF_PROXY_ALL_URL = None
-        mock_config.SSRF_PROXY_HTTP_URL = "http://http-proxy.example.com:8080"
-        mock_config.SSRF_PROXY_HTTPS_URL = "http://https-proxy.example.com:8443"
-        mock_config.HTTP_REQUEST_NODE_SSL_VERIFY = False
-
-        client = create_ssrf_proxy_mcp_http_client()
-
-        assert isinstance(client, httpx.Client)
-        assert client.follow_redirects is True
-
-        # Clean up
-        client.close()
-
-    @patch("core.mcp.utils.dify_config")
-    def test_create_client_without_proxy(self, mock_config):
-        """Test client creation without proxy configuration."""
-        mock_config.SSRF_PROXY_ALL_URL = None
-        mock_config.SSRF_PROXY_HTTP_URL = None
-        mock_config.SSRF_PROXY_HTTPS_URL = None
-        mock_config.HTTP_REQUEST_NODE_SSL_VERIFY = True
-
-        headers = {"X-Custom-Header": "value"}
-        timeout = httpx.Timeout(timeout=30.0, connect=5.0, read=10.0, write=30.0)
-
-        client = create_ssrf_proxy_mcp_http_client(headers=headers, timeout=timeout)
-
-        assert isinstance(client, httpx.Client)
-        assert client.headers["X-Custom-Header"] == "value"
-        assert client.timeout.connect == 5.0
-        assert client.timeout.read == 10.0
-        assert client.follow_redirects is True
-
-        # Clean up
-        client.close()
-
-    @patch("core.mcp.utils.dify_config")
-    def test_create_client_default_params(self, mock_config):
-        """Test client creation with default parameters."""
-        mock_config.SSRF_PROXY_ALL_URL = None
-        mock_config.SSRF_PROXY_HTTP_URL = None
-        mock_config.SSRF_PROXY_HTTPS_URL = None
-        mock_config.HTTP_REQUEST_NODE_SSL_VERIFY = True
-
-        client = create_ssrf_proxy_mcp_http_client()
-
-        assert isinstance(client, httpx.Client)
-        # httpx.Client adds default headers, so we just check it's a Headers object
-        assert isinstance(client.headers, httpx.Headers)
-        # When no timeout is provided, httpx uses its default timeout
-        assert client.timeout is not None
-
-        # Clean up
-        client.close()
-
-
-class TestSSRFProxySSEConnect:
-    """Test ssrf_proxy_sse_connect function."""
-
-    @patch("core.mcp.utils.connect_sse")
-    @patch("core.mcp.utils.create_ssrf_proxy_mcp_http_client")
-    def test_sse_connect_with_provided_client(self, mock_create_client, mock_connect_sse):
-        """Test SSE connection with pre-configured client."""
-        # Setup mocks
-        mock_client = Mock(spec=httpx.Client)
-        mock_event_source = Mock(spec=httpx_sse.EventSource)
-        mock_context = MagicMock()
-        mock_context.__enter__.return_value = mock_event_source
-        mock_connect_sse.return_value = mock_context
-
-        # Call with provided client
-        result = ssrf_proxy_sse_connect(
-            "http://example.com/sse", client=mock_client, method="POST", headers={"Authorization": "Bearer token"}
-        )
-
-        # Verify client creation was not called
-        mock_create_client.assert_not_called()
-
-        # Verify connect_sse was called correctly
-        mock_connect_sse.assert_called_once_with(
-            mock_client, "POST", "http://example.com/sse", headers={"Authorization": "Bearer token"}
-        )
-
-        # Verify result
-        assert result == mock_context
-
-    @patch("core.mcp.utils.connect_sse")
-    @patch("core.mcp.utils.create_ssrf_proxy_mcp_http_client")
-    @patch("core.mcp.utils.dify_config")
-    def test_sse_connect_without_client(self, mock_config, mock_create_client, mock_connect_sse):
-        """Test SSE connection without pre-configured client."""
-        # Setup config
-        mock_config.SSRF_DEFAULT_TIME_OUT = 30.0
-        mock_config.SSRF_DEFAULT_CONNECT_TIME_OUT = 10.0
-        mock_config.SSRF_DEFAULT_READ_TIME_OUT = 60.0
-        mock_config.SSRF_DEFAULT_WRITE_TIME_OUT = 30.0
-
-        # Setup mocks
-        mock_client = Mock(spec=httpx.Client)
-        mock_create_client.return_value = mock_client
-
-        mock_event_source = Mock(spec=httpx_sse.EventSource)
-        mock_context = MagicMock()
-        mock_context.__enter__.return_value = mock_event_source
-        mock_connect_sse.return_value = mock_context
-
-        # Call without client
-        result = ssrf_proxy_sse_connect("http://example.com/sse", headers={"X-Custom": "value"})
-
-        # Verify client was created
-        mock_create_client.assert_called_once()
-        call_args = mock_create_client.call_args
-        assert call_args[1]["headers"] == {"X-Custom": "value"}
-
-        timeout = call_args[1]["timeout"]
-        # httpx.Timeout object has these attributes
-        assert isinstance(timeout, httpx.Timeout)
-        assert timeout.connect == 10.0
-        assert timeout.read == 60.0
-        assert timeout.write == 30.0
-
-        # Verify connect_sse was called
-        mock_connect_sse.assert_called_once_with(
-            mock_client,
-            "GET",  # Default method
-            "http://example.com/sse",
-        )
-
-        # Verify result
-        assert result == mock_context
-
-    @patch("core.mcp.utils.connect_sse")
-    @patch("core.mcp.utils.create_ssrf_proxy_mcp_http_client")
-    def test_sse_connect_with_custom_timeout(self, mock_create_client, mock_connect_sse):
-        """Test SSE connection with custom timeout."""
-        # Setup mocks
-        mock_client = Mock(spec=httpx.Client)
-        mock_create_client.return_value = mock_client
-
-        mock_event_source = Mock(spec=httpx_sse.EventSource)
-        mock_context = MagicMock()
-        mock_context.__enter__.return_value = mock_event_source
-        mock_connect_sse.return_value = mock_context
-
-        custom_timeout = httpx.Timeout(timeout=60.0, read=120.0)
-
-        # Call with custom timeout
-        result = ssrf_proxy_sse_connect("http://example.com/sse", timeout=custom_timeout)
-
-        # Verify client was created with custom timeout
-        mock_create_client.assert_called_once()
-        call_args = mock_create_client.call_args
-        assert call_args[1]["timeout"] == custom_timeout
-
-        # Verify result
-        assert result == mock_context
-
-    @patch("core.mcp.utils.connect_sse")
-    @patch("core.mcp.utils.create_ssrf_proxy_mcp_http_client")
-    def test_sse_connect_error_cleanup(self, mock_create_client, mock_connect_sse):
-        """Test SSE connection cleans up client on error."""
-        # Setup mocks
-        mock_client = Mock(spec=httpx.Client)
-        mock_create_client.return_value = mock_client
-
-        # Make connect_sse raise an exception
-        mock_connect_sse.side_effect = httpx.ConnectError("Connection failed")
-
-        # Call should raise the exception
-        with pytest.raises(httpx.ConnectError):
-            ssrf_proxy_sse_connect("http://example.com/sse")
-
-        # Verify client was cleaned up
-        mock_client.close.assert_called_once()
-
-    @patch("core.mcp.utils.connect_sse")
-    def test_sse_connect_error_no_cleanup_with_provided_client(self, mock_connect_sse):
-        """Test SSE connection doesn't clean up provided client on error."""
-        # Setup mocks
-        mock_client = Mock(spec=httpx.Client)
-
-        # Make connect_sse raise an exception
-        mock_connect_sse.side_effect = httpx.ConnectError("Connection failed")
-
-        # Call should raise the exception
-        with pytest.raises(httpx.ConnectError):
-            ssrf_proxy_sse_connect("http://example.com/sse", client=mock_client)
-
-        # Verify client was NOT cleaned up (because it was provided)
-        mock_client.close.assert_not_called()
-
-
-class TestCreateMCPErrorResponse:
-    """Test create_mcp_error_response function."""
-
-    def test_create_error_response_basic(self):
-        """Test creating basic error response."""
-        generator = create_mcp_error_response(request_id="req-123", code=-32600, message="Invalid Request")
-
-        # Generator should yield bytes
-        assert isinstance(generator, Generator)
-
-        # Get the response
-        response_bytes = next(generator)
-        assert isinstance(response_bytes, bytes)
-
-        # Parse the response
-        response_str = response_bytes.decode("utf-8")
-        response_json = json.loads(response_str)
-
-        assert response_json["jsonrpc"] == "2.0"
-        assert response_json["id"] == "req-123"
-        assert response_json["error"]["code"] == -32600
-        assert response_json["error"]["message"] == "Invalid Request"
-        assert response_json["error"]["data"] is None
-
-        # Generator should be exhausted
-        with pytest.raises(StopIteration):
-            next(generator)
-
-    def test_create_error_response_with_data(self):
-        """Test creating error response with additional data."""
-        error_data = {"field": "username", "reason": "required"}
-
-        generator = create_mcp_error_response(
-            request_id=456,  # Numeric ID
-            code=-32602,
-            message="Invalid params",
-            data=error_data,
-        )
-
-        response_bytes = next(generator)
-        response_json = json.loads(response_bytes.decode("utf-8"))
-
-        assert response_json["id"] == 456
-        assert response_json["error"]["code"] == -32602
-        assert response_json["error"]["message"] == "Invalid params"
-        assert response_json["error"]["data"] == error_data
-
-    def test_create_error_response_without_request_id(self):
-        """Test creating error response without request ID."""
-        generator = create_mcp_error_response(request_id=None, code=-32700, message="Parse error")
-
-        response_bytes = next(generator)
-        response_json = json.loads(response_bytes.decode("utf-8"))
-
-        # Should default to ID 1
-        assert response_json["id"] == 1
-        assert response_json["error"]["code"] == -32700
-        assert response_json["error"]["message"] == "Parse error"
-
-    def test_create_error_response_with_complex_data(self):
-        """Test creating error response with complex error data."""
-        complex_data = {
-            "errors": [{"field": "name", "message": "Too short"}, {"field": "email", "message": "Invalid format"}],
-            "timestamp": "2024-01-01T00:00:00Z",
-        }
-
-        generator = create_mcp_error_response(
-            request_id="complex-req", code=-32602, message="Validation failed", data=complex_data
-        )
-
-        response_bytes = next(generator)
-        response_json = json.loads(response_bytes.decode("utf-8"))
-
-        assert response_json["error"]["data"] == complex_data
-        assert len(response_json["error"]["data"]["errors"]) == 2
-
-    def test_create_error_response_encoding(self):
-        """Test error response with non-ASCII characters."""
-        generator = create_mcp_error_response(
-            request_id="unicode-req",
-            code=-32603,
-            message="内部错误",  # Chinese characters
-            data={"details": "エラー詳細"},  # Japanese characters
-        )
-
-        response_bytes = next(generator)
-
-        # Should be valid UTF-8
-        response_str = response_bytes.decode("utf-8")
-        response_json = json.loads(response_str)
-
-        assert response_json["error"]["message"] == "内部错误"
-        assert response_json["error"]["data"]["details"] == "エラー詳細"
-
-    def test_create_error_response_yields_once(self):
-        """Test that error response generator yields exactly once."""
-        generator = create_mcp_error_response(request_id="test", code=-32600, message="Test")
-
-        # First yield should work
-        first_yield = next(generator)
-        assert isinstance(first_yield, bytes)
-
-        # Second yield should raise StopIteration
-        with pytest.raises(StopIteration):
-            next(generator)
-
-        # Subsequent calls should also raise
-        with pytest.raises(StopIteration):
-            next(generator)

api/tests/unit_tests/core/rag/pipeline/test_queue.py

@@ -0,0 +1,369 @@
+"""
+Unit tests for TenantSelfTaskQueue.
+
+These tests verify the Redis-based task queue functionality for tenant-specific
+task management with proper serialization and deserialization.
+"""
+
+import json
+from unittest.mock import MagicMock, patch
+from uuid import uuid4
+
+import pytest
+
+from core.rag.pipeline.queue import TASK_WRAPPER_PREFIX, TaskWrapper, TenantSelfTaskQueue
+
+
+class TestTaskWrapper:
+    """Test cases for TaskWrapper serialization/deserialization."""
+
+    def test_serialize_simple_data(self):
+        """Test serialization of simple data types."""
+        data = {"key": "value", "number": 42, "list": [1, 2, 3]}
+        wrapper = TaskWrapper(data)
+        
+        serialized = wrapper.serialize()
+        assert isinstance(serialized, str)
+        
+        # Verify it's valid JSON
+        parsed = json.loads(serialized)
+        assert parsed == data
+
+    def test_serialize_complex_data(self):
+        """Test serialization of complex nested data."""
+        data = {
+            "nested": {
+                "deep": {
+                    "value": "test",
+                    "numbers": [1, 2, 3, 4, 5]
+                }
+            },
+            "unicode": "测试中文",
+            "special_chars": "!@#$%^&*()"
+        }
+        wrapper = TaskWrapper(data)
+        
+        serialized = wrapper.serialize()
+        parsed = json.loads(serialized)
+        assert parsed == data
+
+    def test_deserialize_valid_data(self):
+        """Test deserialization of valid JSON data."""
+        original_data = {"key": "value", "number": 42}
+        serialized = json.dumps(original_data, ensure_ascii=False)
+        
+        wrapper = TaskWrapper.deserialize(serialized)
+        assert wrapper.data == original_data
+
+    def test_deserialize_invalid_json(self):
+        """Test deserialization handles invalid JSON gracefully."""
+        invalid_json = "{invalid json}"
+        
+        with pytest.raises(json.JSONDecodeError):
+            TaskWrapper.deserialize(invalid_json)
+
+    def test_serialize_ensure_ascii_false(self):
+        """Test that serialization preserves Unicode characters."""
+        data = {"chinese": "中文测试", "emoji": "🚀"}
+        wrapper = TaskWrapper(data)
+        
+        serialized = wrapper.serialize()
+        assert "中文测试" in serialized
+        assert "🚀" in serialized
+
+
+class TestTenantSelfTaskQueue:
+    """Test cases for TenantSelfTaskQueue functionality."""
+
+    @pytest.fixture
+    def mock_redis_client(self):
+        """Mock Redis client for testing."""
+        mock_redis = MagicMock()
+        return mock_redis
+
+    @pytest.fixture
+    def sample_queue(self, mock_redis_client):
+        """Create a sample TenantSelfTaskQueue instance."""
+        return TenantSelfTaskQueue("tenant-123", "test-key")
+
+    def test_initialization(self, sample_queue):
+        """Test queue initialization with correct key generation."""
+        assert sample_queue.tenant_id == "tenant-123"
+        assert sample_queue.unique_key == "test-key"
+        assert sample_queue.queue == "tenant_self_test-key_task_queue:tenant-123"
+        assert sample_queue.task_key == "tenant_test-key_task:tenant-123"
+        assert sample_queue.DEFAULT_TASK_TTL == 60 * 60
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_get_task_key_exists(self, mock_redis, sample_queue):
+        """Test getting task key when it exists."""
+        mock_redis.get.return_value = "1"
+        
+        result = sample_queue.get_task_key()
+        
+        assert result == "1"
+        mock_redis.get.assert_called_once_with("tenant_test-key_task:tenant-123")
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_get_task_key_not_exists(self, mock_redis, sample_queue):
+        """Test getting task key when it doesn't exist."""
+        mock_redis.get.return_value = None
+        
+        result = sample_queue.get_task_key()
+        
+        assert result is None
+        mock_redis.get.assert_called_once_with("tenant_test-key_task:tenant-123")
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_set_task_waiting_time_default_ttl(self, mock_redis, sample_queue):
+        """Test setting task waiting flag with default TTL."""
+        sample_queue.set_task_waiting_time()
+        
+        mock_redis.setex.assert_called_once_with(
+            "tenant_test-key_task:tenant-123", 
+            3600,  # DEFAULT_TASK_TTL
+            1
+        )
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_set_task_waiting_time_custom_ttl(self, mock_redis, sample_queue):
+        """Test setting task waiting flag with custom TTL."""
+        custom_ttl = 1800
+        sample_queue.set_task_waiting_time(custom_ttl)
+        
+        mock_redis.setex.assert_called_once_with(
+            "tenant_test-key_task:tenant-123", 
+            custom_ttl, 
+            1
+        )
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_delete_task_key(self, mock_redis, sample_queue):
+        """Test deleting task key."""
+        sample_queue.delete_task_key()
+        
+        mock_redis.delete.assert_called_once_with("tenant_test-key_task:tenant-123")
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_push_tasks_string_list(self, mock_redis, sample_queue):
+        """Test pushing string tasks directly."""
+        tasks = ["task1", "task2", "task3"]
+        
+        sample_queue.push_tasks(tasks)
+        
+        mock_redis.lpush.assert_called_once_with(
+            "tenant_self_test-key_task_queue:tenant-123",
+            "task1", "task2", "task3"
+        )
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_push_tasks_mixed_types(self, mock_redis, sample_queue):
+        """Test pushing mixed string and object tasks."""
+        tasks = [
+            "string_task",
+            {"object_task": "data", "id": 123},
+            "another_string"
+        ]
+        
+        sample_queue.push_tasks(tasks)
+        
+        # Verify lpush was called
+        mock_redis.lpush.assert_called_once()
+        call_args = mock_redis.lpush.call_args
+        
+        # Check queue name
+        assert call_args[0][0] == "tenant_self_test-key_task_queue:tenant-123"
+        
+        # Check serialized tasks
+        serialized_tasks = call_args[0][1:]
+        assert len(serialized_tasks) == 3
+        assert serialized_tasks[0] == "string_task"
+        assert serialized_tasks[2] == "another_string"
+        
+        # Check object task is wrapped
+        assert serialized_tasks[1].startswith(TASK_WRAPPER_PREFIX)
+        wrapper_data = serialized_tasks[1][len(TASK_WRAPPER_PREFIX):]
+        parsed_data = json.loads(wrapper_data)
+        assert parsed_data == {"object_task": "data", "id": 123}
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_push_tasks_empty_list(self, mock_redis, sample_queue):
+        """Test pushing empty task list."""
+        sample_queue.push_tasks([])
+        
+        mock_redis.lpush.assert_called_once_with(
+            "tenant_self_test-key_task_queue:tenant-123"
+        )
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_pull_tasks_default_count(self, mock_redis, sample_queue):
+        """Test pulling tasks with default count (1)."""
+        mock_redis.rpop.side_effect = ["task1", None]
+        
+        result = sample_queue.pull_tasks()
+        
+        assert result == ["task1"]
+        assert mock_redis.rpop.call_count == 1
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_pull_tasks_custom_count(self, mock_redis, sample_queue):
+        """Test pulling tasks with custom count."""
+        # First test: pull 3 tasks
+        mock_redis.rpop.side_effect = ["task1", "task2", "task3", None]
+        
+        result = sample_queue.pull_tasks(3)
+        
+        assert result == ["task1", "task2", "task3"]
+        assert mock_redis.rpop.call_count == 3
+
+        # Reset mock for second test
+        mock_redis.reset_mock()
+        mock_redis.rpop.side_effect = ["task1", "task2", None]
+        
+        result = sample_queue.pull_tasks(3)
+        
+        assert result == ["task1", "task2"]
+        assert mock_redis.rpop.call_count == 3
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_pull_tasks_zero_count(self, mock_redis, sample_queue):
+        """Test pulling tasks with zero count returns empty list."""
+        result = sample_queue.pull_tasks(0)
+        
+        assert result == []
+        mock_redis.rpop.assert_not_called()
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_pull_tasks_negative_count(self, mock_redis, sample_queue):
+        """Test pulling tasks with negative count returns empty list."""
+        result = sample_queue.pull_tasks(-1)
+        
+        assert result == []
+        mock_redis.rpop.assert_not_called()
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_pull_tasks_with_wrapped_objects(self, mock_redis, sample_queue):
+        """Test pulling tasks that include wrapped objects."""
+        # Create a wrapped task
+        wrapper = TaskWrapper({"task_id": 123, "data": "test"})
+        wrapped_task = f"{TASK_WRAPPER_PREFIX}{wrapper.serialize()}"
+        
+        mock_redis.rpop.side_effect = [
+            "string_task",
+            wrapped_task.encode('utf-8'),  # Simulate bytes from Redis
+            None
+        ]
+        
+        result = sample_queue.pull_tasks(2)
+        
+        assert len(result) == 2
+        assert result[0] == "string_task"
+        assert result[1] == {"task_id": 123, "data": "test"}
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_pull_tasks_with_invalid_wrapped_data(self, mock_redis, sample_queue):
+        """Test pulling tasks with invalid wrapped data falls back to string."""
+        invalid_wrapped = f"{TASK_WRAPPER_PREFIX}invalid json"
+        mock_redis.rpop.side_effect = [invalid_wrapped, None]
+        
+        result = sample_queue.pull_tasks(1)
+        
+        assert result == [invalid_wrapped]
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_pull_tasks_bytes_decoding(self, mock_redis, sample_queue):
+        """Test pulling tasks handles bytes from Redis correctly."""
+        mock_redis.rpop.side_effect = [
+            b"task1",  # bytes
+            "task2",   # string
+            None
+        ]
+        
+        result = sample_queue.pull_tasks(2)
+        
+        assert result == ["task1", "task2"]
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_get_next_task_success(self, mock_redis, sample_queue):
+        """Test getting next single task."""
+        mock_redis.rpop.side_effect = ["task1", None]
+        
+        result = sample_queue.get_next_task()
+        
+        assert result == "task1"
+        assert mock_redis.rpop.call_count == 1
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_get_next_task_empty_queue(self, mock_redis, sample_queue):
+        """Test getting next task when queue is empty."""
+        mock_redis.rpop.return_value = None
+        
+        result = sample_queue.get_next_task()
+        
+        assert result is None
+        mock_redis.rpop.assert_called_once()
+
+    def test_tenant_isolation(self):
+        """Test that different tenants have isolated queues."""
+        queue1 = TenantSelfTaskQueue("tenant-1", "key")
+        queue2 = TenantSelfTaskQueue("tenant-2", "key")
+        
+        assert queue1.queue != queue2.queue
+        assert queue1.task_key != queue2.task_key
+        assert queue1.queue == "tenant_self_key_task_queue:tenant-1"
+        assert queue2.queue == "tenant_self_key_task_queue:tenant-2"
+
+    def test_key_isolation(self):
+        """Test that different keys have isolated queues."""
+        queue1 = TenantSelfTaskQueue("tenant", "key1")
+        queue2 = TenantSelfTaskQueue("tenant", "key2")
+        
+        assert queue1.queue != queue2.queue
+        assert queue1.task_key != queue2.task_key
+        assert queue1.queue == "tenant_self_key1_task_queue:tenant"
+        assert queue2.queue == "tenant_self_key2_task_queue:tenant"
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_complex_object_serialization_roundtrip(self, mock_redis, sample_queue):
+        """Test complex object serialization and deserialization roundtrip."""
+        complex_task = {
+            "id": uuid4().hex,
+            "data": {
+                "nested": {
+                    "deep": [1, 2, 3],
+                    "unicode": "测试中文",
+                    "special": "!@#$%^&*()"
+                }
+            },
+            "metadata": {
+                "created_at": "2024-01-01T00:00:00Z",
+                "tags": ["tag1", "tag2", "tag3"]
+            }
+        }
+        
+        # Push the complex task
+        sample_queue.push_tasks([complex_task])
+        
+        # Verify it was wrapped
+        call_args = mock_redis.lpush.call_args
+        wrapped_task = call_args[0][1]
+        assert wrapped_task.startswith(TASK_WRAPPER_PREFIX)
+        
+        # Simulate pulling it back
+        mock_redis.rpop.return_value = wrapped_task
+        result = sample_queue.pull_tasks(1)
+        
+        assert len(result) == 1
+        assert result[0] == complex_task
+
+    @patch('core.rag.pipeline.queue.redis_client')
+    def test_large_task_list_handling(self, mock_redis, sample_queue):
+        """Test handling of large task lists."""
+        large_task_list = [f"task_{i}" for i in range(1000)]
+        
+        sample_queue.push_tasks(large_task_list)
+        
+        # Verify all tasks were pushed
+        call_args = mock_redis.lpush.call_args
+        assert len(call_args[0]) == 1001  # queue name + 1000 tasks
+        assert list(call_args[0][1:]) == large_task_list

api/tests/unit_tests/core/workflow/graph_engine/test_iteration_flatten_output.py

@@ -1,96 +0,0 @@
-"""
-Test cases for the Iteration node's flatten_output functionality.
-
-This module tests the iteration node's ability to:
-1. Flatten array outputs when flatten_output=True (default)
-2. Preserve nested array structure when flatten_output=False
-"""
-
-from .test_table_runner import TableTestRunner, WorkflowTestCase
-
-
-def test_iteration_with_flatten_output_enabled():
-    """
-    Test iteration node with flatten_output=True (default behavior).
-
-    The fixture implements an iteration that:
-    1. Iterates over [1, 2, 3]
-    2. For each item, outputs [item, item*2]
-    3. With flatten_output=True, should output [1, 2, 2, 4, 3, 6]
-    """
-    runner = TableTestRunner()
-
-    test_case = WorkflowTestCase(
-        fixture_path="iteration_flatten_output_enabled_workflow",
-        inputs={},
-        expected_outputs={"output": [1, 2, 2, 4, 3, 6]},
-        description="Iteration with flatten_output=True flattens nested arrays",
-        use_auto_mock=False,  # Run code nodes directly
-    )
-
-    result = runner.run_test_case(test_case)
-
-    assert result.success, f"Test failed: {result.error}"
-    assert result.actual_outputs is not None, "Should have outputs"
-    assert result.actual_outputs == {"output": [1, 2, 2, 4, 3, 6]}, (
-        f"Expected flattened output [1, 2, 2, 4, 3, 6], got {result.actual_outputs}"
-    )
-
-
-def test_iteration_with_flatten_output_disabled():
-    """
-    Test iteration node with flatten_output=False.
-
-    The fixture implements an iteration that:
-    1. Iterates over [1, 2, 3]
-    2. For each item, outputs [item, item*2]
-    3. With flatten_output=False, should output [[1, 2], [2, 4], [3, 6]]
-    """
-    runner = TableTestRunner()
-
-    test_case = WorkflowTestCase(
-        fixture_path="iteration_flatten_output_disabled_workflow",
-        inputs={},
-        expected_outputs={"output": [[1, 2], [2, 4], [3, 6]]},
-        description="Iteration with flatten_output=False preserves nested structure",
-        use_auto_mock=False,  # Run code nodes directly
-    )
-
-    result = runner.run_test_case(test_case)
-
-    assert result.success, f"Test failed: {result.error}"
-    assert result.actual_outputs is not None, "Should have outputs"
-    assert result.actual_outputs == {"output": [[1, 2], [2, 4], [3, 6]]}, (
-        f"Expected nested output [[1, 2], [2, 4], [3, 6]], got {result.actual_outputs}"
-    )
-
-
-def test_iteration_flatten_output_comparison():
-    """
-    Run both flatten_output configurations in parallel to verify the difference.
-    """
-    runner = TableTestRunner()
-
-    test_cases = [
-        WorkflowTestCase(
-            fixture_path="iteration_flatten_output_enabled_workflow",
-            inputs={},
-            expected_outputs={"output": [1, 2, 2, 4, 3, 6]},
-            description="flatten_output=True: Flattened output",
-            use_auto_mock=False,  # Run code nodes directly
-        ),
-        WorkflowTestCase(
-            fixture_path="iteration_flatten_output_disabled_workflow",
-            inputs={},
-            expected_outputs={"output": [[1, 2], [2, 4], [3, 6]]},
-            description="flatten_output=False: Nested output",
-            use_auto_mock=False,  # Run code nodes directly
-        ),
-    ]
-
-    suite_result = runner.run_table_tests(test_cases, parallel=True)
-
-    # Assert all tests passed
-    assert suite_result.passed_tests == 2, f"Expected 2 passed tests, got {suite_result.passed_tests}"
-    assert suite_result.failed_tests == 0, f"Expected 0 failed tests, got {suite_result.failed_tests}"
-    assert suite_result.success_rate == 100.0, f"Expected 100% success rate, got {suite_result.success_rate}"

api/tests/unit_tests/libs/test_token.py

@@ -1,10 +1,5 @@
-from unittest.mock import MagicMock
-
-from werkzeug.wrappers import Response
-
 from constants import COOKIE_NAME_ACCESS_TOKEN, COOKIE_NAME_WEBAPP_ACCESS_TOKEN
-from libs import token
-from libs.token import extract_access_token, extract_webapp_access_token, set_csrf_token_to_cookie
+from libs.token import extract_access_token, extract_webapp_access_token
 
 
 class MockRequest:
@@ -28,35 +23,3 @@ def test_extract_access_token():
     for request, expected_console, expected_webapp in test_cases:
         assert extract_access_token(request) == expected_console  # pyright: ignore[reportArgumentType]
         assert extract_webapp_access_token(request) == expected_webapp  # pyright: ignore[reportArgumentType]
-
-
-def test_real_cookie_name_uses_host_prefix_without_domain(monkeypatch):
-    monkeypatch.setattr(token.dify_config, "CONSOLE_WEB_URL", "https://console.example.com", raising=False)
-    monkeypatch.setattr(token.dify_config, "CONSOLE_API_URL", "https://api.example.com", raising=False)
-    monkeypatch.setattr(token.dify_config, "COOKIE_DOMAIN", "", raising=False)
-
-    assert token._real_cookie_name("csrf_token") == "__Host-csrf_token"
-
-
-def test_real_cookie_name_without_host_prefix_when_domain_present(monkeypatch):
-    monkeypatch.setattr(token.dify_config, "CONSOLE_WEB_URL", "https://console.example.com", raising=False)
-    monkeypatch.setattr(token.dify_config, "CONSOLE_API_URL", "https://api.example.com", raising=False)
-    monkeypatch.setattr(token.dify_config, "COOKIE_DOMAIN", ".example.com", raising=False)
-
-    assert token._real_cookie_name("csrf_token") == "csrf_token"
-
-
-def test_set_csrf_cookie_includes_domain_when_configured(monkeypatch):
-    monkeypatch.setattr(token.dify_config, "CONSOLE_WEB_URL", "https://console.example.com", raising=False)
-    monkeypatch.setattr(token.dify_config, "CONSOLE_API_URL", "https://api.example.com", raising=False)
-    monkeypatch.setattr(token.dify_config, "COOKIE_DOMAIN", ".example.com", raising=False)
-
-    response = Response()
-    request = MagicMock()
-
-    set_csrf_token_to_cookie(request, response, "abc123")
-
-    cookies = response.headers.getlist("Set-Cookie")
-    assert any("csrf_token=abc123" in c for c in cookies)
-    assert any("Domain=example.com" in c for c in cookies)
-    assert all("__Host-" not in c for c in cookies)

api/tests/unit_tests/services/test_document_indexing_task_proxy.py

@@ -0,0 +1,332 @@
+
+from unittest.mock import Mock, patch
+
+from core.entities.document_task import DocumentTask
+from core.rag.pipeline.queue import TenantSelfTaskQueue
+from services.document_indexing_task_proxy import DocumentIndexingTaskProxy
+
+
+class DocumentIndexingTaskProxyTestDataFactory:
+    """Factory class for creating test data and mock objects for DocumentIndexingTaskProxy tests."""
+
+    @staticmethod
+    def create_mock_features(
+        billing_enabled: bool = False,
+        plan: str = "sandbox"
+    ) -> Mock:
+        """Create mock features with billing configuration."""
+        features = Mock()
+        features.billing = Mock()
+        features.billing.enabled = billing_enabled
+        features.billing.subscription = Mock()
+        features.billing.subscription.plan = plan
+        return features
+
+    @staticmethod
+    def create_mock_tenant_queue(has_task_key: bool = False) -> Mock:
+        """Create mock TenantSelfTaskQueue."""
+        queue = Mock(spec=TenantSelfTaskQueue)
+        queue.get_task_key.return_value = "task_key" if has_task_key else None
+        queue.push_tasks = Mock()
+        queue.set_task_waiting_time = Mock()
+        return queue
+
+    @staticmethod
+    def create_document_task_proxy(
+        tenant_id: str = "tenant-123",
+        dataset_id: str = "dataset-456",
+        document_ids: list[str] | None = None
+    ) -> DocumentIndexingTaskProxy:
+        """Create DocumentIndexingTaskProxy instance for testing."""
+        if document_ids is None:
+            document_ids = ["doc-1", "doc-2", "doc-3"]
+        return DocumentIndexingTaskProxy(tenant_id, dataset_id, document_ids)
+
+
+class TestDocumentIndexingTaskProxy:
+    """Test cases for DocumentIndexingTaskProxy class."""
+
+    def test_initialization(self):
+        """Test DocumentIndexingTaskProxy initialization."""
+        # Arrange
+        tenant_id = "tenant-123"
+        dataset_id = "dataset-456"
+        document_ids = ["doc-1", "doc-2", "doc-3"]
+
+        # Act
+        proxy = DocumentIndexingTaskProxy(tenant_id, dataset_id, document_ids)
+
+        # Assert
+        assert proxy.tenant_id == tenant_id
+        assert proxy.dataset_id == dataset_id
+        assert proxy.document_ids == document_ids
+        assert isinstance(proxy.tenant_self_task_queue, TenantSelfTaskQueue)
+        assert proxy.tenant_self_task_queue.tenant_id == tenant_id
+        assert proxy.tenant_self_task_queue.unique_key == "document_indexing"
+
+    @patch('services.document_indexing_task_proxy.FeatureService')
+    def test_features_property(self, mock_feature_service):
+        """Test cached_property features."""
+        # Arrange
+        mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features()
+        mock_feature_service.get_features.return_value = mock_features
+        proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy()
+
+        # Act
+        features1 = proxy.features
+        features2 = proxy.features  # Second call should use cached property
+
+        # Assert
+        assert features1 == mock_features
+        assert features2 == mock_features
+        assert features1 is features2  # Should be the same instance due to caching
+        mock_feature_service.get_features.assert_called_once_with("tenant-123")
+
+    @patch('services.document_indexing_task_proxy.normal_document_indexing_task')
+    def test_send_to_direct_queue(self, mock_task):
+        """Test _send_to_direct_queue method."""
+        # Arrange
+        proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy()
+        mock_task.delay = Mock()
+
+        # Act
+        proxy._send_to_direct_queue(mock_task)
+
+        # Assert
+        mock_task.delay.assert_called_once_with(
+            tenant_id="tenant-123",
+            dataset_id="dataset-456",
+            document_ids=["doc-1", "doc-2", "doc-3"]
+        )
+
+    @patch('services.document_indexing_task_proxy.normal_document_indexing_task')
+    def test_send_to_tenant_queue_with_existing_task_key(self, mock_task):
+        """Test _send_to_tenant_queue when task key exists."""
+        # Arrange
+        proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy()
+        proxy.tenant_self_task_queue = DocumentIndexingTaskProxyTestDataFactory.create_mock_tenant_queue(
+            has_task_key=True
+        )
+        mock_task.delay = Mock()
+
+        # Act
+        proxy._send_to_tenant_queue(mock_task)
+
+        # Assert
+        proxy.tenant_self_task_queue.push_tasks.assert_called_once()
+        pushed_tasks = proxy.tenant_self_task_queue.push_tasks.call_args[0][0]
+        assert len(pushed_tasks) == 1
+        assert isinstance(DocumentTask(**pushed_tasks[0]), DocumentTask)
+        assert pushed_tasks[0]["tenant_id"] == "tenant-123"
+        assert pushed_tasks[0]["dataset_id"] == "dataset-456"
+        assert pushed_tasks[0]["document_ids"] == ["doc-1", "doc-2", "doc-3"]
+        mock_task.delay.assert_not_called()
+
+    @patch('services.document_indexing_task_proxy.normal_document_indexing_task')
+    def test_send_to_tenant_queue_without_task_key(self, mock_task):
+        """Test _send_to_tenant_queue when no task key exists."""
+        # Arrange
+        proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy()
+        proxy.tenant_self_task_queue = DocumentIndexingTaskProxyTestDataFactory.create_mock_tenant_queue(
+            has_task_key=False
+        )
+        mock_task.delay = Mock()
+
+        # Act
+        proxy._send_to_tenant_queue(mock_task)
+
+        # Assert
+        proxy.tenant_self_task_queue.set_task_waiting_time.assert_called_once()
+        mock_task.delay.assert_called_once_with(
+            tenant_id="tenant-123",
+            dataset_id="dataset-456",
+            document_ids=["doc-1", "doc-2", "doc-3"]
+        )
+        proxy.tenant_self_task_queue.push_tasks.assert_not_called()
+
+    @patch('services.document_indexing_task_proxy.normal_document_indexing_task')
+    def test_send_to_default_tenant_queue(self, mock_task):
+        """Test _send_to_default_tenant_queue method."""
+        # Arrange
+        proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy()
+        proxy._send_to_tenant_queue = Mock()
+
+        # Act
+        proxy._send_to_default_tenant_queue()
+
+        # Assert
+        proxy._send_to_tenant_queue.assert_called_once_with(mock_task)
+
+    @patch('services.document_indexing_task_proxy.priority_document_indexing_task')
+    def test_send_to_priority_tenant_queue(self, mock_task):
+        """Test _send_to_priority_tenant_queue method."""
+        # Arrange
+        proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy()
+        proxy._send_to_tenant_queue = Mock()
+
+        # Act
+        proxy._send_to_priority_tenant_queue()
+
+        # Assert
+        proxy._send_to_tenant_queue.assert_called_once_with(mock_task)
+
+    @patch('services.document_indexing_task_proxy.priority_document_indexing_task')
+    def test_send_to_priority_direct_queue(self, mock_task):
+        """Test _send_to_priority_direct_queue method."""
+        # Arrange
+        proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy()
+        proxy._send_to_direct_queue = Mock()
+
+        # Act
+        proxy._send_to_priority_direct_queue()
+
+        # Assert
+        proxy._send_to_direct_queue.assert_called_once_with(mock_task)
+
+    @patch('services.document_indexing_task_proxy.FeatureService')
+    def test_dispatch_with_billing_enabled_sandbox_plan(self, mock_feature_service):
+        """Test _dispatch method when billing is enabled with sandbox plan."""
+        # Arrange
+        mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features(
+            billing_enabled=True, plan="sandbox"
+        )
+        mock_feature_service.get_features.return_value = mock_features
+        proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy()
+        proxy._send_to_default_tenant_queue = Mock()
+
+        # Act
+        proxy._dispatch()
+
+        # Assert
+        proxy._send_to_default_tenant_queue.assert_called_once()
+
+    @patch('services.document_indexing_task_proxy.FeatureService')
+    def test_dispatch_with_billing_enabled_non_sandbox_plan(self, mock_feature_service):
+        """Test _dispatch method when billing is enabled with non-sandbox plan."""
+        # Arrange
+        mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features(
+            billing_enabled=True, plan="team"
+        )
+        mock_feature_service.get_features.return_value = mock_features
+        proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy()
+        proxy._send_to_priority_tenant_queue = Mock()
+
+        # Act
+        proxy._dispatch()
+
+        # If billing enabled with non sandbox plan, should send to priority tenant queue
+        proxy._send_to_priority_tenant_queue.assert_called_once()
+
+    @patch('services.document_indexing_task_proxy.FeatureService')
+    def test_dispatch_with_billing_disabled(self, mock_feature_service):
+        """Test _dispatch method when billing is disabled."""
+        # Arrange
+        mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features(
+            billing_enabled=False
+        )
+        mock_feature_service.get_features.return_value = mock_features
+        proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy()
+        proxy._send_to_priority_direct_queue = Mock()
+
+        # Act
+        proxy._dispatch()
+
+        # If billing disabled, for example: self-hosted or enterprise, should send to priority direct queue
+        proxy._send_to_priority_direct_queue.assert_called_once()
+
+    @patch('services.document_indexing_task_proxy.FeatureService')
+    def test_delay_method(self, mock_feature_service):
+        """Test delay method integration."""
+        # Arrange
+        mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features(
+            billing_enabled=True, plan="sandbox"
+        )
+        mock_feature_service.get_features.return_value = mock_features
+        proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy()
+        proxy._send_to_default_tenant_queue = Mock()
+
+        # Act
+        proxy.delay()
+
+        # Assert
+        # If billing enabled with sandbox plan, should send to default tenant queue
+        proxy._send_to_default_tenant_queue.assert_called_once()
+
+    def test_document_task_dataclass(self):
+        """Test DocumentTask dataclass."""
+        # Arrange
+        tenant_id = "tenant-123"
+        dataset_id = "dataset-456"
+        document_ids = ["doc-1", "doc-2"]
+
+        # Act
+        task = DocumentTask(tenant_id=tenant_id, dataset_id=dataset_id, document_ids=document_ids)
+
+        # Assert
+        assert task.tenant_id == tenant_id
+        assert task.dataset_id == dataset_id
+        assert task.document_ids == document_ids
+
+    @patch('services.document_indexing_task_proxy.FeatureService')
+    def test_dispatch_edge_case_empty_plan(self, mock_feature_service):
+        """Test _dispatch method with empty plan string."""
+        # Arrange
+        mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features(
+            billing_enabled=True, plan=""
+        )
+        mock_feature_service.get_features.return_value = mock_features
+        proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy()
+        proxy._send_to_priority_tenant_queue = Mock()
+
+        # Act
+        proxy._dispatch()
+
+        # Assert
+        proxy._send_to_priority_tenant_queue.assert_called_once()
+
+    @patch('services.document_indexing_task_proxy.FeatureService')
+    def test_dispatch_edge_case_none_plan(self, mock_feature_service):
+        """Test _dispatch method with None plan."""
+        # Arrange
+        mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features(
+            billing_enabled=True, plan=None
+        )
+        mock_feature_service.get_features.return_value = mock_features
+        proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy()
+        proxy._send_to_priority_tenant_queue = Mock()
+
+        # Act
+        proxy._dispatch()
+
+        # Assert
+        proxy._send_to_priority_tenant_queue.assert_called_once()
+
+    def test_initialization_with_empty_document_ids(self):
+        """Test initialization with empty document_ids list."""
+        # Arrange
+        tenant_id = "tenant-123"
+        dataset_id = "dataset-456"
+        document_ids = []
+
+        # Act
+        proxy = DocumentIndexingTaskProxy(tenant_id, dataset_id, document_ids)
+
+        # Assert
+        assert proxy.tenant_id == tenant_id
+        assert proxy.dataset_id == dataset_id
+        assert proxy.document_ids == document_ids
+
+    def test_initialization_with_single_document_id(self):
+        """Test initialization with single document_id."""
+        # Arrange
+        tenant_id = "tenant-123"
+        dataset_id = "dataset-456"
+        document_ids = ["doc-1"]
+
+        # Act
+        proxy = DocumentIndexingTaskProxy(tenant_id, dataset_id, document_ids)
+
+        # Assert
+        assert proxy.tenant_id == tenant_id
+        assert proxy.dataset_id == dataset_id
+        assert proxy.document_ids == document_ids

api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py

@@ -0,0 +1,499 @@
+import json
+from unittest.mock import Mock, patch
+
+import pytest
+
+from core.app.entities.rag_pipeline_invoke_entities import RagPipelineInvokeEntity
+from core.rag.pipeline.queue import TenantSelfTaskQueue
+from services.rag_pipeline.rag_pipeline_task_proxy import RagPipelineTaskProxy
+
+
+class RagPipelineTaskProxyTestDataFactory:
+    """Factory class for creating test data and mock objects for RagPipelineTaskProxy tests."""
+
+    @staticmethod
+    def create_mock_features(
+        billing_enabled: bool = False,
+        plan: str = "sandbox"
+    ) -> Mock:
+        """Create mock features with billing configuration."""
+        features = Mock()
+        features.billing = Mock()
+        features.billing.enabled = billing_enabled
+        features.billing.subscription = Mock()
+        features.billing.subscription.plan = plan
+        return features
+
+    @staticmethod
+    def create_mock_tenant_queue(has_task_key: bool = False) -> Mock:
+        """Create mock TenantSelfTaskQueue."""
+        queue = Mock(spec=TenantSelfTaskQueue)
+        queue.get_task_key.return_value = "task_key" if has_task_key else None
+        queue.push_tasks = Mock()
+        queue.set_task_waiting_time = Mock()
+        return queue
+
+    @staticmethod
+    def create_rag_pipeline_invoke_entity(
+        pipeline_id: str = "pipeline-123",
+        user_id: str = "user-456",
+        tenant_id: str = "tenant-789",
+        workflow_id: str = "workflow-101",
+        streaming: bool = True,
+        workflow_execution_id: str | None = None,
+        workflow_thread_pool_id: str | None = None
+    ) -> RagPipelineInvokeEntity:
+        """Create RagPipelineInvokeEntity instance for testing."""
+        return RagPipelineInvokeEntity(
+            pipeline_id=pipeline_id,
+            application_generate_entity={"key": "value"},
+            user_id=user_id,
+            tenant_id=tenant_id,
+            workflow_id=workflow_id,
+            streaming=streaming,
+            workflow_execution_id=workflow_execution_id,
+            workflow_thread_pool_id=workflow_thread_pool_id
+        )
+
+    @staticmethod
+    def create_rag_pipeline_task_proxy(
+        dataset_tenant_id: str = "tenant-123",
+        user_id: str = "user-456",
+        rag_pipeline_invoke_entities: list[RagPipelineInvokeEntity] | None = None
+    ) -> RagPipelineTaskProxy:
+        """Create RagPipelineTaskProxy instance for testing."""
+        if rag_pipeline_invoke_entities is None:
+            rag_pipeline_invoke_entities = [
+                RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity()
+            ]
+        return RagPipelineTaskProxy(dataset_tenant_id, user_id, rag_pipeline_invoke_entities)
+
+    @staticmethod
+    def create_mock_upload_file(file_id: str = "file-123") -> Mock:
+        """Create mock upload file."""
+        upload_file = Mock()
+        upload_file.id = file_id
+        return upload_file
+
+
+class TestRagPipelineTaskProxy:
+    """Test cases for RagPipelineTaskProxy class."""
+
+    def test_initialization(self):
+        """Test RagPipelineTaskProxy initialization."""
+        # Arrange
+        dataset_tenant_id = "tenant-123"
+        user_id = "user-456"
+        rag_pipeline_invoke_entities = [
+            RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity()
+        ]
+
+        # Act
+        proxy = RagPipelineTaskProxy(dataset_tenant_id, user_id, rag_pipeline_invoke_entities)
+
+        # Assert
+        assert proxy.dataset_tenant_id == dataset_tenant_id
+        assert proxy.user_id == user_id
+        assert proxy.rag_pipeline_invoke_entities == rag_pipeline_invoke_entities
+        assert isinstance(proxy.tenant_self_pipeline_task_queue, TenantSelfTaskQueue)
+        assert proxy.tenant_self_pipeline_task_queue.tenant_id == dataset_tenant_id
+        assert proxy.tenant_self_pipeline_task_queue.unique_key == "pipeline"
+
+    def test_initialization_with_empty_entities(self):
+        """Test initialization with empty rag_pipeline_invoke_entities."""
+        # Arrange
+        dataset_tenant_id = "tenant-123"
+        user_id = "user-456"
+        rag_pipeline_invoke_entities = []
+
+        # Act
+        proxy = RagPipelineTaskProxy(dataset_tenant_id, user_id, rag_pipeline_invoke_entities)
+
+        # Assert
+        assert proxy.dataset_tenant_id == dataset_tenant_id
+        assert proxy.user_id == user_id
+        assert proxy.rag_pipeline_invoke_entities == []
+
+    def test_initialization_with_multiple_entities(self):
+        """Test initialization with multiple rag_pipeline_invoke_entities."""
+        # Arrange
+        dataset_tenant_id = "tenant-123"
+        user_id = "user-456"
+        rag_pipeline_invoke_entities = [
+            RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity(pipeline_id="pipeline-1"),
+            RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity(pipeline_id="pipeline-2"),
+            RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity(pipeline_id="pipeline-3")
+        ]
+
+        # Act
+        proxy = RagPipelineTaskProxy(dataset_tenant_id, user_id, rag_pipeline_invoke_entities)
+
+        # Assert
+        assert len(proxy.rag_pipeline_invoke_entities) == 3
+        assert proxy.rag_pipeline_invoke_entities[0].pipeline_id == "pipeline-1"
+        assert proxy.rag_pipeline_invoke_entities[1].pipeline_id == "pipeline-2"
+        assert proxy.rag_pipeline_invoke_entities[2].pipeline_id == "pipeline-3"
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FeatureService')
+    def test_features_property(self, mock_feature_service):
+        """Test cached_property features."""
+        # Arrange
+        mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features()
+        mock_feature_service.get_features.return_value = mock_features
+        proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy()
+
+        # Act
+        features1 = proxy.features
+        features2 = proxy.features  # Second call should use cached property
+
+        # Assert
+        assert features1 == mock_features
+        assert features2 == mock_features
+        assert features1 is features2  # Should be the same instance due to caching
+        mock_feature_service.get_features.assert_called_once_with("tenant-123")
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FileService')
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.db')
+    def test_upload_invoke_entities(self, mock_db, mock_file_service_class):
+        """Test _upload_invoke_entities method."""
+        # Arrange
+        proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy()
+        mock_file_service = Mock()
+        mock_file_service_class.return_value = mock_file_service
+        mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_mock_upload_file("file-123")
+        mock_file_service.upload_text.return_value = mock_upload_file
+
+        # Act
+        result = proxy._upload_invoke_entities()
+
+        # Assert
+        assert result == "file-123"
+        mock_file_service_class.assert_called_once_with(mock_db.engine)
+        
+        # Verify upload_text was called with correct parameters
+        mock_file_service.upload_text.assert_called_once()
+        call_args = mock_file_service.upload_text.call_args
+        json_text, name, user_id, tenant_id = call_args[0]
+        
+        assert name == "rag_pipeline_invoke_entities.json"
+        assert user_id == "user-456"
+        assert tenant_id == "tenant-123"
+        
+        # Verify JSON content
+        parsed_json = json.loads(json_text)
+        assert len(parsed_json) == 1
+        assert parsed_json[0]["pipeline_id"] == "pipeline-123"
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FileService')
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.db')
+    def test_upload_invoke_entities_with_multiple_entities(self, mock_db, mock_file_service_class):
+        """Test _upload_invoke_entities method with multiple entities."""
+        # Arrange
+        entities = [
+            RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity(pipeline_id="pipeline-1"),
+            RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity(pipeline_id="pipeline-2")
+        ]
+        proxy = RagPipelineTaskProxy("tenant-123", "user-456", entities)
+        mock_file_service = Mock()
+        mock_file_service_class.return_value = mock_file_service
+        mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_mock_upload_file("file-456")
+        mock_file_service.upload_text.return_value = mock_upload_file
+
+        # Act
+        result = proxy._upload_invoke_entities()
+
+        # Assert
+        assert result == "file-456"
+        
+        # Verify JSON content contains both entities
+        call_args = mock_file_service.upload_text.call_args
+        json_text = call_args[0][0]
+        parsed_json = json.loads(json_text)
+        assert len(parsed_json) == 2
+        assert parsed_json[0]["pipeline_id"] == "pipeline-1"
+        assert parsed_json[1]["pipeline_id"] == "pipeline-2"
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.rag_pipeline_run_task')
+    def test_send_to_direct_queue(self, mock_task):
+        """Test _send_to_direct_queue method."""
+        # Arrange
+        proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy()
+        proxy.tenant_self_pipeline_task_queue = RagPipelineTaskProxyTestDataFactory.create_mock_tenant_queue()
+        upload_file_id = "file-123"
+        mock_task.delay = Mock()
+
+        # Act
+        proxy._send_to_direct_queue(upload_file_id, mock_task)
+
+        # If sent to direct queue, tenant_self_pipeline_task_queue should not be called
+        proxy.tenant_self_pipeline_task_queue.push_tasks.assert_not_called()
+
+        # Celery should be called directly
+        mock_task.delay.assert_called_once_with(
+            rag_pipeline_invoke_entities_file_id=upload_file_id,
+            tenant_id="tenant-123"
+        )
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.rag_pipeline_run_task')
+    def test_send_to_tenant_queue_with_existing_task_key(self, mock_task):
+        """Test _send_to_tenant_queue when task key exists."""
+        # Arrange
+        proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy()
+        proxy.tenant_self_pipeline_task_queue = RagPipelineTaskProxyTestDataFactory.create_mock_tenant_queue(
+            has_task_key=True
+        )
+        upload_file_id = "file-123"
+        mock_task.delay = Mock()
+
+        # Act
+        proxy._send_to_tenant_queue(upload_file_id, mock_task)
+
+        # If task key exists, should push tasks to the queue
+        proxy.tenant_self_pipeline_task_queue.push_tasks.assert_called_once_with([upload_file_id])
+        # Celery should not be called directly
+        mock_task.delay.assert_not_called()
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.rag_pipeline_run_task')
+    def test_send_to_tenant_queue_without_task_key(self, mock_task):
+        """Test _send_to_tenant_queue when no task key exists."""
+        # Arrange
+        proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy()
+        proxy.tenant_self_pipeline_task_queue = RagPipelineTaskProxyTestDataFactory.create_mock_tenant_queue(
+            has_task_key=False
+        )
+        upload_file_id = "file-123"
+        mock_task.delay = Mock()
+
+        # Act
+        proxy._send_to_tenant_queue(upload_file_id, mock_task)
+
+        # If no task key, should set task waiting time key first
+        proxy.tenant_self_pipeline_task_queue.set_task_waiting_time.assert_called_once()
+        mock_task.delay.assert_called_once_with(
+            rag_pipeline_invoke_entities_file_id=upload_file_id,
+            tenant_id="tenant-123"
+        )
+
+        # The first task should be sent to celery directly, so push tasks should not be called
+        proxy.tenant_self_pipeline_task_queue.push_tasks.assert_not_called()
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.rag_pipeline_run_task')
+    def test_send_to_default_tenant_queue(self, mock_task):
+        """Test _send_to_default_tenant_queue method."""
+        # Arrange
+        proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy()
+        proxy._send_to_tenant_queue = Mock()
+        upload_file_id = "file-123"
+
+        # Act
+        proxy._send_to_default_tenant_queue(upload_file_id)
+
+        # Assert
+        proxy._send_to_tenant_queue.assert_called_once_with(upload_file_id, mock_task)
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.priority_rag_pipeline_run_task')
+    def test_send_to_priority_tenant_queue(self, mock_task):
+        """Test _send_to_priority_tenant_queue method."""
+        # Arrange
+        proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy()
+        proxy._send_to_tenant_queue = Mock()
+        upload_file_id = "file-123"
+
+        # Act
+        proxy._send_to_priority_tenant_queue(upload_file_id)
+
+        # Assert
+        proxy._send_to_tenant_queue.assert_called_once_with(upload_file_id, mock_task)
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.priority_rag_pipeline_run_task')
+    def test_send_to_priority_direct_queue(self, mock_task):
+        """Test _send_to_priority_direct_queue method."""
+        # Arrange
+        proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy()
+        proxy._send_to_direct_queue = Mock()
+        upload_file_id = "file-123"
+
+        # Act
+        proxy._send_to_priority_direct_queue(upload_file_id)
+
+        # Assert
+        proxy._send_to_direct_queue.assert_called_once_with(upload_file_id, mock_task)
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FeatureService')
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FileService')
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.db')
+    def test_dispatch_with_billing_enabled_sandbox_plan(self, mock_db, mock_file_service_class, mock_feature_service):
+        """Test _dispatch method when billing is enabled with sandbox plan."""
+        # Arrange
+        mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(
+            billing_enabled=True, plan="sandbox"
+        )
+        mock_feature_service.get_features.return_value = mock_features
+        proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy()
+        proxy._send_to_default_tenant_queue = Mock()
+        
+        mock_file_service = Mock()
+        mock_file_service_class.return_value = mock_file_service
+        mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_mock_upload_file("file-123")
+        mock_file_service.upload_text.return_value = mock_upload_file
+
+        # Act
+        proxy._dispatch()
+
+        # If billing is enabled with sandbox plan, should send to default tenant queue
+        proxy._send_to_default_tenant_queue.assert_called_once_with("file-123")
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FeatureService')
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FileService')
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.db')
+    def test_dispatch_with_billing_enabled_non_sandbox_plan(
+        self, mock_db, mock_file_service_class, mock_feature_service
+    ):
+        """Test _dispatch method when billing is enabled with non-sandbox plan."""
+        # Arrange
+        mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(
+            billing_enabled=True, plan="team"
+        )
+        mock_feature_service.get_features.return_value = mock_features
+        proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy()
+        proxy._send_to_priority_tenant_queue = Mock()
+        
+        mock_file_service = Mock()
+        mock_file_service_class.return_value = mock_file_service
+        mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_mock_upload_file("file-123")
+        mock_file_service.upload_text.return_value = mock_upload_file
+
+        # Act
+        proxy._dispatch()
+
+        # If billing is enabled with non-sandbox plan, should send to priority tenant queue
+        proxy._send_to_priority_tenant_queue.assert_called_once_with("file-123")
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FeatureService')
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FileService')
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.db')
+    def test_dispatch_with_billing_disabled(self, mock_db, mock_file_service_class, mock_feature_service):
+        """Test _dispatch method when billing is disabled."""
+        # Arrange
+        mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(
+            billing_enabled=False
+        )
+        mock_feature_service.get_features.return_value = mock_features
+        proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy()
+        proxy._send_to_priority_direct_queue = Mock()
+        
+        mock_file_service = Mock()
+        mock_file_service_class.return_value = mock_file_service
+        mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_mock_upload_file("file-123")
+        mock_file_service.upload_text.return_value = mock_upload_file
+
+        # Act
+        proxy._dispatch()
+
+        # If billing is disabled, for example: self-hosted or enterprise, should send to priority direct queue
+        proxy._send_to_priority_direct_queue.assert_called_once_with("file-123")
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FileService')
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.db')
+    def test_dispatch_with_empty_upload_file_id(self, mock_db, mock_file_service_class):
+        """Test _dispatch method when upload_file_id is empty."""
+        # Arrange
+        proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy()
+        
+        mock_file_service = Mock()
+        mock_file_service_class.return_value = mock_file_service
+        mock_upload_file = Mock()
+        mock_upload_file.id = ""  # Empty file ID
+        mock_file_service.upload_text.return_value = mock_upload_file
+
+        # Act & Assert
+        with pytest.raises(ValueError, match="upload_file_id is empty"):
+            proxy._dispatch()
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FeatureService')
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FileService')
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.db')
+    def test_dispatch_edge_case_empty_plan(self, mock_db, mock_file_service_class, mock_feature_service):
+        """Test _dispatch method with empty plan string."""
+        # Arrange
+        mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(
+            billing_enabled=True, plan=""
+        )
+        mock_feature_service.get_features.return_value = mock_features
+        proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy()
+        proxy._send_to_priority_tenant_queue = Mock()
+        
+        mock_file_service = Mock()
+        mock_file_service_class.return_value = mock_file_service
+        mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_mock_upload_file("file-123")
+        mock_file_service.upload_text.return_value = mock_upload_file
+
+        # Act
+        proxy._dispatch()
+
+        # Assert
+        proxy._send_to_priority_tenant_queue.assert_called_once_with("file-123")
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FeatureService')
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FileService')
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.db')
+    def test_dispatch_edge_case_none_plan(self, mock_db, mock_file_service_class, mock_feature_service):
+        """Test _dispatch method with None plan."""
+        # Arrange
+        mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(
+            billing_enabled=True, plan=None
+        )
+        mock_feature_service.get_features.return_value = mock_features
+        proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy()
+        proxy._send_to_priority_tenant_queue = Mock()
+        
+        mock_file_service = Mock()
+        mock_file_service_class.return_value = mock_file_service
+        mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_mock_upload_file("file-123")
+        mock_file_service.upload_text.return_value = mock_upload_file
+
+        # Act
+        proxy._dispatch()
+
+        # Assert
+        proxy._send_to_priority_tenant_queue.assert_called_once_with("file-123")
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FeatureService')
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.FileService')
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.db')
+    def test_delay_method(self, mock_db, mock_file_service_class, mock_feature_service):
+        """Test delay method integration."""
+        # Arrange
+        mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(
+            billing_enabled=True, plan="sandbox"
+        )
+        mock_feature_service.get_features.return_value = mock_features
+        proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy()
+        proxy._dispatch = Mock()
+        
+        mock_file_service = Mock()
+        mock_file_service_class.return_value = mock_file_service
+        mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_mock_upload_file("file-123")
+        mock_file_service.upload_text.return_value = mock_upload_file
+
+        # Act
+        proxy.delay()
+
+        # Assert
+        proxy._dispatch.assert_called_once()
+
+    @patch('services.rag_pipeline.rag_pipeline_task_proxy.logger')
+    def test_delay_method_with_empty_entities(self, mock_logger):
+        """Test delay method with empty rag_pipeline_invoke_entities."""
+        # Arrange
+        proxy = RagPipelineTaskProxy("tenant-123", "user-456", [])
+
+        # Act
+        proxy.delay()
+
+        # Assert
+        mock_logger.warning.assert_called_once_with(
+            "Received empty rag pipeline invoke entities, no tasks delivered: %s %s",
+            "tenant-123",
+            "user-456"
+        )

api/tests/unit_tests/services/test_variable_truncator.py

@@ -21,7 +21,6 @@ from core.file.enums import FileTransferMethod, FileType
 from core.file.models import File
 from core.variables.segments import (
     ArrayFileSegment,
-    ArrayNumberSegment,
     ArraySegment,
     FileSegment,
     FloatSegment,
@@ -31,7 +30,6 @@ from core.variables.segments import (
     StringSegment,
 )
 from services.variable_truncator import (
-    DummyVariableTruncator,
     MaxDepthExceededError,
     TruncationResult,
     UnknownTypeError,
@@ -598,32 +596,3 @@ class TestIntegrationScenarios:
         truncated_mapping, truncated = truncator.truncate_variable_mapping(mapping)
         assert truncated is False
         assert truncated_mapping == mapping
-
-
-def test_dummy_variable_truncator_methods():
-    """Test DummyVariableTruncator methods work correctly."""
-    truncator = DummyVariableTruncator()
-
-    # Test truncate_variable_mapping
-    test_data: dict[str, Any] = {
-        "key1": "value1",
-        "key2": ["item1", "item2"],
-        "large_array": list(range(2000)),
-    }
-    result, is_truncated = truncator.truncate_variable_mapping(test_data)
-
-    assert result == test_data
-    assert not is_truncated
-
-    # Test truncate method
-    segment = StringSegment(value="test string")
-    result = truncator.truncate(segment)
-    assert isinstance(result, TruncationResult)
-    assert result.result == segment
-    assert result.truncated is False
-
-    segment = ArrayNumberSegment(value=list(range(2000)))
-    result = truncator.truncate(segment)
-    assert isinstance(result, TruncationResult)
-    assert result.result == segment
-    assert result.truncated is False

api/tests/unit_tests/services/tools/test_mcp_tools_transform.py

@@ -180,25 +180,6 @@ class TestMCPToolTransform:
         # Set tools data with null description
         mock_provider_full.tools = '[{"name": "tool1", "description": null, "inputSchema": {}}]'
 
-        # Mock the to_entity and to_api_response methods
-        mock_entity = Mock()
-        mock_entity.to_api_response.return_value = {
-            "name": "Test MCP Provider",
-            "type": ToolProviderType.MCP,
-            "is_team_authorization": True,
-            "server_url": "https://*****.com/mcp",
-            "provider_icon": "icon.png",
-            "masked_headers": {"Authorization": "Bearer *****"},
-            "updated_at": 1234567890,
-            "labels": [],
-            "author": "Test User",
-            "description": I18nObject(en_US="Test MCP Provider Description", zh_Hans="Test MCP Provider Description"),
-            "icon": "icon.png",
-            "label": I18nObject(en_US="Test MCP Provider", zh_Hans="Test MCP Provider"),
-            "masked_credentials": {},
-        }
-        mock_provider_full.to_entity.return_value = mock_entity
-
         # Call the method with for_list=True
         result = ToolTransformService.mcp_provider_to_user_provider(mock_provider_full, for_list=True)
 
@@ -217,27 +198,6 @@ class TestMCPToolTransform:
         # Set tools data with description
         mock_provider_full.tools = '[{"name": "tool1", "description": "Tool description", "inputSchema": {}}]'
 
-        # Mock the to_entity and to_api_response methods
-        mock_entity = Mock()
-        mock_entity.to_api_response.return_value = {
-            "name": "Test MCP Provider",
-            "type": ToolProviderType.MCP,
-            "is_team_authorization": True,
-            "server_url": "https://*****.com/mcp",
-            "provider_icon": "icon.png",
-            "masked_headers": {"Authorization": "Bearer *****"},
-            "updated_at": 1234567890,
-            "labels": [],
-            "configuration": {"timeout": "30", "sse_read_timeout": "300"},
-            "original_headers": {"Authorization": "Bearer secret-token"},
-            "author": "Test User",
-            "description": I18nObject(en_US="Test MCP Provider Description", zh_Hans="Test MCP Provider Description"),
-            "icon": "icon.png",
-            "label": I18nObject(en_US="Test MCP Provider", zh_Hans="Test MCP Provider"),
-            "masked_credentials": {},
-        }
-        mock_provider_full.to_entity.return_value = mock_entity
-
         # Call the method with for_list=False
         result = ToolTransformService.mcp_provider_to_user_provider(mock_provider_full, for_list=False)
 
@@ -245,9 +205,8 @@ class TestMCPToolTransform:
         assert isinstance(result, ToolProviderApiEntity)
         assert result.id == "server-identifier-456"  # Should use server_identifier when for_list=False
         assert result.server_identifier == "server-identifier-456"
-        assert result.configuration is not None
-        assert result.configuration.timeout == 30
-        assert result.configuration.sse_read_timeout == 300
+        assert result.timeout == 30
+        assert result.sse_read_timeout == 300
         assert result.original_headers == {"Authorization": "Bearer secret-token"}
         assert len(result.tools) == 1
         assert result.tools[0].description.en_US == "Tool description"

api/uv.lock

@@ -1590,7 +1590,7 @@ dev = [
     { name = "ruff", specifier = "~=0.14.0" },
     { name = "scipy-stubs", specifier = ">=1.15.3.0" },
     { name = "sseclient-py", specifier = ">=1.8.0" },
-    { name = "testcontainers", specifier = "~=4.13.2" },
+    { name = "testcontainers", specifier = "~=4.10.0" },
     { name = "ty", specifier = "~=0.0.1a19" },
     { name = "types-aiofiles", specifier = "~=24.1.0" },
     { name = "types-beautifulsoup4", specifier = "~=4.12.0" },
@@ -5907,7 +5907,7 @@ wheels = [
 
 [[package]]
 name = "testcontainers"
-version = "4.13.2"
+version = "4.10.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "docker" },
@@ -5916,9 +5916,9 @@ dependencies = [
     { name = "urllib3" },
     { name = "wrapt" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/18/51/edac83edab339d8b4dce9a7b659163afb1ea7e011bfed1d5573d495a4485/testcontainers-4.13.2.tar.gz", hash = "sha256:2315f1e21b059427a9d11e8921f85fef322fbe0d50749bcca4eaa11271708ba4", size = 78692, upload-time = "2025-10-07T21:53:07.531Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/a1/49/9c618aff1c50121d183cdfbc3a4a5cf2727a2cde1893efe6ca55c7009196/testcontainers-4.10.0.tar.gz", hash = "sha256:03f85c3e505d8b4edeb192c72a961cebbcba0dd94344ae778b4a159cb6dcf8d3", size = 63327, upload-time = "2025-04-02T16:13:27.582Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/2a/5e/73aa94770f1df0595364aed526f31d54440db5492911e2857318ed326e51/testcontainers-4.13.2-py3-none-any.whl", hash = "sha256:0209baf8f4274b568cde95bef2cadf7b1d33b375321f793790462e235cd684ee", size = 124771, upload-time = "2025-10-07T21:53:05.937Z" },
+    { url = "https://files.pythonhosted.org/packages/1c/0a/824b0c1ecf224802125279c3effff2e25ed785ed046e67da6e53d928de4c/testcontainers-4.10.0-py3-none-any.whl", hash = "sha256:31ed1a81238c7e131a2a29df6db8f23717d892b592fa5a1977fd0dcd0c23fc23", size = 107414, upload-time = "2025-04-02T16:13:25.785Z" },
 ]
 
 [[package]]

dev/start-worker

@@ -7,4 +7,4 @@ cd "$SCRIPT_DIR/.."
 
 uv --directory api run \
 	celery -A app.celery worker \
-	-P gevent -c 1 --loglevel INFO -Q dataset,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation,priority_pipeline,pipeline
+	-P gevent -c 1 --loglevel INFO -Q dataset,priority_dataset,generation,mail,ops_trace,app_deletion,plugin,workflow_storage,conversation,priority_pipeline,pipeline

docker/.env.example

@@ -201,10 +201,6 @@ ENABLE_WEBSITE_JINAREADER=true
 ENABLE_WEBSITE_FIRECRAWL=true
 ENABLE_WEBSITE_WATERCRAWL=true
 
-# Enable inline LaTeX rendering with single dollar signs ($...$) in the web frontend
-# Default is false for security reasons to prevent conflicts with regular text
-NEXT_PUBLIC_ENABLE_SINGLE_DOLLAR_LATEX=false
-
 # ------------------------------
 # Database Configuration
 # The database uses PostgreSQL. Please use the public schema.
@@ -264,18 +260,16 @@ POSTGRES_MAINTENANCE_WORK_MEM=64MB
 POSTGRES_EFFECTIVE_CACHE_SIZE=4096MB
 
 # Sets the maximum allowed duration of any statement before termination.
-# Default is 0 (no timeout).
+# Default is 60000 milliseconds.
 #
 # Reference: https://www.postgresql.org/docs/current/runtime-config-client.html#GUC-STATEMENT-TIMEOUT
-# A value of 0 prevents the server from timing out statements.
-POSTGRES_STATEMENT_TIMEOUT=0
+POSTGRES_STATEMENT_TIMEOUT=60000
 
 # Sets the maximum allowed duration of any idle in-transaction session before termination.
-# Default is 0 (no timeout).
+# Default is 60000 milliseconds.
 #
 # Reference: https://www.postgresql.org/docs/current/runtime-config-client.html#GUC-IDLE-IN-TRANSACTION-SESSION-TIMEOUT
-# A value of 0 prevents the server from terminating idle sessions.
-POSTGRES_IDLE_IN_TRANSACTION_SESSION_TIMEOUT=0
+POSTGRES_IDLE_IN_TRANSACTION_SESSION_TIMEOUT=60000
 
 # ------------------------------
 # Redis Configuration
@@ -320,7 +314,7 @@ REDIS_CLUSTERS_PASSWORD=
 # Celery Configuration
 # ------------------------------
 
-# Use standalone redis as the broker, and redis db 1 for celery broker. (redis_username is usually set by default as empty)
+# Use standalone redis as the broker, and redis db 1 for celery broker. (redis_username is usually set by defualt as empty)
 # Format as follows: `redis://<redis_username>:<redis_password>@<redis_host>:<redis_port>/<redis_database>`.
 # Example: redis://:difyai123456@redis:6379/1
 # If use Redis Sentinel, format as follows: `sentinel://<redis_username>:<redis_password>@<sentinel_host1>:<sentinel_port>/<redis_database>`
@@ -348,11 +342,6 @@ WEB_API_CORS_ALLOW_ORIGINS=*
 # Specifies the allowed origins for cross-origin requests to the console API,
 # e.g. https://cloud.dify.ai or * for all origins.
 CONSOLE_CORS_ALLOW_ORIGINS=*
-# Set COOKIE_DOMAIN when the console frontend and API are on different subdomains.
-# Provide the registrable domain (e.g. example.com); leading dots are optional.
-COOKIE_DOMAIN=
-# The frontend reads NEXT_PUBLIC_COOKIE_DOMAIN to align cookie handling with the API.
-NEXT_PUBLIC_COOKIE_DOMAIN=
 
 # ------------------------------
 # File Storage Configuration
@@ -1255,7 +1244,6 @@ MARKETPLACE_ENABLED=true
 MARKETPLACE_API_URL=https://marketplace.dify.ai
 
 FORCE_VERIFYING_SIGNATURE=true
-ENFORCE_LANGGENIUS_PLUGIN_SIGNATURES=true
 
 PLUGIN_STDIO_BUFFER_SIZE=1024
 PLUGIN_STDIO_MAX_BUFFER_SIZE=5242880