From 1e0ffcf2fd1f8d562cbfa37d4085c3e39415ca21 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sat, 23 May 2026 11:53:25 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=93=A1=20fix:=20Tighten=20Streaming=20Mes?= =?UTF-8?q?sage=20Cache=20Preservation=20(#13271)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: tighten streaming message cache preservation * fix: sync streaming ref after commit --- client/src/components/Chat/ChatView.tsx | 25 +++-- .../components/Chat/Messages/MessageNav.tsx | 13 ++- .../Messages/__tests__/MessageNav.spec.tsx | 3 + .../Messages/__tests__/queries.test.ts | 103 +++++++++++++++++- client/src/data-provider/Messages/queries.ts | 57 ++++++++-- 5 files changed, 177 insertions(+), 24 deletions(-) diff --git a/client/src/components/Chat/ChatView.tsx b/client/src/components/Chat/ChatView.tsx index d2bd7edf0d..27530a5c22 100644 --- a/client/src/components/Chat/ChatView.tsx +++ b/client/src/components/Chat/ChatView.tsx @@ -32,6 +32,7 @@ function LoadingSpinner() { function ChatView({ index = 0 }: { index?: number }) { const { conversationId } = useParams(); const rootSubmission = useRecoilValue(store.submissionByIndex(index)); + const isSubmitting = useRecoilValue(store.isSubmittingFamily(index)); const centerFormOnLanding = useRecoilValue(store.centerFormOnLanding); const methods = useForm({ @@ -40,16 +41,20 @@ function ChatView({ index = 0 }: { index?: number }) { const fileMap = useFileMapContext(); - const { data: messagesTree = null, isLoading } = useGetMessagesByConvoId(conversationId ?? '', { - select: useCallback( - (data: TMessage[]) => { - const dataTree = buildTree({ messages: data, fileMap }); - return dataTree?.length === 0 ? null : (dataTree ?? null); - }, - [fileMap], - ), - enabled: !!fileMap, - }); + const { data: messagesTree = null, isLoading } = useGetMessagesByConvoId( + conversationId ?? '', + { + select: useCallback( + (data: TMessage[]) => { + const dataTree = buildTree({ messages: data, fileMap }); + return dataTree?.length === 0 ? null : (dataTree ?? null); + }, + [fileMap], + ), + enabled: !!fileMap, + }, + { isStreaming: isSubmitting }, + ); const chatHelpers = useChatHelpers(index, conversationId); const addedChatHelpers = useAddedResponse(); diff --git a/client/src/components/Chat/Messages/MessageNav.tsx b/client/src/components/Chat/Messages/MessageNav.tsx index 9ef79197ce..0c2d3b0c62 100644 --- a/client/src/components/Chat/Messages/MessageNav.tsx +++ b/client/src/components/Chat/Messages/MessageNav.tsx @@ -4,7 +4,7 @@ import { ContentTypes } from 'librechat-data-provider'; import { HoverCard, HoverCardTrigger, HoverCardPortal, HoverCardContent } from '@librechat/client'; import type { TMessage, TMessageContentParts } from 'librechat-data-provider'; import { useGetMessagesByConvoId } from '~/data-provider'; -import { useMessagesConversation } from '~/Providers'; +import { useMessagesConversation, useMessagesSubmission } from '~/Providers'; import { useLocalize } from '~/hooks'; import { cn } from '~/utils'; @@ -145,9 +145,14 @@ const chevronButtonClasses = cn( function MessageNav({ scrollableRef }: { scrollableRef: React.RefObject }) { const localize = useLocalize(); const { conversationId } = useMessagesConversation(); - const { data: messages } = useGetMessagesByConvoId(conversationId ?? '', { - enabled: !!conversationId, - }); + const { isSubmitting } = useMessagesSubmission(); + const { data: messages } = useGetMessagesByConvoId( + conversationId ?? '', + { + enabled: !!conversationId, + }, + { isStreaming: isSubmitting }, + ); const messagesById = useMemo(() => { const map = new Map(); if (messages) { diff --git a/client/src/components/Chat/Messages/__tests__/MessageNav.spec.tsx b/client/src/components/Chat/Messages/__tests__/MessageNav.spec.tsx index d9c02baa43..00708e76b7 100644 --- a/client/src/components/Chat/Messages/__tests__/MessageNav.spec.tsx +++ b/client/src/components/Chat/Messages/__tests__/MessageNav.spec.tsx @@ -13,6 +13,7 @@ type TestMessage = { const mockUseGetMessagesByConvoId = jest.fn(); const mockUseMessagesConversation = jest.fn(); +const mockUseMessagesSubmission = jest.fn(); jest.mock('~/data-provider', () => ({ useGetMessagesByConvoId: (...args: unknown[]) => mockUseGetMessagesByConvoId(...args), @@ -20,6 +21,7 @@ jest.mock('~/data-provider', () => ({ jest.mock('~/Providers', () => ({ useMessagesConversation: (...args: unknown[]) => mockUseMessagesConversation(...args), + useMessagesSubmission: (...args: unknown[]) => mockUseMessagesSubmission(...args), })); jest.mock('~/hooks', () => ({ @@ -158,6 +160,7 @@ beforeEach(() => { ).IntersectionObserver = MockIntersectionObserver; jest.useFakeTimers(); mockUseMessagesConversation.mockReturnValue({ conversationId: 'test-convo' }); + mockUseMessagesSubmission.mockReturnValue({ isSubmitting: false }); mockUseGetMessagesByConvoId.mockReturnValue({ data: [] }); }); diff --git a/client/src/data-provider/Messages/__tests__/queries.test.ts b/client/src/data-provider/Messages/__tests__/queries.test.ts index c8fd0e99a1..34b4746020 100644 --- a/client/src/data-provider/Messages/__tests__/queries.test.ts +++ b/client/src/data-provider/Messages/__tests__/queries.test.ts @@ -21,6 +21,7 @@ describe('getStableMessages', () => { message({ messageId: 'user-2', createdAt: undefined, updatedAt: undefined }), message({ messageId: 'user-2_', + parentMessageId: 'user-2', isCreatedByUser: false, createdAt: undefined, updatedAt: undefined, @@ -31,26 +32,122 @@ describe('getStableMessages', () => { pathname: '/c/convo-id', result: [], currentMessages, + isStreaming: true, }); expect(result).toBe(currentMessages); }); - it('keeps cache when a one-message result returns for a larger cache', () => { + it('keeps cache when a prefix result races with a pending assistant tail', () => { const currentMessages = [ message({ messageId: 'persisted-1' }), - message({ messageId: 'persisted-2' }), + message({ messageId: 'user-2', createdAt: undefined, updatedAt: undefined }), + message({ + messageId: 'user-2_', + parentMessageId: 'user-2', + isCreatedByUser: false, + createdAt: undefined, + updatedAt: undefined, + }), ]; const result = getStableMessages({ pathname: '/c/convo-id', result: [currentMessages[0]], currentMessages, + isStreaming: true, }); expect(result).toBe(currentMessages); }); + it('accepts a shorter result when the app is not streaming', () => { + const currentMessages = [ + message({ messageId: 'persisted-1' }), + message({ + messageId: 'user-2_', + parentMessageId: 'user-2', + isCreatedByUser: false, + createdAt: undefined, + updatedAt: undefined, + }), + ]; + const serverMessages = [currentMessages[0]]; + + const result = getStableMessages({ + pathname: '/c/convo-id', + result: serverMessages, + currentMessages, + isStreaming: false, + }); + + expect(result).toBe(serverMessages); + }); + + it('accepts a shorter result when an unhydrated message is not the assistant tail', () => { + const currentMessages = [ + message({ messageId: 'persisted-1' }), + message({ messageId: 'old-unhydrated', createdAt: undefined, updatedAt: undefined }), + message({ messageId: 'persisted-3' }), + ]; + const serverMessages = [currentMessages[0]]; + + const result = getStableMessages({ + pathname: '/c/convo-id', + result: serverMessages, + currentMessages, + isStreaming: true, + }); + + expect(result).toBe(serverMessages); + }); + + it('accepts a shorter result when the server result is not a prefix of cache', () => { + const currentMessages = [ + message({ messageId: 'persisted-1' }), + message({ messageId: 'persisted-2' }), + message({ + messageId: 'persisted-2_', + parentMessageId: 'persisted-2', + isCreatedByUser: false, + createdAt: undefined, + updatedAt: undefined, + }), + ]; + const serverMessages = [message({ messageId: 'different-1' })]; + + const result = getStableMessages({ + pathname: '/c/convo-id', + result: serverMessages, + currentMessages, + isStreaming: true, + }); + + expect(result).toBe(serverMessages); + }); + + it('accepts a shorter result when an unhydrated assistant tail has no parent turn', () => { + const currentMessages = [ + message({ messageId: 'persisted-1' }), + message({ + messageId: 'orphaned-response_', + isCreatedByUser: false, + createdAt: undefined, + updatedAt: undefined, + }), + ]; + const serverMessages = [currentMessages[0]]; + + const result = getStableMessages({ + pathname: '/c/convo-id', + result: serverMessages, + currentMessages, + isStreaming: true, + }); + + expect(result).toBe(serverMessages); + }); + it('accepts fewer persisted messages when the cache is fully hydrated', () => { const currentMessages = [ message({ messageId: 'persisted-1' }), @@ -63,6 +160,7 @@ describe('getStableMessages', () => { pathname: '/c/convo-id', result: serverMessages, currentMessages, + isStreaming: true, }); expect(result).toBe(serverMessages); @@ -75,6 +173,7 @@ describe('getStableMessages', () => { pathname: '/c/new', result: [], currentMessages, + isStreaming: true, }); expect(result).toEqual([]); diff --git a/client/src/data-provider/Messages/queries.ts b/client/src/data-provider/Messages/queries.ts index 0799b4adac..7bd4f3617b 100644 --- a/client/src/data-provider/Messages/queries.ts +++ b/client/src/data-provider/Messages/queries.ts @@ -1,26 +1,46 @@ +import { useLayoutEffect, useRef } from 'react'; import { useLocation } from 'react-router-dom'; import { useQuery, useQueryClient } from '@tanstack/react-query'; -import type { UseQueryOptions, QueryObserverResult } from '@tanstack/react-query'; -import { QueryKeys, dataService } from 'librechat-data-provider'; +import type { UseQueryOptions, QueryObserverResult, QueryClient } from '@tanstack/react-query'; +import { Constants, QueryKeys, dataService } from 'librechat-data-provider'; import type * as t from 'librechat-data-provider'; import { logger } from '~/utils'; type StableMessagesParams = { pathname: string; result: t.TMessage[]; + isStreaming?: boolean; currentMessages?: t.TMessage[]; }; -function hasUnhydratedMessage(messages: t.TMessage[]) { - return messages.some((message) => { - const messageId = message.messageId ?? ''; - return message.createdAt == null || message.updatedAt == null || messageId.endsWith('_'); - }); +type ActiveJobs = { + activeJobIds?: string[]; +}; + +function isUnhydratedMessage(message: t.TMessage) { + const messageId = message.messageId ?? ''; + return message.createdAt == null || message.updatedAt == null || messageId.endsWith('_'); +} + +function hasPendingAssistantTail(messages: t.TMessage[]) { + const lastMessage = messages[messages.length - 1]; + const parentMessageId = lastMessage?.parentMessageId ?? ''; + return ( + lastMessage?.isCreatedByUser !== true && + parentMessageId !== '' && + parentMessageId !== Constants.NO_PARENT && + isUnhydratedMessage(lastMessage) + ); +} + +function isMessagePrefix(result: t.TMessage[], currentMessages: t.TMessage[]) { + return result.every((message, index) => message.messageId === currentMessages[index]?.messageId); } export function getStableMessages({ pathname, result, + isStreaming = false, currentMessages, }: StableMessagesParams): t.TMessage[] { if (pathname.includes('/c/new') || !currentMessages?.length) { @@ -31,19 +51,39 @@ export function getStableMessages({ return result; } - if (result.length === 1 || hasUnhydratedMessage(currentMessages)) { + if ( + isStreaming && + hasPendingAssistantTail(currentMessages) && + isMessagePrefix(result, currentMessages) + ) { return currentMessages; } return result; } +function hasActiveJob(queryClient: QueryClient, id: string) { + if (!id) { + return false; + } + const activeJobs = queryClient.getQueryData([QueryKeys.activeJobs]); + return activeJobs?.activeJobIds?.includes(id) === true; +} + export const useGetMessagesByConvoId = ( id: string, config?: UseQueryOptions, + options?: { isStreaming?: boolean }, ): QueryObserverResult => { const location = useLocation(); const queryClient = useQueryClient(); + const isStreaming = options?.isStreaming === true; + const isStreamingRef = useRef(isStreaming); + + useLayoutEffect(() => { + isStreamingRef.current = isStreaming; + }, [isStreaming]); + return useQuery( [QueryKeys.messages, id], async () => { @@ -53,6 +93,7 @@ export const useGetMessagesByConvoId = ( pathname: location.pathname, result, currentMessages, + isStreaming: isStreamingRef.current || hasActiveJob(queryClient, id), }); if (stableMessages === currentMessages) {