diff --git a/client/src/components/Chat/Input/Composer/Queue.tsx b/client/src/components/Chat/Input/Composer/Queue.tsx index 36eb7b71d4b..569f63af1f6 100644 --- a/client/src/components/Chat/Input/Composer/Queue.tsx +++ b/client/src/components/Chat/Input/Composer/Queue.tsx @@ -18,6 +18,7 @@ import { claimQueuedIntent, releaseQueuedIntent, hasQueuedIntent } from '~/utils import { useQueuedTurnPortal } from '~/components/Chat/Steering/QueuedTurnPortal'; import { escalatingSteerFamily, revealedQueuedTurnFamily } from '~/store/steer'; import EscalateNowButton from '~/components/Chat/Input/EscalateNowButton'; +import { queuedMessagesByConvoId } from '~/store/queue'; import { useLocalize } from '~/hooks'; import { cn } from '~/utils'; import store from '~/store'; @@ -615,7 +616,7 @@ function Queue({ shared constant id would duplicate the element and point every handle's `aria-describedby` at whichever copy the document happened to keep. */ const reorderHintId = useId(); - const queued = useRecoilValue(store.queuedMessagesByConvoId(steering.queueKey)); + const queued = useAtomValue(queuedMessagesByConvoId(steering.queueKey)); const pendingSteers = useRecoilValue(store.pendingSteersByConvoId(conversationId)); const escalating = useAtomValue(escalatingSteerFamily(conversationId)); /* Only one interrupt can be unresolved at a time: a second arm would seal the diff --git a/client/src/components/Chat/Input/Composer/__tests__/Queue.spec.tsx b/client/src/components/Chat/Input/Composer/__tests__/Queue.spec.tsx index 0b528e87d14..9075f2f32e8 100644 --- a/client/src/components/Chat/Input/Composer/__tests__/Queue.spec.tsx +++ b/client/src/components/Chat/Input/Composer/__tests__/Queue.spec.tsx @@ -1,7 +1,7 @@ import React from 'react'; +import { RecoilRoot } from 'recoil'; import { DndProvider } from 'react-dnd'; -import { getDefaultStore } from 'jotai'; -import { RecoilRoot, useSetRecoilState } from 'recoil'; +import { getDefaultStore, useSetAtom } from 'jotai'; import { HTML5Backend } from 'react-dnd-html5-backend'; import { ReasoningEffort } from 'librechat-data-provider'; import { act, render, screen, within, fireEvent } from '@testing-library/react'; @@ -11,10 +11,10 @@ import { QueuedTurnPortalProvider, useQueuedTurnPortal, } from '~/components/Chat/Steering/QueuedTurnPortal'; +import { queuedMessagesByConvoId, resetQueueFamilies } from '~/store/queue'; import { hasQueuedIntent, releaseQueuedIntent } from '~/utils/queueIntent'; import { revealedQueuedTurnFamily } from '~/store/steer'; import Queue from '../Queue'; -import store from '~/store'; jest.mock('~/hooks', () => ({ useLocalize: () => (key: string, options?: Record) => { if (!options) { @@ -140,7 +140,9 @@ function renderQueue( } = {}, ) { return render( - set(store.queuedMessagesByConvoId(CONVO_ID), items)}> + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), items)} + > {/* Mirrors `App`, which mounts the provider around the whole tree. */} @@ -161,6 +163,8 @@ function renderQueue( ); } +beforeEach(() => resetQueueFamilies()); + describe('Queue', () => { beforeEach(() => jest.clearAllMocks()); @@ -620,8 +624,8 @@ describe('Queue', () => { render( - set(store.queuedMessagesByConvoId(CONVO_ID), [ + initializeState={() => + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ queued({ id: 'q1' }), queued({ id: 'q2' }), ]) @@ -635,8 +639,8 @@ describe('Queue', () => { /> - set(store.queuedMessagesByConvoId(CONVO_ID), [ + initializeState={() => + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ queued({ id: 'q3' }), queued({ id: 'q4' }), ]) @@ -669,13 +673,16 @@ describe('Queue', () => { it('forgets its last announcement once the queue empties', () => { let setQueue: (items: QueuedMessage[]) => void = () => undefined; const Driver = () => { - setQueue = useSetRecoilState(store.queuedMessagesByConvoId(CONVO_ID)); + setQueue = useSetAtom(queuedMessagesByConvoId(CONVO_ID)); return null; }; render( - set(store.queuedMessagesByConvoId(CONVO_ID), [queued({ id: 'q1' }), queued({ id: 'q2' })]) + initializeState={() => + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ + queued({ id: 'q1' }), + queued({ id: 'q2' }), + ]) } > diff --git a/client/src/hooks/Chat/__tests__/abort.spec.tsx b/client/src/hooks/Chat/__tests__/abort.spec.tsx index 76094b3e2e6..6d4fd4ee861 100644 --- a/client/src/hooks/Chat/__tests__/abort.spec.tsx +++ b/client/src/hooks/Chat/__tests__/abort.spec.tsx @@ -1,10 +1,12 @@ import React from 'react'; import { Constants } from 'librechat-data-provider'; +import { useAtomValue, useSetAtom, getDefaultStore } from 'jotai'; import { act, renderHook, waitFor } from '@testing-library/react'; 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 '~/store/families'; +import { drainAfterAbortByIndex, runEndByIndex, resetQueueFamilies } from '~/store/queue'; import useChatHelpers from '../useChatHelpers'; import { useAbortCleanup } from '../abort'; import store from '~/store'; @@ -87,6 +89,8 @@ function setup() { return { handles, current }; } +beforeEach(() => resetQueueFamilies()); + describe('useAbortCleanup', () => { it('clears all submissions when the captured submission is still current', async () => { const { handles, current } = setup(); @@ -180,9 +184,9 @@ describe('useChatHelpers stopGenerating (abort steer targeting)', () => { function setupStop(conversationId: string, generationCreatedAt: number | null = 41) { let setDrainAfterAbort: ((value: DrainAfterAbort | false) => void) | undefined; function RunEndProbe() { - observed.runEnd = useRecoilValue(store.runEndByIndex(INDEX)); - observed.drainAfterAbort = useRecoilValue(store.drainAfterAbortByIndex(INDEX)); - setDrainAfterAbort = useSetRecoilState(store.drainAfterAbortByIndex(INDEX)); + observed.runEnd = useAtomValue(runEndByIndex(INDEX)); + observed.drainAfterAbort = useAtomValue(drainAfterAbortByIndex(INDEX)); + setDrainAfterAbort = useSetAtom(drainAfterAbortByIndex(INDEX)); return null; } const queryClient = new QueryClient(); @@ -195,7 +199,7 @@ describe('useChatHelpers stopGenerating (abort steer targeting)', () => { endpoint: 'agents', } as TConversation); // Arm interrupt & send so the abort response writes the drain signal. - set(store.drainAfterAbortByIndex(INDEX), { + getDefaultStore().set(drainAfterAbortByIndex(INDEX), { conversationId, generationCreatedAt: 41, }); diff --git a/client/src/hooks/Chat/__tests__/transport.spec.tsx b/client/src/hooks/Chat/__tests__/transport.spec.tsx index 5ca288581e9..cb031fd8afd 100644 --- a/client/src/hooks/Chat/__tests__/transport.spec.tsx +++ b/client/src/hooks/Chat/__tests__/transport.spec.tsx @@ -1,4 +1,5 @@ import React from 'react'; +import { useAtomValue } from 'jotai'; import { MemoryRouter } from 'react-router-dom'; import { RecoilRoot, useRecoilValue } from 'recoil'; import { QueryKeys } from 'librechat-data-provider'; @@ -14,6 +15,7 @@ import type { import type { MutableSnapshot } from 'recoil'; import type { Transport } from '~/hooks/Chat/contract'; import type { PendingSteer } from '~/store/families'; +import { queuedMessagesByConvoId, resetQueueFamilies } from '~/store/queue'; import { ChatTransportContext } from '~/Providers/ChatTransportContext'; import { useSteerReclaim } from '~/hooks/Chat/useSteerCancel'; import useSteerEscalate from '~/hooks/Chat/useSteerEscalate'; @@ -172,11 +174,13 @@ const renderSteering = (transport: Transport) => sendNow: jest.fn(), stopGenerating: jest.fn(), }), - queue: useRecoilValue(store.queuedMessagesByConvoId('convo-1')), + queue: useAtomValue(queuedMessagesByConvoId('convo-1')), }), { wrapper: createWrapper(transport, seedSteerableRun, seedLiveBranch) }, ); +beforeEach(() => resetQueueFamilies()); + describe('chat transport boundary', () => { describe('send (agents)', () => { it('starts the turn and attaches to its stream through the host transport', async () => { diff --git a/client/src/hooks/Chat/__tests__/useQueueDrain.spec.tsx b/client/src/hooks/Chat/__tests__/useQueueDrain.spec.tsx index 21974dd80c2..a252922d6ce 100644 --- a/client/src/hooks/Chat/__tests__/useQueueDrain.spec.tsx +++ b/client/src/hooks/Chat/__tests__/useQueueDrain.spec.tsx @@ -1,15 +1,23 @@ import React from 'react'; -import { getDefaultStore } from 'jotai'; +import { getDefaultStore, useAtomValue, useSetAtom } from 'jotai'; import { act, renderHook, waitFor } from '@testing-library/react'; import { Constants, ReasoningEffort } from 'librechat-data-provider'; import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; -import { RecoilRoot, useRecoilValue, useSetRecoilState, type MutableSnapshot } from 'recoil'; +import { RecoilRoot, useSetRecoilState, type MutableSnapshot } from 'recoil'; import type { DrainAfterAbort, RunEnd, QueuedMessage, SettledQueuedTurnReceipt, } from '~/store/families'; +import { + settledQueuedTurnReceiptsByConvoId, + queuedMessagesByConvoId, + pendingRunEndByConvoId, + drainAfterAbortByIndex, + runEndByIndex, + resetQueueFamilies, +} from '~/store/queue'; import { claimQueuedIntent, releaseQueuedIntent, @@ -51,23 +59,19 @@ function setup( } = {}; function Harness() { - setters.setRunEnd = useSetRecoilState(store.runEndByIndex(INDEX)); + setters.setRunEnd = useSetAtom(runEndByIndex(INDEX)); setters.setIsSubmitting = useSetRecoilState(store.isSubmittingFamily(INDEX)); - setters.setQueue = useSetRecoilState(store.queuedMessagesByConvoId(CONVO_ID)); + setters.setQueue = useSetAtom(queuedMessagesByConvoId(CONVO_ID)); setters.queueRef = { - current: useRecoilValue(store.queuedMessagesByConvoId(CONVO_ID)), + current: useAtomValue(queuedMessagesByConvoId(CONVO_ID)), }; - setters.setNewConvoQueue = useSetRecoilState( - store.queuedMessagesByConvoId(Constants.NEW_CONVO), - ); - setters.setSettledReceipts = useSetRecoilState( - store.settledQueuedTurnReceiptsByConvoId(CONVO_ID), - ); - setters.setInterruptFlag = useSetRecoilState(store.drainAfterAbortByIndex(INDEX)); - setters.queue = useRecoilValue(store.queuedMessagesByConvoId(CONVO_ID)); - setters.newConvoQueue = useRecoilValue(store.queuedMessagesByConvoId(Constants.NEW_CONVO)); - setters.settledReceipts = useRecoilValue(store.settledQueuedTurnReceiptsByConvoId(CONVO_ID)); - setters.runEnd = useRecoilValue(store.runEndByIndex(INDEX)); + setters.setNewConvoQueue = useSetAtom(queuedMessagesByConvoId(Constants.NEW_CONVO)); + setters.setSettledReceipts = useSetAtom(settledQueuedTurnReceiptsByConvoId(CONVO_ID)); + setters.setInterruptFlag = useSetAtom(drainAfterAbortByIndex(INDEX)); + setters.queue = useAtomValue(queuedMessagesByConvoId(CONVO_ID)); + setters.newConvoQueue = useAtomValue(queuedMessagesByConvoId(Constants.NEW_CONVO)); + setters.settledReceipts = useAtomValue(settledQueuedTurnReceiptsByConvoId(CONVO_ID)); + setters.runEnd = useAtomValue(runEndByIndex(INDEX)); useQueueDrain(INDEX, activeConversationId, ask, revealQueuedTurn); return null; } @@ -111,6 +115,8 @@ const runEnd = (overrides: Partial = {}): RunEnd => ({ ...overrides, }); +beforeEach(() => resetQueueFamilies()); + describe('useQueueDrain', () => { beforeEach(() => { getDefaultStore().set(recoveryDispositionsFamily(CONVO_ID), {}); @@ -127,8 +133,8 @@ describe('useQueueDrain', () => { }; const next: QueuedMessage = { id: 'ordinary', text: 'send after review', createdAt: 2 }; const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), [held]); - set(store.queuedMessagesByConvoId(CONVO_ID), [next]); + getDefaultStore().set(queuedMessagesByConvoId(Constants.NEW_CONVO), [held]); + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [next]); set(store.isSubmittingFamily(INDEX), false); }); const end = runEnd({ startedAsNewConvo: true }); @@ -155,7 +161,7 @@ describe('useQueueDrain', () => { }; const next: QueuedMessage = { id: 'ordinary', text: 'wait for next run', createdAt: 2 }; const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [held, next]); + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [held, next]); set(store.isSubmittingFamily(INDEX), false); }); act(() => setters.setRunEnd!(runEnd({ outcome: 'aborted' }))); @@ -178,9 +184,9 @@ describe('useQueueDrain', () => { }; const firstEnd = runEnd(); const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [item]); + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [item]); set(store.isSubmittingFamily(INDEX), false); - set(store.runEndByIndex(INDEX), firstEnd); + getDefaultStore().set(runEndByIndex(INDEX), firstEnd); }); await waitFor(() => expect(setters.runEnd).toEqual(firstEnd)); act(() => setters.setRunEnd?.(runEnd({ generationCreatedAt: 42 }))); @@ -191,8 +197,8 @@ describe('useQueueDrain', () => { ); it('drains exactly one queued message on clean completion', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ queuedMessage('q1', 'first follow-up'), queuedMessage('q2', 'second follow-up'), ]); @@ -207,8 +213,8 @@ describe('useQueueDrain', () => { }); it('leaves a rejected steer for an explicit send and drains the row behind it', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { ...queuedMessage('q-rejected', 'the server refused this'), needsExplicitSend: true }, queuedMessage('q2', 'second follow-up'), ]); @@ -223,8 +229,8 @@ describe('useQueueDrain', () => { }); it('sends nothing when every queued row awaits an explicit send', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { ...queuedMessage('q-rejected', 'the server refused this'), needsExplicitSend: true }, ]); }); @@ -238,8 +244,8 @@ describe('useQueueDrain', () => { }); it('does not locally drain or renew server-owned Agent rows', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { ...queuedMessage('q-server', 'the server starts this turn'), files: [{ file_id: 'server-held-file' }], @@ -268,8 +274,8 @@ describe('useQueueDrain', () => { server: { id: 'server-queue-1', status: 'queued' as const, revision: 1 }, }; const { ask, setters } = setup( - ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [head]); + () => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [head]); }, CONVO_ID, reveal, @@ -300,8 +306,8 @@ describe('useQueueDrain', () => { server: { id: 'server-queue-2', status: 'queued' as const, revision: 2 }, }; const { setters } = setup( - ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [first, second]); + () => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [first, second]); }, CONVO_ID, reveal, @@ -343,9 +349,9 @@ describe('useQueueDrain', () => { server: { id: 'server-queue-2', status: 'queued' as const, revision: 1 }, }; const { ask, setters } = setup( - ({ set }) => { - set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), [local]); - set(store.queuedMessagesByConvoId(CONVO_ID), [serverRow]); + () => { + getDefaultStore().set(queuedMessagesByConvoId(Constants.NEW_CONVO), [local]); + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [serverRow]); }, CONVO_ID, reveal, @@ -363,8 +369,8 @@ describe('useQueueDrain', () => { it('reveals nothing on a stop, or when the completed run carries no response id', async () => { const reveal = jest.fn(); const { setters } = setup( - ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + () => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { ...queuedMessage('q-server', 'the server starts this turn'), clientRequestId: 'client-request-1', @@ -398,8 +404,8 @@ describe('useQueueDrain', () => { }, localSuccessor, ]; - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), queue); + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), queue); }); act(() => { @@ -424,8 +430,8 @@ describe('useQueueDrain', () => { ...queuedMessage('q-server-2', 'later server-owned turn'), server: { id: 'server-queue-2', status: 'queued', revision: 2 }, }; - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [admittedServer, remainingServer]); + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [admittedServer, remainingServer]); }); act(() => { @@ -452,11 +458,11 @@ describe('useQueueDrain', () => { }); it('discards a predecessor boundary already consumed by server admission', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ queuedMessage('q-local', 'must wait for the admitted run'), ]); - set(store.settledQueuedTurnReceiptsByConvoId(CONVO_ID), [ + getDefaultStore().set(settledQueuedTurnReceiptsByConvoId(CONVO_ID), [ { clientRequestId: 'admission-1', status: 'admitted', @@ -475,11 +481,11 @@ describe('useQueueDrain', () => { }); it('consumes one admission when separate receipts share the same predecessor epoch', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ queuedMessage('q-local', 'wait for both admitted runs'), ]); - set(store.settledQueuedTurnReceiptsByConvoId(CONVO_ID), [ + getDefaultStore().set(settledQueuedTurnReceiptsByConvoId(CONVO_ID), [ { clientRequestId: 'admission-1', status: 'admitted', @@ -519,10 +525,10 @@ describe('useQueueDrain', () => { it('migrates the NEW_CONVO queue before discarding a consumed predecessor boundary', async () => { const queuedBeforeResolution = queuedMessage('q-new', 'wait for the admitted successor'); const queuedAfterResolution = queuedMessage('q-resolved', 'still ordered after migration'); - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), [queuedBeforeResolution]); - set(store.queuedMessagesByConvoId(CONVO_ID), [queuedAfterResolution]); - set(store.settledQueuedTurnReceiptsByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(Constants.NEW_CONVO), [queuedBeforeResolution]); + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [queuedAfterResolution]); + getDefaultStore().set(settledQueuedTurnReceiptsByConvoId(CONVO_ID), [ { clientRequestId: 'admission-1', status: 'admitted', @@ -547,11 +553,11 @@ describe('useQueueDrain', () => { }); it('lets the admitted successor terminal boundary release the next turn', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ queuedMessage('q-local', 'send after the admitted run'), ]); - set(store.settledQueuedTurnReceiptsByConvoId(CONVO_ID), [ + getDefaultStore().set(settledQueuedTurnReceiptsByConvoId(CONVO_ID), [ { clientRequestId: 'admission-1', status: 'admitted', @@ -569,11 +575,11 @@ describe('useQueueDrain', () => { }); it('does not interpret a consumed predecessor as a timestamp range', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ queuedMessage('q-local', 'send after the lower-clock successor'), ]); - set(store.settledQueuedTurnReceiptsByConvoId(CONVO_ID), [ + getDefaultStore().set(settledQueuedTurnReceiptsByConvoId(CONVO_ID), [ { clientRequestId: 'admission-1', status: 'admitted', @@ -601,8 +607,8 @@ describe('useQueueDrain', () => { }); it('parks a mismatched signal instead of draining into the wrong conversation', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [queuedMessage('q1', 'stay put')]); + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [queuedMessage('q1', 'stay put')]); }, 'some-other-convo'); act(() => { @@ -616,9 +622,11 @@ describe('useQueueDrain', () => { }); it('drains a parked signal when the user returns to that conversation', async () => { - const { ask } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [queuedMessage('q1', 'welcome back')]); - set(store.pendingRunEndByConvoId(CONVO_ID), runEnd()); + const { ask } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ + queuedMessage('q1', 'welcome back'), + ]); + getDefaultStore().set(pendingRunEndByConvoId(CONVO_ID), runEnd()); }); await waitFor(() => expect(ask).toHaveBeenCalledTimes(1)); @@ -627,8 +635,8 @@ describe('useQueueDrain', () => { it('passes a queued message`s attachments through as overrideFiles', async () => { const files = [{ file_id: 'f1', filepath: '/uploads/f1.png', type: 'image/png' }]; - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { ...queuedMessage('q1', 'with media'), files }, ]); }); @@ -653,8 +661,8 @@ describe('useQueueDrain', () => { * pause for approval, so a hold taken once at enqueue would lapse before a * deep queue finishes draining. */ it('renews the TTL hold on attachments that stay queued', async () => { - const { setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { ...queuedMessage('q1', 'first'), files: [{ file_id: 'sent', type: 'image/png' }], @@ -683,8 +691,8 @@ describe('useQueueDrain', () => { }); it('does not renew when nothing with attachments stays queued', async () => { - const { setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { ...queuedMessage('q1', 'only one'), files: [{ file_id: 'sent', type: 'image/png' }], @@ -709,8 +717,8 @@ describe('useQueueDrain', () => { file_id: `${prefix}-${i}`, type: 'image/png', })); - const { setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { ...queuedMessage('q1', 'first'), files: manyFiles('sent', 2) }, { ...queuedMessage('q2', 'second'), files: manyFiles('a', 10) }, { ...queuedMessage('q3', 'third'), files: manyFiles('b', 4) }, @@ -735,8 +743,8 @@ describe('useQueueDrain', () => { /** A refused send puts the item back with its run-end signal already * consumed, so nothing else would touch it before the next drain. */ it('renews a restored item when the send is refused', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { ...queuedMessage('q1', 'refused'), files: [{ file_id: 'restored', type: 'image/png' }], @@ -759,8 +767,8 @@ describe('useQueueDrain', () => { it('renews on a heartbeat while items stay queued', async () => { jest.useFakeTimers(); try { - setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { ...queuedMessage('q1', 'waiting'), files: [{ file_id: 'held', type: 'image/png' }], @@ -789,8 +797,8 @@ describe('useQueueDrain', () => { it('emits no heartbeat when the queue holds no attachments', async () => { jest.useFakeTimers(); try { - setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [queuedMessage('q1', 'no files')]); + setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [queuedMessage('q1', 'no files')]); }); act(() => { @@ -806,14 +814,14 @@ describe('useQueueDrain', () => { * run ends, so renewing only the active id would skip them for its whole * duration. */ it('renews the pre-migration NEW_CONVO queue alongside the active one', async () => { - setup(({ set }) => { - set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), [ + setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(Constants.NEW_CONVO), [ { ...queuedMessage('n1', 'queued pre-migration'), files: [{ file_id: 'pending-migrate' }], }, ]); - set(store.queuedMessagesByConvoId(CONVO_ID), [ + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { ...queuedMessage('q1', 'queued after'), files: [{ file_id: 'already-migrated' }], @@ -828,8 +836,8 @@ describe('useQueueDrain', () => { }); it('does not double-count the queue before migration', async () => { - setup(({ set }) => { - set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), [ + setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(Constants.NEW_CONVO), [ { ...queuedMessage('n1', 'new convo'), files: [{ file_id: 'only-once' }], @@ -843,8 +851,8 @@ describe('useQueueDrain', () => { }); it('passes carried quotes, skills, and reasoning through as overrides', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { ...queuedMessage('q1', 'with context'), quotes: ['quoted excerpt'], @@ -872,8 +880,8 @@ describe('useQueueDrain', () => { }); it('keeps the recovery user row stable while forwarding its per-attempt identity', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { ...queuedMessage('source-steer', 'recover once'), clientRequestId: 'attempt-uuid', @@ -899,8 +907,8 @@ describe('useQueueDrain', () => { }); it('does not drain on user abort or error outcomes', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [queuedMessage('q1', 'kept')]); + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [queuedMessage('q1', 'kept')]); }); act(() => { @@ -915,9 +923,11 @@ describe('useQueueDrain', () => { }); it('drains on abort when the interrupt & send flag is armed, then disarms', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [queuedMessage('q1', 'interrupt text')]); - set(store.drainAfterAbortByIndex(INDEX), { + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ + queuedMessage('q1', 'interrupt text'), + ]); + getDefaultStore().set(drainAfterAbortByIndex(INDEX), { conversationId: CONVO_ID, generationCreatedAt: 41, }); @@ -940,9 +950,11 @@ describe('useQueueDrain', () => { }); it('does not let a different generation consume an interrupt arm', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [queuedMessage('q1', 'epoch owner')]); - set(store.drainAfterAbortByIndex(INDEX), { + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ + queuedMessage('q1', 'epoch owner'), + ]); + getDefaultStore().set(drainAfterAbortByIndex(INDEX), { conversationId: CONVO_ID, generationCreatedAt: 41, }); @@ -976,9 +988,13 @@ describe('useQueueDrain', () => { { set(store.isSubmittingFamily(INDEX), true); - set(store.queuedMessagesByConvoId(CONVO_A), [queuedMessage('qa', 'send for A')]); - set(store.queuedMessagesByConvoId(CONVO_B), [queuedMessage('qb', 'send for B')]); - set(store.drainAfterAbortByIndex(INDEX), { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_A), [ + queuedMessage('qa', 'send for A'), + ]); + getDefaultStore().set(queuedMessagesByConvoId(CONVO_B), [ + queuedMessage('qb', 'send for B'), + ]); + getDefaultStore().set(drainAfterAbortByIndex(INDEX), { conversationId: CONVO_A, generationCreatedAt: 11, }); @@ -990,12 +1006,12 @@ describe('useQueueDrain', () => { ); const { result, rerender } = renderHook( ({ activeConversationId }: { activeConversationId: string }) => { - setRunEnd = useSetRecoilState(store.runEndByIndex(INDEX)); + setRunEnd = useSetAtom(runEndByIndex(INDEX)); setIsSubmitting = useSetRecoilState(store.isSubmittingFamily(INDEX)); useQueueDrain(INDEX, activeConversationId, ask); return { - indexEnd: useRecoilValue(store.runEndByIndex(INDEX)), - parkedA: useRecoilValue(store.pendingRunEndByConvoId(CONVO_A)), + indexEnd: useAtomValue(runEndByIndex(INDEX)), + parkedA: useAtomValue(pendingRunEndByConvoId(CONVO_A)), }; }, { wrapper, initialProps: { activeConversationId: CONVO_B } }, @@ -1051,7 +1067,7 @@ describe('useQueueDrain', () => { it('waits for isSubmitting to flip false before draining', async () => { const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [queuedMessage('q1', 'deferred')]); + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [queuedMessage('q1', 'deferred')]); set(store.isSubmittingFamily(INDEX), true); }); @@ -1068,8 +1084,8 @@ describe('useQueueDrain', () => { }); it('migrates a NEW_CONVO-keyed queue when the run started as a new conversation', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(Constants.NEW_CONVO), [ queuedMessage('q1', 'queued before convo existed'), queuedMessage('q2', 'second'), ]); @@ -1085,12 +1101,12 @@ describe('useQueueDrain', () => { }); it('preserves a manually reordered NEW_CONVO queue during migration', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(Constants.NEW_CONVO), [ { id: 'second', text: 'send this first', createdAt: 2 }, { id: 'first', text: 'send this second', createdAt: 1 }, ]); - set(store.queuedMessagesByConvoId(CONVO_ID), [ + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { id: 'third', text: 'send this third', createdAt: 3 }, ]); }); @@ -1105,11 +1121,11 @@ describe('useQueueDrain', () => { }); it('keeps an interrupt queued after URL resolution ahead of pre-migration follow-ups', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(Constants.NEW_CONVO), [ { id: 'ordinary-before-url', text: 'ordinary follow-up', createdAt: 1 }, ]); - set(store.queuedMessagesByConvoId(CONVO_ID), [ + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { id: 'interrupt-after-url', text: 'interrupt next', @@ -1117,7 +1133,7 @@ describe('useQueueDrain', () => { priority: true, }, ]); - set(store.drainAfterAbortByIndex(INDEX), { + getDefaultStore().set(drainAfterAbortByIndex(INDEX), { conversationId: CONVO_ID, generationCreatedAt: 41, }); @@ -1152,10 +1168,10 @@ describe('useQueueDrain', () => { } = {}; function Harness() { - setters.setRunEnd = useSetRecoilState(store.runEndByIndex(INDEX)); - setters.setInterruptFlag = useSetRecoilState(store.drainAfterAbortByIndex(INDEX)); - state.newConvoQueue = useRecoilValue(store.queuedMessagesByConvoId(Constants.NEW_CONVO)); - state.parkedUnderOptimistic = useRecoilValue(store.pendingRunEndByConvoId(OPTIMISTIC_ID)); + setters.setRunEnd = useSetAtom(runEndByIndex(INDEX)); + setters.setInterruptFlag = useSetAtom(drainAfterAbortByIndex(INDEX)); + state.newConvoQueue = useAtomValue(queuedMessagesByConvoId(Constants.NEW_CONVO)); + state.parkedUnderOptimistic = useAtomValue(pendingRunEndByConvoId(OPTIMISTIC_ID)); useQueueDrain(INDEX, Constants.NEW_CONVO as string, ask); return null; } @@ -1176,8 +1192,8 @@ describe('useQueueDrain', () => { } it('leaves the NEW_CONVO queue in place and parks nothing under the optimistic id', async () => { - const { ask, setters, state } = setupNewConvo(({ set }) => { - set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), [ + const { ask, setters, state } = setupNewConvo(() => { + getDefaultStore().set(queuedMessagesByConvoId(Constants.NEW_CONVO), [ queuedMessage('q1', 'queued during first turn'), ]); }); @@ -1202,15 +1218,15 @@ describe('useQueueDrain', () => { }); it('drains under NEW_CONVO when interrupt & send was armed', async () => { - const { ask } = setupNewConvo(({ set }) => { - set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), [ + const { ask } = setupNewConvo(() => { + getDefaultStore().set(queuedMessagesByConvoId(Constants.NEW_CONVO), [ queuedMessage('q1', 'interrupted first turn'), ]); - set(store.drainAfterAbortByIndex(INDEX), { + getDefaultStore().set(drainAfterAbortByIndex(INDEX), { conversationId: String(Constants.NEW_CONVO), generationCreatedAt: 41, }); - set(store.runEndByIndex(INDEX), { + getDefaultStore().set(runEndByIndex(INDEX), { conversationId: Constants.NEW_CONVO as string, outcome: 'aborted', startedAsNewConvo: false, @@ -1227,8 +1243,8 @@ describe('useQueueDrain', () => { }); it('consumes the run-end signal (no double fire on re-render)', async () => { - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ queuedMessage('q1', 'one'), queuedMessage('q2', 'two'), ]); @@ -1261,8 +1277,8 @@ describe('useQueueDrain', () => { it('skips a claimed row and sends the next one instead', async () => { claimQueuedIntent('q-claimed'); - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ queuedMessage('q-claimed', 'being taken back'), queuedMessage('q-next', 'still wanted'), ]); @@ -1279,8 +1295,8 @@ describe('useQueueDrain', () => { it('restores a refused skipped row after its claimed predecessor', async () => { claimQueuedIntent('q-claimed'); - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ queuedMessage('q-claimed', 'being taken back'), queuedMessage('q-next', 'temporarily refused'), queuedMessage('q-tail', 'still last'), @@ -1304,8 +1320,10 @@ describe('useQueueDrain', () => { it('sends nothing when the only queued row is claimed', async () => { claimQueuedIntent('q-only'); - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [queuedMessage('q-only', 'mine for now')]); + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ + queuedMessage('q-only', 'mine for now'), + ]); }); act(() => { @@ -1320,8 +1338,10 @@ describe('useQueueDrain', () => { now can already have called `ask` for this pane in the current task. */ it('refuses to drain while a queued send holds the pane', async () => { const held = acquireQueueSendLock(String(INDEX)); - const { ask, setters } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [queuedMessage('q-lock', 'after the hold')]); + const { ask, setters } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ + queuedMessage('q-lock', 'after the hold'), + ]); }); act(() => { diff --git a/client/src/hooks/Chat/__tests__/useSteerCancel.spec.tsx b/client/src/hooks/Chat/__tests__/useSteerCancel.spec.tsx index 75d5ada469f..b5b5e7b83a0 100644 --- a/client/src/hooks/Chat/__tests__/useSteerCancel.spec.tsx +++ b/client/src/hooks/Chat/__tests__/useSteerCancel.spec.tsx @@ -1,4 +1,5 @@ import React from 'react'; +import { useAtomValue } from 'jotai'; import { RecoilRoot, useRecoilValue } from 'recoil'; import { act, renderHook } from '@testing-library/react'; import { ReasoningEffort } from 'librechat-data-provider'; @@ -8,6 +9,7 @@ import { ComposerRestoreProvider, useComposerRestoreHost, } from '~/Providers/ComposerRestoreContext'; +import { queuedMessagesByConvoId, resetQueueFamilies } from '~/store/queue'; import useSteerCancel, { useSteerMoveToQueue } from '../useSteerCancel'; import store from '~/store'; @@ -49,12 +51,14 @@ function setup(steer: PendingSteer) { () => ({ moveToQueue: useSteerMoveToQueue(CONVO_ID), chips: useRecoilValue(store.pendingSteersByConvoId(CONVO_ID)), - queue: useRecoilValue(store.queuedMessagesByConvoId(CONVO_ID)), + queue: useAtomValue(queuedMessagesByConvoId(CONVO_ID)), }), { wrapper }, ); } +beforeEach(() => resetQueueFamilies()); + describe('useSteerMoveToQueue', () => { beforeEach(() => { jest.clearAllMocks(); @@ -124,7 +128,7 @@ describe('useSteerMoveToQueue', () => { const { result } = renderHook( () => ({ moveToQueue: useSteerMoveToQueue(CONVO_ID), - queue: useRecoilValue(store.queuedMessagesByConvoId(CONVO_ID)), + queue: useAtomValue(queuedMessagesByConvoId(CONVO_ID)), }), { wrapper }, ); @@ -209,7 +213,7 @@ function setupCancel( () => ({ cancel: useSteerCancel(CONVO_ID), chips: useRecoilValue(store.pendingSteersByConvoId(CONVO_ID)), - queue: useRecoilValue(store.queuedMessagesByConvoId(CONVO_ID)), + queue: useAtomValue(queuedMessagesByConvoId(CONVO_ID)), }), { wrapper }, ); diff --git a/client/src/hooks/Chat/__tests__/useSteerConvert.spec.tsx b/client/src/hooks/Chat/__tests__/useSteerConvert.spec.tsx index 6e9af61c66f..2affd7a43c6 100644 --- a/client/src/hooks/Chat/__tests__/useSteerConvert.spec.tsx +++ b/client/src/hooks/Chat/__tests__/useSteerConvert.spec.tsx @@ -1,8 +1,9 @@ import React from 'react'; -import { getDefaultStore } from 'jotai'; import { act, renderHook } from '@testing-library/react'; -import { RecoilRoot, useRecoilValue, useSetRecoilState, type MutableSnapshot } from 'recoil'; +import { getDefaultStore, useAtomValue, useSetAtom } from 'jotai'; +import { RecoilRoot, useRecoilValue, type MutableSnapshot } from 'recoil'; import { recoveryDispositionsFamily } from '~/components/Chat/Steering/recovery'; +import { queuedMessagesByConvoId, resetQueueFamilies } from '~/store/queue'; import useSteerConvert from '../useSteerConvert'; import store from '~/store'; @@ -34,11 +35,11 @@ function setup(initialize?: (snapshot: MutableSnapshot) => void) { ); return renderHook( () => { - const setQueue = useSetRecoilState(store.queuedMessagesByConvoId(CONVO_ID)); + const setQueue = useSetAtom(queuedMessagesByConvoId(CONVO_ID)); return { convert: useSteerConvert(), chips: useRecoilValue(store.pendingSteersByConvoId(CONVO_ID)), - queue: useRecoilValue(store.queuedMessagesByConvoId(CONVO_ID)), + queue: useAtomValue(queuedMessagesByConvoId(CONVO_ID)), applied: useRecoilValue(store.appliedSteerIdsByConvoId(CONVO_ID)), // Mirrors `useQueueDrain` dequeuing the head item after auto-send. drainQueue: () => setQueue((prev) => prev.slice(1)), @@ -48,6 +49,8 @@ function setup(initialize?: (snapshot: MutableSnapshot) => void) { ); } +beforeEach(() => resetQueueFamilies()); + describe('useSteerConvert', () => { beforeEach(() => { getDefaultStore().set(recoveryDispositionsFamily(CONVO_ID), {}); @@ -141,7 +144,7 @@ describe('useSteerConvert', () => { }; const after = { id: 'queue-after', text: 'still queued', createdAt: 20 }; const { result } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [after]); + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [after]); set(store.pendingSteersByConvoId(CONVO_ID), [ { steerId: 'server-replacement-id', @@ -200,7 +203,9 @@ describe('useSteerConvert', () => { clientRequestId: 'attempt', createdAt: 1, }; - const { result } = setup(({ set }) => set(store.queuedMessagesByConvoId(CONVO_ID), [item])); + const { result } = setup(() => + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [item]), + ); act(() => result.current.convert(CONVO_ID, [{ steerId: 'source', text: 'original words' }], { generationProtocolVersion: 1, @@ -253,7 +258,7 @@ describe('useSteerConvert', () => { }; const after = { id: 'queue-after-race', text: 'after', createdAt: 20 }; const { result } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [after]); + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [after]); set(store.pendingSteersByConvoId(CONVO_ID), [ { steerId: 'local-correlation-id', @@ -336,8 +341,8 @@ describe('useSteerConvert', () => { }); it('keeps interrupt front-inserts ahead of chronologically older steers', () => { - const { result } = setup(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { result } = setup(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { id: 'urgent', text: 'interrupt message', createdAt: 100, priority: true }, ]); }); @@ -357,7 +362,7 @@ describe('useSteerConvert', () => { { steerId: 'srv-late', text: 'converted', status: 'pending' as const, createdAt: 9 }, ]); // As if the user had dragged the newest message to the front. - set(store.queuedMessagesByConvoId(CONVO_ID), [ + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { id: 'm3', text: 'third, promoted', createdAt: 3 }, { id: 'm1', text: 'first', createdAt: 1 }, { id: 'm2', text: 'second', createdAt: 2 }, @@ -374,7 +379,7 @@ describe('useSteerConvert', () => { set(store.pendingSteersByConvoId(CONVO_ID), [ { steerId: 'srv-early', text: 'accepted first', status: 'pending' as const, createdAt: 1 }, ]); - set(store.queuedMessagesByConvoId(CONVO_ID), [ + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { id: 'later', text: 'queued afterwards', createdAt: 5 }, ]); }); diff --git a/client/src/hooks/Chat/__tests__/useSteerRecovery.spec.tsx b/client/src/hooks/Chat/__tests__/useSteerRecovery.spec.tsx index ed033ca567c..d454e0faa4b 100644 --- a/client/src/hooks/Chat/__tests__/useSteerRecovery.spec.tsx +++ b/client/src/hooks/Chat/__tests__/useSteerRecovery.spec.tsx @@ -1,5 +1,5 @@ import React from 'react'; -import { getDefaultStore } from 'jotai'; +import { getDefaultStore, useAtomValue } from 'jotai'; import { act, render, renderHook } from '@testing-library/react'; import { RecoilRoot, useRecoilValue, useSetRecoilState, type MutableSnapshot } from 'recoil'; import type { RewakeDrain } from '~/Providers/ComposerRestoreContext'; @@ -7,6 +7,7 @@ import { ComposerRestoreProvider, useComposerRestoreHost, } from '~/Providers/ComposerRestoreContext'; +import { queuedMessagesByConvoId, resetQueueFamilies } from '~/store/queue'; import { pendingSteerCancelClientIdsFamily } from '~/store/steer'; import useSteerRecovery from '../useSteerRecovery'; import store from '~/store'; @@ -62,7 +63,7 @@ function setup(initialize?: (snapshot: MutableSnapshot) => void, rewake?: Rewake () => ({ recovery: useSteerRecovery(CONVO_ID), chips: useRecoilValue(store.pendingSteersByConvoId(CONVO_ID)), - queue: useRecoilValue(store.queuedMessagesByConvoId(CONVO_ID)), + queue: useAtomValue(queuedMessagesByConvoId(CONVO_ID)), applied: useRecoilValue(store.appliedSteerIdsByConvoId(CONVO_ID)), accepted: useRecoilValue(store.acceptedSteerClientIdsByConvoId(CONVO_ID)), }), @@ -70,6 +71,8 @@ function setup(initialize?: (snapshot: MutableSnapshot) => void, rewake?: Rewake ); } +beforeEach(() => resetQueueFamilies()); + describe('useSteerRecovery', () => { beforeEach(() => { jest.clearAllMocks(); @@ -242,7 +245,7 @@ describe('useSteerRecovery', () => { function Tree() { recovery = useSteerRecovery(CONVO_ID); chips = useRecoilValue(store.pendingSteersByConvoId(CONVO_ID)); - queue = useRecoilValue(store.queuedMessagesByConvoId(CONVO_ID)); + queue = useAtomValue(queuedMessagesByConvoId(CONVO_ID)); accepted = useRecoilValue(store.acceptedSteerClientIdsByConvoId(CONVO_ID)); const setChips = useSetRecoilState(store.pendingSteersByConvoId(CONVO_ID)); /* What the optimistic cancel does to a `sending` chip: hide it and @@ -433,6 +436,8 @@ describe('useSteerRecovery', () => { every retry earns that same 409 and Retry would be dead for good. */ it('also routes to the queue on RUN_PAUSED / RUN_REPLACED / STEER_UNSUPPORTED / STEER_QUEUE_FULL', async () => { for (const code of ['RUN_PAUSED', 'RUN_REPLACED', 'STEER_UNSUPPORTED', 'STEER_QUEUE_FULL']) { + /** Each code gets its own render; the queue lives in the shared default store. */ + resetQueueFamilies(); mockMutateAsync.mockRejectedValue({ response: { data: { code } } }); const { result } = setup(({ set }) => { set(store.pendingSteersByConvoId(CONVO_ID), [ @@ -511,7 +516,7 @@ describe('useSteerRecovery', () => { away while the conversation's state stays behind to be read. */ const Observer = () => { chips = useRecoilValue(store.pendingSteersByConvoId(CONVO_ID)); - queue = useRecoilValue(store.queuedMessagesByConvoId(CONVO_ID)); + queue = useAtomValue(queuedMessagesByConvoId(CONVO_ID)); setSubmitting = useSetRecoilState(store.isSubmittingFamily(0)); return null; }; @@ -576,7 +581,7 @@ describe('useSteerRecovery', () => { const Tree = () => { recovery = useSteerRecovery(CONVO_ID); chips = useRecoilValue(store.pendingSteersByConvoId(CONVO_ID)); - queue = useRecoilValue(store.queuedMessagesByConvoId(CONVO_ID)); + queue = useAtomValue(queuedMessagesByConvoId(CONVO_ID)); const setConvo = useSetRecoilState(store.conversationByIndex(0)); navigateAway = () => setConvo({ conversationId: 'a-different-chat' } as never); return null; @@ -978,7 +983,7 @@ describe('useSteerRecovery', () => { set(store.pendingSteersByConvoId(CONVO_ID), [ { steerId: 's-old', text: 'older failed steer', status: 'failed', createdAt: 0 }, ]); - set(store.queuedMessagesByConvoId(CONVO_ID), [ + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { id: 'existing', text: 'first', createdAt: 1 }, ]); }); diff --git a/client/src/hooks/Chat/__tests__/useSteering.spec.tsx b/client/src/hooks/Chat/__tests__/useSteering.spec.tsx index b99d3b24fe9..2ea5afe26da 100644 --- a/client/src/hooks/Chat/__tests__/useSteering.spec.tsx +++ b/client/src/hooks/Chat/__tests__/useSteering.spec.tsx @@ -1,7 +1,13 @@ import React from 'react'; import { act, render, renderHook, waitFor } from '@testing-library/react'; import { RecoilRoot, useRecoilValue, useSetRecoilState, type MutableSnapshot } from 'recoil'; -import { getDefaultStore, Provider as JotaiProvider, createStore, useAtomValue } from 'jotai'; +import { + getDefaultStore, + Provider as JotaiProvider, + createStore, + useAtomValue, + useSetAtom, +} from 'jotai'; import { Constants, ContentTypes, @@ -12,6 +18,15 @@ import { import type { CodeApprovalMode, TConversation, TFile, TMessage } from 'librechat-data-provider'; import type { QueuedMessage } from '~/store/families'; import type { ExtendedFile } from '~/common'; +import { + pendingQueuedTurnEnqueueIdsByConvoId, + settledQueuedTurnReceiptsByConvoId, + queuedMessagesByConvoId, + pendingRunEndByConvoId, + drainAfterAbortByIndex, + runEndByIndex, + resetQueueFamilies, +} from '~/store/queue'; import { getReasoningStateKey, pendingReasoningOverrideFamily, @@ -169,7 +184,7 @@ function setup(params: HookParams = {}, initialize?: (snapshot: MutableSnapshot) } function useQueue(convoId: string) { - return useRecoilValue(store.queuedMessagesByConvoId(convoId)); + return useAtomValue(queuedMessagesByConvoId(convoId)); } /** Puts a message in the queue and then sends it, which is the only order the @@ -185,6 +200,8 @@ function sendFromQueue( current.steering.sendQueuedNow(item); } +beforeEach(() => resetQueueFamilies()); + describe('useSteering', () => { beforeEach(() => { getDefaultStore().set(recoveryDispositionsFamily(CONVO_ID), {}); @@ -219,7 +236,7 @@ describe('useSteering', () => { const wrapper = ({ children }: { children: React.ReactNode }) => ( { - set(store.queuedMessagesByConvoId(CONVO_ID), [item]); + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [item]); set(store.isSubmittingFamily(0), false); })} > @@ -242,7 +259,7 @@ describe('useSteering', () => { return { steering, queue: useQueue(CONVO_ID), - setEnd: useSetRecoilState(store.runEndByIndex(0)), + setEnd: useSetAtom(runEndByIndex(0)), }; }, { wrapper }, @@ -310,7 +327,7 @@ describe('useSteering', () => { const wrapper = ({ children }: { children: React.ReactNode }) => ( { - set(store.queuedMessagesByConvoId(CONVO_ID), [held, first, second]); + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [held, first, second]); set(store.isSubmittingFamily(0), false); })} > @@ -333,8 +350,8 @@ describe('useSteering', () => { return { steering, queue: useQueue(CONVO_ID), - end: useRecoilValue(store.runEndByIndex(0)), - setEnd: useSetRecoilState(store.runEndByIndex(0)), + end: useAtomValue(runEndByIndex(0)), + setEnd: useSetAtom(runEndByIndex(0)), }; }, { wrapper }, @@ -615,8 +632,8 @@ describe('useSteering', () => { stopGenerating: jest.fn(), }), queue: useQueue(CONVO_ID), - settledReceipts: useRecoilValue(store.settledQueuedTurnReceiptsByConvoId(CONVO_ID)), - pendingEnqueueIds: useRecoilValue(store.pendingQueuedTurnEnqueueIdsByConvoId(CONVO_ID)), + settledReceipts: useAtomValue(settledQueuedTurnReceiptsByConvoId(CONVO_ID)), + pendingEnqueueIds: useAtomValue(pendingQueuedTurnEnqueueIdsByConvoId(CONVO_ID)), }), { wrapper }, ); @@ -902,8 +919,8 @@ describe('useSteering', () => { }); }); const { result } = setupServerQueue( - withActiveGeneration(({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [ + withActiveGeneration(() => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { id: 'unrelated-server-row', clientRequestId: 'unrelated-request', @@ -1513,13 +1530,13 @@ describe('useSteering', () => { { set(store.isSubmittingFamily(0), false); - set(store.runEndByIndex(0), { + getDefaultStore().set(runEndByIndex(0), { conversationId: CONVO_ID, outcome: 'completed', endedAt: 200, generationCreatedAt: 41, }); - set(store.queuedMessagesByConvoId(CONVO_ID), [ + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { id: 'client-admitted-race', clientRequestId: 'client-admitted-race', @@ -1549,7 +1566,7 @@ describe('useSteering', () => { useQueueDrain(0, CONVO_ID, ask); return { queue: useQueue(CONVO_ID), - runEnd: useRecoilValue(store.runEndByIndex(0)), + runEnd: useAtomValue(runEndByIndex(0)), }; }, { wrapper }, @@ -1755,10 +1772,10 @@ describe('useSteering', () => { queue: useQueue(CONVO_ID), /** What `useQueueDrain` watches: re-posting it is how this hook asks * the drain to reconsider a queue it already passed over. */ - parkedRunEnd: useRecoilValue(store.pendingRunEndByConvoId(CONVO_ID)), + parkedRunEnd: useAtomValue(pendingRunEndByConvoId(CONVO_ID)), /** Stands in for the drain CONSUMING a signal it has acted on. */ - consumeIndexSignal: useSetRecoilState(store.runEndByIndex(0)), - consumeParkedSignal: useSetRecoilState(store.pendingRunEndByConvoId(CONVO_ID)), + consumeIndexSignal: useSetAtom(runEndByIndex(0)), + consumeParkedSignal: useSetAtom(pendingRunEndByConvoId(CONVO_ID)), }), { wrapper }, ); @@ -1800,8 +1817,8 @@ describe('useSteering', () => { // The drain already consumed its one-shot signal against an empty queue, // so re-post it: the DRAIN sends (FIFO, via `ask`, which does not reset // the composer), never this hook. - const { result, sendNow } = setupWithQueue({ isSubmitting: false }, ({ set }) => { - set(store.runEndByIndex(0), runEnd('completed')); + const { result, sendNow } = setupWithQueue({ isSubmitting: false }, () => { + getDefaultStore().set(runEndByIndex(0), runEnd('completed')); }); act(() => { // The drain ran against an empty queue and consumed the signal; the @@ -1824,8 +1841,8 @@ describe('useSteering', () => { // The run finished with this conversation off-screen, so its signal was // parked rather than delivered on the index. Without watching the parked // carrier too, the outcome would never be seen and the item would strand. - const { result } = setupWithQueue({ isSubmitting: false }, ({ set }) => { - set(store.pendingRunEndByConvoId(CONVO_ID), runEnd('completed')); + const { result } = setupWithQueue({ isSubmitting: false }, () => { + getDefaultStore().set(pendingRunEndByConvoId(CONVO_ID), runEnd('completed')); }); act(() => { result.current.consumeParkedSignal(null); @@ -1844,8 +1861,8 @@ describe('useSteering', () => { (outcome) => { // The drain auto-sends only on a clean completion: a Stop or an error // means the user is taking over, so nothing may smuggle the text out. - const { result, sendNow } = setupWithQueue({ isSubmitting: false }, ({ set }) => { - set(store.runEndByIndex(0), runEnd(outcome)); + const { result, sendNow } = setupWithQueue({ isSubmitting: false }, () => { + getDefaultStore().set(runEndByIndex(0), runEnd(outcome)); }); act(() => { result.current.consumeIndexSignal(null); @@ -1860,8 +1877,8 @@ describe('useSteering', () => { ); it('leaves the item for manual send when the completed run was another chat', () => { - const { result, sendNow } = setupWithQueue({ isSubmitting: false }, ({ set }) => { - set(store.runEndByIndex(0), runEnd('completed', 'convo-elsewhere')); + const { result, sendNow } = setupWithQueue({ isSubmitting: false }, () => { + getDefaultStore().set(runEndByIndex(0), runEnd('completed', 'convo-elsewhere')); }); act(() => { result.current.consumeIndexSignal(null); @@ -1877,9 +1894,9 @@ describe('useSteering', () => { it('keeps older queued follow-ups ahead of the reclaimed steer', () => { // The drain sends ONE item per run end, FIFO. Re-arming (rather than // sending here) is what keeps an older follow-up from being skipped. - const { result } = setupWithQueue({ isSubmitting: false }, ({ set }) => { - set(store.runEndByIndex(0), runEnd('completed')); - set(store.queuedMessagesByConvoId(CONVO_ID), [ + const { result } = setupWithQueue({ isSubmitting: false }, () => { + getDefaultStore().set(runEndByIndex(0), runEnd('completed')); + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [ { id: 'older', text: 'queued first', createdAt: 500 }, ]); }); @@ -1895,9 +1912,9 @@ describe('useSteering', () => { it('does not re-arm while this conversation’s run-end is still unconsumed', () => { // The drain has not run yet, so it will see this item on its own. Arming // a second carrier would drain twice and send two messages. - const { result } = setupWithQueue({ isSubmitting: false }, ({ set }) => { - set(store.pendingRunEndByConvoId(CONVO_ID), runEnd('completed')); - set(store.runEndByIndex(0), runEnd('completed')); + const { result } = setupWithQueue({ isSubmitting: false }, () => { + getDefaultStore().set(pendingRunEndByConvoId(CONVO_ID), runEnd('completed')); + getDefaultStore().set(runEndByIndex(0), runEnd('completed')); }); act(() => { result.current.steering.queueReclaimedSteer(reclaimed); @@ -1931,10 +1948,10 @@ describe('useSteering', () => { sendNow, stopGenerating: jest.fn(), }), - parkedHere: useRecoilValue(store.pendingRunEndByConvoId(CONVO_ID)), + parkedHere: useAtomValue(pendingRunEndByConvoId(CONVO_ID)), queueHere: useQueue(CONVO_ID), /** Stands in for the drain CONSUMING a signal it has acted on. */ - consumeIndexSignal: useSetRecoilState(store.runEndByIndex(0)), + consumeIndexSignal: useSetAtom(runEndByIndex(0)), }), { wrapper, initialProps }, ); @@ -1946,8 +1963,8 @@ describe('useSteering', () => { // they belong to completed, so its queue must still drain on return. const { result, rerender, sendNow } = setupNavigable( { convoId: CONVO_ID, isSubmitting: false }, - ({ set }) => { - set(store.runEndByIndex(0), runEnd('completed')); + () => { + getDefaultStore().set(runEndByIndex(0), runEnd('completed')); }, ); // Captured while still on this chat, resolving after the user left. @@ -2042,8 +2059,8 @@ describe('useSteering', () => { // owns the item, and its own end will drain it. const { result, rerender } = setupNavigable( { convoId: CONVO_ID, isSubmitting: false }, - ({ set }) => { - set(store.runEndByIndex(0), runEnd('completed')); + () => { + getDefaultStore().set(runEndByIndex(0), runEnd('completed')); }, ); act(() => { @@ -2065,8 +2082,8 @@ describe('useSteering', () => { // The index slot is shared. The drain parks a foreign signal under ITS // conversation and then only inspects the active one's queue, so treating // it as proof of an upcoming drain would strand this item. - const { result } = setupWithQueue({ isSubmitting: false }, ({ set }) => { - set(store.runEndByIndex(0), runEnd('completed')); + const { result } = setupWithQueue({ isSubmitting: false }, () => { + getDefaultStore().set(runEndByIndex(0), runEnd('completed')); }); act(() => { result.current.consumeIndexSignal(null); @@ -2098,8 +2115,8 @@ describe('useSteering', () => { generationProtocolVersion: 2, }); }); - const { result } = setupWithQueue({ isSubmitting: false }, ({ set }) => { - set(store.runEndByIndex(0), runEnd('completed')); + const { result } = setupWithQueue({ isSubmitting: false }, () => { + getDefaultStore().set(runEndByIndex(0), runEnd('completed')); }); act(() => { result.current.consumeIndexSignal(null); @@ -2569,7 +2586,7 @@ describe('useSteering', () => { ...params, }), queue: useQueue(CONVO_ID), - setQueue: useSetRecoilState(store.queuedMessagesByConvoId(CONVO_ID)), + setQueue: useSetAtom(queuedMessagesByConvoId(CONVO_ID)), setActiveEpoch: useSetRecoilState(store.activeGenerationCreatedAtByConvoId(CONVO_ID)), setActiveProtocol: useSetRecoilState( store.activeGenerationProtocolVersionByConvoId(CONVO_ID), @@ -2578,7 +2595,7 @@ describe('useSteering', () => { chips: useRecoilValue(store.pendingSteersByConvoId(CONVO_ID)), setChips: useSetRecoilState(store.pendingSteersByConvoId(CONVO_ID)), setAppliedIds: useSetRecoilState(store.appliedSteerIdsByConvoId(CONVO_ID)), - drainFlag: useRecoilValue(store.drainAfterAbortByIndex(0)), + drainFlag: useAtomValue(drainAfterAbortByIndex(0)), }), { wrapper }, ); @@ -2846,9 +2863,13 @@ describe('useSteering', () => { // A draft queued during the 202 round-trip must not drain ahead of a // steer submitted before it, so the ACK'd chip keeps its SUBMIT time, // not the (later) ACK time. - const now = jest.spyOn(Date, 'now').mockReturnValueOnce(1_000).mockReturnValue(9_000); + /** Submit-time reads see 1_000 and everything from the ACK on sees 9_000, whatever else + * reads the clock first. */ + let acknowledged = false; + const now = jest.spyOn(Date, 'now').mockImplementation(() => (acknowledged ? 9_000 : 1_000)); try { mockMutate.mockImplementation((_params, { onSuccess }) => { + acknowledged = true; onSuccess({ steerId: 'srv-t', status: 'queued', @@ -3216,8 +3237,8 @@ describe('useSteering', () => { it('preserves a manual queue order when another message is enqueued', () => { const first: QueuedMessage = { id: 'first', text: 'first', createdAt: 1 }; const second: QueuedMessage = { id: 'second', text: 'second', createdAt: 2 }; - const { result } = setupWithState({}, ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [first, second]); + const { result } = setupWithState({}, () => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [first, second]); }); act(() => { @@ -3262,8 +3283,8 @@ describe('useSteering', () => { text: 'after', createdAt: 2, }; - const { result } = setupWithState({}, ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [before, recovered, after]); + const { result } = setupWithState({}, () => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [before, recovered, after]); }); let discarded = false; @@ -3319,8 +3340,8 @@ describe('useSteering', () => { text: 'after', createdAt: 2, }; - const { result } = setupWithState({}, ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [before, recovered, after]); + const { result } = setupWithState({}, () => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [before, recovered, after]); }); let discarded = false; @@ -3354,8 +3375,8 @@ describe('useSteering', () => { recoverySteerId: 'server-leftover', recoveryClientSteerId: 'client-leftover', }; - const { result } = setupWithState({}, ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [recovered]); + const { result } = setupWithState({}, () => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [recovered]); }); let discarded = true; @@ -3378,8 +3399,8 @@ describe('useSteering', () => { createdAt: 1, recoverySteerId: 'server-leftover', }; - const { result } = setupWithState({}, ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [recovered]); + const { result } = setupWithState({}, () => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [recovered]); }); let discarded = true; @@ -3400,8 +3421,8 @@ describe('useSteering', () => { recoverySteerId: 'source', }; getDefaultStore().set(recoveryDispositionsFamily(CONVO_ID), { source: 'blocked' }); - const { result, sendNow } = setupWithState({ isSubmitting: false }, ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [item]); + const { result, sendNow } = setupWithState({ isSubmitting: false }, () => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [item]); }); act(() => result.current.steering.sendQueuedNow(item)); expect(sendNow).not.toHaveBeenCalled(); @@ -3456,8 +3477,8 @@ describe('useSteering', () => { }; const after = { id: 'q-after', text: 'after', createdAt: 30 }; const originalQueue = [before, selected, after]; - const { result } = setupWithState({}, ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), originalQueue); + const { result } = setupWithState({}, () => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), originalQueue); }); act(() => { @@ -3504,8 +3525,8 @@ describe('useSteering', () => { text: 'queued during request', createdAt: 40, }; - const { result } = setupWithState({}, ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [before, selected, after]); + const { result } = setupWithState({}, () => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [before, selected, after]); }); act(() => { @@ -3537,8 +3558,8 @@ describe('useSteering', () => { text: 'second instruction', createdAt: 20, }; - const { result } = setupWithState({}, ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [first, second]); + const { result } = setupWithState({}, () => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [first, second]); }); act(() => { @@ -3606,8 +3627,8 @@ describe('useSteering', () => { text: 'second accepted', createdAt: 20, }; - const { result } = setupWithState({}, ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [first, second]); + const { result } = setupWithState({}, () => { + getDefaultStore().set(queuedMessagesByConvoId(CONVO_ID), [first, second]); }); act(() => { @@ -3701,7 +3722,7 @@ describe('useSteering', () => { ...params, }), queue: useQueue(CONVO_ID), - setQueue: useSetRecoilState(store.queuedMessagesByConvoId(CONVO_ID)), + setQueue: useSetAtom(queuedMessagesByConvoId(CONVO_ID)), }), { wrapper }, ); @@ -4036,6 +4057,8 @@ describe('useSteering', () => { }); describe('composer quotes + manual skill capture', () => { + /** The harness renders under its own Jotai store, so queue seeds target that store. */ + let contextStore = createStore(); function setupWithContext( params: HookParams = {}, initialize?: (snapshot: MutableSnapshot) => void, @@ -4044,6 +4067,7 @@ describe('useSteering', () => { const sendNow = jest.fn(); const stopGenerating = jest.fn(); const reasoningStore = createStore(); + contextStore = reasoningStore; reasoningStore.set( pendingReasoningOverrideFamily(getReasoningStateKey(params.conversationId ?? CONVO_ID, 0)), initialReasoning, @@ -4071,7 +4095,7 @@ describe('useSteering', () => { ...params, }), queue: useQueue(CONVO_ID), - setQueue: useSetRecoilState(store.queuedMessagesByConvoId(CONVO_ID)), + setQueue: useSetAtom(queuedMessagesByConvoId(CONVO_ID)), chips: useRecoilValue(store.pendingSteersByConvoId(CONVO_ID)), pendingQuotes: useRecoilValue(store.pendingQuotesByConvoId(CONVO_ID)), pendingSkills: useRecoilValue(store.pendingManualSkillsByConvoId(CONVO_ID)), @@ -4222,8 +4246,8 @@ describe('useSteering', () => { createdAt: 1_000, quotes: ['queued excerpt'], }; - const { result } = setupWithContext({}, ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [item]); + const { result } = setupWithContext({}, () => { + contextStore.set(queuedMessagesByConvoId(CONVO_ID), [item]); }); act(() => { result.current.steering.sendQueuedNow(item); @@ -4245,8 +4269,8 @@ describe('useSteering', () => { reasoningOverride: { key: 'reasoning_effort', value: ReasoningEffort.high }, server: { id: 'server-reasoning-live', status: 'queued', revision: 3 }, }; - const { result, sendNow } = setupWithContext({}, ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [item]); + const { result, sendNow } = setupWithContext({}, () => { + contextStore.set(queuedMessagesByConvoId(CONVO_ID), [item]); }); act(() => { @@ -4455,8 +4479,8 @@ describe('useSteering', () => { createdAt: 1_000, quotes: ['queued excerpt'], }; - const { result } = setupWithContext({}, ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [item]); + const { result } = setupWithContext({}, () => { + contextStore.set(queuedMessagesByConvoId(CONVO_ID), [item]); }); act(() => { result.current.steering.sendQueuedNow(item); @@ -4581,8 +4605,8 @@ describe('useSteering', () => { quotes: ['original quote'], }; const after = { id: 'late-after', text: 'after', createdAt: 30 }; - const { result, rerender } = setupWithContext(params, ({ set }) => { - set(store.queuedMessagesByConvoId(CONVO_ID), [before, selected, after]); + const { result, rerender } = setupWithContext(params, () => { + contextStore.set(queuedMessagesByConvoId(CONVO_ID), [before, selected, after]); }); act(() => { diff --git a/client/src/hooks/Chat/useChatHelpers.ts b/client/src/hooks/Chat/useChatHelpers.ts index 41b1d13737a..4f3ae71fa9e 100644 --- a/client/src/hooks/Chat/useChatHelpers.ts +++ b/client/src/hooks/Chat/useChatHelpers.ts @@ -1,8 +1,8 @@ import { useCallback, useMemo, useRef, useState } from 'react'; -import { useSetAtom } from 'jotai'; +import { useSetAtom, useStore } from 'jotai'; import { useQueryClient } from '@tanstack/react-query'; +import { useRecoilState, useRecoilValue, useSetRecoilState } from 'recoil'; import { Constants, QueryKeys, isAssistantsEndpoint } from 'librechat-data-provider'; -import { useRecoilState, useRecoilValue, useSetRecoilState, useRecoilCallback } from 'recoil'; import type { TMessage } from 'librechat-data-provider'; import type { ChatContract } from './contract'; import { @@ -12,6 +12,7 @@ import { } from '~/data-provider'; import { useLatestMessage, useLatestMessageId } from '~/hooks/Messages/useLatestMessage'; import { siblingIdxFamily, siblingKey } from '~/components/Chat/Messages/Thread/state'; +import { drainAfterAbortByIndex, runEndByIndex } from '~/store/queue'; import useChatFunctions from '~/hooks/Chat/useChatFunctions'; import useSteerConvert from '~/hooks/Chat/useSteerConvert'; import { resolveAbortSteerTarget } from '~/utils'; @@ -36,19 +37,19 @@ export default function useChatHelpers(index = 0, paramId?: string): ChatContrac /** Async abort responses can settle after this pane has moved to another * conversation and armed its own interrupt. Clear only the intent owned by * the request that produced the response. */ - const clearInterruptDrain = useRecoilCallback( - ({ snapshot, set }) => - (convoId: string, generationCreatedAt: number) => { - const armed = snapshot.getLoadable(store.drainAfterAbortByIndex(index)).getValue(); - if ( - armed !== false && - armed.conversationId === convoId && - armed.generationCreatedAt === generationCreatedAt - ) { - set(store.drainAfterAbortByIndex(index), false); - } - }, - [index], + const queueStore = useStore(); + const clearInterruptDrain = useCallback( + (convoId: string, generationCreatedAt: number) => { + const armed = queueStore.get(drainAfterAbortByIndex(index)); + if ( + armed !== false && + armed.conversationId === convoId && + armed.generationCreatedAt === generationCreatedAt + ) { + queueStore.set(drainAfterAbortByIndex(index), false); + } + }, + [index, queueStore], ); /** @@ -60,28 +61,27 @@ export default function useChatHelpers(index = 0, paramId?: string): ChatContrac * final DOES arrive later, its signal finds the flag already consumed and * an `aborted` outcome drains nothing, so there is no double fire. */ - const signalInterruptDrain = useRecoilCallback( - ({ snapshot, set }) => - (convoId: string, generationCreatedAt: number, armedConversationId = convoId) => { - const armed = snapshot.getLoadable(store.drainAfterAbortByIndex(index)).getValue(); - const runEnd = snapshot.getLoadable(store.runEndByIndex(index)).getValue(); - const matchesArm = - armed !== false && - armed.conversationId === armedConversationId && - armed.generationCreatedAt === generationCreatedAt; - const alreadySignaled = - runEnd?.conversationId === convoId && runEnd.generationCreatedAt === generationCreatedAt; - if (!matchesArm || alreadySignaled) { - return; - } - set(store.runEndByIndex(index), { - conversationId: convoId, - outcome: 'aborted', - endedAt: Date.now(), - generationCreatedAt, - }); - }, - [index], + const signalInterruptDrain = useCallback( + (convoId: string, generationCreatedAt: number, armedConversationId = convoId) => { + const armed = queueStore.get(drainAfterAbortByIndex(index)); + const runEnd = queueStore.get(runEndByIndex(index)); + const matchesArm = + armed !== false && + armed.conversationId === armedConversationId && + armed.generationCreatedAt === generationCreatedAt; + const alreadySignaled = + runEnd?.conversationId === convoId && runEnd.generationCreatedAt === generationCreatedAt; + if (!matchesArm || alreadySignaled) { + return; + } + queueStore.set(runEndByIndex(index), { + conversationId: convoId, + outcome: 'aborted', + endedAt: Date.now(), + generationCreatedAt, + }); + }, + [index, queueStore], ); const { newConversation } = useNewConvo(index); diff --git a/client/src/hooks/Chat/useQueueDrain.ts b/client/src/hooks/Chat/useQueueDrain.ts index 6e5aad6a1af..7080607fe16 100644 --- a/client/src/hooks/Chat/useQueueDrain.ts +++ b/client/src/hooks/Chat/useQueueDrain.ts @@ -1,10 +1,18 @@ -import { useRef, useEffect, useMemo } from 'react'; +import { useRef, useEffect, useMemo, useCallback } from 'react'; +import { useRecoilValue } from 'recoil'; import { useAtomValue, useStore } from 'jotai'; -import { useRecoilValue, useRecoilCallback } from 'recoil'; import { Constants, DEFAULT_QUEUED_SEND_LOCK_TIMEOUT_MS } from 'librechat-data-provider'; import type { DrainAfterAbort, QueuedMessage, QueuedMessageOrigin, RunEnd } from '~/store/families'; import type { QueueSendLock } from '~/utils/queueIntent'; import type { TAskFunction } from '~/common'; +import { + settledQueuedTurnReceiptsByConvoId, + pendingQueuedTurnEnqueueIdsByConvoId, + pendingRunEndByConvoId, + queuedMessagesByConvoId, + drainAfterAbortByIndex, + runEndByIndex, +} from '~/store/queue'; import { recoveryDispositionsFamily, recoveryDisposition, @@ -87,9 +95,9 @@ export default function useQueueDrain( revealQueuedTurn?: (item: QueuedMessage, end: RunEnd) => void, ) { const jotaiStore = useStore(); - const runEnd = useRecoilValue(store.runEndByIndex(index)); - const parkedRunEnd = useRecoilValue( - store.pendingRunEndByConvoId(activeConversationId ?? Constants.NEW_CONVO), + const runEnd = useAtomValue(runEndByIndex(index)); + const parkedRunEnd = useAtomValue( + pendingRunEndByConvoId(activeConversationId ?? Constants.NEW_CONVO), ); const isSubmitting = useRecoilValue(store.isSubmittingFamily(index)); /** Keyed by pane, because the contended resource is this pane's submission @@ -101,19 +109,19 @@ export default function useQueueDrain( const { data: startupConfig } = useGetStartupConfig(); const sendLockTimeoutMs = startupConfig?.interface?.queuedSendLockTimeoutMs ?? DEFAULT_QUEUED_SEND_LOCK_TIMEOUT_MS; - const ownQueue = useRecoilValue( - store.queuedMessagesByConvoId(activeConversationId ?? Constants.NEW_CONVO), + const ownQueue = useAtomValue( + queuedMessagesByConvoId(activeConversationId ?? Constants.NEW_CONVO), ); /** `drainNext` merges this in, and it outlives the URL update: items queued * during the first turn stay keyed here until that run ends. Renewing only * the active id would skip them for the whole of that run. */ - const newConvoQueue = useRecoilValue(store.queuedMessagesByConvoId(Constants.NEW_CONVO)); + const newConvoQueue = useAtomValue(queuedMessagesByConvoId(Constants.NEW_CONVO)); /** Receipt settlement can consume the parked terminal boundary without * changing whether another server-owned row remains. Subscribe here so the * drain effect observes that durable transition instead of reading it only - * through a callback snapshot whose other dependencies stayed unchanged. */ - const settledQueuedTurnReceipts = useRecoilValue( - store.settledQueuedTurnReceiptsByConvoId(activeConversationId ?? Constants.NEW_CONVO), + * through a callback read whose other dependencies stayed unchanged. */ + const settledQueuedTurnReceipts = useAtomValue( + settledQueuedTurnReceiptsByConvoId(activeConversationId ?? Constants.NEW_CONVO), ); const hasServerOwnedQueue = [...ownQueue, ...newConvoQueue].some((item) => item.server != null); // A held head can leave the queue without changing the terminal signal. @@ -176,252 +184,222 @@ export default function useQueueDrain( * waiting for it to become idle lets its own terminal signal replace the * first one. Matching interrupt intent travels with the exact terminal * epoch, while unrelated intent remains armed for its owner. */ - const parkForeignRunEnd = useRecoilCallback( - ({ snapshot, set }) => - (): boolean => { - const end = snapshot.getLoadable(store.runEndByIndex(index)).getValue(); - if ( - end == null || - end.conversationId == null || - end.conversationId === activeConversationId - ) { - return false; - } - const armed = snapshot.getLoadable(store.drainAfterAbortByIndex(index)).getValue(); - const interruptArmed = matchesInterruptArm(armed, end); - if (interruptArmed) { - set(store.drainAfterAbortByIndex(index), false); - } - set(store.pendingRunEndByConvoId(end.conversationId), { - ...end, - ...((end.interruptArmed === true || interruptArmed) && { - interruptArmed: true, - }), - }); - set(store.runEndByIndex(index), null); - return true; - }, - [index, activeConversationId], - ); + const parkForeignRunEnd = useCallback((): boolean => { + const end = jotaiStore.get(runEndByIndex(index)); + if (end == null || end.conversationId == null || end.conversationId === activeConversationId) { + return false; + } + const armed = jotaiStore.get(drainAfterAbortByIndex(index)); + const interruptArmed = matchesInterruptArm(armed, end); + if (interruptArmed) { + jotaiStore.set(drainAfterAbortByIndex(index), false); + } + jotaiStore.set(pendingRunEndByConvoId(end.conversationId), { + ...end, + ...((end.interruptArmed === true || interruptArmed) && { + interruptArmed: true, + }), + }); + jotaiStore.set(runEndByIndex(index), null); + return true; + }, [index, activeConversationId, jotaiStore]); - // Fully synchronous reads (getLoadable): a useRecoilCallback snapshot is - // only guaranteed valid for the callback's synchronous execution, so no - // awaits may interleave with the reads. - const drainNext = useRecoilCallback( - ({ snapshot, set }) => - (): - | { - kind: 'drain'; - next: QueuedMessage; - conversationId: string; - queuedMessageOrigin: QueuedMessageOrigin; - expectedPredecessorCreatedAt?: number; - } - | { kind: 'reveal'; item: QueuedMessage; end: RunEnd } - | null => { - let end = snapshot.getLoadable(store.runEndByIndex(index)).getValue(); - let fromParked = false; - if ( - end != null && - end.conversationId != null && - end.conversationId !== activeConversationId - ) { - /** - * `ask` is the MOUNTED view's sender: draining another - * conversation's follow-up here would submit it into the wrong - * chat. Park the signal under ITS conversation (freeing the shared - * index slot so a later run cannot overwrite it) and drain when - * the user returns. The armed interrupt flag travels WITH the - * parked signal: leaving it on the index would let another run on - * this pane consume it (or drain the wrong conversation). - */ - const armedNow = snapshot.getLoadable(store.drainAfterAbortByIndex(index)).getValue(); - const interruptArmed = matchesInterruptArm(armedNow, end); - if (interruptArmed) { - set(store.drainAfterAbortByIndex(index), false); - } - set(store.pendingRunEndByConvoId(end.conversationId), { - ...end, - ...((end.interruptArmed === true || interruptArmed) && { - interruptArmed: true, - }), - }); - set(store.runEndByIndex(index), null); - return null; - } - if (end == null && activeConversationId) { - const parked = snapshot - .getLoadable(store.pendingRunEndByConvoId(activeConversationId)) - .getValue(); - if (parked != null) { - end = parked; - fromParked = true; - } - } - if (end == null) { - return null; - } - const indexArmed = snapshot.getLoadable(store.drainAfterAbortByIndex(index)).getValue(); - const matchingIndexArm = matchesInterruptArm(indexArmed, end); - const interruptArmed = matchingIndexArm || end.interruptArmed === true; + // Fully synchronous: every read and write below happens in one task, so no + // other writer can interleave between deciding and consuming the signal. + const drainNext = useCallback((): + | { + kind: 'drain'; + next: QueuedMessage; + conversationId: string; + queuedMessageOrigin: QueuedMessageOrigin; + expectedPredecessorCreatedAt?: number; + } + | { kind: 'reveal'; item: QueuedMessage; end: RunEnd } + | null => { + let end = jotaiStore.get(runEndByIndex(index)); + let fromParked = false; + if (end != null && end.conversationId != null && end.conversationId !== activeConversationId) { + /** + * `ask` is the MOUNTED view's sender: draining another + * conversation's follow-up here would submit it into the wrong + * chat. Park the signal under ITS conversation (freeing the shared + * index slot so a later run cannot overwrite it) and drain when + * the user returns. The armed interrupt flag travels WITH the + * parked signal: leaving it on the index would let another run on + * this pane consume it (or drain the wrong conversation). + */ + const armedNow = jotaiStore.get(drainAfterAbortByIndex(index)); + const interruptArmed = matchesInterruptArm(armedNow, end); + if (interruptArmed) { + jotaiStore.set(drainAfterAbortByIndex(index), false); + } + jotaiStore.set(pendingRunEndByConvoId(end.conversationId), { + ...end, + ...((end.interruptArmed === true || interruptArmed) && { + interruptArmed: true, + }), + }); + jotaiStore.set(runEndByIndex(index), null); + return null; + } + if (end == null && activeConversationId) { + const parked = jotaiStore.get(pendingRunEndByConvoId(activeConversationId)); + if (parked != null) { + end = parked; + fromParked = true; + } + } + if (end == null) { + return null; + } + const indexArmed = jotaiStore.get(drainAfterAbortByIndex(index)); + const matchingIndexArm = matchesInterruptArm(indexArmed, end); + const interruptArmed = matchingIndexArm || end.interruptArmed === true; - const conversationId = end.conversationId; - if (!conversationId) { - return null; - } + const conversationId = end.conversationId; + if (!conversationId) { + return null; + } - const shouldMigrate = - end.startedAsNewConvo === true && conversationId !== Constants.NEW_CONVO; - const newConvoQueue = shouldMigrate - ? snapshot.getLoadable(store.queuedMessagesByConvoId(Constants.NEW_CONVO)).getValue() - : []; - const ownQueue = snapshot - .getLoadable(store.queuedMessagesByConvoId(conversationId)) - .getValue(); - /** Both queues are ordered independently, but migration crosses the - * key boundary: an interrupt queued under the resolved conversation - * must still outrank an ordinary follow-up captured under NEW_CONVO. */ - const merged = shouldMigrate ? mergeQueuedMessages(newConvoQueue, ownQueue) : ownQueue; + const shouldMigrate = end.startedAsNewConvo === true && conversationId !== Constants.NEW_CONVO; + const newConvoQueue = shouldMigrate + ? jotaiStore.get(queuedMessagesByConvoId(Constants.NEW_CONVO)) + : []; + const ownQueue = jotaiStore.get(queuedMessagesByConvoId(conversationId)); + /** Both queues are ordered independently, but migration crosses the + * key boundary: an interrupt queued under the resolved conversation + * must still outrank an ordinary follow-up captured under NEW_CONVO. */ + const merged = shouldMigrate ? mergeQueuedMessages(newConvoQueue, ownQueue) : ownQueue; - const shouldDrain = end.outcome === 'completed' || interruptArmed; - const settledReceipts = snapshot - .getLoadable(store.settledQueuedTurnReceiptsByConvoId(conversationId)) - .getValue(); - const pendingEnqueueIds = snapshot - .getLoadable(store.pendingQueuedTurnEnqueueIdsByConvoId(conversationId)) - .getValue(); - const consumedReceiptIndex = settledReceipts.findIndex( - (receipt) => + const shouldDrain = end.outcome === 'completed' || interruptArmed; + const settledReceipts = jotaiStore.get(settledQueuedTurnReceiptsByConvoId(conversationId)); + const pendingEnqueueIds = jotaiStore.get(pendingQueuedTurnEnqueueIdsByConvoId(conversationId)); + const consumedReceiptIndex = settledReceipts.findIndex( + (receipt) => + receipt.status === 'admitted' && + receipt.boundaryConsumed !== true && + end.generationCreatedAt != null && + receipt.effectivePredecessorCreatedAt === end.generationCreatedAt, + ); + const consumedByServerAdmission = consumedReceiptIndex >= 0; + const consumeEnd = () => { + if (fromParked && activeConversationId) { + jotaiStore.set(pendingRunEndByConvoId(activeConversationId), null); + } else { + jotaiStore.set(runEndByIndex(index), null); + } + if (matchingIndexArm) { + jotaiStore.set(drainAfterAbortByIndex(index), false); + } + }; + if (consumedByServerAdmission) { + /** Admission consumed this terminal boundary on the server. A late + * client terminal observation cannot authorize a second successor. */ + if (shouldMigrate && newConvoQueue.length > 0) { + jotaiStore.set(queuedMessagesByConvoId(Constants.NEW_CONVO), []); + jotaiStore.set(queuedMessagesByConvoId(conversationId), merged); + } + jotaiStore.set(settledQueuedTurnReceiptsByConvoId(conversationId), (previous) => { + let consumed = false; + return previous.flatMap((receipt) => { + if ( + !consumed && receipt.status === 'admitted' && receipt.boundaryConsumed !== true && - end.generationCreatedAt != null && - receipt.effectivePredecessorCreatedAt === end.generationCreatedAt, - ); - const consumedByServerAdmission = consumedReceiptIndex >= 0; - const consumeEnd = () => { - if (fromParked && activeConversationId) { - set(store.pendingRunEndByConvoId(activeConversationId), null); - } else { - set(store.runEndByIndex(index), null); + receipt.effectivePredecessorCreatedAt === end.generationCreatedAt + ) { + consumed = true; + return pendingEnqueueIds.includes(receipt.clientRequestId) + ? [{ ...receipt, boundaryConsumed: true }] + : []; } - if (matchingIndexArm) { - set(store.drainAfterAbortByIndex(index), false); - } - }; - if (consumedByServerAdmission) { - /** Admission consumed this terminal boundary on the server. A late - * client terminal observation cannot authorize a second successor. */ - if (shouldMigrate && newConvoQueue.length > 0) { - set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), []); - set(store.queuedMessagesByConvoId(conversationId), merged); - } - set(store.settledQueuedTurnReceiptsByConvoId(conversationId), (previous) => { - let consumed = false; - return previous.flatMap((receipt) => { - if ( - !consumed && - receipt.status === 'admitted' && - receipt.boundaryConsumed !== true && - receipt.effectivePredecessorCreatedAt === end.generationCreatedAt - ) { - consumed = true; - return pendingEnqueueIds.includes(receipt.clientRequestId) - ? [{ ...receipt, boundaryConsumed: true }] - : []; - } - return [receipt]; - }); - }); - consumeEnd(); - return null; - } - /** A server-owned Agent row means the backend owns the next fresh-turn - * admission. Do not let a legacy/recovered local row race or overtake - * it. Keep this terminal boundary available until the authoritative - * queue snapshot proves whether a server-started successor now owns it. */ - const serverOwnsBoundary = merged.some((item) => item.server != null); - if (serverOwnsBoundary) { - /** Keep the one-shot terminal signal until the authoritative snapshot - * removes the server row. If it was admitted, its own later terminal - * signal orders the remaining local queue; if it was cancelled/dead, - * this signal still lets the legacy successor make progress. The - * head the server is about to admit can already be shown as the - * next user turn; the signal itself stays untouched. */ - const reveal = selectQueuedTurnReveal(end, merged); - return reveal == null ? null : { kind: 'reveal', item: reveal, end }; - } + return [receipt]; + }); + }); + consumeEnd(); + return null; + } + /** A server-owned Agent row means the backend owns the next fresh-turn + * admission. Do not let a legacy/recovered local row race or overtake + * it. Keep this terminal boundary available until the authoritative + * queue snapshot proves whether a server-started successor now owns it. */ + const serverOwnsBoundary = merged.some((item) => item.server != null); + if (serverOwnsBoundary) { + /** Keep the one-shot terminal signal until the authoritative snapshot + * removes the server row. If it was admitted, its own later terminal + * signal orders the remaining local queue; if it was cancelled/dead, + * this signal still lets the legacy successor make progress. The + * head the server is about to admit can already be shown as the + * next user turn; the signal itself stays untouched. */ + const reveal = selectQueuedTurnReveal(end, merged); + return reveal == null ? null : { kind: 'reveal', item: reveal, end }; + } - const head = merged[0]; - if ( - shouldDrain && - head != null && - recoveryDisposition(jotaiStore.get(recoveryDispositionsFamily(conversationId)), head) != - null - ) { - // The first row still owns this boundary. Keep it until dismissal, even when - // a later local row is immediately sendable; a user must choose its fate. - if (shouldMigrate && newConvoQueue.length > 0) { - set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), []); - set(store.queuedMessagesByConvoId(conversationId), merged); - } - return null; - } - // Intent-held and explicitly refused rows are skipped by the redesigned rail. - consumeEnd(); + const head = merged[0]; + if ( + shouldDrain && + head != null && + recoveryDisposition(jotaiStore.get(recoveryDispositionsFamily(conversationId)), head) != null + ) { + // The first row still owns this boundary. Keep it until dismissal, even when + // a later local row is immediately sendable; a user must choose its fate. + if (shouldMigrate && newConvoQueue.length > 0) { + jotaiStore.set(queuedMessagesByConvoId(Constants.NEW_CONVO), []); + jotaiStore.set(queuedMessagesByConvoId(conversationId), merged); + } + return null; + } + // Intent-held and explicitly refused rows are skipped by the redesigned rail. + consumeEnd(); - /** A row the rail is mid-edit or mid-remove on is spoken for: its words - * are already on their way to the composer, and sending them from here - * would deliver the message the user is in the middle of taking back. - * A row swept in from a REJECTED steer is spoken for too: it carries - * words the server refused, and its failure surface offers Retry and - * "Send as new" precisely so the user chooses. - * Both are skipped rather than blocking the whole queue, so an - * untouched follow-up behind them still goes on this run end. */ - const nextIndex = shouldDrain - ? merged.findIndex( - (item) => - !hasQueuedIntent(item.id) && - item.needsExplicitSend !== true && - recoveryDisposition( - jotaiStore.get(recoveryDispositionsFamily(conversationId)), - item, - ) == null, - ) - : -1; - const next = nextIndex >= 0 ? merged[nextIndex] : null; - const remainder = nextIndex >= 0 ? merged.filter((_, at) => at !== nextIndex) : merged; + /** A row the rail is mid-edit or mid-remove on is spoken for: its words + * are already on their way to the composer, and sending them from here + * would deliver the message the user is in the middle of taking back. + * A row swept in from a REJECTED steer is spoken for too: it carries + * words the server refused, and its failure surface offers Retry and + * "Send as new" precisely so the user chooses. + * Both are skipped rather than blocking the whole queue, so an + * untouched follow-up behind them still goes on this run end. */ + const nextIndex = shouldDrain + ? merged.findIndex( + (item) => + !hasQueuedIntent(item.id) && + item.needsExplicitSend !== true && + recoveryDisposition(jotaiStore.get(recoveryDispositionsFamily(conversationId)), item) == + null, + ) + : -1; + const next = nextIndex >= 0 ? merged[nextIndex] : null; + const remainder = nextIndex >= 0 ? merged.filter((_, at) => at !== nextIndex) : merged; - if (shouldMigrate && newConvoQueue.length > 0) { - set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), []); - } - if (remainder.length !== ownQueue.length || shouldMigrate || next != null) { - set(store.queuedMessagesByConvoId(conversationId), remainder); + if (shouldMigrate && newConvoQueue.length > 0) { + jotaiStore.set(queuedMessagesByConvoId(Constants.NEW_CONVO), []); + } + if (remainder.length !== ownQueue.length || shouldMigrate || next != null) { + jotaiStore.set(queuedMessagesByConvoId(conversationId), remainder); + } + return next + ? { + kind: 'drain', + next, + conversationId, + queuedMessageOrigin: { + item: next, + beforeIds: merged.slice(0, nextIndex).map((item) => item.id), + afterIds: merged.slice(nextIndex + 1).map((item) => item.id), + }, + expectedPredecessorCreatedAt: + end.generationCreatedAt ?? next.expectedPredecessorCreatedAt, } - return next - ? { - kind: 'drain', - next, - conversationId, - queuedMessageOrigin: { - item: next, - beforeIds: merged.slice(0, nextIndex).map((item) => item.id), - afterIds: merged.slice(nextIndex + 1).map((item) => item.id), - }, - expectedPredecessorCreatedAt: - end.generationCreatedAt ?? next.expectedPredecessorCreatedAt, - } - : null; - }, - [index, activeConversationId, jotaiStore], - ); + : null; + }, [index, activeConversationId, jotaiStore]); - const restoreQueued = useRecoilCallback( - ({ set }) => - (convoId: string, origin: QueuedMessageOrigin) => { - if (!canRestoreRecovery(jotaiStore.get(recoveryDispositionsFamily(convoId)), origin.item)) - return; - set(store.queuedMessagesByConvoId(convoId), (prev) => insertQueuedOrigin(prev, origin)); - }, + const restoreQueued = useCallback( + (convoId: string, origin: QueuedMessageOrigin) => { + if (!canRestoreRecovery(jotaiStore.get(recoveryDispositionsFamily(convoId)), origin.item)) + return; + jotaiStore.set(queuedMessagesByConvoId(convoId), (prev) => insertQueuedOrigin(prev, origin)); + }, [jotaiStore], ); diff --git a/client/src/hooks/Chat/useSteerCancel.ts b/client/src/hooks/Chat/useSteerCancel.ts index b004cce81ad..ecc2ae9e5d1 100644 --- a/client/src/hooks/Chat/useSteerCancel.ts +++ b/client/src/hooks/Chat/useSteerCancel.ts @@ -1,5 +1,5 @@ import { useCallback } from 'react'; -import { useSetAtom } from 'jotai'; +import { useSetAtom, useStore } from 'jotai'; import { useRecoilCallback, useRecoilValue } from 'recoil'; import type { PendingSteer } from '~/store/families'; import { useComposerRestoreHost } from '~/Providers/ComposerRestoreContext'; @@ -7,6 +7,7 @@ import { appendAppliedSteerIds, carriedSteerContext } from '~/utils'; import { pendingSteerCancelClientIdsFamily } from '~/store/steer'; import useSteerConvert from '~/hooks/Chat/useSteerConvert'; import { useCancelSteerMutation } from '~/data-provider'; +import { queuedMessagesByConvoId } from '~/store/queue'; import store from '~/store'; export type SteerCancelOutcome = 'reclaimed' | 'applied' | 'failed'; @@ -24,6 +25,7 @@ export function useSteerReclaim(conversationId: string) { store.activeGenerationCreatedAtByConvoId(conversationId), ); const setPendingCancelIds = useSetAtom(pendingSteerCancelClientIdsFamily(conversationId)); + const queueStore = useStore(); const settleReclaimed = useRecoilCallback( ({ set }) => (steer: PendingSteer) => { @@ -40,7 +42,7 @@ export function useSteerReclaim(conversationId: string) { ), ); setPendingCancelIds((prev) => prev.filter((id) => !ids.includes(id))); - set(store.queuedMessagesByConvoId(conversationId), (prev) => + queueStore.set(queuedMessagesByConvoId(conversationId), (prev) => prev.filter( (item) => !settled.has(item.id) && @@ -49,7 +51,7 @@ export function useSteerReclaim(conversationId: string) { ), ); }, - [conversationId, setPendingCancelIds], + [conversationId, setPendingCancelIds, queueStore], ); return useCallback( diff --git a/client/src/hooks/Chat/useSteerConvert.ts b/client/src/hooks/Chat/useSteerConvert.ts index df929729189..f26e9a882db 100644 --- a/client/src/hooks/Chat/useSteerConvert.ts +++ b/client/src/hooks/Chat/useSteerConvert.ts @@ -17,6 +17,7 @@ import { canRestoreRecovery, } from '~/components/Chat/Steering/recovery'; import { fetchStreamStatus, getGenerationProtocolVersion } from '~/data-provider'; +import { queuedMessagesByConvoId } from '~/store/queue'; import { useFileMapContext } from '~/Providers'; import store from '~/store'; @@ -119,7 +120,7 @@ export default function useSteerConvert() { set(store.pendingSteersByConvoId(conversationId), (prev) => prev.filter((steer) => !steerIds.has(steer.steerId)), ); - set(store.queuedMessagesByConvoId(conversationId), (prev) => { + jotaiStore.set(queuedMessagesByConvoId(conversationId), (prev) => { /** A legacy status read is destructive: if a v2 live-final path * already created a receipt-bound item before the claim reached an * old replica, that source no longer exists. Downgrade the existing diff --git a/client/src/hooks/Chat/useSteering.ts b/client/src/hooks/Chat/useSteering.ts index eda745b02f6..223b44806d7 100644 --- a/client/src/hooks/Chat/useSteering.ts +++ b/client/src/hooks/Chat/useSteering.ts @@ -1,8 +1,8 @@ import { useCallback, useEffect, useMemo, useRef } from 'react'; import { v4 } from 'uuid'; -import { useAtomValue, useStore } from 'jotai'; import { useToastContext } from '@librechat/client'; -import { useRecoilValue, useSetRecoilState, useRecoilCallback } from 'recoil'; +import { useAtomValue, useSetAtom, useStore } from 'jotai'; +import { useRecoilValue, useRecoilCallback } from 'recoil'; import { Constants, ContentTypes, @@ -41,6 +41,14 @@ import { isDefiniteQueuedTurnsUnsupported, supportsGenerationProtocolV2, } from '~/data-provider'; +import { + settledQueuedTurnReceiptsByConvoId, + pendingQueuedTurnEnqueueIdsByConvoId, + pendingRunEndByConvoId, + queuedMessagesByConvoId, + drainAfterAbortByIndex, + runEndByIndex, +} from '~/store/queue'; import { appendAppliedSteerIds, carriedSteerContext, @@ -612,9 +620,9 @@ export default function useSteering({ const enabled = steerable && index === 0; const queueKey = hasRealConvoId ? conversationId : Constants.NEW_CONVO; const reasoningStateKey = getReasoningStateKey(conversationId, index); - const queuedMessages = useRecoilValue(store.queuedMessagesByConvoId(queueKey)); + const queuedMessages = useAtomValue(queuedMessagesByConvoId(queueKey)); const pendingReasoningOverride = useAtomValue(pendingReasoningOverrideFamily(reasoningStateKey)); - const setQueuedMessages = useSetRecoilState(store.queuedMessagesByConvoId(queueKey)); + const setQueuedMessages = useSetAtom(queuedMessagesByConvoId(queueKey)); const { data: startupConfig } = useGetStartupConfig(); /** Bound transport-outcome reconciliation while still guaranteeing several * list reads after the enqueue promise settles. Focus/remount remains a later @@ -686,86 +694,75 @@ export default function useSteering({ store.activeGenerationProtocolVersionByConvoId(queueKey), ); - const applyQueuedTurnReceipts = useRecoilCallback( - ({ snapshot, set }) => - (receipts: AgentQueuedTurnReceipt[], source: QueuedTurnReceiptSource = 'snapshot') => { - const previousSettled = snapshot - .getLoadable(store.settledQueuedTurnReceiptsByConvoId(queueKey)) - .getValue(); - const previousPending = snapshot - .getLoadable(store.pendingQueuedTurnEnqueueIdsByConvoId(queueKey)) - .getValue(); - const completedRequestIds = new Set( - source === 'enqueue' ? receipts.map((receipt) => receipt.clientRequestId) : [], - ); - const nextPending = previousPending.filter((id) => !completedRequestIds.has(id)); - const settledByRequestId = new Map( - previousSettled.map((receipt) => [receipt.clientRequestId, receipt]), - ); - for (const receipt of receipts) { - const existing = settledByRequestId.get(receipt.clientRequestId); - const settled = mergeSettledQueuedTurnEvidence(existing, receipt, source); - if (settled == null) { - continue; - } - settledByRequestId.set(receipt.clientRequestId, settled); + const applyQueuedTurnReceipts = useCallback( + (receipts: AgentQueuedTurnReceipt[], source: QueuedTurnReceiptSource = 'snapshot') => { + const previousSettled = jotaiStore.get(settledQueuedTurnReceiptsByConvoId(queueKey)); + const previousPending = jotaiStore.get(pendingQueuedTurnEnqueueIdsByConvoId(queueKey)); + const completedRequestIds = new Set( + source === 'enqueue' ? receipts.map((receipt) => receipt.clientRequestId) : [], + ); + const nextPending = previousPending.filter((id) => !completedRequestIds.has(id)); + const settledByRequestId = new Map( + previousSettled.map((receipt) => [receipt.clientRequestId, receipt]), + ); + for (const receipt of receipts) { + const existing = settledByRequestId.get(receipt.clientRequestId); + const settled = mergeSettledQueuedTurnEvidence(existing, receipt, source); + if (settled == null) { + continue; } - const terminalForReconciliation = new Map(settledByRequestId); - const pendingRequestIds = new Set(nextPending); - const nextSettled = [...settledByRequestId.values()].filter( - (receipt) => - (receipt.status === 'admitted' && - receipt.rootPredecessor !== true && - receipt.boundaryConsumed !== true) || - pendingRequestIds.has(receipt.clientRequestId), + settledByRequestId.set(receipt.clientRequestId, settled); + } + const terminalForReconciliation = new Map(settledByRequestId); + const pendingRequestIds = new Set(nextPending); + const nextSettled = [...settledByRequestId.values()].filter( + (receipt) => + (receipt.status === 'admitted' && + receipt.rootPredecessor !== true && + receipt.boundaryConsumed !== true) || + pendingRequestIds.has(receipt.clientRequestId), + ); + if (source === 'enqueue') { + jotaiStore.set(pendingQueuedTurnEnqueueIdsByConvoId(queueKey), nextPending); + } + if ( + nextSettled.length !== previousSettled.length || + nextSettled.some((receipt, index) => receipt !== previousSettled[index]) + ) { + jotaiStore.set(settledQueuedTurnReceiptsByConvoId(queueKey), nextSettled); + } + if (source !== 'direct') { + jotaiStore.set(queuedMessagesByConvoId(queueKey), (previous) => + reconcileServerQueuedTurns( + previous, + receipts, + terminalForReconciliation, + source === 'snapshot', + fileMap, + ), ); - if (source === 'enqueue') { - set(store.pendingQueuedTurnEnqueueIdsByConvoId(queueKey), nextPending); - } - if ( - nextSettled.length !== previousSettled.length || - nextSettled.some((receipt, index) => receipt !== previousSettled[index]) - ) { - set(store.settledQueuedTurnReceiptsByConvoId(queueKey), nextSettled); - } - if (source !== 'direct') { - set(store.queuedMessagesByConvoId(queueKey), (previous) => - reconcileServerQueuedTurns( - previous, - receipts, - terminalForReconciliation, - source === 'snapshot', - fileMap, - ), - ); - } - }, - [fileMap, queueKey], + } + }, + [fileMap, queueKey, jotaiStore], ); - const finishQueuedTurnEnqueue = useRecoilCallback( - ({ snapshot, set }) => - (clientRequestId: string): boolean => { - const settledReceipts = snapshot - .getLoadable(store.settledQueuedTurnReceiptsByConvoId(queueKey)) - .getValue(); - const settled = settledReceipts.find( - (receipt) => receipt.clientRequestId === clientRequestId, - ); - set(store.pendingQueuedTurnEnqueueIdsByConvoId(queueKey), (previous) => - previous.filter((id) => id !== clientRequestId), + const finishQueuedTurnEnqueue = useCallback( + (clientRequestId: string): boolean => { + const settledReceipts = jotaiStore.get(settledQueuedTurnReceiptsByConvoId(queueKey)); + const settled = settledReceipts.find( + (receipt) => receipt.clientRequestId === clientRequestId, + ); + jotaiStore.set(pendingQueuedTurnEnqueueIdsByConvoId(queueKey), (previous) => + previous.filter((id) => id !== clientRequestId), + ); + if (settled != null && (settled.status !== 'admitted' || settled.boundaryConsumed === true)) { + jotaiStore.set(settledQueuedTurnReceiptsByConvoId(queueKey), (previous) => + previous.filter((receipt) => receipt.clientRequestId !== clientRequestId), ); - if ( - settled != null && - (settled.status !== 'admitted' || settled.boundaryConsumed === true) - ) { - set(store.settledQueuedTurnReceiptsByConvoId(queueKey), (previous) => - previous.filter((receipt) => receipt.clientRequestId !== clientRequestId), - ); - } - return settled != null; - }, - [queueKey], + } + return settled != null; + }, + [queueKey, jotaiStore], ); useEffect(() => { @@ -955,8 +952,8 @@ export default function useSteering({ * not the one the words belong to. An entry is dropped once that conversation * starts another run; an older end no longer describes what is happening. */ - const runEnd = useRecoilValue(store.runEndByIndex(index)); - const parkedRunEnd = useRecoilValue(store.pendingRunEndByConvoId(queueKey)); + const runEnd = useAtomValue(runEndByIndex(index)); + const parkedRunEnd = useAtomValue(pendingRunEndByConvoId(queueKey)); const runEndsRef = useRef>(new Map()); const observedRunEnd = [runEnd, parkedRunEnd].find( (end) => end != null && end.conversationId === conversationId, @@ -1096,175 +1093,173 @@ export default function useSteering({ [markFilesUsage], ); - const updateQueuedMessage = useRecoilCallback( - ({ set }) => - (id: string, update: (item: QueuedMessage) => QueuedMessage | null) => { - set(store.queuedMessagesByConvoId(queueKey), (previous) => - previous.flatMap((item) => { - if (item.id !== id) { - return [item]; - } - const next = update(item); - return next == null ? [] : [next]; - }), - ); - }, - [queueKey], + const updateQueuedMessage = useCallback( + (id: string, update: (item: QueuedMessage) => QueuedMessage | null) => { + jotaiStore.set(queuedMessagesByConvoId(queueKey), (previous) => + previous.flatMap((item) => { + if (item.id !== id) { + return [item]; + } + const next = update(item); + return next == null ? [] : [next]; + }), + ); + }, + [queueKey, jotaiStore], ); - const enqueue = useRecoilCallback( - ({ set }) => - ( - text: string, - options?: { - front?: boolean; - files?: TMessage['files']; - quotes?: string[]; - manualSkills?: string[]; - reasoningOverride?: TMessage['reasoningOverride']; - /** Set when the files were ALREADY queued/steered: their TTL was - * held when they first entered the queue (or at the steer 202). */ - skipUsageMark?: boolean; - id?: string; - createdAt?: number; - clientRequestId?: string; - expectedPredecessorCreatedAt?: number; - /** The queue lineage of a row being put back after its parked copy was - * cancelled. The run it waited on may have ended meanwhile, which - * clears the live pair, and the row must stay durable regardless. */ - lineage?: { parentMessageId: string; predecessorCreatedAt: number }; - }, - ) => { - const trimmed = text.trim(); - if (trimmed.length === 0) { - return; - } - const parentMessageId = - options?.lineage?.parentMessageId ?? - (pendingReveal != null - ? pendingReveal.queueParentMessageId - : liveMessageState?.parentMessageId); - const predecessorCreatedAt = - options?.lineage?.predecessorCreatedAt ?? - (pendingReveal != null - ? pendingReveal.queuePredecessorCreatedAt - : activeGenerationCreatedAt); - /** FINAL clears the active epoch before attachment. The revealed - * intent retains the queue's original parent/epoch pair. Its display - * parent and advancing completion boundary are not queue lineage. - * Without an authoritative pair, retain the follow-up locally. */ - const serverOwned = - serverQueueEnabled && parentMessageId != null && predecessorCreatedAt != null; - const generatedClientRequestId = options?.clientRequestId == null; - const clientRequestId = options?.clientRequestId ?? (serverOwned ? v4() : undefined); - const item: QueuedMessage = { - id: options?.id ?? v4(), - text: trimmed, - createdAt: options?.createdAt ?? Date.now(), - ...(clientRequestId != null && { clientRequestId }), - ...(serverOwned && { - parentMessageId, - server: { status: 'sending' }, - }), - ...((options?.expectedPredecessorCreatedAt ?? predecessorCreatedAt) != null && { - expectedPredecessorCreatedAt: - options?.expectedPredecessorCreatedAt ?? predecessorCreatedAt ?? undefined, - }), - ...(options?.files && options.files.length > 0 && { files: options.files }), - ...(options?.quotes && options.quotes.length > 0 && { quotes: options.quotes }), - ...(options?.manualSkills && - options.manualSkills.length > 0 && { - manualSkills: options.manualSkills, - }), - ...(options?.reasoningOverride != null && { - reasoningOverride: options.reasoningOverride, + const enqueue = useCallback( + ( + text: string, + options?: { + front?: boolean; + files?: TMessage['files']; + quotes?: string[]; + manualSkills?: string[]; + reasoningOverride?: TMessage['reasoningOverride']; + /** Set when the files were ALREADY queued/steered: their TTL was + * held when they first entered the queue (or at the steer 202). */ + skipUsageMark?: boolean; + id?: string; + createdAt?: number; + clientRequestId?: string; + expectedPredecessorCreatedAt?: number; + /** The queue lineage of a row being put back after its parked copy was + * cancelled. The run it waited on may have ended meanwhile, which + * clears the live pair, and the row must stay durable regardless. */ + lineage?: { parentMessageId: string; predecessorCreatedAt: number }; + }, + ) => { + const trimmed = text.trim(); + if (trimmed.length === 0) { + return; + } + const parentMessageId = + options?.lineage?.parentMessageId ?? + (pendingReveal != null + ? pendingReveal.queueParentMessageId + : liveMessageState?.parentMessageId); + const predecessorCreatedAt = + options?.lineage?.predecessorCreatedAt ?? + (pendingReveal != null + ? pendingReveal.queuePredecessorCreatedAt + : activeGenerationCreatedAt); + /** FINAL clears the active epoch before attachment. The revealed + * intent retains the queue's original parent/epoch pair. Its display + * parent and advancing completion boundary are not queue lineage. + * Without an authoritative pair, retain the follow-up locally. */ + const serverOwned = + serverQueueEnabled && parentMessageId != null && predecessorCreatedAt != null; + const generatedClientRequestId = options?.clientRequestId == null; + const clientRequestId = options?.clientRequestId ?? (serverOwned ? v4() : undefined); + const item: QueuedMessage = { + id: options?.id ?? v4(), + text: trimmed, + createdAt: options?.createdAt ?? Date.now(), + ...(clientRequestId != null && { clientRequestId }), + ...(serverOwned && { + parentMessageId, + server: { status: 'sending' }, + }), + ...((options?.expectedPredecessorCreatedAt ?? predecessorCreatedAt) != null && { + expectedPredecessorCreatedAt: + options?.expectedPredecessorCreatedAt ?? predecessorCreatedAt ?? undefined, + }), + ...(options?.files && options.files.length > 0 && { files: options.files }), + ...(options?.quotes && options.quotes.length > 0 && { quotes: options.quotes }), + ...(options?.manualSkills && + options.manualSkills.length > 0 && { + manualSkills: options.manualSkills, }), - ...(options?.front && { priority: true }), - }; - set(store.queuedMessagesByConvoId(queueKey), (prev) => insertQueuedMessage(prev, item)); - if (options?.skipUsageMark !== true) { - markQueuedFilesUsage(options?.files); - } - if (!serverOwned || clientRequestId == null) { - return; - } - set(store.pendingQueuedTurnEnqueueIdsByConvoId(queueKey), (previous) => - previous.includes(clientRequestId) ? previous : [...previous, clientRequestId], - ); - const serverFiles = toQueuedTurnFileRefs(item.files); - /** Commit the optimistic Recoil row before mutation callbacks can - * reconcile it. This also makes synchronous test/adaptor completions - * obey the same ordering as a real network response. */ - queueMicrotask(() => - enqueueAgentQueuedTurn( - { - conversationId, - clientRequestId, - parentMessageId, - text: item.text, - ...(serverFiles != null && { files: serverFiles }), - ...(item.quotes != null && item.quotes.length > 0 && { quotes: item.quotes }), - ...(item.manualSkills != null && - item.manualSkills.length > 0 && { - manualSkills: item.manualSkills, - }), - ...(codeApprovalMode != null && { codeApprovalMode }), - ...(item.reasoningOverride != null && { - reasoningOverride: item.reasoningOverride, - }), - ...(item.priority === true && { priority: true }), - ...(item.expectedPredecessorCreatedAt != null && { - expectedPredecessorCreatedAt: item.expectedPredecessorCreatedAt, + ...(options?.reasoningOverride != null && { + reasoningOverride: options.reasoningOverride, + }), + ...(options?.front && { priority: true }), + }; + jotaiStore.set(queuedMessagesByConvoId(queueKey), (prev) => insertQueuedMessage(prev, item)); + if (options?.skipUsageMark !== true) { + markQueuedFilesUsage(options?.files); + } + if (!serverOwned || clientRequestId == null) { + return; + } + jotaiStore.set(pendingQueuedTurnEnqueueIdsByConvoId(queueKey), (previous) => + previous.includes(clientRequestId) ? previous : [...previous, clientRequestId], + ); + const serverFiles = toQueuedTurnFileRefs(item.files); + /** Commit the optimistic Recoil row before mutation callbacks can + * reconcile it. This also makes synchronous test/adaptor completions + * obey the same ordering as a real network response. */ + queueMicrotask(() => + enqueueAgentQueuedTurn( + { + conversationId, + clientRequestId, + parentMessageId, + text: item.text, + ...(serverFiles != null && { files: serverFiles }), + ...(item.quotes != null && item.quotes.length > 0 && { quotes: item.quotes }), + ...(item.manualSkills != null && + item.manualSkills.length > 0 && { + manualSkills: item.manualSkills, }), + ...(codeApprovalMode != null && { codeApprovalMode }), + ...(item.reasoningOverride != null && { + reasoningOverride: item.reasoningOverride, + }), + ...(item.priority === true && { priority: true }), + ...(item.expectedPredecessorCreatedAt != null && { + expectedPredecessorCreatedAt: item.expectedPredecessorCreatedAt, + }), + }, + { + onSuccess: (receipt) => { + applyQueuedTurnReceipts([receipt], 'enqueue'); }, - { - onSuccess: (receipt) => { - applyQueuedTurnReceipts([receipt], 'enqueue'); - }, - onError: (error) => { - if (finishQueuedTurnEnqueue(clientRequestId)) { - return; + onError: (error) => { + if (finishQueuedTurnEnqueue(clientRequestId)) { + return; + } + updateQueuedMessage(item.id, (current) => { + if (isDefiniteQueuedTurnsUnsupported(error)) { + const { + server: _server, + parentMessageId: _parentMessageId, + clientRequestId: fallbackClientRequestId, + ...legacy + } = current; + return { + ...legacy, + ...(!generatedClientRequestId && fallbackClientRequestId != null + ? { clientRequestId: fallbackClientRequestId } + : {}), + }; } - updateQueuedMessage(item.id, (current) => { - if (isDefiniteQueuedTurnsUnsupported(error)) { - const { - server: _server, - parentMessageId: _parentMessageId, - clientRequestId: fallbackClientRequestId, - ...legacy - } = current; - return { - ...legacy, - ...(!generatedClientRequestId && fallbackClientRequestId != null - ? { clientRequestId: fallbackClientRequestId } - : {}), - }; - } - if (!isDefiniteQueuedTurnRejection(error)) { - return { - ...current, - server: { - ...current.server, - status: 'uncertain', - uncertainSince: current.server?.uncertainSince ?? Date.now(), - }, - }; - } - const code = getSteerErrorCode(error); + if (!isDefiniteQueuedTurnRejection(error)) { return { ...current, server: { ...current.server, - status: 'rejected', - ...(code != null && { errorCode: code }), + status: 'uncertain', + uncertainSince: current.server?.uncertainSince ?? Date.now(), }, }; - }); - }, + } + const code = getSteerErrorCode(error); + return { + ...current, + server: { + ...current.server, + status: 'rejected', + ...(code != null && { errorCode: code }), + }, + }; + }); }, - ), - ); - }, + }, + ), + ); + }, [ queueKey, conversationId, @@ -1278,6 +1273,7 @@ export default function useSteering({ applyQueuedTurnReceipts, finishQueuedTurnEnqueue, updateQueuedMessage, + jotaiStore, ], ); @@ -1427,14 +1423,13 @@ export default function useSteering({ would restore text the user already sent. Only autosave knows whether this composer is parked on the pane's pending key or on the conversation key, so the clear belongs there rather than being recomputed here. */ - const removeQueued = useRecoilCallback( - ({ set }) => - (id: string) => { - set(store.queuedMessagesByConvoId(queueKey), (prev) => - prev.filter((item) => item.id !== id), - ); - }, - [queueKey], + const removeQueued = useCallback( + (id: string) => { + jotaiStore.set(queuedMessagesByConvoId(queueKey), (prev) => + prev.filter((item) => item.id !== id), + ); + }, + [queueKey, jotaiStore], ); /** Keeps a row in the queue but out of the run-end drain, for words that @@ -1444,25 +1439,24 @@ export default function useSteering({ [updateQueuedMessage], ); - const downgradeServerQueuedTurn = useRecoilCallback( - ({ snapshot, set }) => - (id: string): boolean => { - const queue = snapshot.getLoadable(store.queuedMessagesByConvoId(queueKey)).getValue(); - let found = false; - const next = queue.map((item) => { - if (item.id !== id) { - return item; - } - found = true; - const { server: _server, parentMessageId: _parentMessageId, ...local } = item; - return local; - }); - if (found) { - set(store.queuedMessagesByConvoId(queueKey), next); + const downgradeServerQueuedTurn = useCallback( + (id: string): boolean => { + const queue = jotaiStore.get(queuedMessagesByConvoId(queueKey)); + let found = false; + const next = queue.map((item) => { + if (item.id !== id) { + return item; } - return found; - }, - [queueKey], + found = true; + const { server: _server, parentMessageId: _parentMessageId, ...local } = item; + return local; + }); + if (found) { + jotaiStore.set(queuedMessagesByConvoId(queueKey), next); + } + return found; + }, + [queueKey, jotaiStore], ); /** @@ -1474,22 +1468,21 @@ export default function useSteering({ * Addressed by id rather than by the index the caller is holding, which a * drain can invalidate between the drag starting and the drop landing. */ - const reorderQueued = useRecoilCallback( - ({ set }) => - (id: string, targetIndex: number) => { - set(store.queuedMessagesByConvoId(queueKey), (prev) => { - const from = prev.findIndex((item) => item.id === id); - const to = Math.min(Math.max(targetIndex, 0), prev.length - 1); - if (from === -1 || from === to) { - return prev; - } - const next = prev.slice(); - const [moved] = next.splice(from, 1); - next.splice(to, 0, moved); - return next; - }); - }, - [queueKey], + const reorderQueued = useCallback( + (id: string, targetIndex: number) => { + jotaiStore.set(queuedMessagesByConvoId(queueKey), (prev) => { + const from = prev.findIndex((item) => item.id === id); + const to = Math.min(Math.max(targetIndex, 0), prev.length - 1); + if (from === -1 || from === to) { + return prev; + } + const next = prev.slice(); + const [moved] = next.splice(from, 1); + next.splice(to, 0, moved); + return next; + }); + }, + [queueKey, jotaiStore], ); /** @@ -1497,26 +1490,25 @@ export default function useSteering({ * Ids that have since drained are skipped rather than resurrected, and * anything queued mid-drag keeps its place at the back. */ - const restoreQueuedOrder = useRecoilCallback( - ({ set }) => - (ids: readonly string[]) => { - set(store.queuedMessagesByConvoId(queueKey), (prev) => { - const byId = new Map(prev.map((item) => [item.id, item])); - const restored: QueuedMessage[] = []; - for (const id of ids) { - const item = byId.get(id); - if (item != null) { - restored.push(item); - byId.delete(id); - } - } - if (restored.length === 0) { - return prev; + const restoreQueuedOrder = useCallback( + (ids: readonly string[]) => { + jotaiStore.set(queuedMessagesByConvoId(queueKey), (prev) => { + const byId = new Map(prev.map((item) => [item.id, item])); + const restored: QueuedMessage[] = []; + for (const id of ids) { + const item = byId.get(id); + if (item != null) { + restored.push(item); + byId.delete(id); } - return [...restored, ...byId.values()]; - }); - }, - [queueKey], + } + if (restored.length === 0) { + return prev; + } + return [...restored, ...byId.values()]; + }); + }, + [queueKey, jotaiStore], ); /** Settle a queued row's terminal recovery source before an Edit/Remove. @@ -1656,37 +1648,36 @@ export default function useSteering({ /** Capture-then-remove, including the item's neighbours, so any refused send * or rejected steer can restore the ORIGINAL item in place even if the run * drains an adjacent entry while the request is in flight. */ - const takeQueued = useRecoilCallback( - ({ snapshot, set }) => - (id: string): QueuedMessageOrigin | undefined => { - const queue = snapshot.getLoadable(store.queuedMessagesByConvoId(queueKey)).getValue(); - const index = queue.findIndex((item) => item.id === id); - if (index < 0) { - return undefined; - } - const item = queue[index]; - let queuedOrigins = queuedOriginsRef.current.get(queueKey); - if (queuedOrigins == null) { - queuedOrigins = new Map(); - queuedOriginsRef.current.set(queueKey, queuedOrigins); - } - let logicalQueue = queue; - for (const pendingOrigin of queuedOrigins.values()) { - logicalQueue = insertQueuedOrigin(logicalQueue, pendingOrigin); - } - const logicalIndex = logicalQueue.findIndex((queued) => queued.id === id); - const origin = { - item, - beforeIds: logicalQueue.slice(0, logicalIndex).map((queued) => queued.id), - afterIds: logicalQueue.slice(logicalIndex + 1).map((queued) => queued.id), - }; - queuedOrigins.set(id, origin); - set(store.queuedMessagesByConvoId(queueKey), (prev) => - prev.filter((item) => item.id !== id), - ); - return origin; - }, - [queueKey], + const takeQueued = useCallback( + (id: string): QueuedMessageOrigin | undefined => { + const queue = jotaiStore.get(queuedMessagesByConvoId(queueKey)); + const index = queue.findIndex((item) => item.id === id); + if (index < 0) { + return undefined; + } + const item = queue[index]; + let queuedOrigins = queuedOriginsRef.current.get(queueKey); + if (queuedOrigins == null) { + queuedOrigins = new Map(); + queuedOriginsRef.current.set(queueKey, queuedOrigins); + } + let logicalQueue = queue; + for (const pendingOrigin of queuedOrigins.values()) { + logicalQueue = insertQueuedOrigin(logicalQueue, pendingOrigin); + } + const logicalIndex = logicalQueue.findIndex((queued) => queued.id === id); + const origin = { + item, + beforeIds: logicalQueue.slice(0, logicalIndex).map((queued) => queued.id), + afterIds: logicalQueue.slice(logicalIndex + 1).map((queued) => queued.id), + }; + queuedOrigins.set(id, origin); + jotaiStore.set(queuedMessagesByConvoId(queueKey), (prev) => + prev.filter((item) => item.id !== id), + ); + return origin; + }, + [queueKey, jotaiStore], ); const releaseQueuedOrigin = useCallback( @@ -1715,16 +1706,15 @@ export default function useSteering({ /** Restore the exact item between its nearest surviving original neighbours, * never duplicated. If every neighbour disappeared, newly front-prioritized * entries stay ahead while ordinary entries created later stay behind. */ - const restoreQueued = useRecoilCallback( - ({ set }) => - (origin: QueuedMessageOrigin) => { - releaseQueuedOrigin(origin); - set(store.queuedMessagesByConvoId(queueKey), (prev) => - canRestoreRecovery(jotaiStore.get(recoveryDispositionsFamily(queueKey)), origin.item) - ? insertQueuedOrigin(prev, origin) - : prev, - ); - }, + const restoreQueued = useCallback( + (origin: QueuedMessageOrigin) => { + releaseQueuedOrigin(origin); + jotaiStore.set(queuedMessagesByConvoId(queueKey), (prev) => + canRestoreRecovery(jotaiStore.get(recoveryDispositionsFamily(queueKey)), origin.item) + ? insertQueuedOrigin(prev, origin) + : prev, + ); + }, [queueKey, releaseQueuedOrigin, jotaiStore], ); @@ -1739,17 +1729,16 @@ export default function useSteering({ * then only inspects the active one's queue; this item would never be looked * at. Park ours alongside it. */ - const rearmDrain = useRecoilCallback( - ({ snapshot, set }) => - (convoId: string, end: RunEnd) => { - const indexArmed = snapshot.getLoadable(store.runEndByIndex(index)).getValue(); - const parkedArmed = snapshot.getLoadable(store.pendingRunEndByConvoId(convoId)).getValue(); - if (isSameRunEpoch(indexArmed, end) || isSameRunEpoch(parkedArmed, end)) { - return; - } - set(store.pendingRunEndByConvoId(convoId), end); - }, - [index], + const rearmDrain = useCallback( + (convoId: string, end: RunEnd) => { + const indexArmed = jotaiStore.get(runEndByIndex(index)); + const parkedArmed = jotaiStore.get(pendingRunEndByConvoId(convoId)); + if (isSameRunEpoch(indexArmed, end) || isSameRunEpoch(parkedArmed, end)) { + return; + } + jotaiStore.set(pendingRunEndByConvoId(convoId), end); + }, + [index, jotaiStore], ); /** @@ -1801,19 +1790,15 @@ export default function useSteering({ [conversationId, convertSteersToQueued, rewakeDrain], ); - const armDrainAfterAbort = useRecoilCallback( - ({ set }) => - () => { - if (activeGenerationCreatedAt == null) { - return; - } - set(store.drainAfterAbortByIndex(index), { - conversationId: queueKey, - generationCreatedAt: activeGenerationCreatedAt, - }); - }, - [index, queueKey, activeGenerationCreatedAt], - ); + const armDrainAfterAbort = useCallback(() => { + if (activeGenerationCreatedAt == null) { + return; + } + jotaiStore.set(drainAfterAbortByIndex(index), { + conversationId: queueKey, + generationCreatedAt: activeGenerationCreatedAt, + }); + }, [index, queueKey, activeGenerationCreatedAt, jotaiStore]); /** POSTs a steer (text + files + quotes; the server merges the quotes into * the model-bound turn at the injection boundary). `context` doubles as the diff --git a/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts b/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts index 7fa4a9ceb8f..8fb0d4f9ca9 100644 --- a/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts +++ b/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts @@ -101,10 +101,11 @@ const mockActiveRunAtom = { key: 'activeRun' }; const mockAbortScrollAtom = { key: 'abortScroll' }; const mockSubmissionAtom = { key: 'submission' }; const mockShowStopButtonAtom = { key: 'showStopButton' }; -const mockRunEndAtom = { key: 'runEnd' }; -const mockDrainAfterAbortAtom = { key: 'drainAfterAbort' }; const mockPendingSteersAtom = { key: 'pendingSteers' }; -const mockQueuedMessagesAtom = { key: 'queuedMessages' }; +/** The follow-up queue a restored submission lands back in. */ +const queuedIn = (conversationId: string) => + getDefaultStore().get(queuedMessagesByConvoId(conversationId)); + const mockSetActiveRun = jest.fn(); const mockSetAbortScroll = jest.fn(); const mockSetSubmission = jest.fn(); @@ -117,10 +118,8 @@ const mockSeedSteerChips = jest.fn(); const mockSettleAppliedSteerParts = jest.fn(); const mockConvertLocalSteersToQueued = jest.fn(); const mockUpdateGenerationEpoch = jest.fn(); -const mockRestoreQueuedSubmission = jest.fn(); let mockRecoilCallbackIndex = 0; const mockRecoilCallbacks = [ - mockRestoreQueuedSubmission, mockResolveSteerChip, mockUpdateSteerChips, mockSeedSteerChips, @@ -144,12 +143,6 @@ const mockUseSetRecoilStateMock = jest.fn((atom: unknown) => { if (atom === mockShowStopButtonAtom) { return mockSetShowStopButton; } - if (atom === mockRunEndAtom) { - return mockSetRunEnd; - } - if (atom === mockDrainAfterAbortAtom) { - return mockSetDrainAfterAbort; - } return jest.fn(); }); function mockUseSetRecoilState(atom: unknown) { @@ -188,6 +181,25 @@ jest.mock('recoil', () => ({ useRecoilCallback: mockUseRecoilCallback, })); +/** The run-end and interrupt-drain signals are Jotai atoms; their pane setters are swapped for + * spies, the same way the Recoil setters above are. */ +jest.mock('jotai', () => { + const actual = jest.requireActual('jotai'); + return { + ...actual, + useSetAtom: (atom: unknown) => { + const queue = jest.requireActual('~/store/queue'); + if (atom === queue.runEndByIndex(0)) { + return mockSetRunEnd; + } + if (atom === queue.drainAfterAbortByIndex(0)) { + return mockSetDrainAfterAbort; + } + return actual.useSetAtom(atom); + }, + }; +}); + jest.mock('~/store', () => ({ __esModule: true, default: { @@ -195,10 +207,7 @@ jest.mock('~/store', () => ({ abortScrollFamily: jest.fn(() => mockAbortScrollAtom), submissionByIndex: jest.fn(() => mockSubmissionAtom), showStopButtonByIndex: jest.fn(() => mockShowStopButtonAtom), - runEndByIndex: jest.fn(() => mockRunEndAtom), - drainAfterAbortByIndex: jest.fn(() => mockDrainAfterAbortAtom), pendingSteersByConvoId: jest.fn(() => mockPendingSteersAtom), - queuedMessagesByConvoId: jest.fn(() => mockQueuedMessagesAtom), }, })); @@ -311,6 +320,7 @@ import useResumableSSE, { ABORT_SWEEP_STATUSES, } from '~/hooks/SSE/useResumableSSE'; import useSSE from '~/hooks/SSE/useSSE'; +import { queuedMessagesByConvoId, resetQueueFamilies } from '~/store/queue'; const CONV_ID = 'conv-abc-123'; @@ -429,7 +439,7 @@ describe('useResumableSSE', () => { mockSettleAppliedSteerParts.mockClear(); mockConvertLocalSteersToQueued.mockClear(); mockUpdateGenerationEpoch.mockClear(); - mockRestoreQueuedSubmission.mockClear(); + resetQueueFamilies(); mockRecoilCallbackIndex = 0; mockConvertSteersToQueued.mockClear(); mockFetchStreamStatus.mockReset(); @@ -3316,7 +3326,7 @@ describe('useResumableSSE', () => { expect(getDefaultStore().get(recoveryDispositionsFamily(CONV_ID))).toEqual({ source: 'blocked', }); - expect(mockRestoreQueuedSubmission).toHaveBeenCalledWith(submission); + expect(queuedIn(CONV_ID)).toEqual([item]); expect(mockConvertSteersToQueued).not.toHaveBeenCalled(); expect(request.post).toHaveBeenCalledTimes(1); unmount(); @@ -3356,7 +3366,7 @@ describe('useResumableSSE', () => { await waitFor(() => expect(mockSetSubmission).toHaveBeenCalledWith(null)); - expect(mockRestoreQueuedSubmission).toHaveBeenCalledWith(submission); + expect(queuedIn(CONV_ID)).toEqual([queuedMessageOrigin.item]); expect(mockSSEInstances).toHaveLength(0); unmount(); }); @@ -3397,8 +3407,9 @@ describe('useResumableSSE', () => { await waitFor(() => expect(mockSetSubmission).toHaveBeenCalledWith(null)); expect(request.post).toHaveBeenCalledTimes(1); - expect(mockRestoreQueuedSubmission).toHaveBeenNthCalledWith(1, submission); - expect(mockRestoreQueuedSubmission).toHaveBeenLastCalledWith(submission, 2000); + expect(queuedIn(CONV_ID)).toEqual([ + { ...queuedMessageOrigin.item, expectedPredecessorCreatedAt: 2000 }, + ]); expect(mockSSEInstances).toHaveLength(0); expect(mockSetQueryData).toHaveBeenCalledWith( ['streamStatus', CONV_ID], @@ -3464,8 +3475,9 @@ describe('useResumableSSE', () => { await waitFor(() => expect(mockSetSubmission).toHaveBeenCalledWith(null)); expect(request.post).toHaveBeenCalledTimes(1); - expect(mockRestoreQueuedSubmission).toHaveBeenNthCalledWith(1, submission); - expect(mockRestoreQueuedSubmission).toHaveBeenLastCalledWith(submission, 2000); + expect(queuedIn(CONV_ID)).toEqual([ + { ...queuedMessageOrigin.item, expectedPredecessorCreatedAt: 2000 }, + ]); expect(mockInvalidateQueries).toHaveBeenCalledWith({ queryKey: [QueryKeys.messages, CONV_ID], refetchType: 'all', @@ -3517,7 +3529,7 @@ describe('useResumableSSE', () => { generationProtocolVersion: 2, allowPreviouslyConvertedIds: ['source-steer'], }); - expect(mockRestoreQueuedSubmission).not.toHaveBeenCalled(); + expect(queuedIn(CONV_ID)).toEqual([]); unmount(); }); diff --git a/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx b/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx index 1cdb74051e9..1eee9a0744c 100644 --- a/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx +++ b/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx @@ -9,6 +9,7 @@ import type { ReactNode } from 'react'; import type { PendingSteer, QueuedMessage } from '~/store/families'; import { siblingIdxFamily, siblingKey } from '~/components/Chat/Messages/Thread/state'; import { pendingApprovalActionFamily } from '~/components/Chat/approval/state'; +import { queuedMessagesByConvoId, resetQueueFamilies } from '~/store/queue'; import { agentQueuedTurnsQueryKey } from '~/data-provider/SSE/queuedTurns'; import { revealedQueuedTurnFamily } from '~/store/steer'; import useResumeOnLoad from '../useResumeOnLoad'; @@ -167,7 +168,7 @@ function renderUseResumeOnLoad({ return null; }; const QueuedMessagesProbe = () => { - const queued = useRecoilValue(store.queuedMessagesByConvoId(conversationId)); + const queued = useAtomValue(queuedMessagesByConvoId(conversationId)); onQueuedMessages?.(queued); return null; }; @@ -209,6 +210,8 @@ function renderUseResumeOnLoad({ }; } +beforeEach(() => resetQueueFamilies()); + describe('useResumeOnLoad', () => { beforeEach(() => { jest.spyOn(console, 'log').mockImplementation(() => undefined); diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index 4a6852ef723..7d4b8db494f 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/client/src/hooks/SSE/useResumableSSE.ts @@ -94,6 +94,7 @@ import useEventHandlers, { buildCreatedInitialResponse, keepLocalCodeApprovalMode, } from './useEventHandlers'; +import { drainAfterAbortByIndex, queuedMessagesByConvoId, runEndByIndex } from '~/store/queue'; import { pendingApprovalActionFamily } from '~/components/Chat/approval/state'; import { useChatTransport } from '~/Providers/ChatTransportContext'; import useSteerConvert from '~/hooks/Chat/useSteerConvert'; @@ -1009,23 +1010,19 @@ export default function useResumableSSE( * id and is therefore shared by every generation within it. */ const prefixStateGenerationIdRef = useRef(null); - const restoreQueuedSubmission = useRecoilCallback( - ({ set }) => - (failedSubmission: TSubmission, expectedPredecessorCreatedAt?: number) => { - const conversationId = failedSubmission.conversation?.conversationId; - const origin = failedSubmission.queuedMessageOrigin as QueuedMessageOrigin | undefined; - if (!conversationId || origin == null) { - return; - } - set(store.queuedMessagesByConvoId(conversationId), (prev) => - canRestoreRecovery( - jotaiStore.get(recoveryDispositionsFamily(conversationId)), - origin.item, - ) - ? insertQueuedOrigin(prev, origin, expectedPredecessorCreatedAt) - : prev, - ); - }, + const restoreQueuedSubmission = useCallback( + (failedSubmission: TSubmission, expectedPredecessorCreatedAt?: number) => { + const conversationId = failedSubmission.conversation?.conversationId; + const origin = failedSubmission.queuedMessageOrigin as QueuedMessageOrigin | undefined; + if (!conversationId || origin == null) { + return; + } + jotaiStore.set(queuedMessagesByConvoId(conversationId), (prev) => + canRestoreRecovery(jotaiStore.get(recoveryDispositionsFamily(conversationId)), origin.item) + ? insertQueuedOrigin(prev, origin, expectedPredecessorCreatedAt) + : prev, + ); + }, [jotaiStore], ); @@ -1080,7 +1077,7 @@ export default function useResumableSSE( * event was waiting for the response placeholder. The applied event * is authoritative, so evict that recovery copy before it can be * drained as a duplicate follow-up. */ - set(store.queuedMessagesByConvoId(conversationId), (prev) => + jotaiStore.set(queuedMessagesByConvoId(conversationId), (prev) => prev.some( (item) => (item.recoverySteerId != null && settledIds.includes(item.recoverySteerId)) || @@ -1358,8 +1355,8 @@ export default function useResumableSSE( [], ); - const setRunEnd = useSetRecoilState(store.runEndByIndex(runIndex)); - const setDrainAfterAbort = useSetRecoilState(store.drainAfterAbortByIndex(runIndex)); + const setRunEnd = useSetAtom(runEndByIndex(runIndex)); + const setDrainAfterAbort = useSetAtom(drainAfterAbortByIndex(runIndex)); const clearDrainAfterAbort = useCallback( (conversationId: string, generationCreatedAt?: number) => { if (generationCreatedAt == null) { diff --git a/client/src/store/families.ts b/client/src/store/families.ts index 466f5575b0f..7f57ec9497d 100644 --- a/client/src/store/families.ts +++ b/client/src/store/families.ts @@ -472,16 +472,6 @@ export type QueuedMessageOrigin = { afterIds: string[]; }; -/** - * Per-conversation client-side queue of follow-up messages. Drained one per - * run completion by `useQueueDrain` (each dequeued message starts a normal - * turn whose own final event drains the next). - */ -const queuedMessagesByConvoId = atomFamily({ - key: 'queuedMessagesByConvoId', - default: [], -}); - export type SettledQueuedTurnReceipt = { clientRequestId: string; status: 'admitted' | 'admitted_pending_boundary' | 'indeterminate' | 'cancelled' | 'dead'; @@ -490,21 +480,6 @@ export type SettledQueuedTurnReceipt = { boundaryConsumed?: boolean; }; -/** Monotonic client knowledge of terminal server queue receipts. Admission - * records preserve boundary multiplicity by request identity. Other terminal - * records exist only while their original enqueue callback is outstanding. */ -const settledQueuedTurnReceiptsByConvoId = atomFamily({ - key: 'settledQueuedTurnReceiptsByConvoId', - default: [], -}); - -/** Enqueue callbacks that can still race newer GET/cancellation evidence. - * Entries retire as soon as that one callback settles. */ -const pendingQueuedTurnEnqueueIdsByConvoId = atomFamily({ - key: 'pendingQueuedTurnEnqueueIdsByConvoId', - default: [], -}); - /** * One-shot run-termination signal written by the SSE final/error handlers and * consumed (reset to null) by `useQueueDrain`. Keyed by chat index like @@ -526,83 +501,11 @@ export type RunEnd = { interruptArmed?: boolean; }; -/** A pane can receive A's terminal frame after the user has navigated to and - * started B. Keep each terminal epoch until the queue drain has either parked - * or consumed it; a single replaceable slot loses A when B finishes first. */ -const runEndsByIndex = atomFamily({ - key: 'runEndsByIndex', - default: [], -}); - -/** Preserve the original nullable one-shot API for stream writers while the - * backing state retains every not-yet-consumed terminal epoch. Writing null - * consumes only the visible (oldest) signal. */ -const runEndByIndex = selectorFamily({ - key: 'runEndByIndex', - get: - (index) => - ({ get }) => - get(runEndsByIndex(index))[0] ?? null, - set: - (index) => - ({ set }, value) => { - if (value instanceof DefaultValue) { - set(runEndsByIndex(index), []); - return; - } - if (value == null) { - set(runEndsByIndex(index), (prev) => prev.slice(1)); - return; - } - set(runEndsByIndex(index), (prev) => [...prev, value]); - }, -}); - -/** Foreign terminal epochs are moved off the shared pane immediately. This - * per-conversation carrier is queued for the same reason as the pane carrier: - * successive epochs cannot overwrite one another while the chat is hidden. */ -const pendingRunEndsByConvoId = atomFamily({ - key: 'pendingRunEndsByConvoId', - default: [], -}); - -const pendingRunEndByConvoId = selectorFamily({ - key: 'pendingRunEndByConvoId', - get: - (conversationId) => - ({ get }) => - get(pendingRunEndsByConvoId(conversationId))[0] ?? null, - set: - (conversationId) => - ({ set }, value) => { - if (value instanceof DefaultValue) { - set(pendingRunEndsByConvoId(conversationId), []); - return; - } - if (value == null) { - set(pendingRunEndsByConvoId(conversationId), (prev) => prev.slice(1)); - return; - } - set(pendingRunEndsByConvoId(conversationId), (prev) => [...prev, value]); - }, -}); - export type DrainAfterAbort = { conversationId: string; generationCreatedAt: number; }; -/** - * One-shot override armed by "interrupt & send": the next `aborted` run-end - * for the exact conversation generation drains the queue exactly once (a - * plain Stop press leaves queued chips for manual send). `false` remains the - * clear value used by stream reconciliation paths. - */ -const drainAfterAbortByIndex = atomFamily({ - key: 'drainAfterAbortByIndex', - default: false, -}); - /** * Server steer ids whose `on_steer_applied` event already landed. The 202 ACK * and the SSE ride different connections, so the applied event can arrive @@ -807,12 +710,6 @@ export default { pendingManualSkillsByConvoId, pendingQuotesByConvoId, pendingSteersByConvoId, - queuedMessagesByConvoId, - settledQueuedTurnReceiptsByConvoId, - pendingQueuedTurnEnqueueIdsByConvoId, - runEndByIndex, - pendingRunEndByConvoId, - drainAfterAbortByIndex, appliedSteerIdsByConvoId, acceptedSteerClientIdsByConvoId, activeGenerationCreatedAtByConvoId, diff --git a/client/src/store/queue.ts b/client/src/store/queue.ts new file mode 100644 index 00000000000..39ad6a67a3c --- /dev/null +++ b/client/src/store/queue.ts @@ -0,0 +1,97 @@ +import { atom } from 'jotai'; +import { atomFamily, RESET } from 'jotai/utils'; +import type { PrimitiveAtom } from 'jotai'; +import type { SettledQueuedTurnReceipt, DrainAfterAbort, QueuedMessage, RunEnd } from './families'; + +/** + * Per-conversation client-side queue of follow-up messages. Drained one per + * run completion by `useQueueDrain` (each dequeued message starts a normal + * turn whose own final event drains the next). + */ +export const queuedMessagesByConvoId = atomFamily((_conversationId: string) => + atom([]), +); + +/** Monotonic client knowledge of terminal server queue receipts. Admission + * records preserve boundary multiplicity by request identity. Other terminal + * records exist only while their original enqueue callback is outstanding. */ +export const settledQueuedTurnReceiptsByConvoId = atomFamily((_conversationId: string) => + atom([]), +); + +/** Enqueue callbacks that can still race newer GET/cancellation evidence. + * Entries retire as soon as that one callback settles. */ +export const pendingQueuedTurnEnqueueIdsByConvoId = atomFamily((_conversationId: string) => + atom([]), +); + +/** The oldest terminal epoch in a queue of them, behind the nullable one-shot API stream writers + * use: writing a signal appends it, writing `null` consumes only the visible (oldest) one, and + * `RESET` clears them all. */ +const runEndQueue = (signals: PrimitiveAtom) => + atom( + (get) => get(signals)[0] ?? null, + (_get, set, value: RunEnd | null | typeof RESET) => { + if (value === RESET) { + set(signals, []); + return; + } + if (value == null) { + set(signals, (prev) => prev.slice(1)); + return; + } + set(signals, (prev) => [...prev, value]); + }, + ); + +/** A pane can receive A's terminal frame after the user has navigated to and + * started B. Keep each terminal epoch until the queue drain has either parked + * or consumed it; a single replaceable slot loses A when B finishes first. */ +const runEndsByIndex = atomFamily((_index: string | number) => atom([])); + +/** One-shot run-termination signal for a pane, written by the SSE final/error handlers and + * consumed by `useQueueDrain`. */ +export const runEndByIndex = atomFamily((index: string | number) => + runEndQueue(runEndsByIndex(index)), +); + +/** Foreign terminal epochs are moved off the shared pane immediately. This + * per-conversation carrier is queued for the same reason as the pane carrier: + * successive epochs cannot overwrite one another while the chat is hidden. */ +const pendingRunEndsByConvoId = atomFamily((_conversationId: string) => atom([])); + +export const pendingRunEndByConvoId = atomFamily((conversationId: string) => + runEndQueue(pendingRunEndsByConvoId(conversationId)), +); + +/** + * One-shot override armed by "interrupt & send": the next `aborted` run-end + * for the exact conversation generation drains the queue exactly once (a + * plain Stop press leaves queued chips for manual send). `false` remains the + * clear value used by stream reconciliation paths. + */ +export const drainAfterAbortByIndex = atomFamily((_index: string | number) => + atom(false), +); + +const clearFamily = (family: { + getParams(): Iterable; + remove(param: Param): void; +}) => { + for (const param of [...family.getParams()]) { + family.remove(param); + } +}; + +/** Drops every per-key queue atom, so the next read of any key starts from its default. Isolates + * tests that share the default store. */ +export function resetQueueFamilies(): void { + clearFamily(queuedMessagesByConvoId); + clearFamily(settledQueuedTurnReceiptsByConvoId); + clearFamily(pendingQueuedTurnEnqueueIdsByConvoId); + clearFamily(pendingRunEndsByConvoId); + clearFamily(pendingRunEndByConvoId); + clearFamily(drainAfterAbortByIndex); + clearFamily(runEndsByIndex); + clearFamily(runEndByIndex); +} diff --git a/e2e/specs/mock/scenarios/queue-owners.spec.ts b/e2e/specs/mock/scenarios/queue-owners.spec.ts new file mode 100644 index 00000000000..acc4271c5d0 --- /dev/null +++ b/e2e/specs/mock/scenarios/queue-owners.spec.ts @@ -0,0 +1,129 @@ +import { expect, test } from '@playwright/test'; +import type { Page } from '@playwright/test'; +import type { TMessage } from 'librechat-data-provider'; +import { + selectMockEndpoint, + getAccessToken, + messagesView, + sendMessage, + replyPrompt, + replyText, + fetchJson, + MOCK_REPLY_TEXT, + MOCK_ENDPOINTS, + NEW_CHAT_PATH, +} from '../helpers'; + +/** + * The follow-up queue, its run-end signals and the interrupt-drain flag are chat-owned Jotai + * state. These scenarios drive the paths that read and write each of them end to end. + */ + +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)}`; + +async function establishConversation(page: Page, label: string): Promise { + const setup = await sendMessage(page, replyPrompt(label)); + expect(setup.ok()).toBeTruthy(); + await expect(messagesView(page).getByText(replyText(label))).toBeVisible({ timeout: 30000 }); + await expect(page).toHaveURL(/\/c\/[0-9a-fA-F-]{36}$/, { timeout: 15000 }); + return new URL(page.url()).pathname.split('/').pop() ?? ''; +} + +async function typeDuringRun(page: Page, text: string) { + const input = messageInput(page); + await input.click(); + await input.fill(text); + await expect(duringRunSendButton(page)).toBeVisible({ timeout: 5000 }); +} + +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. */ + await page.addInitScript(() => { + localStorage.setItem('duringRunDefaultAction', JSON.stringify('steer')); + }); + }); + + test('interrupt and send stops the run and sends the follow-up next @scenario:interrupt-and-send-drains-the-follow-up', async ({ + page, + }) => { + test.setTimeout(120000); + const label = uniqueLabel('interrupt'); + const followUp = `Interrupt follow-up ${label}`; + + await page.goto(NEW_CHAT_PATH, { timeout: 10000 }); + await selectMockEndpoint(page, MOCK_ENDPOINTS[0]); + await establishConversation(page, `interrupt-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('Alt+Enter'); + + await expect(messageTurns(page)).toHaveCount(6, { timeout: 60000 }); + await expect(messageTurns(page).nth(4)).toContainText(followUp); + await expect(messageTurns(page).nth(5)).toContainText(MOCK_REPLY_TEXT, { timeout: 30000 }); + 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. */ + 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); + 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); + }); +});