handle layouting

studio/frontend/package.json

@@ -17,6 +17,7 @@
     "@assistant-ui/react-markdown": "^0.12.1",
     "@assistant-ui/react-streamdown": "^0.1.0",
     "@base-ui/react": "^1.1.0",
+    "@dagrejs/dagre": "^2.0.3",
     "@fontsource-variable/figtree": "^5.2.10",
     "@fontsource-variable/inter": "^5.2.8",
     "@fontsource-variable/space-grotesk": "^5.2.10",

studio/frontend/src/features/canvas-lab/canvas-lab-page.tsx

@@ -6,6 +6,7 @@ import {
   type NodeTypes,
   Panel,
   ReactFlow,
+  useReactFlow,
 } from "@xyflow/react";
 import { type ReactElement, useCallback, useMemo, useState } from "react";
 import { useShallow } from "zustand/react/shallow";
@@ -29,6 +30,38 @@ import { buildCanvasPayload } from "./utils/payload";
 const NODE_TYPES: NodeTypes = { builder: CanvasNode };
 const EDGE_TYPES: EdgeTypes = { canvas: CanvasEdge };
 
+type LayoutControlsProps = {
+  direction: "LR" | "TB";
+  onLayout: () => void;
+  onToggleDirection: () => void;
+};
+
+function LayoutControls({
+  direction,
+  onLayout,
+  onToggleDirection,
+}: LayoutControlsProps): ReactElement {
+  const { fitView } = useReactFlow();
+
+  const handleLayout = useCallback(() => {
+    onLayout();
+    requestAnimationFrame(() => {
+      fitView({ duration: 250 });
+    });
+  }, [fitView, onLayout]);
+
+  return (
+    <Panel position="top-left" className="m-3 flex items-center gap-2">
+      <Button size="sm" variant="secondary" onClick={handleLayout}>
+        Auto layout
+      </Button>
+      <Button size="sm" variant="outline" onClick={onToggleDirection}>
+        {direction}
+      </Button>
+    </Panel>
+  );
+}
+
 export function CanvasLabPage(): ReactElement {
   const {
     nodes,
@@ -37,6 +70,7 @@ export function CanvasLabPage(): ReactElement {
     sheetView,
     activeConfigId,
     dialogOpen,
+    layoutDirection,
     onNodesChange,
     onEdgesChange,
     onConnect,
@@ -49,6 +83,8 @@ export function CanvasLabPage(): ReactElement {
     setSheetView,
     setDialogOpen,
     loadCanvas,
+    setLayoutDirection,
+    applyLayout,
   } = useCanvasLabStore(
     useShallow((state) => ({
       nodes: state.nodes,
@@ -57,6 +93,7 @@ export function CanvasLabPage(): ReactElement {
       sheetView: state.sheetView,
       activeConfigId: state.activeConfigId,
       dialogOpen: state.dialogOpen,
+      layoutDirection: state.layoutDirection,
       onNodesChange: state.onNodesChange,
       onEdgesChange: state.onEdgesChange,
       onConnect: state.onConnect,
@@ -69,6 +106,8 @@ export function CanvasLabPage(): ReactElement {
       setSheetView: state.setSheetView,
       setDialogOpen: state.setDialogOpen,
       loadCanvas: state.loadCanvas,
+      setLayoutDirection: state.setLayoutDirection,
+      applyLayout: state.applyLayout,
     })),
   );
   const [sheetContainer, setSheetContainer] = useState<HTMLDivElement | null>(
@@ -94,6 +133,10 @@ export function CanvasLabPage(): ReactElement {
     [configs],
   );
 
+  const handleToggleDirection = useCallback(() => {
+    setLayoutDirection(layoutDirection === "LR" ? "TB" : "LR");
+  }, [layoutDirection, setLayoutDirection]);
+
   const handlePreview = async (): Promise<void> => {
     setPreviewLoading(true);
     setStatusMessage(null);
@@ -244,6 +287,11 @@ export function CanvasLabPage(): ReactElement {
             fitView={true}
             className="h-full w-full"
           >
+            <LayoutControls
+              direction={layoutDirection}
+              onLayout={applyLayout}
+              onToggleDirection={handleToggleDirection}
+            />
             <Background
               variant={BackgroundVariant.Dots}
               gap={18}

studio/frontend/src/features/canvas-lab/components/canvas-node.tsx

@@ -6,8 +6,8 @@ import {
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import type { NodeProps } from "@xyflow/react";
-import { Handle, Position } from "@xyflow/react";
-import { type ReactElement, memo } from "react";
+import { Handle, Position, useUpdateNodeInternals } from "@xyflow/react";
+import { type ReactElement, memo, useEffect } from "react";
 import type { CanvasNode as CanvasNodeType } from "../types";
 
 const NODE_META = {
@@ -26,10 +26,19 @@ const NODE_META = {
 } as const;
 
 function CanvasNodeBase({
+  id,
   data,
   selected,
 }: NodeProps<CanvasNodeType>): ReactElement {
   const meta = NODE_META[data.kind];
+  const layoutDirection = data.layoutDirection ?? "LR";
+  const isHorizontal = layoutDirection === "LR";
+  const updateNodeInternals = useUpdateNodeInternals();
+
+  useEffect(() => {
+    updateNodeInternals(id);
+  }, [id, layoutDirection, updateNodeInternals]);
+
   return (
     <div
       className={cn(
@@ -57,12 +66,12 @@ function CanvasNodeBase({
       </div>
       <Handle
         type="target"
-        position={Position.Left}
+        position={isHorizontal ? Position.Left : Position.Top}
         className="size-2 border border-border bg-white"
       />
       <Handle
         type="source"
-        position={Position.Right}
+        position={isHorizontal ? Position.Right : Position.Bottom}
         className="size-2 border border-border bg-white"
       />
     </div>

studio/frontend/src/features/canvas-lab/stores/canvas-lab-helpers.ts

@@ -1,4 +1,4 @@
-import type { CanvasNode, NodeConfig } from "../types";
+import type { CanvasNode, LayoutDirection, NodeConfig } from "../types";
 import { nodeDataFromConfig } from "../utils";
 import { removeRef, replaceRef } from "../utils/refs";
 
@@ -22,9 +22,12 @@ export function updateNodeData(
   nodes: CanvasNode[],
   id: string,
   config: NodeConfig,
+  layoutDirection: LayoutDirection,
 ): CanvasNode[] {
   return nodes.map((node) =>
-    node.id === id ? { ...node, data: nodeDataFromConfig(config) } : node,
+    node.id === id
+      ? { ...node, data: nodeDataFromConfig(config, layoutDirection) }
+      : node,
   );
 }
 
@@ -41,12 +44,13 @@ export function findNodeIdByName(
 export function buildNodeUpdate(
   state: NodeUpdateState,
   config: NodeConfig,
+  layoutDirection: LayoutDirection,
 ): NodeUpdateResult {
   const node: CanvasNode = {
     id: config.id,
     type: "builder",
     position: { x: 0, y: state.nextY },
-    data: nodeDataFromConfig(config),
+    data: nodeDataFromConfig(config, layoutDirection),
   };
   return {
     configs: { ...state.configs, [config.id]: config },
@@ -58,6 +62,23 @@ export function buildNodeUpdate(
   };
 }
 
+export function applyLayoutDirectionToNodes(
+  nodes: CanvasNode[],
+  configs: Record<string, NodeConfig>,
+  layoutDirection: LayoutDirection,
+): CanvasNode[] {
+  return nodes.map((node) => {
+    const config = configs[node.id];
+    if (config) {
+      return { ...node, data: nodeDataFromConfig(config, layoutDirection) };
+    }
+    return {
+      ...node,
+      data: { ...node.data, layoutDirection },
+    };
+  });
+}
+
 function updateTemplateFields(
   config: NodeConfig,
   updater: (value: string) => string,
@@ -108,18 +129,12 @@ export function applyRenameToConfig(
     config.sampler_type === "subcategory" &&
     config.subcategory_parent === from
   ) {
-    next =
-      next === config
-        ? {
-            ...config,
-            // biome-ignore lint/style/useNamingConvention: api schema
-            subcategory_parent: to,
-          }
-        : {
-            ...next,
-            // biome-ignore lint/style/useNamingConvention: api schema
-            subcategory_parent: to,
-          };
+    const base = next === config ? config : next;
+    next = {
+      ...base,
+      // biome-ignore lint/style/useNamingConvention: api schema
+      subcategory_parent: to,
+    };
   }
   return next;
 }
@@ -134,22 +149,14 @@ export function applyRemovalToConfig(
     config.sampler_type === "subcategory" &&
     config.subcategory_parent === ref
   ) {
-    next =
-      next === config
-        ? {
-            ...config,
-            // biome-ignore lint/style/useNamingConvention: api schema
-            subcategory_parent: "",
-            // biome-ignore lint/style/useNamingConvention: api schema
-            subcategory_mapping: {},
-          }
-        : {
-            ...next,
-            // biome-ignore lint/style/useNamingConvention: api schema
-            subcategory_parent: "",
-            // biome-ignore lint/style/useNamingConvention: api schema
-            subcategory_mapping: {},
-          };
+    const base = next === config ? config : next;
+    next = {
+      ...base,
+      // biome-ignore lint/style/useNamingConvention: api schema
+      subcategory_parent: "",
+      // biome-ignore lint/style/useNamingConvention: api schema
+      subcategory_mapping: {},
+    };
   }
   return next;
 }

studio/frontend/src/features/canvas-lab/stores/canvas-lab.ts

@@ -11,6 +11,7 @@ import {
 import { create } from "zustand";
 import type {
   CanvasNode,
+  LayoutDirection,
   LlmType,
   NodeConfig,
   SamplerConfig,
@@ -20,10 +21,12 @@ import { getBlockDefinition } from "../blocks/registry";
 import { isCategoryConfig, isSubcategoryConfig } from "../utils";
 import { applyCanvasConnection, isValidCanvasConnection } from "../utils/graph";
 import type { CanvasSnapshot } from "../utils/import";
+import { getLayoutedElements } from "../utils/layout";
 import {
   applyRemovalToConfig,
   applyRemovalToConfigs,
   applyRenameToConfigs,
+  applyLayoutDirectionToNodes,
   buildNodeUpdate,
   findNodeIdByName,
   updateNodeData,
@@ -38,11 +41,14 @@ type CanvasLabState = {
   sheetView: SheetView;
   activeConfigId: string | null;
   dialogOpen: boolean;
+  layoutDirection: LayoutDirection;
   nextId: number;
   nextY: number;
   setSheetView: (view: SheetView) => void;
   setDialogOpen: (open: boolean) => void;
   openConfig: (id: string) => void;
+  setLayoutDirection: (direction: LayoutDirection) => void;
+  applyLayout: () => void;
   addSamplerNode: (type: SamplerType) => void;
   addLlmNode: (type: LlmType) => void;
   addExpressionNode: () => void;
@@ -61,11 +67,34 @@ export const useCanvasLabStore = create<CanvasLabState>((set, get) => ({
   sheetView: "root",
   activeConfigId: null,
   dialogOpen: false,
+  layoutDirection: "LR",
   nextId: 3,
   nextY: 280,
   setSheetView: (view) => set({ sheetView: view }),
   setDialogOpen: (open) => set({ dialogOpen: open }),
   openConfig: (id) => set({ activeConfigId: id, dialogOpen: true }),
+  setLayoutDirection: (direction) =>
+    set((state) => ({
+      layoutDirection: direction,
+      nodes: applyLayoutDirectionToNodes(
+        state.nodes,
+        state.configs,
+        direction,
+      ),
+    })),
+  applyLayout: () =>
+    set((state) => {
+      const { nodes } = getLayoutedElements(state.nodes, state.edges, {
+        direction: state.layoutDirection,
+      });
+      return {
+        nodes: applyLayoutDirectionToNodes(
+          nodes,
+          state.configs,
+          state.layoutDirection,
+        ),
+      };
+    }),
   addSamplerNode: (type) => {
     set((state) => {
       const id = `n${state.nextId}`;
@@ -75,7 +104,7 @@ export const useCanvasLabStore = create<CanvasLabState>((set, get) => ({
         return state;
       }
       const config = definition.createConfig(id, existing);
-      return buildNodeUpdate(state, config);
+      return buildNodeUpdate(state, config, state.layoutDirection);
     });
   },
   addLlmNode: (type) => {
@@ -87,7 +116,7 @@ export const useCanvasLabStore = create<CanvasLabState>((set, get) => ({
         return state;
       }
       const config = definition.createConfig(id, existing);
-      return buildNodeUpdate(state, config);
+      return buildNodeUpdate(state, config, state.layoutDirection);
     });
   },
   addExpressionNode: () => {
@@ -99,20 +128,24 @@ export const useCanvasLabStore = create<CanvasLabState>((set, get) => ({
         return state;
       }
       const config = definition.createConfig(id, existing);
-      return buildNodeUpdate(state, config);
+      return buildNodeUpdate(state, config, state.layoutDirection);
     });
   },
   loadCanvas: (snapshot) =>
-    set({
+    set((state) => ({
       configs: snapshot.configs,
-      nodes: snapshot.nodes,
+      nodes: applyLayoutDirectionToNodes(
+        snapshot.nodes,
+        snapshot.configs,
+        state.layoutDirection,
+      ),
       edges: snapshot.edges,
       nextId: snapshot.nextId,
       nextY: snapshot.nextY,
       activeConfigId: null,
       dialogOpen: false,
       sheetView: "root",
-    }),
+    })),
   updateConfig: (id, patch) => {
     // biome-ignore lint/complexity/noExcessiveCognitiveComplexity: store update
     const applyUpdate = (state: CanvasLabState) => {
@@ -128,7 +161,12 @@ export const useCanvasLabStore = create<CanvasLabState>((set, get) => ({
         ...state.configs,
         [id]: next,
       };
-      const nodes = updateNodeData(state.nodes, id, next);
+      const nodes = updateNodeData(
+        state.nodes,
+        id,
+        next,
+        state.layoutDirection,
+      );
       let edges = state.edges;
 
       const hasParentPatch = Object.prototype.hasOwnProperty.call(

studio/frontend/src/features/canvas-lab/types/index.ts

@@ -14,11 +14,14 @@ export type LlmType = "text" | "structured" | "code";
 
 export type ExpressionDtype = "str" | "int" | "float" | "bool";
 
+export type LayoutDirection = "LR" | "TB";
+
 export type CanvasNodeData = {
   title: string;
   name: string;
   kind: "sampler" | "llm" | "expression";
   subtype: string;
+  layoutDirection?: LayoutDirection;
 };
 
 export type CanvasNode = Node<CanvasNodeData, "builder">;

studio/frontend/src/features/canvas-lab/utils/index.ts

@@ -2,6 +2,7 @@ import type {
   CanvasNodeData,
   ExpressionConfig,
   ExpressionDtype,
+  LayoutDirection,
   LlmConfig,
   LlmType,
   NodeConfig,
@@ -224,13 +225,17 @@ export function labelForExpression(type: ExpressionDtype): string {
   return EXPRESSION_LABELS[type] ?? "Expression";
 }
 
-export function nodeDataFromConfig(config: NodeConfig): CanvasNodeData {
+export function nodeDataFromConfig(
+  config: NodeConfig,
+  layoutDirection: LayoutDirection = "LR",
+): CanvasNodeData {
   if (config.kind === "sampler") {
     return {
       title: "Sampler",
       kind: "sampler",
       subtype: labelForSampler(config.sampler_type),
       name: config.name,
+      layoutDirection,
     };
   }
   if (config.kind === "expression") {
@@ -239,6 +244,7 @@ export function nodeDataFromConfig(config: NodeConfig): CanvasNodeData {
       kind: "expression",
       subtype: labelForExpression(config.dtype),
       name: config.name,
+      layoutDirection,
     };
   }
   return {
@@ -246,6 +252,7 @@ export function nodeDataFromConfig(config: NodeConfig): CanvasNodeData {
     kind: "llm",
     subtype: labelForLlm(config.llm_type),
     name: config.name,
+    layoutDirection,
   };
 }
 

studio/frontend/src/features/canvas-lab/utils/layout.ts

@@ -0,0 +1,56 @@
+import dagre from "@dagrejs/dagre";
+import type { Edge, Node } from "@xyflow/react";
+import type { LayoutDirection } from "../types";
+
+type LayoutOptions = {
+  direction?: LayoutDirection;
+  nodesep?: number;
+  ranksep?: number;
+  nodeWidth?: number;
+  nodeHeight?: number;
+};
+
+export function getLayoutedElements<TNode extends Node>(
+  nodes: TNode[],
+  edges: Edge[],
+  options: LayoutOptions = {},
+): { nodes: TNode[]; edges: Edge[] } {
+  const {
+    direction = "LR",
+    nodesep = 80,
+    ranksep = 80,
+    nodeWidth = 220,
+    nodeHeight = 64,
+  } = options;
+
+  const graph = new dagre.graphlib.Graph();
+  graph.setDefaultEdgeLabel(() => ({}));
+  graph.setGraph({ rankdir: direction, nodesep, ranksep });
+
+  nodes.forEach((node) => {
+    const width = node.measured?.width ?? nodeWidth;
+    const height = node.measured?.height ?? nodeHeight;
+    graph.setNode(node.id, { width, height });
+  });
+
+  edges.forEach((edge) => {
+    graph.setEdge(edge.source, edge.target);
+  });
+
+  dagre.layout(graph);
+
+  const layoutedNodes = nodes.map((node) => {
+    const pos = graph.node(node.id);
+    const width = node.measured?.width ?? nodeWidth;
+    const height = node.measured?.height ?? nodeHeight;
+    return {
+      ...node,
+      position: {
+        x: pos.x - width / 2,
+        y: pos.y - height / 2,
+      },
+    };
+  });
+
+  return { nodes: layoutedNodes, edges };
+}

studio/frontend/vite.config.ts

@@ -6,12 +6,24 @@ import { defineConfig } from "vite";
 // https://vite.dev/config/
 export default defineConfig({
   plugins: [react(), tailwindcss()],
+  optimizeDeps: {
+    include: ["@dagrejs/dagre", "@dagrejs/graphlib"],
+  },
   server: {
     allowedHosts: ["playground.wasimhub.dev"],
   },
   resolve: {
     alias: {
       "@": path.resolve(__dirname, "./src"),
+      "@dagrejs/dagre": path.resolve(
+        __dirname,
+        "./node_modules/@dagrejs/dagre/dist/dagre.cjs.js",
+      ),
+    },
+  },
+  build: {
+    commonjsOptions: {
+      include: [/node_modules/, /@dagrejs\/dagre/, /@dagrejs\/graphlib/],
     },
   },
 });