diff --git a/src/context.ts b/src/context.ts index 3b2abcc..a379c5f 100644 --- a/src/context.ts +++ b/src/context.ts @@ -3,13 +3,13 @@ import { Server } from "@modelcontextprotocol/sdk/server/index.js"; import type { Config } from "../config.d.ts"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { listResources, readResource } from "./mcp/resources.js"; -import { getSession, defaultSessionId } from "./sessionManager.js"; -import type { MCPTool, BrowserSession } from "./types/types.js"; +import { store } from "./tools/helpers/session.js"; +import type { MCPTool, StagehandSession } from "./types/types.js"; export class Context { public readonly config: Config; private server: Server; - public currentSessionId: string = defaultSessionId; + private _currentSession: StagehandSession | null = null; constructor(server: Server, config: Config) { this.server = server; @@ -21,40 +21,35 @@ export class Context { } /** - * Gets the Stagehand instance for the current session from SessionManager + * Gets the Stagehand instance for the current session */ - public async getStagehand( - sessionId: string = this.currentSessionId, - ): Promise { - const session = await getSession(sessionId, this.config); - if (!session) { - throw new Error(`No session found for ID: ${sessionId}`); - } + public async getStagehand(): Promise { + const session = await this.getCurrentSession(); return session.stagehand; } - public async getActivePage(): Promise { - // Get page from session manager - const session = await getSession(this.currentSessionId, this.config); - if (session && session.page && !session.page.isClosed()) { - return session.page; + public async getActivePage(): Promise { + try { + const session = await this.getCurrentSession(); + if (session.page && !session.page.isClosed()) { + return session.page; + } + } catch { + // Session not available } - return null; } - public async getActiveBrowser( - createIfMissing: boolean = true, - ): Promise { - const session = await getSession( - this.currentSessionId, - this.config, - createIfMissing, - ); - if (!session || !session.browser || !session.browser.isConnected()) { - return null; + public async getActiveBrowser(): Promise { + try { + const session = await this.getCurrentSession(); + if (session.browser && session.browser.isConnected()) { + return session.browser; + } + } catch { + // Session not available } - return session.browser; + return null; } async run(tool: MCPTool, args: unknown): Promise { @@ -122,4 +117,15 @@ export class Context { readResource(uri: string) { return readResource(uri); } + + /** + * Gets or creates the current session + */ + private async getCurrentSession(): Promise { + if (!this._currentSession) { + const sessionStore = store(this.config); + this._currentSession = await sessionStore.create(); + } + return this._currentSession; + } } diff --git a/src/program.ts b/src/program.ts index b485493..3f43c2e 100644 --- a/src/program.ts +++ b/src/program.ts @@ -6,7 +6,7 @@ import { fileURLToPath } from "url"; import createServerFunction from "./index.js"; import { ServerList } from "./server.js"; import { startHttpTransport, startStdioTransport } from "./transport.js"; -import * as stagehandStore from "./stagehandStore.js"; +// Session cleanup will be handled by the Context cleanup import { resolveConfig } from "./config.js"; @@ -83,7 +83,7 @@ function setupExitWatchdog(serverList: ServerList) { const handleExit = async () => { setTimeout(() => process.exit(0), 15000); try { - await Promise.all([stagehandStore.removeAll(), serverList.closeAll()]); + await serverList.closeAll(); } catch (error) { console.error("Error during cleanup:", error); } diff --git a/src/sessionManager.ts b/src/sessionManager.ts deleted file mode 100644 index 10df670..0000000 --- a/src/sessionManager.ts +++ /dev/null @@ -1,376 +0,0 @@ -import { Page, BrowserContext } from "@browserbasehq/stagehand"; -import type { Config } from "../config.d.ts"; -import type { Cookie } from "playwright-core"; -import { createStagehandInstance } from "./stagehandStore.js"; -import type { BrowserSession } from "./types/types.js"; - -// Global state for managing browser sessions -const browsers = new Map(); - -// Keep track of the default session explicitly -let defaultBrowserSession: BrowserSession | null = null; - -// Define a specific ID for the default session -export const defaultSessionId = `browserbase_session_main_${Date.now()}`; - -// Keep track of the active session ID. Defaults to the main session. -let activeSessionId: string = defaultSessionId; - -/** - * Sets the active session ID. - * @param id The ID of the session to set as active. - */ -export function setActiveSessionId(id: string): void { - if (browsers.has(id) || id === defaultSessionId) { - activeSessionId = id; - } else { - process.stderr.write( - `[SessionManager] WARN - Set active session failed for non-existent ID: ${id}\n`, - ); - } -} - -/** - * Gets the active session ID. - * @returns The active session ID. - */ -export function getActiveSessionId(): string { - return activeSessionId; -} - -/** - * Adds cookies to a browser context - * @param context Playwright browser context - * @param cookies Array of cookies to add - */ -export async function addCookiesToContext( - context: BrowserContext, - cookies: Cookie[], -): Promise { - if (!cookies || cookies.length === 0) { - return; - } - - try { - process.stderr.write( - `[SessionManager] Adding ${cookies.length} cookies to browser context\n`, - ); - await context.addCookies(cookies); - process.stderr.write( - `[SessionManager] Successfully added cookies to browser context\n`, - ); - } catch (error) { - process.stderr.write( - `[SessionManager] Error adding cookies to browser context: ${ - error instanceof Error ? error.message : String(error) - }\n`, - ); - } -} - -// Function to create a new Browserbase session using Stagehand -export async function createNewBrowserSession( - newSessionId: string, - config: Config, - resumeSessionId?: string, -): Promise { - if (!config.browserbaseApiKey) { - throw new Error("Browserbase API Key is missing in the configuration."); - } - if (!config.browserbaseProjectId) { - throw new Error("Browserbase Project ID is missing in the configuration."); - } - - try { - process.stderr.write( - `[SessionManager] ${resumeSessionId ? "Resuming" : "Creating"} Stagehand session ${newSessionId}...\n`, - ); - - // Create and initialize Stagehand instance using shared function - const stagehand = await createStagehandInstance( - config, - { - ...(resumeSessionId && { browserbaseSessionID: resumeSessionId }), - }, - newSessionId, - ); - - // Get the page and browser from Stagehand - const page = stagehand.page as unknown as Page; - const browser = page.context().browser(); - - if (!browser) { - throw new Error("Failed to get browser from Stagehand page context"); - } - - const browserbaseSessionId = stagehand.browserbaseSessionID; - - process.stderr.write( - `[SessionManager] Stagehand initialized with Browserbase session: ${browserbaseSessionId}\n`, - ); - process.stderr.write( - `[SessionManager] Browserbase Live Debugger URL: https://www.browserbase.com/sessions/${browserbaseSessionId}\n`, - ); - - // Set up disconnect handler - browser.on("disconnected", () => { - process.stderr.write(`[SessionManager] Disconnected: ${newSessionId}\n`); - browsers.delete(newSessionId); - if (defaultBrowserSession && defaultBrowserSession.browser === browser) { - process.stderr.write( - `[SessionManager] Disconnected (default): ${newSessionId}\n`, - ); - defaultBrowserSession = null; - } - if ( - activeSessionId === newSessionId && - newSessionId !== defaultSessionId - ) { - process.stderr.write( - `[SessionManager] WARN - Active session disconnected, resetting to default: ${newSessionId}\n`, - ); - setActiveSessionId(defaultSessionId); - } - }); - - // Add cookies to the context if they are provided in the config - if ( - config.cookies && - Array.isArray(config.cookies) && - config.cookies.length > 0 - ) { - await addCookiesToContext( - page.context() as BrowserContext, - config.cookies, - ); - } - - const sessionObj: BrowserSession = { - browser, - page, - sessionId: browserbaseSessionId!, - stagehand, - }; - - browsers.set(newSessionId, sessionObj); - - if (newSessionId === defaultSessionId) { - defaultBrowserSession = sessionObj; - } - - setActiveSessionId(newSessionId); - process.stderr.write( - `[SessionManager] Session created and active: ${newSessionId}\n`, - ); - - return sessionObj; - } catch (creationError) { - const errorMessage = - creationError instanceof Error - ? creationError.message - : String(creationError); - process.stderr.write( - `[SessionManager] Creating session ${newSessionId} failed: ${errorMessage}\n`, - ); - throw new Error( - `Failed to create/connect session ${newSessionId}: ${errorMessage}`, - ); - } -} - -async function closeBrowserGracefully( - session: BrowserSession | undefined | null, - sessionIdToLog: string, -): Promise { - // Close Stagehand instance which handles browser cleanup - if (session?.stagehand) { - try { - process.stderr.write( - `[SessionManager] Closing Stagehand for session: ${sessionIdToLog}\n`, - ); - await session.stagehand.close(); - process.stderr.write( - `[SessionManager] Successfully closed Stagehand and browser for session: ${sessionIdToLog}\n`, - ); - } catch (closeError) { - process.stderr.write( - `[SessionManager] WARN - Error closing Stagehand for session ${sessionIdToLog}: ${ - closeError instanceof Error ? closeError.message : String(closeError) - }\n`, - ); - } - } -} - -// Internal function to ensure default session -export async function ensureDefaultSessionInternal( - config: Config, -): Promise { - const sessionId = defaultSessionId; - let needsReCreation = false; - - if (!defaultBrowserSession) { - needsReCreation = true; - process.stderr.write( - `[SessionManager] Default session ${sessionId} not found, creating.\n`, - ); - } else if ( - !defaultBrowserSession.browser.isConnected() || - defaultBrowserSession.page.isClosed() - ) { - needsReCreation = true; - process.stderr.write( - `[SessionManager] Default session ${sessionId} is stale, recreating.\n`, - ); - await closeBrowserGracefully(defaultBrowserSession, sessionId); - defaultBrowserSession = null; - browsers.delete(sessionId); - } - - if (needsReCreation) { - try { - defaultBrowserSession = await createNewBrowserSession(sessionId, config); - return defaultBrowserSession; - } catch (creationError) { - // Error during initial creation or recreation - process.stderr.write( - `[SessionManager] Initial/Recreation attempt for default session ${sessionId} failed. Error: ${ - creationError instanceof Error - ? creationError.message - : String(creationError) - }\n`, - ); - // Attempt one more time after a failure - process.stderr.write( - `[SessionManager] Retrying creation of default session ${sessionId} after error...\n`, - ); - try { - defaultBrowserSession = await createNewBrowserSession( - sessionId, - config, - ); - return defaultBrowserSession; - } catch (retryError) { - const finalErrorMessage = - retryError instanceof Error ? retryError.message : String(retryError); - process.stderr.write( - `[SessionManager] Failed to recreate default session ${sessionId} after retry: ${finalErrorMessage}\n`, - ); - throw new Error( - `Failed to ensure default session ${sessionId} after initial error and retry: ${finalErrorMessage}`, - ); - } - } - } - - // If we reached here, the existing default session is considered okay. - setActiveSessionId(sessionId); // Ensure default is marked active - return defaultBrowserSession!; // Non-null assertion: logic ensures it's not null here -} - -// Get a specific session by ID -export async function getSession( - sessionId: string, - config: Config, - createIfMissing: boolean = true, -): Promise { - if (sessionId === defaultSessionId && createIfMissing) { - try { - return await ensureDefaultSessionInternal(config); - } catch { - process.stderr.write( - `[SessionManager] Failed to get default session due to error in ensureDefaultSessionInternal for ${sessionId}. See previous messages for details.\n`, - ); - return null; // Or rethrow if getSession failing for default is critical - } - } - - // For non-default sessions - process.stderr.write(`[SessionManager] Getting session: ${sessionId}\n`); - const sessionObj = browsers.get(sessionId); - - if (!sessionObj) { - process.stderr.write( - `[SessionManager] WARN - Session not found in map: ${sessionId}\n`, - ); - return null; - } - - // Validate the found session - if (!sessionObj.browser.isConnected() || sessionObj.page.isClosed()) { - process.stderr.write( - `[SessionManager] WARN - Found session ${sessionId} is stale, removing.\n`, - ); - await closeBrowserGracefully(sessionObj, sessionId); - browsers.delete(sessionId); - if (activeSessionId === sessionId) { - process.stderr.write( - `[SessionManager] WARN - Invalidated active session ${sessionId}, resetting to default.\n`, - ); - setActiveSessionId(defaultSessionId); - } - return null; - } - - // Session appears valid, make it active - setActiveSessionId(sessionId); - process.stderr.write(`[SessionManager] Using valid session: ${sessionId}\n`); - return sessionObj; -} - -/** - * Clean up a session by removing it from tracking. - * This is called after a browser is closed to ensure proper cleanup. - * @param sessionId The session ID to clean up - */ -export async function cleanupSession(sessionId: string): Promise { - process.stderr.write(`[SessionManager] Cleaning up session: ${sessionId}\n`); - - // Get the session to close it gracefully - const session = browsers.get(sessionId); - if (session) { - await closeBrowserGracefully(session, sessionId); - } - - // Remove from browsers map - browsers.delete(sessionId); - - // Clear default session reference if this was the default - if (sessionId === defaultSessionId && defaultBrowserSession) { - defaultBrowserSession = null; - } - - // Reset active session to default if this was the active one - if (activeSessionId === sessionId) { - process.stderr.write( - `[SessionManager] Cleaned up active session ${sessionId}, resetting to default.\n`, - ); - setActiveSessionId(defaultSessionId); - } -} - -// Function to close all managed browser sessions gracefully -export async function closeAllSessions(): Promise { - process.stderr.write(`[SessionManager] Closing all sessions...\n`); - const closePromises: Promise[] = []; - for (const [id, session] of browsers.entries()) { - process.stderr.write(`[SessionManager] Closing session: ${id}\n`); - closePromises.push( - // Use the helper for consistent logging/error handling - closeBrowserGracefully(session, id), - ); - } - try { - await Promise.all(closePromises); - } 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`, - ); - } - - browsers.clear(); - defaultBrowserSession = null; - setActiveSessionId(defaultSessionId); // Reset active session to default - process.stderr.write(`[SessionManager] All sessions closed and cleared.\n`); -} diff --git a/src/sessions/SessionStore.ts b/src/sessions/SessionStore.ts new file mode 100644 index 0000000..925fd71 --- /dev/null +++ b/src/sessions/SessionStore.ts @@ -0,0 +1,130 @@ +import { randomUUID } from "crypto"; +import type { Config } from "../../config.d.ts"; +import { StagehandSession, CreateSessionParams } from "../types/types.js"; +import { StagehandAdapter } from "./StagehandAdapter.js"; + +/** + * SessionStore manages the lifecycle of Stagehand sessions + * Provides CRUD operations for session management + */ +export class SessionStore { + private store = new Map(); + private config: Config; + private logger: (message: string) => void; + + constructor( + config: Config, + logger: (message: string) => void = console.error, + ) { + this.config = config; + this.logger = logger; + } + + /** + * Create a new session + */ + async create(opts: CreateSessionParams = {}): Promise { + // Generate unique session ID + const id = randomUUID() + "_" + this.config.browserbaseProjectId; + + this.logger(`[SessionStore] Creating new session ${id}...`); + + try { + // Launch Stagehand instance using adapter + const session = await StagehandAdapter.launch(this.config, opts, id); + + // Store the session + this.store.set(id, session); + + this.logger( + `[SessionStore] Session created: ${id} (BB: ${session.metadata?.bbSessionId})`, + ); + this.logger( + `[SessionStore] Live debugger: https://www.browserbase.com/sessions/${session.metadata?.bbSessionId}`, + ); + + // Set up disconnect handler + const disconnectHandler = () => { + this.logger(`[SessionStore] Session disconnected: ${id}`); + this.store.delete(id); + }; + + session.browser.on("disconnected", disconnectHandler); + + // Store the handler for cleanup + session.metadata = { + ...session.metadata, + disconnectHandler, + }; + + return session; + } catch (error) { + this.logger( + `[SessionStore] Failed to create session ${id}: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + throw error; + } + } + + /** + * Get a session by ID + */ + get(id: string): StagehandSession | null { + return this.store.get(id) ?? null; + } + + /** + * List all active sessions + */ + list(): StagehandSession[] { + return Array.from(this.store.values()); + } + + /** + * Remove and close a session + */ + async remove(id: string): Promise { + const session = this.store.get(id); + if (!session) { + this.logger(`[SessionStore] Session not found for removal: ${id}`); + return; + } + + this.logger(`[SessionStore] Removing session: ${id}`); + + try { + if (session.metadata?.disconnectHandler) { + session.browser.off("disconnected", session.metadata.disconnectHandler); + } + + await session.stagehand.close(); + this.logger(`[SessionStore] Session closed: ${id}`); + } catch (error) { + this.logger( + `[SessionStore] Error closing session ${id}: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + } finally { + this.store.delete(id); + } + } + + /** + * Remove all sessions + */ + async removeAll(): Promise { + this.logger(`[SessionStore] Removing all ${this.store.size} sessions...`); + await Promise.all(this.list().map((s) => this.remove(s.id))); + this.logger(`[SessionStore] All sessions removed`); + } + + /** + * Get store size + */ + size(): number { + return this.store.size; + } +} diff --git a/src/sessions/StagehandAdapter.ts b/src/sessions/StagehandAdapter.ts new file mode 100644 index 0000000..ef46adf --- /dev/null +++ b/src/sessions/StagehandAdapter.ts @@ -0,0 +1,80 @@ +import { Stagehand, Page } from "@browserbasehq/stagehand"; +import type { Config } from "../../config.d.ts"; +import { CreateSessionParams, StagehandSession } from "../types/types.js"; + +/** + * Core adapter for launching Stagehand instances + * Handles the initialization and configuration of Stagehand with Browserbase + */ +export class StagehandAdapter { + /** + * Launch a new Stagehand instance with the given configuration and parameters + */ + static async launch( + config: Config, + params: CreateSessionParams = {}, + sessionId: string, + ): Promise { + 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"); + } + + const stagehand = new Stagehand({ + env: "BROWSERBASE", + apiKey, + projectId, + modelName: + params.modelName || config.modelName || "google/gemini-2.0-flash", + modelClientOptions: { + apiKey: config.modelApiKey || process.env.GEMINI_API_KEY, + }, + ...(params.browserbaseSessionID && { + browserbaseSessionID: params.browserbaseSessionID, + }), + browserbaseSessionCreateParams: { + projectId, + 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, + }, + }, + logger: (logLine) => { + console.error(`Stagehand[${sessionId}]: ${logLine.message}`); + }, + }); + + await stagehand.init(); + + const page = stagehand.page as unknown as Page; + const browser = page.context().browser(); + + if (!browser) { + throw new Error("Failed to get browser from Stagehand page context"); + } + + return { + id: sessionId, + stagehand, + page, + browser, + created: Date.now(), + metadata: { + ...params.meta, + bbSessionId: stagehand.browserbaseSessionID, + }, + }; + } +} diff --git a/src/stagehandStore.ts b/src/stagehandStore.ts deleted file mode 100644 index c2e957f..0000000 --- a/src/stagehandStore.ts +++ /dev/null @@ -1,183 +0,0 @@ -import { randomUUID } from "crypto"; -import { Stagehand, Page } from "@browserbasehq/stagehand"; -import { StagehandSession, CreateSessionParams } from "./types/types.js"; -import type { Config } from "../config.d.ts"; - -// Store for all active sessions -const store = new Map(); - -/** - * Create a configured Stagehand instance - */ -export const createStagehandInstance = async ( - config: Config, - params: CreateSessionParams = {}, - sessionId: string, -): Promise => { - 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"); - } - - const stagehand = new Stagehand({ - env: "BROWSERBASE", - apiKey, - projectId, - modelName: - params.modelName || config.modelName || "google/gemini-2.0-flash", - modelClientOptions: { - apiKey: config.modelApiKey || process.env.GEMINI_API_KEY, - }, - ...(params.browserbaseSessionID && { - browserbaseSessionID: params.browserbaseSessionID, - }), - browserbaseSessionCreateParams: { - projectId, - 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, - }, - }, - logger: (logLine) => { - 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 => { - // 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(); - - if (!browser) { - throw new Error("Failed to get browser from Stagehand page context"); - } - - const session: StagehandSession = { - id, - stagehand, - page, - browser, - created: Date.now(), - metadata: { - ...params.meta, - bbSessionId: stagehand.browserbaseSessionID, - }, - }; - - store.set(id, session); - - process.stderr.write( - `[StagehandStore] Session created: ${id} (BB: ${stagehand.browserbaseSessionID})\n`, - ); - process.stderr.write( - `[StagehandStore] Live debugger: https://www.browserbase.com/sessions/${stagehand.browserbaseSessionID}\n`, - ); - - // Set up disconnect handler - const disconnectHandler = () => { - process.stderr.write(`[StagehandStore] Session disconnected: ${id}\n`); - store.delete(id); - }; - - browser.on("disconnected", disconnectHandler); - - // Store the handler for cleanup - session.metadata = { - ...session.metadata, - disconnectHandler, - }; - - return session; -}; - -/** - * Get a session by ID - */ -export const get = (id: string): StagehandSession | null => { - return store.get(id) ?? null; -}; - -/** - * List all active sessions - */ -export const list = (): StagehandSession[] => { - return Array.from(store.values()); -}; - -/** - * Remove and close a session - */ -export const remove = async (id: string): Promise => { - const session = store.get(id); - if (!session) { - process.stderr.write( - `[StagehandStore] Session not found for removal: ${id}\n`, - ); - return; - } - - process.stderr.write(`[StagehandStore] Removing session: ${id}\n`); - - try { - if (session.metadata?.disconnectHandler) { - session.browser.off("disconnected", session.metadata.disconnectHandler); - } - - await session.stagehand.close(); - process.stderr.write(`[StagehandStore] Session closed: ${id}\n`); - } catch (error) { - process.stderr.write( - `[StagehandStore] Error closing session ${id}: ${ - error instanceof Error ? error.message : String(error) - }\n`, - ); - } finally { - store.delete(id); - } -}; - -/** - * Remove all sessions - */ -export const removeAll = async (): Promise => { - process.stderr.write( - `[StagehandStore] Removing all ${store.size} sessions...\n`, - ); - await Promise.all(list().map((s) => remove(s.id))); - process.stderr.write(`[StagehandStore] All sessions removed\n`); -}; - -/** - * Get store size - */ -export const size = (): number => { - return store.size; -}; diff --git a/src/tools/helpers/session.ts b/src/tools/helpers/session.ts new file mode 100644 index 0000000..899e4bc --- /dev/null +++ b/src/tools/helpers/session.ts @@ -0,0 +1,56 @@ +import { SessionStore } from "../../sessions/SessionStore.js"; +import { defaultLogger } from "../../utils/logger.js"; +import { StagehandSession, CreateSessionParams } from "../../types/types.js"; +import type { Config } from "../../../config.d.ts"; +import type { Context } from "../../context.js"; + +/** + * Singleton session store instance + * This provides a global access point to session management + */ +let store: SessionStore | null = null; + +/** + * Initialize or get the singleton store instance + */ +function getStore(config: Config): SessionStore { + if (!store) { + store = new SessionStore(config, defaultLogger); + } + return store; +} + +/** + * Helper function to work with sessions + * If no sessionId is provided, creates a new session + * If sessionId is provided, retrieves existing session or creates new one + */ +export async function withSession( + context: Context, + sessionId?: string, + opts: CreateSessionParams = {}, +): Promise { + const sessionStore = getStore(context.config); + + if (sessionId) { + // Try to get existing session + const existingSession = sessionStore.get(sessionId); + if (existingSession) { + return existingSession; + } + + // If session doesn't exist, create it with the specified ID + // Note: We'll need to modify SessionStore.create to accept a custom ID + throw new Error( + `Session ${sessionId} not found and cannot create with custom ID`, + ); + } + + // Create new session + return await sessionStore.create(opts); +} + +/** + * Export the singleton store getter for direct access + */ +export { getStore as store }; diff --git a/src/tools/multiSession.ts b/src/tools/multiSession.ts index bce5b94..dbc9469 100644 --- a/src/tools/multiSession.ts +++ b/src/tools/multiSession.ts @@ -6,7 +6,7 @@ import { type ToolResult, type InputType, } from "./tool.js"; -import * as stagehandStore from "../stagehandStore.js"; +import { store } from "./helpers/session.js"; import { CreateSessionParams } from "../types/types.js"; import type { Context } from "../context.js"; import navigateTool from "./navigate.js"; @@ -58,8 +58,9 @@ function createMultiSessionAwareTool( ): Promise => { const { sessionId, ...originalParams } = params; - // Get the session - const session = stagehandStore.get(sessionId); + // Get the session using new session management + const sessionStore = store(context.config); + const session = sessionStore.get(sessionId); if (!session) { throw new Error(`Session ${sessionId} not found`); } @@ -110,7 +111,8 @@ export const createSessionTool = defineTool({ meta: name ? { name } : undefined, }; - const session = await stagehandStore.create(context.config, params); + const sessionStore = store(context.config); + const session = await sessionStore.create(params); const bbSessionId = session.metadata?.bbSessionId; if (!bbSessionId) { @@ -154,8 +156,9 @@ export const listSessionsTool = defineTool({ "ONLY WORKS WITH MULTI-SESSION TOOLS! Track all parallel sessions: Critical tool for multi-session management! Shows all active browser sessions with their IDs, names, ages, and Browserbase session IDs. Use this frequently to monitor your parallel automation workflows, verify sessions are running, and get session IDs for session-specific tools. Essential for debugging and resource management in complex multi-browser scenarios.", inputSchema: z.object({}), }, - handle: async (): Promise => { - const sessions = stagehandStore.list(); + handle: async (context: Context): Promise => { + const sessionStore = store(context.config); + const sessions = sessionStore.list(); if (sessions.length === 0) { return { @@ -213,13 +216,14 @@ export const closeSessionTool = defineTool({ ), }), }, - handle: async (_context: Context, { sessionId }): Promise => { - const session = stagehandStore.get(sessionId); + handle: async (context: Context, { sessionId }): Promise => { + const sessionStore = store(context.config); + const session = sessionStore.get(sessionId); if (!session) { throw new Error(`Session ${sessionId} not found`); } - await stagehandStore.remove(sessionId); + await sessionStore.remove(sessionId); return { action: async () => ({ diff --git a/src/tools/session.ts b/src/tools/session.ts index c944d70..cec07c0 100644 --- a/src/tools/session.ts +++ b/src/tools/session.ts @@ -3,16 +3,7 @@ import type { Tool, ToolSchema, ToolResult } from "./tool.js"; import type { Context } from "../context.js"; import type { ToolActionResult } from "../types/types.js"; import { Browserbase } from "@browserbasehq/sdk"; - -// Import SessionManager functions -import { - createNewBrowserSession, - defaultSessionId, - ensureDefaultSessionInternal, - cleanupSession, - getSession, -} from "../sessionManager.js"; -import type { BrowserSession } from "../types/types.js"; +import { store } from "./helpers/session.js"; // --- Tool: Create Session --- const CreateSessionInputSchema = z.object({ @@ -33,88 +24,47 @@ const createSessionSchema: ToolSchema = { inputSchema: CreateSessionInputSchema, }; -// Handle function for CreateSession using SessionManager +// Handle function for CreateSession using new session management async function handleCreateSession( context: Context, params: CreateSessionInput, ): Promise { const action = async (): Promise => { try { - const config = context.config; // Get config from context - let targetSessionId: string; - - if (params.sessionId) { - const projectId = config.browserbaseProjectId || ""; - targetSessionId = `${params.sessionId}_${projectId}`; - process.stderr.write( - `[tool.createSession] Attempting to create/assign session with specified ID: ${targetSessionId}`, - ); - } else { - targetSessionId = defaultSessionId; - } - - let session: BrowserSession; - if (targetSessionId === defaultSessionId) { - session = await ensureDefaultSessionInternal(config); - } else { - // When user provides a sessionId, we want to resume that Browserbase session - session = await createNewBrowserSession( - targetSessionId, - config, - params.sessionId, - ); - } + const sessionStore = store(context.config); - if ( - !session || - !session.browser || - !session.page || - !session.sessionId || - !session.stagehand - ) { - throw new Error( - `SessionManager failed to return a valid session object with actualSessionId for ID: ${targetSessionId}`, - ); - } + // Create session with optional resumption + const session = await sessionStore.create({ + browserbaseSessionID: params.sessionId, + }); - context.currentSessionId = targetSessionId; + // Get debug URL const bb = new Browserbase({ - apiKey: config.browserbaseApiKey, + apiKey: context.config.browserbaseApiKey, }); - const debugUrl = (await bb.sessions.debug(session.sessionId)) - .debuggerFullscreenUrl; - process.stderr.write( - `[tool.connected] Successfully connected to Browserbase session. Internal ID: ${targetSessionId}, Actual ID: ${session.sessionId}`, - ); - - process.stderr.write( - `[SessionManager] Browserbase Live Session View URL: https://www.browserbase.com/sessions/${session.sessionId}`, - ); + const bbSessionId = session.metadata?.bbSessionId; + if (!bbSessionId) { + throw new Error("No Browserbase session ID available"); + } - process.stderr.write( - `[SessionManager] Browserbase Live Debugger URL: ${debugUrl}`, - ); + const debugUrl = (await bb.sessions.debug(bbSessionId)) + .debuggerFullscreenUrl; return { content: [ { type: "text", - text: `Browserbase Live Session View URL: https://www.browserbase.com/sessions/${session.sessionId}\nBrowserbase Live Debugger URL: ${debugUrl}`, + text: `Browserbase Live Session View URL: https://www.browserbase.com/sessions/${bbSessionId}\nBrowserbase Live Debugger URL: ${debugUrl}`, }, ], }; } catch (error: unknown) { const errorMessage = error instanceof Error ? error.message : String(error); - process.stderr.write( - `[tool.createSession] Action failed: ${errorMessage}`, - ); - // Re-throw to be caught by Context.run's error handling for actions throw new Error(`Failed to create Browserbase session: ${errorMessage}`); } }; - // Return the ToolResult structure expected by Context.run return { action: action, waitForNetwork: false, @@ -140,86 +90,25 @@ const closeSessionSchema: ToolSchema = { async function handleCloseSession(context: Context): Promise { const action = async (): Promise => { - // Store the current session ID before it's potentially changed. - const previousSessionId = context.currentSessionId; - let stagehandClosedSuccessfully = false; - let stagehandCloseErrorMessage = ""; - - // Step 1: Attempt to get the session and close Stagehand - let browserbaseSessionId: string | undefined; try { - const session = await getSession( - previousSessionId, - context.config, - false, - ); - - if (session && session.stagehand) { - // Store the actual Browserbase session ID for the replay URL - browserbaseSessionId = session.sessionId; - - process.stderr.write( - `[tool.closeSession] Attempting to close Stagehand for session: ${previousSessionId || "default"} (Browserbase ID: ${browserbaseSessionId})`, - ); - - // Use Stagehand's close method which handles browser cleanup properly - await session.stagehand.close(); - stagehandClosedSuccessfully = true; - - process.stderr.write( - `[tool.closeSession] Stagehand and browser connection for session (${previousSessionId}) closed successfully.`, - ); - - // Clean up the session from tracking - await cleanupSession(previousSessionId); - - if (browserbaseSessionId) { - process.stderr.write( - `[tool.closeSession] View session replay at https://www.browserbase.com/sessions/${browserbaseSessionId}`, - ); - } - } else { - process.stderr.write( - `[tool.closeSession] No Stagehand instance found for session: ${previousSessionId || "default/unknown"}`, - ); - } - } catch (error: unknown) { - stagehandCloseErrorMessage = - error instanceof Error ? error.message : String(error); - process.stderr.write( - `[tool.closeSession] Error retrieving or closing Stagehand (session ID was ${previousSessionId || "default/unknown"}): ${stagehandCloseErrorMessage}`, - ); - } + const sessionStore = store(context.config); - // Step 2: Always reset the context's current session ID to default - const oldContextSessionId = context.currentSessionId; - context.currentSessionId = defaultSessionId; - process.stderr.write( - `[tool.closeSession] Session context reset to default. Previous context session ID was ${oldContextSessionId || "default/unknown"}.`, - ); - - // Step 3: Determine the result message - if (stagehandCloseErrorMessage && !stagehandClosedSuccessfully) { - throw new Error( - `Failed to close the Stagehand session (session ID in context was ${previousSessionId || "default/unknown"}). Error: ${stagehandCloseErrorMessage}. Session context has been reset to default.`, - ); - } - - if (stagehandClosedSuccessfully) { - let successMessage = `Browserbase session (${previousSessionId || "default"}) closed successfully via Stagehand. Context reset to default.`; - if (browserbaseSessionId && previousSessionId !== defaultSessionId) { - successMessage += ` View replay at https://www.browserbase.com/sessions/${browserbaseSessionId}`; - } - return { content: [{ type: "text", text: successMessage }] }; - } + // Close all sessions (since this is for single session workflows) + await sessionStore.removeAll(); - // No Stagehand instance was found - let infoMessage = - "No active Stagehand session found to close. Session context has been reset to default."; - if (previousSessionId && previousSessionId !== defaultSessionId) { - infoMessage = `No active Stagehand session found for session ID '${previousSessionId}'. The context has been reset to default.`; + return { + content: [ + { + type: "text", + text: "Browserbase session closed successfully.", + }, + ], + }; + } catch (error: unknown) { + const errorMessage = + error instanceof Error ? error.message : String(error); + throw new Error(`Failed to close Browserbase session: ${errorMessage}`); } - return { content: [{ type: "text", text: infoMessage }] }; }; return { diff --git a/src/types/types.ts b/src/types/types.ts index b516c99..a289e78 100644 --- a/src/types/types.ts +++ b/src/types/types.ts @@ -22,12 +22,7 @@ export type CreateSessionParams = { meta?: Record; }; -export type BrowserSession = { - browser: Browser; - page: Page; - sessionId: string; - stagehand: Stagehand; -}; +// BrowserSession type removed - use StagehandSession instead export type ToolActionResult = | { content?: (ImageContent | TextContent)[] } diff --git a/src/utils/logger.ts b/src/utils/logger.ts new file mode 100644 index 0000000..05f6b85 --- /dev/null +++ b/src/utils/logger.ts @@ -0,0 +1,7 @@ +/** + * Default logger implementation + * Provides consistent logging across the application + */ +export const defaultLogger = (message: string): void => { + process.stderr.write(`${message}\n`); +};