diff --git a/docs/src/demos/Experiments/Annotation/extension/AnnotationPlugin.ts b/docs/src/demos/Experiments/Annotation/extension/AnnotationPlugin.ts
index 65a7440f..68b08a80 100644
--- a/docs/src/demos/Experiments/Annotation/extension/AnnotationPlugin.ts
+++ b/docs/src/demos/Experiments/Annotation/extension/AnnotationPlugin.ts
@@ -14,6 +14,7 @@ export interface AnnotationPluginOptions {
 
 export const AnnotationPlugin = (options: AnnotationPluginOptions) => new Plugin({
   key: AnnotationPluginKey,
+
   state: {
     init() {
       return new AnnotationState({
@@ -25,6 +26,7 @@ export const AnnotationPlugin = (options: AnnotationPluginOptions) => new Plugin
       return pluginState.apply(transaction, newState)
     },
   },
+
   props: {
     decorations(state) {
       const { decorations } = this.getState(state)
diff --git a/docs/src/demos/Experiments/Annotation/extension/AnnotationState.ts b/docs/src/demos/Experiments/Annotation/extension/AnnotationState.ts
index 154f073f..45c4073a 100644
--- a/docs/src/demos/Experiments/Annotation/extension/AnnotationState.ts
+++ b/docs/src/demos/Experiments/Annotation/extension/AnnotationState.ts
@@ -1,8 +1,9 @@
 // @ts-nocheck
 import * as Y from 'yjs'
-import { EditorState } from 'prosemirror-state'
+import { EditorState, Transaction } from 'prosemirror-state'
 import { Decoration, DecorationSet } from 'prosemirror-view'
 import { ySyncPluginKey, relativePositionToAbsolutePosition, absolutePositionToRelativePosition } from 'y-prosemirror'
+import { AddAnnotationAction, DeleteAnnotationAction } from './annotation'
 import { AnnotationPluginKey } from './AnnotationPlugin'
 
 export interface AnnotationStateOptions {
@@ -39,57 +40,72 @@ export class AnnotationState {
     }
   }
 
+  addAnnotation(action: AddAnnotationAction, state: EditorState) {
+    const ystate = ySyncPluginKey.getState(state)
+    const { type, binding } = ystate
+    const { map, HTMLAttributes } = this.options
+    const { from, to, data } = action
+    const absoluteFrom = absolutePositionToRelativePosition(from, type, binding.mapping)
+    const absoluteTo = absolutePositionToRelativePosition(to, type, binding.mapping)
+
+    map.set(data.id, {
+      from: absoluteFrom,
+      to: absoluteTo,
+      data,
+    })
+
+    const decoration = Decoration.inline(from, to, HTMLAttributes, { data })
+
+    this.decorations = this.decorations.add(state.doc, [decoration])
+  }
+
+  deleteAnnotation(id: number) {
+    const { map } = this.options
+    const decoration = this.findAnnotation(id)
+
+    map.delete(id)
+    this.decorations = this.decorations.remove([decoration])
+  }
+
   annotationsAt(position: number) {
     return this.decorations?.find(position, position)
   }
 
-  apply(transaction: any, state: EditorState) {
+  updateDecorations(state: EditorState) {
     const { map, HTMLAttributes } = this.options
     const ystate = ySyncPluginKey.getState(state)
     const { doc, type, binding } = ystate
-    const action = transaction.getMeta(AnnotationPluginKey)
+
+    const decorations = Array.from(map.keys()).map(id => {
+      const dec = map.get(id)
+      const from = relativePositionToAbsolutePosition(doc, type, dec.from, binding.mapping)
+      const to = relativePositionToAbsolutePosition(doc, type, dec.to, binding.mapping)
+      const decoration = Decoration.inline(from, to, HTMLAttributes, { data: dec.data })
+
+      return decoration
+    })
+
+    this.decorations = DecorationSet.create(state.doc, decorations)
+  }
+
+  apply(transaction: Transaction, state: EditorState) {
+    const ystate = ySyncPluginKey.getState(state)
+    const action = transaction.getMeta(AnnotationPluginKey) as AddAnnotationAction | DeleteAnnotationAction
 
     if (action && action.type) {
-      const { from, to, data } = action
-
       if (action.type === 'addAnnotation') {
-        const absoluteFrom = absolutePositionToRelativePosition(from, type, binding.mapping)
-        const absoluteTo = absolutePositionToRelativePosition(to, type, binding.mapping)
-
-        map.set(data.id, {
-          from: absoluteFrom,
-          to: absoluteTo,
-          data,
-        })
-
-        const decoration = Decoration.inline(from, to, HTMLAttributes, { data })
-
-        this.decorations = this.decorations.add(transaction.doc, [decoration])
+        this.addAnnotation(action, state)
       }
 
       if (action.type === 'deleteAnnotation') {
-        map.delete(action.id)
-
-        const decoration = this.findAnnotation(action.id)
-
-        this.decorations = this.decorations.remove([decoration])
+        this.deleteAnnotation(action.id)
       }
 
       return this
     }
 
     if (ystate.isChangeOrigin) {
-
-      const decorations = Array.from(map.keys()).map(id => {
-        const dec = map.get(id)
-        const from = relativePositionToAbsolutePosition(doc, type, dec.from, binding.mapping)
-        const to = relativePositionToAbsolutePosition(doc, type, dec.to, binding.mapping)
-        const decoration = Decoration.inline(from, to, HTMLAttributes, { data: dec.data })
-
-        return decoration
-      })
-
-      this.decorations = DecorationSet.create(state.doc, decorations)
+      this.updateDecorations(state)
 
       return this
     }
@@ -99,4 +115,5 @@ export class AnnotationState {
 
     return this
   }
+
 }
diff --git a/docs/src/demos/Experiments/Annotation/extension/annotation.ts b/docs/src/demos/Experiments/Annotation/extension/annotation.ts
index 91f6a7fb..7a6c92c0 100644
--- a/docs/src/demos/Experiments/Annotation/extension/annotation.ts
+++ b/docs/src/demos/Experiments/Annotation/extension/annotation.ts
@@ -7,6 +7,18 @@ function randomId() {
   return Math.floor(Math.random() * 0xffffffff)
 }
 
+export interface AddAnnotationAction {
+  type: 'addAnnotation',
+  from: number,
+  to: number,
+  data: AnnotationItem,
+}
+
+export interface DeleteAnnotationAction {
+  id: number,
+  type: 'deleteAnnotation',
+}
+
 export interface AnnotationOptions {
   HTMLAttributes: {
     [key: string]: any
@@ -49,7 +61,7 @@ export const Annotation = Extension.create({
         }
 
         if (dispatch && content) {
-          dispatch(state.tr.setMeta(AnnotationPluginKey, {
+          state.tr.setMeta(AnnotationPluginKey, <AddAnnotationAction>{
             type: 'addAnnotation',
             from: selection.from,
             to: selection.to,
@@ -57,14 +69,17 @@ export const Annotation = Extension.create({
               randomId(),
               content,
             ),
-          }))
+          })
         }
 
         return true
       },
       deleteAnnotation: (id: number): Command => ({ dispatch, state }) => {
         if (dispatch) {
-          dispatch(state.tr.setMeta(AnnotationPluginKey, { type: 'deleteAnnotation', id }))
+          state.tr.setMeta(AnnotationPluginKey, <DeleteAnnotationAction>{
+            type: 'deleteAnnotation',
+            id,
+          })
         }
 
         return true