React中WebSocket状态变更时组件无法实时渲染问题

核心问题 #

你的代码存在三个关键问题导致前端无法实时渲染WebSocket流式响应：

1. WebSocket实例重复创建，旧连接被丢弃
2. 事件监听重复绑定，闭包捕获旧状态值
3. 状态更新逻辑存在冗余，引发不必要的组件渲染

修复后的代码 #

1. 修正WebSocketContext.tsx #

import { createContext, useContext, useEffect, useRef, ReactNode } from 'react';

interface WebSocketContextType extends WebSocket {}

const WebSocketContext = createContext<WebSocketContextType | null>(null);

interface WebSocketProviderProps {
  children: ReactNode;
}

export const WebSocketProvider: React.FC<WebSocketProviderProps> = ({ children }) => {
  const wsRef = useRef<WebSocket | null>(null);

  useEffect(() => {
    // 仅初始化一次WebSocket连接
    const ws = new WebSocket('ws://localhost:5000/ws');
    wsRef.current = ws;

    // 组件卸载时关闭连接
    return () => {
      ws.close();
      wsRef.current = null;
    };
  }, []);

  return (
    <WebSocketContext.Provider value={wsRef.current}>
      {children}
    </WebSocketContext.Provider>
  );
};

export const useWebSocket = () => {
  const context = useContext(WebSocketContext);
  if (!context) {
    throw new Error('useWebSocket必须在WebSocketProvider内部使用');
  }
  return context;
};

2. 修正AsyncChatWindow.tsx #

import { useState, useEffect } from 'react';
import { useWebSocket } from './WebsocketContext';

function AsyncChatWindow() {
  const ws = useWebSocket();

  const [input, setInput] = useState('');
  const [fullInput, setFullInput] = useState('');
  const [output, setOutput] = useState('');

  useEffect(() => {
    if (!ws) return;

    // 定义独立的消息处理函数
    const handleMessage = (event: MessageEvent) => {
      const message = JSON.parse(event.data);
      // 直接更新为后端返回的累计内容
      setOutput(message.content);
      console.log((new Date).getSeconds(), message);
    };

    // 绑定事件监听，避免重复绑定
    ws.addEventListener('message', handleMessage);

    // 清理函数：移除监听
    return () => {
      ws.removeEventListener('message', handleMessage);
    };
  }, [ws]);

  const handleInputChange = (e: React.ChangeEvent<HTMLInputElement>) => {
    setInput(e.target.value);
  };

  const handleSubmit = (e: React.FormEvent<HTMLFormElement>) => {
    e.preventDefault();
    if (ws.readyState === WebSocket.OPEN) {
      ws.send(fullInput);
      setInput('');
      setFullInput('');
    }
  };

  return (
    <div className="flex flex-col bg-gray-500 h-screen">
      <div className="flex flex-col flex-grow overflow-y-auto">
        <div className="flex flex-row justify-center py-4 w-full max-w-[700px] mx-auto my-5 h-auto">
          <p>{output}</p>
        </div>    
      </div>
      <form className="mt-auto w-full bg-red py-2" onSubmit={handleSubmit}>
        <input 
          type="text" 
          value={input} 
          onChange={handleInputChange} 
          placeholder="输入消息..."
        />
        <button type="submit" onClick={() => setFullInput(input)}>发送</button>
      </form>
    </div>
  );
}

export default AsyncChatWindow;

修复说明 #

1. WebSocket实例持久化：用useRef保存WebSocket实例，通过useEffect仅初始化一次连接，避免组件渲染时重复创建连接导致旧连接失效。
2. 事件监听优化：改用addEventListener绑定消息处理函数，并在组件卸载时移除监听，避免重复绑定和闭包捕获旧状态的问题。
3. 简化状态逻辑：移除冗余的messageCount状态，直接通过output的变化触发组件渲染，避免不必要的状态更新。

内容的提问来源于stack exchange，提问作者Iker Ocio Zuazo