Snapshot task worktrees to hidden refs as agents work

internal/agents/snapshots.go

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