From c1898ef7a4d44d4b3b402f33a45e3e6cd84bf076 Mon Sep 17 00:00:00 2001 From: Onur Solmaz <2453968+osolmaz@users.noreply.github.com> Date: Wed, 3 Jun 2026 11:09:26 +0800 Subject: [PATCH] fix(ui): keep live stream below prompt --- ui/src/ui/chat/build-chat-items.test.ts | 25 ++++++++++++++++ ui/src/ui/chat/build-chat-items.ts | 16 +++++++++- ui/src/ui/controllers/chat.test.ts | 34 +++++++++++++++++++++ ui/src/ui/controllers/chat.ts | 40 +++++++++++++++++++++---- 4 files changed, 109 insertions(+), 6 deletions(-) diff --git a/ui/src/ui/chat/build-chat-items.test.ts b/ui/src/ui/chat/build-chat-items.test.ts index ec8f15d4f66b..28b3f5cb3190 100644 --- a/ui/src/ui/chat/build-chat-items.test.ts +++ b/ui/src/ui/chat/build-chat-items.test.ts @@ -413,6 +413,31 @@ describe("buildChatItems", () => { expect(messageRecord(requireGroup(items[1])).content).toBe("Missing timestamp."); }); + it("renders an active stream after the persisted user turn it answers", () => { + const items = buildChatItems( + createProps({ + messages: [ + { + role: "user", + content: [{ type: "text", text: "Persisted prompt." }], + timestamp: 2_000, + }, + ], + stream: "Visible partial answer.", + streamStartedAt: 1_000, + }), + ); + + expect(items).toHaveLength(2); + expect(requireGroup(items[0]).role).toBe("user"); + expect(items[1]).toMatchObject({ + kind: "stream", + text: "Visible partial answer.", + startedAt: 2_001, + isStreaming: true, + }); + }); + it("renders submitted queued sends as user turns before chat.send ACK", () => { const groups = messageGroups({ messages: [{ role: "assistant", content: "Ready.", timestamp: 1 }], diff --git a/ui/src/ui/chat/build-chat-items.ts b/ui/src/ui/chat/build-chat-items.ts index d31c5f03e722..84348a0a80de 100644 --- a/ui/src/ui/chat/build-chat-items.ts +++ b/ui/src/ui/chat/build-chat-items.ts @@ -348,6 +348,19 @@ function chatItemTimestamp(item: ChatItem): number | null { return null; } +function timestampAfterVisibleItems(items: ChatItem[], desiredTimestamp: number): number { + const latestTimestamp = items.reduce((latest, item) => { + const timestamp = chatItemTimestamp(item); + if (timestamp == null) { + return latest; + } + return latest == null || timestamp > latest ? timestamp : latest; + }, null); + return latestTimestamp != null && desiredTimestamp <= latestTimestamp + ? latestTimestamp + 1 + : desiredTimestamp; +} + function sortChatItemsByVisibleTime(items: ChatItem[]): ChatItem[] { return items .map((item, index) => ({ item, index, timestamp: chatItemTimestamp(item) })) @@ -647,13 +660,14 @@ export function buildChatItems(props: BuildChatItemsProps): Array 0) { if (!stripHeartbeatTokenForDisplay(visibleText).shouldSkip) { items.push({ kind: "stream", key, text: visibleText, - startedAt: props.streamStartedAt ?? Date.now(), + startedAt, isStreaming: true, }); } diff --git a/ui/src/ui/controllers/chat.test.ts b/ui/src/ui/controllers/chat.test.ts index 20eda87da283..ee54df937783 100644 --- a/ui/src/ui/controllers/chat.test.ts +++ b/ui/src/ui/controllers/chat.test.ts @@ -2703,6 +2703,40 @@ describe("loadChatHistory retry handling", () => { expect(state.chatStreamStartedAt).toBeNull(); }); + it("timestamps materialized streamed text after the persisted user prompt", async () => { + const persistedUser = { + role: "user", + content: [{ type: "text", text: "first" }], + timestamp: 200, + __openclaw: { seq: 1 }, + }; + const request = vi.fn().mockResolvedValue({ + messages: [persistedUser], + thinkingLevel: "low", + }); + const state = createState({ + connected: true, + client: { request } as unknown as ChatState["client"], + chatMessages: [persistedUser], + chatRunId: null, + chatStream: "Partial answer before history catch-up.", + chatStreamStartedAt: 100, + }); + + await loadChatHistory(state); + + expect(state.chatMessages).toHaveLength(2); + expect(state.chatMessages[0]).toEqual(persistedUser); + expectTextChatMessage( + state.chatMessages[1], + "assistant", + "Partial answer before history catch-up.", + ); + expect(requireRecord(state.chatMessages[1]).timestamp).toBe(201); + expect(state.chatStream).toBeNull(); + expect(state.chatStreamStartedAt).toBeNull(); + }); + it("materializes orphaned segment-only assistant text before clearing caught-up tools", async () => { const persistedUser = { role: "user", diff --git a/ui/src/ui/controllers/chat.ts b/ui/src/ui/controllers/chat.ts index d657e0727c29..35e0df849579 100644 --- a/ui/src/ui/controllers/chat.ts +++ b/ui/src/ui/controllers/chat.ts @@ -513,6 +513,35 @@ function insertMessageAtIndex(messages: unknown[], message: unknown, index: numb return [...messages.slice(0, index), message, ...messages.slice(index)]; } +function timestampForInsertedVisibleStream( + messages: unknown[], + index: number, + desiredTimestamp: number, +): number { + const previousTimestamp = messages + .slice(0, index) + .toReversed() + .map(messageTimestampMs) + .find((timestamp): timestamp is number => timestamp != null); + const nextTimestamp = messages + .slice(index) + .map(messageTimestampMs) + .find((timestamp): timestamp is number => timestamp != null); + if (previousTimestamp != null && desiredTimestamp <= previousTimestamp) { + const afterPrevious = previousTimestamp + 1; + return nextTimestamp != null && afterPrevious >= nextTimestamp + ? previousTimestamp + (nextTimestamp - previousTimestamp) / 2 + : afterPrevious; + } + if (nextTimestamp != null && desiredTimestamp >= nextTimestamp) { + const beforeNext = nextTimestamp - 1; + return previousTimestamp != null && beforeNext <= previousTimestamp + ? previousTimestamp + (nextTimestamp - previousTimestamp) / 2 + : beforeNext; + } + return desiredTimestamp; +} + function appendVisibleStreamStateMessages( messages: unknown[], state: ChatState, @@ -524,15 +553,16 @@ function appendVisibleStreamStateMessages( if (hasAssistantStreamPartReplacement([...nextMessages, ...replacementMessages], part)) { continue; } - const streamMessage = buildAssistantStreamMessage( - part.text, - part.replacementText, - part.timestamp, - ); const toolIndex = part.source === "segment" ? currentToolStreamMessageIndex(nextMessages, state, part.toolCallId) : -1; + const insertIndex = toolIndex >= 0 ? toolIndex : nextMessages.length; + const streamMessage = buildAssistantStreamMessage( + part.text, + part.replacementText, + timestampForInsertedVisibleStream(nextMessages, insertIndex, part.timestamp), + ); nextMessages = toolIndex >= 0 ? insertMessageAtIndex(nextMessages, streamMessage, toolIndex)