Browse Source

Add antigravity quota tool

Alvin Unreal 7 months ago
parent
commit
a00097b009
6 changed files with 391 additions and 0 deletions
  1. 2 0
      src/index.ts
  2. 3 0
      src/tools/index.ts
  3. 184 0
      src/tools/quota/api.ts
  4. 49 0
      src/tools/quota/command.ts
  5. 104 0
      src/tools/quota/index.ts
  6. 49 0
      src/tools/quota/types.ts

+ 2 - 0
src/index.ts

@@ -10,6 +10,7 @@ import {
   grep,
   ast_grep_search,
   ast_grep_replace,
+  antigravity_quota,
 } from "./tools";
 import { loadPluginConfig } from "./config";
 import { createBuiltinMcps } from "./mcp";
@@ -35,6 +36,7 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
       grep,
       ast_grep_search,
       ast_grep_replace,
+      antigravity_quota,
     },
 
     mcp: mcps,

+ 3 - 0
src/tools/index.ts

@@ -12,3 +12,6 @@ export { grep } from "./grep";
 
 // AST-grep tools
 export { ast_grep_search, ast_grep_replace } from "./ast-grep";
+
+// Antigravity quota tool
+export { antigravity_quota } from "./quota";

+ 184 - 0
src/tools/quota/api.ts

@@ -0,0 +1,184 @@
+import * as path from "path";
+import * as os from "os";
+import * as fs from "fs";
+import type {
+  Account,
+  AccountsConfig,
+  TokenResponse,
+  LoadCodeAssistResponse,
+  QuotaResponse,
+  AccountQuotaResult,
+  ModelQuota,
+} from "./types";
+
+// API endpoints
+const GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token";
+const CLOUDCODE_BASE_URL = "https://cloudcode-pa.googleapis.com";
+
+// Timing constants
+const DEFAULT_RESET_MS = 86_400_000; // 24 hours - fallback when API doesn't provide reset time
+const ACCOUNT_FETCH_DELAY_MS = 200; // Delay between account fetches to avoid rate limiting
+const CLOUDCODE_METADATA = {
+  ideType: "ANTIGRAVITY",
+  platform: "PLATFORM_UNSPECIFIED",
+  pluginType: "GEMINI",
+};
+
+// Client credentials (from opencode-antigravity-auth)
+const CLIENT_ID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com";
+const CLIENT_SECRET = "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";
+
+// Config paths
+const isWindows = os.platform() === "win32";
+const configBase = isWindows
+  ? path.join(os.homedir(), "AppData", "Roaming", "opencode")
+  : path.join(os.homedir(), ".config", "opencode");
+
+const xdgData = process.env.XDG_DATA_HOME || path.join(os.homedir(), ".local", "share");
+const dataBase = isWindows ? configBase : path.join(xdgData, "opencode");
+
+export const CONFIG_PATHS = [
+  path.join(configBase, "antigravity-accounts.json"),
+  path.join(dataBase, "antigravity-accounts.json"),
+];
+
+export function loadAccountsConfig(): AccountsConfig | null {
+  for (const p of CONFIG_PATHS) {
+    if (fs.existsSync(p)) {
+      return JSON.parse(fs.readFileSync(p, "utf-8")) as AccountsConfig;
+    }
+  }
+  return null;
+}
+
+async function refreshToken(refreshToken: string): Promise<string> {
+  const params = new URLSearchParams({
+    client_id: CLIENT_ID,
+    client_secret: CLIENT_SECRET,
+    refresh_token: refreshToken,
+    grant_type: "refresh_token",
+  });
+
+  const res = await fetch(GOOGLE_TOKEN_URL, {
+    method: "POST",
+    headers: { "Content-Type": "application/x-www-form-urlencoded" },
+    body: params.toString(),
+  });
+
+  if (!res.ok) throw new Error(`Token refresh failed (${res.status})`);
+  const data = (await res.json()) as TokenResponse;
+  return data.access_token;
+}
+
+async function loadCodeAssist(accessToken: string): Promise<LoadCodeAssistResponse> {
+  const res = await fetch(`${CLOUDCODE_BASE_URL}/v1internal:loadCodeAssist`, {
+    method: "POST",
+    headers: {
+      Authorization: `Bearer ${accessToken}`,
+      "Content-Type": "application/json",
+      "User-Agent": "antigravity",
+    },
+    body: JSON.stringify({ metadata: CLOUDCODE_METADATA }),
+  });
+
+  if (!res.ok) throw new Error(`loadCodeAssist failed (${res.status})`);
+  return (await res.json()) as LoadCodeAssistResponse;
+}
+
+function extractProjectId(project: unknown): string | undefined {
+  if (typeof project === "string" && project) return project;
+  if (project && typeof project === "object" && "id" in project) {
+    const id = (project as { id?: string }).id;
+    if (id) return id;
+  }
+  return undefined;
+}
+
+async function fetchModels(accessToken: string, projectId?: string): Promise<QuotaResponse> {
+  const payload = projectId ? { project: projectId } : {};
+  const res = await fetch(`${CLOUDCODE_BASE_URL}/v1internal:fetchAvailableModels`, {
+    method: "POST",
+    headers: {
+      Authorization: `Bearer ${accessToken}`,
+      "Content-Type": "application/json",
+      "User-Agent": "antigravity",
+    },
+    body: JSON.stringify(payload),
+  });
+
+  if (!res.ok) throw new Error(`fetchModels failed (${res.status})`);
+  return (await res.json()) as QuotaResponse;
+}
+
+function formatDuration(ms: number): string {
+  const seconds = Math.floor(Math.abs(ms) / 1000);
+  const h = Math.floor(seconds / 3600);
+  const m = Math.floor((seconds % 3600) / 60);
+  if (h > 0) return `${h}h${m}m`;
+  return `${m}m`;
+}
+
+// Filter out internal/test models
+const EXCLUDED_PATTERNS = ["chat_", "rev19", "gemini 2.5", "gemini 3 pro image"];
+
+export async function fetchAccountQuota(account: Account): Promise<AccountQuotaResult> {
+  try {
+    const accessToken = await refreshToken(account.refreshToken);
+    let projectId = account.projectId || account.managedProjectId;
+
+    if (!projectId) {
+      const codeAssist = await loadCodeAssist(accessToken);
+      projectId = extractProjectId(codeAssist.cloudaicompanionProject);
+    }
+
+    const quotaRes = await fetchModels(accessToken, projectId);
+    if (!quotaRes.models) {
+      return { email: account.email, success: true, models: [] };
+    }
+
+    const now = Date.now();
+    const models: ModelQuota[] = [];
+
+    for (const [key, info] of Object.entries(quotaRes.models)) {
+      const qi = info.quotaInfo;
+      if (!qi) continue;
+
+      const label = info.displayName || key;
+      const lower = label.toLowerCase();
+      if (EXCLUDED_PATTERNS.some((p) => lower.includes(p))) continue;
+
+      const pct = Math.min(100, Math.max(0, (qi.remainingFraction ?? 0) * 100));
+      let resetMs = DEFAULT_RESET_MS;
+      if (qi.resetTime) {
+        const parsed = new Date(qi.resetTime).getTime();
+        if (!isNaN(parsed)) resetMs = Math.max(0, parsed - now);
+      }
+
+      models.push({
+        name: label,
+        percent: pct,
+        resetIn: formatDuration(resetMs),
+      });
+    }
+
+    // Sort by name
+    models.sort((a, b) => a.name.localeCompare(b.name));
+    return { email: account.email, success: true, models };
+  } catch (err) {
+    return {
+      email: account.email,
+      success: false,
+      error: err instanceof Error ? err.message : String(err),
+      models: [],
+    };
+  }
+}
+
+export async function fetchAllQuotas(accounts: Account[]): Promise<AccountQuotaResult[]> {
+  const results: AccountQuotaResult[] = [];
+  for (let i = 0; i < accounts.length; i++) {
+    if (i > 0) await new Promise((r) => setTimeout(r, ACCOUNT_FETCH_DELAY_MS));
+    results.push(await fetchAccountQuota(accounts[i]));
+  }
+  return results;
+}

+ 49 - 0
src/tools/quota/command.ts

@@ -0,0 +1,49 @@
+import * as path from "path";
+import * as os from "os";
+import * as fs from "fs";
+
+// Define base configuration directory based on OS
+const isWindows = os.platform() === "win32";
+const configBase = isWindows
+  ? path.join(os.homedir(), "AppData", "Roaming", "opencode")
+  : path.join(os.homedir(), ".config", "opencode");
+
+const commandDir = path.join(configBase, "command");
+const commandFile = path.join(commandDir, "antigravity-quota.md");
+
+const commandContent = `---
+description: Check Antigravity quota status for all configured Google accounts
+---
+
+Use the \`antigravity_quota\` tool to check the current quota status.
+
+This will show:
+- API quota remaining for each model (Gemini 3 Pro, Flash, Claude via Antigravity)
+- Per-account breakdown with compact display
+- Time until quota reset
+
+Just call the tool directly:
+\`\`\`
+antigravity_quota()
+\`\`\`
+
+IMPORTANT: Display the tool output EXACTLY as it is returned. Do not summarize, reformat, or modify the output in any way.
+`;
+
+// Try to create the command file for OpenCode context
+try {
+  if (!fs.existsSync(commandDir)) {
+    fs.mkdirSync(commandDir, { recursive: true });
+  }
+  if (!fs.existsSync(commandFile)) {
+    fs.writeFileSync(commandFile, commandContent, "utf-8");
+  } else {
+    const currentContent = fs.readFileSync(commandFile, "utf-8");
+    if (currentContent.includes("model: opencode/grok-code")) {
+      fs.writeFileSync(commandFile, commandContent, "utf-8");
+    }
+  }
+} catch (error) {
+  console.error("Failed to create command file/directory:", error);
+  // Continue execution, as this might not be fatal for the plugin's core function
+}

+ 104 - 0
src/tools/quota/index.ts

@@ -0,0 +1,104 @@
+import { tool } from "@opencode-ai/plugin";
+import { loadAccountsConfig, fetchAllQuotas, CONFIG_PATHS } from "./api";
+
+/**
+ * Compact quota display tool - groups by account, shows progress bars
+ * 
+ * Output format:
+ * ```
+ * user1    G3Pro [████████░░]  80%  2h | G3Flash [██████████] 100%  1h | Sonnet [████░░░░░░]  45%  3h
+ * user2    G3Pro [░░░░░░░░░░]   0%  4h | G3Flash [███████░░░]  72%  1h | Sonnet [██████████] 100%  2h
+ * ```
+ */
+export const antigravity_quota = tool({
+  description: "Check Antigravity API quota for all accounts (compact view with progress bars)",
+  args: {},
+  async execute() {
+    try {
+      const config = await loadAccountsConfig();
+      if (!config) {
+        return `No accounts found. Checked:\n${CONFIG_PATHS.map((p) => `  - ${p}`).join("\n")}`;
+      }
+
+      // Create accounts with default emails if missing (don't mutate original)
+      const accounts = config.accounts.map((acc, i) => ({
+        ...acc,
+        email: acc.email || `account-${i + 1}`,
+      }));
+
+      const results = await fetchAllQuotas(accounts);
+      const errors: string[] = [];
+      const lines: string[] = [];
+
+      // Find the longest email for padding
+      const maxEmailLen = Math.max(...results.map((r) => shortEmail(r.email).length), 6);
+
+      for (const result of results) {
+        if (!result.success) {
+          errors.push(`${shortEmail(result.email)}: ${result.error}`);
+          continue;
+        }
+
+        const email = shortEmail(result.email).padEnd(maxEmailLen);
+        
+        if (result.models.length === 0) {
+          lines.push(`${email}  (no models)`);
+          continue;
+        }
+
+        // Compact model display with progress bar: Name [████░░] XX% Xh
+        const modelParts = result.models.map((m) => {
+          const name = shortModelName(m.name).padEnd(7);
+          const bar = progressBar(m.percent);
+          const pct = m.percent.toFixed(0).padStart(3);
+          return `${name} ${bar} ${pct}% ${m.resetIn.padEnd(4)}`;
+        });
+
+        lines.push(`${email}  ${modelParts.join(" | ")}`);
+      }
+
+      let output = "# Quota\n```\n";
+      if (errors.length > 0) {
+        output += `Errors: ${errors.join(", ")}\n\n`;
+      }
+      output += lines.join("\n");
+      output += "\n```";
+
+      return output;
+    } catch (err) {
+      return `Error: ${err instanceof Error ? err.message : String(err)}`;
+    }
+  },
+});
+
+// ASCII progress bar
+function progressBar(percent: number): string {
+  const width = 10;
+  const filled = Math.round((percent / 100) * width);
+  const empty = width - filled;
+  return `[${"\u2588".repeat(filled)}${"\u2591".repeat(empty)}]`;
+}
+
+// Shorten email to username part
+function shortEmail(email: string): string {
+  return email.split("@")[0] ?? email;
+}
+
+// Shorten model names for compact display
+function shortModelName(name: string): string {
+  const lower = name.toLowerCase();
+  
+  // Common mappings
+  if (lower.includes("claude") && lower.includes("sonnet")) return "Sonnet";
+  if (lower.includes("claude") && lower.includes("opus")) return "Opus";
+  if (lower.includes("claude") && lower.includes("haiku")) return "Haiku";
+  if (lower.includes("claude")) return "Claude";
+  
+  if (lower.includes("gemini 3") && lower.includes("pro")) return "G3Pro";
+  if (lower.includes("gemini 3") && lower.includes("flash")) return "G3Flash";
+  if (lower.includes("gemini") && lower.includes("pro")) return "GemPro";
+  if (lower.includes("gemini") && lower.includes("flash")) return "GemFlash";
+  
+  // Fallback: take first 8 chars
+  return name.slice(0, 8);
+}

+ 49 - 0
src/tools/quota/types.ts

@@ -0,0 +1,49 @@
+export interface Account {
+  email: string;
+  refreshToken: string;
+  projectId?: string;
+  managedProjectId?: string;
+  rateLimitResetTimes: Record<string, number>;
+}
+
+export interface AccountsConfig {
+  accounts: Account[];
+  activeIndex: number;
+}
+
+export interface QuotaInfo {
+  remainingFraction?: number;
+  resetTime?: string;
+}
+
+export interface ModelInfo {
+  displayName?: string;
+  model?: string;
+  quotaInfo?: QuotaInfo;
+  recommended?: boolean;
+}
+
+export interface QuotaResponse {
+  models?: Record<string, ModelInfo>;
+}
+
+export interface TokenResponse {
+  access_token: string;
+}
+
+export interface LoadCodeAssistResponse {
+  cloudaicompanionProject?: unknown;
+}
+
+export interface ModelQuota {
+  name: string;
+  percent: number;
+  resetIn: string;
+}
+
+export interface AccountQuotaResult {
+  email: string;
+  success: boolean;
+  error?: string;
+  models: ModelQuota[];
+}