feat(antigravity): add weekly quota tracking and free-tier handling (#3892)

This commit is contained in:
Christian Gennari
2026-09-09 09:54:40 +07:00
parent 628ff1eab5
commit e3bf94ee25
10 changed files with 811 additions and 9 deletions

View File

@@ -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": {

View File

@@ -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": {

View File

@@ -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": {

View File

@@ -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");

View 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)");
});
});

View 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",
});
});
});