selfheal-lab-apps/harness/run.go

package main

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"strings"
	"sync"
	"time"
)

// RunManager executes one run at a time: for each scenario baseline -> inject -> observe (judge every 30s,
// alerts/diagnoses/actions arrive through the HTTP handlers) -> revert -> verdict -> cooldown (design 5.2, 8).
type RunManager struct {
	app    *App
	mu     sync.Mutex
	active *activeRun
}

type activeRun struct {
	RunID      string
	Campaign   string
	Model      string // LLM for this run's diagnoses ("" = AI VM default); passed to the AI VM on every /diagnose
	ScenarioID string // "" between scenarios
	Scenario   *Scenario
	TInject    time.Time
	cancel     context.CancelFunc
}

type RunRequest struct {
	Campaign        string   `json:"campaign"`
	Model           string   `json:"model"` // optional: LLM id for the diagnoses (e.g. gpt-6-luna, claude-sonnet-5-5); "" = AI VM default
	ScenarioIDs     []string `json:"scenario_ids"`
	BaselineS       *int     `json:"baseline_s"`        // default 300 (design 8)
	CooldownS       *int     `json:"cooldown_s"`        // default 300
	ObserveTimeoutS *int     `json:"observe_timeout_s"` // default 1800
	MinObserveS     *int     `json:"min_observe_s"`     // default 180: earliest t_recover (12 probe successes after inject)
	DriftBefore     *int     `json:"drift_before"`
	GitSHA          string   `json:"git_sha"`
	Notes           string   `json:"notes"`
	SkipRevert      bool     `json:"skip_revert"`
}

// modelRe bounds the model id (empty = AI VM default). The AI VM routes by prefix: claude-* -> Anthropic, else OpenAI.
var modelRe = regexp.MustCompile(`^[A-Za-z0-9._:-]{0,64}$`)

func NewRunManager(a *App) *RunManager { return &RunManager{app: a} }

func (m *RunManager) current() *activeRun {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.active == nil {
		return nil
	}
	c := *m.active
	return &c
}

func newID(prefix string) string {
	b := make([]byte, 4)
	_, _ = rand.Read(b)
	return prefix + time.Now().UTC().Format("20060102-150405") + "-" + hex.EncodeToString(b)
}

func (a *App) handleRunCreate(w http.ResponseWriter, r *http.Request) {
	var req RunRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeErr(w, 400, err)
		return
	}
	if req.Campaign == "" || len(req.ScenarioIDs) == 0 {
		writeErr(w, 400, errors.New("campaign and scenario_ids are required"))
		return
	}
	for _, id := range req.ScenarioIDs {
		if _, ok := a.scen[id]; !ok {
			writeErr(w, 400, fmt.Errorf("unknown scenario %q", id))
			return
		}
	}
	if !modelRe.MatchString(req.Model) {
		writeErr(w, 400, errors.New("model must match [A-Za-z0-9._:-]{1,64}"))
		return
	}
	m := a.runs
	m.mu.Lock()
	if m.active != nil {
		id := m.active.RunID
		m.mu.Unlock()
		writeJSON(w, 409, map[string]string{"error": "a run is active", "run_id": id})
		return
	}
	runID := newID("run-")
	ctx, cancel := context.WithCancel(context.Background())
	m.active = &activeRun{RunID: runID, Campaign: req.Campaign, Model: req.Model, cancel: cancel}
	m.mu.Unlock()
	if err := a.store.CreateRun(r.Context(), runID, req.Campaign, req.DriftBefore, req.GitSHA, req.Notes, req.Model); err != nil {
		m.mu.Lock()
		m.active = nil
		m.mu.Unlock()
		cancel()
		writeErr(w, 500, err)
		return
	}
	mRuns.Set(1)
	go m.execute(ctx, runID, req)
	writeJSON(w, 202, map[string]any{"run_id": runID, "scenario_ids": req.ScenarioIDs, "model": req.Model})
}

func (a *App) handleRunAbort(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")
	m := a.runs
	m.mu.Lock()
	if m.active == nil || m.active.RunID != id {
		m.mu.Unlock()
		writeErr(w, 404, errors.New("run not active"))
		return
	}
	m.active.cancel()
	m.mu.Unlock()
	writeJSON(w, 202, map[string]string{"run_id": id, "status": "aborting"})
}

func (a *App) handleRunFinish(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")
	var body struct {
		DriftAfter *int `json:"drift_after"`
	}
	_ = json.NewDecoder(r.Body).Decode(&body)
	if err := a.store.FinishRun(r.Context(), id, "finished", body.DriftAfter); err != nil {
		writeErr(w, 500, err)
		return
	}
	writeJSON(w, 200, map[string]any{"run_id": id, "drift_after": body.DriftAfter})
}

func (a *App) handleRunList(w http.ResponseWriter, r *http.Request) {
	runs, err := a.store.Runs(r.Context(), "")
	if err != nil {
		writeErr(w, 500, err)
		return
	}
	writeJSON(w, 200, map[string]any{"active": a.runs.current(), "runs": runs})
}

func (a *App) handleRunGet(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")
	runs, err := a.store.Runs(r.Context(), id)
	if err != nil || len(runs) == 0 {
		writeErr(w, 404, errors.New("run not found"))
		return
	}
	inj, _ := a.store.Injections(r.Context(), id)
	ev, _ := a.store.Events(r.Context(), id, "")
	vd, _ := a.store.Verdicts(r.Context(), id)
	writeJSON(w, 200, map[string]any{"run": runs[0], "injections": inj, "events": ev, "verdicts": vd})
}

func dur(p *int, def time.Duration) time.Duration {
	if p == nil {
		return def
	}
	return time.Duration(*p) * time.Second
}

func (m *RunManager) execute(ctx context.Context, runID string, req RunRequest) {
	a := m.app
	status := "finished"
	defer func() {
		m.mu.Lock()
		m.active = nil
		m.mu.Unlock()
		mRuns.Set(0)
		_ = a.store.FinishRun(context.Background(), runID, status, nil)
		log.Printf("run %s %s", runID, status)
	}()
	baseline, cooldown := dur(req.BaselineS, 5*time.Minute), dur(req.CooldownS, 5*time.Minute)
	observeMax, minObserve := dur(req.ObserveTimeoutS, a.cfg.ObserveMax), dur(req.MinObserveS, 3*time.Minute)
	for i, sid := range req.ScenarioIDs {
		sc := a.scen[sid]
		if err := m.runScenario(ctx, runID, sc, baseline, observeMax, minObserve, req.SkipRevert); err != nil {
			if errors.Is(err, context.Canceled) {
				status = "aborted"
				return
			}
			log.Printf("run %s scenario %s: %v", runID, sid, err)
		}
		if i < len(req.ScenarioIDs)-1 && !sleepCtx(ctx, cooldown) {
			status = "aborted"
			return
		}
	}
}

func sleepCtx(ctx context.Context, d time.Duration) bool {
	if d <= 0 {
		return true
	}
	select {
	case <-time.After(d):
		return true
	case <-ctx.Done():
		return false
	}
}

func (m *RunManager) setScenario(sc *Scenario, tInject time.Time) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.active == nil {
		return
	}
	if sc == nil {
		m.active.ScenarioID, m.active.Scenario, m.active.TInject = "", nil, time.Time{}
		return
	}
	m.active.ScenarioID, m.active.Scenario, m.active.TInject = sc.ID, sc, tInject
}

func (m *RunManager) runScenario(ctx context.Context, runID string, sc *Scenario, baseline, observeMax, minObserve time.Duration, skipRevert bool) error {
	a := m.app
	sid := sc.ID
	bg := context.Background()
	a.store.Event(bg, &runID, &sid, "note", "runner", map[string]any{"phase": "baseline", "seconds": baseline.Seconds()})
	if !sleepCtx(ctx, baseline) {
		return ctx.Err()
	}
	// inject
	m.setScenario(sc, time.Now().UTC())
	defer m.setScenario(nil, time.Time{})
	injCtx, cancelInj := context.WithTimeout(ctx, 5*time.Minute)
	params, err := a.runStep(injCtx, sc, sc.Inject, runID, nil)
	cancelInj()
	tInject := time.Now().UTC()
	injID, dbErr := a.store.InsertInjection(bg, runID, sid, sc.Inject.Method, params)
	if dbErr != nil {
		log.Printf("injection insert: %v", dbErr)
	}
	if err != nil {
		a.store.Event(bg, &runID, &sid, "note", "runner", map[string]any{"phase": "inject", "error": err.Error()})
		_ = a.store.InsertVerdict(bg, Verdict{RunID: runID, ScenarioID: sid, Recovered: false, Notes: "injection failed: " + err.Error()})
		return err
	}
	m.setScenario(sc, tInject)
	a.store.Event(bg, &runID, &sid, "note", "runner", map[string]any{"phase": "inject", "t_inject": tInject, "params": params})
	// observe: judge every 30s until recovered (after minObserve) or timeout
	var tRecover *time.Time
	var last SLOState
	deadline := tInject.Add(observeMax)
	for {
		if !sleepCtx(ctx, a.cfg.JudgeEvery) {
			break
		}
		last = a.evaluateSLO(ctx)
		if last.Recovered && time.Since(tInject) >= minObserve {
			t := time.Now().UTC()
			tRecover = &t
			a.store.Event(bg, &runID, &sid, "recovered", "judge", last)
			break
		}
		if time.Now().After(deadline) {
			a.store.Event(bg, &runID, &sid, "timeout", "judge", last)
			break
		}
	}
	aborted := ctx.Err() != nil
	// revert (always, even on abort, unless the scenario is irreversible by design or skip_revert)
	if !skipRevert && sc.Revert.Method != "none" {
		rvCtx, cancelRv := context.WithTimeout(bg, 10*time.Minute)
		rres, rerr := a.runStep(rvCtx, sc, sc.Revert, runID, params)
		cancelRv()
		if rerr != nil {
			a.store.Event(bg, &runID, &sid, "note", "runner", map[string]any{"phase": "revert", "error": rerr.Error(), "result": rres})
		} else {
			a.store.Event(bg, &runID, &sid, "note", "runner", map[string]any{"phase": "revert", "result": rres})
		}
	}
	if injID != 0 {
		_ = a.store.ClearInjection(bg, injID)
	}
	// verdict
	v := a.buildVerdict(bg, runID, sc, tInject, tRecover, last)
	if aborted {
		v.Notes += " aborted;"
	}
	if err := a.store.InsertVerdict(bg, v); err != nil {
		log.Printf("verdict insert: %v", err)
	}
	if aborted {
		return ctx.Err()
	}
	return nil
}

// buildVerdict derives design 5.4 metrics from the events of this scenario.
func (a *App) buildVerdict(ctx context.Context, runID string, sc *Scenario, tInject time.Time, tRecover *time.Time, last SLOState) Verdict {
	v := Verdict{RunID: runID, ScenarioID: sc.ID, Recovered: tRecover != nil}
	if tRecover != nil {
		s := tRecover.Sub(tInject).Seconds()
		v.MTTRs = &s
	}
	evs, _ := a.store.Events(ctx, runID, sc.ID)
	var tAlert, tProposal, tApproval *time.Time
	layer := 1
	misjudge := 0
	var diagScore *float64
	var actionMatch *bool
	for _, e := range evs {
		var p map[string]any
		_ = json.Unmarshal(e.Payload, &p)
		switch e.Kind {
		case "alert":
			if tAlert == nil && e.TS.After(tInject) && p["status"] == "firing" {
				t := e.TS
				tAlert = &t
			}
		case "diagnosis":
			if e, _ := p["error"].(string); e != "" { // AI VM unreachable / malformed: not a diagnosis, not a misjudge
				continue
			}
			if layer < 2 {
				layer = 2
			}
			label, _ := p["root_cause_label"].(string)
			s := scoreDiagnosis(label, sc.GroundTruthLabel)
			diagScore = &s
			if s == 0 {
				misjudge++
			}
			if pa, ok := p["proposed_action"].(string); ok {
				b := pa == sc.GroundTruthAction
				actionMatch = &b
				if sc.GroundTruthAction == "escalate" && pa != "escalate" {
					misjudge++
				}
			}
		case "proposal":
			if tProposal == nil {
				t := e.TS
				tProposal = &t
			}
		case "approval":
			if tApproval == nil {
				t := e.TS
				tApproval = &t
			}
		case "action":
			switch e.Actor {
			case "human":
				layer = 3
			case "ai-approved":
				if layer < 2 {
					layer = 2
				}
			}
		case "misjudge":
			misjudge++
		}
	}
	if tAlert != nil {
		s := tAlert.Sub(tInject).Seconds()
		v.MTTDs = &s
	}
	if tProposal != nil && tApproval != nil {
		s := tApproval.Sub(*tProposal).Seconds()
		v.MTTAs = &s
	}
	v.LayerActual = &layer
	v.DiagnosisScore = diagScore
	v.ActionMatch = actionMatch
	v.MisjudgeCount = misjudge
	b, _ := json.Marshal(last)
	v.Notes = "last_slo=" + string(b)
	return v
}

// design 9: exact label 1.0, same top category (network/compute/db/config) 0.5, else 0
func scoreDiagnosis(got, want string) float64 {
	if got == "" || got == "unknown" {
		return 0
	}
	if got == want {
		return 1
	}
	if category(got) != "" && category(got) == category(want) {
		return 0.5
	}
	return 0
}

func category(label string) string {
	switch {
	case contains(label, "network", "latency", "partition", "loss"):
		return "network"
	case contains(label, "pod", "process", "instance", "memory", "cpu"):
		return "compute"
	case contains(label, "db", "replication", "schema", "primary"):
		return "db"
	case contains(label, "drift", "config", "gitops"):
		return "config"
	}
	return ""
}

func contains(s string, subs ...string) bool {
	for _, x := range subs {
		if strings.Contains(s, x) {
			return true
		}
	}
	return false
}