diff --git a/src/store/execution/__tests__/simpleNodeExecutors.test.ts b/src/store/execution/__tests__/simpleNodeExecutors.test.ts new file mode 100644 index 00000000..aa499b61 --- /dev/null +++ b/src/store/execution/__tests__/simpleNodeExecutors.test.ts @@ -0,0 +1,451 @@ +import { describe, it, expect, vi } from "vitest"; +import { + executeAnnotation, + executePrompt, + executePromptConstructor, + executeOutput, + executeOutputGallery, + executeImageCompare, +} from "../simpleNodeExecutors"; +import type { NodeExecutionContext } from "../types"; +import type { WorkflowNode, WorkflowEdge } from "@/types"; + +function makeCtx( + node: WorkflowNode, + overrides: Partial = {} +): NodeExecutionContext { + return { + node, + getConnectedInputs: vi.fn().mockReturnValue({ + images: [], + videos: [], + audio: [], + text: null, + dynamicInputs: {}, + easeCurve: null, + }), + updateNodeData: vi.fn(), + getFreshNode: vi.fn().mockReturnValue(node), + getEdges: vi.fn().mockReturnValue([]), + getNodes: vi.fn().mockReturnValue([]), + providerSettings: {} as any, + addIncurredCost: vi.fn(), + addToGlobalHistory: vi.fn(), + generationsPath: null, + saveDirectoryPath: null, + get: vi.fn(), + ...overrides, + }; +} + +function makeNode(id: string, type: string, data: Record = {}): WorkflowNode { + return { id, type, position: { x: 0, y: 0 }, data } as WorkflowNode; +} + +describe("executeAnnotation", () => { + it("should set sourceImage from connected image", async () => { + const node = makeNode("ann", "annotation", { outputImage: null }); + const ctx = makeCtx(node, { + getConnectedInputs: vi.fn().mockReturnValue({ + images: ["data:image/png;base64,abc"], + videos: [], + audio: [], + text: null, + dynamicInputs: {}, + easeCurve: null, + }), + }); + + await executeAnnotation(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("ann", { sourceImage: "data:image/png;base64,abc" }); + }); + + it("should pass through image as output when no annotations exist", async () => { + const node = makeNode("ann", "annotation", { outputImage: null }); + const ctx = makeCtx(node, { + getConnectedInputs: vi.fn().mockReturnValue({ + images: ["data:image/png;base64,abc"], + videos: [], + audio: [], + text: null, + dynamicInputs: {}, + easeCurve: null, + }), + }); + + await executeAnnotation(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("ann", { outputImage: "data:image/png;base64,abc" }); + }); + + it("should not overwrite existing outputImage", async () => { + const node = makeNode("ann", "annotation", { outputImage: "existing-annotated-image" }); + const ctx = makeCtx(node, { + getConnectedInputs: vi.fn().mockReturnValue({ + images: ["data:image/png;base64,abc"], + videos: [], + audio: [], + text: null, + dynamicInputs: {}, + easeCurve: null, + }), + }); + + await executeAnnotation(ctx); + + // Should set sourceImage but NOT overwrite outputImage + expect(ctx.updateNodeData).toHaveBeenCalledWith("ann", { sourceImage: "data:image/png;base64,abc" }); + const calls = (ctx.updateNodeData as ReturnType).mock.calls; + const outputCall = calls.find((c: unknown[]) => (c[1] as Record).outputImage !== undefined); + expect(outputCall).toBeUndefined(); + }); + + it("should do nothing when no images connected", async () => { + const node = makeNode("ann", "annotation", { outputImage: null }); + const ctx = makeCtx(node); + + await executeAnnotation(ctx); + + expect(ctx.updateNodeData).not.toHaveBeenCalled(); + }); + + it("should set error on exception", async () => { + const node = makeNode("ann", "annotation", {}); + const ctx = makeCtx(node, { + getConnectedInputs: vi.fn().mockImplementation(() => { + throw new Error("test error"); + }), + }); + + await executeAnnotation(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("ann", { error: "test error" }); + }); +}); + +describe("executePrompt", () => { + it("should update prompt from connected text", async () => { + const node = makeNode("p", "prompt", { prompt: "old" }); + const ctx = makeCtx(node, { + getConnectedInputs: vi.fn().mockReturnValue({ + images: [], + videos: [], + audio: [], + text: "new prompt", + dynamicInputs: {}, + easeCurve: null, + }), + }); + + await executePrompt(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("p", { prompt: "new prompt" }); + }); + + it("should not update prompt when no text connected", async () => { + const node = makeNode("p", "prompt", { prompt: "keep" }); + const ctx = makeCtx(node); + + await executePrompt(ctx); + + expect(ctx.updateNodeData).not.toHaveBeenCalled(); + }); + + it("should set error on exception", async () => { + const node = makeNode("p", "prompt", {}); + const ctx = makeCtx(node, { + getConnectedInputs: vi.fn().mockImplementation(() => { + throw new Error("fail"); + }), + }); + + await executePrompt(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("p", { error: "fail" }); + }); +}); + +describe("executePromptConstructor", () => { + it("should resolve @variables from connected prompt nodes", async () => { + const pcNode = makeNode("pc", "promptConstructor", { + template: "Hello @name, welcome to @place", + outputText: null, + unresolvedVars: [], + }); + const promptNode = makeNode("p1", "prompt", { prompt: "World", variableName: "name" }); + const promptNode2 = makeNode("p2", "prompt", { prompt: "Earth", variableName: "place" }); + + const edges: WorkflowEdge[] = [ + { id: "e1", source: "p1", target: "pc", sourceHandle: "text", targetHandle: "text" } as WorkflowEdge, + { id: "e2", source: "p2", target: "pc", sourceHandle: "text", targetHandle: "text" } as WorkflowEdge, + ]; + + const ctx = makeCtx(pcNode, { + getFreshNode: vi.fn().mockReturnValue(pcNode), + getEdges: vi.fn().mockReturnValue(edges), + getNodes: vi.fn().mockReturnValue([pcNode, promptNode, promptNode2]), + }); + + await executePromptConstructor(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("pc", { + outputText: "Hello World, welcome to Earth", + unresolvedVars: [], + }); + }); + + it("should track unresolved variables", async () => { + const pcNode = makeNode("pc", "promptConstructor", { + template: "Hello @name, welcome to @unknown", + outputText: null, + unresolvedVars: [], + }); + const promptNode = makeNode("p1", "prompt", { prompt: "World", variableName: "name" }); + + const edges: WorkflowEdge[] = [ + { id: "e1", source: "p1", target: "pc", sourceHandle: "text", targetHandle: "text" } as WorkflowEdge, + ]; + + const ctx = makeCtx(pcNode, { + getFreshNode: vi.fn().mockReturnValue(pcNode), + getEdges: vi.fn().mockReturnValue(edges), + getNodes: vi.fn().mockReturnValue([pcNode, promptNode]), + }); + + await executePromptConstructor(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("pc", { + outputText: "Hello World, welcome to @unknown", + unresolvedVars: ["unknown"], + }); + }); + + it("should use fresh node data", async () => { + const staleNode = makeNode("pc", "promptConstructor", { + template: "stale template", + }); + const freshNode = makeNode("pc", "promptConstructor", { + template: "fresh @var", + }); + + const ctx = makeCtx(staleNode, { + getFreshNode: vi.fn().mockReturnValue(freshNode), + getEdges: vi.fn().mockReturnValue([]), + getNodes: vi.fn().mockReturnValue([freshNode]), + }); + + await executePromptConstructor(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("pc", { + outputText: "fresh @var", + unresolvedVars: ["var"], + }); + }); +}); + +describe("executeOutput", () => { + it("should set video content from videos array", async () => { + const node = makeNode("out", "output", {}); + const ctx = makeCtx(node, { + getConnectedInputs: vi.fn().mockReturnValue({ + images: [], + videos: ["data:video/mp4;base64,abc"], + audio: [], + text: null, + dynamicInputs: {}, + easeCurve: null, + }), + }); + + await executeOutput(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("out", { + image: "data:video/mp4;base64,abc", + video: "data:video/mp4;base64,abc", + contentType: "video", + }); + }); + + it("should set image content from images array", async () => { + const node = makeNode("out", "output", {}); + const ctx = makeCtx(node, { + getConnectedInputs: vi.fn().mockReturnValue({ + images: ["data:image/png;base64,img"], + videos: [], + audio: [], + text: null, + dynamicInputs: {}, + easeCurve: null, + }), + }); + + await executeOutput(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("out", { + image: "data:image/png;base64,img", + video: null, + contentType: "image", + }); + }); + + it("should detect video URLs in images array", async () => { + const node = makeNode("out", "output", {}); + const ctx = makeCtx(node, { + getConnectedInputs: vi.fn().mockReturnValue({ + images: ["data:video/mp4;base64,vid"], + videos: [], + audio: [], + text: null, + dynamicInputs: {}, + easeCurve: null, + }), + }); + + await executeOutput(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("out", { + image: "data:video/mp4;base64,vid", + video: "data:video/mp4;base64,vid", + contentType: "video", + }); + }); + + it("should detect fal.media URLs as video", async () => { + const node = makeNode("out", "output", {}); + const ctx = makeCtx(node, { + getConnectedInputs: vi.fn().mockReturnValue({ + images: ["https://fal.media/files/abc123.mp4"], + videos: [], + audio: [], + text: null, + dynamicInputs: {}, + easeCurve: null, + }), + }); + + await executeOutput(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("out", { + image: "https://fal.media/files/abc123.mp4", + video: "https://fal.media/files/abc123.mp4", + contentType: "video", + }); + }); +}); + +describe("executeOutputGallery", () => { + it("should add new images to gallery", async () => { + const node = makeNode("gal", "outputGallery", { images: ["existing.png"] }); + const ctx = makeCtx(node, { + getConnectedInputs: vi.fn().mockReturnValue({ + images: ["new1.png", "new2.png"], + videos: [], + audio: [], + text: null, + dynamicInputs: {}, + easeCurve: null, + }), + }); + + await executeOutputGallery(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("gal", { + images: ["new1.png", "new2.png", "existing.png"], + }); + }); + + it("should not add duplicate images", async () => { + const node = makeNode("gal", "outputGallery", { images: ["existing.png"] }); + const ctx = makeCtx(node, { + getConnectedInputs: vi.fn().mockReturnValue({ + images: ["existing.png", "new.png"], + videos: [], + audio: [], + text: null, + dynamicInputs: {}, + easeCurve: null, + }), + }); + + await executeOutputGallery(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("gal", { + images: ["new.png", "existing.png"], + }); + }); + + it("should not update when no new images", async () => { + const node = makeNode("gal", "outputGallery", { images: ["existing.png"] }); + const ctx = makeCtx(node, { + getConnectedInputs: vi.fn().mockReturnValue({ + images: ["existing.png"], + videos: [], + audio: [], + text: null, + dynamicInputs: {}, + easeCurve: null, + }), + }); + + await executeOutputGallery(ctx); + + expect(ctx.updateNodeData).not.toHaveBeenCalled(); + }); +}); + +describe("executeImageCompare", () => { + it("should set imageA and imageB from connected images", async () => { + const node = makeNode("cmp", "imageCompare", {}); + const ctx = makeCtx(node, { + getConnectedInputs: vi.fn().mockReturnValue({ + images: ["img-a.png", "img-b.png"], + videos: [], + audio: [], + text: null, + dynamicInputs: {}, + easeCurve: null, + }), + }); + + await executeImageCompare(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("cmp", { + imageA: "img-a.png", + imageB: "img-b.png", + }); + }); + + it("should handle single image", async () => { + const node = makeNode("cmp", "imageCompare", {}); + const ctx = makeCtx(node, { + getConnectedInputs: vi.fn().mockReturnValue({ + images: ["img-a.png"], + videos: [], + audio: [], + text: null, + dynamicInputs: {}, + easeCurve: null, + }), + }); + + await executeImageCompare(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("cmp", { + imageA: "img-a.png", + imageB: null, + }); + }); + + it("should handle no images", async () => { + const node = makeNode("cmp", "imageCompare", {}); + const ctx = makeCtx(node); + + await executeImageCompare(ctx); + + expect(ctx.updateNodeData).toHaveBeenCalledWith("cmp", { + imageA: null, + imageB: null, + }); + }); +}); diff --git a/src/store/execution/simpleNodeExecutors.ts b/src/store/execution/simpleNodeExecutors.ts new file mode 100644 index 00000000..c77c8f93 --- /dev/null +++ b/src/store/execution/simpleNodeExecutors.ts @@ -0,0 +1,223 @@ +/** + * Simple Node Executors + * + * Executors for node types that don't call external APIs: + * annotation, prompt, promptConstructor, output, outputGallery, imageCompare. + * + * These are used by executeWorkflow (and some by regenerateNode). + */ + +import type { + AnnotationNodeData, + PromptConstructorNodeData, + PromptNodeData, + OutputNodeData, + OutputGalleryNodeData, + WorkflowNode, +} from "@/types"; +import type { NodeExecutionContext } from "./types"; + +/** + * Annotation node: receives upstream image as source, passes through if no annotations. + */ +export async function executeAnnotation(ctx: NodeExecutionContext): Promise { + const { node, getConnectedInputs, updateNodeData } = ctx; + try { + const { images } = getConnectedInputs(node.id); + const image = images[0] || null; + if (image) { + updateNodeData(node.id, { sourceImage: image }); + // If no annotations, pass through the image + const nodeData = node.data as AnnotationNodeData; + if (!nodeData.outputImage) { + updateNodeData(node.id, { outputImage: image }); + } + } + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + console.error(`[Workflow] Annotation node ${node.id} failed:`, message); + updateNodeData(node.id, { error: message }); + } +} + +/** + * Prompt node: receives upstream text and updates its prompt field. + */ +export async function executePrompt(ctx: NodeExecutionContext): Promise { + const { node, getConnectedInputs, updateNodeData } = ctx; + try { + const { text: connectedText } = getConnectedInputs(node.id); + if (connectedText !== null) { + updateNodeData(node.id, { prompt: connectedText }); + } + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + console.error(`[Workflow] Prompt node ${node.id} failed:`, message); + updateNodeData(node.id, { error: message }); + } +} + +/** + * PromptConstructor node: resolves @variables from connected prompt nodes. + */ +export async function executePromptConstructor(ctx: NodeExecutionContext): Promise { + const { node, updateNodeData, getFreshNode, getEdges, getNodes } = ctx; + try { + // Get fresh node data from store + const freshNode = getFreshNode(node.id); + const nodeData = (freshNode?.data || node.data) as PromptConstructorNodeData; + const template = nodeData.template; + + const edges = getEdges(); + const nodes = getNodes(); + + // Find connected prompt nodes via text edges + const connectedPromptNodes = edges + .filter((e) => e.target === node.id && e.targetHandle === "text") + .map((e) => nodes.find((n) => n.id === e.source)) + .filter((n): n is WorkflowNode => n !== undefined && n.type === "prompt"); + + // Build variable map from connected prompt nodes + const variableMap: Record = {}; + connectedPromptNodes.forEach((promptNode) => { + const promptData = promptNode.data as PromptNodeData; + if (promptData.variableName) { + variableMap[promptData.variableName] = promptData.prompt; + } + }); + + // Find all @variable patterns in template + const varPattern = /@(\w+)/g; + const unresolvedVars: string[] = []; + let resolvedText = template; + + // Replace @variables with values or track unresolved + const matches = template.matchAll(varPattern); + for (const match of matches) { + const varName = match[1]; + if (variableMap[varName] !== undefined) { + resolvedText = resolvedText.replaceAll(`@${varName}`, variableMap[varName]); + } else { + if (!unresolvedVars.includes(varName)) { + unresolvedVars.push(varName); + } + } + } + + updateNodeData(node.id, { + outputText: resolvedText, + unresolvedVars, + }); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + console.error(`[Workflow] PromptConstructor node ${node.id} failed:`, message); + updateNodeData(node.id, { error: message }); + } +} + +/** + * Output node: displays final image/video result. + */ +export async function executeOutput(ctx: NodeExecutionContext): Promise { + const { node, getConnectedInputs, updateNodeData, saveDirectoryPath } = ctx; + const { images, videos } = getConnectedInputs(node.id); + + // Check videos array first (typed data from source) + if (videos.length > 0) { + const videoContent = videos[0]; + updateNodeData(node.id, { + image: videoContent, + video: videoContent, + contentType: "video", + }); + + // Save to /outputs directory if we have a project path + if (saveDirectoryPath) { + const outputNodeData = node.data as OutputNodeData; + const outputsPath = `${saveDirectoryPath}/outputs`; + + fetch("/api/save-generation", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + directoryPath: outputsPath, + video: videoContent, + customFilename: outputNodeData.outputFilename || undefined, + createDirectory: true, + }), + }).catch((err) => { + console.error("Failed to save output:", err); + }); + } + } else if (images.length > 0) { + const content = images[0]; + // Fallback pattern matching for edge cases (video data that ended up in images array) + const isVideoContent = + content.startsWith("data:video/") || + content.includes(".mp4") || + content.includes(".webm") || + content.includes("fal.media"); + + if (isVideoContent) { + updateNodeData(node.id, { + image: content, + video: content, + contentType: "video", + }); + } else { + updateNodeData(node.id, { + image: content, + video: null, + contentType: "image", + }); + } + + // Save to /outputs directory if we have a project path + if (saveDirectoryPath) { + const outputNodeData = node.data as OutputNodeData; + const outputsPath = `${saveDirectoryPath}/outputs`; + + fetch("/api/save-generation", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + directoryPath: outputsPath, + image: isVideoContent ? undefined : content, + video: isVideoContent ? content : undefined, + customFilename: outputNodeData.outputFilename || undefined, + createDirectory: true, + }), + }).catch((err) => { + console.error("Failed to save output:", err); + }); + } + } +} + +/** + * OutputGallery node: accumulates images from upstream nodes. + */ +export async function executeOutputGallery(ctx: NodeExecutionContext): Promise { + const { node, getConnectedInputs, updateNodeData } = ctx; + const { images } = getConnectedInputs(node.id); + const galleryData = node.data as OutputGalleryNodeData; + const existing = new Set(galleryData.images || []); + const newImages = images.filter((img) => !existing.has(img)); + if (newImages.length > 0) { + updateNodeData(node.id, { + images: [...newImages, ...(galleryData.images || [])], + }); + } +} + +/** + * ImageCompare node: takes two upstream images for side-by-side comparison. + */ +export async function executeImageCompare(ctx: NodeExecutionContext): Promise { + const { node, getConnectedInputs, updateNodeData } = ctx; + const { images } = getConnectedInputs(node.id); + updateNodeData(node.id, { + imageA: images[0] || null, + imageB: images[1] || null, + }); +}