Vectrum
Sign In
354 lines · 9.6 KB · Go
package agent

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

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

// settle is how long a test waits before concluding that nothing is going to
// happen. A paused agent proves itself by silence, which can only be measured
// against a clock.
const settle = 150 * time.Millisecond

// pausableAPI announces each request on a channel so a test can tell the
// difference between "not sent yet" and "never sent".
type pausableAPI struct {
	mu        sync.Mutex
	responses []string
	bodies    [][]byte
	n         int
	sent      chan struct{}
	srv       *httptest.Server
}

func newPausableAPI(t *testing.T, responses ...string) *pausableAPI {
	t.Helper()
	f := &pausableAPI{responses: responses, sent: make(chan struct{}, len(responses)+4)}
	f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		f.mu.Lock()
		i := f.n
		f.n++
		f.bodies = append(f.bodies, body)
		f.mu.Unlock()
		f.sent <- struct{}{}
		if i >= len(f.responses) {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		io.WriteString(w, f.responses[i])
	}))
	t.Cleanup(f.srv.Close)
	return f
}

// lastRequest is the body of the most recent request, decoded.
func (f *pausableAPI) lastRequest(t *testing.T) llm.Request {
	t.Helper()
	f.mu.Lock()
	defer f.mu.Unlock()
	if len(f.bodies) == 0 {
		t.Fatal("no request was made")
	}
	var req llm.Request
	if err := json.Unmarshal(f.bodies[len(f.bodies)-1], &req); err != nil {
		t.Fatalf("request body does not decode: %v", err)
	}
	return req
}

func (f *pausableAPI) count() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.n
}

// waitForRequest blocks until one more request arrives, or fails the test.
func (f *pausableAPI) waitForRequest(t *testing.T) {
	t.Helper()
	select {
	case <-f.sent:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for a request")
	}
}

func newPausableAgent(t *testing.T, f *pausableAPI, root string) *Agent {
	t.Helper()
	log := metrics.NewWriter(&bytes.Buffer{}, "test")
	state := NewState()
	reg := tools.NewRegistry(log).Register(tools.NewRead(root, state))
	client := llm.New(llm.WithBaseURL(f.srv.URL), llm.WithAPIKey("k"))
	return New(client, reg, state, log, Options{})
}

// A paused agent sends nothing until it is resumed. This is the whole of what
// pause means: the conversation is intact, it is simply not being spent on.
func TestPauseHoldsTheNextRequestUntilResume(t *testing.T) {
	root := t.TempDir()
	os.WriteFile(filepath.Join(root, "a.go"), []byte("A\n"), 0o644)

	f := newPausableAPI(t,
		msg("tool_use", toolUseBlock("tu_1", "read_file", `{"path":"a.go"}`)),
		msg("end_turn", textBlock("done")),
	)
	a := newPausableAgent(t, f, root)

	paused := make(chan struct{}, 1)
	a.Observe = func(ev Event) {
		if ev.Kind == EventToolResult {
			a.Pause()
		}
		if ev.Kind == EventPaused {
			paused <- struct{}{}
		}
	}

	out := make(chan string, 1)
	go func() {
		text, err := a.Run(context.Background(), "go")
		if err != nil {
			t.Errorf("Run: %v", err)
		}
		out <- text
	}()

	f.waitForRequest(t)
	select {
	case <-paused:
	case <-time.After(2 * time.Second):
		t.Fatal("no EventPaused was emitted")
	}

	time.Sleep(settle)
	if n := f.count(); n != 1 {
		t.Fatalf("a paused agent sent %d requests, want 1", n)
	}

	a.Resume()
	select {
	case text := <-out:
		if text != "done" {
			t.Errorf("after resume the run returned %q, want %q", text, "done")
		}
	case <-time.After(2 * time.Second):
		t.Fatal("the run never finished after Resume")
	}
	if n := f.count(); n != 2 {
		t.Errorf("made %d requests overall, want 2", n)
	}
}

// Pausing mid-turn does not abandon the work in flight: the tool calls the
// model already asked for still run, and their results are still what the
// resumed request carries. Holding anywhere else would leave a tool_use with
// no tool_result, which is an API error, not a pause.
func TestPauseWaitsForTheTurnInFlightToFinish(t *testing.T) {
	root := t.TempDir()
	os.WriteFile(filepath.Join(root, "a.go"), []byte("A\n"), 0o644)

	f := newPausableAPI(t,
		msg("tool_use", toolUseBlock("tu_1", "read_file", `{"path":"a.go"}`)),
		msg("end_turn", textBlock("done")),
	)
	a := newPausableAgent(t, f, root)

	var order []EventKind
	var mu sync.Mutex
	a.Observe = func(ev Event) {
		mu.Lock()
		defer mu.Unlock()
		switch ev.Kind {
		case EventToolStart:
			a.Pause()
			order = append(order, ev.Kind)
		case EventToolResult, EventPaused, EventResumed:
			order = append(order, ev.Kind)
		}
	}

	done := make(chan struct{})
	go func() {
		if _, err := a.Run(context.Background(), "go"); err != nil {
			t.Errorf("Run: %v", err)
		}
		close(done)
	}()

	// Give the hold time to land, then release it.
	time.Sleep(settle)
	a.Resume()
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("the run never finished after Resume")
	}

	mu.Lock()
	defer mu.Unlock()
	want := []EventKind{EventToolStart, EventToolResult, EventPaused, EventResumed}
	if len(order) != len(want) {
		t.Fatalf("events = %v, want %v", order, want)
	}
	for i := range want {
		if order[i] != want[i] {
			t.Fatalf("events = %v, want %v", order, want)
		}
	}
}

// Killing a paused agent is the ordinary case — you pause one to look at it and
// then decide it is not worth continuing. It must not sit on the channel
// forever waiting for a Resume that is never coming.
func TestCancelWhilePausedEndsTheRun(t *testing.T) {
	f := newPausableAPI(t, msg("end_turn", textBlock("hi")))
	a := newPausableAgent(t, f, t.TempDir())
	a.Pause()

	ctx, cancel := context.WithCancel(context.Background())
	errc := make(chan error, 1)
	go func() {
		_, err := a.Run(ctx, "go")
		errc <- err
	}()

	time.Sleep(settle)
	if n := f.count(); n != 0 {
		t.Fatalf("an agent paused before its first turn sent %d requests", n)
	}
	cancel()

	select {
	case err := <-errc:
		if !errors.Is(err, context.Canceled) {
			t.Errorf("Run returned %v, want context.Canceled", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("a cancelled pause never unblocked")
	}
}

// A subagent's report reaches the parent as a user message the front end put
// into the conversation directly, not as something the human typed. Continue
// drives the loop on the history as it stands, so the report is one message
// and whatever the user says next is another.
func TestContinueRunsOnTheHistoryWithoutAddingAMessage(t *testing.T) {
	f := newPausableAPI(t, msg("end_turn", textBlock("acknowledged")))
	a := newPausableAgent(t, f, t.TempDir())

	a.State().Append(llm.Message{
		Role:    llm.RoleUser,
		Content: []llm.ContentBlock{llm.TextBlock("[subagent a1 finished] it is in loop.go")},
	})

	text, err := a.Continue(context.Background())
	if err != nil {
		t.Fatalf("Continue: %v", err)
	}
	if text != "acknowledged" {
		t.Errorf("Continue returned %q", text)
	}

	history := a.State().History()
	if len(history) != 2 {
		t.Fatalf("history has %d messages, want the report and the answer", len(history))
	}
	if history[0].Content[0].Text != "[subagent a1 finished] it is in loop.go" {
		t.Errorf("first message = %+v, want the report untouched", history[0])
	}
}

// Continuing an empty conversation would send a request with no messages,
// which the API rejects; saying so here is clearer than the error that comes
// back from the endpoint.
func TestContinueOnAnEmptyConversationIsRefused(t *testing.T) {
	f := newPausableAPI(t)
	a := newPausableAgent(t, f, t.TempDir())
	if _, err := a.Continue(context.Background()); err == nil {
		t.Fatal("continuing an empty conversation succeeded")
	}
	if f.count() != 0 {
		t.Error("it sent a request anyway")
	}
}

// A paused agent has not spent the turn it is holding before, so nothing
// should be counting it yet: the sidebar would show a turn that never
// happened.
func TestNoTurnIsAnnouncedUntilThePauseIsReleased(t *testing.T) {
	f := newPausableAPI(t, msg("end_turn", textBlock("done")))
	a := newPausableAgent(t, f, t.TempDir())
	a.Pause()

	turns := make(chan int, 8)
	a.Observe = func(ev Event) {
		if ev.Kind == EventTurnStart {
			turns <- ev.Turn
		}
	}

	go a.Run(context.Background(), "go")
	time.Sleep(settle)
	select {
	case turn := <-turns:
		t.Fatalf("turn %d was announced while the agent was still held", turn)
	default:
	}

	a.Resume()
	select {
	case <-turns:
	case <-time.After(2 * time.Second):
		t.Fatal("no turn was announced after the pause was released")
	}
}

// Effort is not a universal parameter. Sending the session's level to a model
// that rejects it fails the whole run with a 400 — which is how a Haiku
// subagent spawned from a high-effort session used to die before its first
// turn.
func TestEffortIsOmittedForAModelThatRejectsIt(t *testing.T) {
	cases := []struct {
		model, effort string
		wantEffort    string
	}{
		{"claude-opus-5", "high", "high"},
		{"claude-haiku-4-5", "high", ""},
		{"claude-opus-4-6", "xhigh", ""},
		{"claude-opus-4-6", "max", "max"},
	}
	for _, c := range cases {
		f := newPausableAPI(t, msg("end_turn", textBlock("ok")))
		a := newPausableAgent(t, f, t.TempDir())
		a.client.SetModel(c.model)
		a.SetEffort(c.effort)

		if _, err := a.Run(context.Background(), "go"); err != nil {
			t.Fatalf("%s: Run: %v", c.model, err)
		}

		req := f.lastRequest(t)
		switch {
		case c.wantEffort == "" && req.OutputConfig != nil:
			t.Errorf("%s at %s: sent output_config %+v, which the model rejects",
				c.model, c.effort, req.OutputConfig)
		case c.wantEffort != "" && (req.OutputConfig == nil || req.OutputConfig.Effort != c.wantEffort):
			t.Errorf("%s at %s: output_config = %+v, want effort %q",
				c.model, c.effort, req.OutputConfig, c.wantEffort)
		}
	}
}