fix(gateway): persist streamed text when webchat final event lacks message

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
This commit is contained in:
SidQin-cyber
2026-03-03 00:45:53 +08:00
committed by Peter Steinberger
parent 0cf533ac61
commit 15226b0b83
3 changed files with 81 additions and 2 deletions

View File

@@ -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", () => {

View File

@@ -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;

View File

@@ -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: [