ag-quota 0.0.2

package/src/config.ts

@@ -0,0 +1,193 @@
+/*
+ * ISC License
+ * Copyright (c) 2026 Philipp
+ */
+
+import { readFile } from "node:fs/promises";
+import { homedir } from "node:os";
+import { join } from "node:path";
+
+/**
+ * Quota data source
+ * - "cloud": Fetch from Cloud Code API (requires opencode-antigravity-auth)
+ * - "local": Fetch from local language server process
+ * - "auto": Try cloud first, fallback to local
+ */
+export type QuotaSource = "cloud" | "local" | "auto";
+
+export interface QuotaIndicator {
+    threshold: number;
+    symbol: string;
+}
+
+/**
+ * Configuration options for the quota plugin
+ */
+export interface QuotaConfig {
+    /**
+     * Where to fetch quota data from.
+     * - "cloud": Use Cloud Code API (requires opencode-antigravity-auth)
+     * - "local": Use local language server process
+     * - "auto": Try cloud first, fallback to local (default)
+     * @default "auto"
+     */
+    quotaSource?: QuotaSource;
+
+    /**
+     * Format string for quota display.
+     * Available placeholders:
+     * - {category} - Category name (Flash, Pro, Claude/GPT)
+     * - {percent} - Quota percentage (e.g., "85.5")
+     * - {resetIn} - Relative time until reset (e.g., "2h 30m")
+     * - {resetAt} - Absolute reset time (e.g., "10:30 PM")
+     * - {model} - Current model ID
+     *
+     * For all quotas mode, the format is applied per category and joined with separator.
+     * @default "{category}: {percent}% ({resetIn})"
+     */
+    format?: string;
+
+    /**
+     * Separator between categories when showing all quotas
+     * @default " | "
+     */
+    separator?: string;
+
+    /**
+     * Show all quota categories or only the current model's quota
+     * @default "all"
+     */
+    displayMode?: "all" | "current";
+
+    /**
+     * Always append quota info, even when unavailable
+     * @default true
+     */
+    alwaysAppend?: boolean;
+
+    /**
+     * The marker string used to separate quota info from the message
+     * @default "> AG Quota:"
+     */
+    quotaMarker?: string;
+
+    /**
+     * Polling interval in milliseconds
+     * @default 30000 (30 seconds)
+     */
+    pollingInterval?: number;
+
+    /**
+     * Array of quota usage percentages (remaining) that trigger alerts
+     * @default [0.5, 0.1, 0.05] (50%, 10%, 5%)
+     */
+    alertThresholds?: number[];
+
+    /**
+     * Visual indicators appended to quota percentages when remaining fraction is low.
+     * The most severe matching indicator is chosen.
+     *
+     * Example: with indicators [{threshold: 0.1, symbol: "⚠️"}, {threshold: 0.05, symbol: "⛔"}]
+     * a remainingFraction of 0.04 will show "⛔".
+     *
+     * @default [{threshold: 0.1, symbol: "⚠️"}, {threshold: 0.05, symbol: "⛔"}]
+     */
+    indicators?: QuotaIndicator[];
+}
+
+const DEFAULT_CONFIG: Required<QuotaConfig> = {
+    quotaSource: "auto",
+    format: "{category}: {percent}% ({resetIn})",
+    separator: " | ",
+    displayMode: "all",
+    alwaysAppend: true,
+    quotaMarker: "> AG Quota:",
+    pollingInterval: 30000,
+    alertThresholds: [0.5, 0.1, 0.05],
+    indicators: [
+        { threshold: 0.2, symbol: "⚠️" },
+        { threshold: 0.05, symbol: "🛑" },
+    ],
+};
+
+/**
+ * Load configuration from file system.
+ * Searches in order:
+ * 1. .opencode/ag-quota.json (project-local)
+ * 2. ~/.config/opencode/ag-quota.json (user global)
+ *
+ * @param projectDir - The project directory to search from
+ * @returns Merged configuration with defaults
+ */
+export async function loadConfig(projectDir?: string): Promise<Required<QuotaConfig>> {
+    const paths: string[] = [];
+
+    // Project-local config
+    if (projectDir) {
+        paths.push(join(projectDir, ".opencode", "ag-quota.json"));
+    } else {
+        paths.push(join(process.cwd(), ".opencode", "ag-quota.json"));
+    }
+
+    // User global config
+    paths.push(join(homedir(), ".config", "opencode", "ag-quota.json"));
+
+    for (const configPath of paths) {
+        try {
+            const content = await readFile(configPath, "utf-8");
+            const userConfig = JSON.parse(content) as QuotaConfig;
+            return { ...DEFAULT_CONFIG, ...userConfig };
+        } catch {
+            // File doesn't exist or is invalid, try next
+            continue;
+        }
+    }
+
+    return DEFAULT_CONFIG;
+}
+
+/**
+ * Format a quota entry using the format string.
+ * Placeholders are only replaced if they exist in the format string.
+ */
+export function formatQuotaEntry(
+    format: string,
+    data: {
+        category: string;
+        percent: string;
+        resetIn: string | null;   // Relative: "2h 30m"
+        resetAt: string | null;   // Absolute: "10:30 PM"
+        model: string;
+    },
+): string {
+    let result = format
+        .replace("{category}", data.category)
+        .replace("{percent}", data.percent)
+        .replace("{model}", data.model);
+
+    // Handle resetIn (relative time) - e.g., "2h 30m"
+    if (format.includes("{resetIn}")) {
+        if (data.resetIn) {
+            result = result.replace("{resetIn}", data.resetIn);
+        } else {
+            result = result
+                .replace(/\s*\(\{resetIn}\)/, "")
+                .replace(/\s*\{resetIn}/, "");
+        }
+    }
+
+    // Handle resetAt (absolute time) - e.g., "10:30 PM"
+    if (format.includes("{resetAt}")) {
+        if (data.resetAt) {
+            result = result.replace("{resetAt}", data.resetAt);
+        } else {
+            result = result
+                .replace(/\s*\(\{resetAt}\)/, "")
+                .replace(/\s*\{resetAt}/, "");
+        }
+    }
+
+    return result;
+}
+
+export { DEFAULT_CONFIG };

package/src/index.test.ts

@@ -0,0 +1,96 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { fetchAntigravityStatus, type ShellRunner } from "./index.js";
+
+type HttpRequestModule = typeof import("node:http");
+type HttpsRequestModule = typeof import("node:https");
+
+type MockedHttpModule = HttpRequestModule & { request: ReturnType<typeof vi.fn> };
+type MockedHttpsModule = HttpsRequestModule & { request: ReturnType<typeof vi.fn> };
+
+vi.mock("node:http", () => ({ request: vi.fn() }));
+vi.mock("node:https", () => ({ request: vi.fn() }));
+
+describe("fetchAntigravityStatus", () => {
+    beforeEach(() => {
+        vi.resetAllMocks();
+    });
+
+    it("should fetch status correctly when discovery succeeds", async () => {
+        const http = (await import("node:http")) as unknown as MockedHttpModule;
+        const https = (await import("node:https")) as unknown as MockedHttpsModule;
+
+        const mockCsrf = "test-csrf-123";
+        const mockPort = 12345;
+
+        const shellRunner: ShellRunner = vi.fn().mockImplementation(async (cmd: string) => {
+            if (cmd.includes("ps aux")) {
+                return `user 123 0.0 0.1 1234 5678 ? Ss 12:00 0:00 /path/to/language_server --csrf_token=${mockCsrf} --extension_server_port=${mockPort}`;
+            }
+            if (cmd.includes("ss -tlnp")) {
+                return `LISTEN 0 128 127.0.0.1:${mockPort} 0.0.0.0:* users:(("language_server",pid=123,fd=4))`;
+            }
+            return "";
+        });
+
+        const mockResponse = {
+            userStatus: {
+                cascadeModelConfigData: {
+                    clientModelConfigs: [
+                        { modelName: "test-model", quotaInfo: { remainingFraction: 0.5 } },
+                    ],
+                },
+            },
+        };
+
+        // Mock https.request to fail (trigger fallback to http)
+        https.request.mockImplementationOnce((_options: unknown, _cb: unknown) => {
+            let errorHandler: (() => void) | null = null;
+            const req = {
+                on: vi.fn().mockImplementation((event: string, handler: () => void) => {
+                    if (event === "error") {
+                        errorHandler = handler;
+                    }
+                    return req;
+                }),
+                write: vi.fn(),
+                end: vi.fn().mockImplementation(() => {
+                    // Trigger error after end() to simulate connection failure
+                    if (errorHandler) errorHandler();
+                }),
+            };
+            return req as unknown as ReturnType<HttpsRequestModule["request"]>;
+        });
+
+        // Mock http.request to succeed
+        http.request.mockImplementationOnce((_options: unknown, cb: unknown) => {
+            const mockRes = {
+                on: vi.fn().mockImplementation((event: string, handler: (arg?: unknown) => void) => {
+                    if (event === "data") handler(Buffer.from(JSON.stringify(mockResponse)));
+                    if (event === "end") handler();
+                    return mockRes;
+                }),
+            };
+            if (typeof cb === "function") {
+                cb(mockRes);
+            }
+            const req = {
+                on: vi.fn().mockReturnThis(),
+                write: vi.fn().mockReturnThis(),
+                end: vi.fn().mockReturnThis(),
+            };
+            return req as unknown as ReturnType<HttpRequestModule["request"]>;
+        });
+
+        const result = await fetchAntigravityStatus(shellRunner);
+
+        expect(result.userStatus).toEqual(mockResponse.userStatus);
+        expect(shellRunner).toHaveBeenCalledWith(expect.stringContaining("ps aux"));
+    });
+
+    it("should throw error when CSRF token is not found", async () => {
+        const shellRunner: ShellRunner = vi.fn().mockResolvedValue("");
+        await expect(fetchAntigravityStatus(shellRunner)).rejects.toThrow(
+            "Antigravity CSRF token not found",
+        );
+    });
+});

package/src/index.ts

@@ -0,0 +1,379 @@
+/*
+ ISC License
+
+ Copyright (c) 2025, Cristian Militaru
+
+ Permission to use, copy, modify, and/or distribute this software for any
+ purpose with or without fee is hereby granted, provided that the above
+ copyright notice and this permission notice appear in all copies.
+
+ THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+ WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+ MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+ ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+ WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+ ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+ OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+ */
+
+import * as http from "node:http";
+import * as https from "node:https";
+
+export const API_ENDPOINTS = {
+    GET_USER_STATUS:
+        "/exa.language_server_pb.LanguageServerService/GetUserStatus",
+};
+
+export interface QuotaInfo {
+    remainingFraction: number;
+    resetTime?: string;
+}
+
+export interface ModelConfig {
+    modelName: string;
+    label?: string;
+    quotaInfo?: QuotaInfo;
+}
+
+export interface UserStatus {
+    cascadeModelConfigData?: {
+        clientModelConfigs?: ModelConfig[];
+    };
+}
+
+export interface UserStatusResponse {
+    userStatus: UserStatus;
+    timestamp: number;
+}
+
+export type ShellRunner = (cmd: string) => Promise<string>;
+
+function makeRequest<T>(
+    port: number,
+    csrfToken: string,
+    path: string,
+    body: object,
+): Promise<T> {
+    return new Promise((resolve, reject) => {
+        const payload = JSON.stringify(body);
+        const options = {
+            hostname: "127.0.0.1",
+            port,
+            path,
+            method: "POST",
+            headers: {
+                "Content-Type": "application/json",
+                "Content-Length": Buffer.byteLength(payload),
+                "X-Codeium-Csrf-Token": csrfToken,
+                "Connect-Protocol-Version": "1",
+            },
+            timeout: 2000,
+        };
+
+        const handleResponse = (response: http.IncomingMessage) => {
+            let data = "";
+            response.on("data", (chunk: Buffer) => {
+                data += chunk.toString();
+            });
+            response.on("end", () => {
+                try {
+                    resolve(JSON.parse(data) as T);
+                } catch {
+                    reject(new Error("JSON parse error"));
+                }
+            });
+        };
+
+        const req = https.request(
+            { ...options, rejectUnauthorized: false },
+            handleResponse,
+        );
+        req.on("error", () => {
+            const reqHttp = http.request(options, handleResponse);
+            reqHttp.on("error", (err) => reject(err));
+            reqHttp.write(payload);
+            reqHttp.end();
+        });
+        req.write(payload);
+        req.end();
+    });
+}
+
+export async function fetchAntigravityStatus(
+    runShell: ShellRunner,
+): Promise<UserStatusResponse> {
+    let procOutput = "";
+    try {
+        procOutput = await runShell(
+            'ps aux | grep -E "csrf_token|language_server" | grep -v grep',
+        );
+    } catch {
+        procOutput = "";
+    }
+
+    const lines = procOutput.split("\n");
+    let csrfToken = "";
+    let cmdLinePort = 0;
+
+    for (const line of lines) {
+        const csrfMatch = line.match(/--csrf_token[=\s]+([\w-]+)/i);
+        if (csrfMatch?.[1]) csrfToken = csrfMatch[1];
+        const portMatch = line.match(/--extension_server_port[=\s]+(\d+)/i);
+        if (portMatch?.[1]) cmdLinePort = parseInt(portMatch[1], 10);
+        if (csrfToken && cmdLinePort) break;
+    }
+
+    if (!csrfToken) {
+        throw new Error(
+            "Antigravity CSRF token not found. Is the Language Server running?",
+        );
+    }
+
+    let netstatOutput = "";
+    try {
+        netstatOutput = await runShell(
+            'ss -tlnp | grep -E "language_server|opencode|node"',
+        );
+    } catch {
+        netstatOutput = "";
+    }
+
+    const portMatches = netstatOutput.match(/:(\d+)/g);
+    let ports = portMatches
+        ? portMatches.map((p: string) => parseInt(p.replace(":", ""), 10))
+        : [];
+
+    if (cmdLinePort && !ports.includes(cmdLinePort)) {
+        ports.unshift(cmdLinePort);
+    }
+
+    ports = Array.from(new Set(ports));
+
+    if (ports.length === 0) {
+        throw new Error(
+            "No listening ports found for Antigravity. Check if the server is active.",
+        );
+    }
+
+    let userStatus: UserStatus | null = null;
+    let lastError: Error | null = null;
+
+    for (const p of ports) {
+        try {
+            const resp = await makeRequest<{ userStatus?: UserStatus }>(
+                p,
+                csrfToken,
+                API_ENDPOINTS.GET_USER_STATUS,
+                { metadata: { ideName: "opencode" } },
+            );
+            if (resp?.userStatus) {
+                userStatus = resp.userStatus;
+                break;
+            }
+        } catch (e) {
+            lastError = e instanceof Error ? e : new Error(String(e));
+            continue;
+        }
+    }
+
+    if (!userStatus) {
+        throw new Error(
+            `Could not communicate with Antigravity API. ${lastError?.message ?? ""}`,
+        );
+    }
+
+    return {
+        userStatus,
+        timestamp: Date.now(),
+    };
+}
+
+/**
+ * Format relative time until target date (e.g., "2h 30m")
+ */
+export function formatRelativeTime(targetDate: Date): string {
+    const now = new Date();
+    const diffMs = targetDate.getTime() - now.getTime();
+    if (diffMs <= 0) return "now";
+
+    const diffMins = Math.floor(diffMs / (1000 * 60));
+    const diffHours = Math.floor(diffMins / 60);
+    const remainingMins = diffMins % 60;
+
+    if (diffHours > 0) {
+        return `${diffHours}h ${remainingMins}m`;
+    }
+    return `${diffMins}m`;
+}
+
+/**
+ * Format absolute time of target date (e.g., "10:30 PM" or "22:30")
+ */
+export function formatAbsoluteTime(targetDate: Date): string {
+    return targetDate.toLocaleTimeString(undefined, {
+        hour: "2-digit",
+        minute: "2-digit",
+    });
+}
+
+// Re-export config utilities
+export {
+    loadConfig,
+    formatQuotaEntry,
+    DEFAULT_CONFIG,
+    type QuotaConfig,
+    type QuotaSource,
+} from "./config";
+
+// Re-export cloud utilities
+export {
+    fetchCloudQuota,
+    type CloudQuotaResult,
+    type CloudAccountInfo,
+} from "./cloud";
+
+// ============================================================================
+// Unified Quota Types and Functions
+// ============================================================================
+
+/**
+ * Unified category quota info (for the three model groups)
+ */
+export interface CategoryQuota {
+    category: "Flash" | "Pro" | "Claude/GPT";
+    remainingFraction: number;
+    resetTime: Date | null;
+}
+
+/**
+ * Unified quota result from either source
+ */
+export interface UnifiedQuotaResult {
+    source: "cloud" | "local";
+    categories: CategoryQuota[];
+    models: ModelConfig[];
+    timestamp: number;
+}
+
+/**
+ * Credentials for cloud quota fetching
+ */
+export interface CloudAuthCredentials {
+    accessToken: string;
+    projectId?: string;
+}
+
+/**
+ * Categorize a model label into one of the three groups.
+ */
+export function categorizeModel(label: string): "Flash" | "Pro" | "Claude/GPT" {
+    const lowerLabel = label.toLowerCase();
+    if (lowerLabel.includes("flash")) {
+        return "Flash";
+    }
+    if (lowerLabel.includes("gemini") || lowerLabel.includes("pro")) {
+        return "Pro";
+    }
+    return "Claude/GPT";
+}
+
+/**
+ * Group models into the three categories, taking the minimum quota per category.
+ */
+export function groupModelsByCategory(models: ModelConfig[]): CategoryQuota[] {
+    const categories: Record<
+        string,
+        { remainingFraction: number; resetTime: Date | null }
+    > = {};
+
+    for (const model of models) {
+        const label = model.label || model.modelName || "";
+        const category = categorizeModel(label);
+        const fraction = model.quotaInfo?.remainingFraction ?? 0;
+        const resetTime = model.quotaInfo?.resetTime
+            ? new Date(model.quotaInfo.resetTime)
+            : null;
+
+        if (
+            !categories[category] ||
+            fraction < categories[category].remainingFraction
+        ) {
+            categories[category] = { remainingFraction: fraction, resetTime };
+        }
+    }
+
+    const result: CategoryQuota[] = [];
+    for (const cat of ["Flash", "Pro", "Claude/GPT"] as const) {
+        if (categories[cat]) {
+            result.push({
+                category: cat,
+                remainingFraction: categories[cat].remainingFraction,
+                resetTime: categories[cat].resetTime,
+            });
+        }
+    }
+
+    return result;
+}
+
+/**
+ * Fetch quota from either cloud or local source.
+ *
+ * @param source - "cloud", "local", or "auto" (try cloud first, fallback to local)
+ * @param shellRunner - Required for local source
+ * @param cloudAuth - Required for cloud source
+ * @returns Unified quota result
+ */
+export async function fetchQuota(
+    source: "cloud" | "local" | "auto",
+    shellRunner?: ShellRunner,
+    cloudAuth?: CloudAuthCredentials,
+): Promise<UnifiedQuotaResult> {
+    // Import cloud module dynamically to avoid issues if not available
+    const { fetchCloudQuota } = await import("./cloud");
+
+    if (source === "cloud" || source === "auto") {
+        // Try cloud first
+        if (cloudAuth) {
+            try {
+                const cloudResult = await fetchCloudQuota(
+                    cloudAuth.accessToken,
+                    cloudAuth.projectId,
+                );
+                const categories = groupModelsByCategory(cloudResult.models);
+                return {
+                    source: "cloud",
+                    categories,
+                    models: cloudResult.models,
+                    timestamp: cloudResult.timestamp,
+                };
+            } catch (error) {
+                if (source === "cloud") {
+                    throw error; // Don't fallback if explicitly requested cloud
+                }
+                // Fall through to local if auto
+            }
+        } else if (source === "cloud") {
+            throw new Error(
+                "Cloud access token not provided. Cannot fetch cloud quota.",
+            );
+        }
+    }
+
+    // Try local
+    if (!shellRunner) {
+        throw new Error("Shell runner required for local quota fetching");
+    }
+
+    const localResult = await fetchAntigravityStatus(shellRunner);
+    const models: ModelConfig[] =
+        localResult.userStatus.cascadeModelConfigData?.clientModelConfigs || [];
+    const categories = groupModelsByCategory(models);
+
+    return {
+        source: "local",
+        categories,
+        models,
+        timestamp: localResult.timestamp,
+    };
+}