Vectrum
Sign In
283 lines · 8.9 KB · Go
// Package hooks is the lifecycle of a session, published.
//
// Codey already tells a front end what is happening through the agent's
// Observe callback, but that stream is about rendering: it carries text
// deltas and tool payloads, there is exactly one of it, and whoever owns the
// screen owns it. This is the other kind of notification — a small, stable set
// of things that happened, broadcast to anything that asked, carrying no
// terminal concerns and no conversation content.
//
// What wants that is anything outside the process: a terminal multiplexer that
// wants to know whether this pane is working or waiting on a human, a
// notifier, a script that logs turns somewhere. The events are therefore named
// for states a bystander cares about rather than for the loop's internals, and
// the set is kept short on purpose. Every one of them is cheap to fire and
// must stay that way: they are fired from the agent's own goroutine, so a
// handler that blocks is a handler that stalls the turn.
//
// The zero value and a nil *Hooks both work and drop everything, so callers
// never guard a Fire with a nil check.
package hooks

import (
	"context"
	"fmt"
	"sync"
	"time"
)

// Event names something that happened. The set is deliberately small: these
// are the moments an observer outside the process can act on, not a trace of
// the loop.
type Event string

const (
	// SessionStart and SessionEnd bracket one conversation. A second tab is a
	// second session and fires its own pair.
	SessionStart Event = "session_start"
	SessionEnd   Event = "session_end"

	// RunStart and RunEnd bracket one prompt: everything from the user's
	// message to the model's final answer, tool calls and all. Between them
	// the agent is busy; after RunEnd it is waiting for a human.
	RunStart Event = "run_start"
	RunEnd   Event = "run_end"

	// TurnStart is one model call within a run.
	TurnStart Event = "turn_start"

	// ToolStart and ToolResult bracket one tool call.
	ToolStart  Event = "tool_start"
	ToolResult Event = "tool_result"

	// Blocked and Unblocked bracket a wait on a human — an approval prompt or
	// a question from the ask tool. This is the distinction a bystander most
	// wants and cannot get by watching the process: an agent that is thinking
	// and an agent that has been sitting on a prompt for ten minutes look
	// exactly alike from outside.
	Blocked   Event = "blocked"
	Unblocked Event = "unblocked"

	// Paused and Resumed are the loop being held at a turn boundary.
	Paused  Event = "paused"
	Resumed Event = "resumed"
)

// Reasons a run is blocked, carried in Payload.Reason.
const (
	ReasonApproval = "approval"
	ReasonQuestion = "question"
)

// Payload is one event. Every field beyond Event and Time is optional, and
// which ones are set depends on the event; a handler must treat an empty field
// as "not applicable here" rather than as an error.
//
// Nothing here carries conversation content. Message is a short human-readable
// summary — which tool wants approval, what a question is about — because the
// point of these events is to reach things outside the process, and what the
// model wrote is not theirs to have.
type Payload struct {
	Event Event
	Time  time.Time

	// Session is the session id this came from, stamped by the Hooks the
	// event was fired on. Tabs mean several sessions in one process.
	Session string
	// Agent is the subagent id, empty for a session's own agent.
	Agent string

	Turn int
	Tool string

	// Reason says what kind of thing happened, where the event alone does not
	// distinguish: ReasonApproval or ReasonQuestion on Blocked.
	Reason string
	// Message is a short summary safe to show a bystander.
	Message string
	// Err is set when the event reports something that failed.
	Err string
}

// Handler receives events. Returning an error does not stop the run or the
// other handlers; it is reported through Hooks.OnError and otherwise ignored,
// because a broken notifier must not break the agent.
//
// Handle is called from the goroutine that fired the event and must return
// promptly. A handler with slow work to do owns its own queue.
type Handler interface {
	Handle(context.Context, Payload) error
}

// HandlerFunc adapts a function to Handler.
type HandlerFunc func(context.Context, Payload) error

// Handle implements Handler.
func (f HandlerFunc) Handle(ctx context.Context, p Payload) error { return f(ctx, p) }

// set is the handler list a Hooks and everything derived from it share.
// Derived hooks exist so a tab can stamp its own session id on every event
// without registering the handlers a second time, and sharing the lock is what
// keeps two tabs from reporting over each other.
type set struct {
	mu       sync.Mutex
	handlers []Handler
	onError  func(Payload, error)
}

// Hooks is a fan-out of lifecycle events to registered handlers, stamped with
// one session's id.
type Hooks struct {
	set     *set
	session string
}

// New returns hooks that stamp every event with a session id. Handlers are
// added with Add; until one is, firing costs a lock and a map-free return.
func New(session string) *Hooks {
	return &Hooks{set: &set{}, session: session}
}

// For returns hooks that stamp a different session id and write to the same
// handlers these do. Adding a handler to either adds it to both; closing
// either closes the handlers once.
func (h *Hooks) For(session string) *Hooks {
	if h == nil || h.set == nil {
		return nil
	}
	return &Hooks{set: h.set, session: session}
}

// Session reports the id stamped on events fired here.
func (h *Hooks) Session() string {
	if h == nil {
		return ""
	}
	return h.session
}

// Add registers handlers, in the order they will be called. Nil handlers are
// skipped, so a detector that returns nil when its thing is not present can be
// passed straight in.
func (h *Hooks) Add(handlers ...Handler) {
	if h == nil || h.set == nil {
		return
	}
	h.set.mu.Lock()
	defer h.set.mu.Unlock()
	for _, handler := range handlers {
		if handler == nil {
			continue
		}
		// A typed nil pointer in a non-nil interface is what a detector
		// returning (*Reporter)(nil) produces, and it would panic on the
		// first event rather than at the call that was wrong.
		if r, ok := handler.(interface{ IsNil() bool }); ok && r.IsNil() {
			continue
		}
		h.set.handlers = append(h.set.handlers, handler)
	}
}

// OnError installs what to do with a handler's error or panic. It is where a
// caller wires the metrics log; the default is to drop them, because a hook
// failing is not something to interrupt a session over and there is no
// guarantee anything is listening.
func (h *Hooks) OnError(f func(Payload, error)) {
	if h == nil || h.set == nil {
		return
	}
	h.set.mu.Lock()
	defer h.set.mu.Unlock()
	h.set.onError = f
}

// Enabled reports whether anything is listening, so a caller can skip building
// a Message nobody will read.
func (h *Hooks) Enabled() bool {
	if h == nil || h.set == nil {
		return false
	}
	h.set.mu.Lock()
	defer h.set.mu.Unlock()
	return len(h.set.handlers) > 0
}

// Fire delivers an event to every handler, in order, on the calling goroutine.
// Time and Session are filled in when the caller left them empty.
//
// It never panics and never returns an error: a hook is a bystander's
// convenience, and the agent it is watching does not stop because the
// bystander fell over.
func (h *Hooks) Fire(ctx context.Context, p Payload) {
	if h == nil || h.set == nil {
		return
	}
	h.set.mu.Lock()
	handlers := h.set.handlers
	onError := h.set.onError
	h.set.mu.Unlock()
	if len(handlers) == 0 {
		return
	}

	if p.Time.IsZero() {
		p.Time = time.Now()
	}
	if p.Session == "" {
		p.Session = h.session
	}
	if ctx == nil {
		ctx = context.Background()
	}

	for _, handler := range handlers {
		if err := safeHandle(ctx, handler, p); err != nil && onError != nil {
			onError(p, err)
		}
	}
}

// safeHandle turns a panicking handler into an ordinary error. A hook is
// third-party code as far as the loop is concerned — the config points it at
// an arbitrary command — and the agent's goroutine is not a good place to find
// out that it was wrong.
func safeHandle(ctx context.Context, handler Handler, p Payload) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("hook handler panicked: %v", r)
		}
	}()
	return handler.Handle(ctx, p)
}

// Close closes every handler that has a Close method, once, and forgets them.
// It is how a reporter gets to say goodbye — releasing a claim, flushing a
// queue — and is safe to call twice and on a nil *Hooks.
func (h *Hooks) Close() error {
	if h == nil || h.set == nil {
		return nil
	}
	h.set.mu.Lock()
	handlers := h.set.handlers
	h.set.handlers = nil
	onError := h.set.onError
	h.set.mu.Unlock()

	var first error
	for _, handler := range handlers {
		c, ok := handler.(interface{ Close() error })
		if !ok {
			continue
		}
		if err := c.Close(); err != nil {
			if onError != nil {
				onError(Payload{Event: SessionEnd}, err)
			}
			if first == nil {
				first = err
			}
		}
	}
	return first
}