From 6bb37b3f5ecb6bdcf585847b227a617bde24477b Mon Sep 17 00:00:00 2001 From: Andrew Wang Date: Sun, 25 Oct 2026 01:38:44 -0400 Subject: [PATCH] fix(antigravity): union multi-account model discovery, skip disabled wire ids, and rank initial failover by quota --- src/codex/catalog/provider-models.ts | 57 ++++++++++++++++++- src/codex/routing-drift.ts | 2 + src/oauth/generic-account-failover.ts | 29 +++++----- src/providers/antigravity-models.ts | 2 +- .../google/antigravity-account-roster.test.ts | 40 +++++++++++++ 5 files changed, 113 insertions(+), 17 deletions(-) diff --git a/src/codex/catalog/provider-models.ts b/src/codex/catalog/provider-models.ts index 58c7a99e88f..2cf95bce213 100644 --- a/src/codex/catalog/provider-models.ts +++ b/src/codex/catalog/provider-models.ts @@ -33,6 +33,7 @@ import { import { buildModelsRequest, getModelsOAuthAccessSnapshot, + getValidAccessSnapshotForAccount, observeActiveOAuthAccessToken, resolveModelsAuthToken, type OAuthActiveTokenObservation, @@ -52,7 +53,7 @@ import { } from "../../providers/service-tier"; import type { FastPolicyAuthority } from "../../providers/fastwire"; import { effectiveGoogleMode, getProviderRegistryEntry, providerMatchesRegistryTransport, registryEntryForProviderDestination } from "../../providers/registry"; -import { parseAntigravityAvailableModels, registerAntigravityDiscoveredWireModels } from "../../providers/antigravity-models"; +import { parseAntigravityAvailableModels, registerAntigravityDiscoveredWireModels, type AntigravityAvailableModel } from "../../providers/antigravity-models"; import { recordAntigravityAccountRoster } from "../../providers/antigravity-account-roster"; import { applyProviderContextCap, providerContextCap, resolveUnknownRoutedContextWindow } from "../../providers/context-cap"; import { clampAutoCompactTokenLimit } from "../../providers/auto-compact-budget"; @@ -791,7 +792,54 @@ export async function fetchProviderModelsWithAuth( return observed(models, "degraded"); } if (antigravity) { - const upstream: CatalogModel[] = antigravity.map(model => ({ + const mergedAntigravity: AntigravityAvailableModel[] = [...antigravity]; + const seenModelIds = new Set(mergedAntigravity.map(model => model.id)); + const siblingRosters = new Map(); + if (resolveAuth.kind === "refreshing" && prov.authMode === "oauth" && auth.oauthAccountId) { + const siblings = (getAccountSet(name)?.accounts ?? []).filter( + account => account.id !== auth.oauthAccountId && account.paused !== true && account.needsReauth !== true, + ); + if (siblings.length > 0) { + await Promise.all(siblings.map(async sibling => { + try { + const snapshot = sibling.credential.access && sibling.credential.expires > Date.now() + 60_000 + ? { accessToken: sibling.credential.access, projectId: sibling.credential.projectId } + : await getValidAccessSnapshotForAccount(name, sibling.id, { requireUsableAccount: true }); + const siblingProject = prov.project ?? snapshot.projectId ?? sibling.credential.projectId; + if (!snapshot.accessToken || !siblingProject) return; + const siblingHeaders = materializeCapturedHeaders(request, snapshot.accessToken); + const siblingRes = request.method === "POST" + ? await providerOutboundPost(name, prov, url, { + headers: siblingHeaders, + body: JSON.stringify({ project: siblingProject }), + signal: AbortSignal.timeout(8000), + }, outboundDependencies) + : await providerOutboundGet(name, prov, url, { + headers: siblingHeaders, + signal: AbortSignal.timeout(8000), + }, outboundDependencies); + if (!siblingRes.ok || await providerRedirectError(siblingRes, url)) return; + const siblingBounded = await readBoundedDiscoveryJson(siblingRes, discovery.maxResponseBytes); + if (!siblingBounded.ok) return; + const siblingModels = parseAntigravityAvailableModels(siblingBounded.value, discovery.maxModels); + if (!siblingModels) return; + siblingRosters.set(sibling.id, siblingModels); + } catch { + // Fail-soft: one sibling's discovery error never degrades the pool catalog. + } + })); + for (const sibling of siblings) { + const siblingModels = siblingRosters.get(sibling.id); + if (!siblingModels) continue; + for (const model of siblingModels) { + if (seenModelIds.has(model.id)) continue; + seenModelIds.add(model.id); + mergedAntigravity.push(model); + } + } + } + } + const upstream: CatalogModel[] = mergedAntigravity.map(model => ({ id: model.id, provider: name, // Only an exact discovered wire map proves a new model's selectable effort ladder. @@ -811,13 +859,16 @@ export async function fetchProviderModelsWithAuth( } // Save exact family routes before publishing synthetic IDs. These synchronous writes // cannot interleave with a credential-generation change; failure keeps the prior cache. - registerAntigravityDiscoveredWireModels(prov.baseUrl, antigravity, { + registerAntigravityDiscoveredWireModels(prov.baseUrl, mergedAntigravity, { provider: name, cacheGeneration, }); // Entitlement is per account (Claude 4.6 vs 5.5, Argon): remember what THIS account // lists so account selection can prefer one that serves the requested model. if (auth.oauthAccountId) recordAntigravityAccountRoster(name, auth.oauthAccountId, antigravity); + for (const [siblingId, siblingModels] of siblingRosters) { + recordAntigravityAccountRoster(name, siblingId, siblingModels); + } if (!setCached(name, forCache, Date.now(), cacheGeneration, catalogAuthority)) { return observed(withConfiguredRetention(configured), "degraded"); } diff --git a/src/codex/routing-drift.ts b/src/codex/routing-drift.ts index f51b65513a7..ed324d2636a 100644 --- a/src/codex/routing-drift.ts +++ b/src/codex/routing-drift.ts @@ -78,6 +78,8 @@ function loopbackEndpoint(url: string): { hostname: string; port: number } | nul try { const parsed = new URL(url); if (!parsed.port) return null; + const pathname = parsed.pathname.replace(/\/+$/, ""); + if (pathname !== "" && pathname !== "/v1") return null; const port = Number(parsed.port); if (!Number.isInteger(port) || port <= 0 || port > 65535) return null; return { hostname: parsed.hostname.replace(/^\[|\]$/g, ""), port }; diff --git a/src/oauth/generic-account-failover.ts b/src/oauth/generic-account-failover.ts index 30fa971eef0..aaa105c34b0 100644 --- a/src/oauth/generic-account-failover.ts +++ b/src/oauth/generic-account-failover.ts @@ -337,8 +337,10 @@ export function routeAntigravityModel( return { kind: "refuse", message: antigravitySelectedAccountUnlistedMessage(model) }; } const ready = new Set(eligibleFailoverAccounts(providerName, now, classifyModelFamilyForQuota(providerName, model))); + const accountRows = new Map((getAccountSet(providerName)?.accounts ?? []).map(account => [account.id, account])); + const quotaRanked = rankAccountsByHeadroom(providerName, serving, model, accountRows); const rank = (id: string) => (ready.has(id) ? 0 : 2) + (accountListsModel(providerName, id, model) === true ? 0 : 1); - const best = serving.reduce((left, right) => rank(right) < rank(left) ? right : left); + const best = quotaRanked.reduce((left, right) => rank(right) < rank(left) ? right : left); return { kind: "reroute", accountId: best }; } @@ -762,7 +764,8 @@ export function preferredInitialAccount( ? kiroModelListingHasRoom(config, id, requestedModelId!) : accountListsModel(providerName, id, requestedModelId!) === true)) : undefined; - if (listingSibling) return listingSibling; + if (providerName === "kiro" && listingSibling) return listingSibling; + const fallbackSibling = listingSibling ?? null; // A configured strategy answers this question itself. Both guards below exist to protect the // QUOTA answer, and both are fatal to the other two: hasHeadroomEvidence refuses every @@ -779,13 +782,13 @@ export function preferredInitialAccount( } candidates = preferKiroModelSupport(providerName, candidates, requestedModelId, config); const picked = candidates.sort((a, b) => accountInFlight("kiro", a) - accountInFlight("kiro", b))[0]; - return picked && picked !== active ? picked : null; + return picked && picked !== active ? picked : fallbackSibling; } if (strategy === "round-robin") { const family = classifyModelFamilyForQuota(providerName, requestedModelId); const eligibleNow = preferKiroModelSupport(providerName, eligibleFailoverAccounts(providerName, now, family), requestedModelId, config); - if (eligibleNow.length === 0) return null; + if (eligibleNow.length === 0) return fallbackSibling; // PEEK, not pick: this proposal is discardable, and advancing the ring for an account the // resolver then rejects would skip a turn for nothing. noteGenericPoolSelection commits. const picked = peekRoundRobinAccount( @@ -793,11 +796,11 @@ export function preferredInitialAccount( eligibleNow, genericStickyLimit(config, providerName), ); - return picked && picked !== active ? picked : null; + return picked && picked !== active ? picked : fallbackSibling; } if (strategy === "fill-first") { const picked = pickFillFirstGenericAccount(config, providerName, active, now, requestedModelId, accountRows, activeAccountOverride !== undefined); - return picked && picked !== active ? picked : null; + return picked && picked !== active ? picked : fallbackSibling; } if (strategy === "reset-first") { // Same candidate set as round-robin: paused, reauth-needed and cooled accounts are out, @@ -806,13 +809,13 @@ export function preferredInitialAccount( const family = classifyModelFamilyForQuota(providerName, requestedModelId); const eligibleNow = preferKiroModelSupport(providerName, eligibleFailoverAccounts(providerName, now, family), requestedModelId, config); - if (eligibleNow.length === 0 || !hasHeadroomEvidence(providerName, eligibleNow, requestedModelId, accountRows)) return null; + if (eligibleNow.length === 0 || !hasHeadroomEvidence(providerName, eligibleNow, requestedModelId, accountRows)) return fallbackSibling; const best = rankAccountsByResetFirst(providerName, eligibleNow, requestedModelId, now, accountRows)[0] ?? null; - return best && best !== active ? best : null; + return best && best !== active ? best : fallbackSibling; } const activeRow = selected.accounts.find(account => account.id === active); - if (activeRow && activeRow.paused !== true && activeRow.needsReauth !== true + if (!activeLacks && activeRow && activeRow.paused !== true && activeRow.needsReauth !== true && !isCooled(providerName, activeRow.id, now, classifyModelFamilyForQuota(providerName, requestedModelId)) && !isAccountQuotaExhausted(providerName, activeRow.id, requestedModelId, activeRow)) return null; @@ -821,7 +824,7 @@ export function preferredInitialAccount( // 429 and the eligible list collapses to one candidate, which any ranking returns // unchanged — an answer that looks ranked but was never measured. The no-op guarantee // for quota-less providers has to be checked on the full roster. - if (!hasHeadroomEvidence(providerName, order, requestedModelId, accountRows)) return null; + if (!hasHeadroomEvidence(providerName, order, requestedModelId, accountRows)) return fallbackSibling; // Cooldowns are respected here, unlike in the presence count: this picks the account to // send to right now, and one inside its 429 window is the single candidate we hold @@ -829,20 +832,20 @@ export function preferredInitialAccount( const eligible = preferKiroModelSupport(providerName, order.filter(id => !isCooled(providerName, id, now, classifyModelFamilyForQuota(providerName, requestedModelId))), requestedModelId, config); - if (eligible.length === 0) return null; + if (eligible.length === 0) return fallbackSibling; // Start the ring at the active account so an unranked outcome reproduces today's choice. const start = active ? order.indexOf(active) : -1; const ring = start >= 0 ? [...order.slice(start), ...order.slice(0, start)] : order; const candidates = ring.filter(id => eligible.includes(id)); - if (candidates.length === 0) return null; + if (candidates.length === 0) return fallbackSibling; const best = rankAccountsByHeadroom(providerName, candidates, requestedModelId, accountRows)[0] ?? null; // Nothing to do when the ranking agrees with the account we would have used anyway. // // A proposal still needs guarded selection commit after credential resolution: a // removal, reauth verdict, or manual choice can arrive during that await. - return best && best !== active ? best : null; + return best && best !== active ? best : fallbackSibling; } /** Earliest remaining cooldown, for a client-facing Retry-After when every account is cooled. */ diff --git a/src/providers/antigravity-models.ts b/src/providers/antigravity-models.ts index 901f07560ec..8d07cfd3e49 100644 --- a/src/providers/antigravity-models.ts +++ b/src/providers/antigravity-models.ts @@ -656,7 +656,7 @@ export function parseAntigravityAvailableModels( const available = new Map>(); for (const wireId of ids) { const info = antigravityRecord(models[wireId]); - if (!info || available.has(wireId)) continue; + if (!info || info.disabled === true || wireId.endsWith("-disabled") || available.has(wireId)) continue; // Compatibility aliases are deliberately routed to NEWER wire ids for saved // selections. CCA keeps serving retired generations (3.5/3.6 Flash tiers) in its // agent list long after they stop being the model anyone should pick, so admitting diff --git a/tests/adapters/google/antigravity-account-roster.test.ts b/tests/adapters/google/antigravity-account-roster.test.ts index 199fc8038b0..ce735e559ca 100644 --- a/tests/adapters/google/antigravity-account-roster.test.ts +++ b/tests/adapters/google/antigravity-account-roster.test.ts @@ -133,6 +133,46 @@ describe("Antigravity exposure follows each account's own discovery", () => { expect(antigravityAccountSupportsModel(PROVIDER, "acct-55", "claude-sonnet-4-6")).toBe(false); expect(antigravityAccountRosterModelIds(PROVIDER, "acct-55")).toEqual(expect.arrayContaining(["claude-opus-5-5", "gemini-4-argon"])); }); + + test("disabled placeholder wire models (gemini-4-argon-disabled) are excluded from discovery and rosters", async () => { + const roster46WithDisabledArgon = { + models: { + ...ROSTER_46.models, + "gemini-4-argon-disabled": { displayName: "Gemini 4 Argon", disabled: true, tagTitle: "Unavailable", maxTokens: 1_048_576 }, + }, + agentModelSorts: [{ groups: [{ modelIds: [ + "gemini-3.8-flash-high", "gemini-3.8-flash-medium", "gemini-3.8-flash-low", + "gemini-4-argon-disabled", "claude-sonnet-4-6", "claude-opus-4-6-thinking", + ] }] }], + }; + const ids = await discoveredIds(roster46WithDisabledArgon, "acct-46-disabled"); + expect(ids).toEqual(["claude-opus-4-6-thinking", "claude-sonnet-4-6", "gemini-3.8-flash"]); + expect(antigravityAccountSupportsModel(PROVIDER, "acct-46-disabled", "gemini-4-argon")).toBe(false); + }); + + test("multi-account discovery unions models across 4.6 and 5.5 accounts while recording per-account rosters", async () => { + homeWithAccounts(["acct-46", "acct-55"], "acct-46"); + globalThis.fetch = (async (_input: RequestInfo | URL, init?: RequestInit) => { + const auth = new Headers(init?.headers).get("authorization"); + return Response.json(auth === "Bearer access-acct-46" ? ROSTER_46 : ROSTER_55); + }) as typeof fetch; + const models = await gatherRoutedModels({ providers: { [PROVIDER]: provider() } } as unknown as OcxConfig); + const ids = models.filter(model => model.provider === PROVIDER).map(model => model.id).sort(); + expect(ids).toEqual([ + "claude-opus-4-6-thinking", + "claude-opus-5-5", + "claude-sonnet-4-6", + "claude-sonnet-5-5", + "gemini-3.8-flash", + "gemini-4-argon", + ]); + expect(antigravityAccountSupportsModel(PROVIDER, "acct-46", "claude-opus-4-6-thinking")).toBe(true); + expect(antigravityAccountSupportsModel(PROVIDER, "acct-46", "claude-opus-5-5")).toBe(false); + expect(antigravityAccountSupportsModel(PROVIDER, "acct-46", "gemini-4-argon")).toBe(false); + expect(antigravityAccountSupportsModel(PROVIDER, "acct-55", "claude-opus-5-5")).toBe(true); + expect(antigravityAccountSupportsModel(PROVIDER, "acct-55", "gemini-4-argon")).toBe(true); + expect(antigravityAccountSupportsModel(PROVIDER, "acct-55", "claude-opus-4-6-thinking")).toBe(false); + }); }); describe("Claude 4.6 and 5.5 both map", () => {