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(""); + }); +});