Vectrum
Sign In
519 lines · 15.7 KB · Go
package subagent

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"sort"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/johannSo/codey/internal/agent"
	"github.com/johannSo/codey/internal/llm"
	"github.com/johannSo/codey/internal/metrics"
	"github.com/johannSo/codey/internal/tools"
)

// settle is how long a test waits before concluding nothing more is coming.
const settle = 150 * time.Millisecond

// fakeAPI answers every child's request. What it answers is derived from the
// task the child was given, so a test can tell two agents' traffic apart.
type fakeAPI struct {
	mu       sync.Mutex
	requests int
	models   []string
	// bodies is every request body, for asking what a child actually sent.
	bodies [][]byte
	// gate, when non-nil, holds every request until it is closed: the way a
	// test keeps an agent running while it inspects the supervisor.
	gate chan struct{}
	// status, when non-zero, is returned instead of a response.
	status int
	// twoTurns makes the first answer a tool call, so the agent has a second
	// turn — which is where a pause lands.
	twoTurns bool
	// callTool makes the first answer a call to a named tool, for asking what
	// happens when a child reaches for one it does not have.
	callTool string
	// loop makes every answer a tool call, so the child never finishes: what a
	// stuck agent looks like, and the only thing a turn cap is there for.
	loop bool
	srv  *httptest.Server
}

func newFakeAPI(t *testing.T) *fakeAPI {
	t.Helper()
	f := &fakeAPI{}
	f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		var req llm.Request
		json.Unmarshal(body, &req)

		f.mu.Lock()
		f.requests++
		f.models = append(f.models, req.Model)
		f.bodies = append(f.bodies, body)
		gate, status := f.gate, f.status
		f.mu.Unlock()

		if gate != nil {
			<-gate
		}
		if status != 0 {
			w.WriteHeader(status)
			io.WriteString(w, `{"type":"error","error":{"type":"overloaded_error","message":"nope"}}`)
			return
		}
		f.mu.Lock()
		two := f.twoTurns
		f.mu.Unlock()
		f.mu.Lock()
		named, loop := f.callTool, f.loop
		f.mu.Unlock()
		if loop {
			io.WriteString(w, toolUse("tu_1", "noop"))
			return
		}
		if len(req.Messages) == 1 {
			switch {
			case named != "":
				io.WriteString(w, toolUse("tu_1", named))
				return
			case two:
				io.WriteString(w, toolUse("tu_1", "noop"))
				return
			}
		}
		io.WriteString(w, endTurn("handled: "+task(req)))
	}))
	t.Cleanup(f.srv.Close)
	return f
}

func (f *fakeAPI) hold() { f.mu.Lock(); f.gate = make(chan struct{}); f.mu.Unlock() }
func (f *fakeAPI) release() {
	f.mu.Lock()
	gate := f.gate
	f.gate = nil
	f.mu.Unlock()
	if gate != nil {
		close(gate)
	}
}
func (f *fakeAPI) fail(status int) { f.mu.Lock(); f.status = status; f.mu.Unlock() }
func (f *fakeAPI) loopForever()    { f.mu.Lock(); f.loop = true; f.mu.Unlock() }
func (f *fakeAPI) count() int      { f.mu.Lock(); defer f.mu.Unlock(); return f.requests }

// requestsSeen decodes every request the fake was sent.
func (f *fakeAPI) requestsSeen(t *testing.T) []llm.Request {
	t.Helper()
	f.mu.Lock()
	defer f.mu.Unlock()
	out := make([]llm.Request, 0, len(f.bodies))
	for _, b := range f.bodies {
		var req llm.Request
		if err := json.Unmarshal(b, &req); err != nil {
			t.Fatalf("request body does not decode: %v", err)
		}
		out = append(out, req)
	}
	return out
}

// effortSent is the effort level the one request carried, or "" for none.
func (f *fakeAPI) effortSent(t *testing.T) string {
	t.Helper()
	reqs := f.requestsSeen(t)
	if len(reqs) != 1 {
		t.Fatalf("want exactly one request, got %d", len(reqs))
	}
	if reqs[0].OutputConfig == nil {
		return ""
	}
	return reqs[0].OutputConfig.Effort
}

func (f *fakeAPI) modelsSeen() []string {
	f.mu.Lock()
	defer f.mu.Unlock()
	return append([]string(nil), f.models...)
}

// task pulls the first user message out of a request: the brief the child was
// spawned on.
func task(req llm.Request) string {
	for _, m := range req.Messages {
		if m.Role == llm.RoleUser && len(m.Content) > 0 {
			return m.Content[0].Text
		}
	}
	return ""
}

func toolUse(id, name string) string {
	return `{"id":"msg","type":"message","role":"assistant","model":"m","content":[` +
		`{"type":"tool_use","id":"` + id + `","name":"` + name + `","input":{}}],` +
		`"stop_reason":"tool_use","usage":{"input_tokens":100,"output_tokens":20}}`
}

func endTurn(text string) string {
	b, _ := json.Marshal(text)
	return `{"id":"msg","type":"message","role":"assistant","model":"m","content":[{"type":"text","text":` +
		string(b) + `}],"stop_reason":"end_turn","usage":{"input_tokens":100,"output_tokens":20}}`
}

func newSupervisor(t *testing.T, f *fakeAPI, opts ...func(*Options)) *Supervisor {
	t.Helper()
	log := metrics.NewWriter(&bytes.Buffer{}, "test")
	root := t.TempDir()
	o := Options{
		Client: llm.New(llm.WithBaseURL(f.srv.URL), llm.WithAPIKey("k"), llm.WithModel("claude-opus-5")),
		Log:    log,
		Tools: func(state *agent.State, log *metrics.Logger) *tools.Registry {
			return tools.NewRegistry(log).Register(
				tools.NewRead(root, state),
				tools.New("noop", "does nothing", func(context.Context, struct{}) (string, error) {
					return "ok", nil
				}),
			)
		},
	}
	for _, fn := range opts {
		fn(&o)
	}
	s := New(context.Background(), o)
	t.Cleanup(s.Shutdown)
	return s
}

// waitFor polls until cond holds, so tests do not race the goroutines they are
// watching. It fails rather than hanging.
func waitFor(t *testing.T, what string, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for %s", what)
}

func statusOf(t *testing.T, s *Supervisor, id string) Status {
	t.Helper()
	info, ok := s.Info(id)
	if !ok {
		t.Fatalf("no agent %q", id)
	}
	return info.Status
}

func TestSpawnRunsTheAgentAndKeepsItsReport(t *testing.T) {
	f := newFakeAPI(t)
	s := newSupervisor(t, f)

	id, err := s.Spawn(Spec{Preset: "explore", Task: "find the approval policy"})
	if err != nil {
		t.Fatalf("Spawn: %v", err)
	}
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })

	info, _ := s.Info(id)
	if info.Report != "handled: find the approval policy" {
		t.Errorf("report = %q, want the agent's final text", info.Report)
	}
	if info.Preset != "explore" || info.Task != "find the approval policy" {
		t.Errorf("info = %+v, want the preset and task it was spawned with", info)
	}
	if info.Turns != 1 {
		t.Errorf("turns = %d, want 1", info.Turns)
	}
	if info.Usage.OutputTokens != 20 {
		t.Errorf("usage = %+v, want the child's token counts", info.Usage)
	}
	if info.Ended.IsZero() || info.Ended.Before(info.Started) {
		t.Errorf("started %v, ended %v", info.Started, info.Ended)
	}
}

// Ids are what the sidebar shows and what the model says back to you, so they
// are short and in spawn order rather than random.
func TestIDsAreShortAndInSpawnOrder(t *testing.T) {
	f := newFakeAPI(t)
	s := newSupervisor(t, f)

	first, _ := s.Spawn(Spec{Preset: "explore", Task: "one"})
	second, _ := s.Spawn(Spec{Preset: "plan", Task: "two"})
	if first != "a1" || second != "a2" {
		t.Errorf("ids = %q, %q; want a1, a2", first, second)
	}

	got := []string{}
	for _, info := range s.Snapshot() {
		got = append(got, info.ID)
	}
	if strings.Join(got, " ") != "a1 a2" {
		t.Errorf("snapshot order = %v, want spawn order", got)
	}
}

// A preset the model made up must come back as something it can act on, not as
// a silent no-op or a crash.
func TestUnknownPresetNamesTheOnesThatExist(t *testing.T) {
	f := newFakeAPI(t)
	s := newSupervisor(t, f)

	_, err := s.Spawn(Spec{Preset: "researcher", Task: "go"})
	if err == nil {
		t.Fatal("spawning an unknown preset succeeded")
	}
	for _, p := range Presets() {
		if !strings.Contains(err.Error(), p.Name) {
			t.Errorf("error %q does not offer the %q preset", err, p.Name)
		}
	}
	if len(s.Snapshot()) != 0 {
		t.Error("a rejected spawn still created an agent")
	}
}

// No preset pins a model: the agent doing the spawning picks one for the job,
// and a spawn that names none runs on the session's own model rather than on
// something the harness chose for it.
func TestTheSpawningAgentChoosesTheModel(t *testing.T) {
	f := newFakeAPI(t)
	s := newSupervisor(t, f)

	unnamed, _ := s.Spawn(Spec{Preset: "explore", Task: "one"})
	chosen, _ := s.Spawn(Spec{Preset: "explore", Task: "two", Model: "claude-haiku-4-5"})
	waitFor(t, "both agents to finish", func() bool {
		return statusOf(t, s, unnamed) == Done && statusOf(t, s, chosen) == Done
	})

	first, _ := s.Info(unnamed)
	second, _ := s.Info(chosen)
	if first.Model != "claude-opus-5" {
		t.Errorf("a spawn naming no model ran on %q, want the session's", first.Model)
	}
	if second.Model != "claude-haiku-4-5" {
		t.Errorf("the chosen model was %q", second.Model)
	}
	// Both agents run at once, so which request arrives first is not
	// something to assert on — only that each asked for its own model.
	models := f.modelsSeen()
	sort.Strings(models)
	want := []string{"claude-haiku-4-5", "claude-opus-5"}
	if strings.Join(models, " ") != strings.Join(want, " ") {
		t.Errorf("requests asked for %v, want %v", models, want)
	}
}

// An agent whose model call fails is finished, not stuck: the status says so
// and the reason is kept for the parent to be told.
func TestAFailedAgentIsRecordedAsFailed(t *testing.T) {
	f := newFakeAPI(t)
	f.fail(http.StatusBadRequest)
	s := newSupervisor(t, f)

	id, _ := s.Spawn(Spec{Preset: "explore", Task: "go"})
	waitFor(t, "the agent to fail", func() bool { return statusOf(t, s, id) == Failed })

	info, _ := s.Info(id)
	if info.Err == "" {
		t.Error("a failed agent kept no reason")
	}
	if info.Ended.IsZero() {
		t.Error("a failed agent never ended")
	}
}

// Every event a child produces reaches the front end tagged with the agent it
// came from, because they all land in one UI.
func TestEventsAreTaggedWithTheAgent(t *testing.T) {
	f := newFakeAPI(t)
	var mu sync.Mutex
	var seen []Event
	s := newSupervisor(t, f, func(o *Options) {
		o.Observe = func(ev Event) {
			mu.Lock()
			defer mu.Unlock()
			seen = append(seen, ev)
		}
	})

	id, _ := s.Spawn(Spec{Preset: "explore", Task: "go"})
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })

	mu.Lock()
	defer mu.Unlock()
	if len(seen) == 0 {
		t.Fatal("no events were observed")
	}
	for _, ev := range seen {
		if ev.ID != id {
			t.Fatalf("event %+v is not tagged with %q", ev, id)
		}
		if ev.Preset != "explore" {
			t.Fatalf("event %+v does not name the preset", ev)
		}
	}
}

// How hard a subagent thinks is the other half of what it costs, so the agent
// spawning it chooses that too.
func TestTheChosenEffortReachesTheChild(t *testing.T) {
	f := newFakeAPI(t)
	s := newSupervisor(t, f, func(o *Options) { o.Base.Effort = "high" })

	id, err := s.Spawn(Spec{Preset: "explore", Task: "go", Model: "claude-opus-5", Effort: "low"})
	if err != nil {
		t.Fatalf("Spawn: %v", err)
	}
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })

	if got := f.effortSent(t); got != "low" {
		t.Errorf("the child asked for effort %q, want the level it was spawned with", got)
	}
	if info, _ := s.Info(id); info.Effort != "low" {
		t.Errorf("Info.Effort = %q", info.Effort)
	}
}

// Choosing nothing inherits the session's level, the way it always did.
func TestASpawnWithNoEffortInheritsTheSessions(t *testing.T) {
	f := newFakeAPI(t)
	s := newSupervisor(t, f, func(o *Options) { o.Base.Effort = "xhigh" })

	id, _ := s.Spawn(Spec{Preset: "explore", Task: "go", Model: "claude-opus-5"})
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })

	if got := f.effortSent(t); got != "xhigh" {
		t.Errorf("the child asked for effort %q, want the session's", got)
	}
}

// A level that is not a level is a typo, and the model needs to be told which
// words exist rather than watching an agent run at something else.
func TestAnUnknownEffortLevelIsRefused(t *testing.T) {
	s := newSupervisor(t, newFakeAPI(t))

	_, err := s.Spawn(Spec{Preset: "explore", Task: "go", Effort: "lowest"})
	if err == nil {
		t.Fatal("an unknown effort level was accepted")
	}
	for _, level := range llm.EffortLevels {
		if !strings.Contains(err.Error(), level) {
			t.Errorf("error %q does not offer %q", err, level)
		}
	}
	if len(s.Snapshot()) != 0 {
		t.Error("a rejected spawn still created an agent")
	}
}

// Haiku takes no effort level at all. The spawn is still a good spawn — the
// model picked a cheap model for a cheap job — so it runs, and the answer says
// plainly that the level was not applied rather than implying it was.
func TestEffortOnAModelThatCannotTakeItIsRunAndReported(t *testing.T) {
	f := newFakeAPI(t)
	s := newSupervisor(t, f)

	id, err := s.Spawn(Spec{Preset: "explore", Task: "go", Model: "claude-haiku-4-5", Effort: "low"})
	if err != nil {
		t.Fatalf("Spawn: %v", err)
	}
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })

	if got := f.effortSent(t); got != "" {
		t.Errorf("the child sent effort %q to a model that rejects it", got)
	}
	info, _ := s.Info(id)
	if info.Effort != "" {
		t.Errorf("Info.Effort = %q, want empty — the level is not in force", info.Effort)
	}
}

// A borrowed Claude Code credential is only accepted alongside its identity
// block: a request without it comes back 429, which reads like a quota problem
// and is not one. Subagents authenticate with the same credential, so they
// carry the same block — and they get their own system prompt, which is
// exactly where it could have been dropped.
func TestASubagentCarriesTheIdentityBlockOfItsSession(t *testing.T) {
	f := newFakeAPI(t)
	s := newSupervisor(t, f, func(o *Options) {
		o.Base.SystemPrefix = agent.ClaudeCodeIdentity
	})

	id, _ := s.Spawn(Spec{Preset: "explore", Task: "go"})
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })

	reqs := f.requestsSeen(t)
	if len(reqs) != 1 {
		t.Fatalf("saw %d requests", len(reqs))
	}
	blocks := reqs[0].System
	if len(blocks) < 2 {
		t.Fatalf("the child sent %d system blocks, want the identity block and its own prompt", len(blocks))
	}
	if blocks[0].Text != agent.ClaudeCodeIdentity {
		t.Errorf("the first system block is %q, want the identity string", blocks[0].Text)
	}
	if !strings.Contains(blocks[1].Text, "You are a subagent") {
		t.Errorf("the child's own brief is missing:\n%s", blocks[1].Text)
	}
}

// Children authenticate exactly as the session does, including a credential
// that is resolved per request and may need refreshing — one source, shared,
// rather than a copy per agent that refreshes on its own.
func TestChildrenAuthenticateThroughTheSessionsCredential(t *testing.T) {
	var mu sync.Mutex
	var calls int
	var bearers []string

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		bearers = append(bearers, r.Header.Get("authorization"))
		mu.Unlock()
		io.WriteString(w, endTurn("ok"))
	}))
	t.Cleanup(srv.Close)

	client := llm.New(llm.WithBaseURL(srv.URL), llm.WithTokenSource(func(context.Context) (string, error) {
		mu.Lock()
		defer mu.Unlock()
		calls++
		return "borrowed-token", nil
	}))
	log := metrics.NewWriter(&bytes.Buffer{}, "test")
	root := t.TempDir()
	s := New(context.Background(), Options{
		Client: client,
		Log:    log,
		Tools: func(state *agent.State, log *metrics.Logger) *tools.Registry {
			return tools.NewRegistry(log).Register(tools.NewRead(root, state))
		},
	})
	t.Cleanup(s.Shutdown)

	id, _ := s.Spawn(Spec{Preset: "explore", Task: "go"})
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })

	mu.Lock()
	defer mu.Unlock()
	if calls == 0 {
		t.Fatal("the child never asked the session's token source for a credential")
	}
	for _, b := range bearers {
		if b != "Bearer borrowed-token" {
			t.Errorf("the child sent authorization %q", b)
		}
	}
}