diff --git a/editors/demo/src/Editoria/Editoria.js b/editors/demo/src/Editoria/Editoria.js
index 8837b278e3d8c3e0e7c6f5a3ea237f63a71cfb22..3e85568eaea88403be4213b9ad9291685b3f3d73 100644
--- a/editors/demo/src/Editoria/Editoria.js
+++ b/editors/demo/src/Editoria/Editoria.js
@@ -55,7 +55,7 @@ const Editoria = () => {
           autoFocus
           placeholder="Type Something..."
           fileUpload={file => renderImage(file)}
-          value={demo}
+          // value={demo}
           // readonly
           layout={layout}
           // onChange={debounce(source => {
diff --git a/editors/demo/src/Editoria/config/config.js b/editors/demo/src/Editoria/config/config.js
index a3d88270db8064ef37483ad944169e5aceb75ee3..e643d45305b693ffd8e57ac99b3b878ec6c43870 100644
--- a/editors/demo/src/Editoria/config/config.js
+++ b/editors/demo/src/Editoria/config/config.js
@@ -26,7 +26,8 @@ import {
   disallowPasteImagesPlugin,
   BlockDropDownToolGroupService,
   AskAiContentService,
-  // YjsService,
+  YjsService,
+  CommentDecoration,
 } from 'wax-prosemirror-services';
 
 import { TablesService, tableEditing, columnResizing } from 'wax-table-service';
@@ -61,9 +62,136 @@ async function DummyPromise(userInput) {
 }
 
 const updateTitle = debounce(title => {
-  console.log(title);
+  // console.log(title);
 }, 100);
 
+const getComments = debounce(comments => {
+  console.log(comments);
+}, 2000);
+
+const setComments = (
+  comments = [
+    {
+      id: 'a1',
+      from: 5,
+      to: 10,
+      data: {
+        type: 'comment',
+        yjsFrom: 5,
+        yjsTo: 10,
+        pmFrom: 5,
+        pmTo: 10,
+        conversation: [
+          {
+            content: '1111',
+            displayName: 'admin',
+            userId: 'b3cfc28e-0f2e-45b5-b505-e66783d4f946',
+            timestamp: 1710501980537,
+          },
+        ],
+        title: '111',
+        group: 'main',
+        viewId: 'main',
+      },
+      endHeight: 362.3579406738281,
+    },
+    {
+      id: 'a2',
+      from: 8,
+      to: 15,
+      data: {
+        type: 'comment',
+        yjsFrom: 8,
+        yjsTo: 15,
+        pmFrom: 8,
+        pmTo: 15,
+        conversation: [
+          {
+            content: '222',
+            displayName: 'admin',
+            userId: 'b3cfc28e-0f2e-45b5-b505-e66783d4f946',
+            timestamp: 1710501987197,
+          },
+        ],
+        title: '222',
+        group: 'main',
+        viewId: 'main',
+      },
+      endHeight: 266.3579406738281,
+    },
+    {
+      id: 'b8d907d4-1859-49a9-abcd-13788d497758',
+      from: {
+        type: {
+          client: 2887320119,
+          clock: 150,
+        },
+        tname: null,
+        item: {
+          client: 2887320119,
+          clock: 185,
+        },
+        assoc: 0,
+      },
+      to: {
+        type: {
+          client: 2887320119,
+          clock: 150,
+        },
+        tname: null,
+        item: {
+          client: 2887320119,
+          clock: 195,
+        },
+        assoc: 0,
+      },
+      data: {
+        yjsFrom: {
+          type: {
+            client: 2887320119,
+            clock: 150,
+          },
+          tname: null,
+          item: {
+            client: 2887320119,
+            clock: 185,
+          },
+          assoc: 0,
+        },
+        yjsTo: {
+          type: {
+            client: 2887320119,
+            clock: 150,
+          },
+          tname: null,
+          item: {
+            client: 2887320119,
+            clock: 195,
+          },
+          assoc: 0,
+        },
+        pmFrom: 164,
+        pmTo: 174,
+        type: 'comment',
+        conversation: [
+          {
+            content: 'dfgdfgd',
+            displayName: 'admin',
+            userId: 'b3cfc28e-0f2e-45b5-b505-e66783d4f946',
+            timestamp: 1713699155995,
+          },
+        ],
+        title: 'dgfdgf',
+        group: 'main',
+        viewId: 'main',
+      },
+      endHeight: 406.734375,
+    },
+  ],
+) => {
+  return comments;
+};
+
 const saveTags = tags => {
   // console.log(tags);
 };
@@ -97,7 +225,7 @@ export default {
         'HighlightToolGroup',
         'TransformToolGroup',
         'CustomTagInline',
-        'Notes',
+        // 'Notes',
         'Lists',
         'Images',
         'SpecialCharacters',
@@ -159,9 +287,13 @@ export default {
     ),
   ],
   ImageService: { showAlt: true },
+
   CommentsService: {
     showTitle: true,
+    getComments,
+    setComments,
   },
+
   CustomTagService: {
     tags: [
       { label: 'custom-tag-label-1', tagType: 'inline' },
@@ -171,11 +303,11 @@ export default {
     ],
     updateTags: saveTags,
   },
-  // YjsService: {
-  //   // eslint-disable-next-line no-restricted-globals
-  //   connectionUrl: 'ws://localhost:4000',
-  //   docIdentifier: 'prosemirror-demo',
-  // },
+  YjsService: {
+    // eslint-disable-next-line no-restricted-globals
+    connectionUrl: 'ws://localhost:4000',
+    docIdentifier: 'prosemirror-demo',
+  },
 
   AskAiContentService: {
     AskAiContentTransformation: DummyPromise,
@@ -183,7 +315,7 @@ export default {
   },
 
   services: [
-    // new YjsService(),
+    new YjsService(),
     new BlockDropDownToolGroupService(),
     new AskAiContentService(),
     new CustomTagService(),
@@ -197,7 +329,7 @@ export default {
     new ImageService(),
     new TablesService(),
     new BaseService(),
-    new NoteService(),
+    // new NoteService(),
     new CodeBlockService(),
     new EditingSuggestingService(),
     new DisplayTextToolGroupService(),
diff --git a/editors/demo/src/Editors.js b/editors/demo/src/Editors.js
index 0e36a5c07cdfbce6f7c2095a8b1470e6f4c781f5..78a82dca42828507ad8168cf3eb705360720ee0d 100644
--- a/editors/demo/src/Editors.js
+++ b/editors/demo/src/Editors.js
@@ -90,7 +90,7 @@ const Editors = () => {
       case 'oen':
         return <OEN />;
       default:
-        return <HHMI />;
+        return <Editoria />;
     }
   };
 
diff --git a/wax-prosemirror-core/src/WaxContext.js b/wax-prosemirror-core/src/WaxContext.js
index 4d3ee58f9760b54b40e156a29e7b0f38a6ed72b1..d97ec379c28c616b9f6b18945d31ae29ab0c433d 100644
--- a/wax-prosemirror-core/src/WaxContext.js
+++ b/wax-prosemirror-core/src/WaxContext.js
@@ -18,7 +18,7 @@ export default props => {
     pmViews: props.view || {},
     activeView: props.activeView || {},
     activeViewId: props.activeViewId || {},
-    options: { fullScreen: false, AiOn: false },
+    options: { fullScreen: false },
     transaction: {},
     setTransaction: tr => {
       Object.assign(context.transaction, tr);
diff --git a/wax-prosemirror-core/src/utilities/commands/Commands.js b/wax-prosemirror-core/src/utilities/commands/Commands.js
index 90f7b1ad8b131c1535c5fccc2059b6fd2aeb0569..138c75ceaca59ce4da8de74256c2e1dc8b7ecf60 100644
--- a/wax-prosemirror-core/src/utilities/commands/Commands.js
+++ b/wax-prosemirror-core/src/utilities/commands/Commands.js
@@ -129,172 +129,6 @@ const isOnSameTextBlock = state => {
   return false;
 };
 
-const createComment = (
-  state,
-  dispatch,
-  group,
-  viewid,
-  conversation = [],
-  title = '',
-  posFrom,
-  posTo,
-) => {
-  const {
-    selection: { $from, $to },
-    tr,
-  } = state;
-  let fromPosition = $from.pos;
-  let toPosition = $to.pos;
-  if ($from.pos === $to.pos) {
-    fromPosition = posFrom;
-    toPosition = posTo;
-  }
-  let footnote = false;
-  let footnoteNode;
-  state.doc.nodesBetween($from.pos, $to.pos, (node, from) => {
-    if (node.type.groups.includes('notes')) {
-      footnote = true;
-      footnoteNode = node;
-    }
-  });
-
-  if (footnote) {
-    if (
-      footnoteNode.content.size + 2 ===
-      state.selection.to - state.selection.from
-    ) {
-      return createCommentOnSingleFootnote(
-        state,
-        dispatch,
-        group,
-        viewid,
-        conversation,
-        title,
-      );
-    }
-    return createCommentOnFootnote(
-      state,
-      dispatch,
-      group,
-      viewid,
-      conversation,
-      title,
-    );
-  }
-  dispatch(
-    state.tr
-      .addMark(
-        fromPosition,
-        toPosition,
-        state.config.schema.marks.comment.create({
-          id: uuidv4(),
-          group,
-          viewid,
-          conversation,
-          title,
-        }),
-      )
-      .setMeta('forceUpdate', true),
-  );
-};
-
-const createCommentOnSingleFootnote = (
-  state,
-  dispatch,
-  group,
-  viewid,
-  conversation,
-  title,
-) => {
-  const { tr } = state;
-  tr.step(
-    new AddMarkStep(
-      state.selection.from,
-      state.selection.from + 1,
-      state.config.schema.marks.comment.create({
-        id: uuidv4(),
-        group,
-        conversation,
-        viewid,
-        title,
-      }),
-    ),
-  ).setMeta('forceUpdate', true);
-  dispatch(tr);
-};
-
-const createCommentOnFootnote = (
-  state,
-  dispatch,
-  group,
-  viewid,
-  conversation,
-  title,
-) => {
-  const {
-    selection: { $from },
-    selection,
-    tr,
-  } = state;
-
-  const { content } = $from.parent;
-  const $pos = state.doc.resolve($from.pos);
-  const commentStart = $from.pos - $pos.textOffset;
-  const commentEnd = commentStart + $pos.parent.child($pos.index()).nodeSize;
-
-  let start = $from.pos;
-  let end = commentEnd;
-  const ranges = [];
-
-  const allFragments = [];
-
-  selection.content().content.content.forEach(node => {
-    node.content.content.forEach(fragment => {
-      allFragments.push(fragment);
-    });
-  });
-
-  allFragments.forEach((contentNode, index) => {
-    start = index === 0 ? start : end;
-    end = index === 0 ? end : end + contentNode.nodeSize;
-    ranges.push({
-      start,
-      end,
-      footnote: contentNode.type.groups.includes('notes'),
-    });
-  });
-
-  const mergedRanges = [];
-  ranges.forEach((item, i) => {
-    if (item.footnote) {
-      mergedRanges[mergedRanges.length - 1].end =
-        mergedRanges[mergedRanges.length - 1].end + 1;
-    } else {
-      mergedRanges.push(item);
-    }
-  });
-
-  const id = uuidv4();
-
-  mergedRanges.forEach(range => {
-    tr.step(
-      new AddMarkStep(
-        range.start,
-        range.end,
-        state.config.schema.marks.comment.create({
-          id,
-          group,
-          conversation,
-          viewid,
-          title,
-        }),
-      ),
-    ).setMeta('forceUpdate', true);
-  });
-
-  dispatch(tr);
-};
-
 const isInTable = state => {
   const { $head } = state.selection;
   for (let d = $head.depth; d > 0; d -= 1)
@@ -328,7 +162,6 @@ export default {
   blockActive,
   customTagBlockActive,
   canInsert,
-  createComment,
   createLink,
   createTable,
   markActive,
diff --git a/wax-prosemirror-services/index.js b/wax-prosemirror-services/index.js
index 95b5a5596a5545d5e9c7daf290e380fdae1736fa..32f70e7d1af089a4e4e96fb75e225417bb7364d3 100644
--- a/wax-prosemirror-services/index.js
+++ b/wax-prosemirror-services/index.js
@@ -36,6 +36,10 @@ ToolGroups
 */
 export { default as DisplayTextToolGroupService } from './src/WaxToolGroups/DisplayTextToolGroupService/DisplayTextToolGroupService';
 export { default as BlockDropDownToolGroupService } from './src/WaxToolGroups/BlockDropDownToolGroupService/BlockDropDownToolGroupService';
-/* Plugins */
 
+/* Plugins */
 export { default as disallowPasteImagesPlugin } from './src/ImageService/plugins/disallowPasteImagesPlugin';
+export {
+  CommentDecorationPlugin,
+  CommentDecorationPluginKey,
+} from './src/CommentsService/plugins/CommentDecorationPlugin';
diff --git a/wax-prosemirror-services/src/AiService/components/ToggleAiComponent.js b/wax-prosemirror-services/src/AiService/components/ToggleAiComponent.js
index 6575a79f97262bbe6a17c86cd3b80a347266cdd2..b730ea0763650ff342529e6f3dbce9e01873872b 100644
--- a/wax-prosemirror-services/src/AiService/components/ToggleAiComponent.js
+++ b/wax-prosemirror-services/src/AiService/components/ToggleAiComponent.js
@@ -37,13 +37,20 @@ const ToggleAiComponent = ({ item }) => {
       enableService.AiOn ? (
         <MenuButton
           active={checked}
-          disabled={!isEditable}
+          disabled={
+            !isEditable || main.state.selection.from === main.state.selection.to
+          }
           iconName={item.icon}
           onMouseDown={onMouseDown}
           title={item.title}
         />
       ) : null,
-    [checked, isDisabled, enableService.AiOn],
+    [
+      checked,
+      isDisabled,
+      enableService.AiOn,
+      main.state.selection.from === main.state.selection.to,
+    ],
   );
 };
 
diff --git a/wax-prosemirror-services/src/CommentsService/CommentsService.js b/wax-prosemirror-services/src/CommentsService/CommentsService.js
index fc32e5b5cfb5aa2c662071a32bd24f2ecdfcb9c1..60668775b01d92c27fb2bf000f6a3eb61eaa9550 100644
--- a/wax-prosemirror-services/src/CommentsService/CommentsService.js
+++ b/wax-prosemirror-services/src/CommentsService/CommentsService.js
@@ -3,18 +3,63 @@ import CommentBubbleComponent from './components/ui/comments/CommentBubbleCompon
 import RightArea from './components/RightArea';
 import commentMark from './schema/commentMark';
 import CommentPlugin from './plugins/CommentPlugin';
-import CopyPasteCommentPlugin from './plugins/CopyPasteCommentPlugin';
+import { CommentDecorationPlugin } from './plugins/CommentDecorationPlugin';
 import './comments.css';
 
-const PLUGIN_KEY = 'commentPlugin';
-
 export default class CommentsService extends Service {
+  allCommentsFromStates = [];
   boot() {
-    this.app.PmPlugins.add(PLUGIN_KEY, CommentPlugin(PLUGIN_KEY));
+    const commentsConfig = this.app.config.get('config.CommentsService');
+
+    this.app.PmPlugins.add(
+      'commentPlugin',
+      CommentPlugin('commentPlugin', this.app.context),
+    );
+
+    const options = {
+      existingComments: () => {
+        const map = this.app.config.get('config.YjsService')
+          ? this.app.context.options.currentYdoc.getMap('comments')
+          : new Map();
+
+        if (commentsConfig.setComments().length > 0) {
+          commentsConfig.setComments().forEach(value => {
+            map.set(value.id, value);
+          });
+        }
+
+        this.app.context.setOption({
+          commentsMap: map,
+        });
+        return map;
+      },
+      context: this.app.context,
+      onSelectionChange: items => {
+        this.allCommentsFromStates = this.allCommentsFromStates.filter(
+          comm =>
+            (items.find(item => item.id === comm.id) || {}).id !== comm.id,
+        );
+        this.allCommentsFromStates = this.allCommentsFromStates.concat([
+          ...items,
+        ]);
+
+        if (this.app.context.options.resolvedComment) {
+          this.allCommentsFromStates = this.allCommentsFromStates.filter(
+            comm => {
+              return comm.id !== this.app.context.options.resolvedComment;
+            },
+          );
+        }
+        commentsConfig.getComments(this.allCommentsFromStates);
+        this.app.context.setOption({ comments: this.allCommentsFromStates });
+      },
+    };
+
     this.app.PmPlugins.add(
-      'copyPasteCommentPlugin',
-      CopyPasteCommentPlugin('copyPasteCommentPlugin', this.app.context),
+      'CommentDecorationPlugin',
+      CommentDecorationPlugin('commentDecorationPlugin', options),
     );
+
     const createOverlay = this.container.get('CreateOverlay');
     const layout = this.container.get('Layout');
     createOverlay(
diff --git a/wax-prosemirror-services/src/CommentsService/comments.css b/wax-prosemirror-services/src/CommentsService/comments.css
index 3f7045dfc3db4b7a2c55e9f98c144a1f513f550e..1cd4d219db77e2a123ee610df53b8569fd8d15f6 100644
--- a/wax-prosemirror-services/src/CommentsService/comments.css
+++ b/wax-prosemirror-services/src/CommentsService/comments.css
@@ -5,7 +5,7 @@
     border-radius: 3px 3px 0 0;
   }
 
-  span.comment .active-comment {
+.active-comment {
        background-color: gold;
       /* color: black; */
   }
\ No newline at end of file
diff --git a/wax-prosemirror-services/src/CommentsService/components/BoxList.js b/wax-prosemirror-services/src/CommentsService/components/BoxList.js
index 1b19e3839d19415b10831285c38266d40417c013..5ce698876f309032e6d07cb7bd9a3e5422b7d6a5 100644
--- a/wax-prosemirror-services/src/CommentsService/components/BoxList.js
+++ b/wax-prosemirror-services/src/CommentsService/components/BoxList.js
@@ -1,5 +1,4 @@
 /* eslint react/prop-types: 0 */
-import { Mark } from 'prosemirror-model';
 import React from 'react';
 import ConnectedComment from './ConnectedComment';
 import ConnectedTrackChange from './ConnectedTrackChange';
@@ -9,14 +8,19 @@ export default ({ commentsTracks, view, position, recalculateTops, users }) => {
   return (
     <>
       {commentsTracks.map((commentTrack, index) => {
-        const id =
-          commentTrack instanceof Mark
-            ? commentTrack.attrs.id
-            : commentTrack.node.attrs.id;
+        let id = '';
+
+        if (commentTrack?.node?.attrs.id) {
+          id = commentTrack.node.attrs.id;
+        } else if (commentTrack?.attrs?.id) {
+          id = commentTrack.attrs.id;
+        } else {
+          id = commentTrack.id;
+        }
 
         const top = position[index] ? position[index][id] : 0;
 
-        if (commentTrack.type && commentTrack.type.name === 'comment') {
+        if (commentTrack.data?.type === 'comment') {
           return (
             <ConnectedComment
               comment={commentTrack}
diff --git a/wax-prosemirror-services/src/CommentsService/components/ConnectedComment.js b/wax-prosemirror-services/src/CommentsService/components/ConnectedComment.js
index b119fc666f5f4c8e40a616bf5cb65b11738925e6..979fef8771c881b170dad6405998e86bff55330e 100644
--- a/wax-prosemirror-services/src/CommentsService/components/ConnectedComment.js
+++ b/wax-prosemirror-services/src/CommentsService/components/ConnectedComment.js
@@ -2,11 +2,11 @@
 /* eslint react/prop-types: 0 */
 import React, { useContext, useMemo, useState, useEffect } from 'react';
 import { TextSelection } from 'prosemirror-state';
-import { last, maxBy, minBy } from 'lodash';
 import styled from 'styled-components';
-import { WaxContext, DocumentHelpers, Commands } from 'wax-prosemirror-core';
+import { WaxContext } from 'wax-prosemirror-core';
 import { override } from '@pubsweet/ui-toolkit';
 import CommentBox from './ui/comments/CommentBox';
+import { CommentDecorationPluginKey } from '../plugins/CommentDecorationPlugin';
 
 const ConnectedCommentStyled = styled.div`
   margin-left: ${props => (props.active ? `${-20}px` : `${50}px`)};
@@ -33,23 +33,14 @@ export default ({ comment, top, commentId, recalculateTops, users }) => {
     app,
     activeView,
     activeViewId,
+    options: { comments, commentsMap },
   } = context;
 
   const [isActive, setIsActive] = useState(false);
   const [clickPost, setClickPost] = useState(false);
 
   const { state, dispatch } = activeView;
-  const viewId = comment.attrs.viewid;
-  let allCommentsWithSameId = [];
-
-  if (pmViews[viewId]) {
-    allCommentsWithSameId = DocumentHelpers.findAllMarksWithSameId(
-      pmViews[viewId].state,
-      comment,
-    );
-  }
-  const commentMark = state.schema.marks.comment;
-
+  const { viewId, conversation } = comment.data;
   const styles = {
     top: `${top}px`,
   };
@@ -65,9 +56,8 @@ export default ({ comment, top, commentId, recalculateTops, users }) => {
   useEffect(() => {
     setIsActive(false);
     recalculateTops();
-    if (activeComment && commentId === activeComment.attrs.id) {
+    if (activeComment && commentId === activeComment.id) {
       setIsActive(true);
-      recalculateTops();
     }
   }, [activeComment]);
 
@@ -84,35 +74,19 @@ export default ({ comment, top, commentId, recalculateTops, users }) => {
       timestamp: Math.floor(Date.now()),
     };
 
-    comment.attrs.title = title || comment.attrs.title;
-    comment.attrs.conversation.push(obj);
+    comment.data.title = title || comment.data.title;
+    comment.data.conversation.push(obj);
 
-    allCommentsWithSameId.forEach(singleComment => {
-      activeView.dispatch(
-        activeView.state.tr.removeMark(
-          singleComment.pos,
-          singleComment.pos + singleComment.node.nodeSize,
-          commentMark,
-        ),
-      );
-    });
-
-    const minPos = minBy(allCommentsWithSameId, 'pos');
-    const maxPos = maxBy(allCommentsWithSameId, 'pos');
-
-    Commands.createComment(
-      activeView.state,
-      activeView.dispatch,
-      comment.attrs.group,
-      comment.attrs.viewid,
-      comment.attrs.conversation,
-      comment.attrs.title,
-      minPos.pos,
-      maxPos.pos + last(allCommentsWithSameId).node.nodeSize,
+    dispatch(
+      state.tr.setMeta(CommentDecorationPluginKey, {
+        type: 'updateComment',
+        id: activeComment.id,
+        data: comment.data,
+      }),
     );
-
     activeView.focus();
-    recalculateTops();
+
+    sendYjsUpdate();
   };
 
   const onClickBox = () => {
@@ -123,12 +97,11 @@ export default ({ comment, top, commentId, recalculateTops, users }) => {
 
     if (viewId !== 'main') context.updateView({}, viewId);
 
-    const maxPos = maxBy(allCommentsWithSameId, 'pos');
-    maxPos.pos += last(allCommentsWithSameId).node.nodeSize;
-
     pmViews[viewId].dispatch(
       pmViews[viewId].state.tr.setSelection(
-        new TextSelection(pmViews[viewId].state.tr.doc.resolve(maxPos.pos)),
+        new TextSelection(
+          pmViews[viewId].state.tr.doc.resolve(comment.data.pmFrom),
+        ),
       ),
     );
 
@@ -137,36 +110,39 @@ export default ({ comment, top, commentId, recalculateTops, users }) => {
   };
 
   const onClickResolve = () => {
-    let maxPos = comment.pos;
-    let minPos = comment.pos;
-
-    allCommentsWithSameId.forEach(singleComment => {
-      const markPosition = DocumentHelpers.findMarkPosition(
-        state,
-        singleComment.pos,
-        'comment',
-      );
-      if (markPosition.from < minPos) minPos = markPosition.from;
-      if (markPosition.to > maxPos) maxPos = markPosition.to;
-    });
-    // if (allCommentsWithSameId.length > 1);
-    // maxPos += last(allCommentsWithSameId).node.nodeSize;
-    recalculateTops();
-    dispatch(state.tr.removeMark(minPos, maxPos, commentMark));
+    context.setOption({ resolvedComment: activeComment.id });
+    dispatch(
+      state.tr.setMeta(CommentDecorationPluginKey, {
+        type: 'deleteComment',
+        id: activeComment.id,
+      }),
+    );
+
+    dispatch(state.tr);
     activeView.focus();
+    sendYjsUpdate();
   };
 
   const onTextAreaBlur = () => {
-    // TODO Save into local storage
-    // if (content !== '') {
-    //   onClickPost(content);
-    // }
-    setTimeout(() => {
-      if (comment.attrs.conversation.length === 0 && !clickPost) {
-        onClickResolve();
-        activeView.focus();
-      }
-    }, 400);
+    if (conversation.length === 0 && !clickPost) {
+      onClickResolve();
+      activeView.focus();
+    }
+  };
+
+  const sendYjsUpdate = () => {
+    if (context.app.config.get('config.YjsService')) {
+      commentsMap.observe(() => {
+        const transaction = context.pmViews.main.state.tr.setMeta(
+          CommentDecorationPluginKey,
+          {
+            type: 'createDecorations',
+          },
+        );
+
+        context.pmViews.main.dispatch(transaction);
+      });
+    }
   };
 
   const MemorizedComponent = useMemo(
@@ -174,12 +150,12 @@ export default ({ comment, top, commentId, recalculateTops, users }) => {
       <ConnectedCommentStyled
         active={isActive}
         data-box={commentId}
-        length={comment.attrs.conversation.length === 0}
+        length={conversation.length === 0}
         style={styles}
       >
         <CommentBox
           active={isActive}
-          commentData={comment.attrs.conversation}
+          commentData={conversation}
           commentId={commentId}
           isReadOnly={isReadOnly}
           key={commentId}
@@ -187,14 +163,13 @@ export default ({ comment, top, commentId, recalculateTops, users }) => {
           onClickPost={onClickPost}
           onClickResolve={onClickResolve}
           onTextAreaBlur={onTextAreaBlur}
-          recalculateTops={recalculateTops}
           showTitle={showTitle}
-          title={comment.attrs.title}
+          title={comment.data.title}
           users={users}
         />
       </ConnectedCommentStyled>
     ),
-    [isActive, top, comment.attrs.conversation.length, users],
+    [isActive, top, conversation.length, users],
   );
   return <>{MemorizedComponent}</>;
 };
diff --git a/wax-prosemirror-services/src/CommentsService/components/ConnectedTrackChange.js b/wax-prosemirror-services/src/CommentsService/components/ConnectedTrackChange.js
index 6b67fb815b9dbe165a0b27645d571ae78a3561be..a8f9af7934402724c0915c56bb1857a9bcd1894e 100644
--- a/wax-prosemirror-services/src/CommentsService/components/ConnectedTrackChange.js
+++ b/wax-prosemirror-services/src/CommentsService/components/ConnectedTrackChange.js
@@ -75,7 +75,6 @@ export default ({ trackChangeId, top, recalculateTops, trackChange }) => {
     recalculateTops();
     if (activeTrackChange && trackChangeId === activeTrackChange.attrs.id) {
       setIsActive(true);
-      recalculateTops();
     }
   }, [activeTrackChange]);
 
diff --git a/wax-prosemirror-services/src/CommentsService/components/RightArea.js b/wax-prosemirror-services/src/CommentsService/components/RightArea.js
index d3644d9b52cacec690047c71d8c5ce807471df55..d2e7869b80b31667f112e13604b117189c8c8150 100644
--- a/wax-prosemirror-services/src/CommentsService/components/RightArea.js
+++ b/wax-prosemirror-services/src/CommentsService/components/RightArea.js
@@ -1,18 +1,21 @@
+/* eslint-disable no-param-reassign */
 /* eslint react/prop-types: 0 */
-import { Mark } from 'prosemirror-model';
 import React, { useContext, useState, useMemo, useCallback } from 'react';
 import useDeepCompareEffect from 'use-deep-compare-effect';
-import { each, uniqBy, sortBy } from 'lodash';
+import { each, uniqBy, sortBy, groupBy } from 'lodash';
 import { WaxContext, DocumentHelpers } from 'wax-prosemirror-core';
 import BoxList from './BoxList';
+import { CommentDecorationPluginKey } from '../plugins/CommentDecorationPlugin';
 
 export default ({ area, users }) => {
+  const context = useContext(WaxContext);
   const {
     pmViews,
     pmViews: { main },
     app,
     activeView,
-  } = useContext(WaxContext);
+    options: { comments, commentsMap },
+  } = context;
 
   const commentPlugin = app.PmPlugins.get('commentPlugin');
   const trakChangePlugin = app.PmPlugins.get('trackChangePlugin');
@@ -39,17 +42,26 @@ export default ({ area, users }) => {
     }
 
     each(marksNodes[area], (markNode, pos) => {
-      const id =
-        markNode instanceof Mark ? markNode.attrs.id : markNode.node.attrs.id;
+      let id = '';
+
+      if (markNode?.node?.attrs.id) {
+        id = markNode.node.attrs.id;
+      } else if (markNode?.attrs?.id) {
+        id = markNode.attrs.id;
+      } else {
+        id = markNode.id;
+      }
+
       let activeTrackChange = null;
       const activeComment = commentPlugin.getState(activeView.state).comment;
+
       if (trakChangePlugin)
         activeTrackChange = trakChangePlugin.getState(activeView.state)
           .trackChange;
 
       let isActive = false;
       if (
-        (activeComment && id === activeComment.attrs.id) ||
+        (activeComment && id === activeComment.id) ||
         (activeTrackChange && id === activeTrackChange.attrs.id)
       )
         isActive = true;
@@ -57,11 +69,45 @@ export default ({ area, users }) => {
       // annotation top
       if (area === 'main') {
         markNodeEl = document.querySelector(`[data-id="${id}"]`);
-        if (markNodeEl)
+        if (!markNodeEl && marksNodes[area][pos - 1]) {
+          markNodeEl = document.querySelector(
+            `[data-id="${marksNodes[area][pos - 1].id}"]`,
+          );
+        }
+
+        if (markNodeEl) {
           annotationTop =
             markNodeEl.getBoundingClientRect().top -
             WaxSurface.top +
             parseInt(WaxSurfaceMarginTop.slice(0, -2), 10);
+        } else if (!isFirstRun) {
+          // comment is deleted from editing surface
+          context.setOption({ resolvedComment: id });
+          context.setOption({
+            comments: comments.filter(comment => {
+              return comment.id !== id;
+            }),
+          });
+          setTimeout(() => {
+            activeView.dispatch(
+              activeView.state.tr.setMeta(CommentDecorationPluginKey, {
+                type: 'deleteComment',
+                id,
+              }),
+            );
+            if (context.app.config.get('config.YjsService')) {
+              commentsMap.observe(() => {
+                const transaction = context.pmViews.main.state.tr.setMeta(
+                  CommentDecorationPluginKey,
+                  {
+                    type: 'createDecorations',
+                  },
+                );
+                context.pmViews.main.dispatch(transaction);
+              });
+            }
+          });
+        }
       } else {
         // Notes
         panelWrapper = document.getElementsByClassName('panelWrapper');
@@ -127,11 +173,15 @@ export default ({ area, users }) => {
           if (doesOverlap) {
             const overlap = boxAboveEnds - currentTop;
             result[i - 1] -= overlap;
+            let previousMarkNode = '';
 
-            const previousMarkNode =
-              marksNodes[area][i - 1] instanceof Mark
-                ? marksNodes[area][i - 1].attrs.id
-                : marksNodes[area][i - 1].node.attrs.id;
+            if (marksNodes[area][i - 1]?.node?.attrs.id) {
+              previousMarkNode = marksNodes[area][i - 1].node.attrs.id;
+            } else if (marksNodes[area][i - 1]?.attrs?.id) {
+              previousMarkNode = marksNodes[area][i - 1].attrs.id;
+            } else {
+              previousMarkNode = marksNodes[area][i - 1].id;
+            }
 
             allCommentsTop[i - 1][previousMarkNode] = result[i - 1];
           }
@@ -152,7 +202,7 @@ export default ({ area, users }) => {
   };
 
   useDeepCompareEffect(() => {
-    setMarksNodes(updateMarks(pmViews));
+    setMarksNodes(updateMarks(pmViews, comments));
     if (isFirstRun) {
       setTimeout(() => {
         setPosition(setTops());
@@ -161,7 +211,7 @@ export default ({ area, users }) => {
     } else {
       setPosition(setTops());
     }
-  }, [updateMarks(pmViews), setTops()]);
+  }, [updateMarks(pmViews, comments), setTops()]);
 
   const CommentTrackComponent = useMemo(
     () => (
@@ -179,7 +229,9 @@ export default ({ area, users }) => {
   return <>{CommentTrackComponent}</>;
 };
 
-const updateMarks = views => {
+const updateMarks = (views, comments) => {
+  const newComments = groupBy(comments, comm => comm.data.group) || [];
+
   if (views.main) {
     const allInlineNodes = [];
 
@@ -197,12 +249,11 @@ const updateMarks = views => {
       if (node.node.marks.length > 0) {
         node.node.marks.filter(mark => {
           if (
-            mark.type.name === 'comment' ||
             mark.type.name === 'insertion' ||
             mark.type.name === 'deletion' ||
             mark.type.name === 'format_change'
           ) {
-            mark.pos = node.pos;
+            mark.from = node.pos;
             finalMarks.push(mark);
           }
         });
@@ -217,7 +268,7 @@ const updateMarks = views => {
 
     const nodesAndMarks = [...uniqBy(finalMarks, 'attrs.id'), ...finalNodes];
 
-    const groupedMarkNodes = {};
+    const groupedMarkNodes = { main: [], notes: [] };
     nodesAndMarks.forEach(markNode => {
       const markNodeAttrs = markNode.attrs
         ? markNode.attrs
@@ -229,9 +280,13 @@ const updateMarks = views => {
         groupedMarkNodes[markNodeAttrs.group].push(markNode);
       }
     });
+    if (newComments?.main?.length > 0)
+      groupedMarkNodes.main = groupedMarkNodes.main.concat(newComments.main);
+    if (newComments?.notes?.length > 0)
+      groupedMarkNodes.notes = groupedMarkNodes.notes.concat(newComments.notes);
 
     return {
-      main: sortBy(groupedMarkNodes.main, ['pos']),
+      main: sortBy(groupedMarkNodes.main, ['data.pmFrom']),
       notes: groupedMarkNodes.notes,
     };
   }
diff --git a/wax-prosemirror-services/src/CommentsService/components/ui/comments/CommentBubbleComponent.js b/wax-prosemirror-services/src/CommentsService/components/ui/comments/CommentBubbleComponent.js
index 2a5ccec585acb0bd6d497d0947bab56902dbe504..876726784da558218389ee73489120ff3cb48653 100644
--- a/wax-prosemirror-services/src/CommentsService/components/ui/comments/CommentBubbleComponent.js
+++ b/wax-prosemirror-services/src/CommentsService/components/ui/comments/CommentBubbleComponent.js
@@ -1,10 +1,22 @@
 /* eslint react/prop-types: 0 */
 import React, { useLayoutEffect, useContext } from 'react';
-import { WaxContext, Commands, DocumentHelpers } from 'wax-prosemirror-core';
+import { WaxContext } from 'wax-prosemirror-core';
+import {
+  ySyncPluginKey,
+  relativePositionToAbsolutePosition,
+  absolutePositionToRelativePosition,
+} from 'y-prosemirror';
 import CommentBubble from './CommentBubble';
+import { CommentDecorationPluginKey } from '../../../plugins/CommentDecorationPlugin';
 
 const CommentBubbleComponent = ({ setPosition, position, group }) => {
-  const { activeView, activeViewId } = useContext(WaxContext);
+  const context = useContext(WaxContext);
+  const {
+    activeView,
+    activeViewId,
+    options: { comments, commentsMap },
+  } = context;
+
   const { state, dispatch } = activeView;
 
   useLayoutEffect(() => {
@@ -21,14 +33,84 @@ const CommentBubbleComponent = ({ setPosition, position, group }) => {
 
   const createComment = event => {
     event.preventDefault();
-    Commands.createComment(state, dispatch, group, activeViewId);
-    activeView.focus();
+    const { selection } = state;
+
+    if (context.app.config.get('config.YjsService')) {
+      return createYjsComments(selection);
+    }
+    dispatch(
+      state.tr.setMeta(CommentDecorationPluginKey, {
+        type: 'addComment',
+        from: selection.from,
+        to: selection.to,
+        yjsFrom: selection.from,
+        yjsTo: selection.to,
+        pmFrom: selection.from,
+        pmTo: selection.to,
+        data: {
+          type: 'comment',
+          conversation: [],
+          title: '',
+          group,
+          viewId: activeViewId,
+        },
+      }),
+    );
+    dispatch(state.tr);
   };
 
-  const isCommentAllowed = () => {
-    const commentMark = activeView.state.schema.marks.comment;
-    const marks = DocumentHelpers.findMark(state, commentMark, true);
+  const createYjsComments = selection => {
+    const ystate = ySyncPluginKey.getState(state);
+    const { doc, type, binding } = ystate;
+    const from = absolutePositionToRelativePosition(
+      selection.from,
+      type,
+      binding.mapping,
+    );
+    const to = absolutePositionToRelativePosition(
+      selection.to,
+      type,
+      binding.mapping,
+    );
+
+    commentsMap.observe(() => {
+      const transaction = context.pmViews.main.state.tr.setMeta(
+        CommentDecorationPluginKey,
+        {
+          type: 'createDecorations',
+        },
+      );
+      context.pmViews.main.dispatch(transaction);
+    });
 
+    dispatch(
+      state.tr.setMeta(CommentDecorationPluginKey, {
+        type: 'addComment',
+        from: relativePositionToAbsolutePosition(
+          doc,
+          type,
+          from,
+          binding.mapping,
+        ),
+        to: relativePositionToAbsolutePosition(doc, type, to, binding.mapping),
+        data: {
+          yjsFrom: selection.from,
+          yjsTo: selection.to,
+          pmFrom: selection.from,
+          pmTo: selection.to,
+          type: 'comment',
+          conversation: [],
+          title: '',
+          group,
+          viewId: activeViewId,
+        },
+      }),
+    );
+
+    dispatch(state.tr);
+  };
+
+  const isCommentAllowed = () => {
     let allowed = true;
     state.doc.nodesBetween(
       state.selection.$from.pos,
@@ -43,13 +125,17 @@ const CommentBubbleComponent = ({ setPosition, position, group }) => {
         }
       },
     );
-    // TODO Overlapping comments . for now don't allow
-    marks.forEach(mark => {
-      if (mark.attrs.group === 'main') allowed = false;
-    });
 
-    // TO DO this is because of a bug and overlay doesn't rerender. Fix in properly in Notes, and remove
-    if (activeViewId !== 'main' && marks.length >= 1) allowed = false;
+    if (
+      comments.find(
+        comm =>
+          comm.data.pmFrom === state.selection.from &&
+          comm.data.pmTo === state.selection.to,
+      )
+    ) {
+      allowed = false;
+    }
+
     return allowed;
   };
 
diff --git a/wax-prosemirror-services/src/CommentsService/plugins/CommentDecoration.js b/wax-prosemirror-services/src/CommentsService/plugins/CommentDecoration.js
new file mode 100644
index 0000000000000000000000000000000000000000..caf61271e4ddc2ca74218cea8281ab298709633e
--- /dev/null
+++ b/wax-prosemirror-services/src/CommentsService/plugins/CommentDecoration.js
@@ -0,0 +1,47 @@
+export default class CommentDecoration {
+  constructor(decoration) {
+    this.decoration = decoration;
+  }
+
+  get displayName() {
+    return this.decoration.type.spec.data.displayName;
+  }
+
+  get tag() {
+    return this.decoration.type.spec.data.tag;
+  }
+
+  get id() {
+    return this.decoration.type.spec.id;
+  }
+
+  get from() {
+    return this.decoration.from;
+  }
+
+  get to() {
+    return this.decoration.to;
+  }
+
+  get selectedText() {
+    return this.decoration.type.spec.data.selectedText;
+  }
+
+  get data() {
+    return this.decoration.type.spec.data;
+  }
+
+  get HTMLAttributes() {
+    return this.decoration.type.attrs;
+  }
+
+  toString() {
+    return JSON.stringify({
+      id: this.id,
+      data: this.data,
+      from: this.from,
+      to: this.to,
+      HTMLAttributes: this.HTMLAttributes,
+    });
+  }
+}
diff --git a/wax-prosemirror-services/src/CommentsService/plugins/CommentDecorationPlugin.js b/wax-prosemirror-services/src/CommentsService/plugins/CommentDecorationPlugin.js
new file mode 100644
index 0000000000000000000000000000000000000000..96cc05e1519cd5257135d31ef472db504e22919e
--- /dev/null
+++ b/wax-prosemirror-services/src/CommentsService/plugins/CommentDecorationPlugin.js
@@ -0,0 +1,46 @@
+import { Plugin, PluginKey } from 'prosemirror-state';
+import CommentState from './CommentState';
+
+let contentSize = 0;
+let allCommentsCount = 0;
+
+export const CommentDecorationPluginKey = new PluginKey(
+  'commentDecorationPlugin',
+);
+export const CommentDecorationPlugin = (name, options) => {
+  return new Plugin({
+    key: CommentDecorationPluginKey,
+    state: {
+      init() {
+        return new CommentState({
+          context: options.context,
+          map: options.existingComments(),
+          onSelectionChange: options.onSelectionChange,
+        });
+      },
+      apply(transaction, pluginState, oldState, newState) {
+        return pluginState.apply(transaction, newState);
+      },
+    },
+    props: {
+      decorations(state) {
+        const { decorations } = this.getState(state);
+        if (
+          contentSize !== state.doc.content.size ||
+          this.getState(state).allCommentsList().length !== allCommentsCount
+        ) {
+          // const annotations = this.getState(state).commentsAt(
+          //   0,
+          //   state.doc.content.size,
+          // );
+          // options.onSelectionChange(annotations);
+
+          options.onSelectionChange(this.getState(state).allCommentsList());
+        }
+        contentSize = state.doc.content.size;
+        allCommentsCount = this.getState(state).allCommentsList().length;
+        return decorations;
+      },
+    },
+  });
+};
diff --git a/wax-prosemirror-services/src/CommentsService/plugins/CommentPlugin.js b/wax-prosemirror-services/src/CommentsService/plugins/CommentPlugin.js
index 3989ab3591e8f44935012e11dc35d2e8daa065e1..7a6c99c92b7d929c5a0f6daccefea33dc95c8180 100644
--- a/wax-prosemirror-services/src/CommentsService/plugins/CommentPlugin.js
+++ b/wax-prosemirror-services/src/CommentsService/plugins/CommentPlugin.js
@@ -1,83 +1,48 @@
-/* eslint-disable */
-
-import { minBy, maxBy, last } from 'lodash';
+/* eslint-disable consistent-return */
+import { inRange, last, sortBy } from 'lodash';
 import { Plugin, PluginKey } from 'prosemirror-state';
 import { Decoration, DecorationSet } from 'prosemirror-view';
-import { DocumentHelpers } from 'wax-prosemirror-core';
 
 const commentPlugin = new PluginKey('commentPlugin');
 
-const getComment = state => {
-  const commentMark = state.schema.marks.comment;
-  const commentOnSelection = DocumentHelpers.findFragmentedMark(
-    state,
-    commentMark,
-  );
-
-  // Don't allow Active comment if selection is not collapsed
-  if (
-    state.selection.from !== state.selection.to &&
-    commentOnSelection &&
-    commentOnSelection.attrs.conversation.length
-  ) {
-    return;
-  }
-
-  if (commentOnSelection) {
-    const commentNodes = DocumentHelpers.findChildrenByMark(
-      state.doc,
-      commentMark,
-      true,
-    );
+const getComment = (state, context) => {
+  const {
+    options: { comments },
+  } = context;
+  if (!comments?.length) return;
 
-    const allCommentsWithSameId = [];
-    commentNodes.map(node => {
-      node.node.marks.filter(mark => {
-        if (
-          mark.type.name === 'comment' &&
-          commentOnSelection.attrs.id === mark.attrs.id
-        ) {
-          allCommentsWithSameId.push(node);
-        }
-      });
-    });
-
-    const minPos = minBy(allCommentsWithSameId, 'pos');
-    const maxPos = maxBy(allCommentsWithSameId, 'pos');
+  let commentData = comments.filter(comment =>
+    inRange(state.selection.from, comment.data.pmFrom, comment.data.pmTo),
+  );
 
+  commentData = sortBy(commentData, ['data.pmFrom']);
+  if (commentData.length > 0) {
     if (
-      state.selection.from ===
-      maxPos.pos + last(allCommentsWithSameId).node.nodeSize
+      (state.selection.from !== state.selection.to &&
+        last(commentData).data.conversation.length === 0) ||
+      (state.selection.from === state.selection.to &&
+        last(commentData).data.conversation.length !== 0)
     ) {
-      state.schema.marks.comment.spec.inclusive = false;
-    } else {
-      state.schema.marks.comment.spec.inclusive = true;
-    }
-    if (allCommentsWithSameId.length > 1) {
-      return {
-        from: minPos.pos,
-        to: maxPos.pos + last(allCommentsWithSameId).node.nodeSize,
-        attrs: commentOnSelection.attrs,
-        contained: commentOnSelection.contained,
-      };
+      return last(commentData);
     }
+    return undefined;
   }
-  return commentOnSelection;
+  return undefined;
 };
 
-export default props => {
+export default (key, context) => {
   return new Plugin({
     key: commentPlugin,
     state: {
       init: (_, state) => {
-        return { comment: getComment(state) };
+        return { comment: getComment(state, context) };
       },
       apply(tr, prev, _, newState) {
-        const comment = getComment(newState);
+        const comment = getComment(newState, context);
         let createDecoration;
         if (comment) {
           createDecoration = DecorationSet.create(newState.doc, [
-            Decoration.inline(comment.from, comment.to, {
+            Decoration.inline(comment.data.pmFrom, comment.data.pmTo, {
               class: 'active-comment',
             }),
           ]);
diff --git a/wax-prosemirror-services/src/CommentsService/plugins/CommentState.js b/wax-prosemirror-services/src/CommentsService/plugins/CommentState.js
new file mode 100644
index 0000000000000000000000000000000000000000..a386e6d70171d1c589a46dec50872ff595eec996
--- /dev/null
+++ b/wax-prosemirror-services/src/CommentsService/plugins/CommentState.js
@@ -0,0 +1,221 @@
+/* eslint-disable no-param-reassign */
+import { v4 as uuidv4 } from 'uuid';
+import { Decoration, DecorationSet } from 'prosemirror-view';
+import {
+  ySyncPluginKey,
+  relativePositionToAbsolutePosition,
+  absolutePositionToRelativePosition,
+} from 'y-prosemirror';
+import CommentDecoration from './CommentDecoration';
+import { CommentDecorationPluginKey } from './CommentDecorationPlugin';
+
+const randomId = () => {
+  return uuidv4();
+};
+
+export default class CommentState {
+  constructor(options) {
+    this.decorations = DecorationSet.empty;
+    this.options = options;
+  }
+
+  addComment(action) {
+    const { map } = this.options;
+    const { from, to, data } = action;
+    const id = randomId();
+    map.set(id, { id, from, to, data });
+  }
+
+  updateComment(action) {
+    const { map } = this.options;
+    const annotationToUpdate = map.get(action.id);
+    if (annotationToUpdate) {
+      annotationToUpdate.data = action.data;
+    }
+  }
+
+  deleteComment(id) {
+    const { map } = this.options;
+    map.delete(id);
+  }
+
+  commentsAt(position, to) {
+    return this.decorations.find(position, to || position).map(decoration => {
+      return new CommentDecoration(decoration);
+    });
+  }
+
+  allCommentsList() {
+    const { map } = this.options;
+    return Array.from(map, ([key, value]) => {
+      return { ...value, id: key };
+    }).filter(value => {
+      return 'from' in value && 'to' in value;
+    });
+  }
+
+  createDecorations(state) {
+    const decorations = [];
+
+    const ystate = ySyncPluginKey.getState(state);
+
+    if (ystate?.binding) {
+      const { doc, type, binding } = ystate;
+      console.log(this.allCommentsList());
+      this.allCommentsList().forEach((annotation, id) => {
+        annotation.data.yjsFrom = absolutePositionToRelativePosition(
+          annotation.data.pmFrom,
+          type,
+          binding.mapping,
+        );
+
+        annotation.data.yjsTo = absolutePositionToRelativePosition(
+          annotation.data.pmTo,
+          type,
+          binding.mapping,
+        );
+
+        const from = relativePositionToAbsolutePosition(
+          doc,
+          type,
+          annotation.data.yjsFrom,
+          binding.mapping,
+        );
+        const to = relativePositionToAbsolutePosition(
+          doc,
+          type,
+          annotation.data.yjsTo,
+          binding.mapping,
+        );
+
+        if (!from || !to) {
+          return;
+        }
+
+        decorations.push(
+          Decoration.inline(
+            from,
+            to,
+            {
+              class: 'comment',
+              'data-id': annotation.id,
+            },
+            {
+              id: annotation.id,
+              data: annotation,
+              inclusiveEnd: true,
+            },
+          ),
+        );
+      });
+    } else {
+      this.allCommentsList().forEach(annotation => {
+        const {
+          data: { pmFrom, pmTo },
+        } = annotation;
+
+        decorations.push(
+          Decoration.inline(
+            pmFrom,
+            pmTo,
+            {
+              class: 'comment',
+              'data-id': annotation.id,
+            },
+            {
+              id: annotation.id,
+              data: annotation,
+              inclusiveEnd: true,
+            },
+          ),
+        );
+      });
+    }
+
+    this.decorations = DecorationSet.create(state.doc, decorations);
+  }
+
+  updateCommentPostions(ystate) {
+    this.options.map.doc.transact(() => {
+      this.decorations.find().forEach(deco => {
+        const { id } = deco.spec;
+        const newFrom = absolutePositionToRelativePosition(
+          deco.from,
+          ystate.type,
+          ystate.binding.mapping,
+        );
+        const newTo = absolutePositionToRelativePosition(
+          deco.to,
+          ystate.type,
+          ystate.binding.mapping,
+        );
+
+        const annotation = this.options.map.get(id);
+
+        annotation.from = newFrom;
+        annotation.to = newTo;
+
+        this.options.map.set(id, annotation);
+      });
+    }, CommentDecorationPluginKey);
+  }
+
+  apply(transaction, state) {
+    const { map } = this.options;
+    const action = transaction.getMeta(CommentDecorationPluginKey);
+    if (action && action.type) {
+      if (action.type === 'addComment') {
+        this.addComment(action);
+      }
+      if (action.type === 'updateComment') {
+        this.updateComment(action);
+      }
+      if (action.type === 'deleteComment') {
+        this.deleteComment(action.id);
+      }
+      if (action.type === 'createDecorations') {
+        this.createDecorations(state);
+      }
+      // this.createDecorations(state);
+      return this;
+    }
+
+    const ystate = ySyncPluginKey.getState(state);
+
+    if (ystate?.isChangeOrigin) {
+      // this.updateCommentPostions(ystate);
+      this.createDecorations(state);
+
+      return this;
+    }
+
+    this.decorations = this.decorations.map(
+      transaction.mapping,
+      transaction.doc,
+    );
+
+    map.forEach((annotation, _) => {
+      if ('from' in annotation && 'to' in annotation) {
+        annotation.data.pmFrom = transaction.mapping.map(
+          annotation.data.pmFrom,
+        );
+        annotation.data.pmTo = transaction.mapping.map(annotation.data.pmTo);
+      }
+    });
+
+    if (ystate?.binding && ystate?.binding.mapping) {
+      this.updateCommentPostions(ystate);
+      return this;
+      // eslint-disable-next-line no-else-return
+    } else {
+      this.options.map.forEach((annotation, _) => {
+        if ('from' in annotation && 'to' in annotation) {
+          annotation.from = transaction.mapping.map(annotation.from);
+          annotation.to = transaction.mapping.map(annotation.to);
+        }
+      });
+      this.createDecorations(state);
+      return this;
+    }
+  }
+}
diff --git a/wax-prosemirror-services/src/YjsService/YjsService.js b/wax-prosemirror-services/src/YjsService/YjsService.js
index 72ed63d39b14236ee98cbda06063a3a33f0ca714..8efdbc131a15fcdd45831dcb2b91bef3754d1348 100644
--- a/wax-prosemirror-services/src/YjsService/YjsService.js
+++ b/wax-prosemirror-services/src/YjsService/YjsService.js
@@ -22,6 +22,7 @@ class YjsService extends Service {
     if (!configProvider || !configYdoc) {
       ydoc = new Y.Doc();
       provider = new WebsocketProvider(connectionUrl, docIdentifier, ydoc);
+      this.app.context.setOption({ currentYdoc: ydoc });
     }
 
     provider.on('sync', args => {