chore: make tab snapshot structured to mimic it in recorder (#799)

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<string, any>;
@@ -77,13 +81,13 @@ export class Response {
     this._includeTabs = true;
   }
 
-  async snapshot(): Promise<string> {
-    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) + '...';
+}

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}`);
     }
 

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<TabEventsInterface> {
   readonly context: Context;
   readonly page: playwright.Page;
@@ -90,14 +99,7 @@ export class Tab extends EventEmitter<TabEventsInterface> {
   }
 
   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<TabEventsInterface> {
     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<string> {
-    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');