fix: improve lazy loading of enableVisualEditing (#1528)

packages/overlays/package.json

@@ -102,7 +102,7 @@
     "root": true
   },
   "dependencies": {
-    "@sanity/visual-editing": "1.8.20"
+    "@sanity/visual-editing": "workspace:*"
   },
   "devDependencies": {
     "@repo/package.config": "workspace:*",

packages/svelte-loader/package.json

@@ -98,7 +98,7 @@
     "@sanity/client": "^6.18.2",
     "@sanity/pkg-utils": "6.8.15",
     "@sanity/preview-url-secret": "^1.6.13",
-    "@sanity/visual-editing": "1.8.20",
+    "@sanity/visual-editing": "workspace:*",
     "@sveltejs/kit": "^2.5.8",
     "@typescript-eslint/eslint-plugin": "^7.9.0",
     "@typescript-eslint/parser": "^7.9.0",

packages/visual-editing/src/ui/enableVisualEditing.tsx

@@ -1,58 +1,32 @@
-import type {Root} from 'react-dom/client'
-
-import {OVERLAY_ID} from '../constants'
 import type {DisableVisualEditing, VisualEditingOptions} from '../types'
 
-let node: HTMLElement | null = null
-let root: Root | null = null
-let cleanup: number | null = null
-
 /**
  * Enables Visual Editing overlay in a page with sourcemap encoding.
  *
  * This will overlay UI on hovered elements that deep-links to Sanity Studio.
  * @public
  */
 export function enableVisualEditing(options: VisualEditingOptions = {}): DisableVisualEditing {
-  if (cleanup) clearTimeout(cleanup)
   const controller = new AbortController()
-
-  // Lazy load everything needed to render the app
-  Promise.all([import('react-dom/client'), import('./VisualEditing')]).then(
-    ([reactClient, {VisualEditing}]) => {
-      if (controller.signal.aborted) return
-
-      if (!node) {
-        node = document.createElement('div')
-        node.id = OVERLAY_ID
-        document.body.appendChild(node)
-      }
-
-      if (!root) {
-        const {createRoot} = 'default' in reactClient ? reactClient.default : reactClient
-        root = createRoot(node)
-      }
-
-      root.render(
-        <>
-          <VisualEditing {...options} />
-        </>,
-      )
-    },
-  )
+  // Lazy load everything, react, the app, all of it
+  import('./renderVisualEditing').then(({renderVisualEditing}) => {
+    const {signal} = controller
+    /**
+     * Due to lazy loading it's possible for the following to happen, and is a consequence of dynamic ESM imports not being cancellable natively:
+     * 1. Userland calls `const disableVisualEditing = enableVisualEditing()` and the dynamic ESM import is started.
+     * 2. The dynamic import uses the network, and it takes a while to load.
+     * 3. The user navigates to a different page in the app that doesn't need Visual Editing, for example a login page.
+     * 4. Since the app is no longer in a state where Visual Editing is needed, the user calls `disableVisualEditing()`.
+     * 5. The dynamic import eventually resolves and this function is called.
+     * When this happens we want to skip calling `renderVisualEditing` since we know it's no longer needed.
+     */
+    if (signal.aborted) return
+
+    // Hand off to the render function with the signal, which will be subscribed to for detecting when to cancel the rendering if needed and unmount the app.
+    renderVisualEditing(signal, options)
+  })
 
   return () => {
     controller.abort()
-    // Handle React StrictMode, delay cleanup for a second in case it's a rerender
-    cleanup = window.setTimeout(() => {
-      if (root) {
-        root.unmount()
-        root = null
-      }
-      if (node) {
-        document.body.removeChild(node)
-        node = null
-      }
-    }, 1000)
   }
 }

packages/visual-editing/src/ui/renderVisualEditing.tsx

@@ -0,0 +1,52 @@
+/**
+ * The purpose of this file is to contain the logic for rendering the <VisualEditing />
+ * component in a way that is easy to lazy load for the `enableVisualEditing` function.
+ */
+
+import {createRoot, type Root} from 'react-dom/client'
+
+import {OVERLAY_ID} from '../constants'
+import type {VisualEditingOptions} from '../types'
+import {VisualEditing} from './VisualEditing'
+
+let node: HTMLElement | null = null
+let root: Root | null = null
+let cleanup: ReturnType<typeof setTimeout> | null = null
+
+export function renderVisualEditing(
+  signal: AbortSignal,
+  {history, refresh, zIndex}: VisualEditingOptions,
+): void {
+  // Cancel pending cleanups, this is useful to avoid overlays blinking as the parent app transition between URLs, or hot module reload is happening
+  if (cleanup) clearTimeout(cleanup)
+  // Setup a cleanup function listener right away, as the signal might abort in-between the next steps
+  signal.addEventListener('abort', () => {
+    // Handle React StrictMode, delay cleanup for a second in case it's a rerender
+    cleanup = setTimeout(() => {
+      if (root) {
+        root.unmount()
+        root = null
+      }
+      if (node) {
+        document.body.removeChild(node)
+        node = null
+      }
+    }, 1000)
+  })
+
+  if (!node) {
+    // eslint-disable-next-line no-warning-comments
+    // @TODO use 'sanity-visual-editing' instead of 'div'
+    node = document.createElement('div')
+    // eslint-disable-next-line no-warning-comments
+    // @TODO after the element is `sanity-visual-editing` instead of `div`, stop setting this ID
+    node.id = OVERLAY_ID
+    document.body.appendChild(node)
+  }
+
+  if (!root) {
+    root = createRoot(node)
+  }
+
+  root.render(<VisualEditing history={history} refresh={refresh} zIndex={zIndex} />)
+}