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

package agent

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"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/openai"
	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/tools"
)

// A mock Responses API: turn 1 asks for get_probe_status, turn 2 submits the diagnosis. Verifies the loop wiring
// (function_call_output with the call_id, previous_response_id chaining, strict submit parsing).
func TestDiagnoseLoop(t *testing.T) {
	var turns int
	var sawOutput, sawPrev bool
	oaiSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/embeddings" {
			w.Write([]byte(`{"data":[{"index":0,"embedding":[0.1,0.2]}]}`))
			return
		}
		var req openai.Request
		_ = json.NewDecoder(r.Body).Decode(&req)
		turns++
		for _, it := range req.Input {
			if it.Type == "function_call_output" && it.CallID == "call_1" && strings.Contains(it.Output, "probe_success") {
				sawOutput = true
			}
		}
		if req.PreviousResponseID == "resp_1" {
			sawPrev = true
		}
		if r.Header.Get("Authorization") != "Bearer test-key" {
			t.Errorf("missing bearer")
		}
		switch turns {
		case 1:
			w.Write([]byte(`{"id":"resp_1","status":"completed","output":[{"type":"function_call","call_id":"call_1","name":"get_probe_status","arguments":"{}"}],"usage":{"input_tokens":10,"output_tokens":5}}`))
		default:
			args := `{\"root_cause_label\":\"db_primary_down\",\"confidence\":0.8,\"evidence\":[\"probe_success db-primary=0\"],\"proposed_action\":\"escalate\",\"target\":\"db-primary\",\"rationale\":\"primary unreachable\"}`
			w.Write([]byte(`{"id":"resp_2","status":"completed","output":[{"type":"function_call","call_id":"call_2","name":"submit_diagnosis","arguments":"` + args + `"}],"usage":{"input_tokens":20,"output_tokens":9}}`))
		}
	}))
	defer oaiSrv.Close()
	mimir := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, "/api/v1/query") {
			w.Write([]byte(`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"target":"db-primary"},"value":[1,"0"]},{"metric":{"target":"alb"},"value":[1,"1"]}]}}`))
			return
		}
		w.Write([]byte(`{"status":"success","data":{"resultType":"matrix","result":[]}}`))
	}))
	defer mimir.Close()
	oai := openai.New(oaiSrv.URL, func() (string, error) { return "test-key", nil })
	tr := tools.New(tools.Config{MimirURL: mimir.URL, LokiURL: mimir.URL, EmbedModel: "e"}, nil, oai)
	enums := diag.Enums{Labels: diag.DefaultLabels, Actions: diag.DefaultActions}
	ag := New(Config{Model: "m", SystemPrompt: "sys", Topology: "topo", SLO: "slo", HitlURL: "http://127.0.0.1:1", MaxTurns: 5}, oai, tr, harness.New("http://127.0.0.1:1"), nil, enums, "pv1")
	res, err := ag.Diagnose(context.Background(), Request{Alert: map[string]any{"alertname": "DbPrimaryDown"}})
	if err != nil {
		t.Fatal(err)
	}
	if res.RootCauseLabel != "db_primary_down" || res.Turns != 2 || !sawOutput || !sawPrev {
		t.Fatalf("unexpected result: %+v sawOutput=%v sawPrev=%v", res, sawOutput, sawPrev)
	}
	ag.Dispose(context.Background(), Request{Alert: map[string]any{}}, res)
	if res.Disposition != "escalated" {
		t.Fatalf("escalate action must escalate, got %s", res.Disposition)
	}
}

func TestTargetFor(t *testing.T) {
	m := targetFor("k8s.rollout_restart", "demo/api")
	if m["namespace"] != "demo" || m["deployment"] != "api" {
		t.Fatalf("%v", m)
	}
	if targetFor("mig.resize", "")["mig"] != "selfheal-legacy-mig" {
		t.Fatal("mig")
	}
}

// A claude-* model id routes the same loop through the Anthropic adapter; the request model wins over Config.Model
// and the summed usage is returned in the result.
func TestDiagnoseRoutesClaudeModel(t *testing.T) {
	var models []string
	anth := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var b map[string]any
		_ = json.NewDecoder(r.Body).Decode(&b)
		models = append(models, b["model"].(string))
		args := `{"root_cause_label":"db_primary_down","confidence":0.8,"evidence":["e"],"proposed_action":"escalate","target":"db-primary","rationale":"r"}`
		w.Write([]byte(`{"id":"msg_a","type":"message","model":"claude-sonnet-5-5","stop_reason":"tool_use","content":[{"type":"tool_use","id":"toolu_a","name":"submit_diagnosis","input":` + args + `}],"usage":{"input_tokens":40,"output_tokens":8}}`))
	}))
	defer anth.Close()
	oaiSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/embeddings" {
			w.Write([]byte(`{"data":[{"index":0,"embedding":[0.1,0.2]}]}`))
			return
		}
		t.Errorf("openai must not be called for a claude-* model: %s", r.URL.Path)
	}))
	defer oaiSrv.Close()
	mimir := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"status":"success","data":{"resultType":"vector","result":[]}}`))
	}))
	defer mimir.Close()
	oai := openai.New(oaiSrv.URL, func() (string, error) { return "test-key", nil })
	tr := tools.New(tools.Config{MimirURL: mimir.URL, LokiURL: mimir.URL, EmbedModel: "e"}, nil, oai)
	enums := diag.Enums{Labels: diag.DefaultLabels, Actions: diag.DefaultActions}
	ag := New(Config{Model: "gpt-default", SystemPrompt: "sys", Topology: "topo", SLO: "slo", HitlURL: "http://127.0.0.1:1", MaxTurns: 3}, oai, tr, harness.New("http://127.0.0.1:1"), nil, enums, "pv1")
	req := Request{Alert: map[string]any{"alertname": "DbPrimaryDown"}, Model: "claude-sonnet-5-5"}
	if _, err := ag.Diagnose(context.Background(), req); err == nil || !strings.Contains(err.Error(), "anthropic provider not configured") {
		t.Fatalf("without SetAnthropic a claude-* model must fail clearly, got %v", err)
	}
	ag.SetAnthropic(anthropic.New(anth.URL, func() (string, error) { return "k", nil }))
	res, err := ag.Diagnose(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}
	if res.Model != "claude-sonnet-5-5" || len(models) != 1 || models[0] != "claude-sonnet-5-5" || res.RootCauseLabel != "db_primary_down" {
		t.Fatalf("routing: res=%+v models=%v", res, models)
	}
	if res.Usage["input_tokens"] != 40 || res.Usage["output_tokens"] != 8 || res.Usage["total_tokens"] != 48 || res.Usage["turns"] != 1 {
		t.Fatalf("usage: %v", res.Usage)
	}
}