From f4a53d6b5f3a3adeabba898ee64fa22828a0c068 Mon Sep 17 00:00:00 2001
From: ryan <ryan@arctel.net>
Date: Sat, 30 May 2026 18:08:06 +0800
Subject: [PATCH] =?UTF-8?q?[=E4=BC=98=E5=8C=96]=20=E5=A2=9E=E5=8A=A0?=
 =?UTF-8?q?=E6=96=AD=E5=BC=80=20WebSocket=20=E5=AE=A2=E6=88=B7=E7=AB=AF?=
 =?UTF-8?q?=E7=9A=84=E5=8A=9F=E8=83=BD=EF=BC=8C=E4=BC=98=E5=8C=96=E8=BF=9E?=
 =?UTF-8?q?=E6=8E=A5=E7=AE=A1=E7=90=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 docker-compose.yaml                  |  2 +-
 openflare_server/controller/agent.go |  6 ++++++
 openflare_server/service/agent_ws.go | 15 +++++++++++++++
 openflare_server/service/node.go     |  1 +
 4 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/docker-compose.yaml b/docker-compose.yaml
index 93a249f5..7640f30e 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -16,7 +16,7 @@ services:
 
     environment:
       OPENFLARE_SERVER_URL: "http://host.docker.internal:3000"
-      OPENFLARE_AGENT_TOKEN: "373956188ddead1df6dd7c86cd330b73"
+      OPENFLARE_AGENT_TOKEN: "07800f31d3f181e65d18dca1407d821c"
       LOG_LEVEL: "debug"
 
     extra_hosts:
diff --git a/openflare_server/controller/agent.go b/openflare_server/controller/agent.go
index 508c4d6a..56c066a0 100644
--- a/openflare_server/controller/agent.go
+++ b/openflare_server/controller/agent.go
@@ -150,6 +150,12 @@ func AgentWebSocket(c *gin.Context) {
 		}()
 
 		slog.Debug("agent ws upgrade succeeded", "node_id", node.NodeID, "remote", c.Request.RemoteAddr)
+
+		go func() {
+			<-client.Done()
+			_ = conn.Close()
+		}()
+
 		go streamAgentWSMessages(c, conn, client)
 
 		for {
diff --git a/openflare_server/service/agent_ws.go b/openflare_server/service/agent_ws.go
index c8a6a555..ea81644b 100644
--- a/openflare_server/service/agent_ws.go
+++ b/openflare_server/service/agent_ws.go
@@ -127,6 +127,21 @@ func UnregisterAgentWSClient(client *AgentWSClient) {
 	slog.Debug("agent ws connection unregistered", "node_id", client.nodeID, "client_count", count)
 }
 
+func DisconnectAgentWSClient(nodeID string) {
+	defaultAgentWSHub.mu.Lock()
+	client := defaultAgentWSHub.clients[nodeID]
+	if client != nil {
+		delete(defaultAgentWSHub.clients, nodeID)
+	}
+	count := len(defaultAgentWSHub.clients)
+	defaultAgentWSHub.mu.Unlock()
+
+	if client != nil {
+		client.Close()
+		slog.Debug("agent ws connection forcefully disconnected", "node_id", nodeID, "client_count", count)
+	}
+}
+
 func IsAgentWSConnected(nodeID string) bool {
 	defaultAgentWSHub.mu.RLock()
 	client := defaultAgentWSHub.clients[nodeID]
diff --git a/openflare_server/service/node.go b/openflare_server/service/node.go
index f001734c..c3f95b55 100644
--- a/openflare_server/service/node.go
+++ b/openflare_server/service/node.go
@@ -130,6 +130,7 @@ func DeleteNode(id uint) error {
 		return err
 	}
 	invalidateAgentTokenCache(node.AgentToken)
+	DisconnectAgentWSClient(node.NodeID)
 	return nil
 }