Skip to content

Commit 0737701

Browse files
committed
2 parents 5c787bb + ec90ea1 commit 0737701

File tree

2 files changed

+123
-25
lines changed

2 files changed

+123
-25
lines changed

ui/components/ChatWindow.tsx

Lines changed: 95 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,9 @@ import crypto from 'crypto';
99
import { toast } from 'sonner';
1010
import { useSearchParams } from 'next/navigation';
1111
import { getSuggestions } from '@/lib/actions';
12-
import Error from 'next/error';
1312
import { Settings } from 'lucide-react';
1413
import SettingsDialog from './SettingsDialog';
14+
import NextError from 'next/error';
1515

1616
export type Message = {
1717
messageId: string;
@@ -34,11 +34,24 @@ const useSocket = (
3434
setIsWSReady: (ready: boolean) => void,
3535
setError: (error: boolean) => void,
3636
) => {
37-
const [ws, setWs] = useState<WebSocket | null>(null);
37+
const wsRef = useRef<WebSocket | null>(null);
38+
const reconnectTimeoutRef = useRef<NodeJS.Timeout>();
39+
const retryCountRef = useRef(0);
40+
const isCleaningUpRef = useRef(false);
41+
const MAX_RETRIES = 3;
42+
const INITIAL_BACKOFF = 1000; // 1 second
43+
44+
const getBackoffDelay = (retryCount: number) => {
45+
return Math.min(INITIAL_BACKOFF * Math.pow(2, retryCount), 10000); // Cap at 10 seconds
46+
};
3847

3948
useEffect(() => {
40-
if (!ws) {
41-
const connectWs = async () => {
49+
const connectWs = async () => {
50+
if (wsRef.current?.readyState === WebSocket.OPEN) {
51+
wsRef.current.close();
52+
}
53+
54+
try {
4255
let chatModel = localStorage.getItem('chatModel');
4356
let chatModelProvider = localStorage.getItem('chatModelProvider');
4457
let embeddingModel = localStorage.getItem('embeddingModel');
@@ -61,7 +74,13 @@ const useSocket = (
6174
'Content-Type': 'application/json',
6275
},
6376
},
64-
).then(async (res) => await res.json());
77+
).then(async (res) => {
78+
if (!res.ok)
79+
throw new Error(
80+
`Failed to fetch models: ${res.status} ${res.statusText}`,
81+
);
82+
return res.json();
83+
});
6584

6685
if (
6786
!chatModel ||
@@ -204,6 +223,7 @@ const useSocket = (
204223
wsURL.search = searchParams.toString();
205224

206225
const ws = new WebSocket(wsURL.toString());
226+
wsRef.current = ws;
207227

208228
const timeoutId = setTimeout(() => {
209229
if (ws.readyState !== 1) {
@@ -219,11 +239,16 @@ const useSocket = (
219239
const interval = setInterval(() => {
220240
if (ws.readyState === 1) {
221241
setIsWSReady(true);
242+
setError(false);
243+
if (retryCountRef.current > 0) {
244+
toast.success('Connection restored.');
245+
}
246+
retryCountRef.current = 0;
222247
clearInterval(interval);
223248
}
224249
}, 5);
225250
clearTimeout(timeoutId);
226-
console.log('[DEBUG] opened');
251+
console.debug(new Date(), 'ws:connected');
227252
}
228253
if (data.type === 'error') {
229254
toast.error(data.data);
@@ -232,24 +257,68 @@ const useSocket = (
232257

233258
ws.onerror = () => {
234259
clearTimeout(timeoutId);
235-
setError(true);
260+
setIsWSReady(false);
236261
toast.error('WebSocket connection error.');
237262
};
238263

239264
ws.onclose = () => {
240265
clearTimeout(timeoutId);
241-
setError(true);
242-
console.log('[DEBUG] closed');
266+
setIsWSReady(false);
267+
console.debug(new Date(), 'ws:disconnected');
268+
if (!isCleaningUpRef.current) {
269+
toast.error('Connection lost. Attempting to reconnect...');
270+
attemptReconnect();
271+
}
243272
};
273+
} catch (error) {
274+
console.debug(new Date(), 'ws:error', error);
275+
setIsWSReady(false);
276+
attemptReconnect();
277+
}
278+
};
244279

245-
setWs(ws);
246-
};
280+
const attemptReconnect = () => {
281+
retryCountRef.current += 1;
247282

248-
connectWs();
249-
}
250-
}, [ws, url, setIsWSReady, setError]);
283+
if (retryCountRef.current > MAX_RETRIES) {
284+
console.debug(new Date(), 'ws:max_retries');
285+
setError(true);
286+
toast.error(
287+
'Unable to connect to server after multiple attempts. Please refresh the page to try again.',
288+
);
289+
return;
290+
}
251291

252-
return ws;
292+
const backoffDelay = getBackoffDelay(retryCountRef.current);
293+
console.debug(
294+
new Date(),
295+
`ws:retry attempt=${retryCountRef.current}/${MAX_RETRIES} delay=${backoffDelay}ms`,
296+
);
297+
298+
if (reconnectTimeoutRef.current) {
299+
clearTimeout(reconnectTimeoutRef.current);
300+
}
301+
302+
reconnectTimeoutRef.current = setTimeout(() => {
303+
connectWs();
304+
}, backoffDelay);
305+
};
306+
307+
connectWs();
308+
309+
return () => {
310+
if (reconnectTimeoutRef.current) {
311+
clearTimeout(reconnectTimeoutRef.current);
312+
}
313+
if (wsRef.current?.readyState === WebSocket.OPEN) {
314+
wsRef.current.close();
315+
isCleaningUpRef.current = true;
316+
console.debug(new Date(), 'ws:cleanup');
317+
}
318+
};
319+
}, [url, setIsWSReady, setError]);
320+
321+
return wsRef.current;
253322
};
254323

255324
const loadMessages = async (
@@ -293,7 +362,7 @@ const loadMessages = async (
293362
return [msg.role, msg.content];
294363
}) as [string, string][];
295364

296-
console.log('[DEBUG] messages loaded');
365+
console.debug(new Date(), 'app:messages_loaded');
297366

298367
document.title = messages[0].content;
299368

@@ -377,7 +446,7 @@ const ChatWindow = ({ id }: { id?: string }) => {
377446
return () => {
378447
if (ws?.readyState === 1) {
379448
ws.close();
380-
console.log('[DEBUG] closed');
449+
console.debug(new Date(), 'ws:cleanup');
381450
}
382451
};
383452
// eslint-disable-next-line react-hooks/exhaustive-deps
@@ -392,12 +461,18 @@ const ChatWindow = ({ id }: { id?: string }) => {
392461
useEffect(() => {
393462
if (isMessagesLoaded && isWSReady) {
394463
setIsReady(true);
395-
console.log('[DEBUG] ready');
464+
console.debug(new Date(), 'app:ready');
465+
} else {
466+
setIsReady(false);
396467
}
397468
}, [isMessagesLoaded, isWSReady]);
398469

399470
const sendMessage = async (message: string, messageId?: string) => {
400471
if (loading) return;
472+
if (!ws || ws.readyState !== WebSocket.OPEN) {
473+
toast.error('Cannot send message while disconnected');
474+
return;
475+
}
401476

402477
setLoading(true);
403478
setMessageAppeared(false);
@@ -408,7 +483,7 @@ const ChatWindow = ({ id }: { id?: string }) => {
408483

409484
messageId = messageId ?? crypto.randomBytes(7).toString('hex');
410485

411-
ws?.send(
486+
ws.send(
412487
JSON.stringify({
413488
type: 'message',
414489
message: {
@@ -571,7 +646,7 @@ const ChatWindow = ({ id }: { id?: string }) => {
571646

572647
return isReady ? (
573648
notFound ? (
574-
<Error statusCode={404} />
649+
<NextError statusCode={404} />
575650
) : (
576651
<div>
577652
{messages.length > 0 ? (

ui/components/SearchVideos.tsx

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
/* eslint-disable @next/next/no-img-element */
22
import { PlayCircle, PlayIcon, PlusIcon, VideoIcon } from 'lucide-react';
3-
import { useState } from 'react';
3+
import { useRef, useState } from 'react';
44
import Lightbox, { GenericSlide, VideoSlide } from 'yet-another-react-lightbox';
55
import 'yet-another-react-lightbox/styles.css';
66
import { Message } from './ChatWindow';
@@ -35,6 +35,8 @@ const Searchvideos = ({
3535
const [loading, setLoading] = useState(false);
3636
const [open, setOpen] = useState(false);
3737
const [slides, setSlides] = useState<VideoSlide[]>([]);
38+
const [currentIndex, setCurrentIndex] = useState(0);
39+
const videoRefs = useRef<(HTMLIFrameElement | null)[]>([]);
3840

3941
return (
4042
<>
@@ -182,18 +184,39 @@ const Searchvideos = ({
182184
open={open}
183185
close={() => setOpen(false)}
184186
slides={slides}
187+
index={currentIndex}
188+
on={{
189+
view: ({ index }) => {
190+
const previousIframe = videoRefs.current[currentIndex];
191+
if (previousIframe?.contentWindow) {
192+
previousIframe.contentWindow.postMessage(
193+
'{"event":"command","func":"pauseVideo","args":""}',
194+
'*',
195+
);
196+
}
197+
198+
setCurrentIndex(index);
199+
},
200+
}}
185201
render={{
186-
slide: ({ slide }) =>
187-
slide.type === 'video-slide' ? (
202+
slide: ({ slide }) => {
203+
const index = slides.findIndex((s) => s === slide);
204+
return slide.type === 'video-slide' ? (
188205
<div className="h-full w-full flex flex-row items-center justify-center">
189206
<iframe
190-
src={slide.iframe_src}
207+
src={`${slide.iframe_src}${slide.iframe_src.includes('?') ? '&' : '?'}enablejsapi=1`}
208+
ref={(el) => {
209+
if (el) {
210+
videoRefs.current[index] = el;
211+
}
212+
}}
191213
className="aspect-video max-h-[95vh] w-[95vw] rounded-2xl md:w-[80vw]"
192214
allowFullScreen
193215
allow="accelerometer; autoplay; encrypted-media; gyroscope; picture-in-picture"
194216
/>
195217
</div>
196-
) : null,
218+
) : null;
219+
},
197220
}}
198221
/>
199222
</>

0 commit comments

Comments
 (0)