Refine Arrow connection UI

ui/package.json

@@ -10,7 +10,7 @@
     "@emotion/styled": "^11.10.5",
     "@mui/icons-material": "^5.10.9",
     "@mui/material": "^5.10.10",
-    "@reactflow/node-resizer": "^1.2.0",
+    "@reactflow/node-resizer": "^2.0.1",
     "@remirror/extension-react-tables": "^2.2.9",
     "@remirror/extension-sub": "^2.0.9",
     "@remirror/extension-text-highlight": "^2.0.10",
@@ -42,7 +42,7 @@
     "react-resizable": "^3.0.4",
     "react-router-dom": "^6.4.2",
     "react-scripts": "5.0.1",
-    "reactflow": "^11.4.0",
+    "reactflow": "^11.5.6",
     "remirror": "^2.0.21",
     "slate": "^0.82.1",
     "slate-history": "^0.66.0",

ui/src/App.css

@@ -87,3 +87,7 @@
 .react-flow__node-scope.selected {
   border-width: 2px;
 }
+
+.react-flow__handle {
+  z-index: 100;
+}

ui/src/components/Canvas.tsx

@@ -42,8 +42,13 @@ import { RichNode } from "./nodes/Rich";
 import { CodeNode } from "./nodes/Code";
 import { ScopeNode } from "./nodes/Scope";
 import { YMap } from "yjs/dist/src/types/YMap";
+import FloatingEdge from "./nodes/FloatingEdge";
+import CustomConnectionLine from "./nodes/CustomConnectionLine";
 
 const nodeTypes = { scope: ScopeNode, code: CodeNode, rich: RichNode };
+const edgeTypes = {
+  floating: FloatingEdge,
+};
 
 /**
  * This hook will load nodes from zustand store into Yjs nodesMap using setNodes.
@@ -529,6 +534,23 @@ function CanvasImpl() {
           minZoom={0.1}
           onPaneContextMenu={onPaneContextMenu}
           nodeTypes={nodeTypes}
+          // custom edge for easy connect
+          edgeTypes={edgeTypes}
+          defaultEdgeOptions={{
+            style: { strokeWidth: 3, stroke: "black" },
+            type: "floating",
+            markerEnd: {
+              type: MarkerType.ArrowClosed,
+              color: "black",
+            },
+          }}
+          connectionLineComponent={CustomConnectionLine}
+          connectionLineStyle={{
+            strokeWidth: 3,
+            stroke: "black",
+          }}
+          // end custom edge
+
           zoomOnScroll={false}
           panOnScroll={true}
           connectionMode={ConnectionMode.Loose}

ui/src/components/nodes/Code.tsx

@@ -221,6 +221,111 @@ export const ResultBlock = memo<any>(function ResultBlock({ id }) {
   );
 });
 
+function FloatingToolbar({ id }) {
+  const store = useContext(RepoContext);
+  if (!store) throw new Error("Missing BearContext.Provider in the tree");
+  const reactFlowInstance = useReactFlow();
+  const devMode = useStore(store, (state) => state.devMode);
+  // const pod = useStore(store, (state) => state.pods[id]);
+  const wsRun = useStore(store, (state) => state.wsRun);
+  const clearResults = useStore(store, (s) => s.clearResults);
+  // right, bottom
+  const [layout, setLayout] = useState("bottom");
+  const getPod = useStore(store, (state) => state.getPod);
+  const isGuest = useStore(store, (state) => state.role === "GUEST");
+  const clonePod = useStore(store, (state) => state.clonePod);
+  const setPaneFocus = useStore(store, (state) => state.setPaneFocus);
+
+  const onCopy = useCallback(
+    (clipboardData: any) => {
+      const pod = clonePod(id);
+      if (!pod) return;
+      clipboardData.setData("text/plain", pod.content);
+      clipboardData.setData(
+        "application/json",
+        JSON.stringify({
+          type: "pod",
+          data: pod,
+        })
+      );
+      setPaneFocus();
+    },
+    [clonePod, id]
+  );
+
+  const cutBegin = useStore(store, (state) => state.cutBegin);
+
+  const onCut = useCallback(
+    (clipboardData: any) => {
+      onCopy(clipboardData);
+      cutBegin(id);
+    },
+    [onCopy, cutBegin, id]
+  );
+
+  return (
+    <Box>
+      {!isGuest && (
+        <Tooltip title="Run (shift-enter)">
+          <IconButton
+            size="small"
+            onClick={() => {
+              clearResults(id);
+              wsRun(id);
+            }}
+          >
+            <PlayCircleOutlineIcon fontSize="inherit" />
+          </IconButton>
+        </Tooltip>
+      )}
+      <CopyToClipboard
+        text="dummy"
+        options={{ debug: true, format: "text/plain", onCopy } as any}
+      >
+        <Tooltip title="Copy">
+          <IconButton size="small" className="copy-button">
+            <ContentCopyIcon fontSize="inherit" className="copy-button" />
+          </IconButton>
+        </Tooltip>
+      </CopyToClipboard>
+      {!isGuest && (
+        <CopyToClipboard
+          text="dummy"
+          options={{ debug: true, format: "text/plain", onCopy: onCut } as any}
+        >
+          <Tooltip title="Cut">
+            <IconButton size="small">
+              <ContentCutIcon fontSize="inherit" />
+            </IconButton>
+          </Tooltip>
+        </CopyToClipboard>
+      )}
+      {!isGuest && (
+        <Tooltip title="Delete">
+          <IconButton
+            size="small"
+            onClick={() => {
+              reactFlowInstance.deleteElements({ nodes: [{ id }] });
+            }}
+          >
+            <DeleteIcon fontSize="inherit" />
+          </IconButton>
+        </Tooltip>
+      )}
+      <Tooltip title="Change layout">
+        <IconButton
+          size="small"
+          onClick={() => {
+            setLayout(layout === "bottom" ? "right" : "bottom");
+          }}
+        >
+          <ViewComfyIcon fontSize="inherit" />
+        </IconButton>
+      </Tooltip>
+    </Box>
+  );
+}
+
 export const CodeNode = memo<NodeProps>(function ({
   data,
   id,
@@ -243,7 +348,7 @@ export const CodeNode = memo<NodeProps>(function ({
   const setPodName = useStore(store, (state) => state.setPodName);
   const setPodGeo = useStore(store, (state) => state.setPodGeo);
   const getPod = useStore(store, (state) => state.getPod);
-  const clonePod = useStore(store, (state) => state.clonePod);
+
   const pod = getPod(id);
   const isGuest = useStore(store, (state) => state.role === "GUEST");
   const isPodFocused = useStore(store, (state) => state.pods[id]?.focus);
@@ -265,7 +370,7 @@ export const CodeNode = memo<NodeProps>(function ({
       state.pods[id]?.stderr
   );
   const nodesMap = useStore(store, (state) => state.ydoc.getMap<Node>("pods"));
-  const setPaneFocus = useStore(store, (state) => state.setPaneFocus);
+
   const onResize = useCallback(
     (e, data) => {
       const { size } = data;
@@ -290,6 +395,7 @@ export const CodeNode = memo<NodeProps>(function ({
     [id, nodesMap, setPodGeo, updateView]
   );
 
+  const [showToolbar, setShowToolbar] = useState(false);
   useEffect(() => {
     if (!data.name) return;
     setPodName({ id, name: data.name });
@@ -298,33 +404,6 @@ export const CodeNode = memo<NodeProps>(function ({
     }
   }, [data.name, setPodName, id]);
 
-  const onCopy = useCallback(
-    (clipboardData: any) => {
-      const pod = clonePod(id);
-      if (!pod) return;
-      clipboardData.setData("text/plain", pod.content);
-      clipboardData.setData(
-        "application/json",
-        JSON.stringify({
-          type: "pod",
-          data: pod,
-        })
-      );
-      setPaneFocus();
-    },
-    [clonePod, id]
-  );
-
-  const cutBegin = useStore(store, (state) => state.cutBegin);
-
-  const onCut = useCallback(
-    (clipboardData: any) => {
-      onCopy(clipboardData);
-      cutBegin(id);
-    },
-    [onCopy, cutBegin, id]
-  );
-
   // if (!pod) throw new Error(`Pod not found: ${id}`);
 
   if (!pod) {
@@ -370,6 +449,12 @@ export const CodeNode = memo<NodeProps>(function ({
           ? "#d6dee6"
           : "#5e92f3",
       }}
+      onMouseEnter={() => {
+        setShowToolbar(true);
+      }}
+      onMouseLeave={() => {
+        setShowToolbar(false);
+      }}
     >
       {/* FIXME this does not support x-axis only resizing. */}
       {/* <NodeResizeControl
@@ -462,7 +547,8 @@ export const CodeNode = memo<NodeProps>(function ({
         </Box>
         <Box
           sx={{
-            display: "flex",
+            // display: "flex",
+            display: showToolbar ? "flex" : "none",
             marginLeft: "10px",
             borderRadius: "4px",
             position: "absolute",
@@ -475,65 +561,7 @@ export const CodeNode = memo<NodeProps>(function ({
           }}
           className="nodrag"
         >
-          {!isGuest && (
-            <Tooltip title="Run (shift-enter)">
-              <IconButton
-                size="small"
-                onClick={() => {
-                  clearResults(id);
-                  wsRun(id);
-                }}
-              >
-                <PlayCircleOutlineIcon fontSize="inherit" />
-              </IconButton>
-            </Tooltip>
-          )}
-          <CopyToClipboard
-            text="dummy"
-            options={{ debug: true, format: "text/plain", onCopy } as any}
-          >
-            <Tooltip title="Copy">
-              <IconButton size="small" className="copy-button">
-                <ContentCopyIcon fontSize="inherit" className="copy-button" />
-              </IconButton>
-            </Tooltip>
-          </CopyToClipboard>
-          {!isGuest && (
-            <CopyToClipboard
-              text="dummy"
-              options={
-                { debug: true, format: "text/plain", onCopy: onCut } as any
-              }
-            >
-              <Tooltip title="Cut">
-                <IconButton size="small">
-                  <ContentCutIcon fontSize="inherit" />
-                </IconButton>
-              </Tooltip>
-            </CopyToClipboard>
-          )}
-          {!isGuest && (
-            <Tooltip title="Delete">
-              <IconButton
-                size="small"
-                onClick={() => {
-                  reactFlowInstance.deleteElements({ nodes: [{ id }] });
-                }}
-              >
-                <DeleteIcon fontSize="inherit" />
-              </IconButton>
-            </Tooltip>
-          )}
-          <Tooltip title="Change layout">
-            <IconButton
-              size="small"
-              onClick={() => {
-                setLayout(layout === "bottom" ? "right" : "bottom");
-              }}
-            >
-              <ViewComfyIcon fontSize="inherit" />
-            </IconButton>
-          </Tooltip>
+          <FloatingToolbar id={id} />
         </Box>
       </Box>
       <Box

ui/src/components/nodes/CustomConnectionLine.tsx

@@ -0,0 +1,32 @@
+import { ConnectionLineComponentProps, getStraightPath } from "reactflow";
+
+function CustomConnectionLine({
+  fromX,
+  fromY,
+  toX,
+  toY,
+  connectionLineStyle,
+}: ConnectionLineComponentProps) {
+  const [edgePath] = getStraightPath({
+    sourceX: fromX,
+    sourceY: fromY,
+    targetX: toX,
+    targetY: toY,
+  });
+
+  return (
+    <g>
+      <path style={connectionLineStyle} fill="none" d={edgePath} />
+      <circle
+        cx={toX}
+        cy={toY}
+        fill="black"
+        r={3}
+        stroke="black"
+        strokeWidth={1.5}
+      />
+    </g>
+  );
+}
+
+export default CustomConnectionLine;

ui/src/components/nodes/FloatingEdge.tsx

@@ -0,0 +1,38 @@
+import { useCallback } from "react";
+import { useStore, getStraightPath, EdgeProps } from "reactflow";
+
+import { getEdgeParams } from "./utils";
+
+function FloatingEdge({ id, source, target, markerEnd, style }: EdgeProps) {
+  const sourceNode = useStore(
+    useCallback((store) => store.nodeInternals.get(source), [source])
+  );
+  const targetNode = useStore(
+    useCallback((store) => store.nodeInternals.get(target), [target])
+  );
+
+  if (!sourceNode || !targetNode) {
+    return null;
+  }
+
+  const { sx, sy, tx, ty } = getEdgeParams(sourceNode, targetNode);
+
+  const [edgePath] = getStraightPath({
+    sourceX: sx,
+    sourceY: sy,
+    targetX: tx,
+    targetY: ty,
+  });
+
+  return (
+    <path
+      id={id}
+      className="react-flow__edge-path"
+      d={edgePath}
+      markerEnd={markerEnd}
+      style={style}
+    />
+  );
+}
+
+export default FloatingEdge;