Skip to content

Commit 699258a

Browse files
refactor: move WaitForHelper to class and use it from inside context (#62)
This does not change any behavior but only refactors the code so the following changes are easier to seem
1 parent 3f38bf1 commit 699258a

File tree

8 files changed

+173
-159
lines changed

8 files changed

+173
-159
lines changed

src/McpContext.ts

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import os from 'node:os';
2020
import path from 'node:path';
2121
import {listPages} from './tools/pages.js';
2222
import {TraceResult} from './trace-processing/parse.js';
23+
import {WaitForHelper} from './WaitForHelper.js';
2324

2425
export interface TextSnapshotNode extends SerializedAXNode {
2526
id: string;
@@ -32,6 +33,9 @@ export interface TextSnapshot {
3233
snapshotId: string;
3334
}
3435

36+
const DEFAULT_TIMEOUT = 5_000;
37+
const NAVIGATION_TIMEOUT = 10_000;
38+
3539
export class McpContext implements Context {
3640
browser: Browser;
3741
logger: Debugger;
@@ -68,10 +72,10 @@ export class McpContext implements Context {
6872
this.#consoleCollector = new PageCollector(
6973
this.browser,
7074
(page, collect) => {
71-
page.on('console', (event: ConsoleMessage) => {
75+
page.on('console', event => {
7276
collect(event);
7377
});
74-
page.on('pageerror', (event: Error) => {
78+
page.on('pageerror', event => {
7579
collect(event);
7680
});
7781
},
@@ -195,10 +199,10 @@ export class McpContext implements Context {
195199
newPage.on('dialog', this.#dialogHandler);
196200

197201
// For waiters 5sec timeout should be sufficient.
198-
newPage.setDefaultTimeout(5_000);
202+
newPage.setDefaultTimeout(DEFAULT_TIMEOUT);
199203
// 10sec should be enough for the load event to be emitted during
200204
// navigations.
201-
newPage.setDefaultNavigationTimeout(10_000);
205+
newPage.setDefaultNavigationTimeout(NAVIGATION_TIMEOUT);
202206
}
203207

204208
async getElementByUid(uid: string): Promise<ElementHandle<Element>> {
@@ -302,4 +306,10 @@ export class McpContext implements Context {
302306
recordedTraces(): TraceResult[] {
303307
return this.#traceResults;
304308
}
309+
310+
waitForEventsAfterAction(action: () => Promise<unknown>): Promise<void> {
311+
const page = this.getSelectedPage();
312+
const waitForHelper = new WaitForHelper(page);
313+
return waitForHelper.waitForEventsAfterAction(action);
314+
}
305315
}

src/McpResponse.ts

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ export class McpResponse implements Response {
2121
#attachedNetworkRequestUrl?: string;
2222
#includeConsoleData: boolean = false;
2323
#textResponseLines: string[] = [];
24-
#formatedConsoleData?: string[];
24+
#formattedConsoleData?: string[];
2525
#images: ImageContentData[] = [];
2626

2727
setIncludePages(value: boolean): void {
@@ -97,7 +97,7 @@ export class McpResponse implements Response {
9797
formattedConsoleMessages = await Promise.all(
9898
consoleMessages.map(message => formatConsoleEvent(message)),
9999
);
100-
this.#formatedConsoleData = formattedConsoleMessages;
100+
this.#formattedConsoleData = formattedConsoleMessages;
101101
}
102102
}
103103

@@ -167,9 +167,9 @@ Call browser_handle_dialog to handle it before continuing.`);
167167
}
168168
}
169169

170-
if (this.#includeConsoleData && this.#formatedConsoleData) {
170+
if (this.#includeConsoleData && this.#formattedConsoleData) {
171171
response.push('## Console messages');
172-
response.push(...this.#formatedConsoleData);
172+
response.push(...this.#formattedConsoleData);
173173
}
174174

175175
const text: TextContent = {

src/WaitForHelper.ts

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
/**
2+
* @license
3+
* Copyright 2025 Google LLC
4+
* SPDX-License-Identifier: Apache-2.0
5+
*/
6+
import type {CdpPage} from 'puppeteer-core/internal/cdp/Page.js';
7+
import {Page, Protocol} from 'puppeteer-core';
8+
import {logger} from './logger.js';
9+
10+
export class WaitForHelper {
11+
#abortController = new AbortController();
12+
#genericTimeout: number;
13+
#stableDomFor: number;
14+
#expectNavigationIn: number;
15+
#page: CdpPage;
16+
17+
constructor(page: Page) {
18+
this.#genericTimeout = 3000;
19+
this.#stableDomFor = 100;
20+
this.#expectNavigationIn = 100;
21+
this.#page = page as unknown as CdpPage;
22+
}
23+
24+
/**
25+
* A wrapper that executes a action and waits for
26+
* a potential navigation, after which it waits
27+
* for the DOM to be stable before returning.
28+
*/
29+
async waitForStableDom(): Promise<void> {
30+
const stableDomObserver = await this.#page.evaluateHandle(timeout => {
31+
let timeoutId: ReturnType<typeof setTimeout>;
32+
function callback() {
33+
clearTimeout(timeoutId);
34+
timeoutId = setTimeout(() => {
35+
domObserver.resolver.resolve();
36+
domObserver.observer.disconnect();
37+
}, timeout);
38+
}
39+
const domObserver = {
40+
resolver: Promise.withResolvers<void>(),
41+
observer: new MutationObserver(callback),
42+
};
43+
// It's possible that the DOM is not gonna change so we
44+
// need to start the timeout initially.
45+
callback();
46+
47+
domObserver.observer.observe(document.body, {
48+
childList: true,
49+
subtree: true,
50+
attributes: true,
51+
});
52+
53+
return domObserver;
54+
}, this.#stableDomFor);
55+
56+
this.#abortController.signal.addEventListener('abort', async () => {
57+
try {
58+
await stableDomObserver.evaluate(observer => {
59+
observer.observer.disconnect();
60+
observer.resolver.resolve();
61+
});
62+
await stableDomObserver.dispose();
63+
} catch {
64+
// Ignored cleanup errors
65+
}
66+
});
67+
68+
return Promise.race([
69+
stableDomObserver.evaluate(async observer => {
70+
return await observer.resolver.promise;
71+
}),
72+
this.timeout(this.#genericTimeout).then(() => {
73+
throw new Error('Timeout');
74+
}),
75+
]);
76+
}
77+
78+
async waitForNavigationStarted() {
79+
// Currently Puppeteer does not have API
80+
// For when a navigation is about to start
81+
const navigationStartedPromise = new Promise<boolean>(resolve => {
82+
const listener = (event: Protocol.Page.FrameStartedNavigatingEvent) => {
83+
if (
84+
[
85+
'historySameDocument',
86+
'historyDifferentDocument',
87+
'sameDocument',
88+
].includes(event.navigationType)
89+
) {
90+
resolve(false);
91+
return;
92+
}
93+
94+
resolve(true);
95+
};
96+
97+
this.#page._client().on('Page.frameStartedNavigating', listener);
98+
this.#abortController.signal.addEventListener('abort', () => {
99+
resolve(false);
100+
this.#page._client().off('Page.frameStartedNavigating', listener);
101+
});
102+
});
103+
104+
return await Promise.race([
105+
navigationStartedPromise,
106+
this.timeout(this.#expectNavigationIn).then(() => false),
107+
]);
108+
}
109+
110+
timeout(time: number): Promise<void> {
111+
return new Promise<void>(res => {
112+
const id = setTimeout(res, time);
113+
this.#abortController.signal.addEventListener('abort', () => {
114+
res();
115+
clearTimeout(id);
116+
});
117+
});
118+
}
119+
120+
async waitForEventsAfterAction(
121+
action: () => Promise<unknown>,
122+
): Promise<void> {
123+
const navigationStartedPromise = this.waitForNavigationStarted();
124+
125+
await action();
126+
127+
try {
128+
const navigationStated = await navigationStartedPromise;
129+
if (navigationStated) {
130+
await this.#page.waitForNavigation({
131+
timeout: this.#genericTimeout,
132+
signal: this.#abortController.signal,
133+
});
134+
}
135+
136+
// Wait for stable dom after navigation so we execute in
137+
// the correct context
138+
await this.waitForStableDom();
139+
} catch (error) {
140+
logger(error);
141+
} finally {
142+
this.#abortController.abort();
143+
}
144+
}
145+
}

src/tools/ToolDefinition.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ export type Context = Readonly<{
7070
data: Uint8Array<ArrayBufferLike>,
7171
mimeType: 'image/png' | 'image/jpeg',
7272
): Promise<{filename: string}>;
73+
waitForEventsAfterAction(action: () => Promise<unknown>): Promise<void>;
7374
}>;
7475

7576
export function defineTool<Schema extends Zod.ZodRawShape>(

src/tools/input.ts

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import z from 'zod';
88
import {defineTool} from './ToolDefinition.js';
99
import {ElementHandle} from 'puppeteer-core';
1010
import {ToolCategories} from './categories.js';
11-
import {waitForEventsAfterAction} from '../waitForHelpers.js';
1211

1312
export const click = defineTool({
1413
name: 'click',
@@ -32,7 +31,7 @@ export const click = defineTool({
3231
const uid = request.params.uid;
3332
const handle = await context.getElementByUid(uid);
3433
try {
35-
await waitForEventsAfterAction(handle.frame.page(), async () => {
34+
await context.waitForEventsAfterAction(async () => {
3635
await handle.asLocator().click({
3736
count: request.params.dblClick ? 2 : 1,
3837
});
@@ -67,7 +66,7 @@ export const hover = defineTool({
6766
const uid = request.params.uid;
6867
const handle = await context.getElementByUid(uid);
6968
try {
70-
await waitForEventsAfterAction(handle.frame.page(), async () => {
69+
await context.waitForEventsAfterAction(async () => {
7170
await handle.asLocator().hover();
7271
});
7372
response.appendResponseLine(`Successfully hovered over the element`);
@@ -96,7 +95,7 @@ export const fill = defineTool({
9695
handler: async (request, response, context) => {
9796
const handle = await context.getElementByUid(request.params.uid);
9897
try {
99-
await waitForEventsAfterAction(handle.frame.page(), async () => {
98+
await context.waitForEventsAfterAction(async () => {
10099
await handle.asLocator().fill(request.params.value);
101100
});
102101
response.appendResponseLine(`Successfully filled out the element`);
@@ -122,7 +121,7 @@ export const drag = defineTool({
122121
const fromHandle = await context.getElementByUid(request.params.from_uid);
123122
const toHandle = await context.getElementByUid(request.params.to_uid);
124123
try {
125-
await waitForEventsAfterAction(fromHandle.frame.page(), async () => {
124+
await context.waitForEventsAfterAction(async () => {
126125
await fromHandle.drag(toHandle);
127126
await new Promise(resolve => setTimeout(resolve, 50));
128127
await toHandle.drop(fromHandle);
@@ -157,7 +156,7 @@ export const fillForm = defineTool({
157156
for (const element of request.params.elements) {
158157
const handle = await context.getElementByUid(element.uid);
159158
try {
160-
await waitForEventsAfterAction(handle.frame.page(), async () => {
159+
await context.waitForEventsAfterAction(async () => {
161160
await handle.asLocator().fill(element.value);
162161
});
163162
} finally {

src/tools/pages.ts

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
import z from 'zod';
88
import {defineTool} from './ToolDefinition.js';
99
import {ToolCategories} from './categories.js';
10-
import {waitForEventsAfterAction} from '../waitForHelpers.js';
1110

1211
export const listPages = defineTool({
1312
name: 'list_pages',
@@ -79,7 +78,7 @@ export const newPage = defineTool({
7978
handler: async (request, response, context) => {
8079
const page = await context.newPage();
8180

82-
await waitForEventsAfterAction(page, async () => {
81+
await context.waitForEventsAfterAction(async () => {
8382
await page.goto(request.params.url);
8483
});
8584

@@ -100,7 +99,7 @@ export const navigatePage = defineTool({
10099
handler: async (request, response, context) => {
101100
const page = context.getSelectedPage();
102101

103-
await waitForEventsAfterAction(page, async () => {
102+
await context.waitForEventsAfterAction(async () => {
104103
await page.goto(request.params.url);
105104
});
106105

src/tools/script.ts

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,7 @@
66
import z from 'zod';
77
import {defineTool} from './ToolDefinition.js';
88
import {ToolCategories} from './categories.js';
9-
import {waitForEventsAfterAction} from '../waitForHelpers.js';
10-
import {JSHandle} from 'puppeteer-core';
9+
import type {JSHandle} from 'puppeteer-core';
1110

1211
export const evaluateScript = defineTool({
1312
name: 'evaluate_script',
@@ -51,7 +50,7 @@ Example with arguments: \`(el) => {
5150
for (const el of request.params.args ?? []) {
5251
args.push(await context.getElementByUid(el.uid));
5352
}
54-
await waitForEventsAfterAction(page, async () => {
53+
await context.waitForEventsAfterAction(async () => {
5554
const result = await page.evaluate(
5655
async (fn, ...args) => {
5756
// @ts-expect-error no types.

0 commit comments

Comments
 (0)