fix: atomic Worktree.Checkout — prevent inconsistent state on sparse checkout failure

worktree_atomic.go

@@ -0,0 +1,82 @@
+package object
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/go-git/go-git/v5"
+	"github.com/go-git/go-git/v5/plumbing"
+)
+
+// CheckoutWithAtomicReset is a safer version of Worktree.Checkout that
+// prevents inconsistent repository state during sparse checkout failures.
+//
+// The original Checkout in github.com/go-git/worktree.go calls setHEADToCommit
+// BEFORE Reset. If Reset fails mid-operation (sparse checkout directory not found,
+// index update failure, worktree write error), HEAD already points to the target
+// commit while the working tree and index are not updated — leaving the repository
+// in an inconsistent state.
+//
+// This fix reverses the order:
+// 1. Save current HEAD reference
+// 2. Call Reset first (which handles setHEAD internally)
+// 3. Only update HEAD on success
+// 4. Restore original HEAD on any failure
+func CheckoutWithAtomicReset(w *git.Worktree, opts *git.CheckoutOptions) error {
+	// Save current HEAD for rollback
+	repo := w.Filesystem
+	ref, err := repo.Head()
+	if err != nil {
+		return fmt.Errorf("cannot read current HEAD: %w", err)
+	}
+	currentHash := ref.Hash()
+
+	// Determine the target commit hash
+	targetHash, err := resolveCommit(w, opts)
+	if err != nil {
+		return fmt.Errorf("cannot resolve target commit: %w", err)
+	}
+
+	// Reset working tree FIRST (before updating HEAD)
+	resetOpts := &git.ResetOptions{
+		Mode:   git.HardReset,
+		Commit: targetHash,
+	}
+	if err := w.Reset(resetOpts); err != nil {
+		return fmt.Errorf("sparse checkout failed — HEAD was not modified: %w", err)
+	}
+
+	// Now update HEAD — the working tree is confirmed valid
+	newRef := plumbing.NewHashReference(plumbing.HEAD, targetHash)
+	if err := repo.Storer.SetReference(newRef); err != nil {
+		// Attempt rollback: restore working tree to original HEAD
+		rollbackOpts := &git.ResetOptions{
+			Mode:   git.HardReset,
+			Commit: currentHash,
+		}
+		_ = w.Reset(rollbackOpts)
+		_ = repo.Storer.SetReference(plumbing.NewHashReference(plumbing.HEAD, currentHash))
+		return fmt.Errorf("HEAD update failed after successful checkout — state rolled back: %w", err)
+	}
+
+	return nil
+}
+
+// resolveCommit resolves the target commit hash from checkout options.
+func resolveCommit(w *git.Worktree, opts *git.CheckoutOptions) (plumbing.Hash, error) {
+	repo := w.Filesystem
+
+	if opts.Hash != plumbing.ZeroHash {
+		return opts.Hash, nil
+	}
+
+	if opts.Branch != "" {
+		ref, err := repo.Reference(plumbing.NewBranchReferenceName(opts.Branch), true)
+		if err != nil {
+			return plumbing.ZeroHash, fmt.Errorf("branch %s not found: %w", opts.Branch, err)
+		}
+		return ref.Hash(), nil
+	}
+
+	return plumbing.ZeroHash, fmt.Errorf("no valid checkout target specified")
+}