implement screenshot search

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));
+  const snapshotServer = new SnapshotServer(traceModel.storage(), sha1 => traceModel.resourceForSha1(sha1), traceModel.contextEntries);
   loadedTraces.set(traceUrl, { traceModel, snapshotServer });
   return traceModel;
 }
@@ -123,12 +123,19 @@ async function doFetch(event: FetchEvent): Promise<Response> {
       const { snapshotServer } = loadedTraces.get(traceUrl!) || {};
       if (!snapshotServer)
         return new Response(null, { status: 404 });
-      const response = snapshotServer.serveSnapshot(relativePath, url.searchParams, url.href);
+      const response = snapshotServer.serveSnapshot(relativePath, url.searchParams, url.href, self.registration.scope, traceUrl!);
       if (isDeployedAsHttps)
         response.headers.set('Content-Security-Policy', 'upgrade-insecure-requests');
       return response;
     }
 
+    if (relativePath.startsWith('/screenshot/')) {
+      const { snapshotServer } = loadedTraces.get(traceUrl!) || {};
+      if (!snapshotServer)
+        return new Response(null, { status: 404 });
+      return snapshotServer.serveClosestScreenshot(relativePath, url.searchParams);
+    }
+
     if (relativePath.startsWith('/sha1/')) {
       // Sha1 for sources is based on the file path, can't load it of a random model.
       const sha1 = relativePath.slice('/sha1/'.length);

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

@@ -78,7 +78,7 @@ export class SnapshotRenderer {
     return this._snapshots[this._index].viewport;
   }
 
-  render(): RenderedFrameSnapshot {
+  render(swScope: string, traceURL: string): RenderedFrameSnapshot {
     const result: string[] = [];
     const visit = (n: NodeSnapshot, snapshotIndex: number, parentTag: string | undefined, parentAttrs: [string, string][] | undefined) => {
       // Text node.
@@ -154,12 +154,16 @@ export class SnapshotRenderer {
     const html = lruCache(this, () => {
       visit(snapshot.html, this._index, undefined, undefined);
 
+      const screenshotURL = new URL(`./screenshot/${snapshot.pageId}`, swScope);
+      screenshotURL.searchParams.set('trace', traceURL);
+      screenshotURL.searchParams.set('name', this.snapshotName!);
+
       const html = result.join('');
       // Hide the document in order to prevent flickering. We will unhide once script has processed shadow.
       const prefix = snapshot.doctype ? `<!DOCTYPE ${snapshot.doctype}>` : '';
       return prefix + [
         '<style>*,*::before,*::after { visibility: hidden }</style>',
-        `<script>${snapshotScript(this._callId, this.snapshotName)}</script>`
+        `<script>${snapshotScript(screenshotURL.toString(), this._callId, this.snapshotName)}</script>`
       ].join('') + html;
     });
 
@@ -242,8 +246,8 @@ function snapshotNodes(snapshot: FrameSnapshot): NodeSnapshot[] {
   return (snapshot as any)._nodes;
 }
 
-function snapshotScript(...targetIds: (string | undefined)[]) {
-  function applyPlaywrightAttributes(unwrapPopoutUrl: (url: string) => string, ...targetIds: (string | undefined)[]) {
+function snapshotScript(screenshotURL: string | undefined, ...targetIds: (string | undefined)[]) {
+  function applyPlaywrightAttributes(unwrapPopoutUrl: (url: string) => string, screenshotURL: string | undefined, ...targetIds: (string | undefined)[]) {
     const kPointerWarningTitle = 'Recorded click position in absolute coordinates did not' +
         ' match the center of the clicked element. This is likely due to a difference between' +
         ' the test runner and the trace viewer operating systems.';
@@ -300,9 +304,8 @@ function snapshotScript(...targetIds: (string | undefined)[]) {
       }
 
       const canvases = root.querySelectorAll('canvas');
-      if (canvases.length > 0) {
-        const sha1 = 'page@52b251b4d0b1412c19639922d9b22cb9-1728986751380.jpeg';
-        fetch(`http://[::1]:58477/trace/sha1/${sha1}`).then(response => response.blob()).then(blob => {
+      if (canvases.length > 0 && screenshotURL) {
+        fetch(screenshotURL).then(response => response.blob()).then(blob => {
           const img = new Image();
           img.onload = () => {
             for (const canvas of canvases) {
@@ -423,7 +426,7 @@ function snapshotScript(...targetIds: (string | undefined)[]) {
     window.addEventListener('DOMContentLoaded', onDOMContentLoaded);
   }
 
-  return `\n(${applyPlaywrightAttributes.toString()})(${unwrapPopoutUrl.toString()}${targetIds.map(id => `, "${id}"`).join('')})`;
+  return `\n(${applyPlaywrightAttributes.toString()})(${unwrapPopoutUrl.toString()}, ${JSON.stringify(screenshotURL)}${targetIds.map(id => `, "${id}"`).join('')})`;
 }
 
 

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

@@ -18,6 +18,16 @@ 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 };
 
@@ -25,21 +35,48 @@ 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>) {
+  constructor(snapshotStorage: SnapshotStorage, resourceLoader: (sha1: string) => Promise<Blob | undefined>, contextEntries: ContextEntry[]) {
     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 {
+  serveSnapshot(pathname: string, searchParams: URLSearchParams, snapshotUrl: string, swScope: string, traceUrl: string): Response {
     const snapshot = this._snapshot(pathname.substring('/snapshot'.length), searchParams);
     if (!snapshot)
       return new Response(null, { status: 404 });
-    const renderedSnapshot = snapshot.render();
+    const renderedSnapshot = snapshot.render(swScope, traceUrl);
     this._snapshotIds.set(snapshotUrl, snapshot);
     return new Response(renderedSnapshot.html, { status: 200, headers: { 'Content-Type': 'text/html; charset=utf-8' } });
   }
 
+  async serveClosestScreenshot(pathname: string, searchParams: URLSearchParams): Promise<Response> {
+    const snapshotRenderer = this._snapshot(pathname.substring('/screenshot'.length), searchParams);
+    if (!snapshotRenderer)
+      return new Response(undefined, { status: 404 });
+
+    const { wallTime, timestamp, pageId } = snapshotRenderer.snapshot();
+    const page = this._pages.get(pageId);
+    if (!page)
+      return new Response(undefined, { status: 404 });
+
+    let sha1 = undefined;
+    if (wallTime && page.screencastFrames[0]?.frameSwapWallTime)
+      sha1 = findClosest(page.screencastFrames, frame => frame.frameSwapWallTime!, wallTime)?.sha1;
+    sha1 ??= findClosest(page.screencastFrames, frame => frame.timestamp, timestamp)?.sha1;
+
+    if (!sha1)
+      return new Response(undefined, { status: 404 });
+
+    const blob = await this._resourceLoader(sha1);
+    if (!blob)
+      return new Response(undefined, { status: 404 });
+
+    return new Response(blob);
+  }
+
   serveSnapshotInfo(pathname: string, searchParams: URLSearchParams): Response {
     const snapshot = this._snapshot(pathname.substring('/snapshotInfo'.length), searchParams);
     return this._respondWithJson(snapshot ? {