Vectrum
Sign In
187 lines · 6.1 KB · Go
package agent

import (
	"context"
	"fmt"
	"strings"

	"harness/internal/llm"
	"harness/internal/metrics"
)

// metricsCompaction records a compaction so its effect on later turns is
// visible in the event log.
func metricsCompaction(inputTokens, before, after int) metrics.Event {
	return metrics.Event{
		Type:        "compaction",
		InputTokens: inputTokens,
		Extra:       map[string]any{"messages_before": before, "messages_after": after},
	}
}

// DefaultContextWindow is the Claude Opus 5 window. Lower it with
// Options.ContextWindow when running a model with less room.
const DefaultContextWindow = 1_000_000

// DefaultCompactAt is the fraction of the window at which to summarize.
// Compaction itself needs headroom — it is a model call with the history in it
// — so waiting until the window is nearly full means the compaction request is
// the one that fails.
const DefaultCompactAt = 0.7

// keepRecentMessages is how much of the tail survives compaction verbatim. The
// model is usually mid-thought, and a summary of the last two minutes of work
// is worse than the work itself.
const keepRecentMessages = 6

const compactPrompt = `Summarize the conversation above for a coding agent that will keep working on the same task with this summary in place of the transcript.

Preserve, in this order:
1. The original task, verbatim enough that nothing about it is lost.
2. Every file touched so far and what was changed in it, with paths.
3. Decisions made and the reasoning that is not re-derivable from the code.
4. Things discovered the hard way: where something actually lives, what a test really asserts, what did not work and why.
5. What is left to do.

Be specific. Names, paths and line numbers are worth more than prose. Omit
pleasantries and anything already visible in the current state of the files.`

// maybeCompact summarizes the older half of the conversation when the last
// request came close to filling the window. It reports whether it compacted.
func (a *Agent) maybeCompact(ctx context.Context, lastInputTokens int) (bool, error) {
	window := a.opts.ContextWindow
	if window <= 0 {
		window = DefaultContextWindow
	}
	at := a.opts.CompactAt
	if at <= 0 {
		at = DefaultCompactAt
	}
	if float64(lastInputTokens) < at*float64(window) {
		return false, nil
	}

	history := a.state.History()
	split := compactionSplit(history)
	if split < 0 {
		return false, nil // nothing safe to fold away yet
	}

	summary, err := a.summarize(ctx, history[:split])
	if err != nil {
		return false, err
	}

	// The task itself is the one thing never summarized: everything downstream
	// is judged against it.
	task := originalTask(history)
	folded := llm.Message{Role: llm.RoleUser, Content: []llm.ContentBlock{
		llm.TextBlock(fmt.Sprintf("[Earlier conversation, compacted]\n\nOriginal task:\n%s\n\nWhat has happened since:\n%s", task, summary)),
	}}

	a.state.Replace(append([]llm.Message{folded}, history[split:]...))
	a.log.Log(metricsCompaction(lastInputTokens, len(history), a.state.Len()))
	a.emit(Event{Kind: EventCompact, Text: summary})
	return true, nil
}

// compactionSplit picks where to cut. The kept tail must begin with an
// assistant message: the folded summary takes the user slot before it, and a
// tail starting with a tool_result user message would reference a tool_use that
// no longer exists, which the API rejects.
func compactionSplit(history []llm.Message) int {
	if len(history) <= keepRecentMessages+2 {
		return -1
	}
	want := len(history) / 2
	if max := len(history) - keepRecentMessages; want > max {
		want = max
	}
	for i := want; i < len(history); i++ {
		if history[i].Role == llm.RoleAssistant {
			return i
		}
	}
	return -1
}

// summarize asks the model to fold the given messages into prose. It runs
// without tools: this call must not be able to touch the repository.
func (a *Agent) summarize(ctx context.Context, history []llm.Message) (string, error) {
	msgs := append([]llm.Message(nil), history...)
	msgs = append(msgs, llm.Message{
		Role:    llm.RoleUser,
		Content: []llm.ContentBlock{llm.TextBlock(compactPrompt)},
	})

	req := llm.Request{
		MaxTokens: 4096,
		Messages:  stripTrailingToolUse(msgs),
		Thinking:  llm.Adaptive(),
	}
	// This is the one request in the agent that is built by hand rather than
	// through (*Agent).request/systemBlocks, and it is deliberately bare — no
	// tools, no harness prompt — so it is easy to forget that an OAuth
	// credential still requires the Claude Code identity string as the FIRST
	// system block. Skip it here and the API does not 401; it returns a 429
	// that (*APIError).Retryable treats as transient, so the request burns
	// three attempts with backoff before the whole turn aborts. Do not "clean
	// up" this lone system block just because there is no other system prompt
	// alongside it — for an OAuth session, this is the whole system prompt.
	if a.opts.SystemPrefix != "" {
		req.System = []llm.SystemBlock{{Type: "text", Text: a.opts.SystemPrefix}}
	}

	resp, err := a.client.CreateMessage(ctx, req)
	if err != nil {
		return "", fmt.Errorf("compaction: %w", err)
	}
	text := strings.TrimSpace(resp.Text())
	if text == "" {
		return "", fmt.Errorf("compaction produced no summary")
	}
	return text, nil
}

// stripTrailingToolUse drops a dangling assistant tool_use at the end, which
// would otherwise demand tool_result blocks the summarization request has no
// way to supply.
func stripTrailingToolUse(msgs []llm.Message) []llm.Message {
	for len(msgs) > 0 {
		last := msgs[len(msgs)-1]
		if last.Role != llm.RoleAssistant {
			break
		}
		var hasToolUse bool
		for _, b := range last.Content {
			if b.Type == llm.BlockToolUse {
				hasToolUse = true
				break
			}
		}
		if !hasToolUse {
			break
		}
		msgs = msgs[:len(msgs)-1]
	}
	return msgs
}

// originalTask returns the first user message's text.
func originalTask(history []llm.Message) string {
	for _, m := range history {
		if m.Role != llm.RoleUser {
			continue
		}
		var b strings.Builder
		for _, blk := range m.Content {
			if blk.Type == llm.BlockText {
				b.WriteString(blk.Text)
			}
		}
		if s := strings.TrimSpace(b.String()); s != "" {
			return s
		}
	}
	return "(unrecorded)"
}