checkout: resolve remote stack by branch name (#477)

* checkout by remote branch name

* address review comments

Author: skarim

README.md

@@ -158,7 +158,7 @@ A bare number is interpreted first as a stack or PR number (repo-scoped identifi
 
 When a remote stack is referenced, the command fetches the stack on GitHub, pulls the branches, and sets up the stack locally. If the stack already exists locally and matches, it switches to the branch. If the local and remote stacks have different compositions, you'll be prompted to resolve the conflict.
 
-When a branch name is provided, the command resolves it against locally tracked stacks only.
+When a branch name is provided, the command checks locally tracked stacks first. If the branch is not tracked locally, it looks for the branch on remote stacks and pulls down the matching stack. If more than one stack matches, use a stack or PR number to choose one explicitly.
 
 When run without arguments in an interactive terminal, first checks whether the current branch belongs to a stack on remote that is not tracked locally, and offers to check it out. If there is no unique match or you decline, it opens a searchable picker listing every stack available to you — both the stacks tracked locally and the stacks that exist only on GitHub. Each row shows the stack number, its bottom and top branch, base branch, a status bar summarizing how many of its pull requests are merged, open, closed, or not yet pushed, and whether the stack is available locally or only on the remote. Filter with the All / Local / Remote tabs or type `/` to search; fully merged stacks are omitted. Selecting a remote-only stack clones it locally before switching to it.
 
@@ -174,7 +174,7 @@ gh stack checkout 42
 # Check out a stack by PR URL
 gh stack checkout https://github.com/owner/repo/pull/42
 
-# Check out a stack by branch name (local only)
+# Check out a stack by branch name
 gh stack checkout feature-auth
 
 # Interactive — pick from all available stacks (local and remote)

cmd/checkout.go

@@ -41,8 +41,9 @@ GitHub API to discover the stack, fetches the branches, and sets up
 the stack locally. If the stack already exists locally and matches,
 it simply switches to the branch.
 
-When a branch name is provided, the command resolves it against
-locally tracked stacks only.
+When a branch name is provided, the command first checks locally tracked
+stacks. If the branch is not tracked locally, it looks for the branch on
+remote stacks and pulls down the matching stack.
 
 When run without arguments, first checks whether the current branch belongs
 to a stack on remote that is not tracked locally, and offers to check
@@ -79,7 +80,7 @@ omitted.`,
 // runCheckout resolves a stack and checks out the target branch.
 // For numeric targets, it tries local lookup first, then falls back to
 // the GitHub API to discover remote stacks, then tries as a branch name.
-// Non-numeric targets use local resolution only.
+// Branch names resolve locally first and then against stacks on GitHub.
 func runCheckout(cfg *config.Config, opts *checkoutOptions) error {
 	gitDir, err := git.GitDir()
 	if err != nil {
@@ -125,14 +126,23 @@ func runCheckout(cfg *config.Config, opts *checkoutOptions) error {
 			return err
 		}
 	} else {
-		// Non-numeric target — resolve against local stacks only
+		// Non-numeric target — resolve locally before checking GitHub.
 		var br *stack.BranchRef
 		s, br, err = resolvePR(cfg, sf, opts.target)
-		if err != nil {
-			cfg.Errorf("%s", err)
-			return ErrNotInStack
+		if err == nil {
+			targetBranch = br.Branch
+		} else {
+			s, targetBranch, err = checkoutRemoteStackByBranch(cfg, sf, gitDir, opts.target)
+			if errors.Is(err, errRemoteBranchNotFound) {
+				cfg.Errorf("no local or remote stack found for %q", opts.target)
+				cfg.Printf("Try a stack or PR number with `%s`",
+					cfg.ColorCyan("gh stack checkout <number>"))
+				return ErrNotInStack
+			}
+			if err != nil {
+				return err
+			}
 		}
-		targetBranch = br.Branch
 	}
 
 	currentBranch, _ := git.CurrentBranch()
@@ -188,7 +198,7 @@ func resolveNumericTarget(cfg *config.Config, sf *stack.StackFile, gitDir string
 	// attempt — the user might have a numeric branch name.
 	remoteErr := err
 
-	// 4. Fall back to branch name lookup (handles numeric branch names).
+	// 4. Fall back to local branch name lookup (handles numeric branch names).
 	stacks := sf.FindAllStacksForBranch(raw)
 	if len(stacks) > 0 {
 		s := stacks[0]
@@ -202,11 +212,77 @@ func resolveNumericTarget(cfg *config.Config, sf *stack.StackFile, gitDir string
 		}
 	}
 
+	// Only a definitive "not in a stack" result can be reinterpreted as a
+	// remote branch name. Preserve API, conflict, and other actionable errors.
+	if !errors.Is(remoteErr, ErrNotInStack) {
+		return nil, "", remoteErr
+	}
+
+	// Finally, try the numeric input as a branch in a stack on GitHub.
+	s, targetBranch, err = checkoutRemoteStackByBranch(cfg, sf, gitDir, raw)
+	if err == nil {
+		return s, targetBranch, nil
+	}
+	if !errors.Is(err, errRemoteBranchNotFound) {
+		return nil, "", err
+	}
+
 	// Nothing worked — return the remote error which has the most
 	// informative message for a numeric input
+	cfg.Errorf("PR #%d is not part of a stack on GitHub", number)
 	return nil, "", remoteErr
 }
 
+var errRemoteBranchNotFound = errors.New("remote branch not found in a stack")
+
+// checkoutRemoteStackByBranch finds the unique active stack on GitHub that
+// contains branch, then imports it through the existing PR checkout path.
+func checkoutRemoteStackByBranch(cfg *config.Config, sf *stack.StackFile, gitDir, branch string) (*stack.Stack, string, error) {
+	client, err := cfg.GitHubClient()
+	if err != nil {
+		cfg.Errorf("failed to create GitHub client: %s", err)
+		return nil, "", ErrAPIFailure
+	}
+
+	remoteStacks, err := client.ListStacks()
+	if err != nil {
+		var httpErr *api.HTTPError
+		if errors.As(err, &httpErr) && httpErr.StatusCode == 404 {
+			warnStacksUnavailable(cfg)
+			return nil, "", ErrStacksUnavailable
+		}
+		cfg.Errorf("failed to list stacks: %v", err)
+		return nil, "", ErrAPIFailure
+	}
+
+	matches := matchingRemoteStacksForBranch(remoteStacks, branch)
+	if len(matches) == 0 {
+		return nil, "", errRemoteBranchNotFound
+	}
+	if len(matches) > 1 {
+		stackNumbers := make([]string, len(matches))
+		for i, match := range matches {
+			stackNumbers[i] = strconv.Itoa(match.Number)
+		}
+		cfg.Errorf("branch %q belongs to multiple stacks on GitHub (%s)",
+			branch, strings.Join(stackNumbers, ", "))
+		cfg.Printf("Use `%s` with a stack or PR number to choose one",
+			cfg.ColorCyan("gh stack checkout <number>"))
+		return nil, "", ErrDisambiguate
+	}
+
+	for _, pr := range matches[0].PRDetails {
+		if pr.Head.Ref == branch && pr.Number > 0 {
+			s, targetBranch, err := checkoutRemoteStack(cfg, sf, gitDir, pr.Number)
+			if errors.Is(err, ErrNotInStack) {
+				cfg.Errorf("PR #%d is not part of a stack on GitHub", pr.Number)
+			}
+			return s, targetBranch, err
+		}
+	}
+	return nil, "", errRemoteBranchNotFound
+}
+
 // checkoutRemoteStack discovers a stack from GitHub for the given PR number,
 // reconciles it with any local state, and returns the resolved stack and
 // target branch name. The stack file is saved before returning.
@@ -224,13 +300,12 @@ func checkoutRemoteStack(cfg *config.Config, sf *stack.StackFile, gitDir string,
 		var httpErr *api.HTTPError
 		if errors.As(err, &httpErr) && httpErr.StatusCode == 404 {
 			warnStacksUnavailable(cfg)
-			return nil, "", ErrAPIFailure
+			return nil, "", ErrStacksUnavailable
 		}
 		cfg.Errorf("failed to list stacks: %v", err)
 		return nil, "", ErrAPIFailure
 	}
 	if remoteStack == nil {
-		cfg.Errorf("PR #%d is not part of a stack on GitHub", prNumber)
 		return nil, "", ErrNotInStack
 	}