🫧 refactor: Clear Drafts and Surface Error on Expired SSE Stream (#12309)

* refactor: error handling in useResumableSSE for 404 responses

- Added logic to clear drafts from localStorage when a 404 error occurs.
- Integrated errorHandler to notify users of the error condition.
- Introduced comprehensive tests to validate the new behavior, ensuring drafts are cleared and error handling is triggered correctly.C

* feat: add STREAM_EXPIRED error handling and message localization

- Introduced handling for STREAM_EXPIRED errors in useResumableSSE, updating errorHandler to provide relevant feedback.
- Added a new error message for STREAM_EXPIRED in translation files for user notifications.
- Updated tests to ensure proper error handling and message verification for STREAM_EXPIRED scenarios.

* refactor: replace clearDraft with clearAllDrafts utility

- Removed the clearDraft function from useResumableSSE and useSSE hooks, replacing it with the new clearAllDrafts utility for better draft management.
- Updated localStorage interactions to ensure both text and file drafts are cleared consistently for a conversation.
- Enhanced code readability and maintainability by centralizing draft clearing logic.

client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts

@@ -0,0 +1,273 @@
+import { renderHook, act } from '@testing-library/react';
+import { Constants, ErrorTypes, LocalStorageKeys } from 'librechat-data-provider';
+import type { TSubmission } from 'librechat-data-provider';
+
+type SSEEventListener = (e: Partial<MessageEvent> & { responseCode?: number }) => void;
+
+interface MockSSEInstance {
+  addEventListener: jest.Mock;
+  stream: jest.Mock;
+  close: jest.Mock;
+  headers: Record<string, string>;
+  _listeners: Record<string, SSEEventListener>;
+  _emit: (event: string, data?: Partial<MessageEvent> & { responseCode?: number }) => void;
+}
+
+const mockSSEInstances: MockSSEInstance[] = [];
+
+jest.mock('sse.js', () => ({
+  SSE: jest.fn().mockImplementation(() => {
+    const listeners: Record<string, SSEEventListener> = {};
+    const instance: MockSSEInstance = {
+      addEventListener: jest.fn((event: string, cb: SSEEventListener) => {
+        listeners[event] = cb;
+      }),
+      stream: jest.fn(),
+      close: jest.fn(),
+      headers: {},
+      _listeners: listeners,
+      _emit: (event, data = {}) => listeners[event]?.(data as MessageEvent),
+    };
+    mockSSEInstances.push(instance);
+    return instance;
+  }),
+}));
+
+const mockSetQueryData = jest.fn();
+const mockQueryClient = { setQueryData: mockSetQueryData };
+
+jest.mock('@tanstack/react-query', () => ({
+  ...jest.requireActual('@tanstack/react-query'),
+  useQueryClient: () => mockQueryClient,
+}));
+
+jest.mock('recoil', () => ({
+  ...jest.requireActual('recoil'),
+  useSetRecoilState: () => jest.fn(),
+}));
+
+jest.mock('~/store', () => ({
+  __esModule: true,
+  default: {
+    activeRunFamily: jest.fn(),
+    abortScrollFamily: jest.fn(),
+    showStopButtonByIndex: jest.fn(),
+  },
+}));
+
+jest.mock('~/hooks/AuthContext', () => ({
+  useAuthContext: () => ({ token: 'test-token', isAuthenticated: true }),
+}));
+
+jest.mock('~/data-provider', () => ({
+  useGetStartupConfig: () => ({ data: { balance: { enabled: false } } }),
+  useGetUserBalance: () => ({ refetch: jest.fn() }),
+  queueTitleGeneration: jest.fn(),
+}));
+
+const mockErrorHandler = jest.fn();
+const mockSetIsSubmitting = jest.fn();
+const mockClearStepMaps = jest.fn();
+
+jest.mock('~/hooks/SSE/useEventHandlers', () =>
+  jest.fn(() => ({
+    errorHandler: mockErrorHandler,
+    finalHandler: jest.fn(),
+    createdHandler: jest.fn(),
+    attachmentHandler: jest.fn(),
+    stepHandler: jest.fn(),
+    contentHandler: jest.fn(),
+    resetContentHandler: jest.fn(),
+    syncStepMessage: jest.fn(),
+    clearStepMaps: mockClearStepMaps,
+    messageHandler: jest.fn(),
+    setIsSubmitting: mockSetIsSubmitting,
+    setShowStopButton: jest.fn(),
+  })),
+);
+
+jest.mock('librechat-data-provider', () => {
+  const actual = jest.requireActual('librechat-data-provider');
+  return {
+    ...actual,
+    createPayload: jest.fn(() => ({
+      payload: { model: 'gpt-4o' },
+      server: '/api/agents/chat',
+    })),
+    removeNullishValues: jest.fn((v: unknown) => v),
+    apiBaseUrl: jest.fn(() => ''),
+    request: {
+      post: jest.fn().mockResolvedValue({ streamId: 'stream-123' }),
+      refreshToken: jest.fn(),
+      dispatchTokenUpdatedEvent: jest.fn(),
+    },
+  };
+});
+
+import useResumableSSE from '~/hooks/SSE/useResumableSSE';
+
+const CONV_ID = 'conv-abc-123';
+
+type PartialSubmission = {
+  conversation: { conversationId?: string };
+  userMessage: Record<string, unknown>;
+  messages: never[];
+  isTemporary: boolean;
+  initialResponse: Record<string, unknown>;
+  endpointOption: { endpoint: string };
+};
+
+const buildSubmission = (overrides: Partial<PartialSubmission> = {}): TSubmission => {
+  const conversationId = overrides.conversation?.conversationId ?? CONV_ID;
+  return {
+    conversation: { conversationId },
+    userMessage: {
+      messageId: 'msg-1',
+      conversationId,
+      text: 'Hello',
+      isCreatedByUser: true,
+      sender: 'User',
+      parentMessageId: '00000000-0000-0000-0000-000000000000',
+    },
+    messages: [],
+    isTemporary: false,
+    initialResponse: {
+      messageId: 'resp-1',
+      conversationId,
+      text: '',
+      isCreatedByUser: false,
+      sender: 'Assistant',
+    },
+    endpointOption: { endpoint: 'agents' },
+    ...overrides,
+  } as unknown as TSubmission;
+};
+
+const buildChatHelpers = () => ({
+  setMessages: jest.fn(),
+  getMessages: jest.fn(() => []),
+  setConversation: jest.fn(),
+  setIsSubmitting: mockSetIsSubmitting,
+  newConversation: jest.fn(),
+  resetLatestMessage: jest.fn(),
+});
+
+const getLastSSE = (): MockSSEInstance => {
+  const sse = mockSSEInstances[mockSSEInstances.length - 1];
+  expect(sse).toBeDefined();
+  return sse;
+};
+
+describe('useResumableSSE - 404 error path', () => {
+  beforeEach(() => {
+    mockSSEInstances.length = 0;
+    localStorage.clear();
+  });
+
+  const seedDraft = (conversationId: string) => {
+    localStorage.setItem(`${LocalStorageKeys.TEXT_DRAFT}${conversationId}`, 'draft text');
+    localStorage.setItem(`${LocalStorageKeys.FILES_DRAFT}${conversationId}`, '[]');
+  };
+
+  const render404Scenario = async (conversationId = CONV_ID) => {
+    const submission = buildSubmission({ conversation: { conversationId } });
+    const chatHelpers = buildChatHelpers();
+
+    const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers));
+
+    await act(async () => {
+      await Promise.resolve();
+    });
+
+    const sse = getLastSSE();
+
+    await act(async () => {
+      sse._emit('error', { responseCode: 404 });
+    });
+
+    return { sse, unmount, chatHelpers };
+  };
+
+  it('clears the text and files draft from localStorage on 404', async () => {
+    seedDraft(CONV_ID);
+    expect(localStorage.getItem(`${LocalStorageKeys.TEXT_DRAFT}${CONV_ID}`)).not.toBeNull();
+    expect(localStorage.getItem(`${LocalStorageKeys.FILES_DRAFT}${CONV_ID}`)).not.toBeNull();
+
+    const { unmount } = await render404Scenario(CONV_ID);
+
+    expect(localStorage.getItem(`${LocalStorageKeys.TEXT_DRAFT}${CONV_ID}`)).toBeNull();
+    expect(localStorage.getItem(`${LocalStorageKeys.FILES_DRAFT}${CONV_ID}`)).toBeNull();
+    unmount();
+  });
+
+  it('calls errorHandler with STREAM_EXPIRED error type on 404', async () => {
+    const { unmount } = await render404Scenario(CONV_ID);
+
+    expect(mockErrorHandler).toHaveBeenCalledTimes(1);
+    const call = mockErrorHandler.mock.calls[0][0];
+    expect(call.data).toBeDefined();
+    const parsed = JSON.parse(call.data.text);
+    expect(parsed.type).toBe(ErrorTypes.STREAM_EXPIRED);
+    expect(call.submission).toEqual(
+      expect.objectContaining({
+        conversation: expect.objectContaining({ conversationId: CONV_ID }),
+      }),
+    );
+    unmount();
+  });
+
+  it('clears both TEXT and FILES drafts for new-convo when conversationId is absent', async () => {
+    localStorage.setItem(`${LocalStorageKeys.TEXT_DRAFT}${Constants.NEW_CONVO}`, 'unsent message');
+    localStorage.setItem(`${LocalStorageKeys.FILES_DRAFT}${Constants.NEW_CONVO}`, '[]');
+
+    const submission = buildSubmission({ conversation: {} });
+    const chatHelpers = buildChatHelpers();
+
+    const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers));
+
+    await act(async () => {
+      await Promise.resolve();
+    });
+
+    const sse = getLastSSE();
+    await act(async () => {
+      sse._emit('error', { responseCode: 404 });
+    });
+
+    expect(localStorage.getItem(`${LocalStorageKeys.TEXT_DRAFT}${Constants.NEW_CONVO}`)).toBeNull();
+    expect(
+      localStorage.getItem(`${LocalStorageKeys.FILES_DRAFT}${Constants.NEW_CONVO}`),
+    ).toBeNull();
+    unmount();
+  });
+
+  it('closes the SSE connection on 404', async () => {
+    const { sse, unmount } = await render404Scenario();
+
+    expect(sse.close).toHaveBeenCalled();
+    unmount();
+  });
+
+  it.each([undefined, 500, 503])(
+    'does not call errorHandler for responseCode %s (reconnect path)',
+    async (responseCode) => {
+      const submission = buildSubmission();
+      const chatHelpers = buildChatHelpers();
+
+      const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers));
+
+      await act(async () => {
+        await Promise.resolve();
+      });
+
+      const sse = getLastSSE();
+
+      await act(async () => {
+        sse._emit('error', { responseCode });
+      });
+
+      expect(mockErrorHandler).not.toHaveBeenCalled();
+      unmount();
+    },
+  );
+});

client/src/hooks/SSE/useResumableSSE.ts

@@ -11,7 +11,6 @@ import {
   apiBaseUrl,
   createPayload,
   ViolationTypes,
-  LocalStorageKeys,
   removeNullishValues,
 } from 'librechat-data-provider';
 import type { TMessage, TPayload, TSubmission, EventSubmission } from 'librechat-data-provider';
@@ -20,18 +19,9 @@ import { useGetStartupConfig, useGetUserBalance, queueTitleGeneration } from '~/
 import type { ActiveJobsResponse } from '~/data-provider';
 import { useAuthContext } from '~/hooks/AuthContext';
 import useEventHandlers from './useEventHandlers';
+import { clearAllDrafts } from '~/utils';
 import store from '~/store';
 
-const clearDraft = (conversationId?: string | null) => {
-  if (conversationId) {
-    localStorage.removeItem(`${LocalStorageKeys.TEXT_DRAFT}${conversationId}`);
-    localStorage.removeItem(`${LocalStorageKeys.FILES_DRAFT}${conversationId}`);
-  } else {
-    localStorage.removeItem(`${LocalStorageKeys.TEXT_DRAFT}${Constants.NEW_CONVO}`);
-    localStorage.removeItem(`${LocalStorageKeys.FILES_DRAFT}${Constants.NEW_CONVO}`);
-  }
-};
-
 type ChatHelpers = Pick<
   EventHandlerParams,
   | 'setMessages'
@@ -176,7 +166,7 @@ export default function useResumableSSE(
               conversationId: data.conversation?.conversationId,
               hasResponseMessage: !!data.responseMessage,
             });
-            clearDraft(currentSubmission.conversation?.conversationId);
+            clearAllDrafts(currentSubmission.conversation?.conversationId);
             try {
               finalHandler(data, currentSubmission as EventSubmission);
             } catch (error) {
@@ -357,7 +347,13 @@ export default function useResumableSSE(
           console.log('[ResumableSSE] Stream not found (404) - job completed or expired');
           sse.close();
           removeActiveJob(currentStreamId);
-          setIsSubmitting(false);
+          clearAllDrafts(currentSubmission.conversation?.conversationId);
+          errorHandler({
+            data: {
+              text: JSON.stringify({ type: ErrorTypes.STREAM_EXPIRED }),
+            } as unknown as Parameters<typeof errorHandler>[0]['data'],
+            submission: currentSubmission as EventSubmission,
+          });
           setShowStopButton(false);
           setStreamId(null);
           reconnectAttemptRef.current = 0;

client/src/hooks/SSE/useSSE.ts

@@ -2,32 +2,16 @@ import { useEffect, useState } from 'react';
 import { v4 } from 'uuid';
 import { SSE } from 'sse.js';
 import { useSetRecoilState } from 'recoil';
-import {
-  request,
-  Constants,
-  /* @ts-ignore */
-  createPayload,
-  LocalStorageKeys,
-  removeNullishValues,
-} from 'librechat-data-provider';
+import { request, createPayload, removeNullishValues } from 'librechat-data-provider';
 import type { TMessage, TPayload, TSubmission, EventSubmission } from 'librechat-data-provider';
 import type { EventHandlerParams } from './useEventHandlers';
 import type { TResData } from '~/common';
 import { useGetStartupConfig, useGetUserBalance } from '~/data-provider';
 import { useAuthContext } from '~/hooks/AuthContext';
 import useEventHandlers from './useEventHandlers';
+import { clearAllDrafts } from '~/utils';
 import store from '~/store';
 
-const clearDraft = (conversationId?: string | null) => {
-  if (conversationId) {
-    localStorage.removeItem(`${LocalStorageKeys.TEXT_DRAFT}${conversationId}`);
-    localStorage.removeItem(`${LocalStorageKeys.FILES_DRAFT}${conversationId}`);
-  } else {
-    localStorage.removeItem(`${LocalStorageKeys.TEXT_DRAFT}${Constants.NEW_CONVO}`);
-    localStorage.removeItem(`${LocalStorageKeys.FILES_DRAFT}${Constants.NEW_CONVO}`);
-  }
-};
-
 type ChatHelpers = Pick<
   EventHandlerParams,
   | 'setMessages'
@@ -120,7 +104,7 @@ export default function useSSE(
       const data = JSON.parse(e.data);
 
       if (data.final != null) {
-        clearDraft(submission.conversation?.conversationId);
+        clearAllDrafts(submission.conversation?.conversationId);
         try {
           finalHandler(data, submission as EventSubmission);
         } catch (error) {