From cd49a2801ac57516c8cb64e3479aa3c467161214 Mon Sep 17 00:00:00 2001
From: xiaowei <xiaowei@hermes.local>
Date: Sat, 30 May 2026 17:15:19 +0800
Subject: [PATCH] fix(ws): add per-conn mutex to prevent concurrent WriteJSON
 panic
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Root cause: Broadcast() and Push() launched multiple goroutines
writing to the same websocket.Conn simultaneously. Gorilla websocket
does not support concurrent WriteMessage/WriteJSON.

Fix: add connMu sync.Map (*websocket.Conn → *sync.Mutex) to WSManager.
Each WriteJSON now goes through writeWithLock() which acquires the
per-conn mutex. Unregister() also acquires the lock before Close()
to prevent write/close races.

No functional changes to API surface.
---
 go/internal/api/routes/ws.go | 30 +++++++++++++++++++++++++-----
 1 file changed, 25 insertions(+), 5 deletions(-)

diff --git a/go/internal/api/routes/ws.go b/go/internal/api/routes/ws.go
index 46d7361..0ffbc22 100644
--- a/go/internal/api/routes/ws.go
+++ b/go/internal/api/routes/ws.go
@@ -15,8 +15,9 @@ var upgrader = websocket.Upgrader{
 
 // WSManager 管理所有 WebSocket 连接，支持按 AgentID 定向推送
 type WSManager struct {
-	mu      sync.RWMutex
-	clients map[string]map[*websocket.Conn]bool // agentID → 连接集合
+	mu       sync.RWMutex
+	clients  map[string]map[*websocket.Conn]bool        // agentID → 连接集合
+	connMu   sync.Map // per-conn mutex: *websocket.Conn → *sync.Mutex
 }
 
 // 全局 WebSocket 管理器
@@ -32,6 +33,7 @@ func (wm *WSManager) Register(agentID string, conn *websocket.Conn) {
 		wm.clients[agentID] = make(map[*websocket.Conn]bool)
 	}
 	wm.clients[agentID][conn] = true
+	wm.connMu.Store(conn, &sync.Mutex{})
 	log.Printf("[ws] agent %s connected (%d total)", agentID, len(wm.clients[agentID]))
 }
 
@@ -45,10 +47,28 @@ func (wm *WSManager) Unregister(agentID string, conn *websocket.Conn) {
 			delete(wm.clients, agentID)
 		}
 	}
-	conn.Close()
+	// 先获取该 conn 的锁，确保没有正在进行的 write，然后关闭
+	if muVal, ok := wm.connMu.Load(conn); ok {
+		mu := muVal.(*sync.Mutex)
+		mu.Lock()
+		conn.Close()
+		mu.Unlock()
+	} else {
+		conn.Close()
+	}
+	wm.connMu.Delete(conn)
 	log.Printf("[ws] agent %s disconnected", agentID)
 }
 
+// writeWithLock 安全地写入 WebSocket JSON（每个连接保证单 writer）
+func (wm *WSManager) writeWithLock(conn *websocket.Conn, msg interface{}) error {
+	muVal, _ := wm.connMu.Load(conn)
+	mu := muVal.(*sync.Mutex)
+	mu.Lock()
+	defer mu.Unlock()
+	return conn.WriteJSON(msg)
+}
+
 // Push 推送到指定 Agent 所有连接
 func (wm *WSManager) Push(agentID string, msgType string, payload interface{}) {
 	wm.mu.RLock()
@@ -66,7 +86,7 @@ func (wm *WSManager) Push(agentID string, msgType string, payload interface{}) {
 
 	for conn := range clients {
 		go func(c *websocket.Conn) {
-			if err := c.WriteJSON(msg); err != nil {
+			if err := wm.writeWithLock(c, msg); err != nil {
 				log.Printf("[ws] write error: %v", err)
 				wm.Unregister(agentID, c)
 			}
@@ -87,7 +107,7 @@ func (wm *WSManager) Broadcast(msgType string, payload interface{}) {
 	for agentID, clients := range wm.clients {
 		for conn := range clients {
 			go func(c *websocket.Conn, aid string) {
-				if err := c.WriteJSON(msg); err != nil {
+				if err := wm.writeWithLock(c, msg); err != nil {
 					log.Printf("[ws] broadcast error: %v", err)
 					wm.Unregister(aid, c)
 				}