Merge remote-tracking branch 'origin/main' into feat/support-agent-sandbox

api/controllers/console/workspace/trigger_providers.py

@@ -1,14 +1,14 @@
 import logging
-from collections.abc import Mapping
 from typing import Any
 
 from flask import make_response, redirect, request
-from flask_restx import Resource, reqparse
-from pydantic import BaseModel, Field, model_validator
+from flask_restx import Resource
+from pydantic import BaseModel, model_validator
 from sqlalchemy.orm import Session
 from werkzeug.exceptions import BadRequest, Forbidden
 
 from configs import dify_config
+from controllers.common.schema import register_schema_models
 from controllers.web.error import NotFoundError
 from core.model_runtime.utils.encoders import jsonable_encoder
 from core.plugin.entities.plugin_daemon import CredentialType
@@ -35,35 +35,38 @@ from ..wraps import (
 logger = logging.getLogger(__name__)
 
 
-class TriggerSubscriptionUpdateRequest(BaseModel):
-    """Request payload for updating a trigger subscription"""
+class TriggerSubscriptionBuilderCreatePayload(BaseModel):
+    credential_type: str = CredentialType.UNAUTHORIZED
 
-    name: str | None = Field(default=None, description="The name for the subscription")
-    credentials: Mapping[str, Any] | None = Field(default=None, description="The credentials for the subscription")
-    parameters: Mapping[str, Any] | None = Field(default=None, description="The parameters for the subscription")
-    properties: Mapping[str, Any] | None = Field(default=None, description="The properties for the subscription")
+
+class TriggerSubscriptionBuilderVerifyPayload(BaseModel):
+    credentials: dict[str, Any]
+
+
+class TriggerSubscriptionBuilderUpdatePayload(BaseModel):
+    name: str | None = None
+    parameters: dict[str, Any] | None = None
+    properties: dict[str, Any] | None = None
+    credentials: dict[str, Any] | None = None
 
     @model_validator(mode="after")
     def check_at_least_one_field(self):
-        if all(v is None for v in (self.name, self.credentials, self.parameters, self.properties)):
+        if all(v is None for v in self.model_dump().values()):
             raise ValueError("At least one of name, credentials, parameters, or properties must be provided")
         return self
 
 
-class TriggerSubscriptionVerifyRequest(BaseModel):
-    """Request payload for verifying subscription credentials."""
-
-    credentials: Mapping[str, Any] = Field(description="The credentials to verify")
+class TriggerOAuthClientPayload(BaseModel):
+    client_params: dict[str, Any] | None = None
+    enabled: bool | None = None
 
 
-console_ns.schema_model(
-    TriggerSubscriptionUpdateRequest.__name__,
-    TriggerSubscriptionUpdateRequest.model_json_schema(ref_template="#/definitions/{model}"),
-)
-
-console_ns.schema_model(
-    TriggerSubscriptionVerifyRequest.__name__,
-    TriggerSubscriptionVerifyRequest.model_json_schema(ref_template="#/definitions/{model}"),
+register_schema_models(
+    console_ns,
+    TriggerSubscriptionBuilderCreatePayload,
+    TriggerSubscriptionBuilderVerifyPayload,
+    TriggerSubscriptionBuilderUpdatePayload,
+    TriggerOAuthClientPayload,
 )
 
 
@@ -132,16 +135,11 @@ class TriggerSubscriptionListApi(Resource):
             raise
 
 
-parser = reqparse.RequestParser().add_argument(
-    "credential_type", type=str, required=False, nullable=True, location="json"
-)
-
-
 @console_ns.route(
     "/workspaces/current/trigger-provider/<path:provider>/subscriptions/builder/create",
 )
 class TriggerSubscriptionBuilderCreateApi(Resource):
-    @console_ns.expect(parser)
+    @console_ns.expect(console_ns.models[TriggerSubscriptionBuilderCreatePayload.__name__])
     @setup_required
     @login_required
     @edit_permission_required
@@ -151,10 +149,10 @@ class TriggerSubscriptionBuilderCreateApi(Resource):
         user = current_user
         assert user.current_tenant_id is not None
 
-        args = parser.parse_args()
+        payload = TriggerSubscriptionBuilderCreatePayload.model_validate(console_ns.payload or {})
 
         try:
-            credential_type = CredentialType.of(args.get("credential_type") or CredentialType.UNAUTHORIZED.value)
+            credential_type = CredentialType.of(payload.credential_type)
             subscription_builder = TriggerSubscriptionBuilderService.create_trigger_subscription_builder(
                 tenant_id=user.current_tenant_id,
                 user_id=user.id,
@@ -182,18 +180,11 @@ class TriggerSubscriptionBuilderGetApi(Resource):
         )
 
 
-parser_api = (
-    reqparse.RequestParser()
-    # The credentials of the subscription builder
-    .add_argument("credentials", type=dict, required=False, nullable=True, location="json")
-)
-
-
 @console_ns.route(
     "/workspaces/current/trigger-provider/<path:provider>/subscriptions/builder/verify-and-update/<path:subscription_builder_id>",
 )
-class TriggerSubscriptionBuilderVerifyAndUpdateApi(Resource):
-    @console_ns.expect(parser_api)
+class TriggerSubscriptionBuilderVerifyApi(Resource):
+    @console_ns.expect(console_ns.models[TriggerSubscriptionBuilderVerifyPayload.__name__])
     @setup_required
     @login_required
     @edit_permission_required
@@ -203,7 +194,7 @@ class TriggerSubscriptionBuilderVerifyAndUpdateApi(Resource):
         user = current_user
         assert user.current_tenant_id is not None
 
-        args = parser_api.parse_args()
+        payload = TriggerSubscriptionBuilderVerifyPayload.model_validate(console_ns.payload or {})
 
         try:
             # Use atomic update_and_verify to prevent race conditions
@@ -213,7 +204,7 @@ class TriggerSubscriptionBuilderVerifyAndUpdateApi(Resource):
                 provider_id=TriggerProviderID(provider),
                 subscription_builder_id=subscription_builder_id,
                 subscription_builder_updater=SubscriptionBuilderUpdater(
-                    credentials=args.get("credentials", None),
+                    credentials=payload.credentials,
                 ),
             )
         except Exception as e:
@@ -221,24 +212,11 @@ class TriggerSubscriptionBuilderVerifyAndUpdateApi(Resource):
             raise ValueError(str(e)) from e
 
 
-parser_update_api = (
-    reqparse.RequestParser()
-    # The name of the subscription builder
-    .add_argument("name", type=str, required=False, nullable=True, location="json")
-    # The parameters of the subscription builder
-    .add_argument("parameters", type=dict, required=False, nullable=True, location="json")
-    # The properties of the subscription builder
-    .add_argument("properties", type=dict, required=False, nullable=True, location="json")
-    # The credentials of the subscription builder
-    .add_argument("credentials", type=dict, required=False, nullable=True, location="json")
-)
-
-
 @console_ns.route(
     "/workspaces/current/trigger-provider/<path:provider>/subscriptions/builder/update/<path:subscription_builder_id>",
 )
 class TriggerSubscriptionBuilderUpdateApi(Resource):
-    @console_ns.expect(parser_update_api)
+    @console_ns.expect(console_ns.models[TriggerSubscriptionBuilderUpdatePayload.__name__])
     @setup_required
     @login_required
     @edit_permission_required
@@ -249,7 +227,7 @@ class TriggerSubscriptionBuilderUpdateApi(Resource):
         assert isinstance(user, Account)
         assert user.current_tenant_id is not None
 
-        args = parser_update_api.parse_args()
+        payload = TriggerSubscriptionBuilderUpdatePayload.model_validate(console_ns.payload or {})
         try:
             return jsonable_encoder(
                 TriggerSubscriptionBuilderService.update_trigger_subscription_builder(
@@ -257,10 +235,10 @@ class TriggerSubscriptionBuilderUpdateApi(Resource):
                     provider_id=TriggerProviderID(provider),
                     subscription_builder_id=subscription_builder_id,
                     subscription_builder_updater=SubscriptionBuilderUpdater(
-                        name=args.get("name", None),
-                        parameters=args.get("parameters", None),
-                        properties=args.get("properties", None),
-                        credentials=args.get("credentials", None),
+                        name=payload.name,
+                        parameters=payload.parameters,
+                        properties=payload.properties,
+                        credentials=payload.credentials,
                     ),
                 )
             )
@@ -295,7 +273,7 @@ class TriggerSubscriptionBuilderLogsApi(Resource):
     "/workspaces/current/trigger-provider/<path:provider>/subscriptions/builder/build/<path:subscription_builder_id>",
 )
 class TriggerSubscriptionBuilderBuildApi(Resource):
-    @console_ns.expect(parser_update_api)
+    @console_ns.expect(console_ns.models[TriggerSubscriptionBuilderUpdatePayload.__name__])
     @setup_required
     @login_required
     @edit_permission_required
@@ -304,7 +282,7 @@ class TriggerSubscriptionBuilderBuildApi(Resource):
         """Build a subscription instance for a trigger provider"""
         user = current_user
         assert user.current_tenant_id is not None
-        args = parser_update_api.parse_args()
+        payload = TriggerSubscriptionBuilderUpdatePayload.model_validate(console_ns.payload or {})
         try:
             # Use atomic update_and_build to prevent race conditions
             TriggerSubscriptionBuilderService.update_and_build_builder(
@@ -313,9 +291,9 @@ class TriggerSubscriptionBuilderBuildApi(Resource):
                 provider_id=TriggerProviderID(provider),
                 subscription_builder_id=subscription_builder_id,
                 subscription_builder_updater=SubscriptionBuilderUpdater(
-                    name=args.get("name", None),
-                    parameters=args.get("parameters", None),
-                    properties=args.get("properties", None),
+                    name=payload.name,
+                    parameters=payload.parameters,
+                    properties=payload.properties,
                 ),
             )
             return 200
@@ -328,7 +306,7 @@ class TriggerSubscriptionBuilderBuildApi(Resource):
     "/workspaces/current/trigger-provider/<path:subscription_id>/subscriptions/update",
 )
 class TriggerSubscriptionUpdateApi(Resource):
-    @console_ns.expect(console_ns.models[TriggerSubscriptionUpdateRequest.__name__])
+    @console_ns.expect(console_ns.models[TriggerSubscriptionBuilderUpdatePayload.__name__])
     @setup_required
     @login_required
     @edit_permission_required
@@ -338,7 +316,7 @@ class TriggerSubscriptionUpdateApi(Resource):
         user = current_user
         assert user.current_tenant_id is not None
 
-        request = TriggerSubscriptionUpdateRequest.model_validate(console_ns.payload)
+        request = TriggerSubscriptionBuilderUpdatePayload.model_validate(console_ns.payload or {})
 
         subscription = TriggerProviderService.get_subscription_by_id(
             tenant_id=user.current_tenant_id,
@@ -568,13 +546,6 @@ class TriggerOAuthCallbackApi(Resource):
         return redirect(f"{dify_config.CONSOLE_WEB_URL}/oauth-callback")
 
 
-parser_oauth_client = (
-    reqparse.RequestParser()
-    .add_argument("client_params", type=dict, required=False, nullable=True, location="json")
-    .add_argument("enabled", type=bool, required=False, nullable=True, location="json")
-)
-
-
 @console_ns.route("/workspaces/current/trigger-provider/<path:provider>/oauth/client")
 class TriggerOAuthClientManageApi(Resource):
     @setup_required
@@ -622,7 +593,7 @@ class TriggerOAuthClientManageApi(Resource):
             logger.exception("Error getting OAuth client", exc_info=e)
             raise
 
-    @console_ns.expect(parser_oauth_client)
+    @console_ns.expect(console_ns.models[TriggerOAuthClientPayload.__name__])
     @setup_required
     @login_required
     @is_admin_or_owner_required
@@ -632,15 +603,15 @@ class TriggerOAuthClientManageApi(Resource):
         user = current_user
         assert user.current_tenant_id is not None
 
-        args = parser_oauth_client.parse_args()
+        payload = TriggerOAuthClientPayload.model_validate(console_ns.payload or {})
 
         try:
             provider_id = TriggerProviderID(provider)
             return TriggerProviderService.save_custom_oauth_client_params(
                 tenant_id=user.current_tenant_id,
                 provider_id=provider_id,
-                client_params=args.get("client_params"),
-                enabled=args.get("enabled"),
+                client_params=payload.client_params,
+                enabled=payload.enabled,
             )
 
         except ValueError as e:
@@ -676,7 +647,7 @@ class TriggerOAuthClientManageApi(Resource):
     "/workspaces/current/trigger-provider/<path:provider>/subscriptions/verify/<path:subscription_id>",
 )
 class TriggerSubscriptionVerifyApi(Resource):
-    @console_ns.expect(console_ns.models[TriggerSubscriptionVerifyRequest.__name__])
+    @console_ns.expect(console_ns.models[TriggerSubscriptionBuilderVerifyPayload.__name__])
     @setup_required
     @login_required
     @edit_permission_required
@@ -686,9 +657,7 @@ class TriggerSubscriptionVerifyApi(Resource):
         user = current_user
         assert user.current_tenant_id is not None
 
-        verify_request: TriggerSubscriptionVerifyRequest = TriggerSubscriptionVerifyRequest.model_validate(
-            console_ns.payload
-        )
+        verify_request = TriggerSubscriptionBuilderVerifyPayload.model_validate(console_ns.payload or {})
 
         try:
             result = TriggerProviderService.verify_subscription_credentials(

api/controllers/console/workspace/workspace.py

@@ -80,6 +80,9 @@ 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 = {

api/controllers/web/login.py

@@ -10,7 +10,12 @@ from controllers.console.auth.error import (
     InvalidEmailError,
 )
 from controllers.console.error import AccountBannedError
-from controllers.console.wraps import only_edition_enterprise, setup_required
+from controllers.console.wraps import (
+    decrypt_code_field,
+    decrypt_password_field,
+    only_edition_enterprise,
+    setup_required,
+)
 from controllers.web import web_ns
 from controllers.web.wraps import decode_jwt_token
 from libs.helper import email
@@ -42,6 +47,7 @@ class LoginApi(Resource):
             404: "Account not found",
         }
     )
+    @decrypt_password_field
     def post(self):
         """Authenticate user and login."""
         parser = (
@@ -181,6 +187,7 @@ class EmailCodeLoginApi(Resource):
             404: "Account not found",
         }
     )
+    @decrypt_code_field
     def post(self):
         parser = (
             reqparse.RequestParser()