Snapshot task worktrees to hidden refs as agents work
This commit is contained in:
@@ -122,6 +122,8 @@ func (m *Manager) ask(ctx context.Context, l *live, raw json.RawMessage) (string
|
||||
l.asked = true
|
||||
m.mu.Unlock()
|
||||
|
||||
m.snapshotSoon(l, api.SnapshotTriggerAskHuman)
|
||||
|
||||
if *args.Blocking {
|
||||
return fmt.Sprintf("Question %s filed in the human's inbox. It is blocking: end your turn now. The answer will arrive as your next message.", id), nil
|
||||
}
|
||||
|
||||
@@ -51,6 +51,9 @@ type Manager struct {
|
||||
deliverMu sync.Mutex
|
||||
capMu sync.Mutex
|
||||
|
||||
snapMu sync.Mutex
|
||||
snaps map[string]*snapState
|
||||
|
||||
bad *badLines
|
||||
watchStop chan struct{}
|
||||
}
|
||||
@@ -85,6 +88,7 @@ func New(st *store.Store, limits Limits, worktreeRoot string) *Manager {
|
||||
live: map[string]*live{},
|
||||
rateLimits: map[api.AgentSessionVendor]map[string]any{},
|
||||
merges: map[string]string{},
|
||||
snaps: map[string]*snapState{},
|
||||
bad: newBadLines(),
|
||||
watchStop: make(chan struct{}),
|
||||
}
|
||||
@@ -565,6 +569,7 @@ func (m *Manager) ensureLive(ctx context.Context, sess api.AgentSession, task ap
|
||||
Hooks: map[string][]claude.HookMatcher{
|
||||
"UserPromptSubmit": {{Func: remindGuidance}},
|
||||
"Stop": {{Func: m.bounceProseQuestion(l)}},
|
||||
"PostToolUse": {{Func: m.snapshotAfterTool(l)}},
|
||||
},
|
||||
OnBadLine: func(line []byte, reason string) { m.recordBadLine(l, line, reason) },
|
||||
}
|
||||
@@ -588,6 +593,7 @@ func (m *Manager) ensureLive(ctx context.Context, sess api.AgentSession, task ap
|
||||
|
||||
m.live[sess.Id] = l
|
||||
go m.pump(l)
|
||||
go m.snapshotTimer(l)
|
||||
return l, nil
|
||||
}
|
||||
|
||||
@@ -712,6 +718,7 @@ func (m *Manager) handle(ctx context.Context, l *live, ev claude.Event) error {
|
||||
if sha != "" {
|
||||
payload["commit"] = sha
|
||||
}
|
||||
m.snapshotSoon(l, api.SnapshotTriggerTurnEnd)
|
||||
switch {
|
||||
case handoff:
|
||||
payload["handoff"] = true
|
||||
|
||||
@@ -187,6 +187,24 @@ func TestTurnAndFollowUp(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
if !has(events, "snapshot.created") {
|
||||
events = append(events, e.waitFor("snapshot", func(ev api.Event) bool { return ev.Kind == "snapshot.created" })...)
|
||||
}
|
||||
snaps, err := e.store.ListSnapshots(e.ctx, e.task.Id)
|
||||
if err != nil || len(snaps) == 0 {
|
||||
t.Fatalf("no snapshots after a turn that wrote a file: %+v %v", snaps, err)
|
||||
}
|
||||
latest := snaps[len(snaps)-1]
|
||||
if ref := gitOut(t, e.repo, "rev-parse", git.SnapshotRef(e.task.Id, latest.N)); ref != latest.CommitSha {
|
||||
t.Fatalf("snapshot ref %s, want %s", ref, latest.CommitSha)
|
||||
}
|
||||
if got := gitOut(t, e.repo, "show", latest.CommitSha+":hi.txt"); strings.TrimSpace(got) != "hi" {
|
||||
t.Fatalf("hi.txt in the latest snapshot: %q", got)
|
||||
}
|
||||
if snaps[0].Trigger != api.SnapshotTriggerToolCall {
|
||||
t.Fatalf("first snapshot trigger %s, want tool_call", snaps[0].Trigger)
|
||||
}
|
||||
|
||||
window := 0
|
||||
for _, ev := range events {
|
||||
if ev.Kind == "session.context" {
|
||||
|
||||
@@ -137,6 +137,10 @@ func (m *Manager) RemoveTask(ctx context.Context, task api.Task, project api.Pro
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
err = m.dropSnapshots(project.RepoPath, task.Id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return git.RemoveWorktree(project.RepoPath, task.WorktreePath, task.Branch)
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,158 @@
|
||||
package agents
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"g.fc.run/misc/byre/internal/api"
|
||||
"g.fc.run/misc/byre/internal/claude"
|
||||
"g.fc.run/misc/byre/internal/git"
|
||||
"g.fc.run/misc/byre/internal/store"
|
||||
)
|
||||
|
||||
const snapshotInterval = 30 * time.Second
|
||||
|
||||
type snapState struct {
|
||||
mu sync.Mutex
|
||||
running bool
|
||||
pending api.SnapshotTrigger
|
||||
closed bool
|
||||
}
|
||||
|
||||
func (m *Manager) snapStateLocked(taskID string) *snapState {
|
||||
st := m.snaps[taskID]
|
||||
if st == nil {
|
||||
st = &snapState{}
|
||||
m.snaps[taskID] = st
|
||||
}
|
||||
return st
|
||||
}
|
||||
|
||||
func (m *Manager) snapshotSoon(l *live, trigger api.SnapshotTrigger) {
|
||||
m.snapMu.Lock()
|
||||
defer m.snapMu.Unlock()
|
||||
|
||||
st := m.snapStateLocked(l.taskID)
|
||||
if st.closed {
|
||||
return
|
||||
}
|
||||
if st.running {
|
||||
st.pending = trigger
|
||||
return
|
||||
}
|
||||
st.running = true
|
||||
go m.runSnapshots(l, st, trigger)
|
||||
}
|
||||
|
||||
func (m *Manager) runSnapshots(l *live, st *snapState, trigger api.SnapshotTrigger) {
|
||||
for {
|
||||
st.mu.Lock()
|
||||
m.snapMu.Lock()
|
||||
closed := st.closed
|
||||
m.snapMu.Unlock()
|
||||
if !closed {
|
||||
err := m.takeSnapshot(context.Background(), l, trigger)
|
||||
if err != nil {
|
||||
slog.Error("take snapshot", "task", l.taskID, "trigger", trigger, "err", err)
|
||||
}
|
||||
}
|
||||
st.mu.Unlock()
|
||||
|
||||
m.snapMu.Lock()
|
||||
trigger = st.pending
|
||||
st.pending = ""
|
||||
if trigger == "" || st.closed {
|
||||
st.running = false
|
||||
m.snapMu.Unlock()
|
||||
return
|
||||
}
|
||||
m.snapMu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Manager) takeSnapshot(ctx context.Context, l *live, trigger api.SnapshotTrigger) error {
|
||||
seq, err := m.store.LastTaskSeq(ctx, l.taskID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
tree, err := git.Snapshot(l.dir)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
last, err := m.store.LastSnapshot(ctx, l.taskID)
|
||||
if err != nil && !errors.Is(err, store.ErrNotFound) {
|
||||
return err
|
||||
}
|
||||
if err == nil && last.TreeSha == tree {
|
||||
return nil
|
||||
}
|
||||
|
||||
n := last.N + 1
|
||||
commit, err := git.CommitSnapshot(l.dir, tree, git.SnapshotRef(l.taskID, n), fmt.Sprintf("byre snapshot %d (%s)", n, trigger), []string{
|
||||
"GIT_AUTHOR_NAME=" + l.gitName,
|
||||
"GIT_AUTHOR_EMAIL=" + l.gitEmail,
|
||||
"GIT_COMMITTER_NAME=" + l.gitName,
|
||||
"GIT_COMMITTER_EMAIL=" + l.gitEmail,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return m.store.CreateSnapshot(ctx, api.Snapshot{
|
||||
Id: store.NewID(),
|
||||
TaskId: l.taskID,
|
||||
AgentSessionId: l.id,
|
||||
N: n,
|
||||
CommitSha: commit,
|
||||
TreeSha: tree,
|
||||
EventSeq: seq,
|
||||
Trigger: trigger,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
}, l.projectID)
|
||||
}
|
||||
|
||||
func (m *Manager) dropSnapshots(repo, taskID string) error {
|
||||
m.snapMu.Lock()
|
||||
st := m.snapStateLocked(taskID)
|
||||
st.closed = true
|
||||
m.snapMu.Unlock()
|
||||
|
||||
st.mu.Lock()
|
||||
defer st.mu.Unlock()
|
||||
return git.DeleteSnapshotRefs(repo, taskID)
|
||||
}
|
||||
|
||||
func (m *Manager) snapshotAfterTool(l *live) claude.HookFunc {
|
||||
return func(ctx context.Context, input json.RawMessage) (map[string]any, error) {
|
||||
m.snapshotSoon(l, api.SnapshotTriggerToolCall)
|
||||
return map[string]any{}, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Manager) snapshotTimer(l *live) {
|
||||
ticker := time.NewTicker(snapshotInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-l.sess.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
sess, err := m.store.GetSession(context.Background(), l.id)
|
||||
if err != nil {
|
||||
slog.Error("check session for snapshot", "session", l.id, "err", err)
|
||||
continue
|
||||
}
|
||||
if sess.State == api.AgentSessionStateRunning {
|
||||
m.snapshotSoon(l, api.SnapshotTriggerTimer)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
package agents
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"g.fc.run/misc/byre/internal/api"
|
||||
"g.fc.run/misc/byre/internal/git"
|
||||
"g.fc.run/misc/byre/internal/store"
|
||||
)
|
||||
|
||||
func TestSnapshots(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
events := make(chan api.Event, 256)
|
||||
st, err := store.Open(filepath.Join(t.TempDir(), "byre.db"), func(ev api.Event) { events <- ev })
|
||||
if err != nil {
|
||||
t.Fatalf("open store: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { st.Close() })
|
||||
|
||||
repo := t.TempDir()
|
||||
gitIn := func(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))
|
||||
}
|
||||
gitIn(repo, "init", "--initial-branch=main")
|
||||
err = os.WriteFile(filepath.Join(repo, "tracked.txt"), []byte("one\n"), 0o644)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
gitIn(repo, "add", "tracked.txt")
|
||||
gitIn(repo, "commit", "--message", "root")
|
||||
|
||||
now := time.Now().UTC()
|
||||
_, err = st.EnsureUser(ctx, "alice", "Alice")
|
||||
if err != nil {
|
||||
t.Fatalf("ensure user: %v", err)
|
||||
}
|
||||
project := api.Project{Id: "p1", Name: "demo", RepoPath: repo, DefaultBranch: "main", Remote: "none", CreatedBy: "alice", CreatedAt: now}
|
||||
err = st.CreateProject(ctx, project)
|
||||
if err != nil {
|
||||
t.Fatalf("create project: %v", err)
|
||||
}
|
||||
task := api.Task{Id: "t1", ProjectId: "p1", Title: "snap", Branch: "byre/t1", WorktreePath: filepath.Join(t.TempDir(), "t1"), Status: api.TaskStatusActive, CreatedBy: "alice", CreatedAt: now}
|
||||
task.BaseRef, err = git.AddWorktree(repo, task.WorktreePath, task.Branch, "main")
|
||||
if err != nil {
|
||||
t.Fatalf("add worktree: %v", err)
|
||||
}
|
||||
err = st.CreateTask(ctx, task)
|
||||
if err != nil {
|
||||
t.Fatalf("create task: %v", err)
|
||||
}
|
||||
err = st.CreateSession(ctx, api.AgentSession{Id: "s1", TaskId: "t1", Vendor: api.Claude, VendorSessionId: "s1", State: api.AgentSessionStateIdle, Role: api.Implementer, CreatedBy: "alice", CreatedAt: now}, "p1")
|
||||
if err != nil {
|
||||
t.Fatalf("create session: %v", err)
|
||||
}
|
||||
|
||||
m := New(st, DefaultLimits, t.TempDir())
|
||||
l := &live{id: "s1", dir: task.WorktreePath, gitName: "Alice Test", gitEmail: "alice@byre.invalid", vendor: api.Claude, taskID: "t1", projectID: "p1"}
|
||||
write := func(name, content string) {
|
||||
t.Helper()
|
||||
err := os.WriteFile(filepath.Join(task.WorktreePath, name), []byte(content), 0o644)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
snapshots := func() []api.Snapshot {
|
||||
t.Helper()
|
||||
snaps, err := st.ListSnapshots(ctx, "t1")
|
||||
if err != nil {
|
||||
t.Fatalf("list snapshots: %v", err)
|
||||
}
|
||||
return snaps
|
||||
}
|
||||
indexPath := gitIn(task.WorktreePath, "rev-parse", "--path-format=absolute", "--git-path", "index")
|
||||
indexBefore, err := os.ReadFile(indexPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
write("tracked.txt", "two\n")
|
||||
write("untracked.txt", "new\n")
|
||||
err = m.takeSnapshot(ctx, l, api.SnapshotTriggerToolCall)
|
||||
if err != nil {
|
||||
t.Fatalf("snapshot: %v", err)
|
||||
}
|
||||
snaps := snapshots()
|
||||
if len(snaps) != 1 || snaps[0].N != 1 || snaps[0].Trigger != api.SnapshotTriggerToolCall || snaps[0].AgentSessionId != "s1" {
|
||||
t.Fatalf("after the first change: %+v", snaps)
|
||||
}
|
||||
first := snaps[0]
|
||||
if ref := gitIn(repo, "rev-parse", "refs/byre/tasks/t1/snap/1"); ref != first.CommitSha {
|
||||
t.Fatalf("snap/1 is %s, want %s", ref, first.CommitSha)
|
||||
}
|
||||
if tree := gitIn(repo, "rev-parse", first.CommitSha+"^{tree}"); tree != first.TreeSha {
|
||||
t.Fatalf("commit tree %s, want %s", tree, first.TreeSha)
|
||||
}
|
||||
if parent := gitIn(repo, "rev-parse", first.CommitSha+"^"); parent != task.BaseRef {
|
||||
t.Fatalf("snapshot parent %s, want the worktree HEAD %s", parent, task.BaseRef)
|
||||
}
|
||||
if got := gitIn(repo, "show", first.CommitSha+":untracked.txt"); got != "new" {
|
||||
t.Fatalf("untracked.txt in snapshot: %q", got)
|
||||
}
|
||||
if got := gitIn(repo, "show", first.CommitSha+":tracked.txt"); got != "two" {
|
||||
t.Fatalf("tracked.txt in snapshot: %q", got)
|
||||
}
|
||||
|
||||
indexAfter, err := os.ReadFile(indexPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(indexBefore) != string(indexAfter) {
|
||||
t.Fatalf("snapshot modified the worktree's index")
|
||||
}
|
||||
if staged := gitIn(task.WorktreePath, "diff", "--cached", "--name-only"); staged != "" {
|
||||
t.Fatalf("snapshot staged files in the real index: %s", staged)
|
||||
}
|
||||
|
||||
err = m.takeSnapshot(ctx, l, api.SnapshotTriggerTimer)
|
||||
if err != nil {
|
||||
t.Fatalf("unchanged snapshot: %v", err)
|
||||
}
|
||||
if snaps := snapshots(); len(snaps) != 1 {
|
||||
t.Fatalf("an unchanged tree made a snapshot: %+v", snaps)
|
||||
}
|
||||
|
||||
write("untracked.txt", "newer\n")
|
||||
m.snapshotSoon(l, api.SnapshotTriggerTurnEnd)
|
||||
deadline := time.After(30 * time.Second)
|
||||
for created := false; !created; {
|
||||
select {
|
||||
case ev := <-events:
|
||||
created = ev.Kind == "snapshot.created" && ev.Payload["n"] == float64(2)
|
||||
case <-deadline:
|
||||
t.Fatalf("no snapshot.created event")
|
||||
}
|
||||
}
|
||||
snaps = snapshots()
|
||||
if len(snaps) != 2 || snaps[1].N != 2 || snaps[1].Trigger != api.SnapshotTriggerTurnEnd || snaps[1].EventSeq < first.EventSeq {
|
||||
t.Fatalf("after the second change: %+v", snaps)
|
||||
}
|
||||
if got := gitIn(repo, "show", "refs/byre/tasks/t1/snap/2:untracked.txt"); got != "newer" {
|
||||
t.Fatalf("untracked.txt in snap/2: %q", got)
|
||||
}
|
||||
|
||||
err = m.RemoveTask(ctx, task, project, "alice")
|
||||
if err != nil {
|
||||
t.Fatalf("remove task: %v", err)
|
||||
}
|
||||
_, err = st.DiscardTask(ctx, "t1", "alice")
|
||||
if err != nil {
|
||||
t.Fatalf("discard: %v", err)
|
||||
}
|
||||
if refs := gitIn(repo, "for-each-ref", "refs/byre/"); refs != "" {
|
||||
t.Fatalf("snapshot refs left after discard:\n%s", refs)
|
||||
}
|
||||
if snaps := snapshots(); len(snaps) != 2 {
|
||||
t.Fatalf("snapshot rows after discard: %+v", snaps)
|
||||
}
|
||||
|
||||
m.snapshotSoon(l, api.SnapshotTriggerToolCall)
|
||||
if refs := gitIn(repo, "for-each-ref", "refs/byre/"); refs != "" {
|
||||
t.Fatalf("snapshot taken after discard:\n%s", refs)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user