diff --git a/editor/js/tools/menus/collectmenu.js b/editor/js/tools/menus/collectmenu.js
index e002016f0887adc84d53efe3f6b67a579363957f..df6ddaf5789f457b0bd01d4100d7e30a9dee3668 100644
--- a/editor/js/tools/menus/collectmenu.js
+++ b/editor/js/tools/menus/collectmenu.js
@@ -1,7 +1,7 @@
 import jQuery from "jquery";
 import { Graph } from "../../graph";
 import { CONTEXT } from "../../state";
-import { state } from "../editor";
+import { state } from "../../editor";
 import ToolMenu from "./toolmenu";
 
 export const COLLECTION_KEY = "collection";
@@ -68,18 +68,24 @@ export class CollectMenu extends ToolMenu {
         }
 
         // Disable previous context
-        if (this.context !== undefined) {
-            this.getDomToContext(this.context).addClass(HIDDEN_CLASS);
+        var dom = this.getDomToContext(this.context);
+        if (dom !== undefined) {
+            dom.addClass(HIDDEN_CLASS);
         }
 
         // Store and activate new context
         this.context = context;
-        if (this.context !== undefined) {
-            this.getDomToContext(this.context).removeClass(HIDDEN_CLASS);
+        dom = this.getDomToContext(this.context);
+        if (dom !== undefined) {
+            dom.removeClass(HIDDEN_CLASS);
         }
     }
 
     getDomToContext(context) {
+        if (context === undefined) {
+            return undefined;
+        }
+
         if (context === CONTEXT.node) {
             return this.find(CONTEXT_NODE);
         } else {