From 923b94ae8e7193bf32447a04f77d9fb13626f21a Mon Sep 17 00:00:00 2001
From: Mike Cao <mike@mikecao.com>
Date: Thu, 7 May 2026 13:59:30 -0700
Subject: [PATCH] Render heatmaps from replay snapshots

---
 .../(reports)/heatmaps/Heatmap.module.css     |  19 +-
 .../(reports)/heatmaps/Heatmap.tsx            | 175 +++++++++-----
 src/queries/sql/heatmap/getHeatmap.ts         | 225 ++++++++++++++++--
 3 files changed, 344 insertions(+), 75 deletions(-)

diff --git a/src/app/(main)/websites/[websiteId]/(reports)/heatmaps/Heatmap.module.css b/src/app/(main)/websites/[websiteId]/(reports)/heatmaps/Heatmap.module.css
index 96b578475..b3c868063 100644
--- a/src/app/(main)/websites/[websiteId]/(reports)/heatmaps/Heatmap.module.css
+++ b/src/app/(main)/websites/[websiteId]/(reports)/heatmaps/Heatmap.module.css
@@ -82,13 +82,28 @@
   max-width: 100%;
 }
 
-.iframe {
+.snapshot {
   position: absolute;
   top: 0;
   left: 0;
-  border: 0;
   transform-origin: top left;
   pointer-events: none;
+  overflow: hidden;
+}
+
+.snapshot :global(.replayer-wrapper) {
+  border: 0;
+  overflow: hidden;
+}
+
+.snapshot :global(iframe) {
+  border: 0;
+  pointer-events: none;
+}
+
+.snapshot :global(.replayer-mouse),
+.snapshot :global(.replayer-mouse-tail) {
+  display: none !important;
 }
 
 .overlay {
diff --git a/src/app/(main)/websites/[websiteId]/(reports)/heatmaps/Heatmap.tsx b/src/app/(main)/websites/[websiteId]/(reports)/heatmaps/Heatmap.tsx
index 493c1327b..db1c54d45 100644
--- a/src/app/(main)/websites/[websiteId]/(reports)/heatmaps/Heatmap.tsx
+++ b/src/app/(main)/websites/[websiteId]/(reports)/heatmaps/Heatmap.tsx
@@ -2,13 +2,14 @@
 import { Column, Grid, Heading, Row, Text } from '@umami/react-zen';
 import { useEffect, useMemo, useRef, useState } from 'react';
 import { LoadingPanel } from '@/components/common/LoadingPanel';
-import { useResultQuery, useWebsite } from '@/components/hooks';
+import { useResultQuery } from '@/components/hooks';
+import { useReplayQuery } from '@/components/hooks/queries/useReplayQuery';
 import { formatLongNumber } from '@/lib/format';
-import type { HeatmapMode, HeatmapPoint, HeatmapResult } from '@/queries/sql';
+import type { HeatmapMode, HeatmapPoint, HeatmapResult, HeatmapSnapshot } from '@/queries/sql';
 import styles from './Heatmap.module.css';
+import 'rrweb/dist/replay/rrweb-replay.css';
 
 const MAX_RENDER_WIDTH = 1024;
-const IFRAME_SANDBOX = 'allow-same-origin allow-scripts allow-forms allow-popups';
 
 function useElementWidth<T extends HTMLElement>() {
   const ref = useRef<T | null>(null);
@@ -26,6 +27,10 @@ function useElementWidth<T extends HTMLElement>() {
   return [ref, width] as const;
 }
 
+interface ReplayData {
+  events: any[];
+}
+
 interface ViewportBucket {
   width: number;
   height: number;
@@ -41,7 +46,6 @@ interface HeatmapProps {
 }
 
 export function Heatmap({ websiteId, urlPath, onUrlPathChange, mode, onModeChange }: HeatmapProps) {
-  const website = useWebsite();
   const { data, error, isLoading } = useResultQuery<HeatmapResult>('heatmap', {
     websiteId,
     urlPath: urlPath || undefined,
@@ -62,16 +66,14 @@ export function Heatmap({ websiteId, urlPath, onUrlPathChange, mode, onModeChang
             mode === 'scroll' ? (
               <ScrollHeatmapView
                 websiteId={websiteId}
-                domain={website?.domain ?? null}
-                urlPath={urlPath}
                 scroll={scroll}
+                snapshot={data?.snapshot ?? null}
               />
             ) : (
               <HeatmapView
                 websiteId={websiteId}
-                domain={website?.domain ?? null}
-                urlPath={urlPath}
                 points={points}
+                snapshot={data?.snapshot ?? null}
               />
             )
           ) : (
@@ -142,27 +144,6 @@ function PageList({
   );
 }
 
-function buildIframeSrc(
-  websiteId: string | null,
-  domain: string | null,
-  urlPath: string,
-): string | null {
-  // Self-record: data was captured from this very umami instance, so iframe the
-  // current origin instead of the website's stored domain (which often points to prod).
-  if (typeof window !== 'undefined' && websiteId && websiteId === process.env.selfRecord) {
-    return `${window.location.origin}${urlPath}`;
-  }
-  if (!domain) return null;
-  if (/^https?:\/\//i.test(domain)) return `${domain}${urlPath}`;
-  const isLocal = /^(localhost|127\.|0\.0\.0\.0|\[::1\])(:|\/|$)/i.test(domain);
-  const protocol = isLocal
-    ? typeof window !== 'undefined'
-      ? window.location.protocol
-      : 'http:'
-    : 'https:';
-  return `${protocol}//${domain}${urlPath}`;
-}
-
 function pickViewport(points: HeatmapPoint[]): ViewportBucket | null {
   if (!points.length) return null;
   const buckets = new Map<string, ViewportBucket>();
@@ -184,14 +165,12 @@ function pickViewport(points: HeatmapPoint[]): ViewportBucket | null {
 
 function HeatmapView({
   websiteId,
-  domain,
-  urlPath,
   points,
+  snapshot,
 }: {
   websiteId: string;
-  domain: string | null;
-  urlPath: string;
   points: HeatmapPoint[];
+  snapshot: HeatmapSnapshot | null;
 }) {
   const [showPage, setShowPage] = useState(true);
 
@@ -216,7 +195,6 @@ function HeatmapView({
   const renderWidth = containerWidth > 0 ? Math.min(containerWidth, MAX_RENDER_WIDTH) : 0;
   const scale = renderWidth ? renderWidth / viewport.width : 0;
   const renderHeight = Math.round(viewport.height * scale);
-  const iframeSrc = buildIframeSrc(websiteId, domain, urlPath);
 
   return (
     <Column gap>
@@ -225,7 +203,7 @@ function HeatmapView({
           {visible.length} positions · {formatLongNumber(visible.reduce((s, p) => s + p.count, 0))}{' '}
           clicks · viewport {viewport.width}×{viewport.height}
         </Text>
-        {iframeSrc && (
+        {snapshot && (
           <button
             type="button"
             className={styles.toggleButton}
@@ -240,14 +218,13 @@ function HeatmapView({
           className={styles.canvas}
           style={{ width: renderWidth || '100%', height: renderHeight || 0 }}
         >
-          {renderWidth > 0 && showPage && iframeSrc && (
-            <iframe
-              className={styles.iframe}
-              src={iframeSrc}
+          {renderWidth > 0 && showPage && snapshot && (
+            <ReplaySnapshot
+              websiteId={websiteId}
+              snapshot={snapshot}
               width={viewport.width}
               height={viewport.height}
-              style={{ transform: `scale(${scale})` }}
-              sandbox={IFRAME_SANDBOX}
+              scale={scale}
             />
           )}
           <div className={styles.overlay}>
@@ -278,14 +255,12 @@ function HeatmapView({
 
 function ScrollHeatmapView({
   websiteId,
-  domain,
-  urlPath,
   scroll,
+  snapshot,
 }: {
   websiteId: string;
-  domain: string | null;
-  urlPath: string;
   scroll: HeatmapResult['scroll'] | undefined;
+  snapshot: HeatmapSnapshot | null;
 }) {
   const [showPage, setShowPage] = useState(true);
 
@@ -299,7 +274,6 @@ function ScrollHeatmapView({
   const renderWidth = containerWidth > 0 ? Math.min(containerWidth, MAX_RENDER_WIDTH) : 0;
   const scale = renderWidth ? renderWidth / viewportW : 0;
   const renderHeight = Math.round(pageH * scale);
-  const iframeSrc = buildIframeSrc(websiteId, domain, urlPath);
 
   // Cumulative reach: % of sessions that scrolled at least to depth D.
   const sortedBuckets = [...buckets].sort((a, b) => a.depth - b.depth);
@@ -334,7 +308,7 @@ function ScrollHeatmapView({
           {formatLongNumber(totalSessions)} sessions · page {viewportW}×{pageH}
           {viewportH ? ` · viewport ${viewportH}` : ''}
         </Text>
-        {iframeSrc && (
+        {snapshot && (
           <button
             type="button"
             className={styles.toggleButton}
@@ -349,14 +323,13 @@ function ScrollHeatmapView({
           className={styles.canvas}
           style={{ width: renderWidth || '100%', height: renderHeight || 0 }}
         >
-          {renderWidth > 0 && showPage && iframeSrc && (
-            <iframe
-              className={styles.iframe}
-              src={iframeSrc}
+          {renderWidth > 0 && showPage && snapshot && (
+            <ReplaySnapshot
+              websiteId={websiteId}
+              snapshot={snapshot}
               width={viewportW}
               height={pageH}
-              style={{ transform: `scale(${scale})` }}
-              sandbox={IFRAME_SANDBOX}
+              scale={scale}
             />
           )}
           <div className={styles.overlay}>
@@ -390,6 +363,102 @@ function ScrollHeatmapView({
   );
 }
 
+function ReplaySnapshot({
+  websiteId,
+  snapshot,
+  width,
+  height,
+  scale,
+}: {
+  websiteId: string;
+  snapshot: HeatmapSnapshot;
+  width: number;
+  height: number;
+  scale: number;
+}) {
+  const containerRef = useRef<HTMLDivElement>(null);
+  const replayerRef = useRef<any>(null);
+  const { data } = useReplayQuery(websiteId, snapshot.replayId) as { data?: ReplayData };
+
+  useEffect(() => {
+    const container = containerRef.current;
+    const events = data?.events;
+    if (!container || !events?.length) return;
+
+    let cancelled = false;
+
+    import('rrweb').then(({ Replayer }) => {
+      if (cancelled || !containerRef.current) return;
+
+      container.innerHTML = '';
+
+      const replayer = new Replayer(events, {
+        root: container,
+        showWarning: false,
+        mouseTail: false,
+        triggerFocus: false,
+        pauseAnimation: true,
+        useVirtualDom: false,
+      });
+
+      replayerRef.current = replayer;
+
+      const freeze = () => {
+        const offset = Math.max(0, snapshot.timestamp - events[0].timestamp);
+        replayer.pause(offset);
+        replayer.disableInteract();
+        resizeReplayFrame(replayer, width, height);
+      };
+
+      requestAnimationFrame(() => {
+        freeze();
+        requestAnimationFrame(freeze);
+      });
+    });
+
+    return () => {
+      cancelled = true;
+      if (replayerRef.current) {
+        replayerRef.current.destroy();
+        replayerRef.current = null;
+      }
+      if (container) {
+        container.innerHTML = '';
+      }
+    };
+  }, [data?.events, height, snapshot.timestamp, width]);
+
+  useEffect(() => {
+    if (replayerRef.current) {
+      resizeReplayFrame(replayerRef.current, width, height);
+    }
+  }, [height, width]);
+
+  return (
+    <div
+      ref={containerRef}
+      className={styles.snapshot}
+      style={{ width, height, transform: `scale(${scale})` }}
+    />
+  );
+}
+
+function resizeReplayFrame(replayer: any, width: number, height: number) {
+  const { iframe, wrapper } = replayer;
+
+  if (wrapper) {
+    wrapper.style.width = `${width}px`;
+    wrapper.style.height = `${height}px`;
+  }
+
+  if (iframe) {
+    iframe.setAttribute('width', String(width));
+    iframe.setAttribute('height', String(height));
+    iframe.style.width = `${width}px`;
+    iframe.style.height = `${height}px`;
+  }
+}
+
 function EmptyState({ message }: { message?: string } = {}) {
   return (
     <Column alignItems="center" justifyContent="center" height="100%" gap>
diff --git a/src/queries/sql/heatmap/getHeatmap.ts b/src/queries/sql/heatmap/getHeatmap.ts
index d172bbe30..c14c4b33d 100644
--- a/src/queries/sql/heatmap/getHeatmap.ts
+++ b/src/queries/sql/heatmap/getHeatmap.ts
@@ -38,10 +38,16 @@ export interface HeatmapScrollBucket {
   sessions: number;
 }
 
+export interface HeatmapSnapshot {
+  replayId: string;
+  timestamp: number;
+}
+
 export interface HeatmapResult {
   mode: HeatmapMode;
   pages: HeatmapPage[];
   points: HeatmapPoint[];
+  snapshot: HeatmapSnapshot | null;
   scroll: {
     buckets: HeatmapScrollBucket[];
     totalSessions: number;
@@ -69,6 +75,29 @@ const emptyScroll = (): HeatmapResult['scroll'] => ({
   viewportH: null,
 });
 
+function pickSnapshotViewport(points: HeatmapPoint[]): { width: number; height: number } | null {
+  const buckets = new Map<string, { width: number; height: number; count: number }>();
+
+  for (const p of points) {
+    const key = `${p.viewportW}x${p.viewportH}`;
+    const existing = buckets.get(key);
+    if (existing) {
+      existing.count += p.count;
+    } else {
+      buckets.set(key, { width: p.viewportW, height: p.viewportH, count: p.count });
+    }
+  }
+
+  let best: { width: number; height: number; count: number } | null = null;
+  for (const bucket of buckets.values()) {
+    if (!best || bucket.count > best.count) {
+      best = bucket;
+    }
+  }
+
+  return best ? { width: best.width, height: best.height } : null;
+}
+
 async function relationalQuery(
   websiteId: string,
   { startDate, endDate, urlPath, mode = 'click' }: HeatmapParameters,
@@ -95,7 +124,7 @@ async function relationalQuery(
   );
 
   if (!urlPath) {
-    return { mode, pages, points: [], scroll: emptyScroll() };
+    return { mode, pages, points: [], snapshot: null, scroll: emptyScroll() };
   }
 
   if (mode === 'scroll') {
@@ -145,17 +174,29 @@ async function relationalQuery(
     );
 
     const dim = dimRows[0];
+    const scroll = {
+      buckets: bucketRows.map(r => ({ depth: Number(r.depth), sessions: Number(r.sessions) })),
+      totalSessions: Number(dim?.totalSessions ?? 0),
+      pageH: dim?.pageH ?? null,
+      viewportW: dim?.viewportW ?? null,
+      viewportH: dim?.viewportH ?? null,
+    };
+    const snapshot = await getRelationalSnapshot(rawQuery, {
+      websiteId,
+      eventType,
+      urlPath,
+      startDate,
+      endDate,
+      viewportW: scroll.viewportW,
+      viewportH: scroll.viewportH,
+    });
+
     return {
       mode,
       pages,
       points: [],
-      scroll: {
-        buckets: bucketRows.map(r => ({ depth: Number(r.depth), sessions: Number(r.sessions) })),
-        totalSessions: Number(dim?.totalSessions ?? 0),
-        pageH: dim?.pageH ?? null,
-        viewportW: dim?.viewportW ?? null,
-        viewportH: dim?.viewportH ?? null,
-      },
+      snapshot,
+      scroll,
     };
   }
 
@@ -185,7 +226,73 @@ async function relationalQuery(
     FUNCTION_NAME,
   );
 
-  return { mode, pages, points: rawPoints, scroll: emptyScroll() };
+  const viewport = pickSnapshotViewport(rawPoints);
+  const snapshot = await getRelationalSnapshot(rawQuery, {
+    websiteId,
+    eventType,
+    urlPath,
+    startDate,
+    endDate,
+    viewportW: viewport?.width ?? null,
+    viewportH: viewport?.height ?? null,
+  });
+
+  return { mode, pages, points: rawPoints, snapshot, scroll: emptyScroll() };
+}
+
+async function getRelationalSnapshot(
+  rawQuery: typeof prisma.rawQuery,
+  {
+    websiteId,
+    eventType,
+    urlPath,
+    startDate,
+    endDate,
+    viewportW,
+    viewportH,
+  }: {
+    websiteId: string;
+    eventType: number;
+    urlPath: string;
+    startDate: Date;
+    endDate: Date;
+    viewportW: number | null;
+    viewportH: number | null;
+  },
+): Promise<HeatmapSnapshot | null> {
+  const viewportFilter =
+    viewportW && viewportH
+      ? `
+      and h.viewport_w = {{viewportW}}
+      and h.viewport_h = {{viewportH}}
+    `
+      : '';
+
+  const rows: { replayId: string; timestamp: number | string }[] = await rawQuery(
+    `
+    select
+      h.visit_id as "replayId",
+      (extract(epoch from h.created_at) * 1000)::bigint as "timestamp"
+    from heatmap_event h
+    inner join (
+      select distinct visit_id
+      from session_replay
+      where website_id = {{websiteId::uuid}}
+    ) sr on sr.visit_id = h.visit_id
+    where h.website_id = {{websiteId::uuid}}
+      and h.event_type = {{eventType}}
+      and h.url_path = {{urlPath}}
+      and h.created_at between {{startDate}} and {{endDate}}
+      ${viewportFilter}
+    order by h.created_at asc
+    limit 1
+    `,
+    { websiteId, eventType, urlPath, startDate, endDate, viewportW, viewportH },
+    FUNCTION_NAME,
+  );
+
+  const row = rows[0];
+  return row ? { replayId: row.replayId, timestamp: Number(row.timestamp) } : null;
 }
 
 async function clickhouseQuery(
@@ -222,7 +329,7 @@ async function clickhouseQuery(
   }));
 
   if (!urlPath) {
-    return { mode, pages, points: [], scroll: emptyScroll() };
+    return { mode, pages, points: [], snapshot: null, scroll: emptyScroll() };
   }
 
   if (mode === 'scroll') {
@@ -274,19 +381,31 @@ async function clickhouseQuery(
     );
 
     const dim = dimRows[0];
+    const scroll = {
+      buckets: bucketRows.map(r => ({ depth: Number(r.depth), sessions: Number(r.sessions) })),
+      totalSessions: Number(dim?.totalSessions ?? 0),
+      pageH: dim?.pageH === null || dim?.pageH === undefined ? null : Number(dim.pageH),
+      viewportW:
+        dim?.viewportW === null || dim?.viewportW === undefined ? null : Number(dim.viewportW),
+      viewportH:
+        dim?.viewportH === null || dim?.viewportH === undefined ? null : Number(dim.viewportH),
+    };
+    const snapshot = await getClickhouseSnapshot(rawQuery, {
+      websiteId,
+      eventType,
+      urlPath,
+      startDate,
+      endDate,
+      viewportW: scroll.viewportW,
+      viewportH: scroll.viewportH,
+    });
+
     return {
       mode,
       pages,
       points: [],
-      scroll: {
-        buckets: bucketRows.map(r => ({ depth: Number(r.depth), sessions: Number(r.sessions) })),
-        totalSessions: Number(dim?.totalSessions ?? 0),
-        pageH: dim?.pageH === null || dim?.pageH === undefined ? null : Number(dim.pageH),
-        viewportW:
-          dim?.viewportW === null || dim?.viewportW === undefined ? null : Number(dim.viewportW),
-        viewportH:
-          dim?.viewportH === null || dim?.viewportH === undefined ? null : Number(dim.viewportH),
-      },
+      snapshot,
+      scroll,
     };
   }
 
@@ -334,5 +453,71 @@ async function clickhouseQuery(
     count: Number(p.count),
   }));
 
-  return { mode, pages, points, scroll: emptyScroll() };
+  const viewport = pickSnapshotViewport(points);
+  const snapshot = await getClickhouseSnapshot(rawQuery, {
+    websiteId,
+    eventType,
+    urlPath,
+    startDate,
+    endDate,
+    viewportW: viewport?.width ?? null,
+    viewportH: viewport?.height ?? null,
+  });
+
+  return { mode, pages, points, snapshot, scroll: emptyScroll() };
+}
+
+async function getClickhouseSnapshot(
+  rawQuery: typeof clickhouse.rawQuery,
+  {
+    websiteId,
+    eventType,
+    urlPath,
+    startDate,
+    endDate,
+    viewportW,
+    viewportH,
+  }: {
+    websiteId: string;
+    eventType: number;
+    urlPath: string;
+    startDate: Date;
+    endDate: Date;
+    viewportW: number | null;
+    viewportH: number | null;
+  },
+): Promise<HeatmapSnapshot | null> {
+  const viewportFilter =
+    viewportW && viewportH
+      ? `
+      and h.viewport_w = {viewportW:UInt32}
+      and h.viewport_h = {viewportH:UInt32}
+    `
+      : '';
+
+  const rows = await rawQuery<{ replayId: string; timestamp: string | number }[]>(
+    `
+    select
+      toString(h.visit_id) as replayId,
+      toUnixTimestamp64Milli(h.created_at) as timestamp
+    from heatmap_event h
+    inner join (
+      select distinct visit_id
+      from session_replay
+      where website_id = {websiteId:UUID}
+    ) sr on sr.visit_id = h.visit_id
+    where h.website_id = {websiteId:UUID}
+      and h.event_type = {eventType:UInt8}
+      and h.url_path = {urlPath:String}
+      and h.created_at between {startDate:DateTime64} and {endDate:DateTime64}
+      ${viewportFilter}
+    order by h.created_at asc
+    limit 1
+    `,
+    { websiteId, eventType, urlPath, startDate, endDate, viewportW, viewportH },
+    FUNCTION_NAME,
+  );
+
+  const row = rows[0];
+  return row ? { replayId: row.replayId, timestamp: Number(row.timestamp) } : null;
 }