diff --git a/corefx/src/main/java/com/core/ui/GraphToolbar.java b/corefx/src/main/java/com/core/ui/GraphToolbar.java
index c24df381..ff4f5dd1 100644
--- a/corefx/src/main/java/com/core/ui/GraphToolbar.java
+++ b/corefx/src/main/java/com/core/ui/GraphToolbar.java
@@ -45,7 +45,6 @@ public class GraphToolbar extends VBox {
     private boolean isEditing = false;
     @FXML private JFXButton runButton;
     @FXML private JFXButton pickingButton;
-    @FXML private JFXButton editingButton;
     @FXML private JFXButton drawingButton;
     @FXML private ComboBox<String> graphModeCombo;
     @FXML private JFXButton nodesButton;
@@ -61,7 +60,6 @@ public class GraphToolbar extends VBox {
         stopIcon.setSize(ICON_SIZE);
 
         setupPickingButton();
-        setupEditingButton();
         setupDrawingButton();
         setupNodesButton();
         setupDevicesButton();
@@ -82,24 +80,20 @@ public class GraphToolbar extends VBox {
             controller.getBottom().getChildren().remove(controller.getAnnotationToolbar());
             controller.getBorderPane().setRight(null);
             setSelected(true, pickingButton);
-            setSelected(false, editingButton, drawingButton, selectedEditButton);
+            setSelected(false, drawingButton, selectedEditButton);
             isEditing = false;
         });
     }
 
-    private void setupEditingButton() {
-        SVGGlyph editIcon = IconUtils.get("mode_edit");
-        editIcon.setSize(ICON_SIZE);
-        editingButton.setGraphic(editIcon);
-        editingButton.setTooltip(new Tooltip("Edit Graph"));
-        editingButton.setOnAction(event -> {
-            controller.getNetworkGraph().setMode(ModalGraphMouse.Mode.EDITING);
-            controller.getBottom().getChildren().remove(controller.getAnnotationToolbar());
-            controller.getBorderPane().setRight(null);
-            setSelected(true, editingButton, selectedEditButton);
-            setSelected(false, drawingButton, pickingButton);
-            isEditing = true;
-        });
+    private void setEditMode() {
+        controller.getNetworkGraph().setMode(ModalGraphMouse.Mode.EDITING);
+        controller.getBottom().getChildren().remove(controller.getAnnotationToolbar());
+        controller.getBorderPane().setRight(null);
+        if (selectedEditButton != null) {
+            setSelected(true, selectedEditButton);
+        }
+        setSelected(false, drawingButton, pickingButton);
+        isEditing = true;
     }
 
     private void setupDrawingButton() {
@@ -112,7 +106,7 @@ public class GraphToolbar extends VBox {
             controller.getBottom().getChildren().add(controller.getAnnotationToolbar());
             controller.getBorderPane().setRight(null);
             setSelected(true, drawingButton);
-            setSelected(false, editingButton, pickingButton, selectedEditButton);
+            setSelected(false, pickingButton, selectedEditButton);
             isEditing = false;
         });
     }
@@ -165,7 +159,9 @@ public class GraphToolbar extends VBox {
         JFXButton previous = selectedEditButton;
         selectedEditButton = button;
         if (isEditing) {
-            setSelected(false, previous);
+            if (previous != null) {
+                setSelected(false, previous);
+            }
             setSelected(true, selectedEditButton);
         }
     }
@@ -187,8 +183,11 @@ public class GraphToolbar extends VBox {
         });
 
         JFXPopup popup = new JFXPopup(nodesList);
-        nodesButton.setOnAction(event -> popup.show(nodesButton, JFXPopup.PopupVPosition.TOP,
-                JFXPopup.PopupHPosition.LEFT, nodesButton.getWidth(), 0));
+        nodesButton.setOnAction(event -> {
+            setEditMode();
+            popup.show(nodesButton, JFXPopup.PopupVPosition.TOP,
+                    JFXPopup.PopupHPosition.LEFT, nodesButton.getWidth(), 0);
+        });
     }
 
     private void setupDevicesButton() {
@@ -208,8 +207,11 @@ public class GraphToolbar extends VBox {
         });
 
         JFXPopup popup = new JFXPopup(devicesList);
-        devicesButton.setOnAction(event -> popup.show(devicesButton, JFXPopup.PopupVPosition.TOP,
-                JFXPopup.PopupHPosition.LEFT, devicesButton.getWidth(), 0));
+        devicesButton.setOnAction(event -> {
+            setEditMode();
+            popup.show(devicesButton, JFXPopup.PopupVPosition.TOP,
+                    JFXPopup.PopupHPosition.LEFT, devicesButton.getWidth(), 0);
+        });
     }
 
     @FXML
@@ -272,7 +274,8 @@ public class GraphToolbar extends VBox {
         if (isRunning) {
             Platform.runLater(() -> {
                 pickingButton.fire();
-                editingButton.setDisable(true);
+                devicesButton.setDisable(true);
+                nodesButton.setDisable(true);
                 runButton.pseudoClassStateChanged(START_CLASS, false);
                 runButton.pseudoClassStateChanged(STOP_CLASS, true);
                 if (runButton.getGraphic() != stopIcon) {
@@ -282,7 +285,8 @@ public class GraphToolbar extends VBox {
             });
         } else {
             Platform.runLater(() -> {
-                editingButton.setDisable(false);
+                devicesButton.setDisable(false);
+                nodesButton.setDisable(false);
                 runButton.pseudoClassStateChanged(START_CLASS, true);
                 runButton.pseudoClassStateChanged(STOP_CLASS, false);
                 if (runButton.getGraphic() != startIcon) {
diff --git a/corefx/src/main/resources/fxml/graph_toolbar.fxml b/corefx/src/main/resources/fxml/graph_toolbar.fxml
index 70d3c35d..f393d122 100644
--- a/corefx/src/main/resources/fxml/graph_toolbar.fxml
+++ b/corefx/src/main/resources/fxml/graph_toolbar.fxml
@@ -8,7 +8,6 @@
    <children>
       <JFXButton id="run-button" fx:id="runButton" onAction="#onRunButtonAction" styleClass="toolbar-button" />
       <JFXButton fx:id="pickingButton" styleClass="toolbar-button" />
-      <JFXButton fx:id="editingButton" styleClass="toolbar-button" />
       <JFXButton fx:id="nodesButton" styleClass="toolbar-button" text="Nodes" />
       <JFXButton fx:id="devicesButton" styleClass="toolbar-button" text="Devices" />
       <JFXButton fx:id="drawingButton" styleClass="toolbar-button" />