Vectrum
Sign In
295 lines · 9.9 KB · Go
package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"harness/internal/llm"
)

func TestCacheBreakpoints(t *testing.T) {
	f := newFakeAPI(t, msg("end_turn", textBlock("hi")))
	a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})
	if _, err := a.Run(context.Background(), "hello"); err != nil {
		t.Fatalf("Run: %v", err)
	}
	req := f.requests[0]
	if req.System[0].CacheControl == nil {
		t.Error("the system prompt should carry a breakpoint: it and the tools never change")
	}
	last := req.Messages[len(req.Messages)-1]
	if last.Content[len(last.Content)-1].CacheControl == nil {
		t.Error("the end of the conversation should carry a rolling breakpoint")
	}
	raw := f.rawRequest(0)
	if n := strings.Count(raw, `"cache_control"`); n != 2 {
		t.Errorf("%d breakpoints in the body, want 2 (the limit is 4):\n%s", n, raw)
	}
}

func TestNoCacheOption(t *testing.T) {
	f := newFakeAPI(t, msg("end_turn", textBlock("hi")))
	a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})
	a.opts.NoCache = true
	if _, err := a.Run(context.Background(), "hello"); err != nil {
		t.Fatalf("Run: %v", err)
	}
	if strings.Contains(f.rawRequest(0), "cache_control") {
		t.Error("NoCache should send no breakpoints")
	}
}

// Breakpoints must not mutate the stored history, or every turn would
// accumulate another one and blow past the four-breakpoint limit.
func TestBreakpointsDoNotLeakIntoHistory(t *testing.T) {
	f := newFakeAPI(t,
		msg("tool_use", toolUseBlock("tu_1", "boom", `{}`)),
		msg("tool_use", toolUseBlock("tu_2", "boom", `{}`)),
		msg("end_turn", textBlock("done")),
	)
	a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})
	if _, err := a.Run(context.Background(), "go"); err != nil {
		t.Fatalf("Run: %v", err)
	}
	for i := range f.requests {
		if n := strings.Count(f.rawRequest(i), `"cache_control"`); n != 2 {
			t.Errorf("request %d has %d breakpoints, want 2", i, n)
		}
	}
	for _, m := range a.State().History() {
		for _, b := range m.Content {
			if b.CacheControl != nil {
				t.Error("stored history should stay free of breakpoints")
			}
		}
	}
}

func TestCompactionSplitLandsOnAnAssistantMessage(t *testing.T) {
	mk := func(roles string) []llm.Message {
		var out []llm.Message
		for _, r := range roles {
			role := llm.RoleUser
			if r == 'a' {
				role = llm.RoleAssistant
			}
			out = append(out, llm.Message{Role: role, Content: []llm.ContentBlock{llm.TextBlock("x")}})
		}
		return out
	}
	tests := []struct {
		name  string
		roles string
		want  int
	}{
		{name: "too short to compact", roles: "uaua", want: -1},
		{name: "cuts at the first assistant past the midpoint", roles: "uauauauauaua", want: 7},
		// The tail can end up one message shorter than the target when the
		// boundary lands on a user message: correctness beats the constant.
		{name: "skips forward past a user message", roles: "uauauauaua", want: 5},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			got := compactionSplit(mk(tc.roles))
			if got != tc.want {
				t.Fatalf("compactionSplit = %d, want %d", got, tc.want)
			}
			if got >= 0 && mk(tc.roles)[got].Role != llm.RoleAssistant {
				t.Error("the kept tail must start with an assistant message, or its tool_results reference nothing")
			}
		})
	}
}

func TestCompactionFoldsHistoryAndKeepsTheTask(t *testing.T) {
	// Six tool-calling turns, then a compaction, then the answer.
	var script []string
	for i := 0; i < 6; i++ {
		script = append(script, msgWithTokens("tool_use", 100, toolUseBlock("tu", "boom", `{}`)))
	}
	script = append(script, msgWithTokens("end_turn", 100, textBlock("all done")))
	f := newFakeAPI(t, script...)

	var logBuf bytes.Buffer
	a := newTestAgent(t, f, t.TempDir(), &logBuf)
	a.opts.ContextWindow = 1000 // 70% of 1000, and each fake turn reports 100
	a.opts.CompactAt = 0.1

	var compacted int
	a.Observe = func(ev Event) {
		if ev.Kind == EventCompact {
			compacted++
		}
	}
	if _, err := a.Run(context.Background(), "fix the failing test in loop.go"); err != nil {
		t.Fatalf("Run: %v", err)
	}
	if compacted == 0 {
		t.Fatal("expected a compaction")
	}

	history := a.State().History()
	if history[0].Role != llm.RoleUser {
		t.Fatal("history must start with a user message")
	}
	folded := history[0].Content[0].Text
	for _, want := range []string{"[Earlier conversation, compacted]", "fix the failing test in loop.go", "the config loads from X not Y"} {
		if !strings.Contains(folded, want) {
			t.Errorf("folded message missing %q:\n%s", want, folded)
		}
	}
	if history[1].Role != llm.RoleAssistant {
		t.Errorf("message after the summary is %q, want assistant", history[1].Role)
	}
	if !strings.Contains(logBuf.String(), `"type":"compaction"`) {
		t.Error("compaction should be logged; it is a cost and cache event")
	}

	// The summarization call must not be able to touch the repository.
	for i, req := range f.requests {
		if len(req.Tools) == 0 && len(req.Messages) > 0 {
			if req.MaxTokens != 4096 {
				t.Errorf("request %d looks like the compaction call but has max_tokens %d", i, req.MaxTokens)
			}
		}
	}
}

// TestSummarizeCarriesSystemPrefixWhenSet guards the compaction bug: the
// summarization request is built by hand, outside (*Agent).request, so it is
// the one path that can silently omit the OAuth identity block. An omitted
// block does not 401 — it 429s, which the client retries three times before
// giving up, so a long OAuth session would die at the compaction threshold
// looking like a quota problem.
func TestSummarizeCarriesSystemPrefixWhenSet(t *testing.T) {
	f := newFakeAPI(t)
	a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})
	a.opts.SystemPrefix = ClaudeCodeIdentity

	history := []llm.Message{{Role: llm.RoleUser, Content: []llm.ContentBlock{llm.TextBlock("hello")}}}
	if _, err := a.summarize(context.Background(), history); err != nil {
		t.Fatalf("summarize: %v", err)
	}
	if len(f.requests) != 1 {
		t.Fatalf("expected 1 request, got %d", len(f.requests))
	}
	got := f.requests[0].System
	if len(got) != 1 || got[0].Text != ClaudeCodeIdentity {
		t.Errorf("System = %+v, want a single block with the Claude Code identity", got)
	}
}

// TestSummarizeCarriesNoSystemWhenPrefixUnset is the other half: a session
// that never borrowed the Claude Code login must not gain a system block it
// never asked for.
func TestSummarizeCarriesNoSystemWhenPrefixUnset(t *testing.T) {
	f := newFakeAPI(t)
	a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})

	history := []llm.Message{{Role: llm.RoleUser, Content: []llm.ContentBlock{llm.TextBlock("hello")}}}
	if _, err := a.summarize(context.Background(), history); err != nil {
		t.Fatalf("summarize: %v", err)
	}
	if len(f.requests) != 1 {
		t.Fatalf("expected 1 request, got %d", len(f.requests))
	}
	if got := f.requests[0].System; len(got) != 0 {
		t.Errorf("System = %+v, want none when SystemPrefix is unset", got)
	}
}

func TestStripTrailingToolUse(t *testing.T) {
	msgs := []llm.Message{
		{Role: llm.RoleUser, Content: []llm.ContentBlock{llm.TextBlock("go")}},
		{Role: llm.RoleAssistant, Content: []llm.ContentBlock{{Type: llm.BlockToolUse, ID: "tu", Name: "x"}}},
	}
	if got := stripTrailingToolUse(msgs); len(got) != 1 {
		t.Errorf("a dangling tool_use must be dropped, got %d messages", len(got))
	}
	keep := []llm.Message{
		{Role: llm.RoleUser, Content: []llm.ContentBlock{llm.TextBlock("go")}},
		{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock("done")}},
	}
	if got := stripTrailingToolUse(keep); len(got) != 2 {
		t.Errorf("a text-only assistant turn should be kept, got %d", len(got))
	}
}

func TestSupersedeStaleReads(t *testing.T) {
	root := t.TempDir()
	path := filepath.Join(root, "a.go")
	os.WriteFile(path, []byte("package a\n\nconst X = 1\n"), 0o644)

	f := newFakeAPI(t,
		msg("tool_use", toolUseBlock("tu_1", "read_file", `{"path":"a.go"}`)),
		msg("tool_use", toolUseBlock("tu_2", "read_file", `{"path":"a.go"}`)),
		msg("tool_use", toolUseBlock("tu_3", "edit_file", `{"path":"a.go","old_str":"const X = 1","new_str":"const X = 2"}`)),
		msg("end_turn", textBlock("done")),
	)
	var logBuf bytes.Buffer
	a := newTestAgent(t, f, root, &logBuf)
	if _, err := a.Run(context.Background(), "bump X"); err != nil {
		t.Fatalf("Run: %v", err)
	}

	var superseded, intact int
	for _, m := range a.State().History() {
		for _, b := range m.Content {
			if b.Type != llm.BlockToolResult {
				continue
			}
			switch {
			case strings.HasPrefix(b.Content, "[superseded:"):
				superseded++
			// Match the numbered read format, so the edit's own diff — which
			// quotes the old line — is not counted as a surviving read.
			case strings.Contains(b.Content, "\tconst X = 1"):
				intact++
			}
		}
	}
	if superseded != 1 {
		t.Errorf("%d results superseded, want 1 (the earlier of two reads)", superseded)
	}
	if intact != 1 {
		t.Errorf("%d intact reads, want 1 (the most recent read stays)", intact)
	}
	if !strings.Contains(logBuf.String(), `"type":"supersede_reads"`) {
		t.Error("superseding should be logged")
	}
}

func TestSupersedeLeavesASingleReadAlone(t *testing.T) {
	root := t.TempDir()
	os.WriteFile(filepath.Join(root, "a.go"), []byte("package a\n\nconst X = 1\n"), 0o644)
	f := newFakeAPI(t,
		msg("tool_use", toolUseBlock("tu_1", "read_file", `{"path":"a.go"}`)),
		msg("tool_use", toolUseBlock("tu_2", "edit_file", `{"path":"a.go","old_str":"const X = 1","new_str":"const X = 2"}`)),
		msg("end_turn", textBlock("done")),
	)
	a := newTestAgent(t, f, root, &bytes.Buffer{})
	if _, err := a.Run(context.Background(), "bump X"); err != nil {
		t.Fatalf("Run: %v", err)
	}
	for _, m := range a.State().History() {
		for _, b := range m.Content {
			if strings.HasPrefix(b.Content, "[superseded:") {
				t.Error("the only read of a file is still the current one; it must not be blanked")
			}
		}
	}
}

func TestReadPathParsing(t *testing.T) {
	if got := readPath(json.RawMessage(`{"path":"x/y.go","offset":3}`)); got != "x/y.go" {
		t.Errorf("readPath = %q", got)
	}
	if got := readPath(json.RawMessage(`not json`)); got != "" {
		t.Errorf("readPath = %q, want empty on malformed input", got)
	}
}