diff --git a/CHANGELOG.md b/CHANGELOG.md index 6c2aa9c9..149f2810 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/apps/client/src/features/ai-chat/components/message-item-memo.test.tsx b/apps/client/src/features/ai-chat/components/message-item-memo.test.tsx new file mode 100644 index 00000000..06c0c5fb --- /dev/null +++ b/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) => `

${text}

`), +})); +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( + + + , + ); + +/** 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( + + + , + ); + + // 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); + }); +}); diff --git a/apps/client/src/features/ai-chat/components/message-item.test.ts b/apps/client/src/features/ai-chat/components/message-item.test.ts new file mode 100644 index 00000000..decc2d2c --- /dev/null +++ b/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 = {}, +) => ({ + 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); + }); +}); diff --git a/apps/client/src/features/ai-chat/components/message-item.tsx b/apps/client/src/features/ai-chat/components/message-item.tsx index 0eabbd87..6bd4374d 100644 --- a/apps/client/src/features/ai-chat/components/message-item.tsx +++ b/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 { diff --git a/apps/client/src/features/ai-chat/utils/message-signature.test.ts b/apps/client/src/features/ai-chat/utils/message-signature.test.ts new file mode 100644 index 00000000..7df210f6 --- /dev/null +++ b/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)); + }); +}); diff --git a/apps/client/src/features/ai-chat/utils/message-signature.ts b/apps/client/src/features/ai-chat/utils/message-signature.ts new file mode 100644 index 00000000..69c93eba --- /dev/null +++ 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 ?? ""}`; +}