From 4d99b57f4dd529f615a082be87fd1eb156f74e63 Mon Sep 17 00:00:00 2001 From: Lorchie Date: Tue, 16 Jun 2026 09:32:25 +0200 Subject: [PATCH] Add node while loop --- .../generate/components/WorkflowPanel.tsx | 15 ++- src/areas/workflows/WorkflowsPage.tsx | 83 +++++++++++++++- src/areas/workflows/nodes/ExtensionNode.tsx | 5 +- src/areas/workflows/nodes/WhileNode.tsx | 99 +++++++++++++++++++ src/areas/workflows/workflowRunStore.ts | 82 ++++++++++++++- src/shared/types/electron.d.ts | 15 ++- 6 files changed, 285 insertions(+), 14 deletions(-) create mode 100644 src/areas/workflows/nodes/WhileNode.tsx diff --git a/src/areas/generate/components/WorkflowPanel.tsx b/src/areas/generate/components/WorkflowPanel.tsx index e2d920b..84a0681 100644 --- a/src/areas/generate/components/WorkflowPanel.tsx +++ b/src/areas/generate/components/WorkflowPanel.tsx @@ -444,6 +444,10 @@ function EmbeddedCanvas({ workflow, allExtensions }: { setNodes((nds) => nds.map((n) => n.id === nodeId ? { ...n, data: { ...n.data, ...patch } } : n, )) + // Push params live so a paused/looping run uses the latest values on the next node start. + if (patch.params) { + useWorkflowRunStore.getState().setLiveNodeParams(nodeId, patch.params as Record) + } }, [setNodes]) const currentMeshUrl = useAppStore((s) => s.currentJob?.outputUrl) @@ -481,7 +485,14 @@ function EmbeddedCanvas({ workflow, allExtensions }: { showToast(firstPreflightIssue) return } - const wf: Workflow = { ...workflow, nodes: nodes as WFNode[], edges: edges as WFEdge[] } + // Persist the edited params so they survive remounts and are the values actually used. + const wf: Workflow = { + ...workflow, + nodes: nodes as WFNode[], + edges: edges as WFEdge[], + updatedAt: new Date().toISOString(), + } + useWorkflowsStore.getState().save(wf) run(wf, allExtensions) }, [firstPreflightIssue, nodes, edges, workflow, allExtensions, run, showToast]) @@ -647,7 +658,7 @@ export default function WorkflowPanel() { {workflow ? ( diff --git a/src/areas/workflows/WorkflowsPage.tsx b/src/areas/workflows/WorkflowsPage.tsx index 99f834c..b152e7e 100644 --- a/src/areas/workflows/WorkflowsPage.tsx +++ b/src/areas/workflows/WorkflowsPage.tsx @@ -30,13 +30,14 @@ import AddToSceneNode from './nodes/AddToSceneNode' import Load3DMeshNode from './nodes/Load3DMeshNode' import PreviewImageNode from './nodes/PreviewImageNode' import WaitNode from './nodes/WaitNode' +import WhileNode from './nodes/WhileNode' import WorkflowEdge from './nodes/WorkflowEdge' // ─── Constants ──────────────────────────────────────────────────────────────── const DRAG_KEY = 'modly/extension-id' const DRAG_NODE_KEY = 'modly/node-type' -const NODE_TYPES = { extensionNode: ExtensionNode, imageNode: ImageNode, textNode: TextNode, outputNode: AddToSceneNode, meshNode: Load3DMeshNode, previewNode: PreviewImageNode, waitNode: WaitNode } +const NODE_TYPES = { extensionNode: ExtensionNode, imageNode: ImageNode, textNode: TextNode, outputNode: AddToSceneNode, meshNode: Load3DMeshNode, previewNode: PreviewImageNode, waitNode: WaitNode, whileNode: WhileNode } const EDGE_TYPES = { workflowEdge: WorkflowEdge } const DEFAULT_EDGE_OPTS = { type: 'workflowEdge' } @@ -159,6 +160,7 @@ const PANEL_BUILTIN_NODES = [ { type: 'outputNode', label: 'Add to Scene', color: '#a78bfa', icon: <> }, { type: 'previewNode', label: 'Preview Views', color: '#38bdf8', icon: <> }, { type: 'waitNode', label: 'Wait', color: '#71717a', icon: <> }, + { type: 'whileNode', label: 'While', color: '#f59e0b', icon: <> }, ] function ExtGroupHeader({ title, author, expanded, onToggle, count }: { title: string; author?: string; expanded: boolean; onToggle: () => void; count: number }) { @@ -403,6 +405,7 @@ const BUILTIN_NODES = [ { type: 'outputNode', label: 'Add to Scene', color: '#a78bfa', description: 'Output node — adds the mesh to the 3D scene' }, { type: 'previewNode', label: 'Preview Views', color: '#38bdf8', description: 'Displays multi-view image outputs in a 2×3 grid' }, { type: 'waitNode', label: 'Wait', color: '#71717a', description: 'Pauses the workflow until you click Continue' }, + { type: 'whileNode', label: 'While', color: '#f59e0b', description: 'Container: wrap nodes to loop them N times or with Continue/Retry' }, ] type PaletteItem = @@ -990,10 +993,16 @@ function WorkflowCanvasInner({ pendingDropPos ?? { x: window.innerWidth / 2, y: window.innerHeight / 2 } ) const newNodeId = newId() - setNodes((nds) => [...nds, { - id: newNodeId, type, position, - data: { extensionId, enabled: true, params: {} } as WFNodeData, - }]) + const isContainer = type === 'whileNode' + setNodes((nds) => { + const node: Node = { + id: newNodeId, type, position, + data: { extensionId, enabled: true, params: {} } as WFNodeData, + ...(isContainer ? { style: { width: 420, height: 240 }, width: 420, height: 240 } : {}), + } + // Containers must sit before their future children in the array → prepend. + return isContainer ? [node, ...nds] : [...nds, node] + }) // If palette was opened from a connection drag, wire the edge automatically const pending = pendingConnectionRef.current @@ -1010,6 +1019,68 @@ function WorkflowCanvasInner({ setPaletteOpen(false) }, [screenToFlowPosition, setNodes, setEdges, pendingDropPos]) + // When a While container is deleted (button or keyboard), detach its children + // to absolute coordinates so they don't get orphaned to the canvas origin. + const onNodesDelete = useCallback((deleted: Node[]) => { + const removedContainers = deleted.filter((n) => n.type === 'whileNode') + if (removedContainers.length === 0) return + setNodes((nds) => nds.map((n) => { + const container = removedContainers.find((c) => c.id === n.parentId) + if (!container) return n + const { parentId: _p, extent: _ext, ...rest } = n + return { ...rest, position: { x: container.position.x + n.position.x, y: container.position.y + n.position.y } } + })) + }, [setNodes]) + + // When a node is dropped, attach/detach it to a While container based on overlap. + // Children get a parentId + parent-relative position (no extent, so they can be dragged back out). + const onNodeDragStop = useCallback((_e: unknown, dragged: Node) => { + if (dragged.type === 'whileNode') return + setNodes((nds) => { + const containers = nds.filter((n) => n.type === 'whileNode') + if (containers.length === 0 && !dragged.parentId) return nds + + const parent = dragged.parentId ? nds.find((n) => n.id === dragged.parentId) : undefined + const absX = (parent?.position.x ?? 0) + dragged.position.x + const absY = (parent?.position.y ?? 0) + dragged.position.y + const w = dragged.measured?.width ?? dragged.width ?? 200 + const h = dragged.measured?.height ?? dragged.height ?? 80 + const cx = absX + w / 2 + const cy = absY + h / 2 + + const container = containers.find((g) => { + const gw = (g.measured?.width ?? g.width ?? (g.style?.width as number)) || 0 + const gh = (g.measured?.height ?? g.height ?? (g.style?.height as number)) || 0 + return cx >= g.position.x && cx <= g.position.x + gw && cy >= g.position.y && cy <= g.position.y + gh + }) + + const newParentId = container?.id + if (newParentId === dragged.parentId) return nds // no change + + let next: Node[] = nds.map((n) => { + if (n.id !== dragged.id) return n + if (container) { + // parentId (no extent) → child moves with the container but can still be dragged out + return { ...n, parentId: container.id, + position: { x: absX - container.position.x, y: absY - container.position.y } } + } + const { parentId: _p, extent: _ext, ...rest } = n // detach + return { ...rest, position: { x: absX, y: absY } } + }) + + // ReactFlow requires the parent to appear before its child in the array. + if (newParentId) { + const cIdx = next.findIndex((n) => n.id === dragged.id) + const pIdx = next.findIndex((n) => n.id === newParentId) + if (pIdx > cIdx) { + const [child] = next.splice(cIdx, 1) + next.splice(next.findIndex((n) => n.id === newParentId) + 1, 0, child) + } + } + return next + }) + }, [setNodes]) + const handleRun = useCallback(() => { if (isRunning) { cancel(); return } if (preflightIssues.length > 0) { @@ -1205,6 +1276,8 @@ function WorkflowCanvasInner({ nodeTypes={NODE_TYPES} edgeTypes={EDGE_TYPES} onNodesChange={onNodesChange} + onNodeDragStop={onNodeDragStop} + onNodesDelete={onNodesDelete} onEdgesChange={onEdgesChange} onConnectStart={onConnectStart} onConnect={onConnect} diff --git a/src/areas/workflows/nodes/ExtensionNode.tsx b/src/areas/workflows/nodes/ExtensionNode.tsx index 7f9e2ff..760e885 100644 --- a/src/areas/workflows/nodes/ExtensionNode.tsx +++ b/src/areas/workflows/nodes/ExtensionNode.tsx @@ -157,7 +157,10 @@ export default function ExtensionNode({ id, data, selected }: { id: string; data }, [isMulti]) const patchParam = useCallback((key: string, val: number | string) => { - updateNodeData(id, { params: { ...data.params, [key]: val } }) + const params = { ...data.params, [key]: val } + updateNodeData(id, { params }) + // Push live so a paused/looping run picks up the change on the next node start. + useWorkflowRunStore.getState().setLiveNodeParams(id, params) }, [id, data.params, updateNodeData]) const paramById = new Map(ext?.params.map((p) => [p.id, p])) diff --git a/src/areas/workflows/nodes/WhileNode.tsx b/src/areas/workflows/nodes/WhileNode.tsx new file mode 100644 index 0000000..3850d5a --- /dev/null +++ b/src/areas/workflows/nodes/WhileNode.tsx @@ -0,0 +1,99 @@ +import { NodeResizer, useReactFlow } from '@xyflow/react' +import type { WFNodeData } from '@shared/types/electron.d' +import { useWorkflowRunStore } from '../workflowRunStore' + +// While container — a resizable frame that wraps the loop-body nodes. +// Drop nodes inside it (they become children) to define the body; the runner +// re-runs that body either N times (iterations) or via the manual buttons. +export default function WhileNode({ id, data, selected }: { id: string; data: WFNodeData; selected?: boolean }) { + const { updateNodeData, deleteElements } = useReactFlow() + const status = useWorkflowRunStore((s) => s.runState.status) + const activeNodeId = useWorkflowRunStore((s) => s.activeNodeId) + const continueRun = useWorkflowRunStore((s) => s.continueRun) + const retryRun = useWorkflowRunStore((s) => s.retryRun) + const isPaused = status === 'paused' && activeNodeId === id + const isRunning = status === 'running' && activeNodeId === id + + return ( +
+ + + {/* ── Header bar (drag handle) ─────────────────────────────────────── */} +
+ + + + + + + While + + {/* Iterations counter — 0 / empty = manual only */} + + +
+ + {/* Continue / Retry — only while paused on this node */} + {isPaused && ( +
+ + +
+ )} + + {/* Delete */} + +
+ + {/* ── Body drop zone (children render here, on top) ────────────────── */} +
+ + Drop nodes here to loop them + +
+
+ ) +} diff --git a/src/areas/workflows/workflowRunStore.ts b/src/areas/workflows/workflowRunStore.ts index e7040fe..a4e154c 100644 --- a/src/areas/workflows/workflowRunStore.ts +++ b/src/areas/workflows/workflowRunStore.ts @@ -26,6 +26,11 @@ const IDLE: WorkflowRunState = { const _cancel = { current: false } const _activeJobId = { current: null as string | null } const _resume = { current: null as (() => void) | null } +// Set by retryRun() before resuming so the paused While node knows to loop back +const _retry = { current: false } +// Live node params — the UI pushes edits here so a paused/looping run re-reads +// the latest values when a node starts (instead of the snapshot from run start). +const _liveParams = { current: new Map>() } function flushResume(): void { const fn = _resume.current @@ -72,6 +77,9 @@ interface WorkflowRunStore { cancel: () => void reset: () => void continueRun: () => void + retryRun: () => void + /** UI → runner: push the latest params for a node so a looping run uses them */ + setLiveNodeParams: (nodeId: string, params: Record) => void } export const useWorkflowRunStore = create((set) => ({ @@ -82,6 +90,8 @@ export const useWorkflowRunStore = create((set) => ({ async run(workflow, allExtensions, overrideImageData?) { _cancel.current = false + // Seed live params from the snapshot; UI edits during the run override these. + _liveParams.current = new Map(workflow.nodes.map((n) => [n.id, { ...(n.data.params ?? {}) }])) const appState = useAppStore.getState() const apiUrl = appState.apiUrl @@ -90,6 +100,28 @@ export const useWorkflowRunStore = create((set) => ({ (n.type === 'extensionNode' || n.type === 'waitNode') && n.data.enabled, ) + // ── While containers → loop ranges over the body nodes they wrap ────── + // A While node's body = the exec nodes parented to it. We re-run that + // contiguous-ish range either N times (data.iterations) or on Retry. + interface LoopInfo { whileId: string; firstIdx: number; lastIdx: number; iterations: number | null } + const loops: LoopInfo[] = [] + for (const w of workflow.nodes) { + if (w.type !== 'whileNode') continue + const idxs = execNodes.reduce((acc, n, idx) => { + if (n.parentId === w.id) acc.push(idx) + return acc + }, []) + if (idxs.length === 0) continue + const iters = Number(w.data?.iterations) + loops.push({ + whileId: w.id, + firstIdx: Math.min(...idxs), + lastIdx: Math.max(...idxs), + iterations: Number.isFinite(iters) && iters > 0 ? Math.floor(iters) : null, + }) + } + const loopCounters = new Map(loops.map((l) => [l.whileId, l.iterations])) + const selectedImagePath = appState.selectedImagePath ?? undefined const selectedImageData = overrideImageData ?? appState.selectedImageData ?? undefined const currentMeshUrl = appState.currentJob?.outputUrl @@ -154,11 +186,38 @@ export const useWorkflowRunStore = create((set) => ({ } } + // End-of-body handler for While containers. Runs after *any* node type + // (including Wait) so a Wait placed last in a loop body still re-loops. + // Returns the index to jump back to, 'cancel', or undefined to continue. + const handleLoopEnd = async (idx: number): Promise => { + const loop = loops.find((l) => l.lastIdx === idx) + if (!loop) return undefined + const remaining = loopCounters.get(loop.whileId) + if (remaining != null && remaining > 1) { + loopCounters.set(loop.whileId, remaining - 1) + set((s) => ({ runState: { ...s.runState, blockStep: `Looping… ${remaining - 1} left` } })) + return loop.firstIdx + } + if (remaining == null) { + _retry.current = false + set({ activeNodeId: loop.whileId }) + set((s) => ({ runState: { ...s.runState, status: 'paused', blockStep: 'Paused — Continue or Retry' } })) + await new Promise((resolve) => { _resume.current = resolve }) + if (_cancel.current) return 'cancel' + set((s) => ({ runState: { ...s.runState, status: 'running' } })) + if (_retry.current) { _retry.current = false; return loop.firstIdx } + } + return undefined // counter exhausted → exit loop + } + for (let i = 0; i < execNodes.length; i++) { if (_cancel.current) { set({ runState: IDLE, activeNodeId: null }); return } const node = execNodes[i] const ext = getWorkflowExtension(node.data.extensionId ?? '', allExtensions) + // Read the freshest params at the moment the node starts (so Retry/loops + // pick up edits made while paused, not the values from run start). + const liveParams = _liveParams.current.get(node.id) ?? node.data.params ?? {} // ── Resolve inputs ──────────────────────────────────────────────── let nodeInputPath: string | undefined @@ -203,6 +262,11 @@ export const useWorkflowRunStore = create((set) => ({ outputType: incomingEdges[0] ? nodeOutputs.get(incomingEdges[0].source)?.outputType : undefined, }) set((s) => ({ runState: { ...s.runState, status: 'running' } })) + + // A Wait can be the last node of a While body → still run the loop check. + const jump = await handleLoopEnd(i) + if (jump === 'cancel') { set({ runState: IDLE, activeNodeId: null }); return } + if (jump !== undefined) { i = jump - 1 } continue } @@ -237,7 +301,7 @@ export const useWorkflowRunStore = create((set) => ({ const schemaDefaults = Object.fromEntries( (ext.params ?? []).map((p) => [p.id, p.default]), ) - const effectiveParams = { ...schemaDefaults, ...(node.data.params ?? {}) } + const effectiveParams = { ...schemaDefaults, ...liveParams } const fd = new FormData() fd.append('image', blob, fname) @@ -305,7 +369,7 @@ export const useWorkflowRunStore = create((set) => ({ const result = await window.electron.extensions.runProcess( extId, { filePath: nodeInputPath, text: nodeInputText, nodeId }, - node.data.params as Record, + liveParams as Record, ) if (!result.success) throw new Error(result.error ?? 'Process extension failed') nodeInputPath = result.result?.filePath ?? nodeInputPath @@ -330,6 +394,11 @@ export const useWorkflowRunStore = create((set) => ({ outputUrl: `/workspace/${norm.slice(workspaceDir.length).replace(/^\//, '')}`, }) } + + // ── While container → end-of-body: auto-loop, or pause for Continue/Retry ── + const jump = await handleLoopEnd(i) + if (jump === 'cancel') { set({ runState: IDLE, activeNodeId: null }); return } + if (jump !== undefined) { i = jump - 1; continue } } // ── Collect image outputs for preview nodes ─────────────────────── @@ -418,4 +487,13 @@ export const useWorkflowRunStore = create((set) => ({ continueRun() { flushResume() }, + + retryRun() { + _retry.current = true + flushResume() + }, + + setLiveNodeParams(nodeId, params) { + _liveParams.current.set(nodeId, params) + }, })) diff --git a/src/shared/types/electron.d.ts b/src/shared/types/electron.d.ts index 94c42fd..cbb9d09 100644 --- a/src/shared/types/electron.d.ts +++ b/src/shared/types/electron.d.ts @@ -79,14 +79,21 @@ export interface WFNodeData { inputType?: 'image' | 'text' enabled: boolean showInGenerate?: boolean + iterations?: number // While container: auto-loop N times (omit/0 = manual only) params: Record } export interface WFNode { - id: string - type: string - position: { x: number; y: number } - data: WFNodeData + id: string + type: string + position: { x: number; y: number } + data: WFNodeData + // Sub-flow / group support (While container) + parentId?: string + extent?: 'parent' + width?: number + height?: number + style?: Record } export interface WFEdge {