Skip to content

Commit c3784d1

Browse files
OrKoNOrKoN
authored
refactor: clear issue aggregator on page navigation (#565)
- fixes a memory leak introduced previously (I could not pinpoint it but McpContext was retained for subsequent tests). - adds a test that issue continue being aggregated on reload. - moves page-specific logic to a class.
1 parent 3b016f1 commit c3784d1

File tree

5 files changed

+198
-73
lines changed

5 files changed

+198
-73
lines changed

‎src/McpContext.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,11 @@ export class McpContext implements Context {
155155
await this.#consoleCollector.init();
156156
}
157157

158+
dispose() {
159+
this.#networkCollector.dispose();
160+
this.#consoleCollector.dispose();
161+
}
162+
158163
static async from(
159164
browser: Browser,
160165
logger: Debugger,

‎src/PageCollector.ts‎

Lines changed: 144 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,11 @@
44
* SPDX-License-Identifier: Apache-2.0
55
*/
66

7+
import type {
8+
AggregatedIssue,
9+
Common,
10+
} from '../node_modules/chrome-devtools-frontend/mcp/mcp.js';
711
import {
8-
type AggregatedIssue,
912
IssueAggregatorEvents,
1013
IssuesManagerEvents,
1114
createIssuesFromProtocolIssue,
@@ -14,7 +17,12 @@ import {
1417

1518
import {FakeIssuesManager} from './DevtoolsUtils.js';
1619
import {logger} from './logger.js';
17-
import type {CDPSession, ConsoleMessage} from './third_party/index.js';
20+
import type {
21+
CDPSession,
22+
ConsoleMessage,
23+
Protocol,
24+
Target,
25+
} from './third_party/index.js';
1826
import {
1927
type Browser,
2028
type Frame,
@@ -79,22 +87,31 @@ export class PageCollector<T> {
7987
this.addPage(page);
8088
}
8189

82-
this.#browser.on('targetcreated', async target => {
83-
const page = await target.page();
84-
if (!page) {
85-
return;
86-
}
87-
this.addPage(page);
88-
});
89-
this.#browser.on('targetdestroyed', async target => {
90-
const page = await target.page();
91-
if (!page) {
92-
return;
93-
}
94-
this.cleanupPageDestroyed(page);
95-
});
90+
this.#browser.on('targetcreated', this.#onTargetCreated);
91+
this.#browser.on('targetdestroyed', this.#onTargetDestroyed);
92+
}
93+
94+
dispose() {
95+
this.#browser.off('targetcreated', this.#onTargetCreated);
96+
this.#browser.off('targetdestroyed', this.#onTargetDestroyed);
9697
}
9798

99+
#onTargetCreated = async (target: Target) => {
100+
const page = await target.page();
101+
if (!page) {
102+
return;
103+
}
104+
this.addPage(page);
105+
};
106+
107+
#onTargetDestroyed = async (target: Target) => {
108+
const page = await target.page();
109+
if (!page) {
110+
return;
111+
}
112+
this.cleanupPageDestroyed(page);
113+
};
114+
98115
public addPage(page: Page) {
99116
this.#initializePage(page);
100117
}
@@ -210,68 +227,130 @@ export class PageCollector<T> {
210227
export class ConsoleCollector extends PageCollector<
211228
ConsoleMessage | Error | AggregatedIssue
212229
> {
230+
#subscribedPages = new WeakMap<Page, PageIssueSubscriber>();
231+
213232
override addPage(page: Page): void {
214-
const subscribed = this.storage.has(page);
215233
super.addPage(page);
216-
if (!subscribed) {
217-
void this.subscribeForIssues(page);
234+
if (!this.#subscribedPages.has(page)) {
235+
const subscriber = new PageIssueSubscriber(page);
236+
this.#subscribedPages.set(page, subscriber);
237+
void subscriber.subscribe();
218238
}
219239
}
220240

221-
async subscribeForIssues(page: Page) {
222-
const seenKeys = new Set<string>();
223-
const mockManager = new FakeIssuesManager();
224-
const aggregator = new IssueAggregator(mockManager);
225-
aggregator.addEventListener(
241+
protected override cleanupPageDestroyed(page: Page): void {
242+
super.cleanupPageDestroyed(page);
243+
this.#subscribedPages.get(page)?.unsubscribe();
244+
this.#subscribedPages.delete(page);
245+
}
246+
}
247+
248+
class PageIssueSubscriber {
249+
#issueManager = new FakeIssuesManager();
250+
#issueAggregator = new IssueAggregator(this.#issueManager);
251+
#seenKeys = new Set<string>();
252+
#seenIssues = new Set<AggregatedIssue>();
253+
#page: Page;
254+
#session: CDPSession;
255+
256+
constructor(page: Page) {
257+
this.#page = page;
258+
// @ts-expect-error use existing CDP client (internal Puppeteer API).
259+
this.#session = this.#page._client() as CDPSession;
260+
}
261+
262+
#resetIssueAggregator() {
263+
this.#issueManager = new FakeIssuesManager();
264+
if (this.#issueAggregator) {
265+
this.#issueAggregator.removeEventListener(
266+
IssueAggregatorEvents.AGGREGATED_ISSUE_UPDATED,
267+
this.#onAggregatedissue,
268+
);
269+
}
270+
this.#issueAggregator = new IssueAggregator(this.#issueManager);
271+
272+
this.#issueAggregator.addEventListener(
226273
IssueAggregatorEvents.AGGREGATED_ISSUE_UPDATED,
227-
event => {
228-
const withId = event.data as WithSymbolId<AggregatedIssue>;
229-
// Emit aggregated issue only if it's a new one
230-
if (withId[stableIdSymbol]) {
231-
return;
232-
}
233-
page.emit('issue', event.data);
234-
},
274+
this.#onAggregatedissue,
235275
);
276+
}
236277

278+
async subscribe() {
279+
this.#resetIssueAggregator();
280+
this.#page.on('framenavigated', this.#onFrameNavigated);
281+
this.#session.on('Audits.issueAdded', this.#onIssueAdded);
237282
try {
238-
// @ts-expect-error use existing CDP client (internal Puppeteer API).
239-
const session = page._client() as CDPSession;
240-
session.on('Audits.issueAdded', data => {
241-
try {
242-
const inspectorIssue = data.issue;
243-
// @ts-expect-error Types of protocol from Puppeteer and CDP are
244-
// incomparable for InspectorIssueCode, one is union, other is enum.
245-
const issue = createIssuesFromProtocolIssue(null, inspectorIssue)[0];
246-
if (!issue) {
247-
logger('No issue mapping for for the issue: ', inspectorIssue.code);
248-
return;
249-
}
250-
251-
const primaryKey = issue.primaryKey();
252-
if (seenKeys.has(primaryKey)) {
253-
return;
254-
}
255-
seenKeys.add(primaryKey);
256-
257-
mockManager.dispatchEventToListeners(
258-
IssuesManagerEvents.ISSUE_ADDED,
259-
{
260-
issue,
261-
// @ts-expect-error We don't care that issues model is null
262-
issuesModel: null,
263-
},
264-
);
265-
} catch (error) {
266-
logger('Error creating a new issue', error);
267-
}
268-
});
269-
270-
await session.send('Audits.enable');
283+
await this.#session.send('Audits.enable');
271284
} catch (error) {
272285
logger('Error subscribing to issues', error);
273286
}
274287
}
288+
289+
unsubscribe() {
290+
this.#seenKeys.clear();
291+
this.#seenIssues.clear();
292+
this.#page.off('framenavigated', this.#onFrameNavigated);
293+
this.#session.off('Audits.issueAdded', this.#onIssueAdded);
294+
if (this.#issueAggregator) {
295+
this.#issueAggregator.removeEventListener(
296+
IssueAggregatorEvents.AGGREGATED_ISSUE_UPDATED,
297+
this.#onAggregatedissue,
298+
);
299+
}
300+
void this.#session.send('Audits.disable').catch(() => {
301+
// might fail.
302+
});
303+
}
304+
305+
#onAggregatedissue = (
306+
event: Common.EventTarget.EventTargetEvent<AggregatedIssue>,
307+
) => {
308+
if (this.#seenIssues.has(event.data)) {
309+
return;
310+
}
311+
this.#seenIssues.add(event.data);
312+
this.#page.emit('issue', event.data);
313+
};
314+
315+
// On navigation, we reset issue aggregation.
316+
#onFrameNavigated = (frame: Frame) => {
317+
// Only split the storage on main frame navigation
318+
if (frame !== frame.page().mainFrame()) {
319+
return;
320+
}
321+
this.#seenKeys.clear();
322+
this.#seenIssues.clear();
323+
this.#resetIssueAggregator();
324+
};
325+
326+
#onIssueAdded = (data: Protocol.Audits.IssueAddedEvent) => {
327+
try {
328+
const inspectorIssue = data.issue;
329+
// @ts-expect-error Types of protocol from Puppeteer and CDP are
330+
// incomparable for InspectorIssueCode, one is union, other is enum.
331+
const issue = createIssuesFromProtocolIssue(null, inspectorIssue)[0];
332+
if (!issue) {
333+
logger('No issue mapping for for the issue: ', inspectorIssue.code);
334+
return;
335+
}
336+
337+
const primaryKey = issue.primaryKey();
338+
if (this.#seenKeys.has(primaryKey)) {
339+
return;
340+
}
341+
this.#seenKeys.add(primaryKey);
342+
this.#issueManager.dispatchEventToListeners(
343+
IssuesManagerEvents.ISSUE_ADDED,
344+
{
345+
issue,
346+
// @ts-expect-error We don't care that issues model is null
347+
issuesModel: null,
348+
},
349+
);
350+
} catch (error) {
351+
logger('Error creating a new issue', error);
352+
}
353+
};
275354
}
276355

277356
export class NetworkCollector extends PageCollector<HTTPRequest> {

‎tests/PageCollector.test.ts‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ import type {
1212
HTTPRequest,
1313
Page,
1414
Target,
15-
CDPSession,
1615
Protocol,
1716
} from 'puppeteer-core';
1817
import sinon from 'sinon';
@@ -60,9 +59,6 @@ function getMockPage(): Page {
6059
mainFrame() {
6160
return mainFrame;
6261
},
63-
createCDPSession() {
64-
return Promise.resolve(cdpSession as unknown as CDPSession);
65-
},
6662
...mockListener(),
6763
// @ts-expect-error internal API.
6864
_client() {

‎tests/tools/console.test.ts‎

Lines changed: 44 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,15 +41,14 @@ describe('console', () => {
4141
});
4242
});
4343

44-
it('lists issues messages', async () => {
44+
it('lists issues', async () => {
4545
await withBrowser(async (response, context) => {
4646
const page = await context.newPage();
4747
const issuePromise = new Promise<void>(resolve => {
48-
page.on('issue', () => {
48+
page.once('issue', () => {
4949
resolve();
5050
});
5151
});
52-
5352
await page.setContent('<input type="text" name="username" />');
5453
await issuePromise;
5554
await listConsoleMessages.handler({params: {}}, response, context);
@@ -63,6 +62,48 @@ describe('console', () => {
6362
});
6463
});
6564

65+
it('lists issues after a page reload', async () => {
66+
await withBrowser(async (response, context) => {
67+
const page = await context.newPage();
68+
const issuePromise = new Promise<void>(resolve => {
69+
page.once('issue', () => {
70+
resolve();
71+
});
72+
});
73+
74+
await page.setContent('<input type="text" name="username" />');
75+
await issuePromise;
76+
await listConsoleMessages.handler({params: {}}, response, context);
77+
{
78+
const formattedResponse = await response.handle('test', context);
79+
const textContent = formattedResponse[0] as {text: string};
80+
assert.ok(
81+
textContent.text.includes(
82+
`msgid=1 [issue] An element doesn't have an autocomplete attribute (count: 1)`,
83+
),
84+
);
85+
}
86+
87+
const anotherIssuePromise = new Promise<void>(resolve => {
88+
page.once('issue', () => {
89+
resolve();
90+
});
91+
});
92+
await page.reload();
93+
await page.setContent('<input type="text" name="username" />');
94+
await anotherIssuePromise;
95+
{
96+
const formattedResponse = await response.handle('test', context);
97+
const textContent = formattedResponse[0] as {text: string};
98+
assert.ok(
99+
textContent.text.includes(
100+
`msgid=2 [issue] An element doesn't have an autocomplete attribute (count: 1)`,
101+
),
102+
);
103+
}
104+
});
105+
});
106+
66107
it('work with primitive unhandled errors', async () => {
67108
await withBrowser(async (response, context) => {
68109
const page = await context.newPage();

‎tests/utils.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import {McpResponse} from '../src/McpResponse.js';
1818
import {stableIdSymbol} from '../src/PageCollector.js';
1919

2020
const browsers = new Map<string, Browser>();
21+
let context: McpContext | undefined;
2122

2223
export async function withBrowser(
2324
cb: (response: McpResponse, context: McpContext) => Promise<void>,
@@ -48,7 +49,10 @@ export async function withBrowser(
4849
}),
4950
);
5051
const response = new McpResponse();
51-
const context = await McpContext.from(
52+
if (context) {
53+
context.dispose();
54+
}
55+
context = await McpContext.from(
5256
browser,
5357
logger('test'),
5458
{

0 commit comments

Comments
 (0)