diff --git a/webui/src/lib/activity-timeline.ts b/webui/src/lib/activity-timeline.ts index b01b2578a..0cb6ff20b 100644 --- a/webui/src/lib/activity-timeline.ts +++ b/webui/src/lib/activity-timeline.ts @@ -55,8 +55,15 @@ export function normalizeActivityTimeline(messages: UIMessage[]): TurnUnit[] { const flushTurn = () => { if (turnMessages.length === 0) return; - const activityMessages: UIMessage[] = []; - const visibleMessages: UIMessage[] = []; + const visibleMessages = visibleMessagesForTurn(turnMessages); + let visibleIndex = 0; + let activityMessages: UIMessage[] = []; + + const flushActivityMessages = () => { + if (!activityMessages.length) return; + pushActivityUnits(units, activityMessages, visibleMessages.slice(visibleIndex)); + activityMessages = []; + }; for (const message of turnMessages) { if (isAgentActivityMember(message)) { @@ -66,19 +73,18 @@ export function normalizeActivityTimeline(messages: UIMessage[]): TurnUnit[] { if (assistantHasInlineReasoning(message)) { activityMessages.push(reasoningOnlyMessageFromAnswer(message)); - visibleMessages.push(stripInlineReasoning(message)); + flushActivityMessages(); + units.push({ type: "message", message: stripInlineReasoning(message) }); + visibleIndex += 1; continue; } - visibleMessages.push(message); - } - - pushActivityUnits(units, activityMessages, visibleMessages); - - for (const message of visibleMessages) { + flushActivityMessages(); units.push({ type: "message", message }); + visibleIndex += 1; } + flushActivityMessages(); turnMessages = []; }; @@ -96,9 +102,19 @@ export function normalizeActivityTimeline(messages: UIMessage[]): TurnUnit[] { return units; } +function visibleMessagesForTurn(messages: UIMessage[]): UIMessage[] { + const visibleMessages: UIMessage[] = []; + for (const message of messages) { + if (isAgentActivityMember(message)) continue; + visibleMessages.push(assistantHasInlineReasoning(message) ? stripInlineReasoning(message) : message); + } + return visibleMessages; +} + function pushActivityUnits(units: TurnUnit[], activityMessages: UIMessage[], visibleMessages: UIMessage[]) { let runMessages: UIMessage[] = []; let runBucket: "file" | "other" | undefined; + let runSegmentId: string | undefined; const flushRun = () => { if (!runMessages.length) return; @@ -110,14 +126,18 @@ function pushActivityUnits(units: TurnUnit[], activityMessages: UIMessage[], vis }); runMessages = []; runBucket = undefined; + runSegmentId = undefined; }; for (const message of activityMessages) { const bucket = isFileEditActivityMessage(message) ? "file" : "other"; - if (runBucket && bucket !== runBucket) { + const segmentId = message.activitySegmentId; + const segmentChanged = !!runSegmentId && !!segmentId && runSegmentId !== segmentId; + if ((runBucket && bucket !== runBucket) || segmentChanged) { flushRun(); } runBucket = bucket; + if (segmentId) runSegmentId = segmentId; runMessages.push(message); } diff --git a/webui/src/tests/thread-messages.test.tsx b/webui/src/tests/thread-messages.test.tsx index feed81496..3d1839f65 100644 --- a/webui/src/tests/thread-messages.test.tsx +++ b/webui/src/tests/thread-messages.test.tsx @@ -102,7 +102,7 @@ describe("ThreadMessages", () => { expect(units[2].type === "activity" ? units[2].messages.map((m) => m.id) : []).toEqual(["r2"]); }); - it("does not split ordinary tool activity just because segment ids changed", () => { + it("splits ordinary tool activity when segment ids changed", () => { const messages: UIMessage[] = [ { id: "r1", @@ -142,15 +142,58 @@ describe("ThreadMessages", () => { const units = buildDisplayUnits(messages); - expect(units).toHaveLength(1); + expect(units).toHaveLength(2); expect(units[0].type === "activity" ? units[0].messages.map((m) => m.id) : []).toEqual([ "r1", "t1", + ]); + expect(units[1].type === "activity" ? units[1].messages.map((m) => m.id) : []).toEqual([ "r2", "t2", ]); }); + it("renders a later tool segment after the visible answer that preceded it", () => { + const messages: UIMessage[] = [ + { + id: "r1", + role: "assistant", + content: "", + reasoning: "I should do a fresh search.", + activitySegmentId: "seg-1", + createdAt: 1, + }, + { + id: "a1", + role: "assistant", + content: "Let me search the latest data.", + createdAt: 2, + }, + { + id: "t1", + role: "tool", + kind: "trace", + content: "Searching query: HKUDS/nanobot GitHub stars", + traces: ["Searching query: HKUDS/nanobot GitHub stars"], + activitySegmentId: "seg-2", + createdAt: 3, + }, + ]; + + const units = buildDisplayUnits(messages); + + expect(units).toHaveLength(3); + expect(units[0].type === "activity" ? units[0].messages.map((m) => m.id) : []).toEqual(["r1"]); + expect(units[1]).toMatchObject({ + type: "message", + message: { + id: "a1", + content: "Let me search the latest data.", + }, + }); + expect(units[2].type === "activity" ? units[2].messages.map((m) => m.id) : []).toEqual(["t1"]); + }); + it("only marks the current activity timeline as live while streaming", () => { const messages: UIMessage[] = [ { @@ -254,7 +297,7 @@ describe("ThreadMessages", () => { expect(screen.getByText("final answer")).toBeInTheDocument(); }); - it("keeps late activity above the live assistant answer while streaming", () => { + it("keeps late activity after the live assistant answer while streaming", () => { const messages: UIMessage[] = [ { id: "t0", @@ -285,11 +328,8 @@ describe("ThreadMessages", () => { const units = buildDisplayUnits(messages); - expect(units).toHaveLength(2); - expect(units[0].type === "activity" ? units[0].messages.map((m) => m.id) : []).toEqual([ - "t0", - "t1", - ]); + expect(units).toHaveLength(3); + expect(units[0].type === "activity" ? units[0].messages.map((m) => m.id) : []).toEqual(["t0"]); expect(units[1]).toMatchObject({ type: "message", message: { @@ -297,15 +337,16 @@ describe("ThreadMessages", () => { content: "partial answer", }, }); + expect(units[2].type === "activity" ? units[2].messages.map((m) => m.id) : []).toEqual(["t1"]); render(); - const activity = screen.getByRole("button", { name: /working/i }); const answer = screen.getByText("partial answer"); - expect(activity.compareDocumentPosition(answer) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + const liveActivity = screen.getByRole("button", { name: /working/i }); + expect(answer.compareDocumentPosition(liveActivity) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); }); - it("keeps late activity above a completed assistant answer", () => { + it("keeps late activity after a completed assistant answer", () => { const messages: UIMessage[] = [ { id: "r1", @@ -335,11 +376,8 @@ describe("ThreadMessages", () => { const units = buildDisplayUnits(messages); - expect(units).toHaveLength(2); - expect(units[0].type === "activity" ? units[0].messages.map((m) => m.id) : []).toEqual([ - "r1", - "t1", - ]); + expect(units).toHaveLength(3); + expect(units[0].type === "activity" ? units[0].messages.map((m) => m.id) : []).toEqual(["r1"]); expect(units[1]).toMatchObject({ type: "message", message: { @@ -347,13 +385,14 @@ describe("ThreadMessages", () => { content: "Hong Kong is hot today.", }, }); + expect(units[2].type === "activity" ? units[2].messages.map((m) => m.id) : []).toEqual(["t1"]); render(); - const activity = screen.getByText("Thought for 2m 41s"); const answer = screen.getByText("Hong Kong is hot today."); - expect(activity.compareDocumentPosition(answer) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); - expect(screen.getAllByText(/thought/i)).toHaveLength(1); + const laterActivity = screen.getAllByText(/thought/i).at(-1); + expect(laterActivity).toBeTruthy(); + expect(answer.compareDocumentPosition(laterActivity!) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); }); it("renders interrupted pre-tool text as activity before the final answer", () => {