From e3a143eaed67e4eef3babaef95724c955cff2bfc Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Thu, 16 May 2024 21:56:00 +1000
Subject: [PATCH] fix(ui): fix jank w/ stale connections

---
 .../src/features/nodes/components/flow/Flow.tsx   | 15 ++++++++++++++-
 .../flow/nodes/Invocation/fields/InputField.tsx   | 10 +++++++++-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx b/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx
index 34c180ec53..c1542caabc 100644
--- a/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx
+++ b/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx
@@ -7,6 +7,8 @@ import { useIsValidConnection } from 'features/nodes/hooks/useIsValidConnection'
 import { useWorkflowWatcher } from 'features/nodes/hooks/useWorkflowWatcher';
 import {
   $cursorPos,
+  $isAddNodePopoverOpen,
+  $pendingConnection,
   $viewport,
   connectionMade,
   edgeAdded,
@@ -33,8 +35,9 @@ import type {
   OnNodesDelete,
   ProOptions,
   ReactFlowProps,
+  ReactFlowState,
 } from 'reactflow';
-import { Background, ReactFlow } from 'reactflow';
+import { Background, ReactFlow, useStore as useReactFlowStore } from 'reactflow';
 
 import CustomConnectionLine from './connectionLines/CustomConnectionLine';
 import InvocationCollapsedEdge from './edges/InvocationCollapsedEdge';
@@ -61,6 +64,8 @@ const proOptions: ProOptions = { hideAttribution: true };
 
 const snapGrid: [number, number] = [25, 25];
 
+const selectCancelConnection = (state: ReactFlowState) => state.cancelConnection;
+
 export const Flow = memo(() => {
   const dispatch = useAppDispatch();
   const nodes = useAppSelector((s) => s.nodes.present.nodes);
@@ -73,6 +78,7 @@ export const Flow = memo(() => {
   const { onConnectStart, onConnect, onConnectEnd } = useConnection();
   const flowWrapper = useRef<HTMLDivElement>(null);
   const isValidConnection = useIsValidConnection();
+  const cancelConnection = useReactFlowStore(selectCancelConnection);
   useWorkflowWatcher();
   const [borderRadius] = useToken('radii', ['base']);
 
@@ -234,6 +240,13 @@ export const Flow = memo(() => {
   }, [dispatch, mayRedo]);
   useHotkeys(['meta+shift+z', 'ctrl+shift+z'], onRedoHotkey);
 
+  const onEscapeHotkey = useCallback(() => {
+    $pendingConnection.set(null);
+    $isAddNodePopoverOpen.set(false);
+    cancelConnection();
+  }, [cancelConnection]);
+  useHotkeys('esc', onEscapeHotkey);
+
   return (
     <ReactFlow
       id="workflow-editor"
diff --git a/invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/InputField.tsx b/invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/InputField.tsx
index 66b0d3f755..cad4dded39 100644
--- a/invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/InputField.tsx
+++ b/invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/InputField.tsx
@@ -95,7 +95,15 @@ const InputField = ({ nodeId, fieldName }: Props) => {
 
   return (
     <InputFieldWrapper shouldDim={shouldDim}>
-      <FormControl isInvalid={isMissingInput} isDisabled={isConnected} orientation="vertical" px={2}>
+      <FormControl
+        isInvalid={isMissingInput}
+        isDisabled={isConnected}
+        // Without pointerEvents prop, disabled inputs don't trigger reactflow events. For example, when making a
+        // connection, the mouse up to end the connection won't fire, leaving the connection in-progress.
+        pointerEvents={isConnected ? 'none' : 'auto'}
+        orientation="vertical"
+        px={2}
+      >
         <Flex flexDir="column" w="full" gap={1} onMouseEnter={onMouseEnter} onMouseLeave={onMouseLeave}>
           <Flex>
             <EditableFieldTitle