opencode-quotas 0.0.1

package/src/providers/antigravity/provider.ts

@@ -0,0 +1,337 @@
+import { formatRelativeTime } from "../../utils/time";
+import { getCloudCredentials } from "./auth";
+import { type IQuotaProvider, type QuotaData } from "../../interfaces";
+import { logger } from "../../logger";
+
+const CLOUDCODE_ENDPOINTS = [
+    "https://daily-cloudcode-pa.sandbox.googleapis.com",
+    "https://autopush-cloudcode-pa.sandbox.googleapis.com",
+    "https://cloudcode-pa.googleapis.com",
+] as const;
+
+const CLOUDCODE_HEADERS = {
+    "Content-Type": "application/json",
+    "User-Agent": "antigravity/1.11.5 windows/amd64",
+    "X-Goog-Api-Client": "google-cloud-sdk vscode_cloudshelleditor/0.1",
+    "Client-Metadata":
+        '{"ideType":"IDE_UNSPECIFIED","platform":"PLATFORM_UNSPECIFIED","pluginType":"GEMINI"}',
+} as const;
+
+export interface QuotaIndicator {
+    threshold: number;
+    symbol: string;
+}
+
+export interface AntigravityConfig {
+    indicators?: QuotaIndicator[];
+    debug?: boolean;
+}
+
+interface CloudQuotaInfo {
+    remainingFraction?: number;
+    resetTime?: string;
+}
+
+interface CloudModelInfo {
+    displayName?: string;
+    model?: string;
+    quotaInfo?: CloudQuotaInfo;
+    supportsImages?: boolean;
+    supportsThinking?: boolean;
+    recommended?: boolean;
+}
+
+interface FetchModelsResponse {
+    models?: Record<string, CloudModelInfo>;
+}
+
+export interface QuotaInfo {
+    remainingFraction: number;
+    resetTime?: string;
+}
+
+export interface ModelConfig {
+    modelName: string;
+    label?: string;
+    quotaInfo?: QuotaInfo;
+}
+
+export interface CloudAccountInfo {
+    email?: string;
+    projectId?: string;
+}
+
+export interface CloudQuotaResult {
+    account: CloudAccountInfo;
+    models: ModelConfig[];
+    timestamp: number;
+}
+
+const DEFAULT_INDICATORS: QuotaIndicator[] = [
+    { threshold: 0.2, symbol: "!" },
+    { threshold: 0.05, symbol: "!!" },
+];
+
+function getIndicatorSymbol(
+    fraction: number,
+    indicators: QuotaIndicator[] = DEFAULT_INDICATORS,
+): string {
+    if (indicators.length === 0) return "";
+    const sorted = [...indicators].sort((a, b) => a.threshold - b.threshold);
+    for (const indicator of sorted) {
+        if (fraction <= indicator.threshold) {
+            return ` ${indicator.symbol}`;
+        }
+    }
+    return "";
+}
+
+async function fetchAvailableModels(
+    accessToken: string,
+    projectId: string | undefined,
+    debugEnabled: boolean,
+): Promise<FetchModelsResponse> {
+    const payload = projectId ? { project: projectId } : {};
+    let lastError: Error | null = null;
+
+    const headers: Record<string, string> = {
+        ...CLOUDCODE_HEADERS,
+        Authorization: `Bearer ${accessToken}`,
+    };
+
+    for (const endpoint of CLOUDCODE_ENDPOINTS) {
+        try {
+            const url = `${endpoint}/v1internal:fetchAvailableModels`;
+            const controller = new AbortController();
+            const timeoutId = setTimeout(() => controller.abort(), 10000);
+
+            try {
+                if (debugEnabled) {
+                    logger.debug("antigravity:request", {
+                        endpoint,
+                        url,
+                        hasProjectId: !!projectId,
+                    });
+                }
+
+                const response = await fetch(url, {
+                    method: "POST",
+                    headers,
+                    body: JSON.stringify(payload),
+                    signal: controller.signal,
+                });
+
+                if (debugEnabled) {
+                    logger.debug("antigravity:response_meta", {
+                        endpoint,
+                        status: response.status,
+                        ok: response.ok,
+                        contentType: response.headers.get("content-type"),
+                    });
+                }
+
+                if (response.status === 401) {
+                    throw new Error("Authorization expired or invalid.");
+                }
+
+                if (response.status === 403) {
+                    throw new Error(
+                        "Access forbidden (403). Check your account permissions.",
+                    );
+                }
+
+                if (!response.ok) {
+                    const text = await response.text();
+                    if (debugEnabled) {
+                        logger.debug("antigravity:error_body", {
+                            endpoint,
+                            status: response.status,
+                            bodyPreview: text.slice(0, 2000),
+                        });
+                    }
+
+                    throw new Error(
+                        `Cloud Code API error ${response.status}: ${text.slice(0, 200)}`,
+                    );
+                }
+
+                const json = (await response.json()) as FetchModelsResponse;
+
+                const modelCount = Object.keys(json.models || {}).length;
+                if (debugEnabled) {
+                    logger.debug("antigravity:fetch_success", { modelCount });
+                }
+
+                if (debugEnabled) {
+                    const sampleKeys = Object.keys(json.models || {}).slice(
+                        0,
+                        8,
+                    );
+                    const sanitizedSample: Record<string, unknown> = {};
+                    for (const key of sampleKeys) {
+                        sanitizedSample[key] = json.models?.[key];
+                    }
+
+                    logger.debug("antigravity:raw_response_sample", {
+                        modelCount,
+                        sampleKeys,
+                        sample: sanitizedSample,
+                    });
+                }
+
+                return json;
+            } finally {
+                clearTimeout(timeoutId);
+            }
+        } catch (error) {
+            lastError =
+                error instanceof Error ? error : new Error(String(error));
+            if (
+                lastError.message.includes("Authorization") ||
+                lastError.message.includes("forbidden") ||
+                lastError.message.includes("invalid_grant")
+            ) {
+                throw lastError;
+            }
+        }
+    }
+
+    throw lastError || new Error("All Cloud Code API endpoints failed");
+}
+
+export async function fetchCloudQuota(
+    accessToken: string,
+    projectId?: string,
+    debugEnabled: boolean = false,
+): Promise<CloudQuotaResult> {
+    if (!accessToken) {
+        throw new Error("Access token is required for cloud quota fetching");
+    }
+
+    const response = await fetchAvailableModels(
+        accessToken,
+        projectId,
+        debugEnabled,
+    );
+
+    const models: ModelConfig[] = [];
+
+    if (response.models) {
+        for (const [modelKey, info] of Object.entries(response.models)) {
+            if (!info.quotaInfo) continue;
+
+            models.push({
+                modelName: info.model || modelKey,
+                label: info.displayName || modelKey,
+                quotaInfo: {
+                    remainingFraction: info.quotaInfo.remainingFraction ?? 0,
+                    resetTime: info.quotaInfo.resetTime,
+                },
+            });
+        }
+    }
+
+    return {
+        account: {
+            projectId,
+        },
+        models,
+        timestamp: Date.now(),
+    };
+}
+
+/**
+ * Creates the Antigravity provider that returns flat, raw quota data.
+ * Grouping and aggregation is handled by the service layer via AggregatedGroups.
+ */
+export function createAntigravityProvider(
+    config: AntigravityConfig = {},
+): IQuotaProvider {
+    return {
+        id: "antigravity",
+        async fetchQuota(): Promise<QuotaData[]> {
+            const debugEnabled = !!config.debug;
+            logger.debug("provider:antigravity:fetch_start", {
+                configDebug: config.debug,
+            });
+
+            // Fetch cloud credentials (Google OAuth)
+            let credentials;
+            try {
+                credentials = await getCloudCredentials();
+                if (debugEnabled) {
+                    logger.debug("provider:antigravity:auth_ok", {
+                        projectId: credentials.projectId,
+                    });
+                }
+            } catch (e) {
+                logger.error("provider:antigravity:auth_failed", e);
+                throw e;
+            }
+
+            // Fetch live quota from Antigravity Cloud API
+            const cloudResult = await fetchCloudQuota(
+                credentials.accessToken,
+                credentials.projectId,
+                debugEnabled,
+            );
+            if (debugEnabled) {
+                logger.debug("provider:antigravity:cloud_ok", {
+                    modelCount: cloudResult.models.length,
+                });
+            }
+
+            // Return flat list of all models with quota info
+            const entries: QuotaData[] = [];
+
+            for (const model of cloudResult.models) {
+                if (
+                    !model.quotaInfo ||
+                    typeof model.quotaInfo.remainingFraction !== "number"
+                ) {
+                    continue;
+                }
+
+                const label = model.label || model.modelName || "";
+                const remainingFraction = model.quotaInfo.remainingFraction;
+                const usedPercent = Math.max(
+                    0,
+                    Math.min(100, (1 - remainingFraction) * 100),
+                );
+
+                // Generate stable raw ID from model name
+                const rawId = `ag-raw-${label
+                    .toLowerCase()
+                    .replace(/[^a-z0-9]+/g, "-")
+                    .replace(/^-|-$/g, "")}`;
+
+                const indicator = getIndicatorSymbol(
+                    remainingFraction,
+                    config.indicators,
+                );
+
+                let reset: string | undefined;
+                if (model.quotaInfo.resetTime) {
+                    reset = `resets in ${formatRelativeTime(new Date(model.quotaInfo.resetTime))}`;
+                }
+
+                entries.push({
+                    id: rawId,
+                    providerName: `Antigravity ${label}`,
+                    used: usedPercent,
+                    limit: 100,
+                    unit: "%",
+                    reset,
+                    info: indicator.trim() || undefined,
+                });
+            }
+
+            if (debugEnabled) {
+                logger.debug("provider:antigravity:fetch_ok", {
+                    count: entries.length,
+                });
+            }
+            return entries;
+        },
+    };
+}

package/src/providers/codex.ts

@@ -0,0 +1,327 @@
+import { readFile } from "node:fs/promises";
+import { AUTH_FILE } from "../utils/paths";
+import { type IQuotaProvider, type QuotaData } from "../interfaces";
+import { logger } from "../logger";
+
+const AUTH_PATH = AUTH_FILE();
+const DEFAULT_BASE_URL = "https://chatgpt.com/backend-api";
+const REQUEST_TIMEOUT_MS = 15_000;
+const MAX_ERROR_BODY_CHARS = 2_000;
+
+type OauthAuth = {
+  type: "oauth";
+  access: string;
+  refresh: string;
+  expires: number;
+  enterpriseUrl?: string;
+};
+
+type ApiAuth = {
+  type: "api";
+  key: string;
+};
+
+type WellKnownAuth = {
+  type: "wellknown";
+  key: string;
+  token: string;
+};
+
+type AuthInfo = OauthAuth | ApiAuth | WellKnownAuth;
+
+type AuthFile = Record<string, AuthInfo>;
+
+type OauthSelection = {
+  providerID: string;
+  access: string;
+  enterpriseUrl?: string;
+};
+
+type RateLimitWindowSnapshot = {
+  used_percent?: number;
+  limit_window_seconds?: number;
+  reset_after_seconds?: number;
+  reset_at?: number;
+};
+
+type RateLimitStatusDetails = {
+  primary_window?: RateLimitWindowSnapshot | null;
+  secondary_window?: RateLimitWindowSnapshot | null;
+};
+
+type CreditStatusDetails = {
+  unlimited?: boolean;
+  balance?: string | null;
+};
+
+type RateLimitStatusPayload = {
+  plan_type?: string;
+  rate_limit?: RateLimitStatusDetails | null;
+  credits?: CreditStatusDetails | null;
+};
+
+async function readAuthFile(): Promise<AuthFile | null> {
+  try {
+    const raw = await readFile(AUTH_PATH, "utf8");
+    const parsed = JSON.parse(raw) as AuthFile;
+    return parsed;
+  } catch (e) {
+    logger.debug("provider:codex:auth_read_failed", { authPath: AUTH_PATH, error: e });
+    // Fallback: try the config directory location
+    try {
+      const configPath = AUTH_PATH.replace("auth.json", "antigravity-accounts.json");
+      // If the replacement is not valid, this will likely fail and we return null
+      const rawConfig = await readFile(configPath, "utf8");
+      return JSON.parse(rawConfig) as AuthFile;
+    } catch {
+      return null;
+    }
+  }
+}
+
+function pickOauthAuth(auth: AuthFile): OauthSelection | null {
+  const preferred = ["opencode", "codex", "openai"];
+  for (const providerID of preferred) {
+    const info = auth[providerID];
+    if (info?.type === "oauth") {
+      return {
+        providerID,
+        access: info.access,
+        enterpriseUrl: info.enterpriseUrl,
+      };
+    }
+  }
+
+  for (const [providerID, info] of Object.entries(auth)) {
+    if (info.type === "oauth") {
+      return {
+        providerID,
+        access: info.access,
+        enterpriseUrl: info.enterpriseUrl,
+      };
+    }
+  }
+
+  return null;
+}
+
+function buildUsageUrl(baseUrl: string): string {
+  const trimmed = baseUrl.replace(/\/+$/, "");
+  if (trimmed.includes("/backend-api")) {
+    return `${trimmed}/wham/usage`;
+  }
+  return `${trimmed}/api/codex/usage`;
+}
+
+async function fetchQuotaPayload(
+  accessToken: string,
+  baseUrl: string,
+): Promise<unknown> {
+  const url = buildUsageUrl(baseUrl);
+  const controller = new AbortController();
+  const timeout = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS);
+  try {
+    const response = await fetch(url, {
+      headers: {
+        Authorization: `Bearer ${accessToken}`,
+      },
+      signal: controller.signal,
+    });
+
+    const bodyText = await response.text();
+    let payload: unknown = null;
+    try {
+      payload = JSON.parse(bodyText);
+    } catch {
+      payload = bodyText;
+    }
+
+    if (!response.ok) {
+      const error = new Error(`quota request failed (${response.status})`);
+      error.cause = {
+        status: response.status,
+        bodyText: bodyText.slice(0, MAX_ERROR_BODY_CHARS),
+      };
+      throw error;
+    }
+
+    return payload;
+  } finally {
+    clearTimeout(timeout);
+  }
+}
+
+function isObject(value: unknown): value is Record<string, unknown> {
+  return Boolean(value) && typeof value === "object" && !Array.isArray(value);
+}
+
+function toNumber(value: unknown): number | null {
+  if (typeof value === "number" && Number.isFinite(value)) {
+    return value;
+  }
+  if (typeof value === "string") {
+    const parsed = Number.parseFloat(value);
+    if (Number.isFinite(parsed)) return parsed;
+  }
+  return null;
+}
+
+function describeWindow(windowSeconds: number | null): string | null {
+  if (!windowSeconds || windowSeconds <= 0) return null;
+  const minutes = Math.round(windowSeconds / 60);
+  if (minutes >= 60 && minutes % 60 === 0) {
+    const hours = minutes / 60;
+    return `${hours}h window`;
+  }
+  return `${minutes}m window`;
+}
+
+function formatRelativeSeconds(seconds: number): string {
+  if (seconds <= 0) return "now";
+  const minutes = Math.floor(seconds / 60);
+  const hours = Math.floor(minutes / 60);
+  const remainingMinutes = minutes % 60;
+  if (hours > 0) {
+    return `${hours}h ${remainingMinutes}m`;
+  }
+  return `${minutes}m`;
+}
+
+function parseRateLimitWindow(
+  id: string,
+  label: string,
+  snapshot: RateLimitWindowSnapshot,
+): QuotaData | null {
+  const usedPercent = toNumber(snapshot.used_percent);
+  if (usedPercent === null) return null;
+
+  // Window info
+  let window: string | undefined;
+  const windowSeconds = toNumber(snapshot.limit_window_seconds);
+  const windowLabel = describeWindow(windowSeconds);
+  if (windowLabel) window = windowLabel;
+
+  // Reset info
+  let reset: string | undefined;
+  const resetAfter = toNumber(snapshot.reset_after_seconds);
+  const resetAt = toNumber(snapshot.reset_at);
+  if (resetAfter !== null) {
+    reset = `resets in ${formatRelativeSeconds(resetAfter)}`;
+  } else if (resetAt !== null) {
+    reset = `resets at ${new Date(resetAt * 1000).toLocaleTimeString()}`;
+  }
+
+  return {
+    id: `codex-${id}`,
+    providerName: `Codex ${label}`,
+    used: Math.max(0, Math.min(100, usedPercent)),
+    limit: 100,
+    unit: "%",
+    window,
+    reset,
+  };
+}
+
+function parseCredits(credits: CreditStatusDetails): QuotaData | null {
+  const base = {
+    id: "codex-credits",
+    providerName: "Codex Credits",
+    unit: "credits",
+  };
+
+  if (credits.unlimited) {
+    return {
+      ...base,
+      used: 0,
+      limit: null,
+      info: "unlimited",
+    };
+  }
+
+  const balance = toNumber(credits.balance ?? null);
+  if (balance === null) return null;
+
+  return {
+    ...base,
+    used: balance,
+    limit: null,
+    info: "balance",
+  };
+}
+
+export function extractCodexQuota(payload: unknown): QuotaData[] {
+  if (!isObject(payload)) return [];
+
+  const rateLimitCandidate = (payload as Record<string, unknown>)["rate_limit"];
+  const rateLimit = isObject(rateLimitCandidate) ? rateLimitCandidate : null;
+
+  const entries: QuotaData[] = [];
+
+  if (rateLimit) {
+    const primary = isObject(rateLimit.primary_window)
+      ? (rateLimit.primary_window as RateLimitWindowSnapshot)
+      : null;
+    const secondary = isObject(rateLimit.secondary_window)
+      ? (rateLimit.secondary_window as RateLimitWindowSnapshot)
+      : null;
+
+    if (primary) {
+      const entry = parseRateLimitWindow("primary", "Primary", primary);
+      if (entry) entries.push(entry);
+    }
+
+    if (secondary) {
+      const entry = parseRateLimitWindow("secondary", "Secondary", secondary);
+      if (entry) entries.push(entry);
+    }
+  }
+
+  const creditsCandidate = (payload as Record<string, unknown>)["credits"];
+  const credits = isObject(creditsCandidate) ? creditsCandidate : null;
+
+  if (credits) {
+    const creditEntry = parseCredits(credits as CreditStatusDetails);
+    if (creditEntry) entries.push(creditEntry);
+  }
+
+  return entries;
+}
+
+export function createCodexProvider(): IQuotaProvider {
+  return {
+    id: "codex",
+    async fetchQuota(): Promise<QuotaData[]> {
+      logger.debug("provider:codex:fetch_start", { authPath: AUTH_PATH });
+
+      const auth = await readAuthFile();
+      if (!auth) {
+        logger.debug("provider:codex:no_auth", { authPath: AUTH_PATH });
+        throw new Error("Codex auth.json not found");
+      }
+
+      const oauth = pickOauthAuth(auth);
+      if (!oauth) {
+        logger.debug("provider:codex:no_oauth", { availableProviders: Object.keys(auth) });
+        throw new Error("Codex OAuth credentials missing");
+      }
+
+      const baseUrl =
+        process.env.OPENCODE_CODEX_BASE_URL ??
+        oauth.enterpriseUrl ??
+        DEFAULT_BASE_URL;
+
+        logger.debug("provider:codex:request", { providerID: oauth.providerID, baseUrl, url: buildUsageUrl(baseUrl) });
+
+      const payload = await fetchQuotaPayload(oauth.access, baseUrl);
+      const entries = extractCodexQuota(payload);
+
+      logger.debug("provider:codex:parse", { count: entries.length });
+
+      if (entries.length === 0) {
+        throw new Error("Codex quota payload did not include rate limits");
+      }
+
+      return entries;
+    },
+  };
+}