feat(react): use refs for nodes and channels (#341)

cngonzalez · web-flow · commit 4a0b7637463e · 2025-03-19T15:20:30.000-04:00
diff --git a/apps/kitchensink-react/src/Comlink/Framed.tsx b/apps/kitchensink-react/src/Comlink/Framed.tsx
@@ -1,16 +1,18 @@
-import {useWindowConnection} from '@sanity/sdk-react'
+import {type ComlinkStatus, useWindowConnection} from '@sanity/sdk-react'
 import {Box, Button, Card, Container, Label, Stack, Text, TextInput} from '@sanity/ui'
 import {type ReactElement, useState} from 'react'
 
 import {FromIFrameMessage, ToIFrameMessage} from './types'
 
 const Framed = (): ReactElement => {
   const [message, setMessage] = useState('')
+  const [status, setStatus] = useState<ComlinkStatus>('idle')
   const [receivedMessages, setReceivedMessages] = useState<string[]>([])
 
-  const {sendMessage, status} = useWindowConnection<FromIFrameMessage, ToIFrameMessage>({
+  const {sendMessage} = useWindowConnection<FromIFrameMessage, ToIFrameMessage>({
     name: 'frame',
     connectTo: 'main-app',
+    onStatus: setStatus,
     onMessage: {
       TO_IFRAME: (data: {message: string}) => {
         setReceivedMessages((prev) => [...prev, data.message])
diff --git a/apps/kitchensink-react/src/Comlink/ParentApp.tsx b/apps/kitchensink-react/src/Comlink/ParentApp.tsx
@@ -1,4 +1,4 @@
-import {useFrameConnection} from '@sanity/sdk-react'
+import {type ComlinkStatus, useFrameConnection} from '@sanity/sdk-react'
 import {Box, Button, Card, Label, Stack, Text, TextInput} from '@sanity/ui'
 import {ReactElement, useEffect, useRef, useState} from 'react'
 
@@ -8,15 +8,17 @@ import {FromIFrameMessage, ToIFrameMessage} from './types'
 const ParentApp = (): ReactElement => {
   const [selectedFrame, setSelectedFrame] = useState<number>(1)
   const [message, setMessage] = useState('')
+  const [status, setStatus] = useState<ComlinkStatus>('idle')
   const [receivedMessages, setReceivedMessages] = useState<Array<{from: string; message: string}>>(
     [],
   )
   const iframeRef = useRef<HTMLIFrameElement>(null)
 
-  const {sendMessage, connect, status} = useFrameConnection<ToIFrameMessage, FromIFrameMessage>({
+  const {sendMessage, connect} = useFrameConnection<ToIFrameMessage, FromIFrameMessage>({
     name: 'main-app',
     connectTo: 'frame',
     targetOrigin: '*',
+    onStatus: setStatus,
     onMessage: {
       FROM_IFRAME: (data: {message: string}) => {
         setReceivedMessages((prev) => [
diff --git a/packages/react/src/_exports/index.ts b/packages/react/src/_exports/index.ts
@@ -59,5 +59,6 @@ export {useQuery} from '../hooks/query/useQuery'
 export {useUsers, type UseUsersParams, type UseUsersResult} from '../hooks/users/useUsers'
 export {REACT_SDK_VERSION} from '../version'
 export {type DatasetsResponse, type SanityProject, type SanityProjectMember} from '@sanity/client'
+export {type Status as ComlinkStatus} from '@sanity/comlink'
 export {type CurrentUser, type DocumentHandle} from '@sanity/sdk'
 export {type SanityDocument, type SortOrderingItem} from '@sanity/types'
diff --git a/packages/react/src/hooks/comlink/useFrameConnection.test.tsx b/packages/react/src/hooks/comlink/useFrameConnection.test.tsx
@@ -61,32 +61,42 @@ describe('useFrameController', () => {
     vi.mocked(getOrCreateController).mockReturnValue(controller)
   })
 
-  it('should initialize with idle status', () => {
-    const {result} = renderHook(() =>
+  it('should call onStatus callback when status changes', () => {
+    const onStatusMock = vi.fn()
+    renderHook(() =>
       useFrameConnection({
         name: 'test',
         connectTo: 'iframe',
         targetOrigin: '*',
+        onStatus: onStatusMock,
       }),
     )
 
-    expect(result.current.status).toBe('idle')
+    act(() => {
+      statusCallback?.({status: 'connected', connection: 'test'})
+    })
+    expect(onStatusMock).toHaveBeenCalledWith('connected')
+
+    act(() => {
+      statusCallback?.({status: 'disconnected', connection: 'test'})
+    })
+    expect(onStatusMock).toHaveBeenCalledWith('disconnected')
   })
 
-  it('should update status to connected when node connects', async () => {
-    const {result} = renderHook(() =>
+  it('should not throw if onStatus is not provided', () => {
+    renderHook(() =>
       useFrameConnection({
         name: 'test',
         connectTo: 'iframe',
         targetOrigin: '*',
       }),
     )
 
-    expect(result.current.status).toBe('idle')
-    act(() => {
-      statusCallback?.({status: 'connected', connection: 'test'})
-    })
-    expect(result.current.status).toBe('connected')
+    expect(() => {
+      act(() => {
+        statusCallback?.({status: 'connected', connection: 'test'})
+      })
+    }).not.toThrow()
   })
 
   it('should register and execute message handlers', () => {
diff --git a/packages/react/src/hooks/comlink/useFrameConnection.ts b/packages/react/src/hooks/comlink/useFrameConnection.ts
@@ -1,12 +1,12 @@
-import {type Status} from '@sanity/comlink'
+import {type ChannelInstance, type Controller, type Status} from '@sanity/comlink'
 import {
   type FrameMessage,
   getOrCreateChannel,
   getOrCreateController,
   releaseChannel,
   type WindowMessage,
 } from '@sanity/sdk'
-import {useCallback, useEffect, useMemo, useState} from 'react'
+import {useCallback, useEffect, useRef} from 'react'
 
 import {useSanityInstance} from '../context/useSanityInstance'
 
@@ -28,6 +28,7 @@ export interface UseFrameConnectionOptions<TWindowMessage extends WindowMessage>
     [K in TWindowMessage['type']]: (data: Extract<TWindowMessage, {type: K}>['data']) => void
   }
   heartbeat?: boolean
+  onStatus?: (status: Status) => void
 }
 
 /**
@@ -40,7 +41,6 @@ export interface FrameConnection<TFrameMessage extends FrameMessage> {
       ? [type: T]
       : [type: T, data: Extract<TFrameMessage, {type: T}>['data']]
   ) => void
-  status: Status
 }
 
 /**
@@ -50,81 +50,58 @@ export function useFrameConnection<
   TFrameMessage extends FrameMessage,
   TWindowMessage extends WindowMessage,
 >(options: UseFrameConnectionOptions<TWindowMessage>): FrameConnection<TFrameMessage> {
-  const {onMessage, targetOrigin, name, connectTo, heartbeat} = options
+  const {onMessage, targetOrigin, name, connectTo, heartbeat, onStatus} = options
   const instance = useSanityInstance()
-  const [status, setStatus] = useState<Status>('idle')
-
-  const controller = useMemo(
-    () => getOrCreateController(instance, targetOrigin),
-    [instance, targetOrigin],
-  )
-
-  const channel = useMemo(
-    () =>
-      getOrCreateChannel(instance, {
-        name,
-        connectTo,
-        heartbeat,
-      }),
-    [instance, name, connectTo, heartbeat],
-  )
+  const controllerRef = useRef<Controller | null>(null)
+  const channelRef = useRef<ChannelInstance<TFrameMessage, TWindowMessage> | null>(null)
 
   useEffect(() => {
-    if (!channel) return
+    const controller = getOrCreateController(instance, targetOrigin)
+    const channel = getOrCreateChannel(instance, {name, connectTo, heartbeat})
+    controllerRef.current = controller
+    channelRef.current = channel
 
-    const unsubscribe = channel.onStatus((event) => {
-      setStatus(event.status)
+    channel.onStatus((event) => {
+      onStatus?.(event.status)
     })
 
-    return unsubscribe
-  }, [channel])
-
-  useEffect(() => {
-    if (!channel || !onMessage) return
-
-    const unsubscribers: Array<() => void> = []
+    const messageUnsubscribers: Array<() => void> = []
 
-    Object.entries(onMessage).forEach(([type, handler]) => {
-      // type assertion, but we've already constrained onMessage to have the correct handler type
-      const unsubscribe = channel.on(type, handler as FrameMessageHandler<TWindowMessage>)
-      unsubscribers.push(unsubscribe)
-    })
+    if (onMessage) {
+      Object.entries(onMessage).forEach(([type, handler]) => {
+        const unsubscribe = channel.on(type, handler as FrameMessageHandler<TWindowMessage>)
+        messageUnsubscribers.push(unsubscribe)
+      })
+    }
 
     return () => {
-      unsubscribers.forEach((unsub) => unsub())
+      // Clean up all subscriptions and stop controller/channel
+      messageUnsubscribers.forEach((unsub) => unsub())
+      releaseChannel(instance, name)
+      channelRef.current = null
+      controllerRef.current = null
     }
-  }, [channel, onMessage])
+  }, [targetOrigin, name, connectTo, heartbeat, onMessage, instance, onStatus])
 
-  const connect = useCallback(
-    (frameWindow: Window) => {
-      const removeTarget = controller?.addTarget(frameWindow)
-      return () => {
-        removeTarget?.()
-      }
-    },
-    [controller],
-  )
+  const connect = useCallback((frameWindow: Window) => {
+    const removeTarget = controllerRef.current?.addTarget(frameWindow)
+    return () => {
+      removeTarget?.()
+    }
+  }, [])
 
   const sendMessage = useCallback(
     <T extends TFrameMessage['type']>(
       type: T,
       data?: Extract<TFrameMessage, {type: T}>['data'],
     ) => {
-      channel?.post(type, data)
+      channelRef.current?.post(type, data)
     },
-    [channel],
+    [],
   )
 
-  // cleanup channel on unmount
-  useEffect(() => {
-    return () => {
-      releaseChannel(instance, name)
-    }
-  }, [name, instance])
-
   return {
     connect,
     sendMessage,
-    status,
   }
 }
diff --git a/packages/react/src/hooks/comlink/useManageFavorite.ts b/packages/react/src/hooks/comlink/useManageFavorite.ts
@@ -1,3 +1,4 @@
+import {type Status} from '@sanity/comlink'
 import {type Events, SDK_CHANNEL_NAME, SDK_NODE_NAME} from '@sanity/message-protocol'
 import {type DocumentHandle, type FrameMessage} from '@sanity/sdk'
 import {useCallback, useState} from 'react'
@@ -48,9 +49,11 @@ interface ManageFavorite {
  */
 export function useManageFavorite({_id, _type}: DocumentHandle): ManageFavorite {
   const [isFavorited, setIsFavorited] = useState(false) // should load this from a comlink fetch
-  const {sendMessage, status} = useWindowConnection<Events.FavoriteMessage, FrameMessage>({
+  const [status, setStatus] = useState<Status>('idle')
+  const {sendMessage} = useWindowConnection<Events.FavoriteMessage, FrameMessage>({
     name: SDK_NODE_NAME,
     connectTo: SDK_CHANNEL_NAME,
+    onStatus: setStatus,
   })
 
   const handleFavoriteAction = useCallback(
diff --git a/packages/react/src/hooks/comlink/useRecordDocumentHistoryEvent.ts b/packages/react/src/hooks/comlink/useRecordDocumentHistoryEvent.ts
@@ -1,6 +1,7 @@
+import {type Status} from '@sanity/comlink'
 import {type Events, SDK_CHANNEL_NAME, SDK_NODE_NAME} from '@sanity/message-protocol'
 import {type DocumentHandle, type FrameMessage} from '@sanity/sdk'
-import {useCallback} from 'react'
+import {useCallback, useState} from 'react'
 
 import {useWindowConnection} from './useWindowConnection'
 
@@ -42,9 +43,11 @@ export function useRecordDocumentHistoryEvent({
   _id,
   _type,
 }: DocumentHandle): DocumentInteractionHistory {
-  const {sendMessage, status} = useWindowConnection<Events.HistoryMessage, FrameMessage>({
+  const [status, setStatus] = useState<Status>('idle')
+  const {sendMessage} = useWindowConnection<Events.HistoryMessage, FrameMessage>({
     name: SDK_NODE_NAME,
     connectTo: SDK_CHANNEL_NAME,
+    onStatus: setStatus,
   })
 
   const recordEvent = useCallback(
diff --git a/packages/react/src/hooks/comlink/useWindowConnection.test.ts b/packages/react/src/hooks/comlink/useWindowConnection.test.ts
@@ -49,30 +49,40 @@ describe('useWindowConnection', () => {
     vi.mocked(getOrCreateNode).mockReturnValue(node as unknown as Node<Message, Message>)
   })
 
-  it('should initialize with idle status', () => {
-    const {result} = renderHook(() =>
+  it('should call onStatus callback when status changes', () => {
+    const onStatusMock = vi.fn()
+    renderHook(() =>
       useWindowConnection<TestMessages, TestMessages>({
         name: 'test',
         connectTo: 'window',
+        onStatus: onStatusMock,
       }),
     )
 
-    expect(result.current.status).toBe('idle')
+    act(() => {
+      statusCallback?.('connected')
+    })
+    expect(onStatusMock).toHaveBeenCalledWith('connected')
+
+    act(() => {
+      statusCallback?.('disconnected')
+    })
+    expect(onStatusMock).toHaveBeenCalledWith('disconnected')
   })
 
-  it('should update status to connected when node connects', async () => {
-    const {result} = renderHook(() =>
+  it('should not throw if onStatus is not provided', () => {
+    renderHook(() =>
       useWindowConnection<TestMessages, TestMessages>({
         name: 'test',
         connectTo: 'window',
       }),
     )
 
-    expect(result.current.status).toBe('idle')
-    act(() => {
-      statusCallback?.('connected')
-    })
-    expect(result.current.status).toBe('connected')
+    expect(() => {
+      act(() => {
+        statusCallback?.('connected')
+      })
+    }).not.toThrow()
   })
 
   it('should register message handlers', () => {
diff --git a/packages/react/src/hooks/comlink/useWindowConnection.ts b/packages/react/src/hooks/comlink/useWindowConnection.ts
