diff --git a/api/main.py b/api/main.py index 586763b..1404bec 100644 --- a/api/main.py +++ b/api/main.py @@ -31,7 +31,7 @@ def filter(self, record): app = FastAPI( title="Modly API", - version="0.3.5", + version="0.3.6", lifespan=lifespan, ) diff --git a/api/routers/generation.py b/api/routers/generation.py index 2c4eb18..bdd0492 100644 --- a/api/routers/generation.py +++ b/api/routers/generation.py @@ -1,6 +1,7 @@ import asyncio import json import threading +import time import traceback import uuid from typing import Dict @@ -16,6 +17,19 @@ _jobs: Dict[str, JobStatus] = {} _cancelled: set = set() _cancel_events: Dict[str, threading.Event] = {} +_completed_at: Dict[str, float] = {} + +_JOB_TTL = 1800 # purge terminal jobs after 30 minutes + + +def _purge_old_jobs() -> None: + cutoff = time.monotonic() - _JOB_TTL + stale = [jid for jid, t in _completed_at.items() if t < cutoff] + for jid in stale: + _jobs.pop(jid, None) + _cancelled.discard(jid) + _cancel_events.pop(jid, None) + _completed_at.pop(jid, None) @router.post("/from-image") @@ -63,6 +77,8 @@ async def generate_from_image( **model_params, } + _purge_old_jobs() + job = JobStatus(job_id=job_id, status="pending", progress=0) _jobs[job_id] = job _cancel_events[job_id] = threading.Event() @@ -91,6 +107,7 @@ async def cancel_job(job_id: str): _cancel_events[job_id].set() if job.status in ("pending", "running"): job.status = "cancelled" + _completed_at[job_id] = time.monotonic() # Kill the active generator subprocess immediately so inference stops now. # _run_generation will catch the resulting exception, see job_id in _cancelled, # and return cleanly without setting an error status. @@ -162,6 +179,7 @@ def progress_cb(pct: int, step: str = "") -> None: job.status = "done" job.progress = 100 + _completed_at[job_id] = time.monotonic() try: rel = output_path.relative_to(WORKSPACE_DIR) job.output_url = f"/workspace/{rel.as_posix()}" @@ -170,6 +188,7 @@ def progress_cb(pct: int, step: str = "") -> None: except GenerationCancelled: job.status = "cancelled" + _completed_at[job_id] = time.monotonic() except Exception as exc: if job_id in _cancelled: return @@ -177,3 +196,4 @@ def progress_cb(pct: int, step: str = "") -> None: print(f"[Generation ERROR] {exc}\n{tb}") job.status = "error" job.error = tb.strip() + _completed_at[job_id] = time.monotonic() diff --git a/electron/main/extension-path-guard.test.ts b/electron/main/extension-path-guard.test.ts new file mode 100644 index 0000000..91d9170 --- /dev/null +++ b/electron/main/extension-path-guard.test.ts @@ -0,0 +1,52 @@ +import assert from 'node:assert/strict' +import path from 'node:path' +import test from 'node:test' + +async function loadGuard() { + return import(new URL('./extension-path-guard.ts', import.meta.url).href) +} + +test('assertSafeExtensionId accepts conservative extension ids', async () => { + const { assertSafeExtensionId } = await loadGuard() + assert.equal(assertSafeExtensionId('mesh-process'), 'mesh-process') + assert.equal(assertSafeExtensionId('image_model.v2'), 'image_model.v2') + assert.equal(assertSafeExtensionId('kimodo-soma-rp'), 'kimodo-soma-rp') +}) + +test('assertSafeExtensionId rejects empty, traversal and absolute ids', async () => { + const { assertSafeExtensionId } = await loadGuard() + assert.throws(() => assertSafeExtensionId(''), /must not be empty/i) + assert.throws(() => assertSafeExtensionId('.'), /is invalid/i) + assert.throws(() => assertSafeExtensionId('..'), /is invalid/i) + assert.throws(() => assertSafeExtensionId('../escape'), /path separators/i) + assert.throws(() => assertSafeExtensionId('..\\escape'), /path separators/i) + assert.throws(() => assertSafeExtensionId('/abs/path'), /absolute path|path separators/i) +}) + +test('assertSafeExtensionId rejects uppercase and unsupported characters', async () => { + const { assertSafeExtensionId } = await loadGuard() + assert.throws(() => assertSafeExtensionId('Mesh-Process'), /must match/i) + assert.throws(() => assertSafeExtensionId('bad id'), /must match/i) + assert.throws(() => assertSafeExtensionId('bad:id'), /must match/i) +}) + +test('resolveExtensionPathWithinRoot confines paths to root', async () => { + const { resolveExtensionPathWithinRoot } = await loadGuard() + const root = path.join('/tmp', 'extensions-root') + assert.equal(resolveExtensionPathWithinRoot(root, 'mesh-process'), path.resolve(root, 'mesh-process')) + assert.throws(() => resolveExtensionPathWithinRoot(root, '../escape'), /path separators/i) +}) + +test('resolvePathWithinRoot rejects canonical escapes', async () => { + const { resolvePathWithinRoot } = await loadGuard() + const root = path.join('/tmp', 'extensions-root') + assert.equal(resolvePathWithinRoot(root, '.modly-backup-safe-123'), path.resolve(root, '.modly-backup-safe-123')) + assert.throws(() => resolvePathWithinRoot(root, '../escape'), /escapes root/i) +}) + +test('buildExtensionBackupPath stays within root and rejects unsafe ids', async () => { + const { buildExtensionBackupPath } = await loadGuard() + const root = path.join('/tmp', 'extensions-root') + assert.equal(buildExtensionBackupPath(root, 'mesh-process', '123'), path.resolve(root, '.modly-backup-mesh-process-123')) + assert.throws(() => buildExtensionBackupPath(root, '../escape', '123'), /path separators/i) +}) diff --git a/electron/main/extension-path-guard.ts b/electron/main/extension-path-guard.ts new file mode 100644 index 0000000..717dbf3 --- /dev/null +++ b/electron/main/extension-path-guard.ts @@ -0,0 +1,53 @@ +import { isAbsolute, relative, resolve as resolvePath } from 'node:path' + +const EXTENSION_ID_PATTERN = /^[a-z0-9][a-z0-9._-]*$/ + +export function assertSafeExtensionId(extensionId: unknown): string { + if (typeof extensionId !== 'string') { + throw new Error('Extension id must be a string') + } + + const trimmed = extensionId.trim() + if (!trimmed) { + throw new Error('Extension id must not be empty') + } + + if (trimmed === '.' || trimmed === '..') { + throw new Error(`Extension id "${extensionId}" is invalid`) + } + + if (isAbsolute(trimmed)) { + throw new Error(`Extension id "${extensionId}" must not be an absolute path`) + } + + if (trimmed.includes('/') || trimmed.includes('\\')) { + throw new Error(`Extension id "${extensionId}" must not contain path separators`) + } + + if (!EXTENSION_ID_PATTERN.test(trimmed)) { + throw new Error(`Extension id "${extensionId}" must match ${EXTENSION_ID_PATTERN}`) + } + + return trimmed +} + +export function resolvePathWithinRoot(rootDir: string, unsafeLeaf: string): string { + const resolvedRoot = resolvePath(rootDir) + const resolvedCandidate = resolvePath(resolvedRoot, unsafeLeaf) + const normalizedRelative = relative(resolvedRoot, resolvedCandidate).replace(/\\/g, '/') + + if (normalizedRelative === '..' || normalizedRelative.startsWith('../') || isAbsolute(normalizedRelative)) { + throw new Error(`Resolved path escapes root: ${unsafeLeaf}`) + } + + return resolvedCandidate +} + +export function resolveExtensionPathWithinRoot(rootDir: string, extensionId: unknown): string { + return resolvePathWithinRoot(rootDir, assertSafeExtensionId(extensionId)) +} + +export function buildExtensionBackupPath(rootDir: string, extensionId: unknown, suffix: string): string { + const safeId = assertSafeExtensionId(extensionId) + return resolvePathWithinRoot(rootDir, `.modly-backup-${safeId}-${suffix}`) +} diff --git a/electron/main/ipc-handlers.ts b/electron/main/ipc-handlers.ts index 24cb9df..757b566 100644 --- a/electron/main/ipc-handlers.ts +++ b/electron/main/ipc-handlers.ts @@ -18,6 +18,7 @@ import { logger } from './logger' import { getProcessRunner, getPythonProcessRunner, getExtPythonExe, terminateProcessRunner, terminateAllProcessRunners } from './process-runner' import { getBuiltinExtensionsDir } from './builtin-sync' import { spawn } from 'child_process' +import { assertSafeExtensionId, buildExtensionBackupPath, resolveExtensionPathWithinRoot } from './extension-path-guard' type WindowGetter = () => BrowserWindow | null @@ -674,6 +675,8 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe const manifest = JSON.parse(manifestRaw) as ParsedManifest if (!manifest.id) throw new Error('manifest.json: required field "id" missing') + const extensionId = assertSafeExtensionId(manifest.id) + manifest.id = extensionId if (!manifest.nodes?.length) throw new Error('manifest.json: required field "nodes" missing or empty') const isProcess = manifest.type === 'process' @@ -698,10 +701,13 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe // 5. Copy to extensions directory (overwrite if already present) const extensionsDir = getSettings(app.getPath('userData')).extensionsDir await mkdir(extensionsDir, { recursive: true }) - const destDir = join(extensionsDir, manifest.id) + const destDir = resolveExtensionPathWithinRoot(extensionsDir, extensionId) + const backupDir = existsSync(destDir) ? buildExtensionBackupPath(extensionsDir, extensionId, String(Date.now())) : null - if (existsSync(destDir)) { - await rmAsync(destDir, { recursive: true, force: true }) + try { + if (backupDir) { + terminateProcessRunner(extensionId) + await rename(destDir, backupDir) } await cp(extractDir, destDir, { recursive: true }) @@ -726,15 +732,10 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe if (existsSync(join(destDir, 'setup.py'))) { emit({ step: 'setting_up', message: 'Setting up Python environment…' }) const { sm: gpuSm, cudaVersion } = await detectGpuInfo() - try { - await runExtensionSetup(destDir, gpuSm, cudaVersion, (line) => { - logger.info(`[ext-setup] ${line}`) - emit({ step: 'setting_up', message: line }) - }) - } catch (err) { - logger.warn(`[ext-setup] setup.py failed: ${err}`) - emit({ step: 'setting_up', message: `Warning: setup failed — ${err}` }) - } + await runExtensionSetup(destDir, gpuSm, cudaVersion, (line) => { + logger.info(`[ext-setup] ${line}`) + emit({ step: 'setting_up', message: line }) + }) } } else if (isProcess) { // 6b. JS process extension: npm install if package.json present @@ -767,14 +768,10 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe if (existsSync(join(destDir, 'setup.py'))) { emit({ step: 'setting_up', message: 'Setting up Python environment…' }) const { sm: gpuSm, cudaVersion } = await detectGpuInfo() - try { - await runExtensionSetup(destDir, gpuSm, cudaVersion, (line) => { - logger.info(`[ext-setup] ${line}`) - emit({ step: 'setting_up', message: line }) - }) - } catch (setupErr: any) { - throw new Error(`Extension setup failed: ${setupErr?.message ?? setupErr}`) - } + await runExtensionSetup(destDir, gpuSm, cudaVersion, (line) => { + logger.info(`[ext-setup] ${line}`) + emit({ step: 'setting_up', message: line }) + }) } try { @@ -782,11 +779,22 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe } catch { /* Python might not be running yet */ } } - emit({ step: 'done', extensionId: manifest.id }) + if (backupDir) { + await rmAsync(backupDir, { recursive: true, force: true }) + } + } catch (installErr) { + await rmAsync(destDir, { recursive: true, force: true }).catch(() => {}) + if (backupDir && existsSync(backupDir)) { + await rename(backupDir, destDir) + } + throw installErr + } + + emit({ step: 'done', extensionId }) const trustedRepos = await fetchTrustedRepos() - const ext = parseExtensionManifest(manifest, manifest.id, trustedRepos) - return { success: true, extensionId: manifest.id, extension: ext } + const ext = parseExtensionManifest(manifest, extensionId, trustedRepos) + return { success: true, extensionId, extension: ext } } catch (err) { emit({ step: 'error', message: String(err) }) @@ -801,16 +809,17 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe // Uninstall an extension — built-ins cannot be uninstalled ipcMain.handle('extensions:uninstall', async (_, extensionId: string) => { const userData = app.getPath('userData') - const builtinPath = join(getBuiltinExtensionsDir(), extensionId) + const safeExtensionId = assertSafeExtensionId(extensionId) + const builtinPath = resolveExtensionPathWithinRoot(getBuiltinExtensionsDir(), safeExtensionId) if (existsSync(builtinPath)) { - return { success: false, error: `"${extensionId}" is a built-in extension and cannot be uninstalled.` } + return { success: false, error: `"${safeExtensionId}" is a built-in extension and cannot be uninstalled.` } } const extensionsDir = getSettings(userData).extensionsDir - const extPath = join(extensionsDir, extensionId) + const extPath = resolveExtensionPathWithinRoot(extensionsDir, safeExtensionId) try { // Terminate process runner if it's a process extension - terminateProcessRunner(extensionId) + terminateProcessRunner(safeExtensionId) await rmAsync(extPath, { recursive: true, force: true }) // Hot-reload Python so it stops using the deleted model extension @@ -826,7 +835,8 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe // Re-run setup.py for a model extension (creates the venv if missing) ipcMain.handle('extensions:repair', async (_, extensionId: string) => { try { - const extDir = join(getSettings(app.getPath('userData')).extensionsDir, extensionId) + const safeExtensionId = assertSafeExtensionId(extensionId) + const extDir = resolveExtensionPathWithinRoot(getSettings(app.getPath('userData')).extensionsDir, safeExtensionId) if (!existsSync(join(extDir, 'setup.py'))) { return { success: false, error: 'No setup.py found for this extension' } } @@ -843,6 +853,7 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe // Trigger Python extension reload (without touching the filesystem) ipcMain.handle('extensions:reload', async () => { + terminateAllProcessRunners() try { const res = await axios.post(`${API_BASE_URL}/extensions/reload`, {}, { timeout: 10_000 }) return { success: true, errors: (res.data as { errors?: Record }).errors ?? {} } diff --git a/electron/main/python-bridge.ts b/electron/main/python-bridge.ts index 4e8d436..f0451ce 100644 --- a/electron/main/python-bridge.ts +++ b/electron/main/python-bridge.ts @@ -56,7 +56,7 @@ export class PythonBridge { MODELS_DIR: this.resolveModelsDir(), WORKSPACE_DIR: this.resolveWorkspaceDir(), EXTENSIONS_DIR: this.resolveExtensionsDir(), - SELECTED_MODEL_ID: process.env['SELECTED_MODEL_ID'] ?? '', + ...(process.env['SELECTED_MODEL_ID'] ? { SELECTED_MODEL_ID: process.env['SELECTED_MODEL_ID'] } : {}), HUGGING_FACE_HUB_TOKEN: this.resolveHfToken(), HF_TOKEN: this.resolveHfToken(), } diff --git a/package.json b/package.json index c81592a..c323a0c 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "modly", - "version": "0.3.5", + "version": "0.3.6", "description": "Local AI-powered 3D mesh generation from images", "main": "./out/main/index.js", "author": "Modly", diff --git a/src/areas/generate/components/WorkflowPanel.tsx b/src/areas/generate/components/WorkflowPanel.tsx index 0f23770..a6a98b2 100644 --- a/src/areas/generate/components/WorkflowPanel.tsx +++ b/src/areas/generate/components/WorkflowPanel.tsx @@ -330,6 +330,39 @@ function TextParamRow({ nodeId, nodes, onPatch }: { nodeId: string; nodes: FlowN ) } +function WaitParamRow({ nodeId }: { nodeId: string }) { + const status = useWorkflowRunStore((s) => s.runState.status) + const activeNodeId = useWorkflowRunStore((s) => s.activeNodeId) + const continueRun = useWorkflowRunStore((s) => s.continueRun) + const isPaused = status === 'paused' && activeNodeId === nodeId + + return ( +
+
+ + + + Wait +
+ {isPaused ? ( + + ) : ( +

+ Pauses the workflow until you click Continue. +

+ )} +
+ ) +} + function ExtensionParamRow({ nodeId, ext, nodes, onPatch }: { nodeId: string; ext: WorkflowExtension; nodes: FlowNode[]; onPatch: PatchFn }) { const [expanded, setExpanded] = useState(true) const node = nodes.find((n) => n.id === nodeId) @@ -411,7 +444,7 @@ function EmbeddedCanvas({ workflow, allExtensions }: { const { setCurrentJob } = useAppStore() const { runState, run, cancel } = useWorkflowRunStore() - const isRunning = runState.status === 'running' + const isRunning = runState.status === 'running' || runState.status === 'paused' // Update AddToScene node when run completes useEffect(() => { @@ -455,7 +488,7 @@ function EmbeddedCanvas({ workflow, allExtensions }: { ) const paramNodes = sortedNodes.filter((n) => - (n.type === 'imageNode' || n.type === 'textNode' || n.type === 'meshNode' || n.type === 'extensionNode') + (n.type === 'imageNode' || n.type === 'textNode' || n.type === 'meshNode' || n.type === 'extensionNode' || n.type === 'waitNode') && (n.data as { showInGenerate?: boolean }).showInGenerate === true, ) @@ -476,6 +509,7 @@ function EmbeddedCanvas({ workflow, allExtensions }: { {node.type === 'imageNode' && } {node.type === 'textNode' && } {node.type === 'meshNode' && } + {node.type === 'waitNode' && } {node.type === 'extensionNode' && (() => { const ext = getWorkflowExtension(node.data.extensionId ?? '', allExtensions) return ext ? : null diff --git a/src/areas/workflows/WorkflowsPage.tsx b/src/areas/workflows/WorkflowsPage.tsx index e2575e6..05641f1 100644 --- a/src/areas/workflows/WorkflowsPage.tsx +++ b/src/areas/workflows/WorkflowsPage.tsx @@ -27,13 +27,14 @@ import TextNode from './nodes/TextNode' import AddToSceneNode from './nodes/AddToSceneNode' import Load3DMeshNode from './nodes/Load3DMeshNode' import PreviewImageNode from './nodes/PreviewImageNode' +import WaitNode from './nodes/WaitNode' 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 } +const NODE_TYPES = { extensionNode: ExtensionNode, imageNode: ImageNode, textNode: TextNode, outputNode: AddToSceneNode, meshNode: Load3DMeshNode, previewNode: PreviewImageNode, waitNode: WaitNode } const EDGE_TYPES = { workflowEdge: WorkflowEdge } const DEFAULT_EDGE_OPTS = { type: 'workflowEdge' } @@ -155,6 +156,7 @@ const PANEL_BUILTIN_NODES = [ { type: 'meshNode', label: 'Load 3D Mesh', color: '#a78bfa', icon: <> }, { type: 'outputNode', label: 'Add to Scene', color: '#a78bfa', icon: <> }, { type: 'previewNode', label: 'Preview Views', color: '#38bdf8', icon: <> }, + { type: 'waitNode', label: 'Wait', color: '#71717a', icon: <> }, ] function ExtGroupHeader({ title, author, expanded, onToggle, count }: { title: string; author?: string; expanded: boolean; onToggle: () => void; count: number }) { @@ -398,6 +400,7 @@ const BUILTIN_NODES = [ { type: 'meshNode', label: 'Load 3D Mesh', color: '#a78bfa', description: 'Load a 3D mesh file or use current model' }, { 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 PaletteItem = @@ -778,7 +781,7 @@ function WorkflowCanvasInner({ }) { const { screenToFlowPosition, updateNodeData, getNode } = useReactFlow() const { runState, run: runWorkflow, cancel } = useWorkflowRunStore() - const isRunning = runState.status === 'running' + const isRunning = runState.status === 'running' || runState.status === 'paused' const [nodes, setNodes, onNodesChange] = useNodesState(workflow.nodes as Node[]) const [edges, setEdges, onEdgesChange] = useEdgesState(workflow.edges as Edge[]) diff --git a/src/areas/workflows/nodes/ExtensionNode.tsx b/src/areas/workflows/nodes/ExtensionNode.tsx index d0498ed..692a467 100644 --- a/src/areas/workflows/nodes/ExtensionNode.tsx +++ b/src/areas/workflows/nodes/ExtensionNode.tsx @@ -152,6 +152,16 @@ export default function ExtensionNode({ id, data, selected }: { id: string; data updateNodeData(id, { params: { ...data.params, [key]: val } }) }, [id, data.params, updateNodeData]) + const paramById = new Map(ext?.params.map((p) => [p.id, p])) + + const isVisible = (param: ParamSchema): boolean => { + if (!param.show_if) return true + return Object.entries(param.show_if).every(([key, expected]) => { + const current = data.params[key] ?? paramById.get(key)?.default + return Array.isArray(expected) ? expected.includes(current as string | number) : current === expected + }) + } + // ── IO subheader ───────────────────────────────────────────────────────── const ioSubheader = isMulti ? ( // Multi-input layout: one row per input, output on first row @@ -242,7 +252,7 @@ export default function ExtensionNode({ id, data, selected }: { id: string; data > {hasParams && (
- {ext!.params.map((param) => { + {ext!.params.filter(isVisible).map((param) => { const val = (data.params[param.id] ?? param.default) as number | string return (
diff --git a/src/areas/workflows/nodes/WaitNode.tsx b/src/areas/workflows/nodes/WaitNode.tsx new file mode 100644 index 0000000..8e30302 --- /dev/null +++ b/src/areas/workflows/nodes/WaitNode.tsx @@ -0,0 +1,52 @@ +import { Handle, Position } from '@xyflow/react' +import type { WFNodeData } from '@shared/types/electron.d' +import { useWorkflowRunStore } from '../workflowRunStore' +import BaseNode from './BaseNode' + +const HANDLE_STYLE = { background: '#71717a', width: 14, height: 14, border: '2.5px solid #18181b' } + +export default function WaitNode({ id, data, selected }: { id: string; data: WFNodeData; selected?: boolean }) { + const status = useWorkflowRunStore((s) => s.runState.status) + const activeNodeId = useWorkflowRunStore((s) => s.activeNodeId) + const continueRun = useWorkflowRunStore((s) => s.continueRun) + const isPaused = status === 'paused' && activeNodeId === id + + return ( + + + + + } + subheader={isPaused ? ( + + ) : undefined} + handles={ + <> + + + + } + > +
+

+ {isPaused ? 'Workflow paused — click Continue to resume.' : 'Pauses the workflow until you click Continue.'} +

+
+ + ) +} diff --git a/src/areas/workflows/workflowRunStore.ts b/src/areas/workflows/workflowRunStore.ts index 44b8c3e..5b86668 100644 --- a/src/areas/workflows/workflowRunStore.ts +++ b/src/areas/workflows/workflowRunStore.ts @@ -8,7 +8,7 @@ import type { Workflow, WFNode, WFEdge } from '@shared/types/electron.d' // ─── Types ──────────────────────────────────────────────────────────────────── export interface WorkflowRunState { - status: 'idle' | 'running' | 'done' | 'error' + status: 'idle' | 'running' | 'paused' | 'done' | 'error' blockIndex: number blockTotal: number blockProgress: number @@ -25,6 +25,14 @@ const IDLE: WorkflowRunState = { // Module-level refs — survive component unmounts / navigation const _cancel = { current: false } const _activeJobId = { current: null as string | null } +const _resume = { current: null as (() => void) | null } + +function flushResume(): void { + const fn = _resume.current + if (!fn) return + _resume.current = null + fn() +} // ─── Topological sort ───────────────────────────────────────────────────────── @@ -60,9 +68,10 @@ interface WorkflowRunStore { /** nodeId → workspace URL for image outputs (populated after each run) */ nodeImageOutputs: Record - run: (workflow: Workflow, allExtensions: WorkflowExtension[]) => Promise - cancel: () => void - reset: () => void + run: (workflow: Workflow, allExtensions: WorkflowExtension[], overrideImageData?: string) => Promise + cancel: () => void + reset: () => void + continueRun: () => void } export const useWorkflowRunStore = create((set) => ({ @@ -77,7 +86,9 @@ export const useWorkflowRunStore = create((set) => ({ const appState = useAppStore.getState() const apiUrl = appState.apiUrl const ordered = topoSort(workflow.nodes, workflow.edges) - const execNodes = ordered.filter((n) => n.type === 'extensionNode' && n.data.enabled) + const execNodes = ordered.filter((n) => + (n.type === 'extensionNode' || n.type === 'waitNode') && n.data.enabled, + ) const selectedImagePath = appState.selectedImagePath ?? '' const selectedImageData = appState.selectedImageData ?? undefined @@ -108,6 +119,7 @@ export const useWorkflowRunStore = create((set) => ({ // nodeId → { filePath, text, outputType } const nodeOutputs = new Map() + const outputNodeIds = new Set(ordered.filter((n) => n.type === 'outputNode').map((n) => n.id)) // Pre-populate source nodes for (const node of ordered) { @@ -182,6 +194,21 @@ export const useWorkflowRunStore = create((set) => ({ runState: { ...s.runState, blockIndex: i, blockProgress: 0, blockStep: 'Starting…' }, })) + // ── Wait node → pause until continueRun(), then passthrough ─────── + if (node.type === 'waitNode') { + set((s) => ({ runState: { ...s.runState, status: 'paused', blockStep: 'Paused — click Continue' } })) + await new Promise((resolve) => { _resume.current = resolve }) + if (_cancel.current) { set({ runState: IDLE, activeNodeId: null }); return } + + nodeOutputs.set(node.id, { + filePath: nodeInputPath, + text: nodeInputText, + outputType: incomingEdges[0] ? nodeOutputs.get(incomingEdges[0].source)?.outputType : undefined, + }) + set((s) => ({ runState: { ...s.runState, status: 'running' } })) + continue + } + // ── Model extensions → HTTP API ─────────────────────────────────── // Process extensions → IPC runProcess const isModelNode = ext?.type === 'model' @@ -272,6 +299,20 @@ export const useWorkflowRunStore = create((set) => ({ // Store output with type for downstream routing const outputType = ext?.output ?? (nodeInputPath ? 'mesh' : undefined) nodeOutputs.set(node.id, { filePath: nodeInputPath, text: nodeInputText, outputType }) + + // If this node feeds an Add-to-Scene, push the mesh to currentJob + // immediately so the 3D viewer loads it without waiting for the rest of the run. + const norm = nodeInputPath?.replace(/\\/g, '/') + if ( + norm?.startsWith(workspaceDir) && + workflow.edges.some((e) => e.source === node.id && outputNodeIds.has(e.target)) + ) { + useAppStore.getState().updateCurrentJob({ + status: 'done', + progress: 100, + outputUrl: `/workspace/${norm.slice(workspaceDir.length).replace(/^\//, '')}`, + }) + } } // ── Collect image outputs for preview nodes ─────────────────────── @@ -289,7 +330,8 @@ export const useWorkflowRunStore = create((set) => ({ let outputUrl: string | undefined let outputPath: string | undefined - const outputNodeDef = ordered.find((n) => n.type === 'outputNode') + // Use the last AddToScene in topo order — its predecessor is the final scene mesh. + const outputNodeDef = [...ordered].reverse().find((n) => n.type === 'outputNode') if (outputNodeDef) { for (const edge of workflow.edges.filter((e) => e.target === outputNodeDef.id)) { const src = nodeOutputs.get(edge.source) @@ -340,6 +382,7 @@ export const useWorkflowRunStore = create((set) => ({ cancel() { _cancel.current = true + flushResume() if (_activeJobId.current) { const apiUrl = useAppStore.getState().apiUrl axios.create({ baseURL: apiUrl }).post(`/generate/cancel/${_activeJobId.current}`).catch(() => {}) @@ -351,4 +394,8 @@ export const useWorkflowRunStore = create((set) => ({ reset() { set({ runState: IDLE, activeNodeId: null, activeWorkflowId: null, nodeImageOutputs: {} }) }, + + continueRun() { + flushResume() + }, })) diff --git a/src/shared/types/electron.d.ts b/src/shared/types/electron.d.ts index acb21ad..f5d5ca8 100644 --- a/src/shared/types/electron.d.ts +++ b/src/shared/types/electron.d.ts @@ -38,6 +38,7 @@ export interface ParamSchema { max?: number step?: number tooltip?: string + show_if?: Record } export interface ProcessExtension {