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
	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)
	}

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

	id := store.NewID()
	e.task = api.Task{
		Id:           id,
		ProjectId:    p.Id,
		Title:        "test",
		Branch:       "byre/" + id,
		WorktreePath: filepath.Join(t.TempDir(), id),
		Status:       api.TaskStatusActive,
		CreatedBy:    login,
		CreatedAt:    time.Now().UTC(),
	}
	e.task.BaseRef, err = git.AddWorktree(repo, e.task.WorktreePath, e.task.Branch, "main")
	if err != nil {
		t.Fatalf("add worktree: %v", err)
	}
	err = e.store.CreateTask(ctx, e.task)
	if err != nil {
		t.Fatalf("create task: %v", err)
	}

	return e
}

func (e *env) manager() *agents.Manager {
	m := agents.New(e.store)
	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)

	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")
	}
}

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))
}