Vectrum
Sign In
360 lines · 12.1 KB · Go
package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

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

// fakeAPI serves a scripted sequence of responses and records every request
// body it received, so tests can assert on what the loop actually sent.
type fakeAPI struct {
	t         *testing.T
	responses []string
	requests  []llm.Request
	srv       *httptest.Server
}

func newFakeAPI(t *testing.T, responses ...string) *fakeAPI {
	f := &fakeAPI{t: t, responses: responses}
	f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		var req llm.Request
		if err := json.Unmarshal(body, &req); err != nil {
			t.Errorf("harness sent a body the API types cannot parse: %v\n%s", err, body)
		}
		f.requests = append(f.requests, req)
		i := len(f.requests) - 1
		if i >= len(f.responses) {
			t.Errorf("harness made %d requests, script only has %d", i+1, len(f.responses))
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		io.WriteString(w, f.responses[i])
	}))
	t.Cleanup(f.srv.Close)
	return f
}

// rawRequest returns the n-th request body re-encoded, for byte-level assertions.
func (f *fakeAPI) rawRequest(n int) string {
	b, err := json.Marshal(f.requests[n])
	if err != nil {
		f.t.Fatal(err)
	}
	return string(b)
}

func msg(stopReason string, blocks ...string) string {
	return `{"id":"msg","type":"message","role":"assistant","model":"claude-opus-5","content":[` +
		strings.Join(blocks, ",") + `],"stop_reason":"` + stopReason + `",` +
		`"usage":{"input_tokens":100,"output_tokens":20,"cache_read_input_tokens":40}}`
}

func textBlock(s string) string { return `{"type":"text","text":"` + s + `"}` }

func toolUseBlock(id, name, input string) string {
	return `{"type":"tool_use","id":"` + id + `","name":"` + name + `","input":` + input + `}`
}

// newTestAgent wires an agent against the fake API, with read_file rooted at
// a temp dir plus a couple of tools that make failure modes easy to trigger.
func newTestAgent(t *testing.T, f *fakeAPI, root string, logBuf *bytes.Buffer) *Agent {
	t.Helper()
	log := metrics.NewWriter(logBuf, "test")
	state := NewState()
	reg := tools.NewRegistry(log).Register(
		tools.NewRead(root, state),
		tools.New("boom", "Always fails", func(context.Context, struct{}) (string, error) {
			return "", errors.New("detonated")
		}),
	)
	client := llm.New(llm.WithBaseURL(f.srv.URL), llm.WithAPIKey("k"))
	return New(client, reg, state, log, Options{})
}

func TestRunReadsAFileAndAnswers(t *testing.T) {
	root := t.TempDir()
	os.WriteFile(filepath.Join(root, "main.go"), []byte("package main\n"), 0o644)

	f := newFakeAPI(t,
		msg("tool_use", textBlock("Let me look."), toolUseBlock("tu_1", "read_file", `{"path":"main.go"}`)),
		msg("end_turn", textBlock("It declares package main.")),
	)
	a := newTestAgent(t, f, root, &bytes.Buffer{})

	got, err := a.Run(context.Background(), "what does main.go do?")
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if got != "It declares package main." {
		t.Errorf("answer = %q", got)
	}
	if len(f.requests) != 2 {
		t.Fatalf("made %d requests, want 2", len(f.requests))
	}

	// Second request: user prompt, assistant turn, then the tool results.
	second := f.requests[1].Messages
	if len(second) != 3 {
		t.Fatalf("second request had %d messages, want 3:\n%s", len(second), f.rawRequest(1))
	}
	if second[2].Role != llm.RoleUser {
		t.Errorf("tool results must come back under the user role, got %q", second[2].Role)
	}
	if n := len(second[2].Content); n != 1 {
		t.Fatalf("want 1 tool_result block, got %d", n)
	}
	res := second[2].Content[0]
	if res.Type != llm.BlockToolResult || res.ToolUseID != "tu_1" || res.IsError {
		t.Errorf("tool_result = %+v, want a successful result for tu_1", res)
	}
	if !strings.Contains(res.Content, "package main") {
		t.Errorf("tool_result content = %q, want the file contents", res.Content)
	}
}

func TestParallelToolCallsReturnInOneUserMessage(t *testing.T) {
	root := t.TempDir()
	os.WriteFile(filepath.Join(root, "a.go"), []byte("A\n"), 0o644)
	os.WriteFile(filepath.Join(root, "b.go"), []byte("B\n"), 0o644)

	f := newFakeAPI(t,
		msg("tool_use",
			toolUseBlock("tu_1", "read_file", `{"path":"a.go"}`),
			toolUseBlock("tu_2", "read_file", `{"path":"b.go"}`)),
		msg("end_turn", textBlock("done")),
	)
	a := newTestAgent(t, f, root, &bytes.Buffer{})
	if _, err := a.Run(context.Background(), "read both"); err != nil {
		t.Fatalf("Run: %v", err)
	}

	last := f.requests[1].Messages[len(f.requests[1].Messages)-1]
	if len(last.Content) != 2 {
		t.Fatalf("want both results in ONE message, got %d messages ending with %d blocks",
			len(f.requests[1].Messages), len(last.Content))
	}
	if last.Content[0].ToolUseID != "tu_1" || last.Content[1].ToolUseID != "tu_2" {
		t.Errorf("results out of order: %s, %s", last.Content[0].ToolUseID, last.Content[1].ToolUseID)
	}
}

func TestFailedToolStillProducesAResult(t *testing.T) {
	tests := []struct {
		name     string
		call     string
		wantText string
	}{
		{name: "tool returns an error", call: toolUseBlock("tu_1", "boom", `{}`), wantText: "detonated"},
		{name: "unknown tool", call: toolUseBlock("tu_1", "telepathy", `{}`), wantText: `unknown tool "telepathy"`},
		{name: "malformed arguments", call: toolUseBlock("tu_1", "read_file", `{"path":42}`), wantText: "invalid arguments"},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			f := newFakeAPI(t, msg("tool_use", tc.call), msg("end_turn", textBlock("ok")))
			a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})
			if _, err := a.Run(context.Background(), "go"); err != nil {
				t.Fatalf("a failing tool is a normal event, not a Run error: %v", err)
			}
			last := f.requests[1].Messages[len(f.requests[1].Messages)-1]
			if len(last.Content) != 1 {
				t.Fatalf("every tool_use id needs exactly one result, got %d", len(last.Content))
			}
			res := last.Content[0]
			if !res.IsError {
				t.Error("is_error must be set so the model knows the call failed")
			}
			if !strings.Contains(res.Content, tc.wantText) {
				t.Errorf("result = %q, want it to contain %q", res.Content, tc.wantText)
			}
		})
	}
}

// Thinking signatures must be echoed back byte-identical or the next request
// is rejected.
func TestThinkingBlocksAreEchoedBackUnchanged(t *testing.T) {
	f := newFakeAPI(t,
		msg("tool_use",
			`{"type":"thinking","thinking":"","signature":"sig-xyz"}`,
			toolUseBlock("tu_1", "boom", `{}`)),
		msg("end_turn", textBlock("fine")),
	)
	a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})
	if _, err := a.Run(context.Background(), "go"); err != nil {
		t.Fatalf("Run: %v", err)
	}
	body := f.rawRequest(1)
	if !strings.Contains(body, `"signature":"sig-xyz"`) {
		t.Errorf("thinking signature was dropped from the resent history:\n%s", body)
	}
	if !strings.Contains(body, `"thinking":""`) {
		t.Errorf("empty thinking text was dropped:\n%s", body)
	}
}

func TestRequestShape(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 len(req.System) != 1 || req.System[0].Text == "" {
		t.Error("system prompt missing")
	}
	if strings.ContainsAny(req.System[0].Text, "0123456789") {
		t.Error("system prompt should hold no digits: timestamps or ids here would break the cached prefix")
	}
	if req.Thinking == nil || req.Thinking.Type != "adaptive" {
		t.Errorf("thinking = %+v, want adaptive (budget_tokens is a 400 on current models)", req.Thinking)
	}
	if len(req.Tools) != 2 || req.Tools[0].Name != "boom" || req.Tools[1].Name != "read_file" {
		t.Errorf("tools = %+v, want them sorted for a byte-stable prefix", req.Tools)
	}
	for _, tl := range req.Tools {
		if len(tl.InputSchema) == 0 {
			t.Errorf("tool %q has no input_schema", tl.Name)
		}
	}
}

func TestMaxTurnsStopsTheLoop(t *testing.T) {
	var script []string
	for i := 0; i < 5; i++ {
		script = append(script, msg("tool_use", toolUseBlock("tu", "boom", `{}`)))
	}
	f := newFakeAPI(t, script...)
	a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})
	a.opts.MaxTurns = 3

	_, err := a.Run(context.Background(), "loop forever")
	if !errors.Is(err, ErrMaxTurns) {
		t.Fatalf("error = %v, want ErrMaxTurns", err)
	}
	if !strings.Contains(err.Error(), "3 turns") {
		t.Errorf("error should say how many turns were spent: %v", err)
	}
	if len(f.requests) != 3 {
		t.Errorf("made %d requests, want exactly the cap of 3", len(f.requests))
	}
}

func TestRefusalIsReportedWithItsCategory(t *testing.T) {
	f := newFakeAPI(t, `{"id":"m","type":"message","role":"assistant","model":"claude-opus-5",
	"content":[],"stop_reason":"refusal","stop_details":{"type":"refusal","category":"cyber","explanation":"nope"},
	"usage":{"input_tokens":5,"output_tokens":0}}`)
	a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})
	_, err := a.Run(context.Background(), "do something")
	if err == nil {
		t.Fatal("want an error")
	}
	if !strings.Contains(err.Error(), "cyber") || !strings.Contains(err.Error(), "nope") {
		t.Errorf("error = %v, want the category and explanation", err)
	}
}

func TestPauseTurnResends(t *testing.T) {
	f := newFakeAPI(t,
		msg("pause_turn", textBlock("thinking")),
		msg("end_turn", textBlock("done")),
	)
	a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})
	got, err := a.Run(context.Background(), "go")
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if got != "done" || len(f.requests) != 2 {
		t.Errorf("got %q after %d requests, want the loop to resend after pause_turn", got, len(f.requests))
	}
}

func TestUsageAndEventsAreRecorded(t *testing.T) {
	root := t.TempDir()
	os.WriteFile(filepath.Join(root, "x.go"), []byte("x\n"), 0o644)
	f := newFakeAPI(t,
		msg("tool_use", toolUseBlock("tu_1", "read_file", `{"path":"x.go"}`)),
		msg("end_turn", textBlock("done")),
	)
	var logBuf bytes.Buffer
	a := newTestAgent(t, f, root, &logBuf)

	var kinds []EventKind
	a.Observe = func(ev Event) { kinds = append(kinds, ev.Kind) }
	if _, err := a.Run(context.Background(), "go"); err != nil {
		t.Fatalf("Run: %v", err)
	}

	usage, turns := a.State().Usage()
	if turns != 2 || usage.InputTokens != 200 || usage.CacheReadInputTokens != 80 {
		t.Errorf("usage = %+v over %d turns, want both turns accumulated", usage, turns)
	}
	if !containsKind(kinds, EventToolStart) || !containsKind(kinds, EventToolResult) || !containsKind(kinds, EventText) {
		t.Errorf("observer saw %v, want tool and text events", kinds)
	}

	log := logBuf.String()
	if n := strings.Count(log, `"type":"turn"`); n != 2 {
		t.Errorf("logged %d turn events, want 2:\n%s", n, log)
	}
	if n := strings.Count(log, `"type":"tool_call"`); n != 1 {
		t.Errorf("logged %d tool_call events, want 1:\n%s", n, log)
	}
	if !strings.Contains(log, `"cache_read_input_tokens":40`) {
		t.Errorf("cache reads must be logged; that is the Phase 4 lever:\n%s", log)
	}
}

func TestReadFileMarksStateSoEditCanCheckLater(t *testing.T) {
	root := t.TempDir()
	path := filepath.Join(root, "tracked.go")
	os.WriteFile(path, []byte("x\n"), 0o644)
	f := newFakeAPI(t,
		msg("tool_use", toolUseBlock("tu_1", "read_file", `{"path":"tracked.go"}`)),
		msg("end_turn", textBlock("done")),
	)
	a := newTestAgent(t, f, root, &bytes.Buffer{})
	if _, err := a.Run(context.Background(), "go"); err != nil {
		t.Fatalf("Run: %v", err)
	}
	if _, ok := a.State().LastRead(path); !ok {
		t.Error("read_file must record the read in State for Phase 2's edit_file")
	}
}

func TestRunPropagatesAPIErrors(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusBadRequest)
		io.WriteString(w, `{"error":{"type":"invalid_request_error","message":"bad model"}}`)
	}))
	defer srv.Close()

	state := NewState()
	a := New(llm.New(llm.WithBaseURL(srv.URL), llm.WithAPIKey("k")),
		tools.NewRegistry(nil), state, nil, Options{})
	if _, err := a.Run(context.Background(), "go"); err == nil || !strings.Contains(err.Error(), "bad model") {
		t.Errorf("error = %v, want the API message surfaced", err)
	}
}

func containsKind(kinds []EventKind, want EventKind) bool {
	for _, k := range kinds {
		if k == want {
			return true
		}
	}
	return false
}