From b358e47d715dc50ece0c1fa37f6434337caff95a Mon Sep 17 00:00:00 2001 From: Pavel Feldman Date: Thu, 3 Apr 2025 10:30:05 -0700 Subject: [PATCH] chore: prep for multiple pages in context (#124) --- src/context.ts | 190 ++++++++++++++++++++++----------------- src/resources/console.ts | 2 +- src/tools/common.ts | 27 +++--- src/tools/screenshot.ts | 32 +++---- src/tools/snapshot.ts | 28 +++--- 5 files changed, 153 insertions(+), 126 deletions(-) diff --git a/src/context.ts b/src/context.ts index 54c7994..4dd4736 100644 --- a/src/context.ts +++ b/src/context.ts @@ -43,47 +43,50 @@ type RunOptions = { export class Context { private _options: ContextOptions; private _browser: playwright.Browser | undefined; - private _page: playwright.Page | undefined; - private _console: playwright.ConsoleMessage[] = []; - private _createPagePromise: Promise | undefined; - private _fileChooser: playwright.FileChooser | undefined; - private _snapshot: PageSnapshot | undefined; + private _browserContext: playwright.BrowserContext | undefined; + private _pages: Page[] = []; + private _currentPage: Page | undefined; + private _createContextPromise: Promise | undefined; constructor(options: ContextOptions) { this._options = options; } async createPage(): Promise { - if (this._createPagePromise) - return this._createPagePromise; - this._createPagePromise = (async () => { - const { browser, page } = await this._createPage(); - page.on('console', event => this._console.push(event)); - page.on('framenavigated', frame => { - if (!frame.parentFrame()) - this._console.length = 0; - }); - page.on('close', () => this._onPageClose()); - page.on('filechooser', chooser => this._fileChooser = chooser); - page.setDefaultNavigationTimeout(60000); - page.setDefaultTimeout(5000); - this._page = page; + if (this._createContextPromise) + return this._createContextPromise; + this._createContextPromise = (async () => { + const { browser, browserContext } = await this._createBrowserContext(); + const pages = browserContext.pages(); + for (const page of pages) + this._onPageCreated(page); + browserContext.on('page', page => this._onPageCreated(page)); + let page = pages[0]; + if (!page) + page = await browserContext.newPage(); + this._currentPage = this._pages[0]; this._browser = browser; + this._browserContext = browserContext; return page; })(); - return this._createPagePromise; + return this._createContextPromise; } - private _onPageClose() { - const browser = this._browser; - const page = this._page; - void page?.context()?.close().then(() => browser?.close()).catch(() => {}); + private _onPageCreated(page: playwright.Page) { + this._pages.push(new Page(page, page => this._onPageClose(page))); + } - this._createPagePromise = undefined; - this._browser = undefined; - this._page = undefined; - this._fileChooser = undefined; - this._console.length = 0; + private _onPageClose(page: Page) { + this._pages = this._pages.filter(p => p !== page); + if (this._currentPage === page) + this._currentPage = this._pages[0]; + const browser = this._browser; + if (this._browserContext && !this._pages.length) { + void this._browserContext.close().then(() => browser?.close()).catch(() => {}); + this._createContextPromise = undefined; + this._browser = undefined; + this._browserContext = undefined; + } } async install(): Promise { @@ -105,24 +108,90 @@ export class Context { }); } - existingPage(): playwright.Page { - if (!this._page) + currentPage(): Page { + if (!this._currentPage) throw new Error('Navigate to a location to create a page'); - return this._page; + return this._currentPage; } - async run(callback: (page: playwright.Page) => Promise, options?: RunOptions): Promise { - const page = this.existingPage(); + async close() { + if (!this._browserContext) + return; + await this._browserContext.close(); + } + + private async _createBrowserContext(): Promise<{ browser?: playwright.Browser, browserContext: playwright.BrowserContext }> { + if (this._options.remoteEndpoint) { + const url = new URL(this._options.remoteEndpoint); + if (this._options.browserName) + url.searchParams.set('browser', this._options.browserName); + if (this._options.launchOptions) + url.searchParams.set('launch-options', JSON.stringify(this._options.launchOptions)); + const browser = await playwright[this._options.browserName ?? 'chromium'].connect(String(url)); + const browserContext = await browser.newContext(); + return { browser, browserContext }; + } + + if (this._options.cdpEndpoint) { + const browser = await playwright.chromium.connectOverCDP(this._options.cdpEndpoint); + const browserContext = browser.contexts()[0]; + return { browser, browserContext }; + } + + const browserContext = await this._launchPersistentContext(); + return { browserContext }; + } + + private async _launchPersistentContext(): Promise { + try { + const browserType = this._options.browserName ? playwright[this._options.browserName] : playwright.chromium; + return await browserType.launchPersistentContext(this._options.userDataDir, this._options.launchOptions); + } catch (error: any) { + if (error.message.includes('Executable doesn\'t exist')) + throw new Error(`Browser specified in your config is not installed. Either install it (likely) or change the config.`); + throw error; + } + } +} + +class Page { + readonly page: playwright.Page; + private _console: playwright.ConsoleMessage[] = []; + private _fileChooser: playwright.FileChooser | undefined; + private _snapshot: PageSnapshot | undefined; + private _onPageClose: (page: Page) => void; + + constructor(page: playwright.Page, onPageClose: (page: Page) => void) { + this.page = page; + this._onPageClose = onPageClose; + page.on('console', event => this._console.push(event)); + page.on('framenavigated', frame => { + if (!frame.parentFrame()) + this._console.length = 0; + }); + page.on('close', () => this._onClose()); + page.on('filechooser', chooser => this._fileChooser = chooser); + page.setDefaultNavigationTimeout(60000); + page.setDefaultTimeout(5000); + } + + private _onClose() { + this._fileChooser = undefined; + this._console.length = 0; + this._onPageClose(this); + } + + async run(callback: (page: Page) => Promise, options?: RunOptions): Promise { try { if (!options?.noClearFileChooser) this._fileChooser = undefined; if (options?.waitForCompletion) - await waitForCompletion(page, () => callback(page)); + await waitForCompletion(this.page, () => callback(this)); else - await callback(page); + await callback(this); } finally { if (options?.captureSnapshot) - this._snapshot = await PageSnapshot.create(page); + this._snapshot = await PageSnapshot.create(this.page); } return { content: [{ @@ -132,14 +201,14 @@ export class Context { }; } - async runAndWait(callback: (page: playwright.Page) => Promise, options?: RunOptions): Promise { + async runAndWait(callback: (page: Page) => Promise, options?: RunOptions): Promise { return await this.run(callback, { waitForCompletion: true, ...options, }); } - async runAndWaitWithSnapshot(callback: (page: playwright.Page) => Promise, options?: RunOptions): Promise { + async runAndWaitWithSnapshot(callback: (page: Page) => Promise, options?: RunOptions): Promise { return await this.run(callback, { captureSnapshot: true, waitForCompletion: true, @@ -157,55 +226,12 @@ export class Context { return this._console; } - async close() { - if (!this._page) - return; - await this._page.close(); - } - async submitFileChooser(paths: string[]) { if (!this._fileChooser) throw new Error('No file chooser visible'); await this._fileChooser.setFiles(paths); this._fileChooser = undefined; } - - private async _createPage(): Promise<{ browser?: playwright.Browser, page: playwright.Page }> { - if (this._options.remoteEndpoint) { - const url = new URL(this._options.remoteEndpoint); - if (this._options.browserName) - url.searchParams.set('browser', this._options.browserName); - if (this._options.launchOptions) - url.searchParams.set('launch-options', JSON.stringify(this._options.launchOptions)); - const browser = await playwright[this._options.browserName ?? 'chromium'].connect(String(url)); - const page = await browser.newPage(); - return { browser, page }; - } - - if (this._options.cdpEndpoint) { - const browser = await playwright.chromium.connectOverCDP(this._options.cdpEndpoint); - const browserContext = browser.contexts()[0]; - let [page] = browserContext.pages(); - if (!page) - page = await browserContext.newPage(); - return { browser, page }; - } - - const context = await this._launchPersistentContext(); - const [page] = context.pages(); - return { page }; - } - - private async _launchPersistentContext(): Promise { - try { - const browserType = this._options.browserName ? playwright[this._options.browserName] : playwright.chromium; - return await browserType.launchPersistentContext(this._options.userDataDir, this._options.launchOptions); - } catch (error: any) { - if (error.message.includes('Executable doesn\'t exist')) - throw new Error(`Browser specified in your config is not installed. Either install it (likely) or change the config.`); - throw error; - } - } } class PageSnapshot { diff --git a/src/resources/console.ts b/src/resources/console.ts index ca9bea8..c93f838 100644 --- a/src/resources/console.ts +++ b/src/resources/console.ts @@ -24,7 +24,7 @@ export const console: Resource = { }, read: async (context, uri) => { - const messages = await context.console(); + const messages = await context.currentPage().console(); const log = messages.map(message => `[${message.type().toUpperCase()}] ${message.text()}`).join('\n'); return [{ uri, diff --git a/src/tools/common.ts b/src/tools/common.ts index c91d111..86b22e3 100644 --- a/src/tools/common.ts +++ b/src/tools/common.ts @@ -37,10 +37,10 @@ export const navigate: ToolFactory = captureSnapshot => ({ handle: async (context, params) => { const validatedParams = navigateSchema.parse(params); await context.createPage(); - return await context.run(async page => { - await page.goto(validatedParams.url, { waitUntil: 'domcontentloaded' }); + return await context.currentPage().run(async page => { + await page.page.goto(validatedParams.url, { waitUntil: 'domcontentloaded' }); // Cap load event to 5 seconds, the page is operational at this point. - await page.waitForLoadState('load', { timeout: 5000 }).catch(() => {}); + await page.page.waitForLoadState('load', { timeout: 5000 }).catch(() => {}); }, { status: `Navigated to ${validatedParams.url}`, captureSnapshot, @@ -57,8 +57,8 @@ export const goBack: ToolFactory = snapshot => ({ inputSchema: zodToJsonSchema(goBackSchema), }, handle: async context => { - return await context.runAndWait(async page => { - await page.goBack(); + return await context.currentPage().runAndWait(async page => { + await page.page.goBack(); }, { status: 'Navigated back', captureSnapshot: snapshot, @@ -75,8 +75,8 @@ export const goForward: ToolFactory = snapshot => ({ inputSchema: zodToJsonSchema(goForwardSchema), }, handle: async context => { - return await context.runAndWait(async page => { - await page.goForward(); + return await context.currentPage().runAndWait(async page => { + await page.page.goForward(); }, { status: 'Navigated forward', captureSnapshot: snapshot, @@ -118,8 +118,8 @@ export const pressKey: (captureSnapshot: boolean) => Tool = captureSnapshot => ( }, handle: async (context, params) => { const validatedParams = pressKeySchema.parse(params); - return await context.runAndWait(async page => { - await page.keyboard.press(validatedParams.key); + return await context.currentPage().runAndWait(async page => { + await page.page.keyboard.press(validatedParams.key); }, { status: `Pressed key ${validatedParams.key}`, captureSnapshot, @@ -136,9 +136,9 @@ export const pdf: Tool = { inputSchema: zodToJsonSchema(pdfSchema), }, handle: async context => { - const page = context.existingPage(); + const page = context.currentPage(); const fileName = path.join(os.tmpdir(), sanitizeForFilePath(`page-${new Date().toISOString()}`)) + '.pdf'; - await page.pdf({ path: fileName }); + await page.page.pdf({ path: fileName }); return { content: [{ type: 'text', @@ -179,8 +179,9 @@ export const chooseFile: ToolFactory = captureSnapshot => ({ }, handle: async (context, params) => { const validatedParams = chooseFileSchema.parse(params); - return await context.runAndWait(async () => { - await context.submitFileChooser(validatedParams.paths); + const page = context.currentPage(); + return await page.runAndWait(async () => { + await page.submitFileChooser(validatedParams.paths); }, { status: `Chose files ${validatedParams.paths.join(', ')}`, captureSnapshot, diff --git a/src/tools/screenshot.ts b/src/tools/screenshot.ts index 46001ec..b9c5bb4 100644 --- a/src/tools/screenshot.ts +++ b/src/tools/screenshot.ts @@ -27,8 +27,8 @@ export const screenshot: Tool = { }, handle: async context => { - const page = context.existingPage(); - const screenshot = await page.screenshot({ type: 'jpeg', quality: 50, scale: 'css' }); + const page = context.currentPage(); + const screenshot = await page.page.screenshot({ type: 'jpeg', quality: 50, scale: 'css' }); return { content: [{ type: 'image', data: screenshot.toString('base64'), mimeType: 'image/jpeg' }], }; @@ -53,8 +53,8 @@ export const moveMouse: Tool = { handle: async (context, params) => { const validatedParams = moveMouseSchema.parse(params); - const page = context.existingPage(); - await page.mouse.move(validatedParams.x, validatedParams.y); + const page = context.currentPage(); + await page.page.mouse.move(validatedParams.x, validatedParams.y); return { content: [{ type: 'text', text: `Moved mouse to (${validatedParams.x}, ${validatedParams.y})` }], }; @@ -74,11 +74,11 @@ export const click: Tool = { }, handle: async (context, params) => { - return await context.runAndWait(async page => { + return await context.currentPage().runAndWait(async page => { const validatedParams = clickSchema.parse(params); - await page.mouse.move(validatedParams.x, validatedParams.y); - await page.mouse.down(); - await page.mouse.up(); + await page.page.mouse.move(validatedParams.x, validatedParams.y); + await page.page.mouse.down(); + await page.page.mouse.up(); }, { status: 'Clicked mouse', }); @@ -101,11 +101,11 @@ export const drag: Tool = { handle: async (context, params) => { const validatedParams = dragSchema.parse(params); - return await context.runAndWait(async page => { - await page.mouse.move(validatedParams.startX, validatedParams.startY); - await page.mouse.down(); - await page.mouse.move(validatedParams.endX, validatedParams.endY); - await page.mouse.up(); + return await context.currentPage().runAndWait(async page => { + await page.page.mouse.move(validatedParams.startX, validatedParams.startY); + await page.page.mouse.down(); + await page.page.mouse.move(validatedParams.endX, validatedParams.endY); + await page.page.mouse.up(); }, { status: `Dragged mouse from (${validatedParams.startX}, ${validatedParams.startY}) to (${validatedParams.endX}, ${validatedParams.endY})`, }); @@ -126,10 +126,10 @@ export const type: Tool = { handle: async (context, params) => { const validatedParams = typeSchema.parse(params); - return await context.runAndWait(async page => { - await page.keyboard.type(validatedParams.text); + return await context.currentPage().runAndWait(async page => { + await page.page.keyboard.type(validatedParams.text); if (validatedParams.submit) - await page.keyboard.press('Enter'); + await page.page.keyboard.press('Enter'); }, { status: `Typed text "${validatedParams.text}"`, }); diff --git a/src/tools/snapshot.ts b/src/tools/snapshot.ts index 029c2a0..f974b19 100644 --- a/src/tools/snapshot.ts +++ b/src/tools/snapshot.ts @@ -28,7 +28,7 @@ export const snapshot: Tool = { }, handle: async context => { - return await context.run(async () => {}, { captureSnapshot: true }); + return await context.currentPage().run(async () => {}, { captureSnapshot: true }); }, }; @@ -46,8 +46,8 @@ export const click: Tool = { handle: async (context, params) => { const validatedParams = elementSchema.parse(params); - return await context.runAndWaitWithSnapshot(async () => { - const locator = context.lastSnapshot().refLocator(validatedParams.ref); + return await context.currentPage().runAndWaitWithSnapshot(async page => { + const locator = page.lastSnapshot().refLocator(validatedParams.ref); await locator.click(); }, { status: `Clicked "${validatedParams.element}"`, @@ -71,9 +71,9 @@ export const drag: Tool = { handle: async (context, params) => { const validatedParams = dragSchema.parse(params); - return await context.runAndWaitWithSnapshot(async () => { - const startLocator = context.lastSnapshot().refLocator(validatedParams.startRef); - const endLocator = context.lastSnapshot().refLocator(validatedParams.endRef); + return await context.currentPage().runAndWaitWithSnapshot(async page => { + const startLocator = page.lastSnapshot().refLocator(validatedParams.startRef); + const endLocator = page.lastSnapshot().refLocator(validatedParams.endRef); await startLocator.dragTo(endLocator); }, { status: `Dragged "${validatedParams.startElement}" to "${validatedParams.endElement}"`, @@ -90,8 +90,8 @@ export const hover: Tool = { handle: async (context, params) => { const validatedParams = elementSchema.parse(params); - return context.runAndWaitWithSnapshot(async () => { - const locator = context.lastSnapshot().refLocator(validatedParams.ref); + return await context.currentPage().runAndWaitWithSnapshot(async page => { + const locator = page.lastSnapshot().refLocator(validatedParams.ref); await locator.hover(); }, { status: `Hovered over "${validatedParams.element}"`, @@ -114,8 +114,8 @@ export const type: Tool = { handle: async (context, params) => { const validatedParams = typeSchema.parse(params); - return await context.runAndWaitWithSnapshot(async () => { - const locator = context.lastSnapshot().refLocator(validatedParams.ref); + return await context.currentPage().runAndWaitWithSnapshot(async page => { + const locator = page.lastSnapshot().refLocator(validatedParams.ref); if (validatedParams.slowly) await locator.pressSequentially(validatedParams.text); else @@ -141,8 +141,8 @@ export const selectOption: Tool = { handle: async (context, params) => { const validatedParams = selectOptionSchema.parse(params); - return await context.runAndWaitWithSnapshot(async () => { - const locator = context.lastSnapshot().refLocator(validatedParams.ref); + return await context.currentPage().runAndWaitWithSnapshot(async page => { + const locator = page.lastSnapshot().refLocator(validatedParams.ref); await locator.selectOption(validatedParams.values); }, { status: `Selected option in "${validatedParams.element}"`, @@ -163,9 +163,9 @@ export const screenshot: Tool = { handle: async (context, params) => { const validatedParams = screenshotSchema.parse(params); - const page = context.existingPage(); + const page = context.currentPage(); const options: playwright.PageScreenshotOptions = validatedParams.raw ? { type: 'png', scale: 'css' } : { type: 'jpeg', quality: 50, scale: 'css' }; - const screenshot = await page.screenshot(options); + const screenshot = await page.page.screenshot(options); return { content: [{ type: 'image', data: screenshot.toString('base64'), mimeType: validatedParams.raw ? 'image/png' : 'image/jpeg' }], };