From d838be8df8f2a528eb9efb79dff0a7794bc2b52e Mon Sep 17 00:00:00 2001 From: Diego Rodrigues de Sa e Souza <8016841+diegosouzapw@users.noreply.github.com> Date: Sat, 11 Jul 2026 07:39:24 -0300 Subject: [PATCH] feat(usage): surface Antigravity weekly quota alongside the 5-hour window (#4017) (#6818) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(usage): surface Antigravity weekly quota alongside the 5-hour window (#4017) Antigravity enforces both a 5-hour and a weekly usage limit, but the agy/antigravity quota widget only exposed the 5-hour window. The weekly limit isn't in the per-model retrieveUserQuota response already fetched — it lives in a separate, undocumented retrieveUserQuotaSummary RPC that groups models into families (Gemini Models, Claude and GPT models) with one weekly bucket per family. Adds a self-contained usage/antigravityWeeklyQuota.ts leaf: a cached, best-effort fetch of that RPC + a pure parser that extracts the weekly-labeled bucket per group (window inferred from bucketId/displayName text, matching the reverse-engineered shape documented by third-party Antigravity clients) into gemini_weekly/ claude_gpt_weekly quota entries, merged into the existing quotas map the widget already renders generically. A failed/unavailable RPC never affects the existing per-model quotas. Live VPS validation attempt (192.168.0.15, real antigravity account): both retrieveUserQuota and retrieveUserQuotaSummary currently return 429 RESOURCE_EXHAUSTED for that account, so the live response shape could not be captured directly. The parser was instead validated via TDD against the bucket shape documented by CodexBar (steipete/CodexBar), a third-party Antigravity client that reverse-engineered the same RPC, and is defensive against both response envelopes it has observed (top-level groups[] and nested quotaSummary.groups[]). * chore(merge): re-sync with release/v3.8.47 (restore CHANGELOG, keep own bullet) * chore(merge): re-sync with release/v3.8.47; move changelog bullet to changelog.d fragment (merge-storm proof) --- .../features/6818-antigravity-weekly-quota.md | 1 + open-sse/services/usage/antigravity.ts | 5 +- .../services/usage/antigravityWeeklyQuota.ts | 184 ++++++++++++++ .../usage/components/ProviderLimits/utils.tsx | 2 + .../antigravity-weekly-quota-4017.test.ts | 240 ++++++++++++++++++ 5 files changed, 431 insertions(+), 1 deletion(-) create mode 100644 changelog.d/features/6818-antigravity-weekly-quota.md create mode 100644 open-sse/services/usage/antigravityWeeklyQuota.ts create mode 100644 tests/unit/antigravity-weekly-quota-4017.test.ts diff --git a/changelog.d/features/6818-antigravity-weekly-quota.md b/changelog.d/features/6818-antigravity-weekly-quota.md new file mode 100644 index 0000000000..84abe02d9d --- /dev/null +++ b/changelog.d/features/6818-antigravity-weekly-quota.md @@ -0,0 +1 @@ +- **feat(usage):** Antigravity/agy quota widget now surfaces the **weekly** window alongside the existing per-model 5-hour window ([#4017](https://github.com/diegosouzapw/OmniRoute/issues/4017)) — the weekly limit isn't part of the per-model `retrieveUserQuota` response the fetcher already calls; it only appears in a separate, undocumented `retrieveUserQuotaSummary` RPC that groups models into families ("Gemini Models", "Claude and GPT models") with one weekly bucket per family. A new `usage/antigravityWeeklyQuota.ts` leaf fetches that RPC (cached, best-effort — a failure or unavailable RPC never breaks the existing per-model quotas) and parses the weekly bucket per group into `gemini_weekly`/`claude_gpt_weekly` quota entries, merged into the same `quotas` map the widget already renders generically. Regression guard: `tests/unit/antigravity-weekly-quota-4017.test.ts` (bucket parsing, the alternate `quotaSummary`-nested envelope, and end-to-end merge via `getUsageForProvider`). diff --git a/open-sse/services/usage/antigravity.ts b/open-sse/services/usage/antigravity.ts index 99f246155c..84f01420e0 100644 --- a/open-sse/services/usage/antigravity.ts +++ b/open-sse/services/usage/antigravity.ts @@ -43,6 +43,7 @@ import { } from "../codeAssistSubscription.ts"; import { toRecord, toNumber, getFieldValue } from "./scalars.ts"; import { type UsageQuota, parseResetTime } from "./quota.ts"; +import { fetchAndParseAntigravityWeeklyQuotas } from "./antigravityWeeklyQuota.ts"; type JsonRecord = Record; type SubscriptionCacheEntry = { @@ -604,9 +605,10 @@ export async function getAntigravityUsage( ); } - const [data, userQuotaData] = await Promise.all([ + const [data, userQuotaData, weeklyQuotas] = await Promise.all([ fetchAntigravityAvailableModelsCached(accessToken, projectId, options), fetchAntigravityUserQuotaCached(accessToken, projectId, options), + fetchAndParseAntigravityWeeklyQuotas(accessToken, projectId, options), // #4017 ]); const dataObj = toRecord(data); if (dataObj.__antigravityForbidden === true) { @@ -717,6 +719,7 @@ export async function getAntigravityUsage( plan: getAntigravityPlanLabel(subscriptionInfo, providerSpecificData), quotas: { ...quotas, + ...weeklyQuotas, ...(creditBalance !== null && { credits: { used: 0, diff --git a/open-sse/services/usage/antigravityWeeklyQuota.ts b/open-sse/services/usage/antigravityWeeklyQuota.ts new file mode 100644 index 0000000000..043e0872ed --- /dev/null +++ b/open-sse/services/usage/antigravityWeeklyQuota.ts @@ -0,0 +1,184 @@ +/** + * usage/antigravityWeeklyQuota.ts — Antigravity weekly-quota fetcher + parser (#4017). + * + * Antigravity enforces both a 5-hour window (already surfaced per-model by + * `getAntigravityUsage()` via `retrieveUserQuota`) and a separate weekly window. + * The weekly window is NOT part of the per-model `retrieveUserQuota` response — + * it lives in a distinct upstream RPC, `v1internal:retrieveUserQuotaSummary`, + * which groups models into families ("Gemini Models", "Claude and GPT models") + * and reports one bucket per family per window (5h + weekly), keyed by a + * `bucketId`/`displayName` pair rather than by individual modelId. There is no + * dedicated window-type field on the bucket — the window is inferred from the + * bucketId/displayName text (matches the reverse-engineered shape documented by + * third-party Antigravity clients, since Google does not publish this API). + * + * This module is a small, self-contained leaf so `usage/antigravity.ts` stays a + * thin caller: fetch (cached, best-effort) + pure parse, mirroring the existing + * `fetchAntigravityUserQuotaCached` pattern. + */ + +import { toRecord, toNumber } from "./scalars.ts"; +import { type UsageQuota, parseResetTime } from "./quota.ts"; + +type JsonRecord = Record; + +interface AntigravityWeeklyQuotaOptions { + forceRefresh?: boolean; +} + +const WEEKLY_QUOTA_CACHE_TTL_MS = 60 * 1000; +const _weeklyQuotaCache = new Map(); +const _weeklyQuotaInflight = new Map>(); + +// Self-contained purge timer — this leaf owns its own cache, so it owns the cleanup too +// (same pattern as usage/antigravity.ts's module-level caches). +const _weeklyQuotaCacheCleanupTimer = setInterval( + () => { + const now = Date.now(); + for (const [key, entry] of _weeklyQuotaCache) { + if (now - entry.fetchedAt > WEEKLY_QUOTA_CACHE_TTL_MS) _weeklyQuotaCache.delete(key); + } + }, + 5 * 60 * 1000 +); +_weeklyQuotaCacheCleanupTimer.unref?.(); + +function buildCacheKey(accessToken: string, projectId?: string | null): string { + return `${accessToken.substring(0, 16)}:${projectId || "default"}`; +} + +/** + * Fetch the weekly-quota-bearing `retrieveUserQuotaSummary` response (cached, best-effort). + * Returns `null` on any failure — callers must treat this as optional data, never a hard + * dependency, since the RPC is undocumented and may not be available for every account/tier. + */ +export async function fetchAntigravityUserQuotaSummaryCached( + accessToken: string, + projectId?: string | null, + options: AntigravityWeeklyQuotaOptions = {} +): Promise { + if (!accessToken || !projectId) return null; + + const cacheKey = buildCacheKey(accessToken, projectId); + const cached = _weeklyQuotaCache.get(cacheKey); + if (!options.forceRefresh && cached && Date.now() - cached.fetchedAt < WEEKLY_QUOTA_CACHE_TTL_MS) { + return cached.data; + } + + const inflight = _weeklyQuotaInflight.get(cacheKey); + if (inflight) return inflight; + + const promise = (async () => { + try { + const response = await fetch( + "https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuotaSummary", + { + method: "POST", + headers: { + Authorization: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ project: projectId }), + signal: AbortSignal.timeout(10000), + } + ); + + if (!response.ok) return null; + + const data = await response.json(); + _weeklyQuotaCache.set(cacheKey, { data, fetchedAt: Date.now() }); + return data; + } catch { + return null; + } + })().finally(() => { + _weeklyQuotaInflight.delete(cacheKey); + }); + + _weeklyQuotaInflight.set(cacheKey, promise); + return promise; +} + +/** Matches a bucket's combined bucketId+displayName text against a window keyword. */ +function bucketMatchesWindow(bucket: JsonRecord, keyword: RegExp): boolean { + const text = `${String(bucket.bucketId || "")} ${String(bucket.displayName || "")}`.toLowerCase(); + return keyword.test(text); +} + +const WEEKLY_KEYWORD = /\bweekly\b/; + +/** Turns a group displayName (e.g. "Gemini Models", "Claude and GPT models") into a quota key. */ +function slugifyGroupWeeklyKey(displayName: string): string | null { + const cleaned = String(displayName || "") + .toLowerCase() + .replace(/\bmodels?\b/g, "") + .replace(/\band\b/g, " ") + .replace(/[^a-z0-9]+/g, "_") + .replace(/^_+|_+$/g, ""); + return cleaned ? `${cleaned}_weekly` : null; +} + +/** + * Parse the raw `retrieveUserQuotaSummary` response into weekly `UsageQuota` entries, + * one per model family group. Tolerant of the two response envelopes third-party + * Antigravity clients have observed (`groups[]` at the top level, or nested under + * `quotaSummary.groups[]`) since the RPC is undocumented and unversioned by Google. + */ +export function parseAntigravityWeeklyQuotas(summaryData: unknown): Record { + const root = toRecord(summaryData); + const rawGroups = Array.isArray(root.groups) + ? root.groups + : Array.isArray(toRecord(root.quotaSummary).groups) + ? (toRecord(root.quotaSummary).groups as unknown[]) + : []; + + const quotas: Record = {}; + + for (const groupValue of rawGroups) { + const group = toRecord(groupValue); + const buckets = Array.isArray(group.buckets) ? group.buckets : []; + const weeklyBucketValue = buckets.find( + (b) => b && typeof b === "object" && bucketMatchesWindow(toRecord(b), WEEKLY_KEYWORD) + ); + if (!weeklyBucketValue) continue; + + const weeklyBucket = toRecord(weeklyBucketValue); + if (weeklyBucket.disabled === true) continue; + + const key = slugifyGroupWeeklyKey(String(group.displayName || "")); + if (!key) continue; + + const rawFraction = toNumber(weeklyBucket.remainingFraction, -1); + if (rawFraction < 0) continue; + + const remainingFraction = Math.max(0, Math.min(1, rawFraction)); + const resetAt = parseResetTime(weeklyBucket.resetTime); + const isUnlimited = !resetAt && remainingFraction >= 1; + const QUOTA_NORMALIZED_BASE = 1000; + const total = QUOTA_NORMALIZED_BASE; + const remaining = Math.round(total * remainingFraction); + + quotas[key] = { + used: isUnlimited ? 0 : Math.max(0, total - remaining), + total: isUnlimited ? 0 : total, + resetAt, + remainingPercentage: isUnlimited ? 100 : remainingFraction * 100, + unlimited: isUnlimited, + fractionReported: true, + quotaSource: "retrieveUserQuota", + displayName: String(group.displayName || "").trim() || undefined, + }; + } + + return quotas; +} + +/** Fetch + parse in one call — the only entry point `usage/antigravity.ts` needs. */ +export async function fetchAndParseAntigravityWeeklyQuotas( + accessToken: string, + projectId: string | undefined | null, + options: AntigravityWeeklyQuotaOptions = {} +): Promise> { + const data = await fetchAntigravityUserQuotaSummaryCached(accessToken, projectId, options); + return parseAntigravityWeeklyQuotas(data); +} diff --git a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.tsx b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.tsx index 7526b87ead..77058cdc82 100644 --- a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.tsx +++ b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.tsx @@ -34,6 +34,8 @@ const QUOTA_LABEL_MAP: Record = { tokens: "Tokens", time_limit: "Time Limit", banked_reset_credits: "Banked Reset Credits", + gemini_weekly: "Gemini Weekly", + claude_gpt_weekly: "Claude & GPT Weekly", }; function toRecord(value: unknown): Record { diff --git a/tests/unit/antigravity-weekly-quota-4017.test.ts b/tests/unit/antigravity-weekly-quota-4017.test.ts new file mode 100644 index 0000000000..894246f1b9 --- /dev/null +++ b/tests/unit/antigravity-weekly-quota-4017.test.ts @@ -0,0 +1,240 @@ +/** + * #4017 — Antigravity weekly-quota widget. + * + * Antigravity enforces both a 5-hour window (already surfaced per-model via + * `retrieveUserQuota`) and a separate weekly window that only appears in the + * `retrieveUserQuotaSummary` RPC, grouped by model family ("Gemini Models", + * "Claude and GPT models") rather than by individual modelId. This guards: + * 1. The pure parser (`parseAntigravityWeeklyQuotas`) against the documented + * bucket shape (bucketId/displayName/remainingFraction/resetTime, window + * inferred from bucketId+displayName text — there is no explicit type field). + * 2. The end-to-end wiring: `getUsageForProvider()` merges the weekly group + * quotas alongside the existing per-model 5h quotas without clobbering them. + */ +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(), "omniroute-ag-weekly-")); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.API_KEY_SECRET = "test-ag-weekly-secret"; + +const core = await import("../../src/lib/db/core.ts"); +const { parseAntigravityWeeklyQuotas } = await import( + "../../open-sse/services/usage/antigravityWeeklyQuota.ts" +); +// Load usage.ts up-front (its index.ts proxyFetch patch runs at module eval) before mocks. +const usageModule = await import("../../open-sse/services/usage.ts"); +const { getUsageForProvider } = usageModule; + +const originalFetch = globalThis.fetch; + +test.after(() => { + globalThis.fetch = originalFetch; + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +const RESET_IN_3_DAYS = new Date(Date.now() + 3 * 24 * 60 * 60 * 1000).toISOString(); +const RESET_IN_2_HOURS = new Date(Date.now() + 2 * 60 * 60 * 1000).toISOString(); + +function requestUrl(input: RequestInfo | URL): string { + if (typeof input === "string") return input; + if (input instanceof URL) return input.toString(); + return input.url; +} + +interface UsageResult { + quotas: Record; +} + +test("parseAntigravityWeeklyQuotas extracts the weekly bucket per model-family group", () => { + const summary = { + groups: [ + { + displayName: "Gemini Models", + buckets: [ + { + bucketId: "gemini-5h", + displayName: "5 Hour Quota", + remainingFraction: 0.4, + resetTime: RESET_IN_2_HOURS, + }, + { + bucketId: "gemini-weekly", + displayName: "Weekly Quota", + remainingFraction: 0.75, + resetTime: RESET_IN_3_DAYS, + }, + ], + }, + { + displayName: "Claude and GPT models", + buckets: [ + { + bucketId: "claude-gpt-weekly", + displayName: "Weekly Quota", + remainingFraction: 0.1, + resetTime: RESET_IN_3_DAYS, + }, + ], + }, + ], + }; + + const quotas = parseAntigravityWeeklyQuotas(summary); + + assert.ok(quotas.gemini_weekly, "gemini weekly bucket extracted"); + assert.equal(quotas.gemini_weekly.remainingPercentage, 75); + assert.equal(quotas.gemini_weekly.resetAt, RESET_IN_3_DAYS); + assert.equal(quotas.gemini_weekly.unlimited, false); + + assert.ok(quotas.claude_gpt_weekly, "claude/gpt weekly bucket extracted"); + assert.equal(quotas.claude_gpt_weekly.remainingPercentage, 10); + + // The 5h bucket in the same group must NOT be picked up as "weekly" — + // only one entry per group, and it must be the one whose text says "weekly". + assert.equal(Object.keys(quotas).length, 2); +}); + +test("parseAntigravityWeeklyQuotas tolerates the quotaSummary-nested envelope", () => { + const summary = { + quotaSummary: { + groups: [ + { + displayName: "Gemini Models", + buckets: [ + { bucketId: "weekly", displayName: "Weekly", remainingFraction: 0.5, resetTime: RESET_IN_3_DAYS }, + ], + }, + ], + }, + }; + + const quotas = parseAntigravityWeeklyQuotas(summary); + assert.ok(quotas.gemini_weekly); + assert.equal(quotas.gemini_weekly.remainingPercentage, 50); +}); + +test("parseAntigravityWeeklyQuotas returns {} for missing/malformed data (best-effort)", () => { + assert.deepEqual(parseAntigravityWeeklyQuotas(null), {}); + assert.deepEqual(parseAntigravityWeeklyQuotas(undefined), {}); + assert.deepEqual(parseAntigravityWeeklyQuotas({}), {}); + assert.deepEqual(parseAntigravityWeeklyQuotas({ groups: [{ displayName: "Gemini Models" }] }), {}); +}); + +test("getUsageForProvider(antigravity) merges weekly group quotas alongside per-model 5h quotas", async () => { + core.resetDbInstance(); + + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = requestUrl(input); + + if (url.includes("retrieveUserQuotaSummary")) { + return { + ok: true, + json: async () => ({ + groups: [ + { + displayName: "Gemini Models", + buckets: [ + { + bucketId: "gemini-weekly", + displayName: "Weekly Quota", + remainingFraction: 0.6, + resetTime: RESET_IN_3_DAYS, + }, + ], + }, + ], + }), + } as Response; + } + + if (url.includes("retrieveUserQuota")) { + return { + ok: true, + json: async () => ({ + buckets: [ + { + modelId: "gemini-3.5-flash-high", + remainingFraction: 0.4, + resetTime: RESET_IN_2_HOURS, + }, + ], + }), + } as Response; + } + + // fetchAvailableModels + return { + ok: true, + json: async () => ({ + models: { + "gemini-3.5-flash-high": { + quotaInfo: { remainingFraction: 1.0, resetTime: RESET_IN_2_HOURS }, + }, + }, + }), + } as Response; + }) as typeof fetch; + + const connection = { + id: "conn-weekly-1", + provider: "antigravity", + accessToken: "fake-token-weekly-unique", + providerSpecificData: {}, + projectId: "test-project", + }; + + const result = await getUsageForProvider(connection, { forceRefresh: true }); + assert.ok(result && "quotas" in result, "should return quotas"); + const quotas = (result as UsageResult).quotas; + + // Existing per-model 5h quota is untouched. + assert.ok(quotas["gemini-3.5-flash-high"], "per-model 5h quota still present"); + assert.equal(quotas["gemini-3.5-flash-high"].quotaSource, "retrieveUserQuota"); + + // New weekly group quota is merged in alongside it. + assert.ok(quotas.gemini_weekly, "weekly group quota merged in"); + assert.equal(quotas.gemini_weekly.remainingPercentage, 60); + assert.equal(quotas.gemini_weekly.resetAt, RESET_IN_3_DAYS); +}); + +test("getUsageForProvider(antigravity) is unaffected when retrieveUserQuotaSummary is unavailable", async () => { + core.resetDbInstance(); + + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = requestUrl(input); + if (url.includes("retrieveUserQuotaSummary")) { + return { ok: false, status: 404, json: async () => ({}) } as Response; + } + if (url.includes("retrieveUserQuota")) { + return { ok: false, status: 404, json: async () => ({}) } as Response; + } + return { + ok: true, + json: async () => ({ + models: { + "gemini-3.5-flash-high": { + quotaInfo: { remainingFraction: 1.0, resetTime: RESET_IN_2_HOURS }, + }, + }, + }), + } as Response; + }) as typeof fetch; + + const connection = { + id: "conn-weekly-2", + provider: "antigravity", + accessToken: "fake-token-weekly-unavailable", + providerSpecificData: {}, + projectId: "test-project", + }; + + const result = await getUsageForProvider(connection, { forceRefresh: true }); + const quotas = (result as UsageResult).quotas; + assert.ok(quotas["gemini-3.5-flash-high"], "per-model quota still present without weekly data"); + assert.equal(quotas.gemini_weekly, undefined, "no weekly key when the RPC is unavailable"); +});