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

// Package tools implements the read-only allow-list of experiment design 6.2. Every tool returns a compact JSON
// string for the model; nothing here writes to any system.
package tools

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/service/ec2"
	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
	elbv2 "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"google.golang.org/api/compute/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"

	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/openai"
	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/rag"
)

type Config struct {
	MimirURL     string // http://10.20.0.10:9009/prometheus
	LokiURL      string // http://10.20.0.10:3100
	K8sAPI       string // https://10.10.0.10:6443
	K8sTokenFile string
	K8sCAFile    string
	AWSRoleARN   string
	AWSRegion    string
	GCPProject   string
	EmbedModel   string
	MaxOutput    int // bytes per tool result handed to the model
}

type Runner struct {
	cfg  Config
	http *http.Client
	rag  *rag.Store
	oai  *openai.Client
}

func New(cfg Config, store *rag.Store, oai *openai.Client) *Runner {
	if cfg.MaxOutput == 0 {
		cfg.MaxOutput = 6000
	}
	return &Runner{cfg: cfg, http: &http.Client{Timeout: 30 * time.Second}, rag: store, oai: oai}
}

// Definitions are the function tools exposed to the model (strict schemas, every property required).
func Definitions() []openai.Tool {
	obj := func(props map[string]any, req ...string) map[string]any {
		if req == nil {
			req = []string{}
		}
		return map[string]any{"type": "object", "properties": props, "required": req, "additionalProperties": false}
	}
	return []openai.Tool{
		{Type: "function", Strict: true, Name: "query_metrics", Description: "Run a PromQL range query against Mimir (Prometheus API) over the last range_minutes. Metrics: probe_success{target=alb|db-primary}, demo_http_requests_total, demo_http_request_duration_seconds_bucket, up{job=legacy-api}, kube_pod_container_status_restarts_total, pg_replication_lag_seconds, pg_is_wal_replay_paused, argocd_app_info, harness_*.",
			Parameters: obj(map[string]any{"promql": map[string]any{"type": "string"}, "range_minutes": map[string]any{"type": "integer", "description": "1 to 120"}}, "promql", "range_minutes")},
		{Type: "function", Strict: true, Name: "query_logs", Description: "Run a LogQL query against Loki over the last range_minutes, newest first. Streams: {job=\"demo-api\"}, {job=\"legacy-api\"}, {namespace=\"demo\"}.",
			Parameters: obj(map[string]any{"logql": map[string]any{"type": "string"}, "range_minutes": map[string]any{"type": "integer"}, "limit": map[string]any{"type": "integer", "description": "1 to 100"}}, "logql", "range_minutes", "limit")},
		{Type: "function", Strict: true, Name: "get_probe_status", Description: "Latest synthetic probe results (probe_success per target: alb = demo-app through the ALB, db-primary = PostgreSQL TCP).",
			Parameters: obj(map[string]any{})},
		{Type: "function", Strict: true, Name: "get_k8s", Description: "List Kubernetes objects in a namespace (read-only). kind: pods | events | deployments | replicasets | nodes. namespace is ignored for nodes.",
			Parameters: obj(map[string]any{"kind": map[string]any{"type": "string", "enum": []string{"pods", "events", "deployments", "replicasets", "nodes"}}, "namespace": map[string]any{"type": "string"}}, "kind", "namespace")},
		{Type: "function", Strict: true, Name: "get_gcp_instances", Description: "GCE instances (selfheal-* pet VMs and legacy-* MIG instances) with status, zone, internal IP, plus the legacy MIG target size and instance health.",
			Parameters: obj(map[string]any{})},
		{Type: "function", Strict: true, Name: "get_aws_state", Description: "AWS state: EC2 instances tagged selfheal-* (state, private IP) and ALB target group health per worker.",
			Parameters: obj(map[string]any{})},
		{Type: "function", Strict: true, Name: "search_runbooks", Description: "Semantic search over the runbooks (one per catalogue action), the topology document and past verdict summaries. Returns the k closest chunks.",
			Parameters: obj(map[string]any{"query": map[string]any{"type": "string"}, "k": map[string]any{"type": "integer", "description": "1 to 8"}}, "query", "k")},
	}
}

// Call executes one tool by name with JSON arguments. Errors are returned as JSON so the model can react.
func (r *Runner) Call(ctx context.Context, name, args string) string {
	var a map[string]any
	if err := json.Unmarshal([]byte(args), &a); err != nil {
		return errJSON(fmt.Errorf("arguments: %w", err))
	}
	var (
		out any
		err error
	)
	switch name {
	case "query_metrics":
		out, err = r.queryMetrics(ctx, str(a, "promql"), num(a, "range_minutes"))
	case "query_logs":
		out, err = r.queryLogs(ctx, str(a, "logql"), num(a, "range_minutes"), num(a, "limit"))
	case "get_probe_status":
		out, err = r.probeStatus(ctx)
	case "get_k8s":
		out, err = r.getK8s(ctx, str(a, "kind"), str(a, "namespace"))
	case "get_gcp_instances":
		out, err = r.gcpInstances(ctx)
	case "get_aws_state":
		out, err = r.awsState(ctx)
	case "search_runbooks":
		out, err = r.SearchRunbooks(ctx, str(a, "query"), int(num(a, "k")))
	default:
		err = fmt.Errorf("unknown tool %s", name)
	}
	if err != nil {
		return errJSON(err)
	}
	b, _ := json.Marshal(out)
	if len(b) > r.cfg.MaxOutput {
		return string(b[:r.cfg.MaxOutput]) + `..."truncated":true}`
	}
	return string(b)
}

func errJSON(err error) string {
	b, _ := json.Marshal(map[string]string{"error": err.Error()})
	return string(b)
}

// ---- Mimir / Loki ----

type promResp struct {
	Status string `json:"status"`
	Data   struct {
		ResultType string `json:"resultType"`
		Result     []struct {
			Metric map[string]string `json:"metric"`
			Values [][]any           `json:"values"`
			Value  []any             `json:"value"`
		} `json:"result"`
	} `json:"data"`
	Error string `json:"error"`
}

func (r *Runner) getJSON(ctx context.Context, u string, out any) error {
	req, _ := http.NewRequestWithContext(ctx, "GET", u, nil)
	resp, err := r.http.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
	if resp.StatusCode != 200 {
		return fmt.Errorf("status %d: %s", resp.StatusCode, truncate(string(b), 300))
	}
	return json.Unmarshal(b, out)
}

func clampMinutes(m int64) int64 {
	if m <= 0 {
		return 10
	}
	if m > 120 {
		return 120
	}
	return m
}

func (r *Runner) queryMetrics(ctx context.Context, promql string, minutes int64) (any, error) {
	minutes = clampMinutes(minutes)
	end := time.Now()
	start := end.Add(-time.Duration(minutes) * time.Minute)
	step := minutes * 60 / 20
	if step < 15 {
		step = 15
	}
	q := url.Values{"query": {promql}, "start": {strconv.FormatInt(start.Unix(), 10)}, "end": {strconv.FormatInt(end.Unix(), 10)}, "step": {strconv.FormatInt(step, 10)}}
	var pr promResp
	if err := r.getJSON(ctx, strings.TrimRight(r.cfg.MimirURL, "/")+"/api/v1/query_range?"+q.Encode(), &pr); err != nil {
		return nil, err
	}
	if pr.Status != "success" {
		return nil, fmt.Errorf("mimir: %s", pr.Error)
	}
	type series struct {
		Metric map[string]string `json:"metric"`
		Points []string          `json:"points"` // "t=HH:MM:SS v=..."
	}
	var out []series
	for i, s := range pr.Data.Result {
		if i >= 20 {
			break
		}
		var pts []string
		vals := s.Values
		if len(vals) > 12 { // keep the last 12 points
			vals = vals[len(vals)-12:]
		}
		for _, v := range vals {
			if len(v) == 2 {
				ts, _ := v[0].(float64)
				pts = append(pts, fmt.Sprintf("t=%s v=%v", time.Unix(int64(ts), 0).UTC().Format("15:04:05"), v[1]))
			}
		}
		out = append(out, series{Metric: s.Metric, Points: pts})
	}
	return map[string]any{"series_count": len(pr.Data.Result), "series": out, "range_minutes": minutes}, nil
}

func (r *Runner) probeStatus(ctx context.Context) (any, error) {
	q := url.Values{"query": {"probe_success"}}
	var pr promResp
	if err := r.getJSON(ctx, strings.TrimRight(r.cfg.MimirURL, "/")+"/api/v1/query?"+q.Encode(), &pr); err != nil {
		return nil, err
	}
	out := map[string]any{}
	for _, s := range pr.Data.Result {
		if len(s.Value) == 2 {
			out[s.Metric["target"]] = s.Value[1]
		}
	}
	return map[string]any{"probe_success": out, "checked_at": time.Now().UTC().Format(time.RFC3339)}, nil
}

type lokiResp struct {
	Status string `json:"status"`
	Data   struct {
		Result []struct {
			Stream map[string]string `json:"stream"`
			Values [][]string        `json:"values"`
		} `json:"result"`
	} `json:"data"`
}

func (r *Runner) queryLogs(ctx context.Context, logql string, minutes, limit int64) (any, error) {
	minutes = clampMinutes(minutes)
	if limit <= 0 || limit > 100 {
		limit = 30
	}
	end := time.Now()
	start := end.Add(-time.Duration(minutes) * time.Minute)
	q := url.Values{"query": {logql}, "start": {strconv.FormatInt(start.UnixNano(), 10)}, "end": {strconv.FormatInt(end.UnixNano(), 10)}, "limit": {strconv.FormatInt(limit, 10)}, "direction": {"backward"}}
	var lr lokiResp
	if err := r.getJSON(ctx, strings.TrimRight(r.cfg.LokiURL, "/")+"/loki/api/v1/query_range?"+q.Encode(), &lr); err != nil {
		return nil, err
	}
	var lines []string
	for _, s := range lr.Data.Result {
		for _, v := range s.Values {
			if len(v) == 2 {
				ns, _ := strconv.ParseInt(v[0], 10, 64)
				lines = append(lines, fmt.Sprintf("%s %s %s", time.Unix(0, ns).UTC().Format("15:04:05"), labelsOf(s.Stream), truncate(v[1], 300)))
			}
		}
	}
	sort.Sort(sort.Reverse(sort.StringSlice(lines)))
	if int64(len(lines)) > limit {
		lines = lines[:limit]
	}
	return map[string]any{"lines": lines, "count": len(lines)}, nil
}

func labelsOf(m map[string]string) string {
	keys := make([]string, 0, len(m))
	for k := range m {
		if k == "job" || k == "pod" || k == "instance" || k == "container" || k == "level" {
			keys = append(keys, k)
		}
	}
	sort.Strings(keys)
	parts := make([]string, 0, len(keys))
	for _, k := range keys {
		parts = append(parts, k+"="+m[k])
	}
	return "{" + strings.Join(parts, ",") + "}"
}

// ---- Kubernetes (SA ai-reader, ClusterRole view) ----

func (r *Runner) k8s() (kubernetes.Interface, error) {
	tok, err := os.ReadFile(r.cfg.K8sTokenFile)
	if err != nil {
		return nil, fmt.Errorf("k8s token: %w", err)
	}
	ca, err := os.ReadFile(r.cfg.K8sCAFile)
	if err != nil {
		return nil, fmt.Errorf("k8s ca: %w", err)
	}
	rc := &rest.Config{Host: r.cfg.K8sAPI, BearerToken: strings.TrimSpace(string(tok)), TLSClientConfig: rest.TLSClientConfig{CAData: ca}, Timeout: 20 * time.Second}
	return kubernetes.NewForConfig(rc)
}

func (r *Runner) getK8s(ctx context.Context, kind, ns string) (any, error) {
	cs, err := r.k8s()
	if err != nil {
		return nil, err
	}
	if ns == "" {
		ns = "demo"
	}
	switch kind {
	case "pods":
		l, err := cs.CoreV1().Pods(ns).List(ctx, metav1.ListOptions{})
		if err != nil {
			return nil, err
		}
		var out []map[string]any
		for _, p := range l.Items {
			restarts := int32(0)
			ready := 0
			var states []string
			for _, c := range p.Status.ContainerStatuses {
				restarts += c.RestartCount
				if c.Ready {
					ready++
				}
				switch {
				case c.State.Waiting != nil:
					states = append(states, c.Name+":"+c.State.Waiting.Reason)
				case c.State.Terminated != nil:
					states = append(states, c.Name+":"+c.State.Terminated.Reason)
				}
			}
			out = append(out, map[string]any{"name": p.Name, "node": p.Spec.NodeName, "phase": string(p.Status.Phase), "ready": fmt.Sprintf("%d/%d", ready, len(p.Spec.Containers)), "restarts": restarts, "states": states, "age_s": int(time.Since(p.CreationTimestamp.Time).Seconds())})
		}
		return out, nil
	case "events":
		l, err := cs.CoreV1().Events(ns).List(ctx, metav1.ListOptions{})
		if err != nil {
			return nil, err
		}
		items := l.Items
		sort.Slice(items, func(i, j int) bool { return items[i].LastTimestamp.After(items[j].LastTimestamp.Time) })
		var out []string
		for i, e := range items {
			if i >= 30 {
				break
			}
			out = append(out, fmt.Sprintf("%s %s %s/%s %s: %s (x%d)", e.LastTimestamp.UTC().Format("15:04:05"), e.Type, e.InvolvedObject.Kind, e.InvolvedObject.Name, e.Reason, truncate(e.Message, 160), e.Count))
		}
		return out, nil
	case "deployments":
		l, err := cs.AppsV1().Deployments(ns).List(ctx, metav1.ListOptions{})
		if err != nil {
			return nil, err
		}
		var out []map[string]any
		for _, d := range l.Items {
			img := ""
			if len(d.Spec.Template.Spec.Containers) > 0 {
				img = d.Spec.Template.Spec.Containers[0].Image
			}
			out = append(out, map[string]any{"name": d.Name, "desired": *d.Spec.Replicas, "ready": d.Status.ReadyReplicas, "updated": d.Status.UpdatedReplicas, "available": d.Status.AvailableReplicas, "image": img, "generation": d.Generation, "observed": d.Status.ObservedGeneration})
		}
		return out, nil
	case "replicasets":
		l, err := cs.AppsV1().ReplicaSets(ns).List(ctx, metav1.ListOptions{})
		if err != nil {
			return nil, err
		}
		var out []map[string]any
		for _, rs := range l.Items {
			out = append(out, map[string]any{"name": rs.Name, "desired": *rs.Spec.Replicas, "ready": rs.Status.ReadyReplicas, "revision": rs.Annotations["deployment.kubernetes.io/revision"]})
		}
		return out, nil
	case "nodes":
		l, err := cs.CoreV1().Nodes().List(ctx, metav1.ListOptions{})
		if err != nil {
			return nil, err
		}
		var out []map[string]any
		for _, n := range l.Items {
			ready := "Unknown"
			for _, c := range n.Status.Conditions {
				if c.Type == "Ready" {
					ready = string(c.Status)
				}
			}
			out = append(out, map[string]any{"name": n.Name, "ready": ready, "unschedulable": n.Spec.Unschedulable, "taints": len(n.Spec.Taints)})
		}
		return out, nil
	}
	return nil, fmt.Errorf("kind %q not allowed", kind)
}

// ---- GCP (VM SA, roles/compute.viewer) ----

func (r *Runner) gcpInstances(ctx context.Context) (any, error) {
	svc, err := compute.NewService(ctx)
	if err != nil {
		return nil, err
	}
	var vms []map[string]any
	err = svc.Instances.AggregatedList(r.cfg.GCPProject).Filter("name eq ^(selfheal-|legacy).*").Pages(ctx, func(l *compute.InstanceAggregatedList) error {
		for _, s := range l.Items {
			for _, in := range s.Instances {
				ip := ""
				if len(in.NetworkInterfaces) > 0 {
					ip = in.NetworkInterfaces[0].NetworkIP
				}
				vms = append(vms, map[string]any{"name": in.Name, "zone": in.Zone[strings.LastIndex(in.Zone, "/")+1:], "status": in.Status, "ip": ip})
			}
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	sort.Slice(vms, func(i, j int) bool { return vms[i]["name"].(string) < vms[j]["name"].(string) })
	var migs []map[string]any
	_ = svc.RegionInstanceGroupManagers.List(r.cfg.GCPProject, "asia-northeast1").Pages(ctx, func(l *compute.RegionInstanceGroupManagerList) error {
		for _, m := range l.Items {
			entry := map[string]any{"name": m.Name, "target_size": m.TargetSize}
			if h, err := svc.RegionInstanceGroupManagers.ListManagedInstances(r.cfg.GCPProject, "asia-northeast1", m.Name).Context(ctx).Do(); err == nil {
				var inst []string
				for _, mi := range h.ManagedInstances {
					hs := "UNKNOWN"
					if len(mi.InstanceHealth) > 0 {
						hs = mi.InstanceHealth[0].DetailedHealthState
					}
					inst = append(inst, fmt.Sprintf("%s status=%s health=%s action=%s", mi.Instance[strings.LastIndex(mi.Instance, "/")+1:], mi.InstanceStatus, hs, mi.CurrentAction))
				}
				entry["instances"] = inst
			}
			migs = append(migs, entry)
		}
		return nil
	})
	return map[string]any{"instances": vms, "migs": migs}, nil
}

// ---- AWS (GCE ID token -> AssumeRoleWithWebIdentity, role selfheal-ai-readonly) ----

func gceIDToken(ctx context.Context, audience string) (string, error) {
	u := "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity?format=full&audience=" + url.QueryEscape(audience)
	req, _ := http.NewRequestWithContext(ctx, "GET", u, nil)
	req.Header.Set("Metadata-Flavor", "Google")
	resp, err := (&http.Client{Timeout: 10 * time.Second}).Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("metadata identity: %s", resp.Status)
	}
	return string(b), nil
}

func (r *Runner) awsCfg(ctx context.Context) (aws.Config, error) {
	if r.cfg.AWSRoleARN == "" {
		return aws.Config{}, fmt.Errorf("AWS_ROLE_ARN not set")
	}
	tok, err := gceIDToken(ctx, r.cfg.AWSRoleARN)
	if err != nil {
		return aws.Config{}, err
	}
	anon := aws.Config{Region: r.cfg.AWSRegion, Credentials: aws.AnonymousCredentials{}}
	out, err := sts.NewFromConfig(anon).AssumeRoleWithWebIdentity(ctx, &sts.AssumeRoleWithWebIdentityInput{
		RoleArn: aws.String(r.cfg.AWSRoleARN), RoleSessionName: aws.String("ai-agent"), WebIdentityToken: aws.String(tok), DurationSeconds: aws.Int32(900),
	})
	if err != nil {
		return aws.Config{}, fmt.Errorf("assume-role-with-web-identity: %w", err)
	}
	c := out.Credentials
	return aws.Config{Region: r.cfg.AWSRegion, Credentials: credentials.NewStaticCredentialsProvider(*c.AccessKeyId, *c.SecretAccessKey, *c.SessionToken)}, nil
}

// AWSCallerIdentity is used by the self-check.
func (r *Runner) AWSCallerIdentity(ctx context.Context) (string, error) {
	cfg, err := r.awsCfg(ctx)
	if err != nil {
		return "", err
	}
	id, err := sts.NewFromConfig(cfg).GetCallerIdentity(ctx, &sts.GetCallerIdentityInput{})
	if err != nil {
		return "", err
	}
	return aws.ToString(id.Arn), nil
}

func (r *Runner) awsState(ctx context.Context) (any, error) {
	cfg, err := r.awsCfg(ctx)
	if err != nil {
		return nil, err
	}
	ec := ec2.NewFromConfig(cfg)
	inst, err := ec.DescribeInstances(ctx, &ec2.DescribeInstancesInput{Filters: []ec2types.Filter{{Name: aws.String("tag:Name"), Values: []string{"selfheal-*"}}}})
	if err != nil {
		return nil, err
	}
	var nodes []map[string]any
	for _, res := range inst.Reservations {
		for _, i := range res.Instances {
			name := ""
			for _, t := range i.Tags {
				if aws.ToString(t.Key) == "Name" {
					name = aws.ToString(t.Value)
				}
			}
			nodes = append(nodes, map[string]any{"name": name, "id": aws.ToString(i.InstanceId), "state": string(i.State.Name), "private_ip": aws.ToString(i.PrivateIpAddress)})
		}
	}
	sort.Slice(nodes, func(a, b int) bool { return nodes[a]["name"].(string) < nodes[b]["name"].(string) })
	lb := elbv2.NewFromConfig(cfg)
	tgs, err := lb.DescribeTargetGroups(ctx, &elbv2.DescribeTargetGroupsInput{Names: []string{"selfheal-tg"}})
	var targets []map[string]any
	if err == nil && len(tgs.TargetGroups) > 0 {
		th, err := lb.DescribeTargetHealth(ctx, &elbv2.DescribeTargetHealthInput{TargetGroupArn: tgs.TargetGroups[0].TargetGroupArn})
		if err == nil {
			for _, t := range th.TargetHealthDescriptions {
				targets = append(targets, map[string]any{"target": aws.ToString(t.Target.Id), "port": aws.ToInt32(t.Target.Port), "state": string(t.TargetHealth.State), "reason": string(t.TargetHealth.Reason)})
			}
		}
	}
	return map[string]any{"ec2": nodes, "alb_targets": targets}, nil
}

// ---- RAG ----

func (r *Runner) SearchRunbooks(ctx context.Context, query string, k int) (any, error) {
	if r.rag == nil || r.oai == nil {
		return nil, fmt.Errorf("rag store not available")
	}
	if k <= 0 || k > 8 {
		k = 4
	}
	vecs, err := r.oai.Embed(ctx, r.cfg.EmbedModel, []string{query})
	if err != nil {
		return nil, err
	}
	hits, err := r.rag.Search(ctx, vecs[0], k)
	if err != nil {
		return nil, err
	}
	return hits, nil
}

func str(m map[string]any, k string) string {
	v, _ := m[k].(string)
	return v
}

func num(m map[string]any, k string) int64 {
	switch v := m[k].(type) {
	case float64:
		return int64(v)
	case int64:
		return v
	case int:
		return int64(v)
	}
	return 0
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "..."
}