package tui
import (
"bytes"
"context"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
"github.com/charmbracelet/bubbles/spinner"
tea "github.com/charmbracelet/bubbletea"
"github.com/johannSo/codey/internal/agent"
"github.com/johannSo/codey/internal/approval"
"github.com/johannSo/codey/internal/llm"
"github.com/johannSo/codey/internal/metrics"
"github.com/johannSo/codey/internal/subagent"
"github.com/johannSo/codey/internal/tools"
)
// scriptedAPI answers the main agent and its children differently, so one
// server can play both halves of a spawn.
func scriptedAPI(t *testing.T) *httptest.Server {
t.Helper()
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)
var system string
for _, b := range req.System {
system += b.Text
}
var conversation string
for _, m := range req.Messages {
for _, b := range m.Content {
conversation += b.Text + b.Content
}
}
switch {
case strings.Contains(system, "You are a subagent"):
io.WriteString(w, liveMsg("end_turn", `{"type":"text","text":"found it in internal/approval/policy.go:70"}`))
case strings.Contains(conversation, "[subagent a1"):
io.WriteString(w, liveMsg("end_turn", `{"type":"text","text":"thanks — that is the line I needed"}`))
case strings.Contains(conversation, "started a1"):
io.WriteString(w, liveMsg("end_turn", `{"type":"text","text":"an agent is on it"}`))
default:
io.WriteString(w, liveMsg("tool_use", `{"type":"tool_use","id":"tu_1","name":"spawn_agent",`+
`"input":{"preset":"explore","task":"find where the approval policy is consulted"}}`))
}
}))
t.Cleanup(srv.Close)
return srv
}
func liveMsg(stop string, blocks ...string) string {
return `{"id":"msg","type":"message","role":"assistant","model":"m","content":[` +
strings.Join(blocks, ",") + `],"stop_reason":"` + stop + `",` +
`"usage":{"input_tokens":100,"output_tokens":20}}`
}
// A whole spawn, end to end: the model asks for an agent, the supervisor runs
// it, the sidebar fills in, the agent's own thread is readable, and its report
// comes back into the conversation without anybody typing.
func TestSpawningFillsTheSidebarAndTheReportComesBack(t *testing.T) {
srv := scriptedAPI(t)
root := t.TempDir()
log := metrics.NewWriter(&bytes.Buffer{}, "test")
client := llm.New(llm.WithBaseURL(srv.URL), llm.WithAPIKey("k"))
policy := approval.New(nil)
policy.YOLO = true // approvals have their own tests; this one is about the wiring
newTools := func(state *agent.State, log *metrics.Logger) *tools.Registry {
return tools.NewRegistry(log).Register(tools.NewRead(root, state))
}
sup := subagent.New(context.Background(), subagent.Options{
Client: client, Log: log, Policy: policy, Tools: newTools,
})
t.Cleanup(sup.Shutdown)
state := agent.NewState()
reg := newTools(state, log)
reg.Register(sup.Tools()...)
a := agent.New(client, reg, state, log, agent.Options{})
a.Approve = policy
m := newModel(context.Background(), Options{
Agent: a, Client: client, Policy: policy, Agents: sup, Root: root,
SessionDir: t.TempDir(), SessionID: "test",
})
m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
// Everything the supervisor and the agent report arrives here, and is
// pumped into the model from this goroutine — which is exactly what
// Bubble Tea's Send does in the real program.
msgs := make(chan tea.Msg, 256)
sup.Subscribe(
func(ev subagent.Event) { msgs <- evSub{ev: ev} },
func() { msgs <- evAgents{} },
)
a.Observe = func(ev agent.Event) {
if ev.Kind == agent.EventToolResult {
msgs <- evToolResult{tool: ev.Tool, result: ev.Result}
}
}
// Send a prompt the way the prompt box does, and run the commands it
// returns off the update loop, as the program's scheduler would.
d := &driver{msgs: msgs}
// Everything the model started has to finish before the test's temporary
// directories go away: a turn saves its session as it ends.
defer d.wg.Wait()
_, cmd := m.submit("find where the approval policy is consulted")
d.run(cmd)
deadline := time.After(10 * time.Second)
var sawReport bool
for !sawReport {
select {
case msg := <-msgs:
_, next := m.Update(msg)
d.run(next)
sawReport = strings.Contains(historyText(a), "[subagent a1")
case <-deadline:
t.Fatalf("the report never came back. transcript:\n%s", m.tr.render())
}
}
// The sidebar knows about it, and the thread is readable.
if len(m.bar.rows) != 1 || m.bar.rows[0].ID != "a1" {
t.Fatalf("sidebar rows = %+v", m.bar.rows)
}
m.Update(key("ctrl+b"))
m.Update(key("enter"))
if view := m.View(); !strings.Contains(view, "internal/approval/policy.go:70") {
t.Errorf("a1's own thread is not readable:\n%s", view)
}
// And the main transcript said so in one line rather than repeating the
// whole report.
if !strings.Contains(m.tr.render(), "a1 (explore) finished") {
t.Errorf("the main transcript never noticed a1 finishing:\n%s", m.tr.render())
}
}
// driver runs Bubble Tea commands off the update loop, the way the program's
// own scheduler does, and keeps count of them so a test can wait for the work
// it started. A batch is a list of commands rather than a message, so it is
// unpacked here exactly as the real scheduler unpacks it.
//
// Spinner ticks are dropped: feeding one back produces another for ever, and
// the animation is not what any of this is testing.
type driver struct {
wg sync.WaitGroup
msgs chan tea.Msg
}
func (d *driver) run(cmd tea.Cmd) {
if cmd == nil {
return
}
d.wg.Add(1)
go func() {
defer d.wg.Done()
switch msg := cmd().(type) {
case nil, tea.QuitMsg, spinner.TickMsg:
case tea.BatchMsg:
for _, c := range msg {
d.run(c)
}
default:
select {
case d.msgs <- msg:
default: // nobody is reading any more; the test is done
}
}
}()
}
func historyText(a *agent.Agent) string {
var s string
for _, m := range a.State().History() {
for _, b := range m.Content {
s += b.Text
}
}
return s
}