byre/internal/agents/snapshots_internal_test.go

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"), store.DefaultPortFirst, 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, "alice")
	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(), 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.SnapshotTriggerMessage, 0)
	if err != nil {
		t.Fatalf("snapshot: %v", err)
	}
	snaps := snapshots()
	if len(snaps) != 1 || snaps[0].N != 1 || snaps[0].Trigger != api.SnapshotTriggerMessage || 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.SnapshotTriggerMessage, 1)
	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, 0)
	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.SnapshotTriggerMessage, 0)
	if refs := gitIn(repo, "for-each-ref", "refs/byre/"); refs != "" {
		t.Fatalf("snapshot taken after discard:\n%s", refs)
	}
}