From 4b5824babcbfef19f4114568a2dad6309e60dd64 Mon Sep 17 00:00:00 2001
From: diegosouzapw <diegosouzapw@users.noreply.github.com>
Date: Sat, 7 Mar 2026 10:03:48 -0300
Subject: [PATCH] fix: pass resolved provider to image handler for custom
 models (#238)

---
 open-sse/handlers/imageGeneration.ts       | 56 +++++++++++++++++++---
 package-lock.json                          |  4 +-
 src/app/api/v1/images/generations/route.ts | 18 ++++++-
 3 files changed, 69 insertions(+), 9 deletions(-)

diff --git a/open-sse/handlers/imageGeneration.ts b/open-sse/handlers/imageGeneration.ts
index a3a417eb45..7e7a15ae7c 100644
--- a/open-sse/handlers/imageGeneration.ts
+++ b/open-sse/handlers/imageGeneration.ts
@@ -30,9 +30,23 @@ import {
  * @param {object} options.body - Request body
  * @param {object} options.credentials - Provider credentials { apiKey, accessToken }
  * @param {object} options.log - Logger
+ * @param {string} [options.resolvedProvider] - Pre-resolved provider ID (from route layer custom model resolution)
  */
-export async function handleImageGeneration({ body, credentials, log }) {
-  const { provider, model } = parseImageModel(body.model);
+export async function handleImageGeneration({ body, credentials, log, resolvedProvider = null }) {
+  let provider, model;
+
+  if (resolvedProvider) {
+    // Provider was already resolved by the route layer (custom model from DB)
+    // Extract model name from the full "provider/model" string
+    provider = resolvedProvider;
+    const modelStr = body.model || "";
+    model = modelStr.startsWith(provider + "/") ? modelStr.slice(provider.length + 1) : modelStr;
+  } else {
+    // Standard path: resolve from built-in image registry
+    const parsed = parseImageModel(body.model);
+    provider = parsed.provider;
+    model = parsed.model;
+  }
 
   if (!provider) {
     return {
@@ -43,12 +57,42 @@ export async function handleImageGeneration({ body, credentials, log }) {
   }
 
   const providerConfig = getImageProvider(provider);
+
+  // For custom models without a built-in provider config, use OpenAI-compatible handler
+  // with a synthetic config based on the provider's credentials
   if (!providerConfig) {
-    return {
-      success: false,
-      status: 400,
-      error: `Unknown image provider: ${provider}`,
+    if (!resolvedProvider) {
+      return {
+        success: false,
+        status: 400,
+        error: `Unknown image provider: ${provider}`,
+      };
+    }
+
+    // Custom model: use OpenAI-compatible format with provider's base URL
+    // The credentials were already resolved by the route layer
+    if (log) {
+      log.info("IMAGE", `Custom model ${provider}/${model} — using OpenAI-compatible handler`);
+    }
+
+    const syntheticConfig = {
+      id: provider,
+      baseUrl:
+        credentials?.baseUrl ||
+        `https://generativelanguage.googleapis.com/v1beta/openai/images/generations`,
+      authType: "apikey",
+      authHeader: "bearer",
+      format: "openai",
     };
+
+    return handleOpenAIImageGeneration({
+      model,
+      provider,
+      providerConfig: syntheticConfig,
+      body,
+      credentials,
+      log,
+    });
   }
 
   // Route to format-specific handler
diff --git a/package-lock.json b/package-lock.json
index 41a9341fca..1433310cd3 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "omniroute",
-  "version": "2.0.6",
+  "version": "2.0.7",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "omniroute",
-      "version": "2.0.6",
+      "version": "2.0.7",
       "hasInstallScript": true,
       "license": "MIT",
       "workspaces": [
diff --git a/src/app/api/v1/images/generations/route.ts b/src/app/api/v1/images/generations/route.ts
index d53207a927..c64da344b3 100644
--- a/src/app/api/v1/images/generations/route.ts
+++ b/src/app/api/v1/images/generations/route.ts
@@ -108,6 +108,7 @@ export async function POST(request) {
 
   // Parse model to get provider
   let { provider } = parseImageModel(body.model);
+  let isCustomModel = false;
 
   // If not in built-in registry, check custom models tagged for images
   if (!provider) {
@@ -121,6 +122,7 @@ export async function POST(request) {
           const fullId = `${providerId}/${model.id}`;
           if (fullId === body.model) {
             provider = providerId;
+            isCustomModel = true;
             break;
           }
         }
@@ -149,9 +151,23 @@ export async function POST(request) {
         `No credentials for image provider: ${provider}`
       );
     }
+  } else if (isCustomModel) {
+    // Custom models need credentials from the provider connection
+    credentials = await getProviderCredentials(provider);
+    if (!credentials) {
+      return errorResponse(
+        HTTP_STATUS.BAD_REQUEST,
+        `No credentials for custom image provider: ${provider}`
+      );
+    }
   }
 
-  const result = await handleImageGeneration({ body, credentials, log });
+  const result = await handleImageGeneration({
+    body,
+    credentials,
+    log,
+    ...(isCustomModel && { resolvedProvider: provider }),
+  });
 
   if (result.success) {
     return new Response(JSON.stringify((result as any).data), {