CHANGELOG.md

@@ -78,6 +78,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- **AI chat: the desktop app no longer freezes at 100% CPU on long agent runs.**
+  `useChat` re-rendered on every streamed token and `MessageItem`/`ReasoningBlock`
+  re-parsed the whole transcript markdown (marked + DOMPurify) on every delta, so
+  per-turn work grew quadratically and saturated the main thread. The stream is now
+  throttled (`experimental_throttle`) to ~20 Hz and each finalized message row /
+  markdown part / reasoning block is memoized, so a long turn no longer re-parses
+  already-finished content. (#182)
 - **Editor: caret/selection landed on the wrong line when clicking inside code
   blocks and footnotes.** The affected NodeViews rendered their non-editable
   chrome (language menu, footnotes heading, footnote number marker) before the

apps/client/src/features/ai-chat/components/message-item-memo.test.tsx

@@ -0,0 +1,81 @@
+import { describe, expect, it, vi } from "vitest";
+import { render } from "@testing-library/react";
+import { MantineProvider } from "@mantine/core";
+import type { UIMessage } from "@ai-sdk/react";
+
+// Stub react-i18next (the component reads `useTranslation`). Mirrors the stub in
+// reasoning-block.test.tsx.
+vi.mock("react-i18next", () => ({
+  useTranslation: () => ({ t: (key: string) => key }),
+}));
+
+// Spy on `renderChatMarkdown` so we can count parse calls per text. We keep every
+// OTHER named export of markdown.ts intact via `importActual`, and override only
+// `renderChatMarkdown` with a `vi.fn()` that returns simple HTML so the component
+// still renders. This is the seam that proves the MarkdownPart memo works: a
+// finalized text part must NOT be re-parsed on a later streamed delta.
+// `vi.hoisted` so the spy exists when the hoisted `vi.mock` factory runs.
+const { renderChatMarkdownSpy } = vi.hoisted(() => ({
+  renderChatMarkdownSpy: vi.fn((text: string) => `<p>${text}</p>`),
+}));
+vi.mock("@/features/ai-chat/utils/markdown.ts", async () => {
+  const actual = await vi.importActual<
+    typeof import("@/features/ai-chat/utils/markdown.ts")
+  >("@/features/ai-chat/utils/markdown.ts");
+  return { ...actual, renderChatMarkdown: renderChatMarkdownSpy };
+});
+
+import MessageItem from "./message-item";
+
+// matchMedia (read by MantineProvider) is stubbed globally in vitest.setup.ts.
+
+const msg = (parts: UIMessage["parts"]): UIMessage =>
+  ({ id: "m1", role: "assistant", parts }) as UIMessage;
+
+const renderRow = (message: UIMessage) =>
+  render(
+    <MantineProvider>
+      <MessageItem message={message} />
+    </MantineProvider>,
+  );
+
+/** Count how many spy calls parsed exactly `text` (filtering by the first arg). */
+const callsFor = (text: string) =>
+  renderChatMarkdownSpy.mock.calls.filter((c) => c[0] === text).length;
+
+describe("MessageItem markdown memoization", () => {
+  it("does not re-parse finalized text parts when only a tail part grows", () => {
+    renderChatMarkdownSpy.mockClear();
+
+    // Two finalized text parts.
+    const first = msg([
+      { type: "text", text: "alpha" },
+      { type: "text", text: "beta" },
+    ]);
+    const { rerender } = renderRow(first);
+
+    // Both finalized parts parsed exactly once on the initial render.
+    expect(callsFor("alpha")).toBe(1);
+    expect(callsFor("beta")).toBe(1);
+
+    // A streamed delta: a NEW message object where only a third tail part grows;
+    // the first two parts' text is byte-identical.
+    const next = msg([
+      { type: "text", text: "alpha" },
+      { type: "text", text: "beta" },
+      { type: "text", text: "gamm" },
+    ]);
+    rerender(
+      <MantineProvider>
+        <MessageItem message={next} />
+      </MantineProvider>,
+    );
+
+    // The finalized parts hit the MarkdownPart memo: still parsed at most once
+    // each across BOTH renders (the resilient invariant). The only new parse is
+    // for the changed/added tail part.
+    expect(callsFor("alpha")).toBe(1);
+    expect(callsFor("beta")).toBe(1);
+    expect(callsFor("gamm")).toBe(1);
+  });
+});

apps/client/src/features/ai-chat/components/message-item.test.ts

@@ -0,0 +1,81 @@
+import { describe, expect, it, vi } from "vitest";
+import type { UIMessage } from "@ai-sdk/react";
+
+// Stub react-i18next: importing the component module pulls in `useTranslation`,
+// and we only exercise the pure `arePropsEqual` comparator (no rendering), so a
+// minimal `t` that echoes the key is enough. Mirrors the stub in
+// reasoning-block.test.tsx.
+vi.mock("react-i18next", () => ({
+  useTranslation: () => ({ t: (key: string) => key }),
+}));
+
+import { arePropsEqual } from "./message-item";
+
+/**
+ * Tests for `arePropsEqual`, the `React.memo` comparator for MessageItem. It must
+ * return false on any visible prop/content change (so the row re-renders) and
+ * true when nothing visible changed (so a finalized row is skipped). A FIXED
+ * message id is used so a content-identical clone yields an equal signature.
+ */
+const msg = (
+  parts: UIMessage["parts"],
+  metadata?: unknown,
+): UIMessage =>
+  ({
+    id: "m1",
+    role: "assistant",
+    parts,
+    metadata,
+  }) as UIMessage;
+
+const props = (
+  message: UIMessage,
+  over: Record<string, unknown> = {},
+) => ({
+  message,
+  showCitations: true,
+  neutralizeInternalLinks: false,
+  assistantName: "AI",
+  ...over,
+});
+
+describe("arePropsEqual", () => {
+  it("returns false when showCitations differs", () => {
+    const m = msg([{ type: "text", text: "answer" }]);
+    expect(
+      arePropsEqual(props(m), props(m, { showCitations: false })),
+    ).toBe(false);
+  });
+
+  it("returns false when neutralizeInternalLinks differs", () => {
+    const m = msg([{ type: "text", text: "answer" }]);
+    expect(
+      arePropsEqual(props(m), props(m, { neutralizeInternalLinks: true })),
+    ).toBe(false);
+  });
+
+  it("returns false when assistantName differs", () => {
+    const m = msg([{ type: "text", text: "answer" }]);
+    expect(
+      arePropsEqual(props(m), props(m, { assistantName: "Other" })),
+    ).toBe(false);
+  });
+
+  it("returns true on the identity fast path (same message object, equal props)", () => {
+    const m = msg([{ type: "text", text: "answer" }]);
+    expect(arePropsEqual(props(m), props(m))).toBe(true);
+  });
+
+  it("returns true for the same content in a different message object", () => {
+    const a = msg([{ type: "text", text: "answer" }]);
+    const b = msg([{ type: "text", text: "answer" }]);
+    expect(a).not.toBe(b);
+    expect(arePropsEqual(props(a), props(b))).toBe(true);
+  });
+
+  it("returns false when content changed in a different message object", () => {
+    const a = msg([{ type: "text", text: "answer" }]);
+    const b = msg([{ type: "text", text: "answer grown" }]);
+    expect(arePropsEqual(props(a), props(b))).toBe(false);
+  });
+});

apps/client/src/features/ai-chat/components/message-item.tsx

@@ -11,6 +11,7 @@ import { assistantMessageHasVisibleContent } from "@/features/ai-chat/utils/mess
 import { renderChatMarkdown } from "@/features/ai-chat/utils/markdown.ts";
 import { resolveAssistantName } from "@/features/ai-chat/utils/assistant-name.ts";
 import { reasoningTokensForPart } from "@/features/ai-chat/utils/reasoning-tokens.ts";
+import { messageSignature } from "@/features/ai-chat/utils/message-signature.ts";
 import { describeChatError } from "@/features/ai-chat/utils/error-message.ts";
 import classes from "@/features/ai-chat/components/ai-chat.module.css";
 
@@ -201,50 +202,11 @@ function MessageItem({
   );
 }
 
-/** Cheap content signature for one message: changes iff something VISIBLE in the
- *  row changed. Streaming is APPEND-ONLY (text parts only grow, parts are only
- *  appended, a tool/text part flips state once), so a per-part [type, text
- *  length, state, error/output presence] tuple + the persisted metadata
- *  (error/finishReason) is a sufficient change signal without comparing full
- *  strings on every delta. */
-function messageSignature(message: UIMessage): string {
-  const parts = message.parts
-    .map((p) => {
-      const any = p as {
-        type: string;
-        text?: string;
-        state?: string;
-        errorText?: string;
-        output?: unknown;
-      };
-      return [
-        any.type,
-        any.text?.length ?? 0,
-        any.state ?? "",
-        any.errorText ? 1 : 0,
-        any.output !== undefined ? 1 : 0,
-      ].join(":");
-    })
-    .join("|");
-  const meta = message.metadata as
-    | { error?: string; finishReason?: string; usage?: { reasoningTokens?: number } }
-    | undefined;
-  // `usage.reasoningTokens` is neither append-only nor part-bound: the authoritative
-  // turn total arrives on the final `finish-step` AFTER the reasoning text length and
-  // state are already frozen. Without it in the signature the row's signature would be
-  // unchanged at that point and the re-render skipped, so the "Thinking · N tokens"
-  // header (reasoningTokensForPart) would keep the live estimate instead of snapping
-  // to the exact figure.
-  return `${message.id}#${message.role}#${parts}#${meta?.error ?? ""}#${
-    meta?.finishReason ?? ""
-  }#${meta?.usage?.reasoningTokens ?? ""}`;
-}
-
 /** Skip re-rendering a message whose visible content is unchanged. The streaming
  *  TAIL message gets a fresh object whose signature changes each delta, so it
  *  still re-renders and streams in; every FINALIZED message is skipped, turning a
  *  per-token whole-transcript re-render into a tail-only one. */
-function arePropsEqual(
+export function arePropsEqual(
   prev: MessageItemProps,
   next: MessageItemProps,
 ): boolean {

apps/client/src/features/ai-chat/utils/message-signature.test.ts

@@ -0,0 +1,129 @@
+import { describe, expect, it } from "vitest";
+import type { UIMessage } from "@ai-sdk/react";
+import { messageSignature } from "@/features/ai-chat/utils/message-signature.ts";
+
+/**
+ * Pure-helper tests for `messageSignature`, the cheap per-message content
+ * signature that drives MessageItem's memo (a streaming row's signature must
+ * change on every delta so it re-renders, while a finalized row's stays stable
+ * so it is skipped). Each test exercises ONE change signal and asserts it flips
+ * the signature; a content-identical clone must keep an EQUAL signature.
+ *
+ * The signature embeds `message.id` and `message.role`, so the `msg` factory
+ * uses a FIXED id/role here (not `Math.random()`): otherwise two messages with
+ * identical content would get different signatures and the negative case would
+ * be impossible to express.
+ */
+const msg = (
+  parts: UIMessage["parts"],
+  metadata?: unknown,
+): UIMessage =>
+  ({
+    id: "m1",
+    role: "assistant",
+    parts,
+    metadata,
+  }) as UIMessage;
+
+describe("messageSignature", () => {
+  it("changes when a text part grows", () => {
+    const before = msg([{ type: "text", text: "alpha" }]);
+    const after = msg([{ type: "text", text: "alpha beta" }]);
+    expect(messageSignature(before)).not.toBe(messageSignature(after));
+  });
+
+  it("changes when a new part is appended", () => {
+    const before = msg([{ type: "text", text: "alpha" }]);
+    const after = msg([
+      { type: "text", text: "alpha" },
+      { type: "text", text: "beta" },
+    ]);
+    expect(messageSignature(before)).not.toBe(messageSignature(after));
+  });
+
+  it("changes when a part's state flips", () => {
+    const before = msg([
+      { type: "tool-getPage", state: "input-streaming" } as never,
+    ]);
+    const after = msg([
+      { type: "tool-getPage", state: "output-available" } as never,
+    ]);
+    expect(messageSignature(before)).not.toBe(messageSignature(after));
+  });
+
+  it("changes when a tool part gains an output", () => {
+    const before = msg([
+      { type: "tool-getPage", state: "output-available" } as never,
+    ]);
+    const after = msg([
+      {
+        type: "tool-getPage",
+        state: "output-available",
+        output: { ok: true },
+      } as never,
+    ]);
+    expect(messageSignature(before)).not.toBe(messageSignature(after));
+  });
+
+  it("changes when a part gains an errorText", () => {
+    const before = msg([
+      { type: "tool-getPage", state: "output-error" } as never,
+    ]);
+    const after = msg([
+      {
+        type: "tool-getPage",
+        state: "output-error",
+        errorText: "boom",
+      } as never,
+    ]);
+    expect(messageSignature(before)).not.toBe(messageSignature(after));
+  });
+
+  it("changes when usage.reasoningTokens arrives on finish-step (text/state already frozen)", () => {
+    // The specifically-commented edge case: the authoritative turn total lands on
+    // the final finish-step AFTER the reasoning text length and state are frozen.
+    // Only the token count appears between these two snapshots, so the signature
+    // MUST still flip — otherwise the "Thinking · N tokens" header would never
+    // snap from the live estimate to the exact figure.
+    const before = msg([
+      { type: "reasoning", text: "thinking", state: "done" } as never,
+    ]);
+    const after = msg(
+      [{ type: "reasoning", text: "thinking", state: "done" } as never],
+      { usage: { reasoningTokens: 42 } },
+    );
+    expect(messageSignature(before)).not.toBe(messageSignature(after));
+  });
+
+  it("changes when metadata.error appears", () => {
+    const before = msg([{ type: "text", text: "answer" }]);
+    const after = msg([{ type: "text", text: "answer" }], { error: "boom" });
+    expect(messageSignature(before)).not.toBe(messageSignature(after));
+  });
+
+  it("changes when metadata.finishReason changes (e.g. to 'aborted')", () => {
+    const before = msg([{ type: "text", text: "answer" }], {
+      finishReason: "stop",
+    });
+    const after = msg([{ type: "text", text: "answer" }], {
+      finishReason: "aborted",
+    });
+    expect(messageSignature(before)).not.toBe(messageSignature(after));
+  });
+
+  it("is UNCHANGED for a content-identical clone (different object, same values)", () => {
+    // A finalized row that is re-created as a fresh object (different parts array
+    // by reference, same parts by value) must keep an EQUAL signature, so the
+    // memo skips re-rendering it.
+    const a = msg([
+      { type: "text", text: "alpha" },
+      { type: "tool-getPage", state: "output-available", output: { ok: true } } as never,
+    ]);
+    const b = msg([
+      { type: "text", text: "alpha" },
+      { type: "tool-getPage", state: "output-available", output: { ok: true } } as never,
+    ]);
+    expect(a).not.toBe(b);
+    expect(messageSignature(a)).toBe(messageSignature(b));
+  });
+});

apps/client/src/features/ai-chat/utils/message-signature.ts

@@ -0,0 +1,40 @@
+import type { UIMessage } from "@ai-sdk/react";
+
+/** Cheap content signature for one message: changes iff something VISIBLE in the
+ *  row changed. Streaming is APPEND-ONLY (text parts only grow, parts are only
+ *  appended, a tool/text part flips state once), so a per-part [type, text
+ *  length, state, error/output presence] tuple + the persisted metadata
+ *  (error/finishReason) is a sufficient change signal without comparing full
+ *  strings on every delta. */
+export function messageSignature(message: UIMessage): string {
+  const parts = message.parts
+    .map((p) => {
+      const any = p as {
+        type: string;
+        text?: string;
+        state?: string;
+        errorText?: string;
+        output?: unknown;
+      };
+      return [
+        any.type,
+        any.text?.length ?? 0,
+        any.state ?? "",
+        any.errorText ? 1 : 0,
+        any.output !== undefined ? 1 : 0,
+      ].join(":");
+    })
+    .join("|");
+  const meta = message.metadata as
+    | { error?: string; finishReason?: string; usage?: { reasoningTokens?: number } }
+    | undefined;
+  // `usage.reasoningTokens` is neither append-only nor part-bound: the authoritative
+  // turn total arrives on the final `finish-step` AFTER the reasoning text length and
+  // state are already frozen. Without it in the signature the row's signature would be
+  // unchanged at that point and the re-render skipped, so the "Thinking · N tokens"
+  // header (reasoningTokensForPart) would keep the live estimate instead of snapping
+  // to the exact figure.
+  return `${message.id}#${message.role}#${parts}#${meta?.error ?? ""}#${
+    meta?.finishReason ?? ""
+  }#${meta?.usage?.reasoningTokens ?? ""}`;
+}