Add branch caching for faster branch picker in create PR view

Co-authored-by: ulugbekna <16353531+ulugbekna@users.noreply.github.com>

Author: Copilot

src/github/createPRViewProvider.ts

@@ -15,7 +15,7 @@ import { IAccount, ILabel, IMilestone, IProject, isITeam, ITeam, MergeMethod, Re
 import { BaseBranchMetadata, PullRequestGitHelper } from './pullRequestGitHelper';
 import { PullRequestModel } from './pullRequestModel';
 import { getDefaultMergeMethod } from './pullRequestOverview';
-import { branchPicks, getAssigneesQuickPickItems, getLabelOptions, getMilestoneFromQuickPick, getProjectFromQuickPick, reviewersQuickPick } from './quickPicks';
+import { branchPicks, cachedBranchPicks, getAssigneesQuickPickItems, getLabelOptions, getMilestoneFromQuickPick, getProjectFromQuickPick, reviewersQuickPick } from './quickPicks';
 import { ISSUE_EXPRESSION, parseIssueExpressionOutput, variableSubstitution } from './utils';
 import { ChangeTemplateReply, DisplayLabel, PreReviewState } from './views';
 import { RemoteInfo } from '../../common/types';
@@ -1010,9 +1010,22 @@ Don't forget to commit your template file to the repository so that it can be us
 		const params = await super.getCreateParams();
 		this.model.baseOwner = params.defaultBaseRemote!.owner;
 		this.model.baseBranch = params.defaultBaseBranch!;
+		// Pre-fetch branches so they're cached when the user opens the branch picker
+		this.prefetchBranches(params.defaultBaseRemote!);
 		return params;
 	}
 
+	private prefetchBranches(baseRemote: RemoteInfo): void {
+		const githubRepository = this._folderRepositoryManager.findRepo(
+			repo => repo.remote.owner === baseRemote.owner && repo.remote.repositoryName === baseRemote.repositoryName,
+		);
+		if (githubRepository) {
+			githubRepository.listBranches(baseRemote.owner, baseRemote.repositoryName, undefined).catch(e => {
+				Logger.debug(`Pre-fetching branches failed: ${e}`, CreatePullRequestViewProvider.ID);
+			});
+		}
+	}
+
 
 	private async remotePicks(isBase: boolean): Promise<(vscode.QuickPickItem & { remote?: RemoteInfo })[]> {
 		const remotes = isBase ? await this._folderRepositoryManager.getActiveGitHubRemotes(await this._folderRepositoryManager.getGitHubRemotes()) : this._folderRepositoryManager.gitHubRepositories.map(repo => repo.remote);
@@ -1141,6 +1154,13 @@ Don't forget to commit your template file to the repository so that it can be us
 		quickPick.show();
 		quickPick.busy = true;
 		if (githubRepository) {
+			// Show cached branches immediately if available, then refresh in the background
+			const cached = cachedBranchPicks(githubRepository, this._folderRepositoryManager, chooseDifferentRemote, isBase);
+			if (cached) {
+				quickPick.items = cached;
+				const activeItem = message.args.currentBranch ? quickPick.items.find(item => item.branch === message.args.currentBranch) : undefined;
+				quickPick.activeItems = activeItem ? [activeItem] : [];
+			}
 			await updateItems(githubRepository, undefined);
 		} else {
 			quickPick.items = await this.remotePicks(isBase);

src/github/githubRepository.ts

@@ -212,6 +212,8 @@ export class GitHubRepository extends Disposable {
 	private _areQueriesLimited: boolean = false;
 	get areQueriesLimited(): boolean { return this._areQueriesLimited; }
 
+	private _branchesCache: Map<string, string[]> = new Map();
+
 	private _onDidAddPullRequest: vscode.EventEmitter<PullRequestModel> = this._register(new vscode.EventEmitter());
 	public readonly onDidAddPullRequest: vscode.Event<PullRequestModel> = this._onDidAddPullRequest.event;
 	private _onDidChangePullRequests: vscode.EventEmitter<PullRequestChangeEvent[]> = this._register(new vscode.EventEmitter());
@@ -1341,6 +1343,14 @@ export class GitHubRepository extends Disposable {
 		return data.repository?.ref?.target.oid;
 	}
 
+	private static branchesCacheKey(owner: string, repositoryName: string): string {
+		return `${owner}/${repositoryName}`;
+	}
+
+	getCachedBranches(owner: string, repositoryName: string): string[] | undefined {
+		return this._branchesCache.get(GitHubRepository.branchesCacheKey(owner, repositoryName));
+	}
+
 	async listBranches(owner: string, repositoryName: string, prefix: string | undefined): Promise<string[]> {
 		const { query, remote, schema } = await this.ensure();
 		Logger.debug(`List branches for ${owner}/${repositoryName} - enter`, this.id);
@@ -1382,6 +1392,10 @@ export class GitHubRepository extends Disposable {
 		if (!branches.includes(defaultBranch)) {
 			branches.unshift(defaultBranch);
 		}
+		// Cache results for unprefixed queries
+		if (!prefix) {
+			this._branchesCache.set(GitHubRepository.branchesCacheKey(owner, repositoryName), branches);
+		}
 		return branches;
 	}
 

src/github/quickPicks.ts

@@ -482,19 +482,7 @@ function getRecentlyUsedBranches(folderRepoManager: FolderRepositoryManager, own
 	return state.branches[repoKey] || [];
 }
 
-export async function branchPicks(githubRepository: GitHubRepository, folderRepoManager: FolderRepositoryManager, changeRepoMessage: string | undefined, isBase: boolean, prefix: string | undefined): Promise<(vscode.QuickPickItem & { remote?: RemoteInfo, branch?: string })[]> {
-	let branches: (string | Ref)[];
-	if (isBase) {
-		// For the base, we only want to show branches from GitHub.
-		branches = await githubRepository.listBranches(githubRepository.remote.owner, githubRepository.remote.repositoryName, prefix);
-	} else {
-		// For the compare, we only want to show local branches.
-		branches = (await folderRepoManager.repository.getBranches({ remote: false })).filter(branch => branch.name);
-	}
-
-
-	const branchNames = branches.map(branch => typeof branch === 'string' ? branch : branch.name!);
-
+function buildBranchPickItems(branchNames: string[], githubRepository: GitHubRepository, folderRepoManager: FolderRepositoryManager, changeRepoMessage: string | undefined, isBase: boolean): (vscode.QuickPickItem & { remote?: RemoteInfo, branch?: string })[] {
 	// Get recently used branches for base branches only
 	let recentBranches: string[] = [];
 	let otherBranches: string[] = branchNames;
@@ -554,4 +542,29 @@ export async function branchPicks(githubRepository: GitHubRepository, folderRepo
 		});
 	}
 	return branchPicks;
+}
+
+export function cachedBranchPicks(githubRepository: GitHubRepository, folderRepoManager: FolderRepositoryManager, changeRepoMessage: string | undefined, isBase: boolean): (vscode.QuickPickItem & { remote?: RemoteInfo, branch?: string })[] | undefined {
+	if (!isBase) {
+		return undefined;
+	}
+	const cached = githubRepository.getCachedBranches(githubRepository.remote.owner, githubRepository.remote.repositoryName);
+	if (!cached) {
+		return undefined;
+	}
+	return buildBranchPickItems(cached, githubRepository, folderRepoManager, changeRepoMessage, isBase);
+}
+
+export async function branchPicks(githubRepository: GitHubRepository, folderRepoManager: FolderRepositoryManager, changeRepoMessage: string | undefined, isBase: boolean, prefix: string | undefined): Promise<(vscode.QuickPickItem & { remote?: RemoteInfo, branch?: string })[]> {
+	let branches: (string | Ref)[];
+	if (isBase) {
+		// For the base, we only want to show branches from GitHub.
+		branches = await githubRepository.listBranches(githubRepository.remote.owner, githubRepository.remote.repositoryName, prefix);
+	} else {
+		// For the compare, we only want to show local branches.
+		branches = (await folderRepoManager.repository.getBranches({ remote: false })).filter(branch => branch.name);
+	}
+
+	const branchNames = branches.map(branch => typeof branch === 'string' ? branch : branch.name!);
+	return buildBranchPickItems(branchNames, githubRepository, folderRepoManager, changeRepoMessage, isBase);
 }

src/test/github/githubRepository.test.ts

@@ -52,4 +52,24 @@ describe('GitHubRepository', function () {
 			// assert(! dotcomRepository.isGitHubDotCom);
 		});
 	});
+
+	describe('getCachedBranches', function () {
+		it('returns undefined when no branches have been cached', function () {
+			const url = 'https://github.com/owner/repo';
+			const remote = new GitHubRemote('origin', url, new Protocol(url), GitHubServerType.GitHubDotCom);
+			const rootUri = Uri.file('C:\\users\\test\\repo');
+			const repo = new GitHubRepository(1, remote, rootUri, credentialStore, telemetry);
+
+			assert.strictEqual(repo.getCachedBranches('owner', 'repo'), undefined);
+		});
+
+		it('returns undefined for a different owner/repo than what was cached', function () {
+			const url = 'https://github.com/owner/repo';
+			const remote = new GitHubRemote('origin', url, new Protocol(url), GitHubServerType.GitHubDotCom);
+			const rootUri = Uri.file('C:\\users\\test\\repo');
+			const repo = new GitHubRepository(1, remote, rootUri, credentialStore, telemetry);
+
+			assert.strictEqual(repo.getCachedBranches('other-owner', 'other-repo'), undefined);
+		});
+	});
 });