From a62eaaf3162c637e40fa3879592b97ac1c24c046 Mon Sep 17 00:00:00 2001
From: Adolanium <94890352+Adolanium@users.noreply.github.com>
Date: Mon, 27 Jul 2026 22:25:30 +0300
Subject: [PATCH] fix(desktop): self-retry transient boundary errors, reactive
 edit composer context

Two correctness holes left by the session-switch perf work (#72504 / #72524):

1. MessageRenderBoundary only cleared a swallowed transient useClientLookup
   error when the structural resetKey changed. Mid-turn, ids/roles/count are
   stable, so a lookup race during a stream left the boundary rendering null
   for the rest of the turn. The boundary now self-retries on a 0ms timer
   (rAF never fires in a parked renderer), bounded to 5 consecutive
   transient catches with the budget reset on recovery; the structural
   resetKey path is unchanged, and non-transient errors still re-throw.

2. cwd / gateway / sessionId were removed from the messageComponents memo
   deps and read through a render-time ref so session switches stop
   reminting the component types. But a mounted UserEditComposer only
   reads that ref when it renders, and a same-session change (cwd remap,
   gateway reconnect) leaves every ThreadMessageList prop referentially
   equal, so the memo'd list bails out and the open composer keeps stale
   values: @-completions, slash completions, and OS-drop uploads target
   the old cwd / gateway / session. Thread now provides the three values
   through a memoized ThreadEditContext; context propagates through the
   bail-out, the component type identity is untouched, and the transcript
   never remounts.
---
 .../message-render-boundary.test.tsx          | 188 ++++++++++++++++-
 .../assistant-ui/message-render-boundary.tsx  |  52 ++++-
 .../assistant-ui/thread/edit-context.test.tsx | 191 ++++++++++++++++++
 .../components/assistant-ui/thread/index.tsx  |  81 +++++---
 4 files changed, 477 insertions(+), 35 deletions(-)
 create mode 100644 apps/desktop/src/components/assistant-ui/thread/edit-context.test.tsx

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 <div>turn content</div>
+    }
+
+    render(
+      <MessageRenderBoundary resetKey="0:m1:user">
+        <MaybeBoom />
+      </MessageRenderBoundary>
+    )
+
+    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(
+      <MessageRenderBoundary resetKey="a">
+        <AlwaysBoom />
+      </MessageRenderBoundary>
+    )
+
+    // 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 <div>turn content</div>
+    }
+
+    const { rerender } = render(
+      <MessageRenderBoundary resetKey="a">
+        <MaybeBoom />
+      </MessageRenderBoundary>
+    )
+
+    // 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(
+        <MessageRenderBoundary resetKey="a">
+          <MaybeBoom />
+        </MessageRenderBoundary>
+      )
+    }
+
+    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(
+      <RecordingBoundary>
+        <MessageRenderBoundary resetKey="a">
+          <Boom error={new Error('boom')} />
+        </MessageRenderBoundary>
+      </RecordingBoundary>
+    )
+
+    // 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<Props, { error: Error | null }> {
   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 <div data-testid="edit-composer">{props.cwd}</div>
+  }
+}))
+
+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<ThreadMessage>({
+    messageRepository: repository,
+    isRunning: false,
+    setMessages: () => {},
+    onNew: noopAsync,
+    onEdit: noopAsync,
+    onCancel: noopAsync,
+    onReload: noopAsync
+  })
+
+  return (
+    <AssistantRuntimeProvider runtime={runtime}>
+      <Thread cwd={cwd} sessionKey={sessionKey} />
+    </AssistantRuntimeProvider>
+  )
+}
+
+describe('thread edit context', () => {
+  it('passes a same-session cwd change to the mounted edit composer', async () => {
+    const { rerender } = render(<Harness cwd="/old" sessionKey="k1" />)
+
+    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(<Harness cwd="/new" sessionKey="k1" />)
+    })
+
+    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(<Harness cwd="/old" sessionKey="k1" />)
+
+    await act(async () => {
+      rerender(<Harness cwd="/new" sessionKey="k2" />)
+    })
+
+    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(<Harness cwd="/old" sessionKey="k1" />)
+
+    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(<Harness cwd="/new" sessionKey="k1" />)
+    })
+
+    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<ThreadEditContextValue>({ 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 <UserEditComposer cwd={editCwd} gateway={editGateway} sessionId={editSessionId} />
       },
@@ -146,25 +161,27 @@ export const Thread = memo(function Thread({
   const loadingIndicator = useMemo(() => <BackgroundResumeNotice />, [])
 
   return (
-    <div className="relative grid h-full min-h-0 max-w-full grid-rows-[minmax(0,1fr)] overflow-hidden bg-transparent contain-[layout_paint]">
-      <ThreadMessageList
-        clampToComposer={clampToComposer}
-        components={messageComponents}
-        emptyPlaceholder={emptyPlaceholder}
-        loadingIndicator={loadingIndicator}
-        sessionKey={sessionKey}
-      />
-      {loading === 'session' && <CenteredThreadSpinner />}
-      <ThreadTimeline />
-      <ConfirmDialog
-        confirmLabel={copy.restoreConfirm}
-        description={copy.restoreBody}
-        destructive
-        onClose={closeRestoreConfirm}
-        onConfirm={confirmRestore}
-        open={Boolean(restoreConfirmTarget)}
-        title={copy.restoreTitle}
-      />
-    </div>
+    <ThreadEditContext.Provider value={editContext}>
+      <div className="relative grid h-full min-h-0 max-w-full grid-rows-[minmax(0,1fr)] overflow-hidden bg-transparent contain-[layout_paint]">
+        <ThreadMessageList
+          clampToComposer={clampToComposer}
+          components={messageComponents}
+          emptyPlaceholder={emptyPlaceholder}
+          loadingIndicator={loadingIndicator}
+          sessionKey={sessionKey}
+        />
+        {loading === 'session' && <CenteredThreadSpinner />}
+        <ThreadTimeline />
+        <ConfirmDialog
+          confirmLabel={copy.restoreConfirm}
+          description={copy.restoreBody}
+          destructive
+          onClose={closeRestoreConfirm}
+          onConfirm={confirmRestore}
+          open={Boolean(restoreConfirmTarget)}
+          title={copy.restoreTitle}
+        />
+      </div>
+    </ThreadEditContext.Provider>
   )
 })