WIP: resuming

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]);
+}