From 9cb5ac63f89210dc506f7cdd64f127265629d781 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 19 Mar 2026 14:51:28 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=AB=A7=20refactor:=20Clear=20Drafts=20and?= =?UTF-8?q?=20Surface=20Error=20on=20Expired=20SSE=20Stream=20(#12309)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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. --- .../src/components/Messages/Content/Error.tsx | 1 + .../SSE/__tests__/useResumableSSE.spec.ts | 273 ++++++++++++++++++ client/src/hooks/SSE/useResumableSSE.ts | 22 +- client/src/hooks/SSE/useSSE.ts | 22 +- client/src/locales/en/translation.json | 1 + client/src/utils/drafts.ts | 9 +- packages/data-provider/src/config.ts | 4 + 7 files changed, 299 insertions(+), 33 deletions(-) create mode 100644 client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts diff --git a/client/src/components/Messages/Content/Error.tsx b/client/src/components/Messages/Content/Error.tsx index ff2f2d7e90..b464ce2f2a 100644 --- a/client/src/components/Messages/Content/Error.tsx +++ b/client/src/components/Messages/Content/Error.tsx @@ -75,6 +75,7 @@ const errorMessages = { return info; }, [ErrorTypes.GOOGLE_TOOL_CONFLICT]: 'com_error_google_tool_conflict', + [ErrorTypes.STREAM_EXPIRED]: 'com_error_stream_expired', [ViolationTypes.BAN]: 'Your account has been temporarily banned due to violations of our service.', [ViolationTypes.ILLEGAL_MODEL_REQUEST]: (json: TGenericError, localize: LocalizeFunction) => { diff --git a/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts b/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts new file mode 100644 index 0000000000..9100f39858 --- /dev/null +++ b/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 & { responseCode?: number }) => void; + +interface MockSSEInstance { + addEventListener: jest.Mock; + stream: jest.Mock; + close: jest.Mock; + headers: Record; + _listeners: Record; + _emit: (event: string, data?: Partial & { responseCode?: number }) => void; +} + +const mockSSEInstances: MockSSEInstance[] = []; + +jest.mock('sse.js', () => ({ + SSE: jest.fn().mockImplementation(() => { + const listeners: Record = {}; + 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; + messages: never[]; + isTemporary: boolean; + initialResponse: Record; + endpointOption: { endpoint: string }; +}; + +const buildSubmission = (overrides: Partial = {}): 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(); + }, + ); +}); diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index 4d4cb4841a..ddfee30120 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/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[0]['data'], + submission: currentSubmission as EventSubmission, + }); setShowStopButton(false); setStreamId(null); reconnectAttemptRef.current = 0; diff --git a/client/src/hooks/SSE/useSSE.ts b/client/src/hooks/SSE/useSSE.ts index ccdb252287..78835f5729 100644 --- a/client/src/hooks/SSE/useSSE.ts +++ b/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) { diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index afd1072b61..9f641fdb16 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -376,6 +376,7 @@ "com_error_no_base_url": "No base URL found. Please provide one and try again.", "com_error_no_user_key": "No key found. Please provide a key and try again.", "com_error_refusal": "Response refused by safety filters. Rewrite your message and try again. If you encounter this frequently while using Claude Sonnet 4.5 or Opus 4.1, you can try Sonnet 4, which has different usage restrictions.", + "com_error_stream_expired": "The response stream has expired or already completed. Please try again.", "com_file_pages": "Pages: {{pages}}", "com_file_source": "File", "com_file_unknown": "Unknown File", diff --git a/client/src/utils/drafts.ts b/client/src/utils/drafts.ts index 1b3172def0..2e47c383b1 100644 --- a/client/src/utils/drafts.ts +++ b/client/src/utils/drafts.ts @@ -1,10 +1,17 @@ import debounce from 'lodash/debounce'; -import { LocalStorageKeys } from 'librechat-data-provider'; +import { Constants, LocalStorageKeys } from 'librechat-data-provider'; export const clearDraft = debounce((id?: string | null) => { localStorage.removeItem(`${LocalStorageKeys.TEXT_DRAFT}${id ?? ''}`); }, 2500); +/** Synchronously removes both text and file drafts for a conversation (or NEW_CONVO fallback) */ +export const clearAllDrafts = (conversationId?: string | null) => { + const key = conversationId || Constants.NEW_CONVO; + localStorage.removeItem(`${LocalStorageKeys.TEXT_DRAFT}${key}`); + localStorage.removeItem(`${LocalStorageKeys.FILES_DRAFT}${key}`); +}; + export const encodeBase64 = (plainText: string): string => { try { const textBytes = new TextEncoder().encode(plainText); diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index e19c69e799..0c8c591488 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -1616,6 +1616,10 @@ export enum ErrorTypes { * Model refused to respond (content policy violation) */ REFUSAL = 'refusal', + /** + * SSE stream 404 — job completed, expired, or was deleted before the subscriber connected + */ + STREAM_EXPIRED = 'stream_expired', } /**