feat(trace): streaming snapshots (#5133)

- Instead of capturing snapshots on demand, we now stream them
  from each frame every 100ms.
- Certain actions can also force snapshots at particular moment using
  "checkpoints".
- Trace viewer is able to show the page snapshot at a particular
  timestamp, or using a "checkpoint" snapshot.
- Small optimization to not process stylesheets if CSSOM was not used.
  There still is a lot of room for improvement.

src/cli/cli.ts

@@ -202,6 +202,8 @@ async function launchContext(options: Options, headless: boolean): Promise<{ bro
   if (contextOptions.isMobile && browserType.name() === 'firefox')
     contextOptions.isMobile = undefined;
 
+  if (process.env.PWTRACE)
+    (contextOptions as any)._traceDir = path.join(process.cwd(), '.trace');
 
   // Proxy
 

src/cli/traceViewer/screenshotGenerator.ts

@@ -19,8 +19,7 @@ import * as path from 'path';
 import * as playwright from '../../..';
 import * as util from 'util';
 import { SnapshotRouter } from './snapshotRouter';
-import { actionById, ActionEntry, ContextEntry, TraceModel } from './traceModel';
-import type { PageSnapshot } from '../../trace/traceTypes';
+import { actionById, ActionEntry, ContextEntry, PageEntry, TraceModel } from './traceModel';
 
 const fsReadFileAsync = util.promisify(fs.readFile.bind(fs));
 const fsWriteFileAsync = util.promisify(fs.writeFile.bind(fs));
@@ -39,11 +38,9 @@ export class ScreenshotGenerator {
   }
 
   generateScreenshot(actionId: string): Promise<Buffer | undefined> {
-    const { context, action } = actionById(this._traceModel, actionId);
-    if (!action.action.snapshot)
-      return Promise.resolve(undefined);
+    const { context, action, page } = actionById(this._traceModel, actionId);
     if (!this._rendering.has(action)) {
-      this._rendering.set(action, this._render(context, action).then(body => {
+      this._rendering.set(action, this._render(context, page, action).then(body => {
         this._rendering.delete(action);
         return body;
       }));
@@ -51,8 +48,8 @@ export class ScreenshotGenerator {
     return this._rendering.get(action)!;
   }
 
-  private async _render(contextEntry: ContextEntry, actionEntry: ActionEntry): Promise<Buffer | undefined> {
-    const imageFileName = path.join(this._traceStorageDir, actionEntry.action.snapshot!.sha1 + '-screenshot.png');
+  private async _render(contextEntry: ContextEntry, pageEntry: PageEntry, actionEntry: ActionEntry): Promise<Buffer | undefined> {
+    const imageFileName = path.join(this._traceStorageDir, actionEntry.action.timestamp + '-screenshot.png');
     try {
       return await fsReadFileAsync(imageFileName);
     } catch (e) {
@@ -70,27 +67,24 @@ export class ScreenshotGenerator {
     });
 
     try {
-      const snapshotPath = path.join(this._traceStorageDir, action.snapshot!.sha1);
-      let snapshot;
-      try {
-        snapshot = await fsReadFileAsync(snapshotPath, 'utf8');
-      } catch (e) {
-        console.log(`Unable to read snapshot at ${snapshotPath}`); // eslint-disable-line no-console
-        return;
-      }
-      const snapshotObject = JSON.parse(snapshot) as PageSnapshot;
       const snapshotRouter = new SnapshotRouter(this._traceStorageDir);
-      snapshotRouter.selectSnapshot(snapshotObject, contextEntry);
+      const snapshots = action.snapshots || [];
+      const snapshotId = snapshots.length ? snapshots[0].snapshotId : undefined;
+      const snapshotTimestamp = action.startTime;
+      const pageUrl = await snapshotRouter.selectSnapshot(contextEntry, pageEntry, snapshotId, snapshotTimestamp);
       page.route('**/*', route => snapshotRouter.route(route));
-      const url = snapshotObject.frames[0].url;
-      console.log('Generating screenshot for ' + action.action, snapshotObject.frames[0].url); // eslint-disable-line no-console
-      await page.goto(url);
+      console.log('Generating screenshot for ' + action.action, pageUrl); // eslint-disable-line no-console
+      await page.goto(pageUrl);
 
-      const element = await page.$(action.selector || '*[__playwright_target__]');
-      if (element) {
-        await element.evaluate(e => {
-          e.style.backgroundColor = '#ff69b460';
-        });
+      try {
+        const element = await page.$(action.selector || '*[__playwright_target__]');
+        if (element) {
+          await element.evaluate(e => {
+            e.style.backgroundColor = '#ff69b460';
+          });
+        }
+      } catch (e) {
+        console.log(e); // eslint-disable-line no-console
       }
       const imageData = await page.screenshot();
       await fsWriteFileAsync(imageFileName, imageData);

src/cli/traceViewer/snapshotRouter.ts

@@ -17,55 +17,117 @@
 import * as fs from 'fs';
 import * as path from 'path';
 import * as util from 'util';
-import type { Route } from '../../..';
+import type { Frame, Route } from '../../..';
 import { parsedURL } from '../../client/clientHelper';
-import type { FrameSnapshot, NetworkResourceTraceEvent, PageSnapshot } from '../../trace/traceTypes';
-import { ContextEntry } from './traceModel';
+import { ContextEntry, PageEntry, trace } from './traceModel';
 
 const fsReadFileAsync = util.promisify(fs.readFile.bind(fs));
 
 export class SnapshotRouter {
   private _contextEntry: ContextEntry | undefined;
   private _unknownUrls = new Set<string>();
-  private _traceStorageDir: string;
-  private _frameBySrc = new Map<string, FrameSnapshot>();
+  private _resourcesDir: string;
+  private _snapshotFrameIdToSnapshot = new Map<string, trace.FrameSnapshot>();
+  private _pageUrl = '';
+  private _frameToSnapshotFrameId = new Map<Frame, string>();
 
-  constructor(traceStorageDir: string) {
-    this._traceStorageDir = traceStorageDir;
+  constructor(resourcesDir: string) {
+    this._resourcesDir = resourcesDir;
   }
 
-  selectSnapshot(snapshot: PageSnapshot, contextEntry: ContextEntry) {
-    this._frameBySrc.clear();
+  // Returns the url to navigate to.
+  async selectSnapshot(contextEntry: ContextEntry, pageEntry: PageEntry, snapshotId?: string, timestamp?: number): Promise<string> {
     this._contextEntry = contextEntry;
-    for (const frameSnapshot of snapshot.frames)
-      this._frameBySrc.set(frameSnapshot.url, frameSnapshot);
+    if (!snapshotId && !timestamp)
+      return 'data:text/html,Snapshot is not available';
+
+    const lastSnapshotEvent = new Map<string, trace.FrameSnapshotTraceEvent>();
+    for (const [frameId, snapshots] of pageEntry.snapshotsByFrameId) {
+      for (const snapshot of snapshots) {
+        const current = lastSnapshotEvent.get(frameId);
+        // Prefer snapshot with exact id.
+        const exactMatch = snapshotId && snapshot.snapshotId === snapshotId;
+        const currentExactMatch = current && snapshotId && current.snapshotId === snapshotId;
+        // If not available, prefer the latest snapshot before the timestamp.
+        const timestampMatch = timestamp && snapshot.timestamp <= timestamp;
+        if (exactMatch || (timestampMatch && !currentExactMatch))
+          lastSnapshotEvent.set(frameId, snapshot);
+      }
+    }
+
+    this._snapshotFrameIdToSnapshot.clear();
+    for (const [frameId, event] of lastSnapshotEvent) {
+      const buffer = await this._readSha1(event.sha1);
+      if (!buffer)
+        continue;
+      try {
+        const snapshot = JSON.parse(buffer.toString('utf8')) as trace.FrameSnapshot;
+        // Request url could come lower case, so we always normalize to lower case.
+        this._snapshotFrameIdToSnapshot.set(frameId.toLowerCase(), snapshot);
+      } catch (e) {
+      }
+    }
+
+    const mainFrameSnapshot = lastSnapshotEvent.get('');
+    if (!mainFrameSnapshot)
+      return 'data:text/html,Snapshot is not available';
+
+    if (!mainFrameSnapshot.frameUrl.startsWith('http'))
+      this._pageUrl = 'http://playwright.snapshot/';
+    else
+      this._pageUrl = mainFrameSnapshot.frameUrl;
+    return this._pageUrl;
   }
 
   async route(route: Route) {
     const url = route.request().url();
-    if (this._frameBySrc.has(url)) {
-      const frameSnapshot = this._frameBySrc.get(url)!;
+    const frame = route.request().frame();
+
+    if (route.request().isNavigationRequest()) {
+      let snapshotFrameId: string | undefined;
+      if (url === this._pageUrl) {
+        snapshotFrameId = '';
+      } else {
+        snapshotFrameId = url.substring(url.indexOf('://') + 3);
+        if (snapshotFrameId.endsWith('/'))
+          snapshotFrameId = snapshotFrameId.substring(0, snapshotFrameId.length - 1);
+        // Request url could come lower case, so we always normalize to lower case.
+        snapshotFrameId = snapshotFrameId.toLowerCase();
+      }
+
+      const snapshot = this._snapshotFrameIdToSnapshot.get(snapshotFrameId);
+      if (!snapshot) {
+        route.fulfill({
+          contentType: 'text/html',
+          body: 'data:text/html,Snapshot is not available',
+        });
+        return;
+      }
+
+      this._frameToSnapshotFrameId.set(frame, snapshotFrameId);
       route.fulfill({
         contentType: 'text/html',
-        body: Buffer.from(frameSnapshot.html),
+        body: snapshot.html,
       });
       return;
     }
 
-    const frameSrc = route.request().frame().url();
-    const frameSnapshot = this._frameBySrc.get(frameSrc);
-    if (!frameSnapshot)
+    const snapshotFrameId = this._frameToSnapshotFrameId.get(frame);
+    if (snapshotFrameId === undefined)
+      return this._routeUnknown(route);
+    const snapshot = this._snapshotFrameIdToSnapshot.get(snapshotFrameId);
+    if (!snapshot)
       return this._routeUnknown(route);
 
     // Find a matching resource from the same context, preferrably from the same frame.
     // Note: resources are stored without hash, but page may reference them with hash.
-    let resource: NetworkResourceTraceEvent | null = null;
+    let resource: trace.NetworkResourceTraceEvent | null = null;
     const resourcesWithUrl = this._contextEntry!.resourcesByUrl.get(removeHash(url)) || [];
     for (const resourceEvent of resourcesWithUrl) {
-      if (resource && resourceEvent.frameId !== frameSnapshot.frameId)
+      if (resource && resourceEvent.frameId !== snapshotFrameId)
         continue;
       resource = resourceEvent;
-      if (resourceEvent.frameId === frameSnapshot.frameId)
+      if (resourceEvent.frameId === snapshotFrameId)
         break;
     }
     if (!resource)
@@ -73,7 +135,7 @@ export class SnapshotRouter {
 
     // This particular frame might have a resource content override, for example when
     // stylesheet is modified using CSSOM.
-    const resourceOverride = frameSnapshot.resourceOverrides.find(o => o.url === url);
+    const resourceOverride = snapshot.resourceOverrides.find(o => o.url === url);
     const overrideSha1 = resourceOverride ? resourceOverride.sha1 : undefined;
     const resourceData = await this._readResource(resource, overrideSha1);
     if (!resourceData)
@@ -98,18 +160,24 @@ export class SnapshotRouter {
     route.abort();
   }
 
-  private async _readResource(event: NetworkResourceTraceEvent, overrideSha1: string | undefined) {
+  private async _readSha1(sha1: string) {
     try {
-      const body = await fsReadFileAsync(path.join(this._traceStorageDir, overrideSha1 || event.sha1));
-      return {
-        contentType: event.contentType,
-        body,
-        headers: event.responseHeaders,
-      };
+      return await fsReadFileAsync(path.join(this._resourcesDir, sha1));
     } catch (e) {
       return undefined;
     }
   }
+
+  private async _readResource(event: trace.NetworkResourceTraceEvent, overrideSha1: string | undefined) {
+    const body = await this._readSha1(overrideSha1 || event.sha1);
+    if (!body)
+      return;
+    return {
+      contentType: event.contentType,
+      body,
+      headers: event.responseHeaders,
+    };
+  }
 }
 
 function removeHash(url: string) {

src/cli/traceViewer/traceModel.ts

@@ -46,6 +46,7 @@ export type PageEntry = {
   actions: ActionEntry[];
   interestingEvents: InterestingPageEvent[];
   resources: trace.NetworkResourceTraceEvent[];
+  snapshotsByFrameId: Map<string, trace.FrameSnapshotTraceEvent[]>;
 }
 
 export type ActionEntry = {
@@ -93,6 +94,7 @@ export function readTraceFile(events: trace.TraceEvent[], traceModel: TraceModel
           actions: [],
           resources: [],
           interestingEvents: [],
+          snapshotsByFrameId: new Map(),
         };
         pageEntries.set(event.pageId, pageEntry);
         contextEntries.get(event.contextId)!.pages.push(pageEntry);
@@ -144,6 +146,13 @@ export function readTraceFile(events: trace.TraceEvent[], traceModel: TraceModel
         pageEntry.interestingEvents.push(event);
         break;
       }
+      case 'snapshot': {
+        const pageEntry = pageEntries.get(event.pageId!)!;
+        if (!pageEntry.snapshotsByFrameId.has(event.frameId))
+          pageEntry.snapshotsByFrameId.set(event.frameId, []);
+        pageEntry.snapshotsByFrameId.get(event.frameId)!.push(event);
+        break;
+      }
     }
 
     const contextEntry = contextEntries.get(event.contextId)!;

src/cli/traceViewer/traceViewer.ts

@@ -21,7 +21,7 @@ import * as util from 'util';
 import { ScreenshotGenerator } from './screenshotGenerator';
 import { SnapshotRouter } from './snapshotRouter';
 import { readTraceFile, TraceModel } from './traceModel';
-import type { ActionTraceEvent, PageSnapshot, TraceEvent } from '../../trace/traceTypes';
+import type { ActionTraceEvent, TraceEvent } from '../../trace/traceTypes';
 
 const fsReadFileAsync = util.promisify(fs.readFile.bind(fs));
 
@@ -92,25 +92,20 @@ class TraceViewer {
     await uiPage.exposeBinding('readFile', async (_, path: string) => {
       return fs.readFileSync(path).toString();
     });
-    await uiPage.exposeBinding('renderSnapshot', async (_, action: ActionTraceEvent) => {
+    await uiPage.exposeBinding('renderSnapshot', async (_, arg: { action: ActionTraceEvent, snapshot: { name: string, snapshotId?: string } }) => {
+      const { action, snapshot } = arg;
       if (!this._document)
         return;
       try {
-        if (!action.snapshot) {
-          const snapshotFrame = uiPage.frames()[1];
-          await snapshotFrame.goto('data:text/html,No snapshot available');
-          return;
-        }
-
-        const snapshot = await fsReadFileAsync(path.join(this._document.resourcesDir, action.snapshot!.sha1), 'utf8');
-        const snapshotObject = JSON.parse(snapshot) as PageSnapshot;
         const contextEntry = this._document.model.contexts.find(entry => entry.created.contextId === action.contextId)!;
-        this._document.snapshotRouter.selectSnapshot(snapshotObject, contextEntry);
+        const pageEntry = contextEntry.pages.find(entry => entry.created.pageId === action.pageId)!;
+        const snapshotTime = snapshot.name === 'before' ? action.startTime : (snapshot.name === 'after' ? action.endTime : undefined);
+        const pageUrl = await this._document.snapshotRouter.selectSnapshot(contextEntry, pageEntry, snapshot.snapshotId, snapshotTime);
 
         // TODO: fix Playwright bug where frame.name is lost (empty).
         const snapshotFrame = uiPage.frames()[1];
         try {
-          await snapshotFrame.goto(snapshotObject.frames[0].url);
+          await snapshotFrame.goto(pageUrl);
         } catch (e) {
           if (!e.message.includes('frame was detached'))
             console.error(e);

src/cli/traceViewer/web/index.tsx

@@ -25,7 +25,7 @@ declare global {
   interface Window {
     getTraceModel(): Promise<TraceModel>;
     readFile(filePath: string): Promise<string>;
-    renderSnapshot(action: trace.ActionTraceEvent): void;
+    renderSnapshot(arg: { action: trace.ActionTraceEvent, snapshot: { name: string, snapshotId?: string } }): void;
   }
 }
 

src/cli/traceViewer/web/ui/propertiesTabbedPane.css

@@ -68,6 +68,28 @@
   font-weight: 600;
 }
 
+.snapshot-tab {
+  display: flex;
+  flex-direction: column;
+  align-items: stretch;
+}
+
+.snapshot-controls {
+  flex: 0 0 24px;
+  display: flex;
+  flex-direction: row;
+  align-items: center;
+}
+
+.snapshot-toggle {
+  padding: 5px 10px;
+  cursor: pointer;
+}
+
+.snapshot-toggle.toggled {
+  background: var(--inactive-focus-ring);
+}
+
 .snapshot-wrapper {
   flex: auto;
   margin: 1px;

src/cli/traceViewer/web/ui/propertiesTabbedPane.tsx

@@ -72,6 +72,16 @@ const SnapshotTab: React.FunctionComponent<{
 }> = ({ actionEntry, snapshotSize }) => {
   const [measure, ref] = useMeasure<HTMLDivElement>();
 
+  let snapshots: { name: string, snapshotId?: string }[] = [];
+
+  snapshots = (actionEntry ? (actionEntry.action.snapshots || []) : []).slice();
+  if (!snapshots.length || snapshots[0].name !== 'before')
+    snapshots.unshift({ name: 'before', snapshotId: undefined });
+  if (snapshots[snapshots.length - 1].name !== 'after')
+    snapshots.push({ name: 'after', snapshotId: undefined });
+
+  const [snapshotIndex, setSnapshotIndex] = React.useState(0);
+
   const iframeRef = React.createRef<HTMLIFrameElement>();
   React.useEffect(() => {
     if (iframeRef.current && !actionEntry)
@@ -80,17 +90,29 @@ const SnapshotTab: React.FunctionComponent<{
 
   React.useEffect(() => {
     if (actionEntry)
-      (window as any).renderSnapshot(actionEntry.action);
-  }, [actionEntry]);
+      (window as any).renderSnapshot({ action: actionEntry.action, snapshot: snapshots[snapshotIndex] });
+  }, [actionEntry, snapshotIndex]);
 
   const scale = Math.min(measure.width / snapshotSize.width, measure.height / snapshotSize.height);
-  return <div ref={ref} className='snapshot-wrapper'>
-    <div className='snapshot-container' style={{
-      width: snapshotSize.width + 'px',
-      height: snapshotSize.height + 'px',
-      transform: `translate(${-snapshotSize.width * (1 - scale) / 2}px, ${-snapshotSize.height * (1 - scale) / 2}px) scale(${scale})`,
-    }}>
-      <iframe ref={iframeRef} id='snapshot' name='snapshot'></iframe>
+  return <div className='snapshot-tab'>
+    <div className='snapshot-controls'>{
+      snapshots.map((snapshot, index) => {
+        return <div
+          key={snapshot.name}
+          className={'snapshot-toggle' + (snapshotIndex === index ? ' toggled' : '')}
+          onClick={() => setSnapshotIndex(index)}>
+          {snapshot.name}
+        </div>
+      })
+    }</div>
+    <div ref={ref} className='snapshot-wrapper'>
+      <div className='snapshot-container' style={{
+        width: snapshotSize.width + 'px',
+        height: snapshotSize.height + 'px',
+        transform: `translate(${-snapshotSize.width * (1 - scale) / 2}px, ${-snapshotSize.height * (1 - scale) / 2}px) scale(${scale})`,
+      }}>
+        <iframe ref={iframeRef} id='snapshot' name='snapshot'></iframe>
+      </div>
     </div>
   </div>;
 };

src/server/browserContext.ts

@@ -67,14 +67,22 @@ export type ActionMetadata = {
 };
 
 export interface ActionListener {
+  onActionCheckpoint(name: string, metadata: ActionMetadata): Promise<void>;
   onAfterAction(result: ProgressResult, metadata: ActionMetadata): Promise<void>;
 }
 
 export async function runAction<T>(task: (controller: ProgressController) => Promise<T>, metadata: ActionMetadata): Promise<T> {
   const controller = new ProgressController();
-  controller.setListener(async result => {
-    for (const listener of metadata.page._browserContext._actionListeners)
-      await listener.onAfterAction(result, metadata);
+  controller.setListener({
+    onProgressCheckpoint: async (name: string): Promise<void> => {
+      for (const listener of metadata.page._browserContext._actionListeners)
+        await listener.onActionCheckpoint(name, metadata);
+    },
+
+    onProgressDone: async (result: ProgressResult): Promise<void> => {
+      for (const listener of metadata.page._browserContext._actionListeners)
+        await listener.onAfterAction(result, metadata);
+    },
   });
   const result = await task(controller);
   return result;

src/server/dom.ts

@@ -378,6 +378,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
       if (options && options.modifiers)
         restoreModifiers = await this._page.keyboard._ensureModifiers(options.modifiers);
       progress.log(`  performing ${actionName} action`);
+      await progress.checkpoint('before');
       await action(point);
       progress.log(`  ${actionName} action done`);
       progress.log('  waiting for scheduled navigations to finish');
@@ -447,6 +448,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
     return this._page._frameManager.waitForSignalsCreatedBy(progress, options.noWaitAfter, async () => {
       progress.throwIfAborted();  // Avoid action that has side-effects.
       progress.log('  selecting specified option(s)');
+      await progress.checkpoint('before');
       const poll = await this._evaluateHandleInUtility(([injected, node, selectOptions]) => injected.waitForOptionsAndSelect(node, selectOptions), selectOptions);
       const pollHandler = new InjectedScriptPollHandler(progress, poll);
       const result = throwFatalDOMError(await pollHandler.finish());
@@ -475,6 +477,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
       if (filled === 'error:notconnected')
         return filled;
       progress.log('  element is visible, enabled and editable');
+      await progress.checkpoint('before');
       if (filled === 'needsinput') {
         progress.throwIfAborted();  // Avoid action that has side-effects.
         if (value)
@@ -521,6 +524,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
     assert(multiple || files.length <= 1, 'Non-multiple file input can only accept single file!');
     await this._page._frameManager.waitForSignalsCreatedBy(progress, options.noWaitAfter, async () => {
       progress.throwIfAborted();  // Avoid action that has side-effects.
+      await progress.checkpoint('before');
       await this._page._delegate.setInputFiles(this as any as ElementHandle<HTMLInputElement>, files);
     });
     await this._page._doSlowMo();
@@ -555,6 +559,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
       if (result !== 'done')
         return result;
       progress.throwIfAborted();  // Avoid action that has side-effects.
+      await progress.checkpoint('before');
       await this._page.keyboard.type(text, options);
       return 'done';
     }, 'input');
@@ -574,6 +579,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
       if (result !== 'done')
         return result;
       progress.throwIfAborted();  // Avoid action that has side-effects.
+      await progress.checkpoint('before');
       await this._page.keyboard.press(key, options);
       return 'done';
     }, 'input');

src/server/frames.ts

@@ -131,8 +131,11 @@ export class FrameManager {
     if (progress)
       progress.cleanupWhenAborted(() => this._signalBarriers.delete(barrier));
     const result = await action();
-    if (source === 'input')
+    if (source === 'input') {
       await this._page._delegate.inputActionEpilogue();
+      if (progress)
+        await progress.checkpoint('after');
+    }
     await barrier.waitFor();
     this._signalBarriers.delete(barrier);
     // Resolve in the next task, after all waitForNavigations.

src/server/progress.ts

@@ -33,6 +33,12 @@ export interface Progress {
   isRunning(): boolean;
   cleanupWhenAborted(cleanup: () => any): void;
   throwIfAborted(): void;
+  checkpoint(name: string): Promise<void>;
+}
+
+export interface ProgressListener {
+  onProgressCheckpoint(name: string): Promise<void>;
+  onProgressDone(result: ProgressResult): Promise<void>;
 }
 
 export async function runAbortableTask<T>(task: (progress: Progress) => Promise<T>, timeout: number): Promise<T> {
@@ -59,7 +65,7 @@ export class ProgressController {
   private _deadline: number = 0;
   private _timeout: number = 0;
   private _logRecording: string[] = [];
-  private _listener?: (result: ProgressResult) => Promise<void>;
+  private _listener?: ProgressListener;
 
   constructor() {
     this._forceAbortPromise = new Promise((resolve, reject) => this._forceAbort = reject);
@@ -71,7 +77,7 @@ export class ProgressController {
     this._logName = logName;
   }
 
-  setListener(listener: (result: ProgressResult) => Promise<void>) {
+  setListener(listener: ProgressListener) {
     this._listener = listener;
   }
 
@@ -103,6 +109,10 @@ export class ProgressController {
         if (this._state === 'aborted')
           throw new AbortedError();
       },
+      checkpoint: async (name: string) => {
+        if (this._listener)
+          await this._listener.onProgressCheckpoint(name);
+      },
     };
 
     const timeoutError = new TimeoutError(`Timeout ${this._timeout}ms exceeded.`);
@@ -114,7 +124,7 @@ export class ProgressController {
       clearTimeout(timer);
       this._state = 'finished';
       if (this._listener) {
-        await this._listener({
+        await this._listener.onProgressDone({
           startTime,
           endTime: monotonicTime(),
           logs: this._logRecording,
@@ -128,7 +138,7 @@ export class ProgressController {
       this._state = 'aborted';
       await Promise.all(this._cleanups.splice(0).map(cleanup => runCleanup(cleanup)));
       if (this._listener) {
-        await this._listener({
+        await this._listener.onProgressDone({
           startTime,
           endTime: monotonicTime(),
           logs: this._logRecording,

src/trace/snapshotter.ts

@@ -18,16 +18,11 @@ import { BrowserContext } from '../server/browserContext';
 import { Page } from '../server/page';
 import * as network from '../server/network';
 import { helper, RegisteredListener } from '../server/helper';
-import { stripFragmentFromUrl } from '../server/network';
-import { Progress, runAbortableTask } from '../server/progress';
 import { debugLogger } from '../utils/debugLogger';
 import { Frame } from '../server/frames';
-import * as js from '../server/javascript';
-import * as types from '../server/types';
-import { SnapshotData, takeSnapshotInFrame } from './snapshotterInjected';
-import { assert, calculateSha1, createGuid } from '../utils/utils';
-import { ElementHandle } from '../server/dom';
-import { FrameSnapshot, PageSnapshot } from './traceTypes';
+import { SnapshotData, frameSnapshotStreamer, kSnapshotBinding, kSnapshotStreamer } from './snapshotterInjected';
+import { calculateSha1 } from '../utils/utils';
+import { FrameSnapshot } from './traceTypes';
 
 export type SnapshotterResource = {
   pageId: string,
@@ -46,6 +41,7 @@ export type SnapshotterBlob = {
 export interface SnapshotterDelegate {
   onBlob(blob: SnapshotterBlob): void;
   onResource(resource: SnapshotterResource): void;
+  onFrameSnapshot(frame: Frame, snapshot: FrameSnapshot, snapshotId?: string): void;
   pageId(page: Page): string;
 }
 
@@ -60,16 +56,63 @@ export class Snapshotter {
     this._eventListeners = [
       helper.addEventListener(this._context, BrowserContext.Events.Page, this._onPage.bind(this)),
     ];
+    this._context.exposeBinding(kSnapshotBinding, false, (source, data: SnapshotData) => {
+      const snapshot: FrameSnapshot = {
+        html: data.html,
+        viewport: data.viewport,
+        resourceOverrides: [],
+        url: data.url,
+      };
+      for (const { url, content } of data.resourceOverrides) {
+        const buffer = Buffer.from(content);
+        const sha1 = calculateSha1(buffer);
+        this._delegate.onBlob({ sha1, buffer });
+        snapshot.resourceOverrides.push({ url, sha1 });
+      }
+      this._delegate.onFrameSnapshot(source.frame, snapshot, data.snapshotId);
+    });
+    this._context._doAddInitScript('(' + frameSnapshotStreamer.toString() + ')()');
   }
 
   dispose() {
     helper.removeEventListeners(this._eventListeners);
   }
 
+  async forceSnapshot(page: Page, snapshotId: string) {
+    await Promise.all([
+      page.frames().forEach(async frame => {
+        try {
+          const context = await frame._mainContext();
+          await context.evaluateInternal(({ kSnapshotStreamer, snapshotId }) => {
+            // Do not block action execution on the actual snapshot.
+            Promise.resolve().then(() => (window as any)[kSnapshotStreamer].forceSnapshot(snapshotId));
+            return undefined;
+          }, { kSnapshotStreamer, snapshotId });
+        } catch (e) {
+        }
+      })
+    ]);
+  }
+
   private _onPage(page: Page) {
     this._eventListeners.push(helper.addEventListener(page, Page.Events.Response, (response: network.Response) => {
       this._saveResource(page, response).catch(e => debugLogger.log('error', e));
     }));
+    this._eventListeners.push(helper.addEventListener(page, Page.Events.FrameAttached, async (frame: Frame) => {
+      try {
+        const frameElement = await frame.frameElement();
+        const parent = frame.parentFrame();
+        if (!parent)
+          return;
+        const context = await parent._mainContext();
+        await context.evaluateInternal(({ kSnapshotStreamer, frameElement, frameId }) => {
+          (window as any)[kSnapshotStreamer].markIframe(frameElement, frameId);
+        }, { kSnapshotStreamer, frameElement, frameId: frame._id });
+        frameElement.dispose();
+      } catch (e) {
+        // Ignore
+      }
+    }));
   }
 
   private async _saveResource(page: Page, response: network.Response) {
@@ -103,121 +146,4 @@ export class Snapshotter {
     if (body)
       this._delegate.onBlob({ sha1, buffer: body });
   }
-
-  async takeSnapshot(page: Page, target: ElementHandle | undefined, timeout: number): Promise<PageSnapshot | null> {
-    assert(page.context() === this._context);
-
-    const frames = page.frames();
-    const frameSnapshotPromises = frames.map(async frame => {
-      // TODO: use different timeout depending on the frame depth/origin
-      // to avoid waiting for too long for some useless frame.
-      const frameResult = await runAbortableTask(progress => this._snapshotFrame(progress, target, frame), timeout).catch(e => null);
-      if (frameResult)
-        return frameResult;
-      const frameSnapshot = {
-        frameId: frame._id,
-        url: stripFragmentFromUrl(frame.url()),
-        html: '<body>Snapshot is not available</body>',
-        resourceOverrides: [],
-      };
-      return { snapshot: frameSnapshot, mapping: new Map<Frame, string>() };
-    });
-
-    const viewportSize = await this._getViewportSize(page, timeout);
-    const results = await Promise.all(frameSnapshotPromises);
-
-    if (!viewportSize)
-      return null;
-
-    const mainFrame = results[0];
-    if (!mainFrame.snapshot.url.startsWith('http'))
-      mainFrame.snapshot.url = 'http://playwright.snapshot/';
-
-    const mapping = new Map<Frame, string>();
-    for (const result of results) {
-      for (const [key, value] of result.mapping)
-        mapping.set(key, value);
-    }
-
-    const childFrames: FrameSnapshot[] = [];
-    for (let i = 1; i < results.length; i++) {
-      const result = results[i];
-      const frame = frames[i];
-      if (!mapping.has(frame))
-        continue;
-      const frameSnapshot = result.snapshot;
-      frameSnapshot.url = mapping.get(frame)!;
-      childFrames.push(frameSnapshot);
-    }
-
-    return {
-      viewportSize,
-      frames: [mainFrame.snapshot, ...childFrames],
-    };
-  }
-
-  private async _getViewportSize(page: Page, timeout: number): Promise<types.Size | null> {
-    return runAbortableTask(async progress => {
-      const viewportSize = page.viewportSize();
-      if (viewportSize)
-        return viewportSize;
-      const context = await page.mainFrame()._utilityContext();
-      return context.evaluateInternal(() => {
-        return {
-          width: Math.max(document.body.offsetWidth, document.documentElement.offsetWidth),
-          height: Math.max(document.body.offsetHeight, document.documentElement.offsetHeight),
-        };
-      });
-    }, timeout).catch(e => null);
-  }
-
-  private async _snapshotFrame(progress: Progress, target: ElementHandle | undefined, frame: Frame): Promise<FrameSnapshotAndMapping | null> {
-    if (!progress.isRunning())
-      return null;
-
-    if (target && (await target.ownerFrame()) !== frame)
-      target = undefined;
-    const context = await frame._utilityContext();
-    const guid = createGuid();
-    const removeNoScript = !frame._page.context()._options.javaScriptEnabled;
-    const result = await js.evaluate(context, false /* returnByValue */, takeSnapshotInFrame, guid, removeNoScript, target) as js.JSHandle;
-    if (!progress.isRunning())
-      return null;
-
-    const properties = await result.getProperties();
-    const data = await properties.get('data')!.jsonValue() as SnapshotData;
-    const frameElements = await properties.get('frameElements')!.getProperties();
-    result.dispose();
-
-    const snapshot: FrameSnapshot = {
-      frameId: frame._id,
-      url: stripFragmentFromUrl(frame.url()),
-      html: data.html,
-      resourceOverrides: [],
-    };
-    const mapping = new Map<Frame, string>();
-
-    for (const { url, content } of data.resourceOverrides) {
-      const buffer = Buffer.from(content);
-      const sha1 = calculateSha1(buffer);
-      this._delegate.onBlob({ sha1, buffer });
-      snapshot.resourceOverrides.push({ url, sha1 });
-    }
-
-    for (let i = 0; i < data.frameUrls.length; i++) {
-      const element = frameElements.get(String(i))!.asElement();
-      if (!element)
-        continue;
-      const frame = await element.contentFrame().catch(e => null);
-      if (frame)
-        mapping.set(frame, data.frameUrls[i]);
-    }
-
-    return { snapshot, mapping };
-  }
 }
-
-type FrameSnapshotAndMapping = {
-  snapshot: FrameSnapshot,
-  mapping: Map<Frame, string>,
-};

src/trace/snapshotterInjected.ts

@@ -17,276 +17,334 @@
 export type SnapshotData = {
   html: string,
   resourceOverrides: { url: string, content: string }[],
-  frameUrls: string[],
+  viewport: { width: number, height: number },
+  url: string,
+  snapshotId?: string,
 };
 
-type SnapshotResult = {
-  data: SnapshotData,
-  frameElements: Element[],
-};
+export const kSnapshotStreamer = '__playwright_snapshot_streamer_';
+export const kSnapshotFrameIdAttribute = '__playwright_snapshot_frameid_';
+export const kSnapshotBinding = '__playwright_snapshot_binding_';
 
-export function takeSnapshotInFrame(guid: string, removeNoScript: boolean, target: Node | undefined): SnapshotResult {
-  const shadowAttribute = 'playwright-shadow-root';
-  const win = window;
-  const doc = win.document;
+export function frameSnapshotStreamer() {
+  const kSnapshotStreamer = '__playwright_snapshot_streamer_';
+  const kSnapshotFrameIdAttribute = '__playwright_snapshot_frameid_';
+  const kSnapshotBinding = '__playwright_snapshot_binding_';
+  const kShadowAttribute = '__playwright_shadow_root_';
 
-  const autoClosing = new Set(['AREA', 'BASE', 'BR', 'COL', 'COMMAND', 'EMBED', 'HR', 'IMG', 'INPUT', 'KEYGEN', 'LINK', 'MENUITEM', 'META', 'PARAM', 'SOURCE', 'TRACK', 'WBR']);
   const escaped = { '&': '&amp;', '<': '&lt;', '>': '&gt;', '"': '&quot;', '\'': '&#39;' };
+  const autoClosing = new Set(['AREA', 'BASE', 'BR', 'COL', 'COMMAND', 'EMBED', 'HR', 'IMG', 'INPUT', 'KEYGEN', 'LINK', 'MENUITEM', 'META', 'PARAM', 'SOURCE', 'TRACK', 'WBR']);
 
-  const escapeAttribute = (s: string): string => {
-    return s.replace(/[&<>"']/ug, char => (escaped as any)[char]);
-  };
-  const escapeText = (s: string): string => {
-    return s.replace(/[&<]/ug, char => (escaped as any)[char]);
-  };
-  const escapeScriptString = (s: string): string => {
-    return s.replace(/'/g, '\\\'');
-  };
+  class Streamer {
+    private _removeNoScript = true;
+    private _needStyleOverrides = false;
+    private _timer: NodeJS.Timeout | undefined;
 
-  const chunks = new Map<string, string>();
-  const frameUrlToFrameElement = new Map<string, Element>();
-  const styleNodeToStyleSheetText = new Map<Node, string>();
-  const styleSheetUrlToContentOverride = new Map<string, string>();
+    constructor() {
+      this._streamSnapshot();
+      this._interceptCSSOM(window.CSSStyleSheet.prototype, 'insertRule');
+      this._interceptCSSOM(window.CSSStyleSheet.prototype, 'deleteRule');
+      this._interceptCSSOM(window.CSSStyleSheet.prototype, 'addRule');
+      this._interceptCSSOM(window.CSSStyleSheet.prototype, 'removeRule');
+      // TODO: should we also intercept setters like CSSRule.cssText and CSSStyleRule.selectorText?
+    }
 
-  let counter = 0;
-  const nextId = (): string => {
-    return guid + (++counter);
-  };
+    private _interceptCSSOM(obj: any, method: string) {
+      const self = this;
+      const native = obj[method] as Function;
+      if (!native)
+        return;
+      obj[method] = function(...args: any[]) {
+        self._needStyleOverrides = true;
+        native.call(this, ...args);
+      };
+    }
 
-  const resolve = (base: string, url: string): string => {
-    if (url === '')
-      return '';
-    try {
-      return new URL(url, base).href;
-    } catch (e) {
+    markIframe(iframeElement: HTMLIFrameElement | HTMLFrameElement, frameId: string) {
+      iframeElement.setAttribute(kSnapshotFrameIdAttribute, frameId);
+    }
+
+    forceSnapshot(snapshotId: string) {
+      this._streamSnapshot(snapshotId);
+    }
+
+    private _streamSnapshot(snapshotId?: string) {
+      if (this._timer) {
+        clearTimeout(this._timer);
+        this._timer = undefined;
+      }
+      const snapshot = this._captureSnapshot(snapshotId);
+      (window as any)[kSnapshotBinding](snapshot).catch((e: any) => {});
+      this._timer = setTimeout(() => this._streamSnapshot(), 100);
+    }
+
+    private _escapeAttribute(s: string): string {
+      return s.replace(/[&<>"']/ug, char => (escaped as any)[char]);
+    }
+
+    private _escapeText(s: string): string {
+      return s.replace(/[&<]/ug, char => (escaped as any)[char]);
+    }
+
+    private _sanitizeUrl(url: string): string {
+      if (url.startsWith('javascript:'))
+        return '';
       return url;
     }
-  };
 
-  const sanitizeUrl = (url: string): string => {
-    if (url.startsWith('javascript:'))
-      return '';
-    return url;
-  };
+    private _sanitizeSrcSet(srcset: string): string {
+      return srcset.split(',').map(src => {
+        src = src.trim();
+        const spaceIndex = src.lastIndexOf(' ');
+        if (spaceIndex === -1)
+          return this._sanitizeUrl(src);
+        return this._sanitizeUrl(src.substring(0, spaceIndex).trim()) + src.substring(spaceIndex);
+      }).join(',');
+    }
 
-  const sanitizeSrcSet = (srcset: string): string => {
-    return srcset.split(',').map(src => {
-      src = src.trim();
-      const spaceIndex = src.lastIndexOf(' ');
-      if (spaceIndex === -1)
-        return sanitizeUrl(src);
-      return sanitizeUrl(src.substring(0, spaceIndex).trim()) + src.substring(spaceIndex);
-    }).join(',');
-  };
-
-  const getSheetBase = (sheet: CSSStyleSheet): string => {
-    let rootSheet = sheet;
-    while (rootSheet.parentStyleSheet)
-      rootSheet = rootSheet.parentStyleSheet;
-    if (rootSheet.ownerNode)
-      return rootSheet.ownerNode.baseURI;
-    return document.baseURI;
-  };
-
-  const getSheetText = (sheet: CSSStyleSheet): string => {
-    const rules: string[] = [];
-    for (const rule of sheet.cssRules)
-      rules.push(rule.cssText);
-    return rules.join('\n');
-  };
-
-  const visitStyleSheet = (sheet: CSSStyleSheet) => {
-    try {
-      for (const rule of sheet.cssRules) {
-        if ((rule as CSSImportRule).styleSheet)
-          visitStyleSheet((rule as CSSImportRule).styleSheet);
+    private _resolveUrl(base: string, url: string): string {
+      if (url === '')
+        return '';
+      try {
+        return new URL(url, base).href;
+      } catch (e) {
+        return url;
       }
-
-      const cssText = getSheetText(sheet);
-      if (sheet.ownerNode && sheet.ownerNode.nodeName === 'STYLE') {
-        // Stylesheets with owner STYLE nodes will be rewritten.
-        styleNodeToStyleSheetText.set(sheet.ownerNode, cssText);
-      } else if (sheet.href !== null) {
-        // Other stylesheets will have resource overrides.
-        const base = getSheetBase(sheet);
-        const url = resolve(base, sheet.href);
-        styleSheetUrlToContentOverride.set(url, cssText);
-      }
-    } catch (e) {
-      // Sometimes we cannot access cross-origin stylesheets.
-    }
-  };
-
-  const visit = (node: Node | ShadowRoot, builder: string[]) => {
-    const nodeName = node.nodeName;
-    const nodeType = node.nodeType;
-
-    if (nodeType === Node.DOCUMENT_TYPE_NODE) {
-      const docType = node as DocumentType;
-      builder.push(`<!DOCTYPE ${docType.name}>`);
-      return;
     }
 
-    if (nodeType === Node.TEXT_NODE) {
-      builder.push(escapeText(node.nodeValue || ''));
-      return;
+    private _getSheetBase(sheet: CSSStyleSheet): string {
+      let rootSheet = sheet;
+      while (rootSheet.parentStyleSheet)
+        rootSheet = rootSheet.parentStyleSheet;
+      if (rootSheet.ownerNode)
+        return rootSheet.ownerNode.baseURI;
+      return document.baseURI;
     }
 
-    if (nodeType !== Node.ELEMENT_NODE &&
-        nodeType !== Node.DOCUMENT_NODE &&
-        nodeType !== Node.DOCUMENT_FRAGMENT_NODE)
-      return;
-
-    if (nodeType === Node.DOCUMENT_NODE || nodeType === Node.DOCUMENT_FRAGMENT_NODE) {
-      const documentOrShadowRoot = node as DocumentOrShadowRoot;
-      for (const sheet of documentOrShadowRoot.styleSheets)
-        visitStyleSheet(sheet);
+    private _getSheetText(sheet: CSSStyleSheet): string {
+      const rules: string[] = [];
+      for (const rule of sheet.cssRules)
+        rules.push(rule.cssText);
+      return rules.join('\n');
     }
 
-    if (nodeName === 'SCRIPT' || nodeName === 'BASE')
-      return;
+    private _captureSnapshot(snapshotId?: string): SnapshotData {
+      const win = window;
+      const doc = win.document;
 
-    if (removeNoScript && nodeName === 'NOSCRIPT')
-      return;
+      const shadowChunks: string[] = [];
+      const styleNodeToStyleSheetText = new Map<Node, string>();
+      const styleSheetUrlToContentOverride = new Map<string, string>();
 
-    if (nodeName === 'STYLE') {
-      const cssText = styleNodeToStyleSheetText.get(node) || node.textContent || '';
-      builder.push('<style>');
-      builder.push(cssText);
-      builder.push('</style>');
-      return;
-    }
+      const visitStyleSheet = (sheet: CSSStyleSheet) => {
+        // TODO: recalculate these upon changes, and only send them once.
+        if (!this._needStyleOverrides)
+          return;
 
-    if (nodeType === Node.ELEMENT_NODE) {
-      const element = node as Element;
-      builder.push('<');
-      builder.push(nodeName);
-      if (node === target)
-        builder.push(' __playwright_target__="true"');
-      for (let i = 0; i < element.attributes.length; i++) {
-        const name = element.attributes[i].name;
-        let value = element.attributes[i].value;
-        if (name === 'value' && (nodeName === 'INPUT' || nodeName === 'TEXTAREA'))
-          continue;
-        if (name === 'checked' || name === 'disabled' || name === 'checked')
-          continue;
-        if (nodeName === 'LINK' && name === 'integrity')
-          continue;
-        if (name === 'src' && (nodeName === 'IFRAME' || nodeName === 'FRAME')) {
-          // TODO: handle srcdoc?
-          let protocol = win.location.protocol;
-          if (!protocol.startsWith('http'))
-            protocol = 'http:';
-          value = protocol + '//' + nextId() + '/';
-          frameUrlToFrameElement.set(value, element);
-        } else if (name === 'src' && (nodeName === 'IMG')) {
-          value = sanitizeUrl(value);
-        } else if (name === 'srcset' && (nodeName === 'IMG')) {
-          value = sanitizeSrcSet(value);
-        } else if (name === 'srcset' && (nodeName === 'SOURCE')) {
-          value = sanitizeSrcSet(value);
-        } else if (name === 'href' && (nodeName === 'LINK')) {
-          value = sanitizeUrl(value);
-        } else if (name.startsWith('on')) {
-          value = '';
+        try {
+          for (const rule of sheet.cssRules) {
+            if ((rule as CSSImportRule).styleSheet)
+              visitStyleSheet((rule as CSSImportRule).styleSheet);
+          }
+
+          const cssText = this._getSheetText(sheet);
+          if (sheet.ownerNode && sheet.ownerNode.nodeName === 'STYLE') {
+            // Stylesheets with owner STYLE nodes will be rewritten.
+            styleNodeToStyleSheetText.set(sheet.ownerNode, cssText);
+          } else if (sheet.href !== null) {
+            // Other stylesheets will have resource overrides.
+            const base = this._getSheetBase(sheet);
+            const url = this._resolveUrl(base, sheet.href);
+            styleSheetUrlToContentOverride.set(url, cssText);
+          }
+        } catch (e) {
+          // Sometimes we cannot access cross-origin stylesheets.
         }
-        builder.push(' ');
-        builder.push(name);
-        builder.push('="');
-        builder.push(escapeAttribute(value));
-        builder.push('"');
-      }
-      if (nodeName === 'INPUT') {
-        builder.push(' value="');
-        builder.push(escapeAttribute((element as HTMLInputElement).value));
-        builder.push('"');
-      }
-      if ((element as any).checked)
-        builder.push(' checked');
-      if ((element as any).disabled)
-        builder.push(' disabled');
-      if ((element as any).readOnly)
-        builder.push(' readonly');
-      if (element.shadowRoot) {
-        const b: string[] = [];
-        visit(element.shadowRoot, b);
-        const chunkId = nextId();
-        chunks.set(chunkId, b.join(''));
-        builder.push(' ');
-        builder.push(shadowAttribute);
-        builder.push('="');
-        builder.push(chunkId);
-        builder.push('"');
-      }
-      builder.push('>');
-    }
-    if (nodeName === 'HEAD') {
-      let baseHref = document.baseURI;
-      let baseTarget: string | undefined;
-      for (let child = node.firstChild; child; child = child.nextSibling) {
-        if (child.nodeName === 'BASE') {
-          baseHref = (child as HTMLBaseElement).href;
-          baseTarget = (child as HTMLBaseElement).target;
-        }
-      }
-      builder.push('<base href="');
-      builder.push(escapeAttribute(baseHref));
-      builder.push('"');
-      if (baseTarget) {
-        builder.push(' target="');
-        builder.push(escapeAttribute(baseTarget));
-        builder.push('"');
-      }
-      builder.push('>');
-    }
-    if (nodeName === 'TEXTAREA') {
-      builder.push(escapeText((node as HTMLTextAreaElement).value));
-    } else {
-      for (let child = node.firstChild; child; child = child.nextSibling)
-        visit(child, builder);
-    }
-    if (node.nodeName === 'BODY' && chunks.size) {
-      builder.push('<script>');
-      const shadowChunks = Array.from(chunks).map(([chunkId, html]) => {
-        return `  ['${chunkId}', '${escapeScriptString(html)}']`;
-      }).join(',\n');
-      const scriptContent = `\n(${applyShadowsInPage.toString()})('${shadowAttribute}', new Map([\n${shadowChunks}\n]))\n`;
-      builder.push(scriptContent);
-      builder.push('</script>');
-    }
-    if (nodeType === Node.ELEMENT_NODE && !autoClosing.has(nodeName)) {
-      builder.push('</');
-      builder.push(nodeName);
-      builder.push('>');
-    }
-  };
+      };
 
-  function applyShadowsInPage(shadowAttribute: string, shadowContent: Map<string, string>) {
-    const visitShadows = (root: Document | ShadowRoot) => {
-      const elements = root.querySelectorAll(`[${shadowAttribute}]`);
-      for (let i = 0; i < elements.length; i++) {
-        const host = elements[i];
-        const chunkId = host.getAttribute(shadowAttribute)!;
-        host.removeAttribute(shadowAttribute);
-        const shadow = host.attachShadow({ mode: 'open' });
-        const html = shadowContent.get(chunkId);
-        if (html) {
-          shadow.innerHTML = html;
-          visitShadows(shadow);
+      const visit = (node: Node | ShadowRoot, builder: string[]) => {
+        const nodeName = node.nodeName;
+        const nodeType = node.nodeType;
+
+        if (nodeType === Node.DOCUMENT_TYPE_NODE) {
+          const docType = node as DocumentType;
+          builder.push(`<!DOCTYPE ${docType.name}>`);
+          return;
         }
+
+        if (nodeType === Node.TEXT_NODE) {
+          builder.push(this._escapeText(node.nodeValue || ''));
+          return;
+        }
+
+        if (nodeType !== Node.ELEMENT_NODE &&
+            nodeType !== Node.DOCUMENT_NODE &&
+            nodeType !== Node.DOCUMENT_FRAGMENT_NODE)
+          return;
+
+        if (nodeType === Node.DOCUMENT_NODE || nodeType === Node.DOCUMENT_FRAGMENT_NODE) {
+          const documentOrShadowRoot = node as DocumentOrShadowRoot;
+          for (const sheet of documentOrShadowRoot.styleSheets)
+            visitStyleSheet(sheet);
+        }
+
+        if (nodeName === 'SCRIPT' || nodeName === 'BASE')
+          return;
+
+        if (this._removeNoScript && nodeName === 'NOSCRIPT')
+          return;
+
+        if (nodeName === 'STYLE') {
+          const cssText = styleNodeToStyleSheetText.get(node) || node.textContent || '';
+          builder.push('<style>');
+          builder.push(cssText);
+          builder.push('</style>');
+          return;
+        }
+
+        if (nodeType === Node.ELEMENT_NODE) {
+          const element = node as Element;
+          builder.push('<');
+          builder.push(nodeName);
+          // if (node === target)
+          //   builder.push(' __playwright_target__="true"');
+          for (let i = 0; i < element.attributes.length; i++) {
+            const name = element.attributes[i].name;
+            if (name === kSnapshotFrameIdAttribute)
+              continue;
+
+            let value = element.attributes[i].value;
+            if (name === 'value' && (nodeName === 'INPUT' || nodeName === 'TEXTAREA'))
+              continue;
+            if (name === 'checked' || name === 'disabled' || name === 'checked')
+              continue;
+            if (nodeName === 'LINK' && name === 'integrity')
+              continue;
+            if (name === 'src' && (nodeName === 'IFRAME' || nodeName === 'FRAME')) {
+              // TODO: handle srcdoc?
+              const frameId = element.getAttribute(kSnapshotFrameIdAttribute);
+              if (frameId) {
+                let protocol = win.location.protocol;
+                if (!protocol.startsWith('http'))
+                  protocol = 'http:';
+                value = protocol + '//' + frameId + '/';
+              } else {
+                value = 'data:text/html,<body>Snapshot is not available</body>';
+              }
+            } else if (name === 'src' && (nodeName === 'IMG')) {
+              value = this._sanitizeUrl(value);
+            } else if (name === 'srcset' && (nodeName === 'IMG')) {
+              value = this._sanitizeSrcSet(value);
+            } else if (name === 'srcset' && (nodeName === 'SOURCE')) {
+              value = this._sanitizeSrcSet(value);
+            } else if (name === 'href' && (nodeName === 'LINK')) {
+              value = this._sanitizeUrl(value);
+            } else if (name.startsWith('on')) {
+              value = '';
+            }
+            builder.push(' ');
+            builder.push(name);
+            builder.push('="');
+            builder.push(this._escapeAttribute(value));
+            builder.push('"');
+          }
+          if (nodeName === 'INPUT') {
+            builder.push(' value="');
+            builder.push(this._escapeAttribute((element as HTMLInputElement).value));
+            builder.push('"');
+          }
+          if ((element as any).checked)
+            builder.push(' checked');
+          if ((element as any).disabled)
+            builder.push(' disabled');
+          if ((element as any).readOnly)
+            builder.push(' readonly');
+          if (element.shadowRoot) {
+            const b: string[] = [];
+            visit(element.shadowRoot, b);
+            const chunkId = shadowChunks.length;
+            shadowChunks.push(b.join(''));
+            builder.push(' ');
+            builder.push(kShadowAttribute);
+            builder.push('="');
+            builder.push('' + chunkId);
+            builder.push('"');
+          }
+          builder.push('>');
+        }
+        if (nodeName === 'HEAD') {
+          let baseHref = document.baseURI;
+          let baseTarget: string | undefined;
+          for (let child = node.firstChild; child; child = child.nextSibling) {
+            if (child.nodeName === 'BASE') {
+              baseHref = (child as HTMLBaseElement).href;
+              baseTarget = (child as HTMLBaseElement).target;
+            }
+          }
+          builder.push('<base href="');
+          builder.push(this._escapeAttribute(baseHref));
+          builder.push('"');
+          if (baseTarget) {
+            builder.push(' target="');
+            builder.push(this._escapeAttribute(baseTarget));
+            builder.push('"');
+          }
+          builder.push('>');
+        }
+        if (nodeName === 'TEXTAREA') {
+          builder.push(this._escapeText((node as HTMLTextAreaElement).value));
+        } else {
+          for (let child = node.firstChild; child; child = child.nextSibling)
+            visit(child, builder);
+        }
+        if (node.nodeName === 'BODY' && shadowChunks.length) {
+          builder.push('<script>');
+          const chunks = shadowChunks.map(html => {
+            return '`' + html.replace(/`/g, '\\\`') + '`';
+          }).join(',\n');
+          const scriptContent = `\n(${applyShadowsInPage.toString()})('${kShadowAttribute}', [\n${chunks}\n])\n`;
+          builder.push(scriptContent);
+          builder.push('</script>');
+        }
+        if (nodeType === Node.ELEMENT_NODE && !autoClosing.has(nodeName)) {
+          builder.push('</');
+          builder.push(nodeName);
+          builder.push('>');
+        }
+      };
+
+      function applyShadowsInPage(shadowAttribute: string, shadowContent: string[]) {
+        const visitShadows = (root: Document | ShadowRoot) => {
+          const elements = root.querySelectorAll(`[${shadowAttribute}]`);
+          for (let i = 0; i < elements.length; i++) {
+            const host = elements[i];
+            const chunkId = host.getAttribute(shadowAttribute)!;
+            host.removeAttribute(shadowAttribute);
+            const shadow = host.attachShadow({ mode: 'open' });
+            const html = shadowContent[+chunkId];
+            if (html) {
+              shadow.innerHTML = html;
+              visitShadows(shadow);
+            }
+          }
+        };
+        visitShadows(document);
       }
-    };
-    visitShadows(document);
+
+      const root: string[] = [];
+      visit(doc, root);
+      return {
+        html: root.join(''),
+        resourceOverrides: Array.from(styleSheetUrlToContentOverride).map(([url, content]) => ({ url, content })),
+        viewport: {
+          width: Math.max(doc.body ? doc.body.offsetWidth : 0, doc.documentElement ? doc.documentElement.offsetWidth : 0),
+          height: Math.max(doc.body ? doc.body.offsetHeight : 0, doc.documentElement ? doc.documentElement.offsetHeight : 0),
+        },
+        url: location.href,
+        snapshotId,
+      };
+    }
   }
 
-  const root: string[] = [];
-  visit(doc, root);
-  return {
-    data: {
-      html: root.join(''),
-      frameUrls: Array.from(frameUrlToFrameElement.keys()),
-      resourceOverrides: Array.from(styleSheetUrlToContentOverride).map(([url, content]) => ({ url, content })),
-    },
-    frameElements: Array.from(frameUrlToFrameElement.values()),
-  };
+  (window as any)[kSnapshotStreamer] = new Streamer();
 }

src/trace/traceTypes.ts

@@ -76,12 +76,9 @@ export type ActionTraceEvent = {
   startTime: number,
   endTime: number,
   logs?: string[],
-  snapshot?: {
-    sha1: string,
-    duration: number,
-  },
   stack?: string,
   error?: string,
+  snapshots?: { name: string, snapshotId: string }[],
 };
 
 export type DialogOpenedEvent = {
@@ -117,6 +114,17 @@ export type LoadEvent = {
   pageId: string,
 };
 
+export type FrameSnapshotTraceEvent = {
+  timestamp: number,
+  type: 'snapshot',
+  contextId: string,
+  pageId: string,
+  frameId: string,  // Empty means main frame.
+  sha1: string,
+  frameUrl: string,
+  snapshotId?: string,
+};
+
 export type TraceEvent =
     ContextCreatedTraceEvent |
     ContextDestroyedTraceEvent |
@@ -128,18 +136,13 @@ export type TraceEvent =
     DialogOpenedEvent |
     DialogClosedEvent |
     NavigationEvent |
-    LoadEvent;
+    LoadEvent |
+    FrameSnapshotTraceEvent;
 
 
 export type FrameSnapshot = {
-  frameId: string,
-  url: string,
   html: string,
   resourceOverrides: { url: string, sha1: string }[],
-};
-
-export type PageSnapshot = {
-  viewportSize?: { width: number, height: number },
-  // First frame is the main frame.
-  frames: FrameSnapshot[],
+  viewport: { width: number, height: number },
+  url: string,
 };

src/trace/tracer.ts

@@ -23,9 +23,7 @@ import * as fs from 'fs';
 import { calculateSha1, createGuid, getFromENV, mkdirIfNeeded, monotonicTime } from '../utils/utils';
 import { Page } from '../server/page';
 import { Snapshotter } from './snapshotter';
-import { ElementHandle } from '../server/dom';
 import { helper, RegisteredListener } from '../server/helper';
-import { DEFAULT_TIMEOUT } from '../utils/timeoutSettings';
 import { ProgressResult } from '../server/progress';
 import { Dialog } from '../server/dialog';
 import { Frame, NavigationEvent } from '../server/frames';
@@ -64,6 +62,14 @@ class Tracer implements ContextListener {
 }
 
 const pageIdSymbol = Symbol('pageId');
+const snapshotsSymbol = Symbol('snapshots');
+
+// TODO: this is a hacky way to pass snapshots between onActionCheckpoint and onAfterAction.
+function snapshotsForMetadata(metadata: ActionMetadata): { name: string, snapshotId: string }[] {
+  if (!(metadata as any)[snapshotsSymbol])
+    (metadata as any)[snapshotsSymbol] = [];
+  return (metadata as any)[snapshotsSymbol];
+}
 
 class ContextTracer implements SnapshotterDelegate, ActionListener {
   private _context: BrowserContext;
@@ -119,31 +125,50 @@ class ContextTracer implements SnapshotterDelegate, ActionListener {
     this._appendTraceEvent(event);
   }
 
+  onFrameSnapshot(frame: Frame, snapshot: trace.FrameSnapshot, snapshotId?: string): void {
+    const buffer = Buffer.from(JSON.stringify(snapshot));
+    const sha1 = calculateSha1(buffer);
+    this._writeArtifact(sha1, buffer);
+    const event: trace.FrameSnapshotTraceEvent = {
+      timestamp: monotonicTime(),
+      type: 'snapshot',
+      contextId: this._contextId,
+      pageId: this.pageId(frame._page),
+      frameId: frame._page.mainFrame() === frame ? '' : frame._id,
+      sha1,
+      frameUrl: snapshot.url,
+      snapshotId,
+    };
+    this._appendTraceEvent(event);
+  }
+
   pageId(page: Page): string {
     return (page as any)[pageIdSymbol];
   }
 
+  async onActionCheckpoint(name: string, metadata: ActionMetadata): Promise<void> {
+    const snapshotId = createGuid();
+    snapshotsForMetadata(metadata).push({ name, snapshotId });
+    await this._snapshotter.forceSnapshot(metadata.page, snapshotId);
+  }
+
   async onAfterAction(result: ProgressResult, metadata: ActionMetadata): Promise<void> {
-    try {
-      const snapshot = await this._takeSnapshot(metadata.page, typeof metadata.target === 'string' ? undefined : metadata.target);
-      const event: trace.ActionTraceEvent = {
-        timestamp: monotonicTime(),
-        type: 'action',
-        contextId: this._contextId,
-        pageId: this.pageId(metadata.page),
-        action: metadata.type,
-        selector: typeof metadata.target === 'string' ? metadata.target : undefined,
-        value: metadata.value,
-        snapshot,
-        startTime: result.startTime,
-        endTime: result.endTime,
-        stack: metadata.stack,
-        logs: result.logs.slice(),
-        error: result.error ? result.error.stack : undefined,
-      };
-      this._appendTraceEvent(event);
-    } catch (e) {
-    }
+    const event: trace.ActionTraceEvent = {
+      timestamp: monotonicTime(),
+      type: 'action',
+      contextId: this._contextId,
+      pageId: this.pageId(metadata.page),
+      action: metadata.type,
+      selector: typeof metadata.target === 'string' ? metadata.target : undefined,
+      value: metadata.value,
+      startTime: result.startTime,
+      endTime: result.endTime,
+      stack: metadata.stack,
+      logs: result.logs.slice(),
+      error: result.error ? result.error.stack : undefined,
+      snapshots: snapshotsForMetadata(metadata),
+    };
+    this._appendTraceEvent(event);
   }
 
   private _onPage(page: Page) {
@@ -237,22 +262,6 @@ class ContextTracer implements SnapshotterDelegate, ActionListener {
     });
   }
 
-  private async _takeSnapshot(page: Page, target: ElementHandle | undefined, timeout: number = 0): Promise<{ sha1: string, duration: number } | undefined> {
-    if (!timeout) {
-      // Never use zero timeout to avoid stalling because of snapshot.
-      // Use 20% of the default timeout.
-      timeout = (page._timeoutSettings.timeout({}) || DEFAULT_TIMEOUT) / 5;
-    }
-    const startTime = monotonicTime();
-    const snapshot = await this._snapshotter.takeSnapshot(page, target, timeout);
-    if (!snapshot)
-      return;
-    const buffer = Buffer.from(JSON.stringify(snapshot));
-    const sha1 = calculateSha1(buffer);
-    this._writeArtifact(sha1, buffer);
-    return { sha1, duration: monotonicTime() - startTime };
-  }
-
   async dispose() {
     this._disposed = true;
     this._context._actionListeners.delete(this);

test/trace.spec.ts

@@ -25,6 +25,7 @@ it('should record trace', test => test.fixme(), async ({browser, testInfo, serve
   const page = await context.newPage();
   const url = server.PREFIX + '/snapshot/snapshot-with-css.html';
   await page.goto(url);
+  await page.click('textarea');
   await context.close();
   const tracePath = path.join(traceDir, fs.readdirSync(traceDir).find(n => n.endsWith('.trace')));
   const traceFileContent = await fs.promises.readFile(tracePath, 'utf8');
@@ -45,6 +46,11 @@ it('should record trace', test => test.fixme(), async ({browser, testInfo, serve
   expect(gotoEvent.pageId).toBe(pageId);
   expect(gotoEvent.value).toBe(url);
 
-  expect(gotoEvent.snapshot).toBeTruthy();
-  expect(fs.existsSync(path.join(traceDir, 'resources', gotoEvent.snapshot!.sha1))).toBe(true);
+  const clickEvent = traceEvents.find(event => event.type === 'action' && event.action === 'click') as trace.ActionTraceEvent;
+  expect(clickEvent).toBeTruthy();
+  expect(clickEvent.snapshots.length).toBe(2);
+  const snapshotId = clickEvent.snapshots[0].snapshotId;
+  const snapshotEvent = traceEvents.find(event => event.type === 'snapshot' && event.snapshotId === snapshotId) as trace.FrameSnapshotTraceEvent;
+
+  expect(fs.existsSync(path.join(traceDir, 'resources', snapshotEvent.sha1))).toBe(true);
 });