Vectrum
Sign In
228 lines · 6.4 KB · Go
package hooks

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"time"
)

// DefaultTimeout bounds one hook command. A hook runs on the agent's own
// goroutine, between a turn and the next thing it was going to do, so a
// command that hangs is a session that hangs. Ten seconds is long enough for
// anything a notification needs to do and short enough that a wedged one is an
// annoyance rather than a lost afternoon.
const DefaultTimeout = 10 * time.Second

// Runner runs shell commands the user configured for an event. It is the
// escape hatch: the built-in handlers cover what Codey knows about, and this
// covers everything else, without anybody needing to write Go.
//
// Each command gets the event as environment variables and as JSON on stdin —
// variables because a one-line hook should be a one-liner, JSON because
// anything more than that wants a parser rather than string surgery.
type Runner struct {
	// Commands maps an event to the command lines fired for it, in order.
	// They run through the system shell, so a hook can be a pipeline.
	Commands map[Event][]string

	// Dir is the working directory commands run in. Empty means the process's.
	Dir string

	// Timeout bounds one command. Zero means DefaultTimeout.
	Timeout time.Duration

	// Env is added to the environment commands inherit, as "K=V" strings.
	Env []string

	// run executes one command line. It is a field so tests do not need a
	// shell, and nil means the real one.
	run func(ctx context.Context, line string, env []string, stdin []byte, dir string) error
}

// NewRunner returns a runner for a map of event name to command lines, as it
// is spelled in the config file. Unknown event names are returned as an error
// rather than dropped: a hook on a misspelt event is a hook that silently
// never fires, which is the worst way for this to fail.
func NewRunner(spec map[string][]string, dir string) (*Runner, error) {
	if len(spec) == 0 {
		return nil, nil
	}
	commands := make(map[Event][]string, len(spec))
	for name, lines := range spec {
		ev := Event(name)
		if !Known(ev) {
			return nil, fmt.Errorf("unknown hook event %q; known events are %s", name, KnownList())
		}
		var kept []string
		for _, line := range lines {
			if line != "" {
				kept = append(kept, line)
			}
		}
		if len(kept) > 0 {
			commands[ev] = kept
		}
	}
	if len(commands) == 0 {
		return nil, nil
	}
	return &Runner{Commands: commands, Dir: dir}, nil
}

// Known reports whether an event name is one Codey fires.
func Known(e Event) bool {
	for _, known := range All {
		if e == known {
			return true
		}
	}
	return false
}

// All is every event, in roughly the order a session produces them. It exists
// for config validation and for documentation that cannot drift from the code.
var All = []Event{
	SessionStart, SessionEnd,
	RunStart, RunEnd,
	TurnStart,
	ToolStart, ToolResult,
	Blocked, Unblocked,
	Paused, Resumed,
}

// KnownList is All as a comma-separated string, for error messages.
func KnownList() string {
	var b bytes.Buffer
	for i, e := range All {
		if i > 0 {
			b.WriteString(", ")
		}
		b.WriteString(string(e))
	}
	return b.String()
}

// Handle implements Handler.
func (r *Runner) Handle(ctx context.Context, p Payload) error {
	if r == nil {
		return nil
	}
	lines := r.Commands[p.Event]
	if len(lines) == 0 {
		return nil
	}

	stdin, err := json.Marshal(payloadJSON(p))
	if err != nil {
		return fmt.Errorf("encoding hook payload: %w", err)
	}
	env := append(os.Environ(), r.Env...)
	env = append(env, payloadEnv(p)...)

	timeout := r.Timeout
	if timeout <= 0 {
		timeout = DefaultTimeout
	}
	run := r.run
	if run == nil {
		run = runShell
	}

	// The event that says a run ended arrives with the run's context already
	// cancelled when the user interrupted it, and a hook that never fires on
	// an interrupt is a hook that cannot clean up after one. The deadline
	// below is the real bound here; cancellation of the turn is not.
	base := context.WithoutCancel(ctx)

	var first error
	for _, line := range lines {
		cmdCtx, cancel := context.WithTimeout(base, timeout)
		err := run(cmdCtx, line, env, stdin, r.Dir)
		cancel()
		if err != nil && first == nil {
			first = fmt.Errorf("hook %s: %q: %w", p.Event, line, err)
		}
	}
	return first
}

// runShell runs one command line through the system shell, with the event on
// stdin. Output is captured and reported only on failure: a hook is not
// allowed to write over the TUI.
func runShell(ctx context.Context, line string, env []string, stdin []byte, dir string) error {
	shell, flag := "/bin/sh", "-c"
	if runtime.GOOS == "windows" {
		shell, flag = "cmd", "/c"
	}
	cmd := exec.CommandContext(ctx, shell, flag, line)
	cmd.Env = env
	cmd.Dir = dir
	cmd.Stdin = bytes.NewReader(stdin)
	var out bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &out
	if err := cmd.Run(); err != nil {
		if trimmed := bytes.TrimSpace(out.Bytes()); len(trimmed) > 0 {
			return fmt.Errorf("%w: %s", err, truncate(string(trimmed), 500))
		}
		return err
	}
	return nil
}

// payloadEnv is the event as environment variables. Empty fields are still
// exported, so a hook script can read one without testing whether it is set.
func payloadEnv(p Payload) []string {
	return []string{
		"CODEY_HOOK_EVENT=" + string(p.Event),
		"CODEY_HOOK_TIME=" + p.Time.UTC().Format(time.RFC3339Nano),
		"CODEY_SESSION=" + p.Session,
		"CODEY_AGENT=" + p.Agent,
		"CODEY_TURN=" + strconv.Itoa(p.Turn),
		"CODEY_TOOL=" + p.Tool,
		"CODEY_REASON=" + p.Reason,
		"CODEY_MESSAGE=" + p.Message,
		"CODEY_ERROR=" + p.Err,
	}
}

// payloadJSON is the wire shape of an event, with snake_case names matching
// the environment variables and the event names themselves.
type jsonPayload struct {
	Event   string `json:"event"`
	Time    string `json:"time"`
	Session string `json:"session,omitempty"`
	Agent   string `json:"agent,omitempty"`
	Turn    int    `json:"turn,omitempty"`
	Tool    string `json:"tool,omitempty"`
	Reason  string `json:"reason,omitempty"`
	Message string `json:"message,omitempty"`
	Err     string `json:"error,omitempty"`
}

func payloadJSON(p Payload) jsonPayload {
	return jsonPayload{
		Event:   string(p.Event),
		Time:    p.Time.UTC().Format(time.RFC3339Nano),
		Session: p.Session,
		Agent:   p.Agent,
		Turn:    p.Turn,
		Tool:    p.Tool,
		Reason:  p.Reason,
		Message: p.Message,
		Err:     p.Err,
	}
}

func truncate(s string, max int) string {
	if len(s) <= max {
		return s
	}
	return s[:max] + "…"
}