feat(ai): add AI VM services (diagnosis agent, pgvector RAG, Discord HITL bot) and harness approval endpoints
실험 설계서 6장(WP-12)의 AI VM 2층 구현.
- cmd/agent: 하네스 /diagnose 수신 → OpenAI Responses API 함수 호출 루프(읽기 전용 도구 7종:
query_metrics, query_logs, get_probe_status, get_k8s, get_gcp_instances, get_aws_state,
search_runbooks) → strict 스키마 submit_diagnosis(6.3) → 하네스 /proposals 등록 + HITL 게시,
또는 신뢰도 0.5 미만·escalate 시 3층 큐 에스컬레이션. 세션·프롬프트 버전·러너북 벡터는
pgvector(ai 스키마)에 저장. 라벨·조치 enum은 하네스 /scenarios·/actions/catalog에서 갱신.
- cmd/hitl: Discord Gateway 봇(intents GUILDS만). 승인/거부 버튼(custom_id approve:/reject:<token>),
3초 내 DEFERRED_UPDATE_MESSAGE, 승인 시 하네스 승인 → /actions/execute → 원 메시지 편집,
15분 무응답 시 에스컬레이션 채널로 이동. 재기동 시 미결 제안 타이머 복원.
- harness: POST /proposals/{token}/approve, /reject 추가. 제안은 진단 시점에 등록하고 승인은
버튼 시점에 기록해 MTTA(승인-제안)가 실제 대기 시간이 되도록 함.
OpenAI·Discord 토큰은 Secret Manager 파일로만 읽고 코드·저장소에 두지 않는다.
This commit is contained in:
seon committed
1 parent
de076a79e4
commit
f30af2a4c3
21 files changed
+2875
-1
No files matched your search
@@ -0,0 +1,358 @@
|
||||
// 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, 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)
|
||||
}
|
||||
|
||||
func (b *Bot) edit(m *rag.HitlMessage, status string, color int, result map[string]any) error {
|
||||
ctx := context.Background()
|
||||
d := map[string]any{}
|
||||
if se, err := b.store.GetHitl(ctx, m.Token); err == nil {
|
||||
m = se
|
||||
}
|
||||
em := embedFor(d, m.ActionID, m.Target, m.RunID, m.ScenarioID, status, color)
|
||||
em.Fields = em.Fields[1:] // the diagnosis fields are already in the original message; keep the edit compact
|
||||
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)
|
||||
}()
|
||||
}
|
||||
Reference in new issue
Block a user