fix(pi): close facade browser resources

packages/integrations/pi/extensions/stagehand.ts

@@ -24,15 +24,25 @@ import {
   SNAPSHOT_TOOL_DESCRIPTION,
   SnapshotInputSchema,
   StagehandFacadeTools,
+  releaseBrowserbaseSession,
   stagehandFacadeConfigFromEnv,
 } from "@browserbasehq/stagehand-integrations/facade";
 
 type FacadeResources = {
   browser: StagehandBrowser;
   stagehand: Stagehand;
   tools: StagehandFacadeTools;
+  releaseSession?: () => Promise<void>;
 };
 
+export class StagehandFacadeCleanupError extends Error {
+  override readonly name = "StagehandFacadeCleanupError";
+
+  constructor() {
+    super("Failed to close the browser session cleanly.");
+  }
+}
+
 // TypeBox mirrors of the wire schemas (pi validates params with TypeBox; the
 // zod validators from the contract re-enforce semantics like code XOR actions
 // at execute time). Kept permissive on action items — the contract validator
@@ -65,24 +75,52 @@ export default function stagehandExtension(pi: ExtensionAPI) {
   // browser launches lazily on first tool call and closes on shutdown.
   let resources: FacadeResources | undefined;
   let resourcesPromise: Promise<FacadeResources> | undefined;
+  let cleanupPromise: Promise<void> = Promise.resolve();
+  const cleanupTargets = new Set<() => Promise<void>>();
 
   async function facadeTools(): Promise<StagehandFacadeTools> {
     if (resources && !resources.browser.closed) return resources.tools;
     resources = undefined;
-    resourcesPromise ??= (async () => {
+    resourcesPromise ??= cleanupPromise.then(retryCleanupTargets).then(async () => {
       const config = stagehandFacadeConfigFromEnv();
       const browser =
         config.browser.type === "browserbase"
           ? await browserbase.launch(config.browser.launchOptions)
           : await localBrowser.launch(config.browser.launchOptions);
+      const sessionId = browser.sessionId;
+      let releaseSession: (() => Promise<void>) | undefined;
+      if (config.browser.type === "browserbase" && sessionId) {
+        const { apiKey, baseUrl } = config.browser.launchOptions;
+        releaseSession = () => releaseBrowserbaseSession({ apiKey, baseUrl, sessionId });
+      }
       try {
         const stagehand = await Stagehand.create({ browser, ...config.stagehand });
-        return { browser, stagehand, tools: new StagehandFacadeTools(stagehand) };
+        let tools: StagehandFacadeTools;
+        tools = new StagehandFacadeTools(stagehand, {
+          close: () => closeResources(tools, true),
+        });
+        return { browser, stagehand, tools, ...(releaseSession ? { releaseSession } : {}) };
       } catch (error) {
-        await browser.close().catch(() => undefined);
-        throw error;
+        const cleanupErrors: unknown[] = [error];
+        let browserCloseFailed = false;
+        await browser.close().catch((cleanupError) => {
+          browserCloseFailed = true;
+          cleanupErrors.push(cleanupError);
+        });
+        if (browserCloseFailed && releaseSession) {
+          await releaseSession().catch((releaseError) => {
+            cleanupTargets.add(releaseSession);
+            cleanupErrors.push(releaseError);
+          });
+        }
+        if (cleanupErrors.length === 1) throw error;
+        throw new AggregateError(
+          cleanupErrors,
+          "Stagehand initialization failed and browser cleanup also failed.",
+          { cause: error },
+        );
       }
-    })();
+    });
     try {
       resources = await resourcesPromise;
       return resources.tools;
@@ -91,19 +129,64 @@ export default function stagehandExtension(pi: ExtensionAPI) {
     }
   }
 
-  async function closeResources(): Promise<void> {
+  async function closeResources(
+    expected?: StagehandFacadeTools,
+    reportErrors = false,
+  ): Promise<void> {
     // A shutdown can race a still-pending launch; wait for it so the browser
     // it produces is closed rather than leaked.
     const pending = resourcesPromise;
-    if (pending) await pending.catch(() => undefined);
-    const current = resources;
+    const launched = await pending?.catch(() => undefined);
+    const current = resources ?? launched;
+    if (expected && current?.tools !== expected) return;
     resources = undefined;
-    if (!current) return;
-    await current.stagehand.close().catch(() => undefined);
-    await current.browser.close().catch(() => undefined);
+    if (!current) {
+      if (!reportErrors) {
+        await cleanupPromise;
+        await retryCleanupTargets().catch(() => undefined);
+      }
+      return;
+    }
+    const closeResult = cleanupPromise.then(() => closeResource(current));
+    cleanupPromise = closeResult.then(
+      () => undefined,
+      () => undefined,
+    );
+    if (!reportErrors) {
+      await closeResult.catch(() => undefined);
+      await retryCleanupTargets().catch(() => undefined);
+      return;
+    }
+    await closeResult;
+  }
+
+  async function closeResource(current: FacadeResources): Promise<void> {
+    const cleanupErrors: unknown[] = [];
+    let browserCloseFailed = false;
+    await current.stagehand.close().catch((error) => cleanupErrors.push(error));
+    await current.browser.close().catch((error) => {
+      browserCloseFailed = true;
+      cleanupErrors.push(error);
+    });
+    if (browserCloseFailed && current.releaseSession) {
+      cleanupTargets.add(current.releaseSession);
+    } else if (current.releaseSession) {
+      cleanupTargets.delete(current.releaseSession);
+    }
+    if (cleanupErrors.length === 0) return;
+    // Pi surfaces rejected tool executions directly. Keep SDK/CDP details out
+    // of the model-visible error, including Error.cause and AggregateError.errors.
+    throw new StagehandFacadeCleanupError();
+  }
+
+  async function retryCleanupTargets(): Promise<void> {
+    for (const releaseSession of cleanupTargets) {
+      await releaseSession();
+      cleanupTargets.delete(releaseSession);
+    }
   }
 
-  pi.on("session_shutdown", closeResources);
+  pi.on("session_shutdown", () => closeResources());
 
   pi.registerTool({
     name: "run",

packages/integrations/pi/tests/lifecycle.test.ts

@@ -0,0 +1,185 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+const mocks = vi.hoisted(() => ({
+  launch: vi.fn(),
+  create: vi.fn(),
+}));
+
+vi.mock("@browserbasehq/stagehand", () => ({
+  browserbase: { launch: mocks.launch },
+  localBrowser: { launch: vi.fn() },
+  Stagehand: { create: mocks.create },
+}));
+
+vi.mock("@browserbasehq/stagehand-integrations/facade", async (importOriginal) => {
+  const original =
+    await importOriginal<typeof import("@browserbasehq/stagehand-integrations/facade")>();
+  return {
+    ...original,
+    stagehandFacadeConfigFromEnv: () => ({
+      browser: { type: "browserbase", launchOptions: { apiKey: "test-key" } },
+      stagehand: {},
+    }),
+    releaseBrowserbaseSession: async (session: { sessionId: string }) => {
+      const response = await fetch(`https://api.browserbase.com/v1/sessions/${session.sessionId}`, {
+        method: "POST",
+      });
+      if (!response.ok) throw new Error("Failed to release the Browserbase session.");
+    },
+    StagehandFacadeTools: class {
+      constructor(
+        _stagehand: unknown,
+        private readonly lifecycle: { close(): Promise<void> },
+      ) {}
+
+      async run(): Promise<string> {
+        await this.lifecycle.close();
+        return "closed";
+      }
+
+      async runActions(): Promise<object> {
+        return {};
+      }
+
+      async snapshot(): Promise<string> {
+        return "snapshot";
+      }
+
+      async screenshot(): Promise<{ data: string; mimeType: "image/png" }> {
+        return { data: "", mimeType: "image/png" };
+      }
+    },
+  };
+});
+
+type RegisteredTool = {
+  name: string;
+  execute(toolCallId: string, params: unknown): Promise<unknown>;
+};
+
+describe("pi stagehand lifecycle", () => {
+  beforeEach(() => {
+    mocks.launch.mockReset();
+    mocks.create.mockReset();
+  });
+
+  afterEach(() => vi.unstubAllGlobals());
+
+  it("surfaces cleanup failures and releases them before launching fresh resources", async () => {
+    const browser = {
+      closed: false,
+      sessionId: "session-one",
+      close: vi.fn(async () => {
+        throw new Error("browser close failed");
+      }),
+    };
+    const freshBrowser = {
+      closed: false,
+      sessionId: "session-two",
+      close: vi.fn(async () => undefined),
+    };
+    mocks.launch.mockResolvedValueOnce(browser).mockResolvedValueOnce(freshBrowser);
+    mocks.create
+      .mockResolvedValueOnce({ close: vi.fn(async () => undefined) })
+      .mockResolvedValueOnce({ close: vi.fn(async () => undefined) });
+    const release = vi
+      .fn()
+      .mockResolvedValueOnce({ ok: false })
+      .mockResolvedValueOnce({ ok: true });
+    vi.stubGlobal("fetch", release);
+    const { StagehandFacadeCleanupError } = await import("../extensions/stagehand.js");
+    const { tools } = await registerExtension();
+    const run = tools.find((tool) => tool.name === "run");
+
+    await expect(run?.execute("call-1", { code: "await browser.close();" })).rejects.toBeInstanceOf(
+      StagehandFacadeCleanupError,
+    );
+    const snapshot = tools.find((tool) => tool.name === "snapshot");
+    await expect(snapshot?.execute("call-2", {})).rejects.toThrow(
+      "Failed to release the Browserbase session.",
+    );
+    expect(mocks.launch).toHaveBeenCalledOnce();
+
+    await expect(snapshot?.execute("call-3", {})).resolves.toBeDefined();
+    expect(release).toHaveBeenCalledTimes(2);
+    expect(mocks.launch).toHaveBeenCalledTimes(2);
+  });
+
+  it("keeps session shutdown cleanup best-effort", async () => {
+    const browser = {
+      closed: false,
+      sessionId: "session-shutdown",
+      close: vi.fn(async () => {
+        throw new Error("browser close failed");
+      }),
+    };
+    mocks.launch.mockResolvedValueOnce(browser);
+    mocks.create.mockResolvedValueOnce({
+      close: vi.fn(async () => {
+        throw new Error("stagehand close failed");
+      }),
+    });
+    const release = vi.fn(async () => ({ ok: true }));
+    vi.stubGlobal("fetch", release);
+    const { tools, shutdown } = await registerExtension();
+    const snapshot = tools.find((tool) => tool.name === "snapshot");
+    await snapshot?.execute("call-1", {});
+
+    await expect(shutdown?.()).resolves.toBeUndefined();
+    expect(browser.close).toHaveBeenCalledOnce();
+    expect(release).toHaveBeenCalledOnce();
+  });
+
+  it("releases the Browserbase session when Stagehand initialization cleanup fails", async () => {
+    const browser = {
+      closed: false,
+      sessionId: "session-init-failed",
+      close: vi.fn(async () => {
+        throw new Error("browser close failed");
+      }),
+    };
+    const freshBrowser = {
+      closed: false,
+      sessionId: "session-after-init-failure",
+      close: vi.fn(async () => undefined),
+    };
+    mocks.launch.mockResolvedValueOnce(browser).mockResolvedValueOnce(freshBrowser);
+    mocks.create
+      .mockRejectedValueOnce(new Error("Stagehand init failed"))
+      .mockResolvedValueOnce({ close: vi.fn(async () => undefined) });
+    const release = vi
+      .fn()
+      .mockResolvedValueOnce({ ok: false })
+      .mockResolvedValueOnce({ ok: true });
+    vi.stubGlobal("fetch", release);
+    const { tools } = await registerExtension();
+    const snapshot = tools.find((tool) => tool.name === "snapshot");
+
+    await expect(snapshot?.execute("call-1", {})).rejects.toThrow(
+      "Stagehand initialization failed and browser cleanup also failed.",
+    );
+    expect(browser.close).toHaveBeenCalledOnce();
+    expect(release).toHaveBeenCalledOnce();
+    expect(mocks.launch).toHaveBeenCalledOnce();
+
+    await expect(snapshot?.execute("call-2", {})).resolves.toBeDefined();
+    expect(release).toHaveBeenCalledTimes(2);
+    expect(mocks.launch).toHaveBeenCalledTimes(2);
+  });
+});
+
+async function registerExtension(): Promise<{
+  tools: RegisteredTool[];
+  shutdown: (() => Promise<void>) | undefined;
+}> {
+  const { default: stagehandExtension } = await import("../extensions/stagehand.js");
+  const tools: RegisteredTool[] = [];
+  let shutdown: (() => Promise<void>) | undefined;
+  stagehandExtension({
+    registerTool: (tool: RegisteredTool) => tools.push(tool),
+    on: (event: string, handler: () => Promise<void>) => {
+      if (event === "session_shutdown") shutdown = handler;
+    },
+  } as never);
+  return { tools, shutdown };
+}