feat(core)!: track node statuses in core

Author: cngonzalez

packages/core/src/_exports/index.ts

@@ -32,7 +32,8 @@ export {
   releaseChannel,
 } from '../comlink/controller/comlinkControllerStore'
 export type {ComlinkNodeState} from '../comlink/node/comlinkNodeStore'
-export {getOrCreateNode, releaseNode} from '../comlink/node/comlinkNodeStore'
+export {getNodeState} from '../comlink/node/getNodeState'
+export {type NodeState} from '../comlink/node/getNodeState'
 export {type FrameMessage, type WindowMessage} from '../comlink/types'
 export {type AuthConfig, type AuthProvider} from '../config/authConfig'
 export {

packages/core/src/comlink/node/actions/getOrCreateNode.test.ts

@@ -12,6 +12,7 @@ vi.mock('@sanity/comlink', () => ({
   createNode: vi.fn(() => ({
     start: vi.fn(),
     stop: vi.fn(),
+    onStatus: vi.fn(),
   })),
 }))
 
@@ -32,9 +33,9 @@ describe('getOrCreateNode', () => {
   }
 
   beforeEach(() => {
-    mockNode = {start: vi.fn(), stop: vi.fn()}
+    mockNode = {start: vi.fn(), stop: vi.fn(), onStatus: vi.fn()}
     vi.mocked(comlink.createNode).mockReturnValue(mockNode as Node<WindowMessage, FrameMessage>)
-    state = createStoreState<ComlinkNodeState>({nodes: new Map()})
+    state = createStoreState<ComlinkNodeState>({nodes: new Map(), subscriptions: new Map()})
     vi.clearAllMocks()
   })
 
@@ -45,7 +46,7 @@ describe('getOrCreateNode', () => {
     expect(node.start).toHaveBeenCalled()
   })
 
-  it('sshould store the node in nodeStore', () => {
+  it('should store the node in nodeStore', () => {
     const node = getOrCreateNode({state, instance}, nodeConfig)
 
     expect(getOrCreateNode({state, instance}, nodeConfig)).toBe(node)
@@ -64,4 +65,43 @@ describe('getOrCreateNode', () => {
       ),
     ).toThrow('Node "test-node" already exists with different options')
   })
+
+  it('should subscribe to status changes and update state', () => {
+    let statusCallback: ((status: string) => void) | undefined
+    const statusUnsubMock = vi.fn()
+    mockNode.onStatus = vi.fn((cb) => {
+      statusCallback = cb
+      return statusUnsubMock
+    })
+
+    getOrCreateNode({state, instance}, nodeConfig)
+
+    expect(mockNode.onStatus).toHaveBeenCalled()
+    expect(typeof statusCallback).toBe('function')
+    expect(state.get().nodes.get(nodeConfig.name)?.statusUnsub).toBe(statusUnsubMock)
+
+    statusCallback?.('connected')
+
+    expect(state.get().nodes.get(nodeConfig.name)?.status).toBe('connected')
+  })
+
+  it('should not update state if node entry is missing when status changes', () => {
+    let statusCallback: ((status: string) => void) | undefined
+    const statusUnsubMock = vi.fn()
+    mockNode.onStatus = vi.fn((cb) => {
+      statusCallback = cb
+      return statusUnsubMock
+    })
+
+    getOrCreateNode({state, instance}, nodeConfig)
+
+    // Remove the node entry before triggering the status callback
+    state.get().nodes.delete(nodeConfig.name)
+
+    // Simulate a status change
+    statusCallback?.('connected')
+
+    // Assert: node entry is still missing, so no update occurred
+    expect(state.get().nodes.has(nodeConfig.name)).toBe(false)
+  })
 })

packages/core/src/comlink/node/actions/getOrCreateNode.ts

@@ -5,12 +5,6 @@ import {type StoreContext} from '../../../store/defineStore'
 import {type FrameMessage, type WindowMessage} from '../../types'
 import {type ComlinkNodeState} from '../comlinkNodeStore'
 
-/**
- * Retrieve or create a node to be used for communication between
- * an application and the controller -- specifically, a node should
- * be created within a frame / window to communicate with the controller.
- * @public
- */
 export const getOrCreateNode = (
   {state}: StoreContext<ComlinkNodeState>,
   options: NodeInput,
@@ -24,21 +18,37 @@ export const getOrCreateNode = (
       throw new Error(`Node "${options.name}" already exists with different options`)
     }
 
-    state.set('incrementNodeRefCount', {
-      nodes: new Map(nodes).set(options.name, {
-        ...existing,
-        refCount: existing.refCount + 1,
-      }),
-    })
-
     existing.node.start()
     return existing.node
   }
 
   const node: Node<WindowMessage, FrameMessage> = createNode(options)
   node.start()
 
-  nodes.set(options.name, {node, options, refCount: 1})
+  // Subscribe to status changes
+  const statusUnsub = node.onStatus((status) => {
+    const currentNodes = state.get().nodes
+    const currentEntry = currentNodes.get(options.name)
+    if (!currentEntry) return
+    const updatedEntry = {
+      ...currentEntry,
+      status,
+      statusUnsub: currentEntry.statusUnsub,
+    }
+    state.set('updateNodeStatus', {
+      nodes: new Map(currentNodes).set(options.name, updatedEntry),
+    })
+  })
+
+  // Set up initial entry with status, error, and statusUnsub
+  const entry = {
+    node,
+    options,
+    status: 'idle' as const,
+    statusUnsub,
+  }
+
+  nodes.set(options.name, entry)
 
   state.set('createNode', {nodes})
 

packages/core/src/comlink/node/actions/releaseNode.test.ts

@@ -4,7 +4,7 @@ import {beforeEach, describe, expect, it, vi} from 'vitest'
 import {createSanityInstance, type SanityInstance} from '../../../store/createSanityInstance'
 import {createStoreState} from '../../../store/createStoreState'
 import {type FrameMessage, type WindowMessage} from '../../types'
-import {type ComlinkNodeState, type NodeEntry} from '../comlinkNodeStore'
+import {type ComlinkNodeState} from '../comlinkNodeStore'
 import {releaseNode} from './releaseNode'
 
 const nodeConfig = {
@@ -18,15 +18,16 @@ describe('releaseNode', () => {
   let mockNode: Partial<Node<WindowMessage, FrameMessage>> & {
     start: ReturnType<typeof vi.fn>
     stop: ReturnType<typeof vi.fn>
+    onStatus: ReturnType<typeof vi.fn>
   }
 
   beforeEach(() => {
     instance = createSanityInstance({
       projectId: 'test-project-id',
       dataset: 'test-dataset',
     })
-    mockNode = {start: vi.fn(), stop: vi.fn()}
-    state = createStoreState<ComlinkNodeState>({nodes: new Map()})
+    mockNode = {start: vi.fn(), stop: vi.fn(), onStatus: vi.fn()}
+    state = createStoreState<ComlinkNodeState>({nodes: new Map(), subscriptions: new Map()})
     vi.clearAllMocks()
   })
 
@@ -40,7 +41,6 @@ describe('releaseNode', () => {
     nodes.set('test-node', {
       node: mockNode as Node<WindowMessage, FrameMessage>,
       options: nodeConfig,
-      refCount: 1,
     })
     state.set('setup', {nodes})
 
@@ -54,83 +54,18 @@ describe('releaseNode', () => {
     expect(state.get().nodes.has('test-node')).toBe(false)
   })
 
-  it('should not stop the node if refCount is still above 0', () => {
-    // Create a node twice to increment refCount
+  it('should call statusUnsub if present when releasing node', () => {
+    const statusUnsub = vi.fn()
     const nodes = new Map()
     nodes.set('test-node', {
       node: mockNode as Node<WindowMessage, FrameMessage>,
       options: nodeConfig,
-      refCount: 2,
+      statusUnsub,
     })
     state.set('setup', {nodes})
 
-    // Release once
     releaseNode({state, instance}, 'test-node')
 
-    // Node should not be stopped
-    expect(mockNode.stop).not.toHaveBeenCalled()
-
-    // Verify refCount is 1
-    const nodeEntry = state.get().nodes.get('test-node') as NodeEntry
-    expect(nodeEntry?.refCount).toBe(1)
-  })
-
-  it('should handle multiple releases gracefully', () => {
-    // Set up a node in the state
-    const nodes = new Map()
-    nodes.set('test-node', {
-      node: mockNode as Node<WindowMessage, FrameMessage>,
-      options: nodeConfig,
-      refCount: 1,
-    })
-    state.set('setup', {nodes})
-
-    // Release multiple times
-    releaseNode({state, instance}, 'test-node')
-    releaseNode({state, instance}, 'test-node')
-    releaseNode({state, instance}, 'test-node')
-
-    // Verify node is removed after first release
-    expect(state.get().nodes.has('test-node')).toBe(false)
-    // Stop should be called exactly once
-    expect(mockNode.stop).toHaveBeenCalledTimes(1)
-  })
-
-  it('should handle releasing non-existent nodes', () => {
-    // Should not throw when releasing non-existent node
-    expect(() => releaseNode({state, instance}, 'non-existent')).not.toThrow()
-  })
-
-  it('should maintain correct state after complex operations', () => {
-    // Set up a node with refCount = 3
-    const nodes = new Map()
-    nodes.set('test-node', {
-      node: mockNode as Node<WindowMessage, FrameMessage>,
-      options: nodeConfig,
-      refCount: 3,
-    })
-    state.set('setup', {nodes})
-
-    // Initial refCount should be 3
-    let nodeEntry = state.get().nodes.get('test-node') as NodeEntry
-    expect(nodeEntry?.refCount).toBe(3)
-
-    // Release twice
-    releaseNode({state, instance}, 'test-node')
-    releaseNode({state, instance}, 'test-node')
-
-    nodeEntry = state.get().nodes.get('test-node') as NodeEntry
-    expect(nodeEntry?.refCount).toBe(1)
-
-    // Verify node hasn't been stopped yet
-    expect(mockNode.stop).not.toHaveBeenCalled()
-
-    // Release final reference
-    releaseNode({state, instance}, 'test-node')
-
-    // Verify node was stopped
-    expect(mockNode.stop).toHaveBeenCalled()
-
-    expect(state.get().nodes.has('test-node')).toBe(false)
+    expect(statusUnsub).toHaveBeenCalled()
   })
 })

packages/core/src/comlink/node/actions/releaseNode.ts

@@ -1,30 +1,17 @@
 import {type StoreContext} from '../../../store/defineStore'
 import {type ComlinkNodeState} from '../comlinkNodeStore'
 
-/**
- * Release a node that was previously created with getOrCreateNode.
- * @public
- */
 export const releaseNode = ({state}: StoreContext<ComlinkNodeState>, name: string): void => {
   const nodes = state.get().nodes
   const existing = nodes.get(name)
 
   if (existing) {
-    const newRefCount = existing.refCount - 1
-
-    if (newRefCount <= 0) {
-      existing.node.stop()
-      nodes.delete(name)
-      state.set('removeNode', {nodes})
-      return
+    if (existing.statusUnsub) {
+      existing.statusUnsub()
     }
-
-    state.set('decrementNodeRefCount', {
-      nodes: new Map(nodes).set(name, {
-        ...existing,
-        refCount: newRefCount,
-      }),
-    })
+    existing.node.stop()
+    nodes.delete(name)
+    state.set('removeNode', {nodes})
     return
   }
 }

packages/core/src/comlink/node/comlinkNodeStore.ts

@@ -1,4 +1,4 @@
-import {type Node, type NodeInput} from '@sanity/comlink'
+import {type Node, type NodeInput, type Status} from '@sanity/comlink'
 
 import {bindActionGlobally} from '../../store/createActionBinder'
 import {defineStore} from '../../store/defineStore'
@@ -14,8 +14,9 @@ export interface NodeEntry {
   node: Node<WindowMessage, FrameMessage>
   // we store options to ensure that channels remain as unique / consistent as possible
   options: NodeInput
-  // we store refCount to ensure nodes are running only as long as they are in use
-  refCount: number
+  // status of the node connection
+  status: Status
+  statusUnsub?: () => void
 }
 
 /**
@@ -24,12 +25,15 @@ export interface NodeEntry {
  */
 export interface ComlinkNodeState {
   nodes: Map<string, NodeEntry>
+  // Map of node name to set of active subscriber symbols
+  subscriptions: Map<string, Set<symbol>>
 }
 
 export const comlinkNodeStore = defineStore<ComlinkNodeState>({
   name: 'nodeStore',
   getInitialState: () => ({
     nodes: new Map(),
+    subscriptions: new Map(),
   }),
 
   initialize({state}) {