fix: address review feedback for undo feature

- Use stable stash commit hash instead of stash@{0} ref
- Restore auto-stashed changes after successful operations
- Use nanosecond precision for snapshot timestamps to avoid collisions
- Fix StashPop to use 'git stash apply' for commit hash support
- Update test names to accurately reflect behavior

Author: boneskull

cmd/cascade.go

@@ -78,13 +78,26 @@ func runCascade(cmd *cobra.Command, args []string) error {
 	}
 
 	// Save undo snapshot (unless dry-run)
+	var stashRef string
 	if !cascadeDryRunFlag {
-		if err := saveUndoSnapshot(g, cfg, branches, nil, "cascade", "gh stack cascade"); err != nil {
-			fmt.Printf("Warning: could not save undo state: %v\n", err)
+		var saveErr error
+		stashRef, saveErr = saveUndoSnapshot(g, cfg, branches, nil, "cascade", "gh stack cascade")
+		if saveErr != nil {
+			fmt.Printf("Warning: could not save undo state: %v\n", saveErr)
 		}
 	}
 
-	return doCascade(g, cfg, branches, cascadeDryRunFlag)
+	err = doCascade(g, cfg, branches, cascadeDryRunFlag)
+
+	// Restore auto-stashed changes after successful operation
+	if err == nil && stashRef != "" {
+		fmt.Println("Restoring auto-stashed changes...")
+		if popErr := g.StashPop(stashRef); popErr != nil {
+			fmt.Printf("Warning: could not restore stashed changes (commit %s): %v\n", stashRef[:7], popErr)
+		}
+	}
+
+	return err
 }
 
 func doCascade(g *git.Git, cfg *config.Config, branches []*tree.Node, dryRun bool) error {
@@ -201,27 +214,30 @@ func doCascadeWithState(g *git.Git, cfg *config.Config, branches []*tree.Node, d
 // It auto-stashes any uncommitted changes if the working tree is dirty.
 // branches: branches that will be modified (rebased)
 // deletedBranches: branches that will be deleted (for sync)
-func saveUndoSnapshot(g *git.Git, cfg *config.Config, branches []*tree.Node, deletedBranches []*tree.Node, operation, command string) error {
+// Returns the stash ref (commit hash) if changes were stashed, empty string otherwise.
+func saveUndoSnapshot(g *git.Git, cfg *config.Config, branches []*tree.Node, deletedBranches []*tree.Node, operation, command string) (string, error) {
 	gitDir := g.GetGitDir()
 
 	// Get current branch for original head
 	currentBranch, err := g.CurrentBranch()
 	if err != nil {
-		return err
+		return "", err
 	}
 
 	// Create snapshot
 	snapshot := undo.NewSnapshot(operation, command, currentBranch)
 
 	// Auto-stash if dirty
+	var stashRef string
 	dirty, err := g.IsDirty()
 	if err != nil {
-		return err
+		return "", err
 	}
 	if dirty {
-		stashRef, stashErr := g.Stash("gh-stack auto-stash before " + operation)
+		var stashErr error
+		stashRef, stashErr = g.Stash("gh-stack auto-stash before " + operation)
 		if stashErr != nil {
-			return fmt.Errorf("failed to stash changes: %w", stashErr)
+			return "", fmt.Errorf("failed to stash changes: %w", stashErr)
 		}
 		if stashRef != "" {
 			snapshot.StashRef = stashRef
@@ -231,52 +247,58 @@ func saveUndoSnapshot(g *git.Git, cfg *config.Config, branches []*tree.Node, del
 
 	// Capture state of branches that will be modified
 	for _, node := range branches {
-		bs, err := captureBranchState(g, cfg, node.Name)
-		if err != nil {
+		bs, captureErr := captureBranchState(g, cfg, node.Name)
+		if captureErr != nil {
 			// Non-fatal: log warning and continue
-			fmt.Printf("Warning: could not capture state for %s: %v\n", node.Name, err)
+			fmt.Printf("Warning: could not capture state for %s: %v\n", node.Name, captureErr)
 			continue
 		}
 		snapshot.Branches[node.Name] = bs
 	}
 
 	// Capture state of branches that will be deleted
 	for _, node := range deletedBranches {
-		bs, err := captureBranchState(g, cfg, node.Name)
-		if err != nil {
-			fmt.Printf("Warning: could not capture state for deleted branch %s: %v\n", node.Name, err)
+		bs, captureErr := captureBranchState(g, cfg, node.Name)
+		if captureErr != nil {
+			fmt.Printf("Warning: could not capture state for deleted branch %s: %v\n", node.Name, captureErr)
 			continue
 		}
 		snapshot.DeletedBranches[node.Name] = bs
 	}
 
 	// Save snapshot
-	return undo.Save(gitDir, snapshot)
+	if saveErr := undo.Save(gitDir, snapshot); saveErr != nil {
+		return "", saveErr
+	}
+	return stashRef, nil
 }
 
 // saveUndoSnapshotByName is like saveUndoSnapshot but takes branch names instead of tree nodes.
 // Useful for sync where we don't always have tree nodes.
-func saveUndoSnapshotByName(g *git.Git, cfg *config.Config, branchNames []string, deletedBranchNames []string, operation, command string) error {
+// Returns the stash ref (commit hash) if changes were stashed, empty string otherwise.
+func saveUndoSnapshotByName(g *git.Git, cfg *config.Config, branchNames []string, deletedBranchNames []string, operation, command string) (string, error) {
 	gitDir := g.GetGitDir()
 
 	// Get current branch for original head
 	currentBranch, err := g.CurrentBranch()
 	if err != nil {
-		return err
+		return "", err
 	}
 
 	// Create snapshot
 	snapshot := undo.NewSnapshot(operation, command, currentBranch)
 
 	// Auto-stash if dirty
+	var stashRef string
 	dirty, err := g.IsDirty()
 	if err != nil {
-		return err
+		return "", err
 	}
 	if dirty {
-		stashRef, stashErr := g.Stash("gh-stack auto-stash before " + operation)
+		var stashErr error
+		stashRef, stashErr = g.Stash("gh-stack auto-stash before " + operation)
 		if stashErr != nil {
-			return fmt.Errorf("failed to stash changes: %w", stashErr)
+			return "", fmt.Errorf("failed to stash changes: %w", stashErr)
 		}
 		if stashRef != "" {
 			snapshot.StashRef = stashRef
@@ -286,26 +308,29 @@ func saveUndoSnapshotByName(g *git.Git, cfg *config.Config, branchNames []string
 
 	// Capture state of branches that will be modified
 	for _, name := range branchNames {
-		bs, err := captureBranchState(g, cfg, name)
-		if err != nil {
-			fmt.Printf("Warning: could not capture state for %s: %v\n", name, err)
+		bs, captureErr := captureBranchState(g, cfg, name)
+		if captureErr != nil {
+			fmt.Printf("Warning: could not capture state for %s: %v\n", name, captureErr)
 			continue
 		}
 		snapshot.Branches[name] = bs
 	}
 
 	// Capture state of branches that will be deleted
 	for _, name := range deletedBranchNames {
-		bs, err := captureBranchState(g, cfg, name)
-		if err != nil {
-			fmt.Printf("Warning: could not capture state for deleted branch %s: %v\n", name, err)
+		bs, captureErr := captureBranchState(g, cfg, name)
+		if captureErr != nil {
+			fmt.Printf("Warning: could not capture state for deleted branch %s: %v\n", name, captureErr)
 			continue
 		}
 		snapshot.DeletedBranches[name] = bs
 	}
 
 	// Save snapshot
-	return undo.Save(gitDir, snapshot)
+	if saveErr := undo.Save(gitDir, snapshot); saveErr != nil {
+		return "", saveErr
+	}
+	return stashRef, nil
 }
 
 // captureBranchState captures the current state of a single branch.

cmd/submit.go

@@ -113,20 +113,33 @@ func runSubmit(cmd *cobra.Command, args []string) error {
 	}
 
 	// Save undo snapshot (unless dry-run)
+	var stashRef string
 	if !submitDryRunFlag {
-		if err := saveUndoSnapshot(g, cfg, branches, nil, "submit", "gh stack submit"); err != nil {
-			fmt.Printf("Warning: could not save undo state: %v\n", err)
+		var saveErr error
+		stashRef, saveErr = saveUndoSnapshot(g, cfg, branches, nil, "submit", "gh stack submit")
+		if saveErr != nil {
+			fmt.Printf("Warning: could not save undo state: %v\n", saveErr)
 		}
 	}
 
 	// Phase 1: Cascade
 	fmt.Println("=== Phase 1: Cascade ===")
-	if err := doCascadeWithState(g, cfg, branches, submitDryRunFlag, state.OperationSubmit, submitUpdateOnlyFlag, submitWebFlag, branchNames); err != nil {
-		return err // Conflict or error - state saved, user can continue
+	if cascadeErr := doCascadeWithState(g, cfg, branches, submitDryRunFlag, state.OperationSubmit, submitUpdateOnlyFlag, submitWebFlag, branchNames); cascadeErr != nil {
+		return cascadeErr // Conflict or error - state saved, user can continue
 	}
 
 	// Phases 2 & 3
-	return doSubmitPushAndPR(g, cfg, root, branches, submitDryRunFlag, submitUpdateOnlyFlag, submitWebFlag)
+	err = doSubmitPushAndPR(g, cfg, root, branches, submitDryRunFlag, submitUpdateOnlyFlag, submitWebFlag)
+
+	// Restore auto-stashed changes after successful operation
+	if err == nil && stashRef != "" {
+		fmt.Println("Restoring auto-stashed changes...")
+		if popErr := g.StashPop(stashRef); popErr != nil {
+			fmt.Printf("Warning: could not restore stashed changes (commit %s): %v\n", stashRef[:7], popErr)
+		}
+	}
+
+	return err
 }
 
 // doSubmitPushAndPR handles push and PR creation/update phases.

cmd/sync.go

@@ -100,10 +100,13 @@ func runSync(cmd *cobra.Command, args []string) error {
 
 	// Save undo snapshot of all tracked branches (unless dry-run)
 	// This captures state before any modifications (fetch, delete, rebase)
+	var stashRef string
 	if !syncDryRunFlag {
 		allBranches, listErr := cfg.ListTrackedBranches()
 		if listErr == nil && len(allBranches) > 0 {
-			if saveErr := saveUndoSnapshotByName(g, cfg, allBranches, nil, "sync", "gh stack sync"); saveErr != nil {
+			var saveErr error
+			stashRef, saveErr = saveUndoSnapshotByName(g, cfg, allBranches, nil, "sync", "gh stack sync")
+			if saveErr != nil {
 				fmt.Printf("Warning: could not save undo state: %v\n", saveErr)
 			}
 		}
@@ -350,6 +353,14 @@ func runSync(cmd *cobra.Command, args []string) error {
 		}
 	}
 
+	// Restore auto-stashed changes after successful operation
+	if stashRef != "" {
+		fmt.Println("Restoring auto-stashed changes...")
+		if popErr := g.StashPop(stashRef); popErr != nil {
+			fmt.Printf("Warning: could not restore stashed changes (commit %s): %v\n", stashRef[:7], popErr)
+		}
+	}
+
 	fmt.Println("\nSync complete!")
 	return nil
 }

cmd/undo.go

@@ -219,11 +219,11 @@ func runUndo(cmd *cobra.Command, args []string) error {
 		}
 	}
 
-	// Restore stash
+	// Restore stash using the specific stash commit hash
 	if snapshot.StashRef != "" {
 		fmt.Println("  Restoring stashed changes...")
-		if err := g.StashPop(); err != nil {
-			fmt.Printf("  Warning: could not cleanly restore stashed changes. Your changes are still in %s.\n", snapshot.StashRef)
+		if err := g.StashPop(snapshot.StashRef); err != nil {
+			fmt.Printf("  Warning: could not cleanly restore stashed changes. Your changes are still in stash (commit %s).\n", snapshot.StashRef[:7])
 		}
 	}
 

e2e/submit_test.go

@@ -146,7 +146,7 @@ func TestSubmitUpdateOnlyDryRun(t *testing.T) {
 	}
 }
 
-func TestSubmitWithDirtyWorkTreeAutoStashes(t *testing.T) {
+func TestSubmitDryRunAllowsDirtyWorkTree(t *testing.T) {
 	env := NewTestEnvWithRemote(t)
 	env.MustRun("init")
 
@@ -158,13 +158,12 @@ func TestSubmitWithDirtyWorkTreeAutoStashes(t *testing.T) {
 	env.WriteFile("dirty.txt", "uncommitted")
 	env.Git("add", "dirty.txt")
 
-	// Submit should auto-stash and succeed (at least the cascade phase)
-	// Note: without GH_TOKEN, PR creation will fail, but cascade phase should work
+	// Submit --dry-run should succeed even with dirty working tree
+	// (dry-run skips the undo snapshot which includes auto-stashing)
 	result := env.Run("submit", "--dry-run")
 
-	// In dry-run mode, auto-stash still happens but no actual rebases
 	if result.Failed() {
-		t.Errorf("submit dry-run should succeed with auto-stash, got: %s", result.Stderr)
+		t.Errorf("submit --dry-run should succeed with dirty tree, got: %s", result.Stderr)
 	}
 }
 

e2e/undo_test.go

@@ -183,7 +183,7 @@ func TestUndoArchivesSnapshot(t *testing.T) {
 	}
 }
 
-func TestUndoWithAutoStash(t *testing.T) {
+func TestCascadeWithAutoStashRestoresAfterSuccess(t *testing.T) {
 	env := NewTestEnv(t)
 	env.MustRun("init")
 
@@ -198,22 +198,19 @@ func TestUndoWithAutoStash(t *testing.T) {
 	// Create uncommitted changes
 	env.WriteFile("uncommitted.txt", "uncommitted content\n")
 
-	// Cascade should auto-stash
+	// Cascade should auto-stash and then restore after success
 	result := env.MustRun("cascade")
 	if !result.ContainsStdout("Auto-stashed") {
 		t.Error("expected auto-stash message")
 	}
+	if !result.ContainsStdout("Restoring auto-stashed") {
+		t.Error("expected restore message after successful cascade")
+	}
 
-	// Working tree should be clean (stashed)
-	env.AssertClean()
-
-	// Undo should restore stash
-	env.MustRun("undo", "--force")
-
-	// Check that uncommitted file is back
+	// Uncommitted file should be restored after successful cascade
 	content, err := os.ReadFile(filepath.Join(env.WorkDir, "uncommitted.txt"))
 	if err != nil {
-		t.Errorf("uncommitted file should be restored: %v", err)
+		t.Errorf("uncommitted file should be present after cascade: %v", err)
 	} else if string(content) != "uncommitted content\n" {
 		t.Errorf("uncommitted file has wrong content: %q", content)
 	}

internal/git/git.go

@@ -293,9 +293,11 @@ func (g *Git) CreateBranchAt(name, sha string) error {
 	return g.runSilent("branch", name, sha)
 }
 
-// Stash creates a stash with the given message and returns the stash reference.
+// Stash creates a stash with the given message and returns the stash commit hash.
 // Returns an empty string if there was nothing to stash.
 // Includes untracked files (-u) to capture all working tree changes.
+// The returned hash is stable and can be used to restore this specific stash
+// even if other stashes are created later.
 func (g *Git) Stash(message string) (string, error) {
 	// Check if there's anything to stash first
 	dirty, err := g.IsDirty()
@@ -307,18 +309,33 @@ func (g *Git) Stash(message string) (string, error) {
 	}
 
 	// Create the stash with -u to include untracked files
-	if err := g.runSilent("stash", "push", "-u", "-m", message); err != nil {
-		return "", err
+	if stashErr := g.runSilent("stash", "push", "-u", "-m", message); stashErr != nil {
+		return "", stashErr
 	}
 
-	// Get the stash reference (stash@{0} after a successful push)
-	return "stash@{0}", nil
+	// Resolve the created stash to a stable identifier (its commit hash)
+	// This is more reliable than stash@{0} which can shift if user creates more stashes
+	out, parseErr := g.run("rev-parse", "stash@{0}")
+	if parseErr != nil {
+		return "", parseErr
+	}
+	return out, nil
 }
 
-// StashPop pops the most recent stash entry.
-// Returns an error if there are conflicts or no stash entries.
-func (g *Git) StashPop() error {
-	return g.runInteractive("stash", "pop")
+// StashPop restores a specific stash entry when a reference is provided.
+// If ref is empty, it pops the most recent stash entry (matching `git stash pop`).
+// When a commit hash is provided (from Stash()), uses apply since pop only accepts stash refs.
+// Note: apply doesn't remove the stash entry, but this is acceptable for undo purposes.
+// Returns an error if there are conflicts or no matching stash entry.
+func (g *Git) StashPop(ref string) error {
+	if ref == "" {
+		return g.runInteractive("stash", "pop")
+	}
+
+	// git stash pop only accepts stash references (stash@{n}), not commit hashes.
+	// Use apply instead which accepts commit hashes.
+	// Note: This leaves the stash entry in place, which is fine for our use case.
+	return g.runInteractive("stash", "apply", ref)
 }
 
 // StashList returns true if there are any stash entries.