@jr200-labs/xstate-nats 0.9.21 → 0.9.22

package/src/machines/root.ts

@@ -1,5 +1,5 @@
 import { ConnectionOptions, NatsConnection } from '@nats-io/nats-core'
-import { assign, sendTo, setup } from 'xstate'
+import { assign, sendTo, setup, stopChild } from 'xstate'
 import { kvManagerLogic, ExternalEvents as KvExternalEvents } from './kv'
 import { subjectManagerLogic, ExternalEvents as SubjectExternalEvents } from './subject'
 import { connectToNats, disconnectNats } from '../actions/connection'
@@ -7,6 +7,13 @@ import { type AuthConfig } from '../actions/types'
 import { InternalStatusEvents as NatsStatusEvents } from '../actions/connection'
 import { withSpan } from '../telemetry'
 import { NatsTrafficResetEvent } from '../traffic'
+import { rejectUnavailableRequest } from '../actions/subject'
+import {
+  CredentialActor,
+  credentialMachine,
+  getCredentials,
+  NatsCredentialConfig,
+} from './credentials'
 
 export interface NatsDiagnosticsConfig {
   lifecycle?: boolean
@@ -15,6 +22,7 @@ export interface NatsDiagnosticsConfig {
 export interface NatsConnectionConfig {
   opts: ConnectionOptions
   auth?: AuthConfig
+  credentials?: NatsCredentialConfig
   maxRetries: number
   diagnostics?: NatsDiagnosticsConfig
 }
@@ -28,6 +36,7 @@ export interface Context {
   kvManagerReady: boolean
   configureAfterClose: boolean
   connectAfterClose: boolean
+  credentialActor?: CredentialActor
 }
 
 // internal events and events from nats connection
@@ -37,6 +46,7 @@ type InternalEvents =
   | { type: 'FAIL'; error: Error }
   | { type: 'RECONNECT' }
   | { type: 'CLOSE' }
+  | { type: 'CREDENTIALS.UNAVAILABLE'; error: Error }
   | NatsStatusEvents
 
 // events which can be sent to the machine from the user
@@ -45,6 +55,7 @@ export type ExternalEvents =
   | { type: 'CONNECT' }
   | { type: 'DISCONNECT' }
   | { type: 'RESET' }
+  | { type: 'CREDENTIALS.RELOAD' }
   | NatsTrafficResetEvent
   | SubjectExternalEvents
   | KvExternalEvents
@@ -72,7 +83,9 @@ function lifecycleAttributes(
     'nats.debug': Boolean(context.natsConfig?.opts.debug),
     'nats.verbose': Boolean(context.natsConfig?.opts.verbose),
     'nats.max_retries': context.natsConfig?.maxRetries,
-    'nats.auth.type': context.natsConfig?.auth?.type,
+    'nats.auth.type':
+      context.credentialActor?.getSnapshot().context.current?.auth?.type ??
+      context.natsConfig?.auth?.type,
     'nats.has_connection': context.connection !== null,
     'nats.subject_manager_ready': context.subjectManagerReady,
     'nats.kv_manager_ready': context.kvManagerReady,
@@ -103,6 +116,17 @@ export const natsMachine = setup({
     events: {} as Events,
   },
   actions: {
+    stopCredentialActor: stopChild(({ context }) => context.credentialActor as CredentialActor),
+    reloadCredentials: ({ context }) => {
+      context.credentialActor?.send({ type: 'CREDENTIALS.RELOAD' })
+    },
+    recordCredentialError: assign({
+      error: ({ event }) => (event.type === 'CREDENTIALS.UNAVAILABLE' ? event.error : undefined),
+    }),
+    rejectUnavailableRequest: ({ event }) => {
+      if (event.type !== 'SUBJECT.REQUEST') return
+      rejectUnavailableRequest(event)
+    },
     doReset: assign({
       natsConfig: (_) => undefined,
       connection: (_) => null,
@@ -112,6 +136,7 @@ export const natsMachine = setup({
       kvManagerReady: (_) => false,
       configureAfterClose: (_) => false,
       connectAfterClose: (_) => false,
+      credentialActor: (_) => undefined,
     }),
     configureNats: assign({
       natsConfig: ({ event }) => {
@@ -127,6 +152,22 @@ export const natsMachine = setup({
       kvManagerReady: (_) => false,
       configureAfterClose: (_) => false,
       connectAfterClose: (_) => false,
+      credentialActor: ({ event, spawn, self }) => {
+        if (event.type !== 'CONFIGURE' || !event.config.credentials) return undefined
+        const onUnavailable = event.config.credentials.onUnavailable
+        return spawn('credentials', {
+          input: {
+            ...event.config.credentials,
+            onUnavailable: (error) => {
+              try {
+                onUnavailable?.(error)
+              } finally {
+                self.send({ type: 'CREDENTIALS.UNAVAILABLE', error })
+              }
+            },
+          },
+        })
+      },
     }),
     configureNatsAfterClose: assign({
       natsConfig: ({ event }) => {
@@ -140,6 +181,22 @@ export const natsMachine = setup({
       subjectManagerReady: (_) => false,
       kvManagerReady: (_) => false,
       configureAfterClose: (_) => true,
+      credentialActor: ({ event, spawn, self }) => {
+        if (event.type !== 'CONFIGURE' || !event.config.credentials) return undefined
+        const onUnavailable = event.config.credentials.onUnavailable
+        return spawn('credentials', {
+          input: {
+            ...event.config.credentials,
+            onUnavailable: (error) => {
+              try {
+                onUnavailable?.(error)
+              } finally {
+                self.send({ type: 'CREDENTIALS.UNAVAILABLE', error })
+              }
+            },
+          },
+        })
+      },
     }),
     reconnectAfterClose: assign({
       connectAfterClose: (_) => true,
@@ -168,6 +225,7 @@ export const natsMachine = setup({
     disconnectNats: disconnectNats,
     subject: subjectManagerLogic,
     kv: kvManagerLogic,
+    credentials: credentialMachine,
   },
 }).createMachine({
   /** @xstate-layout N4IgpgJg5mDOIC5QAoC2BDAxgCwJYDswBKAOnwHsAXAfU3PwDNcoBXAJ0gGIBhAeQDkAYgEkA4gFUASgFEA2gAYAuolAAHcrFyVc9FSAAeiABwBOeSQBsAdgCMRgMwAmRxftWALPYCsAGhABPRHd3c3cvext7IwsveRCrAF8EvzQsPEJSOkZmdi4+fn5pbgAVBWUkEHVNbV0KwwQbGw9LJysLExNHSPkjP0CEdzaSMwsXaPt5LxMbeQsklIwcAmISLMJMbXwoTgh6MBICADdyAGt91KWM1fp1zagEI-JMdBr8MrK9Kq0dfD16m3CVhIUUcVis8isJjCjlmfUQ4XcLUcRiakXcJicc2SIAu6RWazAGwI2zAbDY5DYJFUABsXgwKagSLjlpkboS7g98Mdnq93kpPhpvrVQP9AcCjKDwZDobCAkF3DYSE0TNZ2p1ptF3PMcYs8az8Lc8gAZXgAZTk-IqX1ef0QNlGQK60TMEUGVi8NjhCFciJmGIxRmlYJs2uZV0w1MFWx2ewOXNO511LNWkc0W053JePz55TUgptdUQyK8JDiFiMXhC8hco3k9i9mJIVk17Rs03LYPsoaT4dTxM4pPJlJpdIZTJ7+L76cePOzSg+VvzP1tCGLpfc5cr8mrFlr9blDUrJFi2-kKs6NlB3bSyYjGi4MnNpUteeqy8LCCm9mB3kDtjBPRhA21bAme7bupukzXpcKyDhSnCPtIz65pUS7CgYdoWI0x4RG2IRNJ4jjuF6jgqk2LheJCbQRPYWrahQEBwHoYbEAKb7ofUAC0Fhetx0F6mQVC0PQTCsBwEBsUKvwflYXTkTE9qdq49gmA2jgls2FgbhuowmB4-G3iJOTiZJBYinacSKtYKoWHW7pGIGPEHm6JAhKMXikXYNieHRCw3uGbJElspnvuZDQRI4JAokRHkqjCG69M54TDLMXT2CpERmIk2IsfqhoSYu7HSWFTQqcMTrpVhKkool-SVhYx4uKlMylXYBm9lGUAhRxQSOF69omEq8gAruESTB5kTtZO94Fa+UkruiiqpdCHSOLhTn9E4GktjppH6TlE6kHBbDdcVGENFYZUXg49hVSYNVqeYMIRW07TqR42VJEAA */
@@ -181,12 +239,14 @@ export const natsMachine = setup({
     kvManagerReady: false,
     configureAfterClose: false,
     connectAfterClose: false,
+    credentialActor: undefined,
   },
   invoke: [
     { src: 'subject', id: 'subject', systemId: 'subject' },
     { src: 'kv', id: 'kv' },
   ],
   on: {
+    'CREDENTIALS.RELOAD': { actions: 'reloadCredentials' },
     'NATS_CONNECTION.CLOSE': {
       target: '.closed',
       actions: [
@@ -221,9 +281,10 @@ export const natsMachine = setup({
   states: {
     not_configured: {
       on: {
+        'SUBJECT.REQUEST': { actions: 'rejectUnavailableRequest' },
         CONFIGURE: {
           target: 'configured',
-          actions: ['configureNats'],
+          actions: ['stopCredentialActor', 'configureNats'],
           '*': {
             actions: [
               ({ event }: { event: any }) => {
@@ -237,8 +298,9 @@ export const natsMachine = setup({
     configured: {
       entry: ['clearDeferredCloseActions'],
       on: {
+        'SUBJECT.REQUEST': { actions: 'rejectUnavailableRequest' },
         CONFIGURE: {
-          actions: ['configureNats', lifecycleAction('configured')],
+          actions: ['stopCredentialActor', 'configureNats', lifecycleAction('configured')],
         },
         CONNECT: {
           target: 'connecting',
@@ -246,7 +308,7 @@ export const natsMachine = setup({
         },
         RESET: {
           target: 'not_configured',
-          actions: ['doReset'],
+          actions: ['stopCredentialActor', 'doReset'],
         },
         '*': {
           actions: [
@@ -260,10 +322,11 @@ export const natsMachine = setup({
     connecting: {
       entry: ['clearDeferredCloseActions', lifecycleAction('connecting')],
       on: {
+        'SUBJECT.REQUEST': { actions: 'rejectUnavailableRequest' },
         CONFIGURE: {
           target: 'connecting',
           reenter: true,
-          actions: ['configureNats', lifecycleAction('connecting')],
+          actions: ['stopCredentialActor', 'configureNats', lifecycleAction('connecting')],
         },
         CONNECT: {
           actions: lifecycleAction('connecting'),
@@ -275,6 +338,7 @@ export const natsMachine = setup({
           input: ({ context, self }) => ({
             opts: context.natsConfig!.opts,
             auth: context.natsConfig!.auth,
+            credentialActor: context.credentialActor,
             onStatus: (event: NatsStatusEvents) => self.send(event),
           }),
           onDone: {
@@ -283,6 +347,7 @@ export const natsMachine = setup({
               lifecycleAction('connecting'),
               assign({
                 connection: ({ event }) => event.output,
+                error: (_) => undefined,
                 retries: (_) => 0,
               }),
             ],
@@ -310,9 +375,11 @@ export const natsMachine = setup({
         sendTo('kv', ({ context }) => ({ type: 'KV.CONNECT', connection: context.connection! })),
       ],
       on: {
+        'SUBJECT.REQUEST': { actions: 'rejectUnavailableRequest' },
         CONFIGURE: {
           target: 'closing',
           actions: [
+            'stopCredentialActor',
             'configureNatsAfterClose',
             'reconnectAfterClose',
             lifecycleAction('initialise_managers'),
@@ -341,6 +408,14 @@ export const natsMachine = setup({
             lifecycleAction('initialise_managers'),
           ],
         },
+        'CREDENTIALS.UNAVAILABLE': {
+          target: 'closing',
+          actions: [
+            'recordCredentialError',
+            'clearDeferredCloseActions',
+            lifecycleAction('initialise_managers'),
+          ],
+        },
       },
       always: [
         {
@@ -373,7 +448,12 @@ export const natsMachine = setup({
         },
         CONFIGURE: {
           target: 'closing',
-          actions: ['configureNatsAfterClose', 'reconnectAfterClose', lifecycleAction('connected')],
+          actions: [
+            'stopCredentialActor',
+            'configureNatsAfterClose',
+            'reconnectAfterClose',
+            lifecycleAction('connected'),
+          ],
         },
         DISCONNECT: {
           target: 'closing',
@@ -383,6 +463,14 @@ export const natsMachine = setup({
           target: 'closing',
           actions: ['clearDeferredCloseActions', lifecycleAction('connected')],
         },
+        'CREDENTIALS.UNAVAILABLE': {
+          target: 'closing',
+          actions: [
+            'recordCredentialError',
+            'clearDeferredCloseActions',
+            lifecycleAction('connected'),
+          ],
+        },
         'SUBJECT.*': {
           actions: [
             ({ context, event }: { context: Context; event: any }) => {
@@ -393,7 +481,18 @@ export const natsMachine = setup({
             sendTo(
               'subject',
               ({ event, context }: { event: SubjectExternalEvents; context: Context }) => {
-                return { ...event, connection: context.connection }
+                return {
+                  ...event,
+                  connection: context.connection,
+                  requestHeaders: context.credentialActor
+                    ? () =>
+                        getCredentials(context.credentialActor!).then(
+                          (value) => value.requestHeaders,
+                        )
+                    : event.type === 'SUBJECT.REQUEST'
+                      ? event.requestHeaders
+                      : undefined,
+                }
               },
             ),
           ],
@@ -460,11 +559,12 @@ export const natsMachine = setup({
         },
       },
       on: {
+        'SUBJECT.REQUEST': { actions: 'rejectUnavailableRequest' },
         'NATS_CONNECTION.*': {
           actions: lifecycleAction('closing'),
         },
         CONFIGURE: {
-          actions: ['configureNatsAfterClose', lifecycleAction('closing')],
+          actions: ['stopCredentialActor', 'configureNatsAfterClose', lifecycleAction('closing')],
         },
         CONNECT: {
           actions: ['reconnectAfterClose', lifecycleAction('closing')],
@@ -485,13 +585,14 @@ export const natsMachine = setup({
         }),
       ],
       on: {
+        'SUBJECT.REQUEST': { actions: 'rejectUnavailableRequest' },
         CONFIGURE: {
           target: 'configured',
-          actions: ['configureNats', lifecycleAction('closed')],
+          actions: ['stopCredentialActor', 'configureNats', lifecycleAction('closed')],
         },
         RESET: {
           target: 'not_configured',
-          actions: ['doReset'],
+          actions: ['stopCredentialActor', 'doReset'],
         },
         CONNECT: {
           target: 'connecting',
@@ -509,17 +610,18 @@ export const natsMachine = setup({
     error: {
       entry: ['clearDeferredCloseActions', lifecycleAction('error')],
       on: {
+        'SUBJECT.REQUEST': { actions: 'rejectUnavailableRequest' },
         CONNECT: {
           target: 'connecting',
           actions: lifecycleAction('error'),
         },
         CONFIGURE: {
           target: 'configured',
-          actions: ['configureNats', lifecycleAction('error')],
+          actions: ['stopCredentialActor', 'configureNats', lifecycleAction('error')],
         },
         RESET: {
           target: 'not_configured',
-          actions: ['doReset'],
+          actions: ['stopCredentialActor', 'doReset'],
         },
       },
       '*': {

package/src/machines/subject.test.ts

@@ -1,5 +1,6 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest'
 import { createActor, createMachine, assign, fromPromise, sendTo, setup } from 'xstate'
+import { headers } from '@nats-io/nats-core'
 import { subjectManagerLogic } from './subject'
 
 vi.mock('@nats-io/nats-core', async () => {
@@ -84,6 +85,26 @@ describe('subjectManagerLogic', () => {
     parentActor.stop()
   })
 
+  it('fails requests immediately while disconnected', () => {
+    const parentActor = createActor(createParentMachine())
+    const onRequestResult = vi.fn()
+    parentActor.start()
+
+    parentActor.send({
+      type: 'SUBJECT.REQUEST',
+      subject: 'test.req',
+      payload: {},
+      callback: vi.fn(),
+      onRequestResult,
+    })
+
+    expect(onRequestResult).toHaveBeenCalledWith({
+      ok: false,
+      error: expect.objectContaining({ message: 'NATS connection is not available' }),
+    })
+    parentActor.stop()
+  })
+
   it('should have empty subscriptions initially', () => {
     const parentActor = createActor(createParentMachine())
     parentActor.start()
@@ -289,6 +310,45 @@ describe('subjectManagerLogic', () => {
     parentActor.stop()
   })
 
+  it('gets request headers immediately before sending', async () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+
+    const connection = createMockConnection()
+    parentActor.send({ type: 'SUBJECT.CONNECT', connection })
+    const requestHeaders = vi.fn(async () => {
+      const authHeaders = headers()
+      authHeaders.set('authorization', 'Bearer refreshed-token')
+      return authHeaders
+    })
+    const staleHeaders = headers()
+    staleHeaders.set('authorization', 'Bearer stale-token')
+
+    parentActor.send({
+      type: 'SUBJECT.REQUEST',
+      subject: 'test.req',
+      payload: { data: 1 },
+      opts: { headers: staleHeaders, timeout: 1000 },
+      callback: vi.fn(),
+      requestHeaders,
+    })
+
+    await vi.waitFor(() => {
+      expect(connection.request).toHaveBeenCalledWith(
+        'test.req',
+        { data: 1 },
+        expect.objectContaining({
+          headers: expect.objectContaining({ get: expect.any(Function) }),
+        }),
+      )
+    })
+    expect(requestHeaders).toHaveBeenCalledOnce()
+    expect(connection.request.mock.calls[0][2].headers.get('authorization')).toBe(
+      'Bearer refreshed-token',
+    )
+    parentActor.stop()
+  })
+
   it('should count request upload and reply download bytes', async () => {
     const parentActor = createActor(createParentMachine())
     parentActor.start()

package/src/machines/subject.ts

@@ -1,8 +1,15 @@
-import { Subscription, PublishOptions, NatsConnection, RequestOptions } from '@nats-io/nats-core'
+import {
+  MsgHdrs,
+  Subscription,
+  PublishOptions,
+  NatsConnection,
+  RequestOptions,
+} from '@nats-io/nats-core'
 import { assign, sendParent, setup } from 'xstate'
 import {
   RequestResult,
   SubjectSubscriptionConfig,
+  rejectUnavailableRequest,
   subjectConsolidateState,
   subjectRequest,
   subjectPublish,
@@ -41,6 +48,7 @@ export type ExternalEvents =
       opts?: RequestOptions
       callback: (data: any) => void
       onRequestResult?: (result: RequestResult) => void
+      requestHeaders?: () => Promise<MsgHdrs | undefined>
     }
   | { type: 'SUBJECT.SUBSCRIBE'; config: SubjectSubscriptionConfig }
   | { type: 'SUBJECT.UNSUBSCRIBE'; subject: string }
@@ -69,6 +77,12 @@ export const subjectManagerLogic = setup({
       return context.syncRequired > 0
     },
   },
+  actions: {
+    rejectUnavailableRequest: ({ event }) => {
+      if (event.type !== 'SUBJECT.REQUEST') return
+      rejectUnavailableRequest(event)
+    },
+  },
 }).createMachine({
   /** @xstate-layout N4IgpgJg5mDOIC5QAoC2BDAxgCwJYDswBKAOlgFcAjAKzEwBcB9XCAGzAGIBlAVQCEAUgFEAwgBVGASQByksZICCAGUlchAbQAMAXUSgADgHtYuerkP49IAB6IATJs0kAnAFYAzAEY7rgDQgAT3tXTxJNAA53KOiYrwBfOP80LDxCUgoaOiYCU1x0VlwTfCgOa1h6dHowEnQAMyqAJ2RXRyIOZJwCYjIqWgZmfFz8woIoLV0kECMTMwsrWwQANk0AFhJPNy8ffyCEHzsSV0jY2M8EpIxOtJ7M-swLQgZIbn5hcUZePi4RACVJPg0Ois01yc0mC08jnCJAA7O4Vu4YX5AvYVjDDnZ4YjXOcQB1Ut0Mn0mPd8I8qhAXoJRBIeNJPt8-gDxsDjKDLODEEiDu5FjDPOEkTtEJ4Yc5cfiuulellGKTyc9Pm8JCIlEIFD8PvxGf9ARMDGzZhzQAs4a4SHY7ItXIttiiEN5oTjcfhDBA4FZJWlWTNzMabIgALSLYUIYMSy4E6W3bJsMA+9nzRArOyhy3QiInU4RlJSm7EgZDApFKAJo1JhArTzmzTV23I3bpsLHLNRM6JPGRvNE2XyrKQMt+iswpEkRaChv2eGHHNXQkyu4WWCGAoQSpgRjldeDsEmkWY0IrFYebyTvZomcdr3zmOMMANBqGBo7-0LFaLC0CoX26KHFuthIEiAA */
   initial: 'subject_idle',
@@ -150,6 +164,7 @@ export const subjectManagerLogic = setup({
   states: {
     subject_idle: {
       on: {
+        'SUBJECT.REQUEST': { actions: 'rejectUnavailableRequest' },
         'SUBJECT.DISCONNECTED': {},
         'SUBJECT.CONNECT': {
           target: 'subject_check_sync',
@@ -211,6 +226,7 @@ export const subjectManagerLogic = setup({
                 opts: event.opts,
                 callback: event.callback,
                 onRequestResult: event.onRequestResult,
+                requestHeaders: event.requestHeaders,
                 onDownloadBytes: (bytes) =>
                   self.send({
                     type: 'METRICS.RECORD',
@@ -290,6 +306,7 @@ export const subjectManagerLogic = setup({
     },
     subject_error: {
       on: {
+        'SUBJECT.REQUEST': { actions: 'rejectUnavailableRequest' },
         'SUBJECT.DISCONNECTED': {
           target: 'subject_disconnecting',
         },

package/src/telemetry.ts

@@ -14,12 +14,46 @@
 //     the per-message span so cross-process traces nest automatically.
 
 import type { Context, Span, Tracer } from '@opentelemetry/api'
-import { context as otelContext, propagation, SpanStatusCode, trace } from '@opentelemetry/api'
+import {
+  context as otelContext,
+  metrics,
+  propagation,
+  SpanStatusCode,
+  trace,
+} from '@opentelemetry/api'
 import type { MsgHdrs } from '@nats-io/nats-core'
 import { headers as natsHeaders } from '@nats-io/nats-core'
 import pkg from '../package.json' with { type: 'json' }
 
 const TRACER_NAME = '@jr200-labs/xstate-nats'
+const meter = metrics.getMeter(TRACER_NAME, pkg.version)
+const credentialOperations = meter.createCounter('xstate.nats.credentials.operations', {
+  description: 'Completed NATS credential operations',
+})
+const credentialDuration = meter.createHistogram('xstate.nats.credentials.operation.duration', {
+  description: 'NATS credential operation duration',
+  unit: 'ms',
+})
+const credentialStateTransitions = meter.createCounter(
+  'xstate.nats.credentials.state.transitions',
+  { description: 'NATS credential actor state transitions' },
+)
+
+export function recordCredentialOperation(
+  operation: 'load' | 'refresh',
+  outcome: 'success' | 'error',
+  durationMs: number,
+): void {
+  const attributes = { operation, outcome }
+  credentialOperations.add(1, attributes)
+  credentialDuration.record(durationMs, attributes)
+}
+
+export function recordCredentialState(
+  state: 'loading' | 'ready' | 'refreshing' | 'expired' | 'failed',
+): void {
+  credentialStateTransitions.add(1, { state })
+}
 
 // Do not cache — `trace.getTracer` already returns a lightweight ProxyTracer,
 // and caching across global-provider swaps (e.g. test teardown / re-register)