feat(vertex): dynamic model discovery (surface image models + full catalog)

src/app/api/providers/[id]/models/route.ts

@@ -77,6 +77,7 @@ import {
   isAutoFetchModelsEnabled,
   persistDiscoveredModels,
 } from "@/lib/providerModels/modelDiscovery";
+import { parseGeminiModelsList, type GeminiDiscoveryModel } from "@/lib/providerModels/geminiModelsParser";
 import { getSyncedAvailableModels } from "@/lib/db/models";
 import { fetchCursorAgentModels } from "@/lib/providerModels/cursorAgent";
 
@@ -399,50 +400,7 @@ const PROVIDER_MODELS_CONFIG: Record<string, ProviderModelsConfigEntry> = {
     method: "GET",
     headers: { "Content-Type": "application/json" },
     authQuery: "key", // Use query param for API key
-    parseResponse: (data) => {
-      const METHOD_TO_ENDPOINT: Record<string, string> = {
-        generateContent: "chat",
-        embedContent: "embeddings",
-        predict: "images",
-        predictLongRunning: "images",
-        bidiGenerateContent: "audio",
-        generateAnswer: "chat",
-      };
-      const IGNORED_METHODS = new Set([
-        "countTokens",
-        "countTextTokens",
-        "createCachedContent",
-        "batchGenerateContent",
-        "asyncBatchEmbedContent",
-      ]);
-
-      return (data.models || []).map((m: Record<string, unknown>) => {
-        const methods: string[] = Array.isArray(m.supportedGenerationMethods)
-          ? m.supportedGenerationMethods
-          : [];
-        const endpoints = [
-          ...new Set(
-            methods
-              .filter((method) => !IGNORED_METHODS.has(method))
-              .map((method) => METHOD_TO_ENDPOINT[method] || "chat")
-          ),
-        ];
-        if (endpoints.length === 0) endpoints.push("chat");
-
-        return {
-          ...m,
-          id: ((m.name as string) || (m.id as string) || "").replace(/^models\//, ""),
-          name: (m.displayName as string) || ((m.name as string) || "").replace(/^models\//, ""),
-          supportedEndpoints: endpoints,
-          ...(typeof m.inputTokenLimit === "number" ? { inputTokenLimit: m.inputTokenLimit } : {}),
-          ...(typeof m.outputTokenLimit === "number"
-            ? { outputTokenLimit: m.outputTokenLimit }
-            : {}),
-          ...(typeof m.description === "string" ? { description: m.description } : {}),
-          ...(m.thinking === true ? { supportsThinking: true } : {}),
-        };
-      });
-    },
+    parseResponse: (data) => parseGeminiModelsList(data),
   },
   // gemini-cli handled via retrieveUserQuota (see GET handler)
   huggingface: {
@@ -2084,6 +2042,130 @@ export async function GET(
       });
     }
 
+    if (provider === "vertex" || provider === "vertex-partner") {
+      const cachedResponse = maybeReturnCachedDiscovery();
+      if (cachedResponse) return cachedResponse;
+
+      const autoFetchDisabledResponse = maybeReturnAutoFetchDisabled();
+      if (autoFetchDisabledResponse) return autoFetchDisabledResponse;
+
+      // Vertex AI lists models from the Generative Language `v1beta/models` endpoint, which both
+      // Express-mode API keys (via ?key=) and Service Account JSON (via a minted OAuth Bearer
+      // token) can reach. This surfaces the full live catalog — including image models
+      // (imagen-*, gemini-*-image) absent from the static registry list.
+      const credential = (apiKey || "").trim();
+      let queryKey: string | null = null;
+      let bearerToken: string | null = null;
+      try {
+        const { parseSAFromApiKey, getAccessToken } = await import(
+          "@omniroute/open-sse/executors/vertex.ts"
+        );
+        if (accessToken) {
+          bearerToken = accessToken;
+        } else if (credential) {
+          // A Service Account credential is a JSON object; a Vertex AI Express-mode API key is an
+          // opaque (non-JSON) string. Detect locally so this branch has no dependency on optional
+          // executor helpers.
+          let isServiceAccountJson = false;
+          try {
+            const parsed = JSON.parse(credential);
+            isServiceAccountJson =
+              !!parsed && typeof parsed === "object" && !Array.isArray(parsed);
+          } catch {
+            isServiceAccountJson = false;
+          }
+
+          if (isServiceAccountJson) {
+            bearerToken = await getAccessToken(parseSAFromApiKey(credential));
+          } else {
+            queryKey = credential;
+          }
+        }
+      } catch (error) {
+        // Couldn't resolve a usable credential (e.g. malformed Service Account JSON).
+        const fallback = buildDiscoveryErrorFallbackResponse(error, {
+          cacheWarning: "Vertex credential unavailable — using cached catalog",
+          localWarning: "Vertex credential unavailable — using local catalog",
+        });
+        if (fallback) return fallback;
+      }
+
+      if (!queryKey && !bearerToken) {
+        const fallback = buildDiscoveryFallbackResponse({
+          cacheWarning: "No usable Vertex credential — using cached catalog",
+          localWarning: "No usable Vertex credential — using local catalog",
+        });
+        if (fallback) return fallback;
+        return NextResponse.json(
+          { error: "No usable Vertex AI credential configured for model discovery." },
+          { status: 400 }
+        );
+      }
+
+      const baseUrl = "https://generativelanguage.googleapis.com/v1beta/models?pageSize=1000";
+      const headers: Record<string, string> = { "Content-Type": "application/json" };
+      if (bearerToken) headers["Authorization"] = `Bearer ${bearerToken}`;
+
+      const allModels: GeminiDiscoveryModel[] = [];
+      let pageUrl = queryKey ? `${baseUrl}&key=${encodeURIComponent(queryKey)}` : baseUrl;
+      let pageCount = 0;
+      const MAX_PAGES = 20;
+      const seenTokens = new Set<string>();
+
+      try {
+        while (pageUrl && pageCount < MAX_PAGES) {
+          pageCount++;
+          const response = await safeOutboundFetch(pageUrl, {
+            ...SAFE_OUTBOUND_FETCH_PRESETS.modelsPagination,
+            guard: getProviderOutboundGuard(),
+            proxyConfig: proxy,
+            method: "GET",
+            headers,
+          });
+
+          if (!response.ok) {
+            // Avoid logging the raw upstream body (may contain sensitive data); status is enough.
+            console.log("[models] Vertex model discovery failed", {
+              provider,
+              status: response.status,
+            });
+            const fallback = buildDiscoveryFallbackResponse();
+            if (fallback) return fallback;
+            return NextResponse.json(
+              { error: `Failed to fetch Vertex models: ${response.status}` },
+              { status: response.status }
+            );
+          }
+
+          const data = await response.json();
+          allModels.push(...parseGeminiModelsList(data));
+
+          const nextPageToken = data.nextPageToken;
+          if (!nextPageToken || seenTokens.has(nextPageToken)) break;
+          seenTokens.add(nextPageToken);
+          pageUrl = `${baseUrl}&pageToken=${encodeURIComponent(nextPageToken)}`;
+          if (queryKey) pageUrl += `&key=${encodeURIComponent(queryKey)}`;
+        }
+      } catch (error) {
+        const fallback = buildDiscoveryErrorFallbackResponse(error);
+        if (fallback) return fallback;
+        throw error;
+      }
+
+      if (allModels.length > 0) {
+        return buildApiDiscoveryResponse(allModels);
+      }
+
+      const fallback = buildDiscoveryFallbackResponse();
+      if (fallback) return fallback;
+      return buildResponse({
+        provider,
+        connectionId,
+        models: [],
+        source: "api",
+      });
+    }
+
     if (isAnthropicCompatibleProvider(provider)) {
       const cachedResponse = maybeReturnCachedDiscovery();
       if (cachedResponse) return cachedResponse;

src/lib/providerModels/geminiModelsParser.ts

@@ -0,0 +1,68 @@
+/**
+ * Parses the Google Generative Language `v1beta/models` listing into discovery models.
+ *
+ * Each model's `supportedGenerationMethods` is mapped to OmniRoute endpoints:
+ *   - generateContent / generateAnswer → "chat"
+ *   - predict / predictLongRunning      → "images"   (Imagen models)
+ *   - embedContent                      → "embeddings"
+ *   - bidiGenerateContent               → "audio"
+ *
+ * This is shared by the `gemini` discovery config and the `vertex` discovery branch: Vertex AI
+ * Express keys (and Service Account JSON via a minted OAuth token) list models from the same
+ * endpoint, so image models (imagen-*, gemini-*-image) surface dynamically instead of being
+ * limited to the small static registry list.
+ */
+const METHOD_TO_ENDPOINT: Record<string, string> = {
+  generateContent: "chat",
+  embedContent: "embeddings",
+  predict: "images",
+  predictLongRunning: "images",
+  bidiGenerateContent: "audio",
+  generateAnswer: "chat",
+};
+
+const IGNORED_METHODS = new Set([
+  "countTokens",
+  "countTextTokens",
+  "createCachedContent",
+  "batchGenerateContent",
+  "asyncBatchEmbedContent",
+]);
+
+export interface GeminiDiscoveryModel {
+  id: string;
+  name: string;
+  supportedEndpoints: string[];
+  inputTokenLimit?: number;
+  outputTokenLimit?: number;
+  description?: string;
+  supportsThinking?: boolean;
+  [key: string]: unknown;
+}
+
+export function parseGeminiModelsList(data: any): GeminiDiscoveryModel[] {
+  return (data?.models || []).map((m: Record<string, unknown>) => {
+    const methods: string[] = Array.isArray(m.supportedGenerationMethods)
+      ? (m.supportedGenerationMethods as string[])
+      : [];
+    const endpoints = [
+      ...new Set(
+        methods
+          .filter((method) => !IGNORED_METHODS.has(method))
+          .map((method) => METHOD_TO_ENDPOINT[method] || "chat")
+      ),
+    ];
+    if (endpoints.length === 0) endpoints.push("chat");
+
+    return {
+      ...m,
+      id: ((m.name as string) || (m.id as string) || "").replace(/^models\//, ""),
+      name: (m.displayName as string) || ((m.name as string) || "").replace(/^models\//, ""),
+      supportedEndpoints: endpoints,
+      ...(typeof m.inputTokenLimit === "number" ? { inputTokenLimit: m.inputTokenLimit } : {}),
+      ...(typeof m.outputTokenLimit === "number" ? { outputTokenLimit: m.outputTokenLimit } : {}),
+      ...(typeof m.description === "string" ? { description: m.description } : {}),
+      ...(m.thinking === true ? { supportsThinking: true } : {}),
+    } as GeminiDiscoveryModel;
+  });
+}

tests/unit/gemini-models-parser.test.ts

@@ -0,0 +1,86 @@
+import test from "node:test";
+import assert from "node:assert/strict";
+
+import { parseGeminiModelsList } from "../../src/lib/providerModels/geminiModelsParser";
+
+// A representative slice of the live generativelanguage v1beta/models response — including the
+// image models (gemini-*-image via generateContent, imagen-* via predict) that the Vertex catalog
+// must surface dynamically.
+const SAMPLE = {
+  models: [
+    {
+      name: "models/gemini-2.5-flash",
+      displayName: "Gemini 2.5 Flash",
+      inputTokenLimit: 1048576,
+      outputTokenLimit: 65536,
+      supportedGenerationMethods: ["generateContent", "countTokens", "batchGenerateContent"],
+      thinking: true,
+    },
+    {
+      name: "models/gemini-3-pro-image-preview",
+      displayName: "Gemini 3 Pro Image Preview",
+      supportedGenerationMethods: ["generateContent", "countTokens"],
+    },
+    {
+      name: "models/imagen-4.0-generate-001",
+      displayName: "Imagen 4.0",
+      supportedGenerationMethods: ["predict"],
+    },
+    {
+      name: "models/text-embedding-004",
+      displayName: "Text Embedding 004",
+      supportedGenerationMethods: ["embedContent"],
+    },
+    {
+      name: "models/gemini-live-2.5-flash",
+      displayName: "Gemini Live",
+      supportedGenerationMethods: ["bidiGenerateContent"],
+    },
+  ],
+};
+
+test("parseGeminiModelsList strips the models/ prefix and maps display name", () => {
+  const models = parseGeminiModelsList(SAMPLE);
+  const flash = models.find((m) => m.id === "gemini-2.5-flash");
+  assert.ok(flash, "gemini-2.5-flash should be present");
+  assert.equal(flash!.name, "Gemini 2.5 Flash");
+  assert.equal(flash!.inputTokenLimit, 1048576);
+  assert.equal(flash!.outputTokenLimit, 65536);
+  assert.equal(flash!.supportsThinking, true);
+  assert.deepEqual(flash!.supportedEndpoints, ["chat"]);
+});
+
+test("parseGeminiModelsList maps generateContent image models to the chat endpoint", () => {
+  const models = parseGeminiModelsList(SAMPLE);
+  const proImage = models.find((m) => m.id === "gemini-3-pro-image-preview");
+  assert.ok(proImage, "gemini-3-pro-image-preview should be present");
+  // gemini-*-image models generate images via generateContent (chat-with-modalities path).
+  assert.deepEqual(proImage!.supportedEndpoints, ["chat"]);
+});
+
+test("parseGeminiModelsList maps Imagen predict models to the images endpoint", () => {
+  const models = parseGeminiModelsList(SAMPLE);
+  const imagen = models.find((m) => m.id === "imagen-4.0-generate-001");
+  assert.ok(imagen, "imagen-4.0-generate-001 should be present");
+  assert.deepEqual(imagen!.supportedEndpoints, ["images"]);
+});
+
+test("parseGeminiModelsList maps embedContent and bidiGenerateContent", () => {
+  const models = parseGeminiModelsList(SAMPLE);
+  assert.deepEqual(
+    models.find((m) => m.id === "text-embedding-004")!.supportedEndpoints,
+    ["embeddings"]
+  );
+  assert.deepEqual(
+    models.find((m) => m.id === "gemini-live-2.5-flash")!.supportedEndpoints,
+    ["audio"]
+  );
+});
+
+test("parseGeminiModelsList defaults to chat and tolerates empty/missing input", () => {
+  assert.deepEqual(parseGeminiModelsList({}), []);
+  assert.deepEqual(parseGeminiModelsList(null), []);
+  const [m] = parseGeminiModelsList({ models: [{ name: "models/mystery" }] });
+  assert.equal(m.id, "mystery");
+  assert.deepEqual(m.supportedEndpoints, ["chat"]);
+});