Django+Next.js WebSocket流式传输异常：仅显示首尾消息

技术栈与需求 #

• 后端：Django + Channels（WebSocket）
• 前端：Next.js
• 需求：通过WebSocket逐块流式传输AI响应，实现类似REST API ReadableStream的实时增量更新效果

现有代码 #

后端（Django WebSocket Consumer） #

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."}))

前端（Next.js WebSocket Client） #

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);
  }
};

可正常工作的REST实现对比 #

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;
          });
        }

问题现象 #

ai_response_chunk消息已按预期到达前端，但UI仅显示第一块和最终完整消息，无法实现REST版本的实时增量更新效果。

问题原因与解决方案 #

原因分析 #

1. React状态更新批次处理：React会将连续的状态更新合并为一次渲染，若WebSocket分块到达速度快，中间增量更新会被合并，仅展示最终结果。
2. WebSocket消息合并传输：后端逐块发送的小消息，可能被网络层或Django Channels合并传输，导致前端一次性接收多个分块，直接展示拼接后的最终内容。

解决方案 #

方案1：强制触发独立状态更新（前端调整）

使用setTimeout将状态更新放入宏任务队列，避免React批量合并：

case "ai_response_chunk":
  setTimeout(() => {
    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;
    });
  }, 0);
  break;

方案2：后端添加微小延迟（避免消息合并）

调整分块发送后的延迟，确保每个分块单独传输：

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.01)  # 替换原sleep(0)，添加微小延迟确保分块独立发送

方案3：使用Ref跟踪实时内容（更可靠的状态更新）

通过React Ref存储当前bot消息内容，规避状态更新合并问题：

// 组件内定义Ref
const botContentRef = useRef("");

// 处理WebSocket分块
case "ai_response_chunk":
  botContentRef.current += data.content;
  setMessages((messages) => {
    const updatedMessages = [...messages];
    const lastMessageIndex = updatedMessages.findLastIndex(m => m.role === "bot");
    
    if (lastMessageIndex !== -1) {
      updatedMessages[lastMessageIndex] = {
        ...updatedMessages[lastMessageIndex],
        content: botContentRef.current,
      };
    } else {
      updatedMessages.push({ role: "bot", content: botContentRef.current });
    }

    return updatedMessages;
  });
  break;

// 完成响应后重置Ref
case "ai_response_complete":
  botContentRef.current = "";
  // 其他逻辑
  break;

验证步骤 #

1. 在前端handleWebSocketMessage的ai_response_chunk分支添加console.log(data.content)，确认每个分块是否单独到达。
2. 测试任一方案，观察UI是否实时展示每个分块的增量更新。

内容的提问来源于stack exchange，提问作者Akshat Jaiswal