Restore stacks after incomplete cascade rebases

Roll back branches already rewritten when a later rebase cannot start or
final ancestry verification fails, preventing retries from replaying stale
history. Preserve retryable modify state without repeating completed work,
and add regression coverage for remote-qualified trunk normalization.

Author: skarim

cmd/rebase.go

@@ -226,7 +226,11 @@ func runRebase(cfg *config.Config, opts *rebaseOptions) error {
 
 	if rebaseResult.Err != nil {
 		cfg.Errorf("%v", rebaseResult.Err)
-		_ = git.CheckoutBranch(currentBranch)
+		if rebaseResult.Rebased {
+			restoreRebaseRefs(cfg, currentBranch, originalRefs)
+		} else {
+			_ = git.CheckoutBranch(currentBranch)
+		}
 		return ErrSilent
 	}
 
@@ -266,6 +270,9 @@ func runRebase(cfg *config.Config, opts *rebaseOptions) error {
 
 	if unstacked := verifyStacked(s, trunk.Ref, startIdx, endIdx); len(unstacked) > 0 {
 		reportUnstacked(cfg, trunk.Ref, unstacked)
+		if rebaseResult.Rebased {
+			restoreRebaseRefs(cfg, currentBranch, originalRefs)
+		}
 		return ErrSilent
 	}
 
@@ -408,6 +415,8 @@ func continueRebase(cfg *config.Config, gitDir string) error {
 
 		if result.Err != nil {
 			cfg.Errorf("%v", result.Err)
+			restoreRebaseRefs(cfg, state.OriginalBranch, state.OriginalRefs)
+			clearRebaseState(gitDir)
 			return ErrSilent
 		}
 
@@ -444,6 +453,8 @@ func continueRebase(cfg *config.Config, gitDir string) error {
 	}
 	if unstacked := verifyStacked(s, trunkBase, verifyStart, verifyEnd); len(unstacked) > 0 {
 		reportUnstacked(cfg, trunkRef, unstacked)
+		restoreRebaseRefs(cfg, state.OriginalBranch, state.OriginalRefs)
+		clearRebaseState(gitDir)
 		return ErrSilent
 	}
 

cmd/sync.go

@@ -153,14 +153,15 @@ func runSync(cfg *config.Config, opts *syncOptions) error {
 	// --- Step 3: Cascade rebase ---
 	needsRebase := trunk.Moved || len(updatedBranches) > 0 || stackNeedsRebase(s, trunk.Ref)
 	rebased := false
+	var originalRefs map[string]string
 	if needsRebase {
 		cfg.Printf("")
 		cfg.Printf("Rebasing stack ...")
 
 		// Sync PR state to detect merged PRs before rebasing.
 		_ = syncStackPRs(cfg, s)
 
-		originalRefs, err := resolveOriginalRefs(s)
+		originalRefs, err = resolveOriginalRefs(s)
 		if err != nil {
 			cfg.Warningf("Could not resolve branch SHAs — skipping rebase: %v", err)
 		} else {
@@ -175,7 +176,11 @@ func runSync(cfg *config.Config, opts *syncOptions) error {
 
 			if result.Err != nil {
 				cfg.Errorf("%v", result.Err)
-				_ = git.CheckoutBranch(currentBranch)
+				if result.Rebased {
+					restoreRebaseRefs(cfg, currentBranch, originalRefs)
+				} else {
+					_ = git.CheckoutBranch(currentBranch)
+				}
 				stack.SaveNonBlocking(gitDir, sf)
 				return ErrSilent
 			}
@@ -209,6 +214,9 @@ func runSync(cfg *config.Config, opts *syncOptions) error {
 	if unstacked := verifyStacked(s, trunk.Ref, 0, len(s.Branches)); len(unstacked) > 0 {
 		_ = git.CheckoutBranch(currentBranch)
 		reportUnstacked(cfg, trunk.Ref, unstacked)
+		if rebased && originalRefs != nil {
+			restoreRebaseRefs(cfg, currentBranch, originalRefs)
+		}
 		stack.SaveNonBlocking(gitDir, sf)
 		return ErrSilent
 	}
@@ -402,6 +410,12 @@ func runSync(cfg *config.Config, opts *syncOptions) error {
 func restoreBranches(originalRefs map[string]string) []string {
 	var errors []string
 	for branch, sha := range originalRefs {
+		if !git.BranchExists(branch) {
+			continue
+		}
+		if currentSHA, err := git.RevParse(branch); err == nil && currentSHA == sha {
+			continue
+		}
 		if err := git.CheckoutBranch(branch); err != nil {
 			errors = append(errors, fmt.Sprintf("checkout %s: %s", branch, err))
 			continue
@@ -413,6 +427,12 @@ func restoreBranches(originalRefs map[string]string) []string {
 	return errors
 }
 
+func restoreRebaseRefs(cfg *config.Config, originalBranch string, originalRefs map[string]string) {
+	restoreErrors := restoreBranches(originalRefs)
+	_ = git.CheckoutBranch(originalBranch)
+	reportRestoreStatus(cfg, restoreErrors)
+}
+
 // reportRestoreStatus prints whether branch restoration succeeded or partially failed.
 func reportRestoreStatus(cfg *config.Config, restoreErrors []string) {
 	if len(restoreErrors) > 0 {

cmd/sync_test.go

@@ -475,6 +475,11 @@ func TestSync_RebaseConflict_RestoresAll(t *testing.T) {
 	var checkouts []string
 	currentBranch := "b1"
 	abortCalled := false
+	branchSHAs := map[string]string{
+		"b1": "sha-b1",
+		"b2": "sha-b2",
+		"b3": "sha-b3",
+	}
 
 	mock := newSyncMock(tmpDir, "b1")
 	mock.RevParseFn = func(ref string) (string, error) {
@@ -484,6 +489,9 @@ func TestSync_RebaseConflict_RestoresAll(t *testing.T) {
 		if ref == "origin/main" {
 			return "remote-sha", nil
 		}
+		if sha, ok := branchSHAs[ref]; ok {
+			return sha, nil
+		}
 		return "sha-" + ref, nil
 	}
 	mock.IsAncestorFn = func(a, d string) (bool, error) {
@@ -495,7 +503,10 @@ func TestSync_RebaseConflict_RestoresAll(t *testing.T) {
 		currentBranch = name
 		return nil
 	}
-	mock.RebaseFn = func(string, git.RebaseOpts) error { return nil } // b1 succeeds
+	mock.RebaseFn = func(string, git.RebaseOpts) error {
+		branchSHAs["b1"] = "rebased-b1"
+		return nil
+	}
 	mock.RebaseOntoFn = func(newBase, oldBase, branch string, opts git.RebaseOpts) error {
 		if branch == "b2" {
 			return fmt.Errorf("conflict")
@@ -508,6 +519,7 @@ func TestSync_RebaseConflict_RestoresAll(t *testing.T) {
 	}
 	mock.ResetHardFn = func(ref string) error {
 		resets = append(resets, resetCall{currentBranch, ref})
+		branchSHAs[currentBranch] = ref
 		return nil
 	}
 
@@ -528,14 +540,15 @@ func TestSync_RebaseConflict_RestoresAll(t *testing.T) {
 	assert.Contains(t, output, "Conflict detected")
 	assert.Contains(t, output, "gh stack rebase")
 
-	// All branches should be restored
+	// The branch rewritten before the conflict should be restored. Unchanged
+	// branches are left alone.
 	resetMap := make(map[string]string)
 	for _, r := range resets {
 		resetMap[r.branch] = r.sha
 	}
 	assert.Equal(t, "sha-b1", resetMap["b1"])
-	assert.Equal(t, "sha-b2", resetMap["b2"])
-	assert.Equal(t, "sha-b3", resetMap["b3"])
+	assert.NotContains(t, resetMap, "b2")
+	assert.NotContains(t, resetMap, "b3")
 
 	_ = abortCalled // RebaseAbort is called if IsRebaseInProgress returns true
 }