diff --git a/wax-prosemirror-components/src/components/comments/CommentBox.js b/wax-prosemirror-components/src/components/comments/CommentBox.js
index c394e7a9dbdf79fd80674217393c531d3ea36d6f..d9497b01ba72ab859facd4ab928688a173ba306f 100644
--- a/wax-prosemirror-components/src/components/comments/CommentBox.js
+++ b/wax-prosemirror-components/src/components/comments/CommentBox.js
@@ -28,13 +28,15 @@ const CommentBoxStyled = styled.div`
   }};
 `;
 
-export default ({ mark, view, top, dataComment }) => {
+export default ({ mark, view, top, dataBox }) => {
   const [animate, setAnimate] = useState(false);
   const { view: { main }, app, activeView } = useContext(WaxContext);
   const activeCommentPlugin = app.PmPlugins.get("activeComment");
   const activeComment = activeCommentPlugin.getState(activeView.state).comment;
   let active = false;
-  if (activeComment && mark.attrs.id === activeComment.attrs.id) active = true;
+  const id = mark.attrs ? mark.attrs.id : mark.node.attrs.id;
+
+  if (activeComment && id === activeComment.attrs.id) active = true;
   useEffect(() => {
     setAnimate(true);
   }, []);
@@ -46,7 +48,7 @@ export default ({ mark, view, top, dataComment }) => {
           <CommentBoxStyled
             top={top}
             state={state}
-            data-comment={dataComment}
+            data-box={dataBox}
             active={active}
           />
         )}
diff --git a/wax-prosemirror-components/src/components/comments/CommentsBoxList.js b/wax-prosemirror-components/src/components/comments/CommentsBoxList.js
index 0722446cf0a77567fb5f83a9ea45f235eeb37a42..84514097aa2d991dba242c0b306c450e95cd29f4 100644
--- a/wax-prosemirror-components/src/components/comments/CommentsBoxList.js
+++ b/wax-prosemirror-components/src/components/comments/CommentsBoxList.js
@@ -5,7 +5,7 @@ export default ({ comments, view, position }) => {
   return (
     <Fragment>
       {comments.map((comment, index) => {
-        const { attrs: { id } } = comment;
+        const id = comment.attrs ? comment.attrs.id : comment.node.attrs.id;
         const top = position[index] ? position[index][id] : 0;
         return (
           <CommentBox
@@ -13,7 +13,7 @@ export default ({ comments, view, position }) => {
             mark={comment}
             view={view}
             top={top}
-            dataComment={`comment-${id}`}
+            dataBox={id}
           />
         );
       })}
diff --git a/wax-prosemirror-components/src/components/rightArea/RightArea.js b/wax-prosemirror-components/src/components/rightArea/RightArea.js
index ea583a8def9c69a5e192f85c34454b57fb78c9fa..9c2990918d76f3ed4edfccfa6cc991797652d6fb 100644
--- a/wax-prosemirror-components/src/components/rightArea/RightArea.js
+++ b/wax-prosemirror-components/src/components/rightArea/RightArea.js
@@ -28,17 +28,17 @@ export default ({ area }) => {
 
     each(marks[area], (mark, pos) => {
       const WaxSurface = main.dom.getBoundingClientRect();
-      const { attrs: { id } } = mark;
+      const id = mark.attrs ? mark.attrs.id : mark.node.attrs.id;
+
       const activeComment = activeCommentPlugin.getState(activeView.state)
         .comment;
 
       let isActive = false;
-      if (activeComment && mark.attrs.id === activeComment.attrs.id)
-        isActive = true;
+      if (activeComment && id === activeComment.attrs.id) isActive = true;
 
       //annotation top
       if (area === "main") {
-        markEl = document.querySelector(`span[data-id="${id}"]`);
+        markEl = document.querySelector(`[data-id="${id}"]`);
         if (markEl)
           annotationTop = markEl.getBoundingClientRect().top - WaxSurface.top;
       } else {
@@ -47,14 +47,14 @@ export default ({ area }) => {
           .height;
         markEl = document
           .querySelector("#notes-container")
-          .querySelector(`span[data-id="${id}"]`);
+          .querySelector(`[data-id="${id}"]`);
         if (markEl)
           annotationTop =
             markEl.getBoundingClientRect().top - panelWrapperHeight - 50;
       }
 
       // get height of this mark box
-      const boxEl = document.querySelector(`div[data-comment="comment-${id}"]`);
+      const boxEl = document.querySelector(`div[data-box="${id}"]`);
       if (boxEl) boxHeight = parseInt(boxEl.offsetHeight);
 
       // where the box should move to
@@ -95,7 +95,11 @@ export default ({ area }) => {
             const overlap = boxAboveEnds - currentTop;
             result[i - 1] -= overlap;
 
-            allCommentsTop[i - 1][marks[area][i - 1].attrs.id] = result[i - 1];
+            const temp = marks[area][i - 1].attrs
+              ? marks[area][i - 1].attrs.id
+              : marks[area][i - 1].node.attrs.id;
+
+            allCommentsTop[i - 1][temp] = result[i - 1];
           }
 
           if (!doesOverlap) b = false;