Restore active-page waiting in V4

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,

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(),

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": {

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",

packages/protocol/tests/protocol/schema-registry.test-d.ts

@@ -20,6 +20,17 @@ expectTypeOf<z.output<typeof StagehandMethods.contextActivePage.result>>().toEqu
   url?: string;
   title?: string;
 } | null>();
+expectTypeOf(
+  StagehandMethods.contextAwaitActivePage.name,
+).toEqualTypeOf<"context.await_active_page">();
+expectTypeOf<z.input<typeof StagehandMethods.contextAwaitActivePage.params>>().toEqualTypeOf<{
+  timeout?: number;
+}>();
+expectTypeOf<z.output<typeof StagehandMethods.contextAwaitActivePage.result>>().toEqualTypeOf<{
+  pageId: string;
+  url?: string;
+  title?: string;
+}>();
 expectTypeOf(
   StagehandMethods.contextSetDomainPolicy.name,
 ).toEqualTypeOf<"context.set_domain_policy">();

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<typeof StagehandActParamsSchema>;
 export type StagehandObserveParams = z.infer<typeof StagehandObserveParamsSchema>;
 export type StagehandExtractParams = z.infer<typeof StagehandExtractParamsSchema>;
 export type ContextNewPageParams = z.infer<typeof ContextNewPageParamsSchema>;
+export type ContextAwaitActivePageParams = z.infer<typeof ContextAwaitActivePageParamsSchema>;
 export type ContextCookiesParams = z.infer<typeof ContextCookiesParamsSchema>;
 export type ContextAddCookiesParams = z.infer<typeof ContextAddCookiesParamsSchema>;
 export type ContextClearCookiesParams = z.infer<typeof ContextClearCookiesParamsSchema>;

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",

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",

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 {}),

packages/sdk-python/tests/test_browser_context.py

@@ -22,25 +22,30 @@ 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))
 
     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:

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(

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<Page> {
+    const pageRef = await this.rpcClient.send(
+      StagehandMethods.contextAwaitActivePage,
+      timeoutMs === undefined ? {} : { timeout: timeoutMs },
+    );
+    return new Page(this.rpcClient, pageRef);
+  }
+
   async setActivePage(page: Page): Promise<void> {
     await this.rpcClient.send(StagehandMethods.contextSetActivePage, {
       pageId: page.pageId,

packages/sdk-ts/src/rpcClient.ts

@@ -47,7 +47,6 @@ type PendingRequest = {
   method: RPCMethod;
   resolve(value: unknown): void;
   reject(error: Error): void;
-  timeout: ReturnType<typeof setTimeout>;
 };
 
 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<unknown> {
     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<RPCClie
     cdpConnectTimeoutMs: options.cdpConnectTimeoutMs ?? 10_000,
     commandTimeoutMs,
   });
-  const client = new RPCClient(cdpClient, commandTimeoutMs);
+  const client = new RPCClient(cdpClient);
 
   try {
     await client.send(StagehandMethods.runtimeConfigure, {