byre/internal/agents/snapshots.go

package agents

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"sync"
	"time"

	"g.fc.run/ai/byre/internal/api"
	"g.fc.run/ai/byre/internal/claude"
	"g.fc.run/ai/byre/internal/git"
	"g.fc.run/ai/byre/internal/mcp"
	"g.fc.run/ai/byre/internal/store"
)

type snapRequest struct {
	target  snapTarget
	trigger api.SnapshotTrigger
	seq     int64
	review  bool
	force   bool
}

type snapTarget struct {
	taskID    string
	projectID string
	sessionID string
	dir       string
	gitName   string
	gitEmail  string
}

func liveTarget(l *live) snapTarget {
	return snapTarget{taskID: l.taskID, projectID: l.projectID, sessionID: l.id, dir: l.dir, gitName: l.gitName, gitEmail: l.gitEmail}
}

type snapState struct {
	mu      sync.Mutex
	running bool
	pending *snapRequest
	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(t snapTarget, trigger api.SnapshotTrigger, seq int64, review, force bool) {
	m.snapMu.Lock()
	defer m.snapMu.Unlock()

	st := m.snapStateLocked(t.taskID)
	if st.closed {
		m.dropRequest(t.taskID, &snapRequest{review: review})
		return
	}
	req := &snapRequest{target: t, trigger: trigger, seq: seq, review: review, force: force}
	if st.running {
		if p := st.pending; p != nil && p.force {
			req.force = true
			req.trigger = p.trigger
		}
		m.dropRequest(t.taskID, st.pending)
		st.pending = req
		return
	}
	st.running = true
	go m.runSnapshots(st, req)
}

func (m *Manager) runSnapshots(st *snapState, req *snapRequest) {
	for {
		t := req.target
		unlock := m.lockTask(t.taskID)
		st.mu.Lock()
		m.snapMu.Lock()
		closed := st.closed
		m.snapMu.Unlock()
		if !closed {
			snap, err := m.takeSnapshot(context.Background(), t, req.trigger, req.seq, req.force)
			if err != nil {
				slog.Error("take snapshot", "task", t.taskID, "trigger", req.trigger, "err", err)
			}
			if req.review {
				go func() {
					m.afterSnapshot(snap)
					m.endReviewRequest(t.taskID)
				}()
			}
		} else {
			m.dropRequest(t.taskID, req)
		}
		st.mu.Unlock()
		unlock()

		m.snapMu.Lock()
		req = st.pending
		st.pending = nil
		if req == nil || st.closed {
			m.dropRequest(t.taskID, req)
			st.running = false
			m.snapMu.Unlock()
			return
		}
		m.snapMu.Unlock()
	}
}

func (m *Manager) dropRequest(taskID string, req *snapRequest) {
	if req != nil && req.review {
		m.endReviewRequest(taskID)
	}
}

func (m *Manager) endReviewRequest(taskID string) {
	err := m.store.EndReviewRequest(context.Background(), taskID)
	if err != nil {
		slog.Error("end review request", "task", taskID, "err", err)
	}
}

func (m *Manager) takeSnapshot(ctx context.Context, t snapTarget, trigger api.SnapshotTrigger, seq int64, force bool) (*api.Snapshot, error) {
	if seq == 0 {
		var err error
		seq, err = m.store.LastTaskSeq(ctx, t.taskID)
		if err != nil {
			return nil, err
		}
	}

	tree, err := git.Snapshot(t.dir)
	if err != nil {
		return nil, err
	}

	last, err := m.store.LastSnapshot(ctx, t.taskID)
	if err != nil && !errors.Is(err, store.ErrNotFound) {
		return nil, err
	}
	reviewed, err := m.store.LastSnapshotNot(ctx, t.taskID, api.SnapshotTriggerAnswer)
	if err != nil && !errors.Is(err, store.ErrNotFound) {
		return nil, err
	}
	if err == nil && reviewed.TreeSha == tree && !force {
		return nil, nil
	}

	n := last.N + 1
	commit, err := git.CommitSnapshot(t.dir, tree, git.SnapshotRef(t.taskID, n), fmt.Sprintf("byre snapshot %d (%s)", n, trigger), []string{
		"GIT_AUTHOR_NAME=" + t.gitName,
		"GIT_AUTHOR_EMAIL=" + t.gitEmail,
		"GIT_COMMITTER_NAME=" + t.gitName,
		"GIT_COMMITTER_EMAIL=" + t.gitEmail,
	})
	if err != nil {
		return nil, err
	}

	snap := api.Snapshot{
		Id:             store.NewID(),
		TaskId:         t.taskID,
		AgentSessionId: t.sessionID,
		N:              n,
		CommitSha:      commit,
		TreeSha:        tree,
		EventSeq:       seq,
		Trigger:        trigger,
		CreatedAt:      time.Now().UTC(),
	}
	err = m.store.CreateSnapshot(ctx, snap, t.projectID)
	if err != nil {
		return nil, err
	}
	return &snap, nil
}

func (m *Manager) RequestSnapshot(ctx context.Context, taskID string) error {
	task, _, err := m.activeTask(ctx, taskID)
	if err != nil {
		return err
	}
	discard, err := m.store.GetDiscard(ctx, task.Id)
	if err != nil {
		return err
	}
	if discard.State != store.DiscardNone {
		return fmt.Errorf("%w: task %s is being discarded", ErrConflict, task.Id)
	}
	mergingBy, err := m.store.MergingBy(ctx, task.Id)
	if err != nil {
		return err
	}
	if mergingBy != "" {
		return fmt.Errorf("%w: task %s is merging", ErrConflict, task.Id)
	}
	sessions, err := m.store.ListSessions(ctx, task.Id)
	if err != nil {
		return err
	}
	if len(sessions) == 0 {
		return fmt.Errorf("%w: task %s has no agent session to attribute a snapshot to", ErrConflict, task.Id)
	}
	t, err := m.sessionTarget(ctx, sessions[len(sessions)-1], task)
	if err != nil {
		return err
	}
	err = m.store.BeginReviewRequest(ctx, task.Id)
	if err != nil {
		return err
	}
	m.snapshotSoon(t, api.SnapshotTriggerRequested, 0, true, true)
	return nil
}

func (m *Manager) sessionTarget(ctx context.Context, sess api.AgentSession, task api.Task) (snapTarget, error) {
	creator, err := m.store.GetUser(ctx, sess.CreatedBy)
	if err != nil {
		return snapTarget{}, err
	}
	return snapTarget{taskID: task.Id, projectID: task.ProjectId, sessionID: sess.Id, dir: task.WorktreePath, gitName: creator.GitName, gitEmail: creator.GitEmail}, nil
}

func (m *Manager) snapshotForAnswer(ctx context.Context, t snapTarget, held bool) int {
	if !held {
		unlock := m.lockTask(t.taskID)
		defer unlock()
	}
	m.snapMu.Lock()
	st := m.snapStateLocked(t.taskID)
	m.snapMu.Unlock()
	st.mu.Lock()
	defer st.mu.Unlock()
	m.snapMu.Lock()
	closed := st.closed
	m.snapMu.Unlock()
	if closed {
		return 0
	}

	snap, err := m.takeSnapshot(ctx, t, api.SnapshotTriggerAnswer, 0, true)
	if err != nil {
		slog.Error("take snapshot", "task", t.taskID, "trigger", api.SnapshotTriggerAnswer, "err", err)
		m.recordError(&live{id: t.sessionID, taskID: t.taskID, projectID: t.projectID}, fmt.Sprintf("byre couldn't snapshot the worktree before handing over an answer, so the answer went to the agent without one: %v", err))
		return 0
	}
	return snap.N
}

func (m *Manager) takeSnapshotTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name:        "take_snapshot",
		Description: takeSnapshotDescription + " It snapshots your own task's worktree.",
		InputSchema: json.RawMessage(`{"type": "object", "properties": {}}`),
		Handler: func(ctx context.Context, args json.RawMessage) (string, error) {
			err := m.RequestSnapshot(ctx, l.taskID)
			if err != nil {
				return "", err
			}
			return "Snapshot requested. A reviewer reviews it against the default branch, and any findings come to you as a message from byre.", nil
		},
	}
}

const takeSnapshotDescription = "Snapshot a task's worktree now, uncommitted work included, and start a review of it: a reviewer on the other vendor reviews all of the task's work against the default branch. " +
	"byre already does this at the end of every turn that changes files; call this to review work mid-turn, or to retry a review that failed (it snapshots even when nothing changed since the last snapshot)."

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) {
		mt := m.deliverMidTurn(ctx, l)
		if mt.text == "" {
			return map[string]any{}, nil
		}
		return map[string]any{
			"hookSpecificOutput": map[string]any{
				"hookEventName":     "PostToolUse",
				"additionalContext": mt.text,
			},
		}, nil
	}
}