fix: fix gemini panic

Qianxia666 · Jan 7, 2024 · bba6174 · bba6174
1 parent b91a269
commit bba6174
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 1 deletion.
diff --git a/middleware/recover.go b/middleware/recover.go
@@ -0,0 +1,28 @@
+package middleware
+
+import (
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"net/http"
+	"one-api/common"
+	"runtime/debug"
+)
+
+func RelayPanicRecover() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		defer func() {
+			if err := recover(); err != nil {
+				common.SysError(fmt.Sprintf("panic detected: %v", err))
+				common.SysError(fmt.Sprintf("stacktrace from panic: %s", string(debug.Stack())))
+				c.JSON(http.StatusInternalServerError, gin.H{
+					"error": gin.H{
+						"message": fmt.Sprintf("Panic detected, error: %v. Please submit a issue here: https://github.com/songquanpeng/one-api", err),
+						"type":    "one_api_panic",
+					},
+				})
+				c.Abort()
+			}
+		}()
+		c.Next()
+	}
+}
diff --git a/model/usedata.go b/model/usedata.go
@@ -18,6 +18,12 @@ type QuotaData struct {
 }
 
 func UpdateQuotaData() {
+	// recover
+	defer func() {
+		if r := recover(); r != nil {
+			common.SysLog(fmt.Sprintf("UpdateQuotaData panic: %s", r))
+		}
+	}()
 	for {
 		if common.DataExportEnabled {
 			common.SysLog("正在更新数据看板数据...")

diff --git a/router/relay-router.go b/router/relay-router.go
@@ -17,7 +17,7 @@ func SetRelayRouter(router *gin.Engine) {
 		modelsRouter.GET("/:model", controller.RetrieveModel)
 	}
 	relayV1Router := router.Group("/v1")
-	relayV1Router.Use(middleware.TokenAuth(), middleware.Distribute())
+	relayV1Router.Use(middleware.RelayPanicRecover(), middleware.TokenAuth(), middleware.Distribute())
 	{
 		relayV1Router.POST("/completions", controller.Relay)
 		relayV1Router.POST("/chat/completions", controller.Relay)