refactor: unify nanoBanana execution into shared executor

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
5 months ago · 1b002023ff
2 changed files with 560 additions and 0 deletions
--- a/src/store/execution/tests/nanoBananaExecutor.test.ts
+++ b/src/store/execution/tests/nanoBananaExecutor.test.ts
@ -0,0 +1,347 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { executeNanoBanana } from "../nanoBananaExecutor";
+import type { NodeExecutionContext } from "../types";
+import type { WorkflowNode } from "@/types";
+
+// Mock fetch globally
+const mockFetch = vi.fn();
+vi.stubGlobal("fetch", mockFetch);
+
+// Mock calculateGenerationCost
+vi.mock("@/utils/costCalculator", () => ({
+  calculateGenerationCost: vi.fn().mockReturnValue(0.05),
+}));
+
+function makeNode(data: Record<string, unknown> = {}): WorkflowNode {
+  return {
+    id: "gen-1",
+    type: "nanoBanana",
+    position: { x: 0, y: 0 },
+    data: {
+      outputImage: null,
+      inputImages: [],
+      inputPrompt: null,
+      status: null,
+      error: null,
+      aspectRatio: "1:1",
+      resolution: "1024x1024",
+      model: "nano-banana",
+      useGoogleSearch: false,
+      selectedModel: { provider: "gemini", modelId: "nano-banana", displayName: "Nano Banana" },
+      parameters: {},
+      imageHistory: [],
+      selectedHistoryIndex: 0,
+      ...data,
+    },
+  } as WorkflowNode;
+}
+
+const defaultProviderSettings = {
+  providers: {
+    gemini: { apiKey: "" },
+    replicate: { apiKey: "" },
+    fal: { apiKey: "" },
+    kie: { apiKey: "" },
+    wavespeed: { apiKey: "" },
+    openai: { apiKey: "" },
+  },
+} as any;
+
+function makeCtx(
+  node: WorkflowNode,
+  overrides: Partial<NodeExecutionContext> = {}
+): NodeExecutionContext {
+  return {
+    node,
+    getConnectedInputs: vi.fn().mockReturnValue({
+      images: [],
+      videos: [],
+      audio: [],
+      text: "test prompt",
+      dynamicInputs: {},
+      easeCurve: null,
+    }),
+    updateNodeData: vi.fn(),
+    getFreshNode: vi.fn().mockReturnValue(node),
+    getEdges: vi.fn().mockReturnValue([]),
+    getNodes: vi.fn().mockReturnValue([node]),
+    providerSettings: defaultProviderSettings,
+    addIncurredCost: vi.fn(),
+    addToGlobalHistory: vi.fn(),
+    generationsPath: null,
+    saveDirectoryPath: null,
+    get: vi.fn().mockReturnValue({
+      edges: [],
+      nodes: [node],
+      addToGlobalHistory: vi.fn(),
+      addIncurredCost: vi.fn(),
+      generationsPath: null,
+    }),
+    ...overrides,
+  };
+}
+
+beforeEach(() => {
+  vi.clearAllMocks();
+});
+
+describe("executeNanoBanana", () => {
+  it("should throw when no text input is provided", async () => {
+    const node = makeNode();
+    const ctx = makeCtx(node, {
+      getConnectedInputs: vi.fn().mockReturnValue({
+        images: [],
+        videos: [],
+        audio: [],
+        text: null,
+        dynamicInputs: {},
+        easeCurve: null,
+      }),
+    });
+
+    await expect(executeNanoBanana(ctx)).rejects.toThrow("Missing text input");
+
+    expect(ctx.updateNodeData).toHaveBeenCalledWith("gen-1", {
+      status: "error",
+      error: "Missing text input",
+    });
+  });
+
+  it("should set loading status before API call", async () => {
+    const node = makeNode();
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result" }),
+    });
+
+    const ctx = makeCtx(node);
+    await executeNanoBanana(ctx);
+
+    // Check that loading was set
+    const calls = (ctx.updateNodeData as ReturnType<typeof vi.fn>).mock.calls;
+    const loadingCall = calls.find(
+      (c: unknown[]) => (c[1] as Record<string, unknown>).status === "loading"
+    );
+    expect(loadingCall).toBeDefined();
+  });
+
+  it("should call /api/generate with correct payload", async () => {
+    const node = makeNode();
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result" }),
+    });
+
+    const ctx = makeCtx(node);
+    await executeNanoBanana(ctx);
+
+    expect(mockFetch).toHaveBeenCalledWith(
+      "/api/generate",
+      expect.objectContaining({
+        method: "POST",
+        body: expect.stringContaining('"prompt":"test prompt"'),
+      })
+    );
+  });
+
+  it("should update node with result on success", async () => {
+    const node = makeNode();
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result" }),
+    });
+
+    const ctx = makeCtx(node);
+    await executeNanoBanana(ctx);
+
+    const calls = (ctx.updateNodeData as ReturnType<typeof vi.fn>).mock.calls;
+    const completeCall = calls.find(
+      (c: unknown[]) => (c[1] as Record<string, unknown>).status === "complete"
+    );
+    expect(completeCall).toBeDefined();
+    expect((completeCall![1] as Record<string, unknown>).outputImage).toBe("data:image/png;base64,result");
+  });
+
+  it("should add to global history on success", async () => {
+    const node = makeNode();
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result" }),
+    });
+
+    const ctx = makeCtx(node);
+    await executeNanoBanana(ctx);
+
+    expect(ctx.addToGlobalHistory).toHaveBeenCalledWith(
+      expect.objectContaining({
+        image: "data:image/png;base64,result",
+        prompt: "test prompt",
+      })
+    );
+  });
+
+  it("should track cost for gemini provider", async () => {
+    const node = makeNode();
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result" }),
+    });
+
+    const ctx = makeCtx(node);
+    await executeNanoBanana(ctx);
+
+    expect(ctx.addIncurredCost).toHaveBeenCalledWith(0.05);
+  });
+
+  it("should track cost for fal provider", async () => {
+    const node = makeNode({
+      selectedModel: { provider: "fal", modelId: "fal-model", displayName: "Fal", pricing: { amount: 0.10 } },
+    });
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result" }),
+    });
+
+    const ctx = makeCtx(node, {
+      getFreshNode: vi.fn().mockReturnValue(node),
+    });
+    await executeNanoBanana(ctx);
+
+    expect(ctx.addIncurredCost).toHaveBeenCalledWith(0.10);
+  });
+
+  it("should throw on HTTP error", async () => {
+    const node = makeNode();
+    mockFetch.mockResolvedValueOnce({
+      ok: false,
+      status: 500,
+      statusText: "Internal Server Error",
+      text: () => Promise.resolve('{"error": "Server exploded"}'),
+    });
+
+    const ctx = makeCtx(node);
+    await expect(executeNanoBanana(ctx)).rejects.toThrow("Server exploded");
+
+    const calls = (ctx.updateNodeData as ReturnType<typeof vi.fn>).mock.calls;
+    const errorCall = calls.find(
+      (c: unknown[]) => (c[1] as Record<string, unknown>).status === "error"
+    );
+    expect(errorCall).toBeDefined();
+  });
+
+  it("should throw on API failure (success=false)", async () => {
+    const node = makeNode();
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: () => Promise.resolve({ success: false, error: "Bad prompt" }),
+    });
+
+    const ctx = makeCtx(node);
+    await expect(executeNanoBanana(ctx)).rejects.toThrow("Bad prompt");
+  });
+
+  it("should use text from dynamicInputs.prompt when no direct text", async () => {
+    const node = makeNode();
+    const ctx = makeCtx(node, {
+      getConnectedInputs: vi.fn().mockReturnValue({
+        images: [],
+        videos: [],
+        audio: [],
+        text: null,
+        dynamicInputs: { prompt: "dynamic prompt" },
+        easeCurve: null,
+      }),
+    });
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result" }),
+    });
+
+    await executeNanoBanana(ctx);
+
+    expect(mockFetch).toHaveBeenCalledWith(
+      "/api/generate",
+      expect.objectContaining({
+        body: expect.stringContaining('"prompt":"dynamic prompt"'),
+      })
+    );
+  });
+
+  it("should pass images in request payload", async () => {
+    const node = makeNode();
+    const ctx = makeCtx(node, {
+      getConnectedInputs: vi.fn().mockReturnValue({
+        images: ["data:image/png;base64,img1"],
+        videos: [],
+        audio: [],
+        text: "with image",
+        dynamicInputs: {},
+        easeCurve: null,
+      }),
+    });
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result" }),
+    });
+
+    await executeNanoBanana(ctx);
+
+    const body = JSON.parse(mockFetch.mock.calls[0][1].body);
+    expect(body.images).toEqual(["data:image/png;base64,img1"]);
+  });
+
+  it("should fall back to stored inputs in regenerate mode", async () => {
+    const node = makeNode({
+      inputImages: ["stored-img.png"],
+      inputPrompt: "stored prompt",
+    });
+    const ctx = makeCtx(node, {
+      getConnectedInputs: vi.fn().mockReturnValue({
+        images: [],
+        videos: [],
+        audio: [],
+        text: null,
+        dynamicInputs: {},
+        easeCurve: null,
+      }),
+    });
+    // Enable regenerate mode: fallback to stored inputs
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result" }),
+    });
+
+    await executeNanoBanana(ctx, { useStoredFallback: true });
+
+    const body = JSON.parse(mockFetch.mock.calls[0][1].body);
+    expect(body.images).toEqual(["stored-img.png"]);
+    expect(body.prompt).toBe("stored prompt");
+  });
+
+  it("should push to downstream outputGallery nodes", async () => {
+    const node = makeNode();
+    const galleryNode = {
+      id: "gal-1",
+      type: "outputGallery",
+      data: { images: ["old.png"] },
+    } as WorkflowNode;
+
+    const ctx = makeCtx(node, {
+      getEdges: vi.fn().mockReturnValue([
+        { id: "e1", source: "gen-1", target: "gal-1" },
+      ]),
+      getNodes: vi.fn().mockReturnValue([node, galleryNode]),
+    });
+
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result" }),
+    });
+
+    await executeNanoBanana(ctx);
+
+    expect(ctx.updateNodeData).toHaveBeenCalledWith("gal-1", {
+      images: ["data:image/png;base64,result", "old.png"],
+    });
+  });
+});
--- a/src/store/execution/nanoBananaExecutor.ts
+++ b/src/store/execution/nanoBananaExecutor.ts
@ -0,0 +1,213 @@
+/**
+ * NanoBanana Executor
+ *
+ * Unified executor for nanoBanana (image generation) nodes.
+ * Used by both executeWorkflow and regenerateNode.
+ */
+
+import type {
+  NanoBananaNodeData,
+  OutputGalleryNodeData,
+} from "@/types";
+import { calculateGenerationCost } from "@/utils/costCalculator";
+import { buildGenerateHeaders } from "@/store/utils/buildApiHeaders";
+import type { NodeExecutionContext } from "./types";
+
+export interface NanoBananaOptions {
+  /** When true, falls back to stored inputImages/inputPrompt if no connections provide them. */
+  useStoredFallback?: boolean;
+}
+
+export async function executeNanoBanana(
+  ctx: NodeExecutionContext,
+  options: NanoBananaOptions = {}
+): Promise<void> {
+  const {
+    node,
+    getConnectedInputs,
+    updateNodeData,
+    getFreshNode,
+    getEdges,
+    getNodes,
+    signal,
+    providerSettings,
+    addIncurredCost,
+    addToGlobalHistory,
+    generationsPath,
+    get,
+  } = ctx;
+
+  const { useStoredFallback = false } = options;
+
+  const { images: connectedImages, text: connectedText, dynamicInputs } = getConnectedInputs(node.id);
+
+  // Get fresh node data from store
+  const freshNode = getFreshNode(node.id);
+  const nodeData = (freshNode?.data || node.data) as NanoBananaNodeData;
+
+  // Determine images and text (with optional fallback to stored values)
+  let images: string[];
+  let promptText: string | null;
+
+  if (useStoredFallback) {
+    images = connectedImages.length > 0 ? connectedImages : nodeData.inputImages;
+    promptText = connectedText ?? nodeData.inputPrompt;
+  } else {
+    images = connectedImages;
+    // For dynamic inputs, check if we have at least a prompt
+    const promptFromDynamic = Array.isArray(dynamicInputs.prompt)
+      ? dynamicInputs.prompt[0]
+      : dynamicInputs.prompt;
+    promptText = connectedText || promptFromDynamic || null;
+  }
+
+  if (!promptText) {
+    updateNodeData(node.id, {
+      status: "error",
+      error: "Missing text input",
+    });
+    throw new Error("Missing text input");
+  }
+
+  updateNodeData(node.id, {
+    inputImages: images,
+    inputPrompt: promptText,
+    status: "loading",
+    error: null,
+  });
+
+  const provider = nodeData.selectedModel?.provider || "gemini";
+  const headers = buildGenerateHeaders(provider, providerSettings);
+
+  const requestPayload = {
+    images,
+    prompt: promptText,
+    aspectRatio: nodeData.aspectRatio,
+    resolution: nodeData.resolution,
+    model: nodeData.model,
+    useGoogleSearch: nodeData.useGoogleSearch,
+    selectedModel: nodeData.selectedModel,
+    parameters: nodeData.parameters,
+    dynamicInputs,
+  };
+
+  const response = await fetch("/api/generate", {
+    method: "POST",
+    headers,
+    body: JSON.stringify(requestPayload),
+    ...(signal ? { signal } : {}),
+  });
+
+  if (!response.ok) {
+    const errorText = await response.text();
+    let errorMessage = `HTTP ${response.status}`;
+    try {
+      const errorJson = JSON.parse(errorText);
+      errorMessage = errorJson.error || errorMessage;
+    } catch {
+      if (errorText) errorMessage += ` - ${errorText.substring(0, 200)}`;
+    }
+
+    updateNodeData(node.id, {
+      status: "error",
+      error: errorMessage,
+    });
+    throw new Error(errorMessage);
+  }
+
+  const result = await response.json();
+
+  if (result.success && result.image) {
+    const timestamp = Date.now();
+    const imageId = `${timestamp}`;
+
+    // Save to global history
+    addToGlobalHistory({
+      image: result.image,
+      timestamp,
+      prompt: promptText,
+      aspectRatio: nodeData.aspectRatio,
+      model: nodeData.model,
+    });
+
+    // Add to node's carousel history
+    const newHistoryItem = {
+      id: imageId,
+      timestamp,
+      prompt: promptText,
+      aspectRatio: nodeData.aspectRatio,
+      model: nodeData.model,
+    };
+    const updatedHistory = [newHistoryItem, ...(nodeData.imageHistory || [])];
+
+    updateNodeData(node.id, {
+      outputImage: result.image,
+      status: "complete",
+      error: null,
+      imageHistory: updatedHistory,
+      selectedHistoryIndex: 0,
+    });
+
+    // Push new image to connected downstream outputGallery nodes
+    const edges = getEdges();
+    const nodes = getNodes();
+    edges
+      .filter((e) => e.source === node.id)
+      .forEach((e) => {
+        const target = nodes.find((n) => n.id === e.target);
+        if (target?.type === "outputGallery") {
+          const gData = target.data as OutputGalleryNodeData;
+          updateNodeData(target.id, {
+            images: [result.image, ...(gData.images || [])],
+          });
+        }
+      });
+
+    // Track cost
+    if (nodeData.selectedModel?.provider === "fal" && nodeData.selectedModel?.pricing) {
+      addIncurredCost(nodeData.selectedModel.pricing.amount);
+    } else if (!nodeData.selectedModel || nodeData.selectedModel.provider === "gemini") {
+      const generationCost = calculateGenerationCost(nodeData.model, nodeData.resolution);
+      addIncurredCost(generationCost);
+    }
+
+    // Auto-save to generations folder if configured
+    if (generationsPath) {
+      // Fire-and-forget save
+      fetch("/api/save-generation", {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({
+          directoryPath: generationsPath,
+          image: result.image,
+          prompt: promptText,
+          imageId,
+        }),
+      })
+        .then((res) => res.json())
+        .then((saveResult) => {
+          if (saveResult.success && saveResult.imageId && saveResult.imageId !== imageId) {
+            const currentNode = getNodes().find((n) => n.id === node.id);
+            if (currentNode) {
+              const currentData = currentNode.data as NanoBananaNodeData;
+              const histCopy = [...(currentData.imageHistory || [])];
+              const entryIndex = histCopy.findIndex((h) => h.id === imageId);
+              if (entryIndex !== -1) {
+                histCopy[entryIndex] = { ...histCopy[entryIndex], id: saveResult.imageId };
+                updateNodeData(node.id, { imageHistory: histCopy });
+              }
+            }
+          }
+        })
+        .catch((err) => {
+          console.error("Failed to save generation:", err);
+        });
+    }
+  } else {
+    updateNodeData(node.id, {
+      status: "error",
+      error: result.error || "Generation failed",
+    });
+    throw new Error(result.error || "Generation failed");
+  }
+}