Vectrum
Sign In
397 lines · 12.4 KB · Go
// Package herdr reports Codey's state to herdr, the terminal workspace
// manager, when Codey is running inside one of its panes.
//
// Herdr owns the terminal a coding agent runs in and shows, across every
// project, which agent is working, which is blocked waiting on a human, and
// which is idle. For agents it ships support for it reads that off the screen
// with pattern rules; for anything else it offers a contract, which is what
// this implements: an agent inside a pane inherits HERDR_ENV=1, HERDR_PANE_ID
// and HERDR_BIN_PATH, and reports semantic state by running
//
//	"$HERDR_BIN_PATH" pane report-agent "$HERDR_PANE_ID" \
//	    --source custom:codey --agent codey --state working
//
// releasing the claim on exit. Reporting beats screen-scraping for us in both
// directions: herdr never has to guess from what the TUI drew, and a state it
// could not see at all — that the agent is sitting on an approval prompt
// behind a full-screen view — is stated outright.
//
// Everything here is a no-op outside herdr. Detect returns nil when the
// environment does not say we are in a pane, which is the whole of the
// activation logic: no probing, no socket connection, no cost to anyone who
// has never heard of herdr.
package herdr

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"

	"codey/internal/hooks"
)

// The identity Codey claims in a pane. Source must be stable across runs and
// unique to this integration: herdr keys lifecycle authority on it, so a
// source that changed between reports would leave the old one owning the pane.
const (
	Source = "custom:codey"
	Agent  = "codey"
)

// The states herdr understands. Anything else is refused by its CLI, so they
// are spelled once here rather than at each call.
const (
	StateWorking = "working"
	StateBlocked = "blocked"
	StateIdle    = "idle"
)

// reportTimeout bounds one call to the herdr CLI. Reports are fired from the
// agent's goroutine at turn boundaries, so this is latency the user would
// feel; herdr is a local binary talking to a local socket, and one that has
// not answered in two seconds is one worth giving up on until the next event.
const reportTimeout = 2 * time.Second

// maxMessage is how much of a block reason is worth sending. The message is a
// line in a sidebar, not a report.
const maxMessage = 160

// Reporter maps Codey's lifecycle events to herdr's agent states.
//
// A pane holds one Codey process, and one Codey process can hold several
// sessions at once — tabs. Herdr has one state per pane, so the states are
// tracked per session and reduced: blocked beats working beats idle, because
// what a person scanning a sidebar needs to find is the agent that stopped and
// wants them, and after that the one still going.
type Reporter struct {
	bin  string
	pane string

	// run executes the herdr CLI. A field so the tests do not need herdr
	// installed; nil means the real one.
	run func(ctx context.Context, bin string, args []string) error

	mu sync.Mutex
	// states is what each live session is doing, keyed by session id.
	states map[string]*sessionState
	// session is the native session reference handed to herdr for restore:
	// the first session opened, which is the one `codey -resume` would want.
	session string
	// last is what was last reported, so an unchanged state costs nothing.
	last     string
	lastMsg  string
	lastSess string
	// seq strictly increases across reports. Herdr ignores a stale sequence
	// number from the same source, which is what keeps a slow report for an
	// old state from landing on top of a new one.
	seq      int64
	reported bool
	closed   bool
}

// Detect returns a reporter when this process is running inside a herdr pane,
// and nil when it is not.
func Detect() *Reporter { return detect(os.Getenv) }

// detect is Detect against an arbitrary environment, for tests.
func detect(env func(string) string) *Reporter {
	// HERDR_ENV=1 is the flag herdr documents for "you are inside a pane".
	// Requiring the other two as well is not paranoia: a stale HERDR_ENV
	// exported into a shell that outlived its pane is exactly the case where
	// reporting would go nowhere, slowly, forever.
	if env("HERDR_ENV") != "1" {
		return nil
	}
	pane, bin := env("HERDR_PANE_ID"), env("HERDR_BIN_PATH")
	if pane == "" || bin == "" {
		return nil
	}
	return &Reporter{bin: bin, pane: pane, states: map[string]*sessionState{}}
}

// sessionState is one conversation as herdr should see it.
//
// The two halves are tracked separately because they move independently: the
// loop says whether it is running, and a prompt says whether somebody is being
// waited on, and a subagent can be holding a prompt open long after the
// session's own agent went quiet. Folding them into one field would mean an
// answered prompt had to guess what the agent had been doing before it.
type sessionState struct {
	// base is what the loop last said: working or idle.
	base string
	// blocks are the messages of the prompts currently open, innermost last.
	// A count rather than a flag because a session's own agent and a subagent
	// of it can both be waiting, and the first answer given must not clear the
	// second question.
	blocks []string
}

// state is what this session is, once both halves are taken into account.
func (s *sessionState) state() string {
	if len(s.blocks) > 0 {
		return StateBlocked
	}
	if s.base == "" {
		return StateIdle
	}
	return s.base
}

// message is what it is waiting for, if anything. The innermost prompt wins:
// it is the one on screen.
func (s *sessionState) message() string {
	if len(s.blocks) == 0 {
		return ""
	}
	return s.blocks[len(s.blocks)-1]
}

// Pane reports which pane this is reporting to, for a startup notice.
func (r *Reporter) Pane() string {
	if r == nil {
		return ""
	}
	return r.pane
}

// IsNil lets hooks.Add drop a typed-nil reporter, so a caller can write
// hooks.Add(herdr.Detect()) without a nil check of its own.
func (r *Reporter) IsNil() bool { return r == nil }

// Handle implements hooks.Handler: it folds one lifecycle event into this
// pane's state and reports the result if it changed.
func (r *Reporter) Handle(ctx context.Context, p hooks.Payload) error {
	if r == nil {
		return nil
	}

	// A subagent's own comings and goings are not the pane's state: its parent
	// is running while it works, and already says so. Its prompts are another
	// matter — an approval a subagent needs is a person being waited on just
	// as much as one its parent needs, and the parent may well have gone quiet
	// long before it was asked, since spawning does not block.
	if p.Agent != "" && p.Event != hooks.Blocked && p.Event != hooks.Unblocked {
		return nil
	}

	r.mu.Lock()
	if r.closed {
		r.mu.Unlock()
		return nil
	}
	r.applyLocked(p)
	args, send := r.nextLocked()
	r.mu.Unlock()

	if !send {
		return nil
	}
	return r.exec(ctx, args)
}

// applyLocked folds one event into the session it came from. The caller holds
// the lock.
func (r *Reporter) applyLocked(p hooks.Payload) {
	if p.Event == hooks.SessionEnd && p.Agent == "" {
		delete(r.states, p.Session)
		return
	}

	s := r.states[p.Session]
	if s == nil {
		s = &sessionState{}
		r.states[p.Session] = s
	}

	switch p.Event {
	case hooks.SessionStart:
		// The first session opened is the one `codey -resume` would want, and
		// so the one worth telling herdr about.
		if r.session == "" {
			r.session = p.Session
		}
		s.base = StateIdle
	case hooks.RunStart, hooks.TurnStart, hooks.ToolStart, hooks.ToolResult, hooks.Resumed:
		s.base = StateWorking
	case hooks.RunEnd:
		// The model stopped asking for tools: the answer is on screen and the
		// next move is the user's.
		s.base = StateIdle
	case hooks.Paused:
		// Held at a turn boundary by the user. Not working, and not waiting on
		// an answer either — idle is the honest reading, and it keeps a paused
		// agent out of the list of ones that need attention.
		s.base = StateIdle
	case hooks.Blocked:
		s.blocks = append(s.blocks, blockMessage(p))
	case hooks.Unblocked:
		if len(s.blocks) > 0 {
			s.blocks = s.blocks[:len(s.blocks)-1]
		}
	}
}

// blockMessage is the one-line explanation shown beside a blocked pane.
func blockMessage(p hooks.Payload) string {
	msg := p.Message
	if msg == "" {
		switch p.Reason {
		case hooks.ReasonApproval:
			msg = "waiting for approval"
		case hooks.ReasonQuestion:
			msg = "waiting for an answer"
		default:
			msg = "waiting for input"
		}
	}
	return sanitize(msg)
}

// sanitize makes a message safe to pass as a command-line argument and
// sensible to show in a sidebar: one line, no control characters, bounded.
func sanitize(s string) string {
	s = strings.Map(func(r rune) rune {
		if r == '\n' || r == '\t' || r == '\r' {
			return ' '
		}
		if r < 0x20 || r == 0x7f {
			return -1
		}
		return r
	}, s)
	s = strings.Join(strings.Fields(s), " ")
	if len(s) > maxMessage {
		s = strings.TrimSpace(s[:maxMessage]) + "…"
	}
	return s
}

// nextLocked reduces the per-session states to the one the pane is in and
// builds the command for it, or reports that nothing changed. The caller holds
// the lock.
func (r *Reporter) nextLocked() (args []string, send bool) {
	state, message := reduce(r.states)
	if state == "" {
		// Every session has ended. The claim is released by Close, not here:
		// the process is still alive, and a released source that reported
		// again would be claiming the pane back.
		return nil, false
	}
	if r.reported && state == r.last && message == r.lastMsg && r.session == r.lastSess {
		return nil, false
	}
	r.seq++
	r.last, r.lastMsg, r.lastSess, r.reported = state, message, r.session, true

	args = []string{
		"pane", "report-agent", r.pane,
		"--source", Source,
		"--agent", Agent,
		"--state", state,
		"--seq", strconv.FormatInt(r.seq, 10),
	}
	if message != "" {
		args = append(args, "--message", message)
	}
	// The session reference travels with the state report rather than through
	// a second call: herdr accepts it there, and one process is one report.
	if r.session != "" {
		args = append(args, "--agent-session-id", r.session)
	}
	return args, true
}

// reduce picks the state of a pane holding several sessions. Blocked wins
// because it is the one that needs a person; working beats idle because a pane
// with anything running in it is not idle.
func reduce(states map[string]*sessionState) (state, message string) {
	var working, idle bool
	for _, s := range states {
		switch s.state() {
		case StateBlocked:
			// Which blocked session's message is shown is arbitrary when
			// there are two, since map order is: the state is what a sidebar
			// sorts on, and either message names a real thing being waited on.
			if state != StateBlocked {
				state, message = StateBlocked, s.message()
			}
		case StateWorking:
			working = true
		default:
			idle = true
		}
	}
	if state == StateBlocked {
		return state, message
	}
	switch {
	case working:
		return StateWorking, ""
	case idle:
		return StateIdle, ""
	default:
		return "", ""
	}
}

// Close releases this source's claim on the pane, so herdr stops attributing a
// state to a Codey that is no longer there. It is safe to call twice.
func (r *Reporter) Close() error {
	if r == nil {
		return nil
	}
	r.mu.Lock()
	if r.closed {
		r.mu.Unlock()
		return nil
	}
	r.closed = true
	reported := r.reported
	r.mu.Unlock()

	if !reported {
		// Never claimed the pane, so there is nothing to release.
		return nil
	}
	return r.exec(context.Background(), []string{
		"pane", "release-agent", r.pane,
		"--source", Source,
		"--agent", Agent,
	})
}

// exec runs the herdr CLI with its own deadline, detached from the turn's
// context: an interrupted turn still has a state worth reporting, and the
// report that says so must not be cancelled by the interrupt that caused it.
func (r *Reporter) exec(ctx context.Context, args []string) error {
	run := r.run
	if run == nil {
		run = runCLI
	}
	cmdCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), reportTimeout)
	defer cancel()
	if err := run(cmdCtx, r.bin, args); err != nil {
		return fmt.Errorf("herdr %s: %w", strings.Join(args[:2], " "), err)
	}
	return nil
}

// runCLI is the real herdr invocation. Output is discarded unless the command
// failed: herdr's CLI prints nothing useful on success, and anything it did
// print would land in the middle of the TUI.
func runCLI(ctx context.Context, bin string, args []string) error {
	cmd := exec.CommandContext(ctx, bin, args...)
	out, err := cmd.CombinedOutput()
	if err != nil {
		if trimmed := strings.TrimSpace(string(out)); trimmed != "" {
			return fmt.Errorf("%w: %s", err, trimmed)
		}
		return err
	}
	return nil
}