byre/internal/agents/snapshots.go

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) afterTool(l *live) claude.HookFunc {
	return func(ctx context.Context, input json.RawMessage) (map[string]any, error) {
		m.snapshotSoon(l, api.SnapshotTriggerToolCall)
		text := m.deliverMidTurn(ctx, l)
		if text == "" {
			return map[string]any{}, nil
		}
		return map[string]any{
			"hookSpecificOutput": map[string]any{
				"hookEventName":     "PostToolUse",
				"additionalContext": text,
			},
		}, 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)
			}
		}
	}
}