Vectrum
Sign In
281 lines · 9.0 KB · Go
package subagent

import (
	"context"
	"encoding/json"
	"strings"
	"sync"
	"testing"
	"time"

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

func toolNamed(t *testing.T, s *Supervisor, name string) tools.Tool {
	t.Helper()
	for _, tool := range s.Tools() {
		if tool.Name() == name {
			return tool
		}
	}
	t.Fatalf("no %s tool", name)
	return nil
}

func TestTheParentGetsSpawnAndWaitAndNothingElse(t *testing.T) {
	s := newSupervisor(t, newFakeAPI(t))
	var names []string
	for _, tool := range s.Tools() {
		names = append(names, tool.Name())
	}
	if strings.Join(names, " ") != "spawn_agent agent_wait" {
		t.Errorf("tools = %v", names)
	}
}

// The description is what the model chooses a preset from, so every preset has
// to be in it — and so does the rule that subagents cannot spawn subagents,
// which it would otherwise discover by trying.
func TestTheSpawnDescriptionExplainsThePresetsAndTheRules(t *testing.T) {
	s := newSupervisor(t, newFakeAPI(t))
	desc := toolNamed(t, s, "spawn_agent").Description()

	for _, p := range Presets() {
		if !strings.Contains(desc, p.Name) {
			t.Errorf("description never mentions the %q preset", p.Name)
		}
	}
	for _, phrase := range []string{"background", "approve", "cannot spawn"} {
		if !strings.Contains(strings.ToLower(desc), phrase) {
			t.Errorf("description never says %q:\n%s", phrase, desc)
		}
	}
}

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

	out, err := toolNamed(t, s, "spawn_agent").Execute(context.Background(),
		json.RawMessage(`{"preset":"explore","task":"find the policy"}`))
	if err != nil {
		t.Fatalf("spawn_agent: %v", err)
	}
	if !strings.Contains(out, "a1") {
		t.Errorf("spawn_agent returned %q, which does not name the agent", out)
	}
	f.release()
}

func TestSpawnToolReportsAnUnknownPresetToTheModel(t *testing.T) {
	s := newSupervisor(t, newFakeAPI(t))
	_, err := toolNamed(t, s, "spawn_agent").Execute(context.Background(),
		json.RawMessage(`{"preset":"researcher","task":"go"}`))
	if err == nil {
		t.Fatal("an unknown preset was accepted")
	}
	if !strings.Contains(err.Error(), "explore") {
		t.Errorf("error %q does not offer a preset that exists", err)
	}
}

// The spawning tool call belongs to one turn; the agent it starts does not.
// Interrupting that turn — or simply the turn ending — must not take the
// subagent with it.
func TestASpawnedAgentOutlivesTheTurnThatSpawnedIt(t *testing.T) {
	f := newFakeAPI(t)
	f.hold()
	s := newSupervisor(t, f)

	ctx, cancel := context.WithCancel(context.Background())
	if _, err := toolNamed(t, s, "spawn_agent").Execute(ctx,
		json.RawMessage(`{"preset":"general","task":"keep going"}`)); err != nil {
		t.Fatalf("spawn_agent: %v", err)
	}
	cancel()

	time.Sleep(settle)
	if got := statusOf(t, s, "a1"); !got.Live() {
		t.Fatalf("the agent is %s after the spawning turn ended", got)
	}
	f.release()
	waitFor(t, "the agent to finish anyway", func() bool { return statusOf(t, s, "a1") == Done })
}

func TestAgentWaitHandsBackTheReport(t *testing.T) {
	s := newSupervisor(t, newFakeAPI(t))
	id, _ := s.Spawn(Spec{Preset: "explore", Task: "find it"})

	out, err := toolNamed(t, s, "agent_wait").Execute(context.Background(),
		json.RawMessage(`{"id":"`+id+`"}`))
	if err != nil {
		t.Fatalf("agent_wait: %v", err)
	}
	if !strings.Contains(out, "handled: find it") {
		t.Errorf("agent_wait returned %q, want the agent's report", out)
	}
}

func TestAgentWaitOnAnUnknownIDIsAnError(t *testing.T) {
	s := newSupervisor(t, newFakeAPI(t))
	_, err := toolNamed(t, s, "agent_wait").Execute(context.Background(), json.RawMessage(`{"id":"a9"}`))
	if err == nil {
		t.Fatal("waiting on an unknown agent succeeded")
	}
}

// Nesting is the way a spawn loop becomes a fork bomb with a credit card. A
// subagent's tool set simply does not contain the spawn tools, and calling one
// comes back as an ordinary unknown-tool error it can read.
func TestASubagentCannotSpawnSubagents(t *testing.T) {
	f := newFakeAPI(t)
	f.callTool = "spawn_agent"
	var mu sync.Mutex
	var errors []string
	s := newSupervisor(t, f, func(o *Options) {
		o.Observe = func(ev Event) {
			if ev.Agent.Kind == agent.EventToolResult && ev.Agent.IsError {
				mu.Lock()
				errors = append(errors, ev.Agent.Result)
				mu.Unlock()
			}
		}
	})

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

	mu.Lock()
	defer mu.Unlock()
	if len(errors) != 1 || !strings.Contains(errors[0], "unknown tool") {
		t.Fatalf("a subagent's spawn attempt produced %v", errors)
	}
	if n := len(s.Snapshot()); n != 1 {
		t.Errorf("%d agents exist; the subagent managed to spawn one", n)
	}
}

// The model is the main thing that decides what a subagent costs, and nothing
// else picks it any more — so the description has to say how to choose. It
// names no ids of its own: what this endpoint serves is in the system prompt,
// and a list in two places is a list that goes out of step.
func TestTheSpawnDescriptionExplainsHowToChooseAModel(t *testing.T) {
	s := newSupervisor(t, newFakeAPI(t))
	desc := toolNamed(t, s, "spawn_agent").Description()

	for _, phrase := range []string{"you choose the model", "newest model", "system prompt"} {
		if !strings.Contains(strings.ToLower(desc), phrase) {
			t.Errorf("the description never says %q:\n%s", phrase, desc)
		}
	}
	if strings.Contains(desc, "claude-") {
		t.Errorf("the description hardcodes a model id, which will drift from what the endpoint serves:\n%s", desc)
	}
	// And it must not promise a model per preset, which is what it used to do.
	for _, p := range Presets() {
		line := p.Name + " — " + p.Description
		if !strings.Contains(desc, line) {
			t.Errorf("preset line missing or changed shape: %q", line)
		}
	}
}

// Waiting on an agent whose report the front end has already put into the
// conversation returns a pointer to it, not the text again: the parent is
// paying for context, and it already has this.
func TestWaitingOnAnAlreadyDeliveredReportDoesNotRepeatIt(t *testing.T) {
	f := newFakeAPI(t)
	s := newSupervisor(t, f)
	id, _ := s.Spawn(Spec{Preset: "explore", Task: "find it"})
	waitFor(t, "the agent to finish", func() bool { return statusOf(t, s, id) == Done })

	// The front end delivered it.
	if !s.ClaimReport(id) {
		t.Fatal("could not claim a fresh report")
	}

	out, err := toolNamed(t, s, "agent_wait").Execute(context.Background(), json.RawMessage(`{"id":"`+id+`"}`))
	if err != nil {
		t.Fatalf("agent_wait: %v", err)
	}
	if strings.Contains(out, "handled: find it") {
		t.Errorf("agent_wait repeated a report the conversation already has:\n%s", out)
	}
	for _, want := range []string{id, "already"} {
		if !strings.Contains(out, want) {
			t.Errorf("agent_wait's answer does not say %q:\n%s", want, out)
		}
	}
}

// And when nobody has delivered it, agent_wait is the delivery.
func TestWaitingFirstHandsOverTheWholeReport(t *testing.T) {
	f := newFakeAPI(t)
	s := newSupervisor(t, f)
	id, _ := s.Spawn(Spec{Preset: "explore", Task: "find it"})

	out, err := toolNamed(t, s, "agent_wait").Execute(context.Background(), json.RawMessage(`{"id":"`+id+`"}`))
	if err != nil {
		t.Fatalf("agent_wait: %v", err)
	}
	if !strings.Contains(out, "handled: find it") {
		t.Errorf("agent_wait did not hand over the report:\n%s", out)
	}
	if s.ClaimReport(id) {
		t.Error("the front end would send it again after agent_wait delivered it")
	}
}

// The effort argument is the answer to "spawn these on low", which the tool
// could not express before.
func TestSpawnToolTakesAnEffortLevel(t *testing.T) {
	f := newFakeAPI(t)
	f.hold()
	s := newSupervisor(t, f)

	out, err := toolNamed(t, s, "spawn_agent").Execute(context.Background(),
		json.RawMessage(`{"preset":"explore","task":"find it","model":"claude-opus-5","effort":"low"}`))
	if err != nil {
		t.Fatalf("spawn_agent: %v", err)
	}
	if !strings.Contains(out, "low effort") {
		t.Errorf("the answer does not say what it is running at:\n%s", out)
	}
	if info, _ := s.Info("a1"); info.Effort != "low" {
		t.Errorf("Info.Effort = %q", info.Effort)
	}
	f.release()
}

// Asking for an effort level on a model that has none must not read as though
// it worked.
func TestSpawnToolSaysWhenTheEffortCouldNotBeApplied(t *testing.T) {
	f := newFakeAPI(t)
	f.hold()
	s := newSupervisor(t, f)

	out, err := toolNamed(t, s, "spawn_agent").Execute(context.Background(),
		json.RawMessage(`{"preset":"explore","task":"find it","model":"claude-haiku-4-5","effort":"low"}`))
	if err != nil {
		t.Fatalf("spawn_agent: %v", err)
	}
	for _, want := range []string{"takes no effort level", "not applied"} {
		if !strings.Contains(out, want) {
			t.Errorf("the answer does not say the level was dropped (%q):\n%s", want, out)
		}
	}
	f.release()
}

// And the description has to say the knob exists, or nothing will reach for it.
func TestTheSpawnDescriptionMentionsEffort(t *testing.T) {
	s := newSupervisor(t, newFakeAPI(t))
	desc := strings.ToLower(toolNamed(t, s, "spawn_agent").Description())
	for _, want := range []string{"effort", "low"} {
		if !strings.Contains(desc, want) {
			t.Errorf("the description never mentions %q", want)
		}
	}
}