[feat]: let extract() return element references

MA
12h ago
+709-25

Closes #3406

why

extract() returns data and observe() returns elements, but no call returns both. #3406 asks for "these values, and the element each one came from", so code can pick a value deterministically and then act on its element without the LLM doing the comparison.

what changed

Adds elementRef(), a schema helper that resolves to the element a value came from, the same way z.url() resolves to a link's URL today:

const { data } = await stagehand.extract(
  "extract every account number in the list",
  z.object({ accounts: z.array(z.object({ number: z.string(), element: elementRef() })) }),
);
const match = data.accounts.find((a) => normalize(a.number) === target);
if (match?.element) await page.locator(match.element).click();
  • SDK (sdk-ts): elementRef(description?) returns z.string() tagged with "x-stagehand": "element" in its JSON Schema.
  • Extension: transformSchema swaps marked fields for element-ID fields, alongside URL fields. After inference, extractService resolves them through the snapshot's XPath map to xpath= selectors, the same format as observe(). The marker is never sent to the model provider.
  • Supported shapes: objects, arrays, tuples, unions, $refs, a non-object root, and optional, nullable, default, catch and readonly wrappers. An element field inside a record or recursive schema throws a clear error before inference instead of returning free text.
  • Unresolved IDs become "", like URL fields, and log a warning, like observe().
  • Prompt: one sentence about element IDs is added to the extract system prompt, only when the schema has element fields. Other extract() calls get the same prompt as before.
  • Docs: an "Element references" section in basics/extract.mdx (TypeScript, and Python using the raw marker) and a pointer from the reference page.
  • Changeset: minor for the extension, the TS SDK and the Python SDK (Python ships the extension and can send the raw marker).

notes for reviewers

  • Protocol version not bumped. schema is already free-form JSON, so there's no wire schema change. However, a new SDK talking to an older extension (for example one preinstalled on a Browserbase session) would treat element fields as plain strings and return unresolved model text. Since the extension ships inside each SDK, I left the protocol alone, but I'm happy to add a protocol minor bump if you'd prefer new clients to reject older extensions.
  • Marker repeated in two places. The "x-stagehand": "element" string is defined in both extension/utils.ts and sdk-ts/src/extractSchema.ts (cross-referenced in comments), to keep the protocol package untouched. Happy to move it into packages/protocol instead.
  • Behaviour change on an existing path. URL fields behind .default() or .catch() (e.g. z.url().default("")) were previously not resolved and came back as model text. They now resolve like other URL fields, and their descriptions are kept. Both wrappers are covered by a new service-level test.
  • No Python helper yet. Pydantic emits the raw marker via Field(json_schema_extra={"x-stagehand": "element"}), and the Python SDK forwards it unchanged (checked on the serialized schema; a full Python extract() call wasn't run). A named Element helper could follow if you want parity. The Go SDK wasn't checked.
  • Possible overlap with #3363 and #3050, which also touch the URL-field code. I kept the exported names (injectUrls, transformUrlStringsToNumericIds) unchanged to limit conflicts.

test plan

  • just fmt, just check
  • Extension and SDK unit suites pass. New tests cover nested, array, tuple, union, $ref and wrapper shapes; descriptions on wrappers; the record and recursion errors; unresolved-ID warnings; the conditional system prompt; and URL fields behind .default() / .catch().
  • just test: everything passed except two timing flakes under full parallel load, in suites this change doesn't touch: integrations/core browser tests (Chrome shutdown in afterAll) and scripts/release/scoped-release.test.ts. Both pass when run alone.
  • Live run with local Chrome and anthropic/claude-sonnet-5-5: extracted a 4-row account list (numbers that only match after normalization), all rows resolved to selectors, and clicking the code-selected match hit the right row.

Summary by cubic

Closes #3406. extract() now returns data and, for fields you mark, the element each value came from, surfaced as an xpath= selector you can pass to page.locator(). Before, extract() and observe() each returned only half of this.

Mark a field with elementRef() in TypeScript, or "x-stagehand": "element" in JSON Schema from other SDKs. Unresolved IDs become "", matching URL fields, and log a warning.

Notes

  • The prompt is extended with element-ID guidance only when the schema has element fields; other extract() calls keep the previous prompt.
  • Element fields are supported inside objects, arrays, tuples, unions, $refs, and optional, nullable, default, catch, and readonly wrappers. Records and recursive schemas throw a clear error before inference.
  • The protocol version is unchanged; a new SDK talking to an older preinstalled extension returns plain model text for element fields.
  • URL fields behind .default() or .catch() previously came back as model text; they now resolve like other URL fields.
  • No named Python helper yet; Python uses the raw JSON Schema marker.