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/git"
	"g.fc.run/ai/byre/internal/hookcmd"
	"g.fc.run/ai/byre/internal/ids"
	"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, req.review)
			if err != nil {
				slog.Error("take snapshot", "task", t.taskID, "trigger", req.trigger, "err", err)
			}
			if snap == nil {
				m.dropRequest(t.taskID, req)
			}
		} 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, review 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(ctx, 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.LastReviewable(ctx, t.taskID)
	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(ctx, t.dir, tree, git.SnapshotRef(t.taskID, n), fmt.Sprintf("byre snapshot #%d (%s)", n, trigger), git.IdentityEnv(t.gitName, t.gitEmail))
	if err != nil {
		return nil, err
	}

	snap := api.Snapshot{
		Id:             ids.New(ids.Snapshot),
		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, review)
	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)
	}
	merging, err := m.store.MergeTurnPending(ctx, task.Id)
	if err != nil {
		return err
	}
	if merging {
		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
	}
	unlock := m.lockTask(task.Id)
	_, err = git.CommitAll(ctx, t.dir, "Work before take_snapshot", t.gitName, t.gitEmail)
	unlock()
	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 {
	snap, err := m.cutSnapshot(ctx, t, held, api.SnapshotTriggerAnswer)
	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
	}
	if snap == nil {
		return 0
	}
	return snap.N
}

func (m *Manager) cutSnapshot(ctx context.Context, t snapTarget, held bool, trigger api.SnapshotTrigger) (*api.Snapshot, error) {
	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 nil, nil
	}
	return m.takeSnapshot(ctx, t, trigger, 0, true, false)
}

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 "Worktree committed to the task's branch, and snapshot requested. If the project has a review workflow, its reviewers check it against the default branch; their findings come to you together once the workflow finishes.", nil
		},
	}
}

const takeSnapshotDescription = "Commit and snapshot a task's worktree now: byre commits the whole worktree, uncommitted work included, to the task's branch, and returns once that commit is made. It then snapshots that commit on byre's own ref for the task (refs/byre/tasks/<task>/snap/<n>). If the project has a review workflow, its reviewers check all of the task's work against the default branch, using each stage's model and instructions; findings arrive together once the workflow finishes. " +
	"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(context.Background(), repo, taskID)
}

func (m *Manager) afterTool(l *live) hookcmd.Func {
	return func(ctx context.Context, input json.RawMessage) (map[string]any, error) {
		in := struct {
			HookEventName string `json:"hook_event_name"`
		}{}
		err := json.Unmarshal(input, &in)
		if err != nil {
			return nil, err
		}
		m.holdWhilePaused(ctx, l, hookToolUseID(input))
		mt := m.deliverMidTurn(ctx, l)
		if mt.text == "" {
			return map[string]any{}, nil
		}
		return map[string]any{
			"hookSpecificOutput": map[string]any{
				"hookEventName":     in.HookEventName,
				"additionalContext": mt.text,
			},
		}, nil
	}
}