Snapshot task worktrees to hidden refs as agents work

This commit is contained in:
Ian Gulliver
2026-10-14 22:52:02 -07:00
parent c3ed243573
commit 9039860594
22 changed files with 926 additions and 202 deletions
+2
View File
@@ -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
}
+7
View File
@@ -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
+18
View File
@@ -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" {
+4
View File
@@ -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)
}
+158
View File
@@ -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)
}
}
}
}
+177
View File
@@ -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)
}
}