fix(antigravity): route requests to an account whose plan lists the model

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 (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

src/sse/handlers/chat.ts

@@ -1812,6 +1812,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 }

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 { isRequestScopedMisalignmentRefusal } from "@omniroute/open-sse/utils/providerSafety.ts";
 import { buildAllExpiredCredentials } from "./authExpiredCredentials.ts";
@@ -1008,6 +1009,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);
@@ -1266,8 +1270,18 @@ 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.
+    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)) {

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,
@@ -43,9 +44,115 @@ export function isRequestedModelAdvertised(
   return baseModel !== rawModel && isModelAdvertisedByConnection(baseModel, advertised);
 }
 
+type InventoryAccumulator = {
+  advertised: Map<string, Set<string>>;
+  // 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<string, Set<string>>;
+  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<string>;
+};
+
+function recordPerAccountInventory(
+  connectionId: string,
+  ids: Set<string>,
+  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;
+}
+
+/** 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)) recordPerAccountInventory(connectionId, ids, acc);
+    else acc.advertised.set(connectionId, ids);
+  }
+}
+
+async function collectProviderInventory(
+  providerId: string,
+  connections: ProviderConnectionView[],
+  requestedModel: string,
+  acc: InventoryAccumulator
+): Promise<void> {
+  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);
+}
+
+function isRowInventoryProvider(provider: unknown, requestedModel: string): provider is string {
+  return (
+    typeof provider === "string" &&
+    (isSelfHostedChatProvider(provider) ||
+      (provider === "codex" && !isCodexImageModel(requestedModel)))
+  );
+}
+
 /**
- * Self-hosted hosts and Codex subscription accounts can advertise different
- * inventories. For Codex, once an account inventory advertises the requested
+ * 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<string> {
+  const providers = new Set<string>();
+  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
  * model, unsynchronized accounts are not eligible guesses for it. With no
  * inventories at all, retain the existing bootstrap/offline behavior.
  *
@@ -54,55 +161,51 @@ export function isRequestedModelAdvertised(
  * account can serve it, so Codex accounts are not filtered for it and the
  * upstream decides. Filtering every account answered a misleading
  * provider_not_configured despite active accounts.
+ *
+ * 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). 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<Map<string, Set<string>>> {
   const advertised = new Map<string, Set<string>>();
   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)))
-      )
+  const inventoryProviders = collectInventoryProviders(
+    connections,
+    poolConnections,
+    requestedModel
   );
   if (inventoryProviders.size === 0) return advertised;
 
+  const acc: InventoryAccumulator = {
+    advertised,
+    perAccountInventory: new Map<string, Set<string>>(),
+    perAccountMalformed: false,
+    perAccountPoolIds: new Set([...poolConnections, ...connections].map((c) => c.id)),
+  };
+
   await Promise.all(
-    [...inventoryProviders].map(async (providerId) => {
-      let byConnection: SyncedAvailableModelsByConnection;
-      try {
-        byConnection = await getSyncedAvailableModelsByConnection(providerId);
-      } catch {
-        return;
-      }
-      // Malformed persisted rows: fail open for the whole provider.
-      if (byConnection[SYNCED_AVAILABLE_MODELS_MALFORMED]) 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;
-        advertised.set(connectionId, new Set(models.map((m) => m.id)));
-      }
-    })
+    [...inventoryProviders].map((providerId) =>
+      collectProviderInventory(providerId, connections, requestedModel, acc)
+    )
   );
 
+  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,
+      acc.perAccountInventory
+    )) {
+      advertised.set(connectionId, ids);
+    }
+  }
+
   return advertised;
 }

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;

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<string, Set<string>>
+): Map<string, Set<string>> {
+  for (const models of inventory.values()) {
+    if (models.size > 0 && isModelAdvertisedByConnection(requestedModel, models)) return inventory;
+  }
+  return new Map();
+}

stryker.conf.json

@@ -64,6 +64,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",