diff --git a/apps/desktop/src/components/assistant-ui/message-render-boundary.test.tsx b/apps/desktop/src/components/assistant-ui/message-render-boundary.test.tsx index 8e1b70f7934f3..e4a890ef5addf 100644 --- a/apps/desktop/src/components/assistant-ui/message-render-boundary.test.tsx +++ b/apps/desktop/src/components/assistant-ui/message-render-boundary.test.tsx @@ -1,9 +1,14 @@ -import { cleanup, render, screen } from '@testing-library/react' +import { act, cleanup, render, screen } from '@testing-library/react' +import { Component, type ReactNode } from 'react' import { afterEach, describe, expect, it, vi } from 'vitest' import { MessageRenderBoundary } from './message-render-boundary' -afterEach(cleanup) +afterEach(() => { + cleanup() + vi.useRealTimers() + vi.restoreAllMocks() +}) function Boom({ error }: { error: Error | null }): null { if (error) { @@ -15,6 +20,26 @@ function Boom({ error }: { error: Error | null }): null { const lookupError = new Error('useClientLookup: Index 2 out of bounds (length: 2)') +const outerCaught: Error[] = [] + +// Records what propagates past MessageRenderBoundary, so the tests can tell +// a re-thrown error apart from a swallowed one. +class RecordingBoundary extends Component<{ children: ReactNode }, { error: Error | null }> { + state: { error: Error | null } = { error: null } + + static getDerivedStateFromError(error: Error) { + return { error } + } + + componentDidCatch(error: Error) { + outerCaught.push(error) + } + + render() { + return this.state.error ? null : this.props.children + } +} + describe('MessageRenderBoundary', () => { it('renders children when nothing throws', () => { render( @@ -77,4 +102,163 @@ describe('MessageRenderBoundary', () => { spy.mockRestore() }) + + it('recovers on the retry timer without a resetKey change', () => { + // The mid-turn race: the message list shrinks and regrows while + // ids/roles/count stay stable, so resetKey never changes. The boundary + // must self-retry on a timer instead of rendering null for the rest of + // the turn. + vi.useFakeTimers() + const spy = vi.spyOn(console, 'error').mockImplementation(() => undefined) + + let failing = true + + function MaybeBoom() { + if (failing) { + throw new Error('useClientLookup: index 3 out of bounds') + } + + return
turn content
+ } + + render( + + + + ) + + expect(screen.queryByText('turn content')).toBeNull() + + failing = false + + act(() => { + vi.advanceTimersByTime(0) + }) + + // Recovered through the retry timer alone; resetKey never changed. + expect(screen.getByText('turn content')).toBeTruthy() + spy.mockRestore() + }) + + it('stops retrying after the transient retry cap', () => { + // If the lookup stays out of bounds the boundary must give up instead of + // looping a setState/render cycle forever: initial render plus 5 retries, + // then it stays null and arms no further timer. + vi.useFakeTimers() + const spy = vi.spyOn(console, 'error').mockImplementation(() => undefined) + + let attempts = 0 + + function AlwaysBoom(): null { + attempts += 1 + throw lookupError + } + + const { container } = render( + + + + ) + + // React dev mode replays a failed render once per attempt, and an error + // during the initial mount gets an extra sync retry from the root, so + // measure the per-attempt cost from the first retry instead of guessing. + const mountAttempts = attempts + + act(() => { + vi.advanceTimersByTime(0) + }) + + const perRetry = attempts - mountAttempts + + for (let retry = 0; retry < 4; retry += 1) { + act(() => { + vi.advanceTimersByTime(0) + }) + } + + // Initial render plus 5 retries, then the boundary gives up: it stays + // null and arms no further timer. + expect(attempts).toBe(mountAttempts + perRetry * 5) + expect(vi.getTimerCount()).toBe(0) + expect(container.innerHTML).toBe('') + + act(() => { + vi.advanceTimersByTime(1000) + }) + + expect(attempts).toBe(mountAttempts + perRetry * 5) + spy.mockRestore() + }) + + it('resets the retry budget after a successful recovery', () => { + // The cap bounds a single streak of consecutive transient catches. A + // recovered boundary must get a fresh budget, otherwise enough separate + // races over a long session would permanently blank the turn. + vi.useFakeTimers() + const spy = vi.spyOn(console, 'error').mockImplementation(() => undefined) + + let failing = true + + function MaybeBoom() { + if (failing) { + throw lookupError + } + + return
turn content
+ } + + const { rerender } = render( + + + + ) + + // The mount is the first streak; five more follow. With a lifetime + // budget the sixth streak would find the cap exhausted and stay blank. + // Each rerender needs a fresh element: React bails out on an identical + // element reference and the child would never re-render (or re-throw). + for (let streak = 0; streak < 6; streak += 1) { + expect(screen.queryByText('turn content')).toBeNull() + + failing = false + + act(() => { + vi.advanceTimersByTime(0) + }) + + expect(screen.getByText('turn content')).toBeTruthy() + + failing = true + + rerender( + + + + ) + } + + spy.mockRestore() + }) + + it('does not schedule a retry for non-transient errors', () => { + vi.useFakeTimers() + const spy = vi.spyOn(console, 'error').mockImplementation(() => undefined) + + outerCaught.length = 0 + + render( + + + + + + ) + + // MessageRenderBoundary re-threw, the outer boundary caught it, and no + // retry timer was armed for a failure that cannot heal itself. + expect(outerCaught.map(error => error.message)).toContain('boom') + expect(vi.getTimerCount()).toBe(0) + spy.mockRestore() + }) }) diff --git a/apps/desktop/src/components/assistant-ui/message-render-boundary.tsx b/apps/desktop/src/components/assistant-ui/message-render-boundary.tsx index 4cae22b5c0c0a..fa19368e33163 100644 --- a/apps/desktop/src/components/assistant-ui/message-render-boundary.tsx +++ b/apps/desktop/src/components/assistant-ui/message-render-boundary.tsx @@ -12,6 +12,12 @@ import { Component, type ReactNode } from 'react' const isTransientLookupError = (error: unknown): boolean => error instanceof Error && /(useClientLookup|tapClient(Lookup|Resource)).*out of bounds/.test(error.message) +// Consecutive transient retries before giving up and waiting for a structural +// resetKey change (the pre-retry behavior). The race heals on the next +// consistent store snapshot, so one retry almost always recovers; the cap +// only bounds a pathological loop where the lookup stays out of bounds. +const MAX_TRANSIENT_RETRIES = 5 + interface Props { // Changes whenever the message list mutates STRUCTURALLY (ids/roles/count); // remounting clears the caught error so the next consistent render recovers @@ -27,13 +33,57 @@ interface Props { export class MessageRenderBoundary extends Component { state: { error: Error | null } = { error: null } + private retryTimer: number | null = null + + private transientRetries = 0 + static getDerivedStateFromError(error: Error) { return { error } } - componentDidUpdate(prev: Props) { + componentDidCatch(error: Error) { + // The resetKey path below only recovers on a STRUCTURAL change, but this + // race also fires mid-turn while ids/roles/count are stable: without a + // self-retry the boundary renders null for the rest of the turn (or + // until an unrelated message add/remove). Retry on a timer, not rAF — + // a parked renderer never fires frames, and a timer always runs. + if (!isTransientLookupError(error) || this.transientRetries >= MAX_TRANSIENT_RETRIES) { + return + } + + if (typeof window === 'undefined') { + return + } + + this.transientRetries += 1 + this.retryTimer = window.setTimeout(() => { + this.retryTimer = null + this.setState({ error: null }) + }, 0) + } + + componentDidUpdate(prev: Props, prevState: { error: Error | null }) { if (this.state.error && prev.resetKey !== this.props.resetKey) { this.setState({ error: null }) + + return + } + + if (prevState.error && !this.state.error) { + // Recovered (retry or structural reset): reset the retry budget and + // drop any retry timer the structural reset just made redundant. + this.transientRetries = 0 + + if (this.retryTimer !== null) { + window.clearTimeout(this.retryTimer) + this.retryTimer = null + } + } + } + + componentWillUnmount() { + if (this.retryTimer !== null) { + window.clearTimeout(this.retryTimer) } } diff --git a/apps/desktop/src/components/assistant-ui/thread/edit-context.test.tsx b/apps/desktop/src/components/assistant-ui/thread/edit-context.test.tsx new file mode 100644 index 0000000000000..f9c4f7129f174 --- /dev/null +++ b/apps/desktop/src/components/assistant-ui/thread/edit-context.test.tsx @@ -0,0 +1,191 @@ +// Thread deliberately keeps cwd/gateway/sessionId OUT of the messageComponents +// memo deps: those values change on every session switch, and reminting the +// component types mid-switch remounts the whole outgoing transcript. The +// mounted edit composer still has to see a same-session change (e.g. a cwd +// remap). It used to read the values from a render-time ref, but a mounted +// composer never re-reads the ref when the change leaves every +// ThreadMessageList prop referentially equal (Thread and ThreadMessageList +// are both memo'd, so the wrapper never re-renders). The values now travel +// through ThreadEditContext, whose propagation reaches the mounted consumer +// through the memo bail-out. These tests pin both directions: the composer +// sees the change, and the transcript still does not remount. +import { ExportedMessageRepository } from '@assistant-ui/react' +import { AssistantRuntimeProvider, type ThreadMessage } from '@assistant-ui/react' +import { act, cleanup, fireEvent, render, screen, waitFor } from '@testing-library/react' +import { useState } from 'react' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import { useIncrementalExternalStoreRuntime } from '@/lib/incremental-external-store-runtime' + +import { Thread } from '.' + +interface MockComposerProps { + cwd: string | null + gateway: unknown + sessionId: string | null +} + +const composerRenders = vi.hoisted(() => [] as MockComposerProps[]) + +vi.mock('./user-edit-composer', () => ({ + UserEditComposer: (props: MockComposerProps) => { + composerRenders.push(props) + + return
{props.cwd}
+ } +})) + +const createdAt = new Date('2026-05-01T00:00:00.000Z') + +class TestResizeObserver { + observe() {} + unobserve() {} + disconnect() {} +} + +vi.stubGlobal('ResizeObserver', TestResizeObserver) +vi.stubGlobal('requestAnimationFrame', (callback: FrameRequestCallback) => + window.setTimeout(() => callback(performance.now()), 0) +) +vi.stubGlobal('cancelAnimationFrame', (id: number) => window.clearTimeout(id)) +vi.stubGlobal('CSS', { escape: (str: string) => str }) + +Element.prototype.scrollTo = function scrollTo() {} + +afterEach(() => { + cleanup() +}) + +beforeEach(() => { + composerRenders.length = 0 +}) + +// 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 user-message-edit.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) + +function userMessage(): ThreadMessage { + return { + id: 'user-1', + role: 'user', + content: [{ type: 'text', text: 'edit me please' }], + attachments: [], + createdAt, + metadata: { custom: {} } + } as ThreadMessage +} + +function assistantMessage(): ThreadMessage { + return { + id: 'assistant-1', + role: 'assistant', + content: [{ type: 'text', text: 'done' }], + status: { type: 'complete', reason: 'stop' }, + createdAt, + metadata: { + unstable_state: null, + unstable_annotations: [], + unstable_data: [], + steps: [], + custom: {} + } + } as ThreadMessage +} + +const noopAsync = async () => {} + +// The repository must stay referentially stable across rerenders: a new +// object would make the incremental runtime resync the transcript and +// unmount the open composer, defeating the test. +function Harness({ cwd, sessionKey }: { cwd: string; sessionKey: string }) { + const [repository] = useState(() => ExportedMessageRepository.fromArray([userMessage(), assistantMessage()])) + + const runtime = useIncrementalExternalStoreRuntime({ + messageRepository: repository, + isRunning: false, + setMessages: () => {}, + onNew: noopAsync, + onEdit: noopAsync, + onCancel: noopAsync, + onReload: noopAsync + }) + + return ( + + + + ) +} + +describe('thread edit context', () => { + it('passes a same-session cwd change to the mounted edit composer', async () => { + const { rerender } = render() + + fireEvent.click(await screen.findByRole('button', { name: 'Edit message' })) + await screen.findByTestId('edit-composer') + + expect(composerRenders.at(-1)?.cwd).toBe('/old') + + // Same session, same messages: every ThreadMessageList prop stays + // referentially equal, so only context propagation can reach the + // mounted composer. + await act(async () => { + rerender() + }) + + expect(composerRenders.at(-1)?.cwd).toBe('/new') + expect(screen.getByTestId('edit-composer').textContent).toBe('/new') + }) + + it('still passes the new cwd after a session switch', async () => { + const { rerender } = render() + + await act(async () => { + rerender() + }) + + fireEvent.click(await screen.findByRole('button', { name: 'Edit message' })) + await screen.findByTestId('edit-composer') + + expect(composerRenders.at(-1)?.cwd).toBe('/new') + }) + + it('does not remount the transcript when cwd changes', async () => { + // The perf invariant behind keeping cwd out of the memo deps: a cwd + // change with identical messages must not remint the component types, + // so the mounted message DOM nodes survive the rerender. + const { rerender } = render() + + await waitFor(() => { + expect(screen.getByText('done')).toBeTruthy() + expect(screen.getByText('edit me please')).toBeTruthy() + }) + + const assistantBefore = screen.getByText('done') + const userBefore = screen.getByText('edit me please') + + await act(async () => { + rerender() + }) + + expect(screen.getByText('done')).toBe(assistantBefore) + expect(screen.getByText('edit me please')).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 8be74ba9e1074..e67683eb1ffc1 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 { memo, useCallback, useMemo, useRef, useState } from 'react' +import { createContext, memo, useCallback, useContext, useMemo, useRef, useState } from 'react' import { AssistantMessage } from '@/components/assistant-ui/thread/assistant-message' import { ThreadMessageList } from '@/components/assistant-ui/thread/list' @@ -16,6 +16,22 @@ import { notifyError } from '@/store/notifications' type ThreadLoadingState = 'response' | 'session' +interface ThreadEditContextValue { + cwd: string | null + gateway: HermesGateway | null + sessionId: string | null +} + +// Edit-composer context. The composer only exists while a message is being +// edited, and it mounts deep inside the memo'd ThreadMessageList, so the +// edit context can neither ride the component-map memo deps (that remints +// the component types on every session switch and remounts the outgoing +// transcript) nor sit in a render-time ref (a mounted composer never +// re-reads it when a same-session change leaves every list prop +// referentially equal). Context solves both: the component type stays +// stable, and a changed value propagates straight to the mounted consumer. +const ThreadEditContext = createContext({ cwd: null, gateway: null, sessionId: null }) + interface ThreadProps { clampToComposer?: boolean cwd?: string | null @@ -87,19 +103,18 @@ export const Thread = memo(function Thread({ // Stop button, the restore-confirm affordance). Assigned during render // (the useStoreSelector pattern) so the ref never lags a render. // - // cwd / gateway / sessionId ride the same ref for the same reason, and it - // is load-bearing on the hot path: all three change on EVERY session - // switch, so listing them as deps re-minted these types mid-switch and - // remounted the entire OUTGOING transcript — thousands of renders of a - // thread that was about to be replaced, all of it before the resume RPC - // had even been sent. They are read inside the edit composer (which only - // exists while a message is being edited), never during a plain render, - // so a ref read is always current by the time it matters. + // cwd / gateway / sessionId stay OUT of the memo deps for the same + // reason: all three change on EVERY session switch, so listing them + // re-minted these types mid-switch and remounted the entire OUTGOING + // transcript — thousands of renders of a thread that was about to be + // replaced, all of it before the resume RPC had even been sent. They + // reach the edit composer through ThreadEditContext instead (see above). const callbacksRef = useRef({ onBranchInNewChat, onCancel, onDismissError, onRestoreToMessage }) callbacksRef.current = { onBranchInNewChat, onCancel, onDismissError, onRestoreToMessage } - const editContextRef = useRef({ cwd, gateway, sessionId }) - editContextRef.current = { cwd, gateway, sessionId } + // Only changes identity when one of the three values does, so Thread + // re-renders for unrelated reasons never re-render the composer. + const editContext = useMemo(() => ({ cwd, gateway, sessionId }), [cwd, gateway, sessionId]) const hasBranchInNewChat = Boolean(onBranchInNewChat) const hasCancel = Boolean(onCancel) @@ -118,7 +133,7 @@ export const Thread = memo(function Thread({ ), SystemMessage, UserEditComposer: () => { - const { cwd: editCwd, gateway: editGateway, sessionId: editSessionId } = editContextRef.current + const { cwd: editCwd, gateway: editGateway, sessionId: editSessionId } = useContext(ThreadEditContext) return }, @@ -146,25 +161,27 @@ export const Thread = memo(function Thread({ const loadingIndicator = useMemo(() => , []) return ( -
- - {loading === 'session' && } - - -
+ +
+ + {loading === 'session' && } + + +
+
) })