Skip to content

Commit e9e1e78

Browse files
Kylejeong2Kylejeong2
committed
stagehand tools using browser context similar to browserbase + stagehand init issues fixed
1 parent 2f58ed3 commit e9e1e78

File tree

8 files changed

+548
-99
lines changed

8 files changed

+548
-99
lines changed

‎stagehand/package.json‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,16 @@
77
"homepage": "https://modelcontextprotocol.io",
88
"bugs": "https://github.com/modelcontextprotocol/servers/issues",
99
"type": "module",
10+
"main": "./cli.js",
1011
"bin": {
11-
"mcp-server-stagehand": "dist/index.js"
12+
"mcp-server-stagehand": "cli.js"
1213
},
1314
"files": [
14-
"dist"
15+
"dist",
16+
"cli.js",
17+
"index.d.ts",
18+
"index.js",
19+
"config.d.ts"
1520
],
1621
"scripts": {
1722
"build": "tsc && shx chmod +x dist/*.js",

‎stagehand/src/config.ts‎

Lines changed: 1 addition & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -3,29 +3,10 @@ import fs from 'fs';
33
import path from 'path';
44
import { sanitizeForFilePath } from './tools/utils.js';
55
import type { Cookie } from "playwright-core";
6+
import type { Config } from '../config.js';
67

78
export type ToolCapability = 'core' | string;
89

9-
export interface Config {
10-
browserbaseApiKey?: string;
11-
browserbaseProjectId?: string;
12-
server?: {
13-
port?: number;
14-
host?: string;
15-
};
16-
proxies?: boolean;
17-
advancedStealth?: boolean;
18-
context?: {
19-
contextId?: string;
20-
persist?: boolean;
21-
};
22-
viewPort?: {
23-
browserWidth?: number;
24-
browserHeight?: number;
25-
};
26-
cookies?: Cookie[];
27-
}
28-
2910
// Define Command Line Options Structure
3011
export type CLIOptions = {
3112
browserbaseApiKey?: string;

‎stagehand/src/context.ts‎

Lines changed: 185 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { Stagehand } from "@browserbasehq/stagehand";
22
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
33
import type { Config } from "../config.js";
4-
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
4+
import { CallToolResult, TextContent, ImageContent } from "@modelcontextprotocol/sdk/types.js";
55
import { handleToolCall } from "./tools/tools.js";
66
import { listResources, readResource } from "./resources.js";
77
import {
@@ -12,11 +12,23 @@ import {
1212
setServerInstance,
1313
log
1414
} from "./logging.js";
15+
import {
16+
getSession,
17+
getSessionReadOnly,
18+
defaultSessionId,
19+
type BrowserSession
20+
} from "./sessionManager.js";
21+
22+
export type ToolActionResult =
23+
| { content?: (ImageContent | TextContent)[] }
24+
| undefined
25+
| void;
1526

1627
export class Context {
17-
private stagehand: Stagehand;
28+
private stagehands = new Map<string, Stagehand>();
29+
public readonly config: Config;
1830
private server: Server;
19-
private config: Config;
31+
public currentSessionId: string = defaultSessionId;
2032

2133
constructor(server: Server, config: Config) {
2234
this.server = server;
@@ -28,26 +40,171 @@ export class Context {
2840
setupLogRotation();
2941
registerExitHandlers();
3042
scheduleLogRotation();
43+
}
44+
45+
/**
46+
* Gets the Stagehand instance for the current session, creating one if needed
47+
*/
48+
public async getStagehand(sessionId: string = this.currentSessionId): Promise<Stagehand> {
49+
let stagehand = this.stagehands.get(sessionId);
50+
51+
if (!stagehand) {
52+
// Create a new Stagehand instance for this session
53+
stagehand = new Stagehand({
54+
env: "BROWSERBASE",
55+
logger: (logLine) => {
56+
log(`Stagehand[${sessionId}]: ${logLine.message}`, 'info');
57+
},
58+
});
59+
this.stagehands.set(sessionId, stagehand);
60+
}
61+
62+
await stagehand.init();
63+
64+
return stagehand;
65+
}
66+
67+
/**
68+
* Sets the Stagehand instance for a specific session
69+
*/
70+
public setStagehand(sessionId: string, stagehand: Stagehand): void {
71+
this.stagehands.set(sessionId, stagehand);
72+
}
73+
74+
/**
75+
* Removes the Stagehand instance for a specific session
76+
*/
77+
public async removeStagehand(sessionId: string): Promise<void> {
78+
const stagehand = this.stagehands.get(sessionId);
79+
if (stagehand) {
80+
try {
81+
await stagehand.close();
82+
} catch (error) {
83+
log(`Error closing Stagehand for session ${sessionId}: ${error}`, 'error');
84+
}
85+
this.stagehands.delete(sessionId);
86+
}
87+
}
3188

32-
// Initialize Stagehand
33-
this.stagehand = new Stagehand({
34-
env: "BROWSERBASE",
35-
logger: (logLine) => {
36-
log(`Stagehand: ${logLine.message}`, 'info');
37-
},
38-
});
89+
public async getActivePage(): Promise<BrowserSession["page"] | null> {
90+
// Try to get page from Stagehand first (if available for this session)
91+
const stagehand = this.stagehands.get(this.currentSessionId);
92+
if (stagehand && stagehand.page && !stagehand.page.isClosed()) {
93+
return stagehand.page;
94+
}
95+
96+
// Fallback to session manager
97+
const session = await getSession(this.currentSessionId, this.config);
98+
if (!session || !session.page || session.page.isClosed()) {
99+
try {
100+
const currentSession = await getSession(
101+
this.currentSessionId,
102+
this.config
103+
);
104+
if (
105+
!currentSession ||
106+
!currentSession.page ||
107+
currentSession.page.isClosed()
108+
) {
109+
return null;
110+
}
111+
return currentSession.page;
112+
} catch (refreshError) {
113+
return null;
114+
}
115+
}
116+
return session.page;
117+
}
118+
119+
// Will create a new default session if one doesn't exist
120+
public async getActiveBrowser(): Promise<BrowserSession["browser"] | null> {
121+
const session = await getSession(this.currentSessionId, this.config);
122+
if (!session || !session.browser || !session.browser.isConnected()) {
123+
try {
124+
const currentSession = await getSession(
125+
this.currentSessionId,
126+
this.config
127+
);
128+
if (
129+
!currentSession ||
130+
!currentSession.browser ||
131+
!currentSession.browser.isConnected()
132+
) {
133+
return null;
134+
}
135+
return currentSession.browser;
136+
} catch (refreshError) {
137+
return null;
138+
}
139+
}
140+
return session.browser;
141+
}
142+
143+
/**
144+
* Get the active browser without triggering session creation.
145+
* This is a read-only operation used when we need to check for an existing browser
146+
* without side effects (e.g., during close operations).
147+
* @returns The browser if it exists and is connected, null otherwise
148+
*/
149+
public getActiveBrowserReadOnly(): BrowserSession["browser"] | null {
150+
const session = getSessionReadOnly(this.currentSessionId);
151+
if (!session || !session.browser || !session.browser.isConnected()) {
152+
return null;
153+
}
154+
return session.browser;
155+
}
156+
157+
/**
158+
* Get the active page without triggering session creation.
159+
* This is a read-only operation used when we need to check for an existing page
160+
* without side effects.
161+
* @returns The page if it exists and is not closed, null otherwise
162+
*/
163+
public getActivePageReadOnly(): BrowserSession["page"] | null {
164+
const session = getSessionReadOnly(this.currentSessionId);
165+
if (!session || !session.page || session.page.isClosed()) {
166+
return null;
167+
}
168+
return session.page;
39169
}
40170

41171
async run(tool: any, args: any): Promise<CallToolResult> {
42172
try {
43-
log(`Executing tool: ${tool.name} with args: ${JSON.stringify(args)}`, 'info');
44-
const result = await handleToolCall(tool.name, args, this.stagehand);
45-
log(`Tool ${tool.name} completed successfully`, 'info');
46-
return result;
173+
log(`Executing tool: ${tool.schema.name} with args: ${JSON.stringify(args)}`, 'info');
174+
175+
// Check if this tool has a handle method (new session tools)
176+
// Only use handle method for session create and close tools
177+
if ("handle" in tool && typeof tool.handle === "function" &&
178+
(tool.schema.name === "browserbase_session_create" || tool.schema.name === "browserbase_session_close")) {
179+
const toolResult = await tool.handle(this as any, args);
180+
181+
if (toolResult?.action) {
182+
const actionResult = await toolResult.action();
183+
const content = actionResult?.content || [];
184+
185+
return {
186+
content: Array.isArray(content) ? content : [{ type: "text", text: "Action completed successfully." }],
187+
isError: false,
188+
};
189+
} else {
190+
return {
191+
content: [{ type: "text", text: `${tool.schema.name} completed successfully.` }],
192+
isError: false,
193+
};
194+
}
195+
} else {
196+
const stagehand = await this.getStagehand();
197+
const result = await handleToolCall(tool.schema.name, args, stagehand);
198+
log(`Tool ${tool.schema.name} completed successfully`, 'info');
199+
return result;
200+
}
47201
} catch (error) {
48202
const errorMessage = error instanceof Error ? error.message : String(error);
49-
log(`Tool ${tool.name} failed: ${errorMessage}`, 'error');
50-
throw error;
203+
log(`Tool ${tool.schema?.name || 'unknown'} failed: ${errorMessage}`, 'error');
204+
return {
205+
content: [{ type: "text", text: `Error: ${errorMessage}` }],
206+
isError: true,
207+
};
51208
}
52209
}
53210

@@ -61,10 +218,19 @@ export class Context {
61218

62219
async close() {
63220
try {
64-
await this.stagehand.close();
65-
log('Stagehand context closed successfully', 'info');
221+
// Close all Stagehand instances
222+
for (const [sessionId, stagehand] of this.stagehands.entries()) {
223+
try {
224+
await stagehand.close();
225+
log(`Closed Stagehand for session ${sessionId}`, 'info');
226+
} catch (error) {
227+
log(`Error closing Stagehand for session ${sessionId}: ${error}`, 'error');
228+
}
229+
}
230+
this.stagehands.clear();
231+
log('All Stagehand contexts closed successfully', 'info');
66232
} catch (error) {
67-
log(`Error closing Stagehand context: ${error}`, 'error');
233+
log(`Error closing Stagehand contexts: ${error}`, 'error');
68234
}
69235
}
70236
}

‎stagehand/src/index.ts‎

Lines changed: 1 addition & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -10,17 +10,6 @@ import { zodToJsonSchema } from "zod-to-json-schema";
1010
import { Context } from "./context.js";
1111
import { TOOLS } from "./tools/tools.js";
1212

13-
// Environment variables configuration
14-
const requiredEnvVars = {
15-
BROWSERBASE_API_KEY: process.env.BROWSERBASE_API_KEY,
16-
BROWSERBASE_PROJECT_ID: process.env.BROWSERBASE_PROJECT_ID,
17-
};
18-
19-
// Validate required environment variables
20-
Object.entries(requiredEnvVars).forEach(([name, value]) => {
21-
if (!value) throw new Error(`${name} environment variable is required`);
22-
});
23-
2413
export async function createServer(config: Config): Promise<Server> {
2514
// Create the server
2615
const server = new Server(
@@ -31,6 +20,7 @@ export async function createServer(config: Config): Promise<Server> {
3120
tools: { list: true, call: true },
3221
prompts: { list: true, get: true },
3322
notifications: { resources: { list_changed: true } },
23+
logging: {},
3424
},
3525
}
3626
);

‎stagehand/src/server.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,10 @@ export class ServerList {
1515
}
1616

1717
async close(server: Server) {
18+
await server.close();
1819
const index = this._servers.indexOf(server);
1920
if (index !== -1)
2021
this._servers.splice(index, 1);
21-
await server.close();
2222
}
2323

2424
async closeAll() {

0 commit comments

Comments
 (0)