Skip to content

Commit d5c70be

Browse files
OrKoNOrKoN
committed
refactor: move emulation settings to context
1 parent 0c4f211 commit d5c70be

7 files changed

Lines changed: 125 additions & 162 deletions

File tree

‎src/McpContext.ts‎

Lines changed: 93 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,9 @@ import type {
2727
HTTPRequest,
2828
Page,
2929
SerializedAXNode,
30-
PredefinedNetworkConditions,
3130
Viewport,
3231
} from './third_party/index.js';
32+
import {PredefinedNetworkConditions} from './third_party/index.js';
3333
import {listPages} from './tools/pages.js';
3434
import {takeSnapshot} from './tools/snapshot.js';
3535
import {CLOSE_PAGE_ERROR} from './tools/ToolDefinition.js';
@@ -282,83 +282,127 @@ export class McpContext implements Context {
282282
return this.#networkCollector.getById(this.getSelectedPage(), reqid);
283283
}
284284

285-
setNetworkConditions(conditions: string | null): void {
285+
async emulate(options: {
286+
networkConditions?: string | null;
287+
cpuThrottlingRate?: number | null;
288+
geolocation?: GeolocationOptions | null;
289+
userAgent?: string | null;
290+
colorScheme?: 'dark' | 'light' | 'auto' | null;
291+
viewport?: Viewport | null;
292+
}): Promise<void> {
286293
const page = this.getSelectedPage();
287-
if (conditions === null) {
288-
this.#networkConditionsMap.delete(page);
289-
} else {
290-
this.#networkConditionsMap.set(page, conditions);
294+
let timeoutsNeedUpdate = false;
295+
296+
if (options.networkConditions !== undefined) {
297+
if (options.networkConditions === null || options.networkConditions === 'No emulation') {
298+
await page.emulateNetworkConditions(null);
299+
this.#networkConditionsMap.delete(page);
300+
} else if (options.networkConditions === 'Offline') {
301+
await page.emulateNetworkConditions({
302+
offline: true,
303+
download: 0,
304+
upload: 0,
305+
latency: 0,
306+
});
307+
this.#networkConditionsMap.set(page, 'Offline');
308+
} else if (options.networkConditions in PredefinedNetworkConditions) {
309+
const networkCondition =
310+
PredefinedNetworkConditions[
311+
options.networkConditions as keyof typeof PredefinedNetworkConditions
312+
];
313+
await page.emulateNetworkConditions(networkCondition);
314+
this.#networkConditionsMap.set(page, options.networkConditions);
315+
}
316+
timeoutsNeedUpdate = true;
317+
}
318+
319+
if (options.cpuThrottlingRate !== undefined) {
320+
if (options.cpuThrottlingRate === null) {
321+
await page.emulateCPUThrottling(1);
322+
this.#cpuThrottlingRateMap.delete(page);
323+
} else {
324+
await page.emulateCPUThrottling(options.cpuThrottlingRate);
325+
this.#cpuThrottlingRateMap.set(page, options.cpuThrottlingRate);
326+
}
327+
timeoutsNeedUpdate = true;
328+
}
329+
330+
if (options.geolocation !== undefined) {
331+
if (options.geolocation === null) {
332+
await page.setGeolocation({latitude: 0, longitude: 0});
333+
this.#geolocationMap.delete(page);
334+
} else {
335+
await page.setGeolocation(options.geolocation);
336+
this.#geolocationMap.set(page, options.geolocation);
337+
}
338+
}
339+
340+
if (options.userAgent !== undefined) {
341+
if (options.userAgent === null) {
342+
await page.setUserAgent({userAgent: undefined});
343+
this.#userAgentMap.delete(page);
344+
} else {
345+
await page.setUserAgent({userAgent: options.userAgent});
346+
this.#userAgentMap.set(page, options.userAgent);
347+
}
348+
}
349+
350+
if (options.colorScheme !== undefined) {
351+
if (options.colorScheme === null || options.colorScheme === 'auto') {
352+
await page.emulateMediaFeatures([{name: 'prefers-color-scheme', value: ''}]);
353+
this.#colorSchemeMap.delete(page);
354+
} else {
355+
await page.emulateMediaFeatures([{name: 'prefers-color-scheme', value: options.colorScheme}]);
356+
this.#colorSchemeMap.set(page, options.colorScheme);
357+
}
358+
}
359+
360+
if (options.viewport !== undefined) {
361+
if (options.viewport === null) {
362+
await page.setViewport(null);
363+
this.#viewportMap.delete(page);
364+
} else {
365+
const defaults = {
366+
deviceScaleFactor: 1,
367+
isMobile: false,
368+
hasTouch: false,
369+
isLandscape: false,
370+
};
371+
await page.setViewport({...defaults, ...options.viewport});
372+
this.#viewportMap.set(page, {...defaults, ...options.viewport});
373+
}
374+
}
375+
376+
if (timeoutsNeedUpdate) {
377+
this.#updateSelectedPageTimeouts();
291378
}
292-
this.#updateSelectedPageTimeouts();
293379
}
294380

295381
getNetworkConditions(): string | null {
296382
const page = this.getSelectedPage();
297383
return this.#networkConditionsMap.get(page) ?? null;
298384
}
299385

300-
setCpuThrottlingRate(rate: number): void {
301-
const page = this.getSelectedPage();
302-
this.#cpuThrottlingRateMap.set(page, rate);
303-
this.#updateSelectedPageTimeouts();
304-
}
305-
306386
getCpuThrottlingRate(): number {
307387
const page = this.getSelectedPage();
308388
return this.#cpuThrottlingRateMap.get(page) ?? 1;
309389
}
310390

311-
setGeolocation(geolocation: GeolocationOptions | null): void {
312-
const page = this.getSelectedPage();
313-
if (geolocation === null) {
314-
this.#geolocationMap.delete(page);
315-
} else {
316-
this.#geolocationMap.set(page, geolocation);
317-
}
318-
}
319-
320391
getGeolocation(): GeolocationOptions | null {
321392
const page = this.getSelectedPage();
322393
return this.#geolocationMap.get(page) ?? null;
323394
}
324395

325-
setViewport(viewport: Viewport | null): void {
326-
const page = this.getSelectedPage();
327-
if (viewport === null) {
328-
this.#viewportMap.delete(page);
329-
} else {
330-
this.#viewportMap.set(page, viewport);
331-
}
332-
}
333-
334396
getViewport(): Viewport | null {
335397
const page = this.getSelectedPage();
336398
return this.#viewportMap.get(page) ?? null;
337399
}
338400

339-
setUserAgent(userAgent: string | null): void {
340-
const page = this.getSelectedPage();
341-
if (userAgent === null) {
342-
this.#userAgentMap.delete(page);
343-
} else {
344-
this.#userAgentMap.set(page, userAgent);
345-
}
346-
}
347-
348401
getUserAgent(): string | null {
349402
const page = this.getSelectedPage();
350403
return this.#userAgentMap.get(page) ?? null;
351404
}
352405

353-
setColorScheme(scheme: 'dark' | 'light' | null): void {
354-
const page = this.getSelectedPage();
355-
if (scheme === null) {
356-
this.#colorSchemeMap.delete(page);
357-
} else {
358-
this.#colorSchemeMap.set(page, scheme);
359-
}
360-
}
361-
362406
getColorScheme(): 'dark' | 'light' | null {
363407
const page = this.getSelectedPage();
364408
return this.#colorSchemeMap.get(page) ?? null;

‎src/tools/ToolDefinition.ts‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -116,14 +116,20 @@ export type Context = Readonly<{
116116
selectPage(page: Page): void;
117117
getElementByUid(uid: string): Promise<ElementHandle<Element>>;
118118
getAXNodeByUid(uid: string): TextSnapshotNode | undefined;
119-
setNetworkConditions(conditions: string | null): void;
120-
setCpuThrottlingRate(rate: number): void;
121-
setGeolocation(geolocation: GeolocationOptions | null): void;
122-
setViewport(viewport: Viewport | null): void;
119+
emulate(options: {
120+
networkConditions?: string | null;
121+
cpuThrottlingRate?: number | null;
122+
geolocation?: GeolocationOptions | null;
123+
userAgent?: string | null;
124+
colorScheme?: 'dark' | 'light' | 'auto' | null;
125+
viewport?: Viewport | null;
126+
}): Promise<void>;
127+
getNetworkConditions(): string | null;
128+
getCpuThrottlingRate(): number;
129+
getGeolocation(): GeolocationOptions | null;
123130
getViewport(): Viewport | null;
124-
setUserAgent(userAgent: string | null): void;
125131
getUserAgent(): string | null;
126-
setColorScheme(scheme: 'dark' | 'light' | null): void;
132+
getColorScheme(): 'dark' | 'light' | null;
127133
saveTemporaryFile(
128134
data: Uint8Array<ArrayBufferLike>,
129135
mimeType: 'image/png' | 'image/jpeg' | 'image/webp',

‎src/tools/emulation.ts‎

Lines changed: 1 addition & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -104,97 +104,6 @@ export const emulate = defineTool({
104104
),
105105
},
106106
handler: async (request, _response, context) => {
107-
const page = context.getSelectedPage();
108-
const {
109-
networkConditions,
110-
cpuThrottlingRate,
111-
geolocation,
112-
userAgent,
113-
viewport,
114-
} = request.params;
115-
116-
if (networkConditions) {
117-
if (networkConditions === 'No emulation') {
118-
await page.emulateNetworkConditions(null);
119-
context.setNetworkConditions(null);
120-
} else if (networkConditions === 'Offline') {
121-
await page.emulateNetworkConditions({
122-
offline: true,
123-
download: 0,
124-
upload: 0,
125-
latency: 0,
126-
});
127-
context.setNetworkConditions('Offline');
128-
} else if (networkConditions in PredefinedNetworkConditions) {
129-
const networkCondition =
130-
PredefinedNetworkConditions[
131-
networkConditions as keyof typeof PredefinedNetworkConditions
132-
];
133-
await page.emulateNetworkConditions(networkCondition);
134-
context.setNetworkConditions(networkConditions);
135-
}
136-
}
137-
138-
if (cpuThrottlingRate) {
139-
await page.emulateCPUThrottling(cpuThrottlingRate);
140-
context.setCpuThrottlingRate(cpuThrottlingRate);
141-
}
142-
143-
if (geolocation !== undefined) {
144-
if (geolocation === null) {
145-
await page.setGeolocation({latitude: 0, longitude: 0});
146-
context.setGeolocation(null);
147-
} else {
148-
await page.setGeolocation(geolocation);
149-
context.setGeolocation(geolocation);
150-
}
151-
}
152-
153-
if (userAgent !== undefined) {
154-
if (userAgent === null) {
155-
await page.setUserAgent({
156-
userAgent: undefined,
157-
});
158-
context.setUserAgent(null);
159-
} else {
160-
await page.setUserAgent({
161-
userAgent,
162-
});
163-
context.setUserAgent(userAgent);
164-
}
165-
}
166-
167-
if (request.params.colorScheme) {
168-
if (request.params.colorScheme === 'auto') {
169-
await page.emulateMediaFeatures([
170-
{name: 'prefers-color-scheme', value: ''},
171-
]);
172-
context.setColorScheme(null);
173-
} else {
174-
await page.emulateMediaFeatures([
175-
{
176-
name: 'prefers-color-scheme',
177-
value: request.params.colorScheme,
178-
},
179-
]);
180-
context.setColorScheme(request.params.colorScheme);
181-
}
182-
}
183-
184-
if (viewport !== undefined) {
185-
if (viewport === null) {
186-
await page.setViewport(null);
187-
context.setViewport(null);
188-
} else {
189-
const defaults = {
190-
deviceScaleFactor: 1,
191-
isMobile: false,
192-
hasTouch: false,
193-
isLandscape: false,
194-
};
195-
await page.setViewport({...defaults, ...viewport});
196-
context.setViewport({...defaults, ...viewport});
197-
}
198-
}
107+
await context.emulate(request.params);
199108
},
200109
});

‎tests/McpContext.test.ts‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ describe('McpContext', () => {
5151
await withMcpContext(async (_response, context) => {
5252
const page = await context.newPage();
5353
const timeoutBefore = page.getDefaultTimeout();
54-
context.setCpuThrottlingRate(2);
54+
await context.emulate({cpuThrottlingRate: 2});
5555
const timeoutAfter = page.getDefaultTimeout();
5656
assert(timeoutBefore < timeoutAfter, 'Timeout was less then expected');
5757
});
@@ -61,7 +61,7 @@ describe('McpContext', () => {
6161
await withMcpContext(async (_response, context) => {
6262
const page = await context.newPage();
6363
const timeoutBefore = page.getDefaultNavigationTimeout();
64-
context.setNetworkConditions('Slow 3G');
64+
await context.emulate({networkConditions: 'Slow 3G'});
6565
const timeoutAfter = page.getDefaultNavigationTimeout();
6666
assert(timeoutBefore < timeoutAfter, 'Timeout was less then expected');
6767
});
@@ -71,8 +71,7 @@ describe('McpContext', () => {
7171
await withMcpContext(async (_response, context) => {
7272
const page = await context.newPage();
7373

74-
context.setCpuThrottlingRate(2);
75-
context.setNetworkConditions('Slow 3G');
74+
await context.emulate({cpuThrottlingRate: 2, networkConditions: 'Slow 3G'});
7675
const stub = sinon.spy(context, 'getWaitForHelper');
7776

7877
await context.waitForEventsAfterAction(async () => {

‎tests/McpResponse.test.js.snapshot‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -263,15 +263,18 @@ exports[`McpResponse > adds userAgent emulation setting when it is set 2`] = `
263263
exports[`McpResponse > adds viewport emulation setting when it is set 1`] = `
264264
# test response
265265
## Viewport emulation
266-
Emulating viewport: {"width":400,"height":400,"deviceScaleFactor":1}
266+
Emulating viewport: {"deviceScaleFactor":1,"isMobile":false,"hasTouch":false,"isLandscape":false,"width":400,"height":400}
267267
`;
268268

269269
exports[`McpResponse > adds viewport emulation setting when it is set 2`] = `
270270
{
271271
"viewport": {
272+
"deviceScaleFactor": 1,
273+
"isMobile": false,
274+
"hasTouch": false,
275+
"isLandscape": false,
272276
"width": 400,
273-
"height": 400,
274-
"deviceScaleFactor": 1
277+
"height": 400
275278
}
276279
}
277280
`;

0 commit comments

Comments
 (0)