feat(antigravity): add weekly quota tracking and free-tier handling (#3892)
This commit is contained in:
Christian Gennari
@@ -37,6 +37,7 @@ export default {
|
||||
},
|
||||
usage: {
|
||||
quotaApiUrl: `${ANTIGRAVITY_IDE_BASE_URL}/v1internal:fetchAvailableModels`,
|
||||
quotaSummaryApiUrl: `${ANTIGRAVITY_IDE_BASE_URL}/v1internal:retrieveUserQuotaSummary`,
|
||||
loadProjectApiUrl: "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist",
|
||||
tokenUrl: "https://oauth2.googleapis.com/token",
|
||||
},
|
||||
|
||||
150
open-sse/services/usage/antigravity-weekly.js
Normal file
150
open-sse/services/usage/antigravity-weekly.js
Normal file
@@ -0,0 +1,150 @@
|
||||
/**
|
||||
* Antigravity weekly quota — best-effort retrieval from retrieveUserQuotaSummary.
|
||||
* Failure never breaks existing per-model quota display.
|
||||
*/
|
||||
|
||||
import { U, parseResetTime, fetchWithTimeout } from "./shared.js";
|
||||
import { ANTIGRAVITY_IDE_USER_AGENT, ANTIGRAVITY_IDE_VERSION } from "../../providers/shared.js";
|
||||
|
||||
// — Weekly quota summary config ——————————————————————————————
|
||||
const WEEKLY_CONFIG = {
|
||||
...U("antigravity"),
|
||||
userAgent: ANTIGRAVITY_IDE_USER_AGENT,
|
||||
};
|
||||
|
||||
// — Cache: TTL + in-flight dedup per project ———————————————
|
||||
const WEEKLY_CACHE_TTL_MS = 180_000; // 3 minutes
|
||||
const weeklyCache = new Map(); // cacheKey -> { result, expiresAt } | { promise }
|
||||
|
||||
function cacheKey(accessToken, projectId) {
|
||||
return `${accessToken}::${projectId || ""}`;
|
||||
}
|
||||
|
||||
// Exported for tests only
|
||||
export function _clearWeeklyCache() {
|
||||
weeklyCache.clear();
|
||||
}
|
||||
|
||||
// — Group-name to stable key mapping ——————————————————————
|
||||
const GROUP_MATCHERS = [
|
||||
{ pattern: /gemini/i, key: "gemini_weekly", displayName: "Gemini (Weekly)" },
|
||||
{ pattern: /claude|gpt/i, key: "claude_gpt_weekly", displayName: "Claude & GPT (Weekly)" },
|
||||
];
|
||||
|
||||
/**
|
||||
* Parse a retrieveUserQuotaSummary response into normalized weekly quotas.
|
||||
* Pure function — safe to unit-test without network.
|
||||
*
|
||||
* @param {Object|null} data Raw JSON response
|
||||
* @returns {Object} e.g. { gemini_weekly: { used, total, ... }, claude_gpt_weekly: { ... } }
|
||||
*/
|
||||
export function parseWeeklyQuotaSummary(data) {
|
||||
if (!data || typeof data !== "object") return {};
|
||||
|
||||
// Groups may live at data.groups or data.quotaSummary.groups
|
||||
const groups = Array.isArray(data.groups)
|
||||
? data.groups
|
||||
: Array.isArray(data.quotaSummary?.groups)
|
||||
? data.quotaSummary.groups
|
||||
: null;
|
||||
|
||||
if (!groups) return {};
|
||||
|
||||
const result = {};
|
||||
|
||||
for (const group of groups) {
|
||||
if (!group || typeof group !== "object") continue;
|
||||
const displayName = group.displayName || "";
|
||||
|
||||
const buckets = Array.isArray(group.buckets) ? group.buckets : [];
|
||||
for (const bucket of buckets) {
|
||||
if (!bucket || typeof bucket !== "object") continue;
|
||||
|
||||
// Identify weekly buckets by checking bucketId + displayName for "weekly"
|
||||
const bucketText = `${bucket.bucketId || ""} ${bucket.displayName || ""}`.toLowerCase();
|
||||
if (!bucketText.includes("weekly")) continue;
|
||||
|
||||
// Skip disabled buckets
|
||||
if (bucket.disabled === true) continue;
|
||||
|
||||
const remainingFraction = Number(bucket.remainingFraction);
|
||||
if (!Number.isFinite(remainingFraction)) continue;
|
||||
|
||||
// Match group to a known family
|
||||
for (const matcher of GROUP_MATCHERS) {
|
||||
if (matcher.pattern.test(displayName)) {
|
||||
const total = 1000;
|
||||
const remaining = Math.round(total * remainingFraction);
|
||||
const used = Math.max(0, total - remaining);
|
||||
|
||||
result[matcher.key] = {
|
||||
used,
|
||||
total,
|
||||
resetAt: parseResetTime(bucket.resetTime),
|
||||
remainingPercentage: remainingFraction * 100,
|
||||
unlimited: false,
|
||||
displayName: matcher.displayName,
|
||||
};
|
||||
break; // first matching bucket per family wins
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
* Fetch weekly quota summary — cached, deduped, never throws.
|
||||
*/
|
||||
export async function fetchAntigravityWeeklyQuota(accessToken, projectId, proxyOptions = null) {
|
||||
const key = cacheKey(accessToken, projectId);
|
||||
|
||||
// Serve in-flight or cached
|
||||
const hit = weeklyCache.get(key);
|
||||
if (hit?.promise) return hit.promise;
|
||||
if (hit && hit.expiresAt > Date.now()) return hit.result;
|
||||
|
||||
const promise = (async () => {
|
||||
try {
|
||||
const url = WEEKLY_CONFIG.quotaSummaryApiUrl;
|
||||
if (!url) return {};
|
||||
|
||||
const response = await fetchWithTimeout(url, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Authorization": `Bearer ${accessToken}`,
|
||||
"User-Agent": WEEKLY_CONFIG.userAgent,
|
||||
"Content-Type": "application/json",
|
||||
"X-Client-Name": "antigravity",
|
||||
"X-Client-Version": ANTIGRAVITY_IDE_VERSION,
|
||||
},
|
||||
body: JSON.stringify({
|
||||
...(projectId ? { project: projectId } : {}),
|
||||
}),
|
||||
}, 10000, proxyOptions);
|
||||
|
||||
if (!response.ok) return {};
|
||||
|
||||
const data = await response.json();
|
||||
return parseWeeklyQuotaSummary(data);
|
||||
} catch {
|
||||
return {};
|
||||
}
|
||||
})();
|
||||
|
||||
weeklyCache.set(key, { promise });
|
||||
|
||||
try {
|
||||
const result = await promise;
|
||||
if (result && Object.keys(result).length > 0) {
|
||||
weeklyCache.set(key, { result, expiresAt: Date.now() + WEEKLY_CACHE_TTL_MS });
|
||||
} else {
|
||||
weeklyCache.delete(key);
|
||||
}
|
||||
return result;
|
||||
} catch {
|
||||
weeklyCache.delete(key);
|
||||
return {};
|
||||
}
|
||||
}
|
||||
@@ -5,6 +5,7 @@
|
||||
import { CLIENT_METADATA } from "../../config/appConstants.js";
|
||||
import { ANTIGRAVITY_IDE_USER_AGENT, ANTIGRAVITY_IDE_VERSION, ANTIGRAVITY_OAUTH_CLIENT } from "../../providers/shared.js";
|
||||
import { U, parseResetTime, normalizeCloudCodeProjectId, fetchWithTimeout } from "./shared.js";
|
||||
import { fetchAntigravityWeeklyQuota } from "./antigravity-weekly.js";
|
||||
|
||||
// Antigravity API config (from Quotio) — urls from registry, oauth client + dynamic UA kept here
|
||||
const ANTIGRAVITY_CONFIG = {
|
||||
@@ -157,8 +158,15 @@ export async function getAntigravityUsage(accessToken, providerSpecificData, pro
|
||||
const data = await response.json();
|
||||
const quotas = {};
|
||||
|
||||
// Parse model quotas (inspired by vscode-antigravity-cockpit)
|
||||
if (data.models) {
|
||||
// Detect tier: free-tier accounts only have weekly quotas (no separate 5h window).
|
||||
// On free-tier, fetchAvailableModels returns misleading per-model quota info
|
||||
// (missing remainingFraction defaults to 0, or reflects the weekly limit not a 5h window).
|
||||
const paidTierId = subscriptionInfo?.paidTier?.id;
|
||||
const isFreeTier = !paidTierId || paidTierId === "free-tier";
|
||||
|
||||
// Parse model quotas only for paid-tier accounts.
|
||||
// Free-tier accounts skip this — their only meaningful quota is the weekly limit.
|
||||
if (!isFreeTier && data.models) {
|
||||
// Filter only recommended/important models (must match PROVIDER_MODELS ag ids)
|
||||
const importantModels = [
|
||||
'gemini-3.8-flash-high',
|
||||
@@ -212,6 +220,56 @@ export async function getAntigravityUsage(accessToken, providerSpecificData, pro
|
||||
}
|
||||
}
|
||||
|
||||
// Best-effort weekly quota overlay — never blocks or breaks per-model results
|
||||
try {
|
||||
const weeklyQuotas = await fetchAntigravityWeeklyQuota(
|
||||
accessToken,
|
||||
projectId,
|
||||
proxyOptions
|
||||
);
|
||||
|
||||
// Reconcile weekly quota against model family status:
|
||||
// If every model in a family is locked/exhausted (remainingPercentage === 0)
|
||||
// until a future reset time, the weekly limit cannot be 100% available.
|
||||
// On Google's Free Starter tier, retrieveUserQuotaSummary buggily reports
|
||||
// remainingFraction: 1 even after the starter quota is depleted and all models 429.
|
||||
const entries = Object.entries(quotas);
|
||||
const geminiModels = entries.filter(([k]) => k.startsWith("gemini-") && !k.includes("image"));
|
||||
const claudeModels = entries.filter(([k]) => k.startsWith("claude-"));
|
||||
|
||||
if (weeklyQuotas.gemini_weekly && geminiModels.length > 0) {
|
||||
const allGeminiExhausted = geminiModels.every(([, q]) => (q.remainingPercentage ?? 0) === 0);
|
||||
if (allGeminiExhausted && weeklyQuotas.gemini_weekly.remainingPercentage > 0) {
|
||||
const maxResetAt = geminiModels.reduce((max, [, q]) =>
|
||||
!max || (q.resetAt && new Date(q.resetAt) > new Date(max)) ? q.resetAt : max, null
|
||||
);
|
||||
weeklyQuotas.gemini_weekly.used = weeklyQuotas.gemini_weekly.total;
|
||||
weeklyQuotas.gemini_weekly.remainingPercentage = 0;
|
||||
if (maxResetAt) {
|
||||
weeklyQuotas.gemini_weekly.resetAt = maxResetAt;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (weeklyQuotas.claude_gpt_weekly && claudeModels.length > 0) {
|
||||
const allClaudeExhausted = claudeModels.every(([, q]) => (q.remainingPercentage ?? 0) === 0);
|
||||
if (allClaudeExhausted && weeklyQuotas.claude_gpt_weekly.remainingPercentage > 0) {
|
||||
const maxResetAt = claudeModels.reduce((max, [, q]) =>
|
||||
!max || (q.resetAt && new Date(q.resetAt) > new Date(max)) ? q.resetAt : max, null
|
||||
);
|
||||
weeklyQuotas.claude_gpt_weekly.used = weeklyQuotas.claude_gpt_weekly.total;
|
||||
weeklyQuotas.claude_gpt_weekly.remainingPercentage = 0;
|
||||
if (maxResetAt) {
|
||||
weeklyQuotas.claude_gpt_weekly.resetAt = maxResetAt;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Object.assign(quotas, weeklyQuotas);
|
||||
} catch {
|
||||
// Silently ignore — weekly is best-effort
|
||||
}
|
||||
|
||||
return {
|
||||
plan: subscriptionInfo?.currentTier?.name || "Unknown",
|
||||
quotas,
|
||||
|
||||
@@ -376,10 +376,12 @@ export function parseQuotaData(provider, data) {
|
||||
case "antigravity":
|
||||
if (data.quotas) {
|
||||
const entries = Object.entries(data.quotas);
|
||||
const weeklyKeys = new Set(["gemini_weekly", "claude_gpt_weekly"]);
|
||||
const geminiModels = entries.filter(([k]) => k.startsWith("gemini-") && !k.includes("image"));
|
||||
const claudeModels = entries.filter(([k]) => k.startsWith("claude-"));
|
||||
const imageModels = entries.filter(([k]) => k.includes("image"));
|
||||
const otherModels = entries.filter(([k]) => !k.startsWith("gemini-") && !k.startsWith("claude-") && !k.includes("image"));
|
||||
const weeklyModels = entries.filter(([k]) => weeklyKeys.has(k));
|
||||
const otherModels = entries.filter(([k]) => !k.startsWith("gemini-") && !k.startsWith("claude-") && !k.includes("image") && !weeklyKeys.has(k));
|
||||
|
||||
if (geminiModels.length > 0) {
|
||||
const rep = geminiModels.reduce((min, cur) =>
|
||||
@@ -409,6 +411,17 @@ export function parseQuotaData(provider, data) {
|
||||
});
|
||||
}
|
||||
|
||||
weeklyModels.forEach(([modelKey, quota]) => {
|
||||
normalizedQuotas.push({
|
||||
name: quota.displayName || modelKey,
|
||||
modelKey,
|
||||
used: quota.used || 0,
|
||||
total: quota.total || 0,
|
||||
resetAt: quota.resetAt || null,
|
||||
remainingPercentage: quota.remainingPercentage,
|
||||
});
|
||||
});
|
||||
|
||||
imageModels.forEach(([modelKey, quota]) => {
|
||||
normalizedQuotas.push({
|
||||
name: quota.displayName || modelKey,
|
||||
|
||||
@@ -4,7 +4,7 @@ const proxyAwareFetch = vi.fn(async (url) => ({
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: async () => url.includes(":loadCodeAssist")
|
||||
? { cloudaicompanionProject: "project-1", currentTier: { name: "Pro" } }
|
||||
? { cloudaicompanionProject: "project-1", currentTier: { name: "Pro" }, paidTier: { id: "g1-pro-tier", name: "Google AI Pro" } }
|
||||
: {
|
||||
models: {
|
||||
"gemini-3.6-flash-high": {
|
||||
|
||||
@@ -4,7 +4,7 @@ const proxyAwareFetch = vi.fn(async (url) => ({
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: async () => url.includes(":loadCodeAssist")
|
||||
? { cloudaicompanionProject: "project-1", currentTier: { name: "Pro" } }
|
||||
? { cloudaicompanionProject: "project-1", currentTier: { name: "Pro" }, paidTier: { id: "g1-pro-tier", name: "Google AI Pro" } }
|
||||
: {
|
||||
models: {
|
||||
"gemini-3.7-flash-high": {
|
||||
|
||||
@@ -4,7 +4,7 @@ const proxyAwareFetch = vi.fn(async (url) => ({
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: async () => url.includes(":loadCodeAssist")
|
||||
? { cloudaicompanionProject: "project-1", currentTier: { name: "Pro" } }
|
||||
? { cloudaicompanionProject: "project-1", currentTier: { name: "Pro" }, paidTier: { id: "g1-pro-tier", name: "Google AI Pro" } }
|
||||
: {
|
||||
models: {
|
||||
"gemini-3.8-flash-high": {
|
||||
|
||||
@@ -4,8 +4,10 @@ const proxyAwareFetch = vi.fn(async (url) => ({
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: async () => url.includes(":loadCodeAssist")
|
||||
? { cloudaicompanionProject: "project-1", currentTier: { name: "Pro" } }
|
||||
: { models: {} },
|
||||
? { cloudaicompanionProject: "project-1", currentTier: { name: "Pro" }, paidTier: { id: "g1-pro-tier", name: "Google AI Pro" } }
|
||||
: url.includes(":retrieveUserQuotaSummary")
|
||||
? { groups: [] }
|
||||
: { models: {} },
|
||||
text: async () => "{}",
|
||||
}));
|
||||
|
||||
@@ -21,7 +23,8 @@ describe("Antigravity usage headers", () => {
|
||||
|
||||
await getAntigravityUsage("access-token", {});
|
||||
|
||||
expect(proxyAwareFetch).toHaveBeenCalledTimes(2);
|
||||
// loadCodeAssist + fetchAvailableModels + retrieveUserQuotaSummary
|
||||
expect(proxyAwareFetch).toHaveBeenCalledTimes(3);
|
||||
for (const [, options] of proxyAwareFetch.mock.calls) {
|
||||
expect(options.headers["User-Agent"]).toBe("antigravity/ide/2.11.0 darwin/arm64");
|
||||
expect(options.headers).not.toHaveProperty("x-request-source");
|
||||
|
||||
113
tests/unit/antigravity-weekly-dashboard.test.js
Normal file
113
tests/unit/antigravity-weekly-dashboard.test.js
Normal file
@@ -0,0 +1,113 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { parseQuotaData } from "@/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.js";
|
||||
|
||||
describe("Antigravity dashboard normalization with weekly quotas", () => {
|
||||
const data = {
|
||||
quotas: {
|
||||
"gemini-pro-agent": {
|
||||
displayName: "Gemini 3.1 Pro (High)",
|
||||
used: 200,
|
||||
total: 1000,
|
||||
resetAt: "2026-09-08T00:00:00Z",
|
||||
remainingPercentage: 80,
|
||||
},
|
||||
"claude-opus-4-6-thinking": {
|
||||
displayName: "Claude Opus 4.6 (Thinking)",
|
||||
used: 100,
|
||||
total: 1000,
|
||||
resetAt: "2026-09-08T00:00:00Z",
|
||||
remainingPercentage: 90,
|
||||
},
|
||||
gemini_weekly: {
|
||||
displayName: "Gemini (Weekly)",
|
||||
used: 250,
|
||||
total: 1000,
|
||||
resetAt: "2026-09-15T00:00:00Z",
|
||||
remainingPercentage: 75,
|
||||
},
|
||||
claude_gpt_weekly: {
|
||||
displayName: "Claude & GPT (Weekly)",
|
||||
used: 500,
|
||||
total: 1000,
|
||||
resetAt: "2026-09-14T00:00:00Z",
|
||||
remainingPercentage: 50,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
it("includes weekly rows with correct display names", () => {
|
||||
const quotas = parseQuotaData("antigravity", data);
|
||||
const names = quotas.map((q) => q.name);
|
||||
|
||||
expect(names).toContain("Gemini (Flash / Pro)");
|
||||
expect(names).toContain("Claude (Sonnet / Opus)");
|
||||
expect(names).toContain("Gemini (Weekly)");
|
||||
expect(names).toContain("Claude & GPT (Weekly)");
|
||||
});
|
||||
|
||||
it("uses stable modelKey for weekly rows", () => {
|
||||
const quotas = parseQuotaData("antigravity", data);
|
||||
const keys = quotas.map((q) => q.modelKey);
|
||||
|
||||
expect(keys).toContain("gemini_weekly");
|
||||
expect(keys).toContain("claude_gpt_weekly");
|
||||
});
|
||||
|
||||
it("weekly rows carry correct quota values", () => {
|
||||
const quotas = parseQuotaData("antigravity", data);
|
||||
const geminiWeekly = quotas.find((q) => q.modelKey === "gemini_weekly");
|
||||
const claudeWeekly = quotas.find((q) => q.modelKey === "claude_gpt_weekly");
|
||||
|
||||
expect(geminiWeekly).toMatchObject({
|
||||
used: 250,
|
||||
total: 1000,
|
||||
remainingPercentage: 75,
|
||||
resetAt: "2026-09-15T00:00:00Z",
|
||||
});
|
||||
expect(claudeWeekly).toMatchObject({
|
||||
used: 500,
|
||||
total: 1000,
|
||||
remainingPercentage: 50,
|
||||
resetAt: "2026-09-14T00:00:00Z",
|
||||
});
|
||||
});
|
||||
|
||||
it("weekly rows do NOT appear as otherModels", () => {
|
||||
const quotas = parseQuotaData("antigravity", data);
|
||||
const weeklyRows = quotas.filter((q) =>
|
||||
q.modelKey === "gemini_weekly" || q.modelKey === "claude_gpt_weekly"
|
||||
);
|
||||
expect(weeklyRows).toHaveLength(2);
|
||||
expect(weeklyRows[0].name).toMatch(/Weekly/);
|
||||
expect(weeklyRows[1].name).toMatch(/Weekly/);
|
||||
});
|
||||
|
||||
it("order: gemini family, claude family, weekly, then other", () => {
|
||||
const quotas = parseQuotaData("antigravity", data);
|
||||
const keys = quotas.map((q) => q.modelKey);
|
||||
|
||||
const geminiIdx = keys.indexOf("gemini");
|
||||
const claudeIdx = keys.indexOf("claude");
|
||||
const geminiWeeklyIdx = keys.indexOf("gemini_weekly");
|
||||
const claudeWeeklyIdx = keys.indexOf("claude_gpt_weekly");
|
||||
|
||||
expect(geminiIdx).toBeLessThan(geminiWeeklyIdx);
|
||||
expect(claudeIdx).toBeLessThan(claudeWeeklyIdx);
|
||||
});
|
||||
|
||||
it("works with no weekly keys present (backward compat)", () => {
|
||||
const noWeekly = {
|
||||
quotas: {
|
||||
"gemini-pro-agent": {
|
||||
displayName: "Gemini 3.1 Pro (High)",
|
||||
used: 200,
|
||||
total: 1000,
|
||||
remainingPercentage: 80,
|
||||
},
|
||||
},
|
||||
};
|
||||
const quotas = parseQuotaData("antigravity", noWeekly);
|
||||
expect(quotas).toHaveLength(1);
|
||||
expect(quotas[0].name).toBe("Gemini (Flash / Pro)");
|
||||
});
|
||||
});
|
||||
464
tests/unit/antigravity-weekly-quota.test.js
Normal file
464
tests/unit/antigravity-weekly-quota.test.js
Normal file
@@ -0,0 +1,464 @@
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
|
||||
// Mock proxyAwareFetch before any imports that use it
|
||||
vi.mock("../../open-sse/utils/proxyFetch.js", () => ({
|
||||
proxyAwareFetch: vi.fn(),
|
||||
}));
|
||||
|
||||
import { proxyAwareFetch } from "../../open-sse/utils/proxyFetch.js";
|
||||
import {
|
||||
parseWeeklyQuotaSummary,
|
||||
fetchAntigravityWeeklyQuota,
|
||||
_clearWeeklyCache,
|
||||
} from "../../open-sse/services/usage/antigravity-weekly.js";
|
||||
|
||||
// — Fixtures ——————————————————————————————————————————————
|
||||
const GEMINI_GROUP = {
|
||||
displayName: "Gemini Models",
|
||||
buckets: [
|
||||
{
|
||||
bucketId: "gemini-weekly-bucket",
|
||||
displayName: "Weekly Limit",
|
||||
remainingFraction: 0.75,
|
||||
resetTime: "2026-09-15T00:00:00Z",
|
||||
},
|
||||
{
|
||||
bucketId: "gemini-daily-bucket",
|
||||
displayName: "Daily Limit",
|
||||
remainingFraction: 0.9,
|
||||
resetTime: "2026-09-09T00:00:00Z",
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const CLAUDE_GPT_GROUP = {
|
||||
displayName: "Claude and GPT models",
|
||||
buckets: [
|
||||
{
|
||||
bucketId: "claude-gpt-weekly",
|
||||
displayName: "Weekly Quota",
|
||||
remainingFraction: 0.5,
|
||||
resetTime: "2026-09-14T00:00:00Z",
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const FULL_RESPONSE = { groups: [GEMINI_GROUP, CLAUDE_GPT_GROUP] };
|
||||
|
||||
const NESTED_RESPONSE = {
|
||||
quotaSummary: {
|
||||
groups: [GEMINI_GROUP, CLAUDE_GPT_GROUP],
|
||||
},
|
||||
};
|
||||
|
||||
// — parseWeeklyQuotaSummary ———————————————————————————————
|
||||
describe("parseWeeklyQuotaSummary", () => {
|
||||
it("extracts Gemini weekly quota from top-level groups", () => {
|
||||
const result = parseWeeklyQuotaSummary(FULL_RESPONSE);
|
||||
expect(result.gemini_weekly).toMatchObject({
|
||||
used: 250,
|
||||
total: 1000,
|
||||
remainingPercentage: 75,
|
||||
displayName: "Gemini (Weekly)",
|
||||
unlimited: false,
|
||||
});
|
||||
expect(result.gemini_weekly.resetAt).toBe("2026-09-15T00:00:00.000Z");
|
||||
});
|
||||
|
||||
it("extracts Claude & GPT weekly quota", () => {
|
||||
const result = parseWeeklyQuotaSummary(FULL_RESPONSE);
|
||||
expect(result.claude_gpt_weekly).toMatchObject({
|
||||
used: 500,
|
||||
total: 1000,
|
||||
remainingPercentage: 50,
|
||||
displayName: "Claude & GPT (Weekly)",
|
||||
unlimited: false,
|
||||
});
|
||||
expect(result.claude_gpt_weekly.resetAt).toBe("2026-09-14T00:00:00.000Z");
|
||||
});
|
||||
|
||||
it("handles alternate nested quotaSummary.groups shape", () => {
|
||||
const result = parseWeeklyQuotaSummary(NESTED_RESPONSE);
|
||||
expect(result.gemini_weekly).toBeDefined();
|
||||
expect(result.claude_gpt_weekly).toBeDefined();
|
||||
expect(result.gemini_weekly.remainingPercentage).toBe(75);
|
||||
expect(result.claude_gpt_weekly.remainingPercentage).toBe(50);
|
||||
});
|
||||
|
||||
it("skips non-weekly buckets", () => {
|
||||
const data = {
|
||||
groups: [{
|
||||
displayName: "Gemini Models",
|
||||
buckets: [
|
||||
{
|
||||
bucketId: "gemini-daily-bucket",
|
||||
displayName: "Daily Limit",
|
||||
remainingFraction: 0.9,
|
||||
resetTime: "2026-09-09T00:00:00Z",
|
||||
},
|
||||
],
|
||||
}],
|
||||
};
|
||||
const result = parseWeeklyQuotaSummary(data);
|
||||
expect(result).toEqual({});
|
||||
});
|
||||
|
||||
it("skips disabled weekly buckets", () => {
|
||||
const data = {
|
||||
groups: [{
|
||||
displayName: "Gemini Models",
|
||||
buckets: [{
|
||||
bucketId: "gemini-weekly-bucket",
|
||||
displayName: "Weekly Limit",
|
||||
remainingFraction: 0.75,
|
||||
resetTime: "2026-09-15T00:00:00Z",
|
||||
disabled: true,
|
||||
}],
|
||||
}],
|
||||
};
|
||||
const result = parseWeeklyQuotaSummary(data);
|
||||
expect(result).toEqual({});
|
||||
});
|
||||
|
||||
it("returns empty object for null/undefined input", () => {
|
||||
expect(parseWeeklyQuotaSummary(null)).toEqual({});
|
||||
expect(parseWeeklyQuotaSummary(undefined)).toEqual({});
|
||||
expect(parseWeeklyQuotaSummary("string")).toEqual({});
|
||||
});
|
||||
|
||||
it("returns empty object for response with no groups", () => {
|
||||
expect(parseWeeklyQuotaSummary({})).toEqual({});
|
||||
expect(parseWeeklyQuotaSummary({ groups: "not-array" })).toEqual({});
|
||||
expect(parseWeeklyQuotaSummary({ quotaSummary: {} })).toEqual({});
|
||||
});
|
||||
|
||||
it("handles groups with no buckets gracefully", () => {
|
||||
const data = {
|
||||
groups: [{ displayName: "Gemini Models" }],
|
||||
};
|
||||
expect(parseWeeklyQuotaSummary(data)).toEqual({});
|
||||
});
|
||||
|
||||
it("handles bucket with non-finite remainingFraction", () => {
|
||||
const data = {
|
||||
groups: [{
|
||||
displayName: "Gemini Models",
|
||||
buckets: [{
|
||||
bucketId: "weekly-bucket",
|
||||
displayName: "Weekly",
|
||||
remainingFraction: "not-a-number",
|
||||
}],
|
||||
}],
|
||||
};
|
||||
expect(parseWeeklyQuotaSummary(data)).toEqual({});
|
||||
});
|
||||
|
||||
it("ignores groups that don't match known families", () => {
|
||||
const data = {
|
||||
groups: [{
|
||||
displayName: "Unknown AI Provider",
|
||||
buckets: [{
|
||||
bucketId: "weekly-bucket",
|
||||
displayName: "Weekly",
|
||||
remainingFraction: 0.5,
|
||||
}],
|
||||
}],
|
||||
};
|
||||
expect(parseWeeklyQuotaSummary(data)).toEqual({});
|
||||
});
|
||||
});
|
||||
|
||||
// — fetchAntigravityWeeklyQuota ———————————————————————————
|
||||
describe("fetchAntigravityWeeklyQuota", () => {
|
||||
beforeEach(() => {
|
||||
proxyAwareFetch.mockReset();
|
||||
_clearWeeklyCache();
|
||||
});
|
||||
|
||||
it("fetches and returns parsed weekly quota on success", async () => {
|
||||
proxyAwareFetch.mockResolvedValue({
|
||||
ok: true,
|
||||
json: async () => FULL_RESPONSE,
|
||||
});
|
||||
|
||||
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
||||
expect(result.gemini_weekly).toBeDefined();
|
||||
expect(result.claude_gpt_weekly).toBeDefined();
|
||||
});
|
||||
|
||||
it("returns {} on HTTP 401", async () => {
|
||||
proxyAwareFetch.mockResolvedValue({ ok: false, status: 401 });
|
||||
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
||||
expect(result).toEqual({});
|
||||
});
|
||||
|
||||
it("returns {} on HTTP 403", async () => {
|
||||
proxyAwareFetch.mockResolvedValue({ ok: false, status: 403 });
|
||||
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
||||
expect(result).toEqual({});
|
||||
});
|
||||
|
||||
it("returns {} on HTTP 404", async () => {
|
||||
proxyAwareFetch.mockResolvedValue({ ok: false, status: 404 });
|
||||
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
||||
expect(result).toEqual({});
|
||||
});
|
||||
|
||||
it("returns {} on HTTP 429", async () => {
|
||||
proxyAwareFetch.mockResolvedValue({ ok: false, status: 429 });
|
||||
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
||||
expect(result).toEqual({});
|
||||
});
|
||||
|
||||
it("returns {} on network error", async () => {
|
||||
proxyAwareFetch.mockRejectedValue(new Error("network timeout"));
|
||||
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
||||
expect(result).toEqual({});
|
||||
});
|
||||
|
||||
it("returns {} on malformed JSON response", async () => {
|
||||
proxyAwareFetch.mockResolvedValue({
|
||||
ok: true,
|
||||
json: async () => { throw new SyntaxError("Unexpected token"); },
|
||||
});
|
||||
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
||||
expect(result).toEqual({});
|
||||
});
|
||||
|
||||
it("deduplicates concurrent requests for the same account", async () => {
|
||||
let resolveResponse;
|
||||
proxyAwareFetch.mockReturnValue(new Promise(resolve => {
|
||||
resolveResponse = resolve;
|
||||
}));
|
||||
|
||||
const p1 = fetchAntigravityWeeklyQuota("token", "project-1");
|
||||
const p2 = fetchAntigravityWeeklyQuota("token", "project-1");
|
||||
|
||||
resolveResponse({ ok: true, json: async () => FULL_RESPONSE });
|
||||
|
||||
const [r1, r2] = await Promise.all([p1, p2]);
|
||||
expect(r1).toEqual(r2);
|
||||
expect(proxyAwareFetch).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("serves cached result within TTL", async () => {
|
||||
proxyAwareFetch.mockResolvedValue({
|
||||
ok: true,
|
||||
json: async () => FULL_RESPONSE,
|
||||
});
|
||||
|
||||
await fetchAntigravityWeeklyQuota("token", "project-1");
|
||||
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
||||
|
||||
expect(proxyAwareFetch).toHaveBeenCalledTimes(1);
|
||||
expect(result.gemini_weekly).toBeDefined();
|
||||
});
|
||||
|
||||
it("sends correct headers and body", async () => {
|
||||
proxyAwareFetch.mockResolvedValue({
|
||||
ok: true,
|
||||
json: async () => ({ groups: [] }),
|
||||
});
|
||||
|
||||
await fetchAntigravityWeeklyQuota("token", "project-1");
|
||||
|
||||
expect(proxyAwareFetch).toHaveBeenCalledWith(
|
||||
"https://daily-cloudcode-pa.googleapis.com/v1internal:retrieveUserQuotaSummary",
|
||||
expect.objectContaining({
|
||||
method: "POST",
|
||||
headers: expect.objectContaining({
|
||||
"Authorization": "Bearer token",
|
||||
"User-Agent": "antigravity/ide/2.11.0 darwin/arm64",
|
||||
"Content-Type": "application/json",
|
||||
"X-Client-Name": "antigravity",
|
||||
}),
|
||||
body: JSON.stringify({ project: "project-1" }),
|
||||
}),
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
// — Integration: weekly failure does not affect existing quotas —————
|
||||
describe("weekly quota isolation from existing quota", () => {
|
||||
beforeEach(() => {
|
||||
proxyAwareFetch.mockReset();
|
||||
_clearWeeklyCache();
|
||||
});
|
||||
|
||||
it("existing getAntigravityUsage succeeds even when weekly RPC fails", async () => {
|
||||
proxyAwareFetch.mockImplementation(async (url) => {
|
||||
if (url.includes(":loadCodeAssist")) {
|
||||
return {
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: async () => ({ cloudaicompanionProject: "p1", currentTier: { name: "Pro" }, paidTier: { id: "g1-pro-tier", name: "Google AI Pro" } }),
|
||||
};
|
||||
}
|
||||
if (url.includes(":fetchAvailableModels")) {
|
||||
return {
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: async () => ({
|
||||
models: {
|
||||
"gemini-3.8-flash-high": {
|
||||
displayName: "Gemini 3.8 Flash (High)",
|
||||
quotaInfo: { remainingFraction: 0.85, resetTime: "2026-09-15T00:00:00Z" },
|
||||
},
|
||||
},
|
||||
}),
|
||||
};
|
||||
}
|
||||
if (url.includes(":retrieveUserQuotaSummary")) {
|
||||
throw new Error("weekly endpoint unavailable");
|
||||
}
|
||||
return { ok: false, status: 404 };
|
||||
});
|
||||
|
||||
const { getAntigravityUsage } = await import("../../open-sse/services/usage/google.js");
|
||||
const result = await getAntigravityUsage("token", {});
|
||||
|
||||
expect(result.quotas["gemini-3.8-flash-high"]).toMatchObject({
|
||||
used: 150,
|
||||
total: 1000,
|
||||
remainingPercentage: 85,
|
||||
});
|
||||
expect(result.quotas.gemini_weekly).toBeUndefined();
|
||||
expect(result.quotas.claude_gpt_weekly).toBeUndefined();
|
||||
expect(result.message).toBeUndefined();
|
||||
});
|
||||
|
||||
it("free-tier accounts only show weekly quotas, not per-model short-window quotas", async () => {
|
||||
proxyAwareFetch.mockImplementation(async (url) => {
|
||||
if (url.includes(":loadCodeAssist")) {
|
||||
return {
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: async () => ({ cloudaicompanionProject: "p1", currentTier: { name: "Starter" }, paidTier: { id: "free-tier", name: "Antigravity Starter Quota" } }),
|
||||
};
|
||||
}
|
||||
if (url.includes(":fetchAvailableModels")) {
|
||||
return {
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: async () => ({
|
||||
models: {
|
||||
"gemini-3.8-flash-high": {
|
||||
displayName: "Gemini 3.8 Flash (High)",
|
||||
quotaInfo: { remainingFraction: 1, resetTime: "2026-09-15T00:00:00Z" },
|
||||
},
|
||||
"claude-sonnet-4-6": {
|
||||
displayName: "Claude Sonnet 4.6",
|
||||
// Missing remainingFraction — free tier exhausted
|
||||
quotaInfo: { resetTime: "2026-09-13T12:00:00Z" },
|
||||
},
|
||||
},
|
||||
}),
|
||||
};
|
||||
}
|
||||
if (url.includes(":retrieveUserQuotaSummary")) {
|
||||
return {
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: async () => ({
|
||||
groups: [{
|
||||
displayName: "Gemini Models",
|
||||
buckets: [{
|
||||
bucketId: "gemini-weekly",
|
||||
displayName: "Weekly Limit Remaining",
|
||||
remainingFraction: 1,
|
||||
resetTime: "2026-09-15T00:00:00Z",
|
||||
}],
|
||||
}, {
|
||||
displayName: "Claude and GPT models",
|
||||
buckets: [{
|
||||
bucketId: "3p-weekly",
|
||||
displayName: "Weekly Limit Remaining",
|
||||
remainingFraction: 0,
|
||||
resetTime: "2026-09-13T12:00:00Z",
|
||||
}],
|
||||
}],
|
||||
}),
|
||||
};
|
||||
}
|
||||
return { ok: false, status: 404 };
|
||||
});
|
||||
|
||||
const { getAntigravityUsage } = await import("../../open-sse/services/usage/google.js");
|
||||
const result = await getAntigravityUsage("token", {});
|
||||
|
||||
// Per-model quotas should be absent (free-tier accounts skip model parsing)
|
||||
expect(result.quotas["gemini-3.8-flash-high"]).toBeUndefined();
|
||||
expect(result.quotas["claude-sonnet-4-6"]).toBeUndefined();
|
||||
|
||||
// Only weekly quotas should appear
|
||||
expect(result.quotas.gemini_weekly).toMatchObject({
|
||||
used: 0,
|
||||
total: 1000,
|
||||
remainingPercentage: 100,
|
||||
});
|
||||
expect(result.quotas.claude_gpt_weekly).toMatchObject({
|
||||
used: 1000,
|
||||
total: 1000,
|
||||
remainingPercentage: 0,
|
||||
});
|
||||
});
|
||||
|
||||
it("reconciles weekly quota to 0% when all paid-tier family models are exhausted", async () => {
|
||||
proxyAwareFetch.mockImplementation(async (url) => {
|
||||
if (url.includes(":loadCodeAssist")) {
|
||||
return {
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: async () => ({ cloudaicompanionProject: "p1", currentTier: { name: "Pro" }, paidTier: { id: "g1-pro-tier", name: "Google AI Pro" } }),
|
||||
};
|
||||
}
|
||||
if (url.includes(":fetchAvailableModels")) {
|
||||
return {
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: async () => ({
|
||||
models: {
|
||||
"gemini-3.8-flash-high": {
|
||||
displayName: "Gemini 3.8 Flash (High)",
|
||||
// Exhausted model: no remainingFraction, future resetTime
|
||||
quotaInfo: { resetTime: "2026-09-13T12:00:00Z" },
|
||||
},
|
||||
},
|
||||
}),
|
||||
};
|
||||
}
|
||||
if (url.includes(":retrieveUserQuotaSummary")) {
|
||||
return {
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: async () => ({
|
||||
groups: [{
|
||||
displayName: "Gemini Models",
|
||||
buckets: [{
|
||||
bucketId: "gemini-weekly",
|
||||
displayName: "Weekly Limit Remaining",
|
||||
remainingFraction: 1,
|
||||
resetTime: "2026-09-15T00:00:00Z",
|
||||
}],
|
||||
}],
|
||||
}),
|
||||
};
|
||||
}
|
||||
return { ok: false, status: 404 };
|
||||
});
|
||||
|
||||
const { getAntigravityUsage } = await import("../../open-sse/services/usage/google.js");
|
||||
const result = await getAntigravityUsage("token", {});
|
||||
|
||||
// Per-model quota should show exhausted
|
||||
expect(result.quotas["gemini-3.8-flash-high"].remainingPercentage).toBe(0);
|
||||
// Weekly quota should be reconciled to 0% with the family reset time
|
||||
expect(result.quotas.gemini_weekly).toMatchObject({
|
||||
used: 1000,
|
||||
total: 1000,
|
||||
remainingPercentage: 0,
|
||||
resetAt: "2026-09-13T12:00:00.000Z",
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user