import type { ExtensionAPI, Theme } from "@earendil-works/pi-coding-agent"; import { Key, matchesKey, truncateToWidth, visibleWidth } from "@earendil-works/pi-tui"; import fs from "node:fs"; import os from "node:os"; import path from "node:path"; import * as readline from "node:readline"; import { agentPath } from "../shared/paths.ts"; interface RateWindow { label: string; usedPercent: number; resetDescription?: string; } interface ProviderStatus { indicator: "none" | "minor" | "major" | "critical" | "maintenance" | "unknown"; description?: string; } interface UsageSnapshot { provider: string; displayName: string; windows: RateWindow[]; plan?: string; error?: string; status?: ProviderStatus; } type AuthJson = Record; const STATUS_URLS: Record = { anthropic: "https://status.anthropic.com/api/v2/status.json", codex: "https://status.openai.com/api/v2/status.json", copilot: "https://www.githubstatus.com/api/v2/status.json", }; const SNAPSHOT_CACHE_MS = 15_000; const FAILURE_CACHE_MS = 2_000; const JSONL_FILE_CACHE_MS = 60_000; const MAX_JSONL_FILE_CACHE = 200; const FS_CONCURRENCY = 4; const snapshotCache = new Map }>(); const statusCache = new Map }>(); const jsonlCostCache = new Map }>(); function readJson(file: string): any | undefined { try { if (!fs.existsSync(file)) return undefined; return JSON.parse(fs.readFileSync(file, "utf8")); } catch { return undefined; } } function readPiAuth(): AuthJson { return readJson(agentPath("auth.json")) ?? {}; } async function fetchWithTimeout(url: string, init: RequestInit = {}, timeoutMs = 5000, signal?: AbortSignal): Promise { const controller = new AbortController(); const timeout = setTimeout(() => controller.abort(), timeoutMs); function onExternalAbort() { controller.abort(); } signal?.addEventListener("abort", onExternalAbort, { once: true }); try { return await fetch(url, { ...init, signal: controller.signal }); } finally { clearTimeout(timeout); signal?.removeEventListener("abort", onExternalAbort); } } async function withTimeout(promise: Promise, timeoutMs: number, fallback: T, signal?: AbortSignal): Promise { let timeout: ReturnType | undefined; let rejectOnAbort: (() => void) | undefined; try { return await Promise.race([ promise, new Promise((resolve, reject) => { timeout = setTimeout(() => resolve(fallback), timeoutMs); rejectOnAbort = () => reject(signal?.reason ?? new DOMException("Aborted", "AbortError")); signal?.addEventListener("abort", rejectOnAbort, { once: true }); }), ]); } finally { if (timeout) clearTimeout(timeout); if (rejectOnAbort) signal?.removeEventListener("abort", rejectOnAbort); } } function formatReset(date: Date): string { const diffMs = date.getTime() - Date.now(); if (diffMs <= 0) return "now"; const minutes = Math.floor(diffMs / 60000); if (minutes < 60) return `${minutes}m`; const hours = Math.floor(minutes / 60); if (hours < 24) return `${hours}h${minutes % 60 ? ` ${minutes % 60}m` : ""}`; const days = Math.floor(hours / 24); if (days < 7) return `${days}d ${hours % 24}h`; return new Intl.DateTimeFormat("en-US", { month: "short", day: "numeric" }).format(date); } function statusEmoji(status?: ProviderStatus) { switch (status?.indicator) { case "none": return "✅"; case "minor": return "⚠️"; case "major": return "🟠"; case "critical": return "🔴"; case "maintenance": return "🔧"; default: return ""; } } async function fetchProviderStatus(provider: string, signal?: AbortSignal): Promise { const url = STATUS_URLS[provider]; if (!url) return { indicator: "none" }; try { const res = await fetchWithTimeout(url, {}, 3000, signal); if (!res.ok) return { indicator: "unknown" }; const data = await res.json() as any; return { indicator: (data.status?.indicator ?? "none") as ProviderStatus["indicator"], description: data.status?.description, }; } catch { return { indicator: "unknown" }; } } function cachedUsageSnapshot(key: string, loader: (signal?: AbortSignal) => Promise, signal?: AbortSignal): Promise { const now = Date.now(); const cached = snapshotCache.get(key); if (cached && cached.expiresAt > now) return cached.promise; const promise = loader(signal).then((result) => { const expiresAt = now + (result.error ? FAILURE_CACHE_MS : SNAPSHOT_CACHE_MS); snapshotCache.set(key, { expiresAt, promise: Promise.resolve(result) }); return result; }).catch((error) => { // Network-level errors (not caught inside the loader): cache briefly to avoid // thundering-herd retries, then allow retry. const errorSnapshot: UsageSnapshot = { provider: key, displayName: key, windows: [], error: String(error) }; snapshotCache.set(key, { expiresAt: now + FAILURE_CACHE_MS, promise: Promise.resolve(errorSnapshot) }); return errorSnapshot; }); snapshotCache.set(key, { expiresAt: now + SNAPSHOT_CACHE_MS, promise }); return promise; } function cachedProviderStatus(provider: string, signal?: AbortSignal): Promise { const now = Date.now(); const cached = statusCache.get(provider); if (cached && cached.expiresAt > now) return cached.promise; const promise = fetchProviderStatus(provider, signal).then((result) => { const expiresAt = now + (result.indicator === "unknown" ? FAILURE_CACHE_MS : SNAPSHOT_CACHE_MS); statusCache.set(provider, { expiresAt, promise: Promise.resolve(result) }); return result; }).catch(() => { const result = { indicator: "unknown" as const }; statusCache.set(provider, { expiresAt: now + FAILURE_CACHE_MS, promise: Promise.resolve(result) }); return result; }); statusCache.set(provider, { expiresAt: now + SNAPSHOT_CACHE_MS, promise }); return promise; } async function fetchClaudeUsage(signal?: AbortSignal): Promise { const token = readPiAuth().anthropic?.access; if (!token) return { provider: "anthropic", displayName: "Claude", windows: [], error: "No credentials" }; try { const res = await fetchWithTimeout("https://api.anthropic.com/api/oauth/usage", { headers: { Authorization: `Bearer ${token}`, "anthropic-beta": "oauth-2025-04-20" }, }, 5000, signal); if (!res.ok) return { provider: "anthropic", displayName: "Claude", windows: [], error: `HTTP ${res.status}` }; const data = await res.json() as any; const windows: RateWindow[] = []; for (const [key, label] of [["five_hour", "5h"], ["seven_day", "Week"], ["seven_day_sonnet", "Sonnet"], ["seven_day_opus", "Opus"]]) { const bucket = data[key]; if (bucket?.utilization !== undefined) { windows.push({ label, usedPercent: bucket.utilization, resetDescription: bucket.resets_at ? formatReset(new Date(bucket.resets_at)) : undefined }); } } return { provider: "anthropic", displayName: "Claude", windows }; } catch (error) { return { provider: "anthropic", displayName: "Claude", windows: [], error: String(error) }; } } async function fetchCopilotUsage(signal?: AbortSignal): Promise { const token = readPiAuth()["github-copilot"]?.refresh; if (!token) return { provider: "copilot", displayName: "Copilot", windows: [], error: "No token" }; try { const res = await fetchWithTimeout("https://api.github.com/copilot_internal/user", { headers: { Authorization: `token ${token}`, Accept: "application/json", "Editor-Version": "vscode/1.96.2", "User-Agent": "GitHubCopilotChat/0.26.7", "X-Github-Api-Version": "2025-04-01", }, }, 5000, signal); if (!res.ok) return { provider: "copilot", displayName: "Copilot", windows: [], error: `HTTP ${res.status}` }; const data = await res.json() as any; const windows: RateWindow[] = []; const reset = data.quota_reset_date_utc ? formatReset(new Date(data.quota_reset_date_utc)) : undefined; const premium = data.quota_snapshots?.premium_interactions; if (premium) { const remaining = premium.remaining ?? 0; const entitlement = premium.entitlement ?? 0; windows.push({ label: "Premium", usedPercent: Math.max(0, 100 - (premium.percent_remaining ?? 0)), resetDescription: `${reset ? `${reset} ` : ""}(${remaining}/${entitlement})`, }); } const chat = data.quota_snapshots?.chat; if (chat && !chat.unlimited) windows.push({ label: "Chat", usedPercent: Math.max(0, 100 - (chat.percent_remaining ?? 0)), resetDescription: reset }); return { provider: "copilot", displayName: "Copilot", windows, plan: data.copilot_plan }; } catch (error) { return { provider: "copilot", displayName: "Copilot", windows: [], error: String(error) }; } } const OPENCODE_GO_LIMITS = [ { label: "5h", ms: 5 * 60 * 60 * 1000, limit: 12 }, { label: "Week", ms: 7 * 24 * 60 * 60 * 1000, limit: 30 }, { label: "Month", ms: 30 * 24 * 60 * 60 * 1000, limit: 60 }, ] as const; async function walkJsonlFiles(dir: string, cutoffMs: number, files: string[] = [], signal?: AbortSignal): Promise { if (signal?.aborted) return files; try { const entries = await fs.promises.readdir(dir, { withFileTypes: true }); if (signal?.aborted) return files; for (const entry of entries) { if (signal?.aborted) return files; const entryPath = path.join(dir, entry.name); if (entry.isDirectory()) { await walkJsonlFiles(entryPath, cutoffMs, files, signal); continue; } if (!entry.isFile() || !entry.name.endsWith(".jsonl")) continue; const stat = await fs.promises.stat(entryPath); if (stat.mtimeMs >= cutoffMs) files.push(entryPath); } } catch {} return files; } function entryTimestampMs(entry: any, message: any): number | undefined { const raw = entry?.timestamp ?? message?.timestamp; if (typeof raw === "number") return raw; if (typeof raw === "string") { const parsed = Date.parse(raw); if (!Number.isNaN(parsed)) return parsed; } return undefined; } function pruneJsonlCostCache(now = Date.now()) { for (const [file, cached] of jsonlCostCache) { if (cached.expiresAt <= now) jsonlCostCache.delete(file); } while (jsonlCostCache.size > MAX_JSONL_FILE_CACHE) { const oldestKey = jsonlCostCache.keys().next().value; if (!oldestKey) break; jsonlCostCache.delete(oldestKey); } } async function parseOpenCodeGoCostFile(file: string, stat: fs.Stats, cutoffMs: number, signal?: AbortSignal): Promise> { const now = Date.now(); const cached = jsonlCostCache.get(file); if (cached && cached.mtimeMs === stat.mtimeMs && cached.size === stat.size && cached.expiresAt > now) { return cached.costs.filter((cost) => cost.time >= cutoffMs); } if (signal?.aborted) return []; const costs: Array<{ time: number; cost: number }> = []; const rl = readline.createInterface({ input: fs.createReadStream(file, { encoding: "utf8" }), crlfDelay: Infinity }); const onAbort = () => { rl.close(); }; signal?.addEventListener("abort", onAbort, { once: true }); try { for await (const line of rl) { if (signal?.aborted) break; if (!line.trim()) continue; try { const entry = JSON.parse(line); const message = entry?.message; if (entry?.type !== "message" || message?.role !== "assistant" || message?.provider !== "opencode-go") continue; const cost = Number(message?.usage?.cost?.total ?? message?.cost?.total ?? message?.cost ?? 0); if (!Number.isFinite(cost) || cost <= 0) continue; const time = entryTimestampMs(entry, message); if (time === undefined || time < cutoffMs) continue; costs.push({ time, cost }); } catch {} } } finally { signal?.removeEventListener("abort", onAbort); rl.close(); } jsonlCostCache.set(file, { mtimeMs: stat.mtimeMs, size: stat.size, expiresAt: now + JSONL_FILE_CACHE_MS, costs }); pruneJsonlCostCache(now); return costs; } async function runWithConcurrency(items: T[], concurrency: number, fn: (item: T) => Promise): Promise { const results: R[] = []; const queue = [...items]; async function worker() { while (queue.length > 0) { const item = queue.shift()!; results.push(await fn(item)); } } await Promise.all(Array.from({ length: Math.min(concurrency, items.length) }, () => worker())); return results; } async function collectLocalOpenCodeGoCosts(signal?: AbortSignal): Promise> { const oldestWindow = Math.max(...OPENCODE_GO_LIMITS.map((window) => window.ms)); const cutoffMs = Date.now() - oldestWindow; // Session files are append-only; the 24h buffer avoids missing files whose mtime and entry // timestamp differ slightly because of clock skew, copies, or filesystem precision. const files = await walkJsonlFiles(agentPath("sessions"), cutoffMs - 24 * 60 * 60 * 1000, [], signal); if (signal?.aborted) return []; const results = await runWithConcurrency(files, FS_CONCURRENCY, async (file) => { if (signal?.aborted) return []; try { const stat = await fs.promises.stat(file); if (signal?.aborted) return []; return await parseOpenCodeGoCostFile(file, stat, cutoffMs, signal); } catch { return []; } }); return results.flat(); } function resetFromOldestContribution(costs: Array<{ time: number; cost: number }>, windowMs: number): string | undefined { if (costs.length === 0) return undefined; const oldest = costs.reduce((min, cost) => Math.min(min, cost.time), Number.POSITIVE_INFINITY); return Number.isFinite(oldest) ? formatReset(new Date(oldest + windowMs)) : undefined; } async function fetchOpenCodeGoUsage(signal?: AbortSignal): Promise { try { const token = readPiAuth()["opencode-go"]?.key ?? process.env.OPENCODE_API_KEY; const now = Date.now(); const allCosts = await collectLocalOpenCodeGoCosts(signal); if (signal?.aborted) return { provider: "opencode-go", displayName: "OpenCode Go", windows: [], error: "Cancelled" }; if (!token && allCosts.length === 0) return { provider: "opencode-go", displayName: "OpenCode Go", windows: [], error: "No credentials" }; const windows = OPENCODE_GO_LIMITS.map((window) => { const costs = allCosts.filter((cost) => cost.time >= now - window.ms); const used = costs.reduce((sum, cost) => sum + cost.cost, 0); const reset = resetFromOldestContribution(costs, window.ms); return { label: window.label, usedPercent: (used / window.limit) * 100, resetDescription: `${reset ? `${reset} ` : ""}($${used.toFixed(2)}/$${window.limit})`, }; }); return { provider: "opencode-go", displayName: "OpenCode Go", windows, plan: "Go local estimate" }; } catch (error) { return { provider: "opencode-go", displayName: "OpenCode Go", windows: [], error: String(error) }; } } async function fetchCodexUsage(modelRegistry: any, signal?: AbortSignal): Promise { let accessToken: string | undefined; let accountId: string | undefined; try { accessToken = await Promise.resolve(modelRegistry?.authStorage?.getApiKey?.("openai-codex")); const cred = await Promise.resolve(modelRegistry?.authStorage?.get?.("openai-codex")); accountId = cred?.accountId; } catch {} if (!accessToken) { const auth = readJson(path.join(process.env.CODEX_HOME ?? path.join(os.homedir(), ".codex"), "auth.json")); accessToken = auth?.OPENAI_API_KEY ?? auth?.tokens?.access_token; accountId = auth?.tokens?.account_id; } if (!accessToken) return { provider: "codex", displayName: "Codex", windows: [], error: "No credentials" }; try { const headers: Record = { Authorization: `Bearer ${accessToken}`, "User-Agent": "PiUsageBar", Accept: "application/json" }; if (accountId) headers["ChatGPT-Account-Id"] = accountId; const res = await fetchWithTimeout("https://chatgpt.com/backend-api/wham/usage", { headers }, 5000, signal); if (res.status === 401 || res.status === 403) return { provider: "codex", displayName: "Codex", windows: [], error: "Token expired" }; if (!res.ok) return { provider: "codex", displayName: "Codex", windows: [], error: `HTTP ${res.status}` }; const data = await res.json() as any; const windows: RateWindow[] = []; for (const [key, fallback, name] of [["primary_window", 10800, ""], ["secondary_window", 86400, "Day"]] as const) { const bucket = data.rate_limit?.[key]; if (!bucket) continue; const hours = Math.round((bucket.limit_window_seconds ?? fallback) / 3600); windows.push({ label: name || `${hours}h`, usedPercent: bucket.used_percent ?? 0, resetDescription: bucket.reset_at ? formatReset(new Date(bucket.reset_at * 1000)) : undefined, }); } const balance = data.credits?.balance == null ? undefined : Number(data.credits.balance); const plan = balance === undefined || Number.isNaN(balance) ? data.plan_type : `${data.plan_type ?? "credits"} ($${balance.toFixed(2)})`; return { provider: "codex", displayName: "Codex", windows, plan }; } catch (error) { return { provider: "codex", displayName: "Codex", windows: [], error: String(error) }; } } function padVisible(text: string, width: number) { const clipped = truncateToWidth(text, Math.max(0, width), ""); const length = visibleWidth(clipped); if (length >= width) return clipped; return `${clipped}${" ".repeat(width - length)}`; } class UsageBarComponent { private usages: UsageSnapshot[] = []; private loading = true; private disposed = false; private abortController = new AbortController(); constructor( private readonly tui: { requestRender: () => void }, private readonly theme: Theme, private readonly done: (result: void) => void, private readonly modelRegistry: any, ) { void this.load(); } private async load() { try { const signal = this.abortController.signal; const [claude, copilot, openCodeGo, codex, claudeStatus, copilotStatus, codexStatus] = await Promise.all([ withTimeout(cachedUsageSnapshot("anthropic", (s) => fetchClaudeUsage(s), signal), 6000, { provider: "anthropic", displayName: "Claude", windows: [], error: "Timeout" }, signal), withTimeout(cachedUsageSnapshot("copilot", (s) => fetchCopilotUsage(s), signal), 6000, { provider: "copilot", displayName: "Copilot", windows: [], error: "Timeout" }, signal), withTimeout(cachedUsageSnapshot("opencode-go", (s) => fetchOpenCodeGoUsage(s), signal), 6000, { provider: "opencode-go", displayName: "OpenCode Go", windows: [], error: "Timeout" }, signal), withTimeout(cachedUsageSnapshot("codex", (s) => fetchCodexUsage(this.modelRegistry, s), signal), 6000, { provider: "codex", displayName: "Codex", windows: [], error: "Timeout" }, signal), cachedProviderStatus("anthropic", signal), cachedProviderStatus("copilot", signal), cachedProviderStatus("codex", signal), ]); if (signal.aborted) return; claude.status = claudeStatus; copilot.status = copilotStatus; openCodeGo.status = { indicator: "none" }; codex.status = codexStatus; this.usages = [claude, copilot, openCodeGo, codex].filter((usage) => usage.windows.length > 0 || !["No credentials", "No token"].includes(usage.error ?? "")); } catch (error) { if (this.disposed) return; this.usages = [{ provider: "usage", displayName: "Usage", windows: [], error: String(error) }]; } finally { if (this.disposed) return; this.loading = false; this.tui.requestRender(); } } handleInput(data: string): void { if (matchesKey(data, Key.escape) || matchesKey(data, Key.enter) || matchesKey(data, "q")) { this.abortController.abort(); this.done(undefined); } } invalidate(): void { if (!this.disposed) this.tui.requestRender(); } render(width: number): string[] { if (width < 4) return [truncateToWidth(this.loading ? "..." : "Usage", Math.max(0, width), "")]; const totalW = Math.min(62, width); const innerW = Math.max(0, totalW - 4); const borderText = "─".repeat(Math.max(0, totalW - 2)); const border = this.theme.fg("border", borderText); const lines: string[] = []; const row = (content: string) => `${this.theme.fg("border", "│ ")}${padVisible(content, innerW)}${this.theme.fg("border", " │")}`; const separator = () => `${this.theme.fg("border", "├")}${border}${this.theme.fg("border", "┤")}`; lines.push(this.theme.fg("border", `╭${borderText}╮`)); lines.push(row(this.theme.bold(this.theme.fg("accent", totalW < 24 ? "Usage" : "AI Usage")))); lines.push(separator()); if (this.loading) { lines.push(row("Loading usage and provider status...")); } else if (this.usages.length === 0) { lines.push(row(this.theme.fg("dim", "No configured provider usage found."))); } else { for (const usage of this.usages) { const plan = usage.plan ? this.theme.fg("dim", ` (${usage.plan})`) : ""; const emoji = statusEmoji(usage.status); lines.push(row(`${this.theme.bold(usage.displayName)}${plan}${emoji ? ` ${emoji}` : ""}`)); if (usage.status?.indicator && !["none", "unknown"].includes(usage.status.indicator) && usage.status.description) { lines.push(row(` ${this.theme.fg("warning", usage.status.description.slice(0, 48))}`)); } if (usage.error) { lines.push(row(` ${this.theme.fg("dim", usage.error)}`)); } else { for (const window of usage.windows) { const used = Math.max(0, Math.min(100, window.usedPercent)); const remaining = 100 - used; const percent = `${remaining.toFixed(0).padStart(3)}%`; const labelW = Math.min(8, Math.max(2, innerW - percent.length - 6)); const barW = Math.max(3, Math.min(14, innerW - labelW - percent.length - 4)); const filled = Math.round((used / 100) * barW); const color = remaining <= 10 ? "error" : remaining <= 30 ? "warning" : "success"; const bar = `${this.theme.fg(color, "█".repeat(filled))}${this.theme.fg("dim", "░".repeat(barW - filled))}`; const label = padVisible(window.label, labelW); const reset = innerW >= 30 && window.resetDescription ? this.theme.fg("dim", ` ⏱ ${window.resetDescription}`) : ""; lines.push(row(` ${label} ${bar} ${percent}${reset}`)); } } lines.push(row("")); } } lines.push(separator()); lines.push(row(this.theme.fg("dim", totalW < 28 ? "Esc/q/Enter" : "Esc, Enter, or q to close"))); lines.push(this.theme.fg("border", `╰${borderText}╯`)); return lines.map((line) => truncateToWidth(line, Math.max(0, width), "")); } dispose(): void { this.disposed = true; this.abortController.abort(); } } export default function (pi: ExtensionAPI) { pi.registerCommand("usage", { description: "Show AI provider usage and quota bars", handler: async (_args, ctx) => { if (!ctx.hasUI || ctx.mode !== "tui") { ctx.ui.notify("/usage requires interactive TUI mode", "error"); return; } await ctx.ui.custom((tui, theme, _keybindings, done) => new UsageBarComponent(tui, theme, done, ctx.modelRegistry), { overlay: true }); }, }); }