refactor: dispose listeners on page destroyed (#318)

Author: Lightning00Blade

src/McpContext.ts

@@ -19,6 +19,7 @@ import type {
   PredefinedNetworkConditions,
 } from 'puppeteer-core';
 
+import type {ListenerMap} from './PageCollector.js';
 import {NetworkCollector, PageCollector} from './PageCollector.js';
 import {listPages} from './tools/pages.js';
 import {takeSnapshot} from './tools/snapshot.js';
@@ -94,26 +95,24 @@ export class McpContext implements Context {
     this.browser = browser;
     this.logger = logger;
 
-    this.#networkCollector = new NetworkCollector(
-      this.browser,
-      (page, collect) => {
-        page.on('request', request => {
+    this.#networkCollector = new NetworkCollector(this.browser, collect => {
+      return {
+        request: request => {
           collect(request);
-        });
-      },
-    );
+        },
+      } as ListenerMap;
+    });
 
-    this.#consoleCollector = new PageCollector(
-      this.browser,
-      (page, collect) => {
-        page.on('console', event => {
+    this.#consoleCollector = new PageCollector(this.browser, collect => {
+      return {
+        console: event => {
           collect(event);
-        });
-        page.on('pageerror', event => {
+        },
+        pageerror: event => {
           collect(event);
-        });
-      },
-    );
+        },
+      } as ListenerMap;
+    });
   }
 
   async #init() {

src/PageCollector.ts

@@ -4,11 +4,25 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type {Browser, HTTPRequest, Page} from 'puppeteer-core';
+import {
+  type Browser,
+  type Frame,
+  type Handler,
+  type HTTPRequest,
+  type Page,
+  type PageEvents,
+} from 'puppeteer-core';
+
+export type ListenerMap<EventMap extends PageEvents = PageEvents> = {
+  [K in keyof EventMap]?: (event: EventMap[K]) => void;
+};
 
 export class PageCollector<T> {
   #browser: Browser;
-  #initializer: (page: Page, collector: (item: T) => void) => void;
+  #listenersInitializer: (
+    collector: (item: T) => void,
+  ) => ListenerMap<PageEvents>;
+  #listeners = new WeakMap<Page, ListenerMap>();
   /**
    * The Array in this map should only be set once
    * As we use the reference to it.
@@ -18,10 +32,10 @@ export class PageCollector<T> {
 
   constructor(
     browser: Browser,
-    initializer: (page: Page, collector: (item: T) => void) => void,
+    listeners: (collector: (item: T) => void) => ListenerMap<PageEvents>,
   ) {
     this.#browser = browser;
-    this.#initializer = initializer;
+    this.#listenersInitializer = listeners;
   }
 
   async init() {
@@ -37,6 +51,14 @@ export class PageCollector<T> {
       }
       this.#initializePage(page);
     });
+    this.#browser.on('targetdestroyed', async target => {
+      const page = await target.page();
+      if (!page) {
+        return;
+      }
+      console.log('destro');
+      this.#cleanupPageDestroyed(page);
+    });
   }
 
   public addPage(page: Page) {
@@ -50,34 +72,49 @@ export class PageCollector<T> {
 
     const stored: T[] = [];
     this.storage.set(page, stored);
-
-    page.on('framenavigated', frame => {
+    const listeners = this.#listenersInitializer(value => {
+      stored.push(value);
+    });
+    listeners['framenavigated'] = (frame: Frame) => {
       // Only reset the storage on main frame navigation
       if (frame !== page.mainFrame()) {
         return;
       }
-      this.cleanup(page);
-    });
-    this.#initializer(page, value => {
-      stored.push(value);
-    });
+      this.cleanupAfterNavigation(page);
+    };
+
+    for (const [name, listener] of Object.entries(listeners)) {
+      page.on(name, listener as Handler<unknown>);
+    }
+
+    this.#listeners.set(page, listeners);
   }
 
-  protected cleanup(page: Page) {
+  protected cleanupAfterNavigation(page: Page) {
     const collection = this.storage.get(page);
     if (collection) {
       // Keep the reference alive
       collection.length = 0;
     }
   }
 
+  #cleanupPageDestroyed(page: Page) {
+    const listeners = this.#listeners.get(page);
+    if (listeners) {
+      for (const [name, listener] of Object.entries(listeners)) {
+        page.off(name, listener as Handler<unknown>);
+      }
+    }
+    this.storage.delete(page);
+  }
+
   getData(page: Page): T[] {
     return this.storage.get(page) ?? [];
   }
 }
 
 export class NetworkCollector extends PageCollector<HTTPRequest> {
-  override cleanup(page: Page) {
+  override cleanupAfterNavigation(page: Page) {
     const requests = this.storage.get(page) ?? [];
     if (!requests) {
       return;

tests/PageCollector.test.ts

@@ -8,6 +8,7 @@ import {describe, it} from 'node:test';
 
 import type {Browser, Frame, Page, Target} from 'puppeteer-core';
 
+import type {ListenerMap} from '../src/PageCollector.js';
 import {PageCollector} from '../src/PageCollector.js';
 
 import {getMockRequest} from './utils.js';
@@ -22,6 +23,9 @@ function mockListener() {
         listeners[eventName] = [listener];
       }
     },
+    off(_eventName: string, _listener: (data: unknown) => void) {
+      // no-op
+    },
     emit(eventName: string, data: unknown) {
       for (const listener of listeners[eventName] ?? []) {
         listener(data);
@@ -55,10 +59,12 @@ describe('PageCollector', () => {
     const browser = getMockBrowser();
     const page = (await browser.pages())[0];
     const request = getMockRequest();
-    const collector = new PageCollector(browser, (page, collect) => {
-      page.on('request', req => {
-        collect(req);
-      });
+    const collector = new PageCollector(browser, collect => {
+      return {
+        request: req => {
+          collect(req);
+        },
+      } as ListenerMap;
     });
     await collector.init();
     page.emit('request', request);
@@ -71,10 +77,12 @@ describe('PageCollector', () => {
     const page = (await browser.pages())[0];
     const mainFrame = page.mainFrame();
     const request = getMockRequest();
-    const collector = new PageCollector(browser, (page, collect) => {
-      page.on('request', req => {
-        collect(req);
-      });
+    const collector = new PageCollector(browser, collect => {
+      return {
+        request: req => {
+          collect(req);
+        },
+      } as ListenerMap;
     });
     await collector.init();
     page.emit('request', request);
@@ -89,10 +97,12 @@ describe('PageCollector', () => {
     const browser = getMockBrowser();
     const page = (await browser.pages())[0];
     const request = getMockRequest();
-    const collector = new PageCollector(browser, (page, collect) => {
-      page.on('request', req => {
-        collect(req);
-      });
+    const collector = new PageCollector(browser, collect => {
+      return {
+        request: req => {
+          collect(req);
+        },
+      } as ListenerMap;
     });
     await collector.init();
     page.emit('request', request);
@@ -106,10 +116,12 @@ describe('PageCollector', () => {
     const page = (await browser.pages())[0];
     const mainFrame = page.mainFrame();
     const request = getMockRequest();
-    const collector = new PageCollector(browser, (page, collect) => {
-      page.on('request', req => {
-        collect(req);
-      });
+    const collector = new PageCollector(browser, collect => {
+      return {
+        request: req => {
+          collect(req);
+        },
+      } as ListenerMap;
     });
     await collector.init();
     page.emit('request', request);
@@ -128,10 +140,12 @@ describe('PageCollector', () => {
     const browser = getMockBrowser();
     const page = (await browser.pages())[0];
     const request = getMockRequest();
-    const collector = new PageCollector(browser, (pageListener, collect) => {
-      pageListener.on('request', req => {
-        collect(req);
-      });
+    const collector = new PageCollector(browser, collect => {
+      return {
+        request: req => {
+          collect(req);
+        },
+      } as ListenerMap;
     });
     await collector.init();
     browser.emit('targetcreated', {
@@ -153,4 +167,33 @@ describe('PageCollector', () => {
 
     assert.equal(collector.getData(page).length, 2);
   });
+
+  it('should clear data on page destroy', async () => {
+    const browser = getMockBrowser();
+    const page = (await browser.pages())[0];
+    const request = getMockRequest();
+    const collector = new PageCollector(browser, collect => {
+      return {
+        request: req => {
+          collect(req);
+        },
+      } as ListenerMap;
+    });
+    await collector.init();
+
+    page.emit('request', request);
+
+    assert.equal(collector.getData(page).length, 1);
+
+    browser.emit('targetdestroyed', {
+      page() {
+        return Promise.resolve(page);
+      },
+    } as Target);
+
+    // The page inside part is async so we need to await some time
+    await new Promise<void>(res => res());
+
+    assert.equal(collector.getData(page).length, 0);
+  });
 });