opencode-quotas 0.0.1

package/src/services/prediction-engine.ts

@@ -0,0 +1,133 @@
+import { type IPredictionEngine, type IHistoryService, type HistoryPoint } from "../interfaces";
+
+/**
+ * Configuration options for the prediction engine.
+ */
+export interface PredictionEngineConfig {
+    /**
+     * Default short time window for regression to capture spikes (minutes).
+     * Defaults to 5.
+     */
+    predictionShortWindowMinutes?: number;
+    
+    /**
+     * Idle timeout in milliseconds. If the last data point is older than this,
+     * the prediction returns Infinity. Defaults to 5 minutes.
+     */
+    idleTimeoutMs?: number;
+}
+
+/**
+ * Prediction engine using dual-window linear regression.
+ * 
+ * This implementation uses two time windows to calculate usage slopes:
+ * - Long window: Captures overall trend (default: 60 minutes)
+ * - Short window: Captures recent spikes (default: 5 minutes)
+ * 
+ * The maximum of the two slopes is used for conservative estimation.
+ */
+export class LinearRegressionPredictionEngine implements IPredictionEngine {
+    private readonly historyService: IHistoryService;
+    private readonly config: Required<PredictionEngineConfig>;
+
+    constructor(historyService: IHistoryService, config?: PredictionEngineConfig) {
+        this.historyService = historyService;
+        this.config = {
+            predictionShortWindowMinutes: config?.predictionShortWindowMinutes ?? 5,
+            idleTimeoutMs: config?.idleTimeoutMs ?? 5 * 60 * 1000,
+        };
+    }
+
+    /**
+     * Predicts time to limit in milliseconds using a dual-window linear regression approach.
+     * Returns Infinity if usage is decreasing, stable, or idle.
+     */
+    predictTimeToLimit(
+        quotaId: string, 
+        windowMinutes: number = 60, 
+        shortWindowMinutes?: number
+    ): number {
+        const longWindowMs = windowMinutes * 60 * 1000;
+        const shortWindowMin = shortWindowMinutes ?? this.config.predictionShortWindowMinutes;
+        const shortWindowMs = shortWindowMin * 60 * 1000;
+
+        const history = this.historyService.getHistory(quotaId, longWindowMs);
+        if (history.length < 2) return Infinity;
+
+        // Idle Handling: If the last history point is older than the idle timeout, 
+        // assume usage has stopped.
+        const lastPoint = history[history.length - 1];
+        const now = Date.now();
+        if (now - lastPoint.timestamp > this.config.idleTimeoutMs) {
+            return Infinity;
+        }
+
+        // Long Slope
+        const mLong = this.calculateSlope(history);
+
+        // Short Slope: most recent data in short window or last 15% of points
+        const shortHistory = history.filter(p => p.timestamp > now - shortWindowMs);
+        
+        // Ensure we have enough points in short history, or take the last 15%
+        let effectiveShortHistory = shortHistory;
+        if (effectiveShortHistory.length < 2) {
+            const fifteenPercentCount = Math.max(2, Math.ceil(history.length * 0.15));
+            effectiveShortHistory = history.slice(-fifteenPercentCount);
+        }
+
+        const mShort = this.calculateSlope(effectiveShortHistory);
+
+        // Conservative Estimation: use the maximum slope
+        const m = Math.max(mLong, mShort);
+        
+        if (m <= 0) return Infinity;
+        if (lastPoint.limit === null) return Infinity;
+
+        const remaining = lastPoint.limit - lastPoint.used;
+        if (remaining <= 0) return 0;
+
+        const msFromLastPoint = remaining / m;
+        const elapsedSinceLastPoint = now - lastPoint.timestamp;
+        
+        return Math.max(0, msFromLastPoint - elapsedSinceLastPoint);
+    }
+
+    /**
+     * Calculates the slope (usage per ms) using linear regression for the given history points.
+     */
+    calculateSlope(history: HistoryPoint[]): number {
+        if (history.length < 2) return 0;
+
+        let sumX = 0, sumY = 0, sumXY = 0, sumX2 = 0;
+        const n = history.length;
+        const firstTimestamp = history[0].timestamp;
+
+        for (const p of history) {
+            const x = p.timestamp - firstTimestamp;
+            const y = p.used;
+            sumX += x;
+            sumY += y;
+            sumXY += x * y;
+            sumX2 += x * x;
+        }
+
+        const denominator = (n * sumX2 - sumX * sumX);
+        if (denominator === 0) return 0;
+
+        return (n * sumXY - sumX * sumY) / denominator;
+    }
+}
+
+/**
+ * A null prediction engine that always returns Infinity.
+ * Used when no history service is available.
+ */
+export class NullPredictionEngine implements IPredictionEngine {
+    predictTimeToLimit(
+        _quotaId: string, 
+        _windowMinutes: number = 60, 
+        _shortWindowMinutes?: number
+    ): number {
+        return Infinity;
+    }
+}

package/src/services/quota-service.ts

@@ -0,0 +1,343 @@
+import { DEFAULT_CONFIG } from "../defaults";
+import { 
+    type QuotaConfig, 
+    type QuotaData, 
+    type IQuotaProvider, 
+    type IHistoryService,
+    type IPredictionEngine,
+    type IAggregationService,
+    type AggregatedGroup
+} from "../interfaces";
+import { getQuotaRegistry } from "../registry";
+import { createAntigravityProvider } from "../providers/antigravity";
+import { createCodexProvider } from "../providers/codex";
+import { formatDurationMs } from "../utils/time";
+import { logger } from "../logger";
+import { LinearRegressionPredictionEngine, NullPredictionEngine } from "./prediction-engine";
+import { AggregationService } from "./aggregation-service";
+import { ConfigLoader } from "./config-loader";
+
+export class QuotaService {
+    private config: QuotaConfig;
+    private initialized: boolean = false;
+    private initPromise: Promise<void> | null = null;
+    private historyService?: IHistoryService;
+    private predictionEngine: IPredictionEngine;
+    private aggregationService: IAggregationService;
+
+    constructor(initialConfig?: Partial<QuotaConfig>) {
+        this.config = ConfigLoader.createConfig(initialConfig);
+        
+        // Initialize with null prediction engine until init() is called
+        this.predictionEngine = new NullPredictionEngine();
+        this.aggregationService = new AggregationService(this.predictionEngine);
+    }
+
+    async init(directory: string, historyService?: IHistoryService): Promise<void> {
+        if (this.initialized) return;
+        if (this.initPromise) return this.initPromise;
+
+        this.initPromise = (async () => {
+            this.historyService = historyService;
+
+            // Load config from disk
+            this.config = await ConfigLoader.loadFromDisk(directory, this.config);
+
+            if (this.historyService && this.config.historyMaxAgeHours !== undefined) {
+                this.historyService.setMaxAge(this.config.historyMaxAgeHours);
+            }
+
+            // Initialize prediction engine with history service
+            if (this.historyService) {
+                this.predictionEngine = new LinearRegressionPredictionEngine(
+                    this.historyService,
+                    { predictionShortWindowMinutes: this.config.predictionShortWindowMinutes }
+                );
+            }
+            // Re-initialize aggregation service with the new prediction engine
+            this.aggregationService = new AggregationService(this.predictionEngine);
+
+            // Register providers
+            await this.registerProviders();
+
+            this.initialized = true;
+        })();
+        
+        return this.initPromise;
+    }
+
+    private async registerProviders(): Promise<void> {
+        const registry = getQuotaRegistry();
+
+        // Register Antigravity
+        try {
+            registry.register(
+                createAntigravityProvider({
+                    debug: !!this.config.debug,
+                }),
+            );
+            logger.debug("init:provider_registered", { id: "antigravity" });
+        } catch (e) {
+            logger.error("init:provider_failed", { id: "antigravity", error: e });
+            console.warn("[QuotaService] Failed to initialize Antigravity provider:", e);
+        }
+
+        // Register Codex
+        try {
+            registry.register(createCodexProvider());
+            logger.debug("init:provider_registered", { id: "codex" });
+        } catch (e) {
+            logger.error("init:provider_failed", { id: "codex", error: e });
+            console.warn("[QuotaService] Failed to initialize Codex provider:", e);
+        }
+    }
+
+    getConfig(): QuotaConfig {
+        return this.config;
+    }
+
+    getProviders(): IQuotaProvider[] {
+        return getQuotaRegistry().getAll();
+    }
+
+    /**
+     * Returns the prediction engine used by this service.
+     * Useful for testing or for other services that need prediction capabilities.
+     */
+    getPredictionEngine(): IPredictionEngine {
+        return this.predictionEngine;
+    }
+
+    /**
+     * Returns the aggregation service used by this service.
+     * Useful for testing or for other services that need aggregation capabilities.
+     */
+    getAggregationService(): IAggregationService {
+        return this.aggregationService;
+    }
+
+    async getQuotas(context?: { providerId?: string; modelId?: string }): Promise<QuotaData[]> {
+        const providers = this.getProviders();
+
+        logger.debug(
+            "quota_service:get_quotas_start",
+            { providerCount: providers.length, ids: providers.map((p) => p.id) },
+        );
+        
+        if (providers.length === 0) return [];
+
+        const results = await Promise.all(
+            providers.map(async (p: IQuotaProvider) => {
+                const startedAt = Date.now();
+                try {
+                    logger.debug(
+                        "quota_service:provider_fetch_start",
+                        { id: p.id },
+                    );
+                    const result = await p.fetchQuota();
+                    logger.debug(
+                        "quota_service:provider_fetch_ok",
+                        { id: p.id, count: result.length, durationMs: Date.now() - startedAt },
+                    );
+                    return result;
+                } catch (e) {
+                    logger.error(
+                        "quota_service:provider_fetch_error",
+                        { id: p.id, durationMs: Date.now() - startedAt, error: e },
+                    );
+                    console.error(`Provider ${p.id} failed:`, e);
+                    return [];
+                }
+            })
+        );
+
+        const processed = this.processQuotas(results.flat(), context);
+        logger.debug(
+            "quota_service:get_quotas_end",
+            { totalCount: processed.length },
+        );
+        return processed;
+    }
+
+    processQuotas(data: QuotaData[], context?: { providerId?: string; modelId?: string }): QuotaData[] {
+        let results = [...data];
+
+        // 1. Enrich with predictions (before aggregation so sources have it too)
+        results = results.map(q => {
+            const time = this.predictionEngine.predictTimeToLimit(q.id, 60);
+            if (time !== Infinity) {
+                return {
+                    ...q,
+                    predictedReset: `${formatDurationMs(time)} (predicted)`
+                };
+            }
+            return q;
+        });
+
+        // 2. Apply Aggregation
+        results = this.applyAggregation(results);
+
+        // 3. Filter (Disabled & Model Mapping). If requested, perform model-strict filtering.
+        results = this.filterQuotas(results, context);
+
+        // 4. Sort
+        results = this.sortQuotas(results);
+
+        return results;
+    }
+
+    private applyAggregation(quotas: QuotaData[]): QuotaData[] {
+        if (!this.config.aggregatedGroups || this.config.aggregatedGroups.length === 0) {
+            return quotas;
+        }
+
+        const aggregatedResults: QuotaData[] = [];
+        let remainingQuotas = [...quotas];
+
+        for (const group of this.config.aggregatedGroups) {
+            // Resolve source quotas from explicit sources and patterns
+            const sourceQuotas = this.resolveGroupSources(remainingQuotas, group);
+            if (sourceQuotas.length === 0) continue;
+
+            const strategy = group.strategy || "most_critical";
+            let representative: QuotaData | null = null;
+
+            if (strategy === "most_critical") {
+                representative = this.aggregationService.aggregateMostCritical(
+                    sourceQuotas, 
+                    group.predictionWindowMinutes,
+                    group.predictionShortWindowMinutes
+                );
+            } else if (strategy === "max") {
+                representative = this.aggregationService.aggregateMax(sourceQuotas);
+            } else if (strategy === "min") {
+                representative = this.aggregationService.aggregateMin(sourceQuotas);
+            } else if (strategy === "mean" || strategy === "median") {
+                representative = this.aggregationService.aggregateAverage(
+                    sourceQuotas, 
+                    group.name, 
+                    group.id, 
+                    strategy
+                );
+            }
+
+            if (representative) {
+                // Create a copy for display
+                const displayQuota = { 
+                    ...representative, 
+                    id: group.id, 
+                    providerName: group.name 
+                };
+
+                // Remove matched sources from pool to avoid double aggregation
+                const sourceIds = new Set(sourceQuotas.map(q => q.id));
+                remainingQuotas = remainingQuotas.filter(q => !sourceIds.has(q.id));
+                
+                aggregatedResults.push(displayQuota);
+            }
+        }
+
+        // Return aggregated results. 
+        // If showUnaggregated is false, only return what matched a group.
+        if (this.config.showUnaggregated === false) {
+            return aggregatedResults;
+        }
+
+        return [...remainingQuotas, ...aggregatedResults];
+    }
+
+    /**
+     * Resolves which quotas belong to an AggregatedGroup using explicit sources and patterns.
+     */
+    private resolveGroupSources(quotas: QuotaData[], group: AggregatedGroup): QuotaData[] {
+        const matched: QuotaData[] = [];
+        const matchedIds = new Set<string>();
+
+        // 1. Explicit sources (highest priority)
+        if (group.sources && group.sources.length > 0) {
+            for (const quota of quotas) {
+                if (group.sources.includes(quota.id)) {
+                    matched.push(quota);
+                    matchedIds.add(quota.id);
+                }
+            }
+        }
+
+        // 2. Pattern matching
+        if (group.patterns && group.patterns.length > 0) {
+            for (const quota of quotas) {
+                // Skip if already matched by explicit source
+                if (matchedIds.has(quota.id)) continue;
+
+                // Filter by providerId if specified
+                if (group.providerId) {
+                    const providerMatch = 
+                        quota.providerName.toLowerCase().includes(group.providerId.toLowerCase()) ||
+                        quota.id.toLowerCase().startsWith(group.providerId.toLowerCase());
+                    if (!providerMatch) continue;
+                }
+
+                // Check if any pattern matches
+                const matchTarget = `${quota.id} ${quota.providerName}`.toLowerCase();
+                const patternMatches = group.patterns.some(pattern => {
+                    const lowerPattern = pattern.toLowerCase();
+                    return matchTarget.includes(lowerPattern);
+                });
+
+                if (patternMatches) {
+                    matched.push(quota);
+                    matchedIds.add(quota.id);
+                }
+            }
+        }
+
+        return matched;
+    }
+
+    private filterQuotas(quotas: QuotaData[], context?: { providerId?: string; modelId?: string }): QuotaData[] {
+        let results = [...quotas];
+
+        // Filter out disabled quotas
+        const disabledIds = new Set(this.config.disabled || []);
+        results = results.filter((data) => !disabledIds.has(data.id));
+
+        // If requested, apply model-aware filtering
+        if (this.config.filterByCurrentModel && context && context.providerId && context.modelId) {
+            return this.filterByModel(results, context.providerId, context.modelId);
+        }
+
+        return results;
+    }
+
+    private filterByModel(quotas: QuotaData[], providerId: string, modelId: string): QuotaData[] {
+        const providerLower = providerId.toLowerCase();
+        const modelIdLower = modelId.toLowerCase();
+
+        // Fuzzy token match with scoring
+        const tokens = modelIdLower.split(/[^a-z0-9]+/).filter(Boolean);
+        const scoredMatches = quotas
+            .map(q => {
+                const id = q.id.toLowerCase();
+                const name = q.providerName.toLowerCase();
+                const score = tokens.reduce((acc, t) => acc + (id.includes(t) || name.includes(t) ? 1 : 0), 0);
+                return { q, score };
+            })
+            .filter(m => m.score > 0)
+            .sort((a, b) => b.score - a.score);
+
+        if (scoredMatches.length > 0) {
+            const maxScore = scoredMatches[0].score;
+            return scoredMatches.filter(m => m.score === maxScore).map(m => m.q);
+        }
+
+        // 3) Provider fallback
+        const matchesProvider = quotas.filter(q => 
+            q.providerName.toLowerCase().includes(providerLower)
+        );
+        return matchesProvider.length > 0 ? matchesProvider : [];
+    }
+
+    private sortQuotas(quotas: QuotaData[]): QuotaData[] {
+        return quotas.sort((a, b) => a.providerName.localeCompare(b.providerName));
+    }
+}

package/src/tools/quotas.ts

@@ -0,0 +1,78 @@
+import { tool, type ToolDefinition } from "@opencode-ai/plugin";
+import { type QuotaService } from "../services/quota-service";
+import { type QuotaConfig } from "../interfaces";
+import { renderQuotaTable } from "../ui/quota-table";
+import { logger } from "../logger";
+
+/**
+ * Creates a tool definition for fetching and displaying quota information.
+ * This tool always performs a fresh fetch of quotas from all providers.
+ * 
+ * @returns A ToolDefinition that can be registered under Hooks.tool
+ */
+export function createQuotaTool(
+    quotaService: QuotaService,
+    getConfig: () => QuotaConfig
+): ToolDefinition {
+    return tool({
+        description:
+            "Fetch and display the current API quota usage across all configured providers (Antigravity, Codex, etc.). " +
+            "Always returns fresh, up-to-date quota information. " +
+            "Use this tool when you need to check remaining API capacity, quota limits, or usage statistics.",
+        args: {
+            providerId: tool.schema
+                .string()
+                .optional()
+                .describe(
+                    "Optional: Filter results to show only quotas from a specific provider (e.g., 'antigravity', 'codex')"
+                ),
+            modelId: tool.schema
+                .string()
+                .optional()
+                .describe(
+                    "Optional: Filter results to show only quotas relevant to a specific model"
+                ),
+        },
+        async execute(args) {
+            const config = getConfig();
+            
+            logger.debug("tool:quotas:execute", {
+                providerId: args.providerId,
+                modelId: args.modelId,
+            });
+
+            try {
+                // Always fetch fresh quotas
+                const quotas = await quotaService.getQuotas({
+                    providerId: args.providerId,
+                    modelId: args.modelId,
+                });
+
+                logger.debug("tool:quotas:fetched", {
+                    count: quotas.length,
+                });
+
+                if (quotas.length === 0) {
+                    return "No quota information available. This could mean:\n" +
+                        "- No quota providers are configured\n" +
+                        "- The filter criteria matched no quotas\n" +
+                        "- There was an error fetching quota data";
+                }
+
+                // Render the quota table
+                const lines = renderQuotaTable(quotas, {
+                    progressBarConfig: config.progressBar,
+                    tableConfig: config.table,
+                }).map((l) => l.line);
+
+                const showMode = config.progressBar?.show ?? "used";
+                const modeLabel = showMode === "available" ? "(Remaining)" : "(Used)";
+
+                return `**Quota Status ${modeLabel}**\n\n` + lines.join("\n");
+            } catch (error) {
+                logger.error("tool:quotas:error", { error });
+                return `Error fetching quotas: ${error instanceof Error ? error.message : String(error)}`;
+            }
+        },
+    });
+}