feat(ui): support adding form elements and node fields with dnd

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

@@ -1,12 +1,18 @@
+import { combine } from '@atlaskit/pragmatic-drag-and-drop/combine';
+import { draggable } from '@atlaskit/pragmatic-drag-and-drop/element/adapter';
 import { Flex, FormControl } from '@invoke-ai/ui-library';
+import { firefoxDndFix } from 'features/dnd/util';
 import { FieldHandle } from 'features/nodes/components/flow/nodes/Invocation/fields/FieldHandle';
 import { InputFieldNotesIconButtonEditable } from 'features/nodes/components/flow/nodes/Invocation/fields/InputFieldNotesIconButtonEditable';
 import { InputFieldResetToDefaultValueIconButton } from 'features/nodes/components/flow/nodes/Invocation/fields/InputFieldResetToDefaultValueIconButton';
+import { buildNodeFieldDndData } from 'features/nodes/components/sidePanel/builder/use-builder-dnd';
 import { useInputFieldConnectionState } from 'features/nodes/hooks/useInputFieldConnectionState';
 import { useInputFieldIsConnected } from 'features/nodes/hooks/useInputFieldIsConnected';
 import { useInputFieldIsInvalid } from 'features/nodes/hooks/useInputFieldIsInvalid';
 import { useInputFieldTemplate } from 'features/nodes/hooks/useInputFieldTemplate';
-import { memo, useCallback, useState } from 'react';
+import type { FieldIdentifier } from 'features/nodes/types/field';
+import type { RefObject } from 'react';
+import { memo, useCallback, useEffect, useRef, useState } from 'react';
 
 import { InputFieldAddRemoveLinearViewIconButton } from './InputFieldAddRemoveLinearViewIconButton';
 import { InputFieldRenderer } from './InputFieldRenderer';
@@ -20,6 +26,8 @@ interface Props {
 
 export const InputFieldEditModeNodes = memo(({ nodeId, fieldName }: Props) => {
   const fieldTemplate = useInputFieldTemplate(nodeId, fieldName);
+  const draggableRef = useRef<HTMLDivElement>(null);
+  const dragHandleRef = useRef<HTMLDivElement>(null);
   const [isHovered, setIsHovered] = useState(false);
   const isInvalid = useInputFieldIsInvalid(nodeId, fieldName);
   const isConnected = useInputFieldIsConnected(nodeId, fieldName);
@@ -36,6 +44,8 @@ export const InputFieldEditModeNodes = memo(({ nodeId, fieldName }: Props) => {
     setIsHovered(false);
   }, []);
 
+  const isDragging = useNodeFieldDnd({ nodeId, fieldName }, draggableRef, dragHandleRef);
+
   if (fieldTemplate.input === 'connection' || isConnected) {
     return (
       <InputFieldWrapper>
@@ -62,6 +72,7 @@ export const InputFieldEditModeNodes = memo(({ nodeId, fieldName }: Props) => {
   return (
     <InputFieldWrapper>
       <FormControl
+        ref={draggableRef}
         isInvalid={isInvalid}
         isDisabled={isConnected}
         // Without pointerEvents prop, disabled inputs don't trigger reactflow events. For example, when making a
@@ -69,9 +80,10 @@ export const InputFieldEditModeNodes = memo(({ nodeId, fieldName }: Props) => {
         pointerEvents={isConnected ? 'none' : 'auto'}
         orientation="vertical"
         px={2}
+        opacity={isDragging ? 0.3 : 1}
       >
         <Flex flexDir="column" w="full" gap={1} onMouseEnter={onMouseEnter} onMouseLeave={onMouseLeave}>
-          <Flex gap={1}>
+          <Flex className="nodrag" ref={dragHandleRef} gap={1}>
             <InputFieldTitle nodeId={nodeId} fieldName={fieldName} isInvalid={isInvalid} />
             {isHovered && (
               <>
@@ -99,3 +111,35 @@ export const InputFieldEditModeNodes = memo(({ nodeId, fieldName }: Props) => {
 });
 
 InputFieldEditModeNodes.displayName = 'InputFieldEditModeNodes';
+
+const useNodeFieldDnd = (
+  fieldIdentifier: FieldIdentifier,
+  draggableRef: RefObject<HTMLElement>,
+  dragHandleRef: RefObject<HTMLElement>
+) => {
+  const [isDragging, setIsDragging] = useState(false);
+
+  useEffect(() => {
+    const draggableElement = draggableRef.current;
+    const dragHandleElement = dragHandleRef.current;
+    if (!draggableElement || !dragHandleElement) {
+      return;
+    }
+    return combine(
+      firefoxDndFix(draggableElement),
+      draggable({
+        element: draggableElement,
+        dragHandle: dragHandleElement,
+        getInitialData: () => buildNodeFieldDndData(fieldIdentifier),
+        onDragStart: () => {
+          setIsDragging(true);
+        },
+        onDrop: () => {
+          setIsDragging(false);
+        },
+      })
+    );
+  }, [dragHandleRef, draggableRef, fieldIdentifier]);
+
+  return isDragging;
+};

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/ContainerElementComponent.tsx

@@ -13,7 +13,7 @@ import { TextElementComponent } from 'features/nodes/components/sidePanel/builde
 import { formElementAdded, selectWorkflowFormMode, useElement } from 'features/nodes/store/workflowSlice';
 import type { ContainerElement } from 'features/nodes/types/workflow';
 import {
-  container,
+  buildContainer,
   CONTAINER_CLASS_NAME,
   isContainerElement,
   isDividerElement,
@@ -99,7 +99,7 @@ ContainerElementComponentEditMode.displayName = 'ContainerElementComponentEditMo
 const AddColumnButton = ({ el }: { el: ContainerElement }) => {
   const dispatch = useAppDispatch();
   const onClick = useCallback(() => {
-    const element = container('column', [], el.id);
+    const element = buildContainer('column', [], el.id);
     dispatch(formElementAdded({ element, containerId: el.id }));
   }, [dispatch, el.id]);
   return (
@@ -110,7 +110,7 @@ const AddColumnButton = ({ el }: { el: ContainerElement }) => {
 const AddRowButton = ({ el }: { el: ContainerElement }) => {
   const dispatch = useAppDispatch();
   const onClick = useCallback(() => {
-    const element = container('row', [], el.id);
+    const element = buildContainer('row', [], el.id);
     dispatch(formElementAdded({ element, containerId: el.id }));
   }, [dispatch, el.id]);
   return (

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/DndListDropIndicator.tsx

@@ -2,7 +2,7 @@
 import type { Edge } from '@atlaskit/pragmatic-drag-and-drop-hitbox/types';
 import type { SystemStyleObject } from '@invoke-ai/ui-library';
 import { Box } from '@invoke-ai/ui-library';
-import type { DndListTargetState } from 'features/nodes/components/sidePanel/builder/use-builder-dnd';
+import type { CenterOrEdge } from 'features/nodes/components/sidePanel/builder/center-or-closest-edge';
 
 /**
  * Design decisions for the drop indicator's main line
@@ -103,18 +103,23 @@ function DndDropIndicatorInternal({ edge, gap = '0px' }: DropIndicatorProps) {
   );
 }
 
-export const DndListDropIndicator = ({ dndState, gap }: { dndState: DndListTargetState; gap?: string }) => {
-  if (dndState.type !== 'is-dragging-over') {
+export const DndListDropIndicator = ({
+  activeDropRegion,
+  gap,
+}: {
+  activeDropRegion: CenterOrEdge | null;
+  gap?: string;
+}) => {
+  if (!activeDropRegion) {
     return null;
   }
-
-  if (!dndState.closestCenterOrEdge || dndState.closestCenterOrEdge === 'center') {
+  if (activeDropRegion === 'center') {
     return null;
   }
 
   return (
     <DndDropIndicatorInternal
-      edge={dndState.closestCenterOrEdge}
+      edge={activeDropRegion}
       // This is the gap between items in the list, used to calculate the position of the drop indicator
       gap={gap || 'var(--invoke-space-2)'}
     />

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/FormElementEditModeWrapper.tsx

@@ -37,7 +37,7 @@ const wrapperSx: SystemStyleObject = {
   '&[data-is-dragging="true"]': {
     opacity: 0.3,
   },
-  '&[data-is-dragging-over-center="true"]': {
+  '&[data-active-drop-region="center"]': {
     opacity: 1,
     bg: 'base.700',
   },
@@ -61,7 +61,7 @@ export const FormElementEditModeWrapper = memo(
   ({ element, children, ...rest }: { element: FormElement } & FlexProps) => {
     const draggableRef = useRef<HTMLDivElement>(null);
     const dragHandleRef = useRef<HTMLDivElement>(null);
-    const [dndListState, isDragging] = useDraggableFormElement(element.id, draggableRef, dragHandleRef);
+    const [activeDropRegion, isDragging] = useDraggableFormElement(element.id, draggableRef, dragHandleRef);
     const depth = useDepthContext();
     const dispatch = useAppDispatch();
     const removeElement = useCallback(() => {
@@ -75,9 +75,7 @@ export const FormElementEditModeWrapper = memo(
         sx={wrapperSx}
         className={EDIT_MODE_WRAPPER_CLASS_NAME}
         data-is-dragging={isDragging}
-        data-is-dragging-over-center={
-          dndListState.type === 'is-dragging-over' && dndListState.closestCenterOrEdge === 'center'
-        }
+        data-active-drop-region={activeDropRegion}
         {...rest}
       >
         <Flex ref={dragHandleRef} sx={headerSx} data-depth={depth}>
@@ -98,7 +96,7 @@ export const FormElementEditModeWrapper = memo(
         <Flex w="full" p={4} alignItems="center" gap={4}>
           {children}
         </Flex>
-        <DndListDropIndicator dndState={dndListState} gap="var(--invoke-space-4)" />
+        <DndListDropIndicator activeDropRegion={activeDropRegion} gap="var(--invoke-space-4)" />
       </Flex>
     );
   }

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/WorkflowBuilder.tsx

@@ -1,11 +1,27 @@
-import { Button, Flex } from '@invoke-ai/ui-library';
+import { combine } from '@atlaskit/pragmatic-drag-and-drop/combine';
+import { draggable } from '@atlaskit/pragmatic-drag-and-drop/element/adapter';
+import { Button, ButtonGroup, Flex } from '@invoke-ai/ui-library';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import ScrollableContent from 'common/components/OverlayScrollbars/ScrollableContent';
+import { firefoxDndFix } from 'features/dnd/util';
 import { FormElementComponent } from 'features/nodes/components/sidePanel/builder/ContainerElementComponent';
-import { useMonitorForFormElementDnd } from 'features/nodes/components/sidePanel/builder/use-builder-dnd';
+import {
+  buildAddFormElementDndData,
+  useMonitorForFormElementDnd,
+} from 'features/nodes/components/sidePanel/builder/use-builder-dnd';
 import { formLoaded, formModeToggled, selectWorkflowFormMode } from 'features/nodes/store/workflowSlice';
-import { elements, rootElementId } from 'features/nodes/types/workflow';
-import { memo, useCallback, useEffect } from 'react';
+import type { FormElement } from 'features/nodes/types/workflow';
+import {
+  buildContainer,
+  buildDivider,
+  buildHeading,
+  buildText,
+  elements,
+  rootElementId,
+} from 'features/nodes/types/workflow';
+import type { RefObject } from 'react';
+import { memo, useCallback, useEffect, useRef, useState } from 'react';
+import { assert } from 'tsafe';
 
 export const WorkflowBuilder = memo(() => {
   const dispatch = useAppDispatch();
@@ -19,8 +35,14 @@ export const WorkflowBuilder = memo(() => {
   return (
     <ScrollableContent>
       <Flex w="full" justifyContent="center">
-        <Flex flexDir="column" w={mode === 'view' ? '512px' : 'min-content'} minW="512px">
-          <ToggleModeButton />
+        <Flex flexDir="column" w={mode === 'view' ? '512px' : 'min-content'} minW="512px" gap={2}>
+          <ButtonGroup isAttached={false}>
+            <ToggleModeButton />
+            <AddFormElementDndButton type="container" />
+            <AddFormElementDndButton type="divider" />
+            <AddFormElementDndButton type="heading" />
+            <AddFormElementDndButton type="text" />
+          </ButtonGroup>
           {rootElementId && <FormElementComponent id={rootElementId} />}
         </Flex>
       </Flex>
@@ -41,3 +63,58 @@ const ToggleModeButton = memo(() => {
   return <Button onClick={onClick}>{mode === 'view' ? 'Edit' : 'View'}</Button>;
 });
 ToggleModeButton.displayName = 'ToggleModeButton';
+
+const useAddFormElementDnd = (type: Omit<FormElement['type'], 'node-field'>, draggableRef: RefObject<HTMLElement>) => {
+  const [isDragging, setIsDragging] = useState(false);
+
+  useEffect(() => {
+    const draggableElement = draggableRef.current;
+    if (!draggableElement) {
+      return;
+    }
+    return combine(
+      firefoxDndFix(draggableElement),
+      draggable({
+        element: draggableElement,
+        getInitialData: () => {
+          if (type === 'container') {
+            const element = buildContainer('row', []);
+            return buildAddFormElementDndData(element);
+          }
+          if (type === 'divider') {
+            const element = buildDivider();
+            return buildAddFormElementDndData(element);
+          }
+          if (type === 'heading') {
+            const element = buildHeading('default heading', 1);
+            return buildAddFormElementDndData(element);
+          }
+          if (type === 'text') {
+            const element = buildText('default text', 'sm');
+            return buildAddFormElementDndData(element);
+          }
+          assert(false);
+        },
+        onDragStart: () => {
+          setIsDragging(true);
+        },
+        onDrop: () => {
+          setIsDragging(false);
+        },
+      })
+    );
+  }, [draggableRef, type]);
+
+  return isDragging;
+};
+
+const AddFormElementDndButton = ({ type }: { type: Omit<FormElement['type'], 'node-field'> }) => {
+  const draggableRef = useRef<HTMLButtonElement>(null);
+  const isDragging = useAddFormElementDnd(type, draggableRef);
+
+  return (
+    <Button ref={draggableRef} variant="ghost" pointerEvents="auto" opacity={isDragging ? 0.3 : 1}>
+      {type}
+    </Button>
+  );
+};

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/use-builder-dnd.ts

@@ -14,41 +14,54 @@ import {
   extractClosestCenterOrEdge,
 } from 'features/nodes/components/sidePanel/builder/center-or-closest-edge';
 import { getEditModeWrapperId } from 'features/nodes/components/sidePanel/builder/FormElementEditModeWrapper';
-import { formElementMoved } from 'features/nodes/store/workflowSlice';
+import { formElementAdded, formElementMoved } from 'features/nodes/store/workflowSlice';
+import type { FieldIdentifier } from 'features/nodes/types/field';
 import type { ElementId, FormElement } from 'features/nodes/types/workflow';
-import { isContainerElement } from 'features/nodes/types/workflow';
+import { buildNodeField, isContainerElement } from 'features/nodes/types/workflow';
 import type { RefObject } from 'react';
-import { useEffect, useState } from 'react';
+import { useCallback, useEffect, useState } from 'react';
 import { flushSync } from 'react-dom';
 import { assert } from 'tsafe';
 
-/**
- * States for a dnd list with containers.
- */
-export type DndListTargetState =
-  | {
-      type: 'idle';
-    }
-  | {
-      type: 'preview';
-      container: HTMLElement;
-    }
-  | {
-      type: 'is-dragging';
-    }
-  | {
-      type: 'is-dragging-over';
-      closestCenterOrEdge: CenterOrEdge | null;
-    };
-export const idle: DndListTargetState = { type: 'idle' };
-
-// using a symbol so we can guarantee a key with a unique value
-const uniqueBuilderDndKey = Symbol('builderDnd');
-
-type DndData = {
-  [uniqueBuilderDndKey]: true;
+const uniqueMoveFormElementKey = Symbol('move-form-element');
+type MoveFormElementDndData = {
+  [uniqueMoveFormElementKey]: true;
   element: FormElement;
 };
+const buildMoveFormElementDndData = (element: FormElement): MoveFormElementDndData => ({
+  [uniqueMoveFormElementKey]: true,
+  element,
+});
+const isMoveFormElementDndData = (data: Record<string | symbol, unknown>): data is MoveFormElementDndData => {
+  return uniqueMoveFormElementKey in data;
+};
+
+const uniqueAddFormElementKey = Symbol('add-form-element');
+type AddFormElementDndData = {
+  [uniqueAddFormElementKey]: true;
+  element: FormElement;
+};
+export const buildAddFormElementDndData = (element: FormElement): AddFormElementDndData => ({
+  [uniqueAddFormElementKey]: true,
+  element,
+});
+const isAddFormElementDndData = (data: Record<string | symbol, unknown>): data is AddFormElementDndData => {
+  return uniqueAddFormElementKey in data;
+};
+
+const uniqueNodeFieldKey = Symbol('node-field');
+type NodeFieldDndData = {
+  [uniqueNodeFieldKey]: true;
+  fieldIdentifier: FieldIdentifier;
+};
+export const buildNodeFieldDndData = (fieldIdentifier: FieldIdentifier): NodeFieldDndData => ({
+  [uniqueNodeFieldKey]: true,
+  fieldIdentifier,
+});
+
+const isNodeFieldDndData = (data: Record<string | symbol, unknown>): data is NodeFieldDndData => {
+  return uniqueNodeFieldKey in data;
+};
 
 const getElement = <T extends FormElement>(id: ElementId, guard?: (el: FormElement) => el is T): T => {
   const el = getStore().getState().workflow.form?.elements[id];
@@ -61,88 +74,204 @@ const getElement = <T extends FormElement>(id: ElementId, guard?: (el: FormEleme
   }
 };
 
-const adjustIndexForDrop = (index: number, edge: Exclude<CenterOrEdge, 'center'>) => {
+const adjustIndexForFormElementMoveDrop = (index: number, edge: Exclude<CenterOrEdge, 'center'>) => {
   if (edge === 'left' || edge === 'top') {
     return index - 1;
   }
   return index + 1;
 };
 
+const adjustIndexForNodeFieldDrop = (index: number, edge: Exclude<CenterOrEdge, 'center'>) => {
+  if (edge === 'left' || edge === 'top') {
+    return index;
+  }
+  return index + 1;
+};
+
+const flashElement = (elementId: ElementId) => {
+  const element = document.querySelector(`#${getEditModeWrapperId(elementId)}`);
+  if (element instanceof HTMLElement) {
+    triggerPostMoveFlash(element, colorTokenToCssVar('base.700'));
+  }
+};
+
 export const useMonitorForFormElementDnd = () => {
   const dispatch = useAppDispatch();
 
+  const handleMoveFormElementDrop = useCallback(
+    (sourceData: MoveFormElementDndData, targetData: MoveFormElementDndData) => {
+      if (sourceData.element.id === targetData.element.id) {
+        return;
+      }
+
+      const closestCenterOrEdge = extractClosestCenterOrEdge(targetData);
+
+      if (closestCenterOrEdge === 'center') {
+        // Move the element to the target container - should we double-check that the target is a container?
+        flushSync(() => {
+          dispatch(formElementMoved({ id: sourceData.element.id, containerId: targetData.element.id }));
+        });
+        // Flash the element that was moved
+        flashElement(sourceData.element.id);
+      } else if (closestCenterOrEdge) {
+        // Move the element to the target's parent container at the correct index
+        const { parentId } = targetData.element;
+        assert(parentId !== undefined, 'Target element should have a parent');
+
+        const isReparenting = parentId !== sourceData.element.parentId;
+
+        const parentContainer = getElement(parentId, isContainerElement);
+        const targetIndex = parentContainer.data.children.findIndex((elementId) => elementId === targetData.element.id);
+
+        let index: number | undefined = undefined;
+
+        if (!isReparenting) {
+          const sourceIndex = parentContainer.data.children.findIndex(
+            (elementId) => elementId === sourceData.element.id
+          );
+          if (
+            sourceIndex === targetIndex ||
+            sourceIndex === adjustIndexForFormElementMoveDrop(targetIndex, closestCenterOrEdge)
+          ) {
+            return;
+          }
+          index = targetIndex;
+        } else {
+          index = adjustIndexForFormElementMoveDrop(targetIndex, closestCenterOrEdge);
+        }
+
+        flushSync(() => {
+          dispatch(
+            formElementMoved({
+              id: sourceData.element.id,
+              containerId: parentId,
+              index,
+            })
+          );
+        });
+        // Flash the element that was moved
+        flashElement(sourceData.element.id);
+      } else {
+        // No container, cannot do anything
+        return;
+      }
+    },
+    [dispatch]
+  );
+
+  const handleAddFormElementDrop = useCallback(
+    (sourceData: AddFormElementDndData, targetData: MoveFormElementDndData) => {
+      const closestCenterOrEdge = extractClosestCenterOrEdge(targetData);
+
+      if (closestCenterOrEdge === 'center') {
+        // Move the element to the target container - should we double-check that the target is a container?
+        const { element } = sourceData;
+        flushSync(() => {
+          dispatch(formElementAdded({ element, containerId: targetData.element.id }));
+        });
+        flashElement(element.id);
+      } else if (closestCenterOrEdge) {
+        // Move the element to the target's parent container at the correct index
+        const { parentId } = targetData.element;
+        assert(parentId !== undefined, 'Target element should have a parent');
+        const { element } = sourceData;
+
+        const parentContainer = getElement(parentId, isContainerElement);
+        const targetIndex = parentContainer.data.children.findIndex((elementId) => elementId === targetData.element.id);
+
+        const index = adjustIndexForNodeFieldDrop(targetIndex, closestCenterOrEdge);
+
+        flushSync(() => {
+          dispatch(
+            formElementAdded({
+              element,
+              containerId: parentId,
+              index,
+            })
+          );
+        });
+        flashElement(element.id);
+      } else {
+        // No container, cannot do anything
+        return;
+      }
+    },
+    [dispatch]
+  );
+
+  const handleNodeFieldDrop = useCallback(
+    (sourceData: NodeFieldDndData, targetData: MoveFormElementDndData) => {
+      const closestCenterOrEdge = extractClosestCenterOrEdge(targetData);
+      const { nodeId, fieldName } = sourceData.fieldIdentifier;
+
+      if (closestCenterOrEdge === 'center') {
+        // Move the element to the target container - should we double-check that the target is a container?
+        const element = buildNodeField(nodeId, fieldName, targetData.element.id);
+        flushSync(() => {
+          dispatch(formElementAdded({ element, containerId: targetData.element.id }));
+        });
+        flashElement(element.id);
+      } else if (closestCenterOrEdge) {
+        // Move the element to the target's parent container at the correct index
+        const { parentId } = targetData.element;
+        assert(parentId !== undefined, 'Target element should have a parent');
+        const element = buildNodeField(nodeId, fieldName, parentId);
+
+        const parentContainer = getElement(parentId, isContainerElement);
+        const targetIndex = parentContainer.data.children.findIndex((elementId) => elementId === targetData.element.id);
+
+        const index = adjustIndexForNodeFieldDrop(targetIndex, closestCenterOrEdge);
+
+        flushSync(() => {
+          dispatch(
+            formElementAdded({
+              element,
+              containerId: parentId,
+              index,
+            })
+          );
+        });
+        flashElement(element.id);
+      } else {
+        // No container, cannot do anything
+        return;
+      }
+    },
+    [dispatch]
+  );
+
   useEffect(() => {
     return monitorForElements({
-      canMonitor: ({ source }) => uniqueBuilderDndKey in source.data,
+      canMonitor: ({ source }) =>
+        isMoveFormElementDndData(source.data) ||
+        isNodeFieldDndData(source.data) ||
+        isAddFormElementDndData(source.data),
       onDrop: ({ location, source }) => {
         const target = location.current.dropTargets[0];
         if (!target) {
           return;
         }
 
-        const sourceData = source.data as DndData;
-        const targetData = target.data as DndData;
+        const sourceData = source.data;
+        const targetData = target.data;
 
-        //
-        if (sourceData.element.id === targetData.element.id) {
+        if (isMoveFormElementDndData(targetData) && isMoveFormElementDndData(sourceData)) {
+          handleMoveFormElementDrop(sourceData, targetData);
           return;
         }
 
-        const closestCenterOrEdge = extractClosestCenterOrEdge(targetData);
-
-        if (closestCenterOrEdge === 'center') {
-          // Move the element to the target container - should we double-check that the target is a container?
-          flushSync(() => {
-            dispatch(formElementMoved({ id: sourceData.element.id, containerId: targetData.element.id }));
-          });
-        } else if (closestCenterOrEdge) {
-          // Move the element to the target's parent container at the correct index
-          const { parentId } = targetData.element;
-          assert(parentId !== undefined, 'Target element should have a parent');
-
-          const isReparenting = parentId !== sourceData.element.parentId;
-
-          const parentContainer = getElement(parentId, isContainerElement);
-          const targetIndex = parentContainer.data.children.findIndex(
-            (elementId) => elementId === targetData.element.id
-          );
-
-          let index: number | undefined = undefined;
-
-          if (!isReparenting) {
-            const sourceIndex = parentContainer.data.children.findIndex(
-              (elementId) => elementId === sourceData.element.id
-            );
-            if (sourceIndex === targetIndex || sourceIndex === adjustIndexForDrop(targetIndex, closestCenterOrEdge)) {
-              return;
-            }
-            index = targetIndex;
-          } else {
-            index = adjustIndexForDrop(targetIndex, closestCenterOrEdge);
-          }
-
-          flushSync(() => {
-            dispatch(
-              formElementMoved({
-                id: sourceData.element.id,
-                containerId: parentId,
-                index,
-              })
-            );
-          });
-        } else {
-          // No container, cannot do anything
+        if (isMoveFormElementDndData(targetData) && isAddFormElementDndData(sourceData)) {
+          handleAddFormElementDrop(sourceData, targetData);
           return;
         }
 
-        // Flash the element that was moved
-        const element = document.querySelector(`#${getEditModeWrapperId(sourceData.element.id)}`);
-        if (element instanceof HTMLElement) {
-          triggerPostMoveFlash(element, colorTokenToCssVar('base.700'));
+        if (isMoveFormElementDndData(targetData) && isNodeFieldDndData(sourceData)) {
+          handleNodeFieldDrop(sourceData, targetData);
+          return;
         }
       },
     });
-  }, [dispatch]);
+  }, [handleAddFormElementDrop, handleMoveFormElementDrop, handleNodeFieldDrop]);
 };
 
 export const useDraggableFormElement = (
@@ -150,8 +279,8 @@ export const useDraggableFormElement = (
   draggableRef: RefObject<HTMLElement>,
   dragHandleRef: RefObject<HTMLElement>
 ) => {
-  const [dndListState, setListDndState] = useState<DndListTargetState>(idle);
   const [isDragging, setIsDragging] = useState(false);
+  const [activeDropRegion, setActiveDropRegion] = useState<CenterOrEdge | null>(null);
 
   useEffect(() => {
     const draggableElement = draggableRef.current;
@@ -169,30 +298,25 @@ export const useDraggableFormElement = (
       draggable({
         element: draggableElement,
         dragHandle: dragHandleElement,
-        getInitialData: () => ({
-          [uniqueBuilderDndKey]: true,
-          element: getElement(elementId),
-        }),
+        getInitialData: () => buildMoveFormElementDndData(getElement(elementId)),
         onDragStart: () => {
-          setListDndState({ type: 'is-dragging' });
           setIsDragging(true);
         },
         onDrop: () => {
-          setListDndState(idle);
           setIsDragging(false);
         },
       }),
       dropTargetForElements({
         element: draggableElement,
-        canDrop: ({ source }) => uniqueBuilderDndKey in source.data,
+        canDrop: ({ source }) =>
+          isMoveFormElementDndData(source.data) ||
+          isNodeFieldDndData(source.data) ||
+          isAddFormElementDndData(source.data),
         getData: ({ input }) => {
           const element = getElement(elementId);
           const container = element.parentId ? getElement(element.parentId, isContainerElement) : null;
 
-          const data: DndData = {
-            [uniqueBuilderDndKey]: true,
-            element,
-          };
+          const data = buildMoveFormElementDndData(element);
 
           const allowedCenterOrEdge: CenterOrEdge[] = [];
 
@@ -220,7 +344,7 @@ export const useDraggableFormElement = (
 
           // If the innermost target is not this draggable element, bail. We only want to react when dragging over _this_ element.
           if (!innermostDropTargetElement || innermostDropTargetElement !== draggableElement) {
-            setListDndState(idle);
+            setActiveDropRegion(null);
             return;
           }
 
@@ -228,28 +352,23 @@ export const useDraggableFormElement = (
 
           // Don't allow reparanting to the same container
           if (closestCenterOrEdge === 'center' && source.element === draggableElement) {
-            setListDndState(idle);
+            setActiveDropRegion(null);
             return;
           }
 
           // Only need to update react state if nothing has changed.
           // Prevents re-rendering.
-          setListDndState((current) => {
-            if (current.type === 'is-dragging-over' && current.closestCenterOrEdge === closestCenterOrEdge) {
-              return current;
-            }
-            return { type: 'is-dragging-over', closestCenterOrEdge };
-          });
+          setActiveDropRegion(closestCenterOrEdge);
         },
         onDragLeave: () => {
-          setListDndState(idle);
+          setActiveDropRegion(null);
         },
         onDrop: () => {
-          setListDndState(idle);
+          setActiveDropRegion(null);
         },
       })
     );
   }, [dragHandleRef, draggableRef, elementId]);
 
-  return [dndListState, isDragging] as const;
+  return [activeDropRegion, isDragging] as const;
 };

invokeai/frontend/web/src/features/nodes/store/workflowSlice.ts

@@ -400,6 +400,7 @@ const addElement = (args: {
     return;
   }
 
+  element.parentId = containerId;
   elements[element.id] = element;
 
   if (index === undefined) {

invokeai/frontend/web/src/features/nodes/types/workflow.ts

@@ -108,7 +108,7 @@ const zNodeFieldElement = zElementBase.extend({
 });
 export type NodeFieldElement = z.infer<typeof zNodeFieldElement>;
 export const isNodeFieldElement = (el: FormElement): el is NodeFieldElement => el.type === NODE_FIELD_TYPE;
-const nodeField = (
+export const buildNodeField = (
   nodeId: NodeFieldElement['data']['fieldIdentifier']['nodeId'],
   fieldName: NodeFieldElement['data']['fieldIdentifier']['fieldName'],
   parentId?: NodeFieldElement['parentId']
@@ -123,8 +123,8 @@ const nodeField = (
   };
   return element;
 };
-const _nodeField = (...args: Parameters<typeof nodeField>): NodeFieldElement => {
-  const element = nodeField(...args);
+const _nodeField = (...args: Parameters<typeof buildNodeField>): NodeFieldElement => {
+  const element = buildNodeField(...args);
   addElement(element);
   return element;
 };
@@ -140,7 +140,7 @@ const zHeadingElement = zElementBase.extend({
 });
 export type HeadingElement = z.infer<typeof zHeadingElement>;
 export const isHeadingElement = (el: FormElement): el is HeadingElement => el.type === HEADING_TYPE;
-const heading = (
+export const buildHeading = (
   content: HeadingElement['data']['content'],
   level: HeadingElement['data']['level'],
   parentId?: NodeFieldElement['parentId']
@@ -156,8 +156,8 @@ const heading = (
   };
   return element;
 };
-const _heading = (...args: Parameters<typeof heading>): HeadingElement => {
-  const element = heading(...args);
+const _heading = (...args: Parameters<typeof buildHeading>): HeadingElement => {
+  const element = buildHeading(...args);
   addElement(element);
   return element;
 };
@@ -173,7 +173,7 @@ const zTextElement = zElementBase.extend({
 });
 export type TextElement = z.infer<typeof zTextElement>;
 export const isTextElement = (el: FormElement): el is TextElement => el.type === TEXT_TYPE;
-const text = (
+export const buildText = (
   content: TextElement['data']['content'],
   fontSize: TextElement['data']['fontSize'],
   parentId?: NodeFieldElement['parentId']
@@ -187,11 +187,10 @@ const text = (
       fontSize,
     },
   };
-  addElement(element);
   return element;
 };
-const _text = (...args: Parameters<typeof text>): TextElement => {
-  const element = text(...args);
+const _text = (...args: Parameters<typeof buildText>): TextElement => {
+  const element = buildText(...args);
   addElement(element);
   return element;
 };
@@ -203,17 +202,16 @@ const zDividerElement = zElementBase.extend({
 });
 export type DividerElement = z.infer<typeof zDividerElement>;
 export const isDividerElement = (el: FormElement): el is DividerElement => el.type === DIVIDER_TYPE;
-const divider = (parentId?: NodeFieldElement['parentId']): DividerElement => {
+export const buildDivider = (parentId?: NodeFieldElement['parentId']): DividerElement => {
   const element: DividerElement = {
     id: getPrefixedId(DIVIDER_TYPE, '-'),
     parentId,
     type: DIVIDER_TYPE,
   };
-  addElement(element);
   return element;
 };
-const _divider = (...args: Parameters<typeof divider>): DividerElement => {
-  const element = divider(...args);
+const _divider = (...args: Parameters<typeof buildDivider>): DividerElement => {
+  const element = buildDivider(...args);
   addElement(element);
   return element;
 };
@@ -229,7 +227,7 @@ const zContainerElement = zElementBase.extend({
 });
 export type ContainerElement = z.infer<typeof zContainerElement>;
 export const isContainerElement = (el: FormElement): el is ContainerElement => el.type === CONTAINER_TYPE;
-export const container = (
+export const buildContainer = (
   direction: ContainerElement['data']['direction'],
   children: ContainerElement['data']['children'],
   parentId?: NodeFieldElement['parentId']
@@ -245,8 +243,8 @@ export const container = (
   };
   return element;
 };
-export const _container = (...args: Parameters<typeof container>): ContainerElement => {
-  const element = container(...args);
+export const _container = (...args: Parameters<typeof buildContainer>): ContainerElement => {
+  const element = buildContainer(...args);
   addElement(element);
   return element;
 };
@@ -262,21 +260,25 @@ export type FormElement = z.infer<typeof zFormElement>;
 //   _container('row', [_container('column', []).id, _container('column', []).id, _container('column', []).id]).id,
 // ]).id;
 
-const rootContainer = container('column', []);
+const rootContainer = buildContainer('column', []);
 addElement(rootContainer);
+const rowContainer = buildContainer('row', [], rootContainer.id);
+const rowContainerChildren = [
+  _nodeField('58e748ec-7405-4816-a5ff-c168ee35161a', 'value', rowContainer.id),
+  _nodeField('1b383334-2efc-406d-a000-49c4c5ebccde', 'value', rowContainer.id),
+  _nodeField('7ebda150-63c7-4ccd-a1d2-444459107393', 'value', rowContainer.id),
+];
+rowContainerChildren.forEach((child) => {
+  addElement(child);
+  rowContainer.data.children.push(child.id);
+});
+
 const children = [
-  heading('My Cool Workflow', 1, rootContainer.id),
-  text('This is a description of what my workflow does. It does things.', 'md', rootContainer.id),
-  divider(rootContainer.id),
-  heading('First Section', 2, rootContainer.id),
-  text(
-    'The first section includes fields relevant to the first section. This note describes that fact.',
-    'sm',
-    rootContainer.id
-  ),
-  divider(rootContainer.id),
-  text('These are some text that are definitely super helpful.', 'sm', rootContainer.id),
-  divider(rootContainer.id),
+  buildHeading('My Cool Workflow', 1, rootContainer.id),
+  buildText('This is a description of what my workflow does. It does things.', 'md', rootContainer.id),
+  buildDivider(rootContainer.id),
+  buildText('These are some text that are definitely super helpful.', 'sm', rootContainer.id),
+  rowContainer,
 ];
 children.forEach((child) => {
   addElement(child);