Remove associated message handler when webview is disposed

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) => {

packages/vscode-messenger/tests/messenger.test.ts

@@ -22,6 +22,7 @@ const simpleRequest: RequestType<string, string> = { 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;