diff --git a/src/codex/catalog/provider-fetch.ts b/src/codex/catalog/provider-fetch.ts index add955bb45..9b50636f8f 100644 --- a/src/codex/catalog/provider-fetch.ts +++ b/src/codex/catalog/provider-fetch.ts @@ -79,6 +79,7 @@ import { type ProviderModelsApiItem, type ResolvedProviderModelDiscovery, } from "../../providers/model-discovery"; +import { extractGoogleAiStudioModelItems } from "../../providers/google-ai-studio-model-discovery"; import { applyConfiguredHeadersLast, fetchOllamaShowEnrichment, ollamaShowEnrichable } from "../../providers/ollama-show"; import upstreamModelsSnapshot from "../data/upstream-models.json"; import { createAdmissionGate, ResourceAdmissionError, type AdmissionMetrics } from "../../lib/admission"; @@ -1871,7 +1872,14 @@ async function fetchProviderModelsWithAuth( markProviderDiscoveryOk(name, live.length); return observed(withConfiguredRetention(forCache, { warnDrops: true }), "authoritative"); } - const extracted = extractProviderModelItems(bounded.value, discovery); + const googleAiStudio = effectiveGoogleMode(name, prov) === "ai-studio" + ? extractGoogleAiStudioModelItems(bounded.value, discovery.maxModels) + : undefined; + // Native /v1beta/models wins; a google row served by an OpenAI-compatible + // gateway keeps the generic data[] / top-level-array contract. + const extracted = googleAiStudio?.ok + ? googleAiStudio + : extractProviderModelItems(bounded.value, discovery); if (!extracted.ok) { const { models, fallback, shouldLog } = failedDiscoveryFallback({ reason: "invalid_response" }); const diagnostic: Record = { diff --git a/src/providers/google-ai-studio-model-discovery.ts b/src/providers/google-ai-studio-model-discovery.ts new file mode 100644 index 0000000000..5d9a75c6db --- /dev/null +++ b/src/providers/google-ai-studio-model-discovery.ts @@ -0,0 +1,74 @@ +import { + extractModelEnvelopeRows, + isValidModelDiscoveryModelId, + type ProviderModelItemsResult, + type ProviderModelsApiItem, +} from "./model-discovery"; + +const GOOGLE_MODEL_PREFIX = "models/"; +const MAX_GENERATION_METHODS = 32; +const MAX_GENERATION_METHOD_LENGTH = 64; + +/** Returns the value if it is a positive safe integer; otherwise undefined. */ +function positiveSafeInteger(value: unknown): number | undefined { + return typeof value === "number" && Number.isSafeInteger(value) && value > 0 + ? value + : undefined; +} + +/** + * Extracts and normalizes supported model items from a Google AI Studio + * /v1beta/models response payload. + * + * Validates the native models[] envelope, strips the 'models/' prefix, filters + * to rows supporting 'generateContent', maps input/output token limits, and + * resiliently skips toxic or malformed individual rows. + */ +export function extractGoogleAiStudioModelItems( + value: unknown, + maxModels: number, +): ProviderModelItemsResult { + const envelope = extractModelEnvelopeRows(value, maxModels, ["models"]); + if (!envelope.ok) return envelope; + + const items: ProviderModelsApiItem[] = []; + const seen = new Set(); + for (const raw of envelope.rows) { + if (raw === null || typeof raw !== "object" || Array.isArray(raw)) { + continue; + } + const name = Reflect.get(raw, "name"); + const generationMethods = Reflect.get(raw, "supportedGenerationMethods"); + if (!isValidModelDiscoveryModelId(name)) { + continue; + } + if (generationMethods === undefined) continue; + if ( + !Array.isArray(generationMethods) + || generationMethods.length > MAX_GENERATION_METHODS + || generationMethods.some(method => typeof method !== "string" || method.length > MAX_GENERATION_METHOD_LENGTH) + ) { + continue; + } + if (!generationMethods.includes("generateContent")) continue; + + const id = name.startsWith(GOOGLE_MODEL_PREFIX) + ? name.slice(GOOGLE_MODEL_PREFIX.length) + : name; + if (!isValidModelDiscoveryModelId(id) || seen.has(id)) continue; + seen.add(id); + + const inputTokenLimit = positiveSafeInteger(Reflect.get(raw, "inputTokenLimit")); + const outputTokenLimit = positiveSafeInteger(Reflect.get(raw, "outputTokenLimit")); + items.push({ + id, + owned_by: "google", + ...(inputTokenLimit !== undefined + ? { context_length: inputTokenLimit, max_input_tokens: inputTokenLimit } + : {}), + ...(outputTokenLimit !== undefined ? { max_output_tokens: outputTokenLimit } : {}), + }); + } + return { ok: true, items, rawCount: envelope.rows.length }; +} + diff --git a/src/server/management/provider-routes.ts b/src/server/management/provider-routes.ts index 0fce336631..1439d7899c 100644 --- a/src/server/management/provider-routes.ts +++ b/src/server/management/provider-routes.ts @@ -54,6 +54,7 @@ import { readBoundedDiscoveryJson, resolveProviderModelDiscovery, } from "../../providers/model-discovery"; +import { extractGoogleAiStudioModelItems } from "../../providers/google-ai-studio-model-discovery"; import { routedSlug, slugEquals } from "../../providers/slug-codec"; import { clearAccountQuotaCache, clearProviderQuotaCache, fetchProviderQuotaReports } from "../../providers/quota"; import { clearKeyCooldowns } from "../../providers/key-failover"; @@ -1428,13 +1429,19 @@ export async function handleProviderRoutes(ctx: ManagementContext): Promise : undefined; + const isAiStudio = effectiveGoogleMode(name, prov) === "ai-studio"; + const googleAiStudio = !ccaModels && isAiStudio + ? extractGoogleAiStudioModelItems(bounded.value, discovery.maxModels) + : undefined; const extracted = ccaModels ? undefined + : googleAiStudio?.ok + ? googleAiStudio : Array.isArray(bounded.value) || Array.isArray(record?.data) ? extractProviderModelItems(bounded.value, discovery) : extractModelEnvelopeRows(bounded.value, discovery.maxModels, ["models"]); diff --git a/tests/adapters/google/google-models-listing.test.ts b/tests/adapters/google/google-models-listing.test.ts index 3440ef30a7..755ff6dcc2 100644 --- a/tests/adapters/google/google-models-listing.test.ts +++ b/tests/adapters/google/google-models-listing.test.ts @@ -330,7 +330,7 @@ describe("buildModelsRequest anthropic routing", () => { }); describe("google models listing via catalog", () => { - test("treats a { models } 2xx shape as malformed and degrades to the static seed", async () => { + test("publishes generateContent models from the native models envelope", async () => { clearModelCache("google"); const warning = spyOn(console, "warn").mockImplementation(() => {}); const seen: { url: string; headers: Record }[] = []; @@ -338,8 +338,9 @@ describe("google models listing via catalog", () => { seen.push({ url: String(input), headers: (init?.headers ?? {}) as Record }); return new Response(JSON.stringify({ models: [ - { name: "models/gemini-3-pro", inputTokenLimit: 1048576, supportedGenerationMethods: ["generateContent", "countTokens"] }, + { name: "models/gemini-3-pro", inputTokenLimit: 1048576, outputTokenLimit: 65536, supportedGenerationMethods: ["generateContent", "countTokens"] }, { name: "models/text-embedding-004", supportedGenerationMethods: ["embedContent"] }, + { name: "models/gemini-missing-methods" }, { name: "models/gemini-3-flash", inputTokenLimit: 1048576, supportedGenerationMethods: ["generateContent"] }, ], }), { status: 200, headers: { "content-type": "application/json" } }); @@ -356,16 +357,82 @@ describe("google models listing via catalog", () => { expect(seen).toHaveLength(1); expect(seen[0].url).toBe("https://generativelanguage.googleapis.com/v1beta/models?pageSize=1000"); expect(seen[0].headers["x-goog-api-key"]).toBe("gk-123"); - const ids = models.filter(m => m.provider === "google").map(m => m.id); - expect(ids).toEqual(["gemini-3.1-pro-preview", "gemini-3.5-flash", "gemini-3.5-flash-lite", "gemini-3.6-flash", "gemini-3.7-flash", "gemini-3.8-flash"]); - expect(ids).not.toContain("gemini-3-pro"); - expect(ids).not.toContain("gemini-3-flash"); - expect(getStaleCached("google")).toBeNull(); - expect(warning.mock.calls.flat().join(" ")).toContain("google"); + const live = models.filter(m => m.provider === "google"); + expect(live.map(m => m.id).sort()).toEqual(["gemini-3-flash", "gemini-3-pro"]); + expect(live.find(m => m.id === "gemini-3-pro")).toMatchObject({ + contextWindow: 1_048_576, + maxInputTokens: 1_048_576, + maxOutputTokens: 65_536, + }); } finally { warning.mockRestore(); } }); + + test("skips toxic or malformed rows and preserves valid models from the native models envelope", async () => { + clearModelCache("google"); + const warning = spyOn(console, "warn").mockImplementation(() => {}); + globalThis.fetch = (async () => { + return new Response(JSON.stringify({ + models: [ + null, + "not-an-object", + { name: "models/bad\0name", supportedGenerationMethods: ["generateContent"] }, + { name: "models/ padded ", supportedGenerationMethods: ["generateContent"] }, + { name: "models/gemini-valid", inputTokenLimit: 524288, outputTokenLimit: 8192, supportedGenerationMethods: ["generateContent"] }, + // Same normalized id as the row above: must be deduped, not published twice. + { name: "models/gemini-valid", inputTokenLimit: 1024, supportedGenerationMethods: ["generateContent"] }, + // No `models/` prefix: the name is used verbatim. + { name: "gemini-unprefixed", supportedGenerationMethods: ["generateContent"] }, + { name: "models/invalid-methods", supportedGenerationMethods: "not-an-array" }, + { name: "models/embed-only", supportedGenerationMethods: ["embedContent"] }, + ], + }), { status: 200, headers: { "content-type": "application/json" } }); + }) as typeof fetch; + + try { + const models = await gatherRoutedModels(configWith("google", { + adapter: "google", + authMode: "key", + apiKey: "gk-123", + baseUrl: "https://generativelanguage.googleapis.com", + })); + + const live = models.filter(m => m.provider === "google"); + expect(live.map(m => m.id).sort()).toEqual(["gemini-unprefixed", "gemini-valid"]); + expect(live.filter(m => m.id === "gemini-valid")).toHaveLength(1); + expect(live.find(m => m.id === "gemini-valid")).toMatchObject({ + contextWindow: 524_288, + maxInputTokens: 524_288, + maxOutputTokens: 8_192, + }); + expect(getStaleCached("google")).not.toBeNull(); + } finally { + warning.mockRestore(); + } + }); + + test("falls back to generic parser when a custom google-adapter provider returns data[] envelope", async () => { + clearModelCache("custom-google"); + globalThis.fetch = (async () => { + return new Response(JSON.stringify({ + data: [ + { id: "custom-gemini", owned_by: "custom", context_length: 128000 }, + ], + }), { status: 200, headers: { "content-type": "application/json" } }); + }) as typeof fetch; + + const models = await gatherRoutedModels(configWith("custom-google", { + adapter: "google", + authMode: "key", + apiKey: "gk-custom", + baseUrl: "https://custom-gateway.example/v1", + })); + + const live = models.filter(m => m.provider === "custom-google"); + expect(live.map(m => m.id)).toEqual(["custom-gemini"]); + expect(live[0]?.contextWindow).toBe(128_000); + }); }); describe("models fetch failure cooldown", () => { diff --git a/tests/providers/provider-connection-test.test.ts b/tests/providers/provider-connection-test.test.ts index bbadd8c442..ac28538095 100644 --- a/tests/providers/provider-connection-test.test.ts +++ b/tests/providers/provider-connection-test.test.ts @@ -384,11 +384,11 @@ describe("POST /api/providers/test (WP040 connectivity probe)", () => { expect(body).toMatchObject({ ok: true, models: 390 }); }); - test("Google's models-array response shape is accepted (x-goog-api-key path)", async () => { + test("Google's models-array response counts only generateContent models", async () => { let requestedUrl = ""; globalThis.fetch = (async (input: RequestInfo | URL) => { requestedUrl = String(input); - return new Response(JSON.stringify({ models: [{ name: "models/gemini-3-pro" }, { name: "models/gemini-3-flash" }, { name: "models/gemini-3-lite" }] }), { + return new Response(JSON.stringify({ models: [{ name: "models/gemini-3-pro", supportedGenerationMethods: ["generateContent"] }, { name: "models/gemini-3-flash", supportedGenerationMethods: ["generateContent", "countTokens"] }, { name: "models/text-embedding-004", supportedGenerationMethods: ["embedContent"] }, { name: "models/gemini-missing-methods" }] }), { status: 200, headers: { "content-type": "application/json" }, }); @@ -399,7 +399,41 @@ describe("POST /api/providers/test (WP040 connectivity probe)", () => { const { body } = await probe(config, "google"); expect(requestedUrl).toContain("/v1beta/models"); expect(body.ok).toBe(true); - expect(body.models).toBe(3); + expect(body.models).toBe(2); + }); + + test("Google AI Studio probe skips malformed or toxic rows while counting valid models", async () => { + globalThis.fetch = (async () => new Response(JSON.stringify({ + models: [ + null, + "invalid-row", + { name: "models/badname", supportedGenerationMethods: ["generateContent"] }, + { name: "models/ padded ", supportedGenerationMethods: ["generateContent"] }, + { name: "models/gemini-valid", supportedGenerationMethods: ["generateContent"] }, + ], + }), { + status: 200, + headers: { "content-type": "application/json" }, + })) as typeof fetch; + const config = baseConfig({ + google: { adapter: "google", baseUrl: "https://generativelanguage.googleapis.com", apiKey: "g-key" }, + }); + const { body } = await probe(config, "google"); + expect(body.ok).toBe(true); + expect(body.models).toBe(1); + }); + + test("non-ai-studio providers preserve generic models[] connection-test fallback", async () => { + globalThis.fetch = (async () => new Response(JSON.stringify({ models: [{ id: "m-1" }, { id: "m-2" }] }), { + status: 200, + headers: { "content-type": "application/json" }, + })) as typeof fetch; + const config = baseConfig({ + generic: { adapter: "openai-chat", baseUrl: "https://api.example.test/v1", apiKey: "sk-x" }, + }); + const { body } = await probe(config, "generic"); + expect(body.ok).toBe(true); + expect(body.models).toBe(2); }); test("Together-style top-level /models array is accepted (#617)", async () => {