diff --git a/packages/integrations/pi/extensions/stagehand.ts b/packages/integrations/pi/extensions/stagehand.ts index 9fba1deb1..f29602877 100644 --- a/packages/integrations/pi/extensions/stagehand.ts +++ b/packages/integrations/pi/extensions/stagehand.ts @@ -24,6 +24,7 @@ import { SNAPSHOT_TOOL_DESCRIPTION, SnapshotInputSchema, StagehandFacadeTools, + releaseBrowserbaseSession, stagehandFacadeConfigFromEnv, } from "@browserbasehq/stagehand-integrations/facade"; @@ -31,8 +32,17 @@ type FacadeResources = { browser: StagehandBrowser; stagehand: Stagehand; tools: StagehandFacadeTools; + releaseSession?: () => Promise; }; +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 | undefined; + let cleanupPromise: Promise = Promise.resolve(); + const cleanupTargets = new Set<() => Promise>(); async function facadeTools(): Promise { 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) | 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 { + async function closeResources( + expected?: StagehandFacadeTools, + reportErrors = false, + ): Promise { // 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 { + 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 { + for (const releaseSession of cleanupTargets) { + await releaseSession(); + cleanupTargets.delete(releaseSession); + } } - pi.on("session_shutdown", closeResources); + pi.on("session_shutdown", () => closeResources()); pi.registerTool({ name: "run", diff --git a/packages/integrations/pi/tests/lifecycle.test.ts b/packages/integrations/pi/tests/lifecycle.test.ts new file mode 100644 index 000000000..65a16f1df --- /dev/null +++ b/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(); + 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 }, + ) {} + + async run(): Promise { + await this.lifecycle.close(); + return "closed"; + } + + async runActions(): Promise { + return {}; + } + + async snapshot(): Promise { + 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; +}; + +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) | undefined; +}> { + const { default: stagehandExtension } = await import("../extensions/stagehand.js"); + const tools: RegisteredTool[] = []; + let shutdown: (() => Promise) | undefined; + stagehandExtension({ + registerTool: (tool: RegisteredTool) => tools.push(tool), + on: (event: string, handler: () => Promise) => { + if (event === "session_shutdown") shutdown = handler; + }, + } as never); + return { tools, shutdown }; +}