selfheal-lab-apps/ai/internal/hitl/bot.go

// Package hitl is the Discord human-in-the-loop bot of experiment design 6.5/6.6: it posts layer-2 proposals with
// Approve/Reject buttons, turns a button press into harness approval + execution, expires silent proposals into
// the layer-3 (escalation) channel, and posts escalations with the matching runbook in a thread.
package hitl

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"strings"
	"sync"
	"time"

	"github.com/bwmarrin/discordgo"

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

type Config struct {
	GuildID           string
	ApprovalsChannel  string
	EscalationChannel string
	ApproverRoleID    string // optional: only members with this role may decide
}

// Discord abstracts the discordgo calls so the decision logic is testable without a gateway.
type Discord interface {
	Send(channelID string, msg *discordgo.MessageSend) (*discordgo.Message, error)
	Edit(channelID, messageID string, embed *discordgo.MessageEmbed, components []discordgo.MessageComponent) error
	Thread(channelID, messageID, name, text string) error
}

type Bot struct {
	cfg     Config
	dc      Discord
	harness *harness.Client
	store   *rag.Store
	mu      sync.Mutex
	timers  map[string]*time.Timer
}

func New(cfg Config, dc Discord, hc *harness.Client, store *rag.Store) *Bot {
	return &Bot{cfg: cfg, dc: dc, harness: hc, store: store, timers: map[string]*time.Timer{}}
}

// PostRequest is what the agent sends to POST /post.
type PostRequest struct {
	Token      string         `json:"token"`
	ActionID   string         `json:"action_id"`
	Target     map[string]any `json:"target"`
	Diagnosis  map[string]any `json:"diagnosis"`
	RunID      *string        `json:"run_id"`
	ScenarioID *string        `json:"scenario_id"`
	SessionID  string         `json:"session_id"`
	ExpiresIn  float64        `json:"expires_in_s"`
}

func deref(s *string) string {
	if s == nil {
		return "-"
	}
	return *s
}

func embedFor(d map[string]any, actionID string, target map[string]any, runID, scenarioID *string, status string, color int) *discordgo.MessageEmbed {
	ev, _ := d["evidence"].([]any)
	var evLines []string
	for i, e := range ev {
		if i >= 5 {
			break
		}
		evLines = append(evLines, "- "+fmt.Sprint(e))
	}
	tgt, _ := json.Marshal(target)
	return &discordgo.MessageEmbed{
		Title: fmt.Sprintf("Layer-2 proposal: %s", actionID),
		Color: color,
		Fields: []*discordgo.MessageEmbedField{
			{Name: "Root cause", Value: fmt.Sprintf("%v (confidence %.2f)", d["root_cause_label"], toF(d["confidence"])), Inline: false},
			{Name: "Target", Value: "`" + string(tgt) + "`", Inline: false},
			{Name: "Evidence", Value: nonEmpty(strings.Join(evLines, "\n")), Inline: false},
			{Name: "Rationale", Value: nonEmpty(fmt.Sprint(d["rationale"])), Inline: false},
			{Name: "Run / scenario", Value: deref(runID) + " / " + deref(scenarioID), Inline: true},
			{Name: "Status", Value: status, Inline: true},
		},
		Footer: &discordgo.MessageEmbedFooter{Text: fmt.Sprintf("model %v, prompt %v, session %v", d["model"], d["prompt_version"], d["session_id"])},
	}
}

func nonEmpty(s string) string {
	if strings.TrimSpace(s) == "" {
		return "-"
	}
	if len(s) > 1000 {
		return s[:1000] + "..."
	}
	return s
}

func toF(v any) float64 {
	f, _ := v.(float64)
	return f
}

const (
	colorPending  = 0x5865f2
	colorApproved = 0x57f287
	colorRejected = 0xed4245
	colorExpired  = 0x99aab5
	colorEscalate = 0xfee75c
)

func buttons(token string, disabled bool) []discordgo.MessageComponent {
	return []discordgo.MessageComponent{discordgo.ActionsRow{Components: []discordgo.MessageComponent{
		discordgo.Button{Label: "Approve", Style: discordgo.SuccessButton, CustomID: "approve:" + token, Disabled: disabled},
		discordgo.Button{Label: "Reject", Style: discordgo.DangerButton, CustomID: "reject:" + token, Disabled: disabled},
	}}}
}

// Post publishes the proposal to the approvals channel and arms the expiry timer.
func (b *Bot) Post(ctx context.Context, req PostRequest) error {
	ttl := time.Duration(req.ExpiresIn) * time.Second
	if ttl <= 0 {
		ttl = 15 * time.Minute
	}
	expires := time.Now().Add(ttl)
	em := embedFor(req.Diagnosis, req.ActionID, req.Target, req.RunID, req.ScenarioID, fmt.Sprintf("waiting for approval until <t:%d:T>", expires.Unix()), colorPending)
	msg, err := b.dc.Send(b.cfg.ApprovalsChannel, &discordgo.MessageSend{Embeds: []*discordgo.MessageEmbed{em}, Components: buttons(req.Token, false)})
	if err != nil {
		return err
	}
	if err := b.store.InsertHitl(ctx, rag.HitlMessage{Token: req.Token, ChannelID: msg.ChannelID, MessageID: msg.ID, ActionID: req.ActionID, Target: req.Target, Diagnosis: req.Diagnosis, RunID: req.RunID, ScenarioID: req.ScenarioID, SessionID: req.SessionID, ExpiresAt: expires}); err != nil {
		return err
	}
	b.arm(req.Token, time.Until(expires))
	return nil
}

func (b *Bot) arm(token string, d time.Duration) {
	if d < 0 {
		d = 0
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	if t, ok := b.timers[token]; ok {
		t.Stop()
	}
	b.timers[token] = time.AfterFunc(d, func() { b.expire(token) })
}

func (b *Bot) disarm(token string) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if t, ok := b.timers[token]; ok {
		t.Stop()
		delete(b.timers, token)
	}
}

// RearmPending re-creates expiry timers after a restart.
func (b *Bot) RearmPending(ctx context.Context) {
	list, err := b.store.Pending(ctx)
	if err != nil {
		log.Printf("pending: %v", err)
		return
	}
	for _, m := range list {
		b.arm(m.Token, time.Until(m.ExpiresAt))
	}
	log.Printf("re-armed %d pending proposals", len(list))
}

// Decide handles a button press. by is "username (id)". Returns the status text for logging.
func (b *Bot) Decide(ctx context.Context, token, decision, by string) (string, error) {
	m, err := b.store.GetHitl(ctx, token)
	if err != nil {
		return "", fmt.Errorf("unknown proposal %s", token)
	}
	first, err := b.store.Decide(ctx, token, decision, by, nil)
	if err != nil {
		return "", err
	}
	if !first {
		return "already decided", nil
	}
	b.disarm(token)
	switch decision {
	case "reject":
		if err := b.harness.Reject(ctx, token, by, "rejected in Discord"); err != nil {
			log.Printf("harness reject: %v", err)
		}
		_ = b.edit(m, "rejected by "+by, colorRejected, nil)
		return "rejected", nil
	case "approve":
		if time.Now().After(m.ExpiresAt) {
			_ = b.edit(m, "expired before approval", colorExpired, nil)
			return "expired", nil
		}
		if err := b.harness.Approve(ctx, token, by); err != nil {
			_ = b.edit(m, "approval failed: "+err.Error(), colorRejected, nil)
			_ = b.store.SetHitlResult(ctx, token, map[string]any{"error": err.Error()})
			return "approve failed", err
		}
		_ = b.edit(m, "approved by "+by+", executing...", colorApproved, nil)
		res, err := b.harness.Execute(ctx, token, m.ActionID, m.Target, m.RunID, m.ScenarioID, "approved in Discord by "+by)
		if err != nil {
			_ = b.store.SetHitlResult(ctx, token, map[string]any{"error": err.Error(), "result": res})
			_ = b.edit(m, "approved by "+by+"; execution failed: "+truncate(err.Error(), 300), colorRejected, res)
			return "execute failed", err
		}
		_ = b.store.SetHitlResult(ctx, token, res)
		_ = b.edit(m, "approved by "+by+"; executed", colorApproved, res)
		return "executed", nil
	}
	return "", fmt.Errorf("unknown decision %s", decision)
}

// edit re-renders the original message with the stored diagnosis, the new status and disabled buttons.
func (b *Bot) edit(m *rag.HitlMessage, status string, color int, result map[string]any) error {
	em := embedFor(m.Diagnosis, m.ActionID, m.Target, m.RunID, m.ScenarioID, status, color)
	if result != nil {
		rb, _ := json.Marshal(result)
		em.Fields = append(em.Fields, &discordgo.MessageEmbedField{Name: "Result", Value: "```json\n" + truncate(string(rb), 900) + "\n```"})
	}
	return b.dc.Edit(m.ChannelID, m.MessageID, em, buttons(m.Token, true))
}

func (b *Bot) expire(token string) {
	ctx := context.Background()
	m, err := b.store.GetHitl(ctx, token)
	if err != nil {
		return
	}
	first, err := b.store.Decide(ctx, token, "expired", "timer", nil)
	if err != nil || !first {
		return
	}
	_ = b.edit(m, "no decision within the window; moved to layer 3", colorExpired, nil)
	text := fmt.Sprintf("Proposal `%s` for `%s` (run %s, scenario %s) received no decision within the approval window. Layer 3: follow the runbook or call the harness `/actions/execute` with the operator token.",
		m.ActionID, mustJSON(m.Target), deref(m.RunID), deref(m.ScenarioID))
	if _, err := b.dc.Send(b.cfg.EscalationChannel, &discordgo.MessageSend{Content: text}); err != nil {
		log.Printf("escalation post: %v", err)
	}
}

// EscalateRequest is what the agent sends to POST /escalate (design 6.6).
type EscalateRequest struct {
	Diagnosis  map[string]any `json:"diagnosis"`
	Runbook    string         `json:"runbook"`
	RunID      *string        `json:"run_id"`
	ScenarioID *string        `json:"scenario_id"`
	Alert      map[string]any `json:"alert"`
}

func (b *Bot) Escalate(ctx context.Context, req EscalateRequest) error {
	em := embedFor(req.Diagnosis, "escalate", map[string]any{"target": req.Diagnosis["target"]}, req.RunID, req.ScenarioID, "layer 3: human decision required", colorEscalate)
	em.Title = fmt.Sprintf("Escalation: %v", req.Diagnosis["root_cause_label"])
	if an, ok := req.Alert["alertname"]; ok {
		em.Fields = append(em.Fields, &discordgo.MessageEmbedField{Name: "Alert", Value: fmt.Sprint(an), Inline: true})
	}
	msg, err := b.dc.Send(b.cfg.EscalationChannel, &discordgo.MessageSend{Embeds: []*discordgo.MessageEmbed{em}})
	if err != nil {
		return err
	}
	name := fmt.Sprintf("%s %s", deref(req.ScenarioID), time.Now().UTC().Format("01-02 15:04"))
	text := "Runbook:\n" + nonEmpty(truncate(req.Runbook, 1800))
	if err := b.dc.Thread(msg.ChannelID, msg.ID, name, text); err != nil {
		log.Printf("thread: %v", err)
	}
	return nil
}

// Allowed checks the optional approver role.
func (b *Bot) Allowed(roles []string) bool {
	if b.cfg.ApproverRoleID == "" {
		return true
	}
	for _, r := range roles {
		if r == b.cfg.ApproverRoleID {
			return true
		}
	}
	return false
}

func mustJSON(v any) string {
	b, _ := json.Marshal(v)
	return string(b)
}

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

// ---- discordgo adapter ----

type Session struct{ S *discordgo.Session }

func (d Session) Send(channelID string, msg *discordgo.MessageSend) (*discordgo.Message, error) {
	return d.S.ChannelMessageSendComplex(channelID, msg)
}

func (d Session) Edit(channelID, messageID string, embed *discordgo.MessageEmbed, components []discordgo.MessageComponent) error {
	embeds := []*discordgo.MessageEmbed{embed}
	_, err := d.S.ChannelMessageEditComplex(&discordgo.MessageEdit{Channel: channelID, ID: messageID, Embeds: &embeds, Components: &components})
	return err
}

func (d Session) Thread(channelID, messageID, name, text string) error {
	th, err := d.S.MessageThreadStartComplex(channelID, messageID, &discordgo.ThreadStart{Name: truncate(name, 90), AutoArchiveDuration: 1440})
	if err != nil {
		return err
	}
	_, err = d.S.ChannelMessageSend(th.ID, text)
	return err
}

// HandleInteraction is the discordgo InteractionCreate handler: DEFERRED_UPDATE_MESSAGE within 3 s, then the decision.
func (b *Bot) HandleInteraction(s *discordgo.Session, i *discordgo.InteractionCreate) {
	if i.Type != discordgo.InteractionMessageComponent {
		return
	}
	cid := i.MessageComponentData().CustomID
	parts := strings.SplitN(cid, ":", 2)
	if len(parts) != 2 || (parts[0] != "approve" && parts[0] != "reject") {
		return
	}
	var roles []string
	by := "unknown"
	if i.Member != nil && i.Member.User != nil {
		roles = i.Member.Roles
		by = fmt.Sprintf("%s (%s)", i.Member.User.Username, i.Member.User.ID)
	}
	if !b.Allowed(roles) {
		_ = s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{Type: discordgo.InteractionResponseChannelMessageWithSource,
			Data: &discordgo.InteractionResponseData{Content: "You are not in the approver role.", Flags: discordgo.MessageFlagsEphemeral}})
		return
	}
	if err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{Type: discordgo.InteractionResponseDeferredMessageUpdate}); err != nil {
		log.Printf("interaction respond: %v", err)
	}
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), 11*time.Minute)
		defer cancel()
		st, err := b.Decide(ctx, parts[1], parts[0], by)
		log.Printf("decision %s on %s by %s: %s err=%v", parts[0], parts[1][len(parts[1])-6:], by, st, err)
	}()
}