Websocket in React not rendering component when state is changed

Question

I'm developing a ChatGPT client backend/frontend. Backend is sending via websockets a streaming response writing text asynchronously.

But it seems frontend is not rendering components until whole websocket connection is finished.

I've defined two components: WebsocketContext.tsx

interface WebSocketProviderProps {
  children: React.ReactNode;
}

export const WebSocketProvider: React.FC<WebSocketProviderProps> = ({ children }: WebSocketProviderProps) => {
  const ws = new WebSocket('ws://localhost:5000/ws');

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

export const useWebSocket = (): WebSocketContextType => useContext(WebSocketContext);

and AsyncChatWindow.tsx

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

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

  useEffect(() => {
    if (ws) {
      ws.onmessage = (event) => {
        const message = JSON.parse(event.data);
        setOutput(message.content);
        setMessageCount(prev => prev + 1);
        console.log((new Date).getSeconds());
        console.log(message);
        console.log('Receiving ' + messageCount + ' - Message: '+ output);
      };
    }
  }, [messageCount]);


  useEffect(() => {
    console.log(`Receiving... ${output}`);
  }, [output]);


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

  const handleSubmit = (e: React.FormEvent<HTMLFormElement>) => {
    e.preventDefault();
    console.log(ws?.readyState);
    if (ws && ws.readyState === WebSocket.OPEN) {
      console.log("Sending...");
      ws.send(fullInput);
    }
  };

  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 key={messageCount}>{output}</p>
        </div>    
      </div>
      <form className="mt-auto w-full bg-red py-2" onSubmit={handleSubmit}>
        <input type="text" value={input} onChange={handleInputChange} />
        <button onClick={() => setFullInput(input)}>Enviar</button>
      </form>
    </div>
  );
}

I receive a bunch of console.logs after connection closed like:

• {finished: false, content: 'Hel', role: 'assistant'}
• {finished: false, content: 'Hello! ', role: 'assistant'}
• {finished: false, content: 'Hello! How ca', role: 'assistant'}
• {finished: false, content: 'Hello! How can I ', role: 'assistant'}
• {finished: false, content: 'Hello! How can I assist', role: 'assistant'}
• {finished: false, content: 'Hello! How can I assist you t', role: 'assistant'}
• {finished: false, content: 'Hello! How can I assist you today?', role: 'assistant'}

and, for sure, frontend is finally rendered when connection is closed but not "in real time".

Any suggestion? I've tried different ways like having extra states like lastMessage and so on but nothing works.