From 4c3847bfcb1daf8308cc93f480a17467d068f39b Mon Sep 17 00:00:00 2001 From: shrimbly Date: Sun, 8 Feb 2026 22:51:52 +1300 Subject: [PATCH] refactor: extract API header builder from workflow store Replace 6 duplicated header-building if/else chains (across executeWorkflow and regenerateNode) with buildGenerateHeaders() and buildLlmHeaders() utility functions. Add 13 new tests. Co-Authored-By: Claude Opus 4.6 --- .../utils/__tests__/buildApiHeaders.test.ts | 96 +++++++++++ src/store/utils/buildApiHeaders.ts | 71 ++++++++ src/store/workflowStore.ts | 157 +----------------- 3 files changed, 174 insertions(+), 150 deletions(-) create mode 100644 src/store/utils/__tests__/buildApiHeaders.test.ts create mode 100644 src/store/utils/buildApiHeaders.ts diff --git a/src/store/utils/__tests__/buildApiHeaders.test.ts b/src/store/utils/__tests__/buildApiHeaders.test.ts new file mode 100644 index 00000000..ce13d1eb --- /dev/null +++ b/src/store/utils/__tests__/buildApiHeaders.test.ts @@ -0,0 +1,96 @@ +import { describe, it, expect } from "vitest"; +import { buildGenerateHeaders, buildLlmHeaders } from "../buildApiHeaders"; +import type { ProviderSettings } from "@/types"; + +function makeSettings(overrides: Partial> = {}): ProviderSettings { + const defaults: Record = { + gemini: { id: "gemini", name: "Gemini", enabled: true, apiKey: null }, + replicate: { id: "replicate", name: "Replicate", enabled: true, apiKey: null }, + fal: { id: "fal", name: "fal.ai", enabled: true, apiKey: null }, + kie: { id: "kie", name: "Kie.ai", enabled: true, apiKey: null }, + wavespeed: { id: "wavespeed", name: "WaveSpeed", enabled: true, apiKey: null }, + openai: { id: "openai", name: "OpenAI", enabled: true, apiKey: null }, + }; + for (const [key, val] of Object.entries(overrides)) { + if (defaults[key]) { + defaults[key].apiKey = val.apiKey; + } + } + return { providers: defaults } as unknown as ProviderSettings; +} + +describe("buildGenerateHeaders", () => { + it("should always include Content-Type", () => { + const headers = buildGenerateHeaders("gemini", makeSettings()); + expect(headers["Content-Type"]).toBe("application/json"); + }); + + it("should add Gemini API key header", () => { + const settings = makeSettings({ gemini: { apiKey: "gem-key-123" } }); + const headers = buildGenerateHeaders("gemini", settings); + expect(headers["X-Gemini-API-Key"]).toBe("gem-key-123"); + }); + + it("should add Replicate API key header", () => { + const settings = makeSettings({ replicate: { apiKey: "rep-key" } }); + const headers = buildGenerateHeaders("replicate", settings); + expect(headers["X-Replicate-API-Key"]).toBe("rep-key"); + }); + + it("should add fal.ai API key header", () => { + const settings = makeSettings({ fal: { apiKey: "fal-key" } }); + const headers = buildGenerateHeaders("fal", settings); + expect(headers["X-Fal-API-Key"]).toBe("fal-key"); + }); + + it("should add Kie API key header", () => { + const settings = makeSettings({ kie: { apiKey: "kie-key" } }); + const headers = buildGenerateHeaders("kie", settings); + expect(headers["X-Kie-Key"]).toBe("kie-key"); + }); + + it("should add WaveSpeed API key header", () => { + const settings = makeSettings({ wavespeed: { apiKey: "ws-key" } }); + const headers = buildGenerateHeaders("wavespeed", settings); + expect(headers["X-WaveSpeed-Key"]).toBe("ws-key"); + }); + + it("should not add header when API key is null", () => { + const headers = buildGenerateHeaders("gemini", makeSettings()); + expect(headers["X-Gemini-API-Key"]).toBeUndefined(); + }); + + it("should handle unknown provider gracefully", () => { + const headers = buildGenerateHeaders("unknown", makeSettings()); + expect(headers).toEqual({ "Content-Type": "application/json" }); + }); +}); + +describe("buildLlmHeaders", () => { + it("should always include Content-Type", () => { + const headers = buildLlmHeaders("google", makeSettings()); + expect(headers["Content-Type"]).toBe("application/json"); + }); + + it("should add Gemini API key for google provider", () => { + const settings = makeSettings({ gemini: { apiKey: "gem-key" } }); + const headers = buildLlmHeaders("google", settings); + expect(headers["X-Gemini-API-Key"]).toBe("gem-key"); + }); + + it("should add OpenAI API key for openai provider", () => { + const settings = makeSettings({ openai: { apiKey: "oai-key" } }); + const headers = buildLlmHeaders("openai", settings); + expect(headers["X-OpenAI-API-Key"]).toBe("oai-key"); + }); + + it("should not add header when API key is null", () => { + const headers = buildLlmHeaders("google", makeSettings()); + expect(headers["X-Gemini-API-Key"]).toBeUndefined(); + }); + + it("should handle unknown LLM provider gracefully", () => { + const headers = buildLlmHeaders("unknown", makeSettings()); + expect(headers).toEqual({ "Content-Type": "application/json" }); + }); +}); diff --git a/src/store/utils/buildApiHeaders.ts b/src/store/utils/buildApiHeaders.ts new file mode 100644 index 00000000..6df5099d --- /dev/null +++ b/src/store/utils/buildApiHeaders.ts @@ -0,0 +1,71 @@ +/** + * API Header Builder Utilities + * + * Centralizes the duplicated header-building logic for API calls + * across executeWorkflow and regenerateNode. + */ + +import { ProviderType, ProviderSettings, LLMProvider } from "@/types"; + +/** + * Header name mapping for each provider + */ +const PROVIDER_HEADER_MAP: Record = { + gemini: "X-Gemini-API-Key", + replicate: "X-Replicate-API-Key", + fal: "X-Fal-API-Key", + kie: "X-Kie-Key", + wavespeed: "X-WaveSpeed-Key", + openai: "X-OpenAI-API-Key", +}; + +/** + * Build headers for image/video generation API calls. + * Adds the appropriate API key header based on the provider. + */ +export function buildGenerateHeaders( + provider: ProviderType | string, + providerSettings: ProviderSettings +): Record { + const headers: Record = { + "Content-Type": "application/json", + }; + + const providerKey = provider as ProviderType; + const headerName = PROVIDER_HEADER_MAP[providerKey]; + if (headerName) { + const config = providerSettings.providers[providerKey]; + if (config?.apiKey) { + headers[headerName] = config.apiKey; + } + } + + return headers; +} + +/** + * Build headers for LLM API calls. + * Maps LLM provider names ("google", "openai") to their API key headers. + */ +export function buildLlmHeaders( + llmProvider: LLMProvider | string, + providerSettings: ProviderSettings +): Record { + const headers: Record = { + "Content-Type": "application/json", + }; + + if (llmProvider === "google") { + const geminiConfig = providerSettings.providers.gemini; + if (geminiConfig?.apiKey) { + headers["X-Gemini-API-Key"] = geminiConfig.apiKey; + } + } else if (llmProvider === "openai") { + const openaiConfig = providerSettings.providers.openai; + if (openaiConfig?.apiKey) { + headers["X-OpenAI-API-Key"] = openaiConfig.apiKey; + } + } + + return headers; +} diff --git a/src/store/workflowStore.ts b/src/store/workflowStore.ts index a8db1854..fef60625 100644 --- a/src/store/workflowStore.ts +++ b/src/store/workflowStore.ts @@ -58,6 +58,7 @@ import { GROUP_COLORS, GROUP_COLOR_ORDER, } from "./utils/nodeDefaults"; +import { buildGenerateHeaders, buildLlmHeaders } from "./utils/buildApiHeaders"; export type EdgeStyle = "angular" | "curved"; @@ -1296,37 +1297,8 @@ export const useWorkflowStore = create((set, get) => ({ dynamicInputs, // Pass dynamic inputs for schema-mapped connections }; - // Build headers with API keys for providers - const headers: Record = { - "Content-Type": "application/json", - }; const provider = nodeData.selectedModel?.provider || "gemini"; - if (provider === "gemini") { - const geminiConfig = providerSettingsState.providers.gemini; - if (geminiConfig?.apiKey) { - headers["X-Gemini-API-Key"] = geminiConfig.apiKey; - } - } else if (provider === "replicate") { - const replicateConfig = providerSettingsState.providers.replicate; - if (replicateConfig?.apiKey) { - headers["X-Replicate-API-Key"] = replicateConfig.apiKey; - } - } else if (provider === "fal") { - const falConfig = providerSettingsState.providers.fal; - if (falConfig?.apiKey) { - headers["X-Fal-API-Key"] = falConfig.apiKey; - } - } else if (provider === "kie") { - const kieConfig = providerSettingsState.providers.kie; - if (kieConfig?.apiKey) { - headers["X-Kie-Key"] = kieConfig.apiKey; - } - } else if (provider === "wavespeed") { - const wavespeedConfig = providerSettingsState.providers.wavespeed; - if (wavespeedConfig?.apiKey) { - headers["X-WaveSpeed-Key"] = wavespeedConfig.apiKey; - } - } + const headers = buildGenerateHeaders(provider, providerSettingsState); logger.info('node.execution', `Calling ${provider} API for image generation`, { nodeId: node.id, @@ -1525,37 +1497,8 @@ export const useWorkflowStore = create((set, get) => ({ mediaType: "video" as const, // Signal to API to use queue for long-running video generation }; - // Build headers with API keys for providers - const headers: Record = { - "Content-Type": "application/json", - }; const provider = nodeData.selectedModel.provider; - if (provider === "gemini") { - const geminiConfig = providerSettingsState.providers.gemini; - if (geminiConfig?.apiKey) { - headers["X-Gemini-API-Key"] = geminiConfig.apiKey; - } - } else if (provider === "replicate") { - const replicateConfig = providerSettingsState.providers.replicate; - if (replicateConfig?.apiKey) { - headers["X-Replicate-API-Key"] = replicateConfig.apiKey; - } - } else if (provider === "fal") { - const falConfig = providerSettingsState.providers.fal; - if (falConfig?.apiKey) { - headers["X-Fal-API-Key"] = falConfig.apiKey; - } - } else if (provider === "kie") { - const kieConfig = providerSettingsState.providers.kie; - if (kieConfig?.apiKey) { - headers["X-Kie-Key"] = kieConfig.apiKey; - } - } else if (provider === "wavespeed") { - const wavespeedConfig = providerSettingsState.providers.wavespeed; - if (wavespeedConfig?.apiKey) { - headers["X-WaveSpeed-Key"] = wavespeedConfig.apiKey; - } - } + const headers = buildGenerateHeaders(provider, providerSettingsState); logger.info('node.execution', `Calling ${provider} API for video generation`, { nodeId: node.id, provider, @@ -1735,21 +1678,7 @@ export const useWorkflowStore = create((set, get) => ({ const nodeData = node.data as LLMGenerateNodeData; const providerSettingsState = get().providerSettings; - // Build headers with API keys for LLM providers - const headers: Record = { - "Content-Type": "application/json", - }; - if (nodeData.provider === "google") { - const geminiConfig = providerSettingsState.providers.gemini; - if (geminiConfig?.apiKey) { - headers["X-Gemini-API-Key"] = geminiConfig.apiKey; - } - } else if (nodeData.provider === "openai") { - const openaiConfig = providerSettingsState.providers.openai; - if (openaiConfig?.apiKey) { - headers["X-OpenAI-API-Key"] = openaiConfig.apiKey; - } - } + const headers = buildLlmHeaders(nodeData.provider, providerSettingsState); logger.info('api.llm', 'Calling LLM API', { nodeId: node.id, @@ -2403,36 +2332,7 @@ export const useWorkflowStore = create((set, get) => ({ error: null, }); - // Build headers with API keys for providers - const headers: Record = { - "Content-Type": "application/json", - }; - if (provider === "gemini") { - const geminiConfig = providerSettingsState.providers.gemini; - if (geminiConfig?.apiKey) { - headers["X-Gemini-API-Key"] = geminiConfig.apiKey; - } - } else if (provider === "replicate") { - const replicateConfig = providerSettingsState.providers.replicate; - if (replicateConfig?.apiKey) { - headers["X-Replicate-API-Key"] = replicateConfig.apiKey; - } - } else if (provider === "fal") { - const falConfig = providerSettingsState.providers.fal; - if (falConfig?.apiKey) { - headers["X-Fal-API-Key"] = falConfig.apiKey; - } - } else if (provider === "kie") { - const kieConfig = providerSettingsState.providers.kie; - if (kieConfig?.apiKey) { - headers["X-Kie-Key"] = kieConfig.apiKey; - } - } else if (provider === "wavespeed") { - const wavespeedConfig = providerSettingsState.providers.wavespeed; - if (wavespeedConfig?.apiKey) { - headers["X-WaveSpeed-Key"] = wavespeedConfig.apiKey; - } - } + const headers = buildGenerateHeaders(provider, providerSettingsState); logger.info('node.execution', `Calling ${provider} API for node regeneration`, { nodeId, @@ -2575,21 +2475,7 @@ export const useWorkflowStore = create((set, get) => ({ const providerSettingsState = get().providerSettings; - // Build headers with API keys for LLM providers - const headers: Record = { - "Content-Type": "application/json", - }; - if (nodeData.provider === "google") { - const geminiConfig = providerSettingsState.providers.gemini; - if (geminiConfig?.apiKey) { - headers["X-Gemini-API-Key"] = geminiConfig.apiKey; - } - } else if (nodeData.provider === "openai") { - const openaiConfig = providerSettingsState.providers.openai; - if (openaiConfig?.apiKey) { - headers["X-OpenAI-API-Key"] = openaiConfig.apiKey; - } - } + const headers = buildLlmHeaders(nodeData.provider, providerSettingsState); logger.info('api.llm', 'Calling LLM API for node regeneration', { nodeId, @@ -2694,37 +2580,8 @@ export const useWorkflowStore = create((set, get) => ({ error: null, }); - // Build headers with API keys for providers - const headers: Record = { - "Content-Type": "application/json", - }; const provider = nodeData.selectedModel.provider; - if (provider === "gemini") { - const geminiConfig = providerSettingsState.providers.gemini; - if (geminiConfig?.apiKey) { - headers["X-Gemini-API-Key"] = geminiConfig.apiKey; - } - } else if (provider === "replicate") { - const replicateConfig = providerSettingsState.providers.replicate; - if (replicateConfig?.apiKey) { - headers["X-Replicate-API-Key"] = replicateConfig.apiKey; - } - } else if (provider === "fal") { - const falConfig = providerSettingsState.providers.fal; - if (falConfig?.apiKey) { - headers["X-Fal-API-Key"] = falConfig.apiKey; - } - } else if (provider === "kie") { - const kieConfig = providerSettingsState.providers.kie; - if (kieConfig?.apiKey) { - headers["X-Kie-Key"] = kieConfig.apiKey; - } - } else if (provider === "wavespeed") { - const wavespeedConfig = providerSettingsState.providers.wavespeed; - if (wavespeedConfig?.apiKey) { - headers["X-WaveSpeed-Key"] = wavespeedConfig.apiKey; - } - } + const headers = buildGenerateHeaders(provider, providerSettingsState); logger.info('node.execution', `Calling ${provider} API for video regeneration`, { nodeId, provider,