fix(visual-editing): execute fetches after each node connect event

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

@@ -1,19 +1,11 @@
-import {
-  createCompatibilityActors,
-  type VisualEditingControllerMsg,
-  type VisualEditingNodeMsg,
-} from '@repo/visual-editing-helpers'
-import {createNode, createNodeMachine} from '@sanity/comlink'
-import {useEffect, useState, type FunctionComponent} from 'react'
-import {createActor} from 'xstate'
-import type {VisualEditingNode, VisualEditingOptions} from '../types'
-import {createDatasetMutator} from './comlink'
+import {type FunctionComponent} from 'react'
+import type {VisualEditingOptions} from '../types'
 import {History} from './History'
 import {Meta} from './Meta'
-import {setActor} from './optimistic-state/context'
-import {createSharedListener} from './optimistic-state/machines/createSharedListener'
 import {Overlays} from './Overlays'
 import {Refresh} from './Refresh'
+import {useComlink} from './useComlink'
+import {useDatasetMutator} from './useDatasetMutator'
 
 /**
  * @public
@@ -22,56 +14,8 @@ export const VisualEditing: FunctionComponent<VisualEditingOptions> = (props) =>
   const {components, history, refresh, zIndex} = props
   const inFrame = window.self !== window.top || window.opener
 
-  const [comlink, setComlink] = useState<VisualEditingNode | undefined>(undefined)
-
-  useEffect(() => {
-    if (!inFrame) return
-    const comlink = createNode<VisualEditingNodeMsg, VisualEditingControllerMsg>(
-      {
-        name: 'visual-editing',
-        connectTo: 'presentation',
-      },
-      createNodeMachine<VisualEditingNodeMsg, VisualEditingControllerMsg>().provide({
-        actors: createCompatibilityActors<VisualEditingNodeMsg>(),
-      }),
-    )
-    setComlink(comlink)
-
-    const listener = createSharedListener(comlink)
-    const datasetMutator = createDatasetMutator(comlink)
-    const actor = createActor(datasetMutator, {
-      // @ts-expect-error @todo
-      input: {client: {withConfig: () => {}}, sharedListener: listener},
-    })
-
-    // Fetch features to determine if optimistic updates are supported
-    const controller = new AbortController()
-    comlink
-      .fetch('visual-editing/features', undefined, {
-        signal: controller.signal,
-        suppressWarnings: true,
-      })
-      .then((data) => {
-        if (data.features['optimistic']) {
-          setActor(actor)
-        }
-      })
-      .catch(() => {
-        // eslint-disable-next-line no-console
-        console.warn(
-          '[@sanity/visual-editing] Package version mismatch detected: Please update your Sanity studio to prevent potential compatibility issues.',
-        )
-      })
-
-    actor.start()
-    comlink.start()
-
-    return () => {
-      controller.abort()
-      actor.stop()
-      comlink.stop()
-    }
-  }, [inFrame])
+  const comlink = useComlink(inFrame)
+  useDatasetMutator(comlink)
 
   return (
     <>

packages/visual-editing/src/ui/schema/SchemaProvider.tsx

@@ -90,9 +90,17 @@ export const SchemaProvider: FunctionComponent<
     [comlink],
   )
   useEffect(() => {
-    const controller = new AbortController()
-    fetchSchema(controller.signal)
-    return () => controller.abort()
+    if (!comlink) return
+
+    const schemaFetch = new AbortController()
+    const unsub = comlink.onStatus(() => {
+      fetchSchema(schemaFetch.signal)
+    }, 'connected')
+
+    return () => {
+      schemaFetch.abort()
+      unsub()
+    }
   }, [fetchSchema])
 
   const reportedPathsRef = useRef<UnresolvedPath[]>([])

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

@@ -11,10 +11,11 @@ import type {VisualEditingNode} from '../types'
  * Hook for maintaining a channel between overlays and the presentation tool
  * @internal
  */
-export function useComlink(): VisualEditingNode | undefined {
+export function useComlink(active: boolean = true): VisualEditingNode | undefined {
   const [node, setNode] = useState<VisualEditingNode>()
 
   useEffect(() => {
+    if (!active) return
     const instance = createNode<VisualEditingNodeMsg, VisualEditingControllerMsg>(
       {
         name: 'visual-editing',
@@ -26,8 +27,13 @@ export function useComlink(): VisualEditingNode | undefined {
     )
 
     setNode(instance)
-    return instance.start()
-  }, [])
+    const stop = instance.start()
+
+    return () => {
+      stop()
+      setNode(undefined)
+    }
+  }, [active])
 
   return node
 }

packages/visual-editing/src/ui/useDatasetMutator.ts

@@ -0,0 +1,71 @@
+import {
+  createCompatibilityActors,
+  type VisualEditingControllerMsg,
+  type VisualEditingNodeMsg,
+} from '@repo/visual-editing-helpers'
+import {createNode, createNodeMachine} from '@sanity/comlink'
+import {useEffect, useState} from 'react'
+import {createActor, type ActorRefFrom} from 'xstate'
+import type {VisualEditingNode} from '../types'
+import {createDatasetMutator} from './comlink'
+import {setActor, type MutatorActor} from './optimistic-state/context'
+import {createSharedListener} from './optimistic-state/machines/createSharedListener'
+
+/**
+ * Hook for maintaining a channel between overlays and the presentation tool
+ * @internal
+ */
+export function useDatasetMutator(
+  comlink: VisualEditingNode | undefined,
+): MutatorActor | undefined {
+  const [mutator, setMutator] = useState<MutatorActor>()
+
+  useEffect(() => {
+    if (!comlink) return
+    const listener = createSharedListener(comlink)
+    const datasetMutator = createDatasetMutator(comlink)
+    const mutator = createActor(datasetMutator, {
+      // @ts-expect-error @todo
+      input: {client: {withConfig: () => {}}, sharedListener: listener},
+    })
+
+    setMutator(mutator)
+    mutator.start()
+
+    return () => {
+      mutator.stop()
+      setMutator(undefined)
+    }
+  }, [comlink])
+
+  useEffect(() => {
+    if (!comlink || !mutator) return
+    // Fetch features to determine if optimistic updates are supported
+    const featuresFetch = new AbortController()
+    const unsub = comlink.onStatus(() => {
+      comlink
+        .fetch('visual-editing/features', undefined, {
+          signal: featuresFetch.signal,
+          suppressWarnings: true,
+        })
+        .then((data) => {
+          if (data.features['optimistic']) {
+            setActor(mutator)
+          }
+        })
+        .catch(() => {
+          // eslint-disable-next-line no-console
+          console.warn(
+            '[@sanity/visual-editing] Package version mismatch detected: Please update your Sanity studio to prevent potential compatibility issues.',
+          )
+        })
+    }, 'connected')
+
+    return () => {
+      featuresFetch.abort()
+      unsub()
+    }
+  }, [mutator, comlink])
+
+  return mutator
+}