背景 AI 实时对话 需要流式数据

React Hooks 写法。新建WebSocket.tsx   放在根目录components

import { useCallback, useRef, useState } from 'react';

type MessageHandler = (message: MessageEvent) => void;
type ErrorHandler = (event: Event) => void;

export function useWebSocket(
  url: string,
  onMessage: MessageHandler,
  onError?: ErrorHandler,
  token?: string, // 新增 token 参数
) {
  const [isConnected, setIsConnected] = useState<boolean>(false);
  const [reconnectAttempts, setReconnectAttempts] = useState<number>(0);
  const [messages, setMessages] = useState<string[]>([]); // 用于存储消息
  const websocketRef = useRef<WebSocket | null>(null);
  const reconnectTimeoutRef = useRef<number | null>(null);
  const isManualClose = useRef<boolean>(false); // 改用 useRef 防止异步更新问题

  // 最大重连次数
  const maxReconnectAttempts = 5;
  // 初始重连延迟时间（毫秒）
  const initialReconnectDelay = 1000;

  // 连接 WebSocket
  const connectWebSocket = useCallback(() => {
    isManualClose.current = false; // 每次连接时重置手动关闭状态

    // 如果有 token，则通过 URL 查询参数传递
    const wsUrl = token ? `${url}?token=${token}` : url;

    websocketRef.current = new WebSocket(wsUrl);

    websocketRef.current.onopen = () => {
      console.log('WebSocket连接成功');
      setIsConnected(true);
      setReconnectAttempts(0); // 连接成功后重置重连次数
    };

    websocketRef.current.onmessage = (event: MessageEvent) => {
      onMessage(event);
      setMessages((prevMessages) => [...prevMessages, event.data]); // 将消息保存到状态
    };

    websocketRef.current.onerror = (event: Event) => {
      console.error('WebSocket发生错误:', event);
      if (onError) onError(event); // 如果提供了 onError 回调，调用它
    };

    websocketRef.current.onclose = () => {
      console.log('WebSocket连接关闭');
      setIsConnected(false);

      // 判断是否需要重连，只有在非手动断开时才尝试重连
      if (isManualClose.current) {
        console.log('WebSocket手动关闭，不进行重连');
        return; // 终止重连逻辑
      }

      if (reconnectAttempts < maxReconnectAttempts) {
        const delay = initialReconnectDelay * Math.pow(2, reconnectAttempts); // 指数级递增延迟
        console.log(`Reconnecting in ${delay / 1000} seconds...`);
        reconnectTimeoutRef.current = window.setTimeout(() => {
          setReconnectAttempts((prev) => prev + 1);
          connectWebSocket(); // 重试连接
        }, delay);
      } else {
        console.error('WebSocket重连失败次数过多，停止重连');
      }
    };
  }, [url, token, onMessage, onError, reconnectAttempts]);

  // 手动断开 WebSocket 连接
  const disconnectWebSocket = useCallback(() => {
    if (websocketRef.current) {
      console.log('手动断开 WebSocket 连接');
      isManualClose.current = true; // 使用 ref 直接设置为手动关闭
      websocketRef.current.close();
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current); // 清除任何重连的定时器
      }
    }
  }, []);

  // 发送消息
  const sendMessage = useCallback((message: string) => {
    if (
      websocketRef.current &&
      websocketRef.current.readyState === WebSocket.OPEN
    ) {
      websocketRef.current.send(message);
    } else {
      console.warn('WebSocket is not open.');
    }
  }, []);

  return {
    isConnected, // 是否已连接
    sendMessage, // 发送消息的函数
    disconnectWebSocket, // 手动断开连接的函数
    connectWebSocket, // 手动连接函数
    messages, // 返回存储的消息
  };
}

使用

import { useWebSocket } from '@/components/WebSocket';  //自己的文件路径
import { useEffect, useState } from 'react';
const AiChat = () => {
  const [isDialogue, setIsDialogue] = useState<boolean>(true);


  const {
    isConnected,
    sendMessage,
    disconnectWebSocket,
    connectWebSocket,
    messages,
  } = useWebSocket('wss://service-guide-uat.yadea.com.cn/ws', (message) => {
    console.log('Received message:', message.data);
  });
  const handleConnect = () => {
    console.log('手动触发连接');
    connectWebSocket(); // 手动触发连接
  };
  useEffect(() => {
    connectWebSocket();

    return () => {
      console.log('手动触发关闭');
      disconnectWebSocket();
    };
  }, []);

  return (
    <div className={style.box}>
     测试
    </div>
  );
};

export default AiChat;