Skip to content

Commit 270f666

Browse files
miguelg719miguelg719
authored
stagehand.context pages as stagehand pages and supporting more than one page (#508)
* exposing stagehand.context pages as stagehand pages and supporting more than one page * fixing linting * removed scratchpad file * updated active page logic to switch between different calls * updated deterministic evals * moved types to types folder and fixed linting * fixed build issue * fixed imports in e2e tests * made eval actually deterministic * removing unused import * added delete of uninitialized stagehand error before proxy creation * reviewed changes/updates * missing protected method
1 parent a25a4cb commit 270f666

File tree

6 files changed

+442
-110
lines changed

6 files changed

+442
-110
lines changed

Diff for: ‎.changeset/empty-singers-warn.md

+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"@browserbasehq/stagehand": patch
3+
---
4+
5+
Fixed stagehand to support multiple pages with an enhanced context

Diff for: ‎evals/deterministic/tests/BrowserContext/multiPage.test.ts

+209
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,209 @@
1+
import { test, expect } from "@playwright/test";
2+
import { Stagehand } from "@/dist";
3+
import StagehandConfig from "@/evals/deterministic/stagehand.config";
4+
import { Page } from "@/dist";
5+
6+
import http from "http";
7+
import express from "express";
8+
import { Server as WebSocketServer } from "ws";
9+
10+
test.describe("StagehandContext - Multi-page Support", () => {
11+
let stagehand: Stagehand;
12+
let server: http.Server;
13+
let wss: WebSocketServer;
14+
let serverPort: number;
15+
16+
test.beforeAll(async () => {
17+
// Set up a local Express server
18+
const app = express();
19+
20+
// Serve test pages
21+
app.get("/page1", (_req, res) => {
22+
res.set("Content-Type", "text/html");
23+
res.end(`
24+
<html>
25+
<head><title>Page 1</title></head>
26+
<body>
27+
<h1>Page 1 Content</h1>
28+
<button id="popupBtn" onclick="window.open('/page2', '_blank')">Open Page 2</button>
29+
</body>
30+
</html>
31+
`);
32+
});
33+
34+
app.get("/page2", (_req, res) => {
35+
res.set("Content-Type", "text/html");
36+
res.end(`
37+
<html>
38+
<head><title>Page 2</title></head>
39+
<body>
40+
<h1>Page 2 Content</h1>
41+
</body>
42+
</html>
43+
`);
44+
});
45+
46+
// Create the server on a random free port
47+
server = http.createServer(app);
48+
await new Promise<void>((resolve) => {
49+
server.listen(0, () => resolve());
50+
});
51+
const address = server.address();
52+
if (typeof address === "object" && address !== null) {
53+
serverPort = address.port;
54+
} else {
55+
throw new Error("Failed to get server port");
56+
}
57+
58+
// Set up WebSocket for future tests
59+
wss = new WebSocketServer({ server, path: "/socket" });
60+
wss.on("connection", (ws) => {
61+
console.log("WebSocket client connected");
62+
ws.send("Hello from server WebSocket");
63+
});
64+
});
65+
66+
test.beforeEach(async () => {
67+
stagehand = new Stagehand(StagehandConfig);
68+
await stagehand.init();
69+
});
70+
71+
test.afterEach(async () => {
72+
await stagehand.close();
73+
});
74+
75+
test.afterAll(async () => {
76+
wss?.close();
77+
server?.close();
78+
});
79+
80+
/**
81+
* Test enhanced page capabilities
82+
*/
83+
test("should provide enhanced capabilities for new pages", async () => {
84+
const context = stagehand.context;
85+
const newPage = await context.newPage();
86+
87+
// Verify enhanced methods
88+
expect(typeof newPage.act).toBe("function");
89+
expect(typeof newPage.extract).toBe("function");
90+
expect(typeof newPage.observe).toBe("function");
91+
92+
// Verify basic Playwright functionality
93+
expect(typeof newPage.goto).toBe("function");
94+
expect(typeof newPage.click).toBe("function");
95+
96+
// Test navigation maintains capabilities
97+
await newPage.goto(`http://localhost:${serverPort}/page1`);
98+
expect(typeof newPage.act).toBe("function");
99+
expect(await newPage.title()).toBe("Page 1");
100+
});
101+
102+
/**
103+
* Test context.pages() functionality
104+
*/
105+
test("should return array of enhanced pages via context.pages()", async () => {
106+
const context = stagehand.context;
107+
108+
// Create multiple pages
109+
const page1 = await context.newPage();
110+
const page2 = await context.newPage();
111+
112+
await page1.goto(`http://localhost:${serverPort}/page1`);
113+
await page2.goto(`http://localhost:${serverPort}/page2`);
114+
115+
const pages = context.pages();
116+
expect(pages).toContain(page1);
117+
expect(pages).toContain(page2);
118+
119+
// Verify all pages have enhanced capabilities
120+
for (const page of pages) {
121+
expect(typeof page.act).toBe("function");
122+
expect(typeof page.extract).toBe("function");
123+
expect(typeof page.observe).toBe("function");
124+
}
125+
});
126+
127+
/**
128+
* Test popup handling
129+
*/
130+
test("should handle popups with enhanced capabilities", async () => {
131+
const mainPage = stagehand.page;
132+
let popupPage: Page | null = null;
133+
134+
mainPage.on("popup", (page: Page) => {
135+
popupPage = page;
136+
});
137+
138+
await mainPage.goto(`http://localhost:${serverPort}/page1`);
139+
await mainPage.click("#popupBtn");
140+
141+
// Verify popup has enhanced capabilities
142+
expect(popupPage).not.toBeNull();
143+
expect(typeof popupPage.act).toBe("function");
144+
expect(typeof popupPage.extract).toBe("function");
145+
expect(typeof popupPage.observe).toBe("function");
146+
147+
if (popupPage) {
148+
await popupPage.waitForLoadState();
149+
expect(await popupPage.title()).toBe("Page 2");
150+
}
151+
});
152+
153+
/**
154+
* Test page tracking and cleanup
155+
*/
156+
test("should properly track and cleanup pages", async () => {
157+
const context = stagehand.context;
158+
const initialPages = context.pages().length;
159+
160+
const newPage = await context.newPage();
161+
await newPage.goto(`http://localhost:${serverPort}/page1`);
162+
163+
expect(context.pages().length).toBe(initialPages + 1);
164+
await newPage.close();
165+
expect(context.pages().length).toBe(initialPages);
166+
});
167+
168+
/**
169+
* Test enhanced methods across pages
170+
*/
171+
test("should support enhanced methods across all pages", async () => {
172+
const page1 = await stagehand.context.newPage();
173+
const page2 = await stagehand.context.newPage();
174+
175+
await page1.goto(`http://localhost:${serverPort}/page1`);
176+
await page2.goto(`http://localhost:${serverPort}/page2`);
177+
178+
// Verify both pages have enhanced capabilities
179+
expect(typeof page1.act).toBe("function");
180+
expect(typeof page1.extract).toBe("function");
181+
expect(typeof page1.observe).toBe("function");
182+
183+
expect(typeof page2.act).toBe("function");
184+
expect(typeof page2.extract).toBe("function");
185+
expect(typeof page2.observe).toBe("function");
186+
});
187+
188+
/**
189+
* Test active page tracking
190+
*/
191+
test("should update stagehand.page when creating new pages", async () => {
192+
const initialPage = stagehand.page;
193+
194+
// Create a new page and verify it becomes active
195+
const newPage = await stagehand.context.newPage();
196+
expect(stagehand.page).toBe(newPage);
197+
expect(stagehand.page).not.toBe(initialPage);
198+
199+
// Navigate and verify it's still the active page
200+
await newPage.goto(`http://localhost:${serverPort}/page1`);
201+
expect(stagehand.page).toBe(newPage);
202+
expect(await stagehand.page.title()).toBe("Page 1");
203+
204+
// Create another page and verify it becomes active
205+
const anotherPage = await stagehand.context.newPage();
206+
expect(stagehand.page).toBe(anotherPage);
207+
expect(stagehand.page).not.toBe(newPage);
208+
});
209+
});

Diff for: ‎lib/StagehandContext.ts

+106-10
Original file line numberDiff line numberDiff line change
@@ -1,29 +1,125 @@
1-
import type { BrowserContext as PlaywrightContext } from "@playwright/test";
1+
import type {
2+
BrowserContext as PlaywrightContext,
3+
Page as PlaywrightPage,
4+
} from "@playwright/test";
25
import { Stagehand } from "./index";
6+
import { StagehandPage } from "./StagehandPage";
7+
import { Page } from "../types/page";
8+
import { EnhancedContext } from "../types/context";
39

410
export class StagehandContext {
511
private readonly stagehand: Stagehand;
6-
private readonly intContext: PlaywrightContext;
12+
private readonly intContext: EnhancedContext;
13+
private pageMap: WeakMap<PlaywrightPage, StagehandPage>;
14+
private activeStagehandPage: StagehandPage | null = null;
715

816
private constructor(context: PlaywrightContext, stagehand: Stagehand) {
9-
this.intContext = context;
1017
this.stagehand = stagehand;
18+
this.pageMap = new WeakMap();
19+
20+
// Create proxy around the context
21+
this.intContext = new Proxy(context, {
22+
get: (target, prop) => {
23+
if (prop === "newPage") {
24+
return async (): Promise<Page> => {
25+
const pwPage = await target.newPage();
26+
const stagehandPage = await this.createStagehandPage(pwPage);
27+
// Set as active page when created
28+
this.setActivePage(stagehandPage);
29+
return stagehandPage.page;
30+
};
31+
}
32+
if (prop === "pages") {
33+
return (): Page[] => {
34+
const pwPages = target.pages();
35+
// Convert all pages to StagehandPages synchronously
36+
return pwPages.map((pwPage: PlaywrightPage) => {
37+
let stagehandPage = this.pageMap.get(pwPage);
38+
if (!stagehandPage) {
39+
// Create a new StagehandPage and store it in the map
40+
stagehandPage = new StagehandPage(
41+
pwPage,
42+
this.stagehand,
43+
this,
44+
this.stagehand.llmClient,
45+
this.stagehand.userProvidedInstructions,
46+
this.stagehand.apiClient,
47+
this.stagehand.waitForCaptchaSolves,
48+
);
49+
this.pageMap.set(pwPage, stagehandPage);
50+
}
51+
return stagehandPage.page;
52+
});
53+
};
54+
}
55+
return target[prop as keyof PlaywrightContext];
56+
},
57+
}) as unknown as EnhancedContext;
58+
}
59+
60+
private async createStagehandPage(
61+
page: PlaywrightPage,
62+
): Promise<StagehandPage> {
63+
const stagehandPage = await new StagehandPage(
64+
page,
65+
this.stagehand,
66+
this,
67+
this.stagehand.llmClient,
68+
this.stagehand.userProvidedInstructions,
69+
this.stagehand.apiClient,
70+
this.stagehand.waitForCaptchaSolves,
71+
).init();
72+
this.pageMap.set(page, stagehandPage);
73+
return stagehandPage;
1174
}
1275

1376
static async init(
1477
context: PlaywrightContext,
1578
stagehand: Stagehand,
1679
): Promise<StagehandContext> {
17-
const proxyContext = new Proxy(context, {
18-
get: (target, prop) => {
19-
return target[prop as keyof PlaywrightContext];
20-
},
21-
});
22-
const instance = new StagehandContext(proxyContext, stagehand);
80+
const instance = new StagehandContext(context, stagehand);
81+
82+
// Initialize existing pages
83+
const existingPages = context.pages();
84+
for (const page of existingPages) {
85+
const stagehandPage = await instance.createStagehandPage(page);
86+
// Set the first page as active
87+
if (!instance.activeStagehandPage) {
88+
instance.setActivePage(stagehandPage);
89+
}
90+
}
91+
2392
return instance;
2493
}
2594

26-
public get context(): PlaywrightContext {
95+
public get context(): EnhancedContext {
2796
return this.intContext;
2897
}
98+
99+
public async getStagehandPage(page: PlaywrightPage): Promise<StagehandPage> {
100+
let stagehandPage = this.pageMap.get(page);
101+
if (!stagehandPage) {
102+
stagehandPage = await this.createStagehandPage(page);
103+
}
104+
// Update active page when getting a page
105+
this.setActivePage(stagehandPage);
106+
return stagehandPage;
107+
}
108+
109+
public async getStagehandPages(): Promise<StagehandPage[]> {
110+
const pwPages = this.intContext.pages();
111+
return Promise.all(
112+
pwPages.map((page: PlaywrightPage) => this.getStagehandPage(page)),
113+
);
114+
}
115+
116+
public setActivePage(page: StagehandPage): void {
117+
this.activeStagehandPage = page;
118+
// Update the stagehand's active page reference
119+
this.stagehand["setActivePage"](page);
120+
}
121+
122+
public getActivePage(): StagehandPage | null {
123+
return this.activeStagehandPage;
124+
}
29125
}

0 commit comments

Comments
 (0)