fix(ui): clear chat stream before terminal commits

zhangguiping-xydt · web-flow · commit c2d7b4a486b1 · 2026-06-03T01:45:59.000-07:00
Fix the Control UI WebChat race where terminal assistant messages could be committed while chatStream was still live, causing history and active stream to render the same reply twice. Terminal final/aborted handling now snapshots fallback text, clears the active run/stream through the lifecycle owner, then appends the visible assistant message.\n\nFixes #71992.\n\nVerification: node scripts/run-vitest.mjs run ui/src/ui/controllers/chat.test.ts ui/src/ui/chat/run-lifecycle.test.ts ui/src/ui/chat/build-chat-items.test.ts; node scripts/run-vitest.mjs run ui/src/ui/app-chat.test.ts ui/src/ui/controllers/sessions.test.ts; node scripts/run-vitest.mjs run --config test/vitest/vitest.ui-e2e.config.ts --configLoader runner ui/src/ui/e2e/chat-flow.e2e.test.ts; Blacksmith Testbox tbx_01kt6a4zn7awkdy12d6b0q2d1q / run 26873514898; autoreview clean; PR CI 121 pass / 10 skipped.
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
@@ -1138,45 +1138,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) {
