🛣️ fix: Chat Stream Hangup (#4822)

Embedded sse.js code converted into an external
dependency.
Custom access token refresh logic moved to
useSSE.ts hook.

Closes #4820

Author: thingersoft

client/package.json

@@ -96,6 +96,7 @@
     "remark-gfm": "^4.0.0",
     "remark-math": "^6.0.0",
     "remark-supersub": "^1.0.0",
+    "sse.js": "^2.5.0",
     "tailwind-merge": "^1.9.1",
     "tailwindcss-animate": "^1.0.5",
     "tailwindcss-radix": "^2.8.0",

client/src/hooks/SSE/useSSE.ts

@@ -1,22 +1,23 @@
-import { v4 } from 'uuid';
-import { useSetRecoilState } from 'recoil';
-import { useEffect, useState } from 'react';
+import type { EventSubmission, TMessage, TPayload, TSubmission } from 'librechat-data-provider';
 import {
   /* @ts-ignore */
-  SSE,
   createPayload,
   isAgentsEndpoint,
-  removeNullishValues,
   isAssistantsEndpoint,
+  removeNullishValues,
+  request,
 } from 'librechat-data-provider';
-import { useGetUserBalance, useGetStartupConfig } from 'librechat-data-provider/react-query';
-import type { TMessage, TSubmission, TPayload, EventSubmission } from 'librechat-data-provider';
-import type { EventHandlerParams } from './useEventHandlers';
+import { useGetStartupConfig, useGetUserBalance } from 'librechat-data-provider/react-query';
+import { useEffect, useState } from 'react';
+import { useSetRecoilState } from 'recoil';
+import { SSE } from 'sse.js';
+import { v4 } from 'uuid';
 import type { TResData } from '~/common';
 import { useGenTitleMutation } from '~/data-provider';
 import { useAuthContext } from '~/hooks/AuthContext';
-import useEventHandlers from './useEventHandlers';
 import store from '~/store';
+import type { EventHandlerParams } from './useEventHandlers';
+import useEventHandlers from './useEventHandlers';
 
 type ChatHelpers = Pick<
   EventHandlerParams,
@@ -94,21 +95,21 @@ export default function useSSE(
 
     let textIndex = null;
 
-    const events = new SSE(payloadData.server, {
+    const sse = new SSE(payloadData.server, {
       payload: JSON.stringify(payload),
       headers: { 'Content-Type': 'application/json', Authorization: `Bearer ${token}` },
     });
 
-    events.onattachment = (e: MessageEvent) => {
+    sse.addEventListener('attachment', (e: MessageEvent) => {
       try {
         const data = JSON.parse(e.data);
         attachmentHandler({ data, submission: submission as EventSubmission });
       } catch (error) {
         console.error(error);
       }
-    };
+    });
 
-    events.onmessage = (e: MessageEvent) => {
+    sse.addEventListener('message', (e: MessageEvent) => {
       const data = JSON.parse(e.data);
 
       if (data.final != null) {
@@ -155,14 +156,14 @@ export default function useSSE(
           messageHandler(text, { ...submission, plugin, plugins, userMessage, initialResponse });
         }
       }
-    };
+    });
 
-    events.onopen = () => {
+    sse.addEventListener('open', () => {
       setAbortScroll(false);
       console.log('connection is opened');
-    };
+    });
 
-    events.oncancel = async () => {
+    sse.addEventListener('cancel', async () => {
       const streamKey = (submission as TSubmission | null)?.['initialResponse']?.messageId;
       if (completed.has(streamKey)) {
         setIsSubmitting(false);
@@ -181,9 +182,27 @@ export default function useSSE(
         submission as EventSubmission,
         latestMessages,
       );
-    };
+    });
+
+    sse.addEventListener('error', async (e: MessageEvent) => {
+      /* @ts-ignore */
+      if (e.responseCode === 401) {
+        /* token expired, refresh and retry */
+        try {
+          const refreshResponse = await request.refreshToken();
+          sse.headers = {
+            'Content-Type': 'application/json',
+            Authorization: `Bearer ${refreshResponse.token}`,
+          };
+          request.dispatchTokenUpdatedEvent(refreshResponse.token);
+          sse.stream();
+          return;
+        } catch (error) {
+          /* token refresh failed, continue handling the original 401 */
+          console.log(error);
+        }
+      }
 
-    events.onerror = function (e: MessageEvent) {
       console.log('error in server stream.');
       (startupConfig?.checkBalance ?? false) && balanceQuery.refetch();
 
@@ -197,18 +216,18 @@ export default function useSSE(
       }
 
       errorHandler({ data, submission: { ...submission, userMessage } as EventSubmission });
-    };
+    });
 
     setIsSubmitting(true);
-    events.stream();
+    sse.stream();
 
     return () => {
-      const isCancelled = events.readyState <= 1;
-      events.close();
-      // setSource(null);
+      const isCancelled = sse.readyState <= 1;
+      sse.close();
       if (isCancelled) {
         const e = new Event('cancel');
-        events.dispatchEvent(e);
+        /* @ts-ignore */
+        sse.dispatchEvent(e);
       }
     };
     // eslint-disable-next-line react-hooks/exhaustive-deps

package-lock.json

@@ -8,26 +8,25 @@ export * from './artifacts';
 /* schema helpers  */
 export * from './parsers';
 /* custom/dynamic configurations  */
-export * from './models';
 export * from './generate';
+export * from './models';
 /* RBAC */
 export * from './roles';
 /* types (exports schemas from `./types` as they contain needed in other defs) */
 export * from './types';
 export * from './types/agents';
 export * from './types/assistants';
-export * from './types/queries';
 export * from './types/files';
 export * from './types/mutations';
+export * from './types/queries';
 export * from './types/runs';
 /* query/mutation keys */
 export * from './keys';
 /* api call helpers */
 export * from './headers-helpers';
 export { default as request } from './request';
-import * as dataService from './data-service';
 export { dataService };
+import * as dataService from './data-service';
 /* general helpers */
-export * from './sse';
 export * from './actions';
 export { default as createPayload } from './createPayload';

packages/data-provider/src/index.ts

@@ -1,7 +1,7 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
-import axios, { AxiosRequestConfig, AxiosError } from 'axios';
-import { setTokenHeader } from './headers-helpers';
+import axios, { AxiosError, AxiosRequestConfig } from 'axios';
 import * as endpoints from './api-endpoints';
+import { setTokenHeader } from './headers-helpers';
 
 async function _get<T>(url: string, options?: AxiosRequestConfig): Promise<T> {
   const response = await axios.get(url, { ...options });
@@ -65,6 +65,11 @@ let failedQueue: { resolve: (value?: any) => void; reject: (reason?: any) => voi
 
 const refreshToken = (retry?: boolean) => _post(endpoints.refreshToken(retry));
 
+const dispatchTokenUpdatedEvent = (token: string) => {
+  setTokenHeader(token);
+  window.dispatchEvent(new CustomEvent('tokenUpdated', { detail: token }));
+};
+
 const processQueue = (error: AxiosError | null, token: string | null = null) => {
   failedQueue.forEach((prom) => {
     if (error) {
@@ -109,8 +114,7 @@ axios.interceptors.response.use(
 
         if (token) {
           originalRequest.headers['Authorization'] = 'Bearer ' + token;
-          setTokenHeader(token);
-          window.dispatchEvent(new CustomEvent('tokenUpdated', { detail: token }));
+          dispatchTokenUpdatedEvent(token);
           processQueue(null, token);
           return await axios(originalRequest);
         } else {
@@ -139,4 +143,5 @@ export default {
   deleteWithOptions: _deleteWithOptions,
   patch: _patch,
   refreshToken,
+  dispatchTokenUpdatedEvent,
 };