diff --git a/packages/agent/src/tools/session/SessionTracker.ts b/packages/agent/src/tools/session/SessionTracker.ts new file mode 100644 index 0000000..d26d272 --- /dev/null +++ b/packages/agent/src/tools/session/SessionTracker.ts @@ -0,0 +1,302 @@ +import { chromium } from '@playwright/test'; +import { v4 as uuidv4 } from 'uuid'; + +import { + browserSessions, + BrowserConfig, + Session, + BrowserError, + BrowserErrorCode, +} from './lib/types.js'; + +// Status of a browser session +export enum SessionStatus { + RUNNING = 'running', + COMPLETED = 'completed', + ERROR = 'error', + TERMINATED = 'terminated', +} + +export interface SessionInfo { + id: string; + status: SessionStatus; + startTime: Date; + endTime?: Date; + metadata: { + url?: string; + contentLength?: number; + closedExplicitly?: boolean; + error?: string; + actionType?: string; + closedByCleanup?: boolean; + }; +} + +/** + * Registry to keep track of browser sessions + */ +export class SessionTracker { + private sessions: Map = new Map(); + private readonly defaultConfig: BrowserConfig = { + headless: true, + defaultTimeout: 30000, + }; + + constructor(public ownerAgentId: string | undefined) { + // Store a reference to the instance globally for cleanup + // This allows the CLI to access the instance for cleanup + (globalThis as Record).__BROWSER_MANAGER__ = this; + + // Set up cleanup handlers for graceful shutdown + this.setupGlobalCleanup(); + } + + // Register a new browser session + registerBrowser(url?: string): string { + const id = uuidv4(); + this.sessions.set(id, { + id, + status: SessionStatus.RUNNING, + startTime: new Date(), + metadata: { + url: url || 'about:blank', + }, + }); + return id; + } + + // Update the status of a session + updateSessionStatus( + id: string, + status: SessionStatus, + metadata?: Record, + ): boolean { + const session = this.sessions.get(id); + if (!session) return false; + + session.status = status; + + // If the session is no longer running, set the end time + if (status !== SessionStatus.RUNNING) { + session.endTime = new Date(); + } + + // Update metadata if provided + if (metadata) { + session.metadata = { ...session.metadata, ...metadata }; + } + + return true; + } + + // Get all sessions + getSessions(): SessionInfo[] { + return Array.from(this.sessions.values()); + } + + // Get a session by ID + getSessionById(id: string): SessionInfo | undefined { + return this.sessions.get(id); + } + + // Get sessions by status + getSessionsByStatus(status: SessionStatus): SessionInfo[] { + return this.getSessions().filter((session) => session.status === status); + } + + /** + * Creates a new browser session + * @param config Optional browser configuration + * @returns A promise that resolves to a browser session + */ + async createSession(config?: BrowserConfig): Promise { + try { + const sessionConfig = { ...this.defaultConfig, ...config }; + const browser = await chromium.launch({ + headless: sessionConfig.headless, + }); + + // Create a new context (equivalent to Puppeteer's incognito context) + const context = await browser.newContext({ + viewport: null, + userAgent: + 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36', + }); + + const page = await context.newPage(); + page.setDefaultTimeout(sessionConfig.defaultTimeout ?? 30000); + + const id = this.registerBrowser(); + + const session: Session = { + browser, + page, + id, + }; + + browserSessions.set(id, session); + this.setupCleanup(session); + + return session; + } catch (error) { + throw new BrowserError( + 'Failed to create browser session', + BrowserErrorCode.LAUNCH_FAILED, + error, + ); + } + } + + /** + * Closes a browser session by ID + * @param sessionId The ID of the session to close + */ + async closeSession(sessionId: string): Promise { + const session = browserSessions.get(sessionId); + if (!session) { + throw new BrowserError( + 'Session not found', + BrowserErrorCode.SESSION_ERROR, + ); + } + + try { + // In Playwright, we should close the context which will automatically close its pages + await session.page.context().close(); + await session.browser.close(); + browserSessions.delete(sessionId); + + // Update session status + this.updateSessionStatus(sessionId, SessionStatus.COMPLETED, { + closedExplicitly: true, + }); + } catch (error) { + throw new BrowserError( + 'Failed to close session', + BrowserErrorCode.SESSION_ERROR, + error, + ); + } + } + + /** + * Gets a browser session by ID + * @param sessionId The ID of the session to get + * @returns The browser session + */ + getSession(sessionId: string): Session { + const session = browserSessions.get(sessionId); + if (!session) { + throw new BrowserError( + 'Session not found', + BrowserErrorCode.SESSION_ERROR, + ); + } + return session; + } + + /** + * Cleans up all browser sessions associated with this tracker + * @returns A promise that resolves when cleanup is complete + */ + async cleanup(): Promise { + const cleanupPromises = Array.from(this.sessions.values()).map((session) => + this.cleanupSession(session), + ); + await Promise.all(cleanupPromises); + } + + /** + * Cleans up a browser session + * @param session The browser session to clean up + */ + private async cleanupSession(session: SessionInfo): Promise { + // Only clean up running sessions + if (session.status !== SessionStatus.RUNNING) return; + + const browserSession = browserSessions.get(session.id); + if (!browserSession) return; + + try { + // Close the browser session + await browserSession.page.context().close(); + await browserSession.browser.close(); + browserSessions.delete(session.id); + + // Update session status + this.updateSessionStatus(session.id, SessionStatus.TERMINATED, { + closedByCleanup: true, + }); + } catch { + // Ignore errors during cleanup + } + } + + /** + * Closes all browser sessions + * @returns A promise that resolves when all sessions are closed + */ + async closeAllSessions(): Promise { + const closePromises = Array.from(this.sessions.keys()) + .filter( + (sessionId) => + this.sessions.get(sessionId)?.status === SessionStatus.RUNNING, + ) + .map((sessionId) => this.closeSession(sessionId).catch(() => {})); + await Promise.all(closePromises); + } + + /** + * Sets up cleanup handlers for a browser session + * @param session The browser session to set up cleanup handlers for + */ + private setupCleanup(session: Session): void { + // Handle browser disconnection + session.browser.on('disconnected', () => { + browserSessions.delete(session.id); + this.updateSessionStatus(session.id, SessionStatus.TERMINATED); + }); + } + + /** + * Sets up global cleanup handlers for all browser sessions + */ + private setupGlobalCleanup(): void { + // Use beforeExit for async cleanup + process.on('beforeExit', () => { + this.closeAllSessions().catch((err) => { + console.error('Error closing browser sessions:', err); + }); + }); + + // Use exit for synchronous cleanup (as a fallback) + process.on('exit', () => { + // Can only do synchronous operations here + for (const sessionId of browserSessions.keys()) { + try { + const session = browserSessions.get(sessionId); + if (session) { + // Attempt synchronous close - may not fully work + session.browser.close(); + } + // eslint-disable-next-line unused-imports/no-unused-vars + } catch (e) { + // Ignore errors during exit + } + } + }); + + // Handle SIGINT (Ctrl+C) + process.on('SIGINT', () => { + // eslint-disable-next-line promise/catch-or-return + this.closeAllSessions() + .catch(() => { + return false; + }) + .finally(() => { + // Give a moment for cleanup to complete + setTimeout(() => process.exit(0), 500); + }); + }); + } +}