moving types around, removing unnecessary code, reworking stagehand session create

Author: Kylejeong2

LICENSE

@@ -1,5 +1,5 @@
                                  Apache License
-                           Version 2.0, June 2025
+                           Version 2.0, January 2004
                         http://www.apache.org/licenses/
 
    TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION

eslint.config.js

@@ -15,15 +15,14 @@ export default defineConfig([
     languageOptions: { globals: { ...globals.browser, ...globals.node } },
     ignores: ["dist/**/*"],
   },
-  tseslint.configs.recommended,
+  ...tseslint.configs.recommended,
   {
-    files: ["**/*.{ts,mts,cts}"],
+    files: ["src/types/**/*.ts", "src/mcp/**/*.ts"],
     rules: {
       "@typescript-eslint/no-explicit-any": "off",
       "@typescript-eslint/no-unused-vars": "off",
       "@typescript-eslint/ban-ts-comment": "off",
     },
-    ignores: ["dist/**/*"],
   },
   {
     ignores: ["dist/**/*", "node_modules/**/*"],

src/context.ts

@@ -1,22 +1,10 @@
 import type { Stagehand } from "@browserbasehq/stagehand";
 import { Server } from "@modelcontextprotocol/sdk/server/index.js";
 import type { Config } from "../config.js";
-import {
-  CallToolResult,
-  TextContent,
-  ImageContent,
-} from "@modelcontextprotocol/sdk/types.js";
+import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { listResources, readResource } from "./mcp/resources.js";
-import {
-  getSession,
-  defaultSessionId,
-  type BrowserSession,
-} from "./sessionManager.js";
-
-export type ToolActionResult =
-  | { content?: (ImageContent | TextContent)[] }
-  | undefined
-  | void;
+import { getSession, defaultSessionId } from "./sessionManager.js";
+import type { MCPTool, BrowserSession } from "./types/types.js";
 
 export class Context {
   public readonly config: Config;
@@ -69,15 +57,15 @@ export class Context {
     return session.browser;
   }
 
-  async run(tool: any, args: any): Promise<CallToolResult> {
+  async run(tool: MCPTool, args: unknown): Promise<CallToolResult> {
     try {
       console.error(
         `Executing tool: ${tool.schema.name} with args: ${JSON.stringify(args)}`,
       );
 
       // Check if this tool has a handle method (new tool system)
       if ("handle" in tool && typeof tool.handle === "function") {
-        const toolResult = await tool.handle(this as any, args);
+        const toolResult = await tool.handle(this, args);
 
         if (toolResult?.action) {
           const actionResult = await toolResult.action();

src/index.ts

@@ -3,7 +3,7 @@ dotenv.config();
 
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { z } from "zod";
-import type { Tool } from "./tools/tool.js";
+import type { MCPToolsArray } from "./types/types.js";
 
 import { Context } from "./context.js";
 import type { Config } from "../config.js";
@@ -185,7 +185,7 @@ export default function ({ config }: { config: z.infer<typeof configSchema> }) {
     return prompt;
   });
 
-  const tools: Tool<any>[] = [...TOOLS];
+  const tools: MCPToolsArray = [...TOOLS];
 
   // Register each tool with the Smithery server
   tools.forEach((tool) => {

src/program.ts

@@ -20,7 +20,8 @@ try {
 } catch {
   // Fallback for CommonJS or when import.meta is not available
   __filename =
-    (globalThis as any).__filename || process.cwd() + "/dist/program.js";
+    (globalThis as { __filename: string }).__filename ||
+    process.cwd() + "/dist/program.js";
   __dirname = path.dirname(__filename);
 }
 

src/sessionManager.ts

@@ -1,14 +1,9 @@
-import { Browser, Page } from "playwright-core";
-import { AvailableModel, Stagehand } from "@browserbasehq/stagehand";
+import { Page } from "playwright-core";
+import { BrowserContext } from "@browserbasehq/stagehand";
 import type { Config } from "../config.js";
 import type { Cookie } from "playwright-core";
-
-export type BrowserSession = {
-  browser: Browser;
-  page: Page;
-  sessionId: string;
-  stagehand: Stagehand;
-};
+import { createStagehandInstance } from "./stagehandStore.js";
+import type { BrowserSession } from "./types/types.js";
 
 // Global state for managing browser sessions
 const browsers = new Map<string, BrowserSession>();
@@ -50,7 +45,7 @@ export function getActiveSessionId(): string {
  * @param cookies Array of cookies to add
  */
 export async function addCookiesToContext(
-  context: any,
+  context: BrowserContext,
   cookies: Cookie[],
 ): Promise<void> {
   if (!cookies || cookies.length === 0) {
@@ -92,40 +87,14 @@ export async function createNewBrowserSession(
       `[SessionManager] ${resumeSessionId ? "Resuming" : "Creating"} Stagehand session ${newSessionId}...\n`,
     );
 
-    // Create and initialize Stagehand instance
-    const stagehand = new Stagehand({
-      env: "BROWSERBASE",
-      apiKey: config.browserbaseApiKey,
-      projectId: config.browserbaseProjectId,
-      modelName: (config.modelName ||
-        "google/gemini-2.0-flash") as AvailableModel,
-      modelClientOptions: {
-        apiKey: process.env.GEMINI_API_KEY, // TODO: change this in prod to just use our key
-      },
-      ...(resumeSessionId && { browserbaseSessionID: resumeSessionId }),
-      browserbaseSessionCreateParams: {
-        projectId: config.browserbaseProjectId!,
-        proxies: config.proxies,
-        browserSettings: {
-          viewport: {
-            width: config.viewPort?.browserWidth ?? 1024,
-            height: config.viewPort?.browserHeight ?? 768,
-          },
-          context: config.context?.contextId
-            ? {
-                id: config.context?.contextId,
-                persist: config.context?.persist ?? true,
-              }
-            : undefined,
-          advancedStealth: config.advancedStealth ?? undefined,
-        },
+    // Create and initialize Stagehand instance using shared function
+    const stagehand = await createStagehandInstance(
+      config,
+      {
+        ...(resumeSessionId && { browserbaseSessionID: resumeSessionId }),
       },
-      logger: (logLine) => {
-        console.error(`Stagehand[${newSessionId}]: ${logLine.message}`);
-      },
-    });
-
-    await stagehand.init();
+      newSessionId,
+    );
 
     // Get the page and browser from Stagehand
     const page = stagehand.page as unknown as Page;
@@ -171,7 +140,10 @@ export async function createNewBrowserSession(
       Array.isArray(config.cookies) &&
       config.cookies.length > 0
     ) {
-      await addCookiesToContext(page.context(), config.cookies);
+      await addCookiesToContext(
+        page.context() as BrowserContext,
+        config.cookies,
+      );
     }
 
     const sessionObj: BrowserSession = {
@@ -306,7 +278,7 @@ export async function getSession(
   if (sessionId === defaultSessionId && createIfMissing) {
     try {
       return await ensureDefaultSessionInternal(config);
-    } catch (error) {
+    } catch {
       process.stderr.write(
         `[SessionManager] Failed to get default session due to error in ensureDefaultSessionInternal for ${sessionId}. See previous messages for details.\n`,
       );
@@ -385,7 +357,7 @@ export async function closeAllSessions(): Promise<void> {
   }
   try {
     await Promise.all(closePromises);
-  } catch (_e) {
+  } catch {
     // Individual errors are caught and logged by closeBrowserGracefully
     process.stderr.write(
       `[SessionManager] WARN - Some errors occurred during batch session closing. See individual messages.\n`,

src/stagehandStore.ts

@@ -8,24 +8,20 @@ import type { Config } from "../config.js";
 const store = new Map<string, StagehandSession>();
 
 /**
- * Create a new Stagehand session
+ * Create a configured Stagehand instance
  */
-export const create = async (
+export const createStagehandInstance = async (
   config: Config,
   params: CreateSessionParams = {},
-): Promise<StagehandSession> => {
-  const id = randomUUID() + "_" + Date.now();
-
-  // Merge config with params
+  sessionId: string,
+): Promise<Stagehand> => {
   const apiKey = params.apiKey || config.browserbaseApiKey;
   const projectId = params.projectId || config.browserbaseProjectId;
 
   if (!apiKey || !projectId) {
     throw new Error("Browserbase API Key and Project ID are required");
   }
 
-  process.stderr.write(`[StagehandStore] Creating new session ${id}...\n`);
-
   const stagehand = new Stagehand({
     env: "BROWSERBASE",
     apiKey,
@@ -34,12 +30,12 @@ export const create = async (
       config.modelName ||
       "google/gemini-2.0-flash") as AvailableModel,
     modelClientOptions: {
-      apiKey: process.env.GEMINI_API_KEY, //TODO:
+      apiKey: config.modelApiKey || process.env.GEMINI_API_KEY,
     },
     ...(params.browserbaseSessionID && {
       browserbaseSessionID: params.browserbaseSessionID,
     }),
-    browserbaseSessionCreateParams: params.browserbaseSessionCreateParams || {
+    browserbaseSessionCreateParams: {
       projectId,
       proxies: config.proxies,
       browserSettings: {
@@ -57,11 +53,27 @@ export const create = async (
       },
     },
     logger: (logLine) => {
-      console.error(`Stagehand[${id}]: ${logLine.message}`);
+      console.error(`Stagehand[${sessionId}]: ${logLine.message}`);
     },
   });
 
   await stagehand.init();
+  return stagehand;
+};
+
+/**
+ * Create a new Stagehand session
+ */
+export const create = async (
+  config: Config,
+  params: CreateSessionParams = {},
+): Promise<StagehandSession> => {
+  // Global ID, must be 100% Unique
+  const id = randomUUID() + "_" + config.browserbaseProjectId;
+
+  process.stderr.write(`[StagehandStore] Creating new session ${id}...\n`);
+
+  const stagehand = await createStagehandInstance(config, params, id);
 
   const page = stagehand.page as unknown as Page;
   const browser = page.context().browser();

src/tools/act.ts

@@ -1,7 +1,7 @@
 import { z } from "zod";
 import type { Tool, ToolSchema, ToolResult } from "./tool.js";
 import type { Context } from "../context.js";
-import type { ToolActionResult } from "../context.js";
+import type { ToolActionResult } from "../types/types.js";
 
 const ActInputSchema = z.object({
   action: z

src/tools/extract.ts

@@ -1,7 +1,7 @@
 import { z } from "zod";
 import type { Tool, ToolSchema, ToolResult } from "./tool.js";
 import type { Context } from "../context.js";
-import type { ToolActionResult } from "../context.js";
+import type { ToolActionResult } from "../types/types.js";
 
 const ExtractInputSchema = z.object({
   instruction: z

src/tools/multiSession.ts

@@ -51,7 +51,10 @@ function createMultiSessionAwareTool<TInput extends InputType>(
       description: `${originalTool.schema.description} (for a specific session)`,
       inputSchema: newInputSchema,
     },
-    handle: async (context: Context, params: any): Promise<ToolResult> => {
+    handle: async (
+      context: Context,
+      params: z.infer<typeof newInputSchema>,
+    ): Promise<ToolResult> => {
       const { sessionId, ...originalParams } = params;
 
       // Get the session
@@ -94,22 +97,15 @@ export const createSessionTool = defineTool({
         .describe(
           "Resume an existing Browserbase session by providing its session ID. Use this to continue work in a previously created browser session that may have been paused or disconnected.",
         ),
-      browserbaseSessionCreateParams: z
-        .any()
-        .optional()
-        .describe(
-          "Advanced Browserbase session configuration parameters for customizing browser settings, viewport, user agent, etc. Leave empty for default settings.",
-        ),
     }),
   },
   handle: async (
     context: Context,
-    { name, browserbaseSessionID, browserbaseSessionCreateParams },
+    { name, browserbaseSessionID },
   ): Promise<ToolResult> => {
     try {
       const params: CreateSessionParams = {
         browserbaseSessionID,
-        browserbaseSessionCreateParams,
         meta: name ? { name } : undefined,
       };