selfheal-lab-apps/ai/internal/agent/agent.go

// Package agent runs one diagnosis (experiment design 6.3/6.4): a Responses API tool loop over the read-only tools,
// ending in a strict submit_diagnosis call. After the diagnosis it registers a layer-2 proposal with the harness and
// hands it to the HITL bot, or sends the case to the escalation channel (design 6.5/6.6).
package agent

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

	"github.com/google/uuid"

	"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"
)

type Config struct {
	Model           string
	MaxTurns        int
	MaxOutputTokens int
	SystemPrompt    string // body of prompts/system.md
	Topology        string // body of prompts/topology.md (design 1 IP table)
	SLO             string // body of prompts/slo.md (design 4.2)
	HitlURL         string // http://hitl:8081
	ProposalTTL     time.Duration
	ConfidenceMin   float64 // 0.5 (design 6.3)
}

type Agent struct {
	cfg        Config
	oai        *openai.Client
	tools      *tools.Runner
	harness    *harness.Client
	store      *rag.Store
	enums      diag.Enums
	promptV    string
	http       *http.Client
	jev        *jev.Client // nil = shadow gates off
	jevEnforce bool        // JEV_MODE=enforce: review rules may downgrade a proposal to escalate
}

func New(cfg Config, oai *openai.Client, tr *tools.Runner, hc *harness.Client, store *rag.Store, enums diag.Enums, promptVersion string) *Agent {
	if cfg.MaxTurns == 0 {
		cfg.MaxTurns = 12
	}
	if cfg.MaxOutputTokens == 0 {
		cfg.MaxOutputTokens = 4000
	}
	if cfg.ProposalTTL == 0 {
		cfg.ProposalTTL = 15 * time.Minute
	}
	if cfg.ConfidenceMin == 0 {
		cfg.ConfidenceMin = 0.5
	}
	return &Agent{cfg: cfg, oai: oai, tools: tr, harness: hc, store: store, enums: enums, promptV: promptVersion, http: &http.Client{Timeout: 30 * time.Second}}
}

// EnableShadow turns the Jev gates on (JEV_MODE=shadow records only; enforce also applies the review rules).
func (a *Agent) EnableShadow(c *jev.Client, enforce bool) { a.jev, a.jevEnforce = c, enforce }

// Request is the harness -> AI VM contract (harness alerts.go diagnose()).
type Request struct {
	Alert      map[string]any `json:"alert"`
	RunID      *string        `json:"run_id"`
	ScenarioID *string        `json:"scenario_id"`
}

// Result is returned to the harness: the diagnosis fields at top level (the harness stores the whole object as the
// diagnosis event payload and reads root_cause_label / proposed_action from it) plus bookkeeping.
type Result struct {
	diag.Diagnosis
	SessionID     string         `json:"session_id"`
	Model         string         `json:"model"`
	PromptVersion string         `json:"prompt_version"`
	Turns         int            `json:"turns"`
	ToolCalls     int            `json:"tool_calls"`
	DurationMs    int64          `json:"duration_ms"`
	Disposition   string         `json:"disposition"` // proposal | escalated | none
	ProposalToken string         `json:"proposal_token,omitempty"`
	Jev           map[string]any `json:"jev,omitempty"` // shadow gate answers (recorded, never acted on)
}

type toolCallLog struct {
	Name   string `json:"name"`
	Args   string `json:"args"`
	Output string `json:"output"`
	Ms     int64  `json:"ms"`
}

func (a *Agent) toolDefs() []openai.Tool {
	defs := tools.Definitions()
	defs = append(defs, openai.Tool{Type: "function", Strict: true, Name: "submit_diagnosis",
		Description: "Submit the final diagnosis. Call exactly once, after gathering evidence. This records the result; it does not execute anything.",
		Parameters:  a.enums.Schema()})
	return defs
}

func (a *Agent) userMessage(ctx context.Context, req Request, prefetched string) string {
	if prefetched == "" {
		prefetched = a.Prefetch(ctx)
	}
	alert, _ := json.MarshalIndent(req.Alert, "", "  ")
	var b strings.Builder
	b.WriteString("## Topology\n")
	b.WriteString(a.cfg.Topology)
	b.WriteString("\n\n## SLO\n")
	b.WriteString(a.cfg.SLO)
	b.WriteString("\n\n## Alert payload\n```json\n")
	b.Write(alert)
	b.WriteString("\n```\n\n## Last 10 minutes (pre-fetched)\n")
	b.WriteString(prefetched)
	if req.RunID != nil {
		fmt.Fprintf(&b, "\nrun_id=%s", *req.RunID)
	}
	if req.ScenarioID != nil {
		fmt.Fprintf(&b, " scenario_id=%s", *req.ScenarioID)
	}
	b.WriteString("\n\nInvestigate with the tools, then call submit_diagnosis exactly once.")
	return b.String()
}

// Prefetch returns the four pre-fetched 10-minute series (probe, p95, error ratio, legacy up) as markdown lines.
func (a *Agent) Prefetch(ctx context.Context) string {
	var b strings.Builder
	for _, q := range []struct{ name, promql string }{
		{"probe_success", `probe_success`},
		{"api_p95_latency_s", `histogram_quantile(0.95, sum by (le) (rate(demo_http_request_duration_seconds_bucket[5m])))`},
		{"api_error_ratio", `sum(rate(demo_http_requests_total{status=~"5.."}[5m])) / sum(rate(demo_http_requests_total[5m]))`},
		{"legacy_up", `up{job="legacy-api"}`},
	} {
		out := a.tools.Call(ctx, "query_metrics", fmt.Sprintf(`{"promql":%q,"range_minutes":10}`, q.promql))
		if len(out) > 1200 {
			out = out[:1200] + "..."
		}
		fmt.Fprintf(&b, "- %s: %s\n", q.name, out)
	}
	return b.String()
}

// Diagnose runs the tool loop and returns the validated diagnosis. The session is stored whatever the outcome.
func (a *Agent) Diagnose(ctx context.Context, req Request) (*Result, error) {
	return a.DiagnoseWith(ctx, req, "")
}

// DiagnoseWith runs Diagnose with an already pre-fetched sample block (empty = fetch now).
func (a *Agent) DiagnoseWith(ctx context.Context, req Request, prefetched string) (*Result, error) {
	start := time.Now()
	sid := "ses-" + uuid.NewString()[:8]
	res := &Result{SessionID: sid, Model: a.cfg.Model, PromptVersion: a.promptV}
	var calls []toolCallLog
	input := []openai.Item{
		{Type: "message", Role: "system", Content: a.cfg.SystemPrompt},
		{Type: "message", Role: "user", Content: a.userMessage(ctx, req, prefetched)},
	}
	session := rag.Session{SessionID: sid, RunID: req.RunID, ScenarioID: req.ScenarioID, Model: a.cfg.Model, PromptVersion: a.promptV, Input: map[string]any{"alert": req.Alert, "user_message": input[1].Content}}
	finish := func(err error) {
		session.Output = res
		session.ToolCalls = calls
		session.Duration = time.Since(start)
		if err != nil {
			session.Error = err.Error()
		}
		if a.store != nil {
			if e := a.store.SaveSession(context.Background(), session); e != nil {
				log.Printf("session save: %v", e)
			}
		}
	}

	defs := a.toolDefs()
	var prevID string
	var diagnosis *diag.Diagnosis
	pf := false
	// usage is summed over every turn: with previous_response_id each turn re-bills the whole context, so the
	// last response alone understates the cost (seen 2026-10-08: 12 turns, last turn 11K input tokens).
	usage := map[string]int{"input_tokens": 0, "output_tokens": 0, "total_tokens": 0, "turns": 0}
	for turn := 1; turn <= a.cfg.MaxTurns && diagnosis == nil; turn++ {
		res.Turns = turn
		r := openai.Request{Model: a.cfg.Model, Input: input, Tools: defs, ToolChoice: "auto", MaxOutputTokens: a.cfg.MaxOutputTokens, ParallelToolCalls: &pf}
		if prevID != "" {
			r.PreviousResponseID = prevID
		}
		if turn == a.cfg.MaxTurns { // last chance: force the final call
			r.ToolChoice = map[string]any{"type": "function", "name": "submit_diagnosis"}
		}
		resp, err := a.oai.CreateResponse(ctx, r)
		if err != nil {
			finish(err)
			return nil, err
		}
		usage["input_tokens"] += resp.Usage.InputTokens
		usage["output_tokens"] += resp.Usage.OutputTokens
		usage["total_tokens"] += resp.Usage.TotalTokens
		usage["turns"]++
		prevID = resp.ID
		fcs := resp.FunctionCalls()
		if len(fcs) == 0 {
			// text only: nudge once towards the tool
			input = []openai.Item{{Type: "message", Role: "user", Content: "Use the tools to gather evidence and finish with submit_diagnosis."}}
			continue
		}
		input = input[:0]
		for _, fc := range fcs {
			if fc.Name == "submit_diagnosis" {
				d, err := diag.Parse(fc.Arguments, a.enums)
				if err != nil {
					input = append(input, openai.Item{Type: "function_call_output", CallID: fc.CallID, Output: fmt.Sprintf(`{"error":%q}`, err.Error())})
					calls = append(calls, toolCallLog{Name: fc.Name, Args: fc.Arguments, Output: "rejected: " + err.Error()})
					continue
				}
				diagnosis = &d
				calls = append(calls, toolCallLog{Name: fc.Name, Args: fc.Arguments, Output: "accepted"})
				break
			}
			t0 := time.Now()
			out := a.tools.Call(ctx, fc.Name, fc.Arguments)
			calls = append(calls, toolCallLog{Name: fc.Name, Args: fc.Arguments, Output: out, Ms: time.Since(t0).Milliseconds()})
			input = append(input, openai.Item{Type: "function_call_output", CallID: fc.CallID, Output: out})
		}
	}
	session.Usage = usage
	res.ToolCalls = len(calls)
	res.DurationMs = time.Since(start).Milliseconds()
	if diagnosis == nil {
		err := fmt.Errorf("no valid submit_diagnosis after %d turns", res.Turns)
		finish(err)
		return nil, err
	}
	res.Diagnosis = *diagnosis
	finish(nil)
	return res, nil
}

// Dispose applies design 6.3/6.5/6.6 after the diagnosis: proposal + HITL message, or escalation, or nothing.
func (a *Agent) Dispose(ctx context.Context, req Request, res *Result) {
	d := res.Diagnosis
	diagMap := map[string]any{"root_cause_label": d.RootCauseLabel, "confidence": d.Confidence, "evidence": d.Evidence, "proposed_action": d.ProposedAction, "target": d.Target, "rationale": d.Rationale, "session_id": res.SessionID, "model": res.Model, "prompt_version": res.PromptVersion}
	switch {
	case d.ProposedAction == "none":
		res.Disposition = "none"
		return
	case d.ProposedAction == "escalate" || d.Confidence < a.cfg.ConfidenceMin:
		res.Disposition = "escalated"
		a.escalate(ctx, req, diagMap)
		return
	}
	target := targetFor(d.ProposedAction, d.Target)
	p, err := a.harness.Propose(ctx, req.RunID, req.ScenarioID, d.ProposedAction, target, diagMap, a.cfg.ProposalTTL)
	if err != nil {
		log.Printf("propose: %v", err)
		res.Disposition = "escalated"
		diagMap["note"] = "proposal registration failed: " + err.Error()
		a.escalate(ctx, req, diagMap)
		return
	}
	res.Disposition = "proposal"
	res.ProposalToken = p.Token
	body := map[string]any{"token": p.Token, "action_id": d.ProposedAction, "target": target, "diagnosis": diagMap, "run_id": req.RunID, "scenario_id": req.ScenarioID, "session_id": res.SessionID, "expires_in_s": p.ExpiresIn}
	if err := a.postJSON(ctx, a.cfg.HitlURL+"/post", body); err != nil {
		log.Printf("hitl post: %v", err)
	}
}

func (a *Agent) escalate(ctx context.Context, req Request, diagMap map[string]any) {
	runbook := ""
	if hits, err := a.tools.SearchRunbooks(ctx, fmt.Sprint(diagMap["root_cause_label"])+" "+fmt.Sprint(diagMap["rationale"]), 2); err == nil {
		if hs, ok := hits.([]rag.Hit); ok {
			var parts []string
			for _, h := range hs {
				parts = append(parts, "["+h.Source+"]\n"+h.Chunk)
			}
			runbook = strings.Join(parts, "\n\n")
		}
	}
	body := map[string]any{"diagnosis": diagMap, "runbook": runbook, "run_id": req.RunID, "scenario_id": req.ScenarioID, "alert": req.Alert}
	if err := a.postJSON(ctx, a.cfg.HitlURL+"/escalate", body); err != nil {
		log.Printf("hitl escalate: %v", err)
	}
}

func (a *Agent) postJSON(ctx context.Context, u string, body any) error {
	b, _ := json.Marshal(body)
	req, _ := http.NewRequestWithContext(ctx, "POST", u, bytes.NewReader(b))
	req.Header.Set("Content-Type", "application/json")
	resp, err := a.http.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		return fmt.Errorf("%s: status %d", u, resp.StatusCode)
	}
	return nil
}

// targetFor maps the free-text target of the diagnosis onto the harness action target keys (harness actions.go).
func targetFor(actionID, target string) map[string]any {
	t := strings.TrimSpace(target)
	switch {
	case strings.HasPrefix(actionID, "k8s.rollout"), actionID == "k8s.scale":
		m := map[string]any{"namespace": "demo", "deployment": t}
		if actionID == "k8s.scale" {
			m["replicas"] = 3
		}
		if i := strings.Index(t, "/"); i > 0 {
			m["namespace"], m["deployment"] = t[:i], t[i+1:]
		}
		return m
	case actionID == "k8s.node_drain", actionID == "k8s.node_uncordon":
		return map[string]any{"node": t}
	case actionID == "gce.instance_reset":
		return map[string]any{"instance": t}
	case actionID == "mig.recreate_instance":
		return map[string]any{"instance": t}
	case actionID == "mig.resize":
		return map[string]any{"mig": "selfheal-legacy-mig", "size": 2}
	case actionID == "db.replay_resume":
		return map[string]any{"host": "db-replica"}
	}
	return map[string]any{"target": t}
}