diff --git a/src/components/app/ConnectBanner.tsx b/src/components/app/ConnectBanner.tsx index cb91a8d74..5f1d50688 100644 --- a/src/components/app/ConnectBanner.tsx +++ b/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(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 ( diff --git a/src/components/chat/provider/ai-assistant-provider.tsx b/src/components/chat/provider/ai-assistant-provider.tsx index 05ce31a14..26994a3d5 100644 --- a/src/components/chat/provider/ai-assistant-provider.tsx +++ b/src/components/chat/provider/ai-assistant-provider.tsx @@ -69,7 +69,7 @@ export const AIAssistantProvider = ({ const initialScrollTopRef = useRef(null); const [selectedModelName, setSelectedModelName] = useState('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, diff --git a/src/components/chat/request/writer-request.ts b/src/components/chat/request/writer-request.ts index 24389d17b..8adc8c1e0 100644 --- a/src/components/chat/request/writer-request.ts +++ b/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, }, }), }); diff --git a/src/components/chat/types/writer.ts b/src/components/chat/types/writer.ts index b50b38fda..fb79ee9d4 100644 --- a/src/components/chat/types/writer.ts +++ b/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 {