Compare commits

..

5 Commits

Author SHA1 Message Date
claude code agent 227
cb61274187 test(ai-chat): simplify msg factory and lock signature↔render coupling
Address non-blocking review items on the AI-chat stream-perf PR:

- Drop the unused `metadata` param from the `msg` test factory in
  message-item.test.ts; no caller passed it.
- Add a per-part-kind coupling guard to message-signature.test.ts that, for
  each part kind rendered today (text, reasoning, tool-*) plus the metadata
  banners, asserts that mutating a field the MessageItem render body DRAWS
  flips messageSignature — an executable lock for the load-bearing memo
  invariant documented in message-signature.ts.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-26 16:57:31 +03:00
claude_code
eafd15f0ef docs(ai-chat): document load-bearing invariant of messageSignature memo
PR #182 review (post-fix pass) surfaced two latent correctness risks in the
new MessageItem memo: the per-message signature tracks only [type, text length,
state, error/output presence] + metadata, so a part kind whose VISIBLE content
can change WITHOUT changing those fields would silently freeze a stale row.
Neither is reachable with the current toolset (tool output is set once;
streaming is append-only with a fixed id), so the correct fix is to harden the
documented invariant rather than hash output content on every delta (getPage
returns full page content — hashing it per-delta would tax the hot path this
PR optimizes).

Add a WARNING in messageSignature naming the two future triggers (a tool that
streams `preliminary` output; a client-side regenerate/edit that mutates a
finalized row in place) and the required action (extend the signature).

No behavior change (comment only). vitest src/features/ai-chat 189/189 pass,
tsc clean for the touched files.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-25 23:44:49 +03:00
claude_code
63c26042ba test(review): address PR #182 review — tests + extract messageSignature, CHANGELOG
Resolve the PR #182 code-review (Request changes) on top of the already-merged
develop (the merge commit preserves both the markdown useMemo and the
collapseBlankLines fix in reasoning-block.tsx).

- Extract messageSignature from message-item.tsx into utils/message-signature.ts
  (matches the feature's "pure UIMessage helper + colocated test" convention) and
  export arePropsEqual so the memo seam is unit-testable. No logic change.
- Add utils/message-signature.test.ts covering every change signal (text grows,
  part appended, state flip, output appears, errorText appears, usage.reasoningTokens
  arriving on finish-step, metadata error/finishReason) plus the negative
  content-identical-clone case.
- Add components/message-item.test.ts for arePropsEqual (each prop diff -> false,
  identity fast-path -> true, same-content-different-object -> true, changed -> false).
- Add components/message-item-memo.test.tsx: render-level proof that finalized text
  parts are not re-parsed when only a tail part grows (MarkdownPart memo).
- CHANGELOG: add the user-facing 100% CPU freeze fix under [Unreleased] / Fixed.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-25 22:33:14 +03:00
claude_code
2f058a6e40 Merge remote-tracking branch 'gitea/develop' into fix/ai-chat-stream-perf
# Conflicts:
#	apps/client/src/features/ai-chat/components/reasoning-block.tsx
2026-06-25 22:21:41 +03:00
claude_code
99d0cb8773 perf(ai-chat): throttle stream + memoize markdown to stop CPU spikes on long runs
On long agent runs (dozens of tool calls) the desktop app froze at 100% CPU with
no user interaction: useChat updated state on every streamed token, and
MessageItem/ReasoningBlock re-parsed the whole transcript's markdown (the marked
pipeline + DOMPurify) on every delta. Per-turn work grew quadratically and
saturated the main thread; the SSE stream drove it, so it hung "on its own".

- chat-thread: pass experimental_throttle (50ms) to useChat so the streamed
  messages state re-renders at most ~20 Hz instead of once per token.
- message-item: memoize MessageItem on a cheap per-message content signature
  (the streaming tail still re-renders as it grows; finalized rows are skipped),
  and render each text part via a memoized MarkdownPart so finalized parts are
  not re-parsed. The signature includes usage.reasoningTokens so the
  authoritative "Thinking - N tokens" count still snaps in at finish-step.
- reasoning-block: memoize the markdown render (useMemo on the text) and wrap the
  component in React.memo.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-25 03:26:44 +03:00
12 changed files with 560 additions and 290 deletions

View File

@@ -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
@@ -92,16 +99,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
no longer froze on the previous step's authoritative usage; the current step's
estimate is combined per-component with `max`, so the count rises smoothly and
never jumps backwards. (#163)
- **AI chat: "New chat" pressed during the first turn's stream now resets the
thread instead of leaving the old turn streaming.** While a brand-new,
not-yet-adopted chat streamed its first turn, hitting "New chat" left
`activeChatId === null` (a no-op for the atom), so the reconciler never
remounted and the in-flight thread kept streaming behind the fresh one — and a
late refetch / late `onFinish` from that abandoned thread could yank the user
back into the chat they just left. "New chat" now forces a fresh empty thread
unconditionally and the finished thread's mount key is checked so a late
callback from an abandoned thread no longer adopts or re-arms the fallback.
(#161)
## [0.93.0] - 2026-06-21

View File

@@ -195,7 +195,6 @@ export default function AiChatWindow() {
waitingForHistory,
onTurnFinished,
onServerChatId,
startFreshThread,
cancelPendingAdoption,
} = useChatSession({
activeChatId,
@@ -210,26 +209,18 @@ export default function AiChatWindow() {
// startNewChat/selectChat set the public atom; the hook's render-phase
// reconciler handles the remount when activeChatId actually CHANGES. But
// pressing "New chat" while already in a not-yet-adopted new chat leaves
// activeChatId === null (a no-op for the atom), so the reconciler never fires —
// startFreshThread forces the remount unconditionally (and disarms any armed
// error-path fallback) so a late refetch can't yank the user into a just-failed
// chat after they chose a fresh one (#161).
// pressing "New chat" while already in a new chat leaves activeChatId === null
// (a no-op for the atom), so the reconciler never fires — explicitly disarm any
// armed error-path fallback here so a late refetch can't yank the user into a
// just-failed chat after they chose a fresh one.
const startNewChat = useCallback((): void => {
// Force a fresh thread UNCONDITIONALLY. On a brand-new, not-yet-adopted chat
// (activeChatId === null) whose first turn is still streaming, setActiveChatId
// (null) is a no-op and the render-phase reconciler would not remount — leaving
// the streaming thread in place (#161). startFreshThread guarantees a clean
// remount and already disarms any armed error-path fallback (so a separate
// cancelPendingAdoption() call here is redundant); the abandoned thread's late
// finish is rejected by the threadKey guard in the session hook.
startFreshThread();
cancelPendingAdoption();
setActiveChatId(null);
setHistoryOpen(false);
setDraft("");
// Default the picker back to "Universal assistant" for the fresh chat.
setSelectedRoleId(null);
}, [startFreshThread, setActiveChatId, setDraft, setSelectedRoleId]);
}, [cancelPendingAdoption, setActiveChatId, setDraft, setSelectedRoleId]);
const selectChat = useCallback(
(chatId: string): void => {
@@ -642,7 +633,6 @@ export default function AiChatWindow() {
onRolePicked={(role) => setSelectedRoleId(role.id)}
assistantName={currentRole?.name}
onTurnFinished={onTurnFinished}
threadKey={threadKey}
onServerChatId={onServerChatId}
onLiveTurnTokens={setLiveTurnTokens}
/>

View File

@@ -1,94 +0,0 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
import { render, act } from "@testing-library/react";
import { MantineProvider } from "@mantine/core";
// Capture the options ChatThread passes to useChat so the test can drive the
// hook's terminal callbacks (here: onError) directly, without a real stream. The
// box is created via vi.hoisted so the hoisted vi.mock factory below can close
// over it.
const { useChatBox } = vi.hoisted(() => ({
useChatBox: { options: null as unknown as Record<string, unknown> | null },
}));
// Mock the AI SDK hook: record the options and return an inert, ready store so
// ChatThread renders without any network/streaming machinery.
vi.mock("@ai-sdk/react", () => ({
useChat: (options: Record<string, unknown>) => {
useChatBox.options = options;
return {
messages: [],
sendMessage: vi.fn(),
status: "ready",
stop: vi.fn(),
error: null,
};
},
}));
// Stub react-i18next so `t` returns the key (other component tests use the same
// pattern); ChatThread's rendered chrome is irrelevant to this wiring test.
vi.mock("react-i18next", () => ({
useTranslation: () => ({ t: (key: string) => key }),
}));
// Mock the heavy presentational children to trivial stubs — this test only
// exercises the onError → onTurnFinished wiring, not their rendering.
vi.mock("@/features/ai-chat/components/message-list.tsx", () => ({
default: () => null,
}));
vi.mock("@/features/ai-chat/components/chat-input.tsx", () => ({
default: () => null,
}));
vi.mock("@/features/ai-chat/components/role-cards.tsx", () => ({
default: () => null,
}));
vi.mock("@/features/ai-chat/components/chat-error-alert.tsx", () => ({
default: () => null,
}));
vi.mock("@/features/ai-chat/components/chat-stopped-notice.tsx", () => ({
default: () => null,
}));
import ChatThread from "./chat-thread";
// matchMedia (read by MantineProvider) is stubbed globally in vitest.setup.ts.
describe("ChatThread onError wiring (#161)", () => {
beforeEach(() => {
useChatBox.options = null;
vi.clearAllMocks();
});
it("onError calls onTurnFinished with (undefined, threadKey) so a late error from an abandoned thread is rejected", () => {
const onTurnFinished = vi.fn();
// Silence the deliberate console.error ChatThread logs for devtools.
const consoleError = vi
.spyOn(console, "error")
.mockImplementation(() => {});
render(
<MantineProvider>
<ChatThread
chatId="c1"
onTurnFinished={onTurnFinished}
threadKey="thread-key-1"
/>
</MantineProvider>,
);
const options = useChatBox.options;
expect(options).not.toBeNull();
expect(typeof options?.onError).toBe("function");
// Drive the captured onError exactly as the AI SDK would on a stream error.
act(() => {
(options!.onError as (e: Error) => void)(new Error("stream blew up"));
});
// The thread's own mount key must be forwarded with NO server id, so the
// session hook can reject this finish if the thread has been abandoned.
expect(onTurnFinished).toHaveBeenCalledWith(undefined, "thread-key-1");
consoleError.mockRestore();
});
});

View File

@@ -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;
@@ -60,12 +68,7 @@ interface ChatThreadProps {
* new chat, adopts the freshly created chat id. `serverChatId` is the
* authoritative id the server streamed on the assistant message metadata, or
* undefined on a failed turn — see adopt-chat-id.ts for the full #137 design. */
onTurnFinished: (serverChatId?: string, finishingThreadKey?: string) => void;
/** This thread's mount key (the same value the parent uses as React `key`).
* Forwarded back through onTurnFinished so the session hook can tell a finish
* from THIS still-mounted thread from a late finish of an abandoned thread the
* user already left via New chat / switch (#161). */
threadKey?: string;
onTurnFinished: (serverChatId?: string) => void;
/** Called EARLY (at the stream's `start` chunk) with the authoritative server
* chat id streamed on the assistant message metadata, so a brand-new chat
* adopts its real id WHILE the first turn is still streaming (#174 — makes the
@@ -121,7 +124,6 @@ export default function ChatThread({
onRolePicked,
assistantName,
onTurnFinished,
threadKey,
onServerChatId,
onLiveTurnTokens,
}: ChatThreadProps) {
@@ -252,6 +254,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
@@ -264,7 +268,7 @@ export default function ChatThread({
// Forward the authoritative server chatId (streamed on the assistant
// message metadata) so the parent adopts the REAL created chat id for a new
// chat — see adopt-chat-id.ts for the full #137 design.
onTurnFinished(extractServerChatId(message), threadKey);
onTurnFinished(extractServerChatId(message));
// Show a neutral "stopped" marker for an aborted turn; the red error banner
// (via `error`) already covers isError, and a clean finish clears any marker.
if (isError) setStopNotice(null);
@@ -285,7 +289,7 @@ export default function ChatThread({
// Surface the raw failure in the browser console (devtools) for debugging;
// the UI separately shows a friendly classified banner (see errorView).
console.error("AI chat stream error:", streamError);
onTurnFinished(undefined, threadKey);
onTurnFinished();
},
});

View File

@@ -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);
});
});

View File

@@ -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<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);
});
});

View File

@@ -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 (
<div
className={classes.markdown}
// Sanitized by renderChatMarkdown (DOMPurify) before insertion.
dangerouslySetInnerHTML={{ __html: html }}
/>
);
}
// Fallback when markdown could not render synchronously: raw text.
return (
<Text className={classes.markdown} style={{ whiteSpace: "pre-wrap" }}>
{text}
</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 (
<div
key={index}
className={classes.markdown}
// Sanitized by renderChatMarkdown (DOMPurify) before insertion.
dangerouslySetInnerHTML={{ __html: html }}
/>
);
}
// Fallback when markdown could not render synchronously: raw text.
return (
<Text key={index} className={classes.markdown} style={{ whiteSpace: "pre-wrap" }}>
{part.text}
</Text>
<MarkdownPart
key={index}
text={part.text}
neutralizeInternalLinks={neutralizeInternalLinks}
/>
);
}
@@ -177,3 +201,26 @@ export default function MessageItem({
</Box>
);
}
/** 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);

View File

@@ -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 (
<Box className={classes.reasoningBlock} mb={6}>
@@ -87,3 +91,8 @@ export default function ReasoningBlock({ text, tokens }: ReasoningBlockProps) {
</Box>
);
}
// 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);

View File

@@ -1,5 +1,5 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
import { renderHook, act } from "@testing-library/react";
import { renderHook } from "@testing-library/react";
import { useChatSession } from "./use-chat-session";
import type { UseChatSessionOptions } from "./use-chat-session";
@@ -120,40 +120,16 @@ describe("useChatSession", () => {
expect(setActiveChatId).not.toHaveBeenCalledWith("new");
});
it("cancelPendingAdoption (selectChat) disarms a late refetch from adopting the just-failed chat", () => {
// cancelPendingAdoption is the explicit disarm the window calls from
// selectChat: switching to a chat whose id == null is a no-op for the atom, so
// the render-phase reconciler never fires and only this call disarms an armed
// error-path fallback. (startNewChat no longer routes through here — it calls
// startFreshThread, covered by the next test — but cancelPendingAdoption still
// backs selectChat, so this guard must hold.)
it("startNewChat while already in a new chat: cancelPendingAdoption stops a late refetch adopting the failed chat", () => {
// The Warning path the render-phase reconciler can't catch: pressing "New
// chat" while already in a new chat keeps activeChatId === null (a no-op for
// the atom), so only the explicit cancelPendingAdoption() disarms.
const { result, rerender, setActiveChatId } = setup({
activeChatId: null,
chats: { items: [{ id: "x" }] },
});
result.current.onTurnFinished(undefined); // first turn failed → arm (before=["x"])
result.current.cancelPendingAdoption(); // window calls this from selectChat
// The just-failed row lands in a late refetch; it must NOT be adopted.
rerender({
activeChatId: null,
chats: { items: [{ id: "x" }, { id: "failed" }] },
});
expect(setActiveChatId).not.toHaveBeenCalledWith("failed");
});
it("#161: startFreshThread disarms the armed error-path fallback (New chat during the first turn)", () => {
// Pressing "New chat" while already in a not-yet-adopted new chat keeps
// activeChatId === null, so the render-phase reconciler never fires. The
// window now calls startFreshThread() (NOT cancelPendingAdoption) to force a
// fresh thread; this test pins the load-bearing fact that startFreshThread
// ALSO nulls pendingNewChatRef, so a late refetch of the just-failed row can't
// yank the user back into the abandoned chat.
const { result, rerender, setActiveChatId } = setup({
activeChatId: null,
chats: { items: [{ id: "x" }] },
});
result.current.onTurnFinished(undefined); // first turn failed → arm (before=["x"])
act(() => result.current.startFreshThread()); // "New chat" → fresh thread + disarm
result.current.cancelPendingAdoption(); // window calls this from startNewChat
// The just-failed row lands in a late refetch; it must NOT be adopted.
rerender({
activeChatId: null,
@@ -251,48 +227,6 @@ describe("useChatSession", () => {
expect(result.current.threadKey).toBe("C");
});
it("#161: startFreshThread remounts even when activeChatId stays null (New chat mid-stream)", () => {
// The bug: pressing New chat while still on a brand-new, not-yet-adopted chat
// leaves activeChatId === null, so the render-phase reconciler never fires.
// startFreshThread must remount UNCONDITIONALLY (a new mount key).
const { result } = setup({ activeChatId: null, chats: { items: [] } });
const keyBefore = result.current.threadKey;
act(() => result.current.startFreshThread());
expect(result.current.threadKey).not.toBe(keyBefore);
});
it("#161: a late finish from an ABANDONED thread does not adopt or invalidate messages", () => {
const {
result,
setActiveChatId,
onInvalidateChatList,
onInvalidateChatMessages,
} = setup({ activeChatId: null, chats: { items: [{ id: "x" }] } });
const abandonedKey = result.current.threadKey;
// User pressed New chat mid-stream → fresh thread (new mount key).
act(() => result.current.startFreshThread());
expect(result.current.threadKey).not.toBe(abandonedKey);
// The left-behind thread's onFinish fires late, carrying ITS (now stale) key
// and the server id of the chat the user just left. It must NOT be adopted.
act(() => result.current.onTurnFinished("A", abandonedKey));
expect(setActiveChatId).not.toHaveBeenCalled();
expect(onInvalidateChatMessages).not.toHaveBeenCalled();
// The abandoned chat should still surface in the history list.
expect(onInvalidateChatList).toHaveBeenCalled();
});
it("#161: a finish from the CURRENT thread (matching key) still adopts", () => {
const { result, setActiveChatId } = setup({
activeChatId: null,
chats: { items: [{ id: "x" }] },
});
// Same thread that is mounted reports its finish with the matching key.
act(() =>
result.current.onTurnFinished("A", result.current.threadKey),
);
expect(setActiveChatId).toHaveBeenCalledWith("A");
});
it("waitingForHistory gates the loader only while opening an unloaded existing chat", () => {
// Open an existing chat whose history is still loading => loader on.
const { result, rerender } = setup({

View File

@@ -32,13 +32,8 @@ export interface UseChatSessionResult {
/** Show the history loader instead of the live thread. */
waitingForHistory: boolean;
/** Call when a turn finishes; `serverChatId` is the authoritative streamed id
* (undefined on a failed turn). `finishingThreadKey` is the mount key of the
* thread that produced this callback — when it no longer matches the mounted
* thread (the user pressed New chat / switched mid-stream), the call is from an
* abandoned thread and must NOT adopt or re-arm the fallback. Omitting it (old
* callers / tests) treats the call as belonging to the current thread. Handles
* new-chat id adoption + invalidations. */
onTurnFinished: (serverChatId?: string, finishingThreadKey?: string) => void;
* (undefined on a failed turn). Handles new-chat id adoption + invalidations. */
onTurnFinished: (serverChatId?: string) => void;
/** Call EARLY (at the stream's `start` chunk) with the authoritative streamed
* chat id so a brand-new chat adopts its real id WHILE its first turn is still
* streaming — making `activeChatId`-gated affordances (e.g. the Copy/export
@@ -46,13 +41,6 @@ export interface UseChatSessionResult {
* no list/messages invalidation — that is left to onTurnFinished at the end).
* Idempotent and a no-op once the chat already has an id. */
onServerChatId: (serverChatId?: string) => void;
/** Force a brand-new, empty thread (new mount key, no chat id) UNCONDITIONALLY.
* The render-phase reconciler only remounts when `activeChatId` actually
* changes; pressing "New chat" while already in a not-yet-adopted new chat
* leaves `activeChatId === null` (a no-op for the atom), so the reconciler
* never fires and the stale streaming thread stays mounted (#161). The window
* calls this from startNewChat to guarantee a fresh thread regardless. */
startFreshThread: () => void;
/** Disarm any pending error-path new-chat fallback. The window calls this from
* startNewChat/selectChat so a late refetch can't yank the user back into a
* just-failed chat after they explicitly moved on. */
@@ -97,14 +85,6 @@ export function useChatSession(
const activeChatIdRef = useRef(activeChatId);
activeChatIdRef.current = activeChatId;
// Live mirror of the mounted thread's key, read by onTurnFinished to tell a
// current-thread finish from an ABANDONED one. ai@6's useChat does not abort
// its request on unmount, and its callbacks are proxied so onFinish/onError of
// a thread the user already left (via New chat / switch) still fire AFTER that
// thread unmounts. By then this ref holds the NEW thread's key, so comparing it
// to the key the finishing thread reports rejects the abandoned turn (#161).
const threadKeyRef = useRef<string>("");
// The mounted thread's identity: ONE atomic value tying ChatThread's mount key
// (`thread.key`) to the chat id that mounted thread holds (`thread.chatId`).
// Consolidating these makes the "key vs chat id diverged" state unrepresentable
@@ -118,10 +98,6 @@ export function useChatSession(
: switchThread(activeChatId),
);
// Keep the live mirror pointed at the currently-mounted thread's key so a late
// onTurnFinished can be matched against it (see threadKeyRef above).
threadKeyRef.current = thread.key;
// Error-path fallback for new-chat id adoption. When a brand-new chat's first
// turn errors BEFORE the server's `start` chunk, no authoritative chatId ever
// reaches the client, so the primary metadata adoption cannot run. We then ARM
@@ -139,21 +115,7 @@ export function useChatSession(
// yet) we adopt the server's AUTHORITATIVE streamed id (never the newest in the
// list, which races a second tab — #137; see adopt-chat-id.ts).
const onTurnFinished = useCallback(
(serverChatId?: string, finishingThreadKey?: string) => {
// Reject a finish from an ABANDONED thread. After the user pressed New chat
// (or switched chats) mid-stream, the left-behind thread's onFinish/onError
// still fire (ai@6 does not abort on unmount). Adopting/arming off that late
// callback would yank the user back into the chat they just left (#161).
// `undefined` (legacy callers/tests) is treated as the current thread.
const isCurrentThread =
finishingThreadKey === undefined ||
finishingThreadKey === threadKeyRef.current;
if (!isCurrentThread) {
// Still surface the abandoned chat in the history list, but do NOT adopt,
// arm the fallback, or invalidate per-chat messages (no thread shows it).
onInvalidateChatList();
return;
}
(serverChatId?: string) => {
// Read the live id from the ref, not the closure: on a failed turn this can
// run twice in one turn (onFinish + onError) before any re-render, and the
// primary branch below updates the ref so the second call sees the adopted id.
@@ -296,29 +258,11 @@ export function useChatSession(
pendingNewChatRef.current = null;
}, []);
// Force a fresh, empty thread regardless of the current `activeChatId`. The
// render-phase reconciler only remounts on an activeChatId CHANGE, so "New chat"
// pressed while already in a not-yet-adopted new chat (activeChatId stays null)
// would otherwise leave the in-flight streaming thread mounted (#161). Dispatch
// `reconcile` to chatId:null with a brand-new key so React remounts ChatThread
// (a fresh useChat store). Disarm any armed fallback too. After this dispatch
// thread.chatId is null; the window also sets activeChatId to null, so the
// render-phase reconciler then finds them equal and does not double-remount.
const startFreshThread = useCallback(() => {
pendingNewChatRef.current = null;
dispatch({
type: "reconcile",
chatId: null,
newKey: `new-${generateId()}`,
});
}, []);
return {
threadKey: thread.key,
waitingForHistory,
onTurnFinished,
onServerChatId,
startFreshThread,
cancelPendingAdoption,
};
}

View File

@@ -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));
});
});
});

View File

@@ -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 ?? ""}`;
}