@jr200-labs/xstate-nats 0.6.0

package/src/machines/root.ts

@@ -0,0 +1,286 @@
+import { ConnectionOptions, NatsConnection } from '@nats-io/nats-core'
+import { assign, sendTo, setup } from 'xstate'
+import { kvManagerLogic, ExternalEvents as KvExternalEvents } from './kv'
+import { subjectManagerLogic, ExternalEvents as SubjectExternalEvents } from './subject'
+import { connectToNats, disconnectNats } from '../actions/connection'
+import { type AuthConfig } from '../actions/types'
+import { InternalStatusEvents as NatsStatusEvents } from '../actions/connection'
+
+export interface NatsConnectionConfig {
+  opts: ConnectionOptions
+  auth?: AuthConfig
+  maxRetries: number
+}
+
+export interface Context {
+  connection: NatsConnection | null
+  error?: Error
+  natsConfig?: NatsConnectionConfig
+  retries: number
+  subjectManagerReady: boolean
+  kvManagerReady: boolean
+}
+
+// internal events and events from nats connection
+type InternalEvents =
+  | { type: 'CONNECTED'; connection: NatsConnection }
+  | { type: 'DISCONNECTED' }
+  | { type: 'FAIL'; error: Error }
+  | { type: 'RECONNECT' }
+  | { type: 'CLOSE' }
+  | NatsStatusEvents
+
+// events which can be sent to the machine from the user
+export type ExternalEvents =
+  | { type: 'CONFIGURE'; config: NatsConnectionConfig }
+  | { type: 'CONNECT' }
+  | { type: 'DISCONNECT' }
+  | { type: 'RESET' }
+  | SubjectExternalEvents
+  | KvExternalEvents
+
+type Events = InternalEvents | ExternalEvents
+
+export const natsMachine = setup({
+  types: {
+    context: {} as Context,
+    events: {} as Events,
+  },
+  actions: {
+    doReset: assign({
+      natsConfig: (_) => undefined,
+      connection: (_) => null,
+      error: (_) => undefined,
+      retries: (_) => 0,
+      subjectManagerReady: (_) => false,
+      kvManagerReady: (_) => false,
+    }),
+  },
+  guards: {
+    allManagersReady: ({ context }) => {
+      const hasValidConnection = context.connection !== null
+      const subjectReady = context.subjectManagerReady
+      const kvReady = context.kvManagerReady
+      return hasValidConnection && subjectReady && kvReady
+    },
+  },
+  actors: {
+    connectToNats: connectToNats,
+    disconnectNats: disconnectNats,
+    subject: subjectManagerLogic,
+    kv: kvManagerLogic,
+  },
+}).createMachine({
+  /** @xstate-layout N4IgpgJg5mDOIC5QAoC2BDAxgCwJYDswBKAOnwHsAXAfU3PwDNcoBXAJ0gGIBhAeQDkAYgEkA4gFUASgFEA2gAYAuolAAHcrFyVc9FSAAeiABwBOeSQBsAdgCMRgMwAmRxftWALPYCsAGhABPRHd3c3cvext7IwsveRCrAF8EvzQsPEJSOkZmdi4+fn5pbgAVBWUkEHVNbV0KwwQbGw9LJysLExNHSPkjP0CEdzaSMwsXaPt5LxMbeQsklIwcAmISLMJMbXwoTgh6MBICADdyAGt91KWM1fp1zagEI-JMdBr8MrK9Kq0dfD16m3CVhIUUcVis8isJjCjlmfUQ4XcLUcRiakXcJicc2SIAu6RWazAGwI2zAbDY5DYJFUABsXgwKagSLjlpkboS7g98Mdnq93kpPhpvrVQP9AcCjKDwZDobCAkF3DYSE0TNZ2p1ptF3PMcYs8az8Lc8gAZXgAZTk-IqX1ef0QNlGQK60TMEUGVi8NjhCFciJmGIxRmlYJs2uZV0w1MFWx2ewOXNO511LNWkc0W053JePz55TUgptdUQyK8JDiFiMXhC8hco3k9i9mJIVk17Rs03LYPsoaT4dTxM4pPJlJpdIZTJ7+L76cePOzSg+VvzP1tCGLpfc5cr8mrFlr9blDUrJFi2-kKs6NlB3bSyYjGi4MnNpUteeqy8LCCm9mB3kDtjBPRhA21bAme7bupukzXpcKyDhSnCPtIz65pUS7CgYdoWI0x4RG2IRNJ4jjuF6jgqk2LheJCbQRPYWrahQEBwHoYbEAKb7ofUAC0Fhetx0F6mQVC0PQTCsBwEBsUKvwflYXTkTE9qdq49gmA2jgls2FgbhuowmB4-G3iJOTiZJBYinacSKtYKoWHW7pGIGPEHm6JAhKMXikXYNieHRCw3uGbJElspnvuZDQRI4JAokRHkqjCG69M54TDLMXT2CpERmIk2IsfqhoSYu7HSWFTQqcMTrpVhKkool-SVhYx4uKlMylXYBm9lGUAhRxQSOF69omEq8gAruESTB5kTtZO94Fa+UkruiiqpdCHSOLhTn9E4GktjppH6TlE6kHBbDdcVGENFYZUXg49hVSYNVqeYMIRW07TqR42VJEAA */
+  initial: 'not_configured',
+  context: {
+    connection: null,
+    error: undefined,
+    natsConfig: undefined,
+    retries: 0,
+    subjectManagerReady: false,
+    kvManagerReady: false,
+  },
+  invoke: [
+    { src: 'subject', id: 'subject', systemId: 'subject' },
+    { src: 'kv', id: 'kv' },
+  ],
+  on: {
+    'NATS_CONNECTION.*': {
+      actions: [
+        ({ event }: { event: any }) => {
+          console.log('root received NATS status event', event)
+        },
+      ],
+    },
+  },
+  states: {
+    not_configured: {
+      on: {
+        CONFIGURE: {
+          target: 'configured',
+          actions: [
+            assign({
+              natsConfig: ({ event }) => event.config,
+            }),
+          ],
+          '*': {
+            actions: [
+              ({ event }: { event: any }) => {
+                console.error('root not_configured received unexpected event', event)
+              },
+            ],
+          },
+        },
+      },
+    },
+    configured: {
+      on: {
+        CONNECT: {
+          target: 'connecting',
+        },
+        RESET: {
+          target: 'not_configured',
+          actions: ['doReset'],
+        },
+        '*': {
+          actions: [
+            ({ event }: { event: any }) => {
+              console.error('root configured received unexpected event', event)
+            },
+          ],
+        },
+      },
+    },
+    connecting: {
+      invoke: [
+        {
+          src: 'connectToNats',
+          input: ({ context }) => ({
+            opts: context.natsConfig!.opts,
+            auth: context.natsConfig!.auth,
+          }),
+          onDone: {
+            target: 'initialise_managers',
+            actions: [
+              assign({
+                connection: ({ event }) => event.output,
+                retries: (_) => 0,
+              }),
+            ],
+          },
+          onError: {
+            target: 'error',
+            actions: assign({
+              error: ({ event }) => event.error as Error,
+              retries: ({ context }) => context.retries + 1,
+            }),
+          },
+        },
+      ],
+    },
+    initialise_managers: {
+      entry: [
+        sendTo('subject', ({ context }) => ({
+          type: 'SUBJECT.CONNECT',
+          connection: context.connection!,
+        })),
+        sendTo('kv', ({ context }) => ({ type: 'KV.CONNECT', connection: context.connection! })),
+      ],
+      on: {
+        'SUBJECT.CONNECTED': {
+          actions: [assign({ subjectManagerReady: (_) => true })],
+        },
+        'KV.CONNECTED': {
+          actions: [assign({ kvManagerReady: (_) => true })],
+        },
+      },
+      always: [
+        {
+          guard: 'allManagersReady',
+          target: 'connected',
+        },
+      ],
+    },
+    connected: {
+      entry: [
+        (event) => {
+          console.log('CONNECTED', event.context.connection?.getServer())
+        },
+      ],
+      on: {
+        DISCONNECT: {
+          target: 'closing',
+        },
+        CLOSE: {
+          target: 'closing',
+        },
+        'SUBJECT.*': {
+          actions: [
+            ({ event }: any) => {
+              console.log('xstat-nats forwarding subject event', event)
+            },
+            sendTo(
+              'subject',
+              ({ event, context }: { event: SubjectExternalEvents; context: Context }) => {
+                return { ...event, connection: context.connection }
+              },
+            ),
+          ],
+        },
+        'KV.*': {
+          actions: [
+            ({ event }: any) => {
+              console.log('xstat-nats forwarding kv event', event)
+            },
+            sendTo('kv', ({ event, context }: { event: KvExternalEvents; context: Context }) => {
+              return { ...event, connection: context.connection }
+            }),
+          ],
+        },
+        '*': {
+          actions: [
+            ({ event }: { event: any }) => {
+              console.error('root connected received unexpected event', event)
+            },
+          ],
+        },
+      },
+    },
+    closing: {
+      entry: [
+        (event) => {
+          console.log('CLOSING', event.context.connection?.getServer())
+        },
+        sendTo('subject', { type: 'SUBJECT.DISCONNECTED' }),
+        sendTo('kv', { type: 'KV.DISCONNECTED' }),
+      ],
+      invoke: {
+        src: 'disconnectNats',
+        input: ({ context }) => ({ connection: context.connection }),
+        onDone: {
+          target: 'closed',
+        },
+        onError: {
+          target: 'error',
+          actions: assign({
+            error: ({ event }) => event.error as Error,
+          }),
+        },
+      },
+    },
+    closed: {
+      on: {
+        RESET: {
+          target: 'not_configured',
+          actions: ['doReset'],
+        },
+        CONNECT: {
+          target: 'connecting',
+        },
+      },
+      '*': {
+        actions: [
+          ({ event }: { event: any }) => {
+            console.error('root closing received unexpected event', event)
+          },
+        ],
+      },
+    },
+    error: {
+      on: {
+        RESET: {
+          target: 'not_configured',
+          actions: ['doReset'],
+        },
+      },
+      '*': {
+        actions: [
+          ({ event }: { event: any }) => {
+            console.error('root error received unexpected event', event)
+          },
+        ],
+      },
+    },
+  },
+})

package/src/machines/subject.test.ts

@@ -0,0 +1,272 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest'
+import { createActor, createMachine, assign, fromPromise, sendTo, setup } from 'xstate'
+import { subjectManagerLogic } from './subject'
+
+vi.mock('@nats-io/nats-core', () => ({
+  wsconnect: vi.fn(),
+}))
+
+function createMockConnection() {
+  return {
+    subscribe: vi.fn().mockReturnValue({
+      unsubscribe: vi.fn(),
+      [Symbol.asyncIterator]: () => ({
+        next: () => new Promise(() => {}),
+      }),
+    }),
+    request: vi.fn().mockResolvedValue({ json: () => ({}), string: () => '{}' }),
+    publish: vi.fn(),
+  } as any
+}
+
+// Create a parent machine that wraps the subject manager so sendParent works
+function createParentMachine() {
+  return setup({
+    types: {
+      context: {} as { childState: string },
+      events: {} as any,
+    },
+    actors: {
+      subject: subjectManagerLogic,
+    },
+  }).createMachine({
+    initial: 'active',
+    context: { childState: '' },
+    invoke: {
+      src: 'subject',
+      id: 'subject',
+    },
+    on: {
+      // Forward all events to the subject child
+      'SUBJECT.*': {
+        actions: sendTo('subject', ({ event, context }: any) => {
+          return { ...event, connection: event.connection }
+        }),
+      },
+      'SUBJECT.CONNECTED': {
+        actions: assign({ childState: 'connected' }),
+      },
+      'SUBJECT.DISCONNECTED': {
+        actions: assign({ childState: 'disconnected' }),
+      },
+    },
+    states: {
+      active: {},
+    },
+  })
+}
+
+describe('subjectManagerLogic', () => {
+  beforeEach(() => {
+    vi.spyOn(console, 'log').mockImplementation(() => {})
+    vi.spyOn(console, 'error').mockImplementation(() => {})
+  })
+
+  it('should start in subject_idle state', () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+    const childSnap = parentActor.getSnapshot().children.subject?.getSnapshot()
+    expect(childSnap?.value).toBe('subject_idle')
+    parentActor.stop()
+  })
+
+  it('should have empty subscriptions initially', () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+    const ctx = parentActor.getSnapshot().children.subject?.getSnapshot()?.context
+    expect(ctx?.subscriptions.size).toBe(0)
+    expect(ctx?.subscriptionConfigs.size).toBe(0)
+    expect(ctx?.cachedConnection).toBeNull()
+    expect(ctx?.syncRequired).toBe(0)
+    parentActor.stop()
+  })
+
+  it('should transition to subject_connected on SUBJECT.CONNECT', () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+
+    const connection = createMockConnection()
+    parentActor.send({ type: 'SUBJECT.CONNECT', connection })
+
+    const childSnap = parentActor.getSnapshot().children.subject?.getSnapshot()
+    expect(childSnap?.value).toBe('subject_connected')
+    expect(childSnap?.context.cachedConnection).toBe(connection)
+    parentActor.stop()
+  })
+
+  it('should add subscription config on SUBJECT.SUBSCRIBE', () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+
+    const callback = vi.fn()
+    parentActor.send({
+      type: 'SUBJECT.SUBSCRIBE',
+      config: { subject: 'test.sub', callback },
+    })
+
+    const ctx = parentActor.getSnapshot().children.subject?.getSnapshot()?.context
+    expect(ctx?.subscriptionConfigs.has('test.sub')).toBe(true)
+    expect(ctx?.syncRequired).toBe(1)
+    parentActor.stop()
+  })
+
+  it('should remove subscription config on SUBJECT.UNSUBSCRIBE', () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+
+    parentActor.send({
+      type: 'SUBJECT.SUBSCRIBE',
+      config: { subject: 'test.sub', callback: vi.fn() },
+    })
+    parentActor.send({ type: 'SUBJECT.UNSUBSCRIBE', subject: 'test.sub' })
+
+    const ctx = parentActor.getSnapshot().children.subject?.getSnapshot()?.context
+    expect(ctx?.subscriptionConfigs.has('test.sub')).toBe(false)
+    expect(ctx?.syncRequired).toBe(2)
+    parentActor.stop()
+  })
+
+  it('should clear all subscription configs on SUBJECT.UNSUBSCRIBE_ALL', () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+
+    parentActor.send({
+      type: 'SUBJECT.SUBSCRIBE',
+      config: { subject: 'test.sub1', callback: vi.fn() },
+    })
+    parentActor.send({
+      type: 'SUBJECT.SUBSCRIBE',
+      config: { subject: 'test.sub2', callback: vi.fn() },
+    })
+    parentActor.send({ type: 'SUBJECT.UNSUBSCRIBE_ALL' })
+
+    const ctx = parentActor.getSnapshot().children.subject?.getSnapshot()?.context
+    expect(ctx?.subscriptionConfigs.size).toBe(0)
+    parentActor.stop()
+  })
+
+  it('should sync subscriptions when connecting with pending configs', () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+
+    const callback = vi.fn()
+    parentActor.send({
+      type: 'SUBJECT.SUBSCRIBE',
+      config: { subject: 'test.sub', callback },
+    })
+
+    const connection = createMockConnection()
+    parentActor.send({ type: 'SUBJECT.CONNECT', connection })
+
+    const childSnap = parentActor.getSnapshot().children.subject?.getSnapshot()
+    expect(childSnap?.value).toBe('subject_connected')
+    expect(connection.subscribe).toHaveBeenCalledWith('test.sub', undefined)
+    parentActor.stop()
+  })
+
+  it('should sync when new subscription added while connected', () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+
+    const connection = createMockConnection()
+    parentActor.send({ type: 'SUBJECT.CONNECT', connection })
+
+    parentActor.send({
+      type: 'SUBJECT.SUBSCRIBE',
+      config: { subject: 'new.sub', callback: vi.fn() },
+    })
+
+    const childSnap = parentActor.getSnapshot().children.subject?.getSnapshot()
+    expect(childSnap?.value).toBe('subject_connected')
+    expect(connection.subscribe).toHaveBeenCalledWith('new.sub', undefined)
+    parentActor.stop()
+  })
+
+  it('should handle SUBJECT.PUBLISH when connected', () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+
+    const connection = createMockConnection()
+    parentActor.send({ type: 'SUBJECT.CONNECT', connection })
+
+    parentActor.send({
+      type: 'SUBJECT.PUBLISH',
+      subject: 'test.pub',
+      payload: { msg: 'hello' },
+    })
+
+    expect(connection.publish).toHaveBeenCalledWith('test.pub', { msg: 'hello' }, undefined)
+    parentActor.stop()
+  })
+
+  it('should handle SUBJECT.REQUEST when connected', () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+
+    const connection = createMockConnection()
+    parentActor.send({ type: 'SUBJECT.CONNECT', connection })
+
+    const callback = vi.fn()
+    parentActor.send({
+      type: 'SUBJECT.REQUEST',
+      subject: 'test.req',
+      payload: { data: 1 },
+      callback,
+    })
+
+    expect(connection.request).toHaveBeenCalledWith('test.req', { data: 1 }, undefined)
+    parentActor.stop()
+  })
+
+  it('should accept SUBJECT.SUBSCRIBE in idle state (global handler)', () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+
+    parentActor.send({
+      type: 'SUBJECT.SUBSCRIBE',
+      config: { subject: 'pre.connect', callback: vi.fn() },
+    })
+
+    const ctx = parentActor.getSnapshot().children.subject?.getSnapshot()?.context
+    expect(ctx?.subscriptionConfigs.has('pre.connect')).toBe(true)
+    parentActor.stop()
+  })
+
+  it('should notify parent of SUBJECT.CONNECTED', () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+
+    const connection = createMockConnection()
+    parentActor.send({ type: 'SUBJECT.CONNECT', connection })
+
+    expect(parentActor.getSnapshot().context.childState).toBe('connected')
+    parentActor.stop()
+  })
+
+  it('should handle multiple subscribe/unsubscribe cycles', () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+
+    const connection = createMockConnection()
+    parentActor.send({ type: 'SUBJECT.CONNECT', connection })
+
+    parentActor.send({
+      type: 'SUBJECT.SUBSCRIBE',
+      config: { subject: 'sub1', callback: vi.fn() },
+    })
+    parentActor.send({
+      type: 'SUBJECT.SUBSCRIBE',
+      config: { subject: 'sub2', callback: vi.fn() },
+    })
+
+    const ctx1 = parentActor.getSnapshot().children.subject?.getSnapshot()?.context
+    expect(ctx1?.subscriptionConfigs.size).toBe(2)
+
+    parentActor.send({ type: 'SUBJECT.UNSUBSCRIBE', subject: 'sub1' })
+
+    const ctx2 = parentActor.getSnapshot().children.subject?.getSnapshot()?.context
+    expect(ctx2?.subscriptionConfigs.size).toBe(1)
+    expect(ctx2?.subscriptionConfigs.has('sub2')).toBe(true)
+    parentActor.stop()
+  })
+})