selfheal-lab-apps/harness/store.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
)

// Store wraps the measure DB (experiment design 5.3). Schema is created at start-up; idempotent.
type Store struct{ pool *pgxpool.Pool }

const measureSchema = `
CREATE TABLE IF NOT EXISTS runs (
  run_id       text PRIMARY KEY,
  campaign     text NOT NULL,
  started_at   timestamptz NOT NULL DEFAULT now(),
  finished_at  timestamptz,
  drift_before int,
  drift_after  int,
  git_sha      text,
  status       text NOT NULL DEFAULT 'running',
  notes        text
);
ALTER TABLE runs ADD COLUMN IF NOT EXISTS model text;  -- LLM requested for the run's diagnoses (campaign-level model selection, 2026-10-14)
CREATE TABLE IF NOT EXISTS scenarios (
  scenario_id         text PRIMARY KEY,
  path                text NOT NULL,
  asset               text NOT NULL,
  expected_layer      int  NOT NULL,
  ground_truth_label  text NOT NULL,
  ground_truth_action text NOT NULL,
  injection_spec      jsonb NOT NULL
);
CREATE TABLE IF NOT EXISTS injections (
  id          bigserial PRIMARY KEY,
  run_id      text NOT NULL REFERENCES runs(run_id),
  scenario_id text NOT NULL,
  t_inject    timestamptz,
  t_clear     timestamptz,
  method      text NOT NULL,
  params      jsonb NOT NULL DEFAULT '{}'::jsonb
);
CREATE TABLE IF NOT EXISTS events (
  id          bigserial PRIMARY KEY,
  run_id      text,
  scenario_id text,
  ts          timestamptz NOT NULL DEFAULT now(),
  kind        text NOT NULL,   -- alert | diagnosis | proposal | approval | action | recovered | timeout | misjudge | note
  actor       text NOT NULL,
  payload     jsonb NOT NULL DEFAULT '{}'::jsonb
);
CREATE INDEX IF NOT EXISTS events_run_idx ON events (run_id, scenario_id, ts);
CREATE TABLE IF NOT EXISTS verdicts (
  id              bigserial PRIMARY KEY,
  run_id          text NOT NULL REFERENCES runs(run_id),
  scenario_id     text NOT NULL,
  layer_actual    int,
  recovered       boolean NOT NULL,
  mttd_s          double precision,
  mtta_s          double precision,
  mttr_s          double precision,
  diagnosis_score double precision,
  action_match    boolean,
  misjudge_count  int NOT NULL DEFAULT 0,
  notes           text,
  created_at      timestamptz NOT NULL DEFAULT now()
);
CREATE TABLE IF NOT EXISTS proposals (
  token        text PRIMARY KEY,
  run_id       text,
  scenario_id  text,
  action_id    text NOT NULL,
  target       jsonb NOT NULL DEFAULT '{}'::jsonb,
  created_at   timestamptz NOT NULL DEFAULT now(),
  approved_at  timestamptz,
  approved_by  text,
  executed_at  timestamptz,
  expires_at   timestamptz NOT NULL
);`

func NewStore(ctx context.Context, dsn string) (*Store, error) {
	var pool *pgxpool.Pool
	var err error
	for i := 0; i < 30; i++ { // measure-db container may still be starting
		pool, err = pgxpool.New(ctx, dsn)
		if err == nil {
			if err = pool.Ping(ctx); err == nil {
				break
			}
			pool.Close()
		}
		time.Sleep(2 * time.Second)
	}
	if err != nil {
		return nil, err
	}
	if _, err := pool.Exec(ctx, measureSchema); err != nil {
		return nil, fmt.Errorf("schema: %w", err)
	}
	return &Store{pool: pool}, nil
}

func (s *Store) Ping(ctx context.Context) error { return s.pool.Ping(ctx) }

func (s *Store) UpsertScenarios(ctx context.Context, m map[string]*Scenario) error {
	for _, sc := range m {
		spec, _ := json.Marshal(sc.InjectionSpec())
		_, err := s.pool.Exec(ctx, `INSERT INTO scenarios (scenario_id,path,asset,expected_layer,ground_truth_label,ground_truth_action,injection_spec)
			VALUES ($1,$2,$3,$4,$5,$6,$7) ON CONFLICT (scenario_id) DO UPDATE SET path=EXCLUDED.path, asset=EXCLUDED.asset,
			expected_layer=EXCLUDED.expected_layer, ground_truth_label=EXCLUDED.ground_truth_label,
			ground_truth_action=EXCLUDED.ground_truth_action, injection_spec=EXCLUDED.injection_spec`,
			sc.ID, sc.Path, sc.Asset, sc.ExpectedLayer, sc.GroundTruthLabel, sc.GroundTruthAction, spec)
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) CreateRun(ctx context.Context, id, campaign string, driftBefore *int, gitSHA, notes, model string) error {
	_, err := s.pool.Exec(ctx, `INSERT INTO runs (run_id,campaign,drift_before,git_sha,notes,model) VALUES ($1,$2,$3,$4,$5,$6)`, id, campaign, driftBefore, nz(gitSHA), nz(notes), nz(model))
	return err
}

func (s *Store) FinishRun(ctx context.Context, id, status string, driftAfter *int) error {
	_, err := s.pool.Exec(ctx, `UPDATE runs SET finished_at=now(), status=$2, drift_after=COALESCE($3,drift_after) WHERE run_id=$1`, id, status, driftAfter)
	return err
}

func (s *Store) InsertInjection(ctx context.Context, runID, scenarioID, method string, params map[string]any) (int64, error) {
	p, _ := json.Marshal(params)
	var id int64
	err := s.pool.QueryRow(ctx, `INSERT INTO injections (run_id,scenario_id,t_inject,method,params) VALUES ($1,$2,now(),$3,$4) RETURNING id`, runID, scenarioID, method, p).Scan(&id)
	return id, err
}

func (s *Store) ClearInjection(ctx context.Context, id int64) error {
	_, err := s.pool.Exec(ctx, `UPDATE injections SET t_clear=now() WHERE id=$1`, id)
	return err
}

func (s *Store) Event(ctx context.Context, runID, scenarioID *string, kind, actor string, payload any) {
	p, _ := json.Marshal(payload)
	if _, err := s.pool.Exec(ctx, `INSERT INTO events (run_id,scenario_id,kind,actor,payload) VALUES ($1,$2,$3,$4,$5)`, runID, scenarioID, kind, actor, p); err != nil {
		fmt.Println("event insert failed:", err)
	}
}

type EventRow struct {
	ID         int64           `json:"id"`
	RunID      *string         `json:"run_id"`
	ScenarioID *string         `json:"scenario_id"`
	TS         time.Time       `json:"ts"`
	Kind       string          `json:"kind"`
	Actor      string          `json:"actor"`
	Payload    json.RawMessage `json:"payload"`
}

func (s *Store) Events(ctx context.Context, runID, scenarioID string) ([]EventRow, error) {
	rows, err := s.pool.Query(ctx, `SELECT id,run_id,scenario_id,ts,kind,actor,payload FROM events WHERE run_id=$1 AND ($2='' OR scenario_id=$2) ORDER BY ts, id`, runID, scenarioID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []EventRow
	for rows.Next() {
		var e EventRow
		if err := rows.Scan(&e.ID, &e.RunID, &e.ScenarioID, &e.TS, &e.Kind, &e.Actor, &e.Payload); err != nil {
			return nil, err
		}
		out = append(out, e)
	}
	return out, rows.Err()
}

type Verdict struct {
	RunID          string   `json:"run_id"`
	ScenarioID     string   `json:"scenario_id"`
	LayerActual    *int     `json:"layer_actual"`
	Recovered      bool     `json:"recovered"`
	MTTDs          *float64 `json:"mttd_s"`
	MTTAs          *float64 `json:"mtta_s"`
	MTTRs          *float64 `json:"mttr_s"`
	DiagnosisScore *float64 `json:"diagnosis_score"`
	ActionMatch    *bool    `json:"action_match"`
	MisjudgeCount  int      `json:"misjudge_count"`
	Notes          string   `json:"notes"`
}

func (s *Store) InsertVerdict(ctx context.Context, v Verdict) error {
	_, err := s.pool.Exec(ctx, `INSERT INTO verdicts (run_id,scenario_id,layer_actual,recovered,mttd_s,mtta_s,mttr_s,diagnosis_score,action_match,misjudge_count,notes)
		VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11)`, v.RunID, v.ScenarioID, v.LayerActual, v.Recovered, v.MTTDs, v.MTTAs, v.MTTRs, v.DiagnosisScore, v.ActionMatch, v.MisjudgeCount, nz(v.Notes))
	return err
}

func (s *Store) Verdicts(ctx context.Context, runID string) ([]Verdict, error) {
	rows, err := s.pool.Query(ctx, `SELECT run_id,scenario_id,layer_actual,recovered,mttd_s,mtta_s,mttr_s,diagnosis_score,action_match,misjudge_count,COALESCE(notes,'') FROM verdicts WHERE run_id=$1 ORDER BY id`, runID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Verdict
	for rows.Next() {
		var v Verdict
		if err := rows.Scan(&v.RunID, &v.ScenarioID, &v.LayerActual, &v.Recovered, &v.MTTDs, &v.MTTAs, &v.MTTRs, &v.DiagnosisScore, &v.ActionMatch, &v.MisjudgeCount, &v.Notes); err != nil {
			return nil, err
		}
		out = append(out, v)
	}
	return out, rows.Err()
}

type RunRow struct {
	RunID       string     `json:"run_id"`
	Campaign    string     `json:"campaign"`
	StartedAt   time.Time  `json:"started_at"`
	FinishedAt  *time.Time `json:"finished_at"`
	DriftBefore *int       `json:"drift_before"`
	DriftAfter  *int       `json:"drift_after"`
	GitSHA      *string    `json:"git_sha"`
	Status      string     `json:"status"`
	Notes       *string    `json:"notes"`
	Model       *string    `json:"model"` // nil for runs before campaign-level model selection
}

func (s *Store) Runs(ctx context.Context, runID string) ([]RunRow, error) {
	rows, err := s.pool.Query(ctx, `SELECT run_id,campaign,started_at,finished_at,drift_before,drift_after,git_sha,status,notes,model FROM runs WHERE ($1='' OR run_id=$1) ORDER BY started_at DESC LIMIT 100`, runID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []RunRow
	for rows.Next() {
		var r RunRow
		if err := rows.Scan(&r.RunID, &r.Campaign, &r.StartedAt, &r.FinishedAt, &r.DriftBefore, &r.DriftAfter, &r.GitSHA, &r.Status, &r.Notes, &r.Model); err != nil {
			return nil, err
		}
		out = append(out, r)
	}
	return out, rows.Err()
}

type InjectionRow struct {
	ID         int64           `json:"id"`
	ScenarioID string          `json:"scenario_id"`
	TInject    *time.Time      `json:"t_inject"`
	TClear     *time.Time      `json:"t_clear"`
	Method     string          `json:"method"`
	Params     json.RawMessage `json:"params"`
}

func (s *Store) Injections(ctx context.Context, runID string) ([]InjectionRow, error) {
	rows, err := s.pool.Query(ctx, `SELECT id,scenario_id,t_inject,t_clear,method,params FROM injections WHERE run_id=$1 ORDER BY id`, runID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []InjectionRow
	for rows.Next() {
		var r InjectionRow
		if err := rows.Scan(&r.ID, &r.ScenarioID, &r.TInject, &r.TClear, &r.Method, &r.Params); err != nil {
			return nil, err
		}
		out = append(out, r)
	}
	return out, rows.Err()
}

// proposals (layer 2 approval tokens, design 6.5; the AI VM registers them, Discord approval marks them approved)
func (s *Store) InsertProposal(ctx context.Context, token string, runID, scenarioID *string, actionID string, target map[string]any, ttl time.Duration) error {
	t, _ := json.Marshal(target)
	_, err := s.pool.Exec(ctx, `INSERT INTO proposals (token,run_id,scenario_id,action_id,target,expires_at) VALUES ($1,$2,$3,$4,$5,now()+$6::interval)`,
		token, runID, scenarioID, actionID, t, fmt.Sprintf("%d seconds", int(ttl.Seconds())))
	return err
}

type ProposalRow struct {
	Token      string
	RunID      *string
	ScenarioID *string
	ActionID   string
	ApprovedAt *time.Time
	ApprovedBy *string
	ExecutedAt *time.Time
	ExpiresAt  time.Time
}

func (s *Store) GetProposal(ctx context.Context, token string) (*ProposalRow, error) {
	var p ProposalRow
	err := s.pool.QueryRow(ctx, `SELECT token,run_id,scenario_id,action_id,approved_at,approved_by,executed_at,expires_at FROM proposals WHERE token=$1`, token).
		Scan(&p.Token, &p.RunID, &p.ScenarioID, &p.ActionID, &p.ApprovedAt, &p.ApprovedBy, &p.ExecutedAt, &p.ExpiresAt)
	if err != nil {
		return nil, err
	}
	return &p, nil
}

func (s *Store) ApproveProposal(ctx context.Context, token, by string) error {
	_, err := s.pool.Exec(ctx, `UPDATE proposals SET approved_at=now(), approved_by=$2 WHERE token=$1 AND approved_at IS NULL`, token, by)
	return err
}

func (s *Store) MarkProposalExecuted(ctx context.Context, token string) error {
	_, err := s.pool.Exec(ctx, `UPDATE proposals SET executed_at=now() WHERE token=$1`, token)
	return err
}

func nz(s string) *string {
	if s == "" {
		return nil
	}
	return &s
}