Fix checkpoint v2 rotation push handling

cmd/entire/cli/checkpoint/v2_generation.go

@@ -506,7 +506,7 @@ func (s *V2GitStore) RotateCurrentGenerationIfNeeded(ctx context.Context, maxChe
 		return "", false, fmt.Errorf("rotation: failed to determine next generation number: %w", err)
 	}
 
-	// Phase 1: Archive — create ref pointing to the current commit.
+	// Phase 1: Prepare archive and reset commits without changing refs yet.
 	// If the archive ref already exists, another instance already rotated — skip.
 	archiveRefName := ArchivedGenerationRefName(archiveNumber)
 	if _, refErr := s.repo.Reference(archiveRefName, true); refErr == nil {
@@ -515,22 +515,6 @@ func (s *V2GitStore) RotateCurrentGenerationIfNeeded(ctx context.Context, maxChe
 		)
 		return archiveRefName, false, nil
 	}
-	archiveRef := plumbing.NewHashReference(archiveRefName, currentRef.Hash())
-	if err := s.repo.Storer.SetReference(archiveRef); err != nil {
-		return "", false, fmt.Errorf("rotation: failed to create archived ref %s: %w", archiveRefName, err)
-	}
-
-	// Verify /full/current hasn't been advanced by another writer since we read it.
-	// If it changed, abort — the archive ref is harmless (points to a valid commit)
-	// and the next writer will trigger rotation again.
-	postArchiveRef, err := s.repo.Reference(refName, true)
-	if err != nil {
-		return "", false, fmt.Errorf("rotation: failed to re-read /full/current: %w", err)
-	}
-	if postArchiveRef.Hash() != currentRef.Hash() {
-		logging.Info(ctx, "rotation: /full/current changed during rotation, aborting reset")
-		return archiveRefName, false, nil
-	}
 
 	// Write generation.json to the current tree before archiving.
 	gen := s.computeGenerationTimestamps(currentTreeHash)
@@ -545,13 +529,7 @@ func (s *V2GitStore) RotateCurrentGenerationIfNeeded(ctx context.Context, maxChe
 		return "", false, fmt.Errorf("rotation: failed to create archive commit: %w", err)
 	}
 
-	// Update the archive ref to point to the commit with generation.json
-	archiveRef = plumbing.NewHashReference(archiveRefName, archiveCommitHash)
-	if err := s.repo.Storer.SetReference(archiveRef); err != nil {
-		return "", false, fmt.Errorf("rotation: failed to update archived ref %s: %w", archiveRefName, err)
-	}
-
-	// Phase 2: Create fresh orphan /full/current (empty tree, no generation.json)
+	// Create fresh orphan /full/current (empty tree, no generation.json).
 	emptyTreeHash, err := BuildTreeFromEntries(ctx, s.repo, make(map[string]object.TreeEntry))
 	if err != nil {
 		return "", false, fmt.Errorf("rotation: failed to build empty tree: %w", err)
@@ -562,6 +540,59 @@ func (s *V2GitStore) RotateCurrentGenerationIfNeeded(ctx context.Context, maxChe
 		return "", false, fmt.Errorf("rotation: failed to create orphan commit: %w", err)
 	}
 
+	// Verify /full/current hasn't been advanced by another writer since we read it.
+	// If it changed, abort before recording a publication marker.
+	postArchiveRef, err := s.repo.Reference(refName, true)
+	if err != nil {
+		return "", false, fmt.Errorf("rotation: failed to re-read /full/current: %w", err)
+	}
+	if postArchiveRef.Hash() != currentRef.Hash() {
+		logging.Info(ctx, "rotation: /full/current changed during rotation, aborting reset")
+		return archiveRefName, false, nil
+	}
+
+	publication := PendingV2FullGenerationPublication{
+		ArchiveRefName:           archiveRefName.String(),
+		ArchiveCommitHash:        archiveCommitHash.String(),
+		PreviousFullCurrentHash:  currentRef.Hash().String(),
+		ResetFullCurrentRootHash: orphanCommitHash.String(),
+		QueuedAt:                 time.Now().UTC(),
+	}
+
+	// The commit objects above are not reachable from the v2 refs yet. Record
+	// the pending publication before moving refs so pre-push can recover the
+	// intended rotation if the process stops after either ref update below.
+	if err := s.AppendPendingFullGenerationPublication(ctx, publication); err != nil {
+		return "", false, fmt.Errorf("rotation: failed to record pending full rotation: %w", err)
+	}
+	keepPendingPublication := false
+	defer func() {
+		if keepPendingPublication {
+			return
+		}
+		if removeErr := s.RemovePendingFullGenerationPublications(ctx, []PendingV2FullGenerationPublication{publication}); removeErr != nil {
+			logging.Warn(ctx, "rotation: failed to remove pending full rotation after failed rotation",
+				slog.String("error", removeErr.Error()),
+				slog.String("archive_ref", string(archiveRefName)),
+				slog.String("previous_full_current_hash", currentRef.Hash().String()),
+				slog.String("archive_commit_hash", archiveCommitHash.String()),
+				slog.String("reset_full_current_root_hash", orphanCommitHash.String()),
+			)
+		}
+	}()
+
+	// Phase 2: publish local refs after the pending publication marker exists.
+	if _, refErr := s.repo.Reference(archiveRefName, true); refErr == nil {
+		logging.Info(ctx, "rotation: archive ref already exists, skipping",
+			slog.String("archive_ref", string(archiveRefName)),
+		)
+		return archiveRefName, false, nil
+	}
+	archiveRef := plumbing.NewHashReference(archiveRefName, archiveCommitHash)
+	if err := s.repo.Storer.SetReference(archiveRef); err != nil {
+		return "", false, fmt.Errorf("rotation: failed to update archived ref %s: %w", archiveRefName, err)
+	}
+
 	reset, err := s.resetFullCurrentRefIfUnchanged(ctx, refName, postArchiveRef, orphanCommitHash)
 	if err != nil {
 		return "", false, err
@@ -570,6 +601,7 @@ func (s *V2GitStore) RotateCurrentGenerationIfNeeded(ctx context.Context, maxChe
 		return archiveRefName, false, nil
 	}
 
+	keepPendingPublication = true
 	logging.Info(ctx, "generation rotation complete",
 		slog.Int("archived_generation", archiveNumber),
 		slog.String("archive_ref", string(archiveRefName)),

cmd/entire/cli/checkpoint/v2_generation_test.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"os"
+	"path/filepath"
 	"testing"
 	"time"
 
@@ -447,6 +449,65 @@ func TestRotateGeneration_ArchivesCurrentAndCreatesNewOrphan(t *testing.T) {
 	assert.Empty(t, freshTree.Entries, "fresh tree should be empty (no generation.json)")
 }
 
+func TestRotateGeneration_FailsBeforeResetWhenPendingMarkerCannotBeRecorded(t *testing.T) {
+	t.Parallel()
+	repo := initTestRepo(t)
+	store := NewV2GitStore(repo, "origin")
+	ctx := context.Background()
+
+	populateFullCurrent(t, store, 3, 0)
+
+	worktree, err := repo.Worktree()
+	require.NoError(t, err)
+	blockingPath := filepath.Join(worktree.Filesystem().Root(), ".git", pendingV2FullGenerationPublicationDirName)
+	require.NoError(t, os.WriteFile(blockingPath, []byte("not a directory"), 0o600))
+
+	refName, rotated, err := store.RotateCurrentGenerationIfNeeded(ctx, 3)
+	require.Error(t, err)
+	require.False(t, rotated)
+	require.Empty(t, refName)
+	assert.Contains(t, err.Error(), "failed to record pending full rotation")
+
+	_, currentTreeHash, err := store.GetRefState(plumbing.ReferenceName(paths.V2FullCurrentRefName))
+	require.NoError(t, err)
+	currentCount, err := store.CountCheckpointsInTree(currentTreeHash)
+	require.NoError(t, err)
+	assert.Equal(t, 3, currentCount)
+
+	_, _, err = store.GetRefState(ArchivedGenerationRefName(1))
+	require.Error(t, err)
+}
+
+func TestRemovePendingFullGenerationPublications_PreservesLaterQueuedEntries(t *testing.T) {
+	t.Parallel()
+	repo := initTestRepo(t)
+	store := NewV2GitStore(repo, "origin")
+	ctx := context.Background()
+
+	first := PendingV2FullGenerationPublication{
+		ArchiveRefName:    paths.V2FullRefPrefix + "0000000000001",
+		ArchiveCommitHash: "1111111111111111111111111111111111111111",
+		QueuedAt:          time.Date(2026, 3, 19, 1, 2, 3, 0, time.UTC),
+	}
+	later := PendingV2FullGenerationPublication{
+		ArchiveRefName:    paths.V2FullRefPrefix + "0000000000002",
+		ArchiveCommitHash: "2222222222222222222222222222222222222222",
+		QueuedAt:          time.Date(2026, 3, 19, 4, 5, 6, 0, time.UTC),
+	}
+
+	require.NoError(t, store.AppendPendingFullGenerationPublication(ctx, first))
+	snapshot, err := store.ReadPendingFullGenerationPublications(ctx)
+	require.NoError(t, err)
+	require.Equal(t, []PendingV2FullGenerationPublication{first}, snapshot)
+
+	require.NoError(t, store.AppendPendingFullGenerationPublication(ctx, later))
+	require.NoError(t, store.RemovePendingFullGenerationPublications(ctx, snapshot))
+
+	remaining, err := store.ReadPendingFullGenerationPublications(ctx)
+	require.NoError(t, err)
+	assert.Equal(t, []PendingV2FullGenerationPublication{later}, remaining)
+}
+
 func TestResetFullCurrentRefIfUnchangedRejectsConcurrentChange(t *testing.T) {
 	t.Parallel()
 	repo := initTestRepo(t)