diff --git a/src/providers/antigravity-models.ts b/src/providers/antigravity-models.ts index cc862c6503..05bdd3b19d 100644 --- a/src/providers/antigravity-models.ts +++ b/src/providers/antigravity-models.ts @@ -6,8 +6,8 @@ import { isValidModelDiscoveryModelId, MODEL_DISCOVERY_MAX_MODELS } from "./mode // CLI resolves labels against. The ids below separate CCA wire ids, collapsed picker entries, // and hidden compatibility aliases for saved selections. The CCA envelope's `model` field must // receive the wire id (for example "Gemini 3.1 Pro (High)" => gemini-pro-agent), while the -// picker exposes collapsed base models only when CCA returns every known tier; otherwise each -// returned wire id remains visible so an unavailable tier cannot be selected. +// picker exposes collapsed known base models only when CCA returns every known tier; unknown +// returned wire ids remain visible so they stay directly routable. // ── Wire IDs (what CCA :fetchAvailableModels returns) ── @@ -63,6 +63,41 @@ const ANTIGRAVITY_WIRE_IDS_BY_PICKER_MODEL: Record = Object.en return out; }, {}); +const ANTIGRAVITY_DISCOVERY_EFFORTS = ["low", "medium", "high"] as const; + +function pickerModelIdForDiscoveredWireId( + wireId: string, + available: ReadonlyMap>, +): string { + const explicitPickerId = Object.hasOwn(ANTIGRAVITY_PICKER_MODEL_BY_WIRE_ID, wireId) + ? ANTIGRAVITY_PICKER_MODEL_BY_WIRE_ID[wireId] + : undefined; + if (explicitPickerId) { + const requiredWireIds = Object.hasOwn(ANTIGRAVITY_WIRE_IDS_BY_PICKER_MODEL, explicitPickerId) + ? ANTIGRAVITY_WIRE_IDS_BY_PICKER_MODEL[explicitPickerId] ?? [] + : []; + if (requiredWireIds.every(id => available.has(id))) return explicitPickerId; + } + + // CCA uses a single `-tiered` row for models whose effort levels ride on the + // request's thinkingLevel field. Keep this generic so new tiered models do not + // require another provider-specific ID mapping. + if (wireId.endsWith("-tiered")) { + const baseId = wireId.slice(0, -"-tiered".length); + if (isKnownAntigravityPickerModelId(baseId)) return baseId; + } + + const effortMatch = /^(.*)-(low|medium|high)$/.exec(wireId); + if (effortMatch) { + const baseId = effortMatch[1]!; + if (isKnownAntigravityPickerModelId(baseId) + && ANTIGRAVITY_DISCOVERY_EFFORTS.every(effort => available.has(`${baseId}-${effort}`))) { + return baseId; + } + } + return wireId; +} + // ── Effort ladders per collapsed base model ── // Gemini models: effort → wire model suffix (official agy UI pattern). // Claude Opus: effort → thinkingConfig.thinkingLevel (CLIProxyAPI proven pattern). @@ -141,6 +176,10 @@ export const ANTIGRAVITY_MODELS = [ "gpt-oss-120b-medium", ]; +function isKnownAntigravityPickerModelId(value: string): boolean { + return isValidModelDiscoveryModelId(value) && ANTIGRAVITY_MODELS.includes(value); +} + // Context windows from the upstream `:fetchAvailableModels` maxTokens per model. const ANTIGRAVITY_WIRE_MODEL_CONTEXT_WINDOWS: Record = { "gemini-3.7-flash": 1_048_576, @@ -222,6 +261,7 @@ export function parseAntigravityAvailableModels( if (!Array.isArray(modelIds)) return null; for (const id of modelIds) { if (!isValidModelDiscoveryModelId(id) + || !Object.hasOwn(models, id) || !antigravityRecord(models[id]) || ids.length >= limit) return null; ids.push(id); @@ -235,6 +275,19 @@ export function parseAntigravityAvailableModels( if (ids.length >= limit) return null; ids.push("gemini-3.1-flash-image"); } + // Newer CCA responses identify tiered Flash models through this index instead of + // adding their synthetic wire ids to agentModelSorts. + const tieredModelIds = antigravityRecord(body.tieredModelIds); + const flashTieredIds = tieredModelIds?.flash; + if (Array.isArray(flashTieredIds)) { + for (const id of flashTieredIds) { + if (!isValidModelDiscoveryModelId(id) + || !Object.hasOwn(models, id) + || !antigravityRecord(models[id]) + || ids.length >= limit) return null; + ids.push(id); + } + } const available = new Map>(); for (const wireId of ids) { @@ -242,17 +295,17 @@ export function parseAntigravityAvailableModels( if (!info || available.has(wireId)) continue; // Legacy compatibility aliases are deliberately routed to newer wire ids for saved // selections. They are not safe as independently discovered picker rows. - if (ANTIGRAVITY_MODEL_ALIASES[wireId] && ANTIGRAVITY_MODEL_ALIASES[wireId] !== wireId) continue; + const alias = Object.hasOwn(ANTIGRAVITY_MODEL_ALIASES, wireId) + ? ANTIGRAVITY_MODEL_ALIASES[wireId] + : undefined; + if (alias && alias !== wireId) continue; available.set(wireId, info); } const out: AntigravityAvailableModel[] = []; const seen = new Set(); for (const [wireId, info] of available) { - const pickerId = ANTIGRAVITY_PICKER_MODEL_BY_WIRE_ID[wireId]; - const completePickerSet = pickerId !== undefined - && ANTIGRAVITY_WIRE_IDS_BY_PICKER_MODEL[pickerId]!.every(id => available.has(id)); - const id = completePickerSet ? pickerId! : wireId; + const id = pickerModelIdForDiscoveredWireId(wireId, available); if (seen.has(id)) continue; seen.add(id); out.push({ @@ -265,7 +318,9 @@ export function parseAntigravityAvailableModels( } export function resolveAntigravityWireModelId(modelId: string): string { - return ANTIGRAVITY_MODEL_ALIASES[modelId] ?? modelId; + return Object.hasOwn(ANTIGRAVITY_MODEL_ALIASES, modelId) + ? ANTIGRAVITY_MODEL_ALIASES[modelId] + : modelId; } /** @@ -279,7 +334,7 @@ export function isAntigravitySuffixModelId(modelId: string): boolean { /** The reasoning tier a retired Flash id used to encode, if it is one. */ export function retiredAntigravityFlashTier(modelId: string): string | undefined { - return RETIRED_FLASH_TIERS[modelId]; + return Object.hasOwn(RETIRED_FLASH_TIERS, modelId) ? RETIRED_FLASH_TIERS[modelId] : undefined; } /** @@ -299,7 +354,7 @@ export function resolveAntigravityEffortWireModel( // Rule 0: retired Flash id — Google has taken the wire id offline, so route to the // current generation and carry the tier the retired id encoded. This runs BEFORE the // suffix check because those ids are aliases, and rule 1 would drop the tier. - const retiredTier = RETIRED_FLASH_TIERS[modelId]; + const retiredTier = retiredAntigravityFlashTier(modelId); if (retiredTier) { return { wireModelId: GEMINI_FLASH_CURRENT, diff --git a/tests/google-antigravity-wire.test.ts b/tests/google-antigravity-wire.test.ts index e0b180ed87..83587109cd 100644 --- a/tests/google-antigravity-wire.test.ts +++ b/tests/google-antigravity-wire.test.ts @@ -1,7 +1,7 @@ import { describe, expect, test } from "bun:test"; import { createGoogleAdapter as createGoogleAdapterProduction } from "../src/adapters/google"; import { antigravitySessionId, isLikelyRealThoughtSignature } from "../src/adapters/google-antigravity-wire"; -import { ANTIGRAVITY_MODELS, ANTIGRAVITY_MODEL_EFFORTS, canonicalAntigravityUsageModel, parseAntigravityAvailableModels } from "../src/providers/antigravity-models"; +import { ANTIGRAVITY_MODELS, ANTIGRAVITY_MODEL_EFFORTS, canonicalAntigravityUsageModel, parseAntigravityAvailableModels, resolveAntigravityEffortWireModel, resolveAntigravityWireModelId } from "../src/providers/antigravity-models"; import { MODEL_DISCOVERY_MAX_MODEL_ID_LENGTH, MODEL_DISCOVERY_MAX_MODELS } from "../src/providers/model-discovery"; import type { AdapterEvent, OcxParsedRequest, OcxProviderConfig } from "../src/types"; import { withTestTranslatorBudget } from "./helpers/translator-budget"; @@ -129,6 +129,50 @@ describe("antigravity CCA envelope", () => { agentModelSorts: [{ groups: [{ modelIds }] }], }); + expect(parseAntigravityAvailableModels(payload([ + "gemini-3.7-flash-low", + "gemini-3.7-flash-medium", + "gemini-3.7-flash-high", + ]))?.map(model => model.id)).toEqual(["gemini-3.7-flash"]); + expect(parseAntigravityAvailableModels(payload([ + "future-flash-low", + "future-flash-medium", + "future-flash-high", + ]))?.map(model => model.id)).toEqual([ + "future-flash-low", + "future-flash-medium", + "future-flash-high", + ]); + expect(parseAntigravityAvailableModels(payload([ + "future-flash-low", + "future-flash-high", + ]))?.map(model => model.id)).toEqual([ + "future-flash-low", + "future-flash-high", + ]); + expect(parseAntigravityAvailableModels({ + models: { + "future-flash-tiered": { maxTokens: 1_048_576 }, + }, + agentModelSorts: [{ groups: [{ modelIds: [] }] }], + tieredModelIds: { flash: ["future-flash-tiered"] }, + })?.map(model => model.id)).toEqual(["future-flash-tiered"]); + expect(parseAntigravityAvailableModels({ + models: { + "gemini-3.7-flash-tiered": { maxTokens: 1_048_576 }, + }, + agentModelSorts: [{ groups: [{ modelIds: [] }] }], + tieredModelIds: { flash: ["gemini-3.7-flash-tiered"] }, + })?.map(model => model.id)).toEqual(["gemini-3.7-flash"]); + expect(parseAntigravityAvailableModels({ + models: { "-tiered": { maxTokens: 1_048_576 } }, + agentModelSorts: [{ groups: [{ modelIds: ["-tiered"] }] }], + })?.map(model => model.id)).toEqual(["-tiered"]); + expect(parseAntigravityAvailableModels(payload([ + "-low", + "-medium", + "-high", + ]))?.map(model => model.id)).toEqual(["-low", "-medium", "-high"]); expect(parseAntigravityAvailableModels(payload([ "gemini-3.1-pro-low", "gemini-pro-agent", @@ -140,6 +184,33 @@ describe("antigravity CCA envelope", () => { ]); }); + test("keeps unknown discovered tier IDs directly routable", async () => { + for (const modelId of ["future-flash-tiered", "future-flash-low"]) { + const req = await createGoogleAdapter(effortProvider).buildRequest(parsedWithEffort(modelId, "high")); + const env = JSON.parse(req.body); + expect(env.model).toBe(modelId); + expect(env.request.generationConfig?.thinkingConfig).toBeUndefined(); + } + }); + + test("ignores inherited CCA model and alias properties", () => { + const inheritedModels = Object.create(null) as Record; + Object.defineProperty(inheritedModels, "__proto__", { + value: { maxTokens: 1_048_576 }, + enumerable: true, + }); + const models = Object.create(inheritedModels); + + expect(parseAntigravityAvailableModels({ + models, + agentModelSorts: [{ groups: [{ modelIds: ["__proto__"] }] }], + })).toBeNull(); + expect(resolveAntigravityWireModelId("__proto__")).toBe("__proto__"); + expect(resolveAntigravityEffortWireModel("__proto__", "high")).toEqual({ + wireModelId: "__proto__", + }); + }); + test("rejects malformed and oversized CCA agent-model lists", () => { const payload = (modelIds: unknown[]) => ({ models: Object.fromEntries(modelIds.map(id => [String(id), { maxTokens: 1_048_576 }])), diff --git a/tests/google-models-listing.test.ts b/tests/google-models-listing.test.ts index eb92603dcb..0508058d27 100644 --- a/tests/google-models-listing.test.ts +++ b/tests/google-models-listing.test.ts @@ -87,15 +87,21 @@ describe("Antigravity live model discovery", () => { return Response.json({ models: { "gemini-3.1-pro-low": { maxTokens: 1_048_576, supportsImages: true, supportsThinking: true, thinkingBudget: 1000 }, - "gemini-3.7-flash": { maxTokens: 1_048_576, supportsImages: true, supportsThinking: true, thinkingBudget: 10000 }, + "gemini-3.7-flash-tiered": { maxTokens: 1_048_576, supportsImages: true, supportsThinking: true, thinkingBudget: 10000 }, + "future-flash-tiered": { maxTokens: 1_048_576, supportsImages: true, supportsThinking: true, thinkingBudget: 10000 }, + "future-flash-low": { maxTokens: 1_048_576, supportsImages: true, supportsThinking: true, thinkingBudget: 10000 }, + "future-flash-medium": { maxTokens: 1_048_576, supportsImages: true, supportsThinking: true, thinkingBudget: 10000 }, + "future-flash-high": { maxTokens: 1_048_576, supportsImages: true, supportsThinking: true, thinkingBudget: 10000 }, "future-agent-model": { maxTokens: 333_333, supportsImages: false, supportsThinking: true, thinkingBudget: 7777 }, "gemini-3.1-flash-image": { maxTokens: 555_555, supportsImages: true }, "non-agent-command-model": { maxTokens: 222_222 }, "tab-only-model": { maxTokens: 32_768 }, }, agentModelSorts: [{ groups: [{ modelIds: [ - "future-agent-model", "gemini-3.1-pro-low", "gemini-3.7-flash", + "future-agent-model", "gemini-3.1-pro-low", + "future-flash-low", "future-flash-medium", "future-flash-high", ] }] }], + tieredModelIds: { flash: ["gemini-3.7-flash-tiered", "future-flash-tiered"] }, imageGenerationModelIds: ["gemini-3.1-flash-image"], tabModelIds: ["tab-only-model"], commandModelIds: ["non-agent-command-model"], @@ -120,6 +126,10 @@ describe("Antigravity live model discovery", () => { expect(JSON.parse(String(seen[0]?.init?.body))).toEqual({ project: "configured-project" }); expect(live.map(model => model.id).sort()).toEqual([ "future-agent-model", + "future-flash-high", + "future-flash-low", + "future-flash-medium", + "future-flash-tiered", "gemini-3.1-flash-image", "gemini-3.1-pro-low", "gemini-3.7-flash",