popiart-node-canvas/src/store/execution/__tests__/nanoBananaExecutor.test.ts

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),
}));

const mockWriteBrowserGenerationFile = vi.fn();

vi.mock("@/utils/browserFileSystem", () => ({
  isBrowserFileSystemPath: (path: string | null | undefined) =>
    typeof path === "string" && path.startsWith("browserfs://"),
  joinBrowserFileSystemPath: (basePath: string, folderName: string) =>
    `${basePath.replace(/\/+$/, "")}/${encodeURIComponent(folderName)}`,
  writeBrowserGenerationFile: (...args: unknown[]) => mockWriteBrowserGenerationFile(...args),
}));

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,
      useImageSearch: 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: "" },
    newapiwg: { 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,
    trackSaveGeneration: vi.fn(),
    appendOutputGalleryImage: vi.fn(),
    get: vi.fn().mockReturnValue({
      edges: [],
      nodes: [node],
      addToGlobalHistory: vi.fn(),
      addIncurredCost: vi.fn(),
      generationsPath: null,
    }),
    ...overrides,
  };
}

beforeEach(() => {
  vi.clearAllMocks();
  mockWriteBrowserGenerationFile.mockResolvedValue({ filePath: "browserfs://root/project/generations/img.png" });
});

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("uses NewApiWG default for legacy image nodes when Gemini is not configured", async () => {
    const node = makeNode({ selectedModel: null });
    mockFetch.mockResolvedValueOnce({
      ok: true,
      json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result" }),
    });

    const ctx = makeCtx(node, {
      providerSettings: {
        providers: {
          ...defaultProviderSettings.providers,
          gemini: { apiKey: "" },
          newapiwg: { apiKey: "newapi-key" },
        },
      } as any,
    });

    await executeNanoBanana(ctx);

    const body = JSON.parse(mockFetch.mock.calls[0][1].body);
    expect(body.selectedModel).toMatchObject({
      provider: "newapiwg",
      modelId: "apiyi_nano_banana_2",
    });

    const completeCall = (ctx.updateNodeData as ReturnType<typeof vi.fn>).mock.calls.find(
      (c: unknown[]) => (c[1] as Record<string, unknown>).status === "complete"
    );
    expect(completeCall?.[1]).toMatchObject({
      lastUsedModel: {
        provider: "newapiwg",
        modelId: "apiyi_nano_banana_2",
      },
    });
  });

  it("falls back from the default NewApiWG image model on provider overload", async () => {
    const node = makeNode({
      selectedModel: {
        provider: "newapiwg",
        modelId: "apiyi_nano_banana_2",
        displayName: "apiyi_nano_banana_2",
        capabilities: ["text-to-image", "image-to-image"],
      },
    });

    mockFetch
      .mockResolvedValueOnce({
        ok: false,
        text: () => Promise.resolve(JSON.stringify({
          error: "apiyi_nano_banana_2: provider is temporarily overloaded. Please retry in a moment or switch models.",
        })),
      })
      .mockResolvedValueOnce({
        ok: true,
        json: () => Promise.resolve({ success: true, image: "data:image/png;base64,fallback" }),
      });

    const ctx = makeCtx(node);
    await executeNanoBanana(ctx);

    expect(mockFetch).toHaveBeenCalledTimes(2);
    const fallbackBody = JSON.parse(mockFetch.mock.calls[1][1].body);
    expect(fallbackBody.selectedModel).toMatchObject({
      provider: "newapiwg",
      modelId: "gpt-image-2-all",
    });

    const calls = (ctx.updateNodeData as ReturnType<typeof vi.fn>).mock.calls;
    const stampCall = calls.find(
      (c: unknown[]) => (c[1] as Record<string, unknown>).__usedFallback === true
    );
    expect(stampCall?.[1]).toMatchObject({
      __fallbackModelUsed: "gpt-image-2-all",
      __primaryError: "apiyi_nano_banana_2: provider is temporarily overloaded. Please retry in a moment or switch models.",
    });
  });

  it("does not implicitly fall back from the default NewApiWG image model on non-overload errors", async () => {
    const node = makeNode({
      selectedModel: {
        provider: "newapiwg",
        modelId: "apiyi_nano_banana_2",
        displayName: "apiyi_nano_banana_2",
        capabilities: ["text-to-image", "image-to-image"],
      },
    });

    mockFetch.mockResolvedValueOnce({
      ok: false,
      text: () => Promise.resolve(JSON.stringify({ error: "bad request" })),
    });

    const ctx = makeCtx(node);
    await expect(executeNanoBanana(ctx)).rejects.toThrow("bad request");
    expect(mockFetch).toHaveBeenCalledTimes(1);
  });

  it("reroutes stored Gemini image nodes to NewApiWG when Gemini key is missing", async () => {
    const node = makeNode();
    mockFetch.mockResolvedValueOnce({
      ok: true,
      json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result" }),
    });

    const ctx = makeCtx(node, {
      providerSettings: {
        providers: {
          ...defaultProviderSettings.providers,
          gemini: { apiKey: "" },
          newapiwg: { apiKey: "newapi-key" },
        },
      } as any,
    });

    await executeNanoBanana(ctx);

    const body = JSON.parse(mockFetch.mock.calls[0][1].body);
    expect(body.selectedModel.provider).toBe("newapiwg");
    expect(body.selectedModel.modelId).toBe("apiyi_nano_banana_2");
    expect(ctx.updateNodeData).toHaveBeenCalledWith("gen-1", expect.objectContaining({
      selectedModel: expect.objectContaining({
        provider: "newapiwg",
        modelId: "apiyi_nano_banana_2",
      }),
      fallbackModel: undefined,
      fallbackParameters: undefined,
    }));
  });

  it("reroutes stored Gemini image nodes when NewApiWG is configured via env", async () => {
    const node = makeNode();
    mockFetch.mockResolvedValueOnce({
      ok: true,
      json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result" }),
    });

    const ctx = makeCtx(node, {
      providerEnvStatus: {
        newapiwg: true,
        gemini: false,
        openai: false,
        anthropic: false,
        replicate: false,
        fal: false,
        kie: false,
        wavespeed: false,
      },
    });

    await executeNanoBanana(ctx);

    const request = mockFetch.mock.calls[0];
    const body = JSON.parse(request[1].body);
    expect(body.selectedModel).toMatchObject({
      provider: "newapiwg",
      modelId: "apiyi_nano_banana_2",
    });
    expect(request[1].headers["X-NewApiWG-Key"]).toBeUndefined();
  });

  it("promotes a NewApiWG fallback image model when stored Gemini has no key", async () => {
    const node = makeNode({
      fallbackModel: {
        provider: "newapiwg",
        modelId: "gpt-image-2-all",
        displayName: "gpt-image-2-all",
        capabilities: ["text-to-image", "image-to-image"],
      },
      fallbackParameters: { quality: "auto", size: "1024x1024" },
    });
    mockFetch.mockResolvedValueOnce({
      ok: true,
      json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result" }),
    });

    const ctx = makeCtx(node, {
      providerEnvStatus: {
        newapiwg: true,
        gemini: false,
        openai: false,
        anthropic: false,
        replicate: false,
        fal: false,
        kie: false,
        wavespeed: false,
      },
    });

    await executeNanoBanana(ctx);

    expect(mockFetch).toHaveBeenCalledTimes(1);
    const body = JSON.parse(mockFetch.mock.calls[0][1].body);
    expect(body.selectedModel).toMatchObject({
      provider: "newapiwg",
      modelId: "gpt-image-2-all",
    });
    expect(body.parameters).toEqual({ quality: "auto", size: "1024x1024" });
    expect(ctx.updateNodeData).toHaveBeenCalledWith("gen-1", expect.objectContaining({
      selectedModel: expect.objectContaining({
        provider: "newapiwg",
        modelId: "gpt-image-2-all",
      }),
      parameters: { quality: "auto", size: "1024x1024" },
      fallbackModel: undefined,
      fallbackParameters: undefined,
    }));
  });

  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 mark remote URL image outputs as remote-only before localization", async () => {
    const node = makeNode();
    mockFetch.mockResolvedValueOnce({
      ok: true,
      json: () => Promise.resolve({ success: true, image: "https://cdn.example.com/generated.png" }),
    });

    const ctx = makeCtx(node);
    await executeNanoBanana(ctx);

    const completeCall = (ctx.updateNodeData as ReturnType<typeof vi.fn>).mock.calls.find(
      (c: unknown[]) => (c[1] as Record<string, unknown>).status === "complete"
    );
    expect(completeCall?.[1]).toMatchObject({
      outputImage: "https://cdn.example.com/generated.png",
      outputImageStorageStatus: "remote-only",
    });
  });

  it("should keep generated image data in node carousel history for immediate switching", 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 completeCall = (ctx.updateNodeData as ReturnType<typeof vi.fn>).mock.calls.find(
      (c: unknown[]) => (c[1] as Record<string, unknown>).status === "complete"
    );
    expect(completeCall?.[1]).toMatchObject({
      imageHistory: [
        expect.objectContaining({
          image: "data:image/png;base64,result",
        }),
      ],
    });
  });

  it("should generate multiple image candidates into one node history", async () => {
    const node = makeNode({ imageCount: 4 });
    mockFetch
      .mockResolvedValueOnce({
        ok: true,
        json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result-1" }),
      })
      .mockResolvedValueOnce({
        ok: true,
        json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result-2" }),
      })
      .mockResolvedValueOnce({
        ok: true,
        json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result-3" }),
      })
      .mockResolvedValueOnce({
        ok: true,
        json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result-4" }),
      });

    const ctx = makeCtx(node);
    await executeNanoBanana(ctx);

    expect(mockFetch).toHaveBeenCalledTimes(4);
    expect(ctx.addToGlobalHistory).toHaveBeenCalledTimes(4);

    const completeCall = (ctx.updateNodeData as ReturnType<typeof vi.fn>).mock.calls.find(
      (c: unknown[]) => (c[1] as Record<string, unknown>).status === "complete"
    );
    expect(completeCall?.[1]).toMatchObject({
      outputImage: "data:image/png;base64,result-1",
      selectedHistoryIndex: 0,
      imageHistory: [
        expect.objectContaining({ image: "data:image/png;base64,result-1" }),
        expect.objectContaining({ image: "data:image/png;base64,result-2" }),
        expect.objectContaining({ image: "data:image/png;base64,result-3" }),
        expect.objectContaining({ image: "data:image/png;base64,result-4" }),
      ],
    });
  });

  it("should keep successful candidates when a later image in the batch fails", async () => {
    const node = makeNode({ imageCount: 4 });
    mockFetch
      .mockResolvedValueOnce({
        ok: true,
        json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result-1" }),
      })
      .mockResolvedValueOnce({
        ok: true,
        json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result-2" }),
      })
      .mockResolvedValueOnce({
        ok: true,
        json: () => Promise.resolve({ success: false, error: "Bad prompt" }),
      });

    const ctx = makeCtx(node);
    await executeNanoBanana(ctx);

    expect(mockFetch).toHaveBeenCalledTimes(3);
    const completeCall = (ctx.updateNodeData as ReturnType<typeof vi.fn>).mock.calls.find(
      (c: unknown[]) => (c[1] as Record<string, unknown>).status === "complete"
    );
    expect(completeCall?.[1]).toMatchObject({
      outputImage: "data:image/png;base64,result-1",
      error: "Bad prompt",
      imageHistory: [
        expect.objectContaining({ image: "data:image/png;base64,result-1" }),
        expect.objectContaining({ image: "data:image/png;base64,result-2" }),
      ],
    });
  });

  it("should save browser-local generations through the browser file system", async () => {
    const node = makeNode();
    mockFetch.mockResolvedValueOnce({
      ok: true,
      json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result" }),
    });

    const ctx = makeCtx(node, {
      generationsPath: "browserfs://root/project/generations",
    });
    await executeNanoBanana(ctx);

    expect(mockFetch).toHaveBeenCalledTimes(1);
    expect(mockFetch).toHaveBeenCalledWith("/api/generate", expect.anything());
    expect(mockWriteBrowserGenerationFile).toHaveBeenCalledWith(
      "browserfs://root/project/generations",
      expect.any(String),
      "data:image/png;base64,result"
    );
    expect(ctx.trackSaveGeneration).toHaveBeenCalled();
  });

  it("should localize remote browser-local generations after saving", async () => {
    const node = makeNode();
    const trackedSaves: Promise<void>[] = [];
    const updateNodeData = vi.fn((_: string, data: Record<string, unknown>) => {
      node.data = { ...node.data, ...data };
    });
    mockFetch.mockResolvedValueOnce({
      ok: true,
      json: () => Promise.resolve({ success: true, image: "https://cdn.example.com/generated.png" }),
    });
    mockWriteBrowserGenerationFile.mockResolvedValueOnce({
      filePath: "browserfs://root/project/generations/123.png",
      imageData: "data:image/png;base64,localized",
    });

    const ctx = makeCtx(node, {
      generationsPath: "browserfs://root/project/generations",
      updateNodeData,
      getNodes: vi.fn().mockReturnValue([node]),
      trackSaveGeneration: vi.fn((_: string, promise: Promise<void>) => trackedSaves.push(promise)),
    });
    await executeNanoBanana(ctx);
    await Promise.all(trackedSaves);

    expect(updateNodeData).toHaveBeenCalledWith("gen-1", { outputImageStorageStatus: "localizing" });
    expect(node.data).toMatchObject({
      outputImage: "data:image/png;base64,localized",
      outputImageRef: expect.any(String),
      outputImageStorageStatus: "localized",
    });
  });

  it("should localize remote server-saved generations after saving", async () => {
    const node = makeNode();
    const trackedSaves: Promise<void>[] = [];
    const updateNodeData = vi.fn((_: string, data: Record<string, unknown>) => {
      node.data = { ...node.data, ...data };
    });
    mockFetch.mockImplementation((url: string) => {
      if (url === "/api/generate") {
        return Promise.resolve({
          ok: true,
          json: () => Promise.resolve({ success: true, image: "https://cdn.example.com/generated.png" }),
        });
      }
      if (url === "/api/save-generation") {
        return Promise.resolve({
          json: () => Promise.resolve({ success: true, imageId: "saved-image" }),
        });
      }
      if (url === "/api/load-generation") {
        return Promise.resolve({
          json: () => Promise.resolve({ success: true, image: "data:image/png;base64,localized" }),
        });
      }
      return Promise.reject(new Error(`Unexpected URL: ${url}`));
    });

    const ctx = makeCtx(node, {
      generationsPath: "/tmp/generations",
      updateNodeData,
      getNodes: vi.fn().mockReturnValue([node]),
      trackSaveGeneration: vi.fn((_: string, promise: Promise<void>) => trackedSaves.push(promise)),
    });
    await executeNanoBanana(ctx);
    await Promise.all(trackedSaves);

    expect(mockFetch).toHaveBeenCalledWith("/api/load-generation", expect.objectContaining({
      body: JSON.stringify({ directoryPath: "/tmp/generations", imageId: "saved-image" }),
    }));
    expect(node.data).toMatchObject({
      outputImage: "data:image/png;base64,localized",
      outputImageRef: "saved-image",
      outputImageStorageStatus: "localized",
    });
  });

  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.appendOutputGalleryImage).toHaveBeenCalledWith("gal-1", "data:image/png;base64,result");
  });

  it("should push every generated candidate to downstream outputGallery nodes", async () => {
    const node = makeNode({ imageCount: 2 });
    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-1" }),
      })
      .mockResolvedValueOnce({
        ok: true,
        json: () => Promise.resolve({ success: true, image: "data:image/png;base64,result-2" }),
      });

    await executeNanoBanana(ctx);

    expect(ctx.appendOutputGalleryImage).toHaveBeenNthCalledWith(
      1,
      "gal-1",
      "data:image/png;base64,result-1"
    );
    expect(ctx.appendOutputGalleryImage).toHaveBeenNthCalledWith(
      2,
      "gal-1",
      "data:image/png;base64,result-2"
    );
  });

  it("falls back on primary failure and stamps metadata", async () => {
    const node = makeNode({
      fallbackModel: {
        provider: "replicate",
        modelId: "flux-dev",
        displayName: "Flux Dev",
      },
    });

    // Primary fails, fallback succeeds
    mockFetch
      .mockResolvedValueOnce({
        ok: true,
        json: () => Promise.resolve({ success: false, error: "Primary boom" }),
      })
      .mockResolvedValueOnce({
        ok: true,
        json: () => Promise.resolve({ success: true, image: "data:image/png;base64,fallback" }),
      });

    const ctx = makeCtx(node);
    await executeNanoBanana(ctx);

    expect(mockFetch).toHaveBeenCalledTimes(2);

    // Second fetch should carry the fallback model
    const secondBody = JSON.parse(mockFetch.mock.calls[1][1].body);
    expect(secondBody.selectedModel.modelId).toBe("flux-dev");

    // Metadata stamp should be present in the final updateNodeData call
    const calls = (ctx.updateNodeData as ReturnType<typeof vi.fn>).mock.calls;
    const stampCall = calls.find(
      (c: unknown[]) => (c[1] as Record<string, unknown>).__usedFallback === true
    );
    expect(stampCall).toBeDefined();
    expect((stampCall![1] as Record<string, unknown>).__fallbackModelUsed).toBe("Flux Dev");
    expect((stampCall![1] as Record<string, unknown>).__primaryError).toBe("Primary boom");

    const completeCall = calls.find(
      (c: unknown[]) => (c[1] as Record<string, unknown>).status === "complete"
    );
    expect(completeCall?.[1]).toMatchObject({
      lastUsedModel: {
        provider: "replicate",
        modelId: "flux-dev",
      },
    });
  });
});