feat(ui): custom field types connection validation

In the initial commit, a custom field's original type was added to the *field templates* only as `originalType`. Custom fields' `type` property was `"Custom"`*. This allowed for type safety throughout the UI logic.

*Actually, it was `"Unknown"`, but I changed it to custom for clarity.

Connection validation logic, however, uses the *field instance* of the node/field. Like the templates, *field instances* with custom types have their `type` set to `"Custom"`, but they didn't have an `originalType` property. As a result, all custom fields could be connected to all other custom fields.

To resolve this, we need to add `originalType` to the *field instances*, then switch the validation logic to use this instead of `type`.

This ended up needing a bit of fanagling:

- If we make `originalType` a required property on field instances, existing workflows will break during connection validation, because they won't have this property. We'd need a new layer of logic to migrate the workflows, adding the new `originalType` property.

While this layer is probably needed anyways, typing `originalType` as optional is much simpler. Workflow migration logic can come layer.

(Technically, we could remove all references to field types from the workflow files, and let the templates hold all this information. This feels like a significant change and I'm reluctant to do it now.)

- Because `originalType` is optional, anywhere we care about the type of a field, we need to use it over `type`. So there are a number of `field.originalType ?? field.type` expressions. This is a bit of a gotcha, we'll need to remember this in the future.

- We use `Array.prototype.includes()` often in the workflow editor, e.g. `COLLECTION_TYPES.includes(type)`. In these cases, the const array is of type `FieldType[]`, and `type` is is `FieldType`.

Because we now support custom types, the arg `type` is now widened from `FieldType` to `string`.

This causes a TS error. This behaviour is somewhat controversial (see microsoft/TypeScript#14520). These expressions are now rewritten as `COLLECTION_TYPES.some((t) => t === type)` to satisfy TS. It's logically equivalent.

Author: psychedelicious

invokeai/frontend/web/src/features/nodes/components/flow/AddNodePopover/AddNodePopover.tsx

@@ -73,9 +73,13 @@ const AddNodePopover = () => {
 
             return some(handles, (handle) => {
               const sourceType =
-                handleFilter == 'source' ? fieldFilter : handle.type;
+                handleFilter == 'source'
+                  ? fieldFilter
+                  : handle.originalType ?? handle.type;
               const targetType =
-                handleFilter == 'target' ? fieldFilter : handle.type;
+                handleFilter == 'target'
+                  ? fieldFilter
+                  : handle.originalType ?? handle.type;
 
               return validateSourceAndTargetTypes(sourceType, targetType);
             });

invokeai/frontend/web/src/features/nodes/components/flow/connectionLines/CustomConnectionLine.tsx

@@ -4,14 +4,14 @@ import { useAppSelector } from 'app/store/storeHooks';
 import { colorTokenToCssVar } from 'common/util/colorTokenToCssVar';
 import { memo } from 'react';
 import { ConnectionLineComponentProps, getBezierPath } from 'reactflow';
-import { getFieldColor } from 'features/nodes/components/flow/edges/util/getEdgeColor';
+import { getFieldColor } from '../edges/util/getEdgeColor';
 
 const selector = createSelector(stateSelector, ({ nodes }) => {
-  const { shouldAnimateEdges, connectionStartFieldType, shouldColorEdges } =
+  const { shouldAnimateEdges, currentConnectionFieldType, shouldColorEdges } =
     nodes;
 
   const stroke = shouldColorEdges
-    ? getFieldColor(connectionStartFieldType)
+    ? getFieldColor(currentConnectionFieldType)
     : colorTokenToCssVar('base.500');
 
   let className = 'react-flow__custom_connection-path';

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

@@ -1,12 +1,12 @@
 import { colorTokenToCssVar } from 'common/util/colorTokenToCssVar';
-import { FIELD_COLORS } from 'features/nodes/types/constants';
-import { FieldType } from 'features/nodes/types/field';
+import { FIELDS } from 'features/nodes/types/constants';
+import { FieldType } from 'features/nodes/types/types';
 
-export const getFieldColor = (fieldType: FieldType | null): string => {
+export const getFieldColor = (fieldType: FieldType | string | null): string => {
   if (!fieldType) {
     return colorTokenToCssVar('base.500');
   }
-  const color = FIELD_COLORS[fieldType.name];
+  const color = FIELDS[fieldType]?.color;
 
   return color ? colorTokenToCssVar(color) : colorTokenToCssVar('base.500');
 };

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

@@ -2,7 +2,7 @@ import { createSelector } from '@reduxjs/toolkit';
 import { stateSelector } from 'app/store/store';
 import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
 import { colorTokenToCssVar } from 'common/util/colorTokenToCssVar';
-import { isInvocationNode } from 'features/nodes/types/invocation';
+import { isInvocationNode } from 'features/nodes/types/types';
 import { getFieldColor } from './getEdgeColor';
 
 export const makeEdgeSelector = (

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

@@ -1,8 +1,6 @@
 import { Tooltip } from '@chakra-ui/react';
-import { colorTokenToCssVar } from 'common/util/colorTokenToCssVar';
 import {
   COLLECTION_TYPES,
-  FIELDS,
   HANDLE_TOOLTIP_OPEN_DELAY,
   MODEL_TYPES,
   POLYMORPHIC_TYPES,
@@ -13,6 +11,7 @@ import {
 } from 'features/nodes/types/types';
 import { CSSProperties, memo, useMemo } from 'react';
 import { Handle, HandleType, Position } from 'reactflow';
+import { getFieldColor } from '../../../edges/util/getEdgeColor';
 
 export const handleBaseStyles: CSSProperties = {
   position: 'absolute',
@@ -47,14 +46,14 @@ const FieldHandle = (props: FieldHandleProps) => {
     isConnectionStartField,
     connectionError,
   } = props;
-  const { name, type, originalType } = fieldTemplate;
-  const { color: typeColor } = FIELDS[type];
+  const { name } = fieldTemplate;
+  const type = fieldTemplate.originalType ?? fieldTemplate.type;
 
   const styles: CSSProperties = useMemo(() => {
-    const isCollectionType = COLLECTION_TYPES.includes(type);
-    const isPolymorphicType = POLYMORPHIC_TYPES.includes(type);
-    const isModelType = MODEL_TYPES.includes(type);
-    const color = colorTokenToCssVar(typeColor);
+    const isCollectionType = COLLECTION_TYPES.some((t) => t === type);
+    const isPolymorphicType = POLYMORPHIC_TYPES.some((t) => t === type);
+    const isModelType = MODEL_TYPES.some((t) => t === type);
+    const color = getFieldColor(type);
     const s: CSSProperties = {
       backgroundColor:
         isCollectionType || isPolymorphicType
@@ -97,23 +96,14 @@ const FieldHandle = (props: FieldHandleProps) => {
     isConnectionInProgress,
     isConnectionStartField,
     type,
-    typeColor,
   ]);
 
   const tooltip = useMemo(() => {
-    if (isConnectionInProgress && isConnectionStartField) {
-      return originalType;
-    }
     if (isConnectionInProgress && connectionError) {
-      return connectionError ?? originalType;
+      return connectionError;
     }
-    return originalType;
-  }, [
-    connectionError,
-    isConnectionInProgress,
-    isConnectionStartField,
-    originalType,
-  ]);
+    return type;
+  }, [connectionError, isConnectionInProgress, type]);
 
   return (
     <Tooltip

invokeai/frontend/web/src/features/nodes/hooks/useFieldType.ts.ts

@@ -3,8 +3,8 @@ import { stateSelector } from 'app/store/store';
 import { useAppSelector } from 'app/store/storeHooks';
 import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
 import { useMemo } from 'react';
-import { KIND_MAP } from 'features/nodes/types/constants';
-import { isInvocationNode } from 'features/nodes/types/invocation';
+import { KIND_MAP } from '../types/constants';
+import { isInvocationNode } from '../types/types';
 
 export const useFieldType = (
   nodeId: string,
@@ -21,7 +21,7 @@ export const useFieldType = (
             return;
           }
           const field = node.data[KIND_MAP[kind]][fieldName];
-          return field?.type;
+          return field?.originalType ?? field?.type;
         },
         defaultSelectorOptions
       ),