if session unauthorized, rejoin

.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"
 		},

.devcontainer/post_create_command.sh

@@ -6,7 +6,7 @@ cd web && pnpm install
 pipx install uv
 
 echo "alias start-api=\"cd $WORKSPACE_ROOT/api && uv run python -m flask run --host 0.0.0.0 --port=5001 --debug\"" >> ~/.bashrc
-echo "alias start-worker=\"cd $WORKSPACE_ROOT/api && uv run python -m celery -A app.celery worker -P threads -c 1 --loglevel INFO -Q dataset,generation,mail,ops_trace,app_deletion,plugin,workflow_storage\"" >> ~/.bashrc
+echo "alias start-worker=\"cd $WORKSPACE_ROOT/api && uv run python -m celery -A app.celery worker -P gevent -c 1 --loglevel INFO -Q dataset,generation,mail,ops_trace,app_deletion,plugin,workflow_storage\"" >> ~/.bashrc
 echo "alias start-web=\"cd $WORKSPACE_ROOT/web && pnpm dev\"" >> ~/.bashrc
 echo "alias start-web-prod=\"cd $WORKSPACE_ROOT/web && pnpm build && pnpm start\"" >> ~/.bashrc
 echo "alias start-containers=\"cd $WORKSPACE_ROOT/docker && docker-compose -f docker-compose.middleware.yaml -p dify --env-file middleware.env up -d\"" >> ~/.bashrc

.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

@@ -8,7 +8,8 @@
             "module": "flask",
             "env": {
                 "FLASK_APP": "app.py",
-                "FLASK_ENV": "development"
+                "FLASK_ENV": "development",
+                "GEVENT_SUPPORT": "True"
             },
             "args": [
                 "run",
@@ -27,7 +28,9 @@
             "type": "debugpy",
             "request": "launch",
             "module": "celery",
-            "env": {},
+            "env": {
+                "GEVENT_SUPPORT": "True"
+            },
             "args": [
                 "-A",
                 "app.celery",
@@ -37,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

@@ -159,9 +159,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`.
@@ -374,12 +371,6 @@ UPLOAD_IMAGE_FILE_SIZE_LIMIT=10
 UPLOAD_VIDEO_FILE_SIZE_LIMIT=100
 UPLOAD_AUDIO_FILE_SIZE_LIMIT=50
 
-# Comma-separated list of file extensions blocked from upload for security reasons.
-# Extensions should be lowercase without dots (e.g., exe,bat,sh,dll).
-# Empty by default to allow all file types.
-# Recommended: exe,bat,cmd,com,scr,vbs,ps1,msi,dll
-UPLOAD_FILE_EXTENSION_BLACKLIST=
-
 # Model configuration
 MULTIMODAL_SEND_FORMAT=base64
 PROMPT_GENERATION_MAX_TOKENS=512
@@ -446,9 +437,6 @@ CODE_EXECUTION_SSL_VERIFY=True
 CODE_EXECUTION_POOL_MAX_CONNECTIONS=100
 CODE_EXECUTION_POOL_MAX_KEEPALIVE_CONNECTIONS=20
 CODE_EXECUTION_POOL_KEEPALIVE_EXPIRY=5.0
-CODE_EXECUTION_CONNECT_TIMEOUT=10
-CODE_EXECUTION_READ_TIMEOUT=60
-CODE_EXECUTION_WRITE_TIMEOUT=10
 CODE_MAX_NUMBER=9223372036854775807
 CODE_MIN_NUMBER=-9223372036854775808
 CODE_MAX_STRING_LENGTH=400000
@@ -617,6 +605,3 @@ SWAGGER_UI_PATH=/swagger-ui.html
 # Whether to encrypt dataset IDs when exporting DSL files (default: true)
 # Set to false to export dataset IDs as plain text for easier cross-environment import
 DSL_EXPORT_ENCRYPT_DATASET_ID=true
-
-# Maximum number of segments for dataset segments API (0 for unlimited)
-DATASET_MAX_SEGMENTS_PER_REQUEST=0

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/Dockerfile

@@ -15,11 +15,7 @@ FROM base AS packages
 # RUN sed -i 's@deb.debian.org@mirrors.aliyun.com@g' /etc/apt/sources.list.d/debian.sources
 
 RUN apt-get update \
-    && apt-get install -y --no-install-recommends \
-          # basic environment
-          g++ \
-          # for building gmpy2
-          libmpfr-dev libmpc-dev
+    && apt-get install -y --no-install-recommends gcc g++ libc-dev libffi-dev libgmp-dev libmpfr-dev libmpc-dev
 
 # Install Python dependencies
 COPY pyproject.toml uv.lock ./
@@ -53,9 +49,7 @@ RUN \
     # Install dependencies
     && apt-get install -y --no-install-recommends \
         # basic environment
-        curl nodejs \
-        # for gmpy2 \
-        libgmp-dev libmpfr-dev libmpc-dev \
+        curl nodejs libgmp-dev libmpfr-dev libmpc-dev \
         # For Security
         expat libldap-2.5-0 perl libsqlite3-0 zlib1g \
         # install fonts to support the use of tools like pypdfium2

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 threads -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/app.py

@@ -20,12 +20,23 @@ if is_db_command():
     socketio_app = app
     flask_app = app
 else:
-    # Gunicorn and Celery handle monkey patching automatically in production by
-    # specifying the `gevent` worker class. Manual monkey patching is not required here.
+    # It seems that JetBrains Python debugger does not work well with gevent,
+    # so we need to disable gevent in debug mode.
+    # If you are using debugpy and set GEVENT_SUPPORT=True, you can debug with gevent.
+    # if (flask_debug := os.environ.get("FLASK_DEBUG", "0")) and flask_debug.lower() in {"false", "0", "no"}:
+    # from gevent import monkey
     #
-    # See `api/docker/entrypoint.sh` (lines 33 and 47) for details.
+    # # gevent
+    # monkey.patch_all()
     #
-    # For third-party library patching, refer to `gunicorn.conf.py` and `celery_entrypoint.py`.
+    # from grpc.experimental import gevent as grpc_gevent  # type: ignore
+    #
+    # # grpc gevent
+    # grpc_gevent.init_gevent()
+
+    # import psycogreen.gevent  # type: ignore
+    #
+    # psycogreen.gevent.patch_psycopg()
 
     from app_factory import create_app
 

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
 
@@ -1601,7 +1599,7 @@ def transform_datasource_credentials():
                             "integration_secret": api_key,
                         }
                         datasource_provider = DatasourceProvider(
-                            provider="jinareader",
+                            provider="jina",
                             tenant_id=tenant_id,
                             plugin_id=jina_plugin_id,
                             auth_type=api_key_credential_type.value,

api/configs/feature/__init__.py

@@ -331,42 +331,12 @@ class FileUploadConfig(BaseSettings):
         default=10,
     )
 
-    inner_UPLOAD_FILE_EXTENSION_BLACKLIST: str = Field(
-        description=(
-            "Comma-separated list of file extensions that are blocked from upload. "
-            "Extensions should be lowercase without dots (e.g., 'exe,bat,sh,dll'). "
-            "Empty by default to allow all file types."
-        ),
-        validation_alias=AliasChoices("UPLOAD_FILE_EXTENSION_BLACKLIST"),
-        default="",
-    )
-
-    @computed_field  # type: ignore[misc]
-    @property
-    def UPLOAD_FILE_EXTENSION_BLACKLIST(self) -> set[str]:
-        """
-        Parse and return the blacklist as a set of lowercase extensions.
-        Returns an empty set if no blacklist is configured.
-        """
-        if not self.inner_UPLOAD_FILE_EXTENSION_BLACKLIST:
-            return set()
-        return {
-            ext.strip().lower().strip(".")
-            for ext in self.inner_UPLOAD_FILE_EXTENSION_BLACKLIST.split(",")
-            if ext.strip()
-        }
-
 
 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,
@@ -871,16 +841,6 @@ class MailConfig(BaseSettings):
         default=None,
     )
 
-    ENABLE_TRIAL_APP: bool = Field(
-        description="Enable trial app",
-        default=False,
-    )
-
-    ENABLE_EXPLORE_BANNER: bool = Field(
-        description="Enable explore banner",
-        default=False,
-    )
-
 
 class RagEtlConfig(BaseSettings):
     """
@@ -955,11 +915,6 @@ class DataSetConfig(BaseSettings):
         default=True,
     )
 
-    DATASET_MAX_SEGMENTS_PER_REQUEST: NonNegativeInt = Field(
-        description="Maximum number of segments for dataset segments API (0 for unlimited)",
-        default=0,
-    )
-
 
 class WorkspaceConfig(BaseSettings):
     """
@@ -1159,13 +1114,6 @@ 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,
@@ -1190,7 +1138,6 @@ class FeatureConfig(
     RagEtlConfig,
     RepositoryConfig,
     SecurityConfig,
-    TenantSelfTaskQueueConfig,
     ToolConfig,
     UpdateConfig,
     WorkflowConfig,

api/configs/feature/hosted_service/__init__.py

@@ -8,11 +8,6 @@ class HostedCreditConfig(BaseSettings):
         default="",
     )
 
-    HOSTED_POOL_CREDITS: int = Field(
-        description="Pool credits for hosted service",
-        default=200,
-    )
-
     def get_model_credits(self, model_name: str) -> int:
         """
         Get credit value for a specific model name.
@@ -65,46 +60,19 @@ class HostedOpenAiConfig(BaseSettings):
 
     HOSTED_OPENAI_TRIAL_MODELS: str = Field(
         description="Comma-separated list of available models for trial access",
-        default="gpt-4,"
-        "gpt-4-turbo-preview,"
-        "gpt-4-turbo-2024-04-09,"
-        "gpt-4-1106-preview,"
-        "gpt-4-0125-preview,"
-        "gpt-4-turbo,"
-        "gpt-4.1,"
-        "gpt-4.1-2025-04-14,"
-        "gpt-4.1-mini,"
-        "gpt-4.1-mini-2025-04-14,"
-        "gpt-4.1-nano,"
-        "gpt-4.1-nano-2025-04-14,"
-        "gpt-3.5-turbo,"
+        default="gpt-3.5-turbo,"
+        "gpt-3.5-turbo-1106,"
+        "gpt-3.5-turbo-instruct,"
         "gpt-3.5-turbo-16k,"
         "gpt-3.5-turbo-16k-0613,"
-        "gpt-3.5-turbo-1106,"
         "gpt-3.5-turbo-0613,"
         "gpt-3.5-turbo-0125,"
-        "gpt-3.5-turbo-instruct,"
-        "text-davinci-003,"
-        "chatgpt-4o-latest,"
-        "gpt-4o,"
-        "gpt-4o-2024-05-13,"
-        "gpt-4o-2024-08-06,"
-        "gpt-4o-2024-11-20,"
-        "gpt-4o-audio-preview,"
-        "gpt-4o-audio-preview-2025-06-03,"
-        "gpt-4o-mini,"
-        "gpt-4o-mini-2024-07-18,"
-        "o3-mini,"
-        "o3-mini-2025-01-31,"
-        "gpt-5-mini-2025-08-07,"
-        "gpt-5-mini,"
-        "o4-mini,"
-        "o4-mini-2025-04-16,"
-        "gpt-5-chat-latest,"
-        "gpt-5,"
-        "gpt-5-2025-08-07,"
-        "gpt-5-nano,"
-        "gpt-5-nano-2025-08-07",
+        "text-davinci-003",
+    )
+
+    HOSTED_OPENAI_QUOTA_LIMIT: NonNegativeInt = Field(
+        description="Quota limit for hosted OpenAI service usage",
+        default=200,
     )
 
     HOSTED_OPENAI_PAID_ENABLED: bool = Field(
@@ -119,13 +87,6 @@ class HostedOpenAiConfig(BaseSettings):
         "gpt-4-turbo-2024-04-09,"
         "gpt-4-1106-preview,"
         "gpt-4-0125-preview,"
-        "gpt-4-turbo,"
-        "gpt-4.1,"
-        "gpt-4.1-2025-04-14,"
-        "gpt-4.1-mini,"
-        "gpt-4.1-mini-2025-04-14,"
-        "gpt-4.1-nano,"
-        "gpt-4.1-nano-2025-04-14,"
         "gpt-3.5-turbo,"
         "gpt-3.5-turbo-16k,"
         "gpt-3.5-turbo-16k-0613,"
@@ -133,150 +94,7 @@ class HostedOpenAiConfig(BaseSettings):
         "gpt-3.5-turbo-0613,"
         "gpt-3.5-turbo-0125,"
         "gpt-3.5-turbo-instruct,"
-        "text-davinci-003,"
-        "chatgpt-4o-latest,"
-        "gpt-4o,"
-        "gpt-4o-2024-05-13,"
-        "gpt-4o-2024-08-06,"
-        "gpt-4o-2024-11-20,"
-        "gpt-4o-audio-preview,"
-        "gpt-4o-audio-preview-2025-06-03,"
-        "gpt-4o-mini,"
-        "gpt-4o-mini-2024-07-18,"
-        "o3-mini,"
-        "o3-mini-2025-01-31,"
-        "gpt-5-mini-2025-08-07,"
-        "gpt-5-mini,"
-        "o4-mini,"
-        "o4-mini-2025-04-16,"
-        "gpt-5-chat-latest,"
-        "gpt-5,"
-        "gpt-5-2025-08-07,"
-        "gpt-5-nano,"
-        "gpt-5-nano-2025-08-07",
-    )
-
-
-class HostedGeminiConfig(BaseSettings):
-    """
-    Configuration for fetching Gemini service
-    """
-
-    HOSTED_GEMINI_API_KEY: str | None = Field(
-        description="API key for hosted Gemini service",
-        default=None,
-    )
-
-    HOSTED_GEMINI_API_BASE: str | None = Field(
-        description="Base URL for hosted Gemini API",
-        default=None,
-    )
-
-    HOSTED_GEMINI_API_ORGANIZATION: str | None = Field(
-        description="Organization ID for hosted Gemini service",
-        default=None,
-    )
-
-    HOSTED_GEMINI_TRIAL_ENABLED: bool = Field(
-        description="Enable trial access to hosted Gemini service",
-        default=False,
-    )
-
-    HOSTED_GEMINI_TRIAL_MODELS: str = Field(
-        description="Comma-separated list of available models for trial access",
-        default="gemini-2.5-flash,gemini-2.0-flash,gemini-2.0-flash-lite,",
-    )
-
-    HOSTED_GEMINI_PAID_ENABLED: bool = Field(
-        description="Enable paid access to hosted gemini service",
-        default=False,
-    )
-
-    HOSTED_GEMINI_PAID_MODELS: str = Field(
-        description="Comma-separated list of available models for paid access",
-        default="gemini-2.5-flash,gemini-2.0-flash,gemini-2.0-flash-lite,",
-    )
-
-
-class HostedXAIConfig(BaseSettings):
-    """
-    Configuration for fetching XAI service
-    """
-
-    HOSTED_XAI_API_KEY: str | None = Field(
-        description="API key for hosted XAI service",
-        default=None,
-    )
-
-    HOSTED_XAI_API_BASE: str | None = Field(
-        description="Base URL for hosted XAI API",
-        default=None,
-    )
-
-    HOSTED_XAI_API_ORGANIZATION: str | None = Field(
-        description="Organization ID for hosted XAI service",
-        default=None,
-    )
-
-    HOSTED_XAI_TRIAL_ENABLED: bool = Field(
-        description="Enable trial access to hosted XAI service",
-        default=False,
-    )
-
-    HOSTED_XAI_TRIAL_MODELS: str = Field(
-        description="Comma-separated list of available models for trial access",
-        default="grok-3,grok-3-mini,grok-3-mini-fast",
-    )
-
-    HOSTED_XAI_PAID_ENABLED: bool = Field(
-        description="Enable paid access to hosted XAI service",
-        default=False,
-    )
-
-    HOSTED_XAI_PAID_MODELS: str = Field(
-        description="Comma-separated list of available models for paid access",
-        default="grok-3,grok-3-mini,grok-3-mini-fast",
-    )
-
-
-class HostedDeepseekConfig(BaseSettings):
-    """
-    Configuration for fetching Deepseek service
-    """
-
-    HOSTED_DEEPSEEK_API_KEY: str | None = Field(
-        description="API key for hosted Deepseek service",
-        default=None,
-    )
-
-    HOSTED_DEEPSEEK_API_BASE: str | None = Field(
-        description="Base URL for hosted Deepseek API",
-        default=None,
-    )
-
-    HOSTED_DEEPSEEK_API_ORGANIZATION: str | None = Field(
-        description="Organization ID for hosted Deepseek service",
-        default=None,
-    )
-
-    HOSTED_DEEPSEEK_TRIAL_ENABLED: bool = Field(
-        description="Enable trial access to hosted Deepseek service",
-        default=False,
-    )
-
-    HOSTED_DEEPSEEK_TRIAL_MODELS: str = Field(
-        description="Comma-separated list of available models for trial access",
-        default="deepseek-chat,deepseek-reasoner",
-    )
-
-    HOSTED_DEEPSEEK_PAID_ENABLED: bool = Field(
-        description="Enable paid access to hosted XAI service",
-        default=False,
-    )
-
-    HOSTED_DEEPSEEK_PAID_MODELS: str = Field(
-        description="Comma-separated list of available models for paid access",
-        default="deepseek-chat,deepseek-reasoner",
+        "text-davinci-003",
     )
 
 
@@ -326,30 +144,16 @@ class HostedAnthropicConfig(BaseSettings):
         default=False,
     )
 
+    HOSTED_ANTHROPIC_QUOTA_LIMIT: NonNegativeInt = Field(
+        description="Quota limit for hosted Anthropic service usage",
+        default=600000,
+    )
+
     HOSTED_ANTHROPIC_PAID_ENABLED: bool = Field(
         description="Enable paid access to hosted Anthropic service",
         default=False,
     )
 
-    HOSTED_ANTHROPIC_TRIAL_MODELS: str = Field(
-        description="Comma-separated list of available models for paid access",
-        default="claude-opus-4-20250514,"
-        "claude-sonnet-4-20250514,"
-        "claude-3-5-haiku-20241022,"
-        "claude-3-opus-20240229,"
-        "claude-3-7-sonnet-20250219,"
-        "claude-3-haiku-20240307",
-    )
-    HOSTED_ANTHROPIC_PAID_MODELS: str = Field(
-        description="Comma-separated list of available models for paid access",
-        default="claude-opus-4-20250514,"
-        "claude-sonnet-4-20250514,"
-        "claude-3-5-haiku-20241022,"
-        "claude-3-opus-20240229,"
-        "claude-3-7-sonnet-20250219,"
-        "claude-3-haiku-20240307",
-    )
-
 
 class HostedMinmaxConfig(BaseSettings):
     """
@@ -446,8 +250,5 @@ class HostedServiceConfig(
     HostedModerationConfig,
     # credit config
     HostedCreditConfig,
-    HostedGeminiConfig,
-    HostedXAIConfig,
-    HostedDeepseekConfig,
 ):
     pass

api/configs/middleware/__init__.py

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

api/configs/middleware/vdb/weaviate_config.py

@@ -22,11 +22,6 @@ class WeaviateConfig(BaseSettings):
         default=True,
     )
 
-    WEAVIATE_GRPC_ENDPOINT: str | None = Field(
-        description="URL of the Weaviate gRPC server (e.g., 'grpc://localhost:50051' or 'grpcs://weaviate.example.com:443')",
-        default=None,
-    )
-
     WEAVIATE_BATCH_SIZE: PositiveInt = Field(
         description="Number of objects to be processed in a single batch operation (default is 100)",
         default=100,

api/constants/__init__.py

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

api/constants/languages.py

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

api/controllers/common/errors.py

@@ -25,12 +25,6 @@ class UnsupportedFileTypeError(BaseHTTPException):
     code = 415
 
 
-class BlockedFileExtensionError(BaseHTTPException):
-    error_code = "file_extension_blocked"
-    description = "The file extension is blocked for security reasons."
-    code = 400
-
-
 class TooManyFilesError(BaseHTTPException):
     error_code = "too_many_files"
     description = "Only one file is allowed."

api/controllers/common/helpers.py

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

api/controllers/console/__init__.py

@@ -108,12 +108,10 @@ from .datasets.rag_pipeline import (
 
 # Import explore controllers
 from .explore import (
-    banner,
     installed_app,
     parameter,
     recommended_app,
     saved_message,
-    trial,
 )
 
 # Import tag controllers
@@ -147,7 +145,6 @@ __all__ = [
     "apikey",
     "app",
     "audio",
-    "banner",
     "billing",
     "bp",
     "completion",
@@ -201,7 +198,6 @@ __all__ = [
     "statistic",
     "tags",
     "tool_providers",
-    "trial",
     "version",
     "website",
     "workflow",

api/controllers/console/admin.py

@@ -16,7 +16,7 @@ from controllers.console import api, console_ns
 from controllers.console.wraps import only_edition_cloud
 from extensions.ext_database import db
 from libs.token import extract_access_token
-from models.model import App, ExporleBanner, InstalledApp, RecommendedApp, TrialApp
+from models.model import App, InstalledApp, RecommendedApp
 
 
 def admin_required(view: Callable[P, R]):
@@ -52,8 +52,6 @@ class InsertExploreAppListApi(Resource):
                 "language": fields.String(required=True, description="Language code"),
                 "category": fields.String(required=True, description="App category"),
                 "position": fields.Integer(required=True, description="Display position"),
-                "can_trial": fields.Boolean(required=True, description="Can trial"),
-                "trial_limit": fields.Integer(required=True, description="Trial limit"),
             },
         )
     )
@@ -73,8 +71,6 @@ class InsertExploreAppListApi(Resource):
             .add_argument("language", type=supported_language, required=True, nullable=False, location="json")
             .add_argument("category", type=str, required=True, nullable=False, location="json")
             .add_argument("position", type=int, required=True, nullable=False, location="json")
-            .add_argument("can_trial", type=bool, required=True, nullable=False, location="json")
-            .add_argument("trial_limit", type=int, required=True, nullable=False, location="json")
         )
         args = parser.parse_args()
 
@@ -112,20 +108,6 @@ class InsertExploreAppListApi(Resource):
                 )
 
                 db.session.add(recommended_app)
-                if args["can_trial"]:
-                    trial_app = db.session.execute(
-                        select(TrialApp).where(TrialApp.app_id == args["app_id"])
-                    ).scalar_one_or_none()
-                    if not trial_app:
-                        db.session.add(
-                            TrialApp(
-                                app_id=args["app_id"],
-                                tenant_id=app.tenant_id,
-                                trial_limit=args["trial_limit"],
-                            )
-                        )
-                    else:
-                        trial_app.trial_limit = args["trial_limit"]
 
                 app.is_public = True
                 db.session.commit()
@@ -140,20 +122,6 @@ class InsertExploreAppListApi(Resource):
                 recommended_app.category = args["category"]
                 recommended_app.position = args["position"]
 
-                if args["can_trial"]:
-                    trial_app = db.session.execute(
-                        select(TrialApp).where(TrialApp.app_id == args["app_id"])
-                    ).scalar_one_or_none()
-                    if not trial_app:
-                        db.session.add(
-                            TrialApp(
-                                app_id=args["app_id"],
-                                tenant_id=app.tenant_id,
-                                trial_limit=args["trial_limit"],
-                            )
-                        )
-                    else:
-                        trial_app.trial_limit = args["trial_limit"]
                 app.is_public = True
 
                 db.session.commit()
@@ -199,83 +167,7 @@ class InsertExploreAppApi(Resource):
             for installed_app in installed_apps:
                 session.delete(installed_app)
 
-            trial_app = session.execute(
-                select(TrialApp).where(TrialApp.app_id == recommended_app.app_id)
-            ).scalar_one_or_none()
-            if trial_app:
-                session.delete(trial_app)
-
         db.session.delete(recommended_app)
         db.session.commit()
 
         return {"result": "success"}, 204
-
-
-@console_ns.route("/admin/insert-explore-banner")
-class InsertExploreBanner(Resource):
-    @api.doc("insert_explore_banner")
-    @api.doc(description="Insert an explore banner")
-    @api.expect(
-        api.model(
-            "InsertExploreBannerRequest",
-            {
-                "content": fields.String(required=True, description="Banner content"),
-                "link": fields.String(required=True, description="Banner link"),
-                "sort": fields.Integer(required=True, description="Banner sort"),
-            },
-        )
-    )
-    @api.response(200, "Banner inserted successfully")
-    @admin_required
-    @only_edition_cloud
-    def post(self):
-        parser = reqparse.RequestParser()
-        parser.add_argument("category", type=str, required=True, nullable=False, location="json")
-        parser.add_argument("title", type=str, required=True, nullable=False, location="json")
-        parser.add_argument("description", type=str, required=True, nullable=False, location="json")
-        parser.add_argument("img-src", type=str, required=True, nullable=False, location="json")
-
-        parser.add_argument("language", type=str, required=True, nullable=False, location="json")
-        parser.add_argument("link", type=str, required=True, nullable=False, location="json")
-        parser.add_argument("sort", type=int, required=True, nullable=False, location="json")
-
-        args = parser.parse_args()
-
-        content = {
-            "category": args["category"],
-            "title": args["title"],
-            "description": args["description"],
-            "img-src": args["img-src"],
-        }
-
-        if not args["language"]:
-            args["language"] = "en-US"
-
-        banner = ExporleBanner(
-            content=content,
-            link=args["link"],
-            sort=args["sort"],
-            language=args["language"],
-        )
-        db.session.add(banner)
-        db.session.commit()
-
-        return {"result": "success"}, 200
-
-
-@console_ns.route("/admin/delete-explore-banner/<uuid:banner_id>")
-class DeleteExploreBanner(Resource):
-    @api.doc("delete_explore_banner")
-    @api.doc(description="Delete an explore banner")
-    @api.response(204, "Banner deleted successfully")
-    @admin_required
-    @only_edition_cloud
-    def delete(self, banner_id):
-        banner = db.session.execute(select(ExporleBanner).where(ExporleBanner.id == banner_id)).scalar_one_or_none()
-        if not banner:
-            raise NotFound(f"Banner '{banner_id}' is not found")
-
-        db.session.delete(banner)
-        db.session.commit()
-
-        return {"result": "success"}, 204

api/controllers/console/app/annotation.py

@@ -16,7 +16,6 @@ from fields.annotation_fields import (
     annotation_fields,
     annotation_hit_history_fields,
 )
-from libs.helper import uuid_value
 from libs.login import login_required
 from services.annotation_service import AppAnnotationService
 
@@ -176,10 +175,8 @@ class AnnotationApi(Resource):
         api.model(
             "CreateAnnotationRequest",
             {
-                "message_id": fields.String(description="Message ID (optional)"),
-                "question": fields.String(description="Question text (required when message_id not provided)"),
-                "answer": fields.String(description="Answer text (use 'answer' or 'content')"),
-                "content": fields.String(description="Content text (use 'answer' or 'content')"),
+                "question": fields.String(required=True, description="Question text"),
+                "answer": fields.String(required=True, description="Answer text"),
                 "annotation_reply": fields.Raw(description="Annotation reply data"),
             },
         )
@@ -196,14 +193,11 @@ class AnnotationApi(Resource):
         app_id = str(app_id)
         parser = (
             reqparse.RequestParser()
-            .add_argument("message_id", required=False, type=uuid_value, location="json")
-            .add_argument("question", required=False, type=str, location="json")
-            .add_argument("answer", required=False, type=str, location="json")
-            .add_argument("content", required=False, type=str, location="json")
-            .add_argument("annotation_reply", required=False, type=dict, location="json")
+            .add_argument("question", required=True, type=str, location="json")
+            .add_argument("answer", required=True, type=str, location="json")
         )
         args = parser.parse_args()
-        annotation = AppAnnotationService.up_insert_app_annotation_from_message(args, app_id)
+        annotation = AppAnnotationService.insert_app_annotation_directly(args, app_id)
         return annotation
 
     @setup_required

api/controllers/console/app/conversation.py

@@ -1,5 +1,7 @@
+from datetime import datetime
+
+import pytz
 import sqlalchemy as sa
-from flask import abort
 from flask_restx import Resource, marshal_with, reqparse
 from flask_restx.inputs import int_range
 from sqlalchemy import func, or_
@@ -17,7 +19,7 @@ from fields.conversation_fields import (
     conversation_pagination_fields,
     conversation_with_summary_pagination_fields,
 )
-from libs.datetime_utils import naive_utc_now, parse_time_range
+from libs.datetime_utils import naive_utc_now
 from libs.helper import DatetimeString
 from libs.login import current_account_with_tenant, login_required
 from models import Conversation, EndUser, Message, MessageAnnotation
@@ -88,17 +90,25 @@ class CompletionConversationApi(Resource):
 
         account = current_user
         assert account.timezone is not None
+        timezone = pytz.timezone(account.timezone)
+        utc_timezone = pytz.utc
 
-        try:
-            start_datetime_utc, end_datetime_utc = parse_time_range(args["start"], args["end"], account.timezone)
-        except ValueError as e:
-            abort(400, description=str(e))
+        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_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
 
-        if start_datetime_utc:
             query = query.where(Conversation.created_at >= start_datetime_utc)
 
-        if end_datetime_utc:
-            end_datetime_utc = end_datetime_utc.replace(second=59)
+        if args["end"]:
+            end_datetime = datetime.strptime(args["end"], "%Y-%m-%d %H:%M")
+            end_datetime = end_datetime.replace(second=59)
+
+            end_datetime_timezone = timezone.localize(end_datetime)
+            end_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
+
             query = query.where(Conversation.created_at < end_datetime_utc)
 
         # FIXME, the type ignore in this file
@@ -260,21 +270,29 @@ class ChatConversationApi(Resource):
 
         account = current_user
         assert account.timezone is not None
+        timezone = pytz.timezone(account.timezone)
+        utc_timezone = pytz.utc
 
-        try:
-            start_datetime_utc, end_datetime_utc = parse_time_range(args["start"], args["end"], account.timezone)
-        except ValueError as e:
-            abort(400, description=str(e))
+        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_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
 
-        if start_datetime_utc:
             match args["sort_by"]:
                 case "updated_at" | "-updated_at":
                     query = query.where(Conversation.updated_at >= start_datetime_utc)
                 case "created_at" | "-created_at" | _:
                     query = query.where(Conversation.created_at >= start_datetime_utc)
 
-        if end_datetime_utc:
-            end_datetime_utc = end_datetime_utc.replace(second=59)
+        if args["end"]:
+            end_datetime = datetime.strptime(args["end"], "%Y-%m-%d %H:%M")
+            end_datetime = end_datetime.replace(second=59)
+
+            end_datetime_timezone = timezone.localize(end_datetime)
+            end_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
+
             match args["sort_by"]:
                 case "updated_at" | "-updated_at":
                     query = query.where(Conversation.updated_at <= end_datetime_utc)

api/controllers/console/app/error.py

@@ -115,9 +115,3 @@ class InvokeRateLimitError(BaseHTTPException):
     error_code = "rate_limit_error"
     description = "Rate Limit Error"
     code = 429
-
-
-class NeedAddIdsError(BaseHTTPException):
-    error_code = "need_add_ids"
-    description = "Need to add ids."
-    code = 400

api/controllers/console/app/message.py

@@ -16,6 +16,7 @@ from controllers.console.app.wraps import get_app_model
 from controllers.console.explore.error import AppSuggestedQuestionsAfterAnswerDisabledError
 from controllers.console.wraps import (
     account_initialization_required,
+    cloud_edition_billing_resource_check,
     edit_permission_required,
     setup_required,
 )
@@ -23,11 +24,12 @@ from core.app.entities.app_invoke_entities import InvokeFrom
 from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
 from core.model_runtime.errors.invoke import InvokeError
 from extensions.ext_database import db
-from fields.conversation_fields import message_detail_fields
+from fields.conversation_fields import annotation_fields, message_detail_fields
 from libs.helper import uuid_value
 from libs.infinite_scroll_pagination import InfiniteScrollPagination
 from libs.login import current_account_with_tenant, login_required
 from models.model import AppMode, Conversation, Message, MessageAnnotation, MessageFeedback
+from services.annotation_service import AppAnnotationService
 from services.errors.conversation import ConversationNotExistsError
 from services.errors.message import MessageNotExistsError, SuggestedQuestionsAfterAnswerDisabledError
 from services.message_service import MessageService
@@ -192,6 +194,45 @@ class MessageFeedbackApi(Resource):
         return {"result": "success"}
 
 
+@console_ns.route("/apps/<uuid:app_id>/annotations")
+class MessageAnnotationApi(Resource):
+    @api.doc("create_message_annotation")
+    @api.doc(description="Create message annotation")
+    @api.doc(params={"app_id": "Application ID"})
+    @api.expect(
+        api.model(
+            "MessageAnnotationRequest",
+            {
+                "message_id": fields.String(description="Message ID"),
+                "question": fields.String(required=True, description="Question text"),
+                "answer": fields.String(required=True, description="Answer text"),
+                "annotation_reply": fields.Raw(description="Annotation reply"),
+            },
+        )
+    )
+    @api.response(200, "Annotation created successfully", annotation_fields)
+    @api.response(403, "Insufficient permissions")
+    @marshal_with(annotation_fields)
+    @get_app_model
+    @setup_required
+    @login_required
+    @cloud_edition_billing_resource_check("annotation")
+    @account_initialization_required
+    @edit_permission_required
+    def post(self, app_model):
+        parser = (
+            reqparse.RequestParser()
+            .add_argument("message_id", required=False, type=uuid_value, location="json")
+            .add_argument("question", required=True, type=str, location="json")
+            .add_argument("answer", required=True, type=str, location="json")
+            .add_argument("annotation_reply", required=False, type=dict, location="json")
+        )
+        args = parser.parse_args()
+        annotation = AppAnnotationService.up_insert_app_annotation_from_message(args, app_model.id)
+
+        return annotation
+
+
 @console_ns.route("/apps/<uuid:app_id>/annotations/count")
 class MessageAnnotationCountApi(Resource):
     @api.doc("get_annotation_count")

api/controllers/console/app/statistic.py

@@ -1,7 +1,9 @@
+from datetime import datetime
 from decimal import Decimal
 
+import pytz
 import sqlalchemy as sa
-from flask import abort, jsonify
+from flask import jsonify
 from flask_restx import Resource, fields, reqparse
 
 from controllers.console import api, console_ns
@@ -9,7 +11,6 @@ from controllers.console.app.wraps import get_app_model
 from controllers.console.wraps import account_initialization_required, setup_required
 from core.app.entities.app_invoke_entities import InvokeFrom
 from extensions.ext_database import db
-from libs.datetime_utils import parse_time_range
 from libs.helper import DatetimeString
 from libs.login import current_account_with_tenant, login_required
 from models import AppMode, Message
@@ -55,16 +56,26 @@ WHERE
         arg_dict = {"tz": account.timezone, "app_id": app_model.id, "invoke_from": InvokeFrom.DEBUGGER}
         assert account.timezone is not None
 
-        try:
-            start_datetime_utc, end_datetime_utc = parse_time_range(args["start"], args["end"], account.timezone)
-        except ValueError as e:
-            abort(400, description=str(e))
+        timezone = pytz.timezone(account.timezone)
+        utc_timezone = pytz.utc
+
+        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_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
 
-        if start_datetime_utc:
             sql_query += " AND created_at >= :start"
             arg_dict["start"] = start_datetime_utc
 
-        if end_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_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
+
             sql_query += " AND created_at < :end"
             arg_dict["end"] = end_datetime_utc
 
@@ -109,11 +120,8 @@ class DailyConversationStatistic(Resource):
         )
         args = parser.parse_args()
         assert account.timezone is not None
-
-        try:
-            start_datetime_utc, end_datetime_utc = parse_time_range(args["start"], args["end"], account.timezone)
-        except ValueError as e:
-            abort(400, description=str(e))
+        timezone = pytz.timezone(account.timezone)
+        utc_timezone = pytz.utc
 
         stmt = (
             sa.select(
@@ -126,10 +134,18 @@ class DailyConversationStatistic(Resource):
             .where(Message.app_id == app_model.id, Message.invoke_from != InvokeFrom.DEBUGGER)
         )
 
-        if start_datetime_utc:
+        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_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
             stmt = stmt.where(Message.created_at >= start_datetime_utc)
 
-        if end_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_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
             stmt = stmt.where(Message.created_at < end_datetime_utc)
 
         stmt = stmt.group_by("date").order_by("date")
@@ -182,17 +198,26 @@ WHERE
     AND invoke_from != :invoke_from"""
         arg_dict = {"tz": account.timezone, "app_id": app_model.id, "invoke_from": InvokeFrom.DEBUGGER}
         assert account.timezone is not None
+        timezone = pytz.timezone(account.timezone)
+        utc_timezone = pytz.utc
 
-        try:
-            start_datetime_utc, end_datetime_utc = parse_time_range(args["start"], args["end"], account.timezone)
-        except ValueError as e:
-            abort(400, description=str(e))
+        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_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
 
-        if start_datetime_utc:
             sql_query += " AND created_at >= :start"
             arg_dict["start"] = start_datetime_utc
 
-        if end_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_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
+
             sql_query += " AND created_at < :end"
             arg_dict["end"] = end_datetime_utc
 
@@ -248,17 +273,26 @@ WHERE
     AND invoke_from != :invoke_from"""
         arg_dict = {"tz": account.timezone, "app_id": app_model.id, "invoke_from": InvokeFrom.DEBUGGER}
         assert account.timezone is not None
+        timezone = pytz.timezone(account.timezone)
+        utc_timezone = pytz.utc
 
-        try:
-            start_datetime_utc, end_datetime_utc = parse_time_range(args["start"], args["end"], account.timezone)
-        except ValueError as e:
-            abort(400, description=str(e))
+        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_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
 
-        if start_datetime_utc:
             sql_query += " AND created_at >= :start"
             arg_dict["start"] = start_datetime_utc
 
-        if end_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_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
+
             sql_query += " AND created_at < :end"
             arg_dict["end"] = end_datetime_utc
 
@@ -323,17 +357,26 @@ FROM
             AND m.invoke_from != :invoke_from"""
         arg_dict = {"tz": account.timezone, "app_id": app_model.id, "invoke_from": InvokeFrom.DEBUGGER}
         assert account.timezone is not None
+        timezone = pytz.timezone(account.timezone)
+        utc_timezone = pytz.utc
 
-        try:
-            start_datetime_utc, end_datetime_utc = parse_time_range(args["start"], args["end"], account.timezone)
-        except ValueError as e:
-            abort(400, description=str(e))
+        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_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
 
-        if start_datetime_utc:
             sql_query += " AND c.created_at >= :start"
             arg_dict["start"] = start_datetime_utc
 
-        if end_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_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
+
             sql_query += " AND c.created_at < :end"
             arg_dict["end"] = end_datetime_utc
 
@@ -403,17 +446,26 @@ WHERE
     AND m.invoke_from != :invoke_from"""
         arg_dict = {"tz": account.timezone, "app_id": app_model.id, "invoke_from": InvokeFrom.DEBUGGER}
         assert account.timezone is not None
+        timezone = pytz.timezone(account.timezone)
+        utc_timezone = pytz.utc
 
-        try:
-            start_datetime_utc, end_datetime_utc = parse_time_range(args["start"], args["end"], account.timezone)
-        except ValueError as e:
-            abort(400, description=str(e))
+        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_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
 
-        if start_datetime_utc:
             sql_query += " AND m.created_at >= :start"
             arg_dict["start"] = start_datetime_utc
 
-        if end_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_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
+
             sql_query += " AND m.created_at < :end"
             arg_dict["end"] = end_datetime_utc
 
@@ -473,17 +525,26 @@ WHERE
     AND invoke_from != :invoke_from"""
         arg_dict = {"tz": account.timezone, "app_id": app_model.id, "invoke_from": InvokeFrom.DEBUGGER}
         assert account.timezone is not None
+        timezone = pytz.timezone(account.timezone)
+        utc_timezone = pytz.utc
 
-        try:
-            start_datetime_utc, end_datetime_utc = parse_time_range(args["start"], args["end"], account.timezone)
-        except ValueError as e:
-            abort(400, description=str(e))
+        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_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
 
-        if start_datetime_utc:
             sql_query += " AND created_at >= :start"
             arg_dict["start"] = start_datetime_utc
 
-        if end_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_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
+
             sql_query += " AND created_at < :end"
             arg_dict["end"] = end_datetime_utc
 
@@ -541,17 +602,26 @@ WHERE
     AND invoke_from != :invoke_from"""
         arg_dict = {"tz": account.timezone, "app_id": app_model.id, "invoke_from": InvokeFrom.DEBUGGER}
         assert account.timezone is not None
+        timezone = pytz.timezone(account.timezone)
+        utc_timezone = pytz.utc
 
-        try:
-            start_datetime_utc, end_datetime_utc = parse_time_range(args["start"], args["end"], account.timezone)
-        except ValueError as e:
-            abort(400, description=str(e))
+        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_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
 
-        if start_datetime_utc:
             sql_query += " AND created_at >= :start"
             arg_dict["start"] = start_datetime_utc
 
-        if end_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_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
+
             sql_query += " AND created_at < :end"
             arg_dict["end"] = end_datetime_utc
 

api/controllers/console/app/workflow.py

@@ -5,7 +5,6 @@ from typing import cast
 
 from flask import abort, request
 from flask_restx import Resource, fields, inputs, marshal_with, reqparse
-from pydantic_core import ValidationError
 from sqlalchemy.orm import Session
 from werkzeug.exceptions import Forbidden, InternalServerError, NotFound
 
@@ -103,22 +102,10 @@ class DraftWorkflowApi(Resource):
                 "hash": fields.String(description="Workflow hash for validation"),
                 "environment_variables": fields.List(fields.Raw, required=True, description="Environment variables"),
                 "conversation_variables": fields.List(fields.Raw, description="Conversation variables"),
-                "memory_blocks": fields.List(fields.Raw, description="Memory blocks"),
             },
         )
     )
-    @api.response(
-        200,
-        "Draft workflow synced successfully",
-        api.model(
-            "SyncDraftWorkflowResponse",
-            {
-                "result": fields.String,
-                "hash": fields.String,
-                "updated_at": fields.String,
-            },
-        ),
-    )
+    @api.response(200, "Draft workflow synced successfully", workflow_fields)
     @api.response(400, "Invalid workflow configuration")
     @api.response(403, "Permission denied")
     @edit_permission_required
@@ -139,7 +126,6 @@ class DraftWorkflowApi(Resource):
                 .add_argument("environment_variables", type=list, required=True, location="json")
                 .add_argument("conversation_variables", type=list, required=False, location="json")
                 .add_argument("force_upload", type=bool, required=False, default=False, location="json")
-                .add_argument("memory_blocks", type=list, required=False, location="json")
             )
             args = parser.parse_args()
         elif "text/plain" in content_type:
@@ -157,7 +143,6 @@ class DraftWorkflowApi(Resource):
                     "hash": data.get("hash"),
                     "environment_variables": data.get("environment_variables"),
                     "conversation_variables": data.get("conversation_variables"),
-                    "memory_blocks": data.get("memory_blocks"),
                     "force_upload": data.get("force_upload", False),
                 }
             except json.JSONDecodeError:
@@ -175,10 +160,6 @@ class DraftWorkflowApi(Resource):
             conversation_variables = [
                 variable_factory.build_conversation_variable_from_mapping(obj) for obj in conversation_variables_list
             ]
-            memory_blocks_list = args.get("memory_blocks") or []
-            from core.memory.entities import MemoryBlockSpec
-
-            memory_blocks = [MemoryBlockSpec.model_validate(obj) for obj in memory_blocks_list]
             workflow = workflow_service.sync_draft_workflow(
                 app_model=app_model,
                 graph=args["graph"],
@@ -188,12 +169,9 @@ class DraftWorkflowApi(Resource):
                 environment_variables=environment_variables,
                 conversation_variables=conversation_variables,
                 force_upload=args.get("force_upload", False),
-                memory_blocks=memory_blocks,
             )
         except WorkflowHashNotEqualError:
             raise DraftWorkflowNotSync()
-        except ValidationError as e:
-            return {"message": str(e)}, 400
 
         return {
             "result": "success",
@@ -780,6 +758,7 @@ class WorkflowConfigApi(Resource):
         }
 
 
+@console_ns.route("/apps/<uuid:app_id>/workflows/draft/features")
 class WorkflowFeaturesApi(Resource):
     """Update draft workflow features."""
 
@@ -788,13 +767,13 @@ class WorkflowFeaturesApi(Resource):
     @account_initialization_required
     @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW])
     def post(self, app_model: App):
-        parser = reqparse.RequestParser()
-        parser.add_argument("features", type=dict, required=True, location="json")
+        current_user, _ = current_account_with_tenant()
+
+        parser = reqparse.RequestParser().add_argument("features", type=dict, required=True, location="json")
         args = parser.parse_args()
 
         features = args.get("features")
 
-        # Update draft workflow features
         workflow_service = WorkflowService()
         workflow_service.update_draft_workflow_features(app_model=app_model, features=features, account=current_user)
 
@@ -984,17 +963,17 @@ class DraftWorkflowNodeLastRunApi(Resource):
         return node_exec
 
 
+@console_ns.route("/apps/workflows/online-users")
 class WorkflowOnlineUsersApi(Resource):
     @setup_required
     @login_required
     @account_initialization_required
     @marshal_with(online_user_list_fields)
     def get(self):
-        parser = reqparse.RequestParser()
-        parser.add_argument("workflow_ids", type=str, required=True, location="args")
+        parser = reqparse.RequestParser().add_argument("workflow_ids", type=str, required=True, location="args")
         args = parser.parse_args()
 
-        workflow_ids = [id.strip() for id in args["workflow_ids"].split(",")]
+        workflow_ids = [workflow_id.strip() for workflow_id in args["workflow_ids"].split(",")]
 
         results = []
         for workflow_id in workflow_ids:
@@ -1009,78 +988,3 @@ class WorkflowOnlineUsersApi(Resource):
             results.append({"workflow_id": workflow_id, "users": users})
 
         return {"data": results}
-
-
-api.add_resource(
-    DraftWorkflowApi,
-    "/apps/<uuid:app_id>/workflows/draft",
-)
-api.add_resource(
-    WorkflowConfigApi,
-    "/apps/<uuid:app_id>/workflows/draft/config",
-)
-api.add_resource(
-    WorkflowFeaturesApi,
-    "/apps/<uuid:app_id>/workflows/draft/features",
-)
-api.add_resource(
-    AdvancedChatDraftWorkflowRunApi,
-    "/apps/<uuid:app_id>/advanced-chat/workflows/draft/run",
-)
-api.add_resource(
-    DraftWorkflowRunApi,
-    "/apps/<uuid:app_id>/workflows/draft/run",
-)
-api.add_resource(
-    WorkflowTaskStopApi,
-    "/apps/<uuid:app_id>/workflow-runs/tasks/<string:task_id>/stop",
-)
-api.add_resource(
-    DraftWorkflowNodeRunApi,
-    "/apps/<uuid:app_id>/workflows/draft/nodes/<string:node_id>/run",
-)
-api.add_resource(
-    AdvancedChatDraftRunIterationNodeApi,
-    "/apps/<uuid:app_id>/advanced-chat/workflows/draft/iteration/nodes/<string:node_id>/run",
-)
-api.add_resource(
-    WorkflowDraftRunIterationNodeApi,
-    "/apps/<uuid:app_id>/workflows/draft/iteration/nodes/<string:node_id>/run",
-)
-api.add_resource(
-    AdvancedChatDraftRunLoopNodeApi,
-    "/apps/<uuid:app_id>/advanced-chat/workflows/draft/loop/nodes/<string:node_id>/run",
-)
-api.add_resource(
-    WorkflowDraftRunLoopNodeApi,
-    "/apps/<uuid:app_id>/workflows/draft/loop/nodes/<string:node_id>/run",
-)
-api.add_resource(
-    PublishedWorkflowApi,
-    "/apps/<uuid:app_id>/workflows/publish",
-)
-api.add_resource(
-    PublishedAllWorkflowApi,
-    "/apps/<uuid:app_id>/workflows",
-)
-api.add_resource(
-    DefaultBlockConfigsApi,
-    "/apps/<uuid:app_id>/workflows/default-workflow-block-configs",
-)
-api.add_resource(
-    DefaultBlockConfigApi,
-    "/apps/<uuid:app_id>/workflows/default-workflow-block-configs/<string:block_type>",
-)
-api.add_resource(
-    ConvertToWorkflowApi,
-    "/apps/<uuid:app_id>/convert-to-workflow",
-)
-api.add_resource(
-    WorkflowByIdApi,
-    "/apps/<uuid:app_id>/workflows/<string:workflow_id>",
-)
-api.add_resource(
-    DraftWorkflowNodeLastRunApi,
-    "/apps/<uuid:app_id>/workflows/draft/nodes/<string:node_id>/last-run",
-)
-api.add_resource(WorkflowOnlineUsersApi, "/apps/workflows/online-users")

api/controllers/console/app/workflow_statistic.py

@@ -1,26 +1,23 @@
-from flask import abort, jsonify
+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
 from controllers.console.wraps import account_initialization_required, setup_required
 from extensions.ext_database import db
-from libs.datetime_utils import parse_time_range
 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"})
@@ -40,32 +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
 
-        try:
-            start_date, end_date = parse_time_range(args["start"], args["end"], account.timezone)
-        except ValueError as e:
-            abort(400, description=str(e))
+        if args["start"]:
+            start_datetime = datetime.strptime(args["start"], "%Y-%m-%d %H:%M")
+            start_datetime = start_datetime.replace(second=0)
 
-        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,
-        )
+            start_datetime_timezone = timezone.localize(start_datetime)
+            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_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"})
@@ -85,32 +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
 
-        try:
-            start_date, end_date = parse_time_range(args["start"], args["end"], account.timezone)
-        except ValueError as e:
-            abort(400, description=str(e))
+        if args["start"]:
+            start_datetime = datetime.strptime(args["start"], "%Y-%m-%d %H:%M")
+            start_datetime = start_datetime.replace(second=0)
 
-        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,
-        )
+            start_datetime_timezone = timezone.localize(start_datetime)
+            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_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"})
@@ -130,32 +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
 
-        try:
-            start_date, end_date = parse_time_range(args["start"], args["end"], account.timezone)
-        except ValueError as e:
-            abort(400, description=str(e))
+        if args["start"]:
+            start_datetime = datetime.strptime(args["start"], "%Y-%m-%d %H:%M")
+            start_datetime = start_datetime.replace(second=0)
 
-        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,
-        )
+            start_datetime_timezone = timezone.localize(start_datetime)
+            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_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"})
@@ -175,20 +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
 
-        try:
-            start_date, end_date = parse_time_range(args["start"], args["end"], account.timezone)
-        except ValueError as e:
-            abort(400, description=str(e))
+        if args["start"]:
+            start_datetime = datetime.strptime(args["start"], "%Y-%m-%d %H:%M")
+            start_datetime = start_datetime.replace(second=0)
 
-        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,
-        )
+            start_datetime_timezone = timezone.localize(start_datetime)
+            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_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/app/wraps.py

@@ -23,15 +23,6 @@ def _load_app_model(app_id: str) -> App | None:
     return app_model
 
 
-def _load_app_model_with_trial(app_id: str) -> App | None:
-    app_model = (
-        db.session.query(App)
-        .where(App.id == app_id, App.status == "normal")
-        .first()
-    )
-    return app_model
-
-
 def get_app_model(view: Callable[P, R] | None = None, *, mode: Union[AppMode, list[AppMode], None] = None):
     def decorator(view_func: Callable[P1, R1]):
         @wraps(view_func)
@@ -71,44 +62,3 @@ def get_app_model(view: Callable[P, R] | None = None, *, mode: Union[AppMode, li
         return decorator
     else:
         return decorator(view)
-
-
-def get_app_model_with_trial(view: Callable[P, R] | None = None, *, mode: Union[AppMode, list[AppMode], None] = None):
-    def decorator(view_func: Callable[P, R]):
-        @wraps(view_func)
-        def decorated_view(*args: P.args, **kwargs: P.kwargs):
-            if not kwargs.get("app_id"):
-                raise ValueError("missing app_id in path parameters")
-
-            app_id = kwargs.get("app_id")
-            app_id = str(app_id)
-
-            del kwargs["app_id"]
-
-            app_model = _load_app_model_with_trial(app_id)
-
-            if not app_model:
-                raise AppNotFoundError()
-
-            app_mode = AppMode.value_of(app_model.mode)
-
-            if mode is not None:
-                if isinstance(mode, list):
-                    modes = mode
-                else:
-                    modes = [mode]
-
-                if app_mode not in modes:
-                    mode_values = {m.value for m in modes}
-                    raise AppNotFoundError(f"App mode is not in the supported list: {mode_values}")
-
-            kwargs["app_model"] = app_model
-
-            return view_func(*args, **kwargs)
-
-        return decorated_view
-
-    if view is None:
-        return decorator
-    else:
-        return decorator(view)

api/controllers/console/auth/login.py

@@ -4,7 +4,7 @@ from flask_restx import Resource, reqparse
 
 import services
 from configs import dify_config
-from constants.languages import get_valid_language
+from constants.languages import languages
 from controllers.console import console_ns
 from controllers.console.auth.error import (
     AuthenticationFailedError,
@@ -25,14 +25,12 @@ from controllers.console.wraps import email_password_login_enabled, setup_requir
 from events.tenant_event import tenant_was_created
 from libs.helper import email, extract_remote_ip
 from libs.login import current_account_with_tenant
-from libs.passport import PassportService
 from libs.token import (
-    check_csrf_token,
     clear_access_token_from_cookie,
     clear_csrf_token_from_cookie,
     clear_refresh_token_from_cookie,
     extract_access_token,
-    extract_refresh_token,
+    extract_csrf_token,
     set_access_token_to_cookie,
     set_csrf_token_to_cookie,
     set_refresh_token_to_cookie,
@@ -208,12 +206,10 @@ class EmailCodeLoginApi(Resource):
             .add_argument("email", type=str, required=True, location="json")
             .add_argument("code", type=str, required=True, location="json")
             .add_argument("token", type=str, required=True, location="json")
-            .add_argument("language", type=str, required=False, location="json")
         )
         args = parser.parse_args()
 
         user_email = args["email"]
-        language = args["language"]
 
         token_data = AccountService.get_email_code_login_data(args["token"])
         if token_data is None:
@@ -247,9 +243,7 @@ class EmailCodeLoginApi(Resource):
         if account is None:
             try:
                 account = AccountService.create_account_and_tenant(
-                    email=user_email,
-                    name=user_email,
-                    interface_language=get_valid_language(language),
+                    email=user_email, name=user_email, interface_language=languages[0]
                 )
             except WorkSpaceNotAllowedCreateError:
                 raise NotAllowedCreateWorkspace()
@@ -274,7 +268,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
@@ -300,10 +294,5 @@ class RefreshTokenApi(Resource):
 class LoginStatus(Resource):
     def get(self):
         token = extract_access_token(request)
-        res = True
-        try:
-            validated = PassportService().verify(token=token)
-            check_csrf_token(request=request, user_id=validated.get("user_id"))
-        except:
-            res = False
-        return {"logged_in": res}
+        csrf_token = extract_csrf_token(request)
+        return {"logged_in": bool(token) and bool(csrf_token)}

api/controllers/console/billing/billing.py

@@ -2,7 +2,6 @@ from flask_restx import Resource, reqparse
 
 from controllers.console import console_ns
 from controllers.console.wraps import account_initialization_required, only_edition_cloud, setup_required
-from enums.cloud_plan import CloudPlan
 from libs.login import current_account_with_tenant, login_required
 from services.billing_service import BillingService
 
@@ -17,13 +16,7 @@ class Subscription(Resource):
         current_user, current_tenant_id = current_account_with_tenant()
         parser = (
             reqparse.RequestParser()
-            .add_argument(
-                "plan",
-                type=str,
-                required=True,
-                location="args",
-                choices=[CloudPlan.PROFESSIONAL, CloudPlan.TEAM],
-            )
+            .add_argument("plan", type=str, required=True, location="args", choices=["professional", "team"])
             .add_argument("interval", type=str, required=True, location="args", choices=["month", "year"])
         )
         args = parser.parse_args()

api/controllers/console/datasets/datasets_document.py

@@ -746,7 +746,7 @@ class DocumentApi(DocumentResource):
                 "name": document.name,
                 "created_from": document.created_from,
                 "created_by": document.created_by,
-                "created_at": int(document.created_at.timestamp()),
+                "created_at": document.created_at.timestamp(),
                 "tokens": document.tokens,
                 "indexing_status": document.indexing_status,
                 "completed_at": int(document.completed_at.timestamp()) if document.completed_at else None,
@@ -779,7 +779,7 @@ class DocumentApi(DocumentResource):
                 "name": document.name,
                 "created_from": document.created_from,
                 "created_by": document.created_by,
-                "created_at": int(document.created_at.timestamp()),
+                "created_at": document.created_at.timestamp(),
                 "tokens": document.tokens,
                 "indexing_status": document.indexing_status,
                 "completed_at": int(document.completed_at.timestamp()) if document.completed_at else None,

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/explore/banner.py

@@ -1,43 +0,0 @@
-from flask import request
-from flask_restx import Resource
-
-from controllers.console import api
-from controllers.console.explore.wraps import explore_banner_enabled
-from extensions.ext_database import db
-from models.model import ExporleBanner
-
-
-class BannerApi(Resource):
-    """Resource for banner list."""
-
-    @explore_banner_enabled
-    def get(self):
-        """Get banner list."""
-        language = request.args.get("language", "en-US")
-
-        # Build base query for enabled banners
-        base_query = db.session.query(ExporleBanner).where(ExporleBanner.status == "enabled")
-        
-        # Try to get banners in the requested language
-        banners = base_query.where(ExporleBanner.language == language).order_by(ExporleBanner.sort).all()
-        
-        # Fallback to en-US if no banners found and language is not en-US
-        if not banners and language != "en-US":
-            banners = base_query.where(ExporleBanner.language == "en-US").order_by(ExporleBanner.sort).all()
-        # Convert banners to serializable format
-        result = []
-        for banner in banners:
-            banner_data = {
-                "id": banner.id,
-                "content": banner.content,  # Already parsed as JSON by SQLAlchemy
-                "link": banner.link,
-                "sort": banner.sort,
-                "status": banner.status,
-                "created_at": banner.created_at.isoformat() if banner.created_at else None,
-            }
-            result.append(banner_data)
-
-        return result
-
-
-api.add_resource(BannerApi, "/explore/banners")

api/controllers/console/explore/error.py

@@ -29,25 +29,3 @@ class AppAccessDeniedError(BaseHTTPException):
     error_code = "access_denied"
     description = "App access denied."
     code = 403
-
-
-class TrialAppNotAllowed(BaseHTTPException):
-    """*403* `Trial App Not Allowed`
-
-    Raise if the user has reached the trial app limit.
-    """
-
-    error_code = "trial_app_not_allowed"
-    code = 403
-    description = "the app is not allowed to be trial."
-
-
-class TrialAppLimitExceeded(BaseHTTPException):
-    """*403* `Trial App Limit Exceeded`
-
-    Raise if the user has exceeded the trial app limit.
-    """
-
-    error_code = "trial_app_limit_exceeded"
-    code = 403
-    description = "The user has exceeded the trial app limit."

api/controllers/console/explore/recommended_app.py

@@ -27,7 +27,6 @@ recommended_app_fields = {
     "category": fields.String,
     "position": fields.Integer,
     "is_listed": fields.Boolean,
-    "can_trial": fields.Boolean,
 }
 
 recommended_app_list_fields = {

api/controllers/console/explore/trial.py

@@ -1,514 +0,0 @@
-import logging
-from typing import Any, cast
-
-from flask import request
-from flask_restx import Resource, marshal, marshal_with, reqparse
-from werkzeug.exceptions import Forbidden, InternalServerError, NotFound
-
-import services
-from controllers.common import fields
-from controllers.common.fields import build_site_model
-from controllers.console import api
-from controllers.console.app.error import (
-    AppUnavailableError,
-    AudioTooLargeError,
-    CompletionRequestError,
-    ConversationCompletedError,
-    NeedAddIdsError,
-    NoAudioUploadedError,
-    ProviderModelCurrentlyNotSupportError,
-    ProviderNotInitializeError,
-    ProviderNotSupportSpeechToTextError,
-    ProviderQuotaExceededError,
-    UnsupportedAudioTypeError,
-)
-from controllers.console.app.wraps import get_app_model_with_trial
-from controllers.console.explore.error import (
-    AppSuggestedQuestionsAfterAnswerDisabledError,
-    NotChatAppError,
-    NotCompletionAppError,
-    NotWorkflowAppError,
-)
-from controllers.console.explore.wraps import TrialAppResource, trial_feature_enable
-from controllers.service_api import service_api_ns
-from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError
-from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict
-from core.app.apps.base_app_queue_manager import AppQueueManager
-from core.app.entities.app_invoke_entities import InvokeFrom
-from core.errors.error import (
-    ModelCurrentlyNotSupportError,
-    ProviderTokenNotInitError,
-    QuotaExceededError,
-)
-from core.model_runtime.errors.invoke import InvokeError
-from core.workflow.graph_engine.manager import GraphEngineManager
-from extensions.ext_database import db
-from fields.app_fields import app_detail_fields_with_site
-from fields.dataset_fields import dataset_fields
-from fields.workflow_fields import workflow_fields
-from libs import helper
-from libs.helper import uuid_value
-from libs.login import current_user
-from models import Account
-from models.account import TenantStatus
-from models.model import AppMode, Site
-from models.workflow import Workflow
-from services.app_generate_service import AppGenerateService
-from services.app_service import AppService
-from services.audio_service import AudioService
-from services.dataset_service import DatasetService
-from services.errors.audio import (
-    AudioTooLargeServiceError,
-    NoAudioUploadedServiceError,
-    ProviderNotSupportSpeechToTextServiceError,
-    UnsupportedAudioTypeServiceError,
-)
-from services.errors.conversation import ConversationNotExistsError
-from services.errors.llm import InvokeRateLimitError
-from services.errors.message import (
-    MessageNotExistsError,
-    SuggestedQuestionsAfterAnswerDisabledError,
-)
-from services.message_service import MessageService
-from services.recommended_app_service import RecommendedAppService
-
-logger = logging.getLogger(__name__)
-
-
-class TrialAppWorkflowRunApi(TrialAppResource):
-    def post(self, trial_app):
-        """
-        Run workflow
-        """
-        app_model = trial_app
-        if not app_model:
-            raise NotWorkflowAppError()
-        app_mode = AppMode.value_of(app_model.mode)
-        if app_mode != AppMode.WORKFLOW:
-            raise NotWorkflowAppError()
-
-        parser = reqparse.RequestParser()
-        parser.add_argument("inputs", type=dict, required=True, nullable=False, location="json")
-        parser.add_argument("files", type=list, required=False, location="json")
-        args = parser.parse_args()
-        assert current_user is not None
-        try:
-            app_id = app_model.id
-            user_id = current_user.id
-            response = AppGenerateService.generate(
-                app_model=app_model, user=current_user, args=args, invoke_from=InvokeFrom.EXPLORE, streaming=True
-            )
-            RecommendedAppService.add_trial_app_record(app_id, user_id)
-            return helper.compact_generate_response(response)
-        except ProviderTokenNotInitError as ex:
-            raise ProviderNotInitializeError(ex.description)
-        except QuotaExceededError:
-            raise ProviderQuotaExceededError()
-        except ModelCurrentlyNotSupportError:
-            raise ProviderModelCurrentlyNotSupportError()
-        except InvokeError as e:
-            raise CompletionRequestError(e.description)
-        except InvokeRateLimitError as ex:
-            raise InvokeRateLimitHttpError(ex.description)
-        except ValueError as e:
-            raise e
-        except Exception:
-            logger.exception("internal server error.")
-            raise InternalServerError()
-
-
-class TrialAppWorkflowTaskStopApi(TrialAppResource):
-    def post(self, trial_app, task_id: str):
-        """
-        Stop workflow task
-        """
-        app_model = trial_app
-        if not app_model:
-            raise NotWorkflowAppError()
-        app_mode = AppMode.value_of(app_model.mode)
-        if app_mode != AppMode.WORKFLOW:
-            raise NotWorkflowAppError()
-        assert current_user is not None
-
-        # Stop using both mechanisms for backward compatibility
-        # Legacy stop flag mechanism (without user check)
-        AppQueueManager.set_stop_flag_no_user_check(task_id)
-
-        # New graph engine command channel mechanism
-        GraphEngineManager.send_stop_command(task_id)
-
-        return {"result": "success"}
-
-
-class TrialChatApi(TrialAppResource):
-    @trial_feature_enable
-    def post(self, trial_app):
-        app_model = trial_app
-        app_mode = AppMode.value_of(app_model.mode)
-        if app_mode not in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT}:
-            raise NotChatAppError()
-
-        parser = reqparse.RequestParser()
-        parser.add_argument("inputs", type=dict, required=True, location="json")
-        parser.add_argument("query", type=str, required=True, location="json")
-        parser.add_argument("files", type=list, required=False, location="json")
-        parser.add_argument("conversation_id", type=uuid_value, location="json")
-        parser.add_argument("parent_message_id", type=uuid_value, required=False, location="json")
-        parser.add_argument("retriever_from", type=str, required=False, default="explore_app", location="json")
-        args = parser.parse_args()
-
-        args["auto_generate_name"] = False
-
-        try:
-            if not isinstance(current_user, Account):
-                raise ValueError("current_user must be an Account instance")
-            
-            # Get IDs before they might be detached from session
-            app_id = app_model.id
-            user_id = current_user.id
-            
-            response = AppGenerateService.generate(
-                app_model=app_model, user=current_user, args=args, invoke_from=InvokeFrom.EXPLORE, streaming=True
-            )
-            RecommendedAppService.add_trial_app_record(app_id, user_id)
-            return helper.compact_generate_response(response)
-        except services.errors.conversation.ConversationNotExistsError:
-            raise NotFound("Conversation Not Exists.")
-        except services.errors.conversation.ConversationCompletedError:
-            raise ConversationCompletedError()
-        except services.errors.app_model_config.AppModelConfigBrokenError:
-            logger.exception("App model config broken.")
-            raise AppUnavailableError()
-        except ProviderTokenNotInitError as ex:
-            raise ProviderNotInitializeError(ex.description)
-        except QuotaExceededError:
-            raise ProviderQuotaExceededError()
-        except ModelCurrentlyNotSupportError:
-            raise ProviderModelCurrentlyNotSupportError()
-        except InvokeError as e:
-            raise CompletionRequestError(e.description)
-        except InvokeRateLimitError as ex:
-            raise InvokeRateLimitHttpError(ex.description)
-        except ValueError as e:
-            raise e
-        except Exception:
-            logger.exception("internal server error.")
-            raise InternalServerError()
-
-
-class TrialMessageSuggestedQuestionApi(TrialAppResource):
-    @trial_feature_enable
-    def get(self, trial_app, message_id):
-        app_model = trial_app
-        app_mode = AppMode.value_of(app_model.mode)
-        if app_mode not in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT}:
-            raise NotChatAppError()
-
-        message_id = str(message_id)
-
-        try:
-            if not isinstance(current_user, Account):
-                raise ValueError("current_user must be an Account instance")
-            questions = MessageService.get_suggested_questions_after_answer(
-                app_model=app_model, user=current_user, message_id=message_id, invoke_from=InvokeFrom.EXPLORE
-            )
-        except MessageNotExistsError:
-            raise NotFound("Message not found")
-        except ConversationNotExistsError:
-            raise NotFound("Conversation not found")
-        except SuggestedQuestionsAfterAnswerDisabledError:
-            raise AppSuggestedQuestionsAfterAnswerDisabledError()
-        except ProviderTokenNotInitError as ex:
-            raise ProviderNotInitializeError(ex.description)
-        except QuotaExceededError:
-            raise ProviderQuotaExceededError()
-        except ModelCurrentlyNotSupportError:
-            raise ProviderModelCurrentlyNotSupportError()
-        except InvokeError as e:
-            raise CompletionRequestError(e.description)
-        except Exception:
-            logger.exception("internal server error.")
-            raise InternalServerError()
-
-        return {"data": questions}
-
-
-class TrialChatAudioApi(TrialAppResource):
-    @trial_feature_enable
-    def post(self, trial_app):
-        app_model = trial_app
-
-        file = request.files["file"]
-
-        try:
-            if not isinstance(current_user, Account):
-                raise ValueError("current_user must be an Account instance")
-            
-            # Get IDs before they might be detached from session
-            app_id = app_model.id
-            user_id = current_user.id
-            
-            response = AudioService.transcript_asr(app_model=app_model, file=file, end_user=None)
-            RecommendedAppService.add_trial_app_record(app_id, user_id)
-            return response
-        except services.errors.app_model_config.AppModelConfigBrokenError:
-            logger.exception("App model config broken.")
-            raise AppUnavailableError()
-        except NoAudioUploadedServiceError:
-            raise NoAudioUploadedError()
-        except AudioTooLargeServiceError as e:
-            raise AudioTooLargeError(str(e))
-        except UnsupportedAudioTypeServiceError:
-            raise UnsupportedAudioTypeError()
-        except ProviderNotSupportSpeechToTextServiceError:
-            raise ProviderNotSupportSpeechToTextError()
-        except ProviderTokenNotInitError as ex:
-            raise ProviderNotInitializeError(ex.description)
-        except QuotaExceededError:
-            raise ProviderQuotaExceededError()
-        except ModelCurrentlyNotSupportError:
-            raise ProviderModelCurrentlyNotSupportError()
-        except InvokeError as e:
-            raise CompletionRequestError(e.description)
-        except ValueError as e:
-            raise e
-        except Exception as e:
-            logger.exception("internal server error.")
-            raise InternalServerError()
-
-
-class TrialChatTextApi(TrialAppResource):
-    @trial_feature_enable
-    def post(self, trial_app):
-        app_model = trial_app
-        try:
-            parser = reqparse.RequestParser()
-            parser.add_argument("message_id", type=str, required=False, location="json")
-            parser.add_argument("voice", type=str, location="json")
-            parser.add_argument("text", type=str, location="json")
-            parser.add_argument("streaming", type=bool, location="json")
-            args = parser.parse_args()
-
-            message_id = args.get("message_id", None)
-            text = args.get("text", None)
-            voice = args.get("voice", None)
-            if not isinstance(current_user, Account):
-                raise ValueError("current_user must be an Account instance")
-            
-            # Get IDs before they might be detached from session
-            app_id = app_model.id
-            user_id = current_user.id
-            
-            response = AudioService.transcript_tts(app_model=app_model, text=text, voice=voice, message_id=message_id)
-            RecommendedAppService.add_trial_app_record(app_id, user_id)
-            return response
-        except services.errors.app_model_config.AppModelConfigBrokenError:
-            logger.exception("App model config broken.")
-            raise AppUnavailableError()
-        except NoAudioUploadedServiceError:
-            raise NoAudioUploadedError()
-        except AudioTooLargeServiceError as e:
-            raise AudioTooLargeError(str(e))
-        except UnsupportedAudioTypeServiceError:
-            raise UnsupportedAudioTypeError()
-        except ProviderNotSupportSpeechToTextServiceError:
-            raise ProviderNotSupportSpeechToTextError()
-        except ProviderTokenNotInitError as ex:
-            raise ProviderNotInitializeError(ex.description)
-        except QuotaExceededError:
-            raise ProviderQuotaExceededError()
-        except ModelCurrentlyNotSupportError:
-            raise ProviderModelCurrentlyNotSupportError()
-        except InvokeError as e:
-            raise CompletionRequestError(e.description)
-        except ValueError as e:
-            raise e
-        except Exception as e:
-            logger.exception("internal server error.")
-            raise InternalServerError()
-
-
-class TrialCompletionApi(TrialAppResource):
-    @trial_feature_enable
-    def post(self, trial_app):
-        app_model = trial_app
-        if app_model.mode != "completion":
-            raise NotCompletionAppError()
-
-        parser = reqparse.RequestParser()
-        parser.add_argument("inputs", type=dict, required=True, location="json")
-        parser.add_argument("query", type=str, location="json", default="")
-        parser.add_argument("files", type=list, required=False, location="json")
-        parser.add_argument("response_mode", type=str, choices=["blocking", "streaming"], location="json")
-        parser.add_argument("retriever_from", type=str, required=False, default="explore_app", location="json")
-        args = parser.parse_args()
-
-        streaming = args["response_mode"] == "streaming"
-        args["auto_generate_name"] = False
-
-        try:
-            if not isinstance(current_user, Account):
-                raise ValueError("current_user must be an Account instance")
-            
-            # Get IDs before they might be detached from session
-            app_id = app_model.id
-            user_id = current_user.id
-            
-            response = AppGenerateService.generate(
-                app_model=app_model, user=current_user, args=args, invoke_from=InvokeFrom.EXPLORE, streaming=streaming
-            )
-
-            RecommendedAppService.add_trial_app_record(app_id, user_id)
-            return helper.compact_generate_response(response)
-        except services.errors.conversation.ConversationNotExistsError:
-            raise NotFound("Conversation Not Exists.")
-        except services.errors.conversation.ConversationCompletedError:
-            raise ConversationCompletedError()
-        except services.errors.app_model_config.AppModelConfigBrokenError:
-            logger.exception("App model config broken.")
-            raise AppUnavailableError()
-        except ProviderTokenNotInitError as ex:
-            raise ProviderNotInitializeError(ex.description)
-        except QuotaExceededError:
-            raise ProviderQuotaExceededError()
-        except ModelCurrentlyNotSupportError:
-            raise ProviderModelCurrentlyNotSupportError()
-        except InvokeError as e:
-            raise CompletionRequestError(e.description)
-        except ValueError as e:
-            raise e
-        except Exception:
-            logger.exception("internal server error.")
-            raise InternalServerError()
-
-
-class TrialSitApi(Resource):
-    """Resource for trial app sites."""
-
-    @trial_feature_enable
-    @get_app_model_with_trial
-    @service_api_ns.marshal_with(build_site_model(service_api_ns))
-    def get(self, app_model):
-        """Retrieve app site info.
-
-        Returns the site configuration for the application including theme, icons, and text.
-        """
-        site = db.session.query(Site).where(Site.app_id == app_model.id).first()
-
-        if not site:
-            raise Forbidden()
-
-        assert app_model.tenant
-        if app_model.tenant.status == TenantStatus.ARCHIVE:
-            raise Forbidden()
-
-        return site
-
-
-class TrialAppParameterApi(Resource):
-    """Resource for app variables."""
-
-    @trial_feature_enable
-    @get_app_model_with_trial
-    @marshal_with(fields.parameters_fields)
-    def get(self, app_model):
-        """Retrieve app parameters."""
-
-        if app_model is None:
-            raise AppUnavailableError()
-
-        if app_model.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}:
-            workflow = app_model.workflow
-            if workflow is None:
-                raise AppUnavailableError()
-
-            features_dict = workflow.features_dict
-            user_input_form = workflow.user_input_form(to_old_structure=True)
-        else:
-            app_model_config = app_model.app_model_config
-            if app_model_config is None:
-                raise AppUnavailableError()
-
-            features_dict = app_model_config.to_dict()
-
-            user_input_form = features_dict.get("user_input_form", [])
-
-        return get_parameters_from_feature_dict(features_dict=features_dict, user_input_form=user_input_form)
-
-
-class AppApi(Resource):
-    @trial_feature_enable
-    @get_app_model_with_trial
-    @marshal_with(app_detail_fields_with_site)
-    def get(self, app_model):
-        """Get app detail"""
-
-        app_service = AppService()
-        app_model = app_service.get_app(app_model)
-
-        return app_model
-
-
-class AppWorkflowApi(Resource):
-    @trial_feature_enable
-    @get_app_model_with_trial
-    @marshal_with(workflow_fields)
-    def get(self, app_model):
-        """Get workflow detail"""
-        if not app_model.workflow_id:
-            raise AppUnavailableError()
-
-        workflow = (
-            db.session.query(Workflow)
-            .where(
-                Workflow.id == app_model.workflow_id,
-            )
-            .first()
-        )
-        return workflow
-
-
-class DatasetListApi(Resource):
-    @trial_feature_enable
-    @get_app_model_with_trial
-    def get(self, app_model):
-        page = request.args.get("page", default=1, type=int)
-        limit = request.args.get("limit", default=20, type=int)
-        ids = request.args.getlist("ids")
-      
-        tenant_id = app_model.tenant_id
-        if ids:
-            datasets, total = DatasetService.get_datasets_by_ids(ids, tenant_id)
-        else:
-            raise NeedAddIdsError()
-
-        data = cast(list[dict[str, Any]], marshal(datasets, dataset_fields))
-
-        response = {"data": data, "has_more": len(datasets) == limit, "limit": limit, "total": total, "page": page}
-        return response
-
-
-api.add_resource(TrialChatApi, "/trial-apps/<uuid:app_id>/chat-messages", endpoint="trial_app_chat_completion")
-
-api.add_resource(
-    TrialMessageSuggestedQuestionApi,
-    "/trial-apps/<uuid:app_id>/messages/<uuid:message_id>/suggested-questions",
-    endpoint="trial_app_suggested_question",
-)
-
-api.add_resource(TrialChatAudioApi, "/trial-apps/<uuid:app_id>/audio-to-text", endpoint="trial_app_audio")
-api.add_resource(TrialChatTextApi, "/trial-apps/<uuid:app_id>/text-to-audio", endpoint="trial_app_text")
-
-api.add_resource(TrialCompletionApi, "/trial-apps/<uuid:app_id>/completion-messages", endpoint="trial_app_completion")
-
-api.add_resource(TrialSitApi, "/trial-apps/<uuid:app_id>/site")
-
-api.add_resource(TrialAppParameterApi, "/trial-apps/<uuid:app_id>/parameters", endpoint="trial_app_parameters")
-
-api.add_resource(AppApi, "/trial-apps/<uuid:app_id>", endpoint="trial_app")
-
-api.add_resource(TrialAppWorkflowRunApi, "/trial-apps/<uuid:app_id>/workflows/run", endpoint="trial_app_workflow_run")
-api.add_resource(TrialAppWorkflowTaskStopApi, "/trial-apps/<uuid:app_id>/workflows/tasks/<string:task_id>/stop")
-
-api.add_resource(AppWorkflowApi, "/trial-apps/<uuid:app_id>/workflows", endpoint="trial_app_workflow")
-api.add_resource(DatasetListApi, "/trial-apps/<uuid:app_id>/datasets", endpoint="trial_app_datasets")

api/controllers/console/explore/workflow.py

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

api/controllers/console/explore/wraps.py

@@ -2,16 +2,14 @@ from collections.abc import Callable
 from functools import wraps
 from typing import Concatenate, ParamSpec, TypeVar
 
-from flask import abort
-from flask_login import current_user
 from flask_restx import Resource
 from werkzeug.exceptions import NotFound
 
-from controllers.console.explore.error import AppAccessDeniedError, TrialAppLimitExceeded, TrialAppNotAllowed
+from controllers.console.explore.error import AppAccessDeniedError
 from controllers.console.wraps import account_initialization_required
 from extensions.ext_database import db
 from libs.login import current_account_with_tenant, login_required
-from models import AccountTrialAppRecord, App, InstalledApp, TrialApp
+from models import InstalledApp
 from services.enterprise.enterprise_service import EnterpriseService
 from services.feature_service import FeatureService
 
@@ -73,59 +71,6 @@ def user_allowed_to_access_app(view: Callable[Concatenate[InstalledApp, P], R] |
     return decorator
 
 
-def trial_app_required(view: Callable[Concatenate[App, P], R] | None = None):
-    def decorator(view: Callable[Concatenate[App, P], R]):
-        @wraps(view)
-        def decorated(app_id: str, *args: P.args, **kwargs: P.kwargs):
-            trial_app = db.session.query(TrialApp).where(TrialApp.app_id == str(app_id)).first()
-
-            if trial_app is None:
-                raise TrialAppNotAllowed()
-            app = trial_app.app
-
-            if app is None:
-                raise TrialAppNotAllowed()
-
-            account_trial_app_record = (
-                db.session.query(AccountTrialAppRecord)
-                .where(AccountTrialAppRecord.account_id == current_user.id, AccountTrialAppRecord.app_id == app_id)
-                .first()
-            )
-            if account_trial_app_record:
-                if account_trial_app_record.count >= trial_app.trial_limit:
-                    raise TrialAppLimitExceeded()
-
-            return view(app, *args, **kwargs)
-
-        return decorated
-
-    if view:
-        return decorator(view)
-    return decorator
-
-
-def trial_feature_enable(view: Callable[..., R]) -> Callable[..., R]:
-    @wraps(view)
-    def decorated(*args, **kwargs):
-        features = FeatureService.get_system_features()
-        if not features.enable_trial_app:
-            abort(403, "Trial app feature is not enabled.")
-        return view(*args, **kwargs)
-
-    return decorated
-
-
-def explore_banner_enabled(view: Callable[..., R]) -> Callable[..., R]:
-    @wraps(view)
-    def decorated(*args, **kwargs):
-        features = FeatureService.get_system_features()
-        if not features.enable_explore_banner:
-            abort(403, "Explore banner feature is not enabled.")
-        return view(*args, **kwargs)
-
-    return decorated
-
-
 class InstalledAppResource(Resource):
     # must be reversed if there are multiple decorators
 
@@ -135,13 +80,3 @@ class InstalledAppResource(Resource):
         account_initialization_required,
         login_required,
     ]
-
-
-class TrialAppResource(Resource):
-    # must be reversed if there are multiple decorators
-
-    method_decorators = [
-        trial_app_required,
-        account_initialization_required,
-        login_required,
-    ]

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

@@ -8,7 +8,6 @@ import services
 from configs import dify_config
 from constants import DOCUMENT_EXTENSIONS
 from controllers.common.errors import (
-    BlockedFileExtensionError,
     FilenameNotExistsError,
     FileTooLargeError,
     NoFileUploadedError,
@@ -40,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,
@@ -84,8 +82,6 @@ class FileApi(Resource):
             raise FileTooLargeError(file_too_large_error.description)
         except services.errors.file.UnsupportedFileTypeError:
             raise UnsupportedFileTypeError()
-        except services.errors.file.BlockedFileExtensionError as blocked_extension_error:
-            raise BlockedFileExtensionError(blocked_extension_error.description)
 
         return upload_file, 201
 

api/controllers/console/setup.py

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

api/controllers/console/workspace/account.py

@@ -139,7 +139,7 @@ class AccountAvatarApi(Resource):
 
         avatar_url = file_helpers.get_signed_file_url(args["avatar"])
         return {"avatar_url": avatar_url}
-
+    
     @setup_required
     @login_required
     @account_initialization_required

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/console/workspace/workspace.py

@@ -21,7 +21,6 @@ from controllers.console.wraps import (
     cloud_edition_billing_resource_check,
     setup_required,
 )
-from enums.cloud_plan import CloudPlan
 from extensions.ext_database import db
 from libs.helper import TimestampField
 from libs.login import current_account_with_tenant, login_required
@@ -51,9 +50,6 @@ tenant_fields = {
     "in_trial": fields.Boolean,
     "trial_end_reason": fields.String,
     "custom_config": fields.Raw(attribute="custom_config"),
-    "trial_credits": fields.Integer,
-    "trial_credits_used": fields.Integer,
-    "next_credit_reset_date": fields.Integer,
 }
 
 tenants_fields = {
@@ -87,7 +83,7 @@ class TenantListApi(Resource):
                 "name": tenant.name,
                 "status": tenant.status,
                 "created_at": tenant.created_at,
-                "plan": features.billing.subscription.plan if features.billing.enabled else CloudPlan.SANDBOX,
+                "plan": features.billing.subscription.plan if features.billing.enabled else "sandbox",
                 "current": tenant.id == current_tenant_id if current_tenant_id else False,
             }
 

api/controllers/console/wraps.py

@@ -10,7 +10,6 @@ from flask import abort, request
 
 from configs import dify_config
 from controllers.console.workspace.error import AccountNotInitializedError
-from enums.cloud_plan import CloudPlan
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from libs.login import current_account_with_tenant
@@ -134,7 +133,7 @@ def cloud_edition_billing_knowledge_limit_check(resource: str):
             features = FeatureService.get_features(current_tenant_id)
             if features.billing.enabled:
                 if resource == "add_segment":
-                    if features.billing.subscription.plan == CloudPlan.SANDBOX:
+                    if features.billing.subscription.plan == "sandbox":
                         abort(
                             403,
                             "To unlock this feature and elevate your Dify experience, please upgrade to a paid plan.",

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/controllers/mcp/mcp.py

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

api/controllers/service_api/__init__.py

@@ -19,7 +19,6 @@ from .app import (
     annotation,
     app,
     audio,
-    chatflow_memory,
     completion,
     conversation,
     file,
@@ -41,7 +40,6 @@ __all__ = [
     "annotation",
     "app",
     "audio",
-    "chatflow_memory",
     "completion",
     "conversation",
     "dataset",

api/controllers/service_api/app/chatflow_memory.py

@@ -1,109 +0,0 @@
-from flask_restx import Resource, reqparse
-
-from controllers.service_api import api
-from controllers.service_api.wraps import FetchUserArg, WhereisUserArg, validate_app_token
-from core.memory.entities import MemoryBlock, MemoryCreatedBy
-from core.workflow.runtime.variable_pool import VariablePool
-from models import App, EndUser
-from services.chatflow_memory_service import ChatflowMemoryService
-from services.workflow_service import WorkflowService
-
-
-class MemoryListApi(Resource):
-    @validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.JSON, required=True))
-    def get(self, app_model: App, end_user: EndUser):
-        parser = reqparse.RequestParser()
-        parser.add_argument("conversation_id", required=False, type=str | None, default=None)
-        parser.add_argument("memory_id", required=False, type=str | None, default=None)
-        parser.add_argument("version", required=False, type=int | None, default=None)
-        args = parser.parse_args()
-        conversation_id: str | None = args.get("conversation_id")
-        memory_id = args.get("memory_id")
-        version = args.get("version")
-
-        if conversation_id:
-            result = ChatflowMemoryService.get_persistent_memories_with_conversation(
-                app_model, MemoryCreatedBy(end_user_id=end_user.id), conversation_id, version
-            )
-            session_memories = ChatflowMemoryService.get_session_memories_with_conversation(
-                app_model, MemoryCreatedBy(end_user_id=end_user.id), conversation_id, version
-            )
-            result = [*result, *session_memories]
-        else:
-            result = ChatflowMemoryService.get_persistent_memories(
-                app_model, MemoryCreatedBy(end_user_id=end_user.id), version
-            )
-
-        if memory_id:
-            result = [it for it in result if it.spec.id == memory_id]
-        return [it for it in result if it.spec.end_user_visible]
-
-
-class MemoryEditApi(Resource):
-    @validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.JSON, required=True))
-    def put(self, app_model: App, end_user: EndUser):
-        parser = reqparse.RequestParser()
-        parser.add_argument("id", type=str, required=True)
-        parser.add_argument("conversation_id", type=str | None, required=False, default=None)
-        parser.add_argument("node_id", type=str | None, required=False, default=None)
-        parser.add_argument("update", type=str, required=True)
-        args = parser.parse_args()
-        workflow = WorkflowService().get_published_workflow(app_model)
-        update = args.get("update")
-        conversation_id = args.get("conversation_id")
-        node_id = args.get("node_id")
-        if not isinstance(update, str):
-            return {"error": "Invalid update"}, 400
-        if not workflow:
-            return {"error": "Workflow not found"}, 404
-        memory_spec = next((it for it in workflow.memory_blocks if it.id == args["id"]), None)
-        if not memory_spec:
-            return {"error": "Memory not found"}, 404
-
-        # First get existing memory
-        existing_memory = ChatflowMemoryService.get_memory_by_spec(
-            spec=memory_spec,
-            tenant_id=app_model.tenant_id,
-            app_id=app_model.id,
-            created_by=MemoryCreatedBy(end_user_id=end_user.id),
-            conversation_id=conversation_id,
-            node_id=node_id,
-            is_draft=False,
-        )
-
-        # Create updated memory instance with incremented version
-        updated_memory = MemoryBlock(
-            spec=existing_memory.spec,
-            tenant_id=existing_memory.tenant_id,
-            app_id=existing_memory.app_id,
-            conversation_id=existing_memory.conversation_id,
-            node_id=existing_memory.node_id,
-            value=update,  # New value
-            version=existing_memory.version + 1,  # Increment version for update
-            edited_by_user=True,
-            created_by=existing_memory.created_by,
-        )
-
-        ChatflowMemoryService.save_memory(updated_memory, VariablePool(), False)
-        return "", 204
-
-
-class MemoryDeleteApi(Resource):
-    @validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.JSON, required=True))
-    def delete(self, app_model: App, end_user: EndUser):
-        parser = reqparse.RequestParser()
-        parser.add_argument("id", type=str, required=False, default=None)
-        args = parser.parse_args()
-        memory_id = args.get("id")
-
-        if memory_id:
-            ChatflowMemoryService.delete_memory(app_model, memory_id, MemoryCreatedBy(end_user_id=end_user.id))
-            return "", 204
-        else:
-            ChatflowMemoryService.delete_all_user_memories(app_model, MemoryCreatedBy(end_user_id=end_user.id))
-            return "", 200
-
-
-api.add_resource(MemoryListApi, "/memories")
-api.add_resource(MemoryEditApi, "/memory-edit")
-api.add_resource(MemoryDeleteApi, "/memories")

api/controllers/service_api/dataset/document.py

@@ -592,7 +592,7 @@ class DocumentApi(DatasetApiResource):
                 "name": document.name,
                 "created_from": document.created_from,
                 "created_by": document.created_by,
-                "created_at": int(document.created_at.timestamp()),
+                "created_at": document.created_at.timestamp(),
                 "tokens": document.tokens,
                 "indexing_status": document.indexing_status,
                 "completed_at": int(document.completed_at.timestamp()) if document.completed_at else None,
@@ -625,7 +625,7 @@ class DocumentApi(DatasetApiResource):
                 "name": document.name,
                 "created_from": document.created_from,
                 "created_by": document.created_by,
-                "created_at": int(document.created_at.timestamp()),
+                "created_at": document.created_at.timestamp(),
                 "tokens": document.tokens,
                 "indexing_status": document.indexing_status,
                 "completed_at": int(document.completed_at.timestamp()) if document.completed_at else None,

api/controllers/service_api/dataset/segment.py

@@ -2,7 +2,6 @@ from flask import request
 from flask_restx import marshal, reqparse
 from werkzeug.exceptions import NotFound
 
-from configs import dify_config
 from controllers.service_api import service_api_ns
 from controllers.service_api.app.error import ProviderNotInitializeError
 from controllers.service_api.wraps import (
@@ -108,10 +107,6 @@ class SegmentApi(DatasetApiResource):
         # validate args
         args = segment_create_parser.parse_args()
         if args["segments"] is not None:
-            segments_limit = dify_config.DATASET_MAX_SEGMENTS_PER_REQUEST
-            if segments_limit > 0 and len(args["segments"]) > segments_limit:
-                raise ValueError(f"Exceeded maximum segments limit of {segments_limit}.")
-
             for args_item in args["segments"]:
                 SegmentService.segment_create_args_validate(args_item, document)
             segments = SegmentService.multi_create_segment(args["segments"], document, dataset)

api/controllers/service_api/wraps.py

@@ -13,7 +13,6 @@ from sqlalchemy import select, update
 from sqlalchemy.orm import Session
 from werkzeug.exceptions import Forbidden, NotFound, Unauthorized
 
-from enums.cloud_plan import CloudPlan
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from libs.datetime_utils import naive_utc_now
@@ -67,7 +66,6 @@ def validate_app_token(view: Callable[P, R] | None = None, *, fetch_user_arg: Fe
 
             kwargs["app_model"] = app_model
 
-            # If caller needs end-user context, attach EndUser to current_user
             if fetch_user_arg:
                 if fetch_user_arg.fetch_from == WhereisUserArg.QUERY:
                     user_id = request.args.get("user")
@@ -76,6 +74,7 @@ def validate_app_token(view: Callable[P, R] | None = None, *, fetch_user_arg: Fe
                 elif fetch_user_arg.fetch_from == WhereisUserArg.FORM:
                     user_id = request.form.get("user")
                 else:
+                    # use default-user
                     user_id = None
 
                 if not user_id and fetch_user_arg.required:
@@ -90,28 +89,6 @@ def validate_app_token(view: Callable[P, R] | None = None, *, fetch_user_arg: Fe
                 # Set EndUser as current logged-in user for flask_login.current_user
                 current_app.login_manager._update_request_context_with_user(end_user)  # type: ignore
                 user_logged_in.send(current_app._get_current_object(), user=end_user)  # type: ignore
-            else:
-                # For service API without end-user context, ensure an Account is logged in
-                # so services relying on current_account_with_tenant() work correctly.
-                tenant_owner_info = (
-                    db.session.query(Tenant, Account)
-                    .join(TenantAccountJoin, Tenant.id == TenantAccountJoin.tenant_id)
-                    .join(Account, TenantAccountJoin.account_id == Account.id)
-                    .where(
-                        Tenant.id == app_model.tenant_id,
-                        TenantAccountJoin.role == "owner",
-                        Tenant.status == TenantStatus.NORMAL,
-                    )
-                    .one_or_none()
-                )
-
-                if tenant_owner_info:
-                    tenant_model, account = tenant_owner_info
-                    account.current_tenant = tenant_model
-                    current_app.login_manager._update_request_context_with_user(account)  # type: ignore
-                    user_logged_in.send(current_app._get_current_object(), user=current_user)  # type: ignore
-                else:
-                    raise Unauthorized("Tenant owner account not found or tenant is not active.")
 
             return view_func(*args, **kwargs)
 
@@ -161,7 +138,7 @@ def cloud_edition_billing_knowledge_limit_check(resource: str, api_token_type: s
             features = FeatureService.get_features(api_token.tenant_id)
             if features.billing.enabled:
                 if resource == "add_segment":
-                    if features.billing.subscription.plan == CloudPlan.SANDBOX:
+                    if features.billing.subscription.plan == "sandbox":
                         raise Forbidden(
                             "To unlock this feature and elevate your Dify experience, please upgrade to a paid plan."
                         )

api/controllers/web/__init__.py

@@ -18,7 +18,6 @@ web_ns = Namespace("web", description="Web application API operations", path="/"
 from . import (
     app,
     audio,
-    chatflow_memory,
     completion,
     conversation,
     feature,
@@ -40,7 +39,6 @@ __all__ = [
     "app",
     "audio",
     "bp",
-    "chatflow_memory",
     "completion",
     "conversation",
     "feature",

api/controllers/web/chatflow_memory.py

@@ -1,108 +0,0 @@
-from flask_restx import reqparse
-
-from controllers.web import api
-from controllers.web.wraps import WebApiResource
-from core.memory.entities import MemoryBlock, MemoryCreatedBy
-from core.workflow.runtime.variable_pool import VariablePool
-from models import App, EndUser
-from services.chatflow_memory_service import ChatflowMemoryService
-from services.workflow_service import WorkflowService
-
-
-class MemoryListApi(WebApiResource):
-    def get(self, app_model: App, end_user: EndUser):
-        parser = reqparse.RequestParser()
-        parser.add_argument("conversation_id", required=False, type=str | None, default=None)
-        parser.add_argument("memory_id", required=False, type=str | None, default=None)
-        parser.add_argument("version", required=False, type=int | None, default=None)
-        args = parser.parse_args()
-        conversation_id: str | None = args.get("conversation_id")
-        memory_id = args.get("memory_id")
-        version = args.get("version")
-
-        if conversation_id:
-            result = ChatflowMemoryService.get_persistent_memories_with_conversation(
-                app_model, MemoryCreatedBy(end_user_id=end_user.id), conversation_id, version
-            )
-            session_memories = ChatflowMemoryService.get_session_memories_with_conversation(
-                app_model, MemoryCreatedBy(end_user_id=end_user.id), conversation_id, version
-            )
-            result = [*result, *session_memories]
-        else:
-            result = ChatflowMemoryService.get_persistent_memories(
-                app_model, MemoryCreatedBy(end_user_id=end_user.id), version
-            )
-
-        if memory_id:
-            result = [it for it in result if it.spec.id == memory_id]
-        return [it for it in result if it.spec.end_user_visible]
-
-
-class MemoryEditApi(WebApiResource):
-    def put(self, app_model: App, end_user: EndUser):
-        parser = reqparse.RequestParser()
-        parser.add_argument("id", type=str, required=True)
-        parser.add_argument("conversation_id", type=str | None, required=False, default=None)
-        parser.add_argument("node_id", type=str | None, required=False, default=None)
-        parser.add_argument("update", type=str, required=True)
-        args = parser.parse_args()
-        workflow = WorkflowService().get_published_workflow(app_model)
-        update = args.get("update")
-        conversation_id = args.get("conversation_id")
-        node_id = args.get("node_id")
-        if not isinstance(update, str):
-            return {"error": "Update must be a string"}, 400
-        if not workflow:
-            return {"error": "Workflow not found"}, 404
-        memory_spec = next((it for it in workflow.memory_blocks if it.id == args["id"]), None)
-        if not memory_spec:
-            return {"error": "Memory not found"}, 404
-        if not memory_spec.end_user_editable:
-            return {"error": "Memory not editable"}, 403
-
-        # First get existing memory
-        existing_memory = ChatflowMemoryService.get_memory_by_spec(
-            spec=memory_spec,
-            tenant_id=app_model.tenant_id,
-            app_id=app_model.id,
-            created_by=MemoryCreatedBy(end_user_id=end_user.id),
-            conversation_id=conversation_id,
-            node_id=node_id,
-            is_draft=False,
-        )
-
-        # Create updated memory instance with incremented version
-        updated_memory = MemoryBlock(
-            spec=existing_memory.spec,
-            tenant_id=existing_memory.tenant_id,
-            app_id=existing_memory.app_id,
-            conversation_id=existing_memory.conversation_id,
-            node_id=existing_memory.node_id,
-            value=update,  # New value
-            version=existing_memory.version + 1,  # Increment version for update
-            edited_by_user=True,
-            created_by=existing_memory.created_by,
-        )
-
-        ChatflowMemoryService.save_memory(updated_memory, VariablePool(), False)
-        return "", 204
-
-
-class MemoryDeleteApi(WebApiResource):
-    def delete(self, app_model: App, end_user: EndUser):
-        parser = reqparse.RequestParser()
-        parser.add_argument("id", type=str, required=False, default=None)
-        args = parser.parse_args()
-        memory_id = args.get("id")
-
-        if memory_id:
-            ChatflowMemoryService.delete_memory(app_model, memory_id, MemoryCreatedBy(end_user_id=end_user.id))
-            return "", 204
-        else:
-            ChatflowMemoryService.delete_all_user_memories(app_model, MemoryCreatedBy(end_user_id=end_user.id))
-            return "", 200
-
-
-api.add_resource(MemoryListApi, "/memories")
-api.add_resource(MemoryEditApi, "/memory-edit")
-api.add_resource(MemoryDeleteApi, "/memories")

api/controllers/web/login.py

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

api/controllers/web/passport.py

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

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

@@ -1,11 +1,10 @@
 import logging
 import time
-from collections.abc import Mapping, MutableMapping, Sequence
+from collections.abc import Mapping
 from typing import Any, cast
 
 from sqlalchemy import select
 from sqlalchemy.orm import Session
-from typing_extensions import override
 
 from core.app.apps.advanced_chat.app_config_manager import AdvancedChatAppConfig
 from core.app.apps.base_app_queue_manager import AppQueueManager
@@ -21,16 +20,12 @@ from core.app.entities.queue_entities import (
     QueueTextChunkEvent,
 )
 from core.app.features.annotation_reply.annotation_reply import AnnotationReplyFeature
-from core.memory.entities import MemoryCreatedBy, MemoryScope
-from core.model_runtime.entities import AssistantPromptMessage, UserPromptMessage
 from core.moderation.base import ModerationError
 from core.moderation.input_moderation import InputModeration
 from core.variables.variables import VariableUnion
 from core.workflow.enums import WorkflowType
 from core.workflow.graph_engine.command_channels.redis_channel import RedisChannel
-from core.workflow.graph_engine.layers.base import GraphEngineLayer
 from core.workflow.graph_engine.layers.persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer
-from core.workflow.graph_events import GraphRunSucceededEvent
 from core.workflow.repositories.workflow_execution_repository import WorkflowExecutionRepository
 from core.workflow.repositories.workflow_node_execution_repository import WorkflowNodeExecutionRepository
 from core.workflow.runtime import GraphRuntimeState, VariablePool
@@ -43,8 +38,6 @@ from models import Workflow
 from models.enums import UserFrom
 from models.model import App, Conversation, Message, MessageAnnotation
 from models.workflow import ConversationVariable
-from services.chatflow_history_service import ChatflowHistoryService
-from services.chatflow_memory_service import ChatflowMemoryService
 
 logger = logging.getLogger(__name__)
 
@@ -68,13 +61,11 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
         app: App,
         workflow_execution_repository: WorkflowExecutionRepository,
         workflow_node_execution_repository: WorkflowNodeExecutionRepository,
-        graph_engine_layers: Sequence[GraphEngineLayer] = (),
     ):
         super().__init__(
             queue_manager=queue_manager,
             variable_loader=variable_loader,
             app_id=application_generate_entity.app_config.app_id,
-            graph_engine_layers=graph_engine_layers,
         )
         self.application_generate_entity = application_generate_entity
         self.conversation = conversation
@@ -87,10 +78,6 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
         self._workflow_node_execution_repository = workflow_node_execution_repository
 
     def run(self):
-        ChatflowMemoryService.wait_for_sync_memory_completion(
-            workflow=self._workflow, conversation_id=self.conversation.id
-        )
-
         app_config = self.application_generate_entity.app_config
         app_config = cast(AdvancedChatAppConfig, app_config)
 
@@ -153,7 +140,6 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
                 # Based on the definition of `VariableUnion`,
                 # `list[Variable]` can be safely used as `list[VariableUnion]` since they are compatible.
                 conversation_variables=conversation_variables,
-                memory_blocks=self._fetch_memory_blocks(),
             )
 
             # init graph
@@ -209,39 +195,12 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
         )
 
         workflow_entry.graph_engine.layer(persistence_layer)
-        for layer in self._graph_engine_layers:
-            workflow_entry.graph_engine.layer(layer)
 
         generator = workflow_entry.run()
 
         for event in generator:
             self._handle_event(workflow_entry, event)
 
-        try:
-            self._check_app_memory_updates(variable_pool)
-        except Exception as e:
-            logger.exception("Failed to check app memory updates", exc_info=e)
-
-    @override
-    def _handle_event(self, workflow_entry: WorkflowEntry, event: Any) -> None:
-        super()._handle_event(workflow_entry, event)
-        if isinstance(event, GraphRunSucceededEvent):
-            workflow_outputs = event.outputs
-            if not workflow_outputs:
-                logger.warning("Chatflow output is empty.")
-                return
-            assistant_message = workflow_outputs.get("answer")
-            if not assistant_message:
-                logger.warning("Chatflow output does not contain 'answer'.")
-                return
-            if not isinstance(assistant_message, str):
-                logger.warning("Chatflow output 'answer' is not a string.")
-                return
-            try:
-                self._sync_conversation_to_chatflow_tables(assistant_message)
-            except Exception as e:
-                logger.exception("Failed to sync conversation to memory tables", exc_info=e)
-
     def handle_input_moderation(
         self,
         app_record: App,
@@ -439,67 +398,3 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
 
         # Return combined list
         return existing_variables + new_variables
-
-    def _fetch_memory_blocks(self) -> Mapping[str, str]:
-        """fetch all memory blocks for current app"""
-
-        memory_blocks_dict: MutableMapping[str, str] = {}
-        is_draft = self.application_generate_entity.invoke_from == InvokeFrom.DEBUGGER
-        conversation_id = self.conversation.id
-        memory_block_specs = self._workflow.memory_blocks
-        # Get runtime memory values
-        memories = ChatflowMemoryService.get_memories_by_specs(
-            memory_block_specs=memory_block_specs,
-            tenant_id=self._workflow.tenant_id,
-            app_id=self._workflow.app_id,
-            node_id=None,
-            conversation_id=conversation_id,
-            is_draft=is_draft,
-            created_by=self._get_created_by(),
-        )
-
-        # Build memory_id -> value mapping
-        for memory in memories:
-            if memory.spec.scope == MemoryScope.APP:
-                # App level: use memory_id directly
-                memory_blocks_dict[memory.spec.id] = memory.value
-            else:  # NODE scope
-                node_id = memory.node_id
-                if not node_id:
-                    logger.warning("Memory block %s has no node_id, skip.", memory.spec.id)
-                    continue
-                key = f"{node_id}.{memory.spec.id}"
-                memory_blocks_dict[key] = memory.value
-
-        return memory_blocks_dict
-
-    def _sync_conversation_to_chatflow_tables(self, assistant_message: str):
-        ChatflowHistoryService.save_app_message(
-            prompt_message=UserPromptMessage(content=(self.application_generate_entity.query)),
-            conversation_id=self.conversation.id,
-            app_id=self._workflow.app_id,
-            tenant_id=self._workflow.tenant_id,
-        )
-        ChatflowHistoryService.save_app_message(
-            prompt_message=AssistantPromptMessage(content=assistant_message),
-            conversation_id=self.conversation.id,
-            app_id=self._workflow.app_id,
-            tenant_id=self._workflow.tenant_id,
-        )
-
-    def _check_app_memory_updates(self, variable_pool: VariablePool):
-        is_draft = self.application_generate_entity.invoke_from == InvokeFrom.DEBUGGER
-
-        ChatflowMemoryService.update_app_memory_if_needed(
-            workflow=self._workflow,
-            conversation_id=self.conversation.id,
-            variable_pool=variable_pool,
-            is_draft=is_draft,
-            created_by=self._get_created_by(),
-        )
-
-    def _get_created_by(self) -> MemoryCreatedBy:
-        if self.application_generate_entity.invoke_from in {InvokeFrom.DEBUGGER, InvokeFrom.EXPLORE}:
-            return MemoryCreatedBy(account_id=self.application_generate_entity.user_id)
-        else:
-            return MemoryCreatedBy(end_user_id=self.application_generate_entity.user_id)

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

@@ -1,4 +1,3 @@
-import json
 import logging
 import re
 import time
@@ -61,7 +60,6 @@ from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTas
 from core.app.task_pipeline.message_cycle_manager import MessageCycleManager
 from core.base.tts import AppGeneratorTTSPublisher, AudioTrunk
 from core.model_runtime.entities.llm_entities import LLMUsage
-from core.model_runtime.utils.encoders import jsonable_encoder
 from core.ops.ops_trace_manager import TraceQueueManager
 from core.workflow.enums import WorkflowExecutionStatus
 from core.workflow.nodes import NodeType
@@ -393,14 +391,6 @@ class AdvancedChatAppGenerateTaskPipeline(GraphRuntimeStateSupport):
         if should_direct_answer:
             return
 
-        current_time = time.perf_counter()
-        if self._task_state.first_token_time is None and delta_text.strip():
-            self._task_state.first_token_time = current_time
-            self._task_state.is_streaming_response = True
-
-        if delta_text.strip():
-            self._task_state.last_token_time = current_time
-
         # Only publish tts message at text chunk streaming
         if tts_publisher and queue_message:
             tts_publisher.publish(queue_message)
@@ -782,33 +772,7 @@ class AdvancedChatAppGenerateTaskPipeline(GraphRuntimeStateSupport):
         message.answer = answer_text
         message.updated_at = naive_utc_now()
         message.provider_response_latency = time.perf_counter() - self._base_task_pipeline.start_at
-
-        # Set usage first before dumping metadata
-        if graph_runtime_state and graph_runtime_state.llm_usage:
-            usage = graph_runtime_state.llm_usage
-            message.message_tokens = usage.prompt_tokens
-            message.message_unit_price = usage.prompt_unit_price
-            message.message_price_unit = usage.prompt_price_unit
-            message.answer_tokens = usage.completion_tokens
-            message.answer_unit_price = usage.completion_unit_price
-            message.answer_price_unit = usage.completion_price_unit
-            message.total_price = usage.total_price
-            message.currency = usage.currency
-            self._task_state.metadata.usage = usage
-        else:
-            usage = LLMUsage.empty_usage()
-            self._task_state.metadata.usage = usage
-
-        # Add streaming metrics to usage if available
-        if self._task_state.is_streaming_response and self._task_state.first_token_time:
-            start_time = self._base_task_pipeline.start_at
-            first_token_time = self._task_state.first_token_time
-            last_token_time = self._task_state.last_token_time or first_token_time
-            usage.time_to_first_token = round(first_token_time - start_time, 3)
-            usage.time_to_generate = round(last_token_time - first_token_time, 3)
-
-        metadata = self._task_state.metadata.model_dump()
-        message.message_metadata = json.dumps(jsonable_encoder(metadata))
+        message.message_metadata = self._task_state.metadata.model_dump_json()
         message_files = [
             MessageFile(
                 message_id=message.id,
@@ -826,6 +790,20 @@ class AdvancedChatAppGenerateTaskPipeline(GraphRuntimeStateSupport):
         ]
         session.add_all(message_files)
 
+        if graph_runtime_state and graph_runtime_state.llm_usage:
+            usage = graph_runtime_state.llm_usage
+            message.message_tokens = usage.prompt_tokens
+            message.message_unit_price = usage.prompt_unit_price
+            message.message_price_unit = usage.prompt_price_unit
+            message.answer_tokens = usage.completion_tokens
+            message.answer_unit_price = usage.completion_unit_price
+            message.answer_price_unit = usage.completion_price_unit
+            message.total_price = usage.total_price
+            message.currency = usage.currency
+            self._task_state.metadata.usage = usage
+        else:
+            self._task_state.metadata.usage = LLMUsage.empty_usage()
+
     def _seed_graph_runtime_state_from_queue_manager(self) -> None:
         """Bootstrap the cached runtime state from the queue manager when present."""
         candidate = self._base_task_pipeline.queue_manager.graph_runtime_state

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

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

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

@@ -144,7 +144,7 @@ class AgentChatAppRunner(AppRunner):
             prompt_template_entity=app_config.prompt_template,
             inputs=dict(inputs),
             files=list(files),
-            query=query,
+            query=query or "",
             memory=memory,
         )
 
@@ -172,7 +172,7 @@ class AgentChatAppRunner(AppRunner):
             prompt_template_entity=app_config.prompt_template,
             inputs=dict(inputs),
             files=list(files),
-            query=query,
+            query=query or "",
             memory=memory,
         )
 

api/core/app/apps/base_app_runner.py

@@ -79,7 +79,7 @@ class AppRunner:
         prompt_template_entity: PromptTemplateEntity,
         inputs: Mapping[str, str],
         files: Sequence["File"],
-        query: str = "",
+        query: str | None = None,
         context: str | None = None,
         memory: TokenBufferMemory | None = None,
         image_detail_config: ImagePromptMessageContent.DETAIL | None = None,
@@ -105,7 +105,7 @@ class AppRunner:
                 app_mode=AppMode.value_of(app_record.mode),
                 prompt_template_entity=prompt_template_entity,
                 inputs=inputs,
-                query=query,
+                query=query or "",
                 files=files,
                 context=context,
                 memory=memory,

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/message_based_app_generator.py

@@ -190,7 +190,7 @@ class MessageBasedAppGenerator(BaseAppGenerator):
             override_model_configs=json.dumps(override_model_configs) if override_model_configs else None,
             conversation_id=conversation.id,
             inputs=application_generate_entity.inputs,
-            query=application_generate_entity.query,
+            query=application_generate_entity.query or "",
             message="",
             message_tokens=0,
             message_unit_price=0,

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

@@ -40,15 +40,19 @@ from core.workflow.repositories.draft_variable_repository import DraftVariableSa
 from core.workflow.repositories.workflow_execution_repository import WorkflowExecutionRepository
 from core.workflow.repositories.workflow_node_execution_repository import WorkflowNodeExecutionRepository
 from core.workflow.variable_loader import DUMMY_VARIABLE_LOADER, VariableLoader
-from enums.cloud_plan import CloudPlan
 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.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__)
 
@@ -251,7 +255,7 @@ class PipelineGenerator(BaseAppGenerator):
             json_text = json.dumps(text)
             upload_file = FileService(db.engine).upload_text(json_text, name, user.id, dataset.tenant_id)
             features = FeatureService.get_features(dataset.tenant_id)
-            if features.billing.enabled and features.billing.subscription.plan == CloudPlan.SANDBOX:
+            if features.billing.subscription.plan == "sandbox":
                 tenant_pipeline_task_key = f"tenant_pipeline_task:{dataset.tenant_id}"
                 tenant_self_pipeline_task_queue = f"tenant_self_pipeline_task_queue:{dataset.tenant_id}"
 

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

@@ -135,8 +135,6 @@ class WorkflowAppRunner(WorkflowBasedAppRunner):
         )
 
         workflow_entry.graph_engine.layer(persistence_layer)
-        for layer in self._graph_engine_layers:
-            workflow_entry.graph_engine.layer(layer)
 
         generator = workflow_entry.run()
 

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

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

api/core/app/apps/workflow_app_runner.py

@@ -1,5 +1,5 @@
 import time
-from collections.abc import Mapping, Sequence
+from collections.abc import Mapping
 from typing import Any, cast
 
 from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
@@ -27,7 +27,6 @@ from core.app.entities.queue_entities import (
 )
 from core.workflow.entities import GraphInitParams
 from core.workflow.graph import Graph
-from core.workflow.graph_engine.layers.base import GraphEngineLayer
 from core.workflow.graph_events import (
     GraphEngineEvent,
     GraphRunFailedEvent,
@@ -70,12 +69,10 @@ class WorkflowBasedAppRunner:
         queue_manager: AppQueueManager,
         variable_loader: VariableLoader = DUMMY_VARIABLE_LOADER,
         app_id: str,
-        graph_engine_layers: Sequence[GraphEngineLayer] = (),
     ):
         self._queue_manager = queue_manager
         self._variable_loader = variable_loader
         self._app_id = app_id
-        self._graph_engine_layers = graph_engine_layers
 
     def _init_graph(
         self,

api/core/app/entities/app_invoke_entities.py

@@ -129,7 +129,7 @@ class EasyUIBasedAppGenerateEntity(AppGenerateEntity):
     app_config: EasyUIBasedAppConfig = None  # type: ignore
     model_conf: ModelConfigWithCredentialsEntity
 
-    query: str = ""
+    query: str | None = None
 
     # pydantic configs
     model_config = ConfigDict(protected_namespaces=())

api/core/app/entities/task_entities.py

@@ -48,9 +48,6 @@ class WorkflowTaskState(TaskState):
     """
 
     answer: str = ""
-    first_token_time: float | None = None
-    last_token_time: float | None = None
-    is_streaming_response: bool = False
 
 
 class StreamEvent(StrEnum):

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

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

api/core/app/layers/pause_state_persist_layer.py

@@ -1,71 +0,0 @@
-from sqlalchemy import Engine
-from sqlalchemy.orm import sessionmaker
-
-from core.workflow.graph_engine.layers.base import GraphEngineLayer
-from core.workflow.graph_events.base import GraphEngineEvent
-from core.workflow.graph_events.graph import GraphRunPausedEvent
-from repositories.api_workflow_run_repository import APIWorkflowRunRepository
-from repositories.factory import DifyAPIRepositoryFactory
-
-
-class PauseStatePersistenceLayer(GraphEngineLayer):
-    def __init__(self, session_factory: Engine | sessionmaker, state_owner_user_id: str):
-        """Create a PauseStatePersistenceLayer.
-
-        The `state_owner_user_id` is used when creating state file for pause.
-        It generally should id of the creator of workflow.
-        """
-        if isinstance(session_factory, Engine):
-            session_factory = sessionmaker(session_factory)
-        self._session_maker = session_factory
-        self._state_owner_user_id = state_owner_user_id
-
-    def _get_repo(self) -> APIWorkflowRunRepository:
-        return DifyAPIRepositoryFactory.create_api_workflow_run_repository(self._session_maker)
-
-    def on_graph_start(self) -> None:
-        """
-        Called when graph execution starts.
-
-        This is called after the engine has been initialized but before any nodes
-        are executed. Layers can use this to set up resources or log start information.
-        """
-        pass
-
-    def on_event(self, event: GraphEngineEvent) -> None:
-        """
-        Called for every event emitted by the engine.
-
-        This method receives all events generated during graph execution, including:
-        - Graph lifecycle events (start, success, failure)
-        - Node execution events (start, success, failure, retry)
-        - Stream events for response nodes
-        - Container events (iteration, loop)
-
-        Args:
-            event: The event emitted by the engine
-        """
-        if not isinstance(event, GraphRunPausedEvent):
-            return
-
-        assert self.graph_runtime_state is not None
-        workflow_run_id: str | None = self.graph_runtime_state.system_variable.workflow_execution_id
-        assert workflow_run_id is not None
-        repo = self._get_repo()
-        repo.create_workflow_pause(
-            workflow_run_id=workflow_run_id,
-            state_owner_user_id=self._state_owner_user_id,
-            state=self.graph_runtime_state.dumps(),
-        )
-
-    def on_graph_end(self, error: Exception | None) -> None:
-        """
-        Called when graph execution ends.
-
-        This is called after all nodes have been executed or when execution is
-        aborted. Layers can use this to clean up resources or log final state.
-
-        Args:
-            error: The exception that caused execution to fail, or None if successful
-        """
-        pass

api/core/app/task_pipeline/based_generate_task_pipeline.py

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

api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py

@@ -121,7 +121,7 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline):
         if self._application_generate_entity.app_config.app_mode != AppMode.COMPLETION:
             # start generate conversation name thread
             self._conversation_name_generate_thread = self._message_cycle_manager.generate_conversation_name(
-                conversation_id=self._conversation_id, query=self._application_generate_entity.query
+                conversation_id=self._conversation_id, query=self._application_generate_entity.query or ""
             )
 
         generator = self._wrapper_process_stream_response(trace_manager=self._application_generate_entity.trace_manager)

api/core/app/task_pipeline/message_cycle_manager.py

@@ -140,27 +140,7 @@ class MessageCycleManager:
         if not self._application_generate_entity.app_config.additional_features:
             raise ValueError("Additional features not found")
         if self._application_generate_entity.app_config.additional_features.show_retrieve_source:
-            merged_resources = [r for r in self._task_state.metadata.retriever_resources or [] if r]
-            existing_ids = {(r.dataset_id, r.document_id) for r in merged_resources if r.dataset_id and r.document_id}
-
-            # Add new unique resources from the event
-            for resource in event.retriever_resources or []:
-                if not resource:
-                    continue
-
-                is_duplicate = (
-                    resource.dataset_id
-                    and resource.document_id
-                    and (resource.dataset_id, resource.document_id) in existing_ids
-                )
-
-                if not is_duplicate:
-                    merged_resources.append(resource)
-
-            for i, resource in enumerate(merged_resources, 1):
-                resource.position = i
-
-            self._task_state.metadata.retriever_resources = merged_resources
+            self._task_state.metadata.retriever_resources = event.retriever_resources
 
     def message_file_to_stream_response(self, event: QueueMessageFileEvent) -> MessageFileStreamResponse | None:
         """

api/core/entities/document_task.py

@@ -1,13 +0,0 @@
-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
-from core.tools.utils.encryption import create_provider_encrypter
-
-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"""
-
-        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/entities/provider_configuration.py

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

api/core/file/models.py

@@ -74,10 +74,6 @@ class File(BaseModel):
         storage_key: str | None = None,
         dify_model_identity: str | None = FILE_MODEL_IDENTITY,
         url: str | None = None,
-        # Legacy compatibility fields - explicitly handle known extra fields
-        tool_file_id: str | None = None,
-        upload_file_id: str | None = None,
-        datasource_file_id: str | None = None,
     ):
         super().__init__(
             id=id,

api/core/helper/code_executor/javascript/javascript_transformer.py

@@ -6,7 +6,10 @@ from core.helper.code_executor.template_transformer import TemplateTransformer
 class NodeJsTemplateTransformer(TemplateTransformer):
     @classmethod
     def get_runner_script(cls) -> str:
-        runner_script = dedent(f"""            {cls._code_placeholder}
+        runner_script = dedent(
+            f"""
+            // declare main function
+            {cls._code_placeholder}
 
             // decode and prepare input object
             var inputs_obj = JSON.parse(Buffer.from('{cls._inputs_placeholder}', 'base64').toString('utf-8'))
@@ -18,5 +21,6 @@ class NodeJsTemplateTransformer(TemplateTransformer):
             var output_json = JSON.stringify(output_obj)
             var result = `<<RESULT>>${{output_json}}<<RESULT>>`
             console.log(result)
-            """)
+            """
+        )
         return runner_script

api/core/helper/code_executor/python3/python3_transformer.py

@@ -6,7 +6,9 @@ from core.helper.code_executor.template_transformer import TemplateTransformer
 class Python3TemplateTransformer(TemplateTransformer):
     @classmethod
     def get_runner_script(cls) -> str:
-        runner_script = dedent(f"""            {cls._code_placeholder}
+        runner_script = dedent(f"""
+            # declare main function
+            {cls._code_placeholder}
 
             import json
             from base64 import b64decode

api/core/helper/marketplace.py

@@ -29,18 +29,6 @@ def batch_fetch_plugin_manifests(plugin_ids: list[str]) -> Sequence[MarketplaceP
     return [MarketplacePluginDeclaration.model_validate(plugin) for plugin in response.json()["data"]["plugins"]]
 
 
-def batch_fetch_plugin_by_ids(plugin_ids: list[str]) -> list[dict]:
-    if not plugin_ids:
-        return []
-
-    url = str(marketplace_api_url / "api/v1/plugins/batch")
-    response = httpx.post(url, json={"plugin_ids": plugin_ids}, headers={"X-Dify-Version": dify_config.project.version})
-    response.raise_for_status()
-
-    data = response.json()
-    return data.get("data", {}).get("plugins", [])
-
-
 def batch_fetch_plugin_manifests_ignore_deserialization_error(
     plugin_ids: list[str],
 ) -> Sequence[MarketplacePluginDeclaration]:

api/core/helper/module_import_helper.py

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

api/core/hosting_configuration.py

@@ -56,9 +56,6 @@ class HostingConfiguration:
         self.provider_map[f"{DEFAULT_PLUGIN_ID}/minimax/minimax"] = self.init_minimax()
         self.provider_map[f"{DEFAULT_PLUGIN_ID}/spark/spark"] = self.init_spark()
         self.provider_map[f"{DEFAULT_PLUGIN_ID}/zhipuai/zhipuai"] = self.init_zhipuai()
-        self.provider_map[f"{DEFAULT_PLUGIN_ID}/gemini/google"] = self.init_gemini()
-        self.provider_map[f"{DEFAULT_PLUGIN_ID}/x/x"] = self.init_xai()
-        self.provider_map[f"{DEFAULT_PLUGIN_ID}/deepseek/deepseek"] = self.init_deepseek()
 
         self.moderation_config = self.init_moderation_config()
 
@@ -131,7 +128,7 @@ class HostingConfiguration:
         quotas: list[HostingQuota] = []
 
         if dify_config.HOSTED_OPENAI_TRIAL_ENABLED:
-            hosted_quota_limit = 0
+            hosted_quota_limit = dify_config.HOSTED_OPENAI_QUOTA_LIMIT
             trial_models = self.parse_restrict_models_from_env("HOSTED_OPENAI_TRIAL_MODELS")
             trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trial_models)
             quotas.append(trial_quota)
@@ -159,49 +156,18 @@ class HostingConfiguration:
             quota_unit=quota_unit,
         )
 
-    def init_gemini(self) -> HostingProvider:
-        quota_unit = QuotaUnit.CREDITS
-        quotas: list[HostingQuota] = []
-
-        if dify_config.HOSTED_GEMINI_TRIAL_ENABLED:
-            hosted_quota_limit = 0
-            trial_models = self.parse_restrict_models_from_env("HOSTED_GEMINI_TRIAL_MODELS")
-            trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trial_models)
-            quotas.append(trial_quota)
-
-        if dify_config.HOSTED_GEMINI_PAID_ENABLED:
-            paid_models = self.parse_restrict_models_from_env("HOSTED_GEMINI_PAID_MODELS")
-            paid_quota = PaidHostingQuota(restrict_models=paid_models)
-            quotas.append(paid_quota)
-
-        if len(quotas) > 0:
-            credentials = {
-                "google_api_key": dify_config.HOSTED_GEMINI_API_KEY,
-            }
-
-            if dify_config.HOSTED_GEMINI_API_BASE:
-                credentials["google_base_url"] = dify_config.HOSTED_GEMINI_API_BASE
-
-            return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas)
-
-        return HostingProvider(
-            enabled=False,
-            quota_unit=quota_unit,
-        )
-
-    def init_anthropic(self) -> HostingProvider:
-        quota_unit = QuotaUnit.CREDITS
+    @staticmethod
+    def init_anthropic() -> HostingProvider:
+        quota_unit = QuotaUnit.TOKENS
         quotas: list[HostingQuota] = []
 
         if dify_config.HOSTED_ANTHROPIC_TRIAL_ENABLED:
-            hosted_quota_limit = 0
-            trail_models = self.parse_restrict_models_from_env("HOSTED_ANTHROPIC_TRIAL_MODELS")
-            trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trail_models)
+            hosted_quota_limit = dify_config.HOSTED_ANTHROPIC_QUOTA_LIMIT
+            trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit)
             quotas.append(trial_quota)
 
         if dify_config.HOSTED_ANTHROPIC_PAID_ENABLED:
-            paid_models = self.parse_restrict_models_from_env("HOSTED_ANTHROPIC_PAID_MODELS")
-            paid_quota = PaidHostingQuota(restrict_models=paid_models)
+            paid_quota = PaidHostingQuota()
             quotas.append(paid_quota)
 
         if len(quotas) > 0:
@@ -219,66 +185,6 @@ class HostingConfiguration:
             quota_unit=quota_unit,
         )
 
-    def init_xai(self) -> HostingProvider:
-        quota_unit = QuotaUnit.CREDITS
-        quotas: list[HostingQuota] = []
-
-        if dify_config.HOSTED_XAI_TRIAL_ENABLED:
-            hosted_quota_limit = 0
-            trail_models = self.parse_restrict_models_from_env("HOSTED_XAI_TRIAL_MODELS")
-            trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trail_models)
-            quotas.append(trial_quota)
-
-        if dify_config.HOSTED_XAI_PAID_ENABLED:
-            paid_models = self.parse_restrict_models_from_env("HOSTED_XAI_PAID_MODELS")
-            paid_quota = PaidHostingQuota(restrict_models=paid_models)
-            quotas.append(paid_quota)
-
-        if len(quotas) > 0:
-            credentials = {
-                "api_key": dify_config.HOSTED_XAI_API_KEY,
-            }
-
-            if dify_config.HOSTED_XAI_API_BASE:
-                credentials["endpoint_url"] = dify_config.HOSTED_XAI_API_BASE
-
-            return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas)
-
-        return HostingProvider(
-            enabled=False,
-            quota_unit=quota_unit,
-        )
-
-    def init_deepseek(self) -> HostingProvider:
-        quota_unit = QuotaUnit.CREDITS
-        quotas: list[HostingQuota] = []
-
-        if dify_config.HOSTED_DEEPSEEK_TRIAL_ENABLED:
-            hosted_quota_limit = 0
-            trail_models = self.parse_restrict_models_from_env("HOSTED_DEEPSEEK_TRIAL_MODELS")
-            trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trail_models)
-            quotas.append(trial_quota)
-
-        if dify_config.HOSTED_DEEPSEEK_PAID_ENABLED:
-            paid_models = self.parse_restrict_models_from_env("HOSTED_DEEPSEEK_PAID_MODELS")
-            paid_quota = PaidHostingQuota(restrict_models=paid_models)
-            quotas.append(paid_quota)
-
-        if len(quotas) > 0:
-            credentials = {
-                "api_key": dify_config.HOSTED_DEEPSEEK_API_KEY,
-            }
-
-            if dify_config.HOSTED_DEEPSEEK_API_BASE:
-                credentials["endpoint_url"] = dify_config.HOSTED_DEEPSEEK_API_BASE
-
-            return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas)
-
-        return HostingProvider(
-            enabled=False,
-            quota_unit=quota_unit,
-        )
-
     @staticmethod
     def init_minimax() -> HostingProvider:
         quota_unit = QuotaUnit.TOKENS

api/core/indexing_runner.py

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

api/core/llm_generator/llm_generator.py

@@ -14,12 +14,10 @@ from core.llm_generator.prompts import (
     JAVASCRIPT_CODE_GENERATOR_PROMPT_TEMPLATE,
     LLM_MODIFY_CODE_SYSTEM,
     LLM_MODIFY_PROMPT_SYSTEM,
-    MEMORY_UPDATE_PROMPT,
     PYTHON_CODE_GENERATOR_PROMPT_TEMPLATE,
     SYSTEM_STRUCTURED_OUTPUT_GENERATE,
     WORKFLOW_RULE_CONFIG_PROMPT_GENERATE_TEMPLATE,
 )
-from core.memory.entities import MemoryBlock, MemoryBlockSpec
 from core.model_manager import ModelManager
 from core.model_runtime.entities.llm_entities import LLMResult
 from core.model_runtime.entities.message_entities import PromptMessage, SystemPromptMessage, UserPromptMessage
@@ -30,7 +28,6 @@ from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
 from core.ops.utils import measure_time
 from core.prompt.utils.prompt_template_parser import PromptTemplateParser
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionMetadataKey
-from core.workflow.runtime.variable_pool import VariablePool
 from extensions.ext_database import db
 from extensions.ext_storage import storage
 from models import App, Message, WorkflowNodeExecutionModel
@@ -565,35 +562,3 @@ class LLMGenerator:
                 "Failed to invoke LLM model, model: %s", json.dumps(model_config.get("name")), exc_info=True
             )
             return {"error": f"An unexpected error occurred: {str(e)}"}
-
-    @staticmethod
-    def update_memory_block(
-        tenant_id: str,
-        visible_history: Sequence[tuple[str, str]],
-        variable_pool: VariablePool,
-        memory_block: MemoryBlock,
-        memory_spec: MemoryBlockSpec,
-    ) -> str:
-        model_instance = ModelManager().get_model_instance(
-            tenant_id=tenant_id,
-            provider=memory_spec.model.provider,
-            model=memory_spec.model.name,
-            model_type=ModelType.LLM,
-        )
-        formatted_history = ""
-        for sender, message in visible_history:
-            formatted_history += f"{sender}: {message}\n"
-        filled_instruction = variable_pool.convert_template(memory_spec.instruction).text
-        formatted_prompt = PromptTemplateParser(MEMORY_UPDATE_PROMPT).format(
-            inputs={
-                "formatted_history": formatted_history,
-                "current_value": memory_block.value,
-                "instruction": filled_instruction,
-            }
-        )
-        llm_result = model_instance.invoke_llm(
-            prompt_messages=[UserPromptMessage(content=formatted_prompt)],
-            model_parameters=memory_spec.model.completion_params,
-            stream=False,
-        )
-        return llm_result.message.get_text_content()

api/core/llm_generator/prompts.py

@@ -422,18 +422,3 @@ INSTRUCTION_GENERATE_TEMPLATE_PROMPT = """The output of this prompt is not as ex
 You should edit the prompt according to the IDEAL OUTPUT."""
 
 INSTRUCTION_GENERATE_TEMPLATE_CODE = """Please fix the errors in the {{#error_message#}}."""
-
-MEMORY_UPDATE_PROMPT = """
-Based on the following conversation history, update the memory content:
-
-Conversation history:
-{{formatted_history}}
-
-Current memory:
-{{current_value}}
-
-Update instruction:
-{{instruction}}
-
-Please output only the updated memory content, no other text like greeting:
-"""

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."""