Skip to content

Commit f4bf143

Browse files
OrKoNOrKoN
committed
refactor: clear issue aggregator on page navigation
1 parent 3b016f1 commit f4bf143

File tree

5 files changed

+194
-73
lines changed

5 files changed

+194
-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: 140 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@
44
* SPDX-License-Identifier: Apache-2.0
55
*/
66

7+
import type {AggregatedIssue} from '../node_modules/chrome-devtools-frontend/mcp/mcp.js';
78
import {
8-
type AggregatedIssue,
99
IssueAggregatorEvents,
1010
IssuesManagerEvents,
1111
createIssuesFromProtocolIssue,
@@ -14,7 +14,12 @@ import {
1414

1515
import {FakeIssuesManager} from './DevtoolsUtils.js';
1616
import {logger} from './logger.js';
17-
import type {CDPSession, ConsoleMessage} from './third_party/index.js';
17+
import type {
18+
CDPSession,
19+
ConsoleMessage,
20+
Protocol,
21+
Target,
22+
} from './third_party/index.js';
1823
import {
1924
type Browser,
2025
type Frame,
@@ -79,22 +84,31 @@ export class PageCollector<T> {
7984
this.addPage(page);
8085
}
8186

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-
});
87+
this.#browser.on('targetcreated', this.#onTargetCreated);
88+
this.#browser.on('targetdestroyed', this.#onTargetDestroyed);
89+
}
90+
91+
dispose() {
92+
this.#browser.off('targetcreated', this.#onTargetCreated);
93+
this.#browser.off('targetdestroyed', this.#onTargetDestroyed);
9694
}
9795

96+
#onTargetCreated = async (target: Target) => {
97+
const page = await target.page();
98+
if (!page) {
99+
return;
100+
}
101+
this.addPage(page);
102+
};
103+
104+
#onTargetDestroyed = async (target: Target) => {
105+
const page = await target.page();
106+
if (!page) {
107+
return;
108+
}
109+
this.cleanupPageDestroyed(page);
110+
};
111+
98112
public addPage(page: Page) {
99113
this.#initializePage(page);
100114
}
@@ -210,68 +224,129 @@ export class PageCollector<T> {
210224
export class ConsoleCollector extends PageCollector<
211225
ConsoleMessage | Error | AggregatedIssue
212226
> {
227+
#subscribedPages = new WeakMap<Page, PageIssueSubscriber>();
228+
213229
override addPage(page: Page): void {
214-
const subscribed = this.storage.has(page);
215230
super.addPage(page);
216-
if (!subscribed) {
217-
void this.subscribeForIssues(page);
231+
if (!this.#subscribedPages.has(page)) {
232+
const subscriber = new PageIssueSubscriber(page);
233+
this.#subscribedPages.set(page, subscriber);
234+
void subscriber.subscribe();
218235
}
219236
}
237+
238+
protected override cleanupPageDestroyed(page: Page): void {
239+
super.cleanupPageDestroyed(page);
240+
this.#subscribedPages.get(page)?.unsubscribe();
241+
this.#subscribedPages.delete(page);
242+
}
243+
}
244+
245+
class PageIssueSubscriber {
246+
#issueManager = new FakeIssuesManager();
247+
#issueAggregator = new IssueAggregator(this.#issueManager);
248+
#seenKeys = new Set<string>();
249+
#seenIssues = new Set<AggregatedIssue>();
250+
#page: Page;
251+
#session: CDPSession;
252+
253+
constructor(page: Page) {
254+
this.#page = page;
255+
// @ts-expect-error use existing CDP client (internal Puppeteer API).
256+
this.#session = this.#page._client() as CDPSession;
257+
}
258+
259+
#resetIssueAggregator() {
260+
this.#issueManager = new FakeIssuesManager();
261+
if (this.#issueAggregator) {
262+
this.#issueAggregator.removeEventListener(
263+
IssueAggregatorEvents.AGGREGATED_ISSUE_UPDATED,
264+
this.#onAggregatedissue,
265+
);
266+
}
267+
this.#issueAggregator = new IssueAggregator(this.#issueManager);
220268

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(
269+
this.#issueAggregator.addEventListener(
226270
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-
},
271+
this.#onAggregatedissue,
235272
);
273+
}
236274

275+
async subscribe() {
276+
this.#resetIssueAggregator();
277+
this.#page.on('framenavigated', this.#onFrameNavigated);
278+
this.#session.on('Audits.issueAdded', this.#onIssueAdded);
237279
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');
280+
await this.#session.send('Audits.enable');
271281
} catch (error) {
272282
logger('Error subscribing to issues', error);
273283
}
274284
}
285+
286+
unsubscribe() {
287+
this.#seenKeys.clear();
288+
this.#seenIssues.clear();
289+
this.#page.off('framenavigated', this.#onFrameNavigated);
290+
this.#session.off('Audits.issueAdded', this.#onIssueAdded);
291+
if (this.#issueAggregator) {
292+
this.#issueAggregator.removeEventListener(
293+
IssueAggregatorEvents.AGGREGATED_ISSUE_UPDATED,
294+
this.#onAggregatedissue,
295+
);
296+
}
297+
void this.#session.send('Audits.disable').catch(() => {
298+
// might fail.
299+
});
300+
}
301+
302+
// TODO: expose DevTools types.
303+
#onAggregatedissue = (event: any) => {
304+
if (this.#seenIssues.has(event.data)) {
305+
return;
306+
}
307+
this.#seenIssues.add(event.data);
308+
this.#page.emit('issue', event.data);
309+
};
310+
311+
// On navigation, we reset issue aggregation.
312+
#onFrameNavigated = (frame: Frame) => {
313+
// Only split the storage on main frame navigation
314+
if (frame !== frame.page().mainFrame()) {
315+
return;
316+
}
317+
this.#seenKeys.clear();
318+
this.#seenIssues.clear();
319+
this.#resetIssueAggregator();
320+
};
321+
322+
#onIssueAdded = (data: Protocol.Audits.IssueAddedEvent) => {
323+
try {
324+
const inspectorIssue = data.issue;
325+
// @ts-expect-error Types of protocol from Puppeteer and CDP are
326+
// incomparable for InspectorIssueCode, one is union, other is enum.
327+
const issue = createIssuesFromProtocolIssue(null, inspectorIssue)[0];
328+
if (!issue) {
329+
logger('No issue mapping for for the issue: ', inspectorIssue.code);
330+
return;
331+
}
332+
333+
const primaryKey = issue.primaryKey();
334+
if (this.#seenKeys.has(primaryKey)) {
335+
return;
336+
}
337+
this.#seenKeys.add(primaryKey);
338+
this.#issueManager.dispatchEventToListeners(
339+
IssuesManagerEvents.ISSUE_ADDED,
340+
{
341+
issue,
342+
// @ts-expect-error We don't care that issues model is null
343+
issuesModel: null,
344+
},
345+
);
346+
} catch (error) {
347+
logger('Error creating a new issue', error);
348+
}
349+
};
275350
}
276351

277352
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)