[ENG-1108] Canvas tool lock (#851)

* wip

* lint

* update

* address PR comments

* new behavior

* address PR comment

* lint

* address PR comments

* lint

* address PR comments

* further cleanup

* another fix

Author: trangdoan982

apps/roam/src/components/canvas/Clipboard.tsx

@@ -672,7 +672,6 @@ const ClipboardPageSection = ({
         },
       };
       editor.createShape<DiscourseNodeShape>(shape);
-      editor.setCurrentTool("select");
     },
     [editor, extensionAPI, showNodesOnCanvas],
   );

apps/roam/src/components/canvas/DiscourseNodeUtil.tsx

@@ -110,6 +110,7 @@ export const createNodeShapeTools = (
     return class DiscourseNodeTool extends StateNode {
       static id = n.type;
       static initial = "idle";
+      static isLockable = true;
       shapeType = n.type;
 
       override onEnter = () => {
@@ -130,7 +131,6 @@ export const createNodeShapeTools = (
           props: { fontFamily: "sans", size: "s" },
         });
         this.editor.setEditingShape(shapeId);
-        this.editor.setCurrentTool("select");
       };
     };
   });
@@ -519,6 +519,19 @@ export class BaseDiscourseNodeUtil extends BaseBoxShapeUtil<DiscourseNodeShape>
           this.updateProps(shape.id, shape.type, { h, w, imageUrl });
         };
 
+        const wasToolLocked = this.editor.getInstanceState().isToolLocked;
+
+        const restoreToolState = () => {
+          if (wasToolLocked) {
+            this.editor.updateInstanceState({ isToolLocked: true });
+            this.editor.setCurrentTool(shape.type);
+          } else {
+            this.editor.setCurrentTool("select");
+          }
+          editor.setEditingShape(null);
+          dialogRenderedRef.current = false;
+        };
+
         renderModifyNodeDialog({
           mode: isCreating ? "create" : "edit",
           nodeType: shape.type,
@@ -570,13 +583,14 @@ export class BaseDiscourseNodeUtil extends BaseBoxShapeUtil<DiscourseNodeShape>
                 });
               }
             }
-
-            editor.setEditingShape(null);
-            dialogRenderedRef.current = false;
           },
           onClose: () => {
-            editor.setEditingShape(null);
-            dialogRenderedRef.current = false;
+            if (isCreating) {
+              restoreToolState();
+            } else {
+              editor.setEditingShape(null);
+              dialogRenderedRef.current = false;
+            }
           },
         });
 

apps/roam/src/components/canvas/DiscourseRelationShape/DiscourseRelationTool.tsx

@@ -28,6 +28,7 @@ export const createAllReferencedNodeTools = (
     class ReferencedNodeTool extends StateNode {
       static override initial = "idle";
       static override id = action;
+      static override isLockable = true;
       static override children = (): TLStateNodeConstructor[] => [
         this.Idle,
         this.Pointing,
@@ -277,10 +278,6 @@ export const createAllReferencedNodeTools = (
           this.editor.setCursor({ type: "cross" });
         };
 
-        override onCancel = () => {
-          this.editor.setCurrentTool("select");
-        };
-
         override onKeyUp: TLEventHandlers["onKeyUp"] = (info) => {
           if (info.key === "Enter") {
             if (this.editor.getInstanceState().isReadonly) return null;
@@ -315,6 +312,7 @@ export const createAllRelationShapeTools = (
     class RelationShapeTool extends StateNode {
       static override initial = "idle";
       static override id = name;
+      static override isLockable = true;
       static override children = (): TLStateNodeConstructor[] => [
         this.Idle,
         this.Pointing,
@@ -573,10 +571,6 @@ export const createAllRelationShapeTools = (
           this.editor.setCursor({ type: "cross", rotation: 0 });
         };
 
-        override onCancel = () => {
-          this.editor.setCurrentTool("select");
-        };
-
         override onKeyUp: TLEventHandlers["onKeyUp"] = (info) => {
           if (info.key === "Enter") {
             if (this.editor.getInstanceState().isReadonly) return null;

apps/roam/src/components/canvas/DiscourseRelationShape/DiscourseRelationUtil.tsx

@@ -1086,7 +1086,6 @@ export class BaseDiscourseRelationUtil extends ShapeUtil<DiscourseRelationShape>
   static override props = arrowShapeProps;
 
   cancelAndWarn = (title: string) => {
-    this.editor.setCurrentTool("select");
     dispatchToastEvent({
       id: `tldraw-cancel-and-warn-${title}`,
       title,

apps/roam/src/components/canvas/DiscourseToolPanel.tsx

@@ -190,7 +190,6 @@ const DiscourseGraphPanel = ({
               props: { fontFamily: "sans", size: "s" },
             });
             editor.setEditingShape(shapeId);
-            editor.setCurrentTool("select");
           } else {
             // For relations, just activate the tool
             editor.setCurrentTool(current.item.id);

apps/roam/src/components/canvas/Tldraw.tsx

@@ -693,6 +693,7 @@ const TldrawCanvasShared = ({
   const discourseGraphTool = class DiscourseGraphTool extends StateNode {
     static override id = "discourse-tool";
     static override initial = "idle";
+    static override isLockable = true;
   };
   const discourseNodeTools = createNodeShapeTools(allNodes);
   const discourseRelationTools = createAllRelationShapeTools(allRelationNames);