diff --git a/stories/Undo.stories.tsx b/stories/Undo.stories.tsx
index 4ec547c1..39c95ad5 100644
--- a/stories/Undo.stories.tsx
+++ b/stories/Undo.stories.tsx
@@ -47,12 +47,9 @@ export default {
   }
 } as Meta;
 
-type PropsWithChildrenMock = {
-  initialHistory?: { nodes: NodeData[]; edges: EdgeData[] }[];
-};
+type Props = {}
 
-const Template: Story<PropsWithChildrenMock> = (props) => {
-  const { initialHistory } = props;
+export const Simple: Story<Props> = () => {
   const [nodes, setNodes] = useState<any[]>([
     {
       id: '1',
@@ -84,7 +81,6 @@ const Template: Story<PropsWithChildrenMock> = (props) => {
   const { undo, redo, canUndo, canRedo } = useUndo({
     nodes,
     edges,
-    initialHistory,
     onUndoRedo: (state: UndoRedoEvent) => {
       console.log('Undo / Redo', state);
       setEdges(state.edges);
@@ -133,12 +129,8 @@ const Template: Story<PropsWithChildrenMock> = (props) => {
   );
 };
 
-export const Simple: Story<PropsWithChildrenMock> = Template.bind({});
-Simple.args = {};
-
-export const WithInitialHistory: Story<PropsWithChildrenMock> = Template.bind({});
-WithInitialHistory.args = {
-  initialHistory: [
+export const WithInitialHistory: Story<Props> = () => {
+  const initialHistory: { nodes: NodeData[]; edges: EdgeData[] }[] = [
     {
       nodes: [],
       edges: []
@@ -199,6 +191,83 @@ WithInitialHistory.args = {
         }
       ]
     }
-  ]
-};
+  ];
+  const [nodes, setNodes] = useState<any[]>([
+    {
+      id: '1',
+      text: 'Node 1'
+    },
+    {
+      id: '2',
+      text: 'Node 2'
+    },
+    {
+      id: '3',
+      text: 'Node 3'
+    }
+  ]);
+
+  const [edges, setEdges] = useState<any[]>([
+    {
+      id: '1-2',
+      from: '1',
+      to: '2'
+    },
+    {
+      id: '1-3',
+      from: '1',
+      to: '3'
+    }
+  ]);
+
+  const { undo, redo, canUndo, canRedo } = useUndo({
+    nodes,
+    edges,
+    initialHistory,
+    onUndoRedo: (state: UndoRedoEvent) => {
+      console.log('Undo / Redo', state);
+      setEdges(state.edges);
+      setNodes(state.nodes);
+    }
+  });
+
+  const addNode = () => {
+    setNodes([
+      ...nodes,
+      {
+        id: `a${Math.random()}`,
+        text: `Node ${Math.random()}`
+      }
+    ]);
+  };
 
+  return (
+    <div style={{ position: 'absolute', top: 0, bottom: 0, left: 0, right: 0 }}>
+      <button
+        style={{ position: 'absolute', top: 10, left: 10, zIndex: 999 }}
+        onClick={addNode}
+      >
+        Add Nodes
+      </button>
+      <button
+        style={{ position: 'absolute', top: 10, left: 100, zIndex: 999 }}
+        onClick={undo}
+        disabled={!canUndo}
+      >
+        Undo
+      </button>
+      <button
+        style={{ position: 'absolute', top: 10, left: 160, zIndex: 999 }}
+        onClick={redo}
+        disabled={!canRedo}
+      >
+        Redo
+      </button>
+      <Canvas
+        nodes={nodes}
+        edges={edges}
+        onLayoutChange={layout => console.log('Layout', layout)}
+      />
+    </div>
+  );
+};