Refactor webview contexts for explicit host ownership

src/test/webviews/testHost.ts

@@ -0,0 +1,25 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { afterEach } from 'mocha';
+import { createWebviewHost, WebviewHost, WebviewTransport } from '../../../webviews/common/host';
+
+const hosts = new Set<WebviewHost>();
+
+export function createTestHost(state?: unknown, transport?: WebviewTransport): WebviewHost {
+	const host = createWebviewHost(transport);
+	if (state !== undefined) {
+		host.setState(state);
+	}
+	hosts.add(host);
+	return host;
+}
+
+afterEach(() => {
+	for (const host of hosts) {
+		host.dispose();
+	}
+	hosts.clear();
+});

webviews/activityBarView/app.tsx

@@ -7,10 +7,19 @@ import React, { useContext, useEffect, useState } from 'react';
 import { render } from 'react-dom';
 import { Overview } from './overview';
 import { PullRequest } from '../../src/github/views';
-import PullRequestContext from '../common/context';
+import PullRequestContext, { PRContext } from '../common/context';
+import { createWebviewHost } from '../common/host';
+import { rethrowUnlessDisposed } from '../common/message';
 
 export function main() {
-	render(<Root>{pr => <Overview {...pr} />}</Root>, document.getElementById('app'));
+	const host = createWebviewHost();
+	const context = new PRContext(host);
+	render(
+		<PullRequestContext.Provider value={context}>
+			<Root>{pr => <Overview {...pr} />}</Root>
+		</PullRequestContext.Provider>,
+		document.getElementById('app')
+	);
 }
 
 export function Root({ children }) {
@@ -19,8 +28,11 @@ export function Root({ children }) {
 	useEffect(() => {
 		ctx.onchange = setPR;
 		setPR(ctx.pr);
+		return () => {
+			ctx.onchange = null;
+		};
 	}, []);
-	ctx.postMessage({ command: 'ready' });
-	ctx.postMessage({ command: 'pr.debug', args: 'initialized ' + (pr ? 'with PR' : 'without PR') });
+	ctx.postMessage({ command: 'ready' }).catch(rethrowUnlessDisposed);
+	ctx.postMessage({ command: 'pr.debug', args: 'initialized ' + (pr ? 'with PR' : 'without PR') }).catch(rethrowUnlessDisposed);
 	return pr ? children(pr) : <div className="loading-indicator">Loading...</div>;
 }

webviews/common/context.tsx

@@ -4,9 +4,8 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { createContext } from 'react';
-import { getState, setState, updateState } from './cache';
 import { COMMENT_TEXTAREA_ID } from './constants';
-import { getMessageHandler, MessageHandler } from './message';
+import { WebviewHost } from './host';
 import { CloseResult, DescriptionResult, OpenCommitChangesArgs, OpenLocalFileArgs } from '../../common/views';
 import { IComment } from '../../src/common/comment';
 import { EventType, ReviewEvent, SessionLinkInfo, TimelineEvent } from '../../src/common/timelineEvent';
@@ -32,17 +31,21 @@ function bytesToBase64(bytes: Uint8Array): string {
 const MAX_UPLOAD_SIZE_BYTES = 25 * 1024 * 1024;
 
 export class PRContext {
+	public pr: PullRequest | undefined;
+	public onchange: ((ctx: PullRequest | undefined) => void) | null = null;
 	public preview: OverviewItemPreview | undefined;
 	public onPreviewChange: ((preview: OverviewItemPreview | undefined) => void) | null = null;
+	private readonly _unsubscribe: () => void;
 
-	constructor(
-		public pr: PullRequest | undefined = getState(),
-		public onchange: ((ctx: PullRequest | undefined) => void) | null = null,
-		private _handler: MessageHandler | null = null,
-	) {
-		if (!_handler) {
-			this._handler = getMessageHandler(this.handleMessage);
-		}
+	constructor(private readonly _host: WebviewHost) {
+		this.pr = _host.getState<PullRequest>();
+		this._unsubscribe = _host.onCommand(this.handleMessage);
+	}
+
+	public dispose(): void {
+		this._unsubscribe();
+		this.onchange = null;
+		this.onPreviewChange = null;
 	}
 
 	public setTitle = async (title: string) => {
@@ -194,7 +197,7 @@ export class PRContext {
 		this.postMessage({ command: 'pr.cancel-generate-description' });
 
 	public updateDraft = (id: number, body: string) => {
-		const pullRequest = getState();
+		const pullRequest = this._host.getState<PullRequest>() ?? this.pr;
 		const pendingCommentDrafts = pullRequest.pendingCommentDrafts || Object.create(null);
 		if (body === pendingCommentDrafts[id]) {
 			return;
@@ -550,24 +553,32 @@ export class PRContext {
 		this.preview = undefined;
 		this.onPreviewChange?.(undefined);
 		this.pr = pr;
-		setState(this.pr);
+		const oldPullRequest = this._host.getState<PullRequest>();
+		if (oldPullRequest?.number && oldPullRequest.number === pr?.number) {
+			pr.pendingCommentText = oldPullRequest.pendingCommentText;
+			pr.pendingReviewSummaryText = oldPullRequest.pendingReviewSummaryText;
+		}
+		if (pr) {
+			this._host.setState(pr);
+		}
 		if (this.onchange) {
 			this.onchange(this.pr);
 		}
 		return this;
 	};
 
 	updatePR = (pr: Partial<PullRequest> | undefined) => {
-		updateState(pr);
+		const current = this._host.getState<PullRequest>() ?? this.pr;
+		this._host.setState(Object.assign(current ?? {}, pr));
 		this.pr = this.pr ? { ...this.pr, ...pr } : pr as PullRequest;
 		if (this.onchange) {
 			this.onchange(this.pr);
 		}
 		return this;
 	};
 
-	postMessage(message: any) {
-		return (this._handler?.postMessage(message) ?? Promise.resolve(undefined));
+	async postMessage(message: any): Promise<any> {
+		return this._host.postMessage(message);
 	}
 
 	handleMessage = (message: any) => {
@@ -627,8 +638,8 @@ export class PRContext {
 		}
 	};
 
-	public static instance = new PRContext();
 }
 
-const PullRequestContext = createContext<PRContext>(PRContext.instance);
+// Every render entry point must supply a provider.
+const PullRequestContext = createContext<PRContext>(undefined!);
 export default PullRequestContext;

webviews/common/createContextNew.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { createContext } from 'react';
-import { getMessageHandler, MessageHandler, vscode } from './message';
+import { WebviewHost } from './host';
 import { RemoteInfo } from '../../common/types';
 import { CancelCreatePullRequestNew, ChooseBaseRemoteAndBranchResult, ChooseCompareRemoteAndBranchResult, ChooseRemoteAndBranchArgs, CreateParamsNew, CreatePullRequestNew, ScrollPosition, StackCandidate, TitleAndDescriptionArgs, TitleAndDescriptionResult } from '../../common/views';
 import { compareIgnoreCase } from '../../src/common/utils';
@@ -39,6 +39,10 @@ const defaultCreateParams: CreateParamsNew = {
 	usingTemplate: false
 };
 
+function createDefaultParams(): CreateParamsNew {
+	return { ...defaultCreateParams, labels: [], assignees: [], reviewers: [] };
+}
+
 function remoteChanged(next: RemoteInfo | undefined, current: RemoteInfo | undefined): boolean {
 	return next !== undefined && (next.owner !== current?.owner || next.repositoryName !== current?.repositoryName);
 }
@@ -53,17 +57,19 @@ function stackCandidateChanged(next: Partial<CreateParamsNew>, current: StackCan
 
 export class CreatePRContextNew {
 	public createParams: CreateParamsNew;
+	public onchange: ((ctx: CreateParamsNew) => void) | null = null;
 	private _titleStack: string[] = [];
 	private _descriptionStack: string[] = [];
+	private readonly _unsubscribe: () => void;
 
-	constructor(
-		public onchange: ((ctx: CreateParamsNew) => void) | null = null,
-		private _handler: MessageHandler | null = null,
-	) {
-		this.createParams = vscode.getState() ?? defaultCreateParams;
-		if (!_handler) {
-			this._handler = getMessageHandler(this.handleMessage);
-		}
+	constructor(private readonly _host: WebviewHost) {
+		this.createParams = _host.getState<CreateParamsNew>() ?? createDefaultParams();
+		this._unsubscribe = _host.onCommand(this.handleMessage);
+	}
+
+	public dispose(): void {
+		this._unsubscribe();
+		this.onchange = null;
 	}
 
 	get isCreatable(): boolean {
@@ -116,7 +122,7 @@ export class CreatePRContextNew {
 		// dialog, or the message did not get a response) preserve the user's
 		// in-progress title/description.
 		if (result?.cancelled === true) {
-			vscode.setState(defaultCreateParams);
+			this._host.setState(createDefaultParams());
 		}
 	};
 
@@ -133,8 +139,8 @@ export class CreatePRContextNew {
 	}
 
 	public updateState = (params: Partial<CreateParamsNew>, reset: boolean = false): void => {
-		this.createParams = reset ? { ...defaultCreateParams, ...params } : { ...this.createParams, ...params };
-		vscode.setState(this.createParams);
+		this.createParams = reset ? { ...createDefaultParams(), ...params } : { ...this.createParams, ...params };
+		this._host.setState(this.createParams);
 		if (this.onchange) {
 			this.onchange(this.createParams);
 		}
@@ -324,14 +330,14 @@ export class CreatePRContextNew {
 			// awaiting would discard the user's title/description if the create fails
 			// (for example when the branch push fails and needs `--force`), causing
 			// the description to fall back to the template on the next webview load.
-			vscode.setState(defaultCreateParams);
+			this._host.setState(createDefaultParams());
 		} catch (e) {
 			this.updateState({ createError: (typeof e === 'string') ? e : (e.message ? e.message : 'An unknown error occurred.') });
 		}
 	};
 
 	postMessage = async (message: any): Promise<any> => {
-		return this._handler?.postMessage(message);
+		return this._host.postMessage(message);
 	};
 
 	handleMessage = async (message: { command: string, params?: Partial<CreateParamsNew>, scrollPosition?: ScrollPosition }): Promise<void> => {
@@ -396,7 +402,7 @@ export class CreatePRContextNew {
 
 			case 'reset':
 				if (!message.params) {
-					this.updateState(defaultCreateParams, true);
+					this.updateState(createDefaultParams(), true);
 					return;
 				}
 				message.params.creating = message.params.creating ?? false;
@@ -453,8 +459,8 @@ export class CreatePRContextNew {
 		}
 	};
 
-	public static instance = new CreatePRContextNew();
 }
 
-const PullRequestContextNew = createContext<CreatePRContextNew>(CreatePRContextNew.instance);
+// Every render entry point must supply a provider.
+const PullRequestContextNew = createContext<CreatePRContextNew>(undefined!);
 export default PullRequestContextNew;

webviews/common/host.ts

@@ -0,0 +1,61 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { MessageHandler } from './message';
+
+export interface WebviewHost extends WebviewState {
+	postMessage(message: CommandMessage): Promise<unknown>;
+	onCommand(listener: (command: unknown) => void): () => void;
+	dispose(): void;
+}
+
+export function createWebviewHost(transport: WebviewTransport = vscodeTransport): WebviewHost {
+	return new MessageHandler(transport);
+}
+
+export interface CommandMessage {
+	command: string;
+	args?: unknown;
+}
+
+export interface RequestMessage extends CommandMessage {
+	req: string;
+}
+
+export interface ReplyMessage {
+	seq?: string;
+	err?: string;
+	res?: unknown;
+}
+
+export interface WebviewState {
+	getState<T>(): T | undefined;
+	setState<T>(state: T): void;
+}
+
+export interface WebviewTransport extends WebviewState {
+	postMessage(message: RequestMessage): void;
+	onMessage(listener: (message: ReplyMessage) => void): () => void;
+}
+
+declare function acquireVsCodeApi(): Omit<WebviewTransport, 'onMessage'>;
+
+let api: Omit<WebviewTransport, 'onMessage'> | undefined;
+
+function getApi(): Omit<WebviewTransport, 'onMessage'> {
+	api ??= acquireVsCodeApi();
+	return api;
+}
+
+export const vscodeTransport: WebviewTransport = {
+	getState: <T>() => getApi().getState<T>(),
+	setState: <T>(state: T) => getApi().setState(state),
+	postMessage: message => getApi().postMessage(message),
+	onMessage: listener => {
+		const handler = (event: MessageEvent<ReplyMessage>) => listener(event.data);
+		window.addEventListener('message', handler);
+		return () => window.removeEventListener('message', handler);
+	},
+};