[优化] 增加断开 WebSocket 客户端的功能，优化连接管理

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:

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 {

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]

openflare_server/service/node.go

@@ -130,6 +130,7 @@ func DeleteNode(id uint) error {
 		return err
 	}
 	invalidateAgentTokenCache(node.AgentToken)
+	DisconnectAgentWSClient(node.NodeID)
 	return nil
 }