From 24159fda62420683be283cc4724c7253eab72e7b Mon Sep 17 00:00:00 2001
From: engigu <sayheya@qq.com>
Date: Thu, 1 Jan 2026 12:36:50 +0800
Subject: [PATCH] fix: agent ws connect panic

---
 agent/agent.go                           | 18 ++++++++++++++++++
 internal/controllers/agent_controller.go |  3 +--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/agent/agent.go b/agent/agent.go
index 4016cf9..1442aa7 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -77,6 +77,7 @@ type Agent struct {
 	wsConn        *websocket.Conn
 	wsMu          sync.Mutex
 	stopCh        chan struct{}
+	wsStopCh      chan struct{} // 用于停止当前 WebSocket 相关的 goroutine
 }
 
 // generateMachineID 生成机器识别码
@@ -190,6 +191,7 @@ func (a *Agent) connectWS() error {
 
 	a.wsMu.Lock()
 	a.wsConn = conn
+	a.wsStopCh = make(chan struct{})
 	a.wsMu.Unlock()
 
 	log.Info("WebSocket 已连接")
@@ -202,6 +204,10 @@ func (a *Agent) connectWS() error {
 func (a *Agent) closeWS() {
 	a.wsMu.Lock()
 	defer a.wsMu.Unlock()
+	if a.wsStopCh != nil {
+		close(a.wsStopCh)
+		a.wsStopCh = nil
+	}
 	if a.wsConn != nil {
 		a.wsConn.Close()
 		a.wsConn = nil
@@ -209,6 +215,8 @@ func (a *Agent) closeWS() {
 }
 
 func (a *Agent) readWS() {
+	defer a.closeWS() // 读取结束时关闭连接，停止 heartbeatLoop
+
 	for {
 		a.wsMu.Lock()
 		conn := a.wsConn
@@ -327,10 +335,20 @@ func (a *Agent) heartbeatLoop() {
 	ticker := time.NewTicker(time.Duration(a.config.Interval) * time.Second)
 	defer ticker.Stop()
 
+	a.wsMu.Lock()
+	wsStopCh := a.wsStopCh
+	a.wsMu.Unlock()
+
+	if wsStopCh == nil {
+		return
+	}
+
 	for {
 		select {
 		case <-a.stopCh:
 			return
+		case <-wsStopCh:
+			return
 		case <-ticker.C:
 			a.wsMu.Lock()
 			conn := a.wsConn
diff --git a/internal/controllers/agent_controller.go b/internal/controllers/agent_controller.go
index 662aed4..9eba6cf 100644
--- a/internal/controllers/agent_controller.go
+++ b/internal/controllers/agent_controller.go
@@ -413,9 +413,8 @@ func (c *AgentController) wsReadPump(ac *services.AgentConnection, agent *models
 		c.wsManager.Unregister(agent.ID)
 	}()
 
-	// 检查连接是否有效
+	// 检查连接是否有效（可能是旧连接被新连接替换）
 	if ac == nil || ac.IsClosed() {
-		logger.Warnf("[AgentWS] Agent #%d 连接无效", agent.ID)
 		return
 	}