fix(autolayout): subflow calculation (#2223)

* fix(autolayout): subflow calculation

* cleanup code

* fix missing import

* add back missing import

Author: icecrasher321

apps/sim/lib/workflows/autolayout/core.ts

@@ -17,13 +17,29 @@ import type { BlockState } from '@/stores/workflows/workflow/types'
 
 const logger = createLogger('AutoLayout:Core')
 
+/** Handle names that indicate edges from subflow end */
+const SUBFLOW_END_HANDLES = new Set(['loop-end-source', 'parallel-end-source'])
+
+/**
+ * Checks if an edge comes from a subflow end handle
+ */
+function isSubflowEndEdge(edge: Edge): boolean {
+  return edge.sourceHandle != null && SUBFLOW_END_HANDLES.has(edge.sourceHandle)
+}
+
 /**
  * Assigns layers (columns) to blocks using topological sort.
  * Blocks with no incoming edges are placed in layer 0.
+ * When edges come from subflow end handles, the subflow's internal depth is added.
+ *
+ * @param blocks - The blocks to assign layers to
+ * @param edges - The edges connecting blocks
+ * @param subflowDepths - Optional map of container block IDs to their internal depth (max layers inside)
  */
 export function assignLayers(
   blocks: Record<string, BlockState>,
-  edges: Edge[]
+  edges: Edge[],
+  subflowDepths?: Map<string, number>
 ): Map<string, GraphNode> {
   const nodes = new Map<string, GraphNode>()
 
@@ -40,6 +56,15 @@ export function assignLayers(
     })
   }
 
+  // Build a map of target node -> edges coming into it (to check sourceHandle later)
+  const incomingEdgesMap = new Map<string, Edge[]>()
+  for (const edge of edges) {
+    if (!incomingEdgesMap.has(edge.target)) {
+      incomingEdgesMap.set(edge.target, [])
+    }
+    incomingEdgesMap.get(edge.target)!.push(edge)
+  }
+
   // Build adjacency from edges
   for (const edge of edges) {
     const sourceNode = nodes.get(edge.source)
@@ -79,15 +104,33 @@ export function assignLayers(
     processed.add(nodeId)
 
     // Calculate layer based on max incoming layer + 1
+    // For edges from subflow ends, add the subflow's internal depth (minus 1 to avoid double-counting)
     if (node.incoming.size > 0) {
-      let maxIncomingLayer = -1
+      let maxEffectiveLayer = -1
+      const incomingEdges = incomingEdgesMap.get(nodeId) || []
+
       for (const incomingId of node.incoming) {
         const incomingNode = nodes.get(incomingId)
         if (incomingNode) {
-          maxIncomingLayer = Math.max(maxIncomingLayer, incomingNode.layer)
+          // Find edges from this incoming node to check if it's a subflow end edge
+          const edgesFromSource = incomingEdges.filter((e) => e.source === incomingId)
+          let additionalDepth = 0
+
+          // Check if any edge from this source is a subflow end edge
+          const hasSubflowEndEdge = edgesFromSource.some(isSubflowEndEdge)
+          if (hasSubflowEndEdge && subflowDepths) {
+            // Get the internal depth of the subflow
+            // Subtract 1 because the +1 at the end of layer calculation already accounts for one layer
+            // E.g., if subflow has 2 internal layers (depth=2), we add 1 extra so total offset is 2
+            const depth = subflowDepths.get(incomingId) ?? 1
+            additionalDepth = Math.max(0, depth - 1)
+          }
+
+          const effectiveLayer = incomingNode.layer + additionalDepth
+          maxEffectiveLayer = Math.max(maxEffectiveLayer, effectiveLayer)
         }
       }
-      node.layer = maxIncomingLayer + 1
+      node.layer = maxEffectiveLayer + 1
     }
 
     // Add outgoing nodes when all dependencies processed
@@ -254,12 +297,19 @@ export function calculatePositions(
  * 4. Calculate positions
  * 5. Normalize positions to start from padding
  *
+ * @param blocks - The blocks to lay out
+ * @param edges - The edges connecting blocks
+ * @param options - Layout options including container flag and subflow depths
  * @returns The laid-out nodes with updated positions, and bounding dimensions
  */
 export function layoutBlocksCore(
   blocks: Record<string, BlockState>,
   edges: Edge[],
-  options: { isContainer: boolean; layoutOptions?: LayoutOptions }
+  options: {
+    isContainer: boolean
+    layoutOptions?: LayoutOptions
+    subflowDepths?: Map<string, number>
+  }
 ): { nodes: Map<string, GraphNode>; dimensions: { width: number; height: number } } {
   if (Object.keys(blocks).length === 0) {
     return { nodes: new Map(), dimensions: { width: 0, height: 0 } }
@@ -269,8 +319,8 @@ export function layoutBlocksCore(
     options.layoutOptions ??
     (options.isContainer ? CONTAINER_LAYOUT_OPTIONS : DEFAULT_LAYOUT_OPTIONS)
 
-  // 1. Assign layers
-  const nodes = assignLayers(blocks, edges)
+  // 1. Assign layers (with subflow depth adjustment for subflow end edges)
+  const nodes = assignLayers(blocks, edges, options.subflowDepths)
 
   // 2. Prepare metrics
   prepareBlockMetrics(nodes)

apps/sim/lib/workflows/autolayout/index.ts

@@ -1,8 +1,12 @@
 import { createLogger } from '@/lib/logs/console/logger'
 import { layoutContainers } from '@/lib/workflows/autolayout/containers'
-import { layoutBlocksCore } from '@/lib/workflows/autolayout/core'
+import { assignLayers, layoutBlocksCore } from '@/lib/workflows/autolayout/core'
 import type { Edge, LayoutOptions, LayoutResult } from '@/lib/workflows/autolayout/types'
-import { filterLayoutEligibleBlockIds, getBlocksByParent } from '@/lib/workflows/autolayout/utils'
+import {
+  calculateSubflowDepths,
+  filterLayoutEligibleBlockIds,
+  getBlocksByParent,
+} from '@/lib/workflows/autolayout/utils'
 import type { BlockState } from '@/stores/workflows/workflow/types'
 
 const logger = createLogger('AutoLayout')
@@ -36,10 +40,15 @@ export function applyAutoLayout(
       (edge) => layoutRootIds.includes(edge.source) && layoutRootIds.includes(edge.target)
     )
 
+    // Calculate subflow depths before laying out root blocks
+    // This ensures blocks connected to subflow ends are positioned correctly
+    const subflowDepths = calculateSubflowDepths(blocksCopy, edges, assignLayers)
+
     if (Object.keys(rootBlocks).length > 0) {
       const { nodes } = layoutBlocksCore(rootBlocks, rootEdges, {
         isContainer: false,
         layoutOptions: options,
+        subflowDepths,
       })
 
       for (const node of nodes.values()) {

apps/sim/lib/workflows/autolayout/targeted.ts

@@ -4,9 +4,10 @@ import {
   DEFAULT_HORIZONTAL_SPACING,
   DEFAULT_VERTICAL_SPACING,
 } from '@/lib/workflows/autolayout/constants'
-import { layoutBlocksCore } from '@/lib/workflows/autolayout/core'
+import { assignLayers, layoutBlocksCore } from '@/lib/workflows/autolayout/core'
 import type { Edge, LayoutOptions } from '@/lib/workflows/autolayout/types'
 import {
+  calculateSubflowDepths,
   filterLayoutEligibleBlockIds,
   getBlockMetrics,
   getBlocksByParent,
@@ -48,7 +49,19 @@ export function applyTargetedLayout(
 
   const groups = getBlocksByParent(blocksCopy)
 
-  layoutGroup(null, groups.root, blocksCopy, edges, changedSet, verticalSpacing, horizontalSpacing)
+  // Calculate subflow depths before layout to properly position blocks after subflow ends
+  const subflowDepths = calculateSubflowDepths(blocksCopy, edges, assignLayers)
+
+  layoutGroup(
+    null,
+    groups.root,
+    blocksCopy,
+    edges,
+    changedSet,
+    verticalSpacing,
+    horizontalSpacing,
+    subflowDepths
+  )
 
   for (const [parentId, childIds] of groups.children.entries()) {
     layoutGroup(
@@ -58,7 +71,8 @@ export function applyTargetedLayout(
       edges,
       changedSet,
       verticalSpacing,
-      horizontalSpacing
+      horizontalSpacing,
+      subflowDepths
     )
   }
 
@@ -75,7 +89,8 @@ function layoutGroup(
   edges: Edge[],
   changedSet: Set<string>,
   verticalSpacing: number,
-  horizontalSpacing: number
+  horizontalSpacing: number,
+  subflowDepths: Map<string, number>
 ): void {
   if (childIds.length === 0) return
 
@@ -123,13 +138,15 @@ function layoutGroup(
   }
 
   // Compute layout positions using core function
+  // Only pass subflowDepths for root-level layout (not inside containers)
   const layoutPositions = computeLayoutPositions(
     layoutEligibleChildIds,
     blocks,
     edges,
     parentBlock,
     horizontalSpacing,
-    verticalSpacing
+    verticalSpacing,
+    parentId === null ? subflowDepths : undefined
   )
 
   if (layoutPositions.size === 0) {
@@ -177,7 +194,8 @@ function computeLayoutPositions(
   edges: Edge[],
   parentBlock: BlockState | undefined,
   horizontalSpacing: number,
-  verticalSpacing: number
+  verticalSpacing: number,
+  subflowDepths?: Map<string, number>
 ): Map<string, { x: number; y: number }> {
   const subsetBlocks: Record<string, BlockState> = {}
   for (const id of childIds) {
@@ -200,6 +218,7 @@ function computeLayoutPositions(
       verticalSpacing,
       alignment: 'center',
     },
+    subflowDepths,
   })
 
   // Update parent container dimensions if applicable

apps/sim/lib/workflows/autolayout/utils.ts

@@ -7,7 +7,7 @@ import {
   ROOT_PADDING_X,
   ROOT_PADDING_Y,
 } from '@/lib/workflows/autolayout/constants'
-import type { BlockMetrics, BoundingBox, GraphNode } from '@/lib/workflows/autolayout/types'
+import type { BlockMetrics, BoundingBox, Edge, GraphNode } from '@/lib/workflows/autolayout/types'
 import { BLOCK_DIMENSIONS, CONTAINER_DIMENSIONS } from '@/lib/workflows/blocks/block-dimensions'
 import type { BlockState } from '@/stores/workflows/workflow/types'
 
@@ -265,3 +265,53 @@ export function transferBlockHeights(
     }
   }
 }
+
+/**
+ * Calculates the internal depth (max layer count) for each subflow container.
+ * Used to properly position blocks that connect after a subflow ends.
+ *
+ * @param blocks - All blocks in the workflow
+ * @param edges - All edges in the workflow
+ * @param assignLayersFn - Function to assign layers to blocks
+ * @returns Map of container block IDs to their internal layer depth
+ */
+export function calculateSubflowDepths(
+  blocks: Record<string, BlockState>,
+  edges: Edge[],
+  assignLayersFn: (blocks: Record<string, BlockState>, edges: Edge[]) => Map<string, GraphNode>
+): Map<string, number> {
+  const depths = new Map<string, number>()
+  const { children } = getBlocksByParent(blocks)
+
+  for (const [containerId, childIds] of children.entries()) {
+    if (childIds.length === 0) {
+      depths.set(containerId, 1)
+      continue
+    }
+
+    const childBlocks: Record<string, BlockState> = {}
+    const layoutChildIds = filterLayoutEligibleBlockIds(childIds, blocks)
+    for (const childId of layoutChildIds) {
+      childBlocks[childId] = blocks[childId]
+    }
+
+    const childEdges = edges.filter(
+      (edge) => layoutChildIds.includes(edge.source) && layoutChildIds.includes(edge.target)
+    )
+
+    if (Object.keys(childBlocks).length === 0) {
+      depths.set(containerId, 1)
+      continue
+    }
+
+    const childNodes = assignLayersFn(childBlocks, childEdges)
+    let maxLayer = 0
+    for (const node of childNodes.values()) {
+      maxLayer = Math.max(maxLayer, node.layer)
+    }
+
+    depths.set(containerId, Math.max(maxLayer + 1, 1))
+  }
+
+  return depths
+}