From 346b3f807da40175d9ef2164fad78e2b10880c72 Mon Sep 17 00:00:00 2001 From: TQuentinD Date: Sat, 12 Sep 2026 14:23:07 +0200 Subject: [PATCH] feat(pi): compact snapshot output to cut context The pi integration returned the raw accessibility tree from `snapshot`. On real pages most of it is noise: anonymous layout containers (`div`, `listitem`, `tbody`, `scrollable`, ...) that cannot be clicked or filled, `StaticText` lines that are single-token fragments already aggregated into an ancestor's name, and indentation that grows with nesting depth. In an agent loop that payload is replayed on every later request, so it is paid for repeatedly. `snapshot` now keeps only semantic/actionable roles and re-indents by the kept ancestors. `compact: false` returns the raw tree and `maxChars` hard-caps the payload. Node IDs are copied verbatim and the facade keeps the full xpath map, so `run` actions against compact IDs keep working. Measured on real pages, the compact tree is 21-34% of the raw tree: news.ycombinator.com 34027 -> 7314 chars, github.com/browserbase/stagehand 63080 -> 20274, books.toscrape.com 21553 -> 7436, example.com 316 -> 218. --- .../pi/extensions/snapshot-compaction.ts | 124 ++++++++++++++++++ .../integrations/pi/extensions/stagehand.ts | 31 ++++- .../integrations/pi/tests/extension.test.ts | 18 ++- .../pi/tests/snapshot-compaction.test.ts | 117 +++++++++++++++++ 4 files changed, 283 insertions(+), 7 deletions(-) create mode 100644 packages/integrations/pi/extensions/snapshot-compaction.ts create mode 100644 packages/integrations/pi/tests/snapshot-compaction.test.ts diff --git a/packages/integrations/pi/extensions/snapshot-compaction.ts b/packages/integrations/pi/extensions/snapshot-compaction.ts new file mode 100644 index 000000000..8d2b48166 --- /dev/null +++ b/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 = 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; +} diff --git a/packages/integrations/pi/extensions/stagehand.ts b/packages/integrations/pi/extensions/stagehand.ts index 9fba1deb1..0b261c79d 100644 --- a/packages/integrations/pi/extensions/stagehand.ts +++ b/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; }, }); diff --git a/packages/integrations/pi/tests/extension.test.ts b/packages/integrations/pi/tests/extension.test.ts index 4c896919d..31c7097fc 100644 --- a/packages/integrations/pi/tests/extension.test.ts +++ b/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 }; +}; 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", () => { diff --git a/packages/integrations/pi/tests/snapshot-compaction.test.ts b/packages/integrations/pi/tests/snapshot-compaction.test.ts new file mode 100644 index 000000000..b3d07d17a --- /dev/null +++ b/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(""); + }); +});