From fe06d266e9213ffb8aa7d0d709234aa85c44b1f6 Mon Sep 17 00:00:00 2001
From: -LAN- <laipz8200@outlook.com>
Date: Mon, 25 Aug 2025 09:28:42 +0800
Subject: [PATCH] refactor: better error handler (#24422)

Signed-off-by: -LAN- <laipz8200@outlook.com>
---
 api/extensions/ext_blueprints.py |   8 +-
 api/extensions/ext_login.py      |   4 +
 api/libs/external_api.py         | 176 ++++++++++++++++++-------------
 api/mypy.ini                     |   3 +
 4 files changed, 113 insertions(+), 78 deletions(-)

diff --git a/api/extensions/ext_blueprints.py b/api/extensions/ext_blueprints.py
index a4d013ffc..1024fd9ce 100644
--- a/api/extensions/ext_blueprints.py
+++ b/api/extensions/ext_blueprints.py
@@ -29,7 +29,6 @@ def init_app(app: DifyApp):
         methods=["GET", "PUT", "POST", "DELETE", "OPTIONS", "PATCH"],
         expose_headers=["X-Version", "X-Env"],
     )
-
     app.register_blueprint(web_bp)
 
     CORS(
@@ -40,10 +39,13 @@ def init_app(app: DifyApp):
         methods=["GET", "PUT", "POST", "DELETE", "OPTIONS", "PATCH"],
         expose_headers=["X-Version", "X-Env"],
     )
-
     app.register_blueprint(console_app_bp)
 
-    CORS(files_bp, allow_headers=["Content-Type"], methods=["GET", "PUT", "POST", "DELETE", "OPTIONS", "PATCH"])
+    CORS(
+        files_bp,
+        allow_headers=["Content-Type"],
+        methods=["GET", "PUT", "POST", "DELETE", "OPTIONS", "PATCH"],
+    )
     app.register_blueprint(files_bp)
 
     app.register_blueprint(inner_api_bp)
diff --git a/api/extensions/ext_login.py b/api/extensions/ext_login.py
index 9b18e25ea..9e5c71fb1 100644
--- a/api/extensions/ext_login.py
+++ b/api/extensions/ext_login.py
@@ -20,6 +20,10 @@ login_manager = flask_login.LoginManager()
 @login_manager.request_loader
 def load_user_from_request(request_from_flask_login):
     """Load user based on the request."""
+    # Skip authentication for documentation endpoints
+    if request.path.endswith("/docs") or request.path.endswith("/swagger.json"):
+        return None
+
     auth_header = request.headers.get("Authorization", "")
     auth_token: str | None = None
     if auth_header:
diff --git a/api/libs/external_api.py b/api/libs/external_api.py
index 715b559ab..ed07eb9a9 100644
--- a/api/libs/external_api.py
+++ b/api/libs/external_api.py
@@ -16,98 +16,124 @@ def http_status_message(code):
     return HTTP_STATUS_CODES.get(code, "")
 
 
-class ExternalApi(Api):
-    def handle_error(self, e):
-        """Error handler for the API transforms a raised exception into a Flask
-        response, with the appropriate HTTP status code and body.
+def register_external_error_handlers(api: Api) -> None:
+    """Register error handlers for the API using decorators.
 
-        :param e: the raised Exception object
-        :type e: Exception
+    :param api: The Flask-RestX Api instance
+    """
 
-        """
+    @api.errorhandler(HTTPException)
+    def handle_http_exception(e: HTTPException):
+        """Handle HTTP exceptions."""
         got_request_exception.send(current_app, exception=e)
 
+        if e.response is not None:
+            return e.get_response()
+
         headers = Headers()
-        if isinstance(e, HTTPException):
-            if e.response is not None:
-                resp = e.get_response()
-                return resp
+        status_code = e.code
+        default_data = {
+            "code": re.sub(r"(?<!^)(?=[A-Z])", "_", type(e).__name__).lower(),
+            "message": getattr(e, "description", http_status_message(status_code)),
+            "status": status_code,
+        }
 
-            status_code = e.code
-            default_data = {
-                "code": re.sub(r"(?<!^)(?=[A-Z])", "_", type(e).__name__).lower(),
-                "message": getattr(e, "description", http_status_message(status_code)),
-                "status": status_code,
-            }
+        if (
+            default_data["message"]
+            and default_data["message"] == "Failed to decode JSON object: Expecting value: line 1 column 1 (char 0)"
+        ):
+            default_data["message"] = "Invalid JSON payload received or JSON payload is empty."
 
-            if (
-                default_data["message"]
-                and default_data["message"] == "Failed to decode JSON object: Expecting value: line 1 column 1 (char 0)"
-            ):
-                default_data["message"] = "Invalid JSON payload received or JSON payload is empty."
+        headers = e.get_response().headers
 
-            headers = e.get_response().headers
-        elif isinstance(e, ValueError):
-            status_code = 400
-            default_data = {
-                "code": "invalid_param",
-                "message": str(e),
-                "status": status_code,
-            }
-        elif isinstance(e, AppInvokeQuotaExceededError):
-            status_code = 429
-            default_data = {
-                "code": "too_many_requests",
-                "message": str(e),
-                "status": status_code,
-            }
+        # Handle specific status codes
+        if status_code == 406 and api.default_mediatype is None:
+            supported_mediatypes = list(api.representations.keys())
+            fallback_mediatype = supported_mediatypes[0] if supported_mediatypes else "text/plain"
+            data = {"code": "not_acceptable", "message": default_data.get("message")}
+            resp = api.make_response(data, status_code, headers, fallback_mediatype=fallback_mediatype)
+        elif status_code == 400:
+            if isinstance(default_data.get("message"), dict):
+                param_key, param_value = list(default_data.get("message", {}).items())[0]
+                data = {"code": "invalid_param", "message": param_value, "params": param_key}
+            else:
+                data = default_data
+                if "code" not in data:
+                    data["code"] = "unknown"
+            resp = api.make_response(data, status_code, headers)
         else:
-            status_code = 500
-            default_data = {
-                "message": http_status_message(status_code),
-            }
+            data = default_data
+            if "code" not in data:
+                data["code"] = "unknown"
+            resp = api.make_response(data, status_code, headers)
 
-        # Werkzeug exceptions generate a content-length header which is added
-        # to the response in addition to the actual content-length header
-        # https://github.com/flask-restful/flask-restful/issues/534
+        if status_code == 401:
+            resp = api.unauthorized(resp)
+
+        # Remove duplicate Content-Length header
         remove_headers = ("Content-Length",)
-
         for header in remove_headers:
             headers.pop(header, None)
 
+        return resp
+
+    @api.errorhandler(ValueError)
+    def handle_value_error(e: ValueError):
+        """Handle ValueError exceptions."""
+        got_request_exception.send(current_app, exception=e)
+
+        status_code = 400
+        data = {
+            "code": "invalid_param",
+            "message": str(e),
+            "status": status_code,
+        }
+        return api.make_response(data, status_code)
+
+    @api.errorhandler(AppInvokeQuotaExceededError)
+    def handle_quota_exceeded(e: AppInvokeQuotaExceededError):
+        """Handle AppInvokeQuotaExceededError exceptions."""
+        got_request_exception.send(current_app, exception=e)
+
+        status_code = 429
+        data = {
+            "code": "too_many_requests",
+            "message": str(e),
+            "status": status_code,
+        }
+        return api.make_response(data, status_code)
+
+    @api.errorhandler(Exception)
+    def handle_general_exception(e: Exception):
+        """Handle general exceptions."""
+        got_request_exception.send(current_app, exception=e)
+
+        headers = Headers()
+        status_code = 500
+        default_data = {
+            "message": http_status_message(status_code),
+        }
+
         data = getattr(e, "data", default_data)
 
-        # record the exception in the logs when we have a server error of status code: 500
-        if status_code and status_code >= 500:
-            exc_info: Any = sys.exc_info()
-            if exc_info[1] is None:
-                exc_info = None
-            current_app.log_exception(exc_info)
+        # Log server errors
+        exc_info: Any = sys.exc_info()
+        if exc_info[1] is None:
+            exc_info = None
+        current_app.log_exception(exc_info)
 
-        if status_code == 406 and self.default_mediatype is None:
-            # if we are handling NotAcceptable (406), make sure that
-            # make_response uses a representation we support as the
-            # default mediatype (so that make_response doesn't throw
-            # another NotAcceptable error).
-            supported_mediatypes = list(self.representations.keys())  # only supported application/json
-            fallback_mediatype = supported_mediatypes[0] if supported_mediatypes else "text/plain"
-            data = {"code": "not_acceptable", "message": data.get("message")}
-            resp = self.make_response(data, status_code, headers, fallback_mediatype=fallback_mediatype)
-        elif status_code == 400:
-            if isinstance(data.get("message"), dict):
-                param_key, param_value = list(data.get("message", {}).items())[0]
-                data = {"code": "invalid_param", "message": param_value, "params": param_key}
-            else:
-                if "code" not in data:
-                    data["code"] = "unknown"
+        if "code" not in data:
+            data["code"] = "unknown"
 
-            resp = self.make_response(data, status_code, headers)
-        else:
-            if "code" not in data:
-                data["code"] = "unknown"
+        # Remove duplicate Content-Length header
+        remove_headers = ("Content-Length",)
+        for header in remove_headers:
+            headers.pop(header, None)
 
-            resp = self.make_response(data, status_code, headers)
+        return api.make_response(data, status_code, headers)
 
-        if status_code == 401:
-            resp = self.unauthorized(resp)
-        return resp
+
+class ExternalApi(Api):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        register_external_error_handlers(self)
diff --git a/api/mypy.ini b/api/mypy.ini
index 83d6d149f..44a01068e 100644
--- a/api/mypy.ini
+++ b/api/mypy.ini
@@ -15,5 +15,8 @@ ignore_missing_imports=True
 [mypy-flask_restx]
 ignore_missing_imports=True
 
+[mypy-flask_restx.api]
+ignore_missing_imports=True
+
 [mypy-flask_restx.inputs]
 ignore_missing_imports=True