WIP: resuming

client/src/components/Chat/ChatView.tsx

@@ -7,7 +7,7 @@ import { Constants, buildTree } from 'librechat-data-provider';
 import type { TMessage } from 'librechat-data-provider';
 import type { ChatFormValues } from '~/common';
 import { ChatContext, AddedChatContext, useFileMapContext, ChatFormProvider } from '~/Providers';
-import { useChatHelpers, useAddedResponse, useAdaptiveSSE } from '~/hooks';
+import { useChatHelpers, useAddedResponse, useAdaptiveSSE, useResumeOnLoad } from '~/hooks';
 import ConversationStarters from './Input/ConversationStarters';
 import { useGetMessagesByConvoId } from '~/data-provider';
 import MessagesView from './Messages/MessagesView';
@@ -54,6 +54,9 @@ function ChatView({ index = 0 }: { index?: number }) {
   useAdaptiveSSE(rootSubmission, chatHelpers, false, index);
   useAdaptiveSSE(addedSubmission, addedChatHelpers, true, index + 1);
 
+  // Auto-resume if navigating back to conversation with active job
+  useResumeOnLoad(conversationId, chatHelpers, index);
+
   const methods = useForm<ChatFormValues>({
     defaultValues: { text: '' },
   });

client/src/data-provider/queries/streamStatus.ts

@@ -0,0 +1,40 @@
+import { useQuery } from '@tanstack/react-query';
+import { request } from 'librechat-data-provider';
+
+export interface StreamStatusResponse {
+  active: boolean;
+  streamId?: string;
+  status?: 'running' | 'complete' | 'error' | 'aborted';
+  chunkCount?: number;
+  aggregatedContent?: Array<{ type: string; text?: string }>;
+  createdAt?: number;
+}
+
+/**
+ * Query key for stream status
+ */
+export const streamStatusQueryKey = (conversationId: string) => ['streamStatus', conversationId];
+
+/**
+ * Fetch stream status for a conversation
+ */
+export const fetchStreamStatus = async (conversationId: string): Promise<StreamStatusResponse> => {
+  const response = await request.get(`/api/agents/chat/status/${conversationId}`);
+  return response.data;
+};
+
+/**
+ * React Query hook for checking if a conversation has an active generation stream.
+ * Only fetches when conversationId is provided and resumable streams are enabled.
+ */
+export function useStreamStatus(conversationId: string | undefined, enabled = true) {
+  return useQuery({
+    queryKey: streamStatusQueryKey(conversationId || ''),
+    queryFn: () => fetchStreamStatus(conversationId!),
+    enabled: !!conversationId && enabled,
+    staleTime: 1000, // Consider stale after 1 second
+    refetchOnMount: true,
+    refetchOnWindowFocus: true,
+    retry: false,
+  });
+}

client/src/hooks/SSE/index.ts

@@ -1,6 +1,7 @@
 export { default as useSSE } from './useSSE';
 export { default as useResumableSSE } from './useResumableSSE';
 export { default as useAdaptiveSSE } from './useAdaptiveSSE';
+export { default as useResumeOnLoad } from './useResumeOnLoad';
 export { default as useStepHandler } from './useStepHandler';
 export { default as useContentHandler } from './useContentHandler';
 export { default as useAttachmentHandler } from './useAttachmentHandler';

client/src/hooks/SSE/useResumeOnLoad.ts

@@ -0,0 +1,182 @@
+import { useEffect, useState, useRef } from 'react';
+import { SSE } from 'sse.js';
+import { useSetRecoilState, useRecoilValue } from 'recoil';
+import { request } from 'librechat-data-provider';
+import type { TMessage, EventSubmission } from 'librechat-data-provider';
+import type { EventHandlerParams } from './useEventHandlers';
+import { useAuthContext } from '~/hooks/AuthContext';
+import { useGetStartupConfig, useGetUserBalance } from '~/data-provider';
+import useEventHandlers from './useEventHandlers';
+import store from '~/store';
+
+type ChatHelpers = Pick<
+  EventHandlerParams,
+  | 'setMessages'
+  | 'getMessages'
+  | 'setConversation'
+  | 'setIsSubmitting'
+  | 'newConversation'
+  | 'resetLatestMessage'
+>;
+
+/**
+ * Hook to resume streaming if navigating back to a conversation with active generation.
+ * Checks for active jobs on mount and auto-subscribes if found.
+ */
+export default function useResumeOnLoad(
+  conversationId: string | undefined,
+  chatHelpers: ChatHelpers,
+  runIndex = 0,
+) {
+  const resumableEnabled = useRecoilValue(store.resumableStreams);
+  const { token, isAuthenticated } = useAuthContext();
+  const sseRef = useRef<SSE | null>(null);
+  const checkedConvoRef = useRef<string | null>(null);
+  const [completed, setCompleted] = useState(new Set());
+  const setAbortScroll = useSetRecoilState(store.abortScrollFamily(runIndex));
+  const setShowStopButton = useSetRecoilState(store.showStopButtonByIndex(runIndex));
+
+  const { getMessages, setIsSubmitting } = chatHelpers;
+
+  const { stepHandler, finalHandler, contentHandler } = useEventHandlers({
+    ...chatHelpers,
+    setCompleted,
+    setShowStopButton,
+  });
+
+  const { data: startupConfig } = useGetStartupConfig();
+  const balanceQuery = useGetUserBalance({
+    enabled: !!isAuthenticated && startupConfig?.balance?.enabled,
+  });
+
+  /**
+   * Check for active job when conversation loads
+   */
+  useEffect(() => {
+    if (!resumableEnabled || !conversationId || !token) {
+      checkedConvoRef.current = null;
+      return;
+    }
+
+    // Only check once per conversationId to prevent loops
+    if (checkedConvoRef.current === conversationId) {
+      return;
+    }
+
+    checkedConvoRef.current = conversationId;
+
+    const checkAndResume = async () => {
+      try {
+        const response = await fetch(`/api/agents/chat/status/${conversationId}`, {
+          headers: { Authorization: `Bearer ${token}` },
+        });
+
+        if (!response.ok) {
+          return;
+        }
+
+        const { active, streamId } = await response.json();
+
+        if (!active || !streamId) {
+          return;
+        }
+
+        console.log('[ResumeOnLoad] Found active job, resuming...', { streamId });
+
+        const messages = getMessages() || [];
+        const lastMessage = messages[messages.length - 1];
+        let textIndex: number | null = null;
+
+        const url = `/api/agents/chat/stream/${encodeURIComponent(streamId)}`;
+
+        const sse = new SSE(url, {
+          headers: { Authorization: `Bearer ${token}` },
+          method: 'GET',
+        });
+        sseRef.current = sse;
+
+        sse.addEventListener('open', () => {
+          console.log('[ResumeOnLoad] Reconnected to stream');
+          setAbortScroll(false);
+          setShowStopButton(true);
+          setIsSubmitting(true);
+        });
+
+        sse.addEventListener('message', (e: MessageEvent) => {
+          try {
+            const data = JSON.parse(e.data);
+
+            if (data.final != null) {
+              try {
+                finalHandler(data, { messages } as unknown as EventSubmission);
+              } catch (error) {
+                console.error('[ResumeOnLoad] Error in finalHandler:', error);
+                setIsSubmitting(false);
+                setShowStopButton(false);
+              }
+              (startupConfig?.balance?.enabled ?? false) && balanceQuery.refetch();
+              sse.close();
+              sseRef.current = null;
+              return;
+            }
+
+            if (data.event != null) {
+              stepHandler(data, {
+                messages,
+                userMessage: lastMessage,
+              } as unknown as EventSubmission);
+              return;
+            }
+
+            if (data.type != null) {
+              const { text, index } = data;
+              if (text != null && index !== textIndex) {
+                textIndex = index;
+              }
+              contentHandler({ data, submission: { messages } as unknown as EventSubmission });
+              return;
+            }
+          } catch (error) {
+            console.error('[ResumeOnLoad] Error processing message:', error);
+          }
+        });
+
+        sse.addEventListener('error', async (e: MessageEvent) => {
+          console.log('[ResumeOnLoad] Stream error');
+          sse.close();
+          sseRef.current = null;
+          setIsSubmitting(false);
+          setShowStopButton(false);
+
+          /* @ts-ignore */
+          if (e.responseCode === 401) {
+            try {
+              const refreshResponse = await request.refreshToken();
+              const newToken = refreshResponse?.token ?? '';
+              if (newToken) {
+                request.dispatchTokenUpdatedEvent(newToken);
+              }
+            } catch (error) {
+              console.log('[ResumeOnLoad] Token refresh failed:', error);
+            }
+          }
+        });
+
+        sse.stream();
+      } catch (error) {
+        console.error('[ResumeOnLoad] Error checking job status:', error);
+      }
+    };
+
+    checkAndResume();
+
+    return () => {
+      if (sseRef.current) {
+        sseRef.current.close();
+        sseRef.current = null;
+      }
+    };
+    // Only re-run when conversationId changes
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [conversationId]);
+}