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 <noreply@anthropic.com>

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<ThreadMessage>({
+    messages: MESSAGES,
+    isRunning: false,
+    onNew: async () => {}
+  })
+
+  return (
+    <AssistantRuntimeProvider runtime={runtime}>
+      <Thread onBranchInNewChat={onBranchInNewChat} onCancel={onCancel} />
+    </AssistantRuntimeProvider>
+  )
+}
+
+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
+  // <Thread/>. 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(<Harness onBranchInNewChat={() => {}} 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(<Harness onBranchInNewChat={() => {}} onCancel={() => {}} />)
+    })
+
+    expect(screen.getByText('stable assistant reply')).toBe(assistantBefore)
+    expect(screen.getByText('hello from the user')).toBe(userBefore)
+  })
+})

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: () => (
-        <AssistantMessage onBranchInNewChat={onBranchInNewChat} onDismissError={onDismissError} />
+        <AssistantMessage
+          onBranchInNewChat={
+            hasBranchInNewChat ? messageId => callbacksRef.current.onBranchInNewChat?.(messageId) : undefined
+          }
+          onDismissError={hasDismissError ? messageId => callbacksRef.current.onDismissError?.(messageId) : undefined}
+        />
       ),
       SystemMessage,
       UserEditComposer: () => <UserEditComposer cwd={cwd} gateway={gateway} sessionId={sessionId} />,
       UserMessage: () => (
         <UserMessage
-          onCancel={onCancel}
-          onRequestRestoreConfirm={onRestoreToMessage ? requestRestoreConfirm : undefined}
+          onCancel={hasCancel ? () => 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 ? (