Vectrum
Sign In
240 lines · 6.6 KB · Go
package agent

import (
	"bytes"
	"context"
	"errors"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"

	"github.com/johannSo/codey/internal/hooks"
)

// recorder keeps every payload fired at it. It locks because Paused fires from
// the agent's goroutine while the test reads from its own.
type recorder struct {
	mu       sync.Mutex
	payloads []hooks.Payload
}

func (r *recorder) Handle(_ context.Context, p hooks.Payload) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.payloads = append(r.payloads, p)
	return nil
}

func (r *recorder) events() []hooks.Event {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := make([]hooks.Event, 0, len(r.payloads))
	for _, p := range r.payloads {
		out = append(out, p.Event)
	}
	return out
}

func (r *recorder) find(ev hooks.Event) (hooks.Payload, bool) {
	r.mu.Lock()
	defer r.mu.Unlock()
	for _, p := range r.payloads {
		if p.Event == ev {
			return p, true
		}
	}
	return hooks.Payload{}, false
}

// recordHooks attaches a recorder to the agent and returns it.
func recordHooks(a *Agent) *recorder {
	rec := &recorder{}
	a.Hooks = hooks.New("s1")
	a.Hooks.Add(rec)
	return rec
}

func sameEvents(got, want []hooks.Event) bool {
	if len(got) != len(want) {
		return false
	}
	for i := range want {
		if got[i] != want[i] {
			return false
		}
	}
	return true
}

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

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

	want := []hooks.Event{hooks.RunStart, hooks.TurnStart, hooks.ToolStart, hooks.ToolResult, hooks.TurnStart, hooks.RunEnd}
	if got := rec.events(); !sameEvents(got, want) {
		t.Errorf("events = %v, want %v", got, want)
	}
	if p, _ := rec.find(hooks.ToolStart); p.Tool != "read_file" || p.Turn != 1 {
		t.Errorf("ToolStart = %+v, want read_file on turn 1", p)
	}
	// The session id is what lets one handler tell two tabs apart, and the
	// loop never sets it itself — the Hooks it fires on must.
	if p, _ := rec.find(hooks.RunStart); p.Session != "s1" {
		t.Errorf("RunStart.Session = %q, want the id the hooks were built with", p.Session)
	}
}

// Whether a call failed is the one thing about a tool result a bystander gets
// to know; what it returned is not.
func TestToolResultReportsFailureWithoutTheOutput(t *testing.T) {
	root := t.TempDir()
	os.WriteFile(filepath.Join(root, "a.go"), []byte("package a\n"), 0o644)
	tests := []struct {
		name    string
		call    string
		tool    string
		wantErr string
	}{
		{name: "success leaves Err empty", call: toolUseBlock("tu_1", "read_file", `{"path":"a.go"}`), tool: "read_file"},
		{name: "failure carries the error", call: toolUseBlock("tu_1", "boom", `{}`), tool: "boom", wantErr: "detonated"},
	}
	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, root, &bytes.Buffer{})
			rec := recordHooks(a)
			if _, err := a.Run(context.Background(), "go"); err != nil {
				t.Fatalf("Run: %v", err)
			}

			p, ok := rec.find(hooks.ToolResult)
			if !ok {
				t.Fatalf("no ToolResult among %v", rec.events())
			}
			if p.Tool != tc.tool {
				t.Errorf("Tool = %q, want %q", p.Tool, tc.tool)
			}
			if p.Err != tc.wantErr {
				t.Errorf("Err = %q, want %q", p.Err, tc.wantErr)
			}
			if p.Message != "" {
				t.Errorf("Message = %q, want nothing: a tool's output is not a bystander's business", p.Message)
			}
		})
	}
}

// A run that ends badly is still over. The sidebar that shows an agent as
// working because its run never said otherwise is the bug this guards.
func TestRunEndFiresOnEveryWayOut(t *testing.T) {
	t.Run("turn cap", func(t *testing.T) {
		f := newFakeAPI(t,
			msg("tool_use", toolUseBlock("tu", "boom", `{}`)),
			msg("tool_use", toolUseBlock("tu", "boom", `{}`)),
		)
		a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})
		a.opts.MaxTurns = 2
		rec := recordHooks(a)

		if _, err := a.Run(context.Background(), "go"); !errors.Is(err, ErrMaxTurns) {
			t.Fatalf("error = %v, want ErrMaxTurns", err)
		}
		got := rec.events()
		if len(got) == 0 || got[len(got)-1] != hooks.RunEnd {
			t.Errorf("events = %v, want RunEnd last", got)
		}
	})

	t.Run("cancelled while paused", func(t *testing.T) {
		f := newPausableAPI(t, msg("end_turn", textBlock("hi")))
		a := newPausableAgent(t, f, t.TempDir())
		a.Pause()
		rec := recordHooks(a)

		ctx, cancel := context.WithCancel(context.Background())
		errc := make(chan error, 1)
		go func() {
			_, err := a.Run(ctx, "go")
			errc <- err
		}()
		time.Sleep(settle)
		cancel()
		select {
		case err := <-errc:
			if !errors.Is(err, context.Canceled) {
				t.Fatalf("Run returned %v, want context.Canceled", err)
			}
		case <-time.After(2 * time.Second):
			t.Fatal("a cancelled pause never unblocked")
		}

		// The pause was never released, so Resumed must not appear: a watcher
		// told the agent resumed would look for a turn that never came.
		want := []hooks.Event{hooks.RunStart, hooks.Paused, hooks.RunEnd}
		if got := rec.events(); !sameEvents(got, want) {
			t.Errorf("events = %v, want %v", got, want)
		}
	})
}

func TestPauseAndResumeAreAnnounced(t *testing.T) {
	f := newPausableAPI(t, msg("end_turn", textBlock("done")))
	a := newPausableAgent(t, f, t.TempDir())
	a.Pause()
	rec := recordHooks(a)

	done := make(chan struct{})
	go func() {
		if _, err := a.Run(context.Background(), "go"); err != nil {
			t.Errorf("Run: %v", err)
		}
		close(done)
	}()
	time.Sleep(settle)
	a.Resume()
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("the run never finished after Resume")
	}

	want := []hooks.Event{hooks.RunStart, hooks.Paused, hooks.Resumed, hooks.TurnStart, hooks.RunEnd}
	if got := rec.events(); !sameEvents(got, want) {
		t.Errorf("events = %v, want %v", got, want)
	}
	if p, _ := rec.find(hooks.Paused); p.Turn != 1 {
		t.Errorf("Paused.Turn = %d, want the turn being held", p.Turn)
	}
}

// Most sessions have nothing watching from outside, and the loop must not
// make them pay for it with a nil check at every fire point — or a panic.
func TestNilHooksRunsNormally(t *testing.T) {
	f := newFakeAPI(t,
		msg("tool_use", toolUseBlock("tu_1", "boom", `{}`)),
		msg("end_turn", textBlock("fine")),
	)
	a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})
	if a.Hooks != nil {
		t.Fatal("a fresh agent should have no hooks")
	}
	got, err := a.Run(context.Background(), "go")
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if got != "fine" {
		t.Errorf("answer = %q", got)
	}
}