From 83e0a333b04df8f7272f0f0269363b7bf24fe2ce Mon Sep 17 00:00:00 2001
From: peterkogo <peter@xyflow.com>
Date: Tue, 27 Feb 2024 16:54:21 +0100
Subject: [PATCH] fixed handle connections

---
 .../react/src/components/Handle/index.tsx     | 19 +++++++++++++++----
 packages/react/src/store/index.ts             |  6 ++----
 packages/system/src/styles/init.css           |  5 +++++
 packages/system/src/xyhandle/utils.ts         |  4 ++--
 4 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/packages/react/src/components/Handle/index.tsx b/packages/react/src/components/Handle/index.tsx
index 36778039..25794ee6 100644
--- a/packages/react/src/components/Handle/index.tsx
+++ b/packages/react/src/components/Handle/index.tsx
@@ -16,6 +16,7 @@ import {
   type HandleProps,
   type Connection,
   type HandleType,
+  ConnectionMode,
 } from '@xyflow/system';
 
 import { useStore, useStoreApi } from '../../hooks/useStore';
@@ -36,14 +37,19 @@ const connectingSelector =
       connectionStartHandle: startHandle,
       connectionEndHandle: endHandle,
       connectionClickStartHandle: clickHandle,
+      connectionMode,
     } = state;
-
     return {
       connecting:
         (startHandle?.nodeId === nodeId && startHandle?.handleId === handleId && startHandle?.type === type) ||
         (endHandle?.nodeId === nodeId && endHandle?.handleId === handleId && endHandle?.type === type),
       clickConnecting:
         clickHandle?.nodeId === nodeId && clickHandle?.handleId === handleId && clickHandle?.type === type,
+      isPossibleEndHandle:
+        connectionMode === ConnectionMode.Strict
+          ? startHandle?.type !== type
+          : nodeId !== startHandle?.nodeId || handleId !== startHandle?.handleId,
+      connectionInProcess: !!startHandle,
     };
   };
 
@@ -71,7 +77,10 @@ const HandleComponent = forwardRef<HTMLDivElement, HandleComponentProps>(
     const store = useStoreApi();
     const nodeId = useNodeId();
     const { connectOnClick, noPanClassName, rfId } = useStore(selector, shallow);
-    const { connecting, clickConnecting } = useStore(connectingSelector(nodeId, handleId, type), shallow);
+    const { connecting, clickConnecting, isPossibleEndHandle, connectionInProcess } = useStore(
+      connectingSelector(nodeId, handleId, type),
+      shallow
+    );
 
     if (!nodeId) {
       store.getState().onError?.('010', errorMessages['error010']());
@@ -201,10 +210,12 @@ const HandleComponent = forwardRef<HTMLDivElement, HandleComponentProps>(
             connectable: isConnectable,
             connectablestart: isConnectableStart,
             connectableend: isConnectableEnd,
-            connecting: clickConnecting,
+            connecting: connecting || clickConnecting,
             // this class is used to style the handle when the user is connecting
             connectionindicator:
-              isConnectable && ((isConnectableStart && !connecting) || (isConnectableEnd && connecting)),
+              isConnectable &&
+              (!connectionInProcess || isPossibleEndHandle) &&
+              (connectionInProcess ? isConnectableEnd : isConnectableStart),
           },
         ])}
         onMouseDown={onPointerDown}
diff --git a/packages/react/src/store/index.ts b/packages/react/src/store/index.ts
index 71a51a1e..cae14173 100644
--- a/packages/react/src/store/index.ts
+++ b/packages/react/src/store/index.ts
@@ -291,13 +291,11 @@ const createRFStore = ({
           connectionEndHandle: null,
         }),
       updateConnection: (params) => {
-        const { connectionStatus, connectionStartHandle, connectionEndHandle, connectionPosition } = get();
+        const { connectionPosition } = get();
 
         const currentConnection = {
+          ...params,
           connectionPosition: params.connectionPosition ?? connectionPosition,
-          connectionStatus: params.connectionStatus ?? connectionStatus,
-          connectionStartHandle: params.connectionStartHandle ?? connectionStartHandle,
-          connectionEndHandle: params.connectionEndHandle ?? connectionEndHandle,
         };
 
         set(currentConnection);
diff --git a/packages/system/src/styles/init.css b/packages/system/src/styles/init.css
index 46628842..d74519a6 100644
--- a/packages/system/src/styles/init.css
+++ b/packages/system/src/styles/init.css
@@ -223,9 +223,14 @@ svg.xy-flow__connectionline {
   min-width: 5px;
   min-height: 5px;
 
+  &.connecting {
+    pointer-events: all;
+  }
+
   &.connectionindicator {
     pointer-events: all;
     cursor: crosshair;
+    background: green;
   }
 
   &-bottom {
diff --git a/packages/system/src/xyhandle/utils.ts b/packages/system/src/xyhandle/utils.ts
index b829f691..c10356b1 100644
--- a/packages/system/src/xyhandle/utils.ts
+++ b/packages/system/src/xyhandle/utils.ts
@@ -38,7 +38,7 @@ export function getClosestHandle(
   let closestHandles: ConnectionHandle[] = [];
   let minDistance = Infinity;
 
-  handles.forEach((handle) => {
+  for (const handle of handles) {
     const distance = Math.sqrt(Math.pow(handle.x - pos.x, 2) + Math.pow(handle.y - pos.y, 2));
     if (distance <= connectionRadius) {
       if (distance < minDistance) {
@@ -49,7 +49,7 @@ export function getClosestHandle(
       }
       minDistance = distance;
     }
-  });
+  }
 
   if (!closestHandles.length) {
     return null;