From 628452d5dcafbbe8d630f6bd211a1f9967376b2b Mon Sep 17 00:00:00 2001 From: bufftop25 <147844179+bufftop25@users.noreply.github.com> Date: Sat, 31 Oct 2026 10:28:54 -0300 Subject: [PATCH 1/5] fix(antigravity): route each request to an account whose plan lists the model Antigravity entitlement is per account: on a typical pool only Ultra accounts list the Claude 5.5 models, and Gemini 4 Argon is listed only by some agy accounts. Account selection ignored each account's synced inventory for these hosted providers (the per-connection filter in loadAdvertisedModelsForConnections only covered self-hosted providers and Codex), so a 5.5 or Argon request started on, and failed over to, accounts that do not have the model. Apply the per-connection inventory filter to antigravity and agy too, judging both inventories together because they share one account pool. Fail open when no account advertises the model, so stale or incomplete lists never empty the pool, and keep never-synced accounts eligible. An account the operator pinned for this request (x-omniroute-connection, e.g. the dashboard test) bypasses the inferred filter so an entitlement the catalog under-reports can still be verified; combo-derived pins are still filtered. --- src/sse/handlers/chat.ts | 1 + src/sse/services/auth.ts | 10 ++ src/sse/services/connectionModelInventory.ts | 28 ++- .../services/credentialSelectionOptions.ts | 2 + src/sse/services/perAccountCatalogFilter.ts | 31 ++++ ...tigravity-plan-aware-model-routing.test.ts | 164 ++++++++++++++++++ 6 files changed, 233 insertions(+), 3 deletions(-) create mode 100644 src/sse/services/perAccountCatalogFilter.ts create mode 100644 tests/unit/antigravity-plan-aware-model-routing.test.ts diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index 5c30c08bf427..691b54151380 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -1799,6 +1799,7 @@ async function handleSingleModelChat( preferredConnectionIds: apiKeyInfo?.preferredConnections ?? null, reserveOAuthSession: true, excludeConnectionIds: Array.from(excluded), + ...(hasForcedConnection && !isCombo ? { explicitConnectionPin: true } : {}), ...(agy.on ? { reserveAntigravityLease: true, routingRequestId: agy.requestId } : {}), ...(runtimeOptions.allowRateLimitedConnection ? { allowRateLimitedConnections: true } diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index b97c66fc7311..be9f22c8e07e 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -6,6 +6,7 @@ import { getProviderSearchPool } from "@omniroute/open-sse/services/providerConn import { hydrateConnectionProviderSpecificData } from "./compatibleNodeBaseUrl.ts"; // #13452 import { extractGoogApiKeyHeader } from "./googApiKeyAuth.ts"; import { isCooldownRetrySkipRefusalsEnabled } from "./cooldownAwareRetry.ts"; +import { isPerAccountCatalogProvider } from "./perAccountCatalogFilter.ts"; import { describeUpstreamFailure } from "@/shared/utils/upstreamError"; import { buildAllExpiredCredentials } from "./authExpiredCredentials.ts"; import { pickExpiryFirstConnection } from "./expiryFirstAccountSelection.ts"; @@ -1239,6 +1240,15 @@ export async function getProviderCredentials( connections, requestedModel ); + // An account the operator pinned for this request (dashboard test) is their call: the inferred + // per-account inventory must not veto it, or an entitlement Google under-reports is untestable. + if ( + forcedConnectionId && + options.explicitConnectionPin && + isPerAccountCatalogProvider(provider) + ) { + advertisedModelsByConnection.delete(forcedConnectionId); + } // Filter out unavailable accounts and excluded connection let availableConnections = connections.filter((c) => { if (excludedConnectionIds.has(c.id)) { diff --git a/src/sse/services/connectionModelInventory.ts b/src/sse/services/connectionModelInventory.ts index c88dc9a9a1b8..e517ae5fe64d 100644 --- a/src/sse/services/connectionModelInventory.ts +++ b/src/sse/services/connectionModelInventory.ts @@ -2,6 +2,7 @@ import { splitCodexReasoningSuffix } from "@omniroute/open-sse/executors/codex/r import { isRegisteredImageModel } from "@omniroute/open-sse/config/imageRegistry"; import { isModelAdvertisedByConnection } from "@/domain/connectionModelRules"; import { isSelfHostedChatProvider } from "@/shared/constants/providers"; +import { inventoryOrFailOpen, isPerAccountCatalogProvider } from "./perAccountCatalogFilter.ts"; import { type ProviderConnectionView } from "@/lib/db/providers/lazyConnectionView"; import { getSyncedAvailableModelsByConnection, @@ -60,21 +61,32 @@ export async function loadAdvertisedModelsForConnections( .filter( (p): p is string => typeof p === "string" && - (isSelfHostedChatProvider(p) || (p === "codex" && !isCodexImageModel(requestedModel))) + (isSelfHostedChatProvider(p) || + (p === "codex" && !isCodexImageModel(requestedModel)) || + isPerAccountCatalogProvider(p)) ) ); if (inventoryProviders.size === 0) return advertised; + // Per-account catalogs (Antigravity plans): `antigravity` and `agy` connections share one pool, + // so their inventories are merged and judged together before the fail-open rule applies. + const perAccountInventory = new Map>(); + let perAccountMalformed = false; + await Promise.all( [...inventoryProviders].map(async (providerId) => { let byConnection: SyncedAvailableModelsByConnection; try { byConnection = await getSyncedAvailableModelsByConnection(providerId); } catch { + if (isPerAccountCatalogProvider(providerId)) perAccountMalformed = true; return; } // Malformed persisted rows: fail open for the whole provider. - if (byConnection[SYNCED_AVAILABLE_MODELS_MALFORMED]) return; + if (byConnection[SYNCED_AVAILABLE_MODELS_MALFORMED]) { + if (isPerAccountCatalogProvider(providerId)) perAccountMalformed = true; + return; + } if ( providerId === "codex" && Object.values(byConnection).some((models) => Array.isArray(models) && models.length > 0) @@ -85,10 +97,20 @@ export async function loadAdvertisedModelsForConnections( } for (const [connectionId, models] of Object.entries(byConnection)) { if (!Array.isArray(models) || models.length === 0) continue; - advertised.set(connectionId, new Set(models.map((m) => m.id))); + const ids = new Set(models.map((m) => m.id)); + if (isPerAccountCatalogProvider(providerId)) perAccountInventory.set(connectionId, ids); + else advertised.set(connectionId, ids); } }) ); + if (!perAccountMalformed) { + // When no account advertises the model the lists are stale/incomplete: fail open instead of + // emptying the pool. + for (const [connectionId, ids] of inventoryOrFailOpen(requestedModel, perAccountInventory)) { + advertised.set(connectionId, ids); + } + } + return advertised; } diff --git a/src/sse/services/credentialSelectionOptions.ts b/src/sse/services/credentialSelectionOptions.ts index 7ec6aba78099..344c5ceab55c 100644 --- a/src/sse/services/credentialSelectionOptions.ts +++ b/src/sse/services/credentialSelectionOptions.ts @@ -12,6 +12,8 @@ export interface CredentialSelectionOptions { forcedConnectionId?: string | null; /** API-key-scoped ordered preference. Only ranks connections that survived eligibility gates. */ preferredConnectionIds?: string[] | null; + /** The pin was chosen by the operator for this request (x-omniroute-connection), not derived by a combo. */ + explicitConnectionPin?: boolean; excludeConnectionIds?: string[] | null; sessionKey?: string | null; sessionAffinityTtlMs?: number | null; diff --git a/src/sse/services/perAccountCatalogFilter.ts b/src/sse/services/perAccountCatalogFilter.ts new file mode 100644 index 000000000000..271c93c27ba7 --- /dev/null +++ b/src/sse/services/perAccountCatalogFilter.ts @@ -0,0 +1,31 @@ +import { isModelAdvertisedByConnection } from "@/domain/connectionModelRules"; + +/** + * Hosted providers whose model entitlement differs PER ACCOUNT (plan/tier), so the + * per-connection synced inventory is authoritative for routing. Antigravity: only + * Ultra accounts are offered the 5.5 Claude models; the rest list 4.6 only and answer + * 403/404/429 to a 5.5 request, burning a retry (and a cooldown) per account. + * `agy` connections are served from the same account pool as `antigravity`, and a model can + * be advertised only by them (Gemini 4 Argon), so both inventories are judged together. + */ +const PER_ACCOUNT_CATALOG_PROVIDERS = new Set(["antigravity", "agy"]); + +export function isPerAccountCatalogProvider(providerId: unknown): boolean { + return typeof providerId === "string" && PER_ACCOUNT_CATALOG_PROVIDERS.has(providerId); +} + +/** + * Keep a provider's per-connection inventory only when at least one account really + * advertises the requested model. If none does, the synced lists are stale or + * incomplete (a hidden model, a sync that predates a plan upgrade): fail OPEN and + * return an empty map so selection is unchanged instead of emptying the whole pool. + */ +export function inventoryOrFailOpen( + requestedModel: string, + inventory: Map> +): Map> { + for (const models of inventory.values()) { + if (models.size > 0 && isModelAdvertisedByConnection(requestedModel, models)) return inventory; + } + return new Map(); +} diff --git a/tests/unit/antigravity-plan-aware-model-routing.test.ts b/tests/unit/antigravity-plan-aware-model-routing.test.ts new file mode 100644 index 000000000000..33659e54d905 --- /dev/null +++ b/tests/unit/antigravity-plan-aware-model-routing.test.ts @@ -0,0 +1,164 @@ +/** + * Antigravity entitlement is per ACCOUNT (plan): only Ultra accounts are offered the 5.5 Claude + * models. Routing must (1) skip accounts whose live inventory lacks the requested model, + * (2) fail open when no account advertises it, and (3) let an operator-pinned account through. + */ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-agy-plan-routing-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); +const providersDb = await import("../../src/lib/db/providers.ts"); +const modelsDb = await import("../../src/lib/db/models.ts"); +const auth = await import("../../src/sse/services/auth.ts"); + +const PROVIDER = "antigravity"; +const OPUS55 = "claude-opus-5-5-low"; + +async function resetStorage() { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +test.after(() => { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 }); +}); + +async function account(name: string, priority: number, models: string[]): Promise { + const created = (await providersDb.createProviderConnection({ + provider: PROVIDER, + authType: "oauth", + name, + accessToken: `token-${name}`, + priority, + isActive: true, + })) as { id: string }; + if (models.length > 0) { + await modelsDb.replaceSyncedAvailableModelsForConnection( + PROVIDER, + created.id, + models.map((id) => ({ id, name: id })) + ); + } + return created.id; +} + +const idOf = (selected: unknown): string | null => { + if (!selected || typeof selected !== "object") return null; + const id = (selected as { connectionId?: unknown }).connectionId; + return typeof id === "string" ? id : null; +}; + +const OLD = ["claude-opus-4-6-thinking", "claude-sonnet-4-6"]; +const ULTRA = [...OLD, OPUS55, "claude-opus-5-5-high"]; + +test("a 5.5 request skips the accounts whose plan does not list it", async () => { + await resetStorage(); + const pro1 = await account("pro1", 1, OLD); + const pro2 = await account("pro2", 2, OLD); + const ultra = await account("ultra", 3, ULTRA); + + const first = await auth.getProviderCredentials(PROVIDER, null, null, OPUS55); + assert.equal(idOf(first), ultra, "only the Ultra account advertises the 5.5 model"); + assert.notEqual(idOf(first), pro1); + assert.notEqual(idOf(first), pro2); + + const next = await auth.getProviderCredentials(PROVIDER, ultra, null, OPUS55); + assert.notEqual(idOf(next), pro1, "failover must not wander onto an account without the model"); + assert.notEqual(idOf(next), pro2); +}); + +test("a model every account lists keeps all accounts eligible (4.6 failover intact)", async () => { + await resetStorage(); + const a = await account("a", 1, ULTRA); + const b = await account("b", 2, OLD); + + assert.equal( + idOf(await auth.getProviderCredentials(PROVIDER, null, null, "claude-sonnet-4-6")), + a + ); + assert.equal(idOf(await auth.getProviderCredentials(PROVIDER, a, null, "claude-sonnet-4-6")), b); +}); + +test("fails open when no account advertises the model (stale or incomplete lists)", async () => { + await resetStorage(); + const a = await account("a", 1, OLD); + await account("b", 2, OLD); + + const selected = await auth.getProviderCredentials(PROVIDER, null, null, "claude-opus-9-9-low"); + assert.equal(idOf(selected), a, "an unlisted model must not empty the pool"); +}); + +test("accounts never synced stay eligible (unknown is not 'does not have it')", async () => { + await resetStorage(); + const unsynced = await account("unsynced", 1, []); + await account("ultra", 2, ULTRA); + + const selected = await auth.getProviderCredentials(PROVIDER, null, null, OPUS55); + assert.equal(idOf(selected), unsynced); +}); + +test("inventoryOrFailOpen: keeps the inventory only when some account advertises the model", async () => { + const filter = await import("../../src/sse/services/perAccountCatalogFilter.ts"); + const inv = new Map([ + ["x", new Set(["m1"])], + ["y", new Set(["m2"])], + ]); + assert.equal(filter.inventoryOrFailOpen("m1", inv), inv); + assert.equal(filter.inventoryOrFailOpen("missing", inv).size, 0); + assert.equal(filter.isPerAccountCatalogProvider("antigravity"), true); + assert.equal(filter.isPerAccountCatalogProvider("openai"), false); +}); + +test("an account the operator pinned is not vetoed, but a combo-derived pin still is", async () => { + await resetStorage(); + const pro = await account("pro", 1, OLD); + const ultra = await account("ultra", 2, ULTRA); + + const pinned = await auth.getProviderCredentials(PROVIDER, null, null, OPUS55, { + forcedConnectionId: pro, + explicitConnectionPin: true, + }); + assert.equal(idOf(pinned), pro, "operator pinned this account; it must be tried as asked"); + + const comboTarget = await auth.getProviderCredentials(PROVIDER, null, null, OPUS55, { + forcedConnectionId: pro, + }); + assert.equal( + idOf(comboTarget), + null, + "a combo expansion must not spend calls on a non-covering account" + ); + + const automatic = await auth.getProviderCredentials(PROVIDER, null, null, OPUS55); + assert.equal(idOf(automatic), ultra, "automatic selection still follows the live inventory"); +}); + +test("a model only the agy accounts list (Gemini 4 Argon) skips the antigravity-provider accounts", async () => { + await resetStorage(); + const ARGON = "gemini-4-argon-low"; + const proAccount = await account("pro", 1, OLD); + const agy = (await providersDb.createProviderConnection({ + provider: "agy", + authType: "oauth", + name: "ultra-agy", + accessToken: "token-ultra-agy", + priority: 2, + isActive: true, + })) as { id: string }; + await modelsDb.replaceSyncedAvailableModelsForConnection("agy", agy.id, [ + { id: ARGON, name: ARGON }, + ...OLD.map((id) => ({ id, name: id })), + ]); + + const first = await auth.getProviderCredentials(PROVIDER, null, null, ARGON); + assert.equal(idOf(first), agy.id, "only the agy account advertises Argon"); + assert.notEqual(idOf(first), proAccount); +}); From 4f6dd83dfb54e80bfbe23bf0e8918aecb5c82ea2 Mon Sep 17 00:00:00 2001 From: bufftop25 <147844179+bufftop25@users.noreply.github.com> Date: Sat, 31 Oct 2026 10:30:19 -0300 Subject: [PATCH 2/5] chore(changelog): add the fragment for #18665 --- changelog.d/fixes/18665-antigravity-route-by-account-plan.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog.d/fixes/18665-antigravity-route-by-account-plan.md diff --git a/changelog.d/fixes/18665-antigravity-route-by-account-plan.md b/changelog.d/fixes/18665-antigravity-route-by-account-plan.md new file mode 100644 index 000000000000..a434f96cde50 --- /dev/null +++ b/changelog.d/fixes/18665-antigravity-route-by-account-plan.md @@ -0,0 +1 @@ +- **fix(antigravity):** Route each request to an Antigravity/agy account whose synced model list includes the requested model, so Claude 5.5 and Gemini 4 Argon requests no longer start on or fail over to accounts whose plan lacks them; fails open when no account lists the model, keeps never-synced accounts eligible, and lets an operator-pinned account through ([#18665](https://github.com/cheaper-inference/OmniRoute/pull/18665)) — thanks @bufftop25 From cdea95cf978adabac8dcfbfcba9b9d7246355957 Mon Sep 17 00:00:00 2001 From: bufftop25 <147844179+bufftop25@users.noreply.github.com> Date: Mon, 2 Nov 2026 22:35:03 -0300 Subject: [PATCH 3/5] docs(antigravity): note per-account catalog filtering and combo-pin skip Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com> --- changelog.d/fixes/18665-antigravity-route-by-account-plan.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changelog.d/fixes/18665-antigravity-route-by-account-plan.md b/changelog.d/fixes/18665-antigravity-route-by-account-plan.md index a434f96cde50..64e613fba8d4 100644 --- a/changelog.d/fixes/18665-antigravity-route-by-account-plan.md +++ b/changelog.d/fixes/18665-antigravity-route-by-account-plan.md @@ -1 +1 @@ -- **fix(antigravity):** Route each request to an Antigravity/agy account whose synced model list includes the requested model, so Claude 5.5 and Gemini 4 Argon requests no longer start on or fail over to accounts whose plan lacks them; fails open when no account lists the model, keeps never-synced accounts eligible, and lets an operator-pinned account through ([#18665](https://github.com/cheaper-inference/OmniRoute/pull/18665)) — thanks @bufftop25 +- **fix(antigravity):** Route each request to an Antigravity/agy account whose synced model list includes the requested model, so Claude 5.5 and Gemini 4 Argon requests no longer start on or fail over to accounts whose plan lacks them; fails open when no account lists the model, keeps never-synced accounts eligible, and lets an operator-pinned account through (a combo step pinned to an account whose plan lacks the model is skipped while another account lists it) ([#18665](https://github.com/cheaper-inference/OmniRoute/pull/18665)) — thanks @bufftop25 From ae38f169bff013d6aa8cd0be795041b7724d36f6 Mon Sep 17 00:00:00 2001 From: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:40:07 -0300 Subject: [PATCH 4/5] refactor(sse): split connection model inventory loading into helpers Moves the per-provider inventory collection out of the Promise.all callback in loadAdvertisedModelsForConnections to stay under the cognitive-complexity ratchet, and lists the #18665 test in stryker tap.testFiles. No behavior change. --- src/sse/services/connectionModelInventory.ts | 121 +++++++++++++------ stryker.conf.json | 1 + 2 files changed, 82 insertions(+), 40 deletions(-) diff --git a/src/sse/services/connectionModelInventory.ts b/src/sse/services/connectionModelInventory.ts index 8d2f5198b17b..68167ef476bd 100644 --- a/src/sse/services/connectionModelInventory.ts +++ b/src/sse/services/connectionModelInventory.ts @@ -44,6 +44,74 @@ export function isRequestedModelAdvertised( return baseModel !== rawModel && isModelAdvertisedByConnection(baseModel, advertised); } +type InventoryAccumulator = { + advertised: Map>; + // Per-account catalogs (Antigravity plans): `antigravity` and `agy` connections share one pool, + // so their inventories are merged and judged together before the fail-open rule applies. + perAccountInventory: Map>; + perAccountMalformed: boolean; +}; + +/** Unreadable or malformed persisted rows: a per-account catalog can no longer be trusted. */ +function markPerAccountMalformed(providerId: string, acc: InventoryAccumulator): void { + if (isPerAccountCatalogProvider(providerId)) acc.perAccountMalformed = true; +} + +/** Does any Codex account's non-empty synced inventory advertise the requested model? */ +function isAdvertisedByAnyCodexAccount( + byConnection: SyncedAvailableModelsByConnection, + requestedModel: string +): boolean { + return Object.values(byConnection).some( + (models) => + Array.isArray(models) && + // An empty set means "advertises everything" to the matcher, so an + // empty row must not count as evidence for the requested model. + models.length > 0 && + isRequestedModelAdvertised("codex", requestedModel, new Set(models.map((m) => m.id))) + ); +} + +function recordConnectionInventories( + providerId: string, + byConnection: SyncedAvailableModelsByConnection, + acc: InventoryAccumulator +): void { + for (const [connectionId, models] of Object.entries(byConnection)) { + if (!Array.isArray(models) || models.length === 0) continue; + const ids = new Set(models.map((m) => m.id)); + if (isPerAccountCatalogProvider(providerId)) acc.perAccountInventory.set(connectionId, ids); + else acc.advertised.set(connectionId, ids); + } +} + +async function collectProviderInventory( + providerId: string, + connections: ProviderConnectionView[], + requestedModel: string, + acc: InventoryAccumulator +): Promise { + let byConnection: SyncedAvailableModelsByConnection; + try { + byConnection = await getSyncedAvailableModelsByConnection(providerId); + } catch { + markPerAccountMalformed(providerId, acc); + return; + } + // Malformed persisted rows: fail open for the whole provider. + if (byConnection[SYNCED_AVAILABLE_MODELS_MALFORMED]) { + markPerAccountMalformed(providerId, acc); + return; + } + if (providerId === "codex") { + if (!isAdvertisedByAnyCodexAccount(byConnection, requestedModel)) return; + for (const connection of connections) { + if (connection.provider === providerId) acc.advertised.set(connection.id, new Set()); + } + } + recordConnectionInventories(providerId, byConnection, acc); +} + /** * Self-hosted hosts, Codex subscription accounts and Antigravity/agy accounts (plans list * different models) can advertise different inventories. For Codex, once an account inventory advertises the requested @@ -80,52 +148,25 @@ export async function loadAdvertisedModelsForConnections( ); if (inventoryProviders.size === 0) return advertised; - // Per-account catalogs (Antigravity plans): `antigravity` and `agy` connections share one pool, - // so their inventories are merged and judged together before the fail-open rule applies. - const perAccountInventory = new Map>(); - let perAccountMalformed = false; + const acc: InventoryAccumulator = { + advertised, + perAccountInventory: new Map>(), + perAccountMalformed: false, + }; await Promise.all( - [...inventoryProviders].map(async (providerId) => { - let byConnection: SyncedAvailableModelsByConnection; - try { - byConnection = await getSyncedAvailableModelsByConnection(providerId); - } catch { - if (isPerAccountCatalogProvider(providerId)) perAccountMalformed = true; - return; - } - // Malformed persisted rows: fail open for the whole provider. - if (byConnection[SYNCED_AVAILABLE_MODELS_MALFORMED]) { - if (isPerAccountCatalogProvider(providerId)) perAccountMalformed = true; - return; - } - if (providerId === "codex") { - const advertisedByAnyAccount = Object.values(byConnection).some( - (models) => - Array.isArray(models) && - // An empty set means "advertises everything" to the matcher, so an - // empty row must not count as evidence for the requested model. - models.length > 0 && - isRequestedModelAdvertised(providerId, requestedModel, new Set(models.map((m) => m.id))) - ); - if (!advertisedByAnyAccount) return; - for (const connection of connections) { - if (connection.provider === providerId) advertised.set(connection.id, new Set()); - } - } - for (const [connectionId, models] of Object.entries(byConnection)) { - if (!Array.isArray(models) || models.length === 0) continue; - const ids = new Set(models.map((m) => m.id)); - if (isPerAccountCatalogProvider(providerId)) perAccountInventory.set(connectionId, ids); - else advertised.set(connectionId, ids); - } - }) + [...inventoryProviders].map((providerId) => + collectProviderInventory(providerId, connections, requestedModel, acc) + ) ); - if (!perAccountMalformed) { + if (!acc.perAccountMalformed) { // When no account advertises the model the lists are stale/incomplete: fail open instead of // emptying the pool. - for (const [connectionId, ids] of inventoryOrFailOpen(requestedModel, perAccountInventory)) { + for (const [connectionId, ids] of inventoryOrFailOpen( + requestedModel, + acc.perAccountInventory + )) { advertised.set(connectionId, ids); } } diff --git a/stryker.conf.json b/stryker.conf.json index b5abd93f9f7a..8e345eac65cf 100644 --- a/stryker.conf.json +++ b/stryker.conf.json @@ -55,6 +55,7 @@ "tests/unit/noauth-autocombo-lockout-7623.test.ts", "tests/unit/auth-opencode-zen-noauth-fallback.test.ts", "tests/unit/codex-account-catalog-routing.test.ts", + "tests/unit/antigravity-plan-aware-model-routing.test.ts", "tests/unit/provider-transfer.test.ts", "tests/unit/credential-scope-diagnostics-15889.test.ts", "tests/unit/account-weighted-runtime-16459.test.ts", From 4c4653f4c174f743767f484c0c3b80831e7624d0 Mon Sep 17 00:00:00 2001 From: bufftop25 <147844179+bufftop25@users.noreply.github.com> Date: Wed, 4 Nov 2026 14:10:50 -0300 Subject: [PATCH 5/5] fix(sse): judge antigravity/agy plan inventory over the active pool before pin narrowing The fail-open decision now counts only accounts of the active, key-scoped pool, so a synced row left by a disabled or deleted account can no longer empty the pool. The pool is captured before a combo pin narrows the connections, so a pinned antigravity account still sees the agy inventory for an agy-only model. Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com> --- src/sse/services/auth.ts | 6 +- src/sse/services/connectionModelInventory.ts | 62 ++++++++++++++---- ...tigravity-plan-aware-model-routing.test.ts | 65 +++++++++++++++++++ 3 files changed, 119 insertions(+), 14 deletions(-) diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index 2a82b5984c12..4d6639d48352 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -981,6 +981,9 @@ export async function getProviderCredentials( if (allowedConnections && allowedConnections.length > 0) { connections = connections.filter((conn) => allowedConnections.includes(conn.id)); } + // The active, key-scoped pool before any pin narrows it: per-account catalogs (antigravity/agy) + // judge "does some account list the model" over the whole pool, not just the pinned account. + const inventoryPoolConnections = connections; let explicitProbeKind: "probe" | "suppressed" | "skip" = "skip"; if (forcedConnectionId && !connections.some((c) => c.id === forcedConnectionId)) { const pinnedRaw = await getCachedProviderConnectionById(forcedConnectionId); @@ -1239,7 +1242,8 @@ export async function getProviderCredentials( // the host that does. Empty map = no inventory known = no filtering. const advertisedModelsByConnection = await loadAdvertisedModelsForConnections( connections, - requestedModel + requestedModel, + inventoryPoolConnections ); // An account the operator pinned for this request (dashboard test) is their call: the inferred // per-account inventory must not veto it, or an entitlement Google under-reports is untestable. diff --git a/src/sse/services/connectionModelInventory.ts b/src/sse/services/connectionModelInventory.ts index 68167ef476bd..6f525584a60e 100644 --- a/src/sse/services/connectionModelInventory.ts +++ b/src/sse/services/connectionModelInventory.ts @@ -50,8 +50,19 @@ type InventoryAccumulator = { // so their inventories are merged and judged together before the fail-open rule applies. perAccountInventory: Map>; perAccountMalformed: boolean; + // Accounts of the active shared pool (before any pin narrows it). A synced row left behind by + // a disabled or deleted account must not decide that "some account lists the model". + perAccountPoolIds: Set; }; +function recordPerAccountInventory( + connectionId: string, + ids: Set, + acc: InventoryAccumulator +): void { + if (acc.perAccountPoolIds.has(connectionId)) acc.perAccountInventory.set(connectionId, ids); +} + /** Unreadable or malformed persisted rows: a per-account catalog can no longer be trusted. */ function markPerAccountMalformed(providerId: string, acc: InventoryAccumulator): void { if (isPerAccountCatalogProvider(providerId)) acc.perAccountMalformed = true; @@ -80,7 +91,7 @@ function recordConnectionInventories( for (const [connectionId, models] of Object.entries(byConnection)) { if (!Array.isArray(models) || models.length === 0) continue; const ids = new Set(models.map((m) => m.id)); - if (isPerAccountCatalogProvider(providerId)) acc.perAccountInventory.set(connectionId, ids); + if (isPerAccountCatalogProvider(providerId)) recordPerAccountInventory(connectionId, ids, acc); else acc.advertised.set(connectionId, ids); } } @@ -112,6 +123,33 @@ async function collectProviderInventory( recordConnectionInventories(providerId, byConnection, acc); } +function isRowInventoryProvider(provider: unknown, requestedModel: string): provider is string { + return ( + typeof provider === "string" && + (isSelfHostedChatProvider(provider) || + (provider === "codex" && !isCodexImageModel(requestedModel))) + ); +} + +/** + * Providers whose synced inventory must be read. Per-account catalogs (antigravity/agy) come + * from the whole active pool so a pin on one of them still sees the other's inventory. + */ +function collectInventoryProviders( + connections: ProviderConnectionView[], + poolConnections: ProviderConnectionView[], + requestedModel: string +): Set { + const providers = new Set(); + for (const { provider } of connections) { + if (isRowInventoryProvider(provider, requestedModel)) providers.add(provider); + } + for (const { provider } of [...connections, ...poolConnections]) { + if (isPerAccountCatalogProvider(provider)) providers.add(provider as string); + } + return providers; +} + /** * Self-hosted hosts, Codex subscription accounts and Antigravity/agy accounts (plans list * different models) can advertise different inventories. For Codex, once an account inventory advertises the requested @@ -126,25 +164,22 @@ async function collectProviderInventory( * * Antigravity and agy share one pool and one per-account catalog: accounts whose synced * inventory does not list the requested model are skipped, failing open when no account - * lists it (see perAccountCatalogFilter.ts). + * lists it (see perAccountCatalogFilter.ts). That decision is taken over `poolConnections`, + * the active pool BEFORE a pin narrowed `connections`: a pinned `antigravity` account must + * still see the `agy` inventory, and rows of accounts outside the active pool are ignored. */ export async function loadAdvertisedModelsForConnections( connections: ProviderConnectionView[], - requestedModel: string | null + requestedModel: string | null, + poolConnections: ProviderConnectionView[] = connections ): Promise>> { const advertised = new Map>(); if (!requestedModel) return advertised; - const inventoryProviders = new Set( - connections - .map((c) => c.provider) - .filter( - (p): p is string => - typeof p === "string" && - (isSelfHostedChatProvider(p) || - (p === "codex" && !isCodexImageModel(requestedModel)) || - isPerAccountCatalogProvider(p)) - ) + const inventoryProviders = collectInventoryProviders( + connections, + poolConnections, + requestedModel ); if (inventoryProviders.size === 0) return advertised; @@ -152,6 +187,7 @@ export async function loadAdvertisedModelsForConnections( advertised, perAccountInventory: new Map>(), perAccountMalformed: false, + perAccountPoolIds: new Set([...poolConnections, ...connections].map((c) => c.id)), }; await Promise.all( diff --git a/tests/unit/antigravity-plan-aware-model-routing.test.ts b/tests/unit/antigravity-plan-aware-model-routing.test.ts index 33659e54d905..86230609200d 100644 --- a/tests/unit/antigravity-plan-aware-model-routing.test.ts +++ b/tests/unit/antigravity-plan-aware-model-routing.test.ts @@ -162,3 +162,68 @@ test("a model only the agy accounts list (Gemini 4 Argon) skips the antigravity- assert.equal(idOf(first), agy.id, "only the agy account advertises Argon"); assert.notEqual(idOf(first), proAccount); }); + +test("an inactive or deleted account's leftover synced row does not empty the pool", async () => { + await resetStorage(); + const active = await account("active-pro", 1, OLD); + // A disabled Ultra account keeps its synced row after being switched off. + const disabled = (await providersDb.createProviderConnection({ + provider: PROVIDER, + authType: "oauth", + name: "disabled-ultra", + accessToken: "token-disabled-ultra", + priority: 2, + isActive: false, + })) as { id: string }; + await modelsDb.replaceSyncedAvailableModelsForConnection( + PROVIDER, + disabled.id, + ULTRA.map((id) => ({ id, name: id })) + ); + // A row whose account no longer exists (removed from the pool). + await modelsDb.replaceSyncedAvailableModelsForConnection( + "agy", + "removed-ultra-connection", + ULTRA.map((id) => ({ id, name: id })) + ); + + const selected = await auth.getProviderCredentials(PROVIDER, null, null, OPUS55); + assert.equal( + idOf(selected), + active, + "only active pool accounts decide whether some account lists the model (fail open otherwise)" + ); +}); + +test("a combo pin on an antigravity account still sees the agy inventory for an agy-only model", async () => { + await resetStorage(); + const ARGON = "gemini-4-argon-low"; + const proAccount = await account("pro", 1, OLD); + const agy = (await providersDb.createProviderConnection({ + provider: "agy", + authType: "oauth", + name: "ultra-agy", + accessToken: "token-ultra-agy", + priority: 2, + isActive: true, + })) as { id: string }; + await modelsDb.replaceSyncedAvailableModelsForConnection("agy", agy.id, [ + { id: ARGON, name: ARGON }, + ...OLD.map((id) => ({ id, name: id })), + ]); + + const comboTarget = await auth.getProviderCredentials(PROVIDER, null, null, ARGON, { + forcedConnectionId: proAccount, + }); + assert.equal( + idOf(comboTarget), + null, + "the agy account lists Argon, so a combo step pinned to the antigravity account is skipped" + ); + + const operatorPinned = await auth.getProviderCredentials(PROVIDER, null, null, ARGON, { + forcedConnectionId: proAccount, + explicitConnectionPin: true, + }); + assert.equal(idOf(operatorPinned), proAccount, "an operator pin is still tried as asked"); +});