diff --git a/src/areas/generate/components/WorkflowPanel.tsx b/src/areas/generate/components/WorkflowPanel.tsx index 5d6c14f..1767177 100644 --- a/src/areas/generate/components/WorkflowPanel.tsx +++ b/src/areas/generate/components/WorkflowPanel.tsx @@ -456,6 +456,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) @@ -493,7 +497,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]) @@ -659,7 +670,7 @@ export default function WorkflowPanel() { {workflow ? ( diff --git a/src/areas/workflows/WorkflowsPage.tsx b/src/areas/workflows/WorkflowsPage.tsx index 2f5c849..be337d0 100644 --- a/src/areas/workflows/WorkflowsPage.tsx +++ b/src/areas/workflows/WorkflowsPage.tsx @@ -14,7 +14,7 @@ import { type Edge, type OnConnectStartParams, } from '@xyflow/react' -import { useWorkflowsStore } from '@shared/stores/workflowsStore' +import { useWorkflowsStore, NODE_TYPES_WITHOUT_TARGET, NODE_TYPES_WITHOUT_SOURCE } from '@shared/stores/workflowsStore' import { useExtensionsStore } from '@shared/stores/extensionsStore' import { useNavStore } from '@shared/stores/navStore' import { useAppStore } from '@shared/stores/appStore' @@ -30,17 +30,30 @@ 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' } +// The While container whose bounds contain a flow-space point, if any. Used to +// auto-parent nodes dropped (or created) inside a While so they join its loop body. +function findWhileContainerAt(nodes: Node[], pos: { x: number; y: number }): Node | undefined { + return nodes.find((n) => { + if (n.type !== 'whileNode') return false + const gw = (n.measured?.width ?? n.width ?? (n.style?.width as number)) || 0 + const gh = (n.measured?.height ?? n.height ?? (n.style?.height as number)) || 0 + return pos.x >= n.position.x && pos.x <= n.position.x + gw + && pos.y >= n.position.y && pos.y <= n.position.y + gh + }) +} + // ─── IO badge ───────────────────────────────────────────────────────────────── const IO_STYLES: Record<'image' | 'text' | 'mesh', string> = { @@ -159,6 +172,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 +417,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 = @@ -937,9 +952,13 @@ function WorkflowCanvasInner({ pendingConnectionRef.current = null return } - // Dropped on empty canvas (not on a handle or node body) + // Dropped on empty canvas — or inside a While body — opens the palette. The + // While is a giant node, so don't treat its empty body as "dropped on a node"; + // bail only on a real node or a handle. const target = event.target as Element - if (target.closest('.react-flow__node') || target.closest('.react-flow__handle')) { + const nodeEl = target.closest('.react-flow__node') + const onWhile = nodeEl?.classList.contains('react-flow__node-whileNode') ?? false + if (target.closest('.react-flow__handle') || (nodeEl && !onWhile)) { pendingConnectionRef.current = null return } @@ -959,19 +978,34 @@ function WorkflowCanvasInner({ const nodeType = e.dataTransfer.getData(DRAG_NODE_KEY) if (nodeType) { - setNodes((nds) => [...nds, { - id: newId(), type: nodeType, position, - data: { enabled: true, params: {} } as WFNodeData, - }]) + const isContainer = nodeType === 'whileNode' + setNodes((nds) => { + const parent = isContainer ? undefined : findWhileContainerAt(nds, position) + const node: Node = { + id: newId(), type: nodeType, + position: parent ? { x: position.x - parent.position.x, y: position.y - parent.position.y } : position, + data: { enabled: true, params: {} } as WFNodeData, + ...(isContainer ? { style: { width: 420, height: 240 }, width: 420, height: 240 } : {}), + ...(parent ? { parentId: parent.id } : {}), + } + // Containers must sit before their future children in the array → prepend. + return isContainer ? [node, ...nds] : [...nds, node] + }) return } const extensionId = e.dataTransfer.getData(DRAG_KEY) if (!extensionId) return - setNodes((nds) => [...nds, { - id: newId(), type: 'extensionNode', position, - data: { extensionId, enabled: true, params: {} } as WFNodeData, - }]) + setNodes((nds) => { + const parent = findWhileContainerAt(nds, position) + const node: Node = { + id: newId(), type: 'extensionNode', + position: parent ? { x: position.x - parent.position.x, y: position.y - parent.position.y } : position, + data: { extensionId, enabled: true, params: {} } as WFNodeData, + ...(parent ? { parentId: parent.id } : {}), + } + return [...nds, node] + }) }, [screenToFlowPosition, setNodes]) // Keyboard shortcuts (Space, Ctrl+Z, Ctrl+Y / Ctrl+Shift+Z) @@ -1003,19 +1037,38 @@ 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, - }]) - - // If palette was opened from a connection drag, wire the edge automatically + const isContainer = type === 'whileNode' + setNodes((nds) => { + const parent = isContainer ? undefined : findWhileContainerAt(nds, position) + const node: Node = { + id: newNodeId, type, + position: parent ? { x: position.x - parent.position.x, y: position.y - parent.position.y } : position, + data: { extensionId, enabled: true, params: {} } as WFNodeData, + ...(isContainer ? { style: { width: 420, height: 240 }, width: 420, height: 240 } : {}), + ...(parent ? { parentId: parent.id } : {}), + } + // 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. + // ExtensionNodes use id'd handles (input-0 / output), not the default null + // handle, so the new node's side must reference them or React Flow can't place + // the edge ("Couldn't create edge for target handle id: null"). const pending = pendingConnectionRef.current if (pending?.nodeId) { const isSource = pending.handleType === 'source' - const edge = isSource - ? { id: newId(), source: pending.nodeId, sourceHandle: pending.handleId ?? undefined, target: newNodeId } - : { id: newId(), source: newNodeId, target: pending.nodeId, targetHandle: pending.handleId ?? undefined } - setEdges((eds) => addEdge({ ...edge, ...DEFAULT_EDGE_OPTS }, eds)) + const isExt = type === 'extensionNode' + // Skip wiring when the new node can't take the connection: a source-only node + // (Image/Text/Mesh) as target, or a sink-only node (Add to Scene/Preview) as + // source — those have no matching handle and would orphan the edge. + const canWire = isSource ? !NODE_TYPES_WITHOUT_TARGET.has(type) : !NODE_TYPES_WITHOUT_SOURCE.has(type) + if (canWire) { + const edge = isSource + ? { id: newId(), source: pending.nodeId, sourceHandle: pending.handleId ?? undefined, target: newNodeId, targetHandle: isExt ? 'input-0' : undefined } + : { id: newId(), source: newNodeId, sourceHandle: isExt ? 'output' : undefined, target: pending.nodeId, targetHandle: pending.handleId ?? undefined } + setEdges((eds) => addEdge({ ...edge, ...DEFAULT_EDGE_OPTS }, eds)) + } } pendingConnectionRef.current = null @@ -1023,6 +1076,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) { @@ -1218,6 +1333,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 cc38df9..b03dfb9 100644 --- a/src/areas/workflows/nodes/ExtensionNode.tsx +++ b/src/areas/workflows/nodes/ExtensionNode.tsx @@ -149,7 +149,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..f3a646a --- /dev/null +++ b/src/areas/workflows/nodes/WhileNode.tsx @@ -0,0 +1,108 @@ +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 continueWhile = useWorkflowRunStore((s) => s.continueWhile) + const retryWhile = useWorkflowRunStore((s) => s.retryWhile) + const progress = useWorkflowRunStore((s) => s.whileProgress[id]) + const isPaused = status === 'paused' && activeNodeId === id + const isRunning = status === 'running' && activeNodeId === id + // Lock the iteration count while a run is active so it can't change mid-loop. + const locked = status === 'running' || status === 'paused' + + 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 0f0341d..e253e39 100644 --- a/src/areas/workflows/workflowRunStore.ts +++ b/src/areas/workflows/workflowRunStore.ts @@ -29,6 +29,18 @@ const IDLE: WorkflowRunState = { const _cancel = { current: false } const _activeJobId = { current: null as string | null } +// While container (manual mode) pause/resume — set by continueWhile()/retryWhile(). +const _resume = { current: null as (() => void) | null } +const _retry = { current: false } +// Live node params — the UI pushes edits here so a looping/paused run re-reads the +// latest values when a body node starts (instead of the snapshot from run start). +const _liveParams = { current: new Map>() } + +function flushResume(): void { + const fn = _resume.current + _resume.current = null + if (fn) fn() +} interface NodeOutput { filePath?: string; text?: string; outputType?: string } @@ -85,6 +97,38 @@ function topoSort(nodes: WFNode[], edges: WFEdge[]): WFNode[] { return result } +// ─── While container geometry ────────────────────────────────────────────────── +// Body membership can't rely on parentId alone: React Flow only assigns it when a +// node is dragged into the container, so a While resized around existing nodes (or +// nodes added by palette click) leaves them unparented. We therefore also test +// on-canvas containment at run time. + +interface WhileBounds { x: number; y: number; w: number; h: number } + +function nodeSize(n: WFNode): { w: number; h: number } { + const measured = (n as { measured?: { width?: number; height?: number } }).measured + const styleW = n.style?.width + const styleH = n.style?.height + return { + w: measured?.width ?? n.width ?? (typeof styleW === 'number' ? styleW : 200), + h: measured?.height ?? n.height ?? (typeof styleH === 'number' ? styleH : 80), + } +} + +function whileBounds(w: WFNode): WhileBounds { + const s = nodeSize(w) + return { x: w.position.x, y: w.position.y, w: s.w, h: s.h } +} + +function isInsideWhile(n: WFNode, whileId: string, b: WhileBounds): boolean { + if (n.parentId === whileId) return true + if (n.parentId) return false // explicit child of another container + const s = nodeSize(n) + const cx = n.position.x + s.w / 2 + const cy = n.position.y + s.h / 2 + return cx >= b.x && cx <= b.x + b.w && cy >= b.y && cy <= b.y + b.h +} + // ─── Branch identification ──────────────────────────────────────────────────── // A node belongs to Wait W's branch if its single nearest upstream Wait is W // (dominance). Nodes with no upstream Wait — or with multiple (merges) — execute @@ -144,6 +188,9 @@ async function executeExtensionNode( selectedImagePath, selectedImageData } = ctx const ext = getWorkflowExtension(node.data.extensionId ?? '', allExtensions) + // Freshest params at the moment the node starts (so loop iterations / Retry pick + // up edits made while paused, not the values captured at run start). + const liveParams = _liveParams.current.get(node.id) ?? node.data.params ?? {} const resolveSource = (sourceId: string): NodeOutput | undefined => { const realId = resolveDataSource(sourceId, workflow.edges, nodeMap) @@ -198,7 +245,7 @@ async function executeExtensionNode( 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) @@ -252,7 +299,7 @@ async function executeExtensionNode( const result = await window.electron.extensions.runProcess( extId, { filePath: nodeInputPath, text: nodeInputText, nodeId: nid }, - node.data.params as Record, + liveParams as Record, ) if (!result.success) throw new Error(result.error ?? 'Process extension failed') nodeInputPath = result.result?.filePath ?? nodeInputPath @@ -320,11 +367,19 @@ interface WorkflowRunStore { nodeImageOutputs: Record waitStates: Record runningBranchId: string | null + /** whileId → current iteration / total (total null = manual/unbounded) */ + whileProgress: Record run: (workflow: Workflow, allExtensions: WorkflowExtension[], overrideImageData?: string) => Promise cancel: () => void reset: () => void continueRun: (waitId: string) => Promise + /** While container: resume past the loop (Continue) */ + continueWhile: () => void + /** While container: resume and re-run the loop body once more (Retry) */ + retryWhile: () => 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, get) => { @@ -379,6 +434,7 @@ export const useWorkflowRunStore = create((set, get) => { set((s) => ({ activeNodeId: null, runningBranchId: null, + whileProgress: {}, waitStates: finalWaitStates ?? s.waitStates, nodeImageOutputs: collectImageOutputs(ctx), runState: { @@ -401,16 +457,50 @@ export const useWorkflowRunStore = create((set, get) => { nodeImageOutputs: {}, waitStates: {}, runningBranchId: null, + whileProgress: {}, 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 const { preExecExtNodes, branches, waitIds, parentWait, ordered } = identifyBranches(workflow) const branchSteps = waitIds.reduce((acc, w) => acc + (branches.get(w)?.length ?? 0), 0) - const totalSteps = preExecExtNodes.length + branchSteps + + // ── While containers → loop over their pre-phase body members ───────────── + // A While's body = the extension nodes parented to it (parentId). They run in + // the pre-phase; we re-run those members either N times (data.iterations) or, + // in manual mode, on Continue/Retry. The body need not be contiguous in topo + // order, so re-runs filter by parentId (not an index range) to avoid replaying + // unrelated pre-phase nodes that happen to sort between body members. Body + // nodes gated behind a Wait land in a branch instead and are out of scope. + interface LoopInfo { whileId: string; firstIdx: number; lastIdx: number; bodyIds: Set; iterations: number | null } + const loops: LoopInfo[] = [] + for (const w of workflow.nodes) { + if (w.type !== 'whileNode') continue + const bounds = whileBounds(w) + const idxs = preExecExtNodes.reduce((acc, n, idx) => { + if (isInsideWhile(n, w.id, bounds)) 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), + bodyIds: new Set(idxs.map((i) => preExecExtNodes[i].id)), + iterations: Number.isFinite(iters) && iters > 0 ? Math.floor(iters) : null, + }) + } + const loopCounters = new Map(loops.map((l) => [l.whileId, l.iterations])) + // Auto-mode loops replay their body N times; count the extra passes so the + // progress total reflects the real work (manual loops stay unbounded). + const loopExtraSteps = loops.reduce((acc, l) => acc + (l.iterations != null ? (l.iterations - 1) * l.bodyIds.size : 0), 0) + const totalSteps = preExecExtNodes.length + branchSteps + loopExtraSteps const selectedImagePath = appState.selectedImagePath ?? undefined const selectedImageData = overrideImageData ?? appState.selectedImageData ?? undefined @@ -422,6 +512,7 @@ export const useWorkflowRunStore = create((set, get) => { // Top-level Waits are pending; nested Waits start blocked until their parent finishes. waitStates: Object.fromEntries(waitIds.map((id) => [id, parentWait.get(id) ? 'blocked' as WaitState : 'pending' as WaitState])), runningBranchId: null, + whileProgress: Object.fromEntries(loops.map((l) => [l.whileId, { current: 1, total: l.iterations }])), runState: { status: 'running', blockIndex: 0, blockTotal: totalSteps, blockProgress: 0, blockStep: 'Starting…', @@ -482,15 +573,73 @@ export const useWorkflowRunStore = create((set, get) => { } _ctx.current = ctx + // While the runner is replaying a loop body, this holds the active While id; + // re-iterations then execute only that While's members and skip everything + // else in the range. null = first pass / no active loop (run all nodes once). + let activeLoopId: string | null = null + + // End-of-body handler for While containers. Called after each pre-phase node; + // when the index is a loop's last body node, it either jumps back (auto N× or + // Retry) or pauses for Continue/Retry. Returns the index to resume at, + // 'cancel', or undefined to continue normally. + const bumpWhileProgress = (whileId: string): void => set((s) => { + const prev = s.whileProgress[whileId] + return { whileProgress: { ...s.whileProgress, [whileId]: { current: (prev?.current ?? 1) + 1, total: prev?.total ?? null } } } + } + ) + + const handleLoopEnd = async (idx: number): Promise => { + const loop = loops.find((l) => l.lastIdx === idx) + if (!loop) return undefined + const remaining = loopCounters.get(loop.whileId) + // Auto mode with iterations left → loop back automatically. + if (remaining != null && remaining > 1) { + loopCounters.set(loop.whileId, remaining - 1) + bumpWhileProgress(loop.whileId) + setRunState((s) => ({ ...s, blockStep: `Looping… ${remaining - 1} left` })) + activeLoopId = loop.whileId + return loop.firstIdx + } + // Otherwise the auto counter is exhausted (or it's manual mode): pause on + // the While and wait for Continue (proceed) or Retry (run the body again). + // runningBranchId blocks Wait branches while the pre-phase is parked here. + _retry.current = false + set({ activeNodeId: loop.whileId, runningBranchId: loop.whileId }) + setRunState((s) => ({ ...s, status: 'paused', blockStep: 'Loop finished — Continue or Retry' })) + await new Promise((resolve) => { _resume.current = resolve }) + if (_cancel.current) return 'cancel' + set({ runningBranchId: null }) + setRunState((s) => ({ ...s, status: 'running' })) + if (_retry.current) { + _retry.current = false + bumpWhileProgress(loop.whileId) + activeLoopId = loop.whileId + return loop.firstIdx + } + activeLoopId = null // Continue → resume normal forward execution + return undefined + } + // Pre-phase: nodes that don't belong to any single branch (sources + merges). + let stepsDone = 0 for (let i = 0; i < preExecExtNodes.length; i++) { if (_cancel.current) { _ctx.current = null; set({ runState: IDLE, activeNodeId: null }); return } const node = preExecExtNodes[i] + // During a loop replay, only re-run that While's own body members. + if (activeLoopId) { + const body = loops.find((l) => l.whileId === activeLoopId)?.bodyIds + if (body && !body.has(node.id)) continue + } set((s) => ({ activeNodeId: node.id, - runState: { ...s.runState, blockIndex: i, blockProgress: 0, blockStep: 'Starting…' }, + runState: { ...s.runState, blockIndex: stepsDone, blockProgress: 0, blockStep: 'Starting…' }, })) await executeExtensionNode(node, ctx, setRunState) + stepsDone++ + + const jump = await handleLoopEnd(i) + if (jump === 'cancel') { _ctx.current = null; set({ runState: IDLE, activeNodeId: null }); return } + if (jump !== undefined) { i = jump - 1 } } if (waitIds.length > 0) { @@ -604,19 +753,33 @@ export const useWorkflowRunStore = create((set, get) => { cancel() { _cancel.current = true + flushResume() // unblock a manual While pause so the run can tear down if (_activeJobId.current) { const apiUrl = useAppStore.getState().apiUrl axios.create({ baseURL: apiUrl }).post(`/generate/cancel/${_activeJobId.current}`).catch(() => {}) _activeJobId.current = null } _ctx.current = null - set({ runState: IDLE, activeNodeId: null, activeWorkflowId: null, nodeImageOutputs: {}, waitStates: {}, runningBranchId: null }) + set({ runState: IDLE, activeNodeId: null, activeWorkflowId: null, nodeImageOutputs: {}, waitStates: {}, runningBranchId: null, whileProgress: {} }) useAppStore.getState().setCurrentJob(null) }, reset() { _ctx.current = null - set({ runState: IDLE, activeNodeId: null, activeWorkflowId: null, nodeImageOutputs: {}, waitStates: {}, runningBranchId: null }) + set({ runState: IDLE, activeNodeId: null, activeWorkflowId: null, nodeImageOutputs: {}, waitStates: {}, runningBranchId: null, whileProgress: {} }) + }, + + continueWhile() { + flushResume() + }, + + retryWhile() { + _retry.current = true + flushResume() + }, + + setLiveNodeParams(nodeId, params) { + _liveParams.current.set(nodeId, params) }, } }) diff --git a/src/shared/stores/workflowsStore.ts b/src/shared/stores/workflowsStore.ts index 35ea60c..fc96ef4 100644 --- a/src/shared/stores/workflowsStore.ts +++ b/src/shared/stores/workflowsStore.ts @@ -35,10 +35,38 @@ interface LegacyWorkflow { updatedAt: string } +// Source-only nodes have no target handle; sink-only nodes have no source handle. +// An edge into/out of the wrong side can't resolve a handle and makes React Flow +// warn ("Couldn't create edge for target handle id: null") on every render. +export const NODE_TYPES_WITHOUT_TARGET = new Set(['imageNode', 'textNode', 'meshNode', 'inputNode']) +export const NODE_TYPES_WITHOUT_SOURCE = new Set(['outputNode', 'previewNode']) + +function sanitizeEdges(nodes: WFNode[], edges: WFEdge[]): WFEdge[] { + const typeOf = new Map(nodes.map((n) => [n.id, n.type])) + return edges + .filter((e) => { + const source = typeOf.get(e.source) + const target = typeOf.get(e.target) + if (!source || !target) return false // dangling endpoint + if (NODE_TYPES_WITHOUT_TARGET.has(target)) return false // target can't receive + if (NODE_TYPES_WITHOUT_SOURCE.has(source)) return false // source can't emit + return true + }) + .map((e) => { + // ExtensionNodes use id'd handles (input-0 / output). Repair edges that lost + // them (e.g. older palette auto-wiring) so React Flow can place them instead + // of warning "Couldn't create edge for target handle id: null". + const patch: Partial = {} + if (typeOf.get(e.target) === 'extensionNode' && e.targetHandle == null) patch.targetHandle = 'input-0' + if (typeOf.get(e.source) === 'extensionNode' && e.sourceHandle == null) patch.sourceHandle = 'output' + return Object.keys(patch).length ? { ...e, ...patch } : e + }) +} + function migrateWorkflow(raw: LegacyWorkflow): Workflow { // Already migrated if (raw.nodes && raw.edges) { - return { ...raw, nodes: raw.nodes, edges: raw.edges } as Workflow + return { ...raw, nodes: raw.nodes, edges: sanitizeEdges(raw.nodes, raw.edges) } as Workflow } // Migrate from old blocks format diff --git a/src/shared/types/electron.d.ts b/src/shared/types/electron.d.ts index 182ed32..03ac0e1 100644 --- a/src/shared/types/electron.d.ts +++ b/src/shared/types/electron.d.ts @@ -87,14 +87,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 {