fix(pr): keep targeting the base repository after a fork PR checkout

Repository discovery preferred the current branch's upstream remote,
which for a fork pull request is the pulls/<owner>/<repo> remote that
pr checkout adds, so every later command targeted the fork. pr checkout
now records the pull request number and base remote in
branch.<name>.tea-pr and branch.<name>.tea-base-remote, discovery
honors the recorded base remote, and it skips pulls/* head remotes.

Signed-off-by: bircni <bircni@icloud.com>

internal/app/checkout_test.go

@@ -71,6 +71,7 @@ func TestCheckoutMatchesRemoteThroughEnvironmentSSHHost(t *testing.T) {
 		"status --porcelain=v1 --untracked-files=normal": "", "for-each-ref --format=%(refname) refs/heads/": "",
 		"fetch --no-tags -- origin refs/heads/feature": "", "rev-parse --verify FETCH_HEAD^{commit}": "abc123", "switch --create pr/7 abc123": "",
 		"config branch.pr/7.remote origin": "", "config branch.pr/7.merge refs/heads/feature": "",
+		"config branch.pr/7.tea-pr 7": "", "config branch.pr/7.tea-base-remote origin": "",
 	}}
 	require.NoError(t, Execute(context.Background(), deps, []string{"pr", "checkout", "7", "--repo", "alice/project", "--remote", "origin", "--no-input", "--json"}))
 	require.JSONEq(t, `{"branch":"pr/7","commit":"abc123","detached":false,"remote":"origin","tracking":"origin/feature"}`, out.String())
@@ -97,6 +98,7 @@ func TestSSHCheckoutUsesAccountKeyOnlyForFetch(t *testing.T) {
 		"status --porcelain=v1 --untracked-files=normal": "", "for-each-ref --format=%(refname) refs/heads/": "",
 		fetch: "", "rev-parse --verify FETCH_HEAD^{commit}": "abc123", "switch --create pr/7 abc123": "",
 		"config branch.pr/7.remote origin": "", "config branch.pr/7.merge refs/heads/feature": "",
+		"config branch.pr/7.tea-pr 7": "", "config branch.pr/7.tea-base-remote origin": "",
 	}}
 	deps.Git = runner
 	require.NoError(t, Execute(context.Background(), deps, []string{"pr", "checkout", "7", "--repo", "alice/project", "--remote", "origin", "--no-input", "--json"}))

internal/git/checkout.go

@@ -7,9 +7,16 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"strconv"
 	"strings"
 )
 
+// Branch configuration keys that record the pull request and base remote of a checked-out PR branch.
+const (
+	pullNumberKey     = "tea-pr"
+	pullBaseRemoteKey = "tea-base-remote"
+)
+
 // CheckoutOptions identifies the base remote and how to check out the PR commit.
 type CheckoutOptions struct {
 	// CredentialArgs contains per-operation Git configuration for the fetch only.
@@ -152,5 +159,12 @@ func CheckoutPull(ctx context.Context, runner Runner, options CheckoutOptions) (
 			return CheckoutResult{}, fmt.Errorf("checked out the branch, but could not configure tracking: %w", err)
 		}
 	}
+	if !options.Detached {
+		for _, setting := range [][2]string{{pullNumberKey, strconv.FormatInt(options.Number, 10)}, {pullBaseRemoteKey, options.Remote}} {
+			if _, err := run("config", "branch."+options.Branch+"."+setting[0], setting[1]); err != nil {
+				return CheckoutResult{}, fmt.Errorf("checked out the branch, but could not record its pull request: %w", err)
+			}
+		}
+	}
 	return CheckoutResult{Branch: options.Branch, Commit: commit, Detached: options.Detached, Remote: fetchRemote, Tracking: tracking}, nil
 }

internal/git/checkout_test.go

@@ -68,7 +68,7 @@ func TestPullCheckoutValidatesBeforeChangingGitState(t *testing.T) {
 	result, err := CheckoutPull(context.Background(), runner, options)
 	require.NoError(t, err)
 	require.Equal(t, CheckoutResult{Branch: "pr/7", Commit: "abc123", Remote: "origin"}, result)
-	require.Equal(t, []string{"rev-parse --is-inside-work-tree", "check-ref-format --branch pr/7", "status --porcelain=v1 --untracked-files=normal", "for-each-ref --format=%(refname) refs/heads/", "fetch --no-tags -- origin refs/pull/7/head", "rev-parse --verify FETCH_HEAD^{commit}", "switch --create pr/7 abc123"}, runner.calls)
+	require.Equal(t, []string{"rev-parse --is-inside-work-tree", "check-ref-format --branch pr/7", "status --porcelain=v1 --untracked-files=normal", "for-each-ref --format=%(refname) refs/heads/", "fetch --no-tags -- origin refs/pull/7/head", "rev-parse --verify FETCH_HEAD^{commit}", "switch --create pr/7 abc123", "config branch.pr/7.tea-pr 7", "config branch.pr/7.tea-base-remote origin"}, runner.calls)
 }
 
 func TestPullCheckoutRejectsChangedHead(t *testing.T) {
@@ -103,6 +103,7 @@ func TestDetachedCheckoutDoesNotCreateOrChangeBranches(t *testing.T) {
 	require.Equal(t, CheckoutResult{Detached: true, Commit: "abc123", Remote: "origin"}, result)
 	require.Equal(t, "switch --detach abc123", runner.calls[len(runner.calls)-1])
 	for _, call := range runner.calls {
+		require.False(t, strings.HasPrefix(call, "config "))
 		require.False(t, strings.HasPrefix(call, "for-each-ref"))
 		require.False(t, strings.HasPrefix(call, "check-ref-format"))
 		require.NotContains(t, call, "--create")
@@ -126,7 +127,7 @@ func TestPullCheckoutReusesMatchingBranch(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, "abc123", result.Commit)
 	require.Contains(t, runner.calls, "rev-parse --verify refs/heads/pr/7^{commit}")
-	require.Equal(t, "switch pr/7", runner.calls[len(runner.calls)-1])
+	require.Equal(t, "switch pr/7", runner.calls[len(runner.calls)-3])
 	runner = &checkoutRunner{collision: true}
 	_, err = CheckoutPull(context.Background(), runner, CheckoutOptions{Directory: "/work", Remote: "origin", Number: 7, Branch: "pr/7", ExpectedCommit: "different"})
 	require.ErrorContains(t, err, "already exists with a different commit")
@@ -149,7 +150,7 @@ func TestPullCheckoutUpdatesOnlyByFastForward(t *testing.T) {
 			}
 		} else {
 			require.NoError(t, err)
-			require.Equal(t, []string{"switch pr/7", "merge --ff-only --no-edit abc123"}, runner.calls[len(runner.calls)-2:])
+			require.Equal(t, []string{"switch pr/7", "merge --ff-only --no-edit abc123"}, runner.calls[len(runner.calls)-4:len(runner.calls)-2])
 		}
 	}
 }
@@ -166,6 +167,7 @@ func TestPullCheckoutTracksHeadBranch(t *testing.T) {
 	require.Contains(t, runner.calls, "fetch --no-tags -- pulls/bob refs/heads/feature")
 	require.Contains(t, runner.calls, "config branch.pr/7.remote pulls/bob")
 	require.Contains(t, runner.calls, "config branch.pr/7.merge refs/heads/feature")
+	require.Equal(t, []string{"config branch.pr/7.tea-pr 7", "config branch.pr/7.tea-base-remote origin"}, runner.calls[len(runner.calls)-2:])
 	for _, call := range runner.calls {
 		require.NotContains(t, call, "refs/pull/")
 	}

internal/git/discover.go

@@ -14,12 +14,17 @@ import (
 var ErrAmbiguousRemote = errors.New("ambiguous repository: select --remote or --repo")
 
 // DiscoverRemote names the remote used for repository discovery: preferred when set, otherwise
-// the current branch's upstream remote, then upstream, origin, or the only remote.
+// the base remote pr checkout recorded for the current branch, the current branch's upstream
+// remote unless it is a pulls/<owner>/<repo> head remote, then upstream, origin, or the only remote.
 func DiscoverRemote(ctx context.Context, runner Runner, dir, preferred string) (string, error) {
 	if preferred != "" {
 		return preferred, nil
 	}
-	if _, remote := CurrentUpstreamRemote(ctx, runner, dir); remote != "" {
+	branch, remote := CurrentUpstreamRemote(ctx, runner, dir)
+	if base := pullBaseRemote(ctx, runner, dir, branch); base != "" {
+		return base, nil
+	}
+	if remote != "" && !strings.HasPrefix(remote, "pulls/") {
 		return remote, nil
 	}
 	names, err := runner.Output(ctx, dir, false, "remote")
@@ -61,3 +66,20 @@ func UpstreamRemote(ctx context.Context, runner Runner, dir, branch string) stri
 	}
 	return remote
 }
+
+// pullBaseRemote returns the base remote pr checkout recorded for branch, or "" when there is
+// none or that remote no longer exists.
+func pullBaseRemote(ctx context.Context, runner Runner, dir, branch string) string {
+	if branch == "" {
+		return ""
+	}
+	base, err := runner.Output(ctx, dir, false, "config", "--get", "branch."+branch+"."+pullBaseRemoteKey)
+	if err != nil || base == "" {
+		return ""
+	}
+	names, err := runner.Output(ctx, dir, false, "remote")
+	if err != nil || !slices.Contains(strings.Fields(names), base) {
+		return ""
+	}
+	return base
+}

internal/git/discover_test.go

@@ -23,7 +23,7 @@ func (r discoverRunner) Output(_ context.Context, dir string, prompt bool, args
 }
 
 func TestDiscoverRemote(t *testing.T) {
-	const head, tracking = "symbolic-ref --short HEAD", "config --get branch.main.remote"
+	const head, tracking, base = "symbolic-ref --short HEAD", "config --get branch.main.remote", "config --get branch.main.tea-base-remote"
 	for _, test := range []struct {
 		name      string
 		runner    discoverRunner
@@ -33,7 +33,10 @@ func TestDiscoverRemote(t *testing.T) {
 	}{
 		{name: "preferred", preferred: "fork", want: "fork"},
 		{name: "branch upstream", runner: discoverRunner{head: "main", tracking: "fork", "remote": "origin\nfork"}, want: "fork"},
-		{name: "branch upstream with slashes", runner: discoverRunner{head: "main", tracking: "pulls/bob/project", "remote": "origin\npulls/bob/project"}, want: "pulls/bob/project"},
+		{name: "branch upstream with slashes", runner: discoverRunner{head: "main", tracking: "team/fork", "remote": "origin\nteam/fork"}, want: "team/fork"},
+		{name: "pull head remote", runner: discoverRunner{head: "main", tracking: "pulls/bob/project", "remote": "origin\npulls/bob/project"}, want: "origin"},
+		{name: "recorded pull base", runner: discoverRunner{head: "main", tracking: "fork", base: "upstream", "remote": "origin\nupstream\nfork"}, want: "upstream"},
+		{name: "removed pull base", runner: discoverRunner{head: "main", tracking: "fork", base: "gone", "remote": "origin\nfork"}, want: "fork"},
 		{name: "local branch upstream", runner: discoverRunner{head: "main", tracking: ".", "remote": "origin\nfork"}, want: "origin"},
 		{name: "upstream remote", runner: discoverRunner{"remote": "origin\nupstream"}, want: "upstream"},
 		{name: "origin remote", runner: discoverRunner{"remote": "fork\norigin"}, want: "origin"},