chore: ws auto reconnect (#195)

Author: appflowy

src/components/app/ConnectBanner.tsx

@@ -1,4 +1,4 @@
-import { useCallback, useEffect, useMemo, useState } from 'react';
+import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 
 import { APP_EVENTS } from '@/application/constants';
@@ -7,6 +7,7 @@ import { Button } from '@/components/ui/button';
 import { Progress } from '@/components/ui/progress';
 
 import { useAppHandlers } from './app.hooks';
+import { Log } from '@/utils/log';
 
 // WebSocket readyState enum
 const READY_STATE = {
@@ -19,6 +20,7 @@ const READY_STATE = {
 export function ConnectBanner() {
   const [readyState, setReadyState] = useState<number>(READY_STATE.CONNECTING);
   const [isStableConnection, setIsStableConnection] = useState(false);
+  const autoReconnectAttemptedRef = useRef(false);
   const { eventEmitter } = useAppHandlers();
   const { t } = useTranslation();
 
@@ -65,6 +67,48 @@ export function ConnectBanner() {
     return readyState === READY_STATE.CLOSED;
   }, [readyState]);
 
+  useEffect(() => {
+    if (!isClosed) {
+      autoReconnectAttemptedRef.current = false;
+    }
+  }, [isClosed]);
+
+  // Automatically trigger reconnect when the user returns to the page and the socket is closed
+  useEffect(() => {
+    if (!isClosed || isLoading) return;
+    if (typeof window === 'undefined' || typeof document === 'undefined') return;
+
+    const tryAutoReconnect = () => {
+      if (autoReconnectAttemptedRef.current) return;
+      if (!isClosed || isLoading) return;
+      if (typeof navigator !== 'undefined' && navigator.onLine === false) return;
+      if (document.visibilityState !== 'visible') return;
+
+      autoReconnectAttemptedRef.current = true;
+
+      Log.debug('Trying to auto reconnect');
+      handleReconnect();
+    };
+
+    const handleVisibilityChange = () => {
+      if (document.visibilityState === 'visible') {
+        tryAutoReconnect();
+      }
+    };
+
+    window.addEventListener('focus', tryAutoReconnect);
+    window.addEventListener('online', tryAutoReconnect);
+    document.addEventListener('visibilitychange', handleVisibilityChange);
+
+    tryAutoReconnect();
+
+    return () => {
+      window.removeEventListener('focus', tryAutoReconnect);
+      window.removeEventListener('online', tryAutoReconnect);
+      document.removeEventListener('visibilitychange', handleVisibilityChange);
+    };
+  }, [handleReconnect, isClosed, isLoading]);
+
   // Only hide the banner when the connection is stable
   if (isStableConnection && readyState === READY_STATE.OPEN) {
     return (

src/components/chat/provider/ai-assistant-provider.tsx

@@ -69,7 +69,7 @@ export const AIAssistantProvider = ({
   const initialScrollTopRef = useRef<number | null>(null);
   const [selectedModelName, setSelectedModelName] = useState<string>('Auto');
 
-  const { currentPromptId, updateCurrentPromptId } = usePromptModal();
+  const { currentPromptId, updateCurrentPromptId, prompts } = usePromptModal();
 
   useEffect(() => {
     if (!assistantType) {
@@ -140,6 +140,10 @@ export const AIAssistantProvider = ({
       setError(null);
       try {
         setApplyingState(ApplyingState.analyzing);
+        // Find the selected prompt to get its content for custom_prompt
+        const selectedPrompt = currentPromptId
+          ? prompts.find((p) => p.id === currentPromptId)
+          : undefined;
         const { cancel, streamPromise } = await request.fetchAIAssistant(
           {
             inputText: content,
@@ -148,6 +152,7 @@ export const AIAssistantProvider = ({
             ragIds,
             completionHistory: completionHistoryRef.current,
             promptId: currentPromptId || undefined,
+            customPrompt: selectedPrompt?.content,
             modelName: selectedModelName,
           },
           handleMessageChange
@@ -175,6 +180,7 @@ export const AIAssistantProvider = ({
     [
       currentPromptId,
       handleMessageChange,
+      prompts,
       ragIds,
       request,
       responseFormat,

src/components/chat/request/writer-request.ts

@@ -42,6 +42,7 @@ export class WriterRequest {
     ragIds: string[];
     completionHistory: CompletionResult[];
     promptId?: string;
+    customPrompt?: string;
     modelName?: string;
   }, onMessage: (text: string, comment: string, done?: boolean) => void) => {
     const baseUrl = this.axiosInstance.defaults.baseURL;
@@ -76,6 +77,7 @@ export class WriterRequest {
           rag_ids: payload.ragIds.length === 0 ? [this.viewId] : payload.ragIds,
           completion_history: payload.completionHistory,
           prompt_id: payload.promptId,
+          custom_prompt: payload.customPrompt ? { system: payload.customPrompt } : undefined,
         },
       }),
     });

src/components/chat/types/writer.ts

@@ -5,7 +5,8 @@ export enum AIAssistantType {
   MakeLonger = 4,
   ContinueWriting = 5,
   Explain = 6,
-  AskAIAnything = 7
+  AskAIAnything = 7,
+  CustomPrompt = 8
 }
 
 export enum CompletionRole {