From 6dd44923daa8aa0e81cbfbfa8a567a1eef8bd7b7 Mon Sep 17 00:00:00 2001 From: Pavel Feldman Date: Wed, 30 Jul 2025 20:57:34 -0700 Subject: [PATCH] chore: make tab snapshot structured to mimic it in recorder (#799) --- src/response.ts | 62 +++++++++++++++++++++++++++---- src/sessionLog.ts | 4 +- src/tab.ts | 93 ++++++++++++++++++----------------------------- 3 files changed, 93 insertions(+), 66 deletions(-) diff --git a/src/response.ts b/src/response.ts index 98d77b6..7aa4608 100644 --- a/src/response.ts +++ b/src/response.ts @@ -14,7 +14,11 @@ * limitations under the License. */ -import type { ImageContent, TextContent } from '@modelcontextprotocol/sdk/types.js'; +import { renderModalStates } from './tab.js'; + +import type { TabSnapshot } from './tab.js'; +import type { ModalState } from './tools/tool.js'; +import type { ImageContent, TextContent } from '@modelcontextprotocol/sdk/types.js'; import type { Context } from './context.js'; export class Response { @@ -24,7 +28,7 @@ export class Response { private _context: Context; private _includeSnapshot = false; private _includeTabs = false; - private _snapshot: string | undefined; + private _snapshot: { tabSnapshot?: TabSnapshot, modalState?: ModalState } | undefined; readonly toolName: string; readonly toolArgs: Record; @@ -77,13 +81,13 @@ export class Response { this._includeTabs = true; } - async snapshot(): Promise { - if (this._snapshot !== undefined) + async snapshot(): Promise<{ tabSnapshot?: TabSnapshot, modalState?: ModalState }> { + if (this._snapshot) return this._snapshot; if (this._includeSnapshot && this._context.currentTab()) this._snapshot = await this._context.currentTabOrDie().captureSnapshot(); else - this._snapshot = ''; + this._snapshot = {}; return this._snapshot; } @@ -112,8 +116,14 @@ ${this._code.join('\n')} // Add snapshot if provided. const snapshot = await this.snapshot(); - if (snapshot) - response.push(snapshot, ''); + if (snapshot?.modalState) { + response.push(...renderModalStates(this._context, [snapshot.modalState])); + response.push(''); + } + if (snapshot?.tabSnapshot) { + response.push(renderTabSnapshot(snapshot.tabSnapshot)); + response.push(''); + } // Main response part const content: (TextContent | ImageContent)[] = [ @@ -129,3 +139,41 @@ ${this._code.join('\n')} return { content, isError: this._isError }; } } + +function renderTabSnapshot(tabSnapshot: TabSnapshot): string { + const lines: string[] = []; + + if (tabSnapshot.consoleMessages.length) { + lines.push(`### New console messages`); + for (const message of tabSnapshot.consoleMessages) + lines.push(`- ${trim(message.toString(), 100)}`); + lines.push(''); + } + + if (tabSnapshot.downloads.length) { + lines.push(`### Downloads`); + for (const entry of tabSnapshot.downloads) { + if (entry.finished) + lines.push(`- Downloaded file ${entry.download.suggestedFilename()} to ${entry.outputFile}`); + else + lines.push(`- Downloading file ${entry.download.suggestedFilename()} ...`); + } + lines.push(''); + } + + lines.push(`### Page state`); + lines.push(`- Page URL: ${tabSnapshot.url}`); + lines.push(`- Page Title: ${tabSnapshot.title}`); + lines.push(`- Page Snapshot:`); + lines.push('```yaml'); + lines.push(tabSnapshot.ariaSnapshot); + lines.push('```'); + + return lines.join('\n'); +} + +function trim(text: string, maxLength: number) { + if (text.length <= maxLength) + return text; + return text.slice(0, maxLength) + '...'; +} diff --git a/src/sessionLog.ts b/src/sessionLog.ts index 21a0053..e4099e6 100644 --- a/src/sessionLog.ts +++ b/src/sessionLog.ts @@ -75,9 +75,9 @@ export class SessionLog { } const snapshot = await response.snapshot(); - if (snapshot) { + if (snapshot?.tabSnapshot) { const fileName = `${prefix}.snapshot.yml`; - await fs.promises.writeFile(path.join(this._folder, fileName), snapshot); + await fs.promises.writeFile(path.join(this._folder, fileName), snapshot.tabSnapshot?.ariaSnapshot); lines.push(`- Snapshot: ${fileName}`); } diff --git a/src/tab.ts b/src/tab.ts index 23ca659..ca5e0e9 100644 --- a/src/tab.ts +++ b/src/tab.ts @@ -36,6 +36,15 @@ export type TabEventsInterface = { [TabEvents.modalState]: [modalState: ModalState]; }; +export type TabSnapshot = { + url: string; + title: string; + ariaSnapshot: string; + modalStates: ModalState[]; + consoleMessages: ConsoleMessage[]; + downloads: { download: playwright.Download, finished: boolean, outputFile: string }[]; +}; + export class Tab extends EventEmitter { readonly context: Context; readonly page: playwright.Page; @@ -90,14 +99,7 @@ export class Tab extends EventEmitter { } modalStatesMarkdown(): string[] { - const result: string[] = ['### Modal state']; - if (this._modalStates.length === 0) - result.push('- There is no modal state present'); - for (const state of this._modalStates) { - const tool = this.context.tools.filter(tool => 'clearsModalState' in tool).find(tool => tool.clearsModalState === state.type); - result.push(`- [${state.description}]: can be handled by the "${tool?.schema.name}" tool`); - } - return result; + return renderModalStates(this.context, this.modalStates()); } private _dialogShown(dialog: playwright.Dialog) { @@ -180,53 +182,25 @@ export class Tab extends EventEmitter { return this._requests; } - private _takeRecentConsoleMarkdown(): string[] { - if (!this._recentConsoleMessages.length) - return []; - const result = this._recentConsoleMessages.map(message => { - return `- ${trim(message.toString(), 100)}`; - }); - return [`### New console messages`, ...result, '']; - } - - private _listDownloadsMarkdown(): string[] { - if (!this._downloads.length) - return []; - - const result: string[] = ['### Downloads']; - for (const entry of this._downloads) { - if (entry.finished) - result.push(`- Downloaded file ${entry.download.suggestedFilename()} to ${entry.outputFile}`); - else - result.push(`- Downloading file ${entry.download.suggestedFilename()} ...`); - } - result.push(''); - return result; - } - - async captureSnapshot(): Promise { - const result: string[] = []; - if (this.modalStates().length) { - result.push(...this.modalStatesMarkdown()); - return result.join('\n'); - } - - result.push(...this._takeRecentConsoleMarkdown()); - result.push(...this._listDownloadsMarkdown()); - - await this._raceAgainstModalStates(async () => { + async captureSnapshot(): Promise<{ tabSnapshot?: TabSnapshot, modalState?: ModalState }> { + let tabSnapshot: TabSnapshot | undefined; + const modalState = await this._raceAgainstModalStates(async () => { const snapshot = await (this.page as PageEx)._snapshotForAI(); - result.push( - `### Page state`, - `- Page URL: ${this.page.url()}`, - `- Page Title: ${await this.page.title()}`, - `- Page Snapshot:`, - '```yaml', - snapshot, - '```', - ); + tabSnapshot = { + url: this.page.url(), + title: await this.page.title(), + ariaSnapshot: snapshot, + modalStates: this.modalStates(), + consoleMessages: [], + downloads: this._downloads, + }; }); - return result.join('\n'); + if (tabSnapshot) { + // Assign console message late so that we did not lose any to modal state. + tabSnapshot.consoleMessages = this._recentConsoleMessages; + this._recentConsoleMessages = []; + } + return { tabSnapshot, modalState }; } private _javaScriptBlocked(): boolean { @@ -308,10 +282,15 @@ function pageErrorToConsoleMessage(errorOrValue: Error | any): ConsoleMessage { }; } -function trim(text: string, maxLength: number) { - if (text.length <= maxLength) - return text; - return text.slice(0, maxLength) + '...'; +export function renderModalStates(context: Context, modalStates: ModalState[]): string[] { + const result: string[] = ['### Modal state']; + if (modalStates.length === 0) + result.push('- There is no modal state present'); + for (const state of modalStates) { + const tool = context.tools.filter(tool => 'clearsModalState' in tool).find(tool => tool.clearsModalState === state.type); + result.push(`- [${state.description}]: can be handled by the "${tool?.schema.name}" tool`); + } + return result; } const tabSymbol = Symbol('tabSymbol');