feat: Implement node management features including creation, updating, and deletion

- Added CreateNode, UpdateNode, and DeleteNode functions in the controller for managing nodes. - Introduced NodeInput struct for input validation during node creation and updates. - Enhanced the Node model to include DiscoveryToken and Pending status. - Updated the AgentRegister and AgentHeartbeat functions to utilize the new node management logic. - Refactored the API router to include new routes for node management. - Improved the frontend Node component to support node creation, editing, and deletion with appropriate UI feedback. - Added tests to ensure the new functionality works as expected.
2026-10-09 00:56:37 +08:00 · 2026-03-10 15:33:06 +08:00
parent 861d759f97
commit 05e75549d2
16 changed files with 916 additions and 296 deletions
@@ -2,7 +2,9 @@ package agent

 import (
 	"context"
+	"errors"
 	"log"
+	"strings"
 	"time"

 	"atsflare-agent/internal/config"
@@ -11,8 +13,9 @@ import (
 )

 type HeartbeatService interface {
-	Register(ctx context.Context, payload protocol.NodePayload) error
+	Register(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error)
 	Heartbeat(ctx context.Context, payload protocol.NodePayload) error
+	SetToken(token string)
 }

 type SyncService interface {
@@ -33,21 +36,20 @@ func (r *Runner) Run(ctx context.Context) error {
 		return err
 	}
 	log.Printf("agent runner started: node_id=%s node=%s ip=%s", nodeID, r.Config.NodeName, r.Config.NodeIP)
-	if err = r.HeartbeatService.Register(ctx, r.nodePayload(nodeID)); err != nil {
-		log.Printf("agent register failed: %v", err)
-	} else {
-		log.Printf("agent register succeeded: node_id=%s", nodeID)
-	}
-	if err = r.SyncService.SyncOnStartup(ctx); err != nil {
-		r.recordSyncError(err)
-		log.Printf("agent startup sync failed: %v", err)
-	} else {
-		log.Printf("agent startup sync completed")
-	}
-	if err = r.HeartbeatService.Heartbeat(ctx, r.nodePayload(nodeID)); err != nil {
-		log.Printf("agent startup heartbeat failed: %v", err)
-	} else {
-		log.Printf("agent startup heartbeat succeeded: node_id=%s", nodeID)
+	if r.hasAgentToken() {
+		if err = r.SyncService.SyncOnStartup(ctx); err != nil {
+			r.recordSyncError(err)
+			log.Printf("agent startup sync failed: %v", err)
+		} else {
+			log.Printf("agent startup sync completed")
+		}
+		if err = r.HeartbeatService.Heartbeat(ctx, r.nodePayload(nodeID)); err != nil {
+			log.Printf("agent startup heartbeat failed: %v", err)
+		} else {
+			log.Printf("agent startup heartbeat succeeded: node_id=%s", nodeID)
+		}
+	} else if err = r.tryRegister(ctx, &nodeID); err != nil {
+		log.Printf("agent initial discovery register failed: %v", err)
 	}

 	heartbeatTicker := time.NewTicker(r.Config.HeartbeatInterval)
@@ -61,10 +63,19 @@ func (r *Runner) Run(ctx context.Context) error {
 			log.Printf("agent runner shutting down: %v", ctx.Err())
 			return ctx.Err()
 		case <-heartbeatTicker.C:
+			if !r.hasAgentToken() {
+				if err = r.tryRegister(ctx, &nodeID); err != nil {
+					log.Printf("agent discovery register failed: %v", err)
+				}
+				continue
+			}
 			if err = r.HeartbeatService.Heartbeat(ctx, r.nodePayload(nodeID)); err != nil {
 				log.Printf("agent heartbeat failed: %v", err)
 			}
 		case <-syncTicker.C:
+			if !r.hasAgentToken() {
+				continue
+			}
 			log.Printf("agent sync tick: node_id=%s", nodeID)
 			if err = r.SyncService.SyncOnce(ctx); err != nil {
 				r.recordSyncError(err)
@@ -76,6 +87,47 @@ func (r *Runner) Run(ctx context.Context) error {
 	}
 }

+func (r *Runner) hasAgentToken() bool {
+	return strings.TrimSpace(r.Config.AgentToken) != ""
+}
+
+func (r *Runner) tryRegister(ctx context.Context, nodeID *string) error {
+	if strings.TrimSpace(r.Config.DiscoveryToken) == "" {
+		return errors.New("agent_token 为空且未配置 discovery_token")
+	}
+	log.Printf("agent discovery registration started")
+	response, err := r.HeartbeatService.Register(ctx, r.nodePayload(*nodeID))
+	if err != nil {
+		return err
+	}
+	if response == nil || strings.TrimSpace(response.AgentToken) == "" || strings.TrimSpace(response.NodeID) == "" {
+		return errors.New("discovery register response 缺少 node_id 或 agent_token")
+	}
+	snapshot, err := r.StateStore.Load()
+	if err != nil {
+		return err
+	}
+	snapshot.NodeID = response.NodeID
+	if err = r.StateStore.Save(snapshot); err != nil {
+		return err
+	}
+	r.Config.AgentToken = response.AgentToken
+	r.Config.DiscoveryToken = ""
+	if err = r.Config.Save(); err != nil {
+		return err
+	}
+	r.HeartbeatService.SetToken(response.AgentToken)
+	*nodeID = response.NodeID
+	log.Printf("agent discovery registration succeeded: node_id=%s", response.NodeID)
+	if err = r.SyncService.SyncOnStartup(ctx); err != nil {
+		r.recordSyncError(err)
+		log.Printf("agent post-register startup sync failed: %v", err)
+	} else {
+		log.Printf("agent post-register startup sync completed")
+	}
+	return nil
+}
+
 func (r *Runner) recordSyncError(err error) {
 	if err == nil || r.StateStore == nil {
 		return
@@ -3,6 +3,7 @@ package agent
 import (
 	"context"
 	"errors"
+	"os"
 	"path/filepath"
 	"sync"
 	"testing"
@@ -18,15 +19,17 @@ type fakeHeartbeatService struct {
 	registerCalls  int
 	heartbeatCalls int
 	registerErr    error
+	registerResp   *protocol.RegisterNodeResponse
 	heartbeatErrs  []error
 	onHeartbeat    func(int)
+	lastToken      string
 }

-func (f *fakeHeartbeatService) Register(ctx context.Context, payload protocol.NodePayload) error {
+func (f *fakeHeartbeatService) Register(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error) {
 	f.mu.Lock()
 	defer f.mu.Unlock()
 	f.registerCalls++
-	return f.registerErr
+	return f.registerResp, f.registerErr
 }

 func (f *fakeHeartbeatService) Heartbeat(ctx context.Context, payload protocol.NodePayload) error {
@@ -45,6 +48,12 @@ func (f *fakeHeartbeatService) Heartbeat(ctx context.Context, payload protocol.N
 	return err
 }

+func (f *fakeHeartbeatService) SetToken(token string) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.lastToken = token
+}
+
 type fakeSyncService struct {
 	mu             sync.Mutex
 	startupErr     error
@@ -90,6 +99,7 @@ func TestRunnerKeepsHeartbeatWhenStartupSyncFails(t *testing.T) {
 	}
 	runner := &Runner{
 		Config: &config.Config{
+			AgentToken:        "agent-token",
 			NodeName:          "edge-01",
 			NodeIP:            "10.0.0.8",
 			AgentVersion:      "0.1.0",
@@ -106,8 +116,8 @@ func TestRunnerKeepsHeartbeatWhenStartupSyncFails(t *testing.T) {
 	if !errors.Is(err, context.Canceled) {
 		t.Fatalf("expected context cancellation, got %v", err)
 	}
-	if heartbeatService.registerCalls != 1 {
-		t.Fatalf("expected 1 register call, got %d", heartbeatService.registerCalls)
+	if heartbeatService.registerCalls != 0 {
+		t.Fatalf("expected no discovery register call, got %d", heartbeatService.registerCalls)
 	}
 	if heartbeatService.heartbeatCalls < 2 {
 		t.Fatalf("expected heartbeat loop to continue, got %d heartbeat calls", heartbeatService.heartbeatCalls)
@@ -140,6 +150,7 @@ func TestRunnerDoesNotExitOnHeartbeatOrSyncError(t *testing.T) {
 	}
 	runner := &Runner{
 		Config: &config.Config{
+			AgentToken:        "agent-token",
 			NodeName:          "edge-01",
 			NodeIP:            "10.0.0.8",
 			AgentVersion:      "0.1.0",
@@ -156,8 +167,8 @@ func TestRunnerDoesNotExitOnHeartbeatOrSyncError(t *testing.T) {
 	if !errors.Is(err, context.Canceled) {
 		t.Fatalf("expected context cancellation, got %v", err)
 	}
-	if heartbeatService.registerCalls != 1 {
-		t.Fatalf("expected register attempt, got %d", heartbeatService.registerCalls)
+	if heartbeatService.registerCalls != 0 {
+		t.Fatalf("expected no register attempt, got %d", heartbeatService.registerCalls)
 	}
 	if syncService.syncOnceCalls == 0 {
 		t.Fatal("expected sync loop to continue after heartbeat/register errors")
@@ -170,3 +181,72 @@ func TestRunnerDoesNotExitOnHeartbeatOrSyncError(t *testing.T) {
 		t.Fatalf("expected sync error to be recorded, got %q", snapshot.LastError)
 	}
 }
+
+func TestRunnerDiscoveryRegisterUpdatesTokenAndNodeID(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
+	heartbeatService := &fakeHeartbeatService{
+		registerResp: &protocol.RegisterNodeResponse{
+			NodeID:     "node-server-assigned",
+			AgentToken: "agent-token-issued",
+			Name:       "edge-01",
+		},
+		onHeartbeat: func(callCount int) {
+			if callCount >= 1 {
+				cancel()
+			}
+		},
+	}
+	syncService := &fakeSyncService{}
+	configPath := filepath.Join(t.TempDir(), "agent.json")
+	if err := os.WriteFile(configPath, []byte(`{"server_url":"http://127.0.0.1:3000","discovery_token":"discovery-token","node_name":"edge-01","node_ip":"10.0.0.8"}`), 0o644); err != nil {
+		t.Fatalf("failed to seed config file: %v", err)
+	}
+	cfg, err := config.Load(configPath)
+	if err != nil {
+		t.Fatalf("failed to load config: %v", err)
+	}
+	runner := &Runner{
+		Config: &config.Config{
+			ServerURL:         cfg.ServerURL,
+			DiscoveryToken:    cfg.DiscoveryToken,
+			NodeName:          cfg.NodeName,
+			NodeIP:            cfg.NodeIP,
+			AgentVersion:      "0.1.0",
+			NginxVersion:      "1.25.5",
+			HeartbeatInterval: 10 * time.Millisecond,
+			SyncInterval:      20 * time.Millisecond,
+		},
+		StateStore:       stateStore,
+		HeartbeatService: heartbeatService,
+		SyncService:      syncService,
+	}
+	runner.Config = cfg
+	runner.Config.AgentVersion = "0.1.0"
+	runner.Config.NginxVersion = "1.25.5"
+	runner.Config.HeartbeatInterval = 10 * time.Millisecond
+	runner.Config.SyncInterval = 20 * time.Millisecond
+
+	err = runner.Run(ctx)
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("expected context cancellation, got %v", err)
+	}
+	if heartbeatService.registerCalls == 0 {
+		t.Fatal("expected discovery register to be attempted")
+	}
+	if heartbeatService.lastToken != "agent-token-issued" {
+		t.Fatalf("expected client token to be updated, got %q", heartbeatService.lastToken)
+	}
+	snapshot, loadErr := stateStore.Load()
+	if loadErr != nil {
+		t.Fatalf("failed to load state: %v", loadErr)
+	}
+	if snapshot.NodeID != "node-server-assigned" {
+		t.Fatalf("expected node id to be replaced, got %q", snapshot.NodeID)
+	}
+	if runner.Config.AgentToken != "agent-token-issued" || runner.Config.DiscoveryToken != "" {
+		t.Fatal("expected config token rotation to complete")
+	}
+}