added pi

pi/.pi/agent/extensions/usage-bars/index.ts

@@ -0,0 +1,576 @@
+/**
+ * Usage Extension - Minimal API usage indicator for pi
+ *
+ * Shows Codex (OpenAI), Anthropic (Claude), Z.AI, and optionally
+ * Google Gemini CLI / Antigravity usage as color-coded percentages
+ * in the footer status bar.
+ */
+
+import { DynamicBorder, type ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import {
+  Container,
+  Input,
+  Spacer,
+  Text,
+  getEditorKeybindings,
+  type Focusable,
+} from "@mariozechner/pi-tui";
+import {
+  canShowForProvider,
+  clampPercent,
+  colorForPercent,
+  detectProvider,
+  fetchAllUsages,
+  fetchClaudeUsage,
+  fetchCodexUsage,
+  fetchGoogleUsage,
+  fetchZaiUsage,
+  providerToOAuthProviderId,
+  readAuth,
+  readUsageCache,
+  resolveUsageEndpoints,
+  writeUsageCache,
+  type ProviderKey,
+  type UsageByProvider,
+  type UsageData,
+} from "./core";
+
+const CACHE_TTL_MS = 15 * 60 * 1000;           // reuse cached data for 15 min
+const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000; // 1 hour back-off after 429
+const STATUS_KEY = "usage-bars";
+
+const PROVIDER_LABELS: Record<ProviderKey, string> = {
+  codex: "Codex",
+  claude: "Claude",
+  zai: "Z.AI",
+  gemini: "Gemini",
+  antigravity: "Antigravity",
+};
+
+interface SubscriptionItem {
+  name: string;
+  provider: ProviderKey;
+  data: UsageData | null;
+  isActive: boolean;
+}
+
+class UsageSelectorComponent extends Container implements Focusable {
+  private searchInput: Input;
+  private listContainer: Container;
+  private hintText: Text;
+  private tui: any;
+  private theme: any;
+  private onCancelCallback: () => void;
+  private allItems: SubscriptionItem[] = [];
+  private filteredItems: SubscriptionItem[] = [];
+  private selectedIndex = 0;
+  private loading = true;
+  private activeProvider: ProviderKey | null;
+  private fetchAllFn: () => Promise<UsageByProvider>;
+  private _focused = false;
+
+  get focused(): boolean {
+    return this._focused;
+  }
+
+  set focused(value: boolean) {
+    this._focused = value;
+    this.searchInput.focused = value;
+  }
+
+  constructor(
+    tui: any,
+    theme: any,
+    activeProvider: ProviderKey | null,
+    fetchAll: () => Promise<UsageByProvider>,
+    onCancel: () => void,
+  ) {
+    super();
+    this.tui = tui;
+    this.theme = theme;
+    this.activeProvider = activeProvider;
+    this.fetchAllFn = fetchAll;
+    this.onCancelCallback = onCancel;
+
+    this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
+    this.addChild(new Spacer(1));
+
+    this.hintText = new Text(theme.fg("dim", "Fetching usage from all providers…"), 0, 0);
+    this.addChild(this.hintText);
+    this.addChild(new Spacer(1));
+
+    this.searchInput = new Input();
+    this.addChild(this.searchInput);
+    this.addChild(new Spacer(1));
+
+    this.listContainer = new Container();
+    this.addChild(this.listContainer);
+    this.addChild(new Spacer(1));
+
+    this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
+
+    this.fetchAllFn()
+      .then((results) => {
+        this.loading = false;
+        this.buildItems(results);
+        this.updateList();
+        this.hintText.setText(
+          theme.fg("muted", "Only showing providers with credentials. ") +
+            theme.fg("dim", "✓ = active provider"),
+        );
+        this.tui.requestRender();
+      })
+      .catch(() => {
+        this.loading = false;
+        this.hintText.setText(theme.fg("error", "Failed to fetch usage data"));
+        this.tui.requestRender();
+      });
+
+    this.updateList();
+  }
+
+  private buildItems(results: UsageByProvider) {
+    const providers: Array<{ key: ProviderKey; name: string }> = [
+      { key: "codex", name: "Codex" },
+      { key: "claude", name: "Claude" },
+      { key: "zai", name: "Z.AI" },
+      { key: "gemini", name: "Gemini" },
+      { key: "antigravity", name: "Antigravity" },
+    ];
+
+    this.allItems = [];
+    for (const p of providers) {
+      if (results[p.key] !== null) {
+        this.allItems.push({
+          name: p.name,
+          provider: p.key,
+          data: results[p.key],
+          isActive: this.activeProvider === p.key,
+        });
+      }
+    }
+
+    this.filteredItems = this.allItems;
+    this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
+  }
+
+  private filterItems(query: string) {
+    if (!query) {
+      this.filteredItems = this.allItems;
+    } else {
+      const q = query.toLowerCase();
+      this.filteredItems = this.allItems.filter(
+        (item) => item.name.toLowerCase().includes(q) || item.provider.toLowerCase().includes(q),
+      );
+    }
+    this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
+  }
+
+  private renderBar(pct: number, width = 16): string {
+    const value = clampPercent(pct);
+    const filled = Math.round((value / 100) * width);
+    const color = colorForPercent(value);
+    const full = "█".repeat(Math.max(0, filled));
+    const empty = "░".repeat(Math.max(0, width - filled));
+    return this.theme.fg(color, full) + this.theme.fg("dim", empty);
+  }
+
+  private renderItem(item: SubscriptionItem, isSelected: boolean) {
+    const t = this.theme;
+    const pointer = isSelected ? t.fg("accent", "→ ") : "  ";
+    const activeBadge = item.isActive ? t.fg("success", " ✓") : "";
+    const name = isSelected ? t.fg("accent", t.bold(item.name)) : item.name;
+
+    this.listContainer.addChild(new Text(`${pointer}${name}${activeBadge}`, 0, 0));
+
+    const indent = "    ";
+
+    if (!item.data) {
+      this.listContainer.addChild(new Text(indent + t.fg("dim", "No credentials"), 0, 0));
+    } else if (item.data.error) {
+      this.listContainer.addChild(new Text(indent + t.fg("error", item.data.error), 0, 0));
+    } else {
+      const session = clampPercent(item.data.session);
+      const weekly = clampPercent(item.data.weekly);
+
+      const sessionReset = item.data.sessionResetsIn
+        ? t.fg("dim", `  resets in ${item.data.sessionResetsIn}`)
+        : "";
+      const weeklyReset = item.data.weeklyResetsIn
+        ? t.fg("dim", `  resets in ${item.data.weeklyResetsIn}`)
+        : "";
+
+      this.listContainer.addChild(
+        new Text(
+          indent +
+            t.fg("muted", "Session  ") +
+            this.renderBar(session) +
+            " " +
+            t.fg(colorForPercent(session), `${session}%`.padStart(4)) +
+            sessionReset,
+          0,
+          0,
+        ),
+      );
+
+      this.listContainer.addChild(
+        new Text(
+          indent +
+            t.fg("muted", "Weekly   ") +
+            this.renderBar(weekly) +
+            " " +
+            t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) +
+            weeklyReset,
+          0,
+          0,
+        ),
+      );
+
+      if (typeof item.data.extraSpend === "number" && typeof item.data.extraLimit === "number") {
+        this.listContainer.addChild(
+          new Text(
+            indent +
+              t.fg("muted", "Extra    ") +
+              t.fg("dim", `$${item.data.extraSpend.toFixed(2)} / $${item.data.extraLimit}`),
+            0,
+            0,
+          ),
+        );
+      }
+    }
+
+    this.listContainer.addChild(new Spacer(1));
+  }
+
+  private updateList() {
+    this.listContainer.clear();
+
+    if (this.loading) {
+      this.listContainer.addChild(new Text(this.theme.fg("muted", "  Loading…"), 0, 0));
+      return;
+    }
+
+    if (this.filteredItems.length === 0) {
+      this.listContainer.addChild(new Text(this.theme.fg("muted", "  No matching providers"), 0, 0));
+      return;
+    }
+
+    for (let i = 0; i < this.filteredItems.length; i++) {
+      this.renderItem(this.filteredItems[i]!, i === this.selectedIndex);
+    }
+  }
+
+  handleInput(keyData: string): void {
+    const kb = getEditorKeybindings();
+
+    if (kb.matches(keyData, "selectUp")) {
+      if (this.filteredItems.length === 0) return;
+      this.selectedIndex =
+        this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1;
+      this.updateList();
+      return;
+    }
+
+    if (kb.matches(keyData, "selectDown")) {
+      if (this.filteredItems.length === 0) return;
+      this.selectedIndex =
+        this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1;
+      this.updateList();
+      return;
+    }
+
+    if (kb.matches(keyData, "selectCancel") || kb.matches(keyData, "selectConfirm")) {
+      this.onCancelCallback();
+      return;
+    }
+
+    this.searchInput.handleInput(keyData);
+    this.filterItems(this.searchInput.getValue());
+    this.updateList();
+  }
+}
+
+interface UsageState extends UsageByProvider {
+  lastPoll: number;
+  activeProvider: ProviderKey | null;
+}
+
+export default function (pi: ExtensionAPI) {
+  const endpoints = resolveUsageEndpoints();
+  const state: UsageState = {
+    codex: null,
+    claude: null,
+    zai: null,
+    gemini: null,
+    antigravity: null,
+    lastPoll: 0,
+    activeProvider: null,
+  };
+
+  let pollInFlight: Promise<void> | null = null;
+  let pollQueued = false;
+  let ctx: any = null;
+
+  function renderPercent(theme: any, value: number): string {
+    const v = clampPercent(value);
+    return theme.fg(colorForPercent(v), `${v}%`);
+  }
+
+  function renderBar(theme: any, value: number): string {
+    const v = clampPercent(value);
+    const width = 8;
+    const filled = Math.round((v / 100) * width);
+    const full = "█".repeat(Math.max(0, Math.min(width, filled)));
+    const empty = "░".repeat(Math.max(0, width - filled));
+    return theme.fg(colorForPercent(v), full) + theme.fg("dim", empty);
+  }
+
+  function pickDataForProvider(provider: ProviderKey | null): UsageData | null {
+    if (!provider) return null;
+    return state[provider];
+  }
+
+  function updateStatus() {
+    const active = state.activeProvider;
+    const data = pickDataForProvider(active);
+
+    if (data && !data.error) {
+      pi.events.emit("usage:update", {
+        session: data.session,
+        weekly: data.weekly,
+        sessionResetsIn: data.sessionResetsIn,
+        weeklyResetsIn: data.weeklyResetsIn,
+      });
+    }
+
+    if (!ctx?.hasUI) return;
+
+    if (!active) {
+      ctx.ui.setStatus(STATUS_KEY, undefined);
+      return;
+    }
+
+    const auth = readAuth();
+    if (!canShowForProvider(active, auth, endpoints)) {
+      ctx.ui.setStatus(STATUS_KEY, undefined);
+      return;
+    }
+
+    const theme = ctx.ui.theme;
+    const label = PROVIDER_LABELS[active];
+
+    if (!data) {
+      ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", `${label} usage: loading…`));
+      return;
+    }
+
+    if (data.error) {
+      const cache = readUsageCache();
+      const blockedUntil = active ? (cache?.rateLimitedUntil?.[active] ?? 0) : 0;
+      const backoffNote = blockedUntil > Date.now()
+        ? ` — retry in ${Math.ceil((blockedUntil - Date.now()) / 60000)}m`
+        : "";
+      ctx.ui.setStatus(STATUS_KEY, theme.fg("warning", `${label} usage unavailable (${data.error}${backoffNote})`));
+      return;
+    }
+
+    const session = clampPercent(data.session);
+    const weekly = clampPercent(data.weekly);
+
+    const sessionReset = data.sessionResetsIn ? theme.fg("dim", ` ⟳ ${data.sessionResetsIn}`) : "";
+    const weeklyReset = data.weeklyResetsIn ? theme.fg("dim", ` ⟳ ${data.weeklyResetsIn}`) : "";
+
+    const status =
+      theme.fg("dim", `${label} `) +
+      theme.fg("muted", "S ") +
+      renderBar(theme, session) +
+      " " +
+      renderPercent(theme, session) +
+      sessionReset +
+      theme.fg("muted", " W ") +
+      renderBar(theme, weekly) +
+      " " +
+      renderPercent(theme, weekly) +
+      weeklyReset;
+
+    ctx.ui.setStatus(STATUS_KEY, status);
+  }
+
+  function updateProviderFrom(modelLike: any): boolean {
+    const previous = state.activeProvider;
+    state.activeProvider = detectProvider(modelLike);
+
+    if (previous !== state.activeProvider) {
+      updateStatus();
+      return true;
+    }
+
+    return false;
+  }
+
+  async function runPoll() {
+    const auth = readAuth();
+    const active = state.activeProvider;
+
+    if (!canShowForProvider(active, auth, endpoints) || !auth || !active) {
+      state.lastPoll = Date.now();
+      updateStatus();
+      return;
+    }
+
+    // --- Shared disk cache check ---
+    // All pi sessions read and write the same cache file so that only one
+    // process hits the API per CACHE_TTL_MS window, no matter how many
+    // sessions are open at once.
+    const cache = readUsageCache();
+    const now = Date.now();
+
+    // Respect cross-session rate-limit back-off written by any session.
+    const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
+    if (now < blockedUntil) {
+      // Use whatever data is in the cache so the bar still shows something.
+      if (cache?.data?.[active]) state[active] = cache.data[active]!;
+      state.lastPoll = now;
+      updateStatus();
+      return;
+    }
+
+    // If another session already polled recently, use their result.
+    if (cache && now - cache.timestamp < CACHE_TTL_MS && cache.data?.[active]) {
+      state[active] = cache.data[active]!;
+      state.lastPoll = now;
+      updateStatus();
+      return;
+    }
+
+    // --- Actually hit the API ---
+    // Skip independent token refresh — pi manages OAuth tokens and refreshes
+    // them in memory. A parallel refresh here would cause token rotation
+    // conflicts (Anthropic invalidates the old refresh token on use).
+    let result: UsageData;
+
+    if (active === "codex") {
+      const access = auth["openai-codex"]?.access;
+      result = access
+        ? await fetchCodexUsage(access)
+        : { session: 0, weekly: 0, error: "missing access token (try /login again)" };
+    } else if (active === "claude") {
+      const access = auth.anthropic?.access;
+      result = access
+        ? await fetchClaudeUsage(access)
+        : { session: 0, weekly: 0, error: "missing access token (try /login again)" };
+    } else if (active === "zai") {
+      const token = auth.zai?.access || auth.zai?.key;
+      result = token
+        ? await fetchZaiUsage(token, { endpoints })
+        : { session: 0, weekly: 0, error: "missing token (try /login again)" };
+    } else if (active === "gemini") {
+      const creds = auth["google-gemini-cli"];
+      result = creds?.access
+        ? await fetchGoogleUsage(creds.access, endpoints.gemini, creds.projectId, "gemini", { endpoints })
+        : { session: 0, weekly: 0, error: "missing access token (try /login again)" };
+    } else {
+      const creds = auth["google-antigravity"];
+      result = creds?.access
+        ? await fetchGoogleUsage(creds.access, endpoints.antigravity, creds.projectId, "antigravity", { endpoints })
+        : { session: 0, weekly: 0, error: "missing access token (try /login again)" };
+    }
+
+    state[active] = result;
+
+    // Write result + rate-limit state to shared cache so other sessions
+    // (and our own next timer tick) don't need to re-hit the API.
+    const nextCache: import("./core").UsageCache = {
+      timestamp: now,
+      data: { ...(cache?.data ?? {}), [active]: result },
+      rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) },
+    };
+    if (result.error === "HTTP 429") {
+      nextCache.rateLimitedUntil![active] = now + RATE_LIMITED_BACKOFF_MS;
+    } else {
+      delete nextCache.rateLimitedUntil![active];
+    }
+    writeUsageCache(nextCache);
+
+    state.lastPoll = now;
+    updateStatus();
+  }
+
+  async function poll() {
+    if (pollInFlight) {
+      pollQueued = true;
+      await pollInFlight;
+      return;
+    }
+
+    do {
+      pollQueued = false;
+      pollInFlight = runPoll()
+        .catch(() => {
+          // Never crash extension event handlers on transient polling errors.
+        })
+        .finally(() => {
+          pollInFlight = null;
+        });
+
+      await pollInFlight;
+    } while (pollQueued);
+  }
+
+  pi.on("session_start", async (_event, _ctx) => {
+    ctx = _ctx;
+    updateProviderFrom(_ctx.model);
+    await poll();
+  });
+
+  pi.on("session_shutdown", async (_event, _ctx) => {
+    if (_ctx?.hasUI) {
+      _ctx.ui.setStatus(STATUS_KEY, undefined);
+    }
+  });
+
+  // Refresh usage on every turn (like claude-pulse's UserPromptSubmit hook).
+  // The disk cache means the API is only hit at most once per CACHE_TTL_MS
+  // regardless of how many turns or sessions are active.
+  pi.on("turn_start", async (_event, _ctx) => {
+    ctx = _ctx;
+    updateProviderFrom(_ctx.model);
+    await poll();
+  });
+
+  pi.on("model_select", async (event, _ctx) => {
+    ctx = _ctx;
+    const changed = updateProviderFrom(event.model ?? _ctx.model);
+    if (changed) await poll();
+  });
+
+  pi.registerCommand("usage", {
+    description: "Show API usage for all subscriptions",
+    handler: async (_args, _ctx) => {
+      ctx = _ctx;
+      updateProviderFrom(_ctx.model);
+
+      try {
+        if (_ctx?.hasUI) {
+          await _ctx.ui.custom<void>((tui, theme, _keybindings, done) => {
+            const selector = new UsageSelectorComponent(
+              tui,
+              theme,
+              state.activeProvider,
+              () => fetchAllUsages({ endpoints }),
+              () => done(),
+            );
+            return selector;
+          });
+        }
+      } finally {
+        await poll();
+      }
+    },
+  });
+
+  // Re-poll immediately when the Claude account is switched via /switch-claude
+  pi.events.on("claude-account:switched", () => {
+    void poll();
+  });
+}