From b999f1d204cc9821f99077852af9a4a668f99e06 Mon Sep 17 00:00:00 2001 From: iF2007 Date: Wed, 19 Aug 2026 08:52:14 +0800 Subject: [PATCH] fix: preserve discovered Google effort wire mappings --- src/providers/antigravity-models.ts | 75 +++++++++++++++++++++++++-- tests/google-antigravity-wire.test.ts | 37 +++++++++++-- 2 files changed, 102 insertions(+), 10 deletions(-) diff --git a/src/providers/antigravity-models.ts b/src/providers/antigravity-models.ts index bf155b543c..9b9045cb9b 100644 --- a/src/providers/antigravity-models.ts +++ b/src/providers/antigravity-models.ts @@ -72,6 +72,12 @@ const ANTIGRAVITY_WIRE_IDS_BY_PICKER_MODEL: Record = Object.en }, {}); const ANTIGRAVITY_DISCOVERY_EFFORTS = ["low", "medium", "high"] as const; +type AntigravityDiscoveryEffort = typeof ANTIGRAVITY_DISCOVERY_EFFORTS[number]; +type AntigravityEffortWireModelIds = Partial>; + +function isAntigravityDiscoveryEffort(value: string): value is AntigravityDiscoveryEffort { + return (ANTIGRAVITY_DISCOVERY_EFFORTS as readonly string[]).includes(value); +} function pickerModelIdForDiscoveredWireId( wireId: string, @@ -151,6 +157,25 @@ const ANTIGRAVITY_EFFORT_WIRE_MAP: Record> = { }, }; +function completeDiscoveredEffortWireModelIds( + pickerId: string, + available: ReadonlyMap>, +): AntigravityEffortWireModelIds | undefined { + const explicitEffortMap = ANTIGRAVITY_EFFORT_WIRE_MAP[pickerId]; + if (explicitEffortMap && Object.values(explicitEffortMap).every(wireId => available.has(wireId))) { + return { ...explicitEffortMap }; + } + + if (!isKnownAntigravityPickerModelId(pickerId)) return undefined; + const suffixEffortMap: AntigravityEffortWireModelIds = {}; + for (const effort of ANTIGRAVITY_DISCOVERY_EFFORTS) { + const wireId = `${pickerId}-${effort}`; + if (!available.has(wireId)) return undefined; + suffixEffortMap[effort] = wireId; + } + return suffixEffortMap; +} + // ── Default effort per Gemini base model ── const ANTIGRAVITY_DEFAULT_EFFORT: Record = { "gemini-3.1-pro": "high", @@ -270,6 +295,8 @@ export interface AntigravityAvailableModel { id: string; /** CCA model id used by the agent envelope when `id` comes from display metadata. */ wireModelId: string; + /** Complete effort-to-wire mapping retained for collapsed discovered tier sets. */ + effortWireModelIds?: AntigravityEffortWireModelIds; contextWindow?: number; inputModalities?: string[]; } @@ -286,6 +313,7 @@ function antigravityPositiveInteger(value: unknown): number | undefined { interface DiscoveredWireModelMapping { readonly models: ReadonlyMap; + readonly effortModels: ReadonlyMap; readonly generation?: { provider: string; cacheGeneration: string }; } @@ -327,17 +355,21 @@ export function registerAntigravityDiscoveredWireModels( const key = antigravityBaseUrlKey(baseUrl); if (!key) return; const wireModels = new Map(); - for (const model of models) wireModels.set(model.id, model.wireModelId); + const effortModels = new Map(); + for (const model of models) { + wireModels.set(model.id, model.wireModelId); + if (model.effortWireModelIds) effortModels.set(model.id, { ...model.effortWireModelIds }); + } discoveredWireModelsByBaseUrl.set(key, { models: wireModels, + effortModels, ...(generation ? { generation } : {}), }); } -function discoveredAntigravityWireModelId( - modelId: string, +function discoveredAntigravityMapping( baseUrl: string | undefined, -): string | undefined { +): DiscoveredWireModelMapping | undefined { const key = antigravityBaseUrlKey(baseUrl); if (!key) return undefined; const mapping = discoveredWireModelsByBaseUrl.get(key); @@ -347,7 +379,35 @@ function discoveredAntigravityWireModelId( discoveredWireModelsByBaseUrl.delete(key); return undefined; } - return mapping.models.get(modelId); + return mapping; +} + +function discoveredAntigravityWireModelId( + modelId: string, + baseUrl: string | undefined, +): string | undefined { + return discoveredAntigravityMapping(baseUrl)?.models.get(modelId); +} + +function discoveredAntigravityEffortWireModelId( + modelId: string, + effort: string | undefined, + baseUrl: string | undefined, +): string | undefined { + const effortMap = discoveredAntigravityMapping(baseUrl)?.effortModels.get(modelId); + if (!effortMap) return undefined; + + const requestedEffort = effort ? resolveAntigravityThinkingLevel(effort) : undefined; + if (requestedEffort && isAntigravityDiscoveryEffort(requestedEffort) && effortMap[requestedEffort]) { + return effortMap[requestedEffort]; + } + + const defaultEffort = ANTIGRAVITY_DEFAULT_EFFORT[modelId] + ?? ANTIGRAVITY_THINKING_LEVEL_MODELS[modelId]; + if (defaultEffort && isAntigravityDiscoveryEffort(defaultEffort) && effortMap[defaultEffort]) { + return effortMap[defaultEffort]; + } + return Object.values(effortMap)[0]; } /** @@ -465,9 +525,11 @@ export function parseAntigravityAvailableModels( const id = pickerModelIdForDiscoveredWireId(wireId, info, available); if (seen.has(id)) continue; seen.add(id); + const effortWireModelIds = completeDiscoveredEffortWireModelIds(id, available); out.push({ id, wireModelId: wireId, + ...(effortWireModelIds ? { effortWireModelIds } : {}), ...(antigravityPositiveInteger(info.maxTokens) ? { contextWindow: antigravityPositiveInteger(info.maxTokens) } : {}), // Tri-state, deliberately not a ternary: `true` asserts image support, // `false` asserts against it, and ABSENT is unknown. Collapsing absent into @@ -522,6 +584,9 @@ export function resolveAntigravityEffortWireModel( effort?: string, baseUrl?: string, ): { wireModelId: string; thinkingLevel?: string } { + const discoveredEffortWireModelId = discoveredAntigravityEffortWireModelId(modelId, effort, baseUrl); + if (discoveredEffortWireModelId) return { wireModelId: discoveredEffortWireModelId }; + // A collapsed picker row reports ONE representative wire id (whichever tier CCA // listed first), so live discovery cannot describe a ladder — it can only name a // single rung. Letting it answer for a base model we already have a ladder for diff --git a/tests/google-antigravity-wire.test.ts b/tests/google-antigravity-wire.test.ts index 986b3eef45..8aefb92d46 100644 --- a/tests/google-antigravity-wire.test.ts +++ b/tests/google-antigravity-wire.test.ts @@ -133,11 +133,28 @@ describe("antigravity CCA envelope", () => { agentModelSorts: [{ groups: [{ modelIds }] }], }); - expect(parseAntigravityAvailableModels(payload([ + const rows = 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(rows.map(model => model.id)).toEqual(["gemini-3.7-flash"]); + expect(rows[0]?.wireModelId).toBe("gemini-3.7-flash-low"); + expect(rows[0]?.effortWireModelIds).toEqual({ + low: "gemini-3.7-flash-low", + medium: "gemini-3.7-flash-medium", + high: "gemini-3.7-flash-high", + }); + const baseUrl = "https://cca-tiered-set.example"; + registerAntigravityDiscoveredWireModels(baseUrl, rows); + for (const [effort, wireModelId] of [ + ["low", "gemini-3.7-flash-low"], + ["medium", "gemini-3.7-flash-medium"], + ["high", "gemini-3.7-flash-high"], + ] as const) { + expect(resolveAntigravityEffortWireModel("gemini-3.7-flash", effort, baseUrl)) + .toEqual({ wireModelId }); + } expect(parseAntigravityAvailableModels(payload([ "future-flash-low", "future-flash-medium", @@ -232,14 +249,24 @@ describe("antigravity CCA envelope", () => { ]); // The display label still resolves an id Google renamed on the wire. expect(rows.find(model => model.id === "gemini-nebula")?.wireModelId).toBe("internal-codename-x7"); + expect(rows.find(model => model.id === "gemini-3.1-pro")?.effortWireModelIds).toEqual({ + low: "gemini-3.1-pro-low", + high: "gemini-pro-agent", + }); const baseUrl = "https://cca.example"; registerAntigravityDiscoveredWireModels(baseUrl, rows); - // A discovered representative wire id must NOT override a real effort ladder. + // A complete discovery preserves each discovered suffix for the requested effort. expect(resolveAntigravityEffortWireModel("gemini-3.1-pro", "low", baseUrl)) - .toEqual({ wireModelId: "gemini-3.1-pro-low", thinkingLevel: "low" }); + .toEqual({ wireModelId: "gemini-3.1-pro-low" }); + expect(resolveAntigravityEffortWireModel("gemini-3.1-pro", "high", baseUrl)) + .toEqual({ wireModelId: "gemini-pro-agent" }); expect(resolveAntigravityEffortWireModel("gemini-3.7-flash", "low", baseUrl)) - .toEqual({ wireModelId: "gemini-3.7-flash-tiered", thinkingLevel: "low" }); + .toEqual({ wireModelId: "gemini-3.7-flash-low" }); + expect(resolveAntigravityEffortWireModel("gemini-3.7-flash", "medium", baseUrl)) + .toEqual({ wireModelId: "gemini-3.7-flash-medium" }); + expect(resolveAntigravityEffortWireModel("gemini-3.7-flash", "high", baseUrl)) + .toEqual({ wireModelId: "gemini-3.7-flash-high" }); expect(resolveAntigravityEffortWireModel("gemini-nebula", undefined, baseUrl)) .toEqual({ wireModelId: "internal-codename-x7" }); expect(resolveAntigravityEffortWireModel("claude-sonnet-4-6", "high", baseUrl))