Compare commits

...

2 Commits

Author SHA1 Message Date
jiajia 06221a114b Unify testing with the overload fallback fix 2 months ago
jiajia 088ea822a1 Keep default image generation moving through gateway overloads 2 months ago
  1. 62
      src/store/execution/__tests__/nanoBananaExecutor.test.ts
  2. 16
      src/store/execution/__tests__/runWithFallback.test.ts
  3. 14
      src/store/execution/nanoBananaExecutor.ts
  4. 6
      src/store/execution/runWithFallback.ts
  5. 14
      src/store/utils/defaultImageModel.ts

62
src/store/execution/__tests__/nanoBananaExecutor.test.ts

@ -195,6 +195,68 @@ describe("executeNanoBanana", () => {
}); });
}); });
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 () => { it("reroutes stored Gemini image nodes to NewApiWG when Gemini key is missing", async () => {
const node = makeNode(); const node = makeNode();
mockFetch.mockResolvedValueOnce({ mockFetch.mockResolvedValueOnce({

16
src/store/execution/__tests__/runWithFallback.test.ts

@ -73,6 +73,22 @@ describe("runWithFallback", () => {
expect(runOnce).toHaveBeenCalledTimes(1); expect(runOnce).toHaveBeenCalledTimes(1);
}); });
it("primary fails but predicate rejects fallback: rethrows primary error", async () => {
const runOnce = vi.fn().mockRejectedValueOnce(new Error("primary boom"));
await expect(
runWithFallback({
nodeId: "n1",
primary,
fallback,
updateNodeData,
runOnce,
shouldFallback: () => false,
})
).rejects.toThrow("primary boom");
expect(runOnce).toHaveBeenCalledTimes(1);
});
it("primary fails, fallback succeeds: stamps metadata and clears error", async () => { it("primary fails, fallback succeeds: stamps metadata and clears error", async () => {
const runOnce = vi const runOnce = vi
.fn() .fn()

14
src/store/execution/nanoBananaExecutor.ts

@ -14,6 +14,8 @@ import { buildGenerateHeaders } from "@/store/utils/buildApiHeaders";
import { import {
createDefaultImageModel, createDefaultImageModel,
createNewApiWGDefaultImageModel, createNewApiWGDefaultImageModel,
createNewApiWGDefaultImageFallbackModel,
isNewApiWGDefaultImageModel,
shouldPreferNewApiWGImageModel, shouldPreferNewApiWGImageModel,
} from "@/store/utils/defaultImageModel"; } from "@/store/utils/defaultImageModel";
import { pollGenerateTask } from "./pollTaskCompletion"; import { pollGenerateTask } from "./pollTaskCompletion";
@ -31,6 +33,11 @@ export interface NanoBananaOptions {
useStoredFallback?: boolean; useStoredFallback?: boolean;
} }
function isProviderOverloadError(error: unknown): boolean {
const message = error instanceof Error ? error.message : String(error);
return /provider is temporarily overloaded|system\s+disk\s+overloaded/i.test(message);
}
async function loadSavedGenerationImage(directoryPath: string, imageId: string): Promise<string | null> { async function loadSavedGenerationImage(directoryPath: string, imageId: string): Promise<string | null> {
const response = await fetch("/api/load-generation", { const response = await fetch("/api/load-generation", {
method: "POST", method: "POST",
@ -419,6 +426,7 @@ export async function executeNanoBanana(
nodeData.selectedModel ?? createDefaultImageModel(nodeData.model, providerSettings, providerEnvStatus); nodeData.selectedModel ?? createDefaultImageModel(nodeData.model, providerSettings, providerEnvStatus);
let primaryParameters: Record<string, unknown> | undefined; let primaryParameters: Record<string, unknown> | undefined;
let fallbackModel = nodeData.fallbackModel; let fallbackModel = nodeData.fallbackModel;
let shouldFallback: ((error: unknown) => boolean) | undefined;
if ( if (
primaryModel.provider === "gemini" && primaryModel.provider === "gemini" &&
@ -447,6 +455,11 @@ export async function executeNanoBanana(
}); });
} }
if (!fallbackModel && isNewApiWGDefaultImageModel(primaryModel)) {
fallbackModel = createNewApiWGDefaultImageFallbackModel();
shouldFallback = isProviderOverloadError;
}
await runWithFallback({ await runWithFallback({
nodeId: node.id, nodeId: node.id,
primary: primaryModel, primary: primaryModel,
@ -455,6 +468,7 @@ export async function executeNanoBanana(
fallbackParameters: nodeData.fallbackParameters, fallbackParameters: nodeData.fallbackParameters,
updateNodeData, updateNodeData,
runOnce, runOnce,
shouldFallback,
clearOutput: { outputImage: null }, clearOutput: { outputImage: null },
}); });
} }

6
src/store/execution/runWithFallback.ts

@ -23,6 +23,7 @@ export interface RunWithFallbackOptions {
fallbackParameters?: Record<string, unknown>; fallbackParameters?: Record<string, unknown>;
updateNodeData: (id: string, data: Partial<WorkflowNodeData>) => void; updateNodeData: (id: string, data: Partial<WorkflowNodeData>) => void;
runOnce: (model: SelectedModel, parametersOverride?: Record<string, unknown>) => Promise<void>; runOnce: (model: SelectedModel, parametersOverride?: Record<string, unknown>) => Promise<void>;
shouldFallback?: (error: unknown) => boolean;
/** Data to merge when transitioning to fallback (e.g. { outputImage: null }) */ /** Data to merge when transitioning to fallback (e.g. { outputImage: null }) */
clearOutput?: Partial<WorkflowNodeData>; clearOutput?: Partial<WorkflowNodeData>;
} }
@ -44,7 +45,7 @@ function isSameModel(a: SelectedModel, b: SelectedModel): boolean {
export async function runWithFallback( export async function runWithFallback(
options: RunWithFallbackOptions options: RunWithFallbackOptions
): Promise<void> { ): Promise<void> {
const { nodeId, primary, primaryParameters, fallback, fallbackParameters, updateNodeData, runOnce, clearOutput } = options; const { nodeId, primary, primaryParameters, fallback, fallbackParameters, updateNodeData, runOnce, shouldFallback, clearOutput } = options;
// Clear any prior fallback metadata before we start. // Clear any prior fallback metadata before we start.
updateNodeData(nodeId, { updateNodeData(nodeId, {
@ -68,6 +69,9 @@ export async function runWithFallback(
} }
const primaryErrMsg = errorMessage(primaryError); const primaryErrMsg = errorMessage(primaryError);
if (shouldFallback && !shouldFallback(primaryError)) {
throw primaryError;
}
// Clear error state and stale output, show the fallback is now running. // Clear error state and stale output, show the fallback is now running.
updateNodeData(nodeId, { updateNodeData(nodeId, {

14
src/store/utils/defaultImageModel.ts

@ -2,6 +2,7 @@ import type { ModelType, ProviderEnvStatus, ProviderSettings, ProviderType, Sele
import { MODEL_DISPLAY_NAMES } from "@/types"; import { MODEL_DISPLAY_NAMES } from "@/types";
const NEWAPIWG_DEFAULT_IMAGE_MODEL_ID = "apiyi_nano_banana_2"; const NEWAPIWG_DEFAULT_IMAGE_MODEL_ID = "apiyi_nano_banana_2";
const NEWAPIWG_DEFAULT_IMAGE_FALLBACK_MODEL_ID = "gpt-image-2-all";
const NEWAPIWG_DEFAULT_VIDEO_MODEL_ID = "doubao-seedance-2-0-260128"; const NEWAPIWG_DEFAULT_VIDEO_MODEL_ID = "doubao-seedance-2-0-260128";
const DEFAULT_AUDIO_MODEL_ID = "elevenlabs/turbo-v2.5"; const DEFAULT_AUDIO_MODEL_ID = "elevenlabs/turbo-v2.5";
@ -14,6 +15,19 @@ export function createNewApiWGDefaultImageModel(): SelectedModel {
}; };
} }
export function createNewApiWGDefaultImageFallbackModel(): SelectedModel {
return {
provider: "newapiwg",
modelId: NEWAPIWG_DEFAULT_IMAGE_FALLBACK_MODEL_ID,
displayName: NEWAPIWG_DEFAULT_IMAGE_FALLBACK_MODEL_ID,
capabilities: ["text-to-image", "image-to-image"],
};
}
export function isNewApiWGDefaultImageModel(model: SelectedModel): boolean {
return model.provider === "newapiwg" && model.modelId === NEWAPIWG_DEFAULT_IMAGE_MODEL_ID;
}
export function createNewApiWGDefaultVideoModel(): SelectedModel { export function createNewApiWGDefaultVideoModel(): SelectedModel {
return { return {
provider: "newapiwg", provider: "newapiwg",

Loading…
Cancel
Save