Fix PR webview showing incorrect checkout state

src/github/pullRequestOverview.ts

@@ -558,7 +558,6 @@ export class PullRequestOverviewPanel extends IssueOverviewPanel<PullRequestMode
 			this._repositoryDefaultBranch = defaultBranch!;
 			this.setPanelTitle(this.buildPanelTitle(pullRequestModel.number, pullRequestModel.title));
 
-			const isCurrentlyCheckedOut = this.getCheckedOutPullRequestNumber(pullRequestModel) === pullRequestModel.number;
 			const mergeMethodsAvailability = repositoryAccess!.mergeMethodsAvailability;
 
 			const defaultMergeMethod = getDefaultMergeMethod(mergeMethodsAvailability);
@@ -594,7 +593,7 @@ export class PullRequestOverviewPanel extends IssueOverviewPanel<PullRequestMode
 				...baseContext,
 				canUpdateStack: this.canUpdateStack(pullRequest),
 				canRequestCopilotReview: false,
-				isCurrentlyCheckedOut: isCurrentlyCheckedOut,
+				isCurrentlyCheckedOut: this.getCheckedOutPullRequestNumber(pullRequestModel) === pullRequestModel.number,
 				isRemoteBaseDeleted: pullRequest.isRemoteBaseDeleted,
 				base: `${pullRequest.base.owner}/${pullRequest.remote.repositoryName}:${pullRequest.base.ref}`,
 				isRemoteHeadDeleted: pullRequest.isRemoteHeadDeleted,

src/test/github/pullRequestOverview.test.ts

@@ -38,6 +38,20 @@ import { PullRequest, PullRequestPreview } from '../../github/views';
 
 const EXTENSION_URI = vscode.Uri.joinPath(vscode.Uri.file(__dirname), '../../..');
 
+class TestPullRequestOverviewPanel extends PullRequestOverviewPanel {
+	constructor(telemetry: MockTelemetry, folderRepositoryManager: FolderRepositoryManager) {
+		super(telemetry, EXTENSION_URI, vscode.ViewColumn.One, '#1000', folderRepositoryManager);
+	}
+
+	public override processLinksInBodyHtml(bodyHTML: string | undefined): Promise<string | undefined> {
+		return super.processLinksInBodyHtml(bodyHTML);
+	}
+
+	public override _postMessage(message: { command: string; isCurrentlyCheckedOut?: boolean; pullrequest?: Partial<PullRequest> }): Promise<void> {
+		return super._postMessage(message);
+	}
+}
+
 describe('PullRequestOverview', function () {
 	let sinon: SinonSandbox;
 	let pullRequestManager: FolderRepositoryManager;
@@ -91,6 +105,49 @@ describe('PullRequestOverview', function () {
 		}
 	});
 
+	describe('checkout status', function () {
+		for (const initiallyCheckedOut of [false, true]) {
+			it(`initializes with the latest checkout state when ${initiallyCheckedOut ? 'leaving' : 'entering'} review mode during loading`, async function () {
+				setStacksEnabled(false);
+				const model = new PullRequestModel(credentialStore, telemetry, repo, remote,
+					convertRESTPullRequestToRawPullRequest(new PullRequestBuilder().number(1000).build(), repo));
+				const identity = { owner: remote.owner, repo: remote.repositoryName, number: model.number };
+				pullRequestManager.activePullRequest = initiallyCheckedOut ? model : undefined;
+				sinon.stub(pullRequestManager, 'getCurrentUser').resolves(model.author);
+				sinon.stub(model, 'getTimelineEvents').resolves([]);
+				sinon.stub(model, 'getReviewRequests').resolves([]);
+				sinon.stub(model, 'validateDraftMode').resolves(false);
+				sinon.stub(model, 'getStatusChecks').resolves([{ state: CheckState.Success, statuses: [] }, null]);
+				const panel = new TestPullRequestOverviewPanel(telemetry, pullRequestManager);
+				context.subscriptions.push(panel);
+				const postMessage = sinon.spy(panel, '_postMessage');
+				let releaseBody: (bodyHTML: string | undefined) => void;
+				const blockedBody = new Promise<string | undefined>(resolve => releaseBody = resolve);
+				let bodyProcessingStarted: () => void;
+				const bodyStarted = new Promise<void>(resolve => bodyProcessingStarted = resolve);
+				sinon.stub(panel, 'processLinksInBodyHtml').callsFake(() => {
+					bodyProcessingStarted();
+					return blockedBody;
+				});
+
+				const opening = panel.updateWithIdentity(pullRequestManager, identity, model);
+				await bodyStarted;
+				pullRequestManager.activePullRequest = initiallyCheckedOut ? undefined : model;
+				releaseBody!(model.bodyHTML);
+				await opening;
+
+				const calls = postMessage.getCalls();
+				const checkoutUpdate = calls.find(call => call.args[0].command === 'pr.update-checkout-status');
+				assert(checkoutUpdate);
+				assert.strictEqual(checkoutUpdate.args[0].isCurrentlyCheckedOut, !initiallyCheckedOut);
+				const initialize = calls.find(call => call.args[0].command === 'pr.initialize');
+				assert(initialize);
+				assert.strictEqual(initialize.args[0].pullrequest?.isCurrentlyCheckedOut, !initiallyCheckedOut);
+				assert(calls.indexOf(checkoutUpdate) < calls.indexOf(initialize));
+			});
+		}
+	});
+
 	describe('createOrShow', function () {
 		it('does not load stack membership when stacks are disabled', async function () {
 			setStacksEnabled(false);

webviews/editorWebview/test/app.test.tsx

@@ -139,4 +139,32 @@ describe('Root', function () {
 
 		assert(children.calledWith(pr));
 	});
+
+	it('updates the checkout button when the active pull request changes without clicking checkout', function () {
+		const pr = new PullRequestBuilder().build();
+		pr.isCurrentlyCheckedOut = true;
+		pr.doneCheckoutBranch = 'main';
+		const context = new PRContext(pr);
+		context.setPR(pr);
+		const checkout = sinon.spy(context, 'checkout');
+		const out = render(
+			<PullRequestContext.Provider value={context}>
+				<Root>{pullRequest => <Overview {...pullRequest} />}</Root>
+			</PullRequestContext.Provider>,
+		);
+
+		assert(out.getByText('Checkout \'main\''));
+		act(() => {
+			context.handleMessage({ command: 'pr.update-checkout-status', isCurrentlyCheckedOut: false });
+		});
+		assert(out.getByText('Checkout'));
+		assert.strictEqual(out.queryByText('Checkout \'main\''), null);
+
+		act(() => {
+			context.handleMessage({ command: 'pr.update-checkout-status', isCurrentlyCheckedOut: true });
+		});
+		assert(out.getByText('Checkout \'main\''));
+		assert.strictEqual(out.queryByText('Checkout'), null);
+		sinon.assert.notCalled(checkout);
+	});
 });