Vectrum
Sign In
304 lines · 9.7 KB · Go
package agent

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

	"harness/internal/llm"
	"harness/internal/metrics"
	"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.`

// 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
}

// 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
}

// 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)

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

// 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
	}
	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.tools.Execute(ctx, call.Name, call.Input)
		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
}

// 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 a.opts.Effort != "" {
		req.OutputConfig = &llm.OutputConfig{Effort: a.opts.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())
}