diff --git a/CHANGELOG.md b/CHANGELOG.md index cca22374..64872489 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/chat-thread.tsx b/apps/client/src/features/ai-chat/components/chat-thread.tsx index f021cfa3..af797947 100644 --- a/apps/client/src/features/ai-chat/components/chat-thread.tsx +++ b/apps/client/src/features/ai-chat/components/chat-thread.tsx @@ -29,6 +29,14 @@ import { } from "@/features/ai-chat/utils/queue-helpers.ts"; import classes from "@/features/ai-chat/components/ai-chat.module.css"; +// Throttle how often the streamed `messages` state triggers a re-render. Without +// it, useChat updates state on EVERY token, so the whole transcript's markdown +// (marked + DOMPurify) is re-parsed per token — on a long agent run that grows +// into a quadratic CPU storm that pins the main thread and freezes the UI. +// ~50ms (20 Hz) keeps streaming visually smooth while decoupling re-render cost +// from the token rate. +const STREAM_THROTTLE_MS = 50; + /** The page the user is currently viewing, sent as chat context. */ export interface OpenPageContext { id: string; @@ -254,6 +262,8 @@ export default function ChatThread({ id: chatStoreId, messages: initialMessages, transport, + // See STREAM_THROTTLE_MS — bounds re-render/markdown-reparse frequency. + experimental_throttle: STREAM_THROTTLE_MS, // `onFinish` (ai@6 useChat) fires from a `finally` on EVERY terminal outcome // — success, user Stop/abort (`isAbort`), network drop (`isDisconnect`), and // stream error (`isError`). Keep calling `onTurnFinished()` on all of them 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..dfed46f4 --- /dev/null +++ b/apps/client/src/features/ai-chat/components/message-item.test.ts @@ -0,0 +1,73 @@ +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"]): UIMessage => + ({ id: "m1", role: "assistant", parts }) 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 6436b4d6..6bd4374d 100644 --- a/apps/client/src/features/ai-chat/components/message-item.tsx +++ b/apps/client/src/features/ai-chat/components/message-item.tsx @@ -1,3 +1,4 @@ +import { memo } from "react"; import { Box, Text } from "@mantine/core"; import { useTranslation } from "react-i18next"; import type { UIMessage } from "@ai-sdk/react"; @@ -10,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"; @@ -34,6 +36,39 @@ interface MessageItemProps { assistantName?: string; } +/** + * One assistant text part rendered as sanitized markdown. Memoized on its inputs + * so a finalized text part is NOT re-parsed on every streamed delta: during a + * turn only the actively-growing tail part changes its `text`, so every earlier + * part hits the memo and skips the expensive marked + DOMPurify pass. Props are + * primitives, so React.memo's default shallow compare is exactly right (the + * `text` string is compared by value). + */ +const MarkdownPart = memo(function MarkdownPart({ + text, + neutralizeInternalLinks, +}: { + text: string; + neutralizeInternalLinks: boolean; +}) { + const html = renderChatMarkdown(text, { neutralizeInternalLinks }); + if (html) { + return ( +
+ ); + } + // Fallback when markdown could not render synchronously: raw text. + return ( + + {text} + + ); +}); + /** * Render a single UIMessage by iterating its `parts`: * - `text` parts -> sanitized markdown. @@ -41,12 +76,13 @@ interface MessageItemProps { * Other part kinds (reasoning, sources, files, step-start) are ignored for v1. * User messages render their text as a right-aligned plain bubble. * - * This component is intentionally NOT memoized: `useChat` replaces the streaming - * assistant message with a freshly cloned object on every streamed delta, so the - * `message` prop identity (and its `parts`) changes each tick. Re-rendering the - * text parts on each delta is what makes the answer stream in progressively. + * This component is memoized (see `arePropsEqual` at the bottom) on a cheap + * per-message content signature: the streaming TAIL message's signature changes + * on each delta so it still re-renders and streams in, while finalized rows are + * skipped. Each text part's markdown is itself memoized via `MarkdownPart`, so a + * long turn no longer re-parses the whole transcript on every token. */ -export default function MessageItem({ +function MessageItem({ message, showCitations = true, neutralizeInternalLinks = false, @@ -109,24 +145,12 @@ export default function MessageItem({ // starts with an empty text part before the first token arrives); the // typing indicator covers that gap until real content streams in. if (!part.text.trim()) return null; - const html = renderChatMarkdown(part.text, { - neutralizeInternalLinks, - }); - if (html) { - return ( -
- ); - } - // Fallback when markdown could not render synchronously: raw text. return ( - - {part.text} - + ); } @@ -177,3 +201,26 @@ export default function MessageItem({ ); } + +/** 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. */ +export function arePropsEqual( + prev: MessageItemProps, + next: MessageItemProps, +): boolean { + if ( + prev.showCitations !== next.showCitations || + prev.neutralizeInternalLinks !== next.neutralizeInternalLinks || + prev.assistantName !== next.assistantName + ) { + return false; + } + // Fast path: identical message object (finalized rows keep their identity + // across deltas) — skip without building signatures. + if (prev.message === next.message) return true; + return messageSignature(prev.message) === messageSignature(next.message); +} + +export default memo(MessageItem, arePropsEqual); diff --git a/apps/client/src/features/ai-chat/components/reasoning-block.tsx b/apps/client/src/features/ai-chat/components/reasoning-block.tsx index de35229a..cb3335f4 100644 --- a/apps/client/src/features/ai-chat/components/reasoning-block.tsx +++ b/apps/client/src/features/ai-chat/components/reasoning-block.tsx @@ -1,4 +1,4 @@ -import { useState } from "react"; +import { memo, useMemo, useState } from "react"; import { Box, Collapse, Group, Text, UnstyledButton } from "@mantine/core"; import { IconChevronDown } from "@tabler/icons-react"; import { useTranslation } from "react-i18next"; @@ -27,19 +27,23 @@ interface ReasoningBlockProps { * Providers that don't stream reasoning TEXT still render this block from the * authoritative count alone (header only, empty body) so the cost is visible. */ -export default function ReasoningBlock({ text, tokens }: ReasoningBlockProps) { +function ReasoningBlock({ text, tokens }: ReasoningBlockProps) { const { t } = useTranslation(); const [open, setOpen] = useState(false); // Authoritative count wins; otherwise estimate live from the streamed text. const count = tokens && tokens > 0 ? tokens : estimateTokens(text); const trimmed = text.trim(); - // Collapse the blank-line gaps the model emits between every list item / - // paragraph so the reasoning renders compactly (tight lists, joined - // paragraphs) — see collapseBlankLines. ONLY here, not in the normal answer. - const html = trimmed - ? renderChatMarkdown(collapseBlankLines(trimmed), {}) - : ""; + // Memoize the markdown render so toggling `open` (or a parent re-render caused + // by an unrelated streamed delta) does not re-parse the reasoning text; it + // recomputes only when the reasoning text itself changes (while it streams in). + // collapseBlankLines collapses the blank-line gaps the model emits between every + // list item / paragraph so the reasoning renders compactly (tight lists, joined + // paragraphs) — ONLY here, not in the normal answer. + const html = useMemo( + () => (trimmed ? renderChatMarkdown(collapseBlankLines(trimmed), {}) : ""), + [trimmed], + ); return ( @@ -87,3 +91,8 @@ export default function ReasoningBlock({ text, tokens }: ReasoningBlockProps) { ); } + +// Memoized: re-renders only when `text`/`tokens` change (primitive props, default +// shallow compare), so a parent re-render during streaming of OTHER content does +// not re-run the markdown parse for an already-finalized reasoning block. +export default memo(ReasoningBlock); 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..7c4f7a70 --- /dev/null +++ b/apps/client/src/features/ai-chat/utils/message-signature.test.ts @@ -0,0 +1,241 @@ +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)); + }); +}); + +/** + * Per-part-kind coupling guard for the load-bearing invariant documented at the + * top of message-signature.ts: the signature MUST sample every VISIBLE field the + * MessageItem render body draws, or the memo freezes a stale row. This is an + * executable lock for the part kinds rendered TODAY — read alongside + * `MessageItem` (message-item.tsx) and the `assistantMessageHasVisibleContent` + * helper (message-content.ts), which "mirrors MessageItem's render decisions + * EXACTLY". For each kind, mutating a field the render body DRAWS must flip the + * signature. If a new visible field is rendered without being added here AND to + * the signature, the corresponding assertion below should fail — that is the + * guard. (This intentionally stops short of the render-descriptor refactor: + * adding a part kind or a visible field still requires a human to extend both + * the signature and this block.) + */ +describe("messageSignature ↔ render coupling (per visible part kind)", () => { + describe("text part — render draws part.text (MarkdownPart text={part.text})", () => { + it("flips when the visible text changes", () => { + // Streaming is append-only, so the visible text only grows; the signature + // samples its length, so the growth is the change signal. + const before = msg([{ type: "text", text: "answer" }]); + const after = msg([{ type: "text", text: "answer extended" }]); + expect(messageSignature(before)).not.toBe(messageSignature(after)); + }); + }); + + describe("reasoning part — render draws text + tokens (ReasoningBlock)", () => { + it("flips when the visible reasoning text changes", () => { + const before = msg([ + { type: "reasoning", text: "think", state: "streaming" } as never, + ]); + const after = msg([ + { type: "reasoning", text: "think harder", state: "streaming" } as never, + ]); + expect(messageSignature(before)).not.toBe(messageSignature(after)); + }); + + it("flips when the visible token count (metadata.usage.reasoningTokens) lands", () => { + // The header's "Thinking · N tokens" reads reasoningTokensForPart, fed by + // metadata.usage.reasoningTokens — a VISIBLE field that arrives on the final + // finish-step after text length and state are frozen. + const before = msg([ + { type: "reasoning", text: "think", state: "done" } as never, + ]); + const after = msg( + [{ type: "reasoning", text: "think", state: "done" } as never], + { usage: { reasoningTokens: 99 } }, + ); + expect(messageSignature(before)).not.toBe(messageSignature(after)); + }); + }); + + describe("tool-* part — render draws state/errorText/citations (ToolCallCard)", () => { + it("flips when the run state changes (running ↔ done icon + label)", () => { + // toolRunState(part.state) selects the spinner/check/error icon. + const before = msg([ + { type: "tool-getPage", state: "input-available" } as never, + ]); + const after = msg([ + { type: "tool-getPage", state: "output-available" } as never, + ]); + expect(messageSignature(before)).not.toBe(messageSignature(after)); + }); + + it("flips when output arrives (drives the rendered citation links)", () => { + // toolCitations reads part.output to render the "/p/{id}" anchors. + const before = msg([ + { type: "tool-getPage", state: "output-available" } as never, + ]); + const after = msg([ + { + type: "tool-getPage", + state: "output-available", + output: { id: "page-1", title: "Doc" }, + } as never, + ]); + expect(messageSignature(before)).not.toBe(messageSignature(after)); + }); + + it("flips when errorText appears (the visible red error detail line)", () => { + const before = msg([ + { type: "tool-getPage", state: "output-error" } as never, + ]); + const after = msg([ + { + type: "tool-getPage", + state: "output-error", + errorText: "permission denied", + } as never, + ]); + expect(messageSignature(before)).not.toBe(messageSignature(after)); + }); + }); + + describe("metadata banners — render draws error / aborted notices", () => { + it("flips when metadata.error appears (ChatErrorAlert banner)", () => { + const before = msg([{ type: "text", text: "answer" }]); + const after = msg([{ type: "text", text: "answer" }], { error: "boom" }); + expect(messageSignature(before)).not.toBe(messageSignature(after)); + }); + + it("flips when metadata.finishReason becomes 'aborted' (ChatStoppedNotice)", () => { + 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)); + }); + }); +}); 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..84c37919 --- /dev/null +++ b/apps/client/src/features/ai-chat/utils/message-signature.ts @@ -0,0 +1,44 @@ +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. WARNING — load-bearing for the MessageItem memo: + * if a future part kind's VISIBLE content can change WITHOUT changing [type, + * text length, state, error/output presence] (e.g. a tool that streams + * `preliminary` output, or a client-side regenerate that edits a finalized + * row in place), extend this signature or the memo will freeze a stale row. */ +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 ?? ""}`; +}