perf: fix quadratic text/nth lookup (12.5s to 27ms)

.changeset/perf-text-selectors.md

@@ -0,0 +1,8 @@
+---
+"@browserbasehq/stagehand": patch
+"@browserbasehq/stagehand-go": patch
+"@browserbasehq/stagehand-python": patch
+"@browserbasehq/stagehand-extension": patch
+---
+
+Resolve nth matches directly and remove quadratic text-selector filtering while preserving match order and shadow-root boundaries. Skip frontend node materialization for locator actions.

packages/extension/dom/locatorScripts/counts.ts

@@ -1,3 +1,4 @@
+import { filterInnermostMatches } from "./textMatches.js";
 import { countXPathMatches } from "./xpathResolver.js";
 import { getOpenOrClosedShadowRoot } from "./shadowRoots.js";
 
@@ -176,23 +177,7 @@ export function countTextMatches(rawNeedle: string): TextMatchResult {
     }
   }
 
-  const innermost: typeof matchesList = [];
-  for (const item of matchesList) {
-    const el = item.element;
-    let skip = false;
-    for (const other of matchesList) {
-      if (item === other) continue;
-      try {
-        if (el.contains(other.element)) {
-          skip = true;
-          break;
-        }
-      } catch {
-        // ignore containment errors
-      }
-    }
-    if (!skip) innermost.push(item);
-  }
+  const innermost = filterInnermostMatches(matchesList);
 
   const count = innermost.length;
   const sample = innermost.slice(0, 5).map((item) => ({

packages/extension/dom/locatorScripts/selectors.ts

@@ -1,3 +1,4 @@
+import { filterInnermostMatches } from "./textMatches.js";
 import { resolveXPathAtIndex } from "./xpathResolver.js";
 import { getOpenOrClosedShadowRoot } from "./shadowRoots.js";
 
@@ -172,25 +173,7 @@ export function resolveTextSelector(rawNeedle: string, targetIndexRaw?: number):
     }
   }
 
-  const innermost: typeof matchesList = [];
-  for (const item of matchesList) {
-    const el = item.element;
-    let skip = false;
-    for (const other of matchesList) {
-      if (item === other) continue;
-      try {
-        if (el.contains(other.element)) {
-          skip = true;
-          break;
-        }
-      } catch {
-        // ignore containment errors
-      }
-    }
-    if (!skip) {
-      innermost.push(item);
-    }
-  }
+  const innermost = filterInnermostMatches(matchesList);
 
   const target = innermost[targetIndex];
   return target?.element ?? null;

packages/extension/dom/locatorScripts/textMatches.ts

@@ -0,0 +1,15 @@
+export function filterInnermostMatches<T extends { element: Element }>(matches: T[]): T[] {
+  const matchingElements = new Set(matches.map(({ element }) => element));
+  const matchingAncestors = new Set<Element>();
+  const visited = new Set<Element>();
+  for (const { element } of matches) {
+    let parent = element.parentElement;
+    // parentElement stops at a shadow root, matching Element.contains semantics.
+    while (parent && !visited.has(parent)) {
+      visited.add(parent);
+      if (matchingElements.has(parent)) matchingAncestors.add(parent);
+      parent = parent.parentElement;
+    }
+  }
+  return matches.filter(({ element }) => !matchingAncestors.has(element));
+}

packages/extension/tests/text-matches.test.ts

@@ -0,0 +1,39 @@
+import { describe, expect, it } from "vitest";
+import { filterInnermostMatches } from "../dom/locatorScripts/textMatches.js";
+
+function element(parentElement: Element | null = null): Element {
+  return { parentElement } as Element;
+}
+
+describe("innermost text matches", () => {
+  it("removes matching ancestors across nonmatching elements and preserves order", () => {
+    const root = element();
+    const middle = element(root);
+    const leaf = element(middle);
+    const sibling = element(root);
+    const other = element();
+    const matches = [root, leaf, sibling, other].map((element, i) => ({ element, i }));
+    expect(filterInnermostMatches(matches)).toEqual(matches.slice(1));
+  });
+
+  it("keeps matching shadow hosts alongside matches in their separate tree", () => {
+    const host = element();
+    const shadowChild = element();
+    const nestedChild = element(shadowChild);
+    const matches = [host, shadowChild, nestedChild].map((element) => ({ element }));
+    expect(filterInnermostMatches(matches)).toEqual([matches[0], matches[2]]);
+  });
+
+  it("reads each shared ancestor once instead of comparing all pairs", () => {
+    let ancestorReads = 0;
+    const root = {
+      get parentElement() {
+        ancestorReads += 1;
+        return null;
+      },
+    } as Element;
+    const matches = Array.from({ length: 10_000 }, () => ({ element: element(root) }));
+    expect(filterInnermostMatches(matches)).toEqual(matches);
+    expect(ancestorReads).toBe(1);
+  });
+});

packages/extension/understudy/deepLocator.test.ts

@@ -85,9 +85,7 @@ describe("locator resolution contexts", () => {
       vi.useFakeTimers();
       const { page, root, inner } = createFrames();
       const resolveFrame = vi.spyOn(FrameLocator.prototype, "resolveFrame");
-      const lookups = (["resolveCss", "resolveText", "resolveXPath"] as const).map((method) =>
-        vi.spyOn(FrameSelectorResolver.prototype, method),
-      );
+      const lookup = vi.spyOn(FrameSelectorResolver.prototype, "resolveAtIndex");
       const delegate =
         kind === "frame locator"
           ? frameLocatorFromFrame(page, root, "#outer")
@@ -107,10 +105,10 @@ describe("locator resolution contexts", () => {
         expect(locator.nthIndex).toBe(1);
         await expect(locator.resolveNode(progress)).resolves.toEqual({
           objectId: "node",
-          nodeId: 1,
+          nodeId: null,
         });
         expect(resolveFrame.mock.calls).toEqual([[progress], [progress]]);
-        const calls = lookups.flatMap((lookup) => lookup.mock.calls);
+        const calls = lookup.mock.calls;
         expect(calls).toHaveLength(3);
         expect(calls.every(([, , context]) => context === progress)).toBe(true);
         expect(progress.remainingMs()).toBe(100);

packages/extension/understudy/locator.ts

@@ -38,8 +38,8 @@ const MAX_REMOTE_UPLOAD_BYTES = 50 * 1024 * 1024; // 50MB guard copied from Play
  *
  * Key change:
  * - Prefer **objectId**-based CDP calls (scroll, geometry) to avoid brittle
- *   frontend nodeId mappings. nodeId is resolved on a best-effort basis and
- *   returned for compatibility, but actions do not depend on it.
+ *   frontend nodeId mappings. Actions do not need nodeId, so resolution skips
+ *   DOM.requestNode and leaves nodeId null.
  *
  * Notes:
  * - Resolution is lazy: every action resolves the selector again.
@@ -60,7 +60,7 @@ export class Locator {
     readonly options?: { deep?: boolean; depth?: number },
     nthIndex: number = -1,
   ) {
-    this.selectorResolver = new FrameSelectorResolver(this.frame);
+    this.selectorResolver = new FrameSelectorResolver(this.frame, {}, { includeNodeId: false });
     this.selectorQuery = FrameSelectorResolver.parseSelector(selector);
     const normalized = Number.isFinite(nthIndex) ? Math.floor(nthIndex) : -1;
     this.nthIndex = normalized < 0 ? -1 : normalized;

packages/extension/understudy/locatorResolution.test.ts

@@ -8,6 +8,7 @@ import { Frame } from "./frame.js";
 import { frameLocatorFromFrame } from "./frameLocator.js";
 import { Progress, runWithProgress } from "./progress.js";
 import type { Page } from "./page.js";
+import { FrameSelectorResolver } from "./selectorResolver.js";
 
 function deferred<T = unknown>() {
   let resolve!: (value: T) => void;
@@ -66,6 +67,79 @@ function createPage(...frames: Frame[]) {
   } as unknown as Page;
 }
 
+describe("locator object-only resolution", () => {
+  afterEach(() => vi.restoreAllMocks());
+
+  it.each(["button", "text=button", "xpath=//button"])(
+    "resolves %s without materializing a frontend DOM tree",
+    async (selector) => {
+      const { frame, send } = createFrame("root");
+      await expect(frame.locator(selector).resolveNode()).resolves.toEqual({
+        objectId: "node",
+        nodeId: null,
+      });
+      expect(send).not.toHaveBeenCalledWith("DOM.requestNode", expect.anything());
+    },
+  );
+
+  it("preserves frontend node IDs for low-level resolver callers", async () => {
+    const { frame, send } = createFrame("root");
+    await expect(
+      new FrameSelectorResolver(frame).resolveFirst({ kind: "css", value: "button" }),
+    ).resolves.toEqual({ objectId: "node", nodeId: 1 });
+    expect(send).toHaveBeenCalledWith("DOM.requestNode", { objectId: "node" });
+  });
+
+  it.each(["button", "text=button", "xpath=//button"])(
+    "resolves only the requested index for %s",
+    async (selector) => {
+      const { frame, send, respond } = createFrame("root");
+      let next = 0;
+      send.mockImplementation((method, params) =>
+        method === "Runtime.evaluate"
+          ? Promise.resolve({ result: { objectId: `node-${++next}` } })
+          : respond(method, params),
+      );
+      const progress = new Progress("resolve", 1000);
+      try {
+        await expect(frame.locator(selector).nth(10_000).resolveNode(progress)).resolves.toEqual({
+          objectId: "node-1",
+          nodeId: null,
+        });
+        const evaluations = send.mock.calls.filter(([method]) => method === "Runtime.evaluate");
+        expect(evaluations).toHaveLength(1);
+        expect((evaluations[0][1] as { expression: string }).expression).toContain("10000");
+        expect(send.mock.calls.filter(([method]) => method === "Runtime.releaseObject")).toEqual(
+          [],
+        );
+        expect(send).not.toHaveBeenCalledWith("DOM.requestNode", expect.anything());
+      } finally {
+        progress.dispose();
+      }
+    },
+  );
+
+  it("returns all mask handles without node IDs", async () => {
+    const { frame, send, respond } = createFrame("root");
+    let next = 0;
+    send.mockImplementation((method, params) =>
+      method === "Runtime.evaluate"
+        ? Promise.resolve({ result: ++next <= 2 ? { objectId: `node-${next}` } : {} })
+        : respond(method, params),
+    );
+    const progress = new Progress("mask", 1000);
+    try {
+      await expect(frame.locator("button").resolveNodesForMask(progress)).resolves.toEqual([
+        { objectId: "node-1", nodeId: null },
+        { objectId: "node-2", nodeId: null },
+      ]);
+      expect(send).not.toHaveBeenCalledWith("DOM.requestNode", expect.anything());
+    } finally {
+      progress.dispose();
+    }
+  });
+});
+
 describe("locator resolution deadlines", () => {
   const contexts: Progress[] = [];
   const createProgress = (timeout = 100) => {
@@ -206,7 +280,12 @@ describe("locator resolution deadlines", () => {
             root.frame,
             "iframe",
           ).resolveFrame(context)
-        : root.frame.locator("button").resolveNode(context);
+        : method === "DOM.requestNode"
+          ? new FrameSelectorResolver(root.frame).resolveFirst(
+              { kind: "css", value: "button" },
+              context,
+            )
+          : root.frame.locator("button").resolveNode(context);
     const rejected = expect(pending).rejects.toThrow(TimeoutError);
     await vi.advanceTimersByTimeAsync(40);
     await rejected;
@@ -235,7 +314,7 @@ describe("locator resolution deadlines", () => {
           : Promise.resolve({ result: { objectId: "first" } })
         : respond(method, params),
     );
-    const pending = frame.locator("button").nth(1).resolveNode(createProgress());
+    const pending = frame.locator("button").resolveNodesForMask(createProgress());
     const rejected = expect(pending).rejects.toThrow(TimeoutError);
     await vi.advanceTimersByTimeAsync(100);
     await rejected;
@@ -245,32 +324,15 @@ describe("locator resolution deadlines", () => {
     expect(send).toHaveBeenCalledWith("Runtime.releaseObject", { objectId: "second" });
   });
 
-  it("reports expiry without waiting for a second stalled cleanup", async () => {
-    const { frame, send, respond } = createFrame("root");
+  it("returns null for invalid indices without acquiring remote handles", async () => {
+    const { frame, send } = createFrame("root");
     const resolver = frame.locator("button").selectorResolver;
-    vi.spyOn(resolver, "resolveAll").mockResolvedValue([
-      { objectId: "unselected", nodeId: 1 },
-      { objectId: "selected", nodeId: 2 },
-    ]);
-    const gate = deferred();
-    send.mockImplementation((method, params) =>
-      method === "Runtime.releaseObject" ? gate.promise : respond(method, params),
-    );
-    const settled = vi.fn();
-    const pending = resolver.resolveAtIndex({ kind: "css", value: "button" }, 1, createProgress());
-    void pending.then(settled, settled);
-
-    try {
-      await vi.advanceTimersByTimeAsync(1000);
-      expect(settled).toHaveBeenCalledExactlyOnceWith(expect.any(TimeoutError));
-      expect(send.mock.calls.filter(([method]) => method === "Runtime.releaseObject")).toEqual([
-        ["Runtime.releaseObject", { objectId: "unselected" }],
-        ["Runtime.releaseObject", { objectId: "selected" }],
-      ]);
-    } finally {
-      gate.resolve({});
-      await vi.advanceTimersByTimeAsync(0);
+    for (const index of [-1, 0.5, Infinity, NaN]) {
+      await expect(
+        resolver.resolveAtIndex({ kind: "css", value: "button" }, index),
+      ).resolves.toBeNull();
     }
+    expect(send).not.toHaveBeenCalled();
   });
 
   it("removes the main-world listener & timer when its caller expires", async () => {
@@ -331,7 +393,7 @@ describe("locator resolution deadlines", () => {
         : respond(method, params),
     );
     const pending = runWithProgress({ name: "resolve", timeout: 100 }, (context) =>
-      frame.locator("button").resolveNode(context),
+      new FrameSelectorResolver(frame).resolveFirst({ kind: "css", value: "button" }, context),
     );
     const rejected = expect(pending).rejects.toThrow(TimeoutError);
     await vi.advanceTimersByTimeAsync(100);
@@ -428,7 +490,12 @@ describe("locator resolution deadlines", () => {
               root.frame,
               "iframe",
             ).resolveFrame(progress)
-          : root.frame.locator("button").resolveNode(progress);
+          : command === "DOM.requestNode"
+            ? new FrameSelectorResolver(root.frame).resolveFirst(
+                { kind: "css", value: "button" },
+                progress,
+              )
+            : root.frame.locator("button").resolveNode(progress);
       await expect(pending).rejects.toBe(closed);
     },
   );