From 6d359b878cbce2a8a6bf8e96fdf59d46ae0af9b2 Mon Sep 17 00:00:00 2001
From: Oleg Isonen <oleg008@gmail.com>
Date: Sat, 17 Sep 2022 17:17:55 +0200
Subject: [PATCH] make sure we always catch the iframe when its fully loaded

---
 apps/designer/app/canvas/canvas.tsx           |  4 +--
 .../app/canvas/shared/use-designer-ready.ts   |  2 +-
 apps/designer/app/designer/designer.tsx       | 26 +++++++++++++++----
 apps/designer/app/shared/dom-utils/index.ts   |  1 +
 .../app/shared/dom-utils/on-iframe-load.ts    | 16 ++++++++++++
 5 files changed, 41 insertions(+), 8 deletions(-)
 create mode 100644 apps/designer/app/shared/dom-utils/on-iframe-load.ts

diff --git a/apps/designer/app/canvas/canvas.tsx b/apps/designer/app/canvas/canvas.tsx
index ce422c4fc1..81914d2ae5 100644
--- a/apps/designer/app/canvas/canvas.tsx
+++ b/apps/designer/app/canvas/canvas.tsx
@@ -46,7 +46,7 @@ import { usePublishScrollState } from "./shared/use-publish-scroll-state";
 import { useDragAndDrop } from "./shared/use-drag-drop";
 import { setInstanceChildrenMutable } from "~/shared/tree-utils";
 import { CanvasData } from "~/shared/db";
-import { useDesignerReady } from "./shared/use-designer-ready";
+import { useSubscribeDesignerReady } from "./shared/use-designer-ready";
 
 registerContainers();
 
@@ -119,7 +119,7 @@ export const Canvas = ({ data }: CanvasProps): JSX.Element | null => {
   if (data.tree === null) {
     throw new Error("Tree is null");
   }
-  const isDesignerReady = useDesignerReady();
+  const isDesignerReady = useSubscribeDesignerReady();
   useInitializeBreakpoints(data.breakpoints);
   globalStyles();
   useAllUserProps(data.props);
diff --git a/apps/designer/app/canvas/shared/use-designer-ready.ts b/apps/designer/app/canvas/shared/use-designer-ready.ts
index 81108a26f8..82ea83a792 100644
--- a/apps/designer/app/canvas/shared/use-designer-ready.ts
+++ b/apps/designer/app/canvas/shared/use-designer-ready.ts
@@ -7,7 +7,7 @@ declare module "~/shared/pubsub" {
   }
 }
 
-export const useDesignerReady = () => {
+export const useSubscribeDesignerReady = () => {
   const [isReady, setIsReady] = useState<boolean>(false);
   useSubscribe("designerReady", () => {
     setIsReady(true);
diff --git a/apps/designer/app/designer/designer.tsx b/apps/designer/app/designer/designer.tsx
index f3108aacaa..84b5f56a12 100644
--- a/apps/designer/app/designer/designer.tsx
+++ b/apps/designer/app/designer/designer.tsx
@@ -1,4 +1,4 @@
-import { useCallback } from "react";
+import { useCallback, useEffect, useState } from "react";
 import { type Publish, usePublish, useSubscribe } from "~/shared/pubsub";
 import type { Project } from "@webstudio-is/project";
 import type { Config } from "~/config";
@@ -33,6 +33,7 @@ import {
 import { useClientSettings } from "./shared/client-settings";
 import { Navigator } from "./features/sidebar-left";
 import { PANEL_WIDTH } from "./shared/constants";
+import { onLoad } from "~/shared/dom-utils";
 
 export const links = () => {
   return [
@@ -68,6 +69,22 @@ const useNavigatorLayout = () => {
   return isLoaded ? clientSettings.navigatorLayout : "docked";
 };
 
+const usePublishDesignerReady = (publish: Publish) => {
+  const [ref, setRef] = useState(null);
+
+  useEffect(() => {
+    if (ref === null) return;
+
+    onLoad(ref).then(() => {
+      // We publish this even to let canvas know that we are now listening to the events, otherwise if canvas loads faster than designer, which is possible with SSR,
+      // we can miss the events and designer will just not connect to the canvas.
+      publish({ type: "designerReady" });
+    });
+  }, [ref, publish]);
+
+  return setRef;
+};
+
 type SidePanelProps = {
   children: JSX.Element | Array<JSX.Element>;
   isPreviewMode: boolean;
@@ -226,17 +243,16 @@ export const Designer = ({ config, project }: DesignerProps) => {
   const onRefReadCanvasWidth = useUpdateCanvasWidth();
   const { onRef: onRefReadCanvas, onTransitionEnd } = useReadCanvasRect();
   const [dragAndDropState] = useDragAndDropState();
+  const onRefIframeReady = usePublishDesignerReady(publish);
 
   const iframeRefCallback = useCallback(
     (ref) => {
       publishRef.current = ref;
       onRefReadCanvasWidth(ref);
       onRefReadCanvas(ref);
-      // We publish this even to let canvas know that we are now listening to the events, otherwise if canvas loads faster than designer, which is possible with SSR,
-      // we can miss the events and designer will just not connect to the canvas.
-      publish({ type: "designerReady" });
+      onRefIframeReady(ref);
     },
-    [publishRef, onRefReadCanvasWidth, onRefReadCanvas, publish]
+    [publishRef, onRefReadCanvasWidth, onRefReadCanvas, onRefIframeReady]
   );
 
   return (
diff --git a/apps/designer/app/shared/dom-utils/index.ts b/apps/designer/app/shared/dom-utils/index.ts
index b1120b975a..1304bcc6a9 100644
--- a/apps/designer/app/shared/dom-utils/index.ts
+++ b/apps/designer/app/shared/dom-utils/index.ts
@@ -1 +1,2 @@
 export * from "./instance";
+export * from "./on-iframe-load";
diff --git a/apps/designer/app/shared/dom-utils/on-iframe-load.ts b/apps/designer/app/shared/dom-utils/on-iframe-load.ts
new file mode 100644
index 0000000000..766329bc54
--- /dev/null
+++ b/apps/designer/app/shared/dom-utils/on-iframe-load.ts
@@ -0,0 +1,16 @@
+/**
+ * Resolve a promise when an iframe has readyState complete.
+ * Checks periodically for the readyState to be complete, because somttimes the load event may simply not fire for various reasons.
+ */
+export const onLoad = (iframe: HTMLIFrameElement) =>
+  new Promise((resolve) => {
+    if (iframe === null) return;
+    const check = () => {
+      if (document.readyState === "complete") {
+        clearInterval(intervalId);
+        resolve("complete");
+      }
+    };
+    iframe.addEventListener("load", check);
+    const intervalId = setInterval(check, 100);
+  });