diff --git a/openflare_server/controller/agent.go b/openflare_server/controller/agent.go index 4555c61e..508c4d6a 100644 --- a/openflare_server/controller/agent.go +++ b/openflare_server/controller/agent.go @@ -223,8 +223,13 @@ func handleAgentWSStatus(c *gin.Context, node *model.Node, message service.Agent slog.Debug("agent ws status payload decode failed", "node_id", node.NodeID, "error", err) return } + freshNode, err := model.GetNodeByNodeID(node.NodeID) + if err != nil { + slog.Debug("agent ws status reload node failed", "node_id", node.NodeID, "error", err) + return + } payload.IP = service.ResolveReportedNodeIP(payload.IP, c.Request.RemoteAddr) - response, err := service.HeartbeatNode(node, payload) + response, err := service.HeartbeatNode(freshNode, payload) if err != nil { slog.Debug("agent ws status handling failed", "node_id", node.NodeID, "error", err) return diff --git a/openflare_server/service/node.go b/openflare_server/service/node.go index 903a3d47..f001734c 100644 --- a/openflare_server/service/node.go +++ b/openflare_server/service/node.go @@ -171,6 +171,11 @@ func RequestNodeAgentUpdate(id uint, input NodeAgentUpdateInput) (*NodeView, err return nil, err } refreshAgentTokenCache(node) + if SendAgentWSSettings(node.NodeID, buildAgentSettings(node, true, channel.String(), tagName, node.RestartOpenrestyRequested)) { + slog.Debug("agent manual update pushed via ws", "node_id", node.NodeID, "channel", channel.String(), "tag", tagName) + } else { + slog.Debug("agent manual update waiting for next heartbeat", "node_id", node.NodeID, "channel", channel.String(), "tag", tagName) + } slog.Info("agent manual update requested", "node_id", node.NodeID, "name", node.Name, "channel", channel.String(), "tag", tagName) return buildNodeView(node), nil } diff --git a/openflare_server/service/node_update_test.go b/openflare_server/service/node_update_test.go index cd416797..690f0e06 100644 --- a/openflare_server/service/node_update_test.go +++ b/openflare_server/service/node_update_test.go @@ -120,6 +120,9 @@ func TestRequestNodeAgentPreviewUpdate(t *testing.T) { SetUpdateHTTPClientForTest(originalClient) }) + wsClient := RegisterAgentWSClient(node.NodeID) + defer UnregisterAgentWSClient(wsClient) + updated, err := RequestNodeAgentUpdate(node.ID, NodeAgentUpdateInput{ Channel: "preview", TagName: "v0.5.0-rc.1", @@ -136,6 +139,21 @@ func TestRequestNodeAgentPreviewUpdate(t *testing.T) { if updated.UpdateTag != "v0.5.0-rc.1" { t.Fatalf("unexpected update tag: %s", updated.UpdateTag) } + select { + case message := <-wsClient.Messages(): + if message.Type != AgentWSMessageTypeSettings { + t.Fatalf("expected settings message, got %s", message.Type) + } + settings, ok := message.Payload.(*AgentSettings) + if !ok { + t.Fatalf("expected agent settings payload, got %T", message.Payload) + } + if !settings.UpdateNow || settings.UpdateChannel != "preview" || settings.UpdateTag != "v0.5.0-rc.1" { + t.Fatalf("unexpected pushed settings: %+v", settings) + } + case <-time.After(time.Second): + t.Fatal("expected websocket settings push for manual update") + } } func TestHeartbeatNodeReturnsPreviewUpdateSettings(t *testing.T) {