selfheal-lab-apps/ai/cmd/agent/main.go

// agent: AI VM diagnosis service (experiment design 6). HTTP on :8080. /diagnose is called by the harness.
package main

import (
	"context"
	"encoding/json"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/agent"
	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/anthropic"
	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/diag"
	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/harness"
	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/jev"
	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/openai"
	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/rag"
	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/tools"
)

func env(k, d string) string {
	if v := os.Getenv(k); v != "" {
		return v
	}
	return d
}

func readFile(p string) string {
	b, err := os.ReadFile(p)
	if err != nil {
		log.Fatalf("read %s: %v", p, err)
	}
	return string(b)
}

func keyReader(path string) func() (string, error) {
	return func() (string, error) {
		b, err := os.ReadFile(path)
		if err != nil {
			return "", err
		}
		return strings.TrimSpace(string(b)), nil
	}
}

func main() {
	log.SetFlags(log.LstdFlags | log.Lmicroseconds)
	ctx := context.Background()
	promptDir := env("PROMPT_DIR", "/etc/ai/prompts")
	oai := openai.New(env("OPENAI_BASE_URL", ""), keyReader(env("OPENAI_KEY_FILE", "/run/secrets/openai-api-key")))
	store, err := rag.New(ctx, env("RAG_DSN", "postgres://ai:ai@rag-db:5432/ai?sslmode=disable"))
	if err != nil {
		log.Fatalf("rag: %v", err)
	}
	system := readFile(promptDir + "/system.md")
	pv, err := store.UpsertPrompt(ctx, system)
	if err != nil {
		log.Fatalf("prompt version: %v", err)
	}
	hc := harness.New(env("HARNESS_URL", "http://10.20.0.11:8080"))
	enums := diag.Enums{Labels: diag.DefaultLabels, Actions: diag.DefaultActions}
	if sc, err := hc.Scenarios(ctx); err == nil && len(sc) > 0 {
		seen := map[string]bool{}
		enums.Labels = nil
		for _, s := range sc {
			if s.Label != "" && !seen[s.Label] {
				seen[s.Label] = true
				enums.Labels = append(enums.Labels, s.Label)
			}
		}
	} else if err != nil {
		log.Printf("harness scenarios unavailable, using built-in labels: %v", err)
	}
	if cat, err := hc.Catalog(ctx); err == nil && len(cat) > 0 {
		enums.Actions = nil
		for _, c := range cat {
			if strings.Contains(c.Layer, "2") && c.ActionID != "escalate" {
				enums.Actions = append(enums.Actions, c.ActionID)
			}
		}
	} else if err != nil {
		log.Printf("harness catalog unavailable, using built-in actions: %v", err)
	}
	tcfg := tools.Config{
		MimirURL: env("MIMIR_URL", "http://10.20.0.10:9009/prometheus"), LokiURL: env("LOKI_URL", "http://10.20.0.10:3100"),
		K8sAPI: env("K8S_API", "https://10.10.0.10:6443"), K8sTokenFile: env("K8S_TOKEN_FILE", "/run/secrets/k8s-token"), K8sCAFile: env("K8S_CA_FILE", "/run/secrets/k8s-ca.crt"),
		AWSRoleARN: env("AWS_ROLE_ARN", ""), AWSRegion: env("AWS_REGION", "ap-northeast-1"), GCPProject: env("GCP_PROJECT", ""),
		EmbedModel: env("EMBED_MODEL", "text-embedding-3-small"),
	}
	tr := tools.New(tcfg, store, oai)
	acfg := agent.Config{
		Model: env("OPENAI_MODEL", "gpt-6.1-sol"), SystemPrompt: system,
		Topology: readFile(promptDir + "/topology.md"), SLO: readFile(promptDir + "/slo.md"),
		HitlURL: env("HITL_URL", "http://hitl:8081"),
	}
	ag := agent.New(acfg, oai, tr, hc, store, enums, pv)
	// Claude Messages API for claude-* model ids chosen per campaign; the key may be registered later (empty = diagnoses with
	// a claude-* model fail with "anthropic api key is empty", recorded in the harness diagnosis event).
	anth := anthropic.New(env("ANTHROPIC_BASE_URL", ""), keyReader(env("ANTHROPIC_KEY_FILE", "/run/secrets/anthropic-api-key")))
	// Workload Identity Federation (preferred, no static key): the four ids are plain settings in ai.env; the VM's
	// Google identity token is exchanged for a short-lived access token (internal/anthropic/federation.go).
	anth.SetFederation(&anthropic.Federation{
		RuleID: env("ANTHROPIC_FEDERATION_RULE_ID", ""), OrgID: env("ANTHROPIC_ORGANIZATION_ID", ""),
		ServiceAccountID: env("ANTHROPIC_SERVICE_ACCOUNT_ID", ""), WorkspaceID: env("ANTHROPIC_WORKSPACE_ID", ""),
		IdentityTokenURL: env("ANTHROPIC_IDENTITY_TOKEN_URL", ""),
	})
	ag.SetAnthropic(anth)
	log.Printf("anthropic auth: %s", anth.AuthMode())
	jevMode := env("JEV_MODE", "off") // off | shadow | enforce (shadow records Jev answers next to the diagnosis; enforce also applies the two review rules)
	var jc *jev.Client
	if jevMode == "shadow" || jevMode == "enforce" {
		jc = jev.New(env("JEV_BASE_URL", ""), env("JEV_MODEL", "jev-1.13.0"), keyReader(env("JEV_KEY_FILE", "/run/secrets/typesafe-api-key")))
		ag.EnableShadow(jc, jevMode == "enforce")
	}
	log.Printf("agent up: model=%s prompt=%s labels=%d actions=%d jev=%s", acfg.Model, pv, len(enums.Labels), len(enums.Actions), jevMode)
	if jc != nil && !jc.KeyPresent() {
		log.Printf("jev shadow: typesafe key file empty; gates record 'typesafe api key is empty' until the secret is registered")
	}

	// runbook ingest (needs the OpenAI key; retried in the background until the key exists)
	go func() {
		dir := env("RUNBOOK_DIR", "/etc/ai/runbooks")
		for {
			added, removed, err := store.IngestDir(ctx, dir, oai, tcfg.EmbedModel)
			if err == nil {
				n, _ := store.DocCount(ctx)
				log.Printf("runbooks: +%d -%d, %d chunks stored", added, removed, n)
				return
			}
			log.Printf("runbook ingest deferred: %v", err)
			time.Sleep(5 * time.Minute)
		}
	}()

	var inflight int32
	mux := http.NewServeMux()
	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, r *http.Request) { w.Write([]byte("ok\n")) })
	mux.HandleFunc("GET /selfcheck", func(w http.ResponseWriter, r *http.Request) {
		c, cancel := context.WithTimeout(r.Context(), 40*time.Second)
		defer cancel()
		out := map[string]any{"model": acfg.Model, "prompt_version": pv, "labels": enums.Labels, "actions": enums.Actions}
		key, kerr := keyReader(env("OPENAI_KEY_FILE", "/run/secrets/openai-api-key"))()
		out["openai_key_present"] = kerr == nil && key != ""
		out["anthropic_key_present"] = anth.KeyPresent()
		out["anthropic_auth"] = anth.AuthMode()
		if f := anth.Federation(); f != nil {
			out["anthropic_federation"] = f.Probe(c, anth.BaseURL) // one real token exchange; the Console history page shows the deny reason on failure
		}
		out["jev_mode"] = jevMode
		if jc != nil {
			out["jev_key_present"] = jc.KeyPresent()
			out["jev_model"] = jc.Model()
		}
		out["rag_db"] = errText(store.Ping(c))
		n, _ := store.DocCount(c)
		out["rag_chunks"] = n
		out["mimir"] = tr.Call(c, "get_probe_status", "{}")
		out["k8s"] = short(tr.Call(c, "get_k8s", `{"kind":"nodes","namespace":""}`))
		out["gcp"] = short(tr.Call(c, "get_gcp_instances", "{}"))
		arn, aerr := tr.AWSCallerIdentity(c)
		out["aws"] = map[string]any{"arn": arn, "error": errText(aerr)}
		if _, err := hc.Catalog(c); err != nil {
			out["harness"] = err.Error()
		} else {
			out["harness"] = "ok"
		}
		writeJSON(w, 200, out)
	})
	mux.HandleFunc("POST /diagnose", func(w http.ResponseWriter, r *http.Request) {
		var req agent.Request
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeJSON(w, 400, map[string]string{"error": err.Error()})
			return
		}
		if req.Alert == nil {
			writeJSON(w, 400, map[string]string{"error": "alert required"})
			return
		}
		inflight++
		defer func() { inflight-- }()
		c, cancel := context.WithTimeout(context.Background(), 170*time.Second) // harness waits 3 minutes
		defer cancel()
		prefetched := ag.Prefetch(c)
		tri := ag.ShadowTriage(c, req, prefetched) // nil unless JEV_MODE=shadow; recorded only
		res, err := ag.DiagnoseWith(c, req, prefetched)
		if err != nil {
			ag.SaveShadow(context.Background(), req, "", nil, tri)
			writeJSON(w, 503, map[string]string{"error": err.Error()})
			return
		}
		rev := ag.ShadowReview(c, req, res)
		if tri != nil || rev != nil {
			res.Jev = map[string]any{"mode": jevMode, "triage": tri, "review": rev}
		}
		ag.Enforce(res, rev) // no-op unless JEV_MODE=enforce
		ag.Dispose(c, req, res)
		ag.SaveShadow(context.Background(), req, res.SessionID, res, tri, rev)
		writeJSON(w, 200, res)
	})
	srv := &http.Server{Addr: env("AGENT_ADDR", ":8080"), Handler: logging(mux), ReadHeaderTimeout: 10 * time.Second}
	log.Fatal(srv.ListenAndServe())
}

func errText(err error) string {
	if err == nil {
		return "ok"
	}
	return err.Error()
}

func short(s string) string {
	if len(s) > 400 {
		return s[:400] + "..."
	}
	return s
}

func writeJSON(w http.ResponseWriter, code int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	enc := json.NewEncoder(w)
	enc.SetIndent("", "  ")
	_ = enc.Encode(v)
}

func logging(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/healthz" {
			log.Printf("%s %s from %s", r.Method, r.URL.Path, r.RemoteAddr)
		}
		h.ServeHTTP(w, r)
	})
}