Vectrum
Sign In
147 lines · 4.0 KB · Go
// Package ui renders agent events to a terminal. The dependency runs one way:
// ui imports agent, never the reverse, so the eval runner can drive the loop
// with no terminal attached.
package ui

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"sort"
	"strings"

	"codey/internal/agent"
	"codey/internal/llm"
)

// ANSI codes, blanked when the output is not a terminal or NO_COLOR is set.
type palette struct{ dim, bold, red, cyan, reset string }

var color = func() palette {
	info, err := os.Stdout.Stat()
	interactive := err == nil && info.Mode()&os.ModeCharDevice != 0
	if !interactive || os.Getenv("NO_COLOR") != "" {
		return palette{}
	}
	return palette{dim: "\x1b[2m", bold: "\x1b[1m", red: "\x1b[31m", cyan: "\x1b[36m", reset: "\x1b[0m"}
}()

// Renderer prints the event stream.
type Renderer struct {
	W       io.Writer
	Verbose bool // print full tool results rather than a summary line

	midLine bool // streamed text is mid-line and needs a newline before anything else
}

// endLine closes a partially written line of streamed text.
func (r *Renderer) endLine() {
	if r.midLine {
		fmt.Fprintln(r.W)
		r.midLine = false
	}
}

// Flush closes any open line. Call it when a turn finishes.
func (r *Renderer) Flush() { r.endLine() }

func New(w io.Writer) *Renderer { return &Renderer{W: w} }

// Handle is the agent.Observe callback.
func (r *Renderer) Handle(ev agent.Event) {
	switch ev.Kind {
	case agent.EventText:
		fmt.Fprintf(r.W, "%s\n", strings.TrimSpace(ev.Text))
	case agent.EventTextDelta:
		// Streamed text arrives mid-line; the newline comes when the tool call
		// or the turn ends.
		fmt.Fprint(r.W, ev.Text)
		r.midLine = true
	case agent.EventToolStart:
		r.endLine()
		fmt.Fprintf(r.W, "%s⏺ %s%s%s(%s)%s\n", color.cyan, color.bold, ev.Tool, color.reset, summarizeInput(ev.Input), color.reset)
	case agent.EventToolResult:
		r.endLine()
		r.tool(ev)
	case agent.EventCompact:
		r.endLine()
		fmt.Fprintf(r.W, "%s[compacted the conversation]%s\n", color.dim, color.reset)
	case agent.EventNotice:
		r.endLine()
		fmt.Fprintf(r.W, "%s[%s]%s\n", color.dim, strings.TrimSpace(ev.Text), color.reset)
	}
}

func (r *Renderer) tool(ev agent.Event) {
	// A diff is the one tool result worth showing in full and in colour: it is
	// the record of what actually changed.
	if !ev.IsError && looksLikeDiff(ev.Result) {
		WriteDiff(r.W, ev.Result)
		return
	}
	prefix, c := "  ⎿ ", color.dim
	if ev.IsError {
		prefix, c = "  ⎿ error: ", color.red
	}
	body := strings.TrimRight(ev.Result, "\n")
	if !r.Verbose {
		body = firstLines(body, 4)
	}
	for i, line := range strings.Split(body, "\n") {
		if i == 0 {
			fmt.Fprintf(r.W, "%s%s%s%s\n", c, prefix, line, color.reset)
			continue
		}
		fmt.Fprintf(r.W, "%s     %s%s\n", c, line, color.reset)
	}
}

// summarizeInput renders tool arguments as k=v, sorted so the same call always
// prints the same way.
func summarizeInput(raw json.RawMessage) string {
	var m map[string]any
	if len(raw) == 0 || json.Unmarshal(raw, &m) != nil {
		return string(raw)
	}
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	parts := make([]string, 0, len(keys))
	for _, k := range keys {
		v := fmt.Sprintf("%v", m[k])
		if len(v) > 60 {
			v = v[:60] + "…"
		}
		parts = append(parts, fmt.Sprintf("%s=%s", k, v))
	}
	return strings.Join(parts, ", ")
}

func firstLines(s string, n int) string {
	lines := strings.Split(s, "\n")
	if len(lines) <= n {
		return s
	}
	return strings.Join(lines[:n], "\n") + fmt.Sprintf("\n… %d more lines", len(lines)-n)
}

// Usage prints the running token and cost totals for the session.
func (r *Renderer) Usage(u llm.Usage, turns int, model string) {
	r.endLine()
	line := fmt.Sprintf("%d turns · in %d (%d cached) · out %d",
		turns, u.InputTokens, u.CacheReadInputTokens, u.OutputTokens)
	if cost, known := u.Cost(model); known {
		line += fmt.Sprintf(" · $%.4f", cost)
	}
	fmt.Fprintf(r.W, "%s%s%s\n", color.dim, line, color.reset)
}

// Error prints a failure.
func (r *Renderer) Error(err error) {
	r.endLine()
	fmt.Fprintf(r.W, "%serror: %v%s\n", color.red, err, color.reset)
}