Associate local branches with PRs on first activation (#8227)

* Initial plan

* Initial plan for associating local branches with PRs on first activation

Co-authored-by: alexr00 <38270282+alexr00@users.noreply.github.com>

* Implement branch association logic for first activation

Co-authored-by: alexr00 <38270282+alexr00@users.noreply.github.com>

* Address code review feedback - document chunk size rationale

Co-authored-by: alexr00 <38270282+alexr00@users.noreply.github.com>

* Fix comment - clarify chunk size rationale

Co-authored-by: alexr00 <38270282+alexr00@users.noreply.github.com>

* Limit branch association to 3 most recently used branches

Co-authored-by: alexr00 <38270282+alexr00@users.noreply.github.com>

* Bump it to 10

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: alexr00 <38270282+alexr00@users.noreply.github.com>

Author: Copilot

src/extensionState.ts

@@ -13,6 +13,7 @@ export const NEVER_SHOW_PULL_NOTIFICATION = 'github.pullRequest.pullNotification
 export const REPO_KEYS = 'github.pullRequest.repos';
 export const PREVIOUS_CREATE_METHOD = 'github.pullRequest.previousCreateMethod';
 export const LAST_USED_EMAIL = 'github.pullRequest.lastUsedEmail';
+export const BRANCHES_ASSOCIATED_WITH_PRS = 'github.pullRequest.branchesAssociatedWithPRs';
 
 export interface RepoState {
 	mentionableUsers?: IAccount[];

src/github/folderRepositoryManager.ts

@@ -54,7 +54,7 @@ import { EventType } from '../common/timelineEvent';
 import { Schemes } from '../common/uri';
 import { AsyncPredicate, batchPromiseAll, compareIgnoreCase, formatError, Predicate } from '../common/utils';
 import { PULL_REQUEST_OVERVIEW_VIEW_TYPE } from '../common/webview';
-import { LAST_USED_EMAIL, NEVER_SHOW_PULL_NOTIFICATION, REPO_KEYS, ReposState } from '../extensionState';
+import { BRANCHES_ASSOCIATED_WITH_PRS, LAST_USED_EMAIL, NEVER_SHOW_PULL_NOTIFICATION, REPO_KEYS, ReposState } from '../extensionState';
 import { git } from '../gitProviders/gitCommands';
 import { IThemeWatcher } from '../themeWatcher';
 import { CreatePullRequestHelper } from '../view/createPullRequestHelper';
@@ -576,6 +576,11 @@ export class FolderRepositoryManager extends Disposable {
 			this.getAssignableUsers(repositoriesAdded.length > 0);
 			if (isAuthenticated && activeRemotes.length) {
 				this.state = ReposManagerState.RepositoriesLoaded;
+				// On first activation, associate local branches with PRs
+				// Do this asynchronously to not block the main flow
+				this.associateLocalBranchesWithPRsOnFirstActivation().catch(e => {
+					Logger.error(`Failed to associate branches with PRs: ${e}`, this.id);
+				});
 			} else if (!isAuthenticated) {
 				this.state = ReposManagerState.NeedsAuthentication;
 			}
@@ -953,6 +958,112 @@ export class FolderRepositoryManager extends Disposable {
 		return models.filter(value => value !== undefined) as PullRequestModel[];
 	}
 
+	/**
+	 * On first activation, iterate through local branches and associate them with PRs if they match.
+	 * This helps discover PRs that were created before the extension was installed or in other ways.
+	 */
+	private async associateLocalBranchesWithPRsOnFirstActivation(): Promise<void> {
+		const stateKey = `${BRANCHES_ASSOCIATED_WITH_PRS}.${this.repository.rootUri.fsPath}`;
+		const hasRun = this.context.globalState.get<boolean>(stateKey, false);
+
+		if (hasRun) {
+			Logger.debug('Branch association has already run for this workspace folder', this.id);
+			return;
+		}
+
+		Logger.appendLine('First activation: associating local branches with PRs', this.id);
+
+		const githubRepositories = this._githubRepositories;
+		if (!githubRepositories || !githubRepositories.length || !this.repository.getRefs) {
+			Logger.debug('No GitHub repositories or getRefs not available, skipping branch association', this.id);
+			await this.context.globalState.update(stateKey, true);
+			return;
+		}
+
+		try {
+			// Only check the 3 most recently used branches to minimize API calls
+			const localBranches = (await this.repository.getRefs({
+				pattern: 'refs/heads/',
+				sort: 'committerdate',
+				count: 10
+			}))
+				.filter(r => r.name !== undefined)
+				.map(r => r.name!);
+
+			Logger.debug(`Found ${localBranches.length} local branches to check`, this.id);
+
+			const associationResults: boolean[] = [];
+
+			// Process all branches (max 3) in parallel
+			const chunkResults = await Promise.all(localBranches.map(async branchName => {
+				try {
+					// Check if this branch already has PR metadata
+					const existingMetadata = await PullRequestGitHelper.getMatchingPullRequestMetadataForBranch(
+						this.repository,
+						branchName,
+					);
+
+					if (existingMetadata) {
+						// Branch already has PR metadata, skip
+						return false;
+					}
+
+					// Get the branch to check its upstream
+					const branch = await this.repository.getBranch(branchName);
+					if (!branch.upstream) {
+						// No upstream, can't match to a PR
+						return false;
+					}
+
+					// Try to find a matching PR on GitHub
+					const remoteName = branch.upstream.remote;
+					const upstreamBranchName = branch.upstream.name;
+
+					const githubRepo = githubRepositories.find(
+						repo => repo.remote.remoteName === remoteName,
+					);
+
+					if (!githubRepo) {
+						return false;
+					}
+
+					// Get the metadata of the GitHub repository to find owner
+					const metadata = await githubRepo.getMetadata();
+					if (!metadata?.owner) {
+						return false;
+					}
+
+					// Search for a PR with this head branch
+					const matchingPR = await githubRepo.getPullRequestForBranch(upstreamBranchName, metadata.owner.login);
+
+					if (matchingPR) {
+						Logger.appendLine(`Found PR #${matchingPR.number} for branch ${branchName}, associating...`, this.id);
+						await PullRequestGitHelper.associateBranchWithPullRequest(
+							this.repository,
+							matchingPR,
+							branchName,
+						);
+						return true;
+					}
+					return false;
+				} catch (e) {
+					Logger.debug(`Error checking branch ${branchName}: ${e}`, this.id);
+					// Continue with other branches even if one fails
+					return false;
+				}
+			}));
+			associationResults.push(...chunkResults);
+
+			const associatedCount = associationResults.filter(r => r).length;
+			Logger.appendLine(`Branch association complete: ${associatedCount} branches associated with PRs`, this.id);
+		} catch (e) {
+			Logger.error(`Error during branch association: ${e}`, this.id);
+		} finally {
+			// Mark as complete even if there were errors
+			await this.context.globalState.update(stateKey, true);
+		}
+	}
+
 	async getLabels(issue?: IssueModel, repoInfo?: { owner: string; repo: string }): Promise<ILabel[]> {
 		const repo = issue
 			? issue.githubRepository