byre/internal/agents/manager_test.go

package agents_test

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"g.fc.run/misc/byre/internal/agents"
	"g.fc.run/misc/byre/internal/api"
	"g.fc.run/misc/byre/internal/git"
	"g.fc.run/misc/byre/internal/store"
)

const longCommandPrompt = "Run the bash command `ping -c 120 127.0.0.1` in the foreground (not in the background) and wait for it to finish, then say done."

const (
	testModel = "sonnet"
	login     = "alice@example.com"
	gitName   = "Alice Test"
	gitEmail  = "alice@byre.invalid"
)

type env struct {
	t       *testing.T
	ctx     context.Context
	store   *store.Store
	events  chan api.Event
	user    api.User
	repo    string
	project api.Project
	task    api.Task
}

func newEnv(t *testing.T) *env {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	t.Cleanup(cancel)

	e := &env{t: t, ctx: ctx, events: make(chan api.Event, 4096)}

	var err error
	e.store, err = store.Open(filepath.Join(t.TempDir(), "byre.db"), func(ev api.Event) { e.events <- ev })
	if err != nil {
		t.Fatalf("open store: %v", err)
	}
	t.Cleanup(func() { e.store.Close() })

	_, err = e.store.EnsureUser(ctx, login, "Alice")
	if err != nil {
		t.Fatalf("ensure user: %v", err)
	}
	e.user, err = e.store.UpdateUserGit(ctx, login, gitName, gitEmail)
	if err != nil {
		t.Fatalf("update user: %v", err)
	}

	e.repo = scratchRepo(t)
	e.project = api.Project{Id: store.NewID(), Name: "demo", RepoPath: e.repo, DefaultBranch: "main", Remote: "none", CreatedBy: login, CreatedAt: time.Now().UTC()}
	err = e.store.CreateProject(ctx, e.project)
	if err != nil {
		t.Fatalf("create project: %v", err)
	}

	e.task = e.addTask()
	return e
}

func (e *env) addTask() api.Task {
	e.t.Helper()
	id := store.NewID()
	task := api.Task{
		Id:           id,
		ProjectId:    e.project.Id,
		Title:        "test",
		Branch:       "byre/" + id,
		WorktreePath: filepath.Join(e.t.TempDir(), id),
		Status:       api.TaskStatusActive,
		CreatedBy:    login,
		CreatedAt:    time.Now().UTC(),
	}

	var err error
	task.BaseRef, err = git.AddWorktree(e.repo, task.WorktreePath, task.Branch, "main")
	if err != nil {
		e.t.Fatalf("add worktree: %v", err)
	}
	err = e.store.CreateTask(e.ctx, task)
	if err != nil {
		e.t.Fatalf("create task: %v", err)
	}
	return task
}

func (e *env) manager() *agents.Manager {
	m := agents.New(e.store, agents.DefaultLimits)
	e.t.Cleanup(m.Close)
	return m
}

func (e *env) waitFor(desc string, match func(api.Event) bool) []api.Event {
	e.t.Helper()
	seen := []api.Event{}
	for {
		select {
		case ev := <-e.events:
			seen = append(seen, ev)
			e.t.Logf("event %d %s %v", ev.Seq, ev.Kind, ev.Payload)
			if match(ev) {
				return seen
			}
		case <-e.ctx.Done():
			e.t.Fatalf("timed out waiting for %s", desc)
		}
	}
}

func (e *env) waitIdle() []api.Event {
	e.t.Helper()
	return e.waitFor("session idle", func(ev api.Event) bool {
		return ev.Kind == "session.state" && ev.Payload["state"] == string(api.AgentSessionStateIdle)
	})
}

func kinds(events []api.Event) []string {
	out := []string{}
	for _, ev := range events {
		out = append(out, ev.Kind)
	}
	return out
}

func has(events []api.Event, kind string) bool {
	for _, ev := range events {
		if ev.Kind == kind {
			return true
		}
	}
	return false
}

func TestTurnAndFollowUp(t *testing.T) {
	t.Parallel()
	e := newEnv(t)
	m := e.manager()

	sess, err := m.Create(e.ctx, e.task, e.user, `Create a file hi.txt containing exactly the word hi, then git add and git commit it with the message "add hi".`, testModel)
	if err != nil {
		t.Fatalf("create session: %v", err)
	}
	if sess.State != api.AgentSessionStateRunning {
		t.Fatalf("new session is %s, want running", sess.State)
	}

	events := e.waitIdle()
	for _, kind := range []string{"session.created", "user.message", "agent.tool_call", "agent.tool_result", "agent.report", "turn.ended"} {
		if !has(events, kind) {
			t.Fatalf("missing %s in %v", kind, kinds(events))
		}
	}

	log := gitOut(t, e.task.WorktreePath, "log", "-1", "--format=%an <%ae>")
	if log != gitName+" <"+gitEmail+">" {
		t.Fatalf("agent committed as %q", log)
	}

	_, err = m.Send(e.ctx, sess.Id, login, "Reply with the single word ok.")
	if err != nil {
		t.Fatalf("send: %v", err)
	}

	events = e.waitIdle()
	reply := ""
	for _, ev := range events {
		if ev.Kind == "agent.message" {
			reply += ev.Payload["text"].(string)
		}
	}
	if !strings.Contains(strings.ToLower(reply), "ok") {
		t.Fatalf("follow-up reply %q", reply)
	}

	_, err = m.Send(e.ctx, sess.Id, login, "again")
	if err != nil {
		t.Fatalf("send while idle: %v", err)
	}
	_, err = m.Send(e.ctx, sess.Id, login, "too soon")
	if err == nil {
		t.Fatalf("send while running should conflict")
	}
	e.waitIdle()
}

func TestInterrupt(t *testing.T) {
	t.Parallel()
	e := newEnv(t)
	m := e.manager()

	sess, err := m.Create(e.ctx, e.task, e.user, longCommandPrompt, testModel)
	if err != nil {
		t.Fatalf("create session: %v", err)
	}

	e.waitFor("tool call", func(ev api.Event) bool { return ev.Kind == "agent.tool_call" })
	time.Sleep(2 * time.Second)

	at := time.Now()
	_, err = m.Interrupt(e.ctx, sess.Id)
	if err != nil {
		t.Fatalf("interrupt: %v", err)
	}

	events := e.waitIdle()
	if time.Since(at) > 30*time.Second {
		t.Fatalf("interrupt took %s", time.Since(at))
	}
	for _, ev := range events {
		if ev.Kind != "agent.tool_result" {
			continue
		}
		if ev.Payload["is_error"] != true || strings.Contains(fmt.Sprint(ev.Payload["content"]), "Blocked") {
			t.Fatalf("command finished or was refused instead of being interrupted: %v", ev.Payload)
		}
	}
}

func TestRecoverMidTurn(t *testing.T) {
	t.Parallel()
	e := newEnv(t)
	first := agents.New(e.store, agents.DefaultLimits)

	sess, err := first.Create(e.ctx, e.task, e.user, "Run the bash command `sleep 15` in the foreground (not in the background). After it finishes, create a file done.txt containing exactly the word done.", testModel)
	if err != nil {
		t.Fatalf("create session: %v", err)
	}

	e.waitFor("tool call", func(ev api.Event) bool { return ev.Kind == "agent.tool_call" })
	first.Close()

	stored, err := e.store.GetSession(e.ctx, sess.Id)
	if err != nil {
		t.Fatalf("get session: %v", err)
	}
	if stored.State != api.AgentSessionStateRunning {
		t.Fatalf("killed mid-turn session is %s, want running", stored.State)
	}

	second := e.manager()
	err = second.Recover(e.ctx)
	if err != nil {
		t.Fatalf("recover: %v", err)
	}

	events := e.waitIdle()
	notified := false
	for _, ev := range events {
		if ev.Kind == "user.message" && ev.Actor == "byre" {
			notified = true
		}
	}
	if !notified {
		t.Fatalf("no restart notice in %v", kinds(events))
	}

	out, err := os.ReadFile(filepath.Join(e.task.WorktreePath, "done.txt"))
	if err != nil {
		t.Fatalf("resumed agent did not finish: %v", err)
	}
	if strings.TrimSpace(string(out)) != "done" {
		t.Fatalf("done.txt contains %q", out)
	}
}

func TestStopTask(t *testing.T) {
	t.Parallel()
	e := newEnv(t)
	m := e.manager()

	sess, err := m.Create(e.ctx, e.task, e.user, longCommandPrompt, testModel)
	if err != nil {
		t.Fatalf("create session: %v", err)
	}

	e.waitFor("tool call", func(ev api.Event) bool { return ev.Kind == "agent.tool_call" })

	err = m.StopTask(e.ctx, e.task.Id, login)
	if err != nil {
		t.Fatalf("stop task: %v", err)
	}

	stored, err := e.store.GetSession(e.ctx, sess.Id)
	if err != nil {
		t.Fatalf("get session: %v", err)
	}
	if stored.State != api.AgentSessionStateCompleted || m.LiveCount() != 0 {
		t.Fatalf("after stop: state %s, live %d", stored.State, m.LiveCount())
	}

	_, err = m.Send(e.ctx, sess.Id, login, "hello?")
	if err == nil {
		t.Fatalf("send to completed session should conflict")
	}
}

const colorPrompt = `I need a file color.txt containing exactly one color name. Don't pick it yourself: ` +
	`ask me with ask_human (blocking) whether it should be red or blue, recommending red. ` +
	`When you get my answer, write the file and commit it with the message "add color".`

func (e *env) openQuestions(sessionID string) []api.InboxItem {
	e.t.Helper()
	items, err := e.store.ListInbox(e.ctx, store.InboxFilter{SessionID: sessionID, State: api.Open})
	if err != nil {
		e.t.Fatalf("list inbox: %v", err)
	}
	questions := []api.InboxItem{}
	for _, item := range items {
		if item.Kind == api.InboxItemKindQuestion {
			questions = append(questions, item)
		}
	}
	return questions
}

func (e *env) respond(m *agents.Manager, item api.InboxItem, typ api.InboxResponseType, option, text string, state api.InboxState) {
	e.t.Helper()
	resp := api.InboxResponse{
		Id:          store.NewID(),
		ItemId:      item.Id,
		User:        login,
		Type:        typ,
		Attachments: []api.Attachment{},
		CreatedAt:   time.Now().UTC(),
	}
	if option != "" {
		resp.Option = &option
	}
	if text != "" {
		resp.Text = &text
	}

	err := e.store.Respond(e.ctx, login, []store.Resolution{{Response: resp, State: state}})
	if err != nil {
		e.t.Fatalf("respond: %v", err)
	}
	err = m.Deliver(e.ctx, item.AgentSessionId)
	if err != nil {
		e.t.Fatalf("deliver: %v", err)
	}
}

func TestAskHumanRoundTrip(t *testing.T) {
	t.Parallel()
	e := newEnv(t)
	m := e.manager()

	sess, err := m.Create(e.ctx, e.task, e.user, colorPrompt, testModel)
	if err != nil {
		t.Fatalf("create session: %v", err)
	}
	e.waitIdle()

	questions := e.openQuestions(sess.Id)
	if len(questions) != 1 {
		t.Fatalf("want one open question, got %+v", questions)
	}
	q := questions[0]
	if !q.Blocking || q.Priority != api.Blocked || q.Payload["recommended_option"] != "red" {
		t.Fatalf("unexpected question %+v", q)
	}
	if _, err := os.Stat(filepath.Join(e.task.WorktreePath, "color.txt")); err == nil {
		t.Fatalf("agent wrote color.txt before getting an answer")
	}

	waiting, err := e.store.AwaitingHuman(e.ctx, sess.Id)
	if err != nil || waiting != 1 {
		t.Fatalf("awaiting human: %d %v", waiting, err)
	}

	e.respond(m, q, api.InboxResponseTypeOption, "blue", "", api.Answered)
	events := e.waitIdle()

	delivered := false
	for _, ev := range events {
		if ev.Kind == "user.message" && strings.Contains(ev.Payload["text"].(string), "chose: blue") {
			delivered = true
		}
	}
	if !delivered {
		t.Fatalf("answer was not delivered: %v", kinds(events))
	}

	out, err := os.ReadFile(filepath.Join(e.task.WorktreePath, "color.txt"))
	if err != nil {
		t.Fatalf("agent didn't write color.txt: %v", err)
	}
	if strings.TrimSpace(string(out)) != "blue" {
		t.Fatalf("color.txt contains %q, want blue", out)
	}

	pending, err := e.store.Undelivered(e.ctx, sess.Id)
	if err != nil || len(pending) != 0 {
		t.Fatalf("responses left undelivered: %+v %v", pending, err)
	}
}

func TestFollowUpQuestionThread(t *testing.T) {
	t.Parallel()
	e := newEnv(t)
	m := e.manager()

	sess, err := m.Create(e.ctx, e.task, e.user, colorPrompt, testModel)
	if err != nil {
		t.Fatalf("create session: %v", err)
	}
	e.waitIdle()

	questions := e.openQuestions(sess.Id)
	if len(questions) != 1 {
		t.Fatalf("want one open question, got %+v", questions)
	}
	first := questions[0]

	e.respond(m, first, api.InboxResponseTypeQuestion, "", "Before I pick: which of the two did you recommend, and why?", api.Open)
	e.waitIdle()

	questions = e.openQuestions(sess.Id)
	if len(questions) != 1 {
		t.Fatalf("want one open question after the follow-up, got %+v", questions)
	}
	second := questions[0]
	if second.Id == first.Id || second.ThreadId != first.ThreadId {
		t.Fatalf("follow-up didn't continue the thread: first %s/%s, second %s/%s", first.Id, first.ThreadId, second.Id, second.ThreadId)
	}

	old, err := e.store.GetInboxItem(e.ctx, first.Id)
	if err != nil {
		t.Fatalf("get first question: %v", err)
	}
	if old.State != api.Superseded {
		t.Fatalf("first question is %s, want superseded", old.State)
	}

	e.respond(m, second, api.InboxResponseTypeFreeForm, "", "Use green instead of either.", api.Answered)
	e.waitIdle()

	out, err := os.ReadFile(filepath.Join(e.task.WorktreePath, "color.txt"))
	if err != nil {
		t.Fatalf("agent didn't write color.txt: %v", err)
	}
	if strings.TrimSpace(string(out)) != "green" {
		t.Fatalf("color.txt contains %q, want green", out)
	}
}

func TestConcurrencyQueue(t *testing.T) {
	t.Parallel()
	e := newEnv(t)
	m := agents.New(e.store, agents.Limits{Global: 1, PerVendor: map[api.AgentSessionVendor]int{api.Claude: 1}})
	t.Cleanup(m.Close)

	first, err := m.Create(e.ctx, e.task, e.user, longCommandPrompt, testModel)
	if err != nil {
		t.Fatalf("create first: %v", err)
	}
	e.waitFor("first tool call", func(ev api.Event) bool { return ev.Kind == "agent.tool_call" })

	second, err := m.Create(e.ctx, e.addTask(), e.user, "Reply with the single word ok.", testModel)
	if err != nil {
		t.Fatalf("create second: %v", err)
	}
	if second.State != api.AgentSessionStateQueued {
		t.Fatalf("second session is %s with one slot taken, want queued", second.State)
	}

	usage, err := m.Usage(e.ctx)
	if err != nil || usage.Running != 1 || usage.Queued != 1 {
		t.Fatalf("usage %+v %v", usage, err)
	}

	cancelled, err := m.Interrupt(e.ctx, second.Id)
	if err != nil || cancelled.State != api.AgentSessionStateIdle {
		t.Fatalf("cancel queued turn: %+v %v", cancelled, err)
	}

	second, err = m.Send(e.ctx, second.Id, login, "Reply with the single word ok.")
	if err != nil || second.State != api.AgentSessionStateQueued {
		t.Fatalf("requeue: %+v %v", second, err)
	}

	_, err = m.Interrupt(e.ctx, first.Id)
	if err != nil {
		t.Fatalf("interrupt first: %v", err)
	}

	e.waitFor("first idle", func(ev api.Event) bool {
		return ev.Kind == "session.state" && *ev.AgentSessionId == first.Id && ev.Payload["state"] == string(api.AgentSessionStateIdle)
	})
	events := e.waitFor("second idle", func(ev api.Event) bool {
		return ev.Kind == "session.state" && *ev.AgentSessionId == second.Id && ev.Payload["state"] == string(api.AgentSessionStateIdle)
	})

	replied := false
	for _, ev := range events {
		if ev.Kind == "agent.message" && *ev.AgentSessionId == second.Id && strings.Contains(strings.ToLower(ev.Payload["text"].(string)), "ok") {
			replied = true
		}
	}
	if !replied {
		t.Fatalf("queued turn never ran: %v", kinds(events))
	}
}

func scratchRepo(t *testing.T) string {
	dir := t.TempDir()
	gitOut(t, dir, "init", "--initial-branch=main")
	gitOut(t, dir, "commit", "--allow-empty", "--message", "root")
	return dir
}

func gitOut(t *testing.T, dir string, args ...string) string {
	t.Helper()
	cmd := exec.Command("git", args...)
	cmd.Dir = dir
	cmd.Env = append(os.Environ(),
		"GIT_AUTHOR_NAME=Root",
		"GIT_AUTHOR_EMAIL=root@byre.invalid",
		"GIT_COMMITTER_NAME=Root",
		"GIT_COMMITTER_EMAIL=root@byre.invalid",
	)
	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Fatalf("git %v: %v\n%s", args, err, out)
	}
	return strings.TrimSpace(string(out))
}