fix(desktop): keep message component types stable across Thread re-renders

[travelreader8-stack]

What

Thread builds the component map it hands to the virtualizer inside a useMemo whose deps included the onBranchInNewChat / onCancel callbacks. 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.

This PR routes the callbacks through a ref so the component types survive any parent re-render. Only the callbacks' definedness stays a memo dep, because it gates UI (the user-message Stop button renders only when onCancel is provided). A regression test asserts message DOM nodes keep their identity when callback props change identity — it fails on current main.

Why

This is the structural half of a bug that shipped in v0.15.1: the desktop controller passed an inline arrow for onBranchInNewChat, and the 15-second status-snapshot poll (use-status-snapshot.ts) re-rendered the controller — so any session containing code blocks visibly jumped every 15 seconds. Instrumenting the packaged app over CDP showed two layout shifts per poll cycle (scores 0.39 + 0.47, each preceded by a ~65 ms long task), and a MutationObserver confirmed the message subtrees were being removed and re-inserted with pre.shiki blocks re-highlighting from scratch.

#38333 already fixed that one call site by passing the useCallback'd function directly, but Thread itself was still one inline arrow away from regressing — nothing pinned the behavior. With this change, callback identity churn in any parent can no longer remount messages, and the new test pins it.

How to test

• cd apps/desktop && npx vitest run --environment jsdom src/components/assistant-ui/thread-remount.test.tsx
  • fails on main (message <p> nodes are replaced across a callback-identity-only re-render), passes with this patch
• npx vitest run --environment jsdom src/components/assistant-ui/streaming.test.tsx — 16/17 pass; the one failure (renders an incomplete streaming reasoning fenced code block as a code card) also fails on unpatched main in my environment, so it's unrelated
• Manual repro of the original symptom on v0.15.1: open a session whose history contains fenced code blocks and wait ~15 s — the thread jumps as code blocks collapse/re-expand. With the rebuilt patched app, a 35 s CDP layout-shift observation over multiple poll cycles records zero entries.

Platforms tested

macOS 15 (arm64), Electron 40 — vitest suite plus the rebuilt packaged app with CDP layout-shift instrumentation.

🤖 Generated with Claude Code

[tonydwb]

Code Review Summary

Verdict: Approved

Excellent regression fix with thorough test coverage. The desktop status-snapshot poll (every 15s) was creating new callback identities on every render, which caused React to unmount/remount message DOM nodes — making shiki code blocks visibly collapse and re-expand.

Looks Good

• Fix uses a callbacksRef to store callback references, only tracking boolean definedness in the useMemo deps
• useEffect updates the ref on each render so the latest callbacks are always available
• New test file explicitly tests the regression scenario: rerender with new callback identities but same data should preserve DOM node identity
• Comment explains the root cause (#38333) clearly
• The ref approach is the right pattern when parent re-renders are uncontrollable

---

Reviewed by Hermes Agent

[teknium1]

Thanks for the focused regression test and the clear explanation of why component-type identity matters here. The underlying issue still exists on current main, but this needs a small re-salvage after later Thread work.

Problems

• The PR edits apps/desktop/src/components/assistant-ui/thread.tsx, but commit 7ff6908a5 extracted the live component to apps/desktop/src/components/assistant-ui/thread/index.tsx. The current messageComponents memo is at thread/index.tsx:74-88, so this will not cherry-pick cleanly.
• Current main has two additional callback dependencies in that memo: onDismissError and onRestoreToMessage (thread/index.tsx:88). Ref-proxying only branch/cancel callbacks would still allow callback-identity churn to recreate message component types.

Suggested changes

• Port the fix and regression coverage to the co-located thread/ implementation/tests, and decouple the map from all current callback identities while retaining only definedness dependencies that gate UI.

Automated hermes-sweeper review.

[teknium1]

This memo moved to apps/desktop/src/components/assistant-ui/thread/index.tsx in current main (commit 7ff6908a5). Please port the fix there and include its newer onDismissError and onRestoreToMessage callback dependencies; otherwise either can still recreate the message component types.