⏯️ feat: Resume Streams and Send AI SDK Messages Through the Chat Facade

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 }) => (
+        <QueryClientProvider client={queryClient}>
+          <IsolatedAtomStore
+            seed={(store) => {
+              atoms = store;
+              store.set(resumeRequestsAtom, new Set(['convo-2']));
+            }}
+          >
+            <ChatContext.Provider value={contract}>{children}</ChatContext.Provider>
+          </IsolatedAtomStore>
+        </QueryClientProvider>
+      ),
+    });
+    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({

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 <ChatContext.Provider value={helpers}>{children}</ChatContext.Provider>;
+    }
+
+    /** 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 }) => (
+          <Wrapper>
+            <ChatHost>{children}</ChatHost>
+          </Wrapper>
+        ),
+      });
+    };
+
+    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<string>());
+    });
+
+    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');
+    });
+  });
 });