diff --git a/editors/demo/src/OEN/OEN.js b/editors/demo/src/OEN/OEN.js
index 29b6c46c7aead79e89f1f25a9061220787e0c772..667960f6166f4f593cfe6e8502674f4c2f624b87 100644
--- a/editors/demo/src/OEN/OEN.js
+++ b/editors/demo/src/OEN/OEN.js
@@ -44,9 +44,9 @@ const Oen = () => {
       value={val}
       // readonly
       layout={OenLayout}
-      // onChange={debounce(source => {
-      //   console.log(JSON.stringify(source));
-      // }, 200)}
+      onChange={debounce(source => {
+        console.log(JSON.stringify(source));
+      }, 200)}
       user={user}
     />
   );
diff --git a/wax-prosemirror-components/src/components/OEN/OENToolGroup.js b/wax-prosemirror-components/src/components/OEN/OENToolGroup.js
index 2e385b1ae64d0d99eca0351ae070b78f932138e0..fba1e2316912ab3567ee871b49fba46be288b3ed 100644
--- a/wax-prosemirror-components/src/components/OEN/OENToolGroup.js
+++ b/wax-prosemirror-components/src/components/OEN/OENToolGroup.js
@@ -73,9 +73,9 @@ const OENToolGroup = ({ item }) => {
                         activeView.dispatch,
                         tool.className,
                       );
-                      setTimeout(() => {
-                        main.focus();
-                      });
+                      // setTimeout(() => {
+                      //   main.focus();
+                      // });
                     }}
                     title={tool.displayName}
                   />
diff --git a/wax-prosemirror-services/src/OENContainersService/OENContainersTool.js b/wax-prosemirror-services/src/OENContainersService/OENContainersTool.js
index 14481bde6dab158cde79f6dc700a8bd191d63bac..07165943eb4d77e5f0b7ae5bf7711b332993d4a9 100644
--- a/wax-prosemirror-services/src/OENContainersService/OENContainersTool.js
+++ b/wax-prosemirror-services/src/OENContainersService/OENContainersTool.js
@@ -1,6 +1,7 @@
 import { injectable } from 'inversify';
 import { wrapIn } from 'prosemirror-commands';
-import { Commands } from 'wax-prosemirror-utilities';
+import { NodeSelection } from 'prosemirror-state';
+import { Mapping, RemoveMarkStep, ReplaceStep } from 'prosemirror-transform';
 import Tools from '../lib/Tools';
 
 @injectable()
@@ -11,27 +12,50 @@ export default class OENContainersTool extends Tools {
 
   get run() {
     return (state, dispatch, className) => {
-      console.log(state.selection);
-      const node = className === 'section' ? 'oen_section' : 'oen_container';
+      const { from, to } = state.selection;
+      let isInOenContainer = false;
+      let OENContainer = '';
+      let position = 0;
+      state.doc.nodesBetween(from, to, (node, pos) => {
+        if (node.type.name === 'oen_container') {
+          isInOenContainer = true;
+          OENContainer = node;
+          position = pos;
+          console.log(pos);
+        }
+      });
 
-      wrapIn(state.config.schema.nodes[node], { class: className })(
-        state,
-        dispatch,
-      );
+      if (isInOenContainer) {
+        const map = new Mapping();
+        const newNode = JSON.parse(JSON.stringify(OENContainer));
+        OENContainer.attrs.class = className;
+        console.log('replace', OENContainer);
+        newNode.attrs = {
+          ...newNode.attrs,
+          class: className,
+        };
+        console.log(newNode);
+        state.tr.setSelection(NodeSelection.create(state.doc, position));
+        state.tr.replaceSelectionWith(OENContainer);
+        // state.tr.setNodeMarkup(map.map(position), null, {
+        //   ...OENContainer.attrs,
+        //   class: className,
+        // });
+        dispatch(state.tr);
+      } else {
+        const node = className === 'section' ? 'oen_section' : 'oen_container';
 
-      console.log(className);
+        wrapIn(state.config.schema.nodes[node], { class: className })(
+          state,
+          dispatch,
+        );
+      }
     };
   }
 
   select = (state, activeViewId) => {
-    const { from, to } = state.selection;
-    let status = true;
-    state.doc.nodesBetween(from, to, (node, pos) => {
-      if (node.type.name === 'oen_container') {
-        status = false;
-      }
-    });
-    return status;
+    if (activeViewId !== 'main') return false;
+    return true;
   };
 
   get enable() {
diff --git a/wax-prosemirror-services/src/OENContainersService/oenContainers.css b/wax-prosemirror-services/src/OENContainersService/oenContainers.css
index 726b23d811e1aeae7ba1cf8ff08a2f27c27d0251..6836f45c695b2663e5c9ab2839d8484cb53d320a 100644
--- a/wax-prosemirror-services/src/OENContainersService/oenContainers.css
+++ b/wax-prosemirror-services/src/OENContainersService/oenContainers.css
@@ -11,14 +11,14 @@ div[data-type="content_structure_element"]::before {
     display: block;
 }
 
-.section {
+section {
     position: relative;
     padding: 10px 10px 0 10px;
     margin-top: 20px;
     border: 3px solid #E2EBFF;
 }
 
-.section::before {
+section::before {
     content: 'Section';
     color: #c7d8fc;
     font-weight: bold;
@@ -29,11 +29,11 @@ div[data-type="content_structure_element"]::before {
     left: -3px;
 }
 
-.ProseMirror>section.section:nth-of-type(1) {
+.ProseMirror>section:nth-of-type(1) {
   margin-top: 0px;
 }
 
-.section>div:nth-of-type(1) {
+section>div:nth-of-type(1) {
     padding-top: 0;
 }
 
@@ -73,10 +73,6 @@ content: 'Learning Objectives Section';
     content: 'References Section';
 }
 
-.main-content::before {
-    content: 'Main Content Section';
-}
-
 .bibliography::before {
     content: 'Bibliography Section';
 }