byre/internal/agents/health.go

package agents

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"g.fc.run/ai/byre/internal/api"
)

const (
	StuckAfter     = 5 * time.Minute
	healthInterval = 15 * time.Second
	keptBadLines   = 50
	badLineBytes   = 4096
)

type badLines struct {
	mu     sync.Mutex
	counts map[string]int
	recent []api.ProtocolError
}

func newBadLines() *badLines {
	return &badLines{counts: map[string]int{}}
}

func (b *badLines) add(vendor, sessionID, reason string, line []byte, at time.Time) {
	raw := line
	if len(raw) > badLineBytes {
		raw = raw[:badLineBytes]
	}
	e := api.ProtocolError{
		Vendor:    vendor,
		SessionId: sessionID,
		At:        at.UTC(),
		Reason:    reason,
		Raw:       string(raw),
		Size:      len(line),
	}

	b.mu.Lock()
	defer b.mu.Unlock()
	b.counts[vendor]++
	if len(b.recent) < keptBadLines {
		b.recent = append(b.recent, e)
		return
	}
	copy(b.recent, b.recent[1:])
	b.recent[len(b.recent)-1] = e
}

func (b *badLines) snapshot() api.ProtocolErrors {
	b.mu.Lock()
	defer b.mu.Unlock()
	out := api.ProtocolErrors{Counts: map[string]int{}, Recent: []api.ProtocolError{}}
	for v, n := range b.counts {
		out.Counts[v] = n
	}
	for i := len(b.recent) - 1; i >= 0; i-- {
		out.Recent = append(out.Recent, b.recent[i])
	}
	return out
}

func (m *Manager) recordBadLine(l *live, line []byte, reason string) {
	m.bad.add(string(l.vendor), l.id, reason, line, time.Now())
	m.store.Broadcast(api.Event{Kind: "admin.health", Actor: "byre", Payload: map[string]any{}})
}

func isDecodeError(err error) bool {
	var syntax *json.SyntaxError
	var typ *json.UnmarshalTypeError
	return errors.As(err, &syntax) || errors.As(err, &typ)
}

func (m *Manager) ProtocolErrors() api.ProtocolErrors {
	return m.bad.snapshot()
}

type Health struct {
	Processes []api.AgentProcess
	Zombies   []api.Zombie
}

type procInfo struct {
	sessionID string
	taskID    string
	vendor    api.AgentSessionVendor
	pid       int
	state     api.AgentSessionState
	lastLine  time.Time
	exited    bool
}

func (m *Manager) Health(ctx context.Context) (Health, error) {
	running, err := m.store.SessionsInState(ctx, api.AgentSessionStateRunning)
	if err != nil {
		return Health{}, err
	}

	children, err := claudeChildren(ctx, os.Getpid())
	if err != nil {
		return Health{}, err
	}

	// Snapshot live and starting after running and children, so a start that shows in either is already registered.
	m.mu.Lock()
	all := []*live{}
	for _, l := range m.live {
		all = append(all, l)
	}
	starting := map[string]bool{}
	for id := range m.starting {
		starting[id] = true
	}
	m.mu.Unlock()

	procs := []procInfo{}
	for _, l := range all {
		sess, err := m.store.GetSession(ctx, l.id)
		if err != nil {
			return Health{}, err
		}
		pid := l.sess.pid()
		procs = append(procs, procInfo{
			sessionID: l.id,
			taskID:    l.taskID,
			vendor:    l.vendor,
			pid:       pid,
			state:     sess.State,
			lastLine:  l.sess.lastLine(),
			exited:    l.sess.exited() || !pidAlive(pid),
		})
	}

	return assess(time.Now(), procs, running, children, starting), nil
}

func assess(now time.Time, procs []procInfo, running []api.AgentSession, children []int, starting map[string]bool) Health {
	slices.SortFunc(procs, func(a, b procInfo) int { return strings.Compare(a.sessionID, b.sessionID) })

	h := Health{Processes: []api.AgentProcess{}, Zombies: []api.Zombie{}}
	bySession := map[string]procInfo{}
	tracked := map[int]bool{}

	for _, p := range procs {
		bySession[p.sessionID] = p
		health := api.HealthOk
		switch {
		case p.exited:
			health = api.HealthExited
		case p.state == api.AgentSessionStateRunning && now.Sub(p.lastLine) > StuckAfter:
			health = api.HealthStuck
		}
		h.Processes = append(h.Processes, api.AgentProcess{
			SessionId:   p.sessionID,
			TaskId:      p.taskID,
			Vendor:      string(p.vendor),
			Pid:         p.pid,
			State:       string(p.state),
			LastEventAt: p.lastLine.UTC(),
			Health:      health,
		})
		if p.exited {
			continue
		}
		tracked[p.pid] = true
		if p.state == api.AgentSessionStateCompleted {
			h.Zombies = append(h.Zombies, api.Zombie{
				Kind:      api.ProcessWithoutSession,
				Detail:    fmt.Sprintf("session is completed but its process %d is still alive", p.pid),
				SessionId: &p.sessionID,
				TaskId:    &p.taskID,
				Pid:       &p.pid,
			})
		}
	}

	for _, sess := range running {
		p, ok := bySession[sess.Id]
		if (ok && !p.exited) || starting[sess.Id] {
			continue
		}
		detail := "no agent process"
		if ok {
			detail = fmt.Sprintf("process %d has exited", p.pid)
		}
		h.Zombies = append(h.Zombies, api.Zombie{
			Kind:      api.RunningWithoutProcess,
			Detail:    detail,
			SessionId: &sess.Id,
			TaskId:    &sess.TaskId,
		})
	}

	for _, pid := range children {
		if tracked[pid] || len(starting) > 0 {
			continue
		}
		h.Zombies = append(h.Zombies, api.Zombie{
			Kind:   api.ProcessWithoutSession,
			Detail: fmt.Sprintf("claude process %d is a child of byre but belongs to no session", pid),
			Pid:    &pid,
		})
	}

	return h
}

func pidAlive(pid int) bool {
	err := syscall.Kill(pid, 0)
	return err == nil || errors.Is(err, syscall.EPERM)
}

func claudeChildren(ctx context.Context, parent int) ([]int, error) {
	cctx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	out, err := exec.CommandContext(cctx, "ps", "-A", "-o", "pid=,ppid=,comm=").Output()
	if err != nil {
		return nil, fmt.Errorf("list processes: %w", err)
	}
	return parseChildren(out, parent), nil
}

func parseChildren(out []byte, parent int) []int {
	pids := []int{}
	sc := bufio.NewScanner(bytes.NewReader(out))
	for sc.Scan() {
		fields := strings.Fields(sc.Text())
		if len(fields) < 3 {
			continue
		}
		pid, err := strconv.Atoi(fields[0])
		if err != nil {
			continue
		}
		ppid, err := strconv.Atoi(fields[1])
		if err != nil || ppid != parent {
			continue
		}
		if filepath.Base(strings.Join(fields[2:], " ")) != "claude" {
			continue
		}
		pids = append(pids, pid)
	}
	return pids
}

func problems(h Health) []string {
	out := []string{}
	for _, p := range h.Processes {
		if p.Health != api.HealthOk {
			out = append(out, fmt.Sprintf("session %s process %d is %s", p.SessionId, p.Pid, p.Health))
		}
	}
	for _, z := range h.Zombies {
		id := ""
		if z.SessionId != nil {
			id = "session " + *z.SessionId + ": "
		}
		out = append(out, fmt.Sprintf("%s: %s%s", z.Kind, id, z.Detail))
	}
	slices.Sort(out)
	return out
}

func (m *Manager) Watch() {
	t := time.NewTicker(healthInterval)
	defer t.Stop()

	seen := []string{}
	for {
		select {
		case <-m.watchStop:
			return
		case <-t.C:
		}

		h, err := m.Health(context.Background())
		if err != nil {
			slog.Error("check agent health", "err", err)
			continue
		}
		m.stopStuckReviewers(context.Background(), h)

		now := problems(h)
		if slices.Equal(now, seen) {
			continue
		}
		for _, p := range now {
			if !slices.Contains(seen, p) {
				slog.Warn("agent health problem", "problem", p)
			}
		}
		seen = now
		m.store.Broadcast(api.Event{Kind: "admin.health", Actor: "byre", Payload: map[string]any{}})
	}
}