WebSocket Streaming Not Working in Django + Next.js — Only Getting First and Final Message

I'm building a ai chat app with Django (Channels) for the backend and Next.js for the frontend. The goal is to stream AI-generated responses chunk by chunk over a WebSocket connection — similar to how a REST API with ReadableStream works. However, on the frontend, Its only displaying the first chunk then final completed message, not the intermediate streaming chunks.

Here’s a simplified version of my setup:

Backend (Django WebSocket Consumer): I’m using an async generator (handle_chat_request) to produce message chunks.

async def receive(self, text_data):
    logging.info(f"WebSocket message received: {text_data}")

    try:
        data = json.loads(text_data)
        message_type = data.get("type")

        if message_type == "message":
            jwt_token = data.get("token")
            if not jwt_token:
                await self.send(json.dumps({"type": "error", "message": "Missing Authorization token."}))
                await self.close(code=4001)
                return

            user_message = data.get("content")
            user_id = data.get("user_id")

            combined_response = ""

            # Stream response chunks
            async for chunk in handle_chat_request(user_message, user_id, jwt_token):
                combined_response += chunk
                await self.send(json.dumps({"type": "ai_response_chunk", "content": chunk}))
                await asyncio.sleep(0)  # Yield control to allow chunk sending

            # Send final complete message
            await self.send(json.dumps({"type": "ai_response_complete", "content": combined_response}))
        
    except Exception as e:
        logging.error(f"WebSocket error: {e}")
        await self.send(json.dumps({"type": "error", "message": "An error occurred."}))

Frontend (Next.js WebSocket Client): I’m listening for incoming chunks and trying to append them to the latest bot message:

const handleWebSocketMessage = (event) => {
  try {
    const data = JSON.parse(event.data);

    switch (data.type) {
      case "ai_response_chunk":
        setMessages((messages) => {
          const updatedMessages = [...messages];
          const lastMessageIndex = updatedMessages.findLastIndex(
            (m) => m.role === "bot"
          );
          
          if (lastMessageIndex !== -1) {
            updatedMessages[lastMessageIndex] = {
              ...updatedMessages[lastMessageIndex],
              content: updatedMessages[lastMessageIndex].content + data.content,
            };
          } else {
            updatedMessages.push({ role: "bot", content: data.content });
          }

          return updatedMessages;
        });
        break;

      case "ai_response_complete":
        //some other functionality
        break;

      case "error":
        setError(data.message);
        setLoading(false);
        break;
    }
  } catch (error) {
    console.error("Error parsing WebSocket message:", error);
    setLoading(false);
  }
};

Problem: The "ai_response_chunk" are ariving as expected by the UI only displays the first chunk then the final message content when it is complete, but I don’t see incremental updates like I do with the REST API implementation.

Rest Implementation:

  fetch(`/api/chat`, {
          method: "POST",
          headers: {
            "Content-Type": "application/json",
          },
          body: JSON.stringify({ message: userMessage }),
        }).then(async (response) => {
          if (!response.ok) throw new Error("Chat API failed");
          if (!response.body) throw new Error("ReadableStream not supported");

          const reader = response.body.getReader();
          const decoder = new TextDecoder();
          let done = false;
          let botMessage = "";

          // Add bot message placeholder
          setMessages((messages) => [
            ...messages,
            { content: "", role: "bot" },
          ]);

          while (!done) {
            const { value, done: doneReading } = await reader.read();
            done = doneReading;
            const chunk = decoder.decode(value, { stream: true });
            botMessage += chunk;

            setMessages((messages) => {
              const updatedMessages = [...messages];
              const lastMessageIndex = updatedMessages.findLastIndex(
                (m) => m.role === "bot"
              );
              if (lastMessageIndex !== -1) {
                updatedMessages[lastMessageIndex] = {
                  ...updatedMessages[lastMessageIndex],
                  content: updatedMessages[lastMessageIndex].content + chunk,
                };
              }
              return updatedMessages;
            });
          }
    ```

for chunk in self.sync_stream_chat_response(user_message, user_id, jwt_token):
    chunk_count += 1
    print(f"Sending chunk #{chunk_count} to WebSocket")
    await self.send(json.dumps({"type": "ai_response_chunk", "content": chunk}))
    
    # Add a small delay (e.g., 50ms)
    await asyncio.sleep(0.05)

await self.send(json.dumps({"type": "ai_response_complete"}))