Add opt-in resident worker lifecycle

packages/server/controllers/contextController.ts

@@ -18,7 +18,10 @@ import type {
 import type { HandlerContext } from "../rpcRouter.js";
 import type { StagehandRuntime } from "../runtime.js";
 
-export function createContextController(runtime: StagehandRuntime) {
+export function createContextController(
+  runtime: StagehandRuntime,
+  options: { close?: () => Promise<void> } = {},
+) {
   async function pages(_params: EmptyParams, { logger }: HandlerContext) {
     logger.debug("context.pages", {});
     return runtime.contextPages();
@@ -41,7 +44,11 @@ export function createContextController(runtime: StagehandRuntime) {
 
   async function close(_params: EmptyParams, { logger }: HandlerContext) {
     logger.debug("context.close", {});
-    return runtime.contextClose();
+    if (options.close) {
+      await options.close();
+      return { closed: true };
+    }
+    return await runtime.contextClose();
   }
 
   async function addInitScript(params: ContextAddInitScriptParams, { logger }: HandlerContext) {

packages/server/rpcRouter.ts

@@ -34,6 +34,7 @@ export type HandlerContext = {
 export type RPCRouterOptions = {
   initializeStagehand?: (params: StagehandInitParams) => Promise<StagehandInitResult>;
   closeStagehand?: () => Promise<void>;
+  closeContext?: () => Promise<void>;
 };
 
 export class RPCRouter {
@@ -54,7 +55,10 @@ export class RPCRouter {
       ...(options.initializeStagehand ? { initialize: options.initializeStagehand } : {}),
       ...(options.closeStagehand ? { close: options.closeStagehand } : {}),
     });
-    this.contextController = createContextController(runtime);
+    this.contextController = createContextController(
+      runtime,
+      options.closeContext ? { close: options.closeContext } : {},
+    );
     this.pageController = createPageController(runtime);
     this.locatorController = createLocatorController(runtime);
   }

packages/server/runtime.ts

@@ -216,8 +216,15 @@ export type StagehandBrowserSession = {
 export type StagehandBrowserSessionFactory = (
   cdpUrl: string,
   logger: StagehandLogger,
+  lifecycle?: StagehandBrowserSessionLifecycle,
 ) => Promise<StagehandBrowserSession>;
 
+export type StagehandBrowserSessionLifecycle = {
+  bootstrapMode?: "resident";
+  onConnected?(): void;
+  onDisconnected?(): void;
+};
+
 export type StagehandRuntimeAdapters = {
   browserSessionFactory?: StagehandBrowserSessionFactory;
   emitLog?: StagehandLogEmitter;
@@ -255,7 +262,20 @@ export class StagehandRuntime {
   );
   browserSession?: StagehandBrowserSession;
   pagesById = new Map<string, UnderstudyRuntimePage>();
+  private browserSessionGeneration = 0;
   private initializationInProgress = false;
+  private readonly contextInitScripts: string[] = [];
+  private contextExtraHTTPHeaders?: ContextSetExtraHTTPHeadersParams["headers"];
+  private contextDomainPolicy?: DomainPolicy | null;
+  private readonly pageInitScriptsById = new Map<string, string[]>();
+  private readonly pageExtraHTTPHeadersById = new Map<
+    string,
+    PageSetExtraHTTPHeadersParams["headers"]
+  >();
+  private readonly pageViewportById = new Map<
+    string,
+    Pick<PageSetViewportSizeParams, "width" | "height" | "options">
+  >();
 
   constructor(
     readonly adapters: ResolvedStagehandRuntimeAdapters,
@@ -264,18 +284,37 @@ export class StagehandRuntime {
     this.logger = new StagehandLogger(tracing, adapters.emitLog);
   }
 
-  async replaceBrowserConnection(params: { cdpUrl: string }): Promise<void> {
+  browserConnectionStatus(): { configured: boolean; connected: boolean } {
+    return {
+      configured: this.browserSession !== undefined,
+      connected: this.browserSession?.connected ?? false,
+    };
+  }
+
+  async replaceBrowserConnection(
+    params: { cdpUrl: string },
+    lifecycle?: StagehandBrowserSessionLifecycle,
+  ): Promise<void> {
     const { cdpUrl } = params;
+    const generation = ++this.browserSessionGeneration;
     const previousSession = this.browserSession;
     this.browserSession = undefined;
     this.pagesById.clear();
     await previousSession?.close();
 
+    let browserSession: StagehandBrowserSession | undefined;
     try {
-      this.browserSession = await this.adapters.browserSessionFactory(cdpUrl, this.logger);
+      if (generation !== this.browserSessionGeneration) {
+        throw new Error("Stagehand browser session bootstrap was superseded");
+      }
+      browserSession = await this.adapters.browserSessionFactory(cdpUrl, this.logger, lifecycle);
+      if (generation !== this.browserSessionGeneration) {
+        throw new Error("Stagehand browser session bootstrap was superseded");
+      }
+      this.browserSession = browserSession;
     } catch (error) {
-      await this.browserSession?.close();
-      this.browserSession = undefined;
+      await browserSession?.close();
+      if (generation === this.browserSessionGeneration) this.browserSession = undefined;
       throw error;
     }
   }
@@ -317,6 +356,32 @@ export class StagehandRuntime {
     }
   }
 
+  /** Restores initialization-dependent browser instrumentation after replacing a CDP session. */
+  async restoreInitializedBrowserSession(): Promise<void> {
+    if (this.state.getState().status !== "initialized") return;
+    const session = this.requireBrowserSession();
+    await session.prepareForInitialization?.();
+    for (const source of this.contextInitScripts) await session.addInitScript(source);
+    if (this.contextExtraHTTPHeaders) {
+      await session.setExtraHTTPHeaders(this.contextExtraHTTPHeaders);
+    }
+    if (this.contextDomainPolicy !== undefined) {
+      await session.setDomainPolicy(this.contextDomainPolicy);
+    }
+    await this.contextPages();
+    for (const [pageId, page] of this.pagesById) {
+      for (const source of this.pageInitScriptsById.get(pageId) ?? []) {
+        await page.addInitScript(source);
+      }
+      const headers = this.pageExtraHTTPHeadersById.get(pageId);
+      if (headers) await page.setExtraHTTPHeaders(headers);
+      const viewport = this.pageViewportById.get(pageId);
+      if (viewport) {
+        await page.setViewportSize(viewport.width, viewport.height, viewport.options);
+      }
+    }
+  }
+
   async browserGetVersion(): Promise<BrowserGetVersionResult> {
     return await this.requireBrowserSession().getVersion();
   }
@@ -363,13 +428,16 @@ export class StagehandRuntime {
 
   async contextAddInitScript(params: ContextAddInitScriptParams): Promise<ContextVoidResult> {
     await this.requireBrowserSession().addInitScript(params.source);
+    if (!this.contextInitScripts.includes(params.source))
+      this.contextInitScripts.push(params.source);
     return { ok: true };
   }
 
   async contextSetExtraHTTPHeaders(
     params: ContextSetExtraHTTPHeadersParams,
   ): Promise<ContextVoidResult> {
     await this.requireBrowserSession().setExtraHTTPHeaders(params.headers);
+    this.contextExtraHTTPHeaders = { ...params.headers };
     return { ok: true };
   }
 
@@ -379,6 +447,7 @@ export class StagehandRuntime {
 
   async contextSetDomainPolicy(params: ContextSetDomainPolicyParams): Promise<ContextVoidResult> {
     await this.requireBrowserSession().setDomainPolicy(params.policy);
+    this.contextDomainPolicy = params.policy;
     return { ok: true };
   }
 
@@ -515,11 +584,15 @@ export class StagehandRuntime {
 
   async pageAddInitScript(params: PageAddInitScriptParams): Promise<PageVoidResult> {
     await this.resolvePage(params.pageId).addInitScript(params.source);
+    const sources = this.pageInitScriptsById.get(params.pageId) ?? [];
+    if (!sources.includes(params.source)) sources.push(params.source);
+    this.pageInitScriptsById.set(params.pageId, sources);
     return { ok: true };
   }
 
   async pageSetExtraHTTPHeaders(params: PageSetExtraHTTPHeadersParams): Promise<PageVoidResult> {
     await this.resolvePage(params.pageId).setExtraHTTPHeaders(params.headers);
+    this.pageExtraHTTPHeadersById.set(params.pageId, { ...params.headers });
     return { ok: true };
   }
 
@@ -529,6 +602,11 @@ export class StagehandRuntime {
       params.height,
       params.options,
     );
+    this.pageViewportById.set(params.pageId, {
+      width: params.width,
+      height: params.height,
+      ...(params.options === undefined ? {} : { options: { ...params.options } }),
+    });
     return { ok: true };
   }
 
@@ -595,6 +673,9 @@ export class StagehandRuntime {
     const page = this.resolvePage(params.pageId);
     await page.close();
     this.pagesById.delete(params.pageId);
+    this.pageInitScriptsById.delete(params.pageId);
+    this.pageExtraHTTPHeadersById.delete(params.pageId);
+    this.pageViewportById.delete(params.pageId);
     return { closed: true };
   }
 
@@ -688,9 +769,13 @@ export class StagehandRuntime {
   }
 
   async close(): Promise<void> {
+    ++this.browserSessionGeneration;
     const session = this.browserSession;
     this.browserSession = undefined;
     this.pagesById.clear();
+    this.pageInitScriptsById.clear();
+    this.pageExtraHTTPHeadersById.clear();
+    this.pageViewportById.clear();
     try {
       await session?.close();
     } finally {

packages/server/service-worker-lifecycle/heartbeat-manager.ts

@@ -47,7 +47,6 @@ export function createServiceWorkerHeartbeatManager(
   let creatingDocument: Promise<void> | null = null;
   let heartbeatPort: RuntimePort | null = null;
   let installed = false;
-
   async function ensure(): Promise<void> {
     const offscreen = chromeApi.offscreen;
     if (offscreen == null) return;