fix(editor): keep retained review highlight ranges mapped after edits (#34963)

* fix(editor): keep retained review highlight ranges mapped after edits

* test(editor): assert rendered review highlights instead of plugin state

---------

Co-authored-by: OrcaWin <293788423+OrcaWin@users.noreply.github.com>

src/renderer/src/components/editor/rich-markdown-annotation-highlight.mapping.test.ts

@@ -0,0 +1,97 @@
+// @vitest-environment happy-dom
+import { Editor } from '@tiptap/core'
+import StarterKit from '@tiptap/starter-kit'
+import { afterEach, describe, expect, it } from 'vitest'
+import {
+  createRichMarkdownAnnotationHighlightExtension,
+  richMarkdownAnnotationHighlightPluginKey
+} from './rich-markdown-annotation-highlight'
+
+const editors: Editor[] = []
+
+afterEach(() => {
+  for (const editor of editors.splice(0)) {
+    editor.destroy()
+  }
+})
+
+function createEditor() {
+  const element = document.createElement('div')
+  const editor = new Editor({
+    element,
+    extensions: [StarterKit, createRichMarkdownAnnotationHighlightExtension()],
+    content: '<p>before note after</p>'
+  })
+  editors.push(editor)
+  editor.view.dispatch(
+    editor.state.tr.setMeta(richMarkdownAnnotationHighlightPluginKey, {
+      activeRange: { from: 8, to: 12 },
+      noteRanges: [{ from: 8, to: 12 }]
+    })
+  )
+  return editor
+}
+
+function highlightedText(editor: Editor): string[] {
+  return [...editor.view.dom.querySelectorAll('.rich-markdown-annotation-selection')].map(
+    (element) => element.textContent ?? ''
+  )
+}
+
+describe('review highlight positions across edits and metadata', () => {
+  it('keeps saved note highlights on the same text when clearing a pulse after an edit', () => {
+    const editor = createEditor()
+    editor.view.dispatch(editor.state.tr.insertText('prefix ', 1))
+    expect(highlightedText(editor)).toEqual(['note'])
+    editor.view.dispatch(editor.state.tr.setMeta(richMarkdownAnnotationHighlightPluginKey, null))
+    expect(highlightedText(editor)).toEqual(['note'])
+  })
+
+  it('preserves mapped active and note ranges when metadata updates only one field', () => {
+    const editor = createEditor()
+    editor.view.dispatch(editor.state.tr.insertText('prefix ', 1))
+    editor.view.dispatch(
+      editor.state.tr.setMeta(richMarkdownAnnotationHighlightPluginKey, {
+        noteRanges: [{ from: 15, to: 19 }]
+      })
+    )
+    const active = editor.view.dom.querySelector('.rich-markdown-annotation-selection-active')
+    expect(active?.textContent).toBe('note')
+    expect(highlightedText(editor)).toEqual(['note'])
+  })
+
+  it('drops deleted highlights instead of recreating them over surviving text', () => {
+    const editor = createEditor()
+    editor.view.dispatch(editor.state.tr.delete(8, 12))
+    expect(highlightedText(editor)).toEqual([])
+    editor.view.dispatch(
+      editor.state.tr.setMeta(richMarkdownAnnotationHighlightPluginKey, {
+        activeRange: null
+      })
+    )
+    expect(highlightedText(editor)).toEqual([])
+  })
+
+  it('maps retained notes while accepting new-document metadata in the same edit transaction', () => {
+    const editor = createEditor()
+    editor.view.dispatch(
+      editor.state.tr.insertText('prefix ', 1).setMeta(richMarkdownAnnotationHighlightPluginKey, {
+        activeRange: null
+      })
+    )
+    expect(highlightedText(editor)).toEqual(['note'])
+  })
+
+  it('excludes boundary insertions and keeps undo/redo positions after rebuilding', () => {
+    const editor = createEditor()
+    editor.view.dispatch(editor.state.tr.insertText('end', 12).insertText('start', 8))
+    editor.view.dispatch(editor.state.tr.setMeta(richMarkdownAnnotationHighlightPluginKey, null))
+    expect(highlightedText(editor)).toEqual(['note'])
+    expect(editor.commands.undo()).toBe(true)
+    editor.view.dispatch(editor.state.tr.setMeta(richMarkdownAnnotationHighlightPluginKey, null))
+    expect(highlightedText(editor)).toEqual(['note'])
+    expect(editor.commands.redo()).toBe(true)
+    editor.view.dispatch(editor.state.tr.setMeta(richMarkdownAnnotationHighlightPluginKey, null))
+    expect(highlightedText(editor)).toEqual(['note'])
+  })
+})

src/renderer/src/components/editor/rich-markdown-annotation-highlight.ts

@@ -1,6 +1,6 @@
 import { Extension } from '@tiptap/core'
 import type { Node as ProseMirrorNode } from '@tiptap/pm/model'
-import { Plugin, PluginKey } from '@tiptap/pm/state'
+import { Plugin, PluginKey, type Transaction } from '@tiptap/pm/state'
 import { Decoration, DecorationSet } from '@tiptap/pm/view'
 
 type RichMarkdownAnnotationHighlightState = {
@@ -46,6 +46,18 @@ function createAnnotationDecorations(
   return decorations.length === 0 ? DecorationSet.empty : DecorationSet.create(doc, decorations)
 }
 
+function mapAnnotationRange(
+  range: RichMarkdownAnnotationHighlightRange | null,
+  transaction: Transaction
+): RichMarkdownAnnotationHighlightRange | null {
+  if (!range) {
+    return null
+  }
+  const from = transaction.mapping.map(Math.min(range.from, range.to), 1)
+  const to = transaction.mapping.map(Math.max(range.from, range.to), -1)
+  return from < to ? { from, to } : null
+}
+
 function createRichMarkdownAnnotationHighlightPlugin(): Plugin<RichMarkdownAnnotationHighlightState> {
   return new Plugin<RichMarkdownAnnotationHighlightState>({
     key: richMarkdownAnnotationHighlightPluginKey,
@@ -59,18 +71,26 @@ function createRichMarkdownAnnotationHighlightPlugin(): Plugin<RichMarkdownAnnot
         const meta = tr.getMeta(richMarkdownAnnotationHighlightPluginKey) as
           | RichMarkdownAnnotationHighlightMeta
           | undefined
+        const previousActiveRange = tr.docChanged
+          ? mapAnnotationRange(pluginState.activeRange, tr)
+          : pluginState.activeRange
+        const previousNoteRanges = tr.docChanged
+          ? pluginState.noteRanges.flatMap((range) => {
+              const mapped = mapAnnotationRange(range, tr)
+              return mapped ? [mapped] : []
+            })
+          : pluginState.noteRanges
         if (meta === null) {
           return {
             activeRange: null,
-            noteRanges: pluginState.noteRanges,
-            decorations: createAnnotationDecorations(tr.doc, null, pluginState.noteRanges)
+            noteRanges: previousNoteRanges,
+            decorations: createAnnotationDecorations(tr.doc, null, previousNoteRanges)
           }
         }
         if (meta) {
           const activeRange =
-            meta.activeRange === undefined ? pluginState.activeRange : meta.activeRange
-          const noteRanges =
-            meta.noteRanges === undefined ? pluginState.noteRanges : meta.noteRanges
+            meta.activeRange === undefined ? previousActiveRange : meta.activeRange
+          const noteRanges = meta.noteRanges === undefined ? previousNoteRanges : meta.noteRanges
           return {
             activeRange,
             noteRanges,
@@ -79,7 +99,8 @@ function createRichMarkdownAnnotationHighlightPlugin(): Plugin<RichMarkdownAnnot
         }
         if (tr.docChanged) {
           return {
-            ...pluginState,
+            activeRange: previousActiveRange,
+            noteRanges: previousNoteRanges,
             decorations: pluginState.decorations.map(tr.mapping, tr.doc)
           }
         }