Switch branch quick pick is not showing all branches (#8365)

Fixes #8351

Author: alexr00

src/common/async.ts

@@ -36,13 +36,13 @@ export function throttle<T>(fn: () => Promise<T>): () => Promise<T> {
 	return trigger;
 }
 
-export function debounce(fn: () => any, delay: number): () => void {
+export function debounce<T extends (...args: any[]) => any>(fn: T, delay: number): (...args: Parameters<T>) => void {
 	let timer: NodeJS.Timeout | undefined;
 
-	return () => {
+	return (...args: Parameters<T>) => {
 		if (timer) {
 			clearTimeout(timer);
 		}
-		timer = setTimeout(() => fn(), delay);
+		timer = setTimeout(() => fn(...args), delay);
 	};
 }

src/github/createPRViewProvider.ts

@@ -21,6 +21,7 @@ import { ChangeTemplateReply, DisplayLabel, PreReviewState } from './views';
 import { RemoteInfo } from '../../common/types';
 import { ChooseBaseRemoteAndBranchResult, ChooseCompareRemoteAndBranchResult, ChooseRemoteAndBranchArgs, CreateParamsNew, CreatePullRequestNew, TitleAndDescriptionArgs } from '../../common/views';
 import type { Branch } from '../api/api';
+import { debounce } from '../common/async';
 import { GitHubServerType } from '../common/authentication';
 import { emojify, ensureEmojis } from '../common/emoji';
 import { commands, contexts } from '../common/executeCommands';
@@ -967,10 +968,35 @@ export class CreatePullRequestViewProvider extends BaseCreatePullRequestViewProv
 		const branchPlaceholder = isBase ? vscode.l10n.t('Choose a base branch') : vscode.l10n.t('Choose a branch to merge');
 		const repositoryPlaceholder = isBase ? vscode.l10n.t('Choose a base repository') : vscode.l10n.t('Choose a repository to merge from');
 
+		let updateCounter = 0;
+		const updateItems = async (githubRepository: GitHubRepository, prefix: string | undefined) => {
+			const currentUpdate = ++updateCounter;
+			quickPick.busy = true;
+			const items = await branchPicks(githubRepository, this._folderRepositoryManager, chooseDifferentRemote, isBase, prefix);
+			if (currentUpdate === updateCounter) {
+				quickPick.items = items;
+				quickPick.busy = false;
+			}
+		};
+		const debounced = debounce(updateItems, 300);
+		let onDidChangeValueDisposable: vscode.Disposable | undefined;
+		const addValueChangeListener = () => {
+			if (githubRepository && !onDidChangeValueDisposable) {
+				onDidChangeValueDisposable = quickPick.onDidChangeValue(async value => {
+					return debounced(githubRepository!, value);
+				});
+			}
+		};
+		addValueChangeListener();
+
 		quickPick.placeholder = githubRepository ? branchPlaceholder : remotePlaceholder;
 		quickPick.show();
 		quickPick.busy = true;
-		quickPick.items = githubRepository ? await branchPicks(githubRepository, this._folderRepositoryManager, chooseDifferentRemote, isBase) : await this.remotePicks(isBase);
+		if (githubRepository) {
+			await updateItems(githubRepository, undefined);
+		} else {
+			quickPick.items = await this.remotePicks(isBase);
+		}
 		const activeItem = message.args.currentBranch ? quickPick.items.find(item => item.branch === message.args.currentBranch) : undefined;
 		quickPick.activeItems = activeItem ? [activeItem] : [];
 		quickPick.busy = false;
@@ -989,7 +1015,8 @@ export class CreatePullRequestViewProvider extends BaseCreatePullRequestViewProv
 					const selectedRemote = selectedPick as vscode.QuickPickItem & { remote: RemoteInfo };
 					quickPick.busy = true;
 					githubRepository = this._folderRepositoryManager.findRepo(repo => repo.remote.owner === selectedRemote.remote.owner && repo.remote.repositoryName === selectedRemote.remote.repositoryName)!;
-					quickPick.items = await branchPicks(githubRepository, this._folderRepositoryManager, chooseDifferentRemote, isBase);
+					await updateItems(githubRepository, undefined);
+					addValueChangeListener();
 					quickPick.placeholder = branchPlaceholder;
 					quickPick.busy = false;
 				} else if (selectedPick.branch && selectedPick.remote) {
@@ -1003,6 +1030,7 @@ export class CreatePullRequestViewProvider extends BaseCreatePullRequestViewProv
 		if (!result || !githubRepository) {
 			quickPick.hide();
 			quickPick.dispose();
+			onDidChangeValueDisposable?.dispose();
 			return;
 		}
 
@@ -1011,6 +1039,7 @@ export class CreatePullRequestViewProvider extends BaseCreatePullRequestViewProv
 
 		quickPick.hide();
 		quickPick.dispose();
+		onDidChangeValueDisposable?.dispose();
 		return this._replyMessage(message, chooseResult);
 	}
 

src/github/githubRepository.ts

@@ -1269,7 +1269,7 @@ export class GitHubRepository extends Disposable {
 		return data.repository?.ref?.target.oid;
 	}
 
-	async listBranches(owner: string, repositoryName: string): Promise<string[]> {
+	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);
 
@@ -1289,6 +1289,7 @@ export class GitHubRepository extends Disposable {
 						name: remote.repositoryName,
 						first: 100,
 						after: after,
+						query: prefix ? prefix : null,
 					},
 				});
 

src/github/pullRequestOverview.ts

@@ -26,6 +26,7 @@ import { PullRequestReviewCommon, ReviewContext } from './pullRequestReviewCommo
 import { branchPicks, pickEmail, reviewersQuickPick } from './quickPicks';
 import { parseReviewers } from './utils';
 import { CancelCodingAgentReply, ChangeBaseReply, ChangeReviewersReply, DeleteReviewResult, MergeArguments, MergeResult, PullRequest, ReviewType } from './views';
+import { debounce } from '../common/async';
 import { COPILOT_ACCOUNTS, IComment } from '../common/comment';
 import { COPILOT_REVIEWER, COPILOT_REVIEWER_ACCOUNT, COPILOT_SWE_AGENT, copilotEventToStatus, CopilotPRStatus, mostRecentCopilotEvent } from '../common/copilot';
 import { commands, contexts } from '../common/executeCommands';
@@ -923,14 +924,27 @@ export class PullRequestOverviewPanel extends IssueOverviewPanel<PullRequestMode
 
 	private async changeBaseBranch(message: IRequestMessage<void>): Promise<void> {
 		const quickPick = vscode.window.createQuickPick<vscode.QuickPickItem & { branch?: string }>();
+		let updateCounter = 0;
+		const updateItems = async (prefix: string | undefined) => {
+			const currentUpdate = ++updateCounter;
+			quickPick.busy = true;
+			const items = await branchPicks(this._item.githubRepository, this._folderRepositoryManager, undefined, true, prefix);
+			if (currentUpdate === updateCounter) {
+				quickPick.items = items;
+				quickPick.busy = false;
+			}
+		};
+		const debounced = debounce(updateItems, 300);
+		const onDidChangeValueDisposable = quickPick.onDidChangeValue(async value => {
+			return debounced(value);
+		});
 
 		try {
 			quickPick.busy = true;
 			quickPick.canSelectMany = false;
 			quickPick.placeholder = vscode.l10n.t('Select a new base branch');
 			quickPick.show();
-
-			quickPick.items = await branchPicks(this._item.githubRepository, this._folderRepositoryManager, undefined, true);
+			await updateItems(undefined);
 
 			quickPick.busy = false;
 			const acceptPromise = asPromise<void>(quickPick.onDidAccept).then(() => {
@@ -961,6 +975,7 @@ export class PullRequestOverviewPanel extends IssueOverviewPanel<PullRequestMode
 			vscode.window.showErrorMessage(formatError(e));
 		} finally {
 			quickPick.hide();
+			onDidChangeValueDisposable.dispose();
 			quickPick.dispose();
 		}
 	}

src/github/queriesShared.gql

@@ -945,9 +945,9 @@ query GetBranch($owner: String!, $name: String!, $qualifiedName: String!) {
 	}
 }
 
-query ListBranches($owner: String!, $name: String!, $first: Int!, $after: String) {
+query ListBranches($owner: String!, $name: String!, $first: Int!, $after: String, $query: String) {
 	repository(owner: $owner, name: $name) {
-		refs(first: $first, after: $after, refPrefix: "refs/heads/") {
+		refs(first: $first, after: $after, refPrefix: "refs/heads/", query: $query) {
 			nodes {
 				name
 			}

src/github/quickPicks.ts

@@ -482,11 +482,11 @@ function getRecentlyUsedBranches(folderRepoManager: FolderRepositoryManager, own
 	return state.branches[repoKey] || [];
 }
 
-export async function branchPicks(githubRepository: GitHubRepository, folderRepoManager: FolderRepositoryManager, changeRepoMessage: string | undefined, isBase: boolean): Promise<(vscode.QuickPickItem & { remote?: RemoteInfo, branch?: string })[]> {
+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);
+		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);