@memori.ai/memori-react 8.40.3 → 8.41.0

package/src/helpers/nats/useNatsSession.ts

@@ -1,11 +1,23 @@
 // helpers/nats/useNatsSession.ts - Subscribe to a session's NATS channel.
 //
-// Opens a W3C WebSocket connection (via `wsconnect`), subscribes to the
-// session subject (`sessions.<sessionId>`) and forwards every decoded event to
-// `onMessage`. The cleanup (`nc.close()`) is essential: without it, every
-// `sessionId` change would leave orphan connections / iterators.
-import { useEffect, useRef } from 'react';
+// Opens a W3C WebSocket connection (via `wsconnect`), then either:
+// - core NATS: subscribes to the session subject directly, or
+// - JetStream: consumes from the configured stream (ordered consumer by default).
+//
+// When JetStream is enabled, reconnects on tab visibility / network resume and
+// replays the last message per subject to recover events missed while Safari
+// (or other browsers) suspend the WebSocket in background.
+//
+// The cleanup (`consumerMessages.stop()` + `nc.close()`) is essential: without
+// it, every `sessionId` change would leave orphan connections / iterators.
+import { useEffect, useRef, useState, type MutableRefObject } from 'react';
 import { wsconnect, NatsConnection } from '@nats-io/nats-core';
+import {
+  jetstream,
+  DeliverPolicy,
+  type ConsumerMessages,
+} from '@nats-io/jetstream';
+import type { NatsConfig } from './getNatsConfig';
 
 /**
  * `progress` event: job advancement updates.
@@ -52,6 +64,8 @@ export type NatsSessionEvent =
   | NatsDialogResponseEvent
   | NatsErrorEvent;
 
+type JetStreamConsumeMode = 'live' | 'catch-up';
+
 function readString(
   raw: Record<string, unknown>,
   camelKey: string,
@@ -143,49 +157,198 @@ export function normalizeNatsEvent(raw: Record<string, unknown>): NatsSessionEve
   } as NatsSessionEvent;
 }
 
+function decodeMessage(
+  subject: string,
+  received: number,
+  raw: Record<string, unknown>,
+  onMessage: (event: NatsSessionEvent) => void
+) {
+  const payload = normalizeNatsEvent(raw);
+  console.debug(
+    `[NATS] message #${received} on ${subject} (eventType=${payload.eventType})`,
+    payload
+  );
+  onMessage(payload);
+}
+
+async function consumeCoreSubscription(
+  nc: NatsConnection,
+  subject: string,
+  onMessage: (event: NatsSessionEvent) => void
+) {
+  const sub = nc.subscribe(subject);
+  console.info('[NATS] subscribed to subject', subject);
+  let received = 0;
+  for await (const msg of sub) {
+    received += 1;
+    try {
+      decodeMessage(
+        subject,
+        received,
+        msg.json<Record<string, unknown>>(),
+        onMessage
+      );
+    } catch (err) {
+      console.error('[NATS] message decode error', err);
+    }
+  }
+  console.info(
+    `[NATS] subscription on ${subject} closed (received ${received} messages)`
+  );
+}
+
+async function consumeJetStream(
+  nc: NatsConnection,
+  config: Pick<NatsConfig, 'stream' | 'consumer'>,
+  subject: string,
+  mode: JetStreamConsumeMode,
+  onMessage: (event: NatsSessionEvent) => void,
+  setConsumerMessages: (messages: ConsumerMessages) => void
+) {
+  const js = jetstream(nc);
+  const consumer = config.consumer
+    ? await js.consumers.get(config.stream!, config.consumer)
+    : await js.consumers.get(config.stream!, {
+        filter_subjects: [subject],
+        deliver_policy:
+          mode === 'catch-up'
+            ? DeliverPolicy.LastPerSubject
+            : DeliverPolicy.New,
+      });
+
+  console.info(
+    '[NATS] JetStream consumer ready',
+    config.consumer
+      ? { stream: config.stream, consumer: config.consumer, subject, mode }
+      : { stream: config.stream, subject, mode }
+  );
+
+  const messages = await consumer.consume();
+  setConsumerMessages(messages);
+  let received = 0;
+  for await (const msg of messages) {
+    received += 1;
+    try {
+      decodeMessage(subject, received, msg.json(), onMessage);
+    } catch (err) {
+      console.error('[NATS] message decode error', err);
+    }
+  }
+  console.info(
+    `[NATS] JetStream consumer on ${subject} closed (received ${received} messages)`
+  );
+}
+
+function closeNatsSession(
+  consumerMessagesRef: MutableRefObject<ConsumerMessages | null>,
+  connRef: MutableRefObject<NatsConnection | null>
+) {
+  consumerMessagesRef.current?.stop();
+  consumerMessagesRef.current = null;
+  connRef.current?.close();
+  connRef.current = null;
+}
+
 /**
  * Open a NATS WebSocket connection, subscribe to the session subject and
  * forward decoded events to `onMessage`.
  *
  * `onMessage` is kept in a ref so callers do not need to memoize it: the
- * connection is only (re)opened when `sessionId`, `natsWsUrl` or `natsToken`
+ * connection is only (re)opened when `sessionId`, `config` or `resumeGeneration`
  * change.
  *
  * @param sessionId Current session UUID (subscription is skipped when falsy).
- * @param natsWsUrl WebSocket URL of the NATS server.
- * @param natsToken Bearer token for authentication.
+ * @param config Connection config from `/api/nats` (url + token required).
  * @param onMessage Callback invoked for each decoded event.
  */
 export function useNatsSession(
   sessionId: string | undefined,
-  natsWsUrl: string | undefined,
-  natsToken: string | undefined,
+  config: NatsConfig | undefined,
   onMessage: (event: NatsSessionEvent) => void
 ) {
   const connRef = useRef<NatsConnection | null>(null);
+  const consumerMessagesRef = useRef<ConsumerMessages | null>(null);
   const onMessageRef = useRef(onMessage);
+  const wasHiddenRef = useRef(false);
+  const [resumeGeneration, setResumeGeneration] = useState(0);
 
   // Keep the latest callback without retriggering the connection effect.
   useEffect(() => {
     onMessageRef.current = onMessage;
   }, [onMessage]);
 
+  // Reconnect when the tab returns from background or the network comes back.
   useEffect(() => {
-    if (!sessionId || !natsWsUrl || !natsToken) {
+    if (!sessionId || !config?.stream) {
+      return;
+    }
+
+    const requestResume = () => {
+      if (document.visibilityState !== 'visible') {
+        return;
+      }
+      console.info('[NATS] resuming JetStream after visibility/network change');
+      setResumeGeneration(generation => generation + 1);
+    };
+
+    const onVisibilityChange = () => {
+      if (document.visibilityState === 'hidden') {
+        wasHiddenRef.current = true;
+        return;
+      }
+      if (wasHiddenRef.current) {
+        wasHiddenRef.current = false;
+        requestResume();
+      }
+    };
+
+    const onPageShow = (event: PageTransitionEvent) => {
+      if (event.persisted) {
+        requestResume();
+      }
+    };
+
+    document.addEventListener('visibilitychange', onVisibilityChange);
+    window.addEventListener('online', requestResume);
+    window.addEventListener('pageshow', onPageShow);
+
+    return () => {
+      document.removeEventListener('visibilitychange', onVisibilityChange);
+      window.removeEventListener('online', requestResume);
+      window.removeEventListener('pageshow', onPageShow);
+    };
+  }, [sessionId, config?.stream]);
+
+  useEffect(() => {
+    if (!sessionId || !config?.url || !config.token) {
       console.debug(
         '[NATS] subscription skipped (missing sessionId/url/token)',
-        { sessionId, hasUrl: !!natsWsUrl, hasToken: !!natsToken }
+        {
+          sessionId,
+          hasUrl: !!config?.url,
+          hasToken: !!config?.token,
+        }
       );
       return;
     }
+
     let closed = false;
-    const subject = sessionId;
+    const subject = config.subject ?? `sessions.${sessionId}`;
+    const useJetStream = !!config.stream;
+    const jetStreamMode: JetStreamConsumeMode =
+      useJetStream && resumeGeneration > 0 ? 'catch-up' : 'live';
 
     (async () => {
-      console.info('[NATS] connecting to', natsWsUrl, 'for session', sessionId);
+      console.info(
+        '[NATS] connecting to',
+        config.url,
+        'for session',
+        sessionId,
+        useJetStream ? `(JetStream/${jetStreamMode})` : '(core)'
+      );
       const nc = await wsconnect({
-        servers: [natsWsUrl],
-        token: natsToken,
+        servers: [config.url],
+        token: config.token,
       });
       if (closed) {
         console.debug(
@@ -197,34 +360,35 @@ export function useNatsSession(
       connRef.current = nc;
       console.info('[NATS] connected to', nc.getServer());
 
-      const sub = nc.subscribe(subject);
-      console.info('[NATS] subscribed to subject', subject);
-      let received = 0;
-      for await (const msg of sub) {
-        received += 1;
-        try {
-          const payload = normalizeNatsEvent(
-            msg.json<Record<string, unknown>>()
-          );
-          console.debug(
-            `[NATS] message #${received} on ${subject} (eventType=${payload.eventType})`,
-            payload
-          );
-          onMessageRef.current(payload);
-        } catch (err) {
-          console.error('[NATS] message decode error', err);
-        }
+      const dispatch = (event: NatsSessionEvent) => {
+        onMessageRef.current(event);
+      };
+
+      if (useJetStream) {
+        await consumeJetStream(
+          nc,
+          config,
+          subject,
+          jetStreamMode,
+          dispatch,
+          messages => {
+            consumerMessagesRef.current = messages;
+          }
+        );
+      } else {
+        await consumeCoreSubscription(nc, subject, dispatch);
       }
-      console.info(
-        `[NATS] subscription on ${subject} closed (received ${received} messages)`
-      );
     })().catch(err => console.error('[NATS] connection error', err));
 
     return () => {
       console.info('[NATS] cleanup: closing connection for subject', subject);
       closed = true;
-      connRef.current?.close();
-      connRef.current = null;
+      closeNatsSession(consumerMessagesRef, connRef);
     };
-  }, [sessionId, natsWsUrl, natsToken]);
+  }, [sessionId, config, resumeGeneration]);
+
+  // Reset catch-up generation when the active session changes.
+  useEffect(() => {
+    setResumeGeneration(0);
+  }, [sessionId]);
 }

package/src/version.ts

@@ -1,2 +1,2 @@
 // This file is auto-generated. Do not edit manually.
-export const version = '8.40.3';
+export const version = '8.41.0';