diff --git a/packages/protocol/schema-registry.ts b/packages/protocol/schema-registry.ts index e2d91a3cd..115d3453b 100644 --- a/packages/protocol/schema-registry.ts +++ b/packages/protocol/schema-registry.ts @@ -12,6 +12,7 @@ import { ContextActivePageResultSchema, ContextAddCookiesParamsSchema, ContextAddInitScriptParamsSchema, + ContextAwaitActivePageParamsSchema, ContextClearCookiesParamsSchema, ContextClipboardClearParamsSchema, ContextClipboardCopyParamsSchema, @@ -182,6 +183,11 @@ export const StagehandMethods = { params: EmptyParamsSchema, result: ContextActivePageResultSchema, }, + contextAwaitActivePage: { + name: "context.await_active_page", + params: ContextAwaitActivePageParamsSchema, + result: PageRefSchema, + }, contextSetActivePage: { name: "context.set_active_page", params: ContextSetActivePageParamsSchema, diff --git a/packages/protocol/schemas.ts b/packages/protocol/schemas.ts index d553b8aac..90763e4cb 100644 --- a/packages/protocol/schemas.ts +++ b/packages/protocol/schemas.ts @@ -1375,6 +1375,15 @@ export const ContextNewPageParamsSchema = z .strict() .meta({ id: "ContextNewPageParams" }); +export const ContextAwaitActivePageParamsSchema = z + .object({ + timeout: z.number().int().nonnegative().optional().meta({ + description: "Maximum time in milliseconds to wait for the active page", + }), + }) + .strict() + .meta({ id: "ContextAwaitActivePageParams" }); + export const ContextSetActivePageParamsSchema = z .object({ pageId: z.string(), diff --git a/packages/protocol/stagehand.v4.json b/packages/protocol/stagehand.v4.json index 32294180d..69ca9b57f 100644 --- a/packages/protocol/stagehand.v4.json +++ b/packages/protocol/stagehand.v4.json @@ -188,6 +188,19 @@ "required": ["params", "result"], "additionalProperties": false }, + "context.await_active_page": { + "type": "object", + "properties": { + "params": { + "$ref": "#/$defs/ContextAwaitActivePageParams" + }, + "result": { + "$ref": "#/$defs/PageRef" + } + }, + "required": ["params", "result"], + "additionalProperties": false + }, "context.set_active_page": { "type": "object", "properties": { @@ -893,6 +906,7 @@ "context.pages", "context.new_page", "context.active_page", + "context.await_active_page", "context.set_active_page", "context.close", "context.add_init_script", @@ -2815,6 +2829,18 @@ } ] }, + "ContextAwaitActivePageParams": { + "type": "object", + "properties": { + "timeout": { + "description": "Maximum time in milliseconds to wait for the active page", + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991 + } + }, + "additionalProperties": false + }, "ContextSetActivePageParams": { "type": "object", "properties": { @@ -4746,6 +4772,33 @@ "required": ["jsonrpc", "id", "method", "params"], "additionalProperties": false }, + { + "type": "object", + "properties": { + "jsonrpc": { + "type": "string", + "const": "2.0" + }, + "id": { + "$ref": "#/$defs/JSONRPCRequestId" + }, + "method": { + "type": "string", + "const": "context.await_active_page" + }, + "params": { + "$ref": "#/$defs/ContextAwaitActivePageParams" + }, + "traceparent": { + "type": "string" + }, + "tracestate": { + "type": "string" + } + }, + "required": ["jsonrpc", "id", "method", "params"], + "additionalProperties": false + }, { "type": "object", "properties": { diff --git a/packages/protocol/tests/protocol/object-model-protocol.test.ts b/packages/protocol/tests/protocol/object-model-protocol.test.ts index 52e44be8e..ab0800560 100644 --- a/packages/protocol/tests/protocol/object-model-protocol.test.ts +++ b/packages/protocol/tests/protocol/object-model-protocol.test.ts @@ -300,6 +300,7 @@ describe("Stagehand object-model protocol", () => { "context.pages", "context.new_page", "context.active_page", + "context.await_active_page", "context.set_active_page", "context.close", "context.add_init_script", diff --git a/packages/protocol/tests/protocol/schema-registry.test-d.ts b/packages/protocol/tests/protocol/schema-registry.test-d.ts index 82c8efc86..6b40b4454 100644 --- a/packages/protocol/tests/protocol/schema-registry.test-d.ts +++ b/packages/protocol/tests/protocol/schema-registry.test-d.ts @@ -20,6 +20,17 @@ expectTypeOf>().toEqu url?: string; title?: string; } | null>(); +expectTypeOf( + StagehandMethods.contextAwaitActivePage.name, +).toEqualTypeOf<"context.await_active_page">(); +expectTypeOf>().toEqualTypeOf<{ + timeout?: number; +}>(); +expectTypeOf>().toEqualTypeOf<{ + pageId: string; + url?: string; + title?: string; +}>(); expectTypeOf( StagehandMethods.contextSetDomainPolicy.name, ).toEqualTypeOf<"context.set_domain_policy">(); diff --git a/packages/protocol/types.ts b/packages/protocol/types.ts index 5f22bd7d1..8d9660ff0 100644 --- a/packages/protocol/types.ts +++ b/packages/protocol/types.ts @@ -34,6 +34,7 @@ import type { ContextActivePageResultSchema, ContextAddCookiesParamsSchema, ContextAddInitScriptParamsSchema, + ContextAwaitActivePageParamsSchema, ContextClearCookiesParamsSchema, ContextClipboardClearParamsSchema, ContextClipboardCopyParamsSchema, @@ -328,6 +329,7 @@ export type StagehandActParams = z.infer; export type StagehandObserveParams = z.infer; export type StagehandExtractParams = z.infer; export type ContextNewPageParams = z.infer; +export type ContextAwaitActivePageParams = z.infer; export type ContextCookiesParams = z.infer; export type ContextAddCookiesParams = z.infer; export type ContextClearCookiesParams = z.infer; diff --git a/packages/sdk-python/src/stagehand/_generated/models.py b/packages/sdk-python/src/stagehand/_generated/models.py index 14ff92a52..8a49dc7cc 100644 --- a/packages/sdk-python/src/stagehand/_generated/models.py +++ b/packages/sdk-python/src/stagehand/_generated/models.py @@ -323,6 +323,15 @@ class ContextAddInitScriptParams(WireModel): source: StrictStr +class ContextAwaitActivePageParams(WireModel): + model_config = ConfigDict( + extra="forbid", + validate_by_name=True, + ) + timeout: Annotated[Optional[StrictInt], Field(ge=0, le=9007199254740991)] = None + """Maximum time in milliseconds to wait for the active page""" + + class ContextClearCookiesParams(WireModel): model_config = ConfigDict( extra="forbid", diff --git a/packages/sdk-python/src/stagehand/browser_context.py b/packages/sdk-python/src/stagehand/browser_context.py index e6ef7f3a7..716936c9a 100644 --- a/packages/sdk-python/src/stagehand/browser_context.py +++ b/packages/sdk-python/src/stagehand/browser_context.py @@ -9,6 +9,7 @@ ContextActivePageResult, ContextAddCookiesParams, ContextAddInitScriptParams, + ContextAwaitActivePageParams, ContextClearCookiesParams, ContextCloseResult, ContextCookiesParams, @@ -67,6 +68,15 @@ async def active_page(self) -> Page | None: ) return None if result.root is None else Page(self._rpc_client, result.root) + async def await_active_page(self, timeout: int | None = None) -> Page: + params = ContextAwaitActivePageParams(timeout=timeout) + page_ref = await self._rpc_client.send( + "context.await_active_page", + params, + PageRef, + ) + return Page(self._rpc_client, page_ref) + async def set_active_page(self, page: Page) -> None: await self._rpc_client.send( "context.set_active_page", diff --git a/packages/sdk-python/src/stagehand/rpc_client.py b/packages/sdk-python/src/stagehand/rpc_client.py index 5fa24d088..8bfdb0d7c 100644 --- a/packages/sdk-python/src/stagehand/rpc_client.py +++ b/packages/sdk-python/src/stagehand/rpc_client.py @@ -81,12 +81,8 @@ def __init__(self, error: _JSONRPCError) -> None: class RPCClient: - def __init__(self, transport: _Transport, *, request_timeout_ms: int = 10_000) -> None: - if request_timeout_ms <= 0: - raise ValueError("request_timeout_ms must be positive") - + def __init__(self, transport: _Transport) -> None: self._transport = transport - self._request_timeout_seconds = request_timeout_ms / 1_000 self._next_request_id = 1 self._pending: dict[ int, @@ -142,17 +138,14 @@ async def send( ) try: - async with asyncio.timeout(self._request_timeout_seconds): - await self._transport.send( - cast( - dict[str, object], - request.model_dump(mode="json", exclude_none=True, exclude_unset=True), - ) + await self._transport.send( + cast( + dict[str, object], + request.model_dump(mode="json", exclude_none=True, exclude_unset=True), ) - result = await response - return cast(ResultT, result) - except TimeoutError as error: - raise TimeoutError(f"RPC request timed out: {method}") from error + ) + result = await response + return cast(ResultT, result) finally: self._pending.pop(request_id, None) if not response.done(): @@ -477,7 +470,7 @@ async def connect_rpc_client( command_timeout_ms=command_timeout_ms, cdp_connect_timeout_ms=cdp_connect_timeout_ms, ) - client = RPCClient(cdp, request_timeout_ms=command_timeout_ms) + client = RPCClient(cdp) configure = models.RuntimeConfigureParams( cdp_url=cdp.web_socket_debugger_url, **({"telemetry": telemetry} if telemetry is not None else {}), diff --git a/packages/sdk-python/tests/test_browser_context.py b/packages/sdk-python/tests/test_browser_context.py index 4c534cf19..29971e3d4 100644 --- a/packages/sdk-python/tests/test_browser_context.py +++ b/packages/sdk-python/tests/test_browser_context.py @@ -22,6 +22,7 @@ async def test_browser_context_wraps_generated_page_references() -> None: "context.pages": [PageRef(page_id="page-1")], "context.new_page": PageRef(page_id="page-2"), "context.active_page": PageRef(page_id="page-2"), + "context.await_active_page": PageRef(page_id="page-3"), "context.set_active_page": ContextVoidResult(ok=True), }) context = BrowserContext(cast(RPCClient, recording)) @@ -29,18 +30,22 @@ async def test_browser_context_wraps_generated_page_references() -> None: pages = await context.pages() new_page = await context.new_page(url="https://example.com") active_page = await context.active_page() + awaited_page = await context.await_active_page(timeout=4_000) await context.set_active_page(new_page) assert [page.page_id for page in pages] == ["page-1"] assert new_page.page_id == "page-2" assert active_page is not None and active_page.page_id == "page-2" + assert awaited_page.page_id == "page-3" assert [call[0] for call in recording.calls] == [ "context.pages", "context.new_page", "context.active_page", + "context.await_active_page", "context.set_active_page", ] assert recording.calls[1][1].model_dump(exclude_unset=True) == {"url": "https://example.com"} + assert recording.calls[3][1].model_dump(exclude_unset=True) == {"timeout": 4_000} def test_browser_context_reuses_one_clipboard_wrapper() -> None: diff --git a/packages/sdk-python/tests/test_rpc_client.py b/packages/sdk-python/tests/test_rpc_client.py index 8b14c413c..a1ebc0578 100644 --- a/packages/sdk-python/tests/test_rpc_client.py +++ b/packages/sdk-python/tests/test_rpc_client.py @@ -394,15 +394,7 @@ async def test_error_responses_preserve_the_json_rpc_code_and_data() -> None: @pytest.mark.asyncio -async def test_timeout_and_transport_close_reject_pending_requests() -> None: - timeout_transport = QueueTransport() - timeout_client = RPCClient(timeout_transport, request_timeout_ms=10) - try: - with pytest.raises(TimeoutError, match="RPC request timed out: ping"): - await timeout_client.send("ping", models.EmptyParams(), models.StagehandPingResult) - finally: - await timeout_client.close() - +async def test_transport_close_rejects_pending_requests() -> None: failing_transport = FailingReceiveTransport() failing_client = RPCClient(failing_transport) call = asyncio.create_task( diff --git a/packages/sdk-ts/src/browserContext.ts b/packages/sdk-ts/src/browserContext.ts index db4a54537..e629a408e 100644 --- a/packages/sdk-ts/src/browserContext.ts +++ b/packages/sdk-ts/src/browserContext.ts @@ -44,6 +44,14 @@ export class BrowserContext { return pageRef ? new Page(this.rpcClient, pageRef) : undefined; } + async awaitActivePage(timeoutMs?: number): Promise { + const pageRef = await this.rpcClient.send( + StagehandMethods.contextAwaitActivePage, + timeoutMs === undefined ? {} : { timeout: timeoutMs }, + ); + return new Page(this.rpcClient, pageRef); + } + async setActivePage(page: Page): Promise { await this.rpcClient.send(StagehandMethods.contextSetActivePage, { pageId: page.pageId, diff --git a/packages/sdk-ts/src/rpcClient.ts b/packages/sdk-ts/src/rpcClient.ts index a2a98adaa..3bc853c76 100644 --- a/packages/sdk-ts/src/rpcClient.ts +++ b/packages/sdk-ts/src/rpcClient.ts @@ -47,7 +47,6 @@ type PendingRequest = { method: RPCMethod; resolve(value: unknown): void; reject(error: Error): void; - timeout: ReturnType; }; type RegisteredRequestHandler = { @@ -113,11 +112,9 @@ export class RPCClient { pendingNotifications: StagehandRpcNotification[] = []; closed = false; readonly cdp: CDPTransport; - readonly requestTimeoutMs: number; - constructor(cdp: CDPTransport, requestTimeoutMs: number) { + constructor(cdp: CDPTransport) { this.cdp = cdp; - this.requestTimeoutMs = requestTimeoutMs; this.serviceWorker = cdp.serviceWorker; this.cdp.onmessage = (message) => this.receive(message); this.cdp.onclose = (reason) => this.close(reason); @@ -219,12 +216,7 @@ export class RPCClient { waitForResponse(id: number, method: RPCMethod): Promise { return new Promise((resolve, reject) => { - const timeout = setTimeout(() => { - if (!this.pending.delete(id)) return; - reject(new Error(`RPC request timed out: ${method.name}`)); - }, this.requestTimeoutMs); - - this.pending.set(id, { method, resolve, reject, timeout }); + this.pending.set(id, { method, resolve, reject }); }); } @@ -349,7 +341,6 @@ export class RPCClient { if (!pending) return; this.pending.delete(response.id); - clearTimeout(pending.timeout); if ("error" in response) { pending.reject(new Error(response.error.message, { cause: response.error })); @@ -369,7 +360,6 @@ export class RPCClient { const pending = this.pending.get(id); if (!pending) return; this.pending.delete(id); - clearTimeout(pending.timeout); pending.reject(error); } @@ -413,7 +403,7 @@ export async function connectRPCClient(input: RPCClientOptions): Promise(); constructor() { - super( - { - serviceWorker: { - targetId: "worker-target", - url: "chrome-extension://stagehand/service-worker.js", - title: "Stagehand", - extensionId: "stagehand", - }, - send: async () => {}, - close: () => {}, + super({ + serviceWorker: { + targetId: "worker-target", + url: "chrome-extension://stagehand/service-worker.js", + title: "Stagehand", + extensionId: "stagehand", }, - 1_000, - ); + send: async () => {}, + close: () => {}, + }); this.queueResponse(StagehandMethods.stagehandInit, { initialized: true, pages: [] }); } @@ -179,6 +176,28 @@ describe("Stagehand TS object wrapper", () => { ]); }); + it("waits for and wraps the active page", async () => { + const client = new FakeProtocolClient(); + client.queueResponse(StagehandMethods.contextAwaitActivePage, { + pageId: "new-active-page", + url: "https://example.com/new", + }); + const stagehand = createStagehandWithClientForTest(client); + await stagehand.init(); + + const page = await stagehand.context.awaitActivePage(4_000); + + expect(page).toBeInstanceOf(Page); + expect(page.ref).toStrictEqual({ + pageId: "new-active-page", + url: "https://example.com/new", + }); + expect(client.calls).toStrictEqual([ + stagehandInitCall, + requestCall(StagehandMethods.contextAwaitActivePage, { timeout: 4_000 }), + ]); + }); + it("routes context.setActivePage and context.close", async () => { const client = new FakeProtocolClient(); client.queueResponse(StagehandMethods.contextSetActivePage, { ok: true }); diff --git a/packages/sdk-ts/tests/rpcClient.test.ts b/packages/sdk-ts/tests/rpcClient.test.ts index 46ada7d54..c5602ca30 100644 --- a/packages/sdk-ts/tests/rpcClient.test.ts +++ b/packages/sdk-ts/tests/rpcClient.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, expectTypeOf, it } from "vitest"; +import { describe, expect, expectTypeOf, it, vi } from "vitest"; import { z } from "zod/v4"; import { JSONRPCErrorCodes, type RPCMethod } from "../../protocol/json-rpc/schemas.js"; import type { JSONRPCMessage } from "../../protocol/json-rpc/types.js"; @@ -60,7 +60,7 @@ class ManualCDPTransport implements CDPTransport { describe("RPCClient", () => { it("accepts page methods without SDK wrapper methods", async () => { const cdp = new FakeCDPTransport({ matched: true }); - const client = new RPCClient(cdp, 1_000); + const client = new RPCClient(cdp); const request = client.send(StagehandMethods.pageWaitForSelector, { pageId: "page-1", @@ -97,7 +97,7 @@ describe("RPCClient", () => { }, ], }); - const client = new RPCClient(cdp, 1_000); + const client = new RPCClient(cdp); const request = client.send(StagehandMethods.contextCookies, { urls: ["https://example.com/account"], @@ -144,7 +144,7 @@ describe("RPCClient", () => { page_id: "page-1", url: "https://example.com", }); - const client = new RPCClient(cdp, 1_000); + const client = new RPCClient(cdp); await expect( client.send(StagehandMethods.pageGoto, { @@ -169,7 +169,7 @@ describe("RPCClient", () => { it("rejects invalid method params before sending them over CDP", async () => { const cdp = new FakeCDPTransport({ ok: true, runtime: "service_worker" }); - const client = new RPCClient(cdp, 1_000); + const client = new RPCClient(cdp); await expect(client.send(StagehandMethods.ping, { extra: true } as never)).rejects.toThrow(); @@ -178,7 +178,7 @@ describe("RPCClient", () => { it("lets the worker request client work while the original SDK request is still pending", async () => { const cdp = new ManualCDPTransport(); - const client = new RPCClient(cdp, 1_000); + const client = new RPCClient(cdp); client.onRequest(UppercaseMethod, async ({ value }) => ({ value: value.toUpperCase() })); const originalRequest = client.send(StagehandMethods.ping, {}); @@ -208,7 +208,7 @@ describe("RPCClient", () => { it("validates incoming request parameters before invoking the SDK handler", async () => { const cdp = new ManualCDPTransport(); - const client = new RPCClient(cdp, 1_000); + const client = new RPCClient(cdp); let calls = 0; client.onRequest(UppercaseMethod, async ({ value }) => { calls += 1; @@ -235,7 +235,7 @@ describe("RPCClient", () => { it("validates an SDK handler result before returning it to the worker", async () => { const cdp = new ManualCDPTransport(); - const client = new RPCClient(cdp, 1_000); + const client = new RPCClient(cdp); client.onRequest(UppercaseMethod, async () => ({ value: 42 }) as never); await cdp.receive({ @@ -257,7 +257,7 @@ describe("RPCClient", () => { it("returns method not found when no SDK handler is registered", async () => { const cdp = new ManualCDPTransport(); - new RPCClient(cdp, 1_000); + new RPCClient(cdp); await cdp.receive({ jsonrpc: "2.0", @@ -278,7 +278,7 @@ describe("RPCClient", () => { it("returns a JSON-RPC error when an SDK handler throws", async () => { const cdp = new ManualCDPTransport(); - const client = new RPCClient(cdp, 1_000); + const client = new RPCClient(cdp); client.onRequest(UppercaseMethod, async () => { throw new Error("Client handler failed"); }); @@ -303,7 +303,7 @@ describe("RPCClient", () => { it("rejects a failed request with a plain Error that preserves the JSON-RPC failure", async () => { const cdp = new ManualCDPTransport(); - const client = new RPCClient(cdp, 1_000); + const client = new RPCClient(cdp); const request = client.send(StagehandMethods.ping, {}); const rpcError = { code: JSONRPCErrorCodes.internalError, @@ -320,9 +320,32 @@ describe("RPCClient", () => { }); }); + it("keeps a request pending without a transport response deadline", async () => { + vi.useFakeTimers(); + try { + const cdp = new ManualCDPTransport(); + const client = new RPCClient(cdp); + const request = client.send(StagehandMethods.ping, {}); + + await vi.advanceTimersByTimeAsync(60_001); + await cdp.receive({ + jsonrpc: "2.0", + id: 1, + result: { ok: true, runtime: "service_worker" }, + }); + + await expect(request).resolves.toStrictEqual({ + ok: true, + runtime: "service_worker", + }); + } finally { + vi.useRealTimers(); + } + }); + it("removes incoming SDK request handlers when the RPC client closes", async () => { const cdp = new ManualCDPTransport(); - const client = new RPCClient(cdp, 1_000); + const client = new RPCClient(cdp); let calls = 0; client.onRequest(UppercaseMethod, async ({ value }) => { calls += 1; diff --git a/packages/sdk-ts/tests/stagehand.test.ts b/packages/sdk-ts/tests/stagehand.test.ts index c5c5ef538..698845c9b 100644 --- a/packages/sdk-ts/tests/stagehand.test.ts +++ b/packages/sdk-ts/tests/stagehand.test.ts @@ -18,19 +18,16 @@ class FakeRPCClient extends RPCClient { notificationListeners = new Set<(notification: StagehandRpcNotification) => void>(); constructor() { - super( - { - serviceWorker: { - targetId: "worker-target", - url: "chrome-extension://stagehand/service-worker.js", - title: "Stagehand", - extensionId: "stagehand", - }, - send: async () => {}, - close: () => {}, - }, - 1_000, - ); + super({ + serviceWorker: { + targetId: "worker-target", + url: "chrome-extension://stagehand/service-worker.js", + title: "Stagehand", + extensionId: "stagehand", + }, + send: async () => {}, + close: () => {}, + }); this.queueResponse(StagehandMethods.stagehandInit, { initialized: true, pages: [] }); } diff --git a/packages/server/clients/rpcClient.ts b/packages/server/clients/rpcClient.ts index 5b04d6962..41e671cf9 100644 --- a/packages/server/clients/rpcClient.ts +++ b/packages/server/clients/rpcClient.ts @@ -32,7 +32,6 @@ type PendingRequest = { method: RPCMethod; resolve(value: unknown): void; reject(error: Error): void; - timeout: ReturnType; }; const ERROR_DATA = { @@ -48,7 +47,6 @@ export class RPCClient { constructor( readonly runtime: ChromeRuntimeClient, readonly router: RPCRouter, - readonly requestTimeoutMs = 60_000, ) { this.runtime.onmessage = (message) => this.receive(message); this.runtime.onclose = (reason) => this.close(reason); @@ -133,12 +131,7 @@ export class RPCClient { waitForResponse(id: number, method: RPCMethod): Promise { return new Promise((resolve, reject) => { - const timeout = setTimeout(() => { - if (!this.pending.delete(id)) return; - reject(new Error(`RPC request timed out: ${method.name}`)); - }, this.requestTimeoutMs); - - this.pending.set(id, { method, resolve, reject, timeout }); + this.pending.set(id, { method, resolve, reject }); }); } @@ -248,7 +241,6 @@ export class RPCClient { if (!pending) return; this.pending.delete(response.id); - clearTimeout(pending.timeout); if ("error" in response) { pending.reject(new Error(response.error.message, { cause: response.error })); @@ -268,7 +260,6 @@ export class RPCClient { const pending = this.pending.get(id); if (!pending) return; this.pending.delete(id); - clearTimeout(pending.timeout); pending.reject(error); } diff --git a/packages/server/controllers/contextController.ts b/packages/server/controllers/contextController.ts index 9f4013e1b..a59459306 100644 --- a/packages/server/controllers/contextController.ts +++ b/packages/server/controllers/contextController.ts @@ -1,6 +1,7 @@ import type { ContextAddCookiesParams, ContextAddInitScriptParams, + ContextAwaitActivePageParams, ContextClearCookiesParams, ContextClipboardClearParams, ContextClipboardCopyParams, @@ -34,6 +35,11 @@ export function createContextController(runtime: StagehandRuntime) { return runtime.contextActivePage(); } + async function awaitActivePage(params: ContextAwaitActivePageParams, { logger }: HandlerContext) { + logger.debug("context.await_active_page", {}); + return runtime.contextAwaitActivePage(params); + } + async function setActivePage(params: ContextSetActivePageParams, { logger }: HandlerContext) { logger.debug("context.set_active_page", {}); return runtime.contextSetActivePage(params); @@ -122,6 +128,7 @@ export function createContextController(runtime: StagehandRuntime) { pages, newPage, activePage, + awaitActivePage, setActivePage, close, addInitScript, diff --git a/packages/server/errors.ts b/packages/server/errors.ts index d64dd03fc..9e90f9994 100644 --- a/packages/server/errors.ts +++ b/packages/server/errors.ts @@ -4,3 +4,10 @@ export class TimeoutError extends Error { this.name = "TimeoutError"; } } + +export class PageNotFoundError extends Error { + constructor(message: string) { + super(message); + this.name = "PageNotFoundError"; + } +} diff --git a/packages/server/rpcRouter.ts b/packages/server/rpcRouter.ts index eb9e8eefb..6d89a90ba 100644 --- a/packages/server/rpcRouter.ts +++ b/packages/server/rpcRouter.ts @@ -149,6 +149,11 @@ export class RPCRouter { parseParams(StagehandMethods.contextActivePage, request.params), context, ); + case "context.await_active_page": + return this.contextController.awaitActivePage( + parseParams(StagehandMethods.contextAwaitActivePage, request.params), + context, + ); case "context.set_active_page": return this.contextController.setActivePage( parseParams(StagehandMethods.contextSetActivePage, request.params), diff --git a/packages/server/runtime.ts b/packages/server/runtime.ts index ac8a6e0cd..e1e82e31f 100644 --- a/packages/server/runtime.ts +++ b/packages/server/runtime.ts @@ -4,6 +4,7 @@ import type { ContextActivePageResult, ContextAddCookiesParams, ContextAddInitScriptParams, + ContextAwaitActivePageParams, ContextClearCookiesParams, ContextClipboardClearParams, ContextClipboardCopyParams, @@ -203,6 +204,7 @@ export type StagehandBrowserSession = { pages(): UnderstudyRuntimePage[]; newPage(url?: string): Promise; activePage(): Promise; + awaitActivePage(timeoutMs?: number): Promise; setActivePage(page: UnderstudyRuntimePage): Promise; addInitScript(source: string): Promise; setExtraHTTPHeaders(headers: ContextSetExtraHTTPHeadersParams["headers"]): Promise; @@ -344,6 +346,12 @@ export class StagehandRuntime { return pageRefFromUnderstudyPage(page); } + async contextAwaitActivePage(params: ContextAwaitActivePageParams): Promise { + const page = await this.requireBrowserSession().awaitActivePage(params.timeout); + this.registerPage(page); + return pageRefFromUnderstudyPage(page); + } + async contextSetActivePage(params: ContextSetActivePageParams): Promise { const page = this.resolvePage(params.pageId); await this.requireBrowserSession().setActivePage(page); diff --git a/packages/server/tests/client-llm-client.test.ts b/packages/server/tests/client-llm-client.test.ts index 9f44e21a7..ddaf0e429 100644 --- a/packages/server/tests/client-llm-client.test.ts +++ b/packages/server/tests/client-llm-client.test.ts @@ -55,6 +55,9 @@ describe("client LLM generation", () => { getVersion: async () => ({}), pages: () => [], activePage: async () => undefined, + awaitActivePage: async () => { + throw new Error("Not used by this test"); + }, setActivePage: async () => {}, addInitScript: async () => {}, setExtraHTTPHeaders: async () => {}, diff --git a/packages/server/tests/rpc-client.test.ts b/packages/server/tests/rpc-client.test.ts index f6728b6ee..4b89563bb 100644 --- a/packages/server/tests/rpc-client.test.ts +++ b/packages/server/tests/rpc-client.test.ts @@ -1,6 +1,6 @@ import { ROOT_CONTEXT, TraceFlags, context, trace } from "@opentelemetry/api"; import { StackContextManager } from "@opentelemetry/sdk-trace-web"; -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { JSONRPCRequestSchema } from "../../protocol/json-rpc/schemas.ts"; import { StagehandMethods } from "../../protocol/schema-registry.ts"; import { ChromeRuntimeClient } from "../clients/chromeRuntimeClient.ts"; @@ -37,7 +37,7 @@ describe("worker RPCClient", () => { }, }; runtimeClient = new ChromeRuntimeClient(scope, "sendToHost"); - const client = new RPCClient(runtimeClient, new RPCRouter(runtime), 1_000); + const client = new RPCClient(runtimeClient, new RPCRouter(runtime)); await expect(client.send(StagehandMethods.ping, {})).resolves.toStrictEqual({ ok: true, @@ -77,7 +77,7 @@ describe("worker RPCClient", () => { }, }; runtimeClient = new ChromeRuntimeClient(scope, "sendToHost"); - const client = new RPCClient(runtimeClient, new RPCRouter(runtime), 1_000); + const client = new RPCClient(runtimeClient, new RPCRouter(runtime)); const parentContext = trace.setSpanContext(ROOT_CONTEXT, { traceId: "4bf92f3577b34da6a3ce929d0e0e4736", spanId: "00f067aa0ba902b7", @@ -93,4 +93,42 @@ describe("worker RPCClient", () => { context.disable(); } }); + + it("keeps reverse requests pending without a response deadline", async () => { + vi.useFakeTimers(); + let runtimeClient: ChromeRuntimeClient | undefined; + let requestId: number | undefined; + const runtime = createStagehandRuntime({ + browserSessionFactory: async () => { + throw new Error("Stagehand browser session factory is not configured"); + }, + }); + const scope = { + sendToHost(payload: string): void { + requestId = JSONRPCRequestSchema.parse(JSON.parse(payload)).id; + }, + }; + runtimeClient = new ChromeRuntimeClient(scope, "sendToHost"); + const client = new RPCClient(runtimeClient, new RPCRouter(runtime)); + + try { + const request = client.send(StagehandMethods.ping, {}); + await vi.advanceTimersByTimeAsync(60_001); + await runtimeClient.receive( + JSON.stringify({ + jsonrpc: "2.0", + id: requestId, + result: { ok: true, runtime: "service_worker" }, + }), + ); + + await expect(request).resolves.toStrictEqual({ + ok: true, + runtime: "service_worker", + }); + } finally { + client.close(); + vi.useRealTimers(); + } + }); }); diff --git a/packages/server/tests/runtime-state.test.ts b/packages/server/tests/runtime-state.test.ts index 3775fd57d..4c0298b71 100644 --- a/packages/server/tests/runtime-state.test.ts +++ b/packages/server/tests/runtime-state.test.ts @@ -13,6 +13,9 @@ function createBrowserSession( throw new Error("Not used by this test"); }, activePage: async () => undefined, + awaitActivePage: async () => { + throw new Error("Not used by this test"); + }, setActivePage: async () => {}, addInitScript: async () => {}, setExtraHTTPHeaders: async () => {}, diff --git a/packages/server/tests/stagehand-clients.test.ts b/packages/server/tests/stagehand-clients.test.ts index f5369a060..b8078a7b9 100644 --- a/packages/server/tests/stagehand-clients.test.ts +++ b/packages/server/tests/stagehand-clients.test.ts @@ -101,6 +101,7 @@ class FakeBrowserSession implements StagehandBrowserSession { getVersionCalls = 0; readonly pageRefs: FakeUnderstudyRuntimePage[]; activePageRef: UnderstudyRuntimePage | undefined; + readonly awaitActivePageTimeouts: Array = []; readonly setActivePageCalls: UnderstudyRuntimePage[] = []; readonly contextAddInitScriptCalls: string[] = []; readonly contextSetExtraHTTPHeadersCalls: ContextSetExtraHTTPHeadersParams["headers"][] = []; @@ -146,6 +147,12 @@ class FakeBrowserSession implements StagehandBrowserSession { return this.activePageRef; } + async awaitActivePage(timeoutMs?: number): Promise { + this.awaitActivePageTimeouts.push(timeoutMs); + if (!this.activePageRef) throw new Error("No active page"); + return this.activePageRef; + } + async setActivePage(page: UnderstudyRuntimePage): Promise { this.setActivePageCalls.push(page); this.activePageRef = page; @@ -1022,6 +1029,23 @@ describe("Stagehand worker clients", () => { }, }); + await expect( + handle({ + jsonrpc: "2.0", + id: 90, + method: "context.await_active_page", + params: { timeout: 4_000 }, + }), + ).resolves.toStrictEqual({ + jsonrpc: "2.0", + id: 90, + result: { + page_id: "page-b", + url: "https://example.test/b", + }, + }); + expect(context.awaitActivePageTimeouts).toStrictEqual([4_000]); + await expect( handle({ jsonrpc: "2.0", diff --git a/packages/server/tests/understudy-context-active-page.test.ts b/packages/server/tests/understudy-context-active-page.test.ts index 2d9e16ced..5dd592d7d 100644 --- a/packages/server/tests/understudy-context-active-page.test.ts +++ b/packages/server/tests/understudy-context-active-page.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it, vi } from "vitest"; +import { PageNotFoundError } from "../errors.ts"; import type { ChromeTabTargetController } from "../understudy/chromeTabs.ts"; import { V3Context } from "../understudy/context.ts"; import type { Page } from "../understudy/page.ts"; @@ -34,6 +35,54 @@ describe("V3Context active page", () => { await expect(context.activePage()).resolves.toBeUndefined(); }); + it("returns the newly registered page after Page.windowOpen", async () => { + const listeners = new Map void>(); + const session = { + id: "session-1", + on: vi.fn((event: string, listener: () => void) => { + listeners.set(event, listener); + }), + }; + const chromeTabs: ChromeTabTargetController = { + activeTargetId: vi.fn(async () => "old-target"), + targetIdForTabId: vi.fn(async () => undefined), + tabIdForTargetId: vi.fn(async () => undefined), + activateTarget: vi.fn(async () => {}), + }; + const context = new V3Context( + { getSession: () => session } as never, + {} as never, + chromeTabs, + "BROWSERBASE", + ); + const oldPage = createPage("old-target"); + const newPage = createPage("new-target"); + context.pagesByTarget.set("old-target", oldPage); + context.createdAtByTarget.set("old-target", 0); + context.installFrameEventBridges("session-1", oldPage); + + listeners.get("Page.windowOpen")?.(); + context.pagesByTarget.set("new-target", newPage); + context.createdAtByTarget.set("new-target", Date.now() + 1); + + await expect(context.awaitActivePage(100)).resolves.toBe(newPage); + }); + + it("falls back to the prior active page when popup registration times out", async () => { + const { context } = createContext("page-target"); + const page = createPage("page-target"); + context.pagesByTarget.set("page-target", page); + Reflect.set(context, "_lastPopupSignalAt", Date.now()); + + await expect(context.awaitActivePage(0)).resolves.toBe(page); + }); + + it("throws when no active page becomes available", async () => { + const { context } = createContext(); + + await expect(context.awaitActivePage(0)).rejects.toBeInstanceOf(PageNotFoundError); + }); + it("uses the Chrome-backed active page for implicit clipboard operations", async () => { const { chromeTabs, context } = createContext("page-target"); const page = createPage("page-target"); diff --git a/packages/server/understudy/context.ts b/packages/server/understudy/context.ts index 81879e4c5..c1c975fc5 100644 --- a/packages/server/understudy/context.ts +++ b/packages/server/understudy/context.ts @@ -16,7 +16,7 @@ import type { import { InitScriptSource } from "../types/private/index.js"; import { normalizeInitScriptSource } from "./initScripts.js"; import { ContextClipboard } from "./clipboard.js"; -import { TimeoutError } from "../errors.js"; +import { PageNotFoundError, TimeoutError } from "../errors.js"; import { filterCookies, normalizeCookieParams, @@ -99,6 +99,7 @@ export class V3Context { ) {} readonly _targetSessionListeners = new Set(); + private _lastPopupSignalAt = 0; readonly _domainPolicySessionListeners = new Map< SessionId, (evt: Protocol.Fetch.RequestPausedEvent) => void @@ -271,6 +272,46 @@ export class V3Context { return targetId === undefined ? undefined : this.pagesByTarget.get(targetId); } + private notePopupSignal(): void { + this._lastPopupSignalAt = Date.now(); + } + + /** + * Return the active page, waiting briefly for a newly opened popup target + * to be registered when Chrome has just emitted Page.windowOpen. + */ + public async awaitActivePage(timeoutMs?: number): Promise { + const defaultTimeout = this.env === "BROWSERBASE" ? 4000 : 2000; + const timeout = timeoutMs ?? defaultTimeout; + const recentWindowMs = this.env === "BROWSERBASE" ? 1000 : 300; + const now = Date.now(); + const hasRecentPopup = now - this._lastPopupSignalAt <= recentWindowMs; + + const immediate = await this.activePage(); + if (!hasRecentPopup && immediate) return immediate; + + const deadline = now + timeout; + while (Date.now() < deadline) { + let newestTargetId: TargetId | undefined; + let newestCreatedAt = -1; + for (const [targetId] of this.pagesByTarget) { + const createdAt = this.createdAtByTarget.get(targetId) ?? 0; + if (createdAt > newestCreatedAt) { + newestCreatedAt = createdAt; + newestTargetId = targetId; + } + } + if (newestTargetId) { + const page = this.pagesByTarget.get(newestTargetId); + if (page && newestCreatedAt >= this._lastPopupSignalAt) return page; + } + await new Promise((resolve) => setTimeout(resolve, 25)); + } + + if (immediate) return immediate; + throw new PageNotFoundError("awaitActivePage: no page available"); + } + /** Select the Chrome tab that owns a known understudy Page. */ public async setActivePage(page: Page): Promise { let targetId = page.targetId(); @@ -615,6 +656,7 @@ export class V3Context { this.conn.on("Target.targetCreated", (evt) => { const info = evt.targetInfo; if (info.type === "page" && (info.openerId || info.openerFrameId)) { + this.notePopupSignal(); void this.closePopupIfBlockedByDomainPolicy(info, "targetCreated"); } }); @@ -1139,6 +1181,10 @@ export class V3Context { owner.onNavigatedWithinDocument(evt.frameId, evt.url, session); }, ); + + session.on("Page.windowOpen", () => { + this.notePopupSignal(); + }); } /**