feat(emulation): add geolocation emulation tool

Author: JustasMonkev

src/McpContext.ts

@@ -38,6 +38,11 @@ export interface TextSnapshotNode extends SerializedAXNode {
   children: TextSnapshotNode[];
 }
 
+export interface GeolocationOptions {
+  latitude: number;
+  longitude: number;
+}
+
 export interface TextSnapshot {
   root: TextSnapshotNode;
   idToNode: Map<string, TextSnapshotNode>;
@@ -104,6 +109,7 @@ export class McpContext implements Context {
   #isRunningTrace = false;
   #networkConditionsMap = new WeakMap<Page, string>();
   #cpuThrottlingRateMap = new WeakMap<Page, number>();
+  #geolocationMap = new WeakMap<Page, GeolocationOptions>();
   #dialog?: Dialog;
 
   #nextSnapshotId = 1;
@@ -277,6 +283,20 @@ export class McpContext implements Context {
     return this.#cpuThrottlingRateMap.get(page) ?? 1;
   }
 
+  setGeolocation(geolocation: GeolocationOptions | null): void {
+    const page = this.getSelectedPage();
+    if (geolocation === null) {
+      this.#geolocationMap.delete(page);
+    } else {
+      this.#geolocationMap.set(page, geolocation);
+    }
+  }
+
+  getGeolocation(): GeolocationOptions | null {
+    const page = this.getSelectedPage();
+    return this.#geolocationMap.get(page) ?? null;
+  }
+
   setIsRunningPerformanceTrace(x: boolean): void {
     this.#isRunningTrace = x;
   }

src/tools/ToolDefinition.ts

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type {TextSnapshotNode} from '../McpContext.js';
+import type {TextSnapshotNode, GeolocationOptions} from '../McpContext.js';
 import {zod} from '../third_party/index.js';
 import type {Dialog, ElementHandle, Page} from '../third_party/index.js';
 import type {TraceResult} from '../trace-processing/parse.js';
@@ -98,6 +98,7 @@ export type Context = Readonly<{
   getAXNodeByUid(uid: string): TextSnapshotNode | undefined;
   setNetworkConditions(conditions: string | null): void;
   setCpuThrottlingRate(rate: number): void;
+  setGeolocation(geolocation: GeolocationOptions | null): void;
   saveTemporaryFile(
     data: Uint8Array<ArrayBufferLike>,
     mimeType: 'image/png' | 'image/jpeg' | 'image/webp',

src/tools/emulation.ts

@@ -77,3 +77,57 @@ export const emulate = defineTool({
     }
   },
 });
+
+export const emulateGeolocation = defineTool({
+  name: 'emulate_geolocation',
+  description: `Emulates geolocation on the selected page. Useful for testing location-based features.`,
+  annotations: {
+    category: ToolCategory.EMULATION,
+    readOnlyHint: false,
+  },
+  schema: {
+    latitude: zod
+      .number()
+      .min(-90)
+      .max(90)
+      .optional()
+      .describe(
+        'Latitude between -90 and 90. Omit latitude and longitude to clear the override.',
+      ),
+    longitude: zod
+      .number()
+      .min(-180)
+      .max(180)
+      .optional()
+      .describe(
+        'Longitude between -180 and 180. Omit latitude and longitude to clear the override.',
+      ),
+  },
+  handler: async (request, _response, context) => {
+    const page = context.getSelectedPage();
+    const {latitude, longitude} = request.params;
+
+    if (latitude === undefined && longitude === undefined) {
+      // Clear geolocation override
+      await page.setGeolocation({
+        latitude: 0,
+        longitude: 0,
+      });
+      context.setGeolocation(null);
+    } else if (latitude !== undefined && longitude !== undefined) {
+      // Set geolocation override
+      await page.setGeolocation({
+        latitude,
+        longitude,
+      });
+      context.setGeolocation({
+        latitude,
+        longitude,
+      });
+    } else {
+      throw new Error(
+        'Both latitude and longitude must be provided, or both must be omitted to clear the override.',
+      );
+    }
+  },
+});

tests/tools/emulation.test.ts

@@ -7,7 +7,7 @@
 import assert from 'node:assert';
 import {describe, it} from 'node:test';
 
-import {emulate} from '../../src/tools/emulation.js';
+import {emulate, emulateGeolocation} from '../../src/tools/emulation.js';
 import {withMcpContext} from '../utils.js';
 
 describe('emulation', () => {
@@ -152,4 +152,100 @@ describe('emulation', () => {
       });
     });
   });
+
+  describe('geolocation', () => {
+    it('emulates geolocation with latitude and longitude', async () => {
+      await withMcpContext(async (response, context) => {
+        await emulateGeolocation.handler(
+          {
+            params: {
+              latitude: 48.137154,
+              longitude: 11.576124,
+            },
+          },
+          response,
+          context,
+        );
+
+        const geolocation = context.getGeolocation();
+        assert.strictEqual(geolocation?.latitude, 48.137154);
+        assert.strictEqual(geolocation?.longitude, 11.576124);
+      });
+    });
+
+    it('clears geolocation override when both params are omitted', async () => {
+      await withMcpContext(async (response, context) => {
+        // First set a geolocation
+        await emulateGeolocation.handler(
+          {
+            params: {
+              latitude: 48.137154,
+              longitude: 11.576124,
+            },
+          },
+          response,
+          context,
+        );
+
+        assert.notStrictEqual(context.getGeolocation(), null);
+
+        // Then clear it by omitting both params
+        await emulateGeolocation.handler(
+          {
+            params: {},
+          },
+          response,
+          context,
+        );
+
+        assert.strictEqual(context.getGeolocation(), null);
+      });
+    });
+
+    it('throws error when only latitude is provided', async () => {
+      await withMcpContext(async (response, context) => {
+        await assert.rejects(
+          async () => {
+            await emulateGeolocation.handler(
+              {
+                params: {
+                  latitude: 48.137154,
+                },
+              },
+              response,
+              context,
+            );
+          },
+          {
+            message:
+              'Both latitude and longitude must be provided, or both must be omitted to clear the override.',
+          },
+        );
+      });
+    });
+
+    it('reports correctly for the currently selected page', async () => {
+      await withMcpContext(async (response, context) => {
+        await emulateGeolocation.handler(
+          {
+            params: {
+              latitude: 48.137154,
+              longitude: 11.576124,
+            },
+          },
+          response,
+          context,
+        );
+
+        const geolocation = context.getGeolocation();
+        assert.strictEqual(geolocation?.latitude, 48.137154);
+        assert.strictEqual(geolocation?.longitude, 11.576124);
+
+        const page = await context.newPage();
+        context.selectPage(page);
+
+        assert.strictEqual(context.getGeolocation(), null);
+      });
+    });
+  });
 });