Run Claude sessions on tasks with event streaming and restart recovery
This commit is contained in:
Ian Gulliver
@@ -0,0 +1,321 @@
|
||||
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))
|
||||
}
|
||||
Reference in New Issue
Block a user