diff --git a/client/src/hooks/Chat/__tests__/facade.spec.tsx b/client/src/hooks/Chat/__tests__/facade.spec.tsx index e6bf42298b7..6965cdb82dc 100644 --- a/client/src/hooks/Chat/__tests__/facade.spec.tsx +++ b/client/src/hooks/Chat/__tests__/facade.spec.tsx @@ -1,11 +1,14 @@ import React from 'react'; import { act, render, renderHook } from '@testing-library/react'; -import { QueryKeys, ContentTypes } from 'librechat-data-provider'; +import { QueryKeys, Constants, ContentTypes } from 'librechat-data-provider'; import { QueryClient, QueryClientProvider, useQuery } from '@tanstack/react-query'; import type { TConversation, TMessage, TMessageContentParts } from 'librechat-data-provider'; import type { ChatContract } from '../contract'; +import type { JotaiStore } from 'test/harness'; import { ChatContext } from '~/Providers/ChatContext'; import { useChat, useChatActions } from '../facade'; +import { IsolatedAtomStore } from 'test/harness'; +import { resumeRequestsAtom } from '../resume'; const userMessage: TMessage = { messageId: 'user-1', @@ -637,6 +640,104 @@ describe('useChat', () => { expect(contract.stopGenerating).toHaveBeenCalledTimes(1); }); + it('sends an AI SDK user message as the ask call it describes', () => { + const contract = createContract(); + const { result } = renderChat(contract); + + result.current.sendMessage( + { + parts: [ + { type: 'text', text: 'Hel' }, + { type: 'text', text: 'lo' }, + ], + metadata: { parentMessageId: 'user-1' }, + }, + { isRegenerate: false }, + ); + + expect(contract.ask).toHaveBeenCalledWith( + { text: 'Hello', parentMessageId: 'user-1' }, + { isRegenerate: false }, + ); + }); + + it('attaches an AI SDK message with a null parent at the root', () => { + const contract = createContract(); + const { result } = renderChat(contract); + + result.current.sendMessage({ + parts: [{ type: 'text', text: 'From the top' }], + metadata: { parentMessageId: null }, + }); + + expect(contract.ask).toHaveBeenCalledWith( + { text: 'From the top', parentMessageId: Constants.NO_PARENT }, + undefined, + ); + }); + + it('reports a refused send to the caller', () => { + const contract = createContract({ ask: jest.fn(() => false as const) }); + const { result } = renderChat(contract); + + expect(result.current.sendMessage({ parts: [{ type: 'text', text: 'Hi' }] })).toBe(false); + }); + + it('rejects stop when the stop request fails', async () => { + const failure = new Error('abort failed'); + const contract = createContract({ + isSubmitting: true, + stopGenerating: jest.fn(() => Promise.reject(failure)), + }); + const { result } = renderChat(contract); + + await expect(result.current.stop()).rejects.toBe(failure); + }); + + /** Renders `useChat` under its own atom store, seeded with another pane's pending request. */ + const renderChatWithRequests = (contract: ChatContract) => { + let atoms: JotaiStore | undefined; + const queryClient = new QueryClient(); + const view = renderHook(() => useChat(), { + wrapper: ({ children }) => ( + + { + atoms = store; + store.set(resumeRequestsAtom, new Set(['convo-2'])); + }} + > + {children} + + + ), + }); + const pending = () => [...(atoms?.get(resumeRequestsAtom) ?? [])]; + return { ...view, pending }; + }; + + it('requests a resume of the chat it reads', async () => { + const { result, pending } = renderChatWithRequests(createContract()); + + await result.current.resumeStream(); + await result.current.resumeStream(); + + expect(pending()).toEqual(['convo-2', 'convo-1']); + }); + + it('requests no resume for a chat that has no conversation yet', async () => { + const { result, pending } = renderChatWithRequests( + createContract({ + messagesKey: 'new', + conversation: { conversationId: 'new' } as TConversation, + }), + ); + + await expect(result.current.resumeStream()).resolves.toBeUndefined(); + expect(result.current.id).toBe('new'); + expect(pending()).toEqual(['convo-2']); + }); + it('writes UI messages back onto the stored messages', () => { const answered = response({ text: 'Hello', @@ -716,7 +817,8 @@ describe('useChatActions', () => { it('forwards its actions to the contract', () => { const { result, contract } = renderActions([userMessage, response()]); - expect(result.current.sendMessage).toBe(contract.ask); + result.current.sendMessage({ text: 'Hello' }); + expect(contract.ask).toHaveBeenCalledWith({ text: 'Hello' }, undefined); expect(result.current.stop).toBe(contract.stopGenerating); result.current.regenerate(); expect(contract.regenerate).toHaveBeenCalledWith({ diff --git a/client/src/hooks/Chat/__tests__/transport.spec.tsx b/client/src/hooks/Chat/__tests__/transport.spec.tsx index 8be4e480fa6..0484f8c2e2d 100644 --- a/client/src/hooks/Chat/__tests__/transport.spec.tsx +++ b/client/src/hooks/Chat/__tests__/transport.spec.tsx @@ -1,9 +1,9 @@ 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'; +import { useAtomValue, getDefaultStore } from 'jotai'; +import { QueryKeys, request } from 'librechat-data-provider'; import { act, renderHook, waitFor } from '@testing-library/react'; +import { RecoilRoot, useRecoilValue, useSetRecoilState } from 'recoil'; import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; import type { TEnqueueAgentQueuedTurnRequest, @@ -13,15 +13,22 @@ import type { ChatTransportRequest, } from 'librechat-data-provider'; import type { MutableSnapshot } from 'recoil'; +import type { StreamStatusResponse } from '~/data-provider'; +import type { ChatContract } from '~/hooks/Chat/contract'; import type { Transport } from '~/hooks/Chat/contract'; import type { PendingSteer } from '~/store/families'; import { queuedMessagesByConvoId, resetQueueFamilies } from '~/hooks/Chat/queue'; import { ChatTransportContext } from '~/Providers/ChatTransportContext'; +import { ChatContext, useChatContext } from '~/Providers/ChatContext'; +import { startupConfigKey } from '~/data-provider/Endpoints/queries'; import { useSteerReclaim } from '~/hooks/Chat/useSteerCancel'; import useSteerEscalate from '~/hooks/Chat/useSteerEscalate'; import useResumableSSE from '~/hooks/SSE/useResumableSSE'; +import useResumeOnLoad from '~/hooks/SSE/useResumeOnLoad'; import useChatHelpers from '~/hooks/Chat/useChatHelpers'; +import { resumeRequestsAtom } from '~/hooks/Chat/resume'; import useSteering from '~/hooks/Chat/useSteering'; +import { useChat } from '~/hooks/Chat/facade'; import useSSE from '~/hooks/SSE/useSSE'; import store from '~/store'; @@ -492,4 +499,189 @@ describe('chat transport boundary', () => { expect(result.current.queue).toEqual([expect.objectContaining({ text: 'hold me here' })]); }); }); + + describe('resume', () => { + /** The host builds the pane's chat contract, as the chat view does. */ + function ChatHost({ children }: { children: React.ReactNode }) { + const helpers: ChatContract = useChatHelpers(0, 'convo-1'); + return {children}; + } + + /** The resume-on-load path and the stream hook a chat view mounts, read through `useChat`. */ + const useResumablePane = () => { + const helpers = useChatContext(); + useResumeOnLoad('convo-1', helpers.getMessages, 0, true); + const submission = useRecoilValue(store.submissionByIndex(0)); + useResumableSSE(submission, helpers, false, 0); + const showConversation = useSetRecoilState(store.conversationByIndex(0)); + return { ...useChat(), showConversation }; + }; + + const seedConversation = ({ set }: MutableSnapshot) => + set(store.conversationByIndex(0), { + conversationId: 'convo-1', + endpoint: 'agents', + } as TConversation); + + let status: StreamStatusResponse; + /** A running job as the status route reports it, with the turn it answers. */ + const runningStatus = (): StreamStatusResponse => ({ + active: true, + streamId: 'convo-1', + status: 'running', + createdAt: 2000, + generationProtocolVersion: 2, + resumeState: { + runSteps: [], + aggregatedContent: [], + userMessage: { + messageId: 'msg-1', + parentMessageId: '00000000-0000-0000-0000-000000000000', + conversationId: 'convo-1', + text: 'Hello', + }, + responseMessageId: 'resp-1', + conversationId: 'convo-1', + } as StreamStatusResponse['resumeState'], + }); + const statusReads = () => + (request.get as jest.Mock).mock.calls.filter(([url]) => String(url).includes('/status/')) + .length; + + beforeEach(() => { + status = { active: false }; + /** The status read is the server's answer to "is anything running"; the stream itself + * comes from the fake transport. */ + jest + .spyOn(request, 'get') + .mockImplementation(async (url: string) => + url.includes('/api/agents/chat/status/') ? status : [], + ); + }); + + afterEach(() => { + jest.restoreAllMocks(); + }); + + /** Resume waits for the startup config, which decides a rebuilt turn's retention. */ + const seedStartupConfig = (queryClient: QueryClient) => { + queryClient.setQueryData(startupConfigKey(false), {}); + /** The turn the running generation answers, as history holds it on reload. */ + queryClient.setQueryData([QueryKeys.messages, 'convo-1'], [buildSubmission().userMessage]); + }; + + const renderPane = ( + transport: Transport, + seed: (snapshot: MutableSnapshot) => void = seedConversation, + ) => { + const Wrapper = createWrapper(transport, seed, seedStartupConfig); + return renderHook(useResumablePane, { + wrapper: ({ children }) => ( + + {children} + + ), + }); + }; + + it('reattaches to a running generation through the host transport', async () => { + const fake = createFakeTransport(); + const { result } = renderPane(fake.transport); + await waitFor(() => expect(statusReads()).toBe(1)); + expect(fake.streams).toHaveLength(0); + + status = runningStatus(); + await act(async () => { + await result.current.resumeStream(); + }); + + await waitFor(() => expect(fake.streams).toHaveLength(1)); + expect(fake.transport.stream).toHaveBeenCalledWith({ token: 'test-token' }); + expect(fake.streams[0].url).toContain('/api/agents/chat/stream/convo-1'); + expect(fake.streams[0].url).toContain('resume=true'); + expect(fake.streams[0].url).toContain('generationCreatedAt=2000'); + expect(fake.transport.start).not.toHaveBeenCalled(); + }); + + it('re-reads the status and stays detached when nothing is running', async () => { + const fake = createFakeTransport(); + const { result } = renderPane(fake.transport); + await waitFor(() => expect(statusReads()).toBe(1)); + + await act(async () => { + await result.current.resumeStream(); + }); + + await waitFor(() => expect(statusReads()).toBe(2)); + expect(fake.streams).toHaveLength(0); + expect(result.current.status).toBe('ready'); + }); + + it('holds a request until the pane has loaded the conversation the route names', async () => { + const store$ = getDefaultStore(); + const fake = createFakeTransport(); + /** The pane still shows the Assistants conversation it is navigating away from. */ + const { result } = renderPane(fake.transport, ({ set }) => + set(store.conversationByIndex(0), { + conversationId: 'assistants-convo', + endpoint: 'assistants', + } as TConversation), + ); + status = runningStatus(); + + await act(async () => { + await result.current.resumeStream(); + }); + expect([...store$.get(resumeRequestsAtom)]).toEqual(['convo-1']); + expect(fake.streams).toHaveLength(0); + + act(() => + result.current.showConversation({ + conversationId: 'convo-1', + endpoint: 'agents', + } as TConversation), + ); + + await waitFor(() => expect(fake.streams).toHaveLength(1)); + expect(fake.streams[0].url).toContain('resume=true'); + expect([...store$.get(resumeRequestsAtom)]).toEqual([]); + }); + + it("answers its own conversation's request and leaves another pane's pending", async () => { + const store = getDefaultStore(); + const fake = createFakeTransport(); + const { result } = renderPane(fake.transport); + await waitFor(() => expect(statusReads()).toBe(1)); + + await act(async () => { + /** Another pane asks in the same tick, before any effect runs. */ + store.set(resumeRequestsAtom, (pending) => new Set(pending).add('convo-2')); + await result.current.resumeStream(); + }); + + await waitFor(() => expect(statusReads()).toBe(2)); + expect([...store.get(resumeRequestsAtom)]).toEqual(['convo-2']); + store.set(resumeRequestsAtom, new Set()); + }); + + it('reports a reattached stream that fails as an error', async () => { + const fake = createFakeTransport(); + const { result } = renderPane(fake.transport); + await waitFor(() => expect(statusReads()).toBe(1)); + status = runningStatus(); + await act(async () => { + await result.current.resumeStream(); + }); + await waitFor(() => expect(fake.streams).toHaveLength(1)); + + /** The run ends with the error, and the teardown re-reads the status to confirm it. */ + status = { active: false, status: 'error', createdAt: 2000, generationProtocolVersion: 2 }; + act(() => + fake.streams[0].options.onEvent({ type: 'error', data: { message: 'Generation failed' } }), + ); + + await waitFor(() => expect(result.current.status).toBe('error')); + expect(result.current.error?.message).toContain('Generation failed'); + }); + }); }); diff --git a/client/src/hooks/Chat/facade.ts b/client/src/hooks/Chat/facade.ts index 6db5ac2857b..328a7cb3603 100644 --- a/client/src/hooks/Chat/facade.ts +++ b/client/src/hooks/Chat/facade.ts @@ -1,4 +1,5 @@ import { useRef, useMemo, useCallback, useSyncExternalStore } from 'react'; +import { useStore } from 'jotai'; import { hashQueryKey, useQueryClient } from '@tanstack/react-query'; import { QueryKeys, @@ -7,23 +8,50 @@ import { fromUIMessage, toUIMessage, } from 'librechat-data-provider'; -import type { TMessage, UIMessage, TAttachment, UIMappingOptions } from 'librechat-data-provider'; +import type { + TMessage, + UIMessage, + UITextPart, + TAttachment, + UIMappingOptions, + UIMessageMetadata, +} from 'librechat-data-provider'; import type { QueryCacheNotifyEvent } from '@tanstack/react-query'; +import type { TAskProps, TOptions } from '~/common'; import type { ChatContract } from './contract'; -import type { TAskFunction } from '~/common'; import { isMemoryFailureOutput } from '~/components/Chat/Messages/Content/Parts/MemoryCall'; import { getToolMeta } from '~/components/Chat/Messages/Content/outcome'; import { useChatContext } from '~/Providers/ChatContext'; import { isEmptyContentPart } from '~/utils/messages'; +import { resumeRequestsAtom } from './resume'; import { mapAttachments } from '~/utils/map'; /** AI SDK `ChatStatus`. */ export type ChatStatus = 'submitted' | 'streaming' | 'ready' | 'error'; +/** + * A user turn in the AI SDK's `sendMessage` shape: its text parts, plus the message it attaches + * under. It goes to the facade's own chat, as AI SDK `useChat` does, so it names no conversation. + * File parts are left out because the turn takes its files from the composer, not the message. + */ +export type SendMessageInput = { + parts: UITextPart[]; + metadata?: Partial>; +}; + +/** + * Submits a turn. Takes the contract's `ask` arguments, or a user message in the AI SDK shape, + * which is sent as the `ask` call it describes; `{ text }` is already both. + */ +export type SendMessage = ( + message: TAskProps | SendMessageInput, + options?: TOptions, +) => false | void; + /** * The `useChat` surface of `@ai-sdk/react@4.0.117` (`UseChatHelpers`), over LibreChat's chat - * contract. Members the contract cannot back (`resumeStream`, `addToolOutput`, `clearError`) - * are left out rather than stubbed. + * contract. Members the contract cannot back (`addToolOutput`, `clearError`) are left out rather + * than stubbed. */ export type UseChatHelpers = { /** The conversation the messages are read from (the contract's `messagesKey`); AI SDK's chat id. */ @@ -33,11 +61,19 @@ export type UseChatHelpers = { status: ChatStatus; /** Set when the latest message is an error; LibreChat reports errors as messages. */ error: Error | undefined; - /** Submits a turn: the contract's `ask`, called with the same arguments. */ - sendMessage: TAskFunction; + /** Submits a turn through the contract's `ask`; see {@link SendMessage}. */ + sendMessage: SendMessage; /** Regenerates the response to `messageId`, or the latest message of the branch. */ regenerate: (options?: { messageId?: string }) => void; stop: () => Promise; + /** + * Reattaches this chat to its running generation, if the server has one. The request is + * answered by the resume-on-load path the chat view mounts for its pane, which checks the + * stream status and attaches through the host transport. It resolves once requested. It does + * nothing for a new chat or one already attached, and a request for a conversation no chat view + * shows waits until one does. + */ + resumeStream: () => Promise; /** * Writes messages back to the cache, keeping the stored fields the UI view omits. A message * with no stored counterpart joins the active conversation under the message before it, or @@ -312,9 +348,43 @@ const findLatest = (stored: TMessage[] | undefined, latestMessageId: string | un /** The `useChat` members that need no message list. */ export type ChatActions = Pick< UseChatHelpers, - 'id' | 'status' | 'sendMessage' | 'regenerate' | 'stop' + 'id' | 'status' | 'sendMessage' | 'regenerate' | 'stop' | 'resumeStream' >; +/** `ask` arguments for a `sendMessage` call, which may carry an AI SDK user message. */ +const toAskProps = (message: TAskProps | SendMessageInput): TAskProps => { + if (!('parts' in message)) { + return message; + } + const { parentMessageId } = message.metadata ?? {}; + return { + text: message.parts.map((part) => part.text).join(''), + /** A root message's view carries a `null` parent, which `ask` would read as "append to the + * branch tail"; the turn asked to attach at the root. */ + ...(parentMessageId !== undefined && { + parentMessageId: parentMessageId ?? Constants.NO_PARENT, + }), + }; +}; + +/** `sendMessage` and `resumeStream` for the chat `id`, shared by `useChat` and `useChatActions`. */ +function useTurnActions(ask: ChatContract['ask'], id?: string) { + const jotaiStore = useStore(); + const sendMessage = useCallback( + (message, options) => ask(toAskProps(message), options), + [ask], + ); + const resumeStream = useCallback(async () => { + if (!id || id === Constants.NEW_CONVO) { + return; + } + jotaiStore.set(resumeRequestsAtom, (pending) => + pending.has(id) ? pending : new Set(pending).add(id), + ); + }, [id, jotaiStore]); + return { sendMessage, resumeStream }; +} + /** * `useChat` without `messages`, for controls that submit or read status: it re-renders when the * status changes, not on every stream frame, and never maps the conversation. @@ -338,10 +408,11 @@ export function useChatActions(): ChatActions { const status = useSyncExternalStore(subscribe, readStatus, readStatus); const regenerate = useRegenerate(getMessages, latestMessageId, regenerateTarget); const id = messagesKey || conversation?.conversationId || undefined; + const { sendMessage, resumeStream } = useTurnActions(ask, id); return useMemo( - () => ({ id, status, sendMessage: ask, regenerate, stop: stopGenerating }), - [id, status, ask, regenerate, stopGenerating], + () => ({ id, status, sendMessage, regenerate, stop: stopGenerating, resumeStream }), + [id, status, sendMessage, regenerate, stopGenerating, resumeStream], ); } @@ -404,6 +475,7 @@ export function useChat(): UseChatHelpers { ); const regenerate = useRegenerate(getMessages, latestMessageId, regenerateTarget); + const { sendMessage, resumeStream } = useTurnActions(ask, chatId); const setMessages = useCallback( (update: UIMessage[] | ((messages: UIMessage[]) => UIMessage[])) => { @@ -438,9 +510,10 @@ export function useChat(): UseChatHelpers { messages, status, error, - sendMessage: ask, + sendMessage, regenerate, stop: stopGenerating, + resumeStream, setMessages, }; } diff --git a/client/src/hooks/Chat/resume.ts b/client/src/hooks/Chat/resume.ts new file mode 100644 index 00000000000..ca50ad3d17d --- /dev/null +++ b/client/src/hooks/Chat/resume.ts @@ -0,0 +1,10 @@ +import { atom } from 'jotai'; + +/** + * Conversations with an unanswered request to reattach to their running generation (the + * facade's `resumeStream`). The pane showing a conversation consumes its entry and answers it + * with the same status re-check `useResumeOnLoad` runs when a job is announced, so a request + * never builds a second resume path. Entries are consumed, so the set stays as small as the + * requests in flight. Chat-owned, so it lives with the chat rather than in the app store. + */ +export const resumeRequestsAtom = atom>(new Set()); diff --git a/client/src/hooks/SSE/useResumeOnLoad.ts b/client/src/hooks/SSE/useResumeOnLoad.ts index 8ac28e09cb8..b6ee8e6be37 100644 --- a/client/src/hooks/SSE/useResumeOnLoad.ts +++ b/client/src/hooks/SSE/useResumeOnLoad.ts @@ -1,5 +1,5 @@ import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; -import { useStore } from 'jotai'; +import { useStore, useAtom } from 'jotai'; import { useQueryClient } from '@tanstack/react-query'; import { useSetRecoilState, useRecoilValue, useRecoilCallback } from 'recoil'; import { @@ -43,6 +43,7 @@ import { pendingApprovalActionFamily } from '~/components/Chat/approval/state'; import { agentQueuedTurnsQueryKey } from '~/data-provider/SSE/queuedTurns'; import useSteerConvert from '~/hooks/Chat/useSteerConvert'; import { revealedQueuedTurnFamily } from '~/store/steer'; +import { resumeRequestsAtom } from '~/hooks/Chat/resume'; import { useFileMapContext } from '~/Providers'; import store from '~/store'; @@ -1230,4 +1231,51 @@ export default function useResumeOnLoad( setSubmission, queryClient, ]); + + /** + * An explicit `resumeStream` request takes the announcement's path: the + * status read decides whether anything is running, and the effect above + * builds the resume submission that `useResumableSSE` attaches through the + * host transport. The request is consumed either way: one made while this + * pane is already attached is answered by that attachment. + */ + const [resumeRequests, setResumeRequests] = useAtom(resumeRequestsAtom); + const resumeRequested = !!conversationId && resumeRequests.has(conversationId); + /** The route can name a conversation before this pane has loaded it; until then the endpoint + * that decides resumability is the previous conversation's, so the request waits. */ + const routeConversationLoaded = currentConversation?.conversationId === conversationId; + useEffect(() => { + if (!resumeRequested || !conversationId || !routeConversationLoaded) { + return; + } + setResumeRequests((pending) => { + const next = new Set(pending); + next.delete(conversationId); + return next; + }); + if (!resumableEnabled || conversationId === Constants.NEW_CONVO) { + return; + } + if (hasLiveSubmissionForThisConvo) { + return; + } + /** A finished submission still installed reads as attached to the check above. */ + if (hasActiveSubmissionForThisConvo) { + setSubmission(null); + } + queryClient.invalidateQueries({ queryKey: streamStatusQueryKey(conversationId) }); + queryClient.invalidateQueries({ queryKey: [QueryKeys.messages, conversationId] }); + processedConvoRef.current = null; + setExternalRunArm((arm) => arm + 1); + }, [ + conversationId, + resumeRequested, + routeConversationLoaded, + setResumeRequests, + resumableEnabled, + hasActiveSubmissionForThisConvo, + hasLiveSubmissionForThisConvo, + setSubmission, + queryClient, + ]); }