add TextSnapshot.ts

Author: wolfib

src/McpContext.ts

@@ -26,13 +26,12 @@ import type {
   HTTPRequest,
   Page,
   ScreenRecorder,
-  SerializedAXNode,
   Viewport,
   Target,
   Extension,
 } from './third_party/index.js';
-import type {DevTools, Protocol} from './third_party/index.js';
-import {Locator, type ElementHandle} from './third_party/index.js';
+import type {DevTools} 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 {CLOSE_PAGE_ERROR} from './tools/ToolDefinition.js';
@@ -45,8 +44,6 @@ import type {TraceResult} from './trace-processing/parse.js';
 import type {
   EmulationSettings,
   GeolocationOptions,
-  TextSnapshot,
-  TextSnapshotNode,
   ExtensionServiceWorker,
 } from './types.js';
 import {ensureExtension, saveTemporaryFile} from './utils/files.js';
@@ -92,7 +89,6 @@ export class McpContext implements Context {
   #extensionServiceWorkerMap = new WeakMap<Target, string>();
   #nextExtensionServiceWorkerId = 1;
 
-  #nextSnapshotId = 1;
   #traceResults: TraceResult[] = [];
 
   #locatorClass: typeof Locator;
@@ -687,269 +683,6 @@ export class McpContext implements Context {
   /**
    * Creates a text snapshot of a page.
    */
-  async createTextSnapshot(
-    page: McpPage,
-    verbose = false,
-    devtoolsData: DevToolsData | undefined = undefined,
-    extraHandles: ElementHandle[] = [],
-  ): Promise<void> {
-    const rootNode = await page.pptrPage.accessibility.snapshot({
-      includeIframes: true,
-      interestingOnly: !verbose,
-    });
-    if (!rootNode) {
-      return;
-    }
-
-    const {uniqueBackendNodeIdToMcpId} = page;
-
-    const snapshotId = this.#nextSnapshotId++;
-    // Iterate through the whole accessibility node tree and assign node ids that
-    // will be used for the tree serialization and mapping ids back to nodes.
-    let idCounter = 0;
-    const idToNode = new Map<string, TextSnapshotNode>();
-    const seenUniqueIds = new Set<string>();
-    const seenBackendNodeIds = new Set<number>();
-    const assignIds = (node: SerializedAXNode): TextSnapshotNode => {
-      let id = '';
-      // @ts-expect-error untyped backendNodeId.
-      const backendNodeId: number = node.backendNodeId;
-      // @ts-expect-error untyped loaderId.
-      const uniqueBackendId = `${node.loaderId}_${backendNodeId}`;
-      if (uniqueBackendNodeIdToMcpId.has(uniqueBackendId)) {
-        // Re-use MCP exposed ID if the uniqueId is the same.
-        id = uniqueBackendNodeIdToMcpId.get(uniqueBackendId)!;
-      } else {
-        // Only generate a new ID if we have not seen the node before.
-        id = `${snapshotId}_${idCounter++}`;
-        uniqueBackendNodeIdToMcpId.set(uniqueBackendId, id);
-      }
-      seenUniqueIds.add(uniqueBackendId);
-      seenBackendNodeIds.add(backendNodeId);
-
-      const nodeWithId: TextSnapshotNode = {
-        ...node,
-        id,
-        children: node.children
-          ? node.children.map(child => assignIds(child))
-          : [],
-      };
-
-      // The AXNode for an option doesn't contain its `value`.
-      // Therefore, set text content of the option as value.
-      if (node.role === 'option') {
-        const optionText = node.name;
-        if (optionText) {
-          nodeWithId.value = optionText.toString();
-        }
-      }
-
-      idToNode.set(nodeWithId.id, nodeWithId);
-      return nodeWithId;
-    };
-
-    const rootNodeWithId = assignIds(rootNode);
-
-    await this.#insertExtraNodes(
-      page,
-      idToNode,
-      seenUniqueIds,
-      snapshotId,
-      idCounter,
-      rootNodeWithId,
-      seenBackendNodeIds,
-      extraHandles,
-    );
-
-    const snapshot: TextSnapshot = {
-      root: rootNodeWithId,
-      snapshotId: String(snapshotId),
-      idToNode,
-      hasSelectedElement: false,
-      verbose,
-    };
-    page.textSnapshot = snapshot;
-    const data = devtoolsData ?? (await this.getDevToolsData(page));
-    if (data?.cdpBackendNodeId) {
-      snapshot.hasSelectedElement = true;
-      snapshot.selectedElementUid = page.resolveCdpElementId(
-        data?.cdpBackendNodeId,
-      );
-    }
-
-    // Clean up unique IDs that we did not see anymore.
-    for (const key of uniqueBackendNodeIdToMcpId.keys()) {
-      if (!seenUniqueIds.has(key)) {
-        uniqueBackendNodeIdToMcpId.delete(key);
-      }
-    }
-  }
-
-  // ExtraHandles represent DOM nodes which might not be part of the accessibility tree, e.g. DOM nodes
-  // returned by in-page tools. We insert them into the tree by finding the closest ancestor in the
-  // tree and inserting the node as a child. The ancestor's child nodes are re-parented if necessary.
-  async #insertExtraNodes(
-    page: McpPage,
-    idToNode: Map<string, TextSnapshotNode>,
-    seenUniqueIds: Set<string>,
-    snapshotId: number,
-    idCounter: number,
-    rootNodeWithId: TextSnapshotNode,
-    seenBackendNodeIds: Set<number>,
-    extraHandles: ElementHandle[],
-  ): Promise<void> {
-    const {uniqueBackendNodeIdToMcpId} = page;
-
-    const createExtraNode = async (
-      handle: ElementHandle,
-    ): Promise<TextSnapshotNode | null> => {
-      const backendNodeId = await handle.backendNodeId();
-      if (!backendNodeId || seenBackendNodeIds.has(backendNodeId)) {
-        return null;
-      }
-      const uniqueBackendId = `custom_${backendNodeId}`;
-      if (seenUniqueIds.has(uniqueBackendId)) {
-        return null;
-      }
-      seenBackendNodeIds.add(backendNodeId);
-
-      let id = '';
-      const mcpId = uniqueBackendNodeIdToMcpId.get(uniqueBackendId);
-      if (mcpId !== undefined) {
-        id = mcpId;
-      } else {
-        id = `${snapshotId}_${idCounter++}`;
-        uniqueBackendNodeIdToMcpId.set(uniqueBackendId, id);
-      }
-      seenUniqueIds.add(uniqueBackendId);
-
-      const tagHandle = await handle.getProperty('localName');
-      const tagValue = await tagHandle.jsonValue();
-      const extraNode: TextSnapshotNode = {
-        role: tagValue,
-        id,
-        backendNodeId,
-        children: [],
-        elementHandle: async () => handle,
-      };
-      return extraNode;
-    };
-
-    const findAncestorNode = async (
-      handle: ElementHandle,
-    ): Promise<TextSnapshotNode | null> => {
-      let ancestorHandle = await handle.evaluateHandle(el => el.parentElement);
-
-      while (ancestorHandle) {
-        const ancestorElement = ancestorHandle.asElement();
-        if (!ancestorElement) {
-          await ancestorHandle.dispose();
-          return null;
-        }
-
-        const ancestorBackendId = await ancestorElement.backendNodeId();
-        if (ancestorBackendId) {
-          const ancestorNode = idToNode
-            .values()
-            .find(node => node.backendNodeId === ancestorBackendId);
-          if (ancestorNode) {
-            await ancestorHandle.dispose();
-            return ancestorNode;
-          }
-        }
-
-        const nextHandle = await ancestorElement.evaluateHandle(
-          el => el.parentElement,
-        );
-        await ancestorHandle.dispose();
-        ancestorHandle = nextHandle;
-      }
-      return null;
-    };
-
-    const findDescendantNodes = async (
-      backendNodeId: number,
-    ): Promise<Set<number>> => {
-      const descendantIds = new Set<number>();
-      try {
-        // @ts-expect-error internal API
-        const client = page.pptrPage._client();
-        if (client) {
-          const {node}: {node: Protocol.DOM.Node} = await client.send(
-            'DOM.describeNode',
-            {
-              backendNodeId,
-              depth: -1,
-              pierce: true,
-            },
-          );
-          const collect = (node: Protocol.DOM.Node) => {
-            if (node.backendNodeId && node.backendNodeId !== backendNodeId) {
-              descendantIds.add(node.backendNodeId);
-            }
-            if (node.children) {
-              for (const child of node.children) {
-                collect(child);
-              }
-            }
-          };
-          collect(node);
-        }
-      } catch (e) {
-        this.logger(
-          `Failed to collect descendants for backend node ${backendNodeId}`,
-          e,
-        );
-      }
-      return descendantIds;
-    };
-
-    const moveChildNodes = (
-      attachTarget: TextSnapshotNode,
-      extraNode: TextSnapshotNode,
-      descendantIds: Set<number>,
-    ): number => {
-      let firstMovedIndex = -1;
-      if (descendantIds.size > 0 && attachTarget.children) {
-        const remainingChildren: TextSnapshotNode[] = [];
-        for (const child of attachTarget.children) {
-          if (child.backendNodeId && descendantIds.has(child.backendNodeId)) {
-            if (firstMovedIndex === -1) {
-              firstMovedIndex = remainingChildren.length;
-            }
-            extraNode.children.push(child);
-          } else {
-            remainingChildren.push(child);
-          }
-        }
-        attachTarget.children = remainingChildren;
-      }
-      return firstMovedIndex !== -1
-        ? firstMovedIndex
-        : attachTarget.children
-          ? attachTarget.children.length
-          : 0;
-    };
-
-    if (extraHandles.length) {
-      page.extraHandles = extraHandles;
-    }
-    for (const handle of page.extraHandles) {
-      const extraNode = await createExtraNode(handle);
-      if (!extraNode) {
-        continue;
-      }
-      idToNode.set(extraNode.id, extraNode);
-      const attachTarget = (await findAncestorNode(handle)) || rootNodeWithId;
-      if (extraNode.backendNodeId !== undefined) {
-        const descendantIds = await findDescendantNodes(
-          extraNode.backendNodeId,
-        );
-        const index = moveChildNodes(attachTarget, extraNode, descendantIds);
-        attachTarget.children.splice(index, 0, extraNode);
-      }
-    }
-  }
 
   async saveTemporaryFile(
     data: Uint8Array<ArrayBufferLike>,

src/McpPage.ts

@@ -5,6 +5,7 @@
  */
 
 import {logger} from './logger.js';
+import {TextSnapshot} from './TextSnapshot.js';
 import type {
   Dialog,
   ElementHandle,
@@ -18,7 +19,6 @@ import type {ContextPage, Context, Response} from './tools/ToolDefinition.js';
 import type {
   EmulationSettings,
   GeolocationOptions,
-  TextSnapshot,
   TextSnapshotNode,
 } from './types.js';
 import {
@@ -277,12 +277,9 @@ export class McpPage implements ContextPage {
     }
     const resultWithStashedElements = result.result;
     if (elementHandles.length) {
-      await context.createTextSnapshot(
-        this,
-        false,
-        undefined,
-        elementHandles,
-      );
+      this.textSnapshot = await TextSnapshot.create(this, context, {
+        extraHandles: elementHandles,
+      });
       response.includeSnapshot();
     }
 

src/McpResponse.ts

@@ -15,6 +15,7 @@ import {SnapshotFormatter} from './formatters/SnapshotFormatter.js';
 import type {McpContext} from './McpContext.js';
 import type {McpPage} from './McpPage.js';
 import {UncaughtError} from './PageCollector.js';
+import {TextSnapshot} from './TextSnapshot.js';
 import {DevTools, type Protocol} from './third_party/index.js';
 import type {
   ConsoleMessage,
@@ -443,11 +444,10 @@ export class McpResponse implements Response {
       if (!this.#page) {
         throw new Error('Response must have a page');
       }
-      await context.createTextSnapshot(
-        this.#page,
-        this.#snapshotParams.verbose,
-        this.#devToolsData,
-      );
+      this.#page.textSnapshot = await TextSnapshot.create(this.#page, context, {
+        verbose: this.#snapshotParams.verbose,
+        devtoolsData: this.#devToolsData,
+      });
       const textSnapshot = this.#page.textSnapshot;
       if (textSnapshot) {
         const formatter = new SnapshotFormatter(textSnapshot);