refactor: dispose listeners on page destroyed

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,26 @@
  * 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;
+  // request: (event: PageEvents['request']) => 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 +33,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 +52,13 @@ export class PageCollector<T> {
       }
       this.#initializePage(page);
     });
+    this.#browser.on('targetdestroyed', async target => {
+      const page = await target.page();
+      if (!page) {
+        return;
+      }
+      this.#cleanupPageDestroyed(page);
+    });
   }
 
   public addPage(page: Page) {
@@ -50,34 +72,48 @@ 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>);
+      }
+    }
+  }
+
   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';
@@ -55,10 +56,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 +74,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 +94,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 +113,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 +137,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', {