diff --git a/client/src/hooks/Chat/__tests__/abort.spec.tsx b/client/src/hooks/Chat/__tests__/abort.spec.tsx index 279bd3a9841..2efd0372ff1 100644 --- a/client/src/hooks/Chat/__tests__/abort.spec.tsx +++ b/client/src/hooks/Chat/__tests__/abort.spec.tsx @@ -6,7 +6,12 @@ import { RecoilRoot, useSetRecoilState, useRecoilValue } from 'recoil'; import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; import type { TSubmission, TConversation } from 'librechat-data-provider'; import type { DrainAfterAbort, RunEnd } from '~/hooks/Chat/queue'; -import { drainAfterAbortByIndex, runEndByIndex, resetQueueFamilies } from '~/hooks/Chat/queue'; +import { + stopRequestedByConvoId, + drainAfterAbortByIndex, + resetQueueFamilies, + runEndByIndex, +} from '~/hooks/Chat/queue'; import useChatHelpers from '../useChatHelpers'; import { useAbortCleanup } from '../abort'; import store from '~/store'; @@ -380,6 +385,9 @@ describe('useChatHelpers stopGenerating (abort steer targeting)', () => { expect(mockConvertSteersToQueued).not.toHaveBeenCalled(); expect(observed.runEnd).toBeNull(); expect(observed.drainAfterAbort).toBe(false); + expect(getDefaultStore().get(stopRequestedByConvoId('convo-held'))).toEqual({ + generationCreatedAt: 41, + }); }); it('does not release the abort-drain queue when natural completion already settled the run', async () => { @@ -401,6 +409,8 @@ describe('useChatHelpers stopGenerating (abort steer targeting)', () => { expect(mockConvertSteersToQueued).not.toHaveBeenCalled(); expect(observed.runEnd).toBeNull(); expect(observed.drainAfterAbort).toBe(false); + /** Stop lost the race, so the run's persisted outcome decides the queue if the user leaves. */ + expect(getDefaultStore().get(stopRequestedByConvoId('convo-held'))).toBeNull(); }); it('does not let a late abort response clear another conversation epoch', async () => { diff --git a/client/src/hooks/Chat/__tests__/queue.spec.ts b/client/src/hooks/Chat/__tests__/queue.spec.ts new file mode 100644 index 00000000000..254cdf20d7f --- /dev/null +++ b/client/src/hooks/Chat/__tests__/queue.spec.ts @@ -0,0 +1,234 @@ +import { createStore } from 'jotai'; +import type { TMessage } from 'librechat-data-provider'; +import { + settledQueuedTurnReceiptsByConvoId, + queuedMessagesByConvoId, + stopRequestedByConvoId, + resolveDetachedRunEnd, + canResolveDetachedRun, + isDetachedRunManual, + isStoppedRun, +} from '../queue'; + +const CONVO_ID = 'convo-detached'; +const USER_ID = 'user-1'; + +const response = (overrides: Partial = {}): TMessage => + ({ + messageId: 'response-1', + parentMessageId: USER_ID, + conversationId: CONVO_ID, + isCreatedByUser: false, + text: 'done', + ...overrides, + }) as TMessage; + +const userMessage = { messageId: USER_ID, isCreatedByUser: true } as TMessage; + +const writtenAt = (epoch: number) => new Date(epoch).toISOString(); + +describe('resolveDetachedRunEnd', () => { + it('reports a completed run with its persisted response', () => { + const end = resolveDetachedRunEnd(CONVO_ID, { userMessageId: USER_ID }, [ + userMessage, + response(), + ]); + expect(end).toEqual( + expect.objectContaining({ + conversationId: CONVO_ID, + outcome: 'completed', + responseMessageId: 'response-1', + }), + ); + }); + + it('reports a failed run as an error, so the queue waits for a manual send', () => { + const end = resolveDetachedRunEnd(CONVO_ID, { userMessageId: USER_ID }, [ + response({ error: true }), + ]); + expect(end?.outcome).toBe('error'); + expect(end?.responseMessageId).toBeUndefined(); + }); + + it('reports an unfinished response as aborted', () => { + const end = resolveDetachedRunEnd(CONVO_ID, { userMessageId: USER_ID }, [ + response({ unfinished: true }), + ]); + expect(end?.outcome).toBe('aborted'); + }); + + it('returns nothing while no response to the run is persisted', () => { + expect(resolveDetachedRunEnd(CONVO_ID, { userMessageId: USER_ID }, [userMessage])).toBeNull(); + expect(resolveDetachedRunEnd(CONVO_ID, { userMessageId: USER_ID }, undefined)).toBeNull(); + }); + + it('picks the response the run created among regenerated siblings', () => { + const end = resolveDetachedRunEnd( + CONVO_ID, + { userMessageId: USER_ID, responseMessageId: 'response-2_' }, + [response({ messageId: 'response-1', error: true }), response({ messageId: 'response-2' })], + ); + expect(end).toEqual( + expect.objectContaining({ outcome: 'completed', responseMessageId: 'response-2' }), + ); + }); + + it('matches a persisted response id that itself ends in an underscore', () => { + const end = resolveDetachedRunEnd( + CONVO_ID, + { userMessageId: USER_ID, responseMessageId: 'response-2_' }, + [response({ messageId: 'response-1' }), response({ messageId: 'response-2_' })], + ); + expect(end?.responseMessageId).toBe('response-2_'); + }); + + it('leaves a run unresolved when the response it named is not loaded', () => { + expect( + resolveDetachedRunEnd(CONVO_ID, { userMessageId: USER_ID, responseMessageId: 'response-2' }, [ + response({ messageId: 'response-1' }), + ]), + ).toBeNull(); + }); + + it('does not stand a sibling in for the response a fresh run named', () => { + const run = { + userMessageId: USER_ID, + responseMessageId: 'response-a', + generationCreatedAt: 5000, + }; + expect( + resolveDetachedRunEnd(CONVO_ID, run, [ + response({ messageId: 'response-b', updatedAt: writtenAt(6000) }), + ]), + ).toBeNull(); + }); + + it('treats the padded user id placeholder as naming no response', () => { + const end = resolveDetachedRunEnd( + CONVO_ID, + { userMessageId: USER_ID, responseMessageId: `${USER_ID}_` }, + [response({ messageId: 'server-response' })], + ); + expect(end?.responseMessageId).toBe('server-response'); + }); + + it('leaves a regeneration without its epoch unresolved, since history still holds the reply it replaces', () => { + expect( + resolveDetachedRunEnd( + CONVO_ID, + { userMessageId: USER_ID, responseMessageId: 'response-1_', isRegenerate: true }, + [response({ messageId: 'response-1' })], + ), + ).toBeNull(); + }); + + it('carries the run epoch so the drain can match server admission receipts', () => { + const end = resolveDetachedRunEnd( + CONVO_ID, + { userMessageId: USER_ID, generationCreatedAt: 4200 }, + [response({ updatedAt: writtenAt(4300) })], + ); + expect(end?.generationCreatedAt).toBe(4200); + }); + + it('rejects the reply a regeneration replaces and resolves the one the run wrote', () => { + const run = { + userMessageId: USER_ID, + responseMessageId: 'response-1_', + isRegenerate: true, + generationCreatedAt: 5000, + }; + const replaced = response({ messageId: 'response-1', updatedAt: writtenAt(4000) }); + expect(resolveDetachedRunEnd(CONVO_ID, run, [replaced])).toBeNull(); + const end = resolveDetachedRunEnd(CONVO_ID, run, [ + replaced, + response({ messageId: 'response-2', unfinished: true, updatedAt: writtenAt(6000) }), + ]); + expect(end).toEqual(expect.objectContaining({ outcome: 'aborted', generationCreatedAt: 5000 })); + }); + + it('ignores a row the run never wrote, even under the id it named', () => { + expect( + resolveDetachedRunEnd( + CONVO_ID, + { userMessageId: USER_ID, responseMessageId: 'response-1', generationCreatedAt: 5000 }, + [ + response({ messageId: 'response-1', updatedAt: writtenAt(4999) }), + response({ messageId: 'other' }), + ], + ), + ).toBeNull(); + }); + + it('does not guess between siblings when the run named no response', () => { + expect( + resolveDetachedRunEnd(CONVO_ID, { userMessageId: USER_ID }, [ + response({ messageId: 'response-1' }), + response({ messageId: 'response-2' }), + ]), + ).toBeNull(); + }); +}); + +describe('canResolveDetachedRun', () => { + it('can date a regeneration only by its epoch', () => { + expect(canResolveDetachedRun({ userMessageId: USER_ID })).toBe(true); + expect(canResolveDetachedRun({ userMessageId: USER_ID, isRegenerate: true })).toBe(false); + expect( + canResolveDetachedRun({ userMessageId: USER_ID, isRegenerate: true, generationCreatedAt: 1 }), + ).toBe(true); + }); +}); + +describe('isStoppedRun', () => { + it('holds a Stop to the run it was pressed on', () => { + const stopped = { generationCreatedAt: 1000 }; + expect(isStoppedRun(null, 1000)).toBe(false); + expect(isStoppedRun(stopped, 1000)).toBe(true); + expect(isStoppedRun(stopped, 2000)).toBe(false); + expect(isStoppedRun(stopped, undefined)).toBe(true); + }); +}); + +describe('isDetachedRunManual', () => { + const CONVO = 'convo-manual'; + const epochRun = { userMessageId: USER_ID, generationCreatedAt: 1000 }; + const epochless = { userMessageId: USER_ID }; + + it('resolves an ordinary run from history', () => { + const store = createStore(); + expect(isDetachedRunManual(store.get, CONVO, epochRun)).toBe(false); + expect(isDetachedRunManual(store.get, CONVO, epochless)).toBe(false); + }); + + it('leaves the run the user stopped, and only that run, for a manual send', () => { + const store = createStore(); + store.set(stopRequestedByConvoId(CONVO), { generationCreatedAt: 1000 }); + expect(isDetachedRunManual(store.get, CONVO, epochRun)).toBe(true); + expect(isDetachedRunManual(store.get, CONVO, { ...epochRun, generationCreatedAt: 2000 })).toBe( + false, + ); + expect(isDetachedRunManual(store.get, CONVO, epochless)).toBe(true); + }); + + it('leaves an epochless regeneration for a manual send', () => { + const store = createStore(); + expect(isDetachedRunManual(store.get, CONVO, { ...epochless, isRegenerate: true })).toBe(true); + expect(isDetachedRunManual(store.get, CONVO, { ...epochRun, isRegenerate: true })).toBe(false); + }); + + it('leaves an epochless run whose queue the server shares for a manual send', () => { + const receipts = createStore(); + receipts.set(settledQueuedTurnReceiptsByConvoId(CONVO), [ + { clientRequestId: 'queued-1', status: 'admitted', effectivePredecessorCreatedAt: 41 }, + ]); + expect(isDetachedRunManual(receipts.get, CONVO, epochless)).toBe(true); + expect(isDetachedRunManual(receipts.get, CONVO, epochRun)).toBe(false); + + const serverRow = createStore(); + serverRow.set(queuedMessagesByConvoId(CONVO), [ + { id: 'row-1', text: 'next', createdAt: 1, server: { status: 'queued' } }, + ]); + expect(isDetachedRunManual(serverRow.get, CONVO, epochless)).toBe(true); + }); +}); diff --git a/client/src/hooks/Chat/queue.ts b/client/src/hooks/Chat/queue.ts index bde2bf64685..a648b49200a 100644 --- a/client/src/hooks/Chat/queue.ts +++ b/client/src/hooks/Chat/queue.ts @@ -1,7 +1,7 @@ import { atom } from 'jotai'; import { atomFamily, RESET } from 'jotai/utils'; import type { TMessage } from 'librechat-data-provider'; -import type { PrimitiveAtom } from 'jotai'; +import type { PrimitiveAtom, Getter } from 'jotai'; import type { GenerationProtocolVersion } from '~/data-provider/SSE/protocol'; /** @@ -226,6 +226,141 @@ export const drainAfterAbortByIndex = atomFamily((_index: string | number) => atom(false), ); +/** A run whose stream this pane closed while it was still generating, because the user left its + * chat. The server keeps going and deletes the job once it finishes, so the run's end is learned + * from the persisted response when the user comes back instead of from the stream. */ +export type DetachedRun = { + /** The user message the run answers. */ + userMessageId: string; + /** The response placeholder's id, when the submission carried one. */ + responseMessageId?: string; + /** The run's generation epoch, when the start response installed one. The queue drain matches + * server admission receipts against it. */ + generationCreatedAt?: number; + /** A regeneration rewrites a response that already exists in history, so history cannot tell + * whether the row it finds is the old reply or the new one. */ + isRegenerate?: boolean; +}; + +export const detachedRunByConvoId = atomFamily((_conversationId: string) => + atom(null), +); + +/** The run the user pressed Stop on, named by its generation epoch. A stopped run's response + * persists like a completed one, so a run stopped and then left before its abort event arrived + * must not be resolved as detached and drain the queue. Only a different run resets it, so a + * stopped run the user leaves, reattaches to and leaves again stays stopped. */ +export type StoppedRun = { generationCreatedAt: number }; + +export const stopRequestedByConvoId = atomFamily((_conversationId: string) => + atom(null), +); + +/** Whether a Stop marker belongs to the run with this epoch. A run whose epoch is unknown cannot + * be told apart from the stopped one, so it is treated as stopped. */ +export function isStoppedRun(stopped: StoppedRun | null, generationCreatedAt?: number): boolean { + if (stopped == null) { + return false; + } + return generationCreatedAt == null || stopped.generationCreatedAt === generationCreatedAt; +} + +/** Whether persisted history can ever name the detached run's response. A regeneration leaves the + * previous reply under the same user message, so without the run's epoch to date the rows by, + * no read can tell the old reply from the new one. */ +export function canResolveDetachedRun(run: DetachedRun): boolean { + return run.isRegenerate !== true || run.generationCreatedAt != null; +} + +/** + * Whether a detached run's end must be left for a manual send instead of read from history, which + * retires its marker. Three runs qualify: one the user stopped, whose response persists like a + * completed one (a stopped compaction even settles with `unfinished: false`); a regeneration with + * no epoch, whose new reply history cannot tell from the old one; and a run with no epoch whose + * queue the server shares, since the drain cannot match its end to a server admission receipt. + */ +export function isDetachedRunManual( + get: Getter, + conversationId: string, + run: DetachedRun, +): boolean { + if (!canResolveDetachedRun(run)) { + return true; + } + if (isStoppedRun(get(stopRequestedByConvoId(conversationId)), run.generationCreatedAt)) { + return true; + } + if (run.generationCreatedAt != null) { + return false; + } + return ( + get(settledQueuedTurnReceiptsByConvoId(conversationId)).length > 0 || + get(queuedMessagesByConvoId(conversationId)).some((item) => item.server != null) + ); +} + +const writtenAt = (message: TMessage): number => { + const stamp = message.updatedAt ?? message.createdAt; + return stamp == null ? Number.NaN : Date.parse(stamp); +}; + +/** + * The run end a detached run implies, read from persisted history: `completed` lets the queue + * drain, while a stopped (`unfinished`) or failed response leaves the queue for a manual send, as + * it would had the stream been attached. With the run's epoch, only a response the server wrote at + * or after the run started can be its response, which rules out the reply a regeneration replaces + * and any row the run never touched. Returns `null` while no such response is persisted, so + * nothing is drained on a guess. + */ +export function resolveDetachedRunEnd( + conversationId: string, + run: DetachedRun, + messages: TMessage[] | undefined, +): RunEnd | null { + if (!canResolveDetachedRun(run)) { + return null; + } + const epoch = run.generationCreatedAt; + const responses = (messages ?? []).filter( + (message) => + message.isCreatedByUser === false && + message.parentMessageId === run.userMessageId && + (epoch == null || writtenAt(message) >= epoch), + ); + /** A fresh turn's placeholder is the user message id padded with `_`, which names no response; + * a real id is matched exactly first, since persisted ids may themselves end in `_`. */ + const exact = run.responseMessageId; + const unpadded = exact?.replace(/_+$/, ''); + const namesResponse = exact != null && unpadded !== run.userMessageId; + let response: TMessage | undefined; + if (namesResponse) { + response = + responses.find((message) => message.messageId === exact) ?? + responses.find((message) => message.messageId === unpadded); + } + /** Only a regeneration writes its reply under an id other than the one it named, so the lone + * response written since the epoch stands in for a named response only there. */ + if (response == null && (!namesResponse || run.isRegenerate === true) && responses.length === 1) { + response = responses[0]; + } + if (response == null) { + return null; + } + let outcome: RunEnd['outcome'] = 'completed'; + if (response.error === true) { + outcome = 'error'; + } else if (response.unfinished === true) { + outcome = 'aborted'; + } + return { + conversationId, + outcome, + endedAt: Date.now(), + ...(epoch != null && { generationCreatedAt: epoch }), + ...(outcome === 'completed' && { responseMessageId: response.messageId }), + }; +} + const clearFamily = (family: { getParams(): Iterable; remove(param: Param): void; @@ -246,4 +381,6 @@ export function resetQueueFamilies(): void { clearFamily(drainAfterAbortByIndex); clearFamily(runEndsByIndex); clearFamily(runEndByIndex); + clearFamily(detachedRunByConvoId); + clearFamily(stopRequestedByConvoId); } diff --git a/client/src/hooks/Chat/useChatHelpers.ts b/client/src/hooks/Chat/useChatHelpers.ts index 234757d866b..e86380a6a7b 100644 --- a/client/src/hooks/Chat/useChatHelpers.ts +++ b/client/src/hooks/Chat/useChatHelpers.ts @@ -5,10 +5,10 @@ import { useRecoilState, useRecoilValue, useSetRecoilState } from 'recoil'; import { Constants, QueryKeys, isAssistantsEndpoint } from 'librechat-data-provider'; import type { TMessage } from 'librechat-data-provider'; import type { ChatContract } from './contract'; +import { stopRequestedByConvoId, drainAfterAbortByIndex, runEndByIndex } from '~/hooks/Chat/queue'; import { useLatestMessage, useLatestMessageId } from '~/hooks/Messages/useLatestMessage'; import { useAbortStreamMutation, supportsGenerationProtocolV2 } from '~/data-provider'; import { siblingIdxFamily, siblingKey } from '~/components/Chat/Messages/Thread/state'; -import { drainAfterAbortByIndex, runEndByIndex } from '~/hooks/Chat/queue'; import { useChatSettings } from '~/Providers/ChatSettingsContext'; import useChatFunctions from '~/hooks/Chat/useChatFunctions'; import useSteerConvert from '~/hooks/Chat/useSteerConvert'; @@ -234,6 +234,9 @@ export default function useChatHelpers(index = 0, paramId?: string): ChatContrac // start the NEXT submission) while the abort response is in flight; // the fallback clear below must not tear down that new run. const submissionAtAbort = captureSubmission(); + const stoppedRun = { generationCreatedAt: activeGenerationCreatedAt }; + const stopFamily = stopRequestedByConvoId(conversationId); + queueStore.set(stopFamily, stoppedRun); try { console.log('[useChatHelpers] Calling abort mutation for:', conversationId); const response = await abortStream({ @@ -252,8 +255,12 @@ export default function useChatHelpers(index = 0, paramId?: string): ChatContrac if (canUseV2AbortResponse && response?.settled === true) { /** Natural completion/error won the terminal CAS, so this request did * not emit an abort FINAL and must not release the abort-drain queue. - * The winning terminal event/status owns reconciliation. */ + * The winning terminal event/status owns reconciliation, and the run + * was not stopped, so its persisted outcome may still drain. */ clearInterruptDrain(conversationId, activeGenerationCreatedAt); + if (queueStore.get(stopFamily) === stoppedRun) { + queueStore.set(stopFamily, null); + } queryClient.invalidateQueries({ queryKey: [QueryKeys.messages, conversationId], refetchType: 'all', @@ -341,6 +348,7 @@ export default function useChatHelpers(index = 0, paramId?: string): ChatContrac conversationId, endpoint, endpointType, + queueStore, activeGenerationCreatedAt, activeGenerationProtocolVersion, abortStream, diff --git a/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts b/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts index d641490d051..705d92e9e42 100644 --- a/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts +++ b/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts @@ -325,7 +325,13 @@ import useResumableSSE, { ABORT_SWEEP_STATUSES, } from '~/hooks/SSE/useResumableSSE'; import useSSE from '~/hooks/SSE/useSSE'; -import { queuedMessagesByConvoId, resetQueueFamilies } from '~/hooks/Chat/queue'; +import { + settledQueuedTurnReceiptsByConvoId, + queuedMessagesByConvoId, + stopRequestedByConvoId, + detachedRunByConvoId, + resetQueueFamilies, +} from '~/hooks/Chat/queue'; const CONV_ID = 'conv-abc-123'; @@ -1012,6 +1018,297 @@ describe('useResumableSSE', () => { unmount(); }); + describe('a run the user leaves mid-stream', () => { + const renderLeavable = async () => { + const chatHelpers = buildChatHelpers(); + const rendered = renderHook( + ({ current }: { current: TSubmission | null }) => useResumableSSE(current, chatHelpers), + { initialProps: { current: buildSubmission() as TSubmission | null } }, + ); + await flushMicrotasks(); + expect(mockSSEInstances.length).toBeGreaterThan(0); + return rendered; + }; + const detachedRun = () => getDefaultStore().get(detachedRunByConvoId(CONV_ID)); + + it('remembers the run when navigation clears the submission, so its end is read on return', async () => { + const { rerender, unmount } = await renderLeavable(); + rerender({ current: {} as TSubmission }); + expect(detachedRun()).toEqual({ userMessageId: 'msg-1', responseMessageId: 'resp-1' }); + unmount(); + }); + + it('remembers the run when switching to a saved chat clears the submission to null', async () => { + const { rerender, unmount } = await renderLeavable(); + rerender({ current: null }); + expect(detachedRun()).toEqual({ userMessageId: 'msg-1', responseMessageId: 'resp-1' }); + unmount(); + }); + + it('marks a regeneration and carries a resumed epoch when the user leaves', async () => { + const chatHelpers = buildChatHelpers(); + const resumed = { + ...buildSubmission(), + isRegenerate: true, + resumeStreamId: CONV_ID, + resumeGenerationCreatedAt: 4200, + } as TSubmission; + const { unmount } = renderHook(() => useResumableSSE(resumed, chatHelpers)); + await flushMicrotasks(); + unmount(); + expect(detachedRun()).toEqual( + expect.objectContaining({ isRegenerate: true, generationCreatedAt: 4200 }), + ); + }); + + it('remembers the run when the chat unmounts mid-stream', async () => { + const { unmount } = await renderLeavable(); + unmount(); + expect(detachedRun()).toEqual({ userMessageId: 'msg-1', responseMessageId: 'resp-1' }); + }); + + it('forgets a left run once a new run starts in that conversation', async () => { + const { rerender, unmount } = await renderLeavable(); + rerender({ current: null }); + expect(detachedRun()).not.toBeNull(); + rerender({ + current: buildSubmission({ + userMessage: { + messageId: 'msg-2', + conversationId: CONV_ID, + text: 'Next', + isCreatedByUser: true, + sender: 'User', + parentMessageId: 'resp-1', + }, + }), + }); + expect(detachedRun()).toBeNull(); + unmount(); + }); + + it('records no detached run once the run end already reached the drain', async () => { + mockFetchStreamStatus.mockResolvedValue({ active: false }); + const { rerender, unmount } = await renderLeavable(); + await act(async () => { + getLastSSE()._emit('error', { responseCode: 404 }); + }); + await waitFor(() => expect(mockSetRunEnd).toHaveBeenCalled()); + rerender({ current: {} as TSubmission }); + expect(detachedRun()).toBeNull(); + unmount(); + }); + + const resumeOf = (generationCreatedAt: number) => + ({ + ...buildSubmission(), + resumeStreamId: CONV_ID, + resumeGenerationCreatedAt: generationCreatedAt, + }) as TSubmission; + + it('resolves a reattached run from history when its job is gone before the attach', async () => { + getDefaultStore().set(detachedRunByConvoId(CONV_ID), { + userMessageId: 'msg-1', + responseMessageId: 'resp-1', + generationCreatedAt: 4200, + }); + mockFetchStreamStatus.mockResolvedValue({ active: false }); + mockGetQueryData.mockImplementation((queryKey: unknown[]) => + queryKey[0] === QueryKeys.messages + ? [ + { + messageId: 'resp-1', + parentMessageId: 'msg-1', + conversationId: CONV_ID, + isCreatedByUser: false, + text: 'finished while away', + updatedAt: new Date(5000).toISOString(), + }, + ] + : undefined, + ); + const resumed = resumeOf(4200); + const chatHelpers = buildChatHelpers(); + const { unmount } = renderHook(() => useResumableSSE(resumed, chatHelpers)); + await flushMicrotasks(); + expect(detachedRun()).not.toBeNull(); + await act(async () => { + getLastSSE()._emit('error', { responseCode: 404 }); + }); + await waitFor(() => + expect(mockSetRunEnd).toHaveBeenCalledWith( + expect.objectContaining({ + outcome: 'completed', + responseMessageId: 'resp-1', + generationCreatedAt: 4200, + }), + ), + ); + expect(mockSetRunEnd).toHaveBeenCalledTimes(1); + expect(detachedRun()).toBeNull(); + unmount(); + }); + + it('leaves an epochless run whose queue the server shares for a manual send on 404', async () => { + mockFetchStreamStatus.mockResolvedValue({ active: false }); + mockGetQueryData.mockImplementation((queryKey: unknown[]) => + queryKey[0] === QueryKeys.messages + ? [ + { + messageId: 'resp-1', + parentMessageId: 'msg-1', + conversationId: CONV_ID, + isCreatedByUser: false, + text: 'finished while away', + }, + ] + : undefined, + ); + const resumed = resumeOf(4200); + const chatHelpers = buildChatHelpers(); + const { unmount } = renderHook(() => useResumableSSE(resumed, chatHelpers)); + await flushMicrotasks(); + getDefaultStore().set(detachedRunByConvoId(CONV_ID), { + userMessageId: 'msg-1', + responseMessageId: 'resp-1', + }); + getDefaultStore().set(settledQueuedTurnReceiptsByConvoId(CONV_ID), [ + { clientRequestId: 'queued-1', status: 'admitted', effectivePredecessorCreatedAt: 41 }, + ]); + await act(async () => { + getLastSSE()._emit('error', { responseCode: 404 }); + }); + await waitFor(() => expect(mockSetRunEnd).toHaveBeenCalledTimes(1)); + expect(mockSetRunEnd).toHaveBeenCalledWith( + expect.objectContaining({ conversationId: CONV_ID, outcome: 'aborted' }), + ); + expect(mockSetRunEnd).not.toHaveBeenCalledWith( + expect.objectContaining({ outcome: 'completed' }), + ); + expect(detachedRun()).toBeNull(); + unmount(); + }); + + it('keeps the marker when 404 history has not caught up with the reattached run', async () => { + const left = { + userMessageId: 'msg-1', + responseMessageId: 'resp-1', + generationCreatedAt: 4200, + }; + getDefaultStore().set(detachedRunByConvoId(CONV_ID), left); + mockFetchStreamStatus.mockResolvedValue({ active: false }); + mockGetQueryData.mockImplementation((queryKey: unknown[]) => + queryKey[0] === QueryKeys.messages ? [] : undefined, + ); + const resumed = resumeOf(4200); + const chatHelpers = buildChatHelpers(); + const { unmount } = renderHook(() => useResumableSSE(resumed, chatHelpers)); + await flushMicrotasks(); + await act(async () => { + getLastSSE()._emit('error', { responseCode: 404 }); + }); + await waitFor(() => expect(mockSetRunEnd).toHaveBeenCalledTimes(1)); + expect(mockSetRunEnd).toHaveBeenCalledWith(expect.objectContaining({ outcome: 'aborted' })); + expect(detachedRun()).toEqual(left); + unmount(); + }); + + it('settles a reattached run the user stopped as stopped on 404, even if it persisted as done', async () => { + getDefaultStore().set(detachedRunByConvoId(CONV_ID), { + userMessageId: 'msg-1', + responseMessageId: 'resp-1', + generationCreatedAt: 4200, + }); + getDefaultStore().set(stopRequestedByConvoId(CONV_ID), { generationCreatedAt: 4200 }); + mockFetchStreamStatus.mockResolvedValue({ active: false }); + mockGetQueryData.mockImplementation((queryKey: unknown[]) => + queryKey[0] === QueryKeys.messages + ? [ + { + messageId: 'resp-1', + parentMessageId: 'msg-1', + conversationId: CONV_ID, + isCreatedByUser: false, + text: 'stopped compaction settles as finished', + updatedAt: new Date(5000).toISOString(), + }, + ] + : undefined, + ); + const resumed = resumeOf(4200); + const chatHelpers = buildChatHelpers(); + const { unmount } = renderHook(() => useResumableSSE(resumed, chatHelpers)); + await flushMicrotasks(); + await act(async () => { + getLastSSE()._emit('error', { responseCode: 404 }); + }); + await waitFor(() => expect(mockSetRunEnd).toHaveBeenCalledTimes(1)); + expect(mockSetRunEnd).toHaveBeenCalledWith(expect.objectContaining({ outcome: 'aborted' })); + expect(detachedRun()).toBeNull(); + unmount(); + }); + + it('does not date a fresh run left before its start response with the previous epoch', async () => { + const chatHelpers = buildChatHelpers(); + const rendered = renderHook( + ({ current }: { current: TSubmission | null }) => useResumableSSE(current, chatHelpers), + { initialProps: { current: resumeOf(4200) as TSubmission | null } }, + ); + await flushMicrotasks(); + rendered.rerender({ current: null }); + expect(detachedRun()).toEqual(expect.objectContaining({ generationCreatedAt: 4200 })); + rendered.rerender({ + current: buildSubmission({ + userMessage: { + messageId: 'msg-2', + conversationId: CONV_ID, + text: 'Next', + isCreatedByUser: true, + sender: 'User', + parentMessageId: 'resp-1', + }, + }), + }); + await flushMicrotasks(); + rendered.rerender({ current: null }); + expect(detachedRun()).toEqual(expect.objectContaining({ userMessageId: 'msg-2' })); + expect(detachedRun()?.generationCreatedAt).toBeUndefined(); + rendered.unmount(); + }); + + it('keeps a Stop through a reattach to the stopped run, so leaving again records nothing', async () => { + getDefaultStore().set(stopRequestedByConvoId(CONV_ID), { generationCreatedAt: 4200 }); + const resumed = resumeOf(4200); + const chatHelpers = buildChatHelpers(); + const { unmount } = renderHook(() => useResumableSSE(resumed, chatHelpers)); + await flushMicrotasks(); + unmount(); + expect(detachedRun()).toBeNull(); + expect(getDefaultStore().get(stopRequestedByConvoId(CONV_ID))).toEqual({ + generationCreatedAt: 4200, + }); + }); + + it('drops a Stop once a different run reattaches, and remembers that run when left', async () => { + getDefaultStore().set(stopRequestedByConvoId(CONV_ID), { generationCreatedAt: 4200 }); + const resumed = resumeOf(9000); + const chatHelpers = buildChatHelpers(); + const { unmount } = renderHook(() => useResumableSSE(resumed, chatHelpers)); + await flushMicrotasks(); + expect(getDefaultStore().get(stopRequestedByConvoId(CONV_ID))).toBeNull(); + unmount(); + expect(detachedRun()).toEqual(expect.objectContaining({ generationCreatedAt: 9000 })); + }); + + it('records no detached run for a run the user stopped before leaving', async () => { + const { rerender, unmount } = await renderLeavable(); + getDefaultStore().set(stopRequestedByConvoId(CONV_ID), { generationCreatedAt: 1000 }); + rerender({ current: {} as TSubmission }); + expect(detachedRun()).toBeNull(); + unmount(); + }); + }); + it('authorizes only the exact failed recovery source past the conversion tombstone', async () => { const parked = [ { steerId: 'failed-source', text: 'retry these words', createdAt: 1 }, diff --git a/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx b/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx index 99e4aa09442..3d1f28897f4 100644 --- a/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx +++ b/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx @@ -1,4 +1,4 @@ -import { renderHook, act } from '@testing-library/react'; +import { renderHook, act, waitFor } from '@testing-library/react'; import { RecoilRoot, useRecoilValue, useSetRecoilState } from 'recoil'; import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; import { Provider as JotaiProvider, createStore, useAtomValue } from 'jotai'; @@ -6,10 +6,16 @@ import { Constants, ContentTypes, QueryKeys } from 'librechat-data-provider'; import type { Agents, TMessage, TConversation, TSubmission } from 'librechat-data-provider'; import type { MutableSnapshot } from 'recoil'; import type { ReactNode } from 'react'; -import type { QueuedMessage } from '~/hooks/Chat/queue'; -import type { PendingSteer } from '~/hooks/Chat/queue'; +import type { QueuedMessage, PendingSteer, DetachedRun } from '~/hooks/Chat/queue'; +import { + settledQueuedTurnReceiptsByConvoId, + queuedMessagesByConvoId, + stopRequestedByConvoId, + pendingRunEndByConvoId, + detachedRunByConvoId, + resetQueueFamilies, +} from '~/hooks/Chat/queue'; import { siblingIdxFamily, siblingKey } from '~/components/Chat/Messages/Thread/state'; -import { queuedMessagesByConvoId, resetQueueFamilies } from '~/hooks/Chat/queue'; import { pendingApprovalActionFamily } from '~/components/Chat/approval/state'; import { agentQueuedTurnsQueryKey } from '~/data-provider/SSE/queuedTurns'; import { submissionStartFamily } from '~/store/generation'; @@ -108,6 +114,9 @@ function renderUseResumeOnLoad({ onQueuedMessages, submissionStart, onSubmissionStart, + detachedRun, + seedJotai, + seedQueryClient, }: { messages?: TMessage[]; getMessages?: () => TMessage[] | undefined; @@ -124,15 +133,23 @@ function renderUseResumeOnLoad({ onQueuedMessages?: (queued: QueuedMessage[]) => void; submissionStart?: number; onSubmissionStart?: (submissionStart: number | null) => void; + detachedRun?: DetachedRun; + seedJotai?: (store: ReturnType) => void; + seedQueryClient?: (client: QueryClient) => void; }) { const getMessages = jest.fn(getMessagesOverride ?? (() => messages)); const jotaiStore = createStore(); if (submissionStart != null) { jotaiStore.set(submissionStartFamily(0), submissionStart); } + if (detachedRun != null) { + jotaiStore.set(detachedRunByConvoId(conversationId), detachedRun); + } + seedJotai?.(jotaiStore); const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } }, }); + seedQueryClient?.(queryClient); let setSubmissionState: ((submission: TSubmission | null) => void) | undefined; let setIsSubmittingState: ((value: boolean) => void) | undefined; let setAttachedEpochState: ((value: number | null) => void) | undefined; @@ -2732,6 +2749,317 @@ describe('useResumeOnLoad', () => { expect(observedSteers[observedSteers.length - 1]).toEqual([failedChip]); }); + it('parks the end of a run that finished after the user left, so its queue drains', async () => { + mockUseStreamStatus.mockReturnValue({ + isSuccess: true, + isFetching: false, + data: { active: false }, + }); + const completedResponse = { + messageId: 'response-detached', + parentMessageId: USER_MESSAGE_ID, + conversationId: CONVERSATION_ID, + isCreatedByUser: false, + text: 'finished while away', + } as TMessage; + + const { jotaiStore } = renderUseResumeOnLoad({ + messages: [buildUserMessage(CONVERSATION_ID), completedResponse], + detachedRun: { userMessageId: USER_MESSAGE_ID }, + }); + await act(async () => { + await Promise.resolve(); + }); + + expect(jotaiStore.get(pendingRunEndByConvoId(CONVERSATION_ID))).toEqual( + expect.objectContaining({ + conversationId: CONVERSATION_ID, + outcome: 'completed', + responseMessageId: 'response-detached', + }), + ); + expect(jotaiStore.get(detachedRunByConvoId(CONVERSATION_ID))).toBeNull(); + }); + + it('keeps the detached run until a history refetch shows its response', async () => { + mockUseStreamStatus.mockReturnValue({ + isSuccess: true, + isFetching: false, + data: { active: false }, + }); + let history: TMessage[] = [buildUserMessage(CONVERSATION_ID)]; + const { jotaiStore } = renderUseResumeOnLoad({ + getMessages: () => history, + detachedRun: { userMessageId: USER_MESSAGE_ID }, + }); + /** The first read found no response; it is saved before the history refetch settles. */ + expect(jotaiStore.get(pendingRunEndByConvoId(CONVERSATION_ID))).toBeNull(); + expect(jotaiStore.get(detachedRunByConvoId(CONVERSATION_ID))).not.toBeNull(); + history = [ + ...history, + { + messageId: 'response-late', + parentMessageId: USER_MESSAGE_ID, + conversationId: CONVERSATION_ID, + isCreatedByUser: false, + text: 'saved after the first read', + } as TMessage, + ]; + + await waitFor(() => + expect(jotaiStore.get(pendingRunEndByConvoId(CONVERSATION_ID))).toEqual( + expect.objectContaining({ outcome: 'completed', responseMessageId: 'response-late' }), + ), + ); + expect(jotaiStore.get(detachedRunByConvoId(CONVERSATION_ID))).toBeNull(); + }); + + it('resolves a detached run from refetched history, not the response it left cached', async () => { + mockUseStreamStatus.mockReturnValue({ + isSuccess: true, + isFetching: false, + data: { active: false }, + }); + const historyKey = [QueryKeys.messages, CONVERSATION_ID]; + const leftResponse = { + messageId: 'response-detached', + parentMessageId: USER_MESSAGE_ID, + conversationId: CONVERSATION_ID, + isCreatedByUser: false, + text: 'streamed before leaving', + } as TMessage; + let client: QueryClient | undefined; + const { jotaiStore } = renderUseResumeOnLoad({ + getMessages: () => client?.getQueryData(historyKey), + detachedRun: { userMessageId: USER_MESSAGE_ID, responseMessageId: 'response-detached' }, + seedQueryClient: (queryClient) => { + client = queryClient; + queryClient.setQueryDefaults(historyKey, { + queryFn: async () => [ + buildUserMessage(CONVERSATION_ID), + { ...leftResponse, unfinished: true }, + ], + }); + queryClient.setQueryData(historyKey, [buildUserMessage(CONVERSATION_ID), leftResponse]); + }, + }); + + await waitFor(() => + expect(jotaiStore.get(pendingRunEndByConvoId(CONVERSATION_ID))).toEqual( + expect.objectContaining({ outcome: 'aborted' }), + ), + ); + expect(jotaiStore.get(detachedRunByConvoId(CONVERSATION_ID))).toBeNull(); + }); + + it('keeps the detached run for the next visit when the history refetch fails', async () => { + mockUseStreamStatus.mockReturnValue({ + isSuccess: true, + isFetching: false, + data: { active: false }, + }); + const historyKey = [QueryKeys.messages, CONVERSATION_ID]; + const queryFn = jest.fn(async (): Promise => { + throw new Error('history unavailable'); + }); + let client: QueryClient | undefined; + const { jotaiStore } = renderUseResumeOnLoad({ + getMessages: () => client?.getQueryData(historyKey), + detachedRun: { userMessageId: USER_MESSAGE_ID, responseMessageId: 'response-detached' }, + seedQueryClient: (queryClient) => { + client = queryClient; + queryClient.setQueryDefaults(historyKey, { queryFn }); + queryClient.setQueryData(historyKey, [ + buildUserMessage(CONVERSATION_ID), + { + messageId: 'response-detached', + parentMessageId: USER_MESSAGE_ID, + conversationId: CONVERSATION_ID, + isCreatedByUser: false, + text: 'streamed before leaving', + } as TMessage, + ]); + }, + }); + + await waitFor(() => expect(queryFn).toHaveBeenCalled()); + await act(async () => { + await Promise.resolve(); + }); + expect(jotaiStore.get(pendingRunEndByConvoId(CONVERSATION_ID))).toBeNull(); + expect(jotaiStore.get(detachedRunByConvoId(CONVERSATION_ID))).not.toBeNull(); + }); + + it('leaves a run without an epoch for a manual send when the server shares the queue', async () => { + mockUseStreamStatus.mockReturnValue({ + isSuccess: true, + isFetching: false, + data: { active: false }, + }); + const { jotaiStore } = renderUseResumeOnLoad({ + messages: [ + buildUserMessage(CONVERSATION_ID), + { + messageId: 'response-detached', + parentMessageId: USER_MESSAGE_ID, + conversationId: CONVERSATION_ID, + isCreatedByUser: false, + text: 'finished while away', + } as TMessage, + ], + detachedRun: { userMessageId: USER_MESSAGE_ID }, + seedJotai: (store) => + store.set(settledQueuedTurnReceiptsByConvoId(CONVERSATION_ID), [ + { clientRequestId: 'queued-1', status: 'admitted', effectivePredecessorCreatedAt: 41 }, + ]), + }); + await act(async () => { + await Promise.resolve(); + }); + + expect(jotaiStore.get(pendingRunEndByConvoId(CONVERSATION_ID))).toBeNull(); + expect(jotaiStore.get(detachedRunByConvoId(CONVERSATION_ID))).toBeNull(); + }); + + it('retires a retained marker for a run the user stopped on a later visit', async () => { + mockUseStreamStatus.mockReturnValue({ + isSuccess: true, + isFetching: false, + data: { active: false }, + }); + const { jotaiStore } = renderUseResumeOnLoad({ + messages: [ + buildUserMessage(CONVERSATION_ID), + { + messageId: 'response-detached', + parentMessageId: USER_MESSAGE_ID, + conversationId: CONVERSATION_ID, + isCreatedByUser: false, + text: 'stopped compaction settles as finished', + unfinished: false, + updatedAt: new Date(42).toISOString(), + } as TMessage, + ], + detachedRun: { userMessageId: USER_MESSAGE_ID, generationCreatedAt: 41 }, + seedJotai: (store) => + store.set(stopRequestedByConvoId(CONVERSATION_ID), { generationCreatedAt: 41 }), + }); + await act(async () => { + await Promise.resolve(); + }); + + expect(jotaiStore.get(pendingRunEndByConvoId(CONVERSATION_ID))).toBeNull(); + expect(jotaiStore.get(detachedRunByConvoId(CONVERSATION_ID))).toBeNull(); + }); + + it('parks a shared-queue run end with its epoch, so its admission receipt is matched', async () => { + mockUseStreamStatus.mockReturnValue({ + isSuccess: true, + isFetching: false, + data: { active: false }, + }); + const { jotaiStore } = renderUseResumeOnLoad({ + messages: [ + buildUserMessage(CONVERSATION_ID), + { + messageId: 'response-detached', + parentMessageId: USER_MESSAGE_ID, + conversationId: CONVERSATION_ID, + isCreatedByUser: false, + text: 'finished while away', + updatedAt: new Date(42).toISOString(), + } as TMessage, + ], + detachedRun: { userMessageId: USER_MESSAGE_ID, generationCreatedAt: 41 }, + seedJotai: (store) => + store.set(settledQueuedTurnReceiptsByConvoId(CONVERSATION_ID), [ + { clientRequestId: 'queued-1', status: 'admitted', effectivePredecessorCreatedAt: 41 }, + ]), + }); + await act(async () => { + await Promise.resolve(); + }); + + expect(jotaiStore.get(pendingRunEndByConvoId(CONVERSATION_ID))).toEqual( + expect.objectContaining({ outcome: 'completed', generationCreatedAt: 41 }), + ); + }); + + it('parks nothing for a conversation this pane never left mid-run', async () => { + mockUseStreamStatus.mockReturnValue({ + isSuccess: true, + isFetching: false, + data: { active: false }, + }); + const { jotaiStore } = renderUseResumeOnLoad({ + messages: [ + buildUserMessage(CONVERSATION_ID), + { + messageId: 'response-old', + parentMessageId: USER_MESSAGE_ID, + conversationId: CONVERSATION_ID, + isCreatedByUser: false, + text: 'history', + } as TMessage, + ], + }); + await act(async () => { + await Promise.resolve(); + }); + + expect(jotaiStore.get(pendingRunEndByConvoId(CONVERSATION_ID))).toBeNull(); + }); + + it('keeps the detached run while its job is active, so a job gone before the attach still resolves', async () => { + mockUseStreamStatus.mockReturnValue({ + 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 }, + }, + }, + }); + const detachedRun = { userMessageId: USER_MESSAGE_ID, generationCreatedAt: 4242 }; + const { jotaiStore } = renderUseResumeOnLoad({ + messages: [buildUserMessage(CONVERSATION_ID)], + detachedRun, + }); + await act(async () => { + await Promise.resolve(); + }); + + expect(jotaiStore.get(detachedRunByConvoId(CONVERSATION_ID))).toEqual(detachedRun); + expect(jotaiStore.get(pendingRunEndByConvoId(CONVERSATION_ID))).toBeNull(); + }); + + it('retires a regeneration history can never date without reading history again', async () => { + mockUseStreamStatus.mockReturnValue({ + isSuccess: true, + isFetching: false, + data: { active: false }, + }); + const refetch = jest.spyOn(QueryClient.prototype, 'refetchQueries'); + const { jotaiStore } = renderUseResumeOnLoad({ + messages: [buildUserMessage(CONVERSATION_ID)], + detachedRun: { userMessageId: USER_MESSAGE_ID, isRegenerate: true }, + }); + await act(async () => { + await Promise.resolve(); + }); + + expect(jotaiStore.get(detachedRunByConvoId(CONVERSATION_ID))).toBeNull(); + expect(jotaiStore.get(pendingRunEndByConvoId(CONVERSATION_ID))).toBeNull(); + expect(refetch).not.toHaveBeenCalled(); + refetch.mockRestore(); + }); + it('converts resumeState.pendingSteers to queued when inactive (expired action, unparked queue)', async () => { const observedSteers: PendingSteer[][] = []; const observedQueues: QueuedMessage[][] = []; diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index 5a7a97a9ed2..d0ee4a9c755 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/client/src/hooks/SSE/useResumableSSE.ts @@ -41,8 +41,8 @@ import type { TContextUsageEvent, ChatStreamConnection, } from 'librechat-data-provider'; +import type { QueuedMessageOrigin, DrainAfterAbort, RunEnd } from '~/hooks/Chat/queue'; import type { ActiveJobsResponse, StreamStatusResponse } from '~/data-provider'; -import type { DrainAfterAbort, QueuedMessageOrigin } from '~/hooks/Chat/queue'; import type { GenerationProtocolVersion } from '~/data-provider'; import type { EventHandlerParams } from './useEventHandlers'; import type { TResData, TFinalResData } from '~/common'; @@ -87,6 +87,16 @@ import { fetchConversationMessages, GENERATION_PROTOCOL_VERSION, } from '~/data-provider'; +import { + resolveDetachedRunEnd, + isDetachedRunManual, + stopRequestedByConvoId, + detachedRunByConvoId, + isStoppedRun, + drainAfterAbortByIndex, + queuedMessagesByConvoId, + runEndByIndex, +} from '~/hooks/Chat/queue'; import { recoveryDispositionsFamily, canRestoreRecovery, @@ -96,7 +106,6 @@ import useEventHandlers, { buildCreatedInitialResponse, keepLocalCodeApprovalMode, } from './useEventHandlers'; -import { drainAfterAbortByIndex, queuedMessagesByConvoId, runEndByIndex } from '~/hooks/Chat/queue'; import { pendingApprovalActionFamily } from '~/components/Chat/approval/state'; import { abortScrollFamily, showStopButtonByIndex } from '~/store/generation'; import { withSubmittedCodeDecision } from '~/hooks/Agents/codeDecision'; @@ -1336,6 +1345,8 @@ export default function useResumableSSE( [convertSteersToQueued], ); + /** The epoch this pane last installed, kept so a run left mid-stream can carry it. */ + const liveEpochRef = useRef<{ conversationId: string; createdAt: number } | null>(null); /** Conversation ids are reused by successive agent turns. Keep the exact * live epoch beside the UI controls, and only clear it if the terminal event * still belongs to the epoch that installed it. */ @@ -1352,6 +1363,7 @@ export default function useResumableSSE( return false; } set(state, next); + liveEpochRef.current = next == null ? null : { conversationId, createdAt: next }; set( store.activeGenerationProtocolVersionByConvoId(conversationId), next == null ? 1 : generationProtocolVersion, @@ -1361,7 +1373,32 @@ export default function useResumableSSE( [], ); - const setRunEnd = useSetAtom(runEndByIndex(runIndex)); + const publishRunEnd = useSetAtom(runEndByIndex(runIndex)); + /** Whether the current submission's run end already reached the queue drain. A run that + * ended attached must not also be resolved as detached when the user later leaves. */ + const runEndPublishedRef = useRef(false); + const setRunEnd = useCallback( + (end: RunEnd, options?: { retainDetached?: boolean }) => { + runEndPublishedRef.current = true; + /** The run's end reached the drain here, so a marker this pane left for the same run is + * settled and must not be resolved a second time on a later visit. A provisional end that + * cannot drain the queue keeps the marker, so a later read of history can still settle it. */ + if (end.conversationId != null && options?.retainDetached !== true) { + const detachedFamily = detachedRunByConvoId(end.conversationId); + const detached = jotaiStore.get(detachedFamily); + if ( + detached != null && + (detached.generationCreatedAt == null || + end.generationCreatedAt == null || + detached.generationCreatedAt === end.generationCreatedAt) + ) { + jotaiStore.set(detachedFamily, null); + } + } + publishRunEnd(end); + }, + [jotaiStore, publishRunEnd], + ); const setDrainAfterAbort = useSetAtom(drainAfterAbortByIndex(runIndex)); const clearDrainAfterAbort = useCallback( (conversationId: string, generationCreatedAt?: number) => { @@ -3457,13 +3494,29 @@ export default function useResumableSSE( } // The true outcome is unknown here (job record already cleaned up): // a non-'completed' outcome releases parked interrupt flags without - // auto-sending queued messages the user may not want fired. - setRunEnd({ - conversationId: recoveryConvoId, - outcome: 'aborted', - startedAsNewConvo: optimisticStreamIdsRef.current.has(currentStreamId), - endedAt: Date.now(), - }); + // auto-sending queued messages the user may not want fired. A run this + // pane left and came back to is the exception: the history just read + // names its response by epoch, which is the outcome it would have had. + // History read before the job finished may not hold that response yet, + // so an unresolved marker outlives this end for the next visit to read. + const detachedRun = jotaiStore.get(detachedRunByConvoId(recoveryConvoId)); + const detachedResolvable = + detachedRun != null && + !isDetachedRunManual(jotaiStore.get, recoveryConvoId, detachedRun); + const detachedEnd = + detachedResolvable && persistedMessages != null + ? resolveDetachedRunEnd(recoveryConvoId, detachedRun, persistedMessages) + : null; + setRunEnd( + { + conversationId: recoveryConvoId, + outcome: 'aborted', + endedAt: Date.now(), + ...detachedEnd, + startedAsNewConvo: optimisticStreamIdsRef.current.has(currentStreamId), + }, + { retainDetached: detachedResolvable && detachedEnd == null }, + ); setStreamId(null); optimisticStreamIdsRef.current.delete(currentStreamId); createdStreamIdsRef.current.delete(currentStreamId); @@ -4294,6 +4347,26 @@ export default function useResumableSSE( }); submissionRef.current = submission; + runEndPublishedRef.current = false; + if (submission.conversation?.conversationId != null) { + /** A different run starting here owns the conversation's end from now on: an earlier run + * this pane left is superseded, and a Stop belonged to that earlier run. Reattaching to + * the same generation keeps both, since that run's end is still owed and it was stopped. */ + const startingConvoId = submission.conversation.conversationId; + const stopFamily = stopRequestedByConvoId(startingConvoId); + if (jotaiStore.get(stopFamily)?.generationCreatedAt !== resumeGenerationCreatedAt) { + jotaiStore.set(stopFamily, null); + } + const detachedFamily = detachedRunByConvoId(startingConvoId); + const detached = jotaiStore.get(detachedFamily); + if ( + detached != null && + (resumeGenerationCreatedAt == null || + detached.generationCreatedAt !== resumeGenerationCreatedAt) + ) { + jotaiStore.set(detachedFamily, null); + } + } const startController = new AbortController(); const { signal } = startController; const isCurrentEffect = () => !signal.aborted && submissionRef.current === submission; @@ -4823,6 +4896,37 @@ export default function useResumableSSE( stopForegroundReattachRef.current = null; // Reset reconnect counter before closing (so abort handler doesn't think we're reconnecting) reconnectAttemptRef.current = 0; + /** The pane is moving off a run that is still generating (another chat, a new chat, or + * an unmount): the server keeps going with no subscriber and deletes the job when done, + * so no terminal event will reach the queue drain. Remember the run; on return its end is + * read from history. A run that already published its end, or that the user stopped, is + * not remembered. */ + const closing = submissionRef.current; + const closingConvoId = closing?.conversation?.conversationId; + const closingUserMessageId = closing?.userMessage?.messageId; + if ( + streamRef.current != null && + !runEndPublishedRef.current && + closingConvoId != null && + closingConvoId !== Constants.NEW_CONVO && + closingUserMessageId != null + ) { + const liveEpoch = liveEpochRef.current; + const generationCreatedAt = + liveEpoch?.conversationId === closingConvoId + ? liveEpoch.createdAt + : (closing as TSubmission & { resumeGenerationCreatedAt?: number }) + .resumeGenerationCreatedAt; + const stopped = jotaiStore.get(stopRequestedByConvoId(closingConvoId)); + if (!isStoppedRun(stopped, generationCreatedAt)) { + jotaiStore.set(detachedRunByConvoId(closingConvoId), { + userMessageId: closingUserMessageId, + responseMessageId: closing?.initialResponse?.messageId, + ...(generationCreatedAt != null && { generationCreatedAt }), + ...(closing?.isRegenerate === true && { isRegenerate: true }), + }); + } + } streamRef.current?.abort(); streamRef.current = null; // Clear handler maps to prevent memory leaks and stale state diff --git a/client/src/hooks/SSE/useResumeOnLoad.ts b/client/src/hooks/SSE/useResumeOnLoad.ts index 01bb0c4c6a6..e08d76681c3 100644 --- a/client/src/hooks/SSE/useResumeOnLoad.ts +++ b/client/src/hooks/SSE/useResumeOnLoad.ts @@ -34,6 +34,12 @@ import { extendActiveJobsGrace, ACTIVE_JOBS_SUCCESSOR_GRACE_MS, } from '~/data-provider'; +import { + isDetachedRunManual, + resolveDetachedRunEnd, + pendingRunEndByConvoId, + detachedRunByConvoId, +} from '~/hooks/Chat/queue'; import { getGenerationProtocolVersion, supportsGenerationProtocolV2, @@ -1002,10 +1008,48 @@ export default function useResumeOnLoad( // `unrecoveredSteers` above — same empty-list reconcile as the resume path. settleAppliedSteerParts(conversationId, getMessages()); restoreSteerChips(conversationId, undefined); + /** A run this pane stopped watching when the user left has ended on the server, which + * deleted the job. Its persisted response tells how it ended; parking that end lets the + * queue drain send a follow-up queued during the run, as an attached run would have. + * The cached history still holds the optimistic response the pane left, which lacks the + * persisted stop and error flags, so only a successful refetch of history may resolve it; + * when that fails or predates the saved response, the next visit tries again. */ + const detachedFamily = detachedRunByConvoId(conversationId); + const detachedRun = jotaiStore.get(detachedFamily); + const parkDetachedEnd = () => { + if (detachedRun == null || jotaiStore.get(detachedFamily) !== detachedRun) { + return; + } + if (isDetachedRunManual(jotaiStore.get, conversationId, detachedRun)) { + jotaiStore.set(detachedFamily, null); + return; + } + const end = resolveDetachedRunEnd(conversationId, detachedRun, getMessages()); + if (end == null) { + return; + } + jotaiStore.set(detachedFamily, null); + jotaiStore.set(pendingRunEndByConvoId(conversationId), end); + }; + /** A run left for a manual send is retired without a read instead of refetching on every + * visit; the check repeats after the read, since a receipt can settle meanwhile. */ + if (detachedRun != null && isDetachedRunManual(jotaiStore.get, conversationId, detachedRun)) { + jotaiStore.set(detachedFamily, null); + } else if (detachedRun != null) { + queryClient + .refetchQueries( + { queryKey: [QueryKeys.messages, conversationId], exact: true }, + { throwOnError: true }, + ) + .then(parkDetachedEnd, () => undefined); + } processedConvoRef.current = conversationId; return; } + /** Still generating: a marker left for this run stays until the resume's stream delivers its + * end or its recovery reads it from history, and the resume's start retires a marker left + * for any other run, so the job deleted between this read and the attach is not lost. */ processedConvoRef.current = conversationId; if (handoffGenerationKey != null) { consumedHandoffGenerationRef.current = handoffGenerationKey; diff --git a/e2e/specs/mock/scenarios/queue-owners.spec.ts b/e2e/specs/mock/scenarios/queue-owners.spec.ts index 825c466af5d..ed8fe205bd0 100644 --- a/e2e/specs/mock/scenarios/queue-owners.spec.ts +++ b/e2e/specs/mock/scenarios/queue-owners.spec.ts @@ -1,8 +1,12 @@ import { expect, test } from '@playwright/test'; import type { Page } from '@playwright/test'; +import type { TMessage } from 'librechat-data-provider'; import { selectMockEndpoint, + getAccessToken, messagesView, + requestJson, + fetchJson, sendMessage, replyPrompt, replyText, @@ -10,10 +14,14 @@ import { NEW_CHAT_PATH, } from '../helpers'; -/** Interrupt preserves the response instead of draining a separate follow-up. */ +/** + * Interrupt preserves the response instead of draining a separate follow-up, and a follow-up + * queued in a chat the user left drains once the run ends while the chat is off screen. + */ const messageInput = (page: Page) => page.getByRole('textbox', { name: 'Message input' }); const duringRunSendButton = (page: Page) => page.getByTestId('during-run-send-button'); +const queuedRows = (page: Page) => page.getByTestId('queued-message-row'); const messageTurns = (page: Page) => messagesView(page).locator('.message-render'); const uniqueLabel = (prefix: string) => `${prefix}-${Date.now()}-${Math.floor(Math.random() * 1e6)}`; @@ -33,6 +41,29 @@ async function typeDuringRun(page: Page, text: string) { await expect(duringRunSendButton(page)).toBeVisible({ timeout: 5000 }); } +/** Waits until the server has saved the slow run's complete reply. */ +async function waitForServerToFinish(page: Page, conversationId: string) { + const token = await getAccessToken(page); + await expect + .poll( + async () => { + const messages = await fetchJson( + page, + `/api/messages/${encodeURIComponent(conversationId)}`, + token, + ); + return messages.some( + (message) => + !message.isCreatedByUser && + message.unfinished !== true && + JSON.stringify(message.content ?? message.text ?? '').includes('chunk-159'), + ); + }, + { timeout: 60000 }, + ) + .toBe(true); +} + test.describe('chat-owned queue state', () => { test.beforeEach(async ({ page }) => { /** Steer is the plain-Enter default here, so Cmd/Ctrl+Enter is the queue path. */ @@ -65,4 +96,120 @@ test.describe('chat-owned queue state', () => { await expect(messagesView(page).getByText('chunk-010')).toBeVisible(); await expect(messagesView(page).getByText('chunk-159')).toHaveCount(0); }); + + test('a follow-up queued in a chat the user left sends on return @scenario:parked-run-end-drains-on-return', async ({ + page, + }) => { + test.setTimeout(150000); + const label = uniqueLabel('parked'); + const followUp = `Parked follow-up ${label}`; + + await page.goto(NEW_CHAT_PATH, { timeout: 10000 }); + await selectMockEndpoint(page, MOCK_ENDPOINTS[0]); + const conversationId = await establishConversation(page, `parked-setup-${label}`); + + const run = await sendMessage(page, `E2E_SLOW_REPLY:${label}`); + expect(run.ok()).toBeTruthy(); + await typeDuringRun(page, followUp); + await messageInput(page).press('ControlOrMeta+Enter'); + await expect(queuedRows(page).filter({ hasText: followUp })).toBeVisible({ timeout: 10000 }); + + /** Leave through the router, not a reload, so the in-memory queue survives the visit. */ + await page.evaluate((path) => { + window.history.pushState({}, '', path); + window.dispatchEvent(new PopStateEvent('popstate')); + }, NEW_CHAT_PATH); + await expect(page).toHaveURL(/\/c\/new$/); + + /** The run finishes while its chat is not on screen. */ + await waitForServerToFinish(page, conversationId); + await expect(messagesView(page).getByText(followUp)).toHaveCount(0); + + await page.goBack(); + await expect(page).toHaveURL(new RegExp(`/c/${conversationId}(\\?.*)?$`)); + await expect( + messagesView(page).locator('.user-turn').filter({ hasText: followUp }), + ).toBeVisible({ timeout: 30000 }); + await expect(queuedRows(page).filter({ hasText: followUp })).toHaveCount(0); + }); + + test('a follow-up queued before the user stopped and left the run stays queued on return @scenario:stopped-run-left-keeps-follow-up-queued', async ({ + page, + }) => { + test.setTimeout(150000); + const label = uniqueLabel('stopped'); + const followUp = `Stopped follow-up ${label}`; + + await page.goto(NEW_CHAT_PATH, { timeout: 10000 }); + await selectMockEndpoint(page, MOCK_ENDPOINTS[0]); + const conversationId = await establishConversation(page, `stopped-setup-${label}`); + + const run = await sendMessage(page, `E2E_SLOW_REPLY:${label}`); + expect(run.ok()).toBeTruthy(); + await expect(messagesView(page).getByText('chunk-010')).toBeVisible({ timeout: 15000 }); + await typeDuringRun(page, followUp); + await messageInput(page).press('ControlOrMeta+Enter'); + await expect(queuedRows(page).filter({ hasText: followUp })).toBeVisible({ timeout: 10000 }); + + const stop = page.getByRole('button', { name: 'Stop generating' }); + await stop.click(); + await expect(stop).toBeHidden({ timeout: 15000 }); + + await page.evaluate((path) => { + window.history.pushState({}, '', path); + window.dispatchEvent(new PopStateEvent('popstate')); + }, NEW_CHAT_PATH); + await expect(page).toHaveURL(/\/c\/new$/); + + await page.goBack(); + await expect(page).toHaveURL(new RegExp(`/c/${conversationId}(\\?.*)?$`)); + await expect(queuedRows(page).filter({ hasText: followUp })).toBeVisible({ timeout: 15000 }); + await expect( + messagesView(page).locator('.user-turn').filter({ hasText: followUp }), + ).toHaveCount(0); + }); + + test('a follow-up queued in a chat the user left for a saved chat sends on return @scenario:parked-run-end-drains-after-switching-chats', async ({ + page, + }) => { + const width = page.viewportSize()?.width ?? 0; + test.skip(width < 768, 'the conversation list is in the drawer below md'); + test.setTimeout(150000); + const label = uniqueLabel('switched'); + const followUp = `Switched follow-up ${label}`; + const otherTitle = `Other chat ${label}`; + + /** A saved chat to switch to, titled so its list row can be found. */ + await page.goto(NEW_CHAT_PATH, { timeout: 10000 }); + await selectMockEndpoint(page, MOCK_ENDPOINTS[0]); + const otherId = await establishConversation(page, `switched-other-${label}`); + await requestJson(page, { + path: '/api/convos/update', + token: await getAccessToken(page), + method: 'POST', + body: { arg: { conversationId: otherId, title: otherTitle } }, + }); + + await page.goto(NEW_CHAT_PATH, { timeout: 10000 }); + await selectMockEndpoint(page, MOCK_ENDPOINTS[0]); + const conversationId = await establishConversation(page, `switched-setup-${label}`); + const run = await sendMessage(page, `E2E_SLOW_REPLY:${label}`); + expect(run.ok()).toBeTruthy(); + await typeDuringRun(page, followUp); + await messageInput(page).press('ControlOrMeta+Enter'); + await expect(queuedRows(page).filter({ hasText: followUp })).toBeVisible({ timeout: 10000 }); + + /** Leave the way a user does: pick the other saved chat in the conversation list. */ + await page.getByTestId('convo-item').filter({ hasText: otherTitle }).click(); + await expect(page).toHaveURL(new RegExp(`/c/${otherId}(\\?.*)?$`)); + + await waitForServerToFinish(page, conversationId); + + await page.goBack(); + await expect(page).toHaveURL(new RegExp(`/c/${conversationId}(\\?.*)?$`)); + await expect( + messagesView(page).locator('.user-turn').filter({ hasText: followUp }), + ).toBeVisible({ timeout: 30000 }); + await expect(queuedRows(page).filter({ hasText: followUp })).toHaveCount(0); + }); });