72 lines
2.1 KiB
TypeScript
72 lines
2.1 KiB
TypeScript
'use client';
|
|
|
|
import { useCallback, useEffect, useRef, useState } from 'react';
|
|
import type { Message } from '@/api/chat';
|
|
|
|
type WsEvent =
|
|
| { type: 'connection_established' }
|
|
| { type: 'chat_message'; message: Message }
|
|
| { type: 'error'; error?: string }
|
|
| { type: string; [k: string]: unknown };
|
|
|
|
export function useChatWebSocket(options: {
|
|
chatUuid: string | null;
|
|
enabled?: boolean;
|
|
onMessage?: (m: Message) => void;
|
|
}) {
|
|
const { chatUuid, enabled = true, onMessage } = options;
|
|
const [isConnected, setIsConnected] = useState(false);
|
|
const wsRef = useRef<WebSocket | null>(null);
|
|
const onMessageRef = useRef(onMessage);
|
|
onMessageRef.current = onMessage;
|
|
|
|
const disconnect = useCallback(() => {
|
|
if (wsRef.current) {
|
|
wsRef.current.close();
|
|
wsRef.current = null;
|
|
}
|
|
setIsConnected(false);
|
|
}, []);
|
|
|
|
const connect = useCallback(() => {
|
|
if (!enabled || !chatUuid) return;
|
|
|
|
const token = typeof window !== 'undefined' ? localStorage.getItem('access_token') : null;
|
|
if (!token) return;
|
|
|
|
let apiUrl = process.env.NEXT_PUBLIC_API_URL || window.location.origin;
|
|
apiUrl = apiUrl.replace(/\/api\/?$/, '').replace(/\/api\//, '/');
|
|
apiUrl = apiUrl.replace(/\/$/, '');
|
|
|
|
const wsProtocol = apiUrl.startsWith('https') ? 'wss:' : 'ws:';
|
|
const wsHost = apiUrl.replace(/^https?:\/\//, '');
|
|
const wsUrl = `${wsProtocol}//${wsHost}/ws/chat/${chatUuid}/?token=${token}`;
|
|
|
|
const ws = new WebSocket(wsUrl);
|
|
wsRef.current = ws;
|
|
|
|
ws.onopen = () => setIsConnected(true);
|
|
ws.onclose = () => setIsConnected(false);
|
|
ws.onerror = () => setIsConnected(false);
|
|
ws.onmessage = (ev) => {
|
|
try {
|
|
const data = JSON.parse(ev.data) as WsEvent;
|
|
if (data.type === 'chat_message' && (data as any).message) {
|
|
onMessageRef.current?.((data as any).message as Message);
|
|
}
|
|
} catch {
|
|
// ignore
|
|
}
|
|
};
|
|
}, [chatUuid, enabled]);
|
|
|
|
useEffect(() => {
|
|
disconnect();
|
|
connect();
|
|
return () => disconnect();
|
|
}, [connect, disconnect]);
|
|
|
|
return { isConnected };
|
|
}
|
|
|