main
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
11import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
12import { visibleWidth } from "@mariozechner/pi-tui";
13import * as fs from "node:fs";
14import * as path from "node:path";
15import * as os from "node:os";
16import { execSync } from "node:child_process";
17
18// ============================================================================
19// Types
20// ============================================================================
21
22interface RateWindow {
23 label: string;
24 usedPercent: number;
25 resetDescription?: string;
26 resetsAt?: Date;
27}
28
29interface ProviderStatus {
30 indicator: "none" | "minor" | "major" | "critical" | "maintenance" | "unknown";
31 description?: string;
32}
33
34interface 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
47const 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
53async 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
77async 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
126function 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
154async 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
213function 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
228async 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
320async 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
400type AntigravityAuth = {
401 accessToken: string;
402 refreshToken?: string;
403 expiresAt?: number;
404 projectId?: string;
405};
406
407function 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
431async 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
458async 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
493async 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-6",
570 "claude-sonnet-4-6-thinking",
571 "claude-sonnet-4-5",
572 "claude-sonnet-4-5-thinking",
573 "claude-opus-4-5-thinking",
574 "gpt-oss-120b-medium",
575 ]);
576 if (claudeOrGptOss) {
577 windows.push({ label: "Claude", usedPercent: claudeOrGptOss.usedPercent, resetDescription: claudeOrGptOss.resetDescription });
578 }
579
580 const gemini3Pro = getQuotaInfo(["gemini-3-pro-high", "gemini-3-pro-low", "gemini-3-pro-preview"]);
581 if (gemini3Pro) {
582 windows.push({ label: "G3 Pro", usedPercent: gemini3Pro.usedPercent, resetDescription: gemini3Pro.resetDescription });
583 }
584
585 const gemini3Flash = getQuotaInfo(["gemini-3-flash"]);
586 if (gemini3Flash) {
587 windows.push({ label: "G3 Flash", usedPercent: gemini3Flash.usedPercent, resetDescription: gemini3Flash.resetDescription });
588 }
589
590 if (windows.length === 0) {
591 return { provider: "antigravity", displayName: "Antigravity", windows: [], error: "No quota data" };
592 }
593
594 return { provider: "antigravity", displayName: "Antigravity", windows };
595 } catch (e) {
596 return { provider: "antigravity", displayName: "Antigravity", windows: [], error: String(e) };
597 }
598}
599
600// ============================================================================
601// Codex (OpenAI) Usage
602// ============================================================================
603
604async function fetchCodexUsage(modelRegistry: any): Promise<UsageSnapshot> {
605 // Try to get token from pi's auth storage first
606 let accessToken: string | undefined;
607 let accountId: string | undefined;
608
609 try {
610 // Try openai-codex provider first (pi's built-in)
611 accessToken = await modelRegistry?.authStorage?.getApiKey?.("openai-codex");
612
613 // Get account ID if available from OAuth credentials
614 const cred = modelRegistry?.authStorage?.get?.("openai-codex");
615 if (cred?.type === "oauth") {
616 accountId = (cred as any).accountId;
617 }
618 } catch {}
619
620 // Fallback to ~/.codex/auth.json if not in pi's auth
621 if (!accessToken) {
622 const codexHome = process.env.CODEX_HOME || path.join(os.homedir(), ".codex");
623 const authPath = path.join(codexHome, "auth.json");
624
625 try {
626 if (fs.existsSync(authPath)) {
627 const data = JSON.parse(fs.readFileSync(authPath, "utf-8"));
628
629 if (data.OPENAI_API_KEY) {
630 accessToken = data.OPENAI_API_KEY;
631 } else if (data.tokens?.access_token) {
632 accessToken = data.tokens.access_token;
633 accountId = data.tokens.account_id;
634 }
635 }
636 } catch {}
637 }
638
639 if (!accessToken) {
640 return { provider: "codex", displayName: "Codex", windows: [], error: "No credentials" };
641 }
642
643 try {
644 const controller = new AbortController();
645 setTimeout(() => controller.abort(), 5000);
646
647 const headers: Record<string, string> = {
648 Authorization: `Bearer ${accessToken}`,
649 "User-Agent": "CodexBar",
650 Accept: "application/json",
651 };
652
653 if (accountId) {
654 headers["ChatGPT-Account-Id"] = accountId;
655 }
656
657 const res = await fetch("https://chatgpt.com/backend-api/wham/usage", {
658 method: "GET",
659 headers,
660 signal: controller.signal,
661 });
662
663 if (res.status === 401 || res.status === 403) {
664 return { provider: "codex", displayName: "Codex", windows: [], error: "Token expired" };
665 }
666
667 if (!res.ok) {
668 return { provider: "codex", displayName: "Codex", windows: [], error: `HTTP ${res.status}` };
669 }
670
671 const data = await res.json() as any;
672 const windows: RateWindow[] = [];
673
674 // Primary window (usually 3-hour)
675 if (data.rate_limit?.primary_window) {
676 const pw = data.rate_limit.primary_window;
677 const resetDate = pw.reset_at ? new Date(pw.reset_at * 1000) : undefined;
678 const windowHours = Math.round((pw.limit_window_seconds || 10800) / 3600);
679 windows.push({
680 label: `${windowHours}h`,
681 usedPercent: pw.used_percent || 0,
682 resetDescription: resetDate ? formatReset(resetDate) : undefined,
683 });
684 }
685
686 // Secondary window (usually daily)
687 if (data.rate_limit?.secondary_window) {
688 const sw = data.rate_limit.secondary_window;
689 const resetDate = sw.reset_at ? new Date(sw.reset_at * 1000) : undefined;
690 const windowHours = Math.round((sw.limit_window_seconds || 86400) / 3600);
691 const label = windowHours >= 24 ? "Day" : `${windowHours}h`;
692 windows.push({
693 label,
694 usedPercent: sw.used_percent || 0,
695 resetDescription: resetDate ? formatReset(resetDate) : undefined,
696 });
697 }
698
699 // Credits info
700 let plan = data.plan_type;
701 if (data.credits?.balance !== undefined && data.credits.balance !== null) {
702 const balance = typeof data.credits.balance === 'number'
703 ? data.credits.balance
704 : parseFloat(data.credits.balance) || 0;
705 plan = plan ? `${plan} ($${balance.toFixed(2)})` : `$${balance.toFixed(2)}`;
706 }
707
708 return { provider: "codex", displayName: "Codex", windows, plan };
709 } catch (e) {
710 return { provider: "codex", displayName: "Codex", windows: [], error: String(e) };
711 }
712}
713
714// ============================================================================
715// Kiro (AWS)
716// ============================================================================
717
718function stripAnsi(text: string): string {
719 return text.replace(/\x1B\[[0-9;?]*[A-Za-z]|\x1B\].*?\x07/g, "");
720}
721
722function whichSync(cmd: string): string | null {
723 try {
724 return execSync(`which ${cmd}`, { encoding: "utf-8" }).trim();
725 } catch {
726 return null;
727 }
728}
729
730async function fetchKiroUsage(): Promise<UsageSnapshot> {
731 const kiroBinary = whichSync("kiro-cli");
732 if (!kiroBinary) {
733 return { provider: "kiro", displayName: "Kiro", windows: [], error: "kiro-cli not found" };
734 }
735
736 try {
737 // Check if logged in
738 try {
739 execSync("kiro-cli whoami", { encoding: "utf-8", timeout: 5000 });
740 } catch {
741 return { provider: "kiro", displayName: "Kiro", windows: [], error: "Not logged in" };
742 }
743
744 // Get usage
745 const output = execSync("kiro-cli chat --no-interactive /usage", {
746 encoding: "utf-8",
747 timeout: 10000,
748 env: { ...process.env, TERM: "xterm-256color" }
749 });
750
751 const stripped = stripAnsi(output);
752 const windows: RateWindow[] = [];
753
754 // Parse plan name from "| KIRO FREE" or similar
755 let planName = "Kiro";
756 const planMatch = stripped.match(/\|\s*(KIRO\s+\w+)/i);
757 if (planMatch) {
758 planName = planMatch[1].trim();
759 }
760
761 // Parse credits percentage from "████...█ X%"
762 let creditsPercent = 0;
763 const percentMatch = stripped.match(/█+\s*(\d+)%/);
764 if (percentMatch) {
765 creditsPercent = parseInt(percentMatch[1], 10);
766 }
767
768 // Parse credits used/total from "(X.XX of Y covered in plan)"
769 let creditsUsed = 0;
770 let creditsTotal = 50;
771 const creditsMatch = stripped.match(/\((\d+\.?\d*)\s+of\s+(\d+)\s+covered/);
772 if (creditsMatch) {
773 creditsUsed = parseFloat(creditsMatch[1]);
774 creditsTotal = parseFloat(creditsMatch[2]);
775 if (!percentMatch && creditsTotal > 0) {
776 creditsPercent = (creditsUsed / creditsTotal) * 100;
777 }
778 }
779
780 // Parse reset date from "resets on 01/01"
781 let resetsAt: Date | undefined;
782 const resetMatch = stripped.match(/resets on (\d{2}\/\d{2})/);
783 if (resetMatch) {
784 const [month, day] = resetMatch[1].split("/").map(Number);
785 const now = new Date();
786 const year = now.getFullYear();
787 resetsAt = new Date(year, month - 1, day);
788 if (resetsAt < now) resetsAt.setFullYear(year + 1);
789 }
790
791 windows.push({
792 label: "Credits",
793 usedPercent: creditsPercent,
794 resetDescription: resetsAt ? formatReset(resetsAt) : undefined,
795 });
796
797 // Parse bonus credits
798 const bonusMatch = stripped.match(/Bonus credits:\s*(\d+\.?\d*)\/(\d+)/);
799 if (bonusMatch) {
800 const bonusUsed = parseFloat(bonusMatch[1]);
801 const bonusTotal = parseFloat(bonusMatch[2]);
802 const bonusPercent = bonusTotal > 0 ? (bonusUsed / bonusTotal) * 100 : 0;
803 const expiryMatch = stripped.match(/expires in (\d+) days?/);
804 windows.push({
805 label: "Bonus",
806 usedPercent: bonusPercent,
807 resetDescription: expiryMatch ? `${expiryMatch[1]}d left` : undefined,
808 });
809 }
810
811 return { provider: "kiro", displayName: "Kiro", windows, plan: planName };
812 } catch (e) {
813 return { provider: "kiro", displayName: "Kiro", windows: [], error: String(e) };
814 }
815}
816
817// ============================================================================
818// z.ai
819// ============================================================================
820
821async function fetchZaiUsage(): Promise<UsageSnapshot> {
822 // Check for API key in environment or pi auth
823 let apiKey = process.env.Z_AI_API_KEY;
824
825 if (!apiKey) {
826 // Try pi auth storage
827 try {
828 const authPath = path.join(os.homedir(), ".pi", "agent", "auth.json");
829 if (fs.existsSync(authPath)) {
830 const auth = JSON.parse(fs.readFileSync(authPath, "utf-8"));
831 apiKey = auth["z-ai"]?.access || auth["zai"]?.access;
832 }
833 } catch {}
834 }
835
836 if (!apiKey) {
837 return { provider: "zai", displayName: "z.ai", windows: [], error: "No API key" };
838 }
839
840 try {
841 const controller = new AbortController();
842 setTimeout(() => controller.abort(), 5000);
843
844 const res = await fetch("https://api.z.ai/api/monitor/usage/quota/limit", {
845 method: "GET",
846 headers: {
847 Authorization: `Bearer ${apiKey}`,
848 Accept: "application/json",
849 },
850 signal: controller.signal,
851 });
852
853 if (!res.ok) {
854 return { provider: "zai", displayName: "z.ai", windows: [], error: `HTTP ${res.status}` };
855 }
856
857 const data = await res.json() as any;
858 if (!data.success || data.code !== 200) {
859 return { provider: "zai", displayName: "z.ai", windows: [], error: data.msg || "API error" };
860 }
861
862 const windows: RateWindow[] = [];
863 const limits = data.data?.limits || [];
864
865 for (const limit of limits) {
866 const type = limit.type;
867 const usage = limit.usage || 0;
868 const remaining = limit.remaining || 0;
869 const percent = limit.percentage || 0;
870 const nextReset = limit.nextResetTime ? new Date(limit.nextResetTime) : undefined;
871
872 // Unit: 1=days, 3=hours, 5=minutes
873 let windowLabel = "Limit";
874 if (limit.unit === 1) windowLabel = `${limit.number}d`;
875 else if (limit.unit === 3) windowLabel = `${limit.number}h`;
876 else if (limit.unit === 5) windowLabel = `${limit.number}m`;
877
878 if (type === "TOKENS_LIMIT") {
879 windows.push({
880 label: `Tokens (${windowLabel})`,
881 usedPercent: percent,
882 resetDescription: nextReset ? formatReset(nextReset) : undefined,
883 });
884 } else if (type === "TIME_LIMIT") {
885 windows.push({
886 label: "Monthly",
887 usedPercent: percent,
888 resetDescription: nextReset ? formatReset(nextReset) : undefined,
889 });
890 }
891 }
892
893 const planName = data.data?.planName || data.data?.plan || undefined;
894 return { provider: "zai", displayName: "z.ai", windows, plan: planName };
895 } catch (e) {
896 return { provider: "zai", displayName: "z.ai", windows: [], error: String(e) };
897 }
898}
899
900// ============================================================================
901// Helpers
902// ============================================================================
903
904function formatReset(date: Date): string {
905 const diffMs = date.getTime() - Date.now();
906 if (diffMs < 0) return "now";
907
908 const diffMins = Math.floor(diffMs / 60000);
909 if (diffMins < 60) return `${diffMins}m`;
910
911 const hours = Math.floor(diffMins / 60);
912 const mins = diffMins % 60;
913 if (hours < 24) return mins > 0 ? `${hours}h ${mins}m` : `${hours}h`;
914
915 const days = Math.floor(hours / 24);
916 if (days < 7) return `${days}d ${hours % 24}h`;
917
918 return new Intl.DateTimeFormat("en-US", { month: "short", day: "numeric" }).format(date);
919}
920
921function getStatusEmoji(status?: ProviderStatus): string {
922 if (!status) return "";
923 switch (status.indicator) {
924 case "none": return "✅";
925 case "minor": return "⚠️";
926 case "major": return "🟠";
927 case "critical": return "🔴";
928 case "maintenance": return "🔧";
929 default: return "";
930 }
931}
932
933// ============================================================================
934// UI Component
935// ============================================================================
936
937class UsageComponent {
938 private usages: UsageSnapshot[] = [];
939 private loading = true;
940 private tui: { requestRender: () => void };
941 private theme: any;
942 private onClose: () => void;
943 private modelRegistry: any;
944
945 constructor(tui: { requestRender: () => void }, theme: any, onClose: () => void, modelRegistry: any) {
946 this.tui = tui;
947 this.theme = theme;
948 this.onClose = onClose;
949 this.modelRegistry = modelRegistry;
950 this.load();
951 }
952
953 private async load() {
954 const timeout = <T>(p: Promise<T>, ms: number, fallback: T) =>
955 Promise.race([p, new Promise<T>((r) => setTimeout(() => r(fallback), ms))]);
956
957 // Fetch usage and status in parallel
958 const [claude, copilot, gemini, codex, antigravity, kiro, zai, claudeStatus, copilotStatus, geminiStatus, codexStatus] = await Promise.all([
959 timeout(fetchClaudeUsage(), 6000, { provider: "anthropic", displayName: "Claude", windows: [], error: "Timeout" }),
960 timeout(fetchCopilotUsage(this.modelRegistry), 6000, { provider: "copilot", displayName: "Copilot", windows: [], error: "Timeout" }),
961 timeout(fetchGeminiUsage(this.modelRegistry), 6000, { provider: "gemini", displayName: "Gemini", windows: [], error: "Timeout" }),
962 timeout(fetchCodexUsage(this.modelRegistry), 6000, { provider: "codex", displayName: "Codex", windows: [], error: "Timeout" }),
963 timeout(fetchAntigravityUsage(this.modelRegistry), 6000, { provider: "antigravity", displayName: "Antigravity", windows: [], error: "Timeout" }),
964 timeout(fetchKiroUsage(), 6000, { provider: "kiro", displayName: "Kiro", windows: [], error: "Timeout" }),
965 timeout(fetchZaiUsage(), 6000, { provider: "zai", displayName: "z.ai", windows: [], error: "Timeout" }),
966 timeout(fetchProviderStatus("anthropic"), 3000, { indicator: "unknown" as const }),
967 timeout(fetchProviderStatus("copilot"), 3000, { indicator: "unknown" as const }),
968 timeout(fetchGeminiStatus(), 3000, { indicator: "unknown" as const }),
969 timeout(fetchProviderStatus("codex"), 3000, { indicator: "unknown" as const }),
970 ]);
971
972 // Attach status to usage
973 claude.status = claudeStatus;
974 copilot.status = copilotStatus;
975 gemini.status = geminiStatus;
976 codex.status = codexStatus;
977
978 // Filter out providers with no data and no error (not configured)
979 const allUsages = [claude, copilot, gemini, codex, antigravity, kiro, zai];
980 this.usages = allUsages.filter(u => u.windows.length > 0 || u.error !== "No credentials" && u.error !== "kiro-cli not found" && u.error !== "No API key");
981 this.loading = false;
982 this.tui.requestRender();
983 }
984
985 handleInput(_data: string): void {
986 this.onClose();
987 }
988
989 invalidate(): void {}
990
991 render(width: number): string[] {
992 const t = this.theme;
993 const dim = (s: string) => t.fg("muted", s);
994 const bold = (s: string) => t.bold(s);
995 const accent = (s: string) => t.fg("accent", s);
996
997 // Box dimensions: total width includes borders
998 const totalW = Math.min(55, width - 4);
999 const innerW = totalW - 4; // subtract "│ " and " │"
1000 const hLine = "─".repeat(totalW - 2); // subtract corners
1001
1002 const box = (content: string) => {
1003 const contentW = visibleWidth(content);
1004 const pad = Math.max(0, innerW - contentW);
1005 return dim("│ ") + content + " ".repeat(pad) + dim(" │");
1006 };
1007
1008 const lines: string[] = [];
1009 lines.push(dim(`╭${hLine}╮`));
1010 lines.push(box(bold(accent("AI Usage"))));
1011 lines.push(dim(`├${hLine}┤`));
1012
1013 if (this.loading) {
1014 lines.push(box("Loading..."));
1015 } else {
1016 for (const u of this.usages) {
1017 // Provider header with status emoji and plan
1018 const statusEmoji = getStatusEmoji(u.status);
1019 const planStr = u.plan ? dim(` (${u.plan})`) : "";
1020 const statusStr = statusEmoji ? ` ${statusEmoji}` : "";
1021 lines.push(box(bold(u.displayName) + planStr + statusStr));
1022
1023 // Show incident description if any
1024 if (u.status?.indicator && u.status.indicator !== "none" && u.status.indicator !== "unknown" && u.status.description) {
1025 const desc = u.status.description.length > 40
1026 ? u.status.description.substring(0, 37) + "..."
1027 : u.status.description;
1028 lines.push(box(t.fg("warning", ` ⚡ ${desc}`)));
1029 }
1030
1031 if (u.error) {
1032 lines.push(box(dim(` ${u.error}`)));
1033 } else if (u.windows.length === 0) {
1034 lines.push(box(dim(" No data")));
1035 } else {
1036 for (const w of u.windows) {
1037 const remaining = Math.max(0, 100 - w.usedPercent);
1038 const barW = 12;
1039 const filled = Math.min(barW, Math.round((w.usedPercent / 100) * barW));
1040 const empty = barW - filled;
1041 const color = remaining <= 10 ? "error" : remaining <= 30 ? "warning" : "success";
1042 const bar = t.fg(color, "█".repeat(filled)) + dim("░".repeat(empty));
1043 const reset = w.resetDescription ? dim(` ⏱ ${w.resetDescription}`) : "";
1044 lines.push(box(` ${w.label.padEnd(7)} ${bar} ${remaining.toFixed(0).padStart(3)}%${reset}`));
1045 }
1046 }
1047 lines.push(box(""));
1048 }
1049 }
1050
1051 lines.push(dim(`├${hLine}┤`));
1052 lines.push(box(dim("Press any key to close")));
1053 lines.push(dim(`╰${hLine}╯`));
1054
1055 return lines;
1056 }
1057
1058 dispose(): void {}
1059}
1060
1061// ============================================================================
1062// Hook
1063// ============================================================================
1064
1065export default function (pi: ExtensionAPI) {
1066 pi.registerCommand("usage", {
1067 description: "Show AI provider usage statistics",
1068 handler: async (_args, ctx) => {
1069 if (!ctx.hasUI) {
1070 ctx.ui.notify("Usage requires interactive mode", "error");
1071 return;
1072 }
1073
1074 const modelRegistry = ctx.modelRegistry;
1075 await ctx.ui.custom((tui, theme, _kb, done) => {
1076 return new UsageComponent(tui, theme, () => done(), modelRegistry);
1077 });
1078 },
1079 });
1080}