diff --git a/apps/desktop/src/app/chat/composer/hooks/use-composer-undo.test.tsx b/apps/desktop/src/app/chat/composer/hooks/use-composer-undo.test.tsx new file mode 100644 index 00000000000..bbffc396ecb --- /dev/null +++ b/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, onSync: () => string) { + const api: { current: ReturnType | 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() + + 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() as RefObject + 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() as RefObject + 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() + }) +}) diff --git a/apps/desktop/src/components/assistant-ui/thread/user-edit-composer.tsx b/apps/desktop/src/components/assistant-ui/thread/user-edit-composer.tsx index 0da6e1dea8a..5168a0d8431 100644 --- a/apps/desktop/src/components/assistant-ui/thread/user-edit-composer.tsx +++ b/apps/desktop/src/components/assistant-ui/thread/user-edit-composer.tsx @@ -22,6 +22,7 @@ import { onComposerInsertRequest } from '@/app/chat/composer/focus' import { useAtCompletions } from '@/app/chat/composer/hooks/use-at-completions' +import { useComposerUndo } from '@/app/chat/composer/hooks/use-composer-undo' import { useSlashCompletions } from '@/app/chat/composer/hooks/use-slash-completions' import { dragHasAttachments, @@ -39,6 +40,7 @@ import { } from '@/app/chat/composer/rich-editor' import { detectTrigger, textBeforeCaret, type TriggerState } from '@/app/chat/composer/text-utils' import { ComposerTriggerPopover } from '@/app/chat/composer/trigger-popover' +import { isRedoShortcut, isUndoShortcut } from '@/app/chat/composer/undo-history' import { chipTypedUrlOnSpace, linkifyUrls } from '@/app/chat/composer/url-refs' import { extractDroppedFiles, @@ -215,6 +217,22 @@ export const UserEditComposer: FC = ({ cwd, gateway, sess [aui] ) + // Same stack the main composer owns, for the same reason: the editor mutates + // through `Range` to dodge Chromium's O(n²) editing pipeline, which also + // dodges its undo stack, so a paste was invisible to Cmd+Z. `rememberInitialDraft` + // already marks every mutation site (it's the dirty-edit guard), so the undo + // points ride along with it. + const syncFromEditorRef = useCallback(() => { + const editor = editorRef.current + + return editor ? syncDraftFromEditor(editor) : draftRef.current + }, [syncDraftFromEditor]) + + const { recordUndoPoint, redo, undo, withUndoPoint } = useComposerUndo({ + editorRef, + syncDraftFromEditor: syncFromEditorRef + }) + const refreshTrigger = useCallback(() => { const editor = editorRef.current @@ -278,6 +296,7 @@ export const UserEditComposer: FC = ({ cwd, gateway, sess } rememberInitialDraft() + recordUndoPoint() const serialized = hermesDirectiveFormatter.serialize(item) const starter = serialized.endsWith(':') const text = starter || serialized.endsWith(' ') ? serialized : `${serialized} ` @@ -327,7 +346,7 @@ export const UserEditComposer: FC = ({ cwd, gateway, sess document.execCommand('insertText', false, text) finish() }, - [aui, closeTrigger, refreshTrigger, rememberInitialDraft, requestEditFocus, trigger] + [aui, closeTrigger, recordUndoPoint, refreshTrigger, rememberInitialDraft, requestEditFocus, trigger] ) const insertRefStrings = useCallback( @@ -338,20 +357,23 @@ export const UserEditComposer: FC = ({ cwd, gateway, sess return false } - const nextDraft = insertInlineRefsIntoEditor(editor, refs) + // Bank BEFORE the insert — insertInlineRefsIntoEditor mutates in place, so + // recording after it would snapshot the state we're trying to undo to. + const undone = withUndoPoint(() => insertInlineRefsIntoEditor(editor, refs) !== null) - if (nextDraft === null) { + if (!undone) { return false } rememberInitialDraft() + const nextDraft = composerPlainText(editor) draftRef.current = nextDraft aui.composer().setText(nextDraft) requestEditFocus() return true }, - [aui, rememberInitialDraft, requestEditFocus] + [aui, rememberInitialDraft, requestEditFocus, withUndoPoint] ) const insertDroppedRefs = useCallback( @@ -494,6 +516,19 @@ export const UserEditComposer: FC = ({ cwd, gateway, sess window.setTimeout(refreshTrigger, 0) } + // Native typing/deleting still goes through Chromium's editing pipeline, whose + // undo stack we've taken over — bank the pre-edit state here, while + // `beforeinput` can still see the old text. + const handleBeforeInput = (event: FormEvent) => { + const inputType = (event.nativeEvent as InputEvent).inputType + + if (inputType === 'historyUndo' || inputType === 'historyRedo') { + return + } + + recordUndoPoint({ coalesce: inputType === 'insertText' || inputType === 'deleteContentBackward' }) + } + const handlePaste = (event: ClipboardEvent) => { const pastedText = sanitizeComposerInput(event.clipboardData.getData('text')) @@ -505,6 +540,7 @@ export const UserEditComposer: FC = ({ cwd, gateway, sess event.preventDefault() rememberInitialDraft() + recordUndoPoint() // Links land as `@url:` chips, same as the main composer. insertComposerContentsAtCaret(event.currentTarget, linkifyUrls(pastedText)) @@ -599,6 +635,22 @@ export const UserEditComposer: FC = ({ cwd, gateway, sess } } + // Undo/redo before Escape — we own the stack, and a stray Cmd+Z must never + // fall through to something that cancels the edit outright. + if (isUndoShortcut(event.nativeEvent)) { + event.preventDefault() + undo() + + return + } + + if (isRedoShortcut(event.nativeEvent)) { + event.preventDefault() + redo() + + return + } + if (event.key === 'Escape') { event.preventDefault() aui.composer().cancel() @@ -607,7 +659,7 @@ export const UserEditComposer: FC = ({ cwd, gateway, sess } // A typed link finished with a space chips like a pasted one. - if (chipTypedUrlOnSpace(event)) { + if (withUndoPoint(() => chipTypedUrlOnSpace(event))) { event.preventDefault() rememberInitialDraft() syncDraftFromEditor(event.currentTarget) @@ -682,6 +734,7 @@ export const UserEditComposer: FC = ({ cwd, gateway, sess contentEditable data-placeholder={copy.editMessage} data-slot={RICH_INPUT_SLOT} + onBeforeInput={handleBeforeInput} onBlur={() => window.setTimeout(closeTrigger, 80)} onDragOver={handleDragOver} onDrop={handleDrop}