Vectrum
Sign In
372 lines · 12.3 KB · Go
package agent

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"

	"harness/internal/llm"
	"harness/internal/metrics"
	"harness/internal/skills"
	"harness/internal/tools"
)

// DefaultMaxTurns caps one Run. An agent that loops without converging burns
// money quickly, so the cap is low and the message when it trips is explicit.
const DefaultMaxTurns = 30

// ErrMaxTurns is returned when the turn cap is hit.
var ErrMaxTurns = errors.New("turn limit reached")

// defaultSystem is deliberately static: it is the first thing in the cached
// prompt prefix, so a timestamp or session id here would cost a cache hit on
// every single request.
const defaultSystem = `You are a coding agent operating on a local repository through tools.

Never guess at file paths. If you do not know where something lives, find it
with glob or grep first — one search costs less than a run of wrong guesses,
and a wrong guess tells you nothing about where the file actually is. Once you
know the path, read the specific file or region you need rather than dumping
whole directories.

When you change something, verify it by running the project's own tests rather
than asserting that it works.

If a tool returns an error, read it: the error text usually says exactly what to
do differently. Do not repeat a failing call unchanged.

Be concise. When the task is done, say what you changed and why, briefly.`

// skillListing is appended to the system prompt when skills were discovered.
// It is constant for the session, so the cached prefix still matches request to
// request: what breaks caching is content that varies per request, not content
// that varies per session.
func skillListing(found []skills.Skill) string {
	if len(found) == 0 {
		return ""
	}
	var b strings.Builder
	b.WriteString("\n\nSkills are instruction sets for specific kinds of work. Invoke one with\nthe skill tool when its description matches what you are doing — before\nstarting, not after. A user message beginning with /<name> is a request to\ninvoke that skill; everything after the name is its arguments.\n\nAvailable skills:")
	for _, s := range found {
		fmt.Fprintf(&b, "\n  %s: %s", s.Name, s.Description)
	}
	return b.String()
}

// Options configures an Agent.
type Options struct {
	MaxTurns int
	System   string
	// Effort maps to output_config.effort. Empty means the API default (high).
	Effort string
	// ContextWindow and CompactAt control when the conversation gets folded
	// into a summary. Zero means the defaults.
	ContextWindow int
	CompactAt     float64
	// NoCache disables prompt cache breakpoints, for measuring what they buy.
	NoCache bool
	// Stream uses the SSE endpoint so output appears as it is produced.
	Stream bool
	// Skills are the instruction sets discovered at startup. Their names and
	// descriptions are appended to the system prompt; the bodies are loaded on
	// demand by the skill tool.
	Skills []skills.Skill
}

// EventKind identifies what happened, for an observer to render.
type EventKind string

const (
	EventTurnStart  EventKind = "turn_start"
	EventText       EventKind = "text"
	EventToolStart  EventKind = "tool_start"
	EventToolResult EventKind = "tool_result"
	EventUsage      EventKind = "usage"
	EventCompact    EventKind = "compact"
	// EventTextDelta is emitted instead of EventText while streaming.
	EventTextDelta EventKind = "text_delta"
)

// Event is what the loop reports to an observer. It carries plain data only —
// no terminal concerns — so the eval runner can consume the same stream.
type Event struct {
	Kind    EventKind
	Turn    int
	Text    string
	Tool    string
	Input   json.RawMessage
	Result  string
	IsError bool
	Usage   llm.Usage
}

// Approver decides whether a tool call may proceed. It is an interface here so
// the agent package stays free of any particular policy or interface.
type Approver interface {
	Approve(ctx context.Context, tool string, input json.RawMessage) error
}

// Agent runs the turn loop.
type Agent struct {
	client *llm.Client
	tools  *tools.Registry
	state  *State
	log    *metrics.Logger
	opts   Options

	// Observe is called for every event. It may be nil.
	Observe func(Event)

	// Approve gates mutating tool calls. Nil means everything is allowed,
	// which is only appropriate against a throwaway tree or in a sandbox.
	Approve Approver

	// lastInputTokens is the size of the previous request as the API measured
	// it, cached and uncached parts together.
	lastInputTokens int

	// mu guards the settings a front end can change while the loop is
	// running. Only effort is one today; the loop reads it once per request.
	mu sync.Mutex
}

// SetEffort changes the effort level the turns that follow will ask for. An
// empty level means the API's own default.
func (a *Agent) SetEffort(level string) {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.opts.Effort = level
}

// Effort reports the level currently in force.
func (a *Agent) Effort() string {
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.opts.Effort
}

// New wires an agent. state and log may be freshly created by the caller so the
// eval runner can keep one per run.
func New(client *llm.Client, reg *tools.Registry, state *State, log *metrics.Logger, opts Options) *Agent {
	if opts.MaxTurns <= 0 {
		opts.MaxTurns = DefaultMaxTurns
	}
	if opts.System == "" {
		opts.System = defaultSystem
	}
	opts.System += skillListing(opts.Skills)
	return &Agent{client: client, tools: reg, state: state, log: log, opts: opts}
}

// State exposes the conversation for callers that need to inspect or persist it.
func (a *Agent) State() *State { return a.state }

func (a *Agent) emit(ev Event) {
	if a.Observe != nil {
		a.Observe(ev)
	}
}

// Run sends one user message and drives the tool loop until the model stops
// asking for tools. It returns the final assistant text.
func (a *Agent) Run(ctx context.Context, prompt string) (string, error) {
	a.state.Append(llm.Message{Role: llm.RoleUser, Content: []llm.ContentBlock{llm.TextBlock(prompt)}})

	for turn := 1; turn <= a.opts.MaxTurns; turn++ {
		a.emit(Event{Kind: EventTurnStart, Turn: turn})

		// Compaction is checked before building the request, using the token
		// count the API reported last turn — free and exact, where counting
		// locally would be neither.
		if _, err := a.maybeCompact(ctx, a.lastInputTokens); err != nil {
			return "", err
		}

		resp, err := a.send(ctx, turn)
		if err != nil {
			return "", err
		}
		a.state.AddUsage(resp.Usage)
		a.lastInputTokens = resp.Usage.InputTokens + resp.Usage.CacheReadInputTokens + resp.Usage.CacheCreationInputTokens
		a.log.Log(metrics.Event{
			Type:             "turn",
			Turn:             turn,
			StopReason:       resp.StopReason,
			InputTokens:      resp.Usage.InputTokens,
			OutputTokens:     resp.Usage.OutputTokens,
			CacheReadTokens:  resp.Usage.CacheReadInputTokens,
			CacheWriteTokens: resp.Usage.CacheCreationInputTokens,
		})
		a.emit(Event{Kind: EventUsage, Turn: turn, Usage: resp.Usage})

		// The assistant message goes into history verbatim, thinking blocks and
		// all: their signatures must come back unchanged on the next request.
		a.state.Append(resp.Message())

		// While streaming the text has already gone out delta by delta.
		if text := resp.Text(); text != "" && !a.opts.Stream {
			a.emit(Event{Kind: EventText, Turn: turn, Text: text})
		}

		switch resp.StopReason {
		case llm.StopToolUse:
			results, err := a.dispatch(ctx, turn, resp.ToolUses())
			if err != nil {
				return "", err
			}
			// Every tool_result for this assistant turn goes back in ONE user
			// message. Splitting them across messages is an API error, and even
			// where it is not it teaches the model to stop calling tools in
			// parallel.
			a.state.Append(llm.Message{Role: llm.RoleUser, Content: results})

		case llm.StopPauseTurn:
			// A server-side tool needs another round trip; resend as-is.
			continue

		case llm.StopRefusal:
			return "", refusalError(resp)

		case llm.StopMaxTokens:
			return resp.Text(), fmt.Errorf("response hit max_tokens (%d) and was cut off", a.client.MaxTokens)

		default:
			return resp.Text(), nil
		}
	}

	return "", fmt.Errorf("%w: stopped after %d turns without finishing; narrow the task or raise -max-turns",
		ErrMaxTurns, a.opts.MaxTurns)
}

// send performs one model call, streaming when asked to.
func (a *Agent) send(ctx context.Context, turn int) (*llm.Response, error) {
	req := a.request()
	if !a.opts.Stream {
		return a.client.CreateMessage(ctx, req)
	}
	return a.client.StreamMessage(ctx, req, func(ev llm.StreamEvent) {
		if ev.Kind == llm.StreamText {
			a.emit(Event{Kind: EventTextDelta, Turn: turn, Text: ev.Text})
		}
	})
}

// dispatch runs each tool call in order and builds one tool_result per
// tool_use id. A tool that fails still produces a result — a missing one is an
// API error, and the model can usually recover from the error text.
func (a *Agent) dispatch(ctx context.Context, turn int, calls []llm.ContentBlock) ([]llm.ContentBlock, error) {
	results := make([]llm.ContentBlock, 0, len(calls))
	for _, call := range calls {
		a.emit(Event{Kind: EventToolStart, Turn: turn, Tool: call.Name, Input: call.Input})

		out, err := a.authorizedExecute(ctx, call)
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}

		isErr := err != nil
		payload := out
		if isErr {
			payload = err.Error()
		}
		// A file that was just changed makes every earlier read of it wrong,
		// not merely stale. Blank those results rather than leave the model
		// reasoning from contents that no longer exist.
		if !isErr && (call.Name == "edit_file" || call.Name == "write_file") {
			if p := editedPath(call.Input); p != "" {
				if n := a.state.SupersedeReads(p); n > 0 {
					a.log.Log(metrics.Event{Type: "supersede_reads", Tool: call.Name,
						Extra: map[string]any{"path": p, "results_blanked": n}})
				}
			}
		}

		results = append(results, llm.ToolResultBlock(call.ID, payload, isErr))
		a.emit(Event{Kind: EventToolResult, Turn: turn, Tool: call.Name, Result: payload, IsError: isErr})
	}
	return results, nil
}

// authorizedExecute runs a tool once the approval policy allows it. A denial
// comes back as an ordinary tool error: the model should read it and change
// course, not see a harness crash.
func (a *Agent) authorizedExecute(ctx context.Context, call llm.ContentBlock) (string, error) {
	if a.Approve != nil {
		if err := a.Approve.Approve(ctx, call.Name, call.Input); err != nil {
			if ctx.Err() != nil {
				return "", ctx.Err()
			}
			a.log.Log(metrics.Event{Type: "tool_denied", Tool: call.Name, Err: err.Error()})
			return "", err
		}
	}
	return a.tools.Execute(ctx, call.Name, call.Input)
}

// request builds the body. Order matters for prompt caching: system, then
// tools, then the growing conversation.
func (a *Agent) request() llm.Request {
	defs := make([]llm.ToolDef, 0, len(a.tools.List()))
	for _, t := range a.tools.List() {
		defs = append(defs, llm.ToolDef{
			Name:        t.Name(),
			Description: t.Description(),
			InputSchema: t.InputSchema(),
		})
	}
	system := llm.SystemBlock{Type: "text", Text: a.opts.System}
	messages := a.state.History()

	if !a.opts.NoCache {
		// Two breakpoints. The first caches the tools and system prompt, which
		// never change. The second rolls forward over the conversation each
		// turn, so every turn writes a small delta and reads a large prefix.
		system.CacheControl = llm.Ephemeral()
		if n := len(messages); n > 0 {
			if blocks := messages[n-1].Content; len(blocks) > 0 {
				messages[n-1].Content = append([]llm.ContentBlock(nil), blocks...)
				messages[n-1].Content[len(blocks)-1].CacheControl = llm.Ephemeral()
			}
		}
	}

	req := llm.Request{
		System:   []llm.SystemBlock{system},
		Messages: messages,
		Tools:    defs,
		Thinking: llm.Adaptive(),
	}
	if effort := a.Effort(); effort != "" {
		req.OutputConfig = &llm.OutputConfig{Effort: effort}
	}
	return req
}

// editedPath pulls the path out of an edit_file or write_file call.
func editedPath(input json.RawMessage) string {
	var args struct {
		Path string `json:"path"`
	}
	if json.Unmarshal(input, &args) != nil {
		return ""
	}
	return args.Path
}

func refusalError(resp *llm.Response) error {
	var b strings.Builder
	b.WriteString("the model declined this request")
	if resp.StopDetails != nil {
		if resp.StopDetails.Category != "" {
			fmt.Fprintf(&b, " (%s)", resp.StopDetails.Category)
		}
		if resp.StopDetails.Explanation != "" {
			fmt.Fprintf(&b, ": %s", resp.StopDetails.Explanation)
		}
	}
	return errors.New(b.String())
}