From ea9bce50dd414655424777f4f04756545eb9d5d6 Mon Sep 17 00:00:00 2001
From: Ma <dylanmaaaaa@gmail.com>
Date: Wed, 15 Jul 2026 11:13:44 +0800
Subject: [PATCH] fix(studio): don't close a streaming chat when a terminal
 task snapshot lands
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

根因：stream-events.ts 的 task:snapshot 处理里，终态快照只要通过了
hasInFlightExecution 门槛就无条件置 stream:null、isStreaming:false 并关闭
连接。竞态场景：任务刚结束、消息里还留着 in-flight 任务卡，此刻用户发消息
建立了新连接且聊天轮正在流式——服务端在连接建立时重放终态快照，门槛通过，
正在跑的聊天流被关掉，本轮所有增量事件全部丢失。

修法：接受终态快照时照常 merge 收尾任务卡，但只有聊天轮不在流式
（!isChatStreaming，从会话 runtime 读取）时才关闭连接、置 stream:null 和
isStreaming:false；聊天轮流式中则保持连接和流式状态，等聊天轮请求结束时
由 sendMessage 的收尾逻辑处理。

测试：聊天轮流式中 + 消息有 in-flight 任务卡 + 收到该任务的终态快照 →
任务卡转 completed、连接未关、isStreaming/isChatStreaming 仍为 true；
聊天轮结束后连接正常关闭。
---
 .../store/chat/slices/message/action.test.ts  | 50 +++++++++++++++++++
 .../chat/slices/message/stream-events.ts      | 11 ++--
 2 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/packages/studio/src/store/chat/slices/message/action.test.ts b/packages/studio/src/store/chat/slices/message/action.test.ts
index c305eba5..87a17e1c 100644
--- a/packages/studio/src/store/chat/slices/message/action.test.ts
+++ b/packages/studio/src/store/chat/slices/message/action.test.ts
@@ -538,6 +538,56 @@ describe("chat message actions", () => {
     expect(store.getState().sessions[sessionId]).toMatchObject({ isStreaming: false, stream: null });
   });
 
+  it("keeps the streaming chat open when a terminal task snapshot lands mid-chat", async () => {
+    const store = createTestStore();
+    const sessionId = await setupRunningTaskSession(store);
+
+    let resolveAgent!: (value: unknown) => void;
+    fetchJson.mockClear();
+    fetchJson.mockImplementationOnce(() => new Promise((resolve) => {
+      resolveAgent = resolve;
+    }));
+
+    const sent = store.getState().sendMessage(sessionId, "顺便聊两句");
+    await vi.waitFor(() => expect(fakeEventSources).toHaveLength(2));
+
+    // 竞态：任务刚结束、消息里还有 in-flight 任务卡，此刻聊天轮建立的新连接
+    // 收到服务端重放的终态快照。任务卡要收尾，但正在流式的聊天连接不能被关掉。
+    fakeEventSources[1]?.emit("task:snapshot", {
+      sessionId,
+      execution: {
+        id: "direct-short_run-1",
+        tool: "short_fiction_run",
+        label: "短篇生产",
+        status: "completed",
+        startedAt: 10,
+        completedAt: 40,
+        result: "短篇已完成",
+      },
+    });
+
+    // 任务卡转为 completed
+    expect(findTaskExecution(store, sessionId)).toMatchObject({
+      status: "completed",
+      result: "短篇已完成",
+    });
+    // 聊天轮仍在流式：连接未关、流式状态保持
+    expect(fakeEventSources[1]?.closed).toBe(false);
+    expect(store.getState().sessions[sessionId]).toMatchObject({ isStreaming: true, isChatStreaming: true });
+    expect(store.getState().sessions[sessionId]?.stream).not.toBeNull();
+
+    resolveAgent({ response: "聊完了。", session: { sessionId, sessionKind: "short" } });
+    await sent;
+
+    // 聊天轮自己收尾：任务已完成，连接关闭
+    expect(store.getState().sessions[sessionId]).toMatchObject({
+      isStreaming: false,
+      isChatStreaming: false,
+      stream: null,
+    });
+    expect(fakeEventSources[1]?.closed).toBe(true);
+  });
+
   it("closes the stream after a plain chat round when no production task is running", async () => {
     const store = createTestStore();
     const sessionId = store.getState().createDraftSession(null, "chat");
diff --git a/packages/studio/src/store/chat/slices/message/stream-events.ts b/packages/studio/src/store/chat/slices/message/stream-events.ts
index 2e688023..0bb60259 100644
--- a/packages/studio/src/store/chat/slices/message/stream-events.ts
+++ b/packages/studio/src/store/chat/slices/message/stream-events.ts
@@ -311,14 +311,19 @@ export function attachSessionStreamListeners({
       if (!running && !hasInFlightExecution(get().sessions[sessionId]?.messages ?? [], execution.id)) {
         return;
       }
+      // 聊天轮正在流式时收到终态快照（任务刚结束、新连接建立时服务端重放）：
+      // 只收尾任务卡，连接与流式状态保持不动，由聊天轮自己收尾——否则会把
+      // 正在跑的聊天流关掉，本轮增量全部丢失。
+      const chatStreaming = Boolean(get().sessions[sessionId]?.isChatStreaming);
+      const keepStream = running || chatStreaming;
       set((state) => ({
         sessions: updateSession(state.sessions, sessionId, (runtime) => ({
           messages: mergeTaskExecution(runtime.messages, execution),
-          isStreaming: running,
-          stream: running ? runtime.stream : null,
+          isStreaming: keepStream,
+          stream: keepStream ? runtime.stream : null,
         })),
       }));
-      if (!running) streamEs.close();
+      if (!keepStream) streamEs.close();
     } catch {
       // ignore
     }