mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-03 22:32:42 +00:00
⏱️ refactor: Optimistically Show New Chats In Sidebar (#13298)
* fix: optimistically show new chats in sidebar * fix: reconcile optimistic conversation ids * fix: clean optimistic conversation edge cases
This commit is contained in:
Danny Avila
GitHub
parent
6f7dc1b289
commit
05bb690fde
6 changed files with 539 additions and 20 deletions
|
|
@ -142,7 +142,7 @@ export default function Conversation({
|
|||
}, []);
|
||||
|
||||
const handleNavigation = (ctrlOrMetaKey: boolean) => {
|
||||
if (ctrlOrMetaKey) {
|
||||
if (ctrlOrMetaKey && !isGenerating) {
|
||||
toggleNav();
|
||||
const baseUrl = window.location.origin;
|
||||
const path = `/c/${conversationId}`;
|
||||
|
|
@ -278,7 +278,9 @@ export default function Conversation({
|
|||
// aria-hidden={!(isPopoverActive || isActiveConvo)}
|
||||
>
|
||||
{/* Only render ConvoOptions when user interacts (hover/focus) or for active conversation */}
|
||||
{!renaming && (hasInteracted || isActiveConvo) && <ConvoOptions {...convoOptionsProps} />}
|
||||
{!renaming && !isGenerating && (hasInteracted || isActiveConvo) && (
|
||||
<ConvoOptions {...convoOptionsProps} />
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import { renderHook, act } from '@testing-library/react';
|
||||
import { Constants, LocalStorageKeys } from 'librechat-data-provider';
|
||||
import { Constants, LocalStorageKeys, QueryKeys } from 'librechat-data-provider';
|
||||
import type { TSubmission } from 'librechat-data-provider';
|
||||
|
||||
type SSEEventListener = (e: Partial<MessageEvent> & { responseCode?: number }) => void;
|
||||
|
|
@ -34,12 +34,18 @@ jest.mock('sse.js', () => ({
|
|||
}));
|
||||
|
||||
const mockSetQueryData = jest.fn();
|
||||
const mockGetQueryData = jest.fn();
|
||||
const mockInvalidateQueries = jest.fn();
|
||||
const mockRemoveQueries = jest.fn();
|
||||
const mockFindAll = jest.fn((): Array<{ queryKey: unknown[] }> => []);
|
||||
const mockQueryClient = {
|
||||
setQueryData: mockSetQueryData,
|
||||
getQueryData: mockGetQueryData,
|
||||
invalidateQueries: mockInvalidateQueries,
|
||||
removeQueries: mockRemoveQueries,
|
||||
getQueryCache: () => ({
|
||||
findAll: mockFindAll,
|
||||
}),
|
||||
};
|
||||
|
||||
jest.mock('@tanstack/react-query', () => ({
|
||||
|
|
@ -73,6 +79,7 @@ jest.mock('~/data-provider', () => ({
|
|||
}));
|
||||
|
||||
const mockErrorHandler = jest.fn();
|
||||
const mockCreatedHandler = jest.fn();
|
||||
const mockSetIsSubmitting = jest.fn();
|
||||
const mockClearStepMaps = jest.fn();
|
||||
|
||||
|
|
@ -80,7 +87,7 @@ jest.mock('~/hooks/SSE/useEventHandlers', () =>
|
|||
jest.fn(() => ({
|
||||
errorHandler: mockErrorHandler,
|
||||
finalHandler: jest.fn(),
|
||||
createdHandler: jest.fn(),
|
||||
createdHandler: mockCreatedHandler,
|
||||
attachmentHandler: jest.fn(),
|
||||
stepHandler: jest.fn(),
|
||||
contentHandler: jest.fn(),
|
||||
|
|
@ -169,10 +176,14 @@ describe('useResumableSSE - 404 error path', () => {
|
|||
mockSSEInstances.length = 0;
|
||||
localStorage.clear();
|
||||
mockErrorHandler.mockClear();
|
||||
mockCreatedHandler.mockClear();
|
||||
mockClearStepMaps.mockClear();
|
||||
mockSetIsSubmitting.mockClear();
|
||||
mockSetQueryData.mockClear();
|
||||
mockGetQueryData.mockClear();
|
||||
mockInvalidateQueries.mockClear();
|
||||
mockRemoveQueries.mockClear();
|
||||
mockFindAll.mockClear();
|
||||
});
|
||||
|
||||
const seedDraft = (conversationId: string) => {
|
||||
|
|
@ -251,6 +262,68 @@ describe('useResumableSSE - 404 error path', () => {
|
|||
unmount();
|
||||
});
|
||||
|
||||
it('invalidates the stream conversation id on 404 for a new conversation', async () => {
|
||||
mockFindAll.mockReturnValue([{ queryKey: [QueryKeys.allConversations] }]);
|
||||
const submission = buildSubmission({
|
||||
conversation: {},
|
||||
userMessage: {
|
||||
messageId: 'msg-1',
|
||||
conversationId: null,
|
||||
text: 'Hello',
|
||||
isCreatedByUser: true,
|
||||
sender: 'User',
|
||||
parentMessageId: Constants.NO_PARENT,
|
||||
},
|
||||
initialResponse: {
|
||||
messageId: 'msg-1_',
|
||||
conversationId: null,
|
||||
text: '',
|
||||
isCreatedByUser: false,
|
||||
sender: 'Assistant',
|
||||
},
|
||||
});
|
||||
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(mockInvalidateQueries).toHaveBeenCalledWith({
|
||||
queryKey: [QueryKeys.messages, 'stream-123'],
|
||||
});
|
||||
expect(mockRemoveQueries).toHaveBeenCalledWith({
|
||||
queryKey: ['streamStatus', 'stream-123'],
|
||||
});
|
||||
|
||||
const allConversationWrites = mockSetQueryData.mock.calls.filter(
|
||||
([queryKey]) => Array.isArray(queryKey) && queryKey[0] === QueryKeys.allConversations,
|
||||
);
|
||||
expect(allConversationWrites).toHaveLength(2);
|
||||
|
||||
const removeUpdater = allConversationWrites[1][1] as (data: {
|
||||
pages: { conversations: { conversationId: string }[]; nextCursor: null }[];
|
||||
pageParams: never[];
|
||||
}) => { pages: { conversations: { conversationId: string }[] }[] };
|
||||
const result = removeUpdater({
|
||||
pages: [
|
||||
{
|
||||
conversations: [{ conversationId: 'stream-123' }, { conversationId: 'other' }],
|
||||
nextCursor: null,
|
||||
},
|
||||
],
|
||||
pageParams: [],
|
||||
});
|
||||
expect(result.pages[0].conversations).toEqual([{ conversationId: 'other' }]);
|
||||
unmount();
|
||||
});
|
||||
|
||||
it('closes the SSE connection on 404', async () => {
|
||||
const { sse, unmount } = await render404Scenario();
|
||||
|
||||
|
|
@ -258,6 +331,105 @@ describe('useResumableSSE - 404 error path', () => {
|
|||
unmount();
|
||||
});
|
||||
|
||||
it('seeds sidebar and message caches for a new conversation once the stream id is known', async () => {
|
||||
const submission = buildSubmission({
|
||||
conversation: {},
|
||||
userMessage: {
|
||||
messageId: 'msg-1',
|
||||
conversationId: null,
|
||||
text: 'Hello',
|
||||
isCreatedByUser: true,
|
||||
sender: 'User',
|
||||
parentMessageId: Constants.NO_PARENT,
|
||||
},
|
||||
initialResponse: {
|
||||
messageId: 'msg-1_',
|
||||
conversationId: null,
|
||||
text: '',
|
||||
isCreatedByUser: false,
|
||||
sender: 'Assistant',
|
||||
},
|
||||
});
|
||||
const chatHelpers = buildChatHelpers();
|
||||
|
||||
const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers));
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
expect(mockSetQueryData).toHaveBeenCalledWith(
|
||||
[QueryKeys.conversation, 'stream-123'],
|
||||
expect.any(Function),
|
||||
);
|
||||
expect(mockSetQueryData).toHaveBeenCalledWith(
|
||||
[QueryKeys.messages, 'stream-123'],
|
||||
expect.arrayContaining([
|
||||
expect.objectContaining({ messageId: 'msg-1', conversationId: 'stream-123' }),
|
||||
expect.objectContaining({ messageId: 'msg-1_', conversationId: 'stream-123' }),
|
||||
]),
|
||||
);
|
||||
expect(mockSetQueryData).toHaveBeenCalledWith(
|
||||
[QueryKeys.messages, Constants.NEW_CONVO],
|
||||
expect.arrayContaining([
|
||||
expect.objectContaining({ messageId: 'msg-1', conversationId: 'stream-123' }),
|
||||
]),
|
||||
);
|
||||
expect(mockFindAll).toHaveBeenCalledWith([QueryKeys.allConversations], { exact: false });
|
||||
|
||||
unmount();
|
||||
});
|
||||
|
||||
it('hydrates the submission conversation id before created handlers run', async () => {
|
||||
const submission = buildSubmission({
|
||||
conversation: {},
|
||||
userMessage: {
|
||||
messageId: 'msg-1',
|
||||
conversationId: null,
|
||||
text: 'Hello',
|
||||
isCreatedByUser: true,
|
||||
sender: 'User',
|
||||
parentMessageId: Constants.NO_PARENT,
|
||||
},
|
||||
initialResponse: {
|
||||
messageId: 'msg-1_',
|
||||
conversationId: null,
|
||||
text: '',
|
||||
isCreatedByUser: false,
|
||||
sender: 'Assistant',
|
||||
},
|
||||
});
|
||||
const chatHelpers = buildChatHelpers();
|
||||
|
||||
const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers));
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
const sse = getLastSSE();
|
||||
await act(async () => {
|
||||
sse._emit('message', {
|
||||
data: JSON.stringify({
|
||||
created: true,
|
||||
message: {
|
||||
messageId: 'msg-1',
|
||||
conversationId: 'stream-123',
|
||||
},
|
||||
}),
|
||||
});
|
||||
});
|
||||
|
||||
expect(mockCreatedHandler).toHaveBeenCalledWith(
|
||||
expect.any(Object),
|
||||
expect.objectContaining({
|
||||
conversation: expect.objectContaining({ conversationId: 'stream-123' }),
|
||||
userMessage: expect.objectContaining({ conversationId: 'stream-123' }),
|
||||
}),
|
||||
);
|
||||
unmount();
|
||||
});
|
||||
|
||||
it.each([undefined, 500, 503])(
|
||||
'does not call errorHandler for responseCode %s (reconnect path)',
|
||||
async (responseCode) => {
|
||||
|
|
@ -327,4 +499,66 @@ describe('useResumableSSE - 404 error path', () => {
|
|||
expect(mockErrorHandler).toHaveBeenCalledTimes(1);
|
||||
unmount();
|
||||
});
|
||||
|
||||
it('removes the optimistic sidebar row when a new conversation errors before created', async () => {
|
||||
mockFindAll.mockReturnValue([{ queryKey: [QueryKeys.allConversations] }]);
|
||||
const submission = buildSubmission({
|
||||
conversation: {},
|
||||
userMessage: {
|
||||
messageId: 'msg-1',
|
||||
conversationId: null,
|
||||
text: 'Hello',
|
||||
isCreatedByUser: true,
|
||||
sender: 'User',
|
||||
parentMessageId: Constants.NO_PARENT,
|
||||
},
|
||||
initialResponse: {
|
||||
messageId: 'msg-1_',
|
||||
conversationId: null,
|
||||
text: '',
|
||||
isCreatedByUser: false,
|
||||
sender: 'Assistant',
|
||||
},
|
||||
});
|
||||
const chatHelpers = buildChatHelpers();
|
||||
|
||||
const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers));
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
const sse = getLastSSE();
|
||||
await act(async () => {
|
||||
sse._emit('error', { data: JSON.stringify({ error: 'failed before created' }) });
|
||||
});
|
||||
|
||||
const allConversationWrites = mockSetQueryData.mock.calls.filter(
|
||||
([queryKey]) => Array.isArray(queryKey) && queryKey[0] === QueryKeys.allConversations,
|
||||
);
|
||||
expect(allConversationWrites).toHaveLength(2);
|
||||
|
||||
const removeUpdater = allConversationWrites[1][1] as (data: {
|
||||
pages: { conversations: { conversationId: string }[]; nextCursor: null }[];
|
||||
pageParams: never[];
|
||||
}) => { pages: { conversations: { conversationId: string }[] }[] };
|
||||
const result = removeUpdater({
|
||||
pages: [
|
||||
{
|
||||
conversations: [{ conversationId: 'stream-123' }, { conversationId: 'other' }],
|
||||
nextCursor: null,
|
||||
},
|
||||
],
|
||||
pageParams: [],
|
||||
});
|
||||
expect(result.pages[0].conversations).toEqual([{ conversationId: 'other' }]);
|
||||
expect(mockErrorHandler).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
submission: expect.objectContaining({
|
||||
conversation: expect.objectContaining({ conversationId: 'stream-123' }),
|
||||
}),
|
||||
}),
|
||||
);
|
||||
unmount();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ import {
|
|||
setDraft,
|
||||
scrollToEnd,
|
||||
getAllContentText,
|
||||
addConvoToAllQueries,
|
||||
upsertConvoInAllQueries,
|
||||
updateConvoInAllQueries,
|
||||
removeConvoFromAllQueries,
|
||||
findConversationInInfinite,
|
||||
|
|
@ -354,7 +354,7 @@ export default function useEventHandlers({
|
|||
});
|
||||
|
||||
if (requestMessage.parentMessageId === Constants.NO_PARENT) {
|
||||
addConvoToAllQueries(queryClient, update);
|
||||
upsertConvoInAllQueries(queryClient, update);
|
||||
} else {
|
||||
updateConvoInAllQueries(queryClient, update.conversationId!, (_c) => update, true);
|
||||
}
|
||||
|
|
@ -429,7 +429,7 @@ export default function useEventHandlers({
|
|||
|
||||
if (!isTemporary) {
|
||||
if (parentMessageId === Constants.NO_PARENT) {
|
||||
addConvoToAllQueries(queryClient, update);
|
||||
upsertConvoInAllQueries(queryClient, update);
|
||||
} else {
|
||||
updateConvoInAllQueries(queryClient, update.conversationId!, (_c) => update, true);
|
||||
}
|
||||
|
|
@ -476,6 +476,7 @@ export default function useEventHandlers({
|
|||
isRegenerate = false,
|
||||
isTemporary: _isTemporary = false,
|
||||
} = submission;
|
||||
const serverConversation = conversation as TConversation;
|
||||
|
||||
try {
|
||||
// Handle early abort - aborted during tool loading before any messages saved
|
||||
|
|
@ -608,14 +609,14 @@ export default function useEventHandlers({
|
|||
if (prevState?.model != null && prevState.model !== submissionConvo.model) {
|
||||
update.model = prevState.model;
|
||||
}
|
||||
const cachedConvo = queryClient.getQueryData<TConversation>([
|
||||
QueryKeys.conversation,
|
||||
conversation.conversationId,
|
||||
]);
|
||||
if (!cachedConvo) {
|
||||
queryClient.setQueryData(
|
||||
if (conversation.conversationId) {
|
||||
queryClient.setQueryData<TConversation>(
|
||||
[QueryKeys.conversation, conversation.conversationId],
|
||||
update,
|
||||
(cachedConvo) =>
|
||||
({
|
||||
...cachedConvo,
|
||||
...serverConversation,
|
||||
}) as TConversation,
|
||||
);
|
||||
}
|
||||
return update;
|
||||
|
|
|
|||
|
|
@ -14,7 +14,13 @@ import {
|
|||
ViolationTypes,
|
||||
removeNullishValues,
|
||||
} from 'librechat-data-provider';
|
||||
import type { TMessage, TPayload, TSubmission, EventSubmission } from 'librechat-data-provider';
|
||||
import type {
|
||||
TMessage,
|
||||
TPayload,
|
||||
TSubmission,
|
||||
TConversation,
|
||||
EventSubmission,
|
||||
} from 'librechat-data-provider';
|
||||
import type { EventHandlerParams } from './useEventHandlers';
|
||||
import {
|
||||
useGetUserBalance,
|
||||
|
|
@ -25,7 +31,7 @@ import {
|
|||
import type { ActiveJobsResponse } from '~/data-provider';
|
||||
import { useAuthContext } from '~/hooks/AuthContext';
|
||||
import useEventHandlers from './useEventHandlers';
|
||||
import { clearAllDrafts } from '~/utils';
|
||||
import { clearAllDrafts, removeConvoFromAllQueries, upsertConvoInAllQueries } from '~/utils';
|
||||
import store from '~/store';
|
||||
|
||||
type ChatHelpers = Pick<
|
||||
|
|
@ -35,6 +41,76 @@ type ChatHelpers = Pick<
|
|||
|
||||
const MAX_RETRIES = 5;
|
||||
|
||||
const hasConcreteConversationId = (conversationId?: string | null) =>
|
||||
!!conversationId &&
|
||||
conversationId !== Constants.NEW_CONVO &&
|
||||
conversationId !== Constants.PENDING_CONVO;
|
||||
|
||||
const isInitialNewConversation = (submission: TSubmission) => {
|
||||
const conversationId = submission.conversation?.conversationId;
|
||||
return (
|
||||
submission.userMessage?.parentMessageId === Constants.NO_PARENT &&
|
||||
!hasConcreteConversationId(conversationId)
|
||||
);
|
||||
};
|
||||
|
||||
const shouldHydrateMessage = (message: TMessage) =>
|
||||
!hasConcreteConversationId(message.conversationId);
|
||||
|
||||
const hydrateMessageConversationId = (message: TMessage, conversationId: string): TMessage =>
|
||||
shouldHydrateMessage(message) ? { ...message, conversationId } : message;
|
||||
|
||||
const getOptimisticMessages = (
|
||||
submission: TSubmission,
|
||||
conversationId: string,
|
||||
messages?: TMessage[],
|
||||
): TMessage[] => {
|
||||
const sourceMessages =
|
||||
messages && messages.length > 0
|
||||
? messages
|
||||
: [submission.userMessage, submission.initialResponse].filter(
|
||||
(message): message is TMessage => message != null,
|
||||
);
|
||||
|
||||
return sourceMessages.map((message) => hydrateMessageConversationId(message, conversationId));
|
||||
};
|
||||
|
||||
const buildOptimisticConversation = (
|
||||
submission: TSubmission,
|
||||
conversationId: string,
|
||||
): TConversation => {
|
||||
const now = new Date().toISOString();
|
||||
const messageIds = [
|
||||
submission.userMessage?.messageId,
|
||||
submission.initialResponse?.messageId,
|
||||
].filter((messageId): messageId is string => typeof messageId === 'string' && messageId !== '');
|
||||
|
||||
return {
|
||||
...submission.conversation,
|
||||
conversationId,
|
||||
endpoint: submission.conversation.endpoint ?? null,
|
||||
title: submission.conversation.title ?? 'New Chat',
|
||||
messages: messageIds.length > 0 ? messageIds : submission.conversation.messages,
|
||||
createdAt: submission.conversation.createdAt ?? now,
|
||||
updatedAt: now,
|
||||
} as TConversation;
|
||||
};
|
||||
|
||||
const hydrateSubmissionMessages = (
|
||||
submission: TSubmission,
|
||||
conversationId: string,
|
||||
): TSubmission => ({
|
||||
...submission,
|
||||
conversation: {
|
||||
...submission.conversation,
|
||||
conversationId,
|
||||
},
|
||||
userMessage: hydrateMessageConversationId(submission.userMessage, conversationId),
|
||||
initialResponse: submission.initialResponse
|
||||
? hydrateMessageConversationId(submission.initialResponse, conversationId)
|
||||
: submission.initialResponse,
|
||||
});
|
||||
|
||||
/**
|
||||
* Hook for resumable SSE streams.
|
||||
* Separates generation start (POST) from stream subscription (GET EventSource).
|
||||
|
|
@ -55,6 +131,8 @@ export default function useResumableSSE(
|
|||
const setActiveRunId = useSetRecoilState(store.activeRunFamily(runIndex));
|
||||
|
||||
const { token, isAuthenticated } = useAuthContext();
|
||||
const { setMessages, getMessages, setConversation, setIsSubmitting, newConversation } =
|
||||
chatHelpers;
|
||||
|
||||
/**
|
||||
* Optimistically add a job ID to the active jobs cache.
|
||||
|
|
@ -81,6 +159,38 @@ export default function useResumableSSE(
|
|||
},
|
||||
[queryClient],
|
||||
);
|
||||
|
||||
const addOptimisticConversation = useCallback(
|
||||
(conversationId: string, currentSubmission: TSubmission): TSubmission => {
|
||||
if (!isInitialNewConversation(currentSubmission)) {
|
||||
return currentSubmission;
|
||||
}
|
||||
|
||||
const optimisticConversation = buildOptimisticConversation(currentSubmission, conversationId);
|
||||
const optimisticMessages = getOptimisticMessages(
|
||||
currentSubmission,
|
||||
conversationId,
|
||||
getMessages(),
|
||||
);
|
||||
|
||||
queryClient.setQueryData<TConversation>(
|
||||
[QueryKeys.conversation, conversationId],
|
||||
(current) => current ?? optimisticConversation,
|
||||
);
|
||||
queryClient.setQueryData<TMessage[]>(
|
||||
[QueryKeys.messages, conversationId],
|
||||
optimisticMessages,
|
||||
);
|
||||
queryClient.setQueryData<TMessage[]>(
|
||||
[QueryKeys.messages, Constants.NEW_CONVO],
|
||||
optimisticMessages,
|
||||
);
|
||||
upsertConvoInAllQueries(queryClient, optimisticConversation);
|
||||
|
||||
return hydrateSubmissionMessages(currentSubmission, conversationId);
|
||||
},
|
||||
[getMessages, queryClient],
|
||||
);
|
||||
const [_completed, setCompleted] = useState(new Set());
|
||||
const [streamId, setStreamId] = useState<string | null>(null);
|
||||
const setAbortScroll = useSetRecoilState(store.abortScrollFamily(runIndex));
|
||||
|
|
@ -90,9 +200,8 @@ export default function useResumableSSE(
|
|||
const reconnectAttemptRef = useRef(0);
|
||||
const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);
|
||||
const submissionRef = useRef<TSubmission | null>(null);
|
||||
|
||||
const { setMessages, getMessages, setConversation, setIsSubmitting, newConversation } =
|
||||
chatHelpers;
|
||||
const optimisticStreamIdsRef = useRef(new Set<string>());
|
||||
const createdStreamIdsRef = useRef(new Set<string>());
|
||||
|
||||
const {
|
||||
stepHandler,
|
||||
|
|
@ -161,6 +270,9 @@ export default function useResumableSSE(
|
|||
hasResponseMessage: !!data.responseMessage,
|
||||
});
|
||||
clearAllDrafts(currentSubmission.conversation?.conversationId);
|
||||
if (optimisticStreamIdsRef.current.has(currentStreamId)) {
|
||||
clearAllDrafts(Constants.NEW_CONVO);
|
||||
}
|
||||
try {
|
||||
finalHandler(data, currentSubmission as EventSubmission);
|
||||
} catch (error) {
|
||||
|
|
@ -175,6 +287,8 @@ export default function useResumableSSE(
|
|||
(startupConfig?.balance?.enabled ?? false) && balanceQuery.refetch();
|
||||
sse.close();
|
||||
setStreamId(null);
|
||||
optimisticStreamIdsRef.current.delete(currentStreamId);
|
||||
createdStreamIdsRef.current.delete(currentStreamId);
|
||||
return;
|
||||
}
|
||||
|
||||
|
|
@ -183,6 +297,7 @@ export default function useResumableSSE(
|
|||
messageId: data.message?.messageId,
|
||||
conversationId: data.message?.conversationId,
|
||||
});
|
||||
createdStreamIdsRef.current.add(currentStreamId);
|
||||
const runId = v4();
|
||||
setActiveRunId(runId);
|
||||
userMessage = {
|
||||
|
|
@ -344,14 +459,25 @@ export default function useResumableSSE(
|
|||
sse.close();
|
||||
removeActiveJob(currentStreamId);
|
||||
clearAllDrafts(convoId);
|
||||
if (optimisticStreamIdsRef.current.has(currentStreamId)) {
|
||||
clearAllDrafts(Constants.NEW_CONVO);
|
||||
}
|
||||
clearStepMaps();
|
||||
if (convoId) {
|
||||
queryClient.invalidateQueries({ queryKey: [QueryKeys.messages, convoId] });
|
||||
queryClient.removeQueries({ queryKey: streamStatusQueryKey(convoId) });
|
||||
}
|
||||
if (
|
||||
!createdStreamIdsRef.current.has(currentStreamId) &&
|
||||
optimisticStreamIdsRef.current.has(currentStreamId)
|
||||
) {
|
||||
removeConvoFromAllQueries(queryClient, currentStreamId);
|
||||
}
|
||||
setIsSubmitting(false);
|
||||
setShowStopButton(false);
|
||||
setStreamId(null);
|
||||
optimisticStreamIdsRef.current.delete(currentStreamId);
|
||||
createdStreamIdsRef.current.delete(currentStreamId);
|
||||
reconnectAttemptRef.current = 0;
|
||||
return;
|
||||
}
|
||||
|
|
@ -386,6 +512,12 @@ export default function useResumableSSE(
|
|||
console.log('[ResumableSSE] Server-sent error event received:', e.data);
|
||||
sse.close();
|
||||
removeActiveJob(currentStreamId);
|
||||
if (
|
||||
!createdStreamIdsRef.current.has(currentStreamId) &&
|
||||
optimisticStreamIdsRef.current.has(currentStreamId)
|
||||
) {
|
||||
removeConvoFromAllQueries(queryClient, currentStreamId);
|
||||
}
|
||||
|
||||
try {
|
||||
const errorData = JSON.parse(e.data);
|
||||
|
|
@ -425,6 +557,8 @@ export default function useResumableSSE(
|
|||
setIsSubmitting(false);
|
||||
setShowStopButton(false);
|
||||
setStreamId(null);
|
||||
optimisticStreamIdsRef.current.delete(currentStreamId);
|
||||
createdStreamIdsRef.current.delete(currentStreamId);
|
||||
reconnectAttemptRef.current = 0;
|
||||
return;
|
||||
}
|
||||
|
|
@ -463,9 +597,17 @@ export default function useResumableSSE(
|
|||
errorHandler({ data: undefined, submission: currentSubmission as EventSubmission });
|
||||
// Optimistically remove from active jobs on max retries
|
||||
removeActiveJob(currentStreamId);
|
||||
if (
|
||||
!createdStreamIdsRef.current.has(currentStreamId) &&
|
||||
optimisticStreamIdsRef.current.has(currentStreamId)
|
||||
) {
|
||||
removeConvoFromAllQueries(queryClient, currentStreamId);
|
||||
}
|
||||
setIsSubmitting(false);
|
||||
setShowStopButton(false);
|
||||
setStreamId(null);
|
||||
optimisticStreamIdsRef.current.delete(currentStreamId);
|
||||
createdStreamIdsRef.current.delete(currentStreamId);
|
||||
}
|
||||
});
|
||||
|
||||
|
|
@ -664,7 +806,12 @@ export default function useResumableSSE(
|
|||
if (isNewConvo) {
|
||||
queueTitleGeneration(newStreamId);
|
||||
}
|
||||
subscribeToStream(newStreamId, submission);
|
||||
if (isInitialNewConversation(submission)) {
|
||||
optimisticStreamIdsRef.current.add(newStreamId);
|
||||
}
|
||||
const streamSubmission = addOptimisticConversation(newStreamId, submission);
|
||||
submissionRef.current = streamSubmission;
|
||||
subscribeToStream(newStreamId, streamSubmission);
|
||||
} else {
|
||||
console.error('[ResumableSSE] Failed to get streamId from startGeneration');
|
||||
}
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ import {
|
|||
groupConversationsByDate,
|
||||
updateConvoFieldsInfinite,
|
||||
addConvoToAllQueries,
|
||||
upsertConvoInAllQueries,
|
||||
updateConvoInAllQueries,
|
||||
removeConvoFromAllQueries,
|
||||
addConversationToAllConversationsQueries,
|
||||
|
|
@ -590,6 +591,47 @@ describe('Conversation Utilities', () => {
|
|||
expect(data!.pages[0].conversations.filter((c) => c.conversationId === 'a').length).toBe(1);
|
||||
});
|
||||
|
||||
it('upsertConvoInAllQueries adds missing conversations to the top', () => {
|
||||
upsertConvoInAllQueries(queryClient, convoB);
|
||||
const data = queryClient.getQueryData<InfiniteData<{ conversations: TConversation[] }>>([
|
||||
'allConversations',
|
||||
]);
|
||||
|
||||
expect(data!.pages[0].conversations[0].conversationId).toBe('b');
|
||||
expect(data!.pages[0].conversations[1].conversationId).toBe('a');
|
||||
});
|
||||
|
||||
it('upsertConvoInAllQueries updates existing conversations without duplicating them', () => {
|
||||
upsertConvoInAllQueries(queryClient, {
|
||||
...convoA,
|
||||
title: 'Updated Conversation A',
|
||||
});
|
||||
const data = queryClient.getQueryData<InfiniteData<{ conversations: TConversation[] }>>([
|
||||
'allConversations',
|
||||
]);
|
||||
|
||||
expect(data!.pages[0].conversations).toHaveLength(1);
|
||||
expect(data!.pages[0].conversations[0].title).toBe('Updated Conversation A');
|
||||
});
|
||||
|
||||
it('upsertConvoInAllQueries moves an existing conversation to the top once', () => {
|
||||
const convoC = { conversationId: 'c', updatedAt: '2024-01-03T12:00:00Z' } as TConversation;
|
||||
queryClient.setQueryData(['allConversations'], {
|
||||
pages: [{ conversations: [convoC, convoA], nextCursor: null }],
|
||||
pageParams: [],
|
||||
});
|
||||
|
||||
upsertConvoInAllQueries(queryClient, {
|
||||
...convoA,
|
||||
title: 'Updated Conversation A',
|
||||
});
|
||||
const data = queryClient.getQueryData<InfiniteData<{ conversations: TConversation[] }>>([
|
||||
'allConversations',
|
||||
]);
|
||||
|
||||
expect(data!.pages[0].conversations.map((c) => c.conversationId)).toEqual(['a', 'c']);
|
||||
});
|
||||
|
||||
it('updateConvoInAllQueries updates correct convo', () => {
|
||||
updateConvoInAllQueries(queryClient, 'a', (c) => ({ ...c, model: 'gpt-4' }));
|
||||
const data = queryClient.getQueryData<InfiniteData<any>>(['allConversations']);
|
||||
|
|
|
|||
|
|
@ -347,6 +347,99 @@ export function addConvoToAllQueries(queryClient: QueryClient, newConvo: TConver
|
|||
}
|
||||
}
|
||||
|
||||
export function upsertConvoInAllQueries(
|
||||
queryClient: QueryClient,
|
||||
nextConvo: TConversation,
|
||||
moveToTop = true,
|
||||
) {
|
||||
if (!nextConvo.conversationId) {
|
||||
return;
|
||||
}
|
||||
|
||||
const queries = queryClient
|
||||
.getQueryCache()
|
||||
.findAll([QueryKeys.allConversations], { exact: false });
|
||||
|
||||
for (const query of queries) {
|
||||
queryClient.setQueryData<InfiniteData<ConversationCursorData>>(query.queryKey, (oldData) => {
|
||||
if (!oldData) {
|
||||
return oldData;
|
||||
}
|
||||
|
||||
let pageIdx = -1;
|
||||
let convoIdx = -1;
|
||||
for (let pi = 0; pi < oldData.pages.length; pi++) {
|
||||
const ci = oldData.pages[pi].conversations.findIndex(
|
||||
(c) => c.conversationId === nextConvo.conversationId,
|
||||
);
|
||||
if (ci !== -1) {
|
||||
pageIdx = pi;
|
||||
convoIdx = ci;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
const now = new Date().toISOString();
|
||||
if (pageIdx === -1) {
|
||||
const firstPage = oldData.pages[0] ?? { conversations: [], nextCursor: null };
|
||||
return {
|
||||
...oldData,
|
||||
pages: [
|
||||
{
|
||||
...firstPage,
|
||||
conversations: [
|
||||
{ ...nextConvo, updatedAt: nextConvo.updatedAt ?? now },
|
||||
...firstPage.conversations,
|
||||
],
|
||||
},
|
||||
...oldData.pages.slice(1),
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
const found = oldData.pages[pageIdx].conversations[convoIdx];
|
||||
const updated = {
|
||||
...found,
|
||||
...nextConvo,
|
||||
updatedAt: nextConvo.updatedAt ?? (moveToTop ? now : found.updatedAt),
|
||||
};
|
||||
|
||||
if (!moveToTop || (pageIdx === 0 && convoIdx === 0)) {
|
||||
return {
|
||||
...oldData,
|
||||
pages: oldData.pages.map((page, pi) =>
|
||||
pi === pageIdx
|
||||
? {
|
||||
...page,
|
||||
conversations: page.conversations.map((c, ci) => (ci === convoIdx ? updated : c)),
|
||||
}
|
||||
: page,
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
const pages = oldData.pages.map((page, pi) => {
|
||||
if (pi === 0 && pageIdx === 0) {
|
||||
const conversations = page.conversations.filter((_, ci) => ci !== convoIdx);
|
||||
return { ...page, conversations: [updated, ...conversations] };
|
||||
}
|
||||
if (pi === 0) {
|
||||
return { ...page, conversations: [updated, ...page.conversations] };
|
||||
}
|
||||
if (pi === pageIdx) {
|
||||
return {
|
||||
...page,
|
||||
conversations: page.conversations.filter((_, ci) => ci !== convoIdx),
|
||||
};
|
||||
}
|
||||
return page;
|
||||
});
|
||||
|
||||
return { ...oldData, pages };
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Update
|
||||
export function updateConvoInAllQueries(
|
||||
queryClient: QueryClient,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue