diff --git a/open-sse/services/__tests__/antigravity-quota-family.test.ts b/open-sse/services/__tests__/antigravity-quota-family.test.ts index dbdcc69793..c23d79926b 100644 --- a/open-sse/services/__tests__/antigravity-quota-family.test.ts +++ b/open-sse/services/__tests__/antigravity-quota-family.test.ts @@ -21,15 +21,17 @@ describe("Antigravity account quota-family cooldown", () => { it("maps Gemini variants to Gemini family and Claude/Cloud variants to Claude family", () => { expect(getAntigravityQuotaFamily("gemini-3.5-flash-medium")).toBe("gemini"); expect(getAntigravityQuotaFamily("google/gemini-3.5-flash-low")).toBe("gemini"); + expect(getAntigravityQuotaFamily("agy/gemini-3.5-flash-medium")).toBe("gemini"); expect(getAntigravityQuotaFamily("claude-sonnet-4")).toBe("claude"); expect(getAntigravityQuotaFamily("cloud/claude-opus-4")).toBe("claude"); expect(getAntigravityQuotaFamily("some-new-model")).toBe("other"); }); it("uses family-scoped lock key for Antigravity but preserves exact-model scope elsewhere", () => { - expect(getQuotaScopedModelForProvider(provider, "gemini-3.5-flash-medium")).toBe( + expect(getQuotaScopedModelForProvider("antigravity", "gemini-3.5-flash-medium")).toBe( "family:gemini" ); + expect(getQuotaScopedModelForProvider("agy", "gemini-3.5-flash-medium")).toBe("family:gemini"); expect(getQuotaScopedModelForProvider(provider, "gemini-3.5-flash-low")).toBe("family:gemini"); expect(getQuotaScopedModelForProvider(provider, "claude-sonnet-4")).toBe("family:claude"); expect(getQuotaScopedModelForProvider(provider, "unknown-model")).toBe("unknown-model"); diff --git a/open-sse/services/antigravityQuotaFamily.ts b/open-sse/services/antigravityQuotaFamily.ts index e033291fa0..e9ede18749 100644 --- a/open-sse/services/antigravityQuotaFamily.ts +++ b/open-sse/services/antigravityQuotaFamily.ts @@ -18,7 +18,7 @@ function normalizeModelId(model: string | null | undefined): string { export function getAntigravityQuotaFamily( model: string | null | undefined ): AntigravityQuotaFamily { - const normalized = normalizeModelId(model).replace(/^antigravity\//, ""); + const normalized = normalizeModelId(model).replace(/^(antigravity|agy)\//, ""); const slashIndex = normalized.indexOf("/"); const bare = slashIndex >= 0 ? normalized.slice(slashIndex + 1) : normalized; @@ -42,7 +42,7 @@ export function getQuotaScopedModelForProvider( model: string | null | undefined ): string | null { if (!model) return null; - if (provider !== ANTIGRAVITY_PROVIDER_ID) return model; + if (provider !== "antigravity" && provider !== "agy") return model; const family = getAntigravityQuotaFamily(model); return family === "other" ? model : `family:${family}`; } @@ -51,6 +51,6 @@ export function getQuotaScopeLabelForProvider( provider: string | null | undefined, model: string | null | undefined ): string { - if (provider !== ANTIGRAVITY_PROVIDER_ID) return "model"; + if (provider !== "antigravity" && provider !== "agy") return "model"; return getAntigravityQuotaFamily(model) === "other" ? "model" : "family"; } diff --git a/open-sse/services/rateLimitManager.ts b/open-sse/services/rateLimitManager.ts index fe5bce878d..ab0e31f58f 100644 --- a/open-sse/services/rateLimitManager.ts +++ b/open-sse/services/rateLimitManager.ts @@ -10,6 +10,7 @@ import Bottleneck from "bottleneck"; import { parseRetryAfterFromBody } from "./accountFallback.ts"; +import { getAntigravityQuotaFamily } from "./antigravityQuotaFamily.ts"; import { getProviderCategory } from "../config/providerRegistry.ts"; import { getCodexRateLimitKey } from "../executors/codex.ts"; import { awaitProviderDefaultSlot, setProviderQuotaOverrides } from "./providerDefaultRateLimit.ts"; @@ -460,6 +461,11 @@ function getLimiterKey(provider, connectionId, model = null) { if (provider === "codex" && model) { return `${provider}:${getCodexRateLimitKey(connectionId, model)}`; } + if ((provider === "antigravity" || provider === "agy") && model) { + const family = getAntigravityQuotaFamily(model); + const scope = family === "other" ? model : family; + return `${provider}:${connectionId}:${scope}`; + } // Gemini AI Studio and GitHub Copilot have per-model quotas — use model-scoped // limiter keys so a 429 on one model doesn't pause requests for other models. if ((provider === "gemini" || provider === "github") && model) { diff --git a/src/domain/quotaCache.ts b/src/domain/quotaCache.ts index 79062d53e2..75b037b986 100644 --- a/src/domain/quotaCache.ts +++ b/src/domain/quotaCache.ts @@ -24,6 +24,7 @@ import { } from "@/lib/db/quotaSnapshots"; import { recordProviderQuotaResetEventIfChanged } from "@/lib/db/quotaResetEvents"; import { getCodexQuotaWindowFilterForModel } from "@omniroute/open-sse/config/codexQuotaScopes.ts"; +import { getAntigravityQuotaFamily } from "@omniroute/open-sse/services/antigravityQuotaFamily.ts"; // ─── Types ────────────────────────────────────────────────────────────────── @@ -205,16 +206,38 @@ export function __clearForTests() { cache.clear(); } -export function isQuotaExhaustedForRequest( +function isAntigravityQuotaExhausted( connectionId: string, - provider: string, - requestedModel: string | null = null + entry: QuotaCacheEntry, + requestedModel: string | null ): boolean { - if (!isAccountQuotaExhausted(connectionId)) return false; - if (provider !== "codex" || !requestedModel) return true; - const entry = getQuotaCache(connectionId); - const quotaNames = Object.keys(entry?.quotas || {}); - if (quotaNames.length === 0) return true; + if (!requestedModel) return entry.exhausted; + const quotaNames = Object.keys(entry.quotas || {}); + if (quotaNames.length === 0) return entry.exhausted; + const requestedFamily = getAntigravityQuotaFamily(requestedModel); + const cleanRequestedModel = requestedModel.replace(/^(antigravity|agy)\//, ""); + const matchingWindows = quotaNames.filter((windowName) => { + if (requestedFamily === "other") { + return windowName.replace(/^(antigravity|agy)\//, "") === cleanRequestedModel; + } + return getAntigravityQuotaFamily(windowName) === requestedFamily; + }); + return ( + matchingWindows.length > 0 && + matchingWindows.every( + (windowName) => getQuotaWindowStatus(connectionId, windowName, 100)?.reachedThreshold + ) + ); +} + +function isCodexQuotaExhausted( + connectionId: string, + entry: QuotaCacheEntry, + requestedModel: string | null +): boolean { + if (!requestedModel) return entry.exhausted; + const quotaNames = Object.keys(entry.quotas || {}); + if (quotaNames.length === 0) return entry.exhausted; const filterWindow = getCodexQuotaWindowFilterForModel(requestedModel); const scopedWindowNames = quotaNames.filter((windowName) => filterWindow?.(windowName)); return ( @@ -225,6 +248,40 @@ export function isQuotaExhaustedForRequest( ); } +function isStandardQuotaExhausted(entry: QuotaCacheEntry, now: number): boolean { + if (!entry.exhausted) return false; + const age = now - entry.fetchedAt; + if (!entry.nextResetAt && age > EXHAUSTED_TTL_MS) return false; + return true; +} + +export function isQuotaExhaustedForRequest( + connectionId: string, + provider: string, + requestedModel: string | null = null +): boolean { + const entry = cache.get(connectionId) || hydrateQuotaCacheFromSnapshots(connectionId); + if (!entry) return false; + + const now = Date.now(); + const advanced = advancedWindowResetAt(entry, now); + if (advanced) { + entry.exhausted = false; + return false; + } + + if (provider === "antigravity" || provider === "agy") { + return isAntigravityQuotaExhausted(connectionId, entry, requestedModel); + } + + if (provider === "codex") { + return isCodexQuotaExhausted(connectionId, entry, requestedModel); + } + + // Standard (non-per-model-quota) providers: check connection-wide aggregate + return isStandardQuotaExhausted(entry, now); +} + /** * Store quota data for a connection (called by usage endpoint and background refresh). */ diff --git a/tests/unit/antigravity-quota-skipping.test.ts b/tests/unit/antigravity-quota-skipping.test.ts new file mode 100644 index 0000000000..a9b1f7b79b --- /dev/null +++ b/tests/unit/antigravity-quota-skipping.test.ts @@ -0,0 +1,110 @@ +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(), "omni-quota-antigravity-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const coreDb = await import("../../src/lib/db/core.ts"); +const quotaCache = await import("../../src/domain/quotaCache.ts"); + +test.after(() => { + coreDb.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("isQuotaExhaustedForRequest isolates Claude and Gemini quota families for antigravity & agy", () => { + const connectionId = "conn-antigravity-test"; + + // Simulate Claude Opus being exhausted, while Gemini is NOT. + quotaCache.setQuotaCache(connectionId, "antigravity", { + "claude-opus-4-6-thinking": { remainingPercentage: 0, resetAt: null }, + "gemini-3.5-flash-high": { remainingPercentage: 100, resetAt: null }, + }); + + // Verify that Claude models are considered exhausted. + assert.equal( + quotaCache.isQuotaExhaustedForRequest( + connectionId, + "antigravity", + "antigravity/claude-opus-4-6-thinking" + ), + true, + "Claude Opus should be exhausted" + ); + assert.equal( + quotaCache.isQuotaExhaustedForRequest( + connectionId, + "antigravity", + "antigravity/claude-sonnet-4-6" + ), + true, + "Claude Sonnet should share Claude family quota and be exhausted" + ); + + // Verify that Gemini models are NOT considered exhausted. + assert.equal( + quotaCache.isQuotaExhaustedForRequest( + connectionId, + "antigravity", + "antigravity/gemini-3.5-flash-high" + ), + false, + "Gemini Flash should NOT be exhausted" + ); + assert.equal( + quotaCache.isQuotaExhaustedForRequest( + connectionId, + "antigravity", + "antigravity/gemini-2.5-pro" + ), + false, + "Gemini Pro should share Gemini family quota and NOT be exhausted" + ); + + // Test that 'agy' spelling behaves the exact same way. + const connectionIdAgy = "conn-agy-test"; + quotaCache.setQuotaCache(connectionIdAgy, "agy", { + "claude-opus-4-6-thinking": { remainingPercentage: 0, resetAt: null }, + "gemini-3.5-flash-high": { remainingPercentage: 100, resetAt: null }, + }); + + assert.equal( + quotaCache.isQuotaExhaustedForRequest(connectionIdAgy, "agy", "agy/claude-opus-4-6-thinking"), + true, + "Claude Opus under 'agy' should be exhausted" + ); + assert.equal( + quotaCache.isQuotaExhaustedForRequest(connectionIdAgy, "agy", "agy/gemini-3.5-flash-high"), + false, + "Gemini Flash under 'agy' should NOT be exhausted" + ); + + // Test that unknown models (family 'other') preserve exact-model scoping. + const connectionIdOther = "conn-other-test"; + quotaCache.setQuotaCache(connectionIdOther, "antigravity", { + "unknown-model-a": { remainingPercentage: 0, resetAt: null }, + "unknown-model-b": { remainingPercentage: 100, resetAt: null }, + }); + + assert.equal( + quotaCache.isQuotaExhaustedForRequest( + connectionIdOther, + "antigravity", + "antigravity/unknown-model-a" + ), + true, + "Unknown model A should be exhausted" + ); + assert.equal( + quotaCache.isQuotaExhaustedForRequest( + connectionIdOther, + "antigravity", + "antigravity/unknown-model-b" + ), + false, + "Unknown model B should NOT be exhausted" + ); +});