Fix thought activity ordering

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 = visibleMessagesForTurn(turnMessages);
-    const visibleMessages: UIMessage[] = [];
+    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);
+      flushActivityMessages();
-    }
-
-    pushActivityUnits(units, activityMessages, visibleMessages);
-
-    for (const message of visibleMessages) {
       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);
   }
 

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).toHaveLength(3);
-    expect(units[0].type === "activity" ? units[0].messages.map((m) => m.id) : []).toEqual([
+    expect(units[0].type === "activity" ? units[0].messages.map((m) => m.id) : []).toEqual(["t0"]);
-      "t0",
-      "t1",
-    ]);
     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(<ThreadMessages messages={messages} isStreaming />);
 
-    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).toHaveLength(3);
-    expect(units[0].type === "activity" ? units[0].messages.map((m) => m.id) : []).toEqual([
+    expect(units[0].type === "activity" ? units[0].messages.map((m) => m.id) : []).toEqual(["r1"]);
-      "r1",
-      "t1",
-    ]);
     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(<ThreadMessages messages={messages} isStreaming={false} />);
 
-    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();
+    const laterActivity = screen.getAllByText(/thought/i).at(-1);
-    expect(screen.getAllByText(/thought/i)).toHaveLength(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", () => {