From 10b23b15ae421ed12753773a77b819b032071c3d Mon Sep 17 00:00:00 2001
From: diegosouzapw <diegosouzapw@users.noreply.github.com>
Date: Sat, 7 Mar 2026 06:56:09 -0300
Subject: [PATCH] fix: custom image model routing + Codex OAuth workspace
 isolation (#232, #236)

---
 package-lock.json                             |  4 +--
 .../api/oauth/[provider]/[action]/route.ts    | 36 ++++++++++++++-----
 src/app/api/v1/images/generations/route.ts    | 23 +++++++++++-
 3 files changed, 51 insertions(+), 12 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index 174dca4901..41a9341fca 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "omniroute",
-  "version": "2.0.5",
+  "version": "2.0.6",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "omniroute",
-      "version": "2.0.5",
+      "version": "2.0.6",
       "hasInstallScript": true,
       "license": "MIT",
       "workspaces": [
diff --git a/src/app/api/oauth/[provider]/[action]/route.ts b/src/app/api/oauth/[provider]/[action]/route.ts
index f16aa66105..ead1f38175 100644
--- a/src/app/api/oauth/[provider]/[action]/route.ts
+++ b/src/app/api/oauth/[provider]/[action]/route.ts
@@ -221,9 +221,15 @@ export async function POST(
       let connection: any;
       if (tokenData.email) {
         const existing = await getProviderConnections({ provider });
-        const match = existing.find(
-          (c: any) => c.email === tokenData.email && c.authType === "oauth"
-        );
+        const match = existing.find((c: any) => {
+          if (c.email !== tokenData.email || c.authType !== "oauth") return false;
+          // For Codex, also check workspaceId to avoid overwriting different workspace connections
+          if (provider === "codex" && tokenData.providerSpecificData?.workspaceId) {
+            const existingWorkspace = c.providerSpecificData?.workspaceId;
+            return existingWorkspace === tokenData.providerSpecificData.workspaceId;
+          }
+          return true;
+        });
         const matchId = typeof match?.id === "string" ? match.id : null;
         if (matchId) {
           connection = await updateProviderConnection(matchId, {
@@ -285,9 +291,15 @@ export async function POST(
         let connection: any;
         if (result.tokens.email) {
           const existing = await getProviderConnections({ provider });
-          const match = existing.find(
-            (c: any) => c.email === result.tokens.email && c.authType === "oauth"
-          );
+          const match = existing.find((c: any) => {
+            if (c.email !== result.tokens.email || c.authType !== "oauth") return false;
+            // For Codex, also check workspaceId to avoid overwriting different workspace connections
+            if (provider === "codex" && result.tokens.providerSpecificData?.workspaceId) {
+              const existingWorkspace = c.providerSpecificData?.workspaceId;
+              return existingWorkspace === result.tokens.providerSpecificData.workspaceId;
+            }
+            return true;
+          });
           const matchId = typeof match?.id === "string" ? match.id : null;
           if (matchId) {
             connection = await updateProviderConnection(matchId, {
@@ -399,9 +411,15 @@ export async function POST(
         let connection: any;
         if (tokenData.email) {
           const existing = await getProviderConnections({ provider });
-          const match = existing.find(
-            (c: any) => c.email === tokenData.email && c.authType === "oauth"
-          );
+          const match = existing.find((c: any) => {
+            if (c.email !== tokenData.email || c.authType !== "oauth") return false;
+            // For Codex, also check workspaceId to avoid overwriting different workspace connections
+            if (provider === "codex" && tokenData.providerSpecificData?.workspaceId) {
+              const existingWorkspace = c.providerSpecificData?.workspaceId;
+              return existingWorkspace === tokenData.providerSpecificData.workspaceId;
+            }
+            return true;
+          });
           const matchId = typeof match?.id === "string" ? match.id : null;
           if (matchId) {
             connection = await updateProviderConnection(matchId, {
diff --git a/src/app/api/v1/images/generations/route.ts b/src/app/api/v1/images/generations/route.ts
index 7ea2091d05..d53207a927 100644
--- a/src/app/api/v1/images/generations/route.ts
+++ b/src/app/api/v1/images/generations/route.ts
@@ -107,7 +107,28 @@ export async function POST(request) {
   if (policy.rejection) return policy.rejection;
 
   // Parse model to get provider
-  const { provider } = parseImageModel(body.model);
+  let { provider } = parseImageModel(body.model);
+
+  // If not in built-in registry, check custom models tagged for images
+  if (!provider) {
+    try {
+      const customModelsMap = (await getAllCustomModels()) as Record<string, any>;
+      for (const [providerId, models] of Object.entries(customModelsMap)) {
+        if (!Array.isArray(models)) continue;
+        for (const model of models) {
+          if (!model?.id || !Array.isArray(model.supportedEndpoints)) continue;
+          if (!model.supportedEndpoints.includes("images")) continue;
+          const fullId = `${providerId}/${model.id}`;
+          if (fullId === body.model) {
+            provider = providerId;
+            break;
+          }
+        }
+        if (provider) break;
+      }
+    } catch {}
+  }
+
   if (!provider) {
     return errorResponse(
       HTTP_STATUS.BAD_REQUEST,