diff --git a/src/McpContext.ts b/src/McpContext.ts index cd829e162..4cf6b6610 100644 --- a/src/McpContext.ts +++ b/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() { diff --git a/src/PageCollector.ts b/src/PageCollector.ts index 9b078d554..acbf3239d 100644 --- a/src/PageCollector.ts +++ b/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 = { + [K in keyof EventMap]?: (event: EventMap[K]) => void; +}; export class PageCollector { #browser: Browser; - #initializer: (page: Page, collector: (item: T) => void) => void; + #listenersInitializer: ( + collector: (item: T) => void, + ) => ListenerMap; + #listeners = new WeakMap(); /** * The Array in this map should only be set once * As we use the reference to it. @@ -18,10 +32,10 @@ export class PageCollector { constructor( browser: Browser, - initializer: (page: Page, collector: (item: T) => void) => void, + listeners: (collector: (item: T) => void) => ListenerMap, ) { this.#browser = browser; - this.#initializer = initializer; + this.#listenersInitializer = listeners; } async init() { @@ -37,6 +51,14 @@ export class PageCollector { } 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,20 +72,25 @@ export class PageCollector { 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); + } + + 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 @@ -71,13 +98,23 @@ export class PageCollector { } } + #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); + } + } + this.storage.delete(page); + } + getData(page: Page): T[] { return this.storage.get(page) ?? []; } } export class NetworkCollector extends PageCollector { - override cleanup(page: Page) { + override cleanupAfterNavigation(page: Page) { const requests = this.storage.get(page) ?? []; if (!requests) { return; diff --git a/tests/PageCollector.test.ts b/tests/PageCollector.test.ts index 0e8248ce7..43c82380a 100644 --- a/tests/PageCollector.test.ts +++ b/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(res => res()); + + assert.equal(collector.getData(page).length, 0); + }); });