feat: add websocket heartbeat message for mj page

api/handler/sd_handler.go

@@ -11,6 +11,8 @@ import (
 	"chatplus/utils/resp"
 	"encoding/base64"
 	"fmt"
+	"github.com/gorilla/websocket"
+	"net/http"
 	"time"
 
 	"github.com/gin-gonic/gin"
@@ -36,6 +38,27 @@ func NewSdJobHandler(app *core.AppServer, db *gorm.DB, pool *sd.ServicePool, man
 	return &h
 }
 
+// Client WebSocket 客户端，用于通知任务状态变更
+func (h *SdJobHandler) Client(c *gin.Context) {
+	ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		logger.Error(err)
+		c.Abort()
+		return
+	}
+
+	userId := h.GetInt(c, "user_id", 0)
+	if userId == 0 {
+		logger.Info("Invalid user ID")
+		c.Abort()
+		return
+	}
+
+	client := types.NewWsClient(ws)
+	h.pool.Clients.Put(uint(userId), client)
+	logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
+}
+
 func (h *SdJobHandler) checkLimits(c *gin.Context) bool {
 	user, err := utils.GetLoginUser(c, h.db)
 	if err != nil {