Vectrum
Sign In
195 lines · 6.0 KB · Go
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
}