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

// hitl: Discord approval bot of the AI VM (experiment design 6.5). HTTP on :8081 for the agent; Gateway for buttons.
package main

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

	"github.com/bwmarrin/discordgo"

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

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

func main() {
	log.SetFlags(log.LstdFlags | log.Lmicroseconds)
	ctx := context.Background()
	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)
	}
	cfg := hitl.Config{GuildID: env("DISCORD_GUILD_ID", ""), ApprovalsChannel: env("DISCORD_APPROVALS_CHANNEL_ID", ""), EscalationChannel: env("DISCORD_ESCALATION_CHANNEL_ID", ""), ApproverRoleID: env("DISCORD_APPROVER_ROLE_ID", ""), Lang: env("HITL_LANG", "en")}
	tokB, err := os.ReadFile(env("DISCORD_TOKEN_FILE", "/run/secrets/discord-bot-token"))
	token := strings.TrimSpace(string(tokB))
	connected := false
	var bot *hitl.Bot
	hc := harness.New(env("HARNESS_URL", "http://10.20.0.11:8080"))
	var dg *discordgo.Session
	if err != nil || token == "" {
		log.Printf("discord token missing (%v); running degraded: /post and /escalate return 503 until the secret exists", err)
	} else {
		dg, err = discordgo.New("Bot " + token)
		if err != nil {
			log.Fatalf("discord session: %v", err)
		}
		// Minimum intent set (design 11.7): GUILDS only. INTERACTION_CREATE carries no intent requirement; no privileged intents.
		dg.Identify.Intents = discordgo.IntentsGuilds
		bot = hitl.New(cfg, hitl.Session{S: dg}, hc, store)
		dg.AddHandler(bot.HandleInteraction)
		dg.AddHandler(func(s *discordgo.Session, r *discordgo.Ready) {
			log.Printf("discord ready: user=%s guilds=%d", r.User.Username, len(r.Guilds))
		})
		if err := dg.Open(); err != nil {
			log.Fatalf("discord open: %v", err)
		}
		connected = true
		bot.RearmPending(ctx)
	}

	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) {
		out := map[string]any{"discord_connected": connected, "guild": cfg.GuildID, "approvals_channel": cfg.ApprovalsChannel, "escalation_channel": cfg.EscalationChannel, "approver_role": cfg.ApproverRoleID}
		if dg != nil {
			if ch, err := dg.Channel(cfg.ApprovalsChannel); err == nil {
				out["approvals_channel_name"] = ch.Name
			} else {
				out["approvals_channel_error"] = err.Error()
			}
		}
		pend, _ := store.Pending(r.Context())
		out["pending"] = len(pend)
		writeJSON(w, 200, out)
	})
	mux.HandleFunc("POST /post", func(w http.ResponseWriter, r *http.Request) {
		if !connected {
			writeJSON(w, 503, map[string]string{"error": "discord not connected"})
			return
		}
		var req hitl.PostRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil || req.Token == "" {
			writeJSON(w, 400, map[string]string{"error": "bad request"})
			return
		}
		if err := bot.Post(r.Context(), req); err != nil {
			writeJSON(w, 500, map[string]string{"error": err.Error()})
			return
		}
		writeJSON(w, 201, map[string]string{"status": "posted"})
	})
	mux.HandleFunc("POST /escalate", func(w http.ResponseWriter, r *http.Request) {
		if !connected {
			writeJSON(w, 503, map[string]string{"error": "discord not connected"})
			return
		}
		var req hitl.EscalateRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeJSON(w, 400, map[string]string{"error": "bad request"})
			return
		}
		if err := bot.Escalate(r.Context(), req); err != nil {
			writeJSON(w, 500, map[string]string{"error": err.Error()})
			return
		}
		writeJSON(w, 201, map[string]string{"status": "escalated"})
	})
	srv := &http.Server{Addr: env("HITL_ADDR", ":8081"), Handler: mux, ReadHeaderTimeout: 10 * time.Second}
	go func() { log.Fatal(srv.ListenAndServe()) }()
	log.Printf("hitl up: connected=%v", connected)
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
	<-sig
	if dg != nil {
		_ = dg.Close()
	}
}

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