feat[agent-manager](recovery): port the YAML-driven agent recovery subsystem

Author: Kbayero

agent-manager/Dockerfile

@@ -4,6 +4,8 @@ COPY agent-manager /app/
 COPY ./dependencies/agent/ /dependencies/agent/
 COPY ./dependencies/collector/ /dependencies/collector/
 
+COPY ./recoveries/ /app/recoveries/
+
 # Install jq
 RUN apt-get update && \
     apt-get install -y ca-certificates jq wget && \

agent-manager/agent/agent_imp.go

@@ -117,6 +117,11 @@ func (s *AgentService) RegisterAgent(ctx context.Context, req *AgentRequest) (*A
 	}
 
 	catcher.Info("Agent registered correctly", map[string]any{"hostname": agent.Hostname, "id": agent.ID, "process": "agent-manager"})
+
+	if OnAgentRegisterHook != nil {
+		OnAgentRegisterHook(agent)
+	}
+
 	return &AuthResponse{
 		Id:  uint32(agent.ID),
 		Key: key,
@@ -171,6 +176,10 @@ func (s *AgentService) UpdateAgent(ctx context.Context, req *AgentRequest) (*Aut
 		return nil, status.Errorf(codes.Internal, "failed to update agent: %v", err)
 	}
 
+	if OnAgentUpdateHook != nil {
+		OnAgentUpdateHook(agent)
+	}
+
 	res := &AuthResponse{
 		Id:  uint32(agent.ID),
 		Key: key,
@@ -255,6 +264,10 @@ func (s *AgentService) AgentStream(stream AgentService_AgentStreamServer) error
 	s.AgentStreamMap[idUint] = stream
 	s.AgentStreamMutex.Unlock()
 
+	if OnAgentConnectHook != nil {
+		OnAgentConnectHook(stream.Context(), idUint)
+	}
+
 	for {
 		in, err := stream.Recv()
 		if err == io.EOF {
@@ -288,7 +301,7 @@ func (s *AgentService) AgentStream(stream AgentService_AgentStreamServer) error
 					CmdId:      cmdID,
 					ExecutedAt: msg.Result.ExecutedAt,
 				}
-			} else {
+			} else if OnCommandResultHook == nil || !OnCommandResultHook(msg.Result) {
 				catcher.Error("failed to find result channel for CmdID", nil, map[string]any{"cmdID": cmdID, "process": "agent-manager"})
 			}
 			s.CommandResultChannelM.Unlock()
@@ -338,16 +351,26 @@ func (s *AgentService) ProcessCommand(stream PanelService_ProcessCommandServer)
 			catcher.Error("unable to create a new command history", err, map[string]any{"process": "agent-manager"})
 		}
 
-		err = agentStream.Send(&BidirectionalStream{
-			StreamMessage: &BidirectionalStream_Command{
-				Command: &UtmCommand{
-					AgentId: cmd.AgentId,
-					Command: replaceSecretValues(cmd.Command),
-					CmdId:   cmdID,
-					Shell:   cmd.Shell,
+		var lock sync.Locker
+		if LockStreamHook != nil {
+			lock = LockStreamHook(uint(streamId))
+		}
+		func() {
+			if lock != nil {
+				lock.Lock()
+				defer lock.Unlock()
+			}
+			err = agentStream.Send(&BidirectionalStream{
+				StreamMessage: &BidirectionalStream_Command{
+					Command: &UtmCommand{
+						AgentId: cmd.AgentId,
+						Command: replaceSecretValues(cmd.Command),
+						CmdId:   cmdID,
+						Shell:   cmd.Shell,
+					},
 				},
-			},
-		})
+			})
+		}()
 		if err != nil {
 			return status.Errorf(codes.Internal, "failed to send command to agent: %v", err)
 		}

agent-manager/agent/recovery_hooks.go

@@ -0,0 +1,40 @@
+package agent
+
+import (
+	"context"
+	"sync"
+
+	"github.com/utmstack/UTMStack/agent-manager/models"
+)
+
+var (
+	OnAgentConnectHook  func(ctx context.Context, agentID uint)
+	OnAgentRegisterHook func(agent *models.Agent)
+	OnAgentUpdateHook   func(agent *models.Agent)
+	OnCommandResultHook func(result *CommandResult) bool
+	LockStreamHook      func(agentID uint) sync.Locker
+)
+
+func RegisterRecoveryHooks(
+	onConnect func(ctx context.Context, agentID uint),
+	onRegister func(agent *models.Agent),
+	onUpdate func(agent *models.Agent),
+	onResult func(result *CommandResult) bool,
+	lockStream func(agentID uint) sync.Locker,
+) {
+	if OnAgentConnectHook == nil {
+		OnAgentConnectHook = onConnect
+	}
+	if OnAgentRegisterHook == nil {
+		OnAgentRegisterHook = onRegister
+	}
+	if OnAgentUpdateHook == nil {
+		OnAgentUpdateHook = onUpdate
+	}
+	if OnCommandResultHook == nil {
+		OnCommandResultHook = onResult
+	}
+	if LockStreamHook == nil {
+		LockStreamHook = lockStream
+	}
+}

agent-manager/config/recovery.go

@@ -0,0 +1,43 @@
+package config
+
+import (
+	"os"
+	"strconv"
+)
+
+func envOrDefault(key, fallback string) string {
+	if v := os.Getenv(key); v != "" {
+		return v
+	}
+	return fallback
+}
+
+func envBoolOrDefault(key string, fallback bool) bool {
+	v := os.Getenv(key)
+	if v == "" {
+		return fallback
+	}
+	b, err := strconv.ParseBool(v)
+	if err != nil {
+		return fallback
+	}
+	return b
+}
+
+func envIntOrDefault(key string, fallback int) int {
+	v := os.Getenv(key)
+	if v == "" {
+		return fallback
+	}
+	n, err := strconv.Atoi(v)
+	if err != nil {
+		return fallback
+	}
+	return n
+}
+
+var (
+	RecoveryDir               = envOrDefault("RECOVERY_DIR", "/app/recoveries/")
+	RecoveryDispatchEnabled   = envBoolOrDefault("RECOVERY_DISPATCH_ENABLED", true)
+	RecoveryGlobalConcurrency = envIntOrDefault("RECOVERY_GLOBAL_CONCURRENCY", 50)
+)

agent-manager/database/db.go

@@ -22,6 +22,10 @@ type DB struct {
 	locker sync.RWMutex
 }
 
+func (d *DB) GormDB() *gorm.DB {
+	return d.conn
+}
+
 func (d *DB) Migrate(data ...interface{}) error {
 	d.locker.Lock()
 	defer d.locker.Unlock()

agent-manager/database/migration.go

@@ -4,7 +4,7 @@ import "github.com/utmstack/UTMStack/agent-manager/models"
 
 func MigrateDatabase() error {
 	db := GetDB()
-	err := db.Migrate(&models.Agent{}, &models.AgentCommand{}, &models.LastSeen{}, &models.Collector{}, &models.ConnectionKey{})
+	err := db.Migrate(&models.Agent{}, &models.AgentCommand{}, &models.LastSeen{}, &models.Collector{}, &models.ConnectionKey{}, &models.Recovery{}, &models.RecoveryTarget{})
 	if err != nil {
 		return err
 	}

agent-manager/go.mod

@@ -25,7 +25,7 @@ require (
 	github.com/go-playground/universal-translator v0.18.1 // indirect
 	github.com/go-playground/validator/v10 v10.30.1 // indirect
 	github.com/goccy/go-json v0.10.5 // indirect
-	github.com/goccy/go-yaml v1.19.2 // indirect
+	github.com/goccy/go-yaml v1.19.2
 	github.com/jackc/pgpassfile v1.0.0 // indirect
 	github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
 	github.com/jackc/pgx/v5 v5.8.0 // indirect
@@ -47,7 +47,7 @@ require (
 	golang.org/x/arch v0.24.0 // indirect
 	golang.org/x/crypto v0.49.0 // indirect
 	golang.org/x/net v0.52.0 // indirect
-	golang.org/x/sync v0.20.0 // indirect
+	golang.org/x/sync v0.20.0
 	golang.org/x/sys v0.42.0 // indirect
 	golang.org/x/text v0.35.0 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20260226221140-a57be14db171 // indirect

agent-manager/main.go

@@ -1,25 +1,89 @@
 package main
 
 import (
+	"context"
 	"os"
+	"os/signal"
+	"sync"
+	"syscall"
 	"time"
 
 	"github.com/threatwinds/go-sdk/catcher"
 	"github.com/utmstack/UTMStack/agent-manager/agent"
 	"github.com/utmstack/UTMStack/agent-manager/database"
+	"github.com/utmstack/UTMStack/agent-manager/recovery"
 	"github.com/utmstack/UTMStack/agent-manager/updates"
 )
 
+type recoveryProvider struct {
+	server       *agent.AgentService
+	lastSeenServ *agent.LastSeenService
+}
+
+func (p *recoveryProvider) GetStream(agentID uint) (recovery.AgentStream, bool) {
+	p.server.AgentStreamMutex.Lock()
+	defer p.server.AgentStreamMutex.Unlock()
+	s, ok := p.server.AgentStreamMap[agentID]
+	if !ok {
+		return nil, false
+	}
+	return s, true
+}
+
+func (p *recoveryProvider) IsOnline(agentID uint) bool {
+	if p.lastSeenServ == nil {
+		return false
+	}
+	st, _, err := p.lastSeenServ.GetLastSeenStatus(agentID, "agent")
+	if err != nil {
+		return false
+	}
+	return st == agent.Status_ONLINE
+}
+
 func main() {
 	catcher.Info("Starting Agent Manager", map[string]any{"process": "agent-manager"})
 
-	err := database.MigrateDatabase()
-	if err != nil {
+	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGINT)
+	defer stop()
+
+	if err := database.MigrateDatabase(); err != nil {
 		_ = catcher.Error("failed to migrate database", err, map[string]any{"process": "agent-manager"})
 		time.Sleep(5 * time.Second)
 		os.Exit(1)
 	}
 
+	if err := agent.InitAgentService(); err != nil {
+		_ = catcher.Error("failed to init agent service", err, map[string]any{"process": "agent-manager"})
+		time.Sleep(5 * time.Second)
+		os.Exit(1)
+	}
+
+	go agent.InitCollectorService()
+	agent.InitLastSeenService()
+
+	recovery.SetStreamProvider(&recoveryProvider{
+		server:       agent.AgentServ,
+		lastSeenServ: agent.LastSeenServ,
+	})
+	agent.RegisterRecoveryHooks(
+		recovery.OnAgentConnect,
+		recovery.OnAgentRegister,
+		recovery.OnAgentUpdate,
+		recovery.OnCommandResult,
+		func(agentID uint) sync.Locker { return recovery.StreamMutex.For(agentID) },
+	)
+	if err := recovery.Init(ctx, database.GetDB().GormDB()); err != nil {
+		_ = catcher.Error("failed to init recovery", err, map[string]any{"process": "agent-manager"})
+	}
+
 	go updates.InitUpdatesManager()
-	agent.InitGrpcServer()
+	go agent.StartGrpcServer()
+
+	<-ctx.Done()
+	catcher.Info("Shutdown signal received; draining recovery dispatches", map[string]any{"process": "agent-manager"})
+	if err := recovery.Shutdown(15 * time.Second); err != nil {
+		_ = catcher.Error("recovery shutdown error", err, map[string]any{"process": "agent-manager"})
+	}
+	catcher.Info("Agent Manager shut down cleanly", map[string]any{"process": "agent-manager"})
 }

agent-manager/models/recovery.go

@@ -0,0 +1,48 @@
+package models
+
+import (
+	"time"
+
+	"gorm.io/gorm"
+)
+
+// Recovery represents an imported YAML recovery directive.
+// One row per yaml_id; source of truth is the YAML filesystem.
+type Recovery struct {
+	gorm.Model
+	YamlID           string    `gorm:"type:varchar(128);not null;uniqueIndex"`
+	YamlHash         string    `gorm:"type:char(64);not null"`
+	Name             string    `gorm:"type:varchar(255);not null"`
+	Description      string    `gorm:"type:text"`
+	Shell            string    `gorm:"type:varchar(16);not null"`
+	TargetOS         string    `gorm:"type:varchar(32);not null"`
+	TargetVersionLte string    `gorm:"type:varchar(32)"`
+	SuccessPattern   string    `gorm:"type:varchar(255);not null;default:'RECOVERY_OK'"`
+	MaxConcurrency   int       `gorm:"not null;default:100"`
+	MaxAttempts      int       `gorm:"not null;default:3"`
+	RetryAfterSecs   int       `gorm:"type:int;not null;default:1800"`
+	AckTimeoutSecs   int       `gorm:"type:int;not null;default:300"`
+	ExpiresAt        time.Time `gorm:"not null"`
+	DryRun           bool      `gorm:"not null;default:false"`
+	Script           string    `gorm:"type:text;not null"`
+	DependsOnYamlID  string    `gorm:"type:varchar(128)"`
+	Status           string    `gorm:"type:varchar(16);not null;default:'ACTIVE';index:idx_recoveries_status"`
+	BlockedReason    string    `gorm:"type:varchar(512)"`
+	ImportedAt       time.Time `gorm:"not null;autoCreateTime"`
+}
+
+// RecoveryTarget represents one agent's execution slot for a given recovery.
+// One row per (recovery_id, agent_id) pair.
+type RecoveryTarget struct {
+	gorm.Model
+	RecoveryID    uint     `gorm:"not null;uniqueIndex:idx_recovery_target_unique;index:idx_target_recovery_status,priority:1"`
+	Recovery      Recovery `gorm:"foreignKey:RecoveryID"`
+	AgentID       uint     `gorm:"not null;uniqueIndex:idx_recovery_target_unique;index:idx_target_status_agent,priority:2"`
+	Status        string   `gorm:"type:varchar(16);not null;default:'PENDING';index:idx_target_status_agent,priority:1;index:idx_target_recovery_status,priority:2"`
+	Attempts      int      `gorm:"not null;default:0"`
+	LastAttemptAt *time.Time
+	LastCmdID     string `gorm:"type:varchar(64)"`
+	LastResult    string `gorm:"type:text"`
+	LastError     string `gorm:"type:varchar(512)"`
+	CompletedAt   *time.Time
+}