@jr200-labs/xstate-nats 0.9.21 → 0.9.23

package/src/machines/credentials.test.ts

@@ -0,0 +1,175 @@
+import { TimeoutError } from '@nats-io/nats-core'
+import { createActor, waitFor } from 'xstate'
+import { afterEach, describe, expect, it, vi } from 'vitest'
+import { credentialMachine, getCredentials } from './credentials'
+
+describe('credentialMachine', () => {
+  afterEach(() => vi.useRealTimers())
+
+  it('loads once and shares the ready credentials', async () => {
+    const credentials = { auth: { type: 'token' as const, token: 'one' } }
+    const load = vi.fn(async () => credentials)
+    const actor = createActor(credentialMachine, { input: { adapter: { load } } }).start()
+
+    await expect(Promise.all([getCredentials(actor), getCredentials(actor)])).resolves.toEqual([
+      credentials,
+      credentials,
+    ])
+    expect(load).toHaveBeenCalledOnce()
+    expect(actor.getSnapshot().value).toBe('ready')
+    actor.stop()
+  })
+
+  it('deduplicates requests while refreshing near expiry', async () => {
+    let finishRefresh!: (value: { auth: { type: 'token'; token: string } }) => void
+    const refreshed = { auth: { type: 'token' as const, token: 'two' } }
+    const refresh = vi.fn(
+      () =>
+        new Promise<{ auth: { type: 'token'; token: string } }>((resolve) => {
+          finishRefresh = resolve
+        }),
+    )
+    const actor = createActor(credentialMachine, {
+      input: {
+        adapter: {
+          load: async () => ({
+            auth: { type: 'token' as const, token: 'one' },
+            expiresAt: Date.now() + 30_000,
+          }),
+          refresh,
+        },
+      },
+    }).start()
+
+    await waitFor(actor, (snapshot) => snapshot.matches('refreshing'))
+    const first = getCredentials(actor)
+    const second = getCredentials(actor)
+    finishRefresh(refreshed)
+
+    await expect(Promise.all([first, second])).resolves.toEqual([refreshed, refreshed])
+    expect(refresh).toHaveBeenCalledOnce()
+    actor.stop()
+  })
+
+  it('proactively refreshes before credentials expire', async () => {
+    vi.useFakeTimers()
+    const now = new Date('2026-01-01T00:00:00Z')
+    vi.setSystemTime(now)
+    const refresh = vi.fn(async () => ({
+      auth: { type: 'token' as const, token: 'two' },
+      expiresAt: now.getTime() + 600_000,
+    }))
+    const actor = createActor(credentialMachine, {
+      input: {
+        refreshBeforeExpiryMs: 1_000,
+        adapter: {
+          load: async () => ({
+            auth: { type: 'token' as const, token: 'one' },
+            expiresAt: now.getTime() + 5_000,
+          }),
+          refresh,
+        },
+      },
+    }).start()
+
+    await vi.advanceTimersByTimeAsync(4_000)
+    await vi.waitFor(() => expect(refresh).toHaveBeenCalledOnce())
+    await expect(getCredentials(actor)).resolves.toMatchObject({
+      auth: { token: 'two' },
+    })
+    actor.stop()
+  })
+
+  it('rejects expired credentials when no refresh adapter exists', async () => {
+    vi.useFakeTimers()
+    const actor = createActor(credentialMachine, {
+      input: {
+        adapter: {
+          load: async () => ({
+            auth: { type: 'token' as const, token: 'one' },
+            expiresAt: Date.now() + 10,
+          }),
+        },
+      },
+    }).start()
+
+    await vi.advanceTimersByTimeAsync(10)
+    await expect(getCredentials(actor)).rejects.toThrow('expired')
+    expect(actor.getSnapshot().value).toBe('expired')
+    actor.stop()
+  })
+
+  it('does not release an expired value returned by the initial load', async () => {
+    const actor = createActor(credentialMachine, {
+      input: {
+        adapter: {
+          load: async () => ({
+            auth: { type: 'token' as const, token: 'stale' },
+            expiresAt: Date.now() - 1,
+          }),
+        },
+      },
+    }).start()
+
+    await expect(getCredentials(actor)).rejects.toThrow('expired')
+    expect(actor.getSnapshot().value).toBe('expired')
+    actor.stop()
+  })
+
+  it('fails closed when refresh returns another stale value', async () => {
+    const actor = createActor(credentialMachine, {
+      input: {
+        adapter: {
+          load: async () => ({ expiresAt: Date.now() + 1_000 }),
+          refresh: async () => ({ expiresAt: Date.now() + 1_000 }),
+        },
+      },
+    }).start()
+
+    await expect(getCredentials(actor)).rejects.toThrow('too close to expiry')
+    expect(actor.getSnapshot().value).toBe('failed')
+    actor.stop()
+  })
+
+  it('bounds a hung adapter operation and fails all waiters', async () => {
+    const actor = createActor(credentialMachine, {
+      input: {
+        operationTimeoutMs: 5,
+        adapter: { load: () => new Promise(() => {}) },
+      },
+    }).start()
+
+    const first = getCredentials(actor)
+    const second = getCredentials(actor)
+    await expect(first).rejects.toBeInstanceOf(TimeoutError)
+    await expect(second).rejects.toBeInstanceOf(TimeoutError)
+    await waitFor(actor, (snapshot) => snapshot.matches('failed'))
+    actor.stop()
+  })
+
+  it('can reload after a failure', async () => {
+    const load = vi
+      .fn<() => Promise<{ auth: { type: 'token'; token: string } }>>()
+      .mockRejectedValueOnce(new Error('session unavailable'))
+      .mockResolvedValueOnce({ auth: { type: 'token', token: 'recovered' } })
+    const actor = createActor(credentialMachine, { input: { adapter: { load } } }).start()
+    await waitFor(actor, (snapshot) => snapshot.matches('failed'))
+
+    actor.send({ type: 'CREDENTIALS.RELOAD' })
+
+    await expect(getCredentials(actor)).resolves.toMatchObject({
+      auth: { token: 'recovered' },
+    })
+    expect(load).toHaveBeenCalledTimes(2)
+    actor.stop()
+  })
+
+  it('rejects immediately after the actor is stopped', async () => {
+    const actor = createActor(credentialMachine, {
+      input: { adapter: { load: async () => ({}) } },
+    }).start()
+    actor.stop()
+
+    await expect(getCredentials(actor)).rejects.toThrow('actor is stopped')
+  })
+})

package/src/machines/credentials.ts

@@ -0,0 +1,320 @@
+import { deadline, MsgHdrs } from '@nats-io/nats-core'
+import { ActorRefFrom, assign, fromPromise, setup } from 'xstate'
+import { type AuthConfig } from '../actions/types'
+import { recordCredentialOperation, recordCredentialState, withSpan } from '../telemetry'
+
+export interface NatsCredentials {
+  auth?: AuthConfig
+  requestHeaders?: MsgHdrs
+  /** Epoch milliseconds after which these credentials must not be used. Use the earliest expiry. */
+  expiresAt?: number
+}
+
+export interface NatsCredentialAdapter {
+  load: () => Promise<NatsCredentials>
+  refresh?: (current: NatsCredentials) => Promise<NatsCredentials>
+}
+
+export interface NatsCredentialConfig {
+  adapter: NatsCredentialAdapter
+  /** Refresh this long before expiry. Defaults to 60 seconds. */
+  refreshBeforeExpiryMs?: number
+  /** Maximum duration of one adapter load or refresh. Defaults to 30 seconds. */
+  operationTimeoutMs?: number
+  /** Called when credentials expire or cannot be loaded/refreshed. */
+  onUnavailable?: (error: Error) => void
+  /** Emit sanitized credential lifecycle debug logs. */
+  diagnostics?: boolean
+}
+
+type CredentialResult = { ok: true; credentials: NatsCredentials } | { ok: false; error: Error }
+
+type Waiter = (result: CredentialResult) => void
+
+export interface CredentialContext {
+  config: NatsCredentialConfig
+  current?: NatsCredentials
+  waiters: Waiter[]
+  error?: Error
+}
+
+type Events = { type: 'CREDENTIALS.GET'; reply: Waiter } | { type: 'CREDENTIALS.RELOAD' }
+
+function asError(error: unknown): Error {
+  return error instanceof Error ? error : new Error(String(error))
+}
+
+function operationTimeout(config: NatsCredentialConfig): number {
+  return config.operationTimeoutMs ?? 30_000
+}
+
+function refreshBeforeExpiry(config: NatsCredentialConfig): number {
+  return config.refreshBeforeExpiryMs ?? 60_000
+}
+
+function credentialDiagnostic(
+  config: NatsCredentialConfig,
+  event: string,
+  attributes: Record<string, string | number | boolean>,
+): void {
+  if (config.diagnostics) console.debug('xstate-nats credentials', { event, ...attributes })
+}
+
+async function runCredentialOperation<T>(
+  config: NatsCredentialConfig,
+  operation: 'load' | 'refresh',
+  fn: () => Promise<T>,
+): Promise<T> {
+  const timeout = operationTimeout(config)
+  const startedAt = performance.now()
+  credentialDiagnostic(config, 'operation.start', { operation, timeout_ms: timeout })
+  let originalError: unknown
+
+  try {
+    const result = await withSpan(
+      `xstate.nats.credentials.${operation}`,
+      'xstate.nats.credentials.error',
+      { 'credentials.operation': operation, 'credentials.timeout.ms': timeout },
+      async (span) => {
+        try {
+          const value = await deadline(Promise.resolve().then(fn), timeout)
+          span.setAttribute('credentials.outcome', 'success')
+          return value
+        } catch (error) {
+          originalError = error
+          span.setAttribute('credentials.outcome', 'error')
+          // Keep provider error messages, stacks, tokens, and claims out of traces.
+          throw new Error(`NATS credential ${operation} failed`)
+        }
+      },
+    )
+    const durationMs = performance.now() - startedAt
+    recordCredentialOperation(operation, 'success', durationMs)
+    credentialDiagnostic(config, 'operation.finish', {
+      operation,
+      outcome: 'success',
+      duration_ms: durationMs,
+    })
+    return result
+  } catch {
+    const durationMs = performance.now() - startedAt
+    recordCredentialOperation(operation, 'error', durationMs)
+    credentialDiagnostic(config, 'operation.finish', {
+      operation,
+      outcome: 'error',
+      duration_ms: durationMs,
+    })
+    throw originalError ?? new Error(`NATS credential ${operation} failed`)
+  }
+}
+
+function credentialStateEntry(state: Parameters<typeof recordCredentialState>[0]) {
+  return ({ context }: { context: CredentialContext }) => {
+    recordCredentialState(state)
+    credentialDiagnostic(context.config, 'state.enter', {
+      state,
+      pending_waiters: context.waiters.length,
+    })
+  }
+}
+
+function notify(waiter: Waiter, result: CredentialResult): void {
+  try {
+    waiter(result)
+  } catch (error) {
+    console.error('NATS credential waiter failed', error)
+  }
+}
+
+export const credentialMachine = setup({
+  types: {
+    context: {} as CredentialContext,
+    events: {} as Events,
+    input: {} as NatsCredentialConfig,
+  },
+  actors: {
+    load: fromPromise(({ input }: { input: NatsCredentialConfig }) =>
+      runCredentialOperation(input, 'load', input.adapter.load),
+    ),
+    refresh: fromPromise(
+      ({ input }: { input: { config: NatsCredentialConfig; current: NatsCredentials } }) => {
+        const refresh = input.config.adapter.refresh
+        if (!refresh) throw new Error('NATS credentials cannot be refreshed')
+        return runCredentialOperation(input.config, 'refresh', async () => {
+          const credentials = await refresh(input.current)
+          const expiresAt = credentials.expiresAt
+          if (
+            expiresAt !== undefined &&
+            expiresAt - refreshBeforeExpiry(input.config) <= Date.now()
+          ) {
+            throw new Error('Refreshed NATS credentials are already expired or too close to expiry')
+          }
+          return credentials
+        })
+      },
+    ),
+  },
+  delays: {
+    credentialDeadline: ({ context }) => {
+      const expiresAt = context.current?.expiresAt
+      if (expiresAt === undefined) return 2_147_483_647
+      const advance = context.config.adapter.refresh ? refreshBeforeExpiry(context.config) : 0
+      return Math.max(0, expiresAt - advance - Date.now())
+    },
+  },
+  guards: {
+    canRefresh: ({ context }) => Boolean(context.current && context.config.adapter.refresh),
+    hasExpiry: ({ context }) => context.current?.expiresAt !== undefined,
+    shouldRefreshNow: ({ context }) => {
+      const expiresAt = context.current?.expiresAt
+      return Boolean(
+        context.config.adapter.refresh &&
+        expiresAt !== undefined &&
+        expiresAt - refreshBeforeExpiry(context.config) <= Date.now(),
+      )
+    },
+    isExpired: ({ context }) => {
+      const expiresAt = context.current?.expiresAt
+      return expiresAt !== undefined && expiresAt <= Date.now()
+    },
+    outputShouldRefresh: ({ context, event }) => {
+      if (!('output' in event) || !context.config.adapter.refresh) return false
+      const expiresAt = (event.output as NatsCredentials).expiresAt
+      return (
+        expiresAt !== undefined && expiresAt - refreshBeforeExpiry(context.config) <= Date.now()
+      )
+    },
+    outputExpired: ({ event }) => {
+      if (!('output' in event)) return false
+      const expiresAt = (event.output as NatsCredentials).expiresAt
+      return expiresAt !== undefined && expiresAt <= Date.now()
+    },
+  },
+  actions: {
+    queue: assign({
+      waiters: ({ context, event }) =>
+        event.type === 'CREDENTIALS.GET' ? [...context.waiters, event.reply] : context.waiters,
+    }),
+    replyReady: ({ context, event }) => {
+      if (event.type === 'CREDENTIALS.GET' && context.current) {
+        notify(event.reply, { ok: true, credentials: context.current })
+      }
+    },
+    acceptCredentials: assign(({ context, event }) => {
+      if (!('output' in event)) return {}
+      const credentials = event.output as NatsCredentials
+      context.waiters.forEach((waiter) => notify(waiter, { ok: true, credentials }))
+      return { current: credentials, waiters: [], error: undefined }
+    }),
+    storeCredentials: assign(({ event }) =>
+      'output' in event ? { current: event.output as NatsCredentials, error: undefined } : {},
+    ),
+    rejectWaiters: assign(({ context, event }) => {
+      const error =
+        'error' in event
+          ? asError(event.error)
+          : (context.error ?? new Error('NATS credentials are unavailable'))
+      context.waiters.forEach((waiter) => notify(waiter, { ok: false, error }))
+      return { waiters: [], error }
+    }),
+    markExpired: assign({
+      error: (_) => new Error('NATS credentials expired'),
+    }),
+    notifyUnavailable: ({ context }) => {
+      const error = context.error ?? new Error('NATS credentials are unavailable')
+      try {
+        context.config.onUnavailable?.(error)
+      } catch (callbackError) {
+        console.error('NATS credential unavailable callback failed', callbackError)
+      }
+    },
+  },
+}).createMachine({
+  id: 'credentials',
+  initial: 'loading',
+  context: ({ input }) => ({ config: input, waiters: [] }),
+  states: {
+    loading: {
+      entry: credentialStateEntry('loading'),
+      on: {
+        'CREDENTIALS.GET': { actions: 'queue' },
+      },
+      invoke: {
+        src: 'load',
+        input: ({ context }) => context.config,
+        onDone: [
+          { guard: 'outputShouldRefresh', target: 'refreshing', actions: 'storeCredentials' },
+          { guard: 'outputExpired', target: 'expired', actions: 'storeCredentials' },
+          { target: 'ready', actions: 'acceptCredentials' },
+        ],
+        onError: { target: 'failed', actions: 'rejectWaiters' },
+      },
+    },
+    ready: {
+      entry: credentialStateEntry('ready'),
+      after: {
+        credentialDeadline: [
+          { guard: 'canRefresh', target: 'refreshing' },
+          { guard: 'hasExpiry', target: 'expired' },
+        ],
+      },
+      on: {
+        'CREDENTIALS.GET': [
+          { guard: 'shouldRefreshNow', target: 'refreshing', actions: 'queue' },
+          { guard: 'isExpired', target: 'expired', actions: 'queue' },
+          { actions: 'replyReady' },
+        ],
+        'CREDENTIALS.RELOAD': { target: 'loading' },
+      },
+    },
+    refreshing: {
+      entry: credentialStateEntry('refreshing'),
+      on: {
+        'CREDENTIALS.GET': { actions: 'queue' },
+      },
+      invoke: {
+        src: 'refresh',
+        input: ({ context }) => ({ config: context.config, current: context.current! }),
+        onDone: { target: 'ready', actions: 'acceptCredentials' },
+        onError: { target: 'failed', actions: 'rejectWaiters' },
+      },
+    },
+    expired: {
+      entry: [credentialStateEntry('expired'), 'markExpired', 'rejectWaiters', 'notifyUnavailable'],
+      on: {
+        'CREDENTIALS.GET': {
+          actions: ({ event }) =>
+            notify(event.reply, { ok: false, error: new Error('NATS credentials expired') }),
+        },
+        'CREDENTIALS.RELOAD': { target: 'loading' },
+      },
+    },
+    failed: {
+      entry: [credentialStateEntry('failed'), 'notifyUnavailable'],
+      on: {
+        'CREDENTIALS.GET': {
+          actions: ({ context, event }) =>
+            notify(event.reply, {
+              ok: false,
+              error: context.error ?? new Error('NATS credentials are unavailable'),
+            }),
+        },
+        'CREDENTIALS.RELOAD': { target: 'loading' },
+      },
+    },
+  },
+})
+
+export type CredentialActor = ActorRefFrom<typeof credentialMachine>
+
+export function getCredentials(actor: CredentialActor): Promise<NatsCredentials> {
+  if (actor.getSnapshot().status === 'stopped') {
+    return Promise.reject(new Error('NATS credential actor is stopped'))
+  }
+  return new Promise((resolve, reject) => {
+    actor.send({
+      type: 'CREDENTIALS.GET',
+      reply: (result) => (result.ok ? resolve(result.credentials) : reject(result.error)),
+    })
+  })
+}

package/src/machines/root.test.ts

@@ -2,6 +2,8 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'
 import { createActor, fromPromise, createMachine, sendParent } from 'xstate'
 import { natsMachine } from './root'
 
+const subjectEventSpy = vi.fn()
+
 vi.mock('@nats-io/nats-core', async () => {
   const actual = await vi.importActual<typeof import('@nats-io/nats-core')>('@nats-io/nats-core')
   return {
@@ -29,7 +31,7 @@ const mockSubjectMachine = createMachine({
     connected: {
       entry: sendParent({ type: 'SUBJECT.CONNECTED' }),
       on: {
-        'SUBJECT.*': {},
+        'SUBJECT.*': { actions: ({ event }) => subjectEventSpy(event) },
         'SUBJECT.DISCONNECTED': { target: 'idle' },
       },
     },
@@ -137,6 +139,7 @@ function configureAndConnect(actor: any) {
 
 describe('natsMachine', () => {
   beforeEach(() => {
+    subjectEventSpy.mockClear()
     vi.spyOn(console, 'log').mockImplementation(() => {})
     vi.spyOn(console, 'error').mockImplementation(() => {})
   })
@@ -179,6 +182,90 @@ describe('natsMachine', () => {
     actor.stop()
   })
 
+  it('forwards current credential headers to the subject manager', async () => {
+    const { headers } = await import('@nats-io/nats-core')
+    const requestHeaders = headers()
+    requestHeaders.set('authorization', 'Bearer current')
+    const actor = createActor(createTestMachine())
+    actor.start()
+    actor.send({
+      type: 'CONFIGURE',
+      config: {
+        opts: { servers: ['ws://localhost:4222'] },
+        maxRetries: 3,
+        credentials: {
+          adapter: { load: vi.fn(async () => ({ requestHeaders })) },
+        },
+      },
+    })
+    actor.send({ type: 'CONNECT' })
+    await vi.waitFor(() => expect(actor.getSnapshot().value).toBe('connected'))
+
+    actor.send({
+      type: 'SUBJECT.REQUEST',
+      subject: 'test.request',
+      payload: {},
+      callback: vi.fn(),
+    })
+
+    await vi.waitFor(() => {
+      expect(subjectEventSpy).toHaveBeenCalledWith(
+        expect.objectContaining({
+          requestHeaders: expect.any(Function),
+          type: 'SUBJECT.REQUEST',
+        }),
+      )
+    })
+    const forwarded = subjectEventSpy.mock.calls.at(-1)?.[0]
+    await expect(forwarded.requestHeaders()).resolves.toBe(requestHeaders)
+    actor.stop()
+  })
+
+  it('closes an active connection when credential refresh fails', async () => {
+    const actor = createActor(createTestMachine()).start()
+    actor.send({
+      type: 'CONFIGURE',
+      config: {
+        opts: { servers: ['ws://localhost:4222'] },
+        maxRetries: 3,
+        credentials: {
+          refreshBeforeExpiryMs: 5,
+          adapter: {
+            load: async () => ({ expiresAt: Date.now() + 100 }),
+            refresh: async () => {
+              throw new Error('refresh failed')
+            },
+          },
+        },
+      },
+    })
+    actor.send({ type: 'CONNECT' })
+    await vi.waitFor(() => expect(actor.getSnapshot().value).toBe('connected'))
+
+    await vi.waitFor(() => expect(actor.getSnapshot().value).toBe('closed'))
+    actor.stop()
+  })
+
+  it('fails requests immediately while disconnected', () => {
+    const actor = createActor(createTestMachine())
+    const onRequestResult = vi.fn()
+    actor.start()
+
+    actor.send({
+      type: 'SUBJECT.REQUEST',
+      subject: 'test.request',
+      payload: {},
+      callback: vi.fn(),
+      onRequestResult,
+    })
+
+    expect(onRequestResult).toHaveBeenCalledWith({
+      ok: false,
+      error: expect.objectContaining({ message: 'NATS connection is not available' }),
+    })
+    actor.stop()
+  })
+
   it('should transition to closed on DISCONNECT from connected', async () => {
     const actor = createActor(createTestMachine())
     actor.start()