diff --git a/open-sse/executors/antigravity.ts b/open-sse/executors/antigravity.ts index 1188b1116d..2e70df8cc7 100644 --- a/open-sse/executors/antigravity.ts +++ b/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) { diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index b74ec33890..90191b49ad 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -303,6 +303,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, @@ -3698,7 +3699,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)` ); diff --git a/open-sse/services/__tests__/antigravity-quota-family.test.ts b/open-sse/services/__tests__/antigravity-quota-family.test.ts index c23d79926b..b1983a0909 100644 --- a/open-sse/services/__tests__/antigravity-quota-family.test.ts +++ b/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, diff --git a/open-sse/services/accountFallback.ts b/open-sse/services/accountFallback.ts index 5204c7410d..a09b78924d 100644 --- a/open-sse/services/accountFallback.ts +++ b/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 = { diff --git a/open-sse/services/accountFallback/exactModelLock.ts b/open-sse/services/accountFallback/exactModelLock.ts new file mode 100644 index 0000000000..838d6a9173 --- /dev/null +++ b/open-sse/services/accountFallback/exactModelLock.ts @@ -0,0 +1,158 @@ +/** + * accountFallback/exactModelLock.ts — exact-model (non-family-scoped) lockout key + entry math. + * + * Extracted from services/accountFallback.ts (file-size gate, #8630): pure helpers for the + * opt-in "exact model" lockout scope introduced for Antigravity — a confirmed exhaustion on + * one specific model (e.g. one Claude model) must not lock the whole quota family (Gemini or + * other Claude models on the same account). Pure w.r.t. module state — accountFallback.ts + * still owns the modelLockouts/modelFailureState maps, canonical-provider resolution, and the + * cleanup timer; it calls into these with its own map instances. + */ + +import type { ModelLockoutEntry, ModelFailureState } from "../accountFallback.ts"; + +/** Build the "exact" scoped lockout key — a distinct namespace from the quota-family key. */ +export function buildExactModelLockKey( + canonicalProvider: string, + connectionId: string, + model: string +): string { + return `${canonicalProvider}:${connectionId}:exact:${model.trim().toLowerCase()}`; +} + +/** Dedupe the 3 lockout key shapes callers must check: quota-family, #8050 not_found, exact. */ +export function collectModelLockKeys( + familyKey: string, + notFoundKey: string, + exactKey: string +): string[] { + return Array.from(new Set([familyKey, notFoundKey, exactKey])); +} + +/** + * DI factory for `getModelLockKeys` — accountFallback.ts's own `getModelLockKey` (quota-family + * scoping) and `getCanonicalLockProvider` (alias resolution) are private, so this closes over + * them here rather than duplicating that logic in the leaf. + */ +export function createGetModelLockKeys( + getModelLockKey: ( + provider: string, + connectionId: string, + model: string, + reason?: string | null, + status?: number | null + ) => string, + getCanonicalLockProvider: (provider: string) => string +) { + return function getModelLockKeys(provider: string, connectionId: string, model: string) { + return collectModelLockKeys( + getModelLockKey(provider, connectionId, model), + getModelLockKey(provider, connectionId, model, "not_found", 404), + buildExactModelLockKey(getCanonicalLockProvider(provider), connectionId, model) + ); + }; +} + +/** + * Compute the next ModelLockoutEntry for an exact-model lock, merging with any existing entry + * the same way lockModel() does (extend failureCount on a shorter re-lock instead of shrinking + * the remaining cooldown). Returns null when the caller should leave state untouched. + */ +export function computeExactModelLockEntry( + existing: ModelLockoutEntry | undefined, + reason: string, + cooldownMs: number, + metadata: Partial +): ModelLockoutEntry | null { + const now = Date.now(); + const newUntil = now + cooldownMs; + if (existing && existing.until > newUntil) { + if (!metadata.failureCount || metadata.failureCount <= existing.failureCount) return null; + return { + ...existing, + failureCount: metadata.failureCount, + lastFailureAt: metadata.lastFailureAt ?? existing.lastFailureAt, + resetAfterMs: metadata.resetAfterMs ?? existing.resetAfterMs, + }; + } + return { + reason, + until: newUntil, + lockedAt: now, + failureCount: metadata.failureCount ?? existing?.failureCount ?? 1, + lastFailureAt: metadata.lastFailureAt ?? now, + resetAfterMs: metadata.resetAfterMs ?? existing?.resetAfterMs ?? 0, + }; +} + +/** + * Delete every one of the 3 lockout key shapes from both maps — a success on any one + * of them must clear the lock regardless of which reason originally wrote it. + */ +export function clearMultiKeyLock( + modelLockouts: Map, + modelFailureState: Map, + keys: string[] +): boolean { + let cleared = false; + for (const key of keys) { + cleared = modelLockouts.delete(key) || cleared; + cleared = modelFailureState.delete(key) || cleared; + } + return cleared; +} + +/** True when any of the 3 lockout key shapes is currently active (post-cleanup). */ +export function isAnyKeyLocked( + modelLockouts: Map, + cleanup: (key: string) => void, + keys: string[] +): boolean { + return keys.some((key) => { + cleanup(key); + return modelLockouts.has(key); + }); +} + +/** The active entry with the most remaining time across the 3 lockout key shapes. */ +export function findLatestLockEntry( + modelLockouts: Map, + cleanup: (key: string) => void, + keys: string[] +): ModelLockoutEntry | undefined { + return keys + .map((key) => { + cleanup(key); + return modelLockouts.get(key); + }) + .filter((value): value is ModelLockoutEntry => Boolean(value)) + .sort((a, b) => b.until - a.until)[0]; +} + +/** + * DI factory for the exported `lockExactModel` — accountFallback.ts owns the + * modelLockouts map + cleanup timer/key private functions and closes over them here so + * the full lock-only-this-exact-tuple implementation lives in this leaf, not the god-file. + */ +export function createLockExactModel( + modelLockouts: Map, + ensureCleanupTimer: () => void, + cleanupModelLockKey: (key: string) => void, + getCanonicalLockProvider: (provider: string) => string +) { + return function lockExactModel( + provider: string, + connectionId: string, + model: string | null | undefined, + reason: string, + cooldownMs: number, + metadata: Partial = {} + ): void { + if (!model) return; + ensureCleanupTimer(); + const key = buildExactModelLockKey(getCanonicalLockProvider(provider), connectionId, model); + cleanupModelLockKey(key); + const next = computeExactModelLockEntry(modelLockouts.get(key), reason, cooldownMs, metadata); + if (next) modelLockouts.set(key, next); + }; +} diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index 82067b2688..62885a72a8 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -1778,7 +1778,10 @@ async function handleSingleModelChat( result.status, 0, providerProfile, - { maxCooldownMs: mlSettings.maxCooldownMs } + { + maxCooldownMs: mlSettings.maxCooldownMs, + scope: provider === "antigravity" ? "exact" : undefined, + } ); log.info( diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index 372445bb35..64ea54edf1 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -2029,9 +2029,12 @@ export async function markAccountUnavailable( return { shouldFallback: true, cooldownMs: 0 }; } - const quotaScope = getQuotaScopeLabelForProvider(provider, model); + const usesExactAntigravityLock = provider === "antigravity"; + const quotaScope = usesExactAntigravityLock + ? "model" + : getQuotaScopeLabelForProvider(provider, model); const antigravityFamilyInferredBaseCooldownMs = - provider === "antigravity" && quotaScope === "family" && status === 429 + !usesExactAntigravityLock && provider === "antigravity" && quotaScope === "family" && status === 429 ? ANTIGRAVITY_FAMILY_INFERRED_BASE_COOLDOWN_MS : null; const lockout = recordModelLockoutFailure( @@ -2054,6 +2057,7 @@ export async function markAccountUnavailable( ? fallbackResult.cooldownMs : (fallbackResult.quotaResetHintMs ?? null), maxCooldownMs: mlSettings.maxCooldownMs, + scope: usesExactAntigravityLock ? "exact" : undefined, // #6863 vs #7940: exactCooldownMs above is only ever set from a genuine // upstream signal (Retry-After/reset header or a parsed quotaResetHintMs) — // never a synthetic estimate — so it must bypass maxCooldownMs instead of diff --git a/tests/unit/auth-antigravity-account-retry-v2.test.ts b/tests/unit/auth-antigravity-account-retry-v2.test.ts index e26a0625cc..59aa1e3e71 100644 --- a/tests/unit/auth-antigravity-account-retry-v2.test.ts +++ b/tests/unit/auth-antigravity-account-retry-v2.test.ts @@ -120,7 +120,7 @@ test("round-robin same-model retry treats multi-exclude as fallback LRU and skip assert.equal(selected.connectionId, staleId); }); -test("Antigravity inferred Gemini family cooldown starts around 30s when no upstream hint exists", async () => { +test("Antigravity 429 rate-limited locks only the exact model so siblings stay eligible", async () => { await resetStorage(); const conn = await providersDb.createProviderConnection({ @@ -131,34 +131,52 @@ test("Antigravity inferred Gemini family cooldown starts around 30s when no upst isActive: true, testStatus: "active", }); + const connId = connectionId(conn); - const before = Date.now(); const result = await auth.markAccountUnavailable( - connectionId(conn), + connId, 429, "RESOURCE_EXHAUSTED: Resource has been exhausted (queries per minute limit was reached)", "antigravity", "gemini-3-pro" ); - const elapsedAllowanceMs = Date.now() - before; + // The exact-only lock replaces the previous family inference: only the + // exhausted model is cooled down (short bounded cooldown), and sibling + // models on the same connection stay eligible. See PR #8630. assert.equal(result.shouldFallback, true); assert.ok( - result.cooldownMs >= 30_000 - elapsedAllowanceMs - 500, - `expected inferred cooldown near 30s+, got ${result.cooldownMs}` - ); - assert.ok( - result.cooldownMs <= 65_000, - `expected bounded initial cooldown, got ${result.cooldownMs}` + result.cooldownMs > 0 && result.cooldownMs <= 60_000, + `expected bounded cooldown, got ${result.cooldownMs}` ); - const otherGemini = await auth.getProviderCredentials( + // The exhausted model itself is locked: getProviderCredentials reports + // model-scope cooldown for that exact model on the only connection. + const sameModel = await auth.getProviderCredentials( + "antigravity", + null, + null, + "gemini-3-pro" + ); + assert.ok(sameModel); + assert.ok("allRateLimited" in sameModel && sameModel.allRateLimited); + assert.equal(sameModel.cooldownScope, "model"); + assert.equal(sameModel.cooldownModel, "gemini-3-pro"); + + // Sibling model on the SAME connection stays eligible — the whole point + // of the exact-model lock: a Claude/Gemini 429 must not disable unrelated + // models on the same account. + const siblingModel = await auth.getProviderCredentials( "antigravity", null, null, "gemini-2.5-pro" ); - assert.ok(otherGemini && "allRateLimited" in otherGemini && otherGemini.allRateLimited); - assert.equal(otherGemini.cooldownScope, "model"); - assert.equal(otherGemini.lastErrorCode, 429); + assert.ok(siblingModel && !("allRateLimited" in siblingModel && siblingModel.allRateLimited)); + assert.equal(siblingModel.connectionId, connId); + + // The exact lock clears when a successful request for the same model comes + // back through. clearModelLock is the existing success-path hook. + const { clearModelLock } = await import("../../open-sse/services/accountFallback.ts"); + assert.equal(clearModelLock("antigravity", connId, "gemini-3-pro"), true); });