[v3.8.50] fix(antigravity): lock full quota per exact model

open-sse/executors/antigravity.ts

@@ -13,6 +13,7 @@ import {
   getAntigravityOAuthUserAgent,
 } from "../services/antigravityHeaders.ts";
 import { classify429, decide429, type Decision } from "../services/antigravity429Engine.ts";
+import { lockExactModel } from "../services/accountFallback.ts";
 import {
   shouldRetryWithCredits,
   shouldUseCreditsFirst,
@@ -1424,6 +1425,7 @@ export class AntigravityExecutor extends BaseExecutor {
     const {
       response,
       url,
+      model,
       headers,
       transformedBody,
       credentials,
@@ -1443,10 +1445,9 @@ export class AntigravityExecutor extends BaseExecutor {
       // 1. Try to parse explicit retry time from message
       const parsedRetryMs = this.parseRetryFromErrorMessage(errorMessage);
 
-      // 2. Classify 429, then decide the final retry time BEFORE the credits
-      //    retry so that full_quota_exhausted can skip the credits attempt
-      //    entirely (avoids ~41s hold on an already-exhausted account) and
-      //    persist the cooldown to DB for post-restart routing.
+      // 2. Classify 429, then decide the final retry time BEFORE the credits retry so
+      //    full_quota_exhausted can skip the credits attempt entirely (avoids ~41s hold
+      //    on an already-exhausted account) and locks only this exact model.
       const category = classify429(errorMessage);
       const decision: Decision = decide429(category, parsedRetryMs);
       const retryMs = decision.retryAfterMs;
@@ -1460,10 +1461,9 @@ export class AntigravityExecutor extends BaseExecutor {
         !creditsRetryState.attempted &&
         shouldRetryWithCredits(credentials?.accessToken || "", creditsMode);
 
-      // Retry mode gets one credits attempt before the account cooldown is persisted.
-      // All other full-quota paths fail closed immediately.
+      // Retry mode gets one credits attempt before the exact-model lock is persisted.
       if (decision.kind === "full_quota_exhausted" && retryMs && !creditsRetryEligible) {
-        markConnectionQuotaExhausted(accountId, retryMs);
+        lockExactModel(this.provider, accountId, model, "quota_exhausted", retryMs);
       }
 
       if (category === "quota_exhausted" && creditsAlreadyInjected) {

open-sse/handlers/chatCore.ts

@@ -296,6 +296,7 @@ import {
   markBlocked as markAccountSemaphoreBlocked,
 } from "../services/accountSemaphore.ts";
 import { lockModel, lockModelIfPerModelQuota } from "../services/accountFallback.ts";
+import { lockExactModel } from "../services/accountFallback.ts";
 import {
   generateSignature,
   getCachedResponse,
@@ -3704,7 +3705,8 @@ export async function handleChatCore({
             markAccountSemaphoreBlocked(accountSemaphoreKey, quotaCooldownMs);
           }
           if (isModelScope() && errorConnectionId) {
-            lockModel(provider, errorConnectionId, model, "quota_exhausted", quotaCooldownMs);
+            const lockFn = provider === "antigravity" ? lockExactModel : lockModel;
+            lockFn(provider, errorConnectionId, model, "quota_exhausted", quotaCooldownMs);
             console.warn(
               `[provider] Node ${errorConnectionId} ModelScope model quota exhausted (${statusCode}) for ${model} - ${Math.ceil(quotaCooldownMs / 1000)}s (connection stays active)`
             );

open-sse/services/__tests__/antigravity-quota-family.test.ts

@@ -7,7 +7,10 @@ import {
   clearAllModelLockouts,
   getModelLockoutInfo,
   isModelLocked,
+  lockModelIfPerModelQuota,
+  lockExactModel,
   recordModelLockoutFailure,
+  clearModelLock,
 } from "@omniroute/open-sse/services/accountFallback.ts";
 
 const provider = "antigravity";
@@ -74,6 +77,38 @@ describe("Antigravity account quota-family cooldown", () => {
     expect(isModelLocked(provider, "account-a", "gemini-3.5-flash-low")).toBe(false);
   });
 
+  it("can isolate a confirmed Antigravity quota exhaustion to one exact model", () => {
+    lockExactModel(
+      provider,
+      "account-a",
+      "claude-opus-4-6-thinking",
+      "quota_exhausted",
+      60_000
+    );
+
+    expect(isModelLocked(provider, "account-a", "claude-opus-4-6-thinking")).toBe(true);
+    expect(isModelLocked(provider, "account-a", "claude-sonnet-4-6-thinking")).toBe(false);
+    expect(isModelLocked(provider, "account-a", "gemini-3.5-flash-medium")).toBe(false);
+
+    expect(clearModelLock(provider, "account-a", "claude-opus-4-6-thinking")).toBe(true);
+    expect(isModelLocked(provider, "account-a", "claude-opus-4-6-thinking")).toBe(false);
+  });
+
+  it("uses an exact model lock for Antigravity in the generic per-model quota path", () => {
+    expect(
+      lockModelIfPerModelQuota(
+        provider,
+        "account-a",
+        "claude-opus-4-6-thinking",
+        "quota_exhausted",
+        60_000
+      )
+    ).toBe(true);
+
+    expect(isModelLocked(provider, "account-a", "claude-opus-4-6-thinking")).toBe(true);
+    expect(isModelLocked(provider, "account-a", "claude-sonnet-4-6-thinking")).toBe(false);
+  });
+
   it("honors exact upstream cooldowns and otherwise uses bounded inferred cooldown", () => {
     const upstream = recordModelLockoutFailure(
       provider,

open-sse/services/accountFallback.ts

@@ -58,6 +58,7 @@ import { evictLockoutOverflow } from "./accountFallback/lockoutEviction.ts";
 export { MODEL_LOCKOUT_EVICTION_CAP } from "./accountFallback/lockoutEviction.ts";
 import { capScaledCooldownMs } from "./accountFallback/cooldownCap.ts";
 import { resolveApiKeyForbiddenFallback } from "./accountFallback/nonRetryableUpstream.ts";
+import * as exactModelLock from "./accountFallback/exactModelLock.ts";
 export type ProviderProfile = {
   baseCooldownMs: number;
   useUpstreamRetryHints: boolean;
@@ -442,6 +443,12 @@ function getModelLockKey(
   return `${canonicalProvider}:${connectionId}:${lockModel}`;
 }
 
+const buildExactKey = exactModelLock.buildExactModelLockKey; // see exactModelLock.ts
+const getModelLockKeys = exactModelLock.createGetModelLockKeys(
+  getModelLockKey,
+  getCanonicalLockProvider
+);
+
 function getFailureWindowMs(profile: ProviderProfile | null = null, fallbackMs = 30 * 60 * 1000) {
   const configured = profile?.resetTimeoutMs;
   return typeof configured === "number" && configured > 0 ? configured : fallbackMs;
@@ -559,6 +566,14 @@ export function lockModel(
   });
 }
 
+// Lock only this exact provider/account/model tuple, never a quota family — see exactModelLock.ts.
+export const lockExactModel = exactModelLock.createLockExactModel(
+  modelLockouts,
+  ensureCleanupTimer,
+  cleanupModelLockKey,
+  getCanonicalLockProvider
+);
+
 /**
  * Pick the `exactCooldownMs` to apply to a model lockout (#1308).
  *
@@ -591,6 +606,7 @@ export function recordModelLockoutFailure(
   options: {
     exactCooldownMs?: number | null;
     maxCooldownMs?: number;
+    scope?: "exact" | "quota_family";
     /**
      * #6863 vs #7940: set true only when `exactCooldownMs` came from an actual
      * upstream signal (Retry-After header, X-RateLimit-Reset, or a reset parsed
@@ -606,7 +622,10 @@ export function recordModelLockoutFailure(
   } = {}
 ) {
   ensureCleanupTimer();
-  const key = getModelLockKey(provider, connectionId, model, reason, status);
+  const key =
+    options.scope === "exact"
+      ? buildExactKey(getCanonicalLockProvider(provider), connectionId, model)
+      : getModelLockKey(provider, connectionId, model, reason, status);
   const now = Date.now();
   cleanupModelLockKey(key, now);
 
@@ -656,7 +675,8 @@ export function recordModelLockoutFailure(
     lastCooldownMs: cooldownMs,
   });
 
-  lockModel(provider, connectionId, model, reason, cooldownMs, {
+  const lockFn = options.scope === "exact" ? lockExactModel : lockModel;
+  lockFn(provider, connectionId, model, reason, cooldownMs, {
     failureCount,
     lastFailureAt: now,
     resetAfterMs,
@@ -675,16 +695,11 @@ export function clearModelLock(
   model: string | null | undefined
 ): boolean {
   if (!model) return false;
-  const familyKey = getModelLockKey(provider, connectionId, model);
-  const exactKey = `${getCanonicalLockProvider(provider)}:${connectionId}:${model}`;
-
-  const hadLock1 = modelLockouts.delete(familyKey);
-  const hadFailure1 = modelFailureState.delete(familyKey);
-
-  const hadLock2 = modelLockouts.delete(exactKey);
-  const hadFailure2 = modelFailureState.delete(exactKey);
-
-  return hadLock1 || hadFailure1 || hadLock2 || hadFailure2;
+  return exactModelLock.clearMultiKeyLock(
+    modelLockouts,
+    modelFailureState,
+    getModelLockKeys(provider, connectionId, model)
+  );
 }
 
 /**
@@ -708,6 +723,7 @@ export function hasPerModelQuota(
     return connectionPassthroughModels;
   }
   if (!provider) return false;
+  if (getCanonicalLockProvider(provider) === "antigravity") return true;
   if (getCanonicalLockProvider(provider) === "codex") return true;
   if (provider === "gemini" || provider === "github") return true;
   if (getPassthroughProviders().has(provider)) return true;
@@ -731,7 +747,8 @@ export function lockModelIfPerModelQuota(
   // Skip model-level lock if the entire provider is in circuit-breaker cooldown.
   // The provider cooldown already prevents all requests, so a model lock is redundant.
   if (isProviderInCooldown(provider)) return false;
-  lockModel(provider, connectionId, model, reason, cooldownMs);
+  const lockFn = getCanonicalLockProvider(provider) === "antigravity" ? lockExactModel : lockModel;
+  lockFn(provider, connectionId, model, reason, cooldownMs);
   return true;
 }
 
@@ -800,14 +817,11 @@ export function isModelLocked(
   model: string | null | undefined
 ): boolean {
   if (!model) return false;
-
-  const exactKey = `${getCanonicalLockProvider(provider)}:${connectionId}:${model}`;
-  cleanupModelLockKey(exactKey);
-  if (modelLockouts.has(exactKey)) return true;
-
-  const familyKey = getModelLockKey(provider, connectionId, model);
-  cleanupModelLockKey(familyKey);
-  return modelLockouts.has(familyKey);
+  return exactModelLock.isAnyKeyLocked(
+    modelLockouts,
+    cleanupModelLockKey,
+    getModelLockKeys(provider, connectionId, model)
+  );
 }
 
 /**
@@ -819,32 +833,18 @@ export function getModelLockoutInfo(
   model: string | null | undefined
 ) {
   if (!model) return null;
-
-  const exactKey = `${getCanonicalLockProvider(provider)}:${connectionId}:${model}`;
-  cleanupModelLockKey(exactKey);
-  const exactEntry = modelLockouts.get(exactKey);
-  if (exactEntry) {
-    return {
-      reason: exactEntry.reason,
-      remainingMs: exactEntry.until - Date.now(),
-      lockedAt: new Date(exactEntry.lockedAt).toISOString(),
-      failureCount: exactEntry.failureCount,
-    };
-  }
-
-  const familyKey = getModelLockKey(provider, connectionId, model);
-  cleanupModelLockKey(familyKey);
-  const familyEntry = modelLockouts.get(familyKey);
-  if (familyEntry) {
-    return {
-      reason: familyEntry.reason,
-      remainingMs: familyEntry.until - Date.now(),
-      lockedAt: new Date(familyEntry.lockedAt).toISOString(),
-      failureCount: familyEntry.failureCount,
-    };
-  }
-
-  return null;
+  const entry = exactModelLock.findLatestLockEntry(
+    modelLockouts,
+    cleanupModelLockKey,
+    getModelLockKeys(provider, connectionId, model)
+  );
+  if (!entry) return null;
+  return {
+    reason: entry.reason,
+    remainingMs: entry.until - Date.now(),
+    lockedAt: new Date(entry.lockedAt).toISOString(),
+    failureCount: entry.failureCount,
+  };
 }
 
 export type ModelLockoutInfo = {