[2.0] Wrap StagehandPage and StagehandContext (#319)

* [feat]: start to wrap page

* ignore pnpm lock

* wrap

* temp

* rewrite example

* sample playwright

* working with context

* e2e pass context

* rm example

* uploads/downloads e2e

* ci e2e

* CI env

* changeset

* cleanup context test

.changeset/nervous-dolls-clean.md

@@ -0,0 +1,5 @@
+---
+"@browserbasehq/stagehand": patch
+---
+
+We now wrap playwright page/context within StagehandPage and StagehandContext objects. This helps us augment the Stagehand experience by being able to augment the underlying Playwright

.github/workflows/ci.yml

@@ -55,10 +55,39 @@ jobs:
       - name: Run Build
         run: npm run build
 
-  run-extract-evals:
+  run-e2e-tests:
     needs: [run-lint, run-build]
     runs-on: ubuntu-latest
     timeout-minutes: 50
+    env:
+      OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+      ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+      BROWSERBASE_API_KEY: ${{ secrets.BROWSERBASE_API_KEY }}
+      BROWSERBASE_PROJECT_ID: ${{ secrets.BROWSERBASE_PROJECT_ID }}
+      HEADLESS: true
+
+    steps:
+      - name: Check out repository code
+        uses: actions/checkout@v4
+
+      - name: Set up Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version: "20"
+
+      - name: Install dependencies
+        run: npm install --no-frozen-lockfile
+
+      - name: Install Playwright browsers
+        run: npm exec playwright install --with-deps
+
+      - name: Run E2E Tests
+        run: npm run e2e
+
+  run-extract-evals:
+    needs: [run-lint, run-build, run-e2e-tests]
+    runs-on: ubuntu-latest
+    timeout-minutes: 50
     env:
       OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
       ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}

.gitignore

@@ -16,3 +16,4 @@ evals/public
 evals/playground.ts
 tmp/
 eval-summary.json
+pnpm-lock.yaml

evals/deterministic/playwright.config.ts

@@ -0,0 +1,31 @@
+import { defineConfig, devices } from "@playwright/test";
+
+/**
+ * See https://playwright.dev/docs/test-configuration.
+ */
+export default defineConfig({
+  testDir: "./tests",
+
+  /* Fail the build on CI if you accidentally left test.only in the source code. */
+  /* Run tests in files in parallel */
+  fullyParallel: true,
+  /* Reporter to use. See https://playwright.dev/docs/test-reporters */
+  //   reporter: "html",
+  reporter: "line",
+  /* Retry on CI only */
+  retries: 2,
+
+  /* Shared settings for all the projects below. See https://playwright.dev/docs/api/class-testoptions. */
+  use: {
+    /* Collect trace when retrying the failed test. See https://playwright.dev/docs/trace-viewer */
+    trace: "on-first-retry",
+  },
+
+  /* Configure projects for major browsers */
+  projects: [
+    {
+      name: "chromium",
+      use: { ...devices["Desktop Chrome"] },
+    },
+  ],
+});

evals/deterministic/stagehand.config.ts

@@ -0,0 +1,26 @@
+import type { ConstructorParams, LogLine } from "../../lib";
+
+const StagehandConfig: ConstructorParams = {
+  env: "BROWSERBASE" /* Environment to run Stagehand in */,
+  apiKey: process.env.BROWSERBASE_API_KEY /* API key for authentication */,
+  projectId: process.env.BROWSERBASE_PROJECT_ID /* Project identifier */,
+  verbose: 1 /* Logging verbosity level (0=quiet, 1=normal, 2=verbose) */,
+  debugDom: true /* Enable DOM debugging features */,
+  headless: false /* Run browser in headless mode */,
+  logger: (message: LogLine) =>
+    console.log(
+      `[stagehand::${message.category}] ${message.message}`,
+    ) /* Custom logging function */,
+  domSettleTimeoutMs: 30_000 /* Timeout for DOM to settle in milliseconds */,
+  browserbaseSessionCreateParams: {
+    projectId: process.env.BROWSERBASE_PROJECT_ID!,
+  },
+  enableCaching: true /* Enable caching functionality */,
+  browserbaseSessionID:
+    undefined /* Session ID for resuming Browserbase sessions */,
+  modelName: "gpt-4o" /* Name of the model to use */,
+  modelClientOptions: {
+    apiKey: process.env.OPENAI_API_KEY,
+  } /* Configuration options for the model client */,
+};
+export default StagehandConfig;

evals/deterministic/tests/contexts.test.ts

@@ -0,0 +1,145 @@
+import Browserbase from "@browserbasehq/sdk";
+import { Stagehand } from "../../../lib";
+import { expect, test } from "@playwright/test";
+import StagehandConfig from "../stagehand.config";
+
+// Configuration
+const CONTEXT_TEST_URL = "https://docs.browserbase.com";
+const BROWSERBASE_PROJECT_ID = process.env["BROWSERBASE_PROJECT_ID"]!;
+const BROWSERBASE_API_KEY = process.env["BROWSERBASE_API_KEY"]!;
+
+const bb = new Browserbase({
+  apiKey: BROWSERBASE_API_KEY,
+});
+
+// Helper functions
+function addHour(date: Date): number {
+  const SECOND = 1000;
+  return new Date(date.getTime() + 60 * 60 * 1000).getTime() / SECOND;
+}
+
+async function findCookie(stagehand: Stagehand, name: string) {
+  const defaultContext = stagehand.context;
+  const cookies = await defaultContext?.cookies();
+  return cookies?.find((cookie) => cookie.name === name);
+}
+
+async function createContext() {
+  console.log("Creating a new context...");
+  const context = await bb.contexts.create({
+    projectId: BROWSERBASE_PROJECT_ID,
+  });
+  const contextId = context.id;
+  console.log(`Context created with ID: ${contextId}`);
+  return contextId;
+}
+
+async function setRandomCookie(contextId: string, stagehand: Stagehand) {
+  console.log(
+    `Populating context ${contextId} during session ${stagehand.browserbaseSessionID}`,
+  );
+  const page = stagehand.page;
+
+  await page.goto(CONTEXT_TEST_URL, { waitUntil: "domcontentloaded" });
+
+  const now = new Date();
+  const testCookieName = `bb_${now.getTime().toString()}`;
+  const testCookieValue = now.toISOString();
+
+  await stagehand.context.addCookies([
+    {
+      domain: `.${new URL(CONTEXT_TEST_URL).hostname}`,
+      expires: addHour(now),
+      name: testCookieName,
+      path: "/",
+      value: testCookieValue,
+    },
+  ]);
+
+  expect(findCookie(stagehand, testCookieName)).toBeDefined();
+  console.log(`Set test cookie: ${testCookieName}=${testCookieValue}`);
+  return { testCookieName, testCookieValue };
+}
+
+test.describe("Contexts", () => {
+  test("Persists and re-uses a context", async () => {
+    let contextId: string;
+    let testCookieName: string;
+    let testCookieValue: string;
+    let stagehand: Stagehand;
+
+    await test.step("Create a context", async () => {
+      contextId = await createContext();
+    });
+
+    await test.step("Instantiate Stagehand with the context to persist", async () => {
+      // We will be adding cookies to the context in this session, so we need mark persist=true
+      stagehand = new Stagehand({
+        ...StagehandConfig,
+        browserbaseSessionCreateParams: {
+          projectId: BROWSERBASE_PROJECT_ID,
+          browserSettings: {
+            context: {
+              id: contextId,
+              persist: true,
+            },
+          },
+        },
+      });
+      await stagehand.init();
+    });
+
+    await test.step("Set a random cookie on the page", async () => {
+      ({ testCookieName } = await setRandomCookie(contextId, stagehand));
+
+      const page = stagehand.page;
+      await page.goto("https://www.google.com", {
+        waitUntil: "domcontentloaded",
+      });
+      await page.goBack();
+    });
+
+    await test.step("Validate cookie persistence between pages", async () => {
+      const cookie = await findCookie(stagehand, testCookieName);
+      const found = !!cookie;
+      expect(found).toBe(true);
+      console.log("Cookie persisted between pages:", found);
+
+      await stagehand.close();
+      // Wait for context to persist
+      console.log("Waiting for context to persist...");
+      await new Promise((resolve) => setTimeout(resolve, 5000));
+    });
+
+    await test.step("Create another session with the same context", async () => {
+      // We don't need to persist cookies in this session, so we can mark persist=false
+      const newStagehand = new Stagehand({
+        ...StagehandConfig,
+        browserbaseSessionCreateParams: {
+          projectId: BROWSERBASE_PROJECT_ID,
+          browserSettings: {
+            context: {
+              id: contextId,
+              persist: false,
+            },
+          },
+        },
+      });
+      await newStagehand.init();
+      console.log(
+        `Reusing context ${contextId} during session ${newStagehand.browserbaseSessionID}`,
+      );
+      const newPage = newStagehand.page;
+      await newPage.goto(CONTEXT_TEST_URL, { waitUntil: "domcontentloaded" });
+
+      const foundCookie = await findCookie(newStagehand, testCookieName);
+      console.log("Cookie found in new session:", !!foundCookie);
+      console.log(
+        "Cookie value matches:",
+        foundCookie?.value === testCookieValue,
+      );
+
+      await newStagehand.close();
+    });
+  });
+});

evals/deterministic/tests/downloads.test.ts

@@ -0,0 +1,69 @@
+import { test, expect } from "@playwright/test";
+import AdmZip from "adm-zip";
+import StagehandConfig from "../stagehand.config";
+import { Stagehand } from "../../../lib";
+import Browserbase from "@browserbasehq/sdk";
+
+const downloadRe = /sandstorm-(\d{13})+\.mp3/;
+
+test("Downloads", async () => {
+  const stagehand = new Stagehand(StagehandConfig);
+  await stagehand.init();
+  const page = stagehand.page;
+  const context = stagehand.context;
+
+  const client = await context.newCDPSession(page);
+  await client.send("Browser.setDownloadBehavior", {
+    behavior: "allow",
+    // `downloadPath` gets appended to the browser's default download directory.
+    // set to "downloads", it ends up being "/app/apps/browser/downloads/<file>".
+    downloadPath: "downloads",
+    eventsEnabled: true,
+  });
+
+  await page.goto("https://browser-tests-alpha.vercel.app/api/download-test");
+
+  const [download] = await Promise.all([
+    page.waitForEvent("download"),
+    page.locator("#download").click(),
+  ]);
+
+  const downloadError = await download.failure();
+
+  await stagehand.close();
+
+  if (downloadError !== null) {
+    throw new Error(
+      `Download for session ${stagehand.browserbaseSessionID} failed: ${downloadError}`,
+    );
+  }
+
+  expect(async () => {
+    const bb = new Browserbase();
+    const zipBuffer = await bb.sessions.downloads.list(
+      stagehand.browserbaseSessionID,
+    );
+    if (!zipBuffer) {
+      throw new Error(
+        `Download buffer is empty for session ${stagehand.browserbaseSessionID}`,
+      );
+    }
+
+    const zip = new AdmZip(Buffer.from(await zipBuffer.arrayBuffer()));
+    const zipEntries = zip.getEntries();
+    const mp3Entry = zipEntries.find((entry) =>
+      downloadRe.test(entry.entryName),
+    );
+
+    if (!mp3Entry) {
+      throw new Error(
+        `Session ${stagehand.browserbaseSessionID} is missing a file matching "${downloadRe.toString()}" in its zip entries: ${JSON.stringify(zipEntries.map((entry) => entry.entryName))}`,
+      );
+    }
+
+    const expectedFileSize = 6137541;
+    expect(mp3Entry.header.size).toBe(expectedFileSize);
+  }).toPass({
+    timeout: 30_000,
+  });
+});

evals/deterministic/tests/uploads.test.ts

@@ -0,0 +1,36 @@
+import { join } from "node:path";
+import { test, expect } from "@playwright/test";
+import { Stagehand } from "../../../lib";
+import StagehandConfig from "../stagehand.config";
+
+test.describe("Playwright Upload", () => {
+  let stagehand: Stagehand;
+
+  test.beforeAll(async () => {
+    stagehand = new Stagehand(StagehandConfig);
+    await stagehand.init();
+  });
+
+  test.afterAll(async () => {
+    await stagehand.close();
+  });
+
+  test("uploads a file", async () => {
+    const page = stagehand.page;
+    await page.goto("https://browser-tests-alpha.vercel.app/api/upload-test");
+
+    const fileInput = page.locator("#fileUpload");
+    await fileInput.setInputFiles(
+      join(__dirname, "..", "auxiliary", "logo.png"),
+    );
+
+    const fileNameSpan = page.locator("#fileName");
+    const fileName = await fileNameSpan.innerText();
+
+    const fileSizeSpan = page.locator("#fileSize");
+    const fileSize = Number(await fileSizeSpan.innerText());
+
+    expect(fileName).toBe("logo.png");
+    expect(fileSize).toBeGreaterThan(0);
+  });
+});