refactor(ui): edge rendering

- Fix issues with positioning of labels
- Optimize styling to be less reliant on JS

invokeai/frontend/web/src/features/nodes/components/flow/edges/InvocationCollapsedEdge.tsx

@@ -1,15 +1,40 @@
-import { Badge, Flex } from '@invoke-ai/ui-library';
-import { useStore } from '@nanostores/react';
+import type { SystemStyleObject } from '@invoke-ai/ui-library';
+import { Badge, Box, chakra } from '@invoke-ai/ui-library';
 import type { EdgeProps } from '@xyflow/react';
 import { BaseEdge, EdgeLabelRenderer, getBezierPath } from '@xyflow/react';
 import { useAppSelector } from 'app/store/storeHooks';
-import { useChakraThemeTokens } from 'common/hooks/useChakraThemeTokens';
-import { getEdgeStyles } from 'features/nodes/components/flow/edges/util/getEdgeColor';
-import { makeEdgeSelector } from 'features/nodes/components/flow/edges/util/makeEdgeSelector';
-import { $templates } from 'features/nodes/store/nodesSlice';
+import { buildSelectAreConnectedNodesSelected } from 'features/nodes/components/flow/edges/util/buildEdgeSelectors';
+import { selectShouldAnimateEdges } from 'features/nodes/store/workflowSettingsSlice';
 import type { CollapsedInvocationNodeEdge } from 'features/nodes/types/invocation';
 import { memo, useMemo } from 'react';
 
+const ChakraBaseEdge = chakra(BaseEdge);
+
+const baseEdgeSx: SystemStyleObject = {
+  strokeWidth: '3px !important',
+  stroke: 'base.500 !important',
+  opacity: '0.5 !important',
+  strokeDasharray: 'none',
+  '&[data-selected="true"]': {
+    opacity: '1 !important',
+  },
+  '&[data-selected="true"], &[data-are-connected-nodes-selected="true"]': {
+    strokeDasharray: '5 !important',
+  },
+  '&[data-should-animate-edges="true"]': {
+    animation: 'dashdraw 0.5s linear infinite !important',
+  },
+};
+
+const badgeSx: SystemStyleObject = {
+  bg: 'base.500',
+  opacity: 0.5,
+  shadow: 'base',
+  '&[data-selected="true"]': {
+    opacity: 1,
+  },
+};
+
 const InvocationCollapsedEdge = ({
   sourceX,
   sourceY,
@@ -21,17 +46,15 @@ const InvocationCollapsedEdge = ({
   data,
   selected = false,
   source,
-  sourceHandleId,
   target,
-  targetHandleId,
 }: EdgeProps<CollapsedInvocationNodeEdge>) => {
-  const templates = useStore($templates);
-  const selector = useMemo(
-    () => makeEdgeSelector(templates, source, sourceHandleId, target, targetHandleId),
-    [templates, source, sourceHandleId, target, targetHandleId]
+  const shouldAnimateEdges = useAppSelector(selectShouldAnimateEdges);
+  const selectAreConnectedNodesSelected = useMemo(
+    () => buildSelectAreConnectedNodesSelected(source, target),
+    [source, target]
   );
 
-  const { shouldAnimateEdges, areConnectedNodesSelected } = useAppSelector(selector);
+  const areConnectedNodesSelected = useAppSelector(selectAreConnectedNodesSelected);
 
   const [edgePath, labelX, labelY] = getBezierPath({
     sourceX,
@@ -42,31 +65,29 @@ const InvocationCollapsedEdge = ({
     targetPosition,
   });
 
-  const { base500 } = useChakraThemeTokens();
-
-  const edgeStyles = useMemo(
-    () => getEdgeStyles(base500, selected, shouldAnimateEdges, areConnectedNodesSelected),
-    [areConnectedNodesSelected, base500, selected, shouldAnimateEdges]
-  );
-
   return (
     <>
-      <BaseEdge path={edgePath} markerEnd={markerEnd} style={edgeStyles} />
-      {data?.count && data.count > 1 && (
+      <ChakraBaseEdge
+        path={edgePath}
+        markerEnd={markerEnd}
+        sx={baseEdgeSx}
+        data-selected={selected}
+        data-are-connected-nodes-selected={areConnectedNodesSelected}
+        data-should-animate-edges={shouldAnimateEdges}
+      />
+      {data?.count !== undefined && (
         <EdgeLabelRenderer>
-          <Flex
-            data-testid="asdfasdfasdf"
+          <Box
             position="absolute"
             transform={`translate(-50%, -50%) translate(${labelX}px,${labelY}px)`}
-            className="nodrag nopan"
-            // Unfortunately edge labels do not get the same zIndex treatment as edges do, so we need to manage this ourselves
+            className="edge-label-renderer__custom-edge nodrag nopan" // Unfortunately edge labels do not get the same zIndex treatment as edges do, so we need to manage this ourselves
             // See: https://github.com/xyflow/xyflow/issues/3658
             zIndex={1001}
           >
-            <Badge variant="solid" bg="base.500" opacity={selected ? 0.8 : 0.5} boxShadow="base">
+            <Badge variant="solid" sx={badgeSx} data-selected={selected}>
               {data.count}
             </Badge>
-          </Flex>
+          </Box>
         </EdgeLabelRenderer>
       )}
     </>

invokeai/frontend/web/src/features/nodes/components/flow/edges/InvocationDefaultEdge.tsx

@@ -1,15 +1,61 @@
-import { Flex, Text } from '@invoke-ai/ui-library';
+import type { SystemStyleObject } from '@invoke-ai/ui-library';
+import { chakra, Flex, Text } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import type { EdgeProps } from '@xyflow/react';
 import { BaseEdge, EdgeLabelRenderer, getBezierPath } from '@xyflow/react';
 import { useAppSelector } from 'app/store/storeHooks';
-import { getEdgeStyles } from 'features/nodes/components/flow/edges/util/getEdgeColor';
 import { $templates } from 'features/nodes/store/nodesSlice';
-import { selectShouldShowEdgeLabels } from 'features/nodes/store/workflowSettingsSlice';
+import { selectShouldAnimateEdges, selectShouldShowEdgeLabels } from 'features/nodes/store/workflowSettingsSlice';
 import type { DefaultInvocationNodeEdge } from 'features/nodes/types/invocation';
 import { memo, useMemo } from 'react';
 
-import { makeEdgeSelector } from './util/makeEdgeSelector';
+import {
+  buildSelectAreConnectedNodesSelected,
+  buildSelectEdgeColor,
+  buildSelectEdgeLabel,
+} from './util/buildEdgeSelectors';
+
+const ChakraBaseEdge = chakra(BaseEdge);
+
+const baseEdgeSx: SystemStyleObject = {
+  strokeWidth: '3px !important',
+  opacity: '0.5 !important',
+  strokeDasharray: 'none',
+  '&[data-selected="true"]': {
+    opacity: '1 !important',
+  },
+  '&[data-selected="true"], &[data-are-connected-nodes-selected="true"]': {
+    strokeDasharray: '5 !important',
+  },
+  '&[data-should-animate-edges="true"]': {
+    animation: 'dashdraw 0.5s linear infinite !important',
+  },
+};
+
+const edgeLabelWrapperSx: SystemStyleObject = {
+  pointerEvents: 'all',
+  position: 'absolute',
+  bg: 'base.800',
+  borderRadius: 'base',
+  borderWidth: 1,
+  opacity: 0.5,
+  borderColor: 'transparent',
+  py: 1,
+  px: 3,
+  shadow: 'md',
+  '&[data-selected="true"]': {
+    opacity: 1,
+    borderColor: undefined,
+  },
+};
+
+const edgeLabelTextSx: SystemStyleObject = {
+  fontWeight: 'semibold',
+  color: 'base.300',
+  '&[data-selected="true"]': {
+    color: 'base.100',
+  },
+};
 
 const InvocationDefaultEdge = ({
   sourceX,
@@ -26,13 +72,24 @@ const InvocationDefaultEdge = ({
   targetHandleId,
 }: EdgeProps<DefaultInvocationNodeEdge>) => {
   const templates = useStore($templates);
-  const selector = useMemo(
-    () => makeEdgeSelector(templates, source, sourceHandleId, target, targetHandleId),
+  const shouldAnimateEdges = useAppSelector(selectShouldAnimateEdges);
+  const shouldShowEdgeLabels = useAppSelector(selectShouldShowEdgeLabels);
+
+  const selectAreConnectedNodesSelected = useMemo(
+    () => buildSelectAreConnectedNodesSelected(source, target),
+    [source, target]
+  );
+  const selectStrokeColor = useMemo(
+    () => buildSelectEdgeColor(templates, source, sourceHandleId, target, targetHandleId),
     [templates, source, sourceHandleId, target, targetHandleId]
   );
-
-  const { shouldAnimateEdges, areConnectedNodesSelected, stroke, label } = useAppSelector(selector);
-  const shouldShowEdgeLabels = useAppSelector(selectShouldShowEdgeLabels);
+  const selectEdgeLabel = useMemo(
+    () => buildSelectEdgeLabel(templates, source, sourceHandleId, target, targetHandleId),
+    [templates, source, sourceHandleId, target, targetHandleId]
+  );
+  const areConnectedNodesSelected = useAppSelector(selectAreConnectedNodesSelected);
+  const stroke = useAppSelector(selectStrokeColor);
+  const label = useAppSelector(selectEdgeLabel);
 
   const [edgePath, labelX, labelY] = getBezierPath({
     sourceX,
@@ -43,31 +100,26 @@ const InvocationDefaultEdge = ({
     targetPosition,
   });
 
-  const edgeStyles = useMemo(
-    () => getEdgeStyles(stroke, selected, shouldAnimateEdges, areConnectedNodesSelected),
-    [areConnectedNodesSelected, stroke, selected, shouldAnimateEdges]
-  );
-
   return (
     <>
-      <BaseEdge path={edgePath} markerEnd={markerEnd} style={edgeStyles} />
+      <ChakraBaseEdge
+        path={edgePath}
+        markerEnd={markerEnd}
+        sx={baseEdgeSx}
+        stroke={`${stroke} !important`}
+        data-selected={selected}
+        data-are-connected-nodes-selected={areConnectedNodesSelected}
+        data-should-animate-edges={shouldAnimateEdges}
+      />
       {label && shouldShowEdgeLabels && (
         <EdgeLabelRenderer>
           <Flex
             className="nodrag nopan"
-            pointerEvents="all"
-            position="absolute"
             transform={`translate(-50%, -50%) translate(${labelX}px,${labelY}px)`}
-            bg="base.800"
-            borderRadius="base"
-            borderWidth={1}
-            borderColor={selected ? 'undefined' : 'transparent'}
-            opacity={selected ? 1 : 0.5}
-            py={1}
-            px={3}
-            shadow="md"
+            data-selected={selected}
+            sx={edgeLabelWrapperSx}
           >
-            <Text size="sm" fontWeight="semibold" color={selected ? 'base.100' : 'base.300'}>
+            <Text size="sm" sx={edgeLabelTextSx} data-selected={selected}>
               {label}
             </Text>
           </Flex>

invokeai/frontend/web/src/features/nodes/components/flow/edges/util/buildEdgeSelectors.ts

@@ -0,0 +1,62 @@
+import { createSelector } from '@reduxjs/toolkit';
+import { colorTokenToCssVar } from 'common/util/colorTokenToCssVar';
+import { selectNodesSlice } from 'features/nodes/store/selectors';
+import type { Templates } from 'features/nodes/store/types';
+import { selectWorkflowSettingsSlice } from 'features/nodes/store/workflowSettingsSlice';
+import { isInvocationNode } from 'features/nodes/types/invocation';
+
+import { getFieldColor } from './getEdgeColor';
+
+export const buildSelectAreConnectedNodesSelected = (source: string, target: string) =>
+  createSelector(selectNodesSlice, (nodes): boolean => {
+    const sourceNode = nodes.nodes.find((node) => node.id === source);
+    const targetNode = nodes.nodes.find((node) => node.id === target);
+
+    return Boolean(sourceNode?.selected || targetNode?.selected);
+  });
+
+export const buildSelectEdgeColor = (
+  templates: Templates,
+  source: string,
+  sourceHandleId: string | null | undefined,
+  target: string,
+  targetHandleId: string | null | undefined
+) =>
+  createSelector(selectNodesSlice, selectWorkflowSettingsSlice, (nodes, workflowSettings): string => {
+    const { shouldColorEdges } = workflowSettings;
+    const sourceNode = nodes.nodes.find((node) => node.id === source);
+    const targetNode = nodes.nodes.find((node) => node.id === target);
+
+    if (!sourceNode || !sourceHandleId || !targetNode || !targetHandleId) {
+      return colorTokenToCssVar('base.500');
+    }
+
+    const sourceNodeTemplate = templates[sourceNode.data.type];
+
+    const isInvocationToInvocationEdge = isInvocationNode(sourceNode) && isInvocationNode(targetNode);
+    const outputFieldTemplate = sourceNodeTemplate?.outputs[sourceHandleId];
+    const sourceType = isInvocationToInvocationEdge ? outputFieldTemplate?.type : undefined;
+
+    return sourceType && shouldColorEdges ? getFieldColor(sourceType) : colorTokenToCssVar('base.500');
+  });
+
+export const buildSelectEdgeLabel = (
+  templates: Templates,
+  source: string,
+  sourceHandleId: string | null | undefined,
+  target: string,
+  targetHandleId: string | null | undefined
+) =>
+  createSelector(selectNodesSlice, (nodes): string | null => {
+    const sourceNode = nodes.nodes.find((node) => node.id === source);
+    const targetNode = nodes.nodes.find((node) => node.id === target);
+
+    if (!sourceNode || !sourceHandleId || !targetNode || !targetHandleId) {
+      return null;
+    }
+
+    const sourceNodeTemplate = templates[sourceNode.data.type];
+    const targetNodeTemplate = templates[targetNode.data.type];
+
+    return `${sourceNodeTemplate?.title || sourceNode.data?.label} -> ${targetNodeTemplate?.title || targetNode.data?.label}`;
+  });

invokeai/frontend/web/src/features/nodes/components/flow/edges/util/getEdgeColor.ts

@@ -1,7 +1,6 @@
 import { colorTokenToCssVar } from 'common/util/colorTokenToCssVar';
 import { FIELD_COLORS } from 'features/nodes/types/constants';
 import type { FieldType } from 'features/nodes/types/field';
-import type { CSSProperties } from 'react';
 
 export const getFieldColor = (fieldType: FieldType | null): string => {
   if (!fieldType) {
@@ -11,16 +10,3 @@ export const getFieldColor = (fieldType: FieldType | null): string => {
 
   return color ? colorTokenToCssVar(color) : colorTokenToCssVar('base.500');
 };
-
-export const getEdgeStyles = (
-  stroke: string,
-  selected: boolean,
-  shouldAnimateEdges: boolean,
-  areConnectedNodesSelected: boolean
-): CSSProperties => ({
-  strokeWidth: 3,
-  stroke,
-  opacity: selected ? 1 : 0.5,
-  animation: shouldAnimateEdges ? 'dashdraw 0.5s linear infinite' : undefined,
-  strokeDasharray: selected || areConnectedNodesSelected ? 5 : 'none',
-});

invokeai/frontend/web/src/features/nodes/components/flow/edges/util/makeEdgeSelector.ts

@@ -1,58 +0,0 @@
-import { createMemoizedSelector } from 'app/store/createMemoizedSelector';
-import { colorTokenToCssVar } from 'common/util/colorTokenToCssVar';
-import { deepClone } from 'common/util/deepClone';
-import { selectNodesSlice } from 'features/nodes/store/selectors';
-import type { Templates } from 'features/nodes/store/types';
-import { selectWorkflowSettingsSlice } from 'features/nodes/store/workflowSettingsSlice';
-import { isInvocationNode } from 'features/nodes/types/invocation';
-
-import { getFieldColor } from './getEdgeColor';
-
-const defaultReturnValue = {
-  areConnectedNodesSelected: false,
-  shouldAnimateEdges: false,
-  stroke: colorTokenToCssVar('base.500'),
-  label: '',
-};
-
-export const makeEdgeSelector = (
-  templates: Templates,
-  source: string,
-  sourceHandleId: string | null | undefined,
-  target: string,
-  targetHandleId: string | null | undefined
-) =>
-  createMemoizedSelector(
-    selectNodesSlice,
-    selectWorkflowSettingsSlice,
-    (
-      nodes,
-      workflowSettings
-    ): { areConnectedNodesSelected: boolean; shouldAnimateEdges: boolean; stroke: string; label: string } => {
-      const { shouldAnimateEdges, shouldColorEdges } = workflowSettings;
-      const sourceNode = nodes.nodes.find((node) => node.id === source);
-      const targetNode = nodes.nodes.find((node) => node.id === target);
-
-      const returnValue = deepClone(defaultReturnValue);
-      returnValue.shouldAnimateEdges = shouldAnimateEdges;
-
-      const isInvocationToInvocationEdge = isInvocationNode(sourceNode) && isInvocationNode(targetNode);
-
-      returnValue.areConnectedNodesSelected = Boolean(sourceNode?.selected || targetNode?.selected);
-      if (!sourceNode || !sourceHandleId || !targetNode || !targetHandleId) {
-        return returnValue;
-      }
-
-      const sourceNodeTemplate = templates[sourceNode.data.type];
-      const targetNodeTemplate = templates[targetNode.data.type];
-
-      const outputFieldTemplate = sourceNodeTemplate?.outputs[sourceHandleId];
-      const sourceType = isInvocationToInvocationEdge ? outputFieldTemplate?.type : undefined;
-
-      returnValue.stroke = sourceType && shouldColorEdges ? getFieldColor(sourceType) : colorTokenToCssVar('base.500');
-
-      returnValue.label = `${sourceNodeTemplate?.title || sourceNode.data?.label} -> ${targetNodeTemplate?.title || targetNode.data?.label}`;
-
-      return returnValue;
-    }
-  );

invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/InvocationNodeCollapsedHandles.tsx

@@ -1,40 +1,25 @@
 import { Handle, Position } from '@xyflow/react';
-import { useChakraThemeTokens } from 'common/hooks/useChakraThemeTokens';
 import { useNodeTemplate } from 'features/nodes/hooks/useNodeTemplate';
 import { map } from 'lodash-es';
 import type { CSSProperties } from 'react';
-import { memo, useMemo } from 'react';
+import { memo } from 'react';
 
 interface Props {
   nodeId: string;
 }
 
 const hiddenHandleStyles: CSSProperties = { visibility: 'hidden' };
+const collapsedHandleStyles: CSSProperties = {
+  borderWidth: 0,
+  borderRadius: '3px',
+  width: '1rem',
+  height: '1rem',
+  backgroundColor: 'var(--invoke-colors-base-600)',
+  zIndex: -1,
+};
 
 const InvocationNodeCollapsedHandles = ({ nodeId }: Props) => {
   const template = useNodeTemplate(nodeId);
-  const { base600 } = useChakraThemeTokens();
-
-  const dummyHandleStyles: CSSProperties = useMemo(
-    () => ({
-      borderWidth: 0,
-      borderRadius: '3px',
-      width: '1rem',
-      height: '1rem',
-      backgroundColor: base600,
-      zIndex: -1,
-    }),
-    [base600]
-  );
-
-  const collapsedTargetStyles: CSSProperties = useMemo(
-    () => ({ ...dummyHandleStyles, left: '-0.5rem' }),
-    [dummyHandleStyles]
-  );
-  const collapsedSourceStyles: CSSProperties = useMemo(
-    () => ({ ...dummyHandleStyles, right: '-0.5rem' }),
-    [dummyHandleStyles]
-  );
 
   if (!template) {
     return null;
@@ -47,7 +32,7 @@ const InvocationNodeCollapsedHandles = ({ nodeId }: Props) => {
         id={`${nodeId}-collapsed-target`}
         isConnectable={false}
         position={Position.Left}
-        style={collapsedTargetStyles}
+        style={collapsedHandleStyles}
       />
       {map(template.inputs, (input) => (
         <Handle
@@ -64,7 +49,7 @@ const InvocationNodeCollapsedHandles = ({ nodeId }: Props) => {
         id={`${nodeId}-collapsed-source`}
         isConnectable={false}
         position={Position.Right}
-        style={collapsedSourceStyles}
+        style={collapsedHandleStyles}
       />
       {map(template.outputs, (output) => (
         <Handle