Preserve rewrite checkout during stack selection

Author: skarim

cmd/rebase_test.go

@@ -2491,6 +2491,97 @@ func setupWorktreeRebaseRepo(t *testing.T, conflict bool) worktreeRebaseRepo {
 	return worktreeRebaseRepo{repo, parentDir, childDir}
 }
 
+func setupSharedTrunkRebaseRepo(t *testing.T, conflict bool) worktreeRebaseRepo {
+	t.Helper()
+	repo := setupWorktreeRebaseRepo(t, conflict)
+	issue250Git(t, repo.childDir, "checkout", "--detach")
+	issue250Git(t, repo.dir, "branch", "independent", "main")
+	sf, err := stack.Load(repo.gitDir)
+	require.NoError(t, err)
+	sf.AddStack(stack.Stack{Trunk: stack.BranchRef{Branch: "main"}, Branches: []stack.BranchRef{{Branch: "independent"}}})
+	require.NoError(t, stack.Save(repo.gitDir, sf))
+	return repo
+}
+
+func TestRebase_SharedTrunkSelectionPreservesOriginAndRange(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		opts rebaseOptions
+		want []string
+	}{
+		{name: "whole stack", want: []string{"b1", "b2", "b3"}},
+		{name: "upstack from trunk", opts: rebaseOptions{upstack: true}, want: []string{"b1", "b2", "b3"}},
+		{name: "downstack from trunk", opts: rebaseOptions{downstack: true}, want: []string{"b1"}},
+		{name: "explicit upstack", opts: rebaseOptions{branch: "b2", upstack: true}, want: []string{"b2", "b3"}},
+		{name: "explicit downstack", opts: rebaseOptions{branch: "b2", downstack: true}, want: []string{"b1", "b2"}},
+		{name: "without trunk", opts: rebaseOptions{noTrunk: true}, want: []string{"b2", "b3"}},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			dir := t.TempDir()
+			writeStackFileMulti(t, dir,
+				stack.Stack{Trunk: stack.BranchRef{Branch: "main"}, Branches: []stack.BranchRef{{Branch: "b1"}, {Branch: "b2"}, {Branch: "b3"}}},
+				stack.Stack{Trunk: stack.BranchRef{Branch: "main"}, Branches: []stack.BranchRef{{Branch: "independent"}}},
+			)
+			current := "main"
+			var rebased []string
+			mock := newRebaseMock(dir, current)
+			mock.CurrentBranchFn = func() (string, error) { return current, nil }
+			mock.BranchExistsFn = func(string) (bool, error) { return true, nil }
+			mock.CheckoutBranchFn = func(branch string) error { current = branch; return nil }
+			mock.RebaseFn = func(string, git.RebaseOpts) error { rebased = append(rebased, current); return nil }
+			mock.RebaseOntoFn = func(_, _, branch string, _ git.RebaseOpts) error {
+				rebased = append(rebased, branch)
+				return nil
+			}
+			mock.IsRerereEnabledFn = func() (bool, error) { return true, nil }
+			restore := git.SetOps(mock)
+			defer restore()
+			cfg := issue250TestConfig(t)
+			cfg.ForceInteractive = true
+			cfg.SelectFn = func(_, _ string, _ []string) (int, error) { return 0, nil }
+			tc.opts.remote = "origin"
+
+			require.NoError(t, runRebase(cfg, &tc.opts))
+
+			assert.Equal(t, tc.want, rebased)
+			assert.Equal(t, "main", current)
+			assert.NoFileExists(t, filepath.Join(dir, rebaseStateFile))
+		})
+	}
+}
+
+func TestRebase_SharedTrunkSelectionRecoveryPreservesOrigin(t *testing.T) {
+	for _, abort := range []bool{false, true} {
+		t.Run(fmt.Sprintf("abort=%t", abort), func(t *testing.T) {
+			repo := setupSharedTrunkRebaseRepo(t, true)
+			observerHead := issue250Git(t, repo.childDir, "rev-parse", "HEAD")
+			withIssue250Repo(t, repo.dir)
+			cfg := issue250TestConfig(t)
+			cfg.ForceInteractive = true
+			selections := 0
+			cfg.SelectFn = func(_, _ string, _ []string) (int, error) { selections++; return 0, nil }
+			cfg.ConfirmFn = func(string, bool) (bool, error) { return false, nil }
+			require.ErrorIs(t, runRebase(cfg, &rebaseOptions{remote: "origin"}), ErrConflict)
+			state, err := loadRebaseState(repo.gitDir)
+			require.NoError(t, err)
+			assert.Equal(t, "main", state.OriginalBranch)
+			if !abort {
+				issue250WriteFile(t, repo.dir, "base.txt", "resolved\n")
+				issue250Git(t, repo.dir, "add", "base.txt")
+			}
+			withIssue250Repo(t, repo.parentDir)
+
+			require.NoError(t, runRebase(cfg, &rebaseOptions{abort: abort, cont: !abort}))
+
+			assert.Equal(t, "main", issue250Git(t, repo.dir, "branch", "--show-current"))
+			assert.Equal(t, "parent", issue250Git(t, repo.parentDir, "branch", "--show-current"))
+			assert.Equal(t, observerHead, issue250Git(t, repo.childDir, "rev-parse", "HEAD"))
+			assert.Equal(t, 1, selections)
+			assert.NoFileExists(t, filepath.Join(repo.gitDir, rebaseStateFile))
+		})
+	}
+}
+
 func TestRebase_WorktreesPreserveCheckouts(t *testing.T) {
 	repo := setupWorktreeRebaseRepo(t, false)
 	issue250WriteFile(t, repo.dir, "unrelated.txt", "leave main alone\n")

cmd/sync_test.go

@@ -2634,6 +2634,23 @@ func TestSync_WorktreesPublishesFromLinkedWorktree(t *testing.T) {
 	require.NoError(t, issue250GitMayFail(t, repo.dir, "merge-base", "--is-ancestor", "parent", "child"))
 }
 
+func TestSync_SharedTrunkSelectionRestoresOrigin(t *testing.T) {
+	repo := setupSharedTrunkRebaseRepo(t, false)
+	observerHead := issue250Git(t, repo.childDir, "rev-parse", "HEAD")
+	withIssue250Repo(t, repo.dir)
+	cfg := issue250TestConfig(t)
+	cfg.ForceInteractive = true
+	cfg.SelectFn = func(_, _ string, _ []string) (int, error) { return 0, nil }
+	cfg.ConfirmFn = func(string, bool) (bool, error) { return false, nil }
+
+	require.NoError(t, runSync(cfg, &syncOptions{remote: "origin"}))
+
+	assert.Equal(t, "main", issue250Git(t, repo.dir, "branch", "--show-current"))
+	assert.Equal(t, "parent", issue250Git(t, repo.parentDir, "branch", "--show-current"))
+	assert.Equal(t, observerHead, issue250Git(t, repo.childDir, "rev-parse", "HEAD"))
+	assert.Equal(t, issue250Git(t, repo.dir, "rev-parse", "child"), issue250Git(t, repo.dir, "rev-parse", "origin/child"))
+}
+
 func TestSync_WorktreesConflictRestoresWithoutPush(t *testing.T) {
 	repo := setupWorktreeRebaseRepo(t, true)
 	beforeParent := issue250Git(t, repo.dir, "rev-parse", "parent")

cmd/utils.go

@@ -407,9 +407,9 @@ func handleSaveError(cfg *config.Config, err error) error {
 // resolveStack finds the stack for the given branch, handling ambiguity when
 // a branch (typically a trunk) belongs to multiple stacks. If exactly one
 // stack matches, it is returned directly. If multiple stacks match, the user
-// is prompted to select one and the working tree is switched to the top branch
-// of the selected stack. Returns nil with no error if no stack contains the
-// branch.
+// is prompted to select one. Read-only and rewrite callers preserve their
+// checkout; other mutating callers switch to the selected stack's top branch.
+// Returns nil with no error if no stack contains the branch.
 func resolveStack(sf *stack.StackFile, branch string, cfg *config.Config) (*stack.Stack, error) {
 	stacks := sf.FindAllStacksForBranch(branch)
 
@@ -456,8 +456,9 @@ func resolveStack(sf *stack.StackFile, branch string, cfg *config.Config) (*stac
 	if len(s.Branches) == 0 {
 		return nil, fmt.Errorf("selected stack %q has no branches", s.DisplayChain())
 	}
-	// Read-only selection must remain usable while another operation is paused.
-	if cfg.StackMutation == nil {
+	// Selection must not change a rewrite's origin or bypass its preflight.
+	// Read-only selection also remains available while an operation is paused.
+	if cfg.StackMutation == nil || cfg.StackMutation.NoCheckoutOnSelect {
 		return s, nil
 	}
 

cmd/utils_test.go

@@ -78,6 +78,86 @@ func TestResolveStack_ReadOnlySelectionDoesNotCheckout(t *testing.T) {
 	commandOutput(t, cfg, outR, errR)
 }
 
+func TestResolveStack_RewriteSelectionPreservesCheckout(t *testing.T) {
+	for _, kind := range []string{"rebase", "rebase-continue", "rebase-abort", "sync", "modify", "modify-continue", "modify-abort", "add", "checkout", "submit"} {
+		t.Run(kind, func(t *testing.T) {
+			dir := t.TempDir()
+			writeStackFileMulti(t, dir,
+				stack.Stack{Trunk: stack.BranchRef{Branch: "main"}, Branches: []stack.BranchRef{{Branch: "b1"}, {Branch: "b2"}}},
+				stack.Stack{Trunk: stack.BranchRef{Branch: "main"}, Branches: []stack.BranchRef{{Branch: "independent"}}},
+			)
+			current := "main"
+			var checkouts []string
+			mock := newRebaseMock(dir, current)
+			mock.CurrentBranchFn = func() (string, error) { return current, nil }
+			mock.CheckoutBranchFn = func(branch string) error {
+				checkouts = append(checkouts, branch)
+				current = branch
+				return nil
+			}
+			restore := git.SetOps(mock)
+			defer restore()
+			cfg := issue250TestConfig(t)
+			cfg.ForceInteractive = true
+			cfg.SelectFn = func(_, _ string, choices []string) (int, error) {
+				require.Len(t, choices, 2)
+				return 0, nil
+			}
+			release, err := beginStackMutation(cfg, kind)
+			require.NoError(t, err)
+			defer release()
+
+			result, err := loadStackOptional(cfg, "")
+
+			require.NoError(t, err)
+			assert.Equal(t, []string{"b1", "b2"}, result.Stack.BranchNames())
+			if kind == "add" || kind == "checkout" || kind == "submit" {
+				assert.Equal(t, []string{"b2"}, checkouts, "intentional selection checkout must remain available")
+				assert.Equal(t, "b2", result.CurrentBranch)
+			} else {
+				assert.Empty(t, checkouts, "rewrite selection must not move the origin before preflight")
+				assert.Equal(t, "main", result.CurrentBranch)
+			}
+			assert.Equal(t, result.CurrentBranch, current)
+		})
+	}
+}
+
+func TestStackSelection_RewritePreflightPreservesCheckout(t *testing.T) {
+	for _, command := range []struct {
+		name string
+		run  func(*config.Config) error
+	}{
+		{"rebase", func(cfg *config.Config) error { return runRebase(cfg, &rebaseOptions{remote: "origin"}) }},
+		{"sync", func(cfg *config.Config) error { return runSync(cfg, &syncOptions{remote: "origin"}) }},
+		{"modify", runModify},
+	} {
+		t.Run(command.name, func(t *testing.T) {
+			repo := setupSharedTrunkRebaseRepo(t, false)
+			issue250WriteFile(t, repo.parentDir, "unfinished.txt", "preserve this work\n")
+			beforeRefs := issue250Git(t, repo.dir, "show-ref")
+			beforeCatalog, err := os.ReadFile(filepath.Join(repo.gitDir, "gh-stack"))
+			require.NoError(t, err)
+			withIssue250Repo(t, repo.dir)
+			cfg := issue250TestConfig(t)
+			cfg.ForceInteractive = true
+			cfg.SelectFn = func(_, _ string, _ []string) (int, error) { return 0, nil }
+			cfg.ConfirmFn = func(string, bool) (bool, error) { return false, nil }
+
+			require.Error(t, command.run(cfg))
+
+			assert.Equal(t, "main", issue250Git(t, repo.dir, "branch", "--show-current"))
+			assert.Equal(t, "parent", issue250Git(t, repo.parentDir, "branch", "--show-current"))
+			assert.Equal(t, beforeRefs, issue250Git(t, repo.dir, "show-ref"))
+			afterCatalog, err := os.ReadFile(filepath.Join(repo.gitDir, "gh-stack"))
+			require.NoError(t, err)
+			assert.Equal(t, beforeCatalog, afterCatalog)
+			assert.FileExists(t, filepath.Join(repo.parentDir, "unfinished.txt"))
+			assert.NoFileExists(t, filepath.Join(repo.gitDir, rebaseStateFile))
+		})
+	}
+}
+
 func TestStackMutation_NestedAndReadOnly(t *testing.T) {
 	common := t.TempDir()
 	restore := git.SetOps(&git.MockOps{GitDirFn: func() (string, error) { return common, nil }})

cmd/worktree_utils.go

@@ -78,6 +78,10 @@ func beginStackMutation(cfg *config.Config, kind string) (func(), error) {
 		return nil, err
 	}
 	cfg.StackMutation = &config.StackMutationContext{CommonDir: commonDir, StateDir: stateDir}
+	switch kind {
+	case "rebase", "rebase-continue", "rebase-abort", "sync", "modify", "modify-continue", "modify-abort":
+		cfg.StackMutation.NoCheckoutOnSelect = true
+	}
 	released := false
 	return func() {
 		if !released {

docs/src/content/docs/guides/workflows.md

@@ -67,6 +67,8 @@ gh-stack does not install shell functions or change your shell's directory. Do n
 
 `rebase` and `sync` automatically operate in each affected branch's clean owning worktree. Unoccupied branches are processed in the initiating worktree, whose original checkout is restored afterward. Dirty, busy, missing, or changed affected owners block mutation; unrelated worktrees are left alone. A clean trunk owner can be fast-forwarded, while an unsafe local trunk retains the fetched-remote fallback. gh-stack never auto-stashes, transfers ownership, or creates/removes worktrees.
 
+When a trunk belongs to multiple stacks, selecting one for `rebase`, `sync`, or `modify` does not switch branches. Rebase ranges use the caller's original checkout unless an explicit branch is supplied.
+
 Resolve and stage conflicts in the worktree named by the diagnostic. You can run `gh stack rebase --continue` or `--abort` from any linked worktree: the shared journal routes recovery to the recorded owners. `sync` still restores its cascade on conflicts rather than pushing partial results; completed fetches and earlier fast-forwards are outside that rollback boundary. Recovery retains state and reports any partial failure rather than discarding later edits or claiming a full restoration. Pruning skips branches still occupied in other worktrees.
 
 Finish paused operations before changing gh-stack versions or preview stages. Origin-only journals are explicitly marked; a build that cannot interpret a journal's execution lifecycle must leave it intact. If recovery reports an incompatible lifecycle, use the matching build in the recorded origin to finish or abort it instead of editing or removing the journal.

internal/config/config.go

@@ -67,6 +67,9 @@ type Config struct {
 type StackMutationContext struct {
 	CommonDir string
 	StateDir  string
+
+	// NoCheckoutOnSelect preserves rewrite origins and range anchors.
+	NoCheckoutOnSelect bool
 }
 
 // New creates a new Config with terminal-aware output and color support.