diff --git a/packages/vscode-messenger/src/messenger.ts b/packages/vscode-messenger/src/messenger.ts index 1f24729..2543d6f 100644 --- a/packages/vscode-messenger/src/messenger.ts +++ b/packages/vscode-messenger/src/messenger.ts @@ -76,6 +76,20 @@ export class Messenger implements MessengerAPI { if (!removed) { this.log(`Attempt to remove non-existing registry entry for View: ${viewEntry.id} (type ${viewType})`, 'warn'); } + + // remove all asoociated handlers for the disposed view + for (const [key, handlers] of this.handlerRegistry.entries()) { + + const newHandlers = handlers.filter(handler => !handler.sender || + !isWebviewIdMessageParticipant(handler.sender) || + handler.sender.webviewId !== viewEntry.id); + + if (newHandlers.length === 0 ) { + this.handlerRegistry.delete(key); + } else { + this.handlerRegistry.set(key, newHandlers); + } + } }); view.webview.onDidReceiveMessage(async (msg: unknown) => { diff --git a/packages/vscode-messenger/tests/messenger.test.ts b/packages/vscode-messenger/tests/messenger.test.ts index 0339836..0503d71 100644 --- a/packages/vscode-messenger/tests/messenger.test.ts +++ b/packages/vscode-messenger/tests/messenger.test.ts @@ -22,6 +22,7 @@ const simpleRequest: RequestType = { method: 'request' }; function createWebview(viewType: string) { const view: any = { handlerTimeout: undefined, + disposeCallback: undefined, viewType, webview: { onDidReceiveMessage: (callback: (msg: unknown) => void) => { @@ -51,8 +52,14 @@ function createWebview(viewType: string) { return Promise.resolve(true); } }, - onDidDispose: () => { - view.messages = []; + onDidDispose: (callback: () => void) => { + view.disposeCallback = callback; + return { dispose: () => undefined }; + }, + dispose: () => { + if (view.disposeCallback) { + view.disposeCallback(); + } }, visible: true, messageCallback: undefined, @@ -295,6 +302,27 @@ describe('Extension Messenger', () => { expect(view1.messages[0]).toMatchObject({ id: 'fake_req_id', result: 'handled2:test' }); }); + test('Remove handlers associated with disposed webview', () => { + const messenger = new Messenger(); + const p1 = messenger.registerWebviewView(view1); + const p2 = messenger.registerWebviewView(view2); + + messenger.onNotification(simpleNotification, () => undefined, { sender: p1 }); + messenger.onNotification(simpleNotification, () => undefined, { sender: p2 }); + messenger.onNotification(simpleNotification, () => undefined); + + const beforeDispose = (messenger as any).handlerRegistry.get(simpleNotification.method); + expect(beforeDispose).toBeDefined(); + expect(beforeDispose.length).toBe(3); + + view1.dispose(); + + const afterDispose = (messenger as any).handlerRegistry.get(simpleNotification.method); + expect(afterDispose).toBeDefined(); + expect(afterDispose.length).toBe(2); + expect(afterDispose.every((registration: any) => !registration.sender || registration.sender.webviewId !== p1.webviewId)).toBe(true); + }); + test('Do not handle events for hidden view', async () => { const messenger = new Messenger(); view1.visible = false;