@flamingo-stack/openframe-frontend-core 0.0.195 → 0.0.196

package/src/components/chat/hooks/use-jetstream-dialog-subscription.ts

@@ -0,0 +1,474 @@
+'use client'
+
+import { useCallback, useEffect, useRef, useState } from 'react'
+import {
+  createNatsClient,
+  type JetStreamSubscriptionHandle,
+  type NatsClient,
+} from '../../../nats'
+import {
+  type UseJetStreamDialogSubscriptionOptions,
+  type UseJetStreamDialogSubscriptionReturn,
+  NETWORK_CONFIG,
+} from '../types'
+
+type SharedConnection = {
+  wsUrl: string
+  client: NatsClient
+  connectPromise: Promise<void> | null
+  refCount: number
+  closeTimer: ReturnType<typeof setTimeout> | null
+  retryTimer: ReturnType<typeof setTimeout> | null
+}
+
+let shared: SharedConnection | null = null
+
+const DEFAULT_INACTIVE_THRESHOLD_MS = 5 * 60_000
+const DEFAULT_STREAM_NAME = 'CHAT_CHUNKS'
+
+/**
+ * Subscribe to a chat dialog stream via a JetStream **ephemeral OrderedConsumer**.
+ *
+ * - Subject: `chat.{dialogId}.{topic}`
+ * - When `optStartSeq` is a number, the consumer resumes at `optStartSeq + 1`
+ *   (`DeliverPolicy.ByStartSequence`). When null/undefined, it live-tails
+ *   (`DeliverPolicy.New`).
+ * - On reconnect, the consumer is recreated starting from the highest stream
+ *   sequence we've already observed + 1, so no chunk is replayed or skipped.
+ * - Consumer is ephemeral with `AckPolicy.None` and a 5-minute inactivity
+ *   threshold (overridable via `inactiveThresholdMs`).
+ */
+export function useJetStreamDialogSubscription({
+  enabled,
+  dialogId,
+  streamName = DEFAULT_STREAM_NAME,
+  topic,
+  optStartSeq,
+  onEvent,
+  onConnect,
+  onDisconnect,
+  onSubscribed,
+  onBeforeReconnect,
+  getNatsWsUrl,
+  clientConfig = {},
+  reconnectionBackoff,
+  inactiveThresholdMs,
+}: UseJetStreamDialogSubscriptionOptions): UseJetStreamDialogSubscriptionReturn {
+  const [isConnected, setIsConnected] = useState(false)
+  const [isSubscribed, setIsSubscribed] = useState(false)
+  const [reconnectionCount, setReconnectionCount] = useState(0)
+  const [currentStreamSeq, setCurrentStreamSeq] = useState<number | null>(null)
+
+  const clientRef = useRef<NatsClient | null>(null)
+  const subscriptionRef = useRef<JetStreamSubscriptionHandle | null>(null)
+  const highestStreamSeqRef = useRef<number | null>(null)
+
+  const onEventRef = useRef(onEvent)
+  useEffect(() => {
+    onEventRef.current = onEvent
+  }, [onEvent])
+
+  const onConnectRef = useRef(onConnect)
+  useEffect(() => {
+    onConnectRef.current = onConnect
+  }, [onConnect])
+
+  const onDisconnectRef = useRef(onDisconnect)
+  useEffect(() => {
+    onDisconnectRef.current = onDisconnect
+  }, [onDisconnect])
+
+  const onSubscribedRef = useRef(onSubscribed)
+  useEffect(() => {
+    onSubscribedRef.current = onSubscribed
+  }, [onSubscribed])
+
+  const onBeforeReconnectRef = useRef(onBeforeReconnect)
+  useEffect(() => {
+    onBeforeReconnectRef.current = onBeforeReconnect
+  }, [onBeforeReconnect])
+
+  const getNatsWsUrlRef = useRef(getNatsWsUrl)
+  useEffect(() => {
+    getNatsWsUrlRef.current = getNatsWsUrl
+  }, [getNatsWsUrl])
+
+  const reconnectionBackoffRef = useRef(reconnectionBackoff)
+  useEffect(() => {
+    reconnectionBackoffRef.current = reconnectionBackoff
+  }, [reconnectionBackoff])
+
+  const optStartSeqRef = useRef<number | null | undefined>(optStartSeq)
+  useEffect(() => {
+    optStartSeqRef.current = optStartSeq
+  }, [optStartSeq])
+
+  const inactiveThresholdRef = useRef<number | undefined>(inactiveThresholdMs)
+  useEffect(() => {
+    inactiveThresholdRef.current = inactiveThresholdMs
+  }, [inactiveThresholdMs])
+
+  const hadConnectionBeforeRef = useRef(false)
+
+  const acquireClient = useCallback(
+    (url: string): SharedConnection => {
+      if (shared?.wsUrl !== url) {
+        if (shared) {
+          if (shared.closeTimer) clearTimeout(shared.closeTimer)
+          const old = shared
+          shared = null
+          void old.client.close().catch(() => {})
+        }
+
+        const { name = 'openframe-frontend-jetstream', user = 'machine', pass = '' } = clientConfig
+
+        const client = createNatsClient({
+          servers: url,
+          name,
+          user,
+          pass,
+          connectTimeoutMs: NETWORK_CONFIG.CONNECT_TIMEOUT_MS,
+          reconnect: false,
+          pingIntervalMs: NETWORK_CONFIG.PING_INTERVAL_MS,
+          maxPingOut: NETWORK_CONFIG.MAX_PING_OUT,
+        })
+        shared = {
+          wsUrl: url,
+          client,
+          connectPromise: null,
+          refCount: 0,
+          closeTimer: null,
+          retryTimer: null,
+        }
+      }
+
+      shared.refCount += 1
+      if (shared.closeTimer) {
+        clearTimeout(shared.closeTimer)
+        shared.closeTimer = null
+      }
+      return shared
+    },
+    [clientConfig],
+  )
+
+  const releaseClient = useCallback((url: string) => {
+    if (!shared || shared.wsUrl !== url) return
+
+    shared.refCount = Math.max(0, shared.refCount - 1)
+    if (shared.refCount > 0) return
+
+    shared.closeTimer = setTimeout(() => {
+      const s = shared
+      shared = null
+      if (s) {
+        if (s.retryTimer) {
+          clearTimeout(s.retryTimer)
+          s.retryTimer = null
+        }
+        void s.client.close().catch(() => {})
+      }
+    }, NETWORK_CONFIG.SHARED_CLOSE_DELAY_MS)
+  }, [])
+
+  const currentWsUrlRef = useRef<string>('')
+
+  // Connection lifecycle: acquire/release the shared client based on enabled + URL.
+  useEffect(() => {
+    const wsUrl = getNatsWsUrl()
+    if (!enabled || !wsUrl) {
+      if (currentWsUrlRef.current && clientRef.current) {
+        releaseClient(currentWsUrlRef.current)
+        clientRef.current = null
+        currentWsUrlRef.current = ''
+        setIsConnected(false)
+      }
+      return
+    }
+
+    if (
+      wsUrl === currentWsUrlRef.current &&
+      clientRef.current &&
+      clientRef.current.isConnected()
+    ) {
+      return
+    }
+
+    if (
+      currentWsUrlRef.current &&
+      currentWsUrlRef.current !== wsUrl &&
+      clientRef.current
+    ) {
+      releaseClient(currentWsUrlRef.current)
+      clientRef.current = null
+      setIsConnected(false)
+    }
+
+    currentWsUrlRef.current = wsUrl
+    const sharedConn = acquireClient(wsUrl)
+    const client = sharedConn.client
+
+    clientRef.current = client
+    setIsConnected(false)
+
+    let closed = false
+    let retryAttempt = 0
+
+    function scheduleRetry() {
+      if (closed) return
+      if (shared !== sharedConn) return
+
+      if (sharedConn.retryTimer) {
+        clearTimeout(sharedConn.retryTimer)
+        sharedConn.retryTimer = null
+      }
+
+      const cfg = reconnectionBackoffRef.current ?? {}
+      const fastRetries = cfg.fastRetries ?? 0
+      const fastDelay = cfg.fastRetryDelayMs ?? NETWORK_CONFIG.RETRY_INITIAL_DELAY_MS
+      const baseDelay = cfg.initialDelayMs ?? NETWORK_CONFIG.RETRY_INITIAL_DELAY_MS
+      const maxDelay = cfg.maxDelayMs ?? NETWORK_CONFIG.RETRY_MAX_DELAY_MS
+      const multiplier = cfg.multiplier ?? NETWORK_CONFIG.RETRY_BACKOFF_MULTIPLIER
+
+      const delay =
+        retryAttempt < fastRetries
+          ? fastDelay
+          : Math.min(baseDelay * multiplier ** (retryAttempt - fastRetries), maxDelay)
+      const jitteredDelay = delay * (0.5 + Math.random() * 0.5)
+      retryAttempt++
+
+      sharedConn.retryTimer = setTimeout(async () => {
+        sharedConn.retryTimer = null
+        if (closed) return
+        if (shared !== sharedConn) return
+
+        try {
+          await onBeforeReconnectRef.current?.()
+        } catch {
+          // Token refresh failed; still try to reconnect.
+        }
+
+        if (closed) return
+        if (shared !== sharedConn) return
+
+        const freshUrl = getNatsWsUrlRef.current()
+        if (freshUrl !== wsUrl) return
+
+        try {
+          sharedConn.connectPromise = null
+          sharedConn.connectPromise = client.connect()
+          await sharedConn.connectPromise
+          if (!closed && shared === sharedConn) {
+            retryAttempt = 0
+            setIsConnected(true)
+          }
+        } catch {
+          sharedConn.connectPromise = null
+          if (!closed && shared === sharedConn) {
+            scheduleRetry()
+          }
+        }
+      }, jitteredDelay)
+    }
+
+    const unsubscribeStatus = client.onStatus((event) => {
+      const connected = event.status === 'connected'
+      // `error` is a protocol-level signal (e.g. -ERR Permissions Violation when
+      // CONSUMER.CREATE is denied) that does NOT close the WebSocket. Treating
+      // it as a disconnect causes scheduleRetry() to fire on every -ERR, which
+      // re-runs onBeforeReconnect (auth refresh / `/api/me`) on a loop. Real
+      // transport loss arrives separately as `disconnected` or `closed`.
+      const disconnected = event.status === 'closed' || event.status === 'disconnected'
+      if (connected) {
+        setIsConnected(true)
+        if (hadConnectionBeforeRef.current) {
+          setReconnectionCount((c) => c + 1)
+        }
+        hadConnectionBeforeRef.current = true
+        retryAttempt = 0
+        onConnectRef.current?.()
+      }
+      if (event.status === 'error') {
+        // Subscription-level failures (e.g. consumer.get rejected by JetStream
+        // ACLs) already surface to the subscribe effect via the rejected
+        // promise; log here for diagnostics and let the existing WS connection
+        // keep running.
+        console.warn('[JetStream] NATS protocol error:', event.data)
+        return
+      }
+      if (disconnected) {
+        setIsConnected(false)
+        setIsSubscribed(false)
+
+        if (subscriptionRef.current) {
+          try {
+            subscriptionRef.current.unsubscribe()
+          } catch {
+            // ignore
+          }
+          subscriptionRef.current = null
+        }
+
+        onDisconnectRef.current?.()
+        scheduleRetry()
+      }
+    })
+
+    ;(async () => {
+      try {
+        sharedConn.connectPromise ||= client.connect()
+        await sharedConn.connectPromise
+        if (!closed) {
+          setIsConnected(true)
+          hadConnectionBeforeRef.current = true
+        }
+      } catch {
+        sharedConn.connectPromise = null
+        if (!closed) {
+          setIsConnected(false)
+          onDisconnectRef.current?.()
+          scheduleRetry()
+        }
+      }
+    })()
+
+    return () => {
+      closed = true
+      setIsConnected(false)
+      setIsSubscribed(false)
+      unsubscribeStatus()
+
+      if (sharedConn.retryTimer) {
+        clearTimeout(sharedConn.retryTimer)
+        sharedConn.retryTimer = null
+      }
+
+      if (subscriptionRef.current) {
+        try {
+          subscriptionRef.current.unsubscribe()
+        } catch {
+          // ignore
+        }
+        subscriptionRef.current = null
+      }
+
+      if (clientRef.current && currentWsUrlRef.current) {
+        releaseClient(currentWsUrlRef.current)
+        clientRef.current = null
+        currentWsUrlRef.current = ''
+      }
+    }
+  }, [enabled, getNatsWsUrl, acquireClient, releaseClient])
+
+  // Subscription lifecycle: (re)create the ephemeral JetStream consumer whenever
+  // we transition into a connected state for a dialog, and whenever the dialog
+  // changes. On reconnect we resume from highestStreamSeq + 1.
+  useEffect(() => {
+    if (!enabled || !dialogId || !isConnected) {
+      if (subscriptionRef.current) {
+        try {
+          subscriptionRef.current.unsubscribe()
+        } catch {
+          // ignore
+        }
+        subscriptionRef.current = null
+      }
+      setIsSubscribed(false)
+      return
+    }
+
+    const client = clientRef.current
+    if (!client) return
+
+    const abortController = new AbortController()
+    const decoder = new TextDecoder()
+    const filterSubject = `chat.${dialogId}.${topic}`
+
+    const resumeSeq = highestStreamSeqRef.current
+    const initialOptStart = optStartSeqRef.current
+    const startSeq =
+      resumeSeq != null ? resumeSeq + 1 : initialOptStart != null ? initialOptStart + 1 : undefined
+
+    let cancelled = false
+
+    void (async () => {
+      try {
+        const handle = await client.subscribeJetStreamOrdered(
+          (msg) => {
+            if (cancelled) return
+            const streamSeq = msg.info.streamSequence
+            if (typeof streamSeq === 'number') {
+              if (
+                highestStreamSeqRef.current == null ||
+                streamSeq > highestStreamSeqRef.current
+              ) {
+                highestStreamSeqRef.current = streamSeq
+                setCurrentStreamSeq(streamSeq)
+              }
+            }
+            const cb = onEventRef.current
+            if (!cb) return
+            try {
+              const parsed = JSON.parse(decoder.decode(msg.data)) as Record<string, unknown>
+              if (typeof streamSeq === 'number') {
+                ;(parsed as { streamSeq?: number }).streamSeq = streamSeq
+              }
+              cb(parsed, topic)
+            } catch {
+              // Ignore malformed payloads.
+            }
+          },
+          {
+            streamName,
+            filterSubject,
+            deliverPolicy: startSeq != null ? 'byStartSequence' : 'new',
+            optStartSeq: startSeq,
+            inactiveThresholdMs: inactiveThresholdRef.current ?? DEFAULT_INACTIVE_THRESHOLD_MS,
+            signal: abortController.signal,
+          },
+        )
+
+        if (cancelled) {
+          try {
+            handle.unsubscribe()
+          } catch {
+            // ignore
+          }
+          return
+        }
+
+        subscriptionRef.current = handle
+        setIsSubscribed(true)
+        onSubscribedRef.current?.()
+      } catch {
+        if (!cancelled) {
+          setIsSubscribed(false)
+        }
+      }
+    })()
+
+    return () => {
+      cancelled = true
+      abortController.abort()
+      if (subscriptionRef.current) {
+        try {
+          subscriptionRef.current.unsubscribe()
+        } catch {
+          // ignore
+        }
+        subscriptionRef.current = null
+      }
+      setIsSubscribed(false)
+    }
+  }, [enabled, dialogId, isConnected, streamName, topic, reconnectionCount])
+
+  // Reset the highest-seen sequence whenever the dialog changes so a new dialog
+  // starts from optStartSeq (or DeliverPolicy.New) rather than the previous
+  // dialog's offset.
+  useEffect(() => {
+    highestStreamSeqRef.current = null
+    setCurrentStreamSeq(null)
+  }, [dialogId])
+
+  return { isConnected, isSubscribed, reconnectionCount, currentStreamSeq }
+}

package/src/components/chat/types/api.types.ts

@@ -90,6 +90,51 @@ export interface UseNatsDialogSubscriptionReturn {
   reconnectionCount: number
 }
 
+export interface UseJetStreamDialogSubscriptionOptions {
+  enabled: boolean
+  dialogId: string | null
+  /** JetStream stream name. Default: 'CHAT_CHUNKS'. */
+  streamName?: string
+  /** Single topic to subscribe to. */
+  topic: NatsMessageType
+  /**
+   * Resume from this JetStream sequence + 1. When null/undefined, the consumer starts with
+   * DeliverPolicy.New (live tail only).
+   */
+  optStartSeq?: number | null
+  onEvent?: (payload: unknown, messageType: NatsMessageType) => void
+  onConnect?: () => void
+  onDisconnect?: () => void
+  onSubscribed?: () => void
+  /** Called on disconnect, before reconnect attempt. Use to refresh auth. */
+  onBeforeReconnect?: () => Promise<void> | void
+  /** Build the NATS WebSocket URL (or null when not yet available). */
+  getNatsWsUrl: () => string | null
+  clientConfig?: {
+    name?: string
+    user?: string
+    pass?: string
+  }
+  reconnectionBackoff?: {
+    fastRetries?: number
+    fastRetryDelayMs?: number
+    initialDelayMs?: number
+    maxDelayMs?: number
+    multiplier?: number
+  }
+  /** Consumer inactivity threshold in ms before NATS auto-cleans it. Default: 5 minutes. */
+  inactiveThresholdMs?: number
+}
+
+export interface UseJetStreamDialogSubscriptionReturn {
+  isConnected: boolean
+  isSubscribed: boolean
+  /** Incremented each time the underlying NATS connection reconnects. Starts at 0. */
+  reconnectionCount: number
+  /** Highest JetStream stream sequence observed so far (null before first delivery). */
+  currentStreamSeq: number | null
+}
+
 export interface SegmentsUpdateMetadata {
   /** Segments should be appended to the last assistant message */
   append?: boolean

package/src/components/chat/types/network.types.ts

@@ -34,6 +34,10 @@ export type NetworkConfig = typeof NETWORK_CONFIG
 
 export interface ChunkData {
   sequenceId?: number
+  /** JetStream stream sequence, populated from JsMsg.info.streamSequence when delivered via JetStream. */
+  streamSeq?: number
+  /** Authoritative dialog streaming state at the time this chunk was emitted; consumers should mirror it into their dialog cache. */
+  streamState?: 'IDLE' | 'STREAMING'
   type: string
   text?: string
   integratedToolType?: string

package/src/nats/nats.ts

@@ -1,11 +1,32 @@
 import type {
   ConnectionOptions,
+  Consumer,
+  ConsumerMessages,
+  JsMsg,
   MsgHdrs as NatsHeaders,
   Msg,
   NatsConnection,
   Subscription,
 } from 'nats.ws'
 
+export type JetStreamDeliverPolicy = 'new' | 'byStartSequence'
+
+export interface JetStreamOrderedSubscribeOptions {
+  streamName: string
+  filterSubject: string
+  deliverPolicy: JetStreamDeliverPolicy
+  /** Required when deliverPolicy === 'byStartSequence'. */
+  optStartSeq?: number
+  /** Auto-cleanup the ephemeral consumer after this idle time. Default: 5 minutes. */
+  inactiveThresholdMs?: number
+  /** AbortSignal to tear down the consumer. */
+  signal?: AbortSignal
+}
+
+export interface JetStreamSubscriptionHandle {
+  unsubscribe(): void
+}
+
 export interface NatsClientOptions {
   /**
    * NATS server URL(s), for example:
@@ -147,6 +168,16 @@ export interface NatsClient {
     options?: NatsSubscribeOptions,
   ): NatsSubscriptionHandle
 
+  /**
+   * Subscribe to a JetStream subject via an ephemeral OrderedConsumer (no acks).
+   * Use `optStartSeq` with `deliverPolicy: 'byStartSequence'` to resume from a known offset,
+   * or `deliverPolicy: 'new'` to live-tail.
+   */
+  subscribeJetStreamOrdered(
+    onMessage: (msg: JsMsg) => void | Promise<void>,
+    options: JetStreamOrderedSubscribeOptions,
+  ): Promise<JetStreamSubscriptionHandle>
+
   onStatus(listener: (event: NatsStatusEvent) => void): () => void
 }
 
@@ -473,6 +504,91 @@ export function createNatsClient(options: NatsClientOptions): NatsClient {
     )
   }
 
+  async function subscribeJetStreamOrdered(
+    onMessage: (msg: JsMsg) => void | Promise<void>,
+    opts: JetStreamOrderedSubscribeOptions,
+  ): Promise<JetStreamSubscriptionHandle> {
+    const conn = requireConnection()
+    if (opts.signal?.aborted) {
+      return { unsubscribe() {} }
+    }
+
+    const nats = await importNats()
+    if (opts.signal?.aborted) {
+      return { unsubscribe() {} }
+    }
+
+    const inactiveThresholdNs =
+      (opts.inactiveThresholdMs ?? 5 * 60_000) * 1_000_000
+
+    const js = conn.jetstream()
+    const consumer: Consumer = await js.consumers.get(opts.streamName, {
+      filterSubjects: opts.filterSubject,
+      deliver_policy: nats.DeliverPolicy.StartSequence,
+      opt_start_seq: opts.optStartSeq ?? 0,
+      inactive_threshold: inactiveThresholdNs,
+    })
+
+    const iterRef: { current: ConsumerMessages | null } = { current: null }
+    let closed = false
+
+    const onAbort = () => {
+      void teardown()
+    }
+    opts.signal?.addEventListener('abort', onAbort, { once: true })
+
+    async function teardown(): Promise<void> {
+      if (closed) return
+      closed = true
+      opts.signal?.removeEventListener('abort', onAbort)
+      const iter = iterRef.current
+      iterRef.current = null
+      if (iter) {
+        try {
+          await iter.close()
+        } catch {
+          // ignore
+        }
+      }
+    }
+
+    if (opts.signal?.aborted) {
+      void teardown()
+      return { unsubscribe() {} }
+    }
+
+    ;(async () => {
+      try {
+        const iter = await consumer.consume()
+        if (closed) {
+          try {
+            await iter.close()
+          } catch {
+            // ignore
+          }
+          return
+        }
+        iterRef.current = iter
+        for await (const msg of iter) {
+          if (closed) break
+          try {
+            await onMessage(msg)
+          } catch (e) {
+            emitStatus({ status: 'error', data: e })
+          }
+        }
+      } catch (e) {
+        if (!closed) emitStatus({ status: 'error', data: e })
+      }
+    })().catch((e) => emitStatus({ status: 'error', data: e }))
+
+    return {
+      unsubscribe() {
+        void teardown()
+      },
+    }
+  }
+
   function onStatus(listener: (event: NatsStatusEvent) => void): () => void {
     statusListeners.add(listener)
     return () => statusListeners.delete(listener)
@@ -491,6 +607,7 @@ export function createNatsClient(options: NatsClientOptions): NatsClient {
     subscribeBytes,
     subscribeString,
     subscribeJson,
+    subscribeJetStreamOrdered,
     onStatus,
   }
 }