Decoration loss on undo with async paragraph-level updates

Demo Setup (Minimal Plugin)

import { Extension } from '@tiptap/core'

import { Plugin, PluginKey } from 'prosemirror-state'

import { Decoration, DecorationSet, EditorView } from 'prosemirror-view'

import { Node as ProseMirrorNode } from 'prosemirror-model'

import { postRequest } from '../util/service'




// ─── Constants ────────────────────────────────────────────────────────────────

const DEBOUNCE_DELAY = 250

const REMOVE_DECORATION_META_KEY = 'removeDecoration'

const statcheckKey = new PluginKey('statcheck')




// ─── Module-level state ───────────────────────────────────────────────────────

let latestView: EditorView | null = null

let statcheckEnabled: boolean = false




// ─── Helpers ──────────────────────────────────────────────────────────────────




function debounce(func: Function, wait: number) {

let timeout: ReturnType<typeof setTimeout>

return function () {

clearTimeout(timeout)

timeout = setTimeout(func, wait)

  }

}




/**

 * Build visible text string + char-index → ProseMirror position map

 * for a single paragraph node.

 */

function getVisibleTextAndMap(

paragraphNode: ProseMirrorNode,

paragraphNodeStart: number

): { visibleText: string; mapVisToOrig: number[] } {

let visibleText = ''

const mapVisToOrig: number[] = []




paragraphNode.descendants((node, pos) => {

if (!node.isText) return true

const text = node.text || ''

for (let i = 0; i < text.length; i++) {

visibleText += text[i]

mapVisToOrig.push(paragraphNodeStart + pos + i)

    }

return true

  })




return { visibleText, mapVisToOrig }

}




// ─── API call ─────────────────────────────────────────────────────────────────




async function fetchStatcheckResults(text: string): Promise<any[]> {

if (!text.trim()) return []

try {

const res = await postRequest('/get-arxiv-data', {

text: JSON.stringify(text),

    })

return res?.data ?? []

  } catch {

return []

  }

}




// ─── Build decorations from API results for a paragraph ───────────────────────




function buildDecorations(

results: any[],

visibleText: string,

mapVisToOrig: number[],

paragraphId: string   // <-- added parameter

): Decoration[] {

const decorations: Decoration[] = []




for (const result of results) {

if (!result.error && !result.decision_error) continue




const raw: string = result.raw

if (!raw) continue




let searchFrom = 0

while (true) {

const idx = visibleText.indexOf(raw, searchFrom)

if (idx === -1) break




const visEnd = idx + raw.length

if (idx >= mapVisToOrig.length || visEnd > mapVisToOrig.length) break




const from = mapVisToOrig[idx]

const to = mapVisToOrig[visEnd - 1] + 1




// decision_error (significance flips) → red; plain error → orange

const cssClass =  'stat-warning'




decorations.push(

Decoration.inline(

from,

to,

          {

class: cssClass,

'data-raw': raw,

'data-reported-p': String(result.reported_p),

'data-computed-p': String(result.computed_p),

'data-test-type': result.test_type ?? '',

'data-decision-error': String(result.decision_error),

'data-start-pos': String(from),

'data-end-pos': String(to),

'data-paragraph-id': paragraphId,   // <-- store paragraph ID

          },

          { plugin: 'statcheck' }

        )

      )




searchFrom = idx + 1

    }

  }




return decorations

}




// ─── Check a single paragraph and dispatch its decorations ────────────────────




async function checkAndDecorateParagraph(

visibleText: string,

paragraphId: string,

mapVisToOrig: number[],

view: EditorView

) {

if (!statcheckEnabled) return




const results = await fetchStatcheckResults(visibleText)




if (view.isDestroyed) return




const decorations = buildDecorations(results, visibleText, mapVisToOrig, paragraphId)




view.dispatch(

view.state.tr.setMeta(statcheckKey, {

updateParagraphDecorations: { id: paragraphId, decorations },

    })

  )

}




// ─── Check entire document — single API call ─────────────────────────────────

async function checkEntireDocument(view: EditorView) {

// Step 1: Collect paragraphs

const paragraphs: Array<{

paragraphId: string

visibleText: string

mapVisToOrig: number[]

startInFull: number

endInFull: number

  }> = []




const SEPARATOR = '\n\n'

let cursor = 0




view.state.doc.descendants((node, pos) => {

if (!node.isTextblock) return true

if (!node.textContent.trim()) return true




const { visibleText, mapVisToOrig } = getVisibleTextAndMap(node, pos + 1)

const startInFull = cursor

const endInFull = cursor + visibleText.length




paragraphs.push({

paragraphId: `p-${pos}`,

visibleText,

mapVisToOrig,

startInFull,

endInFull,

    })




cursor = endInFull + SEPARATOR.length

return true

  })




if (paragraphs.length === 0) return




const fullText = paragraphs.map((p) => p.visibleText).join(SEPARATOR)

const results = await fetchStatcheckResults(fullText)




if (view.isDestroyed || !results.length) return




const paraDecoMap = new Map<string, Decoration[]>()




for (const result of results) {

if (!result.error && !result.decision_error) continue




const raw: string = result.raw

if (!raw) continue




let searchFrom = 0

while (true) {

const idxInFull = fullText.indexOf(raw, searchFrom)

if (idxInFull === -1) break




const endInFull = idxInFull + raw.length

const para = paragraphs.find(

        (p) => idxInFull >= p.startInFull && endInFull <= p.endInFull

      )




if (para) {

const localStart = idxInFull - para.startInFull

const localEnd = endInFull - para.startInFull




if (localStart < para.mapVisToOrig.length && localEnd <= para.mapVisToOrig.length) {

const from = para.mapVisToOrig[localStart]

const to = localEnd < para.mapVisToOrig.length

? para.mapVisToOrig[localEnd]

: para.mapVisToOrig[localEnd - 1] + 1




const cssClass = 'stat-warning'




const deco = Decoration.inline(

from,

to,

            {

class: cssClass,

'data-raw': raw,

'data-reported-p': String(result.reported_p),

'data-computed-p': String(result.computed_p),

'data-test-type': result.test_type ?? '',

'data-decision-error': String(result.decision_error),

'data-start-pos': String(from),

'data-end-pos': String(to),

'data-paragraph-id': para.paragraphId,   // <-- store paragraph ID

            },

            { plugin: 'statcheck' }

          )




if (!paraDecoMap.has(para.paragraphId)) paraDecoMap.set(para.paragraphId, [])

paraDecoMap.get(para.paragraphId)!.push(deco)

        }

      }




searchFrom = idxInFull + 1

    }

  }




// Dispatch one transaction per paragraph

for (const para of paragraphs) {

const decorations = paraDecoMap.get(para.paragraphId) ?? []

view.dispatch(

view.state.tr.setMeta(statcheckKey, {

updateParagraphDecorations: { id: para.paragraphId, decorations },

      })

    )

  }

}




// ─── ProseMirror plugin ───────────────────────────────────────────────────────




const createStatcheckPlugin = (editor: any) => {

let debounceTimeoutId: any = null;

let paragraphDecorationMap = new Map();




return new Plugin({

key: statcheckKey,




state: {

init() {

return DecorationSet.empty

      },




apply(tr, value, oldState, newState) {

// 1. Handle special meta that replaces the whole decoration set (clear all)

const meta = tr.getMeta(statcheckKey)




// 2. Always map existing decorations through the transaction

let newPluginStateDecorations = value.map(tr.mapping, tr.doc)




paragraphDecorationMap = new Map();




newPluginStateDecorations?.find(undefined, undefined, spec => spec.plugin === 'statcheck').forEach(deco => {




if (deco.from > newState.doc.content.size) return;




const resolvedPos = newState.doc.resolve(deco.from);

const paragraphNode = resolvedPos.parent;

const paragraphNodePos = resolvedPos.start(resolvedPos.depth);




if (paragraphNode && paragraphNode.isTextblock) {

const paragraphId = `p-${paragraphNodePos}`;

if (!paragraphDecorationMap.has(paragraphId)) {

paragraphDecorationMap.set(paragraphId, []);

            }

paragraphDecorationMap.get(paragraphId).push(deco);

          }

        });




// 3. Handle updates from async checkAndDecorateParagraph (via meta)

if (meta && meta.updateParagraphDecorations) {

const { id: updatedParagraphId, decorations: newDecorationsForParagraph } = meta.updateParagraphDecorations;




// Get the decorations currently associated with this paragraph ID from our map

const oldDecorationsForParagraph = paragraphDecorationMap.get(updatedParagraphId);

// Remove old decorations for this paragraph from the current DecorationSet

if (oldDecorationsForParagraph && oldDecorationsForParagraph.length > 0) {

newPluginStateDecorations = newPluginStateDecorations.remove(oldDecorationsForParagraph);

          }




// Add new decorations for this paragraph to the current DecorationSet

if (newDecorationsForParagraph && newDecorationsForParagraph.length > 0) {

newPluginStateDecorations = newPluginStateDecorations.add(tr.doc, newDecorationsForParagraph);

// Update the map with the NEW decorations for this paragraph

paragraphDecorationMap.set(updatedParagraphId, newDecorationsForParagraph);

          } else {

// If no new decorations, and there were old ones, ensure they are removed from the map

// This handles the case where all errors in a paragraph are fixed.

paragraphDecorationMap.delete(updatedParagraphId);

          }

        } else if (meta?.type == REMOVE_DECORATION_META_KEY) { // remove decoration area

const meta = tr.getMeta(statcheckKey);

if (meta?.decorations) {

return meta.decorations; // replace entire DecorationSet

          }

return oldState.map(tr.mapping, tr.doc); 

        }




// 4. Schedule a new check for the active paragraph only if the document changed

if (!oldState.doc.eq(newState.doc)) {

const selection = newState.selection

let activeParagraphNode: any = null

let activeParagraphNodePos = -1




if (selection.empty) {

const resolvedPos = selection.$head

activeParagraphNode = resolvedPos.parent

activeParagraphNodePos = resolvedPos.start(resolvedPos.depth)

          } else {

const resolvedStart = selection.$from

activeParagraphNode = resolvedStart.parent

activeParagraphNodePos = resolvedStart.start(resolvedStart.depth)

          }




if (activeParagraphNode && activeParagraphNode.isTextblock) {

const domNode = editor.view.nodeDOM(activeParagraphNodePos - 1)

if (domNode instanceof Element && domNode.closest('ref')) {

clearTimeout(debounceTimeoutId)

return newPluginStateDecorations

            }

const currentText = activeParagraphNode.textContent

const paragraphId = `p-${activeParagraphNodePos}`

const { visibleText, mapVisToOrig } = getVisibleTextAndMap(activeParagraphNode, activeParagraphNodePos)




const oldMappedPos = tr.mapping.map(activeParagraphNodePos)

const oldParagraphNode = (oldMappedPos >= 0 && oldMappedPos < oldState.doc.content.size)

? oldState.doc.resolve(oldMappedPos)

: null

const oldText = oldParagraphNode ? oldParagraphNode.parent?.textContent : ''




if (oldText !== currentText) {

clearTimeout(debounceTimeoutId)

debounceTimeoutId = setTimeout(() => {

checkAndDecorateParagraph(visibleText, paragraphId, mapVisToOrig, editor.view)

              }, DEBOUNCE_DELAY)

            }

          } else {

clearTimeout(debounceTimeoutId)

          }

        }




return newPluginStateDecorations

      },

    },




view(view) {

latestView = view

return {}

    },




props: {

decorations(state) {

return this.getState(state)

      },

    },

  })

}




// ─── Tiptap Extension ─────────────────────────────────────────────────────────




export const StatcheckExtension = Extension.create({

name: 'StatcheckExtension',




addOptions() {

return { enabled: true }

  },




addStorage() {

return { enabled: this.options.enabled }

  },




addProseMirrorPlugins() {

return [createStatcheckPlugin(this.editor)]

  },




addCommands():any {

return {

toggleStatcheck:

        () =>

        ({ editor }: any) => {

if (!latestView) {

console.warn('Editor view not available yet.')

return false

          }




if (!statcheckEnabled) {

statcheckEnabled = true

checkEntireDocument(latestView)

          } else {

editor.commands.clearStatcheckHighlights()

          }




return true

        },




clearStatcheckHighlights:

        () =>

        () => {

statcheckEnabled = false




const tr = latestView?.state.tr.setMeta(statcheckKey, {

decorations: DecorationSet.empty,

type: REMOVE_DECORATION_META_KEY,

          })

if (tr) latestView?.dispatch(tr)




return true

        },

    }

  },




onDestroy() {

statcheckEnabled = false

latestView = null

  },

});

Steps to Reproduce

  1. Initialize editor with multiple paragraphs

  2. Apply decorations to entire document (simulate initial run)

  3. Edit one paragraph → async update triggers

  4. Press undo (Ctrl+Z)

Observed Behavior

  • After undo:

    • Most decorations disappear

    • Only the active paragraph gets decorations again

Expected Behavior

  • All decorations should persist via mapping

  • Only modified paragraph should update

Notes

  • Decorations are inline

  • Updates are async (simulated with setTimeout)

  • Paragraphs identified using position-based IDs

Question

What is the recommended way to:

  • Maintain decorations across undo/redo

  • Safely update decorations for only part of the document

  • Avoid losing decorations when async updates are involved

I’m not going to sift through 500+ lines of reproduction code, especially when it includes a TipTap dependency. If you can boil this down to a small, ProseMirror-only setup, I can take a look.

Got it — that makes sense. I’ll reduce this to a minimal ProseMirror-only reproduction without the TipTap layer. I will provide the plugin code below. Thanks for taking a look.

const createStatcheckPlugin = (editor: any) => {

  let debounceTimeoutId: any = null;

let paragraphDecorationMap = new Map();




return new Plugin({

key: statcheckKey,




state: {

init() {

return DecorationSet.empty

      },




apply(tr, value, oldState, newState) {

// 1. Handle special meta that replaces the whole decoration set (clear all)

const meta = tr.getMeta(statcheckKey)




// 2. Always map existing decorations through the transaction

let newPluginStateDecorations = value.map(tr.mapping, tr.doc)




paragraphDecorationMap = new Map();




newPluginStateDecorations?.find(undefined, undefined, spec => spec.plugin === 'proofRead').forEach(deco => {




if (deco.from > newState.doc.content.size) return;




const resolvedPos = newState.doc.resolve(deco.from);

const paragraphNode = resolvedPos.parent;

const paragraphNodePos = resolvedPos.start(resolvedPos.depth);




if (paragraphNode && paragraphNode.isTextblock) {

const paragraphId = `p-${paragraphNodePos}`;

if (!paragraphDecorationMap.has(paragraphId)) {

paragraphDecorationMap.set(paragraphId, []);

            }

paragraphDecorationMap.get(paragraphId).push(deco);

          }

        });




// 3. Handle updates from async checkAndDecorateParagraph (via meta)

if (meta && meta.updateParagraphDecorations) {

const { id: updatedParagraphId, decorations: newDecorationsForParagraph } = meta.updateParagraphDecorations;




// Get the decorations currently associated with this paragraph ID from our map

const oldDecorationsForParagraph = paragraphDecorationMap.get(updatedParagraphId);

console.log(oldDecorationsForParagraph);

// Remove old decorations for this paragraph from the current DecorationSet

if (oldDecorationsForParagraph && oldDecorationsForParagraph.length > 0) {

newPluginStateDecorations = newPluginStateDecorations.remove(oldDecorationsForParagraph);

          }




// Add new decorations for this paragraph to the current DecorationSet

if (newDecorationsForParagraph && newDecorationsForParagraph.length > 0) {

newPluginStateDecorations = newPluginStateDecorations.add(tr.doc, newDecorationsForParagraph);

// Update the map with the NEW decorations for this paragraph

paragraphDecorationMap.set(updatedParagraphId, newDecorationsForParagraph);

          } else {

// If no new decorations, and there were old ones, ensure they are removed from the map

// This handles the case where all errors in a paragraph are fixed.

paragraphDecorationMap.delete(updatedParagraphId);

          }

        } else if (meta?.type == REMOVE_DECORATION_META_KEY) { // remove decoration area

const meta = tr.getMeta(statcheckKey);

if (meta?.decorations) {

return meta.decorations; // replace entire DecorationSet

          }

return oldState.map(tr.mapping, tr.doc); 

        }




// 4. Schedule a new check for the active paragraph only if the document changed

if (!oldState.doc.eq(newState.doc)) {

const selection = newState.selection

let activeParagraphNode: any = null

let activeParagraphNodePos = -1




if (selection.empty) {

const resolvedPos = selection.$head

activeParagraphNode = resolvedPos.parent

activeParagraphNodePos = resolvedPos.start(resolvedPos.depth)

          } else {

const resolvedStart = selection.$from

activeParagraphNode = resolvedStart.parent

activeParagraphNodePos = resolvedStart.start(resolvedStart.depth)

          }




if (activeParagraphNode && activeParagraphNode.isTextblock) {

const domNode = editor.view.nodeDOM(activeParagraphNodePos - 1)

if (domNode instanceof Element && domNode.closest('ref')) {

clearTimeout(debounceTimeoutId)

return newPluginStateDecorations

            }

const currentText = activeParagraphNode.textContent

const paragraphId = `p-${activeParagraphNodePos}`

const { visibleText, mapVisToOrig } = getVisibleTextAndMap(activeParagraphNode, activeParagraphNodePos)




const oldMappedPos = tr.mapping.map(activeParagraphNodePos)

const oldParagraphNode = (oldMappedPos >= 0 && oldMappedPos < oldState.doc.content.size)

? oldState.doc.resolve(oldMappedPos)

: null

const oldText = oldParagraphNode ? oldParagraphNode.parent?.textContent : ''




if (oldText !== currentText) {

clearTimeout(debounceTimeoutId)

debounceTimeoutId = setTimeout(() => {

checkAndDecorateParagraph(visibleText, paragraphId, mapVisToOrig, editor.view)

              }, DEBOUNCE_DELAY)

            }

          } else {

clearTimeout(debounceTimeoutId)

          }

        }




return newPluginStateDecorations

      },

    },




view(view) {

latestView = view

return {}

    },




props: {

decorations(state) {

return this.getState(state)

      },

    },

  })

}

please note that the pargraphId is p-start position of a pargraph

You’re doing imperative things (setting timeouts, changing globals) in your apply reducer function. That is not a good idea. Register a view plugin if you need to do things like that in a plugin.

Thanks for the suggestion — I’ve now moved all side effects (timeouts, async calls) out of apply() into a view plugin, so apply() is now pure. However, the issue still persists:

  • After running statcheck → all decorations render correctly

  • After editing a paragraph → only that paragraph updates (expected)

  • After undo (Ctrl+Z) → most decorations disappear

  • Only the active paragraph retains decorations (if revalidated)

Solution Found

Two issues caused this:

1. Don’t put async calls inside apply() ProseMirror replays apply() during undo, causing the async callback to fire again and overwrite correctly restored decorations. Fix: move all async logic into the view() update hook.

2. Yjs undo is not a real undo Native ProseMirror undo uses inverted steps so value.map() works correctly. But Yjs replaces the entire document in one step, causing value.map() to drop all decorations. Fix: detect Yjs transactions via tr.getMeta('y-sync$')?.isChangeOrigin and rebuild decorations by re-matching raw strings against the new document.