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

	worktrees string
}

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), worktrees: t.TempDir()}

	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.worktrees)
	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()

	startHead := gitOut(t, e.task.WorktreePath, "rev-parse", "HEAD")
	sess, err := m.Create(e.ctx, e.task, e.user, `Create a file hi.txt containing exactly the word 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", "turn.ended", "session.context"} {
		if !has(events, kind) {
			t.Fatalf("missing %s in %v", kind, kinds(events))
		}
	}

	for _, ev := range events {
		if ev.Kind != "turn.ended" {
			continue
		}
		id, _ := ev.Payload["inbox_item_id"].(string)
		if id == "" {
			t.Fatalf("turn.ended has no summary inbox item: %v", ev.Payload)
		}
		item, err := e.store.GetInboxItem(e.ctx, id)
		if err != nil {
			t.Fatalf("get summary item: %v", err)
		}
		if item.Kind != api.Report || item.Payload["text"] != ev.Payload["result"] {
			t.Fatalf("summary item %+v doesn't carry the turn's final message %q", item, ev.Payload["result"])
		}
	}

	window := 0
	for _, ev := range events {
		if ev.Kind == "session.context" {
			window = ev.Payload["window"].(int)
		}
	}
	if window <= 0 {
		t.Fatalf("no context window reported after the turn ended")
	}

	head := strings.TrimSpace(gitOut(t, e.task.WorktreePath, "rev-parse", "HEAD"))
	for _, ev := range events {
		if ev.Kind == "user.message" && ev.Payload["head"] != strings.TrimSpace(startHead) {
			t.Fatalf("user.message head %v, want the HEAD the turn started from %s", ev.Payload["head"], startHead)
		}
		if ev.Kind == "turn.ended" && ev.Payload["commit"] != head {
			t.Fatalf("turn.ended commit %v, want the new HEAD %s", ev.Payload["commit"], head)
		}
	}
	log := gitOut(t, e.task.WorktreePath, "log", "-1", "--format=%an <%ae>|%cn <%ce>", "--name-only")
	if !strings.HasPrefix(log, gitName+" <"+gitEmail+">|"+gitName+" <"+gitEmail+">") || !strings.Contains(log, "hi.txt") {
		t.Fatalf("byre's turn commit is wrong:\n%s", log)
	}
	if status := gitOut(t, e.task.WorktreePath, "status", "--porcelain", "--untracked-files=all"); status != "" {
		t.Fatalf("worktree dirty after the turn commit:\n%s", status)
	}

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

	events = e.waitIdle()
	reply := ""
	switched := false
	contextModel := ""
	for _, ev := range events {
		switch ev.Kind {
		case "agent.message":
			reply += ev.Payload["text"].(string)
		case "session.model":
			switched = ev.Payload["model"] == "haiku"
		case "session.context":
			contextModel = ev.Payload["model"].(string)
		}
	}
	if !strings.Contains(strings.ToLower(reply), "ok") {
		t.Fatalf("follow-up reply %q", reply)
	}
	if !switched || !strings.Contains(contextModel, "haiku") {
		t.Fatalf("model switch: session.model recorded %v, context model %q", switched, contextModel)
	}
	stored, err := e.store.GetSession(e.ctx, sess.Id)
	if err != nil {
		t.Fatalf("get session: %v", err)
	}
	if stored.Model != "haiku" {
		t.Fatalf("session model %q after switch", stored.Model)
	}

	_, 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, e.worktrees)

	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.`

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.Question {
			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.Option, "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 TestCompactCarriesHandoff(t *testing.T) {
	t.Parallel()
	e := newEnv(t)
	m := e.manager()

	old, err := m.Create(e.ctx, e.task, e.user, `Remember the codeword AZURE-17; a later message will ask for it. Reply with the single word ok.`, testModel)
	if err != nil {
		t.Fatalf("create session: %v", err)
	}
	e.waitIdle()

	_, err = m.Compact(e.ctx, old.Id)
	if err != nil {
		t.Fatalf("compact: %v", err)
	}
	events := e.waitFor("new session idle", func(ev api.Event) bool {
		return ev.Kind == "session.state" && ev.Payload["state"] == string(api.AgentSessionStateIdle) &&
			ev.AgentSessionId != nil && *ev.AgentSessionId != old.Id
	})

	for _, ev := range events {
		if ev.Kind == "turn.ended" && *ev.AgentSessionId == old.Id {
			if ev.Payload["handoff"] != true || ev.Payload["inbox_item_id"] != nil {
				t.Fatalf("handoff turn.ended %v: want handoff and no inbox summary", ev.Payload)
			}
		}
	}

	sessions, err := e.store.ListSessions(e.ctx, e.task.Id)
	if err != nil {
		t.Fatalf("list sessions: %v", err)
	}
	if len(sessions) != 2 {
		t.Fatalf("want the old and the new session, got %+v", sessions)
	}
	var next api.AgentSession
	for _, s := range sessions {
		switch s.Id {
		case old.Id:
			if s.State != api.AgentSessionStateCompleted {
				t.Fatalf("old session is %s, want completed", s.State)
			}
		default:
			next = s
		}
	}
	if next.State != api.AgentSessionStateIdle || next.Model != old.Model {
		t.Fatalf("new session %+v, want idle on model %q", next, old.Model)
	}

	_, err = m.Send(e.ctx, next.Id, login, "What was the codeword? Reply with just the codeword.", "")
	if err != nil {
		t.Fatalf("send to new session: %v", err)
	}
	reply := ""
	for _, ev := range e.waitIdle() {
		if ev.Kind == "agent.message" {
			reply += ev.Payload["text"].(string)
		}
	}
	if !strings.Contains(reply, "AZURE-17") {
		t.Fatalf("handoff lost the codeword; new session said %q", reply)
	}
}

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

	gitOut(t, e.task.WorktreePath, "-c", "user.name=x", "-c", "user.email=x@byre.invalid", "commit", "--allow-empty", "--message", "parent work")
	parentHead := strings.TrimSpace(gitOut(t, e.task.WorktreePath, "rev-parse", "HEAD"))

	_, err := m.Create(e.ctx, e.task, e.user, `Call the byre create_task tool once with title "child", base "this_task", and prompt "Reply with the single word ok." Then stop.`, testModel)
	if err != nil {
		t.Fatalf("create session: %v", err)
	}
	e.waitFor("child session", func(ev api.Event) bool {
		return ev.Kind == "session.created" && ev.TaskId != nil && *ev.TaskId != e.task.Id
	})

	tasks, err := e.store.ListTasks(e.ctx, e.project.Id)
	if err != nil {
		t.Fatalf("list tasks: %v", err)
	}
	var child api.Task
	for _, task := range tasks {
		if task.Id != e.task.Id {
			child = task
		}
	}
	if child.Title != "child" || child.CreatedBy != login || filepath.Dir(child.WorktreePath) != filepath.Join(e.worktrees, e.project.Id) {
		t.Fatalf("child task %+v", child)
	}

	sessions, err := e.store.ListSessions(e.ctx, child.Id)
	if err != nil || len(sessions) != 1 || sessions[0].Model != testModel || sessions[0].CreatedBy != login {
		t.Fatalf("child sessions %+v %v", sessions, err)
	}

	if child.BaseRef != parentHead {
		t.Fatalf("child base %s, want the parent task's HEAD %s", child.BaseRef, parentHead)
	}
}

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

	sess, err := m.Create(e.ctx, e.task, e.user, `I want a file hi.txt that says either red or blue. Before you create anything, find out from me which one.`, testModel)
	if err != nil {
		t.Fatalf("create session: %v", err)
	}
	events := e.waitIdle()

	questions := e.openQuestions(sess.Id)
	if len(questions) == 0 {
		t.Fatalf("the prose question never became an inbox question: %v", kinds(events))
	}
}

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.Other, "", "Before I pick: which of the two did you recommend, and why?", api.Answered)
	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)
	}

	e.respond(m, second, api.Other, "", "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}}, e.worktrees)
	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))
}