fix(ui): clear chat stream before terminal commits

zhangguiping-xydt · zhangguiping-xydt · commit 095b9b496f1e · 2026-06-03T00:07:24.000+08:00
diff --git a/ui/src/ui/controllers/chat.test.ts b/ui/src/ui/controllers/chat.test.ts
@@ -82,6 +82,28 @@ function createActiveStreamingState() {
   });
 }
 
+function trackChatMessagesAssignments(state: ChatState) {
+  let chatMessages = state.chatMessages;
+  const assignments: Array<{
+    chatRunId: string | null;
+    chatStream: string | null;
+    messages: unknown[];
+  }> = [];
+  Object.defineProperty(state, "chatMessages", {
+    configurable: true,
+    get: () => chatMessages,
+    set: (messages: unknown[]) => {
+      assignments.push({
+        chatRunId: state.chatRunId,
+        chatStream: state.chatStream,
+        messages,
+      });
+      chatMessages = messages;
+    },
+  });
+  return assignments;
+}
+
 function createOtherRunSilentFinalPayload(text: string): ChatEventPayload {
   return {
     runId: "run-announce",
@@ -733,7 +755,10 @@ describe("handleChatEvent", () => {
       sessionKey: "main",
       state: "final",
     };
+    const assignments = trackChatMessagesAssignments(state);
+
     expect(handleChatEvent(state, payload)).toBe("final");
+    expect(assignments).toMatchObject([{ chatRunId: null, chatStream: null }]);
     expect(state.chatRunId).toBe(null);
     expect(state.chatStream).toBe(null);
     expect(state.chatStreamStartedAt).toBe(null);
@@ -799,7 +824,7 @@ describe("handleChatEvent", () => {
     expect(state.chatStream).toBe(null);
   });
 
-  it("appends final payload message from own run before clearing stream state", () => {
+  it("appends final payload message from own run after clearing stream state", () => {
     const state = createState({
       sessionKey: "main",
       chatRunId: "run-1",
@@ -816,7 +841,10 @@ describe("handleChatEvent", () => {
         timestamp: 101,
       },
     };
+    const assignments = trackChatMessagesAssignments(state);
+
     expect(handleChatEvent(state, payload)).toBe("final");
+    expect(assignments).toMatchObject([{ chatRunId: null, chatStream: null }]);
     expect(state.chatMessages).toEqual([payload.message]);
     expect(state.chatRunId).toBe(null);
     expect(state.chatStream).toBe(null);
@@ -847,8 +875,10 @@ describe("handleChatEvent", () => {
       state: "aborted",
       message: partialMessage,
     };
+    const assignments = trackChatMessagesAssignments(state);
 
     expect(handleChatEvent(state, payload)).toBe("aborted");
+    expect(assignments).toMatchObject([{ chatRunId: null, chatStream: null }]);
     expect(state.chatRunId).toBe(null);
     expect(state.chatStream).toBe(null);
     expect(state.chatStreamStartedAt).toBe(null);
diff --git a/ui/src/ui/controllers/chat.ts b/ui/src/ui/controllers/chat.ts
@@ -1129,45 +1129,50 @@ export function handleChatEvent(state: ChatState, payload?: ChatEventPayload) {
     }
   } else if (payload.state === "final") {
     const finalMessage = normalizeFinalAssistantMessage(payload.message);
-    if (finalMessage && !shouldHideAssistantChatMessage(finalMessage)) {
-      state.chatMessages = [...state.chatMessages, finalMessage];
-    } else if (
-      state.chatStream?.trim() &&
-      !isSilentReplyStream(state.chatStream) &&
-      !isHeartbeatAckStream(state.chatStream)
-    ) {
-      state.chatMessages = [
-        ...state.chatMessages,
-        {
-          role: "assistant",
-          content: [{ type: "text", text: state.chatStream }],
-          timestamp: Date.now(),
-        },
-      ];
-    }
+    const visibleFinalMessage =
+      finalMessage && !shouldHideAssistantChatMessage(finalMessage) ? finalMessage : null;
+    const streamedText = state.chatStream ?? "";
+    const fallbackMessage =
+      !visibleFinalMessage &&
+      streamedText.trim() &&
+      !isSilentReplyStream(streamedText) &&
+      !isHeartbeatAckStream(streamedText)
+        ? {
+            role: "assistant",
+            content: [{ type: "text", text: streamedText }],
+            timestamp: Date.now(),
+          }
+        : null;
     reconcileTerminalRun("done", "done");
+    if (visibleFinalMessage) {
+      state.chatMessages = [...state.chatMessages, visibleFinalMessage];
+    } else if (fallbackMessage) {
+      state.chatMessages = [...state.chatMessages, fallbackMessage];
+    }
   } else if (payload.state === "aborted") {
     const normalizedMessage = normalizeAbortedAssistantMessage(payload.message);
-    if (normalizedMessage && !shouldHideAssistantChatMessage(normalizedMessage)) {
-      state.chatMessages = [...state.chatMessages, normalizedMessage];
-    } else {
-      const streamedText = state.chatStream ?? "";
-      if (
-        streamedText.trim() &&
-        !isSilentReplyStream(streamedText) &&
-        !isHeartbeatAckStream(streamedText)
-      ) {
-        state.chatMessages = [
-          ...state.chatMessages,
-          {
+    const visibleAbortedMessage =
+      normalizedMessage && !shouldHideAssistantChatMessage(normalizedMessage)
+        ? normalizedMessage
+        : null;
+    const streamedText = state.chatStream ?? "";
+    const fallbackMessage =
+      !visibleAbortedMessage &&
+      streamedText.trim() &&
+      !isSilentReplyStream(streamedText) &&
+      !isHeartbeatAckStream(streamedText)
+        ? {
             role: "assistant",
             content: [{ type: "text", text: streamedText }],
             timestamp: Date.now(),
-          },
-        ];
-      }
-    }
+          }
+        : null;
     reconcileTerminalRun("interrupted", "killed");
+    if (visibleAbortedMessage) {
+      state.chatMessages = [...state.chatMessages, visibleAbortedMessage];
+    } else if (fallbackMessage) {
+      state.chatMessages = [...state.chatMessages, fallbackMessage];
+    }
   } else if (payload.state === "error") {
     const errorMessage = hadActiveRunBeforeEvent ? buildErrorAssistantMessage(payload) : null;
     if (errorMessage) {
