perf(webapp): stop the agent transcript re-rendering on every streamed token

The winners map was rebuilt on each render, so every memoized bubble saw a new
prop and re-rendered. Stripped messages were rebuilt too, so the shared
MessageBubble's own memo missed for any tool-calling message. Reuse the map when
the winners are unchanged and cache the stripped message per identity.

Author: kathiekiwi

apps/webapp/app/components/dashboard-agent/DashboardAgentMessages.tsx

@@ -1,17 +1,23 @@
 import type { UIMessage } from "@ai-sdk/react";
-import { memo } from "react";
+import { memo, useMemo, useRef } from "react";
 import { Spinner } from "~/components/primitives/Spinner";
 import { MessageBubble, renderPart } from "~/components/runs/v3/agent/AgentMessageView";
 import { useAutoScrollToBottom } from "~/hooks/useAutoScrollToBottom";
+import { reuseWinners } from "./investigation-winners";
 import { ViewBlocks } from "./view-catalog";
 
-// The shared MessageBubble renders `step-start` parts as a dashed "step"
-// separator — useful in the run inspector / playground, just noise in this
-// simple chat. Drop them before rendering (reference preserved when there are
-// none, so memoization still holds for those messages).
+// The shared MessageBubble renders `step-start` parts as a dashed "step" separator —
+// useful in the run inspector / playground, just noise in this simple chat.
+// Cached so a stripped message keeps its identity across renders and memoization holds.
+const strippedMessages = new WeakMap<UIMessage, UIMessage>();
+
 function stripStepParts(message: UIMessage): UIMessage {
   if (!message.parts?.some((p) => p.type === "step-start")) return message;
-  return { ...message, parts: message.parts.filter((p) => p.type !== "step-start") };
+  const cached = strippedMessages.get(message);
+  if (cached) return cached;
+  const stripped = { ...message, parts: message.parts.filter((p) => p.type !== "step-start") };
+  strippedMessages.set(message, stripped);
+  return stripped;
 }
 
 // A completed render_view tool part carries a `{ blocks }` view spec the agent
@@ -65,6 +71,14 @@ export function winningInvestigationOccurrences(messages: UIMessage[]): Map<stri
   return new Map([...best.entries()].map(([id, w]) => [id, w.occurrence]));
 }
 
+// The stable identity is the point: a fresh `Map` re-renders the whole transcript per token.
+function useInvestigationWinners(messages: UIMessage[]): Map<string, string> {
+  const previous = useRef<Map<string, string>>();
+  const next = useMemo(() => winningInvestigationOccurrences(messages), [messages]);
+  previous.current = reuseWinners(previous.current, next);
+  return previous.current;
+}
+
 function withoutSupersededInvestigations(
   blocks: unknown[],
   occurrence: string,
@@ -127,15 +141,17 @@ export function DashboardAgentMessages({
   error?: Error;
 }) {
   const rootRef = useAutoScrollToBottom([messages, isThinking]);
-  const investigationWinners = winningInvestigationOccurrences(messages);
+  // Must be the exact parts the bubbles render: the winners map keys by part index.
+  const stripped = useMemo(() => messages.map(stripStepParts), [messages]);
+  const investigationWinners = useInvestigationWinners(stripped);
 
   return (
     <div className="flex-1 overflow-y-auto scrollbar-thin scrollbar-track-transparent scrollbar-thumb-surface-control">
       <div ref={rootRef} className="space-y-4 p-4">
-        {messages.map((message) => (
+        {stripped.map((message) => (
           <MemoizedMessageBubble
             key={message.id}
-            message={stripStepParts(message)}
+            message={message}
             investigationWinners={investigationWinners}
           />
         ))}

apps/webapp/app/components/dashboard-agent/investigation-winners.test.ts

@@ -0,0 +1,45 @@
+import { readFileSync } from "node:fs";
+import { join } from "node:path";
+import { describe, expect, it } from "vitest";
+import { reuseWinners, sameOccurrences } from "./investigation-winners";
+
+describe("reuseWinners", () => {
+  it("keeps the previous map when the winners are unchanged", () => {
+    const first = new Map([["inv_1", "msg_a:0"]]);
+    const second = new Map([["inv_1", "msg_a:0"]]);
+    expect(sameOccurrences(first, second)).toBe(true);
+    expect(reuseWinners(first, second)).toBe(first);
+  });
+
+  it("takes the next map when a winner moves", () => {
+    const first = new Map([["inv_1", "msg_a:0"]]);
+    const moved = new Map([["inv_1", "msg_b:2"]]);
+    expect(reuseWinners(first, moved)).toBe(moved);
+  });
+
+  it("takes the next map when an investigation appears", () => {
+    const first = new Map([["inv_1", "msg_a:0"]]);
+    const grown = new Map([
+      ["inv_1", "msg_a:0"],
+      ["inv_2", "msg_b:1"],
+    ]);
+    expect(reuseWinners(first, grown)).toBe(grown);
+  });
+
+  it("takes the next map on the first render", () => {
+    const only = new Map([["inv_1", "msg_a:0"]]);
+    expect(reuseWinners(undefined, only)).toBe(only);
+  });
+});
+
+// Structural: there is no jsdom here, so the wiring is asserted against the source.
+describe("DashboardAgentMessages wiring", () => {
+  const source = readFileSync(join(__dirname, "DashboardAgentMessages.tsx"), "utf8");
+
+  it("stabilises the winners map and the stripped messages it renders", () => {
+    expect(source).toContain("reuseWinners(previous.current, next)");
+    expect(source).toContain("useInvestigationWinners(stripped)");
+    expect(source).toContain("useMemo(() => messages.map(stripStepParts), [messages])");
+    expect(source).toContain("strippedMessages.set(message, stripped)");
+  });
+});

apps/webapp/app/components/dashboard-agent/investigation-winners.ts

@@ -0,0 +1,14 @@
+export function sameOccurrences(a: Map<string, string>, b: Map<string, string>): boolean {
+  if (a.size !== b.size) return false;
+  for (const [id, occurrence] of a) {
+    if (b.get(id) !== occurrence) return false;
+  }
+  return true;
+}
+
+export function reuseWinners(
+  previous: Map<string, string> | undefined,
+  next: Map<string, string>
+): Map<string, string> {
+  return previous && sameOccurrences(previous, next) ? previous : next;
+}