[feat]: expose timeout param in ts & protocol

packages/extension/controllers/locatorController.ts

@@ -1,6 +1,6 @@
 import type {
   LocatorClickParams,
-  LocatorDescriptor,
+  LocatorParams,
   LocatorFillParams,
   LocatorHighlightParams,
   LocatorScrollToParams,
@@ -23,42 +23,42 @@ export function createLocatorController(runtime: StagehandRuntime) {
     return runtime.locatorFill(params);
   }
 
-  async function hover(params: LocatorDescriptor, { logger }: HandlerContext) {
+  async function hover(params: LocatorParams, { logger }: HandlerContext) {
     logger.debug("locator.hover", {});
     return runtime.locatorHover(params);
   }
 
-  async function count(params: LocatorDescriptor, { logger }: HandlerContext) {
+  async function count(params: LocatorParams, { logger }: HandlerContext) {
     logger.debug("locator.count", {});
     return runtime.locatorCount(params);
   }
 
-  async function isChecked(params: LocatorDescriptor, { logger }: HandlerContext) {
+  async function isChecked(params: LocatorParams, { logger }: HandlerContext) {
     logger.debug("locator.is_checked", {});
     return runtime.locatorIsChecked(params);
   }
 
-  async function inputValue(params: LocatorDescriptor, { logger }: HandlerContext) {
+  async function inputValue(params: LocatorParams, { logger }: HandlerContext) {
     logger.debug("locator.input_value", {});
     return runtime.locatorInputValue(params);
   }
 
-  async function isVisible(params: LocatorDescriptor, { logger }: HandlerContext) {
+  async function isVisible(params: LocatorParams, { logger }: HandlerContext) {
     logger.debug("locator.is_visible", {});
     return runtime.locatorIsVisible(params);
   }
 
-  async function innerText(params: LocatorDescriptor, { logger }: HandlerContext) {
+  async function innerText(params: LocatorParams, { logger }: HandlerContext) {
     logger.debug("locator.inner_text", {});
     return runtime.locatorInnerText(params);
   }
 
-  async function innerHtml(params: LocatorDescriptor, { logger }: HandlerContext) {
+  async function innerHtml(params: LocatorParams, { logger }: HandlerContext) {
     logger.debug("locator.inner_html", {});
     return runtime.locatorInnerHtml(params);
   }
 
-  async function textContent(params: LocatorDescriptor, { logger }: HandlerContext) {
+  async function textContent(params: LocatorParams, { logger }: HandlerContext) {
     logger.debug("locator.text_content", {});
     return runtime.locatorTextContent(params);
   }
@@ -68,7 +68,7 @@ export function createLocatorController(runtime: StagehandRuntime) {
     return runtime.locatorScrollTo(params);
   }
 
-  async function centroid(params: LocatorDescriptor, { logger }: HandlerContext) {
+  async function centroid(params: LocatorParams, { logger }: HandlerContext) {
     logger.debug("locator.centroid", {});
     return runtime.locatorCentroid(params);
   }

packages/extension/runtime.ts

@@ -1,3 +1,5 @@
+import { DEFAULT_LOCATOR_TIMEOUT_MS } from "@browserbasehq/stagehand-protocol/schemas";
+import { runWithProgress, type Progress } from "./understudy/progress.js";
 import { ShadowRootEvaluationUnavailableError } from "./errors.js";
 import type {
   ClearCookieOptions,
@@ -33,6 +35,7 @@ import type {
   LocatorCentroidResult,
   LocatorCountResult,
   LocatorDescriptor,
+  LocatorParams,
   LocatorFillParams,
   LocatorFillResult,
   LocatorHighlightParams,
@@ -220,23 +223,30 @@ export type UnderstudyRuntimeClipboard = {
 };
 
 export type UnderstudyRuntimeLocator = {
-  click(options?: LocatorClickParams["options"]): Promise<void> | void;
-  hover(): Promise<void> | void;
-  fill(value: string): Promise<void> | void;
-  count(): Promise<number>;
-  isChecked(): Promise<boolean>;
-  inputValue(): Promise<string>;
-  isVisible(): Promise<boolean>;
-  innerText(): Promise<string>;
-  innerHtml(): Promise<string>;
-  textContent(): Promise<string>;
-  scrollTo(percent: LocatorScrollToParams["percent"]): Promise<void> | void;
-  centroid(): Promise<LocatorCentroidResult>;
-  highlight(options?: LocatorHighlightParams["options"]): Promise<void> | void;
-  sendClickEvent(options?: LocatorSendClickEventParams["options"]): Promise<void> | void;
-  type(text: string, options?: LocatorTypeParams["options"]): Promise<void> | void;
-  selectOption(values: LocatorSelectOptionParams["values"]): Promise<string[]>;
-  setInputFiles(files: SetInputFilesArgument): Promise<void>;
+  click(options?: LocatorClickParams["options"], progress?: Progress): Promise<void> | void;
+  hover(progress?: Progress): Promise<void> | void;
+  fill(value: string, progress?: Progress): Promise<void> | void;
+  count(progress?: Progress): Promise<number>;
+  isChecked(progress?: Progress): Promise<boolean>;
+  inputValue(progress?: Progress): Promise<string>;
+  isVisible(progress?: Progress): Promise<boolean>;
+  innerText(progress?: Progress): Promise<string>;
+  innerHtml(progress?: Progress): Promise<string>;
+  textContent(progress?: Progress): Promise<string>;
+  scrollTo(percent: LocatorScrollToParams["percent"], progress?: Progress): Promise<void> | void;
+  centroid(progress?: Progress): Promise<LocatorCentroidResult>;
+  highlight(options?: LocatorHighlightParams["options"], progress?: Progress): Promise<void> | void;
+  sendClickEvent(
+    options?: LocatorSendClickEventParams["options"],
+    progress?: Progress,
+  ): Promise<void> | void;
+  type(
+    text: string,
+    options?: LocatorTypeParams["options"],
+    progress?: Progress,
+  ): Promise<void> | void;
+  selectOption(values: LocatorSelectOptionParams["values"], progress?: Progress): Promise<string[]>;
+  setInputFiles(files: SetInputFilesArgument, progress?: Progress): Promise<void>;
   nth(index: number): UnderstudyRuntimeLocator;
 };
 
@@ -823,99 +833,144 @@ export class StagehandRuntime {
   }
 
   async locatorClick(params: LocatorClickParams): Promise<LocatorClickResult> {
-    await this.resolveLocator(params).click(params.options);
+    await this.runLocator("locator.click", params, (locator, progress) =>
+      locator.click(params.options, progress),
+    );
     return { clicked: true };
   }
 
-  async locatorHover(params: LocatorDescriptor): Promise<LocatorHoverResult> {
-    await this.resolveLocator(params).hover();
+  async locatorHover(params: LocatorParams): Promise<LocatorHoverResult> {
+    await this.runLocator("locator.hover", params, (locator, progress) => locator.hover(progress));
     return { hovered: true };
   }
 
   async locatorFill(params: LocatorFillParams): Promise<LocatorFillResult> {
-    await this.resolveLocator(params).fill(params.value);
+    await this.runLocator("locator.fill", params, (locator, progress) =>
+      locator.fill(params.value, progress),
+    );
     return { filled: true };
   }
 
-  async locatorCount(params: LocatorDescriptor): Promise<LocatorCountResult> {
-    return await this.resolveLocator(params).count();
+  async locatorCount(params: LocatorParams): Promise<LocatorCountResult> {
+    return await this.runLocator("locator.count", params, (locator, progress) =>
+      locator.count(progress),
+    );
   }
 
-  async locatorIsChecked(params: LocatorDescriptor): Promise<LocatorIsCheckedResult> {
-    return await this.resolveLocator(params).isChecked();
+  async locatorIsChecked(params: LocatorParams): Promise<LocatorIsCheckedResult> {
+    return await this.runLocator("locator.is_checked", params, (locator, progress) =>
+      locator.isChecked(progress),
+    );
   }
 
-  async locatorInputValue(params: LocatorDescriptor): Promise<LocatorInputValueResult> {
-    return await this.resolveLocator(params).inputValue();
+  async locatorInputValue(params: LocatorParams): Promise<LocatorInputValueResult> {
+    return await this.runLocator("locator.input_value", params, (locator, progress) =>
+      locator.inputValue(progress),
+    );
   }
 
-  async locatorIsVisible(params: LocatorDescriptor): Promise<LocatorIsVisibleResult> {
-    return await this.resolveLocator(params).isVisible();
+  async locatorIsVisible(params: LocatorParams): Promise<LocatorIsVisibleResult> {
+    return await this.runLocator("locator.is_visible", params, (locator, progress) =>
+      locator.isVisible(progress),
+    );
   }
 
-  async locatorInnerText(params: LocatorDescriptor): Promise<LocatorInnerTextResult> {
-    return await this.resolveLocator(params).innerText();
+  async locatorInnerText(params: LocatorParams): Promise<LocatorInnerTextResult> {
+    return await this.runLocator("locator.inner_text", params, (locator, progress) =>
+      locator.innerText(progress),
+    );
   }
 
-  async locatorInnerHtml(params: LocatorDescriptor): Promise<LocatorInnerHtmlResult> {
-    return await this.resolveLocator(params).innerHtml();
+  async locatorInnerHtml(params: LocatorParams): Promise<LocatorInnerHtmlResult> {
+    return await this.runLocator("locator.inner_html", params, (locator, progress) =>
+      locator.innerHtml(progress),
+    );
   }
 
-  async locatorTextContent(params: LocatorDescriptor): Promise<LocatorTextContentResult> {
-    return await this.resolveLocator(params).textContent();
+  async locatorTextContent(params: LocatorParams): Promise<LocatorTextContentResult> {
+    return await this.runLocator("locator.text_content", params, (locator, progress) =>
+      locator.textContent(progress),
+    );
   }
 
   async locatorScrollTo(params: LocatorScrollToParams): Promise<LocatorScrollToResult> {
-    await this.resolveLocator(params).scrollTo(params.percent);
+    await this.runLocator("locator.scroll_to", params, (locator, progress) =>
+      locator.scrollTo(params.percent, progress),
+    );
     return { scrolled: true };
   }
 
-  async locatorCentroid(params: LocatorDescriptor): Promise<LocatorCentroidResult> {
-    return await this.resolveLocator(params).centroid();
+  async locatorCentroid(params: LocatorParams): Promise<LocatorCentroidResult> {
+    return await this.runLocator("locator.centroid", params, (locator, progress) =>
+      locator.centroid(progress),
+    );
   }
 
   async locatorHighlight(params: LocatorHighlightParams): Promise<LocatorHighlightResult> {
-    await this.resolveLocator(params).highlight(params.options);
+    await this.runLocator("locator.highlight", params, (locator, progress) =>
+      locator.highlight(params.options, progress),
+    );
     return { highlighted: true };
   }
 
   async locatorSendClickEvent(
     params: LocatorSendClickEventParams,
   ): Promise<LocatorSendClickEventResult> {
-    await this.resolveLocator(params).sendClickEvent(params.options);
+    await this.runLocator("locator.send_click_event", params, (locator, progress) =>
+      locator.sendClickEvent(params.options, progress),
+    );
     return { clicked: true };
   }
 
   async locatorType(params: LocatorTypeParams): Promise<LocatorTypeResult> {
-    await this.resolveLocator(params).type(params.text, params.options);
+    await this.runLocator("locator.type", params, (locator, progress) =>
+      locator.type(params.text, params.options, progress),
+    );
     return { typed: true };
   }
 
   async locatorSelectOption(params: LocatorSelectOptionParams): Promise<LocatorSelectOptionResult> {
-    return await this.resolveLocator(params).selectOption(params.values);
+    return await this.runLocator("locator.select_option", params, (locator, progress) =>
+      locator.selectOption(params.values, progress),
+    );
   }
 
   async locatorSetInputFiles(
     params: LocatorSetInputFilesParams,
   ): Promise<LocatorSetInputFilesResult> {
-    await this.resolveLocator(params).setInputFiles(
-      params.files.map((file) => {
-        const binary = globalThis.atob(file.data);
-        const buffer = new Uint8Array(binary.length);
-        for (let index = 0; index < binary.length; index += 1) {
-          buffer[index] = binary.charCodeAt(index);
-        }
-        return {
-          name: file.name,
-          mimeType: file.mimeType,
-          buffer,
-          lastModified: file.lastModified,
-        };
-      }),
+    await this.runLocator("locator.set_input_files", params, (locator, progress) =>
+      locator.setInputFiles(
+        params.files.map((file) => {
+          progress.throwIfStopped();
+          const binary = globalThis.atob(file.data);
+          const buffer = new Uint8Array(binary.length);
+          for (let index = 0; index < binary.length; index += 1) {
+            buffer[index] = binary.charCodeAt(index);
+          }
+          return {
+            name: file.name,
+            mimeType: file.mimeType,
+            buffer,
+            lastModified: file.lastModified,
+          };
+        }),
+        progress,
+      ),
     );
     return { set: true };
   }
 
+  private runLocator<T>(
+    name: string,
+    params: LocatorParams,
+    action: (locator: UnderstudyRuntimeLocator, progress: Progress) => Promise<T> | T,
+  ): Promise<T> {
+    return runWithProgress(
+      { name, timeout: params.options?.timeout ?? DEFAULT_LOCATOR_TIMEOUT_MS },
+      async (progress) => action(this.resolveLocator(params), progress),
+    );
+  }
+
   async close(): Promise<void> {
     await this.enqueueLifecycle(async () => {
       const session = this.browserSession;

packages/extension/tests/runtime-locator-timeouts.test.ts

@@ -0,0 +1,85 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { StagehandMethods } from "@browserbasehq/stagehand-protocol/schema-registry";
+import { DEFAULT_LOCATOR_TIMEOUT_MS } from "@browserbasehq/stagehand-protocol/schemas";
+import { createStagehandRuntime, type UnderstudyRuntimeLocator } from "../runtime.js";
+import { Progress } from "../understudy/progress.js";
+
+const methods = Object.entries(StagehandMethods).filter(([, method]) =>
+  method.name.startsWith("locator."),
+);
+const fields: Record<string, Record<string, unknown>> = {
+  locatorFill: { value: "hello" },
+  locatorType: { text: "hello" },
+  locatorScrollTo: { percent: 50 },
+  locatorSelectOption: { values: "a" },
+  locatorSetInputFiles: { files: [{ name: "hello.txt", data: "aGVsbG8=" }] },
+};
+
+describe("runtime locator deadlines", () => {
+  beforeEach(() => vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "performance"] }));
+  afterEach(() => {
+    vi.restoreAllMocks();
+    vi.useRealTimers();
+  });
+
+  describe.each(methods)("%s", (key, method) => {
+    it.each([undefined, {}, { timeout: 50 }, { timeout: 0 }])(
+      "owns a deadline before resolution with options %j",
+      async (options) => {
+        const runtime = createStagehandRuntime();
+        let progress: Progress | undefined;
+        let complete!: () => void;
+        const work = new Promise<void>((resolve) => {
+          complete = resolve;
+        });
+        const actionName = key.slice("locator".length);
+        const action = vi.fn((...args: unknown[]) => {
+          progress = args.at(-1) as Progress;
+          expect(progress).toBeInstanceOf(Progress);
+          expect(progress.name).toBe(method.name);
+          return work;
+        });
+        const budget = options?.timeout ?? DEFAULT_LOCATOR_TIMEOUT_MS;
+        vi.spyOn(runtime, "resolveLocator").mockImplementation(() => {
+          expect(vi.getTimerCount()).toBe(budget === 0 ? 0 : 1);
+          // Resolution has already spent part of this call's budget.
+          vi.advanceTimersByTime(10);
+          return {
+            [actionName[0]!.toLowerCase() + actionName.slice(1)]: action,
+          } as unknown as UnderstudyRuntimeLocator;
+        });
+        const invoke = runtime[key as keyof typeof runtime] as (
+          params: unknown,
+        ) => Promise<unknown>;
+        const pending = invoke.call(runtime, {
+          pageId: "page-1",
+          selector: "iframe >> button",
+          ...fields[key],
+          ...(options ? { options } : {}),
+        });
+        const observed = pending.catch((error: unknown) => error);
+        await vi.advanceTimersByTimeAsync(0);
+        expect(action).toHaveBeenCalledOnce();
+        expect(progress!.remainingMs()).toBe(budget === 0 ? Infinity : budget - 10);
+        if (budget === 0) {
+          await vi.advanceTimersByTimeAsync(DEFAULT_LOCATOR_TIMEOUT_MS + 1);
+          expect(progress!.signal.aborted).toBe(false);
+          complete();
+          await pending;
+        } else {
+          await vi.advanceTimersByTimeAsync(budget - 11);
+          expect(progress!.signal.aborted).toBe(false);
+          await vi.advanceTimersByTimeAsync(1);
+          expect(await observed).toMatchObject({
+            name: "TimeoutError",
+            message: expect.stringContaining(method.name),
+          });
+          expect(progress!.signal.aborted).toBe(true);
+          complete();
+        }
+        await vi.advanceTimersByTimeAsync(0);
+        expect(vi.getTimerCount()).toBe(0);
+      },
+    );
+  });
+});