feat(pi): compact snapshot output to cut context

packages/integrations/pi/extensions/snapshot-compaction.ts

@@ -0,0 +1,124 @@
+/**
+ * Context reduction for the Stagehand accessibility tree.
+ *
+ * The raw tree is 3-5x larger than a caller needs: most lines are anonymous
+ * layout containers (`div`, `listitem`, `tbody`, `scrollable`, ...) that cannot
+ * be clicked or filled, most `StaticText` lines are single-token fragments
+ * whose content already appears in an ancestor's accessible name, and
+ * indentation grows without bound with nesting depth.
+ *
+ * Keeping only semantic/actionable nodes and re-indenting by kept ancestors
+ * preserves every bracketed ID that `run` actions can use, because the facade
+ * stores the full xpath map regardless of what is rendered here. Measured on
+ * real pages, the compact tree is 21-34% of the raw tree.
+ */
+
+/** Roles worth keeping: everything a caller can act on or navigate by. */
+export const SEMANTIC_SNAPSHOT_ROLES: ReadonlySet<string> = new Set([
+  "RootWebArea",
+  "alert",
+  "button",
+  "checkbox",
+  "combobox",
+  "dialog",
+  "heading",
+  "image",
+  "img",
+  "link",
+  "listbox",
+  "menuitem",
+  "menuitemcheckbox",
+  "menuitemradio",
+  "option",
+  "radio",
+  "searchbox",
+  "slider",
+  "spinbutton",
+  "switch",
+  "tab",
+  "textbox",
+]);
+
+/**
+ * `StaticText` lines below this length are almost always one-token fragments
+ * (`{`, `,`, `=`, a single word from a code block) whose text is already
+ * aggregated into an ancestor's name.
+ */
+export const MIN_STATIC_TEXT_LENGTH = 12;
+
+/** Names longer than this are truncated so a single node cannot dominate the tree. */
+export const MAX_NAME_LENGTH = 160;
+
+export type CompactSnapshotOptions = {
+  /** Hard cap on the returned text; longer trees are cut with a marker line. */
+  maxChars?: number;
+};
+
+export type SnapshotLine = {
+  /** Nesting depth measured in leading whitespace characters. */
+  indent: number;
+  /** Bracketed snapshot ID, e.g. `1-42`. */
+  id: string;
+  /** Accessibility role, e.g. `button` or `scrollable, html`. */
+  role: string;
+  /** Accessible name, empty when the node has none. */
+  name: string;
+};
+
+/**
+ * Parse one accessibility tree line. Returns undefined for lines that carry no
+ * node — multi-line accessible names produce bare continuation lines, which a
+ * compact tree drops.
+ */
+export function parseSnapshotLine(line: string): SnapshotLine | undefined {
+  if (!line.trim()) return undefined;
+  const match = /^(\s*)\[([^\]]+)\]\s?(.*)$/u.exec(line);
+  if (!match) return undefined;
+  const indent = match[1] ?? "";
+  const id = match[2] ?? "";
+  const body = match[3] ?? "";
+  const separator = body.indexOf(": ");
+  const role = separator === -1 ? body : body.slice(0, separator);
+  const name = separator === -1 ? "" : body.slice(separator + 2);
+  return { indent: indent.length, id, role, name };
+}
+
+export function isKeptSnapshotLine(line: SnapshotLine): boolean {
+  if (SEMANTIC_SNAPSHOT_ROLES.has(line.role)) return true;
+  return line.role === "StaticText" && line.name.trim().length >= MIN_STATIC_TEXT_LENGTH;
+}
+
+/**
+ * Reduce a formatted accessibility tree to the nodes a caller can act on.
+ *
+ * Kept lines are re-indented by their kept ancestors, so two spaces per level
+ * after filtering. Node IDs are copied verbatim: this only removes lines, never
+ * rewrites a node, so IDs stay valid for `run` actions.
+ */
+export function compactSnapshotTree(tree: string, options: CompactSnapshotOptions = {}): string {
+  const out: string[] = [];
+  const ancestorIndents: number[] = [];
+  for (const line of tree.split("\n")) {
+    const parsed = parseSnapshotLine(line);
+    if (!parsed || !isKeptSnapshotLine(parsed)) continue;
+    const name =
+      parsed.name.length > MAX_NAME_LENGTH
+        ? `${parsed.name.slice(0, MAX_NAME_LENGTH)}…`
+        : parsed.name;
+    for (;;) {
+      const top = ancestorIndents[ancestorIndents.length - 1];
+      if (top === undefined || top < parsed.indent) break;
+      ancestorIndents.pop();
+    }
+    const prefix = "  ".repeat(ancestorIndents.length);
+    ancestorIndents.push(parsed.indent);
+    out.push(`${prefix}[${parsed.id}] ${parsed.role}${name ? `: ${name}` : ""}`);
+  }
+
+  let text = out.join("\n");
+  const { maxChars } = options;
+  if (maxChars !== undefined && text.length > maxChars) {
+    text = `${text.slice(0, maxChars)}\n… [compact snapshot truncated at ${maxChars} chars]`;
+  }
+  return text;
+}

packages/integrations/pi/extensions/stagehand.ts

@@ -27,6 +27,8 @@ import {
   stagehandFacadeConfigFromEnv,
 } from "@browserbasehq/stagehand-integrations/facade";
 
+import { compactSnapshotTree } from "./snapshot-compaction.js";
+
 type FacadeResources = {
   browser: StagehandBrowser;
   stagehand: Stagehand;
@@ -51,6 +53,17 @@ const runParameters = Type.Object({
 
 const snapshotParameters = Type.Object({
   includeIframes: Type.Optional(Type.Boolean()),
+  compact: Type.Optional(
+    Type.Boolean({
+      description:
+        "Drop layout containers and fragmented text so the tree costs far fewer tokens. Default true; false returns the raw accessibility tree.",
+    }),
+  ),
+  maxChars: Type.Optional(
+    Type.Number({
+      description: "Hard cap on the returned tree size; longer trees are cut with a marker.",
+    }),
+  ),
 });
 
 const screenshotParameters = Type.Object({
@@ -110,7 +123,10 @@ export default function stagehandExtension(pi: ExtensionAPI) {
     label: "Stagehand run",
     description: RUN_TOOL_DESCRIPTION,
     promptSnippet: "run: execute a JavaScript workflow or snapshot-ID actions in the browser",
-    promptGuidelines: [FACADE_AGENT_INSTRUCTIONS],
+    promptGuidelines: [
+      FACADE_AGENT_INSTRUCTIONS,
+      "Prefer one `run` call that returns only the values you need, e.g. `return await page.evaluate(() => [...document.querySelectorAll('h2')].map((h) => h.textContent))`, over a snapshot-then-read sequence. Snapshot payloads dominate the context; extracted values do not.",
+    ],
     parameters: runParameters,
     executionMode: "sequential",
     async execute(_toolCallId, params) {
@@ -131,16 +147,21 @@ export default function stagehandExtension(pi: ExtensionAPI) {
     name: "snapshot",
     label: "Stagehand snapshot",
     description: SNAPSHOT_TOOL_DESCRIPTION,
-    promptSnippet: "snapshot: inspect the active page and hydrate bracketed element IDs",
+    promptSnippet:
+      "snapshot: list clickable/fillable elements with bracketed IDs (compact by default)",
+    promptGuidelines: [
+      "Use `snapshot` only to discover bracketed element IDs; it is compact by default and drops anonymous layout nodes. Pass `compact: false` only when you genuinely need the raw tree. To read long text or structured data, use `run` and return only the fields you need instead of dumping page content.",
+    ],
     parameters: snapshotParameters,
     executionMode: "sequential",
     async execute(_toolCallId, params) {
-      const input = SnapshotInputSchema.parse(params);
+      const { compact = true, maxChars, ...snapshotInput } = params;
+      const input = SnapshotInputSchema.parse(snapshotInput);
       const tools = await facadeTools();
       const tree = await tools.snapshot(input);
       return {
-        content: [{ type: "text", text: tree }],
-        details: {},
+        content: [{ type: "text", text: compact ? compactSnapshotTree(tree, { maxChars }) : tree }],
+        details: { compact, rawChars: tree.length },
       } satisfies AgentToolResult<unknown>;
     },
   });

packages/integrations/pi/tests/extension.test.ts

@@ -6,7 +6,12 @@ import { describe, expect, it } from "vitest";
 
 import stagehandExtension from "../extensions/stagehand.js";
 
-type Registered = { name: string; description: string; promptGuidelines?: string[] };
+type Registered = {
+  name: string;
+  description: string;
+  promptGuidelines?: string[];
+  parameters?: { properties?: Record<string, unknown> };
+};
 
 function registeredTools(): Registered[] {
   const tools: Registered[] = [];
@@ -33,7 +38,16 @@ describe("pi stagehand extension", () => {
 
   it("forwards the canonical agent instructions as guidelines", () => {
     const run = registeredTools().find((tool) => tool.name === "run");
-    expect(run?.promptGuidelines).toEqual([FACADE_AGENT_INSTRUCTIONS]);
+    expect(run?.promptGuidelines?.[0]).toBe(FACADE_AGENT_INSTRUCTIONS);
+  });
+
+  it("exposes the snapshot context-reduction options", () => {
+    const snapshot = registeredTools().find((tool) => tool.name === "snapshot");
+    expect(Object.keys(snapshot?.parameters?.properties ?? {})).toEqual([
+      "includeIframes",
+      "compact",
+      "maxChars",
+    ]);
   });
 
   it("does not launch a browser at registration time", () => {

packages/integrations/pi/tests/snapshot-compaction.test.ts

@@ -0,0 +1,117 @@
+import { describe, expect, it } from "vitest";
+
+import {
+  MAX_NAME_LENGTH,
+  compactSnapshotTree,
+  isKeptSnapshotLine,
+  parseSnapshotLine,
+} from "../extensions/snapshot-compaction.js";
+
+// A trimmed slice of a real tree: a root, unnamed layout containers, a heading
+// whose text is duplicated by its StaticText child, and a link.
+const REAL_TREE = [
+  "[0-2] RootWebArea: Example Domain",
+  "  [0-4] scrollable, html",
+  "    [0-6] div",
+  "      [0-8] heading: Example Domain",
+  "        [0-9] StaticText: Example Domain",
+  "        [0-10] StaticText: {",
+  "      [0-11] link: Learn more",
+  "",
+].join("\n");
+
+const ids = (tree: string) =>
+  tree
+    .split("\n")
+    .flatMap((line) => parseSnapshotLine(line) ?? [])
+    .map((line) => line.id);
+
+describe("parseSnapshotLine", () => {
+  it("splits indent, id, role, and name", () => {
+    expect(parseSnapshotLine("    [1-42] link: Docs")).toEqual({
+      indent: 4,
+      id: "1-42",
+      role: "link",
+      name: "Docs",
+    });
+  });
+
+  it("keeps roles without a name", () => {
+    expect(parseSnapshotLine("  [1-43] textbox")).toEqual({
+      indent: 2,
+      id: "1-43",
+      role: "textbox",
+      name: "",
+    });
+  });
+
+  it("ignores blank lines and continuation lines of a multi-line name", () => {
+    expect(parseSnapshotLine("   ")).toBeUndefined();
+    expect(parseSnapshotLine("just install")).toBeUndefined();
+  });
+});
+
+describe("isKeptSnapshotLine", () => {
+  it("keeps actionable roles and drops anonymous containers", () => {
+    expect(isKeptSnapshotLine({ indent: 0, id: "1", role: "button", name: "Save" })).toBe(true);
+    expect(isKeptSnapshotLine({ indent: 0, id: "1", role: "div", name: "" })).toBe(false);
+    expect(isKeptSnapshotLine({ indent: 0, id: "1", role: "listitem", name: "" })).toBe(false);
+  });
+
+  it("keeps substantial StaticText and drops fragments", () => {
+    expect(
+      isKeptSnapshotLine({ indent: 0, id: "1", role: "StaticText", name: "Example Domain" }),
+    ).toBe(true);
+    expect(isKeptSnapshotLine({ indent: 0, id: "1", role: "StaticText", name: "{" })).toBe(false);
+  });
+});
+
+describe("compactSnapshotTree", () => {
+  it("keeps only actionable nodes and re-indents them by kept ancestors", () => {
+    expect(compactSnapshotTree(REAL_TREE)).toBe(
+      [
+        "[0-2] RootWebArea: Example Domain",
+        "  [0-8] heading: Example Domain",
+        "    [0-9] StaticText: Example Domain",
+        "  [0-11] link: Learn more",
+      ].join("\n"),
+    );
+  });
+
+  it("shrinks a real tree by more than half", () => {
+    const raw = Array.from({ length: 200 }, (_, index) =>
+      index % 10 === 0
+        ? `  [0-${index}] div`
+        : `      [0-${index}] link: item ${index} with a reasonably long label`,
+    ).join("\n");
+    expect(compactSnapshotTree(raw).length).toBeLessThan(raw.length);
+  });
+
+  it("copies node ids verbatim so run actions stay valid", () => {
+    const compact = compactSnapshotTree(REAL_TREE);
+    expect(ids(compact)).toEqual(["0-2", "0-8", "0-9", "0-11"]);
+    for (const id of ids(compact)) expect(ids(REAL_TREE)).toContain(id);
+  });
+
+  it("truncates long names", () => {
+    const long = "x".repeat(MAX_NAME_LENGTH + 40);
+    const compact = compactSnapshotTree(`[1-1] heading: ${long}`);
+    expect(compact).toBe(`[1-1] heading: ${"x".repeat(MAX_NAME_LENGTH)}…`);
+  });
+
+  it("caps the output when maxChars is set", () => {
+    const compact = compactSnapshotTree(REAL_TREE, { maxChars: 40 });
+    expect(compact.startsWith("[0-2] RootWebArea")).toBe(true);
+    expect(compact.endsWith("… [compact snapshot truncated at 40 chars]")).toBe(true);
+  });
+
+  it("returns the input unchanged in size when nothing is filtered", () => {
+    const onlyKept = "[0-1] RootWebArea: Title\n  [0-2] button: Go";
+    expect(compactSnapshotTree(onlyKept)).toBe(onlyKept);
+  });
+
+  it("returns an empty string for an empty or fully filtered tree", () => {
+    expect(compactSnapshotTree("")).toBe("");
+    expect(compactSnapshotTree("[0-1] div\n  [0-2] listitem")).toBe("");
+  });
+});