Release completed webview message reply callbacks

webviews/common/message.ts

@@ -23,11 +23,10 @@ export const vscode = acquireVsCodeApi();
 export class MessageHandler {
 	private _commandHandler: ((message: any) => void) | null;
 	private lastSentReq: number;
-	private pendingReplies: Record<string, { resolve: (value: any) => void; reject: (reason?: string) => void }>;
+	private readonly _pendingReplies = new Map<string, { resolve: (value: unknown) => void; reject: (reason: unknown) => void }>();
 	constructor(commandHandler: any) {
 		this._commandHandler = commandHandler;
 		this.lastSentReq = 0;
-		this.pendingReplies = Object.create(null);
 		window.addEventListener('message', this.handleMessage.bind(this) as (this: Window, ev: MessageEvent<any>) => any);
 	}
 
@@ -38,14 +37,16 @@ export class MessageHandler {
 	public async postMessage(message: any): Promise<any> {
 		const req = String(++this.lastSentReq);
 		return new Promise<any>((resolve, reject) => {
-			this.pendingReplies[req] = {
-				resolve: resolve,
-				reject: reject,
-			};
-			message = Object.assign(message, {
-				req: req,
-			});
-			vscode.postMessage(message as IRequestMessage<any>);
+			this._pendingReplies.set(req, { resolve, reject });
+			try {
+				message = Object.assign(message, {
+					req: req,
+				});
+				vscode.postMessage(message as IRequestMessage<any>);
+			} catch (error) {
+				this._pendingReplies.delete(req);
+				reject(error);
+			}
 		});
 	}
 
@@ -54,15 +55,16 @@ export class MessageHandler {
 		const message: IReplyMessage = event.data; // The json data that the extension sent
 		if (message.seq) {
 			// this is a reply
-			const pendingReply = this.pendingReplies[message.seq];
+			const pendingReply = this._pendingReplies.get(message.seq);
 			if (pendingReply) {
+				this._pendingReplies.delete(message.seq);
 				if (message.err) {
 					pendingReply.reject(message.err);
 				} else {
 					pendingReply.resolve(message.res);
 				}
-				return;
 			}
+			return;
 		}
 
 		if (this._commandHandler) {

webviews/common/test/message.test.ts

@@ -0,0 +1,47 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import assert from 'assert';
+import { createSandbox } from 'sinon';
+import { MessageHandler, vscode } from '../message';
+
+describe('MessageHandler pending replies', () => {
+	const sandbox = createSandbox();
+	afterEach(() => sandbox.restore());
+
+	it('releases successful and rejected replies while keeping outstanding requests', async () => {
+		const requests: string[] = [];
+		sandbox.stub(vscode, 'postMessage').callsFake((message: { req: string }) => {
+			requests.push(message.req);
+		});
+		const commands = sandbox.spy();
+		const handler = new MessageHandler(commands);
+		const first = handler.postMessage({ command: 'first' });
+		const second = handler.postMessage({ command: 'second' });
+		const secondAssertion = assert.rejects(second, error => error === 'Request failed');
+		assert.strictEqual(handler['_pendingReplies'].size, 2);
+
+		window.dispatchEvent(new MessageEvent('message', { data: { seq: requests[0], res: 'Result' } }));
+		assert.strictEqual(await first, 'Result');
+		assert.strictEqual(handler['_pendingReplies'].size, 1);
+		window.dispatchEvent(new MessageEvent('message', { data: { seq: requests[1], err: 'Request failed' } }));
+		await secondAssertion;
+		assert.strictEqual(handler['_pendingReplies'].size, 0);
+
+		window.dispatchEvent(new MessageEvent('message', { data: { seq: requests[0], res: 'Duplicate' } }));
+		window.dispatchEvent(new MessageEvent('message', { data: { seq: 'unknown', res: 'Unknown' } }));
+		assert.strictEqual(commands.callCount, 0);
+		window.dispatchEvent(new MessageEvent('message', { data: { res: 'Command' } }));
+		assert(commands.calledOnceWithExactly('Command'));
+	});
+
+	it('releases callbacks when sending throws synchronously', async () => {
+		const failure = new Error('Transport failed');
+		sandbox.stub(vscode, 'postMessage').throws(failure);
+		const handler = new MessageHandler(null);
+		await assert.rejects(handler.postMessage({ command: 'failure' }), error => error === failure);
+		assert.strictEqual(handler['_pendingReplies'].size, 0);
+	});
+});