refactor: consistently use McpPage in tools

Author: OrKoN

src/McpContext.ts

@@ -37,7 +37,7 @@ 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';
-import type {Context, DevToolsData} from './tools/ToolDefinition.js';
+import type {Context, DevToolsData, ToolPage} from './tools/ToolDefinition.js';
 import type {TraceResult} from './trace-processing/parse.js';
 import type {
   EmulationSettings,
@@ -117,7 +117,7 @@ export class McpContext implements Context {
     null;
   #focusedPagePerContext = new Map<BrowserContext, Page>();
 
-  #requestPage?: Page;
+  #requestPage?: ToolPage;
 
   #nextPageId = 1;
 
@@ -196,12 +196,12 @@ export class McpContext implements Context {
   // TODO: Refactor away mutable request state (e.g. per-request facade,
   // per-request context object, or another approach). Once resolved, the
   // global toolMutex could become per-BrowserContext for parallel execution.
-  setRequestPage(page?: Page): void {
+  setRequestPage(page?: ToolPage): void {
     this.#requestPage = page;
   }
 
   #resolveTargetPage(): Page {
-    return this.#requestPage ?? this.getSelectedPage();
+    return this.#requestPage?.page ?? this.getSelectedPage();
   }
 
   resolveCdpRequestId(cdpRequestId: string): number | undefined {
@@ -283,7 +283,7 @@ export class McpContext implements Context {
   async newPage(
     background?: boolean,
     isolatedContextName?: string,
-  ): Promise<Page> {
+  ): Promise<ToolPage> {
     let page: Page;
     if (isolatedContextName !== undefined) {
       let ctx = this.#isolatedContexts.get(isolatedContextName);
@@ -299,7 +299,7 @@ export class McpContext implements Context {
     this.selectPage(page);
     this.#networkCollector.addPage(page);
     this.#consoleCollector.addPage(page);
-    return page;
+    return this.#getMcpPage(page);
   }
   async closePage(pageId: number): Promise<void> {
     if (this.#pages.length === 1) {
@@ -489,7 +489,7 @@ export class McpContext implements Context {
   }
 
   getDialog(page?: Page): Dialog | undefined {
-    const targetPage = page ?? this.#requestPage ?? this.#selectedPage;
+    const targetPage = page ?? this.#requestPage?.page ?? this.#selectedPage;
     if (!targetPage) {
       return undefined;
     }
@@ -516,11 +516,17 @@ export class McpContext implements Context {
     return page;
   }
 
-  resolvePageById(pageId?: number): Page {
+  getSelectedMcpPage(): McpPage {
+    const page = this.getSelectedPage();
+    return this.#getMcpPage(page);
+  }
+
+  resolvePageById(pageId?: number): McpPage {
     if (pageId === undefined) {
-      return this.getSelectedPage();
+      return this.getSelectedMcpPage();
     }
-    return this.getPageById(pageId);
+    const page = this.getPageById(pageId);
+    return this.#getMcpPage(page);
   }
 
   getPageById(pageId: number): Page {
@@ -551,7 +557,8 @@ export class McpContext implements Context {
     return this.#selectedPage === page;
   }
 
-  assertPageIsFocused(page: Page): void {
+  assertPageIsFocused(pageToCheck: Page | ToolPage): void {
+    const page = 'page' in pageToCheck ? pageToCheck.page : pageToCheck;
     const ctx = page.browserContext();
     const focused = this.#focusedPagePerContext.get(ctx);
     if (focused && focused !== page) {
@@ -564,7 +571,8 @@ export class McpContext implements Context {
     }
   }
 
-  selectPage(newPage: Page): void {
+  selectPage(pageToSelect: Page | ToolPage): void {
+    const newPage = 'page' in pageToSelect ? pageToSelect.page : pageToSelect;
     const ctx = newPage.browserContext();
     const oldFocused = this.#focusedPagePerContext.get(ctx);
     if (oldFocused && oldFocused !== newPage && !oldFocused.isClosed()) {

src/McpPage.ts

@@ -4,11 +4,19 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type {Dialog, Page, Viewport} from './third_party/index.js';
+import type {
+  Dialog,
+  ElementHandle,
+  Page,
+  Viewport,
+} from './third_party/index.js';
+import {takeSnapshot} from './tools/snapshot.js';
+import type {ToolPage} from './tools/ToolDefinition.js';
 import type {
   EmulationSettings,
   GeolocationOptions,
   TextSnapshot,
+  TextSnapshotNode,
 } from './types.js';
 
 /**
@@ -19,7 +27,7 @@ import type {
  * read/write access. The dialog field is private because it requires an
  * event listener lifecycle managed by the constructor/dispose pair.
  */
-export class McpPage {
+export class McpPage implements ToolPage {
   readonly page: Page;
   readonly id: number;
 
@@ -82,4 +90,39 @@ export class McpPage {
   dispose(): void {
     this.page.off('dialog', this.#dialogHandler);
   }
+
+  async getElementByUid(uid: string): Promise<ElementHandle<Element>> {
+    if (!this.textSnapshot) {
+      throw new Error(
+        `No snapshot found for page ${this.id ?? '?'}. Use ${takeSnapshot.name} to capture one.`,
+      );
+    }
+    const node = this.textSnapshot.idToNode.get(uid);
+    if (!node) {
+      throw new Error(`Element uid "${uid}" not found on page ${this.id}.`);
+    }
+    return this.#resolveElementHandle(node, uid);
+  }
+
+  async #resolveElementHandle(
+    node: TextSnapshotNode,
+    uid: string,
+  ): Promise<ElementHandle<Element>> {
+    const message = `Element with uid ${uid} no longer exists on the page.`;
+    try {
+      const handle = await node.elementHandle();
+      if (!handle) {
+        throw new Error(message);
+      }
+      return handle;
+    } catch (error) {
+      throw new Error(message, {
+        cause: error,
+      });
+    }
+  }
+
+  getAXNodeByUid(uid: string) {
+    return this.textSnapshot?.idToNode.get(uid);
+  }
 }

src/McpResponse.ts

@@ -263,9 +263,11 @@ export class McpResponse implements Response {
       await context.createTextSnapshot(
         this.#snapshotParams.verbose,
         this.#devToolsData,
-        this.#snapshotParams.page,
+        this.#snapshotParams.page?.page,
+      );
+      const textSnapshot = context.getTextSnapshot(
+        this.#snapshotParams.page?.page,
       );
-      const textSnapshot = context.getTextSnapshot(this.#snapshotParams.page);
       if (textSnapshot) {
         const formatter = new SnapshotFormatter(textSnapshot);
         if (this.#snapshotParams.filePath) {

src/server.ts

@@ -181,7 +181,7 @@ export async function createMcpServer(
               const page =
                 serverArgs.experimentalPageIdRouting && params.pageId
                   ? context.resolvePageById(params.pageId)
-                  : context.getSelectedPage();
+                  : context.getSelectedMcpPage();
               context.setRequestPage(page);
               await tool.handler(
                 {

src/tools/ToolDefinition.ts

@@ -60,7 +60,7 @@ export interface ImageContentData {
 export interface SnapshotParams {
   verbose?: boolean;
   filePath?: string;
-  page?: Page;
+  page?: ToolPage;
 }
 
 export interface LighthouseData {
@@ -141,7 +141,10 @@ export type Context = Readonly<{
   getDialog(page?: Page): Dialog | undefined;
   clearDialog(page?: Page): void;
   getPageById(pageId: number): Page;
-  newPage(background?: boolean, isolatedContextName?: string): Promise<Page>;
+  newPage(
+    background?: boolean,
+    isolatedContextName?: string,
+  ): Promise<ToolPage>;
   closePage(pageId: number): Promise<void>;
   selectPage(page: Page): void;
   assertPageIsFocused(page: Page): void;
@@ -191,6 +194,12 @@ export type Context = Readonly<{
   getExtension(id: string): InstalledExtension | undefined;
 }>;
 
+export type ToolPage = Readonly<{
+  readonly page: Page;
+  getAXNodeByUid(uid: string): TextSnapshotNode | undefined;
+  getElementByUid(uid: string): Promise<ElementHandle<Element>>;
+}>;
+
 export function defineTool<Schema extends zod.ZodRawShape>(
   definition: ToolDefinition<Schema>,
 ): ToolDefinition<Schema>;
@@ -223,7 +232,7 @@ interface PageToolDefinition<
   Schema extends zod.ZodRawShape = zod.ZodRawShape,
 > extends BaseToolDefinition<Schema> {
   handler: (
-    request: Request<Schema> & {page: Page},
+    request: Request<Schema> & {page: ToolPage},
     response: Response,
     context: Context,
   ) => Promise<void>;
@@ -233,7 +242,7 @@ export type DefinedPageTool<Schema extends zod.ZodRawShape = zod.ZodRawShape> =
   PageToolDefinition<Schema> & {
     pageScoped: true;
     handler: (
-      request: Request<Schema> & {page: Page},
+      request: Request<Schema> & {page: ToolPage},
       response: Response,
       context: Context,
     ) => Promise<void>;

src/tools/emulation.ts

@@ -105,6 +105,6 @@ export const emulate = definePageTool({
   },
   handler: async (request, _response, context) => {
     const page = request.page;
-    await context.emulate(request.params, page);
+    await context.emulate(request.params, page.page);
   },
 });

src/tools/input.ts

@@ -97,9 +97,9 @@ export const clickAt = definePageTool({
   },
   handler: async (request, response, context) => {
     const page = request.page;
-    context.assertPageIsFocused(page);
+    context.assertPageIsFocused(page.page);
     await context.waitForEventsAfterAction(async () => {
-      await page.mouse.click(request.params.x, request.params.y, {
+      await page.page.mouse.click(request.params.x, request.params.y, {
         clickCount: request.params.dblClick ? 2 : 1,
       });
     });
@@ -239,7 +239,7 @@ export const fill = definePageTool({
         request.params.uid,
         request.params.value,
         context as McpContext,
-        page,
+        page.page,
       );
     });
     response.appendResponseLine(`Successfully filled out the element`);
@@ -262,11 +262,11 @@ export const typeText = definePageTool({
   },
   handler: async (request, response, context) => {
     const page = request.page;
-    context.assertPageIsFocused(page);
+    context.assertPageIsFocused(page.page);
     await context.waitForEventsAfterAction(async () => {
-      await page.keyboard.type(request.params.text);
+      await page.page.keyboard.type(request.params.text);
       if (request.params.submitKey) {
-        await page.keyboard.press(request.params.submitKey as KeyInput);
+        await page.page.keyboard.press(request.params.submitKey as KeyInput);
       }
     });
     response.appendResponseLine(
@@ -333,7 +333,7 @@ export const fillForm = definePageTool({
           element.uid,
           element.value,
           context as McpContext,
-          page,
+          page.page,
         );
       });
     }
@@ -364,7 +364,7 @@ export const uploadFile = definePageTool({
     const {uid, filePath} = request.params;
     const handle = (await context.getElementByUid(
       uid,
-      request.page,
+      request.page.page,
     )) as ElementHandle<HTMLInputElement>;
     try {
       try {
@@ -375,7 +375,7 @@ export const uploadFile = definePageTool({
         // Page.waitForFileChooser() and upload the file this way.
         try {
           const [fileChooser] = await Promise.all([
-            request.page.waitForFileChooser({timeout: 3000}),
+            request.page.page.waitForFileChooser({timeout: 3000}),
             handle.asLocator().click(),
           ]);
           await fileChooser.accept([filePath]);
@@ -412,17 +412,17 @@ export const pressKey = definePageTool({
   },
   handler: async (request, response, context) => {
     const page = request.page;
-    context.assertPageIsFocused(page);
+    context.assertPageIsFocused(page.page);
     const tokens = parseKey(request.params.key);
     const [key, ...modifiers] = tokens;
 
     await context.waitForEventsAfterAction(async () => {
       for (const modifier of modifiers) {
-        await page.keyboard.down(modifier);
+        await page.page.keyboard.down(modifier);
       }
-      await page.keyboard.press(key);
+      await page.page.keyboard.press(key);
       for (const modifier of modifiers.toReversed()) {
-        await page.keyboard.up(modifier);
+        await page.page.keyboard.up(modifier);
       }
     });
 

src/tools/lighthouse.ts

@@ -82,11 +82,11 @@ export const lighthouseAudit = definePageTool({
     let result: RunnerResult | undefined;
     try {
       if (mode === 'navigation') {
-        result = await navigation(page, page.url(), {
+        result = await navigation(page.page, page.page.url(), {
           flags,
         });
       } else {
-        result = await snapshot(page, {
+        result = await snapshot(page.page, {
           flags,
         });
       }

src/tools/memory.ts

@@ -24,7 +24,7 @@ export const takeMemorySnapshot = definePageTool({
   handler: async (request, response, _context) => {
     const page = request.page;
 
-    await page.captureHeapSnapshot({
+    await page.page.captureHeapSnapshot({
       path: request.params.filePath,
     });