From c749df4e0aed91b97d35f33a29dde10c097ba675 Mon Sep 17 00:00:00 2001
From: oudecheng <13802883547@139.com>
Date: Fri, 7 Aug 2026 06:51:00 +0800
Subject: [PATCH] =?UTF-8?q?fix(web):=20WebSocket=20=E6=96=AD=E8=BF=9E?=
 =?UTF-8?q?=E6=97=B6=E6=B8=85=E7=90=86=E6=B5=81=E5=BC=8F=E7=8A=B6=E6=80=81?=
 =?UTF-8?q?=EF=BC=8C=E9=98=B2=E6=AD=A2=E9=87=8D=E8=BF=9E=E5=90=8E=E6=B6=88?=
 =?UTF-8?q?=E6=81=AF=E4=B8=B2=E6=B5=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

流式优化引入的回归风险：当 WebSocket 在流式输出中途断连时，
streamingRef（contentChunks/index）不会被重置。重连后新消息的
delta 会通过 s.index >= 0 分支追加到旧消息的错误位置。

修复链路：useWebSocket 的 onDisconnect 回调 → useChat 的
finishStreaming → useMessages 的 finishStreaming（重置 ref +
flush pending chunks）。

- useMessages.ts: 在 UseMessagesReturn 接口和 return 对象暴露
  finishStreaming（原为内部函数）
- useChat.ts: 在返回对象透传 finishStreaming
- App.tsx: useWebSocket 调用新增 onDisconnect 回调触发 finishStreaming
- useChat.test.ts: 新增用例 17 验证断连后新 delta 不污染旧消息
---
 web/src/App.tsx                   |  2 ++
 web/src/hooks/chat/useMessages.ts |  2 ++
 web/src/hooks/useChat.test.ts     | 25 +++++++++++++++++++++++++
 web/src/hooks/useChat.ts          |  2 ++
 4 files changed, 31 insertions(+)

diff --git a/web/src/App.tsx b/web/src/App.tsx
index 8a2c9ad..54097f3 100644
--- a/web/src/App.tsx
+++ b/web/src/App.tsx
@@ -115,11 +115,13 @@ function App() {
     exitSubAgentView,
     navigateToSubAgentLevel,
     handleStop,
+    finishStreaming,
   } = useChat();
 
   const { status, sendMessage } = useWebSocket({
     url: wsUrl,
     onMessage: handleServerMessage,
+    onDisconnect: () => finishStreaming(),
   });
 
   // 将 sendMessage 注入到 useChat，供 handleServerMessage 内部发送命令
diff --git a/web/src/hooks/chat/useMessages.ts b/web/src/hooks/chat/useMessages.ts
index 8fad370..79f5f77 100644
--- a/web/src/hooks/chat/useMessages.ts
+++ b/web/src/hooks/chat/useMessages.ts
@@ -37,6 +37,7 @@ export interface UseMessagesReturn {
   setIsLoading: Dispatch<SetStateAction<boolean>>;
   handleMessage: (content: string, attachments?: Attachment[]) => void;
   clearMessages: () => void;
+  finishStreaming: () => void;
   handleStop: () => Command;
   /** 处理主视图的消息类 case（task_started, stream_*, tool_*, execution_*, error），返回是否已处理 */
   handleMainViewMessage: (message: WsOutbound) => boolean;
@@ -426,6 +427,7 @@ export function useMessages(options: UseMessagesOptions): UseMessagesReturn {
     setIsLoading,
     handleMessage,
     clearMessages,
+    finishStreaming,
     handleStop,
     handleMainViewMessage,
   };
diff --git a/web/src/hooks/useChat.test.ts b/web/src/hooks/useChat.test.ts
index 515fffa..3384bf3 100644
--- a/web/src/hooks/useChat.test.ts
+++ b/web/src/hooks/useChat.test.ts
@@ -498,4 +498,29 @@ describe('useChat - handleServerMessage characterization', () => {
     expect(result.current.messages[0].content).toBe('final text');
     expect(result.current.messages[0].reasoningContent).toBe('final reasoning');
   });
+
+  it('17. finishStreaming (onDisconnect) mid-stream resets streaming state, preventing post-reconnect leakage', () => {
+    const { result } = renderUseChat();
+    // 流式进行中
+    act(() => {
+      result.current.handleServerMessage({ type: 'stream_delta', id: 's1', delta: 'part1' });
+      result.current.handleServerMessage({ type: 'stream_delta', id: 's1', delta: 'part2' });
+    });
+    expect(result.current.messages).toHaveLength(1);
+    expect(result.current.messages[0].content).toBe('part1part2');
+
+    // 模拟 WebSocket 断连：onDisconnect 回调触发 finishStreaming
+    act(() => result.current.finishStreaming());
+
+    // 重连后，新 delta（不同 id）应作为新消息出现，不应追加到已断连的旧消息
+    act(() =>
+      result.current.handleServerMessage({ type: 'stream_delta', id: 's2', delta: 'fresh' }),
+    );
+    const s2 = result.current.messages.find((m) => m.id === 's2');
+    expect(s2).toBeDefined();
+    expect(s2?.content).toBe('fresh');
+    // 旧消息内容保持不变（未被新 delta 污染）
+    const s1 = result.current.messages.find((m) => m.id === 's1');
+    expect(s1?.content).toBe('part1part2');
+  });
 });
diff --git a/web/src/hooks/useChat.ts b/web/src/hooks/useChat.ts
index 8118156..077bd2a 100644
--- a/web/src/hooks/useChat.ts
+++ b/web/src/hooks/useChat.ts
@@ -59,6 +59,7 @@ interface UseChatReturn {
   handleMessage: (content: string, attachments?: Attachment[]) => void;
   handleCommand: (command: Command) => void;
   clearMessages: () => void;
+  finishStreaming: () => void;
   handleServerMessage: (message: WsOutbound) => void;
   setSendMessage: (fn: (msg: WsInbound) => boolean) => void;
 
@@ -360,6 +361,7 @@ export function useChat(): UseChatReturn {
     handleMessage: messages.handleMessage,
     handleCommand,
     clearMessages: messages.clearMessages,
+    finishStreaming: messages.finishStreaming,
     handleServerMessage,
     setSendMessage: conn.setSendMessage,
     selectTopic,