/ extensions / usage-bar.ts
usage-bar.ts
1 /** 2 * Usage Bar Extension - Shows AI provider usage stats like CodexBar 3 * Run /usage to see usage for Claude, Copilot, Gemini, and Codex 4 * 5 * Features: 6 * - Usage stats with progress bars 7 * - Provider status (outages/incidents) 8 * - Reset countdowns 9 */ 10 11 import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; 12 import { visibleWidth } from "@mariozechner/pi-tui"; 13 import * as fs from "node:fs"; 14 import * as path from "node:path"; 15 import * as os from "node:os"; 16 import { execSync } from "node:child_process"; 17 18 // ============================================================================ 19 // Types 20 // ============================================================================ 21 22 interface RateWindow { 23 label: string; 24 usedPercent: number; 25 resetDescription?: string; 26 resetsAt?: Date; 27 } 28 29 interface ProviderStatus { 30 indicator: "none" | "minor" | "major" | "critical" | "maintenance" | "unknown"; 31 description?: string; 32 } 33 34 interface UsageSnapshot { 35 provider: string; 36 displayName: string; 37 windows: RateWindow[]; 38 plan?: string; 39 error?: string; 40 status?: ProviderStatus; 41 } 42 43 // ============================================================================ 44 // Status Polling 45 // ============================================================================ 46 47 const STATUS_URLS: Record<string, string> = { 48 anthropic: "https://status.anthropic.com/api/v2/status.json", 49 codex: "https://status.openai.com/api/v2/status.json", 50 copilot: "https://www.githubstatus.com/api/v2/status.json", 51 }; 52 53 async function fetchProviderStatus(provider: string): Promise<ProviderStatus> { 54 const url = STATUS_URLS[provider]; 55 if (!url) return { indicator: "none" }; 56 57 try { 58 const controller = new AbortController(); 59 setTimeout(() => controller.abort(), 5000); 60 61 const res = await fetch(url, { signal: controller.signal }); 62 if (!res.ok) return { indicator: "unknown" }; 63 64 const data = await res.json() as any; 65 const indicator = data.status?.indicator || "none"; 66 const description = data.status?.description; 67 68 return { 69 indicator: indicator as ProviderStatus["indicator"], 70 description, 71 }; 72 } catch { 73 return { indicator: "unknown" }; 74 } 75 } 76 77 async function fetchGeminiStatus(): Promise<ProviderStatus> { 78 try { 79 const controller = new AbortController(); 80 setTimeout(() => controller.abort(), 5000); 81 82 const res = await fetch("https://www.google.com/appsstatus/dashboard/incidents.json", { 83 signal: controller.signal, 84 }); 85 if (!res.ok) return { indicator: "unknown" }; 86 87 const incidents = await res.json() as any[]; 88 89 // Look for active Gemini incidents (product ID: npdyhgECDJ6tB66MxXyo) 90 const geminiProductId = "npdyhgECDJ6tB66MxXyo"; 91 const activeIncidents = incidents.filter((inc: any) => { 92 if (inc.end) return false; // Not active 93 const affected = inc.currently_affected_products || inc.affected_products || []; 94 return affected.some((p: any) => p.id === geminiProductId); 95 }); 96 97 if (activeIncidents.length === 0) { 98 return { indicator: "none" }; 99 } 100 101 // Find most severe 102 let worstIndicator: ProviderStatus["indicator"] = "minor"; 103 let description: string | undefined; 104 105 for (const inc of activeIncidents) { 106 const status = inc.most_recent_update?.status || inc.status_impact; 107 if (status === "SERVICE_OUTAGE") { 108 worstIndicator = "critical"; 109 description = inc.external_desc; 110 } else if (status === "SERVICE_DISRUPTION" && worstIndicator !== "critical") { 111 worstIndicator = "major"; 112 description = inc.external_desc; 113 } 114 } 115 116 return { indicator: worstIndicator, description }; 117 } catch { 118 return { indicator: "unknown" }; 119 } 120 } 121 122 // ============================================================================ 123 // Claude Usage 124 // ============================================================================ 125 126 function loadClaudeToken(): string | undefined { 127 // Try pi's auth.json first (has user:profile scope) 128 const piAuthPath = path.join(os.homedir(), ".pi", "agent", "auth.json"); 129 try { 130 if (fs.existsSync(piAuthPath)) { 131 const data = JSON.parse(fs.readFileSync(piAuthPath, "utf-8")); 132 if (data.anthropic?.access) return data.anthropic.access; 133 } 134 } catch {} 135 136 // Fallback to Claude CLI keychain (macOS) 137 try { 138 const keychainData = execSync( 139 'security find-generic-password -s "Claude Code-credentials" -w 2>/dev/null', 140 { encoding: "utf-8", stdio: ["pipe", "pipe", "pipe"] } 141 ).trim(); 142 if (keychainData) { 143 const parsed = JSON.parse(keychainData); 144 const scopes = parsed.claudeAiOauth?.scopes || []; 145 if (scopes.includes("user:profile") && parsed.claudeAiOauth?.accessToken) { 146 return parsed.claudeAiOauth.accessToken; 147 } 148 } 149 } catch {} 150 151 return undefined; 152 } 153 154 async function fetchClaudeUsage(): Promise<UsageSnapshot> { 155 const token = loadClaudeToken(); 156 if (!token) { 157 return { provider: "anthropic", displayName: "Claude", windows: [], error: "No credentials" }; 158 } 159 160 try { 161 const controller = new AbortController(); 162 setTimeout(() => controller.abort(), 5000); 163 164 const res = await fetch("https://api.anthropic.com/api/oauth/usage", { 165 headers: { 166 Authorization: `Bearer ${token}`, 167 "anthropic-beta": "oauth-2025-04-20", 168 }, 169 signal: controller.signal, 170 }); 171 172 if (!res.ok) { 173 return { provider: "anthropic", displayName: "Claude", windows: [], error: `HTTP ${res.status}` }; 174 } 175 176 const data = await res.json() as any; 177 const windows: RateWindow[] = []; 178 179 if (data.five_hour?.utilization !== undefined) { 180 windows.push({ 181 label: "5h", 182 usedPercent: data.five_hour.utilization, 183 resetDescription: data.five_hour.resets_at ? formatReset(new Date(data.five_hour.resets_at)) : undefined, 184 }); 185 } 186 187 if (data.seven_day?.utilization !== undefined) { 188 windows.push({ 189 label: "Week", 190 usedPercent: data.seven_day.utilization, 191 resetDescription: data.seven_day.resets_at ? formatReset(new Date(data.seven_day.resets_at)) : undefined, 192 }); 193 } 194 195 const modelWindow = data.seven_day_sonnet || data.seven_day_opus; 196 if (modelWindow?.utilization !== undefined) { 197 windows.push({ 198 label: data.seven_day_sonnet ? "Sonnet" : "Opus", 199 usedPercent: modelWindow.utilization, 200 }); 201 } 202 203 return { provider: "anthropic", displayName: "Claude", windows }; 204 } catch (e) { 205 return { provider: "anthropic", displayName: "Claude", windows: [], error: String(e) }; 206 } 207 } 208 209 // ============================================================================ 210 // Copilot Usage 211 // ============================================================================ 212 213 function loadCopilotRefreshToken(): string | undefined { 214 // The copilot_internal/user endpoint needs the GitHub OAuth token (ghu_*), 215 // NOT the Copilot session token (tid=*). The refresh token IS the GitHub OAuth token. 216 const authPath = path.join(os.homedir(), ".pi", "agent", "auth.json"); 217 try { 218 if (fs.existsSync(authPath)) { 219 const data = JSON.parse(fs.readFileSync(authPath, "utf-8")); 220 // Use refresh token (GitHub OAuth token ghu_*) for the usage API 221 if (data["github-copilot"]?.refresh) return data["github-copilot"].refresh; 222 } 223 } catch {} 224 225 return undefined; 226 } 227 228 async function fetchCopilotUsage(_modelRegistry: any): Promise<UsageSnapshot> { 229 const token = loadCopilotRefreshToken(); 230 if (!token) { 231 return { provider: "copilot", displayName: "Copilot", windows: [], error: "No token" }; 232 } 233 234 const headersBase = { 235 "Editor-Version": "vscode/1.96.2", 236 "User-Agent": "GitHubCopilotChat/0.26.7", 237 "X-Github-Api-Version": "2025-04-01", 238 Accept: "application/json", 239 }; 240 241 const tryFetch = async (authHeader: string) => { 242 const controller = new AbortController(); 243 setTimeout(() => controller.abort(), 5000); 244 245 const res = await fetch("https://api.github.com/copilot_internal/user", { 246 headers: { 247 ...headersBase, 248 Authorization: authHeader, 249 }, 250 signal: controller.signal, 251 }); 252 return res; 253 }; 254 255 try { 256 // Copilot access tokens (from /login github-copilot) expect Bearer. PATs accept "token". 257 // GitHub OAuth token (ghu_*) requires "token" prefix, not Bearer 258 const attempts = [`token ${token}`]; 259 let lastStatus: number | undefined; 260 let res: Response | undefined; 261 262 for (const auth of attempts) { 263 res = await tryFetch(auth); 264 lastStatus = res.status; 265 if (res.ok) break; 266 if (res.status === 401 || res.status === 403) continue; // try next scheme 267 break; 268 } 269 270 if (!res || !res.ok) { 271 const status = lastStatus ?? 0; 272 return { provider: "copilot", displayName: "Copilot", windows: [], error: `HTTP ${status}` }; 273 } 274 275 const data = await res.json() as any; 276 const windows: RateWindow[] = []; 277 278 // Parse reset date for display 279 const resetDate = data.quota_reset_date_utc ? new Date(data.quota_reset_date_utc) : undefined; 280 const resetDesc = resetDate ? formatReset(resetDate) : undefined; 281 282 // Premium interactions (e.g., Claude, o1 models) - has a cap 283 if (data.quota_snapshots?.premium_interactions) { 284 const pi = data.quota_snapshots.premium_interactions; 285 const remaining = pi.remaining ?? 0; 286 const entitlement = pi.entitlement ?? 0; 287 const usedPercent = Math.max(0, 100 - (pi.percent_remaining || 0)); 288 windows.push({ 289 label: `Premium`, 290 usedPercent, 291 resetDescription: resetDesc ? `${resetDesc} (${remaining}/${entitlement})` : `${remaining}/${entitlement}`, 292 }); 293 } 294 295 // Chat quota - often unlimited, only show if limited 296 if (data.quota_snapshots?.chat && !data.quota_snapshots.chat.unlimited) { 297 const chat = data.quota_snapshots.chat; 298 windows.push({ 299 label: "Chat", 300 usedPercent: Math.max(0, 100 - (chat.percent_remaining || 0)), 301 resetDescription: resetDesc, 302 }); 303 } 304 305 return { 306 provider: "copilot", 307 displayName: "Copilot", 308 windows, 309 plan: data.copilot_plan, 310 }; 311 } catch (e) { 312 return { provider: "copilot", displayName: "Copilot", windows: [], error: String(e) }; 313 } 314 } 315 316 // ============================================================================ 317 // Gemini Usage 318 // ============================================================================ 319 320 async function fetchGeminiUsage(_modelRegistry: any): Promise<UsageSnapshot> { 321 let token: string | undefined; 322 323 // Read directly from pi's auth.json 324 const piAuthPath = path.join(os.homedir(), ".pi", "agent", "auth.json"); 325 try { 326 if (fs.existsSync(piAuthPath)) { 327 const data = JSON.parse(fs.readFileSync(piAuthPath, "utf-8")); 328 token = data["google-gemini-cli"]?.access; 329 } 330 } catch {} 331 332 // Fallback to ~/.gemini/oauth_creds.json 333 if (!token) { 334 const credPath = path.join(os.homedir(), ".gemini", "oauth_creds.json"); 335 try { 336 if (fs.existsSync(credPath)) { 337 const data = JSON.parse(fs.readFileSync(credPath, "utf-8")); 338 token = data.access_token; 339 } 340 } catch {} 341 } 342 343 if (!token) { 344 return { provider: "gemini", displayName: "Gemini", windows: [], error: "No credentials" }; 345 } 346 347 try { 348 const controller = new AbortController(); 349 setTimeout(() => controller.abort(), 5000); 350 351 const res = await fetch("https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota", { 352 method: "POST", 353 headers: { Authorization: `Bearer ${token}`, "Content-Type": "application/json" }, 354 body: "{}", 355 signal: controller.signal, 356 }); 357 358 if (!res.ok) { 359 return { provider: "gemini", displayName: "Gemini", windows: [], error: `HTTP ${res.status}` }; 360 } 361 362 const data = await res.json() as any; 363 const quotas: Record<string, number> = {}; 364 365 for (const bucket of data.buckets || []) { 366 const model = bucket.modelId || "unknown"; 367 const frac = bucket.remainingFraction ?? 1; 368 if (!quotas[model] || frac < quotas[model]) quotas[model] = frac; 369 } 370 371 const windows: RateWindow[] = []; 372 let proMin = 1, flashMin = 1; 373 let hasProModel = false, hasFlashModel = false; 374 375 for (const [model, frac] of Object.entries(quotas)) { 376 if (model.toLowerCase().includes("pro")) { 377 hasProModel = true; 378 if (frac < proMin) proMin = frac; 379 } 380 if (model.toLowerCase().includes("flash")) { 381 hasFlashModel = true; 382 if (frac < flashMin) flashMin = frac; 383 } 384 } 385 386 // Always show windows if model exists (even at 0% usage) 387 if (hasProModel) windows.push({ label: "Pro", usedPercent: (1 - proMin) * 100 }); 388 if (hasFlashModel) windows.push({ label: "Flash", usedPercent: (1 - flashMin) * 100 }); 389 390 return { provider: "gemini", displayName: "Gemini", windows }; 391 } catch (e) { 392 return { provider: "gemini", displayName: "Gemini", windows: [], error: String(e) }; 393 } 394 } 395 396 // ============================================================================ 397 // Antigravity Usage 398 // ============================================================================ 399 400 type AntigravityAuth = { 401 accessToken: string; 402 refreshToken?: string; 403 expiresAt?: number; 404 projectId?: string; 405 }; 406 407 function loadAntigravityAuthFromPiAuthJson(): AntigravityAuth | undefined { 408 const piAuthPath = path.join(os.homedir(), ".pi", "agent", "auth.json"); 409 try { 410 if (!fs.existsSync(piAuthPath)) return undefined; 411 const data = JSON.parse(fs.readFileSync(piAuthPath, "utf-8")); 412 413 // Provider is called "google-antigravity" in pi. 414 const cred = data["google-antigravity"] ?? data["antigravity"] ?? data["anti-gravity"]; 415 if (!cred) return undefined; 416 417 const accessToken = typeof cred.access === "string" ? cred.access : undefined; 418 if (!accessToken) return undefined; 419 420 return { 421 accessToken, 422 refreshToken: typeof cred.refresh === "string" ? cred.refresh : undefined, 423 expiresAt: typeof cred.expires === "number" ? cred.expires : undefined, 424 projectId: typeof cred.projectId === "string" ? cred.projectId : typeof cred.project_id === "string" ? cred.project_id : undefined, 425 }; 426 } catch { 427 return undefined; 428 } 429 } 430 431 async function loadAntigravityAuth(modelRegistry: any): Promise<AntigravityAuth | undefined> { 432 // Prefer model registry auth storage first (may auto-refresh). 433 try { 434 const accessToken = await Promise.resolve(modelRegistry?.authStorage?.getApiKey?.("google-antigravity")); 435 const raw = await Promise.resolve(modelRegistry?.authStorage?.get?.("google-antigravity")); 436 437 const projectId = typeof raw?.projectId === "string" ? raw.projectId : undefined; 438 const refreshToken = typeof raw?.refresh === "string" ? raw.refresh : undefined; 439 const expiresAt = typeof raw?.expires === "number" ? raw.expires : undefined; 440 441 if (typeof accessToken === "string" && accessToken.length > 0) { 442 return { accessToken, projectId, refreshToken, expiresAt }; 443 } 444 } catch {} 445 446 // Fallback to pi auth.json 447 const fromPi = loadAntigravityAuthFromPiAuthJson(); 448 if (fromPi) return fromPi; 449 450 // Last resort: env var (won't have projectId; request will likely fail) 451 if (process.env.ANTIGRAVITY_API_KEY) { 452 return { accessToken: process.env.ANTIGRAVITY_API_KEY }; 453 } 454 455 return undefined; 456 } 457 458 async function refreshAntigravityAccessToken(refreshToken: string): Promise<{ accessToken: string; expiresAt?: number } | null> { 459 try { 460 const controller = new AbortController(); 461 setTimeout(() => controller.abort(), 5000); 462 463 // From the reference snippet in CodexBar issue #129. 464 const clientId = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"; 465 const clientSecret = "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"; 466 467 const res = await fetch("https://oauth2.googleapis.com/token", { 468 method: "POST", 469 headers: { "Content-Type": "application/x-www-form-urlencoded" }, 470 body: new URLSearchParams({ 471 client_id: clientId, 472 client_secret: clientSecret, 473 refresh_token: refreshToken, 474 grant_type: "refresh_token", 475 }).toString(), 476 signal: controller.signal, 477 }); 478 479 if (!res.ok) return null; 480 const data = (await res.json()) as any; 481 const accessToken = typeof data.access_token === "string" ? data.access_token : undefined; 482 if (!accessToken) return null; 483 const expiresIn = typeof data.expires_in === "number" ? data.expires_in : undefined; 484 return { 485 accessToken, 486 expiresAt: expiresIn ? Date.now() + expiresIn * 1000 : undefined, 487 }; 488 } catch { 489 return null; 490 } 491 } 492 493 async function fetchAntigravityUsage(modelRegistry: any): Promise<UsageSnapshot> { 494 const auth = await loadAntigravityAuth(modelRegistry); 495 if (!auth?.accessToken) { 496 return { provider: "antigravity", displayName: "Antigravity", windows: [], error: "No credentials" }; 497 } 498 499 if (!auth.projectId) { 500 return { provider: "antigravity", displayName: "Antigravity", windows: [], error: "Missing projectId" }; 501 } 502 503 let accessToken = auth.accessToken; 504 505 // Refresh if likely expired. 506 if (auth.refreshToken && auth.expiresAt && auth.expiresAt < Date.now() + 5 * 60 * 1000) { 507 const refreshed = await refreshAntigravityAccessToken(auth.refreshToken); 508 if (refreshed?.accessToken) accessToken = refreshed.accessToken; 509 } 510 511 const fetchModels = async (token: string): Promise<Response> => { 512 const controller = new AbortController(); 513 setTimeout(() => controller.abort(), 5000); 514 515 return fetch("https://cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels", { 516 method: "POST", 517 headers: { 518 Authorization: `Bearer ${token}`, 519 "Content-Type": "application/json", 520 "User-Agent": "antigravity/1.12.4", 521 "X-Goog-Api-Client": "google-cloud-sdk vscode_cloudshelleditor/0.1", 522 Accept: "application/json", 523 }, 524 body: JSON.stringify({ project: auth.projectId }), 525 signal: controller.signal, 526 }); 527 }; 528 529 try { 530 let res = await fetchModels(accessToken); 531 532 if ((res.status === 401 || res.status === 403) && auth.refreshToken) { 533 const refreshed = await refreshAntigravityAccessToken(auth.refreshToken); 534 if (refreshed?.accessToken) { 535 accessToken = refreshed.accessToken; 536 res = await fetchModels(accessToken); 537 } 538 } 539 540 if (res.status === 401 || res.status === 403) { 541 return { provider: "antigravity", displayName: "Antigravity", windows: [], error: "Unauthorized" }; 542 } 543 544 if (!res.ok) { 545 return { provider: "antigravity", displayName: "Antigravity", windows: [], error: `HTTP ${res.status}` }; 546 } 547 548 const data = (await res.json()) as any; 549 const models: Record<string, any> = data.models || {}; 550 551 const getQuotaInfo = (modelKeys: string[]): { usedPercent: number; resetDescription?: string } | null => { 552 for (const key of modelKeys) { 553 const qi = models?.[key]?.quotaInfo; 554 if (!qi) continue; 555 // In practice (CodexBar issue #129), some models only provide resetTime. 556 // Treat missing remainingFraction as 0% remaining (100% used), which matches Antigravity's behavior when quota is exhausted. 557 const remainingFraction = typeof qi.remainingFraction === "number" ? qi.remainingFraction : 0; 558 const usedPercent = Math.min(100, Math.max(0, (1 - remainingFraction) * 100)); 559 const resetTime = qi.resetTime ? new Date(qi.resetTime) : undefined; 560 return { usedPercent, resetDescription: resetTime ? formatReset(resetTime) : undefined }; 561 } 562 return null; 563 }; 564 565 // Quota groups from the reference snippet in CodexBar issue #129. 566 const windows: RateWindow[] = []; 567 568 const claudeOrGptOss = getQuotaInfo([ 569 "claude-sonnet-4-5", 570 "claude-sonnet-4-5-thinking", 571 "claude-opus-4-5-thinking", 572 "gpt-oss-120b-medium", 573 ]); 574 if (claudeOrGptOss) { 575 windows.push({ label: "Claude", usedPercent: claudeOrGptOss.usedPercent, resetDescription: claudeOrGptOss.resetDescription }); 576 } 577 578 const gemini3Pro = getQuotaInfo(["gemini-3-pro-high", "gemini-3-pro-low", "gemini-3-pro-preview"]); 579 if (gemini3Pro) { 580 windows.push({ label: "G3 Pro", usedPercent: gemini3Pro.usedPercent, resetDescription: gemini3Pro.resetDescription }); 581 } 582 583 const gemini3Flash = getQuotaInfo(["gemini-3-flash"]); 584 if (gemini3Flash) { 585 windows.push({ label: "G3 Flash", usedPercent: gemini3Flash.usedPercent, resetDescription: gemini3Flash.resetDescription }); 586 } 587 588 if (windows.length === 0) { 589 return { provider: "antigravity", displayName: "Antigravity", windows: [], error: "No quota data" }; 590 } 591 592 return { provider: "antigravity", displayName: "Antigravity", windows }; 593 } catch (e) { 594 return { provider: "antigravity", displayName: "Antigravity", windows: [], error: String(e) }; 595 } 596 } 597 598 // ============================================================================ 599 // Codex (OpenAI) Usage 600 // ============================================================================ 601 602 interface CodexCredential { 603 accessToken: string; 604 accountId?: string; 605 source: string; // Label identifying credential origin (e.g., "pi", "pi:second", ".codex:work") 606 } 607 608 /** 609 * Read all Codex tokens from ~/.pi/agent/auth.json 610 * Finds all keys starting with "openai-codex" (e.g., openai-codex, openai-codex-second, etc.) 611 */ 612 function readAllPiCodexAuths(): Array<{ accessToken: string; accountId?: string; source: string }> { 613 const piAuthPath = path.join(os.homedir(), ".pi", "agent", "auth.json"); 614 const results: Array<{ accessToken: string; accountId?: string; source: string }> = []; 615 616 try { 617 if (!fs.existsSync(piAuthPath)) return results; 618 const data = JSON.parse(fs.readFileSync(piAuthPath, "utf-8")); 619 620 // Find all keys that start with "openai-codex" 621 const codexKeys = Object.keys(data).filter(k => k.startsWith("openai-codex")).sort(); 622 623 for (const key of codexKeys) { 624 const source = data[key]; 625 if (!source) continue; 626 627 let accessToken: string | undefined; 628 let accountId: string | undefined; 629 630 // Pi auth pattern: .access 631 if (typeof source.access === "string") { 632 accessToken = source.access; 633 accountId = source.accountId; 634 } 635 // Fallback: codex schema 636 else if (source.tokens?.access_token) { 637 accessToken = source.tokens.access_token; 638 accountId = source.tokens.account_id; 639 } 640 641 if (accessToken) { 642 // Label with pi: prefix to distinguish from .codex/ files 643 const label = key === "openai-codex" ? "pi" : `pi:${key.replace("openai-codex-", "")}`; 644 results.push({ accessToken, accountId, source: label }); 645 } 646 } 647 } catch {} 648 649 return results; 650 } 651 652 /** 653 * Read Codex token from a ~/.codex/*auth*.json file 654 * Codex files use the pattern: data.tokens.access_token 655 */ 656 function readCodexAuthFile(filePath: string): { accessToken?: string; accountId?: string } { 657 try { 658 if (!fs.existsSync(filePath)) return {}; 659 const data = JSON.parse(fs.readFileSync(filePath, "utf-8")); 660 661 // Codex file pattern: .tokens.access_token 662 if (data.tokens?.access_token) { 663 return { accessToken: data.tokens.access_token, accountId: data.tokens.account_id }; 664 } 665 // Fallback: OPENAI_API_KEY 666 if (typeof data.OPENAI_API_KEY === "string" && data.OPENAI_API_KEY) { 667 return { accessToken: data.OPENAI_API_KEY }; 668 } 669 return {}; 670 } catch { 671 return {}; 672 } 673 } 674 675 /** 676 * Discover all unique Codex credentials from multiple sources: 677 * 1. ~/.pi/agent/auth.json (authoritative, all openai-codex* keys) 678 * 2. modelRegistry.authStorage (runtime auth, may be fresher) 679 * 3. ~/.codex/*auth*.json files 680 * Deduplicates by access_token first, then by usage stats when fetched 681 */ 682 async function discoverCodexCredentials(modelRegistry: any): Promise<CodexCredential[]> { 683 const credentials: CodexCredential[] = []; 684 const seenTokens = new Set<string>(); 685 686 // 1. Primary: from ~/.pi/agent/auth.json (authoritative source) 687 // Read ALL openai-codex* keys (e.g., openai-codex, openai-codex-second, etc.) 688 const piAuths = readAllPiCodexAuths(); 689 for (const piAuth of piAuths) { 690 if (!seenTokens.has(piAuth.accessToken)) { 691 credentials.push({ 692 accessToken: piAuth.accessToken, 693 accountId: piAuth.accountId, 694 source: piAuth.source, 695 }); 696 seenTokens.add(piAuth.accessToken); 697 } 698 } 699 700 // 2. Fallback: modelRegistry.authStorage (may have fresher token or be only source) 701 try { 702 const registryToken = await modelRegistry?.authStorage?.getApiKey?.("openai-codex"); 703 if (registryToken && !seenTokens.has(registryToken)) { 704 const cred = await modelRegistry?.authStorage?.get?.("openai-codex"); 705 const accountId = cred?.type === "oauth" ? cred.accountId : undefined; 706 credentials.push({ 707 accessToken: registryToken, 708 accountId, 709 source: "registry", 710 }); 711 seenTokens.add(registryToken); 712 } 713 } catch {} 714 715 // 3. Additional: scan ~/.codex/ for *auth*.json files 716 const codexHome = process.env.CODEX_HOME || path.join(os.homedir(), ".codex"); 717 try { 718 if (fs.existsSync(codexHome) && fs.statSync(codexHome).isDirectory()) { 719 const files = fs.readdirSync(codexHome); 720 // Only match files starting with "auth" (e.g., auth.json, auth-work.json) to avoid oauth.json etc. 721 const authFiles = files.filter(f => /^auth([_-].+)?\.json$/i.test(f)).sort(); 722 723 for (const authFile of authFiles) { 724 const authPath = path.join(codexHome, authFile); 725 const auth = readCodexAuthFile(authPath); 726 727 // Skip if no token or we've already seen this exact access_token 728 if (!auth.accessToken || seenTokens.has(auth.accessToken)) { 729 continue; 730 } 731 732 seenTokens.add(auth.accessToken); 733 // Label with .codex: prefix (e.g., "auth-xyz.json" -> ".codex:xyz") 734 const nameMatch = authFile.match(/auth[_-]?(.+)?\.json/i); 735 const suffix = nameMatch?.[1] || "auth"; 736 const label = `.codex:${suffix}`; 737 credentials.push({ accessToken: auth.accessToken, accountId: auth.accountId, source: label }); 738 } 739 } 740 } catch {} 741 742 return credentials; 743 } 744 745 async function fetchCodexUsageForCredential(cred: CodexCredential): Promise<UsageSnapshot> { 746 const displayName = `Codex (${cred.source})`; 747 748 try { 749 const controller = new AbortController(); 750 setTimeout(() => controller.abort(), 5000); 751 752 const headers: Record<string, string> = { 753 Authorization: `Bearer ${cred.accessToken}`, 754 "User-Agent": "CodexBar", 755 Accept: "application/json", 756 }; 757 758 if (cred.accountId) { 759 headers["ChatGPT-Account-Id"] = cred.accountId; 760 } 761 762 const res = await fetch("https://chatgpt.com/backend-api/wham/usage", { 763 method: "GET", 764 headers, 765 signal: controller.signal, 766 }); 767 768 if (res.status === 401 || res.status === 403) { 769 return { provider: "codex", displayName, windows: [], error: "Token expired" }; 770 } 771 772 if (!res.ok) { 773 return { provider: "codex", displayName, windows: [], error: `HTTP ${res.status}` }; 774 } 775 776 const data = await res.json() as any; 777 const windows: RateWindow[] = []; 778 779 // Primary window (usually 3-hour) 780 if (data.rate_limit?.primary_window) { 781 const pw = data.rate_limit.primary_window; 782 const resetDate = pw.reset_at ? new Date(pw.reset_at * 1000) : undefined; 783 const windowHours = Math.round((pw.limit_window_seconds || 10800) / 3600); 784 const usedPercent = typeof pw.used_percent === "number" ? pw.used_percent : Number(pw.used_percent) || 0; 785 windows.push({ 786 label: `${windowHours}h`, 787 usedPercent, 788 resetDescription: resetDate ? formatReset(resetDate) : undefined, 789 resetsAt: resetDate, 790 }); 791 } 792 793 // Secondary window (usually weekly) 794 if (data.rate_limit?.secondary_window) { 795 const sw = data.rate_limit.secondary_window; 796 const resetDate = sw.reset_at ? new Date(sw.reset_at * 1000) : undefined; 797 const windowHours = Math.round((sw.limit_window_seconds || 86400) / 3600); 798 const label = windowHours >= 24 ? "Week" : `${windowHours}h`; 799 const usedPercent = typeof sw.used_percent === "number" ? sw.used_percent : Number(sw.used_percent) || 0; 800 windows.push({ 801 label, 802 usedPercent, 803 resetDescription: resetDate ? formatReset(resetDate) : undefined, 804 resetsAt: resetDate, 805 }); 806 } 807 808 // Credits info 809 let plan = data.plan_type; 810 if (data.credits?.balance !== undefined && data.credits.balance !== null) { 811 const balance = typeof data.credits.balance === 'number' 812 ? data.credits.balance 813 : parseFloat(data.credits.balance) || 0; 814 plan = plan ? `${plan} ($${balance.toFixed(2)})` : `$${balance.toFixed(2)}`; 815 } 816 817 return { provider: "codex", displayName, windows, plan }; 818 } catch (e) { 819 return { provider: "codex", displayName, windows: [], error: String(e) }; 820 } 821 } 822 823 /** 824 * Generate a fingerprint from usage stats for deduplication. 825 * Two credentials accessing the same workspace will have identical stats. 826 * Uses absolute timestamps (not relative formatReset strings) for stability. 827 */ 828 function usageFingerprint(snapshot: UsageSnapshot): string | null { 829 if (snapshot.error || snapshot.windows.length === 0) { 830 return null; // Can't fingerprint errors or empty results 831 } 832 // Create a strict fingerprint from all window data using stable values 833 const parts = snapshot.windows.map(w => { 834 const pct = Number.isFinite(w.usedPercent) ? w.usedPercent.toFixed(2) : "NaN"; 835 const resetTs = w.resetsAt ? w.resetsAt.getTime() : ""; 836 return `${w.label}:${pct}:${resetTs}`; 837 }); 838 return parts.sort().join("|"); 839 } 840 841 async function fetchAllCodexUsages(modelRegistry: any): Promise<UsageSnapshot[]> { 842 const credentials = await discoverCodexCredentials(modelRegistry); 843 844 if (credentials.length === 0) { 845 return [{ provider: "codex", displayName: "Codex", windows: [], error: "No credentials" }]; 846 } 847 848 // Fetch usage for all credentials in parallel 849 const results = await Promise.all( 850 credentials.map(cred => fetchCodexUsageForCredential(cred)) 851 ); 852 853 // Deduplicate by usage stats - if two credentials return identical stats, 854 // they access the same workspace and we only show the first one 855 const seenFingerprints = new Set<string>(); 856 const deduplicated: UsageSnapshot[] = []; 857 858 for (const result of results) { 859 const fingerprint = usageFingerprint(result); 860 if (fingerprint === null) { 861 // Keep errors/empty results (they might be transient) 862 deduplicated.push(result); 863 } else if (!seenFingerprints.has(fingerprint)) { 864 seenFingerprints.add(fingerprint); 865 deduplicated.push(result); 866 } 867 // Skip if fingerprint already seen (duplicate workspace) 868 } 869 870 return deduplicated; 871 } 872 873 // ============================================================================ 874 // Kiro (AWS) 875 // ============================================================================ 876 877 function stripAnsi(text: string): string { 878 return text.replace(/\x1B\[[0-9;?]*[A-Za-z]|\x1B\].*?\x07/g, ""); 879 } 880 881 function whichSync(cmd: string): string | null { 882 try { 883 return execSync(`which ${cmd}`, { encoding: "utf-8" }).trim(); 884 } catch { 885 return null; 886 } 887 } 888 889 async function fetchKiroUsage(): Promise<UsageSnapshot> { 890 const kiroBinary = whichSync("kiro-cli"); 891 if (!kiroBinary) { 892 return { provider: "kiro", displayName: "Kiro", windows: [], error: "kiro-cli not found" }; 893 } 894 895 try { 896 // Check if logged in 897 try { 898 execSync("kiro-cli whoami", { encoding: "utf-8", timeout: 5000 }); 899 } catch { 900 return { provider: "kiro", displayName: "Kiro", windows: [], error: "Not logged in" }; 901 } 902 903 // Get usage 904 const output = execSync("kiro-cli chat --no-interactive /usage", { 905 encoding: "utf-8", 906 timeout: 10000, 907 env: { ...process.env, TERM: "xterm-256color" } 908 }); 909 910 const stripped = stripAnsi(output); 911 const windows: RateWindow[] = []; 912 913 // Parse plan name from "| KIRO FREE" or similar 914 let planName = "Kiro"; 915 const planMatch = stripped.match(/\|\s*(KIRO\s+\w+)/i); 916 if (planMatch) { 917 planName = planMatch[1].trim(); 918 } 919 920 // Parse credits percentage from "████...█ X%" 921 let creditsPercent = 0; 922 const percentMatch = stripped.match(/█+\s*(\d+)%/); 923 if (percentMatch) { 924 creditsPercent = parseInt(percentMatch[1], 10); 925 } 926 927 // Parse credits used/total from "(X.XX of Y covered in plan)" 928 let creditsUsed = 0; 929 let creditsTotal = 50; 930 const creditsMatch = stripped.match(/\((\d+\.?\d*)\s+of\s+(\d+)\s+covered/); 931 if (creditsMatch) { 932 creditsUsed = parseFloat(creditsMatch[1]); 933 creditsTotal = parseFloat(creditsMatch[2]); 934 if (!percentMatch && creditsTotal > 0) { 935 creditsPercent = (creditsUsed / creditsTotal) * 100; 936 } 937 } 938 939 // Parse reset date from "resets on 01/01" 940 let resetsAt: Date | undefined; 941 const resetMatch = stripped.match(/resets on (\d{2}\/\d{2})/); 942 if (resetMatch) { 943 const [month, day] = resetMatch[1].split("/").map(Number); 944 const now = new Date(); 945 const year = now.getFullYear(); 946 resetsAt = new Date(year, month - 1, day); 947 if (resetsAt < now) resetsAt.setFullYear(year + 1); 948 } 949 950 windows.push({ 951 label: "Credits", 952 usedPercent: creditsPercent, 953 resetDescription: resetsAt ? formatReset(resetsAt) : undefined, 954 }); 955 956 // Parse bonus credits 957 const bonusMatch = stripped.match(/Bonus credits:\s*(\d+\.?\d*)\/(\d+)/); 958 if (bonusMatch) { 959 const bonusUsed = parseFloat(bonusMatch[1]); 960 const bonusTotal = parseFloat(bonusMatch[2]); 961 const bonusPercent = bonusTotal > 0 ? (bonusUsed / bonusTotal) * 100 : 0; 962 const expiryMatch = stripped.match(/expires in (\d+) days?/); 963 windows.push({ 964 label: "Bonus", 965 usedPercent: bonusPercent, 966 resetDescription: expiryMatch ? `${expiryMatch[1]}d left` : undefined, 967 }); 968 } 969 970 return { provider: "kiro", displayName: "Kiro", windows, plan: planName }; 971 } catch (e) { 972 return { provider: "kiro", displayName: "Kiro", windows: [], error: String(e) }; 973 } 974 } 975 976 // ============================================================================ 977 // z.ai 978 // ============================================================================ 979 980 async function fetchZaiUsage(): Promise<UsageSnapshot> { 981 // Check for API key in environment or pi auth 982 let apiKey = process.env.Z_AI_API_KEY; 983 984 if (!apiKey) { 985 // Try pi auth storage 986 try { 987 const authPath = path.join(os.homedir(), ".pi", "agent", "auth.json"); 988 if (fs.existsSync(authPath)) { 989 const auth = JSON.parse(fs.readFileSync(authPath, "utf-8")); 990 apiKey = auth["z-ai"]?.access || auth["zai"]?.access; 991 } 992 } catch {} 993 } 994 995 if (!apiKey) { 996 return { provider: "zai", displayName: "z.ai", windows: [], error: "No API key" }; 997 } 998 999 try { 1000 const controller = new AbortController(); 1001 setTimeout(() => controller.abort(), 5000); 1002 1003 const res = await fetch("https://api.z.ai/api/monitor/usage/quota/limit", { 1004 method: "GET", 1005 headers: { 1006 Authorization: `Bearer ${apiKey}`, 1007 Accept: "application/json", 1008 }, 1009 signal: controller.signal, 1010 }); 1011 1012 if (!res.ok) { 1013 return { provider: "zai", displayName: "z.ai", windows: [], error: `HTTP ${res.status}` }; 1014 } 1015 1016 const data = await res.json() as any; 1017 if (!data.success || data.code !== 200) { 1018 return { provider: "zai", displayName: "z.ai", windows: [], error: data.msg || "API error" }; 1019 } 1020 1021 const windows: RateWindow[] = []; 1022 const limits = data.data?.limits || []; 1023 1024 for (const limit of limits) { 1025 const type = limit.type; 1026 const usage = limit.usage || 0; 1027 const remaining = limit.remaining || 0; 1028 const percent = limit.percentage || 0; 1029 const nextReset = limit.nextResetTime ? new Date(limit.nextResetTime) : undefined; 1030 1031 // Unit: 1=days, 3=hours, 5=minutes 1032 let windowLabel = "Limit"; 1033 if (limit.unit === 1) windowLabel = `${limit.number}d`; 1034 else if (limit.unit === 3) windowLabel = `${limit.number}h`; 1035 else if (limit.unit === 5) windowLabel = `${limit.number}m`; 1036 1037 if (type === "TOKENS_LIMIT") { 1038 windows.push({ 1039 label: `Tokens (${windowLabel})`, 1040 usedPercent: percent, 1041 resetDescription: nextReset ? formatReset(nextReset) : undefined, 1042 }); 1043 } else if (type === "TIME_LIMIT") { 1044 windows.push({ 1045 label: "Monthly", 1046 usedPercent: percent, 1047 resetDescription: nextReset ? formatReset(nextReset) : undefined, 1048 }); 1049 } 1050 } 1051 1052 const planName = data.data?.planName || data.data?.plan || undefined; 1053 return { provider: "zai", displayName: "z.ai", windows, plan: planName }; 1054 } catch (e) { 1055 return { provider: "zai", displayName: "z.ai", windows: [], error: String(e) }; 1056 } 1057 } 1058 1059 // ============================================================================ 1060 // Helpers 1061 // ============================================================================ 1062 1063 function formatReset(date: Date): string { 1064 const diffMs = date.getTime() - Date.now(); 1065 if (diffMs < 0) return "now"; 1066 1067 const diffMins = Math.floor(diffMs / 60000); 1068 if (diffMins < 60) return `${diffMins}m`; 1069 1070 const hours = Math.floor(diffMins / 60); 1071 const mins = diffMins % 60; 1072 if (hours < 24) return mins > 0 ? `${hours}h ${mins}m` : `${hours}h`; 1073 1074 const days = Math.floor(hours / 24); 1075 if (days < 7) return `${days}d ${hours % 24}h`; 1076 1077 return new Intl.DateTimeFormat("en-US", { month: "short", day: "numeric" }).format(date); 1078 } 1079 1080 function getStatusEmoji(status?: ProviderStatus): string { 1081 if (!status) return ""; 1082 switch (status.indicator) { 1083 case "none": return "✅"; 1084 case "minor": return "⚠️"; 1085 case "major": return "🟠"; 1086 case "critical": return "🔴"; 1087 case "maintenance": return "🔧"; 1088 default: return ""; 1089 } 1090 } 1091 1092 // ============================================================================ 1093 // UI Component 1094 // ============================================================================ 1095 1096 class UsageComponent { 1097 private usages: UsageSnapshot[] = []; 1098 private loading = true; 1099 private tui: { requestRender: () => void }; 1100 private theme: any; 1101 private onClose: () => void; 1102 private modelRegistry: any; 1103 1104 constructor(tui: { requestRender: () => void }, theme: any, onClose: () => void, modelRegistry: any) { 1105 this.tui = tui; 1106 this.theme = theme; 1107 this.onClose = onClose; 1108 this.modelRegistry = modelRegistry; 1109 this.load(); 1110 } 1111 1112 private async load() { 1113 const timeout = <T>(p: Promise<T>, ms: number, fallback: T) => 1114 Promise.race([p, new Promise<T>((r) => setTimeout(() => r(fallback), ms))]); 1115 1116 // Fetch usage and status in parallel 1117 const [claude, copilot, gemini, codexResults, antigravity, kiro, zai, claudeStatus, copilotStatus, geminiStatus, codexStatus] = await Promise.all([ 1118 timeout(fetchClaudeUsage(), 6000, { provider: "anthropic", displayName: "Claude", windows: [], error: "Timeout" }), 1119 timeout(fetchCopilotUsage(this.modelRegistry), 6000, { provider: "copilot", displayName: "Copilot", windows: [], error: "Timeout" }), 1120 timeout(fetchGeminiUsage(this.modelRegistry), 6000, { provider: "gemini", displayName: "Gemini", windows: [], error: "Timeout" }), 1121 timeout(fetchAllCodexUsages(this.modelRegistry), 6000, [{ provider: "codex", displayName: "Codex", windows: [], error: "Timeout" }]), 1122 timeout(fetchAntigravityUsage(this.modelRegistry), 6000, { provider: "antigravity", displayName: "Antigravity", windows: [], error: "Timeout" }), 1123 timeout(fetchKiroUsage(), 6000, { provider: "kiro", displayName: "Kiro", windows: [], error: "Timeout" }), 1124 timeout(fetchZaiUsage(), 6000, { provider: "zai", displayName: "z.ai", windows: [], error: "Timeout" }), 1125 timeout(fetchProviderStatus("anthropic"), 3000, { indicator: "unknown" as const }), 1126 timeout(fetchProviderStatus("copilot"), 3000, { indicator: "unknown" as const }), 1127 timeout(fetchGeminiStatus(), 3000, { indicator: "unknown" as const }), 1128 timeout(fetchProviderStatus("codex"), 3000, { indicator: "unknown" as const }), 1129 ]); 1130 1131 // Attach status to usage 1132 claude.status = claudeStatus; 1133 copilot.status = copilotStatus; 1134 gemini.status = geminiStatus; 1135 // Attach codex status to all codex accounts 1136 for (const codex of codexResults) { 1137 codex.status = codexStatus; 1138 } 1139 1140 // Filter out providers with no data and no error (not configured) 1141 const allUsages = [claude, copilot, gemini, ...codexResults, antigravity, kiro, zai]; 1142 this.usages = allUsages.filter(u => u.windows.length > 0 || u.error !== "No credentials" && u.error !== "kiro-cli not found" && u.error !== "No API key"); 1143 this.loading = false; 1144 this.tui.requestRender(); 1145 } 1146 1147 handleInput(_data: string): void { 1148 this.onClose(); 1149 } 1150 1151 invalidate(): void {} 1152 1153 render(width: number): string[] { 1154 const t = this.theme; 1155 const dim = (s: string) => t.fg("muted", s); 1156 const bold = (s: string) => t.bold(s); 1157 const accent = (s: string) => t.fg("accent", s); 1158 1159 // Box dimensions: total width includes borders 1160 const totalW = Math.min(55, width - 4); 1161 const innerW = totalW - 4; // subtract "│ " and " │" 1162 const hLine = "─".repeat(totalW - 2); // subtract corners 1163 1164 const box = (content: string) => { 1165 const contentW = visibleWidth(content); 1166 const pad = Math.max(0, innerW - contentW); 1167 return dim("│ ") + content + " ".repeat(pad) + dim(" │"); 1168 }; 1169 1170 const lines: string[] = []; 1171 lines.push(dim(`╭${hLine}╮`)); 1172 lines.push(box(bold(accent("Quota Usage")))); 1173 lines.push(dim(`├${hLine}┤`)); 1174 1175 if (this.loading) { 1176 lines.push(box("Loading...")); 1177 } else { 1178 for (const u of this.usages) { 1179 // Provider header with status emoji and plan 1180 const statusEmoji = getStatusEmoji(u.status); 1181 const planStr = u.plan ? dim(` (${u.plan})`) : ""; 1182 const statusStr = (statusEmoji && !u.error) ? ` ${statusEmoji}` : ""; 1183 lines.push(box(bold(u.displayName) + planStr + statusStr)); 1184 1185 // Show incident description if any 1186 if (u.status?.indicator && u.status.indicator !== "none" && u.status.indicator !== "unknown" && u.status.description) { 1187 const desc = u.status.description.length > 40 1188 ? u.status.description.substring(0, 37) + "..." 1189 : u.status.description; 1190 lines.push(box(t.fg("warning", ` ⚡ ${desc}`))); 1191 } 1192 1193 if (u.error) { 1194 lines.push(box(dim(` ${u.error}`))); 1195 } else if (u.windows.length === 0) { 1196 lines.push(box(dim(" No data"))); 1197 } else { 1198 for (const w of u.windows) { 1199 const used = Math.min(100, Math.max(0, w.usedPercent)); 1200 const barW = 12; 1201 const filled = Math.round((used / 100) * barW); 1202 const empty = barW - filled; 1203 const color = used >= 95 ? "error" : used >= 85 ? "warning" : used >= 70 ? "accent" : used >= 50 ? "muted" : "success"; 1204 const bar = t.fg(color, "█".repeat(filled)) + dim("░".repeat(empty)); 1205 const reset = w.resetDescription ? dim(` ⏱ ${w.resetDescription}`) : ""; 1206 lines.push(box(` ${w.label.padEnd(8)} ${bar} ${used.toFixed(0).padStart(3)}%${reset}`)); 1207 } 1208 } 1209 lines.push(box("")); 1210 } 1211 } 1212 1213 lines.push(dim(`├${hLine}┤`)); 1214 lines.push(box(dim("Press any key to close"))); 1215 lines.push(dim(`╰${hLine}╯`)); 1216 1217 return lines; 1218 } 1219 1220 dispose(): void {} 1221 } 1222 1223 // ============================================================================ 1224 // Hook 1225 // ============================================================================ 1226 1227 export default function (pi: ExtensionAPI) { 1228 pi.registerCommand("usage", { 1229 description: "Show AI provider usage statistics", 1230 handler: async (_args, ctx) => { 1231 if (!ctx.hasUI) { 1232 ctx.ui.notify("Usage requires interactive mode", "error"); 1233 return; 1234 } 1235 1236 const modelRegistry = ctx.modelRegistry; 1237 await ctx.ui.custom((tui, theme, _kb, done) => { 1238 return new UsageComponent(tui, theme, () => done(), modelRegistry); 1239 }); 1240 }, 1241 }); 1242 1243 pi.registerShortcut("alt+u", { 1244 description: "Show AI provider usage statistics", 1245 handler: async (ctx) => { 1246 if (!ctx.hasUI) { 1247 ctx.ui.notify("Usage requires interactive mode", "error"); 1248 return; 1249 } 1250 1251 const modelRegistry = ctx.modelRegistry; 1252 await ctx.ui.custom((tui, theme, _kb, done) => { 1253 return new UsageComponent(tui, theme, () => done(), modelRegistry); 1254 }); 1255 }, 1256 }); 1257 }