package serve
import (
"context"
"io"
"strings"
"sync"
"testing"
"time"
"github.com/johannSo/codey/internal/acp"
"github.com/johannSo/codey/internal/agent"
"github.com/johannSo/codey/internal/llm"
"github.com/johannSo/codey/internal/metrics"
"github.com/johannSo/codey/internal/subagent"
"github.com/johannSo/codey/internal/tools"
)
// fakeAgents is a supervisor the test drives by hand: no children, no model
// calls, just the two callbacks and a snapshot the test controls.
type fakeAgents struct {
mu sync.Mutex
observe func(subagent.Event)
changed func()
infos []subagent.Info
reports []subagent.Info
claimed map[string]bool
shutdown bool
}
func newFakeAgents() *fakeAgents { return &fakeAgents{claimed: map[string]bool{}} }
func (f *fakeAgents) Subscribe(observe func(subagent.Event), changed func()) {
f.mu.Lock()
defer f.mu.Unlock()
f.observe, f.changed = observe, changed
}
func (f *fakeAgents) Snapshot() []subagent.Info {
f.mu.Lock()
defer f.mu.Unlock()
return append([]subagent.Info(nil), f.infos...)
}
func (f *fakeAgents) TakeReports() []subagent.Info {
f.mu.Lock()
defer f.mu.Unlock()
out := f.reports
f.reports = nil
return out
}
func (f *fakeAgents) ClaimReport(id string) bool {
f.mu.Lock()
defer f.mu.Unlock()
if f.claimed[id] {
return false
}
f.claimed[id] = true
return true
}
func (f *fakeAgents) Shutdown() {
f.mu.Lock()
defer f.mu.Unlock()
f.shutdown = true
}
// say makes a child speak.
func (f *fakeAgents) say(id, preset, text string) {
f.mu.Lock()
observe := f.observe
f.mu.Unlock()
observe(subagent.Event{ID: id, Preset: preset, Model: "test-model",
Agent: agent.Event{Kind: agent.EventText, Text: text}})
}
// finish ends a child, queues its report, and fires the change callback.
func (f *fakeAgents) finish(info subagent.Info) {
f.mu.Lock()
f.infos = append(f.infos, info)
if info.Status == subagent.Done {
f.reports = append(f.reports, info)
}
changed := f.changed
f.mu.Unlock()
changed()
}
// live adds a child that has not finished.
func (f *fakeAgents) live(info subagent.Info) {
f.mu.Lock()
f.infos = append(f.infos, info)
f.mu.Unlock()
}
// --- the row translation, tested directly --------------------------------
func TestASubagentGetsItsOwnRow(t *testing.T) {
e, r := newRecorded()
sess := &session{emit: e, agentsMoved: make(chan struct{}, 1)}
f := newFakeAgents()
sess.agents = f
sess.watchAgents()
f.say("a1", "explore", strings.Repeat("looking around. ", 40))
if len(r.updates) < 2 {
t.Fatalf("want a row and an update, got %v", r.kinds())
}
open := r.updates[0]
if open.SessionUpdate != acp.UpdateToolCall {
t.Fatalf("first update = %q", open.SessionUpdate)
}
if open.ToolCallID != "agent-a1" {
t.Errorf("toolCallId = %q — it must not collide with a real tool call", open.ToolCallID)
}
if !strings.Contains(open.Title, "a1") || !strings.Contains(open.Title, "explore") {
t.Errorf("title = %q", open.Title)
}
if open.Status != acp.StatusInProgress {
t.Errorf("status = %q", open.Status)
}
if !strings.Contains(contentText(t, r.updates[1]), "looking around") {
t.Errorf("the child's words did not reach its row")
}
}
func TestARowIsOpenedOnceAndClosedOnce(t *testing.T) {
e, r := newRecorded()
sess := &session{emit: e, agentsMoved: make(chan struct{}, 1)}
f := newFakeAgents()
sess.agents = f
sess.watchAgents()
f.say("a1", "explore", "one")
r.clock = r.clock.Add(progressAfter)
f.say("a1", "explore", "two")
info := subagent.Info{ID: "a1", Preset: "explore", Status: subagent.Done, Turns: 3, Report: "found it"}
f.finish(info)
// The change callback fires repeatedly once a child has ended.
f.finish(info)
f.finish(info)
var opens, closes int
for _, u := range r.updates {
switch {
case u.SessionUpdate == acp.UpdateToolCall:
opens++
case u.Status == acp.StatusCompleted || u.Status == acp.StatusFailed:
closes++
}
}
if opens != 1 {
t.Errorf("%d rows opened for one agent", opens)
}
if closes != 1 {
t.Errorf("%d closing updates for one agent", closes)
}
if got := contentText(t, r.updates[len(r.updates)-1]); !strings.Contains(got, "found it") {
t.Errorf("the row closed without the report: %q", got)
}
}
func TestAFailedSubagentClosesAsFailed(t *testing.T) {
e, r := newRecorded()
sess := &session{emit: e, agentsMoved: make(chan struct{}, 1)}
f := newFakeAgents()
sess.agents = f
sess.watchAgents()
f.finish(subagent.Info{ID: "a1", Preset: "plan", Status: subagent.Failed, Err: "the endpoint refused"})
last := r.updates[len(r.updates)-1]
if last.Status != acp.StatusFailed {
t.Errorf("status = %q, want failed", last.Status)
}
if !strings.Contains(contentText(t, last), "the endpoint refused") {
t.Errorf("content = %q", contentText(t, last))
}
// It never said anything, so the row has to be opened before it is closed
// or there is nothing on screen to close.
if r.updates[0].SessionUpdate != acp.UpdateToolCall {
t.Errorf("a silent agent left no row: %v", r.kinds())
}
}
// --- the turn, over the wire ---------------------------------------------
// newAgentsServer wires a server whose sessions all share one fake supervisor.
func newAgentsServer(t *testing.T, f *fakeAPI, agents Agents) *testClient {
t.Helper()
log := metrics.NewWriter(io.Discard, "test")
return newTestServer(t, f, Options{
NewAgent: func(req AgentRequest) (AgentSet, error) {
reg := tools.NewRegistry(log).Register(tools.NewRead(req.Root, req.State))
client := llm.New(llm.WithBaseURL(f.srv.URL), llm.WithAPIKey("k"), llm.WithModel("test-model"))
return AgentSet{
Agent: agent.New(client, reg, req.State, log, agent.Options{}),
Agents: agents,
}, nil
},
})
}
func TestATurnStaysOpenWhileASubagentWorks(t *testing.T) {
// Ending the turn while a child is running would leave its row nowhere to
// appear and its report nobody to be handed to.
f := newFakeAPI(t,
msg("end_turn", textBlock("spawned it")),
msg("end_turn", textBlock("and here is what it found")),
)
agents := newFakeAgents()
agents.live(subagent.Info{ID: "a1", Preset: "explore", Status: subagent.Running})
c := newAgentsServer(t, f, agents)
id := c.handshake(t.TempDir())
done := make(chan acp.PromptResponse, 1)
go func() { done <- c.prompt(id, "find the bug") }()
select {
case <-done:
t.Fatal("the turn ended while a subagent was still running")
case <-time.After(150 * time.Millisecond):
}
// The child finishes and reports. That restarts the turn with its report.
agents.mu.Lock()
agents.infos = nil
agents.mu.Unlock()
agents.finish(subagent.Info{ID: "a1", Preset: "explore", Status: subagent.Done, Turns: 2, Report: "it is the parser"})
select {
case resp := <-done:
if resp.StopReason != acp.StopEndTurn {
t.Errorf("stopReason = %q", resp.StopReason)
}
case <-time.After(5 * time.Second):
t.Fatal("the turn never ended")
}
// The report has to have reached the model, not only the screen.
f.mu.Lock()
defer f.mu.Unlock()
if len(f.requests) != 2 {
t.Fatalf("the parent was asked %d times, want 2", len(f.requests))
}
if !strings.Contains(string(f.requests[1]), "it is the parser") {
t.Error("the subagent's report was never handed to the parent")
}
}
func TestAReportAlreadyClaimedIsNotHandedOverTwice(t *testing.T) {
// agent_wait may have collected it. Putting it in again would charge the
// model twice for something already in its context.
f := newFakeAPI(t, msg("end_turn", textBlock("done")))
agents := newFakeAgents()
agents.claimed["a1"] = true
agents.reports = []subagent.Info{{ID: "a1", Preset: "explore", Status: subagent.Done, Report: "already seen"}}
c := newAgentsServer(t, f, agents)
id := c.handshake(t.TempDir())
c.prompt(id, "go")
f.mu.Lock()
defer f.mu.Unlock()
if len(f.requests) != 1 {
t.Fatalf("the parent was asked %d times, want 1", len(f.requests))
}
}
func TestCancellingATurnStopsWaitingForChildren(t *testing.T) {
f := newFakeAPI(t, msg("end_turn", textBlock("spawned")))
agents := newFakeAgents()
agents.live(subagent.Info{ID: "a1", Preset: "explore", Status: subagent.Running})
c := newAgentsServer(t, f, agents)
id := c.handshake(t.TempDir())
done := make(chan acp.PromptResponse, 1)
go func() { done <- c.prompt(id, "find it") }()
time.Sleep(100 * time.Millisecond)
c.notify(acp.MethodCancel, acp.CancelParams{SessionID: id})
select {
case resp := <-done:
if resp.StopReason != acp.StopCancelled {
t.Errorf("stopReason = %q, want cancelled", resp.StopReason)
}
case <-time.After(5 * time.Second):
t.Fatal("cancelling did not break the wait")
}
}
func TestTheChildrenAreStoppedWhenTheConnectionEnds(t *testing.T) {
f := newFakeAPI(t, msg("end_turn", textBlock("ok")))
agents := newFakeAgents()
toAgent, fromClient := io.Pipe()
fromAgent, toClient := io.Pipe()
go func() { _, _ = io.Copy(io.Discard, fromAgent) }()
log := metrics.NewWriter(io.Discard, "test")
srv := New(Options{
NewAgent: func(req AgentRequest) (AgentSet, error) {
reg := tools.NewRegistry(log).Register(tools.NewRead(req.Root, req.State))
client := llm.New(llm.WithBaseURL(f.srv.URL), llm.WithAPIKey("k"), llm.WithModel("m"))
return AgentSet{Agent: agent.New(client, reg, req.State, log, agent.Options{}), Agents: agents}, nil
},
})
served := make(chan error, 1)
go func() { served <- srv.Run(context.Background(), toAgent, toClient) }()
// Open a session, then hang up.
if _, err := io.WriteString(fromClient, `{"jsonrpc":"2.0","id":1,"method":"session/new","params":{"cwd":`+quote(t.TempDir())+"}}\n"); err != nil {
t.Fatal(err)
}
time.Sleep(100 * time.Millisecond)
_ = fromClient.Close()
select {
case <-served:
case <-time.After(3 * time.Second):
t.Fatal("Run did not return")
}
agents.mu.Lock()
defer agents.mu.Unlock()
if !agents.shutdown {
t.Error("a subagent outlived the connection, still editing a tree nobody is watching")
}
}
func TestSubagentRowsAreBoundedLikeAnyOtherToolOutput(t *testing.T) {
e, r := newRecorded()
sess := &session{emit: e, agentsMoved: make(chan struct{}, 1)}
f := newFakeAgents()
sess.agents = f
sess.watchAgents()
f.say("a1", "explore", strings.Repeat("x", progressMaxBytes*2)+"END")
body := contentText(t, r.updates[len(r.updates)-1])
if len(body) > progressMaxBytes+64 {
t.Errorf("an unbounded %d bytes went over the wire", len(body))
}
// The child's line is newline-terminated, so trim before looking at the end.
if !strings.HasSuffix(strings.TrimSpace(body), "END") {
t.Error("the tail is the useful part and it was the part dropped")
}
}