diff --git a/src/sse/handlers/chat.js b/src/sse/handlers/chat.js index 48ae3578..7aa530d3 100644 --- a/src/sse/handlers/chat.js +++ b/src/sse/handlers/chat.js @@ -7,7 +7,7 @@ import { extractApiKey, isValidApiKey, } from "../services/auth.js"; -import { handleAntigravityQuotaError } from "../services/antigravityQuota.js"; +import { handleAntigravityQuotaError, clearAntigravityStrikes } from "../services/antigravityQuota.js"; import { getSettings } from "@/lib/localDb"; import { getModelInfo, getComboModels } from "../services/model.js"; import { handleChatCore } from "open-sse/handlers/chatCore.js"; @@ -302,6 +302,8 @@ async function handleSingleModelChat(body, modelStr, clientRawRequest = null, re }, onRequestSuccess: async () => { await clearAccountError(credentials.connectionId, credentials, model); + // "Consecutive" strikes: a success clears the breaker for this pair. + clearAntigravityStrikes(credentials.connectionId, model); } }); diff --git a/src/sse/services/antigravityQuota.js b/src/sse/services/antigravityQuota.js index e991614f..f27351f3 100644 --- a/src/sse/services/antigravityQuota.js +++ b/src/sse/services/antigravityQuota.js @@ -17,6 +17,56 @@ const inflightRefresh = new Map(); const MIN_REFRESH_INTERVAL_MS = 30_000; // 30s between refreshes per connection +// Strike-based circuit breaker (#3681): Google's quota API can report remaining +// quota while generation endpoints keep returning 429 (sprint/weekly dual-pool +// mismatch). After STRIKE_THRESHOLD 429s within the window for the same +// connection+model, treat the optimistic quota reading as untrusted and +// cache-block that pair instead of retry-storming upstream. +const STRIKE_WINDOW_MS = 60_000; // strikes older than this reset the count +const STRIKE_THRESHOLD = 3; +const STRIKE_BLOCK_MS = 15 * 60_000; +const strikeCounts = new Map(); // "connectionId|model" → { count, windowStart (anchored at first strike) } +const strikeBlocks = new Map(); // "connectionId|model" → blockedUntil ms + +/** + * Re-apply active strike blocks onto a fresh quotas snapshot so the auth + * pre-filter (which reads this cache) keeps skipping the blocked pair across + * requests until the block expires — same channel as the exhausted-0% path. + */ +function applyActiveStrikeBlocks(connectionId, quotas) { + const now = Date.now(); + for (const [key, until] of strikeBlocks) { + if (!key.startsWith(`${connectionId}|`)) continue; + if (until <= now) { + strikeBlocks.delete(key); + continue; + } + quotas[key.slice(connectionId.length + 1)] = { + remainingPercentage: 0, + resetAt: new Date(until).toISOString(), + }; + } + return quotas; +} + +/** + * Clear strike state for a connection|model after a successful request, so + * "consecutive" strikes means consecutive. Only removes a synthesized cache + * entry (resetAt == our block deadline); a real upstream 0% reading stays. + */ +export function clearAntigravityStrikes(connectionId, model) { + const key = `${connectionId}|${model}`; + strikeCounts.delete(key); + const until = strikeBlocks.get(key); + if (until === undefined) return; + strikeBlocks.delete(key); + const cached = quotaCache.get(connectionId); + if (cached?.[model]?.resetAt === new Date(until).toISOString()) { + delete cached[model]; + quotaCache.set(connectionId, cached); + } +} + /** * Get the quota cache (read-only reference for auth.js pre-filter). */ @@ -69,7 +119,9 @@ async function _doRefresh(connectionId, accessToken, providerSpecificData, now) if (!usage?.quotas || usage.message) return null; // Update in-memory cache. Caller logs CACHE_BLOCK only if requested model is exhausted. - quotaCache.set(connectionId, usage.quotas); + // Strike blocks are re-asserted after every refresh so an optimistic + // upstream reading cannot resurrect a pair we just circuit-broke. + quotaCache.set(connectionId, applyActiveStrikeBlocks(connectionId, usage.quotas)); return usage.quotas; } catch (e) { @@ -89,7 +141,43 @@ export async function handleAntigravityQuotaError(connectionId, status, model, a // Throttle applies to error paths too: one quota request per account/30s. // The first 409/429 populates cache; concurrent or repeated errors reuse it. const quota = (await refreshAntigravityQuota(connectionId, accessToken, providerSpecificData))?.[model]; - if (!quota || quota.remainingPercentage > 0 || !quota.resetAt) return null; + + // Strike breaker: count every 429 whose quota reading is either optimistic + // (remaining > 0) or unavailable (quota API 403/error). 3 within the window + // => the pair is unhealthy regardless of what the API claims; block 15m. + // 409 counts too by design: Antigravity signals pool exhaustion with 409 as + // well (see #3561 — "skip exhausted account/model quota before upstream + // retry" was motivated by 409/429 pairs), and poisoning by transient 409s + // requires 3 of them inside 60 seconds on the same pair. + if (!quota || quota.remainingPercentage > 0) { + const key = `${connectionId}|${model}`; + const now = Date.now(); + const strike = strikeCounts.get(key); + // Fixed window anchored at the FIRST qualifying strike: three 429s must + // all land within 60s of that first one, not within 60s of each other. + const windowStart = strike && now - strike.windowStart <= STRIKE_WINDOW_MS ? strike.windowStart : now; + const count = strike && windowStart === strike.windowStart ? strike.count + 1 : 1; + strikeCounts.set(key, { count, windowStart }); + if (count >= STRIKE_THRESHOLD) { + strikeCounts.delete(key); + const blockedUntil = now + STRIKE_BLOCK_MS; + const reading = quota ? `${Math.round(quota.remainingPercentage)}%` : "unknown"; + log.warn("AG_QUOTA", `${connectionId.slice(0, 8)} | STRIKE_${status} ${model} — ${count}x 429 (quota ${reading}); CACHE_BLOCK 15m`); + // Synthesize a 0% entry in the shared cache so the auth pre-filter skips + // this pair on subsequent requests too, not just the current retry loop + // (the chat handler does not persist modelLock_* for this path). + const cached = quotaCache.get(connectionId) || {}; + cached[model] = { remainingPercentage: 0, resetAt: new Date(blockedUntil).toISOString() }; + quotaCache.set(connectionId, cached); + strikeBlocks.set(key, blockedUntil); + return blockedUntil; + } + return null; + } + + // Healthy-but-exhausted reading: clear strikes and use the exact resetAt. + strikeCounts.delete(`${connectionId}|${model}`); + if (!quota.resetAt) return null; const resetMs = new Date(quota.resetAt).getTime(); if (resetMs <= Date.now()) return null; diff --git a/tests/unit/antigravity-quota-routing.test.js b/tests/unit/antigravity-quota-routing.test.js index c7778a3c..d8c060f6 100644 --- a/tests/unit/antigravity-quota-routing.test.js +++ b/tests/unit/antigravity-quota-routing.test.js @@ -27,7 +27,7 @@ vi.mock("open-sse/services/usage/google.js", () => ({ })); vi.mock("@/sse/utils/logger.js", () => ({ debug: vi.fn(), info: vi.fn(), warn: vi.fn() })); -const { getAntigravityQuotaCache, handleAntigravityQuotaError, refreshAntigravityQuota } = await import("@/sse/services/antigravityQuota.js"); +const { getAntigravityQuotaCache, handleAntigravityQuotaError, refreshAntigravityQuota, clearAntigravityStrikes } = await import("@/sse/services/antigravityQuota.js"); const { getProviderCredentials } = await import("@/sse/services/auth.js"); const MODEL = "claude-opus-4-6-thinking"; @@ -169,4 +169,143 @@ describe("Antigravity quota-aware routing", () => { vi.useRealTimers(); } }); + + it("strike-breaks after 3 optimistic 429s within 60s and cache-blocks 15 minutes", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-08-26T00:00:00.000Z")); + // Quota API lies: reports 90% remaining while generation keeps 429ing. + mocks.getAntigravityUsage.mockResolvedValue({ quotas: { + [MODEL]: { remainingPercentage: 90, resetAt: FUTURE_RESET }, + } }); + + try { + const first = await handleAntigravityQuotaError("ag-strike", 429, MODEL, "token", {}); + expect(first).toBeNull(); + const second = await handleAntigravityQuotaError("ag-strike", 429, MODEL, "token", {}); + expect(second).toBeNull(); + + const third = await handleAntigravityQuotaError("ag-strike", 429, MODEL, "token", {}); + expect(third).toBe(Date.parse("2026-08-26T00:15:00.000Z")); + } finally { + vi.useRealTimers(); + } + }); + + it("resets the strike counter when strikes fall outside the 60s window", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-08-26T00:00:00.000Z")); + mocks.getAntigravityUsage.mockResolvedValue({ quotas: { + [MODEL]: { remainingPercentage: 90, resetAt: FUTURE_RESET }, + } }); + + try { + await handleAntigravityQuotaError("ag-window", 429, MODEL, "token", {}); + await handleAntigravityQuotaError("ag-window", 429, MODEL, "token", {}); + await vi.advanceTimersByTimeAsync(61_000); + const result = await handleAntigravityQuotaError("ag-window", 429, MODEL, "token", {}); + expect(result).toBeNull(); // window lapsed — counter restarted at 1 + } finally { + vi.useRealTimers(); + } + }); + + it("strike-breaks when the quota API is unavailable (null reading) too", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-08-26T00:00:00.000Z")); + // Quota endpoint failing/forbidden => quota unknown. Strikes must still count. + mocks.getAntigravityUsage.mockResolvedValue({ message: "forbidden", quotas: {} }); + + try { + await handleAntigravityQuotaError("ag-null", 429, MODEL, "token", {}); + await handleAntigravityQuotaError("ag-null", 429, MODEL, "token", {}); + const third = await handleAntigravityQuotaError("ag-null", 429, MODEL, "token", {}); + expect(third).toBe(Date.parse("2026-08-26T00:15:00.000Z")); + } finally { + vi.useRealTimers(); + } + }); + + it("persists the block into the shared cache so the next request skips the pair", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-08-26T00:00:00.000Z")); + mocks.getAntigravityUsage.mockResolvedValue({ quotas: { + [MODEL]: { remainingPercentage: 90, resetAt: FUTURE_RESET }, + } }); + + try { + await handleAntigravityQuotaError("ag-persist", 429, MODEL, "token", {}); + await handleAntigravityQuotaError("ag-persist", 429, MODEL, "token", {}); + await handleAntigravityQuotaError("ag-persist", 429, MODEL, "token", {}); + + // The synthesized entry must be visible to the auth pre-filter reading + // the shared cache — and must survive an optimistic upstream refresh. + const cached = getAntigravityQuotaCache().get("ag-persist")?.[MODEL]; + expect(cached).toMatchObject({ remainingPercentage: 0 }); + expect(Date.parse(cached.resetAt)).toBe(Date.parse("2026-08-26T00:15:00.000Z")); + + await refreshAntigravityQuota("ag-persist", "token", {}); + expect(getAntigravityQuotaCache().get("ag-persist")?.[MODEL]).toMatchObject({ + remainingPercentage: 0, + }); + } finally { + vi.useRealTimers(); + } + }); + + it("clears strike state and the synthesized block after a successful request", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-08-26T00:00:00.000Z")); + mocks.getAntigravityUsage.mockResolvedValue({ quotas: { + [MODEL]: { remainingPercentage: 90, resetAt: FUTURE_RESET }, + } }); + + try { + await handleAntigravityQuotaError("ag-clear", 429, MODEL, "token", {}); + await handleAntigravityQuotaError("ag-clear", 429, MODEL, "token", {}); + await handleAntigravityQuotaError("ag-clear", 429, MODEL, "token", {}); + expect(getAntigravityQuotaCache().get("ag-clear")?.[MODEL]?.remainingPercentage).toBe(0); + + clearAntigravityStrikes("ag-clear", MODEL); + // Synthesized entry gone — pair selectable again immediately. + expect(getAntigravityQuotaCache().get("ag-clear")?.[MODEL]).toBeUndefined(); + + // Two more 429s do NOT inherit earlier strikes: no block on the third-in-episode. + await handleAntigravityQuotaError("ag-clear", 429, MODEL, "token", {}); + await expect(handleAntigravityQuotaError("ag-clear", 429, MODEL, "token", {})).resolves.toBeNull(); + } finally { + vi.useRealTimers(); + } + }); + + it("anchors the window at the first strike: 3 strikes spread over 90s do not trip", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-08-26T00:00:00.000Z")); + mocks.getAntigravityUsage.mockResolvedValue({ quotas: { + [MODEL]: { remainingPercentage: 90, resetAt: FUTURE_RESET }, + } }); + + try { + await handleAntigravityQuotaError("ag-anchor", 429, MODEL, "token", {}); // t=0 + await vi.advanceTimersByTimeAsync(45_000); + await handleAntigravityQuotaError("ag-anchor", 429, MODEL, "token", {}); // t=45s + await vi.advanceTimersByTimeAsync(45_000); + // t=90s: within 60s of strike #2 but outside 60s of strike #1 => new window + const result = await handleAntigravityQuotaError("ag-anchor", 429, MODEL, "token", {}); + expect(result).toBeNull(); + } finally { + vi.useRealTimers(); + } + }); + + it("keeps the optimistic path null without touching the quota cache", async () => { + mocks.getAntigravityUsage.mockResolvedValue({ quotas: { + [MODEL]: { remainingPercentage: 90, resetAt: FUTURE_RESET }, + } }); + + await expect(handleAntigravityQuotaError("ag-optimistic", 429, MODEL, "token", {})) + .resolves.toBeNull(); + // Optimistic reading must NOT poison the shared cache (auth pre-filter + // treats cached 0% as exhausted). + expect(getAntigravityQuotaCache().get("ag-optimistic")?.[MODEL]?.remainingPercentage).toBe(90); + }); });