diff --git a/client/src/data-provider/SSE/__tests__/queries.test.ts b/client/src/data-provider/SSE/__tests__/queries.test.ts index dce1e86b57..1c320d3c98 100644 --- a/client/src/data-provider/SSE/__tests__/queries.test.ts +++ b/client/src/data-provider/SSE/__tests__/queries.test.ts @@ -35,8 +35,14 @@ jest.mock('../../Endpoints', () => ({ useGetStartupConfig: jest.fn(() => ({ data: undefined })), })); +import { useQuery } from '@tanstack/react-query'; import { request } from 'librechat-data-provider'; -import { fetchStreamStatus, genTitleQueryKey, queueTitleGeneration } from '../queries'; +import { + fetchStreamStatus, + genTitleQueryKey, + queueTitleGeneration, + useActiveJobs, +} from '../queries'; describe('fetchStreamStatus generation protocol advertisement', () => { it('sends v2 in both the query and header', async () => { @@ -61,6 +67,34 @@ function makeAxiosError(status: number): Error { return err; } +describe('useActiveJobs focus behaviour', () => { + it('refetches on focus unconditionally rather than only when stale', () => { + (useQuery as jest.Mock).mockClear(); + + useActiveJobs(); + + const options = (useQuery as jest.Mock).mock.calls[0][0]; + /** + * A run another client starts inside the `staleTime` window is invisible to + * a plain `true`, which only refetches stale queries — and the interval is + * off while nothing is listed, so nothing else would come back for it. + * Returning to a tab is exactly when a pane needs to know its history moved. + */ + expect(options.refetchOnWindowFocus).toBe('always'); + expect(options.staleTime).toBe(5_000); + }); + + it('polls only while something is listed', () => { + (useQuery as jest.Mock).mockClear(); + + useActiveJobs(); + + const options = (useQuery as jest.Mock).mock.calls[0][0]; + expect(options.refetchInterval({ activeJobIds: ['convo-1'] })).toBe(5_000); + expect(options.refetchInterval({ activeJobIds: [] })).toBe(false); + }); +}); + describe('genTitleQueryKey', () => { it('returns a two-element tuple with the conversationId', () => { expect(genTitleQueryKey('abc-123')).toEqual(['genTitle', 'abc-123']); diff --git a/client/src/data-provider/SSE/queries.ts b/client/src/data-provider/SSE/queries.ts index 16089bf72c..d7d0813e75 100644 --- a/client/src/data-provider/SSE/queries.ts +++ b/client/src/data-provider/SSE/queries.ts @@ -223,7 +223,11 @@ export function useActiveJobs(enabled = true) { enabled, staleTime: 5_000, refetchOnMount: true, - refetchOnWindowFocus: true, + /** Unconditional: the interval is off while nothing is listed, so a run + * another client started during the last `staleTime` window would be + * invisible until some unrelated refetch, and returning to the tab is + * exactly when a pane needs to know whether its history has moved. */ + refetchOnWindowFocus: 'always', refetchInterval: (data) => ((data?.activeJobIds?.length ?? 0) > 0 ? 5_000 : false), retry: false, }); diff --git a/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx b/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx index a776bd4f40..93fb20a020 100644 --- a/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx +++ b/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx @@ -1,6 +1,7 @@ import { renderHook, act } from '@testing-library/react'; -import { Constants, ContentTypes } from 'librechat-data-provider'; import { RecoilRoot, useRecoilValue, useSetRecoilState } from 'recoil'; +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; +import { Constants, ContentTypes, QueryKeys } from 'librechat-data-provider'; import type { TMessage, TConversation, TSubmission } from 'librechat-data-provider'; import type { MutableSnapshot } from 'recoil'; import type { ReactNode } from 'react'; @@ -9,10 +10,13 @@ import useResumeOnLoad from '../useResumeOnLoad'; import store from '~/store'; const mockUseStreamStatus = jest.fn(); +const mockUseActiveJobs = jest.fn(); jest.mock('~/data-provider', () => ({ useStreamStatus: (conversationId: string | undefined, enabled: boolean) => mockUseStreamStatus(conversationId, enabled), + useActiveJobs: (enabled?: boolean) => mockUseActiveJobs(enabled), + streamStatusQueryKey: (conversationId: string) => ['streamStatus', conversationId], })); const CONVERSATION_ID = 'conv-current'; @@ -85,6 +89,9 @@ function renderUseResumeOnLoad({ onQueuedMessages?: (queued: QueuedMessage[]) => void; }) { const getMessages = jest.fn(getMessagesOverride ?? (() => messages)); + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }); let setSubmissionState: ((submission: TSubmission | null) => void) | undefined; const initializeState = (snapshot: MutableSnapshot) => { snapshot.set(store.conversationByIndex(0), buildConversation(conversationId)); @@ -119,16 +126,19 @@ function renderUseResumeOnLoad({ }; const wrapper = ({ children }: { children: ReactNode }) => ( - - - - - - {children} - + + + + + + + {children} + + ); return { + queryClient, getMessages, setSubmission: (nextSubmission: TSubmission | null) => setSubmissionState?.(nextSubmission), ...renderHook(() => useResumeOnLoad(conversationId, getMessages, 0, messagesLoaded), { @@ -146,6 +156,8 @@ describe('useResumeOnLoad', () => { isSuccess: false, isFetching: false, }); + mockUseActiveJobs.mockReset(); + mockUseActiveJobs.mockReturnValue({ data: { activeJobIds: [] }, dataUpdatedAt: 1 }); }); afterEach(() => { @@ -195,6 +207,290 @@ describe('useResumeOnLoad', () => { expect(mockUseStreamStatus).toHaveBeenLastCalledWith(CONVERSATION_ID, false); }); + describe('a run started by another client', () => { + /** Mirrors `ACTIVE_JOB_REARM_INTERVAL_MS` in the hook. */ + const ACTIVE_JOB_REARM_INTERVAL_MS = 5_000; + + const INACTIVE_STATUS = { + isSuccess: true, + isFetching: false, + data: { active: false }, + }; + + const ACTIVE_STATUS = { + isSuccess: true, + isFetching: false, + data: { + active: true, + status: 'running', + createdAt: 4242, + streamId: CONVERSATION_ID, + resumeState: { + aggregatedContent: [{ type: ContentTypes.TEXT, text: 'partial' }], + responseMessageId: RESPONSE_MESSAGE_ID, + userMessage: { messageId: USER_MESSAGE_ID, conversationId: CONVERSATION_ID }, + }, + }, + }; + + it('stops asking about a conversation once it has answered inactive', async () => { + mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); + + const { rerender } = renderUseResumeOnLoad({ messages: [] }); + await act(async () => { + await Promise.resolve(); + }); + + expect(mockUseStreamStatus).toHaveBeenCalledWith(CONVERSATION_ID, true); + + rerender(); + + /** Documents the gap this suite's next case closes: the status query is + * the only thing that could notice another client's run, and it is now + * switched off for the rest of this conversation's mount. */ + expect(mockUseStreamStatus).toHaveBeenLastCalledWith(CONVERSATION_ID, false); + }); + + it('attaches when a job for the viewed conversation becomes active elsewhere', async () => { + const observedSubmissions: Array = []; + mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); + + const { rerender } = renderUseResumeOnLoad({ + messages: [buildUserMessage(CONVERSATION_ID)], + onSubmission: (currentSubmission) => observedSubmissions.push(currentSubmission), + }); + await act(async () => { + await Promise.resolve(); + }); + expect(observedSubmissions[observedSubmissions.length - 1]).toBeNull(); + + /** Another tab (or another device) sends into this same conversation. + * `/chat/active` is scoped to the user, not to the client that started + * the run, so this pane can see it — and it refetches on focus. */ + mockUseStreamStatus.mockClear(); + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [CONVERSATION_ID] }, + dataUpdatedAt: 2, + }); + mockUseStreamStatus.mockReturnValue(ACTIVE_STATUS); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + + /** Re-armed: the status query re-opens off the announcement... */ + expect(mockUseStreamStatus).toHaveBeenCalledWith(CONVERSATION_ID, true); + /** ...and closes again once the attachment exists, rather than staying + * open for the run's whole lifetime. */ + expect(mockUseStreamStatus).toHaveBeenLastCalledWith(CONVERSATION_ID, false); + const attached = observedSubmissions[observedSubmissions.length - 1] as + | (TSubmission & { resumeStreamId?: string; resumeGenerationCreatedAt?: number }) + | null; + expect(attached?.resumeStreamId).toBe(CONVERSATION_ID); + expect(attached?.resumeGenerationCreatedAt).toBe(4242); + }); + + it('re-arms once per run rather than on every poll of the active list', async () => { + mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); + const { rerender } = renderUseResumeOnLoad({ messages: [] }); + await act(async () => { + await Promise.resolve(); + }); + + /** The job is listed, but the status read that answers the announcement + * finds it already finished — a run that ended between the two calls. */ + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [CONVERSATION_ID] }, + dataUpdatedAt: 2, + }); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + expect(mockUseStreamStatus).toHaveBeenCalledWith(CONVERSATION_ID, true); + + mockUseStreamStatus.mockClear(); + /** The list keeps reporting the same job on its heartbeat. Each poll is a + * fresh fetch stamp, so the effect does run — only the elapsed-time gate + * stops it from re-reading status every time. */ + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [CONVERSATION_ID] }, + dataUpdatedAt: 3, + }); + rerender(); + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [CONVERSATION_ID] }, + dataUpdatedAt: 4, + }); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + + expect(mockUseStreamStatus).not.toHaveBeenCalledWith(CONVERSATION_ID, true); + }); + + it('forces fresh status and history reads instead of trusting warm caches', async () => { + mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); + const { rerender, queryClient } = renderUseResumeOnLoad({ messages: [] }); + await act(async () => { + await Promise.resolve(); + }); + + const invalidate = jest.spyOn(queryClient, 'invalidateQueries').mockResolvedValue(undefined); + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [CONVERSATION_ID] }, + dataUpdatedAt: 2, + }); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + + /** An inactive status answered inside its `staleTime` is still fresh, so + * re-enabling the query would replay "nothing running" and the + * announcement would be consumed without ever attaching. */ + expect(invalidate).toHaveBeenCalledWith({ + queryKey: ['streamStatus', CONVERSATION_ID], + }); + /** History has to be authoritative whichever way the status lands: the + * resume submission and `finalHandler` both build on this snapshot. */ + expect(invalidate).toHaveBeenCalledWith({ + queryKey: [QueryKeys.messages, CONVERSATION_ID], + }); + }); + + it('refetches history even when the announced run has already finished', async () => { + const observedSubmissions: Array = []; + mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); + const { rerender, queryClient } = renderUseResumeOnLoad({ + messages: [buildUserMessage(CONVERSATION_ID)], + onSubmission: (currentSubmission) => observedSubmissions.push(currentSubmission), + }); + await act(async () => { + await Promise.resolve(); + }); + + const invalidate = jest.spyOn(queryClient, 'invalidateQueries').mockResolvedValue(undefined); + /** Announced, but the run ends before the status read answers. Nothing + * attaches — so the refetch is the entire repair, and without it the + * turns that other client just completed stay missing here. */ + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [CONVERSATION_ID] }, + dataUpdatedAt: 2, + }); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + + expect(observedSubmissions[observedSubmissions.length - 1]).toBeNull(); + expect(invalidate).toHaveBeenCalledWith({ + queryKey: [QueryKeys.messages, CONVERSATION_ID], + }); + }); + + it('answers a second external run that never left the list', async () => { + const nowSpy = jest.spyOn(Date, 'now'); + let clock = 1_000_000; + nowSpy.mockImplementation(() => clock); + + mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); + const { rerender } = renderUseResumeOnLoad({ messages: [] }); + await act(async () => { + await Promise.resolve(); + }); + + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [CONVERSATION_ID] }, + dataUpdatedAt: 2, + }); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + + mockUseStreamStatus.mockClear(); + /** + * The first run ended and a second began between two polls of the active + * list, so the list reads `[conversationId]` the whole time and never + * shows the gap a latch would need. Only elapsed time distinguishes them. + */ + clock += ACTIVE_JOB_REARM_INTERVAL_MS + 1; + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [CONVERSATION_ID] }, + dataUpdatedAt: 3, + }); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + + expect(mockUseStreamStatus).toHaveBeenCalledWith(CONVERSATION_ID, true); + nowSpy.mockRestore(); + }); + + it('re-arms again for the next run once the previous one leaves the list', async () => { + mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); + const { rerender } = renderUseResumeOnLoad({ messages: [] }); + await act(async () => { + await Promise.resolve(); + }); + + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [CONVERSATION_ID] }, + dataUpdatedAt: 2, + }); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + + mockUseActiveJobs.mockReturnValue({ data: { activeJobIds: [] }, dataUpdatedAt: 1 }); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + + mockUseStreamStatus.mockClear(); + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [CONVERSATION_ID] }, + dataUpdatedAt: 2, + }); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + + expect(mockUseStreamStatus).toHaveBeenCalledWith(CONVERSATION_ID, true); + }); + + it('does not re-arm for a conversation that is not the one being viewed', async () => { + const observedSubmissions: Array = []; + mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); + + const { rerender } = renderUseResumeOnLoad({ + messages: [buildUserMessage(CONVERSATION_ID)], + onSubmission: (currentSubmission) => observedSubmissions.push(currentSubmission), + }); + await act(async () => { + await Promise.resolve(); + }); + + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [STALE_CONVERSATION_ID] }, + dataUpdatedAt: 2, + }); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + + expect(mockUseStreamStatus).toHaveBeenLastCalledWith(CONVERSATION_ID, false); + expect(observedSubmissions[observedSubmissions.length - 1]).toBeNull(); + }); + }); + it('does not replace a null-conversation submission when stream status matches its resume state', async () => { const submission = buildSubmission(null); const observedSubmissions: Array = []; diff --git a/client/src/hooks/SSE/useResumeOnLoad.ts b/client/src/hooks/SSE/useResumeOnLoad.ts index 74a6df54ec..9ff950cc73 100644 --- a/client/src/hooks/SSE/useResumeOnLoad.ts +++ b/client/src/hooks/SSE/useResumeOnLoad.ts @@ -1,6 +1,12 @@ -import { useEffect, useRef } from 'react'; +import { useEffect, useRef, useState } from 'react'; +import { useQueryClient } from '@tanstack/react-query'; import { useSetRecoilState, useRecoilValue, useRecoilCallback } from 'recoil'; -import { Constants, tMessageSchema, isAssistantsEndpoint } from 'librechat-data-provider'; +import { + Constants, + QueryKeys, + tMessageSchema, + isAssistantsEndpoint, +} from 'librechat-data-provider'; import type { TMessage, TConversation, TSubmission, Agents } from 'librechat-data-provider'; import type { GenerationProtocolVersion } from '~/data-provider/SSE/protocol'; import type { StreamStatusResponse } from '~/data-provider'; @@ -16,10 +22,16 @@ import { getGenerationProtocolVersion, supportsGenerationProtocolV2, } from '~/data-provider/SSE/protocol'; +import { useStreamStatus, useActiveJobs, streamStatusQueryKey } from '~/data-provider'; import useSteerConvert from '~/hooks/Chat/useSteerConvert'; -import { useStreamStatus } from '~/data-provider'; import store from '~/store'; +/** + * Matches the active-job list's own poll interval: answering an announcement + * faster than the list can change it would only re-read the same snapshot. + */ +const ACTIVE_JOB_REARM_INTERVAL_MS = 5_000; + function hasSubmissionUserMessage( submission: TSubmission | null, messages: TMessage[] | undefined, @@ -223,6 +235,7 @@ export default function useResumeOnLoad( runIndex = 0, messagesLoaded = true, ) { + const queryClient = useQueryClient(); const setSubmission = useSetRecoilState(store.submissionByIndex(runIndex)); const currentSubmission = useRecoilValue(store.submissionByIndex(runIndex)); const currentConversation = useRecoilValue(store.conversationByIndex(runIndex)); @@ -232,6 +245,21 @@ export default function useResumeOnLoad( const resumableEnabled = !isAssistantsEndpoint(actualEndpoint); // Track conversations we've already processed (either resumed or skipped) const processedConvoRef = useRef(null); + /** + * When this pane last answered an active-job announcement, per conversation. + * A job stays listed for its whole lifetime, so the announcement cannot be + * consumed once and latched: two external runs back to back leave the list + * reading `[conversationId]` continuously, and a latch keyed on observing it + * empty would never release. Rate-limit to the list's own heartbeat instead — + * repeatable, and still one status read per interval at worst. + */ + const answeredActiveJobRef = useRef<{ conversationId: string; at: number } | null>(null); + /** + * Bumped when an announcement is answered. Clearing `processedConvoRef` alone + * cannot restart the check below: a ref mutation neither schedules a render + * nor re-runs an effect, so the arm has to be a value the effect depends on. + */ + const [externalRunArm, setExternalRunArm] = useState(0); /** `generationHandoff` lives in the React Query snapshot until a later * status refetch. Remember the exact epoch already consumed so clearing the * replacement submission on FINAL cannot re-install that stale snapshot and @@ -371,6 +399,33 @@ export default function useResumeOnLoad( const hasStaleSubmissionForDifferentConvo = !!currentSubmission && submissionConvoId != null && submissionConvoId !== conversationId; + /** + * A run this pane did not start — another tab, another device, a scheduled + * trigger — announces itself only through the user-scoped active-job list, + * which already polls while anything is running and refetches on focus. The + * status query is the one thing that could turn that into an attachment, and + * it switches off for the rest of this conversation's mount the moment it has + * answered inactive once. Without a re-arm the pane sits on history it cannot + * see has moved on, and a send from here forks a branch off a stale tail. + * + * Consumed once per run rather than held open: a job stays listed for its + * whole lifetime, and re-opening the query on every poll would turn a + * five-second heartbeat into a five-second status read. + */ + /** + * `dataUpdatedAt` is the trigger, not `activeJobsData`. React Query keeps the + * previous reference when a refetch is deep-equal, and a run that stays + * listed produces exactly that — so the derived boolean below never changes + * and an effect keyed on it would run once and never again. The stamp moves + * on every fetch, which is the heartbeat this needs. + */ + const { data: activeJobsData, dataUpdatedAt: activeJobsUpdatedAt } = + useActiveJobs(resumableEnabled); + const hasActiveJobForThisConvo = + !!conversationId && + conversationId !== Constants.NEW_CONVO && + activeJobsData?.activeJobIds?.includes(conversationId) === true; + const shouldCheck = resumableEnabled && messagesLoaded && // Wait for messages to load before checking @@ -599,6 +654,7 @@ export default function useResumeOnLoad( settleAppliedSteerParts, convertSteersToQueued, setActiveGenerationCreatedAt, + externalRunArm, ]); // Reset processedConvoRef when conversation changes to allow re-checking @@ -611,6 +667,65 @@ export default function useResumeOnLoad( }); processedConvoRef.current = null; consumedHandoffGenerationRef.current = null; + answeredActiveJobRef.current = null; } }, [conversationId]); + + /** + * Answer an active-job announcement for the conversation on screen. + * + * The announcement means "your history may have moved", not merely "re-open + * the status query", so both reads this pane is about to make are forced + * fresh first. Toggling `enabled` alone would not do it: an inactive status + * answered less than `staleTime` ago is still fresh, and React Query would + * hand back that cached "nothing running" and let the effect above record the + * announcement as handled without ever attaching. + * + * History is invalidated for the same reason, and it matters whichever way + * the status read lands. An external client may have completed whole turns + * this pane never saw before starting the one now running: the resume + * submission and `finalHandler` both build on this snapshot, so attaching + * without it grafts the live turn onto a hole. And when the run turns out to + * have already finished, the refetch is the entire repair — the messages + * query disables refetch on focus, mount and reconnect, so nothing else would + * ever collect those turns. The refetch also re-gates `messagesLoaded`, which + * defers the effect above until the authoritative history has landed. + */ + useEffect(() => { + if ( + !resumableEnabled || + !hasActiveJobForThisConvo || + !conversationId || + hasActiveSubmissionForThisConvo + ) { + if (!hasActiveJobForThisConvo) { + answeredActiveJobRef.current = null; + } + return; + } + + const answered = answeredActiveJobRef.current; + const now = Date.now(); + if ( + answered != null && + answered.conversationId === conversationId && + now - answered.at < ACTIVE_JOB_REARM_INTERVAL_MS + ) { + return; + } + answeredActiveJobRef.current = { conversationId, at: now }; + + console.log('[ResumeOnLoad] Active job announced without an attachment', { conversationId }); + queryClient.invalidateQueries({ queryKey: streamStatusQueryKey(conversationId) }); + queryClient.invalidateQueries({ queryKey: [QueryKeys.messages, conversationId] }); + processedConvoRef.current = null; + setExternalRunArm((arm) => arm + 1); + }, [ + conversationId, + resumableEnabled, + hasActiveJobForThisConvo, + hasActiveSubmissionForThisConvo, + activeJobsUpdatedAt, + queryClient, + ]); }