feat(antigravity): add weekly quota tracking and free-tier handling (#3892)
This commit is contained in:
Christian Gennari committed
1 parent
628ff1eab5
commit
e3bf94ee25
10 files changed
+811
-9
No files matched your search
@@ -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