2 changed files with 560 additions and 0 deletions
@ -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"], |
||||
|
}); |
||||
|
}); |
||||
|
}); |
||||
@ -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"); |
||||
|
} |
||||
|
} |
||||
Loading…
Reference in new issue