fix(images): refresh OAuth credentials and rotate accounts after 401

src/app/api/v1/images/generations/route.ts

@@ -19,7 +19,8 @@ import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy";
 import { v1ImageGenerationSchema } from "@/shared/validation/schemas";
 import { isValidationFailure, validateBody } from "@/shared/validation/helpers";
 
-import { getAllCustomModels, resolveProxyForConnection } from "@/lib/localDb";
+import { getAllCustomModels } from "@/lib/db/models";
+import { resolveProxyForConnection } from "@/lib/db/settings";
 import { resolveImageRouteModel } from "@/lib/images/imageRouteModel";
 import { runWithProxyContext } from "@omniroute/open-sse/utils/proxyFetch.ts";
 import { attachOmniRouteMetaHeaders } from "@/domain/omnirouteResponseMeta";
@@ -28,6 +29,7 @@ import { generateRequestId } from "@/shared/utils/requestId";
 import { getSpecialtyModelsResponse } from "@/app/api/v1/_shared/specialtyCatalog";
 import { enforceClientApiRouteAuth } from "@/shared/utils/clientApiRouteAuth";
 import { runWithCallLogApiKeyContext } from "@/lib/usage/callLogApiKeyContext";
+import { executeImageWithCredentialFallback } from "@/sse/services/imageCredentialRetry";
 
 export const dynamic = "force-dynamic";
 
@@ -121,7 +123,7 @@ async function postHandler(request, context) {
   body.model = await resolveImageRouteModel(body.model);
 
   // Parse model to get provider
-  let { provider } = parseImageModel(body.model);
+  let { provider, model: requestedModel } = parseImageModel(body.model);
   let isCustomModel = false;
 
   // If not in built-in registry, check custom models tagged for images
@@ -136,6 +138,7 @@ async function postHandler(request, context) {
           const fullId = `${providerId}/${model.id}`;
           if (fullId === body.model) {
             provider = providerId;
+            requestedModel = model.id;
             isCustomModel = true;
             break;
           }
@@ -184,7 +187,12 @@ async function postHandler(request, context) {
   // Get credentials — skip for local providers (authType: "none")
   let credentials = null;
   if (providerConfig && providerConfig.authType !== "none") {
-    credentials = await getProviderCredentialsWithQuotaPreflight(provider);
+    credentials = await getProviderCredentialsWithQuotaPreflight(
+      provider,
+      null,
+      null,
+      requestedModel
+    );
     if (!credentials) {
       return errorResponse(
         HTTP_STATUS.BAD_REQUEST,
@@ -200,7 +208,12 @@ async function postHandler(request, context) {
       );
     }
   } else if (isCustomModel) {
-    credentials = await getProviderCredentialsWithQuotaPreflight(provider);
+    credentials = await getProviderCredentialsWithQuotaPreflight(
+      provider,
+      null,
+      null,
+      requestedModel
+    );
     if (!credentials) {
       return errorResponse(
         HTTP_STATUS.BAD_REQUEST,
@@ -219,47 +232,59 @@ async function postHandler(request, context) {
     // #6928: best-effort per-connection base-URL override lookup for local
     // no-auth media providers (ComfyUI). A connection is optional here — unlike
     // the authType !== "none" branch above, we never 400 when none exists.
-    const localCredentials = await getProviderCredentialsWithQuotaPreflight(provider);
+    const localCredentials = await getProviderCredentialsWithQuotaPreflight(
+      provider,
+      null,
+      null,
+      requestedModel
+    );
     if (localCredentials && !isAllRateLimitedCredentials(localCredentials)) {
       credentials = localCredentials;
     }
   }
 
-  // Resolve proxy for the connection if credentials exist (#1904)
-  let proxyInfo = null;
-  if (credentials?.connectionId) {
-    try {
-      proxyInfo = await resolveProxyForConnection(credentials.connectionId);
-    } catch {
-      log.debug("PROXY", `Failed to resolve proxy for image provider: ${provider}`);
-    }
-  }
+  const execution = await executeImageWithCredentialFallback({
+    provider,
+    requestedModel,
+    credentials,
+    execute: async (attemptCredentials) => {
+      let proxyInfo = null;
+      if (attemptCredentials?.connectionId) {
+        try {
+          proxyInfo = await resolveProxyForConnection(attemptCredentials.connectionId);
+        } catch {
+          log.debug("PROXY", `Failed to resolve proxy for image provider: ${provider}`);
+        }
+      }
 
-  const generateImage = () =>
-    runWithCallLogApiKeyContext(
-      {
-        apiKeyId: policy.apiKeyInfo?.id ?? null,
-        apiKeyName: policy.apiKeyInfo?.name ?? null,
-      },
-      () =>
-        handleImageGeneration({
-          body,
-          credentials,
-          log,
-          ...(isCustomModel && { resolvedProvider: provider }),
-          signal: request.signal,
-          clientHeaders: publicBaseUrlHeaders(request.headers),
-        })
-    );
+      const generateImage = () =>
+        runWithCallLogApiKeyContext(
+          {
+            apiKeyId: policy.apiKeyInfo?.id ?? null,
+            apiKeyName: policy.apiKeyInfo?.name ?? null,
+          },
+          () =>
+            handleImageGeneration({
+              body,
+              credentials: attemptCredentials,
+              log,
+              ...(isCustomModel && { resolvedProvider: provider }),
+              signal: request.signal,
+              clientHeaders: publicBaseUrlHeaders(request.headers),
+            })
+        );
 
-  // Execute with proxy context when available, direct otherwise (#1904)
-  const result = await (credentials?.connectionId
-    ? runWithProxyContext(proxyInfo?.proxy || null, generateImage).catch((err: any) => ({
-        success: false,
-        status: err.statusCode || 500,
-        error: err.message,
-      }))
-    : generateImage());
+      return attemptCredentials?.connectionId
+        ? runWithProxyContext(proxyInfo?.proxy || null, generateImage).catch((err: any) => ({
+            success: false,
+            status: err.statusCode || 500,
+            error: err.message,
+          }))
+        : generateImage();
+    },
+  });
+  credentials = execution.credentials;
+  const result = execution.result;
 
   if (result.success) {
     await clearRecoveredProviderState(credentials);

src/app/api/v1/providers/[provider]/images/generations/route.ts

@@ -13,6 +13,7 @@ import { v1ImageGenerationSchema } from "@/shared/validation/schemas";
 import { isValidationFailure, validateBody } from "@/shared/validation/helpers";
 import { enforceClientApiRouteAuth } from "@/shared/utils/clientApiRouteAuth";
 import { runWithCallLogApiKeyContext } from "@/lib/usage/callLogApiKeyContext";
+import { executeImageWithCredentialFallback } from "@/sse/services/imageCredentialRetry";
 
 /**
  * Handle CORS preflight
@@ -71,7 +72,13 @@ export async function POST(request, { params }) {
     );
   }
 
-  const credentials = await getProviderCredentialsWithQuotaPreflight(rawProvider);
+  const requestedModel = body.model.slice(rawProvider.length + 1);
+  let credentials = await getProviderCredentialsWithQuotaPreflight(
+    rawProvider,
+    null,
+    null,
+    requestedModel
+  );
   if (!credentials) {
     return errorResponse(
       HTTP_STATUS.BAD_REQUEST,
@@ -87,13 +94,21 @@ export async function POST(request, { params }) {
     );
   }
 
-  const result = await runWithCallLogApiKeyContext(
-    {
-      apiKeyId: policy.apiKeyInfo?.id ?? null,
-      apiKeyName: policy.apiKeyInfo?.name ?? null,
-    },
-    () => handleImageGeneration({ body, credentials, log })
-  );
+  const execution = await executeImageWithCredentialFallback({
+    provider: rawProvider,
+    requestedModel,
+    credentials,
+    execute: (attemptCredentials) =>
+      runWithCallLogApiKeyContext(
+        {
+          apiKeyId: policy.apiKeyInfo?.id ?? null,
+          apiKeyName: policy.apiKeyInfo?.name ?? null,
+        },
+        () => handleImageGeneration({ body, credentials: attemptCredentials, log })
+      ),
+  });
+  credentials = execution.credentials;
+  const result = execution.result;
 
   if (result.success) {
     await clearRecoveredProviderState(credentials);

src/sse/services/imageCredentialRetry.ts

@@ -0,0 +1,119 @@
+import { sanitizeErrorMessage } from "@omniroute/open-sse/utils/error.ts";
+
+import { getProviderCredentialsWithQuotaPreflight } from "./auth";
+import { checkAndRefreshToken } from "./tokenRefresh";
+import * as log from "../utils/logger";
+
+interface ImageGenerationResult {
+  success: boolean;
+  status?: number;
+  error?: unknown;
+  data?: unknown;
+}
+
+interface ImageCredentialRetryOptions {
+  provider: string;
+  requestedModel: string | null;
+  credentials: any;
+  execute: (credentials: any) => Promise<ImageGenerationResult>;
+}
+
+interface ImageCredentialRetryResult {
+  credentials: any;
+  result: ImageGenerationResult;
+}
+
+function connectionIdOf(credentials: any): string | null {
+  const connectionId = credentials?.connectionId;
+  return typeof connectionId === "string" && connectionId.trim().length > 0
+    ? connectionId.trim()
+    : null;
+}
+
+function isCredentialSentinel(credentials: any): boolean {
+  return Boolean(credentials?.allRateLimited || credentials?.allExpired);
+}
+
+async function selectNextCredentials(
+  provider: string,
+  requestedModel: string | null,
+  excludedConnectionIds: Set<string>
+) {
+  return getProviderCredentialsWithQuotaPreflight(provider, null, null, requestedModel, {
+    excludeConnectionIds: Array.from(excludedConnectionIds),
+  });
+}
+
+/**
+ * Keep image requests on the same credential lifecycle as chat requests.
+ *
+ * Each connection is attempted at most once. A refresh failure or upstream 401
+ * excludes only that connection for the current request; it does not mutate the
+ * account into a terminal state because another request may refresh it normally.
+ */
+export async function executeImageWithCredentialFallback({
+  provider,
+  requestedModel,
+  credentials,
+  execute,
+}: ImageCredentialRetryOptions): Promise<ImageCredentialRetryResult> {
+  // Local/no-auth image providers intentionally have no credential row. They
+  // still need one direct attempt, but there is no account identity to refresh
+  // or rotate after a 401.
+  if (!credentials) {
+    return { credentials, result: await execute(credentials) };
+  }
+
+  const excludedConnectionIds = new Set<string>();
+  let currentCredentials = credentials;
+  let lastCredentials = credentials;
+  let lastResult: ImageGenerationResult | null = null;
+
+  while (currentCredentials && !isCredentialSentinel(currentCredentials)) {
+    const connectionId = connectionIdOf(currentCredentials);
+    if (connectionId && excludedConnectionIds.has(connectionId)) break;
+    if (connectionId) excludedConnectionIds.add(connectionId);
+
+    try {
+      currentCredentials = await checkAndRefreshToken(provider, currentCredentials);
+    } catch (error) {
+      log.warn("IMAGE", "Credential refresh failed; trying another image-provider account", {
+        provider,
+        connectionId,
+        error: sanitizeErrorMessage(error instanceof Error ? error : new Error(String(error))),
+      });
+      if (!connectionId) throw error;
+      currentCredentials = await selectNextCredentials(
+        provider,
+        requestedModel,
+        excludedConnectionIds
+      );
+      continue;
+    }
+
+    lastCredentials = currentCredentials;
+    lastResult = await execute(currentCredentials);
+    if (lastResult.success || Number(lastResult.status) !== 401 || !connectionId) {
+      return { credentials: lastCredentials, result: lastResult };
+    }
+
+    log.warn("IMAGE", "Image provider rejected credentials; trying another account", {
+      provider,
+      connectionId,
+    });
+    currentCredentials = await selectNextCredentials(
+      provider,
+      requestedModel,
+      excludedConnectionIds
+    );
+  }
+
+  return {
+    credentials: lastCredentials,
+    result: lastResult || {
+      success: false,
+      status: 401,
+      error: "Authentication failed for all eligible image-provider accounts",
+    },
+  };
+}