From 15226b0b8313bdd4ba783d1c7c3392b6d7eb9bcc Mon Sep 17 00:00:00 2001
From: SidQin-cyber <sidqin0410@gmail.com>
Date: Tue, 3 Mar 2026 00:45:53 +0800
Subject: [PATCH] fix(gateway): persist streamed text when webchat final event
 lacks message
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When an agent streams text and then immediately runs tool calls, the
webchat UI drops the streamed content: the "final" event arrives with
message: undefined (buffer consumed by sub-run), and the client clears
chatStream without saving it to chatMessages.

Before clearing chatStream on a "final" event, check whether the stream
buffer has content. If no finalMessage was provided but the stream is
non-empty, synthesize an assistant message from the buffered text —
mirroring the existing "aborted" handler's preservation logic.

Closes #31895
---
 ui/src/ui/controllers/chat.test.ts | 73 +++++++++++++++++++++++++++++-
 ui/src/ui/controllers/chat.ts      |  9 ++++
 vitest.config.ts                   |  1 +
 3 files changed, 81 insertions(+), 2 deletions(-)

diff --git a/ui/src/ui/controllers/chat.test.ts b/ui/src/ui/controllers/chat.test.ts
index 456d9a537c0..c28b327c325 100644
--- a/ui/src/ui/controllers/chat.test.ts
+++ b/ui/src/ui/controllers/chat.test.ts
@@ -94,12 +94,18 @@ describe("handleChatEvent", () => {
     expect(state.chatMessages).toEqual([]);
   });
 
-  it("processes final from own run and clears state", () => {
+  it("persists streamed text when final event carries no message", () => {
+    const existingMessage = {
+      role: "user",
+      content: [{ type: "text", text: "Hi" }],
+      timestamp: 1,
+    };
     const state = createState({
       sessionKey: "main",
       chatRunId: "run-1",
-      chatStream: "Reply",
+      chatStream: "Here is my reply",
       chatStreamStartedAt: 100,
+      chatMessages: [existingMessage],
     });
     const payload: ChatEventPayload = {
       runId: "run-1",
@@ -110,6 +116,69 @@ describe("handleChatEvent", () => {
     expect(state.chatRunId).toBe(null);
     expect(state.chatStream).toBe(null);
     expect(state.chatStreamStartedAt).toBe(null);
+    expect(state.chatMessages).toHaveLength(2);
+    expect(state.chatMessages[0]).toEqual(existingMessage);
+    expect(state.chatMessages[1]).toMatchObject({
+      role: "assistant",
+      content: [{ type: "text", text: "Here is my reply" }],
+    });
+  });
+
+  it("does not persist empty or whitespace-only stream on final", () => {
+    const state = createState({
+      sessionKey: "main",
+      chatRunId: "run-1",
+      chatStream: "   ",
+      chatStreamStartedAt: 100,
+    });
+    const payload: ChatEventPayload = {
+      runId: "run-1",
+      sessionKey: "main",
+      state: "final",
+    };
+    expect(handleChatEvent(state, payload)).toBe("final");
+    expect(state.chatRunId).toBe(null);
+    expect(state.chatStream).toBe(null);
+    expect(state.chatMessages).toEqual([]);
+  });
+
+  it("does not persist null stream on final with no message", () => {
+    const state = createState({
+      sessionKey: "main",
+      chatRunId: "run-1",
+      chatStream: null,
+      chatStreamStartedAt: 100,
+    });
+    const payload: ChatEventPayload = {
+      runId: "run-1",
+      sessionKey: "main",
+      state: "final",
+    };
+    expect(handleChatEvent(state, payload)).toBe("final");
+    expect(state.chatMessages).toEqual([]);
+  });
+
+  it("prefers final payload message over streamed text", () => {
+    const state = createState({
+      sessionKey: "main",
+      chatRunId: "run-1",
+      chatStream: "Streamed partial",
+      chatStreamStartedAt: 100,
+    });
+    const finalMsg = {
+      role: "assistant",
+      content: [{ type: "text", text: "Complete reply" }],
+      timestamp: 101,
+    };
+    const payload: ChatEventPayload = {
+      runId: "run-1",
+      sessionKey: "main",
+      state: "final",
+      message: finalMsg,
+    };
+    expect(handleChatEvent(state, payload)).toBe("final");
+    expect(state.chatMessages).toEqual([finalMsg]);
+    expect(state.chatStream).toBe(null);
   });
 
   it("appends final payload message from own run before clearing stream state", () => {
diff --git a/ui/src/ui/controllers/chat.ts b/ui/src/ui/controllers/chat.ts
index 5305bde0f65..74d93bdecce 100644
--- a/ui/src/ui/controllers/chat.ts
+++ b/ui/src/ui/controllers/chat.ts
@@ -251,6 +251,15 @@ export function handleChatEvent(state: ChatState, payload?: ChatEventPayload) {
     const finalMessage = normalizeFinalAssistantMessage(payload.message);
     if (finalMessage) {
       state.chatMessages = [...state.chatMessages, finalMessage];
+    } else if (state.chatStream?.trim()) {
+      state.chatMessages = [
+        ...state.chatMessages,
+        {
+          role: "assistant",
+          content: [{ type: "text", text: state.chatStream }],
+          timestamp: Date.now(),
+        },
+      ];
     }
     state.chatStream = null;
     state.chatRunId = null;
diff --git a/vitest.config.ts b/vitest.config.ts
index 8b158848930..424fa3e8427 100644
--- a/vitest.config.ts
+++ b/vitest.config.ts
@@ -40,6 +40,7 @@ export default defineConfig({
       "ui/src/ui/views/agents-utils.test.ts",
       "ui/src/ui/views/usage-render-details.test.ts",
       "ui/src/ui/controllers/agents.test.ts",
+      "ui/src/ui/controllers/chat.test.ts",
     ],
     setupFiles: ["test/setup.ts"],
     exclude: [