[v2]: retry on max depth exceeded

.changeset/funny-terms-yell.md

@@ -0,0 +1,5 @@
+---
+"@browserbasehq/stagehand": patch
+---
+
+Retry DOM.getDocument on max depth exceeded with exponential backoff

evals/deterministic/tests/page/addInitScript.test.ts

@@ -39,32 +39,4 @@ test.describe("StagehandPage - addInitScript", () => {
 
     await stagehand.close();
   });
-
-  test("checks if init scripts are re-added and available even if they've been deleted", async () => {
-    const stagehand = new Stagehand(StagehandConfig);
-    await stagehand.init();
-
-    const page = stagehand.page;
-    await page.goto(
-      "https://browserbase.github.io/stagehand-eval-sites/sites/aigrant/",
-    );
-
-    // delete the __stagehandInjected flag, and delete the
-    // getScrollableElementXpaths function
-    await page.evaluate(() => {
-      delete window.getScrollableElementXpaths;
-      delete window.__stagehandInjected;
-    });
-
-    // attempt to call the getScrollableElementXpaths function
-    // which we previously deleted. page.evaluate should realize
-    // its been deleted and re-inject it
-    const xpaths = await page.evaluate(() => {
-      return window.getScrollableElementXpaths();
-    });
-
-    await stagehand.close();
-    // this is the only scrollable element on the page
-    expect(xpaths).toContain("/html");
-  });
 });

lib/StagehandContext.ts

@@ -8,13 +8,23 @@ import { StagehandPage } from "./StagehandPage";
 import { Page } from "../types/page";
 import { EnhancedContext } from "../types/context";
 import { Protocol } from "devtools-protocol";
+import { scriptContent } from "./dom/build/scriptContent";
+
+const stagehandInitScript = `
+  if (!window.__stagehandInjected) {
+    window.__stagehandInjected = true;
+    ${scriptContent}
+  }
+`;
 
 export class StagehandContext {
   private readonly stagehand: Stagehand;
   private readonly intContext: EnhancedContext;
   private pageMap: WeakMap<PlaywrightPage, StagehandPage>;
   private activeStagehandPage: StagehandPage | null = null;
   private readonly frameIdMap: Map<string, StagehandPage> = new Map();
+  private static readonly contextsWithInitScript =
+    new WeakSet<PlaywrightContext>();
 
   private constructor(context: PlaywrightContext, stagehand: Stagehand) {
     this.stagehand = stagehand;
@@ -81,6 +91,10 @@ export class StagehandContext {
     context: PlaywrightContext,
     stagehand: Stagehand,
   ): Promise<StagehandContext> {
+    if (!StagehandContext.contextsWithInitScript.has(context)) {
+      await context.addInitScript({ content: stagehandInitScript });
+      StagehandContext.contextsWithInitScript.add(context);
+    }
     const instance = new StagehandContext(context, stagehand);
     context.on("page", async (pwPage) => {
       await instance.handleNewPlaywrightPage(pwPage);

lib/StagehandPage.ts

@@ -32,7 +32,6 @@ import {
   ExperimentalApiConflictError,
 } from "../types/stagehandErrors";
 import { StagehandAPIError } from "@/types/stagehandApiErrors";
-import { scriptContent } from "@/lib/dom/build/scriptContent";
 import type { Protocol } from "devtools-protocol";
 
 async function getCurrentRootFrameId(session: CDPSession): Promise<string> {
@@ -165,37 +164,6 @@ export class StagehandPage {
     this.fidOrdinals = new Map([[undefined, 0]]);
   }
 
-  private async ensureStagehandScript(): Promise<void> {
-    try {
-      const injected = await this.rawPage.evaluate(
-        () => !!window.__stagehandInjected,
-      );
-
-      if (injected) return;
-
-      const guardedScript = `if (!window.__stagehandInjected) { \
-window.__stagehandInjected = true; \
-${scriptContent} \
-}`;
-
-      await this.rawPage.addInitScript({ content: guardedScript });
-      await this.rawPage.evaluate(guardedScript);
-    } catch (err) {
-      if (!this.stagehand.isClosed) {
-        this.stagehand.log({
-          category: "dom",
-          message: "Failed to inject Stagehand helper script",
-          level: 1,
-          auxiliary: {
-            error: { value: (err as Error).message, type: "string" },
-            trace: { value: (err as Error).stack, type: "string" },
-          },
-        });
-        throw err;
-      }
-    }
-  }
-
   /** Register the custom selector engine that pierces open/closed shadow roots. */
   private async ensureStagehandSelectorEngine(): Promise<void> {
     const registerFn = () => {
@@ -387,8 +355,6 @@ ${scriptContent} \
             prop === "$$eval"
           ) {
             return async (...args: unknown[]) => {
-              // Make sure helpers exist
-              await this.ensureStagehandScript();
               return (value as (...a: unknown[]) => unknown).apply(
                 target,
                 args,

lib/a11y/utils.ts

@@ -32,6 +32,154 @@ const NBSP_CHARS = new Set<number>([0x00a0, 0x202f, 0x2007, 0xfeff]);
 
 const WORLD_NAME = "stagehand-world";
 
+/** Order of depths we try when DOM.getDocument blows the CBOR stack. */
+const DOM_DEPTH_ATTEMPTS = [-1, 256, 128, 64, 32, 16, 8, 4, 2, 1];
+const DESCRIBE_DEPTH_ATTEMPTS = [-1, 64, 32, 16, 8, 4, 2, 1];
+
+function isCborStackError(message: string): boolean {
+  return message.includes("CBOR: stack limit exceeded");
+}
+
+/** The CDP payload only gives us child stubs when depth was capped. */
+function shouldExpandNode(node: DOMNode): boolean {
+  const declaredChildren = node.childNodeCount ?? 0;
+  const realizedChildren = node.children?.length ?? 0;
+  return declaredChildren > realizedChildren;
+}
+
+/** Replace the shallow node instance with the fully described version. */
+function mergeDomNodes(target: DOMNode, source: DOMNode): void {
+  target.childNodeCount = source.childNodeCount ?? target.childNodeCount;
+  target.children = source.children ?? target.children;
+  target.shadowRoots = source.shadowRoots ?? target.shadowRoots;
+  target.contentDocument = source.contentDocument ?? target.contentDocument;
+}
+
+/** Yield every nested DOMNode collection so we can traverse them uniformly. */
+function collectDomTraversalTargets(node: DOMNode): DOMNode[] {
+  const targets: DOMNode[] = [];
+  if (node.children) targets.push(...node.children);
+  if (node.shadowRoots) targets.push(...node.shadowRoots);
+  if (node.contentDocument) targets.push(node.contentDocument);
+  return targets;
+}
+
+/**
+ * Walk the partial DOM tree and call DOM.describeNode for any node whose
+ * childNodeCount indicates we didn't receive all of its children.
+ */
+async function hydrateDomTree(
+  session: CDPSession,
+  root: DOMNode,
+): Promise<number> {
+  const stack: DOMNode[] = [root];
+  const expandedNodeIds = new Set<number>();
+  const expandedBackendIds = new Set<number>();
+  let describeCalls = 0;
+
+  while (stack.length) {
+    const node = stack.pop()!;
+    const nodeId =
+      typeof node.nodeId === "number" && node.nodeId > 0
+        ? node.nodeId
+        : undefined;
+    const backendId =
+      typeof node.backendNodeId === "number" && node.backendNodeId > 0
+        ? node.backendNodeId
+        : undefined;
+
+    const seenByNode = nodeId ? expandedNodeIds.has(nodeId) : false;
+    const seenByBackend =
+      !nodeId && backendId ? expandedBackendIds.has(backendId) : false;
+    if (seenByNode || seenByBackend) continue;
+    if (nodeId) expandedNodeIds.add(nodeId);
+    else if (backendId) expandedBackendIds.add(backendId);
+
+    const needsExpansion = shouldExpandNode(node);
+    if (needsExpansion && (nodeId || backendId)) {
+      const describeParamsBase = nodeId
+        ? { nodeId }
+        : { backendNodeId: backendId! };
+      let expanded = false;
+      for (const depth of DESCRIBE_DEPTH_ATTEMPTS) {
+        try {
+          const described = (await session.send("DOM.describeNode", {
+            ...describeParamsBase,
+            depth,
+            pierce: true,
+          })) as { node: DOMNode };
+          mergeDomNodes(node, described.node);
+          if (!nodeId && described.node.nodeId && described.node.nodeId > 0) {
+            node.nodeId = described.node.nodeId;
+            expandedNodeIds.add(described.node.nodeId);
+          }
+          describeCalls++;
+          expanded = true;
+          break;
+        } catch (err) {
+          const message = err instanceof Error ? err.message : String(err);
+          if (isCborStackError(message)) {
+            continue;
+          }
+          const identifier = nodeId ?? backendId ?? "unknown";
+          throw new StagehandDomProcessError(
+            `Failed to expand DOM node ${identifier}: ${String(err)}`,
+          );
+        }
+      }
+      if (!expanded) {
+        const identifier = nodeId ?? backendId ?? "unknown";
+        throw new StagehandDomProcessError(
+          `Unable to expand DOM node ${identifier} after describeNode depth retries`,
+        );
+      }
+    }
+
+    for (const child of collectDomTraversalTargets(node)) {
+      stack.push(child);
+    }
+  }
+
+  return describeCalls;
+}
+
+/**
+ * Attempt DOM.getDocument with progressively smaller depths, rehydrating the
+ * missing children after each successful response.
+ */
+async function getDomTreeWithFallback(session: CDPSession): Promise<DOMNode> {
+  let lastCborMessage = "";
+
+  for (const depth of DOM_DEPTH_ATTEMPTS) {
+    try {
+      const params = { depth, pierce: true };
+
+      const { root } = (await session.send("DOM.getDocument", params)) as {
+        root: DOMNode;
+      };
+
+      if (depth !== -1) {
+        await hydrateDomTree(session, root);
+      }
+
+      return root;
+    } catch (err) {
+      const message = err instanceof Error ? err.message : String(err);
+      if (isCborStackError(message)) {
+        lastCborMessage = message;
+        continue;
+      }
+      throw err;
+    }
+  }
+
+  throw new StagehandDomProcessError(
+    lastCborMessage
+      ? `CDP DOM.getDocument failed after adaptive depth retries: ${lastCborMessage}`
+      : "CDP DOM.getDocument failed after adaptive depth retries.",
+  );
+}
+
 /**
  * Clean a string by removing private-use unicode characters, normalizing whitespace,
  * and trimming the result.
@@ -148,11 +296,8 @@ export async function buildBackendIdMaps(
   );
 
   try {
-    // 1. full DOM tree
-    const { root } = (await session.send("DOM.getDocument", {
-      depth: -1,
-      pierce: true,
-    })) as { root: DOMNode };
+    // 1. full DOM tree (with adaptive fallback)
+    const root = await getDomTreeWithFallback(session);
 
     // 2. pick start node + root frame-id
     let startNode: DOMNode = root;

lib/dom/global.d.ts

@@ -10,9 +10,6 @@ declare global {
     __playwright?: unknown;
     __pw_manual?: unknown;
     __PW_inspect?: unknown;
-    getScrollableElementXpaths: (topN?: number) => Promise<string[]>;
-    getNodeFromXpath: (xpath: string) => Node | null;
-    waitForElementScrollEnd: (element: HTMLElement) => Promise<void>;
     readonly __stagehand__?: StagehandBackdoor;
   }
 }

lib/dom/index.ts

@@ -1,2 +1 @@
-export * from "./process";
-export * from "./utils";
+import "./process";