package services import ( "ai_scheduler/internal/biz" "ai_scheduler/internal/config" "ai_scheduler/internal/data/constants" "ai_scheduler/internal/entitys" "ai_scheduler/internal/gateway" "context" "encoding/json" "log" "sync" "time" "github.com/gofiber/fiber/v2" "github.com/gofiber/websocket/v2" ) // ChatHandler 聊天处理器 type ChatService struct { routerBiz *biz.AiRouterBiz Gw *gateway.Gateway mu sync.Mutex ChatHis *biz.ChatHistoryBiz cfg *config.Config } // NewChatHandler 创建聊天处理器 func NewChatService( routerService *biz.AiRouterBiz, chatHis *biz.ChatHistoryBiz, gw *gateway.Gateway, cfg *config.Config, ) *ChatService { return &ChatService{ routerBiz: routerService, Gw: gw, ChatHis: chatHis, cfg: cfg, } } func (h *ChatService) ChatFail(c *websocket.Conn, content string) { err := c.WriteMessage(websocket.TextMessage, []byte(content)) if err != nil { log.Println("发送错误:", err) } _ = c.Close() } // Chat 处理WebSocket聊天连接 // 这是WebSocket处理的主入口函数 func (h *ChatService) Chat(c *websocket.Conn) { ctx, cancel := context.WithCancel(context.Background()) // 创建新的客户端实例 client := gateway.NewClient(c, ctx, cancel) // 心跳检测 go client.InitHeartbeat(time.Duration(h.cfg.Sys.HeartbeatInterval)) // 将客户端添加到网关管理 h.Gw.AddClient(client) // 确保在函数返回时移除客户端并关闭连接 defer func() { h.Gw.Cleanup(client.GetID()) }() // 绑定会话ID uid := c.Query("x-session") if uid != "" { if err := h.Gw.BindUid(client.GetID(), uid); err != nil { log.Println("绑定UID错误:", err) } } // 验证并收集连接数据,后续对话中会使用 if err := client.DataAuth(); err != nil { log.Println("数据验证错误:", err) h.ChatFail(c, err.Error()) return } // 循环读取客户端消息 for { // 读取消息 messageType, message, err := c.ReadMessage() if err != nil { log.Println("读取错误:", err) break } //if string(message) == `{"type":"ping"}` { // client.LastActive = time.Now() // if err := c.WriteMessage(websocket.TextMessage, []byte(`{"type":"pong"}`)); err != nil { // log.Println("Heartbeat response failed", "id", client.GetID(), "err", err) // return // } // continue //} // 处理消息 msg, chatType := h.handleMessageToString(c, messageType, message) if chatType == constants.ConnStatusClosed { break } if chatType == constants.ConnStatusIgnore { continue } log.Printf("收到消息: %s", string(msg)) // 解析请求 var req entitys.ChatSockRequest if err = json.Unmarshal(msg, &req); err != nil { log.Println("JSON parse error:", err) continue } // 路由处理请求 err = h.routerBiz.RouteWithSocket(client, &req) if err != nil { log.Println("处理失败:", err) } } } // handleMessageToString 处理不同类型的WebSocket消息 // 参数: // - c: WebSocket连接 // - msgType: 消息类型 // - msg: 消息内容 // // 返回: // - text: 处理后的文本内容 // - chatType: 连接状态 func (h *ChatService) handleMessageToString(c *websocket.Conn, msgType int, msg any) (text []byte, chatType constants.ConnStatus) { switch msgType { case websocket.TextMessage: return msg.([]byte), constants.ConnStatusNormal case websocket.BinaryMessage: return msg.([]byte), constants.ConnStatusNormal case websocket.CloseMessage: return nil, constants.ConnStatusClosed case websocket.PingMessage: // 可选:回复 Pong c.WriteMessage(websocket.PongMessage, nil) return nil, constants.ConnStatusIgnore case websocket.PongMessage: return nil, constants.ConnStatusIgnore default: return nil, constants.ConnStatusIgnore } return msg.([]byte), constants.ConnStatusIgnore } func (s *ChatService) Useful(c *fiber.Ctx) error { req := &entitys.UseFulRequest{} if err := c.BodyParser(req); err != nil { return err } err := s.ChatHis.Update(c.Context(), req) if err != nil { return err } return nil } func (s *ChatService) UsefulList(c *fiber.Ctx) error { return c.JSON(constants.UseFulMap) }