From 878c05657647c49bf67ed6f74f6278ca801cfeb0 Mon Sep 17 00:00:00 2001
From: psydok <47638600+psydok@users.noreply.github.com>
Date: Mon, 28 Oct 2024 05:27:38 +0500
Subject: [PATCH] Any exception is translated in http response (#2186)

* fix: get error in http request
---
 modules/api/api.py |  2 +-
 webui.py           | 18 +++++++++++++++++-
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/modules/api/api.py b/modules/api/api.py
index 14359736..d4851fd4 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -205,7 +205,7 @@ class Api:
         self.router = APIRouter()
         self.app = app
         self.queue_lock = queue_lock
-        #api_middleware(self.app)  # XXX this will have to be fixed
+        #api_middleware(self.app)  # FIXME: (legacy) this will have to be fixed
         self.add_api_route("/sdapi/v1/txt2img", self.text2imgapi, methods=["POST"], response_model=models.TextToImageResponse)
         self.add_api_route("/sdapi/v1/img2img", self.img2imgapi, methods=["POST"], response_model=models.ImageToImageResponse)
         self.add_api_route("/sdapi/v1/extra-single-image", self.extras_single_image_api, methods=["POST"], response_model=models.ExtrasSingleImageResponse)
diff --git a/webui.py b/webui.py
index ab0e6c2f..247612bc 100644
--- a/webui.py
+++ b/webui.py
@@ -3,6 +3,10 @@ from __future__ import annotations
 import os
 import time
 
+from fastapi import Request
+from fastapi.encoders import jsonable_encoder
+from fastapi.responses import JSONResponse
+
 from modules import timer
 from modules import initialize_util
 from modules import initialize
@@ -23,6 +27,17 @@ initialize.check_versions()
 initialize.initialize()
 
 
+def _handle_exception(request: Request, e: Exception):
+    error_information = vars(e)
+    content = {
+        "error": type(e).__name__,
+        "detail": error_information.get("detail", ""),
+        "body": error_information.get("body", ""),
+        "message": str(e),
+    }
+    return JSONResponse(status_code=int(error_information.get("status_code", 500)), content=jsonable_encoder(content))
+
+
 def create_api(app):
     from modules.api.api import Api
     from modules.call_queue import queue_lock
@@ -35,7 +50,7 @@ def api_only_worker():
     from fastapi import FastAPI
     from modules.shared_cmd_options import cmd_opts
 
-    app = FastAPI()
+    app = FastAPI(exception_handlers={Exception: _handle_exception})
     initialize_util.setup_middleware(app)
     api = create_api(app)
 
@@ -98,6 +113,7 @@ def webui_worker():
             app_kwargs={
                 "docs_url": "/docs",
                 "redoc_url": "/redoc",
+                "exception_handlers": {Exception: _handle_exception},
             },
             root_path=f"/{cmd_opts.subpath}" if cmd_opts.subpath else "",
         )