AetherForge/server/internal/ai/scheduler.go

package ai

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"log"
	"sync"
	"time"
)

// SnapshotProvider supplies live agent telemetry for decision cycles.
type SnapshotProvider interface {
	ConnectedAgentIDs() []string
	AgentSnapshot(agentID string) (AgentSnapshot, bool)
}

// CommandExecutor runs parsed fleet commands.
type CommandExecutor interface {
	Execute(agentID string, cmd Command) (summary string, err error)
}

// ConfigProvider reads current Fleet AI Control settings.
type ConfigProvider interface {
	AIConfig() Config
}

// DecisionStore persists decision audit rows.
type DecisionStore interface {
	InsertAIDecision(agentID, promptHash, response, commandsExecuted string) error
}

// Scheduler runs periodic LLM decisions for online agents.
type Scheduler struct {
	cfg   ConfigProvider
	snap  SnapshotProvider
	exec  CommandExecutor
	store DecisionStore
	stop  chan struct{}
	wg    sync.WaitGroup

	lastRunMu sync.Mutex
	lastRun   map[string]time.Time
}

func NewScheduler(cfg ConfigProvider, snap SnapshotProvider, exec CommandExecutor, store DecisionStore) *Scheduler {
	return &Scheduler{
		cfg:     cfg,
		snap:    snap,
		exec:    exec,
		store:   store,
		stop:    make(chan struct{}),
		lastRun: make(map[string]time.Time),
	}
}

func (s *Scheduler) Start() {
	s.wg.Add(1)
	go s.loop()
}

func (s *Scheduler) Stop() {
	close(s.stop)
	s.wg.Wait()
}

func (s *Scheduler) loop() {
	defer s.wg.Done()
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-s.stop:
			return
		case <-ticker.C:
			s.tick()
		}
	}
}

// Tick runs one scheduler pass (exported for tests).
func (s *Scheduler) Tick() {
	s.tick()
}

func (s *Scheduler) tick() {
	if s.cfg == nil || s.snap == nil {
		return
	}
	cfg := s.cfg.AIConfig()
	if !cfg.Enabled {
		return
	}
	interval := time.Duration(cfg.IntervalSec) * time.Second
	if interval < time.Second {
		interval = 60 * time.Second
	}
	ids := s.snap.ConnectedAgentIDs()
	for i, agentID := range ids {
		if !s.shouldRun(agentID, interval, i) {
			continue
		}
		s.runAgent(context.Background(), agentID, cfg)
	}
}

func (s *Scheduler) shouldRun(agentID string, interval time.Duration, staggerIndex int) bool {
	s.lastRunMu.Lock()
	defer s.lastRunMu.Unlock()
	last, ok := s.lastRun[agentID]
	if !ok {
		offset := time.Duration(staggerIndex%max(1, int(interval/time.Second))) * time.Second
		if offset > 0 {
			s.lastRun[agentID] = time.Now().Add(-interval + offset)
		}
		return true
	}
	return time.Since(last) >= interval
}

func (s *Scheduler) markRun(agentID string) {
	s.lastRunMu.Lock()
	s.lastRun[agentID] = time.Now()
	s.lastRunMu.Unlock()
}

func (s *Scheduler) runAgent(ctx context.Context, agentID string, cfg Config) {
	snap, ok := s.snap.AgentSnapshot(agentID)
	if !ok {
		return
	}
	userPrompt := BuildUserPrompt(snap)
	systemPrompt := SystemPrompt()
	promptHash := hashPrompt(userPrompt)

	decide := Decide
	if DecideFunc != nil {
		decide = DecideFunc
	}
	response, err := decide(ctx, cfg.Endpoint, cfg.Model, systemPrompt, userPrompt)
	if err != nil {
		log.Printf("[fleet-ai] agent %s decide: %v", agentID, err)
		if s.store != nil {
			_ = s.store.InsertAIDecision(agentID, promptHash, "", "error:"+err.Error())
		}
		s.markRun(agentID)
		return
	}

	cmds := ParseCommands(response)
	results := make([]string, 0, len(cmds))
	for _, cmd := range cmds {
		if cmd.Type == CmdNoop {
			results = append(results, "ok")
			continue
		}
		if s.exec == nil {
			results = append(results, "no executor")
			continue
		}
		sum, execErr := s.exec.Execute(agentID, cmd)
		if execErr != nil {
			results = append(results, "err:"+execErr.Error())
		} else {
			results = append(results, sum)
		}
	}
	executed := FormatExecuted(cmds, results)
	if s.store != nil {
		_ = s.store.InsertAIDecision(agentID, promptHash, response, executed)
	}
	s.markRun(agentID)
}

func hashPrompt(prompt string) string {
	h := sha256.Sum256([]byte(prompt))
	return hex.EncodeToString(h[:8])
}

// DecideFunc allows tests to override LLM calls.
var DecideFunc func(ctx context.Context, endpoint, model, systemPrompt, userPrompt string) (string, error)

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}