refactor: extract simple node executors

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
5 months ago · 67b9ffef68
2 changed files with 674 additions and 0 deletions
--- a/src/store/execution/tests/simpleNodeExecutors.test.ts
+++ 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> = {}
+): 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<string, unknown> = {}): 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<typeof vi.fn>).mock.calls;
+    const outputCall = calls.find((c: unknown[]) => (c[1] as Record<string, unknown>).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,
+    });
+  });
+});
--- a/src/store/execution/simpleNodeExecutors.ts
+++ 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<void> {
+  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<void> {
+  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<void> {
+  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<string, string> = {};
+    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<void> {
+  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<void> {
+  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<void> {
+  const { node, getConnectedInputs, updateNodeData } = ctx;
+  const { images } = getConnectedInputs(node.id);
+  updateNodeData(node.id, {
+    imageA: images[0] || null,
+    imageB: images[1] || null,
+  });
+}