openclaw/src/infra/provider-usage.fetch.antigravity.test.ts
Hiren Patel 4de660bec6
[AI Assisted] Usage: add Google Antigravity usage tracking (#1490)
* Usage: add Google Antigravity usage tracking

- Add dedicated fetcher for google-antigravity provider
- Fetch credits and per-model quotas from Cloud Code API
- Report individual model IDs sorted by usage (top 10)
- Include comprehensive debug logging with [antigravity] prefix

* fix: refine antigravity usage tracking (#1490) (thanks @patelhiren)

---------

Co-authored-by: Peter Steinberger <steipete@gmail.com>
2026-01-23 07:17:59 +00:00

579 lines
20 KiB
TypeScript

import { describe, expect, it, vi } from "vitest";
import { fetchAntigravityUsage } from "./provider-usage.fetch.antigravity.js";
const makeResponse = (status: number, body: unknown): Response => {
const payload = typeof body === "string" ? body : JSON.stringify(body);
const headers = typeof body === "string" ? undefined : { "Content-Type": "application/json" };
return new Response(payload, { status, headers });
};
describe("fetchAntigravityUsage", () => {
it("returns 3 windows when both endpoints succeed", async () => {
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(200, {
availablePromptCredits: 750,
planInfo: { monthlyPromptCredits: 1000 },
planType: "Standard",
currentTier: { id: "tier1", name: "Standard Tier" },
});
}
if (url.includes("fetchAvailableModels")) {
return makeResponse(200, {
models: {
"gemini-pro-1.5": {
quotaInfo: {
remainingFraction: 0.6,
resetTime: "2026-01-08T00:00:00Z",
isExhausted: false,
},
},
"gemini-flash-2.0": {
quotaInfo: {
remainingFraction: 0.8,
resetTime: "2026-01-08T00:00:00Z",
isExhausted: false,
},
},
},
});
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(snapshot.provider).toBe("google-antigravity");
expect(snapshot.displayName).toBe("Antigravity");
expect(snapshot.windows).toHaveLength(3);
expect(snapshot.plan).toBe("Standard Tier");
expect(snapshot.error).toBeUndefined();
const creditsWindow = snapshot.windows.find((w) => w.label === "Credits");
expect(creditsWindow?.usedPercent).toBe(25); // (1000 - 750) / 1000 * 100
const proWindow = snapshot.windows.find((w) => w.label === "gemini-pro-1.5");
expect(proWindow?.usedPercent).toBe(40); // (1 - 0.6) * 100
expect(proWindow?.resetAt).toBe(new Date("2026-01-08T00:00:00Z").getTime());
const flashWindow = snapshot.windows.find((w) => w.label === "gemini-flash-2.0");
expect(flashWindow?.usedPercent).toBeCloseTo(20, 1); // (1 - 0.8) * 100
expect(flashWindow?.resetAt).toBe(new Date("2026-01-08T00:00:00Z").getTime());
expect(mockFetch).toHaveBeenCalledTimes(2);
});
it("returns Credits only when loadCodeAssist succeeds but fetchAvailableModels fails", async () => {
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(200, {
availablePromptCredits: 250,
planInfo: { monthlyPromptCredits: 1000 },
currentTier: { name: "Free" },
});
}
if (url.includes("fetchAvailableModels")) {
return makeResponse(403, { error: { message: "Permission denied" } });
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(snapshot.provider).toBe("google-antigravity");
expect(snapshot.windows).toHaveLength(1);
expect(snapshot.plan).toBe("Free");
expect(snapshot.error).toBeUndefined();
const creditsWindow = snapshot.windows[0];
expect(creditsWindow?.label).toBe("Credits");
expect(creditsWindow?.usedPercent).toBe(75); // (1000 - 250) / 1000 * 100
expect(mockFetch).toHaveBeenCalledTimes(2);
});
it("returns model IDs when fetchAvailableModels succeeds but loadCodeAssist fails", async () => {
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(500, "Internal server error");
}
if (url.includes("fetchAvailableModels")) {
return makeResponse(200, {
models: {
"gemini-pro-1.5": {
quotaInfo: { remainingFraction: 0.5, resetTime: "2026-01-08T00:00:00Z" },
},
"gemini-flash-2.0": {
quotaInfo: { remainingFraction: 0.7, resetTime: "2026-01-08T00:00:00Z" },
},
},
});
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(snapshot.provider).toBe("google-antigravity");
expect(snapshot.windows).toHaveLength(2);
expect(snapshot.error).toBeUndefined();
const proWindow = snapshot.windows.find((w) => w.label === "gemini-pro-1.5");
expect(proWindow?.usedPercent).toBe(50); // (1 - 0.5) * 100
const flashWindow = snapshot.windows.find((w) => w.label === "gemini-flash-2.0");
expect(flashWindow?.usedPercent).toBeCloseTo(30, 1); // (1 - 0.7) * 100
expect(mockFetch).toHaveBeenCalledTimes(2);
});
it("uses cloudaicompanionProject string as project id", async () => {
let capturedBody: string | undefined;
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(
async (input, init) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(200, {
availablePromptCredits: 900,
planInfo: { monthlyPromptCredits: 1000 },
cloudaicompanionProject: "projects/alpha",
});
}
if (url.includes("fetchAvailableModels")) {
capturedBody = init?.body?.toString();
return makeResponse(200, { models: {} });
}
return makeResponse(404, "not found");
},
);
await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(capturedBody).toBe(JSON.stringify({ project: "projects/alpha" }));
});
it("uses cloudaicompanionProject object id when present", async () => {
let capturedBody: string | undefined;
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(
async (input, init) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(200, {
availablePromptCredits: 900,
planInfo: { monthlyPromptCredits: 1000 },
cloudaicompanionProject: { id: "projects/beta" },
});
}
if (url.includes("fetchAvailableModels")) {
capturedBody = init?.body?.toString();
return makeResponse(200, { models: {} });
}
return makeResponse(404, "not found");
},
);
await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(capturedBody).toBe(JSON.stringify({ project: "projects/beta" }));
});
it("returns error snapshot when both endpoints fail", async () => {
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(403, { error: { message: "Access denied" } });
}
if (url.includes("fetchAvailableModels")) {
return makeResponse(403, "Forbidden");
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(snapshot.provider).toBe("google-antigravity");
expect(snapshot.windows).toHaveLength(0);
expect(snapshot.error).toBe("Access denied");
expect(mockFetch).toHaveBeenCalledTimes(2);
});
it("returns Token expired when fetchAvailableModels returns 401 and no windows", async () => {
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(500, "Boom");
}
if (url.includes("fetchAvailableModels")) {
return makeResponse(401, { error: { message: "Unauthorized" } });
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(snapshot.error).toBe("Token expired");
expect(snapshot.windows).toHaveLength(0);
});
it("extracts plan info from currentTier.name", async () => {
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(200, {
availablePromptCredits: 500,
planInfo: { monthlyPromptCredits: 1000 },
planType: "Basic",
currentTier: { id: "tier2", name: "Premium Tier" },
});
}
if (url.includes("fetchAvailableModels")) {
return makeResponse(500, "Error");
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(snapshot.plan).toBe("Premium Tier");
});
it("falls back to planType when currentTier.name is missing", async () => {
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(200, {
availablePromptCredits: 500,
planInfo: { monthlyPromptCredits: 1000 },
planType: "Basic Plan",
});
}
if (url.includes("fetchAvailableModels")) {
return makeResponse(500, "Error");
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(snapshot.plan).toBe("Basic Plan");
});
it("includes reset times in model windows", async () => {
const resetTime = "2026-01-10T12:00:00Z";
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(500, "Error");
}
if (url.includes("fetchAvailableModels")) {
return makeResponse(200, {
models: {
"gemini-pro-experimental": {
quotaInfo: { remainingFraction: 0.3, resetTime },
},
},
});
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
const proWindow = snapshot.windows.find((w) => w.label === "gemini-pro-experimental");
expect(proWindow?.resetAt).toBe(new Date(resetTime).getTime());
});
it("parses string numbers correctly", async () => {
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(200, {
availablePromptCredits: "600",
planInfo: { monthlyPromptCredits: "1000" },
});
}
if (url.includes("fetchAvailableModels")) {
return makeResponse(200, {
models: {
"gemini-flash-lite": {
quotaInfo: { remainingFraction: "0.9" },
},
},
});
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(snapshot.windows).toHaveLength(2);
const creditsWindow = snapshot.windows.find((w) => w.label === "Credits");
expect(creditsWindow?.usedPercent).toBe(40); // (1000 - 600) / 1000 * 100
const flashWindow = snapshot.windows.find((w) => w.label === "gemini-flash-lite");
expect(flashWindow?.usedPercent).toBeCloseTo(10, 1); // (1 - 0.9) * 100
});
it("skips internal models", async () => {
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(200, {
availablePromptCredits: 500,
planInfo: { monthlyPromptCredits: 1000 },
cloudaicompanionProject: "projects/internal",
});
}
if (url.includes("fetchAvailableModels")) {
return makeResponse(200, {
models: {
chat_hidden: { quotaInfo: { remainingFraction: 0.1 } },
tab_hidden: { quotaInfo: { remainingFraction: 0.2 } },
"gemini-pro-1.5": { quotaInfo: { remainingFraction: 0.7 } },
},
});
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(snapshot.windows.map((w) => w.label)).toEqual(["Credits", "gemini-pro-1.5"]);
});
it("sorts models by usage and shows individual model IDs", async () => {
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(500, "Error");
}
if (url.includes("fetchAvailableModels")) {
return makeResponse(200, {
models: {
"gemini-pro-1.0": {
quotaInfo: { remainingFraction: 0.8 },
},
"gemini-pro-1.5": {
quotaInfo: { remainingFraction: 0.3 },
},
"gemini-flash-1.5": {
quotaInfo: { remainingFraction: 0.6 },
},
"gemini-flash-2.0": {
quotaInfo: { remainingFraction: 0.9 },
},
},
});
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(snapshot.windows).toHaveLength(4);
// Should be sorted by usage (highest first)
expect(snapshot.windows[0]?.label).toBe("gemini-pro-1.5");
expect(snapshot.windows[0]?.usedPercent).toBe(70); // (1 - 0.3) * 100
expect(snapshot.windows[1]?.label).toBe("gemini-flash-1.5");
expect(snapshot.windows[1]?.usedPercent).toBe(40); // (1 - 0.6) * 100
expect(snapshot.windows[2]?.label).toBe("gemini-pro-1.0");
expect(snapshot.windows[2]?.usedPercent).toBeCloseTo(20, 1); // (1 - 0.8) * 100
expect(snapshot.windows[3]?.label).toBe("gemini-flash-2.0");
expect(snapshot.windows[3]?.usedPercent).toBeCloseTo(10, 1); // (1 - 0.9) * 100
});
it("returns Token expired error on 401 from loadCodeAssist", async () => {
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(401, { error: { message: "Unauthorized" } });
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(snapshot.error).toBe("Token expired");
expect(snapshot.windows).toHaveLength(0);
expect(mockFetch).toHaveBeenCalledTimes(1); // Should stop early on 401
});
it("handles empty models array gracefully", async () => {
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(200, {
availablePromptCredits: 800,
planInfo: { monthlyPromptCredits: 1000 },
});
}
if (url.includes("fetchAvailableModels")) {
return makeResponse(200, { models: {} });
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(snapshot.windows).toHaveLength(1);
const creditsWindow = snapshot.windows[0];
expect(creditsWindow?.label).toBe("Credits");
expect(creditsWindow?.usedPercent).toBe(20);
});
it("handles missing credits fields gracefully", async () => {
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(200, { planType: "Free" });
}
if (url.includes("fetchAvailableModels")) {
return makeResponse(200, {
models: {
"gemini-flash-experimental": {
quotaInfo: { remainingFraction: 0.5 },
},
},
});
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(snapshot.windows).toHaveLength(1);
const flashWindow = snapshot.windows[0];
expect(flashWindow?.label).toBe("gemini-flash-experimental");
expect(flashWindow?.usedPercent).toBe(50);
expect(snapshot.plan).toBe("Free");
});
it("handles invalid reset time gracefully", async () => {
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
return makeResponse(500, "Error");
}
if (url.includes("fetchAvailableModels")) {
return makeResponse(200, {
models: {
"gemini-pro-test": {
quotaInfo: { remainingFraction: 0.4, resetTime: "invalid-date" },
},
},
});
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
const proWindow = snapshot.windows.find((w) => w.label === "gemini-pro-test");
expect(proWindow?.usedPercent).toBe(60);
expect(proWindow?.resetAt).toBeUndefined();
});
it("handles network errors with graceful degradation", async () => {
const mockFetch = vi.fn<Parameters<typeof fetch>, ReturnType<typeof fetch>>(async (input) => {
const url =
typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
if (url.includes("loadCodeAssist")) {
throw new Error("Network failure");
}
if (url.includes("fetchAvailableModels")) {
return makeResponse(200, {
models: {
"gemini-flash-stable": {
quotaInfo: { remainingFraction: 0.85 },
},
},
});
}
return makeResponse(404, "not found");
});
const snapshot = await fetchAntigravityUsage("token-123", 5000, mockFetch);
expect(snapshot.windows).toHaveLength(1);
const flashWindow = snapshot.windows[0];
expect(flashWindow?.label).toBe("gemini-flash-stable");
expect(flashWindow?.usedPercent).toBeCloseTo(15, 1);
expect(snapshot.error).toBeUndefined();
});
});