From d1f0a01a570b5516ec61ea4e4447e46914d6ffd1 Mon Sep 17 00:00:00 2001 From: Hiep Le <69354317+hieptl@users.noreply.github.com> Date: Mon, 2 Feb 2026 19:53:33 +0700 Subject: [PATCH] feat(frontend): add silent WebSocket recovery for V1 conversations (#12677) --- .../conversation-websocket-context.tsx | 23 ++-- .../contexts/websocket-provider-wrapper.tsx | 7 ++ frontend/src/hooks/use-websocket-recovery.ts | 110 ++++++++++++++++++ 3 files changed, 127 insertions(+), 13 deletions(-) create mode 100644 frontend/src/hooks/use-websocket-recovery.ts diff --git a/frontend/src/contexts/conversation-websocket-context.tsx b/frontend/src/contexts/conversation-websocket-context.tsx index f2999c0d03..74d6d8af95 100644 --- a/frontend/src/contexts/conversation-websocket-context.tsx +++ b/frontend/src/contexts/conversation-websocket-context.tsx @@ -7,7 +7,6 @@ import React, { useMemo, useRef, } from "react"; -import { useTranslation } from "react-i18next"; import { useQueryClient } from "@tanstack/react-query"; import { useWebSocket, WebSocketHookOptions } from "#/hooks/use-websocket"; import { useEventStore } from "#/stores/use-event-store"; @@ -72,6 +71,7 @@ export function ConversationWebSocketProvider({ sessionApiKey, subConversations, subConversationIds, + onDisconnect, }: { children: React.ReactNode; conversationId?: string; @@ -79,6 +79,7 @@ export function ConversationWebSocketProvider({ sessionApiKey?: string | null; subConversations?: V1AppConversation[]; subConversationIds?: string[]; + onDisconnect?: () => void; }) { // Separate connection state tracking for each WebSocket const [mainConnectionState, setMainConnectionState] = @@ -126,8 +127,6 @@ export function ConversationWebSocketProvider({ conversationId: string; } | null>(null); - const { t } = useTranslation(); - // Helper function to update metrics from stats event const updateMetricsFromStats = useCallback( (event: ConversationStateUpdateEventStats) => { @@ -631,12 +630,10 @@ export function ConversationWebSocketProvider({ }, onClose: (event: CloseEvent) => { setMainConnectionState("CLOSED"); - // Only show error message if we've previously connected successfully - // This prevents showing errors during initial connection attempts (e.g., when auto-starting a conversation) + // Trigger silent recovery on unexpected disconnect + // Do NOT show error message - recovery happens automatically if (event.code !== 1000 && hasConnectedRefMain.current) { - setErrorMessage( - `${t(I18nKey.STATUS$CONNECTION_LOST)}: ${event.reason || t(I18nKey.STATUS$DISCONNECTED_REFRESH_PAGE)}`, - ); + onDisconnect?.(); } }, onError: () => { @@ -655,6 +652,7 @@ export function ConversationWebSocketProvider({ sessionApiKey, conversationId, conversationUrl, + onDisconnect, ]); // Separate WebSocket options for planning agent connection @@ -703,12 +701,10 @@ export function ConversationWebSocketProvider({ }, onClose: (event: CloseEvent) => { setPlanningConnectionState("CLOSED"); - // Only show error message if we've previously connected successfully - // This prevents showing errors during initial connection attempts (e.g., when auto-starting a conversation) + // Trigger silent recovery on unexpected disconnect + // Do NOT show error message - recovery happens automatically if (event.code !== 1000 && hasConnectedRefPlanning.current) { - setErrorMessage( - `${t(I18nKey.STATUS$CONNECTION_LOST)}: ${event.reason || t(I18nKey.STATUS$DISCONNECTED_REFRESH_PAGE)}`, - ); + onDisconnect?.(); } }, onError: () => { @@ -726,6 +722,7 @@ export function ConversationWebSocketProvider({ removeErrorMessage, sessionApiKey, subConversations, + onDisconnect, ]); // Only attempt WebSocket connection when we have a valid URL diff --git a/frontend/src/contexts/websocket-provider-wrapper.tsx b/frontend/src/contexts/websocket-provider-wrapper.tsx index d278a46551..3aa21e4113 100644 --- a/frontend/src/contexts/websocket-provider-wrapper.tsx +++ b/frontend/src/contexts/websocket-provider-wrapper.tsx @@ -3,6 +3,7 @@ import { WsClientProvider } from "#/context/ws-client-provider"; import { ConversationWebSocketProvider } from "#/contexts/conversation-websocket-context"; import { useActiveConversation } from "#/hooks/query/use-active-conversation"; import { useSubConversations } from "#/hooks/query/use-sub-conversations"; +import { useWebSocketRecovery } from "#/hooks/use-websocket-recovery"; interface WebSocketProviderWrapperProps { children: React.ReactNode; @@ -47,6 +48,10 @@ export function WebSocketProviderWrapper({ (subConversation) => subConversation !== null, ); + // Silent recovery for V1 WebSocket disconnections + const { reconnectKey, handleDisconnect } = + useWebSocketRecovery(conversationId); + if (version === 0) { return ( @@ -58,11 +63,13 @@ export function WebSocketProviderWrapper({ if (version === 1) { return ( {children} diff --git a/frontend/src/hooks/use-websocket-recovery.ts b/frontend/src/hooks/use-websocket-recovery.ts new file mode 100644 index 0000000000..d15358d12e --- /dev/null +++ b/frontend/src/hooks/use-websocket-recovery.ts @@ -0,0 +1,110 @@ +import React from "react"; +import { useQueryClient } from "@tanstack/react-query"; +import { useUnifiedResumeConversationSandbox } from "#/hooks/mutation/use-unified-start-conversation"; +import { useUserProviders } from "#/hooks/use-user-providers"; +import { useErrorMessageStore } from "#/stores/error-message-store"; +import { I18nKey } from "#/i18n/declaration"; + +const MAX_RECOVERY_ATTEMPTS = 3; +const RECOVERY_COOLDOWN_MS = 5000; +const RECOVERY_SETTLED_DELAY_MS = 2000; + +/** + * Hook that handles silent WebSocket recovery by resuming the sandbox + * when a WebSocket disconnection is detected. + * + * @param conversationId - The conversation ID to recover + * @returns reconnectKey - Key to force provider remount (resets connection state) + * @returns handleDisconnect - Callback to trigger recovery on WebSocket disconnect + */ +export function useWebSocketRecovery(conversationId: string) { + // Recovery state (refs to avoid re-renders) + const recoveryAttemptsRef = React.useRef(0); + const recoveryInProgressRef = React.useRef(false); + const lastRecoveryAttemptRef = React.useRef(null); + + // Key to force remount of provider after recovery (resets connection state to "CONNECTING") + const [reconnectKey, setReconnectKey] = React.useState(0); + + const queryClient = useQueryClient(); + const { mutate: resumeConversation } = useUnifiedResumeConversationSandbox(); + const { providers } = useUserProviders(); + const setErrorMessage = useErrorMessageStore( + (state) => state.setErrorMessage, + ); + + // Reset recovery state when conversation changes + React.useEffect(() => { + recoveryAttemptsRef.current = 0; + recoveryInProgressRef.current = false; + lastRecoveryAttemptRef.current = null; + }, [conversationId]); + + // Silent recovery callback - resumes sandbox when WebSocket disconnects + const handleDisconnect = React.useCallback(() => { + // Prevent concurrent recovery attempts + if (recoveryInProgressRef.current) return; + + // Check cooldown + const now = Date.now(); + if ( + lastRecoveryAttemptRef.current && + now - lastRecoveryAttemptRef.current < RECOVERY_COOLDOWN_MS + ) { + return; + } + + // Check max attempts - notify user when recovery is exhausted + if (recoveryAttemptsRef.current >= MAX_RECOVERY_ATTEMPTS) { + setErrorMessage(I18nKey.STATUS$CONNECTION_LOST); + return; + } + + // Start silent recovery + recoveryInProgressRef.current = true; + lastRecoveryAttemptRef.current = now; + recoveryAttemptsRef.current += 1; + + resumeConversation( + { conversationId, providers }, + { + onSuccess: async () => { + // Invalidate and wait for refetch to complete before remounting + // This ensures the provider remounts with fresh data (url: null during startup) + await queryClient.invalidateQueries({ + queryKey: ["user", "conversation", conversationId], + }); + + // Force remount to reset connection state to "CONNECTING" + setReconnectKey((k) => k + 1); + + // Reset recovery state on success + recoveryAttemptsRef.current = 0; + recoveryInProgressRef.current = false; + lastRecoveryAttemptRef.current = null; + }, + onError: () => { + // If this was the last attempt, show error to user + if (recoveryAttemptsRef.current >= MAX_RECOVERY_ATTEMPTS) { + setErrorMessage(I18nKey.STATUS$CONNECTION_LOST); + } + // recoveryInProgressRef will be reset by onSettled + }, + onSettled: () => { + // Allow next attempt after a delay (covers both success and error) + setTimeout(() => { + recoveryInProgressRef.current = false; + }, RECOVERY_SETTLED_DELAY_MS); + }, + }, + ); + }, [ + conversationId, + providers, + resumeConversation, + queryClient, + setErrorMessage, + ]); + + return { reconnectKey, handleDisconnect }; +}