move closestScreenshot logic into renderer

packages/playwright-core/src/server/trace/test/inMemorySnapshotter.ts

@@ -85,7 +85,7 @@ export class InMemorySnapshotter implements SnapshotterDelegate, HarTracerDelega
 
   onFrameSnapshot(snapshot: FrameSnapshot): void {
     ++this._snapshotCount;
-    const renderer = this._storage.addFrameSnapshot(snapshot);
+    const renderer = this._storage.addFrameSnapshot(snapshot, []);
     this._snapshotReadyPromises.get(snapshot.snapshotName || '')?.resolve(renderer);
   }
 

packages/trace-viewer/src/sw/main.ts

@@ -64,7 +64,7 @@ async function loadTrace(traceUrl: string, traceFileName: string | null, clientI
       throw new Error(`Could not load trace from ${traceFileName}. Make sure to upload a valid Playwright trace.`);
     throw new Error(`Could not load trace from ${traceUrl}. Make sure a valid Playwright Trace is accessible over this url.`);
   }
-  const snapshotServer = new SnapshotServer(traceModel.storage(), sha1 => traceModel.resourceForSha1(sha1), traceModel.contextEntries);
+  const snapshotServer = new SnapshotServer(traceModel.storage(), sha1 => traceModel.resourceForSha1(sha1));
   loadedTraces.set(traceUrl, { traceModel, snapshotServer });
   return traceModel;
 }

packages/trace-viewer/src/sw/snapshotRenderer.ts

@@ -16,6 +16,16 @@
 
 import { escapeHTMLAttribute, escapeHTML } from '@isomorphic/stringUtils';
 import type { FrameSnapshot, NodeNameAttributesChildNodesSnapshot, NodeSnapshot, RenderedFrameSnapshot, ResourceSnapshot, SubtreeReferenceSnapshot } from '@trace/snapshot';
+import type { PageEntry } from '../types/entries';
+
+function findClosest<T>(items: T[], metric: (v: T) => number, target: number) {
+  return items.find((item, index) => {
+    if (index === items.length - 1)
+      return true;
+    const next = items[index + 1];
+    return Math.abs(metric(item) - target) < Math.abs(metric(next) - target);
+  });
+}
 
 function isNodeNameAttributesChildNodesSnapshot(n: NodeSnapshot): n is NodeNameAttributesChildNodesSnapshot {
   return Array.isArray(n) && typeof n[0] === 'string';
@@ -60,13 +70,15 @@ export class SnapshotRenderer {
   private _resources: ResourceSnapshot[];
   private _snapshot: FrameSnapshot;
   private _callId: string;
+  private _screencastFrames: PageEntry['screencastFrames'];
 
-  constructor(resources: ResourceSnapshot[], snapshots: FrameSnapshot[], index: number) {
+  constructor(resources: ResourceSnapshot[], snapshots: FrameSnapshot[], screencastFrames: PageEntry['screencastFrames'], index: number) {
     this._resources = resources;
     this._snapshots = snapshots;
     this._index = index;
     this._snapshot = snapshots[index];
     this._callId = snapshots[index].callId;
+    this._screencastFrames = screencastFrames;
     this.snapshotName = snapshots[index].snapshotName;
   }
 
@@ -78,6 +90,14 @@ export class SnapshotRenderer {
     return this._snapshots[this._index].viewport;
   }
 
+  closestScreenshot(): string | undefined {
+    const { wallTime, timestamp } = this.snapshot();
+    const closestFrame = (wallTime && this._screencastFrames[0]?.frameSwapWallTime)
+      ? findClosest(this._screencastFrames, frame => frame.frameSwapWallTime!, wallTime)
+      : findClosest(this._screencastFrames, frame => frame.timestamp, timestamp);
+    return closestFrame?.sha1;
+  }
+
   render(): RenderedFrameSnapshot {
     const result: string[] = [];
     const visit = (n: NodeSnapshot, snapshotIndex: number, parentTag: string | undefined, parentAttrs: [string, string][] | undefined) => {
@@ -437,6 +457,8 @@ function snapshotScript(...targetIds: (string | undefined)[]) {
             const yEnd = boundingRect.bottom / window.innerHeight;
 
             drawWarningBackground(context, canvas);
+
+            // todo: don't show the image if we're in an iframe - we know it's not going to be accurate
             context.drawImage(img, xStart * img.width, yStart * img.height, (xEnd - xStart) * img.width, (yEnd - yStart) * img.height, 0, 0, canvas.width, canvas.height);
             drawWarningIcon(context);
             if (isUnderTest)

packages/trace-viewer/src/sw/snapshotServer.ts

@@ -18,16 +18,6 @@ import type { URLSearchParams } from 'url';
 import type { SnapshotRenderer } from './snapshotRenderer';
 import type { SnapshotStorage } from './snapshotStorage';
 import type { ResourceSnapshot } from '@trace/snapshot';
-import type { ContextEntry, PageEntry } from '../types/entries';
-
-function findClosest<T>(items: T[], metric: (v: T) => number, target: number) {
-  return items.find((item, index) => {
-    if (index === items.length - 1)
-      return true;
-    const next = items[index + 1];
-    return Math.abs(metric(item) - target) < Math.abs(metric(next) - target);
-  });
-}
 
 type Point = { x: number, y: number };
 
@@ -35,12 +25,10 @@ export class SnapshotServer {
   private _snapshotStorage: SnapshotStorage;
   private _resourceLoader: (sha1: string) => Promise<Blob | undefined>;
   private _snapshotIds = new Map<string, SnapshotRenderer>();
-  private _pages: Map<string, PageEntry>;
 
-  constructor(snapshotStorage: SnapshotStorage, resourceLoader: (sha1: string) => Promise<Blob | undefined>, contextEntries: ContextEntry[]) {
+  constructor(snapshotStorage: SnapshotStorage, resourceLoader: (sha1: string) => Promise<Blob | undefined>) {
     this._snapshotStorage = snapshotStorage;
     this._resourceLoader = resourceLoader;
-    this._pages = new Map(contextEntries.flatMap(c => c.pages.map(p => [p.pageId, p])));
   }
 
   serveSnapshot(pathname: string, searchParams: URLSearchParams, snapshotUrl: string): Response {
@@ -55,23 +43,10 @@ export class SnapshotServer {
 
   async serveClosestScreenshot(pathname: string, searchParams: URLSearchParams): Promise<Response> {
     const snapshot = this._snapshot(pathname.substring('/closest-screenshot'.length), searchParams);
-    if (!snapshot)
+    const sha1 = snapshot?.closestScreenshot();
+    if (!sha1)
       return new Response(null, { status: 404 });
-
-    const { wallTime, timestamp, pageId } = snapshot.snapshot();
-    const page = this._pages.get(pageId);
-    if (!page)
-      return new Response(null, { status: 404 });
-
-    const closestFrame = (wallTime && page.screencastFrames[0]?.frameSwapWallTime) ? findClosest(page.screencastFrames, frame => frame.frameSwapWallTime!, wallTime) : findClosest(page.screencastFrames, frame => frame.timestamp, timestamp);
-    if (!closestFrame)
-      return new Response(null, { status: 404 });
-
-    const blob = await this._resourceLoader(closestFrame.sha1);
-    if (!blob)
-      return new Response(null, { status: 404 });
-
-    return new Response(blob);
+    return new Response(await this._resourceLoader(sha1));
   }
 
   serveSnapshotInfo(pathname: string, searchParams: URLSearchParams): Response {

packages/trace-viewer/src/sw/snapshotStorage.ts

@@ -16,6 +16,7 @@
 
 import type { FrameSnapshot, ResourceSnapshot } from '@trace/snapshot';
 import { rewriteURLForCustomProtocol, SnapshotRenderer } from './snapshotRenderer';
+import type { PageEntry } from '../types/entries';
 
 export class SnapshotStorage {
   private _resources: ResourceSnapshot[] = [];
@@ -29,7 +30,7 @@ export class SnapshotStorage {
     this._resources.push(resource);
   }
 
-  addFrameSnapshot(snapshot: FrameSnapshot) {
+  addFrameSnapshot(snapshot: FrameSnapshot, screencastFrames: PageEntry['screencastFrames']) {
     for (const override of snapshot.resourceOverrides)
       override.url = rewriteURLForCustomProtocol(override.url);
     let frameSnapshots = this._frameSnapshots.get(snapshot.frameId);
@@ -43,7 +44,7 @@ export class SnapshotStorage {
         this._frameSnapshots.set(snapshot.pageId, frameSnapshots);
     }
     frameSnapshots.raw.push(snapshot);
-    const renderer = new SnapshotRenderer(this._resources, frameSnapshots.raw, frameSnapshots.raw.length - 1);
+    const renderer = new SnapshotRenderer(this._resources, frameSnapshots.raw, screencastFrames, frameSnapshots.raw.length - 1);
     frameSnapshots.renderers.push(renderer);
     return renderer;
   }

packages/trace-viewer/src/sw/traceModernizer.ts

@@ -159,7 +159,7 @@ export class TraceModernizer {
         contextEntry.resources.push(event.snapshot);
         break;
       case 'frame-snapshot':
-        this._snapshotStorage.addFrameSnapshot(event.snapshot);
+        this._snapshotStorage.addFrameSnapshot(event.snapshot, this._pageEntry(event.snapshot.pageId).screencastFrames);
         break;
     }
     // Make sure there is a page entry for each page, even without screencast frames,