From af4996c585649679e88dee0d64972c115fa49692 Mon Sep 17 00:00:00 2001
From: Ian Gulliver <git@flamingcow.io>
Date: Wed, 21 Oct 2026 12:25:32 -0700
Subject: [PATCH] Centralize snapshot review selection in the store

---
 internal/agents/afterreview.go |  6 +++---
 internal/agents/ready.go       |  2 +-
 internal/agents/review.go      | 18 ++----------------
 internal/store/review.go       |  8 ++++++++
 4 files changed, 14 insertions(+), 20 deletions(-)

diff --git a/internal/agents/afterreview.go b/internal/agents/afterreview.go
index 278a16b..d74f022 100644
--- a/internal/agents/afterreview.go
+++ b/internal/agents/afterreview.go
@@ -41,13 +41,13 @@ func (m *Manager) MergeAfterReview(ctx context.Context, task api.Task, user api.
 	if err != nil {
 		return api.Task{}, err
 	}
-	reviews, err := m.store.ListReviews(ctx, task.Id)
+	reviews, err := m.store.SnapshotReviews(ctx, task.Id, latest.N)
 	if err != nil {
 		return api.Task{}, err
 	}
 	reviewID := ""
 	for _, r := range reviews {
-		if r.SnapshotN == latest.N && r.State == api.ReviewStateRunning {
+		if r.State == api.ReviewStateRunning {
 			reviewID = r.Id
 		}
 	}
@@ -99,7 +99,7 @@ func (m *Manager) reviewFinished(reviewID string, deliver bool) {
 	if m.advanceReviews(ctx, r) {
 		return
 	}
-	reviews, err := m.snapshotReviews(ctx, r.TaskId, r.SnapshotN)
+	reviews, err := m.store.SnapshotReviews(ctx, r.TaskId, r.SnapshotN)
 	if err != nil {
 		slog.Error("list snapshot reviews", "task", r.TaskId, "err", err)
 		return
diff --git a/internal/agents/ready.go b/internal/agents/ready.go
index 0359aa0..ef02c4c 100644
--- a/internal/agents/ready.go
+++ b/internal/agents/ready.go
@@ -201,7 +201,7 @@ func (m *Manager) tryPromptReady(ctx context.Context, taskID string, fromReview
 	if err != nil {
 		return err
 	}
-	reviews, err := m.snapshotReviews(ctx, taskID, latest.N)
+	reviews, err := m.store.SnapshotReviews(ctx, taskID, latest.N)
 	if err != nil {
 		return err
 	}
diff --git a/internal/agents/review.go b/internal/agents/review.go
index e58aed1..c4b3629 100644
--- a/internal/agents/review.go
+++ b/internal/agents/review.go
@@ -164,20 +164,6 @@ func pendingReviewStages(def api.WorkflowDefinition, stage string, byStage map[s
 	return pending
 }
 
-func (m *Manager) snapshotReviews(ctx context.Context, taskID string, n int) ([]api.Review, error) {
-	reviews, err := m.store.ListReviews(ctx, taskID)
-	if err != nil {
-		return nil, err
-	}
-	out := []api.Review{}
-	for _, r := range reviews {
-		if r.SnapshotN == n {
-			out = append(out, r)
-		}
-	}
-	return out, nil
-}
-
 func (m *Manager) reviewDir(projectID, reviewID string) string {
 	return filepath.Join(m.worktreeRoot, projectID, "review-"+reviewID)
 }
@@ -596,13 +582,13 @@ func (m *Manager) respondToFindingTool(l *live) mcp.Tool {
 			if reason == "" {
 				return "", fmt.Errorf("%w: reason must not be empty", ErrInvalid)
 			}
-			reviews, err := m.store.ListReviews(ctx, l.taskID)
+			reviews, err := m.store.SnapshotReviews(ctx, l.taskID, args.Snapshot)
 			if err != nil {
 				return "", err
 			}
 			reviewID := ""
 			for _, r := range reviews {
-				if r.SnapshotN == args.Snapshot && r.Stage == args.Reviewer && r.State == api.ReviewStateDone {
+				if r.Stage == args.Reviewer && r.State == api.ReviewStateDone {
 					reviewID = r.Id
 				}
 			}
diff --git a/internal/store/review.go b/internal/store/review.go
index 28afb92..b148d13 100644
--- a/internal/store/review.go
+++ b/internal/store/review.go
@@ -69,6 +69,14 @@ func (s *Store) ListReviews(ctx context.Context, taskID string) ([]api.Review, e
 	return scanReviews(rows)
 }
 
+func (s *Store) SnapshotReviews(ctx context.Context, taskID string, n int) ([]api.Review, error) {
+	rows, err := s.db.QueryContext(ctx, `SELECT `+reviewColumns+` FROM reviews WHERE task_id = ? AND snapshot_n = ? ORDER BY created_at`, taskID, n)
+	if err != nil {
+		return nil, err
+	}
+	return scanReviews(rows)
+}
+
 func (s *Store) RunningReviews(ctx context.Context) ([]api.Review, error) {
 	rows, err := s.db.QueryContext(ctx, `SELECT `+reviewColumns+` FROM reviews WHERE state = ? ORDER BY created_at`, string(api.ReviewStateRunning))
 	if err != nil {