From f2315cf2d9b54c9c7182d522041adaa78db1b7f1 Mon Sep 17 00:00:00 2001
From: fit2cloud-chenyw <yawen.chen@fit2cloud.com>
Date: Thu, 8 May 2025 15:47:34 +0800
Subject: [PATCH] feat: Unified response model and exception handling

---
 backend/common/core/response_middleware.py | 82 ++++++++++++++++++++++
 backend/main.py                            |  3 +-
 2 files changed, 84 insertions(+), 1 deletion(-)
 create mode 100644 backend/common/core/response_middleware.py

diff --git a/backend/common/core/response_middleware.py b/backend/common/core/response_middleware.py
new file mode 100644
index 00000000..85c1cd85
--- /dev/null
+++ b/backend/common/core/response_middleware.py
@@ -0,0 +1,82 @@
+import json
+import logging
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import JSONResponse
+from starlette.exceptions import HTTPException
+from starlette.requests import Request
+
+class ResponseMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app):
+        super().__init__(app)
+
+    async def dispatch(self, request, call_next):
+        response = await call_next(request)
+        
+        if isinstance(response, JSONResponse):
+            return response
+    
+        if response.headers.get("content-type") == "application/json":
+            try:
+                body = b""
+                async for chunk in response.body_iterator:
+                    body += chunk
+                
+                raw_data = json.loads(body.decode())
+                
+                if isinstance(raw_data, dict) and all(k in raw_data for k in ["code", "data", "msg"]):
+                    return JSONResponse(
+                        content=raw_data,
+                        status_code=response.status_code,
+                        headers={
+                            k: v for k, v in response.headers.items()
+                            if k.lower() not in ("content-length", "content-type")
+                        }
+                    )
+                
+                wrapped_data = {
+                    "code": 0,
+                    "data": raw_data,
+                    "msg": None
+                }
+                
+                return JSONResponse(
+                    content=wrapped_data,
+                    status_code=response.status_code,
+                    headers={
+                        k: v for k, v in response.headers.items()
+                        if k.lower() not in ("content-length", "content-type")
+                    }
+                )
+            except Exception as e:
+                logging.error(f"Response processing error: {str(e)}", exc_info=True)
+                return response
+                
+        return response
+
+
+class exception_handler():
+    @staticmethod
+    async def http_exception_handler(request: Request, exc: HTTPException):
+        return JSONResponse(
+            status_code=exc.status_code,
+            content={
+                "code": exc.status_code,
+                "msg": exc.detail,
+                "data": None
+            }
+        )
+
+
+    @staticmethod
+    async def global_exception_handler(request: Request, exc: Exception):
+        logging.error(f"Error info: {str(exc)}", exc_info=True)
+        
+        return JSONResponse(
+            status_code=500,
+            content={
+                "code": 500,
+                "msg": str(exc),
+                "data": None
+            }
+        )
+
diff --git a/backend/main.py b/backend/main.py
index 6a69bd9f..f062e0f2 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -9,6 +9,7 @@ from starlette.middleware.cors import CORSMiddleware
 from apps.api import api_router
 from apps.system.middleware.auth import TokenMiddleware
 from common.core.config import settings
+from common.core.response_middleware import ResponseMiddleware
 
 def custom_generate_unique_id(route: APIRoute) -> str:
     tag = route.tags[0] if route.tags and len(route.tags) > 0 else ""
@@ -35,10 +36,10 @@ if settings.all_cors_origins:
     )
 
 app.add_middleware(TokenMiddleware)
+app.add_middleware(ResponseMiddleware)
 app.include_router(api_router, prefix=settings.API_V1_STR)
 
 
-
 frontend_dist = os.path.abspath("../frontend/dist")
 if not os.path.exists(frontend_dist):
     logging.warning(f"The front-end build directory does not exist: {frontend_dist}")