Fix "checkout main" on restored PR webview

src/github/externalUriOpener.ts

@@ -45,7 +45,6 @@ class GitHubIssueOrPullRequestExternalUriOpener extends Disposable implements vs
 			return;
 		}
 
-		const folderRepositoryManager = this._folderRepositoryManagerResolver.getManagerForRepository(identity.owner, identity.repo);
 		const requireModel = async <T extends IssueModel>(model: T | undefined): Promise<T> => {
 			if (token.isCancellationRequested) {
 				throw new vscode.CancellationError();
@@ -58,6 +57,7 @@ class GitHubIssueOrPullRequestExternalUriOpener extends Disposable implements vs
 		};
 
 		try {
+			const folderRepositoryManager = this._folderRepositoryManagerResolver.getManagerForRepository(identity.owner, identity.repo);
 			if (identity.kind === 'pullRequest') {
 				const pullRequest = folderRepositoryManager.resolvePullRequest(identity.owner, identity.repo, identity.number, true, 'overview').then(requireModel);
 				await PullRequestOverviewPanel.createOrShow(

src/github/folderRepositoryManagerResolver.ts

@@ -30,6 +30,10 @@ export class FolderRepositoryManagerResolver extends Disposable {
 		if (existingManager) {
 			return existingManager;
 		}
+		return this.getRemoteOnlyManager();
+	}
+
+	getRemoteOnlyManager(): FolderRepositoryManager {
 		if (this._remoteFolderRepositoryManager) {
 			return this._remoteFolderRepositoryManager;
 		}

src/github/issueOverview.ts

@@ -6,6 +6,7 @@
 
 import * as vscode from 'vscode';
 import { CloseResult, OpenLocalFileArgs } from '../../common/views';
+import { RemoteOnlyRepository } from '../api/remoteOnlyRepository';
 import { openItemOnGitHub } from '../commands';
 import { decodeBase64, guessExtensionFromMime, pickFilesForUpload, placeholdersForNames, runFileUploads, runPendingUploads } from './fileUpload';
 import { FolderRepositoryManager } from './folderRepositoryManager';
@@ -14,6 +15,7 @@ import { GithubItemStateEnum, IAccount, IMilestone, IProject, IProjectItem, Repo
 import { IssueModel } from './issueModel';
 import { openIssueOrPullRequestOnGitHub } from './openOnGitHub';
 import { getAssigneesQuickPickItems, getLabelOptions, getMilestoneFromQuickPick, getProjectFromQuickPick } from './quickPicks';
+import type { RepositoriesManager } from './repositoriesManager';
 import { isInCodespaces, processPermalinks, vscodeDevPrLink } from './utils';
 import { ChangeAssigneesReply, DisplayLabel, FileUploadCompletedMessage, Issue, IssuePreview, OverviewItemPreview, ProjectItemsReply, SubmitReviewArgs, SubmitReviewReply, UnresolvedIdentity, UploadFilesReply, UploadPastedFilesArgs } from './views';
 import { COPILOT_ACCOUNTS, IComment } from '../common/comment';
@@ -36,17 +38,49 @@ export class IssueOverviewPanel<TItem extends IssueModel = IssueModel> extends W
 	 * All open panels, keyed by "owner/repo#number".
 	 */
 	protected static _panels: Map<string, IssueOverviewPanel> = new Map();
+	protected static _repositoriesManager: RepositoriesManager | undefined;
 
 	public static readonly viewType: string = 'IssueOverview';
 
 	protected readonly _panel: vscode.WebviewPanel;
 	protected _item: TItem;
 	protected _identity: UnresolvedIdentity;
-	protected _folderRepositoryManager: FolderRepositoryManager;
+	protected readonly _initialFolderRepositoryManager: FolderRepositoryManager;
+	protected _localFolderRepositoryManager: FolderRepositoryManager | undefined;
 	protected _scrollPosition = { x: 0, y: 0 };
 	private _identityUpdateSequence = 0;
 	protected readonly previewLog = { label: 'Issue', id: IssueOverviewPanel.ID };
 
+	protected get _folderRepositoryManager(): FolderRepositoryManager {
+		if (this._localFolderRepositoryManager) {
+			return this._localFolderRepositoryManager;
+		}
+		const manager = this._initialFolderRepositoryManager;
+		if (!this.isDisposed && manager.repository instanceof RemoteOnlyRepository && this._identity) {
+			const localManager = IssueOverviewPanel._repositoriesManager?.getManagerForRepository(this._identity.owner, this._identity.repo);
+			if (localManager && !(localManager.repository instanceof RemoteOnlyRepository)) {
+				this._localFolderRepositoryManager = localManager;
+				this.registerPrListeners();
+				Logger.debug(`Upgraded remote-only manager for ${this._identity.owner}/${this._identity.repo} to ${localManager.repository.rootUri.toString()}`, this.previewLog.id);
+				return localManager;
+			}
+		}
+		return manager;
+	}
+
+	protected static registerRepositoriesManager(context: vscode.ExtensionContext, repositoriesManager: RepositoriesManager): void {
+		IssueOverviewPanel._repositoriesManager = repositoriesManager;
+		context.subscriptions.push(
+			{
+				dispose: () => {
+					if (IssueOverviewPanel._repositoriesManager === repositoriesManager) {
+						IssueOverviewPanel._repositoriesManager = undefined;
+					}
+				}
+			},
+		);
+	}
+
 	protected static _getViewColumn(toTheSide: boolean, panel?: IssueOverviewPanel): number | undefined {
 		const tabViewColumn = vscode.window.tabGroups.activeTabGroup.viewColumn;
 		const activeColumn = toTheSide
@@ -90,7 +124,7 @@ export class IssueOverviewPanel<TItem extends IssueModel = IssueModel> extends W
 			this._panels.set(key, panel);
 		}
 
-		await panel.updateWithIdentity(folderRepositoryManager, identity, issue);
+		await panel.updateWithIdentity(identity, issue);
 	}
 
 	public static refresh(owner: string, repo: string, number: number): void {
@@ -167,7 +201,7 @@ export class IssueOverviewPanel<TItem extends IssueModel = IssueModel> extends W
 			}
 	) {
 		super();
-		this._folderRepositoryManager = folderRepositoryManager;
+		this._initialFolderRepositoryManager = folderRepositoryManager;
 
 		// Create and show a new webview panel
 		this._panel = existingPanel ?? this._register(vscode.window.createWebviewPanel(type, title, column, {
@@ -192,17 +226,6 @@ export class IssueOverviewPanel<TItem extends IssueModel = IssueModel> extends W
 		// This happens when the user closes the panel or when the panel is closed programmatically
 		this._register(this._panel.onDidDispose(() => this.dispose()));
 
-		this._register(this._folderRepositoryManager.onDidChangeActiveIssue(
-			_ => {
-				if (this._folderRepositoryManager && this._item) {
-					const isCurrentlyCheckedOut = this._item.equals(this._folderRepositoryManager.activeIssue);
-					this._postMessage({
-						command: 'pr.update-checkout-status',
-						isCurrentlyCheckedOut: isCurrentlyCheckedOut,
-					});
-				}
-			}));
-
 		this._register(folderRepositoryManager.credentialStore.onDidUpgradeSession(() => {
 			this.updateItem(this._item);
 		}));
@@ -418,13 +441,13 @@ export class IssueOverviewPanel<TItem extends IssueModel = IssueModel> extends W
 	/**
 	 * Update the panel with an unresolved identity and optional model.
 	 * If no model is provided, it will be resolved from the identity.
+	 * Repository manager selection is handled by the panel, not by updates.
 	 */
-	public async updateWithIdentity(foldersManager: FolderRepositoryManager, identity: UnresolvedIdentity, issueModel?: TItem | Promise<TItem | undefined>, progressLocation?: string): Promise<void> {
+	public async updateWithIdentity(identity: UnresolvedIdentity, issueModel?: TItem | Promise<TItem | undefined>, progressLocation?: string): Promise<void> {
 		const updateSequence = ++this._identityUpdateSequence;
 		let loading = true;
 		const isLoading = () => loading && !this.isDisposed && updateSequence === this._identityUpdateSequence;
 		this._identity = identity;
-		this._folderRepositoryManager = foldersManager;
 
 		this._postMessage({
 			command: 'set-scroll',
@@ -442,7 +465,7 @@ export class IssueOverviewPanel<TItem extends IssueModel = IssueModel> extends W
 			void (async () => {
 				try {
 					const start = Date.now();
-					const repository = await foldersManager.createGitHubRepositoryFromOwnerName(identity.owner, identity.repo, false);
+					const repository = await this._folderRepositoryManager.createGitHubRepositoryFromOwnerName(identity.owner, identity.repo, false);
 					if (!isLoading()) {
 						return;
 					}

src/github/overviewRestorer.ts

@@ -33,7 +33,7 @@ export class OverviewRestorer extends Disposable implements vscode.WebviewPanelS
 
 		await this.waitForAuth();
 
-		const folderManager = this._folderRepositoryManagerResolver.getManagerForRepository(state.owner, state.repo);
+		const folderManager = this._folderRepositoryManagerResolver.getRemoteOnlyManager();
 		const identity = { owner: state.owner, repo: state.repo, number: state.number };
 		if (state.isIssue) {
 			const issueModel = await folderManager.resolveIssue(state.owner, state.repo, state.number, true, true);

src/github/pullRequestOverview.ts

@@ -63,7 +63,6 @@ export class PullRequestOverviewPanel extends IssueOverviewPanel<PullRequestMode
 	 * All open PR panels, keyed by "owner/repo#number".
 	 */
 	protected static override _panels: Map<string, PullRequestOverviewPanel> = new Map();
-	private static _repositoriesManager: RepositoriesManager | undefined;
 	private static readonly _updatingStacks = new Set<string>();
 
 	/**
@@ -129,10 +128,9 @@ export class PullRequestOverviewPanel extends IssueOverviewPanel<PullRequestMode
 
 	private postCheckoutStatus(): void {
 		if (this._item) {
-			const checkedOutPullRequestNumber = this.getCheckedOutPullRequestNumber(this._item);
 			this._postMessage({
 				command: 'pr.update-checkout-status',
-				isCurrentlyCheckedOut: checkedOutPullRequestNumber === this._item.number,
+				isCurrentlyCheckedOut: this._item.equals(this._folderRepositoryManager.activePullRequest),
 				canUpdateStack: this.canUpdateStack(this._item),
 			});
 		}
@@ -184,7 +182,7 @@ export class PullRequestOverviewPanel extends IssueOverviewPanel<PullRequestMode
 			this._panels.set(key, panel);
 		}
 
-		await panel.updateWithIdentity(folderRepositoryManager, identity, issue);
+		await panel.updateWithIdentity(identity, issue);
 		if (!panel.isDisposed && panel._item) {
 			/* __GDPR__
 				"pr.openDescription" : {
@@ -263,16 +261,15 @@ export class PullRequestOverviewPanel extends IssueOverviewPanel<PullRequestMode
 	 * and looks up the matching panel.
 	 */
 	public static registerGlobalCommands(context: vscode.ExtensionContext, telemetry: ITelemetry, repositoriesManager: RepositoriesManager): void {
-		this._repositoriesManager = repositoriesManager;
+		IssueOverviewPanel.registerRepositoriesManager(context, repositoriesManager);
 		context.subscriptions.push(
 			repositoriesManager.onDidChangeActivePullRequest(manager => {
 				for (const panel of this._panels.values()) {
-					if (panel._folderRepositoryManager !== manager) {
+					if ((panel._localFolderRepositoryManager ?? panel._initialFolderRepositoryManager) !== manager) {
 						panel.postCheckoutStatus();
 					}
 				}
 			}),
-			{ dispose: () => { this._repositoriesManager = undefined; } },
 			vscode.commands.registerCommand('pr.readyForReviewDescription', async (ctx: ReadyForReviewContext) => {
 				const panel = PullRequestOverviewPanel.findPanel(ctx.owner, ctx.repo, ctx.number);
 				if (panel) {
@@ -343,7 +340,9 @@ export class PullRequestOverviewPanel extends IssueOverviewPanel<PullRequestMode
 
 	protected override registerPrListeners() {
 		disposeAll(this._prListeners);
-		this._prListeners.push(this._folderRepositoryManager.onDidChangeActivePullRequest(() => this.postCheckoutStatus()));
+		// Avoid re-entering the upgrading accessor while rebinding its listeners.
+		const manager = this._localFolderRepositoryManager ?? this._initialFolderRepositoryManager;
+		this._prListeners.push(manager.onDidChangeActivePullRequest(() => this.postCheckoutStatus()));
 
 		if (this._item) {
 			const repository = this._item.githubRepository;
@@ -594,7 +593,7 @@ export class PullRequestOverviewPanel extends IssueOverviewPanel<PullRequestMode
 				...baseContext,
 				canUpdateStack: this.canUpdateStack(pullRequest),
 				canRequestCopilotReview: false,
-				isCurrentlyCheckedOut: this.getCheckedOutPullRequestNumber(pullRequestModel) === pullRequestModel.number,
+				isCurrentlyCheckedOut: pullRequestModel.equals(this._folderRepositoryManager.activePullRequest),
 				isRemoteBaseDeleted: pullRequest.isRemoteBaseDeleted,
 				base: `${pullRequest.base.owner}/${pullRequest.remote.repositoryName}:${pullRequest.base.ref}`,
 				isRemoteHeadDeleted: pullRequest.isRemoteHeadDeleted,
@@ -842,12 +841,11 @@ export class PullRequestOverviewPanel extends IssueOverviewPanel<PullRequestMode
 	}
 
 	public override async updateWithIdentity(
-		folderRepositoryManager: FolderRepositoryManager,
 		identity: UnresolvedIdentity,
 		pullRequestModel?: PullRequestModel | Promise<PullRequestModel>,
 		progressLocation?: string
 	): Promise<void> {
-		await super.updateWithIdentity(folderRepositoryManager, identity, pullRequestModel, progressLocation);
+		await super.updateWithIdentity(identity, pullRequestModel, progressLocation);
 
 		// Notify that this PR overview is now active
 		if (!this.isDisposed && this._item) {

src/test/github/folderRepositoryManagerResolver.test.ts

@@ -0,0 +1,97 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { default as assert } from 'assert';
+import { createSandbox, SinonSandbox } from 'sinon';
+import { Uri } from 'vscode';
+import { GitHubServerType } from '../../common/authentication';
+import { Protocol } from '../../common/protocol';
+import { GitHubRemote } from '../../common/remote';
+import { GitApiImpl } from '../../api/api1';
+import { RemoteOnlyRepository } from '../../api/remoteOnlyRepository';
+import { CredentialStore } from '../../github/credentials';
+import { FolderRepositoryManager } from '../../github/folderRepositoryManager';
+import { FolderRepositoryManagerResolver } from '../../github/folderRepositoryManagerResolver';
+import { RepositoriesManager } from '../../github/repositoriesManager';
+import { CreatePullRequestHelper } from '../../view/createPullRequestHelper';
+import { MockExtensionContext } from '../mocks/mockExtensionContext';
+import { MockRepository } from '../mocks/mockRepository';
+import { MockGitHubRepository } from '../mocks/mockGitHubRepository';
+import { MockTelemetry } from '../mocks/mockTelemetry';
+import { MockThemeWatcher } from '../mocks/mockThemeWatcher';
+
+describe('FolderRepositoryManagerResolver', function () {
+	let context: MockExtensionContext;
+	let telemetry: MockTelemetry;
+	let credentialStore: CredentialStore;
+	let repositoriesManager: RepositoriesManager;
+	let sandbox: SinonSandbox;
+
+	beforeEach(function () {
+		sandbox = createSandbox();
+		context = new MockExtensionContext();
+		telemetry = new MockTelemetry();
+		credentialStore = new CredentialStore(telemetry, context);
+		repositoriesManager = new RepositoriesManager(credentialStore, telemetry);
+	});
+
+	afterEach(function () {
+		repositoriesManager.dispose();
+		credentialStore.dispose();
+		context.dispose();
+		sandbox.restore();
+	});
+
+	function createResolver(): FolderRepositoryManagerResolver {
+		const resolver = new FolderRepositoryManagerResolver(context, repositoriesManager, telemetry);
+		context.subscriptions.push(resolver);
+		return resolver;
+	}
+
+	async function addLocalManager(url: string): Promise<FolderRepositoryManager> {
+		const repository = new MockRepository();
+		repository.rootUri = Uri.file('/workspace');
+		await repository.addRemote('origin', url);
+		const git = new GitApiImpl(repositoriesManager);
+		const helper = new CreatePullRequestHelper();
+		context.subscriptions.push(git, helper);
+		const manager = new FolderRepositoryManager(0, context, repository, telemetry, git, credentialStore, helper, new MockThemeWatcher());
+		const remote = new GitHubRemote('origin', url, new Protocol(url), GitHubServerType.GitHubDotCom);
+		const repo = new MockGitHubRepository(remote, credentialStore, telemetry, sandbox);
+		context.subscriptions.push(repo);
+		sandbox.stub(manager, 'gitHubRepositories').get(() => [repo]);
+		repositoriesManager.insertFolderManager(manager);
+		return manager;
+	}
+
+	for (const url of ['https://github.com/Owner/Repo.git', 'git@github.com:Owner/Repo.git']) {
+		it(`uses the discovered local manager for ${url}`, async function () {
+			const manager = await addLocalManager(url);
+			const resolver = createResolver();
+
+			assert.strictEqual(resolver.getManagerForRepository('owner', 'repo'), manager);
+		});
+	}
+
+	it('reuses the remote-only manager only when no local repository matches', async function () {
+		const resolver = createResolver();
+		const remoteManager = resolver.getManagerForRepository('owner', 'repo');
+		assert.ok(remoteManager.repository instanceof RemoteOnlyRepository);
+		assert.strictEqual(resolver.getManagerForRepository('other', 'repo'), remoteManager);
+		const localManager = await addLocalManager('https://github.com/owner/repo.git');
+
+		assert.strictEqual(resolver.getManagerForRepository('owner', 'repo'), localManager);
+		assert.strictEqual(resolver.getManagerForRepository('other', 'repo'), remoteManager);
+	});
+
+	it('can supply a temporary remote-only manager even when a local repository is already discovered', async function () {
+		const localManager = await addLocalManager('https://github.com/owner/repo.git');
+		const resolver = createResolver();
+
+		assert.strictEqual(resolver.getManagerForRepository('owner', 'repo'), localManager);
+		assert.ok(resolver.getRemoteOnlyManager().repository instanceof RemoteOnlyRepository);
+		assert.strictEqual(resolver.getRemoteOnlyManager(), resolver.getRemoteOnlyManager());
+	});
+});