From f8a554bced3122c10b220c0f847fca968eb88cba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=92=8B=E6=96=B9=E6=98=8E?= Date: Fri, 12 Jun 2026 20:05:58 +0800 Subject: [PATCH] fix(desktop): keep message component types stable across Thread re-renders The component map Thread passes to the virtualizer listed the onBranchInNewChat / onCancel callbacks as useMemo deps. Whenever a parent re-render handed down a fresh callback identity, the memo rebuilt the map and produced new component *types*, so React unmounted and remounted every visible message. Async-rendered parts (shiki code blocks) collapsed and re-expanded on each remount, making the whole thread visibly jump. That is exactly what shipped in v0.15.1: the desktop controller passed an inline arrow for onBranchInNewChat, and the 15s status-snapshot poll re-rendered the controller, so threads with code blocks jumped every 15 seconds (layout-shift scores of 0.39 + 0.47 per cycle, measured via CDP). arrow away from regressing. Route the callbacks through a ref so the component types survive any parent re-render; only the callbacks' definedness stays a dep, because it gates UI (the user-message Stop button). Add a regression test that fails on the old code by asserting message DOM nodes keep their identity when callback props change identity. Tested on macOS arm64 (vitest + rebuilt app, CDP layout-shift instrumentation confirms zero shifts over multiple poll cycles). Co-Authored-By: Claude Fable 5 --- .../assistant-ui/thread-remount.test.tsx | 125 ++++++++++++++++++ .../components/assistant-ui/thread/index.tsx | 33 ++++- 2 files changed, 153 insertions(+), 5 deletions(-) create mode 100644 apps/desktop/src/components/assistant-ui/thread-remount.test.tsx diff --git a/apps/desktop/src/components/assistant-ui/thread-remount.test.tsx b/apps/desktop/src/components/assistant-ui/thread-remount.test.tsx new file mode 100644 index 000000000000..88454751e1b2 --- /dev/null +++ b/apps/desktop/src/components/assistant-ui/thread-remount.test.tsx @@ -0,0 +1,125 @@ +import { AssistantRuntimeProvider, type ThreadMessage, useExternalStoreRuntime } from '@assistant-ui/react' +import { act, render, screen, waitFor } from '@testing-library/react' +import { describe, expect, it, vi } from 'vitest' + +import { Thread } from './thread' + +class NoopResizeObserver { + observe() {} + + unobserve() {} + + disconnect() {} +} + +vi.stubGlobal('ResizeObserver', NoopResizeObserver) +vi.stubGlobal('requestAnimationFrame', (callback: FrameRequestCallback) => + window.setTimeout(() => callback(performance.now()), 0) +) +vi.stubGlobal('cancelAnimationFrame', (id: number) => window.clearTimeout(id)) + +Element.prototype.scrollTo = function scrollTo() {} + +Element.prototype.animate = function animate() { + return { + cancel: () => {}, + finished: Promise.resolve() + } as unknown as Animation +} + +// jsdom returns 0 for offset*; the virtualizer reads those to size its +// viewport. Fall through to client* or a sane default so virtualized +// items render (same stub as streaming.test.tsx). +function stubOffsetDimension( + prop: 'offsetHeight' | 'offsetWidth', + clientProp: 'clientHeight' | 'clientWidth', + fallback: number +) { + const previous = Object.getOwnPropertyDescriptor(HTMLElement.prototype, prop) + + Object.defineProperty(HTMLElement.prototype, prop, { + configurable: true, + get() { + return previous?.get?.call(this) || (this as HTMLElement)[clientProp] || fallback + } + }) +} + +stubOffsetDimension('offsetWidth', 'clientWidth', 800) +stubOffsetDimension('offsetHeight', 'clientHeight', 600) + +const createdAt = new Date('2026-05-01T00:00:00.000Z') + +const MESSAGES: ThreadMessage[] = [ + { + id: 'user-1', + role: 'user', + content: [{ type: 'text', text: 'hello from the user' }], + attachments: [], + createdAt, + metadata: { custom: {} } + } as ThreadMessage, + { + id: 'assistant-1', + role: 'assistant', + content: [{ type: 'text', text: 'stable assistant reply' }], + status: { type: 'complete', reason: 'stop' }, + createdAt, + metadata: { + unstable_state: null, + unstable_annotations: [], + unstable_data: [], + steps: [], + custom: {} + } + } as ThreadMessage +] + +function Harness({ + onBranchInNewChat, + onCancel +}: { + onBranchInNewChat: (messageId: string) => void + onCancel: () => void +}) { + const runtime = useExternalStoreRuntime({ + messages: MESSAGES, + isRunning: false, + onNew: async () => {} + }) + + return ( + + + + ) +} + +describe('thread message mount stability', () => { + // Regression: the desktop controller re-renders every 15s (status + // snapshot poll) and used to pass freshly-created callbacks down to + // . Those callbacks were deps of the `messageComponents` + // useMemo, so new component *types* were created each poll and React + // unmounted/remounted every visible message — shiki re-highlighted + // code blocks and the whole thread visibly jumped. + it('keeps message DOM nodes mounted when callback props get new identities', async () => { + const { rerender } = render( {}} onCancel={() => {}} />) + + await waitFor(() => { + expect(screen.getByText('stable assistant reply')).toBeTruthy() + expect(screen.getByText('hello from the user')).toBeTruthy() + }) + + const assistantBefore = screen.getByText('stable assistant reply') + const userBefore = screen.getByText('hello from the user') + + // Same data, new callback identities — exactly what a parent + // re-render driven by an unrelated state update produces. + await act(async () => { + rerender( {}} onCancel={() => {}} />) + }) + + expect(screen.getByText('stable assistant reply')).toBe(assistantBefore) + expect(screen.getByText('hello from the user')).toBe(userBefore) + }) +}) diff --git a/apps/desktop/src/components/assistant-ui/thread/index.tsx b/apps/desktop/src/components/assistant-ui/thread/index.tsx index 984b33e0e7e5..a5263e137101 100644 --- a/apps/desktop/src/components/assistant-ui/thread/index.tsx +++ b/apps/desktop/src/components/assistant-ui/thread/index.tsx @@ -1,4 +1,4 @@ -import { type FC, useCallback, useMemo, useState } from 'react' +import { type FC, useCallback, useMemo, useRef, useState } from 'react' import { AssistantMessage } from '@/components/assistant-ui/thread/assistant-message' import { ThreadMessageList } from '@/components/assistant-ui/thread/list' @@ -67,21 +67,44 @@ export const Thread: FC<{ setRestoreConfirmTarget({ messageId, ...target }) }, []) + // The values in this map are component *types*: when their identity + // changes, React unmounts and remounts every visible message — async + // re-rendered parts (shiki code blocks) collapse and re-expand, so the + // whole thread visibly jumps. Parents re-render on unrelated state + // (e.g. the 15s status-snapshot poll in the desktop controller) and + // can't be trusted to keep callback identities stable (see #38333), so + // route the callbacks through a ref instead of listing them as memo + // deps. Only their definedness stays a dep — it gates UI (the user + // Stop button, the restore-confirm affordance). Assigned during render + // (the useStoreSelector pattern) so the ref never lags a render. + const callbacksRef = useRef({ onBranchInNewChat, onCancel, onDismissError, onRestoreToMessage }) + callbacksRef.current = { onBranchInNewChat, onCancel, onDismissError, onRestoreToMessage } + + const hasBranchInNewChat = Boolean(onBranchInNewChat) + const hasCancel = Boolean(onCancel) + const hasDismissError = Boolean(onDismissError) + const hasRestoreToMessage = Boolean(onRestoreToMessage) + const messageComponents = useMemo( () => ({ AssistantMessage: () => ( - + callbacksRef.current.onBranchInNewChat?.(messageId) : undefined + } + onDismissError={hasDismissError ? messageId => callbacksRef.current.onDismissError?.(messageId) : undefined} + /> ), SystemMessage, UserEditComposer: () => , UserMessage: () => ( callbacksRef.current.onCancel?.() : undefined} + onRequestRestoreConfirm={hasRestoreToMessage ? requestRestoreConfirm : undefined} /> ) }), - [cwd, gateway, onBranchInNewChat, onCancel, onDismissError, onRestoreToMessage, requestRestoreConfirm, sessionId] + [cwd, gateway, hasBranchInNewChat, hasCancel, hasDismissError, hasRestoreToMessage, requestRestoreConfirm, sessionId] ) const emptyPlaceholder = intro ? (