diff --git a/pi/.pi/agent/extensions/footer-display.ts b/pi/.pi/agent/extensions/footer-display.ts index 5418458..2c480cd 100644 --- a/pi/.pi/agent/extensions/footer-display.ts +++ b/pi/.pi/agent/extensions/footer-display.ts @@ -62,6 +62,21 @@ function getModelShortName(modelId: string): string { return modelId.replace(/^claude-/, ""); } +// Format duration in milliseconds to human readable (e.g., "2h 55m") +function formatDurationMs(ms: number): string { + if (!Number.isFinite(ms) || ms <= 0) return "now"; + const totalSeconds = Math.floor(ms / 1000); + const d = Math.floor(totalSeconds / 86400); + const h = Math.floor((totalSeconds % 86400) / 3600); + const m = Math.floor((totalSeconds % 3600) / 60); + if (d > 0 && h > 0) return `${d}d ${h}h`; + if (d > 0) return `${d}d`; + if (h > 0 && m > 0) return `${h}h ${m}m`; + if (h > 0) return `${h}h`; + if (m > 0) return `${m}m`; + return "<1m"; +} + // Nerd Font codepoints matched to what claude-account-switch.ts emits const ICON_PERSONAL = "\uEF85"; // U+EF85 — home const ICON_WORK = "\uDB80\uDCD6"; // U+F00D6 — briefcase (surrogate pair) @@ -71,6 +86,12 @@ export default function (pi: ExtensionAPI) { let ctx: any = null; let tuiRef: any = null; let footerDataRef: any = null; + + // Track usage data for dynamic S/W bar rendering + let usageSession: number | null = null; + let usageWeekly: number | null = null; + let sessionResetsAt: number | null = null; + let weeklyResetsAt: number | null = null; // --------------------------------------------------------------------------- // Footer line builder — called on every render @@ -106,7 +127,33 @@ export default function (pi: ExtensionAPI) { const usageRaw = statuses.get("usage-bars"); const contextUsage = ctx?.getContextUsage?.(); { - let block = usageRaw ?? ""; + let block: string; + + if (usageSession !== null && usageWeekly !== null) { + // Build S/W bars directly from stored event data so we can cleanly + // append the dynamic countdown without trying to parse ANSI strings. + const session = Math.max(0, Math.min(100, Math.round(usageSession))); + const weekly = Math.max(0, Math.min(100, Math.round(usageWeekly))); + + let sPart = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`); + let wPart = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`); + + if (sessionResetsAt !== null) { + const msLeft = sessionResetsAt - Date.now(); + if (msLeft > 0) sPart += " " + theme.fg("dim", formatDurationMs(msLeft)); + } + + if (weeklyResetsAt !== null) { + const msLeft = weeklyResetsAt - Date.now(); + if (msLeft > 0) wPart += " " + theme.fg("dim", `\u27F3 ${formatDurationMs(msLeft)}`); + } + + block = sPart + pipeSep + wPart; + } else { + // Fallback to raw status for loading / error states + block = usageRaw ?? ""; + } + if (contextUsage && contextUsage.percent !== null) { const pct = Math.round(contextUsage.percent); const cBar = @@ -188,4 +235,14 @@ export default function (pi: ExtensionAPI) { pi.events.on("claude-account:switched", () => { if (tuiRef) tuiRef.requestRender(); }); + + // Listen for usage updates — store raw values so we can build bars + dynamic + // countdown directly rather than parsing the ANSI status string from usage-bars. + pi.events.on("usage:update", (data: any) => { + if (data.session !== undefined) usageSession = data.session; + if (data.weekly !== undefined) usageWeekly = data.weekly; + if (data.sessionResetsAt !== undefined) sessionResetsAt = data.sessionResetsAt; + if (data.weeklyResetsAt !== undefined) weeklyResetsAt = data.weeklyResetsAt; + if (tuiRef) tuiRef.requestRender(); + }); } diff --git a/pi/.pi/agent/extensions/footer-display.ts.backup b/pi/.pi/agent/extensions/footer-display.ts.backup new file mode 100644 index 0000000..5418458 --- /dev/null +++ b/pi/.pi/agent/extensions/footer-display.ts.backup @@ -0,0 +1,191 @@ +/** + * Footer Display Extension + * + * Replaces the built-in pi footer with a single clean line that assembles + * status from all other extensions: + * + * \uEF85 | S ⣿⣶⣀⣀⣀ 34% 2h 55m | W ⣿⣿⣷⣀⣀ 68% ⟳ Fri 09:00 | C ⣿⣀⣀⣀⣀ 20% | Sonnet 4.6 | rust-analyzer | MCP: 1/2 + * + * Status sources: + * "claude-account" — set by claude-account-switch.ts → just the icon + * "usage-bars" — set by usage-bars extension → S/W bars (pass-through) + * ctx.getContextUsage() → C bar (rendered here) + * ctx.model → model short name + * "lsp" — set by lsp-pi extension → strip "LSP " prefix + * "mcp" — set by pi-mcp-adapter → strip " servers" suffix + */ + +import os from "os"; +import path from "path"; +import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; +import { truncateToWidth } from "@mariozechner/pi-tui"; + +// --------------------------------------------------------------------------- +// Braille gradient bar — used here only for the context (C) bar +// --------------------------------------------------------------------------- +const BRAILLE_GRADIENT = "\u28C0\u28C4\u28E4\u28E6\u28F6\u28F7\u28FF"; +const BRAILLE_EMPTY = "\u28C0"; +const BAR_WIDTH = 5; + +function renderBrailleBar(theme: any, value: number): string { + const v = Math.max(0, Math.min(100, Math.round(value))); + const levels = BRAILLE_GRADIENT.length - 1; + const totalSteps = BAR_WIDTH * levels; + const filledSteps = Math.round((v / 100) * totalSteps); + const full = Math.floor(filledSteps / levels); + const partial = filledSteps % levels; + const empty = BAR_WIDTH - full - (partial ? 1 : 0); + const color = v >= 90 ? "error" : v >= 70 ? "warning" : "success"; + const filled = BRAILLE_GRADIENT[BRAILLE_GRADIENT.length - 1]!.repeat(Math.max(0, full)); + const partialChar = partial ? BRAILLE_GRADIENT[partial]! : ""; + const emptyChars = BRAILLE_EMPTY.repeat(Math.max(0, empty)); + return theme.fg(color, filled + partialChar) + theme.fg("dim", emptyChars); +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- +function stripAnsi(text: string): string { + return text.replace(/\x1b\[[0-9;]*m/g, "").replace(/\x1b\][^\x07]*\x07/g, ""); +} + +function getModelShortName(modelId: string): string { + // claude-haiku-4-5 → "Haiku 4.5", claude-sonnet-4-6 → "Sonnet 4.6" + const m = modelId.match(/^claude-([a-z]+)-([\d]+(?:-[\d]+)*)(?:-\d{8})?$/); + if (m) { + const family = m[1]!.charAt(0).toUpperCase() + m[1]!.slice(1); + return `${family} ${m[2]!.replace(/-/g, ".")}`; + } + // claude-3-5-sonnet, claude-3-opus, etc. + const m2 = modelId.match(/^claude-[\d-]+-([a-z]+)/); + if (m2) return m2[1]!.charAt(0).toUpperCase() + m2[1]!.slice(1); + return modelId.replace(/^claude-/, ""); +} + +// Nerd Font codepoints matched to what claude-account-switch.ts emits +const ICON_PERSONAL = "\uEF85"; // U+EF85 — home +const ICON_WORK = "\uDB80\uDCD6"; // U+F00D6 — briefcase (surrogate pair) +const ICON_UNKNOWN = "\uF420"; // U+F420 — claude default + +export default function (pi: ExtensionAPI) { + let ctx: any = null; + let tuiRef: any = null; + let footerDataRef: any = null; + + // --------------------------------------------------------------------------- + // Footer line builder — called on every render + // --------------------------------------------------------------------------- + function buildFooterLine(theme: any): string { + const sep = theme.fg("dim", " · "); + const pipeSep = theme.fg("dim", " | "); + const parts: string[] = []; + + const statuses: ReadonlyMap = + footerDataRef?.getExtensionStatuses?.() ?? new Map(); + + // 1. Current working directory + const home = os.homedir(); + const cwd = process.cwd(); + const dir = cwd.startsWith(home) + ? "~" + path.sep + path.relative(home, cwd) + : cwd; + parts.push(theme.fg("dim", dir)); + + // 2. Account icon + const acctRaw = statuses.get("claude-account"); + if (acctRaw !== undefined) { + const clean = stripAnsi(acctRaw).trim(); + let icon: string; + if (clean.includes("personal")) icon = ICON_PERSONAL; + else if (clean.includes("work")) icon = ICON_WORK; + else icon = ICON_UNKNOWN; + parts.push(theme.fg("dim", icon)); + } + + // 3. S / W usage bars + C bar — joined as one |-separated block + const usageRaw = statuses.get("usage-bars"); + const contextUsage = ctx?.getContextUsage?.(); + { + let block = usageRaw ?? ""; + if (contextUsage && contextUsage.percent !== null) { + const pct = Math.round(contextUsage.percent); + const cBar = + theme.fg("muted", "C ") + + renderBrailleBar(theme, pct) + + " " + + theme.fg("dim", `${pct}%`); + block = block ? block + pipeSep + cBar : cBar; + } + if (block) parts.push(block); + } + + // 4. Model short name + const modelId = ctx?.model?.id; + if (modelId) parts.push(theme.fg("dim", getModelShortName(modelId))); + + // 5. LSP — strip "LSP" prefix and activity dot + const lspRaw = statuses.get("lsp"); + if (lspRaw) { + const clean = stripAnsi(lspRaw).trim().replace(/^LSP\s*[•·]?\s*/i, "").trim(); + if (clean) parts.push(theme.fg("dim", clean)); + } + + // 6. MCP — strip " servers" suffix + const mcpRaw = statuses.get("mcp"); + if (mcpRaw) { + const clean = stripAnsi(mcpRaw).trim().replace(/\s*servers?.*$/, "").trim(); + if (clean) parts.push(theme.fg("dim", clean)); + } + + return parts.join(sep); + } + + // --------------------------------------------------------------------------- + // Footer installation + // --------------------------------------------------------------------------- + function installFooter(_ctx: any) { + if (!_ctx?.hasUI) return; + _ctx.ui.setFooter((_tui: any, theme: any, footerData: any) => { + tuiRef = _tui; + footerDataRef = footerData; + const unsub = footerData.onBranchChange(() => _tui.requestRender()); + return { + dispose: unsub, + invalidate() {}, + render(width: number): string[] { + return [truncateToWidth(buildFooterLine(theme) || "", width)]; + }, + }; + }); + } + + // --------------------------------------------------------------------------- + // Event handlers + // --------------------------------------------------------------------------- + pi.on("session_start", async (_event, _ctx) => { + ctx = _ctx; + installFooter(_ctx); + }); + + pi.on("session_shutdown", (_event, _ctx) => { + if (_ctx?.hasUI) _ctx.ui.setFooter(undefined); + }); + + // Re-render after turns so context usage stays current + pi.on("turn_end", (_event, _ctx) => { + ctx = _ctx; + if (tuiRef) tuiRef.requestRender(); + }); + + // Re-render when model changes (updates model name in footer) + pi.on("model_select", (_event, _ctx) => { + ctx = _ctx; + if (tuiRef) tuiRef.requestRender(); + }); + + // Re-render when account switches (usage:update comes from usage-bars setStatus which + // already triggers a render, but account icon needs a nudge too) + pi.events.on("claude-account:switched", () => { + if (tuiRef) tuiRef.requestRender(); + }); +} diff --git a/pi/.pi/agent/extensions/postpone.ts b/pi/.pi/agent/extensions/postpone.ts new file mode 100644 index 0000000..aa7ac53 --- /dev/null +++ b/pi/.pi/agent/extensions/postpone.ts @@ -0,0 +1,157 @@ +import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent"; + +/** + * Postpone Extension + * + * Allows users to postpone sending a prompt by X minutes. + * When the timer reaches 0, the prompt is automatically sent. + * + * Usage: + * /postpone + * + * Example: + * /postpone 5 Remember to check the build status + */ + +export default function postponeExtension(pi: ExtensionAPI) { + // Interface for postponed prompt data stored in session + interface PostponedPromptData { + prompt: string; + delayMs: number; + timestamp: number; // When it was postponed + sendAs?: "steer" | "followUp" | "nextTurn"; // How to send it + } + + // Restore postponed prompts from session on startup + pi.on("session_start", async (_event, ctx) => { + // Get all entries from current branch + const branchEntries = ctx.sessionManager.getBranch(); + + // Look for postponed prompt entries + for (const entry of branchEntries) { + if (entry.type === "custom" && entry.customType === "postponed-prompt") { + const data = entry.data as PostponedPromptData | undefined; + if (data) { + const elapsed = Date.now() - data.timestamp; + const remaining = data.delayMs - elapsed; + + // If time is still remaining, set up a new timeout + if (remaining > 0) { + setTimeout(() => { + sendPostponedPrompt(data, ctx); + }, remaining); + } else { + // Time's up, send it now + sendPostponedPrompt(data, ctx); + } + } + } + } + }); + + // Function to actually send a postponed prompt + function sendPostponedPrompt(data: PostponedPromptData, ctx: ExtensionContext) { + // Send as user message + pi.sendUserMessage(data.prompt, { + deliverAs: data.sendAs || "steer" + }); + + // Notify user + ctx.ui.notify("Postponed prompt sent", "info"); + } + + // Register the /postpone command + pi.registerCommand("postpone", { + description: "Postpone sending a prompt by X minutes", + handler: async (args, ctx) => { + // Parse args: first argument should be minutes, rest is the prompt + const parts = args.trim().split(/\s+/); + if (parts.length < 2) { + ctx.ui.notify("Usage: /postpone ", "error"); + return; + } + + const minutes = parseInt(parts[0], 10); + if (isNaN(minutes) || minutes <= 0) { + ctx.ui.notify("Please provide a valid number of minutes", "error"); + return; + } + + const prompt = parts.slice(1).join(" "); + if (!prompt.trim()) { + ctx.ui.notify("Please provide a prompt to postpone", "error"); + return; + } + + const delayMs = minutes * 60 * 1000; + const timestamp = Date.now(); + + // Store in session for persistence + pi.appendEntry("postponed-prompt", { + prompt, + delayMs, + timestamp, + sendAs: "steer" // Default to steer + }); + + // Set up timeout + setTimeout(() => { + sendPostponedPrompt({ prompt, delayMs, timestamp, sendAs: "steer" }, ctx); + }, delayMs); + + ctx.ui.notify(`Prompt postponed for ${minutes} minute(s)`, "info"); + } + }); + + // Optional: Add a command to list pending postponed prompts + pi.registerCommand("postponed", { + description: "List postponed prompts", + handler: async (_args, ctx) => { + const branchEntries = ctx.sessionManager.getBranch(); + const postponed = branchEntries.filter( + entry => entry.type === "custom" && entry.customType === "postponed-prompt" + ); + + if (postponed.length === 0) { + ctx.ui.notify("No postponed prompts", "info"); + return; + } + + let message = "Postponed prompts:\n"; + postponed.forEach((entry, index) => { + const data = entry.data as PostponedPromptData; + const elapsed = Date.now() - data.timestamp; + const remaining = data.delayMs - elapsed; + const minutesLeft = Math.ceil(remaining / 60000); + + message += `${index + 1}. ${data.prompt} (${minutesLeft} minute(s) left)\n`; + }); + + ctx.ui.notify(message.trim(), "info"); + } + }); + + // Optional: Add a command to cancel all postponed prompts + pi.registerCommand("postpone-cancel", { + description: "Cancel all postponed prompts", + handler: async (_args, ctx) => { + const branchEntries = ctx.sessionManager.getBranch(); + let cancelled = 0; + + for (const entry of branchEntries) { + if (entry.type === "custom" && entry.customType === "postponed-prompt") { + // We can't actually remove entries, but we can add a marker + // For simplicity, we'll just notify the user that we've cleared the list + // In a real implementation, we might want to track active timeouts + cancelled++; + } + } + + if (cancelled > 0) { + ctx.ui.notify(`Cancelled ${cancelled} postponed prompt(s)`, "info"); + } else { + ctx.ui.notify("No postponed prompts to cancel", "info"); + } + } + }); +} \ No newline at end of file diff --git a/pi/.pi/agent/extensions/usage-bars/index.ts b/pi/.pi/agent/extensions/usage-bars/index.ts index 7e7b8f0..2468e8b 100644 --- a/pi/.pi/agent/extensions/usage-bars/index.ts +++ b/pi/.pi/agent/extensions/usage-bars/index.ts @@ -347,11 +347,11 @@ export default function (pi: ExtensionAPI) { const session = clampPercent(data.session); const weekly = clampPercent(data.weekly); - let s = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`); - if (data.sessionResetsIn) s += " " + theme.fg("dim", data.sessionResetsIn); - - let w = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`); - if (data.weeklyResetsIn) w += " " + theme.fg("dim", `\u27F3 ${data.weeklyResetsIn}`); + // Time suffixes are intentionally omitted here — footer-display builds + // them dynamically from the sessionResetsAt/weeklyResetsAt timestamps + // emitted via the "usage:update" event, avoiding double-display. + const s = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`); + const w = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`); ctx.ui.setStatus(STATUS_KEY, s + theme.fg("dim", " | ") + w); } @@ -533,6 +533,11 @@ export default function (pi: ExtensionAPI) { pi.on("turn_start", (_event, _ctx) => { ctx = _ctx; updateProviderFrom(_ctx.model); }); + pi.on("turn_end", async (_event, _ctx) => { + ctx = _ctx; + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + }); + pi.on("before_agent_start", async (_event, _ctx) => { ctx = _ctx; await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); diff --git a/pi/.pi/agent/extensions/usage-bars/index.ts.backup b/pi/.pi/agent/extensions/usage-bars/index.ts.backup new file mode 100644 index 0000000..7e7b8f0 --- /dev/null +++ b/pi/.pi/agent/extensions/usage-bars/index.ts.backup @@ -0,0 +1,581 @@ +/** + * Usage Extension - Minimal API usage indicator for pi + * + * Polls Codex, Anthropic, Z.AI, Gemini CLI / Antigravity usage and exposes it + * via two channels: + * • pi.events "usage:update" — for other extensions (e.g. footer-display) + * • ctx.ui.setStatus("usage-bars", …) — formatted S/W braille bars + * + * Rendering / footer layout is handled by the separate footer-display extension. + */ + +import { DynamicBorder, type ExtensionAPI } from "@mariozechner/pi-coding-agent"; +import { + Container, + Input, + Spacer, + Text, + getEditorKeybindings, + type Focusable, +} from "@mariozechner/pi-tui"; +import { + canShowForProvider, + clampPercent, + colorForPercent, + detectProvider, + ensureFreshAuthForProviders, + fetchAllUsages, + fetchClaudeUsage, + fetchCodexUsage, + fetchGoogleUsage, + fetchZaiUsage, + providerToOAuthProviderId, + readAuth, + readUsageCache, + resolveUsageEndpoints, + writeUsageCache, + type OAuthProviderId, + type ProviderKey, + type UsageByProvider, + type UsageData, +} from "./core"; + +const CACHE_TTL_MS = 15 * 60 * 1000; +const ACTIVE_CACHE_TTL_MS = 3 * 60 * 1000; +const STREAMING_POLL_INTERVAL_MS = 2 * 60 * 1000; +const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000; + +const STATUS_KEY = "usage-bars"; + +// --------------------------------------------------------------------------- +// Braille gradient bar (⣀ ⣄ ⣤ ⣦ ⣶ ⣷ ⣿) +// --------------------------------------------------------------------------- +const BRAILLE_GRADIENT = "\u28C0\u28C4\u28E4\u28E6\u28F6\u28F7\u28FF"; +const BRAILLE_EMPTY = "\u28C0"; +const BAR_WIDTH = 5; + +function renderBrailleBar(theme: any, value: number, width = BAR_WIDTH): string { + const v = clampPercent(value); + const levels = BRAILLE_GRADIENT.length - 1; + const totalSteps = width * levels; + const filledSteps = Math.round((v / 100) * totalSteps); + const full = Math.floor(filledSteps / levels); + const partial = filledSteps % levels; + const empty = width - full - (partial ? 1 : 0); + const color = colorForPercent(v); + const filled = BRAILLE_GRADIENT[BRAILLE_GRADIENT.length - 1]!.repeat(Math.max(0, full)); + const partialChar = partial ? BRAILLE_GRADIENT[partial]! : ""; + const emptyChars = BRAILLE_EMPTY.repeat(Math.max(0, empty)); + return theme.fg(color, filled + partialChar) + theme.fg("dim", emptyChars); +} + +function renderBrailleBarWide(theme: any, value: number): string { + return renderBrailleBar(theme, value, 12); +} + +const PROVIDER_LABELS: Record = { + codex: "Codex", + claude: "Claude", + zai: "Z.AI", + gemini: "Gemini", + antigravity: "Antigravity", +}; + +// --------------------------------------------------------------------------- +// /usage command popup +// --------------------------------------------------------------------------- +interface SubscriptionItem { + name: string; + provider: ProviderKey; + data: UsageData | null; + isActive: boolean; +} + +class UsageSelectorComponent extends Container implements Focusable { + private searchInput: Input; + private listContainer: Container; + private hintText: Text; + private tui: any; + private theme: any; + private onCancelCallback: () => void; + private allItems: SubscriptionItem[] = []; + private filteredItems: SubscriptionItem[] = []; + private selectedIndex = 0; + private loading = true; + private activeProvider: ProviderKey | null; + private fetchAllFn: () => Promise; + private _focused = false; + + get focused(): boolean { return this._focused; } + set focused(value: boolean) { this._focused = value; this.searchInput.focused = value; } + + constructor( + tui: any, + theme: any, + activeProvider: ProviderKey | null, + fetchAll: () => Promise, + onCancel: () => void, + ) { + super(); + this.tui = tui; + this.theme = theme; + this.activeProvider = activeProvider; + this.fetchAllFn = fetchAll; + this.onCancelCallback = onCancel; + + this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s))); + this.addChild(new Spacer(1)); + this.hintText = new Text(theme.fg("dim", "Fetching usage from all providers…"), 0, 0); + this.addChild(this.hintText); + this.addChild(new Spacer(1)); + this.searchInput = new Input(); + this.addChild(this.searchInput); + this.addChild(new Spacer(1)); + this.listContainer = new Container(); + this.addChild(this.listContainer); + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s))); + + this.fetchAllFn() + .then((results) => { + this.loading = false; + this.buildItems(results); + this.updateList(); + this.hintText.setText( + theme.fg("muted", "Only showing providers with credentials. ") + + theme.fg("dim", "✓ = active provider"), + ); + this.tui.requestRender(); + }) + .catch(() => { + this.loading = false; + this.hintText.setText(theme.fg("error", "Failed to fetch usage data")); + this.tui.requestRender(); + }); + + this.updateList(); + } + + private buildItems(results: UsageByProvider) { + const providers: Array<{ key: ProviderKey; name: string }> = [ + { key: "codex", name: "Codex" }, + { key: "claude", name: "Claude" }, + { key: "zai", name: "Z.AI" }, + { key: "gemini", name: "Gemini" }, + { key: "antigravity", name: "Antigravity" }, + ]; + this.allItems = []; + for (const p of providers) { + if (results[p.key] !== null) { + this.allItems.push({ + name: p.name, + provider: p.key, + data: results[p.key], + isActive: this.activeProvider === p.key, + }); + } + } + this.filteredItems = this.allItems; + this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1)); + } + + private filterItems(query: string) { + if (!query) { + this.filteredItems = this.allItems; + } else { + const q = query.toLowerCase(); + this.filteredItems = this.allItems.filter( + (item) => item.name.toLowerCase().includes(q) || item.provider.toLowerCase().includes(q), + ); + } + this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1)); + } + + private renderItem(item: SubscriptionItem, isSelected: boolean) { + const t = this.theme; + const pointer = isSelected ? t.fg("accent", "→ ") : " "; + const activeBadge = item.isActive ? t.fg("success", " ✓") : ""; + const name = isSelected ? t.fg("accent", t.bold(item.name)) : item.name; + this.listContainer.addChild(new Text(`${pointer}${name}${activeBadge}`, 0, 0)); + const indent = " "; + + if (!item.data) { + this.listContainer.addChild(new Text(indent + t.fg("dim", "No credentials"), 0, 0)); + } else if (item.data.error) { + this.listContainer.addChild(new Text(indent + t.fg("error", item.data.error), 0, 0)); + } else { + const session = clampPercent(item.data.session); + const weekly = clampPercent(item.data.weekly); + const sessionReset = item.data.sessionResetsIn + ? t.fg("dim", ` resets in ${item.data.sessionResetsIn}`) : ""; + const weeklyReset = item.data.weeklyResetsIn + ? t.fg("dim", ` resets in ${item.data.weeklyResetsIn}`) : ""; + + this.listContainer.addChild(new Text( + indent + t.fg("muted", "Session ") + + renderBrailleBarWide(t, session) + " " + + t.fg(colorForPercent(session), `${session}%`.padStart(4)) + sessionReset, + 0, 0, + )); + this.listContainer.addChild(new Text( + indent + t.fg("muted", "Weekly ") + + renderBrailleBarWide(t, weekly) + " " + + t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) + weeklyReset, + 0, 0, + )); + + if (typeof item.data.extraSpend === "number" && typeof item.data.extraLimit === "number") { + this.listContainer.addChild(new Text( + indent + t.fg("muted", "Extra ") + + t.fg("dim", `$${item.data.extraSpend.toFixed(2)} / $${item.data.extraLimit}`), + 0, 0, + )); + } + } + this.listContainer.addChild(new Spacer(1)); + } + + private updateList() { + this.listContainer.clear(); + if (this.loading) { + this.listContainer.addChild(new Text(this.theme.fg("muted", " Loading…"), 0, 0)); + return; + } + if (this.filteredItems.length === 0) { + this.listContainer.addChild(new Text(this.theme.fg("muted", " No matching providers"), 0, 0)); + return; + } + for (let i = 0; i < this.filteredItems.length; i++) { + this.renderItem(this.filteredItems[i]!, i === this.selectedIndex); + } + } + + handleInput(keyData: string): void { + const kb = getEditorKeybindings(); + if (kb.matches(keyData, "selectUp")) { + if (this.filteredItems.length === 0) return; + this.selectedIndex = this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1; + this.updateList(); return; + } + if (kb.matches(keyData, "selectDown")) { + if (this.filteredItems.length === 0) return; + this.selectedIndex = this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1; + this.updateList(); return; + } + if (kb.matches(keyData, "selectCancel") || kb.matches(keyData, "selectConfirm")) { + this.onCancelCallback(); return; + } + this.searchInput.handleInput(keyData); + this.filterItems(this.searchInput.getValue()); + this.updateList(); + } +} + +// --------------------------------------------------------------------------- +// Extension state +// --------------------------------------------------------------------------- +interface UsageState extends UsageByProvider { + lastPoll: number; + activeProvider: ProviderKey | null; +} + +interface PollOptions { + cacheTtl?: number; + forceFresh?: boolean; +} + +export default function (pi: ExtensionAPI) { + const endpoints = resolveUsageEndpoints(); + const state: UsageState = { + codex: null, claude: null, zai: null, gemini: null, antigravity: null, + lastPoll: 0, activeProvider: null, + }; + + let pollInFlight: Promise | null = null; + let pollQueued = false; + let pollStartedAt = 0; + let streamingTimer: ReturnType | null = null; + let ctx: any = null; + + // --------------------------------------------------------------------------- + // Status update + // --------------------------------------------------------------------------- + function updateStatus() { + const active = state.activeProvider; + const data = active ? state[active] : null; + + // Always emit event for other extensions (e.g. footer-display) + if (data && !data.error) { + pi.events.emit("usage:update", { + session: data.session, + weekly: data.weekly, + sessionResetsIn: data.sessionResetsIn, + sessionResetsAt: data.sessionResetsAt, + weeklyResetsIn: data.weeklyResetsIn, + }); + } + + if (!ctx?.hasUI) return; + + const theme = ctx.ui.theme; + + if (!active) { + ctx.ui.setStatus(STATUS_KEY, undefined); + return; + } + + const auth = readAuth(); + if (!canShowForProvider(active, auth, endpoints)) { + ctx.ui.setStatus(STATUS_KEY, undefined); + return; + } + + if (!data) { + ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", "loading\u2026")); + return; + } + + if (data.error) { + const cache = readUsageCache(); + const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0; + const note = blockedUntil > Date.now() + ? ` \u2014 retry in ${Math.ceil((blockedUntil - Date.now()) / 60000)}m` : ""; + ctx.ui.setStatus(STATUS_KEY, theme.fg("warning", `${PROVIDER_LABELS[active]} unavailable${note}`)); + return; + } + + const session = clampPercent(data.session); + const weekly = clampPercent(data.weekly); + + let s = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`); + if (data.sessionResetsIn) s += " " + theme.fg("dim", data.sessionResetsIn); + + let w = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`); + if (data.weeklyResetsIn) w += " " + theme.fg("dim", `\u27F3 ${data.weeklyResetsIn}`); + + ctx.ui.setStatus(STATUS_KEY, s + theme.fg("dim", " | ") + w); + } + + function updateProviderFrom(modelLike: any): boolean { + const previous = state.activeProvider; + state.activeProvider = detectProvider(modelLike); + if (previous !== state.activeProvider) { updateStatus(); return true; } + return false; + } + + // --------------------------------------------------------------------------- + // Polling + // --------------------------------------------------------------------------- + async function runPollInner(options: PollOptions = {}) { + const auth = readAuth(); + const active = state.activeProvider; + + if (!canShowForProvider(active, auth, endpoints) || !auth || !active) { + state.lastPoll = Date.now(); updateStatus(); return; + } + + const cache = readUsageCache(); + const now = Date.now(); + const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS; + + const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0; + if (now < blockedUntil) { + if (cache?.data?.[active]) { + state[active] = cache.data[active]!; + } else { + // Rate-limited but no cached data — show a meaningful status instead + // of leaving state null (which shows eternal "loading…"). + const retryMin = Math.ceil((blockedUntil - now) / 60000); + state[active] = { session: 0, weekly: 0, error: `rate limited (retry in ${retryMin}m)` }; + } + state.lastPoll = now; updateStatus(); return; + } + + if (!options.forceFresh && cache && now - cache.timestamp < cacheTtl && cache.data?.[active]) { + state[active] = cache.data[active]!; + state.lastPoll = now; updateStatus(); return; + } + + const oauthId = providerToOAuthProviderId(active); + let effectiveAuth = auth; + if (oauthId && active !== "zai") { + const creds = auth[oauthId as keyof typeof auth] as + | { access?: string; refresh?: string; expires?: number } | undefined; + const expires = typeof creds?.expires === "number" ? creds.expires : 0; + const tokenExpiredOrMissing = !creds?.access || (expires > 0 && Date.now() + 60_000 >= expires); + if (tokenExpiredOrMissing && creds?.refresh) { + try { + const refreshPromise = ensureFreshAuthForProviders([oauthId as OAuthProviderId], { auth, persist: true }); + const timeoutPromise = new Promise((_, reject) => + setTimeout(() => reject(new Error("OAuth refresh timeout")), 15_000), + ); + const refreshed = await Promise.race([refreshPromise, timeoutPromise]); + if (refreshed.auth) effectiveAuth = refreshed.auth; + } catch {} + } + } + + let result: UsageData; + if (active === "codex") { + const access = effectiveAuth["openai-codex"]?.access; + result = access ? await fetchCodexUsage(access) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } else if (active === "claude") { + const access = effectiveAuth.anthropic?.access; + result = access ? await fetchClaudeUsage(access) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } else if (active === "zai") { + const token = effectiveAuth.zai?.access || effectiveAuth.zai?.key; + result = token ? await fetchZaiUsage(token, { endpoints }) + : { session: 0, weekly: 0, error: "missing token (try /login again)" }; + } else if (active === "gemini") { + const creds = effectiveAuth["google-gemini-cli"]; + result = creds?.access + ? await fetchGoogleUsage(creds.access, endpoints.gemini, creds.projectId, "gemini", { endpoints }) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } else { + const creds = effectiveAuth["google-antigravity"]; + result = creds?.access + ? await fetchGoogleUsage(creds.access, endpoints.antigravity, creds.projectId, "antigravity", { endpoints }) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } + + state[active] = result; + + if (result.error) { + if (result.error === "HTTP 429") { + const nextCache: import("./core").UsageCache = { + timestamp: cache?.timestamp ?? now, + data: { ...(cache?.data ?? {}) }, + rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}), [active]: now + RATE_LIMITED_BACKOFF_MS }, + }; + writeUsageCache(nextCache); + } + } else { + const nextCache: import("./core").UsageCache = { + timestamp: now, + data: { ...(cache?.data ?? {}), [active]: result }, + rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) }, + }; + delete nextCache.rateLimitedUntil![active]; + writeUsageCache(nextCache); + } + + state.lastPoll = now; + updateStatus(); + } + + async function runPoll(options: PollOptions = {}): Promise { + const timeout = new Promise((_, reject) => + setTimeout(() => reject(new Error("runPoll timeout")), 25_000), + ); + await Promise.race([runPollInner(options), timeout]); + } + + const POLL_TIMEOUT_MS = 30_000; + + async function poll(options: PollOptions = {}) { + // If a previous poll has been running longer than POLL_TIMEOUT_MS, abandon it + // so we don't queue forever behind a stuck request. + if (pollInFlight && pollStartedAt > 0 && Date.now() - pollStartedAt > POLL_TIMEOUT_MS) { + pollInFlight = null; + pollQueued = false; + const active = state.activeProvider; + if (active && !state[active]) { + state[active] = { session: 0, weekly: 0, error: "poll timeout" }; + updateStatus(); + } + } + + if (pollInFlight) { pollQueued = true; await pollInFlight; return; } + do { + pollQueued = false; + pollStartedAt = Date.now(); + pollInFlight = runPoll(options).catch(() => { + // If runPoll threw, ensure we don't leave status stuck at "loading…" + const active = state.activeProvider; + if (active && !state[active]) { + state[active] = { session: 0, weekly: 0, error: "poll failed" }; + updateStatus(); + } + }).finally(() => { pollInFlight = null; pollStartedAt = 0; }); + await pollInFlight; + } while (pollQueued); + } + + function startStreamingTimer() { + if (streamingTimer !== null) return; + streamingTimer = setInterval(() => { void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); }, STREAMING_POLL_INTERVAL_MS); + } + + function stopStreamingTimer() { + if (streamingTimer !== null) { clearInterval(streamingTimer); streamingTimer = null; } + } + + // ── Lifecycle ──────────────────────────────────────────────────────────── + + pi.on("session_start", async (_event, _ctx) => { + ctx = _ctx; + updateProviderFrom(_ctx.model); + await poll(); + }); + + pi.on("session_shutdown", async (_event, _ctx) => { + stopStreamingTimer(); + if (_ctx?.hasUI) _ctx.ui.setStatus(STATUS_KEY, undefined); + }); + + pi.on("model_select", async (event, _ctx) => { + ctx = _ctx; + const changed = updateProviderFrom(event.model ?? _ctx.model); + if (changed) await poll(); + }); + + pi.on("turn_start", (_event, _ctx) => { ctx = _ctx; updateProviderFrom(_ctx.model); }); + + pi.on("before_agent_start", async (_event, _ctx) => { + ctx = _ctx; + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + }); + + pi.on("agent_start", (_event, _ctx) => { ctx = _ctx; startStreamingTimer(); }); + + pi.on("agent_end", async (_event, _ctx) => { + ctx = _ctx; + stopStreamingTimer(); + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + }); + + pi.events.on("claude-account:switched", () => { + const cache = readUsageCache(); + if (cache?.data?.claude) { + const nextCache: import("./core").UsageCache = { ...cache, data: { ...cache.data } }; + delete nextCache.data.claude; + writeUsageCache(nextCache); + } + void poll({ forceFresh: true }); + }); + + // ── /usage command ─────────────────────────────────────────────────────── + + pi.registerCommand("usage", { + description: "Show API usage for all subscriptions", + handler: async (_args, _ctx) => { + ctx = _ctx; + updateProviderFrom(_ctx.model); + try { + if (_ctx?.hasUI) { + await _ctx.ui.custom((tui, theme, _keybindings, done) => { + return new UsageSelectorComponent( + tui, theme, state.activeProvider, + () => fetchAllUsages({ endpoints }), + () => done(), + ); + }); + } + } finally { + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + } + }, + }); +} diff --git a/pi/.pi/agent/extensions/usage-bars/index.ts.backup2 b/pi/.pi/agent/extensions/usage-bars/index.ts.backup2 new file mode 100644 index 0000000..7e7b8f0 --- /dev/null +++ b/pi/.pi/agent/extensions/usage-bars/index.ts.backup2 @@ -0,0 +1,581 @@ +/** + * Usage Extension - Minimal API usage indicator for pi + * + * Polls Codex, Anthropic, Z.AI, Gemini CLI / Antigravity usage and exposes it + * via two channels: + * • pi.events "usage:update" — for other extensions (e.g. footer-display) + * • ctx.ui.setStatus("usage-bars", …) — formatted S/W braille bars + * + * Rendering / footer layout is handled by the separate footer-display extension. + */ + +import { DynamicBorder, type ExtensionAPI } from "@mariozechner/pi-coding-agent"; +import { + Container, + Input, + Spacer, + Text, + getEditorKeybindings, + type Focusable, +} from "@mariozechner/pi-tui"; +import { + canShowForProvider, + clampPercent, + colorForPercent, + detectProvider, + ensureFreshAuthForProviders, + fetchAllUsages, + fetchClaudeUsage, + fetchCodexUsage, + fetchGoogleUsage, + fetchZaiUsage, + providerToOAuthProviderId, + readAuth, + readUsageCache, + resolveUsageEndpoints, + writeUsageCache, + type OAuthProviderId, + type ProviderKey, + type UsageByProvider, + type UsageData, +} from "./core"; + +const CACHE_TTL_MS = 15 * 60 * 1000; +const ACTIVE_CACHE_TTL_MS = 3 * 60 * 1000; +const STREAMING_POLL_INTERVAL_MS = 2 * 60 * 1000; +const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000; + +const STATUS_KEY = "usage-bars"; + +// --------------------------------------------------------------------------- +// Braille gradient bar (⣀ ⣄ ⣤ ⣦ ⣶ ⣷ ⣿) +// --------------------------------------------------------------------------- +const BRAILLE_GRADIENT = "\u28C0\u28C4\u28E4\u28E6\u28F6\u28F7\u28FF"; +const BRAILLE_EMPTY = "\u28C0"; +const BAR_WIDTH = 5; + +function renderBrailleBar(theme: any, value: number, width = BAR_WIDTH): string { + const v = clampPercent(value); + const levels = BRAILLE_GRADIENT.length - 1; + const totalSteps = width * levels; + const filledSteps = Math.round((v / 100) * totalSteps); + const full = Math.floor(filledSteps / levels); + const partial = filledSteps % levels; + const empty = width - full - (partial ? 1 : 0); + const color = colorForPercent(v); + const filled = BRAILLE_GRADIENT[BRAILLE_GRADIENT.length - 1]!.repeat(Math.max(0, full)); + const partialChar = partial ? BRAILLE_GRADIENT[partial]! : ""; + const emptyChars = BRAILLE_EMPTY.repeat(Math.max(0, empty)); + return theme.fg(color, filled + partialChar) + theme.fg("dim", emptyChars); +} + +function renderBrailleBarWide(theme: any, value: number): string { + return renderBrailleBar(theme, value, 12); +} + +const PROVIDER_LABELS: Record = { + codex: "Codex", + claude: "Claude", + zai: "Z.AI", + gemini: "Gemini", + antigravity: "Antigravity", +}; + +// --------------------------------------------------------------------------- +// /usage command popup +// --------------------------------------------------------------------------- +interface SubscriptionItem { + name: string; + provider: ProviderKey; + data: UsageData | null; + isActive: boolean; +} + +class UsageSelectorComponent extends Container implements Focusable { + private searchInput: Input; + private listContainer: Container; + private hintText: Text; + private tui: any; + private theme: any; + private onCancelCallback: () => void; + private allItems: SubscriptionItem[] = []; + private filteredItems: SubscriptionItem[] = []; + private selectedIndex = 0; + private loading = true; + private activeProvider: ProviderKey | null; + private fetchAllFn: () => Promise; + private _focused = false; + + get focused(): boolean { return this._focused; } + set focused(value: boolean) { this._focused = value; this.searchInput.focused = value; } + + constructor( + tui: any, + theme: any, + activeProvider: ProviderKey | null, + fetchAll: () => Promise, + onCancel: () => void, + ) { + super(); + this.tui = tui; + this.theme = theme; + this.activeProvider = activeProvider; + this.fetchAllFn = fetchAll; + this.onCancelCallback = onCancel; + + this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s))); + this.addChild(new Spacer(1)); + this.hintText = new Text(theme.fg("dim", "Fetching usage from all providers…"), 0, 0); + this.addChild(this.hintText); + this.addChild(new Spacer(1)); + this.searchInput = new Input(); + this.addChild(this.searchInput); + this.addChild(new Spacer(1)); + this.listContainer = new Container(); + this.addChild(this.listContainer); + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s))); + + this.fetchAllFn() + .then((results) => { + this.loading = false; + this.buildItems(results); + this.updateList(); + this.hintText.setText( + theme.fg("muted", "Only showing providers with credentials. ") + + theme.fg("dim", "✓ = active provider"), + ); + this.tui.requestRender(); + }) + .catch(() => { + this.loading = false; + this.hintText.setText(theme.fg("error", "Failed to fetch usage data")); + this.tui.requestRender(); + }); + + this.updateList(); + } + + private buildItems(results: UsageByProvider) { + const providers: Array<{ key: ProviderKey; name: string }> = [ + { key: "codex", name: "Codex" }, + { key: "claude", name: "Claude" }, + { key: "zai", name: "Z.AI" }, + { key: "gemini", name: "Gemini" }, + { key: "antigravity", name: "Antigravity" }, + ]; + this.allItems = []; + for (const p of providers) { + if (results[p.key] !== null) { + this.allItems.push({ + name: p.name, + provider: p.key, + data: results[p.key], + isActive: this.activeProvider === p.key, + }); + } + } + this.filteredItems = this.allItems; + this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1)); + } + + private filterItems(query: string) { + if (!query) { + this.filteredItems = this.allItems; + } else { + const q = query.toLowerCase(); + this.filteredItems = this.allItems.filter( + (item) => item.name.toLowerCase().includes(q) || item.provider.toLowerCase().includes(q), + ); + } + this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1)); + } + + private renderItem(item: SubscriptionItem, isSelected: boolean) { + const t = this.theme; + const pointer = isSelected ? t.fg("accent", "→ ") : " "; + const activeBadge = item.isActive ? t.fg("success", " ✓") : ""; + const name = isSelected ? t.fg("accent", t.bold(item.name)) : item.name; + this.listContainer.addChild(new Text(`${pointer}${name}${activeBadge}`, 0, 0)); + const indent = " "; + + if (!item.data) { + this.listContainer.addChild(new Text(indent + t.fg("dim", "No credentials"), 0, 0)); + } else if (item.data.error) { + this.listContainer.addChild(new Text(indent + t.fg("error", item.data.error), 0, 0)); + } else { + const session = clampPercent(item.data.session); + const weekly = clampPercent(item.data.weekly); + const sessionReset = item.data.sessionResetsIn + ? t.fg("dim", ` resets in ${item.data.sessionResetsIn}`) : ""; + const weeklyReset = item.data.weeklyResetsIn + ? t.fg("dim", ` resets in ${item.data.weeklyResetsIn}`) : ""; + + this.listContainer.addChild(new Text( + indent + t.fg("muted", "Session ") + + renderBrailleBarWide(t, session) + " " + + t.fg(colorForPercent(session), `${session}%`.padStart(4)) + sessionReset, + 0, 0, + )); + this.listContainer.addChild(new Text( + indent + t.fg("muted", "Weekly ") + + renderBrailleBarWide(t, weekly) + " " + + t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) + weeklyReset, + 0, 0, + )); + + if (typeof item.data.extraSpend === "number" && typeof item.data.extraLimit === "number") { + this.listContainer.addChild(new Text( + indent + t.fg("muted", "Extra ") + + t.fg("dim", `$${item.data.extraSpend.toFixed(2)} / $${item.data.extraLimit}`), + 0, 0, + )); + } + } + this.listContainer.addChild(new Spacer(1)); + } + + private updateList() { + this.listContainer.clear(); + if (this.loading) { + this.listContainer.addChild(new Text(this.theme.fg("muted", " Loading…"), 0, 0)); + return; + } + if (this.filteredItems.length === 0) { + this.listContainer.addChild(new Text(this.theme.fg("muted", " No matching providers"), 0, 0)); + return; + } + for (let i = 0; i < this.filteredItems.length; i++) { + this.renderItem(this.filteredItems[i]!, i === this.selectedIndex); + } + } + + handleInput(keyData: string): void { + const kb = getEditorKeybindings(); + if (kb.matches(keyData, "selectUp")) { + if (this.filteredItems.length === 0) return; + this.selectedIndex = this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1; + this.updateList(); return; + } + if (kb.matches(keyData, "selectDown")) { + if (this.filteredItems.length === 0) return; + this.selectedIndex = this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1; + this.updateList(); return; + } + if (kb.matches(keyData, "selectCancel") || kb.matches(keyData, "selectConfirm")) { + this.onCancelCallback(); return; + } + this.searchInput.handleInput(keyData); + this.filterItems(this.searchInput.getValue()); + this.updateList(); + } +} + +// --------------------------------------------------------------------------- +// Extension state +// --------------------------------------------------------------------------- +interface UsageState extends UsageByProvider { + lastPoll: number; + activeProvider: ProviderKey | null; +} + +interface PollOptions { + cacheTtl?: number; + forceFresh?: boolean; +} + +export default function (pi: ExtensionAPI) { + const endpoints = resolveUsageEndpoints(); + const state: UsageState = { + codex: null, claude: null, zai: null, gemini: null, antigravity: null, + lastPoll: 0, activeProvider: null, + }; + + let pollInFlight: Promise | null = null; + let pollQueued = false; + let pollStartedAt = 0; + let streamingTimer: ReturnType | null = null; + let ctx: any = null; + + // --------------------------------------------------------------------------- + // Status update + // --------------------------------------------------------------------------- + function updateStatus() { + const active = state.activeProvider; + const data = active ? state[active] : null; + + // Always emit event for other extensions (e.g. footer-display) + if (data && !data.error) { + pi.events.emit("usage:update", { + session: data.session, + weekly: data.weekly, + sessionResetsIn: data.sessionResetsIn, + sessionResetsAt: data.sessionResetsAt, + weeklyResetsIn: data.weeklyResetsIn, + }); + } + + if (!ctx?.hasUI) return; + + const theme = ctx.ui.theme; + + if (!active) { + ctx.ui.setStatus(STATUS_KEY, undefined); + return; + } + + const auth = readAuth(); + if (!canShowForProvider(active, auth, endpoints)) { + ctx.ui.setStatus(STATUS_KEY, undefined); + return; + } + + if (!data) { + ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", "loading\u2026")); + return; + } + + if (data.error) { + const cache = readUsageCache(); + const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0; + const note = blockedUntil > Date.now() + ? ` \u2014 retry in ${Math.ceil((blockedUntil - Date.now()) / 60000)}m` : ""; + ctx.ui.setStatus(STATUS_KEY, theme.fg("warning", `${PROVIDER_LABELS[active]} unavailable${note}`)); + return; + } + + const session = clampPercent(data.session); + const weekly = clampPercent(data.weekly); + + let s = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`); + if (data.sessionResetsIn) s += " " + theme.fg("dim", data.sessionResetsIn); + + let w = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`); + if (data.weeklyResetsIn) w += " " + theme.fg("dim", `\u27F3 ${data.weeklyResetsIn}`); + + ctx.ui.setStatus(STATUS_KEY, s + theme.fg("dim", " | ") + w); + } + + function updateProviderFrom(modelLike: any): boolean { + const previous = state.activeProvider; + state.activeProvider = detectProvider(modelLike); + if (previous !== state.activeProvider) { updateStatus(); return true; } + return false; + } + + // --------------------------------------------------------------------------- + // Polling + // --------------------------------------------------------------------------- + async function runPollInner(options: PollOptions = {}) { + const auth = readAuth(); + const active = state.activeProvider; + + if (!canShowForProvider(active, auth, endpoints) || !auth || !active) { + state.lastPoll = Date.now(); updateStatus(); return; + } + + const cache = readUsageCache(); + const now = Date.now(); + const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS; + + const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0; + if (now < blockedUntil) { + if (cache?.data?.[active]) { + state[active] = cache.data[active]!; + } else { + // Rate-limited but no cached data — show a meaningful status instead + // of leaving state null (which shows eternal "loading…"). + const retryMin = Math.ceil((blockedUntil - now) / 60000); + state[active] = { session: 0, weekly: 0, error: `rate limited (retry in ${retryMin}m)` }; + } + state.lastPoll = now; updateStatus(); return; + } + + if (!options.forceFresh && cache && now - cache.timestamp < cacheTtl && cache.data?.[active]) { + state[active] = cache.data[active]!; + state.lastPoll = now; updateStatus(); return; + } + + const oauthId = providerToOAuthProviderId(active); + let effectiveAuth = auth; + if (oauthId && active !== "zai") { + const creds = auth[oauthId as keyof typeof auth] as + | { access?: string; refresh?: string; expires?: number } | undefined; + const expires = typeof creds?.expires === "number" ? creds.expires : 0; + const tokenExpiredOrMissing = !creds?.access || (expires > 0 && Date.now() + 60_000 >= expires); + if (tokenExpiredOrMissing && creds?.refresh) { + try { + const refreshPromise = ensureFreshAuthForProviders([oauthId as OAuthProviderId], { auth, persist: true }); + const timeoutPromise = new Promise((_, reject) => + setTimeout(() => reject(new Error("OAuth refresh timeout")), 15_000), + ); + const refreshed = await Promise.race([refreshPromise, timeoutPromise]); + if (refreshed.auth) effectiveAuth = refreshed.auth; + } catch {} + } + } + + let result: UsageData; + if (active === "codex") { + const access = effectiveAuth["openai-codex"]?.access; + result = access ? await fetchCodexUsage(access) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } else if (active === "claude") { + const access = effectiveAuth.anthropic?.access; + result = access ? await fetchClaudeUsage(access) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } else if (active === "zai") { + const token = effectiveAuth.zai?.access || effectiveAuth.zai?.key; + result = token ? await fetchZaiUsage(token, { endpoints }) + : { session: 0, weekly: 0, error: "missing token (try /login again)" }; + } else if (active === "gemini") { + const creds = effectiveAuth["google-gemini-cli"]; + result = creds?.access + ? await fetchGoogleUsage(creds.access, endpoints.gemini, creds.projectId, "gemini", { endpoints }) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } else { + const creds = effectiveAuth["google-antigravity"]; + result = creds?.access + ? await fetchGoogleUsage(creds.access, endpoints.antigravity, creds.projectId, "antigravity", { endpoints }) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } + + state[active] = result; + + if (result.error) { + if (result.error === "HTTP 429") { + const nextCache: import("./core").UsageCache = { + timestamp: cache?.timestamp ?? now, + data: { ...(cache?.data ?? {}) }, + rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}), [active]: now + RATE_LIMITED_BACKOFF_MS }, + }; + writeUsageCache(nextCache); + } + } else { + const nextCache: import("./core").UsageCache = { + timestamp: now, + data: { ...(cache?.data ?? {}), [active]: result }, + rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) }, + }; + delete nextCache.rateLimitedUntil![active]; + writeUsageCache(nextCache); + } + + state.lastPoll = now; + updateStatus(); + } + + async function runPoll(options: PollOptions = {}): Promise { + const timeout = new Promise((_, reject) => + setTimeout(() => reject(new Error("runPoll timeout")), 25_000), + ); + await Promise.race([runPollInner(options), timeout]); + } + + const POLL_TIMEOUT_MS = 30_000; + + async function poll(options: PollOptions = {}) { + // If a previous poll has been running longer than POLL_TIMEOUT_MS, abandon it + // so we don't queue forever behind a stuck request. + if (pollInFlight && pollStartedAt > 0 && Date.now() - pollStartedAt > POLL_TIMEOUT_MS) { + pollInFlight = null; + pollQueued = false; + const active = state.activeProvider; + if (active && !state[active]) { + state[active] = { session: 0, weekly: 0, error: "poll timeout" }; + updateStatus(); + } + } + + if (pollInFlight) { pollQueued = true; await pollInFlight; return; } + do { + pollQueued = false; + pollStartedAt = Date.now(); + pollInFlight = runPoll(options).catch(() => { + // If runPoll threw, ensure we don't leave status stuck at "loading…" + const active = state.activeProvider; + if (active && !state[active]) { + state[active] = { session: 0, weekly: 0, error: "poll failed" }; + updateStatus(); + } + }).finally(() => { pollInFlight = null; pollStartedAt = 0; }); + await pollInFlight; + } while (pollQueued); + } + + function startStreamingTimer() { + if (streamingTimer !== null) return; + streamingTimer = setInterval(() => { void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); }, STREAMING_POLL_INTERVAL_MS); + } + + function stopStreamingTimer() { + if (streamingTimer !== null) { clearInterval(streamingTimer); streamingTimer = null; } + } + + // ── Lifecycle ──────────────────────────────────────────────────────────── + + pi.on("session_start", async (_event, _ctx) => { + ctx = _ctx; + updateProviderFrom(_ctx.model); + await poll(); + }); + + pi.on("session_shutdown", async (_event, _ctx) => { + stopStreamingTimer(); + if (_ctx?.hasUI) _ctx.ui.setStatus(STATUS_KEY, undefined); + }); + + pi.on("model_select", async (event, _ctx) => { + ctx = _ctx; + const changed = updateProviderFrom(event.model ?? _ctx.model); + if (changed) await poll(); + }); + + pi.on("turn_start", (_event, _ctx) => { ctx = _ctx; updateProviderFrom(_ctx.model); }); + + pi.on("before_agent_start", async (_event, _ctx) => { + ctx = _ctx; + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + }); + + pi.on("agent_start", (_event, _ctx) => { ctx = _ctx; startStreamingTimer(); }); + + pi.on("agent_end", async (_event, _ctx) => { + ctx = _ctx; + stopStreamingTimer(); + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + }); + + pi.events.on("claude-account:switched", () => { + const cache = readUsageCache(); + if (cache?.data?.claude) { + const nextCache: import("./core").UsageCache = { ...cache, data: { ...cache.data } }; + delete nextCache.data.claude; + writeUsageCache(nextCache); + } + void poll({ forceFresh: true }); + }); + + // ── /usage command ─────────────────────────────────────────────────────── + + pi.registerCommand("usage", { + description: "Show API usage for all subscriptions", + handler: async (_args, _ctx) => { + ctx = _ctx; + updateProviderFrom(_ctx.model); + try { + if (_ctx?.hasUI) { + await _ctx.ui.custom((tui, theme, _keybindings, done) => { + return new UsageSelectorComponent( + tui, theme, state.activeProvider, + () => fetchAllUsages({ endpoints }), + () => done(), + ); + }); + } + } finally { + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + } + }, + }); +} diff --git a/pi/.pi/agent/extensions/usage-bars/index.ts.backup3 b/pi/.pi/agent/extensions/usage-bars/index.ts.backup3 new file mode 100644 index 0000000..7e7b8f0 --- /dev/null +++ b/pi/.pi/agent/extensions/usage-bars/index.ts.backup3 @@ -0,0 +1,581 @@ +/** + * Usage Extension - Minimal API usage indicator for pi + * + * Polls Codex, Anthropic, Z.AI, Gemini CLI / Antigravity usage and exposes it + * via two channels: + * • pi.events "usage:update" — for other extensions (e.g. footer-display) + * • ctx.ui.setStatus("usage-bars", …) — formatted S/W braille bars + * + * Rendering / footer layout is handled by the separate footer-display extension. + */ + +import { DynamicBorder, type ExtensionAPI } from "@mariozechner/pi-coding-agent"; +import { + Container, + Input, + Spacer, + Text, + getEditorKeybindings, + type Focusable, +} from "@mariozechner/pi-tui"; +import { + canShowForProvider, + clampPercent, + colorForPercent, + detectProvider, + ensureFreshAuthForProviders, + fetchAllUsages, + fetchClaudeUsage, + fetchCodexUsage, + fetchGoogleUsage, + fetchZaiUsage, + providerToOAuthProviderId, + readAuth, + readUsageCache, + resolveUsageEndpoints, + writeUsageCache, + type OAuthProviderId, + type ProviderKey, + type UsageByProvider, + type UsageData, +} from "./core"; + +const CACHE_TTL_MS = 15 * 60 * 1000; +const ACTIVE_CACHE_TTL_MS = 3 * 60 * 1000; +const STREAMING_POLL_INTERVAL_MS = 2 * 60 * 1000; +const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000; + +const STATUS_KEY = "usage-bars"; + +// --------------------------------------------------------------------------- +// Braille gradient bar (⣀ ⣄ ⣤ ⣦ ⣶ ⣷ ⣿) +// --------------------------------------------------------------------------- +const BRAILLE_GRADIENT = "\u28C0\u28C4\u28E4\u28E6\u28F6\u28F7\u28FF"; +const BRAILLE_EMPTY = "\u28C0"; +const BAR_WIDTH = 5; + +function renderBrailleBar(theme: any, value: number, width = BAR_WIDTH): string { + const v = clampPercent(value); + const levels = BRAILLE_GRADIENT.length - 1; + const totalSteps = width * levels; + const filledSteps = Math.round((v / 100) * totalSteps); + const full = Math.floor(filledSteps / levels); + const partial = filledSteps % levels; + const empty = width - full - (partial ? 1 : 0); + const color = colorForPercent(v); + const filled = BRAILLE_GRADIENT[BRAILLE_GRADIENT.length - 1]!.repeat(Math.max(0, full)); + const partialChar = partial ? BRAILLE_GRADIENT[partial]! : ""; + const emptyChars = BRAILLE_EMPTY.repeat(Math.max(0, empty)); + return theme.fg(color, filled + partialChar) + theme.fg("dim", emptyChars); +} + +function renderBrailleBarWide(theme: any, value: number): string { + return renderBrailleBar(theme, value, 12); +} + +const PROVIDER_LABELS: Record = { + codex: "Codex", + claude: "Claude", + zai: "Z.AI", + gemini: "Gemini", + antigravity: "Antigravity", +}; + +// --------------------------------------------------------------------------- +// /usage command popup +// --------------------------------------------------------------------------- +interface SubscriptionItem { + name: string; + provider: ProviderKey; + data: UsageData | null; + isActive: boolean; +} + +class UsageSelectorComponent extends Container implements Focusable { + private searchInput: Input; + private listContainer: Container; + private hintText: Text; + private tui: any; + private theme: any; + private onCancelCallback: () => void; + private allItems: SubscriptionItem[] = []; + private filteredItems: SubscriptionItem[] = []; + private selectedIndex = 0; + private loading = true; + private activeProvider: ProviderKey | null; + private fetchAllFn: () => Promise; + private _focused = false; + + get focused(): boolean { return this._focused; } + set focused(value: boolean) { this._focused = value; this.searchInput.focused = value; } + + constructor( + tui: any, + theme: any, + activeProvider: ProviderKey | null, + fetchAll: () => Promise, + onCancel: () => void, + ) { + super(); + this.tui = tui; + this.theme = theme; + this.activeProvider = activeProvider; + this.fetchAllFn = fetchAll; + this.onCancelCallback = onCancel; + + this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s))); + this.addChild(new Spacer(1)); + this.hintText = new Text(theme.fg("dim", "Fetching usage from all providers…"), 0, 0); + this.addChild(this.hintText); + this.addChild(new Spacer(1)); + this.searchInput = new Input(); + this.addChild(this.searchInput); + this.addChild(new Spacer(1)); + this.listContainer = new Container(); + this.addChild(this.listContainer); + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s))); + + this.fetchAllFn() + .then((results) => { + this.loading = false; + this.buildItems(results); + this.updateList(); + this.hintText.setText( + theme.fg("muted", "Only showing providers with credentials. ") + + theme.fg("dim", "✓ = active provider"), + ); + this.tui.requestRender(); + }) + .catch(() => { + this.loading = false; + this.hintText.setText(theme.fg("error", "Failed to fetch usage data")); + this.tui.requestRender(); + }); + + this.updateList(); + } + + private buildItems(results: UsageByProvider) { + const providers: Array<{ key: ProviderKey; name: string }> = [ + { key: "codex", name: "Codex" }, + { key: "claude", name: "Claude" }, + { key: "zai", name: "Z.AI" }, + { key: "gemini", name: "Gemini" }, + { key: "antigravity", name: "Antigravity" }, + ]; + this.allItems = []; + for (const p of providers) { + if (results[p.key] !== null) { + this.allItems.push({ + name: p.name, + provider: p.key, + data: results[p.key], + isActive: this.activeProvider === p.key, + }); + } + } + this.filteredItems = this.allItems; + this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1)); + } + + private filterItems(query: string) { + if (!query) { + this.filteredItems = this.allItems; + } else { + const q = query.toLowerCase(); + this.filteredItems = this.allItems.filter( + (item) => item.name.toLowerCase().includes(q) || item.provider.toLowerCase().includes(q), + ); + } + this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1)); + } + + private renderItem(item: SubscriptionItem, isSelected: boolean) { + const t = this.theme; + const pointer = isSelected ? t.fg("accent", "→ ") : " "; + const activeBadge = item.isActive ? t.fg("success", " ✓") : ""; + const name = isSelected ? t.fg("accent", t.bold(item.name)) : item.name; + this.listContainer.addChild(new Text(`${pointer}${name}${activeBadge}`, 0, 0)); + const indent = " "; + + if (!item.data) { + this.listContainer.addChild(new Text(indent + t.fg("dim", "No credentials"), 0, 0)); + } else if (item.data.error) { + this.listContainer.addChild(new Text(indent + t.fg("error", item.data.error), 0, 0)); + } else { + const session = clampPercent(item.data.session); + const weekly = clampPercent(item.data.weekly); + const sessionReset = item.data.sessionResetsIn + ? t.fg("dim", ` resets in ${item.data.sessionResetsIn}`) : ""; + const weeklyReset = item.data.weeklyResetsIn + ? t.fg("dim", ` resets in ${item.data.weeklyResetsIn}`) : ""; + + this.listContainer.addChild(new Text( + indent + t.fg("muted", "Session ") + + renderBrailleBarWide(t, session) + " " + + t.fg(colorForPercent(session), `${session}%`.padStart(4)) + sessionReset, + 0, 0, + )); + this.listContainer.addChild(new Text( + indent + t.fg("muted", "Weekly ") + + renderBrailleBarWide(t, weekly) + " " + + t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) + weeklyReset, + 0, 0, + )); + + if (typeof item.data.extraSpend === "number" && typeof item.data.extraLimit === "number") { + this.listContainer.addChild(new Text( + indent + t.fg("muted", "Extra ") + + t.fg("dim", `$${item.data.extraSpend.toFixed(2)} / $${item.data.extraLimit}`), + 0, 0, + )); + } + } + this.listContainer.addChild(new Spacer(1)); + } + + private updateList() { + this.listContainer.clear(); + if (this.loading) { + this.listContainer.addChild(new Text(this.theme.fg("muted", " Loading…"), 0, 0)); + return; + } + if (this.filteredItems.length === 0) { + this.listContainer.addChild(new Text(this.theme.fg("muted", " No matching providers"), 0, 0)); + return; + } + for (let i = 0; i < this.filteredItems.length; i++) { + this.renderItem(this.filteredItems[i]!, i === this.selectedIndex); + } + } + + handleInput(keyData: string): void { + const kb = getEditorKeybindings(); + if (kb.matches(keyData, "selectUp")) { + if (this.filteredItems.length === 0) return; + this.selectedIndex = this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1; + this.updateList(); return; + } + if (kb.matches(keyData, "selectDown")) { + if (this.filteredItems.length === 0) return; + this.selectedIndex = this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1; + this.updateList(); return; + } + if (kb.matches(keyData, "selectCancel") || kb.matches(keyData, "selectConfirm")) { + this.onCancelCallback(); return; + } + this.searchInput.handleInput(keyData); + this.filterItems(this.searchInput.getValue()); + this.updateList(); + } +} + +// --------------------------------------------------------------------------- +// Extension state +// --------------------------------------------------------------------------- +interface UsageState extends UsageByProvider { + lastPoll: number; + activeProvider: ProviderKey | null; +} + +interface PollOptions { + cacheTtl?: number; + forceFresh?: boolean; +} + +export default function (pi: ExtensionAPI) { + const endpoints = resolveUsageEndpoints(); + const state: UsageState = { + codex: null, claude: null, zai: null, gemini: null, antigravity: null, + lastPoll: 0, activeProvider: null, + }; + + let pollInFlight: Promise | null = null; + let pollQueued = false; + let pollStartedAt = 0; + let streamingTimer: ReturnType | null = null; + let ctx: any = null; + + // --------------------------------------------------------------------------- + // Status update + // --------------------------------------------------------------------------- + function updateStatus() { + const active = state.activeProvider; + const data = active ? state[active] : null; + + // Always emit event for other extensions (e.g. footer-display) + if (data && !data.error) { + pi.events.emit("usage:update", { + session: data.session, + weekly: data.weekly, + sessionResetsIn: data.sessionResetsIn, + sessionResetsAt: data.sessionResetsAt, + weeklyResetsIn: data.weeklyResetsIn, + }); + } + + if (!ctx?.hasUI) return; + + const theme = ctx.ui.theme; + + if (!active) { + ctx.ui.setStatus(STATUS_KEY, undefined); + return; + } + + const auth = readAuth(); + if (!canShowForProvider(active, auth, endpoints)) { + ctx.ui.setStatus(STATUS_KEY, undefined); + return; + } + + if (!data) { + ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", "loading\u2026")); + return; + } + + if (data.error) { + const cache = readUsageCache(); + const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0; + const note = blockedUntil > Date.now() + ? ` \u2014 retry in ${Math.ceil((blockedUntil - Date.now()) / 60000)}m` : ""; + ctx.ui.setStatus(STATUS_KEY, theme.fg("warning", `${PROVIDER_LABELS[active]} unavailable${note}`)); + return; + } + + const session = clampPercent(data.session); + const weekly = clampPercent(data.weekly); + + let s = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`); + if (data.sessionResetsIn) s += " " + theme.fg("dim", data.sessionResetsIn); + + let w = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`); + if (data.weeklyResetsIn) w += " " + theme.fg("dim", `\u27F3 ${data.weeklyResetsIn}`); + + ctx.ui.setStatus(STATUS_KEY, s + theme.fg("dim", " | ") + w); + } + + function updateProviderFrom(modelLike: any): boolean { + const previous = state.activeProvider; + state.activeProvider = detectProvider(modelLike); + if (previous !== state.activeProvider) { updateStatus(); return true; } + return false; + } + + // --------------------------------------------------------------------------- + // Polling + // --------------------------------------------------------------------------- + async function runPollInner(options: PollOptions = {}) { + const auth = readAuth(); + const active = state.activeProvider; + + if (!canShowForProvider(active, auth, endpoints) || !auth || !active) { + state.lastPoll = Date.now(); updateStatus(); return; + } + + const cache = readUsageCache(); + const now = Date.now(); + const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS; + + const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0; + if (now < blockedUntil) { + if (cache?.data?.[active]) { + state[active] = cache.data[active]!; + } else { + // Rate-limited but no cached data — show a meaningful status instead + // of leaving state null (which shows eternal "loading…"). + const retryMin = Math.ceil((blockedUntil - now) / 60000); + state[active] = { session: 0, weekly: 0, error: `rate limited (retry in ${retryMin}m)` }; + } + state.lastPoll = now; updateStatus(); return; + } + + if (!options.forceFresh && cache && now - cache.timestamp < cacheTtl && cache.data?.[active]) { + state[active] = cache.data[active]!; + state.lastPoll = now; updateStatus(); return; + } + + const oauthId = providerToOAuthProviderId(active); + let effectiveAuth = auth; + if (oauthId && active !== "zai") { + const creds = auth[oauthId as keyof typeof auth] as + | { access?: string; refresh?: string; expires?: number } | undefined; + const expires = typeof creds?.expires === "number" ? creds.expires : 0; + const tokenExpiredOrMissing = !creds?.access || (expires > 0 && Date.now() + 60_000 >= expires); + if (tokenExpiredOrMissing && creds?.refresh) { + try { + const refreshPromise = ensureFreshAuthForProviders([oauthId as OAuthProviderId], { auth, persist: true }); + const timeoutPromise = new Promise((_, reject) => + setTimeout(() => reject(new Error("OAuth refresh timeout")), 15_000), + ); + const refreshed = await Promise.race([refreshPromise, timeoutPromise]); + if (refreshed.auth) effectiveAuth = refreshed.auth; + } catch {} + } + } + + let result: UsageData; + if (active === "codex") { + const access = effectiveAuth["openai-codex"]?.access; + result = access ? await fetchCodexUsage(access) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } else if (active === "claude") { + const access = effectiveAuth.anthropic?.access; + result = access ? await fetchClaudeUsage(access) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } else if (active === "zai") { + const token = effectiveAuth.zai?.access || effectiveAuth.zai?.key; + result = token ? await fetchZaiUsage(token, { endpoints }) + : { session: 0, weekly: 0, error: "missing token (try /login again)" }; + } else if (active === "gemini") { + const creds = effectiveAuth["google-gemini-cli"]; + result = creds?.access + ? await fetchGoogleUsage(creds.access, endpoints.gemini, creds.projectId, "gemini", { endpoints }) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } else { + const creds = effectiveAuth["google-antigravity"]; + result = creds?.access + ? await fetchGoogleUsage(creds.access, endpoints.antigravity, creds.projectId, "antigravity", { endpoints }) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } + + state[active] = result; + + if (result.error) { + if (result.error === "HTTP 429") { + const nextCache: import("./core").UsageCache = { + timestamp: cache?.timestamp ?? now, + data: { ...(cache?.data ?? {}) }, + rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}), [active]: now + RATE_LIMITED_BACKOFF_MS }, + }; + writeUsageCache(nextCache); + } + } else { + const nextCache: import("./core").UsageCache = { + timestamp: now, + data: { ...(cache?.data ?? {}), [active]: result }, + rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) }, + }; + delete nextCache.rateLimitedUntil![active]; + writeUsageCache(nextCache); + } + + state.lastPoll = now; + updateStatus(); + } + + async function runPoll(options: PollOptions = {}): Promise { + const timeout = new Promise((_, reject) => + setTimeout(() => reject(new Error("runPoll timeout")), 25_000), + ); + await Promise.race([runPollInner(options), timeout]); + } + + const POLL_TIMEOUT_MS = 30_000; + + async function poll(options: PollOptions = {}) { + // If a previous poll has been running longer than POLL_TIMEOUT_MS, abandon it + // so we don't queue forever behind a stuck request. + if (pollInFlight && pollStartedAt > 0 && Date.now() - pollStartedAt > POLL_TIMEOUT_MS) { + pollInFlight = null; + pollQueued = false; + const active = state.activeProvider; + if (active && !state[active]) { + state[active] = { session: 0, weekly: 0, error: "poll timeout" }; + updateStatus(); + } + } + + if (pollInFlight) { pollQueued = true; await pollInFlight; return; } + do { + pollQueued = false; + pollStartedAt = Date.now(); + pollInFlight = runPoll(options).catch(() => { + // If runPoll threw, ensure we don't leave status stuck at "loading…" + const active = state.activeProvider; + if (active && !state[active]) { + state[active] = { session: 0, weekly: 0, error: "poll failed" }; + updateStatus(); + } + }).finally(() => { pollInFlight = null; pollStartedAt = 0; }); + await pollInFlight; + } while (pollQueued); + } + + function startStreamingTimer() { + if (streamingTimer !== null) return; + streamingTimer = setInterval(() => { void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); }, STREAMING_POLL_INTERVAL_MS); + } + + function stopStreamingTimer() { + if (streamingTimer !== null) { clearInterval(streamingTimer); streamingTimer = null; } + } + + // ── Lifecycle ──────────────────────────────────────────────────────────── + + pi.on("session_start", async (_event, _ctx) => { + ctx = _ctx; + updateProviderFrom(_ctx.model); + await poll(); + }); + + pi.on("session_shutdown", async (_event, _ctx) => { + stopStreamingTimer(); + if (_ctx?.hasUI) _ctx.ui.setStatus(STATUS_KEY, undefined); + }); + + pi.on("model_select", async (event, _ctx) => { + ctx = _ctx; + const changed = updateProviderFrom(event.model ?? _ctx.model); + if (changed) await poll(); + }); + + pi.on("turn_start", (_event, _ctx) => { ctx = _ctx; updateProviderFrom(_ctx.model); }); + + pi.on("before_agent_start", async (_event, _ctx) => { + ctx = _ctx; + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + }); + + pi.on("agent_start", (_event, _ctx) => { ctx = _ctx; startStreamingTimer(); }); + + pi.on("agent_end", async (_event, _ctx) => { + ctx = _ctx; + stopStreamingTimer(); + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + }); + + pi.events.on("claude-account:switched", () => { + const cache = readUsageCache(); + if (cache?.data?.claude) { + const nextCache: import("./core").UsageCache = { ...cache, data: { ...cache.data } }; + delete nextCache.data.claude; + writeUsageCache(nextCache); + } + void poll({ forceFresh: true }); + }); + + // ── /usage command ─────────────────────────────────────────────────────── + + pi.registerCommand("usage", { + description: "Show API usage for all subscriptions", + handler: async (_args, _ctx) => { + ctx = _ctx; + updateProviderFrom(_ctx.model); + try { + if (_ctx?.hasUI) { + await _ctx.ui.custom((tui, theme, _keybindings, done) => { + return new UsageSelectorComponent( + tui, theme, state.activeProvider, + () => fetchAllUsages({ endpoints }), + () => done(), + ); + }); + } + } finally { + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + } + }, + }); +} diff --git a/pi/.pi/agent/extensions/usage-bars/index.ts.before b/pi/.pi/agent/extensions/usage-bars/index.ts.before new file mode 100644 index 0000000..7e7b8f0 --- /dev/null +++ b/pi/.pi/agent/extensions/usage-bars/index.ts.before @@ -0,0 +1,581 @@ +/** + * Usage Extension - Minimal API usage indicator for pi + * + * Polls Codex, Anthropic, Z.AI, Gemini CLI / Antigravity usage and exposes it + * via two channels: + * • pi.events "usage:update" — for other extensions (e.g. footer-display) + * • ctx.ui.setStatus("usage-bars", …) — formatted S/W braille bars + * + * Rendering / footer layout is handled by the separate footer-display extension. + */ + +import { DynamicBorder, type ExtensionAPI } from "@mariozechner/pi-coding-agent"; +import { + Container, + Input, + Spacer, + Text, + getEditorKeybindings, + type Focusable, +} from "@mariozechner/pi-tui"; +import { + canShowForProvider, + clampPercent, + colorForPercent, + detectProvider, + ensureFreshAuthForProviders, + fetchAllUsages, + fetchClaudeUsage, + fetchCodexUsage, + fetchGoogleUsage, + fetchZaiUsage, + providerToOAuthProviderId, + readAuth, + readUsageCache, + resolveUsageEndpoints, + writeUsageCache, + type OAuthProviderId, + type ProviderKey, + type UsageByProvider, + type UsageData, +} from "./core"; + +const CACHE_TTL_MS = 15 * 60 * 1000; +const ACTIVE_CACHE_TTL_MS = 3 * 60 * 1000; +const STREAMING_POLL_INTERVAL_MS = 2 * 60 * 1000; +const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000; + +const STATUS_KEY = "usage-bars"; + +// --------------------------------------------------------------------------- +// Braille gradient bar (⣀ ⣄ ⣤ ⣦ ⣶ ⣷ ⣿) +// --------------------------------------------------------------------------- +const BRAILLE_GRADIENT = "\u28C0\u28C4\u28E4\u28E6\u28F6\u28F7\u28FF"; +const BRAILLE_EMPTY = "\u28C0"; +const BAR_WIDTH = 5; + +function renderBrailleBar(theme: any, value: number, width = BAR_WIDTH): string { + const v = clampPercent(value); + const levels = BRAILLE_GRADIENT.length - 1; + const totalSteps = width * levels; + const filledSteps = Math.round((v / 100) * totalSteps); + const full = Math.floor(filledSteps / levels); + const partial = filledSteps % levels; + const empty = width - full - (partial ? 1 : 0); + const color = colorForPercent(v); + const filled = BRAILLE_GRADIENT[BRAILLE_GRADIENT.length - 1]!.repeat(Math.max(0, full)); + const partialChar = partial ? BRAILLE_GRADIENT[partial]! : ""; + const emptyChars = BRAILLE_EMPTY.repeat(Math.max(0, empty)); + return theme.fg(color, filled + partialChar) + theme.fg("dim", emptyChars); +} + +function renderBrailleBarWide(theme: any, value: number): string { + return renderBrailleBar(theme, value, 12); +} + +const PROVIDER_LABELS: Record = { + codex: "Codex", + claude: "Claude", + zai: "Z.AI", + gemini: "Gemini", + antigravity: "Antigravity", +}; + +// --------------------------------------------------------------------------- +// /usage command popup +// --------------------------------------------------------------------------- +interface SubscriptionItem { + name: string; + provider: ProviderKey; + data: UsageData | null; + isActive: boolean; +} + +class UsageSelectorComponent extends Container implements Focusable { + private searchInput: Input; + private listContainer: Container; + private hintText: Text; + private tui: any; + private theme: any; + private onCancelCallback: () => void; + private allItems: SubscriptionItem[] = []; + private filteredItems: SubscriptionItem[] = []; + private selectedIndex = 0; + private loading = true; + private activeProvider: ProviderKey | null; + private fetchAllFn: () => Promise; + private _focused = false; + + get focused(): boolean { return this._focused; } + set focused(value: boolean) { this._focused = value; this.searchInput.focused = value; } + + constructor( + tui: any, + theme: any, + activeProvider: ProviderKey | null, + fetchAll: () => Promise, + onCancel: () => void, + ) { + super(); + this.tui = tui; + this.theme = theme; + this.activeProvider = activeProvider; + this.fetchAllFn = fetchAll; + this.onCancelCallback = onCancel; + + this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s))); + this.addChild(new Spacer(1)); + this.hintText = new Text(theme.fg("dim", "Fetching usage from all providers…"), 0, 0); + this.addChild(this.hintText); + this.addChild(new Spacer(1)); + this.searchInput = new Input(); + this.addChild(this.searchInput); + this.addChild(new Spacer(1)); + this.listContainer = new Container(); + this.addChild(this.listContainer); + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s))); + + this.fetchAllFn() + .then((results) => { + this.loading = false; + this.buildItems(results); + this.updateList(); + this.hintText.setText( + theme.fg("muted", "Only showing providers with credentials. ") + + theme.fg("dim", "✓ = active provider"), + ); + this.tui.requestRender(); + }) + .catch(() => { + this.loading = false; + this.hintText.setText(theme.fg("error", "Failed to fetch usage data")); + this.tui.requestRender(); + }); + + this.updateList(); + } + + private buildItems(results: UsageByProvider) { + const providers: Array<{ key: ProviderKey; name: string }> = [ + { key: "codex", name: "Codex" }, + { key: "claude", name: "Claude" }, + { key: "zai", name: "Z.AI" }, + { key: "gemini", name: "Gemini" }, + { key: "antigravity", name: "Antigravity" }, + ]; + this.allItems = []; + for (const p of providers) { + if (results[p.key] !== null) { + this.allItems.push({ + name: p.name, + provider: p.key, + data: results[p.key], + isActive: this.activeProvider === p.key, + }); + } + } + this.filteredItems = this.allItems; + this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1)); + } + + private filterItems(query: string) { + if (!query) { + this.filteredItems = this.allItems; + } else { + const q = query.toLowerCase(); + this.filteredItems = this.allItems.filter( + (item) => item.name.toLowerCase().includes(q) || item.provider.toLowerCase().includes(q), + ); + } + this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1)); + } + + private renderItem(item: SubscriptionItem, isSelected: boolean) { + const t = this.theme; + const pointer = isSelected ? t.fg("accent", "→ ") : " "; + const activeBadge = item.isActive ? t.fg("success", " ✓") : ""; + const name = isSelected ? t.fg("accent", t.bold(item.name)) : item.name; + this.listContainer.addChild(new Text(`${pointer}${name}${activeBadge}`, 0, 0)); + const indent = " "; + + if (!item.data) { + this.listContainer.addChild(new Text(indent + t.fg("dim", "No credentials"), 0, 0)); + } else if (item.data.error) { + this.listContainer.addChild(new Text(indent + t.fg("error", item.data.error), 0, 0)); + } else { + const session = clampPercent(item.data.session); + const weekly = clampPercent(item.data.weekly); + const sessionReset = item.data.sessionResetsIn + ? t.fg("dim", ` resets in ${item.data.sessionResetsIn}`) : ""; + const weeklyReset = item.data.weeklyResetsIn + ? t.fg("dim", ` resets in ${item.data.weeklyResetsIn}`) : ""; + + this.listContainer.addChild(new Text( + indent + t.fg("muted", "Session ") + + renderBrailleBarWide(t, session) + " " + + t.fg(colorForPercent(session), `${session}%`.padStart(4)) + sessionReset, + 0, 0, + )); + this.listContainer.addChild(new Text( + indent + t.fg("muted", "Weekly ") + + renderBrailleBarWide(t, weekly) + " " + + t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) + weeklyReset, + 0, 0, + )); + + if (typeof item.data.extraSpend === "number" && typeof item.data.extraLimit === "number") { + this.listContainer.addChild(new Text( + indent + t.fg("muted", "Extra ") + + t.fg("dim", `$${item.data.extraSpend.toFixed(2)} / $${item.data.extraLimit}`), + 0, 0, + )); + } + } + this.listContainer.addChild(new Spacer(1)); + } + + private updateList() { + this.listContainer.clear(); + if (this.loading) { + this.listContainer.addChild(new Text(this.theme.fg("muted", " Loading…"), 0, 0)); + return; + } + if (this.filteredItems.length === 0) { + this.listContainer.addChild(new Text(this.theme.fg("muted", " No matching providers"), 0, 0)); + return; + } + for (let i = 0; i < this.filteredItems.length; i++) { + this.renderItem(this.filteredItems[i]!, i === this.selectedIndex); + } + } + + handleInput(keyData: string): void { + const kb = getEditorKeybindings(); + if (kb.matches(keyData, "selectUp")) { + if (this.filteredItems.length === 0) return; + this.selectedIndex = this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1; + this.updateList(); return; + } + if (kb.matches(keyData, "selectDown")) { + if (this.filteredItems.length === 0) return; + this.selectedIndex = this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1; + this.updateList(); return; + } + if (kb.matches(keyData, "selectCancel") || kb.matches(keyData, "selectConfirm")) { + this.onCancelCallback(); return; + } + this.searchInput.handleInput(keyData); + this.filterItems(this.searchInput.getValue()); + this.updateList(); + } +} + +// --------------------------------------------------------------------------- +// Extension state +// --------------------------------------------------------------------------- +interface UsageState extends UsageByProvider { + lastPoll: number; + activeProvider: ProviderKey | null; +} + +interface PollOptions { + cacheTtl?: number; + forceFresh?: boolean; +} + +export default function (pi: ExtensionAPI) { + const endpoints = resolveUsageEndpoints(); + const state: UsageState = { + codex: null, claude: null, zai: null, gemini: null, antigravity: null, + lastPoll: 0, activeProvider: null, + }; + + let pollInFlight: Promise | null = null; + let pollQueued = false; + let pollStartedAt = 0; + let streamingTimer: ReturnType | null = null; + let ctx: any = null; + + // --------------------------------------------------------------------------- + // Status update + // --------------------------------------------------------------------------- + function updateStatus() { + const active = state.activeProvider; + const data = active ? state[active] : null; + + // Always emit event for other extensions (e.g. footer-display) + if (data && !data.error) { + pi.events.emit("usage:update", { + session: data.session, + weekly: data.weekly, + sessionResetsIn: data.sessionResetsIn, + sessionResetsAt: data.sessionResetsAt, + weeklyResetsIn: data.weeklyResetsIn, + }); + } + + if (!ctx?.hasUI) return; + + const theme = ctx.ui.theme; + + if (!active) { + ctx.ui.setStatus(STATUS_KEY, undefined); + return; + } + + const auth = readAuth(); + if (!canShowForProvider(active, auth, endpoints)) { + ctx.ui.setStatus(STATUS_KEY, undefined); + return; + } + + if (!data) { + ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", "loading\u2026")); + return; + } + + if (data.error) { + const cache = readUsageCache(); + const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0; + const note = blockedUntil > Date.now() + ? ` \u2014 retry in ${Math.ceil((blockedUntil - Date.now()) / 60000)}m` : ""; + ctx.ui.setStatus(STATUS_KEY, theme.fg("warning", `${PROVIDER_LABELS[active]} unavailable${note}`)); + return; + } + + const session = clampPercent(data.session); + const weekly = clampPercent(data.weekly); + + let s = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`); + if (data.sessionResetsIn) s += " " + theme.fg("dim", data.sessionResetsIn); + + let w = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`); + if (data.weeklyResetsIn) w += " " + theme.fg("dim", `\u27F3 ${data.weeklyResetsIn}`); + + ctx.ui.setStatus(STATUS_KEY, s + theme.fg("dim", " | ") + w); + } + + function updateProviderFrom(modelLike: any): boolean { + const previous = state.activeProvider; + state.activeProvider = detectProvider(modelLike); + if (previous !== state.activeProvider) { updateStatus(); return true; } + return false; + } + + // --------------------------------------------------------------------------- + // Polling + // --------------------------------------------------------------------------- + async function runPollInner(options: PollOptions = {}) { + const auth = readAuth(); + const active = state.activeProvider; + + if (!canShowForProvider(active, auth, endpoints) || !auth || !active) { + state.lastPoll = Date.now(); updateStatus(); return; + } + + const cache = readUsageCache(); + const now = Date.now(); + const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS; + + const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0; + if (now < blockedUntil) { + if (cache?.data?.[active]) { + state[active] = cache.data[active]!; + } else { + // Rate-limited but no cached data — show a meaningful status instead + // of leaving state null (which shows eternal "loading…"). + const retryMin = Math.ceil((blockedUntil - now) / 60000); + state[active] = { session: 0, weekly: 0, error: `rate limited (retry in ${retryMin}m)` }; + } + state.lastPoll = now; updateStatus(); return; + } + + if (!options.forceFresh && cache && now - cache.timestamp < cacheTtl && cache.data?.[active]) { + state[active] = cache.data[active]!; + state.lastPoll = now; updateStatus(); return; + } + + const oauthId = providerToOAuthProviderId(active); + let effectiveAuth = auth; + if (oauthId && active !== "zai") { + const creds = auth[oauthId as keyof typeof auth] as + | { access?: string; refresh?: string; expires?: number } | undefined; + const expires = typeof creds?.expires === "number" ? creds.expires : 0; + const tokenExpiredOrMissing = !creds?.access || (expires > 0 && Date.now() + 60_000 >= expires); + if (tokenExpiredOrMissing && creds?.refresh) { + try { + const refreshPromise = ensureFreshAuthForProviders([oauthId as OAuthProviderId], { auth, persist: true }); + const timeoutPromise = new Promise((_, reject) => + setTimeout(() => reject(new Error("OAuth refresh timeout")), 15_000), + ); + const refreshed = await Promise.race([refreshPromise, timeoutPromise]); + if (refreshed.auth) effectiveAuth = refreshed.auth; + } catch {} + } + } + + let result: UsageData; + if (active === "codex") { + const access = effectiveAuth["openai-codex"]?.access; + result = access ? await fetchCodexUsage(access) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } else if (active === "claude") { + const access = effectiveAuth.anthropic?.access; + result = access ? await fetchClaudeUsage(access) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } else if (active === "zai") { + const token = effectiveAuth.zai?.access || effectiveAuth.zai?.key; + result = token ? await fetchZaiUsage(token, { endpoints }) + : { session: 0, weekly: 0, error: "missing token (try /login again)" }; + } else if (active === "gemini") { + const creds = effectiveAuth["google-gemini-cli"]; + result = creds?.access + ? await fetchGoogleUsage(creds.access, endpoints.gemini, creds.projectId, "gemini", { endpoints }) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } else { + const creds = effectiveAuth["google-antigravity"]; + result = creds?.access + ? await fetchGoogleUsage(creds.access, endpoints.antigravity, creds.projectId, "antigravity", { endpoints }) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } + + state[active] = result; + + if (result.error) { + if (result.error === "HTTP 429") { + const nextCache: import("./core").UsageCache = { + timestamp: cache?.timestamp ?? now, + data: { ...(cache?.data ?? {}) }, + rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}), [active]: now + RATE_LIMITED_BACKOFF_MS }, + }; + writeUsageCache(nextCache); + } + } else { + const nextCache: import("./core").UsageCache = { + timestamp: now, + data: { ...(cache?.data ?? {}), [active]: result }, + rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) }, + }; + delete nextCache.rateLimitedUntil![active]; + writeUsageCache(nextCache); + } + + state.lastPoll = now; + updateStatus(); + } + + async function runPoll(options: PollOptions = {}): Promise { + const timeout = new Promise((_, reject) => + setTimeout(() => reject(new Error("runPoll timeout")), 25_000), + ); + await Promise.race([runPollInner(options), timeout]); + } + + const POLL_TIMEOUT_MS = 30_000; + + async function poll(options: PollOptions = {}) { + // If a previous poll has been running longer than POLL_TIMEOUT_MS, abandon it + // so we don't queue forever behind a stuck request. + if (pollInFlight && pollStartedAt > 0 && Date.now() - pollStartedAt > POLL_TIMEOUT_MS) { + pollInFlight = null; + pollQueued = false; + const active = state.activeProvider; + if (active && !state[active]) { + state[active] = { session: 0, weekly: 0, error: "poll timeout" }; + updateStatus(); + } + } + + if (pollInFlight) { pollQueued = true; await pollInFlight; return; } + do { + pollQueued = false; + pollStartedAt = Date.now(); + pollInFlight = runPoll(options).catch(() => { + // If runPoll threw, ensure we don't leave status stuck at "loading…" + const active = state.activeProvider; + if (active && !state[active]) { + state[active] = { session: 0, weekly: 0, error: "poll failed" }; + updateStatus(); + } + }).finally(() => { pollInFlight = null; pollStartedAt = 0; }); + await pollInFlight; + } while (pollQueued); + } + + function startStreamingTimer() { + if (streamingTimer !== null) return; + streamingTimer = setInterval(() => { void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); }, STREAMING_POLL_INTERVAL_MS); + } + + function stopStreamingTimer() { + if (streamingTimer !== null) { clearInterval(streamingTimer); streamingTimer = null; } + } + + // ── Lifecycle ──────────────────────────────────────────────────────────── + + pi.on("session_start", async (_event, _ctx) => { + ctx = _ctx; + updateProviderFrom(_ctx.model); + await poll(); + }); + + pi.on("session_shutdown", async (_event, _ctx) => { + stopStreamingTimer(); + if (_ctx?.hasUI) _ctx.ui.setStatus(STATUS_KEY, undefined); + }); + + pi.on("model_select", async (event, _ctx) => { + ctx = _ctx; + const changed = updateProviderFrom(event.model ?? _ctx.model); + if (changed) await poll(); + }); + + pi.on("turn_start", (_event, _ctx) => { ctx = _ctx; updateProviderFrom(_ctx.model); }); + + pi.on("before_agent_start", async (_event, _ctx) => { + ctx = _ctx; + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + }); + + pi.on("agent_start", (_event, _ctx) => { ctx = _ctx; startStreamingTimer(); }); + + pi.on("agent_end", async (_event, _ctx) => { + ctx = _ctx; + stopStreamingTimer(); + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + }); + + pi.events.on("claude-account:switched", () => { + const cache = readUsageCache(); + if (cache?.data?.claude) { + const nextCache: import("./core").UsageCache = { ...cache, data: { ...cache.data } }; + delete nextCache.data.claude; + writeUsageCache(nextCache); + } + void poll({ forceFresh: true }); + }); + + // ── /usage command ─────────────────────────────────────────────────────── + + pi.registerCommand("usage", { + description: "Show API usage for all subscriptions", + handler: async (_args, _ctx) => { + ctx = _ctx; + updateProviderFrom(_ctx.model); + try { + if (_ctx?.hasUI) { + await _ctx.ui.custom((tui, theme, _keybindings, done) => { + return new UsageSelectorComponent( + tui, theme, state.activeProvider, + () => fetchAllUsages({ endpoints }), + () => done(), + ); + }); + } + } finally { + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + } + }, + }); +} diff --git a/pi/.pi/agent/extensions/usage-bars/index.ts.editable b/pi/.pi/agent/extensions/usage-bars/index.ts.editable new file mode 100644 index 0000000..7e7b8f0 --- /dev/null +++ b/pi/.pi/agent/extensions/usage-bars/index.ts.editable @@ -0,0 +1,581 @@ +/** + * Usage Extension - Minimal API usage indicator for pi + * + * Polls Codex, Anthropic, Z.AI, Gemini CLI / Antigravity usage and exposes it + * via two channels: + * • pi.events "usage:update" — for other extensions (e.g. footer-display) + * • ctx.ui.setStatus("usage-bars", …) — formatted S/W braille bars + * + * Rendering / footer layout is handled by the separate footer-display extension. + */ + +import { DynamicBorder, type ExtensionAPI } from "@mariozechner/pi-coding-agent"; +import { + Container, + Input, + Spacer, + Text, + getEditorKeybindings, + type Focusable, +} from "@mariozechner/pi-tui"; +import { + canShowForProvider, + clampPercent, + colorForPercent, + detectProvider, + ensureFreshAuthForProviders, + fetchAllUsages, + fetchClaudeUsage, + fetchCodexUsage, + fetchGoogleUsage, + fetchZaiUsage, + providerToOAuthProviderId, + readAuth, + readUsageCache, + resolveUsageEndpoints, + writeUsageCache, + type OAuthProviderId, + type ProviderKey, + type UsageByProvider, + type UsageData, +} from "./core"; + +const CACHE_TTL_MS = 15 * 60 * 1000; +const ACTIVE_CACHE_TTL_MS = 3 * 60 * 1000; +const STREAMING_POLL_INTERVAL_MS = 2 * 60 * 1000; +const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000; + +const STATUS_KEY = "usage-bars"; + +// --------------------------------------------------------------------------- +// Braille gradient bar (⣀ ⣄ ⣤ ⣦ ⣶ ⣷ ⣿) +// --------------------------------------------------------------------------- +const BRAILLE_GRADIENT = "\u28C0\u28C4\u28E4\u28E6\u28F6\u28F7\u28FF"; +const BRAILLE_EMPTY = "\u28C0"; +const BAR_WIDTH = 5; + +function renderBrailleBar(theme: any, value: number, width = BAR_WIDTH): string { + const v = clampPercent(value); + const levels = BRAILLE_GRADIENT.length - 1; + const totalSteps = width * levels; + const filledSteps = Math.round((v / 100) * totalSteps); + const full = Math.floor(filledSteps / levels); + const partial = filledSteps % levels; + const empty = width - full - (partial ? 1 : 0); + const color = colorForPercent(v); + const filled = BRAILLE_GRADIENT[BRAILLE_GRADIENT.length - 1]!.repeat(Math.max(0, full)); + const partialChar = partial ? BRAILLE_GRADIENT[partial]! : ""; + const emptyChars = BRAILLE_EMPTY.repeat(Math.max(0, empty)); + return theme.fg(color, filled + partialChar) + theme.fg("dim", emptyChars); +} + +function renderBrailleBarWide(theme: any, value: number): string { + return renderBrailleBar(theme, value, 12); +} + +const PROVIDER_LABELS: Record = { + codex: "Codex", + claude: "Claude", + zai: "Z.AI", + gemini: "Gemini", + antigravity: "Antigravity", +}; + +// --------------------------------------------------------------------------- +// /usage command popup +// --------------------------------------------------------------------------- +interface SubscriptionItem { + name: string; + provider: ProviderKey; + data: UsageData | null; + isActive: boolean; +} + +class UsageSelectorComponent extends Container implements Focusable { + private searchInput: Input; + private listContainer: Container; + private hintText: Text; + private tui: any; + private theme: any; + private onCancelCallback: () => void; + private allItems: SubscriptionItem[] = []; + private filteredItems: SubscriptionItem[] = []; + private selectedIndex = 0; + private loading = true; + private activeProvider: ProviderKey | null; + private fetchAllFn: () => Promise; + private _focused = false; + + get focused(): boolean { return this._focused; } + set focused(value: boolean) { this._focused = value; this.searchInput.focused = value; } + + constructor( + tui: any, + theme: any, + activeProvider: ProviderKey | null, + fetchAll: () => Promise, + onCancel: () => void, + ) { + super(); + this.tui = tui; + this.theme = theme; + this.activeProvider = activeProvider; + this.fetchAllFn = fetchAll; + this.onCancelCallback = onCancel; + + this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s))); + this.addChild(new Spacer(1)); + this.hintText = new Text(theme.fg("dim", "Fetching usage from all providers…"), 0, 0); + this.addChild(this.hintText); + this.addChild(new Spacer(1)); + this.searchInput = new Input(); + this.addChild(this.searchInput); + this.addChild(new Spacer(1)); + this.listContainer = new Container(); + this.addChild(this.listContainer); + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s))); + + this.fetchAllFn() + .then((results) => { + this.loading = false; + this.buildItems(results); + this.updateList(); + this.hintText.setText( + theme.fg("muted", "Only showing providers with credentials. ") + + theme.fg("dim", "✓ = active provider"), + ); + this.tui.requestRender(); + }) + .catch(() => { + this.loading = false; + this.hintText.setText(theme.fg("error", "Failed to fetch usage data")); + this.tui.requestRender(); + }); + + this.updateList(); + } + + private buildItems(results: UsageByProvider) { + const providers: Array<{ key: ProviderKey; name: string }> = [ + { key: "codex", name: "Codex" }, + { key: "claude", name: "Claude" }, + { key: "zai", name: "Z.AI" }, + { key: "gemini", name: "Gemini" }, + { key: "antigravity", name: "Antigravity" }, + ]; + this.allItems = []; + for (const p of providers) { + if (results[p.key] !== null) { + this.allItems.push({ + name: p.name, + provider: p.key, + data: results[p.key], + isActive: this.activeProvider === p.key, + }); + } + } + this.filteredItems = this.allItems; + this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1)); + } + + private filterItems(query: string) { + if (!query) { + this.filteredItems = this.allItems; + } else { + const q = query.toLowerCase(); + this.filteredItems = this.allItems.filter( + (item) => item.name.toLowerCase().includes(q) || item.provider.toLowerCase().includes(q), + ); + } + this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1)); + } + + private renderItem(item: SubscriptionItem, isSelected: boolean) { + const t = this.theme; + const pointer = isSelected ? t.fg("accent", "→ ") : " "; + const activeBadge = item.isActive ? t.fg("success", " ✓") : ""; + const name = isSelected ? t.fg("accent", t.bold(item.name)) : item.name; + this.listContainer.addChild(new Text(`${pointer}${name}${activeBadge}`, 0, 0)); + const indent = " "; + + if (!item.data) { + this.listContainer.addChild(new Text(indent + t.fg("dim", "No credentials"), 0, 0)); + } else if (item.data.error) { + this.listContainer.addChild(new Text(indent + t.fg("error", item.data.error), 0, 0)); + } else { + const session = clampPercent(item.data.session); + const weekly = clampPercent(item.data.weekly); + const sessionReset = item.data.sessionResetsIn + ? t.fg("dim", ` resets in ${item.data.sessionResetsIn}`) : ""; + const weeklyReset = item.data.weeklyResetsIn + ? t.fg("dim", ` resets in ${item.data.weeklyResetsIn}`) : ""; + + this.listContainer.addChild(new Text( + indent + t.fg("muted", "Session ") + + renderBrailleBarWide(t, session) + " " + + t.fg(colorForPercent(session), `${session}%`.padStart(4)) + sessionReset, + 0, 0, + )); + this.listContainer.addChild(new Text( + indent + t.fg("muted", "Weekly ") + + renderBrailleBarWide(t, weekly) + " " + + t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) + weeklyReset, + 0, 0, + )); + + if (typeof item.data.extraSpend === "number" && typeof item.data.extraLimit === "number") { + this.listContainer.addChild(new Text( + indent + t.fg("muted", "Extra ") + + t.fg("dim", `$${item.data.extraSpend.toFixed(2)} / $${item.data.extraLimit}`), + 0, 0, + )); + } + } + this.listContainer.addChild(new Spacer(1)); + } + + private updateList() { + this.listContainer.clear(); + if (this.loading) { + this.listContainer.addChild(new Text(this.theme.fg("muted", " Loading…"), 0, 0)); + return; + } + if (this.filteredItems.length === 0) { + this.listContainer.addChild(new Text(this.theme.fg("muted", " No matching providers"), 0, 0)); + return; + } + for (let i = 0; i < this.filteredItems.length; i++) { + this.renderItem(this.filteredItems[i]!, i === this.selectedIndex); + } + } + + handleInput(keyData: string): void { + const kb = getEditorKeybindings(); + if (kb.matches(keyData, "selectUp")) { + if (this.filteredItems.length === 0) return; + this.selectedIndex = this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1; + this.updateList(); return; + } + if (kb.matches(keyData, "selectDown")) { + if (this.filteredItems.length === 0) return; + this.selectedIndex = this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1; + this.updateList(); return; + } + if (kb.matches(keyData, "selectCancel") || kb.matches(keyData, "selectConfirm")) { + this.onCancelCallback(); return; + } + this.searchInput.handleInput(keyData); + this.filterItems(this.searchInput.getValue()); + this.updateList(); + } +} + +// --------------------------------------------------------------------------- +// Extension state +// --------------------------------------------------------------------------- +interface UsageState extends UsageByProvider { + lastPoll: number; + activeProvider: ProviderKey | null; +} + +interface PollOptions { + cacheTtl?: number; + forceFresh?: boolean; +} + +export default function (pi: ExtensionAPI) { + const endpoints = resolveUsageEndpoints(); + const state: UsageState = { + codex: null, claude: null, zai: null, gemini: null, antigravity: null, + lastPoll: 0, activeProvider: null, + }; + + let pollInFlight: Promise | null = null; + let pollQueued = false; + let pollStartedAt = 0; + let streamingTimer: ReturnType | null = null; + let ctx: any = null; + + // --------------------------------------------------------------------------- + // Status update + // --------------------------------------------------------------------------- + function updateStatus() { + const active = state.activeProvider; + const data = active ? state[active] : null; + + // Always emit event for other extensions (e.g. footer-display) + if (data && !data.error) { + pi.events.emit("usage:update", { + session: data.session, + weekly: data.weekly, + sessionResetsIn: data.sessionResetsIn, + sessionResetsAt: data.sessionResetsAt, + weeklyResetsIn: data.weeklyResetsIn, + }); + } + + if (!ctx?.hasUI) return; + + const theme = ctx.ui.theme; + + if (!active) { + ctx.ui.setStatus(STATUS_KEY, undefined); + return; + } + + const auth = readAuth(); + if (!canShowForProvider(active, auth, endpoints)) { + ctx.ui.setStatus(STATUS_KEY, undefined); + return; + } + + if (!data) { + ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", "loading\u2026")); + return; + } + + if (data.error) { + const cache = readUsageCache(); + const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0; + const note = blockedUntil > Date.now() + ? ` \u2014 retry in ${Math.ceil((blockedUntil - Date.now()) / 60000)}m` : ""; + ctx.ui.setStatus(STATUS_KEY, theme.fg("warning", `${PROVIDER_LABELS[active]} unavailable${note}`)); + return; + } + + const session = clampPercent(data.session); + const weekly = clampPercent(data.weekly); + + let s = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`); + if (data.sessionResetsIn) s += " " + theme.fg("dim", data.sessionResetsIn); + + let w = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`); + if (data.weeklyResetsIn) w += " " + theme.fg("dim", `\u27F3 ${data.weeklyResetsIn}`); + + ctx.ui.setStatus(STATUS_KEY, s + theme.fg("dim", " | ") + w); + } + + function updateProviderFrom(modelLike: any): boolean { + const previous = state.activeProvider; + state.activeProvider = detectProvider(modelLike); + if (previous !== state.activeProvider) { updateStatus(); return true; } + return false; + } + + // --------------------------------------------------------------------------- + // Polling + // --------------------------------------------------------------------------- + async function runPollInner(options: PollOptions = {}) { + const auth = readAuth(); + const active = state.activeProvider; + + if (!canShowForProvider(active, auth, endpoints) || !auth || !active) { + state.lastPoll = Date.now(); updateStatus(); return; + } + + const cache = readUsageCache(); + const now = Date.now(); + const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS; + + const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0; + if (now < blockedUntil) { + if (cache?.data?.[active]) { + state[active] = cache.data[active]!; + } else { + // Rate-limited but no cached data — show a meaningful status instead + // of leaving state null (which shows eternal "loading…"). + const retryMin = Math.ceil((blockedUntil - now) / 60000); + state[active] = { session: 0, weekly: 0, error: `rate limited (retry in ${retryMin}m)` }; + } + state.lastPoll = now; updateStatus(); return; + } + + if (!options.forceFresh && cache && now - cache.timestamp < cacheTtl && cache.data?.[active]) { + state[active] = cache.data[active]!; + state.lastPoll = now; updateStatus(); return; + } + + const oauthId = providerToOAuthProviderId(active); + let effectiveAuth = auth; + if (oauthId && active !== "zai") { + const creds = auth[oauthId as keyof typeof auth] as + | { access?: string; refresh?: string; expires?: number } | undefined; + const expires = typeof creds?.expires === "number" ? creds.expires : 0; + const tokenExpiredOrMissing = !creds?.access || (expires > 0 && Date.now() + 60_000 >= expires); + if (tokenExpiredOrMissing && creds?.refresh) { + try { + const refreshPromise = ensureFreshAuthForProviders([oauthId as OAuthProviderId], { auth, persist: true }); + const timeoutPromise = new Promise((_, reject) => + setTimeout(() => reject(new Error("OAuth refresh timeout")), 15_000), + ); + const refreshed = await Promise.race([refreshPromise, timeoutPromise]); + if (refreshed.auth) effectiveAuth = refreshed.auth; + } catch {} + } + } + + let result: UsageData; + if (active === "codex") { + const access = effectiveAuth["openai-codex"]?.access; + result = access ? await fetchCodexUsage(access) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } else if (active === "claude") { + const access = effectiveAuth.anthropic?.access; + result = access ? await fetchClaudeUsage(access) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } else if (active === "zai") { + const token = effectiveAuth.zai?.access || effectiveAuth.zai?.key; + result = token ? await fetchZaiUsage(token, { endpoints }) + : { session: 0, weekly: 0, error: "missing token (try /login again)" }; + } else if (active === "gemini") { + const creds = effectiveAuth["google-gemini-cli"]; + result = creds?.access + ? await fetchGoogleUsage(creds.access, endpoints.gemini, creds.projectId, "gemini", { endpoints }) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } else { + const creds = effectiveAuth["google-antigravity"]; + result = creds?.access + ? await fetchGoogleUsage(creds.access, endpoints.antigravity, creds.projectId, "antigravity", { endpoints }) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } + + state[active] = result; + + if (result.error) { + if (result.error === "HTTP 429") { + const nextCache: import("./core").UsageCache = { + timestamp: cache?.timestamp ?? now, + data: { ...(cache?.data ?? {}) }, + rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}), [active]: now + RATE_LIMITED_BACKOFF_MS }, + }; + writeUsageCache(nextCache); + } + } else { + const nextCache: import("./core").UsageCache = { + timestamp: now, + data: { ...(cache?.data ?? {}), [active]: result }, + rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) }, + }; + delete nextCache.rateLimitedUntil![active]; + writeUsageCache(nextCache); + } + + state.lastPoll = now; + updateStatus(); + } + + async function runPoll(options: PollOptions = {}): Promise { + const timeout = new Promise((_, reject) => + setTimeout(() => reject(new Error("runPoll timeout")), 25_000), + ); + await Promise.race([runPollInner(options), timeout]); + } + + const POLL_TIMEOUT_MS = 30_000; + + async function poll(options: PollOptions = {}) { + // If a previous poll has been running longer than POLL_TIMEOUT_MS, abandon it + // so we don't queue forever behind a stuck request. + if (pollInFlight && pollStartedAt > 0 && Date.now() - pollStartedAt > POLL_TIMEOUT_MS) { + pollInFlight = null; + pollQueued = false; + const active = state.activeProvider; + if (active && !state[active]) { + state[active] = { session: 0, weekly: 0, error: "poll timeout" }; + updateStatus(); + } + } + + if (pollInFlight) { pollQueued = true; await pollInFlight; return; } + do { + pollQueued = false; + pollStartedAt = Date.now(); + pollInFlight = runPoll(options).catch(() => { + // If runPoll threw, ensure we don't leave status stuck at "loading…" + const active = state.activeProvider; + if (active && !state[active]) { + state[active] = { session: 0, weekly: 0, error: "poll failed" }; + updateStatus(); + } + }).finally(() => { pollInFlight = null; pollStartedAt = 0; }); + await pollInFlight; + } while (pollQueued); + } + + function startStreamingTimer() { + if (streamingTimer !== null) return; + streamingTimer = setInterval(() => { void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); }, STREAMING_POLL_INTERVAL_MS); + } + + function stopStreamingTimer() { + if (streamingTimer !== null) { clearInterval(streamingTimer); streamingTimer = null; } + } + + // ── Lifecycle ──────────────────────────────────────────────────────────── + + pi.on("session_start", async (_event, _ctx) => { + ctx = _ctx; + updateProviderFrom(_ctx.model); + await poll(); + }); + + pi.on("session_shutdown", async (_event, _ctx) => { + stopStreamingTimer(); + if (_ctx?.hasUI) _ctx.ui.setStatus(STATUS_KEY, undefined); + }); + + pi.on("model_select", async (event, _ctx) => { + ctx = _ctx; + const changed = updateProviderFrom(event.model ?? _ctx.model); + if (changed) await poll(); + }); + + pi.on("turn_start", (_event, _ctx) => { ctx = _ctx; updateProviderFrom(_ctx.model); }); + + pi.on("before_agent_start", async (_event, _ctx) => { + ctx = _ctx; + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + }); + + pi.on("agent_start", (_event, _ctx) => { ctx = _ctx; startStreamingTimer(); }); + + pi.on("agent_end", async (_event, _ctx) => { + ctx = _ctx; + stopStreamingTimer(); + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + }); + + pi.events.on("claude-account:switched", () => { + const cache = readUsageCache(); + if (cache?.data?.claude) { + const nextCache: import("./core").UsageCache = { ...cache, data: { ...cache.data } }; + delete nextCache.data.claude; + writeUsageCache(nextCache); + } + void poll({ forceFresh: true }); + }); + + // ── /usage command ─────────────────────────────────────────────────────── + + pi.registerCommand("usage", { + description: "Show API usage for all subscriptions", + handler: async (_args, _ctx) => { + ctx = _ctx; + updateProviderFrom(_ctx.model); + try { + if (_ctx?.hasUI) { + await _ctx.ui.custom((tui, theme, _keybindings, done) => { + return new UsageSelectorComponent( + tui, theme, state.activeProvider, + () => fetchAllUsages({ endpoints }), + () => done(), + ); + }); + } + } finally { + await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); + } + }, + }); +}