import { useCallback, useEffect, useRef, useState } from "react"; import { invoke } from "@tauri-apps/api/core"; import type { Agent } from "../types"; export interface UsageData { planName?: string | null; remaining?: number | null; total?: number | null; used?: number | null; unit?: string | null; isValid?: boolean | null; resetsAt?: string | null; } export interface UsageResult { success: boolean; data?: UsageData[] | null; error?: string | null; } export type UsageStatus = "idle" | "loading" | "success" | "error"; export interface CacheEntry { status: "success" | "error"; data: UsageData[]; /** Raw error text from Rust; null = transient failure (invoke rejected). */ error: string | null; updatedAt: number; } // Module-level cache shared across mounts: re-entering the list within the // stale TTL shows the last result without firing new requests. Both the // compact (card header) and detail (footer) variants read this same cache. const STALE_TTL_MS = 5 * 60 * 1000; const cache = new Map(); // Simple semaphore: at most MAX_CONCURRENT queries in flight at once. const MAX_CONCURRENT = 3; let running = 0; const waiters: Array<() => void> = []; async function acquireSlot(): Promise { if (running >= MAX_CONCURRENT) { await new Promise((resolve) => waiters.push(resolve)); } running += 1; } function releaseSlot(): void { running -= 1; waiters.shift()?.(); } export interface UsageQueryState { status: UsageStatus; data: UsageData[]; error: string | null | undefined; updatedAt: number | null; refresh: (force: boolean) => void; supported: boolean; } export function useUsageQuery( agent: Agent, providerName: string, usageKinds?: string[], autoIntervalMinutes?: number, disabled?: boolean ): UsageQueryState { const supported = (usageKinds?.length ?? 0) > 0; // Cache key includes the agent so a Kimi Code provider and a Pi provider // with the same name do not clobber each other's cached result. const cacheKey = `${agent}:${providerName}`; const [status, setStatus] = useState("idle"); const [data, setData] = useState([]); /** undefined = no error; null = network error; string = Rust error text. */ const [error, setError] = useState(undefined); const [updatedAt, setUpdatedAt] = useState(null); // Generation counter: stale responses (unmounted / superseded query) are ignored. const genRef = useRef(0); const runQuery = useCallback( async (forceRefresh: boolean) => { if (disabled) return; const gen = ++genRef.current; setStatus("loading"); setError(undefined); const finish = (entry: CacheEntry) => { if (genRef.current !== gen) return; cache.set(cacheKey, entry); setStatus(entry.status); setData(entry.data); setError(entry.error); setUpdatedAt(entry.updatedAt); }; try { await acquireSlot(); let result: UsageResult; try { result = await invoke("query_provider_usage", { agent, providerName, forceRefresh, }); } finally { releaseSlot(); } if (genRef.current !== gen) return; if (result.success) { finish({ status: "success", data: result.data ?? [], error: null, updatedAt: Date.now(), }); } else { // Deterministic failure: keep last good data for ghost display. finish({ status: "error", data: cache.get(cacheKey)?.data ?? [], error: result.error ?? "", updatedAt: Date.now(), }); } } catch { // Transient failure (network / timeout): invoke rejected. finish({ status: "error", data: cache.get(cacheKey)?.data ?? [], error: null, updatedAt: Date.now(), }); } }, [agent, providerName, cacheKey, disabled] ); // On mount / provider change: serve fresh cache, otherwise query once. useEffect(() => { if (disabled) { // Toggle back off: stop immediately and return to idle so the footer // hides; bump gen so an in-flight response is dropped (no stale write). genRef.current += 1; setStatus("idle"); setData([]); setError(undefined); setUpdatedAt(null); return; } if (!supported) return; const cached = cache.get(cacheKey); if (cached) { setStatus(cached.status); setData(cached.data); setError(cached.error); setUpdatedAt(cached.updatedAt); if (Date.now() - cached.updatedAt < STALE_TTL_MS) return; } void runQuery(false); }, [supported, cacheKey, runQuery, disabled]); // Ignore late responses after unmount. useEffect(() => { return () => { genRef.current += 1; }; }, []); // Auto query interval — the hook is called once per provider card and both // variants read the same state, so there is no second mount to double-fire. useEffect(() => { if (!supported || disabled) return; const mins = autoIntervalMinutes ?? 0; if (mins <= 0) return; const id = setInterval(() => void runQuery(false), mins * 60_000); return () => clearInterval(id); }, [supported, autoIntervalMinutes, runQuery, disabled]); return { status, data, error, updatedAt, refresh: runQuery, supported, }; }