import type { ExtensionAPI, ExtensionContext, Theme } from "@mariozechner/pi-coding-agent"; import { truncateToWidth, visibleWidth } from "@mariozechner/pi-tui"; import * as fs from "node:fs"; import { homedir } from "node:os"; import { join } from "node:path"; type ProviderName = "anthropic" | "codex" | "gemini" | "opencode-go"; interface RateWindow { label: string; usedPercent: number; resetDescription?: string; resetAt?: string; } interface UsageSnapshot { provider: ProviderName; displayName: string; windows: RateWindow[]; error?: string; fetchedAt: number; lastSuccessAt?: number; fromCache?: boolean; } interface ProviderCache { usage?: UsageSnapshot; lastSuccessAt?: number; } interface CodexRateWindow { reset_at?: number; limit_window_seconds?: number; used_percent?: number; } interface CodexRateLimit { primary_window?: CodexRateWindow; secondary_window?: CodexRateWindow; } interface PiAuthShape { anthropic?: { access?: string }; "google-gemini-cli"?: { access?: string }; "openai-codex"?: { access?: string; accountId?: string }; } const CACHE_TTL_MS = 5 * 60 * 1000; const ANTHROPIC_CACHE_TTL_MS = 30 * 60 * 1000; const REFRESH_MS = 5 * 60 * 1000; const PROVIDER_ORDER: ProviderName[] = ["anthropic", "codex", "gemini", "opencode-go"]; const SHORTCUT_TOGGLE = "ctrl+alt+b"; const SHORTCUT_BAR_STYLE = "ctrl+alt+t"; const showToggleState = async (ctx: ExtensionContext, next: boolean, refresh: () => Promise) => { if (!next) { ctx.ui.setWidget("sub-bar-local", undefined); ctx.ui.notify("sub bar hidden", "info"); return; } ctx.ui.notify("sub bar shown", "info"); await refresh(); }; const OPENCODE_CONFIG_FILE = join(homedir(), ".config", "opencode", "opencode-go-usage.json"); const PI_AUTH_FILE = join(homedir(), ".pi", "agent", "auth.json"); function readJsonFile(path: string): T | undefined { try { if (!fs.existsSync(path)) return undefined; const content = fs.readFileSync(path, "utf-8"); return JSON.parse(content) as T; } catch { return undefined; } } function clampPercent(value: number | undefined): number { if (typeof value !== "number" || Number.isNaN(value)) return 0; return Math.max(0, Math.min(100, Math.round(value))); } function formatDuration(seconds?: number): string | undefined { if (typeof seconds !== "number" || seconds <= 0) return undefined; if (seconds < 60) return `${Math.max(1, Math.round(seconds))}s`; if (seconds < 3600) return `${Math.floor(seconds / 60)}m`; if (seconds < 86400) { const h = Math.floor(seconds / 3600); const m = Math.floor((seconds % 3600) / 60); return m > 0 ? `${h}h ${m}m` : `${h}h`; } const d = Math.floor(seconds / 86400); const h = Math.floor((seconds % 86400) / 3600); return h > 0 ? `${d}d ${h}h` : `${d}d`; } function formatResetDate(resetIso?: string): string | undefined { if (!resetIso) return undefined; const resetDate = new Date(resetIso); if (Number.isNaN(resetDate.getTime())) return undefined; const diffSec = Math.max(0, Math.floor((resetDate.getTime() - Date.now()) / 1000)); return formatDuration(diffSec); } function getErrorAge(lastSuccessAt?: number): string | undefined { if (!lastSuccessAt) return undefined; const diffSec = Math.floor((Date.now() - lastSuccessAt) / 1000); return formatDuration(diffSec); } function getPiAuth(): PiAuthShape { return readJsonFile(PI_AUTH_FILE) ?? {}; } function loadAnthropicToken(): string | undefined { const env = process.env.ANTHROPIC_OAUTH_TOKEN?.trim(); if (env) return env; return getPiAuth().anthropic?.access; } function loadGeminiToken(): string | undefined { const env = ( process.env.GOOGLE_GEMINI_CLI_OAUTH_TOKEN || process.env.GOOGLE_GEMINI_CLI_ACCESS_TOKEN || process.env.GEMINI_OAUTH_TOKEN )?.trim(); if (env) return env; return getPiAuth()["google-gemini-cli"]?.access; } function loadCodexCredentials(): { accessToken?: string; accountId?: string } { const envToken = ( process.env.OPENAI_CODEX_OAUTH_TOKEN || process.env.OPENAI_CODEX_ACCESS_TOKEN || process.env.CODEX_OAUTH_TOKEN || process.env.CODEX_ACCESS_TOKEN )?.trim(); const envAccountId = (process.env.OPENAI_CODEX_ACCOUNT_ID || process.env.CHATGPT_ACCOUNT_ID)?.trim(); if (envToken) { return { accessToken: envToken, accountId: envAccountId || undefined }; } const auth = getPiAuth()["openai-codex"]; if (!auth?.access) return {}; return { accessToken: auth.access, accountId: auth.accountId }; } function loadOpenCodeGoCredentials(): { workspaceId?: string; authCookie?: string } { const envWorkspaceId = process.env.OPENCODE_GO_WORKSPACE_ID?.trim(); const envAuthCookie = process.env.OPENCODE_GO_AUTH_COOKIE?.trim(); if (envWorkspaceId && envAuthCookie) { return { workspaceId: envWorkspaceId, authCookie: envAuthCookie }; } const config = readJsonFile<{ workspaceId?: string; authCookie?: string }>(OPENCODE_CONFIG_FILE); return { workspaceId: config?.workspaceId?.trim(), authCookie: config?.authCookie?.trim(), }; } function modelToProvider(modelProvider?: string): ProviderName | undefined { const id = (modelProvider ?? "").toLowerCase(); if (id.includes("opencode-go")) return "opencode-go"; if (id.includes("anthropic")) return "anthropic"; if (id.includes("gemini") || id.includes("google")) return "gemini"; if (id.includes("codex") || id.includes("openai")) return "codex"; return undefined; } function codexWindowLabel(window: CodexRateWindow | undefined, fallback: string): string { if (!window || typeof window.limit_window_seconds !== "number" || window.limit_window_seconds <= 0) { return fallback; } return formatDuration(window.limit_window_seconds) ?? fallback; } function pushCodexWindow(windows: RateWindow[], fallbackLabel: string, window?: CodexRateWindow): void { if (!window) return; const resetIso = typeof window.reset_at === "number" ? new Date(window.reset_at * 1000).toISOString() : undefined; windows.push({ label: codexWindowLabel(window, fallbackLabel), usedPercent: clampPercent(window.used_percent), resetAt: resetIso, }); } function barForPercent(theme: Theme, usedPercent: number, width = 8, style: "thin" | "thick" = "thick"): string { const safeWidth = Math.max(1, width); const filled = Math.round((clampPercent(usedPercent) / 100) * safeWidth); const empty = Math.max(0, safeWidth - filled); const color = usedPercent >= 85 ? "error" : usedPercent >= 60 ? "warning" : "success"; const filledChar = style === "thin" ? "─" : "█"; const emptyChar = style === "thin" ? "─" : "░"; return `${theme.fg(color, filledChar.repeat(filled))}${theme.fg("dim", emptyChar.repeat(empty))}`; } function padToWidth(text: string, width: number): string { const w = visibleWidth(text); if (w >= width) return truncateToWidth(text, width); return text + " ".repeat(width - w); } function buildColumnWidths(totalWidth: number, columns: number): number[] { if (columns <= 0) return []; const base = Math.max(1, Math.floor(totalWidth / columns)); let remainder = Math.max(0, totalWidth - base * columns); const widths: number[] = []; for (let i = 0; i < columns; i++) { const extra = remainder > 0 ? 1 : 0; widths.push(base + extra); if (remainder > 0) remainder--; } return widths; } function formatUsageTwoLines( theme: Theme, usage: UsageSnapshot, width: number, statusNote?: string, barStyle: "thin" | "thick" = "thick", ): { top: string; bottom?: string } { const provider = theme.bold(theme.fg("accent", usage.displayName)); if (usage.error && usage.windows.length === 0) { const age = getErrorAge(usage.lastSuccessAt); const stale = age ? ` (stale ${age})` : ""; return { top: truncateToWidth(`${provider} ${theme.fg("warning", `⚠ ${usage.error}${stale}`)}`, width) }; } if (usage.windows.length === 0) { return { top: truncateToWidth(provider, width) }; } const prefix = `${provider} ${theme.fg("dim", "•")} `; const prefixWidth = Math.min(visibleWidth(prefix), Math.max(0, width - 1)); const contentWidth = Math.max(1, width - prefixWidth); const gap = " "; const gapWidth = visibleWidth(gap); const windowsCount = usage.windows.length; const totalGapWidth = Math.max(0, (windowsCount - 1) * gapWidth); const columnsArea = Math.max(1, contentWidth - totalGapWidth); const colWidths = buildColumnWidths(columnsArea, windowsCount); const topCols = usage.windows.map((window, index) => { const colWidth = colWidths[index] ?? 1; const pct = clampPercent(window.usedPercent); const resetRaw = formatResetDate(window.resetAt) ?? window.resetDescription ?? ""; const left = `${window.label} ${pct}%`; const right = resetRaw; if (!right) return padToWidth(theme.fg("text", left), colWidth); const leftWidth = visibleWidth(left); const rightWidth = visibleWidth(right); if (leftWidth + 1 + rightWidth <= colWidth) { const spaces = " ".repeat(colWidth - leftWidth - rightWidth); return `${theme.fg("text", left)}${spaces}${theme.fg("dim", right)}`; } const compact = `${left} ${right}`; return padToWidth(theme.fg("text", truncateToWidth(compact, colWidth)), colWidth); }); const bottomCols = usage.windows.map((window, index) => { const colWidth = colWidths[index] ?? 1; const pct = clampPercent(window.usedPercent); return barForPercent(theme, pct, colWidth, barStyle); }); let top = prefix + topCols.join(gap); if (statusNote) { const note = theme.fg("warning", ` ⚠ ${statusNote}`); top = truncateToWidth(top + note, width); } const bottomPrefix = " ".repeat(prefixWidth); const bottom = bottomPrefix + bottomCols.join(gap); return { top: padToWidth(top, width), bottom: padToWidth(bottom, width), }; } async function fetchAnthropicUsage(): Promise { const token = loadAnthropicToken(); if (!token) throw new Error("missing anthropic oauth token"); const response = await fetch("https://api.anthropic.com/api/oauth/usage", { headers: { Authorization: `Bearer ${token}`, "anthropic-beta": "oauth-2025-04-20", }, }); if (!response.ok) throw new Error(`anthropic http ${response.status}`); const data = (await response.json()) as { five_hour?: { utilization?: number; resets_at?: string }; seven_day?: { utilization?: number; resets_at?: string }; extra_usage?: { utilization?: number; used_credits?: number; monthly_limit?: number; is_enabled?: boolean }; }; const windows: RateWindow[] = []; if (data.five_hour?.utilization !== undefined) { windows.push({ label: "5h", usedPercent: clampPercent(data.five_hour.utilization), resetAt: data.five_hour.resets_at }); } if (data.seven_day?.utilization !== undefined) { windows.push({ label: "Week", usedPercent: clampPercent(data.seven_day.utilization), resetAt: data.seven_day.resets_at }); } // hide Anthropic extra_usage window for now (it is confusing/noisy for premium users) if (windows.length === 0) throw new Error("no anthropic usage windows"); const now = Date.now(); return { provider: "anthropic", displayName: "Claude Plan", windows, fetchedAt: now, lastSuccessAt: now, }; } async function fetchCodexUsage(): Promise { const { accessToken, accountId } = loadCodexCredentials(); if (!accessToken) throw new Error("missing codex oauth token"); const headers: Record = { Authorization: `Bearer ${accessToken}`, Accept: "application/json", }; if (accountId) headers["ChatGPT-Account-Id"] = accountId; const response = await fetch("https://chatgpt.com/backend-api/wham/usage", { headers }); if (!response.ok) throw new Error(`codex http ${response.status}`); const data = (await response.json()) as { rate_limit?: CodexRateLimit; }; const windows: RateWindow[] = []; pushCodexWindow(windows, "3h", data.rate_limit?.primary_window); pushCodexWindow(windows, "Day", data.rate_limit?.secondary_window); if (windows.length === 0) throw new Error("no codex usage windows"); const now = Date.now(); return { provider: "codex", displayName: "Codex Plan", windows, fetchedAt: now, lastSuccessAt: now, }; } async function fetchGeminiUsage(): Promise { const token = loadGeminiToken(); if (!token) throw new Error("missing gemini oauth token"); const response = await fetch("https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota", { method: "POST", headers: { Authorization: `Bearer ${token}`, "Content-Type": "application/json", }, body: "{}", }); if (!response.ok) throw new Error(`gemini http ${response.status}`); const data = (await response.json()) as { buckets?: Array<{ modelId?: string; remainingFraction?: number }>; }; let minPro = 1; let minFlash = 1; let hasPro = false; let hasFlash = false; for (const bucket of data.buckets ?? []) { const id = bucket.modelId?.toLowerCase() ?? ""; const remaining = typeof bucket.remainingFraction === "number" ? bucket.remainingFraction : 1; if (id.includes("pro")) { hasPro = true; minPro = Math.min(minPro, remaining); } if (id.includes("flash")) { hasFlash = true; minFlash = Math.min(minFlash, remaining); } } const windows: RateWindow[] = []; if (hasPro) windows.push({ label: "Pro", usedPercent: clampPercent((1 - minPro) * 100) }); if (hasFlash) windows.push({ label: "Flash", usedPercent: clampPercent((1 - minFlash) * 100) }); if (windows.length === 0) throw new Error("no gemini usage windows"); const now = Date.now(); return { provider: "gemini", displayName: "Gemini Plan", windows, fetchedAt: now, lastSuccessAt: now, }; } function parseInlineObject(raw: string): Record { const normalized = raw.replace(/([{,]\s*)([a-zA-Z_][a-zA-Z0-9_]*)(\s*:)/g, "$1\"$2\"$3"); return JSON.parse(normalized) as Record; } async function fetchOpenCodeGoUsage(): Promise { const { workspaceId, authCookie } = loadOpenCodeGoCredentials(); if (!workspaceId || !authCookie) throw new Error(`missing ${OPENCODE_CONFIG_FILE} credentials`); if (!/^wrk_[a-zA-Z0-9]+$/.test(workspaceId)) throw new Error("invalid workspace id format"); const url = `https://opencode.ai/workspace/${encodeURIComponent(workspaceId)}/go`; const response = await fetch(url, { headers: { "User-Agent": "Mozilla/5.0", Accept: "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8", Cookie: `auth=${authCookie}`, }, }); if (!response.ok) throw new Error(`opencode-go http ${response.status}`); const html = await response.text(); const patterns: Record = { Rolling: /rollingUsage:\$R\[\d+\]=(\{[^}]+\})/, Weekly: /weeklyUsage:\$R\[\d+\]=(\{[^}]+\})/, Monthly: /monthlyUsage:\$R\[\d+\]=(\{[^}]+\})/, }; const windows: RateWindow[] = []; for (const [label, pattern] of Object.entries(patterns)) { const match = html.match(pattern); if (!match?.[1]) continue; try { const parsed = parseInlineObject(match[1]); const usagePercent = clampPercent(Number(parsed.usagePercent)); const resetSec = Number(parsed.resetInSec); const resetDescription = formatDuration(Number.isFinite(resetSec) ? resetSec : undefined); const resetAt = Number.isFinite(resetSec) && resetSec > 0 ? new Date(Date.now() + resetSec * 1000).toISOString() : undefined; windows.push({ label, usedPercent: usagePercent, resetDescription, resetAt, }); } catch { // Ignore malformed chunks and keep parsing others } } if (windows.length === 0) throw new Error("could not parse opencode-go usage from page"); const now = Date.now(); return { provider: "opencode-go", displayName: "OpenCode Go", windows, fetchedAt: now, lastSuccessAt: now, }; } async function fetchUsage(provider: ProviderName): Promise { switch (provider) { case "anthropic": return fetchAnthropicUsage(); case "codex": return fetchCodexUsage(); case "gemini": return fetchGeminiUsage(); case "opencode-go": return fetchOpenCodeGoUsage(); } } export default function createSubBarLocal(pi: ExtensionAPI) { let lastCtx: ExtensionContext | undefined; let activeProvider: ProviderName | "auto" = "auto"; let widgetEnabled = true; let barStyle: "thin" | "thick" = "thick"; let refreshTimer: NodeJS.Timeout | undefined; let anthropicRetryAfter = 0; const cache: Partial> = {}; function getSelectedProvider(ctx?: ExtensionContext): ProviderName | undefined { if (activeProvider !== "auto") return activeProvider; return modelToProvider(ctx?.model?.provider); } function render(ctx: ExtensionContext): void { if (!ctx.hasUI) return; if (!widgetEnabled) { ctx.ui.setWidget("sub-bar-local", undefined); return; } const provider = getSelectedProvider(ctx); if (!provider) { ctx.ui.setWidget("sub-bar-local", undefined); return; } const snapshot = cache[provider]?.usage; const setWidgetWithPlacement = (ctx.ui as unknown as { setWidget: (...args: unknown[]) => void }).setWidget; setWidgetWithPlacement( "sub-bar-local", (_tui: unknown, theme: Theme) => ({ render(width: number) { const safeWidth = Math.max(1, width); const topDivider = theme.fg("dim", "─".repeat(safeWidth)); if (!snapshot) { const cooldown = provider === "anthropic" && anthropicRetryAfter > Date.now() ? formatDuration(Math.max(1, Math.floor((anthropicRetryAfter - Date.now()) / 1000))) : undefined; const text = cooldown ? `anthropic usage limited, retry in ${cooldown}` : `sub bar loading ${provider} usage...`; const loading = truncateToWidth(theme.fg("dim", text), safeWidth); return [topDivider, padToWidth(loading, safeWidth)]; } const cooldown = provider === "anthropic" && anthropicRetryAfter > Date.now() ? formatDuration(Math.max(1, Math.floor((anthropicRetryAfter - Date.now()) / 1000))) : undefined; const statusNote = snapshot.error ? snapshot.error : cooldown ? `usage endpoint limited, retry in ${cooldown}` : undefined; const lines = formatUsageTwoLines(theme, snapshot, safeWidth, statusNote, barStyle); const output = [topDivider, lines.top]; if (lines.bottom) output.push(lines.bottom); return output; }, invalidate() {}, }), { placement: "aboveEditor" }, ); } async function refreshCurrent(ctx: ExtensionContext, force = false): Promise { const provider = getSelectedProvider(ctx); if (!provider) { ctx.ui.setWidget("sub-bar-local", undefined); return; } if (provider === "anthropic" && anthropicRetryAfter > Date.now()) { render(ctx); return; } const cached = cache[provider]?.usage; const ttl = provider === "anthropic" ? ANTHROPIC_CACHE_TTL_MS : CACHE_TTL_MS; if (!force && cached && Date.now() - cached.fetchedAt < ttl) { render(ctx); return; } try { const fresh = await fetchUsage(provider); if (provider === "anthropic") { anthropicRetryAfter = 0; } cache[provider] = { usage: { ...fresh, fromCache: false, lastSuccessAt: Date.now(), }, lastSuccessAt: Date.now(), }; } catch (error) { const message = error instanceof Error ? error.message : "fetch failed"; const fallback = cache[provider]?.usage; const isAnthropic429 = provider === "anthropic" && message.includes("429"); if (isAnthropic429) { anthropicRetryAfter = Date.now() + 30 * 60 * 1000; } if (isAnthropic429 && fallback) { cache[provider] = { usage: { ...fallback, fetchedAt: Date.now(), fromCache: true, }, lastSuccessAt: cache[provider]?.lastSuccessAt, }; } else if (isAnthropic429) { delete cache[provider]; } else { cache[provider] = { usage: { provider, displayName: fallback?.displayName ?? provider, windows: fallback?.windows ?? [], error: message, fetchedAt: Date.now(), lastSuccessAt: cache[provider]?.lastSuccessAt, fromCache: Boolean(fallback), }, lastSuccessAt: cache[provider]?.lastSuccessAt, }; } } render(ctx); } function startRefreshLoop(ctx: ExtensionContext): void { if (refreshTimer) clearInterval(refreshTimer); refreshTimer = setInterval(() => { if (!lastCtx) return; void refreshCurrent(lastCtx); }, REFRESH_MS); refreshTimer.unref?.(); void refreshCurrent(ctx, true); } pi.registerCommand("sub:refresh", { description: "Refresh sub bar usage now", handler: async (_args, ctx) => { await refreshCurrent(ctx, true); }, }); pi.registerCommand("sub:provider", { description: "Set sub bar provider (auto|anthropic|codex|gemini|opencode-go)", handler: async (args, ctx) => { const raw = String(args ?? "").trim().toLowerCase(); if (!raw || raw === "auto") { activeProvider = "auto"; ctx.ui.notify("sub bar provider: auto", "info"); await refreshCurrent(ctx, true); return; } if (PROVIDER_ORDER.includes(raw as ProviderName)) { activeProvider = raw as ProviderName; ctx.ui.notify(`sub bar provider: ${activeProvider}`, "info"); await refreshCurrent(ctx, true); return; } ctx.ui.notify("invalid provider. use: auto|anthropic|codex|gemini|opencode-go", "warning"); }, }); pi.registerCommand("sub:toggle", { description: "Toggle sub bar on/off", handler: async (_args, ctx) => { widgetEnabled = !widgetEnabled; await showToggleState(ctx, widgetEnabled, async () => refreshCurrent(ctx, true)); }, }); pi.registerShortcut(SHORTCUT_TOGGLE as import("@mariozechner/pi-tui").KeyId, { description: "Toggle sub bar on/off", handler: async (ctx) => { widgetEnabled = !widgetEnabled; await showToggleState(ctx, widgetEnabled, async () => refreshCurrent(ctx, true)); }, }); pi.registerCommand("sub:bars", { description: "Set sub bar style (thin|thick|toggle)", handler: async (args, ctx) => { const raw = String(args ?? "").trim().toLowerCase(); if (!raw || raw === "toggle") { barStyle = barStyle === "thin" ? "thick" : "thin"; } else if (raw === "thin" || raw === "thick") { barStyle = raw; } else { ctx.ui.notify("invalid style. use: thin|thick|toggle", "warning"); return; } ctx.ui.notify(`sub bar style: ${barStyle}`, "info"); render(ctx); }, }); pi.registerShortcut(SHORTCUT_BAR_STYLE as import("@mariozechner/pi-tui").KeyId, { description: "Toggle sub bar bar style", handler: async (ctx) => { barStyle = barStyle === "thin" ? "thick" : "thin"; ctx.ui.notify(`sub bar style: ${barStyle}`, "info"); render(ctx); }, }); pi.on("session_start", async (_event, ctx) => { lastCtx = ctx; if (!ctx.hasUI) return; render(ctx); startRefreshLoop(ctx); }); pi.on("model_select", async (_event, ctx) => { lastCtx = ctx; if (!ctx.hasUI) return; render(ctx); if (activeProvider === "auto") { await refreshCurrent(ctx, false); } }); pi.on("session_shutdown", async () => { lastCtx = undefined; if (refreshTimer) { clearInterval(refreshTimer); refreshTimer = undefined; } }); }