diff --git a/src/Eagle.ts b/src/Eagle.ts
index 3bf335f2d..6dd94165a 100644
--- a/src/Eagle.ts
+++ b/src/Eagle.ts
@@ -993,8 +993,6 @@ export class Eagle {
     }
 
     createSubgraphFromSelection = () : void => {
-        console.log("createSubgraphFromSelection()");
-
         const eagle = Eagle.getInstance()
         if(eagle.selectedObjects().length === 0){
             Utils.showNotification('Error','At least one node must be selected!', 'warning')
diff --git a/src/LogicalGraph.ts b/src/LogicalGraph.ts
index f5e9a3caa..d27b2420e 100644
--- a/src/LogicalGraph.ts
+++ b/src/LogicalGraph.ts
@@ -55,11 +55,11 @@ export class LogicalGraph {
 
         result.modelData = FileInfo.toOJSJson(graph.fileInfo());
         result.modelData.schemaVersion = Daliuge.SchemaVersion.OJS;
-        result.modelData.numLGNodes = graph.getNodes().length;
+        result.modelData.numLGNodes = graph.nodes().length;
 
         // add nodes
         result.nodeDataArray = [];
-        for (const node of graph.getNodes()){
+        for (const node of graph.nodes()){
             const nodeData : any = Node.toOJSGraphJson(node);
             result.nodeDataArray.push(nodeData);
         }
@@ -267,7 +267,7 @@ export class LogicalGraph {
     getCommentNodes = () : Node[] => {
         const commentNodes: Node[] = [];
 
-        for (const node of this.getNodes()){
+        for (const node of this.nodes()){
             if (node.isComment()){
                 commentNodes.push(node);
             }
@@ -419,14 +419,13 @@ export class LogicalGraph {
     }
 
     findNodeGraphIdByNodeName = (name:string) :string =>{
-        const eagle: Eagle = Eagle.getInstance();
-        let graphNodeId:string
-        eagle.logicalGraph().getNodes().forEach(function(node){
-            if(node.getName() === name){
-                graphNodeId = node.getId()
+        for (const node of this.nodes()){
+            if (node.getName() === name){
+                return node.getId();
             }
-        })
-        return graphNodeId
+        }
+
+        return null;
     }
 
     removeNode = (node: Node) : void => {
@@ -552,7 +551,6 @@ export class LogicalGraph {
             return;
         }
 
-        const nodes : Node[] = this.getNodes();
         let minX : number = Number.MAX_SAFE_INTEGER;
         let minY : number = Number.MAX_SAFE_INTEGER;
         let maxX : number = Number.MIN_SAFE_INTEGER;
@@ -560,7 +558,7 @@ export class LogicalGraph {
         let numChildren : number = 0;
 
         // loop through all nodes, finding all children and determining minimum bounding box to contain all children
-        for (const n of nodes){
+        for (const n of this.nodes()){
             if (n.getParentKey() === node.getKey()){
                 numChildren += 1;
 
@@ -729,7 +727,7 @@ export class LogicalGraph {
 
         // populate index plus depths
         for (let i = 0 ; i < this.nodes().length ; i++){
-            const node = this.getNodes()[i];
+            const node = this.nodes()[i];
 
             const depth = this.findDepthByKey(node.getKey());
 
@@ -743,7 +741,7 @@ export class LogicalGraph {
 
         // write nodes to result in sorted order
         for (const indexPlusDepth of indexPlusDepths){
-            result.push(this.getNodes()[indexPlusDepth.index]);
+            result.push(this.nodes()[indexPlusDepth.index]);
         }
 
         return result;
diff --git a/src/Undo.ts b/src/Undo.ts
index fba473c12..57f3e418f 100644
--- a/src/Undo.ts
+++ b/src/Undo.ts
@@ -124,6 +124,8 @@ export class Undo {
         }
 
         eagle.checkGraph();
+
+        this._updateSelection();
     }
 
     nextSnapshot = (eagle: Eagle) : void => {
@@ -141,6 +143,8 @@ export class Undo {
         }
 
         eagle.checkGraph();
+
+        this._updateSelection();
     }
 
     toString = () : string => {
@@ -174,8 +178,37 @@ export class Undo {
         }
 
         const dataObject: LogicalGraph = snapshot.data();
+        eagle.logicalGraph(dataObject.clone());
+    }
+
+    // if we undo, or redo, then the objects in selectedObject are from the graph prior to the new snapshot
+    // so the references will be to non-existent objects
+    // in this function, we use the ids of the old selectedObjects, and attempt to add the matching objects in the new snapshot to the selectedObjects list
+    _updateSelection = () : void => {
+        const eagle: Eagle = Eagle.getInstance();
+        const objectIds: string[] = [];
+
+        // build a list of the ids of the selected objects
+        for (const object of eagle.selectedObjects()){
+            objectIds.push(object.getId());
+        }
 
-        eagle.logicalGraph(dataObject);
+        // clear selection
+        eagle.setSelection(Eagle.RightWindowMode.Hierarchy, null, Eagle.FileType.Graph);
+
+        // find the objects in the ids list, and add them to the selection
+        for (const id of objectIds){
+            const node = eagle.logicalGraph().findNodeById(id);
+            const edge = eagle.logicalGraph().findEdgeById(id);
+            const object = node || edge;
+
+            // abort if no edge or node exists fot that id
+            if (node === null && edge === null){
+                continue;
+            }
+
+            eagle.editSelection(<Eagle.RightWindowMode>eagle.rightWindow().mode(), object, Eagle.selectedLocation());
+        }
     }
 
     static printTable() : void {
@@ -194,6 +227,8 @@ export class Undo {
                 "current": realCurrent === i ? "->" : "",
                 "description": snapshot.description(),
                 "buffer position": i,
+                "nodes": snapshot.data().getNodes().length,
+                "edges": snapshot.data().getEdges().length
             });
         }