Give the composer its own undo stack

apps/desktop/src/app/chat/composer/hooks/use-composer-draft.ts

@@ -287,7 +287,7 @@
       unsubscribe()
       window.clearTimeout(draftPersistTimerRef.current)
     }
   }, [composerRuntime, queueEditRef])
 
   const insertText = (text: string) => {
     const base = draftRef.current
@@ -391,7 +391,7 @@
       window.removeEventListener('pagehide', flushPendingDraftPersist)
       flushPendingDraftPersist()
     }
   }, [syncDraftFromEditor])
 
   return {
     activeQueueSessionKeyRef,
@@ -408,6 +408,7 @@
     requestMainFocus,
     sessionIdRef,
     setComposerText,
-    stashAt
+    stashAt,
+    syncDraftFromEditor
   }
 }

apps/desktop/src/app/chat/composer/hooks/use-composer-undo.test.tsx

@@ -0,0 +1,193 @@
+import { render } from '@testing-library/react'
+import { createRef, type RefObject } from 'react'
+import { describe, expect, it, vi } from 'vitest'
+
+import { useComposerUndo } from './use-composer-undo'
+
+/** Mount the hook against a real contentEditable, exposing its API. */
+function mountUndo(editorRef: RefObject<HTMLDivElement | null>, onSync: () => string) {
+  const api: { current: ReturnType<typeof useComposerUndo> | null } = { current: null }
+
+  const Harness = () => {
+    // Assigned during render on purpose: the tests drive the API imperatively
+    // right after mount, and this is a harness, not app state.
+    api.current = useComposerUndo({ editorRef, syncDraftFromEditor: onSync })
+
+    return null
+  }
+
+  const view = render(<Harness />)
+
+  return { api, view }
+}
+
+function makeEditor(text: string) {
+  const editor = document.createElement('div')
+  editor.contentEditable = 'true'
+  // jsdom only focuses a contentEditable div when it's explicitly focusable;
+  // the real editor is reachable via the composer's focus bus.
+  editor.tabIndex = 0
+  editor.append(document.createTextNode(text))
+  document.body.append(editor)
+
+  const ref = createRef<HTMLDivElement>() as RefObject<HTMLDivElement | null>
+  ref.current = editor
+
+  return { editor, ref }
+}
+
+const caretAtEnd = (editor: HTMLElement) => {
+  const range = document.createRange()
+  const selection = window.getSelection()!
+  range.selectNodeContents(editor)
+  range.collapse(false)
+  selection.removeAllRanges()
+  selection.addRange(range)
+}
+
+describe('useComposerUndo', () => {
+  it('restores the pre-edit text, which is what a paste destroyed', () => {
+    const { editor, ref } = makeEditor('before')
+    caretAtEnd(editor)
+
+    const { api, view } = mountUndo(ref, () => editor.textContent || '')
+
+    // Bank, then simulate the Range-based paste that Chromium never records.
+    api.current!.recordUndoPoint()
+    editor.append(document.createTextNode(' PASTED'))
+    expect(editor.textContent).toBe('before PASTED')
+
+    api.current!.undo()
+    expect(editor.textContent).toBe('before')
+
+    api.current!.redo()
+    expect(editor.textContent).toBe('before PASTED')
+
+    view.unmount()
+    editor.remove()
+  })
+
+  it('withUndoPoint banks only when the edit actually ran', () => {
+    const { editor, ref } = makeEditor('text')
+    caretAtEnd(editor)
+
+    const { api, view } = mountUndo(ref, () => editor.textContent || '')
+
+    // A guard that declines must not consume an undo slot.
+    expect(api.current!.withUndoPoint(() => false)).toBe(false)
+    expect(api.current!.undo()).toBe(false)
+
+    expect(
+      api.current!.withUndoPoint(() => {
+        editor.append(document.createTextNode('!'))
+
+        return true
+      })
+    ).toBe(true)
+
+    api.current!.undo()
+    expect(editor.textContent).toBe('text')
+
+    view.unmount()
+    editor.remove()
+  })
+
+  it('claims a native historyUndo aimed at the focused editor', () => {
+    const { editor, ref } = makeEditor('kept')
+    editor.focus()
+    caretAtEnd(editor)
+
+    const { api, view } = mountUndo(ref, () => editor.textContent || '')
+
+    api.current!.recordUndoPoint()
+    editor.append(document.createTextNode(' extra'))
+
+    // What Electron's Edit menu `{ role: 'undo' }` produces.
+    const event = new InputEvent('beforeinput', { bubbles: true, cancelable: true, inputType: 'historyUndo' })
+    editor.dispatchEvent(event)
+
+    expect(event.defaultPrevented).toBe(true)
+    expect(editor.textContent).toBe('kept')
+
+    view.unmount()
+    editor.remove()
+  })
+
+  it('ignores a historyUndo while another editor holds focus', () => {
+    const { editor, ref } = makeEditor('mine')
+    const { editor: other } = makeEditor('theirs')
+
+    other.focus()
+
+    const { api, view } = mountUndo(ref, () => editor.textContent || '')
+
+    api.current!.recordUndoPoint()
+    editor.append(document.createTextNode(' changed'))
+
+    const event = new InputEvent('beforeinput', { bubbles: true, cancelable: true, inputType: 'historyUndo' })
+    other.dispatchEvent(event)
+
+    // Not ours to claim — the other surface keeps its native behavior.
+    expect(event.defaultPrevented).toBe(false)
+    expect(editor.textContent).toBe('mine changed')
+
+    view.unmount()
+    editor.remove()
+    other.remove()
+  })
+
+  it('keeps two mounted composers independent', () => {
+    const { editor: main, ref: mainRef } = makeEditor('main')
+    const { editor: edit, ref: editRef } = makeEditor('edit')
+
+    const mainUndo = mountUndo(mainRef, () => main.textContent || '')
+    const editUndo = mountUndo(editRef, () => edit.textContent || '')
+
+    mainUndo.api.current!.recordUndoPoint()
+    main.append(document.createTextNode(' typed'))
+
+    // Undoing in the edit composer must not touch the main composer's text.
+    editUndo.api.current!.undo()
+    expect(main.textContent).toBe('main typed')
+
+    mainUndo.api.current!.undo()
+    expect(main.textContent).toBe('main')
+    expect(edit.textContent).toBe('edit')
+
+    mainUndo.view.unmount()
+    editUndo.view.unmount()
+    main.remove()
+    edit.remove()
+  })
+
+  it('reset drops history so undo cannot cross a draft swap', () => {
+    const { editor, ref } = makeEditor('session A')
+    caretAtEnd(editor)
+
+    const { api, view } = mountUndo(ref, () => editor.textContent || '')
+
+    api.current!.recordUndoPoint()
+    editor.append(document.createTextNode(' edited'))
+    api.current!.resetUndoHistory()
+
+    expect(api.current!.undo()).toBe(false)
+    expect(editor.textContent).toBe('session A edited')
+
+    view.unmount()
+    editor.remove()
+  })
+
+  it('is inert when the editor ref is empty', () => {
+    const ref = createRef<HTMLDivElement>() as RefObject<HTMLDivElement | null>
+    const sync = vi.fn(() => '')
+
+    const { api, view } = mountUndo(ref, sync)
+
+    api.current!.recordUndoPoint()
+
+    expect(api.current!.undo()).toBe(false)
+    expect(sync).not.toHaveBeenCalled()
+
+    view.unmount()
+  })
+})

apps/desktop/src/app/chat/composer/hooks/use-composer-undo.ts

@@ -0,0 +1,124 @@
+import { type RefObject, useCallback, useEffect, useMemo } from 'react'
+
+import {
+  caretOffsetInEditor,
+  composerPlainText,
+  placeCaretAtOffset,
+  renderComposerContents
+} from '../rich-editor'
+import { type ComposerSnapshot, createComposerUndoHistory } from '../undo-history'
+
+interface UseComposerUndoArgs {
+  editorRef: RefObject<HTMLDivElement | null>
+  /** Push a restored snapshot back into draftRef + composer state. */
+  syncDraftFromEditor: () => string
+}
+
+/**
+ * Undo/redo for the rich composer.
+ *
+ * The editor mutates its DOM through `Range` to dodge Chromium's O(n²) editing
+ * pipeline (#45812), which also dodges Chromium's undo stack — so a paste was
+ * invisible to ⌘Z and the keystroke undid whatever edit came before it instead.
+ * We own the stack outright rather than half of it: every edit path records the
+ * pre-edit state here, and the editor claims ⌘Z / ⌘⇧Z itself.
+ */
+export function useComposerUndo({ editorRef, syncDraftFromEditor }: UseComposerUndoArgs) {
+  const history = useMemo(() => createComposerUndoHistory(), [])
+
+  const snapshot = useCallback((): ComposerSnapshot => {
+    const editor = editorRef.current
+
+    if (!editor) {
+      return { caret: 0, text: '' }
+    }
+
+    return { caret: caretOffsetInEditor(editor), text: composerPlainText(editor) }
+  }, [editorRef])
+
+  /** Bank the current state before mutating the editor. `coalesce` marks a
+   *  keystroke, so a run of typing collapses into one undo step. */
+  const recordUndoPoint = useCallback(
+    (options?: { coalesce?: boolean }) => {
+      if (editorRef.current) {
+        history.record(snapshot(), options)
+      }
+    },
+    [editorRef, history, snapshot]
+  )
+
+  const applySnapshot = useCallback(
+    (next: ComposerSnapshot | null) => {
+      const editor = editorRef.current
+
+      if (!next || !editor) {
+        return false
+      }
+
+      renderComposerContents(editor, next.text)
+      placeCaretAtOffset(editor, next.caret)
+      syncDraftFromEditor()
+
+      return true
+    },
+    [editorRef, syncDraftFromEditor]
+  )
+
+  /** Run a conditional edit, banking its pre-edit state only if it actually
+   *  ran. The snapshot has to be taken first (the edit destroys the state we'd
+   *  be saving), but recording unconditionally would clear the redo stack on
+   *  every Backspace that falls through to the native path. */
+  const withUndoPoint = useCallback(
+    (edit: () => boolean) => {
+      const before = snapshot()
+      const ran = edit()
+
+      if (ran) {
+        history.record(before)
+      }
+
+      return ran
+    },
+    [history, snapshot]
+  )
+
+  const undo = useCallback(() => applySnapshot(history.undo(snapshot())), [applySnapshot, history, snapshot])
+  const redo = useCallback(() => applySnapshot(history.redo(snapshot())), [applySnapshot, history, snapshot])
+
+  // A session/draft swap makes prior history meaningless — undoing into another
+  // conversation's text is worse than having no history at all.
+  const resetUndoHistory = useCallback(() => history.reset(), [history])
+
+  // Electron's Edit menu ships `{ role: 'undo' }`, whose accelerator the macOS
+  // menu bar consumes before the web contents sees the keystroke (the same
+  // hazard main.ts documents for ⌘W). It fires the native editing command,
+  // which knows nothing about our stack. Claim it at the document level while
+  // the composer holds focus, so the menu item and the keystroke agree.
+  useEffect(() => {
+    const onBeforeInput = (event: Event) => {
+      const inputType = (event as InputEvent).inputType
+
+      if (inputType !== 'historyUndo' && inputType !== 'historyRedo') {
+        return
+      }
+
+      if (document.activeElement !== editorRef.current) {
+        return
+      }
+
+      event.preventDefault()
+
+      if (inputType === 'historyUndo') {
+        undo()
+      } else {
+        redo()
+      }
+    }
+
+    document.addEventListener('beforeinput', onBeforeInput, true)
+
+    return () => document.removeEventListener('beforeinput', onBeforeInput, true)
+  }, [editorRef, redo, undo])
+
+  return { recordUndoPoint, redo, resetUndoHistory, undo, withUndoPoint }
+}