refactor: move emulation settings to context

src/McpContext.ts

@@ -16,7 +16,6 @@ import {
 } from './DevtoolsUtils.js';
 import type {ListenerMap, UncaughtError} from './PageCollector.js';
 import {NetworkCollector, ConsoleCollector} from './PageCollector.js';
-import {Locator} from './third_party/index.js';
 import type {DevTools} from './third_party/index.js';
 import type {
   Browser,
@@ -27,9 +26,10 @@ import type {
   HTTPRequest,
   Page,
   SerializedAXNode,
-  PredefinedNetworkConditions,
   Viewport,
 } from './third_party/index.js';
+import {Locator} from './third_party/index.js';
+import {PredefinedNetworkConditions} from './third_party/index.js';
 import {listPages} from './tools/pages.js';
 import {takeSnapshot} from './tools/snapshot.js';
 import {CLOSE_PAGE_ERROR} from './tools/ToolDefinition.js';
@@ -64,6 +64,15 @@ export interface TextSnapshot {
   verbose: boolean;
 }
 
+interface EmulationSettings {
+  networkConditions?: string | null;
+  cpuThrottlingRate?: number | null;
+  geolocation?: GeolocationOptions | null;
+  userAgent?: string | null;
+  colorScheme?: 'dark' | 'light' | null;
+  viewport?: Viewport | null;
+}
+
 interface McpContextOptions {
   // Whether the DevTools windows are exposed as pages for debugging of DevTools.
   experimentalDevToolsDebugging: boolean;
@@ -121,12 +130,7 @@ export class McpContext implements Context {
   #extensionRegistry = new ExtensionRegistry();
 
   #isRunningTrace = false;
-  #networkConditionsMap = new WeakMap<Page, string>();
-  #cpuThrottlingRateMap = new WeakMap<Page, number>();
-  #geolocationMap = new WeakMap<Page, GeolocationOptions>();
-  #viewportMap = new WeakMap<Page, Viewport>();
-  #userAgentMap = new WeakMap<Page, string>();
-  #colorSchemeMap = new WeakMap<Page, 'dark' | 'light'>();
+  #emulationSettingsMap = new WeakMap<Page, EmulationSettings>();
   #dialog?: Dialog;
 
   #pageIdMap = new WeakMap<Page, number>();
@@ -282,86 +286,146 @@ export class McpContext implements Context {
     return this.#networkCollector.getById(this.getSelectedPage(), reqid);
   }
 
-  setNetworkConditions(conditions: string | null): void {
+  async emulate(options: {
+    networkConditions?: string | null;
+    cpuThrottlingRate?: number | null;
+    geolocation?: GeolocationOptions | null;
+    userAgent?: string | null;
+    colorScheme?: 'dark' | 'light' | 'auto' | null;
+    viewport?: Viewport | null;
+  }): Promise<void> {
     const page = this.getSelectedPage();
-    if (conditions === null) {
-      this.#networkConditionsMap.delete(page);
-    } else {
-      this.#networkConditionsMap.set(page, conditions);
+    const currentSettings = this.#emulationSettingsMap.get(page) ?? {};
+    const newSettings: EmulationSettings = {...currentSettings};
+    let timeoutsNeedUpdate = false;
+
+    if (options.networkConditions !== undefined) {
+      timeoutsNeedUpdate = true;
+      if (
+        options.networkConditions === null ||
+        options.networkConditions === 'No emulation'
+      ) {
+        await page.emulateNetworkConditions(null);
+        delete newSettings.networkConditions;
+      } else if (options.networkConditions === 'Offline') {
+        await page.emulateNetworkConditions({
+          offline: true,
+          download: 0,
+          upload: 0,
+          latency: 0,
+        });
+        newSettings.networkConditions = 'Offline';
+      } else if (options.networkConditions in PredefinedNetworkConditions) {
+        const networkCondition =
+          PredefinedNetworkConditions[
+            options.networkConditions as keyof typeof PredefinedNetworkConditions
+          ];
+        await page.emulateNetworkConditions(networkCondition);
+        newSettings.networkConditions = options.networkConditions;
+      }
     }
-    this.#updateSelectedPageTimeouts();
-  }
 
-  getNetworkConditions(): string | null {
-    const page = this.getSelectedPage();
-    return this.#networkConditionsMap.get(page) ?? null;
-  }
+    if (options.cpuThrottlingRate !== undefined) {
+      timeoutsNeedUpdate = true;
+      if (options.cpuThrottlingRate === null) {
+        await page.emulateCPUThrottling(1);
+        delete newSettings.cpuThrottlingRate;
+      } else {
+        await page.emulateCPUThrottling(options.cpuThrottlingRate);
+        newSettings.cpuThrottlingRate = options.cpuThrottlingRate;
+      }
+    }
 
-  setCpuThrottlingRate(rate: number): void {
-    const page = this.getSelectedPage();
-    this.#cpuThrottlingRateMap.set(page, rate);
-    this.#updateSelectedPageTimeouts();
-  }
+    if (options.geolocation !== undefined) {
+      if (options.geolocation === null) {
+        await page.setGeolocation({latitude: 0, longitude: 0});
+        delete newSettings.geolocation;
+      } else {
+        await page.setGeolocation(options.geolocation);
+        newSettings.geolocation = options.geolocation;
+      }
+    }
 
-  getCpuThrottlingRate(): number {
-    const page = this.getSelectedPage();
-    return this.#cpuThrottlingRateMap.get(page) ?? 1;
-  }
+    if (options.userAgent !== undefined) {
+      if (options.userAgent === null) {
+        await page.setUserAgent({userAgent: undefined});
+        delete newSettings.userAgent;
+      } else {
+        await page.setUserAgent({userAgent: options.userAgent});
+        newSettings.userAgent = options.userAgent;
+      }
+    }
 
-  setGeolocation(geolocation: GeolocationOptions | null): void {
-    const page = this.getSelectedPage();
-    if (geolocation === null) {
-      this.#geolocationMap.delete(page);
+    if (options.colorScheme !== undefined) {
+      if (options.colorScheme === null || options.colorScheme === 'auto') {
+        await page.emulateMediaFeatures([
+          {name: 'prefers-color-scheme', value: ''},
+        ]);
+        delete newSettings.colorScheme;
+      } else {
+        await page.emulateMediaFeatures([
+          {name: 'prefers-color-scheme', value: options.colorScheme},
+        ]);
+        newSettings.colorScheme = options.colorScheme;
+      }
+    }
+
+    if (options.viewport !== undefined) {
+      if (options.viewport === null) {
+        await page.setViewport(null);
+        delete newSettings.viewport;
+      } else {
+        const defaults = {
+          deviceScaleFactor: 1,
+          isMobile: false,
+          hasTouch: false,
+          isLandscape: false,
+        };
+        const viewport = {...defaults, ...options.viewport};
+        await page.setViewport(viewport);
+        newSettings.viewport = viewport;
+      }
+    }
+
+    if (Object.keys(newSettings).length) {
+      this.#emulationSettingsMap.set(page, newSettings);
     } else {
-      this.#geolocationMap.set(page, geolocation);
+      this.#emulationSettingsMap.delete(page);
     }
-  }
 
-  getGeolocation(): GeolocationOptions | null {
-    const page = this.getSelectedPage();
-    return this.#geolocationMap.get(page) ?? null;
+    if (timeoutsNeedUpdate) {
+      this.#updateSelectedPageTimeouts();
+    }
   }
 
-  setViewport(viewport: Viewport | null): void {
+  getNetworkConditions(): string | null {
     const page = this.getSelectedPage();
-    if (viewport === null) {
-      this.#viewportMap.delete(page);
-    } else {
-      this.#viewportMap.set(page, viewport);
-    }
+    return this.#emulationSettingsMap.get(page)?.networkConditions ?? null;
   }
 
-  getViewport(): Viewport | null {
+  getCpuThrottlingRate(): number {
     const page = this.getSelectedPage();
-    return this.#viewportMap.get(page) ?? null;
+    return this.#emulationSettingsMap.get(page)?.cpuThrottlingRate ?? 1;
   }
 
-  setUserAgent(userAgent: string | null): void {
+  getGeolocation(): GeolocationOptions | null {
     const page = this.getSelectedPage();
-    if (userAgent === null) {
-      this.#userAgentMap.delete(page);
-    } else {
-      this.#userAgentMap.set(page, userAgent);
-    }
+    return this.#emulationSettingsMap.get(page)?.geolocation ?? null;
   }
 
-  getUserAgent(): string | null {
+  getViewport(): Viewport | null {
     const page = this.getSelectedPage();
-    return this.#userAgentMap.get(page) ?? null;
+    return this.#emulationSettingsMap.get(page)?.viewport ?? null;
   }
 
-  setColorScheme(scheme: 'dark' | 'light' | null): void {
+  getUserAgent(): string | null {
     const page = this.getSelectedPage();
-    if (scheme === null) {
-      this.#colorSchemeMap.delete(page);
-    } else {
-      this.#colorSchemeMap.set(page, scheme);
-    }
+    return this.#emulationSettingsMap.get(page)?.userAgent ?? null;
   }
 
   getColorScheme(): 'dark' | 'light' | null {
     const page = this.getSelectedPage();
-    return this.#colorSchemeMap.get(page) ?? null;
+    return this.#emulationSettingsMap.get(page)?.colorScheme ?? null;
   }
 
   setIsRunningPerformanceTrace(x: boolean): void {

src/tools/ToolDefinition.ts

@@ -116,14 +116,20 @@ export type Context = Readonly<{
   selectPage(page: Page): void;
   getElementByUid(uid: string): Promise<ElementHandle<Element>>;
   getAXNodeByUid(uid: string): TextSnapshotNode | undefined;
-  setNetworkConditions(conditions: string | null): void;
-  setCpuThrottlingRate(rate: number): void;
-  setGeolocation(geolocation: GeolocationOptions | null): void;
-  setViewport(viewport: Viewport | null): void;
+  emulate(options: {
+    networkConditions?: string | null;
+    cpuThrottlingRate?: number | null;
+    geolocation?: GeolocationOptions | null;
+    userAgent?: string | null;
+    colorScheme?: 'dark' | 'light' | 'auto' | null;
+    viewport?: Viewport | null;
+  }): Promise<void>;
+  getNetworkConditions(): string | null;
+  getCpuThrottlingRate(): number;
+  getGeolocation(): GeolocationOptions | null;
   getViewport(): Viewport | null;
-  setUserAgent(userAgent: string | null): void;
   getUserAgent(): string | null;
-  setColorScheme(scheme: 'dark' | 'light' | null): void;
+  getColorScheme(): 'dark' | 'light' | null;
   saveTemporaryFile(
     data: Uint8Array<ArrayBufferLike>,
     mimeType: 'image/png' | 'image/jpeg' | 'image/webp',

src/tools/emulation.ts

@@ -104,97 +104,6 @@ export const emulate = defineTool({
       ),
   },
   handler: async (request, _response, context) => {
-    const page = context.getSelectedPage();
-    const {
-      networkConditions,
-      cpuThrottlingRate,
-      geolocation,
-      userAgent,
-      viewport,
-    } = request.params;
-
-    if (networkConditions) {
-      if (networkConditions === 'No emulation') {
-        await page.emulateNetworkConditions(null);
-        context.setNetworkConditions(null);
-      } else if (networkConditions === 'Offline') {
-        await page.emulateNetworkConditions({
-          offline: true,
-          download: 0,
-          upload: 0,
-          latency: 0,
-        });
-        context.setNetworkConditions('Offline');
-      } else if (networkConditions in PredefinedNetworkConditions) {
-        const networkCondition =
-          PredefinedNetworkConditions[
-            networkConditions as keyof typeof PredefinedNetworkConditions
-          ];
-        await page.emulateNetworkConditions(networkCondition);
-        context.setNetworkConditions(networkConditions);
-      }
-    }
-
-    if (cpuThrottlingRate) {
-      await page.emulateCPUThrottling(cpuThrottlingRate);
-      context.setCpuThrottlingRate(cpuThrottlingRate);
-    }
-
-    if (geolocation !== undefined) {
-      if (geolocation === null) {
-        await page.setGeolocation({latitude: 0, longitude: 0});
-        context.setGeolocation(null);
-      } else {
-        await page.setGeolocation(geolocation);
-        context.setGeolocation(geolocation);
-      }
-    }
-
-    if (userAgent !== undefined) {
-      if (userAgent === null) {
-        await page.setUserAgent({
-          userAgent: undefined,
-        });
-        context.setUserAgent(null);
-      } else {
-        await page.setUserAgent({
-          userAgent,
-        });
-        context.setUserAgent(userAgent);
-      }
-    }
-
-    if (request.params.colorScheme) {
-      if (request.params.colorScheme === 'auto') {
-        await page.emulateMediaFeatures([
-          {name: 'prefers-color-scheme', value: ''},
-        ]);
-        context.setColorScheme(null);
-      } else {
-        await page.emulateMediaFeatures([
-          {
-            name: 'prefers-color-scheme',
-            value: request.params.colorScheme,
-          },
-        ]);
-        context.setColorScheme(request.params.colorScheme);
-      }
-    }
-
-    if (viewport !== undefined) {
-      if (viewport === null) {
-        await page.setViewport(null);
-        context.setViewport(null);
-      } else {
-        const defaults = {
-          deviceScaleFactor: 1,
-          isMobile: false,
-          hasTouch: false,
-          isLandscape: false,
-        };
-        await page.setViewport({...defaults, ...viewport});
-        context.setViewport({...defaults, ...viewport});
-      }
-    }
+    await context.emulate(request.params);
   },
 });