first commit

internal/server/ws.go

@@ -0,0 +1,177 @@
+package server
+
+import (
+	"context"
+	"encoding/json"
+	"log"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+
+	"agent/internal/agent"
+
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+)
+
+var upgrader = websocket.Upgrader{
+	ReadBufferSize:  1024,
+	WriteBufferSize: 4096,
+	// 本地单用户 Web，允许任意 Origin
+	CheckOrigin: func(r *http.Request) bool { return true },
+}
+
+// wsInbound 客户端 → 服务端消息。
+// chat：在绑定会话上运行 Agent；switch：切换绑定会话；stop：停止当前生成；ping：保活。
+type wsInbound struct {
+	Type      string `json:"type"`
+	Content   string `json:"content"`
+	SessionID string `json:"session_id"`
+}
+
+// wsConn 封装一个 WebSocket 连接：绑定一个会话、串行处理消息。
+type wsConn struct {
+	server *Server
+	conn   *websocket.Conn
+	sendMu sync.Mutex // 串行化 WriteJSON
+
+	sessionID string
+	unreg     func() // 当前会话的 live 注销函数
+
+	runMu     sync.Mutex
+	runCancel context.CancelFunc
+}
+
+func (s *Server) ws(c *gin.Context) {
+	sessionID := strings.TrimSpace(c.Query("session"))
+	if sessionID == "" || !sessionIDPattern.MatchString(sessionID) {
+		log.Printf("ws: 连接被拒（session 参数无效）session=%q", sessionID)
+		fail(c, http.StatusBadRequest, "缺少有效的 session 参数")
+		return
+	}
+	if _, ok := s.store.Get(sessionID); !ok {
+		log.Printf("ws: 连接被拒（会话不存在）session=%q", sessionID)
+		fail(c, http.StatusNotFound, "会话不存在")
+		return
+	}
+
+	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		return
+	}
+	client := &wsConn{server: s, conn: conn}
+	client.bind(sessionID)
+	client.readPump()
+}
+
+// sendJSON 安全地推送一个事件帧 {type, data}。
+func (cl *wsConn) sendJSON(eventType string, data any) {
+	cl.sendMu.Lock()
+	defer cl.sendMu.Unlock()
+	_ = cl.conn.WriteJSON(map[string]any{"type": eventType, "data": data})
+}
+
+// bind 切换绑定会话：注销旧会话的 live 注册，注册新会话。
+// 传入空串表示仅注销。
+func (cl *wsConn) bind(sessionID string) {
+	if cl.unreg != nil {
+		cl.unreg()
+		cl.unreg = nil
+	}
+	cl.sessionID = sessionID
+	if sessionID == "" {
+		return
+	}
+	cl.unreg = cl.server.agent.RegisterLive(sessionID, func(event agent.Event) {
+		cl.sendJSON(event.Type, event.Data)
+	})
+}
+
+func (cl *wsConn) handleChat(content string) {
+	content = strings.TrimSpace(content)
+	if content == "" {
+		cl.sendJSON("error", map[string]any{"message": "消息内容不能为空"})
+		return
+	}
+	// 未配置 LLM API Key 时直接提示去配置，避免进入无效的模型调用。
+	// 环境变量 LLM_API_KEY（托管模式）与设置页 api-config.json 都计入已配置。
+	if !cl.server.apiCfg.IsAPIKeyConfigured() {
+		cl.sendJSON("error", map[string]any{"message": "尚未配置 LLM API，请先点击左下角设置按钮配置 API Key 与模型"})
+		return
+	}
+	cl.runMu.Lock()
+	if cl.runCancel != nil {
+		cl.runMu.Unlock()
+		cl.sendJSON("error", map[string]any{"message": "Agent 正在执行，请等待完成或先停止"})
+		return
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cl.runCancel = cancel
+	cl.runMu.Unlock()
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		_ = cl.server.agent.Run(ctx, cl.sessionID, content, func(event agent.Event) {
+			cl.sendJSON(event.Type, event.Data)
+		})
+	}()
+	go func() {
+		<-done
+		cl.runMu.Lock()
+		cl.runCancel = nil
+		cl.runMu.Unlock()
+	}()
+}
+
+func (cl *wsConn) handleStop() {
+	cl.runMu.Lock()
+	if cl.runCancel != nil {
+		cl.runCancel()
+	}
+	cl.runMu.Unlock()
+}
+
+func (cl *wsConn) readPump() {
+	defer func() {
+		cl.bind("") // 注销 live
+		_ = cl.conn.Close()
+	}()
+	cl.conn.SetReadLimit(1 << 20) // 1MB
+	_ = cl.conn.SetReadDeadline(time.Now().Add(90 * time.Second))
+	cl.conn.SetPongHandler(func(string) error {
+		_ = cl.conn.SetReadDeadline(time.Now().Add(90 * time.Second))
+		return nil
+	})
+
+	for {
+		// 任何收到消息（含客户端心跳 {type:"ping"}）都刷新 90s 读超时，
+		// 否则浏览器只能发协议层 Pong，应用层 ping 不会触发 SetPongHandler，连接必死在第 90 秒
+		_ = cl.conn.SetReadDeadline(time.Now().Add(90 * time.Second))
+		_, raw, err := cl.conn.ReadMessage()
+		if err != nil {
+			return
+		}
+		var msg wsInbound
+		if err := json.Unmarshal(raw, &msg); err != nil {
+			continue
+		}
+		switch msg.Type {
+		case "chat":
+			cl.handleChat(msg.Content)
+		case "stop":
+			cl.handleStop()
+		case "switch":
+			if msg.SessionID != "" && sessionIDPattern.MatchString(msg.SessionID) {
+				if _, ok := cl.server.store.Get(msg.SessionID); ok {
+					cl.bind(msg.SessionID)
+				}
+			}
+		case "ping":
+			cl.sendJSON("pong", map[string]any{"time": time.Now().UnixMilli()})
+		default:
+			log.Printf("ws: unknown message type %q", msg.Type)
+		}
+	}
+}